History

End a request's sandbox fibers before its database connection closes (#1799)

executeWithPause forks the sandbox as a daemon so a pause can outlive the caller that observed it. That fiber closes over the executor, and the FumaDB client captures its handle at construction rather than resolving it per operation, so the fiber keeps a live reference to the connection the host opened for the scope that built the engine. On the HTTP executor plane that scope is the request, and its finalizer closes the pool as soon as the response is written. A sandbox still parked at that moment issues its next query against a closed pool. Add ExecutionEngine.shutdown, which interrupts the engine's in-flight sandbox fibers and waits for them to unwind, and run it from the shared execution-stack middleware so it happens before the request scope tears the connection down. Nothing is lost by ending it there: on a per-request engine the pause is already unreachable once the response is written, since a resume lands on a different engine and replays the call. The MCP session Durable Object does not use this middleware and keeps its session-lifetime pauses. shutdown is a required member so a decorator that rebuilds the engine object cannot drop it silently.

Rhys Sullivan committed 2 months ago

Commit 6609fa61 parent ebdfb4415 files · +202 −0Browse files

Unified · Split

15 files
  1. apps/cloud/src/api/protected.test.ts +2 −0
  2. apps/cloud/src/engine/execution-gate.ts +3 −0
  3. apps/cloud/src/engine/execution-rate-limit.node.test.ts +2 −0
  4. apps/cloud/src/engine/execution-usage.ts +3 −0
  5. packages/core/analytics/src/engine.test.ts +2 −0
  6. packages/core/api/src/server/execution-stack-middleware.ts +22 −0
  7. packages/core/execution/src/engine.ts +53 −0
  8. packages/core/execution/src/tool-invoker.test.ts +101 −0
  9. packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts +2 −0
  10. packages/hosts/cloudflare/src/mcp/agent-session-model-resume.test.ts +2 −0
  11. packages/hosts/mcp-apps-shell/src/shell-resource.smoke.test.ts +2 −0
  12. packages/hosts/mcp-apps-shell/src/shell/mcp-app.browser.test.ts +2 −0
  13. packages/hosts/mcp/src/artifacts-tools.test.ts +2 −0
  14. packages/hosts/mcp/src/namespace-search-tools.test.ts +2 −0
  15. packages/hosts/mcp/src/tool-server.test.ts +2 −0
apps/cloud/src/api/protected.test.ts+2−0
@@ -20,6 +20,8 @@ const makeBaseEngine = (): ExecutionEngine => · 19 unmodified lines
2020 pausedExecutionCount: () => Effect.succeed(0),
2121 hasPausedExecutions: () => Effect.succeed(false),
2222 getDescription: Effect.succeed("desc"),
23+ // The fake forks nothing, so there is no sandbox fiber to end.
24+ shutdown: Effect.void,
2325 }) as ExecutionEngine;
2426
2527 describe("withExecutionUsageTracking", () => {
apps/cloud/src/engine/execution-gate.ts+3−0
@@ -108,6 +108,9 @@ export const withPreExecutionGate = <E extends Cause.YieldableError>( · 107 unmodified lines
108108 pausedExecutionCount: () => engine.pausedExecutionCount(),
109109 hasPausedExecutions: () => engine.hasPausedExecutions(),
110110 getDescription: engine.getDescription,
111+ // Forwarded, not re-implemented: the wrapped engine owns the sandbox fibers,
112+ // so the host's request-scope teardown has to reach through this decorator.
113+ shutdown: engine.shutdown,
111114 });
112115
113116 // ---------------------------------------------------------------------------
apps/cloud/src/engine/execution-rate-limit.node.test.ts+2−0
@@ -31,6 +31,8 @@ const engineStub: ExecutionEngine = { · 30 unmodified lines
3131 pausedExecutionCount: () => Effect.succeed(0),
3232 hasPausedExecutions: () => Effect.succeed(false),
3333 getDescription: Effect.succeed("stub"),
34+ // The stub forks nothing, so there is no sandbox fiber to end.
35+ shutdown: Effect.void,
3436 };
3537
3638 /** Counter that hands out a caller-controlled sequence of counts. */
apps/cloud/src/engine/execution-usage.ts+3−0
@@ -25,4 +25,7 @@ export const withExecutionUsageTracking = <E extends Cause.YieldableError>( · 24 unmodified lines
2525 pausedExecutionCount: () => engine.pausedExecutionCount(),
2626 hasPausedExecutions: () => engine.hasPausedExecutions(),
2727 getDescription: engine.getDescription,
28+ // Forwarded, not re-implemented: the wrapped engine owns the sandbox fibers,
29+ // so the host's request-scope teardown has to reach through this decorator.
30+ shutdown: engine.shutdown,
2831 });
packages/core/analytics/src/engine.test.ts+2−0
@@ -60,6 +60,8 @@ const makeFakeEngine = ( · 59 unmodified lines
6060 pausedExecutionCount: () => Effect.succeed(0),
6161 hasPausedExecutions: () => Effect.succeed(false),
6262 getDescription: Effect.succeed("fake"),
63+ // The fake forks nothing, so there is no sandbox fiber to end.
64+ shutdown: Effect.void,
6365 });
6466
6567 describe("withExecutionAnalytics", () => {
packages/core/api/src/server/execution-stack-middleware.ts+22−0
@@ -259,6 +259,25 @@ export const makeExecutionStackMiddleware = < · 258 unmodified lines
259259 Effect.provideService(ExecutorService, executor),
260260 Effect.provideService(ExecutionEngineService, engine),
261261 provideExecutorExtensions(executor),
262+ // This engine belongs to THIS request: the stack above was built
263+ // from the host's request-scoped DB handle, and `executeWithPause`
264+ // forks its sandbox as a daemon that would otherwise outlive the
265+ // handler. The host closes the connection when the request scope
266+ // closes — which happens AFTER this effect returns — so ending the
267+ // engine here is what keeps a sandbox fiber from waking up on a
268+ // closed pool.
269+ //
270+ // Nothing is lost by ending it: on a per-request engine the paused
271+ // fiber is already unreachable once the response is written (a
272+ // resume lands on a different engine and replays the call instead),
273+ // so the fiber could only ever have failed. The MCP session Durable
274+ // Object does NOT come through here — it builds its own stack over a
275+ // session-lifetime handle, so its pauses still survive between
276+ // requests, which is the whole point of that plane.
277+ //
278+ // `ensuring`, not `tap`: interruption and failure have to end the
279+ // fiber too, and it must not run before the response is produced.
280+ Effect.ensuring(engine.shutdown),
262281 );
263282 // Provide the boot-captured context; uncaptured deps (cloud's
264283 // request-scoped `DbService`) remain residual and flow through here.
@@ -315,6 +334,9 @@ const readOnlyExecutionEngine: ExecutionEngine<Cause.YieldableError> = { · 50 unmodified lines
315334 getPausedExecution: () => Effect.succeed(null),
316335 pausedExecutionCount: () => Effect.succeed(0),
317336 hasPausedExecutions: () => Effect.succeed(false),
337+ // Nothing is ever forked here — the platform branch cannot execute — so there
338+ // is no sandbox fiber to end.
339+ shutdown: Effect.void,
318340 // oxlint-disable-next-line executor/no-effect-escape-hatch -- boundary: only the MCP tool server reads this, and the MCP plane never serves a platform credential
319341 getDescription: Effect.die(new PlatformEngineUnavailable({ member: "getDescription" })),
320342 };
packages/core/execution/src/engine.ts+53−0
@@ -484,6 +484,21 @@ export type ExecutionEngine<E extends Cause.YieldableError = CodeExecutionError> · 483 unmodified lines
484484 * Get the dynamic tool description (workflow + namespaces).
485485 */
486486 readonly getDescription: Effect.Effect<string>;
487+
488+ /**
489+ * End this engine's in-flight sandbox fibers and wait for them to unwind.
490+ *
491+ * `executeWithPause` forks the sandbox as a daemon so a pause can outlive the
492+ * caller that observed it. That fiber holds the executor, and so the DB handle
493+ * belonging to whichever scope built this engine. A host that builds an engine
494+ * per request MUST call this before that scope's connection is closed, or the
495+ * fiber outlives the pool it queries.
496+ *
497+ * Required, not optional: a decorator that wrapped the engine and quietly
498+ * dropped this member would silently reopen that race, so the type makes
499+ * forwarding it a compile error.
500+ */
501+ readonly shutdown: Effect.Effect<void>;
487502 };
488503
489504 export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecutionError>(
@@ -491,6 +506,17 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 1 unmodified line
491506 ): ExecutionEngine<E> => {
492507 const { executor, codeExecutor, toolDiscoveryProvider = defaultToolDiscoveryProvider } = config;
493508 const pausedExecutions = new Map<string, InternalPausedExecution<E>>();
509+ // Every sandbox fiber `startPausableExecution` still has in flight.
510+ //
511+ // Those fibers are daemons (`Effect.forkDetach`) so a pause can outlive the
512+ // caller that observed it. But they close over `executor`, and the executor
513+ // closes over the FumaDB handle the host opened for whatever scope built THIS
514+ // engine — `makeFumaClient` captures `db` at construction, not per operation.
515+ // A host that builds one engine per HTTP request therefore needs a way to end
516+ // that fiber's life with the request; otherwise it wakes up after the
517+ // request's postgres pool has been closed and every query it makes lands on a
518+ // dead pool. `shutdown` below is that seam.
519+ const liveSandboxFibers = new Set<Fiber.Fiber<ExecuteResult, E>>();
494520 // Outcomes of executions that already settled (resumed to completion, hit a
495521 // new pause, or died while paused). MCP clients retry `resume` when a
496522 // response gets lost in transit; without this cache the retry of an
@@ -613,6 +639,7 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 116 unmodified lines
613639 fiber = yield* Effect.forkDetach(
614640 codeExecutor.execute(code, invoker).pipe(Effect.withSpan("executor.code.exec")),
615641 );
642+ liveSandboxFibers.add(fiber);
616643
617644 // When the fiber settles on its own (sandbox timeout, failure) while
618645 // pauses are still outstanding, drop them: getPausedExecution must not
@@ -624,6 +651,9 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 5 unmodified lines
624651 Fiber.await(sandboxFiber).pipe(
625652 Effect.flatMap((exit) =>
626653 Effect.sync(() => {
654+ // Settled on its own — it can no longer touch the host's DB handle,
655+ // so it is not `shutdown`'s problem any more.
656+ liveSandboxFibers.delete(sandboxFiber);
627657 const outcome = Exit.map(
628658 exit,
629659 (result): ExecutionResult => ({ status: "completed", result }),
@@ -727,10 +757,33 @@ export const createExecutionEngine = <E extends Cause.YieldableError = CodeExecu · 97 unmodified lines
727757 return result;
728758 });
729759
760+ /**
761+ * End this engine's sandbox fibers, and WAIT for them to finish unwinding.
762+ *
763+ * A host calls this when the scope that owns the engine's DB handle is about
764+ * to close. Interruption is awaited rather than fired and forgotten: the point
765+ * is that no sandbox fiber is still able to issue a query by the time the
766+ * host's connection finalizer runs, so returning early would reopen the very
767+ * race this closes.
768+ *
769+ * Paused executions are dropped with the fibers — a pause whose fiber has been
770+ * interrupted can never consume a response, so leaving the entry behind would
771+ * only let a later `resume` hand back a pause that cannot progress.
772+ */
773+ const shutdown: Effect.Effect<void> = Effect.suspend(() => {
774+ const fibers = Array.from(liveSandboxFibers);
775+ liveSandboxFibers.clear();
776+ for (const [id, paused] of pausedExecutions) {
777+ if (fibers.includes(paused.fiber)) pausedExecutions.delete(id);
778+ }
779+ return Fiber.interruptAll(fibers);
780+ });
781+
730782 return {
731783 execute: runInlineExecution,
732784 executeWithPause: startPausableExecution,
733785 resume: resumeExecution,
786+ shutdown,
734787 isExecutionSettled: (executionId) => Effect.sync(() => settledExecutionIds.has(executionId)),
735788 getPausedExecution: (executionId) =>
736789 Effect.sync(() => pausedExecutions.get(executionId) ?? null),
packages/core/execution/src/tool-invoker.test.ts+101−0
@@ -27,6 +27,7 @@ import { · 26 unmodified lines
2727 typeCheckOutputTypeScript,
2828 } from "@executor-js/sdk/testing";
2929 import { makeQuickJsExecutor } from "@executor-js/runtime-quickjs";
30+import type { CodeExecutor, ExecuteResult } from "@executor-js/codemode-core";
3031 import { createExecutionEngine } from "./engine";
3132 import { ExecutionToolError } from "./errors";
3233 import {
@@ -1875,4 +1876,104 @@ describe("pause/resume with multiple elicitations", () => { · 1842 unmodified lines
18751876 expect(resumed.result.error).toBeUndefined();
18761877 expect(resumed.result.result).toMatchObject({ ok: true });
18771878 }, 10000);
1879+
1880+ // -------------------------------------------------------------------------
1881+ // Regression: a paused execution's sandbox fiber must not outlive the
1882+ // database handle it is holding.
1883+ //
1884+ // `executeWithPause` forks the sandbox with `Effect.forkDetach` so a pause can
1885+ // outlive the caller that observed it. What that fiber closes over is the
1886+ // problem: its invoker is built from the engine's `executor`, and the FumaDB
1887+ // client captures its handle at CONSTRUCTION rather than resolving it per
1888+ // operation. So the parked fiber keeps a live reference to whichever
1889+ // connection the host opened for the scope that built this engine.
1890+ //
1891+ // On the cloud `/api/*` plane that scope is the HTTP request, and its
1892+ // finalizer closes the postgres pool as soon as the response is written. A
1893+ // sandbox fiber still parked at that moment issues its next query against a
1894+ // pool that has already had `end()` called on it — which postgres.js answers
1895+ // with CONNECTION_ENDED, and which workerd answers with "Cannot perform I/O on
1896+ // behalf of a different request" once the socket is reached from a later
1897+ // request. Those are the storage faults seen in production on reads like
1898+ // `tool.findMany` and `plugin_storage.findMany`.
1899+ //
1900+ // `engine.shutdown` is the seam the host uses to end the fiber BEFORE closing
1901+ // the connection. Nothing is lost by ending it here: on a per-request engine
1902+ // the pause is already unreachable once the response is written (a resume
1903+ // lands on a different engine and replays the call instead).
1904+ // -------------------------------------------------------------------------
1905+ it.effect(
1906+ "shutdown ends a paused execution's sandbox fiber so it cannot outlive the request",
1907+ () =>
1908+ Effect.gen(function* () {
1909+ const executor = yield* makeElicitingExecutor();
1910+
1911+ // Instrumented in place of the QuickJS runtime so the test can observe
1912+ // the one thing that actually matters: whether the forked sandbox fiber
1913+ // is still alive. It pauses through the real approval-gated tool, so
1914+ // `executeWithPause` returns on the caller's own fiber exactly as it
1915+ // does in the HTTP handler.
1916+ let sandboxEnded = false;
1917+ const probeCodeExecutor: CodeExecutor<ExecutionToolError> = {
1918+ execute: (_code, toolInvoker) =>
1919+ Effect.gen(function* () {
1920+ yield* Effect.orDie(
1921+ toolInvoker.invoke({ path: "api.org.main.singleApproval", args: {} }),
1922+ );
1923+ return { result: "unreachable", logs: [] } satisfies ExecuteResult;
1924+ }).pipe(
1925+ Effect.onInterrupt(() =>
1926+ Effect.sync(() => {
1927+ sandboxEnded = true;
1928+ }),
1929+ ),
1930+ ),
1931+ };
1932+
1933+ const engine = createExecutionEngine({ executor, codeExecutor: probeCodeExecutor });
1934+
1935+ const outcome = yield* engine.executeWithPause("ignored by the probe executor");
1936+ expect(outcome.status).toBe("paused");
1937+ const paused = outcome as Extract<typeof outcome, { status: "paused" }>;
1938+ expect(yield* engine.pausedExecutionCount()).toBe(1);
1939+
1940+ // This is the leak: the response is ready, but the sandbox is still
1941+ // parked on the approval, still holding the executor and its handle.
1942+ expect(sandboxEnded, "the sandbox is still alive while the request is served").toBe(false);
1943+
1944+ // Called from the same fiber that forked the sandbox — the shape the
1945+ // HTTP middleware runs `Effect.ensuring(engine.shutdown)` in. Bounded,
1946+ // because a shutdown that blocked would hold the request open for as
1947+ // long as the sandbox sat parked, which is worse than the leak.
1948+ const ended = yield* Effect.race(
1949+ engine.shutdown.pipe(Effect.as("ended" as const)),
1950+ Effect.sleep("5 seconds").pipe(Effect.as("hung" as const)),
1951+ );
1952+ expect(ended, "shutdown must not block the request it runs in").toBe("ended");
1953+
1954+ // The assertion the fix exists for: the fiber is actually gone by the
1955+ // time shutdown returns, so it can no longer reach the executor — and
1956+ // therefore can no longer query the connection the host closes next.
1957+ expect(sandboxEnded, "shutdown interrupted the parked sandbox fiber").toBe(true);
1958+
1959+ // The pause goes with it: one whose fiber is dead can never consume a
1960+ // response, so leaving it behind would only mislead a later resume.
1961+ expect(yield* engine.pausedExecutionCount()).toBe(0);
1962+ expect(yield* engine.getPausedExecution(paused.execution.id)).toBeNull();
1963+ }),
1964+ { timeout: 15000 },
1965+ );
1966+
1967+ it.effect("shutdown is a no-op when nothing is in flight", () =>
1968+ Effect.gen(function* () {
1969+ const executor = yield* makeElicitingExecutor();
1970+ const engine = createExecutionEngine({ executor, codeExecutor });
1971+
1972+ // Runs on every request, including the overwhelming majority that never
1973+ // pause, so it has to be inert rather than an error path.
1974+ yield* engine.shutdown;
1975+ yield* engine.shutdown;
1976+ expect(yield* engine.pausedExecutionCount()).toBe(0);
1977+ }),
1978+ );
18781979 });
packages/hosts/cloudflare/src/mcp/agent-session-durable-object.test.ts+2−0
@@ -173,6 +173,8 @@ const makeEngine = ( · 172 unmodified lines
173173 pausedExecutionCount: () => Effect.succeed(0),
174174 hasPausedExecutions: () => Effect.succeed(false),
175175 getDescription: Effect.succeed("test engine"),
176+ // The fake forks nothing, so there is no sandbox fiber to end.
177+ shutdown: Effect.void,
176178 },
177179 };
178180 };
packages/hosts/cloudflare/src/mcp/agent-session-model-resume.test.ts+2−0
@@ -263,6 +263,8 @@ const makeEngine = ( · 262 unmodified lines
263263 pausedExecutionCount: () => Effect.succeed(0),
264264 hasPausedExecutions: () => Effect.succeed(false),
265265 getDescription: Effect.succeed("test engine"),
266+ // The fake forks nothing, so there is no sandbox fiber to end.
267+ shutdown: Effect.void,
266268 };
267269 return { calls, engine, resume };
268270 };
packages/hosts/mcp-apps-shell/src/shell-resource.smoke.test.ts+2−0
@@ -28,6 +28,8 @@ const stubEngine: ExecutionEngine<never> = { · 27 unmodified lines
2828 pausedExecutionCount: () => Effect.succeed(0),
2929 hasPausedExecutions: () => Effect.succeed(false),
3030 getDescription: Effect.succeed("smoke"),
31+ // The fake forks nothing, so there is no sandbox fiber to end.
32+ shutdown: Effect.void,
3133 };
3234
3335 // oxlint-disable-next-line executor/no-double-cast -- boundary: MCP SDK ClientCapabilities predates the ext-apps `extensions` field
packages/hosts/mcp-apps-shell/src/shell/mcp-app.browser.test.ts+2−0
@@ -1311,6 +1311,8 @@ const startMcpHarness = async (openApi: OpenApiServer): Promise<McpHarness> => { · 1310 unmodified lines
13111311 const startSchemaElicitationMcpHarness = (): Promise<McpHarness> =>
13121312 startMcpHarnessForEngine({
13131313 getDescription: Effect.succeed("schema elicitation test executor"),
1314+ // The fake forks nothing, so there is no sandbox fiber to end.
1315+ shutdown: Effect.void,
13141316 execute: () => Effect.succeed({ result: null }),
13151317 executeWithPause: () =>
13161318 Effect.succeed(
packages/hosts/mcp/src/artifacts-tools.test.ts+2−0
@@ -54,6 +54,8 @@ const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: { · 53 unmodified lines
5454 pausedExecutionCount: () => Effect.succeed(0),
5555 hasPausedExecutions: () => Effect.succeed(false),
5656 getDescription: Effect.succeed("test executor"),
57+ // The fake forks nothing, so there is no sandbox fiber to end.
58+ shutdown: Effect.void,
5759 });
5860
5961 /**
packages/hosts/mcp/src/namespace-search-tools.test.ts+2−0
@@ -35,6 +35,8 @@ const makeRecordingEngine = (): { · 34 unmodified lines
3535 pausedExecutionCount: () => Effect.succeed(0),
3636 hasPausedExecutions: () => Effect.succeed(false),
3737 getDescription: Effect.succeed("test executor"),
38+ // The fake forks nothing, so there is no sandbox fiber to end.
39+ shutdown: Effect.void,
3840 },
3941 };
4042 };
packages/hosts/mcp/src/tool-server.test.ts+2−0
@@ -56,6 +56,8 @@ const makeStubEngine = <E extends Cause.YieldableError = never>(overrides: { · 55 unmodified lines
5656 pausedExecutionCount: () => Effect.succeed(0),
5757 hasPausedExecutions: () => Effect.succeed(false),
5858 getDescription: Effect.succeed(overrides.description ?? "test executor"),
59+ // The fake forks nothing, so there is no sandbox fiber to end.
60+ shutdown: Effect.void,
5961 });
6062
6163 type TestServerConfig<E extends Cause.YieldableError> = Pick<