From a4eeb6c4b101d17f08956c2e9961aab461edf4b3 Mon Sep 17 00:00:00 2001
From: Calvin <calvinvwei@gmail.com>
Date: Sat, 13 Dec 2025 22:57:37 +0800
Subject: [PATCH] fix: preserve thinking time when stopping reply

Fixes #11886

Signed-off-by: Calvin <calvinvwei@gmail.com>
---
 .../callbacks/baseCallbacks.ts                | 29 +++++++++++++++++--
 .../messageStreaming/callbacks/index.ts       | 15 +++++-----
 .../callbacks/thinkingCallbacks.ts            |  6 ++++
 3 files changed, 41 insertions(+), 9 deletions(-)

diff --git a/src/renderer/src/services/messageStreaming/callbacks/baseCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/baseCallbacks.ts
index ed9bdd5844..146266fd9d 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/baseCallbacks.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/baseCallbacks.ts
@@ -29,10 +29,20 @@ interface BaseCallbacksDependencies {
   assistantMsgId: string
   saveUpdatesToDB: any
   assistant: Assistant
+  getCurrentThinkingInfo?: () => { blockId: string | null; millsec: number }
 }
 
 export const createBaseCallbacks = (deps: BaseCallbacksDependencies) => {
-  const { blockManager, dispatch, getState, topicId, assistantMsgId, saveUpdatesToDB, assistant } = deps
+  const {
+    blockManager,
+    dispatch,
+    getState,
+    topicId,
+    assistantMsgId,
+    saveUpdatesToDB,
+    assistant,
+    getCurrentThinkingInfo
+  } = deps
 
   const startTime = Date.now()
   const notificationService = NotificationService.getInstance()
@@ -111,13 +121,28 @@ export const createBaseCallbacks = (deps: BaseCallbacksDependencies) => {
       if (currentMessage) {
         const allBlockRefs = findAllBlocks(currentMessage)
         const blockState = getState().messageBlocks
+        // 获取当前思考信息（如果有），用于保留实际思考时间
+        const thinkingInfo = getCurrentThinkingInfo?.()
         for (const blockRef of allBlockRefs) {
           const block = blockState.entities[blockRef.id]
           if (block && block.status === MessageBlockStatus.STREAMING && block.id !== possibleBlockId) {
+            // 构建更新对象
+            const changes: Record<string, any> = {
+              status: isErrorTypeAbort ? MessageBlockStatus.PAUSED : MessageBlockStatus.ERROR
+            }
+            // 如果是 thinking block 且有思考时间信息，保留实际思考时间
+            if (
+              block.type === MessageBlockType.THINKING &&
+              thinkingInfo?.blockId === block.id &&
+              thinkingInfo?.millsec &&
+              thinkingInfo.millsec > 0
+            ) {
+              changes.thinking_millsec = thinkingInfo.millsec
+            }
             dispatch(
               updateOneBlock({
                 id: block.id,
-                changes: { status: isErrorTypeAbort ? MessageBlockStatus.PAUSED : MessageBlockStatus.ERROR }
+                changes
               })
             )
           }
diff --git a/src/renderer/src/services/messageStreaming/callbacks/index.ts b/src/renderer/src/services/messageStreaming/callbacks/index.ts
index 2bb1d158bb..1587cb936c 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/index.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/index.ts
@@ -23,6 +23,12 @@ interface CallbacksDependencies {
 export const createCallbacks = (deps: CallbacksDependencies) => {
   const { blockManager, dispatch, getState, topicId, assistantMsgId, saveUpdatesToDB, assistant } = deps
 
+  // 首先创建 thinkingCallbacks ，以便传递 getCurrentThinkingInfo 给 baseCallbacks
+  const thinkingCallbacks = createThinkingCallbacks({
+    blockManager,
+    assistantMsgId
+  })
+
   // 创建基础回调
   const baseCallbacks = createBaseCallbacks({
     blockManager,
@@ -31,13 +37,8 @@ export const createCallbacks = (deps: CallbacksDependencies) => {
     topicId,
     assistantMsgId,
     saveUpdatesToDB,
-    assistant
-  })
-
-  // 创建各类回调
-  const thinkingCallbacks = createThinkingCallbacks({
-    blockManager,
-    assistantMsgId
+    assistant,
+    getCurrentThinkingInfo: thinkingCallbacks.getCurrentThinkingInfo
   })
 
   const toolCallbacks = createToolCallbacks({
diff --git a/src/renderer/src/services/messageStreaming/callbacks/thinkingCallbacks.ts b/src/renderer/src/services/messageStreaming/callbacks/thinkingCallbacks.ts
index aeb160fd02..3c718c4a67 100644
--- a/src/renderer/src/services/messageStreaming/callbacks/thinkingCallbacks.ts
+++ b/src/renderer/src/services/messageStreaming/callbacks/thinkingCallbacks.ts
@@ -19,6 +19,12 @@ export const createThinkingCallbacks = (deps: ThinkingCallbacksDependencies) =>
   let thinking_millsec_now: number = 0
 
   return {
+    // 获取当前思考时间（用于停止回复时保留思考时间）
+    getCurrentThinkingInfo: () => ({
+      blockId: thinkingBlockId,
+      millsec: thinking_millsec_now > 0 ? performance.now() - thinking_millsec_now : 0
+    }),
+
     onThinkingStart: async () => {
       if (blockManager.hasInitialPlaceholder) {
         const changes: Partial<MessageBlock> = {