diff --git a/packages/coding-agent/CHANGELOG.md b/packages/coding-agent/CHANGELOG.md index 100f5dfc5..a42053d2e 100644 --- a/packages/coding-agent/CHANGELOG.md +++ b/packages/coding-agent/CHANGELOG.md @@ -4,6 +4,7 @@ ### Changed +- Compaction now reuses the current prompt cache when the provider can transiently unload tools with native mid-conversation tool changes, falling back to standalone summarization when unsupported or over the context limit. - `/arminsayshi` now plays a 3D version in fullscreen mode, with one cube per pixel of Armin. The 3D pi logo easter egg on header logo click is replaced by it. ## [1.0.0] - 2026-10-01 diff --git a/packages/coding-agent/docs/compaction.md b/packages/coding-agent/docs/compaction.md index a702f7b6e..c4f079319 100644 --- a/packages/coding-agent/docs/compaction.md +++ b/packages/coding-agent/docs/compaction.md @@ -20,7 +20,7 @@ Pi has two summarization mechanisms: | Compaction | Context exceeds threshold, or `/compact` | Summarize old messages to free up context | | Branch summarization | `/tree` navigation | Preserve context when switching branches | -Both use closely related structured formats and track file operations cumulatively. Summarization requests disable prompt-cache writes because these one-off prompts are unlikely to be reused. +Both use closely related structured formats and track file operations cumulatively. Compaction reuses the current transcript and request settings when the provider can transiently unload tools with native mid-conversation tool changes, allowing the summary request to reuse the existing prompt cache. Tool-free contexts use the same path without requiring native tool changes. Unsupported providers, known overflow recovery, and cache-friendly requests that exceed the context window fall back to a standalone summary request. Branch summaries always use standalone requests with prompt-cache writes disabled. ## Compaction @@ -44,9 +44,9 @@ You can also trigger manually with `/compact [instructions]`, where optional ins 1. **Find cut point**: Walk backwards through the finalized session projection, accumulating token estimates until `keepRecentTokens` (default 20k, configurable in `~/.pi/agent/settings.json` or `/.pi/settings.json`) is reached 2. **Extract messages**: Collect projected messages from the previous kept boundary (or session start) up to the cut point -3. **Generate summary**: Call LLM to summarize with structured format, passing the previous summary as iterative context when present -4. **Append entry**: Save `CompactionEntry` with summary and `firstKeptEntryId` -5. **Rebuilds context**: Session rebuilds the context for the next request, using summary + messages from `firstKeptEntryId` onwards +3. **Generate summary**: When supported, insert a transient cutoff marker before the retained messages, transiently remove all tools, and ask the current model to summarize the full current context through that marker. The request uses the current model, reasoning level, routing session, transport, and provider options. Otherwise, use the standalone serialized summarizer. +4. **Append entry**: Save `CompactionEntry` with summary and `firstKeptEntryId`; the transient marker, tool removal, instruction, and raw response are not persisted +5. **Rebuild context**: Session rebuilds the context for the next request, using summary + messages from `firstKeptEntryId` onwards and the normal tool checkpoint ``` Before compaction: @@ -121,7 +121,7 @@ Split user-message span (one span exceeds budget): turnPrefixMessages = [usr, ass, tool, ass, tool, tool] ``` -For split user-message spans, Pi generates two summaries and merges them: +For split user-message spans, cache-friendly compaction places its cutoff marker at `firstKeptEntryId` and produces one summary of everything before it. Standalone fallback generates two summaries and merges them: 1. **History summary**: Previous context (if any) 2. **User-message-span prefix summary**: The early part of the split user-message span @@ -275,7 +275,7 @@ path/to/changed.ts ### Message Serialization -Before summarization, messages are serialized to text via [`serializeConversation()`](https://github.com/earendil-works/pi/blob/main/packages/coding-agent/src/core/compaction/utils.ts): +Standalone compaction fallback and branch summarization serialize messages to text via [`serializeConversation()`](https://github.com/earendil-works/pi/blob/main/packages/coding-agent/src/core/compaction/utils.ts). Cache-friendly compaction instead sends the current transcript with a transient cutoff marker: ``` [User]: What they said diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 63aa588da..6f9bc5100 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -38,6 +38,7 @@ import type { ImageContent, Model, ProviderHeaders, + SimpleStreamOptions, SystemMessage, TextContent, ToolResultMessage, @@ -70,6 +71,7 @@ import { calculateContextTokens, collectEntriesForBranchSummary, compact, + compactCacheFriendly, estimateContextTokens, estimateProjectedContextTokens, estimateTokens, @@ -2675,9 +2677,37 @@ export class AgentSession { customInstructions: string | undefined, signal: AbortSignal, reason: "manual" | "threshold" | "overflow", + allowCacheFriendly = true, ): Promise { // Resolve the request only when Pi summarizes itself: routing may call models or fail. const request = await this._getSummarizationRequestAuth(model, signal); + const retry = this.settingsManager.getRetrySettings(); + const callbacks = this._summarizationRetryCallbacks({ source: "compaction", reason }); + if (allowCacheFriendly && !isVirtualModel(model)) { + const requestOptions: SimpleStreamOptions = { + apiKey: request.apiKey, + headers: request.headers, + env: request.env, + sessionId: this.agent.sessionId, + transport: this.agent.transport, + thinkingBudgets: this.agent.thinkingBudgets, + maxRetryDelayMs: this.agent.maxRetryDelayMs, + onPayload: this.agent.onPayload, + onResponse: this.agent.onResponse, + onProviderStreamEvent: this.agent.onProviderStreamEvent, + ...(request.model.reasoning && request.thinkingLevel !== "off" ? { reasoning: request.thinkingLevel } : {}), + }; + const result = await compactCacheFriendly(preparation, request.model, { + customInstructions, + signal, + streamFn: this.agent.streamFunction, + requestOptions, + retry, + callbacks, + transformContext: this.agent.transformContext, + }); + if (result) return result; + } return compact( preparation, request.model, @@ -2688,8 +2718,8 @@ export class AgentSession { request.thinkingLevel, this.agent.streamFunction, request.env, - this.settingsManager.getRetrySettings(), - this._summarizationRetryCallbacks({ source: "compaction", reason }), + retry, + callbacks, undefined, // sessionId ); } @@ -2966,7 +2996,7 @@ export class AgentSession { // Case 2: the response completed successfully. Compact, but do not retry because // agent.continue() cannot continue from a completed assistant response. if (!willRetry) { - return await this._runAutoCompaction("overflow", false); + return await this._runAutoCompaction("overflow", false, contextOverflow); } if (this._overflowRecoveryAttempted) { @@ -2994,7 +3024,7 @@ export class AgentSession { // Persistently omit the selected final attempt before post-run recovery compaction. this._overflowRecoveryAttempted = true; this._omitRecoveryAttempt(assistantMessage, toolResults); - const retry = await this._runAutoCompaction("overflow", willRetry); + const retry = await this._runAutoCompaction("overflow", willRetry, contextOverflow); if (retry) this._failedResponse = assistantMessage; return retry; } @@ -3047,7 +3077,11 @@ export class AgentSession { * @param willRetry Whether to continue the interrupted turn after overflow compaction * @returns Whether the post-run loop should call `agent.continue()` */ - private async _runAutoCompaction(reason: "overflow" | "threshold", willRetry: boolean): Promise { + private async _runAutoCompaction( + reason: "overflow" | "threshold", + willRetry: boolean, + knownContextOverflow = false, + ): Promise { const model = this.model; const settings = this.settingsManager.getCompactionSettings(model); let abortController: AbortController | undefined; @@ -3118,6 +3152,7 @@ export class AgentSession { undefined, abortController.signal, reason, + !knownContextOverflow, ); summary = compactResult.summary; firstKeptEntryId = compactResult.firstKeptEntryId; diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 6bcfbb2b5..9989fa232 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -9,6 +9,10 @@ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-a import { contentText, getCurrentSystemMessage, + getCurrentTools, + getInitialSystemMessage, + hasToolRedefinitions, + isContextOverflow, normalizeContext, type RetryCallbacks, type RetryPolicy, @@ -113,6 +117,16 @@ export interface CompactionResult { details?: T; } +export interface CacheFriendlyCompactionOptions { + customInstructions?: string; + signal?: AbortSignal; + streamFn?: StreamFn; + requestOptions: SimpleStreamOptions; + retry?: RetryPolicy; + callbacks?: RetryCallbacks; + transformContext?: (messages: AgentMessage[], signal?: AbortSignal) => Promise; +} + // ============================================================================ // Types // ============================================================================ @@ -610,7 +624,7 @@ function createSummarizationOptions( } /** - * Shared choke point for every compaction/branch-summary summarization call. Wraps the + * Shared choke point for standalone compaction and branch-summary calls. Wraps the * single LLM call in {@link retryAssistantCall} so transient stream drops (e.g. * `terminated`, socket close) honor the configured retry policy instead of failing * the whole compaction on the first attempt. Deterministic errors and aborts return @@ -772,6 +786,10 @@ export async function generateSummaryWithUsage( export interface CompactionPreparation { /** UUID of first entry to keep */ firstKeptEntryId: string; + /** Complete projected context before compaction. */ + contextMessages?: AgentMessage[]; + /** Message index in contextMessages where the retained suffix begins. */ + firstKeptMessageIndex?: number; /** Messages that will be summarized and discarded */ messagesToSummarize: AgentMessage[]; /** Messages that will be turned into turn prefix summary (if splitting) */ @@ -923,8 +941,14 @@ export function prepareCompaction( } } + const firstKeptMessageIndex = projectedEntries + .slice(0, cutPoint.firstKeptEntryIndex) + .reduce((count, entry) => count + entry.messages.length, 0); + return { firstKeptEntryId, + contextMessages: projection.messages, + firstKeptMessageIndex, messagesToSummarize, turnPrefixMessages, isSplitTurn: cutPoint.isSplitTurn, @@ -939,6 +963,109 @@ export function prepareCompaction( // Main compaction function // ============================================================================ +const CACHE_FRIENDLY_COMPACTION_BOUNDARY = ` +Messages before this marker must be summarized. Messages after this marker are recent context that will be retained verbatim. Do not repeat the recent messages in the summary. +`; + +const CACHE_FRIENDLY_COMPACTION_PROMPT = `Create a structured context checkpoint for the conversation before the marker. + +The messages after the marker are provided only so you understand where the retained conversation continues. They will remain verbatim after compaction. Do not summarize or repeat them. + +Do not continue the conversation or answer any request from it. Output only the checkpoint summary. + +${SUMMARIZATION_PROMPT}`; + +function supportsNativeToolUnload(model: Model, context: TranscriptContext): boolean { + if (model.api !== "anthropic-messages") return false; + const anthropicModel = model as Model<"anthropic-messages">; + return ( + anthropicModel.compat?.supportsMidConvoSystemMessages === true && + anthropicModel.compat?.supportsMidConvoToolChanges === true && + (getInitialSystemMessage(context.messages)?.toolsAdded?.length ?? 0) > 0 && + !hasToolRedefinitions(context.messages) + ); +} + +/** + * Summarize the current transcript in place so the provider can reuse its prompt cache. + * Returns undefined when the transport cannot unload tools or the expanded request overflows. + */ +export async function compactCacheFriendly( + preparation: CompactionPreparation, + model: Model, + options: CacheFriendlyCompactionOptions, +): Promise { + const { contextMessages, firstKeptMessageIndex } = preparation; + if (!contextMessages || firstKeptMessageIndex === undefined) return undefined; + + const currentContext = normalizeContext({ messages: convertToLlm(contextMessages) }); + const currentTools = getCurrentTools(currentContext.messages); + if (currentTools.length > 0 && !supportsNativeToolUnload(model, currentContext)) return undefined; + const timestamp = Date.now(); + const boundary: AgentMessage = currentTools.length + ? { role: "system", content: CACHE_FRIENDLY_COMPACTION_BOUNDARY, timestamp } + : { role: "user", content: [{ type: "text", text: CACHE_FRIENDLY_COMPACTION_BOUNDARY }], timestamp }; + let prompt = CACHE_FRIENDLY_COMPACTION_PROMPT; + if (options.customInstructions) prompt += `\n\nAdditional focus: ${options.customInstructions}`; + const transientMessages: AgentMessage[] = [ + ...contextMessages.slice(0, firstKeptMessageIndex), + boundary, + ...contextMessages.slice(firstKeptMessageIndex), + ...(currentTools.length > 0 + ? [ + { + role: "system" as const, + content: "", + toolsRemoved: currentTools.map((tool) => ({ name: tool.name })), + timestamp, + }, + ] + : []), + { role: "user", content: [{ type: "text", text: prompt }], timestamp }, + ]; + const transformed = options.transformContext + ? await options.transformContext(transientMessages, options.signal) + : transientMessages; + const context = normalizeContext({ messages: convertToLlm(transformed) }); + const hasBoundary = context.messages.some( + (message) => + message.role !== "assistant" && contentText(message.content).includes(CACHE_FRIENDLY_COMPACTION_BOUNDARY), + ); + const hasPrompt = context.messages.some( + (message) => message.role === "user" && contentText(message.content).includes(CACHE_FRIENDLY_COMPACTION_PROMPT), + ); + if ( + !hasBoundary || + !hasPrompt || + getCurrentTools(context.messages).length > 0 || + (currentTools.length > 0 && !supportsNativeToolUnload(model, context)) + ) { + return undefined; + } + + const requestOptions: SimpleStreamOptions = { ...options.requestOptions, signal: options.signal }; + const produce = async (): Promise => + options.streamFn + ? (await options.streamFn(model, context, requestOptions)).result() + : completeSimple(model, context, requestOptions); + const response = await retryAssistantCall(produce, options.retry, options.signal, options.callbacks); + if (response.stopReason === "error" && isContextOverflow(response)) return undefined; + const failure = getSummarizationFailure(response, "Summarization"); + if (failure) throw new Error(failure); + if (response.content.some((block) => block.type === "toolCall")) { + throw new Error("Summarization attempted to call a tool"); + } + + const { readFiles, modifiedFiles } = computeFileLists(preparation.fileOps); + return { + summary: contentText(response.content) + formatFileOperations(readFiles, modifiedFiles), + firstKeptEntryId: preparation.firstKeptEntryId, + tokensBefore: preparation.tokensBefore, + usage: response.usage, + details: { readFiles, modifiedFiles } as CompactionDetails, + }; +} + const TURN_PREFIX_SUMMARIZATION_PROMPT = `The messages above are earlier context from an ongoing conversation. Later messages are stored separately and do not need to be reconstructed. Create a concise checkpoint of the user's request and the progress shown above. This checkpoint will be placed before the later messages so the conversation can continue with the necessary context. diff --git a/packages/coding-agent/test/suite/agent-session-compaction.test.ts b/packages/coding-agent/test/suite/agent-session-compaction.test.ts index 3f72a3959..1089eb857 100644 --- a/packages/coding-agent/test/suite/agent-session-compaction.test.ts +++ b/packages/coding-agent/test/suite/agent-session-compaction.test.ts @@ -338,6 +338,13 @@ describe("AgentSession compaction characterization", () => { const harness = await createHarness({ settings: { compaction: { keepRecentTokens: 1 } } }); harnesses.push(harness); seedCompactableSession(harness); + harness.sessionManager.appendMessage({ + role: "system", + content: "", + toolsAdded: [{ name: "unsupported_tool", description: "unsupported", parameters: Type.Object({}) }], + timestamp: Date.now(), + }); + harness.session.agent.state.messages = harness.sessionManager.buildSessionContext().messages; const transformContext = vi.fn(async (messages: AgentMessage[]) => messages); harness.session.agent.transformContext = transformContext; @@ -363,6 +370,123 @@ describe("AgentSession compaction characterization", () => { expect(requestOptions?.transport).toBeUndefined(); }); + it("reuses the current context and transiently unloads tools for supported compaction", async () => { + const cacheTool: AgentTool = { + name: "cache_tool", + label: "Cache tool", + description: "A tool that must be unavailable while compacting", + parameters: Type.Object({}), + execute: async () => ({ content: [{ type: "text", text: "unused" }], details: {} }), + }; + const harness = await createHarness({ + settings: { compaction: { keepRecentTokens: 1 } }, + tools: [cacheTool], + initialActiveToolNames: [cacheTool.name], + }); + harnesses.push(harness); + harness.setResponses([fauxAssistantMessage("first answer"), fauxAssistantMessage("second answer")]); + await harness.session.prompt("first request"); + await harness.session.prompt("second request"); + harness.session.agent.state.model = { + ...harness.getModel(), + api: "anthropic-messages", + compat: { supportsMidConvoSystemMessages: true, supportsMidConvoToolChanges: true }, + }; + harness.session.agent.sessionId = "cache-friendly-session"; + harness.session.agent.transport = "websocket"; + + let requestContext: TranscriptContext | undefined; + let requestOptions: SimpleStreamOptions | undefined; + const getStreamCallCount = useSummaryStreamFn(harness, "cache-friendly summary", (context, options) => { + requestContext = context; + requestOptions = options; + }); + const result = await harness.session.compact("preserve cache details"); + + expect(getStreamCallCount()).toBe(1); + expect(result.summary).toContain("cache-friendly summary"); + expect(getCurrentSystemPrompt(requestContext?.messages ?? [])).toContain(harness.session.systemPrompt); + expect(getCurrentTools(requestContext?.messages ?? [])).toEqual([]); + expect( + requestContext?.messages.some((message) => message.role === "system" && message.toolsRemoved?.length), + ).toBe(true); + const serialized = JSON.stringify(requestContext?.messages); + expect(serialized).toContain("first request"); + expect(serialized).toContain("second request"); + expect(serialized).toContain(""); + expect(serialized).toContain("Additional focus: preserve cache details"); + expect(requestOptions).toMatchObject({ + sessionId: "cache-friendly-session", + transport: "websocket", + }); + expect(requestOptions?.cacheRetention).toBeUndefined(); + expect(requestOptions?.maxTokens).toBeUndefined(); + expect(JSON.stringify(harness.session.messages)).not.toContain(""); + }); + + it("uses cache-friendly compaction without native tool changes when no tools are active", async () => { + const harness = await createHarness({ + settings: { compaction: { keepRecentTokens: 1 } }, + initialActiveToolNames: [], + }); + harnesses.push(harness); + seedCompactableSession(harness); + let requestContext: TranscriptContext | undefined; + const getStreamCallCount = useSummaryStreamFn(harness, "tool-free summary", (context) => { + requestContext = context; + }); + + const result = await harness.session.compact(); + + expect(getStreamCallCount()).toBe(1); + expect(result.summary).toContain("tool-free summary"); + expect(getCurrentTools(requestContext?.messages ?? [])).toEqual([]); + expect( + requestContext?.messages.some((message) => message.role === "system" && message.toolsRemoved?.length), + ).toBe(false); + expect(JSON.stringify(requestContext?.messages)).toContain(""); + }); + + it("falls back to standalone compaction when the full cached context overflows", async () => { + const harness = await createHarness({ + settings: { compaction: { keepRecentTokens: 1 } }, + initialActiveToolNames: [], + }); + harnesses.push(harness); + seedCompactableSession(harness); + const contexts: TranscriptContext[] = []; + harness.session.agent.streamFunction = (model, context) => { + contexts.push(context); + const stream = createAssistantMessageEventStream(); + queueMicrotask(() => { + const message = { + ...fauxAssistantMessage(contexts.length === 1 ? "" : "standalone fallback", { + stopReason: contexts.length === 1 ? "error" : "stop", + errorMessage: contexts.length === 1 ? "prompt is too long" : undefined, + }), + api: model.api, + provider: model.provider, + model: model.id, + usage: createUsage(10), + }; + if (message.stopReason === "error") { + stream.push({ type: "error", reason: "error", error: message }); + } else { + stream.push({ type: "done", reason: "stop", message }); + } + }); + return stream; + }; + + const result = await harness.session.compact(); + + expect(result.summary).toContain("standalone fallback"); + expect(contexts).toHaveLength(2); + expect(JSON.stringify(contexts[0]?.messages)).toContain(""); + expect(JSON.stringify(contexts[1]?.messages)).not.toContain(""); + expect(getCurrentSystemPrompt(contexts[1]?.messages ?? [])).not.toBe(harness.session.systemPrompt); + }); + it("persists usage from pi-generated manual compaction", async () => { const harness = await createHarness({ withConfiguredAuth: false }); harnesses.push(harness);