fix(security): bound recursive mention scans (#11228)

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(<MentionPopover {...props} isOpen />, { 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<string[]>((resolve) => {
+      resolveRoot = resolve;
+    });
+    const listFiles = vi.fn((directory: string) =>
+      directory === props.workingDir ? rootEntries : Promise.resolve([])
+    );
+    window.electron.listFiles = listFiles;
+
+    const { rerender } = render(<MentionPopover {...props} isOpen />, {
+      wrapper: IntlTestWrapper,
+    });
+    await waitFor(() => expect(listFiles).toHaveBeenCalledTimes(1));
+
+    rerender(<MentionPopover {...props} isOpen={false} />);
+    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(<MentionPopover {...props} isOpen />, { wrapper: IntlTestWrapper });
+
+    expect(await screen.findByText('README.md')).toBeInTheDocument();
+    expect(await screen.findByText('index.ts')).toBeInTheDocument();
+    expect(listFiles).toHaveBeenCalledTimes(3);
+  });
+});

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<DisplayItemType, number> = {
   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<DisplayItem[]> => {
+      async (
+        dirPath: string,
+        budget: DirectoryScanBudget,
+        relativePath = '',
+        depth = 0
+      ): Promise<DisplayItem[]> => {
         // 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]);