From 46cbc3354c2669f5cdb2e3a348bf7e763bf4e7d6 Mon Sep 17 00:00:00 2001
From: "dayuan.jiang" <jdy.toh@gmail.com>
Date: Fri, 5 Dec 2025 00:26:02 +0900
Subject: [PATCH] fix: add manual token usage reporting to Langfuse for Bedrock
 streaming

Bedrock streaming responses don't auto-report token usage to OpenTelemetry.
This fix manually sets span attributes (ai.usage.promptTokens, gen_ai.usage.input_tokens)
from the AI SDK onFinish callback to ensure Langfuse captures token counts.
---
 app/api/chat/route.ts |  7 ++++++-
 lib/langfuse.ts       | 12 +++++++++++-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/app/api/chat/route.ts b/app/api/chat/route.ts
index 10befbe..03cd939 100644
--- a/app/api/chat/route.ts
+++ b/app/api/chat/route.ts
@@ -182,7 +182,12 @@ ${lastMessageText}
     onFinish: ({ text, usage, providerMetadata }) => {
       console.log('[Cache] Full providerMetadata:', JSON.stringify(providerMetadata, null, 2));
       console.log('[Cache] Usage:', JSON.stringify(usage, null, 2));
-      setTraceOutput(text);
+      // Pass usage to Langfuse (Bedrock streaming doesn't auto-report tokens to telemetry)
+      // AI SDK uses inputTokens/outputTokens, Langfuse expects promptTokens/completionTokens
+      setTraceOutput(text, {
+        promptTokens: usage?.inputTokens,
+        completionTokens: usage?.outputTokens,
+      });
     },
     tools: {
       // Client-side tool that will be executed on the client
diff --git a/lib/langfuse.ts b/lib/langfuse.ts
index 1b74cd3..58d9629 100644
--- a/lib/langfuse.ts
+++ b/lib/langfuse.ts
@@ -43,12 +43,22 @@ export function setTraceInput(params: {
 }
 
 // Update trace with output and end the span
-export function setTraceOutput(output: string) {
+export function setTraceOutput(output: string, usage?: { promptTokens?: number; completionTokens?: number }) {
   if (!isLangfuseEnabled()) return;
 
   updateActiveTrace({ output });
+
   const activeSpan = api.trace.getActiveSpan();
   if (activeSpan) {
+    // Manually set usage attributes since AI SDK Bedrock streaming doesn't provide them
+    if (usage?.promptTokens) {
+      activeSpan.setAttribute('ai.usage.promptTokens', usage.promptTokens);
+      activeSpan.setAttribute('gen_ai.usage.input_tokens', usage.promptTokens);
+    }
+    if (usage?.completionTokens) {
+      activeSpan.setAttribute('ai.usage.completionTokens', usage.completionTokens);
+      activeSpan.setAttribute('gen_ai.usage.output_tokens', usage.completionTokens);
+    }
     activeSpan.end();
   }
 }