diff --git a/src/agentMode/acp/AcpBackendProcess.test.ts b/src/agentMode/acp/AcpBackendProcess.test.ts index a90ada1c..63245eb6 100644 --- a/src/agentMode/acp/AcpBackendProcess.test.ts +++ b/src/agentMode/acp/AcpBackendProcess.test.ts @@ -471,4 +471,77 @@ describe("AcpBackendProcess", () => { expect(req.additionalDirectories).toEqual(["/abs/context-a", "/abs/context-b"]); }); }); + + describe("prompt-result usage fallback", () => { + // Reach the mock connection's `prompt` jest.fn so a test can stub the + // turn-level `usage` the backend reads after `prompt()` resolves. + function promptMock(backend: AcpBackendProcess): jest.Mock { + return (backend as unknown as { connection: { prompt: jest.Mock } }).connection.prompt; + } + + async function makeBackend(): Promise { + const backend = new AcpBackendProcess( + buildApp(), + buildStubBackend(), + "1.0.0", + buildStubDescriptor() + ); + await backend.start(); + return backend; + } + + it("emits a used-only usage_update from the prompt result when no live update was seen", async () => { + const backend = await makeBackend(); + const handler = jest.fn(); + backend.registerSessionHandler("s1", handler); + promptMock(backend).mockResolvedValueOnce({ + stopReason: "end_turn", + usage: { totalTokens: 4200, inputTokens: 100, outputTokens: 20 }, + }); + + await backend.prompt({ sessionId: "s1", prompt: [] }); + + const usageEvents = handler.mock.calls + .map(([e]) => e) + .filter((e) => e.update.sessionUpdate === "usage_update"); + expect(usageEvents).toHaveLength(1); + expect(usageEvents[0].update.usage).toMatchObject({ + usedTokens: 4200, + inputTokens: 100, + outputTokens: 20, + }); + // Fallback carries no window — AgentSession keeps any prior one. + expect(usageEvents[0].update.usage.contextWindow).toBeUndefined(); + }); + + it("suppresses the prompt-result fallback once a live usage_update has been seen", async () => { + const backend = await makeBackend(); + const handler = jest.fn(); + backend.registerSessionHandler("s1", handler); + const client = getVaultClient(backend); + + // A live usage_update reports current context occupancy (used/size). + await client.sessionUpdate({ + sessionId: "s1", + update: { sessionUpdate: "usage_update", used: 5000, size: 200_000 }, + } as unknown as Parameters[0]); + + // The prompt result carries a cumulative total; it must not overwrite the + // live occupancy figure, so no second usage_update should be emitted. + promptMock(backend).mockResolvedValueOnce({ + stopReason: "end_turn", + usage: { totalTokens: 999_999, inputTokens: 1, outputTokens: 1 }, + }); + await backend.prompt({ sessionId: "s1", prompt: [] }); + + const usageEvents = handler.mock.calls + .map(([e]) => e) + .filter((e) => e.update.sessionUpdate === "usage_update"); + expect(usageEvents).toHaveLength(1); + expect(usageEvents[0].update.usage).toMatchObject({ + usedTokens: 5000, + contextWindow: 200_000, + }); + }); + }); }); diff --git a/src/agentMode/acp/AcpBackendProcess.ts b/src/agentMode/acp/AcpBackendProcess.ts index b08ebdf9..81583605 100644 --- a/src/agentMode/acp/AcpBackendProcess.ts +++ b/src/agentMode/acp/AcpBackendProcess.ts @@ -145,6 +145,13 @@ export class AcpBackendProcess implements BackendProcess { // its sessions, so a bare Set would leak ids across sessions and grow // unbounded for the process lifetime. Pruned on session teardown + shutdown. private readonly todoToolCallIdsBySession = new Map>(); + // Sessions that pushed at least one live `usage_update` notification. Its + // `used` is current context occupancy; the prompt-result `usage.totalTokens` + // is a cumulative session total. Once a live update has been seen we suppress + // the coarser prompt-result fallback so it can't overwrite occupancy with the + // cumulative figure. Keyed by session like the other per-session maps; pruned + // on session teardown + shutdown. + private readonly sawLiveUsage = new Set(); constructor( private readonly app: App, @@ -186,6 +193,7 @@ export class AcpBackendProcess implements BackendProcess { this.pendingUpdates.clear(); this.sessionWireState.clear(); this.todoToolCallIdsBySession.clear(); + this.sawLiveUsage.clear(); this.permissionPrompter = null; this.capabilities.clear(); for (const fn of this.exitListeners) { @@ -290,8 +298,9 @@ export class AcpBackendProcess implements BackendProcess { if (this.domainHandlers.get(sessionId) === handler) { this.domainHandlers.delete(sessionId); // Teardown (not per-turn): the handler is unregistered only when the - // AgentSession disposes, so drop this session's todo-id tracker too. + // AgentSession disposes, so drop this session's per-session trackers too. this.todoToolCallIdsBySession.delete(sessionId); + this.sawLiveUsage.delete(sessionId); } }; } @@ -319,6 +328,32 @@ export class AcpBackendProcess implements BackendProcess { sessionId: sessionIdToAcp(params.sessionId), prompt: promptContentToAcp(params.prompt), }); + // Fallback usage source for agents that never push a live `usage_update` + // notification: the prompt result may carry a turn `usage` with no context + // window. `usage.totalTokens` is a cumulative session total (not current + // context occupancy), so once a live `usage_update` has reported occupancy + // for this session we skip the fallback rather than overwrite the finer + // value. AgentSession's precedence rule then keeps the live window too. + const usage = resp.usage; + if (usage && !this.sawLiveUsage.has(params.sessionId)) { + const handler = this.domainHandlers.get(params.sessionId); + if (handler) { + handler({ + sessionId: params.sessionId, + update: { + sessionUpdate: "usage_update", + usage: { + usedTokens: usage.totalTokens, + inputTokens: usage.inputTokens, + outputTokens: usage.outputTokens, + cacheReadTokens: usage.cachedReadTokens ?? undefined, + cacheWriteTokens: usage.cachedWriteTokens ?? undefined, + updatedAt: Date.now(), + }, + }, + }); + } + } return { stopReason: stopReasonFromAcp(resp.stopReason) }; } @@ -523,6 +558,7 @@ export class AcpBackendProcess implements BackendProcess { this.pendingUpdates.clear(); this.sessionWireState.clear(); this.todoToolCallIdsBySession.clear(); + this.sawLiveUsage.clear(); this.permissionPrompter = null; this.capabilities.clear(); if (this.process) { @@ -587,6 +623,10 @@ export class AcpBackendProcess implements BackendProcess { wire.configOptions = u.configOptions; } } + // Record a live occupancy source so the prompt-result fallback stays quiet. + if (update.update.sessionUpdate === "usage_update") { + this.sawLiveUsage.add(sessionId); + } const handler = this.domainHandlers.get(sessionId); if (!handler) { diff --git a/src/agentMode/acp/wireTranslate.test.ts b/src/agentMode/acp/wireTranslate.test.ts index ff38a027..e6f608e2 100644 --- a/src/agentMode/acp/wireTranslate.test.ts +++ b/src/agentMode/acp/wireTranslate.test.ts @@ -182,3 +182,35 @@ describe("acpNotificationToEvents — todowrite → synthesized plan", () => { }); }); }); + +describe("acpNotificationToEvents — usage_update → SessionUsage", () => { + const FIXED_NOW = 1_700_000_000_000; + let nowSpy: jest.SpyInstance; + + beforeEach(() => { + nowSpy = jest.spyOn(Date, "now").mockReturnValue(FIXED_NOW); + }); + afterEach(() => nowSpy.mockRestore()); + + it("maps size→contextWindow and used→usedTokens, ignoring any cost", () => { + const events = acpNotificationToEvents( + notification({ + sessionUpdate: "usage_update", + size: 200_000, + used: 42_000, + // cost is no longer part of the usage model — it must be dropped. + cost: { amount: 0.1234, currency: "USD" }, + }) + ); + expect(events).toHaveLength(1); + expect(events[0].sessionId).toBe(SESSION_ID); + expect(events[0].update).toEqual({ + sessionUpdate: "usage_update", + usage: { + usedTokens: 42_000, + contextWindow: 200_000, + updatedAt: FIXED_NOW, + }, + }); + }); +}); diff --git a/src/agentMode/acp/wireTranslate.ts b/src/agentMode/acp/wireTranslate.ts index 94157fc2..1a88fe77 100644 --- a/src/agentMode/acp/wireTranslate.ts +++ b/src/agentMode/acp/wireTranslate.ts @@ -25,6 +25,7 @@ import type { ToolCallContent as AcpToolCallContent, ToolCallUpdate, ToolKind as AcpToolKind, + UsageUpdate, } from "@agentclientprotocol/sdk"; import type { AgentPlanEntry, @@ -407,6 +408,17 @@ function acpUpdateToSessionUpdate(update: SessionNotification["update"]): Sessio sessionUpdate: "config_option_update", configOptions: configOptionsFromAcp(update.configOptions) ?? [], }; + case "usage_update": { + const usage = update as UsageUpdate; + return { + sessionUpdate: "usage_update", + usage: { + usedTokens: usage.used, + contextWindow: usage.size, + updatedAt: Date.now(), + }, + }; + } default: // Unknown discriminant — fall back to a benign session_info_update with no title. return { sessionUpdate: "session_info_update", title: null }; diff --git a/src/agentMode/sdk/sdkMessageTranslator.test.ts b/src/agentMode/sdk/sdkMessageTranslator.test.ts index d4733889..d0c04158 100644 --- a/src/agentMode/sdk/sdkMessageTranslator.test.ts +++ b/src/agentMode/sdk/sdkMessageTranslator.test.ts @@ -1,18 +1,75 @@ import type { SDKMessage } from "@anthropic-ai/claude-agent-sdk"; +import type { SessionUsage } from "@/agentMode/session/types"; import { createTranslatorState, mapStopReason, translateSdkMessage } from "./sdkMessageTranslator"; const SESSION_ID = "session-test-1"; +type Uuid = `${string}-${string}-${string}-${string}-${string}`; function streamEvent(event: object): SDKMessage { return { type: "stream_event", event, parent_tool_use_id: null, - uuid: "uuid-1" as `${string}-${string}-${string}-${string}-${string}`, + uuid: "uuid-1" as Uuid, session_id: SESSION_ID, } as SDKMessage; } +interface CallUsage { + input_tokens?: number; + output_tokens?: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; +} + +function assistantMsg( + usage: CallUsage, + opts: { model?: string; parentToolUseId?: string | null } = {} +): SDKMessage { + return { + type: "assistant", + message: { content: [], usage, model: opts.model ?? "claude-test" }, + parent_tool_use_id: opts.parentToolUseId ?? null, + uuid: "uuid-assistant" as Uuid, + session_id: SESSION_ID, + } as unknown as SDKMessage; +} + +interface ModelUsageEntry { + contextWindow: number; + inputTokens?: number; + outputTokens?: number; + cacheReadInputTokens?: number; + cacheCreationInputTokens?: number; +} + +function resultMsg(opts: { + usage?: CallUsage; + modelUsage?: Record; + total_cost_usd?: number; +}): SDKMessage { + return { + type: "result", + subtype: "success", + duration_ms: 0, + duration_api_ms: 0, + is_error: false, + num_turns: 1, + result: "ok", + stop_reason: "end_turn", + total_cost_usd: opts.total_cost_usd ?? 0, + usage: opts.usage ?? {}, + modelUsage: opts.modelUsage ?? {}, + permission_denials: [], + uuid: "uuid-result" as Uuid, + session_id: SESSION_ID, + } as unknown as SDKMessage; +} + +function usageOf(out: ReturnType): SessionUsage { + return (out[0].update as { usage: SessionUsage }).usage; +} + describe("translateSdkMessage", () => { it("emits agent_message_chunk for text deltas", () => { const state = createTranslatorState(); @@ -121,32 +178,131 @@ describe("translateSdkMessage", () => { expect(state.toolUseBlocks.size).toBe(0); }); - it("returns [] for `result` (caller resolves the prompt promise separately)", () => { + it("reports occupancy from the last assistant message, not the cumulative result total", () => { const state = createTranslatorState(); - const out = translateSdkMessage( - { - type: "result", - subtype: "success", - duration_ms: 0, - duration_api_ms: 0, - is_error: false, - num_turns: 1, - result: "ok", - stop_reason: "end_turn", - total_cost_usd: 0, - // eslint-disable-next-line @typescript-eslint/no-explicit-any - usage: {} as any, - modelUsage: {}, - permission_denials: [], - uuid: "uuid-2" as `${string}-${string}-${string}-${string}-${string}`, - session_id: SESSION_ID, - }, + // Two tool-loop iterations. The SDK result SUMS these across the turn, but + // occupancy is only the final call's own prompt + reply. + translateSdkMessage( + assistantMsg({ input_tokens: 100, cache_creation_input_tokens: 10_000, output_tokens: 50 }), SESSION_ID, state ); + translateSdkMessage( + assistantMsg({ input_tokens: 200, cache_read_input_tokens: 10_050, output_tokens: 80 }), + SESSION_ID, + state + ); + const out = translateSdkMessage( + resultMsg({ + // Cumulative aggregate — deliberately larger than occupancy. + usage: { + input_tokens: 300, + cache_read_input_tokens: 10_050, + cache_creation_input_tokens: 10_000, + output_tokens: 130, + }, + modelUsage: { "claude-test": { contextWindow: 200_000 } }, + }), + SESSION_ID, + state + ); + expect(out).toHaveLength(1); + expect(out[0].update.sessionUpdate).toBe("usage_update"); + const usage = usageOf(out); + // Final call: 200 + 10_050 + 0 + 80 = 10_330 (occupancy), NOT the cumulative + // 300 + 10_050 + 10_000 + 130 = 20_480. + expect(usage.usedTokens).toBe(10_330); + expect(usage.contextWindow).toBe(200_000); + }); + + it("uses the active model's context window, not the largest in a multi-model turn", () => { + const state = createTranslatorState(); + translateSdkMessage( + assistantMsg({ input_tokens: 1000, output_tokens: 50 }, { model: "main" }), + SESSION_ID, + state + ); + const out = translateSdkMessage( + resultMsg({ + modelUsage: { + // A subagent used a larger window; the ring must divide by the MAIN + // model's window, not the max across the turn. + main: { contextWindow: 200_000 }, + sub: { contextWindow: 1_000_000 }, + }, + }), + SESSION_ID, + state + ); + expect(usageOf(out).contextWindow).toBe(200_000); + }); + + it("ignores subagent assistant usage when sampling occupancy", () => { + const state = createTranslatorState(); + // Subagent turn (parent_tool_use_id set) — a different context; must not + // become the occupancy sample. + translateSdkMessage( + assistantMsg({ input_tokens: 99_999, output_tokens: 0 }, { parentToolUseId: "t1" }), + SESSION_ID, + state + ); + translateSdkMessage( + assistantMsg({ input_tokens: 1000, output_tokens: 20 }, { model: "main" }), + SESSION_ID, + state + ); + const out = translateSdkMessage( + resultMsg({ modelUsage: { main: { contextWindow: 200_000 } } }), + SESSION_ID, + state + ); + expect(usageOf(out).usedTokens).toBe(1020); + }); + + it("emits no usage_update when the turn produced no top-level assistant message", () => { + const state = createTranslatorState(); + const out = translateSdkMessage(resultMsg({}), SESSION_ID, state); expect(out).toEqual([]); }); + it("falls back to the dominant model's window for a synthetic assistant turn", () => { + const state = createTranslatorState(); + // A "" model id keys into no modelUsage entry; the window comes + // from the model that carried the conversation (most tokens), never an aux + // model that happens to have a different window. + translateSdkMessage( + assistantMsg( + { input_tokens: 500, cache_read_input_tokens: 4000, output_tokens: 40 }, + { model: "" } + ), + SESSION_ID, + state + ); + const out = translateSdkMessage( + resultMsg({ + modelUsage: { + "claude-opus-4-8[1m]": { + contextWindow: 1_000_000, + inputTokens: 5000, + cacheReadInputTokens: 400_000, + cacheCreationInputTokens: 100_000, + outputTokens: 6000, + }, + "claude-haiku-4-5-20251001": { + contextWindow: 200_000, + inputTokens: 300, + cacheReadInputTokens: 0, + cacheCreationInputTokens: 0, + outputTokens: 12, + }, + }, + }), + SESSION_ID, + state + ); + expect(usageOf(out).contextWindow).toBe(1_000_000); + }); + it("ignores assistant messages whose tool_use blocks were already streamed", () => { const state = createTranslatorState(); // Pretend the streaming path already saw this tool_use. @@ -784,3 +940,112 @@ describe("session todo-list normalization (TodoWrite / Task tools → plan)", () }); }); }); + +function resultMessage(overrides: { + usage: { + input_tokens: number; + output_tokens: number; + cache_read_input_tokens: number; + cache_creation_input_tokens: number; + }; + modelUsage?: Record; + total_cost_usd?: number; +}): SDKMessage { + return { + type: "result", + subtype: "success", + usage: overrides.usage, + modelUsage: overrides.modelUsage ?? {}, + total_cost_usd: overrides.total_cost_usd ?? 0, + session_id: SESSION_ID, + } as unknown as SDKMessage; +} + +describe("translateSdkMessage — result → usage_update", () => { + const FIXED_NOW = 1_700_000_000_000; + let nowSpy: jest.SpyInstance; + + beforeEach(() => { + nowSpy = jest.spyOn(Date, "now").mockReturnValue(FIXED_NOW); + }); + afterEach(() => nowSpy.mockRestore()); + + it("matches the bare assistant model id to the suffixed modelUsage key", () => { + const state = createTranslatorState(); + // Real runtime shape: the assistant message reports the bare id + // "claude-opus-4-8" while the result keys it "claude-opus-4-8[1m]". An exact + // lookup misses — a prefix match must recover the main model's 1M window and + // NOT fall to the smaller-windowed aux model. Final call occupancy: + // 100 + 5000 + 300 + 20 = 5420. + translateSdkMessage( + assistantMsg( + { + input_tokens: 100, + output_tokens: 20, + cache_read_input_tokens: 5000, + cache_creation_input_tokens: 300, + }, + { model: "claude-opus-4-8" } + ), + SESSION_ID, + state + ); + const out = translateSdkMessage( + resultMessage({ + // Cumulative result usage — deliberately huge; now ignored for occupancy. + usage: { + input_tokens: 9999, + output_tokens: 9999, + cache_read_input_tokens: 9999, + cache_creation_input_tokens: 9999, + }, + modelUsage: { + "claude-opus-4-8[1m]": { contextWindow: 1_000_000 }, + "claude-haiku-4-5-20251001": { contextWindow: 200_000 }, + }, + }), + SESSION_ID, + state + ); + expect(out).toEqual([ + { + sessionId: SESSION_ID, + update: { + sessionUpdate: "usage_update", + usage: { + usedTokens: 5420, + contextWindow: 1_000_000, + inputTokens: 100, + outputTokens: 20, + cacheReadTokens: 5000, + cacheWriteTokens: 300, + updatedAt: FIXED_NOW, + }, + }, + }, + ]); + }); + + it("emits usedTokens with undefined contextWindow when modelUsage lacks the model", () => { + const state = createTranslatorState(); + translateSdkMessage(assistantMsg({ input_tokens: 10, output_tokens: 2 }), SESSION_ID, state); + const out = translateSdkMessage( + resultMessage({ + usage: { + input_tokens: 0, + output_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, + }, + // No matching model entry → no window: count-only, never a wrong ring. + }), + SESSION_ID, + state + ); + expect(out).toHaveLength(1); + const update = out[0].update as { sessionUpdate: string; usage: Record }; + expect(update.sessionUpdate).toBe("usage_update"); + expect(update.usage.usedTokens).toBe(12); + expect(update.usage.contextWindow).toBeUndefined(); + }); +}); diff --git a/src/agentMode/sdk/sdkMessageTranslator.ts b/src/agentMode/sdk/sdkMessageTranslator.ts index 75470704..7f60d5e1 100644 --- a/src/agentMode/sdk/sdkMessageTranslator.ts +++ b/src/agentMode/sdk/sdkMessageTranslator.ts @@ -11,6 +11,7 @@ import type { SessionEvent, SessionId, SessionUpdate, + SessionUsage, ToolCallContent, } from "@/agentMode/session/types"; import { resolveToolName } from "@/agentMode/session/toolName"; @@ -44,6 +45,25 @@ export interface TranslatorState { emittedToolUseIds: Set; /** Session-lived todo/Task accumulator (see claudeTodoPlan.ts). */ claudeTasks: ClaudeTaskPlanState; + /** + * Occupancy sample from the most recent TOP-LEVEL assistant message this turn + * (subagent messages excluded). The `result` message's aggregate `usage` sums + * every API call in the turn — tool loops re-read the whole context from cache + * each iteration — so it overstates current context; the last main-model + * response's own per-call usage is the true occupancy. Reset per query. + */ + lastAssistantUsage?: AssistantUsageSample; +} + +/** Per-call token occupancy captured from one assistant message. */ +interface AssistantUsageSample { + usedTokens: number; + inputTokens: number; + outputTokens: number; + cacheReadTokens: number; + cacheWriteTokens: number; + /** Model that produced this turn — keys into the result's `modelUsage` for the window. */ + model: string; } export function createTranslatorState(claudeTasks?: ClaudeTaskPlanState): TranslatorState { @@ -77,11 +97,73 @@ export function translateSdkMessage( case "user": return translateUserMessage(msg, sessionId, state); case "result": + return translateResultMessage(msg, sessionId, state); default: return []; } } +/** + * A `result` closes a turn. Its aggregate `usage` sums every API call in the + * turn (each tool-loop iteration re-reads the whole context from cache), so it + * is a cumulative bill, not current context occupancy — dividing it by the + * window would peg the meter to 100% after a tool-heavy turn even when the live + * context still fits. We instead report the last top-level assistant message's + * own per-call usage ({@link TranslatorState.lastAssistantUsage}) as occupancy, + * paired with THAT model's window from `modelUsage`. + */ +function translateResultMessage( + msg: SDKResultMessage, + sessionId: SessionId, + state: TranslatorState +): SessionEvent[] { + const sample = state.lastAssistantUsage; + // No main-model turn to measure (e.g. an errored/empty result): leave the + // meter on its prior occupancy rather than invent a cumulative number. + if (!sample) return []; + + const sessionUsage: SessionUsage = { + usedTokens: sample.usedTokens, + contextWindow: windowForModel(msg.modelUsage, sample.model), + inputTokens: sample.inputTokens, + outputTokens: sample.outputTokens, + cacheReadTokens: sample.cacheReadTokens, + cacheWriteTokens: sample.cacheWriteTokens, + updatedAt: Date.now(), + }; + return [event(sessionId, { sessionUpdate: "usage_update", usage: sessionUsage })]; +} + +/** + * The active model's context window — the one that produced the occupancy + * sample. The result keys `modelUsage` with a context-variant/date suffix (e.g. + * `claude-opus-4-8[1m]`, `claude-haiku-4-5-20251001`) while the assistant + * message reports the bare id (`claude-opus-4-8`), so an exact lookup misses on + * real data — match on prefix too. Falling back to the model that accumulated + * the most tokens keeps this correct for a bare `` turn and avoids + * ever taking the max window (a larger-windowed aux/subagent model would + * otherwise deflate the main conversation's percentage). + */ +function windowForModel( + modelUsage: SDKResultMessage["modelUsage"], + sampleModel: string +): number | undefined { + const entries = Object.entries(modelUsage); + if (entries.length === 0) return undefined; + if (sampleModel) { + const exact = modelUsage[sampleModel]; + if (exact) return exact.contextWindow; + const prefixed = entries.find(([id]) => id.startsWith(sampleModel)); + if (prefixed) return prefixed[1].contextWindow; + } + const dominant = entries.reduce((a, b) => (modelTokens(b[1]) > modelTokens(a[1]) ? b : a)); + return dominant[1].contextWindow; +} + +function modelTokens(m: SDKResultMessage["modelUsage"][string]): number { + return m.inputTokens + m.outputTokens + m.cacheReadInputTokens + m.cacheCreationInputTokens; +} + export function mapStopReason(msg: SDKResultMessage): "end_turn" | "cancelled" | "refusal" { if (msg.subtype === "success") return "end_turn"; return "cancelled"; @@ -255,9 +337,24 @@ function translateAssistantMessage( state: TranslatorState ): SessionEvent[] { const out: SessionEvent[] = []; - const content = (msg.message as { content?: unknown }).content; - if (!Array.isArray(content)) return out; + const message = msg.message as { + content?: unknown; + model?: string; + usage?: { + input_tokens?: number; + output_tokens?: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + }; + }; const parentToolUseId = msg.parent_tool_use_id ?? undefined; + // Sample occupancy from the main agent's own response only; a subagent's + // per-call usage measures a different context and must not drive the meter. + if (parentToolUseId === undefined && message.usage) { + state.lastAssistantUsage = assistantUsageSample(message.usage, message.model); + } + const content = message.content; + if (!Array.isArray(content)) return out; for (const block of content) { const b = block as { type?: string; id?: string; name?: string; input?: unknown }; if (b.type !== "tool_use" || !b.id || !b.name) continue; @@ -272,6 +369,36 @@ function translateAssistantMessage( return out; } +/** + * Occupancy for one API call = its full prompt (fresh input + both cache buckets) + * plus the reply it generated. On a cached turn most input arrives as + * `cache_read`, so all three input buckets must be summed to recover the prompt + * size. This is the same formula the SDK result uses — the fix is the *source*: + * one call's usage (occupancy), not the turn's summed total (cumulative). + */ +function assistantUsageSample( + usage: { + input_tokens?: number; + output_tokens?: number; + cache_read_input_tokens?: number; + cache_creation_input_tokens?: number; + }, + model: string | undefined +): AssistantUsageSample { + const inputTokens = usage.input_tokens ?? 0; + const outputTokens = usage.output_tokens ?? 0; + const cacheReadTokens = usage.cache_read_input_tokens ?? 0; + const cacheWriteTokens = usage.cache_creation_input_tokens ?? 0; + return { + usedTokens: inputTokens + cacheReadTokens + cacheWriteTokens + outputTokens, + inputTokens, + outputTokens, + cacheReadTokens, + cacheWriteTokens, + model: model ?? "", + }; +} + /** * Session todo-list normalization (claudeTodoPlan.ts): feed native, TOP-LEVEL * TodoWrite / TaskCreate / TaskUpdate calls into the session accumulator and diff --git a/src/agentMode/session/AgentChatBackend.ts b/src/agentMode/session/AgentChatBackend.ts index ac605ce3..5238b8bf 100644 --- a/src/agentMode/session/AgentChatBackend.ts +++ b/src/agentMode/session/AgentChatBackend.ts @@ -10,6 +10,7 @@ import type { PermissionPrompt, PlanDecisionAction, PromptContent, + SessionUsage, } from "./types"; /** @@ -95,6 +96,13 @@ export interface AgentChatBackend { */ getCurrentTodoList(): AgentTodoListEntry[] | null; + /** + * Latest backend-agnostic token-usage snapshot for the session, or `null` + * when none has been reported yet (fresh session, or a resumed chat with no + * persisted usage). The UI renders this as a context-window indicator. + */ + getSessionUsage(): SessionUsage | null; + /** * True when an ExitPlanMode permission is currently pending. The chat input * disables itself while one is outstanding so the user is funneled to the diff --git a/src/agentMode/session/AgentChatPersistenceManager.test.ts b/src/agentMode/session/AgentChatPersistenceManager.test.ts index 1b0dfb25..89fde06f 100644 --- a/src/agentMode/session/AgentChatPersistenceManager.test.ts +++ b/src/agentMode/session/AgentChatPersistenceManager.test.ts @@ -1,9 +1,11 @@ /* eslint-disable obsidianmd/no-tfile-tfolder-cast -- test fixtures; not real TFiles */ import { AI_SENDER, USER_SENDER } from "@/constants"; +import { readFrontmatterViaAdapter } from "@/utils/vaultAdapterUtils"; import { AgentChatPersistenceManager } from "./AgentChatPersistenceManager"; import { GLOBAL_SCOPE } from "./scope"; import type { AgentChatMessage } from "./types"; -import type { App, TFile } from "obsidian"; +import { TFile } from "obsidian"; +import type { App } from "obsidian"; jest.mock("obsidian", () => ({ Notice: jest.fn(), @@ -282,4 +284,90 @@ describe("AgentChatPersistenceManager", () => { expect(loaded.projectId).toBe(GLOBAL_SCOPE); }); }); + + describe("usage frontmatter", () => { + afterEach(() => { + // Restore the default no-metadata behavior for the adapter helper so a + // per-test override (round-trip-on-omit) doesn't leak into other suites. + (readFrontmatterViaAdapter as jest.Mock).mockResolvedValue(null); + }); + + it("round-trips a SessionUsage snapshot through save/load", async () => { + const messages = [makeMessage(USER_SENDER, "hi")]; + const usage = { + usedTokens: 42_000, + contextWindow: 200_000, + inputTokens: 100, + outputTokens: 20, + cacheReadTokens: 5000, + cacheWriteTokens: 300, + updatedAt: 1_700_000_000_000, + }; + const saved = await manager.saveSession(messages, "claude", { usage }); + const raw = app.files.get(saved!.path)!.contents!; + expect(raw).toContain(`usage: '${JSON.stringify(usage)}'`); + + const loaded = await manager.loadFile(app.files.get(saved!.path) as unknown as TFile); + expect(loaded.usage).toEqual(usage); + }); + + it("round-trips the persisted usage when a later save omits it", async () => { + const messages = [makeMessage(USER_SENDER, "hi")]; + const usage = { usedTokens: 5000, contextWindow: 200_000, updatedAt: 1 }; + const first = await manager.saveSession(messages, "claude", { usage }); + // `resolveExistingFile` gates on `instanceof TFile`; give the stored fake + // the mocked prototype so the resave takes the existing-file path (where + // usage round-trips) instead of treating it as a brand-new write. + Object.setPrototypeOf(app.files.get(first!.path)!, TFile.prototype); + // Mirror production: `readExistingMeta` reads the prior file's frontmatter + // to round-trip fields the caller didn't re-supply. The default mock + // returns null (no metadata), so parse the stored file here — quote-strip + // matches the real adapter helper so the JSON value comes back intact. + (readFrontmatterViaAdapter as jest.Mock).mockImplementation(async (_app, path: string) => { + const raw = app.files.get(path)?.contents ?? ""; + const yaml = raw.match(/^---\n([\s\S]*?)\n---/)?.[1]; + if (!yaml) return null; + const fm: Record = {}; + for (const line of yaml.split("\n")) { + const m = line.match(/^([\w-]+):\s*(.+)/); + if (m) fm[m[1]] = m[2].trim().replace(/^["']|["']$/g, ""); + } + return fm; + }); + // A save with no usage option must not drop the stored snapshot. + const second = await manager.saveSession(messages, "claude", { + existingPath: first!.path, + }); + const loaded = await manager.loadFile(app.files.get(second!.path) as unknown as TFile); + expect(loaded.usage).toEqual(usage); + }); + + it("leaves usage undefined for a chat saved without it", async () => { + const saved = await manager.saveSession([makeMessage(USER_SENDER, "hi")], "claude"); + const raw = app.files.get(saved!.path)!.contents!; + expect(raw).not.toContain("usage:"); + const loaded = await manager.loadFile(app.files.get(saved!.path) as unknown as TFile); + expect(loaded.usage).toBeUndefined(); + }); + + it("ignores malformed usage JSON instead of failing the load", async () => { + const path = "test-folder/agent__badusage.md"; + await app.vault.adapter.write( + path, + [ + "---", + "epoch: 1735732800000", + "mode: agent", + "backendId: claude", + "usage: 'not-json{'", + "---", + "", + "**user**: hi", + ].join("\n") + ); + const loaded = await manager.loadFile(app.files.get(path) as unknown as TFile); + expect(loaded.usage).toBeUndefined(); + expect(loaded.backendId).toBe("claude"); + }); + }); }); diff --git a/src/agentMode/session/AgentChatPersistenceManager.ts b/src/agentMode/session/AgentChatPersistenceManager.ts index 847df394..2af404d5 100644 --- a/src/agentMode/session/AgentChatPersistenceManager.ts +++ b/src/agentMode/session/AgentChatPersistenceManager.ts @@ -22,11 +22,34 @@ import { TFile, type App } from "obsidian"; import { Notice } from "obsidian"; import { coerceProjectId, escapeYamlString, unescapeYamlString } from "./agentChatYaml"; import { GLOBAL_SCOPE } from "./scope"; -import type { AgentChatMessage, BackendId } from "./types"; +import type { AgentChatMessage, BackendId, SessionUsage } from "./types"; const SAFE_FILENAME_BYTE_LIMIT = 100; export const AGENT_FILENAME_PREFIX = "agent__"; +/** + * Parse the frontmatter `usage` field (a JSON string) back into a + * {@link SessionUsage}. Returns `undefined` for absent, non-string, malformed, + * or wrong-shaped values so a corrupt frontmatter never rejects the whole load. + */ +function parseUsageJson(raw: unknown): SessionUsage | undefined { + if (typeof raw !== "string" || raw.trim().length === 0) return undefined; + try { + const parsed = JSON.parse(raw) as unknown; + if ( + typeof parsed === "object" && + parsed !== null && + typeof (parsed as SessionUsage).usedTokens === "number" && + typeof (parsed as SessionUsage).updatedAt === "number" + ) { + return parsed as SessionUsage; + } + } catch { + // Malformed JSON — fall through to undefined. + } + return undefined; +} + /** * Result of `loadFile` — restores display-only Agent Mode messages plus * routing info needed to spawn the right backend session. @@ -51,6 +74,12 @@ export interface LoadedAgentChat { * authority. */ projectId: string; + /** + * Latest token-usage snapshot captured at save time, or `undefined` for chats + * saved before usage was wired up (or with malformed usage frontmatter). Lets + * a resumed session show its last-known usage before the next turn. + */ + usage?: SessionUsage; } interface ExistingMeta { @@ -59,6 +88,7 @@ interface ExistingMeta { lastAccessedAt?: number; sessionId?: string; projectId?: string; + usage?: SessionUsage; } /** @@ -102,6 +132,8 @@ export class AgentChatPersistenceManager { * global chats. */ projectId?: string; + /** Latest token-usage snapshot to persist for resume. */ + usage?: SessionUsage; } ): Promise<{ path: string } | null> { if (messages.length === 0) return null; @@ -136,6 +168,9 @@ export class AgentChatPersistenceManager { // demotes itself to global on a later save. Coerce first so a blank // option falls through to the existing scope instead of clobbering it. projectId: coerceProjectId(options?.projectId) ?? existingMeta.projectId, + // Round-trip the persisted usage when the caller doesn't re-supply it, + // so a save that isn't triggered by a usage change keeps the snapshot. + usage: options?.usage ?? existingMeta.usage, }); if (existingFile && isInVaultCache(this.app, existingFile.path)) { @@ -205,12 +240,13 @@ export class AgentChatPersistenceManager { // HARD CONTRACT: absent/blank projectId → GLOBAL_SCOPE, so legacy `agent__` // chats stay in the global history. Never inferred from the filename. const projectId = frontmatter.projectId?.trim() || GLOBAL_SCOPE; + const usage = parseUsageJson(frontmatter.usage); const messages = this.parseChatBody(body); logInfo( `[AgentChatPersistenceManager] Loaded ${messages.length} messages from ${file.path} (backend=${backendId}, sessionId=${sessionId ?? "none"}, projectId=${projectId})` ); - return { messages, backendId, topic, label, sessionId, projectId }; + return { messages, backendId, topic, label, sessionId, projectId, usage }; } /** @@ -259,6 +295,7 @@ export class AgentChatPersistenceManager { typeof cached.lastAccessedAt === "number" ? cached.lastAccessedAt : undefined, sessionId: typeof cached.sessionId === "string" ? cached.sessionId : undefined, projectId: coerceProjectId(cached.projectId), + usage: parseUsageJson(cached.usage), }; } try { @@ -271,6 +308,7 @@ export class AgentChatPersistenceManager { lastAccessedAt: lastAccessed && Number.isFinite(lastAccessed) ? lastAccessed : undefined, sessionId: typeof fm.sessionId === "string" ? fm.sessionId : undefined, projectId: coerceProjectId(fm.projectId), + usage: parseUsageJson(fm.usage), }; } catch { return {}; @@ -452,6 +490,7 @@ export class AgentChatPersistenceManager { lastAccessedAt?: number; sessionId?: string | null; projectId?: string; + usage?: SessionUsage; }): string { const settings = getSettings(); const lines: string[] = [ @@ -472,6 +511,10 @@ export class AgentChatPersistenceManager { if (args.label) lines.push(`agentLabel: "${escapeYamlString(args.label)}"`); if (args.modelKey) lines.push(`modelKey: "${escapeYamlString(args.modelKey)}"`); if (args.lastAccessedAt) lines.push(`lastAccessedAt: ${args.lastAccessedAt}`); + // Serialized as a single-quoted JSON string: JSON.stringify emits no single + // quotes or raw control chars, so the value round-trips through the YAML + // parser (and our hand-rolled splitFrontmatter) verbatim. + if (args.usage) lines.push(`usage: '${JSON.stringify(args.usage)}'`); lines.push("tags:"); lines.push(` - ${settings.defaultConversationTag}`); lines.push("---"); diff --git a/src/agentMode/session/AgentChatUIState.ts b/src/agentMode/session/AgentChatUIState.ts index a09e41a9..ed2fee45 100644 --- a/src/agentMode/session/AgentChatUIState.ts +++ b/src/agentMode/session/AgentChatUIState.ts @@ -12,6 +12,7 @@ import type { PermissionPrompt, PlanDecisionAction, PromptContent, + SessionUsage, } from "@/agentMode/session/types"; import type { MessageContext } from "@/types/message"; @@ -158,6 +159,10 @@ export class AgentChatUIState implements AgentChatBackend { return this.session.getCurrentTodoList(); } + getSessionUsage(): SessionUsage | null { + return this.session.getSessionUsage(); + } + async resolvePlanProposal( proposalId: string, decision: PlanDecisionAction, diff --git a/src/agentMode/session/AgentSession.test.ts b/src/agentMode/session/AgentSession.test.ts index 9c24805a..1ffb0ec8 100644 --- a/src/agentMode/session/AgentSession.test.ts +++ b/src/agentMode/session/AgentSession.test.ts @@ -385,6 +385,134 @@ describe("AgentSession.restoreLabel", () => { }); }); +describe("AgentSession session usage", () => { + function makeSession(mock: ReturnType) { + return new AgentSession({ + backend: mock.asBackend, + backendSessionId: "acp-1", + internalId: "internal-1", + backendId: "opencode", + }); + } + + it("starts null and stores an incoming usage_update, notifying subscribers", () => { + const mock = makeMockBackend(); + const session = makeSession(mock); + const onMessagesChanged = jest.fn(); + session.subscribe({ onMessagesChanged, onStatusChanged: () => {} }); + expect(session.getSessionUsage()).toBeNull(); + + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 5000, contextWindow: 200_000, updatedAt: 1 }, + }, + }); + + expect(session.getSessionUsage()).toEqual({ + usedTokens: 5000, + contextWindow: 200_000, + updatedAt: 1, + }); + expect(onMessagesChanged).toHaveBeenCalled(); + }); + + it("handles usage without a placeholder (session-scoped, no active turn)", () => { + const mock = makeMockBackend(); + const session = makeSession(mock); + // No sendPrompt → no placeholder. A session-scoped usage must still land. + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 42, updatedAt: 7 }, + }, + }); + expect(session.getSessionUsage()?.usedTokens).toBe(42); + }); + + it("ignores a used-only fallback once an occupancy snapshot exists", () => { + const mock = makeMockBackend(); + const session = makeSession(mock); + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 5000, contextWindow: 200_000, updatedAt: 1 }, + }, + }); + // A later used-only snapshot (no contextWindow) is an ACP prompt-result + // fallback carrying cumulative lifetime totals, not occupancy. It must not + // override the occupancy snapshot — pairing its tokens with the window would + // show a bogus ring. A later live occupancy update would supersede it. + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 6000, updatedAt: 2 }, + }, + }); + expect(session.getSessionUsage()).toEqual({ + usedTokens: 5000, + contextWindow: 200_000, + updatedAt: 1, + }); + }); + + it("lets a fuller snapshot's contextWindow replace an earlier one", () => { + const mock = makeMockBackend(); + const session = makeSession(mock); + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 5000, contextWindow: 200_000, updatedAt: 1 }, + }, + }); + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 6000, contextWindow: 1_000_000, updatedAt: 2 }, + }, + }); + expect(session.getSessionUsage()?.contextWindow).toBe(1_000_000); + }); + + it("seeds usage from persisted history and re-notifies", () => { + const mock = makeMockBackend(); + const session = makeSession(mock); + const onMessagesChanged = jest.fn(); + session.subscribe({ onMessagesChanged, onStatusChanged: () => {} }); + + session.seedSessionUsage({ usedTokens: 1234, contextWindow: 200_000, updatedAt: 9 }); + expect(session.getSessionUsage()).toEqual({ + usedTokens: 1234, + contextWindow: 200_000, + updatedAt: 9, + }); + expect(onMessagesChanged).toHaveBeenCalled(); + + // A live full snapshot supersedes the seed. + mock.emit({ + sessionId: "acp-1", + update: { + sessionUpdate: "usage_update", + usage: { usedTokens: 2000, contextWindow: 200_000, updatedAt: 10 }, + }, + }); + expect(session.getSessionUsage()?.usedTokens).toBe(2000); + }); + + it("seedSessionUsage with undefined is a no-op", () => { + const mock = makeMockBackend(); + const session = makeSession(mock); + session.seedSessionUsage(undefined); + expect(session.getSessionUsage()).toBeNull(); + }); +}); + describe("buildUserDisplayContent", () => { it("returns undefined when there are no images", () => { expect(buildUserDisplayContent("hi")).toBeUndefined(); diff --git a/src/agentMode/session/AgentSession.ts b/src/agentMode/session/AgentSession.ts index fab55303..5f52bd62 100644 --- a/src/agentMode/session/AgentSession.ts +++ b/src/agentMode/session/AgentSession.ts @@ -27,6 +27,7 @@ import { PromptInput, SessionEvent, SessionId, + SessionUsage, StopReason, ToolCallContent, ToolCallDelta, @@ -398,6 +399,10 @@ export class AgentSession { // Signature of the last applied list — multiple equal plan updates (e.g. // opencode's synthesized + occasional real plan channel) must not re-notify. private currentTodoListSignature: string | null = null; + // Latest backend-agnostic token-usage snapshot, or null until the first + // `usage_update` (or a persisted snapshot seeded on resume). Session-scoped: + // not tied to any turn placeholder. + private currentUsage: SessionUsage | null = null; // Monotonic counter for `currentPlan.id` so the React tree can detect a // *new* plan-mode review (vs. an in-place revision that bumps `revision`). private planSeq = 0; @@ -1318,6 +1323,38 @@ export class AgentSession { return this.currentTodoList; } + /** Latest token-usage snapshot, or `null` when the session has none yet. */ + getSessionUsage(): SessionUsage | null { + return this.currentUsage; + } + + /** + * Seed the usage snapshot from persisted frontmatter on resume, so a reopened + * chat shows its last-known usage immediately instead of blank-until-next-turn. + * A live `usage_update` later supersedes it via {@link applyUsageUpdate}. + */ + seedSessionUsage(usage: SessionUsage | undefined): void { + if (!usage) return; + this.currentUsage = usage; + this.notifyMessages(); + } + + /** + * Apply an incoming usage snapshot. A snapshot without a `contextWindow` is a + * cumulative/count-only fallback (e.g. ACP's prompt-result totals, which count + * lifetime tokens across the session), not current-context occupancy. Once we + * already hold an occupancy snapshot (one that carries a window), ignore the + * windowless one rather than pair its lifetime tokens with that window and + * render a bogus percentage ring. A later live occupancy update supersedes it. + */ + private applyUsageUpdate(usage: SessionUsage): void { + if (usage.contextWindow === undefined && this.currentUsage?.contextWindow !== undefined) { + return; + } + this.currentUsage = usage; + this.notifyMessages(); + } + /** * Drop the current plan once the user has decided. The UI gates the card * render on `decision === "pending"`, so a terminal state is never visible @@ -1555,6 +1592,10 @@ export class AgentSession { // Same — the `state_changed` follow-up carries the recomputed state. return; } + if (update.sessionUpdate === "usage_update") { + this.applyUsageUpdate(update.usage); + return; + } // Content chunks can trail past the prompt result on some backends (the // result is flushed before the turn's final chunks — opencode + fast diff --git a/src/agentMode/session/AgentSessionManager.test.ts b/src/agentMode/session/AgentSessionManager.test.ts index 993796ce..b588d059 100644 --- a/src/agentMode/session/AgentSessionManager.test.ts +++ b/src/agentMode/session/AgentSessionManager.test.ts @@ -173,6 +173,7 @@ function makeMockSession(overrides: { setConfigOption: jest.fn(), getLabel: () => null, setLabel: jest.fn(), + getSessionUsage: () => null, subscribe: (l: Parameters[0]) => { listeners.add(l); return () => listeners.delete(l); diff --git a/src/agentMode/session/AgentSessionManager.ts b/src/agentMode/session/AgentSessionManager.ts index b00754c3..6cff962a 100644 --- a/src/agentMode/session/AgentSessionManager.ts +++ b/src/agentMode/session/AgentSessionManager.ts @@ -2502,6 +2502,7 @@ export class AgentSessionManager { } session.loadDisplayMessages(loaded.messages); + session.seedSessionUsage(loaded.usage); if (loaded.label) session.setLabel(loaded.label); this.getSessionState(session.internalId).path = file.path; if (loaded.sessionId) { @@ -2899,6 +2900,7 @@ export class AgentSessionManager { const label = session.getLabel(); const sessionId = session.getBackendSessionId(); + const usage = session.getSessionUsage(); // Skip the write when nothing user-visible has changed since the last // save. Streaming token updates and idempotent label notifications // otherwise rewrite the entire file on every debounce tick. Include @@ -2914,9 +2916,11 @@ export class AgentSessionManager { .map((a) => `${a.status}:${a.text.length}`) .join(",") + `|${last.fanout.summary.status}:${last.fanout.summary.text.length}` : ""; + // Fold in the usage `updatedAt` so a turn that only changed token usage + // (message text/label/sessionId all unchanged) still writes through. const signature = `${label ?? ""}-${sessionId ?? ""}-${messages.length}-${ last?.message ?? "" - }-${fanoutSig}`; + }-${fanoutSig}-${usage?.updatedAt ?? ""}`; const state = this.getSessionState(session.internalId); if (state.signature === signature) { return state.path ? { path: state.path } : null; @@ -2929,6 +2933,7 @@ export class AgentSessionManager { // GLOBAL_SCOPE writes no frontmatter (byte-identical to legacy chats); // a real project id binds the chat to that scope on disk. projectId: session.projectId, + usage: usage ?? undefined, }); if (result) { state.path = result.path; diff --git a/src/agentMode/session/types.ts b/src/agentMode/session/types.ts index 3feec1b7..9d50e7d9 100644 --- a/src/agentMode/session/types.ts +++ b/src/agentMode/session/types.ts @@ -487,7 +487,28 @@ export type SessionUpdate = | { sessionUpdate: "session_info_update"; title?: string | null } | { sessionUpdate: "current_mode_update"; currentModeId: string } | { sessionUpdate: "config_option_update"; configOptions: BackendConfigOption[] } - | { sessionUpdate: "state_changed"; state: BackendState }; + | { sessionUpdate: "state_changed"; state: BackendState } + | { sessionUpdate: "usage_update"; usage: SessionUsage }; + +/** + * Backend-agnostic token-usage snapshot for a session. Every coding-agent + * backend emits its own usage shape (Claude SDK's `result` message, ACP's + * `usage_update` notification / prompt-result `usage`); translators normalize + * to this one type. `usedTokens` is current context occupancy — it can drop + * after auto-compaction. `contextWindow` is absent when the source couldn't + * report it (e.g. an ACP prompt-result `usage`, whose totals are cumulative, + * not occupancy); such a windowless snapshot is treated as count-only and never + * overrides an occupancy snapshot that carries a window. + */ +export interface SessionUsage { + usedTokens: number; + contextWindow?: number; + inputTokens?: number; + outputTokens?: number; + cacheReadTokens?: number; + cacheWriteTokens?: number; + updatedAt: number; +} /** A SessionEvent is the demuxed pair `(sessionId, update)` consumed by handlers. */ export interface SessionEvent { diff --git a/src/agentMode/ui/AgentChatControls.tsx b/src/agentMode/ui/AgentChatControls.tsx index 9e67c619..7f6a7cba 100644 --- a/src/agentMode/ui/AgentChatControls.tsx +++ b/src/agentMode/ui/AgentChatControls.tsx @@ -29,6 +29,12 @@ interface AgentChatControlsProps { onUpdateChatTitle?: (id: string, newTitle: string) => Promise; onDeleteChat?: (id: string) => Promise; onOpenSourceFile?: (id: string) => Promise; + /** + * Context-window usage meter, rendered as the first item in the right-side + * control cluster (left of New Chat). Self-renders `null` until the backend + * reports usage. Omitted in the not-ready state, so nothing renders there. + */ + usageMeter?: React.ReactNode; } /** @@ -47,6 +53,7 @@ export const AgentChatControls: React.FC = ({ onUpdateChatTitle, onDeleteChat, onOpenSourceFile, + usageMeter, }) => { const settings = useSettingsValue(); const historyAvailable = Boolean( @@ -60,6 +67,7 @@ export const AgentChatControls: React.FC = ({ agent (alpha)
+ {usageMeter} {onNewChat && ( diff --git a/src/agentMode/ui/AgentContextMeter.test.tsx b/src/agentMode/ui/AgentContextMeter.test.tsx new file mode 100644 index 00000000..4ed46ed6 --- /dev/null +++ b/src/agentMode/ui/AgentContextMeter.test.tsx @@ -0,0 +1,118 @@ +import type { AgentChatBackend } from "@/agentMode/session/AgentChatBackend"; +import type { SessionUsage } from "@/agentMode/session/types"; +import AgentContextMeter from "@/agentMode/ui/AgentContextMeter"; +import { TooltipProvider } from "@/components/ui/tooltip"; +import { fireEvent, render, screen } from "@testing-library/react"; +import * as React from "react"; + +// Radix Tooltip portals into Obsidian's `activeDocument`; jsdom lacks it. +beforeAll(() => { + (window as unknown as { activeDocument: Document }).activeDocument = window.document; +}); + +/** Minimal backend stub exposing just the getters/subscribe the meter reads. */ +function makeBackend(usage: SessionUsage | null): AgentChatBackend { + return { + getSessionUsage: () => usage, + subscribe: () => () => {}, + } as unknown as AgentChatBackend; +} + +/** The meter's tooltip needs a Radix `TooltipProvider` ancestor (the app mounts + * one at the chat-view root, alongside the sibling control buttons). */ +function renderMeter(backend: AgentChatBackend) { + return render( + + + + ); +} + +describe("AgentContextMeter", () => { + it("renders the % ring plus tooltip numbers when contextWindow is known", () => { + const usage: SessionUsage = { + usedTokens: 50_000, + contextWindow: 200_000, + inputTokens: 40_000, + outputTokens: 8_000, + cacheReadTokens: 1_500, + cacheWriteTokens: 500, + updatedAt: 1, + }; + renderMeter(makeBackend(usage)); + + // The trigger is an icon-sized button with just the ring (no inline % text). + const trigger = screen.getByLabelText("Context usage"); + expect(trigger.textContent).not.toContain("25%"); + + // The tooltip opens on hover/focus, not click. + fireEvent.focus(trigger); + + // Tooltip: "Context window" label + used / total (percent) on one line. + // Radix Tooltip renders the content twice (visible + a visually-hidden a11y + // copy), so assert on all matches rather than a single node. + expect(screen.getAllByText("Context window").length).toBeGreaterThan(0); + // 50k / 200k = 25%, formatted with k/M suffixes. + expect(screen.getAllByText("50.0k / 200.0k (25%)").length).toBeGreaterThan(0); + // The technical breakdown was intentionally dropped. + expect(screen.queryByText(/in ·|out ·| cache/)).toBeNull(); + }); + + it("applies the warning color once usage reaches 85%", () => { + const usage: SessionUsage = { + usedTokens: 170_000, + contextWindow: 200_000, + updatedAt: 1, + }; + renderMeter(makeBackend(usage)); + + const trigger = screen.getByLabelText("Context usage"); + // The warning accent lives on the trigger itself. + expect(trigger.className).toContain("tw-text-warning"); + expect(trigger.className).not.toContain("tw-text-accent"); + + // 170k / 200k = 85%, surfaced in the tooltip stats line (opens on focus). + fireEvent.focus(trigger); + expect(screen.getAllByText("170.0k / 200.0k (85%)").length).toBeGreaterThan(0); + }); + + it("stays on the accent color below the warning threshold", () => { + const usage: SessionUsage = { + usedTokens: 100_000, + contextWindow: 200_000, + updatedAt: 1, + }; + renderMeter(makeBackend(usage)); + + const trigger = screen.getByLabelText("Context usage"); + expect(trigger.className).toContain("tw-text-accent"); + expect(trigger.className).not.toContain("tw-text-warning"); + }); + + it("falls back to the count-only TokenCounter when there is no contextWindow", () => { + const usage: SessionUsage = { + usedTokens: 12_000, + inputTokens: 10_000, + updatedAt: 1, + }; + // TokenCounter also renders a Radix Tooltip, so it needs the provider too. + const { container } = renderMeter(makeBackend(usage)); + + // No ring meter — the fallback chip has no "Context usage" trigger. + expect(screen.queryByLabelText("Context usage")).toBeNull(); + // TokenCounter shows the rounded-thousands chip. + expect(container.textContent).toContain("12k"); + }); + + it("renders nothing (no separator) when usage is null", () => { + const { container } = render(); + expect(container.childElementCount).toBe(0); + }); + + it("renders nothing (no separator, no chip) when usedTokens is 0 and there is no contextWindow", () => { + const usage: SessionUsage = { usedTokens: 0, updatedAt: 1 }; + const { container } = render(); + expect(container.childElementCount).toBe(0); + expect(screen.queryByLabelText("Context usage")).toBeNull(); + }); +}); diff --git a/src/agentMode/ui/AgentContextMeter.tsx b/src/agentMode/ui/AgentContextMeter.tsx new file mode 100644 index 00000000..dd8ee714 --- /dev/null +++ b/src/agentMode/ui/AgentContextMeter.tsx @@ -0,0 +1,126 @@ +import type { AgentChatBackend } from "@/agentMode/session/AgentChatBackend"; +import type { SessionUsage } from "@/agentMode/session/types"; +import { useSessionUsage } from "@/agentMode/ui/hooks/useSessionUsage"; +import { TokenCounter } from "@/components/chat-components/TokenCounter"; +import { Button } from "@/components/ui/button"; +import { Progress } from "@/components/ui/progress"; +import { Tooltip, TooltipContent, TooltipTrigger } from "@/components/ui/tooltip"; +import { cn } from "@/lib/utils"; +import * as React from "react"; + +interface AgentContextMeterProps { + backend: AgentChatBackend; +} + +/** Usage fraction at/above which the ring flips to the warning color. */ +const WARNING_THRESHOLD = 0.85; + +/** SVG donut geometry — sized to match the composer's `tw-size-4` glyphs. */ +const RING_RADIUS = 6; +const RING_CIRCUMFERENCE = 2 * Math.PI * RING_RADIUS; + +/** One-decimal token count with a k/M suffix (e.g. `248.0k`, `1.0M`). */ +function formatTokens(count: number): string { + if (count >= 1_000_000) return `${(count / 1_000_000).toFixed(1)}M`; + if (count >= 1000) return `${(count / 1000).toFixed(1)}k`; + return count.toLocaleString(); +} + +/** SVG donut whose arc fills to `fraction` (0–1). Color comes from `currentColor`. */ +function ContextRing({ fraction }: { fraction: number }) { + const dashOffset = RING_CIRCUMFERENCE * (1 - fraction); + return ( + + ); +} + +/** Full % ring trigger + a horizontal context-window bar in a hover tooltip. */ +function RingMeter({ usage, contextWindow }: { usage: SessionUsage; contextWindow: number }) { + // Guard a non-finite `usedTokens` (e.g. NaN from a malformed upstream value) + // so it can't propagate into the rendered percent or the SVG dashoffset. + const used = Number.isFinite(usage.usedTokens) ? usage.usedTokens : 0; + const fraction = Math.min(1, Math.max(0, used / contextWindow)); + const percent = Math.round(fraction * 100); + const isWarning = fraction >= WARNING_THRESHOLD; + + // Radix Tooltip handles hover/focus open/close and hoverable content natively, + // so no manual open state or close timer is needed. + return ( + + + + + +
+
+ Context window + + {formatTokens(used)} / {formatTokens(contextWindow)} ({percent}%) + +
+ +
+
+
+ ); +} + +/** + * Circular context-window meter for the agent control bar. Renders as a single + * icon-sized button that sits alongside the other row controls (no separator of + * its own). Render ladder from the backend's {@link SessionUsage}: + * + * - a positive `contextWindow` → the % ring with a hover tooltip (the + * percentage and token numbers live inside the tooltip); + * - `usedTokens` known but no window → the legacy count-only `TokenCounter` + * chip; + * - no usage at all → `null` (renders nothing, so the row shows no control). + */ +export default function AgentContextMeter({ backend }: AgentContextMeterProps) { + const usage = useSessionUsage(backend); + if (usage === null) return null; + + const window = usage.contextWindow; + if (typeof window === "number" && Number.isFinite(window) && window > 0) { + return ; + } + + // Count-only fallback: render nothing when there is no usage to show. + if (!(usage.usedTokens > 0)) return null; + return ; +} diff --git a/src/agentMode/ui/AgentHome.tsx b/src/agentMode/ui/AgentHome.tsx index bfea1eef..1f100a97 100644 --- a/src/agentMode/ui/AgentHome.tsx +++ b/src/agentMode/ui/AgentHome.tsx @@ -1,6 +1,7 @@ import AgentChatMessages from "@/agentMode/ui/AgentChatMessages"; import { AgentChatControls } from "@/agentMode/ui/AgentChatControls"; import { AgentChatInput } from "@/agentMode/ui/AgentChatInput"; +import AgentContextMeter from "@/agentMode/ui/AgentContextMeter"; import AgentContextSection, { buildContextSummary } from "@/agentMode/ui/AgentContextSection"; import AgentContextStatusIcon from "@/agentMode/ui/AgentContextStatusIcon"; import { AgentLandingStack } from "@/agentMode/ui/AgentLandingStack"; @@ -905,6 +906,7 @@ const AgentHomeInternal: React.FC = ({ onUpdateChatTitle={handleUpdateChatTitle} onDeleteChat={handleDeleteChat} onOpenSourceFile={handleOpenSourceFile} + usageMeter={} /> {composerNode} diff --git a/src/agentMode/ui/hooks/useSessionUsage.test.tsx b/src/agentMode/ui/hooks/useSessionUsage.test.tsx new file mode 100644 index 00000000..34d48454 --- /dev/null +++ b/src/agentMode/ui/hooks/useSessionUsage.test.tsx @@ -0,0 +1,96 @@ +import type { AgentChatBackend } from "@/agentMode/session/AgentChatBackend"; +import type { SessionUsage } from "@/agentMode/session/types"; +import { useSessionUsage } from "@/agentMode/ui/hooks/useSessionUsage"; +import { act, renderHook } from "@testing-library/react"; + +/** Stand-in exposing only `getSessionUsage` + `subscribe`; rest cast away. */ +function makeFakeBackend(initial: SessionUsage | null = null) { + const state: { usage: SessionUsage | null } = { usage: initial }; + const listeners = new Set<() => void>(); + + const backend = { + subscribe(listener: () => void) { + listeners.add(listener); + return () => listeners.delete(listener); + }, + getSessionUsage: () => state.usage, + } as unknown as AgentChatBackend; + + return { + backend, + state, + listenerCount: () => listeners.size, + emit: () => listeners.forEach((l) => l()), + }; +} + +const usage = (usedTokens: number): SessionUsage => ({ usedTokens, updatedAt: 1 }); + +describe("useSessionUsage", () => { + it("returns the backend's initial usage snapshot", () => { + const fake = makeFakeBackend(usage(10)); + const { result } = renderHook(() => useSessionUsage(fake.backend)); + expect(result.current).toEqual(usage(10)); + }); + + it("returns null when the backend has no usage yet", () => { + const fake = makeFakeBackend(null); + const { result } = renderHook(() => useSessionUsage(fake.backend)); + expect(result.current).toBeNull(); + }); + + it("re-syncs when the backend notifies", () => { + const fake = makeFakeBackend(null); + const { result } = renderHook(() => useSessionUsage(fake.backend)); + + act(() => { + fake.state.usage = usage(42); + fake.emit(); + }); + + expect(result.current).toEqual(usage(42)); + }); + + it("imperatively syncs to a new backend when the prop changes", () => { + const first = makeFakeBackend(usage(1)); + const second = makeFakeBackend(usage(2)); + + const { result, rerender } = renderHook(({ backend }) => useSessionUsage(backend), { + initialProps: { backend: first.backend }, + }); + expect(result.current).toEqual(usage(1)); + + rerender({ backend: second.backend }); + expect(result.current).toEqual(usage(2)); + }); + + it("unsubscribes on backend switch and unmount", () => { + const first = makeFakeBackend(); + const second = makeFakeBackend(); + const { rerender, unmount } = renderHook(({ backend }) => useSessionUsage(backend), { + initialProps: { backend: first.backend }, + }); + expect(first.listenerCount()).toBe(1); + + rerender({ backend: second.backend }); + expect(first.listenerCount()).toBe(0); + expect(second.listenerCount()).toBe(1); + + unmount(); + expect(second.listenerCount()).toBe(0); + }); + + it("ignores notifications fired after unmount", () => { + const fake = makeFakeBackend(null); + const { result, unmount } = renderHook(() => useSessionUsage(fake.backend)); + + unmount(); + expect(() => + act(() => { + fake.state.usage = usage(99); + fake.emit(); + }) + ).not.toThrow(); + expect(result.current).toBeNull(); + }); +}); diff --git a/src/agentMode/ui/hooks/useSessionUsage.ts b/src/agentMode/ui/hooks/useSessionUsage.ts new file mode 100644 index 00000000..552ff5a0 --- /dev/null +++ b/src/agentMode/ui/hooks/useSessionUsage.ts @@ -0,0 +1,36 @@ +import type { AgentChatBackend } from "@/agentMode/session/AgentChatBackend"; +import type { SessionUsage } from "@/agentMode/session/types"; +import { useEffect, useRef, useState } from "react"; + +/** + * Reactive snapshot of the backend's latest token-usage report, kept in sync + * via the same single subscription the rest of the runtime state uses. Returns + * `null` until the backend has reported usage (fresh / resumed session). + * + * Mirrors the subscribe + `isMounted` pattern in + * {@link useAgentChatRuntimeState}: a fresh initial sync on every `backend` + * change (the lazy initializer only ran for the first backend) and a mount + * guard so a notification racing an unmount is a no-op. + */ +export function useSessionUsage(backend: AgentChatBackend): SessionUsage | null { + const [usage, setUsage] = useState(() => backend.getSessionUsage()); + + const isMountedRef = useRef(false); + useEffect(() => { + isMountedRef.current = true; + return () => { + isMountedRef.current = false; + }; + }, []); + + useEffect(() => { + const sync = () => setUsage(backend.getSessionUsage()); + sync(); + return backend.subscribe(() => { + if (!isMountedRef.current) return; + sync(); + }); + }, [backend]); + + return usage; +}