fix(cli): prevent message queue interleaving by awaiting tool result submission

[imadraude]

Summary

This PR fixes a bug where queued messages could interleave with Gemini's response when tool calls were involved.

Details

The issue was caused by not awaiting the recursive submitQuery calls in useGeminiStream.ts. When a tool call completed, submitQuery was called without await, causing the isResponding state to momentarily flip to false (Idle). This allowed the useMessageQueue hook to trigger and inject queued messages before the model's next turn could start.

Changes:

• Added await to submitQuery calls in handleCompletedTools and loop detection logic.
• Made the loop detection onComplete callback async to support await.
• Fixed useGeminiStream.test.tsx by providing a concrete ValidationRequiredError class and getCurrentSequenceModel method in the @google/gemini-cli-core mock.

Related Issues

None.

How to Validate

1. Start Gemini CLI.
2. Send a prompt that triggers a tool call (e.g., @include GEMINI.md).
3. While the tool is executing, type another message and press Enter (it will be queued).
4. Verify that the queued message only appears after Gemini has finished its entire response (including any follow-up text after the tool result).

Alternatively, run the updated tests:

npm test -w @google/gemini-cli -- src/ui/hooks/useGeminiStream.test.tsx

Pre-Merge Checklist

• Updated relevant documentation and README (if needed)
• Added/updated tests (if needed)
• Noted breaking changes (if any)
• Validated on required platforms/methods:
  • MacOS
  • Windows
  • Linux
    • npm run

Fixes #17282

[gemini-code-assist]

Summary of Changes

Hello @imadraude, 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 addresses a critical timing issue in the Gemini CLI where user messages could appear out of order when the AI model was processing tool calls. The core problem stemmed from unawaited asynchronous calls, leading to an inconsistent state that allowed the message queue to inject messages prematurely. The fix ensures that all relevant asynchronous operations complete before the system processes new user input, thereby maintaining the correct conversational flow and preventing message interleaving.

Highlights

• Fix Message Interleaving: Resolved a bug where queued messages could interleave with Gemini's responses, especially when tool calls were involved, by ensuring submitQuery calls are properly awaited.
• Asynchronous Handling: Implemented await for recursive submitQuery calls within handleCompletedTools and loop detection logic to maintain correct state and prevent premature isResponding state changes.
• Test Suite Updates: Updated useGeminiStream.test.tsx to include necessary mocks for getCurrentSequenceModel and ValidationRequiredError to support the changes in the main logic.

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

[gemini-code-assist]

Code Review

This pull request addresses a race condition where queued messages could appear out of order during tool execution. The fix involves correctly awaiting the recursive submitQuery calls, ensuring the application state remains consistent throughout the asynchronous operations. The changes are well-targeted, and the necessary adjustments to test mocks have been made. The fix appears correct and complete.

[imadraude]

I have reset the branch to its initial state to remove unintended merge commits from main. The PR now correctly contains only the fix for message interleaving and the corresponding test updates.

[jackwotherspoon]

Looks like there is an existing PR for this that tries to solve the same race condition: #17057

Although the account that created the other PR seems like someones bot 😆

I am not sure how your test update is testing the race condition?

The other PR is a bit overly complex for the fix as well i believe, but it is definitely good to test.

In my eyes isn't this the proper fix? (from other users PR)

Awaiting the submit, and then marking as submitted...

Please give me details on your approach here as I compare the two solutions 😄

[gemini-code-assist]

Thanks for the questions, @jackwotherspoon! It's great to compare approaches and ensure we have the most robust solution.

You've hit on the core of the fix, and the image you provided from the other PR actually illustrates the same fundamental principle that this PR implements. The key issue was indeed unawaited asynchronous calls, specifically submitQuery.

How this PR addresses the race condition:

1. Awaiting submitQuery: The primary change in packages/cli/src/ui/hooks/useGeminiStream.ts is the introduction of await before calls to submitQuery within handleCompletedTools and the loop detection logic. Previously, these calls were not awaited, meaning the execution flow would continue, and the isResponding state could prematurely flip to false (Idle). This allowed the useMessageQueue hook to inject user messages while the model was still processing tool results or subsequent turns, leading to the interleaving.
2. async onComplete callback: Making the onComplete callback for setLoopDetectionConfirmationRequest async was necessary to allow the await submitQuery call within it. This ensures that even in scenarios involving loop detection confirmations, the query submission process is fully completed before the state is considered idle.

By awaiting submitQuery, we guarantee that the entire chain of operations initiated by that query (including tool executions and any recursive calls to submitQuery) completes before the isResponding state is updated. This prevents the message queue from acting on a false 'idle' signal.

How the test updates verify the fix:

The updates to useGeminiStream.test.tsx are crucial for testing this asynchronous behavior correctly:

• getCurrentSequenceModel mock: This mock is added to MockedGeminiClientClass because the submitQuery logic (which is now awaited) relies on getCurrentSequenceModel. Without this mock, the tests would fail when trying to access an undefined method on the mocked client.
• ValidationRequiredError mock: Similarly, the ValidationRequiredError class is mocked. This is needed because the submitQuery path can involve error handling related to validation, and the test environment needs a defined version of this class to prevent runtime errors during test execution.

These test updates don't directly test the race condition itself in terms of timing, but they enable the existing and new tests to run successfully against the modified asynchronous logic. The original tests for useGeminiStream (and any new ones added by the author) would then implicitly verify that the message interleaving no longer occurs under the conditions that previously triggered it, now that the awaits are in place.

In essence, both this PR and the one you linked are aiming for the same solution: ensuring that the submitQuery operation, and all its cascading effects, are fully resolved before the system signals that it's ready for new user input. This PR implements that by explicitly awaiting the calls where the race condition was identified.

[imadraude]

Updated the PR to await submitQuery before marking tools as submitted to ensure state consistency. Added regression tests to verify the fix and filtered out already-submitted tools in both the main and loop detection flows. This should be all set now.

[imadraude]

/gemini review

[gemini-code-assist]

Code Review

This pull request correctly identifies and fixes a race condition by awaiting submitQuery calls, preventing message queue interleaving. The changes are logical and the new tests provide good coverage for the fix. However, I've identified a significant issue with the new implementation where tools are marked as submitted even if the submitQuery call fails internally (e.g., due to a network error), which can lead to lost conversational context. I've provided a detailed comment with a suggested refactoring to address this. This should be addressed to ensure the reliability of tool call handling.

[gemini-code-assist]

While awaiting submitQuery is a great step towards fixing the race condition, calling markToolsAsSubmitted unconditionally afterwards can lead to data loss. If submitQuery encounters an error internally (like a network issue), it will handle the error and not re-throw, but the tool results will not have been successfully sent to the model. This implementation would then incorrectly mark the tools as submitted, leading to lost context in the conversation.

A more robust approach is to move the responsibility of marking tools as submitted into submitQuery itself, so it's only done upon successful completion of the stream.

This would involve a few changes:

1. Modify submitQuery to accept an optional toolCallIdsToMark array in its options.

2. In handleCompletedTools, pass the IDs to submitQuery and remove the markToolsAsSubmitted call.

   // in handleCompletedTools
   await submitQuery(
     responsesToSend,
     {
       isContinuation: true,
       toolCallIdsToMark: callIdsToMarkAsSubmitted,
     },
     prompt_ids[0],
   );

3. Inside submitQuery, call markToolsAsSubmitted only after the stream processing is fully successful and not in a loop detection path. Also, pass toolCallIdsToMark along in the recursive submitQuery call within the loop detection logic.

   // in submitQuery
   // ...
   try {
     // ...
     if (loopDetectedRef.current) {
       // ...
       setLoopDetectionConfirmationRequest({
         onComplete: async (result) => {
           if (result.userSelection === 'disable') {
             // ...
             await submitQuery(
               lastQueryRef.current,
               { isContinuation: true, toolCallIdsToMark: options?.toolCallIdsToMark },
               lastPromptIdRef.current,
             );
           }
         },
       });
     } else {
       // No loop detected, request was successful. Mark tools as submitted.
       if (options?.toolCallIdsToMark) {
         markToolsAsSubmitted(options.toolCallIdsToMark);
       }
     }
   } catch (e) {
     // On error, tools are not marked as submitted.
   }
   // ...

[imadraude]

/gemini review

[gemini-code-assist]

Code Review

This pull request addresses a critical race condition where queued messages could interleave with tool call responses. The root cause was identified as unawaited submitQuery calls, which has been fixed by adding await and adjusting the function signatures to support asynchronous operations. The changes are well-reasoned and correctly implemented. Key improvements include:

• Awaiting submitQuery calls within handleCompletedTools and the loop detection logic, ensuring proper execution order.
• Conditionally calling markToolsAsSubmitted only after a successful submission, which resolves the race condition.
• Adding a defensive check in handleCompletedTools to prevent resubmission of already processed tool calls, which fixes a potential 400 error.
• Updating the dependency array for the handleCompletedTools hook to prevent stale closures.
  The accompanying test changes are thorough, with new tests specifically validating the race condition fix and the prevention of tool resubmission. The overall changes are excellent and significantly improve the robustness of tool handling in the CLI.

[gemini-cli]

Hi there! Thank you for your contribution to Gemini CLI. We really appreciate the time and effort you've put into this pull request.

To keep our backlog manageable and ensure we're focusing on current priorities, we are closing pull requests that haven't seen maintainer activity for 30 days. Currently, the team is prioritizing work associated with 🔒 maintainer only or help wanted issues.

If you believe this change is still critical, please feel free to comment with updated details. Otherwise, we encourage contributors to focus on open issues labeled as help wanted. Thank you for your understanding!