From caab3434612c835889bfff061935ec73ccd2430b Mon Sep 17 00:00:00 2001 From: Yuval Dinodia <102706514+yetval@users.noreply.github.com> Date: Mon, 15 Jun 2026 13:10:42 -0400 Subject: [PATCH] fix(reply): project preflight compaction gate by next-input on fresh tokens (#91488) The fresh-tokens path of runPreflightCompactionIfNeeded fed the prompt-only entry.totalTokens snapshot straight into the budget threshold check, dropping the current user prompt estimate and the previous turn's output. The sibling memory-flush gate and this function's own stale branch already project base + output + estimate via resolveEffectivePromptTokens, so the preflight gate under-triggered and let over-budget requests through to overflow-retry. Project the fresh persisted base the same way: read transcript output when near the threshold (mirroring the memory-flush gate's buffer) and run the fresh base through resolveEffectivePromptTokens before the threshold check. --- .../reply/agent-runner-memory.test.ts | 75 ++++++++++++++++++- src/auto-reply/reply/agent-runner-memory.ts | 34 ++++++--- 2 files changed, 98 insertions(+), 11 deletions(-) diff --git a/src/auto-reply/reply/agent-runner-memory.test.ts b/src/auto-reply/reply/agent-runner-memory.test.ts index 48c465e9d0a8..051a3ea7af8c 100644 --- a/src/auto-reply/reply/agent-runner-memory.test.ts +++ b/src/auto-reply/reply/agent-runner-memory.test.ts @@ -1361,6 +1361,79 @@ describe("runMemoryFlushIfNeeded", () => { expect(compactCall.authProfileId).toBe("anthropic:claude@martian.engineering"); expect(compactCall.contextTokenBudget).toBe(258_000); }); + it("preflight compacts a fresh session when the current prompt estimate pushes the next request over budget", async () => { + registerMemoryFlushPlanResolverForTest(() => ({ + softThresholdTokens: 0, + forceFlushTranscriptBytes: 1_000_000_000, + reserveTokensFloor: 10, + prompt: "Pre-compaction memory flush.\nNO_REPLY", + systemPrompt: "Write memory to memory/YYYY-MM-DD.md.", + relativePath: "memory/2023-11-14.md", + })); + const sessionEntry: SessionEntry = { + sessionId: "session", + updatedAt: Date.now(), + totalTokens: 985, + totalTokensFresh: true, + compactionCount: 0, + }; + + await runPreflightCompactionIfNeeded({ + cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } }, + followupRun: createTestFollowupRun({ + provider: "anthropic", + model: "claude", + sessionKey: "agent:main:main", + }), + promptForEstimate: "Please summarize the entire design discussion above. ".repeat(8), + defaultModel: "anthropic/claude", + agentCfgContextTokens: 1000, + sessionEntry, + sessionStore: { "agent:main:main": sessionEntry }, + sessionKey: "agent:main:main", + isHeartbeat: false, + replyOperation: createReplyOperation(), + }); + + expect(compactEmbeddedAgentSessionMock).toHaveBeenCalledTimes(1); + }); + it("does not preflight compact a fresh session when only accumulated output tokens are large and the latest output keeps the request under budget", async () => { + registerMemoryFlushPlanResolverForTest(() => ({ + softThresholdTokens: 0, + forceFlushTranscriptBytes: 1_000_000_000, + reserveTokensFloor: 10, + prompt: "Pre-compaction memory flush.\nNO_REPLY", + systemPrompt: "Write memory to memory/YYYY-MM-DD.md.", + relativePath: "memory/2023-11-14.md", + })); + const sessionEntry: SessionEntry = { + sessionId: "session", + updatedAt: Date.now(), + totalTokens: 985, + outputTokens: 50_000, + totalTokensFresh: true, + compactionCount: 0, + }; + + await runPreflightCompactionIfNeeded({ + cfg: { agents: { defaults: { compaction: { memoryFlush: {} } } } }, + followupRun: createTestFollowupRun({ + provider: "anthropic", + model: "claude", + sessionKey: "agent:main:main", + }), + promptForEstimate: "", + defaultModel: "anthropic/claude", + agentCfgContextTokens: 1000, + sessionEntry, + sessionStore: { "agent:main:main": sessionEntry }, + sessionKey: "agent:main:main", + isHeartbeat: false, + replyOperation: createReplyOperation(), + }); + + expect(compactEmbeddedAgentSessionMock).not.toHaveBeenCalled(); + }); it("updates the active preflight run after transcript rotation", async () => { const sessionFile = path.join(rootDir, "session.jsonl"); const successorFile = path.join(rootDir, "session-rotated.jsonl"); @@ -2113,7 +2186,7 @@ describe("runMemoryFlushIfNeeded", () => { const compactCall = requireCompactEmbeddedAgentSessionCall(); expect(compactCall.sessionId).toBe("session"); expect(compactCall.trigger).toBe("budget"); - expect(compactCall.currentTokenCount).toBe(10); + expect(compactCall.currentTokenCount).toBe(12); expect(compactCall.sessionFile).toContain("large-session.jsonl"); }); diff --git a/src/auto-reply/reply/agent-runner-memory.ts b/src/auto-reply/reply/agent-runner-memory.ts index f0959c11cc9b..271f7492f60d 100644 --- a/src/auto-reply/reply/agent-runner-memory.ts +++ b/src/auto-reply/reply/agent-runner-memory.ts @@ -779,10 +779,24 @@ export async function runPreflightCompactionIfNeeded(params: { const promptTokenEstimate = estimatePromptTokensForMemoryFlush( params.promptForEstimate ?? params.followupRun.prompt, ); + const serverCompactionThreshold = resolveResponsesServerCompactionThreshold({ + cfg: params.cfg, + provider: params.followupRun.run.provider, + modelId: params.followupRun.run.model ?? params.defaultModel, + }); + const threshold = Math.max( + contextWindowTokens - reserveTokensFloor - softThresholdTokens, + serverCompactionThreshold ?? 0, + ); + const freshNeedsOutputRead = + typeof freshPersistedTokens === "number" && + typeof promptTokenEstimate === "number" && + threshold > 0 && + freshPersistedTokens + promptTokenEstimate >= threshold - TRANSCRIPT_OUTPUT_READ_BUFFER_TOKENS; const maxActiveTranscriptBytes = resolveMaxActiveTranscriptBytes(params.cfg); const shouldCheckActiveTranscriptBytes = typeof maxActiveTranscriptBytes === "number"; const transcriptUsageTokens = - typeof freshPersistedTokens === "number" + typeof freshPersistedTokens === "number" && !freshNeedsOutputRead ? undefined : await estimatePromptTokensFromSessionTranscript({ sessionId: entry.sessionId, @@ -824,8 +838,17 @@ export async function runPreflightCompactionIfNeeded(params: { promptTokenEstimate, ) : undefined; + const freshProjectedTokenCount = + typeof freshPersistedTokens === "number" + ? resolveEffectivePromptTokens( + freshPersistedTokens, + transcriptOutputTokens, + promptTokenEstimate, + ) + : undefined; const projectedTokenCount = Math.max( usageProjectedTokenCount ?? 0, + freshProjectedTokenCount ?? 0, stalePersistedPromptTokens ?? 0, ); const tokenCountForCompaction = @@ -833,15 +856,6 @@ export async function runPreflightCompactionIfNeeded(params: { ? projectedTokenCount : undefined; - const serverCompactionThreshold = resolveResponsesServerCompactionThreshold({ - cfg: params.cfg, - provider: params.followupRun.run.provider, - modelId: params.followupRun.run.model ?? params.defaultModel, - }); - const threshold = Math.max( - contextWindowTokens - reserveTokensFloor - softThresholdTokens, - serverCompactionThreshold ?? 0, - ); logVerbose( `preflightCompaction check: sessionKey=${params.sessionKey} ` + `tokenCount=${tokenCountForCompaction ?? freshPersistedTokens ?? "undefined"} ` +