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