@od-oneapp/ai-platform 0.1.0

package/src/sdk/middleware/index.ts

@@ -0,0 +1,660 @@
+/**
+ * Middleware System
+ *
+ * AI SDK v6 experimental middleware patterns:
+ * - Circuit breaker for provider failures
+ * - Rate limiting with token bucket
+ * - Provider fallback chains
+ * - Request/response transformation
+ *
+ * @module middleware
+ */
+
+import { logWarn } from '@repo/shared/logs';
+
+import { resolveModelId, type ModelRef } from '../../shared';
+
+/** Middleware function signature */
+export type MiddlewareFn<TContext = unknown> = (
+  ctx: MiddlewareContext<TContext>,
+  next: () => Promise<MiddlewareResult>,
+) => Promise<MiddlewareResult>;
+
+/** Model call options */
+export interface ModelCallOptions {
+  readonly prompt?: string;
+  readonly messages?: unknown[];
+  readonly maxOutputTokens?: number;
+  readonly temperature?: number;
+  readonly [key: string]: unknown;
+}
+
+/** Context passed through middleware chain */
+export interface MiddlewareContext<TContext = unknown> {
+  readonly model: ModelRef;
+  readonly options: ModelCallOptions;
+  readonly custom: TContext | undefined;
+  readonly attempt: number;
+  readonly startTime: number;
+}
+
+/** Result from middleware execution */
+export interface MiddlewareResult {
+  readonly success: boolean;
+  readonly response?: unknown;
+  readonly error?: Error;
+  readonly duration: number;
+  readonly metadata?: Record<string, unknown>;
+}
+
+/** Circuit breaker state */
+export type CircuitState = 'closed' | 'open' | 'half-open';
+
+/** Circuit breaker configuration */
+export interface CircuitBreakerConfig {
+  readonly failureThreshold: number;
+  readonly successThreshold: number;
+  readonly timeout: number;
+  readonly onStateChange?: (state: CircuitState, provider?: string) => void;
+}
+
+/** Rate limiter configuration */
+export interface RateLimiterConfig {
+  readonly tokensPerInterval: number;
+  readonly interval: number;
+  readonly maxBurst?: number;
+}
+
+/** Provider fallback configuration */
+export interface FallbackConfig {
+  readonly providers: ModelRef[];
+  readonly onFallback?: (from: string, to: string, error: Error) => void;
+}
+
+// -----------------------------------------------------------------------------
+// Circuit Breaker
+// -----------------------------------------------------------------------------
+
+interface CircuitBreakerState {
+  state: CircuitState;
+  failures: number;
+  successes: number;
+  lastFailure: number;
+  lastAttempt: number;
+}
+
+/**
+ * Creates a circuit breaker middleware.
+ *
+ * Prevents cascading failures by temporarily blocking requests
+ * to failing providers.
+ *
+ * @example
+ * ```typescript
+ * const breaker = createCircuitBreaker({
+ *   failureThreshold: 5,
+ *   successThreshold: 2,
+ *   timeout: 30_000
+ * });
+ * ```
+ */
+export function createCircuitBreaker(
+  config: CircuitBreakerConfig,
+): MiddlewareFn & { getState: (key?: string) => CircuitState; reset: (key?: string) => void } {
+  const { failureThreshold, successThreshold, timeout, onStateChange } = config;
+  const states = new Map<string, CircuitBreakerState>();
+
+  // Extended middleware with state access
+  type ExtendedMiddleware = MiddlewareFn & {
+    getState: (key?: string) => CircuitState;
+    reset: (key?: string) => void;
+  };
+
+  function getOrCreateState(key: string): CircuitBreakerState {
+    let state = states.get(key);
+    if (!state) {
+      state = {
+        state: 'closed',
+        failures: 0,
+        successes: 0,
+        lastFailure: 0,
+        lastAttempt: 0,
+      };
+      states.set(key, state);
+    }
+    return state;
+  }
+
+  function transitionTo(key: string, breaker: CircuitBreakerState, newState: CircuitState): void {
+    if (breaker.state !== newState) {
+      breaker.state = newState;
+      onStateChange?.(newState, key);
+    }
+  }
+
+  const middleware: MiddlewareFn = async (ctx, next) => {
+    const key = resolveModelId(ctx.model);
+    const breaker = getOrCreateState(key);
+    const now = Date.now();
+
+    // Check if circuit should transition from open to half-open
+    if (breaker.state === 'open' && now - breaker.lastFailure >= timeout) {
+      transitionTo(key, breaker, 'half-open');
+    }
+
+    // Block requests when circuit is open
+    if (breaker.state === 'open') {
+      return {
+        success: false,
+        error: new Error(`Circuit breaker open for ${key}`),
+        duration: 0,
+        metadata: { circuitState: 'open', provider: key },
+      };
+    }
+
+    breaker.lastAttempt = now;
+
+    // Wrap next() to catch thrown exceptions and record as failures
+    let result: MiddlewareResult;
+    try {
+      result = await next();
+    } catch (error) {
+      result = {
+        success: false,
+        error: error instanceof Error ? error : new Error(String(error)),
+        duration: Date.now() - ctx.startTime,
+      };
+    }
+
+    if (result.success) {
+      breaker.failures = 0;
+      breaker.successes++;
+
+      if (breaker.state === 'half-open' && breaker.successes >= successThreshold) {
+        transitionTo(key, breaker, 'closed');
+        breaker.successes = 0;
+      }
+    } else {
+      breaker.successes = 0;
+      breaker.failures++;
+      breaker.lastFailure = now;
+
+      if (breaker.failures >= failureThreshold) {
+        transitionTo(key, breaker, 'open');
+      }
+    }
+
+    return result;
+  };
+
+  const extendedMiddleware = middleware as ExtendedMiddleware;
+  extendedMiddleware.getState = (key = 'default') => getOrCreateState(key).state;
+  extendedMiddleware.reset = (key?: string) => {
+    if (key) {
+      states.delete(key);
+    } else {
+      states.clear();
+    }
+  };
+
+  return extendedMiddleware;
+}
+
+// -----------------------------------------------------------------------------
+// Rate Limiter (Token Bucket)
+// -----------------------------------------------------------------------------
+
+interface TokenBucket {
+  tokens: number;
+  lastRefill: number;
+}
+
+/**
+ * Creates a rate limiter middleware using token bucket algorithm.
+ *
+ * @example
+ * ```typescript
+ * const limiter = createRateLimiter({
+ *   tokensPerInterval: 100,
+ *   interval: 60_000, // 1 minute
+ *   maxBurst: 150
+ * });
+ * ```
+ */
+export function createRateLimiter(
+  config: RateLimiterConfig,
+): MiddlewareFn & { getAvailableTokens: (key?: string) => number } {
+  const { tokensPerInterval, interval, maxBurst = tokensPerInterval } = config;
+  const buckets = new Map<string, TokenBucket>();
+
+  type ExtendedRateLimiter = MiddlewareFn & {
+    getAvailableTokens: (key?: string) => number;
+  };
+
+  function getOrCreateBucket(key: string): TokenBucket {
+    let bucket = buckets.get(key);
+    if (!bucket) {
+      bucket = { tokens: tokensPerInterval, lastRefill: Date.now() };
+      buckets.set(key, bucket);
+    }
+    return bucket;
+  }
+
+  function refillBucket(bucket: TokenBucket): void {
+    const now = Date.now();
+    const elapsed = now - bucket.lastRefill;
+    const refillTokens = Math.floor((elapsed / interval) * tokensPerInterval);
+
+    if (refillTokens > 0) {
+      bucket.tokens = Math.min(maxBurst, bucket.tokens + refillTokens);
+      bucket.lastRefill = now;
+    }
+  }
+
+  const middleware: MiddlewareFn = async (ctx, next) => {
+    const key = resolveModelId(ctx.model);
+    const bucket = getOrCreateBucket(key);
+
+    refillBucket(bucket);
+
+    if (bucket.tokens < 1) {
+      const waitTime = Math.ceil(interval / tokensPerInterval);
+      return {
+        success: false,
+        error: new Error(`Rate limited for ${key}. Retry after ${waitTime}ms`),
+        duration: 0,
+        metadata: { rateLimited: true, retryAfter: waitTime },
+      };
+    }
+
+    bucket.tokens--;
+    return next();
+  };
+
+  const extendedMiddleware = middleware as ExtendedRateLimiter;
+  extendedMiddleware.getAvailableTokens = (key = 'default') => {
+    const bucket = getOrCreateBucket(key);
+    refillBucket(bucket);
+    return bucket.tokens;
+  };
+
+  return extendedMiddleware;
+}
+
+// -----------------------------------------------------------------------------
+// Provider Fallback
+// -----------------------------------------------------------------------------
+
+/**
+ * Creates a provider fallback middleware.
+ *
+ * Automatically falls back to alternative providers when primary fails.
+ *
+ * **IMPORTANT**: Due to pipeline architecture, this middleware should be positioned
+ * as the LAST middleware in the pipeline (closest to the terminal executor) for
+ * fallback to work correctly with each provider. Alternatively, provide an `executor`
+ * function that will be called for each provider attempt.
+ *
+ * @example
+ * ```typescript
+ * // Option 1: Position as last middleware
+ * const pipeline = createPipeline(
+ *   loggingMiddleware,
+ *   createFallbackMiddleware({
+ *     providers: [openai('gpt-4'), anthropic('claude-3')],
+ *     onFallback: (from, to, error) => console.log(`Falling back from ${from} to ${to}`)
+ *   }) // LAST - closest to terminal
+ * );
+ *
+ * // Option 2: Provide custom executor
+ * const fallback = createFallbackMiddleware({
+ *   providers: [openai('gpt-4'), anthropic('claude-3')],
+ *   executor: async (ctx) => executeModelCall(ctx)
+ * });
+ * ```
+ */
+export function createFallbackMiddleware(
+  config: FallbackConfig & {
+    /** Optional executor for each provider attempt. If provided, this is called instead of next(). */
+    executor?: (ctx: MiddlewareContext) => Promise<MiddlewareResult>;
+  },
+): MiddlewareFn {
+  const { providers, onFallback, executor } = config;
+
+  return async (ctx, next) => {
+    let lastError: Error | undefined;
+
+    for (let i = 0; i < providers.length; i++) {
+      const provider = providers[i];
+      if (!provider) continue;
+
+      const isOriginal = i === 0;
+      const previousProvider = isOriginal ? null : providers[i - 1];
+
+      if (!isOriginal && lastError && previousProvider) {
+        onFallback?.(resolveModelId(previousProvider), resolveModelId(provider), lastError);
+      }
+
+      // Create new context with current provider
+      const newCtx: MiddlewareContext = {
+        ...ctx,
+        model: provider,
+        attempt: i + 1,
+      };
+
+      // Execute with current provider
+      // Use executor if provided, otherwise use next()
+      // Note: Without executor, subsequent iterations reuse the same next()
+      // which advances the pipeline rather than re-executing downstream.
+      // This only works correctly when fallback is the last middleware.
+      if (!executor && process.env.NODE_ENV !== 'production') {
+        logWarn(
+          '[createFallbackMiddleware] Warning: Fallback middleware should be positioned last or provide an executor option. Without executor, fallback behavior may not work correctly if other middleware follows.',
+        );
+      }
+      const result = executor
+        ? await executeWithModel(newCtx, () => executor(newCtx))
+        : await executeWithModel(newCtx, next);
+
+      if (result.success) {
+        return {
+          ...result,
+          metadata: { ...result.metadata, provider: resolveModelId(provider), attempt: i + 1 },
+        };
+      }
+
+      lastError = result.error;
+    }
+
+    return {
+      success: false,
+      error: lastError ?? new Error('All providers failed'),
+      duration: Date.now() - ctx.startTime,
+      metadata: { attempts: providers.length },
+    };
+  };
+}
+
+async function executeWithModel(
+  ctx: MiddlewareContext,
+  next: () => Promise<MiddlewareResult>,
+): Promise<MiddlewareResult> {
+  try {
+    return await next();
+  } catch (error) {
+    return {
+      success: false,
+      error: error instanceof Error ? error : new Error(String(error)),
+      duration: Date.now() - ctx.startTime,
+    };
+  }
+}
+
+// -----------------------------------------------------------------------------
+// Middleware Pipeline
+// -----------------------------------------------------------------------------
+
+/**
+ * Creates a middleware pipeline that executes middlewares in order.
+ *
+ * @example
+ * ```typescript
+ * const pipeline = createPipeline(
+ *   createRateLimiter({ tokensPerInterval: 100, interval: 60_000 }),
+ *   createCircuitBreaker({ failureThreshold: 5, successThreshold: 2, timeout: 30_000 }),
+ *   loggingMiddleware
+ * );
+ *
+ * const result = await pipeline.execute(model, options);
+ * ```
+ */
+export function createPipeline<TContext = unknown>(
+  ...middlewares: MiddlewareFn<TContext>[]
+): {
+  execute: (
+    model: ModelRef,
+    options: ModelCallOptions,
+    custom?: TContext,
+  ) => Promise<MiddlewareResult>;
+  use: (middleware: MiddlewareFn<TContext>) => void;
+} {
+  const stack = [...middlewares];
+
+  function execute(
+    model: ModelRef,
+    options: ModelCallOptions,
+    custom?: TContext,
+  ): Promise<MiddlewareResult> {
+    const ctx: MiddlewareContext<TContext> = {
+      model,
+      options,
+      custom,
+      attempt: 1,
+      startTime: Date.now(),
+    };
+
+    let index = 0;
+
+    const next = async (): Promise<MiddlewareResult> => {
+      if (index >= stack.length) {
+        // Terminal: execute actual model call
+        return executeModelCall(ctx);
+      }
+
+      const middleware = stack[index++];
+      if (!middleware) {
+        return next();
+      }
+
+      return middleware(ctx, next);
+    };
+
+    return next();
+  }
+
+  return {
+    execute,
+    use: (middleware: MiddlewareFn<TContext>) => stack.push(middleware),
+  };
+}
+
+async function executeModelCall(ctx: MiddlewareContext): Promise<MiddlewareResult> {
+  const startTime = Date.now();
+
+  try {
+    // This would be replaced with actual model invocation
+    // For now, return success placeholder
+    return {
+      success: true,
+      duration: Date.now() - startTime,
+      metadata: { model: resolveModelId(ctx.model) },
+    };
+  } catch (error) {
+    return {
+      success: false,
+      error: error instanceof Error ? error : new Error(String(error)),
+      duration: Date.now() - startTime,
+    };
+  }
+}
+
+// -----------------------------------------------------------------------------
+// Utility Middlewares
+// -----------------------------------------------------------------------------
+
+/**
+ * Creates a logging middleware.
+ */
+export function createLoggingMiddleware(
+  logger: { log: (message: string, data?: unknown) => void } = console,
+): MiddlewareFn {
+  return async (ctx, next) => {
+    const id = Math.random().toString(36).slice(2, 8);
+    logger.log(`[${id}] Request to ${resolveModelId(ctx.model)}`, { attempt: ctx.attempt });
+
+    const result = await next();
+
+    logger.log(`[${id}] Response from ${resolveModelId(ctx.model)}`, {
+      success: result.success,
+      duration: result.duration,
+    });
+
+    return result;
+  };
+}
+
+/**
+ * Creates a timeout middleware.
+ */
+export function createTimeoutMiddleware(timeoutMs: number): MiddlewareFn {
+  return async (ctx, next) => {
+    const controller = new AbortController();
+    const timeoutId = setTimeout(() => controller.abort(), timeoutMs);
+
+    try {
+      const result = await Promise.race([
+        next(),
+        new Promise<MiddlewareResult>((resolve, reject) => {
+          void resolve;
+          controller.signal.addEventListener('abort', () => {
+            reject(new Error(`Request timeout after ${timeoutMs}ms`));
+          });
+        }),
+      ]);
+      return result;
+    } catch (error) {
+      return {
+        success: false,
+        error: error instanceof Error ? error : new Error(String(error)),
+        duration: Date.now() - ctx.startTime,
+        metadata: { timeout: true },
+      };
+    } finally {
+      clearTimeout(timeoutId);
+    }
+  };
+}
+
+/**
+ * Creates a retry middleware with exponential backoff.
+ *
+ * **IMPORTANT**: Due to pipeline architecture, this middleware should be positioned
+ * as the LAST middleware in the pipeline (closest to the terminal executor) for
+ * retries to work correctly. Alternatively, provide an `executor` function that
+ * will be called for each retry attempt.
+ *
+ * @example
+ * ```typescript
+ * // Option 1: Position as last middleware
+ * const pipeline = createPipeline(
+ *   loggingMiddleware,
+ *   rateLimiter,
+ *   createRetryMiddleware({ maxRetries: 3 }) // LAST - closest to terminal
+ * );
+ *
+ * // Option 2: Provide custom executor for true retries
+ * const retry = createRetryMiddleware({
+ *   maxRetries: 3,
+ *   executor: async (ctx) => executeModelCall(ctx)
+ * });
+ * ```
+ */
+export function createRetryMiddleware(options: {
+  maxRetries: number;
+  baseDelay?: number;
+  maxDelay?: number;
+  shouldRetry?: (error: Error) => boolean;
+  /** Optional executor for true retry behavior. If provided, this is called for each retry. */
+  executor?: (ctx: MiddlewareContext) => Promise<MiddlewareResult>;
+}): MiddlewareFn {
+  const {
+    maxRetries,
+    baseDelay = 1000,
+    maxDelay = 30_000,
+    shouldRetry = () => true,
+    executor,
+  } = options;
+
+  return async (ctx, next) => {
+    let lastError: Error | undefined;
+    // Track if this is the first attempt (use next), or subsequent (use executor if provided)
+    let isFirstAttempt = true;
+
+    for (let attempt = 0; attempt <= maxRetries; attempt++) {
+      if (attempt > 0) {
+        const delay = Math.min(baseDelay * 2 ** (attempt - 1), maxDelay);
+        await new Promise(resolve => setTimeout(resolve, delay));
+      }
+
+      // For first attempt, always use next() to continue pipeline
+      // For subsequent attempts, use executor if provided, otherwise log warning and continue
+      let result: MiddlewareResult;
+      if (isFirstAttempt) {
+        result = await next();
+        isFirstAttempt = false;
+      } else if (executor) {
+        result = await executor({ ...ctx, attempt: attempt + 1 });
+      } else {
+        // Without executor, subsequent next() calls won't retry - they advance pipeline
+        // This is expected when retry is positioned as last middleware
+        result = await next();
+      }
+
+      if (result.success) {
+        return { ...result, metadata: { ...result.metadata, attempts: attempt + 1 } };
+      }
+
+      if (result.error && !shouldRetry(result.error)) {
+        return result;
+      }
+
+      lastError = result.error;
+    }
+
+    return {
+      success: false,
+      error: lastError ?? new Error('Max retries exceeded'),
+      duration: Date.now() - ctx.startTime,
+      metadata: { maxRetriesExceeded: true, attempts: maxRetries + 1 },
+    };
+  };
+}
+
+// -----------------------------------------------------------------------------
+// Model Middleware (AI SDK v6 experimental)
+// -----------------------------------------------------------------------------
+
+export {
+  // Core wrapper
+  wrapLanguageModel,
+  getModelMiddleware,
+  // Pre-built middleware
+  extractReasoningMiddleware,
+  addToolInputExamplesMiddleware,
+  // Composition
+  composeMiddleware,
+  // Utility middleware (named differently to avoid conflict with pipeline middleware)
+  createModelLoggingMiddleware,
+  createModelCachingMiddleware,
+  // Types
+  type WrapLanguageModelConfig,
+  type LanguageModelMiddleware,
+  type ExtractReasoningConfig,
+  type MiddlewareToolExample,
+  type AddToolInputExamplesConfig,
+  type MiddlewareTransformResult,
+} from './model-middleware';
+
+// Circuit breaker (standalone class-based implementation from @repo/ai)
+export {
+  CircuitBreaker,
+  CircuitBreakerRegistry,
+  CircuitOpenError,
+  circuitBreakerRegistry,
+  withCircuitBreaker,
+  type CircuitBreakerConfig as StandaloneCircuitBreakerConfig,
+  type CircuitMetrics,
+  type FallbackBehavior,
+  // Note: CircuitState from circuit-breaker.ts is identical to the one above.
+  // Not re-exported here to avoid conflict.
+} from './circuit-breaker';