fix(codemode): supervise promise lifecycles

packages/codemode/codemode.md

@@ -115,19 +115,12 @@ inside CodeMode.
 
 The runtime currently provides eager, run-once promises for tool calls and async functions; `await`;
 `then`/`catch`/`finally`; and chainable `all`/`allSettled`/`race`/`resolve`/`reject`. `Promise.all` rejects promptly while
-siblings continue, and `Promise.race` leaves losers running as JavaScript does. Tracked work remains supervised, at most
-eight tool calls run concurrently, and successful execution drains unfinished work before closing.
+siblings continue, and `Promise.race` leaves losers running as JavaScript does. Tracked work remains supervised in one
+execution scope, at most eight tool calls run concurrently, and ordinary success or failure drains unfinished work before
+closing. Timeout and external interruption cancel immediately instead.
 
 ### Confirmed defects
 
-- [ ] Keep nested fire-and-forget work alive until the execution drain. A tool call started but not returned inside an
-      async function or promise handler is currently a child of that short-lived fiber and may be interrupted when the
-      parent settles even though the promise remains globally tracked.
-- [ ] Track immediate rejected promises. `Promise.reject(value)` does not enter `pendingSettlements`, so abandoning it
-      produces no unhandled-rejection diagnostic.
-- [ ] Drain handled work after ordinary program failure, then preserve the original failure. Today pending work drains
-      only after success, so a rejecting race winner or a later program throw interrupts race losers and `Promise.all`
-      siblings. Timeout and external interruption should still cancel immediately rather than drain.
 - [ ] Make every `await` continuation asynchronous. Awaiting a plain or already-settled value currently resumes in the
       same scheduling turn and can reorder state mutation relative to JavaScript.
 - [ ] Return rejected promises for invalid `Promise.all`/`allSettled`/`race` inputs instead of throwing during the call.

packages/codemode/src/interpreter/runtime.ts

@@ -1,5 +1,5 @@
 import { parse } from "acorn"
-import { Cause, Effect, Exit, Fiber, Semaphore } from "effect"
+import { Cause, Effect, Exit, Fiber, Scope, Semaphore } from "effect"
 import { DiagnosticCategory, ModuleKind, ScriptTarget, flattenDiagnosticMessageText, transpileModule } from "typescript"
 import {
   copyIn,
@@ -607,9 +607,12 @@ class Interpreter<R> {
   private readonly toolKeys: (path: ReadonlyArray<string>) => ReadonlyArray<string>
   private readonly logs: Array<string>
   private lastValue: unknown
+  // Every promise fiber belongs to the execution rather than the async function or handler
+  // that happened to create it. The execution scope still interrupts all work on teardown.
+  private readonly promiseScope: Scope.Scope
   // Caps how many eagerly forked tool calls run at once (the parallel-call concurrency cap).
   private readonly callPermits: Semaphore.Semaphore
-  // Fiber-backed promises whose settlement no program construct has observed yet. Successful
+  // Fiber-backed promises whose settlement no program construct has observed yet. Ordinary
   // program completion drains these (like a runtime waiting on in-flight work at exit) and
   // surfaces a never-awaited failure as an unhandled-rejection diagnostic.
   private readonly pendingSettlements: Set<SandboxPromise>
@@ -617,6 +620,7 @@ class Interpreter<R> {
   constructor(
     invokeTool: (path: ReadonlyArray<string>, args: Array<unknown>) => Effect.Effect<unknown, unknown, R>,
     toolKeys: (path: ReadonlyArray<string>) => ReadonlyArray<string>,
+    promiseScope: Scope.Scope,
     logs: Array<string> = [],
     shared?: { callPermits: Semaphore.Semaphore; pendingSettlements: Set<SandboxPromise> },
   ) {
@@ -626,6 +630,7 @@ class Interpreter<R> {
     this.toolKeys = toolKeys
     this.logs = logs
     this.lastValue = undefined
+    this.promiseScope = promiseScope
     this.callPermits = shared?.callPermits ?? Semaphore.makeUnsafe(TOOL_CALL_CONCURRENCY)
     this.pendingSettlements = shared?.pendingSettlements ?? new Set<SandboxPromise>()
     globalScope.set("tools", { mutable: false, value: new ToolReference([]) })
@@ -668,7 +673,7 @@ class Interpreter<R> {
     // top-level declarations (`let undefined = 5`, `const Object = ...`) shadow builtins like
     // JS module scope, instead of colliding with the seeded globals.
     this.pushScope()
-    return Effect.gen(function* () {
+    const evaluate = Effect.gen(function* () {
       self.hoistFunctions(program.body)
       let value: unknown = undefined
       let returned = false
@@ -695,8 +700,18 @@ class Interpreter<R> {
       // resolves before crossing the data boundary - `return tools.ns.tool(...)` works
       // without an explicit await, exactly as in JS.
       if (value instanceof SandboxPromise) value = yield* self.settlePromise(value)
-      yield* self.drainPendingSettlements()
       return value
+    })
+    return Effect.gen(function* () {
+      const result = yield* Effect.exit(evaluate)
+      if (Exit.isFailure(result) && Cause.hasInterruptsOnly(result.cause)) {
+        return yield* Effect.failCause(result.cause)
+      }
+
+      const drained = yield* Effect.exit(self.drainPendingSettlements())
+      if (Exit.isFailure(result)) return yield* Effect.failCause(result.cause)
+      if (Exit.isFailure(drained)) return yield* Effect.failCause(drained.cause)
+      return result.value
     }).pipe(Effect.ensuring(Effect.sync(() => self.popScope())))
   }
 
@@ -707,19 +722,22 @@ class Interpreter<R> {
   private drainPendingSettlements(): Effect.Effect<void, unknown, never> {
     const self = this
     return Effect.gen(function* () {
+      let unhandled: InterpreterRuntimeError | undefined
       while (self.pendingSettlements.size > 0) {
         const promise = self.pendingSettlements.values().next().value
         if (promise === undefined) break
         const exit = yield* self.observePromise(promise)
         if (Exit.isSuccess(exit) || Cause.hasInterruptsOnly(exit.cause) || promise.handled) continue
+        if (unhandled !== undefined) continue
         const failure = normalizeError(Cause.squash(exit.cause))
-        throw new InterpreterRuntimeError(
+        unhandled = new InterpreterRuntimeError(
           `Unhandled rejection from an un-awaited promise: ${failure.message}`,
           undefined,
           failure.kind,
           ["Await promises so failures can be caught and handled."],
         )
       }
+      if (unhandled !== undefined) throw unhandled
     })
   }
 
@@ -735,7 +753,7 @@ class Interpreter<R> {
   }
 
   private createPromise(effect: Effect.Effect<unknown, unknown, R>): Effect.Effect<SandboxPromise, never, R> {
-    return Effect.map(Effect.forkChild(effect, { startImmediately: true }), (fiber) => {
+    return Effect.map(Effect.forkIn(effect, this.promiseScope, { startImmediately: true }), (fiber) => {
       const promise = new SandboxPromise(fiber)
       this.pendingSettlements.add(promise)
       return promise
@@ -2200,7 +2218,11 @@ class Interpreter<R> {
       )
     }
     if (ref.name === "reject") {
-      return Effect.sync(() => new SandboxPromise(undefined, Effect.fail(new ProgramThrow(args[0]))))
+      return Effect.sync(() => {
+        const promise = new SandboxPromise(undefined, Effect.fail(new ProgramThrow(args[0])))
+        this.pendingSettlements.add(promise)
+        return promise
+      })
     }
 
     const items = Array.isArray(args[0]) ? args[0] : spreadItems(args[0])
@@ -2306,7 +2328,7 @@ class Interpreter<R> {
   }
 
   private invokeFunction(fn: CodeModeFunction, args: Array<unknown>): Effect.Effect<unknown, unknown, R> {
-    const invocation = new Interpreter(this.invokeTool, this.toolKeys, this.logs, {
+    const invocation = new Interpreter(this.invokeTool, this.toolKeys, this.promiseScope, this.logs, {
       callPermits: this.callPermits,
       pendingSettlements: this.pendingSettlements,
     })
@@ -3555,18 +3577,21 @@ export const executeWithLimits = <const Tools extends Record<string, unknown>>(
     })
   }
 
-  const operation = Effect.gen(function* () {
-    const program = parseProgram(options.code)
-    const interpreter = new Interpreter<Services<Tools>>(tools.invoke, tools.keys, logs)
-    const value = yield* interpreter.run(program)
-    const result = copyOut(copyIn(value, "Execution result"), true) as DataValue
-    return {
-      ok: true,
-      value: result,
-      ...logged(),
-      toolCalls: tools.calls,
-    } satisfies Result
-  }).pipe((program) => {
+  const operation = Effect.scoped(
+    Effect.gen(function* () {
+      const scope = yield* Effect.acquireRelease(Scope.make("parallel"), (scope, exit) => Scope.close(scope, exit))
+      const program = parseProgram(options.code)
+      const interpreter = new Interpreter<Services<Tools>>(tools.invoke, tools.keys, scope, logs)
+      const value = yield* interpreter.run(program)
+      const result = copyOut(copyIn(value, "Execution result"), true) as DataValue
+      return {
+        ok: true,
+        value: result,
+        ...logged(),
+        toolCalls: tools.calls,
+      } satisfies Result
+    }),
+  ).pipe((program) => {
     const timeoutMs = limits.timeoutMs
     if (timeoutMs === undefined) return program
     return program.pipe(

packages/codemode/test/promise.test.ts

@@ -212,6 +212,22 @@ describe("first-class promise values", () => {
     expect(diagnostic.suggestions?.join(" ")).toContain("Await promises")
   })
 
+  test("an unhandled rejection does not stop later work from draining", async () => {
+    const trace = makeTrace()
+    const diagnostic = await error(
+      `
+        tools.host.fail({})
+        tools.host.sleepy({ id: 1, ms: 30 })
+        return "done"
+      `,
+      { trace },
+    )
+    expect(diagnostic.kind).toBe("ToolFailure")
+    expect(diagnostic.message).toContain("Unhandled rejection from an un-awaited promise")
+    expect(trace.completed).toBe(1)
+    expect(trace.interrupted).toBe(0)
+  })
+
   test("a never-awaited failing async function surfaces as an unhandled promise rejection", async () => {
     const diagnostic = await error(`
       const fail = async () => { throw new Error("boom") }
@@ -235,6 +251,58 @@ describe("first-class promise values", () => {
     expect(diagnostic.kind).toBe("ToolFailure")
     expect(diagnostic.message).toContain("Lookup refused")
   })
+
+  test("keeps unreturned calls alive after their async function settles", async () => {
+    const trace = makeTrace()
+    expect(
+      await value(
+        `
+          const start = async () => {
+            tools.host.sleepy({ id: 1, ms: 30 })
+            return "started"
+          }
+          await start()
+          return "done"
+        `,
+        { trace },
+      ),
+    ).toBe("done")
+    expect(trace.completed).toBe(1)
+    expect(trace.interrupted).toBe(0)
+  })
+
+  test("keeps unreturned calls alive after their promise handler settles", async () => {
+    const trace = makeTrace()
+    expect(
+      await value(
+        `
+          await Promise.resolve().then(() => {
+            tools.host.sleepy({ id: 1, ms: 30 })
+          })
+          return "done"
+        `,
+        { trace },
+      ),
+    ).toBe("done")
+    expect(trace.completed).toBe(1)
+    expect(trace.interrupted).toBe(0)
+  })
+
+  test("drains pending work after program failure and preserves the original failure", async () => {
+    const trace = makeTrace()
+    const diagnostic = await error(
+      `
+        tools.host.fail({})
+        tools.host.sleepy({ id: 1, ms: 30 })
+        throw new Error("original")
+      `,
+      { trace },
+    )
+    expect(diagnostic.kind).toBe("ExecutionFailure")
+    expect(diagnostic.message).toBe("Uncaught: original")
+    expect(trace.completed).toBe(1)
+    expect(trace.interrupted).toBe(0)
+  })
 })
 
 describe("promises at data boundaries", () => {
@@ -450,6 +518,18 @@ describe("Promise.race", () => {
     ).toBe("Lookup refused")
   })
 
+  test("a rejecting winner still drains its slow loser", async () => {
+    const trace = makeTrace()
+    const diagnostic = await error(
+      `return await Promise.race([tools.host.fail({}), tools.host.sleepy({ id: 1, ms: 30 })])`,
+      { trace },
+    )
+    expect(diagnostic.kind).toBe("ToolFailure")
+    expect(diagnostic.message).toBe("Lookup refused")
+    expect(trace.completed).toBe(1)
+    expect(trace.interrupted).toBe(0)
+  })
+
   test("a plain value wins over pending promises", async () => {
     const trace = makeTrace()
     expect(
@@ -496,6 +576,16 @@ describe("Promise.resolve / Promise.reject", () => {
     `),
     ).toBe("nope")
   })
+
+  test("an abandoned rejected promise surfaces as an unhandled rejection", async () => {
+    const diagnostic = await error(`
+      Promise.reject(new Error("boom"))
+      return "done"
+    `)
+    expect(diagnostic.kind).toBe("ExecutionFailure")
+    expect(diagnostic.message).toContain("Unhandled rejection from an un-awaited promise")
+    expect(diagnostic.message).toContain("boom")
+  })
 })
 
 describe("Promise combinator values", () => {
@@ -549,6 +639,41 @@ describe("timeout interruption of forked calls", () => {
     expect(result.error.kind).toBe("TimeoutExceeded")
     expect(trace.interrupted).toBe(2)
   })
+
+  test("interrupts promise fibers concurrently during scope teardown", async () => {
+    const cleanup = { active: 0, overlapped: false }
+    const tool = Tool.make({
+      description: "Wait until interrupted",
+      input: Schema.Struct({}),
+      output: Schema.Never,
+      run: () =>
+        Effect.never.pipe(
+          Effect.onInterrupt(() =>
+            Effect.gen(function* () {
+              cleanup.active += 1
+              yield* Effect.sleep(20)
+              cleanup.overlapped ||= cleanup.active > 1
+              cleanup.active -= 1
+            }),
+          ),
+        ),
+    })
+    const result = await Effect.runPromise(
+      CodeMode.execute({
+        tools: { host: { first: tool, second: tool } },
+        code: `
+          tools.host.first({})
+          tools.host.second({})
+          return await Promise.resolve("waiting")
+        `,
+        limits: { timeoutMs: 50 },
+      }),
+    )
+    expect(result.ok).toBe(false)
+    if (result.ok) return
+    expect(result.error.kind).toBe("TimeoutExceeded")
+    expect(cleanup.overlapped).toBe(true)
+  })
 })
 
 describe("promise chaining", () => {