fix(workflows): enforce sandbox execution timeout across async continuations (#285)

extensions/workflows/sandbox-child.cjs

@@ -319,15 +319,17 @@ process.on("message", (message) => {
     }
     initialized = true;
     token = message.token;
+    const seq = typeof message.seq === "number" ? message.seq : undefined;
     if (typeof message.usageJson === "string") usageJson = message.usageJson;
-    run(message.source, message.argsJson, message.maxConcurrency);
+    run(message.source, message.argsJson, message.maxConcurrency, seq);
     return;
   }
   if (message.token !== token || message.kind !== "agentResult") return;
   if (typeof message.usageJson === "string") usageJson = message.usageJson;
   const pending = pendingAgents.get(message.id);
   if (!pending) return;
   pendingAgents.delete(message.id);
+  const seq = typeof message.seq === "number" ? message.seq : undefined;
   if (typeof message.resultJson === "string")
     pending.resolve(message.resultJson);
   else
@@ -336,9 +338,14 @@ process.on("message", (message) => {
         typeof message.error === "string" ? message.error : "Agent IPC failed",
       ),
     );
+  if (seq !== undefined) {
+    setImmediate(() => {
+      send({ kind: "idle", seq });
+    });
+  }
 });
 
-function run(source, argsJson, maxConcurrency) {
+function run(source, argsJson, maxConcurrency, initSeq) {
   try {
     const sandbox = Object.create(null);
     sandbox.__argsJson = argsJson;
@@ -410,6 +417,11 @@ function run(source, argsJson, maxConcurrency) {
       context,
       { timeout: 1000 },
     );
+    if (initSeq !== undefined) {
+      setImmediate(() => {
+        send({ kind: "idle", seq: initSeq });
+      });
+    }
     Promise.resolve(context.__workflowPromise)
       .then((resultJson) => {
         if (typeof resultJson !== "string")

extensions/workflows/sandbox.ts

@@ -19,6 +19,11 @@ const MAX_LOG_MESSAGE_BYTES = 8 * 1024;
  * bypasses the controller entirely.
  */
 export const AGENT_CALL_BACKSTOP_MARGIN = 8;
+/**
+ * Maximum synchronous execution time allowed between async agent boundaries.
+ * Non-yielding code (e.g. while(true){}) after await is terminated by this timeout.
+ */
+const SANDBOX_SYNC_TIMEOUT_MS = 1_000;
 
 export interface SandboxAgentOptions {
   agent_type?: unknown;
@@ -192,6 +197,28 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
     const activeAgentRequests = new Map<number, AbortController>();
     let requestCount = 0;
     let finished = false;
+    let executionSeq = 0;
+    let executionWatchdog: NodeJS.Timeout | undefined;
+
+    const armWatchdog = (timeoutMs = SANDBOX_SYNC_TIMEOUT_MS) => {
+      if (finished) return 0;
+      if (executionWatchdog) clearTimeout(executionWatchdog);
+      const currentSeq = ++executionSeq;
+      executionWatchdog = setTimeout(() => {
+        if (finished || currentSeq !== executionSeq) return;
+        finish(new Error("Script execution timed out"));
+      }, timeoutMs);
+      executionWatchdog.unref?.();
+      return currentSeq;
+    };
+
+    const disarmWatchdog = (seq?: number) => {
+      if (seq !== undefined && seq !== executionSeq) return;
+      if (executionWatchdog) {
+        clearTimeout(executionWatchdog);
+        executionWatchdog = undefined;
+      }
+    };
 
     // The child parses this and falls back to zeros if it is ever unusable, so
     // a broken snapshot degrades `usage()` to a zero reading instead of
@@ -205,6 +232,7 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
     };
 
     const cleanup = () => {
+      disarmWatchdog();
       for (const abortController of activeAgentRequests.values()) {
         abortController.abort(new Error("Workflow stopped"));
       }
@@ -296,6 +324,12 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
         }
         return;
       }
+      if (raw.kind === "idle") {
+        if (typeof raw.seq === "number") {
+          disarmWatchdog(raw.seq);
+        }
+        return;
+      }
       if (raw.kind === "agent") {
         if (
           typeof raw.payloadJson !== "string" ||
@@ -348,10 +382,12 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
               error: "Agent result exceeded the workflow IPC output limit",
             });
           }
+          const seq = armWatchdog();
           child.send({
             token,
             kind: "agentResult",
             id,
+            seq,
             resultJson,
             usageJson: usageJson(),
           });
@@ -398,6 +434,10 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
       finish(new Error("Workflow sandbox sent an unknown IPC message"));
     });
 
+    // Arm the synchronous execution watchdog for the initial script invocation
+    // (covers non-yielding code before and after initial microtask yields).
+    const initSeq = armWatchdog();
+
     child.send(
       {
         kind: "init",
@@ -406,6 +446,7 @@ export function runWorkflowSandbox(options: RunWorkflowSandboxOptions) {
         argsJson,
         maxConcurrency: options.maxConcurrency,
         usageJson: usageJson(),
+        seq: initSeq,
       },
       (error) => {
         if (error) finish(error);

tests/extensions/workflows/sandbox.test.ts

@@ -25,6 +25,36 @@ function run(
   });
 }
 
+async function runWithTimeout(
+  source: string,
+  overrides: Partial<Parameters<typeof runWorkflowSandbox>[0]> = {},
+  timeoutMs = 4000,
+) {
+  const abort = new AbortController();
+  const timer = setTimeout(() => {
+    abort.abort(new Error(`Test deadline of ${timeoutMs}ms exceeded`));
+  }, timeoutMs);
+  timer.unref?.();
+  try {
+    return await runWorkflowSandbox({
+      source,
+      args: undefined,
+      cwd: process.cwd(),
+      signal: overrides.signal ?? abort.signal,
+      onAgent: async (prompt) => ({ ok: true, output: `reply:${prompt}` }),
+      onPhase: () => {},
+      onLog: () => {},
+      usageSnapshot: () => ({ total: 0 }),
+      maxConcurrency: 8,
+      maxAgentCalls: 128,
+      ...overrides,
+    });
+  } finally {
+    clearTimeout(timer);
+    abort.abort();
+  }
+}
+
 test("sandbox exposes only workflow capabilities and validates results", async () => {
   const phases: string[] = [];
   let active = 0;
@@ -201,7 +231,102 @@ test("sandbox source cannot escape the host accounting wrapper", async () => {
 });
 
 test("sandbox VM still rejects non-yielding synchronous code", async () => {
-  await assert.rejects(run(`while (true) {}`), /timed out/);
+  await assert.rejects(runWithTimeout(`while (true) {}`), /timed out/);
+});
+
+test("sandbox rejects non-yielding synchronous code after await Promise.resolve()", async () => {
+  await assert.rejects(
+    runWithTimeout(`await Promise.resolve(); while (true) {}`),
+    /timed out/,
+  );
+});
+
+test("sandbox rejects non-yielding synchronous code after chained microtasks", async () => {
+  await assert.rejects(
+    runWithTimeout(
+      `await Promise.resolve().then(() => Promise.resolve()); while (true) {}`,
+    ),
+    /timed out/,
+  );
+});
+
+test("sandbox rejects non-yielding synchronous code after await agent()", async () => {
+  await assert.rejects(
+    runWithTimeout(`await agent("step"); while (true) {}`, {
+      onAgent: async () => ({ ok: true, output: "done" }),
+    }),
+    /timed out/,
+  );
+});
+
+test("sandbox rejects non-yielding synchronous code between sequential agent calls", async () => {
+  let firstCallSettled = false;
+  await assert.rejects(
+    runWithTimeout(
+      `await agent("first"); while (true) {}; await agent("second");`,
+      {
+        onAgent: async () => {
+          firstCallSettled = true;
+          return { ok: true, output: "first-done" };
+        },
+      },
+    ),
+    /timed out/,
+  );
+  assert.equal(firstCallSettled, true);
+});
+
+test("sandbox rejects non-yielding synchronous code after parallel agent calls", async () => {
+  await assert.rejects(
+    runWithTimeout(
+      `await parallel([() => agent("p1"), () => agent("p2")]); while (true) {}`,
+      {
+        onAgent: async (prompt) => ({ ok: true, output: `${prompt}-done` }),
+      },
+    ),
+    /timed out/,
+  );
+});
+
+test("sandbox rejects non-yielding synchronous code in interleaved parallel branch while another agent is pending", async () => {
+  let fastCalled = false;
+  let slowCalled = false;
+
+  await assert.rejects(
+    runWithTimeout(
+      `
+        await parallel([
+          async () => {
+            await agent("fast");
+            while (true) {}
+          },
+          async () => {
+            await agent("slow");
+            return "slow-done";
+          },
+        ]);
+      `,
+      {
+        onAgent: async (prompt) => {
+          if (prompt === "fast") {
+            fastCalled = true;
+            return { ok: true, output: "fast-ok" };
+          }
+          if (prompt === "slow") {
+            slowCalled = true;
+            await new Promise((resolve) => setTimeout(resolve, 2000));
+            return { ok: true, output: "slow-ok" };
+          }
+          return { ok: true, output: "unknown" };
+        },
+      },
+      4000,
+    ),
+    /timed out/,
+  );
+
+  assert.equal(fastCalled, true);
+  assert.equal(slowCalled, true);
 });
 
 test("workflow sandbox imposes no fixed whole-agent wall timer", async () => {