diff --git a/packages/agent/src/compaction/branch-summarization.ts b/packages/agent/src/compaction/branch-summarization.ts index d8ab3745f..93901d670 100644 --- a/packages/agent/src/compaction/branch-summarization.ts +++ b/packages/agent/src/compaction/branch-summarization.ts @@ -29,6 +29,7 @@ import { SUMMARIZATION_SYSTEM_PROMPT, serializeConversation, stripReadSelector, + truncateToolResultForSummary, upsertFileOperations, } from "./utils"; @@ -206,6 +207,19 @@ function getMessageFromEntry(entry: SessionEntry): AgentMessage | undefined { } } +function estimateBranchSummaryTokens(message: AgentMessage): number { + if (message.role !== "toolResult") return estimateTokens(message); + const text = message.content + .filter((c): c is { type: "text"; text: string } => c.type === "text") + .map(c => c.text) + .join(""); + if (!text) return 0; + return estimateTokens({ + ...message, + content: [{ type: "text", text: truncateToolResultForSummary(text) }], + }); +} + /** * Prepare entries for summarization with token budget. * @@ -251,7 +265,7 @@ export function prepareBranchEntries(entries: SessionEntry[], tokenBudget: numbe // Extract file ops from assistant messages (tool calls) extractFileOpsFromMessage(message, fileOps); - const tokens = estimateTokens(message); + const tokens = estimateBranchSummaryTokens(message); // Check budget before adding if (tokenBudget > 0 && totalTokens + tokens > tokenBudget) { diff --git a/packages/agent/src/compaction/utils.ts b/packages/agent/src/compaction/utils.ts index ddd743344..43b9df03b 100644 --- a/packages/agent/src/compaction/utils.ts +++ b/packages/agent/src/compaction/utils.ts @@ -199,13 +199,12 @@ export function upsertFileOperations( const TOOL_RESULT_MAX_CHARS = 2000; /** - * Truncate text to a maximum character length for summarization. - * Keeps the beginning and appends a truncation marker. + * Truncate tool results to the same representation used in summarization prompts. */ -function truncateForSummary(text: string, maxChars: number): string { - if (text.length <= maxChars) return text; - const truncatedChars = text.length - maxChars; - return `${text.slice(0, maxChars)}\n\n[... ${truncatedChars} more characters truncated]`; +export function truncateToolResultForSummary(text: string): string { + if (text.length <= TOOL_RESULT_MAX_CHARS) return text; + const truncatedChars = text.length - TOOL_RESULT_MAX_CHARS; + return `${text.slice(0, TOOL_RESULT_MAX_CHARS)}\n\n[... ${truncatedChars} more characters truncated]`; } /** @@ -241,7 +240,7 @@ export function serializeConversation(messages: Message[], dialect?: Dialect): s if (!text) continue; processed.push({ ...msg, - content: [{ type: "text", text: truncateForSummary(text, TOOL_RESULT_MAX_CHARS) }], + content: [{ type: "text", text: truncateToolResultForSummary(text) }], }); continue; } @@ -293,7 +292,7 @@ export function serializeConversation(messages: Message[], dialect?: Dialect): s .map(c => c.text) .join(""); if (content) { - const text = truncateForSummary(content, TOOL_RESULT_MAX_CHARS); + const text = truncateToolResultForSummary(content); parts.push(`[Tool Result]: ${text}`); } } diff --git a/packages/agent/test/branch-summarization.test.ts b/packages/agent/test/branch-summarization.test.ts index d51252515..1cbde0fa8 100644 --- a/packages/agent/test/branch-summarization.test.ts +++ b/packages/agent/test/branch-summarization.test.ts @@ -190,4 +190,44 @@ describe("branch summarization", () => { expect(userMessages[0].content).toBe("OLDER_USEFUL_FACT_4076"); expect(messages.some(m => m.role === "toolResult")).toBe(false); }); + + test("large informative tool results are budgeted after summary truncation", () => { + const informativeBlob = `IMPORTANT_LARGE_TOOL_FACT_4112\n${"x".repeat(20_000)}`; + const entries: SessionEntry[] = [ + { + type: "message", + id: "assistant-1", + parentId: null, + timestamp: new Date(0).toISOString(), + message: { + role: "assistant", + content: [{ type: "toolCall", id: "call-read", name: "read", arguments: { path: "big.txt" } }], + api: "mock", + provider: "mock", + model: "mock-model", + usage: ZERO_USAGE, + stopReason: "toolUse", + timestamp: 0, + }, + }, + { + type: "message", + id: "tool-1", + parentId: "assistant-1", + timestamp: new Date(1).toISOString(), + message: { + role: "toolResult", + toolCallId: "call-read", + toolName: "read", + content: [{ type: "text", text: informativeBlob }], + isError: false, + timestamp: 1, + }, + }, + ]; + + const { messages } = prepareBranchEntries(entries, 700); + + expect(messages.some(m => m.role === "toolResult")).toBe(true); + }); });