@ai-sdk/code-mode 1.0.5 → 1.0.7

package/src/runtime/worker.ts

@@ -28,6 +28,7 @@ const pendingBridgeRequests = new Map<
     context: QuickJSAsyncContext;
     deferred: QuickJSDeferredPromise;
     invocationId: string;
+    resetDateNow?: QuickJSHandle;
   }
 >();
 let activeInvocationId: string | undefined;
@@ -35,6 +36,20 @@ let bridgeRequestCounter = 0;
 let embeddedQuickJsWasmModulePromise: Promise<WebAssembly.Module> | undefined;
 
 parentPort.on('message', (message: MainToWorkerMessage) => {
+  if (message.type === 'bridge-drain') {
+    if (message.invocationId !== activeInvocationId) {
+      throw new CodeModeProtocolError(
+        `Bridge drain invocationId mismatch: expected ${activeInvocationId ?? 'no active invocation'}, received ${message.invocationId}.`,
+      );
+    }
+    parentPort?.postMessage({
+      type: 'bridge-drained',
+      invocationId: message.invocationId,
+      drainId: message.drainId,
+    });
+    return;
+  }
+
   if (message.type === 'bridge-response') {
     const pending = pendingBridgeRequests.get(message.requestId);
     if (!pending) {
@@ -59,7 +74,12 @@ parentPort.on('message', (message: MainToWorkerMessage) => {
     }
 
     pendingBridgeRequests.delete(message.requestId);
-    resolveBridgeResponse(pending.context, pending.deferred, message);
+    resolveBridgeResponse(
+      pending.context,
+      pending.deferred,
+      message,
+      pending.resetDateNow,
+    );
     return;
   }
 
@@ -111,6 +131,8 @@ async function execute(message: WorkerRunMessage): Promise<string> {
   const deadline = Date.now() + message.options.timeoutMs;
   let interruptChecks = 0;
   let bridgeFunctions: { invokeTool: QuickJSHandle } | undefined;
+  let determinismHandle: QuickJSHandle | undefined;
+  let resetDateNowHandle: QuickJSHandle | undefined;
 
   runtime.setMemoryLimit(message.options.memoryLimitBytes);
   runtime.setMaxStackSize(message.options.maxStackSizeBytes);
@@ -121,7 +143,12 @@ async function execute(message: WorkerRunMessage): Promise<string> {
 
   try {
     installConsole(context, message.options.maxConsoleOutputBytes);
-    bridgeFunctions = createBridgeFunctions(context, message);
+    bridgeFunctions = createBridgeFunctions(
+      context,
+      message,
+      () => resetDateNowHandle,
+    );
+    determinismHandle = jsToHandle(context, message.determinism);
     const setupSource = buildGuestRuntimeSetupSource();
     const setupEvalResult = await context.evalCodeAsync(
       setupSource,
@@ -139,6 +166,7 @@ async function execute(message: WorkerRunMessage): Promise<string> {
         setupEvalResult.value,
         context.undefined,
         bridgeFunctions.invokeTool,
+        determinismHandle,
       );
       if (setupCallResult.error) {
         const error = context.dump(setupCallResult.error);
@@ -148,6 +176,10 @@ async function execute(message: WorkerRunMessage): Promise<string> {
         throw toError(error);
       }
       if (setupCallResult.value.alive) {
+        resetDateNowHandle = context.getProp(
+          setupCallResult.value,
+          'resetDateNow',
+        );
         setupCallResult.value.dispose();
       }
     } finally {
@@ -210,6 +242,12 @@ async function execute(message: WorkerRunMessage): Promise<string> {
     if (bridgeFunctions?.invokeTool.alive) {
       bridgeFunctions.invokeTool.dispose();
     }
+    if (determinismHandle?.alive) {
+      determinismHandle.dispose();
+    }
+    if (resetDateNowHandle?.alive) {
+      resetDateNowHandle.dispose();
+    }
     context.dispose();
   }
 }
@@ -339,16 +377,19 @@ function dumpConsoleArg(
 function createBridgeFunctions(
   context: QuickJSAsyncContext,
   message: WorkerRunMessage,
+  getResetDateNow: () => QuickJSHandle | undefined,
 ): { invokeTool: QuickJSHandle } {
   const invokeTool = context.newFunction(
     '__codeModeInvokeTool',
     (toolNameHandle: QuickJSHandle, inputJsonHandle: QuickJSHandle) => {
       const toolName = context.getString(toolNameHandle);
       const inputJson = context.getString(inputJsonHandle);
-      return requestHost(context, message.invocationId, {
-        toolName,
-        inputJson,
-      });
+      return requestHost(
+        context,
+        message.invocationId,
+        { toolName, inputJson },
+        getResetDateNow(),
+      );
     },
   );
 
@@ -359,6 +400,7 @@ function requestHost(
   context: QuickJSAsyncContext,
   invocationId: string,
   payload: Record<string, unknown>,
+  resetDateNow: QuickJSHandle | undefined,
 ): QuickJSHandle {
   const requestId = `${invocationId}:bridge-${++bridgeRequestCounter}`;
   const deferred = context.newPromise();
@@ -366,6 +408,7 @@ function requestHost(
     context,
     deferred,
     invocationId,
+    ...(resetDateNow !== undefined ? { resetDateNow } : {}),
   });
   deferred.settled.then(() => {
     context.runtime.executePendingJobs();
@@ -384,7 +427,9 @@ function resolveBridgeResponse(
   context: QuickJSAsyncContext,
   deferred: QuickJSDeferredPromise,
   message: WorkerBridgeResponse,
+  resetDateNow?: QuickJSHandle,
 ): void {
+  resetGuestDateNow(context, resetDateNow, message.dateNowMs);
   if (message.success) {
     const value = context.newString(message.valueJson ?? '');
     deferred.resolve(value);
@@ -396,6 +441,32 @@ function resolveBridgeResponse(
   error.dispose();
 }
 
+function resetGuestDateNow(
+  context: QuickJSAsyncContext,
+  resetDateNow: QuickJSHandle | undefined,
+  dateNowMs: number,
+): void {
+  if (resetDateNow === undefined) {
+    return;
+  }
+  const value = context.newNumber(dateNowMs);
+  try {
+    const result = context.callFunction(resetDateNow, context.undefined, value);
+    if (result.error) {
+      const error = context.dump(result.error);
+      if (result.error.alive) {
+        result.error.dispose();
+      }
+      throw toError(error);
+    }
+    if (result.value.alive) {
+      result.value.dispose();
+    }
+  } finally {
+    value.dispose();
+  }
+}
+
 function drainPendingJobs(context: QuickJSAsyncContext): void {
   while (context.runtime.hasPendingJob()) {
     const pending = context.runtime.executePendingJobs();
@@ -495,3 +566,42 @@ function createBridgeErrorHandle(
   }
   return handle;
 }
+
+function jsToHandle(
+  context: QuickJSAsyncContext,
+  value: unknown,
+): QuickJSHandle {
+  if (value === null || value === undefined) {
+    return context.undefined;
+  }
+  if (typeof value === 'string') {
+    return context.newString(value);
+  }
+  if (typeof value === 'number') {
+    return context.newNumber(value);
+  }
+  if (typeof value === 'boolean') {
+    return value ? context.true : context.false;
+  }
+  if (Array.isArray(value)) {
+    const array = context.newArray();
+    value.forEach((item, index) => {
+      const itemHandle = jsToHandle(context, item);
+      context.setProp(array, index, itemHandle);
+      itemHandle.dispose();
+    });
+    return array;
+  }
+  if (typeof value === 'object') {
+    const object = context.newObject();
+    for (const [key, item] of Object.entries(
+      value as Record<string, unknown>,
+    )) {
+      const itemHandle = jsToHandle(context, item);
+      context.setProp(object, key, itemHandle);
+      itemHandle.dispose();
+    }
+    return object;
+  }
+  return context.undefined;
+}

package/src/tool-invocation.ts

@@ -1,28 +1,57 @@
 import { asSchema } from 'ai';
-import { CodeModeToolError } from './errors.js';
-import type { CodeModeToolExecutionOptions, CodeModeToolSet } from './types.js';
+import { CODE_MODE_TOOL_APPROVAL_KIND } from './approval.js';
+import {
+  CodeModeProtocolError,
+  CodeModeToolApprovalDeniedError,
+  CodeModeToolApprovalRequiredError,
+  CodeModeToolError,
+} from './errors.js';
+import { isCodeModeHostInterruptSignal } from './host-interrupt.js';
+import type {
+  CodeModeInterruptExecutionContext,
+  CodeModeInterruptPayload,
+  CodeModeOptions,
+  CodeModeToolExecutionOptions,
+  CodeModeToolSet,
+} from './types.js';
 import {
   assertJsonSerializable,
   toJsonPayload,
 } from './utils/serialization.js';
 
+export type HostToolInvocationResult =
+  | { type: 'success'; valueJson: string }
+  | {
+      type: 'interrupted';
+      toolName: string;
+      input: unknown;
+      toolCallId: string;
+      payload: CodeModeInterruptPayload;
+    };
+
 export async function invokeHostTool({
   toolName,
   inputJson,
   tools,
   baseExecutionOptions,
+  codeModeOptions,
   maxToolInputBytes,
   maxToolOutputBytes,
   toolCallId,
+  codeModeInterrupt,
+  skipApproval = false,
 }: {
   toolName: string;
   inputJson: string;
   tools: CodeModeToolSet;
   baseExecutionOptions: CodeModeToolExecutionOptions;
+  codeModeOptions: CodeModeOptions;
   maxToolInputBytes: number;
   maxToolOutputBytes: number;
   toolCallId: string;
-}): Promise<string> {
+  codeModeInterrupt?: CodeModeInterruptExecutionContext;
+  skipApproval?: boolean;
+}): Promise<HostToolInvocationResult> {
   throwIfAborted(baseExecutionOptions.abortSignal);
 
   const hostTool = tools[toolName];
@@ -55,28 +84,100 @@ export async function invokeHostTool({
   const executionOptions: CodeModeToolExecutionOptions = {
     ...baseExecutionOptions,
     toolCallId,
+    ...(codeModeInterrupt !== undefined ? { codeModeInterrupt } : {}),
   };
 
-  if (
-    await raceAgainstAbort(
+  const needsApproval =
+    !skipApproval &&
+    (await raceAgainstAbort(
       requiresApproval(hostTool, validation.value, executionOptions),
       executionOptions.abortSignal,
-    )
-  ) {
-    throw new CodeModeToolError(
-      `Tool "${toolName}" requires approval, which code mode does not support yet.`,
-      { toolName, input: validation.value, toolCallId },
+    ));
+
+  if (needsApproval) {
+    if (codeModeOptions.approval?.mode === 'interrupt') {
+      return {
+        type: 'interrupted',
+        toolName,
+        input: validation.value,
+        toolCallId,
+        payload: { kind: CODE_MODE_TOOL_APPROVAL_KIND },
+      };
+    }
+
+    const approval = await raceAgainstAbort(
+      Promise.resolve(
+        codeModeOptions.approval?.onApprovalRequired?.({
+          toolName,
+          input: validation.value,
+          toolCallId,
+        }),
+      ),
+      baseExecutionOptions.abortSignal,
     );
+    if (approval === undefined) {
+      throw new CodeModeToolApprovalRequiredError(
+        toolName,
+        validation.value,
+        toolCallId,
+      );
+    }
+    const approved =
+      typeof approval === 'string'
+        ? approval === 'approved'
+        : approval?.approved;
+    const reason = typeof approval === 'string' ? undefined : approval?.reason;
+    if (typeof approved !== 'boolean') {
+      throw new CodeModeProtocolError(
+        `Tool "${toolName}" approval callback returned a malformed approval decision.`,
+        { toolName, toolCallId },
+      );
+    }
+    if (reason !== undefined && typeof reason !== 'string') {
+      throw new CodeModeProtocolError(
+        `Tool "${toolName}" approval callback returned a malformed approval reason.`,
+        { toolName, toolCallId },
+      );
+    }
+    if (!approved) {
+      throw new CodeModeToolApprovalDeniedError(
+        toolName,
+        validation.value,
+        toolCallId,
+        reason,
+      );
+    }
   }
 
-  const output = await raceAgainstAbort(
-    executeHostTool(hostTool.execute.bind(hostTool), {
-      input: validation.value,
-      options: executionOptions,
-    }),
-    executionOptions.abortSignal,
-  );
-  return toJsonPayload(output, maxToolOutputBytes, `Tool "${toolName}" output`);
+  let output: unknown;
+  try {
+    output = await raceAgainstAbort(
+      executeHostTool(hostTool.execute.bind(hostTool), {
+        input: validation.value,
+        options: executionOptions,
+      }),
+      executionOptions.abortSignal,
+    );
+  } catch (error) {
+    if (isCodeModeHostInterruptSignal(error)) {
+      return {
+        type: 'interrupted',
+        toolName,
+        input: validation.value,
+        toolCallId,
+        payload: error.payload,
+      };
+    }
+    throw error;
+  }
+  return {
+    type: 'success',
+    valueJson: toJsonPayload(
+      output,
+      maxToolOutputBytes,
+      `Tool "${toolName}" output`,
+    ),
+  };
 }
 
 async function requiresApproval(

package/src/types.ts

@@ -14,6 +14,7 @@ export interface CodeModeToolExecutionOptions {
   abortSignal?: AbortSignal;
   experimental_context?: unknown;
   context?: unknown;
+  codeModeInterrupt?: CodeModeInterruptExecutionContext;
 }
 
 /**
@@ -36,6 +37,124 @@ export interface CodeModeToolInput {
   js: string;
 }
 
+export type ApprovalDecision =
+  | 'approved'
+  | 'denied'
+  | { approved: boolean; reason?: string };
+
+export interface CodeModeApprovalRequest {
+  toolName: string;
+  input: unknown;
+  toolCallId: string;
+}
+
+export interface CodeModeApprovalResponse {
+  approvalId: string;
+  approved: boolean;
+  reason?: string;
+}
+
+export interface CodeModeApprovalResolution {
+  approved: boolean;
+  reason?: string;
+}
+
+export interface CodeModeInterruptPayload {
+  kind: string;
+  [key: string]: unknown;
+}
+
+export interface CodeModeInterruptResolution<TResolution = unknown> {
+  interruptId: string;
+  resolution: TResolution;
+}
+
+export interface CodeModeInterruptExecutionContext<
+  TPayload extends CodeModeInterruptPayload = CodeModeInterruptPayload,
+  TResolution = unknown,
+> {
+  interruptId: string;
+  payload: TPayload;
+  resolution: TResolution;
+}
+
+export type CodeModeContinuationLedgerEntry =
+  | {
+      kind: 'tool';
+      name: string;
+      inputJson: string;
+      toolCallId: string;
+      status: 'fulfilled';
+      dateNowMs: number;
+      valueJson: string;
+    }
+  | {
+      kind: 'tool';
+      name: string;
+      inputJson: string;
+      toolCallId: string;
+      status: 'rejected';
+      dateNowMs: number;
+      error: SerializableError;
+    }
+  | {
+      kind: 'tool';
+      name: string;
+      inputJson: string;
+      toolCallId: string;
+      interruptId: string;
+      interruptPayload: CodeModeInterruptPayload;
+      status: 'interrupted';
+    };
+
+export interface CodeModeDeterminismState {
+  dateNowMs: number;
+  randomSeed: string;
+}
+
+export interface CodeModeContinuationAuth {
+  alg: 'HMAC-SHA256';
+  nonce: string;
+  issuedAtMs: number;
+  expiresAtMs: number;
+  signature: string;
+}
+
+export interface CodeModeContinuation {
+  version: 1;
+  js: string;
+  outerToolCallId: string;
+  determinism: CodeModeDeterminismState;
+  ledger: CodeModeContinuationLedgerEntry[];
+  auth: CodeModeContinuationAuth;
+}
+
+export type UnsignedCodeModeContinuation = Omit<CodeModeContinuation, 'auth'>;
+
+export interface CodeModeInterrupt<
+  TPayload extends CodeModeInterruptPayload = CodeModeInterruptPayload,
+> {
+  type: 'code-mode-interrupt';
+  interruptId: string;
+  toolName: string;
+  toolCallId: string;
+  outerToolCallId: string;
+  input: unknown;
+  payload: TPayload;
+  continuation: CodeModeContinuation;
+}
+
+export type CodeModeUnwrappedResult =
+  | { status: 'completed'; output: unknown }
+  | { status: 'interrupted'; interrupt: CodeModeInterrupt };
+
+export interface CodeModeApprovalInterruptPayload extends CodeModeInterruptPayload {
+  kind: 'ai-sdk-code-mode/tool-approval';
+}
+
+export type CodeModeApprovalInterrupt =
+  CodeModeInterrupt<CodeModeApprovalInterruptPayload>;
+
 /**
  * Execution limits applied to each sandbox invocation.
  */
@@ -67,6 +186,24 @@ export interface CodeModeExecutionPolicy {
  */
 export interface CodeModeOptions {
   executionPolicy?: CodeModeExecutionPolicy;
+  continuationSecurity?: CodeModeContinuationSecurityOptions;
+  approval?: {
+    /**
+     * @defaultValue `'callback'`
+     */
+    mode?: 'callback' | 'interrupt';
+    onApprovalRequired?: (
+      request: CodeModeApprovalRequest,
+    ) => Promise<ApprovalDecision> | ApprovalDecision;
+  };
+}
+
+export interface CodeModeContinuationSecurityOptions {
+  signingKey?: string | Uint8Array;
+  /**
+   * @defaultValue `60 * 60 * 1000`
+   */
+  maxAgeMs?: number;
 }
 
 /**
@@ -77,6 +214,8 @@ export interface RunCodeModeInput {
   tools: CodeModeToolSet;
   toolExecutionOptions?: Partial<CodeModeToolExecutionOptions>;
   options?: CodeModeOptions;
+  continuation?: CodeModeContinuation;
+  interruptResolution?: CodeModeInterruptResolution;
 }
 
 /**