diff --git a/README.md b/README.md index 0da0587..d2ce0d9 100644 --- a/README.md +++ b/README.md @@ -13,7 +13,7 @@ key-based auth, quotas, rate limits, failover, and a billing ledger. - **MCP gateway** — `/mcp/{server}` proxies Model Context Protocol servers (Streamable HTTP) behind the same access keys: per-key server entitlement and tool allowlists, `tools/list` filtered to the allowlist, tool results reviewed by the tenant's moderator (mask or block), every tool call, denial and intervention audited, server credentials kept in the gateway's environment as static bearers or OAuth 2.0 client credentials the gateway refreshes itself - **Coding agents drop in** — Claude Code, Codex CLI, VS Code chat (Copilot), Cursor and opencode work with a base URL and an access key; the `anthropic-beta` header rides through to Anthropic-wire upstreams and each client's captured wire shape replays in the live matrix ([Examples](docs/examples.md#coding-agents)) - **Reasoning on every surface** — `reasoning_effort` / `reasoning{}` on `/v1/chat/completions` maps to each family's thinking dialect (Anthropic budget or adaptive by model generation, OpenAI effort, compatible vendors verbatim); reasoning comes back as `reasoning_content` + signed `reasoning_details` and replays into tool loops; `/v1/responses` forwards reasoning items and its native event stream verbatim; `/v1/messages` preserves signed `thinking`/`redacted_thinking` blocks end-to-end, pins reasoning traffic to its requested model, and audits tool-loop continuations against what was actually served (fail-open; tampering is a local 400) -- **Staged request pipeline** — a 4-layer DAG per request: model resolve / quota / cache lookup → account selection (priority, PTU-first, failover) → rate limits + engine call (retry on upstream 5xx, then the model's `fallback_models` chain before any byte is sent) → usage extraction, billing, cache store +- **Staged request pipeline** — a 4-layer DAG per request: model resolve / quota / cache lookup → account selection (priority, PTU-first, round-robin or latency-ranked within a tier, failover) → rate limits + engine call (retry on upstream 5xx, then the model's `fallback_models` chain before any byte is sent) → usage extraction, billing, cache store - **Governance built in** — access-key auth, daily token quotas, QPS / QPM / TPM limits at key, product, and model level, request-level TTL cache, account cooldown and recovery, DLP redaction and blocklist plugins. Admission reserves then settles, so concurrent requests can't overshoot a quota - **Multi-tenant** — keys carry a tenant; tenants get a pooled QPS bucket, a model entitlement allowlist, per-(key, model) quota defaults with an optional fallback-model degrade, key lifecycle (expiry/ban), and tenant-scoped admin tokens. Billing records charged cost and (optionally) vendor cost per row, so margin is queryable per tenant × model - **Per-user billing & enterprise audit** — every ledger row attributes to an effective end user (the key's `owner`, else the request's `x-gw-user` / `user` hint) with a `request_id`, so cost rolls up per user (`/admin/usage/users`) and soft per-user daily budgets apply on every surface; daily and calendar-month cost budgets (charged micro-dollars) cap a tenant pool, each key and each end user, with optional month-to-month rollover of the unspent remainder, raising a webhook alert when reached. Per-tenant content policy adds blocklist action tiers (block / flag / shadow), regex recognizers, secret masking, and an external-moderation seam with an AWS Bedrock Guardrails backend (deny on blocked policies, mask anonymized PII); every hit is recorded without prompt text. An admin-operation trail (key CRUD / config / reload, with source IP) and optional at-rest content retention complete the audit surfaces (`/admin/audit/*`) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 728896d..37ba65e 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -445,6 +445,9 @@ pub struct StabilityConf { /// Window error rate at or above which a model reports `unavailable`. #[serde(default = "default_unavailable_error_rate")] pub unavailable_error_rate: f64, + /// Rank same-priority accounts by their observed call latency (per instance) instead of round-robin. + #[serde(default)] + pub latency_routing: bool, /// Below this many window samples the verdict is `no_data`. #[serde(default = "default_availability_min_samples")] pub availability_min_samples: u64, @@ -458,6 +461,7 @@ impl Default for StabilityConf { availability_window_minutes: default_availability_window_minutes(), unstable_error_rate: default_unstable_error_rate(), unavailable_error_rate: default_unavailable_error_rate(), + latency_routing: false, availability_min_samples: default_availability_min_samples(), } } diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index cc9019a..8aec472 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -306,10 +306,15 @@ impl DagNode for SelectAccount { .find_model(served_model(ctx.request.model_param_v2.as_ref())); ctx.request.prompt_cache = conf.is_some_and(|m| m.prompt_cache); let provider = conf.and_then(|m| m.provider.as_deref()); + let latency = ctx + .cfg + .stability + .latency_routing + .then_some(&ctx.state.latency); let account = ctx .state .pool - .select_healthy(mt, provider, &[], ctx.state.health.as_ref()) + .select_healthy(mt, provider, &[], ctx.state.health.as_ref(), latency) .await; let Some(account) = account else { // unsampled, an exhausted pool would read no_data forever @@ -447,9 +452,10 @@ impl DagNode for CallEngine { } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { let mut engine = gw_engines::get_engine(ctx.request.clone(), ctx.transport.clone())?; + let started = latency_clock(ctx); match engine.run().await { Ok(outcome) => { - note_engine_outcome(ctx, &outcome).await; + note_engine_outcome(ctx, &outcome, started).await; ctx.decide( "call_engine", format!( @@ -473,6 +479,11 @@ impl DagNode for CallEngine { .ok_or_else(|| GatewayError::internal("call_engine without an account"))?; note_failure(ctx, &failed.name).await; let provider = model_provider(ctx); + let latency = ctx + .cfg + .stability + .latency_routing + .then_some(&ctx.state.latency); let next = ctx .state .pool @@ -481,6 +492,7 @@ impl DagNode for CallEngine { provider, std::slice::from_ref(&failed.name), ctx.state.health.as_ref(), + latency, ) .await; let Some(next) = next else { @@ -499,9 +511,10 @@ impl DagNode for CallEngine { ); ctx.request.account = Some(next.clone()); let mut retry = gw_engines::get_engine(ctx.request.clone(), ctx.transport.clone())?; + let started = latency_clock(ctx); match retry.run().await { Ok(mut outcome) => { - note_engine_outcome(ctx, &outcome).await; + note_engine_outcome(ctx, &outcome, started).await; outcome.response.ptu_spillover = spillover; ctx.outcome = Some(outcome); Ok(()) @@ -520,7 +533,19 @@ impl DagNode for CallEngine { } } -async fn note_engine_outcome(ctx: &mut DagContext, outcome: &gw_engines::EngineOutcome) { +/// The call timer for the latency ranker; `None` when latency routing is off, so the default pays nothing. +fn latency_clock(ctx: &DagContext) -> Option { + ctx.cfg + .stability + .latency_routing + .then(std::time::Instant::now) +} + +async fn note_engine_outcome( + ctx: &mut DagContext, + outcome: &gw_engines::EngineOutcome, + started: Option, +) { if outcome.terminal_error.is_some() { ctx.state .avail @@ -539,6 +564,9 @@ async fn note_engine_outcome(ctx: &mut DagContext, outcome: &gw_engines::EngineO .record(requested_model(ctx.request.model_param_v2.as_ref()), true); if let Some(account) = ctx.request.account.as_ref() { ctx.state.health.record_success(&account.name).await; + if let Some(started) = started { + ctx.state.latency.record(&account.name, started.elapsed()); + } } } diff --git a/crates/state/src/latency.rs b/crates/state/src/latency.rs new file mode 100644 index 0000000..5f62806 --- /dev/null +++ b/crates/state/src/latency.rs @@ -0,0 +1,62 @@ +//! Per-account call latency: an in-process EWMA the pool ranks same-priority +//! accounts by under `stability.latency_routing`. Each instance learns its own +//! view; an unknown or stale account ranks first so it gets sampled. + +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use dashmap::DashMap; + +// an account unseen this long is re-probed ahead of every measured one +const STALE_AFTER: Duration = Duration::from_secs(60); +// weight of the newest sample +const ALPHA: f64 = 0.2; + +/// Exponentially weighted call latency per account, in milliseconds; a clone shares the samples. +#[derive(Debug, Default, Clone)] +pub struct Latency { + samples: Arc>, +} + +impl Latency { + pub fn record(&self, name: &str, elapsed: Duration) { + let ms = elapsed.as_secs_f64() * 1e3; + let mut e = crate::slot_mut(&self.samples, name, || (ms, Instant::now())); + e.0 += ALPHA * (ms - e.0); + e.1 = Instant::now(); + } + + /// The ranking key: the EWMA, or 0 for an account never or not recently seen. + pub fn rank(&self, name: &str) -> f64 { + self.samples + .get(name) + .filter(|e| e.1.elapsed() < STALE_AFTER) + .map_or(0.0, |e| e.0) + } + + #[cfg(test)] + pub(crate) fn backdate(&self, name: &str, by: Duration) { + if let Some(mut e) = self.samples.get_mut(name) { + e.1 -= by; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn ewma_tracks_samples_and_stale_or_unknown_accounts_rank_first() { + let l = Latency::default(); + assert_eq!(l.rank("a"), 0.0); + l.record("a", Duration::from_millis(100)); + assert_eq!(l.rank("a"), 100.0); + l.record("a", Duration::from_millis(200)); + assert_eq!(l.rank("a"), 120.0); + l.backdate("a", STALE_AFTER); + assert_eq!(l.rank("a"), 0.0, "a stale sample counts as unknown"); + l.record("a", Duration::from_millis(50)); + assert_eq!(l.rank("a"), 106.0, "the EWMA survives the stale window"); + } +} diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index cdadd28..0ca7cfc 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -22,6 +22,7 @@ pub mod content; pub mod governance; pub mod health; pub mod keystore; +pub mod latency; pub mod store; pub mod thinking_signature; @@ -451,6 +452,7 @@ impl AccountPool { provider: Option<&str>, excluded: &[String], health: &dyn HealthStore, + latency: Option<&latency::Latency>, ) -> Option> { let candidates: Vec<&Arc> = self .accounts @@ -465,7 +467,7 @@ impl AccountPool { .filter(|(_, ok)| !ok) .map(|(a, _)| a.name.as_str()) .collect(); - self.select_with(p, provider, |name| { + self.select_with(p, provider, latency, |name| { excluded.iter().any(|e| e == name) || unhealthy.contains(&name) }) } @@ -482,6 +484,7 @@ impl AccountPool { &self, p: Protocol, provider: Option<&str>, + latency: Option<&latency::Latency>, is_excluded: impl Fn(&str) -> bool, ) -> Option> { let eligible = |a: &&Arc| serves(a, p, provider) && !is_excluded(&a.name); @@ -504,7 +507,16 @@ impl AccountPool { .iter() .filter(|a| eligible(a) && a.is_ptu() == has_ptu && a.priority == best) .collect(); - let idx = self.rr.fetch_add(1, Ordering::Relaxed) % top.len(); + let start = self.rr.fetch_add(1, Ordering::Relaxed) % top.len(); + let idx = match latency { + // the rotating start keeps ties (unknown or equal latency) round-robin + Some(latency) => (0..top.len()) + .map(|k| (start + k) % top.len()) + .map(|i| (latency.rank(&top[i].name), i)) + .min_by(|a, b| a.0.total_cmp(&b.0)) + .map_or(start, |(_, i)| i), + None => start, + }; Some(Arc::clone(top[idx])) } } @@ -747,6 +759,8 @@ pub struct GatewayState { pub alerts: Arc, /// Ten-minute, fail-open replay consistency cache for Anthropic thinking. pub thinking_signatures: ThinkingSignatureAudit, + /// Per-account call latency for `stability.latency_routing`; per instance. + pub latency: latency::Latency, } impl Default for GatewayState { @@ -763,6 +777,7 @@ impl Default for GatewayState { avail: Arc::new(avail::MemoryAvail::default()), alerts: Arc::new(alerts::AlertBus::default()), thinking_signatures: ThinkingSignatureAudit::new(), + latency: latency::Latency::default(), } } } @@ -880,6 +895,7 @@ impl GatewayState { avail: prev.avail.clone(), alerts: prev.alerts.clone(), thinking_signatures: prev.thinking_signatures.clone(), + latency: prev.latency.clone(), }) } } @@ -1038,6 +1054,51 @@ mod tests { GatewayState::from_config(&GatewayConfig::embedded_default().unwrap()) } + #[tokio::test] + async fn pool_ranks_a_tier_by_latency_when_asked() { + let yaml = "listen: {host: h, port: 1}\naccounts: [{name: fast, provider: p, protocols: ['openai-chat']}, {name: slow, provider: p, protocols: ['openai-chat']}]"; + let s = GatewayState::from_config(&GatewayConfig::from_yaml(yaml).unwrap()); + async fn pick(s: &GatewayState, latency: Option<&latency::Latency>) -> String { + s.pool + .select_healthy( + Protocol::OpenaiChat, + Some("p"), + &[], + s.health.as_ref(), + latency, + ) + .await + .unwrap() + .name + .clone() + } + let lat = latency::Latency::default(); + lat.record("fast", Duration::from_millis(20)); + lat.record("slow", Duration::from_millis(200)); + for _ in 0..4 { + assert_eq!(pick(&s, Some(&lat)).await, "fast"); + } + let unranked = [pick(&s, None).await, pick(&s, None).await]; + assert!( + unranked.contains(&"fast".to_owned()) && unranked.contains(&"slow".to_owned()), + "round-robin without the ranker: {unranked:?}" + ); + lat.backdate("slow", Duration::from_secs(60)); + assert_eq!( + pick(&s, Some(&lat)).await, + "slow", + "an account unseen for a minute is probed again" + ); + lat.record("slow", Duration::from_millis(200)); + assert_eq!(pick(&s, Some(&lat)).await, "fast"); + let fresh = latency::Latency::default(); + let probes = [pick(&s, Some(&fresh)).await, pick(&s, Some(&fresh)).await]; + assert!( + probes.contains(&"fast".to_owned()) && probes.contains(&"slow".to_owned()), + "unknown accounts round-robin: {probes:?}" + ); + } + fn ak_info(ak: &str) -> AkInfo { AkInfo { ak_id: access_key_fingerprint(ak).into(), @@ -1348,19 +1409,19 @@ mod tests { let h = s.health.as_ref(); let a = s .pool - .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h) + .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h, None) .await .unwrap(); let b = s .pool - .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h) + .select_healthy(Protocol::OpenaiChat, Some("openai"), &[], h, None) .await .unwrap(); assert_eq!(a.name, "mock-openai-1"); assert_eq!(b.name, "mock-openai-1"); assert_eq!( s.pool - .select_healthy(Protocol::AnthropicMessages, None, &[], h) + .select_healthy(Protocol::AnthropicMessages, None, &[], h, None) .await .unwrap() .name, @@ -1368,7 +1429,7 @@ mod tests { ); assert!( s.pool - .select_healthy(Protocol::Video, Some("nonexistent"), &[], h) + .select_healthy(Protocol::Video, Some("nonexistent"), &[], h, None) .await .is_none() ); @@ -1408,7 +1469,7 @@ mod tests { let h = s.health.as_ref(); let first = s .pool - .select_healthy(Protocol::OpenaiChat, Some("tencent"), &[], h) + .select_healthy(Protocol::OpenaiChat, Some("tencent"), &[], h, None) .await .unwrap(); assert_eq!(first.name, "mock-hunyuan-ptu-down"); @@ -1420,6 +1481,7 @@ mod tests { Some("tencent"), &["mock-hunyuan-ptu-down".into()], h, + None, ) .await .unwrap(); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index e644b75..0a8a1a3 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -399,6 +399,7 @@ async fn realtime_ws( served_conf.and_then(|m| m.provider.as_deref()), &[], snap.state.health.as_ref(), + None, ) .await; let Some(account) = account else { diff --git a/docs/configuration.md b/docs/configuration.md index 847e3b6..06a0d4d 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -256,6 +256,7 @@ stability: unstable_error_rate: 0.1 # window error rate that reports `unstable` unavailable_error_rate: 0.5 # ... and `unavailable` availability_min_samples: 20 # fewer samples than this reports `no_data` + latency_routing: false # rank same-priority accounts by observed call latency instead of round-robin products: - name: myproduct diff --git a/docs/providers.md b/docs/providers.md index 14ece6c..3e30ad9 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -171,7 +171,12 @@ the total budget, while a stalled stream fails at the gap. Multiple accounts can serve the same protocol. Selection is by `priority` (lower first), round-robin within a tie, with PTU-tier accounts preferred over -paygo. On an upstream 5xx the failed account is excluded and another is tried +paygo. With `stability.latency_routing: true` the tie is ranked instead by +each account's observed call latency — an exponentially weighted average of +completed calls, kept per instance — so the fastest account of a tier takes +the traffic while an account never or not recently (60 s) sampled ranks first +and gets probed; equal ranks stay round-robin. The realtime surface keeps +round-robin. On an upstream 5xx the failed account is excluded and another is tried once (a PTU→paygo switch is flagged `ptu_spillover`). Consecutive failures put an account into cooldown (`stability.failure_threshold` / `cooldown_seconds`), and it auto-recovers on expiry. A streaming response that already sent bytes to