diff --git a/Cargo.lock b/Cargo.lock index 2e77991..da0db2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1018,6 +1018,7 @@ dependencies = [ "sqlx", "tempfile", "tokio", + "tokio-util", "tracing", ] @@ -2847,6 +2848,7 @@ dependencies = [ "bytes", "futures-core", "futures-sink", + "futures-util", "pin-project-lite", "tokio", ] diff --git a/Cargo.toml b/Cargo.toml index f759d7f..888fc00 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -35,6 +35,7 @@ undocumented_unsafe_blocks = "deny" [workspace.dependencies] # --- async runtime --- tokio = { version = "1", features = ["rt-multi-thread", "macros", "time", "sync", "net", "signal", "io-util"] } +tokio-util = { version = "0.7", features = ["rt"] } async-trait = "0.1" futures = "0.3" # --- web (server) --- diff --git a/crates/config/src/lib.rs b/crates/config/src/lib.rs index 47642e6..d7d2808 100644 --- a/crates/config/src/lib.rs +++ b/crates/config/src/lib.rs @@ -580,6 +580,9 @@ pub struct StorageConf { /// Prune rolled billing records beyond this many; 0 = unlimited. #[serde(default)] pub ledger_max_rows: u64, + /// Local file holding ledger rows the store refuses until they replay; empty drops them after the retries. + #[serde(default)] + pub ledger_spool_path: String, /// Postgres pool size for the store and key-table pools; 0 = their built-in defaults. #[serde(default)] pub postgres_max_connections: u32, diff --git a/crates/dag/src/context.rs b/crates/dag/src/context.rs index 1f5cb66..4eeff60 100644 --- a/crates/dag/src/context.rs +++ b/crates/dag/src/context.rs @@ -8,7 +8,7 @@ use std::sync::Arc; use gw_config::GatewayConfig; use gw_engines::{EngineOutcome, SharedTransport}; use gw_models::{GResult, GatewayError, GatewayRequest, ModelParamV2}; -use gw_state::admission::TpmReserve; +use gw_state::admission::{BudgetHold, TpmReserve}; use gw_state::{AkInfo, GatewayState, UserBudget}; pub struct DagContext { @@ -28,6 +28,7 @@ pub struct DagContext { /// A fallback model retries this attempt's upstream fault, so its failure is not the outcome. pub fallback_ahead: bool, pub(crate) request_limits_admitted: bool, + pub(crate) reserve_parts: (i64, i64), /// This request's cache key (computed by cache_lookup, reused by cache_store). pub cache_key: Option, /// Governance key for the (AK, model) daily counter; set only when a cap is configured. @@ -40,8 +41,11 @@ pub struct DagContext { pub tpm_reserved: Option, /// The user's budget override resolved at admission, charged at settlement. pub user_budget: Option, + /// Day-budget amounts held at admission (same lifecycle as `quota_reserved`). + pub budget_hold: BudgetHold, /// Outbound DLP buffered this stream; billing waits for the view's delivery result. pub billing_deferred: bool, + pub served_at: Option, } impl DagContext { @@ -63,13 +67,16 @@ impl DagContext { cache_hit: false, fallback_ahead: false, request_limits_admitted: false, + reserve_parts: (0, 0), cache_key: None, model_quota_key: None, quota_reserved: None, quota_at: 0, tpm_reserved: None, user_budget: None, + budget_hold: BudgetHold::default(), billing_deferred: false, + served_at: None, } } @@ -85,6 +92,19 @@ impl DagContext { .ok_or_else(|| GatewayError::internal("model param missing after resolve_model")) } + pub async fn refund_reserves(&mut self) { + let gov = self.state.governance.as_ref(); + tokio::join!( + gov.refund_reserves( + &self.ak.ak_id, + self.quota_reserved.take().unwrap_or(0), + self.tpm_reserved.take(), + self.quota_at, + ), + std::mem::take(&mut self.budget_hold).refund(gov), + ); + } + /// The effective end user: the key's `owner`, else request metadata, else `""`. pub fn effective_user_id(&self) -> &str { self.ak diff --git a/crates/dag/src/nodes.rs b/crates/dag/src/nodes.rs index d980a9f..58056c8 100644 --- a/crates/dag/src/nodes.rs +++ b/crates/dag/src/nodes.rs @@ -240,7 +240,11 @@ impl Write for PromptBytes { } /// Admission estimate: ~bytes/4 prompt plus the requested max_tokens, saturating and capped. -fn reserve_estimate(req: &gw_models::GatewayRequest) -> i64 { +fn reserve_total((prompt, max_out): (i64, i64)) -> i64 { + prompt.saturating_add(max_out).min(MAX_RESERVE) +} + +fn reserve_parts(req: &gw_models::GatewayRequest) -> (i64, i64) { let typed = req.model_param_v2.as_ref().and_then(|p| p.typed.as_ref()); let prompt = if let Some(gw_models::TypedParams::Decisions(p)) = typed { let mut bytes = PromptBytes(0); @@ -256,9 +260,16 @@ fn reserve_estimate(req: &gw_models::GatewayRequest) -> i64 { }) .unwrap_or(DEFAULT_COMPLETION_RESERVE) .clamp(0, MAX_RESERVE); - ((prompt as i64 / 4).max(1)) - .saturating_add(max_out) - .min(MAX_RESERVE) + ((prompt as i64 / 4).clamp(1, MAX_RESERVE), max_out) +} + +fn reserve_units(req: &gw_models::GatewayRequest) -> i64 { + match req.model_param_v2.as_ref().and_then(|p| p.typed.as_ref()) { + Some(gw_models::TypedParams::Image(p)) => p.n.clamp(1, MAX_RESERVE), + Some(gw_models::TypedParams::AudioTts(p)) => p.input.chars().count() as i64, + Some(gw_models::TypedParams::Video(_)) => 0, + _ => 1, + } } /// preprocess/quota_check: reserve the estimate against the AK daily quota; billing settles it. @@ -270,7 +281,8 @@ impl DagNode for QuotaCheck { "quota_check" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let est = reserve_estimate(&ctx.request); + ctx.reserve_parts = reserve_parts(&ctx.request); + let est = reserve_total(ctx.reserve_parts); let at = gw_state::epoch_secs(); let reserved = admission::reserve_daily(ctx.state.governance.as_ref(), &ctx.ak, est, at) .await @@ -407,10 +419,7 @@ impl DagNode for AkTpmLimit { "ak_tpm" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let est = ctx - .quota_reserved - .filter(|&reserved| reserved > 0) - .unwrap_or_else(|| reserve_estimate(&ctx.request)); + let est = reserve_total(ctx.reserve_parts); ctx.tpm_reserved = admission::reserve_tpm(ctx.state.governance.as_ref(), &ctx.ak, est) .await .map_err(key_limit_denied)?; @@ -418,7 +427,7 @@ impl DagNode for AkTpmLimit { } } -/// model_access/budget: the tenant's token and cost budgets (soft caps) for the key and user. +/// model_access/budget: hold the estimate against the day budgets, check the month ones. pub struct BudgetGate; #[async_trait::async_trait] @@ -427,10 +436,34 @@ impl DagNode for BudgetGate { "budget" } async fn execute(&self, ctx: &mut DagContext) -> GResult<()> { - let over = admission::check_budgets(&ctx.state, &ctx.cfg, &ctx.ak, ctx.effective_user_id()) - .await? - .map_err(quota_denied)?; + let (prompt, max_out) = ctx.reserve_parts; + let (cfg, request, tenant) = (&ctx.cfg, &ctx.request, &ctx.ak.tenant); + let est = || { + let served = served_model(request.model_param_v2.as_ref()); + let units = + reserve_units(request).saturating_mul(cfg.unit_price_for_tenant(tenant, served)); + admission::BudgetSpend { + tokens: reserve_total((prompt, max_out)), + cost_micros: gw_models::cost_micros( + prompt, + max_out, + cfg.prices_for_tenant(tenant, served), + ) + .saturating_add(units), + } + }; + let (over, hold) = admission::reserve_budgets( + &ctx.state, + &ctx.cfg, + &ctx.ak, + ctx.effective_user_id(), + est, + ctx.quota_at, + ) + .await? + .map_err(quota_denied)?; ctx.user_budget = Some(over); + ctx.budget_hold = hold; Ok(()) } } @@ -658,37 +691,7 @@ impl DagNode for CostCalc { if resp.aborted && resp.completion_tokens == 0 { return bill_aborted_stream(ctx, None).await; } - 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 mut tokens = match &resp.common_usage { - Some(u) => { - let ti = gw_models::TokenInput { - prompt: u.platform_input, - audio_prompt: u.audio_input, - read_cache: u.read_cache, - write_cache: u.write_cache, - write_cache_1h: u.write_cache_1h, - completion: u.completion, - audio_completion: u.audio_output, - reasoning: u.reason, - }; - 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), - units: 0, - } - } - // total-only vendors (MiniMax v1): meter the total as completion, not zero - None if resp.prompt_tokens == 0 && resp.completion_tokens == 0 => { - BillTokens::weighted(0, resp.total_tokens, &rate) - } - None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), - }; - tokens.units = resp.billed_units; - bill(ctx, tokens, false).await + bill(ctx, usage_tokens(ctx)?, false).await } } @@ -736,7 +739,7 @@ impl BillTokens { } } -/// Close a buffered stream's billing: refund before delivery, bill estimates after a partial one. +/// Close a buffered stream's billing: refund before delivery, else bill its usage or an estimate. pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelivery) -> GResult<()> { if !ctx.billing_deferred { return Ok(()); @@ -753,19 +756,16 @@ pub async fn settle_deferred_stream(ctx: &mut DagContext, delivery: StreamDelive .outcome .as_mut() .ok_or_else(|| GatewayError::internal("stream delivery without an outcome"))?; - blank_stream_response(&mut outcome.response); - bill_aborted_stream(ctx, Some(completion)).await + if outcome.response.completion_tokens == 0 { + blank_stream_response(&mut outcome.response); + return bill_aborted_stream(ctx, Some(completion)).await; + } + outcome.response.message.clear(); + outcome.response.aborted = true; + bill(ctx, usage_tokens(ctx)?, false).await } StreamDelivery::None => { - ctx.state - .governance - .refund_reserves( - &ctx.ak.ak_id, - ctx.quota_reserved.take().unwrap_or(0), - ctx.tpm_reserved.take(), - ctx.quota_at, - ) - .await; + ctx.refund_reserves().await; if let Some(outcome) = ctx.outcome.as_mut() { outcome.response.prompt_tokens = 0; blank_stream_response(&mut outcome.response); @@ -780,6 +780,38 @@ async fn bill_aborted_stream( ctx: &mut DagContext, delivered_completion: Option, ) -> GResult<()> { + ctx.served_at = Some(std::time::Instant::now()); + let mut settle = None; + let usage = match (ctx.request.account.as_deref(), ctx.outcome.as_ref()) { + (Some(account), Some(outcome)) + if gw_engines::generation::records_generation( + account, + &outcome.response.upstream_id, + ) => + { + let guard = settle.insert(ctx.state.begin_settle()); + gw_engines::generation::recorded_usage( + ctx.transport.as_ref(), + account, + &outcome.response.upstream_id, + guard.closing(), + ) + .await + } + _ => None, + }; + if let (Some(usage), Some(outcome)) = (usage, ctx.outcome.as_mut()) { + let common = gw_engines::extract_common_usage(&usage, false); + if let Some(u) = &common { + let r = &mut outcome.response; + (r.prompt_tokens, r.completion_tokens) = (u.prompt_total(), u.completion_total()); + r.total_tokens = r.prompt_tokens.saturating_add(r.completion_tokens); + } + outcome.response.common_usage = common; + outcome.response.raw_usage = Some(usage); + ctx.decide("cost_calc", "aborted stream, billed the provider's record"); + return bill(ctx, usage_tokens(ctx)?, false).await; + } let response = &ctx .outcome .as_ref() @@ -818,6 +850,45 @@ fn blank_stream_response(response: &mut gw_models::GatewayResponse) { response.aborted = true; } +fn usage_tokens(ctx: &DagContext) -> GResult { + let resp = &ctx + .outcome + .as_ref() + .ok_or_else(|| GatewayError::internal("billing without an engine outcome"))? + .response; + 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 mut tokens = match &resp.common_usage { + Some(u) => { + let ti = gw_models::TokenInput { + prompt: u.platform_input, + audio_prompt: u.audio_input, + read_cache: u.read_cache, + write_cache: u.write_cache, + write_cache_1h: u.write_cache_1h, + completion: u.completion, + audio_completion: u.audio_output, + reasoning: u.reason, + }; + 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), + units: 0, + } + } + // total-only vendors (MiniMax v1): meter the total as completion, not zero + None if resp.prompt_tokens == 0 && resp.completion_tokens == 0 => { + BillTokens::weighted(0, resp.total_tokens, &rate) + } + None => BillTokens::weighted(resp.prompt_tokens, resp.completion_tokens, &rate), + }; + tokens.units = resp.billed_units; + Ok(tokens) +} + /// Settle reserves and write the ledger; `estimated` marks an aborted stream's estimated counts. async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> GResult<()> { let ptu_spillover = ctx @@ -848,10 +919,11 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> served_conf.and_then(|m| m.batch_discount).unwrap_or(1.0) }; // locals first: the whole-ctx borrow `effective_user_id` needs can't overlap these takes - let (reserved, tpm_reserved, model_quota_key) = ( + let (reserved, tpm_reserved, model_quota_key, budget_hold) = ( ctx.quota_reserved.take().unwrap_or(0), ctx.tpm_reserved.take(), ctx.model_quota_key.take(), + std::mem::take(&mut ctx.budget_hold), ); let settled = admission::settle_and_bill( ctx.state.as_ref(), @@ -885,6 +957,7 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> model_quota_key, budget_ak: &ctx.ak, user_budget: ctx.user_budget, + budget_hold, }, ) .await; @@ -892,7 +965,7 @@ async fn bill(ctx: &mut DagContext, mut tokens: BillTokens, estimated: bool) -> "cost_calc", format!( "tokens={} cost_micros={}", - settled.total_tokens, settled.cost_micros + settled.tokens, settled.cost_micros ), ); Ok(()) @@ -1095,7 +1168,43 @@ mod tests { } #[test] - fn reserve_estimate_saturates_on_hostile_max_tokens() { + fn reserve_units_counts_what_admission_already_knows() { + use gw_models::params::{ImageParams, TtsParams, VideoParams}; + use gw_models::{GatewayRequest, ModelParamV2, TypedParams}; + let units = |typed: TypedParams| { + let mut param = ModelParamV2::with_name(gw_consts::Protocol::OpenaiChat, "m"); + param.typed = Some(typed); + super::reserve_units(&GatewayRequest { + model_param_v2: Some(param), + ..Default::default() + }) + }; + assert_eq!( + units(TypedParams::Image(ImageParams { + n: 3, + ..Default::default() + })), + 3 + ); + assert_eq!( + units(TypedParams::AudioTts(TtsParams { + input: "héllo".into(), + ..Default::default() + })), + 5 + ); + assert_eq!(units(TypedParams::Video(VideoParams::default())), 0); + assert_eq!( + units(TypedParams::Image(ImageParams { + n: i64::MAX, + ..Default::default() + })), + super::MAX_RESERVE + ); + } + + #[test] + fn reserve_total_saturates_on_hostile_max_tokens() { use gw_models::params::ChatParams; use gw_models::{ChatMsg, GatewayRequest, ModelParamV2, TypedParams}; let mut param = ModelParamV2::with_name(gw_consts::Protocol::OpenaiChat, "m"); @@ -1108,7 +1217,7 @@ mod tests { model_param_v2: Some(param), ..Default::default() }; - let est = super::reserve_estimate(&req); + let est = super::reserve_total(super::reserve_parts(&req)); assert!(est > 0 && est <= super::MAX_RESERVE); } } diff --git a/crates/engines/Cargo.toml b/crates/engines/Cargo.toml index 724c26e..b09c5c6 100644 --- a/crates/engines/Cargo.toml +++ b/crates/engines/Cargo.toml @@ -29,6 +29,7 @@ tokio = { workspace = true } [dev-dependencies] # real-socket wire test for HttpTransport (loopback only, never egress) axum = { workspace = true } +tokio = { workspace = true, features = ["test-util"] } [lints] workspace = true diff --git a/crates/engines/src/generation.rs b/crates/engines/src/generation.rs new file mode 100644 index 0000000..ded339f --- /dev/null +++ b/crates/engines/src/generation.rs @@ -0,0 +1,237 @@ +//! Post-hoc usage of an OpenRouter generation, which OpenRouter records even for a cancelled stream. + +use std::time::Duration; + +use gw_consts::Protocol; +use gw_models::{Account, GResult}; +use serde_json::{Value, json}; + +use crate::base::{parse_json_reply, versioned_url}; +use crate::transport::{Transport, UpstreamRequest}; + +const OPENROUTER_KIND: &str = "openrouter"; +/// Waits before each lookup of an abandoned stream's recorded usage; OpenRouter records one in about 5 s. +const LOOKUP_DELAYS: [Duration; 4] = [ + Duration::from_secs(3), + Duration::from_secs(5), + Duration::from_secs(10), + Duration::from_secs(20), +]; +const LOOKUP_TIMEOUT: Duration = Duration::from_secs(10); + +pub async fn recorded_usage( + transport: &dyn Transport, + account: &Account, + id: &str, + closing: impl Future, +) -> Option { + let key = account.api_key()?; + tokio::pin!(closing); + for delay in LOOKUP_DELAYS { + let lookup = async { + tokio::time::sleep(delay).await; + tokio::time::timeout( + LOOKUP_TIMEOUT, + openrouter_usage(transport, account, id, &key), + ) + .await + }; + tokio::select! { + () = &mut closing => { + tracing::warn!(id, "shutting down before the generation record; billing the estimate"); + return None; + } + found = lookup => match found { + Ok(Ok(Some(usage))) => return Some(usage), + Ok(Ok(None)) => {} + Ok(Err(e)) => tracing::warn!(error = %e, id, "generation usage lookup failed"), + Err(_) => tracing::warn!(id, "generation usage lookup timed out"), + }, + } + } + tracing::warn!( + id, + "no recorded usage for an aborted stream; billing the estimate" + ); + None +} + +pub fn records_generation(account: &Account, id: &str) -> bool { + account.wire_kind() == OPENROUTER_KIND && !account.endpoint.is_empty() && !id.is_empty() +} + +/// The generation's native counts and cost as chat-wire usage; `None` until OpenRouter records it. +async fn openrouter_usage( + transport: &dyn Transport, + account: &Account, + id: &str, + key: &str, +) -> GResult> { + let resp = transport + .send(UpstreamRequest { + protocol: Protocol::OpenaiChat, + method: "GET", + url: versioned_url(account.base_url(""), &format!("generation?id={id}")), + headers: vec![("authorization", format!("Bearer {key}"))], + body: Vec::new(), + stream: false, + account: account.name.clone(), + replay_account: None, + output_cap: 0, + }) + .await?; + if resp.status == 404 { + return Ok(None); + } + let (_, doc) = parse_json_reply(resp)?; + Ok(generation_usage(&doc["data"])) +} + +/// Map `/generation` data to chat-wire usage, reasoning inside completion as in `completion_tokens`. +fn generation_usage(data: &Value) -> Option { + let count = |k: &str| data[k].as_i64(); + let completion = count("native_tokens_completion")?; + Some(json!({ + "prompt_tokens": count("native_tokens_prompt")?, + "completion_tokens": completion, + "prompt_tokens_details": {"cached_tokens": count("native_tokens_cached").unwrap_or(0)}, + "completion_tokens_details": {"reasoning_tokens": count("native_tokens_reasoning").unwrap_or(0).clamp(0, completion)}, + "cost": data["total_cost"], + })) +} + +#[cfg(test)] +mod tests { + use std::sync::Mutex; + + use super::*; + use crate::transport::{UpstreamBody, UpstreamResponse}; + + #[derive(Debug)] + struct Canned(u16, &'static str, Mutex>); + + #[async_trait::async_trait] + impl Transport for Canned { + async fn send(&self, req: UpstreamRequest) -> GResult { + self.2 + .lock() + .unwrap() + .push(format!("{} {}", req.method, req.url)); + Ok(UpstreamResponse { + status: self.0, + body: UpstreamBody::Json(bytes::Bytes::from_static(self.1.as_bytes())), + headers: Default::default(), + }) + } + } + + fn openrouter() -> Account { + Account { + name: "or".into(), + kind: OPENROUTER_KIND.into(), + endpoint: "https://openrouter.ai/api".into(), + api_key_env: "PATH".into(), + ..Default::default() + } + } + + #[tokio::test] + async fn a_recorded_generation_maps_to_chat_usage() { + let t = Canned( + 200, + r#"{"data":{"cancelled":true,"native_tokens_prompt":11,"native_tokens_completion":73,"native_tokens_cached":4,"native_tokens_reasoning":73,"total_cost":9.09e-05}}"#, + Mutex::default(), + ); + let usage = openrouter_usage(&t, &openrouter(), "gen-1792232109-kUA52UZsI2", "k") + .await + .unwrap() + .unwrap(); + assert_eq!( + t.2.lock().unwrap().as_slice(), + ["GET https://openrouter.ai/api/v1/generation?id=gen-1792232109-kUA52UZsI2"] + ); + let common = crate::extract_common_usage(&usage, false).unwrap(); + assert_eq!( + ( + common.prompt_total(), + common.read_cache, + common.completion_total(), + common.reason + ), + (11, 4, 73, 73) + ); + assert_eq!(crate::extract_vendor_cost_micros(&usage), Some(91)); + } + + #[tokio::test] + async fn an_unrecorded_generation_yields_nothing() { + let missing = Canned(404, r#"{"error":{"code":404}}"#, Mutex::default()); + assert!( + openrouter_usage(&missing, &openrouter(), "gen-1", "k") + .await + .unwrap() + .is_none() + ); + let pending = Canned( + 200, + r#"{"data":{"native_tokens_prompt":11}}"#, + Mutex::default(), + ); + assert!( + openrouter_usage(&pending, &openrouter(), "gen-1", "k") + .await + .unwrap() + .is_none() + ); + } + + #[derive(Debug, Default)] + struct Silent(std::sync::atomic::AtomicUsize); + + #[async_trait::async_trait] + impl Transport for Silent { + async fn send(&self, _: UpstreamRequest) -> GResult { + self.0.fetch_add(1, std::sync::atomic::Ordering::SeqCst); + std::future::pending().await + } + } + + #[tokio::test(start_paused = true)] + async fn a_lookup_that_never_answers_times_out_each_try() { + let t = Silent::default(); + let usage = recorded_usage(&t, &openrouter(), "gen-1", std::future::pending()).await; + assert!(usage.is_none()); + assert_eq!( + t.0.load(std::sync::atomic::Ordering::SeqCst), + LOOKUP_DELAYS.len() + ); + } + + #[tokio::test(start_paused = true)] + async fn a_closing_gateway_skips_the_lookup() { + let t = Silent::default(); + let usage = recorded_usage(&t, &openrouter(), "gen-1", std::future::ready(())).await; + assert!(usage.is_none()); + assert_eq!(t.0.load(std::sync::atomic::Ordering::SeqCst), 0); + } + + #[test] + fn only_an_openrouter_account_with_an_id_records_generations() { + assert!(records_generation( + &openrouter(), + "gen-1792232109-kUA52UZsI2" + )); + let other = Account { + kind: "openai".into(), + ..openrouter() + }; + assert!(!records_generation(&other, "gen-1")); + let raw = Account { + kind: String::new(), + provider: OPENROUTER_KIND.into(), + ..openrouter() + }; + assert!(records_generation(&raw, "gen-1")); + assert!(!records_generation(&openrouter(), "")); + } +} diff --git a/crates/engines/src/lib.rs b/crates/engines/src/lib.rs index aed0c23..8f446eb 100644 --- a/crates/engines/src/lib.rs +++ b/crates/engines/src/lib.rs @@ -14,6 +14,7 @@ pub mod engine; pub mod eventstream; pub mod factory; pub mod families; +pub mod generation; pub mod http_transport; pub mod mock_transport; pub mod multipart; diff --git a/crates/engines/src/openai_engine.rs b/crates/engines/src/openai_engine.rs index 39bcb6b..525869c 100644 --- a/crates/engines/src/openai_engine.rs +++ b/crates/engines/src/openai_engine.rs @@ -325,6 +325,11 @@ fn apply_sse_event( if resp.model.is_empty() { resp.model = crate::engine::take_string(&mut v, "/model").unwrap_or_default(); } + if resp.upstream_id.is_empty() + && let Some(id) = crate::engine::take_string(&mut v, "/id") + { + resp.upstream_id = id; + } let mut tool_calls = None; if let Some(delta) = v .get_mut("choices") diff --git a/crates/handler/src/lib.rs b/crates/handler/src/lib.rs index af0c973..1e627c7 100644 --- a/crates/handler/src/lib.rs +++ b/crates/handler/src/lib.rs @@ -270,15 +270,7 @@ impl OnlineHandler { break; }; // a failed pipeline refunds its reservations whole, on the reserve's day bucket - ctx.state - .governance - .refund_reserves( - &ctx.ak.ak_id, - ctx.quota_reserved.take().unwrap_or(0), - ctx.tpm_reserved.take(), - ctx.quota_at, - ) - .await; + ctx.refund_reserves().await; // signed thinking replays only against the model that produced it if is_upstream_fault(&e) && !ctx.request.replays_reasoning_output() @@ -2717,6 +2709,224 @@ mod tests { ); } + #[derive(Debug, Default)] + struct RecordedGenerationStream { + lookups: std::sync::atomic::AtomicUsize, + } + + #[async_trait::async_trait] + impl gw_engines::transport::Transport for RecordedGenerationStream { + async fn send( + &self, + req: gw_engines::transport::UpstreamRequest, + ) -> GResult { + use futures::StreamExt; + use gw_engines::transport::{UpstreamBody, UpstreamResponse}; + if req.method == "GET" { + assert!( + req.url.ends_with("/v1/generation?id=gen-test-1"), + "{}", + req.url + ); + let first = self + .lookups + .fetch_add(1, std::sync::atomic::Ordering::SeqCst) + == 0; + let (status, body) = if first { + (404, r#"{"error":{"code":404}}"#) + } else { + ( + 200, + r#"{"data":{"cancelled":true,"native_tokens_prompt":120,"native_tokens_completion":80,"native_tokens_cached":100,"native_tokens_reasoning":0,"total_cost":0.0005}}"#, + ) + }; + return Ok(UpstreamResponse { + status, + body: UpstreamBody::Json(bytes::Bytes::from_static(body.as_bytes())), + headers: Default::default(), + }); + } + let frames: Vec> = vec![ + Ok(bytes::Bytes::from( + "data: {\"id\":\"gen-test-1\",\"choices\":[{\"delta\":{\"content\":\"partial\"}}]}\n\n", + )), + Err(gw_engines::transport::StreamFault { + timeout: false, + message: "connection reset".to_owned(), + }), + ]; + Ok(UpstreamResponse { + status: 200, + body: UpstreamBody::SseStream(futures::stream::iter(frames).boxed()), + headers: Default::default(), + }) + } + } + + async fn openrouter_handler() -> ( + OnlineHandler, + Arc, + Arc, + ) { + let yaml = "listen: {host: h, port: 1}\nproviders: [{name: openrouter, kind: openrouter, api_key_env: PATH}]\nmodels: [{name: m, provider: openrouter, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 2000, token_rate: {read_cache: 0.1}}]\ntenants: [{name: t1}]\naccess_keys: [{ak: or-ak, 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 transport = Arc::new(RecordedGenerationStream::default()); + let h = OnlineHandler::new(gw_state::SharedConfig::new(cfg, state), transport.clone()); + let ak = h.state().auth.get(&access_key_id("or-ak")).await.unwrap(); + (h, transport, ak) + } + + #[tokio::test(start_paused = true)] + async fn an_aborted_openrouter_stream_bills_the_recorded_generation() { + let (h, transport, ak) = openrouter_handler().await; + + let ctx = h + .run(drained_stream_req("m", "count to a thousand"), ak) + .await + .unwrap(); + + assert!(ctx.served_at.is_some()); + let response = ctx.outcome.expect("outcome").response; + assert!(response.aborted); + assert_eq!( + ( + response.prompt_tokens, + response.completion_tokens, + response.total_tokens + ), + (120, 80, 200), + "the access log and metrics read the record too" + ); + assert_eq!( + transport.lookups.load(std::sync::atomic::Ordering::SeqCst), + 2 + ); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(ledger.len(), 1); + let row = &ledger[0]; + assert!(!row.estimated, "{row:?}"); + assert_eq!( + (row.prompt_tokens, row.completion_tokens), + (120, 80), + "{row:?}" + ); + assert_eq!( + row.cost_micros, + 20 + 10 + 160, + "uncached at full price, cached at 0.1: {row:?}" + ); + assert_eq!(row.vendor_cost_micros, 500, "{row:?}"); + } + + #[tokio::test(start_paused = true)] + async fn shutdown_bills_a_pending_aborted_stream_before_the_flush_returns() { + let (h, transport, ak) = openrouter_handler().await; + let state = h.state(); + let running = tokio::spawn({ + let h = h.clone(); + async move { + h.run(drained_stream_req("m", "count to a thousand"), ak) + .await + } + }); + tokio::time::sleep(std::time::Duration::from_millis(200)).await; + + tokio::time::timeout( + std::time::Duration::from_secs(2), + gw_state::admission::flush_billing(&state), + ) + .await + .expect("the flush waits only for the cancelled lookup"); + + let (_, ledger) = state.store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(ledger.len(), 1, "billed before the flush returned"); + assert!(ledger[0].estimated, "{:?}", ledger[0]); + assert_eq!( + transport.lookups.load(std::sync::atomic::Ordering::SeqCst), + 0 + ); + running.await.unwrap().unwrap(); + } + + #[derive(Debug)] + struct GatedSearch(tokio::sync::Semaphore); + + #[async_trait::async_trait] + impl gw_engines::transport::Transport for GatedSearch { + async fn send( + &self, + _: gw_engines::transport::UpstreamRequest, + ) -> GResult { + self.0.acquire().await.unwrap().forget(); + Ok(gw_engines::transport::UpstreamResponse { + status: 200, + body: gw_engines::transport::UpstreamBody::Json(bytes::Bytes::from_static( + br#"{"items":[{"title":"t","url":"https://u"}],"metadata":{"query":"q","requestId":"r","latencyMs":1}}"#, + )), + headers: Default::default(), + }) + } + } + + #[tokio::test] + async fn an_in_flight_unit_priced_call_holds_the_daily_cost_budget() { + let yaml = "listen: {host: h, port: 1}\nmodels: [{name: web, protocol: search, provider: cloudflare, unit_price_micros: 250}]\naccounts: [{name: cf, provider: cloudflare, protocols: [search]}]\ntenants: [{name: t1, key_daily_cost_quota_micros: 250}]\naccess_keys: [{ak: s-ak, 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 transport = Arc::new(GatedSearch(tokio::sync::Semaphore::new(0))); + let h = OnlineHandler::new(gw_state::SharedConfig::new(cfg, state), transport.clone()); + let ak = h.state().auth.get(&access_key_id("s-ak")).await.unwrap(); + let search = || { + let mut param = ModelParamV2::with_name(Protocol::Search, "web"); + param.typed = Some(gw_models::TypedParams::Search( + gw_models::params::SearchParams { + query: "q".into(), + count: 1, + }, + )); + GatewayRequest { + is_online: true, + model_param_v2: Some(param), + ..Default::default() + } + }; + let first = tokio::spawn({ + let (h, ak) = (h.clone(), ak.clone()); + async move { h.run(search(), ak).await } + }); + let key = format!("cb:ak:{}", ak.ak_id); + tokio::time::timeout(std::time::Duration::from_secs(5), async { + while h.state().governance.quota_used(&key).await == 0 { + tokio::task::yield_now().await; + } + }) + .await + .expect("the first call holds its unit price"); + + let second = tokio::time::timeout( + std::time::Duration::from_secs(5), + h.run(search(), ak.clone()), + ) + .await + .expect("a denial returns without waiting on the upstream"); + + let err = second.err().expect("the in-flight hold leaves no budget"); + assert!( + err.message.contains("daily cost budget exhausted"), + "{}", + err.message + ); + transport.0.add_permits(1); + first.await.unwrap().unwrap(); + let (_, ledger) = h.state().store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!( + (ledger.len(), ledger[0].cost_micros), + (1, 250), + "{ledger:?}" + ); + } + #[derive(Debug)] struct PiiStream; @@ -2732,7 +2942,7 @@ mod tests { "data: {\"choices\":[{\"delta\":{\"content\":\"reach me at jane@corp.com now\"}}]}\n\n", )), Ok(bytes::Bytes::from( - "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}]}\n\n", + "data: {\"choices\":[{\"delta\":{},\"finish_reason\":\"stop\"}],\"usage\":{\"prompt_tokens\":5,\"completion_tokens\":7,\"total_tokens\":12}}\n\n", )), Ok(bytes::Bytes::from("data: [DONE]\n\n")), ]; @@ -2777,7 +2987,7 @@ mod tests { } #[tokio::test] - async fn dlp_partial_replay_bills_only_delivered_text() { + async fn dlp_partial_replay_bills_the_upstream_usage() { let h = OnlineHandler::new(handler().config.clone(), Arc::new(PiiStream)); let mut request = chat_req("gpt-4o", "hello"); request.stream = true; @@ -2793,9 +3003,13 @@ mod tests { .unwrap(); let (_, ledger) = h.state().store.ledger_snapshot(10).await.unwrap(); assert_eq!(ledger.len(), 1); - assert!(ledger[0].estimated); - assert!(ledger[0].completion_tokens > 0); - assert!(ctx.outcome.as_ref().unwrap().response.aborted); + assert!(!ledger[0].estimated, "{:?}", ledger[0]); + assert_eq!( + (ledger[0].prompt_tokens, ledger[0].completion_tokens), + (5, 7) + ); + let response = &ctx.outcome.as_ref().unwrap().response; + assert!(response.aborted && response.message.is_empty()); } #[tokio::test] diff --git a/crates/models/src/response.rs b/crates/models/src/response.rs index 135a353..7370ed4 100644 --- a/crates/models/src/response.rs +++ b/crates/models/src/response.rs @@ -62,6 +62,9 @@ pub struct GatewayResponse { /// raw usage sub-tree from the vendor body/last SSE frame. #[serde(skip)] pub raw_usage: Option, + /// The vendor's id for this generation (OpenRouter `gen-…`), for a post-hoc usage lookup. + #[serde(skip)] + pub upstream_id: String, } impl GatewayResponse { diff --git a/crates/state/Cargo.toml b/crates/state/Cargo.toml index 3e1fd8e..e371a32 100644 --- a/crates/state/Cargo.toml +++ b/crates/state/Cargo.toml @@ -21,6 +21,7 @@ async-trait = { workspace = true } serde_json = { workspace = true } serde = { workspace = true } tokio = { workspace = true } +tokio-util = { workspace = true } futures = { workspace = true } metrics = { workspace = true } chacha20poly1305 = { workspace = true } diff --git a/crates/state/src/admission.rs b/crates/state/src/admission.rs index 1a5d7b4..87ac028 100644 --- a/crates/state/src/admission.rs +++ b/crates/state/src/admission.rs @@ -4,12 +4,17 @@ //! two admission paths cannot drift. Denials carry the user-facing message; //! callers wrap it in their own wire shape. +use std::fs::{File, OpenOptions}; +use std::io::{BufRead, BufReader, Seek, SeekFrom, Write}; +use std::path::{Path, PathBuf}; use std::sync::Arc; use std::time::Duration; use gw_config::{GatewayConfig, TenantConf}; use gw_models::GResult; -use tokio::sync::{mpsc, oneshot}; +use tokio::sync::{Mutex, mpsc, oneshot}; +use tokio_util::sync::{CancellationToken, WaitForCancellationFuture}; +use tokio_util::task::{TaskTracker, task_tracker::TaskTrackerToken}; use crate::store::{BillingInput, BillingRecord, Store, billing_record}; use crate::{AkInfo, GatewayState, Governance, UserBudget, clamp_tokens}; @@ -20,6 +25,9 @@ const LEDGER_RETRY_INITIAL: Duration = Duration::from_millis(100); const LEDGER_RETRY_MAX: Duration = Duration::from_secs(30); // a store that refuses writes must drain the queue rather than wedge it const LEDGER_RETRY_ATTEMPTS: usize = 8; +// with a spool, rows leave the queue sooner, so an outage does not back up into requests +const LEDGER_SPOOL_AFTER_ATTEMPTS: usize = 3; +const LEDGER_REPLAY_INTERVAL: Duration = Duration::from_secs(30); // a month's counter must outlive the next month's rollover read const MONTH_COUNTER_TTL: Duration = Duration::from_secs(62 * 86_400); @@ -38,6 +46,50 @@ pub struct TpmReserve { pub window: Option, } +/// Tokens and cost a call charges the budgets: the admission estimate, or the settled actuals. +#[derive(Debug, Clone, Copy, Default)] +pub struct BudgetSpend { + pub tokens: i64, + pub cost_micros: i64, +} + +impl BudgetSpend { + fn of(self, scope: BudgetScope) -> i64 { + if scope.charges_cost() { + self.cost_micros + } else { + self.tokens + } + } +} + +/// Day-budget amounts held at admission, so concurrent calls count before they settle. +#[derive(Debug, Default)] +pub struct BudgetHold { + held: Vec<(String, i64)>, + at: i64, +} + +impl BudgetHold { + /// Return every held amount, for a call that never reached billing. + pub async fn refund(self, gov: &dyn Governance) { + if self.held.is_empty() { + return; + } + let refunds: Vec<_> = self + .held + .iter() + .map(|(key, amount)| (key.as_str(), -amount)) + .collect(); + gov.quota_settle_all(&refunds, self.at).await; + } + + fn take(&mut self, key: &str) -> Option { + let i = self.held.iter().position(|(k, _)| k == key)?; + Some(self.held.swap_remove(i).1) + } +} + /// One settled call: identity + reserves to close. pub struct SettleInput<'a> { pub billing: BillingInput<'a>, @@ -53,12 +105,20 @@ pub struct SettleInput<'a> { /// The key whose budgets the call charges, and the user override admission resolved. pub budget_ak: &'a AkInfo, pub user_budget: Option, + /// Day-budget amounts admission held; the settle charges the difference. + pub budget_hold: BudgetHold, } -/// What a settled request cost, for the decision trail. -pub struct Settled { - pub total_tokens: i64, - pub cost_micros: i64, +/// A settle in progress: shutdown waits for it to drop and fires `closing` so it bills without waiting. +pub struct SettleGuard { + _tracked: TaskTrackerToken, + closing: CancellationToken, +} + +impl SettleGuard { + pub fn closing(&self) -> WaitForCancellationFuture<'_> { + self.closing.cancelled() + } } #[derive(Debug, Clone)] @@ -66,6 +126,8 @@ pub(crate) struct BillingLedger { store: Arc, queue: Option>, deferred: bool, + settles: TaskTracker, + closing: CancellationToken, } impl BillingLedger { @@ -74,13 +136,29 @@ impl BillingLedger { store, queue: None, deferred: false, + settles: TaskTracker::new(), + closing: CancellationToken::new(), } } - pub(crate) fn repairing(store: Arc) -> Self { + /// A batching writer; with `spool_path`, rows the store keeps refusing wait there instead of being dropped. + pub(crate) fn repairing(store: Arc, spool_path: &str) -> Self { let (queue, mut pending) = mpsc::channel::(LEDGER_QUEUE_CAPACITY); let deferred = store.defers_ledger_writes(); let worker_store = store.clone(); + let spool = (!spool_path.is_empty()).then(|| Arc::new(LedgerSpool::new(spool_path))); + if let Some(spool) = spool.clone() { + let replay_store = store.clone(); + // detached: replays the spool for the life of the process + tokio::spawn(async move { + let mut tick = tokio::time::interval(LEDGER_REPLAY_INTERVAL); + tick.set_missed_tick_behavior(tokio::time::MissedTickBehavior::Delay); + loop { + tick.tick().await; + spool.replay(&replay_store).await; + } + }); + } // the bounded worker owns accepted rows through caller cancellation tokio::spawn(async move { let mut batch = Vec::with_capacity(LEDGER_BATCH_MAX); @@ -96,7 +174,7 @@ impl BillingLedger { next.take(&mut batch, &mut row_acks, &mut ack); } let unacked = (batch.len() - row_acks.len()) as u64; - let committed = Self::commit(&worker_store, &mut batch).await; + let committed = Self::commit(&worker_store, spool.as_deref(), &mut batch).await; if !committed { dropped += unacked; } @@ -111,9 +189,16 @@ impl BillingLedger { } }); Self { - store, queue: Some(queue), deferred, + ..Self::direct(store) + } + } + + pub(crate) fn begin_settle(&self) -> SettleGuard { + SettleGuard { + _tracked: self.settles.token(), + closing: self.closing.clone(), } } @@ -175,7 +260,11 @@ impl BillingLedger { rx.await.unwrap_or(false) } - async fn commit(store: &Arc, batch: &mut Vec) -> bool { + async fn commit( + store: &Arc, + spool: Option<&LedgerSpool>, + batch: &mut Vec, + ) -> bool { if batch.is_empty() { return true; } @@ -187,6 +276,20 @@ impl BillingLedger { break; }; metrics::counter!("gateway_ledger_write_failures_total").increment(batch.len() as u64); + if let Some(spool) = spool + && attempt >= LEDGER_SPOOL_AFTER_ATTEMPTS + { + match spool.append(batch).await { + Ok(()) => { + tracing::warn!(error = %e, rows = batch.len(), "billing ledger write failed; rows spooled for replay"); + committed = true; + break; + } + Err(spool_err) => { + tracing::error!(error = %spool_err, rows = batch.len(), "billing spool write failed") + } + } + } if attempt == LEDGER_RETRY_ATTEMPTS { tracing::error!(error = %e, rows = batch.len(), "billing ledger retries exhausted; rows dropped"); break; @@ -200,7 +303,115 @@ impl BillingLedger { } } -/// One budget: its window, governance counter and cap. +struct LedgerSpool { + path: PathBuf, + replaying: PathBuf, + replayed: PathBuf, + lock: Mutex<()>, +} + +impl LedgerSpool { + fn new(path: &str) -> Self { + let path = PathBuf::from(path); + Self { + replaying: path.with_extension("replaying"), + replayed: path.with_extension("replayed"), + path, + lock: Mutex::new(()), + } + } + + async fn append(&self, rows: &[BillingRecord]) -> std::io::Result<()> { + // a leading newline ends a row a crash tore + let mut lines = vec![b'\n']; + for r in rows { + serde_json::to_writer(&mut lines, r)?; + lines.push(b'\n'); + } + let path = self.path.clone(); + let _held = self.lock.lock().await; + run_blocking(move || { + if let Some(dir) = path.parent() { + std::fs::create_dir_all(dir)?; + } + let mut f = OpenOptions::new().create(true).append(true).open(&path)?; + f.write_all(&lines)?; + f.sync_data() + }) + .await + } + + /// Resumes after the last batch it saved; `request_id` uniqueness covers one committed but not saved. + async fn replay(&self, store: &Arc) { + let claimed = { + let (path, replaying, replayed) = ( + self.path.clone(), + self.replaying.clone(), + self.replayed.clone(), + ); + let _held = self.lock.lock().await; + run_blocking(move || { + if replaying.exists() { + let text = std::fs::read_to_string(&replayed).unwrap_or_default(); + return Ok(Some(text.trim().parse().unwrap_or(0))); + } + if !path.exists() { + return Ok(None); + } + remove_if_present(&replayed)?; + std::fs::rename(&path, &replaying).map(|()| Some(0)) + }) + .await + }; + let mut offset = match claimed { + Ok(Some(offset)) => offset, + Ok(None) => return, + Err(e) => { + tracing::error!(error = %e, path = %self.path.display(), "billing spool claim failed"); + return; + } + }; + let mut total = 0; + loop { + let replaying = self.replaying.clone(); + let (rows, next) = match run_blocking(move || read_spool_batch(&replaying, offset)) + .await + { + Ok(batch) => batch, + Err(e) => { + tracing::error!(error = %e, path = %self.replaying.display(), "billing spool read failed"); + return; + } + }; + if next == offset { + break; + } + if let Err(e) = store.ledger_replay_batch(&rows).await { + tracing::warn!(error = %e, rows = rows.len(), "billing spool replay deferred"); + return; + } + total += rows.len(); + offset = next; + let offset_file = self.replayed.clone(); + if let Err(e) = + run_blocking(move || std::fs::write(offset_file, offset.to_string())).await + { + tracing::warn!(error = %e, "billing spool offset save failed"); + } + } + let (replaying, offset_file) = (self.replaying.clone(), self.replayed.clone()); + if let Err(e) = run_blocking(move || { + std::fs::remove_file(replaying).and_then(|()| remove_if_present(&offset_file)) + }) + .await + { + tracing::error!(error = %e, path = %self.replaying.display(), "billing spool remove failed"); + return; + } + tracing::info!(rows = total, "billing spool replayed"); + } +} + struct Budget { scope: BudgetScope, window: Window, @@ -292,42 +503,101 @@ impl LedgerWrite { } } -/// Commit billing rows the batching writer still holds, before the process exits. +/// Finish in-flight settles, then commit the rows the batching writer holds, before the process exits. pub async fn flush_billing(state: &GatewayState) { + state.billing.closing.cancel(); + state.billing.settles.close(); + state.billing.settles.wait().await; let dropped = state.billing.flush().await; if dropped > 0 { tracing::error!(rows = dropped, "billing rows dropped before shutdown"); } } -/// Daily and monthly budgets (soft caps: check-then-consume, so concurrent -/// turns can overshoot by one): admit while every configured scope is under its cap. -/// Admitted yields the user's override for settlement; the outer error is a failed lookup. -pub async fn check_budgets( +/// Hold the call's estimate against each day budget while spent plus held stays under its cap; month budgets only check. +pub async fn reserve_budgets( state: &GatewayState, cfg: &GatewayConfig, ak: &AkInfo, user: &str, -) -> GResult> { + est: impl FnOnce() -> BudgetSpend, + at: i64, +) -> GResult> { let over = user_budget(state, ak, user).await?; let gov = state.governance.as_ref(); - let budgets = budgets(gov, cfg, ak, user, over).await; - let under = futures::future::join_all(budgets.iter().map(|b| async move { - match b.window { - Window::Day => gov.quota_check(&b.key, b.limit).await, - Window::Month => gov.counter_get(&b.key).await < b.limit, - } - })) - .await; - if let Some((b, _)) = budgets.iter().zip(under).find(|(_, under)| !under) { - return Ok(Err(format!( - "{} {} budget exhausted for {}", - b.window.label(), - b.scope.unit(), - b.scope.subject(ak, user) + let mut budgets = budgets(gov, cfg, ak, user, over).await; + if budgets.is_empty() { + return Ok(Ok(( + over, + BudgetHold { + held: Vec::new(), + at, + }, ))); } - Ok(Ok(over)) + let est = if budgets.iter().any(|b| b.window == Window::Day) { + est() + } else { + BudgetSpend::default() + }; + let holds = |b: &Budget| b.window == Window::Day && est.of(b.scope) > 0; + let hold_args: Vec<_> = budgets + .iter() + .filter(|b| holds(b)) + .map(|b| (b.key.as_str(), est.of(b.scope), b.limit)) + .collect(); + let (reserved, checked) = tokio::join!( + async { + if hold_args.is_empty() { + Ok(true) + } else { + gov.quota_reserve_all(&hold_args, at).await + } + }, + futures::future::join_all(budgets.iter().map(|b| async move { + match b.window { + _ if holds(b) => true, + Window::Day => gov.quota_check(&b.key, b.limit).await, + Window::Month => gov.counter_get(&b.key).await < b.limit, + } + })), + ); + let refused = reserved + .err() + .and_then(|n| budgets.iter().enumerate().filter(|(_, b)| holds(b)).nth(n)) + .map(|(i, _)| i); + let Some(i) = checked + .iter() + .position(|ok| !ok) + .into_iter() + .chain(refused) + .min() + else { + let held = if reserved == Ok(true) { + budgets + .iter_mut() + .filter(|b| holds(b)) + .map(|b| (std::mem::take(&mut b.key), est.of(b.scope))) + .collect() + } else { + Vec::new() + }; + return Ok(Ok((over, BudgetHold { held, at }))); + }; + if reserved == Ok(true) { + let refunds: Vec<_> = hold_args + .iter() + .map(|&(key, amount, _)| (key, -amount)) + .collect(); + gov.quota_settle_all(&refunds, at).await; + } + let b = &budgets[i]; + Ok(Err(format!( + "{} {} budget exhausted for {}", + b.window.label(), + b.scope.unit(), + b.scope.subject(ak, user) + ))) } pub fn model_quota_key(ak_id: &str, model: &str) -> String { @@ -479,12 +749,12 @@ pub async fn settle_and_bill( state: &GatewayState, cfg: &GatewayConfig, s: SettleInput<'_>, -) -> Settled { +) -> BudgetSpend { let gov = state.governance.as_ref(); let total = clamp_tokens(s.billing.total); let record = billing_record(cfg, &s.billing); - let settled = Settled { - total_tokens: record.total_tokens, + let settled = BudgetSpend { + tokens: record.total_tokens, cost_micros: record.cost_micros, }; let settle_daily = gov.quota_settle(s.billing.ak, total - s.reserved, s.reserved_at); @@ -506,8 +776,8 @@ pub async fn settle_and_bill( s.budget_ak, s.billing.user_id, s.user_budget, - settled.total_tokens, - settled.cost_micros, + s.budget_hold, + settled, ); let write_ledger = state.billing.write(record); tokio::join!( @@ -526,18 +796,20 @@ pub fn month_prefixes() -> [String; 2] { [month_prefix(month), month_prefix(previous_month(month))] } -/// Accrue usage under `over`, the override admission resolved (`None` resolves it now); -/// a scope reaching its cap raises a `budget_exhausted` alert (deduped by the bus). +/// Accrue usage under `over`, the override admission resolved (`None` resolves it now), settling +/// what `hold` held; a scope reaching its cap raises a `budget_exhausted` alert (deduped by the bus). async fn consume_budgets( state: &GatewayState, cfg: &GatewayConfig, ak: &AkInfo, user: &str, over: Option, - tokens: i64, - cost_micros: i64, + mut hold: BudgetHold, + spent: BudgetSpend, ) { - if tokens <= 0 && cost_micros <= 0 { + let gov = state.governance.as_ref(); + if spent.tokens <= 0 && spent.cost_micros <= 0 { + hold.refund(gov).await; return; } let over = match over { @@ -547,27 +819,48 @@ async fn consume_budgets( UserBudget::default() }), }; - let gov = state.governance.as_ref(); let budgets = budgets(gov, cfg, ak, user, over).await; - let used = futures::future::join_all(budgets.iter().map(|b| async move { - let amount = if b.scope.charges_cost() { - cost_micros - } else { - tokens - }; - if amount <= 0 { - return None; + if budgets.is_empty() { + return; + } + let (mut settles, mut deltas, mut accrues) = (Vec::new(), Vec::new(), Vec::new()); + for (i, b) in budgets.iter().enumerate() { + let amount = spent.of(b.scope); + match (b.window, hold.take(&b.key)) { + (Window::Day, Some(held)) => { + settles.push(i); + deltas.push((b.key.as_str(), amount - held)); + } + _ if amount <= 0 => {} + _ => accrues.push(i), } - Some(match b.window { - Window::Day => gov.quota_consume(&b.key, amount).await, - Window::Month => gov.counter_add(&b.key, amount, MONTH_COUNTER_TTL).await, - }) - })) - .await; - for (b, used) in budgets.iter().zip(used) { - if let Some(used) = used - && used >= b.limit - { + } + let (settled, accrued) = tokio::join!( + async { + if deltas.is_empty() { + Vec::new() + } else { + gov.quota_settle_all(&deltas, hold.at).await + } + }, + futures::future::join_all(accrues.iter().map(|&i| { + let b = &budgets[i]; + let amount = spent.of(b.scope); + async move { + match b.window { + Window::Day => gov.quota_consume(&b.key, amount).await, + Window::Month => gov.counter_add(&b.key, amount, MONTH_COUNTER_TTL).await, + } + } + })), + ); + let totals = settles + .iter() + .zip(settled) + .chain(accrues.iter().zip(accrued)); + for (&i, used) in totals { + let b = &budgets[i]; + if used >= b.limit { state.alerts.emit( "budget_exhausted", b.scope.subject(ak, user), @@ -682,6 +975,46 @@ fn month_prefix((y, m): (i64, u32)) -> String { format!("m:{y}{m:02}:") } +/// Up to `LEDGER_BATCH_MAX` rows from `offset`, and the offset after them; equal offsets mean the end. +fn read_spool_batch(path: &Path, offset: u64) -> std::io::Result<(Vec, u64)> { + let mut reader = BufReader::new(File::open(path)?); + reader.seek(SeekFrom::Start(offset))?; + let (mut rows, mut next, mut line) = (Vec::new(), offset, Vec::new()); + while rows.len() < LEDGER_BATCH_MAX { + line.clear(); + let n = reader.read_until(b'\n', &mut line)?; + if n == 0 { + break; + } + next += n as u64; + if line.trim_ascii().is_empty() { + continue; + } + match serde_json::from_slice(&line) { + Ok(row) => rows.push(row), + Err(e) => tracing::error!(error = %e, "skipping a torn billing spool line"), + } + } + Ok((rows, next)) +} + +fn remove_if_present(path: &Path) -> std::io::Result<()> { + match std::fs::remove_file(path) { + Err(e) if e.kind() != std::io::ErrorKind::NotFound => Err(e), + _ => Ok(()), + } +} + +async fn run_blocking(f: F) -> std::io::Result +where + T: Send + 'static, + F: FnOnce() -> std::io::Result + Send + 'static, +{ + tokio::task::spawn_blocking(f) + .await + .unwrap_or_else(|e| Err(std::io::Error::other(e))) +} + #[cfg(test)] mod tests { use gw_config::access_key_id; @@ -744,7 +1077,7 @@ mod tests { async fn billing_ledger_repairs_failed_writes() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(2); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); let record = record("req-repair"); ledger.write(record.clone()).await; let (count, rows) = tokio::time::timeout(Duration::from_secs(2), async { @@ -767,7 +1100,7 @@ mod tests { let store = Arc::new(crate::MemoryStore::default()); let ledger = BillingLedger { deferred: true, - ..BillingLedger::repairing(store.clone()) + ..BillingLedger::repairing(store.clone(), "") }; let mut row = record("req-attributed"); row.user_id = "user-42".into(); @@ -785,7 +1118,7 @@ mod tests { store.fail_next_ledger_writes(LEDGER_RETRY_ATTEMPTS); let ledger = BillingLedger { deferred: true, - ..BillingLedger::repairing(store.clone()) + ..BillingLedger::repairing(store.clone(), "") }; let mut row = record("req-fallback"); row.user_id = "user-42".into(); @@ -801,7 +1134,7 @@ mod tests { async fn canceled_attributed_write_is_reported_as_dropped() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(LEDGER_RETRY_ATTEMPTS); - let ledger = BillingLedger::repairing(store); + let ledger = BillingLedger::repairing(store, ""); let (tx, rx) = oneshot::channel(); drop(rx); @@ -820,7 +1153,7 @@ mod tests { async fn billing_ledger_backpressures_at_capacity_then_repairs_every_row() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(usize::MAX); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); let queue = ledger.queue.as_ref().unwrap(); ledger.write(record("req-0")).await; @@ -868,7 +1201,7 @@ mod tests { async fn flush_returns_only_after_queued_rows_commit() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(1); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); ledger.write(record("req-flush")).await; ledger.flush().await; let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); @@ -880,7 +1213,7 @@ mod tests { async fn billing_ledger_drains_a_store_that_never_accepts() { let store = Arc::new(crate::MemoryStore::default()); store.fail_next_ledger_writes(usize::MAX); - let ledger = BillingLedger::repairing(store.clone()); + let ledger = BillingLedger::repairing(store.clone(), ""); tokio::time::timeout(Duration::from_secs(3_600), async { for i in 0..=LEDGER_QUEUE_CAPACITY { @@ -905,6 +1238,176 @@ mod tests { ); } + #[tokio::test] + async fn rows_the_store_keeps_refusing_are_spooled_not_dropped() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("ledger-spool.jsonl"); + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(usize::MAX); + let ledger = BillingLedger::repairing(store.clone(), path.to_str().unwrap()); + + ledger.write(record("req-spooled")).await; + + assert_eq!( + ledger.flush().await, + 0, + "a spooled row is not a dropped row" + ); + let spooled = std::fs::read_to_string(&path).unwrap(); + assert_eq!(spooled.lines().filter(|l| !l.is_empty()).count(), 1); + assert!(spooled.contains("req-spooled")); + } + + #[tokio::test] + async fn spool_replays_into_the_store_once_it_accepts_and_retries_otherwise() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + spool + .append(&[record("req-a"), record("req-b")]) + .await + .unwrap(); + + store.fail_next_ledger_writes(1); + spool.replay(&as_store).await; + assert_eq!(store.ledger_snapshot(usize::MAX).await.unwrap().0, 0); + assert!(spool.replaying.exists(), "a refused replay keeps its rows"); + spool.append(&[record("req-c")]).await.unwrap(); + + spool.replay(&as_store).await; + spool.replay(&as_store).await; + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + let mut ids: Vec<_> = rows.iter().map(|r| r.request_id.as_str()).collect(); + ids.sort_unstable(); + assert_eq!((count, ids), (3, vec!["req-a", "req-b", "req-c"])); + assert!(!spool.replaying.exists() && !spool.path.exists()); + } + + #[tokio::test] + async fn a_torn_spool_row_never_swallows_the_next_one() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + std::fs::write(&spool.path, b"{\"ak\":\"torn").unwrap(); + spool.append(&[record("req-after-crash")]).await.unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + + spool.replay(&as_store).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!(count, 1); + assert_eq!(rows[0].request_id, "req-after-crash"); + } + + #[tokio::test] + async fn an_unwritable_spool_keeps_retrying_the_store() { + let dir = tempfile::tempdir().unwrap(); + let blocker = dir.path().join("not-a-dir"); + std::fs::write(&blocker, b"").unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + store.fail_next_ledger_writes(LEDGER_SPOOL_AFTER_ATTEMPTS + 1); + let path = blocker.join("ledger-spool.jsonl"); + let ledger = BillingLedger::repairing(store.clone(), path.to_str().unwrap()); + + ledger.write(record("req-unspooled")).await; + + assert_eq!(ledger.flush().await, 0); + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!((count, rows[0].request_id.as_str()), (1, "req-unspooled")); + } + + #[tokio::test] + async fn a_spool_dir_that_does_not_exist_yet_is_created() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("gw/ledger/spool.jsonl").to_str().unwrap()); + spool.append(&[record("req-nested")]).await.unwrap(); + assert!( + std::fs::read_to_string(&spool.path) + .unwrap() + .contains("req-nested") + ); + } + + #[tokio::test] + async fn replay_streams_batches_and_rolls_up_rows_billed_before_the_watermark() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + let now = crate::epoch_secs(); + let mut late = record("req-current"); + late.created_at_epoch_secs = now - 600; + store.ledger_add(&late).await.unwrap(); + store.usage_rollup_advance(now).await.unwrap(); + let billed_at = now - 7_200; + let rows: Vec<_> = (0..=LEDGER_BATCH_MAX) + .map(|i| BillingRecord { + created_at_epoch_secs: billed_at, + ..record(format!("req-late-{i}")) + }) + .collect(); + spool.append(&rows).await.unwrap(); + + spool.replay(&as_store).await; + + assert_eq!( + store.ledger_snapshot(usize::MAX).await.unwrap().0, + LEDGER_BATCH_MAX + 2 + ); + let usage = store + .usage_by_user(None, None, billed_at, billed_at + 60) + .await + .unwrap(); + assert_eq!( + usage.iter().map(|u| u.requests).sum::(), + LEDGER_BATCH_MAX as i64 + 1, + "{usage:?}" + ); + assert!(!spool.replaying.exists()); + } + + #[tokio::test] + async fn replay_resumes_after_the_last_saved_batch() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + let mut lines = Vec::new(); + for id in ["req-committed", "req-pending"] { + serde_json::to_writer(&mut lines, &record(id)).unwrap(); + lines.push(b'\n'); + } + let committed = lines.iter().position(|&b| b == b'\n').unwrap() + 1; + std::fs::write(&spool.replaying, &lines).unwrap(); + std::fs::write(&spool.replayed, committed.to_string()).unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + + spool.replay(&as_store).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!((count, rows[0].request_id.as_str()), (1, "req-pending")); + assert!(!spool.replaying.exists() && !spool.replayed.exists()); + } + + #[tokio::test] + async fn a_spool_row_from_an_older_release_still_replays() { + let dir = tempfile::tempdir().unwrap(); + let spool = LedgerSpool::new(dir.path().join("ledger-spool.jsonl").to_str().unwrap()); + std::fs::write( + &spool.path, + br#"{"ak":"ak","tenant":"default","request_id":"req-old","created_at_epoch_secs":1,"model":"m","total_tokens":2}"#, + ) + .unwrap(); + let store = Arc::new(crate::MemoryStore::default()); + let as_store: Arc = store.clone(); + + spool.replay(&as_store).await; + + let (count, rows) = store.ledger_snapshot(usize::MAX).await.unwrap(); + assert_eq!((count, rows[0].request_id.as_str()), (1, "req-old")); + assert_eq!(rows[0].total_tokens, 2); + } + #[test] fn civil_month_follows_the_gregorian_calendar() { assert_eq!(civil_month(0), (1970, 1)); @@ -918,6 +1421,14 @@ mod tests { assert_eq!(previous_month((2026, 9)), (2026, 8)); } + async fn daily_fixture() -> (Arc, Arc, Arc) { + let yaml = "listen: {host: h, port: 1}\ntenants: [{name: t1, key_daily_cost_quota_micros: 100}]\naccess_keys: [{ak: k1, tenant: t1, product: p, qps: 1, daily_token_quota: 1000}]"; + let cfg = Arc::new(GatewayConfig::from_yaml(yaml).unwrap()); + let state = Arc::new(GatewayState::from_config(&cfg)); + let ak = state.auth.get(&access_key_id("k1")).await.unwrap(); + (cfg, state, ak) + } + async fn monthly_fixture( rollover: bool, ) -> (Arc, Arc, Arc) { @@ -974,13 +1485,95 @@ mod tests { ); } + #[tokio::test] + async fn daily_budgets_hold_in_flight_calls_until_they_settle() { + let (cfg, state, ak) = daily_fixture().await; + let gov = state.governance.as_ref(); + let key = format!("cb:ak:{}", ak.ak_id); + let est = BudgetSpend { + tokens: 1, + cost_micros: 60, + }; + let now = crate::epoch_secs(); + let reserve = || reserve_budgets(&state, &cfg, &ak, "", || est, now); + + let (_, first) = reserve().await.unwrap().unwrap(); + let (_, second) = reserve().await.unwrap().unwrap(); + assert_eq!(gov.quota_used(&key).await, 120, "both calls are in flight"); + let denied = reserve().await.unwrap().unwrap_err(); + assert!( + denied.starts_with("daily cost budget exhausted for key:"), + "{denied}" + ); + assert_eq!(gov.quota_used(&key).await, 120, "a denial holds nothing"); + + let spent = BudgetSpend { + tokens: 1, + cost_micros: 10, + }; + consume_budgets(&state, &cfg, &ak, "", None, first, spent).await; + assert_eq!(gov.quota_used(&key).await, 70, "the settle charges actuals"); + second.refund(gov).await; + assert_eq!(gov.quota_used(&key).await, 10, "a refund returns the hold"); + let (_, third) = reserve().await.unwrap().unwrap(); + consume_budgets(&state, &cfg, &ak, "", None, third, BudgetSpend::default()).await; + assert_eq!( + gov.quota_used(&key).await, + 10, + "a zero-cost settle refunds the hold" + ); + } + + #[tokio::test] + async fn a_settle_that_matches_its_hold_still_alerts_at_the_cap() { + let (cfg, state, ak) = daily_fixture().await; + let mut alerts = state.alerts.take_receiver().unwrap(); + let spend = BudgetSpend { + tokens: 1, + cost_micros: 100, + }; + let (_, hold) = reserve_budgets(&state, &cfg, &ak, "", || spend, crate::epoch_secs()) + .await + .unwrap() + .unwrap(); + + consume_budgets(&state, &cfg, &ak, "", None, hold, spend).await; + + let alert = alerts.try_recv().expect("budget_exhausted alert"); + assert_eq!(alert.kind, "budget_exhausted"); + } + + #[tokio::test] + async fn without_a_day_budget_admission_never_prices_the_estimate() { + let (cfg, state, ak) = monthly_fixture(false).await; + reserve_budgets(&state, &cfg, &ak, "", || unreachable!("no day budget"), 0) + .await + .unwrap() + .unwrap(); + } + #[tokio::test] async fn monthly_budget_denies_at_the_cap_without_touching_daily_counters() { let (cfg, state, ak) = monthly_fixture(false).await; let gov = state.governance.as_ref(); - check_budgets(&state, &cfg, &ak, "").await.unwrap().unwrap(); - consume_budgets(&state, &cfg, &ak, "", None, 10, 3).await; - let err = check_budgets(&state, &cfg, &ak, "") + reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default, 0) + .await + .unwrap() + .unwrap(); + consume_budgets( + &state, + &cfg, + &ak, + "", + None, + BudgetHold::default(), + BudgetSpend { + tokens: 10, + cost_micros: 3, + }, + ) + .await; + let err = reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default, 0) .await .unwrap() .unwrap_err(); @@ -991,7 +1584,10 @@ mod tests { assert_eq!(gov.quota_used(&format!("cb:ak:{}", ak.ak_id)).await, 0); gov.quota_reset_all().await; assert!( - check_budgets(&state, &cfg, &ak, "").await.unwrap().is_err(), + reserve_budgets(&state, &cfg, &ak, "", BudgetSpend::default, 0) + .await + .unwrap() + .is_err(), "the daily reset leaves month counters alone" ); } diff --git a/crates/state/src/governance.rs b/crates/state/src/governance.rs index 169edc4..fa82e93 100644 --- a/crates/state/src/governance.rs +++ b/crates/state/src/governance.rs @@ -35,9 +35,16 @@ pub trait Governance: Send + Sync + std::fmt::Debug { limit: i64, at_epoch_secs: i64, ) -> Option; + /// [`Self::quota_reserve`] over every hold at once, all or none: `Err(i)` = `holds[i]` is at its cap, `Ok(false)` = admitted unrecorded. + async fn quota_reserve_all( + &self, + holds: &[(&str, i64, i64)], + at_epoch_secs: i64, + ) -> Result; /// Apply the settle delta (actual - reserved; negative refunds) to the day /// bucket the paired reserve used (`at_epoch_secs`). async fn quota_settle(&self, key: &str, delta: i64, at_epoch_secs: i64); + async fn quota_settle_all(&self, deltas: &[(&str, i64)], at_epoch_secs: i64) -> Vec; /// Tokens spent today under `key`. async fn quota_used(&self, key: &str) -> i64; /// Add to `key`'s spent tokens; returns the day's new total. @@ -105,9 +112,26 @@ impl Governance for MemoryGovernance { async fn quota_reserve(&self, key: &str, amount: i64, limit: i64, _at: i64) -> Option { self.quota.reserve(key, amount, limit).then_some(amount) } + async fn quota_reserve_all(&self, holds: &[(&str, i64, i64)], _at: i64) -> Result { + for (i, &(key, amount, limit)) in holds.iter().enumerate() { + if !self.quota.reserve(key, amount, limit) { + for &(key, amount, _) in &holds[..i] { + self.quota.settle(key, -amount); + } + return Err(i); + } + } + Ok(true) + } async fn quota_settle(&self, key: &str, delta: i64, _at: i64) { self.quota.settle(key, delta); } + async fn quota_settle_all(&self, deltas: &[(&str, i64)], _at: i64) -> Vec { + deltas + .iter() + .map(|&(key, delta)| self.quota.settle(key, delta)) + .collect() + } async fn quota_used(&self, key: &str) -> i64 { self.quota.used(key) } @@ -216,6 +240,30 @@ impl RedisGovernance { } } + async fn incr_all(&self, incrs: &[(String, i64)]) -> redis::RedisResult> { + if incrs.is_empty() { + return Ok(Vec::new()); + } + let mut pipe = redis::pipe(); + for (key, by) in incrs { + pipe.cmd("INCRBY").arg(key).arg(*by); + } + let mut conn = self.conn.clone(); + let totals: Vec = pipe.query_async(&mut conn).await?; + let mut arm = redis::pipe(); + let mut fresh = false; + for ((key, by), total) in incrs.iter().zip(&totals) { + if total == by { + arm.cmd("PEXPIRE").arg(key).arg(QUOTA_TTL_MS).ignore(); + fresh = true; + } + } + if fresh { + arm.query_async::<()>(&mut conn).await?; + } + Ok(totals) + } + /// Increment `key` and set its TTL on first use; returns the post-increment /// count. A failed round-trip returns 0 so limits fail open rather than /// wedging the gateway on a Redis blip. @@ -275,6 +323,34 @@ impl Governance for RedisGovernance { } }) } + async fn quota_reserve_all(&self, holds: &[(&str, i64, i64)], at: i64) -> Result { + let incrs: Vec<_> = holds + .iter() + .map(|&(key, amount, _)| (quota_key_at(key, at), amount)) + .collect(); + let totals = match self.incr_all(&incrs).await { + Ok(totals) => totals, + Err(e) => { + tracing::warn!(error = %e, "redis budget reserve failed; admitting"); + return Ok(false); + } + }; + let Some(i) = holds + .iter() + .zip(&totals) + .position(|(&(_, amount, limit), &v)| v - amount >= limit) + else { + return Ok(true); + }; + let mut undo = redis::pipe(); + for (key, amount) in &incrs { + undo.cmd("DECRBY").arg(key).arg(*amount).ignore(); + } + if let Err(e) = undo.query_async::<()>(&mut self.conn.clone()).await { + tracing::warn!(error = %e, "redis budget rollback failed"); + } + Err(i) + } async fn quota_settle(&self, key: &str, delta: i64, at: i64) { if delta == 0 { return; @@ -288,6 +364,32 @@ impl Governance for RedisGovernance { ) .await; } + async fn quota_settle_all(&self, deltas: &[(&str, i64)], at: i64) -> Vec { + let incrs: Vec<_> = deltas + .iter() + .map(|&(key, delta)| (quota_key_at(key, at), delta)) + .collect(); + let mut totals = match self.incr_all(&incrs).await { + Ok(totals) => totals, + Err(e) => { + tracing::warn!(error = %e, "redis budget settle failed"); + return vec![0; deltas.len()]; + } + }; + let mut floor = redis::pipe(); + let mut floored = false; + for ((key, _), total) in incrs.iter().zip(totals.iter_mut()) { + if *total < 0 { + floor.cmd("SET").arg(key).arg(0).arg("KEEPTTL").ignore(); + *total = 0; + floored = true; + } + } + if floored && let Err(e) = floor.query_async::<()>(&mut self.conn.clone()).await { + tracing::warn!(error = %e, "redis budget floor failed"); + } + totals + } async fn quota_consume(&self, key: &str, tokens: i64) -> i64 { self.incr_window( "a_key(key), @@ -409,6 +511,75 @@ mod tests { .unwrap(); } + async fn exercise_reserve_all(g: &dyn Governance, ns: &str) { + let (a, b, c) = (format!("ra{ns}"), format!("rb{ns}"), format!("rc{ns}")); + let now = crate::epoch_secs(); + let pair = [(a.as_str(), 60, 100), (b.as_str(), 30, 50)]; + assert_eq!(g.quota_reserve_all(&pair, now).await, Ok(true)); + assert_eq!(g.quota_reserve_all(&pair, now).await, Ok(true)); + assert_eq!((g.quota_used(&a).await, g.quota_used(&b).await), (120, 60)); + assert_eq!( + g.quota_reserve_all(&[(a.as_str(), 1, 100), (b.as_str(), 1, 50)], now) + .await, + Err(0) + ); + assert_eq!( + g.quota_reserve_all(&[(c.as_str(), 5, 100), (b.as_str(), 1, 50)], now) + .await, + Err(1) + ); + assert_eq!( + ( + g.quota_used(&a).await, + g.quota_used(&b).await, + g.quota_used(&c).await + ), + (120, 60, 0) + ); + assert_eq!( + g.quota_settle_all(&[(a.as_str(), -200), (b.as_str(), -10)], now) + .await, + vec![0, 50] + ); + assert_eq!(g.quota_settle_all(&[(b.as_str(), 0)], now).await, vec![50]); + assert_eq!(g.quota_reserve_all(&[], now).await, Ok(true)); + assert!(g.quota_settle_all(&[], now).await.is_empty()); + } + + #[tokio::test] + async fn memory_reserve_all_holds_all_or_none() { + exercise_reserve_all(&MemoryGovernance::default(), "").await; + } + + #[tokio::test] + async fn redis_reserve_all_holds_all_or_none() { + let Ok(url) = std::env::var("GW_TEST_REDIS_URL") else { + return; + }; + let g = RedisGovernance::connect(&url).await.expect("redis connect"); + let ns = format!("-{}", std::process::id()); + exercise_reserve_all(&g, &ns).await; + let now = crate::epoch_secs(); + let keys: Vec<_> = ["ra", "rb", "rc"] + .iter() + .map(|k| quota_key_at(&format!("{k}{ns}"), now)) + .collect(); + let mut conn = g.conn.clone(); + for key in &keys { + let ttl: i64 = redis::cmd("PTTL") + .arg(key) + .query_async(&mut conn) + .await + .unwrap(); + assert!(ttl > 0, "{key} ttl {ttl}"); + } + let _: i64 = redis::cmd("DEL") + .arg(&keys) + .query_async(&mut conn) + .await + .unwrap(); + } + #[tokio::test] async fn redis_governance_enforces_limits() { let Ok(url) = std::env::var("GW_TEST_REDIS_URL") else { diff --git a/crates/state/src/lib.rs b/crates/state/src/lib.rs index 53d6cdf..04cf34f 100644 --- a/crates/state/src/lib.rs +++ b/crates/state/src/lib.rs @@ -389,9 +389,11 @@ impl QuotaStore { reserve_on(&mut slot_mut(&self.used, key, || 0), amount, limit) } - /// Apply the settle delta (actual - reserved); never below zero. - pub fn settle(&self, key: &str, delta: i64) { - settle_on(&mut slot_mut(&self.used, key, || 0), delta); + /// Apply the settle delta (actual - reserved); never below zero. Returns the new total. + pub fn settle(&self, key: &str, delta: i64) -> i64 { + let mut e = slot_mut(&self.used, key, || 0); + settle_on(&mut e, delta); + *e } pub fn reset_all(&self) { @@ -865,7 +867,8 @@ impl GatewayState { } } } - state.billing = admission::BillingLedger::repairing(state.store.clone()); + state.billing = + admission::BillingLedger::repairing(state.store.clone(), &st.ledger_spool_path); Ok(state) } @@ -890,6 +893,10 @@ impl GatewayState { streams: prev.streams.clone(), }) } + + pub fn begin_settle(&self) -> admission::SettleGuard { + self.billing.begin_settle() + } } /// A consistent (config, derived-state) pair. A request loads one snapshot and @@ -1075,11 +1082,16 @@ fn settle_on(counter: &mut i64, delta: i64) { #[cfg(test)] pub(crate) async fn scratch_pg(url: &str) -> (String, u128) { + static NEXT: AtomicUsize = AtomicUsize::new(0); let nonce = std::time::SystemTime::now() .duration_since(std::time::UNIX_EPOCH) .unwrap() .as_nanos(); - let db = format!("gwtest_{nonce}"); + let db = format!( + "gwtest_{nonce}_{}_{}", + std::process::id(), + NEXT.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + ); let admin = sqlx::PgPool::connect(url).await.expect("pg admin"); sqlx::query(sqlx::AssertSqlSafe(format!("CREATE DATABASE {db}"))) .execute(&admin) diff --git a/crates/state/src/store.rs b/crates/state/src/store.rs index 908584c..53447e7 100644 --- a/crates/state/src/store.rs +++ b/crates/state/src/store.rs @@ -35,8 +35,8 @@ const ROLLUP_BUCKET_SECS: i64 = 60; const ROLLUP_BACKFILL_SECS: i64 = 20 * 60; /// A minute is rolled only once closed for this long, so no in-flight write or -/// trailing replica lands a row in a rolled minute — a rolled minute's source -/// set can only shrink, which keeps the max-upsert sound. +/// trailing replica lands a row in a rolled minute; a spool replay that does adds +/// the row to the bucket in the same transaction, which keeps the max-upsert sound. const ROLLUP_SETTLE_SECS: i64 = ROLLUP_BUCKET_SECS; /// In-process rollup retention: the monthly-budget window, then a bucket cap. const ROLLUP_RETENTION_SECS: i64 = 62 * 86_400; @@ -60,7 +60,8 @@ const BATCH_ITEM_CHUNK: usize = 13_000; const PG_MAX_CONNECTIONS: u32 = 10; /// One billing entry (recorded locally only; no reporting upstream). -#[derive(Debug, Clone, serde::Serialize)] +#[derive(Debug, Clone, Default, serde::Serialize, serde::Deserialize)] +#[serde(default)] pub struct BillingRecord { pub ak: String, pub product: String, @@ -290,10 +291,10 @@ impl UserUsageRow { self.billed_units = self.billed_units.saturating_add(r.billed_units); } - /// Keep the larger of each counter. A bucket's source rows are append-only - /// and only ever shrink via ledger pruning, so per-column max always picks - /// the more complete snapshot — a recompute from a partially pruned ledger - /// can never shrink a bucket. + /// Keep the larger of each counter. A bucket's source rows only shrink via ledger + /// pruning, and a replayed late row is added to its bucket as it lands, so per-column + /// max always picks the more complete snapshot — a recompute from a partially pruned + /// ledger can never shrink a bucket. fn keep_max(&mut self, o: &UserUsageRow) { self.requests = self.requests.max(o.requests); self.prompt_tokens = self.prompt_tokens.max(o.prompt_tokens); @@ -359,6 +360,8 @@ pub trait Store: Send + Sync + std::fmt::Debug { } Ok(()) } + /// Insert late rows and atomically add the ones inserted to rollup buckets the watermark passed. + async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()>; /// Whether rows may batch off the request path; in-process backends stay synchronous. fn defers_ledger_writes(&self) -> bool { false @@ -386,8 +389,7 @@ pub trait Store: Send + Sync + std::fmt::Debug { until: i64, bucket_secs: i64, ) -> GResult>; - /// Recompute and upsert every rollup bucket in the trailing backfill window; buckets are never - /// deleted, so a missed tick self-heals. Returns the buckets written. + /// Recompute the trailing backfill window's rollup buckets (never deleted: a missed tick self-heals); returns those written. async fn usage_rollup_advance(&self, now_epoch_secs: i64) -> GResult; /// Delete rolled billing rows past `ledger_max_rows`; returns rows deleted. async fn ledger_prune(&self) -> GResult { @@ -576,11 +578,8 @@ impl MemoryStore { pub(crate) fn fail_next_ledger_writes(&self, count: usize) { self.ledger_failures.store(count, Ordering::Relaxed); } -} -#[async_trait::async_trait] -impl Store for MemoryStore { - async fn ledger_add(&self, r: &BillingRecord) -> GResult<()> { + fn injected_ledger_failure(&self) -> GResult<()> { #[cfg(test)] if self .ledger_failures @@ -593,6 +592,14 @@ impl Store for MemoryStore { "injected ledger write failure", )); } + Ok(()) + } +} + +#[async_trait::async_trait] +impl Store for MemoryStore { + async fn ledger_add(&self, r: &BillingRecord) -> GResult<()> { + self.injected_ledger_failure()?; // watermark first: rollup-then-records is the lock order advance uses let watermark = (self.ledger_max_rows > 0).then(|| rollup_watermark(&lock(&self.rollup))); let mut ledger = lock(&self.ledger); @@ -625,6 +632,35 @@ impl Store for MemoryStore { Ok(()) } + async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()> { + self.injected_ledger_failure()?; + let mut rollup = lock(&self.rollup); + let watermark = rollup_watermark(&rollup); + let mut ledger = lock(&self.ledger); + for r in rows { + if !ledger.request_ids.insert(r.request_id.clone()) { + continue; + } + ledger.rows.push(r.clone()); + if r.created_at_epoch_secs < watermark { + rollup + .entry(( + bucket_floor(r.created_at_epoch_secs), + r.tenant.clone(), + r.user_id.clone(), + r.model.clone(), + )) + .or_insert_with(|| UserUsageRow { + user_id: r.user_id.clone(), + model: r.model.clone(), + ..Default::default() + }) + .add_record(r); + } + } + Ok(()) + } + async fn ledger_snapshot(&self, limit: usize) -> GResult<(usize, Vec)> { let ledger = lock(&self.ledger); let records = &ledger.rows; @@ -1378,13 +1414,55 @@ macro_rules! q_series_raw { }; } -macro_rules! q_rollup_insert { +macro_rules! billing_insert { + ($rows:expr) => {{ + // the QueryBuilder numbers placeholders per dialect, so no dialect_sql! here + let mut qb = sqlx::QueryBuilder::new( + "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, + account, prompt_tokens, completion_tokens, total_tokens, cost_micros, + vendor_cost_micros, ptu_spillover, user_id, request_id, + created_at_epoch_secs, estimated, billed_units) ", + ); + qb.push_values($rows, |mut v, r| { + v.push_bind(&r.ak) + .push_bind(&r.product) + .push_bind(&r.tenant) + .push_bind(&r.model) + .push_bind(&r.served_model) + .push_bind(&r.protocol) + .push_bind(&r.account) + .push_bind(r.prompt_tokens) + .push_bind(r.completion_tokens) + .push_bind(r.total_tokens) + .push_bind(r.cost_micros) + .push_bind(r.vendor_cost_micros) + .push_bind(r.ptu_spillover) + .push_bind(&r.user_id) + .push_bind(&r.request_id) + .push_bind(r.created_at_epoch_secs) + .push_bind(r.estimated) + .push_bind(r.billed_units); + }); + qb.push(" ON CONFLICT (request_id) DO NOTHING"); + qb + }}; +} + +macro_rules! rollup_select { ($d:ident) => { concat!( "INSERT INTO usage_rollup (minute_epoch, tenant, user_id, model, requests, \ prompt_tokens, completion_tokens, total_tokens, cost_micros, vendor_cost_micros, \ billed_units) SELECT (created_at_epoch_secs/60)*60, tenant, user_id, model, COUNT(*), ", - usage_sums!($d), + usage_sums!($d) + ) + }; +} + +macro_rules! q_rollup_insert { + ($d:ident) => { + concat!( + rollup_select!($d), " FROM billing WHERE created_at_epoch_secs >= ?1 AND created_at_epoch_secs < ?2 \ GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET requests = ", sql_greatest!($d), "(usage_rollup.requests, excluded.requests), prompt_tokens = ", @@ -1571,34 +1649,7 @@ macro_rules! sql_store_impl { if rows.is_empty() { return Ok(()); } - // the QueryBuilder numbers placeholders per dialect, so no dialect_sql! here - let mut qb = sqlx::QueryBuilder::new( - "INSERT INTO billing (ak, product, tenant, model, served_model, protocol, - account, prompt_tokens, completion_tokens, total_tokens, cost_micros, - vendor_cost_micros, ptu_spillover, user_id, request_id, - created_at_epoch_secs, estimated, billed_units) ", - ); - qb.push_values(rows, |mut v, r| { - v.push_bind(&r.ak) - .push_bind(&r.product) - .push_bind(&r.tenant) - .push_bind(&r.model) - .push_bind(&r.served_model) - .push_bind(&r.protocol) - .push_bind(&r.account) - .push_bind(r.prompt_tokens) - .push_bind(r.completion_tokens) - .push_bind(r.total_tokens) - .push_bind(r.cost_micros) - .push_bind(r.vendor_cost_micros) - .push_bind(r.ptu_spillover) - .push_bind(&r.user_id) - .push_bind(&r.request_id) - .push_bind(r.created_at_epoch_secs) - .push_bind(r.estimated) - .push_bind(r.billed_units); - }); - qb.push(" ON CONFLICT (request_id) DO NOTHING") + billing_insert!(rows) .build() .execute(&self.pool) .await @@ -1606,6 +1657,52 @@ macro_rules! sql_store_impl { Ok(()) } + async fn ledger_replay_batch(&self, rows: &[BillingRecord]) -> GResult<()> { + if rows.is_empty() { + return Ok(()); + } + let mut tx = self + .pool + .begin() + .await + .map_err(|e| crate::sqlx_err("begin replay", e))?; + let mut insert = billing_insert!(rows); + let inserted: Vec = insert + .push(" RETURNING request_id") + .build_query_scalar() + .fetch_all(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("replay billing records", e))?; + if !inserted.is_empty() { + let mut add = sqlx::QueryBuilder::new(concat!( + rollup_select!($dialect), + " FROM billing WHERE created_at_epoch_secs < (", + )); + add.push(ROLLUP_WATERMARK_SQL).push(") AND request_id IN ("); + let mut ids = add.separated(", "); + for id in &inserted { + ids.push_bind(id); + } + add.push( + ") GROUP BY 1, 2, 3, 4 ON CONFLICT (minute_epoch, tenant, user_id, model) DO UPDATE SET \ + requests = usage_rollup.requests + excluded.requests, \ + prompt_tokens = usage_rollup.prompt_tokens + excluded.prompt_tokens, \ + completion_tokens = usage_rollup.completion_tokens + excluded.completion_tokens, \ + total_tokens = usage_rollup.total_tokens + excluded.total_tokens, \ + cost_micros = usage_rollup.cost_micros + excluded.cost_micros, \ + vendor_cost_micros = usage_rollup.vendor_cost_micros + excluded.vendor_cost_micros, \ + billed_units = usage_rollup.billed_units + excluded.billed_units", + ) + .build() + .execute(&mut *tx) + .await + .map_err(|e| crate::sqlx_err("roll up replayed records", e))?; + } + tx.commit() + .await + .map_err(|e| crate::sqlx_err("commit replay", e)) + } + fn defers_ledger_writes(&self) -> bool { true } @@ -3466,6 +3563,69 @@ mod tests { exercise_rollup(&store, "").await; } + async fn exercise_replay(store: &dyn Store, ns: &str, prune: impl AsyncFnOnce(&str)) { + let tenant = format!("rp{ns}"); + let billed = |at: i64| BillingRecord { + tenant: tenant.clone(), + user_id: "late".into(), + created_at_epoch_secs: at, + ..record("m1") + }; + let rolled = billed(400); + store.ledger_add(&rolled).await.unwrap(); + store.usage_rollup_advance(1_500).await.unwrap(); + prune(&rolled.request_id).await; + let late = billed(410); + store + .ledger_replay_batch(std::slice::from_ref(&late)) + .await + .unwrap(); + store + .ledger_replay_batch(std::slice::from_ref(&late)) + .await + .unwrap(); + store.usage_rollup_advance(1_500).await.unwrap(); + let usage = store + .usage_by_user(Some(&tenant), Some("late"), 0, 1_000) + .await + .unwrap(); + assert_eq!( + usage + .iter() + .map(|u| (u.requests, u.cost_micros)) + .collect::>(), + vec![(2, 2 * rolled.cost_micros)], + "{usage:?}" + ); + } + + #[tokio::test] + async fn memory_replay_adds_late_rows_to_a_pruned_bucket() { + let store = MemoryStore::default(); + exercise_replay(&store, "", async |id: &str| { + let mut ledger = lock(&store.ledger); + ledger.rows.retain(|r| r.request_id != id); + ledger.request_ids.remove(id); + }) + .await; + } + + #[tokio::test] + async fn sqlite_replay_adds_late_rows_to_a_pruned_bucket() { + let dir = tempfile::tempdir().unwrap(); + let store = SqliteStore::open(dir.path().join("replay.db").to_str().unwrap()) + .await + .unwrap(); + exercise_replay(&store, "", async |id: &str| { + sqlx::query("DELETE FROM billing WHERE request_id = ?1") + .bind(id) + .execute(&store.pool) + .await + .unwrap(); + }) + .await; + } + #[tokio::test] async fn rollup_survives_ledger_prune() { let store = MemoryStore::with_ledger_cap(1); @@ -3838,6 +3998,14 @@ mod tests { let store = PostgresStore::connect(&own_url).await.expect("pg connect"); let ns = format!("-{nonce}"); exercise_rollup(&store, &ns).await; + exercise_replay(&store, &ns, async |id: &str| { + sqlx::query("DELETE FROM billing WHERE request_id = $1") + .bind(id) + .execute(&store.pool) + .await + .unwrap(); + }) + .await; exercise_erase(&store, &ns).await; exercise_content_list(&store, &ns).await; exercise_terminal_idempotence(&store, &ns).await; diff --git a/crates/views/src/lib.rs b/crates/views/src/lib.rs index ad0b553..176adcd 100644 --- a/crates/views/src/lib.rs +++ b/crates/views/src/lib.rs @@ -585,10 +585,17 @@ async fn realtime_gate( admission::check_model_qpm(gov, cfg, &m.served) .await .map_err(throttled)?; - let user_budget = admission::check_budgets(state, cfg, &ak, ak.attributed_user(hint)) - .await - .map_err(|e| (ErrClass::ServiceUnavailable, e.message))? - .map_err(quota_exceeded)?; + let (user_budget, _) = admission::reserve_budgets( + state, + cfg, + &ak, + ak.attributed_user(hint), + admission::BudgetSpend::default, + gw_state::epoch_secs(), + ) + .await + .map_err(|e| (ErrClass::ServiceUnavailable, e.message))? + .map_err(quota_exceeded)?; if let Some(limit) = admission::model_quota_limit(cfg, &ak, &m.requested) && !gov .quota_check(&admission::model_quota_key(&ak.ak_id, &m.requested), limit) @@ -689,6 +696,7 @@ async fn bill_realtime_turn( model_quota_key, budget_ak: ak, user_budget: admit.user_budget, + budget_hold: Default::default(), }, ) .await; @@ -1225,7 +1233,9 @@ fn log_access(surface: &str, ctx: &DagContext, started: Instant) { ) }) .unwrap_or_default(); - let latency = started.elapsed(); + let latency = ctx + .served_at + .map_or_else(|| started.elapsed(), |at| at.duration_since(started)); let user_id = ctx.effective_user_id(); let decisions = ctx.decisions_line(); let ak_id = &*ctx.ak.ak_id; @@ -1340,14 +1350,14 @@ fn user_hint(hint: Option, field: &Value) -> Option { hint.or_else(|| field.as_str().map(cap_user_hint)) } -/// Bound an attribution hint at `USER_HINT_MAX_BYTES`: it keys governance -/// counters, so an unbounded value from any surface would grow the keyspace. +/// Bound an attribution hint at `USER_HINT_MAX_BYTES` (it keys governance counters) +/// and drop NUL, which Postgres TEXT refuses in every row the user lands in. fn cap_user_hint(s: &str) -> String { let mut end = s.len().min(USER_HINT_MAX_BYTES); while !s.is_char_boundary(end) { end -= 1; } - s[..end].to_owned() + s[..end].replace('\0', "") } /// AK auth: `Authorization: Bearer ` or `x-api-key: `. The error is @@ -4582,6 +4592,7 @@ async fn admit_video_job( model_quota_key: None, budget_ak: submitter.as_deref().unwrap_or(&ak), user_budget: None, + budget_hold: Default::default(), }, ) .await; @@ -6720,6 +6731,17 @@ mod tests { ); } + #[test] + fn user_hints_are_capped_and_drop_nul() { + assert_eq!(cap_user_hint("nul\0user"), "nuluser"); + let long = "é".repeat(USER_HINT_MAX_BYTES); + assert!(cap_user_hint(&long).len() <= USER_HINT_MAX_BYTES); + assert_eq!( + user_hint(None, &serde_json::json!("\u{0}")), + Some(String::new()) + ); + } + #[tokio::test] async fn realtime_attributes_user_from_owner_then_header_hint() { let yaml = "listen: {host: h, port: 1}\nmodels: [{name: rt, protocol: realtime, input_price_per_1k_micros: 1000, output_price_per_1k_micros: 1000}]\naccess_keys: [{ak: k-shared, product: p, qps: 10, daily_token_quota: 100000}, {ak: k-owned, product: p, qps: 10, daily_token_quota: 100000, owner: bob}]"; diff --git a/docs/architecture.md b/docs/architecture.md index 1138f49..991c341 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -62,7 +62,8 @@ default, so the whole pipeline is testable offline: - **`TokenEncoder`** — prompt token estimation capability (tiktoken cl100k_base BPE default, zero-dependency heuristic fallback). `cost_calc` uses it to bill an aborted stream from the delivered text when the vendor's - usage frame never arrives; token-aware PTU sizing is not ported yet. + usage frame never arrives and the account keeps no generation record + (OpenRouter does: `engines::generation`); token-aware PTU sizing is not ported yet. - **`Store`** — durable records (billing ledger, uploaded files, batch jobs). `MemoryStore` by default; `SqliteStore` when `storage.sqlite_path` is configured; `PostgresStore` when `storage.postgres_url` is — shared diff --git a/docs/configuration.md b/docs/configuration.md index 28201ea..bdd4a17 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -34,6 +34,7 @@ storage: shared_cache: false # also share the response cache in Redis (needs redis_url) ledger_max_rows: 100000 # prune oldest billing rows past the cap; 0 = unlimited postgres_max_connections: 10 # store and key-store connection-pool size; 0 = the built-in defaults (10 store, 5 key store) + ledger_spool_path: /var/lib/gw/ledger-spool.jsonl # rows the store refuses wait here for replay; empty = dropped after the retries ``` The billing ledger, uploaded files, and batch jobs live here. In-memory by diff --git a/docs/deployment.md b/docs/deployment.md index 273c479..b2aae4f 100644 --- a/docs/deployment.md +++ b/docs/deployment.md @@ -70,6 +70,7 @@ storage: redis_url: "redis://redis:6379" # shared rate limits + quotas + account health ledger_max_rows: 1000000 # prune oldest billing rows past the cap postgres_max_connections: 10 # store and key-store connection-pool size + ledger_spool_path: /var/lib/gw/ledger-spool.jsonl # keep billing rows through a store outage # sqlite_path: /var/lib/gw/store.db # single-node alternative to postgres_url ``` diff --git a/docs/governance.md b/docs/governance.md index 57ae1db..003bf5f 100644 --- a/docs/governance.md +++ b/docs/governance.md @@ -81,10 +81,10 @@ models: | QPS | pooled per tenant | `tenants[].qps` | | Daily tokens | per access key | `access_keys[].daily_token_quota` (fleet/Redis: rolls at UTC midnight; single-node in-memory: a ~daily background reset) | | Daily tokens | per (key, model) | `tenants[].model_quotas` default, `access_keys[].model_quotas` override (over it: the tenant `fallback_model`, else refused) | -| Daily tokens | per end user | `tenants[].user_daily_token_quota` (soft) | -| Daily cost | pooled per tenant | `tenants[].daily_cost_quota_micros` (micro-dollars of charged price; soft) | -| Daily cost | per access key | `tenants[].key_daily_cost_quota_micros` (soft) | -| Daily cost | per end user | `tenants[].user_daily_cost_quota_micros` (soft), per-user override via `/admin/tenants/{tenant}/users/{user}/budget` | +| Daily tokens | per end user | `tenants[].user_daily_token_quota` (held) | +| Daily cost | pooled per tenant | `tenants[].daily_cost_quota_micros` (micro-dollars of charged price; held) | +| Daily cost | per access key | `tenants[].key_daily_cost_quota_micros` (held) | +| Daily cost | per end user | `tenants[].user_daily_cost_quota_micros` (held), per-user override via `/admin/tenants/{tenant}/users/{user}/budget` | | Monthly cost | pooled per tenant / per key / per end user | `tenants[].monthly_cost_quota_micros`, `key_monthly_cost_quota_micros`, `user_monthly_cost_quota_micros` (UTC calendar month; soft; `monthly_cost_rollover` carries the unspent remainder); the per-user cap takes a per-user override | | TPM | per access key | `access_keys[].tokens_per_minute` | | QPM | per model | `models[].qpm` | @@ -142,8 +142,9 @@ tokens and 1-hour cache writes take their own `token_rate` weights, a scales what `/v1/batches` items cost. A streaming response that breaks after delivery has begun (client disconnect or -upstream failure) is billed for what was delivered: the vendor's usage frame -never arrives, so the token count is estimated from the request and the +upstream failure) never sees the vendor's usage frame. On an OpenRouter account +it is billed from OpenRouter's record of the generation, which counts what the +model generated; elsewhere the token count is estimated from the request and the delivered text. A disconnect *before* any bytes are sent bills nothing. ## MCP tool access @@ -234,8 +235,9 @@ deltas), so **whenever a tenant redacts output — `dlp_redact` or replayed** rather than forwarded token-by-token — DLP trades incremental delivery for a guarantee that no unmasked text reaches the client. Billing is settled after that replay: a disconnect before the first redacted -chunk refunds the reservation, while a partial replay bills only an estimate of -the delivered text. +chunk refunds the reservation, while a partial replay bills the upstream's usage +frame (what the model generated), or an estimate of the delivered text when the +upstream sent none. Turn both `dlp_redact` and `detect_secrets` off to keep incremental delivery; note the embedded demo config ships with `dlp_redact` on. @@ -262,15 +264,25 @@ REST, realtime (the `x-gw-user` hint is captured at WS connect), and batch (each item's `user`, or the submitter's `x-gw-user`, is persisted with the item so a fleet drainer still attributes and budgets it). `GET /admin/usage/users?user=&since=&until=` returns per-(user, model) cost over a billing period (add `format=csv` for -export). `TenantConf.user_daily_token_quota` sets a soft per-user daily cap — -enforced with a hard `400 service_quota_exceeded_exception` once exceeded; -"soft" means check-then-consume, so concurrent in-flight turns can overshoot it -by one before the counter accrues. +export). `TenantConf.user_daily_token_quota` sets a per-user daily cap — +enforced with a hard `400 service_quota_exceeded_exception` once exceeded. +On the REST pipeline every daily budget holds the call's estimate at admission +(the token estimate for the token cap; for the cost caps that estimate's prompt +and `max_tokens` sides at the tenant's price), admitting while spent plus held +is under the cap, so concurrent holds pass it by at most one hold; billing +settles the hold to the charged amount and a failed call returns it. A call that +costs more than it held still carries the total past the cap by the difference: +a call without `max_tokens` holds 256 output tokens. Units known at admission +hold at the unit price (a search, `n` images, the characters of a speech +request; one unit for transcription and rerank); video bills at its poll and +vendor-priced decisions report their cost after the call, so for those the cost +caps only check. Realtime turns and the monthly caps check before the call +and accrue after it ("soft": concurrent turns can overshoot by one each). The per-(key, model) quota is soft the same way: it is a routing trigger (the fallback-model degrade), while the reserved per-key daily quota hard-caps spend. The daily cost budgets — `daily_cost_quota_micros` pooled over the tenant's keys, `key_daily_cost_quota_micros` per key, `user_daily_cost_quota_micros` -per end user — are soft caps of the same kind over the ledger's charged +per end user — are caps of the same kind over the ledger's charged `cost_micros` (list price, the tenant's `model_prices`, or a decisions reply's `usage.cost`, so unit-priced surfaces count too); they apply on every surface, realtime and batch included, and a response-cache hit, which bills nothing, counts nothing. diff --git a/docs/observability.md b/docs/observability.md index 1a7f04b..1bfd6b3 100644 --- a/docs/observability.md +++ b/docs/observability.md @@ -87,9 +87,26 @@ the unit-priced surfaces; also summed in `/admin/usage`), charged `cost_micros` and `vendor_cost_micros`, `created_at_epoch_secs`, the PTU-spillover flag, and an `estimated` flag (set when counts came from an aborted stream rather than a vendor usage payload). +An aborted stream served by an OpenRouter account is billed from OpenRouter's +record of the generation instead (`GET /api/v1/generation`, polled for about +40 s with each lookup capped at 10 s; cancelled generations included), with cache +reads at their weight (the record carries no cache-write count, so a cache write +bills at the prompt weight) and OpenRouter's cost as `vendor_cost_micros`; only when no +record arrives, or the gateway starts shutting down, does it fall back to the +estimate. Shutdown waits for these settles before the final ledger flush, and the +access log's `latency_ms` stops at the end of the stream, not of the lookup. Writes are idempotent by `request_id`; a transient store failure is held in a bounded repair queue and retried with backoff — eight attempts, about 13 s — before the batch is dropped and counted in `gateway_ledger_write_failures_total`. +With `storage.ledger_spool_path` set, the batch instead goes to that local +JSON-lines file after three attempts (about 0.3 s; the directory is created on +first use, and a spool that cannot be written falls back to the remaining +retries), and a replay every 30 s inserts it in batches once the store accepts +writes again; in the same transaction it adds the rows it inserted to any usage +rollup bucket the watermark already passed, so a late row counts even after its +bucket's other rows were pruned, and a repeated replay adds nothing +(`billing spool replayed` in the log). The spool path is per-instance local +disk; replicas must not share one file. A full queue applies backpressure to settlement until the repair worker makes room, so overload alone never discards an accepted row. Per-user usage additionally rolls into durable minute buckets every minute, so