diff --git a/.changeset/quickjs-native-promises.md b/.changeset/quickjs-native-promises.md new file mode 100644 index 000000000..fe48a8c63 --- /dev/null +++ b/.changeset/quickjs-native-promises.md @@ -0,0 +1,5 @@ +--- +'@tanstack/ai-isolate-quickjs': patch +--- + +Bridge asynchronous host tools with native QuickJS promises to prevent repeated awaited tool calls from corrupting the Asyncify runtime. diff --git a/packages/ai-isolate-quickjs/src/isolate-context.ts b/packages/ai-isolate-quickjs/src/isolate-context.ts index e23bcb76a..fc27a4d01 100644 --- a/packages/ai-isolate-quickjs/src/isolate-context.ts +++ b/packages/ai-isolate-quickjs/src/isolate-context.ts @@ -1,31 +1,131 @@ import { wrapCode } from '@tanstack/ai-code-mode' import { isFatalQuickJSLimitError, normalizeError } from './error-normalizer' -import type { QuickJSAsyncContext } from 'quickjs-emscripten' +import type { + QuickJSContext, + QuickJSHandle, + VmCallResult, +} from 'quickjs-emscripten' import type { ExecutionResult, IsolateContext } from '@tanstack/ai-code-mode' /** - * Serializes all QuickJS evalCodeAsync calls across contexts. - * Required because newAsyncContext() reuses a singleton WASM module - * whose asyncify stack can only handle one suspension at a time. + * Preserves the driver's existing execution ordering across contexts. */ let globalExecQueue: Promise = Promise.resolve() +function awaitWithDeadline( + promise: Promise>, + pendingJobFailure: Promise, + deadline: number, +): Promise> { + return new Promise((resolve, reject) => { + let settled = false + const timer = setTimeout( + () => { + settled = true + const error = new Error('Code execution timed out') + error.name = 'TimeoutError' + reject(error) + }, + Math.max(0, deadline - Date.now()), + ) + + promise.then( + (result) => { + if (settled) { + try { + if (result.error) result.error.dispose() + else result.value.dispose() + } catch { + // The context may have been disposed after a timeout. + } + return + } + settled = true + clearTimeout(timer) + resolve(result) + }, + (error) => { + if (settled) return + settled = true + clearTimeout(timer) + reject(error) + }, + ) + + void pendingJobFailure.catch((error: unknown) => { + if (settled) return + settled = true + clearTimeout(timer) + reject(error) + }) + }) +} + /** * IsolateContext implementation using QuickJS WASM */ export class QuickJSIsolateContext implements IsolateContext { - private readonly vm: QuickJSAsyncContext + private readonly vm: QuickJSContext private readonly logs: Array private readonly timeout: number + private rejectPendingJobFailure?: (error: unknown) => void + private readonly pendingHostCalls = new Set<() => void>() + private disposeRequested = false private disposed = false private executing = false - constructor(vm: QuickJSAsyncContext, logs: Array, timeout: number) { + constructor(vm: QuickJSContext, logs: Array, timeout: number) { this.vm = vm this.logs = logs this.timeout = timeout } + runPendingJobs(): void { + try { + const result = this.vm.runtime.executePendingJobs() + if (!result.error) { + result.dispose() + return + } + + let error: unknown + try { + error = result.error.context.dump(result.error) + } finally { + result.dispose() + } + this.rejectPendingJobFailure?.(error) + } catch (error) { + this.rejectPendingJobFailure?.(error) + } + } + + beginHostCall(cancel: () => void): () => void { + this.pendingHostCalls.add(cancel) + let completed = false + + return () => { + if (completed) return + completed = true + this.pendingHostCalls.delete(cancel) + if (this.disposeRequested && this.pendingHostCalls.size === 0) { + this.vm.dispose() + } + } + } + + private requestVmDispose(): void { + if (this.disposeRequested) return + this.disposeRequested = true + this.disposed = true + for (const cancel of [...this.pendingHostCalls]) { + cancel() + } + if (this.pendingHostCalls.size === 0 && this.vm.alive) { + this.vm.dispose() + } + } + async execute(code: string): Promise> { if (this.disposed) { return { @@ -38,8 +138,7 @@ export class QuickJSIsolateContext implements IsolateContext { } } - // Serialize through the global queue to prevent concurrent - // WASM asyncify suspensions across contexts. + // Preserve the driver's existing cross-context execution ordering. let resolve!: () => void const myTurn = new Promise((r) => { resolve = r @@ -64,22 +163,27 @@ export class QuickJSIsolateContext implements IsolateContext { this.executing = true this.logs.length = 0 + const pendingJobFailure = new Promise((_, reject) => { + this.rejectPendingJobFailure = reject + }) - const releaseVmAfterFatalLimit = () => { + const releaseVmAfterTerminalError = () => { if (this.disposed) return try { this.vm.runtime.setInterruptHandler(() => false) } catch { // ignore if runtime is already torn down } - this.disposed = true - this.vm.dispose() + this.requestVmDispose() } const fail = (error: unknown) => { const normalized = normalizeError(error) - if (isFatalQuickJSLimitError(normalized)) { - releaseVmAfterFatalLimit() + if ( + normalized.name === 'TimeoutError' || + isFatalQuickJSLimitError(normalized) + ) { + releaseVmAfterTerminalError() } return { success: false as const, @@ -97,19 +201,21 @@ export class QuickJSIsolateContext implements IsolateContext { }) try { - const result = await this.vm.evalCodeAsync(wrappedCode) + const result = this.vm.evalCode(wrappedCode) let parsedResult: T try { const promiseHandle = this.vm.unwrapResult(result) - // evalCodeAsync returns a Promise handle (our wrapper is an async IIFE). - // Use resolvePromise + executePendingJobs to properly await the - // QuickJS promise without re-entering the WASM asyncify state. + // wrapCode returns an async IIFE, so evalCode yields a QuickJS promise. const nativePromise = this.vm.resolvePromise(promiseHandle) promiseHandle.dispose() - this.vm.runtime.executePendingJobs() - const resolvedResult = await nativePromise + this.runPendingJobs() + const resolvedResult = await awaitWithDeadline( + nativePromise, + pendingJobFailure, + deadline, + ) const valueHandle = this.vm.unwrapResult(resolvedResult) const dumpedResult = this.vm.dump(valueHandle) @@ -134,7 +240,7 @@ export class QuickJSIsolateContext implements IsolateContext { return fail(unwrapError) } } finally { - // fail() may set disposed when releasing the VM after memory/stack limit errors + // fail() may dispose the VM after a timeout or fatal resource limit. // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition -- disposed set in fail() if (!this.disposed) { this.vm.runtime.setInterruptHandler(() => false) @@ -143,6 +249,7 @@ export class QuickJSIsolateContext implements IsolateContext { } catch (error) { return fail(error) } finally { + this.rejectPendingJobFailure = undefined this.executing = false resolve() } @@ -151,14 +258,12 @@ export class QuickJSIsolateContext implements IsolateContext { async dispose(): Promise { if (this.disposed) return - // If an execution is in flight, wait for the global queue to drain - // before disposing the VM. Otherwise the asyncified callback would - // try to access a freed context. + // If an execution is in flight, wait for the global queue to drain so + // pending host callbacks cannot access a disposed context. if (this.executing) { await globalExecQueue } - this.disposed = true - this.vm.dispose() + this.requestVmDispose() } } diff --git a/packages/ai-isolate-quickjs/src/isolate-driver.ts b/packages/ai-isolate-quickjs/src/isolate-driver.ts index 971e0bb96..dfef44035 100644 --- a/packages/ai-isolate-quickjs/src/isolate-driver.ts +++ b/packages/ai-isolate-quickjs/src/isolate-driver.ts @@ -1,10 +1,12 @@ -import { newAsyncContext } from 'quickjs-emscripten' +import { getQuickJS } from 'quickjs-emscripten' import { QuickJSIsolateContext } from './isolate-context' import type { IsolateConfig, IsolateContext, IsolateDriver, + ToolBinding, } from '@tanstack/ai-code-mode' +import type { QuickJSContext } from 'quickjs-emscripten' /** Default memory limit in MB (matches Node isolate driver default). */ const DEFAULT_MEMORY_LIMIT_MB = 128 @@ -34,6 +36,80 @@ export interface QuickJSIsolateDriverConfig { maxStackSize?: number } +async function invokeBinding( + binding: ToolBinding, + argsJson: string, +): Promise { + try { + const value = await binding.execute(JSON.parse(argsJson)) + return JSON.stringify({ success: true, value }) + } catch (error) { + return JSON.stringify({ + success: false, + error: error instanceof Error ? error.message : String(error), + }) + } +} + +function injectBinding( + vm: QuickJSContext, + context: QuickJSIsolateContext, + name: string, + binding: ToolBinding, +): void { + const toolFn = vm.newFunction(name, (argsHandle) => { + const argsJson = vm.getString(argsHandle) + const deferred = vm.newPromise() + let completeHostCall = () => {} + const cancelHostCall = () => { + if (deferred.alive) deferred.dispose() + completeHostCall() + } + completeHostCall = context.beginHostCall(cancelHostCall) + + void invokeBinding(binding, argsJson).then((payloadJson) => { + if (!vm.alive || !deferred.alive) { + completeHostCall() + return + } + + const payloadHandle = vm.newString(payloadJson) + deferred.resolve(payloadHandle) + payloadHandle.dispose() + + void deferred.settled + .then(() => { + if (vm.runtime.alive) { + context.runPendingJobs() + } + }) + .finally(completeHostCall) + }) + + return deferred.handle + }) + + vm.setProp(vm.global, `__${name}_impl`, toolFn) + toolFn.dispose() + + const wrapperResult = vm.evalCode(` + async function ${name}(input) { + const resultJson = await __${name}_impl(JSON.stringify(input)); + const result = JSON.parse(resultJson); + if (!result.success) { + throw new Error(result.error); + } + return result.value; + } + `) + if (wrapperResult.error) { + const errorStr = vm.dump(wrapperResult.error) + wrapperResult.error.dispose() + throw new Error(`Failed to create wrapper for ${name}: ${errorStr}`) + } + wrapperResult.value.dispose() +} + /** * Create a QuickJS WASM isolate driver * @@ -82,8 +158,8 @@ export function createQuickJSIsolateDriver( const memoryLimitMb = isolateConfig.memoryLimit ?? defaultMemoryLimit const maxStackSizeBytes = defaultMaxStackSize - // Create async QuickJS context (supports async host functions) - const vm = await newAsyncContext() + const QuickJS = await getQuickJS() + const vm = QuickJS.newContext() // Enforce heap and stack limits so OOM/stack overflow surface as JS errors // instead of growing WASM memory until the host process OOMs. @@ -126,60 +202,15 @@ export function createQuickJSIsolateDriver( infoFn.dispose() consoleObj.dispose() - // Inject each tool binding as an async function - for (const [name, binding] of Object.entries(isolateConfig.bindings)) { - // Create async function that calls back to host - // newAsyncifiedFunction receives QuickJS handles as arguments - const toolFn = vm.newAsyncifiedFunction(name, async (argsHandle) => { - try { - // Get the input argument - argsHandle is a QuickJS handle - const argsJson = vm.getString(argsHandle) - const args = JSON.parse(argsJson) - - // Execute the tool on the host - const result = await binding.execute(args) - - // Return result as JSON string handle - const returnHandle = vm.newString( - JSON.stringify({ success: true, value: result }), - ) - return returnHandle - } catch (error) { - const errorMessage = - error instanceof Error ? error.message : String(error) - const returnHandle = vm.newString( - JSON.stringify({ success: false, error: errorMessage }), - ) - return returnHandle - } - }) + const context = new QuickJSIsolateContext(vm, logs, timeout) - // Set on global - the VM keeps its own reference - vm.setProp(vm.global, `__${name}_impl`, toolFn) - toolFn.dispose() - - // Create wrapper that parses input and output - // Function names match the binding keys (e.g., external_fetchWeather) - const wrapperCode = ` - async function ${name}(input) { - const resultJson = await __${name}_impl(JSON.stringify(input)); - const result = JSON.parse(resultJson); - if (!result.success) { - throw new Error(result.error); - } - return result.value; - } - ` - const wrapperResult = vm.evalCode(wrapperCode) - if (wrapperResult.error) { - const errorStr = vm.dump(wrapperResult.error) - wrapperResult.error.dispose() - throw new Error(`Failed to create wrapper for ${name}: ${errorStr}`) - } - wrapperResult.value.dispose() + // Return native QuickJS promises from synchronous host functions so host + // work does not suspend and resume the WASM stack through Asyncify. + for (const [name, binding] of Object.entries(isolateConfig.bindings)) { + injectBinding(vm, context, name, binding) } - return new QuickJSIsolateContext(vm, logs, timeout) + return context }, } } diff --git a/packages/ai-isolate-quickjs/tests/isolate-driver.test.ts b/packages/ai-isolate-quickjs/tests/isolate-driver.test.ts index eb1506afb..04213d643 100644 --- a/packages/ai-isolate-quickjs/tests/isolate-driver.test.ts +++ b/packages/ai-isolate-quickjs/tests/isolate-driver.test.ts @@ -96,21 +96,46 @@ describe('createQuickJSIsolateDriver', () => { expect(result.value).toBe(5) }) - it('supports multiple tools in one execution', async () => { + it('supports multiple sequential tool awaits and context reuse', async () => { const getA = makeBinding('getA', async () => 'A') const getB = makeBinding('getB', async () => 'B') + const getC = makeBinding('getC', async () => 'C') const driver = createQuickJSIsolateDriver() const context = await driver.createContext({ - bindings: { getA, getB }, + bindings: { getA, getB, getC }, }) const result = await context.execute(` - return (await getA({})) + (await getB({})); + return (await getA({})) + (await getB({})) + (await getC({})); `) + const reuse = await context.execute('return 1 + 1') expect(result.success).toBe(true) - expect(result.value).toBe('AB') + expect(result.value).toBe('ABC') + expect(reuse).toMatchObject({ success: true, value: 2 }) + }) + + it('supports concurrent host calls with Promise.all', async () => { + let activeCalls = 0 + let maximumActiveCalls = 0 + const get = makeBinding('get', async () => { + activeCalls += 1 + maximumActiveCalls = Math.max(maximumActiveCalls, activeCalls) + await new Promise((resolve) => setTimeout(resolve, 5)) + activeCalls -= 1 + return activeCalls + }) + + const context = await createQuickJSIsolateDriver().createContext({ + bindings: { get }, + }) + const result = await context.execute(` + return Promise.all([get({}), get({}), get({})]); + `) + + expect(result.success).toBe(true) + expect(maximumActiveCalls).toBe(3) }) it('surfaces tool execution errors', async () => { @@ -150,6 +175,41 @@ describe('createQuickJSIsolateDriver', () => { expect(result.success).toBe(false) expect(result.error?.name).toBeDefined() }) + + it('returns a timeout when a host tool settles after the deadline', async () => { + const delayed = makeBinding('delayed', async () => { + await new Promise((resolve) => setTimeout(resolve, 100)) + return 'late' + }) + const context = await createQuickJSIsolateDriver({ + timeout: 20, + }).createContext({ bindings: { delayed } }) + + const result = await context.execute('return await delayed({})') + const reuse = await context.execute('return 1 + 1') + + expect(result.success).toBe(false) + expect(result.error?.name).toBe('TimeoutError') + expect(reuse.error?.name).toBe('DisposedError') + }) + + it('disposes the VM when a timed-out host tool never settles', async () => { + const hanging = makeBinding('hanging', () => new Promise(() => {})) + const context = await createQuickJSIsolateDriver({ + timeout: 20, + }).createContext({ bindings: { hanging } }) + + const result = await context.execute('return await hanging({})') + const vm = ( + context as unknown as { + vm: { alive: boolean } + } + ).vm + + expect(result.error?.name).toBe('TimeoutError') + expect(vm.alive).toBe(false) + await expect(context.dispose()).resolves.toBeUndefined() + }) }) describe('execute - error handling', () => {