@posthog/agent 2.3.297 → 2.3.302

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@posthog/agent",
-  "version": "2.3.297",
+  "version": "2.3.302",
   "repository": "https://github.com/PostHog/code",
   "description": "TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog",
   "exports": {
@@ -86,8 +86,8 @@
     "tsx": "^4.20.6",
     "typescript": "^5.5.0",
     "vitest": "^2.1.8",
-    "@posthog/shared": "1.0.0",
-    "@posthog/git": "1.0.0"
+    "@posthog/git": "1.0.0",
+    "@posthog/shared": "1.0.0"
   },
   "dependencies": {
     "@agentclientprotocol/sdk": "0.16.1",

package/src/acp-extensions.test.ts

@@ -0,0 +1,72 @@
+import { describe, expect, it } from "vitest";
+import {
+  isMethod,
+  isNotification,
+  POSTHOG_METHODS,
+  POSTHOG_NOTIFICATIONS,
+} from "./acp-extensions";
+
+describe("isNotification", () => {
+  it("matches the exact notification name", () => {
+    expect(
+      isNotification(
+        POSTHOG_NOTIFICATIONS.TURN_COMPLETE,
+        POSTHOG_NOTIFICATIONS.TURN_COMPLETE,
+      ),
+    ).toBe(true);
+  });
+
+  it("matches the double-underscore prefix variant", () => {
+    expect(
+      isNotification(
+        `_${POSTHOG_NOTIFICATIONS.TURN_COMPLETE}`,
+        POSTHOG_NOTIFICATIONS.TURN_COMPLETE,
+      ),
+    ).toBe(true);
+  });
+
+  it("returns false for a different notification", () => {
+    expect(
+      isNotification(
+        POSTHOG_NOTIFICATIONS.USAGE_UPDATE,
+        POSTHOG_NOTIFICATIONS.TURN_COMPLETE,
+      ),
+    ).toBe(false);
+  });
+
+  it("returns false for undefined", () => {
+    expect(isNotification(undefined, POSTHOG_NOTIFICATIONS.TURN_COMPLETE)).toBe(
+      false,
+    );
+  });
+});
+
+describe("isMethod", () => {
+  it("matches the exact method name", () => {
+    expect(
+      isMethod(
+        POSTHOG_METHODS.REFRESH_SESSION,
+        POSTHOG_METHODS.REFRESH_SESSION,
+      ),
+    ).toBe(true);
+  });
+
+  it("matches the double-underscore prefix variant", () => {
+    expect(
+      isMethod(
+        `_${POSTHOG_METHODS.REFRESH_SESSION}`,
+        POSTHOG_METHODS.REFRESH_SESSION,
+      ),
+    ).toBe(true);
+  });
+
+  it("returns false for unrelated method strings", () => {
+    expect(isMethod("session/prompt", POSTHOG_METHODS.REFRESH_SESSION)).toBe(
+      false,
+    );
+  });
+
+  it("returns false for undefined", () => {
+    expect(isMethod(undefined, POSTHOG_METHODS.REFRESH_SESSION)).toBe(false);
+  });
+});

package/src/acp-extensions.ts

@@ -68,17 +68,48 @@ export const POSTHOG_NOTIFICATIONS = {
   PERMISSION_RESPONSE: "_posthog/permission_response",
 } as const;
 
-type NotificationMethod =
+/**
+ * Custom request methods for PostHog-specific operations that need a response
+ * (request/response, not fire-and-forget). Used with
+ * ClientSideConnection.extMethod() on the sender and Agent.extMethod() on the
+ * receiver.
+ */
+export const POSTHOG_METHODS = {
+  /**
+   * Client requests a session refresh between turns. Payload may include
+   * `mcpServers` to trigger a resume-with-new-options reinit; future fields
+   * can extend this without adding new methods. Returns once the refresh has
+   * completed so the caller can safely send the next prompt.
+   */
+  REFRESH_SESSION: "_posthog/refresh_session",
+} as const;
+
+type PosthogNotification =
   (typeof POSTHOG_NOTIFICATIONS)[keyof typeof POSTHOG_NOTIFICATIONS];
 
+type PosthogMethod = (typeof POSTHOG_METHODS)[keyof typeof POSTHOG_METHODS];
+
 /**
- * Check if an ACP method matches a PostHog notification, handling the
- * possible `__posthog/` double-prefix from extNotification().
+ * Does `method` match `expected`? Shared by notification and method matchers.
+ * Handles the `__posthog/` double-prefix that extNotification() can produce.
  */
+function matchesExt(method: string | undefined, expected: string): boolean {
+  if (!method) return false;
+  return method === expected || method === `_${expected}`;
+}
+
+/** Dispatcher check for incoming `extNotification` calls on the agent side. */
 export function isNotification(
   method: string | undefined,
-  notification: NotificationMethod,
+  expected: PosthogNotification,
 ): boolean {
-  if (!method) return false;
-  return method === notification || method === `_${notification}`;
+  return matchesExt(method, expected);
+}
+
+/** Dispatcher check for incoming `extMethod` calls on the agent side. */
+export function isMethod(
+  method: string | undefined,
+  expected: PosthogMethod,
+): boolean {
+  return matchesExt(method, expected);
 }

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

@@ -0,0 +1,292 @@
+import type { AgentSideConnection } from "@agentclientprotocol/sdk";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+import { POSTHOG_METHODS } from "../../acp-extensions";
+import { Pushable } from "../../utils/streams";
+
+type InitResult = {
+  result: "success";
+  commands?: unknown[];
+  models?: unknown[];
+};
+
+type SdkQueryHandle = {
+  interrupt: ReturnType<typeof vi.fn>;
+  setModel: ReturnType<typeof vi.fn>;
+  setMcpServers: ReturnType<typeof vi.fn>;
+  supportedCommands: ReturnType<typeof vi.fn>;
+  initializationResult: ReturnType<typeof vi.fn>;
+  [Symbol.asyncIterator]: () => AsyncIterator<never>;
+};
+
+let nextInitPromise: Promise<InitResult> = Promise.resolve({
+  result: "success",
+  commands: [],
+  models: [],
+});
+
+function makeQueryHandle(): SdkQueryHandle {
+  return {
+    interrupt: vi.fn().mockResolvedValue(undefined),
+    setModel: vi.fn().mockResolvedValue(undefined),
+    setMcpServers: vi.fn().mockResolvedValue(undefined),
+    supportedCommands: vi.fn().mockResolvedValue([]),
+    initializationResult: vi.fn().mockImplementation(() => nextInitPromise),
+    [Symbol.asyncIterator]: async function* () {
+      /* never yields */
+    } as never,
+  };
+}
+
+const lastQueryCall: { options?: Record<string, unknown> } = {};
+const createdQueries: SdkQueryHandle[] = [];
+
+vi.mock("@anthropic-ai/claude-agent-sdk", () => ({
+  query: vi.fn((params: { options: Record<string, unknown> }) => {
+    lastQueryCall.options = params.options;
+    const handle = makeQueryHandle();
+    createdQueries.push(handle);
+    return handle;
+  }),
+}));
+
+const fetchMcpToolMetadataMock = vi.fn().mockResolvedValue(undefined);
+vi.mock("./mcp/tool-metadata", () => ({
+  fetchMcpToolMetadata: fetchMcpToolMetadataMock,
+  getConnectedMcpServerNames: vi.fn().mockReturnValue([]),
+}));
+
+// Import after the mocks so ClaudeAcpAgent resolves the mocked SDK
+const { ClaudeAcpAgent } = await import("./claude-agent");
+type Agent = InstanceType<typeof ClaudeAcpAgent>;
+
+function makeAgent(): Agent {
+  const client = {
+    sessionUpdate: vi.fn().mockResolvedValue(undefined),
+    extNotification: vi.fn().mockResolvedValue(undefined),
+  } as unknown as AgentSideConnection;
+  return new ClaudeAcpAgent(client);
+}
+
+function installFakeSession(
+  agent: Agent,
+  sessionId: string,
+  overrides: Partial<{ modelId: string }> = {},
+) {
+  const oldQuery = makeQueryHandle();
+  const input = new Pushable();
+  const endSpy = vi.spyOn(input, "end");
+  const abortController = new AbortController();
+
+  const session = {
+    query: oldQuery,
+    queryOptions: {
+      sessionId,
+      cwd: "/tmp/repo",
+      model: "claude-sonnet-4-6",
+      mcpServers: { posthog: { type: "http", url: "https://old" } },
+      abortController,
+    },
+    input,
+    cancelled: false,
+    settingsManager: { dispose: vi.fn() },
+    permissionMode: "default",
+    abortController,
+    accumulatedUsage: {
+      inputTokens: 42,
+      outputTokens: 17,
+      cachedReadTokens: 0,
+      cachedWriteTokens: 0,
+    },
+    configOptions: [],
+    promptRunning: false,
+    pendingMessages: new Map(),
+    nextPendingOrder: 0,
+    cwd: "/tmp/repo",
+    notificationHistory: [{ foo: "bar" }],
+    taskRunId: "run-1",
+    modelId: overrides.modelId,
+  } as unknown as Parameters<typeof Object.assign>[0];
+
+  (agent as unknown as { session: unknown }).session = session;
+  (agent as unknown as { sessionId: string }).sessionId = sessionId;
+
+  return { session, oldQuery, endSpy, abortController };
+}
+
+const freshMcpServers = [
+  {
+    name: "posthog",
+    type: "http" as const,
+    url: "https://fresh",
+    headers: [{ name: "x-foo", value: "bar" }],
+  },
+];
+
+describe("ClaudeAcpAgent.extMethod refresh_session", () => {
+  beforeEach(() => {
+    lastQueryCall.options = undefined;
+    createdQueries.length = 0;
+    nextInitPromise = Promise.resolve({
+      result: "success",
+      commands: [],
+      models: [],
+    });
+    fetchMcpToolMetadataMock.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 payload has no refreshable fields", async () => {
+    const agent = makeAgent();
+    installFakeSession(agent, "s-empty");
+
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {}),
+    ).rejects.toThrow(/requires at least one refreshable field/);
+  });
+
+  it("rejects when mcpServers is not an array", async () => {
+    const agent = makeAgent();
+    installFakeSession(agent, "s-malformed");
+
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: "not-an-array",
+      }),
+    ).rejects.toThrow(/mcpServers must be an array/);
+  });
+
+  it("rejects refresh while a prompt is in flight", async () => {
+    const agent = makeAgent();
+    const { session } = installFakeSession(agent, "s-1");
+    (session as unknown as { promptRunning: boolean }).promptRunning = true;
+
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: freshMcpServers,
+      }),
+    ).rejects.toThrow(/prompt turn is in flight/);
+  });
+
+  it("rejects when session model does not support MCP injection", async () => {
+    const agent = makeAgent();
+    installFakeSession(agent, "s-haiku", { modelId: "claude-haiku-4-5" });
+
+    await expect(
+      agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: freshMcpServers,
+      }),
+    ).rejects.toThrow(/does not support MCP injection/);
+  });
+
+  it("throws when initialization of the new query times out", async () => {
+    vi.useFakeTimers();
+    try {
+      const agent = makeAgent();
+      installFakeSession(agent, "s-timeout");
+      // Never resolves — withTimeout must win the race.
+      nextInitPromise = new Promise<InitResult>(() => {});
+
+      const promise = agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+        mcpServers: freshMcpServers,
+      });
+      // Drop the rejection on the floor so an unhandled-rejection warning
+      // doesn't race the assertion below.
+      promise.catch(() => {});
+
+      await vi.advanceTimersByTimeAsync(30_001);
+
+      await expect(promise).rejects.toThrow(
+        /Session refresh timed out for s-timeout/,
+      );
+    } finally {
+      vi.useRealTimers();
+    }
+  });
+
+  it("swaps query/input/options and preserves session state", async () => {
+    const agent = makeAgent();
+    const { session, oldQuery, endSpy } = installFakeSession(agent, "s-2");
+
+    const result = await agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+      mcpServers: freshMcpServers,
+    });
+
+    expect(result).toEqual({ refreshed: true });
+    expect(oldQuery.interrupt).toHaveBeenCalledTimes(1);
+    expect(endSpy).toHaveBeenCalledTimes(1);
+
+    // New query was built with resume identity (not sessionId) and new servers
+    expect(lastQueryCall.options).toMatchObject({
+      resume: "s-2",
+      forkSession: false,
+      mcpServers: {
+        posthog: {
+          type: "http",
+          url: "https://fresh",
+          headers: { "x-foo": "bar" },
+        },
+      },
+    });
+    expect(lastQueryCall.options?.sessionId).toBeUndefined();
+
+    // Session fields swapped to the new instances
+    const updated = session as unknown as {
+      query: SdkQueryHandle;
+      input: unknown;
+      queryOptions: Record<string, unknown>;
+      accumulatedUsage: { inputTokens: number };
+      notificationHistory: unknown[];
+    };
+    expect(updated.query).toBe(createdQueries[0]);
+    expect(updated.query).not.toBe(oldQuery);
+    expect(updated.input).toBeInstanceOf(Pushable);
+    expect(updated.queryOptions).toBe(lastQueryCall.options);
+
+    // Preserves session-level state (usage, notification history)
+    expect(updated.accumulatedUsage.inputTokens).toBe(42);
+    expect(updated.notificationHistory).toEqual([{ foo: "bar" }]);
+  });
+
+  it("aborts the old controller and allocates a fresh one for the new query", async () => {
+    const agent = makeAgent();
+    const { session, abortController: oldController } = installFakeSession(
+      agent,
+      "s-abort",
+    );
+
+    await agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+      mcpServers: freshMcpServers,
+    });
+
+    expect(oldController.signal.aborted).toBe(true);
+
+    const updated = session as unknown as {
+      abortController: AbortController;
+      queryOptions: { abortController: AbortController };
+    };
+    expect(updated.abortController).not.toBe(oldController);
+    expect(updated.abortController.signal.aborted).toBe(false);
+    expect(updated.queryOptions.abortController).toBe(updated.abortController);
+    expect(lastQueryCall.options?.abortController).toBe(
+      updated.abortController,
+    );
+  });
+
+  it("re-fetches MCP tool metadata for the new query", async () => {
+    const agent = makeAgent();
+    installFakeSession(agent, "s-metadata");
+
+    await agent.extMethod(POSTHOG_METHODS.REFRESH_SESSION, {
+      mcpServers: freshMcpServers,
+    });
+
+    expect(fetchMcpToolMetadataMock).toHaveBeenCalledTimes(1);
+    expect(fetchMcpToolMetadataMock.mock.calls[0][0]).toBe(createdQueries[0]);
+  });
+});

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

@@ -38,13 +38,19 @@ import {
   type CanUseTool,
   getSessionMessages,
   listSessions,
+  type McpServerConfig,
+  type Options,
   type Query,
   query,
   type SDKUserMessage,
 } from "@anthropic-ai/claude-agent-sdk";
 import { v7 as uuidv7 } from "uuid";
 import packageJson from "../../../package.json" with { type: "json" };
-import { POSTHOG_NOTIFICATIONS } from "../../acp-extensions";
+import {
+  isMethod,
+  POSTHOG_METHODS,
+  POSTHOG_NOTIFICATIONS,
+} from "../../acp-extensions";
 import { unreachable, withTimeout } from "../../utils/common";
 import { Logger } from "../../utils/logger";
 import { Pushable } from "../../utils/streams";
@@ -652,6 +658,120 @@ export class ClaudeAcpAgent extends BaseAcpAgent {
     await this.session.query.interrupt();
   }
 
+  /**
+   * Refresh the session between turns. Currently the only refreshable field
+   * is `mcpServers` — a resume-with-new-options reinit that bakes the servers
+   * into query() options (preserving conversation history via resume).
+   *
+   * This is an `extMethod` (request/response), not `extNotification`, so the
+   * caller can await completion before sending the next prompt. The sandbox
+   * agent-server uses this on pre-prompt TTL checks.
+   *
+   * Why resume+rebuild instead of query.setMcpServers()?
+   * setMcpServers() does NOT always overwrite servers installed by local/plugin
+   * config — it can non-deterministically surface either the config-provided
+   * server or the plugin-installed one. In the sandbox, repos may have Claude
+   * plugins with their own MCPs, and we want the CLI-supplied set to fully win.
+   * Passing mcpServers via query() options (as a "managed"/static set) has that
+   * overwrite guarantee, so we tear down the current Query and construct a new
+   * one with resume.
+   *
+   * 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 — parseMcpServers does 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",
+      );
+    }
+
+    const mcpServers = parseMcpServers(
+      params as Pick<NewSessionRequest, "mcpServers">,
+    );
+    await this.refreshSession(mcpServers);
+    return { refreshed: true };
+  }
+
+  private async refreshSession(
+    mcpServers: Record<string, McpServerConfig>,
+  ): Promise<void> {
+    const prev = this.session;
+    if (prev.promptRunning) {
+      throw new RequestError(
+        -32002,
+        "Cannot refresh session while a prompt turn is in flight",
+      );
+    }
+    if (prev.modelId && !supportsMcpInjection(prev.modelId)) {
+      throw new RequestError(
+        -32002,
+        `Model ${prev.modelId} does not support MCP injection; cannot refresh`,
+      );
+    }
+
+    this.logger.info("Refreshing session with fresh MCP servers", {
+      serverCount: Object.keys(mcpServers).length,
+      sessionId: this.sessionId,
+    });
+
+    // Abort FIRST so any stuck in-flight HTTP request unblocks — otherwise
+    // interrupt() can deadlock waiting on an API call that never returns.
+    // We allocate a fresh controller for the new Query below so aborting
+    // the old one doesn't poison it.
+    prev.abortController.abort();
+    await prev.query.interrupt();
+    prev.input.end();
+
+    // Reuse every option from the running session; swap mcpServers, re-root
+    // identity on `resume` instead of `sessionId`, and give the new Query a
+    // fresh AbortController.
+    const newAbortController = new AbortController();
+    const { sessionId: _drop, ...rest } = prev.queryOptions;
+    const newOptions: Options = {
+      ...rest,
+      mcpServers,
+      resume: this.sessionId,
+      forkSession: false,
+      abortController: newAbortController,
+    };
+
+    const newInput = new Pushable<SDKUserMessage>();
+    const newQuery = query({ prompt: newInput, options: newOptions });
+
+    prev.query = newQuery;
+    prev.input = newInput;
+    prev.queryOptions = newOptions;
+    prev.abortController = newAbortController;
+
+    const result = await withTimeout(
+      newQuery.initializationResult(),
+      SESSION_VALIDATION_TIMEOUT_MS,
+    );
+    if (result.result === "timeout") {
+      throw new Error(`Session refresh timed out for ${this.sessionId}`);
+    }
+
+    // Re-fetch MCP tool metadata + slash commands — the server list changed.
+    this.deferBackgroundFetches(newQuery);
+  }
+
   async unstable_setSessionModel(
     params: SetSessionModelRequest,
   ): Promise<SetSessionModelResponse | undefined> {

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

@@ -60,20 +60,32 @@ describe("CodexAcpAgent", () => {
     vi.clearAllMocks();
   });
 
-  function createAgent(): CodexAcpAgent {
+  function createAgent(overrides: Partial<AgentSideConnection> = {}): {
+    agent: CodexAcpAgent;
+    client: AgentSideConnection & {
+      extNotification: ReturnType<typeof vi.fn>;
+      sessionUpdate: ReturnType<typeof vi.fn>;
+    };
+  } {
     const client = {
       extNotification: vi.fn(),
-    } as unknown as AgentSideConnection;
-
-    return new CodexAcpAgent(client, {
+      sessionUpdate: vi.fn(),
+      ...overrides,
+    } as unknown as AgentSideConnection & {
+      extNotification: ReturnType<typeof vi.fn>;
+      sessionUpdate: ReturnType<typeof vi.fn>;
+    };
+
+    const agent = new CodexAcpAgent(client, {
       codexProcessOptions: {
         cwd: process.cwd(),
       },
     });
+    return { agent, client };
   }
 
   it("applies the requested initial mode for a new session", async () => {
-    const agent = createAgent();
+    const { agent } = createAgent();
     mockCodexConnection.newSession.mockResolvedValue({
       sessionId: "session-1",
       modes: { currentModeId: "auto", availableModes: [] },
@@ -96,7 +108,7 @@ describe("CodexAcpAgent", () => {
   });
 
   it("preserves the live session mode when loading an existing session", async () => {
-    const agent = createAgent();
+    const { agent } = createAgent();
     mockCodexConnection.loadSession.mockResolvedValue({
       modes: { currentModeId: "read-only", availableModes: [] },
       configOptions: [],
@@ -114,4 +126,53 @@ describe("CodexAcpAgent", () => {
         .sessionState.permissionMode,
     ).toBe("read-only");
   });
+
+  it("broadcasts user prompt as user_message_chunk before delegating to codex-acp", async () => {
+    const { agent, client } = createAgent();
+    // Seed an active session so prompt() has the state it expects.
+    mockCodexConnection.newSession.mockResolvedValue({
+      sessionId: "session-1",
+      modes: { currentModeId: "auto", availableModes: [] },
+      configOptions: [],
+    } satisfies Partial<NewSessionResponse>);
+    await agent.newSession({
+      cwd: process.cwd(),
+    } as never);
+
+    const callOrder: string[] = [];
+    client.sessionUpdate.mockImplementation(async () => {
+      callOrder.push("sessionUpdate");
+    });
+    mockCodexConnection.prompt.mockImplementation(async () => {
+      callOrder.push("prompt");
+      return { stopReason: "end_turn" };
+    });
+
+    await agent.prompt({
+      sessionId: "session-1",
+      prompt: [
+        { type: "text", text: "first chunk" },
+        { type: "text", text: "second chunk" },
+      ],
+    } as never);
+
+    expect(client.sessionUpdate).toHaveBeenCalledTimes(2);
+    expect(client.sessionUpdate).toHaveBeenNthCalledWith(1, {
+      sessionId: "session-1",
+      update: {
+        sessionUpdate: "user_message_chunk",
+        content: { type: "text", text: "first chunk" },
+      },
+    });
+    expect(client.sessionUpdate).toHaveBeenNthCalledWith(2, {
+      sessionId: "session-1",
+      update: {
+        sessionUpdate: "user_message_chunk",
+        content: { type: "text", text: "second chunk" },
+      },
+    });
+    // Broadcast must land before the prompt reaches codex-acp so the user
+    // turn is persisted even if the underlying prompt fails.
+    expect(callOrder).toEqual(["sessionUpdate", "sessionUpdate", "prompt"]);
+  });
 });