diff --git a/ui/desktop/src/components/MentionPopover.resourceLimit.test.tsx b/ui/desktop/src/components/MentionPopover.resourceLimit.test.tsx new file mode 100644 index 000000000..5a103b936 --- /dev/null +++ b/ui/desktop/src/components/MentionPopover.resourceLimit.test.tsx @@ -0,0 +1,74 @@ +import { act, render, screen, waitFor } from '@testing-library/react'; +import { describe, expect, it, vi } from 'vitest'; +import { IntlTestWrapper } from '../i18n/test-utils'; +import MentionPopover from './MentionPopover'; + +vi.mock('../acp/autocomplete', () => ({ + listAgentMentionItems: vi.fn().mockResolvedValue([]), + listSlashCommandItems: vi.fn().mockResolvedValue([]), +})); + +const props = { + onClose: vi.fn(), + onSelect: vi.fn(), + position: { x: 0, y: 0 }, + query: '', + isSlashCommand: false, + selectedIndex: -1, + onSelectedIndexChange: vi.fn(), + workingDir: '/workspace', +}; + +describe('MentionPopover scan resource limits', () => { + it('caps filesystem operations across a high-fanout tree', async () => { + const listFiles = vi.fn(async (directory: string) => { + const depth = directory.slice(props.workingDir.length).split('/').filter(Boolean).length; + return depth < 4 ? ['dir-a', 'dir-b', 'dir-c', 'dir-d'] : []; + }); + window.electron.listFiles = listFiles; + + render(, { wrapper: IntlTestWrapper }); + + await waitFor(() => expect(screen.queryByText('Scanning files...')).not.toBeInTheDocument()); + expect(listFiles.mock.calls.length).toBeLessThanOrEqual(100); + }); + + it('stops scheduling filesystem operations after the popover closes', async () => { + let resolveRoot!: (entries: string[]) => void; + const rootEntries = new Promise((resolve) => { + resolveRoot = resolve; + }); + const listFiles = vi.fn((directory: string) => + directory === props.workingDir ? rootEntries : Promise.resolve([]) + ); + window.electron.listFiles = listFiles; + + const { rerender } = render(, { + wrapper: IntlTestWrapper, + }); + await waitFor(() => expect(listFiles).toHaveBeenCalledTimes(1)); + + rerender(); + await act(async () => { + resolveRoot(['child']); + await rootEntries; + }); + + expect(listFiles).toHaveBeenCalledTimes(1); + }); + + it('keeps returning files and directories in ordinary trees', async () => { + const listFiles = vi.fn(async (directory: string) => { + if (directory === props.workingDir) return ['src', 'README.md']; + if (directory === `${props.workingDir}/src`) return ['index.ts']; + return []; + }); + window.electron.listFiles = listFiles; + + render(, { wrapper: IntlTestWrapper }); + + expect(await screen.findByText('README.md')).toBeInTheDocument(); + expect(await screen.findByText('index.ts')).toBeInTheDocument(); + expect(listFiles).toHaveBeenCalledTimes(3); + }); +}); diff --git a/ui/desktop/src/components/MentionPopover.tsx b/ui/desktop/src/components/MentionPopover.tsx index ccdcadb5e..c70e8e5de 100644 --- a/ui/desktop/src/components/MentionPopover.tsx +++ b/ui/desktop/src/components/MentionPopover.tsx @@ -38,6 +38,27 @@ const i18n = defineMessages({ type CommandItemType = 'Builtin' | 'Recipe' | 'Skill' | 'Agent'; type DisplayItemType = CommandItemType | 'Directory' | 'File'; +const MAX_SCAN_OPERATIONS = 100; +const MAX_SCAN_RESULTS = 500; + +interface DirectoryScanBudget { + remainingOperations: number; + remainingResults: number; + isCancelled: () => boolean; +} + +const reserveScanOperation = (budget: DirectoryScanBudget): boolean => { + if (budget.isCancelled() || budget.remainingOperations === 0) return false; + budget.remainingOperations--; + return true; +}; + +const reserveScanResult = (budget: DirectoryScanBudget): boolean => { + if (budget.isCancelled() || budget.remainingResults === 0) return false; + budget.remainingResults--; + return true; +}; + const typeOrder: Record = { Agent: 0, Directory: 1, @@ -165,12 +186,19 @@ const MentionPopover = forwardRef< const currentWorkingDir = workingDir ?? getInitialWorkingDir(); const scanDirectoryFromRoot = useCallback( - async (dirPath: string, relativePath = '', depth = 0): Promise => { + async ( + dirPath: string, + budget: DirectoryScanBudget, + relativePath = '', + depth = 0 + ): Promise => { // Increase depth limit for better file discovery - if (depth > 5) return []; + if (depth > 5 || budget.isCancelled() || budget.remainingResults === 0) return []; try { + if (!reserveScanOperation(budget)) return []; const items = await window.electron.listFiles(dirPath); + if (budget.isCancelled()) return []; const results: DisplayItem[] = []; // Common directories to prioritize or skip @@ -231,6 +259,7 @@ const MentionPopover = forwardRef< const itemLimit = depth === 0 ? 50 : depth === 1 ? 40 : 30; for (const item of sortedItems.slice(0, itemLimit)) { + if (budget.isCancelled() || budget.remainingResults === 0) return results; const fullPath = `${dirPath}/${item}`; const itemRelativePath = relativePath ? `${relativePath}/${item}` : item; @@ -328,6 +357,7 @@ const MentionPopover = forwardRef< // If it has a known file extension, treat it as a file if (hasExtension && ext && commonExtensions.includes(ext)) { + if (!reserveScanResult(budget)) return results; results.push({ extra: fullPath, name: item, @@ -347,6 +377,7 @@ const MentionPopover = forwardRef< 'makefile', ]; if (!hasExtension && knownFiles.includes(item.toLowerCase())) { + if (!reserveScanResult(budget)) return results; results.push({ extra: fullPath, name: item, @@ -358,8 +389,11 @@ const MentionPopover = forwardRef< // Otherwise, try to determine if it's a directory try { + if (!reserveScanOperation(budget)) return results; await window.electron.listFiles(fullPath); + if (budget.isCancelled()) return results; + if (!reserveScanResult(budget)) return results; results.push({ name: item, extra: fullPath, @@ -369,7 +403,12 @@ const MentionPopover = forwardRef< // Recursively scan directories more aggressively if (depth < 4 || priorityDirs.includes(item)) { - const subFiles = await scanDirectoryFromRoot(fullPath, itemRelativePath, depth + 1); + const subFiles = await scanDirectoryFromRoot( + fullPath, + budget, + itemRelativePath, + depth + 1 + ); results.push(...subFiles); } } catch { @@ -498,7 +537,11 @@ const MentionPopover = forwardRef< // Fetch agents from server and scan files in parallel const [agentItems, scannedFiles] = await Promise.all([ listAgentMentionItems(currentWorkingDir, sessionId ?? undefined).catch(() => []), - scanDirectoryFromRoot(currentWorkingDir || getDefaultStartPath()), + scanDirectoryFromRoot(currentWorkingDir || getDefaultStartPath(), { + remainingOperations: MAX_SCAN_OPERATIONS, + remainingResults: MAX_SCAN_RESULTS, + isCancelled: () => cancelled, + }), ]); if (cancelled) return; setItems([...agentItems, ...scannedFiles]);