fix(codemode): preserve async callback semantics (#35603)

packages/codemode/src/interpreter/model.ts

@@ -45,6 +45,7 @@ export class CodeModeFunction {
     readonly parameters: ReadonlyArray<AstNode>,
     readonly body: AstNode,
     readonly capturedScopes: ReadonlyArray<Map<string, Binding>>,
+    readonly async: boolean,
   ) {}
 }
 
@@ -153,7 +154,8 @@ export const unsupportedSyntax = (kind: string, node: AstNode): InterpreterRunti
     [supportedSyntaxMessage],
   )
 
-export const isRecord = (value: unknown): value is Record<string, unknown> => typeof value === "object" && value !== null
+export const isRecord = (value: unknown): value is Record<string, unknown> =>
+  typeof value === "object" && value !== null
 
 export const asNode = (value: unknown, context: string): AstNode => {
   if (!isRecord(value) || typeof value.type !== "string") {

packages/codemode/src/interpreter/runtime.ts

@@ -612,12 +612,13 @@ class Interpreter<R> {
   // Fiber-backed promises whose settlement no program construct has observed yet. Successful
   // 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 = new Set<SandboxPromise>()
+  private readonly pendingSettlements: Set<SandboxPromise>
 
   constructor(
     invokeTool: (path: ReadonlyArray<string>, args: Array<unknown>) => Effect.Effect<unknown, unknown, R>,
     toolKeys: (path: ReadonlyArray<string>) => ReadonlyArray<string>,
     logs: Array<string> = [],
+    shared?: { callPermits: Semaphore.Semaphore; pendingSettlements: Set<SandboxPromise> },
   ) {
     const globalScope = new Map<string, Binding>()
     this.scopes = [globalScope]
@@ -625,7 +626,8 @@ class Interpreter<R> {
     this.toolKeys = toolKeys
     this.logs = logs
     this.lastValue = undefined
-    this.callPermits = Semaphore.makeUnsafe(TOOL_CALL_CONCURRENCY)
+    this.callPermits = shared?.callPermits ?? Semaphore.makeUnsafe(TOOL_CALL_CONCURRENCY)
+    this.pendingSettlements = shared?.pendingSettlements ?? new Set<SandboxPromise>()
     globalScope.set("tools", { mutable: false, value: new ToolReference([]) })
     globalScope.set("Promise", { mutable: false, value: new PromiseNamespace() })
     globalScope.set("undefined", { mutable: false, value: undefined })
@@ -705,15 +707,17 @@ class Interpreter<R> {
   private drainPendingSettlements(): Effect.Effect<void, unknown, never> {
     const self = this
     return Effect.gen(function* () {
-      for (const promise of [...self.pendingSettlements]) {
+      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)) continue
         const failure = normalizeError(Cause.squash(exit.cause))
         throw new InterpreterRuntimeError(
-          `Unhandled rejection from an un-awaited tool call: ${failure.message}`,
+          `Unhandled rejection from an un-awaited promise: ${failure.message}`,
           undefined,
           failure.kind,
-          ["Await tool calls - `const result = await tools.ns.tool(...)` - so failures can be caught and handled."],
+          ["Await promises so failures can be caught and handled."],
         )
       }
     })
@@ -727,17 +731,15 @@ class Interpreter<R> {
     path: ReadonlyArray<string>,
     args: Array<unknown>,
   ): Effect.Effect<SandboxPromise, never, R> {
-    const self = this
-    return Effect.map(
-      Effect.forkChild(this.callPermits.withPermit(Effect.suspend(() => self.invokeTool(path, args))), {
-        startImmediately: true,
-      }),
-      (fiber) => {
-        const promise = new SandboxPromise(fiber)
-        self.pendingSettlements.add(promise)
-        return promise
-      },
-    )
+    return this.createPromise(this.callPermits.withPermit(Effect.suspend(() => this.invokeTool(path, args))))
+  }
+
+  private createPromise(effect: Effect.Effect<unknown, unknown, R>): Effect.Effect<SandboxPromise, never, R> {
+    return Effect.map(Effect.forkChild(effect, { startImmediately: true }), (fiber) => {
+      const promise = new SandboxPromise(fiber)
+      this.pendingSettlements.add(promise)
+      return promise
+    })
   }
 
   // The promise's settlement as an Exit, marking it observed for unhandled-rejection tracking.
@@ -858,6 +860,7 @@ class Interpreter<R> {
       getArray(node, "params").map((parameter, index) => asNode(parameter, `params[${index}]`)),
       getNode(node, "body"),
       this.scopes.slice(),
+      node.async === true,
     )
   }
 
@@ -2307,43 +2310,42 @@ class Interpreter<R> {
   }
 
   private invokeFunction(fn: CodeModeFunction, args: Array<unknown>): Effect.Effect<unknown, unknown, R> {
-    const self = this
-    return Effect.suspend(() => {
-      const savedScopes = self.scopes
-      self.scopes = [...fn.capturedScopes, new Map<string, Binding>()]
-      const run = Effect.gen(function* () {
-        // Seed every parameter name into the scope as a TDZ slot first, so a default that
-        // references another parameter resolves to that (uninitialized) param rather than
-        // silently falling through to an outer binding of the same name - matching JS.
-        const paramScope = self.currentScope()
-        for (const parameter of fn.parameters) {
-          for (const name of collectPatternNames(parameter)) {
-            paramScope.set(name, { mutable: true, value: undefined, initialized: false })
-          }
-        }
-        for (const [index, parameter] of fn.parameters.entries()) {
-          if (parameter.type === "RestElement") {
-            yield* self.declarePattern(getNode(parameter, "argument"), args.slice(index), true, parameter)
-            break
-          }
-          yield* self.declarePattern(parameter, args[index], true, parameter)
-        }
-
-        if (fn.body.type === "BlockStatement") {
-          const result = yield* self.evaluateStatement(fn.body)
-          return result.kind === "return" || result.kind === "value" ? result.value : undefined
-        }
-
-        return yield* self.evaluateExpression(fn.body)
-      })
-      return run.pipe(
-        Effect.ensuring(
-          Effect.sync(() => {
-            self.scopes = savedScopes
-          }),
-        ),
-      )
+    const invocation = new Interpreter(this.invokeTool, this.toolKeys, this.logs, {
+      callPermits: this.callPermits,
+      pendingSettlements: this.pendingSettlements,
     })
+    invocation.scopes = [...fn.capturedScopes, new Map<string, Binding>()]
+    const run = Effect.gen(function* () {
+      // Seed every parameter name into the scope as a TDZ slot first, so a default that
+      // references another parameter resolves to that (uninitialized) param rather than
+      // silently falling through to an outer binding of the same name - matching JS.
+      const paramScope = invocation.currentScope()
+      for (const parameter of fn.parameters) {
+        for (const name of collectPatternNames(parameter)) {
+          paramScope.set(name, { mutable: true, value: undefined, initialized: false })
+        }
+      }
+      for (const [index, parameter] of fn.parameters.entries()) {
+        if (parameter.type === "RestElement") {
+          yield* invocation.declarePattern(getNode(parameter, "argument"), args.slice(index), true, parameter)
+          break
+        }
+        yield* invocation.declarePattern(parameter, args[index], true, parameter)
+      }
+
+      if (fn.body.type === "BlockStatement") {
+        const result = yield* invocation.evaluateStatement(fn.body)
+        return result.kind === "return" || result.kind === "value" ? result.value : undefined
+      }
+
+      return yield* invocation.evaluateExpression(fn.body)
+    })
+    if (!fn.async) return run
+    return this.createPromise(
+      Effect.flatMap(run, (value) =>
+        value instanceof SandboxPromise ? invocation.settlePromise(value) : Effect.succeed(value),
+      ),
+    )
   }
 
   private invokeIntrinsic(
@@ -2432,13 +2434,19 @@ class Interpreter<R> {
       else value.replaceAll(pattern, collect)
     }
 
+    const self = this
     return Effect.gen(function* () {
       const output: Array<string> = []
       let end = 0
       for (const match of matches) {
+        const replacement = yield* apply(match.args)
+        const resolved =
+          args[1] instanceof CodeModeFunction && args[1].async && replacement instanceof SandboxPromise
+            ? yield* self.settlePromise(replacement)
+            : replacement
         output.push(
           value.slice(end, match.offset),
-          coerceToString(boundedData(yield* apply(match.args), `String.${name} replacer result`)),
+          coerceToString(boundedData(resolved, `String.${name} replacer result`)),
         )
         end = match.offset + match.match.length
       }

packages/codemode/test/promise.test.ts

@@ -75,6 +75,42 @@ const error = async (code: string, options: { trace?: Trace; limits?: CodeMode.E
 }
 
 describe("first-class promise values", () => {
+  test("async functions return promises with isolated concurrent invocations", async () => {
+    expect(
+      await value(`
+        const load = async (id) => {
+          const result = await tools.host.sleepy({ id, ms: 20 })
+          return [id, result]
+        }
+        const first = load(1)
+        const second = load(2)
+        return [first instanceof Promise, second instanceof Promise, await Promise.all([first, second])]
+      `),
+    ).toEqual([
+      true,
+      true,
+      [
+        [1, 1],
+        [2, 2],
+      ],
+    ])
+  })
+
+  test("async function errors reject instead of throwing at the call site", async () => {
+    expect(
+      await value(`
+        const fail = async () => { throw new Error("boom") }
+        const promise = fail()
+        try {
+          await promise
+          return "no"
+        } catch (error) {
+          return error.message
+        }
+      `),
+    ).toBe("boom")
+  })
+
   test("an un-awaited tool call starts eagerly, in call order, before any await", async () => {
     const trace = makeTrace()
     const result = await value(
@@ -163,9 +199,33 @@ describe("first-class promise values", () => {
       return "done"
     `)
     expect(diagnostic.kind).toBe("ToolFailure")
-    expect(diagnostic.message).toContain("Unhandled rejection from an un-awaited tool call")
+    expect(diagnostic.message).toContain("Unhandled rejection from an un-awaited promise")
+    expect(diagnostic.message).toContain("Lookup refused")
+    expect(diagnostic.suggestions?.join(" ")).toContain("Await promises")
+  })
+
+  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") }
+      fail()
+      return "done"
+    `)
+    expect(diagnostic.kind).toBe("ExecutionFailure")
+    expect(diagnostic.message).toContain("Unhandled rejection from an un-awaited promise")
+    expect(diagnostic.message).toContain("boom")
+  })
+
+  test("drains promises started by an async function after an await", async () => {
+    const diagnostic = await error(`
+      const run = async () => {
+        await tools.host.sleepy({ id: 1 })
+        tools.host.fail({})
+      }
+      run()
+      return "done"
+    `)
+    expect(diagnostic.kind).toBe("ToolFailure")
     expect(diagnostic.message).toContain("Lookup refused")
-    expect(diagnostic.suggestions?.join(" ")).toContain("await tools.ns.tool(...)")
   })
 })
 
@@ -232,6 +292,19 @@ describe("Promise.all over arbitrary arrays", () => {
     expect(trace.maxActive).toBeGreaterThan(1)
   })
 
+  test("runs async map callbacks concurrently", async () => {
+    const trace = makeTrace()
+    const result = await value(
+      `
+        const ids = [1, 2, 3, 4]
+        return await Promise.all(ids.map(async (id) => await tools.host.sleepy({ id, ms: 40 })))
+      `,
+      { trace },
+    )
+    expect(result).toEqual([1, 2, 3, 4])
+    expect(trace.maxActive).toBeGreaterThan(1)
+  })
+
   test("caps live tool-call concurrency at the fixed internal constant (8)", async () => {
     const trace = makeTrace()
     const result = await value(