openclaw/src/agents/code-mode.worker.ts

/**
 * QuickJS worker for Code Mode guest execution and suspended VM snapshots.
 */
import { parentPort, workerData } from "node:worker_threads";
import { isRecord } from "@openclaw/normalization-core/record-coerce";
import { EvalFlags, JSException, QuickJS, type JSValueHandle } from "quickjs-wasi";
import { CODE_MODE_CONTROLLER_SOURCE } from "./code-mode-controller-source.js";
import { toCodeModeJsonSafe as toJsonSafe } from "./code-mode-json.js";
import type { CodeModeApiVirtualFile } from "./code-mode-namespaces.js";
import type {
  CodeModeConfig,
  CodeModeNamespaceDescriptor,
  CodeModeWorkerPayload,
  CodeModeWorkerThreadResult as CodeModeWorkerResult,
  PendingBridgeRequest,
  SettledBridgeRequest,
} from "./code-mode-worker-types.js";
class CodeModeWorkerFailure extends Error {
  readonly code: Extract<CodeModeWorkerResult, { status: "failed" }>["code"];

  constructor(
    code: Extract<CodeModeWorkerResult, { status: "failed" }>["code"],
    message: string,
    options?: ErrorOptions,
  ) {
    super(message, options);
    this.name = "CodeModeWorkerFailure";
    this.code = code;
  }
}

class CodeModeWorkerFailureWithOutput extends CodeModeWorkerFailure {
  readonly output: unknown[];

  constructor(
    code: Extract<CodeModeWorkerResult, { status: "failed" }>["code"],
    message: string,
    output: unknown[],
    options?: ErrorOptions,
  ) {
    super(code, message, options);
    this.name = "CodeModeWorkerFailureWithOutput";
    this.output = output;
  }
}

function isQuickJsInterruptedError(error: unknown): boolean {
  return error instanceof JSException && error.message === "interrupted";
}

type VmRun = {
  vm: QuickJS;
  didTimeout: () => boolean;
};

// QuickJS error stacks are backtrace frames only ("    at file:line:col"), with
// no leading "Name: message" header like V8. Returning .stack alone therefore
// dropped the actual cause, surfacing failures to the model as a bare location
// (e.g. "at openclaw-code-mode:user.js:2:37"). Lead with name+message so the
// model can self-correct, and keep the frames for location.
function formatQuickJsError(name: string, message: string, stack: string | undefined): string {
  const header = message ? `${name}: ${message}` : name;
  if (!stack || stack.split(/\r?\n/, 1)[0] === header) {
    return header;
  }
  return `${header}\n${stack}`;
}

function errorMessage(error: unknown): string {
  if (error instanceof JSException) {
    return formatQuickJsError(error.name, error.message, error.stack);
  }
  if (error instanceof Error) {
    return error.message || String(error);
  }
  return String(error);
}

function buildUserSource(code: string): string {
  return `globalThis.__openclawResult = (async () => {\n${code}\n})()`;
}

function createHostRequestHandler(params: {
  vm: QuickJS;
  pendingRequests: PendingBridgeRequest[];
  config: CodeModeConfig;
}): (
  this: JSValueHandle,
  method: JSValueHandle,
  argsJson: JSValueHandle,
  bridgeId?: JSValueHandle,
) => JSValueHandle {
  return (methodHandle, argsHandle, bridgeIdHandle) => {
    if (params.pendingRequests.length >= params.config.maxPendingToolCalls) {
      throw new Error("too many pending code mode tool calls");
    }
    const method = methodHandle.toString();
    if (
      method !== "search" &&
      method !== "describe" &&
      method !== "call" &&
      method !== "callValue" &&
      method !== "nodes" &&
      method !== "yield" &&
      method !== "namespace" &&
      method !== "agentSpawn" &&
      method !== "agentWait" &&
      method !== "skillsList" &&
      method !== "skillsRead" &&
      method !== "swarmNote"
    ) {
      throw new Error("unsupported code mode bridge method");
    }
    let args: unknown;
    try {
      args = JSON.parse(argsHandle.toString()) as unknown;
    } catch {
      args = [];
    }
    // Snapshotted method counters keep launch identity independent of unrelated bridge traffic.
    // Snapshots are process-local, so every resumable guest comes from the ID-aware source above.
    const id = bridgeIdHandle?.toString();
    if (!id?.startsWith(`bridge:${method}:`) || !/^bridge:[A-Za-z]+:[1-9]\d*$/u.test(id)) {
      throw new Error("invalid code mode bridge id");
    }
    if (params.pendingRequests.some((request) => request.id === id)) {
      throw new Error("duplicate code mode bridge id");
    }
    // The guest receives only an opaque id. Host-side tool execution and policy
    // happen after the worker returns a waiting snapshot.
    params.pendingRequests.push({
      id,
      method,
      args: Array.isArray(args) ? args : [],
    });
    return params.vm.newString(id);
  };
}

async function createVm(params: {
  wasmModule: WebAssembly.Module;
  catalog: unknown[];
  apiFiles: CodeModeApiVirtualFile[];
  namespaces: CodeModeNamespaceDescriptor[];
  swarmEnabled: boolean;
  config: CodeModeConfig;
  pendingRequests: PendingBridgeRequest[];
}): Promise<VmRun> {
  const startedAt = Date.now();
  let timedOut = false;
  const deadlineReached = () => Date.now() - startedAt >= params.config.timeoutMs;
  const vm = await QuickJS.create({
    wasm: params.wasmModule,
    memoryLimit: params.config.memoryLimitBytes,
    timezoneOffset: 0,
    interruptHandler: () => {
      timedOut = deadlineReached();
      return timedOut;
    },
  });
  vm.hostToHandle(params.catalog).consume((handle) =>
    vm.global.setProp("__openclawCatalog", handle),
  );
  vm.hostToHandle(params.namespaces).consume((handle) =>
    vm.global.setProp("__openclawNamespaces", handle),
  );
  vm.hostToHandle(params.apiFiles).consume((handle) =>
    vm.global.setProp("__openclawApiFiles", handle),
  );
  vm.hostToHandle(params.swarmEnabled).consume((handle) =>
    vm.global.setProp("__openclawSwarmEnabled", handle),
  );
  vm.newFunction(
    "__openclawHostRequest",
    createHostRequestHandler({
      vm,
      pendingRequests: params.pendingRequests,
      config: params.config,
    }),
  ).consume((hostRequest) => vm.global.setProp("__openclawHostRequest", hostRequest));
  vm.evalCode(CODE_MODE_CONTROLLER_SOURCE, "openclaw-code-mode:controller.js").dispose();
  return { vm, didTimeout: () => timedOut || deadlineReached() };
}

async function restoreVm(params: {
  wasmModule: WebAssembly.Module;
  snapshotBytes: Uint8Array;
  config: CodeModeConfig;
  pendingRequests: PendingBridgeRequest[];
}): Promise<VmRun> {
  const startedAt = Date.now();
  let timedOut = false;
  const deadlineReached = () => Date.now() - startedAt >= params.config.timeoutMs;
  const snapshot = QuickJS.deserializeSnapshot(params.snapshotBytes);
  const vm = await QuickJS.restore(snapshot, {
    wasm: params.wasmModule,
    memoryLimit: params.config.memoryLimitBytes,
    timezoneOffset: 0,
    interruptHandler: () => {
      timedOut = deadlineReached();
      return timedOut;
    },
  });
  vm.registerHostCallback(
    "__openclawHostRequest",
    createHostRequestHandler({
      vm,
      pendingRequests: params.pendingRequests,
      config: params.config,
    }),
  );
  return { vm, didTimeout: () => timedOut || deadlineReached() };
}

function takeOutput(vm: QuickJS): unknown[] {
  return vm.global.getProp("__openclawTakeOutput").consume((take) =>
    vm.callFunction(take, vm.undefined).consume((output) => {
      const dumped = vm.dump(output);
      return Array.isArray(dumped) ? (dumped as unknown[]) : [];
    }),
  );
}

function takeOutputSafely(vm: QuickJS): unknown[] {
  try {
    return takeOutput(vm);
  } catch {
    return [];
  }
}

function enforceWorkerOutputLimit(
  value: unknown,
  config: CodeModeConfig,
  consumedBytes = 0,
): number {
  const bytes = Buffer.byteLength(JSON.stringify(toJsonSafe(value)) ?? "null", "utf8");
  if (consumedBytes + bytes > config.maxOutputBytes) {
    throw new CodeModeWorkerFailure("output_limit_exceeded", "code mode output limit exceeded");
  }
  return bytes;
}

function throwWorkerFailureWithOutput(params: {
  error: unknown;
  didTimeout: () => boolean;
  output: unknown[];
  vm: QuickJS;
  config: CodeModeConfig;
}): never {
  const timedOut = params.didTimeout() || isQuickJsInterruptedError(params.error);
  const failureOutput = params.output.length > 0 ? params.output : takeOutputSafely(params.vm);
  if (
    params.error instanceof CodeModeWorkerFailure &&
    params.error.code === "output_limit_exceeded"
  ) {
    throw new CodeModeWorkerFailureWithOutput(params.error.code, params.error.message, [], {
      cause: params.error,
    });
  }
  try {
    enforceWorkerOutputLimit(failureOutput, params.config);
  } catch (error) {
    if (error instanceof CodeModeWorkerFailure) {
      throw new CodeModeWorkerFailureWithOutput(error.code, error.message, [], { cause: error });
    }
    throw error;
  }
  if (timedOut) {
    throw new CodeModeWorkerFailureWithOutput(
      "timeout",
      "code mode timeout exceeded",
      failureOutput,
      { cause: params.error },
    );
  }
  if (params.error instanceof CodeModeWorkerFailure) {
    throw new CodeModeWorkerFailureWithOutput(
      params.error.code,
      params.error.message,
      failureOutput,
      { cause: params.error },
    );
  }
  if (failureOutput.length > 0) {
    throw new CodeModeWorkerFailureWithOutput(
      "internal_error",
      errorMessage(params.error),
      failureOutput,
      { cause: params.error },
    );
  }
  throw params.error;
}

async function readCompletedResult(vm: QuickJS, resultHandle: JSValueHandle): Promise<unknown> {
  if (!resultHandle.isPromise) {
    return toJsonSafe(vm.dump(resultHandle));
  }
  const settled = await vm.resolvePromise(resultHandle);
  if ("error" in settled) {
    return settled.error.consume((error) => {
      // vm.dump rebuilds a host Error carrying the QuickJS name/message/stack;
      // format it like the synchronous path so async rejections keep their cause
      // and location instead of collapsing to the bare message.
      const dumped = vm.dump(error);
      // Node module globals are deliberately absent from the WASI guest. Keep
      // aliases fail-closed at that runtime boundary rather than guessing source
      // provenance or installing a host-backed loader.
      if (
        dumped instanceof Error &&
        dumped.name === "ReferenceError" &&
        /^(?:require|module|process) is not defined$/u.test(dumped.message)
      ) {
        throw new CodeModeWorkerFailure("invalid_input", "code mode module access is disabled.");
      }
      const text =
        dumped instanceof Error
          ? formatQuickJsError(dumped.name, dumped.message, dumped.stack)
          : errorMessage(dumped);
      throw new Error(text);
    });
  }
  return settled.value.consume((value) => toJsonSafe(vm.dump(value)));
}

function waitingResult(params: {
  vm: QuickJS;
  pendingRequests: PendingBridgeRequest[];
  settlementMode: Extract<CodeModeWorkerResult, { status: "waiting" }>["settlementMode"];
  output: unknown[];
  config: CodeModeConfig;
}): CodeModeWorkerResult {
  const snapshotBytes = QuickJS.serializeSnapshot(params.vm.snapshot());
  if (snapshotBytes.byteLength > params.config.maxSnapshotBytes) {
    throw new CodeModeWorkerFailure("snapshot_limit_exceeded", "code mode snapshot limit exceeded");
  }
  return {
    status: "waiting",
    snapshotBytes,
    pendingRequests: params.pendingRequests,
    settlementMode: params.settlementMode,
    output: params.output,
  };
}

async function runVmExecution(params: {
  vm: QuickJS;
  didTimeout: () => boolean;
  pendingRequests: PendingBridgeRequest[];
  config: CodeModeConfig;
  prepare: () => void;
}): Promise<CodeModeWorkerResult> {
  let output: unknown[] = [];
  try {
    params.prepare();
    params.vm.executePendingJobs();
    output = takeOutput(params.vm);
    const outputBytes = enforceWorkerOutputLimit(output, params.config);
    const resultHandle = params.vm.global.getProp("__openclawResult");
    try {
      const promisePending = resultHandle.isPromise && resultHandle.promiseState === 0;
      if (promisePending && params.pendingRequests.length === 0) {
        throw new Error("code mode promise is pending without host work");
      }
      const requiredPendingRequestIds = params.pendingRequests.map((request) => request.id);
      if (promisePending || requiredPendingRequestIds.length > 0) {
        // Native await does not expose Promise ownership. Every dispatched
        // call remains required, including detached calls and race branches.
        return waitingResult({
          vm: params.vm,
          pendingRequests: params.pendingRequests,
          settlementMode: promisePending
            ? { kind: "awaiting" }
            : { kind: "draining", requiredRequestIds: requiredPendingRequestIds },
          output,
          config: params.config,
        });
      }
      const value = await readCompletedResult(params.vm, resultHandle);
      enforceWorkerOutputLimit(value, params.config, output.length > 0 ? outputBytes : 0);
      return {
        status: "completed",
        value,
        output,
      };
    } finally {
      resultHandle.dispose();
    }
  } catch (error) {
    return throwWorkerFailureWithOutput({
      error,
      didTimeout: params.didTimeout,
      output,
      vm: params.vm,
      config: params.config,
    });
  } finally {
    params.vm.dispose();
  }
}

async function runExec(input: Extract<CodeModeWorkerPayload, { kind: "exec" }>) {
  const pendingRequests: PendingBridgeRequest[] = [];
  const { vm, didTimeout } = await createVm({
    wasmModule: input.wasmModule,
    catalog: input.catalog,
    apiFiles: input.apiFiles ?? [],
    namespaces: input.namespaces,
    swarmEnabled: input.swarmEnabled === true,
    config: input.config,
    pendingRequests,
  });
  return runVmExecution({
    vm,
    didTimeout,
    pendingRequests,
    config: input.config,
    prepare: () => {
      vm.evalCode(
        buildUserSource(input.source),
        "openclaw-code-mode:user.js",
        EvalFlags.ASYNC,
      ).dispose();
    },
  });
}

async function runResume(input: Extract<CodeModeWorkerPayload, { kind: "resume" }>) {
  // Restored promises keep their original bridge ids; do not redispatch calls
  // that are still running when a faster sibling resumes this snapshot.
  const pendingRequests: PendingBridgeRequest[] = [...(input.pendingRequests ?? [])];
  const { vm, didTimeout } = await restoreVm({
    wasmModule: input.wasmModule,
    snapshotBytes: input.snapshotBytes,
    config: input.config,
    pendingRequests,
  });
  return runVmExecution({
    vm,
    didTimeout,
    pendingRequests,
    config: input.config,
    prepare: () => {
      vm.global.getProp("__openclawSettleBridge").consume((settle) => {
        for (const request of input.settledRequests) {
          const id = vm.newString(request.id);
          const payload = vm.newString(JSON.stringify(request.ok ? request.value : request.error));
          try {
            vm.callFunction(
              settle,
              vm.undefined,
              id,
              request.ok ? vm.true : vm.false,
              payload,
            ).dispose();
          } finally {
            id.dispose();
            payload.dispose();
          }
        }
      });
    },
  });
}

function isQuickJsWasmModule(value: unknown): value is WebAssembly.Module {
  return Object.prototype.toString.call(value) === "[object WebAssembly.Module]";
}

async function main(): Promise<CodeModeWorkerResult> {
  const input = workerData as unknown;
  if (!isRecord(input) || !isRecord(input.config) || !isQuickJsWasmModule(input.wasmModule)) {
    return {
      status: "failed",
      error: "invalid code mode worker input",
      code: "invalid_input",
      failurePhase: "input",
      bridgeDispatchStarted: false,
      output: [],
    };
  }
  try {
    if (input.kind === "exec" && typeof input.source === "string") {
      return await runExec({
        kind: "exec",
        wasmModule: input.wasmModule,
        source: input.source,
        config: input.config as CodeModeConfig,
        catalog: Array.isArray(input.catalog) ? input.catalog : [],
        apiFiles: Array.isArray(input.apiFiles) ? (input.apiFiles as CodeModeApiVirtualFile[]) : [],
        namespaces: Array.isArray(input.namespaces)
          ? (input.namespaces as CodeModeNamespaceDescriptor[])
          : [],
        swarmEnabled: input.swarmEnabled === true,
      });
    }
    if (input.kind === "resume" && input.snapshotBytes instanceof Uint8Array) {
      return await runResume({
        kind: "resume",
        wasmModule: input.wasmModule,
        snapshotBytes: input.snapshotBytes,
        config: input.config as CodeModeConfig,
        settledRequests: Array.isArray(input.settledRequests)
          ? (input.settledRequests as SettledBridgeRequest[])
          : [],
        pendingRequests: Array.isArray(input.pendingRequests)
          ? (input.pendingRequests as PendingBridgeRequest[])
          : [],
      });
    }
    return {
      status: "failed",
      error: "invalid code mode worker input",
      code: "invalid_input",
      failurePhase: "input",
      bridgeDispatchStarted: false,
      output: [],
    };
  } catch (error) {
    const timedOut = isQuickJsInterruptedError(error);
    const code = timedOut
      ? "timeout"
      : error instanceof CodeModeWorkerFailure
        ? error.code
        : "internal_error";
    return {
      status: "failed",
      error: timedOut ? "code mode timeout exceeded" : errorMessage(error),
      code,
      failurePhase: code === "invalid_input" ? "input" : "guest",
      bridgeDispatchStarted: false,
      output: error instanceof CodeModeWorkerFailureWithOutput ? error.output : [],
    };
  }
}

if (parentPort) {
  Reflect.apply(Reflect.get(parentPort, "postMessage") as (message: unknown) => void, parentPort, [
    await main(),
  ]);
}