feat: allow editing queued messages with up arrow key (#10392)

Co-authored-by: Akhil Appana <akhilapp@google.com>

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,
     ],
   );
 

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."

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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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(
+        <InputPrompt {...props} />,
+      );
+      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;

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<InputPromptProps> = ({
   isEmbeddedShellFocused,
   setQueueErrorMessage,
   streamingState,
+  popAllMessages,
 }) => {
   const kittyProtocol = useKittyKeyboardProtocol();
   const isShellFocused = useShellFocusState();
@@ -288,6 +290,23 @@ export const InputPrompt: React.FC<InputPromptProps> = ({
     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<InputPromptProps> = ({
         }
 
         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<InputPromptProps> = ({
           (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<InputPromptProps> = ({
       commandSearchActive,
       commandSearchCompletion,
       kittyProtocol.supported,
+      tryLoadQueuedMessages,
     ],
   );
 

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');
   });
 });

packages/cli/src/ui/components/QueuedMessageDisplay.tsx

@@ -21,13 +21,16 @@ export const QueuedMessageDisplay = ({
 
   return (
     <Box flexDirection="column" marginTop={1}>
+      <Box paddingLeft={2}>
+        <Text dimColor>Queued (press ↑ to edit):</Text>
+      </Box>
       {messageQueue
         .slice(0, MAX_DISPLAYED_QUEUED_MESSAGES)
         .map((message, index) => {
           const preview = message.replace(/\s+/g, ' ');
 
           return (
-            <Box key={index} paddingLeft={2} width="100%">
+            <Box key={index} paddingLeft={4} width="100%">
               <Text dimColor wrap="truncate">
                 {preview}
               </Text>
@@ -35,7 +38,7 @@ export const QueuedMessageDisplay = ({
           );
         })}
       {messageQueue.length > MAX_DISPLAYED_QUEUED_MESSAGES && (
-        <Box paddingLeft={2}>
+        <Box paddingLeft={4}>
           <Text dimColor>
             ... (+
             {messageQueue.length - MAX_DISPLAYED_QUEUED_MESSAGES} more)

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<UIActions | null>(null);

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([]);
+    });
+  });
 });

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,
   };
 }