Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion README.md
Original file line number Diff line number Diff line change
Expand Up @@ -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/*`)
Expand Down
4 changes: 4 additions & 0 deletions crates/config/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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(),
}
}
Expand Down
36 changes: 32 additions & 4 deletions crates/dag/src/nodes.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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!(
Expand All @@ -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
Expand All @@ -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 {
Expand All @@ -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(())
Expand All @@ -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<std::time::Instant> {
ctx.cfg
.stability
.latency_routing
.then(std::time::Instant::now)
}

async fn note_engine_outcome(
ctx: &mut DagContext,
outcome: &gw_engines::EngineOutcome,
started: Option<std::time::Instant>,
) {
if outcome.terminal_error.is_some() {
ctx.state
.avail
Expand All @@ -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());
}
}
}

Expand Down
62 changes: 62 additions & 0 deletions crates/state/src/latency.rs
Original file line number Diff line number Diff line change
@@ -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<DashMap<String, (f64, Instant)>>,
}

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");
}
}
76 changes: 69 additions & 7 deletions crates/state/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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;

Expand Down Expand Up @@ -451,6 +452,7 @@ impl AccountPool {
provider: Option<&str>,
excluded: &[String],
health: &dyn HealthStore,
latency: Option<&latency::Latency>,
) -> Option<Arc<Account>> {
let candidates: Vec<&Arc<Account>> = self
.accounts
Expand All @@ -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)
})
}
Expand All @@ -482,6 +484,7 @@ impl AccountPool {
&self,
p: Protocol,
provider: Option<&str>,
latency: Option<&latency::Latency>,
is_excluded: impl Fn(&str) -> bool,
) -> Option<Arc<Account>> {
let eligible = |a: &&Arc<Account>| serves(a, p, provider) && !is_excluded(&a.name);
Expand All @@ -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]))
}
}
Expand Down Expand Up @@ -747,6 +759,8 @@ pub struct GatewayState {
pub alerts: Arc<alerts::AlertBus>,
/// 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 {
Expand All @@ -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(),
}
}
}
Expand Down Expand Up @@ -880,6 +895,7 @@ impl GatewayState {
avail: prev.avail.clone(),
alerts: prev.alerts.clone(),
thinking_signatures: prev.thinking_signatures.clone(),
latency: prev.latency.clone(),
})
}
}
Expand Down Expand Up @@ -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(),
Expand Down Expand Up @@ -1348,27 +1409,27 @@ 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,
"mock-anthropic-1"
);
assert!(
s.pool
.select_healthy(Protocol::Video, Some("nonexistent"), &[], h)
.select_healthy(Protocol::Video, Some("nonexistent"), &[], h, None)
.await
.is_none()
);
Expand Down Expand Up @@ -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");
Expand All @@ -1420,6 +1481,7 @@ mod tests {
Some("tencent"),
&["mock-hunyuan-ptu-down".into()],
h,
None,
)
.await
.unwrap();
Expand Down
1 change: 1 addition & 0 deletions crates/views/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand Down
1 change: 1 addition & 0 deletions docs/configuration.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
7 changes: 6 additions & 1 deletion docs/providers.md
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down