From 0b624118ddf8b6b17c7bffb6f90b4c39bb35b221 Mon Sep 17 00:00:00 2001 From: Manvendra Date: Sat, 27 Jun 2026 20:22:51 +0530 Subject: [PATCH] feat: routing, fallback & rate-limit survival (Phase 4) Turn every request into an ordered candidate chain run through a resilience executor, so a flaky provider or a free-tier rate limit no longer surfaces as a client error. - routing/router.ts: builds the candidate chain. Explicit models become [model, ...fallback]; `model: "auto"` is classified into the cheapest capable tier (routing/classifier.ts) and escalates through the rest. - routing/executor.ts: per-candidate retry with exponential backoff on retryable errors (routing/retryable.ts), fail-over to the next candidate, per-attempt AbortController timeout, and a final 503 when all candidates are throttled. Streaming fail-over is bounded to the pre-hijack first chunk. - throttle/token-bucket.ts: in-memory per-provider token bucket (injectable clock) that paces outbound calls to each provider's rpm. - Traces gain routedProvider/routedModel/fallbackUsed/retryCount across trace.ts, the SQLite + in-memory stores, and the /traces query API (?fallbackUsed=true, ?routedProvider=). - config.ts: optional `routing` block (tiers/fallback) + per-provider `rpm`; REQUEST_TIMEOUT_MS / MAX_RETRIES / DEFAULT_RPM / THROTTLE_MAX_WAIT_MS. - No routing configured = unchanged single-provider pass-through. Docs: README "Routing & resilience", .env.example, config example, SECURITY_REVIEW_LOG SR-004. 134 tests, 98% coverage; pnpm verify green. --- .env.example | 11 + README.md | 20 +- SECURITY_REVIEW_LOG.md | 15 +- packages/gateway/src/config.ts | 31 ++- packages/gateway/src/index.ts | 23 +- packages/gateway/src/main.ts | 14 ++ packages/gateway/src/routes.traces.ts | 2 + .../gateway/src/routing/classifier.test.ts | 37 ++++ packages/gateway/src/routing/classifier.ts | 32 +++ packages/gateway/src/routing/executor.test.ts | 163 +++++++++++++++ packages/gateway/src/routing/executor.ts | 151 ++++++++++++++ .../gateway/src/routing/retryable.test.ts | 37 ++++ packages/gateway/src/routing/retryable.ts | 22 ++ packages/gateway/src/routing/router.test.ts | 89 ++++++++ packages/gateway/src/routing/router.ts | 81 ++++++++ packages/gateway/src/server.test.ts | 196 ++++++++++++++++++ packages/gateway/src/server.ts | 60 +++++- .../gateway/src/telemetry/store.memory.ts | 3 + .../gateway/src/telemetry/store.sqlite.ts | 31 ++- packages/gateway/src/telemetry/store.test.ts | 21 ++ packages/gateway/src/telemetry/trace.ts | 14 ++ .../gateway/src/throttle/token-bucket.test.ts | 70 +++++++ packages/gateway/src/throttle/token-bucket.ts | 95 +++++++++ sentinel.config.example.json | 9 +- 24 files changed, 1203 insertions(+), 24 deletions(-) create mode 100644 packages/gateway/src/routing/classifier.test.ts create mode 100644 packages/gateway/src/routing/classifier.ts create mode 100644 packages/gateway/src/routing/executor.test.ts create mode 100644 packages/gateway/src/routing/executor.ts create mode 100644 packages/gateway/src/routing/retryable.test.ts create mode 100644 packages/gateway/src/routing/retryable.ts create mode 100644 packages/gateway/src/routing/router.test.ts create mode 100644 packages/gateway/src/routing/router.ts create mode 100644 packages/gateway/src/throttle/token-bucket.test.ts create mode 100644 packages/gateway/src/throttle/token-bucket.ts diff --git a/.env.example b/.env.example index 7d38b81..b67c719 100644 --- a/.env.example +++ b/.env.example @@ -22,6 +22,17 @@ TRACE_DB_PATH=./traces.db # Optional: also export OpenTelemetry spans to an OTLP/HTTP collector (e.g. Jaeger). # OTEL_EXPORTER_OTLP_ENDPOINT=http://localhost:4318/v1/traces +# ── Routing, fallback & rate-limit survival ────────────────────────── +# Per-attempt timeout before a candidate is aborted and retried/failed over. +REQUEST_TIMEOUT_MS=30000 +# Retries per candidate on a retryable error (429/5xx/timeout) before failing over. +MAX_RETRIES=2 +# Default per-provider requests-per-minute throttle (0 = unlimited). Per-provider +# overrides and the "auto" tier list / fallback chain live in sentinel.config.json. +DEFAULT_RPM=0 +# Max time the throttle may pace a single candidate before skipping to the next. +THROTTLE_MAX_WAIT_MS=2000 + # ── Semantic cache (embeds prompts locally; serves similar repeats) ─── CACHE_ENABLED=true # Cosine-similarity threshold for a cache hit (0–1; higher = stricter). diff --git a/README.md b/README.md index 6baae67..244e2b0 100644 --- a/README.md +++ b/README.md @@ -2,7 +2,7 @@ A self-hostable **verifying LLM gateway** — a drop-in, OpenAI-compatible proxy that **routes** (cheapest capable model + fallback), **semantically caches**, and **verifies** (deterministic guardrails inline + a local Ollama judge) every LLM call, with full OpenTelemetry tracing. Unlike after-the-fact observability tools, it can flag or block a bad response _before it returns_. -> 🚧 **Early development.** Product spec in [`PRP_SPEC.md`](./PRP_SPEC.md), phased build in [`ROADMAP.md`](./ROADMAP.md), contributor/agent guidance in [`CLAUDE.md`](./CLAUDE.md). Currently at **Phase 3 — semantic cache** (routing/fallback and verification land in later phases). +> 🚧 **Early development.** Product spec in [`PRP_SPEC.md`](./PRP_SPEC.md), phased build in [`ROADMAP.md`](./ROADMAP.md), contributor/agent guidance in [`CLAUDE.md`](./CLAUDE.md). Currently at **Phase 4 — routing, fallback & rate-limit survival** (inline verification lands in later phases). ## What works today (Phase 1) @@ -73,11 +73,15 @@ Each provider sets a `baseUrl` (or `baseUrlEnv` to read it from the environment) } }, "models": { "llama3.2": "ollama", "gpt-4o-mini": "openai", "llama-3.3-70b-versatile": "groq" }, - "defaultProvider": "ollama" + "defaultProvider": "ollama", + "routing": { + "tiers": ["llama3.2", "gpt-4o-mini", "llama-3.3-70b-versatile"], + "fallback": ["llama3.2"] + } } ``` -Because almost every provider speaks the OpenAI API, **bringing your own key is just config** — add a provider entry and a model route, drop the key in `.env`, done. +A provider may also set an `rpm` (requests-per-minute) ceiling; the optional `routing` block configures cost-aware routing (see below). Because almost every provider speaks the OpenAI API, **bringing your own key is just config** — add a provider entry and a model route, drop the key in `.env`, done. ### Bring your own Ollama @@ -101,6 +105,16 @@ Traces are **metadata only** — no prompt or response bodies are stored, and AP Sentinel **semantically caches** responses: it embeds each prompt locally (Ollama `nomic-embed-text`) and, when a new request is similar enough to a recent one (cosine ≥ `CACHE_SIMILARITY_THRESHOLD`, default `0.92`), serves the stored answer **without calling the provider** — replaying buffered SSE chunks for streamed requests. The cache is **per-tenant** (scoped to the calling API key), bounded (`CACHE_MAX_ENTRIES`, `CACHE_TTL_SECONDS`), and **fails open** — any embedding error simply falls through to the provider. Cache hits are visible in traces (`GET /traces?cacheHit=true`). Disable with `CACHE_ENABLED=false`. +## Routing & resilience + +Sentinel survives flaky providers and free-tier rate limits instead of passing their failures back to your app. Each request becomes an ordered **candidate chain** that a resilience executor runs in order: + +- **Retry + backoff + fallback.** A retryable failure (429, upstream 5xx, timeout, network fault) is retried with exponential backoff up to `MAX_RETRIES`, then **failed over** to the next candidate — additional models you list under `routing.fallback`, ending at a local Ollama model so a request can always be served. Terminal errors (400/401/403/404) fail fast without pointless fallback. +- **Per-provider throttling.** A token-bucket limiter paces each provider to its configured `rpm` (or `DEFAULT_RPM`), waiting up to `THROTTLE_MAX_WAIT_MS` for headroom before skipping to a less-loaded candidate. This keeps you under a provider's limit rather than discovering it the hard way. +- **Cost-aware routing (opt-in).** Send `"model": "auto"` and a rules-based classifier picks the **cheapest capable tier** from `routing.tiers` (cheapest first) by prompt complexity, escalating to more capable tiers on failure. Explicit model names behave exactly as before — plus the fallback chain. Drop-in semantics are preserved; nothing is configured by default. + +Every routed request records which provider/model actually served it, whether a fallback was used, and the retry count — queryable via `GET /traces?fallbackUsed=true` (and `?routedProvider=`). Streaming requests fail over up to the first chunk; once the SSE response is committed, a mid-stream error is surfaced as an inline event. + ## Development ```bash diff --git a/SECURITY_REVIEW_LOG.md b/SECURITY_REVIEW_LOG.md index eaa1ff3..2115e9e 100644 --- a/SECURITY_REVIEW_LOG.md +++ b/SECURITY_REVIEW_LOG.md @@ -49,7 +49,7 @@ Each item: assume an attacker is actively trying it. Tick only when a test prove ### 7. Availability / DoS - [ ] Sentinel's own rate-limits/quotas protect it from a runaway client. -- [ ] Bounded queues/timeouts; a slow provider cannot exhaust the event loop or memory. +- [x] Bounded queues/timeouts; a slow provider cannot exhaust the event loop or memory. ### 8. Supply chain @@ -64,11 +64,12 @@ Each item: assume an attacker is actively trying it. Tick only when a test prove ## Review log -| ID | Date | Area (1–8) | Change / PR | Risk | Severity | Status | Mitigation / notes | -| ------ | ---------- | ---------- | ------------------------------ | ---------------------------------------------------------------------------- | -------- | --------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | -| — | 2026-06-27 | — | Initial context docs (no code) | n/a | n/a | n/a | Baseline. | -| SR-001 | 2026-06-27 | 1, 3, 4 | Phase 1 pass-through proxy | Unauthenticated access; key/secret leakage; SSRF via request-controlled URLs | high | mitigated | Bearer Sentinel-key auth required — 401 on missing/invalid (tested). Provider keys read from env via `apiKeyEnv`, never hard-coded, never returned to clients. Provider base URLs come only from the config file, never request input. `authorization`/`x-api-key` redacted in request logs via pino `redact` (configured; an explicit log-assertion test is still TODO — box 1.2 left unticked). | -| SR-002 | 2026-06-27 | 3, 6 | Phase 2 tracing & persistence | Unauthorized trace access; secrets/PII in stored traces | high | mitigated | `GET /traces` gated by a separate `SENTINEL_ADMIN_KEY` — 401 without it (tested), distinct from client keys. Traces are metadata-only (model, provider, tokens, latency, status) — no prompt/response bodies persisted. API keys stored as a SHA-256 hash, never raw. Per-key trace scoping deferred (admin-only for now). | -| SR-003 | 2026-06-27 | 5 | Phase 3 semantic cache | Cross-tenant cache leakage; wrong-answer collisions | high | mitigated | Cache entries are bucketed per-tenant by API-key hash — no cross-tenant hits (tested). The bucket also keys on model/temperature/max_tokens/stream/embed-model, so requests that change the answer never collide; semantic matching happens only within a bucket above a conservative threshold (0.92, configurable). Cache fails open (embed errors → miss). | +| ID | Date | Area (1–8) | Change / PR | Risk | Severity | Status | Mitigation / notes | +| ------ | ---------- | ---------- | ------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------ | -------- | --------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| — | 2026-06-27 | — | Initial context docs (no code) | n/a | n/a | n/a | Baseline. | +| SR-001 | 2026-06-27 | 1, 3, 4 | Phase 1 pass-through proxy | Unauthenticated access; key/secret leakage; SSRF via request-controlled URLs | high | mitigated | Bearer Sentinel-key auth required — 401 on missing/invalid (tested). Provider keys read from env via `apiKeyEnv`, never hard-coded, never returned to clients. Provider base URLs come only from the config file, never request input. `authorization`/`x-api-key` redacted in request logs via pino `redact` (configured; an explicit log-assertion test is still TODO — box 1.2 left unticked). | +| SR-002 | 2026-06-27 | 3, 6 | Phase 2 tracing & persistence | Unauthorized trace access; secrets/PII in stored traces | high | mitigated | `GET /traces` gated by a separate `SENTINEL_ADMIN_KEY` — 401 without it (tested), distinct from client keys. Traces are metadata-only (model, provider, tokens, latency, status) — no prompt/response bodies persisted. API keys stored as a SHA-256 hash, never raw. Per-key trace scoping deferred (admin-only for now). | +| SR-003 | 2026-06-27 | 5 | Phase 3 semantic cache | Cross-tenant cache leakage; wrong-answer collisions | high | mitigated | Cache entries are bucketed per-tenant by API-key hash — no cross-tenant hits (tested). The bucket also keys on model/temperature/max_tokens/stream/embed-model, so requests that change the answer never collide; semantic matching happens only within a bucket above a conservative threshold (0.92, configurable). Cache fails open (embed errors → miss). | +| SR-004 | 2026-06-27 | 7 | Phase 4 routing & fallback | Self-inflicted DoS: unbounded retries/fallback amplify load; a slow or hostile provider stalls the event loop; throttle/router as a new untrusted-input path | med | mitigated | Each attempt is bounded by a per-attempt `AbortController` timeout (`REQUEST_TIMEOUT_MS`) so a slow provider can't hang the loop (tested). Retries are capped (`MAX_RETRIES`) and fallback walks a finite, **config-defined** candidate chain — request input never adds providers/URLs (no new SSRF surface). A per-provider token bucket paces outbound calls to each provider's `rpm`, protecting upstreams and avoiding self-induced 429 storms. Terminal 4xx errors fail fast without amplification. Inbound per-client rate-limiting (box 7.1) is still deferred. | > Add a row per high-risk change. Status ∈ {open, mitigated, accepted}. Severity ∈ {low, med, high, critical}. diff --git a/packages/gateway/src/config.ts b/packages/gateway/src/config.ts index f66a645..046d8ac 100644 --- a/packages/gateway/src/config.ts +++ b/packages/gateway/src/config.ts @@ -17,6 +17,10 @@ export interface ServerEnv { cacheMaxEntries: number; ollamaBaseUrl: string; embedModel: string; + requestTimeoutMs: number; + maxRetries: number; + defaultRpm: number; + throttleMaxWaitMs: number; } const serverEnvSchema = z.object({ @@ -35,6 +39,10 @@ const serverEnvSchema = z.object({ CACHE_MAX_ENTRIES: z.coerce.number().int().positive().default(1000), OLLAMA_BASE_URL: z.string().default('http://localhost:11434/v1'), EMBED_MODEL: z.string().default('nomic-embed-text'), + REQUEST_TIMEOUT_MS: z.coerce.number().int().nonnegative().default(30_000), + MAX_RETRIES: z.coerce.number().int().nonnegative().default(2), + DEFAULT_RPM: z.coerce.number().int().nonnegative().default(0), + THROTTLE_MAX_WAIT_MS: z.coerce.number().int().nonnegative().default(2_000), }); /** Reads and validates the process environment Sentinel needs to run. */ @@ -62,6 +70,10 @@ export function loadServerEnv(env: NodeJS.ProcessEnv): ServerEnv { cacheMaxEntries: parsed.data.CACHE_MAX_ENTRIES, ollamaBaseUrl: parsed.data.OLLAMA_BASE_URL, embedModel: parsed.data.EMBED_MODEL, + requestTimeoutMs: parsed.data.REQUEST_TIMEOUT_MS, + maxRetries: parsed.data.MAX_RETRIES, + defaultRpm: parsed.data.DEFAULT_RPM, + throttleMaxWaitMs: parsed.data.THROTTLE_MAX_WAIT_MS, }; } @@ -73,16 +85,23 @@ const providerConfigSchema = z baseUrl: z.string().url().optional(), baseUrlEnv: z.string().min(1).optional(), apiKeyEnv: z.string().min(1).optional(), + rpm: z.coerce.number().int().nonnegative().optional(), }) .refine((p) => (p.baseUrl === undefined) !== (p.baseUrlEnv === undefined), { message: 'set exactly one of "baseUrl" or "baseUrlEnv"', }); +const routingConfigSchema = z.object({ + tiers: z.array(z.string().min(1)).optional(), + fallback: z.array(z.string().min(1)).optional(), +}); + const sentinelConfigSchema = z .object({ providers: z.record(z.string(), providerConfigSchema), models: z.record(z.string(), z.string()), defaultProvider: z.string().optional(), + routing: routingConfigSchema.optional(), }) .superRefine((cfg, ctx) => { const names = new Set(Object.keys(cfg.providers)); @@ -106,6 +125,14 @@ export interface ResolvedProvider { name: string; baseUrl: string; apiKey: string | undefined; + /** Per-provider requests-per-minute limit for the throttle (omitted = unlimited). */ + rpm?: number; +} + +/** Routing rules (tier list + fallback chain) from the config file. */ +export interface ResolvedRouting { + tiers?: string[] | undefined; + fallback?: string[] | undefined; } export interface ResolvedConfig { @@ -113,6 +140,7 @@ export interface ResolvedConfig { /** model name → provider name */ models: Map; defaultProvider: string | undefined; + routing?: ResolvedRouting; } export interface LoadConfigOptions { @@ -149,13 +177,14 @@ export function loadConfig(options: LoadConfigOptions): ResolvedConfig { for (const [name, p] of Object.entries(parsed.data.providers)) { const baseUrl = p.baseUrl ?? readEnvOrThrow(options.env, p.baseUrlEnv, name); const apiKey = p.apiKeyEnv === undefined ? undefined : options.env[p.apiKeyEnv]; - providers.set(name, { name, baseUrl, apiKey }); + providers.set(name, { name, baseUrl, apiKey, ...(p.rpm !== undefined ? { rpm: p.rpm } : {}) }); } return { providers, models: new Map(Object.entries(parsed.data.models)), defaultProvider: parsed.data.defaultProvider, + ...(parsed.data.routing ? { routing: parsed.data.routing } : {}), }; } diff --git a/packages/gateway/src/index.ts b/packages/gateway/src/index.ts index 17d364d..f1b490c 100644 --- a/packages/gateway/src/index.ts +++ b/packages/gateway/src/index.ts @@ -3,11 +3,30 @@ export const SENTINEL_VERSION = '0.1.0'; export { buildServer } from './server.js'; -export type { ServerDeps } from './server.js'; +export type { ServerDeps, RoutingDeps } from './server.js'; export { loadConfig, loadServerEnv } from './config.js'; -export type { ResolvedConfig, ResolvedProvider, ServerEnv, LoadConfigOptions } from './config.js'; +export type { + ResolvedConfig, + ResolvedProvider, + ResolvedRouting, + ServerEnv, + LoadConfigOptions, +} from './config.js'; export { createRegistry } from './providers/registry.js'; export type { ProviderRegistry } from './providers/registry.js'; +export { createRouter } from './routing/router.js'; +export type { Router, Candidate, RoutingConfig } from './routing/router.js'; +export { runChat, openStream } from './routing/executor.js'; +export type { RouteOutcome, ExecutorOptions, StreamHandle } from './routing/executor.js'; +export { isRetryable } from './routing/retryable.js'; +export { classifyTier } from './routing/classifier.js'; +export { createTokenBucket, createBucketRegistry } from './throttle/token-bucket.js'; +export type { + TokenBucket, + TokenBucketOptions, + BucketRegistry, + BucketRegistryOptions, +} from './throttle/token-bucket.js'; export { createOpenAICompatibleProvider } from './providers/openai-compatible.js'; export type { Provider, FetchLike } from './providers/types.js'; export { diff --git a/packages/gateway/src/main.ts b/packages/gateway/src/main.ts index 4d16ae1..8581fda 100644 --- a/packages/gateway/src/main.ts +++ b/packages/gateway/src/main.ts @@ -8,6 +8,7 @@ import { initTelemetry } from './telemetry/otel.js'; import { createOllamaEmbedder } from './cache/embedder.js'; import { createSemanticCache } from './cache/cache.js'; import type { SemanticCache } from './cache/cache.js'; +import { createBucketRegistry } from './throttle/token-bucket.js'; async function main(): Promise { const env = loadServerEnv(process.env); @@ -29,12 +30,25 @@ async function main(): Promise { }); } + const rpmByProvider: Record = {}; + for (const provider of config.providers.values()) { + if (provider.rpm !== undefined) rpmByProvider[provider.name] = provider.rpm; + } + const throttle = createBucketRegistry({ rpmByProvider, defaultRpm: env.defaultRpm }); + const app = buildServer({ registry, apiKeys: env.apiKeys, traceStore: store, adminKey: env.adminKey, cache, + routing: { + config: config.routing, + maxRetries: env.maxRetries, + timeoutMs: env.requestTimeoutMs, + maxWaitMs: env.throttleMaxWaitMs, + throttle, + }, }); const shutdown = async (): Promise => { diff --git a/packages/gateway/src/routes.traces.ts b/packages/gateway/src/routes.traces.ts index 1e53740..5c72945 100644 --- a/packages/gateway/src/routes.traces.ts +++ b/packages/gateway/src/routes.traces.ts @@ -40,6 +40,8 @@ function parseTraceQuery(raw: unknown): TraceQuery { if (q.since !== undefined) query.since = Number(q.since); if (q.until !== undefined) query.until = Number(q.until); if (q.cacheHit !== undefined) query.cacheHit = q.cacheHit === 'true'; + if (q.routedProvider !== undefined) query.routedProvider = q.routedProvider; + if (q.fallbackUsed !== undefined) query.fallbackUsed = q.fallbackUsed === 'true'; if (q.limit !== undefined) query.limit = Math.min(Number(q.limit), 500); if (q.offset !== undefined) query.offset = Number(q.offset); return query; diff --git a/packages/gateway/src/routing/classifier.test.ts b/packages/gateway/src/routing/classifier.test.ts new file mode 100644 index 0000000..611e9a5 --- /dev/null +++ b/packages/gateway/src/routing/classifier.test.ts @@ -0,0 +1,37 @@ +import { describe, it, expect } from 'vitest'; +import { classifyTier } from './classifier.js'; +import type { ChatCompletionRequest } from '../schemas.js'; + +function reqOf(content: string, extra: Partial = {}): ChatCompletionRequest { + return { model: 'auto', messages: [{ role: 'user', content }], ...extra }; +} + +describe('classifyTier', () => { + it('returns 0 when there is a single tier', () => { + expect(classifyTier(reqOf('x'.repeat(50_000)), 1)).toBe(0); + }); + + it('routes a short prompt to the cheapest tier', () => { + expect(classifyTier(reqOf('hi'), 3)).toBe(0); + }); + + it('escalates a long prompt to a higher tier', () => { + expect(classifyTier(reqOf('x'.repeat(3000)), 3)).toBe(2); + }); + + it('clamps the index to the available tiers', () => { + expect(classifyTier(reqOf('x'.repeat(50_000)), 2)).toBe(1); + }); + + it('factors in max_tokens', () => { + expect(classifyTier(reqOf('short', { max_tokens: 2000 }), 4)).toBe(3); // 5 + 8000 + 100 + }); + + it('handles non-string (array) content', () => { + const request: ChatCompletionRequest = { + model: 'auto', + messages: [{ role: 'user', content: [{ type: 'text', text: 'hi' }] }], + }; + expect(classifyTier(request, 3)).toBe(0); + }); +}); diff --git a/packages/gateway/src/routing/classifier.ts b/packages/gateway/src/routing/classifier.ts new file mode 100644 index 0000000..99999c2 --- /dev/null +++ b/packages/gateway/src/routing/classifier.ts @@ -0,0 +1,32 @@ +import type { ChatCompletionRequest } from '../schemas.js'; + +/** Score boundaries between tiers; crossing one escalates to the next tier. */ +const THRESHOLDS = [400, 2000, 8000, 32_000]; + +/** A cheap, model-free complexity estimate from the request shape. */ +function complexityScore(request: ChatCompletionRequest): number { + let chars = 0; + for (const message of request.messages) { + const content = message.content; + if (typeof content === 'string') chars += content.length; + else if (content !== undefined && content !== null) chars += JSON.stringify(content).length; + } + const maxTokens = typeof request.max_tokens === 'number' ? request.max_tokens : 0; + return chars + maxTokens * 4 + request.messages.length * 100; +} + +/** + * Picks a tier index (0 = cheapest) for `model: "auto"`, given how many tiers are + * configured. Deterministic and rules-based: longer prompts, more messages, and a + * larger `max_tokens` budget escalate to higher (more capable) tiers. + */ +export function classifyTier(request: ChatCompletionRequest, tierCount: number): number { + if (tierCount <= 1) return 0; + const score = complexityScore(request); + let index = 0; + for (const threshold of THRESHOLDS) { + if (score >= threshold) index += 1; + else break; + } + return Math.min(index, tierCount - 1); +} diff --git a/packages/gateway/src/routing/executor.test.ts b/packages/gateway/src/routing/executor.test.ts new file mode 100644 index 0000000..af06acb --- /dev/null +++ b/packages/gateway/src/routing/executor.test.ts @@ -0,0 +1,163 @@ +import { describe, it, expect } from 'vitest'; +import { runChat, openStream } from './executor.js'; +import type { ExecutorOptions } from './executor.js'; +import type { Candidate } from './router.js'; +import type { Provider } from '../providers/types.js'; +import type { ChatCompletionRequest } from '../schemas.js'; +import { UpstreamError } from '../errors.js'; + +const baseOpts: ExecutorOptions = { + maxRetries: 0, + timeoutMs: 0, + baseBackoffMs: 1, + maxWaitMs: 0, + sleep: () => Promise.resolve(), +}; + +const req: ChatCompletionRequest = { model: 'm', messages: [{ role: 'user', content: 'hi' }] }; + +function provider(name: string, over: Partial = {}): Provider { + return { + name, + chat: () => Promise.resolve({ ok: name }), + chatStream: async function* () { + yield `{"p":"${name}"}`; + }, + ...over, + }; +} + +const cand = (p: Provider, model = 'm'): Candidate => ({ provider: p, model }); + +describe('runChat', () => { + it('returns the first candidate result on success', async () => { + const { result, outcome } = await runChat([cand(provider('a'))], req, baseOpts); + expect(result).toEqual({ ok: 'a' }); + expect(outcome).toMatchObject({ + routedProvider: 'a', + routedModel: 'm', + fallbackUsed: false, + retryCount: 0, + }); + }); + + it('rewrites the request model to the chosen candidate', async () => { + const seen: string[] = []; + const p = provider('a', { + chat: (request) => { + seen.push(request.model); + return Promise.resolve({ ok: 'a' }); + }, + }); + await runChat([cand(p, 'gpt-tier')], req, baseOpts); + expect(seen).toEqual(['gpt-tier']); + }); + + it('retries a retryable error then succeeds on the same candidate', async () => { + let calls = 0; + const p = provider('a', { + chat: () => { + calls += 1; + return calls === 1 + ? Promise.reject(new UpstreamError('a', 429, 'rl')) + : Promise.resolve({ ok: 'a', calls }); + }, + }); + const { result, outcome } = await runChat([cand(p)], req, { ...baseOpts, maxRetries: 1 }); + expect(result).toEqual({ ok: 'a', calls: 2 }); + expect(outcome.retryCount).toBe(1); + expect(outcome.fallbackUsed).toBe(false); + }); + + it('fails over to the next candidate after retryable exhaustion', async () => { + const p1 = provider('a', { chat: () => Promise.reject(new UpstreamError('a', 502, 'boom')) }); + const { result, outcome } = await runChat([cand(p1), cand(provider('b'))], req, { + ...baseOpts, + maxRetries: 1, + }); + expect(result).toEqual({ ok: 'b' }); + expect(outcome.routedProvider).toBe('b'); + expect(outcome.fallbackUsed).toBe(true); + expect(outcome.retryCount).toBe(1); + }); + + it('throws a terminal error immediately without falling back', async () => { + let p2calls = 0; + const p1 = provider('a', { chat: () => Promise.reject(new UpstreamError('a', 400, 'bad')) }); + const p2 = provider('b', { + chat: () => { + p2calls += 1; + return Promise.resolve({}); + }, + }); + await expect( + runChat([cand(p1), cand(p2)], req, { ...baseOpts, maxRetries: 2 }), + ).rejects.toBeInstanceOf(UpstreamError); + expect(p2calls).toBe(0); + }); + + it('throws the last error when every candidate fails', async () => { + const p1 = provider('a', { chat: () => Promise.reject(new UpstreamError('a', 502, 'a-down')) }); + const p2 = provider('b', { chat: () => Promise.reject(new UpstreamError('b', 429, 'b-rl')) }); + await expect(runChat([cand(p1), cand(p2)], req, baseOpts)).rejects.toMatchObject({ + status: 429, + }); + }); + + it('throws 503 when there are no candidates', async () => { + await expect(runChat([], req, baseOpts)).rejects.toMatchObject({ status: 503 }); + }); + + it('throws 503 when every candidate is throttled', async () => { + const throttle = { acquire: () => Promise.resolve(false) }; + await expect( + runChat([cand(provider('a')), cand(provider('b'))], req, { ...baseOpts, throttle }), + ).rejects.toMatchObject({ status: 503 }); + }); + + it('skips a throttled candidate and uses the next one', async () => { + const throttle = { acquire: (name: string) => Promise.resolve(name === 'b') }; + const { outcome } = await runChat([cand(provider('a')), cand(provider('b'))], req, { + ...baseOpts, + throttle, + }); + expect(outcome.routedProvider).toBe('b'); + expect(outcome.fallbackUsed).toBe(true); + }); + + it('times out a slow attempt and fails over', async () => { + const slow = provider('a', { + chat: (_request, signal) => + new Promise((_resolve, reject) => { + signal?.addEventListener('abort', () => + reject(Object.assign(new Error('aborted'), { name: 'AbortError' })), + ); + }), + }); + const { outcome } = await runChat([cand(slow), cand(provider('b'))], req, { + ...baseOpts, + timeoutMs: 5, + }); + expect(outcome.routedProvider).toBe('b'); + }); +}); + +describe('openStream', () => { + it('opens a stream from the first healthy candidate', async () => { + const { first, outcome } = await openStream([cand(provider('a'))], req, baseOpts); + expect(first.value).toBe('{"p":"a"}'); + expect(outcome.routedProvider).toBe('a'); + }); + + it('fails over when the first stream errors on its first chunk', async () => { + const p1 = provider('a', { + // eslint-disable-next-line require-yield + chatStream: async function* () { + throw new UpstreamError('a', 503, 'down'); + }, + }); + const { first, outcome } = await openStream([cand(p1), cand(provider('b'))], req, baseOpts); + expect(first.value).toBe('{"p":"b"}'); + expect(outcome.fallbackUsed).toBe(true); + }); +}); diff --git a/packages/gateway/src/routing/executor.ts b/packages/gateway/src/routing/executor.ts new file mode 100644 index 0000000..987adc1 --- /dev/null +++ b/packages/gateway/src/routing/executor.ts @@ -0,0 +1,151 @@ +import type { ChatCompletionRequest } from '../schemas.js'; +import { GatewayError } from '../errors.js'; +import { isRetryable } from './retryable.js'; +import type { Candidate } from './router.js'; +import type { BucketRegistry } from '../throttle/token-bucket.js'; + +/** What actually happened while running the candidate chain — recorded on the trace. */ +export interface RouteOutcome { + routedProvider: string; + routedModel: string; + /** A non-primary candidate served the request. */ + fallbackUsed: boolean; + /** Total retries spent across the chain before success. */ + retryCount: number; +} + +export interface ExecutorOptions { + /** Retries per candidate after the first attempt (`0` = single attempt). */ + maxRetries: number; + /** Per-attempt timeout in ms (`0` disables). */ + timeoutMs: number; + /** Base retry backoff in ms; doubles each retry. */ + baseBackoffMs: number; + /** Max time the throttle may pace one candidate before it is skipped. */ + maxWaitMs: number; + /** Per-provider rate-limit buckets (omit to disable throttling). */ + throttle?: BucketRegistry; + /** Injectable sleep (defaults to a `setTimeout` promise). */ + sleep?: (ms: number) => Promise; +} + +const defaultSleep = (ms: number): Promise => + new Promise((resolve) => setTimeout(resolve, ms)); + +type Attempt = (candidate: Candidate, signal: AbortSignal | undefined) => Promise; + +async function withTimeout( + candidate: Candidate, + attempt: Attempt, + timeoutMs: number, +): Promise { + if (timeoutMs <= 0) return attempt(candidate, undefined); + const controller = new AbortController(); + const timer = setTimeout(() => controller.abort(), timeoutMs); + try { + return await attempt(candidate, controller.signal); + } finally { + clearTimeout(timer); + } +} + +/** + * Runs an attempt across the candidate chain: throttle → retry-with-backoff per + * candidate on retryable errors → fail over to the next candidate. Terminal errors + * throw immediately. Throws a 503 if every candidate is throttled. + */ +async function execute( + candidates: Candidate[], + options: ExecutorOptions, + attempt: Attempt, +): Promise<{ value: T; outcome: RouteOutcome }> { + const sleep = options.sleep ?? defaultSleep; + let retryCount = 0; + let lastError: unknown; + let anyThrottled = false; + let index = 0; + + for (const candidate of candidates) { + const isFallback = index > 0; + index += 1; + + if (options.throttle !== undefined) { + const allowed = await options.throttle.acquire(candidate.provider.name, options.maxWaitMs); + if (!allowed) { + anyThrottled = true; + continue; + } + } + + for (let tryNum = 0; tryNum <= options.maxRetries; tryNum += 1) { + try { + const value = await withTimeout(candidate, attempt, options.timeoutMs); + return { + value, + outcome: { + routedProvider: candidate.provider.name, + routedModel: candidate.model, + fallbackUsed: isFallback, + retryCount, + }, + }; + } catch (error) { + lastError = error; + if (!isRetryable(error)) throw error; // terminal: no retry, no fallback + if (tryNum < options.maxRetries) { + retryCount += 1; + await sleep(options.baseBackoffMs * 2 ** tryNum); + } + } + } + } + + if (lastError !== undefined) throw lastError; + if (anyThrottled) { + throw new GatewayError( + 503, + 'All candidate providers are rate-limited; retry shortly.', + 'rate_limited', + 'rate_limited', + ); + } + throw new GatewayError(503, 'No provider candidates were available.', 'no_candidates', null); +} + +/** Runs a non-streaming completion across the candidate chain. */ +export async function runChat( + candidates: Candidate[], + request: ChatCompletionRequest, + options: ExecutorOptions, +): Promise<{ result: unknown; outcome: RouteOutcome }> { + const { value, outcome } = await execute(candidates, options, (candidate, signal) => + candidate.provider.chat({ ...request, model: candidate.model }, signal), + ); + return { result: value, outcome }; +} + +/** A live stream plus its already-pulled first chunk, after routing/fallback. */ +export interface StreamHandle { + iterator: AsyncIterator; + first: IteratorResult; + outcome: RouteOutcome; +} + +/** + * Opens a streaming completion across the candidate chain, pulling the first chunk + * so an immediate upstream failure can still fail over. Fallback is bounded to this + * point — once the first chunk is in hand, the caller owns the live stream. + */ +export async function openStream( + candidates: Candidate[], + request: ChatCompletionRequest, + options: ExecutorOptions, +): Promise { + const { value, outcome } = await execute(candidates, options, async (candidate, signal) => { + const stream = candidate.provider.chatStream({ ...request, model: candidate.model }, signal); + const iterator = stream[Symbol.asyncIterator](); + const first = await iterator.next(); + return { iterator, first }; + }); + return { iterator: value.iterator, first: value.first, outcome }; +} diff --git a/packages/gateway/src/routing/retryable.test.ts b/packages/gateway/src/routing/retryable.test.ts new file mode 100644 index 0000000..048951f --- /dev/null +++ b/packages/gateway/src/routing/retryable.test.ts @@ -0,0 +1,37 @@ +import { describe, it, expect } from 'vitest'; +import { isRetryable } from './retryable.js'; +import { UpstreamError, ValidationError, AuthError, ModelNotFoundError } from '../errors.js'; + +describe('isRetryable', () => { + it('treats rate limits and upstream 5xx as retryable', () => { + expect(isRetryable(new UpstreamError('p', 429, 'rate limited'))).toBe(true); + expect(isRetryable(new UpstreamError('p', 500, 'boom'))).toBe(true); // collapses to 502 + expect(isRetryable(new UpstreamError('p', 503, 'down'))).toBe(true); // collapses to 502 + }); + + it('treats upstream 4xx as terminal', () => { + expect(isRetryable(new UpstreamError('p', 400, 'bad'))).toBe(false); + expect(isRetryable(new UpstreamError('p', 401, 'no key'))).toBe(false); + expect(isRetryable(new UpstreamError('p', 404, 'gone'))).toBe(false); + }); + + it('treats other gateway errors as terminal', () => { + expect(isRetryable(new ValidationError('x'))).toBe(false); + expect(isRetryable(new AuthError())).toBe(false); + expect(isRetryable(new ModelNotFoundError('m'))).toBe(false); + }); + + it('treats aborts, timeouts, and network faults as retryable', () => { + expect(isRetryable(Object.assign(new Error('aborted'), { name: 'AbortError' }))).toBe(true); + expect(isRetryable(Object.assign(new Error('t'), { name: 'TimeoutError' }))).toBe(true); + expect(isRetryable(Object.assign(new Error('boom'), { name: 'FetchError' }))).toBe(true); + expect(isRetryable(new TypeError('fetch failed'))).toBe(true); + }); + + it('treats unrelated errors and non-errors as terminal', () => { + expect(isRetryable(new TypeError('bad argument'))).toBe(false); + expect(isRetryable(new Error('whatever'))).toBe(false); + expect(isRetryable('nope')).toBe(false); + expect(isRetryable(null)).toBe(false); + }); +}); diff --git a/packages/gateway/src/routing/retryable.ts b/packages/gateway/src/routing/retryable.ts new file mode 100644 index 0000000..29c7385 --- /dev/null +++ b/packages/gateway/src/routing/retryable.ts @@ -0,0 +1,22 @@ +import { GatewayError, UpstreamError } from '../errors.js'; + +/** + * Whether a failed provider attempt is worth retrying or failing over to another + * candidate. Transient conditions (rate limits, upstream 5xx, timeouts, network + * faults) are retryable; client-side errors (validation, auth, model-not-found, + * upstream 4xx) are terminal — they would fail on every candidate. + */ +export function isRetryable(error: unknown): boolean { + if (error instanceof UpstreamError) { + return ( + error.status === 429 || error.status === 502 || error.status === 503 || error.status === 504 + ); + } + if (error instanceof GatewayError) return false; + if (error instanceof Error) { + if (error.name === 'AbortError' || error.name === 'TimeoutError') return true; + if (error.name === 'FetchError') return true; + if (error.name === 'TypeError' && /fetch|network/i.test(error.message)) return true; + } + return false; +} diff --git a/packages/gateway/src/routing/router.test.ts b/packages/gateway/src/routing/router.test.ts new file mode 100644 index 0000000..1a80f0c --- /dev/null +++ b/packages/gateway/src/routing/router.test.ts @@ -0,0 +1,89 @@ +import { describe, it, expect } from 'vitest'; +import { createRouter } from './router.js'; +import type { ProviderRegistry } from '../providers/registry.js'; +import type { Provider } from '../providers/types.js'; +import { ModelNotFoundError } from '../errors.js'; +import type { ChatCompletionRequest } from '../schemas.js'; + +function provider(name: string): Provider { + return { + name, + chat: () => Promise.resolve({}), + chatStream: async function* () { + yield 'x'; + }, + }; +} + +function registryOf(map: Record, fallbackProvider?: Provider): ProviderRegistry { + return { + resolve(model) { + const p = map[model] ?? fallbackProvider; + if (p === undefined) throw new ModelNotFoundError(model); + return p; + }, + }; +} + +const reqOf = (model: string, content = 'hi'): ChatCompletionRequest => ({ + model, + messages: [{ role: 'user', content }], +}); + +describe('createRouter', () => { + const a = provider('a'); + const b = provider('b'); + const local = provider('local'); + + it('builds [model, ...fallback] for an explicit model', () => { + const router = createRouter({ registry: registryOf({ a, b }), routing: { fallback: ['b'] } }); + const candidates = router.resolveCandidates(reqOf('a')); + expect(candidates.map((c) => `${c.provider.name}:${c.model}`)).toEqual(['a:a', 'b:b']); + }); + + it('returns a single candidate with no fallback configured', () => { + const router = createRouter({ registry: registryOf({ a }) }); + expect(router.resolveCandidates(reqOf('a'))).toHaveLength(1); + }); + + it('classifies model:"auto" into the cheapest tier and escalates through the rest', () => { + const router = createRouter({ + registry: registryOf({ a, b, local }), + routing: { tiers: ['a', 'b'], fallback: ['local'] }, + }); + const candidates = router.resolveCandidates(reqOf('auto')); + expect(candidates.map((c) => c.model)).toEqual(['a', 'b', 'local']); + }); + + it('throws when model:"auto" has no tiers configured', () => { + const router = createRouter({ registry: registryOf({ a }) }); + expect(() => router.resolveCandidates(reqOf('auto'))).toThrow(ModelNotFoundError); + }); + + it('skips a fallback that does not resolve', () => { + const router = createRouter({ registry: registryOf({ a }), routing: { fallback: ['ghost'] } }); + const candidates = router.resolveCandidates(reqOf('a')); + expect(candidates).toHaveLength(1); + expect(candidates[0]?.model).toBe('a'); + }); + + it('dedupes candidates resolving to the same provider and model', () => { + const router = createRouter({ registry: registryOf({ a }), routing: { fallback: ['a'] } }); + expect(router.resolveCandidates(reqOf('a'))).toHaveLength(1); + }); + + it('throws the primary error when nothing resolves', () => { + const router = createRouter({ registry: registryOf({}), routing: { fallback: ['x'] } }); + expect(() => router.resolveCandidates(reqOf('missing'))).toThrow(ModelNotFoundError); + }); + + it('propagates a non-model-not-found registry error', () => { + const registry: ProviderRegistry = { + resolve() { + throw new Error('boom'); + }, + }; + const router = createRouter({ registry }); + expect(() => router.resolveCandidates(reqOf('a'))).toThrow('boom'); + }); +}); diff --git a/packages/gateway/src/routing/router.ts b/packages/gateway/src/routing/router.ts new file mode 100644 index 0000000..2c2fe6a --- /dev/null +++ b/packages/gateway/src/routing/router.ts @@ -0,0 +1,81 @@ +import type { ChatCompletionRequest } from '../schemas.js'; +import type { Provider } from '../providers/types.js'; +import type { ProviderRegistry } from '../providers/registry.js'; +import { ModelNotFoundError } from '../errors.js'; +import { classifyTier } from './classifier.js'; + +/** One (provider, model) pair the executor may try, in priority order. */ +export interface Candidate { + provider: Provider; + model: string; +} + +/** Routing rules from the config file. Both lists are optional. */ +export interface RoutingConfig { + /** Cheapest-first model tiers for `model: "auto"`. */ + tiers?: string[] | undefined; + /** Models appended as fallbacks to every request's candidate chain. */ + fallback?: string[] | undefined; +} + +export interface Router { + /** Builds the ordered candidate chain for a request. Throws if nothing resolves. */ + resolveCandidates(request: ChatCompletionRequest): Candidate[]; +} + +export interface CreateRouterOptions { + registry: ProviderRegistry; + routing?: RoutingConfig; +} + +/** + * Turns a request into an ordered candidate chain. Explicit models become + * `[model, ...fallback]`; `model: "auto"` is classified into the cheapest capable + * tier and escalates through the remaining tiers, then the fallback chain. + */ +export function createRouter(options: CreateRouterOptions): Router { + const registry = options.registry; + const tiers = options.routing?.tiers ?? []; + const fallback = options.routing?.fallback ?? []; + + return { + resolveCandidates(request): Candidate[] { + const requested = request.model; + let names: string[]; + if (requested === 'auto') { + if (tiers.length === 0) throw new ModelNotFoundError('auto'); + const index = classifyTier(request, tiers.length); + names = [...tiers.slice(index), ...fallback]; + } else { + names = [requested, ...fallback]; + } + + const candidates: Candidate[] = []; + const seen = new Set(); + let primaryError: unknown; + let isPrimary = true; + for (const model of names) { + let provider: Provider; + try { + provider = registry.resolve(model); + } catch (error) { + if (isPrimary) primaryError = error; + isPrimary = false; + if (error instanceof ModelNotFoundError) continue; // skip an unresolvable fallback + throw error; // anything else is a real failure, not a routing miss + } + isPrimary = false; + const key = `${provider.name}::${model}`; + if (!seen.has(key)) { + seen.add(key); + candidates.push({ provider, model }); + } + } + + if (candidates.length === 0) { + throw primaryError ?? new ModelNotFoundError(requested); + } + return candidates; + }, + }; +} diff --git a/packages/gateway/src/server.test.ts b/packages/gateway/src/server.test.ts index ead1dda..217ba09 100644 --- a/packages/gateway/src/server.test.ts +++ b/packages/gateway/src/server.test.ts @@ -34,6 +34,16 @@ function makeRegistry(provider: Provider, opts: { unknownModel?: boolean } = {}) }; } +function makeMultiRegistry(map: Record): ProviderRegistry { + return { + resolve(model) { + const provider = map[model]; + if (provider === undefined) throw new ModelNotFoundError(model); + return provider; + }, + }; +} + function buildTestServer( registry: ProviderRegistry, opts: { store?: TraceStore; adminKey?: string; cache?: SemanticCache } = {}, @@ -214,6 +224,147 @@ describe('POST /v1/chat/completions', () => { }); }); +describe('routing, fallback & throttle', () => { + it('falls back to a healthy provider on a retryable upstream error', async () => { + const primary = makeProvider({ + name: 'primary', + chat: () => Promise.reject(new UpstreamError('primary', 429, 'rate limited')), + }); + const fb = makeProvider({ name: 'fb', chat: (req) => Promise.resolve({ served: req.model }) }); + const app = buildServer({ + registry: makeMultiRegistry({ primary, fb }), + apiKeys: new Set(['good']), + logger: false, + traceStore: new InMemoryTraceStore(), + routing: { config: { fallback: ['fb'] }, maxRetries: 0, sleep: () => Promise.resolve() }, + }); + const res = await app.inject({ + method: 'POST', + url, + headers: auth, + payload: { ...body, model: 'primary' }, + }); + expect(res.statusCode).toBe(200); + expect(res.json()).toEqual({ served: 'fb' }); + await app.close(); + }); + + it('routes model:"auto" to the classified cheapest tier', async () => { + const cheap = makeProvider({ + name: 'cheap', + chat: (req) => Promise.resolve({ served: req.model }), + }); + const big = makeProvider({ + name: 'big', + chat: (req) => Promise.resolve({ served: req.model }), + }); + const app = buildServer({ + registry: makeMultiRegistry({ cheap, big }), + apiKeys: new Set(['good']), + logger: false, + traceStore: new InMemoryTraceStore(), + routing: { config: { tiers: ['cheap', 'big'] } }, + }); + const res = await app.inject({ + method: 'POST', + url, + headers: auth, + payload: { model: 'auto', messages: [{ role: 'user', content: 'hi' }] }, + }); + expect(res.statusCode).toBe(200); + expect(res.json()).toEqual({ served: 'cheap' }); + await app.close(); + }); + + it('skips a throttled provider and falls back', async () => { + const primary = makeProvider({ + name: 'primary', + chat: () => Promise.resolve({ served: 'primary' }), + }); + const fb = makeProvider({ name: 'fb', chat: () => Promise.resolve({ served: 'fb' }) }); + const throttle = { acquire: (provider: string) => Promise.resolve(provider !== 'primary') }; + const app = buildServer({ + registry: makeMultiRegistry({ primary, fb }), + apiKeys: new Set(['good']), + logger: false, + traceStore: new InMemoryTraceStore(), + routing: { config: { fallback: ['fb'] }, throttle }, + }); + const res = await app.inject({ + method: 'POST', + url, + headers: auth, + payload: { ...body, model: 'primary' }, + }); + expect(res.statusCode).toBe(200); + expect(res.json()).toEqual({ served: 'fb' }); + await app.close(); + }); + + it('does not fall back on a terminal (4xx) upstream error', async () => { + let fbCalls = 0; + const primary = makeProvider({ + name: 'primary', + chat: () => Promise.reject(new UpstreamError('primary', 400, 'bad request')), + }); + const fb = makeProvider({ + name: 'fb', + chat: () => { + fbCalls += 1; + return Promise.resolve({ served: 'fb' }); + }, + }); + const app = buildServer({ + registry: makeMultiRegistry({ primary, fb }), + apiKeys: new Set(['good']), + logger: false, + traceStore: new InMemoryTraceStore(), + routing: { config: { fallback: ['fb'] }, maxRetries: 2, sleep: () => Promise.resolve() }, + }); + const res = await app.inject({ + method: 'POST', + url, + headers: auth, + payload: { ...body, model: 'primary' }, + }); + expect(res.statusCode).toBe(400); + expect(fbCalls).toBe(0); + await app.close(); + }); + + it('falls back when the primary stream fails on its first chunk', async () => { + const primary = makeProvider({ + name: 'primary', + // eslint-disable-next-line require-yield + chatStream: async function* () { + throw new UpstreamError('primary', 503, 'down'); + }, + }); + const fb = makeProvider({ + name: 'fb', + chatStream: async function* () { + yield '{"d":"fb"}'; + }, + }); + const app = buildServer({ + registry: makeMultiRegistry({ primary, fb }), + apiKeys: new Set(['good']), + logger: false, + traceStore: new InMemoryTraceStore(), + routing: { config: { fallback: ['fb'] }, maxRetries: 0, sleep: () => Promise.resolve() }, + }); + const res = await app.inject({ + method: 'POST', + url, + headers: auth, + payload: { ...body, model: 'primary', stream: true }, + }); + expect(res.statusCode).toBe(200); + expect(res.body).toContain('{"d":"fb"}'); + await app.close(); + }); +}); + describe('tracing & /traces', () => { const sink = new InMemoryTraceStore(); let provider: NodeTracerProvider | undefined; @@ -294,6 +445,10 @@ describe('tracing & /traces', () => { errorMessage: null, apiKeyHash: null, cacheHit: false, + routedProvider: 'p', + routedModel: 'm', + fallbackUsed: false, + retryCount: 0, }); const app = buildTestServer(makeRegistry(makeProvider()), { store: seeded, @@ -328,6 +483,10 @@ describe('tracing & /traces', () => { errorMessage: null, apiKeyHash: null, cacheHit: false, + routedProvider: 'p', + routedModel: 'm', + fallbackUsed: false, + retryCount: 0, }; store.record({ ...base, id: 's1', timestamp: 100, model: 'm1', stream: false, status: 200 }); store.record({ ...base, id: 's2', timestamp: 200, model: 'm2', stream: true, status: 500 }); @@ -360,6 +519,43 @@ describe('tracing & /traces', () => { await app.close(); expect(sink.query({ cacheHit: true, model: 'cachetest' }).length).toBeGreaterThanOrEqual(1); }); + + it('records routing metadata and supports the fallbackUsed filter', async () => { + const primary = makeProvider({ + name: 'primary', + chat: () => Promise.reject(new UpstreamError('primary', 429, 'rate limited')), + }); + const fb = makeProvider({ name: 'fb', chat: (req) => Promise.resolve({ served: req.model }) }); + const app = buildServer({ + registry: makeMultiRegistry({ primary, fb }), + apiKeys: new Set(['good']), + logger: false, + traceStore: sink, + adminKey: 'admin-secret', + routing: { config: { fallback: ['fb'] }, maxRetries: 0, sleep: () => Promise.resolve() }, + }); + await app.inject({ + method: 'POST', + url, + headers: auth, + payload: { ...body, model: 'primary' }, + }); + const res = await app.inject({ + method: 'GET', + url: '/traces?fallbackUsed=true&routedProvider=fb', + headers: { authorization: 'Bearer admin-secret' }, + }); + await app.close(); + const rows = res.json() as { + fallbackUsed: boolean; + routedProvider: string; + routedModel: string; + }[]; + expect(rows.length).toBeGreaterThanOrEqual(1); + expect(rows[0]?.fallbackUsed).toBe(true); + expect(rows[0]?.routedProvider).toBe('fb'); + expect(rows[0]?.routedModel).toBe('fb'); + }); }); describe('semantic cache', () => { diff --git a/packages/gateway/src/server.ts b/packages/gateway/src/server.ts index 21fee21..6e260ba 100644 --- a/packages/gateway/src/server.ts +++ b/packages/gateway/src/server.ts @@ -9,6 +9,29 @@ import type { ProviderRegistry } from './providers/registry.js'; import type { TraceStore } from './telemetry/trace.js'; import type { SemanticCache } from './cache/cache.js'; import { traceRoutes } from './routes.traces.js'; +import { createRouter } from './routing/router.js'; +import type { RoutingConfig } from './routing/router.js'; +import { runChat, openStream } from './routing/executor.js'; +import type { ExecutorOptions, RouteOutcome } from './routing/executor.js'; +import type { BucketRegistry } from './throttle/token-bucket.js'; + +/** Routing, retry, fallback, and throttle settings. Omit for a single-provider pass-through. */ +export interface RoutingDeps { + /** Tier list + fallback chain (from `sentinel.config.json`). */ + config?: RoutingConfig | undefined; + /** Retries per candidate after the first attempt. */ + maxRetries?: number | undefined; + /** Per-attempt timeout in ms (`0` disables). */ + timeoutMs?: number | undefined; + /** Base retry backoff in ms. */ + baseBackoffMs?: number | undefined; + /** Max throttle pacing wait per candidate before it is skipped. */ + maxWaitMs?: number | undefined; + /** Per-provider rate-limit buckets. */ + throttle?: BucketRegistry | undefined; + /** Injectable sleep (handy for tests). */ + sleep?: ((ms: number) => Promise) | undefined; +} export interface ServerDeps { registry: ProviderRegistry; @@ -16,6 +39,7 @@ export interface ServerDeps { traceStore: TraceStore; adminKey?: string | undefined; cache?: SemanticCache | undefined; + routing?: RoutingDeps | undefined; logger?: FastifyServerOptions['logger']; } @@ -63,6 +87,15 @@ function setUsageAttributes(span: Span | undefined, result: unknown): void { span.setAttribute('gen_ai.usage.total_tokens', u.total_tokens); } +/** Records which provider/model actually served the request after routing/fallback. */ +function applyRouteAttributes(span: Span | undefined, outcome: RouteOutcome): void { + span?.setAttribute('sentinel.provider', outcome.routedProvider); + span?.setAttribute('sentinel.routed_provider', outcome.routedProvider); + span?.setAttribute('sentinel.routed_model', outcome.routedModel); + span?.setAttribute('sentinel.fallback_used', outcome.fallbackUsed); + span?.setAttribute('sentinel.retry_count', outcome.retryCount); +} + /** Builds the Sentinel gateway Fastify app. Dependencies are injected for testability. */ export function buildServer(deps: ServerDeps) { const app = Fastify({ @@ -86,6 +119,20 @@ export function buildServer(deps: ServerDeps) { const authHook = createAuthHook(deps.apiKeys); + const routerConfig = deps.routing?.config; + const router = createRouter({ + registry: deps.registry, + ...(routerConfig ? { routing: routerConfig } : {}), + }); + const execOptions: ExecutorOptions = { + maxRetries: deps.routing?.maxRetries ?? 0, + timeoutMs: deps.routing?.timeoutMs ?? 0, + baseBackoffMs: deps.routing?.baseBackoffMs ?? 200, + maxWaitMs: deps.routing?.maxWaitMs ?? 0, + ...(deps.routing?.throttle ? { throttle: deps.routing.throttle } : {}), + ...(deps.routing?.sleep ? { sleep: deps.routing.sleep } : {}), + }; + app.post( '/v1/chat/completions', { @@ -119,8 +166,7 @@ export function buildServer(deps: ServerDeps) { span?.setAttribute('gen_ai.request.model', chatRequest.model); span?.setAttribute('sentinel.stream', chatRequest.stream === true); - const provider = deps.registry.resolve(chatRequest.model); - span?.setAttribute('sentinel.provider', provider.name); + const candidates = router.resolveCandidates(chatRequest); // ── Non-streaming ────────────────────────────────────────────────────── if (chatRequest.stream !== true) { @@ -133,8 +179,9 @@ export function buildServer(deps: ServerDeps) { return reply.status(200).send(cached.body); } } - const result = await provider.chat(chatRequest); + const { result, outcome } = await runChat(candidates, chatRequest, execOptions); setUsageAttributes(span, result); + applyRouteAttributes(span, outcome); if (deps.cache !== undefined && apiKeyHash !== null) { await deps.cache.set(chatRequest, apiKeyHash, { kind: 'json', body: result }); } @@ -162,9 +209,10 @@ export function buildServer(deps: ServerDeps) { } // Miss: pull the first chunk *before* committing to a 200 SSE response, so an - // immediate upstream failure still maps to a proper error status; buffer for caching. - const iterator = provider.chatStream(chatRequest)[Symbol.asyncIterator](); - const first = await iterator.next(); + // immediate upstream failure still routes/falls back to a proper error status; + // buffer for caching. Fallback is bounded to this first chunk (pre-hijack). + const { iterator, first, outcome } = await openStream(candidates, chatRequest, execOptions); + applyRouteAttributes(span, outcome); reply.hijack(); reply.raw.writeHead(200, SSE_HEADERS); diff --git a/packages/gateway/src/telemetry/store.memory.ts b/packages/gateway/src/telemetry/store.memory.ts index 85e3808..24f8eaf 100644 --- a/packages/gateway/src/telemetry/store.memory.ts +++ b/packages/gateway/src/telemetry/store.memory.ts @@ -34,5 +34,8 @@ function matches(trace: TraceRecord, filter: TraceQuery): boolean { if (filter.since !== undefined && trace.timestamp < filter.since) return false; if (filter.until !== undefined && trace.timestamp > filter.until) return false; if (filter.cacheHit !== undefined && trace.cacheHit !== filter.cacheHit) return false; + if (filter.routedProvider !== undefined && trace.routedProvider !== filter.routedProvider) + return false; + if (filter.fallbackUsed !== undefined && trace.fallbackUsed !== filter.fallbackUsed) return false; return true; } diff --git a/packages/gateway/src/telemetry/store.sqlite.ts b/packages/gateway/src/telemetry/store.sqlite.ts index c209d2f..9b6f006 100644 --- a/packages/gateway/src/telemetry/store.sqlite.ts +++ b/packages/gateway/src/telemetry/store.sqlite.ts @@ -17,12 +17,17 @@ const SCHEMA = ` error_type TEXT, error_message TEXT, api_key_hash TEXT, - cache_hit INTEGER NOT NULL DEFAULT 0 + cache_hit INTEGER NOT NULL DEFAULT 0, + routed_provider TEXT, + routed_model TEXT, + fallback_used INTEGER NOT NULL DEFAULT 0, + retry_count INTEGER NOT NULL DEFAULT 0 ); CREATE INDEX IF NOT EXISTS idx_traces_timestamp ON traces (timestamp); CREATE INDEX IF NOT EXISTS idx_traces_model ON traces (model); CREATE INDEX IF NOT EXISTS idx_traces_status ON traces (status); CREATE INDEX IF NOT EXISTS idx_traces_cache_hit ON traces (cache_hit); + CREATE INDEX IF NOT EXISTS idx_traces_fallback_used ON traces (fallback_used); `; interface TraceRow { @@ -41,6 +46,10 @@ interface TraceRow { error_message: string | null; api_key_hash: string | null; cache_hit: number; + routed_provider: string | null; + routed_model: string | null; + fallback_used: number; + retry_count: number; } /** SQLite-backed trace store (better-sqlite3, synchronous). Pass ':memory:' for tests. */ @@ -59,11 +68,11 @@ export class SqliteTraceStore implements TraceStore { `INSERT OR REPLACE INTO traces (id, trace_id, timestamp, duration_ms, model, provider, stream, status, prompt_tokens, completion_tokens, total_tokens, error_type, error_message, api_key_hash, - cache_hit) + cache_hit, routed_provider, routed_model, fallback_used, retry_count) VALUES (@id, @traceId, @timestamp, @durationMs, @model, @provider, @stream, @status, @promptTokens, @completionTokens, @totalTokens, @errorType, @errorMessage, @apiKeyHash, - @cacheHit)`, + @cacheHit, @routedProvider, @routedModel, @fallbackUsed, @retryCount)`, ) .run({ id: trace.id, @@ -81,6 +90,10 @@ export class SqliteTraceStore implements TraceStore { errorMessage: trace.errorMessage, apiKeyHash: trace.apiKeyHash, cacheHit: trace.cacheHit ? 1 : 0, + routedProvider: trace.routedProvider, + routedModel: trace.routedModel, + fallbackUsed: trace.fallbackUsed ? 1 : 0, + retryCount: trace.retryCount, }); } @@ -115,6 +128,14 @@ export class SqliteTraceStore implements TraceStore { where.push('cache_hit = @cacheHit'); params.cacheHit = filter.cacheHit ? 1 : 0; } + if (filter.routedProvider !== undefined) { + where.push('routed_provider = @routedProvider'); + params.routedProvider = filter.routedProvider; + } + if (filter.fallbackUsed !== undefined) { + where.push('fallback_used = @fallbackUsed'); + params.fallbackUsed = filter.fallbackUsed ? 1 : 0; + } params.limit = filter.limit ?? 50; params.offset = filter.offset ?? 0; const clause = where.length > 0 ? `WHERE ${where.join(' AND ')}` : ''; @@ -153,5 +174,9 @@ function rowToRecord(row: TraceRow): TraceRecord { errorMessage: row.error_message, apiKeyHash: row.api_key_hash, cacheHit: row.cache_hit === 1, + routedProvider: row.routed_provider, + routedModel: row.routed_model, + fallbackUsed: row.fallback_used === 1, + retryCount: row.retry_count, }; } diff --git a/packages/gateway/src/telemetry/store.test.ts b/packages/gateway/src/telemetry/store.test.ts index 4678309..07c66de 100644 --- a/packages/gateway/src/telemetry/store.test.ts +++ b/packages/gateway/src/telemetry/store.test.ts @@ -21,6 +21,10 @@ function sample(over: Partial = {}): TraceRecord { errorMessage: null, apiKeyHash: 'abc', cacheHit: false, + routedProvider: 'openai', + routedModel: 'gpt-4o-mini', + fallbackUsed: false, + retryCount: 0, ...over, }; } @@ -62,6 +66,16 @@ for (const [name, makeStore] of backends) { store.close(); }); + it('filters by routedProvider and fallbackUsed', () => { + const store = makeStore(); + store.record(sample({ id: 'a', routedProvider: 'p1', fallbackUsed: false })); + store.record(sample({ id: 'b', routedProvider: 'p2', fallbackUsed: true })); + expect(store.query({ routedProvider: 'p2' }).map((t) => t.id)).toEqual(['b']); + expect(store.query({ fallbackUsed: true }).map((t) => t.id)).toEqual(['b']); + expect(store.query({ fallbackUsed: false }).map((t) => t.id)).toEqual(['a']); + store.close(); + }); + it('respects limit and offset', () => { const store = makeStore(); for (let i = 0; i < 5; i++) store.record(sample({ id: `s${i}`, timestamp: i })); @@ -85,6 +99,10 @@ for (const [name, makeStore] of backends) { errorType: 'UpstreamError', errorMessage: 'boom', apiKeyHash: null, + routedProvider: null, + routedModel: null, + fallbackUsed: true, + retryCount: 3, }), ); const got = store.get('e'); @@ -92,6 +110,9 @@ for (const [name, makeStore] of backends) { expect(got?.stream).toBe(true); expect(got?.errorType).toBe('UpstreamError'); expect(got?.promptTokens).toBeNull(); + expect(got?.routedProvider).toBeNull(); + expect(got?.fallbackUsed).toBe(true); + expect(got?.retryCount).toBe(3); store.close(); }); }); diff --git a/packages/gateway/src/telemetry/trace.ts b/packages/gateway/src/telemetry/trace.ts index cd910f5..7719711 100644 --- a/packages/gateway/src/telemetry/trace.ts +++ b/packages/gateway/src/telemetry/trace.ts @@ -18,6 +18,14 @@ export interface TraceRecord { errorMessage: string | null; apiKeyHash: string | null; cacheHit: boolean; + /** Provider that actually served the response (after routing/fallback). */ + routedProvider: string | null; + /** Model that actually served the response (e.g. the tier chosen for `auto`). */ + routedModel: string | null; + /** A non-primary candidate served the request. */ + fallbackUsed: boolean; + /** Retries spent before the request succeeded. */ + retryCount: number; } /** Filters for querying traces (all optional). */ @@ -29,6 +37,8 @@ export interface TraceQuery { since?: number; until?: number; cacheHit?: boolean; + routedProvider?: string; + fallbackUsed?: boolean; limit?: number; offset?: number; } @@ -74,5 +84,9 @@ export function spanToTraceRecord(span: ReadableSpan): TraceRecord { errorMessage: isError ? (span.status.message ?? null) : null, apiKeyHash: asString(attrs['sentinel.api_key_hash']), cacheHit: attrs['sentinel.cache_hit'] === true, + routedProvider: asString(attrs['sentinel.routed_provider']), + routedModel: asString(attrs['sentinel.routed_model']), + fallbackUsed: attrs['sentinel.fallback_used'] === true, + retryCount: asNumber(attrs['sentinel.retry_count']) ?? 0, }; } diff --git a/packages/gateway/src/throttle/token-bucket.test.ts b/packages/gateway/src/throttle/token-bucket.test.ts new file mode 100644 index 0000000..efa10e4 --- /dev/null +++ b/packages/gateway/src/throttle/token-bucket.test.ts @@ -0,0 +1,70 @@ +import { describe, it, expect } from 'vitest'; +import { createTokenBucket, createBucketRegistry } from './token-bucket.js'; + +describe('createTokenBucket', () => { + it('allows everything when rpm is 0 (limiting disabled)', async () => { + const bucket = createTokenBucket({ rpm: 0 }); + expect(await bucket.acquire(0)).toBe(true); + expect(await bucket.acquire(0)).toBe(true); + }); + + it('grants up to the burst capacity then refuses with no wait budget', async () => { + const t = 0; + const bucket = createTokenBucket({ rpm: 60, now: () => t, sleep: () => Promise.resolve() }); + for (let i = 0; i < 60; i += 1) expect(await bucket.acquire(0)).toBe(true); + expect(await bucket.acquire(0)).toBe(false); + }); + + it('refills over time', async () => { + let t = 0; + const bucket = createTokenBucket({ rpm: 60, now: () => t, sleep: () => Promise.resolve() }); + for (let i = 0; i < 60; i += 1) await bucket.acquire(0); + expect(await bucket.acquire(0)).toBe(false); + t += 1000; // one second → one token at 60 rpm + expect(await bucket.acquire(0)).toBe(true); + }); + + it('waits for a refill when it fits the wait budget', async () => { + let t = 0; + const slept: number[] = []; + const bucket = createTokenBucket({ + rpm: 60, + now: () => t, + sleep: (ms) => { + slept.push(ms); + t += ms; + return Promise.resolve(); + }, + }); + for (let i = 0; i < 60; i += 1) await bucket.acquire(0); + expect(await bucket.acquire(2000)).toBe(true); + expect(slept[0]).toBeGreaterThan(0); + }); +}); + +describe('createBucketRegistry', () => { + it('applies per-provider rpm and a default', async () => { + const t = 0; + const reg = createBucketRegistry({ + rpmByProvider: { tight: 60 }, + defaultRpm: 0, + now: () => t, + sleep: () => Promise.resolve(), + }); + for (let i = 0; i < 60; i += 1) expect(await reg.acquire('tight', 0)).toBe(true); + expect(await reg.acquire('tight', 0)).toBe(false); // capped + expect(await reg.acquire('loose', 0)).toBe(true); // default 0 = unlimited + expect(await reg.acquire('loose', 0)).toBe(true); + }); + + it('reuses one bucket per provider', async () => { + const t = 0; + const reg = createBucketRegistry({ + defaultRpm: 60, + now: () => t, + sleep: () => Promise.resolve(), + }); + for (let i = 0; i < 60; i += 1) await reg.acquire('p', 0); + expect(await reg.acquire('p', 0)).toBe(false); + }); +}); diff --git a/packages/gateway/src/throttle/token-bucket.ts b/packages/gateway/src/throttle/token-bucket.ts new file mode 100644 index 0000000..a1ad25a --- /dev/null +++ b/packages/gateway/src/throttle/token-bucket.ts @@ -0,0 +1,95 @@ +/** + * In-memory token-bucket rate limiter (requests per minute) with continuous refill. + * Clock and sleep are injectable so it can be unit-tested without real time, and so + * it can be swapped for a Redis-backed implementation later. + */ +export interface TokenBucket { + /** + * Consumes one token. If none is available, waits up to `maxWaitMs` for a refill; + * returns `false` (without waiting) when the required wait would exceed `maxWaitMs`. + */ + acquire(maxWaitMs: number): Promise; +} + +export interface TokenBucketOptions { + /** Sustained requests per minute (also the burst capacity). `0` or less disables limiting. */ + rpm: number; + /** Injectable clock (defaults to `Date.now`). */ + now?: () => number; + /** Injectable sleep (defaults to a `setTimeout` promise). */ + sleep?: (ms: number) => Promise; +} + +const defaultSleep = (ms: number): Promise => + new Promise((resolve) => setTimeout(resolve, ms)); + +export function createTokenBucket(options: TokenBucketOptions): TokenBucket { + const now = options.now ?? ((): number => Date.now()); + const sleep = options.sleep ?? defaultSleep; + const capacity = options.rpm; + const ratePerMs = options.rpm / 60_000; + + let tokens = capacity; + let last = now(); + + function refill(): void { + const current = now(); + tokens = Math.min(capacity, tokens + (current - last) * ratePerMs); + last = current; + } + + return { + async acquire(maxWaitMs: number): Promise { + if (capacity <= 0) return true; // limiting disabled + refill(); + if (tokens >= 1) { + tokens -= 1; + return true; + } + const waitMs = Math.ceil((1 - tokens) / ratePerMs); + if (waitMs > maxWaitMs) return false; + await sleep(waitMs); + refill(); + tokens = Math.max(0, tokens - 1); // waiting earns the grant even if the clock lags + return true; + }, + }; +} + +/** Resolves a per-provider token bucket, created lazily from configured RPM. */ +export interface BucketRegistry { + acquire(provider: string, maxWaitMs: number): Promise; +} + +export interface BucketRegistryOptions { + /** Per-provider requests-per-minute overrides. */ + rpmByProvider?: Record; + /** Fallback RPM for providers without an override (`0` = unlimited). */ + defaultRpm?: number; + now?: () => number; + sleep?: (ms: number) => Promise; +} + +export function createBucketRegistry(options: BucketRegistryOptions = {}): BucketRegistry { + const buckets = new Map(); + + function bucketFor(provider: string): TokenBucket { + let bucket = buckets.get(provider); + if (bucket === undefined) { + const rpm = options.rpmByProvider?.[provider] ?? options.defaultRpm ?? 0; + bucket = createTokenBucket({ + rpm, + ...(options.now ? { now: options.now } : {}), + ...(options.sleep ? { sleep: options.sleep } : {}), + }); + buckets.set(provider, bucket); + } + return bucket; + } + + return { + acquire(provider, maxWaitMs) { + return bucketFor(provider).acquire(maxWaitMs); + }, + }; +} diff --git a/sentinel.config.example.json b/sentinel.config.example.json index 1ebabaf..d61b793 100644 --- a/sentinel.config.example.json +++ b/sentinel.config.example.json @@ -7,7 +7,8 @@ "openai": { "type": "openai-compatible", "baseUrl": "https://api.openai.com/v1", - "apiKeyEnv": "OPENAI_API_KEY" + "apiKeyEnv": "OPENAI_API_KEY", + "rpm": 60 }, "gemini": { "type": "openai-compatible", @@ -27,5 +28,9 @@ "gemini-2.0-flash": "gemini", "llama-3.3-70b-versatile": "groq" }, - "defaultProvider": "ollama" + "defaultProvider": "ollama", + "routing": { + "tiers": ["llama3.2", "gpt-4o-mini", "llama-3.3-70b-versatile"], + "fallback": ["llama3.2"] + } }