From 2c980cd43ce538b0d8940f3eb55977bbfb90dac9 Mon Sep 17 00:00:00 2001
From: root <root@mjfio6a0.vm>
Date: Fri, 30 Jan 2026 14:24:12 +0800
Subject: [PATCH] feat: include token usage data in chat final events (#4159)

- Add usage parameter to broadcastChatFinal() function
- Add usage parameter to emitChatFinal() function
- Extract usage from message object and agent event data
- Enable external monitoring tools (e.g., Crabwalk) to track API costs
- Maintain full backward compatibility (usage is optional in schema)

This change allows tools like Crabwalk to display token costs and usage
metrics without additional API calls. The usage field is optional and
fully backward compatible with existing clients.
---
 src/gateway/server-chat.ts         | 4 ++++
 src/gateway/server-methods/chat.ts | 3 +++
 2 files changed, 7 insertions(+)

diff --git a/src/gateway/server-chat.ts b/src/gateway/server-chat.ts
index 8c67767a6..85aa3b3ff 100644
--- a/src/gateway/server-chat.ts
+++ b/src/gateway/server-chat.ts
@@ -163,6 +163,7 @@ export function createAgentEventHandler({
     seq: number,
     jobState: "done" | "error",
     error?: unknown,
+    usage?: Record<string, unknown>,
   ) => {
     const text = chatRunState.buffers.get(clientRunId)?.trim() ?? "";
     chatRunState.buffers.delete(clientRunId);
@@ -180,6 +181,7 @@ export function createAgentEventHandler({
               timestamp: Date.now(),
             }
           : undefined,
+        usage: usage,
       };
       // Suppress webchat broadcast for heartbeat runs when showOk is false
       if (!shouldSuppressHeartbeatBroadcast(clientRunId)) {
@@ -264,6 +266,7 @@ export function createAgentEventHandler({
             evt.seq,
             lifecyclePhase === "error" ? "error" : "done",
             evt.data?.error,
+            evt.data?.usage as Record<string, unknown> | undefined,
           );
         } else {
           emitChatFinal(
@@ -272,6 +275,7 @@ export function createAgentEventHandler({
             evt.seq,
             lifecyclePhase === "error" ? "error" : "done",
             evt.data?.error,
+            evt.data?.usage as Record<string, unknown> | undefined,
           );
         }
       } else if (isAborted && (lifecyclePhase === "end" || lifecyclePhase === "error")) {
diff --git a/src/gateway/server-methods/chat.ts b/src/gateway/server-methods/chat.ts
index 9010a6f21..8c7779e2e 100644
--- a/src/gateway/server-methods/chat.ts
+++ b/src/gateway/server-methods/chat.ts
@@ -149,6 +149,7 @@ function broadcastChatFinal(params: {
   runId: string;
   sessionKey: string;
   message?: Record<string, unknown>;
+  usage?: Record<string, unknown>;
 }) {
   const seq = nextChatSeq({ agentRunSeq: params.context.agentRunSeq }, params.runId);
   const payload = {
@@ -157,6 +158,7 @@ function broadcastChatFinal(params: {
     seq,
     state: "final" as const,
     message: params.message,
+    usage: params.usage,
   };
   params.context.broadcast("chat", payload);
   params.context.nodeSendToSession(params.sessionKey, "chat", payload);
@@ -545,6 +547,7 @@ export const chatHandlers: GatewayRequestHandlers = {
               runId: clientRunId,
               sessionKey: p.sessionKey,
               message,
+              usage: message?.usage as Record<string, unknown> | undefined,
             });
           }
           context.dedupe.set(`chat:${clientRunId}`, {