@ai-sdk/harness 1.0.39 → 1.0.41

package/src/agent/internal/turn-telemetry.ts

@@ -42,9 +42,9 @@ export interface TurnTelemetry {
    * model the runtime resolved to (overriding the session's configured id).
    * Idempotent — the first call wins.
    */
-  start(modelId?: string): void;
+  start(modelId?: string): Promise<void>;
   /** Open a step span lazily, before the first content of a step. */
-  ensureStepOpen(): void;
+  ensureStepOpen(): Promise<void>;
   /** Close the current step (on a harness `finish-step`). */
   stepFinish(info: {
     finishReason: unknown;
@@ -52,13 +52,18 @@ export interface TurnTelemetry {
     providerMetadata?: unknown;
     /** The model's output content for this step (text/reasoning/tool-calls). */
     content?: TurnContentPart[];
-  }): void;
+  }): Promise<void>;
   /** A tool execution began (on a `tool-call`). */
   toolStart(call: {
     toolCallId: string;
     toolName: string;
     input: unknown;
-  }): void;
+  }): Promise<void>;
+  /** Execute a host tool through each telemetry integration's context wrapper. */
+  executeTool<T>(input: {
+    toolCallId: string;
+    execute: () => PromiseLike<T>;
+  }): Promise<T>;
   /**
    * A tool execution completed (on its `tool-result` or after host execution).
    * Idempotent per `toolCallId` — the first caller wins, so provider-executed
@@ -67,21 +72,24 @@ export interface TurnTelemetry {
   toolEnd(
     toolCallId: string,
     output: { ok: true; output: unknown } | { ok: false; error: unknown },
-  ): void;
+  ): Promise<void>;
   /** The turn ended (on a harness `finish`). */
-  end(info: { finishReason: unknown; usage: unknown }): void;
+  end(info: { finishReason: unknown; usage: unknown }): Promise<void>;
   /** The turn failed. */
-  error(err: unknown): void;
+  error(err: unknown): Promise<void>;
 }
 
 const NOOP: TurnTelemetry = {
-  start() {},
-  ensureStepOpen() {},
-  stepFinish() {},
-  toolStart() {},
-  toolEnd() {},
-  end() {},
-  error() {},
+  async start() {},
+  async ensureStepOpen() {},
+  async stepFinish() {},
+  async toolStart() {},
+  async executeTool({ execute }) {
+    return await execute();
+  },
+  async toolEnd() {},
+  async end() {},
+  async error() {},
 };
 
 export function createTurnTelemetry(opts: {
@@ -122,10 +130,10 @@ export function createTurnTelemetry(opts: {
 
   // onStart — open the operation (root) span. Deferred until `start()` so the
   // runtime-resolved model can be attached to the operation span + trace label.
-  const fireStart = (): void => {
+  const fireStart = async (): Promise<void> => {
     if (started) return;
     started = true;
-    dispatcher.onStart?.(
+    await dispatcher.onStart?.(
       cast<'onStart'>({
         callId,
         operationId: 'ai.harness',
@@ -147,17 +155,17 @@ export function createTurnTelemetry(opts: {
     );
   };
 
-  const start = (overrideModelId?: string): void => {
+  const start = async (overrideModelId?: string): Promise<void> => {
     if (started) return;
     if (overrideModelId) modelId = overrideModelId;
-    fireStart();
+    await fireStart();
   };
 
-  const ensureStepOpen = (): void => {
-    if (!started) fireStart();
+  const ensureStepOpen = async (): Promise<void> => {
+    if (!started) await fireStart();
     if (stepOpen || ended) return;
     stepOpen = true;
-    dispatcher.onStepStart?.(
+    await dispatcher.onStepStart?.(
       cast<'onStepStart'>({
         callId,
         provider,
@@ -175,7 +183,7 @@ export function createTurnTelemetry(opts: {
     );
     // Open the inference (language-model call) span — the gen_ai home for the
     // step's input and (on end) output messages.
-    dispatcher.onLanguageModelCallStart?.(
+    await dispatcher.onLanguageModelCallStart?.(
       cast<'onLanguageModelCallStart'>({
         callId,
         provider,
@@ -187,12 +195,12 @@ export function createTurnTelemetry(opts: {
   };
 
   /** Close the inference span with the step's output content. */
-  const inferenceEnd = (info: {
+  const inferenceEnd = async (info: {
     finishReason: unknown;
     usage: unknown;
     content: TurnContentPart[];
-  }): void => {
-    dispatcher.onLanguageModelCallEnd?.(
+  }): Promise<void> => {
+    await dispatcher.onLanguageModelCallEnd?.(
       cast<'onLanguageModelCallEnd'>({
         callId,
         finishReason: info.finishReason,
@@ -203,9 +211,9 @@ export function createTurnTelemetry(opts: {
     );
   };
 
-  const closeOpenTools = (): void => {
+  const closeOpenTools = async (): Promise<void> => {
     for (const call of openTools.values()) {
-      dispatcher.onToolExecutionEnd?.(
+      await dispatcher.onToolExecutionEnd?.(
         cast<'onToolExecutionEnd'>({
           callId,
           toolExecutionMs: 0,
@@ -229,16 +237,16 @@ export function createTurnTelemetry(opts: {
     start,
     ensureStepOpen,
 
-    stepFinish(info) {
+    async stepFinish(info) {
       if (!stepOpen) return;
       const content = info.content ?? [];
-      closeOpenTools();
-      inferenceEnd({
+      await closeOpenTools();
+      await inferenceEnd({
         finishReason: info.finishReason,
         usage: info.usage,
         content,
       });
-      dispatcher.onStepEnd?.(
+      await dispatcher.onStepEnd?.(
         cast<'onStepEnd'>({
           callId,
           stepNumber,
@@ -258,10 +266,11 @@ export function createTurnTelemetry(opts: {
       stepNumber += 1;
     },
 
-    toolStart(call) {
-      ensureStepOpen();
+    async toolStart(call) {
+      await ensureStepOpen();
+      if (openTools.has(call.toolCallId)) return;
       openTools.set(call.toolCallId, call);
-      dispatcher.onToolExecutionStart?.(
+      await dispatcher.onToolExecutionStart?.(
         cast<'onToolExecutionStart'>({
           callId,
           messages: [],
@@ -277,11 +286,16 @@ export function createTurnTelemetry(opts: {
       );
     },
 
-    toolEnd(toolCallId, output) {
+    async executeTool({ toolCallId, execute }) {
+      if (dispatcher.executeTool == null) return await execute();
+      return await dispatcher.executeTool({ callId, toolCallId, execute });
+    },
+
+    async toolEnd(toolCallId, output) {
       const call = openTools.get(toolCallId);
       if (call == null) return;
       openTools.delete(toolCallId);
-      dispatcher.onToolExecutionEnd?.(
+      await dispatcher.onToolExecutionEnd?.(
         cast<'onToolExecutionEnd'>({
           callId,
           toolExecutionMs: 0,
@@ -301,17 +315,17 @@ export function createTurnTelemetry(opts: {
       );
     },
 
-    end(info) {
+    async end(info) {
       if (ended) return;
-      if (!started) fireStart();
+      if (!started) await fireStart();
       if (stepOpen) {
-        closeOpenTools();
-        inferenceEnd({
+        await closeOpenTools();
+        await inferenceEnd({
           finishReason: info.finishReason,
           usage: info.usage,
           content: [],
         });
-        dispatcher.onStepEnd?.(
+        await dispatcher.onStepEnd?.(
           cast<'onStepEnd'>({
             callId,
             stepNumber,
@@ -330,7 +344,7 @@ export function createTurnTelemetry(opts: {
         stepOpen = false;
       }
       ended = true;
-      dispatcher.onEnd?.(
+      await dispatcher.onEnd?.(
         cast<'onEnd'>({
           callId,
           operationId: 'ai.harness',
@@ -350,12 +364,12 @@ export function createTurnTelemetry(opts: {
       );
     },
 
-    error(err) {
+    async error(err) {
       if (ended) return;
-      if (!started) fireStart();
-      closeOpenTools();
+      if (!started) await fireStart();
+      await closeOpenTools();
       ended = true;
-      dispatcher.onError?.(err);
+      await dispatcher.onError?.(err);
     },
   };
 }

package/src/bridge/index.ts

@@ -234,6 +234,7 @@ export async function runBridge<TStart extends { type: 'start' }>(
   let currentBoundPort = 0;
   let currentTurnState: BridgeState = 'init';
   let activeSocket: WebSocket | undefined;
+  let activeSocketReadyForLiveEvents = false;
   let isFirstTurn = true;
   let turnAbort: AbortController | undefined;
   let currentUserMessages: string[] | undefined;
@@ -389,7 +390,10 @@ export async function runBridge<TStart extends { type: 'start' }>(
     eventLog.push({ seq, line });
     diskBuffer += `${line}\n`;
     scheduleEventFlush();
-    if (activeSocket?.readyState === WS_OPEN) {
+    if (
+      activeSocketReadyForLiveEvents &&
+      activeSocket?.readyState === WS_OPEN
+    ) {
       try {
         activeSocket.send(line);
       } catch {
@@ -522,6 +526,8 @@ export async function runBridge<TStart extends { type: 'start' }>(
   ): Promise<void> => {
     switch (msg.type) {
       case 'start': {
+        if (activeSocket !== ws) return;
+        activeSocketReadyForLiveEvents = true;
         const firstTurn = isFirstTurn;
         isFirstTurn = false;
         eventLog = []; // clear previous turn; keep seqCounter monotonic
@@ -643,7 +649,9 @@ export async function runBridge<TStart extends { type: 'start' }>(
         }
         return;
       case 'resume':
+        if (activeSocket !== ws) return;
         replay(ws, msg.lastSeenEventId);
+        activeSocketReadyForLiveEvents = true;
         return;
       case 'shutdown':
         currentTurnState = 'done';
@@ -721,6 +729,7 @@ export async function runBridge<TStart extends { type: 'start' }>(
     // (the host reconnecting after a drop). The previous socket's close is a
     // no-op below because it is no longer `activeSocket`.
     activeSocket = ws;
+    activeSocketReadyForLiveEvents = false;
 
     // Announce liveness the instant we accept. Some sandbox runtimes complete
     // the host-side WS handshake before the connection is forwarded here; the
@@ -754,6 +763,7 @@ export async function runBridge<TStart extends { type: 'start' }>(
       // log for replay when the host reconnects.
       if (activeSocket === ws) {
         activeSocket = undefined;
+        activeSocketReadyForLiveEvents = false;
       }
     });
 

package/src/utils/sandbox-channel.ts

@@ -75,6 +75,30 @@ type Listener<TOut extends { type: string }, T extends EventTypeOf<TOut>> = (
   event: Extract<TOut, { type: T }>,
 ) => void;
 
+/*
+ * The agent and utilities entrypoints bundle this module separately. A global
+ * symbol lets the agent recognize metadata attached by the channel's bundle
+ * copy, while the non-enumerable property leaves protocol payloads unchanged.
+ */
+const sandboxChannelEventCheckpointSymbol = Symbol.for(
+  'vercel.ai.harness.sandboxChannelEventCheckpoint',
+);
+
+type SandboxChannelEventCheckpoint = {
+  pin: () => () => void;
+};
+
+export function pinSandboxChannelEventCheckpoint(
+  event: unknown,
+): (() => void) | undefined {
+  if (event == null || typeof event !== 'object') return undefined;
+  return (
+    event as {
+      [sandboxChannelEventCheckpointSymbol]?: SandboxChannelEventCheckpoint;
+    }
+  )[sandboxChannelEventCheckpointSymbol]?.pin();
+}
+
 const sleep = (ms: number): Promise<void> =>
   new Promise(resolve => {
     const t = setTimeout(resolve, ms);
@@ -136,6 +160,9 @@ export class SandboxChannel<
    * replayed to the next process on `resume`.
    */
   private suspended = false;
+  private pinnedSuspensionCursor:
+    | { eventId: number; token: object }
+    | undefined;
   /** Channel is fully torn down; `send` throws and `onClose` has fired. */
   private terminal = false;
   private _lastSeenEventId = 0;
@@ -251,6 +278,9 @@ export class SandboxChannel<
   }
 
   interrupt(options?: { timeoutMs?: number }): Promise<void> {
+    if (this.pinnedSuspensionCursor != null) {
+      return Promise.resolve();
+    }
     const timeoutMs = options?.timeoutMs ?? 5000;
     return new Promise<void>((resolve, reject) => {
       let settled = false;
@@ -313,19 +343,24 @@ export class SandboxChannel<
    * aborts it) and accumulates events past the cursor for the next process to
    * `resume`. Unlike {@link close}, the consumer's active turn is wound down
    * cleanly — adapters distinguish a suspend from an unexpected drop via the
-   * `'suspended'` close reason and resolve `done` successfully.
+   * `'suspended'` close reason and resolve `done` successfully. When an event
+   * checkpoint is pinned, the returned cursor points to that event so any
+   * already-dispatched tail is replayed by the next process.
    */
   suspend(): Promise<number> {
     return new Promise<number>(resolve => {
+      const pinnedSuspensionCursor = this.pinnedSuspensionCursor?.eventId;
       if (this.terminal) {
-        resolve(this._lastSeenEventId);
+        resolve(pinnedSuspensionCursor ?? this._lastSeenEventId);
         return;
       }
       // Stop counting/dispatching further inbound frames immediately, and
       // suppress reconnect so the socket close finalises.
       this.suspended = true;
       this.closing = true;
-      this.onClose(() => resolve(this._lastSeenEventId));
+      this.onClose(() =>
+        resolve(pinnedSuspensionCursor ?? this._lastSeenEventId),
+      );
       // Queue the close behind any already-dispatched frames so everything
       // delivered to the consumer is reflected in the final cursor.
       this.enqueue(() => {
@@ -465,6 +500,9 @@ export class SandboxChannel<
       schema: this.outboundSchema,
     });
     if (validated.success) {
+      if (seq !== undefined) {
+        this.attachEventCheckpoint({ event: validated.value, eventId: seq });
+      }
       this.dispatch(validated.value);
     } else {
       this.dispatch({
@@ -506,6 +544,26 @@ export class SandboxChannel<
     }
   }
 
+  private attachEventCheckpoint(options: {
+    event: TOut;
+    eventId: number;
+  }): void {
+    if (!Object.isExtensible(options.event)) return;
+    Object.defineProperty(options.event, sandboxChannelEventCheckpointSymbol, {
+      value: {
+        pin: () => {
+          const token = {};
+          this.pinnedSuspensionCursor = { eventId: options.eventId, token };
+          return () => {
+            if (this.pinnedSuspensionCursor?.token === token) {
+              this.pinnedSuspensionCursor = undefined;
+            }
+          };
+        },
+      } satisfies SandboxChannelEventCheckpoint,
+    });
+  }
+
   private finalizeClose(code: number, reason: string): void {
     if (this.terminal) return;
     this.terminal = true;