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

Commit 4c8db13

Browse files
roomoteroomote
committed
feat: add context limits for folder mentions to prevent overflow
This change adds limits to prevent context window explosion when using @folder mentions: - MAX_FOLDER_FILES_TO_READ (10): Limits the number of files read from a folder - MAX_FOLDER_CONTENT_SIZE (100KB): Limits total content size from folder mentions When limits are reached, users receive informative messages suggesting to use individual @file mentions instead. Fixes #11137
1 parent f97a5c2 commit 4c8db13

2 files changed

Lines changed: 310 additions & 1 deletion

File tree

‎src/core/mentions/__tests__/folder-limits.spec.ts‎

Lines changed: 251 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,251 @@
1+
// npx vitest src/core/mentions/__tests__/folder-limits.spec.ts
2+
3+
import * as path from "path"
4+
import { MAX_FOLDER_FILES_TO_READ, MAX_FOLDER_CONTENT_SIZE } from "../index"
5+
6+
// Mock vscode
7+
vi.mock("vscode", () => ({
8+
window: {
9+
showErrorMessage: vi.fn(),
10+
},
11+
languages: {
12+
getDiagnostics: vi.fn().mockReturnValue([]),
13+
},
14+
}))
15+
16+
// Mock i18n
17+
vi.mock("../../../i18n", () => ({
18+
t: vi.fn((key: string) => key),
19+
}))
20+
21+
// Mock isbinaryfile
22+
vi.mock("isbinaryfile", () => ({
23+
isBinaryFile: vi.fn().mockResolvedValue(false),
24+
}))
25+
26+
// Mock fs/promises
27+
const mockReaddir = vi.fn()
28+
const mockStat = vi.fn()
29+
const mockReadFile = vi.fn()
30+
31+
vi.mock("fs/promises", () => ({
32+
default: {
33+
readdir: (...args: any[]) => mockReaddir(...args),
34+
stat: (...args: any[]) => mockStat(...args),
35+
readFile: (...args: any[]) => mockReadFile(...args),
36+
access: vi.fn().mockResolvedValue(undefined),
37+
},
38+
readdir: (...args: any[]) => mockReaddir(...args),
39+
stat: (...args: any[]) => mockStat(...args),
40+
readFile: (...args: any[]) => mockReadFile(...args),
41+
access: vi.fn().mockResolvedValue(undefined),
42+
}))
43+
44+
// Mock extract-text module
45+
vi.mock("../../../integrations/misc/extract-text", () => ({
46+
extractTextFromFileWithMetadata: vi.fn().mockImplementation(async (filePath: string) => {
47+
// Return a predictable result for testing
48+
return {
49+
content: `1 | // Content of ${path.basename(filePath)}\n2 | const x = 1;`,
50+
totalLines: 2,
51+
returnedLines: 2,
52+
wasTruncated: false,
53+
}
54+
}),
55+
}))
56+
57+
describe("Folder Mention Content Limits", () => {
58+
beforeEach(() => {
59+
vi.clearAllMocks()
60+
})
61+
62+
describe("Constants", () => {
63+
it("should export MAX_FOLDER_FILES_TO_READ constant", () => {
64+
expect(MAX_FOLDER_FILES_TO_READ).toBeDefined()
65+
expect(typeof MAX_FOLDER_FILES_TO_READ).toBe("number")
66+
expect(MAX_FOLDER_FILES_TO_READ).toBeGreaterThan(0)
67+
})
68+
69+
it("should export MAX_FOLDER_CONTENT_SIZE constant", () => {
70+
expect(MAX_FOLDER_CONTENT_SIZE).toBeDefined()
71+
expect(typeof MAX_FOLDER_CONTENT_SIZE).toBe("number")
72+
expect(MAX_FOLDER_CONTENT_SIZE).toBeGreaterThan(0)
73+
})
74+
75+
it("should have reasonable default values", () => {
76+
// MAX_FOLDER_FILES_TO_READ should be reasonable (e.g., 10)
77+
expect(MAX_FOLDER_FILES_TO_READ).toBeLessThanOrEqual(50)
78+
expect(MAX_FOLDER_FILES_TO_READ).toBeGreaterThanOrEqual(5)
79+
80+
// MAX_FOLDER_CONTENT_SIZE should be reasonable (e.g., 100KB)
81+
expect(MAX_FOLDER_CONTENT_SIZE).toBeGreaterThanOrEqual(50_000)
82+
expect(MAX_FOLDER_CONTENT_SIZE).toBeLessThanOrEqual(500_000)
83+
})
84+
})
85+
86+
describe("parseMentions with folder limits", () => {
87+
// These tests import parseMentions dynamically to ensure mocks are applied
88+
it("should limit the number of files read from a folder", async () => {
89+
// Create many file entries to exceed the limit
90+
const numFiles = MAX_FOLDER_FILES_TO_READ + 5
91+
const entries = Array.from({ length: numFiles }, (_, i) => ({
92+
name: `file${i}.ts`,
93+
isFile: () => true,
94+
isDirectory: () => false,
95+
}))
96+
97+
mockStat.mockResolvedValue({
98+
isFile: () => false,
99+
isDirectory: () => true,
100+
})
101+
mockReaddir.mockResolvedValue(entries)
102+
103+
// Import the module after mocks are set up
104+
const { parseMentions } = await import("../index")
105+
const mockUrlContentFetcher = {
106+
launchBrowser: vi.fn(),
107+
urlToMarkdown: vi.fn(),
108+
closeBrowser: vi.fn(),
109+
} as any
110+
111+
const result = await parseMentions("Check @/test-folder/", "/workspace", mockUrlContentFetcher)
112+
113+
// Should have content blocks with folder content
114+
expect(result.contentBlocks.length).toBeGreaterThan(0)
115+
const folderBlock = result.contentBlocks.find((b) => b.type === "folder")
116+
expect(folderBlock).toBeDefined()
117+
118+
// Should contain truncation notice
119+
expect(folderBlock?.content).toContain("Content Truncated")
120+
expect(folderBlock?.content).toContain(`Only ${MAX_FOLDER_FILES_TO_READ} files were read`)
121+
})
122+
123+
it("should limit total content size", async () => {
124+
// Create a few files that together exceed the content size limit
125+
const numFiles = 3
126+
const entries = Array.from({ length: numFiles }, (_, i) => ({
127+
name: `file${i}.ts`,
128+
isFile: () => true,
129+
isDirectory: () => false,
130+
}))
131+
132+
mockStat.mockResolvedValue({
133+
isFile: () => false,
134+
isDirectory: () => true,
135+
})
136+
mockReaddir.mockResolvedValue(entries)
137+
138+
// Mock extractTextFromFileWithMetadata to return large content
139+
const { extractTextFromFileWithMetadata } = await import("../../../integrations/misc/extract-text")
140+
vi.mocked(extractTextFromFileWithMetadata).mockImplementation(async () => {
141+
// Return content that's about half of MAX_FOLDER_CONTENT_SIZE
142+
const largeContent = "x".repeat(Math.ceil(MAX_FOLDER_CONTENT_SIZE / 2))
143+
return {
144+
content: largeContent,
145+
totalLines: 1000,
146+
returnedLines: 1000,
147+
wasTruncated: false,
148+
}
149+
})
150+
151+
// Import the module after mocks are set up
152+
const { parseMentions } = await import("../index")
153+
const mockUrlContentFetcher = {
154+
launchBrowser: vi.fn(),
155+
urlToMarkdown: vi.fn(),
156+
closeBrowser: vi.fn(),
157+
} as any
158+
159+
const result = await parseMentions("Check @/test-folder/", "/workspace", mockUrlContentFetcher)
160+
161+
// Should have content blocks with folder content
162+
expect(result.contentBlocks.length).toBeGreaterThan(0)
163+
const folderBlock = result.contentBlocks.find((b) => b.type === "folder")
164+
expect(folderBlock).toBeDefined()
165+
166+
// Should contain truncation notice due to size
167+
expect(folderBlock?.content).toContain("Content Truncated")
168+
expect(folderBlock?.content).toContain("KB to prevent context window overflow")
169+
})
170+
171+
it("should not add truncation notice when within limits", async () => {
172+
// Create a few small files within limits
173+
const numFiles = 3
174+
const entries = Array.from({ length: numFiles }, (_, i) => ({
175+
name: `file${i}.ts`,
176+
isFile: () => true,
177+
isDirectory: () => false,
178+
}))
179+
180+
mockStat.mockResolvedValue({
181+
isFile: () => false,
182+
isDirectory: () => true,
183+
})
184+
mockReaddir.mockResolvedValue(entries)
185+
186+
// Mock extractTextFromFileWithMetadata to return small content
187+
const { extractTextFromFileWithMetadata } = await import("../../../integrations/misc/extract-text")
188+
vi.mocked(extractTextFromFileWithMetadata).mockImplementation(async (filePath: string) => {
189+
return {
190+
content: `1 | // Small file ${path.basename(filePath)}`,
191+
totalLines: 1,
192+
returnedLines: 1,
193+
wasTruncated: false,
194+
}
195+
})
196+
197+
// Import the module after mocks are set up
198+
const { parseMentions } = await import("../index")
199+
const mockUrlContentFetcher = {
200+
launchBrowser: vi.fn(),
201+
urlToMarkdown: vi.fn(),
202+
closeBrowser: vi.fn(),
203+
} as any
204+
205+
const result = await parseMentions("Check @/small-folder/", "/workspace", mockUrlContentFetcher)
206+
207+
// Should have content blocks with folder content
208+
expect(result.contentBlocks.length).toBeGreaterThan(0)
209+
const folderBlock = result.contentBlocks.find((b) => b.type === "folder")
210+
expect(folderBlock).toBeDefined()
211+
212+
// Should NOT contain truncation notice
213+
expect(folderBlock?.content).not.toContain("Content Truncated")
214+
})
215+
216+
it("should still show folder listing even when files are skipped", async () => {
217+
// Create many file entries
218+
const numFiles = MAX_FOLDER_FILES_TO_READ + 10
219+
const entries = Array.from({ length: numFiles }, (_, i) => ({
220+
name: `file${i}.ts`,
221+
isFile: () => true,
222+
isDirectory: () => false,
223+
}))
224+
225+
mockStat.mockResolvedValue({
226+
isFile: () => false,
227+
isDirectory: () => true,
228+
})
229+
mockReaddir.mockResolvedValue(entries)
230+
231+
// Import the module after mocks are set up
232+
const { parseMentions } = await import("../index")
233+
const mockUrlContentFetcher = {
234+
launchBrowser: vi.fn(),
235+
urlToMarkdown: vi.fn(),
236+
closeBrowser: vi.fn(),
237+
} as any
238+
239+
const result = await parseMentions("Check @/large-folder/", "/workspace", mockUrlContentFetcher)
240+
241+
const folderBlock = result.contentBlocks.find((b) => b.type === "folder")
242+
expect(folderBlock).toBeDefined()
243+
244+
// Should contain the folder listing with all files listed
245+
// (the listing shows all files, only content reading is limited)
246+
expect(folderBlock?.content).toContain("Folder listing:")
247+
expect(folderBlock?.content).toContain("file0.ts")
248+
expect(folderBlock?.content).toContain(`file${numFiles - 1}.ts`)
249+
})
250+
})
251+
})

‎src/core/mentions/index.ts‎

Lines changed: 59 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,18 @@ import { extractTextFromFileWithMetadata, type ExtractTextResult } from "../../i
1313
import { diagnosticsToProblemsString } from "../../integrations/diagnostics"
1414
import { DEFAULT_LINE_LIMIT } from "../prompts/tools/native-tools/read_file"
1515

16+
/**
17+
* Maximum number of files to read from a folder mention.
18+
* This prevents context window explosion when mentioning large directories.
19+
*/
20+
export const MAX_FOLDER_FILES_TO_READ = 10
21+
22+
/**
23+
* Maximum total content size (in characters) to read from a folder mention.
24+
* This is approximately 100KB which should be safe for most context windows.
25+
*/
26+
export const MAX_FOLDER_CONTENT_SIZE = 100_000
27+
1628
import { UrlContentFetcher } from "../../services/browser/UrlContentFetcher"
1729

1830
import { FileContextTracker } from "../context-tracking/FileContextTracker"
@@ -390,6 +402,12 @@ async function getFileOrFolderContentWithMetadata(
390402
const fileReadResults: string[] = []
391403
const LOCK_SYMBOL = "🔒"
392404

405+
// Track limits to prevent context window explosion
406+
let filesRead = 0
407+
let totalContentSize = 0
408+
let limitReached: "files" | "size" | null = null
409+
let skippedFilesCount = 0
410+
393411
for (let index = 0; index < entries.length; index++) {
394412
const entry = entries[index]
395413
const isLast = index === entries.length - 1
@@ -410,13 +428,44 @@ async function getFileOrFolderContentWithMetadata(
410428
if (entry.isFile()) {
411429
folderListing += `${linePrefix}${displayName}\n`
412430
if (!isIgnored) {
431+
// Check if we've hit the file limit
432+
if (filesRead >= MAX_FOLDER_FILES_TO_READ) {
433+
if (!limitReached) {
434+
limitReached = "files"
435+
}
436+
skippedFilesCount++
437+
continue
438+
}
439+
440+
// Check if we've hit the content size limit
441+
if (totalContentSize >= MAX_FOLDER_CONTENT_SIZE) {
442+
if (!limitReached) {
443+
limitReached = "size"
444+
}
445+
skippedFilesCount++
446+
continue
447+
}
448+
413449
const filePath = path.join(mentionPath, entry.name)
414450
const absoluteFilePath = path.resolve(absPath, entry.name)
415451
try {
416452
const isBinary = await isBinaryFile(absoluteFilePath).catch(() => false)
417453
if (!isBinary) {
418454
const result = await extractTextFromFileWithMetadata(absoluteFilePath)
419-
fileReadResults.push(formatFileReadResult(filePath.toPosix(), result))
455+
const fileContent = formatFileReadResult(filePath.toPosix(), result)
456+
457+
// Check if adding this file would exceed the size limit
458+
if (totalContentSize + fileContent.length > MAX_FOLDER_CONTENT_SIZE) {
459+
if (!limitReached) {
460+
limitReached = "size"
461+
}
462+
skippedFilesCount++
463+
continue
464+
}
465+
466+
fileReadResults.push(fileContent)
467+
filesRead++
468+
totalContentSize += fileContent.length
420469
}
421470
} catch (error) {
422471
// Skip files that can't be read
@@ -435,6 +484,15 @@ async function getFileOrFolderContentWithMetadata(
435484
content += `\n\n--- File Contents ---\n\n${fileReadResults.join("\n\n")}`
436485
}
437486

487+
// Add truncation notice if limits were hit
488+
if (limitReached) {
489+
const limitMessage =
490+
limitReached === "files"
491+
? `\n\n--- Content Truncated ---\nNote: Only ${MAX_FOLDER_FILES_TO_READ} files were read to prevent context window overflow. ${skippedFilesCount} additional file(s) were skipped.\nTo read specific files, use individual @file mentions instead of @folder.`
492+
: `\n\n--- Content Truncated ---\nNote: Content was limited to approximately ${Math.round(MAX_FOLDER_CONTENT_SIZE / 1000)}KB to prevent context window overflow. ${skippedFilesCount} additional file(s) were skipped.\nTo read specific files, use individual @file mentions instead of @folder.`
493+
content += limitMessage
494+
}
495+
438496
return {
439497
type: "folder",
440498
path: mentionPath,

0 commit comments

Comments
 (0)