From cff1cf52c6c73ef873dcb6148238ed68f69e1ead Mon Sep 17 00:00:00 2001 From: Vegard Stikbakke Date: Tue, 18 Aug 2026 14:25:28 +0200 Subject: [PATCH] feat(coding-agent): add cache-friendly compaction primitives Add provider-context prefix preparation and optional summary request settings needed by cache-friendly compaction. Existing AgentSession callers continue to pass undefined, so standalone summarization remains the default until the feature is wired in. --- packages/ai/src/index.ts | 1 + packages/coding-agent/docs/compaction.md | 2 + .../coding-agent/src/core/agent-session.ts | 2 +- .../src/core/compaction/compaction.ts | 255 +++++++++++++++--- packages/coding-agent/src/index.ts | 1 + .../test/compaction-summary-reasoning.test.ts | 227 +++++++++++++++- packages/coding-agent/test/compaction.test.ts | 85 +++++- 7 files changed, 528 insertions(+), 45 deletions(-) diff --git a/packages/ai/src/index.ts b/packages/ai/src/index.ts index 82bac9738..64142c999 100644 --- a/packages/ai/src/index.ts +++ b/packages/ai/src/index.ts @@ -37,6 +37,7 @@ export * from "./providers/faux.ts"; export * from "./session-resources.ts"; export * from "./types.ts"; export * from "./utils/diagnostics.ts"; +export { estimateContextTokens } from "./utils/estimate.ts"; export * from "./utils/event-stream.ts"; export * from "./utils/json-parse.ts"; export * from "./utils/overflow.ts"; diff --git a/packages/coding-agent/docs/compaction.md b/packages/coding-agent/docs/compaction.md index 3a9dff1c6..1ed15fcd6 100644 --- a/packages/coding-agent/docs/compaction.md +++ b/packages/coding-agent/docs/compaction.md @@ -281,7 +281,9 @@ pi.on("session_before_compact", async (event, ctx) => { const { preparation, branchEntries, customInstructions, reason, willRetry, signal } = event; // preparation.messagesToSummarize - messages to summarize + // preparation.sourceMessages - optional active-context history prefix, including any previous summary // preparation.turnPrefixMessages - split turn prefix (if isSplitTurn) + // preparation.turnPrefixSourceMessages - optional active-context prefix through the split turn // preparation.previousSummary - previous compaction summary // preparation.fileOps - extracted file operations // preparation.tokensBefore - context tokens before compaction diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 74280d564..4d02d1943 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -1813,7 +1813,7 @@ export class AgentSession { env, this.settingsManager.getRetrySettings(), this._summarizationRetryCallbacks({ source: "compaction", reason }), - undefined, // sessionId + undefined, // cacheFriendly ); } diff --git a/packages/coding-agent/src/core/compaction/compaction.ts b/packages/coding-agent/src/core/compaction/compaction.ts index 5fbeefc2a..a2a1c063f 100644 --- a/packages/coding-agent/src/core/compaction/compaction.ts +++ b/packages/coding-agent/src/core/compaction/compaction.ts @@ -6,7 +6,14 @@ */ import type { AgentMessage, StreamFn, ThinkingLevel } from "@earendil-works/pi-agent-core"; -import { contentText, type RetryCallbacks, type RetryPolicy, retryAssistantCall, uuidv7 } from "@earendil-works/pi-ai"; +import { + contentText, + estimateContextTokens as estimateProviderContextTokens, + type RetryCallbacks, + type RetryPolicy, + retryAssistantCall, + uuidv7, +} from "@earendil-works/pi-ai"; import type { AssistantMessage, Context, Model, SimpleStreamOptions, Usage } from "@earendil-works/pi-ai/compat"; import { completeSimple } from "@earendil-works/pi-ai/compat"; import { convertToLlm } from "../messages.ts"; @@ -95,6 +102,27 @@ function getMessageFromEntryForCompaction(entry: SessionEntry): AgentMessage | u return sessionEntryToContextMessages(entry)[0]; } +/** Build an active-context prefix, placing the previous compaction summary before its retained messages. */ +function collectSourceMessages( + entries: SessionEntry[], + startIndex: number, + endIndex: number, + previousCompactionIndex: number, +): AgentMessage[] { + const messages: AgentMessage[] = []; + if (previousCompactionIndex >= 0) { + messages.push(...sessionEntryToContextMessages(entries[previousCompactionIndex])); + } + for (let i = startIndex; i < endIndex; i++) { + // The latest compaction was moved to the front above. Keep older compaction + // entries because buildSessionContext retains them in the active provider prefix. + if (i !== previousCompactionIndex) { + messages.push(...sessionEntryToContextMessages(entries[i])); + } + } + return messages; +} + /** Result from compact() - SessionManager adds uuid/parentUuid when saving */ export interface CompactionResult { summary: string; @@ -140,6 +168,19 @@ export interface CompactionSettings { keepRecentTokens: number; } +/** Active provider contexts and request settings used to preserve cacheable compaction prefixes. */ +export interface CacheFriendlySummaryOptions { + /** Exact provider context prefix containing the history to summarize. */ + sourceContext?: Context; + /** Exact provider context prefix containing a split turn's prefix. */ + turnPrefixSourceContext?: Context; + /** Provider request settings copied from the active agent request path. */ + requestOptions?: Pick< + SimpleStreamOptions, + "sessionId" | "onPayload" | "onResponse" | "transport" | "thinkingBudgets" | "maxRetryDelayMs" + >; +} + export const DEFAULT_COMPACTION_SETTINGS: CompactionSettings = { enabled: true, reserveTokens: 16384, @@ -549,6 +590,10 @@ const UPDATE_SUMMARIZATION_PROMPT = `The messages above are NEW conversation mes ${UPDATE_SUMMARIZATION_INSTRUCTIONS}`; +const SOURCE_CONTEXT_UPDATE_SUMMARIZATION_PROMPT = `The messages above contain an existing structured summary of earlier conversation history followed by NEW conversation messages. + +${UPDATE_SUMMARIZATION_INSTRUCTIONS}`; + function createSummarizationOptions( model: Model, maxTokens: number, @@ -557,9 +602,18 @@ function createSummarizationOptions( env: Record | undefined, signal: AbortSignal | undefined, thinkingLevel: ThinkingLevel | undefined, - sessionId: string | undefined, + requestOptions: CacheFriendlySummaryOptions["requestOptions"] | undefined, + cacheRetention: SimpleStreamOptions["cacheRetention"] | undefined, ): SimpleStreamOptions { - const options: SimpleStreamOptions = { maxTokens, signal, apiKey, headers, env, sessionId }; + const options: SimpleStreamOptions = { + ...requestOptions, + maxTokens, + signal, + apiKey, + headers, + env, + cacheRetention, + }; const refusalFallbacks = getAnthropicSummarizationFallback(model); if (refusalFallbacks) { options.refusalFallbacks = refusalFallbacks; @@ -585,12 +639,15 @@ export async function completeSummarization( retry?: RetryPolicy, callbacks?: RetryCallbacks, ): Promise { - // Avoid cache writes for one-off summaries. Reuse caller-supplied routing when available; - // callers without a session ID, including branch summaries, receive a fresh routing ID. + // Standalone one-off summaries default to no prompt caching. Cache-friendly callers + // explicitly request short retention so providers can reuse the active prefix. + // Callers without a session ID, including standalone branch summaries, receive a fresh routing ID. const requestOptions: SimpleStreamOptions = { ...options, - cacheRetention: "none", + cacheRetention: options.cacheRetention ?? "none", sessionId: options.sessionId ?? uuidv7(), + // Anthropic invalidates the messages cache when tool_choice changes. Its 20-content-block lookup + // can still reuse an earlier user-message checkpoint, but long tool-heavy turns may need reprocessing. toolChoice: "none", }; const produce = async (): Promise => @@ -618,7 +675,7 @@ export async function generateSummary( env?: Record, retry?: RetryPolicy, callbacks?: RetryCallbacks, - sessionId?: string, + cacheFriendly?: Pick, ): Promise { return ( await generateSummaryWithUsage( @@ -635,25 +692,47 @@ export async function generateSummary( env, retry, callbacks, - sessionId, + cacheFriendly, ) ).text; } -/** Build the provider context for a standalone summary request. */ -function buildSummarizationContext(promptText: string): Context { +/** Build a standalone summary request or append its instruction to an existing provider context. */ +function buildSummarizationContext(promptText: string, sourceContext?: Context): Context { + const instructionMessage = { + role: "user" as const, + content: [{ type: "text" as const, text: promptText }], + timestamp: Date.now(), + }; + + if (sourceContext) { + return { + ...sourceContext, + messages: [...sourceContext.messages, instructionMessage], + }; + } + return { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, - messages: [ - { - role: "user", - content: [{ type: "text", text: promptText }], - timestamp: Date.now(), - }, - ], + messages: [instructionMessage], }; } +/** + * Extra room for provider framing and tokenizer variance omitted by the heuristic context estimate. + * This matches the 4096-token margin used when normal simple requests clamp maxTokens to their context window. + */ +const CACHE_FRIENDLY_CONTEXT_SAFETY_TOKENS = 4096; + +/** Whether the source context leaves room for the requested summary output and provider safety margin. */ +function cacheFriendlyContextFits(model: Model, context: Context, maxTokens: number): boolean { + return ( + model.contextWindow <= 0 || + estimateProviderContextTokens(context).tokens + maxTokens + CACHE_FRIENDLY_CONTEXT_SAFETY_TOKENS <= + model.contextWindow + ); +} + /** Generate or update a conversation summary and return its provider usage. */ export async function generateSummaryWithUsage( currentMessages: AgentMessage[], @@ -669,31 +748,55 @@ export async function generateSummaryWithUsage( env?: Record, retry?: RetryPolicy, callbacks?: RetryCallbacks, - sessionId?: string, + cacheFriendly?: Pick, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.8 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, ); + // Provider-visible history prefix to reuse instead of serializing messages into a standalone prompt. + let sourceContext = cacheFriendly?.sourceContext; - // Use update prompt if we have a previous summary, otherwise initial prompt - let basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; + // Cache-friendly source contexts already contain the previous compaction summary, + // but still need iterative-update instructions so prior information is preserved. + let basePrompt = previousSummary + ? sourceContext + ? SOURCE_CONTEXT_UPDATE_SUMMARIZATION_PROMPT + : UPDATE_SUMMARIZATION_PROMPT + : SUMMARIZATION_PROMPT; if (customInstructions) { basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; } - // Serialize conversation to text so model doesn't try to continue it - // Convert to LLM messages first (handles custom types like bashExecution, custom, etc.) - const llmMessages = convertToLlm(currentMessages); - const conversationText = serializeConversation(llmMessages); + // A valid prefix can still be too large to leave the intended summary budget, + // especially during overflow recovery or after switching to a smaller-context model. + // In that case, use standalone serialization, which truncates large tool results. + if ( + sourceContext && + !cacheFriendlyContextFits(model, buildSummarizationContext(basePrompt, sourceContext), maxTokens) + ) { + sourceContext = undefined; + basePrompt = previousSummary ? UPDATE_SUMMARIZATION_PROMPT : SUMMARIZATION_PROMPT; + if (customInstructions) { + basePrompt = `${basePrompt}\n\nAdditional focus: ${customInstructions}`; + } + } - // Build the prompt with conversation wrapped in tags - let promptText = `\n${conversationText}\n\n\n`; - if (previousSummary) { + // Source contexts already contain the conversation. Standalone requests serialize it + // so the model treats those messages as data rather than continuing the conversation. + let promptText = ""; + if (!sourceContext) { + const llmMessages = convertToLlm(currentMessages); + const conversationText = serializeConversation(llmMessages); + promptText = `\n${conversationText}\n\n\n`; + } + if (previousSummary && !sourceContext) { promptText += `\n${previousSummary}\n\n\n`; } promptText += basePrompt; + const sourceRequestOptions = sourceContext ? cacheFriendly?.requestOptions : undefined; + const sourceCacheRetention = sourceContext ? "short" : undefined; const completionOptions = createSummarizationOptions( model, maxTokens, @@ -702,12 +805,13 @@ export async function generateSummaryWithUsage( env, signal, thinkingLevel, - sessionId, + sourceRequestOptions, + sourceCacheRetention, ); const response = await completeSummarization( model, - buildSummarizationContext(promptText), + buildSummarizationContext(promptText, sourceContext), completionOptions, streamFn, retry, @@ -735,8 +839,15 @@ export interface CompactionPreparation { firstKeptEntryId: string; /** Messages that will be summarized and discarded */ messagesToSummarize: AgentMessage[]; + /** + * Active-context prefix for the history summary. + * Includes the previous compaction summary before messages retained by that compaction. + */ + sourceMessages?: AgentMessage[]; /** Messages that will be turned into turn prefix summary (if splitting) */ turnPrefixMessages: AgentMessage[]; + /** Active-context prefix through the split-turn prefix, or empty when not splitting. */ + turnPrefixSourceMessages?: AgentMessage[]; /** Whether this is a split turn (cut point in middle of turn) */ isSplitTurn: boolean; tokensBefore: number; @@ -794,6 +905,8 @@ export function prepareCompaction( if (msg) messagesToSummarize.push(msg); } + const sourceMessages = collectSourceMessages(pathEntries, boundaryStart, historyEnd, prevCompactionIndex); + // Messages for turn prefix summary (if splitting a turn) const turnPrefixMessages: AgentMessage[] = []; if (cutPoint.isSplitTurn) { @@ -802,6 +915,9 @@ export function prepareCompaction( if (msg) turnPrefixMessages.push(msg); } } + const turnPrefixSourceMessages = cutPoint.isSplitTurn + ? collectSourceMessages(pathEntries, boundaryStart, cutPoint.firstKeptEntryIndex, prevCompactionIndex) + : []; if (messagesToSummarize.length === 0 && turnPrefixMessages.length === 0) { return undefined; @@ -820,7 +936,9 @@ export function prepareCompaction( return { firstKeptEntryId, messagesToSummarize, + sourceMessages, turnPrefixMessages, + turnPrefixSourceMessages, isSplitTurn: cutPoint.isSplitTurn, tokensBefore, previousSummary, @@ -848,13 +966,30 @@ Summarize the prefix to provide context for the retained suffix: Be concise. Focus on what's needed to understand the kept suffix.`; +const SOURCE_CONTEXT_TURN_PREFIX_SUMMARIZATION_PROMPT = `The final turn in the source conversation was too large to keep in full. Its SUFFIX (recent work) is retained. + +The source conversation may also contain complete earlier turns for background. Summarize only the final, incomplete turn. It begins with the last user-role request before this instruction. Do not summarize earlier turns except for details needed to understand this final turn's prefix. + +Summarize the prefix to provide context for the retained suffix: + +## Original Request +[What did the user ask for in this turn?] + +## Early Progress +- [Key decisions and work done in the prefix] + +## Context for Suffix +- [Information needed to understand the retained recent work] + +Be concise. Focus on what's needed to understand the kept suffix.`; + /** * Generate summaries for compaction using prepared data. * Returns CompactionResult - SessionManager adds uuid/parentUuid when saving. * * @param preparation - Pre-calculated preparation from prepareCompaction() * @param customInstructions - Optional custom focus for the summary - * @param sessionId - Optional routing session ID forwarded without enabling prompt caching + * @param cacheFriendly - Active provider contexts and request settings for cache-friendly summarization */ export async function compact( preparation: CompactionPreparation, @@ -868,7 +1003,7 @@ export async function compact( env?: Record, retry?: RetryPolicy, callbacks?: RetryCallbacks, - sessionId?: string, + cacheFriendly?: CacheFriendlySummaryOptions, ): Promise { const { firstKeptEntryId, @@ -903,7 +1038,10 @@ export async function compact( env, retry, callbacks, - sessionId, + { + sourceContext: cacheFriendly?.sourceContext, + requestOptions: cacheFriendly?.requestOptions, + }, ); historyText = historyResult.text; historyUsage = historyResult.usage; @@ -920,7 +1058,10 @@ export async function compact( streamFn, retry, callbacks, - sessionId, + { + sourceContext: cacheFriendly?.turnPrefixSourceContext, + requestOptions: cacheFriendly?.requestOptions, + }, ); // Merge into single summary summary = `${historyText}\n\n---\n\n**Turn Context (split turn):**\n\n${turnPrefixResult.text}`; @@ -941,7 +1082,10 @@ export async function compact( env, retry, callbacks, - sessionId, + { + sourceContext: cacheFriendly?.sourceContext, + requestOptions: cacheFriendly?.requestOptions, + }, ); summary = result.text; summaryUsage = result.usage; @@ -979,20 +1123,51 @@ async function generateTurnPrefixSummary( streamFn?: StreamFn, retry?: RetryPolicy, callbacks?: RetryCallbacks, - sessionId?: string, + cacheFriendly?: Pick, ): Promise<{ text: string; usage: Usage }> { const maxTokens = Math.min( Math.floor(0.5 * reserveTokens), model.maxTokens > 0 ? model.maxTokens : Number.POSITIVE_INFINITY, ); // Smaller budget for turn prefix - const llmMessages = convertToLlm(messages); - const conversationText = serializeConversation(llmMessages); - const promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; + // Reuse the provider-visible split-turn prefix only when it leaves room for the summary; + // otherwise serialize and truncate the messages in a standalone prompt. + let sourceContext = cacheFriendly?.sourceContext; + if ( + sourceContext && + !cacheFriendlyContextFits( + model, + buildSummarizationContext(SOURCE_CONTEXT_TURN_PREFIX_SUMMARIZATION_PROMPT, sourceContext), + maxTokens, + ) + ) { + sourceContext = undefined; + } + let promptText: string; + if (sourceContext) { + promptText = SOURCE_CONTEXT_TURN_PREFIX_SUMMARIZATION_PROMPT; + } else { + const llmMessages = convertToLlm(messages); + const conversationText = serializeConversation(llmMessages); + promptText = `\n${conversationText}\n\n\n${TURN_PREFIX_SUMMARIZATION_PROMPT}`; + } + + const sourceRequestOptions = sourceContext ? cacheFriendly?.requestOptions : undefined; + const sourceCacheRetention = sourceContext ? "short" : undefined; const response = await completeSummarization( model, - buildSummarizationContext(promptText), - createSummarizationOptions(model, maxTokens, apiKey, headers, env, signal, thinkingLevel, sessionId), + buildSummarizationContext(promptText, sourceContext), + createSummarizationOptions( + model, + maxTokens, + apiKey, + headers, + env, + signal, + thinkingLevel, + sourceRequestOptions, + sourceCacheRetention, + ), streamFn, retry, callbacks, diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index 6d64d70c8..7c9c596a4 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -28,6 +28,7 @@ export { readStoredCredential } from "./core/auth-storage.ts"; export { type BranchPreparation, type BranchSummaryResult, + type CacheFriendlySummaryOptions, type CollectEntriesResult, type CompactionResult, type CutPointResult, diff --git a/packages/coding-agent/test/compaction-summary-reasoning.test.ts b/packages/coding-agent/test/compaction-summary-reasoning.test.ts index 340d5c1f6..9bafc9d5d 100644 --- a/packages/coding-agent/test/compaction-summary-reasoning.test.ts +++ b/packages/coding-agent/test/compaction-summary-reasoning.test.ts @@ -115,7 +115,7 @@ describe("generateSummary reasoning options", () => { expect(sessionIds[0]).not.toBe(sessionIds[1]); }); - it("honors a caller-supplied routing session without prompt caching", async () => { + it("honors caller-supplied cache retention and routing", async () => { await completeSummarization( createModel(false), { systemPrompt: "Summarize", messages: [] }, @@ -124,7 +124,7 @@ describe("generateSummary reasoning options", () => { expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ sessionId: "current-routing-session", - cacheRetention: "none", + cacheRetention: "long", toolChoice: "none", }); }); @@ -148,6 +148,229 @@ describe("generateSummary reasoning options", () => { expect(prompt).toContain(""); }); + it("appends instructions to a cache-friendly source context", async () => { + const sourceContext: Context = { + systemPrompt: "You are a coding agent.", + messages: [{ role: "user", content: "Previous summary and original request", timestamp: 1 }], + tools: [], + }; + const onPayload = async (payload: unknown) => payload; + const onResponse = async () => {}; + const preparation: CompactionPreparation = { + firstKeptEntryId: "entry-keep", + messagesToSummarize: messages, + turnPrefixMessages: [], + isSplitTurn: false, + tokensBefore: 100, + previousSummary: "Previous summary", + fileOps: { read: new Set(), written: new Set(), edited: new Set() }, + settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, + }; + + await compact( + preparation, + createModel(false), + "test-key", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { + sourceContext, + requestOptions: { + sessionId: "routing-session", + onPayload, + onResponse, + transport: "websocket", + thinkingBudgets: { low: 1234 }, + maxRetryDelayMs: 4321, + }, + }, + ); + + const requestContext = completeSimpleMock.mock.calls[0][1] as Context; + expect(requestContext.systemPrompt).toBe(sourceContext.systemPrompt); + expect(requestContext.tools).toBe(sourceContext.tools); + expect(requestContext.messages.slice(0, -1)).toEqual(sourceContext.messages); + const instruction = JSON.stringify(requestContext.messages.at(-1)); + expect(instruction).toContain("existing structured summary of earlier conversation history"); + expect(instruction).toContain("PRESERVE all existing information from the previous summary"); + expect(instruction).not.toContain(""); + expect(instruction).not.toContain(""); + expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ + cacheRetention: "short", + sessionId: "routing-session", + toolChoice: "none", + transport: "websocket", + thinkingBudgets: { low: 1234 }, + maxRetryDelayMs: 4321, + }); + expect(completeSimpleMock.mock.calls[0][2]?.onPayload).toBe(onPayload); + expect(completeSimpleMock.mock.calls[0][2]?.onResponse).toBe(onResponse); + expect(sourceContext.messages).toHaveLength(1); + }); + + it("falls back to standalone summarization when a cache-friendly source cannot leave the summary budget", async () => { + const oversizedToolResult: AgentMessage = { + role: "toolResult", + toolCallId: "tool-call-1", + toolName: "read", + content: [{ type: "text", text: "x".repeat(800_000) }], + isError: false, + timestamp: 1, + }; + const sourceContext: Context = { + systemPrompt: "You are a coding agent.", + messages: [oversizedToolResult], + tools: [], + }; + const preparation: CompactionPreparation = { + firstKeptEntryId: "entry-keep", + messagesToSummarize: [oversizedToolResult], + turnPrefixMessages: [], + isSplitTurn: false, + tokensBefore: 250_000, + fileOps: { read: new Set(), written: new Set(), edited: new Set() }, + settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, + }; + + await compact( + preparation, + createModel(false), + "test-key", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { sourceContext, requestOptions: { sessionId: "routing-session" } }, + ); + + const requestContext = completeSimpleMock.mock.calls[0][1] as Context; + const prompt = JSON.stringify(requestContext.messages); + expect(requestContext.systemPrompt).not.toBe(sourceContext.systemPrompt); + expect(requestContext.tools).toBeUndefined(); + expect(prompt).toContain(""); + expect(prompt).toContain("more characters truncated"); + expect(prompt).not.toContain("x".repeat(3000)); + expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ cacheRetention: "none", toolChoice: "none" }); + expect(completeSimpleMock.mock.calls[0][2]?.sessionId).not.toBe("routing-session"); + }); + + it("falls back to standalone summarization for an oversized split-turn source", async () => { + const splitUser = { role: "user" as const, content: "Large final request", timestamp: 1 }; + const oversizedToolResult: AgentMessage = { + role: "toolResult", + toolCallId: "tool-call-1", + toolName: "read", + content: [{ type: "text", text: "x".repeat(800_000) }], + isError: false, + timestamp: 2, + }; + const turnPrefixSourceContext: Context = { + systemPrompt: "You are a coding agent.", + messages: [splitUser, oversizedToolResult], + tools: [], + }; + const preparation: CompactionPreparation = { + firstKeptEntryId: "entry-keep", + messagesToSummarize: [], + turnPrefixMessages: [splitUser, oversizedToolResult], + isSplitTurn: true, + tokensBefore: 250_000, + fileOps: { read: new Set(), written: new Set(), edited: new Set() }, + settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, + }; + + await compact( + preparation, + createModel(false), + "test-key", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { turnPrefixSourceContext, requestOptions: { sessionId: "routing-session" } }, + ); + + const requestContext = completeSimpleMock.mock.calls[0][1] as Context; + const prompt = JSON.stringify(requestContext.messages); + expect(requestContext.systemPrompt).not.toBe(turnPrefixSourceContext.systemPrompt); + expect(requestContext.tools).toBeUndefined(); + expect(prompt).toContain("This is the PREFIX of a turn that was too large to keep"); + expect(prompt).toContain("more characters truncated"); + expect(prompt).not.toContain("source conversation may also contain complete earlier turns"); + expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ cacheRetention: "none", toolChoice: "none" }); + expect(completeSimpleMock.mock.calls[0][2]?.sessionId).not.toBe("routing-session"); + }); + + it("limits a cache-friendly split-turn summary to the final incomplete turn", async () => { + const earlierUser = { role: "user" as const, content: "Earlier request", timestamp: 1 }; + const earlierAssistant: AssistantMessage = { + ...mockSummaryResponse, + content: [{ type: "text", text: "Earlier work completed" }], + timestamp: 2, + }; + const splitUser = { role: "user" as const, content: "Large final request", timestamp: 3 }; + const earlyAssistant: AssistantMessage = { + ...mockSummaryResponse, + content: [{ type: "text", text: "Early work in final turn" }], + timestamp: 4, + }; + const turnPrefixSourceContext: Context = { + systemPrompt: "You are a coding agent.", + messages: [earlierUser, earlierAssistant, splitUser, earlyAssistant], + tools: [], + }; + const preparation: CompactionPreparation = { + firstKeptEntryId: "entry-keep", + messagesToSummarize: [], + turnPrefixMessages: [splitUser, earlyAssistant], + isSplitTurn: true, + tokensBefore: 100, + fileOps: { read: new Set(), written: new Set(), edited: new Set() }, + settings: { enabled: true, reserveTokens: 2000, keepRecentTokens: 20 }, + }; + + await compact( + preparation, + createModel(false), + "test-key", + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + undefined, + { turnPrefixSourceContext }, + ); + + const requestContext = completeSimpleMock.mock.calls[0][1] as Context; + expect(requestContext.messages.slice(0, -1)).toEqual(turnPrefixSourceContext.messages); + const instruction = JSON.stringify(requestContext.messages.at(-1)); + expect(instruction).toContain("Summarize only the final, incomplete turn"); + expect(instruction).toContain("last user-role request before this instruction"); + expect(instruction).toContain("Do not summarize earlier turns"); + expect(instruction).not.toContain("Earlier request"); + expect(completeSimpleMock.mock.calls[0][2]).toMatchObject({ + cacheRetention: "short", + toolChoice: "none", + }); + }); + it("rejects tool calls from conversation summaries", async () => { completeSimpleMock.mockResolvedValueOnce(mockToolCallResponse); diff --git a/packages/coding-agent/test/compaction.test.ts b/packages/coding-agent/test/compaction.test.ts index 12c9d0be2..8fe754ce6 100644 --- a/packages/coding-agent/test/compaction.test.ts +++ b/packages/coding-agent/test/compaction.test.ts @@ -5,6 +5,7 @@ import { readFileSync } from "fs"; import { join } from "path"; import { beforeEach, describe, expect, it } from "vitest"; import { + type CompactionPreparation, type CompactionSettings, calculateContextTokens, compact, @@ -150,6 +151,17 @@ function createCustomMessageEntry(content: string): CustomMessageEntry { return entry; } +function requireSourceMessages( + preparation: CompactionPreparation | undefined, +): asserts preparation is CompactionPreparation & { + sourceMessages: AgentMessage[]; + turnPrefixSourceMessages: AgentMessage[]; +} { + if (!preparation?.sourceMessages || !preparation.turnPrefixSourceMessages) { + throw new Error("Expected compaction source messages"); + } +} + function extractText(messages: AgentMessage[]): string { return messages .map((message) => { @@ -461,6 +473,44 @@ describe("buildSessionContext", () => { }); }); +describe("prepareCompaction source messages", () => { + it("provides the active history prefix for a normal compaction", () => { + const u1 = createMessageEntry(createUserMessage("old user ".repeat(20))); + const a1 = createMessageEntry(createAssistantMessage("old assistant ".repeat(20))); + const u2 = createMessageEntry(createUserMessage("recent user")); + const a2 = createMessageEntry(createAssistantMessage("ok")); + const settings: CompactionSettings = { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 3 }; + + const preparation = prepareCompaction([u1, a1, u2, a2], settings); + + expect(preparation).toBeDefined(); + requireSourceMessages(preparation); + expect(preparation.isSplitTurn).toBe(false); + expect(preparation.sourceMessages).toEqual(preparation.messagesToSummarize); + expect(extractText(preparation.sourceMessages)).toContain("old user"); + expect(extractText(preparation.sourceMessages)).not.toContain("recent user"); + expect(preparation.turnPrefixSourceMessages).toEqual([]); + }); + + it("provides the active prefix through a split turn", () => { + const user = createMessageEntry(createUserMessage("large request ".repeat(20))); + const earlyAssistant = createMessageEntry(createAssistantMessage("early work ".repeat(20))); + const keptAssistant = createMessageEntry(createAssistantMessage("kept")); + const settings: CompactionSettings = { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 1 }; + + const preparation = prepareCompaction([user, earlyAssistant, keptAssistant], settings); + + expect(preparation).toBeDefined(); + requireSourceMessages(preparation); + expect(preparation.isSplitTurn).toBe(true); + expect(preparation.messagesToSummarize).toEqual([]); + expect(preparation.sourceMessages).toEqual([]); + expect(preparation.turnPrefixSourceMessages).toEqual(preparation.turnPrefixMessages); + expect(extractText(preparation.turnPrefixSourceMessages)).toContain("large request"); + expect(extractText(preparation.turnPrefixSourceMessages)).not.toContain("kept"); + }); +}); + describe("prepareCompaction with previous compaction", () => { it("should skip repeated compactions when kept messages still fit", () => { const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)")); @@ -479,6 +529,33 @@ describe("prepareCompaction with previous compaction", () => { expect(preparation).toBeUndefined(); }); + it("preserves retained older summaries in the active source prefix", () => { + const u1 = createMessageEntry(createUserMessage("user 1")); + const a1 = createMessageEntry(createAssistantMessage("assistant 1")); + const u2 = createMessageEntry(createUserMessage("user 2 ".repeat(20))); + const a2 = createMessageEntry(createAssistantMessage("assistant 2 ".repeat(20))); + const compaction1 = createCompactionEntry("First summary", u1.id); + const u3 = createMessageEntry(createUserMessage("user 3 ".repeat(20))); + const a3 = createMessageEntry(createAssistantMessage("assistant 3 ".repeat(20))); + const compaction2 = createCompactionEntry("Second summary", u2.id); + const u4 = createMessageEntry(createUserMessage("user 4 ".repeat(20))); + const a4 = createMessageEntry(createAssistantMessage("kept")); + const settings: CompactionSettings = { ...DEFAULT_COMPACTION_SETTINGS, keepRecentTokens: 1 }; + const pathEntries = [u1, a1, u2, a2, compaction1, u3, a3, compaction2, u4, a4]; + + const preparation = prepareCompaction(pathEntries, settings); + + expect(preparation).toBeDefined(); + requireSourceMessages(preparation); + const activeMessages = buildSessionContext(pathEntries).messages; + expect(preparation.sourceMessages).toEqual(activeMessages.slice(0, preparation.sourceMessages.length)); + expect(preparation.turnPrefixSourceMessages).toEqual( + activeMessages.slice(0, preparation.turnPrefixSourceMessages.length), + ); + const sourceSummaries = preparation.sourceMessages.filter((message) => message.role === "compactionSummary"); + expect(sourceSummaries.map((message) => message.summary)).toEqual(["Second summary", "First summary"]); + }); + it("should re-summarize previously kept messages when the recent window moves past them", () => { const u1 = createMessageEntry(createUserMessage("user msg 1 (summarized by compaction1)".repeat(4))); const a1 = createMessageEntry(createAssistantMessage("assistant msg 1".repeat(4))); @@ -497,11 +574,15 @@ describe("prepareCompaction with previous compaction", () => { const preparation = prepareCompaction([u1, a1, u2, a2, u3, a3, compaction1, u4, a4], settings); expect(preparation).toBeDefined(); - const summarizedText = extractText(preparation!.messagesToSummarize); + requireSourceMessages(preparation); + const summarizedText = extractText(preparation.messagesToSummarize); expect(summarizedText).toContain("user msg 2 - kept by compaction1"); expect(summarizedText).toContain("user msg 3 - kept by compaction1"); expect(summarizedText).not.toContain("First summary"); - expect(preparation!.previousSummary).toBe("First summary"); + expect(preparation.previousSummary).toBe("First summary"); + expect(preparation.sourceMessages[0]?.role).toBe("compactionSummary"); + expect(extractText(preparation.sourceMessages)).toContain("First summary"); + expect(preparation.sourceMessages.slice(1)).toEqual(preparation.messagesToSummarize); }); });