diff --git a/src/agents/embedded-agent-runner/run/preemptive-compaction.test.ts b/src/agents/embedded-agent-runner/run/preemptive-compaction.test.ts index 806226c6cabd..f2274486158b 100644 --- a/src/agents/embedded-agent-runner/run/preemptive-compaction.test.ts +++ b/src/agents/embedded-agent-runner/run/preemptive-compaction.test.ts @@ -455,4 +455,133 @@ describe("preemptive-compaction", () => { expect(result.route).toBe("truncate_tool_results_only"); expect(result.shouldCompact).toBe(false); }); + + it("estimates CJK tool results at roughly one token per character", () => { + const cjkText = "中".repeat(85_000); + const toolResultTokens = estimateLlmBoundaryTokenPressure({ + messages: [makeToolResultMessage(cjkText)], + systemPrompt: "sys", + prompt: "continue", + }); + const assistantTokens = estimateLlmBoundaryTokenPressure({ + messages: [makeAssistantHistory(cjkText)], + systemPrompt: "sys", + prompt: "continue", + }); + const result = shouldPreemptivelyCompactBeforePrompt({ + messages: [makeToolResultMessage(cjkText)], + systemPrompt: "sys", + prompt: "continue", + contextTokenBudget: 128_000, + reserveTokens: 20_000, + }); + + expect(toolResultTokens).toBeGreaterThanOrEqual(assistantTokens); + expect(toolResultTokens - assistantTokens).toBeLessThanOrEqual(5); + expect(result.estimatedPromptTokens).toBe(toolResultTokens); + expect(result.promptBudgetBeforeReserve).toBeGreaterThan(result.estimatedPromptTokens); + expect(result.route).toBe("fits"); + expect(result.shouldCompact).toBe(false); + expect(result.overflowTokens).toBe(0); + }); + + it("avoids false overflow when CJK is less than half of a tool result", () => { + const mixedContent = "中".repeat(40_000) + "a".repeat(60_000); + const result = shouldPreemptivelyCompactBeforePrompt({ + messages: [makeToolResultMessage(mixedContent)], + systemPrompt: "sys", + prompt: "continue", + contextTokenBudget: 100_000, + reserveTokens: 20_000, + }); + + expect(result.estimatedPromptTokens).toBeLessThan(result.promptBudgetBeforeReserve); + expect(result.route).toBe("fits"); + expect(result.shouldCompact).toBe(false); + }); + + it("keeps mixed-script estimates monotonic across the former CJK cutoff", () => { + const estimate = (cjkChars: number) => + estimateLlmBoundaryTokenPressure({ + messages: [makeToolResultMessage("中".repeat(cjkChars) + "a".repeat(10_000 - cjkChars))], + systemPrompt: "sys", + prompt: "continue", + }); + + const belowCutoff = estimate(4_999); + const atCutoff = estimate(5_000); + const aboveCutoff = estimate(5_001); + + expect(atCutoff).toBeGreaterThanOrEqual(belowCutoff); + expect(aboveCutoff).toBeGreaterThanOrEqual(atCutoff); + expect(aboveCutoff - belowCutoff).toBeLessThanOrEqual(2); + }); + + it("keeps the conservative ratio for non-CJK tool results", () => { + const latinText = "alpha beta gamma delta epsilon ".repeat(1000); + const toolResultTokens = estimateLlmBoundaryTokenPressure({ + messages: [makeToolResultMessage(latinText)], + systemPrompt: "sys", + prompt: "continue", + }); + const assistantTokens = estimateLlmBoundaryTokenPressure({ + messages: [makeAssistantHistory(latinText)], + systemPrompt: "sys", + prompt: "continue", + }); + + expect(toolResultTokens).toBeGreaterThan(assistantTokens * 1.5); + expect(toolResultTokens).toBeLessThan(assistantTokens * 2.5); + }); + + it("applies the CJK-aware ratio to JSON tool-result payloads", () => { + const cjkPayload = { + summary: "中文内容".repeat(5_000), + note: "更多中文文本".repeat(2_000), + }; + const messages = [makeJsonToolResultMessage(cjkPayload)]; + + const estimatedPromptTokens = estimateLlmBoundaryTokenPressure({ + messages, + systemPrompt: "sys", + prompt: "continue", + }); + + expect(estimatedPromptTokens).toBeLessThan(90_000); + + const result = shouldPreemptivelyCompactBeforePrompt({ + messages, + systemPrompt: "sys", + prompt: "continue", + contextTokenBudget: 128_000, + reserveTokens: 20_000, + }); + + expect(result.route).toBe("fits"); + expect(result.shouldCompact).toBe(false); + expect(result.overflowTokens).toBe(0); + }); + + it("does not throw when tool-result content cannot be serialized", () => { + const circular: Record = { self: undefined }; + circular.self = circular; + const message = { + role: "toolResult", + toolCallId: "call_circular", + toolName: "bad_tool", + content: circular, + isError: false, + timestamp: timestamp++, + } as unknown as AgentMessage; + + const result = shouldPreemptivelyCompactBeforePrompt({ + messages: [message], + systemPrompt: "sys", + prompt: "continue", + contextTokenBudget: 128_000, + reserveTokens: 20_000, + }); + + expect(Number.isFinite(result.estimatedPromptTokens)).toBe(true); + }); }); diff --git a/src/agents/embedded-agent-runner/run/preemptive-compaction.ts b/src/agents/embedded-agent-runner/run/preemptive-compaction.ts index 8a811d1de1bb..ce7250efa0b7 100644 --- a/src/agents/embedded-agent-runner/run/preemptive-compaction.ts +++ b/src/agents/embedded-agent-runner/run/preemptive-compaction.ts @@ -112,18 +112,50 @@ function estimateContentBlockTokenPressure( return CONTENT_BLOCK_OVERHEAD_TOKENS + estimateJsonPayloadTokenPressure(block, charsPerToken); } +function estimateToolResultStringTokenPressure(text: string): number { + const conservativeToolResultEstimate = Math.ceil(text.length / TOOL_RESULT_CHARS_PER_TOKEN); + const cjkAwareEstimate = estimateStringTokenPressure(text); + return Math.max(conservativeToolResultEstimate, cjkAwareEstimate); +} + +function estimateToolResultJsonTokenPressure(value: unknown): number { + try { + const serialized = JSON.stringify(value); + return typeof serialized === "string" ? estimateToolResultStringTokenPressure(serialized) : 1; + } catch { + return 256; + } +} + +function estimateToolResultBlockTokenPressure(block: unknown): number { + if (typeof block === "string") { + return estimateToolResultStringTokenPressure(block); + } + if (!isRecord(block)) { + return estimateToolResultJsonTokenPressure(block); + } + + if (block.type === "text" && typeof block.text === "string") { + return CONTENT_BLOCK_OVERHEAD_TOKENS + estimateToolResultStringTokenPressure(block.text); + } + if (block.type === "thinking" && typeof block.thinking === "string") { + return CONTENT_BLOCK_OVERHEAD_TOKENS + estimateToolResultStringTokenPressure(block.thinking); + } + if (block.type === "image") { + return IMAGE_BLOCK_TOKENS; + } + return CONTENT_BLOCK_OVERHEAD_TOKENS + estimateToolResultJsonTokenPressure(block); +} + function estimateToolResultContentTokenPressure(content: unknown): number { if (typeof content === "string") { - return estimateStringTokenPressure(content, TOOL_RESULT_CHARS_PER_TOKEN); + return estimateToolResultStringTokenPressure(content); } if (Array.isArray(content)) { - return content.reduce( - (sum, block) => sum + estimateContentBlockTokenPressure(block, TOOL_RESULT_CHARS_PER_TOKEN), - 0, - ); + return content.reduce((sum, block) => sum + estimateToolResultBlockTokenPressure(block), 0); } if (content !== undefined) { - return estimateJsonPayloadTokenPressure(content, TOOL_RESULT_CHARS_PER_TOKEN); + return estimateToolResultJsonTokenPressure(content); } return 0; }