Skip to content
This repository was archived by the owner on May 15, 2026. It is now read-only.

Commit 21cda01

Browse files
committed
feat(condense): add smart code folding with tree-sitter signatures
At context condensation time, use tree-sitter to generate folded code signatures (function definitions, class declarations) for files read during the conversation. Each file is included as its own <system-reminder> block in the condensed summary, preserving structural awareness without consuming excessive tokens. - Add getFilesReadByRoo() method to FileContextTracker - Create generateFoldedFileContext() using tree-sitter parsing - Update summarizeConversation() to accept array of file sections - Each file gets its own content block in the summary message - Add comprehensive test coverage (12 tests)
1 parent c7910a9 commit 21cda01

5 files changed

Lines changed: 565 additions & 2 deletions

File tree

Lines changed: 335 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,335 @@
1+
// npx vitest src/core/condense/__tests__/foldedFileContext.spec.ts
2+
3+
import * as path from "path"
4+
import { Anthropic } from "@anthropic-ai/sdk"
5+
import type { ModelInfo } from "@roo-code/types"
6+
import { TelemetryService } from "@roo-code/telemetry"
7+
import { BaseProvider } from "../../../api/providers/base-provider"
8+
9+
// Mock the tree-sitter module
10+
vi.mock("../../../services/tree-sitter", () => ({
11+
parseSourceCodeDefinitionsForFile: vi.fn(),
12+
}))
13+
14+
import { generateFoldedFileContext } from "../foldedFileContext"
15+
import { parseSourceCodeDefinitionsForFile } from "../../../services/tree-sitter"
16+
17+
const mockedParseSourceCodeDefinitions = vi.mocked(parseSourceCodeDefinitionsForFile)
18+
19+
describe("foldedFileContext", () => {
20+
beforeEach(() => {
21+
vi.clearAllMocks()
22+
})
23+
24+
describe("generateFoldedFileContext", () => {
25+
it("should return empty content for empty file list", async () => {
26+
const result = await generateFoldedFileContext([], { cwd: "/test" })
27+
28+
expect(result.content).toBe("")
29+
expect(result.sections).toEqual([])
30+
expect(result.filesProcessed).toBe(0)
31+
expect(result.filesSkipped).toBe(0)
32+
expect(result.characterCount).toBe(0)
33+
})
34+
35+
it("should generate folded context for a TypeScript file with its own system-reminder block", async () => {
36+
const mockDefinitions = `1--5 | export interface User
37+
7--12 | export function createUser(name: string): User
38+
14--28 | export class UserService`
39+
40+
mockedParseSourceCodeDefinitions.mockResolvedValue(mockDefinitions)
41+
42+
const result = await generateFoldedFileContext(["/test/user.ts"], { cwd: "/test" })
43+
44+
// Each file should be wrapped in its own <system-reminder> block
45+
expect(result.content).toContain("<system-reminder>")
46+
expect(result.content).toContain("</system-reminder>")
47+
expect(result.content).toContain("## File Context: /test/user.ts")
48+
expect(result.content).toContain("interface User")
49+
expect(result.content).toContain("function createUser")
50+
expect(result.content).toContain("class UserService")
51+
expect(result.filesProcessed).toBe(1)
52+
expect(result.filesSkipped).toBe(0)
53+
})
54+
55+
it("should generate folded context for a JavaScript file with its own system-reminder block", async () => {
56+
const mockDefinitions = `1--3 | function greet(name)
57+
5--15 | class Calculator`
58+
59+
mockedParseSourceCodeDefinitions.mockResolvedValue(mockDefinitions)
60+
61+
const result = await generateFoldedFileContext(["/test/utils.js"], { cwd: "/test" })
62+
63+
expect(result.content).toContain("<system-reminder>")
64+
expect(result.content).toContain("## File Context: /test/utils.js")
65+
expect(result.content).toContain("function greet")
66+
expect(result.content).toContain("class Calculator")
67+
expect(result.filesProcessed).toBe(1)
68+
})
69+
70+
it("should skip files when parseSourceCodeDefinitions returns undefined", async () => {
71+
// First file succeeds, second returns undefined
72+
mockedParseSourceCodeDefinitions
73+
.mockResolvedValueOnce("1--3 | export const x = 1")
74+
.mockResolvedValueOnce(undefined)
75+
76+
const result = await generateFoldedFileContext(["/test/existing.ts", "/test/unsupported.txt"], {
77+
cwd: "/test",
78+
})
79+
80+
expect(result.filesProcessed).toBe(1)
81+
expect(result.filesSkipped).toBe(1)
82+
})
83+
84+
it("should skip files when parseSourceCodeDefinitions throws an error", async () => {
85+
mockedParseSourceCodeDefinitions
86+
.mockResolvedValueOnce("1--3 | export const x = 1")
87+
.mockRejectedValueOnce(new Error("File not found"))
88+
89+
const result = await generateFoldedFileContext(["/test/existing.ts", "/test/non-existent.ts"], {
90+
cwd: "/test",
91+
})
92+
93+
expect(result.filesProcessed).toBe(1)
94+
expect(result.filesSkipped).toBe(1)
95+
})
96+
97+
it("should respect character budget limit", async () => {
98+
// Create multiple files that would exceed a small budget
99+
const longDefinitions = `1--3 | export function longFunctionName1()
100+
5--7 | export function longFunctionName2()
101+
9--11 | export function longFunctionName3()`
102+
103+
mockedParseSourceCodeDefinitions.mockResolvedValue(longDefinitions)
104+
105+
const result = await generateFoldedFileContext(["/test/file1.ts", "/test/file2.ts", "/test/file3.ts"], {
106+
cwd: "/test",
107+
maxCharacters: 200, // Small budget
108+
})
109+
110+
expect(result.characterCount).toBeLessThanOrEqual(200)
111+
// Some files should be skipped due to budget limit
112+
expect(result.filesSkipped).toBeGreaterThan(0)
113+
})
114+
115+
it("should handle Python files with its own system-reminder block", async () => {
116+
const mockDefinitions = `1--2 | def greet(name)
117+
4--12 | class Person`
118+
119+
mockedParseSourceCodeDefinitions.mockResolvedValue(mockDefinitions)
120+
121+
const result = await generateFoldedFileContext(["/test/person.py"], { cwd: "/test" })
122+
123+
expect(result.content).toContain("<system-reminder>")
124+
expect(result.content).toContain("## File Context: /test/person.py")
125+
expect(result.content).toContain("def greet")
126+
expect(result.content).toContain("class Person")
127+
expect(result.filesProcessed).toBe(1)
128+
})
129+
130+
it("should include file path in the File Context header", async () => {
131+
mockedParseSourceCodeDefinitions.mockResolvedValue("1--3 | export function helper()")
132+
133+
const result = await generateFoldedFileContext(["/test/src/utils/helpers.ts"], { cwd: "/test" })
134+
135+
// The path should appear in the File Context header
136+
expect(result.content).toContain("## File Context: /test/src/utils/helpers.ts")
137+
})
138+
139+
it("should generate separate system-reminder blocks for multiple files", async () => {
140+
mockedParseSourceCodeDefinitions
141+
.mockResolvedValueOnce("1--3 | export async function fetchData(url: string): Promise<any>")
142+
.mockResolvedValueOnce("1--4 | export interface DataModel")
143+
144+
const result = await generateFoldedFileContext(["/test/api.ts", "/test/models.ts"], { cwd: "/test" })
145+
146+
// Each file should have its own <system-reminder> block
147+
const systemReminderMatches = result.content.match(/<system-reminder>/g)
148+
expect(systemReminderMatches).toHaveLength(2)
149+
150+
// sections array should have separate entries for each file
151+
expect(result.sections).toHaveLength(2)
152+
expect(result.sections[0]).toContain("## File Context: /test/api.ts")
153+
expect(result.sections[1]).toContain("## File Context: /test/models.ts")
154+
155+
expect(result.content).toContain("## File Context: /test/api.ts")
156+
expect(result.content).toContain("## File Context: /test/models.ts")
157+
expect(result.content).toContain("fetchData")
158+
expect(result.content).toContain("interface DataModel")
159+
expect(result.filesProcessed).toBe(2)
160+
})
161+
162+
it("should truncate content when approaching character limit", async () => {
163+
// Create a definition that would fit but is close to the limit
164+
const longDefinitions = "1--3 | " + "x".repeat(300)
165+
166+
mockedParseSourceCodeDefinitions.mockResolvedValue(longDefinitions)
167+
168+
const result = await generateFoldedFileContext(["/test/file1.ts", "/test/file2.ts"], {
169+
cwd: "/test",
170+
maxCharacters: 350, // First file will fit, second will be truncated
171+
})
172+
173+
// Content should include truncation marker if truncation happened
174+
expect(result.filesProcessed + result.filesSkipped).toBe(2)
175+
})
176+
})
177+
178+
describe("summarizeConversation with foldedFileContext", () => {
179+
beforeEach(() => {
180+
if (!TelemetryService.hasInstance()) {
181+
TelemetryService.createInstance([])
182+
}
183+
})
184+
185+
// Mock API handler for testing
186+
class MockApiHandler extends BaseProvider {
187+
createMessage(): any {
188+
const mockStream = {
189+
async *[Symbol.asyncIterator]() {
190+
yield { type: "text", text: "Mock summary of the conversation" }
191+
yield { type: "usage", inputTokens: 100, outputTokens: 50, totalCost: 0.01 }
192+
},
193+
}
194+
return mockStream
195+
}
196+
197+
getModel(): { id: string; info: ModelInfo } {
198+
return {
199+
id: "test-model",
200+
info: {
201+
contextWindow: 100000,
202+
maxTokens: 50000,
203+
supportsPromptCache: true,
204+
supportsImages: false,
205+
inputPrice: 0,
206+
outputPrice: 0,
207+
description: "Test model",
208+
},
209+
}
210+
}
211+
212+
override async countTokens(content: Array<Anthropic.Messages.ContentBlockParam>): Promise<number> {
213+
let tokens = 0
214+
for (const block of content) {
215+
if (block.type === "text") {
216+
tokens += Math.ceil(block.text.length / 4)
217+
}
218+
}
219+
return tokens
220+
}
221+
}
222+
223+
it("should include folded file context with each file as a separate content block", async () => {
224+
const { summarizeConversation } = await import("../index")
225+
226+
const mockApiHandler = new MockApiHandler()
227+
const taskId = "test-task-id"
228+
229+
const messages: any[] = [
230+
{ role: "user", content: "First message" },
231+
{ role: "assistant", content: "Second message" },
232+
{ role: "user", content: "Third message" },
233+
{ role: "assistant", content: "Fourth message" },
234+
{ role: "user", content: "Fifth message" },
235+
{ role: "assistant", content: "Sixth message" },
236+
{ role: "user", content: "Seventh message" },
237+
]
238+
239+
// Folded file context sections - each file in its own <system-reminder> block
240+
const foldedFileContextSections = [
241+
`<system-reminder>
242+
## File Context: src/user.ts
243+
1--5 | export interface User
244+
7--12 | export function createUser(name: string): User
245+
14--28 | export class UserService
246+
</system-reminder>`,
247+
`<system-reminder>
248+
## File Context: src/api.ts
249+
1--3 | export async function fetchData(url: string): Promise<any>
250+
</system-reminder>`,
251+
]
252+
253+
const result = await summarizeConversation(
254+
messages,
255+
mockApiHandler,
256+
"System prompt",
257+
taskId,
258+
false,
259+
undefined, // customCondensingPrompt
260+
undefined, // metadata
261+
undefined, // environmentDetails
262+
foldedFileContextSections, // Array of sections, each file in its own block
263+
)
264+
265+
// Verify the summary was created
266+
expect(result.summary).toBeDefined()
267+
expect(result.messages.length).toBeGreaterThan(0)
268+
269+
// Find the summary message
270+
const summaryMessage = result.messages.find((msg: any) => msg.isSummary)
271+
expect(summaryMessage).toBeDefined()
272+
273+
// Each file should have its own content block
274+
const contentArray = summaryMessage!.content as any[]
275+
276+
// Find the content blocks containing file contexts
277+
const userFileBlock = contentArray.find(
278+
(block: any) => block.type === "text" && block.text?.includes("## File Context: src/user.ts"),
279+
)
280+
const apiFileBlock = contentArray.find(
281+
(block: any) => block.type === "text" && block.text?.includes("## File Context: src/api.ts"),
282+
)
283+
284+
expect(userFileBlock).toBeDefined()
285+
expect(apiFileBlock).toBeDefined()
286+
287+
// Each file block should have its own <system-reminder> tags
288+
expect(userFileBlock.text).toContain("<system-reminder>")
289+
expect(userFileBlock.text).toContain("export interface User")
290+
291+
expect(apiFileBlock.text).toContain("<system-reminder>")
292+
expect(apiFileBlock.text).toContain("fetchData")
293+
})
294+
295+
it("should not include file context section when foldedFileContextSections is empty", async () => {
296+
const { summarizeConversation } = await import("../index")
297+
298+
const mockApiHandler = new MockApiHandler()
299+
const taskId = "test-task-id-2"
300+
301+
const messages: any[] = [
302+
{ role: "user", content: "First message" },
303+
{ role: "assistant", content: "Second message" },
304+
{ role: "user", content: "Third message" },
305+
{ role: "assistant", content: "Fourth message" },
306+
{ role: "user", content: "Fifth message" },
307+
{ role: "assistant", content: "Sixth message" },
308+
{ role: "user", content: "Seventh message" },
309+
]
310+
311+
const result = await summarizeConversation(
312+
messages,
313+
mockApiHandler,
314+
"System prompt",
315+
taskId,
316+
false,
317+
undefined, // customCondensingPrompt
318+
undefined, // metadata
319+
undefined, // environmentDetails
320+
[], // Empty foldedFileContextSections array
321+
)
322+
323+
// Find the summary message
324+
const summaryMessage = result.messages.find((msg: any) => msg.isSummary)
325+
expect(summaryMessage).toBeDefined()
326+
327+
// The summary content should NOT contain any file context blocks
328+
const contentArray = summaryMessage!.content as any[]
329+
const fileContextBlock = contentArray.find(
330+
(block: any) => block.type === "text" && block.text?.includes("## File Context"),
331+
)
332+
expect(fileContextBlock).toBeUndefined()
333+
})
334+
})
335+
})

0 commit comments

Comments
 (0)