mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-20 02:03:16 +00:00
parent
7b96041068
commit
6b4b920425
2 changed files with 75 additions and 2 deletions
|
|
@ -232,6 +232,7 @@ export class AgentSession {
|
||||||
// Compaction state
|
// Compaction state
|
||||||
private _compactionAbortController: AbortController | undefined = undefined;
|
private _compactionAbortController: AbortController | undefined = undefined;
|
||||||
private _autoCompactionAbortController: AbortController | undefined = undefined;
|
private _autoCompactionAbortController: AbortController | undefined = undefined;
|
||||||
|
private _overflowRecoveryAttempted = false;
|
||||||
|
|
||||||
// Branch summarization state
|
// Branch summarization state
|
||||||
private _branchSummaryAbortController: AbortController | undefined = undefined;
|
private _branchSummaryAbortController: AbortController | undefined = undefined;
|
||||||
|
|
@ -368,6 +369,7 @@ export class AgentSession {
|
||||||
// When a user message starts, check if it's from either queue and remove it BEFORE emitting
|
// When a user message starts, check if it's from either queue and remove it BEFORE emitting
|
||||||
// This ensures the UI sees the updated queue state
|
// This ensures the UI sees the updated queue state
|
||||||
if (event.type === "message_start" && event.message.role === "user") {
|
if (event.type === "message_start" && event.message.role === "user") {
|
||||||
|
this._overflowRecoveryAttempted = false;
|
||||||
const messageText = this._getUserMessageText(event.message);
|
const messageText = this._getUserMessageText(event.message);
|
||||||
if (messageText) {
|
if (messageText) {
|
||||||
// Check steering queue first
|
// Check steering queue first
|
||||||
|
|
@ -415,9 +417,13 @@ export class AgentSession {
|
||||||
if (event.message.role === "assistant") {
|
if (event.message.role === "assistant") {
|
||||||
this._lastAssistantMessage = event.message;
|
this._lastAssistantMessage = event.message;
|
||||||
|
|
||||||
|
const assistantMsg = event.message as AssistantMessage;
|
||||||
|
if (assistantMsg.stopReason !== "error") {
|
||||||
|
this._overflowRecoveryAttempted = false;
|
||||||
|
}
|
||||||
|
|
||||||
// Reset retry counter immediately on successful assistant response
|
// Reset retry counter immediately on successful assistant response
|
||||||
// This prevents accumulation across multiple LLM calls within a turn
|
// This prevents accumulation across multiple LLM calls within a turn
|
||||||
const assistantMsg = event.message as AssistantMessage;
|
|
||||||
if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
|
if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
|
||||||
this._emit({
|
this._emit({
|
||||||
type: "auto_retry_end",
|
type: "auto_retry_end",
|
||||||
|
|
@ -1679,6 +1685,19 @@ export class AgentSession {
|
||||||
|
|
||||||
// Case 1: Overflow - LLM returned context overflow error
|
// Case 1: Overflow - LLM returned context overflow error
|
||||||
if (sameModel && !errorIsFromBeforeCompaction && isContextOverflow(assistantMessage, contextWindow)) {
|
if (sameModel && !errorIsFromBeforeCompaction && isContextOverflow(assistantMessage, contextWindow)) {
|
||||||
|
if (this._overflowRecoveryAttempted) {
|
||||||
|
this._emit({
|
||||||
|
type: "auto_compaction_end",
|
||||||
|
result: undefined,
|
||||||
|
aborted: false,
|
||||||
|
willRetry: false,
|
||||||
|
errorMessage:
|
||||||
|
"Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.",
|
||||||
|
});
|
||||||
|
return;
|
||||||
|
}
|
||||||
|
|
||||||
|
this._overflowRecoveryAttempted = true;
|
||||||
// Remove the error message from agent state (it IS saved to session for history,
|
// Remove the error message from agent state (it IS saved to session for history,
|
||||||
// but we don't want it in context for the retry)
|
// but we don't want it in context for the retry)
|
||||||
const messages = this.agent.state.messages;
|
const messages = this.agent.state.messages;
|
||||||
|
|
|
||||||
|
|
@ -2,7 +2,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs";
|
||||||
import { tmpdir } from "node:os";
|
import { tmpdir } from "node:os";
|
||||||
import { join } from "node:path";
|
import { join } from "node:path";
|
||||||
import { Agent } from "@mariozechner/pi-agent-core";
|
import { Agent } from "@mariozechner/pi-agent-core";
|
||||||
import { getModel } from "@mariozechner/pi-ai";
|
import { type AssistantMessage, getModel } from "@mariozechner/pi-ai";
|
||||||
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
import { afterEach, beforeEach, describe, expect, it, vi } from "vitest";
|
||||||
import { AgentSession } from "../src/core/agent-session.js";
|
import { AgentSession } from "../src/core/agent-session.js";
|
||||||
import { AuthStorage } from "../src/core/auth-storage.js";
|
import { AuthStorage } from "../src/core/auth-storage.js";
|
||||||
|
|
@ -94,4 +94,58 @@ describe("AgentSession auto-compaction queue resume", () => {
|
||||||
|
|
||||||
expect(continueSpy).toHaveBeenCalledTimes(1);
|
expect(continueSpy).toHaveBeenCalledTimes(1);
|
||||||
});
|
});
|
||||||
|
|
||||||
|
it("should not compact repeatedly after overflow recovery already attempted", async () => {
|
||||||
|
const model = session.model!;
|
||||||
|
const overflowMessage: AssistantMessage = {
|
||||||
|
role: "assistant",
|
||||||
|
content: [{ type: "text", text: "" }],
|
||||||
|
api: model.api,
|
||||||
|
provider: model.provider,
|
||||||
|
model: model.id,
|
||||||
|
usage: {
|
||||||
|
input: 0,
|
||||||
|
output: 0,
|
||||||
|
cacheRead: 0,
|
||||||
|
cacheWrite: 0,
|
||||||
|
totalTokens: 0,
|
||||||
|
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
|
||||||
|
},
|
||||||
|
stopReason: "error",
|
||||||
|
errorMessage: "prompt is too long",
|
||||||
|
timestamp: Date.now(),
|
||||||
|
};
|
||||||
|
|
||||||
|
const runAutoCompactionSpy = vi
|
||||||
|
.spyOn(
|
||||||
|
session as unknown as {
|
||||||
|
_runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise<void>;
|
||||||
|
},
|
||||||
|
"_runAutoCompaction",
|
||||||
|
)
|
||||||
|
.mockResolvedValue();
|
||||||
|
|
||||||
|
const events: Array<{ type: string; errorMessage?: string }> = [];
|
||||||
|
session.subscribe((event) => {
|
||||||
|
if (event.type === "auto_compaction_end") {
|
||||||
|
events.push({ type: event.type, errorMessage: event.errorMessage });
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
const checkCompaction = (
|
||||||
|
session as unknown as {
|
||||||
|
_checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise<void>;
|
||||||
|
}
|
||||||
|
)._checkCompaction.bind(session);
|
||||||
|
|
||||||
|
await checkCompaction(overflowMessage);
|
||||||
|
await checkCompaction({ ...overflowMessage, timestamp: Date.now() + 1 });
|
||||||
|
|
||||||
|
expect(runAutoCompactionSpy).toHaveBeenCalledTimes(1);
|
||||||
|
expect(events).toContainEqual({
|
||||||
|
type: "auto_compaction_end",
|
||||||
|
errorMessage:
|
||||||
|
"Context overflow recovery failed after one compact-and-retry attempt. Try reducing context or switching to a larger-context model.",
|
||||||
|
});
|
||||||
|
});
|
||||||
});
|
});
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue