Add @ command handling to useGeminiStream (#217)

* First integration of at commands into useGeminiStream.ts

* feat: Integrate @ command for file/directory reading

   - Adds support for `@<path>` commands in the CLI UI to read file or directory contents using the `read_many_files` tool.
   - Refactors `useGeminiStream` hook to handle slash, passthrough, and @ commands before sending queries to the Gemini API.
   - Improves history item ID generation to prevent React duplicate key warnings.

* fix: Handle additional text after @ command path

   - Modifies the `@` command processor to parse text following the file/directory path (e.g., `@README.md explain this`).
   - Includes both the fetched file content and the subsequent text in the query sent to the Gemini API.
   - Resolves the TODO item in `atCommandProcessor.ts`.

* feat: Allow @ command anywhere in query and fix build

   - Update `atCommandProcessor` to correctly parse `@<path>` commands regardless of their position in the input string using regex. This enables queries like "Explain @README.md to me".
   - Fix build error in `useGeminiStream` by importing the missing `findSafeSplitPoint` function.

* rename isPotentiallyAtCommand to isAtCommand

* respond to review comments.

3 files changed, 117 insertions, 68 deletions

diff --git a/packages/cli/src/ui/hooks/atCommandProcessor.ts b/packages/cli/src/ui/hooks/atCommandProcessor.ts
index 314c969d..a075157d 100644
--- a/packages/cli/src/ui/hooks/atCommandProcessor.ts
+++ b/packages/cli/src/ui/hooks/atCommandProcessor.ts
@@ -12,7 +12,7 @@ import {
   ToolCallStatus,
 } from '../types.js';
 
-// Helper function to add history items (could be moved to a shared util if needed elsewhere)
+// Helper function to add history items
 const addHistoryItem = (
   setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
   itemData: Omit<HistoryItem, 'id'>,
@@ -25,7 +25,7 @@ const addHistoryItem = (
 };
 
 interface HandleAtCommandParams {
-  query: string; // Raw user input
+  query: string; // Raw user input, potentially containing '@'
   config: Config;
   setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>;
   setDebugMessage: React.Dispatch<React.SetStateAction<string>>;
@@ -34,17 +34,18 @@ interface HandleAtCommandParams {
 }
 
 interface HandleAtCommandResult {
-  processedQuery: PartListUnion; // Query to potentially send to Gemini
+  processedQuery: PartListUnion | null; // Query for Gemini (null on error/no-proceed)
   shouldProceed: boolean; // Whether the main hook should continue processing
 }
 
 /**
- * Processes user input that might start with the '@' command to read files/directories.
- * If it's an '@' command, it attempts to read the specified path, updates the UI
- * with the tool call status, and prepares the query to be sent to the LLM.
+ * Processes user input potentially containing an '@<path>' command.
+ * It finds the first '@<path>', reads the specified path, updates the UI,
+ * and prepares the query for the LLM, incorporating the file content
+ * and surrounding text.
  *
- * @returns An object containing the potentially modified query and a flag
- *          indicating if the main hook should proceed with the Gemini API call.
+ * @returns An object containing the potentially modified query (or null)
+ *          and a flag indicating if the main hook should proceed.
  */
 export async function handleAtCommand({
   query,
@@ -56,42 +57,50 @@ export async function handleAtCommand({
 }: HandleAtCommandParams): Promise<HandleAtCommandResult> {
   const trimmedQuery = query.trim();
 
-  if (!trimmedQuery.startsWith('@')) {
-    // Not an '@' command, proceed as normal
-    // Add the user message here before returning
+  // Regex to find the first occurrence of @ followed by non-whitespace chars
+  // It captures the text before, the @path itself (including @), and the text after.
+  const atCommandRegex = /^(.*?)(@\S+)(.*)$/s; // s flag for dot to match newline
+  const match = trimmedQuery.match(atCommandRegex);
+
+  if (!match) {
+    // This should technically not happen if isPotentiallyAtCommand was true,
+    // but handle defensively.
+    const errorTimestamp = getNextMessageId(userMessageTimestamp);
     addHistoryItem(
       setHistory,
-      { type: 'user', text: query },
-      userMessageTimestamp,
+      { type: 'error', text: 'Error: Could not parse @ command.' },
+      errorTimestamp,
     );
-    // Use property shorthand for processedQuery
-    return { processedQuery: query, shouldProceed: true };
+    return { processedQuery: null, shouldProceed: false };
   }
 
-  // --- It is an '@' command ---
-  const filePath = trimmedQuery.substring(1);
+  const textBefore = match[1].trim();
+  const atPath = match[2]; // Includes the '@'
+  const textAfter = match[3].trim();
+
+  const pathPart = atPath.substring(1); // Remove the leading '@'
+
+  // Add user message for the full original @ command
+  addHistoryItem(
+    setHistory,
+    { type: 'user', text: query }, // Use original full query for history
+    userMessageTimestamp,
+  );
 
-  if (!filePath) {
-    // Handle case where it's just "@" - treat as normal input
+  if (!pathPart) {
+    // Handle case where it's just "@" or "@ " - treat as error/don't proceed
+    const errorTimestamp = getNextMessageId(userMessageTimestamp);
     addHistoryItem(
       setHistory,
-      { type: 'user', text: query },
-      userMessageTimestamp,
+      { type: 'error', text: 'Error: No path specified after @.' },
+      errorTimestamp,
     );
-    // Use property shorthand for processedQuery
-    return { processedQuery: query, shouldProceed: true }; // Send the "@" to the model
+    return { processedQuery: null, shouldProceed: false };
   }
 
   const toolRegistry = config.getToolRegistry();
   const readManyFilesTool = toolRegistry.getTool('read_many_files');
 
-  // Add user message first, so it appears before potential errors/tool UI
-  addHistoryItem(
-    setHistory,
-    { type: 'user', text: query },
-    userMessageTimestamp,
-  );
-
   if (!readManyFilesTool) {
     const errorTimestamp = getNextMessageId(userMessageTimestamp);
     addHistoryItem(
@@ -99,23 +108,21 @@ export async function handleAtCommand({
       { type: 'error', text: 'Error: read_many_files tool not found.' },
       errorTimestamp,
     );
-    // Use property shorthand for processedQuery
-    return { processedQuery: query, shouldProceed: false }; // Don't proceed if tool is missing
+    return { processedQuery: null, shouldProceed: false }; // Don't proceed if tool is missing
   }
 
   // --- Path Handling for @ command ---
-  let pathSpec = filePath;
+  let pathSpec = pathPart; // Use the extracted path part
   // Basic check: If no extension or ends with '/', assume directory and add globstar.
-  if (!filePath.includes('.') || filePath.endsWith('/')) {
-    pathSpec = filePath.endsWith('/') ? `${filePath}**` : `${filePath}/**`;
+  if (!pathPart.includes('.') || pathPart.endsWith('/')) {
+    pathSpec = pathPart.endsWith('/') ? `${pathPart}**` : `${pathPart}/**`;
   }
   const toolArgs = { paths: [pathSpec] };
   const contentLabel =
-    pathSpec === filePath ? filePath : `directory ${filePath}`; // Adjust label
+    pathSpec === pathPart ? pathPart : `directory ${pathPart}`; // Adjust label
   // --- End Path Handling ---
 
   let toolCallDisplay: IndividualToolCallDisplay;
-  let processedQuery: PartListUnion = query; // Default to original query
 
   try {
     setDebugMessage(`Reading via @ command: ${pathSpec}`);
@@ -132,15 +139,19 @@ export async function handleAtCommand({
       confirmationDetails: undefined,
     };
 
-    // Prepend file content to the query sent to the model
-    processedQuery = [
-      {
-        text: `--- Content from: ${contentLabel} ---
-${fileContent}
---- End Content ---`,
-      },
-      // TODO: Handle cases like "@README.md explain this" by appending the rest of the query
-    ];
+    // Construct the query for Gemini, combining parts
+    const processedQueryParts = [];
+    if (textBefore) {
+      processedQueryParts.push({ text: textBefore });
+    }
+    processedQueryParts.push({
+      text: `\n--- Content from: ${contentLabel} ---\n${fileContent}\n--- End Content ---`,
+    });
+    if (textAfter) {
+      processedQueryParts.push({ text: textAfter });
+    }
+
+    const processedQuery: PartListUnion = processedQueryParts;
 
     // Add the tool group UI
     const toolGroupId = getNextMessageId(userMessageTimestamp);
@@ -153,7 +164,6 @@ ${fileContent}
       toolGroupId,
     );
 
-    // Use property shorthand for processedQuery
     return { processedQuery, shouldProceed: true }; // Proceed to Gemini
   } catch (error) {
     // Construct error UI
@@ -177,7 +187,6 @@ ${fileContent}
       toolGroupId,
     );
 
-    // Use property shorthand for processedQuery
-    return { processedQuery: query, shouldProceed: false }; // Don't proceed on error
+    return { processedQuery: null, shouldProceed: false }; // Don't proceed on error
   }
 }
diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts
index 63a02fca..7e1f2177 100644
--- a/packages/cli/src/ui/hooks/useGeminiStream.ts
+++ b/packages/cli/src/ui/hooks/useGeminiStream.ts
@@ -27,9 +27,11 @@ import {
   IndividualToolCallDisplay,
   ToolCallStatus,
 } from '../types.js';
-import { findSafeSplitPoint } from '../utils/markdownUtilities.js';
+import { isAtCommand } from '../utils/commandUtils.js'; // Import the @ command checker
 import { useSlashCommandProcessor } from './slashCommandProcessor.js';
 import { usePassthroughProcessor } from './passthroughCommandProcessor.js';
+import { handleAtCommand } from './atCommandProcessor.js'; // Import the @ command handler
+import { findSafeSplitPoint } from '../utils/markdownUtilities.js'; // Import the split point finder
 
 const addHistoryItem = (
   setHistory: React.Dispatch<React.SetStateAction<HistoryItem[]>>,
@@ -61,6 +63,7 @@ export const useGeminiStream = (
 
   // ID Generation Callback
   const getNextMessageId = useCallback((baseTimestamp: number): number => {
+    // Increment *before* adding to ensure uniqueness against the base timestamp
     messageIdCounterRef.current += 1;
     return baseTimestamp + messageIdCounterRef.current;
   }, []);
@@ -144,32 +147,63 @@ export const useGeminiStream = (
       if (typeof query === 'string' && query.trim().length === 0) return;
 
       const userMessageTimestamp = Date.now();
+      messageIdCounterRef.current = 0; // Reset counter for this new submission
+      let queryToSendToGemini: PartListUnion | null = null;
 
       if (typeof query === 'string') {
-        setDebugMessage(`User query: '${query}'`);
+        const trimmedQuery = query.trim();
+        setDebugMessage(`User query: '${trimmedQuery}'`);
 
         // 1. Check for Slash Commands
-        if (handleSlashCommand(query)) {
-          return; // Command was handled, exit early
+        if (handleSlashCommand(trimmedQuery)) {
+          return; // Handled, exit
         }
 
         // 2. Check for Passthrough Commands
-        if (handlePassthroughCommand(query)) {
-          return; // Command was handled, exit early
+        if (handlePassthroughCommand(trimmedQuery)) {
+          return; // Handled, exit
         }
 
-        // 3. Add user message if not handled by slash/passthrough
-        addHistoryItem(
-          setHistory,
-          { type: 'user', text: query },
-          userMessageTimestamp,
-        );
+        // 3. Check for @ Commands using the utility function
+        if (isAtCommand(trimmedQuery)) {
+          const atCommandResult = await handleAtCommand({
+            query: trimmedQuery,
+            config,
+            setHistory,
+            setDebugMessage,
+            getNextMessageId,
+            userMessageTimestamp,
+          });
+
+          if (!atCommandResult.shouldProceed) {
+            return; // @ command handled it (e.g., error) or decided not to proceed
+          }
+          queryToSendToGemini = atCommandResult.processedQuery;
+          // User message and tool UI were added by handleAtCommand
+        } else {
+          // 4. It's a normal query for Gemini
+          addHistoryItem(
+            setHistory,
+            { type: 'user', text: trimmedQuery },
+            userMessageTimestamp,
+          );
+          queryToSendToGemini = trimmedQuery;
+        }
       } else {
-        // For function responses (PartListUnion that isn't a string),
-        // we don't add a user message here. The tool call/response UI handles it.
+        // 5. It's a function response (PartListUnion that isn't a string)
+        // Tool call/response UI handles history. Always proceed.
+        queryToSendToGemini = query;
+      }
+
+      // --- Proceed to Gemini API call ---
+      if (queryToSendToGemini === null) {
+        // Should only happen if @ command failed and returned null query
+        setDebugMessage(
+          'Query processing resulted in null, not sending to Gemini.',
+        );
+        return;
       }
 
-      // 4. Proceed to Gemini API call
       const client = geminiClientRef.current;
       if (!client) {
         setInitError('Gemini client is not available.');
@@ -188,7 +222,6 @@ export const useGeminiStream = (
 
       setStreamingState(StreamingState.Responding);
       setInitError(null);
-      messageIdCounterRef.current = 0; // Reset counter for new submission
       const chat = chatSessionRef.current;
       let currentToolGroupId: number | null = null;
 
@@ -196,8 +229,12 @@ export const useGeminiStream = (
         abortControllerRef.current = new AbortController();
         const signal = abortControllerRef.current.signal;
 
-        // Use the original query for the Gemini call
-        const stream = client.sendMessageStream(chat, query, signal);
+        // Use the determined query for the Gemini call
+        const stream = client.sendMessageStream(
+          chat,
+          queryToSendToGemini,
+          signal,
+        );
 
         // Process the stream events from the server logic
         let currentGeminiText = ''; // To accumulate message content
@@ -488,6 +525,9 @@ export const useGeminiStream = (
       updateGeminiMessage,
       handleSlashCommand,
       handlePassthroughCommand,
+      // handleAtCommand is implicitly included via its direct call
+      setDebugMessage, // Added dependency for handleAtCommand & passthrough
+      setStreamingState, // Added dependency for handlePassthroughCommand
       updateAndAddGeminiMessageContent,
     ],
   );
diff --git a/packages/cli/src/ui/utils/commandUtils.ts b/packages/cli/src/ui/utils/commandUtils.ts
index 4f731bb3..89e207d9 100644
--- a/packages/cli/src/ui/utils/commandUtils.ts
+++ b/packages/cli/src/ui/utils/commandUtils.ts
@@ -12,7 +12,7 @@
  * @param query The input query string.
  * @returns True if the query looks like an '@' command, false otherwise.
  */
-export const isPotentiallyAtCommand = (query: string): boolean =>
+export const isAtCommand = (query: string): boolean =>
   // Check if starts with @ OR has a space, then @, then a non-space character.
   query.startsWith('@') || /\s@\S/.test(query);