@graphorin/provider 0.6.1 → 0.7.0

package/src/middleware/with-fallback.ts

@@ -0,0 +1,157 @@
+/**
+ * `withFallback` - request-level fallback chain. When the wrapped
+ * provider raises a fallback-eligible error, the middleware retries
+ * the same request against each alternate provider in `fallbacks`.
+ * The model concept is preserved: callers wire `[primary, alt1, alt2]`
+ * pointing at providers serving the same logical model. Cross-model
+ * fallbacks are the agent-level concern (Phase 12).
+ *
+ * @packageDocumentation
+ */
+
+import type { Provider, ProviderEvent, ProviderResponse } from '@graphorin/core';
+
+import { isAbortError } from '../internal/abort.js';
+import { defineProviderMiddleware } from './compose.js';
+
+/**
+ * Options for {@link withFallback}.
+ *
+ * @stable
+ */
+export interface WithFallbackOptions {
+  /** Alternate providers tried in order. */
+  readonly fallbacks: ReadonlyArray<Provider>;
+  /** Predicate deciding whether an error should trigger a fallback. */
+  readonly shouldFallback?: (err: unknown) => boolean;
+  /** Optional log sink; defaults to `console.warn`. */
+  readonly logger?: (message: string, meta?: object) => void;
+}
+
+/**
+ * @stable
+ */
+export const withFallback = defineProviderMiddleware<WithFallbackOptions>({
+  kind: 'withFallback',
+  factory: (opts: WithFallbackOptions) => {
+    const shouldFallback = opts.shouldFallback ?? defaultShouldFallback;
+    const logger = opts.logger ?? defaultLogger;
+    if (!Array.isArray(opts.fallbacks) || opts.fallbacks.length === 0) {
+      throw new TypeError('withFallback: fallbacks array must contain at least one Provider.');
+    }
+    return (next: Provider): Provider => {
+      const candidates: ReadonlyArray<Provider> = [next, ...opts.fallbacks];
+      return {
+        name: next.name,
+        modelId: next.modelId,
+        capabilities: next.capabilities,
+        ...(next.acceptsSensitivity !== undefined
+          ? { acceptsSensitivity: next.acceptsSensitivity }
+          : {}),
+        stream(req) {
+          return fallbackStream(candidates, req, shouldFallback, logger);
+        },
+        async generate(req) {
+          let lastErr: unknown;
+          for (let i = 0; i < candidates.length; i++) {
+            const cand = candidates[i];
+            if (cand === undefined) continue;
+            try {
+              return await cand.generate(req);
+            } catch (err) {
+              if (req.signal?.aborted) throw err;
+              lastErr = err;
+              if (i === candidates.length - 1 || !shouldFallback(err)) throw err;
+              logger(`[graphorin/provider] withFallback: '${cand.name}' failed, trying next.`, {
+                error: (err as Error).message,
+              });
+            }
+          }
+          throw lastErr ?? new Error('withFallback: no fallback succeeded.');
+        },
+        ...(next.countTokens ? { countTokens: next.countTokens.bind(next) } : {}),
+      };
+    };
+  },
+});
+
+async function* fallbackStream(
+  candidates: ReadonlyArray<Provider>,
+  req: import('@graphorin/core').ProviderRequest,
+  shouldFallback: (err: unknown) => boolean,
+  logger: (message: string, meta?: object) => void,
+): AsyncIterable<ProviderEvent> {
+  let lastErr: unknown;
+  for (let i = 0; i < candidates.length; i++) {
+    const cand = candidates[i];
+    if (cand === undefined) continue;
+    let yieldedAny = false;
+    try {
+      for await (const event of cand.stream(req)) {
+        yieldedAny = true;
+        yield event;
+      }
+      return;
+    } catch (err) {
+      if (req.signal?.aborted) throw err;
+      lastErr = err;
+      if (yieldedAny) {
+        // Once we have streamed events, falling back would split the
+        // logical response across providers - propagate instead.
+        throw err;
+      }
+      if (i === candidates.length - 1 || !shouldFallback(err)) throw err;
+      logger(`[graphorin/provider] withFallback: '${cand.name}' failed, trying next.`, {
+        error: (err as Error).message,
+      });
+    }
+  }
+  throw lastErr ?? new Error('withFallback: no fallback succeeded.');
+}
+
+function defaultShouldFallback(err: unknown): boolean {
+  if (err === null || typeof err !== 'object') return false;
+  // An aborted request must not trigger a fallback, even as a `status: 0`
+  // network error (PS-2). The loop also short-circuits on `req.signal?.aborted`.
+  if (isAbortError(err)) return false;
+  const e = err as { kind?: string; errorKind?: string; status?: number };
+  // `ProviderHttpError.kind` is always 'provider-http' (stable discriminant);
+  // the mapped canonical kind rides on `errorKind`. Consult both so a 429
+  // fails over as a rate limit and a 400/context overflow never does.
+  const kinds = [e.kind, e.errorKind];
+  if (
+    kinds.includes('unauthorized') ||
+    kinds.includes('invalid-request') ||
+    kinds.includes('context-length') ||
+    kinds.includes('content-filter')
+  ) {
+    return false;
+  }
+  if (
+    kinds.includes('transient') ||
+    kinds.includes('rate-limit') ||
+    kinds.includes('rate-limit-exceeded') ||
+    kinds.includes('capacity')
+  ) {
+    return true;
+  }
+  // A 429 on the primary is exactly the case a same-model fallback chain
+  // exists for - the alternate deployment has its own rate budget.
+  if (typeof e.status === 'number' && e.status === 429) return true;
+  if (typeof e.status === 'number' && e.status >= 500) return true;
+  // PS-2: the headline fallback scenario - the primary provider is down
+  // (a local server refusing connections) surfaces as `status: 0`. Treat a
+  // network failure as fallback-eligible so the chain advances to the next
+  // provider (abort already excluded above).
+  if (typeof e.status === 'number' && e.status === 0) return true;
+  return false;
+}
+
+function defaultLogger(message: string, meta?: object): void {
+  if (meta !== undefined) console.warn(message, meta);
+  else console.warn(message);
+}
+
+// Cast to silence the "declared but unused" lint when only the types
+// are needed by consumers.
+export type _FallbackResponseGuard = ProviderResponse;

package/src/middleware/with-rate-limit.ts

@@ -0,0 +1,306 @@
+/**
+ * `withRateLimit` - token-bucket rate limiter keyed per `provider × model`.
+ * Configurable mode: `'throw'` (default) raises
+ * {@link RateLimitExceededError} on overflow; `'queue'` waits for the
+ * next available slot.
+ *
+ * W-145: an optional second dimension, `tokensPerMinute`, budgets
+ * MODEL tokens alongside requests. For agentic workloads the binding
+ * provider limit is TPM, not RPM - a 150k-token compacted transcript
+ * used to occupy the same single slot as a 200-token reranker call,
+ * so bursts of long-context steps sailed through the RPM gate and
+ * surfaced as provider 429s that only `withRetry` absorbed.
+ *
+ * @packageDocumentation
+ */
+
+import type { Provider, ProviderRequest } from '@graphorin/core';
+
+import { RateLimitExceededError } from '../errors/errors.js';
+import { defineProviderMiddleware } from './compose.js';
+
+/**
+ * Options for {@link withRateLimit}.
+ *
+ * @stable
+ */
+export interface WithRateLimitOptions {
+  /** Allowed requests per minute. */
+  readonly requestsPerMinute: number;
+  /** Burst size - defaults to `requestsPerMinute / 4` (rounded up to >= 1). */
+  readonly burst?: number;
+  /**
+   * Optional token budget per minute (W-145). When set, each request
+   * additionally reserves its estimated token weight from a second
+   * bucket whose capacity is the full minute budget; a request whose
+   * weight exceeds the remaining budget waits (queue mode) or throws
+   * with the TPM-aware `retryAfterMs` (throw mode) even when the RPM
+   * bucket has room. Unset: behaviour is byte-identical to the
+   * RPM-only limiter.
+   */
+  readonly tokensPerMinute?: number;
+  /**
+   * Estimator for a request's token weight (only consulted when
+   * `tokensPerMinute` is set). The default is the deliberate cheap
+   * heuristic `ceil(textChars / 4) + (maxTokens ?? 0)` - synchronous
+   * and allocation-free, because this runs in the request hot path and
+   * must not add latency or network. Wire the counter from
+   * `@graphorin/provider/counters` (`createDefaultCounter`) here when
+   * you need provider-accurate weights.
+   */
+  readonly estimateTokens?: (req: ProviderRequest) => number;
+  /** What to do on overflow. Default `'throw'`. */
+  readonly mode?: 'throw' | 'queue';
+  /** Test hook overriding `Date.now`. */
+  readonly nowImpl?: () => number;
+  /** Test hook overriding `setTimeout`-based wait. */
+  readonly sleepImpl?: (ms: number, signal?: AbortSignal) => Promise<void>;
+}
+
+interface Waiter {
+  resolve: () => void;
+  reject: (err: unknown) => void;
+  signal?: AbortSignal;
+  cancelled: boolean;
+  /** TPM weight this waiter must reserve (0 when TPM is unset). */
+  weight: number;
+}
+
+interface BucketState {
+  tokens: number;
+  lastRefill: number;
+  /** W-145: remaining TPM budget (only meaningful when TPM is set). */
+  tpmTokens: number;
+  tpmLastRefill: number;
+  /** FIFO waiters in `'queue'` mode (PS-18). */
+  queue: Waiter[];
+  /** Whether the single drain loop for this bucket is running. */
+  draining: boolean;
+}
+
+/**
+ * @stable
+ */
+export const withRateLimit = defineProviderMiddleware<WithRateLimitOptions>({
+  kind: 'withRateLimit',
+  factory: (opts: WithRateLimitOptions) => {
+    if (!(opts.requestsPerMinute > 0)) {
+      throw new RangeError('withRateLimit: requestsPerMinute must be > 0.');
+    }
+    if (opts.tokensPerMinute !== undefined && !(opts.tokensPerMinute > 0)) {
+      throw new RangeError('withRateLimit: tokensPerMinute must be > 0 when set.');
+    }
+    const burst = Math.max(opts.burst ?? Math.ceil(opts.requestsPerMinute / 4), 1);
+    const refillPerMs = opts.requestsPerMinute / 60_000;
+    const tpm = opts.tokensPerMinute;
+    const tpmRefillPerMs = tpm === undefined ? 0 : tpm / 60_000;
+    const estimate = opts.estimateTokens ?? estimateTokensHeuristic;
+    const mode = opts.mode ?? 'throw';
+    const now = opts.nowImpl ?? (() => Date.now());
+    const sleep = opts.sleepImpl ?? defaultSleep;
+    const buckets = new Map<string, BucketState>();
+    return (next: Provider): Provider => {
+      const key = `${next.name}::${next.modelId}`;
+      const bucketFor = (): BucketState => {
+        let bucket = buckets.get(key);
+        if (bucket === undefined) {
+          bucket = {
+            tokens: burst,
+            lastRefill: now(),
+            tpmTokens: tpm ?? 0,
+            tpmLastRefill: now(),
+            queue: [],
+            draining: false,
+          };
+          buckets.set(key, bucket);
+        }
+        return bucket;
+      };
+      const refill = (bucket: BucketState): void => {
+        const ts = now();
+        bucket.tokens = Math.min(burst, bucket.tokens + (ts - bucket.lastRefill) * refillPerMs);
+        bucket.lastRefill = ts;
+        if (tpm !== undefined) {
+          // TPM bucket capacity = the full minute budget, so any
+          // (clamped) weight is eventually satisfiable.
+          bucket.tpmTokens = Math.min(
+            tpm,
+            bucket.tpmTokens + (ts - bucket.tpmLastRefill) * tpmRefillPerMs,
+          );
+          bucket.tpmLastRefill = ts;
+        }
+      };
+      /**
+       * A request's TPM weight, clamped to the bucket capacity - a
+       * single request estimated above the whole minute budget must
+       * degrade to "wait for a full bucket", not deadlock the queue
+       * forever behind an unsatisfiable reservation.
+       */
+      const weightFor = (req: ProviderRequest): number => {
+        if (tpm === undefined) return 0;
+        const raw = estimate(req);
+        return Math.min(Math.max(0, Number.isFinite(raw) ? raw : tpm), tpm);
+      };
+      /** Both dimensions available for this weight? */
+      const available = (bucket: BucketState, weight: number): boolean =>
+        bucket.tokens >= 1 && (tpm === undefined || bucket.tpmTokens >= weight);
+      /** Wait until BOTH dimensions can serve `weight`, in ms. */
+      const waitMsFor = (bucket: BucketState, weight: number): number => {
+        const reqWait = bucket.tokens >= 1 ? 0 : (1 - bucket.tokens) / refillPerMs;
+        const tpmWait =
+          tpm === undefined || bucket.tpmTokens >= weight
+            ? 0
+            : (weight - bucket.tpmTokens) / tpmRefillPerMs;
+        return Math.ceil(Math.max(reqWait, tpmWait));
+      };
+      /** Deduct BOTH dimensions atomically (caller checked `available`). */
+      const take = (bucket: BucketState, weight: number): void => {
+        bucket.tokens -= 1;
+        if (tpm !== undefined) bucket.tpmTokens -= weight;
+      };
+      // PS-18: a single drain loop per bucket grants queued waiters ONE slot at
+      // a time, sleeping until the next token actually refills. Previously every
+      // concurrent waiter computed the same wait, slept once, and then all
+      // passed - letting N requests through on ~1 token (a burst). FIFO order is
+      // preserved and grants track the real refill schedule. W-145: a grant now
+      // requires BOTH the request token and the head's TPM weight.
+      const drain = (bucket: BucketState): void => {
+        if (bucket.draining) return;
+        bucket.draining = true;
+        void (async () => {
+          while (bucket.queue.length > 0) {
+            const head = bucket.queue[0];
+            if (head === undefined) break;
+            if (head.cancelled) {
+              bucket.queue.shift();
+              continue;
+            }
+            refill(bucket);
+            if (available(bucket, head.weight)) {
+              take(bucket, head.weight);
+              bucket.queue.shift();
+              head.resolve();
+              continue;
+            }
+            const waitMs = waitMsFor(bucket, head.weight);
+            try {
+              await sleep(waitMs);
+            } catch {
+              // The shared drain sleep carries no per-waiter signal; individual
+              // aborts are handled on the waiter via its listener.
+            }
+          }
+          bucket.draining = false;
+        })();
+      };
+      const acquire = async (signal?: AbortSignal, weight = 0): Promise<void> => {
+        if (signal?.aborted === true) throw signal.reason ?? new Error('aborted');
+        const bucket = bucketFor();
+        refill(bucket);
+        // Fast path: both dimensions free AND nobody queued ahead (FIFO).
+        if (available(bucket, weight) && bucket.queue.length === 0) {
+          take(bucket, weight);
+          return;
+        }
+        if (mode === 'throw') {
+          throw new RateLimitExceededError(waitMsFor(bucket, weight));
+        }
+        await new Promise<void>((resolve, reject) => {
+          const waiter: Waiter = {
+            resolve,
+            reject,
+            ...(signal ? { signal } : {}),
+            cancelled: false,
+            weight,
+          };
+          if (signal !== undefined) {
+            const onAbort = (): void => {
+              waiter.cancelled = true;
+              reject(signal.reason ?? new Error('aborted'));
+            };
+            signal.addEventListener('abort', onAbort, { once: true });
+            const settle = resolve;
+            waiter.resolve = (): void => {
+              signal.removeEventListener('abort', onAbort);
+              settle();
+            };
+          }
+          bucket.queue.push(waiter);
+          drain(bucket);
+        });
+      };
+      return {
+        name: next.name,
+        modelId: next.modelId,
+        capabilities: next.capabilities,
+        ...(next.acceptsSensitivity !== undefined
+          ? { acceptsSensitivity: next.acceptsSensitivity }
+          : {}),
+        stream(req) {
+          return rateLimitedStream(next, req, acquire, weightFor(req));
+        },
+        async generate(req) {
+          await acquire(req.signal, weightFor(req));
+          return next.generate(req);
+        },
+        ...(next.countTokens ? { countTokens: next.countTokens.bind(next) } : {}),
+      };
+    };
+  },
+});
+
+/**
+ * Default TPM weight heuristic: `ceil(textChars / 4) + (maxTokens ?? 0)`.
+ * Counts string contents and the `text` of content parts plus the
+ * system message; images / other parts are not weighed (heuristic by
+ * design - see {@link WithRateLimitOptions.estimateTokens}).
+ */
+function estimateTokensHeuristic(req: ProviderRequest): number {
+  let chars = req.systemMessage?.length ?? 0;
+  for (const message of req.messages) {
+    const content: unknown = message.content;
+    if (typeof content === 'string') {
+      chars += content.length;
+    } else if (Array.isArray(content)) {
+      for (const part of content) {
+        const text = (part as { readonly text?: unknown }).text;
+        if (typeof text === 'string') chars += text.length;
+      }
+    }
+  }
+  return Math.ceil(chars / 4) + (req.maxTokens ?? 0);
+}
+
+async function* rateLimitedStream(
+  next: Provider,
+  req: import('@graphorin/core').ProviderRequest,
+  acquire: (signal?: AbortSignal, weight?: number) => Promise<void>,
+  weight: number,
+): AsyncIterable<import('@graphorin/core').ProviderEvent> {
+  await acquire(req.signal, weight);
+  for await (const event of next.stream(req)) {
+    yield event;
+  }
+}
+
+function defaultSleep(ms: number, signal?: AbortSignal): Promise<void> {
+  return new Promise<void>((resolve, reject) => {
+    if (signal?.aborted) {
+      reject(signal.reason ?? new Error('aborted'));
+      return;
+    }
+    const timer = setTimeout(() => {
+      cleanup();
+      resolve();
+    }, ms);
+    const onAbort = () => {
+      cleanup();
+      reject(signal?.reason ?? new Error('aborted'));
+    };
+    function cleanup(): void {
+      clearTimeout(timer);
+      signal?.removeEventListener('abort', onAbort);
+    }
+    signal?.addEventListener('abort', onAbort, { once: true });
+  });
+}