Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
117 changes: 117 additions & 0 deletions packages/cli/src/ui/hooks/useAtCompletion.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import { describe, it, expect, beforeEach, vi, afterEach } from 'vitest';
import { act, useState } from 'react';
import * as path from 'node:path';
import { renderHook } from '../../test-utils/render.js';
import { waitFor } from '../../test-utils/async.js';
import { useAtCompletion } from './useAtCompletion.js';
Expand Down Expand Up @@ -589,4 +590,120 @@ describe('useAtCompletion', () => {
]);
});
});

describe('Multi-directory workspace support', () => {
const multiDirTmpDirs: string[] = [];

afterEach(async () => {
await Promise.all(multiDirTmpDirs.map((dir) => cleanupTmpDir(dir)));
multiDirTmpDirs.length = 0;
});

it('should include files from workspace directories beyond cwd', async () => {
const cwdStructure: FileSystemStructure = { 'main.txt': '' };
const addedDirStructure: FileSystemStructure = { 'added-file.txt': '' };
const cwdDir = await createTmpDir(cwdStructure);
multiDirTmpDirs.push(cwdDir);
const addedDir = await createTmpDir(addedDirStructure);
multiDirTmpDirs.push(addedDir);

const multiDirConfig = {
...mockConfig,
getWorkspaceContext: vi.fn().mockReturnValue({
getDirectories: () => [cwdDir, addedDir],
onDirectoriesChanged: vi.fn(() => () => {}),
}),
} as unknown as Config;

const { result } = renderHook(() =>
useTestHarnessForAtCompletion(true, '', multiDirConfig, cwdDir),
);

await waitFor(() => {
const values = result.current.suggestions.map((s) => s.value);
expect(values).toContain('main.txt');
expect(values).toContain(
escapePath(path.join(addedDir, 'added-file.txt')),
);
});
});

it('should pick up newly added directories via onDirectoriesChanged', async () => {
const cwdStructure: FileSystemStructure = { 'original.txt': '' };
const addedStructure: FileSystemStructure = { 'new-file.txt': '' };
const cwdDir = await createTmpDir(cwdStructure);
multiDirTmpDirs.push(cwdDir);
const addedDir = await createTmpDir(addedStructure);
multiDirTmpDirs.push(addedDir);

let dirChangeListener: (() => void) | null = null;
const directories = [cwdDir];

const dynamicConfig = {
...mockConfig,
getWorkspaceContext: vi.fn().mockReturnValue({
getDirectories: () => [...directories],
onDirectoriesChanged: vi.fn((listener: () => void) => {
dirChangeListener = listener;
return () => {
dirChangeListener = null;
};
}),
}),
} as unknown as Config;

const { result } = renderHook(() =>
useTestHarnessForAtCompletion(true, '', dynamicConfig, cwdDir),
);

await waitFor(() => {
const values = result.current.suggestions.map((s) => s.value);
expect(values).toContain('original.txt');
expect(values.every((v) => !v.includes('new-file.txt'))).toBe(true);
});

directories.push(addedDir);
act(() => {
dirChangeListener?.();
});

await waitFor(() => {
const values = result.current.suggestions.map((s) => s.value);
expect(values).toContain(
escapePath(path.join(addedDir, 'new-file.txt')),
);
});
});

it('should show same-named files from different directories without false deduplication', async () => {
const dir1Structure: FileSystemStructure = { 'readme.md': '' };
const dir2Structure: FileSystemStructure = { 'readme.md': '' };
const dir1 = await createTmpDir(dir1Structure);
multiDirTmpDirs.push(dir1);
const dir2 = await createTmpDir(dir2Structure);
multiDirTmpDirs.push(dir2);

const multiDirConfig = {
...mockConfig,
getWorkspaceContext: vi.fn().mockReturnValue({
getDirectories: () => [dir1, dir2],
onDirectoriesChanged: vi.fn(() => () => {}),
}),
} as unknown as Config;

const { result } = renderHook(() =>
useTestHarnessForAtCompletion(true, 'readme', multiDirConfig, dir1),
);

await waitFor(() => {
const values = result.current.suggestions.map((s) => s.value);
const readmeEntries = values.filter((v) => v.includes('readme.md'));
expect(readmeEntries.length).toBe(2);
expect(readmeEntries).toContain('readme.md');
expect(readmeEntries).toContain(
escapePath(path.join(dir2, 'readme.md')),
);
});
});
});
});
123 changes: 95 additions & 28 deletions packages/cli/src/ui/hooks/useAtCompletion.ts
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import { useEffect, useReducer, useRef } from 'react';
import { setTimeout as setTimeoutPromise } from 'node:timers/promises';
import * as path from 'node:path';
import type { Config, FileSearch } from '@google/gemini-cli-core';
import {
FileSearchFactory,
Expand Down Expand Up @@ -203,7 +204,8 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
setIsLoadingSuggestions,
} = props;
const [state, dispatch] = useReducer(atCompletionReducer, initialState);
const fileSearch = useRef<FileSearch | null>(null);
const fileSearchMap = useRef<Map<string, FileSearch>>(new Map());
const initEpoch = useRef(0);
const searchAbortController = useRef<AbortController | null>(null);
const slowSearchTimer = useRef<NodeJS.Timeout | null>(null);

Expand All @@ -215,10 +217,26 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
setIsLoadingSuggestions(state.isLoading);
}, [state.isLoading, setIsLoadingSuggestions]);

useEffect(() => {
const resetFileSearchState = () => {
fileSearchMap.current.clear();
initEpoch.current += 1;
dispatch({ type: 'RESET' });
};

useEffect(() => {
resetFileSearchState();
}, [cwd, config]);

useEffect(() => {
const workspaceContext = config?.getWorkspaceContext?.();
if (!workspaceContext) return;

const unsubscribe =
workspaceContext.onDirectoriesChanged(resetFileSearchState);

return unsubscribe;
}, [config]);

// Reacts to user input (`pattern`) ONLY.
useEffect(() => {
if (!enabled) {
Expand Down Expand Up @@ -250,38 +268,64 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
// The "Worker" that performs async operations based on status.
useEffect(() => {
const initialize = async () => {
const currentEpoch = initEpoch.current;
try {
const searcher = FileSearchFactory.create({
projectRoot: cwd,
ignoreDirs: [],
fileDiscoveryService: new FileDiscoveryService(
cwd,
config?.getFileFilteringOptions(),
),
cache: true,
cacheTtl: 30, // 30 seconds
enableRecursiveFileSearch:
config?.getEnableRecursiveFileSearch() ?? true,
enableFuzzySearch:
config?.getFileFilteringEnableFuzzySearch() ?? true,
maxFiles: config?.getFileFilteringOptions()?.maxFileCount,
});
await searcher.initialize();
fileSearch.current = searcher;
const directories = config
?.getWorkspaceContext?.()
?.getDirectories() ?? [cwd];

const initPromises: Array<Promise<void>> = [];

for (const dir of directories) {
if (fileSearchMap.current.has(dir)) continue;

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

This check appears to be redundant. The resetFileSearchState function, which is called before initialize is triggered (e.g., from onDirectoriesChanged or changes to cwd/config), clears the fileSearchMap. Therefore, this condition will always be false. Removing this check would simplify the code and make its behavior clearer, reflecting the full re-initialization strategy that is currently implemented. This also has performance implications, as it suggests an incremental update that isn't happening, potentially masking the inefficiency of re-initializing all searchers on every change.

References
  1. When initializing an object that requires fetching remote settings, prefer to gather all settings first, then perform a single initialization. Avoid creating and initializing a temporary object just to fetch settings, as this can lead to resource leaks and unnecessary overhead. This aligns with the principle of performing a single, clean initialization after gathering all necessary settings, avoiding redundant logic that suggests an incremental update strategy and can lead to unnecessary overhead and confusion.

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The check is not redundant. There are two distinct reset paths, both of which exist in the original upstream code (not introduced by this PR):

  1. Hard reset (resetFileSearchState()) - clears the map. Triggered by cwd/config changes and onDirectoriesChanged.
  2. Soft reset (dispatch({ type: 'RESET' }) alone, lines 248/253) - resets reducer state to IDLE but preserves the search cache. Triggered when the user exits @ mode (clears the @ prefix or disables completion).

The original single-FileSearch code had the same pattern: the [cwd, config] effect dispatched RESET but the soft-reset paths (disabled/pattern-null) also only dispatched RESET neither path nulled out the fileSearch ref. The ref was reused on re-entry, avoiding a redundant filesystem crawl.

Our has(dir) check preserves this existing caching behavior for the multi-directory case. After a soft reset, the map retains cached FileSearch instances. When the user re-enters @ mode, has(dir) skips re-initialization for already-indexed directories. This is the same cache optimization the original code relied on.


const searcher = FileSearchFactory.create({
projectRoot: dir,
ignoreDirs: [],
fileDiscoveryService: new FileDiscoveryService(
dir,
config?.getFileFilteringOptions(),
),
cache: true,
cacheTtl: 30,
enableRecursiveFileSearch:
config?.getEnableRecursiveFileSearch() ?? true,
enableFuzzySearch:
config?.getFileFilteringEnableFuzzySearch() ?? true,
maxFiles: config?.getFileFilteringOptions()?.maxFileCount,
});

initPromises.push(
searcher.initialize().then(() => {
if (initEpoch.current === currentEpoch) {
fileSearchMap.current.set(dir, searcher);
}
}),
);
}

await Promise.all(initPromises);

if (initEpoch.current !== currentEpoch) return;

dispatch({ type: 'INITIALIZE_SUCCESS' });
if (state.pattern !== null) {
dispatch({ type: 'SEARCH', payload: state.pattern });
}
} catch (_) {
dispatch({ type: 'ERROR' });
if (initEpoch.current === currentEpoch) {
dispatch({ type: 'ERROR' });
}
}
};

const search = async () => {
if (!fileSearch.current || state.pattern === null) {
if (fileSearchMap.current.size === 0 || state.pattern === null) {
return;
}

const currentPattern = state.pattern;

if (slowSearchTimer.current) {
clearTimeout(slowSearchTimer.current);
}
Expand Down Expand Up @@ -310,10 +354,26 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
})();

try {
const results = await fileSearch.current.search(state.pattern, {
signal: controller.signal,
maxResults: MAX_SUGGESTIONS_TO_SHOW * 3,
});
const directories = config
?.getWorkspaceContext?.()
?.getDirectories() ?? [cwd];
const cwdRealpath = directories[0];

const allSearchPromises = [...fileSearchMap.current.entries()].map(
async ([dir, searcher]): Promise<string[]> => {
const results = await searcher.search(currentPattern, {
signal: controller.signal,
maxResults: MAX_SUGGESTIONS_TO_SHOW * 3,
});

if (dir !== cwdRealpath) {
return results.map((p: string) => path.join(dir, p));
}
return results;
},
);

const allResults = await Promise.all(allSearchPromises);

if (slowSearchTimer.current) {
clearTimeout(slowSearchTimer.current);
Expand All @@ -323,15 +383,17 @@ export function useAtCompletion(props: UseAtCompletionProps): void {
return;
}

const fileSuggestions = results.map((p) => ({
const mergedResults = allResults.flat();

const fileSuggestions = mergedResults.map((p) => ({
label: p,
value: escapePath(p),
}));

const resourceCandidates = buildResourceCandidates(config);
const resourceSuggestions = (
await searchResourceCandidates(
state.pattern ?? '',
currentPattern ?? '',
resourceCandidates,
)
).map((suggestion) => ({
Expand All @@ -342,10 +404,15 @@ export function useAtCompletion(props: UseAtCompletionProps): void {

const agentCandidates = buildAgentCandidates(config);
const agentSuggestions = await searchAgentCandidates(
state.pattern ?? '',
currentPattern ?? '',
agentCandidates,
);

// Re-check after resource/agent searches which are not abort-aware
if (controller.signal.aborted) {
return;
}

const combinedSuggestions = [
...agentSuggestions,
...fileSuggestions,
Expand Down
Loading