From 889200d400c4dec60de0d7b5cdd77261bbb63edb Mon Sep 17 00:00:00 2001 From: Allen Hutchison Date: Tue, 29 Apr 2025 15:39:36 -0700 Subject: [PATCH] Add @ command handling to useGeminiStream (#217) * First integration of at commands into useGeminiStream.ts * feat: Integrate @ command for file/directory reading - Adds support for `@` 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 `@` 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. --- .../cli/src/ui/hooks/atCommandProcessor.ts | 107 ++++++++++-------- packages/cli/src/ui/hooks/useGeminiStream.ts | 76 ++++++++++--- packages/cli/src/ui/utils/commandUtils.ts | 2 +- 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>, itemData: Omit, @@ -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>; setDebugMessage: React.Dispatch>; @@ -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 '@' command. + * It finds the first '@', 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 { 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(); - if (!filePath) { - // Handle case where it's just "@" - treat as normal input + 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 (!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>, @@ -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);