diff --git a/CHANGELOG.md b/CHANGELOG.md index a0751c1301d..aeed216403d 100755 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,7 @@ ## [Windows] Unreleased +- OpenAI API: each Admin refresh also reads the last hour of 1-minute usage from the Admin Usage API (input plus output tokens, cached input included, Batch API traffic kept out of limit headroom) and carries a Tokens/min grid with a reporting frontier, so buckets that arrive late show as pending instead of zero. The backend can also load 15-minute, 1-hour or 6-hour windows, read per-project rate-limit caps, and measure headroom against manual per-model organization limits stored in settings. The per-refresh request can be turned off, and a failure never affects the daily usage. The chart, menu bar token and Settings controls come in a later release (upstream #4495). - **Claude**: show the plan renewal or expiry date from claude.ai billing under the quota rows when your cookie session supplies it, and make saved limit-reset credits, Extra usage (including a prepaid balance), and plan dates independently selectable under Visible usage items. Existing settings keep all three visible; Admin API spend is unaffected (upstream #4324, #4484). - Devin: use Devin's official transparent hexagon mark for the provider icon in Settings and the CLI dashboard, instead of the settings-gear badge and the letter fallback (upstream #4458). - Kimi (China): use the current Kimi Desktop session from its Local Storage (`%APPDATA%\kimi-desktop`) for web usage and monthly membership enrichment, and check that session again before each request, so a signed-out, replaced, or expired Desktop token is never sent or recovered from an older record (upstream #4416, #3536). diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs index 178d4b8bd8a..2fa559fb0fe 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge.rs @@ -4,9 +4,13 @@ mod openai_usage_tests; pub(crate) mod pace; mod quota_block; mod status; +mod throughput; +#[cfg(test)] +mod throughput_tests; pub(crate) use openai_usage::OpenAiApiUsageSnapshot; pub use quota_block::MonthlyLimitBlockSnapshot; pub(crate) use status::{compact_tray_status_label, friendly_provider_error}; +pub(crate) use throughput::ProviderThroughputSnapshot; use super::*; use codexbar::core::BlockedWindows; @@ -381,6 +385,9 @@ pub struct ProviderUsageSnapshot { /// Per-day OpenAI Admin API history for the daily usage chart. #[serde(default, skip_serializing_if = "Option::is_none")] pub open_ai_api_usage: Option, + /// Last-hour Tokens/min grid (OpenAI Admin API 1-minute buckets). + #[serde(default, skip_serializing_if = "Option::is_none")] + pub throughput: Option, #[serde(skip_serializing_if = "Option::is_none", default)] pub session_equivalent_forecast: Option, /// Recorded remaining-quota burndown for the selected window @@ -713,6 +720,7 @@ impl ProviderUsageSnapshot { fetch_duration_ms: None, wayfinder_usage: result.wayfinder_usage.clone(), open_ai_api_usage: result.open_ai_api_usage.as_ref().map(Into::into), + throughput: result.throughput.as_ref().map(Into::into), session_equivalent_forecast, quota_burndown, } @@ -773,6 +781,7 @@ impl ProviderUsageSnapshot { fetch_duration_ms: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, quota_burndown: None, } diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge/throughput.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge/throughput.rs new file mode 100644 index 00000000000..9bfa143730b --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge/throughput.rs @@ -0,0 +1,91 @@ +//! Bridge DTOs for the last-hour Tokens/min grid (`throughput`, upstream 0.75.0 +//! `ProviderThroughputSnapshot`, #4495). +//! +//! Mirrors `ProviderThroughputSnapshot` in `src/types/bridge.ts`. Times are epoch seconds. +//! Minutes after `frontier` are `isPending` and never count toward peaks or utilization. + +use codexbar::core::{ProviderThroughput, ThroughputMinute, ThroughputModel}; +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ProviderThroughputSnapshot { + #[serde(default)] + pub minutes: Vec, + #[serde(default)] + pub frontier: Option, + #[serde(default)] + pub latest_tokens_per_minute: Option, + pub peak_tokens_per_minute: u64, + #[serde(default)] + pub peak_at: Option, + pub five_minute_average_tokens_per_minute: f64, + #[serde(default)] + pub models: Vec, + pub window_minutes: u32, + pub fetched_at: i64, + #[serde(default)] + pub project_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ThroughputMinuteSnapshot { + pub start: i64, + pub total_tokens: u64, + pub requests: u64, + pub is_pending: bool, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(rename_all = "camelCase")] +pub struct ThroughputModelSnapshot { + pub name: String, + pub latest_tpm: u64, + pub peak_tpm: u64, + pub total_tokens: u64, + pub peak_limit_counted_tokens: u64, + #[serde(default)] + pub limit_counted_tokens_by_minute: Vec, +} + +impl From<&ProviderThroughput> for ProviderThroughputSnapshot { + fn from(throughput: &ProviderThroughput) -> Self { + Self { + minutes: throughput.minutes.iter().map(Into::into).collect(), + frontier: throughput.frontier, + latest_tokens_per_minute: throughput.latest_tokens_per_minute, + peak_tokens_per_minute: throughput.peak_tokens_per_minute, + peak_at: throughput.peak_at, + five_minute_average_tokens_per_minute: throughput.five_minute_average_tokens_per_minute, + models: throughput.models.iter().map(Into::into).collect(), + window_minutes: throughput.window_minutes, + fetched_at: throughput.fetched_at, + project_id: throughput.project_id.clone(), + } + } +} + +impl From<&ThroughputMinute> for ThroughputMinuteSnapshot { + fn from(minute: &ThroughputMinute) -> Self { + Self { + start: minute.start, + total_tokens: minute.total_tokens, + requests: minute.requests, + is_pending: minute.is_pending, + } + } +} + +impl From<&ThroughputModel> for ThroughputModelSnapshot { + fn from(model: &ThroughputModel) -> Self { + Self { + name: model.name.clone(), + latest_tpm: model.latest_tpm, + peak_tpm: model.peak_tpm, + total_tokens: model.total_tokens, + peak_limit_counted_tokens: model.peak_limit_counted_tokens, + limit_counted_tokens_by_minute: model.limit_counted_tokens_by_minute.clone(), + } + } +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/bridge/throughput_tests.rs b/apps/desktop-tauri/src-tauri/src/commands/bridge/throughput_tests.rs new file mode 100644 index 00000000000..be8042a4b41 --- /dev/null +++ b/apps/desktop-tauri/src-tauri/src/commands/bridge/throughput_tests.rs @@ -0,0 +1,94 @@ +use crate::commands::ProviderUsageSnapshot; +use codexbar::core::{ + ProviderFetchResult, ProviderId, ProviderThroughput, RateWindow, ThroughputMinute, + ThroughputModel, UsageSnapshot, instantiate_provider, +}; + +fn throughput() -> ProviderThroughput { + ProviderThroughput { + minutes: vec![ + ThroughputMinute { + start: 1_700_179_020, + total_tokens: 120, + requests: 3, + is_pending: false, + }, + ThroughputMinute { + start: 1_700_179_080, + total_tokens: 40, + requests: 1, + is_pending: true, + }, + ], + frontier: Some(1_700_179_020), + latest_tokens_per_minute: Some(120), + peak_tokens_per_minute: 120, + peak_at: Some(1_700_179_020), + five_minute_average_tokens_per_minute: 24.0, + models: vec![ThroughputModel { + name: "gpt-5".to_string(), + latest_tpm: 120, + peak_tpm: 120, + total_tokens: 160, + peak_limit_counted_tokens: 100, + limit_counted_tokens_by_minute: vec![100, 40], + }], + window_minutes: 60, + fetched_at: 1_700_179_105, + project_id: None, + } +} + +fn snapshot(throughput: Option) -> ProviderUsageSnapshot { + let metadata = instantiate_provider(ProviderId::OpenAIApi) + .metadata() + .clone(); + let mut result = + ProviderFetchResult::new(UsageSnapshot::new(RateWindow::new(0.0)), "admin-api"); + result.throughput = throughput; + ProviderUsageSnapshot::from_fetch_result(ProviderId::OpenAIApi, &metadata, &result, None, None) +} + +#[test] +fn fetch_result_throughput_reaches_the_bridge_as_camel_case_epoch_seconds() { + let json = serde_json::to_value(snapshot(Some(throughput()))).unwrap(); + assert_eq!( + json["throughput"], + serde_json::json!({ + "minutes": [ + {"start": 1_700_179_020, "totalTokens": 120, "requests": 3, "isPending": false}, + {"start": 1_700_179_080, "totalTokens": 40, "requests": 1, "isPending": true}, + ], + "frontier": 1_700_179_020, + "latestTokensPerMinute": 120, + "peakTokensPerMinute": 120, + "peakAt": 1_700_179_020, + "fiveMinuteAverageTokensPerMinute": 24.0, + "models": [{ + "name": "gpt-5", + "latestTpm": 120, + "peakTpm": 120, + "totalTokens": 160, + "peakLimitCountedTokens": 100, + "limitCountedTokensByMinute": [100, 40], + }], + "windowMinutes": 60, + "fetchedAt": 1_700_179_105, + "projectId": null, + }) + ); +} + +#[test] +fn missing_throughput_is_omitted_and_round_trips() { + let json = serde_json::to_value(snapshot(None)).unwrap(); + assert!(json.get("throughput").is_none()); + + let original = snapshot(Some(throughput())); + let decoded: ProviderUsageSnapshot = + serde_json::from_value(serde_json::to_value(&original).unwrap()).unwrap(); + assert_eq!(decoded.throughput, original.throughput); + + let legacy: ProviderUsageSnapshot = serde_json::from_value(json).unwrap(); + assert!(legacy.throughput.is_none()); +} diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers.rs b/apps/desktop-tauri/src-tauri/src/commands/providers.rs index 622b0149d9b..7dbc58021e5 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers.rs @@ -480,6 +480,7 @@ pub(crate) fn build_fetch_context( && token_account_kind == Some(codexbar::core::TokenAccountKind::ApiKey)), browser_cookie_import, optional_details_enabled: settings.optional_details_enabled(id), + live_usage_enabled: settings.live_usage_enabled(id), ..FetchContext::default() } } @@ -1843,6 +1844,7 @@ mod reset_backfill_tests { fetch_duration_ms: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, quota_burndown: None, } diff --git a/apps/desktop-tauri/src-tauri/src/commands/providers/reset_backfill.rs b/apps/desktop-tauri/src-tauri/src/commands/providers/reset_backfill.rs index 5a41b498a6b..77dd2b77a8b 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/providers/reset_backfill.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/providers/reset_backfill.rs @@ -151,6 +151,7 @@ mod tests { wayfinder_usage: None, quota_burndown: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, } } diff --git a/apps/desktop-tauri/src-tauri/src/commands/tests.rs b/apps/desktop-tauri/src-tauri/src/commands/tests.rs index d1011d0fba7..4dc5385b3de 100644 --- a/apps/desktop-tauri/src-tauri/src/commands/tests.rs +++ b/apps/desktop-tauri/src-tauri/src/commands/tests.rs @@ -1389,6 +1389,7 @@ fn usage_item_descriptors_keep_raw_ids_and_redact_titles() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -1447,6 +1448,7 @@ fn usage_item_descriptors_include_detail_sections() { pace_authoritative: true, account_identity: None, open_ai_api_usage: None, + throughput: None, last_good_owner: None, }; let snapshot = @@ -1503,6 +1505,7 @@ fn usage_item_descriptors_skip_card_hidden_windows() { pace_authoritative: true, account_identity: None, open_ai_api_usage: None, + throughput: None, last_good_owner: None, }; let mut snapshot = @@ -1564,6 +1567,7 @@ fn claude_snapshot_with_sections(source_label: &str) -> ProviderUsageSnapshot { pace_authoritative: true, account_identity: None, last_good_owner: None, + throughput: None, }; let mut snapshot = ProviderUsageSnapshot::from_fetch_result( ProviderId::Claude, @@ -1879,6 +1883,7 @@ fn provider_cache_upsert_replaces_existing_provider() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -1910,6 +1915,7 @@ fn provider_cache_prunes_disabled_providers() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -1954,6 +1960,7 @@ fn claude_transient_auth_failure_preserves_first_last_good_snapshot() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -1999,6 +2006,7 @@ fn codex_transient_transport_failure_helper_uses_typed_policy() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2038,6 +2046,7 @@ fn claude_repeated_auth_failure_surfaces_error() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2089,6 +2098,7 @@ fn claude_cloudflare_challenge_retains_prior_usage_while_surfaceing_guidance() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2151,6 +2161,7 @@ fn claude_cloudflare_challenge_keeps_prior_usage_when_guidance_surfaces() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2210,6 +2221,7 @@ fn claude_cli_parse_failure_keeps_last_good_every_time() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2265,6 +2277,7 @@ fn claude_hard_credentials_missing_does_not_preserve_stale() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2480,6 +2493,7 @@ fn japanese_provider_snapshot_localizes_weekly_label() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2522,6 +2536,7 @@ fn japanese_provider_snapshot_localizes_pace_reserve_description() { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -2962,6 +2977,7 @@ fn limit_reset_snapshot() -> ProviderUsageSnapshot { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), diff --git a/apps/desktop-tauri/src-tauri/src/powertoys.rs b/apps/desktop-tauri/src-tauri/src/powertoys.rs index 5b59a047dbd..26e4e1c0357 100644 --- a/apps/desktop-tauri/src-tauri/src/powertoys.rs +++ b/apps/desktop-tauri/src-tauri/src/powertoys.rs @@ -224,6 +224,7 @@ mod tests { wayfinder_usage: None, quota_burndown: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, }); let value = serde_json::to_value(snapshot).unwrap(); @@ -283,6 +284,7 @@ mod tests { fetch_duration_ms: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, quota_burndown: None, }); diff --git a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs index 071d98a4d3b..02ecc6c6732 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_bridge.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_bridge.rs @@ -787,6 +787,7 @@ mod tests { wayfinder_usage: None, quota_burndown: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, } } diff --git a/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs b/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs index 2e5a4c5e29f..9dabb12bad0 100644 --- a/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs +++ b/apps/desktop-tauri/src-tauri/src/tray_presentation_tests.rs @@ -78,6 +78,7 @@ fn fake_snapshot_with( pace_authoritative: true, account_organization: None, open_ai_api_usage: None, + throughput: None, tray_status_label: None, fetch_duration_ms: None, wayfinder_usage: None, diff --git a/apps/desktop-tauri/src-tauri/src/usage_metric.rs b/apps/desktop-tauri/src-tauri/src/usage_metric.rs index 55fce4db586..84bf2feef0a 100644 --- a/apps/desktop-tauri/src-tauri/src/usage_metric.rs +++ b/apps/desktop-tauri/src-tauri/src/usage_metric.rs @@ -496,6 +496,7 @@ mod tests { wayfinder_usage: None, quota_burndown: None, open_ai_api_usage: None, + throughput: None, session_equivalent_forecast: None, } } diff --git a/apps/desktop-tauri/src/types/bridge.ts b/apps/desktop-tauri/src/types/bridge.ts index ab9bf01824a..427b0362fec 100644 --- a/apps/desktop-tauri/src/types/bridge.ts +++ b/apps/desktop-tauri/src/types/bridge.ts @@ -991,6 +991,8 @@ export interface ProviderUsageSnapshot { wayfinderUsage?: WayfinderUsageSnapshot | null; /** Per-UTC-day OpenAI Admin API history; only the `openaiapi` Admin path sets it. */ openAiApiUsage?: OpenAiApiUsageSnapshot | null; + /** Last-hour Tokens/min grid; only the `openaiapi` Admin path sets it. */ + throughput?: ProviderThroughputSnapshot | null; sessionEquivalentForecast?: SessionEquivalentForecastSnapshot | null; /** Recorded remaining-quota burndown; Codex and Claude only. */ quotaBurndown?: QuotaBurndownSnapshot | null; @@ -1070,6 +1072,45 @@ export interface OpenAiApiUsageSnapshot { daily: OpenAiApiDailyUsageSnapshot[]; } +/** One minute of the Tokens/min grid. */ +export interface ThroughputMinuteSnapshot { + /** Minute start, epoch seconds. */ + start: number; + totalTokens: number; + requests: number; + /** After the reporting frontier: not reported yet, never a real zero or a peak. */ + isPending: boolean; +} + +/** Per-model throughput, ranked by reported volume, then pending volume, then name. */ +export interface ThroughputModelSnapshot { + name: string; + latestTpm: number; + peakTpm: number; + totalTokens: number; + /** Peak non-pending tokens that count toward a rate limit (Batch API excluded). */ + peakLimitCountedTokens: number; + /** Limit-counted tokens aligned with `ProviderThroughputSnapshot.minutes`. */ + limitCountedTokensByMinute: number[]; +} + +export interface ProviderThroughputSnapshot { + /** Regular ascending one-minute grid. */ + minutes: ThroughputMinuteSnapshot[]; + /** Newest complete minute with reported traffic, epoch seconds. */ + frontier: number | null; + latestTokensPerMinute: number | null; + peakTokensPerMinute: number; + peakAt: number | null; + fiveMinuteAverageTokensPerMinute: number; + models: ThroughputModelSnapshot[]; + windowMinutes: number; + /** Epoch seconds. */ + fetchedAt: number; + /** Project scope; `null` covers the whole organization (only then can limits apply). */ + projectId: string | null; +} + export interface RefreshCompletePayload { providerCount: number; errorCount: number; diff --git a/rust/src/cli/diagnose.rs b/rust/src/cli/diagnose.rs index 623389af5c5..d1609daa362 100644 --- a/rust/src/cli/diagnose.rs +++ b/rust/src/cli/diagnose.rs @@ -194,6 +194,7 @@ async fn collect_provider_diagnostic( // Diagnostics keep the short optional-join grace (upstream #2583 gate). requires_optional_usage_completeness: false, optional_details_enabled: settings.optional_details_enabled(provider_id), + live_usage_enabled: false, }; let fetch_result = provider.fetch_usage(&ctx).await; diff --git a/rust/src/cli/fetch_context.rs b/rust/src/cli/fetch_context.rs index 4e3b9719536..3fd05d12f8b 100644 --- a/rust/src/cli/fetch_context.rs +++ b/rust/src/cli/fetch_context.rs @@ -39,6 +39,7 @@ mod tests { browser_cookie_import: false, requires_optional_usage_completeness: false, optional_details_enabled: false, + live_usage_enabled: false, } } diff --git a/rust/src/cli/guard.rs b/rust/src/cli/guard.rs index da5e9f138b6..e276e141ea8 100644 --- a/rust/src/cli/guard.rs +++ b/rust/src/cli/guard.rs @@ -329,6 +329,7 @@ async fn fetch_guard_outcome( // Guard checks keep the short optional-join grace. requires_optional_usage_completeness: false, optional_details_enabled: false, + live_usage_enabled: false, }; match provider.fetch_usage(&ctx).await { diff --git a/rust/src/cli/hooks.rs b/rust/src/cli/hooks.rs index 7592f4bdbcf..9b4d6d9506d 100644 --- a/rust/src/cli/hooks.rs +++ b/rust/src/cli/hooks.rs @@ -307,6 +307,7 @@ async fn hooks_watch_observation( // Hook watches keep the short optional-join grace. requires_optional_usage_completeness: false, optional_details_enabled: settings.optional_details_enabled(provider_id), + live_usage_enabled: false, }; if ctx.api_key.is_none() { diff --git a/rust/src/cli/serve/dashboard/source.rs b/rust/src/cli/serve/dashboard/source.rs index 9bf69fe277a..34ef383a026 100644 --- a/rust/src/cli/serve/dashboard/source.rs +++ b/rust/src/cli/serve/dashboard/source.rs @@ -186,6 +186,7 @@ fn provider_fetch_context() -> FetchContext { browser_cookie_import: false, requires_optional_usage_completeness: false, optional_details_enabled: false, + live_usage_enabled: false, } } @@ -320,6 +321,7 @@ async fn collect_claude_accounts( browser_cookie_import: false, requires_optional_usage_completeness: false, optional_details_enabled: false, + live_usage_enabled: false, }; populate_api_region_from_settings(ProviderId::Claude, settings, &mut ctx); set.spawn(async move { diff --git a/rust/src/cli/serve/data.rs b/rust/src/cli/serve/data.rs index ae212bdc464..ffc05631c94 100644 --- a/rust/src/cli/serve/data.rs +++ b/rust/src/cli/serve/data.rs @@ -147,6 +147,7 @@ pub(super) async fn usage_response( // for the full completeness window. requires_optional_usage_completeness: false, optional_details_enabled: false, + live_usage_enabled: false, }; let rows = collect_usage_rows( diff --git a/rust/src/cli/usage.rs b/rust/src/cli/usage.rs index 01854c1f560..726f5d74435 100755 --- a/rust/src/cli/usage.rs +++ b/rust/src/cli/usage.rs @@ -326,6 +326,7 @@ fn build_usage_fetch_context(args: &UsageArgs, source_mode: SourceMode) -> Fetch // Per-provider opt-ins live in settings, which this shared context // does not load; the optional breakdowns stay off here. optional_details_enabled: false, + live_usage_enabled: false, } } diff --git a/rust/src/core/mod.rs b/rust/src/core/mod.rs index 5646aa66bbd..1c10ad2a6c8 100755 --- a/rust/src/core/mod.rs +++ b/rust/src/core/mod.rs @@ -31,6 +31,7 @@ mod redactor; mod session_equivalent_forecast; mod session_quota; mod sqlite; +mod throughput; mod timezone; mod token_accounts; mod usage_pace; @@ -68,6 +69,7 @@ pub use redactor::*; pub use session_equivalent_forecast::*; pub use session_quota::*; pub use sqlite::*; +pub use throughput::*; pub use timezone::{local_timezone_name, try_local_timezone_name}; pub use token_accounts::*; pub use usage_pace::*; diff --git a/rust/src/core/provider.rs b/rust/src/core/provider.rs index 43bf067aaac..90caa7b400f 100755 --- a/rust/src/core/provider.rs +++ b/rust/src/core/provider.rs @@ -1011,6 +1011,11 @@ pub struct FetchContext { /// (LiteLLM model activity, Claude workspace spend). Scoped to the provider /// being fetched; false everywhere the setting is not consulted. pub optional_details_enabled: bool, + + /// OpenAI API: also fetch the last hour of 1-minute usage for the Tokens/min chart + /// (upstream `OPENAI_LIVE_MINUTES`). The app sets it from settings (default on); + /// paths with no throughput consumer leave it off to save the request. + pub live_usage_enabled: bool, } impl Default for FetchContext { @@ -1033,6 +1038,7 @@ impl Default for FetchContext { browser_cookie_import: false, requires_optional_usage_completeness: false, optional_details_enabled: false, + live_usage_enabled: false, } } } diff --git a/rust/src/core/throughput.rs b/rust/src/core/throughput.rs new file mode 100644 index 00000000000..8a60887b2c9 --- /dev/null +++ b/rust/src/core/throughput.rs @@ -0,0 +1,47 @@ +//! Compact per-minute token throughput behind a provider's Tokens/min chart and +//! menu-bar token (upstream 0.75.0 `ProviderThroughputSnapshot`, #4495). +//! +//! Times are epoch seconds. Minutes form a regular ascending grid; minutes after the +//! reporting `frontier` are `is_pending` (the provider has not reported them yet) and +//! never count toward peaks or limit utilization. + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ProviderThroughput { + pub minutes: Vec, + /// Newest complete minute with reported traffic; `None` when no minute is complete yet. + pub frontier: Option, + /// Tokens in the frontier minute; `None` without a frontier. + pub latest_tokens_per_minute: Option, + pub peak_tokens_per_minute: u64, + pub peak_at: Option, + pub five_minute_average_tokens_per_minute: f64, + /// Ranked by reported volume, then pending volume, then name. + pub models: Vec, + pub window_minutes: u32, + pub fetched_at: i64, + /// Project the traffic is filtered to; `None` covers the whole organization. + pub project_id: Option, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ThroughputMinute { + pub start: i64, + pub total_tokens: u64, + pub requests: u64, + pub is_pending: bool, +} + +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ThroughputModel { + pub name: String, + pub latest_tpm: u64, + pub peak_tpm: u64, + pub total_tokens: u64, + /// Peak non-pending tokens that count toward a rate limit (Batch API excluded). + pub peak_limit_counted_tokens: u64, + /// Limit-counted tokens aligned with [`ProviderThroughput::minutes`]; pending + /// values are retained but excluded from utilization. + pub limit_counted_tokens_by_minute: Vec, +} diff --git a/rust/src/core/usage_snapshot.rs b/rust/src/core/usage_snapshot.rs index 3c772f384d1..8c1df3ab089 100755 --- a/rust/src/core/usage_snapshot.rs +++ b/rust/src/core/usage_snapshot.rs @@ -794,6 +794,11 @@ pub struct ProviderFetchResult { #[serde(skip)] pub open_ai_api_usage: Option, + /// Last-hour tokens-per-minute grid (OpenAI Admin API 1-minute buckets). Transient + /// like [`Self::open_ai_api_usage`]: projected through the bridge, never serialized. + #[serde(skip)] + pub throughput: Option, + /// Transient non-quota inventory for provider-specific display. /// /// The field is intentionally skipped by serde: it belongs to the current @@ -847,6 +852,7 @@ impl ProviderFetchResult { cost: None, wayfinder_usage: None, open_ai_api_usage: None, + throughput: None, inventory: Vec::new(), reset_credits: None, display_details: Vec::new(), @@ -903,6 +909,12 @@ impl ProviderFetchResult { self } + /// Attach the per-minute token throughput. + pub fn with_throughput(mut self, throughput: super::ProviderThroughput) -> Self { + self.throughput = Some(throughput); + self + } + /// Attach one display-only inventory item without exposing redemption IDs. pub fn with_inventory_item(mut self, item: ProviderInventoryItem) -> Self { self.inventory.push(item); @@ -994,6 +1006,32 @@ mod tests { assert!(decoded.open_ai_api_usage.is_none()); } + #[test] + fn throughput_is_transient_and_not_serialized() { + let usage = UsageSnapshot::new(RateWindow::new(25.0)); + let result = ProviderFetchResult::new(usage, "admin-api").with_throughput( + super::super::ProviderThroughput { + minutes: Vec::new(), + frontier: None, + latest_tokens_per_minute: None, + peak_tokens_per_minute: 0, + peak_at: None, + five_minute_average_tokens_per_minute: 0.0, + models: Vec::new(), + window_minutes: 60, + fetched_at: 1_700_000_000, + project_id: None, + }, + ); + + assert!(result.throughput.is_some()); + let encoded = serde_json::to_value(&result).unwrap(); + assert!(encoded.get("throughput").is_none()); + + let decoded: ProviderFetchResult = serde_json::from_value(encoded).unwrap(); + assert!(decoded.throughput.is_none()); + } + #[test] fn display_details_reject_invalid_shapes_and_duplicate_ids() { let usage = UsageSnapshot::new(RateWindow::new(25.0)); diff --git a/rust/src/providers/openaiapi/live.rs b/rust/src/providers/openaiapi/live.rs new file mode 100644 index 00000000000..499a544c1d5 --- /dev/null +++ b/rust/src/providers/openaiapi/live.rs @@ -0,0 +1,819 @@ +//! OpenAI Live Usage: tokens per minute from the Admin API's 1-minute completion buckets +//! (upstream 0.75.0 #4495: `OpenAILiveUsage`, `OpenAILiveUsageBuilder`, +//! `OpenAILiveUsageFetcher`, and the `liveUsage()` step of `Resources/Plugins/openai.js`). +//! +//! Two fetches read `GET /v1/organization/usage/completions` with `bucket_width=1m`, +//! `group_by=model` and `group_by=batch`, in the same project scope as the daily view: +//! - Each refresh makes one unpaginated `limit=60` request for the last hour +//! ([`OpenAIApiProvider::fetch_refresh_throughput`]). Any failure drops only the +//! Tokens/min data; the daily usage stays usable. +//! - The on-demand Live Usage window ([`OpenAIApiProvider::fetch_live_usage`]) pages through +//! a 15-minute, 1-hour or 6-hour window and reports typed errors. +//! +//! Token totals are `input + output`; cached input is a subset of input. Batch API traffic +//! is included in totals but excluded from limit-counted tokens because it has its own +//! rate-limit pool. Buckets commonly arrive 2-3 minutes late, so every view stops at a +//! shared reporting frontier and treats later minutes as pending instead of zero. + +use chrono::{DateTime, Utc}; +use serde::Deserialize; +use std::collections::{BTreeMap, HashMap, HashSet}; + +use super::{OpenAIApiProvider, clean_project_id}; +use crate::core::{ProviderThroughput, ThroughputMinute, ThroughputModel}; + +/// Windows the Live Usage window offers, in minutes. +pub const LIVE_WINDOW_MINUTES: [u32; 3] = [15, 60, 360]; +const REFRESH_WINDOW_MINUTES: u32 = 60; +/// Largest integer a JavaScript number holds exactly; upstream rejects anything above it. +const MAX_SAFE_COUNT: u64 = 9_007_199_254_740_991; +/// Upstream card bound for each per-model value of the refresh payload. +const MAX_CARD_COUNT: u64 = 900_719_925_474; +/// 9999-12-31T23:59:59Z. +const MAX_EPOCH_SECONDS: i64 = 253_402_300_799; +const MAX_WINDOW_PAGES: usize = 20; +const MAX_RESULTS_PER_BUCKET: usize = 400; +const MAX_MODELS_PER_BUCKET: usize = 200; +const MAX_REFRESH_ENTRIES: usize = 10_000; +/// The frontier is searched within the last 10 complete minutes. +const FRONTIER_LOOKBACK_SECONDS: i64 = 9 * 60; +const FIVE_MINUTES_SECONDS: i64 = 5 * 60; +const OTHER_MODELS: &str = "Other"; +const UNKNOWN_MODEL: &str = "unknown"; + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum LiveUsageError { + #[error("Configure an OpenAI organization Admin API key in Settings → Providers → OpenAI.")] + MissingCredential, + #[error( + "Live usage needs an OpenAI organization Admin API key (sk-admin-…). Project and service-account keys cannot read organization usage." + )] + AdminKeyRequired, + #[error("OpenAI live usage is rate limited.")] + RateLimited, + #[error("OpenAI live usage API failed (HTTP {0}).")] + ApiFailure(u16), + #[error("Could not parse the OpenAI live usage response.")] + ParseFailure, + #[error("OpenAI live usage supports 15-minute, 1-hour, and 6-hour windows.")] + InvalidWindow, + #[error("Could not reach the OpenAI live usage API.")] + Network, +} + +/// One model's usage in one minute. `batch_*` is the Batch API share of the totals. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LiveModelUsage { + pub name: String, + pub input_tokens: u64, + pub cached_input_tokens: u64, + pub output_tokens: u64, + pub requests: u64, + pub batch_tokens: u64, + pub batch_requests: u64, +} + +impl LiveModelUsage { + pub fn total_tokens(&self) -> u64 { + self.input_tokens + self.output_tokens + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LiveMinute { + /// Epoch seconds, minute-aligned. + pub start: i64, + /// Ascending by name. + pub models: Vec, +} + +/// A validated, regular minute grid ending at the minute containing `fetched_at`. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OpenAiLiveUsage { + pub minutes: Vec, + pub window_minutes: u32, + /// Epoch seconds. + pub fetched_at: i64, + /// Project the usage is scoped to; `None` covers the whole organization. + pub project_id: Option, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum LiveModelFilter<'a> { + AllModels, + Model(&'a str), +} + +impl LiveModelFilter<'_> { + fn includes(self, name: &str) -> bool { + match self { + Self::AllModels => true, + Self::Model(model) => model == name, + } + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LivePoint { + pub start: i64, + pub total_tokens: u64, + pub uncached_input_tokens: u64, + pub cached_input_tokens: u64, + pub output_tokens: u64, + pub requests: u64, + pub is_in_progress: bool, + pub is_pending: bool, + pub batch_tokens: u64, + pub batch_requests: u64, +} + +impl LivePoint { + /// Cached input counts toward TPM; Batch API traffic has its own pool and is excluded. + pub fn limit_counted_tokens(&self) -> u64 { + self.total_tokens - self.batch_tokens + } + + pub fn limit_counted_requests(&self) -> u64 { + self.requests - self.batch_requests + } +} + +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct LiveModelStack { + pub minute: i64, + /// A top model's name, or `Other` for the rest. + pub model_or_other: String, + pub tokens: u64, + pub is_in_progress: bool, + pub is_pending: bool, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct LiveSummary { + pub latest_tpm: u64, + pub latest_minute_start: Option, + pub five_minute_average_tpm: f64, + pub peak_tpm: u64, + pub peak_minute_start: Option, + pub total_tokens: u64, + pub input_tokens: u64, + pub cached_input_tokens: u64, + pub output_tokens: u64, + pub requests: u64, + /// Batch API share of `total_tokens`. + pub batch_tokens: u64, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct LiveModelRow { + pub name: String, + pub latest_tpm: u64, + pub peak_tpm: u64, + pub total_tokens: u64, + pub requests: u64, + /// Manual organization limit, when one resolves for this model. + pub tpm_limit: Option, + pub rpm_limit: Option, + /// Highest per-project cap for context only; never drives utilization. + pub project_cap: Option, + pub peak_tpm_percent: Option, + pub latest_tpm_percent: Option, +} + +impl OpenAiLiveUsage { + fn is_in_progress(&self, start: i64) -> bool { + start + 60 > self.fetched_at + } + + pub(super) fn is_pending(frontier: Option, start: i64) -> bool { + frontier.is_none_or(|frontier| start > frontier) + } + + /// Newest complete minute with any reported tokens within the last 10 complete + /// minutes, else the last complete minute, so lagging buckets never read as a drop to + /// zero while genuine idle periods still do. + pub fn reporting_frontier(&self) -> Option { + let complete: Vec<&LiveMinute> = self + .minutes + .iter() + .filter(|minute| minute.start + 60 <= self.fetched_at) + .collect(); + let last = complete.last()?; + let cutoff = last.start - FRONTIER_LOOKBACK_SECONDS; + Some( + complete + .iter() + .rev() + .find(|minute| { + minute.start >= cutoff + && minute.models.iter().any(|model| model.total_tokens() > 0) + }) + .map_or(last.start, |minute| minute.start), + ) + } + + /// Ranked by reported (non-pending) volume so partial buckets cannot evict visible + /// models from the top series; pending-only volume, then the name, break ties. + pub fn model_names(&self) -> Vec { + let frontier = self.reporting_frontier(); + let mut reported: HashMap<&str, u64> = HashMap::new(); + let mut totals: HashMap<&str, u64> = HashMap::new(); + for minute in &self.minutes { + let is_reported = frontier.is_some_and(|frontier| minute.start <= frontier); + for model in &minute.models { + *totals.entry(&model.name).or_default() += model.total_tokens(); + if is_reported { + *reported.entry(&model.name).or_default() += model.total_tokens(); + } + } + } + let mut names: Vec<&str> = totals.keys().copied().collect(); + names.sort_by(|left, right| { + let reported_of = |name: &str| reported.get(name).copied().unwrap_or(0); + reported_of(right) + .cmp(&reported_of(left)) + .then_with(|| totals[right].cmp(&totals[left])) + .then_with(|| left.cmp(right)) + }); + names.into_iter().map(ToOwned::to_owned).collect() + } + + pub fn series(&self, filter: LiveModelFilter<'_>) -> Vec { + let frontier = self.reporting_frontier(); + self.minutes + .iter() + .map(|minute| { + let mut point = LivePoint { + start: minute.start, + total_tokens: 0, + uncached_input_tokens: 0, + cached_input_tokens: 0, + output_tokens: 0, + requests: 0, + is_in_progress: self.is_in_progress(minute.start), + is_pending: Self::is_pending(frontier, minute.start), + batch_tokens: 0, + batch_requests: 0, + }; + for model in minute.models.iter().filter(|m| filter.includes(&m.name)) { + point.total_tokens += model.total_tokens(); + point.uncached_input_tokens += model.input_tokens - model.cached_input_tokens; + point.cached_input_tokens += model.cached_input_tokens; + point.output_tokens += model.output_tokens; + point.requests += model.requests; + point.batch_tokens += model.batch_tokens; + point.batch_requests += model.batch_requests; + } + point + }) + .collect() + } + + /// Per-minute stacks for the top `top_n` models plus `Other`. + pub fn stacked_by_model(&self, top_n: usize) -> Vec { + let frontier = self.reporting_frontier(); + let names = self.model_names(); + let top: Vec<&str> = names.iter().take(top_n).map(String::as_str).collect(); + let has_other = names.len() > top.len(); + let mut labels = top.clone(); + if has_other { + labels.push(OTHER_MODELS); + } + self.minutes + .iter() + .flat_map(|minute| { + let mut totals: HashMap<&str, u64> = HashMap::new(); + for model in &minute.models { + let label = if top.contains(&model.name.as_str()) { + model.name.as_str() + } else { + OTHER_MODELS + }; + *totals.entry(label).or_default() += model.total_tokens(); + } + labels + .iter() + .map(|label| LiveModelStack { + minute: minute.start, + model_or_other: (*label).to_string(), + tokens: totals.get(label).copied().unwrap_or(0), + is_in_progress: self.is_in_progress(minute.start), + is_pending: Self::is_pending(frontier, minute.start), + }) + .collect::>() + }) + .collect() + } + + pub fn summary(&self, filter: LiveModelFilter<'_>) -> LiveSummary { + let points = self.series(filter); + let frontier = self.reporting_frontier(); + let latest = points.iter().find(|point| Some(point.start) == frontier); + let recent: Vec<&LivePoint> = frontier.map_or_else(Vec::new, |frontier| { + points + .iter() + .filter(|point| { + point.start <= frontier && point.start > frontier - FIVE_MINUTES_SECONDS + }) + .collect() + }); + // Pending minutes can hold partial data, so peaks stop at the frontier; a tie keeps + // the earliest minute. + let mut peak: Option<&LivePoint> = None; + for point in points.iter().filter(|point| !point.is_pending) { + if peak.is_none_or(|best| point.total_tokens > best.total_tokens) { + peak = Some(point); + } + } + let sum = |field: fn(&LivePoint) -> u64| points.iter().map(field).sum::(); + LiveSummary { + latest_tpm: latest.map_or(0, |point| point.total_tokens), + latest_minute_start: frontier, + five_minute_average_tpm: if recent.is_empty() { + 0.0 + } else { + recent + .iter() + .map(|point| point.total_tokens as f64) + .sum::() + / recent.len() as f64 + }, + peak_tpm: peak.map_or(0, |point| point.total_tokens), + peak_minute_start: peak.map(|point| point.start), + total_tokens: sum(|point| point.total_tokens), + input_tokens: sum(|point| point.uncached_input_tokens + point.cached_input_tokens), + cached_input_tokens: sum(|point| point.cached_input_tokens), + output_tokens: sum(|point| point.output_tokens), + requests: sum(|point| point.requests), + batch_tokens: sum(|point| point.batch_tokens), + } + } + + /// One row per model in [`Self::model_names`] order, without limits. + pub fn model_rows(&self) -> Vec { + self.model_names() + .into_iter() + .map(|name| { + let summary = self.summary(LiveModelFilter::Model(&name)); + LiveModelRow { + name, + latest_tpm: summary.latest_tpm, + peak_tpm: summary.peak_tpm, + total_tokens: summary.total_tokens, + requests: summary.requests, + tpm_limit: None, + rpm_limit: None, + project_cap: None, + peak_tpm_percent: None, + latest_tpm_percent: None, + } + }) + .collect() + } + + /// The compact snapshot the menu card, tray token and bridge carry. + pub fn throughput(&self) -> ProviderThroughput { + let frontier = self.reporting_frontier(); + let summary = self.summary(LiveModelFilter::AllModels); + ProviderThroughput { + minutes: self + .series(LiveModelFilter::AllModels) + .into_iter() + .map(|point| ThroughputMinute { + start: point.start, + total_tokens: point.total_tokens, + requests: point.requests, + is_pending: point.is_pending, + }) + .collect(), + frontier, + latest_tokens_per_minute: frontier.map(|_| summary.latest_tpm), + peak_tokens_per_minute: summary.peak_tpm, + peak_at: summary.peak_minute_start, + five_minute_average_tokens_per_minute: summary.five_minute_average_tpm, + models: self + .model_names() + .into_iter() + .map(|name| { + let model_summary = self.summary(LiveModelFilter::Model(&name)); + let points = self.series(LiveModelFilter::Model(&name)); + ThroughputModel { + latest_tpm: model_summary.latest_tpm, + peak_tpm: model_summary.peak_tpm, + total_tokens: model_summary.total_tokens, + peak_limit_counted_tokens: points + .iter() + .filter(|point| !point.is_pending) + .map(LivePoint::limit_counted_tokens) + .max() + .unwrap_or(0), + limit_counted_tokens_by_minute: points + .iter() + .map(LivePoint::limit_counted_tokens) + .collect(), + name, + } + }) + .collect(), + window_minutes: self.window_minutes, + fetched_at: self.fetched_at, + project_id: self.project_id.clone(), + } + } +} + +fn checked_add(left: u64, right: u64) -> Result { + left.checked_add(right) + .filter(|sum| *sum <= MAX_SAFE_COUNT) + .ok_or(LiveUsageError::ParseFailure) +} + +/// Upstream `OpenAILiveUsageBuilder.build`: validates every bucket, drops minutes outside +/// the window, merges duplicate model rows, and fills missing minutes with empty ones. +pub(super) fn build( + buckets: Vec, + window_minutes: u32, + fetched_at: i64, + project_id: Option, +) -> Result { + if !LIVE_WINDOW_MINUTES.contains(&window_minutes) + || !(0..=MAX_EPOCH_SECONDS).contains(&fetched_at) + { + return Err(LiveUsageError::ParseFailure); + } + let current = fetched_at / 60 * 60; + let first = current - (i64::from(window_minutes) - 1) * 60; + let mut grid: HashMap> = HashMap::new(); + let mut tokens = 0_u64; + let mut requests = 0_u64; + for bucket in buckets { + if !(0..=MAX_EPOCH_SECONDS).contains(&bucket.start) { + return Err(LiveUsageError::ParseFailure); + } + let minute = bucket.start / 60 * 60; + for model in bucket.models { + if model.cached_input_tokens > model.input_tokens + || model.batch_requests > model.requests + { + return Err(LiveUsageError::ParseFailure); + } + let total = checked_add(model.input_tokens, model.output_tokens)?; + if model.batch_tokens > total { + return Err(LiveUsageError::ParseFailure); + } + if minute < first || minute > current { + continue; + } + tokens = checked_add(tokens, total)?; + requests = checked_add(requests, model.requests)?; + let models = grid.entry(minute).or_default(); + match models.get_mut(&model.name) { + Some(prior) => { + prior.input_tokens = checked_add(prior.input_tokens, model.input_tokens)?; + prior.cached_input_tokens = + checked_add(prior.cached_input_tokens, model.cached_input_tokens)?; + prior.output_tokens = checked_add(prior.output_tokens, model.output_tokens)?; + prior.requests = checked_add(prior.requests, model.requests)?; + prior.batch_tokens = checked_add(prior.batch_tokens, model.batch_tokens)?; + prior.batch_requests = checked_add(prior.batch_requests, model.batch_requests)?; + } + None => { + models.insert(model.name.clone(), model); + } + } + } + } + Ok(OpenAiLiveUsage { + minutes: (0..i64::from(window_minutes)) + .map(|index| { + let start = first + index * 60; + LiveMinute { + start, + models: grid + .remove(&start) + .map(|models| models.into_values().collect()) + .unwrap_or_default(), + } + }) + .collect(), + window_minutes, + fetched_at, + project_id, + }) +} + +#[derive(Debug, Deserialize)] +pub(super) struct LivePage { + data: Vec, + #[serde(default)] + has_more: Option, + #[serde(default)] + next_page: Option, +} + +#[derive(Debug, Deserialize)] +struct LiveBucket { + start_time: i64, + #[serde(default)] + end_time: Option, + results: Vec, +} + +#[derive(Debug, Deserialize)] +struct LiveResult { + model: Option, + batch: Option, + input_tokens: Option, + input_cached_tokens: Option, + output_tokens: Option, + num_model_requests: Option, + input_audio_tokens: Option, + output_audio_tokens: Option, +} + +impl LiveResult { + /// Absent counts are zero; negative counts, cached above input, or an input+output + /// total beyond the safe-integer range are parse failures. + fn usage(&self) -> Result { + let count = |value: Option| { + u64::try_from(value.unwrap_or(0)).map_err(|_| LiveUsageError::ParseFailure) + }; + let input = count(self.input_tokens)?; + let cached = count(self.input_cached_tokens)?; + let output = count(self.output_tokens)?; + let requests = count(self.num_model_requests)?; + // input/output already aggregate text, audio and image; audio is only a breakdown. + count(self.input_audio_tokens)?; + count(self.output_audio_tokens)?; + let total = checked_add(input, output)?; + if cached > input { + return Err(LiveUsageError::ParseFailure); + } + let name = self + .model + .as_deref() + .map(str::trim) + .filter(|name| !name.is_empty()) + .unwrap_or(UNKNOWN_MODEL) + .to_string(); + let batch = self.batch == Some(true); + Ok(LiveModelUsage { + name, + input_tokens: input, + cached_input_tokens: cached, + output_tokens: output, + requests, + batch_tokens: if batch { total } else { 0 }, + batch_requests: if batch { requests } else { 0 }, + }) + } +} + +fn live_query( + start: i64, + end: i64, + limit: u32, + project_id: Option<&str>, + page: Option<&str>, +) -> Vec<(&'static str, String)> { + let mut params = vec![ + ("start_time", start.to_string()), + ("end_time", end.to_string()), + ("bucket_width", "1m".to_string()), + ("group_by", "model".to_string()), + ("group_by", "batch".to_string()), + ("limit", limit.to_string()), + ]; + if let Some(project_id) = clean_project_id(project_id) { + params.push(("project_ids", project_id)); + } + if let Some(page) = page { + params.push(("page", page.to_string())); + } + params +} + +fn status_error(status: reqwest::StatusCode) -> LiveUsageError { + match status.as_u16() { + 401 | 403 => LiveUsageError::AdminKeyRequired, + 429 => LiveUsageError::RateLimited, + code => LiveUsageError::ApiFailure(code), + } +} + +/// Upstream plugin `liveUsage()` bounds for the one-page refresh response. +fn refresh_buckets( + page: LivePage, + start: i64, + fetched_at: i64, +) -> Result, LiveUsageError> { + if page.data.len() > REFRESH_WINDOW_MINUTES as usize || page.has_more.is_none() { + return Err(LiveUsageError::ParseFailure); + } + let mut entries = 0_usize; + let mut total_tokens = 0_u64; + let mut total_requests = 0_u64; + page.data + .into_iter() + .map(|bucket| { + if bucket.start_time < start + || bucket.start_time > fetched_at + || bucket.start_time % 60 != 0 + || bucket.end_time.is_none_or(|end| end <= bucket.start_time) + || bucket.results.len() > MAX_RESULTS_PER_BUCKET + { + return Err(LiveUsageError::ParseFailure); + } + let mut merged: BTreeMap = BTreeMap::new(); + for result in &bucket.results { + let row = result.usage()?; + total_tokens = checked_add(total_tokens, row.total_tokens())?; + total_requests = checked_add(total_requests, row.requests)?; + match merged.get_mut(&row.name) { + Some(prior) => { + prior.input_tokens = checked_add(prior.input_tokens, row.input_tokens)?; + prior.cached_input_tokens = + checked_add(prior.cached_input_tokens, row.cached_input_tokens)?; + prior.output_tokens = checked_add(prior.output_tokens, row.output_tokens)?; + prior.requests = checked_add(prior.requests, row.requests)?; + prior.batch_tokens = checked_add(prior.batch_tokens, row.batch_tokens)?; + prior.batch_requests = + checked_add(prior.batch_requests, row.batch_requests)?; + } + None => { + merged.insert(row.name.clone(), row); + } + } + } + if merged.len() > MAX_MODELS_PER_BUCKET { + return Err(LiveUsageError::ParseFailure); + } + entries += merged.len(); + if entries > MAX_REFRESH_ENTRIES { + return Err(LiveUsageError::ParseFailure); + } + let models: Vec = merged.into_values().collect(); + let within_card_bounds = models.iter().all(|model| { + [ + model.input_tokens, + model.cached_input_tokens, + model.output_tokens, + model.requests, + model.batch_tokens, + model.batch_requests, + ] + .into_iter() + .all(|value| value <= MAX_CARD_COUNT) + }); + if !within_card_bounds { + return Err(LiveUsageError::ParseFailure); + } + Ok(LiveMinute { + start: bucket.start_time, + models, + }) + }) + .collect() +} + +impl OpenAIApiProvider { + async fn get_live_page( + &self, + query: &[(&'static str, String)], + api_key: &str, + ) -> Result<(reqwest::StatusCode, String), LiveUsageError> { + let response = self + .client + .get(&self.endpoints.completions) + .query(query) + .bearer_auth(api_key) + .header("Accept", "application/json") + .send() + .await + .map_err(|_| LiveUsageError::Network)?; + let status = response.status(); + let body = response.text().await.map_err(|_| LiveUsageError::Network)?; + Ok((status, body)) + } + + /// The per-refresh last-hour Tokens/min grid, or `None` when it is unavailable. One + /// request without retries, as upstream; the daily data never depends on it. + pub(super) async fn fetch_refresh_throughput( + &self, + api_key: &str, + project_id: Option<&str>, + now: DateTime, + ) -> Option { + match self.fetch_refresh_live(api_key, project_id, now).await { + Ok(live) => Some(live.throughput()), + Err(error) => { + tracing::debug!(%error, "OpenAI live usage unavailable; daily usage is unaffected"); + None + } + } + } + + pub(super) async fn fetch_refresh_live( + &self, + api_key: &str, + project_id: Option<&str>, + now: DateTime, + ) -> Result { + let fetched_at = now.timestamp(); + let start = fetched_at / 60 * 60 - (i64::from(REFRESH_WINDOW_MINUTES) - 1) * 60; + let query = live_query(start, fetched_at, REFRESH_WINDOW_MINUTES, project_id, None); + let (status, body) = self.get_live_page(&query, api_key).await?; + if status != reqwest::StatusCode::OK { + return Err(status_error(status)); + } + let page: LivePage = + serde_json::from_str(&body).map_err(|_| LiveUsageError::ParseFailure)?; + let buckets = refresh_buckets(page, start, fetched_at)?; + build( + buckets, + REFRESH_WINDOW_MINUTES, + fetched_at, + clean_project_id(project_id), + ) + } + + /// The on-demand Live Usage window: every page of a 15-, 60- or 360-minute window. + /// The key resolves like refreshes do (Preferences, credential store, environment); + /// `project_id` uses the provider's configured project scope. + pub async fn fetch_live_usage( + &self, + api_key: Option<&str>, + project_id: Option<&str>, + window_minutes: u32, + ) -> Result { + if !LIVE_WINDOW_MINUTES.contains(&window_minutes) { + return Err(LiveUsageError::InvalidWindow); + } + let credential = + Self::credential(api_key).map_err(|_| LiveUsageError::MissingCredential)?; + self.fetch_live_window(&credential.key, project_id, window_minutes, Utc::now()) + .await + } + + pub(super) async fn fetch_live_window( + &self, + api_key: &str, + project_id: Option<&str>, + window_minutes: u32, + now: DateTime, + ) -> Result { + if !LIVE_WINDOW_MINUTES.contains(&window_minutes) { + return Err(LiveUsageError::InvalidWindow); + } + let fetched_at = now.timestamp(); + let first = fetched_at / 60 * 60 - (i64::from(window_minutes) - 1) * 60; + let mut buckets = Vec::new(); + let mut cursor: Option = None; + let mut seen = HashSet::new(); + for _ in 0..MAX_WINDOW_PAGES { + let query = live_query( + first, + fetched_at, + window_minutes, + project_id, + cursor.as_deref(), + ); + let (status, body) = self.get_live_page(&query, api_key).await?; + if !status.is_success() { + return Err(status_error(status)); + } + let page: LivePage = + serde_json::from_str(&body).map_err(|_| LiveUsageError::ParseFailure)?; + for bucket in page.data { + if bucket.start_time < 0 { + return Err(LiveUsageError::ParseFailure); + } + buckets.push(LiveMinute { + start: bucket.start_time, + models: bucket + .results + .iter() + .map(LiveResult::usage) + .collect::>()?, + }); + } + if !page.has_more.unwrap_or(false) { + return build( + buckets, + window_minutes, + fetched_at, + clean_project_id(project_id), + ); + } + let next = page + .next_page + .filter(|next| !next.is_empty()) + .ok_or(LiveUsageError::ParseFailure)?; + if !seen.insert(next.clone()) { + return Err(LiveUsageError::ParseFailure); + } + cursor = Some(next); + } + Err(LiveUsageError::ParseFailure) + } +} diff --git a/rust/src/providers/openaiapi/live_tests.rs b/rust/src/providers/openaiapi/live_tests.rs new file mode 100644 index 00000000000..6d4d1177e1c --- /dev/null +++ b/rust/src/providers/openaiapi/live_tests.rs @@ -0,0 +1,1859 @@ +//! Fixtures ported from upstream 0.75.0 `OpenAILiveUsageTests`, `OpenAIThroughputTests` +//! and `OpenAIOrgRateLimitsTests` (#4495). + +use super::live::{ + LiveMinute, LiveModelFilter, LiveModelUsage, LiveUsageError, OpenAiLiveUsage, build, +}; +use super::org_limits::{ + OrgLimitResolver, OrgRateLimit, OrgRateLimits, OrgRateLimitsError, throughput_limit_utilization, +}; +use super::*; +use crate::settings::ModelRateLimit; +use mockito::{Matcher, Server, ServerGuard}; + +const COSTS_PATH: &str = "/v1/organization/costs"; +const COMPLETIONS_PATH: &str = "/v1/organization/usage/completions"; +const PROJECTS_PATH: &str = "/v1/organization/projects"; +const EMPTY_PAGE: &str = r#"{"object":"page","data":[],"has_more":false,"next_page":null}"#; +/// Upstream fixture clock: 25 seconds into the minute starting at 120000. +const NOW: i64 = 120_025; +/// A realistic refresh clock so the daily ranges stay positive: 25 seconds into `BASE`. +const BASE: i64 = 1_700_179_200; +const REFRESH_NOW: i64 = BASE + 25; +const REFRESH_START: i64 = BASE - 59 * 60; + +fn at(seconds: i64) -> DateTime { + Utc.timestamp_opt(seconds, 0).single().unwrap() +} + +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!("{}/v1/dashboard/billing/credit_grants", server.url()), + costs: format!("{}{COSTS_PATH}", server.url()), + completions: format!("{}{COMPLETIONS_PATH}", server.url()), + projects: format!("{}{PROJECTS_PATH}", server.url()), + }; + provider.retry = RetryPolicy { + base_delay: Duration::ZERO, + max_delay: Duration::ZERO, + }; + provider +} + +fn admin() -> ApiCredential { + ApiCredential { + key: "test-admin-key".to_string(), + is_admin: true, + } +} + +fn page(buckets: &str, more: bool, cursor: Option<&str>) -> String { + let next = cursor.map_or_else(|| "null".to_string(), |cursor| format!("\"{cursor}\"")); + format!(r#"{{"data":[{buckets}],"has_more":{more},"next_page":{next},"unknown":true}}"#) +} + +/// The exact first-page query `live_query` sends (parameter order included). +fn live_query_string(start: i64, end: i64, limit: u32, project: Option<&str>) -> String { + let mut query = format!( + "start_time={start}&end_time={end}&bucket_width=1m&group_by=model&group_by=batch&limit={limit}" + ); + if let Some(project) = project { + query.push_str("&project_ids="); + query.push_str(&project.replace('&', "%26").replace(' ', "+")); + } + query +} + +async fn mock( + server: &mut ServerGuard, + path: &str, + query: Matcher, + status: usize, + body: &str, +) -> mockito::Mock { + server + .mock("GET", path) + .match_query(query) + .match_header("authorization", "Bearer test-admin-key") + .with_status(status) + .with_body(body) + .create_async() + .await +} + +fn model(name: &str, input: u64, cached: u64, output: u64, requests: u64) -> LiveModelUsage { + LiveModelUsage { + name: name.to_string(), + input_tokens: input, + cached_input_tokens: cached, + output_tokens: output, + requests, + batch_tokens: 0, + batch_requests: 0, + } +} + +fn snapshot(minutes: Vec, project_id: Option<&str>) -> OpenAiLiveUsage { + OpenAiLiveUsage { + minutes, + window_minutes: 15, + fetched_at: NOW, + project_id: project_id.map(str::to_string), + } +} + +/// Upstream `frontierSnapshot`: 15 minutes from 119160 with one model per active minute. +fn frontier_snapshot(active: &[(i64, &str, u64)]) -> OpenAiLiveUsage { + snapshot( + (0..15_i64) + .map(|index| LiveMinute { + start: 119_160 + index * 60, + models: active + .iter() + .filter(|(minute, _, _)| *minute == index) + .map(|(_, name, tokens)| model(name, *tokens, 0, 0, 1)) + .collect(), + }) + .collect(), + None, + ) +} + +// ---- On-demand Live Usage window ------------------------------------------------------- + +#[tokio::test] +async fn window_requests_minute_buckets_model_and_batch_grouping_project_scope_and_pages() { + let mut server = Server::new_async().await; + let first_query = live_query_string(119_160, NOW, 15, Some("proj_&test")); + let first = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(first_query.clone()), + 200, + &page("", true, Some("next & page")), + ) + .await + .expect(1); + let second = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(format!("{first_query}&page=next+%26+page")), + 200, + &page("", false, None), + ) + .await + .expect(1); + + let live = provider(&server) + .fetch_live_window("test-admin-key", Some("proj_&test"), 15, at(NOW)) + .await + .unwrap(); + + first.assert_async().await; + second.assert_async().await; + assert_eq!(live.project_id.as_deref(), Some("proj_&test")); + assert_eq!(live.minutes.len(), 15); + assert_eq!(live.minutes.first().unwrap().start, 119_160); + assert_eq!(live.minutes.last().unwrap().start, 120_000); +} + +#[tokio::test] +async fn window_parsing_merges_pages_floors_buckets_drops_outside_minutes_and_defaults_counts() { + let mut server = Server::new_async().await; + let first_query = live_query_string(119_160, NOW, 15, None); + let _first = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(first_query.clone()), + 200, + &page( + r#"{"start_time":119947,"results":[ + {"model":"model-a","input_tokens":100,"input_cached_tokens":20,"output_tokens":50,"num_model_requests":2}, + {"model":"model-b","output_tokens":30}, {"model":null}, {"model":""}]}, + {"start_time":119100,"results":[{"input_tokens":999}]}, + {"start_time":120060,"results":[{"input_tokens":999}]}"#, + true, + Some("two"), + ), + ) + .await; + let _second = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(format!("{first_query}&page=two")), + 200, + &page( + r#"{"start_time":119940,"results":[{"model":"model-a","input_tokens":5,"output_tokens":5}]}, + {"start_time":120000,"results":[{"model":"model-a","input_tokens":500}]}"#, + false, + None, + ), + ) + .await; + + let live = provider(&server) + .fetch_live_window("test-admin-key", None, 15, at(NOW)) + .await + .unwrap(); + + assert_eq!(live.model_names(), ["model-a", "model-b", "unknown"]); + let all = live.series(LiveModelFilter::AllModels); + assert_eq!(all.len(), 15); + assert!(all[..13].iter().all(|point| point.total_tokens == 0)); + assert_eq!(all[13].total_tokens, 190); + assert_eq!(all[13].uncached_input_tokens, 85); + assert_eq!(all[13].cached_input_tokens, 20); + assert_eq!(all[13].output_tokens, 85); + assert_eq!(all[13].requests, 2); + assert!(!all[13].is_in_progress); + assert!(all[14].is_in_progress); + assert_eq!( + live.series(LiveModelFilter::Model("model-b"))[13].total_tokens, + 30 + ); + assert!( + live.series(LiveModelFilter::Model("absent")) + .iter() + .all(|point| point.total_tokens == 0) + ); + let summary = live.summary(LiveModelFilter::AllModels); + assert_eq!(summary.latest_tpm, 190); + assert_eq!(summary.latest_minute_start, Some(all[13].start)); + assert_eq!(summary.five_minute_average_tpm, 38.0); + // The in-progress minute's partial 500 tokens count toward totals but never become the peak. + assert_eq!(summary.peak_tpm, 190); + assert_eq!(summary.peak_minute_start, Some(all[13].start)); + assert_eq!(summary.total_tokens, 690); + assert_eq!(summary.input_tokens, 605); + assert_eq!(summary.cached_input_tokens, 20); + assert_eq!(summary.output_tokens, 85); + assert_eq!(summary.requests, 2); + let row = &live.model_rows()[0]; + assert_eq!( + ( + row.name.as_str(), + row.latest_tpm, + row.peak_tpm, + row.total_tokens, + row.requests + ), + ("model-a", 160, 160, 660, 2) + ); +} + +#[tokio::test] +async fn window_model_ordering_breaks_ties_by_name_and_stacks_five_models_plus_other() { + let mut server = Server::new_async().await; + let entries = (1..=7) + .map(|index| format!(r#"{{"model":"model-{index}","input_tokens":{index}}}"#)) + .collect::>() + .join(","); + let _page = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Any, + 200, + &page( + &format!(r#"{{"start_time":119940,"results":[{entries}]}}"#), + false, + None, + ), + ) + .await; + + let live = provider(&server) + .fetch_live_window("test-admin-key", None, 15, at(NOW)) + .await + .unwrap(); + + assert_eq!( + live.model_names(), + [ + "model-7", "model-6", "model-5", "model-4", "model-3", "model-2", "model-1" + ] + ); + let stacks = live.stacked_by_model(5); + assert_eq!(stacks.len(), 15 * 6); + let other: u64 = stacks + .iter() + .filter(|stack| stack.model_or_other == "Other") + .map(|stack| stack.tokens) + .sum(); + assert_eq!(other, 3); + assert_eq!( + stacks.iter().map(|stack| stack.tokens).sum::(), + live.summary(LiveModelFilter::AllModels).total_tokens + ); + + let tied = snapshot( + vec![LiveMinute { + start: NOW, + models: vec![model("z", 1, 0, 0, 0), model("a", 1, 0, 0, 0)], + }], + None, + ); + assert_eq!(tied.model_names(), ["a", "z"]); +} + +#[tokio::test] +async fn window_http_errors_map_to_typed_errors() { + for (status, expected) in [ + (401, LiveUsageError::AdminKeyRequired), + (403, LiveUsageError::AdminKeyRequired), + (429, LiveUsageError::RateLimited), + (500, LiveUsageError::ApiFailure(500)), + ] { + let mut server = Server::new_async().await; + let _page = mock(&mut server, COMPLETIONS_PATH, Matcher::Any, status, "{}").await; + let error = provider(&server) + .fetch_live_window("test-admin-key", None, 60, at(NOW)) + .await + .unwrap_err(); + assert_eq!(error, expected, "{status}"); + } +} + +#[tokio::test] +async fn window_malformed_responses_and_invalid_counts_are_parse_errors() { + for body in [ + "not json", + "{}", + r#"{"data":null}"#, + r#"{"data":[{"results":[]}]}"#, + r#"{"data":[{"start_time":120000}]}"#, + r#"{"data":[{"start_time":120000,"results":[{"input_tokens":-1}]}]}"#, + r#"{"data":[{"start_time":120000,"results":[{"output_tokens":1.5}]}]}"#, + r#"{"data":[{"start_time":120000,"results":[{"num_model_requests":true}]}]}"#, + r#"{"data":[{"start_time":120000,"results":[{"input_audio_tokens":-1}]}]}"#, + r#"{"data":[{"start_time":120000,"results":[{"input_tokens":9223372036854775807,"output_tokens":1}]}]}"#, + r#"{"data":[{"start_time":119940,"results":[{"batch":"yes"}]}]}"#, + r#"{"data":[{"start_time":119940,"results":[{"batch":1}]}]}"#, + r#"{"data":[{"start_time":119940,"results":[{"batch":[]}]}]}"#, + r#"{"data":[{"start_time":119940,"results":[{"batch":{}}]}]}"#, + r#"{"data":[{"start_time":119940,"results":[{"batch":true,"input_tokens":9007199254740991,"output_tokens":1}]}]}"#, + r#"{"data":[{"start_time":119940,"results":[{"batch":true,"input_tokens":9223372036854775807,"output_tokens":1}]}]}"#, + ] { + let mut server = Server::new_async().await; + let _page = mock(&mut server, COMPLETIONS_PATH, Matcher::Any, 200, body).await; + let error = provider(&server) + .fetch_live_window("test-admin-key", None, 60, at(NOW)) + .await + .unwrap_err(); + assert_eq!(error, LiveUsageError::ParseFailure, "{body}"); + } +} + +#[tokio::test] +async fn window_pagination_rejects_missing_repeated_and_excessive_cursors() { + let first_query = live_query_string(116_460, NOW, 60, None); + + let mut server = Server::new_async().await; + let _missing = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(first_query.clone()), + 200, + &page("", true, None), + ) + .await; + let error = provider(&server) + .fetch_live_window("test-admin-key", None, 60, at(NOW)) + .await + .unwrap_err(); + assert_eq!(error, LiveUsageError::ParseFailure); + + let mut server = Server::new_async().await; + let _first = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(first_query.clone()), + 200, + &page("", true, Some("same")), + ) + .await; + let repeated = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(format!("{first_query}&page=same")), + 200, + &page("", true, Some("same")), + ) + .await + .expect(1); + let error = provider(&server) + .fetch_live_window("test-admin-key", None, 60, at(NOW)) + .await + .unwrap_err(); + assert_eq!(error, LiveUsageError::ParseFailure); + repeated.assert_async().await; + + let mut server = Server::new_async().await; + let first = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(first_query.clone()), + 200, + &page("", true, Some("page-0")), + ) + .await + .expect(1); + let mut later = Vec::new(); + for index in 0..19 { + let next = format!("page-{}", index + 1); + later.push( + mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(format!("{first_query}&page=page-{index}")), + 200, + &page("", true, Some(&next)), + ) + .await + .expect(1), + ); + } + let beyond = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(format!("{first_query}&page=page-19")), + 200, + EMPTY_PAGE, + ) + .await + .expect(0); + let error = provider(&server) + .fetch_live_window("test-admin-key", None, 60, at(NOW)) + .await + .unwrap_err(); + assert_eq!(error, LiveUsageError::ParseFailure); + // Twenty requests in total: the first page plus nineteen cursors, never a 21st. + first.assert_async().await; + for mock in &later { + mock.assert_async().await; + } + beyond.assert_async().await; +} + +#[tokio::test] +async fn window_unsupported_ranges_never_request_the_network() { + let mut server = Server::new_async().await; + let never = mock(&mut server, COMPLETIONS_PATH, Matcher::Any, 200, EMPTY_PAGE) + .await + .expect(0); + let provider = provider(&server); + for window in [0, 30, 61] { + assert_eq!( + provider + .fetch_live_window("test-admin-key", None, window, at(NOW)) + .await + .unwrap_err(), + LiveUsageError::InvalidWindow + ); + assert_eq!( + provider + .fetch_live_usage(Some("test-admin-key"), None, window) + .await + .unwrap_err(), + LiveUsageError::InvalidWindow + ); + } + never.assert_async().await; +} + +#[tokio::test] +async fn window_unscoped_empty_windows_have_exact_ascending_grids() { + for minutes in [15_u32, 60, 360] { + let mut server = Server::new_async().await; + let first = NOW / 60 * 60 - (i64::from(minutes) - 1) * 60; + let unscoped = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(live_query_string(first, NOW, minutes, None)), + 200, + &page("", false, None), + ) + .await + .expect(1); + let live = provider(&server) + .fetch_live_window("test-admin-key", Some(" "), minutes, at(NOW)) + .await + .unwrap(); + unscoped.assert_async().await; + assert_eq!(live.minutes.len(), minutes as usize); + assert_eq!(live.project_id, None); + let summary = live.summary(LiveModelFilter::AllModels); + assert_eq!(summary.total_tokens, 0); + assert_eq!(summary.five_minute_average_tpm, 0.0); + assert!( + live.minutes + .windows(2) + .all(|pair| pair[1].start - pair[0].start == 60) + ); + assert_eq!( + live.series(LiveModelFilter::AllModels) + .iter() + .filter(|point| point.is_in_progress) + .count(), + 1 + ); + } +} + +#[tokio::test] +async fn window_batch_decoding_keeps_audio_inside_totals_and_defaults_to_standard() { + for batch in ["true", "false", "null", "missing"] { + let field = if batch == "missing" { + String::new() + } else { + format!(r#""batch":{batch},"#) + }; + let body = page( + &format!( + r#"{{"start_time":119940,"results":[{{{field}"model":"model-fixture", + "input_tokens":100,"input_cached_tokens":60,"output_tokens":20,"num_model_requests":3, + "input_audio_tokens":40,"output_audio_tokens":10}}]}}"# + ), + false, + None, + ); + let mut server = Server::new_async().await; + let _page = mock(&mut server, COMPLETIONS_PATH, Matcher::Any, 200, &body).await; + let live = provider(&server) + .fetch_live_window("test-admin-key", None, 15, at(NOW)) + .await + .unwrap(); + let point = live + .series(LiveModelFilter::AllModels) + .into_iter() + .find(|point| point.total_tokens > 0) + .unwrap(); + let is_batch = batch == "true"; + assert_eq!(point.total_tokens, 120, "{batch}"); + assert_eq!(point.batch_tokens, if is_batch { 120 } else { 0 }); + assert_eq!(point.batch_requests, if is_batch { 3 } else { 0 }); + assert_eq!(point.limit_counted_tokens(), if is_batch { 0 } else { 120 }); + assert_eq!(point.limit_counted_requests(), if is_batch { 0 } else { 3 }); + } +} + +// ---- Snapshot math --------------------------------------------------------------------- + +#[test] +fn reporting_frontier_skips_lagging_trailing_zeros_across_all_models() { + let live = frontier_snapshot(&[(8, "model-a", 80), (11, "model-b", 110)]); + assert_eq!(live.reporting_frontier(), Some(live.minutes[11].start)); + assert_eq!(live.summary(LiveModelFilter::AllModels).latest_tpm, 110); + let pending: Vec = live + .series(LiveModelFilter::AllModels) + .iter() + .filter(|point| point.is_pending) + .map(|point| point.start) + .collect(); + let tail: Vec = live.minutes[12..] + .iter() + .map(|minute| minute.start) + .collect(); + assert_eq!(pending, tail); + assert_eq!( + live.stacked_by_model(5) + .iter() + .filter(|stack| stack.is_pending) + .count(), + 6 + ); +} + +#[test] +fn peaks_ignore_partial_minutes_after_the_reporting_frontier() { + let live = frontier_snapshot(&[(8, "model-a", 80), (14, "model-a", 900)]); + assert_eq!(live.reporting_frontier(), Some(live.minutes[8].start)); + let summary = live.summary(LiveModelFilter::AllModels); + assert_eq!(summary.peak_tpm, 80); + assert_eq!(summary.peak_minute_start, Some(live.minutes[8].start)); + let limits = OrgLimitResolver::new(&[ModelRateLimit::new("model-a", 1000, Some(10))], None); + let utilization = live + .utilization(LiveModelFilter::Model("model-a"), &limits) + .unwrap(); + assert_eq!(utilization.peak_tpm_percent, 8.0); + assert_eq!(utilization.peak_rpm_percent, Some(10.0)); + let throughput = live.throughput(); + assert_eq!(throughput.peak_tokens_per_minute, 80); + assert_eq!(throughput.models[0].peak_limit_counted_tokens, 80); +} + +#[test] +fn model_ranking_ignores_pending_only_volume_for_the_top_series() { + let mut minutes: Vec = (0..15) + .map(|index| LiveMinute { + start: 119_160 + index * 60, + models: Vec::new(), + }) + .collect(); + minutes[8].models = (1..=6) + .map(|index| model(&format!("model-{index}"), 100 - index, 0, 0, 1)) + .collect(); + minutes[14].models = vec![model("pending-model", 10_000, 0, 0, 1)]; + let live = snapshot(minutes, None); + let names = live.model_names(); + assert_eq!( + names[..5], + ["model-1", "model-2", "model-3", "model-4", "model-5"] + ); + assert_eq!(names.last().map(String::as_str), Some("pending-model")); + let top: HashSet = live + .stacked_by_model(5) + .into_iter() + .map(|stack| stack.model_or_other) + .collect(); + assert!(top.contains("model-5")); + assert!(!top.contains("pending-model")); + assert_eq!(live.throughput().models[0].name, "model-1"); +} + +#[test] +fn reporting_frontier_resets_to_real_zeros_after_ten_idle_minutes() { + let live = frontier_snapshot(&[(3, "model-a", 100)]); + assert_eq!(live.reporting_frontier(), Some(live.minutes[13].start)); + let summary = live.summary(LiveModelFilter::AllModels); + assert_eq!(summary.latest_tpm, 0); + assert_eq!(summary.five_minute_average_tpm, 0.0); + assert_eq!( + live.series(LiveModelFilter::AllModels) + .iter() + .filter(|point| point.is_pending) + .count(), + 1 + ); + // The oldest of the ten complete minutes is still eligible. + let boundary = frontier_snapshot(&[(4, "model-a", 100)]); + assert_eq!( + boundary.reporting_frontier(), + Some(boundary.minutes[4].start) + ); +} + +#[test] +fn frontier_is_shared_across_filters_so_an_idle_model_reports_zero() { + let live = frontier_snapshot(&[(8, "model-a", 80), (11, "model-b", 110)]); + let idle = live.summary(LiveModelFilter::Model("model-a")); + assert_eq!(idle.latest_tpm, 0); + assert_eq!(idle.latest_minute_start, Some(live.minutes[11].start)); + assert_eq!( + live.summary(LiveModelFilter::Model("absent")) + .latest_minute_start, + idle.latest_minute_start + ); +} + +#[test] +fn five_minute_average_ends_at_the_reporting_frontier_including_idle_minutes() { + let live = frontier_snapshot(&[ + (6, "model-a", 999), + (7, "model-a", 50), + (8, "model-a", 100), + (11, "model-b", 150), + ]); + assert_eq!( + live.summary(LiveModelFilter::AllModels) + .five_minute_average_tpm, + 60.0 + ); + assert_eq!( + live.summary(LiveModelFilter::Model("model-a")) + .five_minute_average_tpm, + 30.0 + ); + assert_eq!( + live.summary(LiveModelFilter::Model("model-b")) + .five_minute_average_tpm, + 30.0 + ); + let limits = OrgLimitResolver::new(&[ModelRateLimit::new("model-b", 300, None)], None); + assert_eq!( + live.utilization(LiveModelFilter::Model("model-b"), &limits) + .unwrap() + .latest_tpm_percent, + 50.0 + ); +} + +fn batch_row(input: u64, cached: u64, output: u64, requests: u64, batch: bool) -> LiveModelUsage { + LiveModelUsage { + batch_tokens: if batch { input + output } else { 0 }, + batch_requests: if batch { requests } else { 0 }, + ..model("model-fixture", input, cached, output, requests) + } +} + +#[test] +fn batch_traffic_stays_in_totals_while_utilization_uses_standard_peaks() { + let live = build( + vec![ + LiveMinute { + start: 119_880, + models: vec![ + batch_row(900, 800, 100, 90, true), + batch_row(80, 70, 20, 2, false), + ], + }, + LiveMinute { + start: 119_940, + models: vec![ + batch_row(40, 30, 10, 5, true), + batch_row(160, 150, 40, 4, false), + model("model-second", 90, 0, 0, 1), + ], + }, + ], + 15, + NOW, + None, + ) + .unwrap(); + let points = live.series(LiveModelFilter::Model("model-fixture")); + let tail = &points[points.len() - 3..]; + let field = |get: fn(&super::live::LivePoint) -> u64| tail.iter().map(get).collect::>(); + assert_eq!(field(|point| point.total_tokens), [1100, 250, 0]); + assert_eq!(field(|point| point.batch_tokens), [1000, 50, 0]); + assert_eq!(field(|point| point.batch_requests), [90, 5, 0]); + assert_eq!(field(|point| point.limit_counted_tokens()), [100, 200, 0]); + assert_eq!(field(|point| point.limit_counted_requests()), [2, 4, 0]); + assert_eq!(live.minutes[12].models.len(), 1); + assert_eq!(live.model_names(), ["model-fixture", "model-second"]); + assert_eq!(live.reporting_frontier(), Some(119_940)); + let summary = live.summary(LiveModelFilter::Model("model-fixture")); + assert_eq!(summary.total_tokens, 1350); + assert_eq!(summary.batch_tokens, 1050); + assert_eq!(summary.input_tokens, 1180); + assert_eq!(summary.cached_input_tokens, 1050); + assert_eq!(summary.output_tokens, 170); + assert_eq!(summary.requests, 101); + assert_eq!(summary.latest_tpm, 250); + assert_eq!(summary.peak_tpm, 1100); + assert_eq!(summary.peak_minute_start, Some(119_880)); + let all = live.summary(LiveModelFilter::AllModels); + assert_eq!((all.total_tokens, all.batch_tokens), (1440, 1050)); + assert_eq!( + live.stacked_by_model(5) + .iter() + .map(|stack| stack.tokens) + .sum::(), + 1440 + ); + assert_eq!(live.model_rows()[0].total_tokens, 1350); + assert_eq!(live.model_rows()[0].requests, 101); + + let limits = OrgLimitResolver::new( + &[ + ModelRateLimit::new("model-fixture", 400, Some(8)), + ModelRateLimit::new("model-second", 100, None), + ], + None, + ); + let utilization = live + .utilization(LiveModelFilter::Model("model-fixture"), &limits) + .unwrap(); + assert_eq!(utilization.latest_tpm_percent, 50.0); + assert_eq!(utilization.peak_tpm_percent, 50.0); + assert_eq!(utilization.peak_rpm_percent, Some(50.0)); + let rows = live.model_rows_with_limits(&limits); + assert_eq!(rows[0].peak_tpm_percent, Some(50.0)); + assert_eq!(rows[0].latest_tpm_percent, Some(50.0)); + let highest = live.highest_utilization(&limits).unwrap(); + assert_eq!(highest.name, "model-second"); + assert_eq!(highest.peak_tpm_percent, Some(90.0)); + + let throughput = live.throughput(); + let tail = &throughput.minutes[throughput.minutes.len() - 3..]; + assert_eq!( + tail.iter() + .map(|minute| minute.total_tokens) + .collect::>(), + [1100, 340, 0] + ); + assert_eq!( + tail.iter() + .map(|minute| minute.requests) + .collect::>(), + [92, 10, 0] + ); + assert_eq!(throughput.models[0].peak_tpm, 1100); + assert_eq!(throughput.models[0].peak_limit_counted_tokens, 200); +} + +#[test] +fn builder_rejects_invalid_batch_subsets_and_bad_windows() { + for (tokens, requests) in [(121, 0), (0, 4)] { + let row = LiveModelUsage { + batch_tokens: tokens, + batch_requests: requests, + ..model("model-fixture", 100, 60, 20, 3) + }; + let error = build( + vec![LiveMinute { + start: NOW - 85, + models: vec![row], + }], + 15, + NOW, + None, + ) + .unwrap_err(); + assert_eq!(error, LiveUsageError::ParseFailure); + } + assert!(build(Vec::new(), 30, NOW, None).is_err()); + assert!(build(Vec::new(), 15, -1, None).is_err()); + let cached_above_input = LiveMinute { + start: NOW - 85, + models: vec![model("m", 1, 2, 0, 0)], + }; + assert!(build(vec![cached_above_input], 15, NOW, None).is_err()); +} + +#[test] +fn compact_throughput_keeps_every_model_with_aligned_limit_traffic_including_pending() { + let models: Vec = (1..=7_u64) + .map(|index| { + model( + &format!("model-{index}"), + index * 10_000_000, + index * 9_500_000, + 1_000_000, + index, + ) + }) + .collect(); + let live = build( + vec![ + LiveMinute { + start: NOW - 85, + models, + }, + LiveMinute { + start: NOW - 25, + models: vec![LiveModelUsage { + batch_tokens: 50, + ..model("model-1", 200, 0, 0, 1) + }], + }, + ], + 60, + NOW, + None, + ) + .unwrap(); + let throughput = live.throughput(); + let names: Vec<&str> = throughput.models.iter().map(|m| m.name.as_str()).collect(); + assert_eq!( + names, + [ + "model-7", "model-6", "model-5", "model-4", "model-3", "model-2", "model-1" + ] + ); + for compact in &throughput.models { + let expected: Vec = live + .series(LiveModelFilter::Model(&compact.name)) + .iter() + .map(|point| point.limit_counted_tokens()) + .collect(); + assert_eq!(compact.limit_counted_tokens_by_minute, expected); + assert_eq!( + compact.limit_counted_tokens_by_minute.len(), + throughput.minutes.len() + ); + } + assert!(throughput.minutes.last().unwrap().is_pending); + assert_eq!( + throughput + .models + .last() + .unwrap() + .limit_counted_tokens_by_minute + .last(), + Some(&150) + ); + + let manual = [ModelRateLimit::new("model-7", 40_000_000, None)]; + let limits = OrgLimitResolver::new(&manual, None); + let utilization = live + .utilization(LiveModelFilter::Model("model-7"), &limits) + .unwrap(); + assert_eq!( + live.summary(LiveModelFilter::Model("model-7")).peak_tpm, + 71_000_000 + ); + // Cached input counts toward TPM: 71M of a 40M limit. + assert_eq!(utilization.peak_tpm_percent, 177.5); + assert_eq!(utilization.latest_tpm_percent, 177.5); + assert_eq!( + live.highest_utilization(&limits).unwrap().peak_tpm_percent, + Some(177.5) + ); + assert_eq!( + live.model_rows_with_limits(&limits)[0].peak_tpm_percent, + Some(177.5) + ); + + let tile = throughput_limit_utilization(&throughput, &manual).unwrap(); + assert_eq!(tile.label, "model-7"); + assert_eq!(tile.peak_tokens, 71_000_000); + assert_eq!(tile.tokens_per_minute_limit, 40_000_000); + assert_eq!(tile.percent, 177.5); +} + +#[test] +fn throughput_json_round_trips_through_serde() { + let throughput = frontier_snapshot(&[(8, "model-a", 80)]).throughput(); + let encoded = serde_json::to_string(&throughput).unwrap(); + let decoded: crate::core::ProviderThroughput = serde_json::from_str(&encoded).unwrap(); + assert_eq!(decoded, throughput); +} + +// ---- Per-refresh Tokens/min ------------------------------------------------------------ + +/// Upstream `livePage`, shifted onto the refresh clock. +fn refresh_live_page() -> String { + format!( + r#"{{"data":[ + {{"start_time":{a},"end_time":{b},"results":[{{ + "model":"model-fixture","input_tokens":96000000,"input_cached_tokens":92400000, + "output_tokens":1000000,"num_model_requests":50}}]}}, + {{"start_time":{b},"end_time":{c},"results":[{{ + "model":"model-fixture","input_tokens":74000000,"input_cached_tokens":70400000, + "output_tokens":1000000,"num_model_requests":50}}]}} + ],"has_more":false}}"#, + a = BASE - 240, + b = BASE - 180, + c = BASE - 120, + ) +} + +/// One refresh bucket at `BASE - 180` with the given results. +fn refresh_bucket(results: &str) -> String { + format!( + r#"{{"data":[{{"start_time":{start},"end_time":{end},"results":[{results}]}}],"has_more":false}}"#, + start = BASE - 180, + end = BASE - 120, + ) +} + +struct RefreshServer { + server: ServerGuard, + live: mockito::Mock, +} + +async fn refresh_server(live_status: usize, live_body: &str) -> RefreshServer { + let mut server = Server::new_async().await; + let _costs = mock(&mut server, COSTS_PATH, Matcher::Any, 200, EMPTY_PAGE).await; + let _daily = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::UrlEncoded("bucket_width".into(), "1d".into()), + 200, + EMPTY_PAGE, + ) + .await; + let live = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::UrlEncoded("bucket_width".into(), "1m".into()), + live_status, + live_body, + ) + .await; + RefreshServer { server, live } +} + +async fn refresh(server: &RefreshServer, project: Option<&str>, live: bool) -> ProviderFetchResult { + provider(&server.server) + .fetch_refresh(&admin(), project, at(REFRESH_NOW), live) + .await + .expect("the daily admin usage should succeed") +} + +#[tokio::test] +async fn refresh_requests_one_scoped_minute_page_and_maps_lagging_throughput() { + let mut server = Server::new_async().await; + let _costs = mock(&mut server, COSTS_PATH, Matcher::Any, 200, EMPTY_PAGE).await; + let _daily = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::UrlEncoded("bucket_width".into(), "1d".into()), + 200, + EMPTY_PAGE, + ) + .await; + let live = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::Exact(live_query_string( + REFRESH_START, + REFRESH_NOW, + 60, + Some("proj_fixture & scope"), + )), + 200, + &refresh_live_page(), + ) + .await + .expect(1); + + let result = provider(&server) + .fetch_refresh( + &admin(), + Some("proj_fixture & scope"), + at(REFRESH_NOW), + true, + ) + .await + .unwrap(); + + live.assert_async().await; + assert!(result.open_ai_api_usage.is_some()); + let throughput = result.throughput.expect("live minutes attach on refresh"); + assert_eq!(throughput.minutes.len(), 60); + assert_eq!(throughput.window_minutes, 60); + assert_eq!(throughput.fetched_at, REFRESH_NOW); + assert_eq!(throughput.frontier, Some(BASE - 180)); + assert_eq!(throughput.latest_tokens_per_minute, Some(75_000_000)); + assert_eq!(throughput.peak_tokens_per_minute, 97_000_000); + assert_eq!(throughput.peak_at, Some(BASE - 240)); + assert_eq!( + throughput.five_minute_average_tokens_per_minute, + 34_400_000.0 + ); + assert!( + throughput.minutes[57..] + .iter() + .all(|minute| minute.is_pending) + ); + assert_eq!(throughput.models[0].peak_limit_counted_tokens, 97_000_000); + assert_eq!( + throughput.project_id.as_deref(), + Some("proj_fixture & scope") + ); + // Project-scoped usage never reports organization limit utilization. + assert!( + throughput_limit_utilization( + &throughput, + &[ModelRateLimit::new("model-fixture", 1, None)] + ) + .is_none() + ); + + // The refresh and the on-demand 60-minute window agree on the same buckets. + let window = provider(&server) + .fetch_live_window( + "test-admin-key", + Some("proj_fixture & scope"), + 60, + at(REFRESH_NOW), + ) + .await + .unwrap(); + assert_eq!(window.throughput(), throughput); +} + +#[tokio::test] +async fn refresh_live_failures_keep_daily_data_and_the_disabled_setting_skips_the_request() { + for (status, body) in [ + (500, "{}"), + (429, "{}"), + (401, "{}"), + (200, "malformed"), + (200, r#"{"data":false}"#), + ] { + let server = refresh_server(status, body).await; + let result = refresh(&server, None, true).await; + assert!(result.throughput.is_none(), "{status} {body}"); + assert!(result.open_ai_api_usage.is_some()); + assert_eq!(result.source_label, "admin-api"); + } + + let mut server = refresh_server(200, &refresh_live_page()).await; + server.live = server.live.expect(0); + let result = refresh(&server, None, false).await; + assert!(result.throughput.is_none()); + assert!(result.open_ai_api_usage.is_some()); + server.live.assert_async().await; +} + +#[tokio::test] +async fn refresh_balance_fallback_never_requests_live_minutes() { + let mut server = Server::new_async().await; + let _costs = mock(&mut server, COSTS_PATH, Matcher::Any, 403, "{}").await; + let _grants = server + .mock("GET", "/v1/dashboard/billing/credit_grants") + .with_body(r#"{"total_granted":100.0,"total_used":25.0,"total_available":75.0}"#) + .create_async() + .await; + let live = mock( + &mut server, + COMPLETIONS_PATH, + Matcher::UrlEncoded("bucket_width".into(), "1m".into()), + 200, + &refresh_live_page(), + ) + .await + .expect(0); + + let result = provider(&server) + .fetch_refresh(&admin(), None, at(REFRESH_NOW), true) + .await + .unwrap(); + + assert_eq!(result.source_label, "billing-api"); + assert!(result.throughput.is_none()); + live.assert_async().await; +} + +#[tokio::test] +async fn refresh_throughput_does_not_count_audio_breakdowns_twice() { + let server = refresh_server( + 200, + &refresh_bucket( + r#"{"model":"model-fixture","input_tokens":100,"input_audio_tokens":40,"input_cached_tokens":60, + "output_tokens":15,"output_audio_tokens":5,"num_model_requests":2}"#, + ), + ) + .await; + let throughput = refresh(&server, None, true).await.throughput.unwrap(); + assert_eq!(throughput.peak_tokens_per_minute, 115); + assert_eq!(throughput.models[0].peak_limit_counted_tokens, 115); +} + +#[tokio::test] +async fn refresh_batch_rows_merge_by_trimmed_model_name() { + let server = refresh_server( + 200, + &refresh_bucket( + r#"{"model":" model-fixture ","batch":true,"input_tokens":100,"input_cached_tokens":60, + "output_tokens":20,"num_model_requests":3,"input_audio_tokens":40,"output_audio_tokens":10}, + {"model":"model-fixture","batch":false,"input_tokens":200,"input_cached_tokens":150, + "output_tokens":30,"num_model_requests":2}, + {"model":"model-second","batch":null,"input_tokens":40,"num_model_requests":1}"#, + ), + ) + .await; + let throughput = refresh(&server, None, true).await.throughput.unwrap(); + assert_eq!(throughput.models[0].name, "model-fixture"); + assert_eq!(throughput.models[0].peak_tpm, 350); + assert_eq!(throughput.models[0].peak_limit_counted_tokens, 230); + assert_eq!(throughput.peak_tokens_per_minute, 390); +} + +#[tokio::test] +async fn refresh_invalid_live_payloads_discard_only_throughput() { + let bucket = |start: i64, end: &str, results: &str| { + format!(r#"{{"start_time":{start}{end},"results":[{results}]}}"#) + }; + let row = r#"{"model":"m","input_tokens":10}"#; + let aligned = BASE - 180; + let too_many_buckets = (0..61) + .map(|_| bucket(aligned, &format!(r#","end_time":{}"#, aligned + 60), row)) + .collect::>() + .join(","); + let too_many_results = vec![row; 401].join(","); + let mut bodies = vec![ + format!(r#"{{"data":[{too_many_buckets}],"has_more":false}}"#), + format!( + r#"{{"data":[{}]}}"#, + bucket(aligned, &format!(r#","end_time":{}"#, aligned + 60), row) + ), + format!( + r#"{{"data":[{}],"has_more":false}}"#, + bucket( + aligned + 1, + &format!(r#","end_time":{}"#, aligned + 60), + row + ) + ), + format!( + r#"{{"data":[{}],"has_more":false}}"#, + bucket( + REFRESH_START - 60, + &format!(r#","end_time":{REFRESH_START}"#), + row + ) + ), + format!( + r#"{{"data":[{}],"has_more":false}}"#, + bucket(BASE + 60, &format!(r#","end_time":{}"#, BASE + 120), row) + ), + format!( + r#"{{"data":[{}],"has_more":false}}"#, + bucket(aligned, "", row) + ), + format!( + r#"{{"data":[{}],"has_more":false}}"#, + bucket(aligned, &format!(r#","end_time":{aligned}"#), row) + ), + format!( + r#"{{"data":[{}],"has_more":false}}"#, + bucket( + aligned, + &format!(r#","end_time":{}"#, aligned + 60), + &too_many_results + ) + ), + refresh_bucket(r#"{"model":"m","input_tokens":1,"input_cached_tokens":2}"#), + refresh_bucket(r#"{"model":"m","input_tokens":900719925475}"#), + ]; + for batch in [r#""yes""#, "1", "[]", "{}"] { + bodies.push(refresh_bucket(&format!( + r#"{{"model":"m","batch":{batch},"input_tokens":10}}"# + ))); + } + for body in bodies { + let server = refresh_server(200, &body).await; + let result = refresh(&server, None, true).await; + assert!( + result.throughput.is_none(), + "{}", + &body[..body.len().min(160)] + ); + assert!(result.open_ai_api_usage.is_some()); + } +} + +#[tokio::test] +async fn refresh_model_bounds_count_merged_entries() { + for count in [200, 201] { + let rows: Vec = (0..count) + .flat_map(|index| { + [true, false].map(|batch| { + format!( + r#"{{"model":"model-fixture-{index}","batch":{batch},"input_tokens":10}}"# + ) + }) + }) + .collect(); + // 201 unique models fit the raw cap but exceed the merged cap. + let raw: Vec = if count == 200 { + rows + } else { + rows.into_iter().step_by(2).collect() + }; + let server = refresh_server(200, &refresh_bucket(&raw.join(","))).await; + let result = refresh(&server, None, true).await; + if count == 200 { + assert_eq!(result.throughput.unwrap().peak_tokens_per_minute, 4000); + } else { + assert!(result.throughput.is_none()); + } + assert!(result.open_ai_api_usage.is_some()); + } +} + +// ---- Organization rate limits ---------------------------------------------------------- + +fn list_page(rows: &str, more: bool, last_id: Option<&str>) -> String { + let last = last_id.map_or_else(|| "null".to_string(), |id| format!("\"{id}\"")); + format!(r#"{{"object":"list","data":[{rows}],"has_more":{more},"last_id":{last}}}"#) +} + +fn limit_row(model: &str, tpm: u64, rpm: u64) -> String { + format!( + r#"{{"model":"{model}","max_tokens_per_1_minute":{tpm},"max_requests_per_1_minute":{rpm}}}"# + ) +} + +fn rate_limits_path(project: &str) -> String { + format!("{PROJECTS_PATH}/{project}/rate_limits") +} + +#[tokio::test] +async fn org_discovery_paginates_both_endpoints_and_takes_the_max_cap() { + let mut server = Server::new_async().await; + let mut mocks = vec![ + mock( + &mut server, + PROJECTS_PATH, + Matcher::Exact("limit=100".into()), + 200, + &list_page( + r#"{"id":"proj_a","status":"active"},{"id":"proj_archived","status":"archived"}"#, + true, + Some("proj_page & 2"), + ), + ) + .await, + mock( + &mut server, + PROJECTS_PATH, + Matcher::Exact("limit=100&after=proj_page+%26+2".into()), + 200, + &list_page(r#"{"id":"proj_b","status":"active"}"#, false, None), + ) + .await, + mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Exact("limit=100".into()), + 200, + &list_page(&limit_row("gpt-5", 100, 20), true, Some("rl_next")), + ) + .await, + mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Exact("limit=100&after=rl_next".into()), + 200, + &list_page(&limit_row("gpt-4o-mini", 50, 5), false, None), + ) + .await, + mock( + &mut server, + &rate_limits_path("proj_b"), + Matcher::Exact("limit=100".into()), + 200, + &list_page(&limit_row("gpt-5", 200, 10), false, None), + ) + .await, + ]; + let never = mock( + &mut server, + &rate_limits_path("proj_archived"), + Matcher::Any, + 200, + &list_page("", false, None), + ) + .await + .expect(0); + + let limits = provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap(); + + for mock in mocks.drain(..) { + mock.expect(1).assert_async().await; + } + never.assert_async().await; + assert_eq!(limits.projects_sampled, 2); + assert!(!limits.truncated); + assert_eq!(limits.fetched_at, NOW); + assert_eq!( + limits.by_model.get("gpt-5"), + Some(&OrgRateLimit { + tokens_per_minute: 200, + requests_per_minute: Some(20) + }) + ); + assert_eq!( + limits.by_model.get("gpt-4o-mini"), + Some(&OrgRateLimit { + tokens_per_minute: 50, + requests_per_minute: Some(5) + }) + ); +} + +#[tokio::test] +async fn org_projects_without_a_status_are_active() { + let mut server = Server::new_async().await; + let _projects = mock( + &mut server, + PROJECTS_PATH, + Matcher::Any, + 200, + &list_page( + r#"{"id":"proj_a"},{"id":"proj_b","status":null},{"id":"proj_c","status":"archived"}"#, + false, + None, + ), + ) + .await; + let _a = mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Any, + 200, + &list_page(&limit_row("gpt-5", 100, 2), false, None), + ) + .await; + let _b = mock( + &mut server, + &rate_limits_path("proj_b"), + Matcher::Any, + 200, + &list_page(&limit_row("gpt-5", 300, 4), false, None), + ) + .await; + + let limits = provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap(); + assert_eq!(limits.projects_sampled, 2); + assert_eq!( + limits.by_model.get("gpt-5"), + Some(&OrgRateLimit { + tokens_per_minute: 300, + requests_per_minute: Some(4) + }) + ); +} + +#[tokio::test] +async fn org_project_sample_truncates_after_twenty_five_active_projects() { + for (total, truncated) in [(30, true), (25, false)] { + let mut server = Server::new_async().await; + let projects = (0..total) + .map(|index| format!(r#"{{"id":"proj_{index}","status":"active"}}"#)) + .collect::>() + .join(","); + let _projects = mock( + &mut server, + PROJECTS_PATH, + Matcher::Exact("limit=100".into()), + 200, + &list_page(&projects, false, None), + ) + .await; + let rate_limits = server + .mock( + "GET", + Matcher::Regex(r"^/v1/organization/projects/proj_\d+/rate_limits$".into()), + ) + .match_query(Matcher::Any) + .with_body(list_page("", false, None)) + .expect(25) + .create_async() + .await; + + let limits = provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap(); + rate_limits.assert_async().await; + assert_eq!(limits.projects_sampled, 25); + assert_eq!(limits.truncated, truncated, "{total}"); + } +} + +#[tokio::test] +async fn org_malformed_project_responses_fail_parsing() { + for body in [ + "not json", + "{}", + r#"{"data":[],"has_more":true}"#, + r#"{"data":[{"id":"proj_a"}],"has_more":false}"#, + ] { + let mut server = Server::new_async().await; + let _projects = mock(&mut server, PROJECTS_PATH, Matcher::Any, 200, body).await; + let _limits = mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Any, + 200, + "{}", + ) + .await; + let error = provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap_err(); + assert_eq!(error, OrgRateLimitsError::ParseFailure, "{body}"); + } +} + +async fn one_project_with_rows(rows: &str) -> Result { + let mut server = Server::new_async().await; + let _projects = mock( + &mut server, + PROJECTS_PATH, + Matcher::Any, + 200, + &list_page(r#"{"id":"proj_a","status":"active"}"#, false, None), + ) + .await; + let _limits = mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Any, + 200, + &list_page(rows, false, None), + ) + .await; + provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await +} + +#[tokio::test] +async fn org_rows_without_token_or_request_limits_are_tolerated() { + let limits = one_project_with_rows( + r#"{"model":"gpt-5","max_tokens_per_1_minute":100}, + {"model":"gpt-5","max_tokens_per_1_minute":80,"max_requests_per_1_minute":3}, + {"model":"sora-2","max_requests_per_1_minute":5}, + {"model":"gpt-4o","max_tokens_per_1_minute":0,"max_requests_per_1_minute":0}"#, + ) + .await + .unwrap(); + assert_eq!(limits.by_model.len(), 1); + assert_eq!( + limits.by_model.get("gpt-5"), + Some(&OrgRateLimit { + tokens_per_minute: 100, + requests_per_minute: Some(3) + }) + ); +} + +#[tokio::test] +async fn org_required_rate_limit_fields_are_validated() { + for row in [ + r#"{"max_tokens_per_1_minute":100,"max_requests_per_1_minute":2}"#, + r#"{"model":"","max_tokens_per_1_minute":100,"max_requests_per_1_minute":2}"#, + r#"{"model":"gpt-5","max_tokens_per_1_minute":-1,"max_requests_per_1_minute":2}"#, + r#"{"model":"gpt-5","max_tokens_per_1_minute":1.5,"max_requests_per_1_minute":2}"#, + r#"{"model":"gpt-5","max_tokens_per_1_minute":100,"max_requests_per_1_minute":-2}"#, + ] { + assert_eq!( + one_project_with_rows(row).await.unwrap_err(), + OrgRateLimitsError::ParseFailure, + "{row}" + ); + } +} + +#[tokio::test] +async fn org_repeated_and_excessive_cursors_are_rejected() { + let mut server = Server::new_async().await; + let _first = mock( + &mut server, + PROJECTS_PATH, + Matcher::Exact("limit=100".into()), + 200, + &list_page("", true, Some("same")), + ) + .await; + let _again = mock( + &mut server, + PROJECTS_PATH, + Matcher::Exact("limit=100&after=same".into()), + 200, + &list_page("", true, Some("same")), + ) + .await; + assert_eq!( + provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap_err(), + OrgRateLimitsError::ParseFailure + ); + + let mut server = Server::new_async().await; + let _projects = mock( + &mut server, + PROJECTS_PATH, + Matcher::Any, + 200, + &list_page(r#"{"id":"proj_a","status":"active"}"#, false, None), + ) + .await; + let _first = mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Exact("limit=100".into()), + 200, + &list_page("", true, Some("page-0")), + ) + .await; + let mut later = Vec::new(); + for index in 0..19 { + let next = format!("page-{}", index + 1); + later.push( + mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Exact(format!("limit=100&after=page-{index}")), + 200, + &list_page("", true, Some(&next)), + ) + .await + .expect(1), + ); + } + let beyond = mock( + &mut server, + &rate_limits_path("proj_a"), + Matcher::Exact("limit=100&after=page-19".into()), + 200, + &list_page("", false, None), + ) + .await + .expect(0); + assert_eq!( + provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap_err(), + OrgRateLimitsError::ParseFailure + ); + for mock in &later { + mock.assert_async().await; + } + beyond.assert_async().await; +} + +#[tokio::test] +async fn org_http_failures_have_typed_errors() { + for (status, expected) in [ + (401, OrgRateLimitsError::AdminKeyRequired), + (403, OrgRateLimitsError::AdminKeyRequired), + (429, OrgRateLimitsError::RateLimited), + (500, OrgRateLimitsError::ApiFailure(500)), + ] { + let mut server = Server::new_async().await; + let _projects = mock(&mut server, PROJECTS_PATH, Matcher::Any, status, "{}").await; + assert_eq!( + provider(&server) + .fetch_org_rate_limits_with("test-admin-key", at(NOW)) + .await + .unwrap_err(), + expected, + "{status}" + ); + } +} + +#[test] +fn org_manual_limits_and_project_caps_match_by_exact_then_dated_snapshot_family() { + let cap = |tpm| OrgRateLimit { + tokens_per_minute: tpm, + requests_per_minute: None, + }; + let project_caps = OrgRateLimits { + by_model: [ + ("gpt-5-2025-08-07".to_string(), cap(1000)), + ("gpt-4o".to_string(), cap(200)), + ("gpt-4o-mini".to_string(), cap(300)), + ] + .into_iter() + .collect(), + projects_sampled: 2, + truncated: false, + fetched_at: NOW, + }; + let resolver = OrgLimitResolver::new( + &[ModelRateLimit::new("gpt-5", 100, None)], + Some(project_caps), + ); + assert_eq!(resolver.project_cap("gpt-5-2025-08-07"), Some(cap(1000))); + assert_eq!(resolver.resolve("gpt-5-2025-08-07"), Some(cap(100))); + assert_eq!(resolver.manual_limit_key("gpt-5-2025-08-07"), Some("gpt-5")); + assert_eq!(resolver.manual_limit_key("gpt-4o"), None); + assert_eq!(resolver.resolve("gpt-4o"), None); + assert_eq!(resolver.project_cap("gpt-4o"), Some(cap(200))); + assert_eq!( + resolver.project_cap("gpt-4o-mini-2024-07-18"), + Some(cap(300)) + ); + assert_eq!(resolver.resolve("absent"), None); + + let manual = OrgLimitResolver::new( + &[ + ModelRateLimit::new("gpt-4", 1, None), + ModelRateLimit::new("gpt-4o", 2, None), + ModelRateLimit::new("gpt-4o-mini", 3, None), + ModelRateLimit::new("gpt-4o-mini-2024-07-18", 4, None), + ], + None, + ); + assert_eq!(manual.resolve("gpt-4o-mini-2024-07-18"), Some(cap(4))); + assert_eq!( + manual.manual_limit_key("gpt-4o-mini-2024-07-18"), + Some("gpt-4o-mini-2024-07-18") + ); + assert_eq!( + manual.manual_limit_key("gpt-4o-mini-2024-08-01"), + Some("gpt-4o-mini") + ); + assert_eq!(manual.resolve("gpt-4o-mini-next"), None); + assert_eq!(manual.resolve("gpt-4-0613"), Some(cap(1))); + let gpt4 = OrgLimitResolver::new(&[ModelRateLimit::new("gpt-4", 1, None)], None); + assert_eq!(gpt4.resolve("gpt-4o-mini"), None); + // A family limit covers dated snapshots only, never distinct sibling models. + let gpt5 = OrgLimitResolver::new(&[ModelRateLimit::new("gpt-5", 100, None)], None); + assert_eq!(gpt5.resolve("gpt-5-mini"), None); + assert_eq!(gpt5.resolve("gpt-5-mini-2025-08-07"), None); + assert_eq!(gpt5.resolve("gpt-5-2025-08-07"), Some(cap(100))); +} + +#[test] +fn org_utilization_uses_complete_minute_peaks_and_rpm_with_a_deterministic_highest_model() { + let live = snapshot( + vec![ + LiveMinute { + start: NOW - 90, + models: vec![model("b", 80, 0, 20, 4), model("a", 50, 0, 0, 1)], + }, + LiveMinute { + start: NOW - 30, + models: vec![ + model("b", 150, 0, 0, 2), + model("a", 75, 0, 0, 1), + model("no-limit", 1000, 0, 0, 1), + ], + }, + ], + None, + ); + let limits = OrgLimitResolver::new( + &[ + ModelRateLimit::new("a", 100, None), + ModelRateLimit::new("b", 200, Some(8)), + ], + None, + ); + let b = live + .utilization(LiveModelFilter::Model("b"), &limits) + .unwrap(); + assert_eq!(b.latest_tpm_percent, 50.0); + // The in-progress minute (150 TPM, 75%) is pending, so the peak is the complete minute. + assert_eq!(b.peak_tpm_percent, 50.0); + assert_eq!(b.peak_rpm_percent, Some(50.0)); + assert!( + live.utilization(LiveModelFilter::AllModels, &limits) + .is_none() + ); + assert!( + live.utilization(LiveModelFilter::Model("no-limit"), &limits) + .is_none() + ); + assert_eq!(live.highest_utilization(&limits).unwrap().name, "a"); + let rows = live.model_rows_with_limits(&limits); + let row = rows.iter().find(|row| row.name == "b").unwrap(); + assert_eq!(row.tpm_limit, Some(200)); + assert_eq!(row.rpm_limit, Some(8)); + assert_eq!(row.project_cap, None); + assert_eq!(row.peak_tpm_percent, Some(50.0)); + assert_eq!(row.latest_tpm_percent, Some(50.0)); + assert_eq!( + live.utilization(LiveModelFilter::Model("a"), &limits) + .unwrap() + .peak_rpm_percent, + None + ); +} + +#[test] +fn org_shared_manual_limits_pool_minute_traffic_while_exact_limits_stay_independent() { + let minutes = vec![ + LiveMinute { + start: NOW - 150, + models: vec![ + LiveModelUsage { + batch_tokens: 40, + batch_requests: 4, + ..model("gpt-5", 90, 50, 10, 10) + }, + model("gpt-5-2025-08-07", 60, 0, 0, 6), + model("gpt-5-2025-08-08", 70, 0, 0, 7), + ], + }, + LiveMinute { + start: NOW - 90, + models: vec![ + model("gpt-5", 20, 0, 0, 1), + model("gpt-5-2025-08-07", 30, 0, 0, 2), + ], + }, + LiveMinute { + start: NOW - 30, + models: vec![ + model("gpt-5", 1000, 0, 0, 100), + model("gpt-5-2025-08-07", 1000, 0, 0, 100), + ], + }, + ]; + let manual = [ + ModelRateLimit::new("gpt-5", 100, Some(10)), + ModelRateLimit::new("gpt-5-2025-08-08", 200, Some(20)), + ]; + let limits = OrgLimitResolver::new(&manual, None); + let live = snapshot(minutes.clone(), None); + for name in ["gpt-5", "gpt-5-2025-08-07"] { + let utilization = live + .utilization(LiveModelFilter::Model(name), &limits) + .unwrap(); + assert_eq!(utilization.pooled_models, ["gpt-5", "gpt-5-2025-08-07"]); + assert_eq!(utilization.latest_tpm_percent, 50.0); + assert_eq!(utilization.peak_tpm_percent, 120.0); + assert_eq!(utilization.peak_rpm_percent, Some(120.0)); + } + let exact = live + .utilization(LiveModelFilter::Model("gpt-5-2025-08-08"), &limits) + .unwrap(); + assert_eq!(exact.pooled_models, ["gpt-5-2025-08-08"]); + assert_eq!(exact.peak_tpm_percent, 35.0); + assert_eq!(exact.latest_tpm_percent, 0.0); + assert_eq!(exact.peak_rpm_percent, Some(35.0)); + assert_eq!( + live.highest_utilization(&limits).unwrap().peak_tpm_percent, + Some(120.0) + ); + let rows = live.model_rows_with_limits(&limits); + assert_eq!( + rows[..2] + .iter() + .map(|row| row.peak_tpm_percent) + .collect::>(), + [Some(120.0), Some(120.0)] + ); + + // The menu tile pools the alias with its snapshot under the configured ID. + let tile = throughput_limit_utilization(&live.throughput(), &manual).unwrap(); + assert_eq!(tile.label, "gpt-5"); + assert_eq!(tile.peak_tokens, 120); + assert_eq!(tile.tokens_per_minute_limit, 100); + assert_eq!(tile.percent, 120.0); + assert!(throughput_limit_utilization(&live.throughput(), &[]).is_none()); + + let scoped = snapshot(minutes, Some("proj_fixture")); + assert!( + scoped + .utilization(LiveModelFilter::Model("gpt-5"), &limits) + .is_none() + ); + assert!(throughput_limit_utilization(&scoped.throughput(), &manual).is_none()); +} + +#[test] +fn org_traffic_above_every_project_cap_never_implies_an_org_limit() { + let caps = OrgRateLimits { + by_model: [( + "gpt-4o".to_string(), + OrgRateLimit { + tokens_per_minute: 30_000_000, + requests_per_minute: None, + }, + )] + .into_iter() + .collect(), + projects_sampled: 1, + truncated: false, + fetched_at: NOW, + }; + let live = snapshot( + vec![LiveMinute { + start: NOW - 90, + models: vec![model("gpt-4o", 75_000_000, 0, 0, 10)], + }], + None, + ); + let limits = OrgLimitResolver::new(&[], Some(caps)); + let row = &live.model_rows_with_limits(&limits)[0]; + assert_eq!(row.project_cap, Some(30_000_000)); + assert_eq!(row.tpm_limit, None); + assert_eq!(row.peak_tpm_percent, None); + assert!(live.highest_utilization(&limits).is_none()); +} diff --git a/rust/src/providers/openaiapi/mod.rs b/rust/src/providers/openaiapi/mod.rs index a246405027f..08a56b570cc 100644 --- a/rust/src/providers/openaiapi/mod.rs +++ b/rust/src/providers/openaiapi/mod.rs @@ -19,6 +19,11 @@ //! 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. +//! +//! Upstream v0.75.0 (#4495) adds OpenAI Live Usage: when +//! [`FetchContext::live_usage_enabled`] is set, a successful Admin refresh also attaches the +//! last hour of tokens per minute ([`live`]); manual organization rate-limit headroom and +//! project caps live in [`org_limits`]. use async_trait::async_trait; use chrono::{DateTime, TimeZone, Utc}; @@ -29,6 +34,8 @@ use std::collections::{HashMap, HashSet}; use std::time::Duration; mod history; +pub mod live; +pub mod org_limits; use crate::core::{ CostSnapshot, FetchContext, Provider, ProviderError, ProviderFetchResult, ProviderId, @@ -38,6 +45,7 @@ use crate::core::{ const OPENAI_CREDIT_GRANTS_URL: &str = "https://api.openai.com/v1/dashboard/billing/credit_grants"; 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_ORG_PROJECTS_URL: &str = "https://api.openai.com/v1/organization/projects"; const OPENAI_API_CREDENTIAL_TARGET: &str = "codexbar-openaiapi"; /// Days of history requested from the Admin API (period label `Last 30 days`). @@ -156,6 +164,7 @@ struct Endpoints { credit_grants: String, costs: String, completions: String, + projects: String, } impl Endpoints { @@ -164,6 +173,7 @@ impl Endpoints { credit_grants: OPENAI_CREDIT_GRANTS_URL.to_string(), costs: OPENAI_ORG_COSTS_URL.to_string(), completions: OPENAI_ORG_COMPLETIONS_URL.to_string(), + projects: OPENAI_ORG_PROJECTS_URL.to_string(), } } } @@ -661,6 +671,27 @@ fn number_value(value: &serde_json::Value) -> Option { value.is_finite().then_some(value) } +impl OpenAIApiProvider { + async fn fetch_refresh( + &self, + credential: &ApiCredential, + project_id: Option<&str>, + now: DateTime, + live_usage_enabled: bool, + ) -> Result { + let mut result = self + .fetch_admin_or_balance(credential, project_id, now) + .await?; + // The balance fallback has no organization usage, so no live minutes either. + if live_usage_enabled && result.source_label == "admin-api" { + result.throughput = self + .fetch_refresh_throughput(&credential.key, project_id, now) + .await; + } + Ok(result) + } +} + impl Default for OpenAIApiProvider { fn default() -> Self { Self::new() @@ -682,8 +713,13 @@ impl Provider for OpenAIApiProvider { SourceMode::Auto | SourceMode::OAuth => { 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 + self.fetch_refresh( + &credential, + project_id.as_deref(), + Utc::now(), + ctx.live_usage_enabled, + ) + .await } SourceMode::Web | SourceMode::Cli => { Err(ProviderError::UnsupportedSource(ctx.source_mode)) @@ -746,4 +782,6 @@ fn _assert_datetime_send(_: DateTime) {} #[cfg(test)] mod history_tests; #[cfg(test)] +mod live_tests; +#[cfg(test)] mod tests; diff --git a/rust/src/providers/openaiapi/org_limits.rs b/rust/src/providers/openaiapi/org_limits.rs new file mode 100644 index 00000000000..ab8903ba920 --- /dev/null +++ b/rust/src/providers/openaiapi/org_limits.rs @@ -0,0 +1,477 @@ +//! OpenAI organization rate-limit headroom for Live Usage (upstream 0.75.0 #4495: +//! `OpenAIOrgRateLimits`, `OpenAIOrgLimitResolver`, `OpenAIOrgRateLimitsFetcher`, and the +//! menu-card `throughputPresentation`). +//! +//! OpenAI has no organization-level limits endpoint. Only the user's manual limits +//! ([`ModelRateLimit`]) drive utilization. Per-project caps are context: they are read from +//! up to 25 active projects, take the maximum TPM and RPM independently per model, and +//! never drive a percentage. Project-scoped usage is only one slice of the organization +//! traffic the limits apply to, so it never reports utilization. +//! +//! Utilization counts all input (cached included) plus output and excludes Batch API +//! traffic. Models that resolve to the same configured limit (an alias and its dated +//! snapshots) share it, and their traffic is summed per minute. + +use chrono::{DateTime, Utc}; +use serde::Deserialize; +use serde::de::DeserializeOwned; +use std::collections::{BTreeMap, HashSet}; + +use super::OpenAIApiProvider; +use super::live::{LiveModelFilter, LiveModelRow, OpenAiLiveUsage}; +use crate::core::ProviderThroughput; +use crate::settings::ModelRateLimit; + +const MAX_PROJECTS: usize = 25; +const MAX_LIST_PAGES: usize = 20; +const LIST_PAGE_LIMIT: &str = "100"; + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct OrgRateLimit { + pub tokens_per_minute: u64, + pub requests_per_minute: Option, +} + +/// Project caps discovered across the organization. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct OrgRateLimits { + pub by_model: BTreeMap, + pub projects_sampled: usize, + /// More than [`MAX_PROJECTS`] active projects exist; the sample is labeled truncated. + pub truncated: bool, + /// Epoch seconds. + pub fetched_at: i64, +} + +#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)] +pub enum OrgRateLimitsError { + #[error("Configure an OpenAI organization Admin API key.")] + MissingCredential, + #[error("Organization rate limits require an Admin API key with project read permissions.")] + AdminKeyRequired, + #[error("OpenAI organization rate limits are temporarily rate limited.")] + RateLimited, + #[error("OpenAI organization rate limits API failed (HTTP {0}).")] + ApiFailure(u16), + #[error("Could not parse the OpenAI organization rate limits response.")] + ParseFailure, + #[error("Could not reach the OpenAI organization rate limits API.")] + Network, +} + +/// Resolves a usage model ID to its manual limit and project cap. +#[derive(Debug, Clone, Default)] +pub struct OrgLimitResolver { + manual: Vec, + project_caps: Option, +} + +impl OrgLimitResolver { + pub fn new(manual: &[ModelRateLimit], project_caps: Option) -> Self { + Self { + manual: ModelRateLimit::validated(manual.to_vec()), + project_caps, + } + } + + pub fn resolve(&self, model: &str) -> Option { + let key = self.manual_limit_key(model)?; + self.manual + .iter() + .find(|row| row.model == key) + .map(|row| OrgRateLimit { + tokens_per_minute: row.tokens_per_minute, + requests_per_minute: row.requests_per_minute, + }) + } + + /// The configured model ID whose limit applies to `model`. + pub fn manual_limit_key(&self, model: &str) -> Option<&str> { + match_key(model, self.manual.iter().map(|row| row.model.as_str())) + } + + pub fn project_cap(&self, model: &str) -> Option { + let caps = &self.project_caps.as_ref()?.by_model; + match_key(model, caps.keys().map(String::as_str)).and_then(|key| caps.get(key).copied()) + } +} + +/// Exact IDs first; otherwise the longest key followed by `-` and a dated snapshot suffix, +/// so `gpt-5` covers `gpt-5-2025-08-07` and `gpt-4` covers `gpt-4-0613`, while distinct +/// models such as `gpt-5-mini` or `gpt-4o` keep their own limits. +fn match_key<'a>(model: &str, keys: impl Iterator + Clone) -> Option<&'a str> { + if let Some(exact) = keys.clone().find(|key| *key == model) { + return Some(exact); + } + let mut best: Option<&'a str> = None; + for key in keys { + let matches = model + .strip_prefix(key) + .and_then(|rest| rest.strip_prefix('-')) + .is_some_and(is_snapshot_suffix); + if matches && best.is_none_or(|best| key.len() > best.len()) { + best = Some(key); + } + } + best +} + +/// `YYYY-MM-DD` or the older `MMDD`. +fn is_snapshot_suffix(suffix: &str) -> bool { + let bytes = suffix.as_bytes(); + let digits = |range: std::ops::Range| bytes[range].iter().all(u8::is_ascii_digit); + match bytes.len() { + 4 => digits(0..4), + 10 => digits(0..4) && bytes[4] == b'-' && digits(5..7) && bytes[7] == b'-' && digits(8..10), + _ => false, + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct OrgLimitUtilization { + pub limit: OrgRateLimit, + /// Every model on the snapshot that shares the configured limit. + pub pooled_models: Vec, + pub latest_tpm_percent: f64, + pub peak_tpm_percent: f64, + pub peak_rpm_percent: Option, +} + +impl OpenAiLiveUsage { + /// Manual-limit utilization for one model's pool; `None` for all models, project-scoped + /// usage, or a model without a manual limit. + pub fn utilization( + &self, + filter: LiveModelFilter<'_>, + limits: &OrgLimitResolver, + ) -> Option { + if self.project_id.is_some() { + return None; + } + let LiveModelFilter::Model(name) = filter else { + return None; + }; + let key = limits.manual_limit_key(name)?; + let limit = limits.resolve(name)?; + let pooled_models: Vec = self + .model_names() + .into_iter() + .filter(|model| limits.manual_limit_key(model) == Some(key)) + .collect(); + let members: HashSet<&str> = pooled_models.iter().map(String::as_str).collect(); + let frontier = self.reporting_frontier(); + let points: Vec<(i64, u64, u64, bool)> = self + .minutes + .iter() + .map(|minute| { + let pooled = minute + .models + .iter() + .filter(|model| members.contains(model.name.as_str())); + let (tokens, requests) = pooled.fold((0u64, 0u64), |(tokens, requests), model| { + ( + tokens.saturating_add( + model.total_tokens().saturating_sub(model.batch_tokens), + ), + requests + .saturating_add(model.requests.saturating_sub(model.batch_requests)), + ) + }); + ( + minute.start, + tokens, + requests, + Self::is_pending(frontier, minute.start), + ) + }) + .collect(); + let latest_tokens = points + .iter() + .find(|point| Some(point.0) == frontier) + .map_or(0, |point| point.1); + let reported = points.iter().filter(|point| !point.3); + let peak_tokens = reported.clone().map(|point| point.1).max().unwrap_or(0); + let peak_requests = reported.map(|point| point.2).max().unwrap_or(0); + let tpm = limit.tokens_per_minute as f64; + Some(OrgLimitUtilization { + limit, + pooled_models, + latest_tpm_percent: latest_tokens as f64 / tpm * 100.0, + peak_tpm_percent: peak_tokens as f64 / tpm * 100.0, + peak_rpm_percent: limit + .requests_per_minute + .filter(|rpm| *rpm > 0) + .map(|rpm| peak_requests as f64 / rpm as f64 * 100.0), + }) + } + + /// [`Self::model_rows`] with manual limits, utilization and project caps. + pub fn model_rows_with_limits(&self, limits: &OrgLimitResolver) -> Vec { + self.model_rows() + .into_iter() + .map(|mut row| { + let utilization = self.utilization(LiveModelFilter::Model(&row.name), limits); + row.tpm_limit = utilization.as_ref().map(|u| u.limit.tokens_per_minute); + row.rpm_limit = utilization + .as_ref() + .and_then(|u| u.limit.requests_per_minute); + row.project_cap = limits + .project_cap(&row.name) + .map(|cap| cap.tokens_per_minute); + row.peak_tpm_percent = utilization.as_ref().map(|u| u.peak_tpm_percent); + row.latest_tpm_percent = utilization.as_ref().map(|u| u.latest_tpm_percent); + row + }) + .collect() + } + + /// The row with the highest peak utilization; ties pick the alphabetically first name. + pub fn highest_utilization(&self, limits: &OrgLimitResolver) -> Option { + let mut best: Option<(f64, LiveModelRow)> = None; + for row in self.model_rows_with_limits(limits) { + let Some(percent) = row.peak_tpm_percent else { + continue; + }; + let better = best.as_ref().is_none_or(|(best_percent, best_row)| { + percent > *best_percent || (percent == *best_percent && row.name < best_row.name) + }); + if better { + best = Some((percent, row)); + } + } + best.map(|(_, row)| row) + } +} + +/// The menu card's limit tile: the highest manual-limit utilization across every model. +#[derive(Debug, Clone, PartialEq)] +pub struct ThroughputLimitUtilization { + /// The model, or the configured limit ID when several models share it. + pub label: String, + /// Peak non-pending limit-counted tokens of the pool in one minute. + pub peak_tokens: u64, + pub tokens_per_minute_limit: u64, + pub percent: f64, +} + +/// Upstream `OpenAIAPIProviderImplementation.throughputPresentation` limit line: `None` for +/// project-scoped snapshots (the snapshot's own scope, which can outlive a settings change) +/// or when no model has a manual limit. A tie keeps the earlier pool in snapshot order. +pub fn throughput_limit_utilization( + snapshot: &ProviderThroughput, + manual: &[ModelRateLimit], +) -> Option { + if snapshot.project_id.is_some() { + return None; + } + let resolver = OrgLimitResolver::new(manual, None); + let mut groups: Vec<(&str, Vec<&crate::core::ThroughputModel>)> = Vec::new(); + for model in &snapshot.models { + let Some(key) = resolver.manual_limit_key(&model.name) else { + continue; + }; + match groups.iter_mut().find(|(group, _)| *group == key) { + Some((_, models)) => models.push(model), + None => groups.push((key, vec![model])), + } + } + let mut best: Option = None; + for (key, models) in groups { + let Some(limit) = resolver.resolve(&models[0].name) else { + continue; + }; + let aligned = models + .iter() + .all(|model| model.limit_counted_tokens_by_minute.len() == snapshot.minutes.len()); + let peak = if aligned { + snapshot + .minutes + .iter() + .enumerate() + .filter(|(_, minute)| !minute.is_pending) + .map(|(index, _)| { + models + .iter() + .map(|model| model.limit_counted_tokens_by_minute[index]) + .sum::() + }) + .max() + .unwrap_or(0) + } else { + models + .iter() + .map(|model| model.peak_limit_counted_tokens) + .max() + .unwrap_or(0) + }; + let percent = peak as f64 / limit.tokens_per_minute as f64 * 100.0; + if best.as_ref().is_none_or(|best| percent > best.percent) { + best = Some(ThroughputLimitUtilization { + label: if models.len() > 1 { + key.to_string() + } else { + models[0].name.clone() + }, + peak_tokens: peak, + tokens_per_minute_limit: limit.tokens_per_minute, + percent, + }); + } + } + best +} + +#[derive(Debug, Deserialize)] +struct ListPage { + data: Vec, + has_more: bool, + last_id: Option, +} + +#[derive(Debug, Deserialize)] +struct Project { + id: String, + status: Option, +} + +#[derive(Debug, Deserialize)] +struct ProjectRateLimit { + model: String, + max_tokens_per_1_minute: Option, + max_requests_per_1_minute: Option, +} + +impl OpenAIApiProvider { + /// Project caps across the organization for the Live Usage window. Discovery is always + /// organization-wide, even when usage is project-scoped. + pub async fn fetch_org_rate_limits( + &self, + api_key: Option<&str>, + ) -> Result { + let credential = + Self::credential(api_key).map_err(|_| OrgRateLimitsError::MissingCredential)?; + self.fetch_org_rate_limits_with(&credential.key, Utc::now()) + .await + } + + pub(super) async fn fetch_org_rate_limits_with( + &self, + api_key: &str, + now: DateTime, + ) -> Result { + let projects_url = reqwest::Url::parse(&self.endpoints.projects) + .map_err(|_| OrgRateLimitsError::ParseFailure)?; + // A missing or null status counts as active; archived projects are excluded by default. + let (projects, truncated) = self + .list_pages::( + projects_url.clone(), + api_key, + Some(MAX_PROJECTS), + |project| project.status.as_deref().unwrap_or("active") == "active", + ) + .await?; + let mut by_model: BTreeMap = BTreeMap::new(); + for project in &projects { + let mut url = projects_url.clone(); + url.path_segments_mut() + .map_err(|_| OrgRateLimitsError::ParseFailure)? + .push(&project.id) + .push("rate_limits"); + let (rows, _) = self + .list_pages::(url, api_key, None, |_| true) + .await?; + for row in rows { + let model = row.model.trim(); + let tpm = row.max_tokens_per_1_minute; + let rpm = row.max_requests_per_1_minute; + if model.is_empty() || tpm.is_some_and(|v| v < 0) || rpm.is_some_and(|v| v < 0) { + return Err(OrgRateLimitsError::ParseFailure); + } + // Some models omit a token limit (video models, for example). + let Some(tpm) = tpm.and_then(|v| u64::try_from(v).ok()).filter(|v| *v > 0) else { + continue; + }; + let rpm = rpm.and_then(|v| u64::try_from(v).ok()).filter(|v| *v > 0); + let entry = by_model.entry(model.to_string()).or_insert(OrgRateLimit { + tokens_per_minute: 0, + requests_per_minute: None, + }); + entry.tokens_per_minute = entry.tokens_per_minute.max(tpm); + entry.requests_per_minute = entry.requests_per_minute.max(rpm); + } + } + Ok(OrgRateLimits { + by_model, + projects_sampled: projects.len(), + truncated, + fetched_at: now.timestamp(), + }) + } + + /// Follows `has_more` / `last_id` (`after=`) for at most [`MAX_LIST_PAGES`] pages. + /// With `max_items`, stops once that many included rows exist and reports truncation. + async fn list_pages( + &self, + url: reqwest::Url, + api_key: &str, + max_items: Option, + include: impl Fn(&T) -> bool, + ) -> Result<(Vec, bool), OrgRateLimitsError> { + let mut rows = Vec::new(); + let mut cursor: Option = None; + let mut seen = HashSet::new(); + for _ in 0..MAX_LIST_PAGES { + let mut query = vec![("limit", LIST_PAGE_LIMIT.to_string())]; + if let Some(cursor) = &cursor { + query.push(("after", cursor.clone())); + } + let response = self + .client + .get(url.clone()) + .query(&query) + .bearer_auth(api_key) + .header("Accept", "application/json") + .send() + .await + .map_err(|_| OrgRateLimitsError::Network)?; + let status = response.status(); + if !status.is_success() { + return Err(match status.as_u16() { + 401 | 403 => OrgRateLimitsError::AdminKeyRequired, + 429 => OrgRateLimitsError::RateLimited, + code => OrgRateLimitsError::ApiFailure(code), + }); + } + let body = response + .text() + .await + .map_err(|_| OrgRateLimitsError::Network)?; + let page: ListPage = + serde_json::from_str(&body).map_err(|_| OrgRateLimitsError::ParseFailure)?; + for row in page.data { + if !include(&row) { + continue; + } + if max_items.is_some_and(|max| rows.len() == max) { + return Ok((rows, true)); + } + rows.push(row); + } + if !page.has_more { + return Ok((rows, false)); + } + if max_items.is_some_and(|max| rows.len() == max) { + return Ok((rows, true)); + } + let next = page + .last_id + .filter(|next| !next.is_empty()) + .ok_or(OrgRateLimitsError::ParseFailure)?; + if !seen.insert(next.clone()) { + return Err(OrgRateLimitsError::ParseFailure); + } + cursor = Some(next); + } + Err(OrgRateLimitsError::ParseFailure) + } +} diff --git a/rust/src/providers/openaiapi/tests.rs b/rust/src/providers/openaiapi/tests.rs index 27be966306e..323a0bdc70e 100644 --- a/rust/src/providers/openaiapi/tests.rs +++ b/rust/src/providers/openaiapi/tests.rs @@ -23,6 +23,7 @@ fn provider(server: &ServerGuard) -> OpenAIApiProvider { credit_grants: format!("{}{GRANTS_PATH}", server.url()), costs: format!("{}{COSTS_PATH}", server.url()), completions: format!("{}{COMPLETIONS_PATH}", server.url()), + projects: format!("{}/v1/organization/projects", server.url()), }; provider.retry = RetryPolicy { base_delay: Duration::ZERO, diff --git a/rust/src/settings.rs b/rust/src/settings.rs index f553cdbba56..dbd8873fa99 100755 --- a/rust/src/settings.rs +++ b/rust/src/settings.rs @@ -41,6 +41,7 @@ mod api_keys; mod config_path; mod cost_time_zone; mod manual_cookies; +mod model_rate_limits; mod optional_details; mod preferences_document; mod provider_workspace; @@ -57,6 +58,7 @@ pub use config_path::{ }; pub use cost_time_zone::*; pub use manual_cookies::*; +pub use model_rate_limits::{ModelRateLimit, rate_limit_number}; pub use optional_details::provider_has_optional_details; pub use preferences_document::*; pub use provider_workspace::*; diff --git a/rust/src/settings/model_rate_limits.rs b/rust/src/settings/model_rate_limits.rs new file mode 100644 index 00000000000..20a367493a7 --- /dev/null +++ b/rust/src/settings/model_rate_limits.rs @@ -0,0 +1,428 @@ +//! OpenAI Live Usage settings (upstream 0.75.0 #4495): the per-refresh Tokens/min fetch +//! toggle and the manual per-model organization rate limits. +//! +//! OpenAI has no organization-level limits endpoint, so headroom is measured only against +//! limits the user copies from the organization Limits page. On Windows the OpenAI API +//! provider has no token accounts, so the provider-level list is the only credential scope. + +use serde::{Deserialize, Deserializer, Serialize}; + +use super::Settings; +use crate::core::ProviderId; + +/// One manual organization rate limit. Rows are trimmed and validated on load and on +/// every write; see [`ModelRateLimit::validated`]. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct ModelRateLimit { + pub model: String, + #[serde(alias = "tokensPerMinute")] + pub tokens_per_minute: u64, + #[serde( + default, + alias = "requestsPerMinute", + skip_serializing_if = "Option::is_none" + )] + pub requests_per_minute: Option, +} + +/// Wire shape used only while loading, so negative or zero values are dropped rather than +/// failing the whole settings file. +#[derive(Deserialize)] +struct RawModelRateLimit { + model: String, + #[serde(alias = "tokensPerMinute")] + tokens_per_minute: i64, + #[serde(default, alias = "requestsPerMinute")] + requests_per_minute: Option, +} + +impl ModelRateLimit { + pub fn new( + model: impl Into, + tokens_per_minute: u64, + requests_per_minute: Option, + ) -> Self { + Self { + model: model.into(), + tokens_per_minute, + requests_per_minute, + } + } + + /// Upstream `ProviderModelRateLimit.validated`: model IDs are trimmed, empty IDs and + /// zero-TPM rows are dropped, a zero RPM becomes absent, and the last valid row for a + /// duplicate ID wins while keeping its position among the survivors. + pub fn validated(rows: Vec) -> Vec { + let mut seen = std::collections::HashSet::new(); + let mut kept: Vec = rows + .into_iter() + .rev() + .filter_map(|row| { + let model = row.model.trim().to_string(); + if model.is_empty() || row.tokens_per_minute == 0 || !seen.insert(model.clone()) { + return None; + } + Some(Self { + model, + tokens_per_minute: row.tokens_per_minute, + requests_per_minute: row.requests_per_minute.filter(|rpm| *rpm > 0), + }) + }) + .collect(); + kept.reverse(); + kept + } +} + +/// Lenient loader for `ProviderConfig::model_rate_limits`: a malformed list or row is +/// dropped instead of resetting every setting, and the survivors are validated. +pub(super) fn deserialize_model_rate_limits<'de, D>( + deserializer: D, +) -> Result>, D::Error> +where + D: Deserializer<'de>, +{ + let value = Option::::deserialize(deserializer)?; + let Some(serde_json::Value::Array(items)) = value else { + return Ok(None); + }; + let rows = items + .into_iter() + .filter_map(|item| serde_json::from_value::(item).ok()) + .filter_map(|row| { + Some(ModelRateLimit { + model: row.model, + tokens_per_minute: u64::try_from(row.tokens_per_minute).ok()?, + requests_per_minute: row + .requests_per_minute + .and_then(|rpm| u64::try_from(rpm).ok()), + }) + }) + .collect(); + let rows = ModelRateLimit::validated(rows); + Ok((!rows.is_empty()).then_some(rows)) +} + +/// Providers with the OpenAI Live Usage settings. +fn supports_live_usage(id: ProviderId) -> bool { + id == ProviderId::OpenAIApi +} + +impl Settings { + /// Whether each refresh also fetches the last hour of 1-minute usage (upstream + /// `OPENAI_LIVE_MINUTES`). On by default for OpenAI API; false for every other provider. + pub fn live_usage_enabled(&self, id: ProviderId) -> bool { + supports_live_usage(id) + && self + .provider_configs + .get(&id) + .and_then(|config| config.live_usage_enabled) + .unwrap_or(true) + } + + pub fn set_live_usage_enabled(&mut self, id: ProviderId, enabled: bool) { + if supports_live_usage(id) { + // The default is on, so only an opt-out is persisted. + self.provider_config_mut(id).live_usage_enabled = (!enabled).then_some(false); + } + } + + /// Validated manual organization rate limits for `id`; empty when unset. + pub fn model_rate_limits(&self, id: ProviderId) -> &[ModelRateLimit] { + if !supports_live_usage(id) { + return &[]; + } + self.provider_configs + .get(&id) + .and_then(|config| config.model_rate_limits.as_deref()) + .unwrap_or(&[]) + } + + /// Replace `id`'s manual limits with the validated rows. An empty list clears them. + pub fn set_model_rate_limits(&mut self, id: ProviderId, limits: Vec) { + if supports_live_usage(id) { + let limits = ModelRateLimit::validated(limits); + self.provider_config_mut(id).model_rate_limits = (!limits.is_empty()).then_some(limits); + } + } +} + +/// Upstream `ProviderRateLimitNumber`: the Settings editor's number parser/formatter. +pub mod rate_limit_number { + /// Accepts plain (`40000000`) or comma-grouped (`40,000,000`) integers, optionally with + /// a decimal and a `k`/`M` suffix (`150k`, `1.5M`). A fraction without a suffix, or one + /// that does not land on a whole count, is rejected rather than rounded. + pub fn parse(text: &str) -> Option { + let text = text.trim(); + let (body, multiplier) = match text.chars().last()? { + 'k' | 'K' => (&text[..text.len() - 1], 1_000_u64), + 'm' | 'M' => (&text[..text.len() - 1], 1_000_000_u64), + _ => (text, 1_u64), + }; + let (integer, fraction) = match body.split_once('.') { + Some((integer, fraction)) => (integer, Some(fraction)), + None => (body, None), + }; + if fraction.is_some() && multiplier == 1 { + return None; + } + let digits = integer_digits(integer)?; + let fraction = fraction.unwrap_or(""); + if !fraction.chars().all(|ch| ch.is_ascii_digit()) + || (body.contains('.') && fraction.is_empty()) + { + return None; + } + let scale = 10_u64.checked_pow(u32::try_from(fraction.len()).ok()?)?; + let whole: u64 = digits.parse().ok()?; + let fraction_value: u64 = if fraction.is_empty() { + 0 + } else { + fraction.parse().ok()? + }; + let numerator = u128::from(whole) * u128::from(scale) + u128::from(fraction_value); + let scaled = numerator.checked_mul(u128::from(multiplier))?; + if scaled % u128::from(scale) != 0 { + return None; + } + let value = u64::try_from(scaled / u128::from(scale)).ok()?; + (value > 0 && value <= i64::MAX as u64).then_some(value) + } + + /// `[0-9]+` or `[0-9]{1,3}(,[0-9]{3})+`, returned without commas. + fn integer_digits(text: &str) -> Option { + if text.is_empty() { + return None; + } + if !text.contains(',') { + return text + .chars() + .all(|ch| ch.is_ascii_digit()) + .then(|| text.to_string()); + } + let mut groups = text.split(','); + let head = groups.next()?; + if head.is_empty() || head.len() > 3 || !head.chars().all(|ch| ch.is_ascii_digit()) { + return None; + } + let mut digits = head.to_string(); + for group in groups { + if group.len() != 3 || !group.chars().all(|ch| ch.is_ascii_digit()) { + return None; + } + digits.push_str(group); + } + Some(digits) + } + + /// Shortens only exactly representable values (`40000000` -> `40M`, `1500000` -> + /// `1.5M`), so reopening Settings never changes a stored limit. + pub fn format(value: u64) -> String { + for (divisor, suffix) in [(1_000_000_u64, "M"), (1_000_u64, "K")] { + if value >= divisor && value.is_multiple_of(divisor / 10) { + let whole = value / divisor; + let tenths = value % divisor / (divisor / 10); + return if tenths == 0 { + format!("{whole}{suffix}") + } else { + format!("{whole}.{tenths}{suffix}") + }; + } + } + value.to_string() + } +} + +#[cfg(test)] +mod tests { + use super::rate_limit_number::{format, parse}; + use super::*; + use crate::settings::ProviderConfig; + + fn row(model: &str, tpm: u64, rpm: Option) -> ModelRateLimit { + ModelRateLimit::new(model, tpm, rpm) + } + + #[test] + fn validated_trims_drops_invalid_rows_and_keeps_last_duplicate() { + let rows = ModelRateLimit::validated(vec![ + row(" gpt-5 ", 1_000, Some(10)), + row("", 5, None), + row("gpt-4o", 0, None), + row("o3", 2_000, Some(0)), + row("gpt-5", 3_000, None), + ]); + assert_eq!( + rows, + vec![row("o3", 2_000, None), row("gpt-5", 3_000, None)] + ); + } + + #[test] + fn live_usage_defaults_on_for_openai_only_and_persists_only_the_opt_out() { + let mut settings = Settings::default(); + assert!(settings.live_usage_enabled(ProviderId::OpenAIApi)); + assert!(!settings.live_usage_enabled(ProviderId::Codex)); + + settings.set_live_usage_enabled(ProviderId::OpenAIApi, false); + assert!(!settings.live_usage_enabled(ProviderId::OpenAIApi)); + let config = settings.provider_config(ProviderId::OpenAIApi).unwrap(); + assert_eq!(config.live_usage_enabled, Some(false)); + + settings.set_live_usage_enabled(ProviderId::OpenAIApi, true); + assert!(settings.live_usage_enabled(ProviderId::OpenAIApi)); + let config = settings.provider_config(ProviderId::OpenAIApi).unwrap(); + assert_eq!(config.live_usage_enabled, None); + + settings.set_live_usage_enabled(ProviderId::Codex, true); + assert!(settings.provider_config(ProviderId::Codex).is_none()); + } + + #[test] + fn model_rate_limits_set_validates_and_empty_clears() { + let mut settings = Settings::default(); + settings.set_model_rate_limits( + ProviderId::OpenAIApi, + vec![row(" gpt-5 ", 40_000_000, Some(15_000)), row(" ", 1, None)], + ); + assert_eq!( + settings.model_rate_limits(ProviderId::OpenAIApi), + &[row("gpt-5", 40_000_000, Some(15_000))] + ); + assert!(settings.model_rate_limits(ProviderId::Codex).is_empty()); + + settings.set_model_rate_limits(ProviderId::OpenAIApi, Vec::new()); + assert!(settings.model_rate_limits(ProviderId::OpenAIApi).is_empty()); + let config = settings.provider_config(ProviderId::OpenAIApi).unwrap(); + assert!(config.model_rate_limits.is_none()); + } + + #[test] + fn provider_config_round_trips_and_omits_defaults() { + let encoded = serde_json::to_value(ProviderConfig::default()).unwrap(); + assert!(encoded.get("model_rate_limits").is_none()); + assert!(encoded.get("live_usage_enabled").is_none()); + + let config = ProviderConfig { + live_usage_enabled: Some(false), + model_rate_limits: Some(vec![ + row("gpt-5", 40_000_000, Some(15_000)), + row("o3", 2_000, None), + ]), + ..ProviderConfig::default() + }; + let encoded = serde_json::to_value(&config).unwrap(); + assert_eq!( + encoded["model_rate_limits"], + serde_json::json!([ + {"model": "gpt-5", "tokens_per_minute": 40_000_000, "requests_per_minute": 15_000}, + {"model": "o3", "tokens_per_minute": 2_000}, + ]) + ); + let decoded: ProviderConfig = serde_json::from_value(encoded).unwrap(); + assert_eq!(decoded, config); + } + + #[test] + fn legacy_config_without_live_fields_loads_with_defaults() { + let decoded: ProviderConfig = + serde_json::from_str(r#"{"workspace_id":"proj_abc","api_token":"sk-admin"}"#).unwrap(); + assert_eq!(decoded.live_usage_enabled, None); + assert_eq!(decoded.model_rate_limits, None); + assert_eq!(decoded.workspace_id.as_deref(), Some("proj_abc")); + } + + #[test] + fn malformed_limits_are_dropped_without_failing_the_config() { + let decoded: ProviderConfig = serde_json::from_str( + r#"{ + "workspace_id": "proj_abc", + "model_rate_limits": [ + {"model": " gpt-5 ", "tokensPerMinute": 40000000, "requestsPerMinute": 15000}, + {"model": "neg", "tokens_per_minute": -5}, + {"model": "rpm", "tokens_per_minute": 10, "requests_per_minute": -1}, + {"model": 7, "tokens_per_minute": 10}, + "garbage", + {"model": "zero", "tokens_per_minute": 0} + ] + }"#, + ) + .unwrap(); + assert_eq!( + decoded.model_rate_limits, + Some(vec![ + row("gpt-5", 40_000_000, Some(15_000)), + row("rpm", 10, None), + ]) + ); + assert_eq!(decoded.workspace_id.as_deref(), Some("proj_abc")); + + for bad in [ + r#"{"model_rate_limits": "x"}"#, + r#"{"model_rate_limits": null}"#, + ] { + let decoded: ProviderConfig = serde_json::from_str(bad).unwrap(); + assert_eq!(decoded.model_rate_limits, None); + } + } + + #[test] + fn settings_file_round_trips_live_usage_fields() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("settings.json"); + let mut settings = Settings::default(); + settings.set_live_usage_enabled(ProviderId::OpenAIApi, false); + settings.set_model_rate_limits( + ProviderId::OpenAIApi, + vec![row("gpt-4o-mini", 150_000, Some(500))], + ); + settings.save_to_path(&path).unwrap(); + + let loaded = Settings::load_from_path(Some(&path)); + assert!(!loaded.live_usage_enabled(ProviderId::OpenAIApi)); + assert_eq!( + loaded.model_rate_limits(ProviderId::OpenAIApi), + &[row("gpt-4o-mini", 150_000, Some(500))] + ); + } + + #[test] + fn rate_limit_number_parses_upstream_forms() { + assert_eq!(parse("40000000"), Some(40_000_000)); + assert_eq!(parse(" 40,000,000 "), Some(40_000_000)); + assert_eq!(parse("40M"), Some(40_000_000)); + assert_eq!(parse("150k"), Some(150_000)); + assert_eq!(parse("1.5M"), Some(1_500_000)); + assert_eq!(parse("1.25k"), Some(1_250)); + for invalid in [ + "", + "0", + "1.5", + "1.0005k", + "40,00", + "1,0000", + ",100", + "1.", + "-5", + "5x", + "k", + "1.5.5M", + "99999999999999999999", + ] { + assert_eq!(parse(invalid), None, "{invalid}"); + } + } + + #[test] + fn rate_limit_number_formats_only_exact_values() { + assert_eq!(format(40_000_000), "40M"); + assert_eq!(format(1_500_000), "1.5M"); + assert_eq!(format(150_000), "150K"); + assert_eq!(format(1_250), "1250"); + assert_eq!(format(1_234_567), "1234567"); + assert_eq!(format(999), "999"); + for value in [40_000_000, 1_500_000, 150_000, 1_250, 1_234_567] { + assert_eq!(parse(&format(value)), Some(value)); + } + } +} diff --git a/rust/src/settings/types.rs b/rust/src/settings/types.rs index 031ce2f7e7d..06ce0763dee 100644 --- a/rust/src/settings/types.rs +++ b/rust/src/settings/types.rs @@ -606,6 +606,18 @@ pub struct ProviderConfig { /// publish this denominator; it is only used when a valid value is set. #[serde(skip_serializing_if = "Option::is_none")] pub seat_credit_entitlement: Option, + /// OpenAI API-only: `Some(false)` opts out of the per-refresh Tokens/min fetch. + /// `None` keeps upstream's default of on. + #[serde(skip_serializing_if = "Option::is_none")] + pub live_usage_enabled: Option, + /// OpenAI API-only: manual per-model organization rate limits for Live Usage + /// headroom. Validated on load; `None` when no valid row is stored. + #[serde( + default, + deserialize_with = "super::model_rate_limits::deserialize_model_rate_limits", + skip_serializing_if = "Option::is_none" + )] + pub model_rate_limits: Option>, } /// How the single tray icon is drawn (Mac "Menu bar style").