@falai/agent 2.4.3 → 2.6.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}`
@@ -103,38 +170,81 @@ export interface StreamRetryOptions {
   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* — e.g. an
- * empty completion that throws "No response", or an error while establishing
- * the stream. Once any chunk has been yielded the stream is committed and
- * further errors propagate, so a retry can never double-emit deltas that the
- * consumer has already received.
+ * 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: () => AsyncGenerator<T>,
+  factory: (signal: AbortSignal) => AsyncGenerator<T>,
   options: StreamRetryOptions = {}
 ): AsyncGenerator<T> {
   const {
     maxRetries = 3,
-    delay = (attempt: number) => Math.min(1000 * Math.pow(2, attempt), 5000),
+    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 {
-      for await (const chunk of factory()) {
+      // 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 chunk;
+        yield result.value;
       }
-      return;
     } catch (error: unknown) {
       // Can't retry once deltas are out, and don't retry past the budget.
       if (yielded || attempt === maxRetries) {
@@ -147,6 +257,18 @@ export async function* withStreamRetry<T>(
         `[${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);
+      }
     }
   }
 }

package/src/utils/streamingMessage.ts

@@ -0,0 +1,220 @@
+/**
+ * Incremental extraction of the `message` field from a streamed structured
+ * JSON response.
+ *
+ * Providers that enforce a JSON schema stream the raw wrapper object
+ * (`{"message":"Hel` → `{"message":"Hello"}`), not the message text. Without
+ * this, every streaming consumer has to re-implement partial-JSON unwrapping to
+ * recover the user-facing tokens. These helpers do it once, at the framework
+ * boundary, so `delta`/`accumulated` carry clean message text and the parsed
+ * object is surfaced only when complete.
+ *
+ * The extractor targets the top-level `message` string field specifically
+ * (tracking object depth and string context, so a nested decoy `"message"` key
+ * or a `"message"` substring inside another value is never mistaken for it),
+ * and is tolerant of truncation at any byte — a dangling escape sequence is
+ * held back rather than emitted half-decoded. Input that is not a JSON object
+ * (e.g. a provider streaming plain text) is passed through verbatim.
+ */
+
+const WHITESPACE = " \t\n\r";
+
+const ESCAPE_CHARS: Record<string, string> = {
+  '"': '"',
+  "\\": "\\",
+  "/": "/",
+  b: "\b",
+  f: "\f",
+  n: "\n",
+  r: "\r",
+  t: "\t",
+};
+
+interface StringRead {
+  /** Decoded value, or decoded prefix when the closing quote has not arrived. */
+  value: string;
+  /** Index just past the closing quote when complete; the input length otherwise. */
+  end: number;
+  /** Whether the closing quote was seen. */
+  complete: boolean;
+}
+
+/**
+ * Read a JSON string token whose opening quote is at `s[start]`. Decodes
+ * escapes. When the closing quote has not arrived, returns the decoded prefix
+ * with any trailing incomplete escape (`\` or a partial `\uXXXX`) held back, so
+ * a half-decoded character is never produced.
+ */
+function readJsonString(s: string, start: number): StringRead {
+  let out = "";
+  const n = s.length;
+  let i = start + 1; // skip opening quote
+
+  while (i < n) {
+    const ch = s[i];
+
+    if (ch === '"') {
+      return { value: out, end: i + 1, complete: true };
+    }
+
+    if (ch === "\\") {
+      const esc = s[i + 1];
+      if (esc === undefined) {
+        // Dangling backslash — wait for the rest of the escape.
+        return { value: out, end: n, complete: false };
+      }
+      if (esc === "u") {
+        if (i + 6 > n) {
+          // Incomplete \uXXXX — hold it back.
+          return { value: out, end: n, complete: false };
+        }
+        const code = parseInt(s.slice(i + 2, i + 6), 16);
+        if (Number.isNaN(code)) {
+          return { value: out, end: n, complete: false };
+        }
+        out += String.fromCharCode(code);
+        i += 6;
+      } else if (esc in ESCAPE_CHARS) {
+        out += ESCAPE_CHARS[esc];
+        i += 2;
+      } else {
+        // Not a valid JSON escape; pass the character through leniently.
+        out += esc;
+        i += 2;
+      }
+      continue;
+    }
+
+    out += ch;
+    i++;
+  }
+
+  return { value: out, end: n, complete: false };
+}
+
+interface ValueSkip {
+  end: number;
+  complete: boolean;
+}
+
+/**
+ * Skip one JSON value (string, object, array, or primitive) starting at `s[i]`.
+ * Tolerant of truncation: an unfinished value reports `complete: false`.
+ */
+function skipJsonValue(s: string, i: number): ValueSkip {
+  const n = s.length;
+  if (i >= n) return { end: n, complete: false };
+
+  const ch = s[i];
+
+  if (ch === '"') {
+    const r = readJsonString(s, i);
+    return { end: r.end, complete: r.complete };
+  }
+
+  if (ch === "{" || ch === "[") {
+    let depth = 0;
+    let j = i;
+    while (j < n) {
+      const c = s[j];
+      if (c === '"') {
+        const r = readJsonString(s, j);
+        if (!r.complete) return { end: n, complete: false };
+        j = r.end;
+        continue;
+      }
+      if (c === "{" || c === "[") depth++;
+      else if (c === "}" || c === "]") {
+        depth--;
+        if (depth === 0) return { end: j + 1, complete: true };
+      }
+      j++;
+    }
+    return { end: n, complete: false };
+  }
+
+  // Primitive (number, true, false, null): runs until a structural delimiter.
+  // Hitting end-of-input first means it may still be streaming.
+  let j = i;
+  while (j < n && !`,}]${WHITESPACE}`.includes(s[j])) j++;
+  return j < n ? { end: j, complete: true } : { end: n, complete: false };
+}
+
+/**
+ * Extract the decoded value of the top-level `message` string field from a
+ * (possibly partial) JSON object string, returning the text available so far.
+ *
+ * Returns `""` while the `message` value has not begun streaming, and passes
+ * `accumulated` through unchanged when it is not a JSON object.
+ */
+export function extractMessageSoFar(accumulated: string): string {
+  const s = accumulated;
+  const n = s.length;
+
+  let i = 0;
+  while (i < n && WHITESPACE.includes(s[i])) i++;
+
+  // Not a JSON object — a plain-text stream; emit verbatim.
+  if (i >= n || s[i] !== "{") return accumulated;
+  i++; // skip '{'
+
+  while (i < n) {
+    while (i < n && (WHITESPACE.includes(s[i]) || s[i] === ",")) i++;
+    if (i >= n) return "";
+    if (s[i] === "}") return ""; // object closed without a message
+    if (s[i] !== '"') return ""; // key not (fully) arrived
+
+    const key = readJsonString(s, i);
+    if (!key.complete) return ""; // key still streaming
+    i = key.end;
+
+    while (i < n && WHITESPACE.includes(s[i])) i++;
+    if (i >= n || s[i] !== ":") return ""; // colon not arrived
+    i++;
+    while (i < n && WHITESPACE.includes(s[i])) i++;
+    if (i >= n) return "";
+
+    if (key.value === "message") {
+      // Found it. Only a string value yields text; null/other → no message yet.
+      if (s[i] !== '"') return "";
+      return readJsonString(s, i).value;
+    }
+
+    // A field before `message`: skip its value. If it is still streaming we
+    // cannot have reached `message` yet.
+    const skipped = skipJsonValue(s, i);
+    if (!skipped.complete) return "";
+    i = skipped.end;
+  }
+
+  return "";
+}
+
+/**
+ * Stateful wrapper over {@link extractMessageSoFar} for a single stream: feed
+ * each chunk's accumulated JSON and get back the clean message-so-far plus the
+ * newly revealed delta.
+ *
+ * Each push re-scans the full accumulated buffer (O(n) per chunk, O(n²) over a
+ * stream) — deliberately kept simple and stateless: at LLM response sizes (KBs)
+ * the cost is negligible, and it avoids carrying cross-chunk parser/escape state.
+ */
+export class StreamingMessageDecoder {
+  private previous = "";
+
+  /**
+   * @param accumulated The provider chunk's full accumulated output so far.
+   * @returns `message` (clean text so far) and `delta` (the new text since the
+   *          previous push).
+   */
+  push(accumulated: string): { message: string; delta: string } {
+    const message = extractMessageSoFar(accumulated);
+    // Decoding is monotonic (each push extends the prefix); the guard is a
+    // belt-and-braces reset for any non-prefix anomaly.
+    const delta = message.startsWith(this.previous)
+      ? message.slice(this.previous.length)
+      : message;
+    this.previous = message;
+    return { message, delta };
+  }
+}