fix(ai-isolate-quickjs): bridge host tools with native promises

.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.

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<void> = Promise.resolve()
 
+function awaitWithDeadline(
+  promise: Promise<VmCallResult<QuickJSHandle>>,
+  pendingJobFailure: Promise<never>,
+  deadline: number,
+): Promise<VmCallResult<QuickJSHandle>> {
+  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<string>
   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<string>, timeout: number) {
+  constructor(vm: QuickJSContext, logs: Array<string>, 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<T = unknown>(code: string): Promise<ExecutionResult<T>> {
     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<void>((r) => {
       resolve = r
@@ -64,22 +163,27 @@ export class QuickJSIsolateContext implements IsolateContext {
 
     this.executing = true
     this.logs.length = 0
+    const pendingJobFailure = new Promise<never>((_, 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<void> {
     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()
   }
 }