diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 98a41de7..b892d932 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -10,6 +10,7 @@ ### Fixed - Fixed auto-compaction retriggering on the first prompt after compaction due to stale pre-compaction assistant usage ([#1860](https://github.com/badlogic/pi-mono/issues/1860) by [@joelhooks](https://github.com/joelhooks)) +- Fixed sessions never auto-compacting when hitting persistent API errors (e.g. 529 overloaded) by estimating context size from the last successful response ([#1834](https://github.com/badlogic/pi-mono/issues/1834)) - Fixed `/new` leaving startup header content, including the changelog, visible after starting a fresh session ([#1880](https://github.com/badlogic/pi-mono/issues/1880)) - Fixed misleading docs and example implying that returning `{ isError: true }` from a tool's `execute` function marks the execution as failed; errors must be signaled by throwing ([#1881](https://github.com/badlogic/pi-mono/issues/1881)) - Fixed model switches through non-reasoning models to preserve the saved default thinking level instead of persisting a capability-forced `off` clamp ([#1864](https://github.com/badlogic/pi-mono/issues/1864)) diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index db587150..5efacdc6 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1731,11 +1731,17 @@ export class AgentSession { return; } - // Case 2: Threshold - turn succeeded but context is getting large - // Skip if this was an error (non-overflow errors don't have usage data) - if (assistantMessage.stopReason === "error") return; - - const contextTokens = calculateContextTokens(assistantMessage.usage); + // Case 2: Threshold - context is getting large + // For error messages (no usage data), estimate from last successful response. + // This ensures sessions that hit persistent API errors (e.g. 529) can still compact. + let contextTokens: number; + if (assistantMessage.stopReason === "error") { + const estimate = estimateContextTokens(this.agent.state.messages); + if (estimate.lastUsageIndex === null) return; // No usage data at all + contextTokens = estimate.tokens; + } else { + contextTokens = calculateContextTokens(assistantMessage.usage); + } if (shouldCompact(contextTokens, contextWindow, settings)) { await this._runAutoCompaction("threshold", false); } diff --git a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts index f6bfcb3b..48381cc8 100644 --- a/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts +++ b/packages/coding-agent/test/agent-session-auto-compaction-queue.test.ts @@ -26,7 +26,24 @@ vi.mock("../src/core/compaction/index.js", () => ({ tokensBefore: 100, details: {}, }), - estimateContextTokens: () => ({ tokens: 0, usageTokens: 0, trailingTokens: 0, lastUsageIndex: -1 }), + estimateContextTokens: ( + messages: Array<{ + role: string; + usage?: { input: number; output: number; cacheRead: number; cacheWrite: number; totalTokens?: number }; + stopReason?: string; + }>, + ) => { + // Walk backwards to find last non-error, non-aborted assistant with usage + for (let i = messages.length - 1; i >= 0; i--) { + const msg = messages[i]; + if (msg.role === "assistant" && msg.stopReason !== "error" && msg.stopReason !== "aborted" && msg.usage) { + const tokens = + msg.usage.totalTokens ?? msg.usage.input + msg.usage.output + msg.usage.cacheRead + msg.usage.cacheWrite; + return { tokens, usageTokens: tokens, trailingTokens: 0, lastUsageIndex: i }; + } + } + return { tokens: 0, usageTokens: 0, trailingTokens: 0, lastUsageIndex: null }; + }, generateBranchSummary: async () => ({ summary: "", aborted: false, readFiles: [], modifiedFiles: [] }), prepareCompaction: () => ({ dummy: true }), shouldCompact: ( @@ -216,4 +233,122 @@ describe("AgentSession auto-compaction queue resume", () => { expect(runAutoCompactionSpy).not.toHaveBeenCalled(); }); + + it("should trigger threshold compaction for error messages using last successful usage", async () => { + const model = session.model!; + + // A successful assistant message with high token usage (near context limit) + const successfulAssistant: AssistantMessage = { + role: "assistant", + content: [{ type: "text", text: "large successful response" }], + api: model.api, + provider: model.provider, + model: model.id, + usage: { + input: 180_000, + output: 10_000, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 190_000, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason: "stop", + timestamp: Date.now(), + }; + + // An error message (e.g. 529 overloaded) with no useful usage data + const errorAssistant: 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: "529 overloaded", + timestamp: Date.now() + 1000, + }; + + // Put both messages into agent state so estimateContextTokens can find the successful one + session.agent.replaceMessages([ + { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, + successfulAssistant, + { role: "user", content: [{ type: "text", text: "another prompt" }], timestamp: Date.now() + 500 }, + errorAssistant, + ]); + + const runAutoCompactionSpy = vi + .spyOn( + session as unknown as { + _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; + }, + "_runAutoCompaction", + ) + .mockResolvedValue(); + + const checkCompaction = ( + session as unknown as { + _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise; + } + )._checkCompaction.bind(session); + + await checkCompaction(errorAssistant); + + expect(runAutoCompactionSpy).toHaveBeenCalledWith("threshold", false); + }); + + it("should not trigger threshold compaction for error messages when no prior usage exists", async () => { + const model = session.model!; + + // An error message with no prior successful assistant in context + const errorAssistant: 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: "529 overloaded", + timestamp: Date.now(), + }; + + session.agent.replaceMessages([ + { role: "user", content: [{ type: "text", text: "hello" }], timestamp: Date.now() - 1000 }, + errorAssistant, + ]); + + const runAutoCompactionSpy = vi + .spyOn( + session as unknown as { + _runAutoCompaction: (reason: "overflow" | "threshold", willRetry: boolean) => Promise; + }, + "_runAutoCompaction", + ) + .mockResolvedValue(); + + const checkCompaction = ( + session as unknown as { + _checkCompaction: (assistantMessage: AssistantMessage, skipAbortedCheck?: boolean) => Promise; + } + )._checkCompaction.bind(session); + + await checkCompaction(errorAssistant); + + expect(runAutoCompactionSpy).not.toHaveBeenCalled(); + }); });