diff --git a/Cargo.lock b/Cargo.lock index 2db4957..f8f5f2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1007,7 +1007,12 @@ dependencies = [ name = "gw-task" version = "0.1.1" dependencies = [ + "futures", + "gw-config", + "gw-consts", "gw-state", + "reqwest", + "serde_json", "tokio", "tracing", ] diff --git a/conf/gateway.yaml b/conf/gateway.yaml index 385b4ec..aa60cbb 100644 --- a/conf/gateway.yaml +++ b/conf/gateway.yaml @@ -142,6 +142,13 @@ models: protocol: ernie input_price_per_1k_micros: 300 output_price_per_1k_micros: 900 + - name: text-moderation # content moderation (OpenAI moderations shape) + protocol: moderations + provider: openai + - name: rerank-mini # document rerank (Cohere/Jina-compatible shape) + protocol: rerank + provider: openai + input_price_per_1k_micros: 20 - name: qpm-mini # e2e: model-level QPM limit (2 per minute) protocol: openai-chat provider: openai @@ -188,7 +195,7 @@ accounts: priority: 1 cost_input_price_per_1k_micros: 100 # e2e: margin accounting (vendor cost) cost_output_price_per_1k_micros: 400 - protocols: ["openai-chat", "completions"] + protocols: ["openai-chat", "completions", "moderations", "rerank"] - name: mock-openai-2 provider: openai priority: 2 @@ -276,3 +283,10 @@ stability: products: - name: qpmprod qpm: 2 + +# Abuse auto-suspension and the alert webhook (uncomment to enable): +# abuse: +# tiers: [{rejects: 20, suspend_hours: 2}, {rejects: 30, suspend_hours: 24}] +# alerts: +# webhook_url_env: GW_ALERT_WEBHOOK +# dedup_seconds: 300 diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 8b167d5..b2294c5 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -29,7 +29,7 @@ pub enum ConfigError { #[error("model `{model}` references unknown protocol `{wire}`")] UnknownModelMapping { model: String, wire: String }, #[error( - "provider `{provider}` has unknown kind `{kind}` (known: openai, anthropic, gemini, deepseek, openrouter)" + "provider `{provider}` has unknown kind `{kind}` (known: openai, anthropic, gemini, deepseek, openrouter, moonshot, siliconflow)" )] UnknownProviderKind { provider: String, kind: String }, #[error("model `{model}` references unknown provider `{provider}`")] @@ -143,6 +143,30 @@ pub struct VariantConf { pub weight: u32, } +/// Cumulative-weight pick over a model's variants, keyed by a stable hash so +/// every instance (REST DAG and realtime handshake alike) maps the same key +/// to the same bucket with no shared state. +pub fn pick_variant<'a>(variants: &'a [VariantConf], key: &str) -> &'a VariantConf { + let total: u64 = variants.iter().map(|v| u64::from(v.weight)).sum(); + let mut roll = fnv1a(key) % total.max(1); + for v in variants { + if roll < u64::from(v.weight) { + return v; + } + roll -= u64::from(v.weight); + } + // unreachable (weights validated >= 1); quiets the type checker + &variants[0] +} + +/// FNV-1a 64: deterministic across processes and releases (std's hasher is +/// neither), which the fleet-consistent sticky mapping depends on. +fn fnv1a(s: &str) -> u64 { + s.bytes().fold(0xcbf2_9ce4_8422_2325, |h, b| { + (h ^ u64::from(b)).wrapping_mul(0x0000_0100_0000_01b3) + }) +} + /// Per-component billing weights relative to the model's unit prices /// (e.g. cache reads at 0.1, cache writes at 1.25). Missing fields stay 1.0. #[derive(Debug, Clone, Copy, Deserialize)] @@ -355,6 +379,55 @@ impl Default for StabilityConf { } } +/// One automatic-suspension tier: this many admission rejections in a day +/// suspend the key for `suspend_hours`. +#[derive(Debug, Clone, Deserialize)] +pub struct AbuseTier { + pub rejects: i64, + pub suspend_hours: u64, +} + +/// Automatic abuse suspension; empty tiers = disabled (no counting at all). +/// Tiers must ascend in both fields (validated) — the highest reject +/// threshold met wins. Counts REST admission rejections only: the realtime +/// gate denies per turn without feeding this counter. +#[derive(Debug, Clone, Default, Deserialize)] +pub struct AbuseConf { + #[serde(default)] + pub tiers: Vec, +} + +/// Outbound alert webhook. Advisory: delivery failures are logged and dropped. +#[derive(Debug, Clone, Deserialize)] +pub struct AlertsConf { + /// Env var naming the webhook URL; empty = alerts disabled. + #[serde(default)] + pub webhook_url_env: String, + /// Repeat alerts for the same (kind, subject) are muted this long. + #[serde(default = "default_alert_dedup_seconds")] + pub dedup_seconds: u64, +} + +impl AlertsConf { + /// The webhook target, if configured and the env var is set non-empty. + pub fn webhook_url(&self) -> Option { + token_from_env(&self.webhook_url_env) + } +} + +impl Default for AlertsConf { + fn default() -> Self { + Self { + webhook_url_env: String::new(), + dedup_seconds: default_alert_dedup_seconds(), + } + } +} + +fn default_alert_dedup_seconds() -> u64 { + 300 +} + /// Durable-record backend selection. #[derive(Debug, Clone, Default, Deserialize)] pub struct StorageConf { @@ -515,6 +588,7 @@ fn provider_preset(kind: &str) -> Option { "responses", "completions", "realtime", + "moderations", ], default_model_wire: "openai-chat", }, @@ -539,6 +613,16 @@ fn provider_preset(kind: &str) -> Option { wires: &["openai-chat"], default_model_wire: "openai-chat", }, + "moonshot" => ProviderPreset { + endpoint: "https://api.moonshot.cn", + wires: &["openai-chat"], + default_model_wire: "openai-chat", + }, + "siliconflow" => ProviderPreset { + endpoint: "https://api.siliconflow.cn", + wires: &["openai-chat", "embeddings", "rerank"], + default_model_wire: "openai-chat", + }, _ => return None, }) } @@ -568,6 +652,12 @@ pub struct GatewayConfig { /// Admin surface gate (dynamic config reload / key management). #[serde(default)] pub admin: AdminConf, + /// Automatic abuse suspension tiers (empty = off). + #[serde(default)] + pub abuse: AbuseConf, + /// Outbound alert webhook (unset = off). + #[serde(default)] + pub alerts: AlertsConf, /// Trust `x-real-ip` / `x-forwarded-for` for the audit source IP. Off by /// default: the audit records the real TCP peer, which a client can't forge. /// Enable only when a trusted proxy fronts the gateway and sets those headers. @@ -719,11 +809,6 @@ impl GatewayConfig { variant: v.model.clone(), reason, }; - // realtime pins its model at handshake, never traversing - // VariantSelect — accepting this would silently serve the parent - if m.protocol() == Some(Protocol::Realtime) { - return Err(bad("variants are not supported on realtime models")); - } if v.weight == 0 { return Err(bad("weight must be >= 1")); } @@ -792,6 +877,14 @@ impl GatewayConfig { return Err(neg_limit(format!("product {}", p.name))); } } + for (i, t) in self.abuse.tiers.iter().enumerate() { + let ascending = self.abuse.tiers[..i] + .iter() + .all(|p| p.rejects < t.rejects && p.suspend_hours <= t.suspend_hours); + if t.rejects <= 0 || t.suspend_hours == 0 || !ascending { + return Err(neg_limit("abuse tier (ascending rejects/hours)".to_owned())); + } + } let st = &self.stability; let bad_rate = |v: f64| !v.is_finite() || !(0.0..=1.0).contains(&v); // upper bound mirrors the store's one-hour retention — a longer window @@ -845,6 +938,16 @@ impl GatewayConfig { }); } } + // a colon in an ak would collide with the prefixed governance keyspaces + // (`abuse:{ak}` above all — a key named `abuse:X` could force-suspend X) + for k in &self.access_keys { + if k.ak.contains(':') { + return Err(ConfigError::DuplicateName { + kind: "access_key (':' not allowed)", + name: k.ak.clone(), + }); + } + } // a typo'd tenant would silently fall back to the unrestricted default — reject at load for k in &self.access_keys { if !self.is_known_tenant(&k.tenant) { @@ -1087,9 +1190,12 @@ listen: {host: 127.0.0.1, port: 0} providers: - {name: deepseek, kind: deepseek, api_key_env: DEEPSEEK_KEY} - {name: openrouter, kind: openrouter, api_key_env: OPENROUTER_KEY} + - {name: kimi, kind: moonshot, api_key_env: MOONSHOT_KEY} + - {name: sf, kind: siliconflow, api_key_env: SF_KEY} models: - {name: deepseek-chat, provider: deepseek} - {name: some-model, provider: openrouter} + - {name: kimi-k2, provider: kimi} "#; let cfg = GatewayConfig::from_yaml(yaml).unwrap(); assert_eq!( @@ -1098,6 +1204,14 @@ models: ); let ds = cfg.accounts.iter().find(|a| a.name == "deepseek").unwrap(); assert_eq!(ds.endpoint, "https://api.deepseek.com"); + assert_eq!( + cfg.find_model("kimi-k2").unwrap().protocol(), + Some(Protocol::OpenaiChat) + ); + let kimi = cfg.accounts.iter().find(|a| a.name == "kimi").unwrap(); + assert_eq!(kimi.endpoint, "https://api.moonshot.cn"); + let sf = cfg.accounts.iter().find(|a| a.name == "sf").unwrap(); + assert!(sf.protocols.iter().any(|w| w == "rerank")); let orr = cfg .accounts .iter() @@ -1356,11 +1470,30 @@ tenants: [{name: t1}, {name: t1}] let rt = "listen: {host: h, port: 1}\nmodels: [{name: r1, protocol: realtime, variants: [{model: r2, weight: 1}]}, {name: r2, protocol: realtime}]"; assert!( - matches!( - GatewayConfig::from_yaml(rt), - Err(ConfigError::BadVariant { .. }) - ), - "realtime variants are rejected at load" + GatewayConfig::from_yaml(rt).is_ok(), + "realtime variants pick at the session handshake" + ); + + let variants = [ + VariantConf { + model: "a".into(), + weight: 9, + }, + VariantConf { + model: "b".into(), + weight: 1, + }, + ]; + let first = pick_variant(&variants, "user-1").model.clone(); + for _ in 0..10 { + assert_eq!(pick_variant(&variants, "user-1").model, first, "sticky"); + } + let hits = (0..1000) + .filter(|i| pick_variant(&variants, &format!("user-{i}")).model == "b") + .count(); + assert!( + (40..250).contains(&hits), + "10% weight took {hits}/1000 keys" ); for bad in [ @@ -1381,6 +1514,33 @@ tenants: [{name: t1}, {name: t1}] let ok = "listen: {host: h, port: 1}\nstability: {availability_window_minutes: 60}"; assert!(GatewayConfig::from_yaml(ok).is_ok()); + for bad in [ + "abuse: {tiers: [{rejects: 0, suspend_hours: 2}]}", + "abuse: {tiers: [{rejects: 30, suspend_hours: 2}, {rejects: 20, suspend_hours: 24}]}", + "abuse: {tiers: [{rejects: 20, suspend_hours: 24}, {rejects: 30, suspend_hours: 2}]}", + ] { + let yaml = format!("listen: {{host: h, port: 1}}\n{bad}"); + assert!( + matches!( + GatewayConfig::from_yaml(&yaml), + Err(ConfigError::NegativeLimit { .. }) + ), + "non-ascending or zero abuse tier is rejected: {bad}" + ); + } + let colon_ak = "listen: {host: h, port: 1}\naccess_keys: [{ak: 'abuse:x', product: p, qps: 1, daily_token_quota: 1}]"; + assert!( + matches!( + GatewayConfig::from_yaml(colon_ak), + Err(ConfigError::DuplicateName { .. }) + ), + "':' in an ak collides with governance prefixes" + ); + let tiers = "listen: {host: h, port: 1}\nabuse: {tiers: [{rejects: 20, suspend_hours: 2}, {rejects: 30, suspend_hours: 24}]}\nalerts: {webhook_url_env: GW_ALERT_URL}"; + let cfg = GatewayConfig::from_yaml(tiers).unwrap(); + assert_eq!(cfg.abuse.tiers.len(), 2); + assert_eq!(cfg.alerts.dedup_seconds, 300, "default dedup window"); + let neg_quota = "listen: {host: h, port: 1}\naccess_keys: [{ak: k1, product: p, qps: 1, daily_token_quota: -5}]"; assert!( matches!( diff --git a/crates/consts/src/protocol.rs b/crates/consts/src/protocol.rs index cd522de..28d11b7 100644 --- a/crates/consts/src/protocol.rs +++ b/crates/consts/src/protocol.rs @@ -47,6 +47,10 @@ pub enum Protocol { AwsLlama, /// Alibaba DashScope native (input.messages/parameters/output.choices) Dashscope, + /// content moderation (OpenAI moderations shape) + Moderations, + /// document rerank (Cohere/Jina-compatible shape) + Rerank, } impl Protocol { @@ -70,6 +74,8 @@ impl Protocol { Protocol::AwsCohere, Protocol::AwsLlama, Protocol::Dashscope, + Protocol::Moderations, + Protocol::Rerank, ]; pub const fn as_str(self) -> &'static str { @@ -93,6 +99,8 @@ impl Protocol { Protocol::AwsCohere => "aws-cohere", Protocol::AwsLlama => "aws-llama", Protocol::Dashscope => "dashscope", + Protocol::Moderations => "moderations", + Protocol::Rerank => "rerank", } } @@ -116,7 +124,7 @@ mod tests { for &p in Protocol::ALL { assert_eq!(Protocol::from_wire(p.as_str()), Some(p)); } - assert_eq!(Protocol::ALL.len(), 19); + assert_eq!(Protocol::ALL.len(), 21); assert!(Protocol::from_wire("nope").is_none()); } } diff --git a/crates/dag/src/executor.rs b/crates/dag/src/executor.rs index e6e0412..48e2fed 100644 --- a/crates/dag/src/executor.rs +++ b/crates/dag/src/executor.rs @@ -143,6 +143,7 @@ mod tests { tokens_per_minute: None, expires_at_epoch_secs: None, banned: false, + suspended_until_epoch_secs: None, model_quotas: Default::default(), }, ) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 14883a2..1e085c8 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -55,21 +55,20 @@ impl DagNode for ModelQuotaGate { if under { return Ok(()); } - let fallback = ctx - .cfg - .find_tenant(&ctx.ak.tenant) - .and_then(|t| t.fallback_model.clone()); - match fallback { - Some(fb) if fb != requested => { - ctx.decide( - "model_quota", - format!("{requested} over {limit}, serving {fb}"), - ); - if let Some(param) = ctx.request.model_param_v2.as_mut() { - param.fallback_from = Some(requested); - param.model_name = fb; - } + let (cfg, tenant) = (&ctx.cfg, &ctx.ak.tenant); + let swapped = ctx + .request + .model_param_v2 + .as_mut() + .map(|p| admission::swap_to_fallback(cfg, tenant, p)); + match swapped { + Some(admission::FallbackSwap::Swapped(from, fb)) => { + ctx.decide("model_quota", format!("{from} over {limit}, serving {fb}")) } + Some(admission::FallbackSwap::AlreadyServing) => ctx.decide( + "model_quota", + format!("{requested} over {limit}, already the fallback"), + ), _ => ctx.decide( "model_quota", format!("{requested} over {limit}, no fallback"), @@ -182,7 +181,7 @@ impl DagNode for VariantSelect { "" => ctx.request.request_id.as_str(), user => user, }; - let target = pick_variant(&conf.variants, key).model.clone(); + let target = gw_config::pick_variant(&conf.variants, key).model.clone(); if target == param.model_name { ctx.decide("variant_select", format!("{target} (self)")); return Ok(()); @@ -196,32 +195,6 @@ impl DagNode for VariantSelect { } } -/// Cumulative-weight pick, keyed by a stable hash so every instance maps the -/// same key to the same bucket with no shared state. -fn pick_variant<'a>( - variants: &'a [gw_config::VariantConf], - key: &str, -) -> &'a gw_config::VariantConf { - let total: u64 = variants.iter().map(|v| u64::from(v.weight)).sum(); - let mut roll = fnv1a(key) % total.max(1); - for v in variants { - if roll < u64::from(v.weight) { - return v; - } - roll -= u64::from(v.weight); - } - // unreachable (weights validated >= 1); quiets the type checker - &variants[0] -} - -/// FNV-1a 64: deterministic across processes and releases (std's hasher is -/// neither), which the fleet-consistent sticky mapping depends on. -fn fnv1a(s: &str) -> u64 { - s.bytes().fold(0xcbf2_9ce4_8422_2325, |h, b| { - (h ^ u64::from(b)).wrapping_mul(0x0000_0100_0000_01b3) - }) -} - /// preprocess/cache_lookup: request-level TTL cache. On a hit the outcome is /// produced directly and the downstream nodes all short-circuit. pub struct CacheLookup; @@ -544,20 +517,7 @@ impl DagNode for CallEngine { .account .clone() .ok_or_else(|| GatewayError::internal("call_engine without an account"))?; - if ctx - .state - .health - .record_failure(&failed.name, threshold, cooldown) - .await - { - ctx.decide( - "account_health", - format!( - "{} entered cooldown ({}s)", - failed.name, ctx.cfg.stability.cooldown_seconds - ), - ); - } + note_failure(ctx, &failed.name, threshold, cooldown).await; let provider = model_provider(ctx); let next = ctx .state @@ -602,10 +562,7 @@ impl DagNode for CallEngine { ctx.state .avail .record(requested_model(ctx.request.model_param_v2.as_ref()), false); - ctx.state - .health - .record_failure(&next.name, threshold, cooldown) - .await; + note_failure(ctx, &next.name, threshold, cooldown).await; Err(e) } } @@ -615,6 +572,35 @@ impl DagNode for CallEngine { } } +/// Record an account failure; on the cooldown transition, alert and note the +/// decision — one implementation for both engine attempts, so the fleet-backed +/// `record_failure` call and its alerting can't drift apart. +async fn note_failure( + ctx: &mut DagContext, + account: &str, + threshold: usize, + cooldown: std::time::Duration, +) { + if !ctx + .state + .health + .record_failure(account, threshold, cooldown) + .await + { + return; + } + let secs = ctx.cfg.stability.cooldown_seconds; + ctx.state.alerts.emit( + "account_cooldown", + account.to_owned(), + format!("{threshold} consecutive failures; cooling {secs}s"), + ); + ctx.decide( + "account_health", + format!("{account} entered cooldown ({secs}s)"), + ); +} + /// post_process/common_usage: RawUsageJSON -> CommonUsage. pub struct CommonUsageNode; @@ -923,31 +909,6 @@ mod tests { assert_eq!(t.total(), 100); } - #[test] - fn variant_pick_is_sticky_and_weighted() { - let variants = vec![ - gw_config::VariantConf { - model: "a".into(), - weight: 9, - }, - gw_config::VariantConf { - model: "b".into(), - weight: 1, - }, - ]; - let first = super::pick_variant(&variants, "user-1").model.clone(); - for _ in 0..10 { - assert_eq!(super::pick_variant(&variants, "user-1").model, first); - } - let hits = (0..1000) - .filter(|i| super::pick_variant(&variants, &format!("user-{i}")).model == "b") - .count(); - assert!( - (40..250).contains(&hits), - "10% weight took {hits}/1000 keys" - ); - } - #[test] fn reserve_estimate_saturates_on_hostile_max_tokens() { use gw_models::params::ChatParams; diff --git a/crates/engines/src/factory.rs b/crates/engines/src/factory.rs index 0a098d2..cbd9399 100644 --- a/crates/engines/src/factory.rs +++ b/crates/engines/src/factory.rs @@ -10,8 +10,8 @@ use crate::bespoke::{CohereEngine, DashScopeEngine, ErnieEngine, LlamaEngine, Mi use crate::claude_engine::ClaudeEngine; use crate::engine::ModelEngine; use crate::families::{ - AudioEngine, AudioKind, CompletionsEngine, EmbeddingsEngine, ImageEngine, PassthroughEngine, - ResponsesEngine, SearchEngine, VertexEngine, VideoEngine, + AudioEngine, AudioKind, CompletionsEngine, EmbeddingsEngine, ImageEngine, ModerationsEngine, + PassthroughEngine, RerankEngine, ResponsesEngine, SearchEngine, VertexEngine, VideoEngine, }; use crate::openai_engine::OpenAiEngine; use crate::transport::SharedTransport; @@ -38,6 +38,8 @@ pub fn get_engine( Protocol::Audio => Box::new(AudioEngine::new(request, transport, AudioKind::Other)), Protocol::Video => Box::new(VideoEngine::new(request, transport)), Protocol::Search => Box::new(SearchEngine::new(request, transport)), + Protocol::Moderations => Box::new(ModerationsEngine::new(request, transport)), + Protocol::Rerank => Box::new(RerankEngine::new(request, transport)), Protocol::Passthrough => Box::new(PassthroughEngine::new(request, transport)), Protocol::Ernie => Box::new(ErnieEngine::new(request, transport)), Protocol::MinimaxV1 => Box::new(MinimaxV1Engine::new(request, transport)), diff --git a/crates/engines/src/families.rs b/crates/engines/src/families.rs index 585086b..c68fc0c 100644 --- a/crates/engines/src/families.rs +++ b/crates/engines/src/families.rs @@ -436,13 +436,20 @@ impl ModelEngine for AudioEngine { ("/v1/audio/speech", b) } AudioKind::Stt => { - let (audio, language) = match ¶m.typed { - Some(TypedParams::AudioStt(p)) => (p.audio_b64.as_str(), p.language.as_deref()), - _ => ("", None), + let (audio, language, translate) = match ¶m.typed { + Some(TypedParams::AudioStt(p)) => { + (p.audio_b64.as_str(), p.language.as_deref(), p.translate) + } + _ => ("", None, false), }; require_non_empty(audio, "stt audio_b64")?; + let path = if translate { + "/v1/audio/translations" + } else { + "/v1/audio/transcriptions" + }; ( - "/v1/audio/transcriptions", + path, json!({"model": param.model_name, "audio_b64": audio, "language": language}), ) } @@ -541,6 +548,101 @@ impl ModelEngine for SearchEngine { } } +base_engine!(ModerationsEngine); + +#[async_trait::async_trait] +impl ModelEngine for ModerationsEngine { + /// OpenAI moderations shape: `{model, input: [..]}` → per-input verdicts. + async fn run(&self) -> GResult { + let param = self.base.param()?; + let Some(TypedParams::Moderation(p)) = ¶m.typed else { + return Err(GatewayError::bad_request("moderations params are required")); + }; + if p.input.is_empty() { + return Err(GatewayError::bad_request( + "moderations input must not be empty", + )); + } + let body = json!({"model": param.model_name, "input": p.input}); + let (status, v) = self + .base + .round_trip( + &format!( + "{}/v1/moderations", + self.base.base_url("mock://api.openai.com") + ), + body, + ) + .await?; + let flagged = v["results"] + .as_array() + .map(|rs| { + rs.iter() + .filter(|r| r["flagged"].as_bool().unwrap_or(false)) + .count() + }) + .unwrap_or(0); + Ok(family_outcome( + format!("{flagged} flagged"), + ¶m.model_name, + v, + status, + )) + } +} + +base_engine!(RerankEngine); + +#[async_trait::async_trait] +impl ModelEngine for RerankEngine { + /// Cohere/Jina-compatible rerank: `{model, query, documents, top_n?}` → + /// `{results: [{index, relevance_score}]}`. + async fn run(&self) -> GResult { + let param = self.base.param()?; + let Some(TypedParams::Rerank(p)) = ¶m.typed else { + return Err(GatewayError::bad_request("rerank params are required")); + }; + require_non_empty(&p.query, "rerank query")?; + if p.documents.is_empty() { + return Err(GatewayError::bad_request( + "rerank documents must not be empty", + )); + } + let mut body = json!({ + "model": param.model_name, + "query": p.query, + "documents": p.documents, + }); + if let Some(n) = p.top_n { + body["top_n"] = json!(n); + } + let (status, v) = self + .base + .round_trip( + &format!("{}/v1/rerank", self.base.base_url("mock://api.vendor.com")), + body, + ) + .await?; + let n = v["results"].as_array().map(Vec::len).unwrap_or(0); + let tokens = rerank_tokens(&v); + let mut out = family_outcome(format!("{n} results"), ¶m.model_name, v, status); + out.response.prompt_tokens = tokens; + out.response.total_tokens = tokens; + Ok(out) + } +} + +/// Rerank usage across the two wire dialects: Jina-style `usage.total_tokens`, +/// else Cohere/SiliconFlow-style `meta.tokens.{input,output}_tokens`. +fn rerank_tokens(v: &Value) -> i64 { + let total = crate::engine::tok(&v["usage"]["total_tokens"]); + if total > 0 { + return total; + } + crate::engine::tok(&v["meta"]["tokens"]["input_tokens"]) + .saturating_add(crate::engine::tok(&v["meta"]["tokens"]["output_tokens"])) +} + base_engine!(PassthroughEngine); #[async_trait::async_trait] @@ -971,6 +1073,76 @@ mod tests { assert!(out.response.response_v2.is_some()); } + #[tokio::test] + async fn moderations_flags_and_validates() { + let e = ModerationsEngine::new( + req( + Protocol::Moderations, + "text-moderation", + Some(TypedParams::Moderation(gw_models::ModerationParams { + input: vec!["fine".into(), "really unsafe".into()], + })), + ), + t(), + ); + let out = e.run().await.unwrap(); + assert_eq!(out.response.message, "1 flagged"); + assert_eq!( + out.response.response_v2.unwrap()["results"][1]["flagged"], + true + ); + + for typed in [ + None, + Some(TypedParams::Moderation(gw_models::ModerationParams { + input: vec![], + })), + ] { + let e = ModerationsEngine::new(req(Protocol::Moderations, "m", typed), t()); + assert_eq!(e.run().await.unwrap_err().http_status, 400); + } + } + + #[tokio::test] + async fn rerank_orders_and_validates() { + let params = |query: &str, documents: Vec<&str>| { + Some(TypedParams::Rerank(gw_models::RerankParams { + query: query.into(), + documents: documents.into_iter().map(str::to_owned).collect(), + top_n: Some(1), + })) + }; + let e = RerankEngine::new( + req( + Protocol::Rerank, + "rerank-mini", + params("rust gateway", vec!["cooking", "a rust gateway"]), + ), + t(), + ); + let out = e.run().await.unwrap(); + assert_eq!(out.response.message, "1 results"); + let v = out.response.response_v2.unwrap(); + assert_eq!(v["results"][0]["index"], 1, "the matching doc ranks first"); + assert!(out.response.total_tokens > 0, "usage flows from the vendor"); + + for typed in [None, params("", vec!["doc"]), params("query", vec![])] { + let e = RerankEngine::new(req(Protocol::Rerank, "m", typed), t()); + assert_eq!(e.run().await.unwrap_err().http_status, 400); + } + } + + #[test] + fn rerank_tokens_reads_both_dialects() { + assert_eq!(rerank_tokens(&json!({"usage": {"total_tokens": 7}})), 7); + assert_eq!( + rerank_tokens(&json!({"meta": {"tokens": {"input_tokens": 5, "output_tokens": 2}}})), + 7, + "SiliconFlow/Cohere meta.tokens shape bills too" + ); + assert_eq!(rerank_tokens(&json!({})), 0); + } + #[tokio::test] async fn audio_tts_and_stt() { let tts = AudioEngine::new( @@ -1002,6 +1174,7 @@ mod tests { Some(TypedParams::AudioStt(SttParams { audio_b64: "TU9DSw==".into(), language: Some("en".into()), + translate: false, })), ), t(), diff --git a/crates/engines/src/lib.rs b/crates/engines/src/lib.rs index f121cd8..b61c1cd 100644 --- a/crates/engines/src/lib.rs +++ b/crates/engines/src/lib.rs @@ -27,8 +27,8 @@ pub use claude_engine::ClaudeEngine; pub use engine::{EngineOutcome, ModelEngine, StreamChunk, vendor_error}; pub use factory::get_engine; pub use families::{ - AudioEngine, AudioKind, CompletionsEngine, EmbeddingsEngine, ImageEngine, PassthroughEngine, - ResponsesEngine, SearchEngine, VertexEngine, VideoEngine, + AudioEngine, AudioKind, CompletionsEngine, EmbeddingsEngine, ImageEngine, ModerationsEngine, + PassthroughEngine, RerankEngine, ResponsesEngine, SearchEngine, VertexEngine, VideoEngine, }; pub use openai_engine::{OpenAiEngine, merge_tool_call_fragments}; pub use sse::SseDecoder; diff --git a/crates/engines/src/transport.rs b/crates/engines/src/transport.rs index c7c4cf9..d738411 100644 --- a/crates/engines/src/transport.rs +++ b/crates/engines/src/transport.rs @@ -467,12 +467,59 @@ impl MockTransport { Self::ok_json(json!({"created": MOCK_CREATED, "data": data})) } + fn moderations_reply(&self, req: &UpstreamRequest) -> GResult { + let body = Self::parse(&req.body, "moderations")?; + let results: Vec = body["input"] + .as_array() + .cloned() + .unwrap_or_default() + .iter() + .map(|i| { + // deterministic: the literal token "unsafe" flags + let flagged = i.as_str().is_some_and(|s| s.contains("unsafe")); + json!({"flagged": flagged, "categories": {"unsafe": flagged}}) + }) + .collect(); + Self::ok_json(json!({"id": "modr-mock", "model": body["model"], "results": results})) + } + + fn rerank_reply(&self, req: &UpstreamRequest) -> GResult { + let body = Self::parse(&req.body, "rerank")?; + let query = body["query"].as_str().unwrap_or_default().to_lowercase(); + let mut scored: Vec<(usize, f64)> = body["documents"] + .as_array() + .cloned() + .unwrap_or_default() + .iter() + .enumerate() + .map(|(i, d)| { + // deterministic relevance: shared-word count with the query + let doc = d.as_str().unwrap_or_default().to_lowercase(); + let hits = query.split_whitespace().filter(|w| doc.contains(w)).count(); + (i, hits as f64) + }) + .collect(); + scored.sort_by(|a, b| b.1.total_cmp(&a.1).then(a.0.cmp(&b.0))); + let top_n = body["top_n"].as_u64().unwrap_or(scored.len() as u64) as usize; + let results: Vec = scored + .into_iter() + .take(top_n) + .map(|(index, relevance_score)| json!({"index": index, "relevance_score": relevance_score})) + .collect(); + Self::ok_json(json!({ + "results": results, + "usage": {"total_tokens": (query.len() as i64 / 4).max(1)} + })) + } + fn audio_reply(&self, req: &UpstreamRequest) -> GResult { let body = Self::parse(&req.body, "audio")?; if req.url.ends_with("/audio/transcriptions") { Self::ok_json( json!({"text": "[mock-stt] transcribed audio", "language": body["language"]}), ) + } else if req.url.ends_with("/audio/translations") { + Self::ok_json(json!({"text": "[mock-stt] translated audio"})) } else if req.url.ends_with("/audio/speech") { let chars = body["input"].as_str().map(|s| s.len()).unwrap_or(0) as i64; Self::ok_json(json!({"audio_b64": MOCK_B64, "characters": chars})) @@ -626,6 +673,10 @@ impl Transport for MockTransport { self.audio_reply(&req) } else if u.contains("/videos") { self.video_reply(&req) + } else if u.contains("/moderations") { + self.moderations_reply(&req) + } else if u.contains("/rerank") { + self.rerank_reply(&req) } else if u.contains("/search") { self.search_reply(&req) } else if u.contains("/responses") { diff --git a/crates/engines/tests/request_construction.rs b/crates/engines/tests/request_construction.rs index 3672dce..e09fa0d 100644 --- a/crates/engines/tests/request_construction.rs +++ b/crates/engines/tests/request_construction.rs @@ -690,6 +690,7 @@ async fn stt_request_shape() { TypedParams::AudioStt(SttParams { audio_b64: "TU9DSw==".into(), language: Some("en".into()), + translate: false, }), ); let _ = AudioEngine::new(req, t.clone(), AudioKind::Stt) diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 49801ca..e130482 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -17,6 +17,7 @@ use gw_dag::DagContext; use gw_engines::http_transport::UpstreamPolicy; use gw_engines::{EngineOutcome, SharedTransport}; use gw_models::{Block, GResult, GatewayError, GatewayRequest, GatewayResponse}; +use gw_state::admission; use gw_state::{AkInfo, GatewayState, SharedConfig}; pub use gw_models::BatchItem; @@ -128,14 +129,67 @@ impl OnlineHandler { let inbound = (sec.moderate || retention.is_some()).then(|| plugins::inbound_text(&mut ctx.request)); - if sec.moderate - && let Some(block) = self - .moderate(&ctx, sec, inbound.as_deref().unwrap_or_default()) + if sec.moderate { + match self + .moderation(sec, inbound.as_deref().unwrap_or_default()) .await - { - ctx.decide("moderation", "denied"); - ctx.outcome = Some(content_filter_outcome(block)); - return Ok(ctx); + { + Moderation::Allow => {} + Moderation::Mask(spans) => { + let masked = plugins::apply_mask_spans(&mut ctx.request, &spans); + if masked > 0 { + ctx.decide("moderation", format!("masked {masked} span(s)")); + emit_security_event(&ctx, "moderation", "mask", masked as i64).await; + } + } + Moderation::Degrade => { + let tenant = &ctx.ak.tenant; + let swapped = ctx + .request + .model_param_v2 + .as_mut() + .map(|p| admission::swap_to_fallback(&snap.cfg, tenant, p)); + match swapped { + Some(admission::FallbackSwap::Swapped(from, fb)) => { + ctx.decide("moderation", format!("degraded {from} -> {fb}")); + emit_security_event(&ctx, "moderation", "degrade", 1).await; + } + Some(admission::FallbackSwap::AlreadyServing) => { + ctx.decide("moderation", "already serving the fallback"); + emit_security_event(&ctx, "moderation", "degrade", 1).await; + } + _ => { + emit_security_event(&ctx, "moderation", "block", 1).await; + deny_moderation( + &mut ctx, + "degrade without a fallback: denied", + "content requires degraded serving; no fallback model configured", + gw_consts::ErrCode::EMPTY_RESP.value() as i32, + ); + return Ok(ctx); + } + } + } + Moderation::Deny(reason) => { + emit_security_event(&ctx, "moderation", "block", 1).await; + deny_moderation( + &mut ctx, + "denied", + reason, + gw_consts::ErrCode::EMPTY_RESP.value() as i32, + ); + return Ok(ctx); + } + Moderation::Unavailable => { + deny_moderation( + &mut ctx, + "moderator unavailable: denied", + MODERATION_UNAVAILABLE, + gw_consts::ErrCode::SYSTEM_ERROR.value() as i32, + ); + return Ok(ctx); + } + } } let redacted = plugins::dlp_redact_request(sec, &mut ctx.request); @@ -162,6 +216,9 @@ impl OnlineHandler { ctx.quota_at, ) .await; + if e.code == gw_consts::ErrCode::STOP_LIMIT_MSG { + note_abuse(&ctx).await; + } return Err(e); } @@ -214,35 +271,17 @@ impl OnlineHandler { /// Run the wired moderator over raw text; `Some(reason)` to deny, `None` to /// allow. The seam the realtime surface uses (it has no `DagContext`); the /// caller records the security event on its own surface. - pub async fn moderate_text(&self, sec: &gw_config::SecurityConf, text: &str) -> Option { - match self.moderation(sec, text).await { - Moderation::Allow => None, - Moderation::Deny(reason) => Some(reason), - Moderation::Unavailable => Some(MODERATION_UNAVAILABLE.to_owned()), - } - } - - /// Run the wired moderator over the request's pre-DLP inbound `text`; - /// `Some(Block)` to deny. Records a security event on a moderator deny. - async fn moderate( - &self, - ctx: &DagContext, - sec: &gw_config::SecurityConf, - text: &str, - ) -> Option { + /// [`Self::moderation`] narrowed to the realtime surface: a live session + /// can't switch models, so `Degrade` denies there. + pub async fn moderate_rt(&self, sec: &gw_config::SecurityConf, text: &str) -> RtModeration { match self.moderation(sec, text).await { - Moderation::Allow => None, - Moderation::Deny(reason) => { - emit_security_event(ctx, "moderation", "block", 1).await; - Some(Block::blocked( - reason, - gw_consts::ErrCode::EMPTY_RESP.value() as i32, - )) - } - Moderation::Unavailable => Some(Block::blocked( - MODERATION_UNAVAILABLE, - gw_consts::ErrCode::SYSTEM_ERROR.value() as i32, - )), + Moderation::Allow => RtModeration::Allow, + Moderation::Mask(spans) => RtModeration::Mask(spans), + Moderation::Degrade => RtModeration::Deny( + "content requires degraded serving; not available on a live session".to_owned(), + ), + Moderation::Deny(reason) => RtModeration::Deny(reason), + Moderation::Unavailable => RtModeration::Deny(MODERATION_UNAVAILABLE.to_owned()), } } @@ -251,6 +290,8 @@ impl OnlineHandler { async fn moderation(&self, sec: &gw_config::SecurityConf, text: &str) -> Moderation { match self.moderator.review(text).await { Ok(moderation::Verdict::Allow) => Moderation::Allow, + Ok(moderation::Verdict::Mask(spans)) => Moderation::Mask(spans), + Ok(moderation::Verdict::Degrade) => Moderation::Degrade, Ok(moderation::Verdict::Deny(reason)) => Moderation::Deny(reason), Err(e) => { tracing::warn!(error = %e, fail_open = sec.moderation_fail_open, "moderator error"); @@ -292,10 +333,19 @@ impl OnlineHandler { /// fail-closed posture (fail-open resolves to `Allow`). enum Moderation { Allow, + Mask(Vec>), + Degrade, Deny(String), Unavailable, } +/// The realtime-surface subset of [`Moderation`]: no degrade mid-session. +pub enum RtModeration { + Allow, + Mask(Vec>), + Deny(String), +} + /// A per-request correlation id: `req--`, time-sortable and /// unique within the process (the seq disambiguates same-millisecond requests). pub fn new_request_id() -> String { @@ -306,6 +356,69 @@ pub fn new_request_id() -> String { format!("req-{ms}-{}", REQ_SEQ.fetch_add(1, Ordering::Relaxed)) } +/// Count one admission rejection against the key's daily abuse counter and +/// suspend it when a configured tier trips. Deny-path only, so the serving +/// path never pays for it; the day-scoped quota counter is fleet-shared. +async fn note_abuse(ctx: &DagContext) { + let tiers = &ctx.cfg.abuse.tiers; + if tiers.is_empty() { + return; + } + let now = gw_state::epoch_secs(); + // fresh read: a concurrent rejection may have suspended the key already + match ctx.state.auth.authenticate(&ctx.ak.ak).await { + Some(fresh) if fresh.status_at(now) != gw_state::KeyStatus::Suspended => {} + _ => return, + } + // ':' is banned in ak names, so the prefix cannot collide with a real key + let counter = format!("abuse:{}", ctx.ak.ak); + ctx.state.governance.quota_consume(&counter, 1).await; + let rejects = ctx.state.governance.quota_used(&counter).await; + let Some(tier) = tiers + .iter() + .filter(|t| rejects >= t.rejects) + .max_by_key(|t| t.rejects) + else { + return; + }; + let until = now + (tier.suspend_hours * 3600) as i64; + let patch = gw_state::KeyPatch { + suspended_until_epoch_secs: Some(Some(until)), + ..Default::default() + }; + if let Err(e) = ctx.state.auth.patch(&ctx.ak.ak, &patch).await { + tracing::warn!(error = %e, ak = %ctx.ak.ak, "abuse suspension patch failed"); + return; + } + let summary = format!( + "{rejects} admission rejects today; suspended {}h", + tier.suspend_hours + ); + ctx.state + .store + .admin_audit_add(&gw_state::AdminAudit { + created_at_epoch_secs: now, + actor: "system".to_owned(), + scope: "global".to_owned(), + action: "abuse_suspend".to_owned(), + target: ctx.ak.ak.clone(), + summary: summary.clone(), + source_ip: String::new(), + }) + .await + .unwrap_or_else(|e| tracing::warn!(error = %e, "abuse audit write failed")); + ctx.state + .alerts + .emit("abuse_suspend", ctx.ak.ak.clone(), summary); +} + +/// Record a moderation denial: decision line + content-filter outcome. Event +/// emission stays at the call sites (`Unavailable` deliberately records none). +fn deny_moderation(ctx: &mut DagContext, decision: &str, reason: impl Into, code: i32) { + ctx.decide("moderation", decision.to_owned()); + ctx.outcome = Some(content_filter_outcome(Block::blocked(reason, code))); +} + /// The 200-with-`content_filter` outcome every pre-stage denial returns. fn content_filter_outcome(block: Block) -> EngineOutcome { EngineOutcome { @@ -608,8 +721,33 @@ mod tests { } } + #[derive(Debug)] + struct MaskModerator; + + #[async_trait::async_trait] + impl moderation::Moderator for MaskModerator { + async fn review(&self, text: &str) -> Result { + match text.find("secret") { + Some(i) => Ok(moderation::Verdict::Mask( + std::iter::once(i..i + "secret".len()).collect(), + )), + None => Ok(moderation::Verdict::Allow), + } + } + } + + #[derive(Debug)] + struct DegradeModerator; + + #[async_trait::async_trait] + impl moderation::Moderator for DegradeModerator { + async fn review(&self, _text: &str) -> Result { + Ok(moderation::Verdict::Degrade) + } + } + #[tokio::test] - async fn moderate_text_allow_deny_and_failure_posture() { + async fn moderate_rt_maps_verdicts_and_failure_posture() { let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); let state = Arc::new(GatewayState::from_config(&cfg)); let base = OnlineHandler::new( @@ -617,23 +755,227 @@ mod tests { Arc::new(gw_engines::MockTransport), ); let mut sec = gw_config::SecurityConf::default(); - assert_eq!(base.moderate_text(&sec, "x").await, None, "default allows"); + assert!(matches!( + base.moderate_rt(&sec, "x").await, + RtModeration::Allow + )); let deny = base.clone().with_moderator(Arc::new(DenyModerator)); - assert_eq!(deny.moderate_text(&sec, "x").await.as_deref(), Some("nope")); + assert!(matches!(deny.moderate_rt(&sec, "x").await, RtModeration::Deny(r) if r == "nope")); + let mask = base.clone().with_moderator(Arc::new(MaskModerator)); + assert!(matches!( + mask.moderate_rt(&sec, "a secret").await, + RtModeration::Mask(spans) if spans.len() == 1 && spans[0] == (2..8) + )); + let degrade = base.clone().with_moderator(Arc::new(DegradeModerator)); + assert!( + matches!(degrade.moderate_rt(&sec, "x").await, RtModeration::Deny(r) if r.contains("live session")), + "degrade maps to deny on the realtime surface" + ); let err = base.with_moderator(Arc::new(ErrModerator)); sec.moderation_fail_open = true; + assert!(matches!( + err.moderate_rt(&sec, "x").await, + RtModeration::Allow + )); + sec.moderation_fail_open = false; + assert!(matches!( + err.moderate_rt(&sec, "x").await, + RtModeration::Deny(_) + )); + } + + #[tokio::test] + async fn moderation_mask_redacts_before_the_engine() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(MaskModerator)); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let ctx = h + .run(chat_req("gpt-4o", "tell secret now"), key) + .await + .unwrap(); + assert!( + ctx.outcome + .expect("outcome") + .response + .message + .contains("you said: tell [MASKED] now"), + "the engine saw the masked text" + ); + let events = h.state().store.security_events(None, 10).await.unwrap(); + assert!( + events + .iter() + .any(|e| e.rule == "moderation" && e.action == "mask") + ); + } + + #[tokio::test] + async fn moderation_degrade_swaps_to_tenant_fallback() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true}\nmodels: [{name: pub-m, protocol: openai-chat}, {name: fb-m, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, models: [pub-m, fb-m], fallback_model: fb-m}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DegradeModerator)); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let ctx = h.run(chat_req("pub-m", "hi"), key).await.unwrap(); assert_eq!( - err.moderate_text(&sec, "x").await, - None, - "error + fail-open allows" + ctx.outcome.expect("outcome").response.model, + "pub-m", + "response echoes the requested name" + ); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(ledger[0].model, "pub-m"); + assert_eq!(ledger[0].served_model, "fb-m"); + let events = h.state().store.security_events(None, 10).await.unwrap(); + assert!( + events + .iter() + .any(|e| e.rule == "moderation" && e.action == "degrade") + ); + } + + #[tokio::test] + async fn moderation_degrade_serves_when_already_on_the_fallback() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true}\nmodels: [{name: pub-m, protocol: openai-chat}, {name: fb-m, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, models: [pub-m, fb-m], fallback_model: fb-m}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DegradeModerator)); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let ctx = h.run(chat_req("fb-m", "hi"), key).await.unwrap(); + let out = ctx.outcome.expect("outcome"); + assert_eq!( + out.response.finish_reason, "stop", + "a request already on the fallback serves, not denies: {}", + out.response.message ); - sec.moderation_fail_open = false; assert!( - err.moderate_text(&sec, "x").await.is_some(), - "error + fail-closed denies" + ctx.decisions + .iter() + .any(|(n, w)| *n == "moderation" && w.contains("already serving")), + "{:?}", + ctx.decisions ); } + #[tokio::test] + async fn moderation_degrade_without_fallback_denies() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true}\nmodels: [{name: pub-m, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(DegradeModerator)); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let ctx = h.run(chat_req("pub-m", "hi"), key).await.unwrap(); + let out = ctx.outcome.expect("outcome"); + assert_eq!(out.response.finish_reason, "content_filter"); + assert!( + out.response + .message + .contains("no fallback model configured") + ); + } + + #[tokio::test] + async fn abuse_tiers_suspend_after_repeated_rejections() { + // daily_token_quota 1: every request rejects at reserve time (429) + let yaml = "listen: {host: h, port: 1}\nabuse: {tiers: [{rejects: 2, suspend_hours: 2}]}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 1}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let mut alerts = h.state().alerts.take_receiver().expect("receiver"); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let reject = |k: AkInfo| h.run(chat_req("gpt-4o", "hi"), k); + assert!( + reject(key.clone()).await.is_ok(), + "first request admits and burns the 1-token quota" + ); + assert_eq!( + reject(key.clone()).await.err().map(|e| e.http_status), + Some(429) + ); + let now = gw_state::epoch_secs(); + let fresh = h.state().auth.authenticate("k1").await.unwrap(); + assert_eq!( + fresh.status_at(now), + gw_state::KeyStatus::Active, + "one reject is under the tier" + ); + assert_eq!(reject(key).await.err().map(|e| e.http_status), Some(429)); + let fresh = h.state().auth.authenticate("k1").await.unwrap(); + assert_eq!( + fresh.status_at(now), + gw_state::KeyStatus::Suspended, + "second reject trips the 2-reject tier" + ); + let until = fresh.suspended_until_epoch_secs.expect("deadline set"); + assert!( + (until - now - 2 * 3600).abs() < 60, + "2h suspension: {until}" + ); + let trail = h.state().store.admin_audit_list(10).await.unwrap(); + assert!( + trail + .iter() + .any(|e| e.action == "abuse_suspend" && e.actor == "system" && e.target == "k1"), + "audit row: {trail:?}" + ); + let ev = alerts.recv().await.expect("alert emitted"); + assert_eq!((ev.kind, ev.subject.as_str()), ("abuse_suspend", "k1")); + } + + #[tokio::test] + async fn quota_fallback_skips_variant_select() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: pub-m, protocol: openai-chat, variants: [{model: canary-m, weight: 1}]}, {name: canary-m, protocol: openai-chat}, {name: fb-m, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, models: [pub-m, canary-m, fb-m], fallback_model: fb-m, model_quotas: {pub-m: 1}}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let h = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ); + let key = h.state().auth.authenticate("k1").await.unwrap(); + h.run(chat_req("pub-m", "burn the tiny quota"), key.clone()) + .await + .unwrap(); + let ctx = h + .run(chat_req("pub-m", "over quota now"), key) + .await + .unwrap(); + assert!( + ctx.decisions + .iter() + .any(|(n, w)| *n == "model_quota" && w.contains("serving fb-m")), + "quota gate degraded the request: {:?}", + ctx.decisions + ); + assert!( + !ctx.decisions.iter().any(|(n, _)| *n == "variant_select"), + "an already-degraded request skips the variant split" + ); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + let rec = ledger.last().expect("two billed requests"); + assert_eq!(rec.model, "pub-m"); + assert_eq!(rec.served_model, "fb-m"); + } + #[tokio::test] async fn per_user_budget_denies_over_the_cap() { let yaml = "listen: {host: h, port: 1}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, user_daily_token_quota: 5}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 100, daily_token_quota: 100000}]"; @@ -795,6 +1137,14 @@ mod tests { "estimated tokens, not zero: {:?}", ledger[0] ); + let avail = &h.state().avail; + avail.flush().await; + let minute = gw_state::epoch_secs() / 60; + assert_eq!( + avail.window("gpt-4o", minute - 5, minute).await, + (0, 0), + "an aborted stream is neither an availability success nor an error" + ); } #[derive(Debug)] @@ -1060,6 +1410,7 @@ mod tests { tokens_per_minute: None, expires_at_epoch_secs: None, banned: false, + suspended_until_epoch_secs: None, model_quotas: Default::default(), }; h.state() diff --git a/crates/handler/src/moderation.rs b/crates/handler/src/moderation.rs index d3edd3e..89bb252 100644 --- a/crates/handler/src/moderation.rs +++ b/crates/handler/src/moderation.rs @@ -4,12 +4,20 @@ //! turns `moderate` on. External review is latency-bearing and pull-driven, so //! this ships as a trait, not a built-in integration. +use std::ops::Range; use std::sync::Arc; /// A moderator's decision on one request's text. #[derive(Debug, Clone, PartialEq, Eq)] pub enum Verdict { Allow, + /// Redact these byte ranges of the reviewed text, then serve. Offsets + /// address the exact string `review` received; the caller maps them back + /// onto the request's text slots. + Mask(Vec>), + /// Serve via the tenant's fallback model. Denies when no fallback is + /// configured or the surface can't switch models mid-session (realtime). + Degrade, /// Deny with a user-facing reason. Deny(String), } diff --git a/crates/handler/src/plugins.rs b/crates/handler/src/plugins.rs index 7f59e2e..d0b21c3 100644 --- a/crates/handler/src/plugins.rs +++ b/crates/handler/src/plugins.rs @@ -126,6 +126,100 @@ fn push_text(out: &mut String, s: &str) { } } +/// Apply moderator mask spans (byte ranges into the reviewed text) back onto +/// the request's text slots. Returns the number of replaced ranges. +pub fn apply_mask_spans(request: &mut GatewayRequest, spans: &[std::ops::Range]) -> usize { + let mut masker = SpanMasker::new(spans); + for msg in &mut request.message { + for_each_message_text(msg, &mut |s| masker.apply(s)); + } + if let Some(param) = request.model_param_v2.as_mut() { + for_each_param_text(param, &mut |s| masker.apply(s)); + } + masker.hits +} + +/// [`apply_mask_spans`] for one realtime frame (spans address the frame's +/// collected text). +pub fn apply_mask_spans_frame( + frame: &mut serde_json::Value, + spans: &[std::ops::Range], +) -> usize { + let mut masker = SpanMasker::new(spans); + gw_engines::realtime::visit_frame_text(frame, &mut |s| masker.apply(s)) +} + +/// Replays the `push_text` walk (empty slots skipped, '\n' joiners) to map +/// global byte spans back onto individual slots. Spans are merged up front and +/// snapped to char boundaries so a sloppy external service can't split UTF-8. +struct SpanMasker { + spans: Vec>, + base: usize, + seen_any: bool, + hits: usize, +} + +impl SpanMasker { + fn new(spans: &[std::ops::Range]) -> Self { + let mut spans: Vec<_> = spans.iter().filter(|r| r.start < r.end).cloned().collect(); + spans.sort_by_key(|r| r.start); + spans.dedup_by(|next, prev| { + if next.start <= prev.end { + prev.end = prev.end.max(next.end); + true + } else { + false + } + }); + Self { + spans, + base: 0, + seen_any: false, + hits: 0, + } + } + + fn apply(&mut self, s: &mut String) -> usize { + if s.is_empty() { + return 0; + } + if self.seen_any { + self.base += 1; + } + self.seen_any = true; + let (start, end) = (self.base, self.base + s.len()); + self.base = end; + let mut replaced = 0; + for span in self.spans.iter().rev() { + if span.start >= end || span.end <= start { + continue; + } + let lo = snap_floor(s, span.start.saturating_sub(start)); + let hi = snap_ceil(s, span.end.saturating_sub(start).min(s.len())); + s.replace_range(lo..hi, "[MASKED]"); + replaced += 1; + } + self.hits += replaced; + replaced + } +} + +fn snap_floor(s: &str, mut i: usize) -> usize { + i = i.min(s.len()); + while i > 0 && !s.is_char_boundary(i) { + i -= 1; + } + i +} + +fn snap_ceil(s: &str, mut i: usize) -> usize { + i = i.min(s.len()); + while i < s.len() && !s.is_char_boundary(i) { + i += 1; + } + i +} + /// Case-insensitive blocklist test; terms are pre-lowercased at config load. /// ASCII text matches without allocating; non-ASCII falls back to a lowercase copy. fn blocklist_hit(sec: &SecurityConf, text: &str) -> bool { @@ -219,6 +313,8 @@ fn for_each_typed_text( T::Image(p) => f(&mut p.prompt), T::Video(p) => f(&mut p.prompt), T::Search(p) => f(&mut p.query), + T::Moderation(p) => p.input.iter_mut().map(&mut *f).sum(), + T::Rerank(p) => f(&mut p.query) + p.documents.iter_mut().map(&mut *f).sum::(), T::AudioStt(_) => 0, } } @@ -508,6 +604,36 @@ mod tests { use super::*; use gw_models::ChatMsg; + #[test] + fn mask_spans_map_across_slots_like_inbound_text() { + let mut req = GatewayRequest { + message: vec![ + ChatMsg::text("user", "first part"), + ChatMsg::text("user", "the secret word"), + ], + ..Default::default() + }; + let joined = inbound_text(&mut req.clone()); + let i = joined.find("secret").unwrap(); + let span = i..i + "secret".len(); + let hits = apply_mask_spans(&mut req, std::slice::from_ref(&span)); + assert_eq!(hits, 1); + assert_eq!(req.message[0].content, "first part", "untouched slot"); + assert_eq!(req.message[1].content, "the [MASKED] word"); + } + + #[test] + fn mask_spans_snap_to_char_boundaries_and_merge() { + let mut req = GatewayRequest { + message: vec![ChatMsg::text("user", "码字abc")], + ..Default::default() + }; + // splits the second CJK char and overlaps a second range — must not panic + let hits = apply_mask_spans(&mut req, &[2..4, 4..7]); + assert_eq!(hits, 1, "overlapping ranges merge to one replacement"); + assert!(req.message[0].content.contains("[MASKED]")); + } + fn sec() -> SecurityConf { SecurityConf { blocklist: vec!["forbiddenword".into()], diff --git a/crates/models/src/lib.rs b/crates/models/src/lib.rs index eee0e0e..8c52fe7 100644 --- a/crates/models/src/lib.rs +++ b/crates/models/src/lib.rs @@ -18,8 +18,8 @@ pub use cost::{ }; pub use error::{GResult, GatewayError}; pub use params::{ - ChatParams, EmbeddingParams, ImageParams, SearchParams, SttParams, TtsParams, TypedParams, - VideoParams, + ChatParams, EmbeddingParams, ImageParams, ModerationParams, RerankParams, SearchParams, + SttParams, TtsParams, TypedParams, VideoParams, }; pub use request::domain::{Account, ChatMsg}; pub use request::{BatchItem, GatewayRequest, ModelParamV2}; diff --git a/crates/models/src/params.rs b/crates/models/src/params.rs index 4b98b0f..e908b48 100644 --- a/crates/models/src/params.rs +++ b/crates/models/src/params.rs @@ -23,6 +23,10 @@ pub enum TypedParams { Video(VideoParams), /// bing / brave / serp / google custom search Search(SearchParams), + /// content moderation (openai moderations shape) + Moderation(ModerationParams), + /// document rerank (cohere / jina / voyage compatible) + Rerank(RerankParams), } #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -94,6 +98,9 @@ pub struct SttParams { pub audio_b64: String, #[serde(skip_serializing_if = "Option::is_none")] pub language: Option, + /// `/v1/audio/translations` (translate-to-English) instead of transcription. + #[serde(default)] + pub translate: bool, } #[derive(Debug, Clone, Default, Serialize, Deserialize)] @@ -112,6 +119,19 @@ pub struct SearchParams { pub count: i64, } +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct ModerationParams { + pub input: Vec, +} + +#[derive(Debug, Clone, Default, Serialize, Deserialize)] +pub struct RerankParams { + pub query: String, + pub documents: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_n: Option, +} + fn one() -> i64 { 1 } diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 4073d29..eafc984 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -97,6 +97,8 @@ async fn main() -> anyhow::Result<()> { let transport = select_transport()?; let postgres_url = cfg.storage.postgres_url.clone(); let shared = gw_state::SharedConfig::new(cfg, state); + let alert_task = gw_task::spawn_alert_dispatch(shared.clone()); + let avail_alert_task = gw_task::spawn_avail_alerts(shared.clone(), gw_task::AVAIL_ALERT_PERIOD); let loader: gw_views::ConfigLoader = match &config_store { Some(store) => { let store = store.clone(); @@ -206,6 +208,8 @@ async fn main() -> anyhow::Result<()> { purge_task.abort(); rollup_task.abort(); avail_task.abort(); + alert_task.abort(); + avail_alert_task.abort(); tracing::info!("gw drained and exiting"); Ok(()) } diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index ce7c864..5b32c2d 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -1840,6 +1840,92 @@ async fn realtime_entitlement_blocks_unentitled_tenant() { assert!(tokio_tungstenite::connect_async(ok).await.is_ok()); } +#[tokio::test] +async fn realtime_variant_session_serves_canary_bills_public_name() { + use futures::{SinkExt, StreamExt}; + use tokio_tungstenite::tungstenite::Message; + use tokio_tungstenite::tungstenite::client::IntoClientRequest; + + let yaml = r#" +listen: {host: 127.0.0.1, port: 0} +access_keys: + - {ak: ak-rt, product: rt, qps: 100, daily_token_quota: 1000000} +accounts: + - {name: rt-acc, provider: openai, protocols: ["realtime"]} +models: + - {name: rt-pub, protocol: realtime, variants: [{model: rt-canary, weight: 1}]} + - {name: rt-canary, protocol: realtime} +"#; + let cfg = Arc::new(gw_config::GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(gw_state::GatewayState::from_config(&cfg)); + let application = gw_views::app(gw_views::AppState::new( + cfg, + state.clone(), + Arc::new(gw_engines::MockTransport), + )); + let addr = serve_app(application).await; + + let mut req = format!("ws://{addr}/v1/realtime?model=rt-pub") + .into_client_request() + .unwrap(); + req.headers_mut() + .insert("authorization", "Bearer ak-rt".parse().unwrap()); + req.headers_mut() + .insert("x-gw-user", "alice".parse().unwrap()); + let (mut ws, _) = tokio_tungstenite::connect_async(req) + .await + .expect("ws connect"); + + let first = ws.next().await.unwrap().unwrap(); + let v: Value = serde_json::from_str(first.to_text().unwrap()).unwrap(); + assert_eq!( + v["session"]["model"], "rt-pub", + "clients see the public name" + ); + + ws.send(Message::text( + serde_json::json!({"type":"input_text","text":"canary hello"}).to_string(), + )) + .await + .unwrap(); + let mut assembled = String::new(); + while let Some(Ok(msg)) = ws.next().await { + let v: Value = serde_json::from_str(msg.to_text().unwrap()).unwrap(); + match v["type"].as_str().unwrap() { + "response.delta" => assembled.push_str(v["delta"].as_str().unwrap()), + "response.done" => break, + other => panic!("unexpected event {other}"), + } + } + assert!( + assembled.contains("[mock-realtime:rt-canary]"), + "the variant serves the session: {assembled}" + ); + + let rec = 'ledger: { + for _ in 0..100 { + let (_, ledger) = state.store.ledger_snapshot(usize::MAX).await.unwrap(); + if let Some(rec) = ledger.into_iter().next() { + break 'ledger rec; + } + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + } + panic!("turn was never billed"); + }; + assert_eq!( + rec.model, "rt-pub", + "quota/billing rows carry the public name" + ); + assert_eq!(rec.served_model, "rt-canary"); + state.avail.flush().await; + let minute = gw_state::epoch_secs() / 60; + assert_eq!( + state.avail.window("rt-pub", minute - 5, minute).await, + (1, 0), + "the turn samples availability under the public name" + ); +} + #[tokio::test] async fn realtime_refuses_ungovernable_provider() { use tokio_tungstenite::tungstenite::client::IntoClientRequest; diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index de249a8..3ac9b9b 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -85,6 +85,38 @@ pub fn model_quota_limit(cfg: &GatewayConfig, ak: &AkInfo, model: &str) -> Optio }) } +/// Outcome of a tenant-fallback swap attempt. `AlreadyServing` is distinct +/// from `Unconfigured`: a request already on the fallback has nowhere further +/// to degrade but IS degraded — callers must not deny it as unconfigured. +pub enum FallbackSwap { + /// Swapped; carries `(requested, fallback)` for the decision trail. + Swapped(String, String), + AlreadyServing, + Unconfigured, +} + +/// Swap `param` to the tenant's fallback model, threading `fallback_from` so +/// billing/echo semantics follow. The one fallback swap the quota gate and +/// moderation degrade share. +pub fn swap_to_fallback( + cfg: &GatewayConfig, + tenant: &str, + param: &mut gw_models::ModelParamV2, +) -> FallbackSwap { + let Some(fb) = cfg + .find_tenant(tenant) + .and_then(|t| t.fallback_model.clone()) + else { + return FallbackSwap::Unconfigured; + }; + if fb == param.model_name { + return FallbackSwap::AlreadyServing; + } + let from = std::mem::replace(&mut param.model_name, fb.clone()); + param.fallback_from = Some(from.clone()); + FallbackSwap::Swapped(from, fb) +} + /// Pooled tenant QPS, when the tenant configures one. pub async fn check_tenant_rate( gov: &dyn Governance, diff --git a/crates/state/src/alerts.rs b/crates/state/src/alerts.rs new file mode 100644 index 0000000..422b5b6 --- /dev/null +++ b/crates/state/src/alerts.rs @@ -0,0 +1,82 @@ +//! Advisory alert bus: emit sites push fire-and-forget events, one dispatcher +//! (spawned by the server) drains them to the configured webhook. Bounded and +//! lossy by design — an alert must never block or slow the serving path. + +use tokio::sync::mpsc; + +/// Bounded: a stuck dispatcher drops alerts instead of growing memory. +const ALERT_QUEUE: usize = 256; + +/// One outbound alert. +#[derive(Debug)] +pub struct AlertEvent { + pub kind: &'static str, + pub subject: String, + pub detail: String, + pub at_epoch_secs: i64, +} + +/// The emit side is sync and never fails; the receiver is taken once by the +/// dispatch task (survives config reloads — the bus is a preserved seam). +pub struct AlertBus { + tx: mpsc::Sender, + rx: std::sync::Mutex>>, +} + +impl AlertBus { + pub fn emit(&self, kind: &'static str, subject: String, detail: String) { + let ev = AlertEvent { + kind, + subject, + detail, + at_epoch_secs: crate::epoch_secs(), + }; + if self.tx.try_send(ev).is_err() { + tracing::debug!(kind, "alert queue full or unclaimed; dropped"); + } + } + + /// The single consumer end; `None` after the first take. + pub fn take_receiver(&self) -> Option> { + self.rx + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .take() + } +} + +impl Default for AlertBus { + fn default() -> Self { + let (tx, rx) = mpsc::channel(ALERT_QUEUE); + Self { + tx, + rx: std::sync::Mutex::new(Some(rx)), + } + } +} + +impl std::fmt::Debug for AlertBus { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("AlertBus") + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn emits_are_lossy_and_receiver_single() { + let bus = AlertBus::default(); + let mut rx = bus.take_receiver().expect("first take"); + assert!(bus.take_receiver().is_none(), "single consumer"); + bus.emit("abuse_suspend", "k1".into(), "2 rejects".into()); + let ev = rx.recv().await.expect("delivered"); + assert_eq!((ev.kind, ev.subject.as_str()), ("abuse_suspend", "k1")); + for _ in 0..(ALERT_QUEUE + 10) { + bus.emit("account_cooldown", "a".into(), String::new()); + } + drop(rx); + bus.emit("account_cooldown", "a".into(), String::new()); + } +} diff --git a/crates/state/src/avail.rs b/crates/state/src/avail.rs index 74b678c..64fa0c6 100644 --- a/crates/state/src/avail.rs +++ b/crates/state/src/avail.rs @@ -23,6 +23,17 @@ pub enum AvailState { NoData, } +impl AvailState { + pub fn as_str(self) -> &'static str { + match self { + AvailState::Available => "available", + AvailState::Unstable => "unstable", + AvailState::Unavailable => "unavailable", + AvailState::NoData => "no_data", + } + } +} + /// Classify a window's counts against the configured error-rate thresholds. pub fn classify( ok: u64, diff --git a/crates/state/src/keystore.rs b/crates/state/src/keystore.rs index 95d2dce..46e6286 100644 --- a/crates/state/src/keystore.rs +++ b/crates/state/src/keystore.rs @@ -82,6 +82,12 @@ impl PostgresKeyStore { .execute(&pool) .await .map_err(|e| crate::sqlx_err("migrate access_keys owner", e))?; + sqlx::query( + "ALTER TABLE access_keys ADD COLUMN IF NOT EXISTS suspended_until_epoch_secs BIGINT", + ) + .execute(&pool) + .await + .map_err(|e| crate::sqlx_err("migrate access_keys suspension", e))?; Ok(Self { pool, cache: moka::sync::Cache::builder() @@ -101,7 +107,7 @@ impl PostgresKeyStore { async fn fetch(&self, ak: &str) -> Result, sqlx::Error> { let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys WHERE ak = $1", + expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs FROM access_keys WHERE ak = $1", ) .bind(ak) .fetch_optional(&self.pool) @@ -149,7 +155,7 @@ impl KeyStore for PostgresKeyStore { .map_err(|e| crate::sqlx_err("begin patch", e))?; let row = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys + expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs FROM access_keys WHERE ak = $1 FOR UPDATE", ) .bind(ak) @@ -161,7 +167,8 @@ impl KeyStore for PostgresKeyStore { info.apply_patch(patch); sqlx::query( "UPDATE access_keys SET qps = $2, daily_token_quota = $3, - tokens_per_minute = $4, expires_at_epoch_secs = $5, banned = $6 + tokens_per_minute = $4, expires_at_epoch_secs = $5, banned = $6, + suspended_until_epoch_secs = $7 WHERE ak = $1", ) .bind(ak) @@ -170,6 +177,7 @@ impl KeyStore for PostgresKeyStore { .bind(info.tokens_per_minute) .bind(info.expires_at_epoch_secs) .bind(info.banned) + .bind(info.suspended_until_epoch_secs) .execute(&mut *tx) .await .map_err(|e| crate::sqlx_err("apply patch", e))?; @@ -199,7 +207,7 @@ impl KeyStore for PostgresKeyStore { ) -> GResult> { let rows = sqlx::query( "SELECT ak, product, tenant, qps, daily_token_quota, tokens_per_minute, - expires_at_epoch_secs, banned, model_quotas, owner FROM access_keys + expires_at_epoch_secs, banned, model_quotas, owner, suspended_until_epoch_secs FROM access_keys WHERE ($1::text IS NULL OR tenant = $1) ORDER BY ak LIMIT $2 OFFSET $3", ) .bind(tenant) @@ -240,6 +248,8 @@ impl KeyStore for PostgresKeyStore { /// The one INSERT..ON CONFLICT for the key table. The sticky-source CASE is a /// no-op when `source` is Config, so the reload path shares it unchanged. +/// `suspended_until_epoch_secs` is deliberately absent: an upsert never clears +/// a runtime suspension — only a patch touches it. async fn upsert( exec: impl sqlx::PgExecutor<'_>, info: &AkInfo, @@ -290,6 +300,7 @@ fn row_to_info(row: &sqlx::postgres::PgRow) -> AkInfo { serde_json::from_str(row.get::<&str, _>(8)).unwrap_or_default(), ), owner: row.get(9), + suspended_until_epoch_secs: row.get(10), } } @@ -315,6 +326,7 @@ mod tests { tokens_per_minute: None, expires_at_epoch_secs: None, banned: false, + suspended_until_epoch_secs: None, model_quotas: Default::default(), } } diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 7455b37..15d5cb7 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -14,6 +14,7 @@ use gw_consts::Protocol; use gw_models::Account; pub mod admission; +pub mod alerts; pub mod avail; pub mod configstore; pub mod content; @@ -22,6 +23,7 @@ pub mod health; pub mod keystore; pub mod store; +pub use alerts::{AlertBus, AlertEvent}; pub use avail::{AvailState, AvailStore, classify}; pub use configstore::{CONFIG_CHANNEL, PostgresConfigStore}; pub use content::{ContentRecord, sealing_available}; @@ -52,17 +54,27 @@ pub struct AkInfo { pub expires_at_epoch_secs: Option, /// A banned key stays in the table but fails auth with a distinct 403. pub banned: bool, + /// Abuse auto-suspension deadline; the key self-recovers when it passes. + /// Runtime state — a config reload must not clear it. + pub suspended_until_epoch_secs: Option, /// Per-model daily token caps overriding the tenant defaults; empty = none. /// Arc'd: `AkInfo` is cloned per request and the map is write-rare. pub model_quotas: Arc>, } impl AkInfo { - /// Lifecycle state at `now` (unix seconds). Ban wins over expiry. + /// Lifecycle state at `now` (unix seconds). Ban wins over suspension, + /// suspension over expiry; an elapsed suspension self-recovers. pub fn status_at(&self, now_epoch_secs: i64) -> KeyStatus { if self.banned { return KeyStatus::Banned; } + if self + .suspended_until_epoch_secs + .is_some_and(|t| now_epoch_secs < t) + { + return KeyStatus::Suspended; + } match self.expires_at_epoch_secs { Some(t) if now_epoch_secs >= t => KeyStatus::Expired, _ => KeyStatus::Active, @@ -96,6 +108,9 @@ impl AkInfo { if let Some(v) = patch.banned { self.banned = v; } + if let Some(v) = patch.suspended_until_epoch_secs { + self.suspended_until_epoch_secs = v; + } } } @@ -111,6 +126,7 @@ impl From<&gw_config::AkConf> for AkInfo { tokens_per_minute: k.tokens_per_minute, expires_at_epoch_secs: k.expires_at_epoch_secs, banned: k.banned, + suspended_until_epoch_secs: None, model_quotas: Arc::new(k.model_quotas.clone()), } } @@ -125,6 +141,7 @@ pub struct KeyPatch { pub tokens_per_minute: Option>, pub expires_at_epoch_secs: Option>, pub banned: Option, + pub suspended_until_epoch_secs: Option>, } /// Key lifecycle state: expired and banned keys authenticate to distinct 403s. @@ -133,6 +150,7 @@ pub enum KeyStatus { Active, Banned, Expired, + Suspended, } /// Where a key came from: the config file (re-applied on reload) or the admin @@ -159,7 +177,7 @@ impl AkAuth { /// Insert or replace a key. Config ownership is sticky: an admin write to a /// config-declared key keeps it `Config` (so a reload can still revoke it), /// while a config write claims an admin key. Atomic via the entry lock. - pub fn put(&self, info: AkInfo, source: KeySource) { + pub fn put(&self, mut info: AkInfo, source: KeySource) { use dashmap::mapref::entry::Entry; match self.keys.entry(info.ak.clone()) { Entry::Occupied(mut e) => { @@ -169,6 +187,11 @@ impl AkAuth { } else { source }; + // an upsert never clears a runtime suspension — only a patch + // does (mirrors the Postgres upsert, which omits the column) + if info.suspended_until_epoch_secs.is_none() { + info.suspended_until_epoch_secs = e.get().0.suspended_until_epoch_secs; + } e.insert((info, source)); } Entry::Vacant(e) => { @@ -667,6 +690,8 @@ pub struct GatewayState { pub cache: Arc, /// Per-model availability counters; Redis for fleet-wide aggregation. pub avail: Arc, + /// Advisory alert bus; the server's dispatch task drains it. + pub alerts: Arc, } impl Default for GatewayState { @@ -679,6 +704,7 @@ impl Default for GatewayState { health: Arc::new(AccountHealth::default()), cache: Arc::new(MemoryResponseCache::default()), avail: Arc::new(avail::MemoryAvail::default()), + alerts: Arc::new(alerts::AlertBus::default()), } } } @@ -783,6 +809,7 @@ impl GatewayState { health: prev.health.clone(), cache: prev.cache.clone(), avail: prev.avail.clone(), + alerts: prev.alerts.clone(), }) } } @@ -918,10 +945,47 @@ mod tests { tokens_per_minute: None, expires_at_epoch_secs: None, banned: false, + suspended_until_epoch_secs: None, model_quotas: Default::default(), } } + #[test] + fn suspension_orders_below_ban_and_self_recovers() { + let mut ak = ak_info("k"); + ak.suspended_until_epoch_secs = Some(100); + assert_eq!(ak.status_at(50), KeyStatus::Suspended); + assert_eq!(ak.status_at(100), KeyStatus::Active, "deadline passed"); + ak.banned = true; + assert_eq!(ak.status_at(50), KeyStatus::Banned, "ban wins"); + ak.banned = false; + ak.expires_at_epoch_secs = Some(40); + assert_eq!( + ak.status_at(50), + KeyStatus::Suspended, + "suspension reported over expiry" + ); + } + + #[test] + fn config_reapply_keeps_runtime_suspension() { + let auth = AkAuth::default(); + auth.put(ak_info("k"), KeySource::Config); + auth.patch( + "k", + &KeyPatch { + suspended_until_epoch_secs: Some(Some(i64::MAX)), + ..Default::default() + }, + ); + auth.put(ak_info("k"), KeySource::Config); + assert_eq!( + auth.authenticate("k").unwrap().suspended_until_epoch_secs, + Some(i64::MAX), + "re-applied config key keeps the suspension" + ); + } + #[test] fn attributed_user_prefers_nonempty_owner_else_fallback() { let mut ak = ak_info("k"); diff --git a/crates/task/Cargo.toml b/crates/task/Cargo.toml index c9ca1c1..8c9d753 100644 --- a/crates/task/Cargo.toml +++ b/crates/task/Cargo.toml @@ -7,11 +7,16 @@ publish.workspace = true description = "Local background tasks: quota reset, content purge, usage rollup." [dependencies] +futures = { workspace = true } +gw-consts = { workspace = true } gw-state = { workspace = true } +reqwest = { workspace = true } +serde_json = { workspace = true } tokio = { workspace = true } tracing = { workspace = true } [dev-dependencies] +gw-config = { workspace = true } tokio = { workspace = true } [lints] diff --git a/crates/task/src/lib.rs b/crates/task/src/lib.rs index 13d8dc8..b167e3b 100644 --- a/crates/task/src/lib.rs +++ b/crates/task/src/lib.rs @@ -1,11 +1,13 @@ //! Local background tasks: periodic AK daily quota reset, retained-content -//! purge, and the usage rollup. Batch job execution lives in -//! gw-handler::offline (spawned on submit) and needs no separate poller. +//! purge, the usage rollup, availability flush/alerting, and the alert +//! webhook dispatcher. Batch job execution lives in gw-handler::offline +//! (spawned on submit) and needs no separate poller. +use std::collections::HashMap; use std::sync::Arc; -use std::time::Duration; +use std::time::{Duration, Instant}; -use gw_state::GatewayState; +use gw_state::{GatewayState, SharedConfig}; /// The production period: once a day. pub const DAILY: Duration = Duration::from_secs(24 * 60 * 60); @@ -15,6 +17,8 @@ pub const PURGE_PERIOD: Duration = Duration::from_secs(60 * 60); pub const ROLLUP_PERIOD: Duration = Duration::from_secs(60); /// Buffered per-model availability counts are flushed to the store this often. pub const AVAIL_FLUSH_PERIOD: Duration = Duration::from_secs(2); +/// Model availability is re-classified for alerting this often. +pub const AVAIL_ALERT_PERIOD: Duration = Duration::from_secs(60); /// Spawn the daily quota reset loop. Returns the join handle (abort to stop). /// `period` is configurable so tests don't wait 24h. @@ -97,10 +101,190 @@ pub fn spawn_avail_flush( }) } +/// Drain the alert bus and POST each event to the configured webhook, muting +/// repeats of the same (kind, subject) within the dedup window. Exits when the +/// bus's receiver was already taken or every sender is gone. +pub fn spawn_alert_dispatch(shared: SharedConfig) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let Some(mut rx) = shared.load().state.alerts.take_receiver() else { + return; + }; + let client = reqwest::Client::new(); + let mut sent: HashMap = HashMap::new(); + while let Some(ev) = rx.recv().await { + let conf = shared.load().cfg.alerts.clone(); + let Some(url) = conf.webhook_url() else { + continue; + }; + let dedup = Duration::from_secs(conf.dedup_seconds); + if !should_send(&mut sent, format!("{}:{}", ev.kind, ev.subject), dedup) { + continue; + } + let body = serde_json::json!({ + "kind": ev.kind, + "subject": ev.subject, + "detail": ev.detail, + "at_epoch_secs": ev.at_epoch_secs, + }); + let post = client + .post(&url) + .header("content-type", "application/json") + .body(body.to_string()) + .timeout(Duration::from_secs(10)) + .send() + .await; + match post { + Ok(resp) if resp.status().is_success() => {} + Ok(resp) => { + tracing::warn!(status = %resp.status(), kind = ev.kind, "alert webhook rejected") + } + Err(e) => tracing::warn!(error = %e, kind = ev.kind, "alert webhook unreachable"), + } + } + }) +} + +/// Spawn the availability alert sweep: re-classify every model each period +/// and emit an alert on a state transition. +pub fn spawn_avail_alerts(shared: SharedConfig, period: Duration) -> tokio::task::JoinHandle<()> { + tokio::spawn(async move { + let mut tick = tokio::time::interval(period); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Skip); + tick.tick().await; + let mut last = HashMap::new(); + loop { + tick.tick().await; + avail_alert_sweep(&shared, &mut last).await; + } + }) +} + +/// One sweep round, factored out so tests drive it directly. +async fn avail_alert_sweep( + shared: &SharedConfig, + last: &mut HashMap, +) { + let snap = shared.load(); + let st = &snap.cfg.stability; + let until = gw_state::epoch_secs() / 60; + let since = until - (st.availability_window_minutes - 1); + let counts = futures::future::join_all(snap.cfg.models.iter().map(|m| async { + ( + &m.name, + snap.state.avail.window(&m.name, since, until).await, + ) + })) + .await; + for (name, (ok, err)) in counts { + let cur = gw_state::classify( + ok, + err, + st.availability_min_samples, + st.unstable_error_rate, + st.unavailable_error_rate, + ); + if let Some(prev) = last.insert(name.clone(), cur) + && prev != cur + { + snap.state.alerts.emit( + "model_availability", + name.clone(), + format!("{} -> {}", prev.as_str(), cur.as_str()), + ); + } + } +} + +/// Whether `key` is outside its dedup window (records the send when so). +/// Bounded: past 1024 entries the elapsed ones are swept, so a long-lived +/// dispatcher with many distinct subjects can't grow without limit. +fn should_send(sent: &mut HashMap, key: String, window: Duration) -> bool { + let now = Instant::now(); + if sent.len() > 1024 { + sent.retain(|_, at| now.duration_since(*at) < window); + } + match sent.get(&key) { + Some(at) if now.duration_since(*at) < window => false, + _ => { + sent.insert(key, now); + true + } + } +} + #[cfg(test)] mod tests { use super::*; + #[test] + fn dedup_window_mutes_repeats() { + let mut sent = HashMap::new(); + let w = Duration::from_secs(60); + assert!(should_send(&mut sent, "a:x".into(), w)); + assert!(!should_send(&mut sent, "a:x".into(), w), "muted repeat"); + assert!(should_send(&mut sent, "a:y".into(), w), "distinct subject"); + assert!( + should_send(&mut sent, "a:x".into(), Duration::ZERO), + "elapsed window resends" + ); + } + + #[tokio::test] + async fn avail_sweep_alerts_on_transition_only() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: m, protocol: openai-chat}, {name: rt, protocol: realtime}]\nstability: {availability_min_samples: 2}"; + let cfg = std::sync::Arc::new(gw_config::GatewayConfig::from_yaml(yaml).unwrap()); + let state = std::sync::Arc::new(GatewayState::from_config(&cfg)); + let shared = SharedConfig::new(cfg, state.clone()); + let mut rx = state.alerts.take_receiver().expect("receiver"); + let mut last = HashMap::new(); + avail_alert_sweep(&shared, &mut last).await; + avail_alert_sweep(&shared, &mut last).await; + assert!(rx.try_recv().is_err(), "steady no_data emits nothing"); + for _ in 0..4 { + state.avail.record("m", false); + } + state.avail.flush().await; + avail_alert_sweep(&shared, &mut last).await; + let ev = rx.try_recv().expect("transition alert"); + assert_eq!((ev.kind, ev.subject.as_str()), ("model_availability", "m")); + assert_eq!(ev.detail, "no_data -> unavailable"); + avail_alert_sweep(&shared, &mut last).await; + assert!(rx.try_recv().is_err(), "no repeat while the state holds"); + } + + #[tokio::test] + async fn alert_dispatch_posts_to_the_webhook() { + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let served = tokio::spawn(async move { + let (mut sock, _) = listener.accept().await.unwrap(); + let mut buf = vec![0u8; 4096]; + let n = sock.read(&mut buf).await.unwrap(); + sock.write_all(b"HTTP/1.1 200 OK\r\ncontent-length: 0\r\n\r\n") + .await + .unwrap(); + String::from_utf8_lossy(&buf[..n]).into_owned() + }); + // SAFETY: unique var name for this test; no concurrent reader of it. + unsafe { + std::env::set_var("GW_TEST_ALERT_URL", format!("http://{addr}/hook")); + } + let yaml = "listen: {host: h, port: 1}\nalerts: {webhook_url_env: GW_TEST_ALERT_URL}"; + let cfg = std::sync::Arc::new(gw_config::GatewayConfig::from_yaml(yaml).unwrap()); + let state = std::sync::Arc::new(GatewayState::from_config(&cfg)); + let shared = SharedConfig::new(cfg, state.clone()); + let task = spawn_alert_dispatch(shared); + state + .alerts + .emit("abuse_suspend", "k9".into(), "3 rejects".into()); + let request = served.await.unwrap(); + task.abort(); + assert!(request.starts_with("POST /hook")); + assert!(request.contains(r#""kind":"abuse_suspend""#), "{request}"); + assert!(request.contains(r#""subject":"k9""#)); + } + #[tokio::test] async fn quota_reset_clears_counters() { let state = Arc::new(GatewayState::default()); diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index d4aef21..ccf5990 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -111,6 +111,9 @@ pub fn app(state: AppState) -> Router { .route("/v1/images/edits", post(images_edits)) .route("/v1/audio/speech", post(audio_speech)) .route("/v1/audio/transcriptions", post(audio_transcriptions)) + .route("/v1/audio/translations", post(audio_translations)) + .route("/v1/moderations", post(moderations)) + .route("/v1/rerank", post(rerank)) .route("/v1/batches", post(batches_submit)) .route("/v1/batches/{id}", get(batches_get)) .route("/v1/files", post(files_upload)) @@ -219,26 +222,52 @@ async fn realtime_ws( format!("model `{model}` is not entitled for tenant `{}`", ak.tenant), ); } + // client attribution hint captured at connect (no per-turn body user field) + let hint = user_header(&headers).unwrap_or_default(); + // variant split pins for the whole session, sticky by the attributed user + // (per-connection spread when anonymous) — the REST semantics, one level up + let served = match model_conf { + Some(conf) if !conf.variants.is_empty() => { + let sticky = ak.attributed_user(&hint); + let key = if sticky.is_empty() { + gw_handler::new_request_id() + } else { + sticky.to_owned() + }; + gw_config::pick_variant(&conf.variants, &key).model.clone() + } + _ => model.clone(), + }; + let served_conf = if served == model { + model_conf + } else { + snap.cfg.find_model(&served) + }; let account = snap .state .pool .select_healthy( mt, - model_conf.and_then(|m| m.provider.as_deref()), + served_conf.and_then(|m| m.provider.as_deref()), &[], snap.state.health.as_ref(), ) .await; let Some(account) = account else { + // an exhausted pool is a client-visible failure of the public name + snap.state.avail.record(&model, false); return error_response(503, format!("no healthy upstream account serves `{model}`")); }; + let m = RtModel { + requested: model, + served, + from_config: served_conf.is_some(), + }; // select "realtime" so subprotocol-offering clients get a valid handshake let ws = ws.protocols(["realtime"]); - // client attribution hint captured at connect (no per-turn body user field) - let hint = user_header(&headers).unwrap_or_default(); if account.endpoint.is_empty() { ws.on_upgrade(move |socket| { - realtime_session(socket, s, ak, model, mt, account.name.clone(), hint) + realtime_session(socket, s, ak, m, mt, account.name.clone(), hint) }) } else if gw_engines::realtime::is_gemini_realtime(&account.provider) { // no pre-generation gate signal in this dialect — refuse rather than bill after the fact @@ -250,10 +279,23 @@ async fn realtime_ws( ), ) } else { - ws.on_upgrade(move |socket| realtime_bridge(socket, s, ak, model, mt, account, hint)) + ws.on_upgrade(move |socket| realtime_bridge(socket, s, ak, m, mt, account, hint)) } } +/// A realtime session's model identity: the public name the client asked for +/// and the variant actually served (equal without a split). Entitlement and +/// the per-(AK, model) counter judge `requested`; capacity (QPM), pricing, +/// the upstream URL, and availability's success/error samples follow the REST +/// semantics — samples attribute to `requested`. +struct RtModel { + requested: String, + served: String, + /// Whether `served` came from config at handshake — only then can a + /// reload invalidate it (wire-direct sessions have no config row). + from_config: bool, +} + /// A turn admitted by [`realtime_gate`]: the freshly re-authenticated key, /// the reserves taken, the admission day (the paired settle/refund lands on /// the same bucket), and the admission snapshot (settlement must not drift @@ -294,7 +336,7 @@ impl RealtimeAdmit { async fn realtime_gate( s: &AppState, ak: &AkInfo, - model: &str, + m: &RtModel, hint: &str, ) -> Result { let snap = s.handler.config.load(); @@ -305,24 +347,34 @@ async fn realtime_gate( } _ => return Err(format!("access key {} is no longer valid", ak.ak)), }; - if !cfg.tenant_allows_model(&ak.tenant, model) { + if !cfg.tenant_allows_model(&ak.tenant, &m.requested) { + return Err(format!( + "model `{}` is not entitled for tenant `{}`", + m.requested, ak.tenant + )); + } + // the session pinned `served` at handshake; a reload may have removed it, + // and pricing an unknown model silently bills zero — deny so the client + // reconnects and picks against the live config. Wire-direct sessions + // (`?model=realtime`, never in config) are exempt: nothing to vanish. + if m.from_config && cfg.find_model(&m.served).is_none() { return Err(format!( - "model `{model}` is not entitled for tenant `{}`", - ak.tenant + "model `{}` is no longer configured; reconnect", + m.served )); } let gov = state.governance.as_ref(); admission::check_tenant_rate(gov, cfg, &ak.tenant).await?; admission::check_ak_rate(gov, &ak).await?; admission::check_product_qpm(gov, cfg, &ak.product).await?; - admission::check_model_qpm(gov, cfg, model).await?; + admission::check_model_qpm(gov, cfg, &m.served).await?; admission::check_user_budget(gov, cfg, &ak.tenant, ak.attributed_user(hint)).await?; - if let Some(limit) = admission::model_quota_limit(cfg, &ak, model) + if let Some(limit) = admission::model_quota_limit(cfg, &ak, &m.requested) && !gov - .quota_check(&admission::model_quota_key(&ak.ak, model), limit) + .quota_check(&admission::model_quota_key(&ak.ak, &m.requested), limit) .await { - return Err(format!("model quota exhausted for `{model}`")); + return Err(format!("model quota exhausted for `{}`", m.requested)); } let at = gw_state::epoch_secs(); admission::reserve_daily(gov, &ak, REALTIME_TURN_RESERVE, at).await?; @@ -350,7 +402,7 @@ async fn realtime_gate( /// frame (cancelled/empty turn) refunds the reserves and writes nothing. async fn bill_realtime_turn( admit: &RealtimeAdmit, - model: &str, + m: &RtModel, mt: gw_consts::Protocol, account: &str, it: i64, @@ -365,11 +417,11 @@ async fn bill_realtime_turn( } let (cfg, state) = (&admit.snap.cfg, &admit.snap.state); // pipeline parity: the same shared weighting the DAG estimate paths use - let rate = gw_state::model_token_rate(cfg, model); + let rate = gw_state::model_token_rate(cfg, &m.served); let (bp, bc) = gw_models::weighted_pair(it, ot, &rate); let total = gw_state::clamp_tokens(bp.saturating_add(bc)); - let model_quota_key = admission::model_quota_limit(cfg, ak, model) - .map(|_| admission::model_quota_key(&ak.ak, model)); + let model_quota_key = admission::model_quota_limit(cfg, ak, &m.requested) + .map(|_| admission::model_quota_key(&ak.ak, &m.requested)); admission::settle_and_bill( state.governance.as_ref(), state.store.as_ref(), @@ -381,8 +433,8 @@ async fn bill_realtime_turn( tenant: &ak.tenant, user_id: admit.user.as_str(), request_id: &admit.request_id, - requested_model: model, - served_model: model, + requested_model: &m.requested, + served_model: &m.served, protocol: mt.as_str(), account, prompt: it, @@ -408,6 +460,7 @@ async fn bill_realtime_turn( total, ) .await; + state.avail.record(&m.requested, true); metrics::counter!("gateway_tokens_total", "kind" => "prompt").increment(it as u64); metrics::counter!("gateway_tokens_total", "kind" => "completion").increment(ot as u64); } @@ -417,7 +470,7 @@ async fn realtime_session( mut socket: axum::extract::ws::WebSocket, s: AppState, ak: AkInfo, - model: String, + rtm: RtModel, mt: gw_consts::Protocol, account: String, hint: String, @@ -427,7 +480,7 @@ async fn realtime_session( let _ = socket .send(send(json!({"type":"session.created", - "session":{"model": model, "account": account}}))) + "session":{"model": rtm.requested, "account": account}}))) .await; while let Some(Ok(msg)) = socket.recv().await { @@ -450,7 +503,7 @@ async fn realtime_session( } match ev["type"].as_str().unwrap_or_default() { "input_text" => { - let admit = match realtime_gate(&s, &ak, &model, &hint).await { + let admit = match realtime_gate(&s, &ak, &rtm, &hint).await { Ok(a) => a, Err(denied) => { let _ = socket @@ -460,7 +513,7 @@ async fn realtime_session( } }; let input = ev["text"].as_str().unwrap_or_default().to_owned(); - let reply = format!("[mock-realtime:{model}] you said: {input}"); + let reply = format!("[mock-realtime:{}] you said: {input}", rtm.served); let (it, ot) = ( (input.len() as i64 / 4).max(1) + 3, (reply.len() as i64 / 4).max(1), @@ -480,7 +533,7 @@ async fn realtime_session( .send(send(json!({"type":"response.done", "usage":{"input_tokens": it, "output_tokens": ot}}))) .await; - bill_realtime_turn(&admit, &model, mt, &account, it, ot).await; + bill_realtime_turn(&admit, &rtm, mt, &account, it, ot).await; } "session.close" => { let _ = socket.send(send(json!({"type":"session.closed"}))).await; @@ -531,7 +584,7 @@ async fn realtime_bridge( mut client: axum::extract::ws::WebSocket, s: AppState, ak: AkInfo, - model: String, + rtm: RtModel, mt: gw_consts::Protocol, account: Arc, hint: String, @@ -547,10 +600,11 @@ async fn realtime_bridge( let ws_base = base .replacen("https://", "wss://", 1) .replacen("http://", "ws://", 1); - let url = format!("{ws_base}/v1/realtime?model={model}"); + let url = format!("{ws_base}/v1/realtime?model={}", rtm.served); let mut req = match url.into_client_request() { Ok(r) => r, Err(e) => { + s.handler.state().avail.record(&rtm.requested, false); let _ = client .send(CMsg::Text( send_err(format!("bad upstream url: {e}")) @@ -568,6 +622,7 @@ async fn realtime_bridge( let upstream = match tokio_tungstenite::connect_async(req).await { Ok((u, _)) => u, Err(e) => { + s.handler.state().avail.record(&rtm.requested, false); let _ = client .send(CMsg::Text( send_err(format!("upstream realtime connect failed: {e}")) @@ -626,7 +681,7 @@ async fn realtime_bridge( // upstream rejects the duplicate, and a raced accept is // caught by the response.created gate below if is_response_create(&frame) && pending.is_none() { - match realtime_gate(&s, &ak, &model, &hint).await { + match realtime_gate(&s, &ak, &rtm, &hint).await { Ok(admit) => { pending = Some(admit); generations += 1; @@ -659,7 +714,12 @@ async fn realtime_bridge( Some(Ok(UMsg::Binary(b))) => { (serde_json::from_slice::(&b).ok(), false, None, Some(b)) } - Some(Ok(UMsg::Close(_))) | Some(Err(_)) | None => break, + Some(Err(_)) => { + // upstream died mid-session — a client-visible model error + s.handler.state().avail.record(&rtm.requested, false); + break; + } + Some(Ok(UMsg::Close(_))) | None => break, Some(Ok(_)) => continue, // ping/pong handled by the ws stacks }; let mut relay = true; @@ -676,7 +736,7 @@ async fn realtime_bridge( // server-VAD: OpenAI auto-starts a turn with no client // response.create — gate it here like a manual one else if realtime_turn_started(&account.provider, &v) && pending.is_none() { - match realtime_gate(&s, &ak, &model, &hint).await { + match realtime_gate(&s, &ak, &rtm, &hint).await { Ok(admit) => pending = Some(admit), Err(denied) => { let _ = up_tx @@ -694,7 +754,7 @@ async fn realtime_bridge( // a boundary with no gated turn bills unreserved match pending.take() { Some(a) => { - bill_realtime_turn(&a, &model, mt, &account.name, it, ot).await + bill_realtime_turn(&a, &rtm, mt, &account.name, it, ot).await } None if it.saturating_add(ot) > 0 => { // re-authenticate so billing uses the key's current @@ -718,7 +778,7 @@ async fn realtime_bridge( }; bill_realtime_turn( &unreserved, - &model, + &rtm, mt, &account.name, it, @@ -779,7 +839,7 @@ async fn realtime_bridge( if generations > 0 && recognized == 0 { tracing::warn!( account = %account.name, - model = %model, + model = %rtm.requested, generations, "realtime bridge relayed generations but saw no usage frame — vendor dialect not recognized?" ); @@ -940,6 +1000,9 @@ fn check_key_status(info: &AkInfo) -> Result<(), (u16, &'static str)> { gw_state::KeyStatus::Active => Ok(()), gw_state::KeyStatus::Banned => Err((403, "access key is banned")), gw_state::KeyStatus::Expired => Err((403, "access key has expired")), + gw_state::KeyStatus::Suspended => { + Err((403, "access key is suspended for abuse; retry later")) + } } } @@ -1070,8 +1133,28 @@ async fn rt_inbound_policy( if let Some(block) = scan.block { return Err(block.message); } - if let Some(reason) = realtime_moderate(s, sec, ak, hint, &text).await { - return Err(reason); + if sec.moderate && !text.is_empty() { + match s.handler.moderate_rt(sec, &text).await { + gw_handler::RtModeration::Allow => {} + gw_handler::RtModeration::Mask(spans) => { + let masked = gw_handler::plugins::apply_mask_spans_frame(frame, &spans); + if masked > 0 { + write_rt_event( + s, + ak, + ak.attributed_user(hint), + "moderation", + "mask", + masked as i64, + ) + .await; + } + } + gw_handler::RtModeration::Deny(reason) => { + write_rt_event(s, ak, ak.attributed_user(hint), "moderation", "block", 1).await; + return Err(reason); + } + } } let redacted = gw_handler::plugins::dlp_redact_realtime_frame(sec, frame); if redacted > 0 { @@ -1102,25 +1185,6 @@ async fn emit_rt_hits( } } -/// Moderate a realtime frame's inbound `text` (collected by the frame scan) -/// via the wired moderator — parity with the REST surface, so a moderated -/// tenant can't be bypassed over the WebSocket. `Some(reason)` denies the -/// frame; records a moderation event. -async fn realtime_moderate( - s: &AppState, - sec: &gw_config::SecurityConf, - ak: &AkInfo, - hint: &str, - text: &str, -) -> Option { - if !sec.moderate || text.is_empty() { - return None; - } - let reason = s.handler.moderate_text(sec, text).await?; - write_rt_event(s, ak, ak.attributed_user(hint), "moderation", "block", 1).await; - Some(reason) -} - /// The caller IP for the admin audit trail, resolved at request entry — before /// any config mutation the handler performs — so the op that flips /// `trust_proxy_headers` is audited under the policy in effect when it @@ -1255,6 +1319,7 @@ fn ak_public_json(k: &AkInfo) -> Value { "tokens_per_minute": k.tokens_per_minute, "expires_at_epoch_secs": k.expires_at_epoch_secs, "banned": k.banned, + "suspended_until_epoch_secs": k.suspended_until_epoch_secs, }) } @@ -1330,6 +1395,11 @@ async fn admin_key_create( let (Some(ak), Some(product)) = (body["ak"].as_str(), body["product"].as_str()) else { return error_response(400, "ak and product are required"); }; + // same ban as config load: a ':' would collide with the prefixed + // governance keyspaces (`abuse:{ak}` could force-suspend another key) + if ak.is_empty() || ak.contains(':') { + return error_response(400, "ak must be non-empty and must not contain ':'"); + } let default_tenant = match &scope { AdminScope::Global => gw_config::DEFAULT_TENANT, AdminScope::Tenant(t) => t.as_str(), @@ -1358,6 +1428,7 @@ async fn admin_key_create( tokens_per_minute: body["tokens_per_minute"].as_i64(), expires_at_epoch_secs: body["expires_at_epoch_secs"].as_i64(), banned: body["banned"].as_bool().unwrap_or(false), + suspended_until_epoch_secs: None, model_quotas: Arc::new( body["model_quotas"] .as_object() @@ -1417,6 +1488,7 @@ async fn admin_key_patch( tokens_per_minute: tri("tokens_per_minute"), expires_at_epoch_secs: tri("expires_at_epoch_secs"), banned: body["banned"].as_bool(), + suspended_until_epoch_secs: tri("suspended_until_epoch_secs"), }; let patched = s.handler.state().auth.patch(&ak, &patch).await; match patched { @@ -1531,21 +1603,18 @@ async fn admin_key_list( } /// GET /admin/models/status — per-model availability over the configured -/// window, judged from minute-bucketed success/error counts. A tenant admin -/// sees only its entitled models. Realtime models are excluded: never -/// sampled, listing them would read as a permanent no_data. +/// window, judged from minute-bucketed success/error counts (REST terminal +/// outcomes; realtime samples per billed turn and on session-fatal upstream +/// errors). A tenant admin sees only its entitled models. async fn admin_models_status(State(s): State, scope: AdminScope) -> Response { let cfg = s.handler.cfg(); let st = &cfg.stability; let until = gw_state::epoch_secs() / 60; let since = until - (st.availability_window_minutes - 1); let state = s.handler.state(); - let entitled = cfg.models.iter().filter(|m| { - m.protocol() != Some(gw_consts::Protocol::Realtime) - && match &scope { - AdminScope::Tenant(t) => cfg.tenant_allows_model(t, &m.name), - AdminScope::Global => true, - } + let entitled = cfg.models.iter().filter(|m| match &scope { + AdminScope::Tenant(t) => cfg.tenant_allows_model(t, &m.name), + AdminScope::Global => true, }); let avail = &state.avail; let counts = futures::future::join_all( @@ -2590,6 +2659,22 @@ async fn run_family( /// A pre-stage content block answers 400 with the block message — these /// surfaces have no in-band content_filter shape, and falling through would /// misreport the block as an engine failure. +/// An `input`-style field that may be a string or an array of strings +/// (the OpenAI embeddings/moderations shape). +fn string_or_string_array(v: Option) -> Vec { + match v { + Some(Value::String(x)) => vec![x], + Some(Value::Array(a)) => a + .into_iter() + .filter_map(|v| match v { + Value::String(s) => Some(s), + _ => None, + }) + .collect(), + _ => vec![], + } +} + fn response_v2_or_500(outcome: Option, engine: &str) -> Response { match outcome { Some(o) if o.block.block => error_response(400, o.response.message), @@ -2811,17 +2896,7 @@ async fn embeddings( ) -> Response { let started = Instant::now(); let model = body["model"].as_str().unwrap_or_default().to_owned(); - let input: Vec = match body.get_mut("input").map(Value::take) { - Some(Value::String(x)) => vec![x], - Some(Value::Array(a)) => a - .into_iter() - .filter_map(|v| match v { - Value::String(s) => Some(s), - _ => None, - }) - .collect(), - _ => vec![], - }; + let input = string_or_string_array(body.get_mut("input").map(Value::take)); if model.is_empty() || input.is_empty() { return error_response(400, "model and input are required"); } @@ -2833,7 +2908,7 @@ async fn embeddings( &s, ak, model, - gw_consts::Protocol::OpenaiChat, + gw_consts::Protocol::Embeddings, typed, vec![], user_hint(&headers, &body["user"]), @@ -2870,7 +2945,7 @@ async fn images_generations( &s, ak, model, - gw_consts::Protocol::OpenaiChat, + gw_consts::Protocol::Image, typed, vec![], user_hint(&headers, &body["user"]), @@ -2910,7 +2985,7 @@ async fn images_edits( &s, ak, model, - gw_consts::Protocol::OpenaiChat, + gw_consts::Protocol::Image, typed, vec![], user_hint(&headers, &body["user"]), @@ -2947,7 +3022,7 @@ async fn audio_speech( &s, ak, model, - gw_consts::Protocol::OpenaiChat, + gw_consts::Protocol::Tts, typed, vec![], user_hint(&headers, &body["user"]), @@ -2988,6 +3063,29 @@ async fn audio_transcriptions( headers: HeaderMap, Authed(ak): Authed, Json(body): Json, +) -> Response { + audio_transcribe(s, headers, ak, body, false).await +} + +/// POST /v1/audio/translations — the transcriptions shape, translated to +/// English by the upstream (OpenAI translations semantics). +async fn audio_translations( + State(s): State, + headers: HeaderMap, + Authed(ak): Authed, + Json(body): Json, +) -> Response { + audio_transcribe(s, headers, ak, body, true).await +} + +/// The shared STT body: transcriptions and translations differ only in the +/// upstream path the `translate` flag selects. +async fn audio_transcribe( + s: AppState, + headers: HeaderMap, + ak: AkInfo, + body: Value, + translate: bool, ) -> Response { let started = Instant::now(); let model = body["model"].as_str().unwrap_or_default().to_owned(); @@ -2998,12 +3096,13 @@ async fn audio_transcriptions( let typed = TypedParams::AudioStt(SttParams { audio_b64: audio, language: body["language"].as_str().map(str::to_owned), + translate, }); let ctx = match run_family( &s, ak, model, - gw_consts::Protocol::OpenaiChat, + gw_consts::Protocol::Stt, typed, vec![], user_hint(&headers, &body["user"]), @@ -3013,7 +3112,12 @@ async fn audio_transcriptions( Ok(ctx) => ctx, Err(resp) => return resp, }; - log_access("audio_transcriptions", &ctx, started); + let surface = if translate { + "audio_translations" + } else { + "audio_transcriptions" + }; + log_access(surface, &ctx, started); match ctx.outcome { Some(o) if o.block.block => error_response(400, o.response.message), Some(o) => (StatusCode::OK, Json(json!({ "text": o.response.message }))).into_response(), @@ -3021,6 +3125,85 @@ async fn audio_transcriptions( } } +/// POST /v1/moderations — OpenAI moderations shape; input may be a string or +/// an array of strings. +async fn moderations( + State(s): State, + headers: HeaderMap, + Authed(ak): Authed, + Json(mut body): Json, +) -> Response { + let started = Instant::now(); + let model = body["model"].as_str().unwrap_or_default().to_owned(); + let input = string_or_string_array(body.get_mut("input").map(Value::take)); + if model.is_empty() || input.is_empty() { + return error_response(400, "model and input are required"); + } + let typed = TypedParams::Moderation(gw_models::ModerationParams { input }); + let ctx = match run_family( + &s, + ak, + model, + gw_consts::Protocol::Moderations, + typed, + vec![], + user_hint(&headers, &body["user"]), + ) + .await + { + Ok(ctx) => ctx, + Err(resp) => return resp, + }; + log_access("moderations", &ctx, started); + response_v2_or_500(ctx.outcome, "moderations") +} + +/// POST /v1/rerank — Cohere/Jina-compatible: `{model, query, documents, top_n?}`. +async fn rerank( + State(s): State, + headers: HeaderMap, + Authed(ak): Authed, + Json(mut body): Json, +) -> Response { + let started = Instant::now(); + let model = body["model"].as_str().unwrap_or_default().to_owned(); + let query = body["query"].as_str().unwrap_or_default().to_owned(); + let documents: Vec = match body.get_mut("documents").map(Value::take) { + Some(Value::Array(a)) => a + .into_iter() + .filter_map(|v| match v { + Value::String(s) => Some(s), + _ => None, + }) + .collect(), + _ => vec![], + }; + if model.is_empty() || query.is_empty() || documents.is_empty() { + return error_response(400, "model, query, and documents are required"); + } + let typed = TypedParams::Rerank(gw_models::RerankParams { + query, + documents, + top_n: body["top_n"].as_i64(), + }); + let ctx = match run_family( + &s, + ak, + model, + gw_consts::Protocol::Rerank, + typed, + vec![], + user_hint(&headers, &body["user"]), + ) + .await + { + Ok(ctx) => ctx, + Err(resp) => return resp, + }; + log_access("rerank", &ctx, started); + response_v2_or_500(ctx.outcome, "rerank") +} + fn parse_batch_messages(v: &Value) -> Vec { v["messages"] .as_array() @@ -3221,6 +3404,14 @@ mod tests { )) } + fn rt(name: &str) -> RtModel { + RtModel { + requested: name.to_owned(), + served: name.to_owned(), + from_config: true, + } + } + async fn body_json(resp: Response) -> Value { let bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) .await @@ -3439,10 +3630,11 @@ mod tests { assert_eq!(of("m-ok")["requests"], 4); assert_eq!(of("m-bad")["state"], "unavailable"); assert_eq!(of("m-bad")["errors"], 4); - assert!( - !rows.iter().any(|r| r["model"] == "rt-m"), - "realtime models are not listed (never sampled)" - ); + let rt_row = rows + .iter() + .find(|r| r["model"] == "rt-m") + .expect("realtime models are listed (turn-sampled)"); + assert_eq!(rt_row["state"], "no_data", "no turns billed yet"); let resp = router.clone().oneshot(status("t1-tok")).await.unwrap(); let j = body_json(resp).await; @@ -3692,11 +3884,10 @@ mod tests { let cfg = app.handler.cfg(); let sec = cfg.security_for(&ak.tenant); + let mut frame = json!({"type":"input_text","text":"hello there"}); assert_eq!( - realtime_moderate(&app, sec, &ak, "", "hello there") - .await - .as_deref(), - Some("blocked by moderator") + rt_inbound_policy(&app, &ak, "", &mut frame).await, + Err("blocked by moderator".to_owned()) ); let mut secret = json!({"type":"input_text","text":"sk-abcdefghijklmnopqrstuvwxyz012345"}); let n = gw_handler::plugins::dlp_redact_realtime_frame(sec, &mut secret); @@ -3720,6 +3911,110 @@ mod tests { ); } + #[derive(Debug)] + struct FrameMaskModerator; + + #[async_trait::async_trait] + impl gw_handler::moderation::Moderator for FrameMaskModerator { + async fn review(&self, text: &str) -> Result { + match text.find("secret") { + Some(i) => Ok(gw_handler::moderation::Verdict::Mask( + std::iter::once(i..i + "secret".len()).collect(), + )), + None => Ok(gw_handler::moderation::Verdict::Allow), + } + } + } + + #[tokio::test] + async fn realtime_moderation_masks_the_frame() { + let yaml = "listen: {host: h, port: 1}\nsecurity: {moderate: true}\nmodels: [{name: rt, protocol: realtime}]\naccess_keys: [{ak: k1, product: p, qps: 10, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let handler = OnlineHandler::new( + gw_state::SharedConfig::new(cfg, state), + Arc::new(gw_engines::MockTransport), + ) + .with_moderator(Arc::new(FrameMaskModerator)); + let offline = OfflineHandler::new(handler.clone()); + let app = AppState { + handler, + offline, + loader: None, + config_store: None, + }; + let ak = app.handler.state().auth.authenticate("k1").await.unwrap(); + let mut frame = json!({"type":"input_text","text":"tell secret now"}); + assert_eq!(rt_inbound_policy(&app, &ak, "", &mut frame).await, Ok(0)); + assert_eq!( + frame["text"], "tell [MASKED] now", + "the mask lands in the frame before it is forwarded" + ); + let events = app + .handler + .state() + .store + .security_events(None, 10) + .await + .unwrap(); + assert!( + events + .iter() + .any(|e| e.rule == "moderation" && e.action == "mask") + ); + } + + #[tokio::test] + async fn moderations_rerank_and_translations_roundtrip() { + let post = |path: &str, body: &str| { + Request::builder() + .method("POST") + .uri(path.to_owned()) + .header("content-type", "application/json") + .header("authorization", "Bearer ak-demo-123") + .body(Body::from(body.to_owned())) + .unwrap() + }; + let resp = test_app() + .oneshot(post( + "/v1/moderations", + r#"{"model":"text-moderation","input":["fine text","really unsafe text"]}"#, + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let j = body_json(resp).await; + assert_eq!(j["results"][0]["flagged"], false); + assert_eq!(j["results"][1]["flagged"], true); + + let resp = test_app() + .oneshot(post( + "/v1/rerank", + r#"{"model":"rerank-mini","query":"rust gateway","documents":["a rust gateway","cooking pasta"],"top_n":1}"#, + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let j = body_json(resp).await; + let results = j["results"].as_array().unwrap(); + assert_eq!(results.len(), 1, "top_n honored"); + assert_eq!(results[0]["index"], 0, "the matching document ranks first"); + + let resp = test_app() + .oneshot(post( + "/v1/audio/translations", + r#"{"model":"whisper-1","audio_b64":"AAAA"}"#, + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let j = body_json(resp).await; + assert!( + j["text"].as_str().unwrap().contains("translated"), + "the translations path reached the translations mock: {j}" + ); + } + #[tokio::test] async fn chat_non_stream_ok() { let resp = test_app() @@ -3768,21 +4063,43 @@ mod tests { let gov = || s.handler.state().governance.clone(); let used = || async { gov().quota_used(&ak.ak).await }; - let a1 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); + let a1 = realtime_gate(&s, &ak, &rt("gpt-4o"), "") + .await + .expect("admit"); assert_eq!(used().await, REALTIME_TURN_RESERVE, "reserved up front"); - bill_realtime_turn(&a1, "gpt-4o", gw_consts::Protocol::Realtime, "acc", 30, 70).await; + bill_realtime_turn( + &a1, + &rt("gpt-4o"), + gw_consts::Protocol::Realtime, + "acc", + 30, + 70, + ) + .await; assert_eq!(used().await, 100, "settled to actual (30 + 70)"); - let a2 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); + let a2 = realtime_gate(&s, &ak, &rt("gpt-4o"), "") + .await + .expect("admit"); assert_eq!(used().await, 100 + REALTIME_TURN_RESERVE); gov().quota_settle(&a2.ak.ak, -a2.reserved, a2.at).await; assert_eq!(used().await, 100, "dropped turn refunded whole"); - let a3 = realtime_gate(&s, &ak, "gpt-4o", "").await.expect("admit"); + let a3 = realtime_gate(&s, &ak, &rt("gpt-4o"), "") + .await + .expect("admit"); assert_eq!(used().await, 100 + REALTIME_TURN_RESERVE); let ledger_before = s.handler.state().store.ledger_snapshot(1).await.unwrap().0; - bill_realtime_turn(&a3, "gpt-4o", gw_consts::Protocol::Realtime, "acc", 0, 0).await; + bill_realtime_turn( + &a3, + &rt("gpt-4o"), + gw_consts::Protocol::Realtime, + "acc", + 0, + 0, + ) + .await; assert_eq!(used().await, 100, "zero-usage turn refunds its reserve"); let ledger_after = s.handler.state().store.ledger_snapshot(1).await.unwrap().0; assert_eq!( @@ -3798,8 +4115,18 @@ mod tests { let state = Arc::new(GatewayState::from_config(&cfg)); let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); let ak = s.handler.state().auth.authenticate("k1").await.unwrap(); - let a = realtime_gate(&s, &ak, "rt-m", "").await.expect("admit"); - bill_realtime_turn(&a, "rt-m", gw_consts::Protocol::Realtime, "acc", 100, 100).await; + let a = realtime_gate(&s, &ak, &rt("rt-m"), "") + .await + .expect("admit"); + bill_realtime_turn( + &a, + &rt("rt-m"), + gw_consts::Protocol::Realtime, + "acc", + 100, + 100, + ) + .await; let (_, ledger) = s .handler .state() @@ -3819,6 +4146,34 @@ mod tests { ); } + #[tokio::test] + async fn realtime_turn_denied_when_served_model_reloaded_away() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: rt-pub, protocol: realtime, variants: [{model: rt-canary, weight: 1}]}, {name: rt-canary, protocol: realtime}]\naccounts: [{name: a1, provider: openai, protocols: ['realtime']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let ak = s.handler.state().auth.authenticate("k1").await.unwrap(); + let pinned = RtModel { + requested: "rt-pub".into(), + served: "rt-canary".into(), + from_config: true, + }; + assert!(realtime_gate(&s, &ak, &pinned, "").await.is_ok()); + // canary rollback: the reload drops the variant the session pinned + let without = "listen: {host: h, port: 1}\nmodels: [{name: rt-pub, protocol: realtime}]\naccounts: [{name: a1, provider: openai, protocols: ['realtime']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]"; + s.handler + .reload(GatewayConfig::from_yaml(without).unwrap()) + .await + .unwrap(); + let denied = realtime_gate(&s, &ak, &pinned, "").await.err(); + assert!( + denied + .as_deref() + .is_some_and(|e| e.contains("no longer configured")), + "a pinned variant removed by reload must deny the turn, not bill zero: {denied:?}" + ); + } + #[tokio::test] async fn realtime_gate_reserves_tpm_and_rolls_back_on_denial() { let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); @@ -3833,14 +4188,14 @@ mod tests { .unwrap(); let gov = s.handler.state().governance.clone(); - let a1 = realtime_gate(&s, &ak, "gpt-4o", "") + let a1 = realtime_gate(&s, &ak, &rt("gpt-4o"), "") .await .expect("first admits"); assert_eq!(a1.tpm_reserved, Some(REALTIME_TURN_RESERVE)); let daily_before = gov.quota_used(&ak.ak).await; assert!( - realtime_gate(&s, &ak, "gpt-4o", "").await.is_err(), + realtime_gate(&s, &ak, &rt("gpt-4o"), "").await.is_err(), "second turn denied by the TPM reserve" ); assert_eq!( @@ -3862,12 +4217,20 @@ mod tests { let s = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); let ak = s.handler.state().auth.authenticate("k-rt").await.unwrap(); - let admit = realtime_gate(&s, &ak, "rt", "").await.expect("admit"); + let admit = realtime_gate(&s, &ak, &rt("rt"), "").await.expect("admit"); s.handler .reload(GatewayConfig::from_yaml(&price(2_000_000)).unwrap()) .await .unwrap(); - bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 100, 100).await; + bill_realtime_turn( + &admit, + &rt("rt"), + gw_consts::Protocol::Realtime, + "acc", + 100, + 100, + ) + .await; let (_, records) = s.handler.state().store.ledger_snapshot(1).await.unwrap(); assert_eq!( @@ -3890,11 +4253,19 @@ mod tests { .authenticate("k-shared") .await .unwrap(); - let admit = realtime_gate(&s, &shared, "rt", "alice") + let admit = realtime_gate(&s, &shared, &rt("rt"), "alice") .await .expect("admit"); assert_eq!(admit.user, "alice", "ownerless key attributes to the hint"); - bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 40, 60).await; + bill_realtime_turn( + &admit, + &rt("rt"), + gw_consts::Protocol::Realtime, + "acc", + 40, + 60, + ) + .await; let owned = s .handler @@ -3903,14 +4274,22 @@ mod tests { .authenticate("k-owned") .await .unwrap(); - let admit = realtime_gate(&s, &owned, "rt", "mallory") + let admit = realtime_gate(&s, &owned, &rt("rt"), "mallory") .await .expect("admit"); assert_eq!( admit.user, "bob", "owner is authoritative over a spoofed hint" ); - bill_realtime_turn(&admit, "rt", gw_consts::Protocol::Realtime, "acc", 10, 20).await; + bill_realtime_turn( + &admit, + &rt("rt"), + gw_consts::Protocol::Realtime, + "acc", + 10, + 20, + ) + .await; let (_, records) = s.handler.state().store.ledger_snapshot(2).await.unwrap(); let users: std::collections::HashSet<&str> = diff --git a/docs/api.md b/docs/api.md index 4e58d38..9cef96f 100644 --- a/docs/api.md +++ b/docs/api.md @@ -39,8 +39,16 @@ user. See [Governance](governance.md#per-user-attribution-and-billing). | POST | `/v1/images/edits` | source image + optional mask (base64) | | POST | `/v1/audio/speech` | TTS, returns audio bytes | | POST | `/v1/audio/transcriptions` | STT, JSON carries base64 audio | +| POST | `/v1/audio/translations` | STT translated to English (same request shape) | +| POST | `/v1/moderations` | content moderation; `input` string or array, native results pass through | | GET | `/v1/models` | configured public model names | +## Rerank + +| Method | Path | Notes | +|--------|------|-------| +| POST | `/v1/rerank` | Cohere/Jina-compatible: `{model, query, documents, top_n?}` → `{results: [{index, relevance_score}]}` | + ### Chat completions ```bash @@ -137,11 +145,11 @@ regardless. | PUT | `/admin/config` | validate + publish a new config document to the fleet config store; every instance reloads via the change feed (global token; needs `storage.postgres_url`) | | GET | `/admin/keys` | list keys, `?offset=&limit=` paged (default 200; a tenant token sees only its own tenant's) | | POST | `/admin/keys` | create/replace a key: `{ak, product, tenant?, owner?, qps, daily_token_quota, tokens_per_minute?, expires_at_epoch_secs?, banned?, model_quotas?}` (`owner` binds the key to one end user — authoritative for attribution) | -| PATCH | `/admin/keys/{ak}` | update any of `qps` / `daily_token_quota` / `tokens_per_minute` / `expires_at_epoch_secs` (null clears) / `banned` | +| PATCH | `/admin/keys/{ak}` | update any of `qps` / `daily_token_quota` / `tokens_per_minute` / `expires_at_epoch_secs` (null clears) / `banned` / `suspended_until_epoch_secs` (null lifts an abuse suspension early) | | DELETE | `/admin/keys/{ak}` | revoke a key | | GET | `/admin/usage` | ledger rollup by tenant × model (requests, tokens, charged `cost_micros`, `vendor_cost_micros` for margin); `?tenant=` filter for the global token | | GET | `/admin/usage/users` | per-user cost rollup (user × model) over a billing period: `?since=&until=` (unix secs), `?user=` filter, `?format=csv` export; tenant-scoped | -| GET | `/admin/models/status` | per-model availability over the recent window (`available` / `unstable` / `unavailable` / `no_data`), judged from client-visible request outcomes against `stability.*` thresholds; attributes to the requested public name under a `variants` split; realtime models are not sampled and not listed; tenant-scoped | +| GET | `/admin/models/status` | per-model availability over the recent window (`available` / `unstable` / `unavailable` / `no_data`), judged from client-visible outcomes against `stability.*` thresholds; attributes to the requested public name under a `variants` split; realtime models sample per billed turn and on session-fatal upstream errors; tenant-scoped | | GET | `/admin/audit/events` | content-safety hits (blocklist / regex / DLP / moderation) recorded without prompt text; `?limit=`; tenant-scoped | | GET | `/admin/audit/ops` | admin-operation trail (key CRUD, config publish, reload) with actor, target, and source IP; `?limit=`; global token only | | GET | `/admin/audit/content/{request_id}` | retained prompt/response for one request, unsealed when `GW_CONTENT_KEY` is set (sealed rows without it return `content: null`); tenant-scoped | diff --git a/docs/configuration.md b/docs/configuration.md index 2f4632c..23c16d3 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -115,17 +115,18 @@ models: `token_rate` weights scale cost and quota consumption per token component; the ledger's prompt/completion columns stay vendor-reported, while `total_tokens` is the weighted platform total. `variants` splits a public name across other -declared same-protocol models (one level, no realtime): entitlement and the -per-(AK, model) daily counter judge the public name, billing prices the served -variant, and the response echoes the requested name. Selection hashes the -effective user, so a user sticks to one backend across the fleet. +declared same-protocol models (one level): entitlement and the per-(AK, model) +daily counter judge the public name, billing prices the served variant, and +the response echoes the requested name. Selection hashes the effective user, +so a user sticks to one backend across the fleet; a realtime session picks +its variant once at the handshake and pins it for the whole session. ### `providers` — first-class provider presets ```yaml providers: - name: openai - kind: openai # openai | anthropic | gemini | deepseek | openrouter + kind: openai # openai | anthropic | gemini | deepseek | openrouter | moonshot | siliconflow api_key_env: OPENAI_API_KEY # endpoint / timeout_seconds / connect_retries / secret_key_env may be # set here too and are inherited by the synthesized account @@ -177,6 +178,8 @@ security: # global default; a tenant may override it whole regex_rules: # named recognizers, each with its own action - {name: ssn, pattern: '\d{3}-\d{2}-\d{4}', action: block} moderate: false # route inbound text through the wired external moderator + # (its verdicts: allow / mask spans / degrade to the + # tenant fallback model / deny) moderation_fail_open: false # on a moderator error: admit (true) or deny (false) stability: @@ -190,6 +193,15 @@ stability: products: - name: myproduct qpm: 120 # product-level request rate + +abuse: # automatic suspension; omit = off + tiers: # highest tier at or under the day's reject count wins + - {rejects: 20, suspend_hours: 2} + - {rejects: 30, suspend_hours: 24} + +alerts: # outbound webhook; omit = off + webhook_url_env: GW_ALERT_WEBHOOK # env var naming the URL (secrets stay out of config) + dedup_seconds: 300 # mute repeats of the same (kind, subject) ``` Every rule that fires (block / flag / DLP / moderation) is recorded without the