From 518230ffa861ac04f3699ba74cbd11cc6cfe3dbf Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Wed, 30 Sep 2026 14:19:56 +0700 Subject: [PATCH 1/3] Port upstream 0.64.0: LiteLLM key-bound identity, user/team budgets, private-network HTTP --- rust/src/locale/en-US.ftl | 2 +- rust/src/providers/litellm/endpoint.rs | 172 +++++++++++++++ rust/src/providers/litellm/info.rs | 271 +++++++++++++++++++++++ rust/src/providers/litellm/mod.rs | 219 ++++++------------- rust/src/providers/litellm/tests.rs | 278 ++++++++++++++++++++++++ rust/src/settings/provider_workspace.rs | 28 ++- 6 files changed, 812 insertions(+), 158 deletions(-) create mode 100644 rust/src/providers/litellm/endpoint.rs create mode 100644 rust/src/providers/litellm/info.rs create mode 100644 rust/src/providers/litellm/tests.rs diff --git a/rust/src/locale/en-US.ftl b/rust/src/locale/en-US.ftl index fb96ecd4b8..ec1baafc9b 100644 --- a/rust/src/locale/en-US.ftl +++ b/rust/src/locale/en-US.ftl @@ -784,7 +784,7 @@ OpenAiProjectIdHelp = Leave blank for organization-wide usage. Set a project ID LiteLlmApiTitle = LiteLLM API LiteLlmBaseUrlLabel = Base URL LiteLlmBaseUrlPlaceholder = https://litellm.example.com -LiteLlmBaseUrlHelp = Used with the saved API key for LiteLLM /key/info. +LiteLlmBaseUrlHelp = Used with the saved API key for LiteLLM key, user, and team info. Use HTTPS, or HTTP on a loopback or private-network address. DevinApiTitle = Devin API DevinOrganizationLabel = Organization DevinOrganizationPlaceholder = org/acme diff --git a/rust/src/providers/litellm/endpoint.rs b/rust/src/providers/litellm/endpoint.rs new file mode 100644 index 0000000000..b175ed137e --- /dev/null +++ b/rust/src/providers/litellm/endpoint.rs @@ -0,0 +1,172 @@ +//! LiteLLM base-URL policy and management-route URLs. +//! +//! Upstream `litellm.ts` declares the `LITELLM_BASE_URL` endpoint with the +//! `https-or-private-network-http` policy: HTTPS anywhere, plain HTTP only for +//! loopback, RFC 1918, link-local, IPv6 unique-local, and `.local` hosts. + +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; + +use reqwest::Url; + +use crate::core::ProviderError; + +const INVALID_BASE: &str = "LiteLLM base URL must use HTTPS, or HTTP on a loopback or private-network address, without embedded credentials."; + +/// Validate a LiteLLM base URL. A scheme-less value is treated as HTTPS. +pub(crate) fn validated_base_url(raw: &str) -> Result { + let trimmed = raw.trim(); + if trimmed.is_empty() { + return Err(ProviderError::Other("LiteLLM base URL is empty".into())); + } + let lower = trimmed.to_ascii_lowercase(); + if ["%2f", "%5c", "%3f", "%23", "%40", "%3a"] + .iter() + .any(|encoded| lower.contains(encoded)) + { + return Err(ProviderError::Other( + "LiteLLM base URL must not contain encoded host delimiters".into(), + )); + } + let candidate = if trimmed.contains("://") { + trimmed.to_string() + } else { + format!("https://{trimmed}") + }; + let url = Url::parse(&candidate) + .map_err(|e| ProviderError::Other(format!("Invalid LiteLLM base URL: {e}")))?; + let host = url + .host_str() + .ok_or_else(|| ProviderError::Other("LiteLLM base URL must include a host".into()))?; + let scheme_ok = match url.scheme() { + "https" => true, + "http" => is_private_network_host(host), + _ => false, + }; + if !scheme_ok + || !url.username().is_empty() + || url.password().is_some() + || host.contains('%') + || host.chars().any(|c| c.is_control() || c.is_whitespace()) + { + return Err(ProviderError::Other(INVALID_BASE.into())); + } + Ok(url) +} + +/// Build `{base}/{path}` for a management route. A trailing `/v1` on the base +/// is dropped, and the base path and query are otherwise preserved. `query` +/// replaces the base query when given. +pub(super) fn management_url( + base: &str, + path: &str, + query: Option<(&str, &str)>, +) -> Result { + let mut url = validated_base_url(base)?; + let trimmed = url.path().trim_end_matches('/'); + let root = trimmed.strip_suffix("/v1").unwrap_or(trimmed).to_string(); + url.set_path(&format!("{root}/{path}")); + url.set_fragment(None); + if let Some((key, value)) = query { + url.query_pairs_mut().clear().append_pair(key, value); + } + Ok(url) +} + +fn is_private_network_host(host: &str) -> bool { + let normalized = host.trim_end_matches('.').to_ascii_lowercase(); + if normalized == "localhost" + || normalized.ends_with(".localhost") + || normalized.ends_with(".local") + { + return true; + } + let ip_candidate = normalized + .strip_prefix('[') + .and_then(|value| value.strip_suffix(']')) + .unwrap_or(&normalized); + match ip_candidate.parse::() { + Ok(IpAddr::V4(ip)) => is_private_ipv4(ip), + Ok(IpAddr::V6(ip)) => is_private_ipv6(ip), + Err(_) => false, + } +} + +fn is_private_ipv4(ip: Ipv4Addr) -> bool { + ip.is_loopback() || ip.is_private() || ip.is_link_local() +} + +fn is_private_ipv6(ip: Ipv6Addr) -> bool { + if let Some(mapped) = ip.to_ipv4_mapped() { + return is_private_ipv4(mapped); + } + ip.is_loopback() + || (ip.segments()[0] & 0xfe00) == 0xfc00 + || (ip.segments()[0] & 0xffc0) == 0xfe80 +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn allows_https_anywhere_and_private_network_http() { + for value in [ + "https://litellm.example.com", + "litellm.example.com", + "http://localhost:4000", + "http://127.0.0.1:4000", + "http://[::1]:4000", + "http://10.1.2.3", + "http://172.16.0.9", + "http://192.168.1.20:4000", + "http://169.254.10.10", + "http://[fd12:3456::1]", + "http://[fe80::1]", + "http://proxy.local:4000", + ] { + assert!(validated_base_url(value).is_ok(), "rejected {value}"); + } + } + + #[test] + fn rejects_public_http_credentials_and_encoded_delimiters() { + for value in [ + "", + "http://litellm.example.com", + "http://8.8.8.8", + "http://172.32.0.1", + "http://[2001:db8::1]", + "http://example.com.evil.test", + "ftp://10.0.0.1", + "https://user:pass@litellm.example.com", + "http://user@10.0.0.1", + "https://example.com%2f.evil.test", + ] { + assert!(validated_base_url(value).is_err(), "accepted {value}"); + } + } + + #[test] + fn management_url_strips_v1_and_keeps_subpath() { + let url = management_url("https://h.example.com/litellm/v1/", "key/info", None).unwrap(); + assert_eq!(url.as_str(), "https://h.example.com/litellm/key/info"); + let url = management_url("http://10.0.0.2:4000/v1", "key/info", None).unwrap(); + assert_eq!(url.as_str(), "http://10.0.0.2:4000/key/info"); + } + + #[test] + fn management_url_encodes_query_and_replaces_base_query() { + let url = management_url( + "https://h.example.com?token=abc", + "user/info", + Some(("user_id", "a b&c")), + ) + .unwrap(); + assert_eq!( + url.as_str(), + "https://h.example.com/user/info?user_id=a+b%26c" + ); + let kept = management_url("https://h.example.com?token=abc", "key/info", None).unwrap(); + assert_eq!(kept.as_str(), "https://h.example.com/key/info?token=abc"); + } +} diff --git a/rust/src/providers/litellm/info.rs b/rust/src/providers/litellm/info.rs new file mode 100644 index 0000000000..c308c41ede --- /dev/null +++ b/rust/src/providers/litellm/info.rs @@ -0,0 +1,271 @@ +//! LiteLLM management-route payloads and their projection into a usage result. +//! +//! Wire shapes follow upstream `litellm.ts`: `/key/info` names the key's +//! `user_id` / `team_id`, then `/user/info` or `/team/info` supplies the +//! budgets. Returned IDs must match the key's IDs before anything is shown. + +use chrono::{DateTime, NaiveDateTime, Utc}; +use serde::Deserialize; +use serde_json::Value; + +use crate::core::{ + CostSnapshot, ProviderError, ProviderFetchResult, RateWindow, SubscriptionMetadata, + UsageSnapshot, +}; + +#[derive(Deserialize)] +pub(super) struct KeyInfoResponse { + info: KeyInfo, +} + +#[derive(Deserialize)] +struct KeyInfo { + user_id: Option, + team_id: Option, + expires: Option, +} + +#[derive(Deserialize)] +pub(super) struct UserInfoResponse { + user_id: Option, + user_info: UserInfo, + teams: Option>, +} + +#[derive(Deserialize)] +struct UserInfo { + user_id: Option, + user_email: Option, + user_alias: Option, + spend: Option, + max_budget: Option, + budget_reset_at: Option, + metadata: Option, +} + +#[derive(Deserialize)] +struct UserMetadata { + preferred_username: Option, +} + +#[derive(Deserialize)] +pub(super) struct TeamInfoResponse { + team_id: Option, + team_info: Budget, +} + +#[derive(Deserialize)] +struct Budget { + team_id: Option, + team_alias: Option, + spend: Option, + max_budget: Option, + budget_reset_at: Option, +} + +/// Identity and routing data read from `/key/info`. +pub(super) struct KeyBinding { + pub user_id: Option, + pub team_id: Option, + expires: Option>, +} + +/// A spend/budget pair with its optional reset instant. +struct Spend { + spend: f64, + limit: Option, + reset: Option>, +} + +impl Spend { + fn budget(&self) -> Option { + self.limit.filter(|limit| *limit > 0.0) + } + + fn window(&self, label: Option<&str>) -> Option { + let limit = self.budget()?; + let mut window = RateWindow::new(self.spend / limit * 100.0); + window.resets_at = self.reset; + let detail = format!("${:.2} / ${limit:.2}", self.spend); + window.reset_description = Some(match label { + Some(label) => format!("{label}: {detail}"), + None => detail, + }); + Some(window) + } +} + +struct TeamBudget { + alias: Option, + spend: Spend, +} + +impl TeamBudget { + fn from_wire(budget: &Budget) -> Self { + Self { + alias: budget.team_alias.clone(), + spend: Spend { + spend: budget.spend.unwrap_or(0.0), + limit: budget.max_budget, + reset: parse_date(budget.budget_reset_at.as_deref()), + }, + } + } + + fn window(&self) -> Option { + let label = match &self.alias { + Some(alias) => format!("Team {alias}"), + None => "Team".to_string(), + }; + self.spend.window(Some(&label)) + } +} + +pub(super) fn parse_error(message: impl std::fmt::Display) -> ProviderError { + ProviderError::Parse(format!("LiteLLM parse error: {message}")) +} + +pub(super) fn bind_key(response: KeyInfoResponse) -> Result { + let info = response.info; + let user_id = nonempty(info.user_id); + let team_id = nonempty(info.team_id); + if user_id.is_none() && team_id.is_none() { + return Err(parse_error( + "LiteLLM key info did not include a user_id or team_id.", + )); + } + Ok(KeyBinding { + user_id, + team_id, + expires: parse_date(info.expires.as_deref()), + }) +} + +/// Project a user-bound key: personal budget plus the key's matching team. +pub(super) fn result_from_user( + key: &KeyBinding, + user_id: &str, + response: UserInfoResponse, +) -> Result { + let user = response.user_info; + let response_id = user.user_id.as_deref().or(response.user_id.as_deref()); + if response_id.is_some_and(|id| id != user_id) { + return Err(parse_error("user_id did not match /key/info")); + } + let preferred = user + .metadata + .and_then(|metadata| metadata.preferred_username) + .and_then(|value| value.as_str().map(str::to_owned)); + let email = nonempty(user.user_email) + .or_else(|| nonempty(user.user_alias)) + .or_else(|| nonempty(preferred)); + let mut team = None; + for wire in response.teams.unwrap_or_default() { + let id = wire + .team_id + .as_deref() + .ok_or_else(|| parse_error("missing team_id"))?; + if team.is_none() && key.team_id.as_deref() == Some(id) { + team = Some(TeamBudget::from_wire(&wire)); + } + } + let personal = Spend { + spend: user.spend.unwrap_or(0.0), + limit: user.max_budget, + reset: parse_date(user.budget_reset_at.as_deref()), + }; + let primary = personal + .window(None) + .unwrap_or_else(|| RateWindow::new(0.0)); + let mut snapshot = UsageSnapshot::new(primary); + if let Some(email) = email { + snapshot = snapshot.with_email(email); + } + if let Some(team) = &team { + if let Some(alias) = &team.alias { + snapshot = snapshot.with_organization(alias); + } + if let Some(window) = team.window() { + snapshot = snapshot.with_extra_rate_window("team", "Team budget", window); + } + } + Ok(finish(snapshot, key, &personal, "Personal")) +} + +/// Project a team-only key: the team budget is the sole usage window. +pub(super) fn result_from_team( + key: &KeyBinding, + team_id: &str, + response: TeamInfoResponse, +) -> Result { + let response_id = response + .team_info + .team_id + .as_deref() + .map(str::trim) + .filter(|id| !id.is_empty()) + .or(response + .team_id + .as_deref() + .filter(|id| !id.trim().is_empty())); + if response_id.is_some_and(|id| id != team_id) { + return Err(parse_error("team_id did not match /key/info")); + } + let team = TeamBudget::from_wire(&response.team_info); + let window = team.window().unwrap_or_else(|| RateWindow::new(0.0)); + let mut snapshot = UsageSnapshot::new(window).with_primary_label("Team budget"); + if let Some(alias) = &team.alias { + snapshot = snapshot.with_organization(alias); + } + Ok(finish(snapshot, key, &team.spend, "Team")) +} + +fn finish( + snapshot: UsageSnapshot, + key: &KeyBinding, + spend: &Spend, + scope: &str, +) -> ProviderFetchResult { + let mut snapshot = snapshot.with_login_method("api"); + if key.expires.is_some() { + snapshot = + snapshot.with_subscription(Some(SubscriptionMetadata::new(None, key.expires, None))); + } + let mut result = ProviderFetchResult::new(snapshot, "api"); + let limit = spend.budget(); + if spend.spend > 0.0 || limit.is_some() { + let kind = if limit.is_some() { "budget" } else { "spend" }; + let mut cost = CostSnapshot::new(spend.spend, "USD", format!("{scope} {kind}")); + if let Some(limit) = limit { + cost = cost.with_limit(limit); + } + if let Some(reset) = spend.reset { + cost = cost.with_resets_at(reset); + } + result = result.with_cost(cost); + } + result +} + +fn nonempty(value: Option) -> Option { + value + .map(|value| value.trim().to_string()) + .filter(|value| !value.is_empty()) +} + +/// Parse an ISO-8601 instant; naive timestamps are read as UTC and anything +/// unparseable is dropped, matching upstream's tolerant `date()` helper. +fn parse_date(value: Option<&str>) -> Option> { + let raw = value?.trim(); + if raw.is_empty() { + return None; + } + DateTime::parse_from_rfc3339(raw) + .map(|date| date.with_timezone(&Utc)) + .ok() + .or_else(|| { + NaiveDateTime::parse_from_str(raw, "%Y-%m-%dT%H:%M:%S%.f") + .ok() + .map(|date| date.and_utc()) + }) +} diff --git a/rust/src/providers/litellm/mod.rs b/rust/src/providers/litellm/mod.rs index d27a573668..c4e479a67a 100644 --- a/rust/src/providers/litellm/mod.rs +++ b/rust/src/providers/litellm/mod.rs @@ -1,13 +1,27 @@ use async_trait::async_trait; -use reqwest::{Client, Url}; -use serde_json::Value; +use reqwest::{Client, StatusCode, Url}; +use serde::de::DeserializeOwned; use crate::core::{ - CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, - ProviderMetadata, RateWindow, SourceMode, UsageSnapshot, + FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, ProviderMetadata, + SourceMode, +}; +use crate::providers::{BoundedBodyError, read_bounded_response}; + +mod endpoint; +mod info; +#[cfg(test)] +mod tests; + +use endpoint::management_url; +pub(crate) use endpoint::validated_base_url; +use info::{ + KeyInfoResponse, TeamInfoResponse, UserInfoResponse, bind_key, parse_error, result_from_team, + result_from_user, }; const CREDENTIAL_TARGET: &str = "codexbar-litellm"; +const MAX_RESPONSE_BYTES: usize = 1024 * 1024; pub struct LiteLLMProvider { metadata: ProviderMetadata, @@ -38,6 +52,40 @@ impl LiteLLMProvider { } } +impl LiteLLMProvider { + async fn get_json(&self, url: Url, key: &str) -> Result { + let route = url.path().to_string(); + let response = self + .client + .get(url) + .bearer_auth(key) + .header("Accept", "application/json") + .send() + .await?; + let status = response.status(); + if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN { + return Err(ProviderError::AuthRequired); + } + if status == StatusCode::TOO_MANY_REQUESTS { + return Err(ProviderError::Other( + "LiteLLM rate limited the request (HTTP 429).".into(), + )); + } + if !status.is_success() { + return Err(ProviderError::Other(format!( + "LiteLLM {route} returned status {status}" + ))); + } + let body = read_bounded_response(response, MAX_RESPONSE_BYTES) + .await + .map_err(|error| match error { + BoundedBodyError::TooLarge => parse_error("response too large"), + BoundedBodyError::Read(error) => ProviderError::Network(error), + })?; + serde_json::from_slice(&body).map_err(|e| parse_error(format!("{route}: {e}"))) + } +} + impl Default for LiteLLMProvider { fn default() -> Self { Self::new() @@ -58,28 +106,23 @@ impl Provider for LiteLLMProvider { match ctx.source_mode { SourceMode::Auto | SourceMode::OAuth => { let (base, key) = resolve_base_and_key(ctx)?; - let response = self - .client - .get(management_url(&base, "key/info")?) - .bearer_auth(key) - .header("Accept", "application/json") - .send() + let key_info: KeyInfoResponse = self + .get_json(management_url(&base, "key/info", None)?, &key) .await?; - if response.status() == reqwest::StatusCode::UNAUTHORIZED - || response.status() == reqwest::StatusCode::FORBIDDEN - { - return Err(ProviderError::AuthRequired); + let binding = bind_key(key_info)?; + if let Some(user_id) = binding.user_id.as_deref() { + let url = management_url(&base, "user/info", Some(("user_id", user_id)))?; + let response: UserInfoResponse = self.get_json(url, &key).await?; + result_from_user(&binding, user_id, response) + } else if let Some(team_id) = binding.team_id.as_deref() { + let url = management_url(&base, "team/info", Some(("team_id", team_id)))?; + let response: TeamInfoResponse = self.get_json(url, &key).await?; + result_from_team(&binding, team_id, response) + } else { + Err(parse_error( + "LiteLLM key info did not include a user_id or team_id.", + )) } - if !response.status().is_success() { - return Err(ProviderError::Other(format!( - "LiteLLM key/info returned status {}", - response.status() - ))); - } - let value: Value = response.json().await.map_err(|e| { - ProviderError::Parse(format!("Failed to parse LiteLLM key/info: {e}")) - })?; - Ok(result_from_key_info(&value)) } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -117,131 +160,3 @@ fn resolve_base_and_key(ctx: &FetchContext) -> Result<(String, String), Provider })?; Ok((base, key)) } - -fn management_url(base: &str, path: &str) -> Result { - let mut url = crate::providers::validated_https_url(base, "LiteLLM base")?; - if url.path().trim_end_matches('/').ends_with("/v1") { - let stripped = url - .path() - .trim_end_matches('/') - .trim_end_matches("/v1") - .to_string(); - url.set_path(&stripped); - } - url.join(path) - .map_err(|e| ProviderError::Other(format!("Invalid LiteLLM URL: {e}"))) -} - -fn result_from_key_info(value: &Value) -> ProviderFetchResult { - let root = value - .get("info") - .or_else(|| value.get("key")) - .unwrap_or(value); - let spend = number(root, &["spend", "spend_usd", "spendUSD"]).unwrap_or(0.0); - let limit = number(root, &["max_budget", "maxBudget", "budget", "limit"]); - let percent = limit - .filter(|v| *v > 0.0) - .map_or(0.0, |limit| spend / limit * 100.0); - let mut primary = RateWindow::new(percent); - if let Some(limit) = limit.filter(|value| *value > 0.0) { - primary.reset_description = Some(budget_detail(spend, limit)); - } - let mut snapshot = UsageSnapshot::new(primary).with_login_method(format!("Spend ${spend:.2}")); - if let Some(team) = root.get("team_info").or_else(|| root.get("teamInfo")) - && let Some(team_spend) = number(team, &["spend", "team_spend", "teamSpend"]) - { - let team_limit = number(team, &["max_budget", "budget", "limit"]); - let team_percent = team_limit - .filter(|v| *v > 0.0) - .map_or(0.0, |limit| team_spend / limit * 100.0); - let mut team_window = RateWindow::new(team_percent); - if let Some(team_limit) = team_limit.filter(|value| *value > 0.0) { - let alias = string(team, &["team_alias", "teamAlias", "alias"]) - .map(|value| format!("Team {value}: ")) - .unwrap_or_default(); - team_window.reset_description = - Some(format!("{alias}{}", budget_detail(team_spend, team_limit))); - } - snapshot = snapshot.with_extra_rate_window("team", "Team budget", team_window); - } - let mut result = ProviderFetchResult::new(snapshot, "api"); - if spend > 0.0 { - let mut cost = CostSnapshot::new(spend, "USD", "Spend"); - if let Some(limit) = limit { - cost = cost.with_limit(limit); - } - result = result.with_cost(cost); - } - result -} - -fn number(value: &Value, keys: &[&str]) -> Option { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_f64)) -} - -fn string(value: &Value, keys: &[&str]) -> Option { - keys.iter() - .find_map(|key| value.get(*key).and_then(Value::as_str)) - .map(str::trim) - .filter(|value| !value.is_empty()) - .map(ToOwned::to_owned) -} - -fn budget_detail(spend: f64, budget: f64) -> String { - format!("${spend:.2} / ${budget:.2}") -} - -#[cfg(test)] -mod tests { - use super::*; - - #[test] - fn parses_spend_budget() { - let result = - result_from_key_info(&serde_json::json!({"info":{"spend":25.0,"max_budget":100.0}})); - assert_eq!(result.usage.primary.used_percent, 25.0); - assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("$25.00 / $100.00") - ); - } - - #[test] - fn preserves_team_budget_detail_with_alias() { - let result = result_from_key_info(&serde_json::json!({ - "info": { - "team_info": { - "team_alias": "Platform", - "spend": 70.0, - "max_budget": 1000.0 - } - } - })); - assert_eq!(result.usage.extra_rate_windows.len(), 1); - assert_eq!( - result.usage.extra_rate_windows[0] - .window - .reset_description - .as_deref(), - Some("Team Platform: $70.00 / $1000.00") - ); - } - - #[test] - fn saved_base_url_uses_only_app_saved_key() { - let mut ctx = FetchContext { - workspace_id: Some("https://litellm.example.com".to_string()), - ..Default::default() - }; - assert!(matches!( - resolve_base_and_key(&ctx), - Err(ProviderError::AuthRequired) - )); - - ctx.api_key = Some("sk-app".to_string()); - let (base, key) = resolve_base_and_key(&ctx).unwrap(); - assert_eq!(base, "https://litellm.example.com"); - assert_eq!(key, "sk-app"); - } -} diff --git a/rust/src/providers/litellm/tests.rs b/rust/src/providers/litellm/tests.rs new file mode 100644 index 0000000000..3cd74e01ae --- /dev/null +++ b/rust/src/providers/litellm/tests.rs @@ -0,0 +1,278 @@ +use serde_json::{Value, json}; + +use super::info::{ + KeyBinding, KeyInfoResponse, TeamInfoResponse, UserInfoResponse, bind_key, result_from_team, + result_from_user, +}; +use super::*; + +fn binding(info: Value) -> KeyBinding { + bind_key(serde_json::from_value::(json!({ "info": info })).unwrap()).unwrap() +} + +fn user_result(key: &KeyBinding, body: Value) -> Result { + let user_id = key.user_id.clone().unwrap(); + result_from_user( + key, + &user_id, + serde_json::from_value::(body).unwrap(), + ) +} + +fn team_result(key: &KeyBinding, body: Value) -> Result { + let team_id = key.team_id.clone().unwrap(); + result_from_team( + key, + &team_id, + serde_json::from_value::(body).unwrap(), + ) +} + +fn assert_parse_error(result: Result, expected: &str) { + match result { + Err(ProviderError::Parse(message)) => { + assert!(message.contains(expected), "unexpected message: {message}") + } + Err(other) => panic!("expected parse error, got {other}"), + Ok(_) => panic!("expected parse error"), + } +} + +#[test] +fn key_info_without_user_or_team_id_fails() { + let response: KeyInfoResponse = + serde_json::from_value(json!({"info": {"user_id": " ", "spend": 1.0}})).unwrap(); + match bind_key(response) { + Err(ProviderError::Parse(message)) => assert!( + message.contains("LiteLLM key info did not include a user_id or team_id."), + "unexpected message: {message}" + ), + _ => panic!("expected parse error"), + } +} + +#[test] +fn key_info_requires_info_object() { + assert!(serde_json::from_value::(json!({"user_id": "u"})).is_err()); +} + +#[test] +fn personal_budget_is_primary_with_identity() { + let key = binding(json!({"user_id": "user-1", "expires": "2026-12-31T00:00:00Z"})); + let result = user_result( + &key, + json!({ + "user_id": "user-1", + "user_info": { + "user_id": "user-1", + "user_email": "dev@example.com", + "spend": 25.0, + "max_budget": 100.0, + "budget_reset_at": "2026-10-01T00:00:00Z" + }, + "teams": [] + }), + ) + .unwrap(); + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("$25.00 / $100.00") + ); + assert!(result.usage.primary.resets_at.is_some()); + assert_eq!( + result.usage.account_email.as_deref(), + Some("dev@example.com") + ); + assert_eq!(result.usage.login_method.as_deref(), Some("api")); + assert!(result.usage.extra_rate_windows.is_empty()); + assert!( + result + .usage + .subscription + .as_ref() + .is_some_and(|sub| sub.expires_at.is_some()) + ); + let cost = result.cost.expect("personal cost"); + assert_eq!(cost.used, 25.0); + assert_eq!(cost.limit, Some(100.0)); + assert_eq!(cost.period, "Personal budget"); +} + +#[test] +fn identity_falls_back_to_alias_then_preferred_username() { + let key = binding(json!({"user_id": "user-1"})); + let alias = user_result( + &key, + json!({"user_info": {"user_alias": "alias", "metadata": {"preferred_username": "pref"}}}), + ) + .unwrap(); + assert_eq!(alias.usage.account_email.as_deref(), Some("alias")); + let pref = user_result( + &key, + json!({"user_info": {"user_email": " ", "metadata": {"preferred_username": "pref"}}}), + ) + .unwrap(); + assert_eq!(pref.usage.account_email.as_deref(), Some("pref")); +} + +#[test] +fn matching_team_budget_is_a_separate_row() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-b"})); + let result = user_result( + &key, + json!({ + "user_info": {"user_id": "user-1", "spend": 3.0}, + "teams": [ + {"team_id": "team-a", "team_alias": "Other", "spend": 1.0, "max_budget": 10.0}, + {"team_id": "team-b", "team_alias": "Platform", "spend": 70.0, "max_budget": 1000.0} + ] + }), + ) + .unwrap(); + assert_eq!(result.usage.extra_rate_windows.len(), 1); + let team = &result.usage.extra_rate_windows[0].window; + assert!((team.used_percent - 7.0).abs() < 1e-9); + assert_eq!( + team.reset_description.as_deref(), + Some("Team Platform: $70.00 / $1000.00") + ); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Platform") + ); + let cost = result.cost.expect("spend-only cost"); + assert_eq!(cost.period, "Personal spend"); + assert_eq!(cost.limit, None); +} + +#[test] +fn team_without_a_matching_entry_is_omitted() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-x"})); + let result = user_result( + &key, + json!({ + "user_info": {"spend": 1.0}, + "teams": [{"team_id": "team-a", "spend": 1.0, "max_budget": 10.0}] + }), + ) + .unwrap(); + assert!(result.usage.extra_rate_windows.is_empty()); + assert_eq!(result.usage.account_organization, None); +} + +#[test] +fn mismatched_user_id_is_rejected() { + let key = binding(json!({"user_id": "user-1"})); + assert_parse_error( + user_result(&key, json!({"user_info": {"user_id": "user-2"}})), + "user_id did not match /key/info", + ); + assert_parse_error( + user_result(&key, json!({"user_id": "user-2", "user_info": {}})), + "user_id did not match /key/info", + ); +} + +#[test] +fn team_entries_without_team_id_are_rejected() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-a"})); + assert_parse_error( + user_result(&key, json!({"user_info": {}, "teams": [{"spend": 1.0}]})), + "missing team_id", + ); +} + +#[test] +fn wrongly_typed_fields_fail_to_parse() { + assert!( + serde_json::from_value::(json!({"user_info": {"spend": "12"}})).is_err() + ); + assert!(serde_json::from_value::(json!({"teams": []})).is_err()); +} + +#[test] +fn team_only_key_shows_team_budget_as_sole_window() { + let key = binding(json!({"team_id": "team-a"})); + let result = team_result( + &key, + json!({ + "team_id": "team-a", + "team_info": { + "team_id": "team-a", + "team_alias": "Platform", + "spend": 70.0, + "max_budget": 1000.0, + "budget_reset_at": "2026-10-01T00:00:00" + } + }), + ) + .unwrap(); + assert!((result.usage.primary.used_percent - 7.0).abs() < 1e-9); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("Team Platform: $70.00 / $1000.00") + ); + assert!(result.usage.primary.resets_at.is_some()); + assert_eq!(result.usage.primary_label.as_deref(), Some("Team budget")); + assert!(result.usage.extra_rate_windows.is_empty()); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Platform") + ); + assert_eq!(result.usage.account_email, None); + assert_eq!(result.cost.expect("team cost").period, "Team budget"); +} + +#[test] +fn mismatched_team_id_is_rejected() { + let key = binding(json!({"team_id": "team-a"})); + assert_parse_error( + team_result( + &key, + json!({"team_id": "team-b", "team_info": {"spend": 1.0}}), + ), + "team_id did not match /key/info", + ); + assert_parse_error( + team_result(&key, json!({"team_info": {"team_id": "team-b"}})), + "team_id did not match /key/info", + ); +} + +#[test] +fn spend_above_budget_clamps_percent_and_zero_budget_is_spend_only() { + let key = binding(json!({"user_id": "user-1"})); + let over = user_result( + &key, + json!({"user_info": {"spend": 150.0, "max_budget": 100.0}}), + ) + .unwrap(); + assert_eq!(over.usage.primary.used_percent, 100.0); + let unbudgeted = user_result( + &key, + json!({"user_info": {"spend": 4.0, "max_budget": 0.0}}), + ) + .unwrap(); + assert_eq!(unbudgeted.usage.primary.reset_description, None); + assert_eq!(unbudgeted.cost.expect("cost").period, "Personal spend"); + let empty = user_result(&key, json!({"user_info": {}})).unwrap(); + assert!(empty.cost.is_none()); +} + +#[test] +fn saved_base_url_uses_only_app_saved_key() { + let mut ctx = FetchContext { + workspace_id: Some("https://litellm.example.com".to_string()), + ..Default::default() + }; + assert!(matches!( + resolve_base_and_key(&ctx), + Err(ProviderError::AuthRequired) + )); + + ctx.api_key = Some("sk-app".to_string()); + let (base, key) = resolve_base_and_key(&ctx).unwrap(); + assert_eq!(base, "https://litellm.example.com"); + assert_eq!(key, "sk-app"); +} diff --git a/rust/src/settings/provider_workspace.rs b/rust/src/settings/provider_workspace.rs index 338747c8ef..e0e48fb4be 100644 --- a/rust/src/settings/provider_workspace.rs +++ b/rust/src/settings/provider_workspace.rs @@ -55,12 +55,19 @@ pub fn validate_provider_workspace_value( Err("Helmcode tenant must be 'helmcode' or 'nanBuilders'".to_string()) } } - ProviderId::LiteLLM => validate_token_endpoint(trimmed, "LiteLLM base URL", |_| true), + ProviderId::LiteLLM => validate_litellm_base_url(trimmed), ProviderId::Sub2Api => validate_sub2api_base_url(trimmed), _ => Ok(trimmed.to_string()), } } +fn validate_litellm_base_url(raw: &str) -> Result { + match crate::providers::litellm::validated_base_url(raw) { + Ok(url) => Ok(url.to_string().trim_end_matches('/').to_string()), + Err(err) => Err(err.to_string()), + } +} + fn validate_sub2api_base_url(raw: &str) -> Result { match crate::providers::sub2api::validated_sub2api_base_url(raw) { Ok(url) => Ok(url.to_string().trim_end_matches('/').to_string()), @@ -195,7 +202,7 @@ mod tests { } #[test] - fn validates_token_endpoint_hosts() { + fn validates_litellm_base_url_policy() { assert_eq!( validate_provider_workspace_value( ProviderId::LiteLLM, @@ -204,12 +211,23 @@ mod tests { .unwrap(), "https://litellm.example.com/v1" ); + for value in [ + "http://127.0.0.1:4000", + "http://10.0.0.5:4000", + "http://192.168.1.4", + "http://[::1]:4000", + "http://proxy.local:4000", + "https://10.0.0.5", + ] { + assert!( + validate_provider_workspace_value(ProviderId::LiteLLM, value).is_ok(), + "rejected {value}" + ); + } for value in [ "http://litellm.example.com", + "http://8.8.8.8", "https://user@litellm.example.com", - "https://127.0.0.1", - "https://10.0.0.5", - "https://[::1]", "https://example.com%2f.evil.test", ] { assert!( From 242af8a27f2cc5e40d2fa411f8fcc693f752e9ea Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Thu, 1 Oct 2026 15:27:01 +0700 Subject: [PATCH 2/3] Align LiteLLM budgets with upstream lanes Personal and team budgets fill the primary and secondary lanes, amounts are detail lines, Automatic prefers the team budget unless one is exhausted, and the private-network HTTP check matches upstream isPrivateNetworkHost. --- .../src-tauri/src/usage_metric.rs | 76 ++++ docs/PROVIDERS.md | 6 + rust/src/core/provider.rs | 7 + rust/src/providers/litellm/endpoint.rs | 16 +- rust/src/providers/litellm/info.rs | 118 ++++-- rust/src/providers/litellm/mod.rs | 12 +- rust/src/providers/litellm/tests.rs | 354 ++++++++++++++++-- 7 files changed, 521 insertions(+), 68 deletions(-) diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index 65a2408a15..49e0b95fe7 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -137,6 +137,20 @@ fn automatic_window( return None; } + if policy.prefers_secondary_window { + let primary = non_informational(Some(&snapshot.primary)); + let secondary = non_informational(snapshot.secondary.as_ref()); + let preferred = primary + .into_iter() + .chain(secondary) + .find(|window| automatic_window_is_exhausted(window)) + .or(secondary) + .or(primary); + if let Some(window) = preferred { + return Some(window.clone()); + } + } + let mut windows = Vec::with_capacity(4 + snapshot.extra_rate_windows.len()); windows.push(&snapshot.primary); windows.extend(snapshot.secondary.iter()); @@ -181,6 +195,9 @@ struct AutomaticMetricPolicy { /// is a dead end for Automatic selection. False for providers whose /// fallback lanes (seat credits) should still be considered. missing_core_is_terminal: bool, + /// Whether the secondary lane represents the provider unless a core lane + /// is exhausted (upstream's LiteLLM team-budget resolver). + prefers_secondary_window: bool, } fn automatic_metric_policy(provider: Option) -> AutomaticMetricPolicy { @@ -190,12 +207,16 @@ fn automatic_metric_policy(provider: Option) -> AutomaticMetricPolic let missing_core_is_terminal = |id: ProviderId| { codexbar::core::instantiate_provider(id).automatic_metric_missing_core_is_terminal() }; + let prefers_secondary = |id: ProviderId| { + codexbar::core::instantiate_provider(id).automatic_metric_prefers_secondary_window() + }; match provider { Some(ProviderId::Antigravity) => AutomaticMetricPolicy { prefers_available_window: true, prioritizes_exhausted_window: false, uses_extra_windows: false, missing_core_is_terminal: true, + prefers_secondary_window: false, }, // Cursor's monthly Auto lane is the semantic weekly pace. Grok Bot is // a named extra allowance and must stay available through the explicit @@ -205,18 +226,21 @@ fn automatic_metric_policy(provider: Option) -> AutomaticMetricPolic prioritizes_exhausted_window: prioritizes(ProviderId::Cursor), uses_extra_windows: false, missing_core_is_terminal: true, + prefers_secondary_window: false, }, Some(id) => AutomaticMetricPolicy { prefers_available_window: false, prioritizes_exhausted_window: prioritizes(id), uses_extra_windows: true, missing_core_is_terminal: missing_core_is_terminal(id), + prefers_secondary_window: prefers_secondary(id), }, None => AutomaticMetricPolicy { prefers_available_window: false, prioritizes_exhausted_window: true, uses_extra_windows: true, missing_core_is_terminal: false, + prefers_secondary_window: false, }, } } @@ -682,4 +706,56 @@ mod tests { ); assert!(restored.hidden_usage_item_ids.is_empty()); } + + fn litellm_budgets(personal: f64, team: Option) -> ProviderUsageSnapshot { + let mut snapshot = snapshot(); + snapshot.provider_id = "litellm".to_string(); + snapshot.primary = window(personal); + snapshot.secondary = team.map(window); + snapshot + } + + #[test] + fn litellm_automatic_prefers_the_team_budget_over_a_fuller_personal_budget() { + let snapshot = litellm_budgets(60.0, Some(7.0)); + + assert_eq!( + selected_usage_window(&snapshot, &Settings::default()).used_percent, + 7.0 + ); + let (selected, companion) = selected_usage_icon_windows(&snapshot, &Settings::default()); + assert_eq!(selected.used_percent, 7.0); + assert_eq!(companion.map(|window| window.used_percent), Some(60.0)); + } + + #[test] + fn litellm_automatic_shows_an_exhausted_budget_first() { + let personal_exhausted = litellm_budgets(100.0, Some(7.0)); + assert_eq!( + selected_usage_window(&personal_exhausted, &Settings::default()).used_percent, + 100.0 + ); + + let team_exhausted = litellm_budgets(40.0, Some(100.0)); + assert_eq!( + selected_usage_window(&team_exhausted, &Settings::default()).used_percent, + 100.0 + ); + } + + #[test] + fn litellm_automatic_uses_the_only_budget_and_explicit_choices_still_win() { + let personal_only = litellm_budgets(25.0, None); + assert_eq!( + selected_usage_window(&personal_only, &Settings::default()).used_percent, + 25.0 + ); + + let mut settings = Settings::default(); + settings.set_provider_metric(ProviderId::LiteLLM, MetricPreference::Session); + assert_eq!( + selected_usage_window(&litellm_budgets(60.0, Some(7.0)), &settings).used_percent, + 60.0 + ); + } } diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 237f7dab3d..9fd31a751d 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -93,6 +93,12 @@ z.ai Coding Plans accept both `TOKENS_LIMIT` and `CREDIT_LIMIT` rows. The shorte Upstream's independent **WidgetKit** provider-widget configuration has no Windows analogue in this repository. Win-CodexBar has no WidgetKit extension; provider cards and tray entries are already independent Windows/Tauri surfaces. +### LiteLLM budgets + +LiteLLM reads a virtual key's own budgets from the proxy's management routes. Set the base URL and key in Settings → Providers → LiteLLM, or use `LITELLM_BASE_URL` and `LITELLM_API_KEY`. A trailing `/v1` is dropped. The base URL must use HTTPS unless it names `localhost`, a `.local` host, or a loopback, RFC 1918, link-local or IPv6 unique-local address, and it must not embed credentials, because the key is sent as a bearer token. + +The provider calls `GET /key/info`, then `GET /user/info?user_id=…` for a user-bound key or `GET /team/info?team_id=…` for a team-only key, and rejects a response whose user or team ID differs from the key's. The personal budget fills the primary lane and the key's matching team budget fills the secondary lane. When only one budget exists, it takes the primary lane under its own label, and a key without any budget shows "No budget set". Amounts such as `$25.00 / $100.00` are detail lines, shown apart from a real reset date. The Automatic tray and float bar metric shows the team budget unless a budget is exhausted. Spend without a budget stays visible as API spend, and no pace is derived from budget resets. + ## Upstream doc warning Upstream `docs/providers.md` is a large auto-strategy matrix (60+ providers) for the macOS app. Use it as **inspiration** when porting a provider. For runtime truth on Windows: diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index cc954839f7..4add100e20 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -892,6 +892,13 @@ pub trait Provider: Send + Sync { true } + /// Whether Automatic metric selection shows the secondary lane whenever + /// neither core lane is exhausted, instead of the fuller lane. An + /// exhausted primary or secondary lane still wins, primary first. + fn automatic_metric_prefers_secondary_window(&self) -> bool { + false + } + /// Whether browser-cookie discovery/recovery is owned by the provider. fn owns_browser_cookie_resolution(&self) -> bool { false diff --git a/rust/src/providers/litellm/endpoint.rs b/rust/src/providers/litellm/endpoint.rs index b175ed137e..808509e40b 100644 --- a/rust/src/providers/litellm/endpoint.rs +++ b/rust/src/providers/litellm/endpoint.rs @@ -72,11 +72,17 @@ pub(super) fn management_url( Ok(url) } +/// Upstream `isPrivateNetworkHost`: `localhost`, `.local` names, and IP +/// literals that are loopback, RFC 1918, link-local, or IPv6 unique-local. +/// Other names (including `*.localhost`) and IPv4-mapped IPv6 literals stay +/// HTTPS-only, because a bearer key would otherwise cross the network in +/// plain text if the name resolved somewhere public. fn is_private_network_host(host: &str) -> bool { let normalized = host.trim_end_matches('.').to_ascii_lowercase(); if normalized == "localhost" - || normalized.ends_with(".localhost") - || normalized.ends_with(".local") + || normalized + .strip_suffix(".local") + .is_some_and(|label| !label.is_empty()) { return true; } @@ -96,9 +102,6 @@ fn is_private_ipv4(ip: Ipv4Addr) -> bool { } fn is_private_ipv6(ip: Ipv6Addr) -> bool { - if let Some(mapped) = ip.to_ipv4_mapped() { - return is_private_ipv4(mapped); - } ip.is_loopback() || (ip.segments()[0] & 0xfe00) == 0xfc00 || (ip.segments()[0] & 0xffc0) == 0xfe80 @@ -141,6 +144,9 @@ mod tests { "https://user:pass@litellm.example.com", "http://user@10.0.0.1", "https://example.com%2f.evil.test", + "http://app.localhost:4000", + "http://[::ffff:10.0.0.1]", + "http://.local:4000", ] { assert!(validated_base_url(value).is_err(), "accepted {value}"); } diff --git a/rust/src/providers/litellm/info.rs b/rust/src/providers/litellm/info.rs index c308c41ede..10065361be 100644 --- a/rust/src/providers/litellm/info.rs +++ b/rust/src/providers/litellm/info.rs @@ -3,8 +3,15 @@ //! Wire shapes follow upstream `litellm.ts`: `/key/info` names the key's //! `user_id` / `team_id`, then `/user/info` or `/team/info` supplies the //! budgets. Returned IDs must match the key's IDs before anything is shown. +//! +//! Upstream reports the personal budget as the primary window and the key's +//! team budget as the secondary window, and either may be absent. A Windows +//! snapshot always has a primary lane, so the windows that exist fill the +//! lanes in that order (personal, then team) with their own labels, and a key +//! with no budget at all shows an informational "No budget set" lane. +//! Budget amounts are detail lines, never reset wording (upstream #3631). -use chrono::{DateTime, NaiveDateTime, Utc}; +use chrono::{DateTime, NaiveDate, NaiveDateTime, Utc}; use serde::Deserialize; use serde_json::Value; @@ -23,6 +30,11 @@ struct KeyInfo { user_id: Option, team_id: Option, expires: Option, + /// Type-checked like upstream, but not otherwise used. + #[serde(rename = "key_name")] + _key_name: Option, + #[serde(rename = "spend")] + _spend: Option, } #[derive(Deserialize)] @@ -61,6 +73,9 @@ struct Budget { spend: Option, max_budget: Option, budget_reset_at: Option, + /// Type-checked like upstream, but not otherwise used. + #[serde(rename = "budget_duration")] + _budget_duration: Option, } /// Identity and routing data read from `/key/info`. @@ -82,16 +97,19 @@ impl Spend { self.limit.filter(|limit| *limit > 0.0) } + /// The budget window, or `None` without a positive budget. The amount is + /// a detail line, so a missing reset never reads as "Resets $x / $y" and + /// a real reset never hides the amount. fn window(&self, label: Option<&str>) -> Option { let limit = self.budget()?; let mut window = RateWindow::new(self.spend / limit * 100.0); window.resets_at = self.reset; - let detail = format!("${:.2} / ${limit:.2}", self.spend); + let detail = format!("{} / {}", usd(self.spend), usd(limit)); window.reset_description = Some(match label { Some(label) => format!("{label}: {detail}"), None => detail, }); - Some(window) + Some(window.with_description_as_detail()) } } @@ -103,7 +121,7 @@ struct TeamBudget { impl TeamBudget { fn from_wire(budget: &Budget) -> Self { Self { - alias: budget.team_alias.clone(), + alias: nonempty(budget.team_alias.clone()), spend: Spend { spend: budget.spend.unwrap_or(0.0), limit: budget.max_budget, @@ -121,6 +139,11 @@ impl TeamBudget { } } +/// Lane label for a team budget shown in the primary lane. +const TEAM_BUDGET: &str = "Team budget"; +/// Informational lane text for a key without a personal or team budget. +const NO_BUDGET: &str = "No budget set"; + pub(super) fn parse_error(message: impl std::fmt::Display) -> ProviderError { ProviderError::Parse(format!("LiteLLM parse error: {message}")) } @@ -174,20 +197,23 @@ pub(super) fn result_from_user( limit: user.max_budget, reset: parse_date(user.budget_reset_at.as_deref()), }; - let primary = personal - .window(None) - .unwrap_or_else(|| RateWindow::new(0.0)); - let mut snapshot = UsageSnapshot::new(primary); + let team_window = team.as_ref().and_then(TeamBudget::window); + let mut snapshot = match (personal.window(None), team_window) { + (Some(personal), team) => { + let mut snapshot = UsageSnapshot::new(personal); + if let Some(team) = team { + snapshot = snapshot.with_secondary(team); + } + snapshot + } + (None, Some(team)) => UsageSnapshot::new(team).with_primary_label(TEAM_BUDGET), + (None, None) => UsageSnapshot::new(RateWindow::informational(NO_BUDGET)), + }; if let Some(email) = email { snapshot = snapshot.with_email(email); } - if let Some(team) = &team { - if let Some(alias) = &team.alias { - snapshot = snapshot.with_organization(alias); - } - if let Some(window) = team.window() { - snapshot = snapshot.with_extra_rate_window("team", "Team budget", window); - } + if let Some(alias) = team.as_ref().and_then(|team| team.alias.as_deref()) { + snapshot = snapshot.with_organization(alias); } Ok(finish(snapshot, key, &personal, "Personal")) } @@ -198,22 +224,16 @@ pub(super) fn result_from_team( team_id: &str, response: TeamInfoResponse, ) -> Result { - let response_id = response - .team_info - .team_id - .as_deref() - .map(str::trim) - .filter(|id| !id.is_empty()) - .or(response - .team_id - .as_deref() - .filter(|id| !id.trim().is_empty())); + let response_id = trimmed_id(response.team_info.team_id.as_deref()) + .or_else(|| trimmed_id(response.team_id.as_deref())); if response_id.is_some_and(|id| id != team_id) { return Err(parse_error("team_id did not match /key/info")); } let team = TeamBudget::from_wire(&response.team_info); - let window = team.window().unwrap_or_else(|| RateWindow::new(0.0)); - let mut snapshot = UsageSnapshot::new(window).with_primary_label("Team budget"); + let window = team + .window() + .unwrap_or_else(|| RateWindow::informational(NO_BUDGET)); + let mut snapshot = UsageSnapshot::new(window).with_primary_label(TEAM_BUDGET); if let Some(alias) = &team.alias { snapshot = snapshot.with_organization(alias); } @@ -231,13 +251,18 @@ fn finish( snapshot = snapshot.with_subscription(Some(SubscriptionMetadata::new(None, key.expires, None))); } - let mut result = ProviderFetchResult::new(snapshot, "api"); + // Budget resets carry no window length, so the shell's weekly pace + // defaults would invent a pace upstream never shows. + let mut result = ProviderFetchResult::new(snapshot, "api").with_non_authoritative_pace(); let limit = spend.budget(); if spend.spend > 0.0 || limit.is_some() { let kind = if limit.is_some() { "budget" } else { "spend" }; let mut cost = CostSnapshot::new(spend.spend, "USD", format!("{scope} {kind}")); - if let Some(limit) = limit { - cost = cost.with_limit(limit); + match limit { + Some(limit) => cost = cost.with_limit(limit), + // Upstream's "API spend" card: spend without a budget is the + // only usage signal, so it stays visible like other API spend. + None => cost = cost.always_visible(), } if let Some(reset) = spend.reset { cost = cost.with_resets_at(reset); @@ -247,14 +272,37 @@ fn finish( result } +/// Upstream `ctx.format.usd`: two decimals with thousands separators, and +/// `-$` for negative amounts. +fn usd(amount: f64) -> String { + let fixed = format!("{:.2}", amount.abs()); + let (whole, cents) = fixed.split_once('.').unwrap_or((&fixed, "00")); + let mut grouped = String::with_capacity(whole.len() + whole.len() / 3); + for (index, digit) in whole.chars().enumerate() { + if index > 0 && (whole.len() - index) % 3 == 0 { + grouped.push(','); + } + grouped.push(digit); + } + let sign = if amount < 0.0 { "-$" } else { "$" }; + format!("{sign}{grouped}.{cents}") +} + +fn trimmed_id(id: Option<&str>) -> Option<&str> { + id.map(str::trim).filter(|id| !id.is_empty()) +} + fn nonempty(value: Option) -> Option { value .map(|value| value.trim().to_string()) .filter(|value| !value.is_empty()) } -/// Parse an ISO-8601 instant; naive timestamps are read as UTC and anything -/// unparseable is dropped, matching upstream's tolerant `date()` helper. +/// Parse an ISO-8601 instant and drop anything unparseable, like upstream's +/// tolerant `date()` helper. A date-only value is UTC midnight, as in +/// JavaScript. A timestamp without an offset is also read as UTC, because +/// LiteLLM stores its budget times in UTC; upstream's JavaScript `Date` would +/// read it in the machine's local zone instead. fn parse_date(value: Option<&str>) -> Option> { let raw = value?.trim(); if raw.is_empty() { @@ -268,4 +316,10 @@ fn parse_date(value: Option<&str>) -> Option> { .ok() .map(|date| date.and_utc()) }) + .or_else(|| { + NaiveDate::parse_from_str(raw, "%Y-%m-%d") + .ok() + .and_then(|date| date.and_hms_opt(0, 0, 0)) + .map(|date| date.and_utc()) + }) } diff --git a/rust/src/providers/litellm/mod.rs b/rust/src/providers/litellm/mod.rs index c4e479a67a..5516596484 100644 --- a/rust/src/providers/litellm/mod.rs +++ b/rust/src/providers/litellm/mod.rs @@ -34,10 +34,10 @@ impl LiteLLMProvider { metadata: ProviderMetadata { id: ProviderId::LiteLLM, display_name: "LiteLLM", - session_label: "Budget", - weekly_label: "Spend", + session_label: "Personal budget", + weekly_label: "Team budget", supports_opus: false, - supports_credits: true, + supports_credits: false, default_enabled: false, is_primary: false, dashboard_url: None, @@ -102,6 +102,12 @@ impl Provider for LiteLLMProvider { &self.metadata } + /// Upstream's LiteLLM menu bar resolver: the team budget is enforced for + /// the key, so Automatic shows it unless a budget is already exhausted. + fn automatic_metric_prefers_secondary_window(&self) -> bool { + true + } + async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { SourceMode::Auto | SourceMode::OAuth => { diff --git a/rust/src/providers/litellm/tests.rs b/rust/src/providers/litellm/tests.rs index 3cd74e01ae..111b68446a 100644 --- a/rust/src/providers/litellm/tests.rs +++ b/rust/src/providers/litellm/tests.rs @@ -1,5 +1,8 @@ +use chrono::{TimeZone, Utc}; use serde_json::{Value, json}; +use crate::core::RateWindow; + use super::info::{ KeyBinding, KeyInfoResponse, TeamInfoResponse, UserInfoResponse, bind_key, result_from_team, result_from_user, @@ -28,6 +31,27 @@ fn team_result(key: &KeyBinding, body: Value) -> Result ProviderFetchResult { + let key = binding(json!({"user_id": "user-1"})); + user_result( + &key, + json!({"user_info": {"spend": spend, "max_budget": max_budget}}), + ) + .unwrap() +} + +fn personal_reset(budget_reset_at: &str) -> Option> { + let key = binding(json!({"user_id": "user-1"})); + user_result( + &key, + json!({"user_info": {"spend": 1.0, "max_budget": 10.0, "budget_reset_at": budget_reset_at}}), + ) + .unwrap() + .usage + .primary + .resets_at +} + fn assert_parse_error(result: Result, expected: &str) { match result { Err(ProviderError::Parse(message)) => { @@ -38,6 +62,13 @@ fn assert_parse_error(result: Result, expect } } +fn assert_no_budget(window: &RateWindow) { + assert!(window.is_informational); + assert!(!window.usage_known()); + assert_eq!(window.reset_description.as_deref(), Some("No budget set")); + assert_eq!(window.resets_at, None); +} + #[test] fn key_info_without_user_or_team_id_fails() { let response: KeyInfoResponse = @@ -56,6 +87,16 @@ fn key_info_requires_info_object() { assert!(serde_json::from_value::(json!({"user_id": "u"})).is_err()); } +#[test] +fn metadata_names_budget_lanes_and_prefers_the_team_lane() { + let provider = LiteLLMProvider::new(); + let metadata = provider.metadata(); + assert_eq!(metadata.session_label, "Personal budget"); + assert_eq!(metadata.weekly_label, "Team budget"); + assert!(!metadata.supports_credits); + assert!(provider.automatic_metric_prefers_secondary_window()); +} + #[test] fn personal_budget_is_primary_with_identity() { let key = binding(json!({"user_id": "user-1", "expires": "2026-12-31T00:00:00Z"})); @@ -74,18 +115,26 @@ fn personal_budget_is_primary_with_identity() { }), ) .unwrap(); - assert_eq!(result.usage.primary.used_percent, 25.0); + let primary = &result.usage.primary; + assert_eq!(primary.used_percent, 25.0); assert_eq!( - result.usage.primary.reset_description.as_deref(), + primary.reset_description.as_deref(), Some("$25.00 / $100.00") ); - assert!(result.usage.primary.resets_at.is_some()); + assert!(primary.description_is_detail); + assert!(!primary.is_informational); + assert_eq!( + primary.resets_at, + Some(Utc.with_ymd_and_hms(2026, 10, 1, 0, 0, 0).unwrap()) + ); + assert_eq!(result.usage.primary_label, None); + assert!(result.usage.secondary.is_none()); + assert!(result.usage.extra_rate_windows.is_empty()); assert_eq!( result.usage.account_email.as_deref(), Some("dev@example.com") ); assert_eq!(result.usage.login_method.as_deref(), Some("api")); - assert!(result.usage.extra_rate_windows.is_empty()); assert!( result .usage @@ -93,10 +142,12 @@ fn personal_budget_is_primary_with_identity() { .as_ref() .is_some_and(|sub| sub.expires_at.is_some()) ); + assert!(!result.pace_authoritative); let cost = result.cost.expect("personal cost"); assert_eq!(cost.used, 25.0); assert_eq!(cost.limit, Some(100.0)); assert_eq!(cost.period, "Personal budget"); + assert!(!cost.always_visible); } #[test] @@ -114,36 +165,97 @@ fn identity_falls_back_to_alias_then_preferred_username() { ) .unwrap(); assert_eq!(pref.usage.account_email.as_deref(), Some("pref")); + let non_string = user_result( + &key, + json!({"user_info": {"metadata": {"preferred_username": 7}}}), + ) + .unwrap(); + assert_eq!(non_string.usage.account_email, None); } #[test] -fn matching_team_budget_is_a_separate_row() { +fn personal_and_team_budgets_fill_primary_and_secondary() { let key = binding(json!({"user_id": "user-1", "team_id": "team-b"})); let result = user_result( &key, json!({ - "user_info": {"user_id": "user-1", "spend": 3.0}, + "user_info": {"user_id": "user-1", "spend": 25.0, "max_budget": 100.0}, "teams": [ {"team_id": "team-a", "team_alias": "Other", "spend": 1.0, "max_budget": 10.0}, - {"team_id": "team-b", "team_alias": "Platform", "spend": 70.0, "max_budget": 1000.0} + { + "team_id": "team-b", + "team_alias": "Platform", + "spend": 70.0, + "max_budget": 1000.0, + "budget_reset_at": "2026-11-01", + "budget_duration": "30d" + } ] }), ) .unwrap(); - assert_eq!(result.usage.extra_rate_windows.len(), 1); - let team = &result.usage.extra_rate_windows[0].window; + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!( + result.usage.primary.reset_description.as_deref(), + Some("$25.00 / $100.00") + ); + let team = result.usage.secondary.as_ref().expect("team lane"); assert!((team.used_percent - 7.0).abs() < 1e-9); assert_eq!( team.reset_description.as_deref(), - Some("Team Platform: $70.00 / $1000.00") + Some("Team Platform: $70.00 / $1,000.00") ); + assert!(team.description_is_detail); + assert_eq!( + team.resets_at, + Some(Utc.with_ymd_and_hms(2026, 11, 1, 0, 0, 0).unwrap()) + ); + // The metadata labels ("Personal budget" / "Team budget") apply. + assert_eq!(result.usage.primary_label, None); + assert_eq!(result.usage.secondary_label, None); + assert!(result.usage.extra_rate_windows.is_empty()); assert_eq!( result.usage.account_organization.as_deref(), Some("Platform") ); + let cost = result.cost.expect("personal cost"); + assert_eq!(cost.period, "Personal budget"); + assert_eq!(cost.limit, Some(100.0)); +} + +#[test] +fn team_budget_fills_the_primary_lane_without_a_personal_budget() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-b"})); + let result = user_result( + &key, + json!({ + "user_info": {"user_id": "user-1", "spend": 3.0}, + "teams": [ + {"team_id": "team-a", "team_alias": "Other", "spend": 1.0, "max_budget": 10.0}, + {"team_id": "team-b", "team_alias": "Platform", "spend": 70.0, "max_budget": 1000.0} + ] + }), + ) + .unwrap(); + let primary = &result.usage.primary; + assert!((primary.used_percent - 7.0).abs() < 1e-9); + assert_eq!( + primary.reset_description.as_deref(), + Some("Team Platform: $70.00 / $1,000.00") + ); + assert!(primary.description_is_detail); + assert_eq!(result.usage.primary_label.as_deref(), Some("Team budget")); + assert!(result.usage.secondary.is_none()); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Platform") + ); + // The cost stays scoped to the key's own (personal) spend. let cost = result.cost.expect("spend-only cost"); + assert_eq!(cost.used, 3.0); assert_eq!(cost.period, "Personal spend"); assert_eq!(cost.limit, None); + assert!(cost.always_visible); } #[test] @@ -157,10 +269,61 @@ fn team_without_a_matching_entry_is_omitted() { }), ) .unwrap(); + assert_no_budget(&result.usage.primary); + assert_eq!(result.usage.primary_label, None); + assert!(result.usage.secondary.is_none()); assert!(result.usage.extra_rate_windows.is_empty()); assert_eq!(result.usage.account_organization, None); } +#[test] +fn team_entries_match_the_key_team_id_exactly() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-a"})); + let result = user_result( + &key, + json!({ + "user_info": {"spend": 5.0, "max_budget": 50.0}, + "teams": [{"team_id": " team-a ", "team_alias": "Padded", "spend": 1.0, "max_budget": 10.0}] + }), + ) + .unwrap(); + assert!(result.usage.secondary.is_none()); + assert_eq!(result.usage.account_organization, None); +} + +#[test] +fn team_without_budget_or_alias_adds_no_lane_or_organization() { + let key = binding(json!({"user_id": "user-1", "team_id": "team-a"})); + let unbudgeted = user_result( + &key, + json!({ + "user_info": {"spend": 5.0, "max_budget": 50.0}, + "teams": [{"team_id": "team-a", "team_alias": "Platform", "spend": 9.0}] + }), + ) + .unwrap(); + assert!(unbudgeted.usage.secondary.is_none()); + assert_eq!( + unbudgeted.usage.account_organization.as_deref(), + Some("Platform") + ); + + let blank_alias = user_result( + &key, + json!({ + "user_info": {"spend": 5.0, "max_budget": 50.0}, + "teams": [{"team_id": "team-a", "team_alias": " ", "spend": 9.0, "max_budget": 90.0}] + }), + ) + .unwrap(); + let team = blank_alias.usage.secondary.as_ref().expect("team lane"); + assert_eq!( + team.reset_description.as_deref(), + Some("Team: $9.00 / $90.00") + ); + assert_eq!(blank_alias.usage.account_organization, None); +} + #[test] fn mismatched_user_id_is_rejected() { let key = binding(json!({"user_id": "user-1"})); @@ -189,6 +352,34 @@ fn wrongly_typed_fields_fail_to_parse() { serde_json::from_value::(json!({"user_info": {"spend": "12"}})).is_err() ); assert!(serde_json::from_value::(json!({"teams": []})).is_err()); + assert!( + serde_json::from_value::(json!({"user_info": {}, "teams": {}})).is_err() + ); + assert!( + serde_json::from_value::(json!({ + "user_info": {}, + "teams": [{"team_id": "team-a", "budget_duration": 30}] + })) + .is_err() + ); + assert!( + serde_json::from_value::(json!({"info": {"user_id": "u", "key_name": 5}})) + .is_err() + ); + assert!( + serde_json::from_value::(json!({"info": {"user_id": "u", "spend": "1"}})) + .is_err() + ); + assert!( + serde_json::from_value::(json!({"team_info": {"budget_duration": 30}})) + .is_err() + ); + assert!( + serde_json::from_value::(json!({ + "info": {"user_id": "u", "key_name": "ci", "spend": 1.5} + })) + .is_ok() + ); } #[test] @@ -208,20 +399,49 @@ fn team_only_key_shows_team_budget_as_sole_window() { }), ) .unwrap(); - assert!((result.usage.primary.used_percent - 7.0).abs() < 1e-9); + let primary = &result.usage.primary; + assert!((primary.used_percent - 7.0).abs() < 1e-9); assert_eq!( - result.usage.primary.reset_description.as_deref(), - Some("Team Platform: $70.00 / $1000.00") + primary.reset_description.as_deref(), + Some("Team Platform: $70.00 / $1,000.00") + ); + assert!(primary.description_is_detail); + assert_eq!( + primary.resets_at, + Some(Utc.with_ymd_and_hms(2026, 10, 1, 0, 0, 0).unwrap()) ); - assert!(result.usage.primary.resets_at.is_some()); assert_eq!(result.usage.primary_label.as_deref(), Some("Team budget")); + assert!(result.usage.secondary.is_none()); assert!(result.usage.extra_rate_windows.is_empty()); assert_eq!( result.usage.account_organization.as_deref(), Some("Platform") ); assert_eq!(result.usage.account_email, None); - assert_eq!(result.cost.expect("team cost").period, "Team budget"); + assert!(!result.pace_authoritative); + let cost = result.cost.expect("team cost"); + assert_eq!(cost.period, "Team budget"); + assert_eq!(cost.limit, Some(1000.0)); + assert!(!cost.always_visible); +} + +#[test] +fn team_only_key_without_a_budget_reports_spend_only() { + let key = binding(json!({"team_id": "team-a"})); + let result = team_result( + &key, + json!({"team_info": {"team_id": "team-a", "spend": 12.5}}), + ) + .unwrap(); + assert_no_budget(&result.usage.primary); + assert_eq!(result.usage.primary_label.as_deref(), Some("Team budget")); + assert!(result.usage.secondary.is_none()); + assert_eq!(result.usage.account_organization, None); + let cost = result.cost.expect("team spend"); + assert_eq!(cost.used, 12.5); + assert_eq!(cost.period, "Team spend"); + assert_eq!(cost.limit, None); + assert!(cost.always_visible); } #[test] @@ -238,28 +458,106 @@ fn mismatched_team_id_is_rejected() { team_result(&key, json!({"team_info": {"team_id": "team-b"}})), "team_id did not match /key/info", ); + // A blank nested id falls back to the root id. + assert_parse_error( + team_result( + &key, + json!({"team_id": "team-b", "team_info": {"team_id": " "}}), + ), + "team_id did not match /key/info", + ); +} + +#[test] +fn team_ids_in_team_info_are_trimmed() { + let key = binding(json!({"team_id": "team-a"})); + assert!( + team_result( + &key, + json!({"team_id": " team-a ", "team_info": {"spend": 1.0}}) + ) + .is_ok() + ); + assert!(team_result(&key, json!({"team_info": {"team_id": "team-a "}})).is_ok()); + assert!(team_result(&key, json!({"team_info": {}})).is_ok()); } #[test] fn spend_above_budget_clamps_percent_and_zero_budget_is_spend_only() { - let key = binding(json!({"user_id": "user-1"})); - let over = user_result( - &key, - json!({"user_info": {"spend": 150.0, "max_budget": 100.0}}), - ) - .unwrap(); + let over = personal_budget(150.0, 100.0); assert_eq!(over.usage.primary.used_percent, 100.0); - let unbudgeted = user_result( - &key, - json!({"user_info": {"spend": 4.0, "max_budget": 0.0}}), - ) - .unwrap(); - assert_eq!(unbudgeted.usage.primary.reset_description, None); - assert_eq!(unbudgeted.cost.expect("cost").period, "Personal spend"); + assert_eq!( + over.usage.primary.reset_description.as_deref(), + Some("$150.00 / $100.00") + ); + + for limit in [0.0, -5.0] { + let unbudgeted = personal_budget(4.0, limit); + assert_no_budget(&unbudgeted.usage.primary); + let cost = unbudgeted.cost.expect("cost"); + assert_eq!(cost.period, "Personal spend"); + assert_eq!(cost.limit, None); + assert!(cost.always_visible); + } + + let key = binding(json!({"user_id": "user-1"})); let empty = user_result(&key, json!({"user_info": {}})).unwrap(); + assert_no_budget(&empty.usage.primary); assert!(empty.cost.is_none()); } +#[test] +fn amounts_use_grouped_dollars_like_upstream() { + assert_eq!( + personal_budget(1_234_567.891, 2_000_000.0) + .usage + .primary + .reset_description + .as_deref(), + Some("$1,234,567.89 / $2,000,000.00") + ); + assert_eq!( + personal_budget(999.999, 1000.0) + .usage + .primary + .reset_description + .as_deref(), + Some("$1,000.00 / $1,000.00") + ); + let credit = personal_budget(-5.0, 10.0); + assert_eq!(credit.usage.primary.used_percent, 0.0); + assert_eq!( + credit.usage.primary.reset_description.as_deref(), + Some("-$5.00 / $10.00") + ); + assert_eq!( + personal_budget(0.0, 0.5) + .usage + .primary + .reset_description + .as_deref(), + Some("$0.00 / $0.50") + ); +} + +#[test] +fn budget_reset_dates_accept_offsets_naive_times_and_dates() { + assert_eq!( + personal_reset("2026-10-01T05:30:00+02:00"), + Some(Utc.with_ymd_and_hms(2026, 10, 1, 3, 30, 0).unwrap()) + ); + assert_eq!( + personal_reset("2026-10-01T05:30:00"), + Some(Utc.with_ymd_and_hms(2026, 10, 1, 5, 30, 0).unwrap()) + ); + assert_eq!( + personal_reset(" 2026-10-01 "), + Some(Utc.with_ymd_and_hms(2026, 10, 1, 0, 0, 0).unwrap()) + ); + assert_eq!(personal_reset("next month"), None); + assert_eq!(personal_reset(""), None); +} + #[test] fn saved_base_url_uses_only_app_saved_key() { let mut ctx = FetchContext { From f376f5aef1205e536f50d0bd37f5c488874f11c9 Mon Sep 17 00:00:00 2001 From: RCD <90105158+Finesssee@users.noreply.github.com> Date: Thu, 1 Oct 2026 16:10:41 +0700 Subject: [PATCH 3/3] Show LiteLLM parse errors without a doubled prefix --- rust/src/providers/litellm/info.rs | 10 ++++++---- rust/src/providers/litellm/mod.rs | 8 +++----- rust/src/providers/litellm/tests.rs | 14 +++++++++++--- 3 files changed, 20 insertions(+), 12 deletions(-) diff --git a/rust/src/providers/litellm/info.rs b/rust/src/providers/litellm/info.rs index 10065361be..0eb8b49bc9 100644 --- a/rust/src/providers/litellm/info.rs +++ b/rust/src/providers/litellm/info.rs @@ -144,18 +144,20 @@ const TEAM_BUDGET: &str = "Team budget"; /// Informational lane text for a key without a personal or team budget. const NO_BUDGET: &str = "No budget set"; +/// Upstream fails with `LiteLLM parse error: `. `ProviderError::Parse` +/// already displays as "Parse error: …", so only the provider name is added. pub(super) fn parse_error(message: impl std::fmt::Display) -> ProviderError { - ProviderError::Parse(format!("LiteLLM parse error: {message}")) + ProviderError::Parse(format!("LiteLLM {message}")) } +pub(super) const MISSING_KEY_IDS: &str = "key info did not include a user_id or team_id."; + pub(super) fn bind_key(response: KeyInfoResponse) -> Result { let info = response.info; let user_id = nonempty(info.user_id); let team_id = nonempty(info.team_id); if user_id.is_none() && team_id.is_none() { - return Err(parse_error( - "LiteLLM key info did not include a user_id or team_id.", - )); + return Err(parse_error(MISSING_KEY_IDS)); } Ok(KeyBinding { user_id, diff --git a/rust/src/providers/litellm/mod.rs b/rust/src/providers/litellm/mod.rs index 5516596484..d075eeeea4 100644 --- a/rust/src/providers/litellm/mod.rs +++ b/rust/src/providers/litellm/mod.rs @@ -16,8 +16,8 @@ mod tests; use endpoint::management_url; pub(crate) use endpoint::validated_base_url; use info::{ - KeyInfoResponse, TeamInfoResponse, UserInfoResponse, bind_key, parse_error, result_from_team, - result_from_user, + KeyInfoResponse, MISSING_KEY_IDS, TeamInfoResponse, UserInfoResponse, bind_key, parse_error, + result_from_team, result_from_user, }; const CREDENTIAL_TARGET: &str = "codexbar-litellm"; @@ -125,9 +125,7 @@ impl Provider for LiteLLMProvider { let response: TeamInfoResponse = self.get_json(url, &key).await?; result_from_team(&binding, team_id, response) } else { - Err(parse_error( - "LiteLLM key info did not include a user_id or team_id.", - )) + Err(parse_error(MISSING_KEY_IDS)) } } SourceMode::Web | SourceMode::Cli => { diff --git a/rust/src/providers/litellm/tests.rs b/rust/src/providers/litellm/tests.rs index 111b68446a..9d39105a5d 100644 --- a/rust/src/providers/litellm/tests.rs +++ b/rust/src/providers/litellm/tests.rs @@ -74,9 +74,9 @@ fn key_info_without_user_or_team_id_fails() { let response: KeyInfoResponse = serde_json::from_value(json!({"info": {"user_id": " ", "spend": 1.0}})).unwrap(); match bind_key(response) { - Err(ProviderError::Parse(message)) => assert!( - message.contains("LiteLLM key info did not include a user_id or team_id."), - "unexpected message: {message}" + Err(error @ ProviderError::Parse(_)) => assert_eq!( + error.to_string(), + "Parse error: LiteLLM key info did not include a user_id or team_id." ), _ => panic!("expected parse error"), } @@ -335,6 +335,14 @@ fn mismatched_user_id_is_rejected() { user_result(&key, json!({"user_id": "user-2", "user_info": {}})), "user_id did not match /key/info", ); + // The card shows the error's Display text, so "parse error" appears once. + let Err(error) = user_result(&key, json!({"user_info": {"user_id": "user-2"}})) else { + panic!("expected parse error"); + }; + assert_eq!( + error.to_string(), + "Parse error: LiteLLM user_id did not match /key/info" + ); } #[test]