@ai-sdk/harness 1.0.117 → 1.0.119

package/src/utils/sandbox-channel.ts

@@ -4,6 +4,7 @@ import {
   type FlexibleSchema,
 } from '@ai-sdk/provider-utils';
 import type { WebSocket } from 'ws';
+import { sleep } from './sleep';
 
 /**
  * Diagnostic event surfaced by {@link SandboxChannel} during its connection
@@ -22,7 +23,10 @@ export type SandboxChannelDebugEvent =
     };
 
 export interface SandboxChannelReconnectOptions {
-  /** Give up reconnecting after this many milliseconds. Default 30_000. */
+  /**
+   * Give up reconnecting after this many milliseconds, including connection
+   * establishment and backoff delays. Default 30_000.
+   */
   readonly maxElapsedMs?: number;
   /** First backoff delay. Default 50. */
   readonly initialDelayMs?: number;
@@ -30,6 +34,10 @@ export interface SandboxChannelReconnectOptions {
   readonly maxDelayMs?: number;
 }
 
+export interface SandboxChannelConnectOptions {
+  readonly abortSignal: AbortSignal;
+}
+
 export interface SandboxChannelOptions<TOut> {
   /**
    * Open a fresh WebSocket to the bridge and resolve once it is ready to carry
@@ -38,7 +46,7 @@ export interface SandboxChannelOptions<TOut> {
    * every transient reconnect. Must reject if the connection cannot be
    * established.
    */
-  connect: () => Promise<WebSocket>;
+  connect: (options: SandboxChannelConnectOptions) => Promise<WebSocket>;
 
   /** Schema validating inbound (bridge → host) frames. */
   outboundSchema: FlexibleSchema<TOut>;
@@ -99,12 +107,68 @@ export function pinSandboxChannelEventCheckpoint(
   )[sandboxChannelEventCheckpointSymbol]?.pin();
 }
 
-const sleep = (ms: number): Promise<void> =>
-  new Promise(resolve => {
-    const t = setTimeout(resolve, ms);
-    (t as { unref?: () => void }).unref?.();
+function getAbortReason({
+  abortSignal,
+}: {
+  abortSignal: AbortSignal;
+}): unknown {
+  return abortSignal.reason ?? new Error('SandboxChannel connection aborted');
+}
+
+function closeWebSocket({ ws }: { ws: WebSocket }): void {
+  try {
+    const terminable = ws as WebSocket & { terminate?: () => void };
+    if (terminable.terminate != null) {
+      terminable.terminate();
+    } else {
+      ws.close();
+    }
+  } catch {
+    try {
+      ws.close();
+    } catch {
+      // best-effort
+    }
+  }
+}
+
+async function awaitWebSocketConnection({
+  connection,
+  abortSignal,
+}: {
+  connection: Promise<WebSocket>;
+  abortSignal: AbortSignal;
+}): Promise<WebSocket> {
+  if (abortSignal.aborted) {
+    void connection.then(ws => closeWebSocket({ ws })).catch(() => {});
+    throw getAbortReason({ abortSignal });
+  }
+
+  let onAbort: (() => void) | undefined;
+  const aborted = new Promise<never>((_, reject) => {
+    onAbort = () => reject(getAbortReason({ abortSignal }));
+    abortSignal.addEventListener('abort', onAbort, { once: true });
   });
 
+  try {
+    const ws = await Promise.race([connection, aborted]);
+    if (abortSignal.aborted) {
+      closeWebSocket({ ws });
+      throw getAbortReason({ abortSignal });
+    }
+    return ws;
+  } catch (error) {
+    if (abortSignal.aborted) {
+      void connection.then(ws => closeWebSocket({ ws })).catch(() => {});
+    }
+    throw error;
+  } finally {
+    if (onAbort != null) {
+      abortSignal.removeEventListener('abort', onAbort);
+    }
+  }
+}
+
 /**
  * Host-side typed wrapper around the bridge WebSocket connection.
  *
@@ -135,7 +199,9 @@ export class SandboxChannel<
   >();
   private readonly onReconnectHandlers = new Set<() => void>();
 
-  private readonly connectThunk: () => Promise<WebSocket>;
+  private readonly connectThunk: (
+    options: SandboxChannelConnectOptions,
+  ) => Promise<WebSocket>;
   private readonly outboundSchema: FlexibleSchema<TOut>;
   private readonly onDebug:
     | ((event: SandboxChannelDebugEvent) => void)
@@ -169,6 +235,8 @@ export class SandboxChannel<
   private _lastSeenEventId = 0;
   private readonly pendingSends: string[] = [];
   private dispatchChain: Promise<void> = Promise.resolve();
+  private activeConnectAbortController: AbortController | undefined;
+  private reconnectAbortController: AbortController | undefined;
 
   constructor(options: SandboxChannelOptions<TOut>) {
     this.connectThunk = options.connect;
@@ -206,17 +274,27 @@ export class SandboxChannel<
     if (this.terminal) {
       throw new Error('SandboxChannel: cannot open a closed channel.');
     }
-    const ws = await this.connectThunk();
-    this.wire(ws);
-    this.ws = ws;
-    this.connected = true;
-    if (opts?.resume) {
-      this.rawSend(
-        JSON.stringify({
-          type: 'resume',
-          lastSeenEventId: this._lastSeenEventId,
-        }),
-      );
+    const abortController = new AbortController();
+    this.activeConnectAbortController = abortController;
+    try {
+      const ws = await this.connect({
+        abortSignal: abortController.signal,
+      });
+      this.wire(ws);
+      this.ws = ws;
+      this.connected = true;
+      if (opts?.resume) {
+        this.rawSend(
+          JSON.stringify({
+            type: 'resume',
+            lastSeenEventId: this._lastSeenEventId,
+          }),
+        );
+      }
+    } finally {
+      if (this.activeConnectAbortController === abortController) {
+        this.activeConnectAbortController = undefined;
+      }
     }
   }
 
@@ -272,11 +350,13 @@ export class SandboxChannel<
    */
   beginClose(): void {
     this.closing = true;
+    this.abortActiveConnect();
   }
 
   close(): void {
     if (this.terminal) return;
     this.closing = true;
+    this.abortActiveConnect();
     try {
       this.ws?.close();
     } catch {
@@ -310,6 +390,7 @@ export class SandboxChannel<
       // suppress reconnect so the socket close finalises.
       this.suspended = true;
       this.closing = true;
+      this.abortActiveConnect();
       this.onClose(() =>
         resolve(pinnedSuspensionCursor ?? this._lastSeenEventId),
       );
@@ -362,62 +443,143 @@ export class SandboxChannel<
 
   private async reconnectLoop(): Promise<void> {
     if (this.terminal || this.closing) return;
-    const start = Date.now();
+    const deadline = Date.now() + this.maxElapsedMs;
     let attempt = 0;
     let delay = this.initialDelayMs;
-    while (!this.terminal && !this.closing) {
-      attempt++;
-      this.onDebug?.({
-        event: 'reconnect-attempt',
-        attempt,
-        lastSeenEventId: this._lastSeenEventId,
-      });
-      try {
-        const ws = await this.connectThunk();
-        if (this.terminal || this.closing) {
-          try {
-            ws.close();
-          } catch {
-            // best-effort
-          }
+    const reconnectAbortController = new AbortController();
+    this.reconnectAbortController = reconnectAbortController;
+    try {
+      while (!this.terminal && !this.closing) {
+        const remaining = deadline - Date.now();
+        if (remaining <= 0) {
+          this.failReconnect({
+            attempts: attempt,
+            cause: new Error('Reconnect deadline expired'),
+          });
           return;
         }
-        this.wire(ws);
-        this.ws = ws;
-        this.connected = true;
-        // Ask the bridge to replay everything we have not seen, then flush any
-        // host → bridge frames produced while we were disconnected.
-        this.rawSend(
-          JSON.stringify({
-            type: 'resume',
-            lastSeenEventId: this._lastSeenEventId,
-          }),
-        );
-        this.flushPending();
-        for (const handler of this.onReconnectHandlers) handler();
+
+        attempt++;
         this.onDebug?.({
-          event: 'reconnected',
+          event: 'reconnect-attempt',
           attempt,
           lastSeenEventId: this._lastSeenEventId,
         });
-        return;
-      } catch (cause) {
-        if (Date.now() - start >= this.maxElapsedMs) {
-          this.finalizeClose(1006, 'reconnect failed');
+        const abortController = new AbortController();
+        this.activeConnectAbortController = abortController;
+        const abortConnect = () =>
+          abortController.abort(reconnectAbortController.signal.reason);
+        reconnectAbortController.signal.addEventListener(
+          'abort',
+          abortConnect,
+          {
+            once: true,
+          },
+        );
+        const timeout = setTimeout(() => {
+          abortController.abort(new Error('Reconnect deadline expired'));
+        }, remaining);
+        (timeout as { unref?: () => void }).unref?.();
+
+        try {
+          const ws = await this.connect({
+            abortSignal: abortController.signal,
+          });
+          if (this.terminal || this.closing) {
+            closeWebSocket({ ws });
+            return;
+          }
+          this.wire(ws);
+          this.ws = ws;
+          this.connected = true;
+          // Ask the bridge to replay everything we have not seen, then flush any
+          // host → bridge frames produced while we were disconnected.
+          this.rawSend(
+            JSON.stringify({
+              type: 'resume',
+              lastSeenEventId: this._lastSeenEventId,
+            }),
+          );
+          this.flushPending();
+          for (const handler of this.onReconnectHandlers) handler();
           this.onDebug?.({
-            event: 'reconnect-failed',
-            attempts: attempt,
+            event: 'reconnected',
+            attempt,
             lastSeenEventId: this._lastSeenEventId,
-            cause,
           });
           return;
+        } catch (cause) {
+          if (
+            this.terminal ||
+            this.closing ||
+            reconnectAbortController.signal.aborted
+          ) {
+            return;
+          }
+          const remainingAfterFailure = deadline - Date.now();
+          if (remainingAfterFailure <= 0) {
+            this.failReconnect({ attempts: attempt, cause });
+            return;
+          }
+          await sleep({
+            ms: Math.min(delay, remainingAfterFailure),
+            abortSignal: reconnectAbortController.signal,
+          });
+          delay = Math.min(delay * 1.5, this.maxDelayMs);
+        } finally {
+          clearTimeout(timeout);
+          reconnectAbortController.signal.removeEventListener(
+            'abort',
+            abortConnect,
+          );
+          if (this.activeConnectAbortController === abortController) {
+            this.activeConnectAbortController = undefined;
+          }
         }
-        await sleep(delay);
-        delay = Math.min(delay * 1.5, this.maxDelayMs);
+      }
+    } finally {
+      if (this.reconnectAbortController === reconnectAbortController) {
+        this.reconnectAbortController = undefined;
       }
     }
   }
 
+  private connect({
+    abortSignal,
+  }: SandboxChannelConnectOptions): Promise<WebSocket> {
+    return awaitWebSocketConnection({
+      connection: Promise.resolve().then(() =>
+        this.connectThunk({ abortSignal }),
+      ),
+      abortSignal,
+    });
+  }
+
+  private abortActiveConnect(): void {
+    const controller = this.activeConnectAbortController;
+    this.activeConnectAbortController = undefined;
+    controller?.abort(new Error('SandboxChannel connection aborted'));
+    this.reconnectAbortController?.abort(
+      new Error('SandboxChannel connection aborted'),
+    );
+  }
+
+  private failReconnect({
+    attempts,
+    cause,
+  }: {
+    attempts: number;
+    cause: unknown;
+  }): void {
+    this.finalizeClose(1006, 'reconnect failed');
+    this.onDebug?.({
+      event: 'reconnect-failed',
+      attempts,
+      lastSeenEventId: this._lastSeenEventId,
+      cause,
+    });
+  }
+
   private rawSend(text: string): void {
     if (this.connected && this.ws) {
       this.ws.send(text);
@@ -519,6 +681,7 @@ export class SandboxChannel<
 
   private finalizeClose(code: number, reason: string): void {
     if (this.terminal) return;
+    this.abortActiveConnect();
     this.terminal = true;
     this.connected = false;
     for (const h of this.onCloseHandlers) h(code, reason);

package/src/utils/sleep.ts

@@ -0,0 +1,30 @@
+export function sleep({
+  ms,
+  abortSignal,
+}: {
+  ms: number;
+  abortSignal?: AbortSignal;
+}): Promise<void> {
+  return new Promise(resolve => {
+    if (abortSignal?.aborted) {
+      resolve();
+      return;
+    }
+
+    let settled = false;
+    let onAbort = () => {};
+    const finish = () => {
+      if (settled) return;
+      settled = true;
+      abortSignal?.removeEventListener('abort', onAbort);
+      resolve();
+    };
+    const timer = setTimeout(finish, ms);
+    (timer as { unref?: () => void }).unref?.();
+    onAbort = () => {
+      clearTimeout(timer);
+      finish();
+    };
+    abortSignal?.addEventListener('abort', onAbort, { once: true });
+  });
+}