From 444c13e1e31fdedadcbab6b56851ec425865bc74 Mon Sep 17 00:00:00 2001
From: Phantom <eurfelux@gmail.com>
Date: Sun, 30 Nov 2025 15:35:23 +0800
Subject: [PATCH] fix: correct trace token usage (#11575)

fix: correct ai sdk token usage mapping
---
 .../src/aiCore/trace/AiSdkSpanAdapter.ts      |  4 +-
 .../trace/__tests__/AiSdkSpanAdapter.test.ts  | 53 +++++++++++++++++++
 2 files changed, 55 insertions(+), 2 deletions(-)
 create mode 100644 src/renderer/src/aiCore/trace/__tests__/AiSdkSpanAdapter.test.ts

diff --git a/src/renderer/src/aiCore/trace/AiSdkSpanAdapter.ts b/src/renderer/src/aiCore/trace/AiSdkSpanAdapter.ts
index f3df504de8..0c0e08a03d 100644
--- a/src/renderer/src/aiCore/trace/AiSdkSpanAdapter.ts
+++ b/src/renderer/src/aiCore/trace/AiSdkSpanAdapter.ts
@@ -245,8 +245,8 @@ export class AiSdkSpanAdapter {
       'gen_ai.usage.output_tokens'
     ]
 
-    const completionTokens = attributes[inputsTokenKeys.find((key) => attributes[key]) || '']
-    const promptTokens = attributes[outputTokenKeys.find((key) => attributes[key]) || '']
+    const promptTokens = attributes[inputsTokenKeys.find((key) => attributes[key]) || '']
+    const completionTokens = attributes[outputTokenKeys.find((key) => attributes[key]) || '']
 
     if (completionTokens !== undefined || promptTokens !== undefined) {
       const usage: TokenUsage = {
diff --git a/src/renderer/src/aiCore/trace/__tests__/AiSdkSpanAdapter.test.ts b/src/renderer/src/aiCore/trace/__tests__/AiSdkSpanAdapter.test.ts
new file mode 100644
index 0000000000..4cd6241e64
--- /dev/null
+++ b/src/renderer/src/aiCore/trace/__tests__/AiSdkSpanAdapter.test.ts
@@ -0,0 +1,53 @@
+import type { Span } from '@opentelemetry/api'
+import { SpanKind, SpanStatusCode } from '@opentelemetry/api'
+import { describe, expect, it, vi } from 'vitest'
+
+import { AiSdkSpanAdapter } from '../AiSdkSpanAdapter'
+
+vi.mock('@logger', () => ({
+  loggerService: {
+    withContext: () => ({
+      debug: vi.fn(),
+      error: vi.fn(),
+      info: vi.fn(),
+      warn: vi.fn()
+    })
+  }
+}))
+
+describe('AiSdkSpanAdapter', () => {
+  const createMockSpan = (attributes: Record<string, unknown>): Span => {
+    const span = {
+      spanContext: () => ({
+        traceId: 'trace-id',
+        spanId: 'span-id'
+      }),
+      _attributes: attributes,
+      _events: [],
+      name: 'test span',
+      status: { code: SpanStatusCode.OK },
+      kind: SpanKind.CLIENT,
+      startTime: [0, 0] as [number, number],
+      endTime: [0, 1] as [number, number],
+      ended: true,
+      parentSpanId: '',
+      links: []
+    }
+    return span as unknown as Span
+  }
+
+  it('maps prompt and completion usage tokens to the correct fields', () => {
+    const attributes = {
+      'ai.usage.promptTokens': 321,
+      'ai.usage.completionTokens': 654
+    }
+
+    const span = createMockSpan(attributes)
+    const result = AiSdkSpanAdapter.convertToSpanEntity({ span })
+
+    expect(result.usage).toBeDefined()
+    expect(result.usage?.prompt_tokens).toBe(321)
+    expect(result.usage?.completion_tokens).toBe(654)
+    expect(result.usage?.total_tokens).toBe(975)
+  })
+})