@ai-sdk/code-mode 1.0.5 → 1.0.7

package/src/runtime/manager.ts

@@ -1,6 +1,17 @@
 import { AsyncResource } from 'node:async_hooks';
 import { Buffer } from 'node:buffer';
+import { randomBytes } from 'node:crypto';
 import { Worker } from 'node:worker_threads';
+import {
+  isCodeModeApprovalInterruptPayload,
+  normalizeApprovalResolution,
+} from '../approval.js';
+import {
+  resolveCodeModeContinuationSecurity,
+  signCodeModeContinuation,
+  verifyCodeModeContinuation,
+  type ResolvedCodeModeContinuationSecurity,
+} from '../continuation-capability.js';
 import {
   CodeModeAbortedError,
   CodeModeBridgeLimitError,
@@ -8,6 +19,7 @@ import {
   CodeModeDetachedBridgeRequestError,
   CodeModeProtocolError,
   CodeModeTimeoutError,
+  CodeModeToolApprovalDeniedError,
   deserializeError,
   serializeBridgeErrorForGuest,
 } from '../errors.js';
@@ -15,6 +27,9 @@ import { invokeHostTool } from '../tool-invocation.js';
 import { normalizeOptions } from '../utils/options.js';
 import { assertSourceSize, transformSource } from '../utils/source-cache.js';
 import type {
+  CodeModeContinuationLedgerEntry,
+  CodeModeInterrupt,
+  CodeModeInterruptExecutionContext,
   CodeModeToolExecutionOptions,
   NormalizedCodeModeOptions,
   RunCodeModeInput,
@@ -22,7 +37,6 @@ import type {
 import { getMaxWorkers } from './max-workers.js';
 import type {
   MainToWorkerMessage,
-  WorkerBridgeResponse,
   WorkerReadyMessage,
   WorkerResultMessage,
   WorkerToolRequest,
@@ -48,6 +62,9 @@ export async function runManagedCodeMode(
   input: RunCodeModeInput,
 ): Promise<unknown> {
   const normalizedOptions = normalizeOptions(input.options);
+  const continuationSecurity = resolveCodeModeContinuationSecurity(
+    input.options?.continuationSecurity,
+  );
   const maxWorkers = getMaxWorkers({
     memoryLimitBytes: normalizedOptions.memoryLimitBytes,
     activeWorkers: activeInvocations,
@@ -63,11 +80,19 @@ export async function runManagedCodeMode(
     }
 
     assertSourceSize(input.js, normalizedOptions.maxSourceBytes);
+    const js = transformSource(input.js);
+    assertContinuationInput({
+      js,
+      continuation: input.continuation,
+      interruptResolution: input.interruptResolution,
+      continuationSecurity,
+    });
     const run = startWorkerRun({
       ...input,
-      js: transformSource(input.js),
+      js,
       normalizedOptions,
       maxWorkers,
+      continuationSecurity,
     });
     return await run.result;
   } finally {
@@ -79,11 +104,16 @@ function startWorkerRun({
   js,
   tools,
   toolExecutionOptions,
+  options,
+  continuation,
+  interruptResolution,
   normalizedOptions,
   maxWorkers,
+  continuationSecurity,
 }: RunCodeModeInput & {
   normalizedOptions: NormalizedCodeModeOptions;
   maxWorkers: number;
+  continuationSecurity: ResolvedCodeModeContinuationSecurity;
 }): ManagedWorkerRun {
   const invocationId = `code-mode-${++invocationCounter}`;
   const pooledWorker = acquireWorker(maxWorkers);
@@ -99,12 +129,16 @@ function startWorkerRun({
 
   const outerAbortSignal = toolExecutionOptions?.abortSignal;
   const invocationAbortController = new AbortController();
+  let nestedToolCounter = getMaxNestedToolCounter(continuation?.ledger ?? []);
   const forwardedContext =
     toolExecutionOptions?.context ?? toolExecutionOptions?.experimental_context;
   const forwardedExperimentalContext =
     toolExecutionOptions?.experimental_context ?? toolExecutionOptions?.context;
   const baseExecutionOptions: CodeModeToolExecutionOptions = {
-    toolCallId: toolExecutionOptions?.toolCallId ?? invocationId,
+    toolCallId:
+      toolExecutionOptions?.toolCallId ??
+      continuation?.outerToolCallId ??
+      invocationId,
     messages: toolExecutionOptions?.messages ?? [],
     abortSignal: invocationAbortController.signal,
     ...(forwardedExperimentalContext !== undefined
@@ -119,7 +153,13 @@ function startWorkerRun({
   let workerCleanedUp = false;
   let totalBridgeRequests = 0;
   let inFlightBridgeRequests = 0;
+  let interruptEntryIndex: number | undefined;
+  let interruptDrainCounter = 0;
+  let pendingInterruptDrainId: string | undefined;
+  let interruptResolutionConsumed = interruptResolution === undefined;
   const seenWorkerRequestIds = new Set<string>();
+  const bridgeLedger = cloneLedger(continuation?.ledger ?? []);
+  const determinism = continuation?.determinism ?? createDeterminismState();
 
   let resolveResult!: (value: unknown) => void;
   let rejectResult!: (reason?: unknown) => void;
@@ -219,6 +259,30 @@ function startWorkerRun({
       return;
     }
 
+    if (message.type === 'bridge-drained') {
+      if (
+        pendingInterruptDrainId === undefined ||
+        message.drainId !== pendingInterruptDrainId
+      ) {
+        failTerminal(
+          new CodeModeProtocolError(
+            `Worker sent unexpected bridge drain acknowledgement ${message.drainId}.`,
+            {
+              invocationId,
+              expectedDrainId: pendingInterruptDrainId,
+              receivedDrainId: message.drainId,
+            },
+          ),
+        );
+        return;
+      }
+      pendingInterruptDrainId = undefined;
+      if (inFlightBridgeRequests === 0) {
+        settleInterrupt();
+      }
+      return;
+    }
+
     const bridgeIndex = markWorkerRequest(message);
     if (bridgeIndex !== undefined) {
       void handleToolRequest(message, bridgeIndex);
@@ -251,6 +315,7 @@ function startWorkerRun({
     type: 'run',
     invocationId,
     js,
+    determinism,
     options: {
       timeoutMs: normalizedOptions.timeoutMs,
       memoryLimitBytes: normalizedOptions.memoryLimitBytes,
@@ -321,37 +386,286 @@ function startWorkerRun({
     message: WorkerToolRequest,
     bridgeIndex: number,
   ): Promise<void> {
+    const replayEntry = bridgeLedger[bridgeIndex - 1];
     try {
-      const valueJson = await invokeHostTool({
+      if (replayEntry !== undefined) {
+        assertReplayEntryMatches(replayEntry, message, bridgeIndex);
+
+        if (replayEntry.status === 'fulfilled') {
+          postBridgeResponse({
+            type: 'bridge-response',
+            invocationId,
+            requestId: message.requestId,
+            success: true,
+            dateNowMs: replayEntry.dateNowMs,
+            valueJson: replayEntry.valueJson,
+          });
+          return;
+        }
+
+        if (replayEntry.status === 'rejected') {
+          postBridgeResponse({
+            type: 'bridge-response',
+            invocationId,
+            requestId: message.requestId,
+            success: false,
+            dateNowMs: replayEntry.dateNowMs,
+            error: replayEntry.error,
+          });
+          return;
+        }
+
+        if (interruptResolution?.interruptId !== replayEntry.interruptId) {
+          requestInterrupt(bridgeIndex - 1);
+          return;
+        }
+
+        interruptResolutionConsumed = true;
+        const isApproval = isCodeModeApprovalInterruptPayload(
+          replayEntry.interruptPayload,
+        );
+        if (isApproval) {
+          const decision = normalizeApprovalResolution(
+            interruptResolution.resolution,
+          );
+          if (!decision.approved) {
+            const error = new CodeModeToolApprovalDeniedError(
+              replayEntry.name,
+              fromJsonPayload(message.inputJson),
+              replayEntry.toolCallId,
+              decision.reason,
+            );
+            const dateNowMs = Date.now();
+            const guestError = serializeBridgeErrorForGuest(error, 'tool');
+            bridgeLedger[bridgeIndex - 1] = {
+              kind: 'tool',
+              name: message.toolName,
+              inputJson: message.inputJson,
+              toolCallId: replayEntry.toolCallId,
+              status: 'rejected',
+              dateNowMs,
+              error: guestError,
+            };
+            postBridgeResponse({
+              type: 'bridge-response',
+              invocationId,
+              requestId: message.requestId,
+              success: false,
+              dateNowMs,
+              error: guestError,
+            });
+            return;
+          }
+        }
+
+        const codeModeInterrupt: CodeModeInterruptExecutionContext | undefined =
+          isApproval
+            ? undefined
+            : {
+                interruptId: replayEntry.interruptId,
+                payload: replayEntry.interruptPayload,
+                resolution: interruptResolution.resolution,
+              };
+        const outcome = await invokeHostTool({
+          toolName: message.toolName,
+          inputJson: message.inputJson,
+          tools,
+          baseExecutionOptions,
+          codeModeOptions: options ?? {},
+          maxToolInputBytes: normalizedOptions.maxToolInputBytes,
+          maxToolOutputBytes: normalizedOptions.maxToolOutputBytes,
+          toolCallId: replayEntry.toolCallId,
+          ...(codeModeInterrupt !== undefined ? { codeModeInterrupt } : {}),
+          skipApproval: true,
+        });
+        if (outcome.type === 'interrupted') {
+          bridgeLedger[bridgeIndex - 1] = {
+            kind: 'tool',
+            name: message.toolName,
+            inputJson: message.inputJson,
+            toolCallId: replayEntry.toolCallId,
+            interruptId: `${replayEntry.toolCallId}:interrupt`,
+            interruptPayload: outcome.payload,
+            status: 'interrupted',
+          };
+          requestInterrupt(bridgeIndex - 1);
+          return;
+        }
+
+        const dateNowMs = Date.now();
+        bridgeLedger[bridgeIndex - 1] = {
+          kind: 'tool',
+          name: message.toolName,
+          inputJson: message.inputJson,
+          toolCallId: replayEntry.toolCallId,
+          status: 'fulfilled',
+          dateNowMs,
+          valueJson: outcome.valueJson,
+        };
+        postBridgeResponse({
+          type: 'bridge-response',
+          invocationId,
+          requestId: message.requestId,
+          success: true,
+          dateNowMs,
+          valueJson: outcome.valueJson,
+        });
+        return;
+      }
+
+      const toolCallId = `${baseExecutionOptions.toolCallId}:tool-${++nestedToolCounter}`;
+      const outcome = await invokeHostTool({
         toolName: message.toolName,
         inputJson: message.inputJson,
         tools,
         baseExecutionOptions,
+        codeModeOptions: options ?? {},
         maxToolInputBytes: normalizedOptions.maxToolInputBytes,
         maxToolOutputBytes: normalizedOptions.maxToolOutputBytes,
-        toolCallId: `${baseExecutionOptions.toolCallId}:tool-${bridgeIndex}`,
+        toolCallId,
       });
+      if (outcome.type === 'interrupted') {
+        bridgeLedger[bridgeIndex - 1] = {
+          kind: 'tool',
+          name: message.toolName,
+          inputJson: message.inputJson,
+          toolCallId,
+          interruptId: `${toolCallId}:interrupt`,
+          interruptPayload: outcome.payload,
+          status: 'interrupted',
+        };
+        requestInterrupt(bridgeIndex - 1);
+        return;
+      }
+
+      const dateNowMs = Date.now();
+      bridgeLedger[bridgeIndex - 1] = {
+        kind: 'tool',
+        name: message.toolName,
+        inputJson: message.inputJson,
+        toolCallId,
+        status: 'fulfilled',
+        dateNowMs,
+        valueJson: outcome.valueJson,
+      };
       postBridgeResponse({
         type: 'bridge-response',
         invocationId,
         requestId: message.requestId,
         success: true,
-        valueJson,
+        dateNowMs,
+        valueJson: outcome.valueJson,
       });
     } catch (error) {
+      if (error instanceof CodeModeProtocolError) {
+        failTerminal(error);
+        return;
+      }
+      const dateNowMs = Date.now();
+      const guestError = serializeBridgeErrorForGuest(error, 'tool');
+      const toolCallId =
+        replayEntry?.toolCallId ??
+        `${baseExecutionOptions.toolCallId}:tool-${nestedToolCounter}`;
+      if (
+        replayEntry === undefined ||
+        (replayEntry.status === 'interrupted' &&
+          interruptResolution?.interruptId === replayEntry.interruptId)
+      ) {
+        bridgeLedger[bridgeIndex - 1] = {
+          kind: 'tool',
+          name: message.toolName,
+          inputJson: message.inputJson,
+          toolCallId,
+          status: 'rejected',
+          dateNowMs,
+          error: guestError,
+        };
+      }
       postBridgeResponse({
         type: 'bridge-response',
         invocationId,
         requestId: message.requestId,
         success: false,
-        error: serializeBridgeErrorForGuest(error, 'tool'),
+        dateNowMs,
+        error: guestError,
       });
     } finally {
       inFlightBridgeRequests--;
+      settleInterruptIfReady();
+    }
+  }
+
+  function requestInterrupt(entryIndex: number): void {
+    interruptEntryIndex ??= entryIndex;
+  }
+
+  function settleInterruptIfReady(): void {
+    if (
+      interruptEntryIndex === undefined ||
+      terminalReached ||
+      inFlightBridgeRequests > 0 ||
+      pendingInterruptDrainId !== undefined ||
+      !interruptResolutionConsumed
+    ) {
+      return;
+    }
+    pendingInterruptDrainId = `${invocationId}:drain-${++interruptDrainCounter}`;
+    postBridgeResponse({
+      type: 'bridge-drain',
+      invocationId,
+      drainId: pendingInterruptDrainId,
+    });
+  }
+
+  function settleInterrupt(): void {
+    if (
+      interruptEntryIndex === undefined ||
+      terminalReached ||
+      inFlightBridgeRequests > 0 ||
+      !interruptResolutionConsumed
+    ) {
+      return;
+    }
+
+    const entry = bridgeLedger[interruptEntryIndex];
+    if (entry === undefined || entry.status !== 'interrupted') {
+      failTerminal(
+        new CodeModeProtocolError(
+          'Code mode interruption references a missing or resolved ledger entry.',
+          { invocationId, interruptEntryIndex },
+        ),
+      );
+      return;
     }
+
+    const continuationState = signCodeModeContinuation(
+      {
+        version: 1,
+        js,
+        outerToolCallId: baseExecutionOptions.toolCallId,
+        determinism: { ...determinism },
+        ledger: cloneLedger(bridgeLedger),
+      },
+      continuationSecurity,
+    );
+    const interrupt: CodeModeInterrupt = {
+      type: 'code-mode-interrupt',
+      interruptId: entry.interruptId,
+      toolName: entry.name,
+      toolCallId: entry.toolCallId,
+      outerToolCallId: continuationState.outerToolCallId,
+      input: fromJsonPayload(entry.inputJson),
+      payload: structuredClone(entry.interruptPayload),
+      continuation: continuationState,
+    };
+
+    terminalReached = true;
+    abortInvocation(interrupt);
+    cleanupWorker(false);
+    settleCaller(() => resolveResult(interrupt));
   }
 
-  function postBridgeResponse(message: WorkerBridgeResponse): void {
+  function postBridgeResponse(message: MainToWorkerMessage): void {
     if (terminalReached) {
       return;
     }
@@ -367,6 +681,10 @@ function startWorkerRun({
     if (terminalReached) {
       return;
     }
+    if (interruptEntryIndex !== undefined) {
+      settleInterruptIfReady();
+      return;
+    }
     if (resultMessage === undefined) {
       failTerminal(
         new CodeModeProtocolError(
@@ -377,6 +695,31 @@ function startWorkerRun({
       return;
     }
 
+    if (!interruptResolutionConsumed) {
+      failTerminal(
+        new CodeModeProtocolError(
+          'Code mode continuation completed without consuming its interrupt resolution.',
+          { interruptId: interruptResolution?.interruptId },
+        ),
+      );
+      return;
+    }
+    if (
+      continuation !== undefined &&
+      totalBridgeRequests < bridgeLedger.length
+    ) {
+      failTerminal(
+        new CodeModeProtocolError(
+          'Code mode continuation returned before replaying the full bridge ledger.',
+          {
+            replayedBridgeRequests: totalBridgeRequests,
+            ledgerEntries: bridgeLedger.length,
+          },
+        ),
+      );
+      return;
+    }
+
     terminalReached = true;
     const finalResultMessage = resultMessage;
     if (inFlightBridgeRequests > 0) {
@@ -403,6 +746,128 @@ function startWorkerRun({
   }
 }
 
+function assertContinuationInput({
+  js,
+  continuation,
+  interruptResolution,
+  continuationSecurity,
+}: {
+  js: string;
+  continuation: RunCodeModeInput['continuation'];
+  interruptResolution: RunCodeModeInput['interruptResolution'];
+  continuationSecurity: ResolvedCodeModeContinuationSecurity;
+}): void {
+  if (continuation === undefined) {
+    if (interruptResolution !== undefined) {
+      throw new CodeModeProtocolError(
+        'A code-mode interrupt resolution was provided without continuation state.',
+      );
+    }
+    return;
+  }
+
+  verifyCodeModeContinuation(continuation, continuationSecurity);
+  if (continuation.version !== 1) {
+    throw new CodeModeProtocolError(
+      'Unsupported code-mode continuation version.',
+    );
+  }
+  if (continuation.js !== js) {
+    throw new CodeModeProtocolError(
+      'Code mode continuation source does not match the resumed source.',
+    );
+  }
+  if (
+    !Number.isInteger(continuation.determinism.dateNowMs) ||
+    !/^[0-9a-f]{32}$/i.test(continuation.determinism.randomSeed)
+  ) {
+    throw new CodeModeProtocolError(
+      'Code mode continuation determinism state is malformed.',
+    );
+  }
+  for (const entry of continuation.ledger) {
+    if (
+      entry.kind !== 'tool' ||
+      typeof entry.name !== 'string' ||
+      typeof entry.inputJson !== 'string' ||
+      typeof entry.toolCallId !== 'string'
+    ) {
+      throw new CodeModeProtocolError(
+        'Code mode continuation ledger is malformed.',
+      );
+    }
+  }
+  if (interruptResolution !== undefined) {
+    const matches = continuation.ledger.filter(
+      entry =>
+        entry.status === 'interrupted' &&
+        entry.interruptId === interruptResolution.interruptId,
+    );
+    if (matches.length !== 1) {
+      throw new CodeModeProtocolError(
+        'Interrupt resolution does not match exactly one pending continuation ledger entry.',
+        {
+          interruptId: interruptResolution.interruptId,
+          matches: matches.length,
+        },
+      );
+    }
+  }
+}
+
+function assertReplayEntryMatches(
+  entry: CodeModeContinuationLedgerEntry,
+  message: WorkerToolRequest,
+  bridgeIndex: number,
+): void {
+  if (
+    entry.name !== message.toolName ||
+    entry.inputJson !== message.inputJson
+  ) {
+    throw new CodeModeProtocolError(
+      'Code mode continuation replay diverged from its bridge ledger.',
+      {
+        bridgeIndex,
+        expectedToolName: entry.name,
+        receivedToolName: message.toolName,
+      },
+    );
+  }
+}
+
+function createDeterminismState(): {
+  dateNowMs: number;
+  randomSeed: string;
+} {
+  return {
+    dateNowMs: Date.now(),
+    randomSeed: randomBytes(16).toString('hex'),
+  };
+}
+
+function cloneLedger(
+  ledger: CodeModeContinuationLedgerEntry[],
+): CodeModeContinuationLedgerEntry[] {
+  return structuredClone(ledger);
+}
+
+function getMaxNestedToolCounter(
+  ledger: CodeModeContinuationLedgerEntry[],
+): number {
+  let max = 0;
+  for (const entry of ledger) {
+    const match = /:tool-(\d+)$/u.exec(entry.toolCallId);
+    if (match !== null) {
+      max = Math.max(max, Number(match[1]));
+    }
+  }
+  return max;
+}
+
+function fromJsonPayload(valueJson: string): unknown {
+  return valueJson === '' ? undefined : JSON.parse(valueJson);
+}
+
 function acquireWorker(maxPoolSize: number): PooledWorker {
   let pooledWorker = idleWorkers.pop();
   while (pooledWorker?.destroyed) {

package/src/runtime/protocol.ts

@@ -1,9 +1,14 @@
-import type { NormalizedCodeModeOptions, SerializableError } from '../types.js';
+import type {
+  CodeModeDeterminismState,
+  NormalizedCodeModeOptions,
+  SerializableError,
+} from '../types.js';
 
 export interface WorkerRunMessage {
   type: 'run';
   invocationId: string;
   js: string;
+  determinism: CodeModeDeterminismState;
   options: Pick<
     NormalizedCodeModeOptions,
     | 'timeoutMs'
@@ -27,10 +32,23 @@ export interface WorkerBridgeResponse {
   invocationId: string;
   requestId: string;
   success: boolean;
+  dateNowMs: number;
   valueJson?: string;
   error?: SerializableError;
 }
 
+export interface WorkerBridgeDrainRequest {
+  type: 'bridge-drain';
+  invocationId: string;
+  drainId: string;
+}
+
+export interface WorkerBridgeDrainedMessage {
+  type: 'bridge-drained';
+  invocationId: string;
+  drainId: string;
+}
+
 export interface WorkerResultMessage {
   type: 'result';
   invocationId: string;
@@ -46,7 +64,11 @@ export interface WorkerReadyMessage {
 
 export type WorkerToMainMessage =
   | WorkerToolRequest
+  | WorkerBridgeDrainedMessage
   | WorkerResultMessage
   | WorkerReadyMessage;
 
-export type MainToWorkerMessage = WorkerRunMessage | WorkerBridgeResponse;
+export type MainToWorkerMessage =
+  | WorkerRunMessage
+  | WorkerBridgeResponse
+  | WorkerBridgeDrainRequest;