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

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

Author: HahaBill

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

@@ -39,7 +39,7 @@ import { CloudService, ExtensionBridgeService } from "@roo-code/cloud"
 
 // api
 import { ApiHandler, ApiHandlerCreateMessageMetadata, buildApiHandler } from "../../api"
-import { ApiStream } from "../../api/transform/stream"
+import { ApiStream, GroundingSource } from "../../api/transform/stream"
 
 // shared
 import { findLastIndex } from "../../shared/array"
@@ -1746,7 +1746,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 = []
@@ -1767,6 +1767,7 @@ export class Task extends EventEmitter<TaskEvents> implements TaskLike {
 				const stream = this.attemptApiRequest()
 				let assistantMessage = ""
 				let reasoningMessage = ""
+				let pendingGroundingSources: GroundingSource[] = []
 				this.isStreaming = true
 
 				try {
@@ -1793,6 +1794,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
 
@@ -2086,9 +2094,30 @@ 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 = `Sources: ${citationLinks.join(", ")}`
+
+						await this.say("text", sourcesText, undefined, false, undefined, undefined, {
+							isNonInteractive: true,
+						})
+					}
+
+					// Strip grounding sources from assistant message before persisting to API history
+					// This prevents state persistence issues while maintaining user experience
+					let cleanAssistantMessage = assistantMessage
+					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, "")
+							.replace(/Sources?:\s*[\s\S]*?(?=\n\n|\n$|$)/g, "")
+							.trim()
+					}
+
 					await this.addToApiConversationHistory({
 						role: "assistant",
-						content: [{ type: "text", text: assistantMessage }],
+						content: [{ type: "text", text: cleanAssistantMessage }],
 					})
 
 					TelemetryService.instance.captureConversationMessage(this.taskId, "assistant")