diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx index 6a97507042..b153d2121f 100644 --- a/packages/cli/src/ui/AppContainer.tsx +++ b/packages/cli/src/ui/AppContainer.tsx @@ -621,12 +621,17 @@ Logging in with Google... Please restart Gemini CLI to continue. onApprovalModeChange: handleApprovalModeChange, }); - const { messageQueue, addMessage, clearQueue, getQueuedMessagesText } = - useMessageQueue({ - isConfigInitialized, - streamingState, - submitQuery, - }); + const { + messageQueue, + addMessage, + clearQueue, + getQueuedMessagesText, + popAllMessages, + } = useMessageQueue({ + isConfigInitialized, + streamingState, + submitQuery, + }); cancelHandlerRef.current = useCallback(() => { const pendingHistoryItems = [ @@ -1306,6 +1311,7 @@ Logging in with Google... Please restart Gemini CLI to continue. onWorkspaceMigrationDialogClose, handleProQuotaChoice, setQueueErrorMessage, + popAllMessages, }), [ handleThemeSelect, @@ -1332,6 +1338,7 @@ Logging in with Google... Please restart Gemini CLI to continue. onWorkspaceMigrationDialogClose, handleProQuotaChoice, setQueueErrorMessage, + popAllMessages, ], ); diff --git a/packages/cli/src/ui/components/Composer.tsx b/packages/cli/src/ui/components/Composer.tsx index af68096349..46f765d981 100644 --- a/packages/cli/src/ui/components/Composer.tsx +++ b/packages/cli/src/ui/components/Composer.tsx @@ -148,6 +148,7 @@ export const Composer = () => { focus={true} vimHandleInput={uiActions.vimHandleInput} isEmbeddedShellFocused={uiState.embeddedShellFocused} + popAllMessages={uiActions.popAllMessages} placeholder={ vimEnabled ? " Press 'i' for INSERT mode and 'Esc' for NORMAL mode." diff --git a/packages/cli/src/ui/components/InputPrompt.test.tsx b/packages/cli/src/ui/components/InputPrompt.test.tsx index c5021f68da..5198dff2f5 100644 --- a/packages/cli/src/ui/components/InputPrompt.test.tsx +++ b/packages/cli/src/ui/components/InputPrompt.test.tsx @@ -2110,6 +2110,174 @@ describe('InputPrompt', () => { }); }); + describe('queued message editing', () => { + it('should load all queued messages when up arrow is pressed with empty input', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = ''; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + + expect(mockPopAllMessages).toHaveBeenCalled(); + const callback = mockPopAllMessages.mock.calls[0][0]; + + act(() => { + callback('Message 1\n\nMessage 2\n\nMessage 3'); + }); + expect(props.buffer.setText).toHaveBeenCalledWith( + 'Message 1\n\nMessage 2\n\nMessage 3', + ); + unmount(); + }); + + it('should not load queued messages when input is not empty', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = 'some text'; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + expect(mockPopAllMessages).not.toHaveBeenCalled(); + expect(mockInputHistory.navigateUp).toHaveBeenCalled(); + unmount(); + }); + + it('should handle undefined messages from popAllMessages', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = ''; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + + expect(mockPopAllMessages).toHaveBeenCalled(); + const callback = mockPopAllMessages.mock.calls[0][0]; + act(() => { + callback(undefined); + }); + + expect(props.buffer.setText).not.toHaveBeenCalled(); + expect(mockInputHistory.navigateUp).toHaveBeenCalled(); + unmount(); + }); + + it('should work with NAVIGATION_UP key as well', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = ''; + props.buffer.allVisualLines = ['']; + props.buffer.visualCursor = [0, 0]; + props.buffer.visualScrollRow = 0; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + expect(mockPopAllMessages).toHaveBeenCalled(); + unmount(); + }); + + it('should handle single queued message', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = ''; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + + const callback = mockPopAllMessages.mock.calls[0][0]; + act(() => { + callback('Single message'); + }); + + expect(props.buffer.setText).toHaveBeenCalledWith('Single message'); + unmount(); + }); + + it('should only check for queued messages when buffer text is trimmed empty', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = ' '; // Whitespace only + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + + expect(mockPopAllMessages).toHaveBeenCalled(); + unmount(); + }); + + it('should not call popAllMessages if it is not provided', async () => { + props.popAllMessages = undefined; + props.buffer.text = ''; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + + expect(mockInputHistory.navigateUp).toHaveBeenCalled(); + unmount(); + }); + + it('should navigate input history on fresh start when no queued messages exist', async () => { + const mockPopAllMessages = vi.fn(); + props.popAllMessages = mockPopAllMessages; + props.buffer.text = ''; + + const { stdin, unmount } = renderWithProviders( + , + ); + await wait(); + + stdin.write('\u001B[A'); + await wait(); + + expect(mockPopAllMessages).toHaveBeenCalled(); + + const callback = mockPopAllMessages.mock.calls[0][0]; + act(() => { + callback(undefined); + }); + + expect(mockInputHistory.navigateUp).toHaveBeenCalled(); + expect(props.buffer.setText).not.toHaveBeenCalled(); + + unmount(); + }); + }); + describe('snapshots', () => { it('should render correctly in shell mode', async () => { props.shellModeActive = true; diff --git a/packages/cli/src/ui/components/InputPrompt.tsx b/packages/cli/src/ui/components/InputPrompt.tsx index 18d623fb58..35ee31e1db 100644 --- a/packages/cli/src/ui/components/InputPrompt.tsx +++ b/packages/cli/src/ui/components/InputPrompt.tsx @@ -75,6 +75,7 @@ export interface InputPromptProps { isEmbeddedShellFocused?: boolean; setQueueErrorMessage: (message: string | null) => void; streamingState: StreamingState; + popAllMessages?: (onPop: (messages: string | undefined) => void) => void; } // The input content, input container, and input suggestions list may have different widths @@ -113,6 +114,7 @@ export const InputPrompt: React.FC = ({ isEmbeddedShellFocused, setQueueErrorMessage, streamingState, + popAllMessages, }) => { const kittyProtocol = useKittyKeyboardProtocol(); const isShellFocused = useShellFocusState(); @@ -288,6 +290,23 @@ export const InputPrompt: React.FC = ({ resetCommandSearchCompletionState, ]); + // Helper function to handle loading queued messages into input + // Returns true if we should continue with input history navigation + const tryLoadQueuedMessages = useCallback(() => { + if (buffer.text.trim() === '' && popAllMessages) { + popAllMessages((allMessages) => { + if (allMessages) { + buffer.setText(allMessages); + } else { + // No queued messages, proceed with input history + inputHistory.navigateUp(); + } + }); + return true; // We handled the up arrow key + } + return false; + }, [buffer, popAllMessages, inputHistory]); + // Handle clipboard image pasting with Ctrl+V const handleClipboardImage = useCallback(async () => { try { @@ -597,6 +616,12 @@ export const InputPrompt: React.FC = ({ } if (keyMatchers[Command.HISTORY_UP](key)) { + // Check for queued messages first when input is empty + // If no queued messages, inputHistory.navigateUp() is called inside tryLoadQueuedMessages + if (tryLoadQueuedMessages()) { + return; + } + // Only navigate history if popAllMessages doesn't exist inputHistory.navigateUp(); return; } @@ -610,6 +635,12 @@ export const InputPrompt: React.FC = ({ (buffer.allVisualLines.length === 1 || (buffer.visualCursor[0] === 0 && buffer.visualScrollRow === 0)) ) { + // Check for queued messages first when input is empty + // If no queued messages, inputHistory.navigateUp() is called inside tryLoadQueuedMessages + if (tryLoadQueuedMessages()) { + return; + } + // Only navigate history if popAllMessages doesn't exist inputHistory.navigateUp(); return; } @@ -753,6 +784,7 @@ export const InputPrompt: React.FC = ({ commandSearchActive, commandSearchCompletion, kittyProtocol.supported, + tryLoadQueuedMessages, ], ); diff --git a/packages/cli/src/ui/components/QueuedMessageDisplay.test.tsx b/packages/cli/src/ui/components/QueuedMessageDisplay.test.tsx index e041092fe1..8c5f026495 100644 --- a/packages/cli/src/ui/components/QueuedMessageDisplay.test.tsx +++ b/packages/cli/src/ui/components/QueuedMessageDisplay.test.tsx @@ -21,6 +21,7 @@ describe('QueuedMessageDisplay', () => { ); const output = lastFrame(); + expect(output).toContain('Queued (press ↑ to edit):'); expect(output).toContain('First message'); }); @@ -36,6 +37,7 @@ describe('QueuedMessageDisplay', () => { ); const output = lastFrame(); + expect(output).toContain('Queued (press ↑ to edit):'); expect(output).toContain('First queued message'); expect(output).toContain('Second queued message'); expect(output).toContain('Third queued message'); @@ -55,6 +57,7 @@ describe('QueuedMessageDisplay', () => { ); const output = lastFrame(); + expect(output).toContain('Queued (press ↑ to edit):'); expect(output).toContain('Message 1'); expect(output).toContain('Message 2'); expect(output).toContain('Message 3'); @@ -71,6 +74,7 @@ describe('QueuedMessageDisplay', () => { ); const output = lastFrame(); + expect(output).toContain('Queued (press ↑ to edit):'); expect(output).toContain('Message with multiple whitespace'); }); }); diff --git a/packages/cli/src/ui/components/QueuedMessageDisplay.tsx b/packages/cli/src/ui/components/QueuedMessageDisplay.tsx index a42e9feab1..edf10b5593 100644 --- a/packages/cli/src/ui/components/QueuedMessageDisplay.tsx +++ b/packages/cli/src/ui/components/QueuedMessageDisplay.tsx @@ -21,13 +21,16 @@ export const QueuedMessageDisplay = ({ return ( + + Queued (press ↑ to edit): + {messageQueue .slice(0, MAX_DISPLAYED_QUEUED_MESSAGES) .map((message, index) => { const preview = message.replace(/\s+/g, ' '); return ( - + {preview} @@ -35,7 +38,7 @@ export const QueuedMessageDisplay = ({ ); })} {messageQueue.length > MAX_DISPLAYED_QUEUED_MESSAGES && ( - + ... (+ {messageQueue.length - MAX_DISPLAYED_QUEUED_MESSAGES} more) diff --git a/packages/cli/src/ui/contexts/UIActionsContext.tsx b/packages/cli/src/ui/contexts/UIActionsContext.tsx index 555e2235d9..99d3c5a852 100644 --- a/packages/cli/src/ui/contexts/UIActionsContext.tsx +++ b/packages/cli/src/ui/contexts/UIActionsContext.tsx @@ -46,6 +46,7 @@ export interface UIActions { onWorkspaceMigrationDialogClose: () => void; handleProQuotaChoice: (choice: 'auth' | 'continue') => void; setQueueErrorMessage: (message: string | null) => void; + popAllMessages: (onPop: (messages: string | undefined) => void) => void; } export const UIActionsContext = createContext(null); diff --git a/packages/cli/src/ui/hooks/useMessageQueue.test.ts b/packages/cli/src/ui/hooks/useMessageQueue.test.ts index 33dbf3211c..d28f5fb250 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.test.ts +++ b/packages/cli/src/ui/hooks/useMessageQueue.test.ts @@ -232,4 +232,160 @@ describe('useMessageQueue', () => { expect(mockSubmitQuery).toHaveBeenCalledWith('Second batch'); expect(mockSubmitQuery).toHaveBeenCalledTimes(2); }); + + describe('popAllMessages', () => { + it('should pop all messages and return them joined with double newlines', () => { + const { result } = renderHook(() => + useMessageQueue({ + isConfigInitialized: true, + streamingState: StreamingState.Responding, + submitQuery: mockSubmitQuery, + }), + ); + + // Add multiple messages + act(() => { + result.current.addMessage('Message 1'); + result.current.addMessage('Message 2'); + result.current.addMessage('Message 3'); + }); + + expect(result.current.messageQueue).toEqual([ + 'Message 1', + 'Message 2', + 'Message 3', + ]); + + // Pop all messages + let poppedMessages: string | undefined; + act(() => { + result.current.popAllMessages((messages) => { + poppedMessages = messages; + }); + }); + + expect(poppedMessages).toBe('Message 1\n\nMessage 2\n\nMessage 3'); + expect(result.current.messageQueue).toEqual([]); + }); + + it('should return undefined when queue is empty', () => { + const { result } = renderHook(() => + useMessageQueue({ + isConfigInitialized: true, + streamingState: StreamingState.Responding, + submitQuery: mockSubmitQuery, + }), + ); + + let poppedMessages: string | undefined = 'not-undefined'; + act(() => { + result.current.popAllMessages((messages) => { + poppedMessages = messages; + }); + }); + + expect(poppedMessages).toBeUndefined(); + expect(result.current.messageQueue).toEqual([]); + }); + + it('should handle single message correctly', () => { + const { result } = renderHook(() => + useMessageQueue({ + isConfigInitialized: true, + streamingState: StreamingState.Responding, + submitQuery: mockSubmitQuery, + }), + ); + + act(() => { + result.current.addMessage('Single message'); + }); + + let poppedMessages: string | undefined; + act(() => { + result.current.popAllMessages((messages) => { + poppedMessages = messages; + }); + }); + + expect(poppedMessages).toBe('Single message'); + expect(result.current.messageQueue).toEqual([]); + }); + + it('should clear the entire queue after popping', () => { + const { result } = renderHook(() => + useMessageQueue({ + isConfigInitialized: true, + streamingState: StreamingState.Responding, + submitQuery: mockSubmitQuery, + }), + ); + + act(() => { + result.current.addMessage('Message 1'); + result.current.addMessage('Message 2'); + }); + + act(() => { + result.current.popAllMessages(() => {}); + }); + + // Queue should be empty + expect(result.current.messageQueue).toEqual([]); + expect(result.current.getQueuedMessagesText()).toBe(''); + + // Popping again should return undefined + let secondPop: string | undefined = 'not-undefined'; + act(() => { + result.current.popAllMessages((messages) => { + secondPop = messages; + }); + }); + + expect(secondPop).toBeUndefined(); + }); + + it('should work correctly with state updates', () => { + const { result } = renderHook(() => + useMessageQueue({ + isConfigInitialized: true, + streamingState: StreamingState.Responding, + submitQuery: mockSubmitQuery, + }), + ); + + // Add messages + act(() => { + result.current.addMessage('First'); + result.current.addMessage('Second'); + }); + + // Pop all messages + let firstPop: string | undefined; + act(() => { + result.current.popAllMessages((messages) => { + firstPop = messages; + }); + }); + + expect(firstPop).toBe('First\n\nSecond'); + + // Add new messages after popping + act(() => { + result.current.addMessage('Third'); + result.current.addMessage('Fourth'); + }); + + // Pop again + let secondPop: string | undefined; + act(() => { + result.current.popAllMessages((messages) => { + secondPop = messages; + }); + }); + + expect(secondPop).toBe('Third\n\nFourth'); + expect(result.current.messageQueue).toEqual([]); + }); + }); }); diff --git a/packages/cli/src/ui/hooks/useMessageQueue.ts b/packages/cli/src/ui/hooks/useMessageQueue.ts index 517040fecf..23f0218574 100644 --- a/packages/cli/src/ui/hooks/useMessageQueue.ts +++ b/packages/cli/src/ui/hooks/useMessageQueue.ts @@ -18,6 +18,7 @@ export interface UseMessageQueueReturn { addMessage: (message: string) => void; clearQueue: () => void; getQueuedMessagesText: () => string; + popAllMessages: (onPop: (messages: string | undefined) => void) => void; } /** @@ -51,6 +52,23 @@ export function useMessageQueue({ return messageQueue.join('\n\n'); }, [messageQueue]); + // Pop all messages from the queue and return them as a single string + const popAllMessages = useCallback( + (onPop: (messages: string | undefined) => void) => { + setMessageQueue((prev) => { + if (prev.length === 0) { + onPop(undefined); + return prev; + } + // Join all messages with double newlines, same as when they're sent + const allMessages = prev.join('\n\n'); + onPop(allMessages); + return []; // Clear the entire queue + }); + }, + [], + ); + // Process queued messages when streaming becomes idle useEffect(() => { if ( @@ -71,5 +89,6 @@ export function useMessageQueue({ addMessage, clearQueue, getQueuedMessagesText, + popAllMessages, }; }