Skip to content
This repository was archived by the owner on Aug 6, 2026. It is now read-only.

Commit 76ba72c

Browse files
dmarticusdmarticus
authored
fix(agent-builder): resolve target revision for set_secret/connect_mcp punch-outs (#3323)
1 parent 4d3addd commit 76ba72c

5 files changed

Lines changed: 334 additions & 19 deletions

File tree

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -58,10 +58,14 @@ function buildAgentBuilderContext(
5858
): Record<string, unknown> {
5959
const agent = "slug" in page ? page.slug : undefined;
6060
const sessionId = page.kind === "agent-session" ? page.sessionId : undefined;
61+
const revisionId = page.kind === "agent-config" ? page.revision : undefined;
6162
return {
6263
page: page.kind,
6364
agent,
6465
session_id: sessionId,
66+
// The revision open in the configuration pane — the default target for
67+
// revision-scoped punch-outs (`set_secret`, `connect_mcp`).
68+
revision_id: revisionId,
6569
follow_enabled: followEnabled,
6670
// The project the user is currently in — the agent threads this into the
6771
// `project_id` arg of every `@posthog/*` tool (it's tenant-neutral and acts

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

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
import { renderHook } from "@testing-library/react";
2+
import { beforeEach, describe, expect, it, vi } from "vitest";
3+
4+
vi.mock("@tanstack/react-router", () => ({
5+
useNavigate: () => vi.fn(),
6+
}));
7+
8+
const client = {
9+
getAgentApplication: vi.fn(),
10+
listAgentRevisions: vi.fn(),
11+
getAgentRevision: vi.fn(),
12+
};
13+
14+
vi.mock("@posthog/ui/features/auth/authClient", () => ({
15+
useAuthenticatedClient: () => client,
16+
}));
17+
18+
import type { ClientToolCallData } from "@posthog/core/agent-chat/identifiers";
19+
import { useAgentBuilderStore } from "./agentBuilderStore";
20+
import { useAgentBuilderClientTools } from "./useAgentBuilderClientTools";
21+
22+
function call(
23+
tool_id: string,
24+
args: Record<string, unknown>,
25+
): ClientToolCallData {
26+
return { call_id: "call-1", tool_id, args };
27+
}
28+
29+
function handler() {
30+
return renderHook(() => useAgentBuilderClientTools()).result.current;
31+
}
32+
33+
describe("useAgentBuilderClientTools revision resolution", () => {
34+
beforeEach(() => {
35+
client.getAgentApplication.mockReset();
36+
client.listAgentRevisions.mockReset();
37+
client.getAgentRevision.mockReset();
38+
useAgentBuilderStore.setState({
39+
page: { kind: "unknown" },
40+
pendingSecret: null,
41+
pendingMcpConnect: null,
42+
});
43+
});
44+
45+
it("uses an explicit revision_id once it verifies as belonging to the agent", async () => {
46+
client.getAgentRevision.mockResolvedValue({ id: "rev-explicit" });
47+
48+
const outcome = await handler()(
49+
call("set_secret", {
50+
agent_slug: "my-agent",
51+
secret: "API_KEY",
52+
revision_id: "rev-explicit",
53+
}),
54+
);
55+
56+
expect(outcome).toEqual({ defer: true });
57+
expect(useAgentBuilderStore.getState().pendingSecret).toMatchObject({
58+
agentSlug: "my-agent",
59+
secret: "API_KEY",
60+
revisionId: "rev-explicit",
61+
});
62+
expect(client.getAgentRevision).toHaveBeenCalledWith(
63+
"my-agent",
64+
"rev-explicit",
65+
);
66+
expect(client.getAgentApplication).not.toHaveBeenCalled();
67+
expect(client.listAgentRevisions).not.toHaveBeenCalled();
68+
});
69+
70+
it("rejects an explicit revision_id that does not belong to the agent", async () => {
71+
// The nested revision route 404s (→ null) for another agent's revision.
72+
client.getAgentRevision.mockResolvedValue(null);
73+
74+
const outcome = await handler()(
75+
call("set_secret", {
76+
agent_slug: "my-agent",
77+
secret: "API_KEY",
78+
revision_id: "rev-of-other-agent",
79+
}),
80+
);
81+
82+
expect(outcome).toEqual({
83+
error: "revision_not_found: rev-of-other-agent on my-agent",
84+
});
85+
expect(useAgentBuilderStore.getState().pendingSecret).toBeNull();
86+
});
87+
88+
it("falls back to the revision open on this agent's config page", async () => {
89+
useAgentBuilderStore.setState({
90+
page: { kind: "agent-config", slug: "my-agent", revision: "rev-page" },
91+
});
92+
93+
const outcome = await handler()(
94+
call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
95+
);
96+
97+
expect(outcome).toEqual({ defer: true });
98+
expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
99+
"rev-page",
100+
);
101+
expect(client.getAgentApplication).not.toHaveBeenCalled();
102+
});
103+
104+
it("rotation targets live even while a draft config page is open", async () => {
105+
useAgentBuilderStore.setState({
106+
page: { kind: "agent-config", slug: "my-agent", revision: "rev-draft" },
107+
});
108+
client.getAgentApplication.mockResolvedValue({ live_revision: "rev-live" });
109+
client.listAgentRevisions.mockResolvedValue([
110+
{ id: "rev-draft", state: "draft" },
111+
{ id: "rev-live", state: "ready" },
112+
]);
113+
114+
await handler()(
115+
call("set_secret", {
116+
agent_slug: "my-agent",
117+
secret: "API_KEY",
118+
mode: "rotate",
119+
}),
120+
);
121+
122+
expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
123+
"rev-live",
124+
);
125+
});
126+
127+
it("ignores the page revision when it belongs to a different agent", async () => {
128+
useAgentBuilderStore.setState({
129+
page: { kind: "agent-config", slug: "other-agent", revision: "rev-page" },
130+
});
131+
client.getAgentApplication.mockResolvedValue({ live_revision: null });
132+
client.listAgentRevisions.mockResolvedValue([
133+
{ id: "rev-draft", state: "draft" },
134+
]);
135+
136+
await handler()(
137+
call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
138+
);
139+
140+
expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
141+
"rev-draft",
142+
);
143+
});
144+
145+
it.each([
146+
// A new secret targets the draft being authored (secrets only copy
147+
// forward at draft creation), even when a live revision exists.
148+
{ mode: "set", expected: "rev-draft" },
149+
// A rotation targets what's running.
150+
{ mode: "rotate", expected: "rev-live" },
151+
])(
152+
"set_secret mode=$mode resolves to $expected via the API",
153+
async ({ mode, expected }) => {
154+
client.getAgentApplication.mockResolvedValue({
155+
live_revision: "rev-live",
156+
});
157+
client.listAgentRevisions.mockResolvedValue([
158+
{ id: "rev-draft", state: "draft" },
159+
{ id: "rev-live", state: "ready" },
160+
]);
161+
162+
await handler()(
163+
call("set_secret", { agent_slug: "my-agent", secret: "API_KEY", mode }),
164+
);
165+
166+
expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
167+
expected,
168+
);
169+
},
170+
);
171+
172+
it("set_secret falls back to the newest revision when no live or draft exists", async () => {
173+
client.getAgentApplication.mockResolvedValue({ live_revision: null });
174+
client.listAgentRevisions.mockResolvedValue([
175+
{ id: "rev-newest", state: "ready" },
176+
{ id: "rev-older", state: "ready" },
177+
]);
178+
179+
await handler()(
180+
call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
181+
);
182+
183+
expect(useAgentBuilderStore.getState().pendingSecret?.revisionId).toBe(
184+
"rev-newest",
185+
);
186+
});
187+
188+
it("errors when the agent has no revisions at all", async () => {
189+
client.getAgentApplication.mockResolvedValue({ live_revision: null });
190+
client.listAgentRevisions.mockResolvedValue([]);
191+
192+
const outcome = await handler()(
193+
call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
194+
);
195+
196+
expect(outcome).toEqual({ error: "no_target_revision: my-agent" });
197+
expect(useAgentBuilderStore.getState().pendingSecret).toBeNull();
198+
});
199+
200+
it("errors when revision lookup fails", async () => {
201+
client.getAgentApplication.mockRejectedValue(new Error("network"));
202+
client.listAgentRevisions.mockRejectedValue(new Error("network"));
203+
204+
const outcome = await handler()(
205+
call("set_secret", { agent_slug: "my-agent", secret: "API_KEY" }),
206+
);
207+
208+
expect(outcome).toEqual({ error: "no_target_revision: my-agent" });
209+
});
210+
211+
it("connect_mcp prefers the newest draft (spec edits are draft-only)", async () => {
212+
client.getAgentApplication.mockResolvedValue({
213+
live_revision: "rev-live",
214+
});
215+
client.listAgentRevisions.mockResolvedValue([
216+
{ id: "rev-draft", state: "draft" },
217+
{ id: "rev-live", state: "ready" },
218+
]);
219+
220+
const outcome = await handler()(
221+
call("connect_mcp", { agent_slug: "my-agent", url: "https://mcp.test" }),
222+
);
223+
224+
expect(outcome).toEqual({ defer: true });
225+
expect(useAgentBuilderStore.getState().pendingMcpConnect).toMatchObject({
226+
agentSlug: "my-agent",
227+
revisionId: "rev-draft",
228+
url: "https://mcp.test",
229+
});
230+
});
231+
});

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

Lines changed: 79 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import { useAuthenticatedClient } from "@posthog/ui/features/auth/authClient";
12
import { useNavigate } from "@tanstack/react-router";
23
import { useCallback, useRef } from "react";
34
import type { ClientToolHandler } from "../hooks/useAgentChat";
@@ -33,6 +34,7 @@ export const AGENT_BUILDER_CLIENT_TOOLS = [
3334
*/
3435
export function useAgentBuilderClientTools(): ClientToolHandler {
3536
const navigate = useNavigate();
37+
const client = useAuthenticatedClient();
3638
const followMode = useAgentBuilderStore((s) => s.followMode);
3739
const setPendingSecret = useAgentBuilderStore((s) => s.setPendingSecret);
3840
const setPendingMcpConnect = useAgentBuilderStore(
@@ -47,25 +49,83 @@ export function useAgentBuilderClientTools(): ClientToolHandler {
4749
pageRef.current = page;
4850

4951
return useCallback(
50-
(data) => {
52+
async (data) => {
5153
const args = (data.args ?? {}) as Record<string, unknown>;
5254
const str = (v: unknown) => (typeof v === "string" ? v : undefined);
5355

56+
// Env keys and spec edits are revision-scoped, but the punch-out tool
57+
// schemas don't define `revision_id`, so the agent usually omits it.
58+
// Resolve the target: explicit arg → the revision the user is viewing on
59+
// this agent's config page → API fallback. Preference matters because
60+
// secrets only copy forward at draft creation and spec PATCHes only land
61+
// on drafts: a *new* secret or MCP connection targets the draft being
62+
// authored, while a rotation targets what's running (live).
63+
const resolveRevision = async (
64+
agentSlug: string,
65+
prefer: "live" | "draft",
66+
): Promise<string | undefined> => {
67+
const p = pageRef.current;
68+
// The viewed revision only stands in for authoring flows: a rotation
69+
// must reach what's running even while the user is viewing a draft.
70+
if (
71+
prefer === "draft" &&
72+
p.kind === "agent-config" &&
73+
p.slug === agentSlug &&
74+
p.revision
75+
) {
76+
return p.revision;
77+
}
78+
try {
79+
// A revision's `state` stays "ready" when promoted — live is the
80+
// application's `live_revision` pointer, not a revision state.
81+
const [app, revisions] = await Promise.all([
82+
client.getAgentApplication(agentSlug),
83+
client.listAgentRevisions(agentSlug),
84+
]);
85+
const live = app?.live_revision ?? undefined;
86+
const draft = revisions.find((r) => r.state === "draft")?.id;
87+
const newest = revisions[0]?.id;
88+
return prefer === "draft"
89+
? (draft ?? live ?? newest)
90+
: (live ?? draft ?? newest);
91+
} catch {
92+
return undefined;
93+
}
94+
};
95+
96+
// An explicit `revision_id` must belong to `agent_slug` before we park
97+
// the punch-out. The nested env/spec routes reject mismatches
98+
// server-side anyway, but that failure would only surface at submit —
99+
// after the user has already typed a secret into a doomed form.
100+
const verifyExplicitRevision = async (
101+
agentSlug: string,
102+
revisionId: string,
103+
): Promise<boolean> => {
104+
const revision = await client
105+
.getAgentRevision(agentSlug, revisionId)
106+
.catch(() => null);
107+
return revision != null;
108+
};
109+
54110
// set_secret — interactive punch-out. Park the call (defer) and render a
55-
// form; the dock PUTs the key and wakes the session on submit. Env keys
56-
// are revision-scoped, so resolve the target revision from the tool args,
57-
// falling back to the revision the user is currently viewing in the
58-
// agent-config page.
111+
// form; the dock PUTs the key and wakes the session on submit.
59112
if (data.tool_id === "set_secret") {
60113
const agentSlug = str(args.agent_slug);
61114
const secret = str(args.secret);
62115
if (!agentSlug) return { error: "missing_arg: agent_slug" };
63116
if (!secret) return { error: "missing_arg: secret" };
64-
const p = pageRef.current;
65-
const pageRevision = p.kind === "agent-config" ? p.revision : undefined;
66-
const revisionId = str(args.revision_id) ?? pageRevision;
67-
if (!revisionId) return { error: "missing_arg: revision_id" };
68117
const mode = args.mode === "rotate" ? "rotate" : "set";
118+
const explicit = str(args.revision_id);
119+
if (explicit && !(await verifyExplicitRevision(agentSlug, explicit))) {
120+
return { error: `revision_not_found: ${explicit} on ${agentSlug}` };
121+
}
122+
const revisionId =
123+
explicit ??
124+
(await resolveRevision(
125+
agentSlug,
126+
mode === "rotate" ? "live" : "draft",
127+
));
128+
if (!revisionId) return { error: `no_target_revision: ${agentSlug}` };
69129
setPendingSecret({
70130
callId: data.call_id,
71131
agentSlug,
@@ -80,15 +140,18 @@ export function useAgentBuilderClientTools(): ClientToolHandler {
80140
// connect_mcp — interactive punch-out. Park the call and render a prefilled
81141
// connect form; the dock runs the native OAuth/api-key connect (auth never
82142
// touches the agent), writes the resulting mcps[].connection onto the
83-
// target agent's spec, and wakes the session. Like set_secret, the target
84-
// revision comes from the args or the current agent-config page.
143+
// target agent's spec, and wakes the session. Same revision resolution as
144+
// set_secret.
85145
if (data.tool_id === "connect_mcp") {
86146
const agentSlug = str(args.agent_slug);
87147
if (!agentSlug) return { error: "missing_arg: agent_slug" };
88-
const p = pageRef.current;
89-
const pageRevision = p.kind === "agent-config" ? p.revision : undefined;
90-
const revisionId = str(args.revision_id) ?? pageRevision;
91-
if (!revisionId) return { error: "missing_arg: revision_id" };
148+
const explicit = str(args.revision_id);
149+
if (explicit && !(await verifyExplicitRevision(agentSlug, explicit))) {
150+
return { error: `revision_not_found: ${explicit} on ${agentSlug}` };
151+
}
152+
const revisionId =
153+
explicit ?? (await resolveRevision(agentSlug, "draft"));
154+
if (!revisionId) return { error: `no_target_revision: ${agentSlug}` };
92155
setPendingMcpConnect({
93156
callId: data.call_id,
94157
agentSlug,
@@ -194,6 +257,6 @@ export function useAgentBuilderClientTools(): ClientToolHandler {
194257
return { result: { focused: false, reason: "unknown_focus_target" } };
195258
}
196259
},
197-
[navigate, setPendingSecret, setPendingMcpConnect],
260+
[navigate, client, setPendingSecret, setPendingMcpConnect],
198261
);
199262
}

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

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -454,7 +454,12 @@ export function AgentConfigurationPane({
454454
) : null;
455455

456456
return (
457-
<AgentDetailLayout idOrSlug={idOrSlug} activeTab="configuration" fill>
457+
<AgentDetailLayout
458+
idOrSlug={idOrSlug}
459+
activeTab="configuration"
460+
fill
461+
configRevision={revisionId}
462+
>
458463
{!revisionId ? (
459464
<div className="p-6">
460465
<AgentDetailEmptyState

0 commit comments

Comments
 (0)