diff --git a/conf/gateway.yaml b/conf/gateway.yaml index ede32b6..385b4ec 100644 --- a/conf/gateway.yaml +++ b/conf/gateway.yaml @@ -152,6 +152,11 @@ models: input_price_per_1k_micros: 150 output_price_per_1k_micros: 600 cache_ttl_seconds: 60 + # Optional per-model knobs (uncomment to use): + # billing weights per token component (cache reads cheaper, writes dearer) + # token_rate: {read_cache: 0.1, write_cache: 1.25} + # weighted canary split across declared same-protocol models, sticky per user + # variants: [{model: cached-mini, weight: 90}, {model: qpm-mini, weight: 10}] - name: spark-lite # e2e: account cooldown (sole account stays down -> cooldown -> recovers on expiry) protocol: openai-chat provider: iflytek @@ -257,10 +262,15 @@ security: blocklist: - forbiddenword -# Account stability: 3 consecutive failures enter a 2s cooldown (e2e uses a short cooldown) +# Account stability: 3 consecutive failures enter a 2s cooldown (e2e uses a short cooldown). +# The availability_* fields drive /admin/models/status (defaults shown; window max 60). stability: failure_threshold: 3 cooldown_seconds: 2 + availability_window_minutes: 5 + unstable_error_rate: 0.1 + unavailable_error_rate: 0.5 + availability_min_samples: 20 # Product-level governance products: diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 29b4590..8b167d5 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -50,6 +50,16 @@ pub enum ConfigError { BadFallbackModel { tenant: String, model: String }, #[error("`{owner}` sets a negative price")] NegativePrice { owner: String }, + #[error("model `{model}` token_rate `{field}` must be finite and >= 0")] + BadTokenRate { model: String, field: &'static str }, + #[error("stability: bad {field}")] + BadStability { field: &'static str }, + #[error("model `{model}` variant `{variant}`: {reason}")] + BadVariant { + model: String, + variant: String, + reason: &'static str, + }, #[error("`{owner}` sets a negative or non-finite limit")] NegativeLimit { owner: String }, #[error("storage.shared_cache needs storage.redis_url")] @@ -109,6 +119,15 @@ pub struct ModelConf { /// Request-level cache TTL; None = this model isn't cached. #[serde(default)] pub cache_ttl_seconds: Option, + /// Billing weights per token component; None = every component at 1.0. + #[serde(default)] + pub token_rate: Option, + /// Weighted routing split across other declared models (canary/gray); + /// empty = this name serves itself. A self-referencing entry keeps a + /// share on this model. Once routed, the variant's own qpm/cache/ + /// token_rate apply — model-level limits guard backend capacity. + #[serde(default)] + pub variants: Vec, } impl ModelConf { @@ -117,6 +136,46 @@ impl ModelConf { } } +/// One weighted routing target of a public model name. +#[derive(Debug, Clone, Deserialize)] +pub struct VariantConf { + pub model: String, + pub weight: u32, +} + +/// 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)] +pub struct TokenRateConf { + #[serde(default = "weight_one")] + pub prompt: f64, + #[serde(default = "weight_one")] + pub read_cache: f64, + #[serde(default = "weight_one")] + pub write_cache: f64, + #[serde(default = "weight_one")] + pub completion: f64, + #[serde(default = "weight_one")] + pub reasoning: f64, +} + +impl TokenRateConf { + /// `(field name, value)` pairs for validation. + fn fields(&self) -> [(&'static str, f64); 5] { + [ + ("prompt", self.prompt), + ("read_cache", self.read_cache), + ("write_cache", self.write_cache), + ("completion", self.completion), + ("reasoning", self.reasoning), + ] + } +} + +fn weight_one() -> f64 { + 1.0 +} + /// Upstream account slot (mock credentials unless a live endpoint is configured). #[derive(Debug, Clone, Deserialize)] pub struct AccountConf { @@ -249,11 +308,36 @@ pub struct StabilityConf { /// Cooldown duration (seconds); auto-recovers on expiry. #[serde(default = "default_cooldown_seconds")] pub cooldown_seconds: u64, + /// Minutes of per-model success/error counts the status API judges over. + /// Capped at 60: the availability store retains one hour of buckets. + #[serde(default = "default_availability_window_minutes")] + pub availability_window_minutes: i64, + /// Window error rate at or above which a model reports `unstable`. + #[serde(default = "default_unstable_error_rate")] + pub unstable_error_rate: f64, + /// Window error rate at or above which a model reports `unavailable`. + #[serde(default = "default_unavailable_error_rate")] + pub unavailable_error_rate: f64, + /// Below this many window samples the verdict is `no_data`. + #[serde(default = "default_availability_min_samples")] + pub availability_min_samples: u64, } fn default_failure_threshold() -> usize { 3 } +fn default_availability_window_minutes() -> i64 { + 5 +} +fn default_unstable_error_rate() -> f64 { + 0.1 +} +fn default_unavailable_error_rate() -> f64 { + 0.5 +} +fn default_availability_min_samples() -> u64 { + 20 +} fn default_cooldown_seconds() -> u64 { 30 } @@ -263,6 +347,10 @@ impl Default for StabilityConf { Self { failure_threshold: default_failure_threshold(), cooldown_seconds: default_cooldown_seconds(), + availability_window_minutes: default_availability_window_minutes(), + unstable_error_rate: default_unstable_error_rate(), + unavailable_error_rate: default_unavailable_error_rate(), + availability_min_samples: default_availability_min_samples(), } } } @@ -607,12 +695,53 @@ impl GatewayConfig { // negative prices would make cost accounting non-monotonic (the usage // rollup's max-upsert relies on per-column monotone sums) let neg = |i: i64, o: i64| i < 0 || o < 0; + let neg_or_nan = |v: f64| !v.is_finite() || v < 0.0; for m in &self.models { if neg(m.input_price_per_1k_micros, m.output_price_per_1k_micros) { return Err(ConfigError::NegativePrice { owner: format!("model {}", m.name), }); } + // a NaN weight would poison every downstream cost computation + if let Some(rate) = &m.token_rate { + for (field, v) in rate.fields() { + if neg_or_nan(v) { + return Err(ConfigError::BadTokenRate { + model: m.name.clone(), + field, + }); + } + } + } + for (i, v) in m.variants.iter().enumerate() { + let bad = |reason| ConfigError::BadVariant { + model: m.name.clone(), + 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")); + } + // a duplicate would silently double the target's share + if m.variants[..i].iter().any(|p| p.model == v.model) { + return Err(bad("duplicate target")); + } + let Some(target) = self.models.iter().find(|t| t.name == v.model) else { + return Err(bad("unknown model")); + }; + if target.protocol() != m.protocol() { + return Err(bad("protocol differs from the parent")); + } + // one level only: nested splits would make routing unauditable + if target.name != m.name && !target.variants.is_empty() { + return Err(bad("variant target declares variants itself")); + } + } } for a in &self.accounts { if neg( @@ -636,9 +765,8 @@ impl GatewayConfig { // a negative quota/qps is a typo that would silently deny (or never // limit); a NaN qps would bypass the rate bucket — reject both at load let neg_limit = |owner: String| ConfigError::NegativeLimit { owner }; - let bad_qps = |v: f64| !v.is_finite() || v < 0.0; for k in &self.access_keys { - if bad_qps(k.qps) + if neg_or_nan(k.qps) || k.daily_token_quota < 0 || k.tokens_per_minute.is_some_and(|v| v < 0) || k.model_quotas.values().any(|v| *v < 0) @@ -647,7 +775,7 @@ impl GatewayConfig { } } for t in &self.tenants { - if t.qps.is_some_and(bad_qps) + if t.qps.is_some_and(neg_or_nan) || t.user_daily_token_quota.is_some_and(|v| v < 0) || t.model_quotas.values().any(|v| *v < 0) { @@ -664,6 +792,25 @@ impl GatewayConfig { return Err(neg_limit(format!("product {}", p.name))); } } + 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 + // would silently judge over already-evicted buckets + if !(1..=60).contains(&st.availability_window_minutes) { + return Err(ConfigError::BadStability { + field: "availability_window_minutes (1..=60)", + }); + } + if bad_rate(st.unstable_error_rate) || bad_rate(st.unavailable_error_rate) { + return Err(ConfigError::BadStability { + field: "error rates (must be finite, within [0, 1])", + }); + } + if st.unstable_error_rate > st.unavailable_error_rate { + return Err(ConfigError::BadStability { + field: "unstable_error_rate above unavailable_error_rate", + }); + } for m in &self.models { if m.protocol().is_none() { return Err(ConfigError::UnknownModelMapping { @@ -1160,6 +1307,80 @@ tenants: [{name: t1}, {name: t1}] "negative prices are rejected at load" ); + for bad in [".nan", "-0.5"] { + let bad_rate = format!( + "listen: {{host: h, port: 1}}\nmodels: [{{name: m1, protocol: openai-chat, token_rate: {{read_cache: {bad}}}}}]" + ); + assert!( + matches!( + GatewayConfig::from_yaml(&bad_rate), + Err(ConfigError::BadTokenRate { + field: "read_cache", + .. + }) + ), + "token_rate `{bad}` is rejected at load" + ); + } + + let rate = "listen: {host: h, port: 1}\nmodels: [{name: m1, protocol: openai-chat, token_rate: {read_cache: 0.1, write_cache: 1.25}}]"; + let cfg = GatewayConfig::from_yaml(rate).unwrap(); + let tr = cfg.find_model("m1").unwrap().token_rate.unwrap(); + assert_eq!((tr.read_cache, tr.write_cache), (0.1, 1.25)); + assert_eq!((tr.prompt, tr.completion, tr.reasoning), (1.0, 1.0, 1.0)); + + for (variants, reason) in [ + ("[{model: nope, weight: 1}]", "unknown target"), + ("[{model: m2, weight: 0}]", "zero weight"), + ("[{model: emb, weight: 1}]", "cross-protocol"), + ("[{model: nested, weight: 1}]", "nested variants"), + ( + "[{model: m2, weight: 5}, {model: m2, weight: 5}]", + "duplicate target", + ), + ] { + let bad = format!( + "listen: {{host: h, port: 1}}\nmodels: [{{name: m1, protocol: openai-chat, variants: {variants}}}, {{name: m2, protocol: openai-chat}}, {{name: emb, protocol: embeddings}}, {{name: nested, protocol: openai-chat, variants: [{{model: m2, weight: 1}}]}}]" + ); + assert!( + matches!( + GatewayConfig::from_yaml(&bad), + Err(ConfigError::BadVariant { .. }) + ), + "{reason} is rejected at load" + ); + } + let split = "listen: {host: h, port: 1}\nmodels: [{name: m1, protocol: openai-chat, variants: [{model: m1, weight: 9}, {model: m2, weight: 1}]}, {name: m2, protocol: openai-chat}]"; + let cfg = GatewayConfig::from_yaml(split).unwrap(); + assert_eq!(cfg.find_model("m1").unwrap().variants.len(), 2); + + 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" + ); + + for bad in [ + "availability_window_minutes: 0", + "availability_window_minutes: 61", + "unstable_error_rate: 1.5", + "unstable_error_rate: 0.6, unavailable_error_rate: 0.5", + ] { + let yaml = format!("listen: {{host: h, port: 1}}\nstability: {{{bad}}}"); + assert!( + matches!( + GatewayConfig::from_yaml(&yaml), + Err(ConfigError::BadStability { .. }) + ), + "stability `{bad}` is rejected at load" + ); + } + let ok = "listen: {host: h, port: 1}\nstability: {availability_window_minutes: 60}"; + assert!(GatewayConfig::from_yaml(ok).is_ok()); + 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/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 17929c2..14883a2 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -150,6 +150,78 @@ impl DagNode for TenantEntitlement { } } +/// preprocess/variant_select: weighted split of a public model name across its +/// configured variant targets, sticky by effective user (per-request spread +/// when anonymous). Runs after entitlement — the public name is what a tenant +/// is entitled to — and before the cache, so each variant caches separately. +/// A request the quota gate already degraded is left alone. +pub struct VariantSelect; + +#[async_trait::async_trait] +impl DagNode for VariantSelect { + fn name(&self) -> &'static str { + "variant_select" + } + fn deps(&self) -> &'static [&'static str] { + &["tenant_entitlement"] + } + async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { + let Some(param) = ctx.request.model_param_v2.as_ref() else { + return Ok(()); + }; + if param.fallback_from.is_some() { + return Ok(()); + } + let Some(conf) = ctx.cfg.find_model(¶m.model_name) else { + return Ok(()); + }; + if conf.variants.is_empty() { + return Ok(()); + } + let key = match ctx.effective_user_id() { + "" => ctx.request.request_id.as_str(), + user => user, + }; + let target = pick_variant(&conf.variants, key).model.clone(); + if target == param.model_name { + ctx.decide("variant_select", format!("{target} (self)")); + return Ok(()); + } + let decision = format!("{} -> {target}", param.model_name); + if let Some(param) = ctx.request.model_param_v2.as_mut() { + param.fallback_from = Some(std::mem::replace(&mut param.model_name, target)); + } + ctx.decide("variant_select", decision); + Ok(()) + } +} + +/// 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; @@ -160,7 +232,7 @@ impl DagNode for CacheLookup { "cache_lookup" } fn deps(&self) -> &'static [&'static str] { - &["tenant_entitlement"] + &["variant_select"] } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { let param = match ctx.request.model_param_v2.as_ref() { @@ -278,14 +350,19 @@ impl DagNode for SelectAccount { .state .pool .select_healthy(mt, provider, &[], ctx.state.health.as_ref()) - .await - .ok_or_else(|| { - GatewayError::new( - ErrCode::SYSTEM_ERROR, - 503, - format!("no healthy upstream account serves model type `{mt}`"), - ) - })?; + .await; + let Some(account) = account else { + // an exhausted pool is a client-visible model failure; unsampled, + // a sustained outage would sit below min_samples as no_data forever + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), false); + return Err(GatewayError::new( + ErrCode::SYSTEM_ERROR, + 503, + format!("no healthy upstream account serves model type `{mt}`"), + )); + }; ctx.decide("select_account", account.name.clone()); ctx.request.account = Some(account); Ok(()) @@ -417,6 +494,10 @@ impl DagNode for UserBudgetGate { /// model_access/call_engine: factory dispatch + engine execution + failover. /// On an upstream 5xx the failed account is excluded and reselected once (a /// PTU → paygo spill sets `ptu_spillover`); a second failure propagates as-is. +/// Availability samples record the client-visible terminal outcome only — an +/// attempt recovered by failover is the account health layer's business, not +/// a model error — and attribute to the requested public name, so a variant +/// split doesn't scatter a model's health across backend names. pub struct CallEngine; #[async_trait::async_trait] @@ -434,10 +515,13 @@ impl DagNode for CallEngine { match engine.run().await { Ok(outcome) => { // an aborted stream is neither a success nor an account fault - if !outcome.response.aborted - && let Some(a) = ctx.request.account.as_ref() - { - ctx.state.health.record_success(&a.name).await; + if !outcome.response.aborted { + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), true); + if let Some(a) = ctx.request.account.as_ref() { + ctx.state.health.record_success(&a.name).await; + } } ctx.decide( "call_engine", @@ -486,6 +570,9 @@ impl DagNode for CallEngine { ) .await; let Some(next) = next else { + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), false); return Err(first_err); }; let spillover = failed.is_ptu() && !next.is_ptu(); @@ -500,12 +587,21 @@ impl DagNode for CallEngine { let retry = gw_engines::get_engine(ctx.request.clone(), ctx.transport.clone())?; match retry.run().await { Ok(mut outcome) => { - ctx.state.health.record_success(&next.name).await; + // same aborted-stream exclusion as the first attempt + if !outcome.response.aborted { + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), true); + ctx.state.health.record_success(&next.name).await; + } outcome.response.ptu_spillover = spillover; ctx.outcome = Some(outcome); Ok(()) } Err(e) => { + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), false); ctx.state .health .record_failure(&next.name, threshold, cooldown) @@ -572,12 +668,18 @@ impl DagNode for CostCalc { gw_models::estimate_prompt_tokens(&ctx.request.message, tools, model_name, enc) }; let ct = enc.encode_len(&resp.message) as i64; + let rate = gw_state::model_token_rate( + &ctx.cfg, + served_model(ctx.request.model_param_v2.as_ref()), + ); + let tokens = BillTokens::weighted(pt, ct, &rate); ctx.decide("cost_calc", format!("aborted stream, billed {pt}+{ct}")); - return bill(ctx, pt, ct, pt.saturating_add(ct), true).await; + return bill(ctx, tokens, true).await; } - // default rate is 1:1; the formula carries future weighted rates. + let rate = + gw_state::model_token_rate(&ctx.cfg, served_model(ctx.request.model_param_v2.as_ref())); // saturating sums so a malformed usage subtree can't overflow the totals - let (prompt, completion, total) = match &resp.common_usage { + let tokens = match &resp.common_usage { Some(u) => { let ti = gw_models::TokenInput { prompt: u.platform_input, @@ -586,33 +688,54 @@ impl DagNode for CostCalc { completion: u.completion, reasoning: u.reason, }; - let rate = gw_models::TokenRate::default(); - ( - u.prompt_total(), - u.completion_total(), - gw_models::platform_total(&ti, &rate), - ) + BillTokens { + prompt: u.prompt_total(), + completion: u.completion_total(), + billable_prompt: gw_models::weighted_prompt(&ti, &rate), + billable_completion: gw_models::weighted_completion(&ti, &rate), + } } - None => ( - resp.prompt_tokens, - resp.completion_tokens, - resp.prompt_tokens.saturating_add(resp.completion_tokens), - ), + None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), }; - bill(ctx, prompt, completion, total, false).await + bill(ctx, tokens, false).await + } +} + +/// Token counts for one bill: vendor-reported sides plus the weighted billable +/// sides; the platform total is always the billable sum, so quota consumption +/// and cost cannot drift. +struct BillTokens { + prompt: i64, + completion: i64, + billable_prompt: i64, + billable_completion: i64, +} + +impl BillTokens { + /// Prompt/completion-only counts with the model's weights applied — the + /// paths without a usage payload (estimates, malformed usage) must price + /// identically to the happy path or a cut stream changes effective pricing. + fn weighted(prompt: i64, completion: i64, rate: &gw_models::TokenRate) -> Self { + let (billable_prompt, billable_completion) = + gw_models::weighted_pair(prompt, completion, rate); + Self { + prompt, + completion, + billable_prompt, + billable_completion, + } + } + + fn total(&self) -> i64 { + self.billable_prompt + .saturating_add(self.billable_completion) } } /// Settle reserves and write the ledger for one served request via the shared /// [`admission::settle_and_bill`] orchestration. `estimated` marks a bill from /// an aborted stream's estimated counts rather than a vendor usage payload. -async fn bill( - ctx: &mut DagContext, - prompt: i64, - completion: i64, - total: i64, - estimated: bool, -) -> GResult<()> { +async fn bill(ctx: &mut DagContext, tokens: BillTokens, estimated: bool) -> GResult<()> { let ptu_spillover = ctx .outcome .as_ref() @@ -621,10 +744,8 @@ async fn bill( let param = ctx.request.model_param_v2.as_ref(); // cost bills at the served (post-fallback) model's price; the (AK, model) // counter accrues against the requested name - let served = param.map(|p| p.model_name.as_str()).unwrap_or_default(); - let requested = param - .and_then(|p| p.fallback_from.as_deref()) - .unwrap_or(served); + let served = served_model(param); + let requested = requested_model(param); // locals first: the whole-ctx borrow `effective_user_id` needs can't overlap these takes let (reserved, tpm_reserved, model_quota_key) = ( ctx.quota_reserved.take().unwrap_or(0), @@ -646,9 +767,11 @@ async fn bill( served_model: served, protocol: param.map(|p| p.protocol.as_str()).unwrap_or_default(), account: ctx.request.account_name(), - prompt, - completion, - total, + prompt: tokens.prompt, + completion: tokens.completion, + billable_prompt: tokens.billable_prompt, + billable_completion: tokens.billable_completion, + total: tokens.total(), ptu_spillover, estimated, }, @@ -733,6 +856,7 @@ pub fn default_layers() -> Vec { Box::new(ModelQuotaGate), Box::new(ResolveModel), Box::new(TenantEntitlement), + Box::new(VariantSelect), Box::new(CacheLookup), Box::new(QuotaCheck), ], @@ -770,8 +894,60 @@ fn model_provider(ctx: &DagContext) -> Option<&str> { ctx.cfg.find_model(name).and_then(|m| m.provider.as_deref()) } +/// The model name the request currently targets (post-fallback). Takes the +/// param, not the whole context, so `bill` can call it across its ctx takes. +fn served_model(param: Option<&gw_models::ModelParamV2>) -> &str { + param.map(|p| p.model_name.as_str()).unwrap_or_default() +} + +/// The public name the caller requested (pre-fallback/variant) — availability +/// attributes here: callers and the tenant status view know this name, not the +/// backend the split routed to. +fn requested_model(param: Option<&gw_models::ModelParamV2>) -> &str { + param + .map(|p| p.fallback_from.as_deref().unwrap_or(p.model_name.as_str())) + .unwrap_or_default() +} + #[cfg(test)] mod tests { + #[test] + fn bill_tokens_weighted_prices_estimate_paths() { + let rate = gw_models::TokenRate { + prompt_weight: 0.5, + ..Default::default() + }; + let t = super::BillTokens::weighted(100, 50, &rate); + assert_eq!((t.prompt, t.completion), (100, 50)); + assert_eq!((t.billable_prompt, t.billable_completion), (50, 50)); + 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/handler/src/lib.rs b/crates/handler/src/lib.rs index e856d60..49801ca 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -461,6 +461,103 @@ mod tests { } } + /// MockTransport with the usage rewritten to 100 prompt tokens, 80 cached. + #[derive(Debug)] + struct CachedUsageTransport; + + #[async_trait::async_trait] + impl gw_engines::Transport for CachedUsageTransport { + async fn send( + &self, + req: gw_engines::UpstreamRequest, + ) -> GResult { + let mut resp = gw_engines::MockTransport.send(req).await?; + if let gw_engines::UpstreamBody::Json(b) = &mut resp.body { + let mut v: serde_json::Value = serde_json::from_slice(b).unwrap(); + v["usage"]["prompt_tokens"] = 100.into(); + v["usage"]["prompt_tokens_details"] = serde_json::json!({"cached_tokens": 80}); + *b = serde_json::to_vec(&v).unwrap(); + } + Ok(resp) + } + } + + #[tokio::test] + async fn token_rate_discounts_cached_prompt_cost() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: m-cache, protocol: openai-chat, input_price_per_1k_micros: 1000, token_rate: {read_cache: 0.1}}]\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(CachedUsageTransport), + ); + let key = h.state().auth.authenticate("k1").await.unwrap(); + h.run(chat_req("m-cache", "hi"), key).await.unwrap(); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + let rec = &ledger[0]; + assert_eq!(rec.prompt_tokens, 100, "raw column: 20 fresh + 80 cached"); + assert_eq!( + rec.cost_micros, 28, + "20 + 80*0.1 billable at 1000 micros/1k" + ); + assert_eq!(rec.total_tokens, 28 + rec.completion_tokens); + } + + #[tokio::test] + async fn failover_recovery_records_one_availability_success() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: m, protocol: openai-chat, provider: p}]\naccounts: [{name: a-down, provider: p, priority: 1, protocols: ['openai-chat']}, {name: a-up, provider: p, priority: 2, protocols: ['openai-chat']}]\nstability: {failure_threshold: 100}\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), + ); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let ctx = h.run(chat_req("m", "hi"), key).await.unwrap(); + assert!( + ctx.decisions + .iter() + .any(|(_, w)| w.contains("failover a-down -> a-up")), + "request must have gone through failover: {:?}", + ctx.decisions + ); + let avail = &h.state().avail; + avail.flush().await; + let minute = gw_state::epoch_secs() / 60; + assert_eq!(avail.window("m", minute - 5, minute).await, (1, 0)); + } + + #[tokio::test] + async fn variant_split_bills_requested_serves_target() { + // tenant entitled only to the public name: entitlement precedes the swap + 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}]\naccounts: [{name: a1, provider: openai, protocols: ['openai-chat']}]\ntenants: [{name: t1, models: [pub-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), + ); + let key = h.state().auth.authenticate("k1").await.unwrap(); + let ctx = h.run(chat_req("pub-m", "hi"), key).await.unwrap(); + assert_eq!( + ctx.outcome.expect("outcome").response.model, + "pub-m", + "response echoes the requested public 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, "canary-m"); + let avail = &h.state().avail; + avail.flush().await; + let minute = gw_state::epoch_secs() / 60; + assert_eq!( + avail.window("pub-m", minute - 5, minute).await, + (1, 0), + "availability attributes to the requested public name" + ); + assert_eq!(avail.window("canary-m", minute - 5, minute).await, (0, 0)); + } + #[derive(Debug)] struct DenyModerator; diff --git a/crates/models/src/cost.rs b/crates/models/src/cost.rs index 2fc87ad..8f9402f 100644 --- a/crates/models/src/cost.rs +++ b/crates/models/src/cost.rs @@ -43,15 +43,6 @@ pub struct TokenInput { pub reasoning: i64, } -/// Cache-normalized prompt (clamped at 0). -fn normalize_prompt(input: &TokenInput, rate: &TokenRate) -> i64 { - let mut prompt = input.prompt; - if rate.prompt_includes_cache { - prompt -= input.read_cache + input.write_cache; - } - prompt.max(0) -} - /// Cost in micro-dollars for one call at per-1k-token prices. Saturating, so a /// malformed/hostile token count can't overflow the multiply into a wrong bill. pub fn cost_micros(prompt: i64, completion: i64, price_per_1k: (i64, i64)) -> i64 { @@ -59,15 +50,47 @@ pub fn cost_micros(prompt: i64, completion: i64, price_per_1k: (i64, i64)) -> i6 .saturating_add(completion.saturating_mul(price_per_1k.1) / 1000) } -/// Weighted platform-total token count. -pub fn platform_total(input: &TokenInput, rate: &TokenRate) -> i64 { - let prompt = normalize_prompt(input, rate) as f64; - let total = prompt * rate.prompt_weight +/// Weighted input-side tokens: prompt plus cache reads/writes. +pub fn weighted_prompt(input: &TokenInput, rate: &TokenRate) -> i64 { + let sum = normalize_prompt(input, rate) as f64 * rate.prompt_weight + input.read_cache as f64 * rate.read_cache_weight - + input.write_cache as f64 * rate.write_cache_weight - + input.completion as f64 * rate.completion_weight + + input.write_cache as f64 * rate.write_cache_weight; + round_tokens(sum) +} + +/// Weighted output-side tokens: completion plus reasoning. +pub fn weighted_completion(input: &TokenInput, rate: &TokenRate) -> i64 { + let sum = input.completion as f64 * rate.completion_weight + input.reasoning as f64 * rate.reasoning_weight; - if total < 0.0 { 0 } else { total.round() as i64 } + round_tokens(sum) +} + +/// Weighted (prompt, completion) for the paths carrying no cache/reasoning +/// components (estimates, realtime turns) — one implementation, so no billing +/// surface can skip the rate. +pub fn weighted_pair(prompt: i64, completion: i64, rate: &TokenRate) -> (i64, i64) { + let input = TokenInput { + prompt, + completion, + ..Default::default() + }; + ( + weighted_prompt(&input, rate), + weighted_completion(&input, rate), + ) +} + +/// Cache-normalized prompt (clamped at 0). +fn normalize_prompt(input: &TokenInput, rate: &TokenRate) -> i64 { + let mut prompt = input.prompt; + if rate.prompt_includes_cache { + prompt -= input.read_cache + input.write_cache; + } + prompt.max(0) +} + +fn round_tokens(sum: f64) -> i64 { + if sum < 0.0 { 0 } else { sum.round() as i64 } } #[cfg(test)] @@ -86,7 +109,9 @@ mod tests { #[test] fn default_rate_is_plain_sum() { - assert_eq!(platform_total(&sample(), &TokenRate::default()), 20); + let rate = TokenRate::default(); + assert_eq!(weighted_prompt(&sample(), &rate), 13); + assert_eq!(weighted_completion(&sample(), &rate), 7); } #[test] @@ -95,7 +120,7 @@ mod tests { prompt_includes_cache: true, ..Default::default() }; - assert_eq!(platform_total(&sample(), &rate), 17); + assert_eq!(weighted_prompt(&sample(), &rate), 10); } #[test] @@ -105,7 +130,35 @@ mod tests { completion_weight: 0.5, ..Default::default() }; - assert_eq!(platform_total(&sample(), &rate), 23); + assert_eq!(weighted_prompt(&sample(), &rate), 18); + assert_eq!(weighted_completion(&sample(), &rate), 5); + } + + #[test] + fn cache_discount_weights_bill_each_side() { + let rate = TokenRate { + read_cache_weight: 0.1, + write_cache_weight: 1.25, + ..Default::default() + }; + let input = TokenInput { + prompt: 100, + read_cache: 1000, + write_cache: 40, + completion: 50, + reasoning: 10, + }; + assert_eq!(weighted_prompt(&input, &rate), 250); + assert_eq!(weighted_completion(&input, &rate), 60); + } + + #[test] + fn weighted_pair_carries_flat_counts() { + let rate = TokenRate { + prompt_weight: 0.5, + ..Default::default() + }; + assert_eq!(weighted_pair(100, 50, &rate), (50, 50)); } #[test] @@ -120,6 +173,6 @@ mod tests { write_cache: 5, ..Default::default() }; - assert_eq!(platform_total(&input, &rate), 10); + assert_eq!(weighted_prompt(&input, &rate), 10); } } diff --git a/crates/models/src/lib.rs b/crates/models/src/lib.rs index 208e61f..eee0e0e 100644 --- a/crates/models/src/lib.rs +++ b/crates/models/src/lib.rs @@ -13,7 +13,9 @@ pub mod token_estimate; pub mod usage; pub use block::Block; -pub use cost::{TokenInput, TokenRate, cost_micros, platform_total}; +pub use cost::{ + TokenInput, TokenRate, cost_micros, weighted_completion, weighted_pair, weighted_prompt, +}; pub use error::{GResult, GatewayError}; pub use params::{ ChatParams, EmbeddingParams, ImageParams, SearchParams, SttParams, TtsParams, TypedParams, diff --git a/crates/models/src/request.rs b/crates/models/src/request.rs index 08ea4b1..b3b6411 100644 --- a/crates/models/src/request.rs +++ b/crates/models/src/request.rs @@ -62,8 +62,8 @@ pub mod domain { pub protocol: gw_consts::Protocol, /// public model name from the caller; empty if a wire type was sent directly. pub model_name: String, - /// original caller model when a quota fallback swapped `model_name`; - /// the response echoes it and the ledger records both. + /// original caller model when a quota fallback or variant split swapped + /// `model_name`; the response echoes it and the ledger records both. pub fallback_from: Option, pub typed: Option, /// untyped vendor extras, passed through verbatim. diff --git a/crates/server/src/main.rs b/crates/server/src/main.rs index 1de19f6..4073d29 100644 --- a/crates/server/src/main.rs +++ b/crates/server/src/main.rs @@ -91,6 +91,7 @@ async fn main() -> anyhow::Result<()> { let quota_task = gw_task::spawn_quota_reset(state.clone(), gw_task::DAILY); let purge_task = gw_task::spawn_content_purge(state.clone(), gw_task::PURGE_PERIOD); let rollup_task = gw_task::spawn_usage_rollup(state.clone(), gw_task::ROLLUP_PERIOD); + let avail_task = gw_task::spawn_avail_flush(state.clone(), gw_task::AVAIL_FLUSH_PERIOD); let distributed_batches = state.store.distributed_batches(); let transport = select_transport()?; @@ -204,6 +205,7 @@ async fn main() -> anyhow::Result<()> { quota_task.abort(); purge_task.abort(); rollup_task.abort(); + avail_task.abort(); tracing::info!("gw drained and exiting"); Ok(()) } diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index 5ee9f99..ce7c864 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -1156,6 +1156,8 @@ async fn chat_non_stream_full_pipeline_bills_the_ledger() { assert_eq!(rec["ak"], "ak-demo-123"); assert_eq!(rec["model"], "gpt-4o"); assert_eq!(rec["account"], "mock-openai-1"); + // equal only while gpt-4o sets no token_rate: ledger total is weighted, + // wire usage stays vendor-raw — they diverge by design under weights assert_eq!(rec["total_tokens"].as_i64().unwrap(), total); assert!(rec["cost_micros"].as_i64().unwrap() > 0); } diff --git a/crates/state/src/avail.rs b/crates/state/src/avail.rs new file mode 100644 index 0000000..74b678c --- /dev/null +++ b/crates/state/src/avail.rs @@ -0,0 +1,249 @@ +//! Per-model availability from minute-bucketed success/error counts. The +//! claim path never touches the network: the in-process store writes its +//! ring directly, and the Redis store buffers samples until the background +//! flush drains them. + +use std::collections::VecDeque; + +use async_trait::async_trait; +use dashmap::DashMap; + +/// Buckets older than this fall off the memory ring / expire in Redis. History +/// beyond the classification window belongs to metrics, not this store. +const RETAIN_MINUTES: i64 = 60; + +/// Availability verdict over the recent window. +#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize)] +#[serde(rename_all = "snake_case")] +pub enum AvailState { + Available, + Unstable, + Unavailable, + /// Too little traffic in the window to judge. + NoData, +} + +/// Classify a window's counts against the configured error-rate thresholds. +pub fn classify( + ok: u64, + err: u64, + min_samples: u64, + unstable: f64, + unavailable: f64, +) -> AvailState { + let total = ok + err; + if total < min_samples.max(1) { + return AvailState::NoData; + } + let rate = err as f64 / total as f64; + if rate >= unavailable { + AvailState::Unavailable + } else if rate >= unstable { + AvailState::Unstable + } else { + AvailState::Available + } +} + +/// Minute-bucket availability store. `record` runs on the claim path and must +/// stay network-free; the flush task calls `flush` so a buffering backend can +/// drain (sub-period skew across a minute boundary is accepted). +#[async_trait] +pub trait AvailStore: Send + Sync + std::fmt::Debug { + fn record(&self, model: &str, ok: bool); + /// Drain buffered samples; no-op for stores that write directly. + async fn flush(&self) {} + /// Summed (ok, err) over buckets in `[since_minute, until_minute]`. + async fn window(&self, model: &str, since_minute: i64, until_minute: i64) -> (u64, u64); +} + +/// In-process ring for single-node deployments; records write it directly. +#[derive(Debug, Default)] +pub struct MemoryAvail { + rings: DashMap>, +} + +impl MemoryAvail { + fn add(&self, model: &str, minute: i64, ok: u64, err: u64) { + let mut ring = crate::slot_mut(&self.rings, model, VecDeque::new); + match ring.back_mut() { + Some(b) if b.0 == minute => { + b.1 += ok; + b.2 += err; + } + _ => ring.push_back((minute, ok, err)), + } + while ring.front().is_some_and(|b| b.0 <= minute - RETAIN_MINUTES) { + ring.pop_front(); + } + } +} + +#[async_trait] +impl AvailStore for MemoryAvail { + fn record(&self, model: &str, ok: bool) { + let minute = crate::epoch_secs() / 60; + self.add(model, minute, u64::from(ok), u64::from(!ok)); + } + + async fn window(&self, model: &str, since_minute: i64, until_minute: i64) -> (u64, u64) { + self.rings + .get(model) + .map(|ring| { + ring.iter() + .filter(|b| b.0 >= since_minute && b.0 <= until_minute) + .fold((0, 0), |(o, e), b| (o + b.1, e + b.2)) + }) + .unwrap_or((0, 0)) + } +} + +/// Fleet-shared buckets: samples buffer in-process and every instance's flush +/// increments the same keys. A Redis outage drops the increment — +/// availability is advisory, matching the governance fail-open posture. +pub struct RedisAvail { + conn: redis::aio::ConnectionManager, + buffer: DashMap, +} + +impl RedisAvail { + pub async fn connect(url: &str) -> Result { + Ok(Self { + conn: crate::redis_connect(url).await?, + buffer: DashMap::new(), + }) + } +} + +impl std::fmt::Debug for RedisAvail { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("RedisAvail") + } +} + +#[async_trait] +impl AvailStore for RedisAvail { + fn record(&self, model: &str, ok: bool) { + let mut e = crate::slot_mut(&self.buffer, model, || (0, 0)); + if ok { + e.0 += 1; + } else { + e.1 += 1; + } + } + + /// Increments landing mid-drain are either taken now or survive for the + /// next flush. + async fn flush(&self) { + let minute = crate::epoch_secs() / 60; + let models: Vec = self.buffer.iter().map(|e| e.key().clone()).collect(); + for m in models { + let Some((m, (ok, err))) = self.buffer.remove(&m) else { + continue; + }; + let mut conn = self.conn.clone(); + let mut pipe = redis::pipe(); + for (key, n) in [ + (count_key(&m, minute, "ok"), ok), + (count_key(&m, minute, "err"), err), + ] { + if n > 0 { + pipe.cmd("INCRBY").arg(&key).arg(n).ignore(); + pipe.cmd("EXPIRE") + .arg(&key) + .arg(RETAIN_MINUTES * 60) + .ignore(); + } + } + if let Err(e) = pipe.query_async::<()>(&mut conn).await { + tracing::warn!(error = %e, model = %m, "redis avail unavailable; counts dropped"); + } + } + } + + async fn window(&self, model: &str, since_minute: i64, until_minute: i64) -> (u64, u64) { + let keys: Vec = (since_minute..=until_minute) + .flat_map(|m| [count_key(model, m, "ok"), count_key(model, m, "err")]) + .collect(); + let mut conn = self.conn.clone(); + let vals: Vec> = + match redis::cmd("MGET").arg(&keys).query_async(&mut conn).await { + Ok(v) => v, + Err(e) => { + tracing::warn!(error = %e, model, "redis avail unavailable; empty window"); + return (0, 0); + } + }; + vals.chunks_exact(2).fold((0, 0), |(o, e), pair| { + (o + pair[0].unwrap_or(0), e + pair[1].unwrap_or(0)) + }) + } +} + +fn count_key(model: &str, minute: i64, kind: &str) -> String { + format!("gw:avail:{model}:{minute}:{kind}") +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn classify_thresholds() { + assert_eq!(classify(0, 0, 20, 0.1, 0.5), AvailState::NoData); + assert_eq!(classify(19, 0, 20, 0.1, 0.5), AvailState::NoData); + assert_eq!(classify(20, 0, 20, 0.1, 0.5), AvailState::Available); + assert_eq!(classify(18, 2, 20, 0.1, 0.5), AvailState::Unstable); + assert_eq!(classify(10, 10, 20, 0.1, 0.5), AvailState::Unavailable); + assert_eq!(classify(0, 20, 20, 0.1, 0.5), AvailState::Unavailable); + assert_eq!(classify(0, 0, 0, 0.1, 0.5), AvailState::NoData); + } + + #[tokio::test] + async fn memory_records_and_windows() { + let store = MemoryAvail::default(); + for _ in 0..3 { + store.record("m", true); + } + store.record("m", false); + store.flush().await; + let minute = crate::epoch_secs() / 60; + assert_eq!(store.window("m", minute - 5, minute).await, (3, 1)); + } + + #[tokio::test] + async fn memory_ring_prunes_and_bounds_window() { + let store = MemoryAvail::default(); + store.add("m", 100, 5, 1); + store.add("m", 101, 2, 0); + assert_eq!(store.window("m", 100, 101).await, (7, 1)); + assert_eq!(store.window("m", 101, 101).await, (2, 0)); + store.add("m", 100 + RETAIN_MINUTES, 1, 0); + assert_eq!( + store.window("m", 100, 100).await, + (0, 0), + "old bucket pruned" + ); + assert_eq!(store.window("m", 0, i64::MAX).await, (3, 0)); + } + + #[tokio::test] + async fn redis_avail_buffers_flushes_and_windows() { + let Ok(url) = std::env::var("GW_TEST_REDIS_URL") else { + return; + }; + let store = RedisAvail::connect(&url).await.expect("redis connect"); + let model = format!("m-{}", std::process::id()); + store.record(&model, true); + store.record(&model, true); + store.record(&model, false); + let minute = crate::epoch_secs() / 60; + assert_eq!( + store.window(&model, minute - 1, minute + 1).await, + (0, 0), + "buffered until flush" + ); + store.flush().await; + assert_eq!(store.window(&model, minute - 1, minute + 1).await, (2, 1)); + } +} diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index b3798f7..7455b37 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 avail; pub mod configstore; pub mod content; pub mod governance; @@ -21,6 +22,7 @@ pub mod health; pub mod keystore; pub mod store; +pub use avail::{AvailState, AvailStore, classify}; pub use configstore::{CONFIG_CHANNEL, PostgresConfigStore}; pub use content::{ContentRecord, sealing_available}; pub use governance::{Governance, MemoryGovernance, RedisGovernance}; @@ -663,6 +665,8 @@ pub struct GatewayState { pub health: Arc, /// Request-level response cache: in-process by default, Redis when shared. pub cache: Arc, + /// Per-model availability counters; Redis for fleet-wide aggregation. + pub avail: Arc, } impl Default for GatewayState { @@ -674,6 +678,7 @@ impl Default for GatewayState { store: Arc::new(MemoryStore::default()), health: Arc::new(AccountHealth::default()), cache: Arc::new(MemoryResponseCache::default()), + avail: Arc::new(avail::MemoryAvail::default()), } } } @@ -752,6 +757,15 @@ impl GatewayState { tracing::error!(error = %e, "redis health connect failed; staying in-process") } } + match avail::RedisAvail::connect(&cfg.storage.redis_url).await { + Ok(a) => { + state.avail = Arc::new(a); + tracing::info!("model availability = redis (fleet-wide counts)"); + } + Err(e) => { + tracing::error!(error = %e, "redis avail connect failed; staying in-process") + } + } } Ok(state) } @@ -768,6 +782,7 @@ impl GatewayState { store: prev.store.clone(), health: prev.health.clone(), cache: prev.cache.clone(), + avail: prev.avail.clone(), }) } } diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 4057b7f..d672b60 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -78,6 +78,8 @@ pub struct BillingRecord { pub account: String, pub prompt_tokens: i64, pub completion_tokens: i64, + /// Weighted platform total (what quota metering consumed); equals + /// prompt + completion only when the served model has no `token_rate`. pub total_tokens: i64, pub cost_micros: i64, /// What the serving account's vendor charged us (zero = untracked). @@ -182,6 +184,10 @@ pub struct BillingInput<'a> { pub account: &'a str, pub prompt: i64, pub completion: i64, + /// Weighted billable sides (the served model's token_rate applied); equal + /// to `prompt`/`completion` when the model has no rate configured. + pub billable_prompt: i64, + pub billable_completion: i64, pub total: i64, pub ptu_spillover: bool, /// Counts are estimated (aborted stream), not vendor-reported. @@ -193,15 +199,38 @@ pub fn clamp_tokens(n: i64) -> i64 { n.clamp(0, MAX_METERED_TOKENS) } +/// The served model's billing weights; identity when unconfigured. The +/// `prompt_includes_cache` normalization already happened at usage extraction. +pub fn model_token_rate(cfg: &gw_config::GatewayConfig, model: &str) -> gw_models::TokenRate { + match cfg.find_model(model).and_then(|m| m.token_rate) { + Some(r) => gw_models::TokenRate { + prompt_includes_cache: false, + prompt_weight: r.prompt, + read_cache_weight: r.read_cache, + write_cache_weight: r.write_cache, + completion_weight: r.completion, + reasoning_weight: r.reasoning, + }, + None => gw_models::TokenRate::default(), + } +} + /// Price one call into a [`BillingRecord`]: the tenant's price for the served /// model, vendor cost from the serving account. Shared by the request pipeline /// and the realtime surface so the two can't drift; token counts are clamped. +/// Cost multiplies the weighted billable sides; the prompt/completion columns +/// keep the vendor-reported counts, while `total_tokens` is the weighted +/// platform total that quota metering consumed. pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> BillingRecord { let (prompt, completion, total) = ( clamp_tokens(b.prompt), clamp_tokens(b.completion), clamp_tokens(b.total), ); + let (billable_prompt, billable_completion) = ( + clamp_tokens(b.billable_prompt), + clamp_tokens(b.billable_completion), + ); let charged = cfg.prices_for_tenant(b.tenant, b.served_model); let vendor = cfg .accounts @@ -228,8 +257,8 @@ pub fn billing_record(cfg: &gw_config::GatewayConfig, b: &BillingInput) -> Billi prompt_tokens: prompt, completion_tokens: completion, total_tokens: total, - cost_micros: gw_models::cost_micros(prompt, completion, charged), - vendor_cost_micros: gw_models::cost_micros(prompt, completion, vendor), + cost_micros: gw_models::cost_micros(billable_prompt, billable_completion, charged), + vendor_cost_micros: gw_models::cost_micros(billable_prompt, billable_completion, vendor), ptu_spillover: b.ptu_spillover, estimated: b.estimated, } @@ -2371,6 +2400,53 @@ mod tests { assert_eq!(pg_numbered("no placeholders"), "no placeholders"); } + #[test] + fn model_token_rate_maps_config_weights() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: m, protocol: openai-chat, token_rate: {read_cache: 0.1, write_cache: 1.25}}]"; + let cfg = gw_config::GatewayConfig::from_yaml(yaml).unwrap(); + let rate = model_token_rate(&cfg, "m"); + assert_eq!( + (rate.read_cache_weight, rate.write_cache_weight), + (0.1, 1.25) + ); + assert!(!rate.prompt_includes_cache); + assert_eq!( + model_token_rate(&cfg, "absent"), + gw_models::TokenRate::default() + ); + } + + #[test] + fn billing_record_bills_weighted_sides_keeps_raw_columns() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: m, protocol: openai-chat, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 2000}]"; + let cfg = gw_config::GatewayConfig::from_yaml(yaml).unwrap(); + let rec = billing_record( + &cfg, + &BillingInput { + ak: "k", + product: "demo", + tenant: "default", + user_id: "u1", + request_id: "req-1", + requested_model: "m", + served_model: "m", + protocol: "openai-chat", + account: "acc", + prompt: 1140, + completion: 60, + billable_prompt: 250, + billable_completion: 60, + total: 310, + ptu_spillover: false, + estimated: false, + }, + ); + assert_eq!(rec.prompt_tokens, 1140); + assert_eq!(rec.completion_tokens, 60); + assert_eq!(rec.total_tokens, 310); + assert_eq!(rec.cost_micros, 250 + 120); + } + #[test] fn billing_record_clamps_hostile_usage() { let cfg = gw_config::GatewayConfig::embedded_default().unwrap(); @@ -2388,6 +2464,8 @@ mod tests { account: "acc", prompt: i64::MAX, completion: i64::MAX, + billable_prompt: i64::MAX, + billable_completion: i64::MAX, total: i64::MAX, ptu_spillover: false, estimated: false, @@ -2819,8 +2897,12 @@ mod tests { async fn exercise_erasure_markers(store: &dyn Store, ns: &str) { let (t1, u1, u2) = (format!("t1{ns}"), format!("u1{ns}"), format!("u2{ns}")); let now = crate::epoch_millis(); + // wide margin: a 5ms one flaked under parallel-suite load assert!( - store.user_erased_since(&t1, &u1, now - 5).await.unwrap(), + store + .user_erased_since(&t1, &u1, now - 60_000) + .await + .unwrap(), "marker visible to a batch started before the erasure" ); assert!( diff --git a/crates/task/src/lib.rs b/crates/task/src/lib.rs index a4085b0..13d8dc8 100644 --- a/crates/task/src/lib.rs +++ b/crates/task/src/lib.rs @@ -13,6 +13,8 @@ pub const DAILY: Duration = Duration::from_secs(24 * 60 * 60); pub const PURGE_PERIOD: Duration = Duration::from_secs(60 * 60); /// Ledger minutes are rolled into the durable usage buckets this often. 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); /// Spawn the daily quota reset loop. Returns the join handle (abort to stop). /// `period` is configurable so tests don't wait 24h. @@ -78,6 +80,23 @@ pub fn spawn_usage_rollup( }) } +/// Spawn the availability flush loop: drains the in-process per-model counters +/// into the minute-bucket store, keeping the claim path free of network hops. +pub fn spawn_avail_flush( + state: Arc, + 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; + loop { + tick.tick().await; + state.avail.flush().await; + } + }) +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 8c260be..d4aef21 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -124,6 +124,7 @@ pub fn app(state: AppState) -> Router { .route("/admin/keys", post(admin_key_create).get(admin_key_list)) .route("/admin/usage", get(admin_usage)) .route("/admin/usage/users", get(admin_usage_users)) + .route("/admin/models/status", get(admin_models_status)) .route("/admin/audit/events", get(admin_security_events)) .route("/admin/audit/ops", get(admin_audit_ops)) .route("/admin/audit/content/{request_id}", get(admin_content_get)) @@ -358,12 +359,15 @@ async fn bill_realtime_turn( let ak = &admit.ak; // clamp parts and total so a hostile count can't overflow shared counters let (it, ot) = (gw_state::clamp_tokens(it), gw_state::clamp_tokens(ot)); - let total = gw_state::clamp_tokens(it.saturating_add(ot)); - if total == 0 { + if it.saturating_add(ot) == 0 { admit.refund().await; return; } 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 (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)); admission::settle_and_bill( @@ -383,6 +387,8 @@ async fn bill_realtime_turn( account, prompt: it, completion: ot, + billable_prompt: bp, + billable_completion: bc, total, ptu_spillover: false, estimated: false, @@ -1524,6 +1530,50 @@ async fn admin_key_list( Json(resp).into_response() } +/// 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. +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 avail = &state.avail; + let counts = futures::future::join_all( + entitled.map(|m| async move { (&m.name, avail.window(&m.name, since, until).await) }), + ) + .await; + let rows: Vec = counts + .into_iter() + .map(|(name, (ok, err))| { + let verdict = gw_state::classify( + ok, + err, + st.availability_min_samples, + st.unstable_error_rate, + st.unavailable_error_rate, + ); + json!({ + "model": name, + "state": verdict, + "requests": ok + err, + "errors": err, + "window_minutes": st.availability_window_minutes, + }) + }) + .collect(); + Json(json!({ "models": rows })).into_response() +} + /// GET /admin/usage — ledger rollup by (tenant, requested model). A tenant /// admin sees only its own tenant; the global admin may filter with ?tenant=. async fn admin_usage( @@ -3339,6 +3389,115 @@ mod tests { assert!(store.file_get(&f.id).await.unwrap().is_none()); } + #[tokio::test] + async fn models_status_classifies_and_scopes() { + // high threshold + disjoint providers: cooldown/round-robin must not skew samples + let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_AVAIL_ADMIN}\nmodels: [{name: m-ok, protocol: openai-chat, provider: openai}, {name: m-bad, protocol: openai-chat, provider: downp}, {name: rt-m, protocol: realtime}]\naccounts: [{name: a-up, provider: openai, protocols: ['openai-chat']}, {name: a-down, provider: downp, protocols: ['openai-chat']}]\ntenants: [{name: t1, models: [m-ok], admin_token_env: GW_TEST_AVAIL_T1}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]\nstability: {availability_min_samples: 3, failure_threshold: 100}"; + // SAFETY: unique var names for this test; no concurrent reader of them. + unsafe { + std::env::set_var("GW_TEST_AVAIL_ADMIN", "root-tok"); + std::env::set_var("GW_TEST_AVAIL_T1", "t1-tok"); + } + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let avail = app_state.handler.state().avail.clone(); + let router = app(app_state); + let chat = |model: &str| { + chat_req( + Some("k1"), + &format!(r#"{{"model":"{model}","messages":[{{"role":"user","content":"x"}}]}}"#), + ) + }; + for _ in 0..4 { + let ok = router.clone().oneshot(chat("m-ok")).await.unwrap(); + assert_eq!(ok.status(), StatusCode::OK); + let bad = router.clone().oneshot(chat("m-bad")).await.unwrap(); + assert_eq!(bad.status(), StatusCode::SERVICE_UNAVAILABLE); + } + avail.flush().await; + + let status = |token: &'static str| { + Request::builder() + .method("GET") + .uri("/admin/models/status") + .header("authorization", format!("Bearer {token}")) + .body(Body::empty()) + .unwrap() + }; + let resp = router.clone().oneshot(status("root-tok")).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let j = body_json(resp).await; + let rows = j["models"].as_array().unwrap(); + let of = |name: &str| { + rows.iter() + .find(|r| r["model"] == name) + .unwrap_or_else(|| panic!("row for {name}")) + .clone() + }; + assert_eq!(of("m-ok")["state"], "available"); + 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 resp = router.clone().oneshot(status("t1-tok")).await.unwrap(); + let j = body_json(resp).await; + let names: Vec<_> = j["models"] + .as_array() + .unwrap() + .iter() + .map(|r| r["model"].as_str().unwrap().to_owned()) + .collect(); + assert_eq!(names, ["m-ok"], "tenant admin sees only entitled models"); + } + + #[tokio::test] + async fn sustained_pool_exhaustion_converges_to_unavailable() { + // past the 2nd request the pool is cooling: those 503s must sample too + let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_OUT_ADMIN}\nmodels: [{name: m-out, protocol: openai-chat, provider: downp}]\naccounts: [{name: a-down, provider: downp, protocols: ['openai-chat']}]\naccess_keys: [{ak: k1, product: p, qps: 100, daily_token_quota: 100000}]\nstability: {failure_threshold: 2, cooldown_seconds: 300, availability_min_samples: 5}"; + // SAFETY: unique var name for this test; no concurrent reader of it. + unsafe { + std::env::set_var("GW_TEST_OUT_ADMIN", "out-tok"); + } + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let app_state = AppState::new(cfg, state, Arc::new(gw_engines::MockTransport)); + let avail = app_state.handler.state().avail.clone(); + let router = app(app_state); + for _ in 0..6 { + let resp = router + .clone() + .oneshot(chat_req( + Some("k1"), + r#"{"model":"m-out","messages":[{"role":"user","content":"x"}]}"#, + )) + .await + .unwrap(); + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + } + avail.flush().await; + let resp = router + .oneshot( + Request::builder() + .method("GET") + .uri("/admin/models/status") + .header("authorization", "Bearer out-tok") + .body(Body::empty()) + .unwrap(), + ) + .await + .unwrap(); + let j = body_json(resp).await; + let row = &j["models"][0]; + assert_eq!(row["model"], "m-out"); + assert_eq!(row["errors"], 6, "cooldown-era 503s sample too"); + assert_eq!(row["state"], "unavailable"); + } + #[tokio::test] async fn content_erase_is_tenant_scoped_and_audited() { let yaml = "listen: {host: h, port: 1}\nadmin: {token_env: GW_TEST_ERASE_ADMIN}\nmodels: [{name: gpt-4o, protocol: openai-chat}]\ntenants: [{name: t1}, {name: t2, admin_token_env: GW_TEST_ERASE_T2}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 10, daily_token_quota: 1000}]"; @@ -3632,6 +3791,34 @@ mod tests { ); } + #[tokio::test] + async fn realtime_billing_applies_token_rate() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: rt-m, protocol: realtime, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 1000, token_rate: {completion: 0.5}}]\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 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 (_, ledger) = s + .handler + .state() + .store + .ledger_snapshot(usize::MAX) + .await + .unwrap(); + let rec = &ledger[0]; + assert_eq!(rec.prompt_tokens, 100); + assert_eq!(rec.completion_tokens, 100); + assert_eq!(rec.total_tokens, 150, "100 + 100*0.5 weighted"); + assert_eq!(rec.cost_micros, 150, "billable 100+50 at 1000/1k each"); + assert_eq!( + s.handler.state().governance.quota_used(&ak.ak).await, + 150, + "quota settles the weighted total" + ); + } + #[tokio::test] async fn realtime_gate_reserves_tpm_and_rolls_back_on_denial() { let cfg = Arc::new(GatewayConfig::embedded_default().unwrap()); diff --git a/docs/api.md b/docs/api.md index 89b6ec3..4e58d38 100644 --- a/docs/api.md +++ b/docs/api.md @@ -141,6 +141,7 @@ regardless. | 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/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 ec616bf..2f4632c 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -104,8 +104,22 @@ models: output_price_per_1k_micros: 10000 qpm: 60 # optional model-level rate limit cache_ttl_seconds: 60 # optional request-level response cache + token_rate: # optional per-component billing weights + read_cache: 0.1 # cache reads at 10% of the input price + write_cache: 1.25 # (prompt/completion/reasoning default 1.0) + variants: # optional weighted canary split, sticky per user + - {model: gpt-4o, weight: 90} # self-reference keeps a share here + - {model: gpt-4o-next, weight: 10} ``` +`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. + ### `providers` — first-class provider presets ```yaml @@ -168,6 +182,10 @@ security: # global default; a tenant may override it whole stability: failure_threshold: 3 # consecutive failures before an account cools down cooldown_seconds: 300 + availability_window_minutes: 5 # /admin/models/status judgment window (max 60) + unstable_error_rate: 0.1 # window error rate that reports `unstable` + unavailable_error_rate: 0.5 # ... and `unavailable` + availability_min_samples: 20 # fewer samples than this reports `no_data` products: - name: myproduct