From 754dd57177105142ccc1b954aba60206c52beb9d Mon Sep 17 00:00:00 2001 From: Kit Langton Date: Mon, 6 Jul 2026 17:29:11 -0400 Subject: [PATCH] fix(codemode): preserve async callback semantics (#35603) --- packages/codemode/src/interpreter/model.ts | 4 +- packages/codemode/src/interpreter/runtime.ts | 114 ++++++++++--------- packages/codemode/test/promise.test.ts | 77 ++++++++++++- 3 files changed, 139 insertions(+), 56 deletions(-) diff --git a/packages/codemode/src/interpreter/model.ts b/packages/codemode/src/interpreter/model.ts index 0869a89cb28..08d47cd37e8 100644 --- a/packages/codemode/src/interpreter/model.ts +++ b/packages/codemode/src/interpreter/model.ts @@ -45,6 +45,7 @@ export class CodeModeFunction { readonly parameters: ReadonlyArray, readonly body: AstNode, readonly capturedScopes: ReadonlyArray>, + readonly async: boolean, ) {} } @@ -153,7 +154,8 @@ export const unsupportedSyntax = (kind: string, node: AstNode): InterpreterRunti [supportedSyntaxMessage], ) -export const isRecord = (value: unknown): value is Record => typeof value === "object" && value !== null +export const isRecord = (value: unknown): value is Record => + typeof value === "object" && value !== null export const asNode = (value: unknown, context: string): AstNode => { if (!isRecord(value) || typeof value.type !== "string") { diff --git a/packages/codemode/src/interpreter/runtime.ts b/packages/codemode/src/interpreter/runtime.ts index f34e3949e12..b1d822281ac 100644 --- a/packages/codemode/src/interpreter/runtime.ts +++ b/packages/codemode/src/interpreter/runtime.ts @@ -612,12 +612,13 @@ class Interpreter { // 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() + private readonly pendingSettlements: Set constructor( invokeTool: (path: ReadonlyArray, args: Array) => Effect.Effect, toolKeys: (path: ReadonlyArray) => ReadonlyArray, logs: Array = [], + shared?: { callPermits: Semaphore.Semaphore; pendingSettlements: Set }, ) { const globalScope = new Map() this.scopes = [globalScope] @@ -625,7 +626,8 @@ class Interpreter { 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() 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 { private drainPendingSettlements(): Effect.Effect { 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 { path: ReadonlyArray, args: Array, ): Effect.Effect { - 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): Effect.Effect { + 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 { 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 { } private invokeFunction(fn: CodeModeFunction, args: Array): Effect.Effect { - const self = this - return Effect.suspend(() => { - const savedScopes = self.scopes - self.scopes = [...fn.capturedScopes, new Map()] - 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()] + 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 { else value.replaceAll(pattern, collect) } + const self = this return Effect.gen(function* () { const output: Array = [] 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 } diff --git a/packages/codemode/test/promise.test.ts b/packages/codemode/test/promise.test.ts index 545d463abfa..b56291daf87 100644 --- a/packages/codemode/test/promise.test.ts +++ b/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(