From 63c7b25f0d256f1eef170406b2fe4b9e9d29541f Mon Sep 17 00:00:00 2001
From: "Ton Hoang Nguyen (Bill)" <32552798+HahaBill@users.noreply.github.com>
Date: Tue, 8 Jul 2025 16:17:12 +0100
Subject: [PATCH] feat: adding citation for streams (generateContextStream)

---
 src/api/providers/gemini.ts | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

diff --git a/src/api/providers/gemini.ts b/src/api/providers/gemini.ts
index e9687c15cc..49d4c0e653 100644
--- a/src/api/providers/gemini.ts
+++ b/src/api/providers/gemini.ts
@@ -93,11 +93,19 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 			const result = await this.client.models.generateContentStream(params)
 
 			let lastUsageMetadata: GenerateContentResponseUsageMetadata | undefined
+			let accumulatedText = ""
+			let pendingGroundingMetadata: GroundingMetadata | undefined
+			let hasGroundingEnabled = this.options.enableGrounding
 
 			for await (const chunk of result) {
 				// Process candidates and their parts to separate thoughts from content
 				if (chunk.candidates && chunk.candidates.length > 0) {
 					const candidate = chunk.candidates[0]
+
+					if (candidate.groundingMetadata) {
+						pendingGroundingMetadata = candidate.groundingMetadata
+					}
+
 					if (candidate.content && candidate.content.parts) {
 						for (const part of candidate.content.parts) {
 							if (part.thought) {
@@ -108,7 +116,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 							} else {
 								// This is regular content
 								if (part.text) {
-									yield { type: "text", text: part.text }
+									accumulatedText += part.text
+
+									if (!hasGroundingEnabled) {
+										yield { type: "text", text: part.text }
+									}
 								}
 							}
 						}
@@ -117,7 +129,11 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 
 				// Fallback to the original text property if no candidates structure
 				else if (chunk.text) {
-					yield { type: "text", text: chunk.text }
+					accumulatedText += chunk.text
+
+					if (!hasGroundingEnabled) {
+						yield { type: "text", text: chunk.text }
+					}
 				}
 
 				if (chunk.usageMetadata) {
@@ -125,6 +141,16 @@ export class GeminiHandler extends BaseProvider implements SingleCompletionHandl
 				}
 			}
 
+			if (hasGroundingEnabled && accumulatedText) {
+				let finalText = accumulatedText
+
+				if (pendingGroundingMetadata) {
+					finalText = this.processGroundingCitations(accumulatedText, pendingGroundingMetadata)
+				}
+
+				yield { type: "text", text: finalText }
+			}
+
 			if (lastUsageMetadata) {
 				const inputTokens = lastUsageMetadata.promptTokenCount ?? 0
 				const outputTokens = lastUsageMetadata.candidatesTokenCount ?? 0