fix(agent-builder): resolve target revision for set_secret/connect_mcp punch-outs

packages/ui/src/features/agent-applications/agent-builder/AgentBuilderDock.tsx

@@ -58,10 +58,14 @@ function buildAgentBuilderContext(
 ): Record<string, unknown> {
   const agent = "slug" in page ? page.slug : undefined;
   const sessionId = page.kind === "agent-session" ? page.sessionId : undefined;
+  const revisionId = page.kind === "agent-config" ? page.revision : undefined;
   return {
     page: page.kind,
     agent,
     session_id: sessionId,
+    // The revision open in the configuration pane — the default target for
+    // revision-scoped punch-outs (`set_secret`, `connect_mcp`).
+    revision_id: revisionId,
     follow_enabled: followEnabled,
     // The project the user is currently in — the agent threads this into the
     // `project_id` arg of every `@posthog/*` tool (it's tenant-neutral and acts

packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.test.ts

@@ -0,0 +1,231 @@
+import { renderHook } from "@testing-library/react";
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+vi.mock("@tanstack/react-router", () => ({
+  useNavigate: () => vi.fn(),
+}));
+
+const client = {
+  getAgentApplication: vi.fn(),
+  listAgentRevisions: vi.fn(),
+  getAgentRevision: vi.fn(),
+};
+
+vi.mock("@posthog/ui/features/auth/authClient", () => ({
+  useAuthenticatedClient: () => client,
+}));
+
+import type { ClientToolCallData } from "@posthog/core/agent-chat/identifiers";
+import { useAgentBuilderStore } from "./agentBuilderStore";
+import { useAgentBuilderClientTools } from "./useAgentBuilderClientTools";
+
+function call(
+  tool_id: string,
+  args: Record<string, unknown>,
+): ClientToolCallData {
+  return { call_id: "call-1", tool_id, args };
+}
+
+function handler() {
+  return renderHook(() => useAgentBuilderClientTools()).result.current;
+}
+
+describe("useAgentBuilderClientTools revision resolution", () => {
+  beforeEach(() => {
+    client.getAgentApplication.mockReset();
+    client.listAgentRevisions.mockReset();
+    client.getAgentRevision.mockReset();
+    useAgentBuilderStore.setState({
+      page: { kind: "unknown" },
+      pendingSecret: null,
+      pendingMcpConnect: null,
+    });
+  });
+
+  it("uses an explicit revision_id once it verifies as belonging to the agent", async () => {
+    client.getAgentRevision.mockResolvedValue({ id: "rev-explicit" });
+
+    const outcome = await handler()(
+      call("set_secret", {
+        agent_slug: "my-agent",
+        secret: "API_KEY",
+        revision_id: "rev-explicit",
+      }),
+    );
+
+    expect(outcome).toEqual({ defer: true });
+    expect(useAgentBuilderStore.getState().pendingSecret).toMatchObject({
+      agentSlug: "my-agent",
+      secret: "API_KEY",
+      revisionId: "rev-explicit",
+    });
+    expect(client.getAgentRevision).toHaveBeenCalledWith(
+      "my-agent",
+      "rev-explicit",
+    );
+    expect(client.getAgentApplication).not.toHaveBeenCalled();
+    expect(client.listAgentRevisions).not.toHaveBeenCalled();
+  });
+
+  it("rejects an explicit revision_id that does not belong to the agent", async () => {
+    // The nested revision route 404s (→ null) for another agent's revision.
+    client.getAgentRevision.mockResolvedValue(null);
+
+    const outcome = await handler()(
+      call("set_secret", {
+        agent_slug: "my-agent",
+        secret: "API_KEY",
+        revision_id: "rev-of-other-agent",
+      }),
+    );
+
+    expect(outcome).toEqual({
+      error: "revision_not_found: rev-of-other-agent on my-agent",
+    });
+    expect(useAgentBuilderStore.getState().pendingSecret).toBeNull();
+  });
+
+  it("falls back to the revision open on this agent's config page", async () => {
+    useAgentBuilderStore.setState({
+      page: { kind: "agent-config", slug: "my-agent", revision: "rev-page" },
+    });
+
+    const outcome = await handler()(
+      call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
+    );
+
+    expect(outcome).toEqual({ defer: true });
+    expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
+      "rev-page",
+    );
+    expect(client.getAgentApplication).not.toHaveBeenCalled();
+  });
+
+  it("rotation targets live even while a draft config page is open", async () => {
+    useAgentBuilderStore.setState({
+      page: { kind: "agent-config", slug: "my-agent", revision: "rev-draft" },
+    });
+    client.getAgentApplication.mockResolvedValue({ live_revision: "rev-live" });
+    client.listAgentRevisions.mockResolvedValue([
+      { id: "rev-draft", state: "draft" },
+      { id: "rev-live", state: "ready" },
+    ]);
+
+    await handler()(
+      call("set_secret", {
+        agent_slug: "my-agent",
+        secret: "API_KEY",
+        mode: "rotate",
+      }),
+    );
+
+    expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
+      "rev-live",
+    );
+  });
+
+  it("ignores the page revision when it belongs to a different agent", async () => {
+    useAgentBuilderStore.setState({
+      page: { kind: "agent-config", slug: "other-agent", revision: "rev-page" },
+    });
+    client.getAgentApplication.mockResolvedValue({ live_revision: null });
+    client.listAgentRevisions.mockResolvedValue([
+      { id: "rev-draft", state: "draft" },
+    ]);
+
+    await handler()(
+      call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
+    );
+
+    expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
+      "rev-draft",
+    );
+  });
+
+  it.each([
+    // A new secret targets the draft being authored (secrets only copy
+    // forward at draft creation), even when a live revision exists.
+    { mode: "set", expected: "rev-draft" },
+    // A rotation targets what's running.
+    { mode: "rotate", expected: "rev-live" },
+  ])(
+    "set_secret mode=$mode resolves to $expected via the API",
+    async ({ mode, expected }) => {
+      client.getAgentApplication.mockResolvedValue({
+        live_revision: "rev-live",
+      });
+      client.listAgentRevisions.mockResolvedValue([
+        { id: "rev-draft", state: "draft" },
+        { id: "rev-live", state: "ready" },
+      ]);
+
+      await handler()(
+        call("set_secret", { agent_slug: "my-agent", secret: "API_KEY", mode }),
+      );
+
+      expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
+        expected,
+      );
+    },
+  );
+
+  it("set_secret falls back to the newest revision when no live or draft exists", async () => {
+    client.getAgentApplication.mockResolvedValue({ live_revision: null });
+    client.listAgentRevisions.mockResolvedValue([
+      { id: "rev-newest", state: "ready" },
+      { id: "rev-older", state: "ready" },
+    ]);
+
+    await handler()(
+      call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
+    );
+
+    expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
+      "rev-newest",
+    );
+  });
+
+  it("errors when the agent has no revisions at all", async () => {
+    client.getAgentApplication.mockResolvedValue({ live_revision: null });
+    client.listAgentRevisions.mockResolvedValue([]);
+
+    const outcome = await handler()(
+      call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
+    );
+
+    expect(outcome).toEqual({ error: "no_target_revision: my-agent" });
+    expect(useAgentBuilderStore.getState().pendingSecret).toBeNull();
+  });
+
+  it("errors when revision lookup fails", async () => {
+    client.getAgentApplication.mockRejectedValue(new Error("network"));
+    client.listAgentRevisions.mockRejectedValue(new Error("network"));
+
+    const outcome = await handler()(
+      call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
+    );
+
+    expect(outcome).toEqual({ error: "no_target_revision: my-agent" });
+  });
+
+  it("connect_mcp prefers the newest draft (spec edits are draft-only)", async () => {
+    client.getAgentApplication.mockResolvedValue({
+      live_revision: "rev-live",
+    });
+    client.listAgentRevisions.mockResolvedValue([
+      { id: "rev-draft", state: "draft" },
+      { id: "rev-live", state: "ready" },
+    ]);
+
+    const outcome = await handler()(
+      call("connect_mcp", { agent_slug: "my-agent", url: "https://mcp.test" }),
+    );
+
+    expect(outcome).toEqual({ defer: true });
+    expect(useAgentBuilderStore.getState().pendingMcpConnect).toMatchObject({
+      agentSlug: "my-agent",
+      revisionId: "rev-draft",
+      url: "https://mcp.test",
+    });
+  });
+});

packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.ts

@@ -1,3 +1,4 @@
+import { useAuthenticatedClient } from "@posthog/ui/features/auth/authClient";
 import { useNavigate } from "@tanstack/react-router";
 import { useCallback, useRef } from "react";
 import type { ClientToolHandler } from "../hooks/useAgentChat";
@@ -33,6 +34,7 @@ export const AGENT_BUILDER_CLIENT_TOOLS = [
  */
 export function useAgentBuilderClientTools(): ClientToolHandler {
   const navigate = useNavigate();
+  const client = useAuthenticatedClient();
   const followMode = useAgentBuilderStore((s) => s.followMode);
   const setPendingSecret = useAgentBuilderStore((s) => s.setPendingSecret);
   const setPendingMcpConnect = useAgentBuilderStore(
@@ -47,25 +49,83 @@ export function useAgentBuilderClientTools(): ClientToolHandler {
   pageRef.current = page;
 
   return useCallback(
-    (data) => {
+    async (data) => {
       const args = (data.args ?? {}) as Record<string, unknown>;
       const str = (v: unknown) => (typeof v === "string" ? v : undefined);
 
+      // Env keys and spec edits are revision-scoped, but the punch-out tool
+      // schemas don't define `revision_id`, so the agent usually omits it.
+      // Resolve the target: explicit arg → the revision the user is viewing on
+      // this agent's config page → API fallback. Preference matters because
+      // secrets only copy forward at draft creation and spec PATCHes only land
+      // on drafts: a *new* secret or MCP connection targets the draft being
+      // authored, while a rotation targets what's running (live).
+      const resolveRevision = async (
+        agentSlug: string,
+        prefer: "live" | "draft",
+      ): Promise<string | undefined> => {
+        const p = pageRef.current;
+        // The viewed revision only stands in for authoring flows: a rotation
+        // must reach what's running even while the user is viewing a draft.
+        if (
+          prefer === "draft" &&
+          p.kind === "agent-config" &&
+          p.slug === agentSlug &&
+          p.revision
+        ) {
+          return p.revision;
+        }
+        try {
+          // A revision's `state` stays "ready" when promoted — live is the
+          // application's `live_revision` pointer, not a revision state.
+          const [app, revisions] = await Promise.all([
+            client.getAgentApplication(agentSlug),
+            client.listAgentRevisions(agentSlug),
+          ]);
+          const live = app?.live_revision ?? undefined;
+          const draft = revisions.find((r) => r.state === "draft")?.id;
+          const newest = revisions[0]?.id;
+          return prefer === "draft"
+            ? (draft ?? live ?? newest)
+            : (live ?? draft ?? newest);
+        } catch {
+          return undefined;
+        }
+      };
+
+      // An explicit `revision_id` must belong to `agent_slug` before we park
+      // the punch-out. The nested env/spec routes reject mismatches
+      // server-side anyway, but that failure would only surface at submit —
+      // after the user has already typed a secret into a doomed form.
+      const verifyExplicitRevision = async (
+        agentSlug: string,
+        revisionId: string,
+      ): Promise<boolean> => {
+        const revision = await client
+          .getAgentRevision(agentSlug, revisionId)
+          .catch(() => null);
+        return revision != null;
+      };
+
       // set_secret — interactive punch-out. Park the call (defer) and render a
-      // form; the dock PUTs the key and wakes the session on submit. Env keys
-      // are revision-scoped, so resolve the target revision from the tool args,
-      // falling back to the revision the user is currently viewing in the
-      // agent-config page.
+      // form; the dock PUTs the key and wakes the session on submit.
       if (data.tool_id === "set_secret") {
         const agentSlug = str(args.agent_slug);
         const secret = str(args.secret);
         if (!agentSlug) return { error: "missing_arg: agent_slug" };
         if (!secret) return { error: "missing_arg: secret" };
-        const p = pageRef.current;
-        const pageRevision = p.kind === "agent-config" ? p.revision : undefined;
-        const revisionId = str(args.revision_id) ?? pageRevision;
-        if (!revisionId) return { error: "missing_arg: revision_id" };
         const mode = args.mode === "rotate" ? "rotate" : "set";
+        const explicit = str(args.revision_id);
+        if (explicit && !(await verifyExplicitRevision(agentSlug, explicit))) {
+          return { error: `revision_not_found: ${explicit} on ${agentSlug}` };
+        }
+        const revisionId =
+          explicit ??
+          (await resolveRevision(
+            agentSlug,
+            mode === "rotate" ? "live" : "draft",
+          ));
+        if (!revisionId) return { error: `no_target_revision: ${agentSlug}` };
         setPendingSecret({
           callId: data.call_id,
           agentSlug,
@@ -80,15 +140,18 @@ export function useAgentBuilderClientTools(): ClientToolHandler {
       // connect_mcp — interactive punch-out. Park the call and render a prefilled
       // connect form; the dock runs the native OAuth/api-key connect (auth never
       // touches the agent), writes the resulting mcps[].connection onto the
-      // target agent's spec, and wakes the session. Like set_secret, the target
-      // revision comes from the args or the current agent-config page.
+      // target agent's spec, and wakes the session. Same revision resolution as
+      // set_secret.
       if (data.tool_id === "connect_mcp") {
         const agentSlug = str(args.agent_slug);
         if (!agentSlug) return { error: "missing_arg: agent_slug" };
-        const p = pageRef.current;
-        const pageRevision = p.kind === "agent-config" ? p.revision : undefined;
-        const revisionId = str(args.revision_id) ?? pageRevision;
-        if (!revisionId) return { error: "missing_arg: revision_id" };
+        const explicit = str(args.revision_id);
+        if (explicit && !(await verifyExplicitRevision(agentSlug, explicit))) {
+          return { error: `revision_not_found: ${explicit} on ${agentSlug}` };
+        }
+        const revisionId =
+          explicit ?? (await resolveRevision(agentSlug, "draft"));
+        if (!revisionId) return { error: `no_target_revision: ${agentSlug}` };
         setPendingMcpConnect({
           callId: data.call_id,
           agentSlug,
@@ -194,6 +257,6 @@ export function useAgentBuilderClientTools(): ClientToolHandler {
           return { result: { focused: false, reason: "unknown_focus_target" } };
       }
     },
-    [navigate, setPendingSecret, setPendingMcpConnect],
+    [navigate, client, setPendingSecret, setPendingMcpConnect],
   );
 }

packages/ui/src/features/agent-applications/components/AgentConfigurationPane.tsx

@@ -451,7 +451,12 @@ export function AgentConfigurationPane({
     ) : null;
 
   return (
-    <AgentDetailLayout idOrSlug={idOrSlug} activeTab="configuration" fill>
+    <AgentDetailLayout
+      idOrSlug={idOrSlug}
+      activeTab="configuration"
+      fill
+      configRevision={revisionId}
+    >
       {!revisionId ? (
         <div className="p-6">
           <AgentDetailEmptyState