Fix MCP elicitation detection and multi-resume hang

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<ExecuteResult>;
+  /** Deferred the caller completes with the user's response to resume the fiber. */
+  readonly response: Deferred.Deferred<typeof ElicitationResponse.Type>;
+  /** The fiber running the sandboxed code — stays alive across pause/resume cycles. */
+  readonly fiber: Fiber.Fiber<ExecuteResult, unknown>;
+  /** Ref to the current pause signal — swapped by resume() before unblocking. */
+  readonly pauseSignalRef: Ref.Ref<Deferred.Deferred<PausedExecution>>;
 };
 
 export type ResumeResponse = {
@@ -267,9 +271,10 @@ export type ExecutionEngine = {
   readonly executeWithPause: (code: string) => Promise<ExecutionResult>;
 
   /**
-   * 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<ExecuteResult | null>;
+  readonly resume: (executionId: string, response: ResumeResponse) => Promise<ExecutionResult | null>;
 
   /**
    * Get the dynamic tool description (workflow + namespaces).
@@ -286,61 +291,109 @@ export const createExecutionEngine = (config: ExecutionEngineConfig): ExecutionE
   const pausedExecutions = new Map<string, PausedExecution>();
   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<PausedExecution>((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<ExecuteResult, unknown>,
+    pauseSignal: Deferred.Deferred<PausedExecution>,
+  ): Effect.Effect<ExecutionResult> =>
+    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<typeof ElicitationResponse.Type>((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<ExecutionResult> =>
+    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<PausedExecution>(),
+      );
+
+      // Will be set once the fiber is forked.
+      let fiber: Fiber.Fiber<ExecuteResult, unknown>;
+
+      const elicitationHandler: ElicitationHandler = (ctx) =>
+        Effect.gen(function* () {
+          const responseDeferred = yield* Deferred.make<typeof ElicitationResponse.Type>();
           const id = `exec_${++nextId}`;
+
           const paused: PausedExecution = {
             id,
             elicitationContext: ctx,
-            resolve: (response) => resume(Effect.succeed(response)),
-            completion: undefined as unknown as Promise<ExecuteResult>,
+            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<ExecuteResult> }).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<ExecutionResult | null> =>
+    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<PausedExecution>();
+      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)),
   };
 };

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<never>((_, reject) =>
+              setTimeout(
+                () =>
+                  reject(
+                    new Error(
+                      "resume hung — second elicitation not surfaced",
+                    ),
+                  ),
+                5000,
+              ),
+            ),
+          ]),
+        );
+
+        expect(outcome2).not.toBeNull();
+      }),
+    { timeout: 10000 },
+  );
+});