@falai/agent 2.4.2 → 2.5.0

package/src/utils/retry.ts

@@ -11,6 +11,34 @@ export interface RetryOptions<T> {
   onFailure?: (error: unknown) => boolean;
 }
 
+/** Provider timeout (ms) + retry count, after defaults are applied. */
+export interface RetryConfig {
+  timeout: number;
+  retries: number;
+}
+
+const DEFAULT_RETRY_CONFIG: RetryConfig = {
+  timeout: 60000,
+  retries: 3,
+};
+
+/**
+ * Normalize a provider's optional retry config against the defaults. `timeout`
+ * uses `||` (a 0ms timeout is degenerate — it aborts every call immediately — so
+ * fall back to the default), while `retries` uses `??` so an explicit
+ * `retries: 0` (disable retries) is honored rather than treated as unset. Single
+ * definition so this distinction can't drift between providers.
+ */
+export function resolveRetryConfig(input?: {
+  timeout?: number;
+  retries?: number;
+}): RetryConfig {
+  return {
+    timeout: input?.timeout || DEFAULT_RETRY_CONFIG.timeout,
+    retries: input?.retries ?? DEFAULT_RETRY_CONFIG.retries,
+  };
+}
+
 export async function retry<T>(options: RetryOptions<T>): Promise<T> {
   const { operation, maxRetries, delay, onRetry, onFailure } = options;
 
@@ -43,8 +71,44 @@ export async function retry<T>(options: RetryOptions<T>): Promise<T> {
   throw lastError ?? new Error("Operation failed");
 }
 
+/** Capped exponential backoff (1s, 2s, 4s, … max 5s) shared by the retry helpers. */
+const defaultBackoff = (attempt: number): number =>
+  Math.min(1000 * Math.pow(2, attempt), 5000);
+
+/**
+ * Combine abort signals into one that aborts as soon as any input aborts.
+ * `undefined` inputs are skipped, so a caller can pass an optional caller-signal
+ * alongside a required per-attempt one. Returns the sole signal when only one is
+ * present (no wrapper), or `undefined` when none are.
+ *
+ * Prefers the platform `AbortSignal.any`, which cleans up its listeners via weak
+ * refs — so merging onto a long-lived caller signal that's reused across many
+ * calls can't accumulate listeners. Falls back to a manual controller on older
+ * runtimes (< Node 20.3); there the listeners live until a source aborts or is
+ * garbage-collected, which is fine for the usual per-request signal.
+ */
+export function combineAbortSignals(
+  ...signals: Array<AbortSignal | undefined>
+): AbortSignal | undefined {
+  const present = signals.filter((s): s is AbortSignal => s != null);
+  if (present.length <= 1) return present[0];
+  if (typeof AbortSignal.any === "function") return AbortSignal.any(present);
+
+  const controller = new AbortController();
+  for (const signal of present) {
+    if (signal.aborted) {
+      controller.abort(signal.reason);
+      break;
+    }
+    signal.addEventListener("abort", () => controller.abort(signal.reason), {
+      once: true,
+    });
+  }
+  return controller.signal;
+}
+
 export const withTimeoutAndRetry = async <T>(
-  operation: () => Promise<T>,
+  operation: (signal: AbortSignal) => Promise<T>,
   timeoutMs: number = 60000,
   maxRetries: number = 3,
   operationName: string = "AI operation"
@@ -56,8 +120,11 @@ export const withTimeoutAndRetry = async <T>(
     }, timeoutMs);
 
     try {
+      // Hand the operation the timeout signal so the in-flight upstream call is
+      // actually cancelled when the deadline fires — otherwise the abandoned
+      // attempt keeps running while the retry stacks a second concurrent call.
       const result = await Promise.race([
-        operation(),
+        operation(controller.signal),
         new Promise<never>((_, reject) => {
           controller.signal.addEventListener("abort", () => {
             reject(new Error(`Operation timed out after ${timeoutMs}ms`));
@@ -76,14 +143,14 @@ export const withTimeoutAndRetry = async <T>(
   return retry<T>({
     operation: createTimeoutOperation(),
     maxRetries,
-    delay: (attempt: number) => Math.min(1000 * Math.pow(2, attempt), 5000),
+    delay: defaultBackoff,
     onRetry: (attempt: number, error: unknown) => {
       const message = error instanceof Error ? error.message : String(error);
       logger.error(
         `[${operationName}] Failed attempt ${attempt + 1}:`,
         message
       );
-      const delay = Math.min(1000 * Math.pow(2, attempt), 5000);
+      const delay = defaultBackoff(attempt);
       logger.debug(`[${operationName}] Retrying in ${delay}ms...`);
       logger.debug(
         `[${operationName}] Attempt ${attempt + 2}/${maxRetries + 1}`
@@ -95,3 +162,113 @@ export const withTimeoutAndRetry = async <T>(
     },
   });
 };
+
+export interface StreamRetryOptions {
+  /** Maximum number of retries after the first attempt. Defaults to 3. */
+  maxRetries?: number;
+  /** Backoff before the next attempt, in ms. Defaults to capped exponential. */
+  delay?: (attempt: number) => number;
+  /** Label used in retry logs. */
+  operationName?: string;
+  /**
+   * Max ms to wait for the *first* chunk (time-to-first-token) before treating
+   * the attempt as failed. Guards a provider that opens a stream and then
+   * stalls. Only the first chunk is bounded — later chunks are unbounded so a
+   * long but healthy stream is never cut off. `0`/omitted disables the deadline.
+   */
+  firstChunkTimeoutMs?: number;
+}
+
+/** Reject if `next` hasn't settled within `timeoutMs`; always clears its timer. */
+function raceFirstChunk<T>(
+  next: Promise<IteratorResult<T>>,
+  timeoutMs: number,
+  operationName: string
+): Promise<IteratorResult<T>> {
+  let timer: ReturnType<typeof setTimeout> | undefined;
+  const timeout = new Promise<never>((_, reject) => {
+    timer = setTimeout(
+      () =>
+        reject(
+          new Error(
+            `[${operationName}] Stream timed out: no first chunk within ${timeoutMs}ms`
+          )
+        ),
+      timeoutMs
+    );
+  });
+  return Promise.race([next, timeout]).finally(() => clearTimeout(timer));
+}
+
+/**
+ * Streaming analog of {@link withTimeoutAndRetry}. Re-runs an async-generator
+ * factory as long as it fails *before yielding its first chunk* — an empty
+ * completion that throws "No response", an error while establishing the stream,
+ * or (with `firstChunkTimeoutMs`) a stall before the first token. Once any chunk
+ * has been yielded the stream is committed and further errors propagate, so a
+ * retry can never double-emit deltas the consumer has already received.
+ *
+ * This mirrors the non-streaming path, where the provider throws on an empty
+ * completion inside `withTimeoutAndRetry` and is retried on the same model
+ * before the caller falls through to backup models.
+ */
+export async function* withStreamRetry<T>(
+  factory: (signal: AbortSignal) => AsyncGenerator<T>,
+  options: StreamRetryOptions = {}
+): AsyncGenerator<T> {
+  const {
+    maxRetries = 3,
+    delay = defaultBackoff,
+    operationName = "AI stream",
+    firstChunkTimeoutMs,
+  } = options;
+
+  for (let attempt = 0; attempt <= maxRetries; attempt++) {
+    // One controller per attempt: aborting it on abandon cancels the upstream
+    // SDK call so a retry can't stack a second concurrent stream.
+    const controller = new AbortController();
+    const iterator = factory(controller.signal);
+    let yielded = false;
+    let completed = false;
+    try {
+      // Drive the iterator by hand so the first chunk can be raced against a
+      // deadline; for-await would give no hook for that.
+      for (;;) {
+        const result =
+          !yielded && firstChunkTimeoutMs
+            ? await raceFirstChunk(iterator.next(), firstChunkTimeoutMs, operationName)
+            : await iterator.next();
+        if (result.done) {
+          completed = true;
+          return;
+        }
+        yielded = true;
+        yield result.value;
+      }
+    } catch (error: unknown) {
+      // Can't retry once deltas are out, and don't retry past the budget.
+      if (yielded || attempt === maxRetries) {
+        throw error;
+      }
+      const message = error instanceof Error ? error.message : String(error);
+      logger.error(`[${operationName}] Failed attempt ${attempt + 1}:`, message);
+      const delayMs = delay(attempt);
+      logger.debug(
+        `[${operationName}] Retrying in ${delayMs}ms... (attempt ${attempt + 2}/${maxRetries + 1})`
+      );
+      await new Promise((resolve) => setTimeout(resolve, delayMs));
+    } finally {
+      // Abandoning an attempt (first-chunk timeout, pre-yield error before a
+      // retry, or a consumer that breaks early): abort the upstream SDK call
+      // first so it's actually torn down — not left running while we retry —
+      // then dispose the iterator. Aborting also unblocks a stalled
+      // iterator.next() so return() can run the generator's cleanup. Both are
+      // fire-and-forget: a wedged generator's return() may never settle, so we
+      // must not await it.
+      if (!completed) {
+        controller.abort();
+        void iterator.return?.(undefined)?.catch(() => undefined);
+      }
+    }
+  }
+}