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

Co-authored-by: daniel-lxs <[email protected]>
Co-authored-by: Matt Rubens <[email protected]>

Author: 3 people

src/api/providers/__tests__/gemini-handler.spec.ts

@@ -85,7 +85,7 @@ describe("GeminiHandler backend support", () => {
 							groundingMetadata: {
 								groundingChunks: [
 									{ web: null }, // Missing URI
-									{ web: { uri: "https://example.com" } }, // Valid
+									{ web: { uri: "https://example.com", title: "Example Site" } }, // Valid
 									{}, // Missing web property entirely
 								],
 							},
@@ -105,13 +105,20 @@ describe("GeminiHandler backend support", () => {
 				messages.push(chunk)
 			}
 
-			// Should only include valid citations
-			const sourceMessage = messages.find((m) => m.type === "text" && m.text?.includes("[2]"))
-			expect(sourceMessage).toBeDefined()
-			if (sourceMessage && "text" in sourceMessage) {
-				expect(sourceMessage.text).toContain("https://example.com")
-				expect(sourceMessage.text).not.toContain("[1]")
-				expect(sourceMessage.text).not.toContain("[3]")
+			// Should have the text response
+			const textMessage = messages.find((m) => m.type === "text")
+			expect(textMessage).toBeDefined()
+			if (textMessage && "text" in textMessage) {
+				expect(textMessage.text).toBe("test response")
+			}
+
+			// Should have grounding chunk with only valid sources
+			const groundingMessage = messages.find((m) => m.type === "grounding")
+			expect(groundingMessage).toBeDefined()
+			if (groundingMessage && "sources" in groundingMessage) {
+				expect(groundingMessage.sources).toHaveLength(1)
+				expect(groundingMessage.sources[0].url).toBe("https://example.com")
+				expect(groundingMessage.sources[0].title).toBe("Example Site")
 			}
 		})
 

src/api/providers/gemini.ts

@@ -15,7 +15,7 @@ import { safeJsonParse } from "../../shared/safeJsonParse"
 
 import { convertAnthropicContentToGemini, convertAnthropicMessageToGemini } from "../transform/gemini-format"
 import { t } from "i18next"
-import type { ApiStream } from "../transform/stream"
+import type { ApiStream, GroundingSource } from "../transform/stream"
 import { getModelParams } from "../transform/model-params"
 
 import type { SingleCompletionHandler, ApiHandlerCreateMessageMetadata } from "../index"
@@ -132,9 +132,9 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 			}
 
 			if (pendingGroundingMetadata) {
-				const citations = this.extractCitationsOnly(pendingGroundingMetadata)
-				if (citations) {
-					yield { type: "text", text: `\n\n${t("common:errors.gemini.sources")} ${citations}` }
+				const sources = this.extractGroundingSources(pendingGroundingMetadata)
+				if (sources.length > 0) {
+					yield { type: "grounding", sources }
 				}
 			}
 
@@ -175,28 +175,38 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 		return { id: id.endsWith(":thinking") ? id.replace(":thinking", "") : id, info, ...params }
 	}
 
-	private extractCitationsOnly(groundingMetadata?: GroundingMetadata): string | null {
+	private extractGroundingSources(groundingMetadata?: GroundingMetadata): GroundingSource[] {
 		const chunks = groundingMetadata?.groundingChunks
 
 		if (!chunks) {
-			return null
+			return []
 		}
 
-		const citationLinks = chunks
-			.map((chunk, i) => {
+		return chunks
+			.map((chunk): GroundingSource | null => {
 				const uri = chunk.web?.uri
+				const title = chunk.web?.title || uri || "Unknown Source"
+
 				if (uri) {
-					return `[${i + 1}](${uri})`
+					return {
+						title,
+						url: uri,
+					}
 				}
 				return null
 			})
-			.filter((link): link is string => link !== null)
+			.filter((source): source is GroundingSource => source !== null)
+	}
+
+	private extractCitationsOnly(groundingMetadata?: GroundingMetadata): string | null {
+		const sources = this.extractGroundingSources(groundingMetadata)
 
-		if (citationLinks.length > 0) {
-			return citationLinks.join(", ")
+		if (sources.length === 0) {
+			return null
 		}
 
-		return null
+		const citationLinks = sources.map((source, i) => `[${i + 1}](${source.url})`)
+		return citationLinks.join(", ")
 	}
 
 	async completePrompt(prompt: string): Promise<string> {

src/api/transform/stream.ts

@@ -1,6 +1,11 @@
 export type ApiStream = AsyncGenerator<ApiStreamChunk>
 
-export type ApiStreamChunk = ApiStreamTextChunk | ApiStreamUsageChunk | ApiStreamReasoningChunk | ApiStreamError
+export type ApiStreamChunk =
+	| ApiStreamTextChunk
+	| ApiStreamUsageChunk
+	| ApiStreamReasoningChunk
+	| ApiStreamGroundingChunk
+	| ApiStreamError
 
 export interface ApiStreamError {
 	type: "error"
@@ -27,3 +32,14 @@ export interface ApiStreamUsageChunk {
 	reasoningTokens?: number
 	totalCost?: number
 }
+
+export interface ApiStreamGroundingChunk {
+	type: "grounding"
+	sources: GroundingSource[]
+}
+
+export interface GroundingSource {
+	title: string
+	url: string
+	snippet?: string
+}

src/core/task/Task.ts

@@ -41,7 +41,7 @@ import { CloudService, BridgeOrchestrator } from "@roo-code/cloud"
 
 // api
 import { ApiHandler, ApiHandlerCreateMessageMetadata, buildApiHandler } from "../../api"
-import { ApiStream } from "../../api/transform/stream"
+import { ApiStream, GroundingSource } from "../../api/transform/stream"
 import { maybeRemoveImageBlocks } from "../../api/transform/image-cleaning"
 
 // shared
@@ -1897,7 +1897,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 					this.didFinishAbortingStream = true
 				}
 
-				// Reset streaming state.
+				// Reset streaming state for each new API request
 				this.currentStreamingContentIndex = 0
 				this.currentStreamingDidCheckpoint = false
 				this.assistantMessageContent = []
@@ -1918,6 +1918,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 				const stream = this.attemptApiRequest()
 				let assistantMessage = ""
 				let reasoningMessage = ""
+				let pendingGroundingSources: GroundingSource[] = []
 				this.isStreaming = true
 
 				try {
@@ -1944,6 +1945,13 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 								cacheReadTokens += chunk.cacheReadTokens ?? 0
 								totalCost = chunk.totalCost
 								break
+							case "grounding":
+								// Handle grounding sources separately from regular content
+								// to prevent state persistence issues - store them separately
+								if (chunk.sources && chunk.sources.length > 0) {
+									pendingGroundingSources.push(...chunk.sources)
+								}
+								break
 							case "text": {
 								assistantMessage += chunk.text
 
@@ -2237,6 +2245,16 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 				let didEndLoop = false
 
 				if (assistantMessage.length > 0) {
+					// 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 = `${t("common:gemini.sources")} ${citationLinks.join(", ")}`
+
+						await this.say("text", sourcesText, undefined, false, undefined, undefined, {
+							isNonInteractive: true,
+						})
+					}
+
 					await this.addToApiConversationHistory({
 						role: "assistant",
 						content: [{ type: "text", text: assistantMessage }],