fix: Tackling Race/State condition issue by Changing the Code Design for Gemini Grounding Sources #7434

HahaBill · 2025-08-26T22:47:02Z

Related GitHub Issue

Closes: #6372

Description

Before we were collecting sources in from the stream chunks in api/gemini.ts. I tried to resolve the issue by finding a culprit responsible for the race condition and state inconsistencies. It turns out that there were leakages from previous message to new messages. I tried very hard to resolve this but then I told myself that maybe it's better to rethink the code design pattern.

I came up with this with a help of Roo Code:

Basically, we are decoupling gemini stream chunks and grounding metadata.

Summary:

adding new type ApiStreamGroundingChunk to the stream type
collecting sources in the Task.ts instead -> decoupling

Test Procedure

I test this with gemini-2.5-flash. For some reason, there are so many errors with gemini-2.5-pro.

Go to the profile and pick Google Gemini
Then pick gemini-2.5-flash
Then enable Grounding with Google Search
Start asking recent/latest info about some topics
It should generate sources correctly without any source leakage from previous messages
It works way better than before - code implementation looks cleaner

Pre-Submission Checklist

Issue Linked: This PR is linked to an approved GitHub Issue (see "Related GitHub Issue" above).
Scope: My changes are focused on the linked issue (one major feature/fix per PR).
Self-Review: I have performed a thorough self-review of my code.
Testing: New and/or updated tests have been added to cover my changes (if applicable).
Documentation Impact: I have considered if my changes require documentation updates (see "Documentation Updates" section below).
Contribution Guidelines: I have read and agree to the Contributor Guidelines.

Discord username: HaHaBill (belikebill.)

Important

Redesigns code to decouple gemini stream chunks and grounding metadata, preventing state persistence issues and improving clarity.

Behavior:
- Decouples gemini stream chunks and grounding metadata in gemini.ts.
- Moves source collection to Task.ts to prevent state persistence issues.
- Introduces ApiStreamGroundingChunk type in stream.ts.
Functions:
- extractGroundingSources() in gemini.ts now returns GroundingSource[] instead of a string.
- extractCitationsOnly() in gemini.ts uses extractGroundingSources().
Classes:
- Task in Task.ts now handles grounding sources separately and strips them from messages before saving.
Types:
- Adds ApiStreamGroundingChunk and GroundingSource to stream.ts.

^{This description was created by}^{for a250a7e. You can customize this summary. It will automatically update as commits are pushed.}

- adding new type `ApiStreamGroundingChunk` to the stream type - collecting sources in the `Task.ts` instead -> decoupling

ellipsis-dev · 2025-08-26T22:49:18Z

src/core/task/Task.ts

+					if (pendingGroundingSources.length > 0) {
+						// Remove any grounding source references that might have been integrated into the message
+						cleanAssistantMessage = assistantMessage
+							.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "")


When cleaning the assistant message to remove integrated grounding source references, the regex patterns are nontrivial. Consider adding an inline comment with example inputs to clarify what patterns are being stripped to improve future maintainability.

Suggested change

.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "")

\t\t\t\t\t\t.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "") // e.g. "[1] Example Source: https://example.com"

roomote

Thank you for this well-thought-out solution to the Gemini grounding sources race condition issue! The architectural change to decouple stream chunks from grounding metadata is a solid approach that properly addresses the state persistence problems.

The implementation looks clean and follows good separation of concerns. I've left a few suggestions inline to improve maintainability and robustness, with the most important being the addition of test coverage to prevent regression of this race condition fix.

roomote · 2025-08-26T22:52:37Z

src/core/task/Task.ts

+					if (pendingGroundingSources.length > 0) {
+						// Remove any grounding source references that might have been integrated into the message
+						cleanAssistantMessage = assistantMessage
+							.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "")


I agree with the previous suggestion - these regex patterns would benefit from inline comments explaining what they match. Consider adding comments like:

Suggested change

.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "")

cleanAssistantMessage = assistantMessage

.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "") // e.g., "[1] Example Source: https://example.com"

.replace(/Sources?:\s*[\s\S]*?(?=\n\n|\n$|$)/g, "") // e.g., "Sources: [1](url1), [2](url2)"

roomote · 2025-08-26T22:52:37Z

src/api/providers/gemini.ts

+					return {
+						title,
+						url: uri,
+					}
 				}
 				return null


Consider adding debug logging when sources are filtered out to help with troubleshooting malformed grounding metadata during development. This would make it easier to identify issues with the grounding chunks.

roomote · 2025-08-26T22:52:37Z

src/api/transform/stream.ts

@@ -27,3 +32,14 @@ export interface ApiStreamUsageChunk {
 	reasoningTokens?: number
 	totalCost?: number
 }
+
+export interface ApiStreamGroundingChunk {


Consider adding JSDoc documentation for the new grounding types to clarify their purpose. For example:

/** * Represents grounding metadata from search results or citations. * Used to decouple source information from the main content stream. */ export interface ApiStreamGroundingChunk { type: "grounding" sources: GroundingSource[] }

roomote · 2025-08-26T22:52:37Z

src/core/task/Task.ts

+							case "grounding":
+								// Handle grounding sources separately from regular content
+								// to prevent state persistence issues - store them separately
+								if (chunk.sources && chunk.sources.length > 0) {


Good separation of concerns here! The grounding sources are now properly handled independently from the assistant message content, which should prevent the race condition issues mentioned in #6372.

roomote · 2025-08-26T22:52:37Z

src/core/task/Task.ts

@@ -2086,9 +2094,30 @@
 				let didEndLoop = false

 				if (assistantMessage.length > 0) {
+					// Display grounding sources to the user if they exist
+					if (pendingGroundingSources.length > 0) {


Nice implementation detail - storing grounding sources separately and only displaying them to the user while keeping the API history clean. This prevents the state persistence issues that were causing the race condition.

- Fixed failing test to expect grounding chunk instead of text with citations - Added inline comments to regex patterns for source stripping in Task.ts - Added test coverage for grounding source handling to prevent regression

daniel-lxs

Thank you @HahaBill!

HahaBill · 2025-08-26T23:48:42Z

Thank you @HahaBill!

Anytime :)

mrubens · 2025-08-27T02:31:15Z

src/core/task/Task.ts

+					// Display grounding sources to the user if they exist
+					if (pendingGroundingSources.length > 0) {
+						const citationLinks = pendingGroundingSources.map((source, i) => `[${i + 1}](${source.url})`)
+						const sourcesText = `Sources: ${citationLinks.join(", ")}`


Should probably internationalize "Sources"

Right, I will add it back.

mrubens · 2025-08-27T02:32:05Z

src/core/task/Task.ts

+					if (pendingGroundingSources.length > 0) {
+						// Remove any grounding source references that might have been integrated into the message
+						cleanAssistantMessage = assistantMessage
+							.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "") // e.g., "[1] Example Source: https://example.com"


I'm not sure about this. Does it have the potential to remove some content from the message unrelated to grounding? Is this a necessary thing to do in general?

I see, those are valid concerns. I personally think that it is unlikely unless users have custom prompts saying that the model should include in the response the grounding-like format content in some parts of the response.

But this might be premature optimization/edge-case handling so we can remove this for now.

Overall, these are answers to your questions:

There's a small non-zero risk.

No, it's not necessary. This might be a premature optimization.

HahaBill · 2025-08-27T14:46:24Z

Alright, I made some additional changes:
commit 7e9a8c4c00741a7b3343dc8855679818218575fe

feat: Tackling Race/State condition issue by Changing the Code Design

a250a7e

- adding new type `ApiStreamGroundingChunk` to the stream type - collecting sources in the `Task.ts` instead -> decoupling

HahaBill requested review from mrubens, cte and jr as code owners August 26, 2025 22:47

github-project-automation bot added this to Roo Code Roadmap and Roo Code Roadmap Aug 26, 2025

github-project-automation bot moved this to New in Roo Code Roadmap Aug 26, 2025

github-project-automation bot moved this to Triage in Roo Code Roadmap Aug 26, 2025

dosubot bot added size:M This PR changes 30-99 lines, ignoring generated files. bug Something isn't working labels Aug 26, 2025

HahaBill changed the title ~~feat: Tackling Race/State condition issue by Changing the Code Design~~ fix: Tackling Race/State condition issue by Changing the Code Design for Gemini Grounding Sources Aug 26, 2025

ellipsis-dev bot reviewed Aug 26, 2025

View reviewed changes

daniel-lxs moved this from Triage to PR [Needs Prelim Review] in Roo Code Roadmap Aug 26, 2025

roomote bot reviewed Aug 26, 2025

View reviewed changes

hannesrudolph added the PR - Needs Preliminary Review label Aug 26, 2025

fix: update gemini grounding test to match new architecture

edd6f3d

- Fixed failing test to expect grounding chunk instead of text with citations - Added inline comments to regex patterns for source stripping in Task.ts - Added test coverage for grounding source handling to prevent regression

dosubot bot added size:L This PR changes 100-499 lines, ignoring generated files. and removed size:M This PR changes 30-99 lines, ignoring generated files. labels Aug 26, 2025

daniel-lxs moved this from PR [Needs Prelim Review] to PR [Needs Review] in Roo Code Roadmap Aug 26, 2025

daniel-lxs approved these changes Aug 26, 2025

View reviewed changes

dosubot bot added the lgtm This PR has been approved by a maintainer label Aug 26, 2025

hannesrudolph added PR - Needs Review and removed PR - Needs Preliminary Review labels Aug 26, 2025

mrubens reviewed Aug 27, 2025

View reviewed changes

fix: making changes from the code review

7e9a8c4

Merge branch 'RooCodeInc:main' into i/fix-gemini-sources-as-separate-msg

4636584

mrubens approved these changes Sep 5, 2025

View reviewed changes

Merge branch 'main' into i/fix-gemini-sources-as-separate-msg

d428028

mrubens merged commit c206da4 into RooCodeInc:main Sep 5, 2025
9 checks passed

github-project-automation bot moved this from New to Done in Roo Code Roadmap Sep 5, 2025

github-project-automation bot moved this from PR [Needs Review] to Done in Roo Code Roadmap Sep 5, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: Tackling Race/State condition issue by Changing the Code Design for Gemini Grounding Sources #7434

fix: Tackling Race/State condition issue by Changing the Code Design for Gemini Grounding Sources #7434

Uh oh!

HahaBill commented Aug 26, 2025 •

edited by ellipsis-dev bot

Loading

Uh oh!

ellipsis-dev bot Aug 26, 2025

Uh oh!

roomote bot left a comment

Uh oh!

roomote bot Aug 26, 2025

Uh oh!

roomote bot Aug 26, 2025

Uh oh!

roomote bot Aug 26, 2025

Uh oh!

roomote bot Aug 26, 2025

Uh oh!

roomote bot Aug 26, 2025

Uh oh!

daniel-lxs left a comment

Uh oh!

HahaBill commented Aug 26, 2025

Uh oh!

mrubens Aug 27, 2025

Uh oh!

HahaBill Aug 27, 2025

Uh oh!

mrubens Aug 27, 2025

Uh oh!

HahaBill Aug 27, 2025

Uh oh!

HahaBill commented Aug 27, 2025

Uh oh!

Uh oh!

Uh oh!

	.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "")
	\t\t\t\t\t\t.replace(/\[\d+\]\s+[^:\n]+:\s+https?:\/\/[^\s\n]+/g, "") // e.g. "[1] Example Source: https://example.com"

fix: Tackling Race/State condition issue by Changing the Code Design for Gemini Grounding Sources #7434

fix: Tackling Race/State condition issue by Changing the Code Design for Gemini Grounding Sources #7434

Uh oh!

Conversation

HahaBill commented Aug 26, 2025 • edited by ellipsis-dev bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Related GitHub Issue

Description

Summary:

Test Procedure

Pre-Submission Checklist

Uh oh!

ellipsis-dev bot Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

roomote bot left a comment

Choose a reason for hiding this comment

Uh oh!

roomote bot Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

roomote bot Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

roomote bot Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

roomote bot Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

roomote bot Aug 26, 2025

Choose a reason for hiding this comment

Uh oh!

daniel-lxs left a comment

Choose a reason for hiding this comment

Uh oh!

HahaBill commented Aug 26, 2025

Uh oh!

mrubens Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

HahaBill Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

mrubens Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

HahaBill Aug 27, 2025

Choose a reason for hiding this comment

Uh oh!

HahaBill commented Aug 27, 2025

Uh oh!

Uh oh!

Uh oh!

HahaBill commented Aug 26, 2025 •

edited by ellipsis-dev bot

Loading