From a924059dafc743d754914f9cb127836e25d06f69 Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Sat, 20 Jun 2026 23:20:10 +0700 Subject: [PATCH] Show live context usage during active turns --- .../providers/claude/agent.redesign.test.ts | 53 +++-- .../agent/providers/claude/agent.test.ts | 28 +++ .../server/agent/providers/claude/agent.ts | 41 +++- .../daemon-e2e/claude-live-usage.e2e.test.ts | 187 ++++++++++++++++++ 4 files changed, 282 insertions(+), 27 deletions(-) create mode 100644 packages/server/src/server/daemon-e2e/claude-live-usage.e2e.test.ts diff --git a/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts b/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts index fd32340e7..83be9f5ba 100644 --- a/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.redesign.test.ts @@ -14,6 +14,7 @@ interface QueryMock { close: ReturnType; setPermissionMode: ReturnType; setModel: ReturnType; + getContextUsage: ReturnType; supportedModels: ReturnType; supportedCommands: ReturnType; rewindFiles: ReturnType; @@ -53,6 +54,7 @@ function createBaseQueryMock(nextImpl: QueryMock["next"]): QueryMock { close: vi.fn(() => undefined), setPermissionMode: vi.fn(async () => undefined), setModel: vi.fn(async () => undefined), + getContextUsage: vi.fn(async () => undefined), supportedModels: vi.fn(async () => [{ value: "opus", displayName: "Opus" }]), supportedCommands: vi.fn(async () => []), rewindFiles: vi.fn(async () => ({ canRewind: true })), @@ -1078,39 +1080,52 @@ test("reuses one autonomous run for unbound stream_event bursts with no foregrou const internal: { turnState: "idle" | "foreground" | "autonomous"; nextTurnOrdinal: number; - routeSdkMessageFromPump: (message: Record) => void; + routeSdkMessageFromPump: ( + message: Record, + activeQuery: QueryMock, + ) => Promise; autonomousTurn: { id: string } | null; } = asInternals(session); + const queryMock = createBaseQueryMock(vi.fn(async () => ({ done: true, value: undefined }))); internal.turnState = "idle"; - internal.routeSdkMessageFromPump({ - type: "stream_event", - event: { - type: "content_block_delta", - delta: { type: "text_delta", text: "AUTO " }, + await internal.routeSdkMessageFromPump( + { + type: "stream_event", + event: { + type: "content_block_delta", + delta: { type: "text_delta", text: "AUTO " }, + }, }, - }); + queryMock, + ); const firstRunId = internal.autonomousTurn?.id ?? null; expect(firstRunId).toBe("autonomous-turn-1"); expect(internal.nextTurnOrdinal).toBe(2); - internal.routeSdkMessageFromPump({ - type: "stream_event", - event: { - type: "content_block_delta", - delta: { type: "text_delta", text: "WAKE" }, + await internal.routeSdkMessageFromPump( + { + type: "stream_event", + event: { + type: "content_block_delta", + delta: { type: "text_delta", text: "WAKE" }, + }, }, - }); + queryMock, + ); expect(internal.autonomousTurn?.id).toBe(firstRunId); expect(internal.nextTurnOrdinal).toBe(2); - internal.routeSdkMessageFromPump({ - type: "result", - subtype: "success", - usage: buildUsage(), - total_cost_usd: 0, - }); + await internal.routeSdkMessageFromPump( + { + type: "result", + subtype: "success", + usage: buildUsage(), + total_cost_usd: 0, + }, + queryMock, + ); expect(internal.autonomousTurn).toBeNull(); await session.close(); 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 a01e12a3d..b156bb961 100644 --- a/packages/server/src/server/agent/providers/claude/agent.test.ts +++ b/packages/server/src/server/agent/providers/claude/agent.test.ts @@ -561,6 +561,7 @@ describe("ClaudeAgentSession features", () => { return: queryReturn, applyFlagSettings: vi.fn(async () => undefined), setModel: vi.fn(async () => undefined), + getContextUsage: vi.fn(async () => undefined), [Symbol.asyncIterator](): AsyncIterator { return { next: async () => { @@ -1046,6 +1047,7 @@ describe("ClaudeAgentSession context window usage", () => { interface QueryFactoryForTurnsOptions { currentContextUsageByTurn?: Array | undefined>; + model?: string; } async function createSessionForTest(): Promise { @@ -1069,6 +1071,7 @@ describe("ClaudeAgentSession context window usage", () => { return await client.createSession({ provider: "claude", cwd: process.cwd(), + model: options?.model, }); } @@ -1353,6 +1356,7 @@ describe("ClaudeAgentSession context window usage", () => { close: () => undefined, setPermissionMode: async () => undefined, setModel: async () => undefined, + getContextUsage: async () => undefined, supportedModels: async () => [], supportedCommands: async () => [ { @@ -1745,6 +1749,30 @@ describe("ClaudeAgentSession context window usage", () => { } }); + test("selected Claude models seed active context window usage with max tokens", async () => { + const session = await createSessionForTurns( + [[createInitMessage(), createMessageStartEvent(), createSuccessResult()]], + { model: "claude-sonnet-4-6" }, + ); + + try { + const events = await collectStreamEvents(session); + + expect(events).toContainEqual( + expect.objectContaining({ + type: "usage_updated", + provider: "claude", + usage: { + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 150, + }, + }), + ); + } finally { + await session.close(); + } + }); + test("message_delta stream events update per-request usage", async () => { const session = await createSessionForTurns([ [ diff --git a/packages/server/src/server/agent/providers/claude/agent.ts b/packages/server/src/server/agent/providers/claude/agent.ts index c6fd80ffc..16d07e336 100644 --- a/packages/server/src/server/agent/providers/claude/agent.ts +++ b/packages/server/src/server/agent/providers/claude/agent.ts @@ -1603,6 +1603,23 @@ function extractContextWindowSize(modelUsage: unknown): number | undefined { return maxContextWindow; } +function resolveInitialContextWindowSize(modelId: string | null | undefined): number | undefined { + const normalized = typeof modelId === "string" ? modelId.trim().toLowerCase() : ""; + if (!normalized) { + return undefined; + } + if (normalized.includes("[1m]") || normalized.includes("context-1m")) { + return 1_000_000; + } + if (normalized.includes("claude-fable-5")) { + return 1_000_000; + } + if (/(?:^|[~/_-])(?:claude[-_ ]*)?(opus|sonnet|haiku)(?:$|[-_ ./])/.test(normalized)) { + return 200_000; + } + return undefined; +} + function readStreamRequestInputTokens(event: Record): number | undefined { const messageUsage = toObjectRecord(toObjectRecord(event.message)?.usage); if (!messageUsage) { @@ -1722,11 +1739,19 @@ class ClaudeContextUsageState { private streamRequestOutputTokens: number | undefined; private completedResultTurns = 0; + constructor(initialContextWindowMaxTokens?: number) { + this.contextWindowMaxTokens = initialContextWindowMaxTokens; + } + beginTurn(): void { this.streamRequestInputTokens = undefined; this.streamRequestOutputTokens = undefined; } + setInitialContextWindowMaxTokens(contextWindowMaxTokens: number | undefined): void { + this.contextWindowMaxTokens = contextWindowMaxTokens; + } + recordModelUsage(modelUsage: unknown): number | undefined { const contextWindowMaxTokens = extractContextWindowSize(modelUsage); if (contextWindowMaxTokens !== undefined) { @@ -1879,7 +1904,7 @@ class ClaudeAgentSession implements AgentSession { private pendingInterruptAbort = false; private foregroundHasVisibleActivity = false; private activeTurnHasAssistantText = false; - private readonly contextUsage = new ClaudeContextUsageState(); + private readonly contextUsage: ClaudeContextUsageState; private userMessageIds: string[] = []; private readonly emittedUserMessageIds = new Set(); private readonly rewindTurnAnchors: ClaudeRewindTurnAnchor[] = []; @@ -1897,6 +1922,9 @@ class ClaudeAgentSession implements AgentSession { this.logger = options.logger.child({ agentId: this.agentId }); this.queryFactory = options.queryFactory; this.resolveBinary = options.resolveBinary; + this.contextUsage = new ClaudeContextUsageState( + resolveInitialContextWindowSize(this.config.model), + ); const handle = options.handle; if (handle) { @@ -2140,6 +2168,9 @@ class ClaudeAgentSession implements AgentSession { if (!claudeModelSupportsFastMode(this.config.model) && this.config.featureValues?.fast_mode) { await this.applyFastModeFeature(false, activeQuery); } + this.contextUsage.setInitialContextWindowMaxTokens( + resolveInitialContextWindowSize(this.config.model), + ); this.lastOptionsModel = normalizedModelId ?? this.lastOptionsModel; this.lastRuntimeModel = null; this.cachedRuntimeInfo = null; @@ -3455,14 +3486,8 @@ class ClaudeAgentSession implements AgentSession { private async queryCurrentContextUsage( activeQuery: Query, ): Promise { - const query = activeQuery as Query & { - getContextUsage?: () => Promise; - }; - if (typeof query.getContextUsage !== "function") { - return undefined; - } try { - const usage = await withTimeout(query.getContextUsage(), 3_000, "timeout"); + const usage = await withTimeout(activeQuery.getContextUsage(), 3_000, "timeout"); return readCurrentContextUsage(usage); } catch (error) { this.logger.debug({ err: error }, "Claude context usage query failed"); diff --git a/packages/server/src/server/daemon-e2e/claude-live-usage.e2e.test.ts b/packages/server/src/server/daemon-e2e/claude-live-usage.e2e.test.ts new file mode 100644 index 000000000..4d4ceb3a5 --- /dev/null +++ b/packages/server/src/server/daemon-e2e/claude-live-usage.e2e.test.ts @@ -0,0 +1,187 @@ +import { mkdtempSync } from "node:fs"; +import { tmpdir } from "node:os"; +import path from "node:path"; +import pino from "pino"; +import { describe, expect, test, vi } from "vitest"; + +import { ClaudeAgentClient } from "../agent/providers/claude/agent.js"; +import { DaemonClient } from "../test-utils/daemon-client.js"; +import { createTestPaseoDaemon } from "../test-utils/paseo-daemon.js"; + +function deferred(): { promise: Promise; resolve: () => void } { + let resolve!: () => void; + const promise = new Promise((innerResolve) => { + resolve = innerResolve; + }); + return { promise, resolve }; +} + +function createControlledClaudeQueryFactory(resultGate: Promise) { + return vi.fn(({ prompt }: { prompt: AsyncIterable }) => { + const queuedMessages: Array> = []; + const waiters: Array<() => void> = []; + let closed = false; + + function wakeNextWaiter() { + waiters.shift()?.(); + } + + function enqueue(message: Record) { + queuedMessages.push(message); + wakeNextWaiter(); + } + + void (async () => { + for await (const _message of prompt) { + enqueue({ + type: "system", + subtype: "init", + session_id: "session-live-usage", + permissionMode: "default", + model: "claude-sonnet-4-6", + }); + enqueue({ + 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-live-usage", + }); + enqueue({ + type: "stream_event", + event: { + type: "message_delta", + usage: { + output_tokens: 25, + }, + }, + session_id: "session-live-usage", + }); + await resultGate; + enqueue({ + type: "result", + subtype: "success", + duration_ms: 100, + duration_api_ms: 75, + is_error: false, + num_turns: 1, + result: "done", + stop_reason: null, + total_cost_usd: 0.25, + usage: { + input_tokens: 10, + cache_read_input_tokens: 5, + output_tokens: 7, + }, + modelUsage: { + "claude-sonnet-4-6": { contextWindow: 200_000 }, + }, + permission_denials: [], + uuid: "result-live-usage", + session_id: "session-live-usage", + }); + break; + } + closed = true; + wakeNextWaiter(); + })(); + + return { + next: vi.fn(async () => { + for (;;) { + if (queuedMessages.length > 0 || closed) { + break; + } + await new Promise((resolve) => { + waiters.push(resolve); + }); + } + if (queuedMessages.length === 0) { + return { done: true, value: undefined }; + } + return { done: false, value: queuedMessages.shift() }; + }), + interrupt: vi.fn(async () => undefined), + return: vi.fn(async () => { + closed = true; + wakeNextWaiter(); + return undefined; + }), + close: vi.fn(() => { + closed = true; + wakeNextWaiter(); + }), + setPermissionMode: vi.fn(async () => undefined), + setModel: vi.fn(async () => undefined), + getContextUsage: vi.fn(async () => undefined), + supportedModels: vi.fn(async () => []), + supportedCommands: vi.fn(async () => []), + rewindFiles: vi.fn(async () => ({ canRewind: true })), + [Symbol.asyncIterator]() { + return this; + }, + }; + }); +} + +describe("daemon E2E (claude live usage)", () => { + test("publishes renderable active context usage through agent_update", async () => { + const logger = pino({ level: "silent" }); + const cwd = mkdtempSync(path.join(tmpdir(), "paseo-claude-live-usage-")); + const resultGate = deferred(); + const daemon = await createTestPaseoDaemon({ + agentClients: { + claude: new ClaudeAgentClient({ + logger, + queryFactory: createControlledClaudeQueryFactory(resultGate.promise), + resolveBinary: async () => "/test/claude/bin", + }), + }, + logger, + }); + const client = new DaemonClient({ url: `ws://127.0.0.1:${daemon.port}/ws` }); + + try { + await client.connect(); + await client.fetchAgents({ subscribe: { subscriptionId: "claude-live-usage" } }); + + const agent = await client.createAgent({ + provider: "claude", + model: "claude-sonnet-4-6", + cwd, + title: "Claude live usage", + }); + + await client.sendMessage(agent.id, "Report live usage"); + + const activeSnapshot = await client.waitForAgentUpsert( + agent.id, + (snapshot) => + snapshot.status === "running" && + snapshot.lastUsage?.contextWindowMaxTokens === 200_000 && + snapshot.lastUsage.contextWindowUsedTokens === 175, + 10_000, + ); + + expect(activeSnapshot.lastUsage).toEqual({ + contextWindowMaxTokens: 200_000, + contextWindowUsedTokens: 175, + }); + + resultGate.resolve(); + const finalState = await client.waitForFinish(agent.id, 10_000); + expect(finalState.status).toBe("idle"); + } finally { + resultGate.resolve(); + await client.close().catch(() => undefined); + await daemon.close(); + } + }); +});