From c666361611edb92626c601de30a442622cc9cdc8 Mon Sep 17 00:00:00 2001
From: Phantom <59059173+EurFelux@users.noreply.github.com>
Date: Mon, 11 Aug 2025 17:05:05 +0800
Subject: [PATCH] fix: trace usage (#9018)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

* fix(openai): 移除不必要的类型断言并更新类型定义

更新OpenAIApiClient中的usage处理逻辑，移除不必要的类型断言
在sdk.ts中更新OpenAISdkRawChunk类型定义，明确包含可能的cost字段

* fix(openai): 修复流式响应中完成信号触发逻辑

调整完成信号的触发条件，不再区分 OpenRouter 和其他提供商
统一等待 usage 信息后再触发完成信号，以统一适配 OpenAI Chat Completions API

* fix(sdk类型): 将OpenAISdkRawChunk中的usage字段改为可选

* refactor(openai): 移除对OpenRouter的特殊处理逻辑
---
 .../aiCore/clients/openai/OpenAIApiClient.ts  | 47 ++++++-------------
 src/renderer/src/types/sdk.ts                 |  6 ++-
 2 files changed, 18 insertions(+), 35 deletions(-)

diff --git a/src/renderer/src/aiCore/clients/openai/OpenAIApiClient.ts b/src/renderer/src/aiCore/clients/openai/OpenAIApiClient.ts
index 183fcbb7cf..60173551b4 100644
--- a/src/renderer/src/aiCore/clients/openai/OpenAIApiClient.ts
+++ b/src/renderer/src/aiCore/clients/openai/OpenAIApiClient.ts
@@ -758,12 +758,10 @@ export class OpenAIAPIClient extends OpenAIBaseClient<
     let accumulatingText = false
     return (context: ResponseChunkTransformerContext) => ({
       async transform(chunk: OpenAISdkRawChunk, controller: TransformStreamDefaultController<GenericChunk>) {
-        const isOpenRouter = context.provider?.id === 'openrouter'
-
         // 持续更新usage信息
         logger.silly('chunk', chunk)
         if (chunk.usage) {
-          const usage = chunk.usage as any // OpenRouter may include additional fields like cost
+          const usage = chunk.usage
           lastUsageInfo = {
             prompt_tokens: usage.prompt_tokens || 0,
             completion_tokens: usage.completion_tokens || 0,
@@ -771,19 +769,11 @@ export class OpenAIAPIClient extends OpenAIBaseClient<
             // Handle OpenRouter specific cost fields
             ...(usage.cost !== undefined ? { cost: usage.cost } : {})
           }
-
-          // For OpenRouter, if we've seen finish_reason and now have usage, emit completion signals
-          if (isOpenRouter && hasFinishReason && !isFinished) {
-            emitCompletionSignals(controller)
-            return
-          }
         }
 
-        // For OpenRouter, if this chunk only contains usage without choices, emit completion signals
-        if (isOpenRouter && chunk.usage && (!chunk.choices || chunk.choices.length === 0)) {
-          if (!isFinished) {
-            emitCompletionSignals(controller)
-          }
+        // if we've already seen finish_reason, emit completion signals. No matter whether we get usage or not.
+        if (hasFinishReason && !isFinished) {
+          emitCompletionSignals(controller)
           return
         }
 
@@ -832,16 +822,12 @@ export class OpenAIAPIClient extends OpenAIBaseClient<
 
             if (!contentSource) {
               if ('finish_reason' in choice && choice.finish_reason) {
-                // For OpenRouter, don't emit completion signals immediately after finish_reason
-                // Wait for the usage chunk that comes after
-                if (isOpenRouter) {
-                  hasFinishReason = true
-                  // If we already have usage info, emit completion signals now
-                  if (lastUsageInfo && lastUsageInfo.total_tokens > 0) {
-                    emitCompletionSignals(controller)
-                  }
-                } else {
-                  // For other providers, emit completion signals immediately
+                // OpenAI Chat Completions API 在启用 stream_options: { include_usage: true } 以后
+                // 包含 usage 的 chunk 会在包含 finish_reason: stop 的 chunk 之后
+                // 所以试图等到拿到 usage 之后再发出结束信号
+                hasFinishReason = true
+                // If we already have usage info, emit completion signals now
+                if (lastUsageInfo && lastUsageInfo.total_tokens > 0) {
                   emitCompletionSignals(controller)
                 }
               }
@@ -947,16 +933,11 @@ export class OpenAIAPIClient extends OpenAIBaseClient<
                 })
               }
 
-              // For OpenRouter, don't emit completion signals immediately after finish_reason
+              // Don't emit completion signals immediately after finish_reason
               // Wait for the usage chunk that comes after
-              if (isOpenRouter) {
-                hasFinishReason = true
-                // If we already have usage info, emit completion signals now
-                if (lastUsageInfo && lastUsageInfo.total_tokens > 0) {
-                  emitCompletionSignals(controller)
-                }
-              } else {
-                // For other providers, emit completion signals immediately
+              hasFinishReason = true
+              // If we already have usage info, emit completion signals now
+              if (lastUsageInfo && lastUsageInfo.total_tokens > 0) {
                 emitCompletionSignals(controller)
               }
             }
diff --git a/src/renderer/src/types/sdk.ts b/src/renderer/src/types/sdk.ts
index 12661c818a..a5413e54fb 100644
--- a/src/renderer/src/types/sdk.ts
+++ b/src/renderer/src/types/sdk.ts
@@ -86,11 +86,13 @@ export type ReasoningEffortOptionalParams = {
 }
 
 export type OpenAISdkParams = OpenAIParamsWithoutReasoningEffort & ReasoningEffortOptionalParams
+
+// OpenRouter may include additional fields like cost
 export type OpenAISdkRawChunk =
-  | OpenAI.Chat.Completions.ChatCompletionChunk
+  | (OpenAI.Chat.Completions.ChatCompletionChunk & { usage?: OpenAI.CompletionUsage & { cost?: number } })
   | ({
       _request_id?: string | null | undefined
-    } & OpenAI.ChatCompletion)
+    } & OpenAI.ChatCompletion & { usage?: OpenAI.CompletionUsage & { cost?: number } })
 
 export type OpenAISdkRawOutput = Stream<OpenAI.Chat.Completions.ChatCompletionChunk> | OpenAI.ChatCompletion
 export type OpenAISdkRawContentSource =