From f8120c2ebbdc121b860e6147631fa8dd7710be79 Mon Sep 17 00:00:00 2001
From: one <wangan.cs@gmail.com>
Date: Wed, 20 Aug 2025 13:25:22 +0800
Subject: [PATCH] fix: web search references missing caused by early reset
 (#9328)

---
 .../messageStreaming/callbacks/citationCallbacks.ts         | 6 ++++--
 .../src/services/messageStreaming/callbacks/index.ts        | 3 ++-
 .../services/messageStreaming/callbacks/textCallbacks.ts    | 4 +++-
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/src/renderer/src/services/messageStreaming/callbacks/citationCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/citationCallbacks.ts
index 9ba743b2cd..03d0f23d48 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/citationCallbacks.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/citationCallbacks.ts
@@ -40,7 +40,6 @@ export const createCitationCallbacks = (deps: CitationCallbacksDependencies) =>
           status: MessageBlockStatus.SUCCESS
         }
         blockManager.smartBlockUpdate(citationBlockId, changes, MessageBlockType.CITATION, true)
-        citationBlockId = null
       } else {
         logger.error('[onExternalToolComplete] citationBlockId is null. Cannot update.')
       }
@@ -121,6 +120,9 @@ export const createCitationCallbacks = (deps: CitationCallbacksDependencies) =>
     },
 
     // 暴露给外部的方法，用于textCallbacks中获取citationBlockId
-    getCitationBlockId: () => citationBlockId
+    getCitationBlockId: () => citationBlockId,
+    setCitationBlockId: (id: string | null) => {
+      citationBlockId = id
+    }
   }
 }
diff --git a/src/renderer/src/services/messageStreaming/callbacks/index.ts b/src/renderer/src/services/messageStreaming/callbacks/index.ts
index 2b6fc5968a..ada8ed7b67 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/index.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/index.ts
@@ -59,7 +59,8 @@ export const createCallbacks = (deps: CallbacksDependencies) => {
     blockManager,
     getState,
     assistantMsgId,
-    getCitationBlockId: citationCallbacks.getCitationBlockId
+    getCitationBlockId: citationCallbacks.getCitationBlockId,
+    setCitationBlockId: citationCallbacks.setCitationBlockId
   })
 
   // 组合所有回调
diff --git a/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts
index 657cf6f0f3..53621b625f 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/textCallbacks.ts
@@ -12,10 +12,11 @@ interface TextCallbacksDependencies {
   getState: any
   assistantMsgId: string
   getCitationBlockId: () => string | null
+  setCitationBlockId: (id: string | null) => void
 }
 
 export const createTextCallbacks = (deps: TextCallbacksDependencies) => {
-  const { blockManager, getState, assistantMsgId, getCitationBlockId } = deps
+  const { blockManager, getState, assistantMsgId, getCitationBlockId, setCitationBlockId } = deps
 
   // 内部维护的状态
   let mainTextBlockId: string | null = null
@@ -62,6 +63,7 @@ export const createTextCallbacks = (deps: TextCallbacksDependencies) => {
         }
         blockManager.smartBlockUpdate(mainTextBlockId, changes, MessageBlockType.MAIN_TEXT, true)
         mainTextBlockId = null
+        setCitationBlockId(null)
       } else {
         logger.warn(
           `[onTextComplete] Received text.complete but last block was not MAIN_TEXT (was ${blockManager.lastBlockType}) or lastBlockId is null.`