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.
This commit is contained in:
root 2026-01-30 14:24:12 +08:00
parent 87267fad4f
commit 2c980cd43c
2 changed files with 7 additions and 0 deletions

View File

@ -163,6 +163,7 @@ export function createAgentEventHandler({
seq: number, seq: number,
jobState: "done" | "error", jobState: "done" | "error",
error?: unknown, error?: unknown,
usage?: Record<string, unknown>,
) => { ) => {
const text = chatRunState.buffers.get(clientRunId)?.trim() ?? ""; const text = chatRunState.buffers.get(clientRunId)?.trim() ?? "";
chatRunState.buffers.delete(clientRunId); chatRunState.buffers.delete(clientRunId);
@ -180,6 +181,7 @@ export function createAgentEventHandler({
timestamp: Date.now(), timestamp: Date.now(),
} }
: undefined, : undefined,
usage: usage,
}; };
// Suppress webchat broadcast for heartbeat runs when showOk is false // Suppress webchat broadcast for heartbeat runs when showOk is false
if (!shouldSuppressHeartbeatBroadcast(clientRunId)) { if (!shouldSuppressHeartbeatBroadcast(clientRunId)) {
@ -264,6 +266,7 @@ export function createAgentEventHandler({
evt.seq, evt.seq,
lifecyclePhase === "error" ? "error" : "done", lifecyclePhase === "error" ? "error" : "done",
evt.data?.error, evt.data?.error,
evt.data?.usage as Record<string, unknown> | undefined,
); );
} else { } else {
emitChatFinal( emitChatFinal(
@ -272,6 +275,7 @@ export function createAgentEventHandler({
evt.seq, evt.seq,
lifecyclePhase === "error" ? "error" : "done", lifecyclePhase === "error" ? "error" : "done",
evt.data?.error, evt.data?.error,
evt.data?.usage as Record<string, unknown> | undefined,
); );
} }
} else if (isAborted && (lifecyclePhase === "end" || lifecyclePhase === "error")) { } else if (isAborted && (lifecyclePhase === "end" || lifecyclePhase === "error")) {

View File

@ -149,6 +149,7 @@ function broadcastChatFinal(params: {
runId: string; runId: string;
sessionKey: string; sessionKey: string;
message?: Record<string, unknown>; message?: Record<string, unknown>;
usage?: Record<string, unknown>;
}) { }) {
const seq = nextChatSeq({ agentRunSeq: params.context.agentRunSeq }, params.runId); const seq = nextChatSeq({ agentRunSeq: params.context.agentRunSeq }, params.runId);
const payload = { const payload = {
@ -157,6 +158,7 @@ function broadcastChatFinal(params: {
seq, seq,
state: "final" as const, state: "final" as const,
message: params.message, message: params.message,
usage: params.usage,
}; };
params.context.broadcast("chat", payload); params.context.broadcast("chat", payload);
params.context.nodeSendToSession(params.sessionKey, "chat", payload); params.context.nodeSendToSession(params.sessionKey, "chat", payload);
@ -545,6 +547,7 @@ export const chatHandlers: GatewayRequestHandlers = {
runId: clientRunId, runId: clientRunId,
sessionKey: p.sessionKey, sessionKey: p.sessionKey,
message, message,
usage: message?.usage as Record<string, unknown> | undefined,
}); });
} }
context.dedupe.set(`chat:${clientRunId}`, { context.dedupe.set(`chat:${clientRunId}`, {