diff --git a/packages/llmz/package.json b/packages/llmz/package.json index f74d780c4cb..aea03736450 100644 --- a/packages/llmz/package.json +++ b/packages/llmz/package.json @@ -2,7 +2,7 @@ "name": "llmz", "type": "module", "description": "LLMz - An LLM-native Typescript VM built on top of Zui", - "version": "0.0.80", + "version": "0.0.81", "types": "./dist/index.d.ts", "main": "./dist/index.cjs", "module": "./dist/index.js", diff --git a/packages/llmz/src/index.ts b/packages/llmz/src/index.ts index 35934bbd84b..cc85f35c464 100644 --- a/packages/llmz/src/index.ts +++ b/packages/llmz/src/index.ts @@ -127,7 +127,7 @@ export const init = async () => { await import('./tool.js') await import('./exit.js') await import('./jsx.js') - await import('./vm.js') + await import('./vm/index.js') await import('./utils.js') await import('./truncator.js') await import('./typings.js') diff --git a/packages/llmz/src/llmz.ts b/packages/llmz/src/llmz.ts index bf40c91d78c..cef09d9e109 100644 --- a/packages/llmz/src/llmz.ts +++ b/packages/llmz/src/llmz.ts @@ -35,7 +35,7 @@ import { truncateWrappedContent } from './truncator.js' import { Trace } from './types.js' import { init, stripInvalidIdentifiers } from './utils.js' -import { runAsyncFunction } from './vm.js' +import { runAsyncFunction } from './vm/index.js' const getErrorMessage = (err: unknown) => (err instanceof Error ? err.message : JSON.stringify(err)) diff --git a/packages/llmz/src/quickjs-variant.ts b/packages/llmz/src/quickjs-variant.ts index d93b4e8e543..59ff1d38f92 100644 --- a/packages/llmz/src/quickjs-variant.ts +++ b/packages/llmz/src/quickjs-variant.ts @@ -11,7 +11,14 @@ const getVariant = async (): Promise => { return module.default as unknown as QuickJSSyncVariant } -export const BundledReleaseSyncVariant: QuickJSSyncVariant = { +export type QuickJSSyncVariantEx = QuickJSSyncVariant & { + _wasmSource?: any + _wasmLoadedSuccessfully?: any + _wasmSize?: any + _wasmLoadError?: any +} + +export const BundledReleaseSyncVariant: QuickJSSyncVariantEx = { type: 'sync', importFFI: async () => { const variant = await getVariant() diff --git a/packages/llmz/src/vm.ts b/packages/llmz/src/vm.ts deleted file mode 100644 index 6b8ec402077..00000000000 --- a/packages/llmz/src/vm.ts +++ /dev/null @@ -1,1285 +0,0 @@ -/** - * LLMz VM Implementation - * - * Supports two execution drivers: - * 1. QuickJS (quickjs-emscripten) - Sandboxed execution with memory limits and timeout - * 2. Node (native) - Direct execution without sandbox, for environments where QuickJS is not available - */ - -/* oxlint-disable max-depth */ - -import { isFunction, mapValues, maxBy } from 'lodash-es' -import { newQuickJSWASMModuleFromVariant, shouldInterruptAfterDeadline } from 'quickjs-emscripten-core' -import { SourceMapConsumer } from 'source-map-js' - -import { compile, CompiledCode, Identifiers } from './compiler/index.js' -import { CodeExecutionError, InvalidCodeError, Signals, SnapshotSignal, VMSignal } from './errors.js' -import { createJsxComponent, JsxComponent } from './jsx.js' -import { BundledReleaseSyncVariant } from './quickjs-variant.js' -import { cleanStackTrace } from './stack-traces.js' -import { Trace, Traces, VMExecutionResult } from './types.js' - -const MAX_VM_EXECUTION_TIME = 60_000 - -type Driver = 'quickjs' | 'node' - -// These are the identifiers that we want to exclude from the variable tracking system -const NO_TRACKING = [ - Identifiers.CommentFnIdentifier, - Identifiers.ToolCallTrackerFnIdentifier, - Identifiers.ToolTrackerRetIdentifier, - Identifiers.VariableTrackingFnIdentifier, - Identifiers.JSXFnIdentifier, - Identifiers.ConsoleObjIdentifier, -] as const - -function getCompiledCode(code: string, traces: Trace[] = []): CompiledCode { - try { - return compile(code) - } catch (err: any) { - traces.push({ - type: 'invalid_code_exception', - message: err?.message ?? 'Unknown error', - code, - started_at: Date.now(), - }) - throw new InvalidCodeError(err.message, code) - } -} - -export async function runAsyncFunction( - context: any, - code: string, - traces: Trace[] = [], - signal: AbortSignal | null = null, - timeout: number = MAX_VM_EXECUTION_TIME -): Promise { - const transformed = getCompiledCode(code, traces) - const lines_executed = new Map() - const variables: { [k: string]: any } = {} - - const consumer = new SourceMapConsumer({ - version: transformed.map!.version.toString(), - mappings: transformed.map!.mappings, - names: transformed.map!.names!, - sources: [transformed.map!.file!], - sourcesContent: [transformed.code!], - file: transformed.map!.file!, - sourceRoot: transformed.map!.sourceRoot!, - }) - - context ??= {} - - // Remove variables that will be tracked - for (const name of Array.from(transformed.variables)) { - delete context[name] - } - - // Determine which driver to use - try QuickJS first, fallback to node if it fails - let DRIVER: Driver = 'quickjs' - - // Check if user explicitly disabled QuickJS - if (typeof process !== 'undefined' && process?.env?.USE_QUICKJS === 'false') { - DRIVER = 'node' - } - - // ============================================================================ - // QuickJS Driver - // ============================================================================ - if (DRIVER === 'quickjs') { - // Try to load QuickJS - if it fails, fallback to node driver - try { - // Setup tracking functions - context[Identifiers.CommentFnIdentifier] = (comment: string, line: number) => { - // Filter out internal markers from traces - if (comment.includes('__LLMZ_USER_CODE_START__') || comment.includes('__LLMZ_USER_CODE_END__')) { - return - } - traces.push({ type: 'comment', comment, line, started_at: Date.now() }) - } - - // Find the actual offset by locating the user code start marker - // Use codeWithMarkers which still has the markers before postProcessing removes them - const codeWithMarkers = (transformed as any).codeWithMarkers || transformed.code - const markerLines = codeWithMarkers.split('\n') - const USER_CODE_START_MARKER = '/* __LLMZ_USER_CODE_START__ */' - let userCodeStartLine = -1 - for (let i = 0; i < markerLines.length; i++) { - if (markerLines[i]?.includes(USER_CODE_START_MARKER)) { - userCodeStartLine = i + 1 // Line numbers are 1-indexed - break - } - } - - context[Identifiers.LineTrackingFnIdentifier] = (line: number) => { - // Map the transformed code line back to the original source line - const originalLine = consumer.originalPositionFor({ - line, - column: 0, - }) - const mappedLine = originalLine.line ?? line - - // Calculate offset: the marker line in transformed code corresponds to line 0 of user code - // So user line = mapped line - marker line - const userCodeLine = Math.max(1, mappedLine - userCodeStartLine) - - lines_executed.set(userCodeLine, (lines_executed.get(userCodeLine) ?? 0) + 1) - } - - context[Identifiers.JSXFnIdentifier] = (tool: any, props: any, ...children: any[]) => - createJsxComponent({ - type: tool, - props, - children, - }) - - context[Identifiers.VariableTrackingFnIdentifier] = (name: string, getter: () => any) => { - if (NO_TRACKING.includes(name)) { - return - } - variables[name] = () => { - try { - const value = getter() - if (typeof value === 'function') { - return '[[non-primitive]]' - } - return value - } catch { - return '[[non-primitive]]' - } - } - } - - let currentToolCall: SnapshotSignal['toolCall'] | undefined - context[Identifiers.ToolCallTrackerFnIdentifier] = ( - callId: number, - type: 'start' | 'end', - outputOrError?: any - ) => { - const temp = Signals.maybeDeserializeError(outputOrError?.message) - if (type === 'end' && temp instanceof SnapshotSignal && temp?.toolCall) { - currentToolCall = { - ...temp.toolCall, - assignment: transformed.toolCalls.get(callId)?.assignment, - } - } - } - - context[Identifiers.ConsoleObjIdentifier] = { - log: (...args: any[]) => { - const message = args.shift() - traces.push({ type: 'log', message, args, started_at: Date.now() }) - }, - } - - context[Identifiers.AsyncIterYieldFnIdentifier] = async function (value: JsxComponent) { - const startedAt = Date.now() - try { - if (typeof value.type !== 'string' || value.type.trim().length === 0) { - throw new Error('A yield statement must yield a valid tool') - } - - const toolName = Object.keys(context).find((x) => x.toUpperCase() === value.type.toUpperCase()) - - if (!toolName) { - throw new Error(`Yield tool "${value.type}", but tool is not found`) - } - - await context[toolName](value) - } finally { - traces.push({ type: 'yield', value, started_at: startedAt, ended_at: Date.now() }) - } - } - - // Initialize QuickJS using our bundled variant - // This includes the WASM file directly in llmz's dist/ to avoid path resolution issues - const QuickJS = await newQuickJSWASMModuleFromVariant(BundledReleaseSyncVariant) - const runtime = QuickJS.newRuntime() - - // Set memory limit (128MB) - runtime.setMemoryLimit(128 * 1024 * 1024) - - // Set interrupt handler for timeout and abort signal - const startTime = Date.now() - const timeoutHandler = shouldInterruptAfterDeadline(startTime + timeout) - - runtime.setInterruptHandler(() => { - // Check if execution was aborted - if (signal?.aborted) { - return true // Interrupt execution - } - // Check if timeout exceeded - return timeoutHandler(runtime) - }) - - const vm = runtime.newContext() - - // Track which properties need to be copied back - const trackedProperties = new Set() - const referenceProperties = new Set() - - // Track all pending promises - we need to resolve them synchronously before disposing VM - const pendingPromises: Array<{ - hostPromise: Promise - deferredPromise: any - }> = [] - - // Helper to convert JS value to QuickJS handle - const toVmValue = (value: any): any => { - if (typeof value === 'string') { - return vm.newString(value) - } else if (typeof value === 'number') { - return vm.newNumber(value) - } else if (typeof value === 'boolean') { - return value ? vm.true : vm.false - } else if (value === null) { - return vm.null - } else if (value === undefined) { - return vm.undefined - } else if (Array.isArray(value)) { - // Create a proper array with prototype methods using evalCode - // We build the array literal directly to preserve methods like .map() - const items = value.map((item) => { - if (typeof item === 'string') { - return JSON.stringify(item) - } else if (typeof item === 'number' || typeof item === 'boolean') { - return String(item) - } else if (item === null) { - return 'null' - } else if (item === undefined) { - return 'undefined' - } else if (typeof item === 'object') { - return JSON.stringify(item) - } - return 'undefined' - }) - const arrayLiteral = `[${items.join(',')}]` - const result = vm.evalCode(arrayLiteral) - if ('error' in result) { - result.error?.dispose() - return vm.undefined - } - const arrHandle = result.value - return arrHandle - } else if (typeof value === 'object') { - const obj = vm.newObject() - for (const [k, v] of Object.entries(value)) { - if (typeof v !== 'function') { - const propHandle = toVmValue(v) - vm.setProp(obj, k, propHandle) - disposeIfNeeded(propHandle) - } - } - return obj - } - return vm.undefined - } - - const disposeIfNeeded = (handle: any) => { - if (handle !== vm.true && handle !== vm.false && handle !== vm.null && handle !== vm.undefined) { - handle.dispose() - } - } - - // Helper to bridge functions - handles both sync and async - const bridgeFunction = (fn: Function, _fnName: string = 'anonymous') => { - return (...argHandles: any[]) => { - const args = argHandles.map((h: any) => vm.dump(h)) - try { - const result = fn(...args) - // Check if it's a promise - create a QuickJS deferred promise - if (result && typeof result.then === 'function') { - // Create a QuickJS deferred promise - const promise = vm.newPromise() - - // Track this promise so we can await and resolve it synchronously - pendingPromises.push({ - hostPromise: result, - deferredPromise: promise, - }) - - // Schedule executePendingJobs when the promise settles - void promise.settled.then(() => { - if (runtime.alive) { - runtime.executePendingJobs() - } - }) - - // Return the promise handle - return promise.handle - } - // Synchronous result - return toVmValue(result) - } catch (err) { - // Serialize the error properly (especially for VMSignal and other special errors) - // VMSignal and other error classes auto-serialize themselves in their constructor - // so err.message already contains the JSON-serialized error data - const serialized = err instanceof Error ? err.message : String(err) - - // Re-throw the error as a plain Error with the serialized message - // QuickJS-emscripten will catch this and convert it to a QuickJS error - throw new Error(serialized) - } - } - } - - try { - // Bridge context values to QuickJS - for (const [key, value] of Object.entries(context)) { - const descriptor = Object.getOwnPropertyDescriptor(context, key) - - if (descriptor && (descriptor.get || descriptor.set)) { - // Handle getter/setter properties on globalThis - referenceProperties.add(key) - trackedProperties.add(key) - - // Create getter function if exists - let getterCode = 'undefined' - if (descriptor.get) { - const getterBridge = vm.newFunction(`get_${key}`, () => { - try { - const hostValue = (context as any)[key] - return toVmValue(hostValue) - } catch (err: any) { - const serialized = err instanceof Error ? err.message : String(err) - throw new Error(serialized) - } - }) - const getterName = `__getter_${key}__` - vm.setProp(vm.global, getterName, getterBridge) - getterBridge.dispose() - getterCode = getterName - } - - // Create setter function if exists - let setterCode = 'undefined' - if (descriptor.set) { - const setterBridge = vm.newFunction(`set_${key}`, (valueHandle: any) => { - try { - const jsValue = vm.dump(valueHandle) - ;(context as any)[key] = jsValue - return vm.undefined - } catch (err: any) { - const serialized = err instanceof Error ? err.message : String(err) - throw new Error(serialized) - } - }) - const setterName = `__setter_${key}__` - vm.setProp(vm.global, setterName, setterBridge) - setterBridge.dispose() - setterCode = setterName - } - - // Use evalCode to define the property with getter/setter on globalThis - const definePropertyCode = ` - Object.defineProperty(globalThis, '${key}', { - enumerable: true, - configurable: ${descriptor.configurable !== false}, - get: ${getterCode}, - set: ${setterCode} - }); - ` - const result = vm.evalCode(definePropertyCode) - if ('error' in result) { - result.error?.dispose() - } else { - result.value.dispose() - } - continue - } - - if (typeof value === 'function') { - // Bridge functions - supports both sync and async - const fnHandle = vm.newFunction(key, bridgeFunction(value, key)) - vm.setProp(vm.global, key, fnHandle) - fnHandle.dispose() - } else if (Array.isArray(value)) { - // Bridge arrays - use toVmValue which creates proper arrays with methods - trackedProperties.add(key) - const arrayHandle = toVmValue(value) - vm.setProp(vm.global, key, arrayHandle) - disposeIfNeeded(arrayHandle) - } else if (typeof value === 'object' && value !== null) { - trackedProperties.add(key) - - // Bridge objects - handle nested functions and properties - const objHandle = vm.newObject() - const props = new Set([...Object.keys(value), ...Object.getOwnPropertyNames(value)]) - const getterSetterProps: Array<{ - prop: string - descriptor: PropertyDescriptor - }> = [] - - for (const prop of props) { - const propDescriptor = Object.getOwnPropertyDescriptor(value, prop) - - if (propDescriptor && (propDescriptor.get || propDescriptor.set)) { - // Defer getter/setter setup until after object is on global - referenceProperties.add(`${key}.${prop}`) - getterSetterProps.push({ prop, descriptor: propDescriptor }) - } else if (typeof (value as any)[prop] === 'function') { - // Bridge nested functions - supports both sync and async - const propFnHandle = vm.newFunction(prop, bridgeFunction((value as any)[prop], `${key}.${prop}`)) - vm.setProp(objHandle, prop, propFnHandle) - propFnHandle.dispose() - } else { - const propHandle = toVmValue((value as any)[prop]) - vm.setProp(objHandle, prop, propHandle) - disposeIfNeeded(propHandle) - } - } - - vm.setProp(vm.global, key, objHandle) - objHandle.dispose() - - // Now set up getter/setter properties (after object is on global) - for (const { prop, descriptor } of getterSetterProps) { - // Create getter function if exists - let getterCode = 'undefined' - if (descriptor.get) { - const getterBridge = vm.newFunction(`get_${prop}`, () => { - try { - const hostValue = (context as any)[key][prop] - return toVmValue(hostValue) - } catch (err: any) { - const serialized = err instanceof Error ? err.message : String(err) - throw new Error(serialized) - } - }) - const getterName = `__getter_${key}_${prop}__` - vm.setProp(vm.global, getterName, getterBridge) - getterBridge.dispose() - getterCode = getterName - } - - // Create setter function if exists - let setterCode = 'undefined' - if (descriptor.set) { - const setterBridge = vm.newFunction(`set_${prop}`, (valueHandle: any) => { - try { - const jsValue = vm.dump(valueHandle) - ;(context as any)[key][prop] = jsValue - return vm.undefined - } catch (err: any) { - const serialized = err instanceof Error ? err.message : String(err) - throw new Error(serialized) - } - }) - const setterName = `__setter_${key}_${prop}__` - vm.setProp(vm.global, setterName, setterBridge) - setterBridge.dispose() - setterCode = setterName - } - - // Use evalCode to define the property with getter/setter - const definePropertyCode = ` - Object.defineProperty(${key}, '${prop}', { - enumerable: true, - configurable: ${descriptor.configurable !== false}, - get: ${getterCode}, - set: ${setterCode} - }); - ` - const result = vm.evalCode(definePropertyCode) - if ('error' in result) { - result.error?.dispose() - } else { - result.value.dispose() - } - } - - // Apply object constraints - if (Object.isSealed(value)) { - const sealResult = vm.evalCode(`Object.seal(globalThis['${key}']);`) - if ('error' in sealResult) { - sealResult.error?.dispose() - } else { - sealResult.value.dispose() - } - } - if (!Object.isExtensible(value)) { - const preventResult = vm.evalCode(`Object.preventExtensions(globalThis['${key}']);`) - if ('error' in preventResult) { - preventResult.error?.dispose() - } else { - preventResult.value.dispose() - } - } - } else { - // Bridge primitives - trackedProperties.add(key) - const valueHandle = toVmValue(value) - vm.setProp(vm.global, key, valueHandle) - disposeIfNeeded(valueHandle) - } - } - - // Setup variable tracking bridge - const varTrackFnHandle = vm.newFunction( - Identifiers.VariableTrackingFnIdentifier, - (nameHandle, getterHandle) => { - const name = vm.getString(nameHandle) - - if (NO_TRACKING.includes(name)) { - return - } - - // Try to get the value - try { - const valueResult = vm.callFunction(getterHandle, vm.undefined) - if ('error' in valueResult) { - variables[name] = '[[non-primitive]]' - valueResult.error?.dispose() - return - } - const value = vm.dump(valueResult.value) - valueResult.value.dispose() - - // In QuickJS, functions are dumped as their source code strings - // Check if it's a function type or looks like function source - if (typeof value === 'function' || (typeof value === 'string' && value.includes('=>'))) { - variables[name] = '[[non-primitive]]' - } else { - variables[name] = value - } - } catch { - variables[name] = '[[non-primitive]]' - } - } - ) - vm.setProp(vm.global, Identifiers.VariableTrackingFnIdentifier, varTrackFnHandle) - varTrackFnHandle.dispose() - - // Wrap code in async generator - QuickJS pattern with global variables - // This avoids promise resolution issues by storing results in globalThis - const scriptCode = ` -"use strict"; -globalThis.__llmz_result = undefined; -globalThis.__llmz_result_set = false; -globalThis.__llmz_error = null; -globalThis.__llmz_error_stack = null; -globalThis.__llmz_yields = []; - -(async () => { - try { - async function* __fn__() { -${transformed.code} - } - - const fn = __fn__(); - let iteration = 0; - const maxIterations = 10000; // Safety limit - - while (iteration < maxIterations) { - const { value, done } = await fn.next(); - - if (done) { - globalThis.__llmz_result = value; - globalThis.__llmz_result_set = true; - break; - } - - // Store yielded value - globalThis.__llmz_yields.push(value); - - // Call yield handler - await ${Identifiers.AsyncIterYieldFnIdentifier}(value); - - iteration++; - } - - if (iteration >= maxIterations) { - throw new Error('Maximum iterations exceeded'); - } - } catch (err) { - // Store both the error message (which may contain serialized signal data) - // and the stack trace from QuickJS - // If err is a string (as thrown from promise rejection), use it directly - // Otherwise extract the message property - globalThis.__llmz_error = typeof err === 'string' ? err : String(err.message || err || ''); - // Force the stack to be converted to a string in QuickJS context - globalThis.__llmz_error_stack = '' + (err.stack || ''); - } -})(); -`.trim() - - // Helper to copy context back from VM - uses vm.dump which handles deep cloning - const copyBackContextFromVM = () => { - for (const key of trackedProperties) { - if (referenceProperties.has(key)) { - // Skip reference properties (getters/setters) - they update in real-time - continue - } - - try { - // Get the entire value from VM (vm.dump handles deep cloning) - const valueResult = vm.evalCode(`globalThis['${key}']`) - const vmValue = vm.dump(valueResult.unwrap()) - valueResult.unwrap().dispose() - // Update the context with the cloned value - try { - context[key] = vmValue - } catch { - // Ignore read-only property errors - } - } catch { - // Ignore errors when copying back - } - } - } - - // Execute code - result is undefined (async IIFE returns nothing) - const execResult = vm.evalCode(scriptCode, '') - if ('error' in execResult) { - if (execResult.error) { - const err = vm.dump(execResult.error) - execResult.error.dispose() - throw new Error(err?.message || 'Execution failed') - } - throw new Error('Execution failed') - } - execResult.value.dispose() - - // CRITICAL: Execute pending microtasks and host promises in a loop - // QuickJS doesn't have automatic event loop, so we need to manually pump it - const maxIterations = 1000 - let iteration = 0 - - while (iteration < maxIterations) { - // Execute all pending QuickJS microtasks - let hasJobs = false - const maxJobs = 10000 - for (let i = 0; i < maxJobs; i++) { - const pending = runtime.executePendingJobs?.(-1) - const jobCount = pending === undefined ? 0 : pending.unwrap() - if (jobCount <= 0) break - hasJobs = true - } - - // Resolve all pending host promises - const currentPromises = [...pendingPromises] - pendingPromises.length = 0 // Clear the array for new promises - - if (currentPromises.length > 0) { - // Check if aborted before processing promises - if (signal?.aborted) { - // Reject all pending promises with abort error - const reason = (signal as any).reason - const abortMessage = - reason instanceof Error - ? `${reason.name}: ${reason.message}` - : reason - ? String(reason) - : 'Execution was aborted' - for (const { deferredPromise } of currentPromises) { - const errValue = vm.newString(abortMessage) - deferredPromise.reject(errValue) - errValue.dispose() - } - runtime.executePendingJobs() - // Break out of the event loop - we're aborting - break - } - - // Create abort handler that rejects all pending promises immediately - let abortListener: (() => void) | null = null - if (signal) { - abortListener = () => { - const reason = (signal as any).reason - const abortMessage = - reason instanceof Error - ? `${reason.name}: ${reason.message}` - : reason - ? String(reason) - : 'Execution was aborted' - // Reject all pending promises immediately - for (const { deferredPromise } of currentPromises) { - const errValue = vm.newString(abortMessage) - deferredPromise.reject(errValue) - errValue.dispose() - } - runtime.executePendingJobs() - } - signal.addEventListener('abort', abortListener) - } - - try { - await Promise.all( - currentPromises.map(async ({ hostPromise, deferredPromise }) => { - // If abort was triggered, skip resolution - if (signal?.aborted) { - return - } - - try { - const value = await hostPromise - // Double-check abort wasn't triggered during await - if (signal?.aborted) { - return - } - const vmValue = toVmValue(value) - deferredPromise.resolve(vmValue) - disposeIfNeeded(vmValue) - } catch (err: any) { - // If abort was triggered, the abort listener already rejected the promise - if (signal?.aborted) { - return - } - - const serialized = err instanceof Error ? err.message : String(err) - - // Create an Error object in QuickJS with the serialized message - const createErrorResult = vm.evalCode(`new Error(${JSON.stringify(serialized)})`) - if ('error' in createErrorResult) { - // Fallback to string if error creation fails - const errValue = vm.newString(serialized) - deferredPromise.reject(errValue) - errValue.dispose() - } else { - const errorHandle = createErrorResult.value - deferredPromise.reject(errorHandle) - errorHandle.dispose() - } - } - }) - ) - } finally { - // Clean up abort listener - if (signal && abortListener) { - signal.removeEventListener('abort', abortListener) - } - } - - // After resolving promises, execute pending jobs to continue the async IIFE - runtime.executePendingJobs() - - // Check if abort was triggered during promise resolution - if (signal?.aborted) { - break - } - } - - // If no jobs and no promises, we're done - if (!hasJobs && pendingPromises.length === 0) { - break - } - - iteration++ - } - - if (iteration >= maxIterations) { - throw new Error('Maximum event loop iterations exceeded') - } - - // NOW check for errors - AFTER all async work is complete - const errorResult = vm.evalCode('globalThis.__llmz_error') - const errorValue = vm.dump(errorResult.unwrap()) - errorResult.unwrap().dispose() - - // Check if aborted - take precedence over other errors - if (signal?.aborted) { - const reason = (signal as any).reason - if (reason instanceof Error) { - throw reason - } - throw new Error(reason ? String(reason) : 'Execution was aborted') - } - - if (errorValue !== null && errorValue !== '') { - // Copy back context even on error - try { - copyBackContextFromVM() - } catch { - // Ignore errors when copying context back after an error - } - - // Get the QuickJS stack trace as well - const errorStackResult = vm.evalCode('globalThis.__llmz_error_stack') - const errorStack = vm.dump(errorStackResult.unwrap()) || '' - errorStackResult.unwrap().dispose() - - // The error value is a string that may contain serialized signal data - // Deserialize to check if it's a VMSignal - const deserializedError = Signals.maybeDeserializeError(errorValue) - - // If it's a VMSignal, set its stack and throw it - if (deserializedError instanceof VMSignal) { - deserializedError.stack = errorStack - throw deserializedError - } - - // Otherwise create an Error with the serialized message and QuickJS stack - const error = new Error(errorValue) - error.stack = errorStack - throw error - } - - // Copy context back from VM before reading result - copyBackContextFromVM() - - // Get result value from global AFTER promises have settled - const resultSetResult = vm.evalCode('globalThis.__llmz_result_set') - const resultSet = vm.dump(resultSetResult.unwrap()) - resultSetResult.unwrap().dispose() - - let returnValue: any = undefined - if (resultSet) { - const resultResult = vm.evalCode('globalThis.__llmz_result') - returnValue = vm.dump(resultResult.unwrap()) - resultResult.unwrap().dispose() - } - - // Deserialize any signals - returnValue = Signals.maybeDeserializeError(returnValue) - - return { - success: true, - variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), - signal: returnValue instanceof VMSignal ? returnValue : undefined, - lines_executed: Array.from(lines_executed), - return_value: returnValue, - } satisfies VMExecutionResult - } catch (err: any) { - // Check if execution was aborted - if (signal?.aborted) { - // Get abort reason if available - const reason = (signal as any).reason - const abortError = - reason instanceof Error ? reason : new Error(reason ? String(reason) : 'Execution was aborted') - return handleErrorQuickJS( - abortError, - code, - consumer, - traces, - variables, - lines_executed, - userCodeStartLine, - currentToolCall - ) - } - - // Also resolve pending promises on error before disposing - await Promise.all( - pendingPromises.map(async ({ hostPromise, deferredPromise }) => { - try { - const value = await hostPromise - const vmValue = toVmValue(value) - deferredPromise.resolve(vmValue) - disposeIfNeeded(vmValue) - } catch (err2: any) { - const serialized = err2 instanceof Error ? err2.message : String(err2) - const errValue = vm.newString(serialized) - deferredPromise.reject(errValue) - errValue.dispose() - } - }) - ).catch(() => {}) - return handleErrorQuickJS( - err, - code, - consumer, - traces, - variables, - lines_executed, - userCodeStartLine, - currentToolCall - ) - } finally { - try { - vm.dispose() - } catch {} - try { - runtime.dispose() - } catch {} - } - } catch (quickjsError: any) { - // QuickJS failed to load or initialize - fallback to node driver - const debugInfo = { - error: quickjsError?.message || String(quickjsError), - errorStack: quickjsError?.stack, - wasmSource: (BundledReleaseSyncVariant as any)._wasmSource, - wasmLoadedSuccessfully: (BundledReleaseSyncVariant as any)._wasmLoadedSuccessfully, - wasmSize: (BundledReleaseSyncVariant as any)._wasmSize, - wasmLoadError: (BundledReleaseSyncVariant as any)._wasmLoadError, - nodeVersion: typeof process !== 'undefined' && process.version ? process.version : 'undefined', - platform: typeof process !== 'undefined' && process.platform ? process.platform : 'undefined', - } - - console.warn('QuickJS failed to load, falling back to node driver.') - console.warn('Error:', quickjsError?.message || quickjsError) - console.warn('Debug info:', JSON.stringify(debugInfo, null, 2)) - DRIVER = 'node' - } - } - - // ============================================================================ - // Node Driver (No VM) - // ============================================================================ - if (DRIVER === 'node') { - context[Identifiers.CommentFnIdentifier] = (comment: string, line: number) => - traces.push({ type: 'comment', comment, line, started_at: Date.now() }) - - context[Identifiers.LineTrackingFnIdentifier] = (line: number) => { - lines_executed.set(line, (lines_executed.get(line) ?? 0) + 1) - } - - context[Identifiers.JSXFnIdentifier] = (tool: any, props: any, ...children: any[]) => - createJsxComponent({ - type: tool, - props, - children, - }) - - context[Identifiers.VariableTrackingFnIdentifier] = (name: string, getter: () => any) => { - if (NO_TRACKING.includes(name)) { - return - } - variables[name] = () => { - try { - const value = getter() - if (typeof value === 'function') { - return '[[non-primitive]]' - } - return value - } catch { - return '[[non-primitive]]' - } - } - } - - let currentToolCall: SnapshotSignal['toolCall'] | undefined - context[Identifiers.ToolCallTrackerFnIdentifier] = (callId: number, type: 'start' | 'end', outputOrError?: any) => { - const temp = Signals.maybeDeserializeError(outputOrError?.message) - if (type === 'end' && temp instanceof SnapshotSignal && temp?.toolCall) { - currentToolCall = { - ...temp.toolCall, - assignment: transformed.toolCalls.get(callId)?.assignment, - } - } - } - - context[Identifiers.ConsoleObjIdentifier] = { - log: (...args: any[]) => { - const message = args.shift() - traces.push({ type: 'log', message, args, started_at: Date.now() }) - }, - } - - context[Identifiers.AsyncIterYieldFnIdentifier] = async function (value: JsxComponent) { - const startedAt = Date.now() - try { - if (typeof value.type !== 'string' || value.type.trim().length === 0) { - throw new Error('A yield statement must yield a valid tool') - } - - const toolName = Object.keys(context).find((x) => x.toUpperCase() === value.type.toUpperCase()) - - if (!toolName) { - throw new Error(`Yield tool "${value.type}", but tool is not found`) - } - - await context[toolName](value) - } finally { - traces.push({ type: 'yield', value, started_at: startedAt, ended_at: Date.now() }) - } - } - - const AsyncFunction: (...args: unknown[]) => (...args: unknown[]) => AsyncGenerator = - async function* () {}.constructor as any - - return await (async () => { - // We need to track the top-level properties of the context object - const descriptors = Object.getOwnPropertyDescriptors(context) - - const topLevelProperties = Object.keys(descriptors).filter( - (x) => - !NO_TRACKING.includes(x) && - descriptors[x] && - typeof descriptors[x].value !== 'function' && - typeof descriptors[x].value !== 'object' - ) - - const __report = (name: string, value: unknown) => { - if (context[name] !== undefined && context[name] !== value) { - context[name] = value - } - } - - context.__report = __report - - const reportAll = topLevelProperties.map((x) => `__report("${x}", ${x})`).join(';') - - // we're wrapping the function creation inside a promise closure to catch sync errors in the promise catch clause below - const assigner = `let __${Identifiers.LineTrackingFnIdentifier} = ${Identifiers.LineTrackingFnIdentifier}; ${Identifiers.LineTrackingFnIdentifier} = function(line) { ${reportAll}; __${Identifiers.LineTrackingFnIdentifier}(line);}` - const wrapper = `"use strict"; try { ${assigner};${transformed.code} } finally { ${reportAll} };` - - const fn = AsyncFunction(...Object.keys(context), wrapper) - const res = fn(...Object.values(context)) - - do { - const { value, done } = await res.next() - if (done) { - return value - } - await context[Identifiers.AsyncIterYieldFnIdentifier](value) - // oxlint-disable-next-line no-constant-condition - } while (true) - })() - .then((res) => { - res = Signals.maybeDeserializeError(res) - return { - success: true, - variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), - signal: res instanceof VMSignal ? res : undefined, - lines_executed: Array.from(lines_executed), - return_value: res, - } satisfies VMExecutionResult - }) - .catch((err) => handleErrorNode(err, code, consumer, traces, variables, lines_executed, currentToolCall)) - .catch((err) => handleCatch(err, traces, variables, lines_executed)) - } - - throw new Error(`Unknown driver: ${DRIVER}`) -} - -// ============================================================================ -// QuickJS Error Handler -// ============================================================================ -const handleErrorQuickJS = ( - err: Error, - code: string, - _consumer: SourceMapConsumer, - traces: Traces.Trace[], - variables: { [k: string]: any }, - lines_executed: Map, - userCodeStartLine: number, - currentToolCall?: SnapshotSignal['toolCall'] | undefined -): VMExecutionResult => { - err = Signals.maybeDeserializeError(err) - const lines = code.split('\n') - const stackTrace = err.stack || '' - const LINE_OFFSET = 1 - - // Parse QuickJS stack traces: "at doThrow (:16)" or "at (:58)" - const regex = /:(\d+)/g - - // QuickJS line numbers include the wrapper code before transformed.code - // The wrapper has 10 lines before transformed.code starts (counting from scriptCode template) - const QUICKJS_WRAPPER_OFFSET = 10 - - const matches = Array.from(stackTrace.matchAll(regex)).map((x) => { - // Adjust for the wrapper offset to get the line in transformed.code - const quickjsLine = Number(x[1]) - const transformedCodeLine = quickjsLine - QUICKJS_WRAPPER_OFFSET - - // Don't use source map for stack traces - it's unreliable due to Babel's line squashing - // Instead, rely on retainLines keeping line numbers approximately correct - // The transformed code lines should roughly correspond to the marker code lines - // Add +1 because QuickJS reports the line where the IIFE starts, but the actual call is on the next line - const line = Math.max(1, transformedCodeLine - userCodeStartLine + 1) - const actualLine = lines[line - LINE_OFFSET] ?? '' - const whiteSpacesCount = actualLine.length - actualLine.trimStart().length - const minColumn = whiteSpacesCount - - return { - line, - column: minColumn, - } - }) - - // If no matches found in stack trace (e.g., promise rejection from native function), - // use the last executed line from line tracking - if (matches.length === 0 && lines_executed.size > 0) { - const lastLine = Math.max(...Array.from(lines_executed.keys())) - const actualLine = lines[lastLine - LINE_OFFSET] ?? '' - const whiteSpacesCount = actualLine.length - actualLine.trimStart().length - - matches.push({ - line: lastLine, - column: whiteSpacesCount, - }) - } - - const lastLine = maxBy(matches, (x) => x.line ?? 0)?.line ?? 0 - - let debugUserCode = '' - let truncatedCode = '' - let truncated = false - const appendCode = (line: string) => { - debugUserCode += line - if (!truncated) { - truncatedCode += line - } - } - - for (let i = 0; i < lines.length; i++) { - const VM_OFFSET = 2 - const DISPLAY_OFFSET = 0 - const line = lines[i] - const correctedStackLineIndex = i + LINE_OFFSET + VM_OFFSET - const match = matches.find((x) => x.line + VM_OFFSET === correctedStackLineIndex) - const paddedLineNumber = String(correctedStackLineIndex - VM_OFFSET - DISPLAY_OFFSET).padStart(3, '0') - - if (match) { - appendCode(`> ${paddedLineNumber} | ${line}\n`) - appendCode(` ${' '.repeat(paddedLineNumber.length + match.column)}^^^^^^^^^^\n`) - if (match.line >= lastLine) { - truncated = true - } - } else { - appendCode(` ${paddedLineNumber} | ${line}\n`) - } - } - - debugUserCode = cleanStackTrace(debugUserCode).trim() - truncatedCode = cleanStackTrace(truncatedCode).trim() - - if (err instanceof VMSignal) { - // Add properties to VMSignal (these properties exist on the error in the original vm.ts) - const signalError = err as VMSignal & { - stack: string - truncatedCode: string - variables: any - toolCall?: SnapshotSignal['toolCall'] - } - signalError.stack = debugUserCode - signalError.truncatedCode = truncatedCode - signalError.variables = mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)) - signalError.toolCall = currentToolCall - - return { - success: true, - variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), - signal: err, - lines_executed: Array.from(lines_executed), - } - } else { - traces.push({ - type: 'code_execution_exception', - position: [matches[0]?.line ?? 0, matches[0]?.column ?? 0], - message: err.message, - stackTrace: debugUserCode, - started_at: Date.now(), - }) - - // Create error and deserialize to get proper message format - const codeError = new CodeExecutionError(err.message, code, debugUserCode) - const deserializedError = Signals.maybeDeserializeError(codeError) - - return { - success: false, - variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), - error: deserializedError, - traces, - lines_executed: Array.from(lines_executed), - } - } -} - -// ============================================================================ -// Node Error Handler -// ============================================================================ -const handleErrorNode = ( - err: Error, - code: string, - consumer: SourceMapConsumer, - traces: Traces.Trace[], - variables: { [k: string]: () => any }, - _lines_executed: Map, - currentToolCall?: SnapshotSignal['toolCall'] | undefined -) => { - err = Signals.maybeDeserializeError(err) - const lines = code.split('\n') - const stackTrace = err.stack || '' - const LINE_OFFSET = 1 - - const regex = /:(\d+):(\d+)/g - // :13:269) - // ~~~~~~~~~~~~~~~~~~ - - const matches = [...stackTrace.matchAll(regex)].map((x) => { - const originalLine = consumer.originalPositionFor({ - line: Number(x[1]), - column: Number(x[2]), - }) - const line = originalLine.line ?? Number(x[1]) - const actualLine = lines[line - LINE_OFFSET] ?? '' - const whiteSpacesCount = actualLine.length - actualLine.trimStart().length - const minColumn = Math.max(whiteSpacesCount, originalLine.column) - - return { - line, - column: Math.min(minColumn, Number(x[2])), - } - }) - - const lastLine = maxBy(matches, (x) => x.line ?? 0)?.line ?? 0 - - let debugUserCode = '' - let truncatedCode = '' - let truncated = false - const appendCode = (line: string) => { - debugUserCode += line - if (!truncated) { - truncatedCode += line - } - } - - for (let i = 0; i < lines.length; i++) { - const VM_OFFSET = 2 - const DISPLAY_OFFSET = 0 - const line = lines[i] - const correctedStackLineIndex = i + LINE_OFFSET + VM_OFFSET // 1 for the array index starting at 0, then 2 is the number of lines added by the sandbox - const match = matches.find((x) => x.line + VM_OFFSET === correctedStackLineIndex) - // add line number padded with 000 to the lines below - const paddedLineNumber = String(correctedStackLineIndex - VM_OFFSET - DISPLAY_OFFSET).padStart(3, '0') - - if (match) { - appendCode(`> ${paddedLineNumber} | ${line}\n`) - appendCode(` ${' '.repeat(paddedLineNumber.length + match.column)}^^^^^^^^^^\n`) - if (match.line >= lastLine) { - // we want to keep the code up to the location of the top-level error (which is not necessarily the first error we will traverse in the stack trace) - truncated = true - } - } else { - appendCode(` ${paddedLineNumber} | ${line}\n`) - } - } - - debugUserCode = cleanStackTrace(debugUserCode).trim() - truncatedCode = cleanStackTrace(truncatedCode).trim() - - if (err instanceof VMSignal) { - err.stack = debugUserCode - err.truncatedCode = truncatedCode - err.variables = mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)) - err.toolCall = currentToolCall - throw err - } else { - traces.push({ - type: 'code_execution_exception', - position: [matches[0]?.line ?? 0, matches[0]?.column ?? 0], - message: err.message, - stackTrace: debugUserCode, - started_at: Date.now(), - }) - throw new CodeExecutionError(err.message, code, debugUserCode) - } -} - -const handleCatch = ( - err: Error, - traces: Traces.Trace[], - variables: { [k: string]: () => any }, - lines_executed: Map -) => { - err = Signals.maybeDeserializeError(err) - return { - success: err instanceof VMSignal ? true : false, - variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), - error: err, - signal: err instanceof VMSignal ? err : undefined, - traces, - lines_executed: Array.from(lines_executed), - } satisfies VMExecutionResult -} diff --git a/packages/llmz/src/vm/drivers/node.ts b/packages/llmz/src/vm/drivers/node.ts new file mode 100644 index 00000000000..23e377b3880 --- /dev/null +++ b/packages/llmz/src/vm/drivers/node.ts @@ -0,0 +1,73 @@ +import { isFunction, mapValues } from 'lodash-es' + +import { Identifiers } from '../../compiler/index.js' +import { Signals, VMSignal } from '../../errors.js' +import type { JsxComponent } from '../../jsx.js' +import type { VMExecutionResult } from '../../types.js' +import { handleCatch, handleErrorNode } from '../errors.js' +import { instrumentContext, NO_TRACKING } from '../instrument.js' +import type { DriverExecutionContext, VMDriver } from '../types.js' + +// Unsandboxed execution via Node's AsyncGeneratorFunction constructor. +// No isolation — shares the same heap. Used as fallback when QuickJS WASM can't load. +export class NodeDriver implements VMDriver { + public async execute(ctx: DriverExecutionContext): Promise { + const { transformed, consumer, context, traces, code, lines_executed, variables } = ctx + + const state = instrumentContext(context, transformed, traces, variables, lines_executed, consumer, 0) + + // No built-in AsyncGeneratorFunction type in TS — extract the constructor at runtime + type AsyncGeneratorCtor = (...args: unknown[]) => (...args: unknown[]) => AsyncGenerator + const AsyncFunction: AsyncGeneratorCtor = async function* () {}.constructor as AsyncGeneratorCtor + + return await (async () => { + const descriptors = Object.getOwnPropertyDescriptors(context) + + const topLevelProperties = Object.keys(descriptors).filter( + (x) => + !NO_TRACKING.includes(x) && + descriptors[x] && + typeof descriptors[x].value !== 'function' && + typeof descriptors[x].value !== 'object' + ) + + const __report = (name: string, value: unknown) => { + if (context[name] !== undefined && context[name] !== value) { + context[name] = value + } + } + + context.__report = __report + + // Inject __report calls into the line tracker so primitive context values sync back on every line + const reportAll = topLevelProperties.map((x) => `__report("${x}", ${x})`).join(';') + + const assigner = `let __${Identifiers.LineTrackingFnIdentifier} = ${Identifiers.LineTrackingFnIdentifier}; ${Identifiers.LineTrackingFnIdentifier} = function(line) { ${reportAll}; __${Identifiers.LineTrackingFnIdentifier}(line);}` + const wrapper = `"use strict"; try { ${assigner};${transformed.code} } finally { ${reportAll} };` + + const fn = AsyncFunction(...Object.keys(context), wrapper) + const res = fn(...Object.values(context)) + + do { + const { value, done } = await res.next() + if (done) { + return value + } + await context[Identifiers.AsyncIterYieldFnIdentifier](value) + // oxlint-disable-next-line no-constant-condition + } while (true) + })() + .then((res) => { + res = Signals.maybeDeserializeError(res) + return { + success: true, + variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), + signal: res instanceof VMSignal ? res : undefined, + lines_executed: Array.from(lines_executed), + return_value: res, + } satisfies VMExecutionResult + }) + .catch((err) => handleErrorNode(err, code, consumer, traces, variables, lines_executed, state.currentToolCall)) + .catch((err) => handleCatch(err, traces, variables, lines_executed)) + } +} diff --git a/packages/llmz/src/vm/drivers/quickjs.ts b/packages/llmz/src/vm/drivers/quickjs.ts new file mode 100644 index 00000000000..60f155bdd0b --- /dev/null +++ b/packages/llmz/src/vm/drivers/quickjs.ts @@ -0,0 +1,598 @@ +import { isFunction, mapValues } from 'lodash-es' +import { + newQuickJSWASMModuleFromVariant, + QuickJSContext, + type QuickJSHandle, + shouldInterruptAfterDeadline, +} from 'quickjs-emscripten-core' + +import { Identifiers } from '../../compiler/index.js' +import { Signals, VMSignal } from '../../errors.js' +import { BundledReleaseSyncVariant } from '../../quickjs-variant.js' +import type { VMExecutionResult } from '../../types.js' +import { handleErrorQuickJS } from '../errors.js' +import { NO_TRACKING, findUserCodeStartLine, instrumentContext } from '../instrument.js' +import type { DriverExecutionContext, VMContext, VMDriver } from '../types.js' + +// Sandboxed execution via QuickJS WASM. All host values must be manually marshalled +// across the boundary — QuickJS has its own heap, separate from Node.js. +export class QuickJSDriver implements VMDriver { + public async execute(ctx: DriverExecutionContext): Promise { + const { transformed, consumer, context, traces, signal, timeout, code, lines_executed, variables } = ctx + + const userCodeStartLine = findUserCodeStartLine(transformed) + const state = instrumentContext( + context, + transformed, + traces, + variables, + lines_executed, + consumer, + userCodeStartLine + ) + + const QuickJS = await newQuickJSWASMModuleFromVariant(BundledReleaseSyncVariant) + const runtime = QuickJS.newRuntime() + runtime.setMemoryLimit(128 * 1024 * 1024) + + const startTime = Date.now() + const timeoutHandler = shouldInterruptAfterDeadline(startTime + timeout) + runtime.setInterruptHandler(() => { + if (signal?.aborted) { + return true + } + return timeoutHandler(runtime) + }) + + const vm = runtime.newContext() + const trackedProperties = new Set() + const referenceProperties = new Set() + const pendingPromises: Array<{ hostPromise: Promise; deferredPromise: any }> = [] + + // Convert a host JS value into a QuickJS handle (the WASM equivalent) + const toVmValue = (value: any): QuickJSHandle => { + if (typeof value === 'string') { + return vm.newString(value) + } else if (typeof value === 'number') { + return vm.newNumber(value) + } else if (typeof value === 'boolean') { + return value ? vm.true : vm.false + } else if (value === null) { + return vm.null + } else if (value === undefined) { + return vm.undefined + } else if (Array.isArray(value)) { + const items = value.map((item) => { + if (typeof item === 'string') return JSON.stringify(item) + else if (typeof item === 'number' || typeof item === 'boolean') return String(item) + else if (item === null) return 'null' + else if (item === undefined) return 'undefined' + else if (typeof item === 'object') return JSON.stringify(item) + return 'undefined' + }) + const result = vm.evalCode(`[${items.join(',')}]`) + if ('error' in result) { + result.error?.dispose() + return vm.undefined + } + return result.value + } else if (typeof value === 'object') { + const obj = vm.newObject() + for (const [k, v] of Object.entries(value)) { + if (typeof v !== 'function') { + const propHandle = toVmValue(v) + vm.setProp(obj, k, propHandle) + disposeIfNeeded(propHandle) + } + } + return obj + } + return vm.undefined + } + + // Singleton handles (true/false/null/undefined) must not be disposed + const disposeIfNeeded = (handle: QuickJSHandle) => { + if (handle !== vm.true && handle !== vm.false && handle !== vm.null && handle !== vm.undefined) { + ;(handle as any).dispose() + } + } + + // Wrap a host function so QuickJS can call it: unmarshal args, call host, marshal result back. + // Async results become QuickJS deferred promises, resolved by _pumpEventLoop. + const bridgeFunction = (fn: Function, _fnName: string = 'anonymous') => { + return (...argHandles: any[]) => { + const args = argHandles.map((h: any) => vm.dump(h)) + try { + const result = fn(...args) + if (result && typeof result.then === 'function') { + const promise = vm.newPromise() + pendingPromises.push({ hostPromise: result, deferredPromise: promise }) + void promise.settled.then(() => { + if (runtime.alive) { + runtime.executePendingJobs() + } + }) + return promise.handle + } + return toVmValue(result) + } catch (err) { + const serialized = err instanceof Error ? err.message : String(err) + throw new Error(serialized) + } + } + } + + try { + bridgeContextToVM(context, vm, trackedProperties, referenceProperties, toVmValue, disposeIfNeeded, bridgeFunction) + + setupVariableTrackingBridge(vm, variables) + + const scriptCode = buildScriptCode(transformed.code) + + const copyBackContextFromVM = () => { + for (const key of trackedProperties) { + if (referenceProperties.has(key)) continue + try { + const valueResult = vm.evalCode(`globalThis['${key}']`) + const vmValue = vm.dump(valueResult.unwrap()) + valueResult.unwrap().dispose() + try { + context[key] = vmValue + } catch { + // Ignore read-only property errors + } + } catch { + // Ignore errors when copying back + } + } + } + + const execResult = vm.evalCode(scriptCode, '') + if ('error' in execResult) { + if (execResult.error) { + const err = vm.dump(execResult.error) + execResult.error.dispose() + throw new Error(err?.message || 'Execution failed') + } + throw new Error('Execution failed') + } + execResult.value.dispose() + + await this._pumpEventLoop(runtime, vm, pendingPromises, signal, toVmValue, disposeIfNeeded) + + const errorResult = vm.evalCode('globalThis.__llmz_error') + const errorValue = vm.dump(errorResult.unwrap()) + errorResult.unwrap().dispose() + + if (signal?.aborted) { + const reason = (signal as any).reason + if (reason instanceof Error) throw reason + throw new Error(reason ? String(reason) : 'Execution was aborted') + } + + if (errorValue !== null && errorValue !== '') { + try { + copyBackContextFromVM() + } catch {} + + const errorStackResult = vm.evalCode('globalThis.__llmz_error_stack') + const errorStack = vm.dump(errorStackResult.unwrap()) || '' + errorStackResult.unwrap().dispose() + + const deserializedError = Signals.maybeDeserializeError(errorValue) + + if (deserializedError instanceof VMSignal) { + deserializedError.stack = errorStack + throw deserializedError + } + + const error = new Error(errorValue) + error.stack = errorStack + throw error + } + + copyBackContextFromVM() + + const resultSetResult = vm.evalCode('globalThis.__llmz_result_set') + const resultSet = vm.dump(resultSetResult.unwrap()) + resultSetResult.unwrap().dispose() + + let returnValue: any = undefined + if (resultSet) { + const resultResult = vm.evalCode('globalThis.__llmz_result') + returnValue = vm.dump(resultResult.unwrap()) + resultResult.unwrap().dispose() + } + + returnValue = Signals.maybeDeserializeError(returnValue) + + return { + success: true, + variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), + signal: returnValue instanceof VMSignal ? returnValue : undefined, + lines_executed: Array.from(lines_executed), + return_value: returnValue, + } satisfies VMExecutionResult + } catch (err: any) { + if (signal?.aborted) { + const reason = (signal as any).reason + const abortError = + reason instanceof Error ? reason : new Error(reason ? String(reason) : 'Execution was aborted') + return handleErrorQuickJS( + abortError, + code, + consumer, + traces, + variables, + lines_executed, + userCodeStartLine, + ctx.currentToolCall ?? state.currentToolCall + ) + } + + await Promise.all( + pendingPromises.map(async ({ hostPromise, deferredPromise }) => { + try { + const value = await hostPromise + const vmValue = toVmValue(value) + deferredPromise.resolve(vmValue) + disposeIfNeeded(vmValue) + } catch (err2: any) { + const serialized = err2 instanceof Error ? err2.message : String(err2) + const errValue = vm.newString(serialized) + deferredPromise.reject(errValue) + errValue.dispose() + } + }) + ).catch(() => {}) + + return handleErrorQuickJS( + err, + code, + consumer, + traces, + variables, + lines_executed, + userCodeStartLine, + ctx.currentToolCall ?? state.currentToolCall + ) + } finally { + try { + vm.dispose() + } catch {} + try { + runtime.dispose() + } catch {} + } + } + + // QuickJS has no event loop — we manually drain pending microtasks and resolve + // host promises in a loop until all async work completes or the signal aborts. + private async _pumpEventLoop( + runtime: any, + vm: any, + pendingPromises: Array<{ hostPromise: Promise; deferredPromise: any }>, + signal: AbortSignal | null, + toVmValue: (value: any) => QuickJSHandle, + disposeIfNeeded: (handle: QuickJSHandle) => void + ) { + const maxIterations = 1000 + let iteration = 0 + + while (iteration < maxIterations) { + let hasJobs = false + const maxJobs = 10000 + for (let i = 0; i < maxJobs; i++) { + const pending = runtime.executePendingJobs?.(-1) + const jobCount = pending === undefined ? 0 : pending.unwrap() + if (jobCount <= 0) break + hasJobs = true + } + + const currentPromises = [...pendingPromises] + pendingPromises.length = 0 + + if (currentPromises.length > 0) { + if (signal?.aborted) { + const reason = (signal as any).reason + const abortMessage = + reason instanceof Error + ? `${reason.name}: ${reason.message}` + : reason + ? String(reason) + : 'Execution was aborted' + for (const { deferredPromise } of currentPromises) { + const errValue = vm.newString(abortMessage) + deferredPromise.reject(errValue) + errValue.dispose() + } + runtime.executePendingJobs() + break + } + + let abortListener: (() => void) | null = null + if (signal) { + abortListener = () => { + const reason = (signal as any).reason + const abortMessage = + reason instanceof Error + ? `${reason.name}: ${reason.message}` + : reason + ? String(reason) + : 'Execution was aborted' + for (const { deferredPromise } of currentPromises) { + const errValue = vm.newString(abortMessage) + deferredPromise.reject(errValue) + errValue.dispose() + } + runtime.executePendingJobs() + } + signal.addEventListener('abort', abortListener) + } + + try { + await Promise.all( + currentPromises.map(async ({ hostPromise, deferredPromise }) => { + if (signal?.aborted) return + try { + const value = await hostPromise + if (signal?.aborted) return + const vmValue = toVmValue(value) + deferredPromise.resolve(vmValue) + disposeIfNeeded(vmValue) + } catch (err: any) { + if (signal?.aborted) return + const serialized = err instanceof Error ? err.message : String(err) + const createErrorResult = vm.evalCode(`new Error(${JSON.stringify(serialized)})`) + if ('error' in createErrorResult) { + const errValue = vm.newString(serialized) + deferredPromise.reject(errValue) + errValue.dispose() + } else { + const errorHandle = createErrorResult.value + deferredPromise.reject(errorHandle) + errorHandle.dispose() + } + } + }) + ) + } finally { + if (signal && abortListener) { + signal.removeEventListener('abort', abortListener) + } + } + + runtime.executePendingJobs() + + if (signal?.aborted) break + } + + if (!hasJobs && pendingPromises.length === 0) break + iteration++ + } + + if (iteration >= maxIterations) { + throw new Error('Maximum event loop iterations exceeded') + } + } +} + +// Marshal all context entries (functions, objects, arrays, primitives, getter/setters) +// onto QuickJS globalThis so generated code can access them. +function bridgeContextToVM( + // TODO: rename these and their associated concepts, these types ain't making sense + context: VMContext, + vm: QuickJSContext, + trackedProperties: Set, + referenceProperties: Set, + toVmValue: (value: any) => QuickJSHandle, + disposeIfNeeded: (handle: QuickJSHandle) => void, + bridgeFunction: (fn: Function, name?: string) => (...args: any[]) => any +) { + for (const [key, value] of Object.entries(context)) { + const descriptor = Object.getOwnPropertyDescriptor(context, key) + + if (descriptor && (descriptor.get || descriptor.set)) { + referenceProperties.add(key) + trackedProperties.add(key) + bridgeGetterSetter(vm, key, undefined, descriptor, context, toVmValue) + continue + } + + if (typeof value === 'function') { + const fnHandle = vm.newFunction(key, bridgeFunction(value, key)) + vm.setProp(vm.global, key, fnHandle) + fnHandle.dispose() + } else if (Array.isArray(value)) { + trackedProperties.add(key) + const arrayHandle = toVmValue(value) + vm.setProp(vm.global, key, arrayHandle) + disposeIfNeeded(arrayHandle) + } else if (typeof value === 'object' && value !== null) { + trackedProperties.add(key) + const objHandle = vm.newObject() + const props = new Set([...Object.keys(value), ...Object.getOwnPropertyNames(value)]) + const getterSetterProps: Array<{ prop: string; descriptor: PropertyDescriptor }> = [] + + for (const prop of props) { + const propDescriptor = Object.getOwnPropertyDescriptor(value, prop) + if (propDescriptor && (propDescriptor.get || propDescriptor.set)) { + referenceProperties.add(`${key}.${prop}`) + getterSetterProps.push({ prop, descriptor: propDescriptor }) + } else if (typeof (value as any)[prop] === 'function') { + const propFnHandle = vm.newFunction(prop, bridgeFunction((value as any)[prop], `${key}.${prop}`)) + vm.setProp(objHandle, prop, propFnHandle) + propFnHandle.dispose() + } else { + const propHandle = toVmValue((value as any)[prop]) + vm.setProp(objHandle, prop, propHandle) + disposeIfNeeded(propHandle) + } + } + + vm.setProp(vm.global, key, objHandle) + objHandle.dispose() + + for (const { prop, descriptor } of getterSetterProps) { + bridgeGetterSetter(vm, key, prop, descriptor, context, toVmValue) + } + + if (Object.isSealed(value)) { + const sealResult = vm.evalCode(`Object.seal(globalThis['${key}']);`) + if ('error' in sealResult) sealResult.error?.dispose() + else sealResult.value.dispose() + } + if (!Object.isExtensible(value)) { + const preventResult = vm.evalCode(`Object.preventExtensions(globalThis['${key}']);`) + if ('error' in preventResult) preventResult.error?.dispose() + else preventResult.value.dispose() + } + } else { + trackedProperties.add(key) + const valueHandle = toVmValue(value) + vm.setProp(vm.global, key, valueHandle) + disposeIfNeeded(valueHandle) + } + } +} + +// Bridge a getter/setter property across the host-QuickJS boundary using Object.defineProperty +function bridgeGetterSetter( + vm: QuickJSContext, + key: string, + prop: string | undefined, + descriptor: PropertyDescriptor, + context: VMContext, + toVmValue: (value: any) => QuickJSHandle +) { + const target = prop ? `${key}` : 'globalThis' + const propName = prop ?? key + const prefix = prop ? `${key}_${prop}` : key + + let getterCode = 'undefined' + if (descriptor.get) { + const getterBridge = vm.newFunction(`get_${propName}`, () => { + try { + const hostValue = prop ? context[key][prop] : context[key] + return toVmValue(hostValue) + } catch (err: any) { + throw new Error(err instanceof Error ? err.message : String(err)) + } + }) + const getterName = `__getter_${prefix}__` + vm.setProp(vm.global, getterName, getterBridge) + getterBridge.dispose() + getterCode = getterName + } + + let setterCode = 'undefined' + if (descriptor.set) { + const setterBridge = vm.newFunction(`set_${propName}`, (valueHandle: any) => { + try { + const jsValue = vm.dump(valueHandle) + if (prop) { + context[key][prop] = jsValue + } else { + context[key] = jsValue + } + return vm.undefined + } catch (err: any) { + throw new Error(err instanceof Error ? err.message : String(err)) + } + }) + const setterName = `__setter_${prefix}__` + vm.setProp(vm.global, setterName, setterBridge) + setterBridge.dispose() + setterCode = setterName + } + + const definePropertyCode = ` + Object.defineProperty(${target}, '${propName}', { + enumerable: true, + configurable: ${descriptor.configurable !== false}, + get: ${getterCode}, + set: ${setterCode} + }); + ` + const result = vm.evalCode(definePropertyCode) + if ('error' in result) result.error?.dispose() + else result.value.dispose() +} + +// QuickJS-specific variable tracking: uses vm.callFunction to invoke getter handles inside the VM +function setupVariableTrackingBridge(vm: QuickJSContext, variables: Record) { + const varTrackFnHandle = vm.newFunction( + Identifiers.VariableTrackingFnIdentifier, + (nameHandle: any, getterHandle: any) => { + const name = vm.getString(nameHandle) + if (NO_TRACKING.includes(name)) return + + try { + const valueResult = vm.callFunction(getterHandle, vm.undefined) + if ('error' in valueResult) { + variables[name] = '[[non-primitive]]' + valueResult.error?.dispose() + return + } + const value = vm.dump(valueResult.value) + valueResult.value.dispose() + + if (typeof value === 'function' || (typeof value === 'string' && value.includes('=>'))) { + variables[name] = '[[non-primitive]]' + } else { + variables[name] = value + } + } catch { + variables[name] = '[[non-primitive]]' + } + } + ) + vm.setProp(vm.global, Identifiers.VariableTrackingFnIdentifier, varTrackFnHandle) + varTrackFnHandle.dispose() +} + +// Wraps transformed code in an async generator IIFE that stores results/errors on globalThis. +// QuickJS can't return values from async code directly, so we read them back after the event loop. +function buildScriptCode(transformedCode: string): string { + return ` +"use strict"; +globalThis.__llmz_result = undefined; +globalThis.__llmz_result_set = false; +globalThis.__llmz_error = null; +globalThis.__llmz_error_stack = null; +globalThis.__llmz_yields = []; + +(async () => { + try { + async function* __fn__() { +${transformedCode} + } + + const fn = __fn__(); + let iteration = 0; + const maxIterations = 10000; + + while (iteration < maxIterations) { + const { value, done } = await fn.next(); + + if (done) { + globalThis.__llmz_result = value; + globalThis.__llmz_result_set = true; + break; + } + + globalThis.__llmz_yields.push(value); + await ${Identifiers.AsyncIterYieldFnIdentifier}(value); + iteration++; + } + + if (iteration >= maxIterations) { + throw new Error('Maximum iterations exceeded'); + } + } catch (err) { + globalThis.__llmz_error = typeof err === 'string' ? err : String(err.message || err || ''); + globalThis.__llmz_error_stack = '' + (err.stack || ''); + } +})(); +`.trim() +} diff --git a/packages/llmz/src/vm/errors.ts b/packages/llmz/src/vm/errors.ts new file mode 100644 index 00000000000..307939f3424 --- /dev/null +++ b/packages/llmz/src/vm/errors.ts @@ -0,0 +1,208 @@ +import { isFunction, mapValues, maxBy } from 'lodash-es' +import type { SourceMapConsumer } from 'source-map-js' + +import { CodeExecutionError, Signals, SnapshotSignal, VMSignal } from '../errors.js' +import { cleanStackTrace } from '../stack-traces.js' +import type { Traces, VMExecutionResult } from '../types.js' + +// Parse QuickJS stack traces (":16") and map line numbers back through +// the wrapper offset and user code start line to produce annotated source output. +export const handleErrorQuickJS = ( + err: Error, + code: string, + _consumer: SourceMapConsumer, + traces: Traces.Trace[], + variables: { [k: string]: any }, + lines_executed: Map, + userCodeStartLine: number, + currentToolCall?: SnapshotSignal['toolCall'] | undefined +): VMExecutionResult => { + err = Signals.maybeDeserializeError(err) + const lines = code.split('\n') + const stackTrace = err.stack || '' + const LINE_OFFSET = 1 + + const regex = /:(\d+)/g + const QUICKJS_WRAPPER_OFFSET = 10 + + const matches = Array.from(stackTrace.matchAll(regex)).map((x) => { + const quickjsLine = Number(x[1]) + const transformedCodeLine = quickjsLine - QUICKJS_WRAPPER_OFFSET + const line = Math.max(1, transformedCodeLine - userCodeStartLine + 1) + const actualLine = lines[line - LINE_OFFSET] ?? '' + const whiteSpacesCount = actualLine.length - actualLine.trimStart().length + return { line, column: whiteSpacesCount } + }) + + if (matches.length === 0 && lines_executed.size > 0) { + const lastLine = Math.max(...Array.from(lines_executed.keys())) + const actualLine = lines[lastLine - LINE_OFFSET] ?? '' + const whiteSpacesCount = actualLine.length - actualLine.trimStart().length + matches.push({ line: lastLine, column: whiteSpacesCount }) + } + + return formatError(err, lines, matches, traces, variables, lines_executed, currentToolCall, code) +} + +// Parse Node VM stack traces (":13:269") and use source maps to map +// back to original line/column positions in the LLM-generated code. +export const handleErrorNode = ( + err: Error, + code: string, + consumer: SourceMapConsumer, + traces: Traces.Trace[], + variables: { [k: string]: () => any }, + _lines_executed: Map, + currentToolCall?: SnapshotSignal['toolCall'] | undefined +) => { + err = Signals.maybeDeserializeError(err) + const lines = code.split('\n') + const stackTrace = err.stack || '' + const LINE_OFFSET = 1 + + const regex = /:(\d+):(\d+)/g + + const matches = [...stackTrace.matchAll(regex)].map((x) => { + const originalLine = consumer.originalPositionFor({ + line: Number(x[1]), + column: Number(x[2]), + }) + const line = originalLine.line ?? Number(x[1]) + const actualLine = lines[line - LINE_OFFSET] ?? '' + const whiteSpacesCount = actualLine.length - actualLine.trimStart().length + const minColumn = Math.max(whiteSpacesCount, originalLine.column) + return { line, column: Math.min(minColumn, Number(x[2])) } + }) + + const { debugUserCode, truncatedCode } = buildDebugCode(lines, matches) + + if (err instanceof VMSignal) { + err.stack = debugUserCode + err.truncatedCode = truncatedCode + err.variables = mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)) + err.toolCall = currentToolCall + throw err + } else { + traces.push({ + type: 'code_execution_exception', + position: [matches[0]?.line ?? 0, matches[0]?.column ?? 0], + message: err.message, + stackTrace: debugUserCode, + started_at: Date.now(), + }) + throw new CodeExecutionError(err.message, code, debugUserCode) + } +} + +// Last-resort catch: wraps any error (including VMSignals) into a VMExecutionResult +export const handleCatch = ( + err: Error, + traces: Traces.Trace[], + variables: { [k: string]: () => any }, + lines_executed: Map +) => { + err = Signals.maybeDeserializeError(err) + return { + success: err instanceof VMSignal ? true : false, + variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), + error: err, + signal: err instanceof VMSignal ? err : undefined, + traces, + lines_executed: Array.from(lines_executed), + } satisfies VMExecutionResult +} + +// Build annotated source listing with "> 003 | code" markers and "^^^^^^^^^^" carets +// at error positions. truncatedCode stops after the deepest error line. +function buildDebugCode(lines: string[], matches: Array<{ line: number; column: number }>) { + const LINE_OFFSET = 1 + const lastLine = maxBy(matches, (x) => x.line ?? 0)?.line ?? 0 + + let debugUserCode = '' + let truncatedCode = '' + let truncated = false + const appendCode = (line: string) => { + debugUserCode += line + if (!truncated) { + truncatedCode += line + } + } + + for (let i = 0; i < lines.length; i++) { + const VM_OFFSET = 2 + const DISPLAY_OFFSET = 0 + const line = lines[i] + const correctedStackLineIndex = i + LINE_OFFSET + VM_OFFSET + const match = matches.find((x) => x.line + VM_OFFSET === correctedStackLineIndex) + const paddedLineNumber = String(correctedStackLineIndex - VM_OFFSET - DISPLAY_OFFSET).padStart(3, '0') + + if (match) { + appendCode(`> ${paddedLineNumber} | ${line}\n`) + appendCode(` ${' '.repeat(paddedLineNumber.length + match.column)}^^^^^^^^^^\n`) + if (match.line >= lastLine) { + truncated = true + } + } else { + appendCode(` ${paddedLineNumber} | ${line}\n`) + } + } + + return { + debugUserCode: cleanStackTrace(debugUserCode).trim(), + truncatedCode: cleanStackTrace(truncatedCode).trim(), + } +} + +// Shared formatter: annotates the error with debug code, then returns a VMExecutionResult. +// VMSignals are treated as successful (agent control flow), other errors as failures. +function formatError( + err: Error, + lines: string[], + matches: Array<{ line: number; column: number }>, + traces: Traces.Trace[], + variables: { [k: string]: any }, + lines_executed: Map, + currentToolCall: SnapshotSignal['toolCall'] | undefined, + code: string +): VMExecutionResult { + const { debugUserCode, truncatedCode } = buildDebugCode(lines, matches) + + if (err instanceof VMSignal) { + const signalError = err as VMSignal & { + stack: string + truncatedCode: string + variables: any + toolCall?: SnapshotSignal['toolCall'] + } + signalError.stack = debugUserCode + signalError.truncatedCode = truncatedCode + signalError.variables = mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)) + signalError.toolCall = currentToolCall + + return { + success: true, + variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), + signal: err, + lines_executed: Array.from(lines_executed), + } + } else { + traces.push({ + type: 'code_execution_exception', + position: [matches[0]?.line ?? 0, matches[0]?.column ?? 0], + message: err.message, + stackTrace: debugUserCode, + started_at: Date.now(), + }) + + const codeError = new CodeExecutionError(err.message, code, debugUserCode) + const deserializedError = Signals.maybeDeserializeError(codeError) + + return { + success: false, + variables: mapValues(variables, (getter) => (isFunction(getter) ? getter() : getter)), + error: deserializedError, + traces, + lines_executed: Array.from(lines_executed), + } + } +} diff --git a/packages/llmz/src/vm/index.ts b/packages/llmz/src/vm/index.ts new file mode 100644 index 00000000000..3e9c17af3c6 --- /dev/null +++ b/packages/llmz/src/vm/index.ts @@ -0,0 +1,107 @@ +import { SourceMapConsumer } from 'source-map-js' + +import { compile } from '../compiler/index.js' +import { InvalidCodeError } from '../errors.js' +import { BundledReleaseSyncVariant } from '../quickjs-variant.js' +import type { Trace, VMExecutionResult } from '../types.js' +import { NodeDriver } from './drivers/node.js' +import { QuickJSDriver } from './drivers/quickjs.js' +import type { VMContext, VMDriver } from './types.js' + +const MAX_VM_EXECUTION_TIME = 60_000 + +export async function runAsyncFunction( + context: VMContext, + code: string, + traces: Trace[] = [], + signal: AbortSignal | null = null, + timeout: number = MAX_VM_EXECUTION_TIME +): Promise { + const transformed = (() => { + try { + return compile(code) + } catch (err: any) { + traces.push({ + type: 'invalid_code_exception', + message: err?.message ?? 'Unknown error', + code, + started_at: Date.now(), + }) + throw new InvalidCodeError(err.message, code) + } + })() + + const lines_executed = new Map() + const variables: Record = {} + + // TODO: transformed.map (the result of compile above) needs typing, + // once that's done, we can remove the null assertions here + const consumer = new SourceMapConsumer({ + version: transformed.map!.version.toString(), + mappings: transformed.map!.mappings, + names: transformed.map!.names!, + sources: [transformed.map!.file!], + sourcesContent: [transformed.code!], + file: transformed.map!.file!, + sourceRoot: transformed.map!.sourceRoot!, + }) + + context ??= {} + + // Remove variables that the compiler will track — avoids stale values in the context + for (const name of Array.from(transformed.variables)) { + delete context[name] + } + + let driver: VMDriver + + const useQuickJS = typeof process === 'undefined' || process?.env?.USE_QUICKJS !== 'false' + + if (useQuickJS) { + try { + driver = new QuickJSDriver() + return await driver.execute({ + transformed, + consumer, + context, + traces, + signal, + timeout, + code, + lines_executed, + variables, + currentToolCall: undefined, + }) + } catch (quickjsError: any) { + // QuickJS WASM failed to load — fall back to unsandboxed Node driver + const debugInfo = { + error: quickjsError?.message || String(quickjsError), + errorStack: quickjsError?.stack, + wasmSource: BundledReleaseSyncVariant._wasmSource, + wasmLoadedSuccessfully: BundledReleaseSyncVariant._wasmLoadedSuccessfully, + wasmSize: BundledReleaseSyncVariant._wasmSize, + wasmLoadError: BundledReleaseSyncVariant._wasmLoadError, + nodeVersion: typeof process !== 'undefined' && process.version ? process.version : 'undefined', + platform: typeof process !== 'undefined' && process.platform ? process.platform : 'undefined', + } + + console.warn('QuickJS failed to load, falling back to node driver.') + console.warn('Error:', quickjsError?.message || quickjsError) + console.warn('Debug info:', JSON.stringify(debugInfo, null, 2)) + } + } + + driver = new NodeDriver() + return await driver.execute({ + transformed, + consumer, + context, + traces, + signal, + timeout, + code, + lines_executed, + variables, + currentToolCall: undefined, + }) +} diff --git a/packages/llmz/src/vm/instrument.ts b/packages/llmz/src/vm/instrument.ts new file mode 100644 index 00000000000..76f41aec807 --- /dev/null +++ b/packages/llmz/src/vm/instrument.ts @@ -0,0 +1,123 @@ +import type { SourceMapConsumer } from 'source-map-js' + +import { type CompiledCode, Identifiers } from '../compiler/index.js' +import { USER_CODE_START_MARKER } from '../compiler/plugins/async-iterator.js' + +const USER_CODE_MARKER_TAG_START = '__LLMZ_USER_CODE_START__' +const USER_CODE_MARKER_TAG_END = '__LLMZ_USER_CODE_END__' +import { Signals, SnapshotSignal } from '../errors.js' +import { createJsxComponent, JsxComponent } from '../jsx.js' +import type { Trace } from '../types.js' +import { VMContext } from './types.js' + +// Internal identifiers injected by the compiler — excluded from variable tracking +export const NO_TRACKING = [ + Identifiers.CommentFnIdentifier, + Identifiers.ToolCallTrackerFnIdentifier, + Identifiers.ToolTrackerRetIdentifier, + Identifiers.VariableTrackingFnIdentifier, + Identifiers.JSXFnIdentifier, + Identifiers.ConsoleObjIdentifier, +] as const + +export type InstrumentationState = { + currentToolCall: SnapshotSignal['toolCall'] | undefined +} + +// Injects tracking functions (comments, lines, variables, tools, console, yield) into the context. +// Shared by both QuickJS and Node drivers. +export function instrumentContext( + context: VMContext, + transformed: CompiledCode, + traces: Trace[], + variables: Record, + lines_executed: Map, + consumer: SourceMapConsumer, + userCodeStartLine: number +): InstrumentationState { + const state: InstrumentationState = { currentToolCall: undefined } + + context[Identifiers.CommentFnIdentifier] = (comment: string, line: number) => { + if (comment.includes(USER_CODE_MARKER_TAG_START) || comment.includes(USER_CODE_MARKER_TAG_END)) { + return + } + traces.push({ type: 'comment', comment, line, started_at: Date.now() }) + } + + context[Identifiers.LineTrackingFnIdentifier] = (line: number) => { + const originalLine = consumer.originalPositionFor({ line, column: 0 }) + const mappedLine = originalLine.line ?? line + const userCodeLine = Math.max(1, mappedLine - userCodeStartLine) + lines_executed.set(userCodeLine, (lines_executed.get(userCodeLine) ?? 0) + 1) + } + + context[Identifiers.JSXFnIdentifier] = (tool: string, props: Object, ...children: any[]) => + createJsxComponent({ type: tool, props, children }) + + context[Identifiers.VariableTrackingFnIdentifier] = (name: string, getter: () => any) => { + if (NO_TRACKING.includes(name)) { + return + } + variables[name] = () => { + try { + const value = getter() + if (typeof value === 'function') { + return '[[non-primitive]]' + } + return value + } catch { + return '[[non-primitive]]' + } + } + } + + context[Identifiers.ToolCallTrackerFnIdentifier] = (callId: number, type: 'start' | 'end', outputOrError?: Error) => { + const temp = Signals.maybeDeserializeError(outputOrError?.message) + if (type === 'end' && temp instanceof SnapshotSignal && temp?.toolCall) { + state.currentToolCall = { + ...temp.toolCall, + assignment: transformed.toolCalls.get(callId)?.assignment, + } + } + } + + context[Identifiers.ConsoleObjIdentifier] = { + log: (...args: any[]) => { + const message = args.shift() + traces.push({ type: 'log', message, args, started_at: Date.now() }) + }, + } + + context[Identifiers.AsyncIterYieldFnIdentifier] = async function (value: JsxComponent) { + const startedAt = Date.now() + try { + if (typeof value.type !== 'string' || value.type.trim().length === 0) { + throw new Error('A yield statement must yield a valid tool') + } + + const toolName = Object.keys(context).find((x) => x.toUpperCase() === value.type.toUpperCase()) + + if (!toolName) { + throw new Error(`Yield tool "${value.type}", but tool is not found`) + } + + await context[toolName](value) + } finally { + traces.push({ type: 'yield', value, started_at: startedAt, ended_at: Date.now() }) + } + } + + return state +} + +// Locates the __LLMZ_USER_CODE_START__ marker to calculate line offsets for stack traces +export function findUserCodeStartLine(transformed: CompiledCode): number { + const codeWithMarkers = transformed.codeWithMarkers || transformed.code + const markerLines = codeWithMarkers.split('\n') + for (let i = 0; i < markerLines.length; i++) { + if (markerLines[i]?.includes(USER_CODE_START_MARKER)) { + return i + 1 + } + } + return -1 +} diff --git a/packages/llmz/src/vm/types.ts b/packages/llmz/src/vm/types.ts new file mode 100644 index 00000000000..0774036f5a0 --- /dev/null +++ b/packages/llmz/src/vm/types.ts @@ -0,0 +1,25 @@ +import type { SourceMapConsumer } from 'source-map-js' + +import type { CompiledCode } from '../compiler/index.js' +import type { SnapshotSignal } from '../errors.js' +import type { Trace, VMExecutionResult } from '../types.js' + +export type VMContext = Record + +export type DriverExecutionContext = { + transformed: CompiledCode + consumer: SourceMapConsumer + context: VMContext + traces: Trace[] + signal: AbortSignal | null + timeout: number + code: string + lines_executed: Map + variables: Record + currentToolCall: SnapshotSignal['toolCall'] | undefined +} + +// Any execution driver (QuickJS, Node, future drivers) must implement this type +export type VMDriver = { + execute(ctx: DriverExecutionContext): Promise +} diff --git a/packages/llmz/src/vm-jsx.test.ts b/packages/llmz/src/vm/vm-jsx.test.ts similarity index 99% rename from packages/llmz/src/vm-jsx.test.ts rename to packages/llmz/src/vm/vm-jsx.test.ts index 77f1b42ee64..30fe6089711 100644 --- a/packages/llmz/src/vm-jsx.test.ts +++ b/packages/llmz/src/vm/vm-jsx.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it, vi } from 'vitest' -import { runAsyncFunction } from './vm.js' +import { runAsyncFunction } from './index.js' const user = { name: 'John', age: 30, email: 'john@test.com' } diff --git a/packages/llmz/src/vm.test.ts b/packages/llmz/src/vm/vm.test.ts similarity index 99% rename from packages/llmz/src/vm.test.ts rename to packages/llmz/src/vm/vm.test.ts index d6ba92a1834..8f8af8b787d 100644 --- a/packages/llmz/src/vm.test.ts +++ b/packages/llmz/src/vm/vm.test.ts @@ -1,8 +1,8 @@ import { assert, describe, expect, it, vi } from 'vitest' -import { CodeExecutionError, InvalidCodeError, VMSignal } from './errors.js' -import { Trace, Traces } from './types.js' -import { runAsyncFunction } from './vm.js' +import { CodeExecutionError, InvalidCodeError, VMSignal } from '../errors.js' +import { Trace, Traces } from '../types.js' +import { runAsyncFunction } from './index.js' describe('llmz/vm', () => { it('stack traces points to original source map code', async () => { diff --git a/packages/zai/package.json b/packages/zai/package.json index 6fdc8e178f8..a11845ab973 100644 --- a/packages/zai/package.json +++ b/packages/zai/package.json @@ -1,7 +1,7 @@ { "name": "@botpress/zai", "description": "Zui AI (zai) – An LLM utility library written on top of Zui and the Botpress API", - "version": "2.6.23", + "version": "2.6.24", "main": "./dist/index.js", "types": "./dist/index.d.ts", "exports": { diff --git a/packages/zai/tsdown.config.ts b/packages/zai/tsdown.config.ts index 95bcfb0a5c1..1067f7a4a25 100644 --- a/packages/zai/tsdown.config.ts +++ b/packages/zai/tsdown.config.ts @@ -6,7 +6,9 @@ export default defineConfig({ outDir: 'dist', platform: 'neutral', clean: true, - unbundle: true, + // Keep the declaration entrypoint bundled so operation module augmentations + // like zai.extract/check/summarize are visible from @botpress/zai. + unbundle: false, format: 'cjs', target: undefined, })