From 96aba33077bfde753ff872217b2b05b49838c9d1 Mon Sep 17 00:00:00 2001
From: MyPrototypeWhat <daoquqiexing@gmail.com>
Date: Fri, 12 Dec 2025 18:23:49 +0800
Subject: [PATCH] fix: correct token calculation in prompt tool use plugin
 (#11867)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix: correct token calculation in prompt tool use plugin

- Fix duplicate token accumulation in recursive stream handling
- Accumulate usage for finish-step without tool calls
- Filter out recursive stream's start/finish events (only one per conversation)
- Make accumulateUsage method public for reuse

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

* refactor: simplify pipeRecursiveStream method in StreamEventManager

- Removed unnecessary context parameter from pipeRecursiveStream method
- Streamlined the invocation of pipeRecursiveStream in recursive call handling

This change enhances code clarity and reduces complexity in stream management.

---------

Co-authored-by: Claude Opus 4.5 <noreply@anthropic.com>
---
 .../toolUsePlugin/StreamEventManager.ts       | 24 +++++++------------
 .../toolUsePlugin/promptToolUsePlugin.ts      |  5 +++-
 2 files changed, 12 insertions(+), 17 deletions(-)

diff --git a/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/StreamEventManager.ts b/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/StreamEventManager.ts
index 555d4929d..c30c2015f 100644
--- a/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/StreamEventManager.ts
+++ b/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/StreamEventManager.ts
@@ -62,7 +62,7 @@ export class StreamEventManager {
     const recursiveResult = await context.recursiveCall(recursiveParams)
 
     if (recursiveResult && recursiveResult.fullStream) {
-      await this.pipeRecursiveStream(controller, recursiveResult.fullStream, context)
+      await this.pipeRecursiveStream(controller, recursiveResult.fullStream)
     } else {
       console.warn('[MCP Prompt] No fullstream found in recursive result:', recursiveResult)
     }
@@ -74,11 +74,7 @@ export class StreamEventManager {
   /**
    * 将递归流的数据传递到当前流
    */
-  private async pipeRecursiveStream(
-    controller: StreamController,
-    recursiveStream: ReadableStream,
-    context?: AiRequestContext
-  ): Promise<void> {
+  private async pipeRecursiveStream(controller: StreamController, recursiveStream: ReadableStream): Promise<void> {
     const reader = recursiveStream.getReader()
     try {
       while (true) {
@@ -86,18 +82,14 @@ export class StreamEventManager {
         if (done) {
           break
         }
+        if (value.type === 'start') {
+          continue
+        }
+
         if (value.type === 'finish') {
-          // 迭代的流不发finish，但需要累加其 usage
-          if (value.usage && context?.accumulatedUsage) {
-            this.accumulateUsage(context.accumulatedUsage, value.usage)
-          }
           break
         }
-        // 对于 finish-step 类型，累加其 usage
-        if (value.type === 'finish-step' && value.usage && context?.accumulatedUsage) {
-          this.accumulateUsage(context.accumulatedUsage, value.usage)
-        }
-        // 将递归流的数据传递到当前流
+
         controller.enqueue(value)
       }
     } finally {
@@ -159,7 +151,7 @@ export class StreamEventManager {
   /**
    * 累加 usage 数据
    */
-  private accumulateUsage(target: any, source: any): void {
+  accumulateUsage(target: any, source: any): void {
     if (!target || !source) return
 
     // 累加各种 token 类型
diff --git a/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/promptToolUsePlugin.ts b/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/promptToolUsePlugin.ts
index 274fdcee5..22e8b5a60 100644
--- a/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/promptToolUsePlugin.ts
+++ b/packages/aiCore/src/core/plugins/built-in/toolUsePlugin/promptToolUsePlugin.ts
@@ -411,7 +411,10 @@ export const createPromptToolUsePlugin = (config: PromptToolUseConfig = {}) => {
               }
             }
 
-            // 如果没有执行工具调用，直接传递原始finish-step事件
+            // 如果没有执行工具调用，累加 usage 后透传 finish-step 事件
+            if (chunk.usage && context.accumulatedUsage) {
+              streamEventManager.accumulateUsage(context.accumulatedUsage, chunk.usage)
+            }
             controller.enqueue(chunk)
 
             // 清理状态