From dd67450fd9bddc82db3e71273b535e07b2672804 Mon Sep 17 00:00:00 2001 From: Michael Gendy <50384638+Mng-dev-ai@users.noreply.github.com> Date: Thu, 16 Apr 2026 18:55:28 +0200 Subject: [PATCH] fix: emit usage updates from Claude stream events (#506) ## Summary - emit ACP `usage_update` notifications from top-level Claude `stream_event` message usage - keep tracking a usage snapshot across `message_start` and `message_delta` so partial deltas still produce a complete total - preserve the existing `result`-path update as the final exact/cost-bearing usage event ## Why Claude SDK stream events already carry incremental usage on `message_delta`, but `claude-agent-acp` currently ignores that path and only emits `usage_update` from the final `result`. This makes usage appear stale until the turn ends in ACP clients that expect live context usage updates. ## Verification - built the repo with `npm run build` - ran a direct ACP stdio client against the patched local build - observed `usage_update` notifications emitted before the final result-path usage update Patched local smoke-test output included this sequence: - `usage_update used=20973 size=200000` - `agent_message_chunk text=TEST_OK` - `usage_update used=20989 size=200000` - `usage_update used=20989 size=1000000 cost=0.1314875` --------- Co-authored-by: Ben Brandt --- src/acp-agent.ts | 174 +++++++-- src/tests/acp-agent.test.ts | 478 +++++++++++++++++++++++ src/tests/session-config-options.test.ts | 3 + 3 files changed, 626 insertions(+), 29 deletions(-) diff --git a/src/acp-agent.ts b/src/acp-agent.ts index 6c4fb342..91d07b32 100644 --- a/src/acp-agent.ts +++ b/src/acp-agent.ts @@ -54,7 +54,6 @@ import { Query, query, SDKPartialAssistantMessage, - SDKResultMessage, SDKUserMessage, SlashCommand, } from "@anthropic-ai/claude-agent-sdk"; @@ -108,6 +107,22 @@ type AccumulatedUsage = { cachedWriteTokens: number; }; +type UsageSnapshot = { + input_tokens: number; + output_tokens: number; + cache_read_input_tokens: number; + cache_creation_input_tokens: number; +}; + +const ZERO_USAGE = Object.freeze({ + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); + +const DEFAULT_CONTEXT_WINDOW = 200000; + type Session = { query: Query; input: Pushable; @@ -126,6 +141,12 @@ type Session = { nextPendingOrder: number; abortController: AbortController; emitRawSDKMessages: boolean | SDKMessageFilter[]; + /** Context window size of the last top-level assistant model, carried across + * prompts so mid-stream usage_update notifications report a correct `size` + * before the turn's first result message arrives. Defaults to + * DEFAULT_CONTEXT_WINDOW, refreshed from each result's modelUsage, and + * invalidated when the user switches the session's model. */ + contextWindowSize: number; }; /** Compute a stable fingerprint of the session-defining params so we can @@ -545,8 +566,8 @@ export class ClaudeAcpAgent implements Agent { }; let lastAssistantTotalUsage: number | null = null; + let lastAssistantUsage: UsageSnapshot | null = null; let lastAssistantModel: string | null = null; - let lastContextWindowSize: number = 200000; const userMessage = promptToClaude(params); @@ -628,12 +649,13 @@ export class ClaudeAcpAgent implements Agent { // "944k/1m" right after the user sees "Compacting completed", // which is confusing and wrong. lastAssistantTotalUsage = 0; + lastAssistantUsage = null; await this.client.sessionUpdate({ sessionId: message.session_id, update: { sessionUpdate: "usage_update", used: 0, - size: lastContextWindowSize, + size: session.contextWindowSize, }, }); await this.client.sessionUpdate({ @@ -691,8 +713,13 @@ export class ClaudeAcpAgent implements Agent { const matchingModelUsage = lastAssistantModel ? getMatchingModelUsage(message.modelUsage, lastAssistantModel) : null; - const contextWindowSize = matchingModelUsage?.contextWindow ?? 200000; - lastContextWindowSize = contextWindowSize; + // Only overwrite when we have an authoritative value — a miss + // (e.g. a turn with no top-level assistant message) would + // otherwise discard the window learned on a prior turn and + // leave the next prompt's mid-stream updates reporting 200k. + if (matchingModelUsage) { + session.contextWindowSize = matchingModelUsage.contextWindow; + } // Send usage_update notification if (lastAssistantTotalUsage !== null) { @@ -701,7 +728,7 @@ export class ClaudeAcpAgent implements Agent { update: { sessionUpdate: "usage_update", used: lastAssistantTotalUsage, - size: contextWindowSize, + size: session.contextWindowSize, cost: { amount: message.total_cost_usd, currency: "USD", @@ -775,6 +802,57 @@ export class ClaudeAcpAgent implements Agent { break; } case "stream_event": { + if ( + message.parent_tool_use_id === null && + (message.event.type === "message_start" || message.event.type === "message_delta") + ) { + if (message.event.type === "message_start") { + lastAssistantUsage = snapshotFromUsage(message.event.message.usage); + const model = message.event.message.model; + if (model && model !== "") { + lastAssistantModel = model; + // Only upgrade from the default — once a `result` has given + // us an authoritative window, trust it over the heuristic. + // Model switches invalidate the cached window via + // `syncSessionConfigState`, which resets us back to the + // default so this branch runs again for the new model. + if (session.contextWindowSize === DEFAULT_CONTEXT_WINDOW) { + const inferred = inferContextWindowFromModel(model); + if (inferred !== null) { + session.contextWindowSize = inferred; + } + } + } + } else { + const usage = message.event.usage; + const prev: Readonly = lastAssistantUsage ?? ZERO_USAGE; + // Per Anthropic API, message_delta usage fields are *cumulative*; + // nullable fields (input_tokens and the cache fields) fall back + // to the prior snapshot when the server omits them from this + // delta. Only output_tokens is guaranteed non-null. + lastAssistantUsage = { + input_tokens: usage.input_tokens ?? prev.input_tokens, + output_tokens: usage.output_tokens, + cache_read_input_tokens: + usage.cache_read_input_tokens ?? prev.cache_read_input_tokens, + cache_creation_input_tokens: + usage.cache_creation_input_tokens ?? prev.cache_creation_input_tokens, + }; + } + + const nextUsage = totalTokens(lastAssistantUsage); + if (nextUsage !== lastAssistantTotalUsage) { + lastAssistantTotalUsage = nextUsage; + await this.client.sessionUpdate({ + sessionId: params.sessionId, + update: { + sessionUpdate: "usage_update", + used: nextUsage, + size: session.contextWindowSize, + }, + }); + } + } for (const notification of streamEventToAcpNotifications( message, params.sessionId, @@ -817,29 +895,16 @@ export class ClaudeAcpAgent implements Agent { } } - // Store latest assistant usage (excluding subagents) - // Sum all token types as a proxy for post-turn context occupancy: - // current turn's output will become next turn's input. - // Note: per the Anthropic API, input_tokens excludes cache tokens — - // cache_read and cache_creation are reported separately, so summing - // all four fields is not double-counting. - if ((message.message as any).usage && message.parent_tool_use_id === null) { - const messageWithUsage = message.message as unknown as SDKResultMessage; - lastAssistantTotalUsage = - messageWithUsage.usage.input_tokens + - messageWithUsage.usage.output_tokens + - messageWithUsage.usage.cache_read_input_tokens + - messageWithUsage.usage.cache_creation_input_tokens; - } - // Track the current top-level model for context window size lookup - // (exclude subagent messages to stay in sync with lastAssistantTotalUsage) - if ( - message.type === "assistant" && - message.parent_tool_use_id === null && - message.message.model && - message.message.model !== "" - ) { - lastAssistantModel = message.message.model; + // Snapshot the latest top-level assistant usage and model so the + // next `result` can emit a usage_update tied to the right context + // window. Subagent messages are excluded to keep the snapshot + // aligned with what the user's current selection is producing. + if (message.type === "assistant" && message.parent_tool_use_id === null) { + lastAssistantUsage = snapshotFromUsage(message.message.usage); + lastAssistantTotalUsage = totalTokens(lastAssistantUsage); + if (message.message.model && message.message.model !== "") { + lastAssistantModel = message.message.model; + } } // Slash commands like /compact can generate invalid output... doesn't match @@ -1334,6 +1399,13 @@ export class ClaudeAcpAgent implements Agent { if (configId === "mode") { session.modes = { ...session.modes, currentModeId: value }; } else if (configId === "model") { + if (session.models.currentModelId !== value) { + // The cached context window was learned for the previous model; reset + // to the new model's heuristic so mid-stream updates between now and + // the next `result` reflect the user's selection instead of the old + // model's window. + session.contextWindowSize = inferContextWindowFromModel(value) ?? DEFAULT_CONTEXT_WINDOW; + } session.models = { ...session.models, currentModelId: value }; } } @@ -1642,6 +1714,8 @@ export class ClaudeAcpAgent implements Agent { nextPendingOrder: 0, abortController, emitRawSDKMessages: sessionMeta?.claudeCode?.emitRawSDKMessages ?? false, + contextWindowSize: + inferContextWindowFromModel(models.currentModelId) ?? DEFAULT_CONTEXT_WINDOW, }; return { @@ -1678,6 +1752,38 @@ function sessionUsage(session: Session) { }; } +/** Sum all four fields as a proxy for post-turn context occupancy: the current + * turn's output becomes next turn's input. Per the Anthropic API, input_tokens + * excludes cache tokens — cache_read and cache_creation are reported + * separately — so summing all four is not double-counting. */ +function totalTokens(usage: UsageSnapshot): number { + return ( + usage.input_tokens + + usage.output_tokens + + usage.cache_read_input_tokens + + usage.cache_creation_input_tokens + ); +} + +/** Project a nullable API usage object into our non-null snapshot shape. + * Both SDK message_start and assistant message `usage` have `number | null` + * cache fields; we coerce absent values to 0 so `totalTokens` never hits + * NaN. Delta events have different semantics (cumulative + prev fallback) + * and are handled inline. */ +function snapshotFromUsage(usage: { + input_tokens: number; + output_tokens: number; + cache_read_input_tokens: number | null; + cache_creation_input_tokens: number | null; +}): UsageSnapshot { + return { + input_tokens: usage.input_tokens, + output_tokens: usage.output_tokens, + cache_read_input_tokens: usage.cache_read_input_tokens ?? 0, + cache_creation_input_tokens: usage.cache_creation_input_tokens ?? 0, + }; +} + function createEnvForGateway(gatewayMeta?: GatewayAuthMeta) { if (!gatewayMeta) { return {}; @@ -2302,6 +2408,16 @@ function commonPrefixLength(a: string, b: string) { return i; } +/** Best-effort first guess of a model's context window from its ID, used only + * until a `result` message arrives with the authoritative `modelUsage` value. + * Anthropic 1M-context variants encode "1m" as a distinct token in the SDK + * model ID (e.g., "claude-opus-4-6-1m"), which `\b1m\b` catches without also + * matching things like "10m" or embedded substrings. */ +function inferContextWindowFromModel(model: string): number | null { + if (/\b1m\b/i.test(model)) return 1_000_000; + return null; +} + function getMatchingModelUsage(modelUsage: Record, currentModel: string) { let bestKey: string | null = null; let bestLen = 0; diff --git a/src/tests/acp-agent.test.ts b/src/tests/acp-agent.test.ts index 4b3bbe84..bd290a44 100644 --- a/src/tests/acp-agent.test.ts +++ b/src/tests/acp-agent.test.ts @@ -1350,6 +1350,7 @@ describe("stop reason propagation", () => { nextPendingOrder: 0, abortController: new AbortController(), emitRawSDKMessages: false, + contextWindowSize: 200000, }; } @@ -1491,6 +1492,7 @@ describe("stop reason propagation", () => { pendingMessages: new Map(), nextPendingOrder: 0, emitRawSDKMessages: false, + contextWindowSize: 200000, }; const response = await agent.prompt({ @@ -1566,6 +1568,7 @@ describe("session/close", () => { nextPendingOrder: 0, abortController: new AbortController(), emitRawSDKMessages: false, + contextWindowSize: 200000, }; return agent.sessions[sessionId]!; } @@ -1660,6 +1663,7 @@ describe("getOrCreateSession param change detection", () => { nextPendingOrder: 0, abortController: new AbortController(), emitRawSDKMessages: false, + contextWindowSize: 200000, }; return agent.sessions[sessionId]!; } @@ -1819,6 +1823,23 @@ describe("usage_update computation", () => { }; } + function createStreamEvent( + eventType: "message_start" | "message_delta", + payload: Record, + parentToolUseId: string | null = null, + ) { + return { + type: "stream_event" as const, + parent_tool_use_id: parentToolUseId, + uuid: randomUUID(), + session_id: "test-session", + event: + eventType === "message_start" + ? { type: "message_start" as const, message: payload } + : { type: "message_delta" as const, ...payload }, + }; + } + function createMockAgentWithCapture() { const updates: any[] = []; const mockClient = { @@ -1875,6 +1896,7 @@ describe("usage_update computation", () => { nextPendingOrder: 0, abortController: new AbortController(), emitRawSDKMessages: false, + contextWindowSize: 200000, }; } @@ -1915,6 +1937,461 @@ describe("usage_update computation", () => { expect(usageUpdate.update.used).toBe(1800); }); + it("stream_event message_start emits usage_update before result", async () => { + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-20250514", + usage: { + input_tokens: 1000, + output_tokens: 500, + cache_read_input_tokens: 200, + cache_creation_input_tokens: 100, + }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-20250514": { + inputTokens: 1000, + outputTokens: 500, + cacheReadInputTokens: 200, + cacheCreationInputTokens: 100, + webSearchRequests: 0, + costUSD: 0.01, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(2); + expect(usageUpdates[0].update.used).toBe(1800); + // First prompt of a session has no prior result to learn the window from, + // so the mid-stream update falls back to the default context window. + expect(usageUpdates[0].update.size).toBe(200000); + expect(usageUpdates[0].update.cost).toBeUndefined(); + expect(usageUpdates[1].update.used).toBe(1800); + expect(usageUpdates[1].update.size).toBe(1000000); + expect(usageUpdates[1].update.cost).toBeDefined(); + }); + + it("stream_event message_delta patches previous snapshot", async () => { + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-20250514", + usage: { + input_tokens: 1000, + output_tokens: 0, + cache_read_input_tokens: 200, + cache_creation_input_tokens: 100, + }, + }), + createStreamEvent("message_delta", { + usage: { output_tokens: 500 }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-20250514": { + inputTokens: 1000, + outputTokens: 500, + cacheReadInputTokens: 200, + cacheCreationInputTokens: 100, + webSearchRequests: 0, + costUSD: 0.01, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(3); + expect(usageUpdates[0].update.used).toBe(1300); + expect(usageUpdates[0].update.cost).toBeUndefined(); + expect(usageUpdates[1].update.used).toBe(1800); + expect(usageUpdates[1].update.cost).toBeUndefined(); + expect(usageUpdates[2].update.used).toBe(1800); + expect(usageUpdates[2].update.cost).toBeDefined(); + }); + + it("mid-stream size is inferred from a 1M model name before the first result", async () => { + // On the very first prompt there is no learned context window yet, so the + // mid-stream update would otherwise fall back to 200k. A "-1m" suffix in + // the SDK model ID is enough signal to emit 1_000_000 up front. + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-6-1m", + usage: { + input_tokens: 2000, + output_tokens: 1000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-6-1m": { + inputTokens: 2000, + outputTokens: 1000, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.02, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(2); + expect(usageUpdates[0].update.size).toBe(1000000); + expect(usageUpdates[1].update.size).toBe(1000000); + }); + + it("duplicate stream_event totals do not re-emit usage_update", async () => { + // A message_delta whose cumulative totals match the prior snapshot should + // not trigger a duplicate usage_update — only the result adds cost on top. + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-20250514", + usage: { + input_tokens: 1000, + output_tokens: 500, + cache_read_input_tokens: 200, + cache_creation_input_tokens: 100, + }, + }), + createStreamEvent("message_delta", { + usage: { output_tokens: 500 }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-20250514": { + inputTokens: 1000, + outputTokens: 500, + cacheReadInputTokens: 200, + cacheCreationInputTokens: 100, + webSearchRequests: 0, + costUSD: 0.01, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(2); + expect(usageUpdates[0].update.used).toBe(1800); + expect(usageUpdates[0].update.cost).toBeUndefined(); + expect(usageUpdates[1].update.used).toBe(1800); + expect(usageUpdates[1].update.cost).toBeDefined(); + }); + + it("mid-stream size uses the session's learned context window", async () => { + // Session state persists the model's context window across prompts, so a + // mid-stream update in a later prompt reports the real size immediately + // instead of snapping back to the 200k default before the result arrives. + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-20250514", + usage: { + input_tokens: 2000, + output_tokens: 1000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-20250514": { + inputTokens: 2000, + outputTokens: 1000, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.02, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + // Simulate a prior prompt having learned the 1M window for this model. + agent.sessions["test-session"].contextWindowSize = 1000000; + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(2); + expect(usageUpdates[0].update.size).toBe(1000000); + expect(usageUpdates[1].update.size).toBe(1000000); + }); + + it("switching to a 1M model seeds the context window from the heuristic", async () => { + // The heuristic runs at config-change time so mid-stream updates in the + // next prompt already report 1M — without waiting for message_start or + // the next `result` to correct us. + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-6-1m", + usage: { + input_tokens: 2000, + output_tokens: 1000, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-6-1m": { + inputTokens: 2000, + outputTokens: 1000, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.02, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + const session = agent.sessions["test-session"]; + expect(session.contextWindowSize).toBe(200000); + + (agent as any).syncSessionConfigState(session, "model", "claude-opus-4-6-1m"); + expect(session.contextWindowSize).toBe(1000000); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(2); + expect(usageUpdates[0].update.size).toBe(1000000); + expect(usageUpdates[1].update.size).toBe(1000000); + }); + + it("result with no matching modelUsage preserves the learned window", async () => { + // A turn whose `result.modelUsage` doesn't contain the current top-level + // model (e.g. no top-level assistant message, or only a subagent ran) must + // not clobber the window learned on a prior turn — otherwise the next + // prompt's mid-stream updates regress to the 200k default. + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createResultMessageWithModel({ + modelUsage: { + "claude-haiku-4-5-20251001": { + inputTokens: 100, + outputTokens: 50, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.001, + contextWindow: 200000, + maxOutputTokens: 8192, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + const session = agent.sessions["test-session"]; + session.contextWindowSize = 1000000; + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + expect(session.contextWindowSize).toBe(1000000); + // The emit itself falls back to session.contextWindowSize, which is + // unchanged from the learned value. + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + // No lastAssistantTotalUsage was set (no top-level assistant / stream + // event), so the result branch skips its emit entirely. + expect(usageUpdates).toHaveLength(0); + }); + + it("switching the session's model invalidates the learned context window", async () => { + // When the user switches models mid-session, the window learned for the + // previous model would otherwise persist into the next prompt's first + // mid-stream update. syncSessionConfigState should reset it so the next + // turn's first update falls back to the heuristic (here: 200k default). + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-sonnet-4-6", + usage: { + input_tokens: 1000, + output_tokens: 500, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }), + createResultMessageWithModel({ + modelUsage: { + "claude-sonnet-4-6": { + inputTokens: 1000, + outputTokens: 500, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.01, + contextWindow: 200000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + const session = agent.sessions["test-session"]; + session.contextWindowSize = 1000000; + session.models = { ...session.models, currentModelId: "claude-opus-4-6-1m" }; + + // User flips the selector to a 200k model. + (agent as any).syncSessionConfigState(session, "model", "claude-sonnet-4-6"); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(2); + expect(usageUpdates[0].update.size).toBe(200000); + expect(usageUpdates[1].update.size).toBe(200000); + }); + + it("non-usage stream events do not re-emit usage_update", async () => { + // content_block_* and message_stop carry no usage fields; they must not + // trigger duplicate emits between the real message_start / message_delta + // / result updates. + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent("message_start", { + model: "claude-opus-4-20250514", + usage: { + input_tokens: 1000, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }), + { + type: "stream_event" as const, + parent_tool_use_id: null, + uuid: randomUUID(), + session_id: "test-session", + event: { type: "content_block_start", index: 0, content_block: { type: "text", text: "" } }, + }, + { + type: "stream_event" as const, + parent_tool_use_id: null, + uuid: randomUUID(), + session_id: "test-session", + event: { type: "content_block_delta", index: 0, delta: { type: "text_delta", text: "hi" } }, + }, + { + type: "stream_event" as const, + parent_tool_use_id: null, + uuid: randomUUID(), + session_id: "test-session", + event: { type: "content_block_stop", index: 0 }, + }, + createStreamEvent("message_delta", { + usage: { output_tokens: 200 }, + }), + { + type: "stream_event" as const, + parent_tool_use_id: null, + uuid: randomUUID(), + session_id: "test-session", + event: { type: "message_stop" }, + }, + createResultMessageWithModel({ + modelUsage: { + "claude-opus-4-20250514": { + inputTokens: 1000, + outputTokens: 200, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.01, + contextWindow: 1000000, + maxOutputTokens: 16384, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + // Exactly three: message_start (1000), message_delta (1200), result (1200 + cost). + expect(usageUpdates).toHaveLength(3); + expect(usageUpdates[0].update.used).toBe(1000); + expect(usageUpdates[1].update.used).toBe(1200); + expect(usageUpdates[2].update.used).toBe(1200); + expect(usageUpdates[2].update.cost).toBeDefined(); + }); + + it("subagent stream_event does not emit usage_update", async () => { + const { agent, updates } = createMockAgentWithCapture(); + injectSession(agent, [ + createStreamEvent( + "message_start", + { + model: "claude-haiku-4-5-20251001", + usage: { + input_tokens: 500, + output_tokens: 100, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + }, + "tool_use_123", + ), + createResultMessageWithModel({ + modelUsage: { + "claude-haiku-4-5-20251001": { + inputTokens: 500, + outputTokens: 100, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + webSearchRequests: 0, + costUSD: 0.001, + contextWindow: 200000, + maxOutputTokens: 8192, + }, + }, + }), + { type: "system", subtype: "session_state_changed", state: "idle" }, + ]); + + await agent.prompt({ sessionId: "test-session", prompt: [{ type: "text", text: "test" }] }); + + const usageUpdates = updates.filter((u: any) => u.update?.sessionUpdate === "usage_update"); + expect(usageUpdates).toHaveLength(0); + }); + it("size reflects the current model's context window, not min across all", async () => { const { agent, updates } = createMockAgentWithCapture(); injectSession(agent, [ @@ -2255,6 +2732,7 @@ describe("emitRawSDKMessages", () => { nextPendingOrder: 0, abortController: new AbortController(), emitRawSDKMessages, + contextWindowSize: 200000, }; } diff --git a/src/tests/session-config-options.test.ts b/src/tests/session-config-options.test.ts index 35f75f9c..7d42c9b2 100644 --- a/src/tests/session-config-options.test.ts +++ b/src/tests/session-config-options.test.ts @@ -93,7 +93,10 @@ describe("session config options", () => { cancelled: false, permissionMode: "default", settingsManager: {}, + modes: structuredClone(MOCK_MODES), + models: structuredClone(MOCK_MODELS), configOptions: structuredClone(MOCK_CONFIG_OPTIONS), + contextWindowSize: 200000, }; }