feat(agent): emit per-category context token breakdown

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<string, number>,
+ 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);
+ });
+});

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<ContextUsage, "breakdown"> | 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<ContextUsage, "breakdown"> | 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;
+}

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,

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<string, unknown> = {};
+ 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);
+ });
+});