diff --git a/rust/src/providers/openaiapi/mod.rs b/rust/src/providers/openaiapi/mod.rs index 91e5abf697..bc92d1ebd8 100644 --- a/rust/src/providers/openaiapi/mod.rs +++ b/rust/src/providers/openaiapi/mod.rs @@ -1,13 +1,28 @@ //! OpenAI API usage provider. //! //! Tracks organization usage from the Admin API, with the older platform credit -//! balance endpoint as a fallback for project/user keys. +//! balance endpoint as a fallback for keys that are not project-scoped Admin keys. +//! +//! Upstream v0.66.0 `Resources/Plugins/openai.js` parity: +//! - `/v1/organization/costs` (`group_by=line_item`) and +//! `/v1/organization/usage/completions` (`group_by=model`) use `bucket_width=1d`, +//! `limit <= 31` and UTC-day-aligned 31-day ranges, following `has_more` / `next_page` +//! (at most 100 pages per range; a missing or repeated cursor is an error). +//! - Token totals are `input + input_audio + output + output_audio`; cached input is a +//! subset of input and is never added on top. +//! - Each Admin GET gets one transient retry (see [`RetryPolicy`]). +//! +//! The history window is fixed at 30 days. Upstream's `OPENAI_HISTORY_DAYS` (1-365) is +//! deferred until a Windows setting exists for it; [`usage_ranges`] already takes the day +//! count, so honoring it later only needs to pass the setting through. use async_trait::async_trait; -use chrono::{DateTime, Duration, TimeZone, Utc}; +use chrono::{DateTime, TimeZone, Utc}; use reqwest::Client; +use reqwest::header::{HeaderValue, RETRY_AFTER}; use serde::Deserialize; -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; +use std::time::Duration; use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, @@ -19,6 +34,16 @@ const OPENAI_ORG_COSTS_URL: &str = "https://api.openai.com/v1/organization/costs const OPENAI_ORG_COMPLETIONS_URL: &str = "https://api.openai.com/v1/organization/usage/completions"; const OPENAI_API_CREDENTIAL_TARGET: &str = "codexbar-openaiapi"; +/// Days of history requested from the Admin API (period label `Last 30 days`). +const HISTORY_DAYS: u32 = 30; +/// Endpoint bucket limit: each request covers at most this many daily buckets. +const MAX_BUCKETS_PER_REQUEST: u32 = 31; +const SECONDS_PER_DAY: i64 = 86_400; +/// Upstream `pages()` stops with a parse failure after this many pages per range. +const MAX_PAGES_PER_RANGE: usize = 100; +/// Upstream fetches with `timeoutSeconds: 20`. +const REQUEST_TIMEOUT: Duration = Duration::from_secs(20); + #[derive(Debug, Deserialize)] struct CreditGrantsResponse { total_granted: f64, @@ -37,9 +62,12 @@ struct CreditGrant { expires_at: Option, } +/// One page of an Admin API listing. `has_more` is required; a page without it is malformed. #[derive(Debug, Deserialize)] -struct CostsResponse { - data: Vec, +struct Page { + data: Vec, + has_more: bool, + next_page: Option, } #[derive(Debug, Deserialize)] @@ -64,11 +92,6 @@ struct CostAmount { value: serde_json::Value, } -#[derive(Debug, Deserialize)] -struct CompletionsUsageResponse { - data: Vec, -} - #[derive(Debug, Deserialize)] struct CompletionsUsageBucket { start_time: i64, @@ -84,6 +107,10 @@ struct CompletionsUsageBucket { struct CompletionsUsageResult { model: Option, input_tokens: Option, + #[allow( + dead_code, + reason = "cached input is a subset of input_tokens, so it is tracked but never added to totals" + )] input_cached_tokens: Option, output_tokens: Option, input_audio_tokens: Option, @@ -91,9 +118,95 @@ struct CompletionsUsageResult { num_model_requests: Option, } +impl CompletionsUsageResult { + /// Upstream `tokens = input + input_audio + output + output_audio`. + fn total_tokens(&self) -> i64 { + [ + self.input_tokens, + self.input_audio_tokens, + self.output_tokens, + self.output_audio_tokens, + ] + .into_iter() + .map(|tokens| tokens.unwrap_or(0)) + .sum() + } +} + +/// One `start_time..end_time` request window of at most [`MAX_BUCKETS_PER_REQUEST`] days. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +struct UsageRange { + start: i64, + end: i64, + limit: u32, +} + +/// Transient retry for the Admin GETs, mirroring upstream `transientIdempotent` +/// (one retry, base delay 1 s, `Retry-After` honored up to 10 s). +#[derive(Debug, Clone, Copy)] +struct RetryPolicy { + base_delay: Duration, + max_delay: Duration, +} + +impl RetryPolicy { + const DEFAULT: Self = Self { + base_delay: Duration::from_secs(1), + max_delay: Duration::from_secs(10), + }; + + /// Upstream `retryableStatusCodes`. + fn retries_status(status: reqwest::StatusCode) -> bool { + matches!(status.as_u16(), 408 | 429 | 500 | 502 | 503 | 504) + } + + /// Numeric `Retry-After` seconds capped at `max_delay`; anything else uses `base_delay`. + fn delay(&self, retry_after: Option<&HeaderValue>) -> Duration { + retry_after + .and_then(|value| value.to_str().ok()) + .and_then(|value| value.trim().parse::().ok()) + .filter(|seconds| seconds.is_finite() && *seconds >= 0.0) + .map_or(self.base_delay, |seconds| { + Duration::from_secs_f64(seconds.min(self.max_delay.as_secs_f64())) + }) + } +} + +#[derive(Debug, Clone)] +struct Endpoints { + credit_grants: String, + costs: String, + completions: String, +} + +impl Endpoints { + fn production() -> Self { + Self { + credit_grants: OPENAI_CREDIT_GRANTS_URL.to_string(), + costs: OPENAI_ORG_COSTS_URL.to_string(), + completions: OPENAI_ORG_COMPLETIONS_URL.to_string(), + } + } +} + +/// Resolved key plus whether it came from an Admin key source (Preferences, keychain, +/// `OPENAI_ADMIN_KEY`) rather than a plain `OPENAI_API_KEY`-style variable. +struct ApiCredential { + key: String, + is_admin: bool, +} + +/// Upstream `OpenAIAPIUsageCredential.allowsLegacyBalanceFallback`: the credit-grants +/// endpoint is not project-filtered, so a project-scoped Admin key never falls back to it. +fn allows_legacy_balance_fallback(project_id: Option<&str>, is_admin: bool) -> bool { + project_id.is_none() || !is_admin +} + pub struct OpenAIApiProvider { metadata: ProviderMetadata, client: Client, + endpoints: Endpoints, + retry: RetryPolicy, } impl OpenAIApiProvider { @@ -113,21 +226,23 @@ impl OpenAIApiProvider { tertiary_label_key: None, }, client: crate::core::credentialed_http_client_builder() - .timeout(std::time::Duration::from_secs(15)) + .timeout(REQUEST_TIMEOUT) .build() .unwrap_or_else(|_| Client::new()), + endpoints: Endpoints::production(), + retry: RetryPolicy::DEFAULT, } } - fn api_key(api_key: Option<&str>) -> Result { + fn credential(api_key: Option<&str>) -> Result { resolve_api_key( api_key, OPENAI_API_CREDENTIAL_TARGET, &[ - "OPENAI_ADMIN_KEY", - "OPENAI_ADMIN_API_KEY", - "OPENAI_API_KEY", - "OPENAI_PLATFORM_API_KEY", + ("OPENAI_ADMIN_KEY", true), + ("OPENAI_ADMIN_API_KEY", true), + ("OPENAI_API_KEY", false), + ("OPENAI_PLATFORM_API_KEY", false), ], ) } @@ -135,7 +250,7 @@ impl OpenAIApiProvider { async fn fetch_api(&self, api_key: &str) -> Result { let response = self .client - .get(OPENAI_CREDIT_GRANTS_URL) + .get(&self.endpoints.credit_grants) .bearer_auth(api_key) .header("Accept", "application/json") .send() @@ -164,49 +279,117 @@ impl OpenAIApiProvider { &self, api_key: &str, project_id: Option<&str>, + now: DateTime, ) -> Result { - let now = Utc::now(); - let (start, end) = admin_time_range(now)?; let project_id = clean_project_id(project_id); - - let costs: CostsResponse = self - .fetch_admin_json( - OPENAI_ORG_COSTS_URL, - &admin_query( - ("start_time", start.to_string()), - ("end_time", end.to_string()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by", "line_item".to_string()), - project_id.as_deref(), - ), + let ranges = usage_ranges(now, HISTORY_DAYS); + + let costs: Vec = self + .fetch_pages( + &self.endpoints.costs, + AdminQuery { + group_by: "line_item", + project_id: project_id.as_deref(), + label: "costs", + }, + &ranges, api_key, - "costs", ) .await?; - - let completions: CompletionsUsageResponse = self - .fetch_admin_json( - OPENAI_ORG_COMPLETIONS_URL, - &admin_query( - ("start_time", start.to_string()), - ("end_time", end.to_string()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by", "model".to_string()), - project_id.as_deref(), - ), + let completions: Vec = self + .fetch_pages( + &self.endpoints.completions, + AdminQuery { + group_by: "model", + project_id: project_id.as_deref(), + label: "completions", + }, + &ranges, api_key, - "completions", ) .await?; - Ok(result_from_admin_usage( - &costs, - &completions, - now, - project_id.as_deref(), - )) + result_from_admin_usage(&costs, &completions, now, project_id.as_deref()) + } + + /// Admin usage first; the legacy balance endpoint only when + /// [`allows_legacy_balance_fallback`] permits it. + async fn fetch_admin_or_balance( + &self, + credential: &ApiCredential, + project_id: Option<&str>, + now: DateTime, + ) -> Result { + let admin_error = match self + .fetch_admin_usage(&credential.key, project_id, now) + .await + { + Ok(result) => return Ok(result), + Err(admin_error) => admin_error, + }; + if !allows_legacy_balance_fallback(project_id, credential.is_admin) { + return Err(admin_error); + } + match self.fetch_api(&credential.key).await { + Ok(result) => Ok(ProviderFetchResult { + source_label: "billing-api".to_string(), + ..result + }), + Err(balance_error) => { + if matches!(admin_error, ProviderError::AuthRequired) { + Err(balance_error) + } else { + Err(admin_error) + } + } + } + } + + /// Fetch every page of every range, following `has_more` / `next_page`. + async fn fetch_pages( + &self, + url: &str, + query: AdminQuery<'_>, + ranges: &[UsageRange], + api_key: &str, + ) -> Result, ProviderError> { + let label = query.label; + let mut buckets = Vec::new(); + for range in ranges { + let mut page: Option = None; + let mut seen = HashSet::new(); + for count in 0..MAX_PAGES_PER_RANGE { + let params = admin_query(range, &query, page.as_deref()); + let decoded: Page = self.fetch_admin_json(url, ¶ms, api_key, label).await?; + buckets.extend(decoded.data); + if !decoded.has_more { + break; + } + let cursor = decoded + .next_page + .as_deref() + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .ok_or_else(|| { + ProviderError::Parse(format!( + "OpenAI API {label} pagination cursor missing" + )) + })? + .to_string(); + if !seen.insert(cursor.clone()) { + return Err(ProviderError::Parse(format!( + "OpenAI API {label} pagination cursor repeated" + ))); + } + if count + 1 == MAX_PAGES_PER_RANGE { + return Err(ProviderError::Parse(format!( + "OpenAI API {label} pagination exceeded {MAX_PAGES_PER_RANGE} pages" + ))); + } + page = Some(cursor); + } + } + Ok(buckets) } async fn fetch_admin_json( @@ -216,14 +399,7 @@ impl OpenAIApiProvider { api_key: &str, label: &str, ) -> Result { - let response = self - .client - .get(url) - .query(query) - .bearer_auth(api_key) - .header("Accept", "application/json") - .send() - .await?; + let response = self.get_with_retry(url, query, api_key).await?; if response.status() == reqwest::StatusCode::UNAUTHORIZED || response.status() == reqwest::StatusCode::FORBIDDEN { @@ -244,6 +420,40 @@ impl OpenAIApiProvider { .await .map_err(|e| ProviderError::Parse(format!("Failed to parse OpenAI API {label}: {e}"))) } + + /// GET with at most one retry on a transient failure: 408/429/5xx, a timeout, or a + /// refused connection. Auth failures, TLS errors and other transport errors are + /// returned at once, and dropping the future (cancellation) never triggers a retry. + async fn get_with_retry( + &self, + url: &str, + query: &[(&str, String)], + api_key: &str, + ) -> Result { + let send = || { + self.client + .get(url) + .query(query) + .bearer_auth(api_key) + .header("Accept", "application/json") + .send() + }; + match send().await { + Ok(response) if RetryPolicy::retries_status(response.status()) => { + let delay = self.retry.delay(response.headers().get(RETRY_AFTER)); + tokio::time::sleep(delay).await; + } + Ok(response) => return Ok(response), + Err(error) => { + let error = ProviderError::Network(error); + if !error.is_transport_failure() { + return Err(error); + } + tokio::time::sleep(self.retry.delay(None)).await; + } + } + Ok(send().await?) + } } fn result_from_grants(grants: &CreditGrantsResponse) -> ProviderFetchResult { @@ -288,85 +498,52 @@ fn result_from_grants(grants: &CreditGrantsResponse) -> ProviderFetchResult { } fn result_from_admin_usage( - costs: &CostsResponse, - completions: &CompletionsUsageResponse, + costs: &[CostBucket], + completions: &[CompletionsUsageBucket], now: DateTime, project_id: Option<&str>, -) -> ProviderFetchResult { - let cost_total: f64 = costs - .data - .iter() - .flat_map(|bucket| &bucket.results) - .map(|result| { - result - .amount - .as_ref() - .and_then(|a| number_value(&a.value)) - .unwrap_or(0.0) - }) - .sum(); - let request_total: i64 = completions - .data - .iter() - .flat_map(|bucket| &bucket.results) - .map(|result| result.num_model_requests.unwrap_or(0)) - .sum(); - let token_total: i64 = completions - .data - .iter() - .flat_map(|bucket| &bucket.results) - .map(|result| { - result.input_tokens.unwrap_or(0) - + result.input_cached_tokens.unwrap_or(0) - + result.output_tokens.unwrap_or(0) - + result.input_audio_tokens.unwrap_or(0) - + result.output_audio_tokens.unwrap_or(0) - }) - .sum(); - - let first_bucket = costs - .data - .first() - .map(|b| b.start_time) - .or_else(|| completions.data.first().map(|b| b.start_time)); - let start = first_bucket.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); +) -> Result { + let mut cost_total = 0.0; + let mut line_item_costs: HashMap = HashMap::new(); + for result in costs.iter().flat_map(|bucket| &bucket.results) { + let amount = cost_amount(result)?; + cost_total += amount; + let line_item = result + .line_item + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + .unwrap_or("API"); + *line_item_costs.entry(line_item.to_string()).or_default() += amount; + } + let mut request_total: i64 = 0; + let mut token_total: i64 = 0; let mut model_tokens: HashMap = HashMap::new(); - for result in completions.data.iter().flat_map(|bucket| &bucket.results) { + for result in completions.iter().flat_map(|bucket| &bucket.results) { + let tokens = result.total_tokens(); + request_total += result.num_model_requests.unwrap_or(0); + token_total += tokens; let model = result .model .as_deref() .map(str::trim) .filter(|s| !s.is_empty()) .unwrap_or("Responses and Chat Completions"); - let tokens = result.input_tokens.unwrap_or(0) - + result.input_cached_tokens.unwrap_or(0) - + result.output_tokens.unwrap_or(0) - + result.input_audio_tokens.unwrap_or(0) - + result.output_audio_tokens.unwrap_or(0); *model_tokens.entry(model.to_string()).or_default() += tokens; } - let mut line_item_costs: HashMap = HashMap::new(); - for result in costs.data.iter().flat_map(|bucket| &bucket.results) { - let line_item = result - .line_item - .as_deref() - .map(str::trim) - .filter(|s| !s.is_empty()) - .unwrap_or("API"); - *line_item_costs.entry(line_item.to_string()).or_default() += result - .amount - .as_ref() - .and_then(|a| number_value(&a.value)) - .unwrap_or(0.0); - } + let first_bucket = costs + .first() + .map(|b| b.start_time) + .or_else(|| completions.first().map(|b| b.start_time)); + let start = first_bucket.and_then(|ts| Utc.timestamp_opt(ts, 0).single()); let mut usage = UsageSnapshot::new(RateWindow::with_details( 0.0, None, start, - Some(format!("${cost_total:.2} over last 30 days")), + Some(format!("${cost_total:.2} over last {HISTORY_DAYS} days")), )) .with_extra_rate_window( "requests", @@ -409,36 +586,82 @@ fn result_from_admin_usage( ); } - ProviderFetchResult::new(usage, "admin-api").with_cost(CostSnapshot::new( - cost_total, - "USD", - "Last 30 days", - )) + Ok( + ProviderFetchResult::new(usage, "admin-api").with_cost(CostSnapshot::new( + cost_total, + "USD", + format!("Last {HISTORY_DAYS} days"), + )), + ) +} + +/// A missing, null or blank amount counts as zero; a present but non-numeric or +/// non-finite one (`NaN`, `Infinity`, `1e309`) is a parse failure, as upstream. +fn cost_amount(result: &CostResult) -> Result { + let Some(amount) = &result.amount else { + return Ok(0.0); + }; + match &amount.value { + serde_json::Value::Null => Ok(0.0), + serde_json::Value::String(text) if text.trim().is_empty() => Ok(0.0), + value => number_value(value).ok_or_else(|| { + ProviderError::Parse("OpenAI API costs amount must be numeric".to_string()) + }), + } +} + +/// Query parameters shared by every page of one Admin API listing. +struct AdminQuery<'a> { + group_by: &'static str, + project_id: Option<&'a str>, + label: &'static str, } fn admin_query( - start: (&'static str, String), - end: (&'static str, String), - bucket_width: (&'static str, String), - limit: (&'static str, String), - group_by: (&'static str, String), - project_id: Option<&str>, + range: &UsageRange, + query: &AdminQuery<'_>, + page: Option<&str>, ) -> Vec<(&'static str, String)> { - let mut query = vec![start, end, bucket_width, limit, group_by]; - if let Some(project_id) = clean_project_id(project_id) { - query.push(("project_ids", project_id)); + let mut params = vec![ + ("start_time", range.start.to_string()), + ("end_time", range.end.to_string()), + ("bucket_width", "1d".to_string()), + ("limit", range.limit.to_string()), + ("group_by", query.group_by.to_string()), + ]; + if let Some(project_id) = clean_project_id(query.project_id) { + params.push(("project_ids", project_id)); } - query + if let Some(page) = page { + params.push(("page", page.to_string())); + } + params } -fn admin_time_range(now: DateTime) -> Result<(i64, i64), ProviderError> { - let start = (now.date_naive() - Duration::days(29)) - .and_hms_opt(0, 0, 0) - .ok_or_else(|| ProviderError::Parse("Invalid OpenAI usage start date".to_string()))? - .and_utc() - .timestamp(); - Ok((start, now.timestamp())) +/// Upstream `ranges()`: UTC-day-aligned windows of at most [`MAX_BUCKETS_PER_REQUEST`] days +/// covering the last `history_days` days including today. Upstream ends the final window +/// at tomorrow's midnight; the API rejects a future `end_time` +/// (`end_time must not be in the future`), so the end is clamped to `now`. +fn usage_ranges(now: DateTime, history_days: u32) -> Vec { + let now = now.timestamp(); + let today = now - now.rem_euclid(SECONDS_PER_DAY); + let mut start = today - (i64::from(history_days) - 1) * SECONDS_PER_DAY; + let mut remaining = history_days; + let mut ranges = Vec::new(); + while remaining > 0 { + let limit = remaining.min(MAX_BUCKETS_PER_REQUEST); + let end = start + i64::from(limit) * SECONDS_PER_DAY; + ranges.push(UsageRange { + start, + end: end.min(now), + limit, + }); + start = end; + remaining -= limit; + } + ranges } + fn clean_project_id(project_id: Option<&str>) -> Option { project_id .map(str::trim) @@ -493,26 +716,10 @@ impl Provider for OpenAIApiProvider { async fn fetch_usage(&self, ctx: &FetchContext) -> Result { match ctx.source_mode { SourceMode::Auto | SourceMode::OAuth => { - let api_key = Self::api_key(ctx.api_key.as_deref())?; - match self - .fetch_admin_usage(&api_key, ctx.workspace_id.as_deref()) + let credential = Self::credential(ctx.api_key.as_deref())?; + let project_id = clean_project_id(ctx.workspace_id.as_deref()); + self.fetch_admin_or_balance(&credential, project_id.as_deref(), Utc::now()) .await - { - Ok(result) => Ok(result), - Err(admin_error) => match self.fetch_api(&api_key).await { - Ok(result) => Ok(ProviderFetchResult { - source_label: "billing-api".to_string(), - ..result - }), - Err(balance_error) => { - if matches!(admin_error, ProviderError::AuthRequired) { - Err(balance_error) - } else { - Err(admin_error) - } - } - }, - } } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -525,32 +732,44 @@ impl Provider for OpenAIApiProvider { } } +/// Resolve the key from Preferences, the credential store, then the environment. +/// Each `env_names` entry carries whether that variable holds an Admin key. fn resolve_api_key( explicit: Option<&str>, credential_target: &str, - env_names: &[&str], -) -> Result { + env_names: &[(&str, bool)], +) -> Result { if let Some(key) = explicit && !key.trim().is_empty() { - return Ok(key.trim().to_string()); + return Ok(ApiCredential { + key: key.trim().to_string(), + is_admin: true, + }); } if let Ok(entry) = keyring::Entry::new(credential_target, "api_key") && let Ok(key) = entry.get_password() && !key.trim().is_empty() { - return Ok(key); + return Ok(ApiCredential { + key, + is_admin: true, + }); } - for env in env_names { + for (env, is_admin) in env_names { if let Ok(key) = std::env::var(env) && !key.trim().is_empty() { - return Ok(key); + return Ok(ApiCredential { + key, + is_admin: *is_admin, + }); } } + let names: Vec<&str> = env_names.iter().map(|(name, _)| *name).collect(); Err(ProviderError::NotInstalled(format!( "API key not found. Set {} in Preferences or environment.", - env_names.join(" / ") + names.join(" / ") ))) } @@ -561,120 +780,4 @@ fn resolve_api_key( fn _assert_datetime_send(_: DateTime) {} #[cfg(test)] -mod tests { - use super::*; - - #[test] - fn openai_api_credit_snapshot_formats_available_balance() { - let result = result_from_grants(&CreditGrantsResponse { - total_granted: 100.0, - total_used: 25.0, - total_available: 75.0, - grants: None, - }); - assert_eq!(result.usage.primary.used_percent, 25.0); - assert_eq!(result.cost.unwrap().remaining(), Some(75.0)); - } - - #[test] - fn openai_admin_usage_accepts_numeric_string_cost_amounts() { - let costs = CostsResponse { - data: vec![CostBucket { - start_time: 1_797_638_400, - end_time: 1_797_724_800, - results: vec![CostResult { - amount: Some(CostAmount { - value: serde_json::json!("12.50"), - }), - line_item: Some("API".to_string()), - }], - }], - }; - let completions = CompletionsUsageResponse { - data: vec![CompletionsUsageBucket { - start_time: 1_797_638_400, - end_time: 1_797_724_800, - results: vec![CompletionsUsageResult { - model: Some("gpt-5.1".to_string()), - input_tokens: Some(100), - input_cached_tokens: Some(25), - output_tokens: Some(50), - input_audio_tokens: None, - output_audio_tokens: None, - num_model_requests: Some(7), - }], - }], - }; - let now = Utc.timestamp_opt(1_797_724_800, 0).single().unwrap(); - let result = result_from_admin_usage(&costs, &completions, now, Some("proj_demo")); - assert_eq!(result.cost.unwrap().used, 12.5); - assert_eq!( - result.usage.account_organization.as_deref(), - Some("Project: proj_demo") - ); - assert_eq!( - result.usage.login_method.as_deref(), - Some("Admin API: proj_demo") - ); - assert!(result.usage.secondary.is_none()); - let requests = result - .usage - .extra_rate_windows - .iter() - .find(|window| window.id == "requests") - .unwrap(); - assert!(requests.window.is_informational); - assert_eq!( - requests.window.reset_description.as_deref(), - Some("7 requests") - ); - assert!( - result.usage.extra_rate_windows.iter().any(|window| window - .window - .reset_description - .as_deref() - == Some("175 tokens")) - ); - } - - #[test] - fn openai_admin_usage_rejects_nonfinite_cost_amounts() { - assert_eq!(number_value(&serde_json::json!("NaN")), None); - assert_eq!(number_value(&serde_json::json!("Infinity")), None); - } - - #[test] - fn openai_admin_query_scopes_project_ids_when_configured() { - let query = admin_query( - ("start_time", "1".to_string()), - ("end_time", "2".to_string()), - ("bucket_width", "1d".to_string()), - ("limit", "31".to_string()), - ("group_by", "model".to_string()), - Some(" proj_123 "), - ); - assert!(query.contains(&("project_ids", "proj_123".to_string()))); - } - - #[test] - fn openai_admin_time_range_does_not_end_in_the_future() { - let now = Utc.timestamp_opt(1_783_981_862, 0).single().unwrap(); - let (start, end) = admin_time_range(now).unwrap(); - - assert_eq!(end, now.timestamp()); - assert!(start < end); - } - - #[test] - fn openai_response_error_detail_prefers_api_message() { - assert_eq!( - response_error_detail(r#"{"error":{"message":"end_time must not be in the future"}}"#), - "end_time must not be in the future" - ); - assert_eq!( - response_error_detail("plain text failure"), - "plain text failure" - ); - assert_eq!(response_error_detail(" "), ""); - } -} +mod tests; diff --git a/rust/src/providers/openaiapi/tests.rs b/rust/src/providers/openaiapi/tests.rs new file mode 100644 index 0000000000..37493fef4b --- /dev/null +++ b/rust/src/providers/openaiapi/tests.rs @@ -0,0 +1,795 @@ +use super::*; +use mockito::{Matcher, Server, ServerGuard}; + +const COSTS_PATH: &str = "/v1/organization/costs"; +const COMPLETIONS_PATH: &str = "/v1/organization/usage/completions"; +const GRANTS_PATH: &str = "/v1/dashboard/billing/credit_grants"; +const EMPTY_PAGE: &str = r#"{"object":"page","data":[],"has_more":false,"next_page":null}"#; +/// 2023-11-17T00:00:00Z, the fixture clock used by upstream `OpenAIAPIUsageFetcherTests`. +const NOW: i64 = 1_700_179_200; + +fn fixed_now(offset_seconds: i64) -> DateTime { + Utc.timestamp_opt(NOW + offset_seconds, 0).single().unwrap() +} + +/// A provider pointed at the mock server, with retry delays removed so tests stay fast. +fn provider(server: &ServerGuard) -> OpenAIApiProvider { + let mut provider = OpenAIApiProvider::new(); + provider.client = Client::builder() + .no_proxy() + .build() + .expect("the test client should build"); + provider.endpoints = Endpoints { + credit_grants: format!("{}{GRANTS_PATH}", server.url()), + costs: format!("{}{COSTS_PATH}", server.url()), + completions: format!("{}{COMPLETIONS_PATH}", server.url()), + }; + provider.retry = RetryPolicy { + base_delay: Duration::ZERO, + max_delay: Duration::ZERO, + }; + provider +} + +fn admin_credential() -> ApiCredential { + ApiCredential { + key: "sk-test".to_string(), + is_admin: true, + } +} + +async fn mock_page( + server: &mut ServerGuard, + path: &str, + query: Matcher, + body: &str, +) -> mockito::Mock { + server + .mock("GET", path) + .match_query(query) + .with_status(200) + .with_body(body) + .create_async() + .await +} + +fn extra_window<'a>(result: &'a ProviderFetchResult, id: &str) -> &'a str { + result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == id) + .and_then(|window| window.window.reset_description.as_deref()) + .unwrap_or_else(|| panic!("missing extra window {id}")) +} + +fn cost_bucket(amount: serde_json::Value) -> CostBucket { + CostBucket { + start_time: 1_797_638_400, + end_time: 1_797_724_800, + results: vec![CostResult { + amount: Some(CostAmount { value: amount }), + line_item: Some("API".to_string()), + }], + } +} + +#[test] +fn openai_api_credit_snapshot_formats_available_balance() { + let result = result_from_grants(&CreditGrantsResponse { + total_granted: 100.0, + total_used: 25.0, + total_available: 75.0, + grants: None, + }); + assert_eq!(result.usage.primary.used_percent, 25.0); + assert_eq!(result.cost.unwrap().remaining(), Some(75.0)); +} + +#[test] +fn openai_admin_usage_accepts_numeric_string_cost_amounts() { + let costs = [cost_bucket(serde_json::json!("12.50"))]; + let completions = [CompletionsUsageBucket { + start_time: 1_797_638_400, + end_time: 1_797_724_800, + results: vec![CompletionsUsageResult { + model: Some("gpt-5.1".to_string()), + input_tokens: Some(100), + input_cached_tokens: Some(25), + output_tokens: Some(50), + input_audio_tokens: None, + output_audio_tokens: None, + num_model_requests: Some(7), + }], + }]; + let now = Utc.timestamp_opt(1_797_724_800, 0).single().unwrap(); + let result = result_from_admin_usage(&costs, &completions, now, Some("proj_demo")).unwrap(); + assert_eq!(result.cost.as_ref().unwrap().used, 12.5); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Project: proj_demo") + ); + assert_eq!( + result.usage.login_method.as_deref(), + Some("Admin API: proj_demo") + ); + assert!(result.usage.secondary.is_none()); + let requests = result + .usage + .extra_rate_windows + .iter() + .find(|window| window.id == "requests") + .unwrap(); + assert!(requests.window.is_informational); + assert_eq!( + requests.window.reset_description.as_deref(), + Some("7 requests") + ); + // Cached input is a subset of input: 100 + 50, not 100 + 25 + 50. + assert_eq!(extra_window(&result, "tokens"), "150 tokens"); + assert_eq!(extra_window(&result, "model-0"), "150 tokens"); +} + +#[test] +fn openai_admin_tokens_sum_text_and_audio_without_cached() { + // Upstream `tokens = input + input_audio + output + output_audio`. + let completions = [CompletionsUsageBucket { + start_time: 1_700_000_000, + end_time: 1_700_086_400, + results: vec![ + CompletionsUsageResult { + model: Some("gpt-5.2".to_string()), + input_tokens: Some(1000), + input_cached_tokens: Some(250), + output_tokens: Some(500), + input_audio_tokens: Some(40), + output_audio_tokens: Some(10), + num_model_requests: Some(7), + }, + CompletionsUsageResult { + model: Some("gpt-5.2-codex".to_string()), + input_tokens: Some(300), + input_cached_tokens: None, + output_tokens: Some(200), + input_audio_tokens: None, + output_audio_tokens: None, + num_model_requests: Some(3), + }, + ], + }]; + let result = result_from_admin_usage(&[], &completions, fixed_now(0), None).unwrap(); + assert_eq!(extra_window(&result, "tokens"), "2050 tokens"); + assert_eq!(extra_window(&result, "requests"), "10 requests"); + assert_eq!(extra_window(&result, "model-0"), "1550 tokens"); + assert_eq!(extra_window(&result, "model-1"), "500 tokens"); +} + +#[test] +fn openai_admin_usage_sums_costs_and_ranks_line_items() { + // Cases from upstream `parses admin costs and completions usage into daily summaries`. + let costs = [ + CostBucket { + start_time: 1_700_000_000, + end_time: 1_700_086_400, + results: vec![ + CostResult { + amount: Some(CostAmount { + value: serde_json::json!(12.50), + }), + line_item: Some("Text tokens".to_string()), + }, + CostResult { + amount: Some(CostAmount { + value: serde_json::json!("2.25"), + }), + line_item: Some("Web search tool calls".to_string()), + }, + ], + }, + CostBucket { + start_time: 1_700_086_400, + end_time: 1_700_172_800, + results: vec![CostResult { + amount: Some(CostAmount { + value: serde_json::json!(4.00), + }), + line_item: Some("Text tokens".to_string()), + }], + }, + ]; + let result = result_from_admin_usage(&costs, &[], fixed_now(0), None).unwrap(); + assert_eq!(result.cost.as_ref().unwrap().used, 18.75); + assert_eq!(result.cost.as_ref().unwrap().period, "Last 30 days"); + assert_eq!(extra_window(&result, "line-item-0"), "$16.50"); + assert_eq!(extra_window(&result, "line-item-1"), "$2.25"); +} + +#[test] +fn openai_admin_usage_rejects_nonfinite_cost_amounts() { + assert_eq!(number_value(&serde_json::json!("NaN")), None); + assert_eq!(number_value(&serde_json::json!("Infinity")), None); + for value in ["NaN", "Infinity", "-Infinity", "1e309", "-1e309"] { + let costs = [cost_bucket(serde_json::json!(value))]; + let error = result_from_admin_usage(&costs, &[], fixed_now(0), None) + .err() + .unwrap_or_else(|| panic!("{value} should be rejected")); + assert!(matches!(error, ProviderError::Parse(_)), "{value}: {error}"); + } +} + +#[test] +fn openai_admin_usage_treats_blank_amounts_as_zero() { + let costs = [ + cost_bucket(serde_json::Value::Null), + cost_bucket(serde_json::json!(" ")), + CostBucket { + start_time: 1, + end_time: 2, + results: vec![CostResult { + amount: None, + line_item: None, + }], + }, + ]; + let result = result_from_admin_usage(&costs, &[], fixed_now(0), None).unwrap(); + assert_eq!(result.cost.unwrap().used, 0.0); +} + +#[test] +fn openai_admin_query_scopes_project_ids_and_page_when_configured() { + let range = UsageRange { + start: 1, + end: 2, + limit: 31, + }; + let query = AdminQuery { + group_by: "model", + project_id: Some(" proj_123 "), + label: "completions", + }; + let params = admin_query(&range, &query, Some("cursor_2")); + assert!(params.contains(&("project_ids", "proj_123".to_string()))); + assert!(params.contains(&("page", "cursor_2".to_string()))); + assert!(params.contains(&("bucket_width", "1d".to_string()))); + assert!(params.contains(&("group_by", "model".to_string()))); + + let unscoped = admin_query( + &range, + &AdminQuery { + project_id: None, + ..query + }, + None, + ); + assert!(unscoped.iter().all(|(name, _)| *name != "project_ids")); + assert!(unscoped.iter().all(|(name, _)| *name != "page")); +} + +#[test] +fn openai_admin_time_range_does_not_end_in_the_future() { + let now = Utc.timestamp_opt(1_783_981_862, 0).single().unwrap(); + let ranges = usage_ranges(now, HISTORY_DAYS); + + assert_eq!(ranges.len(), 1); + assert_eq!(ranges[0].end, now.timestamp()); + assert_eq!( + ranges[0].start, + 1_783_981_862 - 1_783_981_862 % 86_400 - 29 * 86_400 + ); + assert_eq!(ranges[0].limit, 30); + assert!(ranges[0].start < ranges[0].end); +} + +#[test] +fn openai_usage_ranges_page_long_history_within_endpoint_bucket_limit() { + // Upstream `admin usage fetch pages long history within endpoint bucket limit`: 90 days + // is three ranges of 31, 31 and 28 buckets, contiguous and aligned to UTC day starts. + let ranges = usage_ranges(fixed_now(0), 90); + let limits: Vec = ranges.iter().map(|range| range.limit).collect(); + assert_eq!(limits, [31, 31, 28]); + assert!(limits.iter().all(|limit| *limit <= 31)); + assert_eq!(ranges[0].start, NOW - 89 * 86_400); + assert!(ranges.iter().all(|range| range.start % 86_400 == 0)); + assert_eq!(ranges[1].start, ranges[0].start + 31 * 86_400); + assert_eq!(ranges[2].start, ranges[1].start + 31 * 86_400); + assert_eq!(ranges[0].end, ranges[1].start); + assert_eq!(ranges[2].end, NOW); + + assert_eq!(usage_ranges(fixed_now(3_600), 1).len(), 1); + assert_eq!(usage_ranges(fixed_now(3_600), 1)[0].start, NOW); + assert_eq!(usage_ranges(fixed_now(3_600), 31).len(), 1); + assert_eq!(usage_ranges(fixed_now(3_600), 32).len(), 2); +} + +#[test] +fn openai_response_error_detail_prefers_api_message() { + assert_eq!( + response_error_detail(r#"{"error":{"message":"end_time must not be in the future"}}"#), + "end_time must not be in the future" + ); + assert_eq!( + response_error_detail("plain text failure"), + "plain text failure" + ); + assert_eq!(response_error_detail(" "), ""); +} + +#[test] +fn openai_retry_policy_honors_retry_after_up_to_ten_seconds() { + let policy = RetryPolicy::DEFAULT; + let header = |value: &'static str| HeaderValue::from_static(value); + + assert_eq!(policy.delay(None), Duration::from_secs(1)); + assert_eq!(policy.delay(Some(&header("3"))), Duration::from_secs(3)); + assert_eq!(policy.delay(Some(&header(" 2 "))), Duration::from_secs(2)); + assert_eq!(policy.delay(Some(&header("0"))), Duration::ZERO); + assert_eq!(policy.delay(Some(&header("600"))), Duration::from_secs(10)); + for invalid in ["soon", "-5", "inf", "NaN", ""] { + assert_eq!( + policy.delay(Some(&header(invalid))), + Duration::from_secs(1), + "{invalid:?}" + ); + } +} + +#[test] +fn openai_retry_policy_retries_only_transient_statuses() { + for status in [408, 429, 500, 502, 503, 504] { + assert!(RetryPolicy::retries_status( + reqwest::StatusCode::from_u16(status).unwrap() + )); + } + for status in [200, 400, 401, 403, 404, 422] { + assert!(!RetryPolicy::retries_status( + reqwest::StatusCode::from_u16(status).unwrap() + )); + } +} + +#[test] +fn openai_balance_fallback_is_blocked_only_for_project_scoped_admin_keys() { + assert!(allows_legacy_balance_fallback(None, true)); + assert!(allows_legacy_balance_fallback(None, false)); + assert!(allows_legacy_balance_fallback(Some("proj_abc"), false)); + assert!(!allows_legacy_balance_fallback(Some("proj_abc"), true)); +} + +#[test] +fn openai_credential_kind_follows_the_key_source() { + let credential = resolve_api_key( + Some(" sk-admin-configured "), + "codexbar-openaiapi-test-unused", + &[("CODEXBAR_TEST_OPENAI_UNSET", false)], + ) + .unwrap(); + assert_eq!(credential.key, "sk-admin-configured"); + assert!(credential.is_admin); +} + +#[tokio::test] +async fn openai_admin_usage_filters_costs_and_completions_by_project() { + let mut server = Server::new_async().await; + let query = |group_by: &str| { + Matcher::AllOf(vec![ + Matcher::UrlEncoded("project_ids".into(), "proj_abc".into()), + Matcher::UrlEncoded("group_by".into(), group_by.into()), + Matcher::UrlEncoded("bucket_width".into(), "1d".into()), + Matcher::UrlEncoded("limit".into(), "30".into()), + ]) + }; + let costs = mock_page(&mut server, COSTS_PATH, query("line_item"), EMPTY_PAGE) + .await + .expect(1); + let completions = mock_page(&mut server, COMPLETIONS_PATH, query("model"), EMPTY_PAGE) + .await + .expect(1); + + let result = provider(&server) + .fetch_admin_usage("sk-test", Some(" proj_abc "), fixed_now(3_600)) + .await + .unwrap(); + + costs.assert_async().await; + completions.assert_async().await; + assert_eq!(result.source_label, "admin-api"); + assert_eq!( + result.usage.account_organization.as_deref(), + Some("Project: proj_abc") + ); +} + +#[tokio::test] +async fn openai_admin_usage_pages_each_range_of_a_long_history() { + let mut server = Server::new_async().await; + let costs = mock_page(&mut server, COSTS_PATH, Matcher::Any, EMPTY_PAGE) + .await + .expect(3); + let completions = mock_page(&mut server, COMPLETIONS_PATH, Matcher::Any, EMPTY_PAGE) + .await + .expect(3); + + let provider = provider(&server); + let ranges = usage_ranges(fixed_now(0), 90); + let buckets: Vec = provider + .fetch_pages( + &provider.endpoints.costs, + AdminQuery { + group_by: "line_item", + project_id: None, + label: "costs", + }, + &ranges, + "sk-test", + ) + .await + .unwrap(); + assert!(buckets.is_empty()); + let _: Vec = provider + .fetch_pages( + &provider.endpoints.completions, + AdminQuery { + group_by: "model", + project_id: None, + label: "completions", + }, + &ranges, + "sk-test", + ) + .await + .unwrap(); + + costs.assert_async().await; + completions.assert_async().await; +} + +const COSTS_PAGE_1: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.costs.result","amount":{"value":1.25,"currency":"usd"},"line_item":"Text tokens"}]}],"has_more":true,"next_page":"costs_page_2"}"#; +const COSTS_PAGE_2: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.costs.result","amount":{"value":2.75,"currency":"usd"},"line_item":"Web search tool calls"}]}],"has_more":false,"next_page":null}"#; +const COMPLETIONS_PAGE_1: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.usage.completions.result","input_tokens":10,"output_tokens":5,"num_model_requests":1,"model":"gpt-5.2"}]}],"has_more":true,"next_page":"completions_page_2"}"#; +const COMPLETIONS_PAGE_2: &str = r#"{"object":"page","data":[{"object":"bucket","start_time":1700000000,"end_time":1700086400,"results":[{"object":"organization.usage.completions.result","input_tokens":20,"output_tokens":10,"num_model_requests":2,"model":"gpt-5.2"}]}],"has_more":false,"next_page":null}"#; + +#[tokio::test] +async fn openai_admin_usage_follows_costs_and_completions_pagination_cursors() { + let mut server = Server::new_async().await; + let page = |name: &str| Matcher::UrlEncoded("page".into(), name.into()); + let mocks = vec![ + server + .mock("GET", COSTS_PATH) + .match_query(Matcher::AllOf(vec![Matcher::Regex( + "project_ids=proj_abc$".into(), + )])) + .with_body(COSTS_PAGE_1) + .expect(1) + .create_async() + .await, + server + .mock("GET", COSTS_PATH) + .match_query(Matcher::AllOf(vec![ + page("costs_page_2"), + Matcher::UrlEncoded("project_ids".into(), "proj_abc".into()), + ])) + .with_body(COSTS_PAGE_2) + .expect(1) + .create_async() + .await, + server + .mock("GET", COMPLETIONS_PATH) + .match_query(Matcher::AllOf(vec![Matcher::Regex( + "project_ids=proj_abc$".into(), + )])) + .with_body(COMPLETIONS_PAGE_1) + .expect(1) + .create_async() + .await, + server + .mock("GET", COMPLETIONS_PATH) + .match_query(Matcher::AllOf(vec![ + page("completions_page_2"), + Matcher::UrlEncoded("project_ids".into(), "proj_abc".into()), + ])) + .with_body(COMPLETIONS_PAGE_2) + .expect(1) + .create_async() + .await, + ]; + + let result = provider(&server) + .fetch_admin_usage("sk-test", Some("proj_abc"), fixed_now(3_600)) + .await + .unwrap(); + + for mock in &mocks { + mock.assert_async().await; + } + assert_eq!(result.cost.as_ref().unwrap().used, 4.0); + assert_eq!(extra_window(&result, "requests"), "3 requests"); + assert_eq!(extra_window(&result, "tokens"), "45 tokens"); +} + +async fn fetch_with_page_body(body: &str) -> ProviderError { + let mut server = Server::new_async().await; + let _costs = mock_page(&mut server, COSTS_PATH, Matcher::Any, body).await; + let _completions = mock_page(&mut server, COMPLETIONS_PATH, Matcher::Any, EMPTY_PAGE).await; + provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("the page should be rejected") +} + +#[tokio::test] +async fn openai_admin_usage_rejects_repeated_pagination_cursor() { + let error = fetch_with_page_body( + r#"{"object":"page","data":[],"has_more":true,"next_page":"same_page"}"#, + ) + .await; + assert!( + matches!(&error, ProviderError::Parse(message) if message.contains("cursor repeated")), + "{error}" + ); +} + +#[tokio::test] +async fn openai_admin_usage_rejects_missing_pagination_cursor() { + for body in [ + r#"{"object":"page","data":[],"has_more":true,"next_page":null}"#, + r#"{"object":"page","data":[],"has_more":true,"next_page":" "}"#, + r#"{"object":"page","data":[],"has_more":true}"#, + ] { + let error = fetch_with_page_body(body).await; + assert!( + matches!(&error, ProviderError::Parse(message) if message.contains("cursor missing")), + "{body}: {error}" + ); + } +} + +#[tokio::test] +async fn openai_admin_usage_rejects_malformed_pages() { + for body in [ + r#"{"object":"page","has_more":false,"next_page":null}"#, + r#"{"object":"page","data":[],"next_page":null}"#, + r#"{"object":"page","data":[],"has_more":false,"next_page":7}"#, + ] { + let error = fetch_with_page_body(body).await; + assert!(matches!(error, ProviderError::Parse(_)), "{body}: {error}"); + } +} + +#[tokio::test] +async fn openai_admin_usage_rejects_pagination_beyond_one_hundred_pages() { + let mut server = Server::new_async().await; + let costs = server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_body_from_request(|request| { + let path = request.path_and_query(); + let current = path + .split("page=p") + .nth(1) + .and_then(|rest| rest.split('&').next()) + .and_then(|number| number.parse::().ok()) + .unwrap_or(0); + format!( + r#"{{"data":[],"has_more":true,"next_page":"p{}"}}"#, + current + 1 + ) + .into_bytes() + }) + .expect(MAX_PAGES_PER_RANGE) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("the endless cursor chain should be rejected"); + + costs.assert_async().await; + assert!( + matches!(&error, ProviderError::Parse(message) if message.contains("exceeded 100 pages")), + "{error}" + ); +} + +#[tokio::test] +async fn openai_admin_usage_retries_transient_completions_failure_once() { + let mut server = Server::new_async().await; + let costs = mock_page(&mut server, COSTS_PATH, Matcher::Any, EMPTY_PAGE) + .await + .expect(1); + let failure = server + .mock("GET", COMPLETIONS_PATH) + .match_query(Matcher::Any) + .with_status(503) + .expect(1) + .create_async() + .await; + let success = mock_page( + &mut server, + COMPLETIONS_PATH, + Matcher::Any, + r#"{"object":"page","data":[{"start_time":1700000000,"end_time":1700086400,"results":[{"input_tokens":10,"output_tokens":5,"num_model_requests":1,"model":"gpt-5.2"}]}],"has_more":false,"next_page":null}"#, + ) + .await + .expect(1); + + let result = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .unwrap(); + + costs.assert_async().await; + failure.assert_async().await; + success.assert_async().await; + assert_eq!(extra_window(&result, "tokens"), "15 tokens"); + assert_eq!(extra_window(&result, "requests"), "1 requests"); +} + +#[tokio::test] +async fn openai_admin_usage_retries_once_then_reports_the_status() { + let mut server = Server::new_async().await; + let costs = server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_status(429) + .with_header("Retry-After", "600") + .with_body(r#"{"error":{"message":"slow down"}}"#) + .expect(2) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("a persistent 429 should fail"); + + costs.assert_async().await; + assert!( + matches!(&error, ProviderError::Other(message) + if message.contains("429") && message.contains("slow down")), + "{error}" + ); +} + +#[tokio::test] +async fn openai_admin_usage_never_retries_auth_failures() { + for status in [401, 403] { + let mut server = Server::new_async().await; + let costs = server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_status(status) + .expect(1) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_usage("sk-test", None, fixed_now(3_600)) + .await + .expect_err("auth failures should surface"); + + costs.assert_async().await; + assert!(matches!(error, ProviderError::AuthRequired), "{status}"); + } +} + +const GRANTS_BODY: &str = r#"{"total_granted":100.0,"total_used":25.0,"total_available":75.0}"#; + +async fn mock_admin_status(server: &mut ServerGuard, status: usize) -> mockito::Mock { + server + .mock("GET", COSTS_PATH) + .match_query(Matcher::Any) + .with_status(status) + .create_async() + .await +} + +#[tokio::test] +async fn openai_unscoped_key_falls_back_to_balance_on_admin_auth_failure() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 403).await; + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(1) + .create_async() + .await; + + let result = provider(&server) + .fetch_admin_or_balance(&admin_credential(), None, fixed_now(3_600)) + .await + .unwrap(); + + grants.assert_async().await; + assert_eq!(result.source_label, "billing-api"); + assert_eq!( + result.usage.login_method.as_deref(), + Some("API balance: $75.00") + ); +} + +#[tokio::test] +async fn openai_unscoped_key_falls_back_to_balance_during_admin_outage() { + let mut server = Server::new_async().await; + let costs = mock_admin_status(&mut server, 500).await.expect(2); + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(1) + .create_async() + .await; + + let result = provider(&server) + .fetch_admin_or_balance(&admin_credential(), None, fixed_now(3_600)) + .await + .unwrap(); + + costs.assert_async().await; + grants.assert_async().await; + assert_eq!(result.source_label, "billing-api"); +} + +#[tokio::test] +async fn openai_project_scoped_admin_key_never_uses_the_unfiltered_balance() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 403).await; + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(0) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_or_balance(&admin_credential(), Some("proj_abc"), fixed_now(3_600)) + .await + .expect_err("a project-scoped admin failure should not fall back"); + + grants.assert_async().await; + assert!(matches!(error, ProviderError::AuthRequired)); +} + +#[tokio::test] +async fn openai_project_scoped_plain_key_keeps_the_balance_fallback() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 403).await; + let grants = server + .mock("GET", GRANTS_PATH) + .with_body(GRANTS_BODY) + .expect(1) + .create_async() + .await; + let credential = ApiCredential { + key: "sk-plain".to_string(), + is_admin: false, + }; + + let result = provider(&server) + .fetch_admin_or_balance(&credential, Some("proj_abc"), fixed_now(3_600)) + .await + .unwrap(); + + grants.assert_async().await; + assert_eq!(result.source_label, "billing-api"); +} + +#[tokio::test] +async fn openai_balance_failure_keeps_the_admin_outage_error() { + let mut server = Server::new_async().await; + let _costs = mock_admin_status(&mut server, 500).await; + let _grants = server + .mock("GET", GRANTS_PATH) + .with_status(404) + .create_async() + .await; + + let error = provider(&server) + .fetch_admin_or_balance(&admin_credential(), None, fixed_now(3_600)) + .await + .expect_err("both endpoints failed"); + + assert!( + matches!(&error, ProviderError::Other(message) if message.contains("costs returned status 500")), + "{error}" + ); +}