feat(app): add @ workspace file autocomplete in agent chat

2026-02-18 13:25:55 +07:00
parent b1955cf74d
commit d30bb84741
16 changed files with 7993 additions and 7680 deletions
@@ -1,125 +1,290 @@
-import { useCallback, useMemo } from "react";
-import type { AutocompleteOption } from "@/components/ui/autocomplete";
-import { useAgentCommandsQuery, type DraftCommandConfig } from "./use-agent-commands-query";
-import { orderAutocompleteOptions } from "@/components/ui/autocomplete-utils";
-import { useAutocomplete } from "./use-autocomplete";
+import { useCallback, useMemo } from 'react'
+import { useQuery } from '@tanstack/react-query'
+import type { AutocompleteOption } from '@/components/ui/autocomplete'
+import { useAgentCommandsQuery, type DraftCommandConfig } from './use-agent-commands-query'
+import { orderAutocompleteOptions } from '@/components/ui/autocomplete-utils'
+import { useAutocomplete } from './use-autocomplete'
+import { useSessionStore } from '@/stores/session-store'
+import {
+  applyFileMentionReplacement,
+  findActiveFileMention,
+  type FileMentionRange,
+} from '@/utils/file-mention-autocomplete'

 interface UseAgentAutocompleteInput {
-  userInput: string;
-  setUserInput: (nextValue: string) => void;
-  serverId: string;
-  agentId: string;
-  draftConfig?: DraftCommandConfig;
-  onAutocompleteApplied?: () => void;
+  userInput: string
+  cursorIndex: number
+  setUserInput: (nextValue: string) => void
+  serverId: string
+  agentId: string
+  draftConfig?: DraftCommandConfig
+  onAutocompleteApplied?: () => void
 }

+type AgentAutocompleteOption =
+  | (AutocompleteOption & { type: 'command' })
+  | (AutocompleteOption & {
+      type: 'workspace_entry'
+      entryPath: string
+      mention: FileMentionRange
+    })
+
 interface AgentAutocompleteResult {
-  isVisible: boolean;
-  options: AutocompleteOption[];
-  selectedIndex: number;
-  isLoading: boolean;
-  errorMessage?: string;
-  onSelectOption: (option: AutocompleteOption) => void;
-  onKeyPress: (event: { key: string; preventDefault: () => void }) => boolean;
+  isVisible: boolean
+  options: AutocompleteOption[]
+  selectedIndex: number
+  isLoading: boolean
+  errorMessage?: string
+  loadingText: string
+  emptyText: string
+  onSelectOption: (option: AutocompleteOption) => void
+  onKeyPress: (event: { key: string; preventDefault: () => void }) => boolean
+}
+
+interface DirectorySuggestionEntry {
+  path: string
+  kind: 'file' | 'directory'
 }

 function normalizeDraftCommandConfig(
  draftConfig?: DraftCommandConfig
 ): DraftCommandConfig | undefined {
  if (!draftConfig) {
-    return undefined;
+    return undefined
  }

-  const cwd = draftConfig.cwd.trim();
+  const cwd = draftConfig.cwd.trim()
  if (!cwd) {
-    return undefined;
+    return undefined
  }

-  const modeId = draftConfig.modeId?.trim() ?? "";
-  const model = draftConfig.model?.trim() ?? "";
-  const thinkingOptionId = draftConfig.thinkingOptionId?.trim() ?? "";
+  const modeId = draftConfig.modeId?.trim() ?? ''
+  const model = draftConfig.model?.trim() ?? ''
+  const thinkingOptionId = draftConfig.thinkingOptionId?.trim() ?? ''
  return {
    provider: draftConfig.provider,
    cwd,
    ...(modeId ? { modeId } : {}),
    ...(model ? { model } : {}),
    ...(thinkingOptionId ? { thinkingOptionId } : {}),
-  };
+  }
 }

-export function useAgentAutocomplete(
-  input: UseAgentAutocompleteInput
-): AgentAutocompleteResult {
+function mapDirectorySuggestionsToEntries(payload: {
+  entries?: Array<{ path: string; kind: string }>
+  directories?: string[]
+}): DirectorySuggestionEntry[] {
+  if (Array.isArray(payload.entries) && payload.entries.length > 0) {
+    return payload.entries.flatMap((entry) => {
+      if (
+        !entry ||
+        typeof entry.path !== 'string' ||
+        (entry.kind !== 'file' && entry.kind !== 'directory')
+      ) {
+        return []
+      }
+      return [{ path: entry.path, kind: entry.kind }]
+    })
+  }
+
+  return (payload.directories ?? []).map((path) => ({
+    path,
+    kind: 'directory' as const,
+  }))
+}
+
+export function useAgentAutocomplete(input: UseAgentAutocompleteInput): AgentAutocompleteResult {
  const {
    userInput,
+    cursorIndex,
    setUserInput,
    serverId,
    agentId,
    draftConfig,
    onAutocompleteApplied,
-  } = input;
+  } = input

-  const showAutocomplete = userInput.startsWith("/") && !userInput.includes(" ");
-  const filterQuery = showAutocomplete ? userInput.slice(1) : "";
+  const showCommandAutocomplete = userInput.startsWith('/') && !userInput.includes(' ')
+  const commandFilterQuery = showCommandAutocomplete ? userInput.slice(1) : ''
+
+  const activeFileMention = useMemo(
+    () =>
+      findActiveFileMention({
+        text: userInput,
+        cursorIndex,
+      }),
+    [cursorIndex, userInput]
+  )
+  const showFileAutocomplete = activeFileMention !== null
+  const fileFilterQuery = activeFileMention?.query ?? ''

  const normalizedDraftConfig = useMemo(
    () => normalizeDraftCommandConfig(draftConfig),
    [draftConfig]
-  );
+  )

-  const isRealAgent = Boolean(agentId) && !agentId.startsWith("__");
-  const queryDraftConfig = isRealAgent ? undefined : normalizedDraftConfig;
-  const canLoadOptions = Boolean(serverId) && (isRealAgent || !!queryDraftConfig);
-  const isVisible = showAutocomplete && canLoadOptions;
+  const isRealAgent = Boolean(agentId) && !agentId.startsWith('__')
+  const queryDraftConfig = isRealAgent ? undefined : normalizedDraftConfig
+  const canLoadCommands = Boolean(serverId) && (isRealAgent || !!queryDraftConfig)

-  const { commands, isLoading, isError, error } = useAgentCommandsQuery({
+  const agentCwd = useSessionStore(
+    (state) => state.sessions[serverId]?.agents?.get(agentId)?.cwd ?? ''
+  )
+  const autocompleteCwd = useMemo(() => {
+    if (isRealAgent) {
+      return agentCwd.trim()
+    }
+    return queryDraftConfig?.cwd ?? ''
+  }, [agentCwd, isRealAgent, queryDraftConfig])
+
+  const client = useSessionStore((state) => state.sessions[serverId]?.client ?? null)
+  const isConnected = useSessionStore(
+    (state) => state.sessions[serverId]?.connection.isConnected ?? false
+  )
+
+  const mode: 'command' | 'file' | null = showFileAutocomplete
+    ? 'file'
+    : showCommandAutocomplete
+      ? 'command'
+      : null
+  const isVisible =
+    mode === 'command'
+      ? canLoadCommands
+      : mode === 'file'
+        ? Boolean(serverId) && autocompleteCwd.length > 0
+        : false
+
+  const {
+    commands,
+    isLoading: isCommandsLoading,
+    isError,
+    error,
+  } = useAgentCommandsQuery({
    serverId,
    agentId,
-    enabled: canLoadOptions,
+    enabled: mode === 'command' && canLoadCommands,
    draftConfig: queryDraftConfig,
-  });
+  })

-  const options = useMemo<AutocompleteOption[]>(() => {
+  const fileSuggestionsQuery = useQuery({
+    queryKey: ['directorySuggestions', serverId, autocompleteCwd, fileFilterQuery, true, true],
+    queryFn: async (): Promise<DirectorySuggestionEntry[]> => {
+      if (!client) {
+        throw new Error('Daemon client unavailable')
+      }
+      const response = await client.getDirectorySuggestions({
+        cwd: autocompleteCwd,
+        query: fileFilterQuery,
+        limit: 50,
+        includeFiles: true,
+        includeDirectories: true,
+      })
+      if (response.error) {
+        throw new Error(response.error)
+      }
+      return mapDirectorySuggestionsToEntries(response)
+    },
+    enabled:
+      mode === 'file' &&
+      Boolean(serverId) &&
+      autocompleteCwd.length > 0 &&
+      Boolean(client) &&
+      isConnected,
+    retry: false,
+    staleTime: 15_000,
+  })
+
+  const options = useMemo<AgentAutocompleteOption[]>(() => {
    if (!isVisible) {
-      return [];
+      return []
    }

-    const filterLower = filterQuery.toLowerCase();
-    const matches = commands.filter((cmd) =>
-      cmd.name.toLowerCase().includes(filterLower)
-    );
-    const orderedMatches = orderAutocompleteOptions(matches);
-    return orderedMatches.map((cmd) => ({
-      id: cmd.name,
-      label: `/${cmd.name}`,
-      detail: cmd.argumentHint || undefined,
-      description: cmd.description,
-    }));
-  }, [commands, filterQuery, isVisible]);
+    if (mode === 'command') {
+      const filterLower = commandFilterQuery.toLowerCase()
+      const matches = commands.filter((cmd) => cmd.name.toLowerCase().includes(filterLower))
+      const orderedMatches = orderAutocompleteOptions(matches)
+      return orderedMatches.map((cmd) => ({
+        type: 'command' as const,
+        id: cmd.name,
+        label: `/${cmd.name}`,
+        detail: cmd.argumentHint || undefined,
+        description: cmd.description,
+        kind: 'command',
+      }))
+    }
+
+    if (mode === 'file' && activeFileMention) {
+      return (fileSuggestionsQuery.data ?? []).map((entry) => ({
+        type: 'workspace_entry' as const,
+        id: `${entry.kind}:${entry.path}`,
+        label: entry.path,
+        detail: entry.kind === 'directory' ? 'Directory' : 'File',
+        kind: entry.kind,
+        entryPath: entry.path,
+        mention: activeFileMention,
+      }))
+    }
+
+    return []
+  }, [activeFileMention, commandFilterQuery, commands, fileSuggestionsQuery.data, isVisible, mode])

  const onSelectOption = useCallback(
    (option: AutocompleteOption) => {
-      setUserInput(`/${option.id} `);
-      onAutocompleteApplied?.();
+      const selected = option as AgentAutocompleteOption
+      if (selected.type === 'command') {
+        setUserInput(`/${selected.id} `)
+        onAutocompleteApplied?.()
+        return
+      }
+
+      const nextInput = applyFileMentionReplacement({
+        text: userInput,
+        mention: selected.mention,
+        relativePath: selected.entryPath,
+      })
+      setUserInput(nextInput)
+      onAutocompleteApplied?.()
    },
-    [onAutocompleteApplied, setUserInput]
-  );
+    [onAutocompleteApplied, setUserInput, userInput]
+  )

  const { selectedIndex, onKeyPress } = useAutocomplete({
    isVisible,
    options,
-    query: filterQuery,
+    query: mode === 'command' ? commandFilterQuery : fileFilterQuery,
    onSelectOption,
-    onEscape: () => setUserInput(""),
-  });
+    onEscape: mode === 'command' ? () => setUserInput('') : undefined,
+  })
+
+  const isLoading =
+    mode === 'command'
+      ? isCommandsLoading
+      : mode === 'file'
+        ? fileSuggestionsQuery.isPending || fileSuggestionsQuery.isLoading
+        : false
+  const errorMessage =
+    mode === 'command'
+      ? isError
+        ? (error?.message ?? 'Failed to load')
+        : undefined
+      : mode === 'file'
+        ? fileSuggestionsQuery.error instanceof Error
+          ? fileSuggestionsQuery.error.message
+          : undefined
+        : undefined
+
+  const loadingText = mode === 'file' ? 'Searching workspace...' : 'Loading commands...'
+  const emptyText = mode === 'file' ? 'No files or directories found' : 'No commands found'

  return {
    isVisible,
    options,
    selectedIndex,
    isLoading,
-    errorMessage: isError ? (error?.message ?? "Failed to load") : undefined,
+    errorMessage,
+    loadingText,
+    emptyText,
    onSelectOption,
    onKeyPress,
-  };
+  }
 }