feat: emit usage_updated events for real-time token usage updates

This commit is contained in:
Mohamed Boudra
2026-04-08 08:59:02 +07:00
parent 5ac7b3f7c7
commit 5c93fbc955
8 changed files with 498 additions and 7 deletions
@@ -2124,6 +2124,10 @@ export class AgentManager {
void this.refreshRuntimeInfo(agent);
}
break;
case "usage_updated":
agent.lastUsage = event.usage;
this.emitState(agent);
break;
case "timeline":
// Skip provider-replayed user_message items during history hydration.
if (options?.fromHistory && event.item.type === "user_message") {
@@ -301,6 +301,7 @@ export type AgentStreamEvent =
| { type: "thread_started"; sessionId: string; provider: AgentProvider }
| { type: "turn_started"; provider: AgentProvider; turnId?: string }
| { type: "turn_completed"; provider: AgentProvider; usage?: AgentUsage; turnId?: string }
| { type: "usage_updated"; provider: AgentProvider; usage: AgentUsage; turnId?: string }
| {
type: "turn_failed";
provider: AgentProvider;
@@ -468,6 +468,124 @@ describe("ClaudeAgentSession context window usage", () => {
});
});
test("task_progress emits a usage_updated event", async () => {
const session = await createSessionForTest();
const events = session.translateMessageToEvents({
type: "system",
subtype: "task_progress",
task_id: "task-1",
description: "Processing",
usage: {
total_tokens: 999,
tool_uses: 1,
duration_ms: 50,
},
uuid: "task-progress-1",
session_id: "session-1",
});
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 999,
},
});
});
test("task_notification emits a usage_updated event", async () => {
const session = await createSessionForTest();
const events = session.translateMessageToEvents({
type: "system",
subtype: "task_notification",
uuid: "task-note-1",
task_id: "task-1",
status: "running",
summary: "Background task still running",
usage: {
total_tokens: 777,
tool_uses: 1,
duration_ms: 50,
},
session_id: "session-1",
} as any);
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 777,
},
});
});
test("message_start stream events emit usage_updated with per-request usage", async () => {
const session = await createSessionForTest();
const events = session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
},
},
session_id: "session-1",
} as any);
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 150,
},
});
});
test("message_delta stream events update per-request usage", async () => {
const session = await createSessionForTest();
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
},
},
session_id: "session-1",
} as any);
const events = session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 25,
},
},
session_id: "session-1",
} as any);
expect(events).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 175,
},
});
});
test("task_progress usage takes priority over derived result usage", async () => {
const session = await createSessionForTest();
@@ -638,4 +756,134 @@ describe("ClaudeAgentSession context window usage", () => {
contextWindowUsedTokens: 22,
});
});
test("convertUsage uses per-request stream usage when no task_progress is available", async () => {
const session = await createSessionForTest();
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
},
},
session_id: "session-1",
} as any);
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 25,
},
},
session_id: "session-1",
} as any);
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 175,
});
});
test("per-request stream usage is not cumulative across API calls in a turn", async () => {
const session = await createSessionForTest();
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 100,
cache_creation_input_tokens: 20,
cache_read_input_tokens: 30,
},
},
},
session_id: "session-1",
} as any);
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 25,
},
},
session_id: "session-1",
} as any);
const secondStartEvents = session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_start",
message: {
usage: {
input_tokens: 40,
cache_creation_input_tokens: 5,
cache_read_input_tokens: 10,
},
},
},
session_id: "session-1",
} as any);
expect(secondStartEvents).toContainEqual({
type: "usage_updated",
provider: "claude",
usage: {
contextWindowUsedTokens: 55,
},
});
session.translateMessageToEvents({
type: "stream_event",
event: {
type: "message_delta",
usage: {
output_tokens: 7,
},
},
session_id: "session-1",
} as any);
const usage = session.convertUsage({
type: "result",
subtype: "success",
usage: {
input_tokens: 10,
cache_read_input_tokens: 5,
output_tokens: 7,
},
total_cost_usd: 0.12,
});
expect(usage).toEqual({
inputTokens: 10,
cachedInputTokens: 5,
outputTokens: 7,
totalCostUsd: 0.12,
contextWindowUsedTokens: 62,
});
});
});
@@ -1225,16 +1225,67 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined {
return maxContextWindow;
}
function readContextWindowUsedTokensFromTaskProgress(
message: SDKTaskProgressMessage,
function readUsageTotalTokens(
usage: unknown,
): number | undefined {
const totalTokens = message.usage?.total_tokens;
if (!usage || typeof usage !== "object") {
return undefined;
}
const totalTokens = (usage as { total_tokens?: unknown }).total_tokens;
if (typeof totalTokens !== "number" || !Number.isFinite(totalTokens) || totalTokens < 0) {
return undefined;
}
return totalTokens;
}
function readContextWindowUsedTokensFromTaskProgress(
message: SDKTaskProgressMessage,
): number | undefined {
return readUsageTotalTokens(message.usage);
}
function readUsageFromTaskNotification(message: { usage?: unknown }): number | undefined {
return readUsageTotalTokens(message.usage);
}
function readStreamRequestInputTokens(event: Record<string, unknown>): number | undefined {
const messageUsage = (event.message as { usage?: unknown } | undefined)?.usage;
if (!messageUsage || typeof messageUsage !== "object") {
return undefined;
}
const usage = messageUsage as {
input_tokens?: unknown;
cache_creation_input_tokens?: unknown;
cache_read_input_tokens?: unknown;
};
const inputTokens =
typeof usage.input_tokens === "number" && Number.isFinite(usage.input_tokens)
? usage.input_tokens
: undefined;
const cacheCreationInputTokens =
typeof usage.cache_creation_input_tokens === "number" &&
Number.isFinite(usage.cache_creation_input_tokens)
? usage.cache_creation_input_tokens
: 0;
const cacheReadInputTokens =
typeof usage.cache_read_input_tokens === "number" &&
Number.isFinite(usage.cache_read_input_tokens)
? usage.cache_read_input_tokens
: 0;
if (typeof inputTokens !== "number" || inputTokens < 0) {
return undefined;
}
return inputTokens + cacheCreationInputTokens + cacheReadInputTokens;
}
function readStreamRequestOutputTokens(event: Record<string, unknown>): number | undefined {
const outputTokens = (event.usage as { output_tokens?: unknown } | undefined)?.output_tokens;
if (typeof outputTokens !== "number" || !Number.isFinite(outputTokens) || outputTokens < 0) {
return undefined;
}
return outputTokens;
}
class ClaudeAgentSession implements AgentSession {
readonly provider: "claude" = "claude";
readonly capabilities = CLAUDE_CAPABILITIES;
@@ -1278,6 +1329,9 @@ class ClaudeAgentSession implements AgentSession {
private lastForegroundPromptText: string | null = null;
private foregroundHasVisibleActivity = false;
private lastContextWindowUsedTokens: number | undefined;
private lastContextWindowMaxTokens: number | undefined;
private lastStreamRequestInputTokens: number | undefined;
private lastStreamRequestOutputTokens: number | undefined;
private userMessageIds: string[] = [];
private recentStderr = "";
private closed = false;
@@ -2672,10 +2726,18 @@ class ClaudeAgentSession implements AgentSession {
provider: "claude",
});
}
const usage = readUsageFromTaskNotification(message);
if (typeof usage === "number") {
this.lastContextWindowUsedTokens = usage;
events.push(this.createUsageUpdatedEvent(usage));
}
} else if (message.subtype === "task_progress") {
this.lastContextWindowUsedTokens =
readContextWindowUsedTokensFromTaskProgress(message) ??
this.lastContextWindowUsedTokens;
if (typeof this.lastContextWindowUsedTokens === "number") {
events.push(this.createUsageUpdatedEvent(this.lastContextWindowUsedTokens));
}
}
break;
case "user": {
@@ -2743,6 +2805,10 @@ class ClaudeAgentSession implements AgentSession {
break;
}
case "stream_event": {
const usageUpdatedEvent = this.trackStreamEventUsage(message.event);
if (usageUpdatedEvent) {
events.push(usageUpdatedEvent);
}
const timelineItems = this.mapPartialEvent(message.event, {
suppressAssistantText: options?.suppressAssistantText ?? false,
suppressReasoning: options?.suppressReasoning ?? false,
@@ -2911,12 +2977,21 @@ class ClaudeAgentSession implements AgentSession {
modelUsage ?? message.modelUsage,
);
if (contextWindowMaxTokens !== undefined) {
this.lastContextWindowMaxTokens = contextWindowMaxTokens;
usage.contextWindowMaxTokens = contextWindowMaxTokens;
} else if (this.lastContextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens;
}
if (typeof this.lastContextWindowUsedTokens === "number") {
// task_progress.total_tokens is the accurate context window fill level.
// Prefer it over result.usage which contains accumulated session totals.
usage.contextWindowUsedTokens = this.lastContextWindowUsedTokens;
} else if (
typeof this.lastStreamRequestInputTokens === "number" &&
typeof this.lastStreamRequestOutputTokens === "number"
) {
usage.contextWindowUsedTokens =
this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens;
} else if (message.usage) {
// Fallback: derive from result.usage when no task_progress has been
// received yet. These values are accumulated across all API calls, but
@@ -2937,6 +3012,54 @@ class ClaudeAgentSession implements AgentSession {
return usage;
}
private createUsageUpdatedEvent(contextWindowUsedTokens: number): AgentStreamEvent {
const usage: AgentUsage = {
contextWindowUsedTokens,
};
if (this.lastContextWindowMaxTokens !== undefined) {
usage.contextWindowMaxTokens = this.lastContextWindowMaxTokens;
}
return {
type: "usage_updated",
provider: "claude",
usage,
};
}
private trackStreamEventUsage(event: unknown): AgentStreamEvent | null {
if (!event || typeof event !== "object") {
return null;
}
const streamEvent = event as Record<string, unknown>;
const eventType = readTrimmedString(streamEvent.type);
if (eventType === "message_start") {
const inputTokens = readStreamRequestInputTokens(streamEvent);
if (typeof inputTokens !== "number") {
return null;
}
this.lastStreamRequestInputTokens = inputTokens;
this.lastStreamRequestOutputTokens = 0;
} else if (eventType === "message_delta") {
const outputTokens = readStreamRequestOutputTokens(streamEvent);
if (typeof outputTokens !== "number") {
return null;
}
this.lastStreamRequestOutputTokens = outputTokens;
} else {
return null;
}
if (
typeof this.lastStreamRequestInputTokens !== "number" ||
typeof this.lastStreamRequestOutputTokens !== "number"
) {
return null;
}
return this.createUsageUpdatedEvent(
this.lastStreamRequestInputTokens + this.lastStreamRequestOutputTokens,
);
}
private handlePermissionRequest: CanUseTool = async (
toolName,
input,
@@ -577,6 +577,52 @@ describe("Codex app-server provider", () => {
});
});
test("emits usage_updated on token usage updates and keeps usage on turn completion", () => {
const session = createSession();
const events: AgentStreamEvent[] = [];
session.subscribe((event) => events.push(event));
(session as any).handleNotification("thread/tokenUsage/updated", {
tokenUsage: {
model_context_window: 200000,
last: {
total_tokens: 50000,
inputTokens: 30000,
cachedInputTokens: 5000,
outputTokens: 15000,
},
},
});
(session as any).handleNotification("turn/completed", {
turn: { status: "completed", error: null },
});
expect(events).toContainEqual({
type: "usage_updated",
provider: "codex",
turnId: "test-turn",
usage: {
inputTokens: 30000,
cachedInputTokens: 5000,
outputTokens: 15000,
contextWindowMaxTokens: 200000,
contextWindowUsedTokens: 50000,
},
});
expect(events.at(-1)).toEqual({
type: "turn_completed",
provider: "codex",
turnId: "test-turn",
usage: {
inputTokens: 30000,
cachedInputTokens: 5000,
outputTokens: 15000,
contextWindowMaxTokens: 200000,
contextWindowUsedTokens: 50000,
},
});
});
test("approving a synthetic Codex plan permission disables plan and fast mode and starts implementation", async () => {
const session = createSession({
featureValues: { plan_mode: true, fast_mode: true },
@@ -3506,6 +3506,13 @@ class CodexAppServerAgentSession implements AgentSession {
if (parsed.kind === "token_usage_updated") {
this.latestUsage = toAgentUsage(parsed.tokenUsage);
if (this.latestUsage) {
this.notifySubscribers({
type: "usage_updated",
provider: CODEX_PROVIDER,
usage: this.latestUsage,
});
}
return;
}
@@ -450,16 +450,16 @@ function mergeOpenCodeStepFinishUsage(
const cost = readPositiveFiniteNumber(part.cost);
if (inputTokens !== undefined) {
usage.inputTokens = (usage.inputTokens ?? 0) + inputTokens;
usage.inputTokens = inputTokens;
}
if (cacheReadTokens !== undefined) {
usage.cachedInputTokens = (usage.cachedInputTokens ?? 0) + cacheReadTokens;
usage.cachedInputTokens = cacheReadTokens;
}
if (outputTokens !== undefined) {
usage.outputTokens = (usage.outputTokens ?? 0) + outputTokens;
usage.outputTokens = outputTokens;
}
if (totalTokens > 0) {
usage.contextWindowUsedTokens = (usage.contextWindowUsedTokens ?? 0) + totalTokens;
usage.contextWindowUsedTokens = totalTokens;
}
if (cost !== undefined) {
usage.totalCostUsd = (usage.totalCostUsd ?? 0) + cost;
@@ -1144,6 +1144,13 @@ export function translateOpenCodeEvent(
}
} else if (partType === "step-finish") {
mergeOpenCodeStepFinishUsage(state.accumulatedUsage, part);
if (hasNormalizedOpenCodeUsage(state.accumulatedUsage)) {
events.push({
type: "usage_updated",
provider: "opencode",
usage: { ...state.accumulatedUsage },
});
}
}
break;
}
@@ -255,6 +255,61 @@ describe("translateOpenCodeEvent", () => {
]);
});
it("emits usage_updated after step-finish parts", () => {
const state = createState();
state.accumulatedUsage.contextWindowMaxTokens = 400_000;
const events = translateOpenCodeEvent(
{
type: "message.part.updated",
properties: {
part: {
id: "step-finish-1",
sessionID: "session-1",
messageID: "message-usage-1",
type: "step-finish",
reason: "stop",
cost: 0.25,
tokens: {
total: 999_999,
input: 30_000,
output: 12_000,
reasoning: 10_000,
cache: {
read: 2_000,
write: 1_000,
},
},
},
},
},
state,
);
expect(events).toEqual([
{
type: "usage_updated",
provider: "opencode",
usage: {
contextWindowMaxTokens: 400_000,
contextWindowUsedTokens: 55_000,
cachedInputTokens: 2_000,
inputTokens: 30_000,
outputTokens: 12_000,
totalCostUsd: 0.25,
},
},
]);
expect(state.accumulatedUsage).toEqual({
contextWindowMaxTokens: 400_000,
contextWindowUsedTokens: 55_000,
cachedInputTokens: 2_000,
inputTokens: 30_000,
outputTokens: 12_000,
totalCostUsd: 0.25,
});
});
it("emits reasoning from message.part.delta events", () => {
const state = createState();