From eaef0a8029cc0fa9a8afee0c35a6dc2af3f35183 Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 21:35:04 +0800 Subject: [PATCH 01/11] feat: per-model token_rate weights price cache/reasoning tokens Cost and quota consumption multiply the served model's configured component weights (cache reads/writes, reasoning); ledger token columns keep the vendor-reported counts. Weights validated finite and non-negative at load. --- crates/config/src/lib.rs | 71 +++++++++++++++++++++++++ crates/dag/src/nodes.rs | 107 +++++++++++++++++++++++++++++--------- crates/handler/src/lib.rs | 41 +++++++++++++++ crates/models/src/cost.rs | 61 +++++++++++++++++++--- crates/models/src/lib.rs | 4 +- crates/state/src/store.rs | 47 ++++++++++++++++- crates/views/src/lib.rs | 2 + 7 files changed, 298 insertions(+), 35 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 29b4590..4670b76 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -50,6 +50,8 @@ 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("`{owner}` sets a negative or non-finite limit")] NegativeLimit { owner: String }, #[error("storage.shared_cache needs storage.redis_url")] @@ -109,6 +111,9 @@ 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, } impl ModelConf { @@ -117,6 +122,39 @@ impl ModelConf { } } +/// 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. + pub 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 { @@ -613,6 +651,17 @@ impl GatewayConfig { 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 !v.is_finite() || v < 0.0 { + return Err(ConfigError::BadTokenRate { + model: m.name.clone(), + field, + }); + } + } + } } for a in &self.accounts { if neg( @@ -1160,6 +1209,28 @@ 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)); + 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..aa7c218 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -573,11 +573,10 @@ impl DagNode for CostCalc { }; let ct = enc.encode_len(&resp.message) as i64; ctx.decide("cost_calc", format!("aborted stream, billed {pt}+{ct}")); - return bill(ctx, pt, ct, pt.saturating_add(ct), true).await; + return bill(ctx, BillTokens::flat(pt, ct), true).await; } - // default rate is 1:1; the formula carries future weighted rates. // 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 +585,73 @@ 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), - ) + let served = ctx + .request + .model_param_v2 + .as_ref() + .map(|p| p.model_name.as_str()) + .unwrap_or_default(); + let rate = token_rate(&ctx.cfg, served); + 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::flat(resp.prompt_tokens, resp.completion_tokens), }; - bill(ctx, prompt, completion, total, false).await + bill(ctx, tokens, false).await + } +} + +/// The served model's billing weights; identity when unconfigured. The +/// `prompt_includes_cache` normalization already happened at usage extraction. +fn 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(), + } +} + +/// 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 { + /// Unweighted counts: billable equals reported. + fn flat(prompt: i64, completion: i64) -> Self { + Self { + prompt, + completion, + billable_prompt: prompt, + billable_completion: 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() @@ -646,9 +685,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, }, @@ -772,6 +813,22 @@ fn model_provider(ctx: &DagContext) -> Option<&str> { #[cfg(test)] mod tests { + #[test] + fn 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 = super::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!( + super::token_rate(&cfg, "absent"), + gw_models::TokenRate::default() + ); + } + #[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..a081c35 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -461,6 +461,47 @@ 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]; + // raw usage stays truthful: 20 fresh + 80 cached + assert_eq!(rec.prompt_tokens, 100); + // billable prompt = 20 + 80*0.1 = 28 at 1000 micros/1k input + assert_eq!(rec.cost_micros, 28); + assert_eq!(rec.total_tokens, 28 + rec.completion_tokens); + } + #[derive(Debug)] struct DenyModerator; diff --git a/crates/models/src/cost.rs b/crates/models/src/cost.rs index 2fc87ad..5577a28 100644 --- a/crates/models/src/cost.rs +++ b/crates/models/src/cost.rs @@ -59,15 +59,29 @@ 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 platform-total token count; always the sum of the two sides so +/// quota consumption and per-side billing cannot drift. +pub fn platform_total(input: &TokenInput, rate: &TokenRate) -> i64 { + weighted_prompt(input, rate).saturating_add(weighted_completion(input, rate)) +} + +fn round_tokens(sum: f64) -> i64 { + if sum < 0.0 { 0 } else { sum.round() as i64 } } #[cfg(test)] @@ -108,6 +122,39 @@ mod tests { assert_eq!(platform_total(&sample(), &rate), 23); } + #[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); + assert_eq!(platform_total(&input, &rate), 310); + } + + #[test] + fn total_is_sum_of_sides() { + let rate = TokenRate { + prompt_weight: 1.5, + completion_weight: 0.5, + ..Default::default() + }; + let input = sample(); + assert_eq!( + platform_total(&input, &rate), + weighted_prompt(&input, &rate) + weighted_completion(&input, &rate) + ); + } + #[test] fn negative_prompt_clamped_to_zero() { let rate = TokenRate { diff --git a/crates/models/src/lib.rs b/crates/models/src/lib.rs index 208e61f..d07bfa9 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, platform_total, weighted_completion, weighted_prompt, +}; pub use error::{GResult, GatewayError}; pub use params::{ ChatParams, EmbeddingParams, ImageParams, SearchParams, SttParams, TtsParams, TypedParams, diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 4057b7f..f17e76d 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -182,6 +182,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. @@ -196,12 +200,18 @@ pub fn clamp_tokens(n: i64) -> i64 { /// 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 token columns keep the +/// vendor-reported counts so usage stays truthful. 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 +238,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 +2381,37 @@ mod tests { assert_eq!(pg_numbered("no placeholders"), "no placeholders"); } + #[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 +2429,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, diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 8c260be..390a126 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -383,6 +383,8 @@ async fn bill_realtime_turn( account, prompt: it, completion: ot, + billable_prompt: it, + billable_completion: ot, total, ptu_spillover: false, estimated: false, From abd5a849a1f008011eef9cbb286d68ae3692ff0b Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 21:52:08 +0800 Subject: [PATCH 02/11] feat: per-model availability status API MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit CallEngine outcomes (upstream 5xx vs success; 4xx and client aborts say nothing about model health) accumulate in an in-process buffer, flushed to minute buckets every 2s — memory ring locally, shared Redis keys fleet-wide — so the claim path never gains a network hop. GET /admin/models/status classifies the configured window into available/unstable/unavailable/no_data; tenant admins see only entitled models. --- crates/config/src/lib.rs | 47 +++++++ crates/dag/src/nodes.rs | 23 +++- crates/server/src/main.rs | 2 + crates/state/src/avail.rs | 274 ++++++++++++++++++++++++++++++++++++++ crates/state/src/lib.rs | 17 +++ crates/task/src/lib.rs | 19 +++ crates/views/src/lib.rs | 99 ++++++++++++++ 7 files changed, 477 insertions(+), 4 deletions(-) create mode 100644 crates/state/src/avail.rs diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 4670b76..1e17b97 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -52,6 +52,8 @@ pub enum ConfigError { 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("`{owner}` sets a negative or non-finite limit")] NegativeLimit { owner: String }, #[error("storage.shared_cache needs storage.redis_url")] @@ -287,11 +289,35 @@ 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. + #[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 } @@ -301,6 +327,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(), } } } @@ -713,6 +743,23 @@ 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); + if st.availability_window_minutes < 1 { + return Err(ConfigError::BadStability { + field: "availability_window_minutes", + }); + } + 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 { diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index aa7c218..b1b4d39 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -434,10 +434,11 @@ 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(served_model(ctx), true); + if let Some(a) = ctx.request.account.as_ref() { + ctx.state.health.record_success(&a.name).await; + } } ctx.decide( "call_engine", @@ -450,6 +451,7 @@ impl DagNode for CallEngine { Ok(()) } Err(first_err) if first_err.http_status >= 500 => { + ctx.state.avail.record(served_model(ctx), false); let mt = ctx .request .protocol() @@ -500,12 +502,16 @@ 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.avail.record(served_model(ctx), true); ctx.state.health.record_success(&next.name).await; outcome.response.ptu_spillover = spillover; ctx.outcome = Some(outcome); Ok(()) } Err(e) => { + if e.http_status >= 500 { + ctx.state.avail.record(served_model(ctx), false); + } ctx.state .health .record_failure(&next.name, threshold, cooldown) @@ -811,6 +817,15 @@ 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). +fn served_model(ctx: &DagContext) -> &str { + ctx.request + .model_param_v2 + .as_ref() + .map(|p| p.model_name.as_str()) + .unwrap_or_default() +} + #[cfg(test)] mod tests { #[test] 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/state/src/avail.rs b/crates/state/src/avail.rs new file mode 100644 index 0000000..19cc0be --- /dev/null +++ b/crates/state/src/avail.rs @@ -0,0 +1,274 @@ +//! Per-model availability from minute-bucketed success/error counts. The +//! claim path only bumps an in-process buffer; a background task flushes it +//! to the store (Redis when the fleet shares one), so recording never adds a +//! network hop to a request. + +use std::collections::VecDeque; +use std::sync::Arc; + +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 + } +} + +/// Claim-path recorder: counts accumulate in-process and [`Self::flush`] +/// folds them into the store under the flush-time minute (sub-period skew +/// across a minute boundary is accepted). +pub struct AvailTracker { + buffer: DashMap, + store: Arc, +} + +impl AvailTracker { + pub fn new(store: Arc) -> Self { + Self { + buffer: DashMap::new(), + store, + } + } + + /// Bump the model's (ok, err) pair; allocates only on a model's first + /// record since the last flush. + pub fn record(&self, model: &str, ok: bool) { + if let Some(mut e) = self.buffer.get_mut(model) { + bump(e.value_mut(), ok); + return; + } + bump( + self.buffer + .entry(model.to_owned()) + .or_insert((0, 0)) + .value_mut(), + ok, + ); + } + + /// Drain the buffer into the store. Increments landing mid-drain are + /// either taken now or survive for the next flush. + pub 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 { + if let Some((m, (ok, err))) = self.buffer.remove(&m) { + self.store.add(&m, minute, ok, err).await; + } + } + } + + /// Summed (ok, err) over `[since_minute, until_minute]`. + pub async fn window(&self, model: &str, since_minute: i64, until_minute: i64) -> (u64, u64) { + self.store.window(model, since_minute, until_minute).await + } +} + +impl std::fmt::Debug for AvailTracker { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("AvailTracker") + } +} + +fn bump(counts: &mut (u64, u64), ok: bool) { + if ok { + counts.0 += 1; + } else { + counts.1 += 1; + } +} + +/// Minute-bucket store behind the flush task. +#[async_trait] +pub trait AvailStore: Send + Sync + std::fmt::Debug { + /// Fold one flush increment into `minute`'s bucket. + async fn add(&self, model: &str, minute: i64, ok: u64, err: u64); + /// 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. +#[derive(Debug, Default)] +pub struct MemoryAvail { + rings: DashMap>, +} + +#[async_trait] +impl AvailStore for MemoryAvail { + async fn add(&self, model: &str, minute: i64, ok: u64, err: u64) { + let mut ring = self.rings.entry(model.to_owned()).or_default(); + 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 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: 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, +} + +impl RedisAvail { + pub async fn connect(url: &str) -> Result { + Ok(Self { + conn: crate::redis_connect(url).await?, + }) + } +} + +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 { + async fn add(&self, model: &str, minute: i64, ok: u64, err: u64) { + let mut conn = self.conn.clone(); + let mut pipe = redis::pipe(); + for (key, n) in [ + (count_key(model, minute, "ok"), ok), + (count_key(model, 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, "redis avail unavailable; counts dropped"); + } + } + + async fn window(&self, model: &str, since_minute: i64, until_minute: i64) -> (u64, u64) { + let mut keys = Vec::new(); + for minute in since_minute..=until_minute { + keys.push(count_key(model, minute, "ok")); + keys.push(count_key(model, minute, "err")); + } + 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(2).fold((0, 0), |(o, e), pair| { + ( + o + pair.first().copied().flatten().unwrap_or(0), + e + pair.get(1).copied().flatten().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); + // min_samples 0 still needs at least one sample + assert_eq!(classify(0, 0, 0, 0.1, 0.5), AvailState::NoData); + } + + #[tokio::test] + async fn tracker_buffers_flushes_and_windows() { + let tracker = AvailTracker::new(Arc::new(MemoryAvail::default())); + for _ in 0..3 { + tracker.record("m", true); + } + tracker.record("m", false); + let minute = crate::epoch_secs() / 60; + assert_eq!(tracker.window("m", minute - 5, minute).await, (0, 0)); + tracker.flush().await; + assert_eq!(tracker.window("m", minute - 5, minute).await, (3, 1)); + // second flush with an empty buffer adds nothing + tracker.flush().await; + assert_eq!(tracker.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).await; + store.add("m", 101, 2, 0).await; + 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).await; + assert_eq!( + store.window("m", 100, 100).await, + (0, 0), + "old bucket pruned" + ); + assert_eq!(store.window("m", 0, i64::MAX).await, (3, 0)); + } +} diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index b3798f7..380abbc 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, AvailTracker, 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,9 @@ impl Default for GatewayState { store: Arc::new(MemoryStore::default()), health: Arc::new(AccountHealth::default()), cache: Arc::new(MemoryResponseCache::default()), + avail: Arc::new(avail::AvailTracker::new(Arc::new( + avail::MemoryAvail::default(), + ))), } } } @@ -752,6 +759,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(avail::AvailTracker::new(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 +784,7 @@ impl GatewayState { store: prev.store.clone(), health: prev.health.clone(), cache: prev.cache.clone(), + avail: prev.avail.clone(), }) } } 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 390a126..d06411b 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)) @@ -1526,6 +1527,41 @@ 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. +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 avail = &s.handler.state().avail; + let mut rows = Vec::new(); + for m in &cfg.models { + if let AdminScope::Tenant(t) = &scope + && !cfg.tenant_allows_model(t, &m.name) + { + continue; + } + let (ok, err) = avail.window(&m.name, since, until).await; + let state = gw_state::classify( + ok, + err, + st.availability_min_samples, + st.unstable_error_rate, + st.unavailable_error_rate, + ); + rows.push(json!({ + "model": m.name, + "state": state, + "requests": ok + err, + "errors": err, + "window_minutes": st.availability_window_minutes, + })); + } + 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( @@ -3341,6 +3377,69 @@ mod tests { assert!(store.file_get(&f.id).await.unwrap().is_none()); } + #[tokio::test] + async fn models_status_classifies_and_scopes() { + // failure_threshold high so account cooldown never blocks a sample; + // models bound to disjoint providers so round-robin can't cross-route + 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}]\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); + + 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 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}]"; From 0c17e75dbc49f9f92c0c16e0709859b707ce6f70 Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 22:01:32 +0800 Subject: [PATCH 03/11] feat: weighted variant routing with sticky user hashing A public model name may declare a weighted split across other declared models (same protocol, one level, validated at load). Selection hashes the effective user (request id when anonymous) with FNV-1a so every fleet instance maps a user to the same variant with no shared state. The swap reuses the fallback seam: entitlement and the (AK, model) counter judge the public name, billing prices the served variant, the response echoes the requested name, and each variant caches separately. --- crates/config/src/lib.rs | 59 +++++++++++++++++++++ crates/dag/src/nodes.rs | 104 +++++++++++++++++++++++++++++++++++++- crates/handler/src/lib.rs | 23 +++++++++ 3 files changed, 185 insertions(+), 1 deletion(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 1e17b97..d5c77a3 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -54,6 +54,12 @@ pub enum ConfigError { 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")] @@ -116,6 +122,11 @@ pub struct ModelConf { /// 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. + #[serde(default)] + pub variants: Vec, } impl ModelConf { @@ -124,6 +135,13 @@ 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)] @@ -692,6 +710,26 @@ impl GatewayConfig { } } } + for v in &m.variants { + let bad = |reason| ConfigError::BadVariant { + model: m.name.clone(), + variant: v.model.clone(), + reason, + }; + if v.weight == 0 { + return Err(bad("weight must be >= 1")); + } + 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( @@ -1278,6 +1316,27 @@ tenants: [{name: t1}, {name: t1}] 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"), + ] { + 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 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 b1b4d39..a048665 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -150,6 +150,82 @@ 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); + } + // weights are validated >= 1, so the loop always returns; this arm only + // 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 { + let mut h: u64 = 0xcbf2_9ce4_8422_2325; + for b in s.as_bytes() { + h ^= u64::from(*b); + h = h.wrapping_mul(0x0000_0100_0000_01b3); + } + h +} + /// 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 +236,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() { @@ -780,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), ], @@ -844,6 +921,31 @@ mod tests { ); } + #[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 a081c35..06b708a 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -502,6 +502,29 @@ mod tests { assert_eq!(rec.total_tokens, 28 + rec.completion_tokens); } + #[tokio::test] + async fn variant_split_bills_requested_serves_target() { + // tenant entitled only to the public name — proves entitlement is + // judged before the variant 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"); + } + #[derive(Debug)] struct DenyModerator; From 0f2a88e1b5e65900b1b8773010024a0e3c7e9290 Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 22:32:06 +0800 Subject: [PATCH 04/11] =?UTF-8?q?review:=20converge=20the=20round=20?= =?UTF-8?q?=E2=80=94=20shared=20token=5Frate=20on=20every=20billing=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Round findings applied: realtime turns and the estimate/no-usage branches now weight through the same model_token_rate helper as the pipeline (a cut stream or WS turn must price like a completed one); availability records the client-visible terminal outcome instead of per-attempt (failover recovery is account-health business); availability_window_minutes capped at the store's one-hour retention; duplicate variant targets rejected; status endpoint fans out with join_all; small cleanups from the simplify lenses. total_tokens documented as the weighted platform total. --- conf/gateway.yaml | 12 ++++- crates/config/src/lib.rs | 22 ++++++-- crates/dag/src/nodes.rs | 76 +++++++++++----------------- crates/handler/src/lib.rs | 12 ++--- crates/state/src/avail.rs | 15 +++--- crates/state/src/store.rs | 39 +++++++++++++- crates/views/src/lib.rs | 104 +++++++++++++++++++++++++++----------- 7 files changed, 182 insertions(+), 98 deletions(-) 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 d5c77a3..4af81d8 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -124,7 +124,8 @@ pub struct ModelConf { 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. + /// 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, } @@ -160,7 +161,7 @@ pub struct TokenRateConf { impl TokenRateConf { /// `(field name, value)` pairs for validation. - pub fn fields(&self) -> [(&'static str, f64); 5] { + fn fields(&self) -> [(&'static str, f64); 5] { [ ("prompt", self.prompt), ("read_cache", self.read_cache), @@ -308,6 +309,7 @@ pub struct StabilityConf { #[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`. @@ -710,7 +712,7 @@ impl GatewayConfig { } } } - for v in &m.variants { + for (i, v) in m.variants.iter().enumerate() { let bad = |reason| ConfigError::BadVariant { model: m.name.clone(), variant: v.model.clone(), @@ -719,6 +721,10 @@ impl GatewayConfig { 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")); }; @@ -783,9 +789,11 @@ impl GatewayConfig { } let st = &self.stability; let bad_rate = |v: f64| !v.is_finite() || !(0.0..=1.0).contains(&v); - if st.availability_window_minutes < 1 { + // 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", + field: "availability_window_minutes (1..=60)", }); } if bad_rate(st.unstable_error_rate) || bad_rate(st.unavailable_error_rate) { @@ -1321,6 +1329,10 @@ tenants: [{name: t1}, {name: t1}] ("[{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}}]}}]" diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index a048665..509c584 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -493,6 +493,9 @@ 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. pub struct CallEngine; #[async_trait::async_trait] @@ -527,7 +530,6 @@ impl DagNode for CallEngine { Ok(()) } Err(first_err) if first_err.http_status >= 500 => { - ctx.state.avail.record(served_model(ctx), false); let mt = ctx .request .protocol() @@ -564,6 +566,7 @@ impl DagNode for CallEngine { ) .await; let Some(next) = next else { + ctx.state.avail.record(served_model(ctx), false); return Err(first_err); }; let spillover = failed.is_ptu() && !next.is_ptu(); @@ -585,9 +588,7 @@ impl DagNode for CallEngine { Ok(()) } Err(e) => { - if e.http_status >= 500 { - ctx.state.avail.record(served_model(ctx), false); - } + ctx.state.avail.record(served_model(ctx), false); ctx.state .health .record_failure(&next.name, threshold, cooldown) @@ -654,9 +655,12 @@ 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)); + let tokens = BillTokens::weighted(pt, ct, &rate); ctx.decide("cost_calc", format!("aborted stream, billed {pt}+{ct}")); - return bill(ctx, BillTokens::flat(pt, ct), true).await; + return bill(ctx, tokens, true).await; } + let rate = gw_state::model_token_rate(&ctx.cfg, served_model(ctx)); // saturating sums so a malformed usage subtree can't overflow the totals let tokens = match &resp.common_usage { Some(u) => { @@ -667,13 +671,6 @@ impl DagNode for CostCalc { completion: u.completion, reasoning: u.reason, }; - let served = ctx - .request - .model_param_v2 - .as_ref() - .map(|p| p.model_name.as_str()) - .unwrap_or_default(); - let rate = token_rate(&ctx.cfg, served); BillTokens { prompt: u.prompt_total(), completion: u.completion_total(), @@ -681,28 +678,12 @@ impl DagNode for CostCalc { billable_completion: gw_models::weighted_completion(&ti, &rate), } } - None => BillTokens::flat(resp.prompt_tokens, resp.completion_tokens), + None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), }; bill(ctx, tokens, false).await } } -/// The served model's billing weights; identity when unconfigured. The -/// `prompt_includes_cache` normalization already happened at usage extraction. -fn 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(), - } -} - /// 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. @@ -714,13 +695,20 @@ struct BillTokens { } impl BillTokens { - /// Unweighted counts: billable equals reported. - fn flat(prompt: i64, completion: i64) -> Self { + /// 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 ti = gw_models::TokenInput { + prompt, + completion, + ..Default::default() + }; Self { prompt, completion, - billable_prompt: prompt, - billable_completion: completion, + billable_prompt: gw_models::weighted_prompt(&ti, rate), + billable_completion: gw_models::weighted_completion(&ti, rate), } } @@ -906,19 +894,15 @@ fn served_model(ctx: &DagContext) -> &str { #[cfg(test)] mod tests { #[test] - fn 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 = super::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!( - super::token_rate(&cfg, "absent"), - gw_models::TokenRate::default() - ); + 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] diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 06b708a..cf218c2 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -495,17 +495,17 @@ mod tests { 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]; - // raw usage stays truthful: 20 fresh + 80 cached - assert_eq!(rec.prompt_tokens, 100); - // billable prompt = 20 + 80*0.1 = 28 at 1000 micros/1k input - assert_eq!(rec.cost_micros, 28); + 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 variant_split_bills_requested_serves_target() { - // tenant entitled only to the public name — proves entitlement is - // judged before the variant swap + // 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)); diff --git a/crates/state/src/avail.rs b/crates/state/src/avail.rs index 19cc0be..714284d 100644 --- a/crates/state/src/avail.rs +++ b/crates/state/src/avail.rs @@ -211,11 +211,8 @@ impl AvailStore for RedisAvail { return (0, 0); } }; - vals.chunks(2).fold((0, 0), |(o, e), pair| { - ( - o + pair.first().copied().flatten().unwrap_or(0), - e + pair.get(1).copied().flatten().unwrap_or(0), - ) + vals.chunks_exact(2).fold((0, 0), |(o, e), pair| { + (o + pair[0].unwrap_or(0), e + pair[1].unwrap_or(0)) }) } } @@ -236,7 +233,6 @@ mod tests { 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); - // min_samples 0 still needs at least one sample assert_eq!(classify(0, 0, 0, 0.1, 0.5), AvailState::NoData); } @@ -251,9 +247,12 @@ mod tests { assert_eq!(tracker.window("m", minute - 5, minute).await, (0, 0)); tracker.flush().await; assert_eq!(tracker.window("m", minute - 5, minute).await, (3, 1)); - // second flush with an empty buffer adds nothing tracker.flush().await; - assert_eq!(tracker.window("m", minute - 5, minute).await, (3, 1)); + assert_eq!( + tracker.window("m", minute - 5, minute).await, + (3, 1), + "empty-buffer flush adds nothing" + ); } #[tokio::test] diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index f17e76d..d2ffc33 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). @@ -197,11 +199,28 @@ 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 token columns keep the -/// vendor-reported counts so usage stays truthful. +/// 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), @@ -2381,6 +2400,22 @@ 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}]"; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index d06411b..4903463 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -359,12 +359,24 @@ 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); + // same weights as the request pipeline; realtime usage carries no cache + // components, so only the prompt/completion weights can bite + let rate = gw_state::model_token_rate(cfg, model); + let ti = gw_models::TokenInput { + prompt: it, + completion: ot, + ..Default::default() + }; + let (bp, bc) = ( + gw_models::weighted_prompt(&ti, &rate), + gw_models::weighted_completion(&ti, &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( @@ -384,8 +396,8 @@ async fn bill_realtime_turn( account, prompt: it, completion: ot, - billable_prompt: it, - billable_completion: ot, + billable_prompt: bp, + billable_completion: bc, total, ptu_spillover: false, estimated: false, @@ -1535,30 +1547,35 @@ async fn admin_models_status(State(s): State, scope: AdminScope) -> Re let st = &cfg.stability; let until = gw_state::epoch_secs() / 60; let since = until - (st.availability_window_minutes - 1); - let avail = &s.handler.state().avail; - let mut rows = Vec::new(); - for m in &cfg.models { - if let AdminScope::Tenant(t) = &scope - && !cfg.tenant_allows_model(t, &m.name) - { - continue; - } - let (ok, err) = avail.window(&m.name, since, until).await; - let state = gw_state::classify( - ok, - err, - st.availability_min_samples, - st.unstable_error_rate, - st.unavailable_error_rate, - ); - rows.push(json!({ - "model": m.name, - "state": state, - "requests": ok + err, - "errors": err, - "window_minutes": st.availability_window_minutes, - })); - } + let state = s.handler.state(); + 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( + 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() } @@ -3379,8 +3396,7 @@ mod tests { #[tokio::test] async fn models_status_classifies_and_scopes() { - // failure_threshold high so account cooldown never blocks a sample; - // models bound to disjoint providers so round-robin can't cross-route + // 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}]\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 { @@ -3733,6 +3749,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()); From 25545087661d9bd39f58884f6aff611241252ce2 Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 22:32:06 +0800 Subject: [PATCH 05/11] review: glue cost.rs private helpers below the public surface --- crates/models/src/cost.rs | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/crates/models/src/cost.rs b/crates/models/src/cost.rs index 5577a28..68e3e29 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 { @@ -80,6 +71,15 @@ pub fn platform_total(input: &TokenInput, rate: &TokenRate) -> i64 { weighted_prompt(input, rate).saturating_add(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 } } From 8ccdf869947cefff2bd8f647f8aa3c4473454e37 Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 22:51:24 +0800 Subject: [PATCH 06/11] review: pin failover availability semantic and stability validation Regression tests from the verification pass: a failover-recovered request samples exactly one availability success; BadStability rejects window bounds and inverted/out-of-range rates. Plus a WHY note on the e2e wire-vs-ledger total assertion (equal only without token_rate). --- crates/config/src/lib.rs | 18 ++++++++++++++++++ crates/handler/src/lib.rs | 26 ++++++++++++++++++++++++++ crates/server/tests/e2e.rs | 3 +++ 3 files changed, 47 insertions(+) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 4af81d8..34d22a9 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -1349,6 +1349,24 @@ tenants: [{name: t1}, {name: t1}] let cfg = GatewayConfig::from_yaml(split).unwrap(); assert_eq!(cfg.find_model("m1").unwrap().variants.len(), 2); + 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/handler/src/lib.rs b/crates/handler/src/lib.rs index cf218c2..9bba600 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -503,6 +503,32 @@ mod tests { assert_eq!(rec.total_tokens, 28 + rec.completion_tokens); } + #[tokio::test] + async fn failover_recovery_records_one_availability_success() { + // a-down (priority 1) always 503s, a-up recovers via failover; the + // model must sample exactly one success, no error + 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 diff --git a/crates/server/tests/e2e.rs b/crates/server/tests/e2e.rs index 5ee9f99..77bc093 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -1156,6 +1156,9 @@ 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"); + // holds because gpt-4o sets no token_rate: the ledger total is the weighted + // platform total, the wire usage stays vendor-raw — they diverge by design + // once weights are configured assert_eq!(rec["total_tokens"].as_i64().unwrap(), total); assert!(rec["cost_micros"].as_i64().unwrap() > 0); } From 9507a23ce73594279f49c988c3c260a7d02a5e5c Mon Sep 17 00:00:00 2001 From: CMGS Date: Fri, 17 Jul 2026 23:35:19 +0800 Subject: [PATCH 07/11] fix: close the realtime/variant blind spots MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit External review found two fail-silent gaps: legal variants config on a realtime model routed 100% to the parent (sessions never traverse VariantSelect) — now rejected at load; and availability attributed samples to the served backend name, leaving a tenant's entitled public name at permanent no_data under a variant split — samples now attribute to the requested public name, matching the client-visible-outcome semantic. Realtime models are excluded from /admin/models/status (never sampled; listing them read as a permanent no_data). Docs catch up: the status endpoint in api.md, token_rate/variants/availability in configuration.md. Also deflakes sqlite_erase_roundtrip (5ms marker margin lost races under parallel-suite load). --- crates/config/src/lib.rs | 15 +++++++++++++++ crates/dag/src/nodes.rs | 22 +++++++++++++++++----- crates/handler/src/lib.rs | 9 +++++++++ crates/state/src/store.rs | 7 ++++++- crates/views/src/lib.rs | 19 ++++++++++++++----- docs/api.md | 1 + docs/configuration.md | 18 ++++++++++++++++++ 7 files changed, 80 insertions(+), 11 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 34d22a9..bdbba25 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -718,6 +718,12 @@ impl GatewayConfig { variant: v.model.clone(), reason, }; + // realtime sessions never traverse VariantSelect (they pin a + // model at handshake) — accepting the config would silently + // serve 100% 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")); } @@ -1349,6 +1355,15 @@ tenants: [{name: t1}, {name: t1}] 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", diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 509c584..ffb53fe 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -495,7 +495,8 @@ impl DagNode for UserBudgetGate { /// 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. +/// 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] @@ -514,7 +515,7 @@ impl DagNode for CallEngine { Ok(outcome) => { // an aborted stream is neither a success nor an account fault if !outcome.response.aborted { - ctx.state.avail.record(served_model(ctx), true); + ctx.state.avail.record(requested_model(ctx), true); if let Some(a) = ctx.request.account.as_ref() { ctx.state.health.record_success(&a.name).await; } @@ -566,7 +567,7 @@ impl DagNode for CallEngine { ) .await; let Some(next) = next else { - ctx.state.avail.record(served_model(ctx), false); + ctx.state.avail.record(requested_model(ctx), false); return Err(first_err); }; let spillover = failed.is_ptu() && !next.is_ptu(); @@ -581,14 +582,14 @@ 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.avail.record(served_model(ctx), true); + ctx.state.avail.record(requested_model(ctx), 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(served_model(ctx), false); + ctx.state.avail.record(requested_model(ctx), false); ctx.state .health .record_failure(&next.name, threshold, cooldown) @@ -891,6 +892,17 @@ fn served_model(ctx: &DagContext) -> &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(ctx: &DagContext) -> &str { + ctx.request + .model_param_v2 + .as_ref() + .map(|p| p.fallback_from.as_deref().unwrap_or(p.model_name.as_str())) + .unwrap_or_default() +} + #[cfg(test)] mod tests { #[test] diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 9bba600..6a5e53b 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -549,6 +549,15 @@ mod tests { 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)] diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index d2ffc33..a6b20ae 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -2897,8 +2897,13 @@ 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: the erasure ran earlier in this test, and a tight one + // (5ms) 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/views/src/lib.rs b/crates/views/src/lib.rs index 4903463..c11c7c3 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -1541,16 +1541,21 @@ 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. +/// sees only its entitled models. Realtime models are excluded: sessions +/// don't traverse the sampled engine path, and listing them would show a +/// permanent (misleading) 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| match &scope { - AdminScope::Tenant(t) => cfg.tenant_allows_model(t, &m.name), - AdminScope::Global => true, + 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( @@ -3397,7 +3402,7 @@ mod tests { #[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}]\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}"; + 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"); @@ -3444,6 +3449,10 @@ 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 resp = router.clone().oneshot(status("t1-tok")).await.unwrap(); let j = body_json(resp).await; 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 From 51db9a7286b3643ca3e510de680f1a286ea6000f Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 18 Jul 2026 00:46:36 +0800 Subject: [PATCH 08/11] fix: sample pool-exhaustion 503s so outages converge to unavailable MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All avail.record sites sat in CallEngine, so once an account tripped into cooldown every subsequent 503 came from SelectAccount unsampled — a sustained outage held ~failure_threshold samples per window, below min_samples, reading no_data forever (zero-account models: permanently). SelectAccount now records one failure against the requested public name before erroring; cache hits are unaffected (the executor short-circuits them before this node). --- crates/dag/src/nodes.rs | 20 +++++++++++-------- crates/views/src/lib.rs | 44 +++++++++++++++++++++++++++++++++++++++++ 2 files changed, 56 insertions(+), 8 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index ffb53fe..05d70f3 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -354,14 +354,18 @@ 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 (all accounts cooling, or none configured) is a + // client-visible model failure; without this sample a sustained + // outage sits below min_samples and reads as no_data forever + ctx.state.avail.record(requested_model(ctx), 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(()) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index c11c7c3..5c116ec 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -3465,6 +3465,50 @@ mod tests { assert_eq!(names, ["m-ok"], "tenant admin sees only entitled models"); } + #[tokio::test] + async fn sustained_pool_exhaustion_converges_to_unavailable() { + // low threshold: requests past the 2nd die in SelectAccount during the + // cooldown — those 503s must sample too, or the outage reads no_data + 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}]"; From 2e705bda17220013068eea2ab9d3cf6114af86e0 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 18 Jul 2026 01:06:59 +0800 Subject: [PATCH 09/11] review: simplify lenses + comment tightening Applied: weighted_pair hoisted to gw_models (the dag estimate paths and realtime shared a hand-built TokenInput pattern); platform_total deleted (no production caller survived the BillTokens refactor); served_model/ requested_model narrowed to take the model param so bill() reuses them instead of re-deriving; token_rate validation shares the finite/non-negative predicate with the qps checks; fallback_from doc names both writers. Skipped with reasons: per-request ModelConf memo (tens of ns vs a stale-memo hazard at two rewrite points; rides the perf pass), spawn-loop dedup (touches three pre-existing tasks; loc-justify adjudicates), realtime pre-clamp removal (would desync user-budget consumption from the DAG's clamped totals). Test comments trimmed to the zero-budget, multi-line WHYs compressed. --- crates/config/src/lib.rs | 13 ++++--- crates/dag/src/nodes.rs | 70 ++++++++++++++++++------------------ crates/handler/src/lib.rs | 2 -- crates/models/src/cost.rs | 40 ++++++++++++--------- crates/models/src/lib.rs | 2 +- crates/models/src/request.rs | 4 +-- crates/server/tests/e2e.rs | 5 ++- crates/state/src/store.rs | 3 +- crates/views/src/lib.rs | 21 +++-------- 9 files changed, 76 insertions(+), 84 deletions(-) diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index bdbba25..8b167d5 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -695,6 +695,7 @@ 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 { @@ -704,7 +705,7 @@ impl GatewayConfig { // a NaN weight would poison every downstream cost computation if let Some(rate) = &m.token_rate { for (field, v) in rate.fields() { - if !v.is_finite() || v < 0.0 { + if neg_or_nan(v) { return Err(ConfigError::BadTokenRate { model: m.name.clone(), field, @@ -718,9 +719,8 @@ impl GatewayConfig { variant: v.model.clone(), reason, }; - // realtime sessions never traverse VariantSelect (they pin a - // model at handshake) — accepting the config would silently - // serve 100% parent + // 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")); } @@ -765,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) @@ -776,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) { diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 05d70f3..027053e 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -210,8 +210,7 @@ fn pick_variant<'a>( } roll -= u64::from(v.weight); } - // weights are validated >= 1, so the loop always returns; this arm only - // quiets the type checker + // unreachable (weights validated >= 1); quiets the type checker &variants[0] } @@ -356,10 +355,11 @@ impl DagNode for SelectAccount { .select_healthy(mt, provider, &[], ctx.state.health.as_ref()) .await; let Some(account) = account else { - // an exhausted pool (all accounts cooling, or none configured) is a - // client-visible model failure; without this sample a sustained - // outage sits below min_samples and reads as no_data forever - ctx.state.avail.record(requested_model(ctx), false); + // 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, @@ -519,7 +519,9 @@ impl DagNode for CallEngine { Ok(outcome) => { // an aborted stream is neither a success nor an account fault if !outcome.response.aborted { - ctx.state.avail.record(requested_model(ctx), true); + 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; } @@ -571,7 +573,9 @@ impl DagNode for CallEngine { ) .await; let Some(next) = next else { - ctx.state.avail.record(requested_model(ctx), false); + 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(); @@ -586,14 +590,18 @@ 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.avail.record(requested_model(ctx), true); + 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), false); + ctx.state + .avail + .record(requested_model(ctx.request.model_param_v2.as_ref()), false); ctx.state .health .record_failure(&next.name, threshold, cooldown) @@ -660,12 +668,16 @@ 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)); + 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, tokens, true).await; } - let rate = gw_state::model_token_rate(&ctx.cfg, served_model(ctx)); + 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 tokens = match &resp.common_usage { Some(u) => { @@ -704,16 +716,13 @@ impl BillTokens { /// 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 ti = gw_models::TokenInput { - prompt, - completion, - ..Default::default() - }; + let (billable_prompt, billable_completion) = + gw_models::weighted_pair(prompt, completion, rate); Self { prompt, completion, - billable_prompt: gw_models::weighted_prompt(&ti, rate), - billable_completion: gw_models::weighted_completion(&ti, rate), + billable_prompt, + billable_completion, } } @@ -735,10 +744,8 @@ async fn bill(ctx: &mut DagContext, tokens: BillTokens, estimated: bool) -> GRes 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), @@ -887,22 +894,17 @@ 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). -fn served_model(ctx: &DagContext) -> &str { - ctx.request - .model_param_v2 - .as_ref() - .map(|p| p.model_name.as_str()) - .unwrap_or_default() +/// 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(ctx: &DagContext) -> &str { - ctx.request - .model_param_v2 - .as_ref() +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() } diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index 6a5e53b..49801ca 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -505,8 +505,6 @@ mod tests { #[tokio::test] async fn failover_recovery_records_one_availability_success() { - // a-down (priority 1) always 503s, a-up recovers via failover; the - // model must sample exactly one success, no error 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)); diff --git a/crates/models/src/cost.rs b/crates/models/src/cost.rs index 68e3e29..8f9402f 100644 --- a/crates/models/src/cost.rs +++ b/crates/models/src/cost.rs @@ -65,10 +65,19 @@ pub fn weighted_completion(input: &TokenInput, rate: &TokenRate) -> i64 { round_tokens(sum) } -/// Weighted platform-total token count; always the sum of the two sides so -/// quota consumption and per-side billing cannot drift. -pub fn platform_total(input: &TokenInput, rate: &TokenRate) -> i64 { - weighted_prompt(input, rate).saturating_add(weighted_completion(input, rate)) +/// 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). @@ -100,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] @@ -109,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] @@ -119,7 +130,8 @@ 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] @@ -138,21 +150,15 @@ mod tests { }; assert_eq!(weighted_prompt(&input, &rate), 250); assert_eq!(weighted_completion(&input, &rate), 60); - assert_eq!(platform_total(&input, &rate), 310); } #[test] - fn total_is_sum_of_sides() { + fn weighted_pair_carries_flat_counts() { let rate = TokenRate { - prompt_weight: 1.5, - completion_weight: 0.5, + prompt_weight: 0.5, ..Default::default() }; - let input = sample(); - assert_eq!( - platform_total(&input, &rate), - weighted_prompt(&input, &rate) + weighted_completion(&input, &rate) - ); + assert_eq!(weighted_pair(100, 50, &rate), (50, 50)); } #[test] @@ -167,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 d07bfa9..eee0e0e 100644 --- a/crates/models/src/lib.rs +++ b/crates/models/src/lib.rs @@ -14,7 +14,7 @@ pub mod usage; pub use block::Block; pub use cost::{ - TokenInput, TokenRate, cost_micros, platform_total, weighted_completion, weighted_prompt, + TokenInput, TokenRate, cost_micros, weighted_completion, weighted_pair, weighted_prompt, }; pub use error::{GResult, GatewayError}; pub use params::{ 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/tests/e2e.rs b/crates/server/tests/e2e.rs index 77bc093..ce7c864 100644 --- a/crates/server/tests/e2e.rs +++ b/crates/server/tests/e2e.rs @@ -1156,9 +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"); - // holds because gpt-4o sets no token_rate: the ledger total is the weighted - // platform total, the wire usage stays vendor-raw — they diverge by design - // once weights are configured + // 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/store.rs b/crates/state/src/store.rs index a6b20ae..d672b60 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -2897,8 +2897,7 @@ 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: the erasure ran earlier in this test, and a tight one - // (5ms) flaked under parallel-suite load + // wide margin: a 5ms one flaked under parallel-suite load assert!( store .user_erased_since(&t1, &u1, now - 60_000) diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index 5c116ec..d4aef21 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -364,18 +364,9 @@ async fn bill_realtime_turn( return; } let (cfg, state) = (&admit.snap.cfg, &admit.snap.state); - // same weights as the request pipeline; realtime usage carries no cache - // components, so only the prompt/completion weights can bite + // pipeline parity: the same shared weighting the DAG estimate paths use let rate = gw_state::model_token_rate(cfg, model); - let ti = gw_models::TokenInput { - prompt: it, - completion: ot, - ..Default::default() - }; - let (bp, bc) = ( - gw_models::weighted_prompt(&ti, &rate), - gw_models::weighted_completion(&ti, &rate), - ); + 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)); @@ -1541,9 +1532,8 @@ 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: sessions -/// don't traverse the sampled engine path, and listing them would show a -/// permanent (misleading) no_data. +/// 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; @@ -3467,8 +3457,7 @@ mod tests { #[tokio::test] async fn sustained_pool_exhaustion_converges_to_unavailable() { - // low threshold: requests past the 2nd die in SelectAccount during the - // cooldown — those 503s must sample too, or the outage reads no_data + // 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 { From 1f8b82505cf59dc16e5a5a501a1a5fa461c128a1 Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 18 Jul 2026 01:23:23 +0800 Subject: [PATCH 10/11] fix: exclude aborted streams from the failover-retry success record MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The retry-success arm recorded availability and account health without the first attempt's aborted-stream guard, so a client cutting the stream during a failover retry read as a model success — against CallEngine's own documented exclusion. --- crates/dag/src/nodes.rs | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index 027053e..e648ef6 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -590,10 +590,13 @@ 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 - .avail - .record(requested_model(ctx.request.model_param_v2.as_ref()), true); - 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(()) From 1fcc836cadcfa02c97d7460d458f82cb9591eb6f Mon Sep 17 00:00:00 2001 From: CMGS Date: Sat, 18 Jul 2026 10:00:09 +0800 Subject: [PATCH 11/11] review: apply the loc-justify cut-list AvailTracker dissolved into the AvailStore impls per the RedisHealth embedded-cache precedent: MemoryAvail records straight into its ring (no buffer between an in-process write and an in-process store), RedisAvail owns the claim-path buffer its flush drains, and the pure-passthrough window wrapper is gone. The buffer/ring bumps reuse slot_mut instead of re-rolling its miss-path idiom; fnv1a folds; the Redis window keys build via flat_map. Adds the GW_TEST_REDIS_URL-gated RedisAvail test matching the health/governance pattern. prompt_includes_cache kept (vendor- inclusive prompt accounting is a plausible near-term need). --- crates/dag/src/nodes.rs | 9 +- crates/state/src/avail.rs | 208 +++++++++++++++++--------------------- crates/state/src/lib.rs | 10 +- 3 files changed, 99 insertions(+), 128 deletions(-) diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index e648ef6..14883a2 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -217,12 +217,9 @@ fn pick_variant<'a>( /// 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 { - let mut h: u64 = 0xcbf2_9ce4_8422_2325; - for b in s.as_bytes() { - h ^= u64::from(*b); - h = h.wrapping_mul(0x0000_0100_0000_01b3); - } - h + 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 diff --git a/crates/state/src/avail.rs b/crates/state/src/avail.rs index 714284d..74b678c 100644 --- a/crates/state/src/avail.rs +++ b/crates/state/src/avail.rs @@ -1,10 +1,9 @@ //! Per-model availability from minute-bucketed success/error counts. The -//! claim path only bumps an in-process buffer; a background task flushes it -//! to the store (Redis when the fleet shares one), so recording never adds a -//! network hop to a request. +//! 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 std::sync::Arc; use async_trait::async_trait; use dashmap::DashMap; @@ -46,89 +45,27 @@ pub fn classify( } } -/// Claim-path recorder: counts accumulate in-process and [`Self::flush`] -/// folds them into the store under the flush-time minute (sub-period skew -/// across a minute boundary is accepted). -pub struct AvailTracker { - buffer: DashMap, - store: Arc, -} - -impl AvailTracker { - pub fn new(store: Arc) -> Self { - Self { - buffer: DashMap::new(), - store, - } - } - - /// Bump the model's (ok, err) pair; allocates only on a model's first - /// record since the last flush. - pub fn record(&self, model: &str, ok: bool) { - if let Some(mut e) = self.buffer.get_mut(model) { - bump(e.value_mut(), ok); - return; - } - bump( - self.buffer - .entry(model.to_owned()) - .or_insert((0, 0)) - .value_mut(), - ok, - ); - } - - /// Drain the buffer into the store. Increments landing mid-drain are - /// either taken now or survive for the next flush. - pub 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 { - if let Some((m, (ok, err))) = self.buffer.remove(&m) { - self.store.add(&m, minute, ok, err).await; - } - } - } - - /// Summed (ok, err) over `[since_minute, until_minute]`. - pub async fn window(&self, model: &str, since_minute: i64, until_minute: i64) -> (u64, u64) { - self.store.window(model, since_minute, until_minute).await - } -} - -impl std::fmt::Debug for AvailTracker { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str("AvailTracker") - } -} - -fn bump(counts: &mut (u64, u64), ok: bool) { - if ok { - counts.0 += 1; - } else { - counts.1 += 1; - } -} - -/// Minute-bucket store behind the flush task. +/// 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 { - /// Fold one flush increment into `minute`'s bucket. - async fn add(&self, model: &str, minute: i64, ok: u64, err: u64); + 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. +/// In-process ring for single-node deployments; records write it directly. #[derive(Debug, Default)] pub struct MemoryAvail { rings: DashMap>, } -#[async_trait] -impl AvailStore for MemoryAvail { - async fn add(&self, model: &str, minute: i64, ok: u64, err: u64) { - let mut ring = self.rings.entry(model.to_owned()).or_default(); +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; @@ -140,6 +77,14 @@ impl AvailStore for MemoryAvail { 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 @@ -153,17 +98,19 @@ impl AvailStore for MemoryAvail { } } -/// Fleet-shared buckets: every instance's flush increments the same keys. -/// A Redis outage drops the increment — availability is advisory, matching -/// the governance fail-open posture. +/// 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(), }) } } @@ -176,32 +123,48 @@ impl std::fmt::Debug for RedisAvail { #[async_trait] impl AvailStore for RedisAvail { - async fn add(&self, model: &str, minute: i64, ok: u64, err: u64) { - let mut conn = self.conn.clone(); - let mut pipe = redis::pipe(); - for (key, n) in [ - (count_key(model, minute, "ok"), ok), - (count_key(model, minute, "err"), err), - ] { - if n > 0 { - pipe.cmd("INCRBY").arg(&key).arg(n).ignore(); - pipe.cmd("EXPIRE") - .arg(&key) - .arg(RETAIN_MINUTES * 60) - .ignore(); - } + 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; } - if let Err(e) = pipe.query_async::<()>(&mut conn).await { - tracing::warn!(error = %e, model, "redis avail unavailable; counts dropped"); + } + + /// 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 mut keys = Vec::new(); - for minute in since_minute..=until_minute { - keys.push(count_key(model, minute, "ok")); - keys.push(count_key(model, minute, "err")); - } + 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 { @@ -237,32 +200,25 @@ mod tests { } #[tokio::test] - async fn tracker_buffers_flushes_and_windows() { - let tracker = AvailTracker::new(Arc::new(MemoryAvail::default())); + async fn memory_records_and_windows() { + let store = MemoryAvail::default(); for _ in 0..3 { - tracker.record("m", true); + store.record("m", true); } - tracker.record("m", false); + store.record("m", false); + store.flush().await; let minute = crate::epoch_secs() / 60; - assert_eq!(tracker.window("m", minute - 5, minute).await, (0, 0)); - tracker.flush().await; - assert_eq!(tracker.window("m", minute - 5, minute).await, (3, 1)); - tracker.flush().await; - assert_eq!( - tracker.window("m", minute - 5, minute).await, - (3, 1), - "empty-buffer flush adds nothing" - ); + 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).await; - store.add("m", 101, 2, 0).await; + 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).await; + store.add("m", 100 + RETAIN_MINUTES, 1, 0); assert_eq!( store.window("m", 100, 100).await, (0, 0), @@ -270,4 +226,24 @@ mod tests { ); 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 380abbc..7455b37 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -22,7 +22,7 @@ pub mod health; pub mod keystore; pub mod store; -pub use avail::{AvailState, AvailTracker, classify}; +pub use avail::{AvailState, AvailStore, classify}; pub use configstore::{CONFIG_CHANNEL, PostgresConfigStore}; pub use content::{ContentRecord, sealing_available}; pub use governance::{Governance, MemoryGovernance, RedisGovernance}; @@ -666,7 +666,7 @@ pub struct GatewayState { /// 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, + pub avail: Arc, } impl Default for GatewayState { @@ -678,9 +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::AvailTracker::new(Arc::new( - avail::MemoryAvail::default(), - ))), + avail: Arc::new(avail::MemoryAvail::default()), } } } @@ -761,7 +759,7 @@ impl GatewayState { } match avail::RedisAvail::connect(&cfg.storage.redis_url).await { Ok(a) => { - state.avail = Arc::new(avail::AvailTracker::new(Arc::new(a))); + state.avail = Arc::new(a); tracing::info!("model availability = redis (fleet-wide counts)"); } Err(e) => {