diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 37dc325180f..1465a176399 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -697,6 +697,7 @@ export const AppContainer = (props: AppContainerProps) => { }, [config, historyManager, settings.merged]); const cancelHandlerRef = useRef<() => void>(() => {}); + const midTurnDrainRef = useRef<(() => string[]) | null>(null); const { streamingState, @@ -728,6 +729,7 @@ export const AppContainer = (props: AppContainerProps) => { setEmbeddedShellFocused, terminalWidth, terminalHeight, + midTurnDrainRef, ); // Track whether suggestions are visible for Tab key handling @@ -751,6 +753,18 @@ export const AppContainer = (props: AppContainerProps) => { submitQuery, }); + // Bridge message queue to mid-turn drain via ref. + // Sync ref on every render so the drain callback always reads latest state. + const messageQueueRef = useRef(messageQueue); + messageQueueRef.current = messageQueue; + midTurnDrainRef.current = () => { + const queue = messageQueueRef.current; + if (queue.length === 0) return []; + messageQueueRef.current = []; + clearQueue(); + return [...queue]; + }; + // Callback for handling final submit (must be after addMessage from useMessageQueue) const handleFinalSubmit = useCallback( (submittedValue: string) => { diff --git a/packages/cli/src/ui/hooks/useGeminiStream.ts b/packages/cli/src/ui/hooks/useGeminiStream.ts index 4dd2e952193..979f84748e4 100644 --- a/packages/cli/src/ui/hooks/useGeminiStream.ts +++ b/packages/cli/src/ui/hooks/useGeminiStream.ts @@ -172,6 +172,7 @@ export const useGeminiStream = ( setShellInputFocused: (value: boolean) => void, terminalWidth: number, terminalHeight: number, + midTurnDrainRef?: React.RefObject<(() => string[]) | null>, ) => { const [initError, setInitError] = useState(null); const abortControllerRef = useRef(null); @@ -1561,6 +1562,23 @@ export const useGeminiStream = ( return; } + // Mid-turn queue drain: inject queued user messages alongside tool + // results so the model sees them in the next API call. + // Skip if the turn was cancelled — messages stay in queue for next turn. + const drained = + turnCancelledRef.current || abortControllerRef.current?.signal.aborted + ? [] + : (midTurnDrainRef?.current?.() ?? []); + if (drained.length > 0) { + for (const msg of drained) { + responsesToSend.push({ + text: `\n[User message received during tool execution]: ${msg}`, + }); + // Record in UI history so the transcript stays complete. + addItem({ type: MessageType.USER, text: msg }, Date.now()); + } + } + submitQuery(responsesToSend, SendMessageType.ToolResult, prompt_ids[0]); }, [ @@ -1571,6 +1589,8 @@ export const useGeminiStream = ( performMemoryRefresh, modelSwitchedFromQuotaError, config, + midTurnDrainRef, + addItem, ], ); diff --git a/packages/cli/src/ui/hooks/useMessageQueue.ts b/packages/cli/src/ui/hooks/useMessageQueue.ts index 517040fecfa..00af564b38f 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.ts +++ b/packages/cli/src/ui/hooks/useMessageQueue.ts @@ -4,7 +4,7 @@ * SPDX-License-Identifier: Apache-2.0 */ -import { useCallback, useEffect, useState } from 'react'; +import { useCallback, useEffect, useRef, useState } from 'react'; import { StreamingState } from '../types.js'; export interface UseMessageQueueOptions { @@ -18,6 +18,12 @@ export interface UseMessageQueueReturn { addMessage: (message: string) => void; clearQueue: () => void; getQueuedMessagesText: () => string; + /** + * Atomically drain all queued messages. Returns the drained messages + * and clears both the synchronous ref and React state. Safe to call + * from non-React contexts (e.g., tool completion callbacks). + */ + drainQueue: () => string[]; } /** @@ -31,17 +37,22 @@ export function useMessageQueue({ submitQuery, }: UseMessageQueueOptions): UseMessageQueueReturn { const [messageQueue, setMessageQueue] = useState([]); + // Synchronous ref mirrors React state so non-React callbacks (e.g., + // mid-turn drain in handleCompletedTools) always see the latest queue. + const queueRef = useRef([]); // Add a message to the queue const addMessage = useCallback((message: string) => { const trimmedMessage = message.trim(); if (trimmedMessage.length > 0) { - setMessageQueue((prev) => [...prev, trimmedMessage]); + queueRef.current = [...queueRef.current, trimmedMessage]; + setMessageQueue(queueRef.current); } }, []); // Clear the entire queue const clearQueue = useCallback(() => { + queueRef.current = []; setMessageQueue([]); }, []); @@ -51,6 +62,15 @@ export function useMessageQueue({ return messageQueue.join('\n\n'); }, [messageQueue]); + // Atomically drain all queued messages (synchronous, safe from callbacks). + const drainQueue = useCallback((): string[] => { + const drained = queueRef.current; + if (drained.length === 0) return []; + queueRef.current = []; + setMessageQueue([]); + return drained; + }, []); + // Process queued messages when streaming becomes idle useEffect(() => { if ( @@ -61,15 +81,22 @@ export function useMessageQueue({ // Combine all messages with double newlines for clarity const combinedMessage = messageQueue.join('\n\n'); // Clear the queue and submit - setMessageQueue([]); + clearQueue(); submitQuery(combinedMessage); } - }, [isConfigInitialized, streamingState, messageQueue, submitQuery]); + }, [ + isConfigInitialized, + streamingState, + messageQueue, + submitQuery, + clearQueue, + ]); return { messageQueue, addMessage, clearQueue, getQueuedMessagesText, + drainQueue, }; }