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, ) => { 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 | undefined, ); } else { emitChatFinal( @@ -272,6 +275,7 @@ export function createAgentEventHandler({ evt.seq, lifecyclePhase === "error" ? "error" : "done", evt.data?.error, + evt.data?.usage as Record | 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; + usage?: Record; }) { 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 | undefined, }); } context.dedupe.set(`chat:${clientRunId}`, {