mirror of
https://github.com/getcompanion-ai/co-mono.git
synced 2026-04-15 11:02:17 +00:00
parent
7b96041068
commit
6b4b920425
2 changed files with 75 additions and 2 deletions
|
|
@ -232,6 +232,7 @@ export class AgentSession {
|
|||
// Compaction state
|
||||
private _compactionAbortController: AbortController | undefined = undefined;
|
||||
private _autoCompactionAbortController: AbortController | undefined = undefined;
|
||||
private _overflowRecoveryAttempted = false;
|
||||
|
||||
// Branch summarization state
|
||||
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
|
||||
// This ensures the UI sees the updated queue state
|
||||
if (event.type === "message_start" && event.message.role === "user") {
|
||||
this._overflowRecoveryAttempted = false;
|
||||
const messageText = this._getUserMessageText(event.message);
|
||||
if (messageText) {
|
||||
// Check steering queue first
|
||||
|
|
@ -415,9 +417,13 @@ export class AgentSession {
|
|||
if (event.message.role === "assistant") {
|
||||
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
|
||||
// This prevents accumulation across multiple LLM calls within a turn
|
||||
const assistantMsg = event.message as AssistantMessage;
|
||||
if (assistantMsg.stopReason !== "error" && this._retryAttempt > 0) {
|
||||
this._emit({
|
||||
type: "auto_retry_end",
|
||||
|
|
@ -1679,6 +1685,19 @@ export class AgentSession {
|
|||
|
||||
// Case 1: Overflow - LLM returned context overflow error
|
||||
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,
|
||||
// but we don't want it in context for the retry)
|
||||
const messages = this.agent.state.messages;
|
||||
|
|
|
|||
|
|
@ -2,7 +2,7 @@ import { existsSync, mkdirSync, rmSync } from "node:fs";
|
|||
import { tmpdir } from "node:os";
|
||||
import { join } from "node:path";
|
||||
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 { AgentSession } from "../src/core/agent-session.js";
|
||||
import { AuthStorage } from "../src/core/auth-storage.js";
|
||||
|
|
@ -94,4 +94,58 @@ describe("AgentSession auto-compaction queue resume", () => {
|
|||
|
||||
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