diff --git a/packages/core/execution/src/engine.ts b/packages/core/execution/src/engine.ts index 0e1196813c..412b598e08 100644 --- a/packages/core/execution/src/engine.ts +++ b/packages/core/execution/src/engine.ts @@ -1,4 +1,4 @@ -import { Effect } from "effect"; +import { Deferred, Effect, Fiber, Ref } from "effect"; import type { Executor, @@ -35,8 +35,12 @@ export type ExecutionResult = export type PausedExecution = { readonly id: string; readonly elicitationContext: ElicitationContext; - readonly resolve: (response: typeof ElicitationResponse.Type) => void; - readonly completion: Promise; + /** Deferred the caller completes with the user's response to resume the fiber. */ + readonly response: Deferred.Deferred; + /** The fiber running the sandboxed code — stays alive across pause/resume cycles. */ + readonly fiber: Fiber.Fiber; + /** Ref to the current pause signal — swapped by resume() before unblocking. */ + readonly pauseSignalRef: Ref.Ref>; }; export type ResumeResponse = { @@ -267,9 +271,10 @@ export type ExecutionEngine = { readonly executeWithPause: (code: string) => Promise; /** - * Resume a paused execution. + * Resume a paused execution. Returns a completed result, a new pause, or + * null if the executionId was not found. */ - readonly resume: (executionId: string, response: ResumeResponse) => Promise; + readonly resume: (executionId: string, response: ResumeResponse) => Promise; /** * Get the dynamic tool description (workflow + namespaces). @@ -286,61 +291,109 @@ export const createExecutionEngine = (config: ExecutionEngineConfig): ExecutionE const pausedExecutions = new Map(); let nextId = 0; - return { - execute: async (code, options) => { - const invoker = makeFullInvoker(executor, { - onElicitation: options.onElicitation, - }); - return runEffect(codeExecutor.execute(code, invoker)); - }, - - executeWithPause: async (code) => { - // Signal from the elicitation handler to the race below. - let signalPause: ((paused: PausedExecution) => void) | null = null; - const pausePromise = new Promise((resolve) => { - signalPause = resolve; - }); + /** + * Race a running fiber against a pause signal. Returns when either + * the fiber completes or an elicitation handler fires (whichever + * comes first). Re-used by both executeWithPause and resume. + */ + const awaitCompletionOrPause = ( + fiber: Fiber.Fiber, + pauseSignal: Deferred.Deferred, + ): Effect.Effect => + Effect.race( + Fiber.join(fiber).pipe( + Effect.orDie, + Effect.map((result): ExecutionResult => ({ status: "completed", result })), + ), + Deferred.await(pauseSignal).pipe( + Effect.map((paused): ExecutionResult => ({ status: "paused", execution: paused })), + ), + ); - const elicitationHandler: ElicitationHandler = (ctx: ElicitationContext) => - Effect.async((resume) => { + /** + * Start an execution in the pause/resume mode. Forks the sandbox + * onto its own fiber and waits for either completion or the first + * elicitation pause. + */ + const startPausableExecution = (code: string): Effect.Effect => + Effect.gen(function* () { + // Ref holds the current pause signal. The elicitation handler reads + // it each time it fires, so resume() can swap in a fresh Deferred + // before unblocking the fiber. + const pauseSignalRef = yield* Ref.make( + yield* Deferred.make(), + ); + + // Will be set once the fiber is forked. + let fiber: Fiber.Fiber; + + const elicitationHandler: ElicitationHandler = (ctx) => + Effect.gen(function* () { + const responseDeferred = yield* Deferred.make(); const id = `exec_${++nextId}`; + const paused: PausedExecution = { id, elicitationContext: ctx, - resolve: (response) => resume(Effect.succeed(response)), - completion: undefined as unknown as Promise, + response: responseDeferred, + fiber: fiber!, + pauseSignalRef, }; pausedExecutions.set(id, paused); - signalPause!(paused); - }); - const invoker = makeFullInvoker(executor, { onElicitation: elicitationHandler }); - const completionPromise = runEffect(codeExecutor.execute(code, invoker)); + const currentSignal = yield* Ref.get(pauseSignalRef); + yield* Deferred.succeed(currentSignal, paused); - // Race: either the execution completes, or it pauses for elicitation. - const result = await Promise.race([ - completionPromise.then((r) => ({ kind: "completed" as const, result: r })), - pausePromise.then((p) => ({ kind: "paused" as const, execution: p })), - ]); + // Suspend until resume() completes responseDeferred. + return yield* Deferred.await(responseDeferred); + }); - if (result.kind === "completed") { - return { status: "completed", result: result.result }; - } + const invoker = makeFullInvoker(executor, { onElicitation: elicitationHandler }); + fiber = yield* Effect.fork(codeExecutor.execute(code, invoker)); - // Execution paused — attach the completion promise and return - (result.execution as { completion: Promise }).completion = completionPromise; - return { status: "paused", execution: result.execution }; - }, + const initialSignal = yield* Ref.get(pauseSignalRef); + return yield* awaitCompletionOrPause(fiber, initialSignal); + }); - resume: async (executionId, response) => { + /** + * Resume a paused execution. Swaps in a fresh pause signal, completes + * the response Deferred to unblock the fiber, then races completion + * against the next pause. + */ + const resumeExecution = ( + executionId: string, + response: ResumeResponse, + ): Effect.Effect => + Effect.gen(function* () { const paused = pausedExecutions.get(executionId); if (!paused) return null; - pausedExecutions.delete(executionId); - paused.resolve({ action: response.action, content: response.content }); - return paused.completion; + + // Swap in a fresh pause signal BEFORE unblocking the fiber, so the + // next elicitation handler call signals this new Deferred. + const nextSignal = yield* Deferred.make(); + yield* Ref.set(paused.pauseSignalRef, nextSignal); + + yield* Deferred.succeed(paused.response, { + action: response.action, + content: response.content, + }); + + return yield* awaitCompletionOrPause(paused.fiber, nextSignal); + }); + + return { + execute: async (code, options) => { + const invoker = makeFullInvoker(executor, { + onElicitation: options.onElicitation, + }); + return runEffect(codeExecutor.execute(code, invoker)); }, + executeWithPause: (code) => runEffect(startPausableExecution(code)), + + resume: (executionId, response) => runEffect(resumeExecution(executionId, response)), + getDescription: () => runEffect(buildExecuteDescription(executor)), }; }; diff --git a/packages/core/execution/src/tool-invoker.test.ts b/packages/core/execution/src/tool-invoker.test.ts index 52aa41665c..d4fa506a51 100644 --- a/packages/core/execution/src/tool-invoker.test.ts +++ b/packages/core/execution/src/tool-invoker.test.ts @@ -8,6 +8,7 @@ import { inMemoryToolsPlugin, makeTestConfig, tool, + type ToolId, } from "@executor/sdk"; import { createExecutionEngine } from "./engine"; import { describeTool, searchTools } from "./tool-invoker"; @@ -21,6 +22,9 @@ const ContactInput = Schema.Struct({ email: Schema.String, }); +import type { ExecutionResult } from "./engine"; +import { FormElicitation } from "@executor/sdk"; + const acceptAll = () => Effect.succeed(new ElicitationResponse({ action: "accept" })); const makeSearchExecutor = () => @@ -273,3 +277,100 @@ describe("tool discovery", () => { }), ); }); + +// --------------------------------------------------------------------------- +// pause/resume — multiple elicitations in a single execution +// --------------------------------------------------------------------------- + +describe("pause/resume with multiple elicitations", () => { + const makeElicitingExecutor = () => + Effect.gen(function* () { + const config = makeTestConfig({ + plugins: [ + inMemoryToolsPlugin({ + namespace: "api", + tools: [ + tool({ + name: "multiApproval", + description: "A tool that elicits twice", + inputSchema: EmptyInput, + handler: (_args, ctx) => + Effect.gen(function* () { + const r1 = yield* ctx.elicit( + new FormElicitation({ + message: "First approval", + requestedSchema: {}, + }), + ); + const r2 = yield* ctx.elicit( + new FormElicitation({ + message: "Second approval", + requestedSchema: {}, + }), + ); + return { first: r1, second: r2 }; + }), + }), + ], + }), + ] as const, + }); + + yield* config.sources.registerRuntime( + new Source({ + id: "api", + name: "API", + kind: "in-memory", + runtime: true, + canRemove: false, + canRefresh: false, + }), + ); + + return yield* createExecutor(config); + }); + + it.effect( + "resume does not hang when execution hits a second elicitation", + () => + Effect.gen(function* () { + const executor = yield* makeElicitingExecutor(); + const engine = createExecutionEngine({ executor }); + + const code = 'return await tools.api.multiApproval({});'; + + // First executeWithPause — should pause on first elicitation + const outcome1 = yield* Effect.promise(() => + engine.executeWithPause(code), + ); + expect(outcome1.status).toBe("paused"); + if (outcome1.status !== "paused") throw new Error("expected pause"); + expect(outcome1.execution.elicitationContext.request.message).toBe( + "First approval", + ); + + // Resume first pause — execution continues to second elicitation. + // resume() must not hang; it should return (either a new paused + // result or the completion). + const outcome2 = yield* Effect.promise(() => + Promise.race([ + engine.resume(outcome1.execution.id, { action: "accept" }), + new Promise((_, reject) => + setTimeout( + () => + reject( + new Error( + "resume hung — second elicitation not surfaced", + ), + ), + 5000, + ), + ), + ]), + ); + + expect(outcome2).not.toBeNull(); + }), + { timeout: 10000 }, + ); +}); diff --git a/packages/hosts/mcp/src/server.test.ts b/packages/hosts/mcp/src/server.test.ts index 4a5703d91a..59cd80b61f 100644 --- a/packages/hosts/mcp/src/server.test.ts +++ b/packages/hosts/mcp/src/server.test.ts @@ -313,11 +313,11 @@ describe("MCP host server — client with elicitation", () => { }); // --------------------------------------------------------------------------- -// Tests — client with form-only elicitation (falls back to pause/resume) +// Tests — client with form-only elicitation (uses managed elicitation) // --------------------------------------------------------------------------- describe("MCP host server — client with form-only elicitation", () => { - it("resume tool is visible when client only supports form elicitation", async () => { + it("resume tool is hidden when client supports form elicitation", async () => { const engine = makeStubEngine({}); const { client, close } = await connect(engine, { elicitation: { form: {} }, @@ -327,17 +327,16 @@ describe("MCP host server — client with form-only elicitation", () => { const { tools } = await client.listTools(); const names = tools.map((t) => t.name); expect(names).toContain("execute"); - expect(names).toContain("resume"); + expect(names).not.toContain("resume"); } finally { await close(); } }); - it("uses pause/resume path when client only supports form", async () => { + it("uses managed elicitation path when client supports form", async () => { const engine = makeStubEngine({ - executeWithPause: async () => ({ - status: "completed", - result: { result: "form-only-path" }, + execute: async (code) => ({ + result: `managed: ${code}`, }), }); @@ -350,7 +349,49 @@ describe("MCP host server — client with form-only elicitation", () => { name: "execute", arguments: { code: "test" }, }); - expect(result.content).toEqual([{ type: "text", text: "form-only-path" }]); + expect(result.content).toEqual([{ type: "text", text: "managed: test" }]); + } finally { + await close(); + } + }); + + it("UrlElicitation falls back to form when client lacks url support", async () => { + let receivedMessage: string | undefined; + + const engine = makeStubEngine({ + execute: async (_code, { onElicitation }) => { + const response = await Effect.runPromise( + onElicitation({ + toolId: "t" as any, + args: {}, + request: new UrlElicitation({ + message: "Please authenticate", + url: "https://auth.example.com/oauth", + elicitationId: "elic-1", + }), + }), + ); + return { result: response.action }; + }, + }); + + const { client, close } = await connect(engine, { + elicitation: { form: {} }, // no url support + }); + + client.setRequestHandler(ElicitRequestSchema, async (request) => { + receivedMessage = (request.params as Record).message as string; + return { action: "accept" as const, content: {} }; + }); + + try { + const result = await client.callTool({ + name: "execute", + arguments: { code: "oauth" }, + }); + expect(result.content).toEqual([{ type: "text", text: "accept" }]); + expect(receivedMessage).toContain("https://auth.example.com/oauth"); + expect(receivedMessage).toContain("Please authenticate"); } finally { await close(); } @@ -446,7 +487,7 @@ describe("MCP host server — client without elicitation (pause/resume)", () => const engine = makeStubEngine({ resume: async (executionId, response) => { if (executionId === "exec_1" && response.action === "accept") { - return { result: "resumed-ok" }; + return { status: "completed", result: { result: "resumed-ok" } }; } return null; }, @@ -478,7 +519,7 @@ describe("MCP host server — client without elicitation (pause/resume)", () => const engine = makeStubEngine({ resume: async (_id, response) => { receivedContent = response.content; - return { result: "ok" }; + return { status: "completed", result: { result: "ok" } }; }, }); @@ -505,7 +546,7 @@ describe("MCP host server — client without elicitation (pause/resume)", () => const engine = makeStubEngine({ resume: async (_id, response) => { receivedContent = response.content; - return { result: "ok" }; + return { status: "completed", result: { result: "ok" } }; }, }); @@ -652,7 +693,7 @@ describe("MCP host server — resume content parsing", () => { const engine = makeStubEngine({ resume: async (_id, response) => { receivedContent = response.content; - return { result: "ok" }; + return { status: "completed", result: { result: "ok" } }; }, }); @@ -680,7 +721,7 @@ describe("MCP host server — resume content parsing", () => { const engine = makeStubEngine({ resume: async (_id, response) => { receivedContent = response.content; - return { result: "ok" }; + return { status: "completed", result: { result: "ok" } }; }, }); diff --git a/packages/hosts/mcp/src/server.ts b/packages/hosts/mcp/src/server.ts index 80615726b9..f6f4bd214f 100644 --- a/packages/hosts/mcp/src/server.ts +++ b/packages/hosts/mcp/src/server.ts @@ -28,60 +28,84 @@ export type ExecutorMcpServerConfig = // Elicitation bridge // --------------------------------------------------------------------------- -const supportsManagedElicitation = (server: McpServer): boolean => { +const getElicitationSupport = ( + server: McpServer, +): { form: boolean; url: boolean } => { const capabilities = server.server.getClientCapabilities(); - if (capabilities === undefined || !capabilities.elicitation) return false; + if (capabilities === undefined || !capabilities.elicitation) + return { form: false, url: false }; const elicitation = capabilities.elicitation as Record; - return Boolean(elicitation.form) && Boolean(elicitation.url); + return { form: Boolean(elicitation.form), url: Boolean(elicitation.url) }; }; +const supportsManagedElicitation = (server: McpServer): boolean => + getElicitationSupport(server).form; + type ElicitInputParams = - | { mode?: "form"; message: string; requestedSchema: { readonly [key: string]: unknown } } + | { + mode?: "form"; + message: string; + requestedSchema: { readonly [key: string]: unknown }; + } | { mode: "url"; message: string; url: string; elicitationId: string }; -const elicitationRequestToParams: (request: ElicitationRequest) => ElicitInputParams = - Match.type().pipe( - Match.tag("UrlElicitation", (req) => ({ - mode: "url" as const, - message: req.message, - url: req.url, - elicitationId: req.elicitationId, - })), - Match.tag("FormElicitation", (req) => ({ - message: req.message, - // The MCP SDK validates requestedSchema as a JSON Schema with - // `type: "object"` and `properties`. For approval-only elicitations - // where no fields are needed, provide a minimal valid schema. - requestedSchema: - Object.keys(req.requestedSchema).length === 0 - ? { type: "object" as const, properties: {} } - : req.requestedSchema, - })), - Match.exhaustive, - ); - -const makeMcpElicitationHandler = (server: McpServer): ElicitationHandler => +const elicitationRequestToParams: ( + request: ElicitationRequest, +) => ElicitInputParams = Match.type().pipe( + Match.tag("UrlElicitation", (req) => ({ + mode: "url" as const, + message: req.message, + url: req.url, + elicitationId: req.elicitationId, + })), + Match.tag("FormElicitation", (req) => ({ + message: req.message, + // The MCP SDK validates requestedSchema as a JSON Schema with + // `type: "object"` and `properties`. For approval-only elicitations + // where no fields are needed, provide a minimal valid schema. + requestedSchema: + Object.keys(req.requestedSchema).length === 0 + ? { type: "object" as const, properties: {} } + : req.requestedSchema, + })), + Match.exhaustive, +); + +const makeMcpElicitationHandler = + (server: McpServer): ElicitationHandler => (ctx: ElicitationContext): Effect.Effect => { - const params = elicitationRequestToParams(ctx.request); - - return Effect.promise(async (): Promise => { - try { - const response = await server.server.elicitInput( - params as Parameters[0], - ); - - return { - action: response.action, - content: response.content, - }; - } catch (err) { - console.error( - "[executor] elicitInput failed — falling back to cancel.", - err instanceof Error ? err.message : err, - ); - return { action: "cancel" }; - } - }); + const { url: supportsUrl } = getElicitationSupport(server); + + // If client doesn't support url mode, fall back to a form asking the user + // to visit the URL manually and confirm when done. + const params = + ctx.request._tag === "UrlElicitation" && !supportsUrl + ? { + message: `${ctx.request.message}\n\nPlease visit this URL:\n${ctx.request.url}\n\nClick accept once you have completed the flow.`, + requestedSchema: { type: "object" as const, properties: {} }, + } + : elicitationRequestToParams(ctx.request); + + return Effect.promise( + async (): Promise => { + try { + const response = await server.server.elicitInput( + params as Parameters[0], + ); + + return { + action: response.action, + content: response.content, + }; + } catch (err) { + console.error( + "[executor] elicitInput failed — falling back to cancel.", + err instanceof Error ? err.message : err, + ); + return { action: "cancel" }; + } + }, + ); }; // --------------------------------------------------------------------------- @@ -116,7 +140,8 @@ const toMcpPausedResult = ( export const createExecutorMcpServer = async ( config: ExecutorMcpServerConfig, ): Promise => { - const engine = "engine" in config ? config.engine : createExecutionEngine(config); + const engine = + "engine" in config ? config.engine : createExecutionEngine(config); const description = await engine.getDescription(); const server = new McpServer( @@ -138,7 +163,9 @@ export const createExecutorMcpServer = async ( : toMcpPausedResult(formatPausedExecution(outcome.execution)); }; - const parseJsonContent = (raw: string): Record | undefined => { + const parseJsonContent = ( + raw: string, + ): Record | undefined => { if (raw === "{}") return undefined; let parsed: unknown; try { @@ -146,7 +173,9 @@ export const createExecutorMcpServer = async ( } catch { return undefined; } - return typeof parsed === "object" && parsed !== null && !Array.isArray(parsed) + return typeof parsed === "object" && + parsed !== null && + !Array.isArray(parsed) ? (parsed as Record) : undefined; }; @@ -170,23 +199,36 @@ export const createExecutorMcpServer = async ( "Never call this without user approval unless they explicitly state otherwise.", ].join("\n"), inputSchema: { - executionId: z.string().describe("The execution ID from the paused result"), - action: z.enum(["accept", "decline", "cancel"]).describe("How to respond to the interaction"), - content: z.string().describe("Optional JSON-encoded response content for form elicitations").default("{}"), + executionId: z + .string() + .describe("The execution ID from the paused result"), + action: z + .enum(["accept", "decline", "cancel"]) + .describe("How to respond to the interaction"), + content: z + .string() + .describe( + "Optional JSON-encoded response content for form elicitations", + ) + .default("{}"), }, }, async ({ executionId, action, content: rawContent }) => { const content = parseJsonContent(rawContent); - const result = await engine.resume(executionId, { action, content }); + const outcome = await engine.resume(executionId, { action, content }); - if (!result) { + if (!outcome) { return { - content: [{ type: "text", text: `No paused execution: ${executionId}` }], + content: [ + { type: "text", text: `No paused execution: ${executionId}` }, + ], isError: true, }; } - return toMcpResult(formatExecuteResult(result)); + return outcome.status === "completed" + ? toMcpResult(formatExecuteResult(outcome.result)) + : toMcpPausedResult(formatPausedExecution(outcome.execution)); }, );