Skip to content

Commit e6de349

Browse files
committed
fix(context): truncat type definition to match max read line
1 parent 97f9686 commit e6de349

File tree

5 files changed

+540
-2
lines changed

5 files changed

+540
-2
lines changed
Lines changed: 351 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,351 @@
1+
// npx vitest src/core/tools/__tests__/listCodeDefinitionNamesTool.spec.ts
2+
3+
import { describe, it, expect, vi, beforeEach } from "vitest"
4+
import { listCodeDefinitionNamesTool } from "../listCodeDefinitionNamesTool"
5+
import { Task } from "../../task/Task"
6+
import { ToolUse } from "../../../shared/tools"
7+
import * as treeSitter from "../../../services/tree-sitter"
8+
import fs from "fs/promises"
9+
10+
// Mock the tree-sitter service
11+
vi.mock("../../../services/tree-sitter", () => ({
12+
parseSourceCodeDefinitionsForFile: vi.fn(),
13+
parseSourceCodeForDefinitionsTopLevel: vi.fn(),
14+
}))
15+
16+
// Mock fs module
17+
vi.mock("fs/promises", () => ({
18+
default: {
19+
stat: vi.fn(),
20+
},
21+
}))
22+
23+
describe("listCodeDefinitionNamesTool", () => {
24+
let mockTask: Partial<Task>
25+
let mockAskApproval: any
26+
let mockHandleError: any
27+
let mockPushToolResult: any
28+
let mockRemoveClosingTag: any
29+
30+
beforeEach(() => {
31+
vi.clearAllMocks()
32+
33+
mockTask = {
34+
cwd: "/test/path",
35+
consecutiveMistakeCount: 0,
36+
recordToolError: vi.fn(),
37+
sayAndCreateMissingParamError: vi.fn(),
38+
ask: vi.fn(),
39+
fileContextTracker: {
40+
trackFileContext: vi.fn(),
41+
},
42+
providerRef: {
43+
deref: vi.fn(() => ({
44+
getState: vi.fn(async () => ({ maxReadFileLine: -1 })),
45+
})),
46+
},
47+
rooIgnoreController: undefined,
48+
} as any
49+
50+
mockAskApproval = vi.fn(async () => true)
51+
mockHandleError = vi.fn()
52+
mockPushToolResult = vi.fn()
53+
mockRemoveClosingTag = vi.fn((tag: string, value: string) => value)
54+
})
55+
56+
describe("truncateDefinitionsToLineLimit", () => {
57+
it("should not truncate when maxReadFileLine is -1 (no limit)", async () => {
58+
const mockDefinitions = `# test.ts
59+
10--20 | function foo() {
60+
30--40 | function bar() {
61+
50--60 | function baz() {`
62+
63+
vi.mocked(treeSitter.parseSourceCodeDefinitionsForFile).mockResolvedValue(mockDefinitions)
64+
65+
vi.mocked(fs.stat).mockResolvedValue({
66+
isFile: () => true,
67+
isDirectory: () => false,
68+
} as any)
69+
70+
mockTask.providerRef = {
71+
deref: vi.fn(() => ({
72+
getState: vi.fn(async () => ({ maxReadFileLine: -1 })),
73+
})),
74+
} as any
75+
76+
const block: ToolUse = {
77+
type: "tool_use",
78+
name: "list_code_definition_names",
79+
params: { path: "test.ts" },
80+
partial: false,
81+
}
82+
83+
await listCodeDefinitionNamesTool(
84+
mockTask as Task,
85+
block,
86+
mockAskApproval,
87+
mockHandleError,
88+
mockPushToolResult,
89+
mockRemoveClosingTag,
90+
)
91+
92+
expect(mockPushToolResult).toHaveBeenCalledWith(mockDefinitions)
93+
})
94+
95+
it("should not truncate when maxReadFileLine is 0 (definitions only mode)", async () => {
96+
const mockDefinitions = `# test.ts
97+
10--20 | function foo() {
98+
30--40 | function bar() {
99+
50--60 | function baz() {`
100+
101+
vi.mocked(treeSitter.parseSourceCodeDefinitionsForFile).mockResolvedValue(mockDefinitions)
102+
103+
vi.mocked(fs.stat).mockResolvedValue({
104+
isFile: () => true,
105+
isDirectory: () => false,
106+
} as any)
107+
108+
mockTask.providerRef = {
109+
deref: vi.fn(() => ({
110+
getState: vi.fn(async () => ({ maxReadFileLine: 0 })),
111+
})),
112+
} as any
113+
114+
const block: ToolUse = {
115+
type: "tool_use",
116+
name: "list_code_definition_names",
117+
params: { path: "test.ts" },
118+
partial: false,
119+
}
120+
121+
await listCodeDefinitionNamesTool(
122+
mockTask as Task,
123+
block,
124+
mockAskApproval,
125+
mockHandleError,
126+
mockPushToolResult,
127+
mockRemoveClosingTag,
128+
)
129+
130+
expect(mockPushToolResult).toHaveBeenCalledWith(mockDefinitions)
131+
})
132+
133+
it("should truncate definitions when maxReadFileLine is set", async () => {
134+
const mockDefinitions = `# test.ts
135+
10--20 | function foo() {
136+
30--40 | function bar() {
137+
50--60 | function baz() {`
138+
139+
vi.mocked(treeSitter.parseSourceCodeDefinitionsForFile).mockResolvedValue(mockDefinitions)
140+
141+
vi.mocked(fs.stat).mockResolvedValue({
142+
isFile: () => true,
143+
isDirectory: () => false,
144+
} as any)
145+
146+
mockTask.providerRef = {
147+
deref: vi.fn(() => ({
148+
getState: vi.fn(async () => ({ maxReadFileLine: 25 })),
149+
})),
150+
} as any
151+
152+
const block: ToolUse = {
153+
type: "tool_use",
154+
name: "list_code_definition_names",
155+
params: { path: "test.ts" },
156+
partial: false,
157+
}
158+
159+
await listCodeDefinitionNamesTool(
160+
mockTask as Task,
161+
block,
162+
mockAskApproval,
163+
mockHandleError,
164+
mockPushToolResult,
165+
mockRemoveClosingTag,
166+
)
167+
168+
// Should only include definitions starting at or before line 25
169+
const expectedResult = `# test.ts
170+
10--20 | function foo() {`
171+
172+
expect(mockPushToolResult).toHaveBeenCalledWith(expectedResult)
173+
})
174+
175+
it("should include definitions that start within limit even if they end beyond it", async () => {
176+
const mockDefinitions = `# test.ts
177+
10--50 | function foo() {
178+
60--80 | function bar() {`
179+
180+
vi.mocked(treeSitter.parseSourceCodeDefinitionsForFile).mockResolvedValue(mockDefinitions)
181+
182+
vi.mocked(fs.stat).mockResolvedValue({
183+
isFile: () => true,
184+
isDirectory: () => false,
185+
} as any)
186+
187+
mockTask.providerRef = {
188+
deref: vi.fn(() => ({
189+
getState: vi.fn(async () => ({ maxReadFileLine: 30 })),
190+
})),
191+
} as any
192+
193+
const block: ToolUse = {
194+
type: "tool_use",
195+
name: "list_code_definition_names",
196+
params: { path: "test.ts" },
197+
partial: false,
198+
}
199+
200+
await listCodeDefinitionNamesTool(
201+
mockTask as Task,
202+
block,
203+
mockAskApproval,
204+
mockHandleError,
205+
mockPushToolResult,
206+
mockRemoveClosingTag,
207+
)
208+
209+
// Should include foo (starts at 10) but not bar (starts at 60)
210+
const expectedResult = `# test.ts
211+
10--50 | function foo() {`
212+
213+
expect(mockPushToolResult).toHaveBeenCalledWith(expectedResult)
214+
})
215+
216+
it("should handle single-line definitions", async () => {
217+
const mockDefinitions = `# test.ts
218+
10 | const foo = 1
219+
20 | const bar = 2
220+
30 | const baz = 3`
221+
222+
vi.mocked(treeSitter.parseSourceCodeDefinitionsForFile).mockResolvedValue(mockDefinitions)
223+
224+
vi.mocked(fs.stat).mockResolvedValue({
225+
isFile: () => true,
226+
isDirectory: () => false,
227+
} as any)
228+
229+
mockTask.providerRef = {
230+
deref: vi.fn(() => ({
231+
getState: vi.fn(async () => ({ maxReadFileLine: 25 })),
232+
})),
233+
} as any
234+
235+
const block: ToolUse = {
236+
type: "tool_use",
237+
name: "list_code_definition_names",
238+
params: { path: "test.ts" },
239+
partial: false,
240+
}
241+
242+
await listCodeDefinitionNamesTool(
243+
mockTask as Task,
244+
block,
245+
mockAskApproval,
246+
mockHandleError,
247+
mockPushToolResult,
248+
mockRemoveClosingTag,
249+
)
250+
251+
// Should include foo and bar but not baz
252+
const expectedResult = `# test.ts
253+
10 | const foo = 1
254+
20 | const bar = 2`
255+
256+
expect(mockPushToolResult).toHaveBeenCalledWith(expectedResult)
257+
})
258+
259+
it("should preserve header line when truncating", async () => {
260+
const mockDefinitions = `# test.ts
261+
100--200 | function foo() {`
262+
263+
vi.mocked(treeSitter.parseSourceCodeDefinitionsForFile).mockResolvedValue(mockDefinitions)
264+
265+
vi.mocked(fs.stat).mockResolvedValue({
266+
isFile: () => true,
267+
isDirectory: () => false,
268+
} as any)
269+
270+
mockTask.providerRef = {
271+
deref: vi.fn(() => ({
272+
getState: vi.fn(async () => ({ maxReadFileLine: 50 })),
273+
})),
274+
} as any
275+
276+
const block: ToolUse = {
277+
type: "tool_use",
278+
name: "list_code_definition_names",
279+
params: { path: "test.ts" },
280+
partial: false,
281+
}
282+
283+
await listCodeDefinitionNamesTool(
284+
mockTask as Task,
285+
block,
286+
mockAskApproval,
287+
mockHandleError,
288+
mockPushToolResult,
289+
mockRemoveClosingTag,
290+
)
291+
292+
// Should keep header but exclude all definitions beyond line 50
293+
const expectedResult = `# test.ts`
294+
295+
expect(mockPushToolResult).toHaveBeenCalledWith(expectedResult)
296+
})
297+
})
298+
299+
it("should handle missing path parameter", async () => {
300+
const block: ToolUse = {
301+
type: "tool_use",
302+
name: "list_code_definition_names",
303+
params: {},
304+
partial: false,
305+
}
306+
307+
mockTask.sayAndCreateMissingParamError = vi.fn(async () => "Missing parameter: path")
308+
309+
await listCodeDefinitionNamesTool(
310+
mockTask as Task,
311+
block,
312+
mockAskApproval,
313+
mockHandleError,
314+
mockPushToolResult,
315+
mockRemoveClosingTag,
316+
)
317+
318+
expect(mockTask.consecutiveMistakeCount).toBe(1)
319+
expect(mockTask.recordToolError).toHaveBeenCalledWith("list_code_definition_names")
320+
expect(mockPushToolResult).toHaveBeenCalledWith("Missing parameter: path")
321+
})
322+
323+
it("should handle directory path", async () => {
324+
const mockDefinitions = "# Directory definitions"
325+
326+
vi.mocked(treeSitter.parseSourceCodeForDefinitionsTopLevel).mockResolvedValue(mockDefinitions)
327+
328+
vi.mocked(fs.stat).mockResolvedValue({
329+
isFile: () => false,
330+
isDirectory: () => true,
331+
} as any)
332+
333+
const block: ToolUse = {
334+
type: "tool_use",
335+
name: "list_code_definition_names",
336+
params: { path: "src" },
337+
partial: false,
338+
}
339+
340+
await listCodeDefinitionNamesTool(
341+
mockTask as Task,
342+
block,
343+
mockAskApproval,
344+
mockHandleError,
345+
mockPushToolResult,
346+
mockRemoveClosingTag,
347+
)
348+
349+
expect(mockPushToolResult).toHaveBeenCalledWith(mockDefinitions)
350+
})
351+
})

0 commit comments

Comments
 (0)