fix(cli): prevent race condition when restoring prompt after context overflow

[SandyTao520]

TLDR

Fixes a bug where the input prompt was restored with stale text (the previous message) instead of the current overflowing message due to a race condition between the ContextWindowWillOverflow event and the userMessages state update.

Dive Deeper

The ContextWindowWillOverflow event triggers onCancelSubmit synchronously. However, userMessages (used to retrieve the text to restore) is updated asynchronously via a useEffect observing history changes. This PR introduces a pendingRestorePrompt state to defer the restoration until userMessages reflects the latest history item.

Reviewer Test Plan

1. Set a low token limit or use a very long prompt that triggers ContextWindowWillOverflow.
2. Send the prompt.
3. Verify that when the overflow warning appears and the prompt is restored, it contains the current long prompt, not the previous one.
4. Run the regression test: npm run test --workspace packages/cli -- src/ui/AppContainer.test.tsx

Testing Matrix

	🍏	🪟	🐧
npm run	✅	❓	❓
npx	✅	❓	❓
Docker	❓	❓	❓
Podman	❓	-	-
Seatbelt	❓	-	-

Linked issues / bugs

[gemini-code-assist]

Summary of Changes

Hello @SandyTao520, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request resolves a critical race condition within the CLI that caused the input prompt to display an incorrect, outdated message after a context window overflow warning. By implementing a deferred restoration mechanism, the system now correctly waits for the latest user message state to synchronize before restoring the prompt, ensuring users always see their current input.

Highlights

• Race Condition Fix: Addresses a race condition where the CLI's input prompt would restore with stale text (the previous message) instead of the current overflowing message after a context overflow warning.
• Deferred Prompt Restoration: Introduces a pendingRestorePrompt state to defer the prompt restoration process until the userMessages state, which updates asynchronously, is fully synchronized with the latest history.
• Synchronous Event vs. Asynchronous State: Resolves the conflict between the synchronous ContextWindowWillOverflow event triggering onCancelSubmit and the asynchronous update of userMessages via a useEffect.
• New Test Case: Adds a dedicated test case in AppContainer.test.tsx to specifically validate the fix for the prompt restoration race condition, ensuring the correct prompt is restored.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

Size Change: +722 B (0%)

Total Size: 21.1 MB

ℹ️ View Unchanged

Filename	Size	Change
./bundle/gemini.js	21.1 MB	+722 B (0%)
./bundle/sandbox-macos-permissive-closed.sb	1.03 kB	0 B
./bundle/sandbox-macos-permissive-open.sb	890 B	0 B
./bundle/sandbox-macos-permissive-proxied.sb	1.31 kB	0 B
./bundle/sandbox-macos-restrictive-closed.sb	3.29 kB	0 B
./bundle/sandbox-macos-restrictive-open.sb	3.36 kB	0 B
./bundle/sandbox-macos-restrictive-proxied.sb	3.56 kB	0 B

_{compressed-size-action}

[gemini-code-assist]

Code Review

This pull request aims to fix a race condition when restoring a user's prompt after a context window overflow. The approach of introducing a pendingRestorePrompt state and a useEffect to handle the asynchronous nature of the state updates is sound. However, the implementation of the new useEffect hook introduces a potential deadlock. It relies on the userMessages state to be in sync with historyManager.history, but the effect that syncs them is asynchronous and lacks error handling. If that sync fails, the prompt restoration will never occur, leaving the UI in a stuck state. I've proposed a more robust implementation for the useEffect that removes this fragile dependency and makes the logic self-contained.

[SandyTao520]

Great catch on the race condition! The solution correctly handles the asynchronous nature of the userMessages state update by deferring the prompt restoration until the state is synchronized.

The implementation of pendingRestorePrompt and the associated useEffect is a clean way to resolve this within the React lifecycle without needing to plumb the prompt text through multiple layers of event handlers.

The added regression test in AppContainer.test.tsx effectively mocks the race condition (stale userMessages) and verifies that the fix waits for the update.

Code quality looks good and follows the project patterns.

Verification:

• Verified waitFor is imported from ../test-utils/async.js.
• Verified act usage in tests.
• Logic handles the case where shouldRestorePrompt is false correctly by clearing the pending state.

LGTM! 🚀

[SandyTao520]

Thank you for the detailed analysis.

While I acknowledge the theoretical risk if logger.getPreviousUserMessages() fails, that would represent a broader system failure affecting the entire chat history display. In that scenario, the prompt restoration failing is a secondary concern.

I prefer to maintain the current implementation to keep the restoration logic centralized within cancelHandlerRef (DRY) rather than duplicating the queue handling logic inside the useEffect. The current approach effectively resolves the race condition for the happy path, which is the primary goal.

[jacob314]

The current implementation effectively mitigates the race condition by waiting for the state to settle.

However, I see an opportunity to simplify this significantly and remove the race condition entirely by passing the prompt text directly.

In useGeminiStream.ts, the lastQueryRef already holds the text of the prompt being processed.
We could update onCancelSubmit to accept an optional restoredText: string argument.

In useGeminiStream.ts:

const handleContextWindowWillOverflowEvent = useCallback(
  (estimatedRequestTokenCount: number, remainingTokenCount: number) => {
    // Pass the actual text that caused the overflow
    const textToRestore = typeof lastQueryRef.current === 'string' ? lastQueryRef.current : undefined;
    onCancelSubmit(textToRestore);
    // ...

In AppContainer.tsx:

const onCancelSubmit = useCallback((restoredText?: string) => {
  // ...
  cancelHandlerRef.current(restoredText);
}, []);

// Update cancelHandlerRef to use the provided text directly
// instead of looking up userMessages.at(-1)

This approach:

1. Eliminates the race condition: We use the source of truth (lastQueryRef) instead of waiting for a derived state (userMessages) to sync.
2. Reduces Complexity: Removes the need for pendingRestorePrompt state and the useEffect in AppContainer.
3. Handles Edge Cases: Correctly handles cases where userMessages might be stale for other reasons.

Would you consider this refactor? It would likely require updating the test case to verify that the passed text is restored, rather than simulating the history timing.

[jacob314]

Approved after offline discussion.