diff --git a/packages/ui/src/features/agent-applications/agent-builder/AgentBuilderDock.tsx b/packages/ui/src/features/agent-applications/agent-builder/AgentBuilderDock.tsx index 683e608447..fdb975bbe4 100644 --- a/packages/ui/src/features/agent-applications/agent-builder/AgentBuilderDock.tsx +++ b/packages/ui/src/features/agent-applications/agent-builder/AgentBuilderDock.tsx @@ -58,10 +58,14 @@ function buildAgentBuilderContext( ): Record { 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 diff --git a/packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.test.ts b/packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.test.ts new file mode 100644 index 0000000000..62e250600b --- /dev/null +++ b/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, +): 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", + }); + }); +}); diff --git a/packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.ts b/packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.ts index f296f0be4d..dcbc19adea 100644 --- a/packages/ui/src/features/agent-applications/agent-builder/useAgentBuilderClientTools.ts +++ b/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; 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 => { + 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 => { + 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], ); } diff --git a/packages/ui/src/features/agent-applications/components/AgentConfigurationPane.tsx b/packages/ui/src/features/agent-applications/components/AgentConfigurationPane.tsx index 12d40b3329..1ad2aca012 100644 --- a/packages/ui/src/features/agent-applications/components/AgentConfigurationPane.tsx +++ b/packages/ui/src/features/agent-applications/components/AgentConfigurationPane.tsx @@ -451,7 +451,12 @@ export function AgentConfigurationPane({ ) : null; return ( - + {!revisionId ? (