diff --git a/packages/server/src/server/agent/agent-manager.ts b/packages/server/src/server/agent/agent-manager.ts index f95f29205..cd84027ec 100644 --- a/packages/server/src/server/agent/agent-manager.ts +++ b/packages/server/src/server/agent/agent-manager.ts @@ -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") { diff --git a/packages/server/src/server/agent/agent-sdk-types.ts b/packages/server/src/server/agent/agent-sdk-types.ts index 795e375a0..cc0850cc1 100644 --- a/packages/server/src/server/agent/agent-sdk-types.ts +++ b/packages/server/src/server/agent/agent-sdk-types.ts @@ -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; diff --git a/packages/server/src/server/agent/providers/claude-agent.test.ts b/packages/server/src/server/agent/providers/claude-agent.test.ts index b6863a3f7..8da63128c 100644 --- a/packages/server/src/server/agent/providers/claude-agent.test.ts +++ b/packages/server/src/server/agent/providers/claude-agent.test.ts @@ -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, + }); + }); }); diff --git a/packages/server/src/server/agent/providers/claude-agent.ts b/packages/server/src/server/agent/providers/claude-agent.ts index e780586e2..f936e9124 100644 --- a/packages/server/src/server/agent/providers/claude-agent.ts +++ b/packages/server/src/server/agent/providers/claude-agent.ts @@ -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): 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): 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; + 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, diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts index 813a201c1..0cd41a189 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.test.ts @@ -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 }, diff --git a/packages/server/src/server/agent/providers/codex-app-server-agent.ts b/packages/server/src/server/agent/providers/codex-app-server-agent.ts index 141a05d96..35baca8d2 100644 --- a/packages/server/src/server/agent/providers/codex-app-server-agent.ts +++ b/packages/server/src/server/agent/providers/codex-app-server-agent.ts @@ -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; } diff --git a/packages/server/src/server/agent/providers/opencode-agent.ts b/packages/server/src/server/agent/providers/opencode-agent.ts index 6f004d11f..fcb93de7d 100644 --- a/packages/server/src/server/agent/providers/opencode-agent.ts +++ b/packages/server/src/server/agent/providers/opencode-agent.ts @@ -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; } diff --git a/packages/server/src/server/agent/providers/opencode/event-translator.test.ts b/packages/server/src/server/agent/providers/opencode/event-translator.test.ts index b68a4f009..72986fe28 100644 --- a/packages/server/src/server/agent/providers/opencode/event-translator.test.ts +++ b/packages/server/src/server/agent/providers/opencode/event-translator.test.ts @@ -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();