@posthog/agent 2.3.678 → 2.3.703

package/src/adapters/claude/mcp/tool-metadata.ts

@@ -116,6 +116,12 @@ export function getConnectedMcpServerNames(): string[] {
   return [...names];
 }
 
+/** Snapshot of every tool currently in the metadata cache. Used by the
+ *  context-breakdown estimator to size the MCP category. */
+export function getCachedMcpTools(): McpToolMetadata[] {
+  return [...mcpToolMetadataCache.values()];
+}
+
 export function getMcpToolApprovalState(
   toolName: string,
 ): McpToolApprovalState | undefined {

package/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,8 @@ export type Session = BaseSession & {
   pendingMessages: Map<string, PendingMessage>;
   nextPendingOrder: number;
   emitRawSDKMessages: boolean | SDKMessageFilter[];
+  /** Refreshed at session init and on MCP/skill changes. */
+  contextBreakdownBaseline?: ContextBreakdownBaseline;
 };
 
 export type ToolUseCache = {

package/src/adapters/codex/codex-agent.ts

@@ -64,6 +64,12 @@ import {
   nodeWritableToWebWritable,
 } from "../../utils/streams";
 import { BaseAcpAgent, type BaseSession } from "../base-acp-agent";
+import {
+  buildBreakdown,
+  type ContextBreakdownBaseline,
+  emptyBaseline,
+  estimateTokens,
+} from "../claude/context-breakdown";
 import { classifyAgentError } from "../error-classification";
 import {
   enabledLocalTools,
@@ -76,6 +82,7 @@ import { normalizeCodexConfigOptions } from "./models";
 import {
   type CodexSessionState,
   createSessionState,
+  resetSessionState,
   resetUsage,
 } from "./session-state";
 import { CodexSettingsManager } from "./settings";
@@ -161,6 +168,19 @@ function classifyPromptError(error: unknown): unknown {
   );
 }
 
+// codex-rs/protocol/src/protocol.rs BASELINE_TOKENS — the always-resident
+// floor (MCP schemas, skills, preset prompt) we can't attribute per-source.
+const CODEX_BASELINE_TOKENS = 12000;
+
+function buildCodexBaseline(
+  meta: NewSessionMeta | undefined,
+): ContextBreakdownBaseline {
+  const baseline = emptyBaseline();
+  baseline.systemPrompt =
+    CODEX_BASELINE_TOKENS + estimateTokens(meta?.systemPrompt);
+  return baseline;
+}
+
 const CODEX_NATIVE_MODE: Record<CodeExecutionMode, CodexNativeMode> = {
   auto: "auto",
   default: "auto",
@@ -392,8 +412,10 @@ export class CodexAcpAgent extends BaseAcpAgent {
       response.configOptions,
     );
 
-    // Initialize session state
-    this.sessionState = createSessionState(response.sessionId, params.cwd, {
+    // Initialize session state. Mutate in place — codex-client closure-
+    // captured this object in the constructor and writes contextUsed/
+    // accumulatedUsage to it on every upstream usage_update.
+    resetSessionState(this.sessionState, response.sessionId, params.cwd, {
       taskRunId: meta?.taskRunId,
       taskId: resolveTaskId(meta),
       modeId: response.modes?.currentModeId ?? "auto",
@@ -402,6 +424,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
     });
     this.sessionId = response.sessionId;
     this.sessionState.configOptions = response.configOptions ?? [];
+    this.sessionState.contextBreakdownBaseline = buildCodexBaseline(meta);
 
     await this.applyInitialPermissionMode(
       response.sessionId,
@@ -445,7 +468,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
     // notifications (TURN_COMPLETE, USAGE_UPDATE) after a reload. newSession
     // and unstable_resumeSession both do this; loadSession historically did
     // not, which silently broke task-completion tracking on re-attach.
-    this.sessionState = createSessionState(params.sessionId, params.cwd, {
+    resetSessionState(this.sessionState, params.sessionId, params.cwd, {
       taskRunId: meta?.taskRunId,
       taskId: resolveTaskId(meta),
       modeId: response.modes?.currentModeId ?? "auto",
@@ -453,6 +476,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
     });
     this.sessionId = params.sessionId;
     this.sessionState.configOptions = response.configOptions ?? [];
+    this.sessionState.contextBreakdownBaseline = buildCodexBaseline(meta);
 
     if (meta?.taskRunId) {
       await this.client.extNotification(POSTHOG_NOTIFICATIONS.SDK_SESSION, {
@@ -491,7 +515,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
       loadResponse.modes?.currentModeId,
       meta?.permissionMode,
     );
-    this.sessionState = createSessionState(params.sessionId, params.cwd, {
+    resetSessionState(this.sessionState, params.sessionId, params.cwd, {
       taskRunId: meta?.taskRunId,
       taskId: resolveTaskId(meta),
       modeId: loadResponse.modes?.currentModeId ?? "auto",
@@ -499,6 +523,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
     });
     this.sessionId = params.sessionId;
     this.sessionState.configOptions = loadResponse.configOptions ?? [];
+    this.sessionState.contextBreakdownBaseline = buildCodexBaseline(meta);
 
     if (meta?.taskRunId) {
       await this.client.extNotification(POSTHOG_NOTIFICATIONS.SDK_SESSION, {
@@ -538,7 +563,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
     );
 
     const requestedPermissionMode = toCodexPermissionMode(meta?.permissionMode);
-    this.sessionState = createSessionState(newResponse.sessionId, params.cwd, {
+    resetSessionState(this.sessionState, newResponse.sessionId, params.cwd, {
       taskRunId: meta?.taskRunId,
       taskId: resolveTaskId(meta),
       modeId: newResponse.modes?.currentModeId ?? "auto",
@@ -546,6 +571,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
     });
     this.sessionId = newResponse.sessionId;
     this.sessionState.configOptions = newResponse.configOptions ?? [];
+    this.sessionState.contextBreakdownBaseline = buildCodexBaseline(meta);
 
     await this.applyInitialPermissionMode(
       newResponse.sessionId,
@@ -730,6 +756,21 @@ export class CodexAcpAgent extends BaseAcpAgent {
       }
     }
 
+    // Emit the per-source breakdown so the renderer's ContextBreakdownPopover
+    // has data to show. This fires regardless of `taskRunId` (local sessions
+    // need it too) and regardless of `response.usage` (codex-acp doesn't
+    // populate it — context size comes from upstream usage_update events,
+    // tracked on sessionState.contextUsed).
+    if (this.sessionState.contextUsed !== undefined) {
+      await this.client.extNotification(POSTHOG_NOTIFICATIONS.USAGE_UPDATE, {
+        sessionId: params.sessionId,
+        breakdown: buildBreakdown(
+          this.sessionState.contextBreakdownBaseline ?? emptyBaseline(),
+          this.sessionState.contextUsed,
+        ),
+      });
+    }
+
     return response;
   }
 

package/src/adapters/codex/codex-client.test.ts

@@ -14,7 +14,7 @@ vi.mock("../../enrichment/file-enricher", () => ({
 }));
 
 import { createCodexClient } from "./codex-client";
-import { createSessionState } from "./session-state";
+import { createSessionState, resetSessionState } from "./session-state";
 
 function makeUpstream(response: ReadTextFileResponse): AgentSideConnection & {
   readTextFile: ReturnType<typeof vi.fn>;
@@ -288,3 +288,51 @@ describe("createCodexClient onStructuredOutput", () => {
     expect(upstream.sessionUpdate).toHaveBeenCalledTimes(1);
   });
 });
+
+describe("createCodexClient usage_update propagation", () => {
+  const logger = new Logger({ debug: false, prefix: "[test]" });
+
+  function makeUpstream(): AgentSideConnection {
+    return {
+      sessionUpdate: vi.fn(async () => {}),
+      requestPermission: vi.fn(),
+      readTextFile: vi.fn(),
+      writeTextFile: vi.fn(),
+      createTerminal: vi.fn(),
+      terminalOutput: vi.fn(),
+      releaseTerminal: vi.fn(),
+      waitForTerminalExit: vi.fn(),
+      killTerminal: vi.fn(),
+      extMethod: vi.fn(),
+      extNotification: vi.fn(),
+    } as unknown as AgentSideConnection;
+  }
+
+  // Regression: codex-client closure-captures the sessionState reference in
+  // its factory. CodexAcpAgent constructs the client once at startup with the
+  // initial "" sessionId state, then resetSessionState() mutates that same
+  // object on every newSession/loadSession/etc. If the agent ever reassigned
+  // `this.sessionState`, contextUsed writes would land on an orphan and the
+  // breakdown notification would never fire.
+  test("writes contextUsed to the same state object after resetSessionState", async () => {
+    const sessionState = createSessionState("", "/tmp");
+    const upstream = makeUpstream();
+    const client = createCodexClient(upstream, logger, sessionState);
+
+    resetSessionState(sessionState, "real-session", "/tmp/repo", {
+      taskRunId: "run-1",
+    });
+
+    await client.sessionUpdate?.({
+      sessionId: "real-session",
+      update: {
+        sessionUpdate: "usage_update",
+        used: 123_456,
+        size: 200_000,
+      },
+    } as unknown as SessionNotification);
+
+    expect(sessionState.contextUsed).toBe(123_456);
+    expect(sessionState.contextSize).toBe(200_000);
+  });
+});

package/src/adapters/codex/session-state.ts

@@ -1,10 +1,6 @@
-/**
- * Session state tracking for Codex proxy agent.
- * Tracks usage accumulation, model/mode state, and config options.
- */
-
 import type { SessionConfigOption } from "@agentclientprotocol/sdk";
 import type { PermissionMode } from "../../execution-mode";
+import type { ContextBreakdownBaseline } from "../claude/context-breakdown";
 
 export interface CodexUsage {
   inputTokens: number;
@@ -22,6 +18,7 @@ export interface CodexSessionState {
   accumulatedUsage: CodexUsage;
   contextSize?: number;
   contextUsed?: number;
+  contextBreakdownBaseline?: ContextBreakdownBaseline;
   permissionMode: PermissionMode;
   taskRunId?: string;
   taskId?: string;
@@ -56,6 +53,40 @@ export function createSessionState(
   };
 }
 
+// codex-client closure-captures the original sessionState reference, so we
+// must mutate in place across newSession/loadSession/resumeSession/forkSession
+// — reassigning would orphan it and silently break usage propagation.
+export function resetSessionState(
+  state: CodexSessionState,
+  sessionId: string,
+  cwd: string,
+  opts?: {
+    taskRunId?: string;
+    taskId?: string;
+    modeId?: string;
+    modelId?: string;
+    permissionMode?: PermissionMode;
+  },
+): void {
+  state.sessionId = sessionId;
+  state.cwd = cwd;
+  state.modeId = opts?.modeId ?? "auto";
+  state.modelId = opts?.modelId;
+  state.configOptions = [];
+  state.accumulatedUsage = {
+    inputTokens: 0,
+    outputTokens: 0,
+    cachedReadTokens: 0,
+    cachedWriteTokens: 0,
+  };
+  state.contextSize = undefined;
+  state.contextUsed = undefined;
+  state.contextBreakdownBaseline = undefined;
+  state.permissionMode = opts?.permissionMode ?? "auto";
+  state.taskRunId = opts?.taskRunId;
+  state.taskId = opts?.taskId;
+}
+
 export function resetUsage(state: CodexSessionState): void {
   state.accumulatedUsage = {
     inputTokens: 0,