From dcb16d86f690c86876dd4b24889ef243ff5de0c8 Mon Sep 17 00:00:00 2001 From: can1357 Date: Fri, 19 Jun 2026 06:57:53 +0200 Subject: [PATCH] perf(compaction): kept per-turn tool-result pruning inside the warm prompt-cache prefix - Added `keepBoundaryId` and `cacheWarmSuffixTokens` guards to `pruneToolOutputs` and `pruneSupersededToolResults` so superseded/useless results sitting in the already-sent cached prefix are no longer rewritten mid-session, with new `computeMessageSuffixTokens`/`resolveBoundaryIndex` helpers. - Added a `keepBoundaryId` floor to `collectShakeRegions` so shake skips entries summarized away by the latest compaction. - Threaded `firstKeptEntryId`, `PRUNE_CACHE_WARM_SUFFIX_TOKENS` (8k) and `PRUNE_IDLE_FLUSH_MS` (90m, above the 1h cache TTL) from `#pruneToolOutputs`, `#pruneStaleToolResults` and `shake` in `agent-session.ts`. - Added six boundary tests in `supersede-prune.test.ts` covering warm-prefix protection, tail-case pruning, and the pre-boundary floor. --- packages/agent/src/compaction/pruning.ts | 112 +++++++++++++-- packages/agent/src/compaction/shake.ts | 19 +++ packages/agent/test/supersede-prune.test.ts | 133 ++++++++++++++++++ .../coding-agent/src/session/agent-session.ts | 37 ++++- 4 files changed, 284 insertions(+), 17 deletions(-) diff --git a/packages/agent/src/compaction/pruning.ts b/packages/agent/src/compaction/pruning.ts index 58c3dccca..1ba9a7ac0 100644 --- a/packages/agent/src/compaction/pruning.ts +++ b/packages/agent/src/compaction/pruning.ts @@ -30,6 +30,24 @@ export interface PruneConfig { supersedeKey?: SupersedeKeyFn; /** Useless-flagged results bypass the protect window (see {@link USELESS_NOTICE}). Default true. */ pruneUseless?: boolean; + /** + * Compaction boundary: the `firstKeptEntryId` of the latest compaction on + * the branch. Entries at indices BEFORE this id are summarized away and never + * sent to the model, so mutating them only churns persisted history without + * shrinking the prompt — they are skipped. Undefined = no compaction (the + * whole branch is sent). + */ + keepBoundaryId?: string; + /** + * Prompt-cache guard. When set, a tool result whose all-message suffix + * (tokens of every message after it) EXCEEDS this is part of the warm, + * already-sent cache prefix: mutating it forces the provider to re-write the + * whole suffix (cacheWrite premium). Such results — including superseded and + * useless ones, which otherwise bypass {@link protectTokens} — are left for + * compaction/shake (which rebuild the cache anyway) to reclaim. Undefined = + * no cache guard (legacy: superseded/useless prune at any depth). + */ + cacheWarmSuffixTokens?: number; } export const DEFAULT_PRUNE_CONFIG: PruneConfig = { @@ -66,10 +84,22 @@ export interface SupersedePruneConfig { pruneUseless?: boolean; /** Prune a candidate now when all messages after it total at most this many estimated tokens. Default 8 000. */ suffixTokenLimit?: number; - /** Prune all candidates when the last message is at least this old (prompt cache is cold anyway). Default 30 min. */ + /** + * Prune all candidates when the last message is at least this old: the + * provider prompt cache is then cold, so re-writing it is free. MUST exceed + * the cache retention (Anthropic "long" = 1h) or a still-warm prefix is busted + * by the flush. Default 30 min — callers on long retention override it. + */ idleFlushMs?: number; /** Clock override for tests. */ now?: number; + /** + * Compaction boundary (`firstKeptEntryId` of the latest compaction). Entries + * before it are summarized away and never sent, so they are skipped in every + * path — including the idle flush — to avoid pointless history churn. + * Undefined = no compaction (the whole branch is sent). + */ + keepBoundaryId?: string; /** Tool-result protection matchers (same contract as {@link PruneConfig.protectedTools}). */ protectedTools: ProtectedToolMatcher[]; } @@ -103,6 +133,35 @@ function estimatePrunedSavings(tokens: number, notice: string): number { return Math.max(0, tokens - noticeTokens); } +/** + * For each entry index, the estimated token total of all *message* entries + * strictly after it — how much prompt-cache content the provider must re-write + * (cacheWrite premium) if that entry is mutated in place. Used to keep prune + * mutations inside the cheap-to-recache tail. + */ +function computeMessageSuffixTokens(entries: readonly SessionEntry[]): number[] { + const suffix = new Array(entries.length); + let accumulated = 0; + for (let i = entries.length - 1; i >= 0; i--) { + suffix[i] = accumulated; + const entry = entries[i]; + if (entry.type === "message") accumulated += estimateTokens(entry.message as AgentMessage); + } + return suffix; +} + +/** + * Resolve the array index of the compaction boundary (`keepBoundaryId`). Entries + * before this index are summarized away by the latest compaction and never sent, + * so prune passes must not mutate them. Returns 0 when there is no boundary (no + * compaction → whole branch is sent) or the id is absent from `entries`. + */ +function resolveBoundaryIndex(entries: readonly SessionEntry[], keepBoundaryId: string | undefined): number { + if (keepBoundaryId === undefined) return 0; + const index = entries.findIndex(entry => entry.id === keepBoundaryId); + return index < 0 ? 0 : index; +} + interface SupersedeCandidate { entry: SessionMessageEntry; message: ToolResultMessage; @@ -183,7 +242,8 @@ function collectUselessResults( * flagged contextually useless. Cheap, incremental, and prompt-cache-aware: a * candidate is pruned now only when the suffix after it is small (tail case — * the read→edit→read loop) or when the context has been idle long enough that - * the provider cache is cold anyway (then ALL candidates flush). + * the provider cache is cold anyway (then all still-sent candidates flush). + * Never mutates entries before `keepBoundaryId` (summarized away — not sent). */ export function pruneSupersededToolResults(entries: SessionEntry[], config: SupersedePruneConfig): PruneResult { const toolCallsById = collectToolCallsById(entries); @@ -209,20 +269,24 @@ export function pruneSupersededToolResults(entries: SessionEntry[], config: Supe const idle = lastMessageTimestamp !== undefined && now - lastMessageTimestamp >= (config.idleFlushMs ?? DEFAULT_IDLE_FLUSH_MS); + const boundaryIndex = resolveBoundaryIndex(entries, config.keepBoundaryId); + let toPrune: SupersedeCandidate[]; if (idle) { - toPrune = candidates; + // Provider cache is cold (idle exceeds the retention TTL), so re-writing + // the sent region costs nothing. Entries before the compaction boundary + // are summarized away and never sent — skip them to avoid pointless churn. + toPrune = candidates.filter(candidate => candidate.index >= boundaryIndex); } else { const suffixTokenLimit = config.suffixTokenLimit ?? DEFAULT_SUFFIX_TOKEN_LIMIT; // suffixTokens[i] = estimated tokens of all messages strictly after entry i. - const suffixTokens = new Array(entries.length); - let accumulated = 0; - for (let i = entries.length - 1; i >= 0; i--) { - suffixTokens[i] = accumulated; - const entry = entries[i]; - if (entry.type === "message") accumulated += estimateTokens(entry.message as AgentMessage); - } - toPrune = candidates.filter(candidate => suffixTokens[candidate.index] <= suffixTokenLimit); + // Mutating a candidate re-writes its suffix in the warm cache, so prune only + // when that suffix is small (cheap-to-recache tail) and the candidate sits + // at/after the compaction boundary. + const suffixTokens = computeMessageSuffixTokens(entries); + toPrune = candidates.filter( + candidate => candidate.index >= boundaryIndex && suffixTokens[candidate.index] <= suffixTokenLimit, + ); } if (toPrune.length === 0) return { prunedCount: 0, tokensSaved: 0 }; @@ -262,6 +326,11 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = ) : undefined; + const boundaryIndex = resolveBoundaryIndex(entries, config.keepBoundaryId); + const cacheWarmSuffixTokens = config.cacheWarmSuffixTokens; + // All-message suffix per index, only when the cache guard is armed. + const messageSuffix = cacheWarmSuffixTokens === undefined ? undefined : computeMessageSuffixTokens(entries); + for (let i = entries.length - 1; i >= 0; i--) { const entry = entries[i]; const message = getToolResultMessage(entry); @@ -275,10 +344,23 @@ export function pruneToolOutputs(entries: SessionEntry[], config: PruneConfig = continue; } - // Superseded and useless results are pruned first: they bypass the - // protect window (a stale copy of re-read content — or a result the - // tool itself flagged as carrying no information — is dead weight at - // any age). + // Prompt-cache guard: a result whose all-message suffix exceeds the + // warm-cache window sits in the already-sent cached prefix — mutating it + // re-writes the whole suffix (cacheWrite premium). Entries before the + // compaction boundary are summarized away (never sent). Both are skipped + // before any prune decision, so superseded/useless cannot reach a deep, + // still-cached copy; compaction/shake reclaim those when they rebuild. + const inWarmPrefix = + messageSuffix !== undefined && cacheWarmSuffixTokens !== undefined && messageSuffix[i] > cacheWarmSuffixTokens; + if (inWarmPrefix || i < boundaryIndex) { + accumulatedTokens += tokens; + continue; + } + + // Superseded and useless results bypass the age-based protect window + // (a stale re-read copy, or a result the tool flagged as uninformative, + // is dead weight at any age) — but only within the cache-warm tail: the + // guard above already excluded deeper, still-cached copies. const superseded = supersededMessages?.has(message) ?? false; const useless = uselessMessages?.has(message) ?? false; const tooSmall = tokens < MIN_PRUNE_TOKENS; diff --git a/packages/agent/src/compaction/shake.ts b/packages/agent/src/compaction/shake.ts index f4535318d..7e0f2c5ed 100644 --- a/packages/agent/src/compaction/shake.ts +++ b/packages/agent/src/compaction/shake.ts @@ -31,6 +31,14 @@ export interface ShakeConfig { protectedTools: ProtectedToolMatcher[]; /** Minimum token size for a fenced/XML block to be eligible. */ fenceMinTokens: number; + /** + * Compaction boundary (`firstKeptEntryId` of the latest compaction). Entries + * before it are summarized away and never sent, so they are skipped — shaking + * them only churns persisted history. Undefined = no compaction (whole branch + * is sent). Note: shake still elides the warm cached prefix at/after the + * boundary — that is its job as a compaction-class reducer. + */ + keepBoundaryId?: string; } /** Auto-shake config: protects the live tail, conservative thresholds. */ @@ -289,9 +297,20 @@ export function collectShakeRegions(entries: SessionEntry[], config: ShakeConfig const toolCallsById = collectToolCallsById(entries); + // Entries before the compaction boundary are summarized away and never sent — + // shaking them only churns persisted history (no prompt/cache effect). + const boundaryIndex = + config.keepBoundaryId === undefined + ? 0 + : Math.max( + 0, + entries.findIndex(entry => entry.id === config.keepBoundaryId), + ); + const regions: ShakeRegion[] = []; for (let i = 0; i < n; i++) { const entry = entries[i]; + if (i < boundaryIndex) continue; const toolResult = getToolResultMessage(entry); // Useless-flagged results carry no information once consumed; they are // eligible even inside the protect-recent window. diff --git a/packages/agent/test/supersede-prune.test.ts b/packages/agent/test/supersede-prune.test.ts index f71f63526..dcc2c2b95 100644 --- a/packages/agent/test/supersede-prune.test.ts +++ b/packages/agent/test/supersede-prune.test.ts @@ -522,3 +522,136 @@ describe("pruneToolOutputs — small-result floor", () => { expect(resultMessage(bigResult).prunedAt).toBeDefined(); }); }); + +describe("cache-stable boundary — warm prefix protection", () => { + // (a) The primary bug: in pruneToolOutputs a superseded result bypasses the + // protect window and is rewritten at any depth. With the cache guard armed it + // must be left alone when it sits in the warm, already-sent cached prefix. + test("(a) deep superseded result is rewritten WITHOUT the guard but kept WITH it", () => { + const build = (): { entries: SessionEntry[]; result1: SessionMessageEntry; result2: SessionMessageEntry } => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const big = textEntry(BIG_TEXT, T0 + 500); // pushes result1 deep into the suffix + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); // tail, supersedes result1 + return { entries: [call1, result1, big, call2, result2], result1, result2 }; + }; + const base = { + protectTokens: 1_000_000, // everything inside the (age) protect window + minimumSavings: 0, + protectedTools: [], + supersedeKey: readToolSupersedeKey, + }; + + // Legacy (no cacheWarmSuffixTokens): superseded result1 bypasses the window -> pruned. + const legacy = build(); + const legacyRun = pruneToolOutputs(legacy.entries, base); + expect(legacyRun.prunedCount).toBe(1); + expect(resultText(legacy.result1)).toBe(SUPERSEDED_NOTICE); + + // Guard armed: result1's all-message suffix (BIG_TEXT + call2 + result2) far + // exceeds the window, so it is part of the warm cached prefix and is kept. + const guarded = build(); + const guardedRun = pruneToolOutputs(guarded.entries, { ...base, cacheWarmSuffixTokens: 200 }); + expect(guardedRun.prunedCount).toBe(0); + expect(resultText(guarded.result1)).toBe(FILE_CONTENT); + expect(resultMessage(guarded.result1).prunedAt).toBeUndefined(); + }); + + test("(a) deep useless result is kept when the cache guard is armed", () => { + const [call1, result1] = uselessPair("search", NO_MATCH_TEXT, T0); + const big = textEntry(BIG_TEXT, T0 + 500); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const entries: SessionEntry[] = [call1, result1, big, call2, result2]; + + const result = pruneToolOutputs(entries, { + protectTokens: 1_000_000, + minimumSavings: 0, + protectedTools: [], + pruneUseless: true, + cacheWarmSuffixTokens: 200, + }); + + expect(result.prunedCount).toBe(0); + expect(resultText(result1)).toBe(NO_MATCH_TEXT); + expect(resultMessage(result1).prunedAt).toBeUndefined(); + expect(resultText(result2)).toBe(FILE_CONTENT); + }); + + // (b) The legit case must still fire: a superseded copy in the cheap-to-recache + // tail (suffix below the window) is still reclaimed. + test("(b) tail-case superseded result still prunes with the guard armed", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const entries: SessionEntry[] = [call1, result1, call2, result2]; + + const result = pruneToolOutputs(entries, { + protectTokens: 1_000_000, + minimumSavings: 0, + protectedTools: [], + supersedeKey: readToolSupersedeKey, + cacheWarmSuffixTokens: 100_000, // result1's suffix is far below this -> tail -> prunable + }); + + expect(result.prunedCount).toBe(1); + expect(resultText(result1)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result2)).toBe(FILE_CONTENT); + }); + + test("(b) supersede pass still prunes the tail case with keepBoundaryId set", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); + const entries: SessionEntry[] = [call1, result1, call2, result2]; + + const result = pruneSupersededToolResults(entries, cfg({ keepBoundaryId: call1.id, now: T0 + 1_000 })); + + expect(result.prunedCount).toBe(1); + expect(resultText(result1)).toBe(SUPERSEDED_NOTICE); + expect(resultText(result2)).toBe(FILE_CONTENT); + }); + + // (c) Entries before firstKeptEntryId are summarized away — never sent — so no + // pass may mutate them, not even the idle full-flush. + test("(c) idle flush never mutates entries before keepBoundaryId", () => { + const [call1, result1] = readPair("src/foo.ts", FILE_CONTENT, T0); // idx 0,1 — before boundary + const [call2, result2] = readPair("src/foo.ts", FILE_CONTENT, T0 + 1_000); // idx 2,3 — boundary at call2 + const [call3, result3] = readPair("src/foo.ts", FILE_CONTENT, T0 + 2_000); // idx 4,5 — latest + const big = textEntry(BIG_TEXT, T0 + 3_000); + const entries: SessionEntry[] = [call1, result1, call2, result2, call3, result3, big]; + + // Cold cache (idle > threshold) with suffixTokenLimit 0: only the idle path can fire. + const result = pruneSupersededToolResults( + entries, + cfg({ + keepBoundaryId: call2.id, + suffixTokenLimit: 0, + idleFlushMs: 30 * 60_000, + now: T0 + 3_000 + 31 * 60_000, + }), + ); + + expect(result.prunedCount).toBe(1); + expect(resultText(result1)).toBe(FILE_CONTENT); // before boundary -> untouched + expect(resultMessage(result1).prunedAt).toBeUndefined(); + expect(resultText(result2)).toBe(SUPERSEDED_NOTICE); // at/after boundary -> flushed + expect(resultText(result3)).toBe(FILE_CONTENT); // latest -> kept + }); + + test("(c) pruneToolOutputs never mutates entries before keepBoundaryId", () => { + const [call1, result1] = readPair("src/old.ts", FILE_CONTENT, T0); // idx 0,1 — before boundary + const [call2, result2] = readPair("src/new.ts", FILE_CONTENT, T0 + 1_000); // idx 2,3 — boundary at call2 + const entries: SessionEntry[] = [call1, result1, call2, result2]; + + // protectTokens 0 -> the age path would prune both; the window is wide so the + // guard does not protect either; only keepBoundaryId shields result1. + pruneToolOutputs(entries, { + protectTokens: 0, + minimumSavings: 0, + protectedTools: [], + keepBoundaryId: call2.id, + cacheWarmSuffixTokens: 1_000_000, + }); + + expect(resultText(result1)).toBe(FILE_CONTENT); // before boundary -> untouched + expect(resultMessage(result1).prunedAt).toBeUndefined(); + expect(resultMessage(result2).prunedAt).toBeDefined(); // at/after boundary, in tail -> pruned + }); +}); diff --git a/packages/coding-agent/src/session/agent-session.ts b/packages/coding-agent/src/session/agent-session.ts index 773e8b801..299010d27 100644 --- a/packages/coding-agent/src/session/agent-session.ts +++ b/packages/coding-agent/src/session/agent-session.ts @@ -368,6 +368,23 @@ const COMPACTION_CHECK_CONTINUATION: CompactionCheckResult = { deferredHandoff: false, continuationScheduled: true, }; + +/** + * Per-turn prune cache window. A tool result whose all-message suffix exceeds + * this is in the warm, already-sent prompt-cache prefix: re-writing it costs the + * cacheWrite premium on the whole suffix. Per-turn passes only reclaim inside + * this tail (matches the supersede pass's default `suffixTokenLimit`); deeper + * stale/age victims are left to compaction/shake, which rebuild the cache anyway. + */ +const PRUNE_CACHE_WARM_SUFFIX_TOKENS = 8_000; + +/** + * Idle gap after which the supersede pass may flush the whole sent region (the + * provider cache is cold, so re-writing it is free). MUST exceed the maximum + * Anthropic prompt-cache TTL — "long" retention (the OAuth default) is 1h — or a + * still-warm prefix is busted by the flush. 90 min leaves margin over the 1h TTL. + */ +const PRUNE_IDLE_FLUSH_MS = 90 * 60_000; export type CommandMetadataChangedListener = () => void | Promise; export type AsyncJobSnapshotItem = Pick; @@ -7309,11 +7326,16 @@ export class AgentSession { async #pruneToolOutputs(): Promise<{ prunedCount: number; tokensSaved: number } | undefined> { const branchEntries = this.sessionManager.getBranch(); + const keepBoundaryId = getLatestCompactionEntry(branchEntries)?.firstKeptEntryId; const result = pruneToolOutputs( branchEntries, this.#withPlanProtection({ ...DEFAULT_PRUNE_CONFIG, pruneUseless: this.settings.getGroup("compaction").dropUseless, + // Cache-stable boundary: never re-write the warm, already-sent prefix + // (deep stale/age victims) or summarized-away entries every turn. + keepBoundaryId, + cacheWarmSuffixTokens: PRUNE_CACHE_WARM_SUFFIX_TOKENS, }), ); if (result.prunedCount === 0) { @@ -7341,12 +7363,17 @@ export class AgentSession { const { supersedeReads, dropUseless } = this.settings.getGroup("compaction"); if (!supersedeReads && !dropUseless) return undefined; const branchEntries = this.sessionManager.getBranch(); + const keepBoundaryId = getLatestCompactionEntry(branchEntries)?.firstKeptEntryId; const result = pruneSupersededToolResults( branchEntries, this.#withPlanProtection({ supersedeKey: supersedeReads ? readToolSupersedeKey : undefined, pruneUseless: dropUseless, protectedTools: [...DEFAULT_PRUNE_CONFIG.protectedTools], + // Never re-write summarized-away entries; only flush the whole sent + // region once the cache is genuinely cold (idle exceeds the 1h TTL). + keepBoundaryId, + idleFlushMs: PRUNE_IDLE_FLUSH_MS, }), ); if (result.prunedCount === 0) { @@ -7430,8 +7457,14 @@ export class AgentSession { return { mode, toolResultsDropped: 0, blocksDropped: 0, imagesDropped: removed, tokensFreed: 0 }; } - const config = this.#withPlanProtection(opts.config ?? AGGRESSIVE_SHAKE_CONFIG); - const regions = collectShakeRegions(this.sessionManager.getBranch(), config); + const branchEntries = this.sessionManager.getBranch(); + const config = this.#withPlanProtection({ + ...(opts.config ?? AGGRESSIVE_SHAKE_CONFIG), + // Skip entries summarized away by the latest compaction — shaking them + // only churns persisted history with no prompt/cache effect. + keepBoundaryId: getLatestCompactionEntry(branchEntries)?.firstKeptEntryId, + }); + const regions = collectShakeRegions(branchEntries, config); if (regions.length === 0) { return { mode, toolResultsDropped: 0, blocksDropped: 0, tokensFreed: 0 }; }