diff --git a/apps/code/src/renderer/features/sessions/hooks/useContextUsage.test.ts b/apps/code/src/renderer/features/sessions/hooks/useContextUsage.test.ts new file mode 100644 index 000000000..88c37ffa0 --- /dev/null +++ b/apps/code/src/renderer/features/sessions/hooks/useContextUsage.test.ts @@ -0,0 +1,79 @@ +import type { AcpMessage } from "@shared/types/session-events"; +import { describe, expect, it } from "vitest"; +import { extractContextUsage } from "./useContextUsage"; + +function usageUpdateEvent(used: number, size: number): AcpMessage { + return { + type: "acp_message", + ts: 1, + message: { + jsonrpc: "2.0", + method: "session/update", + params: { + sessionId: "s1", + update: { sessionUpdate: "usage_update", used, size }, + }, + }, + }; +} + +function breakdownEvent( + breakdown: Record, + method = "_posthog/usage_update", +): AcpMessage { + return { + type: "acp_message", + ts: 1, + message: { jsonrpc: "2.0", method, params: { sessionId: "s1", breakdown } }, + }; +} + +describe("extractContextUsage", () => { + it("returns null with no usage event", () => { + expect(extractContextUsage([])).toBeNull(); + }); + + it("derives aggregate from the latest session/update", () => { + const result = extractContextUsage([usageUpdateEvent(50_000, 200_000)]); + expect(result?.used).toBe(50_000); + expect(result?.size).toBe(200_000); + expect(result?.percentage).toBe(25); + expect(result?.breakdown).toBeNull(); + }); + + it("merges breakdown from a _posthog/usage_update notification", () => { + const result = extractContextUsage([ + usageUpdateEvent(50_000, 200_000), + breakdownEvent({ + systemPrompt: 4000, + tools: 500, + rules: 0, + skills: 0, + mcp: 0, + subagents: 0, + conversation: 45_500, + }), + ]); + expect(result?.breakdown?.systemPrompt).toBe(4000); + expect(result?.breakdown?.conversation).toBe(45_500); + }); + + it("tolerates the double-underscore method prefix from extNotification", () => { + const result = extractContextUsage([ + usageUpdateEvent(50_000, 200_000), + breakdownEvent( + { + systemPrompt: 4000, + tools: 0, + rules: 0, + skills: 0, + mcp: 0, + subagents: 0, + conversation: 46_000, + }, + "__posthog/usage_update", + ), + ]); + expect(result?.breakdown?.systemPrompt).toBe(4000); + }); +}); diff --git a/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts b/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts index b340e9abb..8c1e2dcc8 100644 --- a/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts +++ b/apps/code/src/renderer/features/sessions/hooks/useContextUsage.ts @@ -1,11 +1,26 @@ import type { AcpMessage } from "@shared/types/session-events"; import { useMemo } from "react"; +// Shape mirrors `ContextBreakdown` emitted by the agent in +// `_posthog/usage_update` (see packages/agent/src/adapters/claude/context-breakdown.ts). +// Kept local to avoid a renderer dependency on the agent package; if the shape +// drifts, lift it into @posthog/shared. +export interface ContextBreakdown { + systemPrompt: number; + tools: number; + rules: number; + skills: number; + mcp: number; + subagents: number; + conversation: number; +} + export interface ContextUsage { used: number; size: number; percentage: number; cost: { amount: number; currency: string } | null; + breakdown: ContextBreakdown | null; } /** @@ -18,42 +33,75 @@ export function useContextUsage(events: AcpMessage[]): ContextUsage | null { } export function extractContextUsage(events: AcpMessage[]): ContextUsage | null { + let aggregate: Omit | null = null; + let breakdown: ContextBreakdown | null = null; + for (let i = events.length - 1; i >= 0; i--) { const msg = events[i].message; + if (!aggregate) { + aggregate = extractAggregate(msg); + } + if (!breakdown) { + breakdown = extractBreakdown(msg); + } + if (aggregate && breakdown) break; + } + + if (!aggregate) return null; + return { ...aggregate, breakdown }; +} + +function extractAggregate( + msg: AcpMessage["message"], +): Omit | null { + if ( + "method" in msg && + msg.method === "session/update" && + !("id" in msg) && + "params" in msg + ) { + const params = msg.params as + | { + update?: { + sessionUpdate?: string; + used?: number; + size?: number; + cost?: { amount: number; currency: string } | null; + }; + } + | undefined; + const update = params?.update; if ( - "method" in msg && - msg.method === "session/update" && - !("id" in msg) && - "params" in msg + update?.sessionUpdate === "usage_update" && + typeof update.used === "number" && + typeof update.size === "number" ) { - const params = msg.params as - | { - update?: { - sessionUpdate?: string; - used?: number; - size?: number; - cost?: { amount: number; currency: string } | null; - }; - } - | undefined; - const update = params?.update; - if ( - update?.sessionUpdate === "usage_update" && - typeof update.used === "number" && - typeof update.size === "number" - ) { - const percentage = - update.size > 0 - ? Math.min(100, Math.round((update.used / update.size) * 100)) - : 0; - return { - used: update.used, - size: update.size, - percentage, - cost: update.cost ?? null, - }; - } + const percentage = + update.size > 0 + ? Math.min(100, Math.round((update.used / update.size) * 100)) + : 0; + return { + used: update.used, + size: update.size, + percentage, + cost: update.cost ?? null, + }; } } return null; } + +function extractBreakdown(msg: AcpMessage["message"]): ContextBreakdown | null { + if (!("method" in msg) || !("params" in msg)) return null; + // Method may be received as either `_posthog/usage_update` or + // `__posthog/usage_update` depending on how the transport stringifies it + // (see acp-extensions.ts:matchesExt). + if ( + msg.method !== "_posthog/usage_update" && + msg.method !== "__posthog/usage_update" + ) { + return null; + } + const params = msg.params as { breakdown?: ContextBreakdown } | undefined; + return params?.breakdown ?? null; +} diff --git a/packages/agent/src/adapters/claude/claude-agent.ts b/packages/agent/src/adapters/claude/claude-agent.ts index cb712669c..a8824c8fb 100644 --- a/packages/agent/src/adapters/claude/claude-agent.ts +++ b/packages/agent/src/adapters/claude/claude-agent.ts @@ -68,6 +68,11 @@ import { Pushable } from "../../utils/streams"; import { BaseAcpAgent } from "../base-acp-agent"; import { LOCAL_TOOLS_MCP_NAME } from "../local-tools"; import { resolveTaskId } from "../session-meta"; +import { + buildBreakdown, + emptyBaseline, + estimateSystemPrompt, +} from "./context-breakdown"; import { promptToClaude } from "./conversion/acp-to-sdk"; import { handleResultMessage, @@ -556,6 +561,14 @@ export class ClaudeAcpAgent extends BaseAcpAgent { }); } + // Sum the result's own input categories rather than reusing + // `lastAssistantTotalUsage` (which is the streamed delta from the + // outermost model only). For subagent turns the two can diverge; + // the breakdown is indicative either way. + const breakdownInputTokens = + (message.usage.input_tokens ?? 0) + + (message.usage.cache_read_input_tokens ?? 0) + + (message.usage.cache_creation_input_tokens ?? 0); await this.client.extNotification( POSTHOG_NOTIFICATIONS.USAGE_UPDATE, { @@ -567,6 +580,10 @@ export class ClaudeAcpAgent extends BaseAcpAgent { cachedWriteTokens: message.usage.cache_creation_input_tokens, }, cost: message.total_cost_usd, + breakdown: buildBreakdown( + this.session.contextBreakdownBaseline ?? emptyBaseline(), + breakdownInputTokens, + ), }, ); @@ -1221,6 +1238,10 @@ export class ClaudeAcpAgent extends BaseAcpAgent { pendingMessages: new Map(), nextPendingOrder: 0, emitRawSDKMessages: meta?.claudeCode?.emitRawSDKMessages ?? false, + contextBreakdownBaseline: { + ...emptyBaseline(), + systemPrompt: estimateSystemPrompt(systemPrompt), + }, // Custom properties cwd, diff --git a/packages/agent/src/adapters/claude/context-breakdown.test.ts b/packages/agent/src/adapters/claude/context-breakdown.test.ts new file mode 100644 index 000000000..cf8660e73 --- /dev/null +++ b/packages/agent/src/adapters/claude/context-breakdown.test.ts @@ -0,0 +1,84 @@ +import { describe, expect, it } from "vitest"; +import { + buildBreakdown, + emptyBaseline, + estimateJsonTokens, + estimateSystemPrompt, + estimateTokens, +} from "./context-breakdown"; + +describe("estimateTokens", () => { + it("returns 0 for empty input", () => { + expect(estimateTokens("")).toBe(0); + expect(estimateTokens(undefined)).toBe(0); + expect(estimateTokens(null)).toBe(0); + }); + + it("scales roughly with input length", () => { + expect(estimateTokens("a".repeat(35))).toBe(10); + expect(estimateTokens("a".repeat(350))).toBe(100); + }); +}); + +describe("estimateJsonTokens", () => { + it("counts JSON representation of objects", () => { + const tokens = estimateJsonTokens({ name: "Read", schema: { foo: 1 } }); + expect(tokens).toBeGreaterThan(0); + }); + + it("returns 0 for non-serializable values", () => { + const circular: Record = {}; + circular.self = circular; + expect(estimateJsonTokens(circular)).toBe(0); + }); +}); + +describe("estimateSystemPrompt", () => { + it("includes the Claude preset budget when preset is used", () => { + const noAppend = estimateSystemPrompt({ type: "preset" }); + expect(noAppend).toBeGreaterThan(0); + }); + + it("adds the append portion on top of the preset", () => { + const append = "a".repeat(350); + const result = estimateSystemPrompt({ type: "preset", append }); + const presetOnly = estimateSystemPrompt({ type: "preset" }); + expect(result - presetOnly).toBe(100); + }); + + it("counts a raw string verbatim with no preset overhead", () => { + expect(estimateSystemPrompt("a".repeat(350))).toBe(100); + }); + + it("treats undefined as the bare preset", () => { + expect(estimateSystemPrompt(undefined)).toBe( + estimateSystemPrompt({ type: "preset" }), + ); + }); +}); + +describe("buildBreakdown", () => { + it("derives conversation from input - stable sum", () => { + const baseline = { + ...emptyBaseline(), + systemPrompt: 4000, + tools: 500, + }; + const result = buildBreakdown(baseline, 10_000); + expect(result.systemPrompt).toBe(4000); + expect(result.tools).toBe(500); + expect(result.conversation).toBe(5500); + }); + + it("floors conversation at 0 when stable pieces exceed input", () => { + const baseline = { ...emptyBaseline(), systemPrompt: 5000 }; + expect(buildBreakdown(baseline, 1000).conversation).toBe(0); + }); + + it("includes zero categories", () => { + const result = buildBreakdown(emptyBaseline(), 100); + expect(result.mcp).toBe(0); + expect(result.skills).toBe(0); + expect(result.subagents).toBe(0); + }); +}); diff --git a/packages/agent/src/adapters/claude/context-breakdown.ts b/packages/agent/src/adapters/claude/context-breakdown.ts new file mode 100644 index 000000000..2775cbc28 --- /dev/null +++ b/packages/agent/src/adapters/claude/context-breakdown.ts @@ -0,0 +1,105 @@ +/** + * Per-source context-window token breakdown for the renderer's + * `ContextBreakdownPopover`. Anthropic doesn't break down `input_tokens` by + * source, so we tokenize the pieces we control client-side using a cheap + * character-ratio estimator (~3.5 chars/token). Numbers are indicative, not + * invoice-grade — used only for relative-share UX. + */ + +export type ContextCategory = + | "systemPrompt" + | "tools" + | "rules" + | "skills" + | "mcp" + | "subagents" + | "conversation"; + +export type ContextBreakdown = Record; + +// Rough estimate of Claude's bundled `claude_code` preset system prompt. The +// preset content is opaque to us so we add this constant when the systemPrompt +// uses the preset — otherwise it'd show up as Conversation and skew the chart. +const CLAUDE_PRESET_ESTIMATE_TOKENS = 4000; + +const CHARS_PER_TOKEN = 3.5; + +export function estimateTokens(text: string | undefined | null): number { + if (!text) return 0; + return Math.max(0, Math.round(text.length / CHARS_PER_TOKEN)); +} + +export function estimateJsonTokens(value: unknown): number { + try { + return estimateTokens(JSON.stringify(value)); + } catch { + return 0; + } +} + +export interface ContextBreakdownBaseline { + systemPrompt: number; + tools: number; + rules: number; + skills: number; + mcp: number; + subagents: number; +} + +export function emptyBaseline(): ContextBreakdownBaseline { + return { + systemPrompt: 0, + tools: 0, + rules: 0, + skills: 0, + mcp: 0, + subagents: 0, + }; +} + +/** + * Estimate tokens for whatever shape `Options["systemPrompt"]` ended up being: + * a raw string, a `{ type: "preset", append }` object, or undefined. + */ +export function estimateSystemPrompt(systemPrompt: unknown): number { + if (!systemPrompt) return CLAUDE_PRESET_ESTIMATE_TOKENS; + if (typeof systemPrompt === "string") return estimateTokens(systemPrompt); + if (typeof systemPrompt === "object") { + const obj = systemPrompt as { type?: string; append?: unknown }; + const appendTokens = + typeof obj.append === "string" ? estimateTokens(obj.append) : 0; + if (obj.type === "preset") { + return CLAUDE_PRESET_ESTIMATE_TOKENS + appendTokens; + } + return appendTokens; + } + return 0; +} + +/** + * Derive the per-source breakdown from a stable baseline + the current turn's + * input-token total. The conversation bucket is whatever is left after the + * stable pieces are subtracted; it's floored at 0 to absorb estimation drift. + */ +export function buildBreakdown( + baseline: ContextBreakdownBaseline, + currentInputTokens: number, +): ContextBreakdown { + const stableSum = + baseline.systemPrompt + + baseline.tools + + baseline.rules + + baseline.skills + + baseline.mcp + + baseline.subagents; + const conversation = Math.max(0, currentInputTokens - stableSum); + return { + systemPrompt: baseline.systemPrompt, + tools: baseline.tools, + rules: baseline.rules, + skills: baseline.skills, + mcp: baseline.mcp, + subagents: baseline.subagents, + conversation, + }; +} diff --git a/packages/agent/src/adapters/claude/types.ts b/packages/agent/src/adapters/claude/types.ts index da75bb62d..b79fdb2bd 100644 --- a/packages/agent/src/adapters/claude/types.ts +++ b/packages/agent/src/adapters/claude/types.ts @@ -10,6 +10,7 @@ import type { } from "@anthropic-ai/claude-agent-sdk"; import type { Pushable } from "../../utils/streams"; import type { BaseSession } from "../base-acp-agent"; +import type { ContextBreakdownBaseline } from "./context-breakdown"; import type { McpToolApprovals } from "./mcp/tool-metadata"; import type { SettingsManager } from "./session/settings"; import type { CodeExecutionMode } from "./tools"; @@ -65,6 +66,10 @@ export type Session = BaseSession & { pendingMessages: Map; nextPendingOrder: number; emitRawSDKMessages: boolean | SDKMessageFilter[]; + /** Per-source token estimates for stable pieces (system prompt, tools, etc.) + * used by the renderer's context-breakdown popover. Refreshed at session + * init and on MCP/skill changes. */ + contextBreakdownBaseline?: ContextBreakdownBaseline; }; export type ToolUseCache = {