/** * @license * Copyright 2025 Google LLC * SPDX-License-Identifier: Apache-2.0 */ 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, escapePath, FileDiscoveryService, } from '@google/gemini-cli-core'; import type { Suggestion } from '../components/SuggestionsDisplay.js'; import { MAX_SUGGESTIONS_TO_SHOW } from '../components/SuggestionsDisplay.js'; import { CommandKind } from '../commands/types.js'; import { AsyncFzf } from 'fzf'; const DEFAULT_SEARCH_TIMEOUT_MS = 5000; export enum AtCompletionStatus { IDLE = 'idle', INITIALIZING = 'initializing', READY = 'ready', SEARCHING = 'searching', ERROR = 'error', } interface AtCompletionState { status: AtCompletionStatus; suggestions: Suggestion[]; isLoading: boolean; pattern: string | null; } type AtCompletionAction = | { type: 'INITIALIZE' } | { type: 'INITIALIZE_SUCCESS' } | { type: 'SEARCH'; payload: string } | { type: 'SEARCH_SUCCESS'; payload: Suggestion[] } | { type: 'SET_LOADING'; payload: boolean } | { type: 'ERROR' } | { type: 'RESET' }; const initialState: AtCompletionState = { status: AtCompletionStatus.IDLE, suggestions: [], isLoading: false, pattern: null, }; function atCompletionReducer( state: AtCompletionState, action: AtCompletionAction, ): AtCompletionState { switch (action.type) { case 'INITIALIZE': return { ...state, status: AtCompletionStatus.INITIALIZING, isLoading: true, }; case 'INITIALIZE_SUCCESS': return { ...state, status: AtCompletionStatus.READY, isLoading: false }; case 'SEARCH': // Keep old suggestions, don't set loading immediately return { ...state, status: AtCompletionStatus.SEARCHING, pattern: action.payload, }; case 'SEARCH_SUCCESS': return { ...state, status: AtCompletionStatus.READY, suggestions: action.payload, isLoading: false, }; case 'SET_LOADING': // Only show loading if we are still in a searching state if (state.status === AtCompletionStatus.SEARCHING) { return { ...state, isLoading: action.payload, suggestions: [] }; } return state; case 'ERROR': return { ...state, status: AtCompletionStatus.ERROR, isLoading: false, suggestions: [], }; case 'RESET': return initialState; default: return state; } } export interface UseAtCompletionProps { enabled: boolean; pattern: string; config: Config | undefined; cwd: string; setSuggestions: (suggestions: Suggestion[]) => void; setIsLoadingSuggestions: (isLoading: boolean) => void; } interface ResourceSuggestionCandidate { searchKey: string; suggestion: Suggestion; } function buildResourceCandidates( config?: Config, ): ResourceSuggestionCandidate[] { const registry = config?.getResourceRegistry?.(); if (!registry) { return []; } const resources = registry.getAllResources().map((resource) => { // Use serverName:uri format to disambiguate resources from different MCP servers const prefixedUri = `${resource.serverName}:${resource.uri}`; return { // Include prefixedUri in searchKey so users can search by the displayed format searchKey: `${prefixedUri} ${resource.name ?? ''}`.toLowerCase(), suggestion: { label: prefixedUri, value: prefixedUri, }, } satisfies ResourceSuggestionCandidate; }); return resources; } function buildAgentCandidates(config?: Config): Suggestion[] { const registry = config?.getAgentRegistry?.(); if (!registry) { return []; } return registry.getAllDefinitions().map((def) => ({ label: def.name, value: def.name, commandKind: CommandKind.AGENT, })); } async function searchResourceCandidates( pattern: string, candidates: ResourceSuggestionCandidate[], ): Promise { if (candidates.length === 0) { return []; } const normalizedPattern = pattern.toLowerCase(); if (!normalizedPattern) { return candidates .slice(0, MAX_SUGGESTIONS_TO_SHOW) .map((candidate) => candidate.suggestion); } const fzf = new AsyncFzf(candidates, { selector: (candidate: ResourceSuggestionCandidate) => candidate.searchKey, }); // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment const results = await fzf.find(normalizedPattern, { limit: MAX_SUGGESTIONS_TO_SHOW * 3, }); // eslint-disable-next-line @typescript-eslint/no-unsafe-return return results.map( (result: { item: ResourceSuggestionCandidate }) => result.item.suggestion, ); } async function searchAgentCandidates( pattern: string, candidates: Suggestion[], ): Promise { if (candidates.length === 0) { return []; } const normalizedPattern = pattern.toLowerCase(); if (!normalizedPattern) { return candidates.slice(0, MAX_SUGGESTIONS_TO_SHOW); } const fzf = new AsyncFzf(candidates, { selector: (s: Suggestion) => s.label, }); // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment const results = await fzf.find(normalizedPattern, { limit: MAX_SUGGESTIONS_TO_SHOW, }); // eslint-disable-next-line @typescript-eslint/no-unsafe-return return results.map((r: { item: Suggestion }) => r.item); } export function useAtCompletion(props: UseAtCompletionProps): void { const { enabled, pattern, config, cwd, setSuggestions, setIsLoadingSuggestions, } = props; const [state, dispatch] = useReducer(atCompletionReducer, initialState); const fileSearchMap = useRef>(new Map()); const initEpoch = useRef(0); const searchAbortController = useRef(null); const slowSearchTimer = useRef(null); useEffect(() => { setSuggestions(state.suggestions); }, [state.suggestions, setSuggestions]); useEffect(() => { setIsLoadingSuggestions(state.isLoading); }, [state.isLoading, setIsLoadingSuggestions]); 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) { // reset when first getting out of completion suggestions if ( state.status === AtCompletionStatus.READY || state.status === AtCompletionStatus.ERROR ) { dispatch({ type: 'RESET' }); } return; } if (pattern === null) { dispatch({ type: 'RESET' }); return; } if (state.status === AtCompletionStatus.IDLE) { dispatch({ type: 'INITIALIZE' }); } else if ( (state.status === AtCompletionStatus.READY || state.status === AtCompletionStatus.SEARCHING) && pattern.toLowerCase() !== state.pattern // Only search if the pattern has changed ) { dispatch({ type: 'SEARCH', payload: pattern.toLowerCase() }); } }, [enabled, pattern, state.status, state.pattern]); // The "Worker" that performs async operations based on status. useEffect(() => { const initialize = async () => { const currentEpoch = initEpoch.current; try { const directories = config ?.getWorkspaceContext?.() ?.getDirectories() ?? [cwd]; const initPromises: Array> = []; for (const dir of directories) { if (fileSearchMap.current.has(dir)) continue; 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 (_) { if (initEpoch.current === currentEpoch) { dispatch({ type: 'ERROR' }); } } }; const search = async () => { if (fileSearchMap.current.size === 0 || state.pattern === null) { return; } const currentPattern = state.pattern; if (slowSearchTimer.current) { clearTimeout(slowSearchTimer.current); } const controller = new AbortController(); searchAbortController.current = controller; slowSearchTimer.current = setTimeout(() => { dispatch({ type: 'SET_LOADING', payload: true }); }, 200); const timeoutMs = config?.getFileFilteringOptions()?.searchTimeout ?? DEFAULT_SEARCH_TIMEOUT_MS; // eslint-disable-next-line @typescript-eslint/no-floating-promises (async () => { try { await setTimeoutPromise(timeoutMs, undefined, { signal: controller.signal, }); controller.abort(); } catch { // ignore } })(); try { const directories = config ?.getWorkspaceContext?.() ?.getDirectories() ?? [cwd]; const cwdRealpath = directories[0]; const allSearchPromises = [...fileSearchMap.current.entries()].map( async ([dir, searcher]): Promise => { 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); } if (controller.signal.aborted) { return; } const mergedResults = allResults.flat(); const fileSuggestions = mergedResults.map((p) => ({ label: p, value: escapePath(p), })); const resourceCandidates = buildResourceCandidates(config); const resourceSuggestions = ( await searchResourceCandidates( currentPattern ?? '', resourceCandidates, ) ).map((suggestion) => ({ ...suggestion, label: suggestion.label.replace(/^@/, ''), value: suggestion.value.replace(/^@/, ''), })); const agentCandidates = buildAgentCandidates(config); const agentSuggestions = await searchAgentCandidates( currentPattern ?? '', agentCandidates, ); // Re-check after resource/agent searches which are not abort-aware if (controller.signal.aborted) { return; } const combinedSuggestions = [ ...agentSuggestions, ...fileSuggestions, ...resourceSuggestions, ]; dispatch({ type: 'SEARCH_SUCCESS', payload: combinedSuggestions }); } catch (error) { if (!(error instanceof Error && error.name === 'AbortError')) { dispatch({ type: 'ERROR' }); } } finally { controller.abort(); } }; if (state.status === AtCompletionStatus.INITIALIZING) { // eslint-disable-next-line @typescript-eslint/no-floating-promises initialize(); } else if (state.status === AtCompletionStatus.SEARCHING) { // eslint-disable-next-line @typescript-eslint/no-floating-promises search(); } return () => { searchAbortController.current?.abort(); if (slowSearchTimer.current) { clearTimeout(slowSearchTimer.current); } }; }, [state.status, state.pattern, config, cwd]); }