@posthog/agent 2.3.678 → 2.3.696

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,