From 06394789f4602e72a869dfa255488956a1b8bdfb Mon Sep 17 00:00:00 2001 From: Mega Mind <146339422+M3gA-Mind@users.noreply.github.com> Date: Mon, 27 Jul 2026 21:42:44 +0530 Subject: [PATCH] fix(cost): exempt BYOK inference from the managed-credit budget (#5218) --- src/openhuman/config/schema/identity_cost.rs | 11 +- src/openhuman/cost/mod.rs | 2 + src/openhuman/cost/route.rs | 229 ++++++++++++++++++ src/openhuman/cost/tracker.rs | 110 ++++++++- src/openhuman/cost/tracker_tests.rs | 231 ++++++++++++++++++- src/openhuman/tinyagents/middleware.rs | 21 +- 6 files changed, 591 insertions(+), 13 deletions(-) create mode 100644 src/openhuman/cost/route.rs diff --git a/src/openhuman/config/schema/identity_cost.rs b/src/openhuman/config/schema/identity_cost.rs index b6e20de9a..ae79e5750 100644 --- a/src/openhuman/config/schema/identity_cost.rs +++ b/src/openhuman/config/schema/identity_cost.rs @@ -27,11 +27,18 @@ pub struct CostConfig { #[serde(default = "default_cost_enabled")] pub enabled: bool, - /// Daily spending limit in USD (default: 10.00) + /// Daily spending limit in USD (default: 10.00). + /// + /// Applies to **managed (OpenHuman-credit) inference only** — see + /// [`crate::openhuman::cost::route`]. Bring-your-own-key and local + /// inference is billed by the user's own provider, so it is recorded for + /// the dashboard but never counted against this limit and can never + /// refuse a request (#5016). #[serde(default = "default_daily_limit")] pub daily_limit_usd: f64, - /// Monthly spending limit in USD (default: 100.00) + /// Monthly spending limit in USD (default: 100.00). Managed-route only, + /// on the same terms as [`Self::daily_limit_usd`]. #[serde(default = "default_monthly_limit")] pub monthly_limit_usd: f64, diff --git a/src/openhuman/cost/mod.rs b/src/openhuman/cost/mod.rs index c325ae656..45a2838bb 100644 --- a/src/openhuman/cost/mod.rs +++ b/src/openhuman/cost/mod.rs @@ -1,5 +1,6 @@ pub mod catalog; mod global; +pub mod route; mod rpc; mod schemas; pub mod tools; @@ -7,6 +8,7 @@ pub mod tracker; pub mod types; pub use global::{init_global, record_embedding_usage, record_provider_usage, try_global}; +pub use route::{route_for_model, CostRoute}; pub use schemas::{ all_controller_schemas as all_cost_controller_schemas, all_registered_controllers as all_cost_registered_controllers, diff --git a/src/openhuman/cost/route.rs b/src/openhuman/cost/route.rs new file mode 100644 index 000000000..863de01ff --- /dev/null +++ b/src/openhuman/cost/route.rs @@ -0,0 +1,229 @@ +//! Which billing route a recorded cost belongs to (issue #5016). +//! +//! The local `[cost]` daily/monthly limits exist to cap spend against +//! **OpenHuman-managed credits**. They were being enforced against *every* +//! recorded call, including bring-your-own-key (BYOK) and local inference that +//! OpenHuman never bills for. A BYOK user therefore accumulated phantom spend +//! — priced locally from [`super::catalog`] because their provider echoes no +//! `charged_amount_usd` — until they tripped the default $10/day cap and got +//! "You're out of credits", despite OpenHuman having charged them nothing. +//! +//! The route is derived from the recorded model id rather than threaded +//! through as a new parameter, which matters for two reasons: +//! +//! 1. Every recording site already has the model id, so no call site has to +//! learn about routing. +//! 2. Records persisted by older builds carry a model id too, so history +//! classifies correctly with **no migration and no schema change** — a +//! stored-route field would have had to guess a default for every existing +//! record and would keep mis-gating real users for up to a month. + +/// The billing route a cost record belongs to. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CostRoute { + /// Served by the OpenHuman managed backend and paid for with OpenHuman + /// credits. Counts toward — and is gated by — the local `[cost]` limits. + Managed, + /// Served by the user's own key (OpenRouter, Anthropic, a self-hosted + /// OpenAI-compatible gateway, …) or a local model. OpenHuman bills nothing + /// for it, so it is recorded for the dashboard but never gated. + Byok, +} + +impl CostRoute { + /// Whether spend on this route counts toward the local `[cost]` budget. + pub fn counts_toward_budget(self) -> bool { + matches!(self, CostRoute::Managed) + } +} + +/// Model ids the managed backend serves. These are the backend's own tier +/// slugs (`crate::openhuman::config::MODEL_*_V1`) — a BYOK +/// provider is always addressed by its real model id (`minimax/minimax-m3`, +/// `anthropic/claude-sonnet-4-20250514`, `llama3:8b`), never by a tier slug, +/// because the tier vocabulary only means something to the managed backend. +/// +/// Kept as a local list, deliberately: this is the set of ids that imply +/// *managed billing*, which is a narrower question than "is this a known tier +/// constant". A new tier must be added here consciously, and the +/// `managed_tier_slugs_stay_in_sync` test fails if one is added upstream +/// without doing so. +const MANAGED_MODEL_SLUGS: &[&str] = &[ + "chat-v1", + "reasoning-v1", + "reasoning-quick-v1", + "agentic-v1", + "burst-v1", + "coding-v1", + "vision-v1", + "summarization-v1", +]; + +/// Classify a recorded model id into its billing route. +/// +/// Unknown ids classify as [`CostRoute::Byok`], which is the safe direction: +/// the failure mode is under-enforcing a *local* convenience cap, not +/// over-charging. Real managed-credit exhaustion is still enforced +/// server-side, where the backend returns its own out-of-credits error — that +/// path is untouched by this classification. +pub fn route_for_model(model: &str) -> CostRoute { + let normalized = normalize_model_id(model); + if MANAGED_MODEL_SLUGS.contains(&normalized.as_str()) { + CostRoute::Managed + } else { + CostRoute::Byok + } +} + +/// Lower-case, trim, and strip the decorations a model id can pick up on its +/// way to a cost record: a `hint:` prefix from the tier-resolution helpers and +/// an `openhuman/` provider qualifier. +/// +/// Stripping loops until neither prefix applies, so the two decorations are +/// **order-independent**. A single fixed-order pass classified +/// `openhuman/hint:chat-v1` as BYOK (it stripped `openhuman/`, leaving +/// `hint:chat-v1`, which is not a managed slug) while `hint:openhuman/chat-v1` +/// classified as Managed — meaning managed spend could silently stop counting +/// toward the cap depending only on which decoration a recording site applied +/// first. +fn normalize_model_id(model: &str) -> String { + let mut current = model.trim().to_ascii_lowercase(); + loop { + let stripped = current + .strip_prefix("hint:") + .or_else(|| current.strip_prefix("openhuman/")); + match stripped { + Some(rest) => current = rest.to_string(), + None => return current, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Both decorations, in BOTH orders, must normalize to the same slug. + /// A fixed-order strip classified `openhuman/hint:…` as BYOK, so managed + /// spend silently stopped counting toward the cap depending only on which + /// prefix a recording site applied first (review, #5016). + #[test] + fn decoration_prefixes_are_order_independent() { + for id in [ + "chat-v1", + "hint:chat-v1", + "openhuman/chat-v1", + "hint:openhuman/chat-v1", + "openhuman/hint:chat-v1", + " HINT:OpenHuman/Chat-V1 ", + ] { + assert_eq!( + route_for_model(id), + CostRoute::Managed, + "{id} must classify as managed" + ); + } + } + + #[test] + fn decorations_do_not_promote_a_byok_model_to_managed() { + for id in [ + "openhuman/hint:llama3:8b", + "hint:openhuman/anthropic/claude-sonnet-4-20250514", + "ollama:chat-v1-not-a-slug", + ] { + assert_eq!(route_for_model(id), CostRoute::Byok, "{id} must stay BYOK"); + } + } + + use crate::openhuman::config::{ + MODEL_AGENTIC_V1, MODEL_BURST_V1, MODEL_CHAT_V1, MODEL_CODING_V1, MODEL_REASONING_QUICK_V1, + MODEL_REASONING_V1, MODEL_SUMMARIZATION_V1, MODEL_VISION_V1, + }; + + #[test] + fn managed_tier_slugs_are_managed() { + for slug in MANAGED_MODEL_SLUGS { + assert_eq!( + route_for_model(slug), + CostRoute::Managed, + "tier slug {slug} must bill as managed" + ); + } + } + + /// The local slug list must not drift from the backend tier constants. If a + /// new tier is introduced upstream, this fails until it is classified here + /// — otherwise managed spend on the new tier would silently stop counting + /// toward the budget. + #[test] + fn managed_tier_slugs_stay_in_sync() { + let upstream = [ + MODEL_CHAT_V1, + MODEL_REASONING_V1, + MODEL_REASONING_QUICK_V1, + MODEL_AGENTIC_V1, + MODEL_BURST_V1, + MODEL_CODING_V1, + MODEL_VISION_V1, + MODEL_SUMMARIZATION_V1, + ]; + for tier in upstream { + assert!( + MANAGED_MODEL_SLUGS.contains(&tier), + "backend tier {tier} is not classified as managed" + ); + } + assert_eq!( + MANAGED_MODEL_SLUGS.len(), + upstream.len(), + "MANAGED_MODEL_SLUGS has an entry with no matching backend tier constant" + ); + } + + #[test] + fn byok_and_local_models_are_not_managed() { + // The exact models from #5016 / #5127 (OpenRouter + a self-hosted + // OpenAI-compatible gateway) and other common BYOK / local shapes. + for model in [ + "minimax/minimax-m3", + "anthropic/claude-sonnet-4-20250514", + "openai/gpt-4o", + "llama3:8b", + "ollama:gemma3:1b-it-qat", + "lmstudio:qwen2.5-coder", + "", + ] { + assert_eq!( + route_for_model(model), + CostRoute::Byok, + "{model} must not bill as managed" + ); + } + } + + #[test] + fn normalizes_case_whitespace_and_prefixes() { + assert_eq!(route_for_model(" Chat-V1 "), CostRoute::Managed); + assert_eq!(route_for_model("hint:chat-v1"), CostRoute::Managed); + assert_eq!( + route_for_model("openhuman/reasoning-v1"), + CostRoute::Managed + ); + } + + #[test] + fn a_byok_model_merely_containing_a_tier_name_is_not_managed() { + // Substring matching would misclassify these and silently re-introduce + // the phantom limit for the user. + for model in ["vendor/chat-v1-turbo", "my-chat-v1", "chat-v1x"] { + assert_eq!(route_for_model(model), CostRoute::Byok, "{model}"); + } + } + + #[test] + fn only_managed_counts_toward_budget() { + assert!(CostRoute::Managed.counts_toward_budget()); + assert!(!CostRoute::Byok.counts_toward_budget()); + } +} diff --git a/src/openhuman/cost/tracker.rs b/src/openhuman/cost/tracker.rs index bd5fe5543..e716d8445 100644 --- a/src/openhuman/cost/tracker.rs +++ b/src/openhuman/cost/tracker.rs @@ -1,3 +1,4 @@ +use super::route::route_for_model; use super::types::{ BudgetCheck, BudgetStatus, CostDashboard, CostRecord, CostSummary, DailyCostEntry, ModelStats, TokenUsage, UsagePeriod, @@ -51,6 +52,17 @@ impl CostTracker { } /// Check if a request is within budget. + /// + /// Only **managed-route** spend is considered (#5016). The local `[cost]` + /// limits cap spend against OpenHuman credits; bring-your-own-key and + /// local inference is billed by the user's own provider, so counting it + /// here produced a phantom limit — a pure-BYOK user accrued locally + /// *estimated* spend they were never charged for and got "You're out of + /// credits" at the default $10/day. See [`super::route`]. + /// + /// A pure-BYOK user therefore has zero managed spend and can never trip + /// this gate. Real managed-credit exhaustion is unaffected: it is enforced + /// server-side by the backend, which returns its own billing error. pub fn check_budget(&self, estimated_cost_usd: f64) -> Result { if !self.config.enabled { return Ok(BudgetCheck::Allowed); @@ -63,7 +75,23 @@ impl CostTracker { } let mut storage = self.lock_storage(); - let (daily_cost, monthly_cost) = storage.get_aggregated_costs()?; + let (daily_cost, monthly_cost) = storage.get_aggregated_managed_costs()?; + // The all-routes totals exist purely to make the managed-vs-BYOK split + // visible in a debug log. `tracing` evaluates field expressions eagerly, + // so compute them only when that level is actually enabled rather than + // on every budget check in production. + if tracing::enabled!(tracing::Level::DEBUG) { + let (daily_all, monthly_all) = storage.get_aggregated_costs()?; + tracing::debug!( + daily_managed_usd = daily_cost, + monthly_managed_usd = monthly_cost, + daily_all_routes_usd = daily_all, + monthly_all_routes_usd = monthly_all, + daily_limit_usd = self.config.daily_limit_usd, + monthly_limit_usd = self.config.monthly_limit_usd, + "[cost] budget check against managed-route spend only (BYOK excluded, #5016)" + ); + } // Check daily limit let projected_daily = daily_cost + estimated_cost_usd; @@ -183,12 +211,20 @@ impl CostTracker { storage.get_cost_for_date(date) } - /// Get the monthly cost for a specific month. + /// Get the monthly cost for a specific month, across all routes. pub fn get_monthly_cost(&self, year: i32, month: u32) -> Result { let storage = self.lock_storage(); storage.get_cost_for_month(year, month) } + /// Get the **managed-route** monthly cost for a specific month — the + /// portion of spend the local `[cost]` budget applies to (#5016). BYOK and + /// local inference is excluded because OpenHuman never bills for it. + pub fn get_managed_monthly_cost(&self, year: i32, month: u32) -> Result { + let storage = self.lock_storage(); + storage.get_managed_cost_for_month(year, month) + } + /// Get a daily cost/token history covering the last `days` calendar days, /// ending on today (UTC) inclusive. Days with no recorded usage are /// returned as zero-filled entries so callers can render the chart bars @@ -300,9 +336,16 @@ impl CostTracker { let budget_limit_monthly_usd = self.config.monthly_limit_usd.max(0.0); let now = Utc::now(); + // `month_to_date_usd` stays an all-route total: users want to see what + // their BYOK providers cost them. Budget utilisation and status, + // however, describe a limit that only gates *managed* spend (#5016) — + // driving them off the total made a pure-BYOK user's gauge fill to + // 100% against a cap that can never fire, which is the phantom limit + // reported in the issue. let month_to_date_usd = self.get_monthly_cost(now.year(), now.month())?; + let managed_month_to_date_usd = self.get_managed_monthly_cost(now.year(), now.month())?; let budget_utilization = if budget_limit_monthly_usd > 0.0 { - (month_to_date_usd / budget_limit_monthly_usd).clamp(0.0, 1.0) + (managed_month_to_date_usd / budget_limit_monthly_usd).clamp(0.0, 1.0) } else { 0.0 }; @@ -312,7 +355,7 @@ impl CostTracker { } else { let warn = warn_threshold.clamp(0.0, 1.0); let alert = alert_threshold.clamp(0.0, 1.0).max(warn); - let utilization_raw = month_to_date_usd / budget_limit_monthly_usd; + let utilization_raw = managed_month_to_date_usd / budget_limit_monthly_usd; if utilization_raw >= alert { BudgetStatus::Exceeded } else if utilization_raw >= warn { @@ -415,6 +458,12 @@ struct CostStorage { path: PathBuf, daily_cost_usd: f64, monthly_cost_usd: f64, + /// Managed-route subset of `daily_cost_usd` — the only spend the local + /// `[cost]` budget may gate (#5016). BYOK/local spend is still counted in + /// the totals above so the dashboard shows it, but never here. + daily_managed_cost_usd: f64, + /// Managed-route subset of `monthly_cost_usd`. See above. + monthly_managed_cost_usd: f64, cached_day: NaiveDate, cached_year: i32, cached_month: u32, @@ -433,6 +482,8 @@ impl CostStorage { path: path.to_path_buf(), daily_cost_usd: 0.0, monthly_cost_usd: 0.0, + daily_managed_cost_usd: 0.0, + monthly_managed_cost_usd: 0.0, cached_day: now.date_naive(), cached_year: now.year(), cached_month: now.month(), @@ -491,21 +542,34 @@ impl CostStorage { fn rebuild_aggregates(&mut self, day: NaiveDate, year: i32, month: u32) -> Result<()> { let mut daily_cost = 0.0; let mut monthly_cost = 0.0; + let mut daily_managed_cost = 0.0; + let mut monthly_managed_cost = 0.0; self.for_each_record(|record| { let timestamp = record.usage.timestamp.naive_utc(); + // Classified from the record's own model id, so records written by + // builds that predate #5016 are routed correctly with no migration. + let counts = route_for_model(&record.usage.model).counts_toward_budget(); if timestamp.date() == day { daily_cost += record.usage.cost_usd; + if counts { + daily_managed_cost += record.usage.cost_usd; + } } if timestamp.year() == year && timestamp.month() == month { monthly_cost += record.usage.cost_usd; + if counts { + monthly_managed_cost += record.usage.cost_usd; + } } })?; self.daily_cost_usd = daily_cost; self.monthly_cost_usd = monthly_cost; + self.daily_managed_cost_usd = daily_managed_cost; + self.monthly_managed_cost_usd = monthly_managed_cost; self.cached_day = day; self.cached_year = year; self.cached_month = month; @@ -547,22 +611,39 @@ impl CostStorage { self.ensure_period_cache_current()?; let timestamp = record.usage.timestamp.naive_utc(); + let counts = route_for_model(&record.usage.model).counts_toward_budget(); if timestamp.date() == self.cached_day { self.daily_cost_usd += record.usage.cost_usd; + if counts { + self.daily_managed_cost_usd += record.usage.cost_usd; + } } if timestamp.year() == self.cached_year && timestamp.month() == self.cached_month { self.monthly_cost_usd += record.usage.cost_usd; + if counts { + self.monthly_managed_cost_usd += record.usage.cost_usd; + } } Ok(()) } - /// Get aggregated costs for current day and month. + /// Get aggregated costs for current day and month, across **all** routes. + /// This is the dashboard/summary figure — use + /// [`Self::get_aggregated_managed_costs`] for anything that gates a + /// request. fn get_aggregated_costs(&mut self) -> Result<(f64, f64)> { self.ensure_period_cache_current()?; Ok((self.daily_cost_usd, self.monthly_cost_usd)) } + /// Get aggregated **managed-route** costs for the current day and month — + /// the only spend the local `[cost]` budget may gate (#5016). + fn get_aggregated_managed_costs(&mut self) -> Result<(f64, f64)> { + self.ensure_period_cache_current()?; + Ok((self.daily_managed_cost_usd, self.monthly_managed_cost_usd)) + } + /// Get cost for a specific date. fn get_cost_for_date(&self, date: NaiveDate) -> Result { let mut cost = 0.0; @@ -589,6 +670,25 @@ impl CostStorage { Ok(cost) } + + /// Get **managed-route** cost for a specific month — the figure the budget + /// gauge should reflect, since only managed spend can trip the limit + /// (#5016). + fn get_managed_cost_for_month(&self, year: i32, month: u32) -> Result { + let mut cost = 0.0; + + self.for_each_record(|record| { + let timestamp = record.usage.timestamp.naive_utc(); + if timestamp.year() == year + && timestamp.month() == month + && route_for_model(&record.usage.model).counts_toward_budget() + { + cost += record.usage.cost_usd; + } + })?; + + Ok(cost) + } } #[cfg(test)] diff --git a/src/openhuman/cost/tracker_tests.rs b/src/openhuman/cost/tracker_tests.rs index fc8db7021..704fe6c98 100644 --- a/src/openhuman/cost/tracker_tests.rs +++ b/src/openhuman/cost/tracker_tests.rs @@ -9,6 +9,14 @@ fn enabled_config() -> CostConfig { } } +/// A managed-backend tier slug — spend on this route is billed to OpenHuman +/// credits and so is the only kind the local budget may gate (#5016). +const MANAGED_MODEL: &str = "chat-v1"; + +/// A bring-your-own-key model id, as reported in #5016 (OpenRouter). Spend +/// here is billed by the user's own provider and must never gate a request. +const BYOK_MODEL: &str = "minimax/minimax-m3"; + #[test] fn cost_tracker_initialization() { let tmp = TempDir::new().unwrap(); @@ -54,8 +62,9 @@ fn budget_exceeded_daily_limit() { let tracker = CostTracker::new(config, tmp.path()).unwrap(); - // Record a usage that exceeds the limit - let usage = TokenUsage::new("test/model", 10000, 5000, 1.0, 2.0); // ~0.02 USD + // Record managed-route usage that exceeds the limit. Only managed spend + // gates a request (#5016), so the model id has to be a backend tier slug. + let usage = TokenUsage::new(MANAGED_MODEL, 10000, 5000, 1.0, 2.0); // ~0.02 USD tracker.record_usage(usage).unwrap(); let check = tracker.check_budget(0.01).unwrap(); @@ -225,7 +234,7 @@ fn budget_monthly_exceeded() { }; let tracker = CostTracker::new(config, tmp.path()).unwrap(); - let usage = TokenUsage::new("test/model", 10000, 5000, 1.0, 2.0); + let usage = TokenUsage::new(MANAGED_MODEL, 10000, 5000, 1.0, 2.0); tracker.record_usage(usage).unwrap(); let check = tracker.check_budget(0.01).unwrap(); @@ -238,6 +247,216 @@ fn budget_monthly_exceeded() { )); } +// ── BYOK budget exemption (#5016 / #5127) ────────────────────────────── +// +// The reported bug: a user with no OpenHuman credits configured, routing all +// inference through their own OpenRouter key, accumulated locally *estimated* +// spend until they tripped the default $10/day cap and were told "You're out +// of credits" — for inference OpenHuman never billed them for. + +#[test] +fn byok_spend_never_exceeds_the_daily_limit() { + let tmp = TempDir::new().unwrap(); + let config = CostConfig { + enabled: true, + daily_limit_usd: 0.01, + ..Default::default() + }; + let tracker = CostTracker::new(config, tmp.path()).unwrap(); + + // Far past the $0.01 daily cap — and irrelevant, because it is BYOK. + let usage = TokenUsage::new(BYOK_MODEL, 10_000_000, 5_000_000, 1.0, 2.0); + tracker.record_usage(usage).unwrap(); + + // Estimate 0.0, matching what `CostBudgetMiddleware::before_model` actually + // passes. Charging the whole $0.01 limit to the *current* request would trip + // the 80% warning on that request's own projected cost, which says nothing + // about whether the recorded BYOK history leaked into the budget. + let check = tracker.check_budget(0.0).unwrap(); + assert!( + matches!(check, BudgetCheck::Allowed), + "BYOK spend must never gate a request, got {check:?}" + ); + + // `Exceeded` is the only variant that actually blocks a request, so pin it + // separately: even a request that would consume the entire remaining limit + // must not be blocked by BYOK history. + assert!( + !matches!( + tracker.check_budget(0.01).unwrap(), + BudgetCheck::Exceeded { .. } + ), + "BYOK history must never push a request over the managed cap" + ); +} + +#[test] +fn byok_spend_never_exceeds_the_monthly_limit() { + let tmp = TempDir::new().unwrap(); + let config = CostConfig { + enabled: true, + daily_limit_usd: 1000.0, + monthly_limit_usd: 0.01, + ..Default::default() + }; + let tracker = CostTracker::new(config, tmp.path()).unwrap(); + + let usage = TokenUsage::new(BYOK_MODEL, 10_000_000, 5_000_000, 1.0, 2.0); + tracker.record_usage(usage).unwrap(); + + // Estimate 0.0, as `CostBudgetMiddleware::before_model` passes: charging the + // whole $0.01 limit to the current request would trip the 80% warning on + // that request's own cost, which says nothing about the BYOK history. + let check = tracker.check_budget(0.0).unwrap(); + assert!( + matches!(check, BudgetCheck::Allowed), + "BYOK spend must never gate a request, got {check:?}" + ); + assert!( + !matches!( + tracker.check_budget(0.01).unwrap(), + BudgetCheck::Exceeded { .. } + ), + "BYOK history must never push a request over the managed monthly cap" + ); +} + +#[test] +fn byok_spend_does_not_trip_the_warning_threshold_either() { + let tmp = TempDir::new().unwrap(); + let config = CostConfig { + enabled: true, + daily_limit_usd: 10.0, + warn_at_percent: 80, + monthly_limit_usd: 1000.0, + ..Default::default() + }; + let tracker = CostTracker::new(config, tmp.path()).unwrap(); + + let mut usage = TokenUsage::new(BYOK_MODEL, 1000, 500, 1.0, 1.0); + usage.cost_usd = 9.5; // 95% of the daily limit, if it counted + tracker.record_usage(usage).unwrap(); + + let check = tracker.check_budget(0.0).unwrap(); + assert!( + matches!(check, BudgetCheck::Allowed), + "BYOK spend must not raise a budget warning, got {check:?}" + ); +} + +#[test] +fn byok_spend_is_still_recorded_for_the_dashboard() { + // Exempting BYOK from the *budget* must not hide it from usage reporting: + // the user in #5016 explicitly wanted to understand the counter. + let tmp = TempDir::new().unwrap(); + let tracker = CostTracker::new(enabled_config(), tmp.path()).unwrap(); + + let mut usage = TokenUsage::new(BYOK_MODEL, 1000, 500, 1.0, 1.0); + usage.cost_usd = 4.25; + tracker.record_usage(usage).unwrap(); + + let summary = tracker.get_summary().unwrap(); + assert_eq!(summary.request_count, 1); + assert!((summary.daily_cost_usd - 4.25).abs() < 0.0001); + assert!((summary.session_cost_usd - 4.25).abs() < 0.0001); +} + +#[test] +fn managed_spend_still_gates_when_byok_spend_is_also_present() { + // A mixed user: BYOK for chat, managed for background workloads. Only the + // managed portion may push them over the limit. + let tmp = TempDir::new().unwrap(); + let config = CostConfig { + enabled: true, + daily_limit_usd: 5.0, + monthly_limit_usd: 1000.0, + ..Default::default() + }; + let tracker = CostTracker::new(config, tmp.path()).unwrap(); + + let mut byok = TokenUsage::new(BYOK_MODEL, 1000, 500, 1.0, 1.0); + byok.cost_usd = 100.0; // dwarfs the limit, and must be ignored + tracker.record_usage(byok).unwrap(); + + let mut managed = TokenUsage::new(MANAGED_MODEL, 1000, 500, 1.0, 1.0); + managed.cost_usd = 2.0; // under the $5 limit on its own + tracker.record_usage(managed).unwrap(); + + assert!( + matches!(tracker.check_budget(0.0).unwrap(), BudgetCheck::Allowed), + "managed spend is under the limit; BYOK spend must not push it over" + ); + + let mut more_managed = TokenUsage::new(MANAGED_MODEL, 1000, 500, 1.0, 1.0); + more_managed.cost_usd = 4.0; // 2.0 + 4.0 = 6.0 > 5.0 + tracker.record_usage(more_managed).unwrap(); + + assert!( + matches!( + tracker.check_budget(0.0).unwrap(), + BudgetCheck::Exceeded { .. } + ), + "managed spend over the limit must still gate" + ); +} + +#[test] +fn legacy_byok_records_are_exempt_after_an_aggregate_rebuild() { + // Records persisted by builds that predate #5016 carry no route field. The + // route is derived from the model id they already store, so a tracker that + // rebuilds its aggregates from disk classifies them correctly with no + // migration — this is what unblocks an affected user on upgrade rather + // than making them wait out the window. + let tmp = TempDir::new().unwrap(); + let today = Utc::now(); + write_raw_record(tmp.path(), &dated_record("legacy", BYOK_MODEL, 50.0, today)); + + let config = CostConfig { + enabled: true, + daily_limit_usd: 10.0, + monthly_limit_usd: 10.0, + ..Default::default() + }; + let tracker = CostTracker::new(config, tmp.path()).unwrap(); + + assert!( + matches!(tracker.check_budget(0.0).unwrap(), BudgetCheck::Allowed), + "pre-existing BYOK history must not keep an upgraded user blocked" + ); + // …while still showing up in the usage figures. + let now = Utc::now(); + let monthly = tracker.get_monthly_cost(now.year(), now.month()).unwrap(); + assert!((monthly - 50.0).abs() < 0.0001); + let managed_monthly = tracker + .get_managed_monthly_cost(now.year(), now.month()) + .unwrap(); + assert!(managed_monthly.abs() < f64::EPSILON); +} + +#[test] +fn dashboard_budget_gauge_reflects_managed_spend_only() { + // The phantom "$10/day limit" in the issue was also visible as a budget + // gauge filling up from BYOK spend against a cap that could never fire. + let tmp = TempDir::new().unwrap(); + let today = Utc::now(); + write_raw_record(tmp.path(), &dated_record("s1", BYOK_MODEL, 95.0, today)); + + let config = CostConfig { + enabled: true, + monthly_limit_usd: 100.0, + ..Default::default() + }; + let tracker = CostTracker::new(config, tmp.path()).unwrap(); + let dash = tracker.get_dashboard("USD", 0.8, 0.95).unwrap(); + + // Usage is still reported… + assert!((dash.month_to_date_usd - 95.0).abs() < 0.0001); + assert!((dash.period_total_usd - 95.0).abs() < 0.0001); + // …but the budget gauge stays empty, because none of it is gateable. + assert_eq!(dash.budget_status, BudgetStatus::Normal); + assert!(dash.budget_utilization.abs() < f64::EPSILON); +} + #[test] fn get_daily_cost_for_today() { let tmp = TempDir::new().unwrap(); @@ -369,7 +588,9 @@ fn get_dashboard_computes_period_total_and_monthly_pace() { fn get_dashboard_budget_status_warning_and_exceeded() { let tmp = TempDir::new().unwrap(); let today = Utc::now(); - write_raw_record(tmp.path(), &dated_record("s1", "model-a", 85.0, today)); + // Managed-route spend: the budget gauge only tracks what can be gated + // (#5016), so these have to be managed tier ids to move the status. + write_raw_record(tmp.path(), &dated_record("s1", MANAGED_MODEL, 85.0, today)); let config = CostConfig { enabled: true, @@ -380,7 +601,7 @@ fn get_dashboard_budget_status_warning_and_exceeded() { let warn_dash = tracker.get_dashboard("USD", 0.8, 0.95).unwrap(); assert_eq!(warn_dash.budget_status, BudgetStatus::Warning); - write_raw_record(tmp.path(), &dated_record("s1", "model-a", 15.0, today)); + write_raw_record(tmp.path(), &dated_record("s1", MANAGED_MODEL, 15.0, today)); let tracker2 = CostTracker::new(config, tmp.path()).unwrap(); let alert_dash = tracker2.get_dashboard("USD", 0.8, 0.95).unwrap(); assert_eq!(alert_dash.budget_status, BudgetStatus::Exceeded); diff --git a/src/openhuman/tinyagents/middleware.rs b/src/openhuman/tinyagents/middleware.rs index 4326221d6..7d62556d1 100644 --- a/src/openhuman/tinyagents/middleware.rs +++ b/src/openhuman/tinyagents/middleware.rs @@ -1936,12 +1936,31 @@ impl Middleware<()> for CostBudgetMiddleware { &self, _ctx: &mut RunContext<()>, _state: &(), - _request: &mut ModelRequest, + request: &mut ModelRequest, ) -> TaResult<()> { use crate::openhuman::cost::types::BudgetCheck; let Some(tracker) = crate::openhuman::cost::try_global() else { return Ok(()); }; + + // #5016: exempt the CURRENT request when it is BYOK, not just BYOK + // history. Excluding BYOK from the managed totals alone still refuses a + // mixed-route user's own-key calls once their managed spend has + // legitimately crossed the cap — managed exhaustion would disable the + // provider OpenHuman never bills for, which is the whole bug. Classify + // this call's route and skip the gate when OpenHuman is not the biller. + if let Some(model) = request.model.as_deref() { + let route = crate::openhuman::cost::route::route_for_model(model); + if !route.counts_toward_budget() { + tracing::debug!( + %model, + ?route, + "[tinyagents::mw] BYOK/local route — skipping the managed budget gate (#5016)" + ); + return Ok(()); + } + } + // Pass 0.0 to test whether we are *already* over budget before spending // more (rather than projecting this call's cost, which needs a token // estimate).