Skip to content

Commit ee6d6db

Browse files
committed
feat: add tool result compression to prevent context window exhaustion
- Add compressToolResult utility with intelligent truncation - Integrate compression into pushToolResult function - Scale compression limits based on model context window - Add comprehensive tests for compression functionality - Fixes #6463: Tool results that exceed context limits now get compressed
1 parent b7410dc commit ee6d6db

File tree

3 files changed

+304
-2
lines changed

3 files changed

+304
-2
lines changed

src/core/assistant-message/presentAssistantMessage.ts

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import { TelemetryService } from "@roo-code/telemetry"
66

77
import { defaultModeSlug, getModeBySlug } from "../../shared/modes"
88
import type { ToolParamName, ToolResponse } from "../../shared/tools"
9+
import { compressToolResult, getCompressionLimitsForContextWindow } from "../tools/compressToolResult"
910

1011
import { fetchInstructionsTool } from "../tools/fetchInstructionsTool"
1112
import { listFilesTool } from "../tools/listFilesTool"
@@ -249,9 +250,34 @@ export async function presentAssistantMessage(cline: Task) {
249250
cline.userMessageContent.push({ type: "text", text: `${toolDescription()} Result:` })
250251

251252
if (typeof content === "string") {
252-
cline.userMessageContent.push({ type: "text", text: content || "(tool did not return anything)" })
253+
// Get compression limits based on the model's context window
254+
const modelInfo = cline.api.getModel().info
255+
const contextWindow = modelInfo.contextWindow
256+
const { characterLimit, lineLimit } = getCompressionLimitsForContextWindow(contextWindow)
257+
258+
// Compress the tool result if it's too large
259+
const processedContent = compressToolResult(
260+
content || "(tool did not return anything)",
261+
characterLimit,
262+
lineLimit,
263+
)
264+
cline.userMessageContent.push({ type: "text", text: processedContent })
253265
} else {
254-
cline.userMessageContent.push(...content)
266+
// For non-string content (arrays of blocks), we still need to process text blocks
267+
const processedContent = content.map((block) => {
268+
if (block.type === "text" && block.text) {
269+
const modelInfo = cline.api.getModel().info
270+
const contextWindow = modelInfo.contextWindow
271+
const { characterLimit, lineLimit } = getCompressionLimitsForContextWindow(contextWindow)
272+
273+
return {
274+
...block,
275+
text: compressToolResult(block.text, characterLimit, lineLimit),
276+
}
277+
}
278+
return block
279+
})
280+
cline.userMessageContent.push(...processedContent)
255281
}
256282

257283
// Once a tool result has been collected, ignore all other tool
Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,162 @@
1+
import { describe, it, expect } from "vitest"
2+
import {
3+
compressToolResult,
4+
shouldCompressToolResult,
5+
getCompressionLimitsForContextWindow,
6+
DEFAULT_TOOL_RESULT_CHARACTER_LIMIT,
7+
DEFAULT_TOOL_RESULT_LINE_LIMIT,
8+
} from "../compressToolResult"
9+
10+
describe("compressToolResult", () => {
11+
describe("compressToolResult function", () => {
12+
it("should return original result when within limits", () => {
13+
const shortResult = "This is a short result"
14+
const compressed = compressToolResult(shortResult)
15+
expect(compressed).toBe(shortResult)
16+
})
17+
18+
it("should return empty string for empty input", () => {
19+
expect(compressToolResult("")).toBe("")
20+
expect(compressToolResult(null as any)).toBe(null)
21+
expect(compressToolResult(undefined as any)).toBe(undefined)
22+
})
23+
24+
it("should compress result when exceeding character limit", () => {
25+
const longResult = "A".repeat(60000) // Exceeds default 50000 char limit
26+
const compressed = compressToolResult(longResult)
27+
28+
expect(compressed).not.toBe(longResult)
29+
expect(compressed.length).toBeLessThan(longResult.length)
30+
expect(compressed).toContain("[Tool result compressed:")
31+
expect(compressed).toContain("characters omitted")
32+
})
33+
34+
it("should compress result when exceeding line limit", () => {
35+
const manyLines = Array(1500).fill("line").join("\n") // Exceeds default 1000 line limit
36+
const compressed = compressToolResult(manyLines)
37+
38+
expect(compressed).not.toBe(manyLines)
39+
expect(compressed.split("\n").length).toBeLessThan(manyLines.split("\n").length)
40+
expect(compressed).toContain("[Tool result compressed:")
41+
expect(compressed).toContain("lines omitted")
42+
})
43+
44+
it("should use custom limits when provided", () => {
45+
const result = "A".repeat(200)
46+
const compressed = compressToolResult(result, 100, 10) // Custom limits
47+
48+
expect(compressed).not.toBe(result)
49+
expect(compressed).toContain("[Tool result compressed:")
50+
})
51+
52+
it("should preserve structure with compression note at beginning", () => {
53+
const longResult = "A".repeat(60000)
54+
const compressed = compressToolResult(longResult)
55+
56+
expect(compressed.startsWith("[Tool result compressed:")).toBe(true)
57+
expect(compressed).toContain("Original 60000 characters")
58+
})
59+
60+
it("should handle mixed character and line limits", () => {
61+
// Create content that exceeds both limits
62+
const longLines = Array(1500).fill("A".repeat(100)).join("\n")
63+
const compressed = compressToolResult(longLines)
64+
65+
expect(compressed).not.toBe(longLines)
66+
expect(compressed).toContain("[Tool result compressed:")
67+
})
68+
})
69+
70+
describe("shouldCompressToolResult function", () => {
71+
it("should return false for short results", () => {
72+
const shortResult = "Short result"
73+
expect(shouldCompressToolResult(shortResult)).toBe(false)
74+
})
75+
76+
it("should return true for results exceeding character limit", () => {
77+
const longResult = "A".repeat(60000)
78+
expect(shouldCompressToolResult(longResult)).toBe(true)
79+
})
80+
81+
it("should return true for results exceeding line limit", () => {
82+
const manyLines = Array(1500).fill("line").join("\n")
83+
expect(shouldCompressToolResult(manyLines)).toBe(true)
84+
})
85+
86+
it("should return false for empty results", () => {
87+
expect(shouldCompressToolResult("")).toBe(false)
88+
expect(shouldCompressToolResult(null as any)).toBe(false)
89+
expect(shouldCompressToolResult(undefined as any)).toBe(false)
90+
})
91+
92+
it("should respect custom limits", () => {
93+
const result = "A".repeat(200)
94+
expect(shouldCompressToolResult(result, 100, 10)).toBe(true)
95+
expect(shouldCompressToolResult(result, 300, 10)).toBe(false)
96+
})
97+
})
98+
99+
describe("getCompressionLimitsForContextWindow function", () => {
100+
it("should return appropriate limits for small context windows", () => {
101+
const limits = getCompressionLimitsForContextWindow(8000) // Small context window
102+
103+
expect(limits.characterLimit).toBeGreaterThanOrEqual(DEFAULT_TOOL_RESULT_CHARACTER_LIMIT)
104+
expect(limits.lineLimit).toBeGreaterThanOrEqual(DEFAULT_TOOL_RESULT_LINE_LIMIT)
105+
})
106+
107+
it("should return larger limits for large context windows", () => {
108+
const smallLimits = getCompressionLimitsForContextWindow(8000)
109+
const largeLimits = getCompressionLimitsForContextWindow(200000) // Large context window
110+
111+
expect(largeLimits.characterLimit).toBeGreaterThanOrEqual(smallLimits.characterLimit)
112+
expect(largeLimits.lineLimit).toBeGreaterThanOrEqual(smallLimits.lineLimit)
113+
})
114+
115+
it("should never return limits below defaults", () => {
116+
const limits = getCompressionLimitsForContextWindow(1000) // Very small context window
117+
118+
expect(limits.characterLimit).toBeGreaterThanOrEqual(DEFAULT_TOOL_RESULT_CHARACTER_LIMIT)
119+
expect(limits.lineLimit).toBeGreaterThanOrEqual(DEFAULT_TOOL_RESULT_LINE_LIMIT)
120+
})
121+
122+
it("should scale limits proportionally with context window", () => {
123+
const limits1 = getCompressionLimitsForContextWindow(50000)
124+
const limits2 = getCompressionLimitsForContextWindow(100000)
125+
126+
// Larger context window should allow larger tool results
127+
expect(limits2.characterLimit).toBeGreaterThanOrEqual(limits1.characterLimit)
128+
})
129+
130+
it("should cap limits at reasonable maximums", () => {
131+
const limits = getCompressionLimitsForContextWindow(1000000) // Extremely large context window
132+
133+
// Should not exceed 2x the default limits
134+
expect(limits.characterLimit).toBeLessThanOrEqual(DEFAULT_TOOL_RESULT_CHARACTER_LIMIT * 2)
135+
})
136+
})
137+
138+
describe("integration with truncateOutput", () => {
139+
it("should preserve beginning and end of content", () => {
140+
const longResult = "START" + "A".repeat(60000) + "END"
141+
const compressed = compressToolResult(longResult)
142+
143+
// Should contain compression note plus truncated content
144+
expect(compressed).toContain("[Tool result compressed:")
145+
// The truncated content should preserve structure from truncateOutput
146+
expect(compressed).toContain("START")
147+
expect(compressed).toContain("END")
148+
})
149+
150+
it("should handle line-based truncation", () => {
151+
const lines = Array(1500)
152+
.fill(0)
153+
.map((_, i) => `Line ${i + 1}`)
154+
const longResult = lines.join("\n")
155+
const compressed = compressToolResult(longResult)
156+
157+
expect(compressed).toContain("[Tool result compressed:")
158+
expect(compressed).toContain("Line 1") // Should preserve beginning
159+
expect(compressed).toContain("lines omitted") // Should indicate truncation
160+
})
161+
})
162+
})
Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
import { truncateOutput } from "../../integrations/misc/extract-text"
2+
import { DEFAULT_TERMINAL_OUTPUT_CHARACTER_LIMIT } from "@roo-code/types"
3+
4+
/**
5+
* Default character limit for tool results to prevent context window exhaustion
6+
* This is set to be conservative to ensure tool results don't consume too much context
7+
*/
8+
export const DEFAULT_TOOL_RESULT_CHARACTER_LIMIT = 50000
9+
10+
/**
11+
* Default line limit for tool results
12+
*/
13+
export const DEFAULT_TOOL_RESULT_LINE_LIMIT = 1000
14+
15+
/**
16+
* Compresses a tool result if it exceeds the specified limits.
17+
* Uses the same truncation logic as terminal output compression to maintain consistency.
18+
*
19+
* @param result The tool result string to potentially compress
20+
* @param characterLimit Maximum number of characters allowed (defaults to DEFAULT_TOOL_RESULT_CHARACTER_LIMIT)
21+
* @param lineLimit Maximum number of lines allowed (defaults to DEFAULT_TOOL_RESULT_LINE_LIMIT)
22+
* @returns The original result if within limits, or a compressed version with truncation indicators
23+
*/
24+
export function compressToolResult(
25+
result: string,
26+
characterLimit: number = DEFAULT_TOOL_RESULT_CHARACTER_LIMIT,
27+
lineLimit: number = DEFAULT_TOOL_RESULT_LINE_LIMIT,
28+
): string {
29+
// If result is empty or null, return as-is
30+
if (!result || result.length === 0) {
31+
return result
32+
}
33+
34+
// Check if compression is needed
35+
const needsCharacterCompression = characterLimit > 0 && result.length > characterLimit
36+
const needsLineCompression = lineLimit > 0 && result.split("\n").length > lineLimit
37+
38+
// If no compression is needed, return original result
39+
if (!needsCharacterCompression && !needsLineCompression) {
40+
return result
41+
}
42+
43+
// Use the existing truncateOutput function which handles both character and line limits
44+
// and provides intelligent truncation with context preservation
45+
const compressedResult = truncateOutput(
46+
result,
47+
lineLimit > 0 ? lineLimit : undefined,
48+
characterLimit > 0 ? characterLimit : undefined,
49+
)
50+
51+
// Add a note about compression if the result was actually truncated
52+
if (compressedResult !== result) {
53+
const originalLength = result.length
54+
const originalLines = result.split("\n").length
55+
const compressedLength = compressedResult.length
56+
const compressedLines = compressedResult.split("\n").length
57+
58+
// Add compression info at the beginning to make it clear to the model
59+
const compressionNote = `[Tool result compressed: Original ${originalLength} characters, ${originalLines} lines → Compressed to ${compressedLength} characters, ${compressedLines} lines to prevent context window exhaustion]\n\n`
60+
61+
return compressionNote + compressedResult
62+
}
63+
64+
return compressedResult
65+
}
66+
67+
/**
68+
* Determines if a tool result should be compressed based on its size and the model's context window.
69+
* This can be used to make compression decisions before actually compressing.
70+
*
71+
* @param result The tool result to check
72+
* @param characterLimit Character limit threshold
73+
* @param lineLimit Line limit threshold
74+
* @returns true if the result exceeds the limits and should be compressed
75+
*/
76+
export function shouldCompressToolResult(
77+
result: string,
78+
characterLimit: number = DEFAULT_TOOL_RESULT_CHARACTER_LIMIT,
79+
lineLimit: number = DEFAULT_TOOL_RESULT_LINE_LIMIT,
80+
): boolean {
81+
if (!result || result.length === 0) {
82+
return false
83+
}
84+
85+
const exceedsCharacterLimit = characterLimit > 0 && result.length > characterLimit
86+
const exceedsLineLimit = lineLimit > 0 && result.split("\n").length > lineLimit
87+
88+
return exceedsCharacterLimit || exceedsLineLimit
89+
}
90+
91+
/**
92+
* Gets appropriate compression limits based on the model's context window size.
93+
* Larger context windows can accommodate larger tool results.
94+
*
95+
* @param contextWindow The model's context window size in tokens
96+
* @returns Object with characterLimit and lineLimit appropriate for the context window
97+
*/
98+
export function getCompressionLimitsForContextWindow(contextWindow: number): {
99+
characterLimit: number
100+
lineLimit: number
101+
} {
102+
// Conservative approach: tool results should not exceed a small percentage of context window
103+
// Assuming roughly 4 characters per token on average
104+
const maxToolResultTokens = Math.floor(contextWindow * 0.1) // 10% of context window
105+
const characterLimit = Math.min(maxToolResultTokens * 4, DEFAULT_TOOL_RESULT_CHARACTER_LIMIT * 2) // Cap at 2x default
106+
107+
// Line limit scales with character limit
108+
const lineLimit = Math.floor(characterLimit / 50) // Assume ~50 chars per line on average
109+
110+
return {
111+
characterLimit: Math.max(characterLimit, DEFAULT_TOOL_RESULT_CHARACTER_LIMIT), // Never go below default
112+
lineLimit: Math.max(lineLimit, DEFAULT_TOOL_RESULT_LINE_LIMIT), // Never go below default
113+
}
114+
}

0 commit comments

Comments
 (0)