From b1258dd52c9ca5710e7c96dc183c8961ff0830a2 Mon Sep 17 00:00:00 2001
From: Sandy Tao <sandytao520@icloud.com>
Date: Thu, 20 Nov 2025 13:54:16 +0800
Subject: [PATCH] fix(cli): prevent race condition when restoring prompt after
 context overflow (#13473)

---
 packages/cli/src/ui/AppContainer.test.tsx | 57 +++++++++++++++++++++++
 packages/cli/src/ui/AppContainer.tsx      | 26 ++++++++++-
 2 files changed, 82 insertions(+), 1 deletion(-)

diff --git a/packages/cli/src/ui/AppContainer.test.tsx b/packages/cli/src/ui/AppContainer.test.tsx
index d37f3961ce..2adecff6cd 100644
--- a/packages/cli/src/ui/AppContainer.test.tsx
+++ b/packages/cli/src/ui/AppContainer.test.tsx
@@ -1834,5 +1834,62 @@ describe('AppContainer State Management', () => {
 
       unmount();
     });
+
+    it('correctly restores prompt even if userMessages is stale (race condition fix)', async () => {
+      // Setup initial history with one message
+      const initialHistory = [{ type: 'user', text: 'Previous Prompt' }];
+      mockedUseHistory.mockReturnValue({
+        history: initialHistory,
+        addItem: vi.fn(),
+        updateItem: vi.fn(),
+        clearItems: vi.fn(),
+        loadHistory: vi.fn(),
+      });
+
+      // Mock logger to resolve so userMessages gets populated
+      mockedUseLogger.mockReturnValue({
+        getPreviousUserMessages: vi.fn().mockResolvedValue([]),
+      });
+
+      const { unmount, rerender } = renderAppContainer();
+
+      // Wait for userMessages to be populated with 'Previous Prompt'
+      await waitFor(() =>
+        expect(capturedUIState.userMessages).toContain('Previous Prompt'),
+      );
+
+      // Simulate a new prompt being added (e.g., user sent it, but it overflowed)
+      const newPrompt = 'Current Prompt that Overflowed';
+      const newHistory = [...initialHistory, { type: 'user', text: newPrompt }];
+
+      mockedUseHistory.mockReturnValue({
+        history: newHistory,
+        addItem: vi.fn(),
+        updateItem: vi.fn(),
+        clearItems: vi.fn(),
+        loadHistory: vi.fn(),
+      });
+
+      // Rerender to reflect the history change.
+      // This triggers the effect to update userMessages, but it's async.
+      rerender(getAppContainer());
+
+      const { onCancelSubmit } = extractUseGeminiStreamArgs(
+        mockedUseGeminiStream.mock.lastCall!,
+      );
+
+      // Call onCancelSubmit immediately (simulating the race condition where
+      // the overflow event comes in before the effect updates userMessages)
+      act(() => {
+        onCancelSubmit(true);
+      });
+
+      // With the fix, it should wait for userMessages to update and then set the new prompt
+      await waitFor(() => {
+        expect(mockSetText).toHaveBeenCalledWith(newPrompt);
+      });
+
+      unmount();
+    });
   });
 });
diff --git a/packages/cli/src/ui/AppContainer.tsx b/packages/cli/src/ui/AppContainer.tsx
index 2b51f98613..47765de699 100644
--- a/packages/cli/src/ui/AppContainer.tsx
+++ b/packages/cli/src/ui/AppContainer.tsx
@@ -170,6 +170,7 @@ export const AppContainer = (props: AppContainerProps) => {
     null,
   );
   const [copyModeEnabled, setCopyModeEnabled] = useState(false);
+  const [pendingRestorePrompt, setPendingRestorePrompt] = useState(false);
 
   const [shellModeActive, setShellModeActive] = useState(false);
   const [modelSwitchedFromQuotaError, setModelSwitchedFromQuotaError] =
@@ -666,9 +667,32 @@ Logging in with Google... Please restart Gemini CLI to continue.
   );
 
   const onCancelSubmit = useCallback((shouldRestorePrompt?: boolean) => {
-    cancelHandlerRef.current(shouldRestorePrompt);
+    if (shouldRestorePrompt) {
+      setPendingRestorePrompt(true);
+    } else {
+      setPendingRestorePrompt(false);
+      cancelHandlerRef.current(false);
+    }
   }, []);
 
+  useEffect(() => {
+    if (pendingRestorePrompt) {
+      const lastHistoryUserMsg = historyManager.history.findLast(
+        (h) => h.type === 'user',
+      );
+      const lastUserMsg = userMessages.at(-1);
+
+      if (
+        !lastHistoryUserMsg ||
+        (typeof lastHistoryUserMsg.text === 'string' &&
+          lastHistoryUserMsg.text === lastUserMsg)
+      ) {
+        cancelHandlerRef.current(true);
+        setPendingRestorePrompt(false);
+      }
+    }
+  }, [pendingRestorePrompt, userMessages, historyManager.history]);
+
   const {
     streamingState,
     submitQuery,