@posthog/agent 2.3.312 → 2.3.316

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@posthog/agent",
-  "version": "2.3.312",
+  "version": "2.3.316",
   "repository": "https://github.com/PostHog/code",
   "description": "TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog",
   "exports": {

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

@@ -0,0 +1,289 @@
+import { Readable, Writable } from "node:stream";
+import type { AgentSideConnection, McpServer } from "@agentclientprotocol/sdk";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { POSTHOG_METHODS } from "../../acp-extensions";
+
+type MockCodexConnection = {
+  initialize: ReturnType<typeof vi.fn>;
+  newSession: ReturnType<typeof vi.fn>;
+  loadSession: ReturnType<typeof vi.fn>;
+  setSessionMode: ReturnType<typeof vi.fn>;
+  listSessions: ReturnType<typeof vi.fn>;
+  prompt: ReturnType<typeof vi.fn>;
+  setSessionConfigOption: ReturnType<typeof vi.fn>;
+  cancel: ReturnType<typeof vi.fn>;
+};
+
+type SpawnHandle = {
+  process: { pid: number };
+  stdin: Writable;
+  stdout: Readable;
+  kill: ReturnType<typeof vi.fn>;
+};
+
+const hoisted = vi.hoisted(() => {
+  // Everything the mock factories depend on must live here — vi.mock()
+  // invocations are hoisted above any other top-level code.
+  const createdConnections: MockCodexConnection[] = [];
+  const spawnedProcesses: SpawnHandle[] = [];
+
+  const makeConnection = (): MockCodexConnection => ({
+    initialize: vi.fn().mockResolvedValue({
+      protocolVersion: 1,
+      agentCapabilities: {},
+    }),
+    newSession: vi.fn(),
+    loadSession: vi.fn().mockResolvedValue({
+      modes: { currentModeId: "auto", availableModes: [] },
+      configOptions: [],
+    }),
+    setSessionMode: vi.fn().mockResolvedValue({}),
+    listSessions: vi.fn(),
+    prompt: vi.fn(),
+    setSessionConfigOption: vi.fn(),
+    cancel: vi.fn().mockResolvedValue(undefined),
+  });
+
+  const clientSideConnectionCtor = vi.fn(() => {
+    const conn = makeConnection();
+    createdConnections.push(conn);
+    return conn;
+  });
+
+  const spawnCodexProcessMock = vi.fn(() => {
+    const handle: SpawnHandle = {
+      process: { pid: 1000 + spawnedProcesses.length },
+      stdin: new Writable({
+        write(_chunk, _encoding, callback) {
+          callback();
+        },
+      }),
+      stdout: new Readable({ read() {} }),
+      kill: vi.fn(),
+    };
+    spawnedProcesses.push(handle);
+    return handle;
+  });
+
+  return {
+    createdConnections,
+    spawnedProcesses,
+    clientSideConnectionCtor,
+    spawnCodexProcessMock,
+  };
+});
+
+const createdConnections = hoisted.createdConnections;
+const spawnedProcesses = hoisted.spawnedProcesses;
+const clientSideConnectionCtor = hoisted.clientSideConnectionCtor;
+
+vi.mock("@agentclientprotocol/sdk", async () => {
+  const actual = await vi.importActual("@agentclientprotocol/sdk");
+  return {
+    ...actual,
+    ClientSideConnection: hoisted.clientSideConnectionCtor,
+    ndJsonStream: vi.fn(() => ({}) as object),
+  };
+});
+
+vi.mock("./spawn", () => ({
+  spawnCodexProcess: hoisted.spawnCodexProcessMock,
+}));
+
+vi.mock("./settings", () => ({
+  CodexSettingsManager: vi.fn().mockImplementation((cwd: string) => ({
+    initialize: vi.fn().mockResolvedValue(undefined),
+    dispose: vi.fn(),
+    getCwd: () => cwd,
+    setCwd: vi.fn(),
+    getSettings: () => ({ mcpServerNames: [] }),
+  })),
+}));
+
+import { CodexAcpAgent } from "./codex-agent";
+
+type PrivateAgent = {
+  session: {
+    abortController: AbortController;
+    settingsManager: { dispose: ReturnType<typeof vi.fn> };
+    notificationHistory: unknown[];
+    promptRunning: boolean;
+  };
+  sessionId: string;
+  sessionState: {
+    sessionId: string;
+    cwd: string;
+    accumulatedUsage: {
+      inputTokens: number;
+      outputTokens: number;
+      cachedReadTokens: number;
+      cachedWriteTokens: number;
+    };
+    configOptions: unknown[];
+    taskRunId?: string;
+  };
+  codexProcess: SpawnHandle;
+  codexConnection: MockCodexConnection;
+  lastInitRequest?: { protocolVersion: number };
+};
+
+function makeAgent(): CodexAcpAgent {
+  const client = {
+    extNotification: vi.fn().mockResolvedValue(undefined),
+  } as unknown as AgentSideConnection;
+  return new CodexAcpAgent(client, {
+    codexProcessOptions: { cwd: "/tmp/repo" },
+  });
+}
+
+function primeSession(
+  agent: CodexAcpAgent,
+  sessionId: string,
+): {
+  oldProcess: SpawnHandle;
+  oldConnection: MockCodexConnection;
+  priv: PrivateAgent;
+} {
+  const priv = agent as unknown as PrivateAgent;
+  priv.sessionId = sessionId;
+  priv.sessionState = {
+    sessionId,
+    cwd: "/tmp/repo",
+    accumulatedUsage: {
+      inputTokens: 42,
+      outputTokens: 17,
+      cachedReadTokens: 0,
+      cachedWriteTokens: 0,
+    },
+    configOptions: [{ id: "opt", value: "x" }],
+    taskRunId: "run-1",
+  };
+  priv.session.notificationHistory = [{ foo: "bar" }];
+  priv.lastInitRequest = { protocolVersion: 1 };
+  return {
+    oldProcess: priv.codexProcess,
+    oldConnection: priv.codexConnection,
+    priv,
+  };
+}
+
+describe("CodexAcpAgent.extMethod refresh_session", () => {
+  beforeEach(() => {
+    spawnedProcesses.length = 0;
+    createdConnections.length = 0;
+    clientSideConnectionCtor.mockClear();
+  });
+
+  it("returns methodNotFound for unknown extension methods", async () => {
+    const agent = makeAgent();
+    await expect(agent.extMethod("_posthog/nope", {})).rejects.toThrow(
+      /Method not found/i,
+    );
+  });
+
+  it("rejects when mcpServers is missing", async () => {
+    const agent = makeAgent();
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {}),
+    ).rejects.toThrow(/at least one refreshable field/);
+  });
+
+  it("rejects when mcpServers is not an array", async () => {
+    const agent = makeAgent();
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: "nope" as unknown,
+      }),
+    ).rejects.toThrow(/mcpServers must be an array/);
+  });
+
+  it("rejects refresh while a prompt is in flight", async () => {
+    const agent = makeAgent();
+    const { priv } = primeSession(agent, "s-1");
+    priv.session.promptRunning = true;
+
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: [
+          { name: "posthog", type: "http", url: "https://new", headers: [] },
+        ],
+      }),
+    ).rejects.toThrow(/prompt turn is in flight/);
+  });
+
+  it("respawns the subprocess, re-initializes, and rehydrates with new MCP servers", async () => {
+    const agent = makeAgent();
+    const { oldProcess, oldConnection, priv } = primeSession(agent, "s-2");
+    const oldAbortController = priv.session.abortController;
+    const oldSettingsManager = priv.session.settingsManager;
+
+    const mcpServers: McpServer[] = [
+      {
+        name: "posthog",
+        type: "http",
+        url: "https://fresh",
+        headers: [{ name: "x-foo", value: "bar" }],
+      },
+    ];
+
+    const result = await agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+      mcpServers,
+    });
+
+    expect(result).toEqual({ refreshed: true });
+
+    // Old subprocess torn down, old connection cancelled.
+    expect(oldConnection.cancel).toHaveBeenCalledWith({ sessionId: "s-2" });
+    expect(oldProcess.kill).toHaveBeenCalledTimes(1);
+    expect(oldAbortController.signal.aborted).toBe(true);
+    expect(oldSettingsManager.dispose).toHaveBeenCalledTimes(1);
+
+    // A fresh subprocess was spawned and a new ClientSideConnection wired up.
+    expect(spawnedProcesses).toHaveLength(2);
+    expect(createdConnections).toHaveLength(2);
+    const newConnection = createdConnections[1];
+    if (!newConnection) throw new Error("expected a second connection");
+
+    // ACP handshake replayed against the new subprocess.
+    expect(newConnection.initialize).toHaveBeenCalledWith({
+      protocolVersion: 1,
+    });
+    expect(newConnection.loadSession).toHaveBeenCalledWith({
+      sessionId: "s-2",
+      cwd: "/tmp/repo",
+      mcpServers,
+    });
+
+    // References swapped to the new instances.
+    expect(priv.codexProcess).toBe(spawnedProcesses[1]);
+    expect(priv.codexConnection).toBe(newConnection);
+    expect(priv.session.abortController).not.toBe(oldAbortController);
+    expect(priv.session.settingsManager).not.toBe(oldSettingsManager);
+
+    // Session-level state preserved across refresh.
+    expect(priv.sessionState.accumulatedUsage.inputTokens).toBe(42);
+    expect(priv.sessionState.accumulatedUsage.outputTokens).toBe(17);
+    expect(priv.sessionState.configOptions).toEqual([
+      { id: "opt", value: "x" },
+    ]);
+    expect(priv.sessionState.taskRunId).toBe("run-1");
+    expect(priv.session.notificationHistory).toEqual([{ foo: "bar" }]);
+  });
+
+  it("does not fail refresh when cancel() throws on the stale connection", async () => {
+    const agent = makeAgent();
+    const { oldConnection } = primeSession(agent, "s-3");
+    oldConnection.cancel.mockRejectedValueOnce(new Error("already dead"));
+
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: [
+          { name: "posthog", type: "http", url: "https://x", headers: [] },
+        ],
+      }),
+    ).resolves.toEqual({ refreshed: true });
+
+    expect(spawnedProcesses).toHaveLength(2);
+    expect(createdConnections[1]?.loadSession).toHaveBeenCalled();
+  });
+});

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

@@ -21,11 +21,13 @@ import {
   type ListSessionsResponse,
   type LoadSessionRequest,
   type LoadSessionResponse,
+  type McpServer,
   type NewSessionRequest,
   type NewSessionResponse,
   ndJsonStream,
   type PromptRequest,
   type PromptResponse,
+  RequestError,
   type ResumeSessionRequest,
   type ResumeSessionResponse,
   type SetSessionConfigOptionRequest,
@@ -34,7 +36,11 @@ import {
   type SetSessionModeResponse,
 } from "@agentclientprotocol/sdk";
 import packageJson from "../../../package.json" with { type: "json" };
-import { POSTHOG_NOTIFICATIONS } from "../../acp-extensions";
+import {
+  isMethod,
+  POSTHOG_METHODS,
+  POSTHOG_NOTIFICATIONS,
+} from "../../acp-extensions";
 import {
   type CodeExecutionMode,
   type CodexNativeMode,
@@ -84,6 +90,7 @@ export interface CodexAcpAgentOptions {
 
 type CodexSession = BaseSession & {
   settingsManager: CodexSettingsManager;
+  promptRunning: boolean;
 };
 
 function toCodexPermissionMode(mode?: string): PermissionMode {
@@ -156,6 +163,11 @@ export class CodexAcpAgent extends BaseAcpAgent {
    * single-owner.
    */
   private promptMutex: Promise<unknown> = Promise.resolve();
+  private readonly codexProcessOptions: CodexProcessOptions;
+  private readonly processCallbacks?: ProcessSpawnedCallback;
+  // Snapshot of the initialize() request so refreshSession can replay the
+  // same handshake against a respawned codex-acp subprocess.
+  private lastInitRequest?: InitializeRequest;
 
   constructor(client: AgentSideConnection, options: CodexAcpAgentOptions) {
     super(client);
@@ -166,6 +178,9 @@ export class CodexAcpAgent extends BaseAcpAgent {
     const cwd = options.codexProcessOptions.cwd ?? process.cwd();
     const settingsManager = new CodexSettingsManager(cwd);
 
+    this.codexProcessOptions = options.codexProcessOptions;
+    this.processCallbacks = options.processCallbacks;
+
     // Spawn the codex-acp subprocess
     this.codexProcess = spawnCodexProcess({
       ...options.codexProcessOptions,
@@ -185,6 +200,7 @@ export class CodexAcpAgent extends BaseAcpAgent {
       settingsManager,
       notificationHistory: [],
       cancelled: false,
+      promptRunning: false,
     };
 
     this.sessionState = createSessionState("", cwd);
@@ -203,6 +219,9 @@ export class CodexAcpAgent extends BaseAcpAgent {
     // Initialize settings
     await this.session.settingsManager.initialize();
 
+    // Snapshot the handshake so refreshSession can replay it after respawn.
+    this.lastInitRequest = request;
+
     // Forward to codex-acp
     const response = await this.codexConnection.initialize(request);
 
@@ -427,9 +446,13 @@ export class CodexAcpAgent extends BaseAcpAgent {
     // injected PR context is not rendered as a user message.
     await this.broadcastUserMessage(params);
 
-    const response = await this.codexConnection.prompt(
-      prependPrContext(params),
-    );
+    this.session.promptRunning = true;
+    let response: PromptResponse;
+    try {
+      response = await this.codexConnection.prompt(prependPrContext(params));
+    } finally {
+      this.session.promptRunning = false;
+    }
 
     // Usage is already accumulated via sessionUpdate notifications in
     // codex-client.ts. Do NOT also add response.usage here or tokens
@@ -491,6 +514,125 @@ export class CodexAcpAgent extends BaseAcpAgent {
     }
   }
 
+  /**
+   * Refresh the session between turns. Currently the only refreshable field
+   * is `mcpServers`. Unlike Claude (where we rebuild an in-process Query with
+   * `resume`), Codex runs as a `codex-acp` subprocess whose MCP set is bound
+   * at `newSession`/`loadSession` time and whose user-local MCPs are disabled
+   * via spawn-time `-c mcp_servers.<name>.enabled=false` CLI args. To
+   * guarantee the caller-supplied set fully wins, we respawn the subprocess
+   * and rehydrate the session via `loadSession` — codex-acp persists sessions
+   * to disk, so conversation history is preserved.
+   *
+   * This is an `extMethod` (request/response), not `extNotification`, so the
+   * caller can await completion before sending the next prompt.
+   *
+   * Caller contract: only call REFRESH_SESSION between turns (no prompt in flight).
+   */
+  async extMethod(
+    method: string,
+    params: Record<string, unknown>,
+  ): Promise<Record<string, unknown>> {
+    if (!isMethod(method, POSTHOG_METHODS.REFRESH_SESSION)) {
+      throw RequestError.methodNotFound(method);
+    }
+
+    // Trust boundary: refresh is only safe when the caller is trusted infra
+    // (e.g. the sandbox agent-server). Do not route this method from
+    // untrusted clients — mcpServers contents are forwarded verbatim to
+    // codex-acp with no URL/command validation.
+    if (params.mcpServers === undefined) {
+      throw new RequestError(
+        -32602,
+        "refresh_session requires at least one refreshable field (e.g. mcpServers)",
+      );
+    }
+    if (!Array.isArray(params.mcpServers)) {
+      throw new RequestError(
+        -32602,
+        "refresh_session: mcpServers must be an array",
+      );
+    }
+
+    await this.refreshSession(params.mcpServers as McpServer[]);
+    return { refreshed: true };
+  }
+
+  private async refreshSession(mcpServers: McpServer[]): Promise<void> {
+    const prev = this.session;
+    if (prev.promptRunning) {
+      throw new RequestError(
+        -32002,
+        "Cannot refresh session while a prompt turn is in flight",
+      );
+    }
+
+    this.logger.info("Refreshing Codex session with fresh MCP servers", {
+      serverCount: mcpServers.length,
+      sessionId: this.sessionId,
+    });
+
+    // Abort FIRST so any stuck in-flight ACP request unblocks — otherwise
+    // cancel() can deadlock waiting on a codex-acp call that never returns.
+    prev.abortController.abort();
+    try {
+      await this.codexConnection.cancel({ sessionId: this.sessionId });
+    } catch (err) {
+      this.logger.warn("cancel() during refresh failed (non-fatal)", {
+        error: err,
+      });
+    }
+    this.codexProcess.kill();
+
+    // Respawn with the same options and a fresh settings manager rooted at
+    // the current cwd (so the `mcp_servers.<name>.enabled=false` args are
+    // regenerated from the latest ~/.codex/config.toml).
+    const cwd = prev.settingsManager.getCwd();
+    const newSettingsManager = new CodexSettingsManager(cwd);
+    await newSettingsManager.initialize();
+
+    const newProcess = spawnCodexProcess({
+      ...this.codexProcessOptions,
+      cwd,
+      settings: newSettingsManager.getSettings(),
+      logger: this.logger,
+      processCallbacks: this.processCallbacks,
+    });
+
+    const codexReadable = nodeReadableToWebReadable(newProcess.stdout);
+    const codexWritable = nodeWritableToWebWritable(newProcess.stdin);
+    const codexStream = ndJsonStream(codexWritable, codexReadable);
+
+    const newAbortController = new AbortController();
+    const newConnection = new ClientSideConnection(
+      (_agent) =>
+        createCodexClient(this.client, this.logger, this.sessionState),
+      codexStream,
+    );
+
+    // Re-run ACP init on the new subprocess, then rehydrate the session with
+    // the new MCP set. loadSession is codex-acp's equivalent of Claude's
+    // `resume` — conversation history is restored from disk.
+    const initRequest: InitializeRequest = this.lastInitRequest ?? {
+      protocolVersion: 1,
+    };
+    await newConnection.initialize(initRequest);
+    await newConnection.loadSession({
+      sessionId: this.sessionId,
+      cwd: this.sessionState.cwd,
+      mcpServers,
+    });
+
+    // Swap everything at once so closeSession/prompt/cancel target the new
+    // subprocess going forward. Preserve sessionState (accumulatedUsage,
+    // taskRunId, configOptions) untouched.
+    this.codexProcess = newProcess;
+    this.codexConnection = newConnection;
+    prev.settingsManager.dispose();
+    prev.settingsManager = newSettingsManager;
+    prev.abortController = newAbortController;
+  }
+
   async setSessionMode(
     params: SetSessionModeRequest,
   ): Promise<SetSessionModeResponse> {