History

enable MCP approval elicitation during runTaskNow

Rhys Sullivan committed 9 months ago

Commit 41cd3931 parent fe7314e3 files · +328 −4Browse files

Unified · Split

3 files
  1. executor/packages/core/src/mcp-server.ts +12 −2
  2. executor/packages/core/src/mcp/server-approval.test.ts +197 −0
  3. executor/packages/core/src/mcp/server-approval.ts +119 −2
executor/packages/core/src/mcp-server.ts+12−2
@@ -4,7 +4,11 @@ import type { AnySchema } from "@modelcontextprotocol/sdk/server/zod-compat.js"; · 3 unmodified lines
4import { z } from "zod";4import { z } from "zod";
5import type { McpExecutorService, ApprovalPrompt } from "./mcp/server-contracts";5import type { McpExecutorService, ApprovalPrompt } from "./mcp/server-contracts";
6import type { Id } from "../../database/convex/_generated/dataModel.d.ts";6import type { Id } from "../../database/convex/_generated/dataModel.d.ts";
7import { waitForTerminalTask, createMcpApprovalPrompt } from "./mcp/server-approval";7import {
8 waitForTerminalTask,
9 createMcpApprovalPrompt,
10 runTaskNowWithApprovalElicitation,
11} from "./mcp/server-approval";
8import { buildRunCodeDescription, summarizeTask } from "./mcp/server-formatting";12import { buildRunCodeDescription, summarizeTask } from "./mcp/server-formatting";
9import { getTaskTerminalState, textContent } from "./mcp/server-utils";13import { getTaskTerminalState, textContent } from "./mcp/server-utils";
10 14
@@ -88,7 +92,13 @@ function createRunCodeTool( · 77 unmodified lines
88 }92 }
89 93
90 if (service.runTaskNow) {94 if (service.runTaskNow) {
91 const runOutcome = await service.runTaskNow(created.task.id);95 const runOutcome = await runTaskNowWithApprovalElicitation(
96 service,
97 created.task.id,
98 () => service.runTaskNow!(created.task.id),
99 onApprovalPrompt,
100 { workspaceId: context.workspaceId, accountId: context.accountId },
101 );
92 const task = runOutcome?.task ?? await service.getTask(created.task.id, context.workspaceId);102 const task = runOutcome?.task ?? await service.getTask(created.task.id, context.workspaceId);
93 if (!task) {103 if (!task) {
94 return {104 return {
executor/packages/core/src/mcp/server-approval.test.ts+197−0
@@ -0,0 +1,197 @@
1import { expect, test } from "bun:test";
2import { Client } from "@modelcontextprotocol/sdk/client/index.js";
3import { InMemoryTransport } from "@modelcontextprotocol/sdk/inMemory.js";
4import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
5import { ElicitRequestSchema } from "@modelcontextprotocol/sdk/types.js";
6import type { Id } from "../../../database/convex/_generated/dataModel";
7import type { PendingApprovalRecord, TaskExecutionOutcome, TaskRecord, ToolDescriptor } from "../types";
8import {
9 createMcpApprovalPrompt,
10 runTaskNowWithApprovalElicitation,
11} from "./server-approval";
12import type { McpExecutorService } from "./server-contracts";
13
14function createTaskRecord(overrides?: Partial<TaskRecord>): TaskRecord {
15 const now = Date.now();
16 return {
17 id: "task_test",
18 code: "console.log('test')",
19 runtimeId: "local-bun",
20 status: "running",
21 timeoutMs: 15_000,
22 metadata: {},
23 workspaceId: "ws_test" as Id<"workspaces">,
24 accountId: "account_test" as Id<"accounts">,
25 createdAt: now,
26 updatedAt: now,
27 ...overrides,
28 };
29}
30
31test("runTaskNowWithApprovalElicitation resolves pending approvals while runTaskNow is waiting", async () => {
32 const workspaceId = "ws_test" as Id<"workspaces">;
33 const accountId = "account_test" as Id<"accounts">;
34 const taskId = "task_test";
35
36 let approvalStatus: "pending" | "approved" | "denied" = "pending";
37 let resolveCalls = 0;
38
39 const pendingApproval: PendingApprovalRecord = {
40 id: "approval_test",
41 taskId,
42 toolPath: "github.issues.create",
43 input: { title: "Ship elicitation", apiKey: "secret-value" },
44 status: "pending",
45 createdAt: Date.now(),
46 task: {
47 id: taskId,
48 status: "running",
49 runtimeId: "local-bun",
50 timeoutMs: 15_000,
51 createdAt: Date.now(),
52 },
53 };
54
55 const service: McpExecutorService = {
56 createTask: async () => ({ task: createTaskRecord({ status: "queued" }) }),
57 getTask: async () => createTaskRecord(),
58 subscribe: () => () => {},
59 bootstrapAnonymousContext: async () => ({
60 sessionId: "anon_session_test",
61 workspaceId,
62 clientId: "mcp",
63 accountId,
64 createdAt: Date.now(),
65 lastSeenAt: Date.now(),
66 }),
67 listTools: async (): Promise<ToolDescriptor[]> => [],
68 listPendingApprovals: async () => (approvalStatus === "pending" ? [pendingApproval] : []),
69 resolveApproval: async (input) => {
70 resolveCalls += 1;
71 approvalStatus = input.decision;
72 return { ok: true };
73 },
74 };
75
76 const runTaskNow = async (): Promise<TaskExecutionOutcome> => {
77 while (approvalStatus === "pending") {
78 await new Promise((resolve) => setTimeout(resolve, 10));
79 }
80
81 if (approvalStatus === "denied") {
82 return {
83 task: createTaskRecord({
84 status: "denied",
85 error: "Denied",
86 completedAt: Date.now(),
87 updatedAt: Date.now(),
88 }),
89 };
90 }
91
92 return {
93 task: createTaskRecord({
94 status: "completed",
95 exitCode: 0,
96 completedAt: Date.now(),
97 updatedAt: Date.now(),
98 }),
99 result: { approved: true },
100 };
101 };
102
103 const outcome = await runTaskNowWithApprovalElicitation(
104 service,
105 taskId,
106 runTaskNow,
107 async () => ({ decision: "approved" }),
108 { workspaceId, accountId },
109 );
110
111 expect(resolveCalls).toBe(1);
112 expect(outcome?.task.status).toBe("completed");
113});
114
115test("MCP client can connect and answer approval elicitation prompts", async () => {
116 const server = new McpServer(
117 { name: "approval-elicit-test-server", version: "0.0.1" },
118 { capabilities: {} },
119 );
120
121 server.registerTool(
122 "approval_prompt_demo",
123 {
124 description: "Runs approval prompt via MCP elicitation",
125 inputSchema: {},
126 },
127 async () => {
128 const prompt = createMcpApprovalPrompt(server);
129 const decision = await prompt(
130 {
131 id: "approval_demo",
132 taskId: "task_demo",
133 toolPath: "github.issues.create",
134 input: {
135 title: "Create issue",
136 apiKey: "super-secret-token",
137 nested: { authorization: "Bearer token" },
138 },
139 status: "pending",
140 createdAt: Date.now(),
141 task: {
142 id: "task_demo",
143 status: "running",
144 runtimeId: "local-bun",
145 timeoutMs: 15_000,
146 createdAt: Date.now(),
147 },
148 },
149 {
150 workspaceId: "ws_demo" as Id<"workspaces">,
151 accountId: "account_demo" as Id<"accounts">,
152 },
153 );
154
155 return {
156 content: [{ type: "text", text: JSON.stringify(decision) }],
157 };
158 },
159 );
160
161 let elicitationMessage = "";
162 const client = new Client(
163 { name: "approval-elicit-test-client", version: "0.0.1" },
164 { capabilities: { elicitation: { form: {} } } },
165 );
166
167 client.setRequestHandler(ElicitRequestSchema, async (request) => {
168 elicitationMessage = request.params.message;
169 return {
170 action: "accept",
171 content: {
172 decision: "approved",
173 reason: "Looks good",
174 },
175 };
176 });
177
178 const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
179 await Promise.all([server.connect(serverTransport), client.connect(clientTransport)]);
180
181 try {
182 const result = (await client.callTool({ name: "approval_prompt_demo", arguments: {} })) as {
183 content: Array<{ type: string; text?: string }>;
184 };
185
186 const text = result.content.find((item) => item.type === "text")?.text ?? "";
187 expect(text).toContain('"decision":"approved"');
188 expect(text).toContain('"reason":"Looks good"');
189
190 expect(elicitationMessage).toContain("[redacted]");
191 expect(elicitationMessage).not.toContain("super-secret-token");
192 expect(elicitationMessage).not.toContain("Bearer token");
193 } finally {
194 await client.close();
195 await server.close();
196 }
197});
executor/packages/core/src/mcp/server-approval.ts+119−2
@@ -1,6 +1,6 @@
1import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";1import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
2import { z } from "zod";2import { z } from "zod";
3import type { PendingApprovalRecord, TaskRecord } from "../types";3import type { PendingApprovalRecord, TaskExecutionOutcome, TaskRecord } from "../types";
4import type { Id } from "../../../database/convex/_generated/dataModel";4import type { Id } from "../../../database/convex/_generated/dataModel";
5import type {5import type {
6 ApprovalPrompt,6 ApprovalPrompt,
@@ -23,10 +23,39 @@ const subscriptionEventPayloadSchema = z.object({ · 16 unmodified lines
23 pendingApprovalCount: z.coerce.number().optional(),23 pendingApprovalCount: z.coerce.number().optional(),
24});24});
25 25
26const REDACTED_PLACEHOLDER = "[redacted]";
27const sensitiveInputKeyPattern = /(authorization|api[-_]?key|token|secret|password|cookie|credential)/i;
28
29function sanitizeApprovalInput(input: unknown, seen = new WeakSet<object>()): unknown {
30 if (Array.isArray(input)) {
31 return input.map((entry) => sanitizeApprovalInput(entry, seen));
32 }
33
34 if (input && typeof input === "object") {
35 const objectInput = input as Record<string, unknown>;
36 if (seen.has(objectInput)) {
37 return "[circular]";
38 }
39 seen.add(objectInput);
40
41 const sanitized: Record<string, unknown> = {};
42 for (const [key, value] of Object.entries(objectInput)) {
43 if (sensitiveInputKeyPattern.test(key)) {
44 sanitized[key] = REDACTED_PLACEHOLDER;
45 continue;
46 }
47 sanitized[key] = sanitizeApprovalInput(value, seen);
48 }
49 return sanitized;
50 }
51
52 return input;
53}
54
26function formatApprovalInput(input: unknown, maxLength = 2000): string {55function formatApprovalInput(input: unknown, maxLength = 2000): string {
27 let serialized: string;56 let serialized: string;
28 try {57 try {
29 serialized = JSON.stringify(input ?? {}, null, 2);58 serialized = JSON.stringify(sanitizeApprovalInput(input ?? {}), null, 2);
30 } catch {59 } catch {
31 serialized = String(input);60 serialized = String(input);
32 }61 }
@@ -236,3 +265,91 @@ export function waitForTerminalTask( · 203 unmodified lines
236 }).catch(() => {});265 }).catch(() => {});
237 });266 });
238}267}
268
269export async function runTaskNowWithApprovalElicitation(
270 service: McpExecutorService,
271 taskId: string,
272 runTaskNow: () => Promise<TaskExecutionOutcome | null>,
273 onApprovalPrompt?: ApprovalPrompt,
274 approvalContext?: ApprovalPromptContext,
275): Promise<TaskExecutionOutcome | null> {
276 const hasApprovalSupport = Boolean(
277 onApprovalPrompt
278 && approvalContext
279 && service.listPendingApprovals
280 && service.resolveApproval,
281 );
282
283 if (!hasApprovalSupport || !service.listPendingApprovals || !service.resolveApproval || !onApprovalPrompt || !approvalContext) {
284 return await runTaskNow();
285 }
286 const listPendingApprovals = service.listPendingApprovals;
287 const resolveApproval = service.resolveApproval;
288
289 let settled = false;
290 let elicitationEnabled = true;
291 let loggedElicitationFallback = false;
292 const seenApprovalIds = new Set<string>();
293
294 const logElicitationFallback = (reason: string) => {
295 if (loggedElicitationFallback) return;
296 loggedElicitationFallback = true;
297 console.warn(`[executor] MCP approval elicitation unavailable, using out-of-band approvals: ${reason}`);
298 };
299
300 const maybeHandleApprovals = async () => {
301 if (!elicitationEnabled) {
302 return;
303 }
304
305 const approvals = await listPendingApprovals(approvalContext.workspaceId);
306 const pending = approvals.filter((approval) => approval.taskId === taskId && !seenApprovalIds.has(approval.id));
307 if (pending.length === 0) {
308 return;
309 }
310
311 for (const approval of pending) {
312 let decision: ApprovalPromptDecision | null;
313 try {
314 decision = await onApprovalPrompt(approval, approvalContext);
315 } catch (error) {
316 elicitationEnabled = false;
317 logElicitationFallback(error instanceof Error ? error.message : String(error));
318 return;
319 }
320
321 if (!decision) {
322 elicitationEnabled = false;
323 logElicitationFallback("client did not provide elicitation response support");
324 return;
325 }
326
327 await resolveApproval({
328 workspaceId: approvalContext.workspaceId,
329 approvalId: approval.id,
330 decision: decision.decision,
331 reason: decision.reason,
332 reviewerId: approvalContext.accountId,
333 });
334 seenApprovalIds.add(approval.id);
335 }
336 };
337
338 const pollApprovals = async () => {
339 while (!settled) {
340 await maybeHandleApprovals().catch(() => {});
341 if (settled) {
342 break;
343 }
344 await new Promise((resolve) => setTimeout(resolve, 750));
345 }
346 };
347
348 const approvalPoller = pollApprovals();
349 try {
350 return await runTaskNow();
351 } finally {
352 settled = true;
353 void approvalPoller;
354 }
355}

Keyboard shortcuts

General

Show keyboard shortcuts
?
Focus the current search
/

Issues and pull requests

Open filters
f

Issues, pull requests and Inbox

Select next item
j↓
Select previous item
k↑
Open selected item
Enter
Dismiss filters or clear selection
Esc

Repository code

Go to file
t