From 3630133efd1badf7a5d57fe5764c590a8a4659b0 Mon Sep 17 00:00:00 2001
From: MyPrototypeWhat <daoquqiexing@gmail.com>
Date: Thu, 14 Aug 2025 16:27:23 +0800
Subject: [PATCH] fix: update MessageKnowledgeSearch to use knowledgeReferences

- Modified MessageKnowledgeSearch component to display additional context from toolInput.
- Updated the fetch complete message to reflect the count of knowledgeReferences instead of toolOutput.
- Adjusted the mapping of results to iterate over knowledgeReferences for rendering.
---
 .../home/Messages/Tools/MessageKnowledgeSearch.tsx   |  6 +++---
 .../messageStreaming/callbacks/toolCallbacks.ts      | 12 ++++++++++++
 src/renderer/src/utils/messageUtils/create.ts        |  2 +-
 3 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/src/renderer/src/pages/home/Messages/Tools/MessageKnowledgeSearch.tsx b/src/renderer/src/pages/home/Messages/Tools/MessageKnowledgeSearch.tsx
index d44c59e895..a7ebeb648c 100644
--- a/src/renderer/src/pages/home/Messages/Tools/MessageKnowledgeSearch.tsx
+++ b/src/renderer/src/pages/home/Messages/Tools/MessageKnowledgeSearch.tsx
@@ -16,14 +16,14 @@ export function MessageKnowledgeSearchToolTitle({ toolResponse }: { toolResponse
       text={
         <PrepareToolWrapper>
           {i18n.t('message.searching')}
-          <span>{toolInput?.rewrite ?? toolInput?.query ?? ''}</span>
+          <span>{toolInput?.additionalContext ?? ''}</span>
         </PrepareToolWrapper>
       }
     />
   ) : (
     <MessageWebSearchToolTitleTextWrapper type="secondary">
       <FileSearch size={16} style={{ color: 'unset' }} />
-      {i18n.t('message.websearch.fetch_complete', { count: toolOutput.length ?? 0 })}
+      {i18n.t('message.websearch.fetch_complete', { count: toolOutput.knowledgeReferences.length ?? 0 })}
     </MessageWebSearchToolTitleTextWrapper>
   )
 }
@@ -33,7 +33,7 @@ export function MessageKnowledgeSearchToolBody({ toolResponse }: { toolResponse:
 
   return toolResponse.status === 'done' ? (
     <MessageWebSearchToolBodyUlWrapper>
-      {toolOutput.map((result) => (
+      {toolOutput.knowledgeReferences.map((result) => (
         <li key={result.id}>
           <span>{result.id}</span>
           <span>{result.content}</span>
diff --git a/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts
index ec5a288595..3cf1b186d9 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/toolCallbacks.ts
@@ -114,6 +114,18 @@ export const createToolCallbacks = (deps: ToolCallbacksDependencies) => {
           citationBlockId = citationBlock.id
           blockManager.handleBlockTransition(citationBlock, MessageBlockType.CITATION)
         }
+        if (toolResponse.tool.name === 'builtin_knowledge_search' && toolResponse.response?.knowledgeReferences) {
+          const citationBlock = createCitationBlock(
+            assistantMsgId,
+            { knowledge: toolResponse.response.knowledgeReferences },
+            {
+              status: MessageBlockStatus.SUCCESS
+            }
+          )
+          citationBlockId = citationBlock.id
+          blockManager.handleBlockTransition(citationBlock, MessageBlockType.CITATION)
+        }
+        // TODO: 处理 memory 引用
       } else {
         logger.warn(
           `[onToolCallComplete] Received unhandled tool status: ${toolResponse.status} for ID: ${toolResponse.id}`
diff --git a/src/renderer/src/utils/messageUtils/create.ts b/src/renderer/src/utils/messageUtils/create.ts
index e0081c8c6b..ed1c7e8efc 100644
--- a/src/renderer/src/utils/messageUtils/create.ts
+++ b/src/renderer/src/utils/messageUtils/create.ts
@@ -135,7 +135,7 @@ export function createThinkingBlock(
   return {
     ...baseBlock,
     content,
-    thinking_millsec: overrides.thinking_millsec
+    thinking_millsec: overrides.thinking_millsec || 0
   }
 }