From 626210c783107989f99f423cad01e7b9137466f7 Mon Sep 17 00:00:00 2001 From: Woonggi Min Date: Sat, 25 Jul 2026 00:40:03 +0900 Subject: [PATCH] feat(coding-agent): add staged context compaction --- packages/coding-agent/src/changes.md | 20 ++ .../coding-agent/src/core/agent-session.ts | 227 +++++++++++++++--- packages/coding-agent/src/core/changes.md | 18 ++ .../src/core/compaction/changes.md | 21 ++ .../src/core/compaction/compaction.ts | 221 +++++++++++++++-- .../extensions/builtin/compaction/changes.md | 10 + .../extensions/builtin/compaction/index.ts | 14 +- .../builtin/compaction/speculative.ts | 118 ++++----- .../src/core/extensions/changes.md | 17 ++ .../coding-agent/src/core/extensions/types.ts | 16 ++ .../compaction/speculative-compaction.test.ts | 54 ++--- .../staged-compaction-planner.test.ts | 120 +++++++++ .../agent-session-staged-compaction.test.ts | 213 ++++++++++++++++ 13 files changed, 915 insertions(+), 154 deletions(-) create mode 100644 packages/coding-agent/test/compaction/staged-compaction-planner.test.ts create mode 100644 packages/coding-agent/test/suite/agent-session-staged-compaction.test.ts diff --git a/packages/coding-agent/src/changes.md b/packages/coding-agent/src/changes.md index c673600d7..49b0b2edd 100644 --- a/packages/coding-agent/src/changes.md +++ b/packages/coding-agent/src/changes.md @@ -1,3 +1,23 @@ +## Staged context compaction recovery (2026-07-25) + +### What changed + +- Oversized one-shot compactions now recover through bounded, durable complete-turn stages while successful one-shot + compactions retain their existing path. +- Intermediate checkpoints are append-only and restart-safe; no-fit, later-stage failure, non-progress, stage-count, + and time bounds fail explicitly without retrying the original provider prompt. + +### Why extension system couldn't handle this alone + +Only `AgentSession` owns provider retry admission, lifecycle generations, message revision guards, and durable +`SessionManager.appendCompaction()` ordering across multiple summary calls. + +### Expected merge conflict zones + +- MEDIUM: `core/agent-session.ts` compaction execution and recovery helpers. +- MEDIUM: `core/compaction/compaction.ts` preparation helpers. +- HIGH: `core/extensions/types.ts` public compaction event interfaces. + ## Manual compaction keeps agent lifecycle subscription through abort (2026-07-24) - `core/agent-session.ts`: manual or extension-initiated compaction claims its synchronous admission/barrier first, diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 590d0bcbb..f63f99068 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -58,13 +58,16 @@ import { sleep } from "../utils/sleep.ts"; import { formatNoApiKeyFoundMessage, formatNoModelSelectedMessage } from "./auth-guidance.ts"; import { type BashResult, executeBashWithOperations } from "./bash-executor.ts"; import { + type CompactionPreparation, type CompactionResult, + CompactionSummaryOverflowError, calculateContextTokens, collectEntriesForBranchSummary, compact, estimateContextTokens, estimateTokens, generateBranchSummary, + planStagedCompactionChunk, prepareCompaction, shouldCompact, } from "./compaction/index.ts"; @@ -110,6 +113,7 @@ import type { ApplyCompactionResult, CompactionReason, CompactionRejectionCause, + CompactionStage, ModelSelectSource, } from "./extensions/types.ts"; import { type BashExecutionMessage, type CustomMessage, filterContextExcludedMessages } from "./messages.ts"; @@ -179,8 +183,8 @@ export type AgentSessionEvent = steering: readonly string[]; followUp: readonly string[]; } - | { type: "compaction_start"; reason: CompactionReason } - | { type: "compaction_progress"; reason: CompactionReason; delta?: string; text?: string } + | { type: "compaction_start"; reason: CompactionReason; stage?: CompactionStage } + | { type: "compaction_progress"; reason: CompactionReason; delta?: string; text?: string; stage?: CompactionStage } | { type: "entry_appended"; entry: SessionEntry } | { type: "session_info_changed"; name: string | undefined } | ExtensionToolHookLifecycleEvent @@ -196,6 +200,7 @@ export type AgentSessionEvent = accepted?: boolean; rejectionCause?: CompactionRejectionCause; errorMessage?: string; + stage?: CompactionStage; } | { type: "auto_retry_start"; attempt: number; maxAttempts: number; delayMs: number; errorMessage: string } | { type: "auto_retry_end"; success: boolean; attempt: number; finalError?: string } @@ -283,6 +288,10 @@ interface CompactionExecutionRequest { precomputed?: CompactionResult; allowSummaryOnly?: boolean; agentMessagesAtStart?: readonly AgentMessage[]; + preparation?: CompactionPreparation; + stage?: CompactionStage; + allowIntermediateOverflow?: boolean; + suppressWouldOverflowFeedback?: boolean; } type CompactionExecutionResult = @@ -292,6 +301,7 @@ type CompactionExecutionResult = result: CompactionResult; compactionEntry: CompactionEntry; fromExtension: boolean; + stageFinal: boolean; } | { accepted: false; @@ -391,6 +401,21 @@ class RequiredCompactionError extends Error { } } +const STAGED_COMPACTION_MAX_STAGES = 8; +const STAGED_COMPACTION_MAX_DURATION_MS = 5 * 60 * 1000; +const STAGED_COMPACTION_INPUT_RATIO = 0.4; + +class StagedCompactionError extends Error { + constructor(message: string) { + super(message); + this.name = "StagedCompactionError"; + } +} + +function getStagedCompactionInputBudget(model: Model): number { + return Math.max(1, Math.floor(model.contextWindow * STAGED_COMPACTION_INPUT_RATIO)); +} + class MissingModelAccessError extends Error { constructor() { super("AgentSession requires modelRuntime or modelRegistry"); @@ -3305,7 +3330,7 @@ export class AgentSession { this._disconnectFromAgent(); disconnected = true; this._emit({ type: "compaction_start", reason: "manual" }); - const execution = await this._executeCompaction({ + const execution = await this._executeCompactionWithStagedRecovery({ controller, owner: "compaction", reason: "manual", @@ -3366,7 +3391,7 @@ export class AgentSession { this._claimCompactionController(controller, "compaction"); try { - const execution = await this._executeCompaction({ + const execution = await this._executeCompactionWithStagedRecovery({ controller, owner: "compaction", reason: options.reason, @@ -3465,6 +3490,103 @@ export class AgentSession { this._releaseCompactionController(options.signal); } + private async _executeCompactionWithStagedRecovery( + request: CompactionExecutionRequest, + ): Promise { + const initialExecution = await this._executeCompaction({ ...request, suppressWouldOverflowFeedback: true }); + if ( + initialExecution.accepted || + initialExecution.rejectionCause !== "would-overflow" || + request.preparation !== undefined + ) { + return initialExecution; + } + + return await this._executeStagedCompaction(request); + } + + private async _executeStagedCompaction(request: CompactionExecutionRequest): Promise { + const model = this.model; + if (!model) throw new Error(formatNoModelSelectedMessage()); + const startedAt = Date.now(); + const inputBudgetTokens = getStagedCompactionInputBudget(model); + let previousCursor: string | undefined; + let timedOut = false; + const timeout = setTimeout(() => { + timedOut = true; + request.controller.abort(); + }, STAGED_COMPACTION_MAX_DURATION_MS); + + try { + for (let stageIndex = 1; stageIndex <= STAGED_COMPACTION_MAX_STAGES; stageIndex++) { + if (Date.now() - startedAt >= STAGED_COMPACTION_MAX_DURATION_MS) { + throw new StagedCompactionError( + `Staged compaction exceeded ${STAGED_COMPACTION_MAX_DURATION_MS}ms after ${stageIndex - 1} stages`, + ); + } + if ( + !this._ownsCompactionController(request.controller, request.owner) || + request.controller.signal.aborted + ) { + throw new CompactionCancelledError(); + } + + const plan = planStagedCompactionChunk( + this.sessionManager.getBranch(), + this.settingsManager.getCompactionSettings(), + inputBudgetTokens, + ); + if (plan.status === "no-fit") { + const entry = plan.entryId ? ` at entry ${plan.entryId}` : ""; + throw new StagedCompactionError( + plan.reason === "single-group-too-large" + ? `Staged compaction cannot fit the oldest complete turn${entry}: estimated ${plan.requiredTokens} tokens exceeds the ${plan.budgetTokens}-token stage budget` + : `Staged compaction cannot advance${entry}: no older complete turn remains outside the required continuation suffix`, + ); + } + + const cursor = plan.chunk.preparation.firstKeptEntryId; + if (cursor === previousCursor) { + throw new StagedCompactionError(`Staged compaction made no cursor progress at entry ${cursor}`); + } + const stage: CompactionStage = { index: stageIndex, maxStages: STAGED_COMPACTION_MAX_STAGES }; + this._emit({ type: "compaction_start", reason: request.reason, stage }); + this._emit({ + type: "compaction_progress", + reason: request.reason, + stage, + text: `Compacting stage ${stageIndex}/${STAGED_COMPACTION_MAX_STAGES}`, + }); + + const execution = await this._executeCompaction({ + ...request, + precomputed: undefined, + preparation: plan.chunk.preparation, + stage, + allowIntermediateOverflow: stageIndex < STAGED_COMPACTION_MAX_STAGES, + agentMessagesAtStart: this.agent.state.messages.slice(), + }); + if (!execution.accepted || execution.stageFinal) { + return execution; + } + previousCursor = cursor; + } + + throw new StagedCompactionError( + `Staged compaction reached the ${STAGED_COMPACTION_MAX_STAGES}-stage limit before fitting the context`, + ); + } catch (error) { + if (timedOut) { + throw new StagedCompactionError( + `Staged compaction exceeded ${STAGED_COMPACTION_MAX_DURATION_MS}ms before recovery completed`, + ); + } + throw error; + } finally { + clearTimeout(timeout); + } + } + private async _executeCompaction(request: CompactionExecutionRequest): Promise { const model = this.model; if (!model) throw new Error(formatNoModelSelectedMessage()); @@ -3501,12 +3623,9 @@ export class AgentSession { let fromExtension = request.precomputed !== undefined; if (!compactionResult) { - const preparation = prepareCompaction( - pathEntries, - settings, - request.reason === "overflow", - request.allowSummaryOnly, - ); + const preparation = + request.preparation ?? + prepareCompaction(pathEntries, settings, request.reason === "overflow", request.allowSummaryOnly); if (!preparation) { const lastEntry = pathEntries[pathEntries.length - 1]; @@ -3527,6 +3646,7 @@ export class AgentSession { branchEntries: pathEntries, customInstructions: request.customInstructions, signal, + ...(request.stage ? { stage: request.stage } : {}), })) as SessionBeforeCompactResult | undefined; for (const message of this.agent.state.messages) { if ( @@ -3560,19 +3680,33 @@ export class AgentSession { if (!compactionResult) { const { apiKey, headers, extraBody, env } = await this._getCompactionRequestAuth(model); - compactionResult = await compact( - preparation, - model, - apiKey, - headers, - request.customInstructions, - signal, - extraBody, - thinkingLevel, - this.agent.streamFn, - env, - this.agent.transformContext, - ); + try { + compactionResult = await compact( + preparation, + model, + apiKey, + headers, + request.customInstructions, + signal, + extraBody, + thinkingLevel, + this.agent.streamFn, + env, + this.agent.transformContext, + ); + } catch (error) { + if (error instanceof CompactionSummaryOverflowError) { + return await this._rejectCompaction( + request, + requestId, + operationId, + "would-overflow", + false, + error.message, + ); + } + throw error; + } } } @@ -3609,9 +3743,22 @@ export class AgentSession { return await this._rejectCompaction(request, requestId, operationId, "stale-revision", false); } - if (this._wouldCompactionOverflow(pathEntries, compactionResult, fromExtension, model)) { + if (request.stage && Math.ceil(compactionResult.summary.length / 4) >= getStagedCompactionInputBudget(model)) { + return await this._rejectCompaction( + request, + requestId, + operationId, + "would-overflow", + false, + "staged compaction summary would overflow the input budget for the next checkpoint", + ); + } + + const wouldOverflow = this._wouldCompactionOverflow(pathEntries, compactionResult, fromExtension, model); + if (wouldOverflow && !(request.stage && request.allowIntermediateOverflow)) { return await this._rejectCompaction(request, requestId, operationId, "would-overflow", false); } + const stageFinal = !wouldOverflow; const compactionEntryId = this.sessionManager.appendCompaction( compactionResult.summary, @@ -3663,14 +3810,17 @@ export class AgentSession { throw new CompactionCancelledError(); } + const willRetry = request.stage ? request.willRetry && stageFinal : request.willRetry; + const completedStage = request.stage ? { ...request.stage, final: stageFinal } : undefined; this._emit({ type: "compaction_end", reason: request.reason, result: compactionResult, aborted: false, - willRetry: request.willRetry, + willRetry, requestId, accepted: true, + ...(completedStage ? { stage: completedStage } : {}), }); await this._extensionRunner.emit({ @@ -3680,10 +3830,18 @@ export class AgentSession { accepted: true, compactionEntry: savedEntry, fromExtension, - willRetry: request.willRetry, + willRetry, + ...(completedStage ? { stage: completedStage } : {}), }); - return { accepted: true, requestId, result: compactionResult, compactionEntry: savedEntry, fromExtension }; + return { + accepted: true, + requestId, + result: compactionResult, + compactionEntry: savedEntry, + fromExtension, + stageFinal, + }; } catch (error) { if (error instanceof CompactionExecutionError) { throw error; @@ -3762,7 +3920,9 @@ export class AgentSession { aborted: boolean, extensionReason?: string, ): Promise { - // Per plan Section 1: rejection must never be silent. The compaction_end event + // User-visible rejection must never be silent. The one exception is the + // internal one-shot would-overflow signal immediately consumed by staged + // recovery; its staged terminal event owns the visible outcome. Otherwise the compaction_end event // carries a non-empty human-readable errorMessage (unless the user aborted, where // the aborted branch already renders "Compaction cancelled"). session_compact is // also emitted with accepted:false so the compaction extension's circuit-breaker @@ -3784,6 +3944,9 @@ export class AgentSession { ) { throw new CompactionCancelledError(); } + if (request.suppressWouldOverflowFeedback && rejectionCause === "would-overflow") { + return { accepted: false, requestId, rejectionCause }; + } this._emit({ type: "compaction_end", reason: request.reason, @@ -3794,6 +3957,7 @@ export class AgentSession { accepted: false, rejectionCause, errorMessage, + ...(request.stage ? { stage: request.stage } : {}), }); await this._extensionRunner.emit({ type: "session_compact", @@ -3803,6 +3967,7 @@ export class AgentSession { rejectionCause, fromExtension: false, willRetry: false, + ...(request.stage ? { stage: request.stage } : {}), }); return { accepted: false, requestId, rejectionCause }; } @@ -4119,7 +4284,7 @@ export class AgentSession { this._emit({ type: "compaction_start", reason }); try { - const execution = await this._executeCompaction({ + const execution = await this._executeCompactionWithStagedRecovery({ controller, owner: "compaction", reason, @@ -4252,7 +4417,7 @@ export class AgentSession { if (!this._ownsCompactionController(autoCompactionController, "auto")) return false; this._emit({ type: "compaction_start", reason }); - const execution = await this._executeCompaction({ + const execution = await this._executeCompactionWithStagedRecovery({ controller: autoCompactionController, owner: "auto", reason, @@ -4595,7 +4760,7 @@ export class AgentSession { this._disconnectFromAgent(); disconnected = true; this._emit({ type: "compaction_start", reason: "extension" }); - const execution = await this._executeCompaction({ + const execution = await this._executeCompactionWithStagedRecovery({ controller, owner: "compaction", reason: "extension", diff --git a/packages/coding-agent/src/core/changes.md b/packages/coding-agent/src/core/changes.md index ec5e08714..b34b4632d 100644 --- a/packages/coding-agent/src/core/changes.md +++ b/packages/coding-agent/src/core/changes.md @@ -1,5 +1,23 @@ # changes +## Durable staged compaction recovery (2026-07-25) + +### What changed + +- `agent-session.ts` preserves the existing one-shot compaction attempt, then falls back on `would-overflow` to at + most eight complete-turn stages within five minutes. +- Each accepted intermediate stage appends a normal compaction entry before the next request. A later error or abort + therefore preserves the last accepted checkpoint, while the original provider retry is marked `willRetry` only on + the stage that actually brings active context below budget. +- Staged execution rejects a non-advancing cursor, an oversized single turn, and a summary that leaves no budget for + the next checkpoint. The eighth stage must finish below budget or is rejected without another append. + +### Why + +One-shot overflow recovery could not recover when its own summarization request or resulting checkpoint remained too +large. The bounded append-only sequence makes partial progress restart-safe without changing successful one-shot +behavior. + ## Session-owned compaction lifecycle (2026-07-23) ### What changed diff --git a/packages/coding-agent/src/core/compaction/changes.md b/packages/coding-agent/src/core/compaction/changes.md index b497758b0..5b007ed6c 100644 --- a/packages/coding-agent/src/core/compaction/changes.md +++ b/packages/coding-agent/src/core/compaction/changes.md @@ -1,5 +1,26 @@ # changes.md — compaction +## Bounded staged compaction planner (2026-07-25) + +### What changed + +- `compaction.ts` adds a pure `planStagedCompactionChunk()` planner. It selects a token-budgeted contiguous prefix of + complete oldest turns, retains the newest turn, keeps assistant tool calls with their results, and incorporates the + latest accepted summary into each next-stage budget. +- Planning returns an explicit `no-fit` result when the oldest complete turn cannot fit or no older turn remains, and + every ready plan advances `firstKeptEntryId`. +- Local summarization now identifies summary-request context overflow with `CompactionSummaryOverflowError` so + `AgentSession` can switch to durable staged recovery instead of treating it as an unrelated provider failure. + +### Why + +A one-shot summary can exceed the model context before any checkpoint is written. Retrying after silently deleting +old summary input would persist the original checkpoint boundary without actually summarizing that deleted history. + +### Expected merge conflict zones + +- MEDIUM: `compaction.ts` around preparation and summarization error handling. + ## Lifecycle ownership and required-admission safety (2026-07-23) ### What changed diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 001b29de3..b2b122df5 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -7,7 +7,7 @@ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; -import { streamSimple } from "@earendil-works/pi-ai/compat"; +import { isContextOverflow, streamSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm, filterContextExcludedMessages, isContextExcludedCustomMessage } from "../messages.ts"; import { buildSessionContext, @@ -31,6 +31,14 @@ type SummarizationOptions = SimpleStreamOptions & { readonly env?: Record; }; +/** Signals that only the summarization input, not the main agent turn, overflowed. */ +export class CompactionSummaryOverflowError extends Error { + constructor(message: string) { + super(message); + this.name = "CompactionSummaryOverflowError"; + } +} + // ============================================================================ // File Operation Tracking // ============================================================================ @@ -715,6 +723,12 @@ export async function generateSummary( streamFn, ); + if (isContextOverflow(response, model.contextWindow)) { + throw new CompactionSummaryOverflowError( + response.errorMessage || "Compaction summary request exceeded the context window", + ); + } + if (response.stopReason === "error") { throw new Error(`Summarization failed: ${response.errorMessage || "Unknown error"}`); } @@ -749,16 +763,39 @@ export interface CompactionPreparation { settings: CompactionSettings; } -export function prepareCompaction( - pathEntries: SessionEntry[], - settings: CompactionSettings, - forceProgress = false, - allowSummaryOnly = false, -): CompactionPreparation | undefined { - if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") { - return undefined; - } +export interface StagedCompactionChunk { + /** Preparation for exactly one durable, contiguous-prefix checkpoint. */ + preparation: CompactionPreparation; + /** First context-visible entry included in this chunk. */ + firstSummarizedEntryId: string; + /** Last context-visible entry included in this chunk. */ + lastSummarizedEntryId: string; + /** Estimated previous-summary plus new-message tokens sent to the summarizer. */ + estimatedInputTokens: number; +} +export type StagedCompactionPlan = + | { status: "ready"; chunk: StagedCompactionChunk } + | { + status: "no-fit"; + reason: "single-group-too-large" | "no-older-complete-turn"; + budgetTokens: number; + requiredTokens: number; + entryId?: string; + }; + +interface StagedTurnGroup { + startIndex: number; + endIndex: number; + messages: AgentMessage[]; + tokens: number; +} + +function getPreviousCompactionBoundary(pathEntries: SessionEntry[]): { + boundaryStart: number; + previousSummary: string | undefined; + prevCompactionIndex: number; +} { let prevCompactionIndex = -1; for (let i = pathEntries.length - 1; i >= 0; i--) { if (pathEntries[i].type === "compaction") { @@ -767,14 +804,158 @@ export function prepareCompaction( } } - let previousSummary: string | undefined; - let boundaryStart = 0; - if (prevCompactionIndex >= 0) { - const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; - previousSummary = prevCompaction.summary; - const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); - boundaryStart = firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1; + if (prevCompactionIndex < 0) { + return { boundaryStart: 0, previousSummary: undefined, prevCompactionIndex }; } + + const prevCompaction = pathEntries[prevCompactionIndex] as CompactionEntry; + const firstKeptEntryIndex = pathEntries.findIndex((entry) => entry.id === prevCompaction.firstKeptEntryId); + return { + boundaryStart: firstKeptEntryIndex >= 0 ? firstKeptEntryIndex : prevCompactionIndex + 1, + previousSummary: prevCompaction.summary, + prevCompactionIndex, + }; +} + +function buildStagedTurnGroups(pathEntries: SessionEntry[], boundaryStart: number): StagedTurnGroup[] { + const groups: StagedTurnGroup[] = []; + let current: StagedTurnGroup | undefined; + + for (let index = boundaryStart; index < pathEntries.length; index++) { + const messages = contextMessagesForCompactionEntry(pathEntries[index]); + for (const message of messages) { + if (isTurnStartMessage(message) && current && current.messages.length > 0) { + groups.push(current); + current = undefined; + } + current ??= { startIndex: index, endIndex: index + 1, messages: [], tokens: 0 }; + current.endIndex = index + 1; + current.messages.push(message); + current.tokens += estimateTokens(message); + } + } + + if (current && current.messages.length > 0) { + groups.push(current); + } + return groups; +} + +/** + * Plan one bounded staged-compaction chunk without mutating session state. + * + * Chunks always contain a contiguous prefix of complete turns. The newest turn + * is retained, and assistant tool calls remain grouped with every following + * tool result because a chunk boundary is only placed at a user-like turn + * start. A successful caller can therefore append the returned checkpoint and + * call this function again against the new branch; the cursor must advance to + * `preparation.firstKeptEntryId` on every accepted stage. + */ +export function planStagedCompactionChunk( + pathEntries: SessionEntry[], + settings: CompactionSettings, + inputBudgetTokens: number, +): StagedCompactionPlan { + const budgetTokens = Math.max(1, Math.floor(inputBudgetTokens)); + const { boundaryStart, previousSummary, prevCompactionIndex } = getPreviousCompactionBoundary(pathEntries); + const groups = buildStagedTurnGroups(pathEntries, boundaryStart); + const previousSummaryTokens = previousSummary ? Math.ceil(previousSummary.length / 4) : 0; + + if (groups.length < 2) { + return { + status: "no-fit", + reason: "no-older-complete-turn", + budgetTokens, + requiredTokens: previousSummaryTokens + (groups[0]?.tokens ?? 0), + ...(groups[0] ? { entryId: pathEntries[groups[0].startIndex]?.id } : {}), + }; + } + + const availableMessageTokens = budgetTokens - previousSummaryTokens; + const oldestGroup = groups[0]; + if (availableMessageTokens < oldestGroup.tokens) { + return { + status: "no-fit", + reason: "single-group-too-large", + budgetTokens, + requiredTokens: previousSummaryTokens + oldestGroup.tokens, + entryId: pathEntries[oldestGroup.startIndex]?.id, + }; + } + + let selectedGroupCount = 0; + let selectedMessageTokens = 0; + // Always retain the newest complete turn as the continuation suffix. + while (selectedGroupCount < groups.length - 1) { + const candidate = groups[selectedGroupCount]; + if (selectedMessageTokens + candidate.tokens > availableMessageTokens) break; + selectedMessageTokens += candidate.tokens; + selectedGroupCount++; + } + + if (selectedGroupCount === 0) { + return { + status: "no-fit", + reason: "single-group-too-large", + budgetTokens, + requiredTokens: previousSummaryTokens + oldestGroup.tokens, + entryId: pathEntries[oldestGroup.startIndex]?.id, + }; + } + + const selectedGroups = groups.slice(0, selectedGroupCount); + const firstSelectedGroup = selectedGroups[0]; + const lastSelectedGroup = selectedGroups[selectedGroups.length - 1]; + const firstKeptGroup = groups[selectedGroupCount]; + const firstKeptEntry = pathEntries[firstKeptGroup.startIndex]; + const firstSummarizedEntry = pathEntries[firstSelectedGroup.startIndex]; + const lastSummarizedEntry = pathEntries[lastSelectedGroup.endIndex - 1]; + if (!firstKeptEntry?.id || !firstSummarizedEntry?.id || !lastSummarizedEntry?.id) { + return { + status: "no-fit", + reason: "no-older-complete-turn", + budgetTokens, + requiredTokens: previousSummaryTokens + selectedMessageTokens, + }; + } + + const messagesToSummarize = selectedGroups.flatMap((group) => group.messages); + const fileOps = extractFileOperations(messagesToSummarize, pathEntries, prevCompactionIndex); + const tokensBefore = estimateContextTokens( + filterContextExcludedMessages(buildSessionContext(pathEntries).messages), + ).tokens; + + return { + status: "ready", + chunk: { + preparation: { + firstKeptEntryId: firstKeptEntry.id, + messagesToSummarize, + turnPrefixMessages: [], + isSplitTurn: false, + tokensBefore, + previousSummary, + fileOps, + settings, + }, + firstSummarizedEntryId: firstSummarizedEntry.id, + lastSummarizedEntryId: lastSummarizedEntry.id, + estimatedInputTokens: previousSummaryTokens + selectedMessageTokens, + }, + }; +} + +export function prepareCompaction( + pathEntries: SessionEntry[], + settings: CompactionSettings, + forceProgress = false, + allowSummaryOnly = false, +): CompactionPreparation | undefined { + if (pathEntries.length > 0 && pathEntries[pathEntries.length - 1].type === "compaction") { + return undefined; + } + + const { boundaryStart, previousSummary, prevCompactionIndex } = getPreviousCompactionBoundary(pathEntries); const boundaryEnd = pathEntries.length; const tokensBefore = estimateContextTokens( @@ -1009,6 +1190,12 @@ async function generateTurnPrefixSummary( streamFn, ); + if (isContextOverflow(response, model.contextWindow)) { + throw new CompactionSummaryOverflowError( + response.errorMessage || "Turn prefix summary request exceeded the context window", + ); + } + if (response.stopReason === "error") { throw new Error(`Turn prefix summarization failed: ${response.errorMessage || "Unknown error"}`); } diff --git a/packages/coding-agent/src/core/extensions/builtin/compaction/changes.md b/packages/coding-agent/src/core/extensions/builtin/compaction/changes.md index 74a527664..0ab50b486 100644 --- a/packages/coding-agent/src/core/extensions/builtin/compaction/changes.md +++ b/packages/coding-agent/src/core/extensions/builtin/compaction/changes.md @@ -1,3 +1,13 @@ +## Durable staged overflow handoff (2026-07-25) + +- Local summary context overflow is surfaced as structured `would-overflow` feedback instead of deleting oldest + messages from the summarizer input while retaining the original checkpoint boundary. +- Staged requests after stage 1 bypass a second per-turn admission check, and accepted-stage accounting increments + only once for the logical compaction. Other accepted-stage effects remain append-driven so a later abort still leaves + a valid latest checkpoint. + +Expected upstream conflict zones: `builtin/compaction/index.ts` event policy and `speculative.ts` summary error path. + ## Canonical remote compaction provenance and route ownership (2026-07-24) - `openai-remote-model.ts`: provenance now hashes the normalized endpoint and every final header by default. The only excluded volatile transport headers are `content-length`, `user-agent`, `request-id`, `x-request-id`, and `x-client-request-id`; raw values are never persisted. This binds non-Codex checkpoints to authorization plus final tenant/workspace routing headers. diff --git a/packages/coding-agent/src/core/extensions/builtin/compaction/index.ts b/packages/coding-agent/src/core/extensions/builtin/compaction/index.ts index c4da6f59f..2b8bf5f51 100644 --- a/packages/coding-agent/src/core/extensions/builtin/compaction/index.ts +++ b/packages/coding-agent/src/core/extensions/builtin/compaction/index.ts @@ -336,7 +336,9 @@ export default function compactionExtension(pi: ExtensionAPI): void { pi.on("session_before_compact", async (event, ctx) => { invalidateSpeculativeCompaction(); - if (cap.shouldRejectByCap(state, { reason: event.reason }).cancel) { + // A staged sequence is one logical compaction. Admit it against the cap on + // stage 1, then let later durable checkpoints finish that admitted recovery. + if ((!event.stage || event.stage.index === 1) && cap.shouldRejectByCap(state, { reason: event.reason }).cancel) { return { cancel: true, rejectionCause: "per-turn-cap", @@ -390,7 +392,11 @@ export default function compactionExtension(pi: ExtensionAPI): void { // the real provider message; ctx.ui.notify would be a duplicate toast. pendingMetadata.delete(event.requestId); if (error instanceof SummaryGenerationError) { - return { cancel: true, reason: error.message }; + return { + cancel: true, + ...(error.kind === "context-overflow" ? { rejectionCause: "would-overflow" as const } : {}), + reason: error.message, + }; } const message = error instanceof Error ? error.message : String(error); return { cancel: true, reason: `compaction generator failed: ${message}` }; @@ -422,7 +428,9 @@ export default function compactionExtension(pi: ExtensionAPI): void { const branchEntries = ctx.sessionManager.getBranch(); const firstKeptIndex = branchEntries.findIndex((entry) => entry.id === event.compactionEntry.firstKeptEntryId); const keptEntries = firstKeptIndex === -1 ? [] : branchEntries.slice(firstKeptIndex); - state = cap.incrementAccepted(state); + if (!event.stage || event.stage.index === 1) { + state = cap.incrementAccepted(state); + } state = breaker.recordSuccess(state); state = updateLastYield(state, event.compactionEntry); resetOnSessionCompact(degradationState); diff --git a/packages/coding-agent/src/core/extensions/builtin/compaction/speculative.ts b/packages/coding-agent/src/core/extensions/builtin/compaction/speculative.ts index 01c48cb54..6d7fbe6d4 100644 --- a/packages/coding-agent/src/core/extensions/builtin/compaction/speculative.ts +++ b/packages/coding-agent/src/core/extensions/builtin/compaction/speculative.ts @@ -80,9 +80,9 @@ function approxTokens(text: string): number { * manual route and degrade to "unavailable" on automatic routes. */ export class SummaryGenerationError extends Error { - readonly kind: "auth" | "empty-summary"; + readonly kind: "auth" | "context-overflow" | "empty-summary"; - constructor(kind: "auth" | "empty-summary", message: string) { + constructor(kind: "auth" | "context-overflow" | "empty-summary", message: string) { super(message); this.name = "SummaryGenerationError"; this.kind = kind; @@ -296,16 +296,6 @@ function pruneOldMessagesToBudget(messages: AgentMessage[], targetTokens: number return pruned; } -function removeOldestHistoryItemForOverflowRetry(messages: AgentMessage[]): AgentMessage[] | undefined { - if (messages.length <= 1) return undefined; - const boundaryIndex = findLastUserLikeIndex(messages); - return ( - removeFirstOldToolPair(messages, boundaryIndex)?.messages ?? - removeFirstOldMessage(messages, boundaryIndex)?.messages ?? - (messages.length > 1 ? messages.slice(1) : undefined) - ); -} - function estimateTotalTokens(messages: AgentMessage[]): number { let total = 0; for (const message of messages) total += estimateTokens(message); @@ -396,7 +386,7 @@ export async function runExtensionCompaction( throw new SummaryGenerationError("auth", `summarization credentials unavailable: ${detail}`); } - let messages = pruneToolResults( + const messages = pruneToolResults( [...snapshot.preparation.messagesToSummarize, ...snapshot.preparation.turnPrefixMessages], snapshot.contextWindow, ); @@ -406,66 +396,60 @@ export async function runExtensionCompaction( customInstructions: snapshot.customInstructions, }); - while (true) { - if (signal?.aborted) return undefined; - const response = await generateSummaryMessage({ - context, - messages, - onProgress, - prompt, - signal, - snapshot, - auth: { - apiKey: auth.apiKey, - headers: auth.headers, - extraBody: auth.extraBody, - }, - }); - if (!response) return undefined; - - if (isAssistantMessage(response) && isContextOverflow(response, snapshot.contextWindow)) { - const retryMessages = removeOldestHistoryItemForOverflowRetry(messages); - if (!retryMessages || retryMessages.length === messages.length) { - break; - } - messages = retryMessages; - continue; - } - - if (isAssistantMessage(response) && response.stopReason === "aborted") { - // A partial summary from an aborted stream must never be applied. - return undefined; - } + if (signal?.aborted) return undefined; + const response = await generateSummaryMessage({ + context, + messages, + onProgress, + prompt, + signal, + snapshot, + auth: { + apiKey: auth.apiKey, + headers: auth.headers, + extraBody: auth.extraBody, + }, + }); + if (!response) return undefined; - if (isAssistantMessage(response) && response.stopReason === "error") { - // Surface the real provider failure instead of silently degrading - // into a generic "Compaction cancelled". - throw new Error(response.errorMessage || "Compaction summary request failed"); - } + if (isAssistantMessage(response) && isContextOverflow(response, snapshot.contextWindow)) { + throw new SummaryGenerationError( + "context-overflow", + "compaction summary request exceeded the context window; retrying with staged compaction", + ); + } - const summary = getSummaryText(response); - if (!summary) { - const stopReason = isAssistantMessage(response) ? response.stopReason : "unknown"; - throw new SummaryGenerationError( - "empty-summary", - `summarization response contained no text (stopReason: ${stopReason})`, - ); - } + if (isAssistantMessage(response) && response.stopReason === "aborted") { + // A partial summary from an aborted stream must never be applied. + return undefined; + } - // Informational only: the core rejects an applied compaction that would - // still overflow (_wouldCompactionOverflow). Rejecting here based on the - // size of the *discarded* input made large sessions uncompactable. - const tokenEstimate = estimateContextTokens(convertToLlm(messages)).tokens + approxTokens(summary); + if (isAssistantMessage(response) && response.stopReason === "error") { + // Surface the real provider failure instead of silently degrading + // into a generic "Compaction cancelled". + throw new Error(response.errorMessage || "Compaction summary request failed"); + } - return { - summary, - firstKeptEntryId: snapshot.preparation.firstKeptEntryId, - tokensBefore: snapshot.preparation.tokensBefore, - details: { schema: SUMMARY_SCHEMA, promptVariant: snapshot.promptVariant, tokenEstimate }, - }; + const summary = getSummaryText(response); + if (!summary) { + const stopReason = isAssistantMessage(response) ? response.stopReason : "unknown"; + throw new SummaryGenerationError( + "empty-summary", + `summarization response contained no text (stopReason: ${stopReason})`, + ); } - throw new Error("Compaction summary request exceeded the context window after retrying with a smaller input"); + // Informational only: the core rejects an applied compaction that would + // still overflow (_wouldCompactionOverflow). Rejecting here based on the + // size of the *discarded* input made large sessions uncompactable. + const tokenEstimate = estimateContextTokens(convertToLlm(messages)).tokens + approxTokens(summary); + + return { + summary, + firstKeptEntryId: snapshot.preparation.firstKeptEntryId, + tokensBefore: snapshot.preparation.tokensBefore, + details: { schema: SUMMARY_SCHEMA, promptVariant: snapshot.promptVariant, tokenEstimate }, + }; } export async function applyGeneratedCompaction( diff --git a/packages/coding-agent/src/core/extensions/changes.md b/packages/coding-agent/src/core/extensions/changes.md index db710deec..ed0c359e2 100644 --- a/packages/coding-agent/src/core/extensions/changes.md +++ b/packages/coding-agent/src/core/extensions/changes.md @@ -1,5 +1,22 @@ # Core Extensions Changes +## 2026-07-25 - Optional staged compaction metadata + +### What changed + +- `SessionBeforeCompactEvent` and both `SessionCompactEvent` variants accept optional `stage` metadata with a + one-based index, hard maximum, and terminal `final` flag. +- Existing handlers remain source-compatible because the field is additive and ordinary one-shot compactions omit it. + +### Why + +Builtin compaction policy must treat several durable checkpoints as one admitted logical compaction while observers +can still correlate intermediate versus final stages. + +### Expected merge conflict zones + +- HIGH: `types.ts` around the public compaction event interfaces. + ## 2026-07-23 - Compaction feedback operation handles ### What changed diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index 9979a49d8..0bb48d7f8 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -98,6 +98,16 @@ export type CompactionRejectionCause = | "per-turn-cap" | "stale-revision"; +/** Optional metadata for one stage of a bounded logical compaction. */ +export interface CompactionStage { + /** One-based stage number. */ + index: number; + /** Hard upper bound on provider calls made by the staged recovery sequence. */ + maxStages: number; + /** Present on terminal stage events after the resulting context is measured. */ + final?: boolean; +} + // ============================================================================ // UI Context // ============================================================================ @@ -716,6 +726,8 @@ export interface SessionBeforeCompactEvent { branchEntries: SessionEntry[]; customInstructions?: string; signal: AbortSignal; + /** Present only when core is recovering from an oversized one-shot compaction. */ + stage?: CompactionStage; } /** @@ -740,6 +752,8 @@ export interface SessionCompactAcceptedEvent { fromExtension: boolean; /** True when the aborted turn is retried after this compaction (overflow recovery) */ willRetry: boolean; + /** Present only for an accepted durable stage of one logical compaction. */ + stage?: CompactionStage; } export interface SessionCompactRejectedEvent { @@ -757,6 +771,8 @@ export interface SessionCompactRejectedEvent { compactionEntry?: undefined; fromExtension: false; willRetry: false; + /** Present only when a staged attempt was rejected. */ + stage?: CompactionStage; } /** Fired before an extension runtime is torn down due to quit, reload, or session replacement. */ diff --git a/packages/coding-agent/test/compaction/speculative-compaction.test.ts b/packages/coding-agent/test/compaction/speculative-compaction.test.ts index e7bcc8249..335be9096 100644 --- a/packages/coding-agent/test/compaction/speculative-compaction.test.ts +++ b/packages/coding-agent/test/compaction/speculative-compaction.test.ts @@ -478,7 +478,7 @@ describe("speculative compaction", () => { expect(deltas.join("")).toBe("live summary"); }); - it("retries a compaction summary request with a smaller input after a context-window failure", async () => { + it("surfaces a context-window failure for durable staged recovery", async () => { // Given const context = createContext(); const snapshot = createSpeculativeCompactionSnapshot(context, { generation: 1 }); @@ -488,18 +488,17 @@ describe("speculative compaction", () => { errorMessage: "Your input exceeds the context window of this model. Please adjust your input and try again.", }), - fauxAssistantMessage("retry summary"), ]); - // When - const result = snapshot ? await runExtensionCompaction(context, snapshot) : undefined; - - // Then - expect(result?.summary).toBe("retry summary"); - expect(context.registration.getCallLog()).toHaveLength(2); + // When / Then + await expect(snapshot ? runExtensionCompaction(context, snapshot) : undefined).rejects.toMatchObject({ + name: "SummaryGenerationError", + kind: "context-overflow", + }); + expect(context.registration.getCallLog()).toHaveLength(1); }); - it("keeps pruning and retrying after repeated compaction summary context-window failures", async () => { + it("does not drop old messages from a checkpoint after a summary context-window failure", async () => { // Given const context = createContext(); context.getCompactionSettings = () => ({ ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 1 }); @@ -515,36 +514,19 @@ describe("speculative compaction", () => { errorMessage: "Your input exceeds the context window of this model. Please adjust your input and try again.", }), - fauxAssistantMessage("", { - stopReason: "error", - errorMessage: - "Your input exceeds the context window of this model. Please adjust your input and try again.", - }), - fauxAssistantMessage("eventually compacted"), ]); - // When - const result = snapshot ? await runExtensionCompaction(context, snapshot) : undefined; - - // Then - expect(result?.summary).toBe("eventually compacted"); - const requestTexts = context.registration.getCallLog().map((entry) => { - const firstMessage = entry.context.messages[0]; - if (!firstMessage) return ""; - const content = firstMessage.content; - if (typeof content === "string") return content; - return content - .filter((part) => part.type === "text") - .map((part) => part.text) - .join("\n"); + // When / Then + await expect(snapshot ? runExtensionCompaction(context, snapshot) : undefined).rejects.toMatchObject({ + name: "SummaryGenerationError", + kind: "context-overflow", }); - expect(requestTexts).toHaveLength(3); - expect(requestTexts[0]).toContain("first user"); - expect(requestTexts[1]).not.toContain("first user"); - expect(requestTexts[1]).toContain("first assistant"); - expect(requestTexts[2]).not.toContain("first assistant"); - expect(requestTexts[2]).toContain("second user"); - expect(requestTexts[2]).not.toContain("kept recent user"); + const calls = context.registration.getCallLog(); + expect(calls).toHaveLength(1); + const requestText = calls[0]?.context.messages.map(messageText).join("\n") ?? ""; + expect(requestText).toContain("first user"); + expect(requestText).toContain("second user"); + expect(requestText).not.toContain("kept recent user"); }); it("sends the conversation as native messages with a trailing summarization instruction", async () => { diff --git a/packages/coding-agent/test/compaction/staged-compaction-planner.test.ts b/packages/coding-agent/test/compaction/staged-compaction-planner.test.ts new file mode 100644 index 000000000..528fc9999 --- /dev/null +++ b/packages/coding-agent/test/compaction/staged-compaction-planner.test.ts @@ -0,0 +1,120 @@ +import type { AgentMessage } from "@earendil-works/pi-agent-core"; +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { beforeEach, describe, expect, it } from "vitest"; +import { DEFAULT_COMPACTION_SETTINGS, planStagedCompactionChunk } from "../../src/core/compaction/index.ts"; +import type { SessionEntry, SessionMessageEntry } from "../../src/core/session-manager.ts"; + +let nextId = 0; +let parentId: string | null = null; + +function entry(message: AgentMessage): SessionMessageEntry { + const id = `entry-${nextId++}`; + const result: SessionMessageEntry = { + type: "message", + id, + parentId, + timestamp: new Date().toISOString(), + message, + }; + parentId = id; + return result; +} + +function user(text: string): AgentMessage { + return { role: "user", content: [{ type: "text", text }], timestamp: Date.now() }; +} + +function assistant(content: AssistantMessage["content"], stopReason: AssistantMessage["stopReason"] = "stop") { + return { + role: "assistant" as const, + content, + api: "faux-completion" as const, + provider: "faux", + model: "faux-model", + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + stopReason, + timestamp: Date.now(), + } satisfies AssistantMessage; +} + +function plan(entries: SessionEntry[], budgetTokens: number) { + return planStagedCompactionChunk(entries, { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 1 }, budgetTokens); +} + +describe("staged compaction chunk planner", () => { + beforeEach(() => { + nextId = 0; + parentId = null; + }); + + it("selects the largest contiguous prefix of complete turns that fits", () => { + const firstUser = entry(user("a".repeat(80))); + const firstAssistant = entry(assistant([{ type: "text", text: "b".repeat(80) }])); + const secondUser = entry(user("c".repeat(80))); + const secondAssistant = entry(assistant([{ type: "text", text: "d".repeat(80) }])); + const thirdUser = entry(user("keep newest")); + const entries = [firstUser, firstAssistant, secondUser, secondAssistant, thirdUser]; + + const result = plan(entries, 45); + + expect(result.status).toBe("ready"); + if (result.status !== "ready") return; + expect(result.chunk.firstSummarizedEntryId).toBe(firstUser.id); + expect(result.chunk.lastSummarizedEntryId).toBe(firstAssistant.id); + expect(result.chunk.preparation.firstKeptEntryId).toBe(secondUser.id); + expect(result.chunk.preparation.messagesToSummarize).toEqual([firstUser.message, firstAssistant.message]); + }); + + it("keeps a tool call, its result, and the rest of their turn in one chunk", () => { + const firstUser = entry(user("run tool")); + const toolCall = entry( + assistant([{ type: "toolCall", id: "call-1", name: "read", arguments: { path: "a.ts" } }], "toolUse"), + ); + const toolResult = entry({ + role: "toolResult", + toolCallId: "call-1", + toolName: "read", + content: [{ type: "text", text: "file contents" }], + isError: false, + timestamp: Date.now(), + }); + const finalAssistant = entry(assistant([{ type: "text", text: "done" }])); + const nextUser = entry(user("next turn")); + + const result = plan([firstUser, toolCall, toolResult, finalAssistant, nextUser], 1_000); + + expect(result.status).toBe("ready"); + if (result.status !== "ready") return; + expect(result.chunk.preparation.firstKeptEntryId).toBe(nextUser.id); + expect(result.chunk.preparation.messagesToSummarize).toEqual([ + firstUser.message, + toolCall.message, + toolResult.message, + finalAssistant.message, + ]); + }); + + it("returns an explicit no-fit result for one oversized oldest entry", () => { + const oversizedUser = entry(user("x".repeat(4_000))); + const newestUser = entry(user("keep me")); + + const result = plan([oversizedUser, newestUser], 500); + + expect(result).toMatchObject({ + status: "no-fit", + reason: "single-group-too-large", + entryId: oversizedUser.id, + budgetTokens: 500, + }); + if (result.status === "no-fit") { + expect(result.requiredTokens).toBeGreaterThan(result.budgetTokens); + } + }); +}); diff --git a/packages/coding-agent/test/suite/agent-session-staged-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-staged-compaction.test.ts new file mode 100644 index 000000000..aafc9d3cc --- /dev/null +++ b/packages/coding-agent/test/suite/agent-session-staged-compaction.test.ts @@ -0,0 +1,213 @@ +import type { AssistantMessage } from "@earendil-works/pi-ai"; +import { fauxAssistantMessage } from "@earendil-works/pi-ai"; +import { afterEach, describe, expect, it } from "vitest"; +import type { ExtensionAPI } from "../../src/core/extensions/index.ts"; +import { createHarness, type Harness } from "./harness.ts"; + +const CONTEXT_WINDOW = 2_000; +const RESERVE_TOKENS = 400; +const OVERSIZED_ONE_SHOT_SUMMARY = "oversized summary ".repeat(500); + +function assistant(text: string): AssistantMessage { + return { + ...fauxAssistantMessage(text), + usage: { + input: 0, + output: 0, + cacheRead: 0, + cacheWrite: 0, + totalTokens: 0, + cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 }, + }, + }; +} + +function seedTurns(harness: Harness, turnCount: number, charsPerMessage = 400): void { + const baseTimestamp = Date.now() - turnCount * 2_000; + for (let index = 0; index < turnCount; index++) { + harness.sessionManager.appendMessage({ + role: "user", + content: [{ type: "text", text: `user-${index}:${"u".repeat(charsPerMessage)}` }], + timestamp: baseTimestamp + index * 2_000, + }); + harness.sessionManager.appendMessage( + Object.assign(assistant(`assistant-${index}:${"a".repeat(charsPerMessage)}`), { + timestamp: baseTimestamp + index * 2_000 + 1_000, + }), + ); + } + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; +} + +function stagedSummaryExtension(options?: { failAtStage?: number; stageCalls?: number[] }) { + return (pi: ExtensionAPI): void => { + pi.on("session_before_compact", async (event) => { + if (!event.stage) { + return { + compaction: { + summary: OVERSIZED_ONE_SHOT_SUMMARY, + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + }; + } + + options?.stageCalls?.push(event.stage.index); + if (event.stage.index === options?.failAtStage) { + return { cancel: true, reason: `stage ${event.stage.index} failed` }; + } + + return { + compaction: { + summary: `accepted stage ${event.stage.index}`, + firstKeptEntryId: event.preparation.firstKeptEntryId, + tokensBefore: event.preparation.tokensBefore, + }, + }; + }); + }; +} + +async function runOverflowRecovery(harness: Harness): Promise { + const run = Reflect.get(harness.session, "_runAutoCompaction"); + if (typeof run !== "function") throw new Error("AgentSession._runAutoCompaction is unavailable"); + return await run.call(harness.session, "overflow", true); +} + +async function createStagedHarness(extension: (pi: ExtensionAPI) => void): Promise { + return await createHarness({ + models: [{ id: "faux-staged", contextWindow: CONTEXT_WINDOW }], + settings: { + compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: RESERVE_TOKENS }, + }, + extensionFactories: [extension], + }); +} + +describe("AgentSession staged compaction recovery", () => { + const harnesses: Harness[] = []; + + afterEach(() => { + while (harnesses.length > 0) harnesses.pop()?.cleanup(); + }); + + it("durably advances through multiple stages and finishes below budget", async () => { + const stageCalls: number[] = []; + const harness = await createStagedHarness(stagedSummaryExtension({ stageCalls })); + harnesses.push(harness); + seedTurns(harness, 14); + + const result = await harness.session.compact(); + + const entries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + const acceptedStages = harness + .eventsOfType("compaction_end") + .filter((event) => event.accepted && event.stage !== undefined); + expect(stageCalls.length).toBeGreaterThanOrEqual(2); + expect(entries).toHaveLength(stageCalls.length); + expect(new Set(entries.map((entry) => entry.firstKeptEntryId)).size).toBe(entries.length); + expect(acceptedStages.at(-1)?.stage?.final).toBe(true); + expect(acceptedStages.slice(0, -1).every((event) => event.stage?.final === false)).toBe(true); + expect(result.summary).toBe(`accepted stage ${stageCalls.at(-1)}`); + expect(result.estimatedTokensAfter).toBeLessThanOrEqual(CONTEXT_WINDOW - RESERVE_TOKENS); + }); + + it("recovers when the one-shot summarization request itself overflows", async () => { + const harness = await createHarness({ + models: [{ id: "faux-staged", contextWindow: CONTEXT_WINDOW }], + settings: { + compaction: { enabled: true, keepRecentTokens: 1, reserveTokens: RESERVE_TOKENS }, + }, + }); + harnesses.push(harness); + seedTurns(harness, 14); + harness.setResponses([ + fauxAssistantMessage("", { + stopReason: "error", + errorMessage: "Your input exceeds the context window of this model", + }), + ...Array.from({ length: 8 }, (_value, index) => fauxAssistantMessage(`core stage ${index + 1}`)), + ]); + + const result = await harness.session.compact(); + + const entries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + expect(harness.faux.state.callCount).toBeGreaterThanOrEqual(3); + expect(entries.length).toBe(harness.faux.state.callCount - 1); + expect(result.summary).toMatch(/^core stage /); + expect(result.estimatedTokensAfter).toBeLessThanOrEqual(CONTEXT_WINDOW - RESERVE_TOKENS); + }); + + it("marks only the final accepted stage for the original prompt retry", async () => { + const stageCalls: number[] = []; + const harness = await createStagedHarness(stagedSummaryExtension({ stageCalls })); + harnesses.push(harness); + seedTurns(harness, 14); + harness.sessionManager.appendMessage({ + ...assistant(""), + stopReason: "error", + errorMessage: "prompt is too long", + timestamp: Date.now(), + }); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; + let continuationCount = 0; + Reflect.set(harness.session, "_scheduleContinuationAfterCurrentEvent", () => { + continuationCount++; + }); + + const recovered = await runOverflowRecovery(harness); + + const acceptedStages = harness + .eventsOfType("compaction_end") + .filter((event) => event.accepted && event.stage !== undefined); + expect(recovered).toBe(true); + expect(acceptedStages.length).toBeGreaterThanOrEqual(2); + expect(acceptedStages.slice(0, -1).every((event) => event.willRetry === false)).toBe(true); + expect(acceptedStages.at(-1)?.willRetry).toBe(true); + expect(continuationCount).toBe(1); + }); + + it("fails explicitly when the oldest complete turn cannot fit one stage", async () => { + const stageCalls: number[] = []; + const harness = await createStagedHarness(stagedSummaryExtension({ stageCalls })); + harnesses.push(harness); + seedTurns(harness, 1, 4_000); + seedTurns(harness, 1, 20); + + await expect(harness.session.compact()).rejects.toThrow(/cannot fit the oldest complete turn/i); + expect(stageCalls).toEqual([]); + expect(harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction")).toHaveLength(0); + }); + + it("keeps an accepted checkpoint when a later stage fails", async () => { + const stageCalls: number[] = []; + const harness = await createStagedHarness(stagedSummaryExtension({ failAtStage: 2, stageCalls })); + harnesses.push(harness); + seedTurns(harness, 14); + + await expect(harness.session.compact()).rejects.toThrow(); + + const entries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + expect(stageCalls).toEqual([1, 2]); + expect(entries).toHaveLength(1); + expect(entries[0]?.summary).toBe("accepted stage 1"); + expect(harness.sessionManager.buildSessionContext().messages[0]).toMatchObject({ + role: "compactionSummary", + summary: "accepted stage 1", + }); + }); + + it("stops at the hard stage bound without retrying the original prompt", async () => { + const stageCalls: number[] = []; + const harness = await createStagedHarness(stagedSummaryExtension({ stageCalls })); + harnesses.push(harness); + seedTurns(harness, 50); + + await expect(harness.session.compact()).rejects.toThrow(); + + const entries = harness.sessionManager.getEntries().filter((entry) => entry.type === "compaction"); + expect(stageCalls).toEqual([1, 2, 3, 4, 5, 6, 7, 8]); + expect(entries).toHaveLength(7); + expect(harness.eventsOfType("compaction_start").filter((event) => event.stage)).toHaveLength(8); + }); +});