1use std::collections::{BTreeMap, HashMap, HashSet};
2use std::sync::{Arc, Mutex, OnceLock};
3use std::time::{Duration, Instant};
4
5use anyhow::{Context, Result};
6use futures_util::stream::{FuturesUnordered, StreamExt};
7use reqwest::Client;
8use serde::{Deserialize, Serialize};
9use tracing::{info, warn};
10
11use crate::balance::{
12 BalanceSnapshotStatus, ProviderBalanceSnapshot, ProviderUsageAlert, ProviderUsageAlertKind,
13 ProviderUsageModelStat, ProviderUsageRateSnapshot, ProviderUsageWindow,
14};
15use crate::config::{ProxyConfig, ServiceConfigManager, proxy_home_dir};
16use crate::lb::LbState;
17use crate::policy_actions::{PolicyAction, PolicyActionKind};
18use crate::pricing::UsdAmount;
19use crate::provider_signals::{
20 ProviderSignal, ProviderSignalKind, ProviderSignalSource, ProviderSignalTarget,
21};
22use crate::runtime_identity::ProviderEndpointKey;
23use crate::state::ProxyState;
24use crate::usage_forecast::next_reset_at_ms;
25
26#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq, Hash)]
27#[serde(rename_all = "snake_case")]
28enum ProviderKind {
29 BudgetHttpJson,
31 YescodeProfile,
33 #[serde(
35 rename = "openai_balance_http_json",
36 alias = "open_ai_balance_http_json",
37 alias = "relay_balance_http_json"
38 )]
39 OpenAiBalanceHttpJson,
40 #[serde(rename = "sub2api_usage", alias = "sub2api_usage_http_json")]
42 Sub2ApiUsage,
43 #[serde(rename = "sub2api_auth_me", alias = "sub2api_auth_me_http_json")]
45 Sub2ApiAuthMe,
46 #[serde(
48 rename = "new_api_token_usage",
49 alias = "new_api_token_usage_http_json"
50 )]
51 NewApiTokenUsage,
52 NewApiUserSelf,
54 #[serde(
56 rename = "rightcode_account_summary",
57 alias = "right_code_account_summary",
58 alias = "rightcode"
59 )]
60 RightCodeAccountSummary,
61 #[serde(
63 rename = "openai_organization_costs",
64 alias = "openai_org_costs",
65 alias = "openai_costs"
66 )]
67 OpenAiOrganizationCosts,
68}
69
70impl ProviderKind {
71 fn source_name(&self) -> &'static str {
72 match self {
73 ProviderKind::BudgetHttpJson => "usage_provider:budget_http_json",
74 ProviderKind::YescodeProfile => "usage_provider:yescode_profile",
75 ProviderKind::OpenAiBalanceHttpJson => "usage_provider:openai_balance_http_json",
76 ProviderKind::Sub2ApiUsage => "usage_provider:sub2api_usage",
77 ProviderKind::Sub2ApiAuthMe => "usage_provider:sub2api_auth_me",
78 ProviderKind::NewApiTokenUsage => "usage_provider:new_api_token_usage",
79 ProviderKind::NewApiUserSelf => "usage_provider:new_api_user_self",
80 ProviderKind::RightCodeAccountSummary => "usage_provider:rightcode_account_summary",
81 ProviderKind::OpenAiOrganizationCosts => "usage_provider:openai_organization_costs",
82 }
83 }
84
85 fn default_endpoint(&self) -> Option<&'static str> {
86 match self {
87 ProviderKind::OpenAiBalanceHttpJson => Some("{{base_url}}/user/balance"),
88 ProviderKind::Sub2ApiUsage => Some("{{base_url}}/v1/usage"),
89 ProviderKind::Sub2ApiAuthMe => Some("{{base_url}}/api/v1/auth/me"),
90 ProviderKind::NewApiTokenUsage => Some("{{base_url}}/api/usage/token/"),
91 ProviderKind::NewApiUserSelf => Some("{{base_url}}/api/user/self"),
92 ProviderKind::RightCodeAccountSummary => {
93 Some("https://www.right.codes/account/summary")
94 }
95 ProviderKind::OpenAiOrganizationCosts => {
96 Some("{{base_url}}/v1/organization/costs?start_time={{unix_days_ago:30}}&limit=30")
97 }
98 _ => None,
99 }
100 }
101}
102
103#[derive(Debug, Deserialize, Serialize, Default, Clone)]
104#[serde(default)]
105struct UsageProviderExtractConfig {
106 #[serde(skip_serializing_if = "Vec::is_empty")]
107 remaining_balance_paths: Vec<String>,
108 #[serde(skip_serializing_if = "Vec::is_empty")]
109 subscription_balance_paths: Vec<String>,
110 #[serde(skip_serializing_if = "Vec::is_empty")]
111 paygo_balance_paths: Vec<String>,
112 #[serde(skip_serializing_if = "Vec::is_empty")]
113 monthly_budget_paths: Vec<String>,
114 #[serde(skip_serializing_if = "Vec::is_empty")]
115 monthly_spent_paths: Vec<String>,
116 #[serde(skip_serializing_if = "Vec::is_empty")]
117 exhausted_paths: Vec<String>,
118 #[serde(skip_serializing_if = "Option::is_none")]
119 remaining_divisor: Option<u64>,
120 #[serde(skip_serializing_if = "Option::is_none")]
121 monthly_budget_divisor: Option<u64>,
122 #[serde(skip_serializing_if = "Option::is_none")]
123 monthly_spent_divisor: Option<u64>,
124 #[serde(skip_serializing_if = "bool_is_false")]
125 derive_budget_from_remaining_and_spent: bool,
126 #[serde(skip_serializing_if = "bool_is_false")]
127 derive_remaining_from_budget_and_spent: bool,
128}
129
130impl UsageProviderExtractConfig {
131 fn is_empty(&self) -> bool {
132 self.remaining_balance_paths.is_empty()
133 && self.subscription_balance_paths.is_empty()
134 && self.paygo_balance_paths.is_empty()
135 && self.monthly_budget_paths.is_empty()
136 && self.monthly_spent_paths.is_empty()
137 && self.exhausted_paths.is_empty()
138 && self.remaining_divisor.is_none()
139 && self.monthly_budget_divisor.is_none()
140 && self.monthly_spent_divisor.is_none()
141 && !self.derive_budget_from_remaining_and_spent
142 && !self.derive_remaining_from_budget_and_spent
143 }
144}
145
146#[derive(Debug, Deserialize, Serialize)]
147struct UsageProviderConfig {
148 id: String,
149 kind: ProviderKind,
150 domains: Vec<String>,
151 #[serde(default)]
152 endpoint: String,
153 #[serde(default)]
154 token_env: Option<String>,
155 #[serde(default, skip_serializing_if = "bool_is_false")]
156 require_token_env: bool,
157 #[serde(default)]
158 poll_interval_secs: Option<u64>,
159 #[serde(
160 default = "default_refresh_on_request",
161 skip_serializing_if = "bool_is_true"
162 )]
163 refresh_on_request: bool,
164 #[serde(
165 default = "default_trust_exhaustion_for_routing",
166 skip_serializing_if = "bool_is_true"
167 )]
168 trust_exhaustion_for_routing: bool,
169 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
170 headers: BTreeMap<String, String>,
171 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
172 variables: BTreeMap<String, String>,
173 #[serde(default, skip_serializing_if = "UsageProviderExtractConfig::is_empty")]
174 extract: UsageProviderExtractConfig,
175}
176
177#[derive(Debug, Deserialize, Serialize, Default)]
178struct UsageProvidersFile {
179 #[serde(default)]
180 providers: Vec<UsageProviderConfig>,
181}
182
183#[derive(Debug, Clone)]
184struct UpstreamRef {
185 station_name: String,
186 index: usize,
187 provider_endpoint: Option<ProviderEndpointKey>,
188}
189
190#[derive(Debug, Clone)]
191struct UsageProviderTarget {
192 upstream: UpstreamRef,
193 base_url: String,
194 provider_id: Option<String>,
195}
196
197#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
198struct UsageProviderTargetKey {
199 station_name: String,
200 upstream_index: usize,
201}
202
203#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
204pub struct UsageProviderRefreshSummary {
205 pub providers_configured: usize,
206 pub providers_matched: usize,
207 pub upstreams_matched: usize,
208 pub attempted: usize,
209 pub refreshed: usize,
210 pub failed: usize,
211 pub missing_token: usize,
212 #[serde(skip_serializing_if = "usize_is_zero")]
213 pub auto_attempted: usize,
214 #[serde(skip_serializing_if = "usize_is_zero")]
215 pub auto_refreshed: usize,
216 #[serde(skip_serializing_if = "usize_is_zero")]
217 pub auto_failed: usize,
218 #[serde(skip_serializing_if = "usize_is_zero")]
219 pub deduplicated: usize,
220}
221
222#[derive(Debug, Clone, Copy, PartialEq, Eq)]
223enum UsageProviderRefreshOutcome {
224 Refreshed,
225 Failed,
226 MissingToken,
227}
228
229struct RefreshProviderTargetParams<'a> {
230 client: &'a Client,
231 provider: &'a UsageProviderConfig,
232 target: &'a UsageProviderTarget,
233 cfg: &'a ProxyConfig,
234 lb_states: &'a Arc<Mutex<HashMap<String, LbState>>>,
235 state: &'a Arc<ProxyState>,
236 service_name: &'a str,
237 interval_secs: u64,
238}
239
240static LAST_USAGE_POLL: OnceLock<Mutex<HashMap<String, Instant>>> = OnceLock::new();
242static REQUEST_BALANCE_QUEUE: OnceLock<Mutex<HashMap<ProviderEndpointKey, Instant>>> =
243 OnceLock::new();
244static AUTO_PROBE_KIND_HINTS: OnceLock<Mutex<HashMap<String, ProviderKind>>> = OnceLock::new();
245static AUTO_PROBE_KIND_FAILURES: OnceLock<Mutex<HashMap<AutoProbeKindFailureKey, Instant>>> =
246 OnceLock::new();
247static USAGE_PROVIDER_TARGET_SUPPRESSIONS: OnceLock<
248 Mutex<HashMap<ProviderTargetSuppressionKey, ProviderTargetSuppression>>,
249> = OnceLock::new();
250
251const DEFAULT_POLL_INTERVAL_SECS: u64 = 10 * 60;
252const MIN_POLL_INTERVAL_SECS: u64 = 2 * 60;
254pub const REQUEST_BALANCE_REFRESH_DELAY: Duration = Duration::from_secs(60);
255const BALANCE_REFRESH_CONCURRENCY: usize = 6;
256const BALANCE_HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(6);
257const BALANCE_HTTP_ERROR_BODY_LIMIT: usize = 2_048;
258const AUTO_PROBE_KIND_FAILURE_TTL: Duration = Duration::from_secs(10 * 60);
259const USAGE_PROVIDER_TERMINAL_FAILURE_TTL: Duration = Duration::from_secs(6 * 60 * 60);
260const USAGE_PROVIDER_EXHAUSTED_SUPPRESSION_TTL: Duration = Duration::from_secs(6 * 60 * 60);
261const USAGE_PROVIDER_DAILY_RESET_SUPPRESSION_GRACE: Duration = Duration::from_secs(5 * 60);
262const LOW_BALANCE_ALERT_THRESHOLD_USD: &str = "10";
263const EXPIRING_SOON_WINDOW_SECS: u64 = 7 * 24 * 60 * 60;
264const AUTO_PROVIDER_ID_PREFIX: &str = "auto:balance:";
265const AUTO_PROBE_KINDS: [ProviderKind; 5] = [
266 ProviderKind::RightCodeAccountSummary,
267 ProviderKind::Sub2ApiUsage,
268 ProviderKind::NewApiTokenUsage,
269 ProviderKind::NewApiUserSelf,
270 ProviderKind::OpenAiBalanceHttpJson,
271];
272
273#[derive(Debug, Clone, PartialEq, Eq, Hash)]
274struct AutoProbeKindFailureKey {
275 provider_id: String,
276 target: AutoProbeTargetKey,
277 kind: ProviderKind,
278}
279
280#[derive(Debug, Clone, PartialEq, Eq, Hash)]
281struct ProviderTargetSuppressionKey {
282 provider_id: String,
283 target: AutoProbeTargetKey,
284}
285
286#[derive(Debug, Clone)]
287struct ProviderTargetSuppression {
288 until: Instant,
289 reason: String,
290 routing_exhausted: bool,
291}
292
293#[derive(Debug, Clone)]
294struct ProviderTargetSuppressionDecision {
295 reason: String,
296 routing_exhausted: bool,
297 ttl: Duration,
298}
299
300#[derive(Debug, Clone, PartialEq, Eq, Hash)]
301struct AutoProbeTargetKey {
302 station_name: String,
303 upstream_index: usize,
304 provider_endpoint_key: Option<String>,
305 base_url: String,
306}
307
308#[derive(Debug, Clone, Copy, PartialEq, Eq)]
309enum RequestBalanceQueueDue {
310 Due,
311 NotDue(Duration),
312 Missing,
313}
314
315#[derive(Debug, Clone, Copy, PartialEq, Eq)]
316enum RequestBalancePollOutcome {
317 Attempted,
318 Deferred(Duration),
319 Skipped,
320}
321
322fn bool_is_false(value: &bool) -> bool {
323 !*value
324}
325
326fn bool_is_true(value: &bool) -> bool {
327 *value
328}
329
330fn usize_is_zero(value: &usize) -> bool {
331 *value == 0
332}
333
334fn default_refresh_on_request() -> bool {
335 true
336}
337
338fn default_trust_exhaustion_for_routing() -> bool {
339 true
340}
341
342fn unix_now_ms() -> u64 {
343 std::time::SystemTime::now()
344 .duration_since(std::time::UNIX_EPOCH)
345 .map(|d| d.as_millis() as u64)
346 .unwrap_or(0)
347}
348
349fn unix_now_secs() -> u64 {
350 std::time::SystemTime::now()
351 .duration_since(std::time::UNIX_EPOCH)
352 .map(|d| d.as_secs())
353 .unwrap_or(0)
354}
355
356fn stale_after_ms(fetched_at_ms: u64, interval_secs: u64) -> Option<u64> {
357 fetched_at_ms.checked_add(interval_secs.saturating_mul(3).saturating_mul(1_000))
358}
359
360fn snapshot_refresh_interval_secs(provider: &UsageProviderConfig) -> u64 {
361 let interval_secs = provider
362 .poll_interval_secs
363 .unwrap_or(DEFAULT_POLL_INTERVAL_SECS);
364 if interval_secs == 0 {
365 DEFAULT_POLL_INTERVAL_SECS
366 } else {
367 interval_secs.max(MIN_POLL_INTERVAL_SECS)
368 }
369}
370
371fn effective_poll_interval_secs(provider: &UsageProviderConfig) -> Option<u64> {
372 if !provider.refresh_on_request {
373 return None;
374 }
375
376 let interval_secs = provider
377 .poll_interval_secs
378 .unwrap_or(DEFAULT_POLL_INTERVAL_SECS);
379 if interval_secs == 0 {
380 return None;
381 }
382 Some(interval_secs.max(MIN_POLL_INTERVAL_SECS))
383}
384
385fn remaining_poll_cooldown(last: Instant, interval_secs: u64, now: Instant) -> Option<Duration> {
386 let interval = Duration::from_secs(interval_secs);
387 let elapsed = now.saturating_duration_since(last);
388 interval.checked_sub(elapsed).filter(|d| !d.is_zero())
389}
390
391fn usage_providers_path() -> std::path::PathBuf {
392 proxy_home_dir().join("usage_providers.json")
393}
394
395fn service_manager<'a>(cfg: &'a ProxyConfig, service_name: &str) -> &'a ServiceConfigManager {
396 match service_name {
397 "claude" => &cfg.claude,
398 _ => &cfg.codex,
399 }
400}
401
402fn default_provider_config(
403 id: &str,
404 kind: ProviderKind,
405 domains: Vec<&str>,
406 endpoint: &str,
407 extract: UsageProviderExtractConfig,
408) -> UsageProviderConfig {
409 UsageProviderConfig {
410 id: id.to_string(),
411 kind,
412 domains: domains.into_iter().map(str::to_string).collect(),
413 endpoint: endpoint.to_string(),
414 token_env: None,
415 require_token_env: false,
416 poll_interval_secs: Some(DEFAULT_POLL_INTERVAL_SECS),
417 refresh_on_request: true,
418 trust_exhaustion_for_routing: true,
419 headers: BTreeMap::new(),
420 variables: BTreeMap::new(),
421 extract,
422 }
423}
424
425fn default_rightcode_provider_config(id: &str) -> UsageProviderConfig {
426 let mut provider = default_provider_config(
427 id,
428 ProviderKind::RightCodeAccountSummary,
429 vec!["www.right.codes", "right.codes"],
430 "https://www.right.codes/account/summary",
431 UsageProviderExtractConfig::default(),
432 );
433 provider.trust_exhaustion_for_routing = false;
437 provider
438}
439
440fn host_from_base_url(base_url: &str) -> Option<String> {
441 reqwest::Url::parse(base_url)
442 .ok()
443 .and_then(|url| url.host_str().map(|host| host.to_ascii_lowercase()))
444}
445
446fn is_official_openai_base_url(base_url: &str) -> bool {
447 host_from_base_url(base_url).as_deref() == Some("api.openai.com")
448}
449
450fn is_rightcode_base_url(base_url: &str) -> bool {
451 matches!(
452 host_from_base_url(base_url).as_deref(),
453 Some("www.right.codes" | "right.codes")
454 )
455}
456
457fn provider_id_component(value: &str) -> String {
458 let component = value
459 .chars()
460 .map(|ch| {
461 if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '.' {
462 ch
463 } else {
464 '-'
465 }
466 })
467 .collect::<String>()
468 .trim_matches('-')
469 .to_string();
470 if component.is_empty() {
471 "station".to_string()
472 } else {
473 component
474 }
475}
476
477fn auto_provider_id(target: &UsageProviderTarget) -> String {
478 if let Some(provider_id) = target
479 .provider_id
480 .as_deref()
481 .map(str::trim)
482 .filter(|value| !value.is_empty())
483 {
484 return provider_id.to_string();
485 }
486 format!(
487 "{}{}:{}",
488 AUTO_PROVIDER_ID_PREFIX,
489 provider_id_component(&target.upstream.station_name),
490 target.upstream.index
491 )
492}
493
494fn auto_usage_provider(target: &UsageProviderTarget, kind: ProviderKind) -> UsageProviderConfig {
495 let mut provider = UsageProviderConfig {
496 id: auto_provider_id(target),
497 kind,
498 domains: host_from_base_url(&target.base_url)
499 .into_iter()
500 .collect::<Vec<_>>(),
501 endpoint: String::new(),
502 token_env: None,
503 require_token_env: false,
504 poll_interval_secs: Some(DEFAULT_POLL_INTERVAL_SECS),
505 refresh_on_request: true,
506 trust_exhaustion_for_routing: true,
507 headers: BTreeMap::new(),
508 variables: BTreeMap::new(),
509 extract: UsageProviderExtractConfig::default(),
510 };
511 if matches!(kind, ProviderKind::RightCodeAccountSummary) {
512 provider.trust_exhaustion_for_routing = false;
513 }
514 provider
515}
516
517fn auto_target_matches_provider_id_filter(
518 target: &UsageProviderTarget,
519 provider_id_filter: Option<&str>,
520) -> bool {
521 match provider_id_filter {
522 Some(filter) => auto_provider_id(target) == filter,
523 None => true,
524 }
525}
526
527fn first_auto_probe_kind(target: &UsageProviderTarget) -> ProviderKind {
528 if is_rightcode_base_url(&target.base_url) {
529 ProviderKind::RightCodeAccountSummary
530 } else {
531 ProviderKind::Sub2ApiUsage
532 }
533}
534
535fn auto_probe_target_key(target: &UsageProviderTarget) -> AutoProbeTargetKey {
536 AutoProbeTargetKey {
537 station_name: target.upstream.station_name.clone(),
538 upstream_index: target.upstream.index,
539 provider_endpoint_key: target
540 .upstream
541 .provider_endpoint
542 .as_ref()
543 .map(ProviderEndpointKey::stable_key),
544 base_url: normalized_balance_base_url(&target.base_url)
545 .unwrap_or_else(|| target.base_url.clone()),
546 }
547}
548
549fn auto_probe_kind_order(provider_id: &str, target: &UsageProviderTarget) -> Vec<ProviderKind> {
550 let now = Instant::now();
551 let target_key = auto_probe_target_key(target);
552 let mut ordered = Vec::new();
553 if let Some(kind) = remembered_auto_probe_kind(provider_id) {
554 ordered.push(kind);
555 }
556 ordered.push(first_auto_probe_kind(target));
557 ordered.extend(AUTO_PROBE_KINDS);
558
559 let mut seen = HashSet::new();
560 ordered
561 .into_iter()
562 .filter(|kind| {
563 if *kind == ProviderKind::RightCodeAccountSummary
564 && !is_rightcode_base_url(&target.base_url)
565 {
566 return false;
567 }
568 seen.insert(*kind)
569 && !auto_probe_kind_failure_active(provider_id, &target_key, *kind, now)
570 })
571 .collect()
572}
573
574fn remembered_auto_probe_kind(provider_id: &str) -> Option<ProviderKind> {
575 AUTO_PROBE_KIND_HINTS
576 .get_or_init(|| Mutex::new(HashMap::new()))
577 .lock()
578 .ok()
579 .and_then(|map| map.get(provider_id).copied())
580}
581
582fn remember_auto_probe_kind_success(
583 provider_id: &str,
584 target: &UsageProviderTarget,
585 kind: ProviderKind,
586) {
587 if let Ok(mut hints) = AUTO_PROBE_KIND_HINTS
588 .get_or_init(|| Mutex::new(HashMap::new()))
589 .lock()
590 {
591 hints.insert(provider_id.to_string(), kind);
592 }
593 if let Ok(mut failures) = AUTO_PROBE_KIND_FAILURES
594 .get_or_init(|| Mutex::new(HashMap::new()))
595 .lock()
596 {
597 failures.remove(&AutoProbeKindFailureKey {
598 provider_id: provider_id.to_string(),
599 target: auto_probe_target_key(target),
600 kind,
601 });
602 }
603 clear_usage_provider_target_suppression(provider_id, target);
604}
605
606fn remember_auto_probe_kind_failure(
607 provider_id: &str,
608 target: &UsageProviderTarget,
609 kind: ProviderKind,
610 now: Instant,
611) {
612 if let Ok(mut failures) = AUTO_PROBE_KIND_FAILURES
613 .get_or_init(|| Mutex::new(HashMap::new()))
614 .lock()
615 {
616 failures.insert(
617 AutoProbeKindFailureKey {
618 provider_id: provider_id.to_string(),
619 target: auto_probe_target_key(target),
620 kind,
621 },
622 now,
623 );
624 }
625}
626
627fn auto_probe_kind_failure_active(
628 provider_id: &str,
629 target: &AutoProbeTargetKey,
630 kind: ProviderKind,
631 now: Instant,
632) -> bool {
633 let key = AutoProbeKindFailureKey {
634 provider_id: provider_id.to_string(),
635 target: target.clone(),
636 kind,
637 };
638 let Ok(mut failures) = AUTO_PROBE_KIND_FAILURES
639 .get_or_init(|| Mutex::new(HashMap::new()))
640 .lock()
641 else {
642 return false;
643 };
644 let Some(failed_at) = failures.get(&key).copied() else {
645 return false;
646 };
647 if now.duration_since(failed_at) < AUTO_PROBE_KIND_FAILURE_TTL {
648 true
649 } else {
650 failures.remove(&key);
651 false
652 }
653}
654
655fn usage_provider_target_suppression_key(
656 provider_id: &str,
657 target: &UsageProviderTarget,
658) -> ProviderTargetSuppressionKey {
659 ProviderTargetSuppressionKey {
660 provider_id: provider_id.to_string(),
661 target: auto_probe_target_key(target),
662 }
663}
664
665fn remember_usage_provider_target_suppression(
666 provider_id: &str,
667 target: &UsageProviderTarget,
668 ttl: Duration,
669 reason: impl Into<String>,
670 routing_exhausted: bool,
671 now: Instant,
672) {
673 if let Ok(mut suppressions) = USAGE_PROVIDER_TARGET_SUPPRESSIONS
674 .get_or_init(|| Mutex::new(HashMap::new()))
675 .lock()
676 {
677 suppressions.insert(
678 usage_provider_target_suppression_key(provider_id, target),
679 ProviderTargetSuppression {
680 until: now + ttl,
681 reason: reason.into(),
682 routing_exhausted,
683 },
684 );
685 }
686}
687
688fn clear_usage_provider_target_suppression(provider_id: &str, target: &UsageProviderTarget) {
689 if let Ok(mut suppressions) = USAGE_PROVIDER_TARGET_SUPPRESSIONS
690 .get_or_init(|| Mutex::new(HashMap::new()))
691 .lock()
692 {
693 suppressions.remove(&usage_provider_target_suppression_key(provider_id, target));
694 }
695}
696
697#[cfg(test)]
698fn clear_usage_provider_target_suppressions_for_provider(provider_id: &str) {
699 if let Some(suppressions) = USAGE_PROVIDER_TARGET_SUPPRESSIONS.get()
700 && let Ok(mut suppressions) = suppressions.lock()
701 {
702 suppressions.retain(|key, _| key.provider_id != provider_id);
703 }
704}
705
706fn usage_provider_target_suppression_active(
707 provider_id: &str,
708 target: &UsageProviderTarget,
709 now: Instant,
710) -> Option<ProviderTargetSuppression> {
711 let key = usage_provider_target_suppression_key(provider_id, target);
712 let Ok(mut suppressions) = USAGE_PROVIDER_TARGET_SUPPRESSIONS
713 .get_or_init(|| Mutex::new(HashMap::new()))
714 .lock()
715 else {
716 return None;
717 };
718 let suppression = suppressions.get(&key).cloned()?;
719 if now < suppression.until {
720 Some(suppression)
721 } else {
722 suppressions.remove(&key);
723 None
724 }
725}
726
727fn auto_openai_official_provider(target: &UsageProviderTarget) -> UsageProviderConfig {
728 let mut provider = auto_usage_provider(target, ProviderKind::OpenAiOrganizationCosts);
729 provider.token_env = Some("OPENAI_ADMIN_KEY".to_string());
730 provider.require_token_env = true;
731 provider.refresh_on_request = false;
732 provider.trust_exhaustion_for_routing = false;
733 provider
734}
735
736fn default_providers() -> UsageProvidersFile {
737 let openrouter_extract = UsageProviderExtractConfig {
738 monthly_budget_paths: vec!["data.total_credits".to_string()],
739 monthly_spent_paths: vec!["data.total_usage".to_string()],
740 derive_remaining_from_budget_and_spent: true,
741 ..Default::default()
742 };
743
744 let novita_extract = UsageProviderExtractConfig {
745 remaining_balance_paths: vec!["availableBalance".to_string()],
746 remaining_divisor: Some(10_000),
747 ..Default::default()
748 };
749
750 let mut openai_official = default_provider_config(
751 "openai-official-costs",
752 ProviderKind::OpenAiOrganizationCosts,
753 vec!["api.openai.com"],
754 "https://api.openai.com/v1/organization/costs?start_time={{unix_days_ago:30}}&limit=30",
755 UsageProviderExtractConfig::default(),
756 );
757 openai_official.token_env = Some("OPENAI_ADMIN_KEY".to_string());
758 openai_official.require_token_env = true;
759 openai_official.refresh_on_request = false;
760 openai_official.trust_exhaustion_for_routing = false;
761
762 UsageProvidersFile {
763 providers: vec![
764 default_rightcode_provider_config("rightcode"),
765 default_provider_config(
766 "packycode",
767 ProviderKind::BudgetHttpJson,
768 vec!["packycode.com"],
769 "https://www.packycode.com/api/backend/users/info",
770 UsageProviderExtractConfig::default(),
771 ),
772 default_provider_config(
773 "yescode",
774 ProviderKind::YescodeProfile,
775 vec!["yes.vg"],
777 "https://co.yes.vg/api/v1/auth/profile",
778 UsageProviderExtractConfig::default(),
779 ),
780 default_provider_config(
781 "deepseek",
782 ProviderKind::OpenAiBalanceHttpJson,
783 vec!["api.deepseek.com"],
784 "https://api.deepseek.com/user/balance",
785 UsageProviderExtractConfig::default(),
786 ),
787 default_provider_config(
788 "stepfun",
789 ProviderKind::OpenAiBalanceHttpJson,
790 vec!["api.stepfun.ai", "api.stepfun.com"],
791 "https://api.stepfun.com/v1/accounts",
792 UsageProviderExtractConfig::default(),
793 ),
794 default_provider_config(
795 "siliconflow",
796 ProviderKind::OpenAiBalanceHttpJson,
797 vec!["api.siliconflow.cn", "api.siliconflow.com"],
798 "{{base_url}}/v1/user/info",
799 UsageProviderExtractConfig::default(),
800 ),
801 default_provider_config(
802 "openrouter",
803 ProviderKind::OpenAiBalanceHttpJson,
804 vec!["openrouter.ai"],
805 "https://openrouter.ai/api/v1/credits",
806 openrouter_extract,
807 ),
808 default_provider_config(
809 "novita",
810 ProviderKind::OpenAiBalanceHttpJson,
811 vec!["api.novita.ai"],
812 "https://api.novita.ai/v3/user/balance",
813 novita_extract,
814 ),
815 openai_official,
816 ],
817 }
818}
819
820fn load_providers() -> UsageProvidersFile {
821 let path = usage_providers_path();
822 if let Ok(text) = std::fs::read_to_string(&path)
823 && let Ok(file) = serde_json::from_str::<UsageProvidersFile>(&text)
824 {
825 return file;
826 }
827
828 let default = default_providers();
830 if let Ok(text) = serde_json::to_string_pretty(&default) {
831 if let Some(parent) = path.parent() {
832 let _ = std::fs::create_dir_all(parent);
833 }
834 let _ = std::fs::write(&path, text);
835 }
836 default
837}
838
839fn domain_matches(base_url: &str, domains: &[String]) -> bool {
840 let url = match reqwest::Url::parse(base_url) {
841 Ok(u) => u,
842 Err(_) => return false,
843 };
844 let host = match url.host_str() {
845 Some(h) => h,
846 None => return false,
847 };
848 let host = host.to_ascii_lowercase();
849 for d in domains {
850 let domain = d.trim().to_ascii_lowercase();
851 if host == domain || host.ends_with(&format!(".{}", domain)) {
852 return true;
853 }
854 }
855 false
856}
857
858fn matching_provider_targets(
859 cfg: &ProxyConfig,
860 service_name: &str,
861 provider: &UsageProviderConfig,
862 station_name_filter: Option<&str>,
863) -> Vec<UsageProviderTarget> {
864 let mut stations: Vec<_> = service_manager(cfg, service_name)
865 .stations()
866 .iter()
867 .collect();
868 stations.sort_by_key(|(name, _)| name.as_str());
869
870 let mut targets = Vec::new();
871 for (station_name, service) in stations {
872 if station_name_filter.is_some_and(|filter| filter != station_name.as_str()) {
873 continue;
874 }
875 for (index, upstream) in service.upstreams.iter().enumerate() {
876 if domain_matches(&upstream.base_url, &provider.domains) {
877 targets.push(UsageProviderTarget {
878 upstream: UpstreamRef {
879 station_name: station_name.clone(),
880 index,
881 provider_endpoint: upstream.provider_endpoint_key(service_name),
882 },
883 base_url: upstream.base_url.clone(),
884 provider_id: upstream.tags.get("provider_id").cloned(),
885 });
886 }
887 }
888 }
889
890 targets
891}
892
893fn usage_provider_targets(
894 cfg: &ProxyConfig,
895 service_name: &str,
896 station_name_filter: Option<&str>,
897) -> Vec<UsageProviderTarget> {
898 let mut stations: Vec<_> = service_manager(cfg, service_name)
899 .stations()
900 .iter()
901 .collect();
902 stations.sort_by_key(|(name, _)| name.as_str());
903
904 let mut targets = Vec::new();
905 for (station_name, service) in stations {
906 if station_name_filter.is_some_and(|filter| filter != station_name.as_str()) {
907 continue;
908 }
909 for (index, upstream) in service.upstreams.iter().enumerate() {
910 targets.push(UsageProviderTarget {
911 upstream: UpstreamRef {
912 station_name: station_name.clone(),
913 index,
914 provider_endpoint: upstream.provider_endpoint_key(service_name),
915 },
916 base_url: upstream.base_url.clone(),
917 provider_id: upstream.tags.get("provider_id").cloned(),
918 });
919 }
920 }
921
922 targets
923}
924
925fn target_key(target: &UsageProviderTarget) -> UsageProviderTargetKey {
926 UsageProviderTargetKey {
927 station_name: target.upstream.station_name.clone(),
928 upstream_index: target.upstream.index,
929 }
930}
931
932fn enqueue_request_balance_refresh(key: ProviderEndpointKey) -> Option<Duration> {
933 let now = Instant::now();
934 let queue = REQUEST_BALANCE_QUEUE.get_or_init(|| Mutex::new(HashMap::new()));
935 let mut queue = match queue.lock() {
936 Ok(queue) => queue,
937 Err(_) => return None,
938 };
939
940 match queue.get(&key).copied() {
941 Some(due_at) if due_at > now => None,
942 Some(_) => Some(Duration::ZERO),
943 None => {
944 queue.insert(key, now + REQUEST_BALANCE_REFRESH_DELAY);
945 Some(REQUEST_BALANCE_REFRESH_DELAY)
946 }
947 }
948}
949
950fn schedule_request_balance_refresh_at(key: ProviderEndpointKey, due_at: Instant) {
951 let queue = REQUEST_BALANCE_QUEUE.get_or_init(|| Mutex::new(HashMap::new()));
952 if let Ok(mut queue) = queue.lock() {
953 queue.insert(key, due_at);
954 }
955}
956
957fn take_request_balance_refresh_if_due(key: &ProviderEndpointKey) -> RequestBalanceQueueDue {
958 let now = Instant::now();
959 let queue = REQUEST_BALANCE_QUEUE.get_or_init(|| Mutex::new(HashMap::new()));
960 let mut queue = match queue.lock() {
961 Ok(queue) => queue,
962 Err(_) => return RequestBalanceQueueDue::Missing,
963 };
964
965 match queue.get(key).copied() {
966 Some(due_at) if due_at <= now => {
967 queue.remove(key);
968 RequestBalanceQueueDue::Due
969 }
970 Some(due_at) => RequestBalanceQueueDue::NotDue(due_at.saturating_duration_since(now)),
971 None => RequestBalanceQueueDue::Missing,
972 }
973}
974
975#[cfg(test)]
976pub fn request_balance_refresh_queued_for_provider_endpoint(
977 provider_endpoint: &ProviderEndpointKey,
978) -> bool {
979 let Some(queue) = REQUEST_BALANCE_QUEUE.get() else {
980 return false;
981 };
982 match queue.lock() {
983 Ok(guard) => guard.contains_key(provider_endpoint),
984 Err(error) => error.into_inner().contains_key(provider_endpoint),
985 }
986}
987
988fn usage_provider_target_for_provider_endpoint(
989 cfg: &ProxyConfig,
990 service_name: &str,
991 provider_endpoint: &ProviderEndpointKey,
992) -> Option<UsageProviderTarget> {
993 service_manager(cfg, service_name)
994 .stations()
995 .iter()
996 .filter_map(|(station_name, service)| {
997 service
998 .upstreams
999 .iter()
1000 .enumerate()
1001 .find_map(|(index, upstream)| {
1002 let upstream_endpoint = upstream.provider_endpoint_key(service_name)?;
1003 if upstream_endpoint != *provider_endpoint {
1004 return None;
1005 }
1006 Some(UsageProviderTarget {
1007 upstream: UpstreamRef {
1008 station_name: station_name.clone(),
1009 index,
1010 provider_endpoint: Some(upstream_endpoint),
1011 },
1012 base_url: upstream.base_url.clone(),
1013 provider_id: upstream.tags.get("provider_id").cloned(),
1014 })
1015 })
1016 })
1017 .next()
1018}
1019
1020trait UsageProviderUpstreamIdentityExt {
1021 fn provider_endpoint_key(&self, service_name: &str) -> Option<ProviderEndpointKey>;
1022}
1023
1024impl UsageProviderUpstreamIdentityExt for crate::config::UpstreamConfig {
1025 fn provider_endpoint_key(&self, service_name: &str) -> Option<ProviderEndpointKey> {
1026 let provider_id = self.tags.get("provider_id")?.trim();
1027 let endpoint_id = self.tags.get("endpoint_id")?.trim();
1028 if provider_id.is_empty() || endpoint_id.is_empty() {
1029 return None;
1030 }
1031 Some(ProviderEndpointKey::new(
1032 service_name.to_string(),
1033 provider_id.to_string(),
1034 endpoint_id.to_string(),
1035 ))
1036 }
1037}
1038
1039#[cfg(test)]
1040fn configured_target_keys(
1041 cfg: &ProxyConfig,
1042 service_name: &str,
1043 providers: &[UsageProviderConfig],
1044 station_name_filter: Option<&str>,
1045) -> HashSet<UsageProviderTargetKey> {
1046 providers
1047 .iter()
1048 .flat_map(|provider| {
1049 matching_provider_targets(cfg, service_name, provider, station_name_filter)
1050 })
1051 .map(|target| target_key(&target))
1052 .collect()
1053}
1054
1055fn resolve_token(
1056 provider: &UsageProviderConfig,
1057 upstreams: &[UpstreamRef],
1058 cfg: &ProxyConfig,
1059 service_name: &str,
1060) -> Option<String> {
1061 if let Some(env_name) = &provider.token_env
1063 && let Ok(v) = std::env::var(env_name)
1064 && !v.trim().is_empty()
1065 {
1066 return Some(v);
1067 }
1068
1069 if provider.require_token_env {
1070 return None;
1071 }
1072
1073 for uref in upstreams {
1075 if let Some(service) = service_manager(cfg, service_name).station(&uref.station_name)
1076 && let Some(up) = service.upstreams.get(uref.index)
1077 {
1078 if let Some(token) = up.auth.resolve_auth_token() {
1079 return Some(token);
1080 }
1081 if let Some(token) = up.auth.resolve_api_key() {
1082 return Some(token);
1083 }
1084 }
1085 }
1086 None
1087}
1088
1089fn normalized_balance_base_url(base_url: &str) -> Option<String> {
1090 let mut url = reqwest::Url::parse(base_url).ok()?;
1091 url.set_query(None);
1092 url.set_fragment(None);
1093 let path = url.path().trim_end_matches('/').to_string();
1094 if path.eq_ignore_ascii_case("/v1") {
1095 url.set_path("");
1096 } else if path.to_ascii_lowercase().ends_with("/v1") {
1097 let new_path = &path[..path.len().saturating_sub(3)];
1098 url.set_path(if new_path.is_empty() { "/" } else { new_path });
1099 }
1100 Some(url.as_str().trim_end_matches('/').to_string())
1101}
1102
1103fn base_path_prefixes(base_url: &str) -> Vec<String> {
1104 let Some(normalized) = normalized_balance_base_url(base_url) else {
1105 return Vec::new();
1106 };
1107 let Ok(url) = reqwest::Url::parse(&normalized) else {
1108 return Vec::new();
1109 };
1110 let parts = url
1111 .path_segments()
1112 .map(|segments| {
1113 segments
1114 .filter(|segment| !segment.is_empty())
1115 .collect::<Vec<_>>()
1116 })
1117 .unwrap_or_default();
1118 let mut prefixes = Vec::new();
1119 for len in (1..=parts.len()).rev() {
1120 prefixes.push(format!("/{}", parts[..len].join("/")));
1121 }
1122 if prefixes.is_empty() {
1123 prefixes.push("/".to_string());
1124 }
1125 prefixes
1126}
1127
1128fn path_prefixes_match(provider_prefixes: &[String], available_prefixes: &[String]) -> bool {
1129 if provider_prefixes.is_empty() || available_prefixes.is_empty() {
1130 return false;
1131 }
1132 provider_prefixes.iter().any(|provider_prefix| {
1133 available_prefixes.iter().any(|available_prefix| {
1134 provider_prefix == available_prefix
1135 || provider_prefix
1136 .strip_prefix(available_prefix)
1137 .is_some_and(|suffix| suffix.starts_with('/'))
1138 })
1139 })
1140}
1141
1142fn render_provider_template(
1143 template: &str,
1144 base_url: &str,
1145 upstream_base_url: &str,
1146 token: &str,
1147 variables: &BTreeMap<String, String>,
1148) -> String {
1149 let mut out = template
1150 .replace("{{baseUrl}}", base_url)
1151 .replace("{{base_url}}", base_url)
1152 .replace("{{upstreamBaseUrl}}", upstream_base_url)
1153 .replace("{{upstream_base_url}}", upstream_base_url)
1154 .replace("{{apiKey}}", token)
1155 .replace("{{accessToken}}", token)
1156 .replace("{{token}}", token);
1157
1158 out = out
1159 .replace("{{unix_now}}", &unix_now_secs().to_string())
1160 .replace("{{unix_now_ms}}", &unix_now_ms().to_string());
1161
1162 while let Some(start) = out.find("{{unix_days_ago:") {
1163 let Some(end_offset) = out[start..].find("}}") else {
1164 break;
1165 };
1166 let end = start + end_offset + 2;
1167 let days_str = out[start + "{{unix_days_ago:".len()..end - 2].trim();
1168 let replacement = days_str
1169 .parse::<u64>()
1170 .ok()
1171 .map(|days| unix_now_secs().saturating_sub(days.saturating_mul(24 * 60 * 60)))
1172 .map(|secs| secs.to_string())
1173 .unwrap_or_default();
1174 out.replace_range(start..end, &replacement);
1175 }
1176
1177 while let Some(start) = out.find("{{env:") {
1178 let Some(end_offset) = out[start..].find("}}") else {
1179 break;
1180 };
1181 let end = start + end_offset + 2;
1182 let env_name = out[start + 6..end - 2].trim();
1183 let value = std::env::var(env_name).unwrap_or_default();
1184 out.replace_range(start..end, &value);
1185 }
1186
1187 for (name, value_template) in variables {
1188 let value = render_provider_template(
1189 value_template,
1190 base_url,
1191 upstream_base_url,
1192 token,
1193 &BTreeMap::new(),
1194 );
1195 out = out.replace(&format!("{{{{{name}}}}}"), &value);
1196 }
1197
1198 out
1199}
1200
1201fn resolve_endpoint(
1202 provider: &UsageProviderConfig,
1203 upstream_base_url: &str,
1204 token: &str,
1205) -> Result<String> {
1206 let base_url = normalized_balance_base_url(upstream_base_url)
1207 .ok_or_else(|| anyhow::anyhow!("invalid upstream base_url for balance endpoint"))?;
1208 let endpoint = if provider.endpoint.trim().is_empty() {
1209 provider
1210 .kind
1211 .default_endpoint()
1212 .unwrap_or_default()
1213 .to_string()
1214 } else {
1215 provider.endpoint.trim().to_string()
1216 };
1217 if endpoint.is_empty() {
1218 anyhow::bail!(
1219 "usage provider '{}' has no endpoint and kind {:?} has no default endpoint",
1220 provider.id,
1221 provider.kind
1222 );
1223 }
1224
1225 let rendered = render_provider_template(
1226 &endpoint,
1227 &base_url,
1228 upstream_base_url,
1229 token,
1230 &provider.variables,
1231 );
1232 if rendered.starts_with("http://") || rendered.starts_with("https://") {
1233 return Ok(rendered);
1234 }
1235
1236 let path = if rendered.starts_with('/') {
1237 rendered
1238 } else {
1239 format!("/{rendered}")
1240 };
1241 Ok(format!("{base_url}{path}"))
1242}
1243
1244fn endpoint_origin(endpoint: &str) -> String {
1245 reqwest::Url::parse(endpoint)
1246 .ok()
1247 .and_then(|url| {
1248 let host = url.host_str()?;
1249 let origin = match url.port() {
1250 Some(port) => format!("{}://{}:{}", url.scheme(), host, port),
1251 None => format!("{}://{}", url.scheme(), host),
1252 };
1253 Some(origin)
1254 })
1255 .unwrap_or_else(|| "unknown-origin".to_string())
1256}
1257
1258async fn poll_provider_http_json(
1259 client: &Client,
1260 provider: &UsageProviderConfig,
1261 upstream_base_url: &str,
1262 token: &str,
1263) -> Result<serde_json::Value> {
1264 let endpoint = resolve_endpoint(provider, upstream_base_url, token)?;
1265 let origin = endpoint_origin(&endpoint);
1266 let base_url = normalized_balance_base_url(upstream_base_url).unwrap_or_default();
1267 let mut req = client
1268 .get(endpoint)
1269 .timeout(BALANCE_HTTP_REQUEST_TIMEOUT)
1270 .header("Accept", "application/json")
1271 .header(
1272 "User-Agent",
1273 concat!("codex-helper/", env!("CARGO_PKG_VERSION")),
1274 );
1275
1276 match provider.kind {
1277 ProviderKind::YescodeProfile => {
1278 req = req.header("X-API-Key", token);
1279 }
1280 _ => {
1281 req = req.header("Authorization", format!("Bearer {}", token));
1282 }
1283 }
1284
1285 for (name, template) in &provider.headers {
1286 let value = render_provider_template(
1287 template,
1288 &base_url,
1289 upstream_base_url,
1290 token,
1291 &provider.variables,
1292 );
1293 if !value.trim().is_empty() {
1294 req = req.header(name.as_str(), value);
1295 }
1296 }
1297
1298 let resp = req.send().await.with_context(|| {
1299 format!(
1300 "usage provider request failed for {} via {:?}",
1301 origin, provider.kind
1302 )
1303 })?;
1304
1305 let status = resp.status();
1306 let content_type = resp
1307 .headers()
1308 .get(reqwest::header::CONTENT_TYPE)
1309 .and_then(|value| value.to_str().ok())
1310 .map(str::to_string)
1311 .unwrap_or_else(|| "unknown".to_string());
1312 if !status.is_success() {
1313 let text = resp.text().await.with_context(|| {
1314 format!(
1315 "usage provider error response read failed from {} via {:?}",
1316 origin, provider.kind
1317 )
1318 })?;
1319 let detail = usage_provider_http_error_detail(&text)
1320 .map(|detail| format!(": {detail}"))
1321 .unwrap_or_default();
1322 anyhow::bail!(
1323 "usage provider HTTP {} from {} via {:?}{}",
1324 status,
1325 origin,
1326 provider.kind,
1327 detail
1328 );
1329 }
1330 let text = resp.text().await.with_context(|| {
1331 format!(
1332 "usage provider response read failed from {} via {:?}",
1333 origin, provider.kind
1334 )
1335 })?;
1336 serde_json::from_str(&text).with_context(|| {
1337 format!(
1338 "usage provider returned non-JSON response from {} via {:?} (content-type {}, {} bytes)",
1339 origin,
1340 provider.kind,
1341 content_type,
1342 text.len()
1343 )
1344 })
1345}
1346
1347fn truncate_error_detail(value: &str, max_chars: usize) -> String {
1348 let mut out = String::new();
1349 for (idx, ch) in value.chars().enumerate() {
1350 if idx >= max_chars {
1351 out.push_str("...");
1352 return out;
1353 }
1354 out.push(ch);
1355 }
1356 out
1357}
1358
1359fn compact_error_detail(value: &str) -> Option<String> {
1360 let compact = value.split_whitespace().collect::<Vec<_>>().join(" ");
1361 if compact.is_empty() {
1362 None
1363 } else {
1364 Some(truncate_error_detail(
1365 &compact,
1366 BALANCE_HTTP_ERROR_BODY_LIMIT,
1367 ))
1368 }
1369}
1370
1371fn json_error_detail(value: &serde_json::Value) -> Option<String> {
1372 let code = first_string_from_paths(
1373 value,
1374 &[
1375 "code",
1376 "error.code",
1377 "error.type",
1378 "type",
1379 "data.code",
1380 "data.error.code",
1381 ],
1382 );
1383 let message = first_string_from_paths(
1384 value,
1385 &[
1386 "message",
1387 "msg",
1388 "detail",
1389 "error.message",
1390 "error_description",
1391 "error",
1392 "data.message",
1393 "data.error.message",
1394 ],
1395 );
1396
1397 match (code, message) {
1398 (Some(code), Some(message)) if !message.eq_ignore_ascii_case(&code) => {
1399 Some(format!("{code}: {message}"))
1400 }
1401 (Some(code), _) => Some(code),
1402 (_, Some(message)) => Some(message),
1403 _ => None,
1404 }
1405}
1406
1407fn usage_provider_http_error_detail(text: &str) -> Option<String> {
1408 let trimmed = text.trim();
1409 if trimmed.is_empty() {
1410 return None;
1411 }
1412
1413 if let Ok(value) = serde_json::from_str::<serde_json::Value>(trimmed)
1414 && let Some(detail) = json_error_detail(&value)
1415 {
1416 return compact_error_detail(&detail);
1417 }
1418
1419 compact_error_detail(trimmed)
1420}
1421
1422fn amount_from_json(value: &serde_json::Value) -> Option<UsdAmount> {
1423 let raw = match value {
1424 serde_json::Value::Number(number) => number.to_string(),
1425 serde_json::Value::String(text) => text.trim().to_string(),
1426 _ => return None,
1427 };
1428 UsdAmount::from_decimal_str(raw.as_str())
1429}
1430
1431fn decimal_string_from_json(value: &serde_json::Value) -> Option<String> {
1432 match value {
1433 serde_json::Value::Number(number) => Some(number.to_string()),
1434 serde_json::Value::String(text) => {
1435 let text = text.trim();
1436 if text.is_empty() {
1437 None
1438 } else {
1439 Some(text.to_string())
1440 }
1441 }
1442 _ => None,
1443 }
1444}
1445
1446fn amount_from_json_with_divisor(
1447 value: &serde_json::Value,
1448 divisor: Option<u64>,
1449) -> Option<UsdAmount> {
1450 let amount = amount_from_json(value)?;
1451 match divisor {
1452 Some(divisor) => amount.checked_div_u64(divisor),
1453 None => Some(amount),
1454 }
1455}
1456
1457fn json_value_at_path<'a>(
1458 value: &'a serde_json::Value,
1459 path: &str,
1460) -> Option<&'a serde_json::Value> {
1461 let mut current = value;
1462 for segment in path
1463 .split('.')
1464 .map(str::trim)
1465 .filter(|segment| !segment.is_empty())
1466 {
1467 current = match current {
1468 serde_json::Value::Array(items) => {
1469 let index = segment.parse::<usize>().ok()?;
1470 items.get(index)?
1471 }
1472 _ => current.get(segment)?,
1473 };
1474 }
1475 Some(current)
1476}
1477
1478fn first_amount_from_paths(
1479 value: &serde_json::Value,
1480 custom_paths: &[String],
1481 default_paths: &[&str],
1482 divisor: Option<u64>,
1483) -> Option<UsdAmount> {
1484 custom_paths
1485 .iter()
1486 .map(String::as_str)
1487 .chain(default_paths.iter().copied())
1488 .find_map(|path| {
1489 json_value_at_path(value, path)
1490 .and_then(|value| amount_from_json_with_divisor(value, divisor))
1491 })
1492}
1493
1494fn bool_from_json(value: &serde_json::Value) -> Option<bool> {
1495 match value {
1496 serde_json::Value::Bool(value) => Some(*value),
1497 serde_json::Value::Number(number) => number.as_i64().map(|value| value != 0),
1498 serde_json::Value::String(text) => match text.trim().to_ascii_lowercase().as_str() {
1499 "true" | "yes" | "1" | "exhausted" => Some(true),
1500 "false" | "no" | "0" | "ok" => Some(false),
1501 _ => None,
1502 },
1503 _ => None,
1504 }
1505}
1506
1507fn first_bool_from_paths(
1508 value: &serde_json::Value,
1509 custom_paths: &[String],
1510 default_paths: &[&str],
1511) -> Option<bool> {
1512 custom_paths
1513 .iter()
1514 .map(String::as_str)
1515 .chain(default_paths.iter().copied())
1516 .find_map(|path| json_value_at_path(value, path).and_then(bool_from_json))
1517}
1518
1519fn first_decimal_string_from_paths(
1520 value: &serde_json::Value,
1521 default_paths: &[&str],
1522) -> Option<String> {
1523 default_paths
1524 .iter()
1525 .copied()
1526 .find_map(|path| json_value_at_path(value, path).and_then(decimal_string_from_json))
1527}
1528
1529fn string_from_json(value: &serde_json::Value) -> Option<String> {
1530 match value {
1531 serde_json::Value::String(text) => {
1532 let text = text.trim();
1533 if text.is_empty() {
1534 None
1535 } else {
1536 Some(text.to_string())
1537 }
1538 }
1539 _ => None,
1540 }
1541}
1542
1543fn first_string_from_paths(value: &serde_json::Value, default_paths: &[&str]) -> Option<String> {
1544 default_paths
1545 .iter()
1546 .copied()
1547 .find_map(|path| json_value_at_path(value, path).and_then(string_from_json))
1548}
1549
1550fn u64_from_json(value: &serde_json::Value) -> Option<u64> {
1551 match value {
1552 serde_json::Value::Number(number) => number.as_u64(),
1553 serde_json::Value::String(text) => {
1554 let text = text.trim();
1555 if text.is_empty() {
1556 None
1557 } else {
1558 text.parse::<u64>().ok()
1559 }
1560 }
1561 _ => None,
1562 }
1563}
1564
1565fn seconds_from_json(value: &serde_json::Value) -> Option<u64> {
1566 match value {
1567 serde_json::Value::Number(number) => number.as_f64().map(|value| value.max(0.0) as u64),
1568 serde_json::Value::String(text) => {
1569 let text = text.trim();
1570 if text.is_empty() {
1571 None
1572 } else if let Ok(value) = text.parse::<f64>() {
1573 Some(value.max(0.0) as u64)
1574 } else {
1575 parse_timestamp_secs(text)
1576 }
1577 }
1578 _ => None,
1579 }
1580}
1581
1582fn first_secs_from_paths(value: &serde_json::Value, default_paths: &[&str]) -> Option<u64> {
1583 default_paths
1584 .iter()
1585 .copied()
1586 .find_map(|path| json_value_at_path(value, path).and_then(seconds_from_json))
1587}
1588
1589fn parse_timestamp_secs(value: &str) -> Option<u64> {
1590 parse_rfc3339_like_secs(value).or_else(|| {
1591 httpdate::parse_http_date(value).ok().and_then(|time| {
1592 time.duration_since(std::time::UNIX_EPOCH)
1593 .ok()
1594 .map(|duration| duration.as_secs())
1595 })
1596 })
1597}
1598
1599fn parse_rfc3339_like_secs(value: &str) -> Option<u64> {
1600 let value = value.trim();
1601 let datetime_sep = value.find('T').or_else(|| value.find(' '))?;
1602 let (datetime, offset_secs) = if let Some(datetime) = value.strip_suffix('Z') {
1603 (datetime, 0_i64)
1604 } else {
1605 let offset_pos = value[datetime_sep + 1..]
1606 .rfind(['+', '-'])
1607 .map(|pos| datetime_sep + 1 + pos)?;
1608 let (datetime, offset) = value.split_at(offset_pos);
1609 (datetime, parse_rfc3339_offset_secs(offset)?)
1610 };
1611
1612 let (date, time) = datetime.split_at(datetime_sep);
1613 let time = time.get(1..)?;
1614 let mut date_parts = date.split('-');
1615 let year = date_parts.next()?.parse::<i32>().ok()?;
1616 let month = date_parts.next()?.parse::<u32>().ok()?;
1617 let day = date_parts.next()?.parse::<u32>().ok()?;
1618 if date_parts.next().is_some() {
1619 return None;
1620 }
1621
1622 let mut time_parts = time.split(':');
1623 let hour = time_parts.next()?.parse::<u32>().ok()?;
1624 let minute = time_parts.next()?.parse::<u32>().ok()?;
1625 let second_raw = time_parts.next().unwrap_or("0");
1626 if time_parts.next().is_some() {
1627 return None;
1628 }
1629 let second = second_raw
1630 .split('.')
1631 .next()
1632 .and_then(|value| value.parse::<u32>().ok())?;
1633 if !(1..=12).contains(&month) || day == 0 || hour > 23 || minute > 59 || second > 60 {
1634 return None;
1635 }
1636
1637 let local_secs = days_from_civil(year, month, day)
1638 .checked_mul(86_400)?
1639 .checked_add(i64::from(hour) * 3_600 + i64::from(minute) * 60 + i64::from(second))?;
1640 local_secs
1641 .checked_sub(offset_secs)
1642 .and_then(|utc_secs| u64::try_from(utc_secs).ok())
1643}
1644
1645fn parse_rfc3339_offset_secs(offset: &str) -> Option<i64> {
1646 let sign = match offset.as_bytes().first().copied()? {
1647 b'+' => 1_i64,
1648 b'-' => -1_i64,
1649 _ => return None,
1650 };
1651 let raw = offset.get(1..)?;
1652 let (hours, minutes) = raw
1653 .split_once(':')
1654 .unwrap_or_else(|| raw.split_at(raw.len().min(2)));
1655 let hours = hours.parse::<i64>().ok()?;
1656 let minutes = minutes.parse::<i64>().ok()?;
1657 if hours > 23 || minutes > 59 {
1658 return None;
1659 }
1660 Some(sign * (hours * 3_600 + minutes * 60))
1661}
1662
1663fn days_from_civil(year: i32, month: u32, day: u32) -> i64 {
1664 let year = i64::from(year) - if month <= 2 { 1 } else { 0 };
1665 let era = if year >= 0 { year } else { year - 399 } / 400;
1666 let yoe = year - era * 400;
1667 let month = i64::from(month);
1668 let doy = (153 * (month + if month > 2 { -3 } else { 9 }) + 2) / 5 + i64::from(day) - 1;
1669 let doe = yoe * 365 + yoe / 4 - yoe / 100 + doy;
1670 era * 146_097 + doe - 719_468
1671}
1672
1673fn first_u64_from_paths(value: &serde_json::Value, default_paths: &[&str]) -> Option<u64> {
1674 default_paths
1675 .iter()
1676 .copied()
1677 .find_map(|path| json_value_at_path(value, path).and_then(u64_from_json))
1678}
1679
1680fn array_from_json_path<'a>(
1681 value: &'a serde_json::Value,
1682 path: &str,
1683) -> Option<&'a Vec<serde_json::Value>> {
1684 json_value_at_path(value, path).and_then(|value| value.as_array())
1685}
1686
1687fn amount_to_string(amount: UsdAmount) -> String {
1688 amount.format_usd()
1689}
1690
1691#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1692struct QuotaWindowSnapshot {
1693 period: &'static str,
1694 remaining: UsdAmount,
1695 used: UsdAmount,
1696 limit: UsdAmount,
1697}
1698
1699#[derive(Debug, Clone, PartialEq, Eq)]
1700struct RateLimitWindowSnapshot {
1701 period: String,
1702 reset_at_ms: Option<u64>,
1703}
1704
1705fn base_snapshot(
1706 provider: &UsageProviderConfig,
1707 upstream: &UpstreamRef,
1708 fetched_at_ms: u64,
1709 stale_after_ms: Option<u64>,
1710) -> ProviderBalanceSnapshot {
1711 let mut snapshot = ProviderBalanceSnapshot::new(
1712 provider.id.clone(),
1713 upstream.station_name.clone(),
1714 upstream.index,
1715 provider.kind.source_name(),
1716 fetched_at_ms,
1717 stale_after_ms,
1718 );
1719 if let Some(provider_endpoint) = &upstream.provider_endpoint {
1720 snapshot.provider_endpoint_key = Some(provider_endpoint.stable_key());
1721 }
1722 snapshot.exhaustion_affects_routing = provider.trust_exhaustion_for_routing;
1723 snapshot
1724}
1725
1726fn snapshot_error(
1727 provider: &UsageProviderConfig,
1728 upstream: &UpstreamRef,
1729 fetched_at_ms: u64,
1730 stale_after_ms: Option<u64>,
1731 message: impl Into<String>,
1732) -> ProviderBalanceSnapshot {
1733 base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms).with_error(message)
1734}
1735
1736fn budget_snapshot_from_json(
1737 provider: &UsageProviderConfig,
1738 upstream: &UpstreamRef,
1739 value: &serde_json::Value,
1740 fetched_at_ms: u64,
1741 stale_after_ms: Option<u64>,
1742) -> ProviderBalanceSnapshot {
1743 let monthly_budget = first_amount_from_paths(
1744 value,
1745 &provider.extract.monthly_budget_paths,
1746 &["monthly_budget_usd", "data.monthly_budget_usd"],
1747 provider.extract.monthly_budget_divisor,
1748 );
1749 let monthly_spent = first_amount_from_paths(
1750 value,
1751 &provider.extract.monthly_spent_paths,
1752 &["monthly_spent_usd", "data.monthly_spent_usd"],
1753 provider.extract.monthly_spent_divisor,
1754 );
1755 let exhausted = match (monthly_budget, monthly_spent) {
1756 (Some(budget), Some(spent)) if !budget.is_zero() => Some(spent >= budget),
1757 (Some(_), Some(_)) => Some(false),
1758 _ => None,
1759 };
1760
1761 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1762 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
1763 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1764 snapshot.exhausted = exhausted;
1765 snapshot.refresh_status(fetched_at_ms);
1766 snapshot
1767}
1768
1769fn yescode_snapshot_from_json(
1770 provider: &UsageProviderConfig,
1771 upstream: &UpstreamRef,
1772 value: &serde_json::Value,
1773 fetched_at_ms: u64,
1774 stale_after_ms: Option<u64>,
1775) -> ProviderBalanceSnapshot {
1776 let subscription_balance = first_amount_from_paths(
1777 value,
1778 &provider.extract.subscription_balance_paths,
1779 &["subscription_balance", "data.subscription_balance"],
1780 provider.extract.remaining_divisor,
1781 );
1782 let paygo_balance = first_amount_from_paths(
1783 value,
1784 &provider.extract.paygo_balance_paths,
1785 &[
1786 "pay_as_you_go_balance",
1787 "paygo_balance",
1788 "data.pay_as_you_go_balance",
1789 "data.paygo_balance",
1790 ],
1791 provider.extract.remaining_divisor,
1792 );
1793 let total_balance = match (subscription_balance, paygo_balance) {
1794 (Some(subscription), Some(paygo)) => Some(subscription.saturating_add(paygo)),
1795 (Some(subscription), None) => Some(subscription),
1796 (None, Some(paygo)) => Some(paygo),
1797 (None, None) => None,
1798 };
1799
1800 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1801 snapshot.total_balance_usd = total_balance.map(amount_to_string);
1802 snapshot.subscription_balance_usd = subscription_balance.map(amount_to_string);
1803 snapshot.paygo_balance_usd = paygo_balance.map(amount_to_string);
1804 snapshot.exhausted = total_balance.map(UsdAmount::is_zero);
1805 snapshot.refresh_status(fetched_at_ms);
1806 snapshot
1807}
1808
1809fn balance_http_snapshot_from_json(
1810 provider: &UsageProviderConfig,
1811 upstream: &UpstreamRef,
1812 value: &serde_json::Value,
1813 fetched_at_ms: u64,
1814 stale_after_ms: Option<u64>,
1815) -> ProviderBalanceSnapshot {
1816 let remaining_balance = first_amount_from_paths(
1817 value,
1818 &provider.extract.remaining_balance_paths,
1819 &[
1820 "balance",
1821 "remaining",
1822 "remain",
1823 "available",
1824 "available_balance",
1825 "credit",
1826 "credits",
1827 "total_balance",
1828 "total_balance_usd",
1829 "totalBalance",
1830 "availableBalance",
1831 "available_balance_usd",
1832 "balance_infos.0.total_balance",
1833 "data.balance",
1834 "data.remaining",
1835 "data.available",
1836 "data.available_balance",
1837 "data.credit",
1838 "data.credits",
1839 "data.total_balance",
1840 "data.totalBalance",
1841 "data.availableBalance",
1842 ],
1843 provider.extract.remaining_divisor,
1844 );
1845 let subscription_balance = first_amount_from_paths(
1846 value,
1847 &provider.extract.subscription_balance_paths,
1848 &[
1849 "subscription_balance",
1850 "subscription_balance_usd",
1851 "subscriptionBalance",
1852 "data.subscription_balance",
1853 "data.subscription_balance_usd",
1854 "data.subscriptionBalance",
1855 ],
1856 provider.extract.remaining_divisor,
1857 );
1858 let paygo_balance = first_amount_from_paths(
1859 value,
1860 &provider.extract.paygo_balance_paths,
1861 &[
1862 "pay_as_you_go_balance",
1863 "paygo_balance",
1864 "paygo",
1865 "paygoBalance",
1866 "chargeBalance",
1867 "voucherBalance",
1868 "data.pay_as_you_go_balance",
1869 "data.paygo_balance",
1870 "data.paygo",
1871 "data.paygoBalance",
1872 "data.chargeBalance",
1873 "data.voucherBalance",
1874 ],
1875 provider.extract.remaining_divisor,
1876 );
1877 let component_remaining = match (subscription_balance, paygo_balance) {
1878 (Some(subscription), Some(paygo)) => Some(subscription.saturating_add(paygo)),
1879 (Some(subscription), None) => Some(subscription),
1880 (None, Some(paygo)) => Some(paygo),
1881 (None, None) => None,
1882 };
1883 let monthly_spent = first_amount_from_paths(
1884 value,
1885 &provider.extract.monthly_spent_paths,
1886 &[
1887 "monthly_spent_usd",
1888 "spent",
1889 "used",
1890 "used_balance",
1891 "usedBalance",
1892 "total_usage",
1893 "data.monthly_spent_usd",
1894 "data.spent",
1895 "data.used",
1896 "data.used_balance",
1897 "data.usedBalance",
1898 "data.total_usage",
1899 ],
1900 provider.extract.monthly_spent_divisor,
1901 );
1902 let monthly_budget = first_amount_from_paths(
1903 value,
1904 &provider.extract.monthly_budget_paths,
1905 &[
1906 "monthly_budget_usd",
1907 "budget",
1908 "limit",
1909 "quota_total",
1910 "creditLimit",
1911 "total_credits",
1912 "data.monthly_budget_usd",
1913 "data.budget",
1914 "data.limit",
1915 "data.quota_total",
1916 "data.creditLimit",
1917 "data.total_credits",
1918 ],
1919 provider.extract.monthly_budget_divisor,
1920 )
1921 .or_else(|| {
1922 if provider.extract.derive_budget_from_remaining_and_spent {
1923 match (remaining_balance.or(component_remaining), monthly_spent) {
1924 (Some(remaining), Some(spent)) => Some(remaining.saturating_add(spent)),
1925 _ => None,
1926 }
1927 } else {
1928 None
1929 }
1930 });
1931 let total_balance = remaining_balance.or(component_remaining).or_else(|| {
1932 match (
1933 provider.extract.derive_remaining_from_budget_and_spent,
1934 monthly_budget,
1935 monthly_spent,
1936 ) {
1937 (true, Some(budget), Some(spent)) => Some(budget.saturating_sub(spent)),
1938 _ => None,
1939 }
1940 });
1941 let exhausted = first_bool_from_paths(
1942 value,
1943 &provider.extract.exhausted_paths,
1944 &[
1945 "exhausted",
1946 "quota_exhausted",
1947 "balance_exhausted",
1948 "data.exhausted",
1949 "data.quota_exhausted",
1950 "data.balance_exhausted",
1951 ],
1952 )
1953 .or_else(|| total_balance.map(UsdAmount::is_zero));
1954
1955 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1956 snapshot.total_balance_usd = total_balance.map(amount_to_string);
1957 snapshot.subscription_balance_usd = subscription_balance.map(amount_to_string);
1958 snapshot.paygo_balance_usd = paygo_balance.map(amount_to_string);
1959 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
1960 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1961 snapshot.exhausted = exhausted;
1962 snapshot.refresh_status(fetched_at_ms);
1963 snapshot
1964}
1965
1966fn has_any_json_path(value: &serde_json::Value, paths: &[&str]) -> bool {
1967 paths
1968 .iter()
1969 .any(|path| json_value_at_path(value, path).is_some())
1970}
1971
1972fn populate_sub2api_usage_fields(
1973 snapshot: &mut ProviderBalanceSnapshot,
1974 value: &serde_json::Value,
1975) {
1976 snapshot.plan_name = first_string_from_paths(
1977 value,
1978 &["planName", "plan_name", "data.planName", "data.plan_name"],
1979 );
1980 let remaining_balance = sub2api_remaining_balance(value);
1981 snapshot.total_balance_usd = snapshot
1982 .total_balance_usd
1983 .take()
1984 .or_else(|| remaining_balance.map(amount_to_string));
1985 snapshot.total_used_usd = first_amount_from_paths(
1986 value,
1987 &[],
1988 &[
1989 "usage.total.total_cost_usd",
1990 "usage.total.total_cost",
1991 "usage.total.cost",
1992 "data.usage.total.total_cost_usd",
1993 "data.usage.total.total_cost",
1994 "data.usage.total.cost",
1995 ],
1996 None,
1997 )
1998 .map(amount_to_string);
1999 snapshot.today_used_usd = first_amount_from_paths(
2000 value,
2001 &[],
2002 &[
2003 "usage.today.total_cost_usd",
2004 "usage.today.total_cost",
2005 "usage.today.cost",
2006 "data.usage.today.total_cost_usd",
2007 "data.usage.today.total_cost",
2008 "data.usage.today.cost",
2009 ],
2010 None,
2011 )
2012 .map(amount_to_string);
2013 snapshot.total_requests = first_u64_from_paths(
2014 value,
2015 &[
2016 "usage.total.request_count",
2017 "usage.total.requests",
2018 "usage.total.count",
2019 "data.usage.total.request_count",
2020 "data.usage.total.requests",
2021 "data.usage.total.count",
2022 ],
2023 );
2024 snapshot.today_requests = first_u64_from_paths(
2025 value,
2026 &[
2027 "usage.today.request_count",
2028 "usage.today.requests",
2029 "usage.today.count",
2030 "data.usage.today.request_count",
2031 "data.usage.today.requests",
2032 "data.usage.today.count",
2033 ],
2034 );
2035 snapshot.total_tokens = first_u64_from_paths(
2036 value,
2037 &[
2038 "usage.total.total_tokens",
2039 "usage.total.tokens",
2040 "usage.total.input_tokens",
2041 "usage.total.prompt_tokens",
2042 "data.usage.total.total_tokens",
2043 "data.usage.total.tokens",
2044 ],
2045 );
2046 snapshot.today_tokens = first_u64_from_paths(
2047 value,
2048 &[
2049 "usage.today.total_tokens",
2050 "usage.today.tokens",
2051 "usage.today.input_tokens",
2052 "usage.today.prompt_tokens",
2053 "data.usage.today.total_tokens",
2054 "data.usage.today.tokens",
2055 ],
2056 );
2057 snapshot.usage_rate = sub2api_usage_rate(value);
2058 snapshot.usage_windows = sub2api_usage_windows(value);
2059 snapshot.usage_model_stats = sub2api_model_stats(value);
2060 snapshot.subscription_expires_at = first_string_from_paths(
2061 value,
2062 &[
2063 "subscription.expires_at",
2064 "data.subscription.expires_at",
2065 "subscription.expiresAt",
2066 "data.subscription.expiresAt",
2067 ],
2068 );
2069 snapshot.usage_alerts = sub2api_usage_alerts(value);
2070}
2071
2072fn sub2api_remaining_balance(value: &serde_json::Value) -> Option<UsdAmount> {
2073 let remaining = first_amount_from_paths(value, &[], &["remaining", "data.remaining"], None)?;
2074 if sub2api_has_subscription_windows(value)
2075 && sub2api_window_remaining_amounts(value).contains(&remaining)
2076 {
2077 return None;
2078 }
2079 Some(remaining)
2080}
2081
2082fn sub2api_has_subscription_windows(value: &serde_json::Value) -> bool {
2083 has_any_json_path(
2084 value,
2085 &[
2086 "subscription.daily_usage_usd",
2087 "subscription.daily_limit_usd",
2088 "subscription.weekly_usage_usd",
2089 "subscription.weekly_limit_usd",
2090 "subscription.monthly_usage_usd",
2091 "subscription.monthly_limit_usd",
2092 "data.subscription.daily_usage_usd",
2093 "data.subscription.daily_limit_usd",
2094 "data.subscription.weekly_usage_usd",
2095 "data.subscription.weekly_limit_usd",
2096 "data.subscription.monthly_usage_usd",
2097 "data.subscription.monthly_limit_usd",
2098 ],
2099 )
2100}
2101
2102fn sub2api_window_remaining_amounts(value: &serde_json::Value) -> Vec<UsdAmount> {
2103 ["daily", "weekly", "monthly"]
2104 .into_iter()
2105 .filter_map(|period| {
2106 let used = first_amount_from_paths(
2107 value,
2108 &[],
2109 &[
2110 &format!("subscription.{period}_usage_usd"),
2111 &format!("data.subscription.{period}_usage_usd"),
2112 ],
2113 None,
2114 );
2115 let limit = first_amount_from_paths(
2116 value,
2117 &[],
2118 &[
2119 &format!("subscription.{period}_limit_usd"),
2120 &format!("data.subscription.{period}_limit_usd"),
2121 ],
2122 None,
2123 );
2124 match (limit, used) {
2125 (Some(limit), Some(used)) if !limit.is_zero() => Some(limit.saturating_sub(used)),
2126 _ => None,
2127 }
2128 })
2129 .collect()
2130}
2131
2132fn sub2api_usage_rate(value: &serde_json::Value) -> Option<ProviderUsageRateSnapshot> {
2133 let rate = ProviderUsageRateSnapshot {
2134 average_duration_ms: first_decimal_string_from_paths(
2135 value,
2136 &[
2137 "usage.average_duration_ms",
2138 "data.usage.average_duration_ms",
2139 "average_duration_ms",
2140 "data.average_duration_ms",
2141 ],
2142 ),
2143 rpm: first_decimal_string_from_paths(value, &["usage.rpm", "data.usage.rpm", "rpm"]),
2144 tpm: first_decimal_string_from_paths(value, &["usage.tpm", "data.usage.tpm", "tpm"]),
2145 };
2146 (!rate.is_empty()).then_some(rate)
2147}
2148
2149fn sub2api_usage_windows(value: &serde_json::Value) -> Vec<ProviderUsageWindow> {
2150 ["daily", "weekly", "monthly"]
2151 .into_iter()
2152 .filter_map(|period| {
2153 let used = first_amount_from_paths(
2154 value,
2155 &[],
2156 &[
2157 &format!("subscription.{period}_usage_usd"),
2158 &format!("data.subscription.{period}_usage_usd"),
2159 ],
2160 None,
2161 );
2162 let limit = first_amount_from_paths(
2163 value,
2164 &[],
2165 &[
2166 &format!("subscription.{period}_limit_usd"),
2167 &format!("data.subscription.{period}_limit_usd"),
2168 ],
2169 None,
2170 );
2171 if used.is_none() && limit.is_none() {
2172 return None;
2173 }
2174 let unlimited = limit.map(|limit| limit.is_zero());
2175 let remaining = match (limit, used) {
2176 (Some(limit), Some(used)) if !limit.is_zero() => Some(limit.saturating_sub(used)),
2177 _ => None,
2178 };
2179 Some(ProviderUsageWindow {
2180 period: period.to_string(),
2181 used_usd: used.map(amount_to_string),
2182 limit_usd: limit.map(amount_to_string),
2183 remaining_usd: remaining.map(amount_to_string),
2184 unlimited,
2185 })
2186 })
2187 .collect()
2188}
2189
2190fn sub2api_rate_limit_window_from_json(
2191 value: &serde_json::Value,
2192) -> Option<RateLimitWindowSnapshot> {
2193 let period = first_string_from_paths(value, &["window", "period", "name"])?;
2194 let limit = first_u64_from_paths(value, &["limit"]);
2195 if limit == Some(0) {
2196 return None;
2197 }
2198 let remaining = first_u64_from_paths(value, &["remaining"])?;
2199 if remaining > 0 {
2200 return None;
2201 }
2202 let reset_at_ms = first_secs_from_paths(value, &["reset_at", "resets_at", "resetAt"])
2203 .map(|secs| secs.saturating_mul(1000));
2204 Some(RateLimitWindowSnapshot {
2205 period: format!("rate_limit:{period}"),
2206 reset_at_ms,
2207 })
2208}
2209
2210fn sub2api_limiting_rate_limit_window(
2211 value: &serde_json::Value,
2212) -> Option<RateLimitWindowSnapshot> {
2213 ["rate_limits", "data.rate_limits"]
2214 .into_iter()
2215 .find_map(|path| array_from_json_path(value, path))
2216 .and_then(|items| {
2217 items
2218 .iter()
2219 .filter_map(sub2api_rate_limit_window_from_json)
2220 .max_by_key(|window| window.reset_at_ms.unwrap_or(0))
2221 })
2222}
2223
2224fn sub2api_model_stats(value: &serde_json::Value) -> Vec<ProviderUsageModelStat> {
2225 [
2226 "model_stats",
2227 "data.model_stats",
2228 "modelStats",
2229 "data.modelStats",
2230 ]
2231 .into_iter()
2232 .find_map(|path| array_from_json_path(value, path))
2233 .map(|items| {
2234 items
2235 .iter()
2236 .filter_map(sub2api_model_stat_from_json)
2237 .collect::<Vec<_>>()
2238 })
2239 .unwrap_or_default()
2240}
2241
2242fn sub2api_model_stat_from_json(value: &serde_json::Value) -> Option<ProviderUsageModelStat> {
2243 let model = first_string_from_paths(value, &["model", "model_name", "name"])?;
2244 let input_cost = first_amount_from_paths(value, &[], &["input_cost_usd", "input_cost"], None);
2245 let output_cost =
2246 first_amount_from_paths(value, &[], &["output_cost_usd", "output_cost"], None);
2247 let total_cost =
2248 first_amount_from_paths(value, &[], &["total_cost_usd", "total_cost", "cost"], None)
2249 .or_else(|| match (input_cost, output_cost) {
2250 (Some(input), Some(output)) => Some(input.saturating_add(output)),
2251 _ => None,
2252 });
2253 let input_tokens = first_u64_from_paths(value, &["input_tokens", "prompt_tokens"]);
2254 let output_tokens = first_u64_from_paths(value, &["output_tokens", "completion_tokens"]);
2255 let total_tokens =
2256 first_u64_from_paths(value, &["total_tokens", "tokens"]).or_else(|| {
2257 match (input_tokens, output_tokens) {
2258 (Some(input), Some(output)) => input.checked_add(output),
2259 _ => None,
2260 }
2261 });
2262 Some(ProviderUsageModelStat {
2263 model,
2264 request_count: first_u64_from_paths(value, &["request_count", "requests", "count"]),
2265 input_tokens,
2266 output_tokens,
2267 total_tokens,
2268 input_cost_usd: input_cost.map(amount_to_string),
2269 output_cost_usd: output_cost.map(amount_to_string),
2270 total_cost_usd: total_cost.map(amount_to_string),
2271 })
2272}
2273
2274fn sub2api_usage_alerts(value: &serde_json::Value) -> Vec<ProviderUsageAlert> {
2275 let mut alerts = Vec::new();
2276 if let (Some(used), Some(limit)) = (
2277 first_amount_from_paths(
2278 value,
2279 &[],
2280 &[
2281 "subscription.daily_usage_usd",
2282 "data.subscription.daily_usage_usd",
2283 ],
2284 None,
2285 ),
2286 first_amount_from_paths(
2287 value,
2288 &[],
2289 &[
2290 "subscription.daily_limit_usd",
2291 "data.subscription.daily_limit_usd",
2292 ],
2293 None,
2294 ),
2295 ) && !limit.is_zero()
2296 {
2297 let used_femto = used.femto_usd();
2298 let limit_femto = limit.femto_usd();
2299 if used_femto.saturating_mul(100) >= limit_femto.saturating_mul(95) {
2300 alerts.push(ProviderUsageAlert {
2301 kind: ProviderUsageAlertKind::DailyUsage95,
2302 message: "daily usage is at or above 95%".to_string(),
2303 });
2304 } else if used_femto.saturating_mul(100) >= limit_femto.saturating_mul(80) {
2305 alerts.push(ProviderUsageAlert {
2306 kind: ProviderUsageAlertKind::DailyUsage80,
2307 message: "daily usage is at or above 80%".to_string(),
2308 });
2309 }
2310 }
2311
2312 if let Some(remaining) = sub2api_remaining_balance(value)
2313 && let Some(threshold) = UsdAmount::from_decimal_str(LOW_BALANCE_ALERT_THRESHOLD_USD)
2314 && remaining <= threshold
2315 {
2316 alerts.push(ProviderUsageAlert {
2317 kind: ProviderUsageAlertKind::LowBalance,
2318 message: "remaining balance is low".to_string(),
2319 });
2320 }
2321
2322 if let Some(expires_at_secs) = first_secs_from_paths(
2323 value,
2324 &[
2325 "subscription.expires_at",
2326 "data.subscription.expires_at",
2327 "subscription.expiresAt",
2328 "data.subscription.expiresAt",
2329 ],
2330 ) {
2331 let now = unix_now_secs();
2332 if expires_at_secs <= now {
2333 alerts.push(ProviderUsageAlert {
2334 kind: ProviderUsageAlertKind::SubscriptionExpired,
2335 message: "subscription has expired".to_string(),
2336 });
2337 } else if expires_at_secs <= now.saturating_add(EXPIRING_SOON_WINDOW_SECS) {
2338 alerts.push(ProviderUsageAlert {
2339 kind: ProviderUsageAlertKind::SubscriptionExpiringSoon,
2340 message: "subscription expires within 7 days".to_string(),
2341 });
2342 }
2343 }
2344
2345 alerts.sort_by_key(|alert| alert.kind);
2346 alerts.dedup_by_key(|alert| alert.kind);
2347 alerts
2348}
2349
2350fn sub2api_subscription_limit_snapshot(
2351 value: &serde_json::Value,
2352 period: &'static str,
2353 limit_paths: &[&str],
2354 usage_paths: &[&str],
2355) -> Option<QuotaWindowSnapshot> {
2356 let budget = first_amount_from_paths(value, &[], limit_paths, None)?;
2357 if budget.is_zero() {
2358 return None;
2359 }
2360 let spent = first_amount_from_paths(value, &[], usage_paths, None).unwrap_or(UsdAmount::ZERO);
2361 let remaining = budget.saturating_sub(spent);
2362 Some(QuotaWindowSnapshot {
2363 period,
2364 remaining,
2365 used: spent,
2366 limit: budget,
2367 })
2368}
2369
2370fn sub2api_limiting_subscription_window(value: &serde_json::Value) -> Option<QuotaWindowSnapshot> {
2371 let windows = [
2372 sub2api_subscription_limit_snapshot(
2373 value,
2374 "daily",
2375 &[
2376 "subscription.daily_limit_usd",
2377 "data.subscription.daily_limit_usd",
2378 ],
2379 &[
2380 "subscription.daily_usage_usd",
2381 "data.subscription.daily_usage_usd",
2382 ],
2383 ),
2384 sub2api_subscription_limit_snapshot(
2385 value,
2386 "weekly",
2387 &[
2388 "subscription.weekly_limit_usd",
2389 "data.subscription.weekly_limit_usd",
2390 ],
2391 &[
2392 "subscription.weekly_usage_usd",
2393 "data.subscription.weekly_usage_usd",
2394 ],
2395 ),
2396 sub2api_subscription_limit_snapshot(
2397 value,
2398 "monthly",
2399 &[
2400 "subscription.monthly_limit_usd",
2401 "data.subscription.monthly_limit_usd",
2402 ],
2403 &[
2404 "subscription.monthly_usage_usd",
2405 "data.subscription.monthly_usage_usd",
2406 ],
2407 ),
2408 ];
2409
2410 windows
2411 .into_iter()
2412 .flatten()
2413 .min_by_key(|window| window.remaining)
2414}
2415
2416fn sub2api_usage_snapshot_from_json(
2417 provider: &UsageProviderConfig,
2418 upstream: &UpstreamRef,
2419 value: &serde_json::Value,
2420 fetched_at_ms: u64,
2421 stale_after_ms: Option<u64>,
2422) -> ProviderBalanceSnapshot {
2423 if json_value_at_path(value, "isValid").and_then(bool_from_json) == Some(false) {
2424 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
2425 .with_error("sub2api usage response reported invalid API key");
2426 }
2427
2428 let mode = first_string_from_paths(value, &["mode", "data.mode"]);
2429 let has_subscription = has_any_json_path(value, &["subscription", "data.subscription"]);
2430
2431 if mode.as_deref() == Some("quota_limited") {
2432 let rate_limit_window = sub2api_limiting_rate_limit_window(value);
2433 let quota_remaining = first_amount_from_paths(
2434 value,
2435 &provider.extract.remaining_balance_paths,
2436 &[
2437 "quota.remaining",
2438 "data.quota.remaining",
2439 "remaining",
2440 "data.remaining",
2441 ],
2442 provider.extract.remaining_divisor,
2443 );
2444 let quota_limit = first_amount_from_paths(
2445 value,
2446 &provider.extract.monthly_budget_paths,
2447 &["quota.limit", "data.quota.limit"],
2448 provider.extract.monthly_budget_divisor,
2449 );
2450 let quota_used = first_amount_from_paths(
2451 value,
2452 &provider.extract.monthly_spent_paths,
2453 &["quota.used", "data.quota.used"],
2454 provider.extract.monthly_spent_divisor,
2455 );
2456 let quota_exhausted = first_bool_from_paths(
2457 value,
2458 &provider.extract.exhausted_paths,
2459 &[
2460 "exhausted",
2461 "data.exhausted",
2462 "quota_exhausted",
2463 "data.quota_exhausted",
2464 ],
2465 )
2466 .or_else(|| quota_remaining.map(UsdAmount::is_zero));
2467 let exhausted = Some(quota_exhausted.unwrap_or(false) || rate_limit_window.is_some());
2468
2469 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
2470 if let Some(rate_limit_window) = rate_limit_window.clone()
2471 && quota_exhausted != Some(true)
2472 {
2473 snapshot.quota_period = Some(rate_limit_window.period);
2474 snapshot.quota_resets_at_ms = rate_limit_window.reset_at_ms;
2475 } else {
2476 snapshot.quota_period = Some("quota".to_string());
2477 snapshot.quota_remaining_usd = quota_remaining.map(amount_to_string);
2478 snapshot.quota_limit_usd = quota_limit.map(amount_to_string);
2479 snapshot.quota_used_usd = quota_used.map(amount_to_string);
2480 snapshot.monthly_budget_usd = quota_limit.map(amount_to_string);
2481 snapshot.monthly_spent_usd = quota_used.map(amount_to_string);
2482 }
2483 snapshot.exhausted = exhausted;
2484 populate_sub2api_usage_fields(&mut snapshot, value);
2485 snapshot.refresh_status(fetched_at_ms);
2486 return snapshot;
2487 }
2488
2489 if mode.as_deref() == Some("unrestricted") && has_subscription {
2490 let limiting_window = sub2api_limiting_subscription_window(value);
2491 let exhausted = first_bool_from_paths(
2492 value,
2493 &provider.extract.exhausted_paths,
2494 &[
2495 "exhausted",
2496 "data.exhausted",
2497 "quota_exhausted",
2498 "data.quota_exhausted",
2499 ],
2500 )
2501 .or_else(|| limiting_window.map(|window| window.remaining.is_zero()))
2502 .or(Some(false));
2503
2504 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
2505 if let Some(window) = limiting_window {
2506 snapshot.quota_period = Some(window.period.to_string());
2507 snapshot.quota_remaining_usd = Some(amount_to_string(window.remaining));
2508 snapshot.quota_limit_usd = Some(amount_to_string(window.limit));
2509 snapshot.quota_used_usd = Some(amount_to_string(window.used));
2510 snapshot.monthly_budget_usd = Some(amount_to_string(window.limit));
2511 snapshot.monthly_spent_usd = Some(amount_to_string(window.used));
2512 }
2513 snapshot.exhaustion_affects_routing = false;
2514 snapshot.exhausted = exhausted;
2515 populate_sub2api_usage_fields(&mut snapshot, value);
2516 snapshot.refresh_status(fetched_at_ms);
2517 return snapshot;
2518 }
2519
2520 let mut snapshot =
2521 balance_http_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms);
2522 populate_sub2api_usage_fields(&mut snapshot, value);
2523 snapshot.refresh_status(fetched_at_ms);
2524 snapshot
2525}
2526
2527fn sub2api_auth_me_snapshot_from_json(
2528 provider: &UsageProviderConfig,
2529 upstream: &UpstreamRef,
2530 value: &serde_json::Value,
2531 fetched_at_ms: u64,
2532 stale_after_ms: Option<u64>,
2533) -> ProviderBalanceSnapshot {
2534 if json_value_at_path(value, "code")
2535 .and_then(|value| value.as_i64())
2536 .is_some_and(|code| code != 0)
2537 {
2538 let message = json_value_at_path(value, "message")
2539 .and_then(|value| value.as_str())
2540 .unwrap_or("sub2api auth/me response reported failure");
2541 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
2542 .with_error(message.to_string());
2543 }
2544
2545 balance_http_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
2546}
2547
2548fn rightcode_available_prefixes(value: &serde_json::Value) -> Vec<String> {
2549 array_from_json_path(value, "available_prefixes")
2550 .into_iter()
2551 .flatten()
2552 .filter_map(|value| value.as_str())
2553 .map(str::trim)
2554 .filter(|value| !value.is_empty())
2555 .map(ToOwned::to_owned)
2556 .collect()
2557}
2558
2559fn rightcode_subscription_window(value: &serde_json::Value) -> Option<QuotaWindowSnapshot> {
2560 let limit = json_value_at_path(value, "total_quota").and_then(amount_from_json)?;
2561 if limit < UsdAmount::from_decimal_str("10").unwrap_or(UsdAmount::ZERO) {
2562 return None;
2563 }
2564 let raw_remaining = json_value_at_path(value, "remaining_quota").and_then(amount_from_json)?;
2565 let reset_today = json_value_at_path(value, "reset_today").and_then(bool_from_json);
2566 let remaining = if reset_today == Some(true) {
2567 raw_remaining
2568 } else {
2569 raw_remaining.saturating_add(limit)
2570 };
2571 let used = limit.saturating_sub(remaining);
2572 Some(QuotaWindowSnapshot {
2573 period: "daily",
2574 remaining,
2575 used,
2576 limit,
2577 })
2578}
2579
2580fn rightcode_account_summary_snapshot_from_json(
2581 provider: &UsageProviderConfig,
2582 upstream: &UpstreamRef,
2583 value: &serde_json::Value,
2584 upstream_base_url: &str,
2585 fetched_at_ms: u64,
2586 stale_after_ms: Option<u64>,
2587) -> ProviderBalanceSnapshot {
2588 let balance = json_value_at_path(value, "balance").and_then(amount_from_json);
2589 let provider_prefixes = base_path_prefixes(upstream_base_url);
2590 let mut matched_windows = Vec::new();
2591 let mut matched_plan_names = Vec::new();
2592
2593 if let Some(subscriptions) = array_from_json_path(value, "subscriptions") {
2594 for subscription in subscriptions {
2595 let available_prefixes = rightcode_available_prefixes(subscription);
2596 if !path_prefixes_match(&provider_prefixes, &available_prefixes) {
2597 continue;
2598 }
2599 let Some(window) = rightcode_subscription_window(subscription) else {
2600 continue;
2601 };
2602 matched_windows.push(window);
2603 if let Some(name) = json_value_at_path(subscription, "name").and_then(string_from_json)
2604 {
2605 matched_plan_names.push(name);
2606 }
2607 }
2608 }
2609
2610 if balance.is_none() && matched_windows.is_empty() {
2611 return snapshot_error(
2612 provider,
2613 upstream,
2614 fetched_at_ms,
2615 stale_after_ms,
2616 "rightcode account summary missing balance and matching subscription quota fields",
2617 );
2618 }
2619
2620 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
2621 snapshot.total_balance_usd = balance.map(amount_to_string);
2622
2623 if !matched_windows.is_empty() {
2624 let mut remaining = UsdAmount::ZERO;
2625 let mut used = UsdAmount::ZERO;
2626 let mut limit = UsdAmount::ZERO;
2627 for window in matched_windows {
2628 remaining = remaining.saturating_add(window.remaining);
2629 used = used.saturating_add(window.used);
2630 limit = limit.saturating_add(window.limit);
2631 }
2632 snapshot.quota_period = Some("daily".to_string());
2633 snapshot.quota_remaining_usd = Some(amount_to_string(remaining));
2634 snapshot.quota_used_usd = Some(amount_to_string(used));
2635 snapshot.quota_limit_usd = Some(amount_to_string(limit));
2636 if !matched_plan_names.is_empty() {
2637 matched_plan_names.sort();
2638 matched_plan_names.dedup();
2639 snapshot.plan_name = Some(matched_plan_names.join(", "));
2640 }
2641 snapshot.exhausted = Some(remaining.is_zero() && balance.is_none_or(UsdAmount::is_zero));
2642 } else {
2643 snapshot.exhausted = balance.map(UsdAmount::is_zero);
2644 }
2645
2646 snapshot.refresh_status(fetched_at_ms);
2647 snapshot
2648}
2649
2650fn new_api_token_usage_snapshot_from_json(
2651 provider: &UsageProviderConfig,
2652 upstream: &UpstreamRef,
2653 value: &serde_json::Value,
2654 fetched_at_ms: u64,
2655 stale_after_ms: Option<u64>,
2656) -> ProviderBalanceSnapshot {
2657 if json_value_at_path(value, "success").and_then(bool_from_json) == Some(false)
2658 || json_value_at_path(value, "code").and_then(bool_from_json) == Some(false)
2659 {
2660 let message = json_value_at_path(value, "message")
2661 .and_then(|value| value.as_str())
2662 .unwrap_or("new api token usage response reported failure");
2663 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
2664 .with_error(message.to_string());
2665 }
2666
2667 let mut effective = provider.extract.clone();
2668 if effective.remaining_balance_paths.is_empty() {
2669 effective.remaining_balance_paths = vec![
2670 "data.total_available".to_string(),
2671 "data.remain_quota".to_string(),
2672 "total_available".to_string(),
2673 "remain_quota".to_string(),
2674 ];
2675 }
2676 if effective.monthly_spent_paths.is_empty() {
2677 effective.monthly_spent_paths = vec![
2678 "data.total_used".to_string(),
2679 "data.used_quota".to_string(),
2680 "total_used".to_string(),
2681 "used_quota".to_string(),
2682 ];
2683 }
2684 if effective.monthly_budget_paths.is_empty() {
2685 effective.monthly_budget_paths = vec![
2686 "data.total_granted".to_string(),
2687 "total_granted".to_string(),
2688 ];
2689 }
2690 effective.remaining_divisor = effective.remaining_divisor.or(Some(500_000));
2691 effective.monthly_spent_divisor = effective.monthly_spent_divisor.or(Some(500_000));
2692 effective.monthly_budget_divisor = effective.monthly_budget_divisor.or(Some(500_000));
2693
2694 let unlimited_quota =
2695 first_bool_from_paths(value, &[], &["data.unlimited_quota", "unlimited_quota"])
2696 == Some(true);
2697 let remaining_balance = first_amount_from_paths(
2698 value,
2699 &effective.remaining_balance_paths,
2700 &[],
2701 effective.remaining_divisor,
2702 );
2703 let monthly_spent = first_amount_from_paths(
2704 value,
2705 &effective.monthly_spent_paths,
2706 &[],
2707 effective.monthly_spent_divisor,
2708 );
2709 let monthly_budget = first_amount_from_paths(
2710 value,
2711 &effective.monthly_budget_paths,
2712 &[],
2713 effective.monthly_budget_divisor,
2714 )
2715 .or_else(|| match (remaining_balance, monthly_spent) {
2716 (Some(remaining), Some(spent)) => Some(remaining.saturating_add(spent)),
2717 _ => None,
2718 });
2719 let exhausted = if unlimited_quota {
2720 Some(false)
2721 } else {
2722 first_bool_from_paths(
2723 value,
2724 &effective.exhausted_paths,
2725 &["data.exhausted", "exhausted"],
2726 )
2727 .or_else(|| remaining_balance.map(UsdAmount::is_zero))
2728 };
2729
2730 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
2731 snapshot.plan_name = first_string_from_paths(value, &["data.name", "name"]);
2732 snapshot.unlimited_quota = Some(unlimited_quota);
2733 if !unlimited_quota {
2734 snapshot.quota_period = Some("token".to_string());
2735 snapshot.quota_remaining_usd = remaining_balance.map(amount_to_string);
2736 snapshot.quota_limit_usd = monthly_budget.map(amount_to_string);
2737 snapshot.quota_used_usd = monthly_spent.map(amount_to_string);
2738 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
2739 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
2740 } else {
2741 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
2742 }
2743 snapshot.exhausted = exhausted;
2744 snapshot.refresh_status(fetched_at_ms);
2745 snapshot
2746}
2747
2748fn new_api_snapshot_from_json(
2749 provider: &UsageProviderConfig,
2750 upstream: &UpstreamRef,
2751 value: &serde_json::Value,
2752 fetched_at_ms: u64,
2753 stale_after_ms: Option<u64>,
2754) -> ProviderBalanceSnapshot {
2755 if json_value_at_path(value, "success").and_then(bool_from_json) == Some(false) {
2756 let message = json_value_at_path(value, "message")
2757 .and_then(|value| value.as_str())
2758 .unwrap_or("new api balance response reported failure");
2759 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
2760 .with_error(message.to_string());
2761 }
2762
2763 let mut effective = provider.extract.clone();
2764 if effective.remaining_balance_paths.is_empty() {
2765 effective.remaining_balance_paths = vec!["data.quota".to_string(), "quota".to_string()];
2766 }
2767 if effective.monthly_spent_paths.is_empty() {
2768 effective.monthly_spent_paths =
2769 vec!["data.used_quota".to_string(), "used_quota".to_string()];
2770 }
2771 effective.remaining_divisor = effective.remaining_divisor.or(Some(500_000));
2772 effective.monthly_spent_divisor = effective.monthly_spent_divisor.or(Some(500_000));
2773 effective.monthly_budget_divisor = effective.monthly_budget_divisor.or(Some(500_000));
2774
2775 let remaining_balance = first_amount_from_paths(
2776 value,
2777 &effective.remaining_balance_paths,
2778 &[],
2779 effective.remaining_divisor,
2780 );
2781 let monthly_spent = first_amount_from_paths(
2782 value,
2783 &effective.monthly_spent_paths,
2784 &[],
2785 effective.monthly_spent_divisor,
2786 );
2787 let monthly_budget = first_amount_from_paths(
2788 value,
2789 &effective.monthly_budget_paths,
2790 &["data.total_quota", "total_quota"],
2791 effective.monthly_budget_divisor,
2792 )
2793 .or_else(|| match (remaining_balance, monthly_spent) {
2794 (Some(remaining), Some(spent)) => Some(remaining.saturating_add(spent)),
2795 _ => None,
2796 });
2797 let unlimited_quota =
2798 first_bool_from_paths(value, &[], &["data.unlimited_quota", "unlimited_quota"])
2799 == Some(true);
2800 let exhausted = if unlimited_quota {
2801 Some(false)
2802 } else {
2803 first_bool_from_paths(
2804 value,
2805 &effective.exhausted_paths,
2806 &["data.exhausted", "exhausted"],
2807 )
2808 .or_else(|| remaining_balance.map(UsdAmount::is_zero))
2809 };
2810
2811 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
2812 snapshot.unlimited_quota = Some(unlimited_quota);
2813 if !unlimited_quota {
2814 snapshot.quota_period = Some("quota".to_string());
2815 snapshot.quota_remaining_usd = remaining_balance.map(amount_to_string);
2816 snapshot.quota_limit_usd = monthly_budget.map(amount_to_string);
2817 snapshot.quota_used_usd = monthly_spent.map(amount_to_string);
2818 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
2819 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
2820 } else {
2821 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
2822 }
2823 snapshot.exhausted = exhausted;
2824 snapshot.refresh_status(fetched_at_ms);
2825 snapshot
2826}
2827
2828fn openai_cost_result_usd_amount(result: &serde_json::Value) -> Option<UsdAmount> {
2829 let amount = json_value_at_path(result, "amount.value").and_then(amount_from_json)?;
2830 let currency = json_value_at_path(result, "amount.currency").and_then(|value| value.as_str());
2831 match currency {
2832 Some(currency) if currency.eq_ignore_ascii_case("usd") => Some(amount),
2833 None => Some(amount),
2834 _ => None,
2835 }
2836}
2837
2838fn openai_organization_costs_total(value: &serde_json::Value) -> Option<UsdAmount> {
2839 let buckets = json_value_at_path(value, "data")?.as_array()?;
2840 let mut total = UsdAmount::ZERO;
2841
2842 for bucket in buckets {
2843 let Some(results) =
2844 json_value_at_path(bucket, "results").and_then(|value| value.as_array())
2845 else {
2846 continue;
2847 };
2848 for result in results {
2849 if let Some(amount) = openai_cost_result_usd_amount(result) {
2850 total = total.saturating_add(amount);
2851 }
2852 }
2853 }
2854
2855 Some(total)
2856}
2857
2858fn openai_organization_costs_snapshot_from_json(
2859 provider: &UsageProviderConfig,
2860 upstream: &UpstreamRef,
2861 value: &serde_json::Value,
2862 fetched_at_ms: u64,
2863 stale_after_ms: Option<u64>,
2864) -> ProviderBalanceSnapshot {
2865 let spent = first_amount_from_paths(
2866 value,
2867 &provider.extract.monthly_spent_paths,
2868 &[],
2869 provider.extract.monthly_spent_divisor,
2870 )
2871 .or_else(|| openai_organization_costs_total(value));
2872
2873 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
2874 snapshot.monthly_spent_usd = spent.map(amount_to_string);
2875 snapshot.exhausted = None;
2876 snapshot.exhaustion_affects_routing = false;
2877 snapshot.refresh_status(fetched_at_ms);
2878 snapshot
2879}
2880
2881fn balance_exhaustion_policy_action(
2882 endpoint_key: ProviderEndpointKey,
2883 observed_at_ms: u64,
2884) -> Option<PolicyAction> {
2885 const BALANCE_EXHAUSTION_COOLDOWN_SECS: u64 = 24 * 60 * 60;
2886 let mut signal = ProviderSignal::high_confidence_route_facing(
2887 ProviderSignalKind::Balance,
2888 ProviderSignalSource::BalanceSnapshot,
2889 ProviderSignalTarget::ProviderEndpoint {
2890 provider_endpoint_key: endpoint_key,
2891 },
2892 observed_at_ms,
2893 );
2894 signal.reset_after_secs = Some(BALANCE_EXHAUSTION_COOLDOWN_SECS);
2895 signal.reason = Some("balance_exhausted".to_string());
2896 PolicyAction::cooldown_from_signal(signal, observed_at_ms, 0, observed_at_ms)
2897}
2898
2899async fn sync_balance_policy_action_for_endpoint(
2900 state: &Arc<ProxyState>,
2901 service_name: &str,
2902 endpoint_key: ProviderEndpointKey,
2903 exhausted: bool,
2904) {
2905 if exhausted {
2906 if let Some(action) =
2907 balance_exhaustion_policy_action(endpoint_key, crate::logging::now_ms())
2908 {
2909 state.upsert_owned_policy_action(service_name, action).await;
2910 }
2911 } else {
2912 state
2913 .clear_owned_policy_action(
2914 service_name,
2915 &endpoint_key,
2916 PolicyActionKind::Cooldown,
2917 ProviderSignalKind::Balance,
2918 ProviderSignalSource::BalanceSnapshot,
2919 )
2920 .await;
2921 }
2922}
2923
2924async fn update_usage_exhausted(
2925 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
2926 state: &Arc<ProxyState>,
2927 cfg: &ProxyConfig,
2928 service_name: &str,
2929 upstreams: &[UpstreamRef],
2930 exhausted: bool,
2931) {
2932 if let Ok(mut map) = lb_states.lock() {
2933 for uref in upstreams {
2934 let service = match service_manager(cfg, service_name).station(&uref.station_name) {
2935 Some(s) => s,
2936 None => continue,
2937 };
2938
2939 let entry = map
2940 .entry(uref.station_name.clone())
2941 .or_insert_with(LbState::default);
2942 entry.ensure_layout(service.name.as_str(), &service.upstreams);
2943 if uref.index < entry.usage_exhausted.len() {
2944 entry.usage_exhausted[uref.index] = exhausted;
2945 }
2946 }
2947 }
2948
2949 for uref in upstreams {
2950 if let Some(endpoint_key) = uref.provider_endpoint.clone() {
2951 sync_balance_policy_action_for_endpoint(
2952 state,
2953 service_name,
2954 endpoint_key.clone(),
2955 exhausted,
2956 )
2957 .await;
2958 state
2959 .set_provider_endpoint_usage_exhausted(service_name, endpoint_key, exhausted)
2960 .await;
2961 }
2962 }
2963}
2964
2965fn provider_hosts_for_diagnostics(
2966 cfg: &ProxyConfig,
2967 service_name: &str,
2968 provider: &UsageProviderConfig,
2969) -> Vec<String> {
2970 let mut hosts: Vec<String> = Vec::new();
2971 for service in service_manager(cfg, service_name).stations().values() {
2972 for upstream in &service.upstreams {
2973 if domain_matches(&upstream.base_url, &provider.domains)
2974 && let Ok(url) = reqwest::Url::parse(&upstream.base_url)
2975 && let Some(host) = url.host_str()
2976 {
2977 hosts.push(host.to_string());
2978 }
2979 }
2980 }
2981 hosts.sort();
2982 hosts.dedup();
2983 hosts
2984}
2985
2986fn warn_if_provider_spans_hosts(
2987 cfg: &ProxyConfig,
2988 service_name: &str,
2989 provider: &UsageProviderConfig,
2990) {
2991 let hosts = provider_hosts_for_diagnostics(cfg, service_name, provider);
2992 if hosts.len() > 1 {
2993 warn!(
2994 "usage provider '{}' is associated with multiple hosts: {:?}; \
2995将按统一额度处理这些 upstream,如需区分配额请拆分为多个 provider 配置",
2996 provider.id, hosts
2997 );
2998 }
2999}
3000
3001fn snapshot_from_provider_json(
3002 provider: &UsageProviderConfig,
3003 upstream: &UpstreamRef,
3004 value: &serde_json::Value,
3005 upstream_base_url: &str,
3006 fetched_at_ms: u64,
3007 stale_after_ms: Option<u64>,
3008) -> ProviderBalanceSnapshot {
3009 match provider.kind {
3010 ProviderKind::BudgetHttpJson => {
3011 budget_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
3012 }
3013 ProviderKind::YescodeProfile => {
3014 yescode_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
3015 }
3016 ProviderKind::OpenAiBalanceHttpJson => balance_http_snapshot_from_json(
3017 provider,
3018 upstream,
3019 value,
3020 fetched_at_ms,
3021 stale_after_ms,
3022 ),
3023 ProviderKind::Sub2ApiUsage => sub2api_usage_snapshot_from_json(
3024 provider,
3025 upstream,
3026 value,
3027 fetched_at_ms,
3028 stale_after_ms,
3029 ),
3030 ProviderKind::Sub2ApiAuthMe => sub2api_auth_me_snapshot_from_json(
3031 provider,
3032 upstream,
3033 value,
3034 fetched_at_ms,
3035 stale_after_ms,
3036 ),
3037 ProviderKind::NewApiTokenUsage => new_api_token_usage_snapshot_from_json(
3038 provider,
3039 upstream,
3040 value,
3041 fetched_at_ms,
3042 stale_after_ms,
3043 ),
3044 ProviderKind::NewApiUserSelf => {
3045 new_api_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
3046 }
3047 ProviderKind::RightCodeAccountSummary => rightcode_account_summary_snapshot_from_json(
3048 provider,
3049 upstream,
3050 value,
3051 upstream_base_url,
3052 fetched_at_ms,
3053 stale_after_ms,
3054 ),
3055 ProviderKind::OpenAiOrganizationCosts => openai_organization_costs_snapshot_from_json(
3056 provider,
3057 upstream,
3058 value,
3059 fetched_at_ms,
3060 stale_after_ms,
3061 ),
3062 }
3063}
3064
3065async fn refresh_provider_target(
3066 params: RefreshProviderTargetParams<'_>,
3067) -> UsageProviderRefreshOutcome {
3068 let RefreshProviderTargetParams {
3069 client,
3070 provider,
3071 target,
3072 cfg,
3073 lb_states,
3074 state,
3075 service_name,
3076 interval_secs,
3077 } = params;
3078
3079 let upstreams = vec![target.upstream.clone()];
3080 let fetched_at_ms = unix_now_ms();
3081 let stale_after_ms = stale_after_ms(fetched_at_ms, interval_secs);
3082 if let Some(suppression) =
3083 usage_provider_target_suppression_active(&provider.id, target, Instant::now())
3084 {
3085 update_usage_exhausted(
3086 lb_states,
3087 state,
3088 cfg,
3089 service_name,
3090 &upstreams,
3091 suppression.routing_exhausted,
3092 )
3093 .await;
3094 warn!(
3095 "usage provider '{}' skipped {}[{}]: balance refresh suppressed: {}",
3096 provider.id, target.upstream.station_name, target.upstream.index, suppression.reason
3097 );
3098 return UsageProviderRefreshOutcome::Failed;
3099 }
3100 if let Some(decision) = existing_usage_provider_target_suppression_decision(
3101 state,
3102 cfg,
3103 service_name,
3104 &provider.id,
3105 target,
3106 fetched_at_ms,
3107 )
3108 .await
3109 {
3110 remember_usage_provider_target_suppression(
3111 &provider.id,
3112 target,
3113 decision.ttl,
3114 decision.reason.clone(),
3115 decision.routing_exhausted,
3116 Instant::now(),
3117 );
3118 update_usage_exhausted(
3119 lb_states,
3120 state,
3121 cfg,
3122 service_name,
3123 &upstreams,
3124 decision.routing_exhausted,
3125 )
3126 .await;
3127 warn!(
3128 "usage provider '{}' skipped {}[{}]: existing balance snapshot suppresses refresh: {}",
3129 provider.id, target.upstream.station_name, target.upstream.index, decision.reason
3130 );
3131 return UsageProviderRefreshOutcome::Failed;
3132 }
3133
3134 let Some(token) = resolve_token(provider, &upstreams, cfg, service_name) else {
3135 let snapshot = if provider.kind == ProviderKind::OpenAiOrganizationCosts {
3136 base_snapshot(provider, &upstreams[0], fetched_at_ms, stale_after_ms)
3137 } else {
3138 base_snapshot(provider, &upstreams[0], fetched_at_ms, stale_after_ms)
3139 .with_error("no usable token; checked provider token_env and upstream auth")
3140 };
3141 state
3142 .record_provider_balance_snapshot(service_name, snapshot)
3143 .await;
3144 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
3145 if provider.kind == ProviderKind::OpenAiOrganizationCosts {
3146 warn!(
3147 "usage provider '{}' is missing OPENAI_ADMIN_KEY; OpenAI official costs stay unknown",
3148 provider.id
3149 );
3150 } else {
3151 warn!(
3152 "usage provider '{}' has no usable token (checked token_env and associated upstream auth_token); \
3153跳过本次用量查询,请检查 usage_providers.json 和 ~/.codex-helper/config.json",
3154 provider.id
3155 );
3156 }
3157 return UsageProviderRefreshOutcome::MissingToken;
3158 };
3159
3160 match poll_provider_http_json(client, provider, &target.base_url, &token).await {
3161 Ok(value) => {
3162 let snapshot = snapshot_from_provider_json(
3163 provider,
3164 &upstreams[0],
3165 &value,
3166 &target.base_url,
3167 fetched_at_ms,
3168 stale_after_ms,
3169 );
3170 let snapshot_error = usage_provider_snapshot_error(&snapshot).map(str::to_string);
3171 let suppression_decision =
3172 usage_provider_suppression_decision_from_snapshot(&snapshot, cfg, fetched_at_ms);
3173 let exhausted_for_lb = suppression_decision
3174 .as_ref()
3175 .is_some_and(|decision| decision.routing_exhausted);
3176 if let Some(decision) = suppression_decision.as_ref() {
3177 remember_usage_provider_target_suppression(
3178 &provider.id,
3179 target,
3180 decision.ttl,
3181 decision.reason.as_str(),
3182 decision.routing_exhausted,
3183 Instant::now(),
3184 );
3185 } else {
3186 clear_usage_provider_target_suppression(&provider.id, target);
3187 }
3188 update_usage_exhausted(
3189 lb_states,
3190 state,
3191 cfg,
3192 service_name,
3193 &upstreams,
3194 exhausted_for_lb,
3195 )
3196 .await;
3197 state
3198 .record_provider_balance_snapshot(service_name, snapshot)
3199 .await;
3200 if let Some(error) = snapshot_error {
3201 warn!(
3202 "usage provider '{}' returned error snapshot for {}[{}]: {}",
3203 provider.id, target.upstream.station_name, target.upstream.index, error
3204 );
3205 return UsageProviderRefreshOutcome::Failed;
3206 }
3207 info!(
3208 "usage provider '{}' refreshed {}[{}], exhausted = {}, routing_trusted = {}",
3209 provider.id,
3210 target.upstream.station_name,
3211 target.upstream.index,
3212 exhausted_for_lb,
3213 provider.trust_exhaustion_for_routing
3214 );
3215 UsageProviderRefreshOutcome::Refreshed
3216 }
3217 Err(err) => {
3218 let error = err.to_string();
3219 let terminal_failure = usage_provider_error_is_terminal(&error);
3220 if terminal_failure {
3221 remember_usage_provider_target_suppression(
3222 &provider.id,
3223 target,
3224 USAGE_PROVIDER_TERMINAL_FAILURE_TTL,
3225 error.clone(),
3226 true,
3227 Instant::now(),
3228 );
3229 }
3230 state
3231 .record_provider_balance_snapshot(
3232 service_name,
3233 base_snapshot(provider, &upstreams[0], fetched_at_ms, stale_after_ms)
3234 .with_error(error.clone()),
3235 )
3236 .await;
3237 update_usage_exhausted(
3238 lb_states,
3239 state,
3240 cfg,
3241 service_name,
3242 &upstreams,
3243 terminal_failure,
3244 )
3245 .await;
3246 warn!(
3247 "usage provider '{}' poll failed for {}[{}]: {}",
3248 provider.id, target.upstream.station_name, target.upstream.index, error
3249 );
3250 UsageProviderRefreshOutcome::Failed
3251 }
3252 }
3253}
3254
3255struct ConfiguredRefreshJob<'a> {
3256 provider: &'a UsageProviderConfig,
3257 target: UsageProviderTarget,
3258 interval_secs: u64,
3259}
3260
3261struct AutoRefreshJob {
3262 target: UsageProviderTarget,
3263}
3264
3265async fn run_configured_refresh_job<'a>(
3266 client: &'a Client,
3267 job: ConfiguredRefreshJob<'a>,
3268 cfg: &'a ProxyConfig,
3269 lb_states: &'a Arc<Mutex<HashMap<String, LbState>>>,
3270 state: &'a Arc<ProxyState>,
3271 service_name: &'a str,
3272) -> (String, UsageProviderRefreshOutcome) {
3273 let provider_id = job.provider.id.clone();
3274 let outcome = refresh_provider_target(RefreshProviderTargetParams {
3275 client,
3276 provider: job.provider,
3277 target: &job.target,
3278 cfg,
3279 lb_states,
3280 state,
3281 service_name,
3282 interval_secs: job.interval_secs,
3283 })
3284 .await;
3285 (provider_id, outcome)
3286}
3287
3288async fn run_auto_refresh_job(
3289 client: &Client,
3290 job: AutoRefreshJob,
3291 cfg: &ProxyConfig,
3292 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
3293 state: &Arc<ProxyState>,
3294 service_name: &str,
3295) -> UsageProviderRefreshOutcome {
3296 auto_probe_provider_target(client, &job.target, cfg, lb_states, state, service_name).await
3297}
3298
3299async fn run_configured_refresh_jobs<'a>(
3300 client: &'a Client,
3301 jobs: Vec<ConfiguredRefreshJob<'a>>,
3302 cfg: &'a ProxyConfig,
3303 lb_states: &'a Arc<Mutex<HashMap<String, LbState>>>,
3304 state: &'a Arc<ProxyState>,
3305 service_name: &'a str,
3306) -> Vec<(String, UsageProviderRefreshOutcome)> {
3307 let mut pending = jobs.into_iter();
3308 let mut running = FuturesUnordered::new();
3309 let mut results = Vec::new();
3310 let concurrency = BALANCE_REFRESH_CONCURRENCY.max(1);
3311
3312 for _ in 0..concurrency {
3313 let Some(job) = pending.next() else {
3314 break;
3315 };
3316 running.push(run_configured_refresh_job(
3317 client,
3318 job,
3319 cfg,
3320 lb_states,
3321 state,
3322 service_name,
3323 ));
3324 }
3325
3326 while let Some(result) = running.next().await {
3327 results.push(result);
3328 if let Some(job) = pending.next() {
3329 running.push(run_configured_refresh_job(
3330 client,
3331 job,
3332 cfg,
3333 lb_states,
3334 state,
3335 service_name,
3336 ));
3337 }
3338 }
3339
3340 results
3341}
3342
3343async fn run_auto_refresh_jobs(
3344 client: &Client,
3345 jobs: Vec<AutoRefreshJob>,
3346 cfg: &ProxyConfig,
3347 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
3348 state: &Arc<ProxyState>,
3349 service_name: &str,
3350) -> Vec<UsageProviderRefreshOutcome> {
3351 let mut pending = jobs.into_iter();
3352 let mut running = FuturesUnordered::new();
3353 let mut results = Vec::new();
3354 let concurrency = BALANCE_REFRESH_CONCURRENCY.max(1);
3355
3356 for _ in 0..concurrency {
3357 let Some(job) = pending.next() else {
3358 break;
3359 };
3360 running.push(run_auto_refresh_job(
3361 client,
3362 job,
3363 cfg,
3364 lb_states,
3365 state,
3366 service_name,
3367 ));
3368 }
3369
3370 while let Some(result) = running.next().await {
3371 results.push(result);
3372 if let Some(job) = pending.next() {
3373 running.push(run_auto_refresh_job(
3374 client,
3375 job,
3376 cfg,
3377 lb_states,
3378 state,
3379 service_name,
3380 ));
3381 }
3382 }
3383
3384 results
3385}
3386
3387fn auto_snapshot_is_usable(snapshot: &ProviderBalanceSnapshot) -> bool {
3388 snapshot.error.is_none()
3389 && matches!(
3390 snapshot.status,
3391 BalanceSnapshotStatus::Ok | BalanceSnapshotStatus::Exhausted
3392 )
3393}
3394
3395fn normalized_error_text(value: &str) -> String {
3396 value
3397 .chars()
3398 .map(|ch| match ch {
3399 '_' | '-' | '.' | '/' | ':' | '[' | ']' | '(' | ')' | ',' | ';' => ' ',
3400 _ => ch.to_ascii_lowercase(),
3401 })
3402 .collect::<String>()
3403}
3404
3405fn usage_provider_error_is_terminal(error: &str) -> bool {
3406 let normalized = normalized_error_text(error);
3407 let terminal_markers = [
3408 "user inactive",
3409 "user account is not active",
3410 "account is not active",
3411 "account inactive",
3412 "account disabled",
3413 "user disabled",
3414 "api key disabled",
3415 "api key is disabled",
3416 "api key inactive",
3417 "api key is not active",
3418 "key inactive",
3419 "key disabled",
3420 "invalid api key",
3421 "invalid token",
3422 "invalid bearer token",
3423 "token invalid",
3424 "unauthorized api key",
3425 "insufficient balance",
3426 "balance insufficient",
3427 "insufficient quota",
3428 "quota exhausted",
3429 "quota exceeded",
3430 "no balance",
3431 "余额不足",
3432 "额度不足",
3433 "配额不足",
3434 "账户未激活",
3435 "账号未激活",
3436 "用户未激活",
3437 "账户已禁用",
3438 "账号已禁用",
3439 "用户已禁用",
3440 "密钥无效",
3441 "令牌无效",
3442 ];
3443 terminal_markers
3444 .iter()
3445 .any(|marker| normalized.contains(marker))
3446}
3447
3448fn quota_period_is_current_day(period: &str) -> bool {
3449 let normalized = period.trim().to_ascii_lowercase();
3450 matches!(
3451 normalized.as_str(),
3452 "daily" | "day" | "today" | "current_day" | "current-day" | "1d" | "24h" | "今日" | "今天"
3453 )
3454}
3455
3456fn quota_period_is_refreshable_window(period: &str) -> bool {
3457 let normalized = period.trim().to_ascii_lowercase();
3458 quota_period_is_current_day(&normalized)
3459 || matches!(
3460 normalized.as_str(),
3461 "weekly" | "week" | "7d" | "monthly" | "month"
3462 )
3463 || normalized.starts_with("rate_limit:")
3464}
3465
3466fn snapshot_has_current_day_quota_exhaustion(snapshot: &ProviderBalanceSnapshot) -> bool {
3467 snapshot.exhausted == Some(true)
3468 && snapshot
3469 .quota_period
3470 .as_deref()
3471 .is_some_and(quota_period_is_current_day)
3472}
3473
3474fn snapshot_has_refreshable_window_exhaustion(snapshot: &ProviderBalanceSnapshot) -> bool {
3475 snapshot.exhausted == Some(true)
3476 && snapshot
3477 .quota_period
3478 .as_deref()
3479 .is_some_and(quota_period_is_refreshable_window)
3480}
3481
3482fn duration_millis_u64(duration: Duration) -> u64 {
3483 u64::try_from(duration.as_millis()).unwrap_or(u64::MAX)
3484}
3485
3486fn duration_until_ms(deadline_ms: u64, now_ms: u64) -> Option<Duration> {
3487 deadline_ms
3488 .checked_sub(now_ms)
3489 .filter(|remaining_ms| *remaining_ms > 0)
3490 .map(Duration::from_millis)
3491}
3492
3493fn snapshot_freshness_ttl(snapshot: &ProviderBalanceSnapshot, now_ms: u64) -> Option<Duration> {
3494 snapshot
3495 .stale_after_ms
3496 .and_then(|stale_after_ms| duration_until_ms(stale_after_ms, now_ms))
3497}
3498
3499fn current_day_quota_suppression_ttl(
3500 snapshot: &ProviderBalanceSnapshot,
3501 cfg: &ProxyConfig,
3502 now_ms: u64,
3503) -> Option<Duration> {
3504 let reset_at_ms = next_reset_at_ms(
3505 snapshot.fetched_at_ms,
3506 cfg.ui.usage_forecast.reset_utc_offset.as_str(),
3507 cfg.ui.usage_forecast.reset_time.as_str(),
3508 )?;
3509 let suppress_until_ms = reset_at_ms.saturating_add(duration_millis_u64(
3510 USAGE_PROVIDER_DAILY_RESET_SUPPRESSION_GRACE,
3511 ));
3512 duration_until_ms(suppress_until_ms, now_ms)
3513}
3514
3515fn usage_provider_snapshot_suppression_ttl(
3516 snapshot: &ProviderBalanceSnapshot,
3517 cfg: &ProxyConfig,
3518 now_ms: u64,
3519) -> Option<Duration> {
3520 if let Some(reset_at_ms) = snapshot.quota_resets_at_ms {
3521 let suppress_until_ms = reset_at_ms.saturating_add(duration_millis_u64(
3522 USAGE_PROVIDER_DAILY_RESET_SUPPRESSION_GRACE,
3523 ));
3524 return duration_until_ms(suppress_until_ms, now_ms)
3525 .or_else(|| snapshot_freshness_ttl(snapshot, now_ms));
3526 }
3527
3528 if snapshot_has_current_day_quota_exhaustion(snapshot) {
3529 return current_day_quota_suppression_ttl(snapshot, cfg, now_ms)
3530 .or_else(|| snapshot_freshness_ttl(snapshot, now_ms));
3531 }
3532
3533 if snapshot_has_refreshable_window_exhaustion(snapshot) {
3534 return snapshot_freshness_ttl(snapshot, now_ms);
3535 }
3536
3537 if snapshot.stale_at(now_ms) {
3538 None
3539 } else {
3540 Some(USAGE_PROVIDER_EXHAUSTED_SUPPRESSION_TTL)
3541 }
3542}
3543
3544fn usage_provider_snapshot_suppression_reason(
3545 snapshot: &ProviderBalanceSnapshot,
3546) -> Option<String> {
3547 if snapshot.status_at(snapshot.fetched_at_ms) != BalanceSnapshotStatus::Exhausted {
3548 return None;
3549 }
3550
3551 if snapshot_has_current_day_quota_exhaustion(snapshot) {
3552 let period = snapshot.quota_period.as_deref().unwrap_or("daily");
3553 return Some(format!(
3554 "{period} package quota exhausted for current period"
3555 ));
3556 }
3557
3558 if snapshot_has_refreshable_window_exhaustion(snapshot) {
3559 let period = snapshot.quota_period.as_deref().unwrap_or("usage");
3560 return Some(format!(
3561 "{period} usage window exhausted for current period"
3562 ));
3563 }
3564
3565 if snapshot.routing_exhausted() {
3566 return Some("balance exhausted".to_string());
3567 }
3568
3569 None
3570}
3571
3572fn usage_provider_snapshot_error(snapshot: &ProviderBalanceSnapshot) -> Option<&str> {
3573 snapshot
3574 .error
3575 .as_deref()
3576 .map(str::trim)
3577 .filter(|value| !value.is_empty())
3578}
3579
3580fn usage_provider_snapshot_terminal_error(snapshot: &ProviderBalanceSnapshot) -> Option<&str> {
3581 usage_provider_snapshot_error(snapshot).filter(|error| usage_provider_error_is_terminal(error))
3582}
3583
3584fn usage_provider_suppression_decision_from_snapshot(
3585 snapshot: &ProviderBalanceSnapshot,
3586 cfg: &ProxyConfig,
3587 now_ms: u64,
3588) -> Option<ProviderTargetSuppressionDecision> {
3589 if let Some(error) = usage_provider_snapshot_terminal_error(snapshot) {
3590 if snapshot.stale_at(now_ms) {
3591 return None;
3592 }
3593 return Some(ProviderTargetSuppressionDecision {
3594 reason: error.to_string(),
3595 routing_exhausted: true,
3596 ttl: USAGE_PROVIDER_TERMINAL_FAILURE_TTL,
3597 });
3598 }
3599
3600 let reason = usage_provider_snapshot_suppression_reason(snapshot)?;
3601 usage_provider_snapshot_suppression_ttl(snapshot, cfg, now_ms).map(|ttl| {
3602 ProviderTargetSuppressionDecision {
3603 reason,
3604 routing_exhausted: true,
3605 ttl,
3606 }
3607 })
3608}
3609
3610fn provider_balance_snapshot_matches_target(
3611 snapshot: &ProviderBalanceSnapshot,
3612 provider_id: &str,
3613 target: &UsageProviderTarget,
3614) -> bool {
3615 if snapshot.provider_id != provider_id {
3616 return false;
3617 }
3618 if snapshot.station_name.as_deref() != Some(target.upstream.station_name.as_str()) {
3619 return false;
3620 }
3621 if snapshot.upstream_index != Some(target.upstream.index) {
3622 return false;
3623 }
3624
3625 match (
3626 snapshot.provider_endpoint_key.as_deref(),
3627 target
3628 .upstream
3629 .provider_endpoint
3630 .as_ref()
3631 .map(ProviderEndpointKey::stable_key),
3632 ) {
3633 (Some(snapshot_key), Some(target_key)) => snapshot_key == target_key,
3634 _ => true,
3635 }
3636}
3637
3638async fn existing_usage_provider_target_suppression_decision(
3639 state: &Arc<ProxyState>,
3640 cfg: &ProxyConfig,
3641 service_name: &str,
3642 provider_id: &str,
3643 target: &UsageProviderTarget,
3644 now_ms: u64,
3645) -> Option<ProviderTargetSuppressionDecision> {
3646 let view = state.get_provider_balance_view(service_name).await;
3647 view.get(&target.upstream.station_name)
3648 .and_then(|snapshots| {
3649 snapshots
3650 .iter()
3651 .filter(|snapshot| {
3652 provider_balance_snapshot_matches_target(snapshot, provider_id, target)
3653 })
3654 .max_by_key(|snapshot| snapshot.fetched_at_ms)
3655 .and_then(|snapshot| {
3656 usage_provider_suppression_decision_from_snapshot(snapshot, cfg, now_ms)
3657 })
3658 })
3659}
3660
3661fn auto_probe_error_summary(probe_errors: &[String]) -> Option<String> {
3662 (!probe_errors.is_empty()).then(|| format!("attempts failed: {}", probe_errors.join("; ")))
3663}
3664
3665async fn auto_probe_provider_target(
3666 client: &Client,
3667 target: &UsageProviderTarget,
3668 cfg: &ProxyConfig,
3669 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
3670 state: &Arc<ProxyState>,
3671 service_name: &str,
3672) -> UsageProviderRefreshOutcome {
3673 let upstreams = vec![target.upstream.clone()];
3674 let fetched_at_ms = unix_now_ms();
3675 let interval_secs = DEFAULT_POLL_INTERVAL_SECS;
3676 let stale_after_ms = stale_after_ms(fetched_at_ms, interval_secs);
3677
3678 if is_official_openai_base_url(&target.base_url) {
3679 let provider = auto_openai_official_provider(target);
3680 let Some(token) = resolve_token(&provider, &upstreams, cfg, service_name) else {
3681 state
3682 .record_provider_balance_snapshot(
3683 service_name,
3684 base_snapshot(&provider, &target.upstream, fetched_at_ms, stale_after_ms),
3685 )
3686 .await;
3687 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
3688 warn!(
3689 "OpenAI organization costs require OPENAI_ADMIN_KEY; balance stays unknown for {}[{}]",
3690 target.upstream.station_name, target.upstream.index
3691 );
3692 return UsageProviderRefreshOutcome::MissingToken;
3693 };
3694
3695 return match poll_provider_http_json(client, &provider, &target.base_url, &token).await {
3696 Ok(value) => {
3697 let snapshot = snapshot_from_provider_json(
3698 &provider,
3699 &upstreams[0],
3700 &value,
3701 &target.base_url,
3702 fetched_at_ms,
3703 stale_after_ms,
3704 );
3705 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false)
3706 .await;
3707 state
3708 .record_provider_balance_snapshot(service_name, snapshot)
3709 .await;
3710 UsageProviderRefreshOutcome::Refreshed
3711 }
3712 Err(err) => {
3713 state
3714 .record_provider_balance_snapshot(
3715 service_name,
3716 base_snapshot(&provider, &upstreams[0], fetched_at_ms, stale_after_ms)
3717 .with_error(err.to_string()),
3718 )
3719 .await;
3720 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false)
3721 .await;
3722 warn!(
3723 "OpenAI organization costs poll failed for {}[{}]: {}",
3724 target.upstream.station_name, target.upstream.index, err
3725 );
3726 UsageProviderRefreshOutcome::Failed
3727 }
3728 };
3729 }
3730
3731 let first_provider = auto_usage_provider(target, first_auto_probe_kind(target));
3732 let provider_id = first_provider.id.clone();
3733 if let Some(suppression) =
3734 usage_provider_target_suppression_active(&provider_id, target, Instant::now())
3735 {
3736 update_usage_exhausted(
3737 lb_states,
3738 state,
3739 cfg,
3740 service_name,
3741 &upstreams,
3742 suppression.routing_exhausted,
3743 )
3744 .await;
3745 warn!(
3746 "auto usage provider '{}' skipped {}[{}]: balance refresh suppressed: {}",
3747 first_provider.id,
3748 target.upstream.station_name,
3749 target.upstream.index,
3750 suppression.reason
3751 );
3752 return UsageProviderRefreshOutcome::Failed;
3753 }
3754 if let Some(decision) = existing_usage_provider_target_suppression_decision(
3755 state,
3756 cfg,
3757 service_name,
3758 &provider_id,
3759 target,
3760 fetched_at_ms,
3761 )
3762 .await
3763 {
3764 remember_usage_provider_target_suppression(
3765 &provider_id,
3766 target,
3767 decision.ttl,
3768 decision.reason.clone(),
3769 decision.routing_exhausted,
3770 Instant::now(),
3771 );
3772 update_usage_exhausted(
3773 lb_states,
3774 state,
3775 cfg,
3776 service_name,
3777 &upstreams,
3778 decision.routing_exhausted,
3779 )
3780 .await;
3781 warn!(
3782 "auto usage provider '{}' skipped {}[{}]: existing balance snapshot suppresses refresh: {}",
3783 first_provider.id, target.upstream.station_name, target.upstream.index, decision.reason
3784 );
3785 return UsageProviderRefreshOutcome::Failed;
3786 }
3787
3788 let Some(token) = resolve_token(&first_provider, &upstreams, cfg, service_name) else {
3789 state
3790 .record_provider_balance_snapshot(
3791 service_name,
3792 base_snapshot(
3793 &first_provider,
3794 &target.upstream,
3795 fetched_at_ms,
3796 stale_after_ms,
3797 )
3798 .with_error("no usable token; checked upstream auth"),
3799 )
3800 .await;
3801 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
3802 return UsageProviderRefreshOutcome::MissingToken;
3803 };
3804
3805 let probe_order = auto_probe_kind_order(&provider_id, target);
3806 if probe_order.is_empty() {
3807 let error = "all balance probe kinds are temporarily suppressed";
3808 warn!(
3809 "auto usage provider '{}' skipped {}[{}]: {}",
3810 first_provider.id, target.upstream.station_name, target.upstream.index, error
3811 );
3812 state
3813 .record_provider_balance_snapshot(
3814 service_name,
3815 base_snapshot(
3816 &first_provider,
3817 &target.upstream,
3818 fetched_at_ms,
3819 stale_after_ms,
3820 )
3821 .with_error(error),
3822 )
3823 .await;
3824 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
3825 return UsageProviderRefreshOutcome::Failed;
3826 }
3827
3828 let mut probe_errors = Vec::new();
3829 for kind in probe_order {
3830 let provider = auto_usage_provider(target, kind);
3831 match poll_provider_http_json(client, &provider, &target.base_url, &token).await {
3832 Ok(value) => {
3833 let snapshot = snapshot_from_provider_json(
3834 &provider,
3835 &upstreams[0],
3836 &value,
3837 &target.base_url,
3838 fetched_at_ms,
3839 stale_after_ms,
3840 );
3841 if auto_snapshot_is_usable(&snapshot) {
3842 remember_auto_probe_kind_success(&provider_id, target, kind);
3843 let suppression_decision = usage_provider_suppression_decision_from_snapshot(
3844 &snapshot,
3845 cfg,
3846 fetched_at_ms,
3847 );
3848 let exhausted_for_lb = suppression_decision
3849 .as_ref()
3850 .is_some_and(|decision| decision.routing_exhausted);
3851 if let Some(decision) = suppression_decision.as_ref() {
3852 remember_usage_provider_target_suppression(
3853 &provider_id,
3854 target,
3855 decision.ttl,
3856 decision.reason.as_str(),
3857 decision.routing_exhausted,
3858 Instant::now(),
3859 );
3860 } else {
3861 clear_usage_provider_target_suppression(&provider_id, target);
3862 }
3863 update_usage_exhausted(
3864 lb_states,
3865 state,
3866 cfg,
3867 service_name,
3868 &upstreams,
3869 exhausted_for_lb,
3870 )
3871 .await;
3872 state
3873 .record_provider_balance_snapshot(service_name, snapshot)
3874 .await;
3875 info!(
3876 "auto usage provider '{}' refreshed {}[{}] via {:?}, exhausted = {}",
3877 provider.id,
3878 target.upstream.station_name,
3879 target.upstream.index,
3880 kind,
3881 exhausted_for_lb
3882 );
3883 return UsageProviderRefreshOutcome::Refreshed;
3884 }
3885 remember_auto_probe_kind_failure(&provider_id, target, kind, Instant::now());
3886 let error = snapshot.error.unwrap_or_else(|| {
3887 format!("auto probe {:?} returned no usable balance fields", kind)
3888 });
3889 probe_errors.push(format!("{:?}: {}", kind, error));
3890 }
3891 Err(err) => {
3892 remember_auto_probe_kind_failure(&provider_id, target, kind, Instant::now());
3893 probe_errors.push(format!("{:?}: {}", kind, err));
3894 }
3895 }
3896 }
3897
3898 if let Some(error) = auto_probe_error_summary(&probe_errors) {
3899 warn!(
3900 "auto usage provider '{}' found no usable balance endpoint for {}[{}]: {}",
3901 first_provider.id, target.upstream.station_name, target.upstream.index, error
3902 );
3903 state
3904 .record_provider_balance_snapshot(
3905 service_name,
3906 base_snapshot(
3907 &first_provider,
3908 &target.upstream,
3909 fetched_at_ms,
3910 stale_after_ms,
3911 )
3912 .with_error(error.clone()),
3913 )
3914 .await;
3915 let terminal_failure = usage_provider_error_is_terminal(error.as_str());
3916 if terminal_failure {
3917 remember_usage_provider_target_suppression(
3918 &provider_id,
3919 target,
3920 USAGE_PROVIDER_TERMINAL_FAILURE_TTL,
3921 error.clone(),
3922 true,
3923 Instant::now(),
3924 );
3925 }
3926 update_usage_exhausted(
3927 lb_states,
3928 state,
3929 cfg,
3930 service_name,
3931 &upstreams,
3932 terminal_failure,
3933 )
3934 .await;
3935 }
3936 UsageProviderRefreshOutcome::Failed
3937}
3938
3939pub async fn refresh_balances_for_service(
3940 client: &Client,
3941 cfg: Arc<ProxyConfig>,
3942 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
3943 state: Arc<ProxyState>,
3944 service_name: &str,
3945 station_name_filter: Option<&str>,
3946 provider_id_filter: Option<&str>,
3947) -> UsageProviderRefreshSummary {
3948 if cfg!(test) {
3950 return UsageProviderRefreshSummary::default();
3951 }
3952
3953 let station_name_filter = station_name_filter
3954 .map(str::trim)
3955 .filter(|value| !value.is_empty());
3956 let provider_id_filter = provider_id_filter
3957 .map(str::trim)
3958 .filter(|value| !value.is_empty());
3959 let providers_file = load_providers();
3960 let mut summary = UsageProviderRefreshSummary {
3961 providers_configured: providers_file.providers.len(),
3962 ..UsageProviderRefreshSummary::default()
3963 };
3964
3965 let poll_map = LAST_USAGE_POLL.get_or_init(|| Mutex::new(HashMap::new()));
3966 let mut configured_jobs = Vec::new();
3967 let mut configured_job_keys = HashSet::new();
3968 for provider in &providers_file.providers {
3969 if provider_id_filter.is_some_and(|filter| filter != provider.id.as_str()) {
3970 continue;
3971 }
3972
3973 let targets = matching_provider_targets(&cfg, service_name, provider, station_name_filter);
3974 if targets.is_empty() {
3975 continue;
3976 }
3977
3978 summary.providers_matched += 1;
3979 summary.upstreams_matched += targets.len();
3980 warn_if_provider_spans_hosts(&cfg, service_name, provider);
3981
3982 let interval_secs = snapshot_refresh_interval_secs(provider);
3983 for target in targets {
3984 summary.attempted += 1;
3985 configured_job_keys.insert(target_key(&target));
3986 configured_jobs.push(ConfiguredRefreshJob {
3987 provider,
3988 target,
3989 interval_secs,
3990 });
3991 }
3992 }
3993
3994 let mut refreshed_provider_ids = HashSet::new();
3995 if !configured_jobs.is_empty() {
3996 for (provider_id, outcome) in run_configured_refresh_jobs(
3997 client,
3998 configured_jobs,
3999 &cfg,
4000 &lb_states,
4001 &state,
4002 service_name,
4003 )
4004 .await
4005 {
4006 match outcome {
4007 UsageProviderRefreshOutcome::Refreshed => {
4008 summary.refreshed += 1;
4009 refreshed_provider_ids.insert(provider_id);
4010 }
4011 UsageProviderRefreshOutcome::Failed => summary.failed += 1,
4012 UsageProviderRefreshOutcome::MissingToken => summary.missing_token += 1,
4013 }
4014 }
4015 }
4016
4017 let mut auto_jobs = Vec::new();
4018 for target in usage_provider_targets(&cfg, service_name, station_name_filter) {
4019 if configured_job_keys.contains(&target_key(&target)) {
4020 continue;
4021 }
4022 if !auto_target_matches_provider_id_filter(&target, provider_id_filter) {
4023 continue;
4024 }
4025
4026 summary.attempted += 1;
4027 summary.auto_attempted += 1;
4028 auto_jobs.push(AutoRefreshJob { target });
4029 }
4030
4031 if !auto_jobs.is_empty() {
4032 for outcome in
4033 run_auto_refresh_jobs(client, auto_jobs, &cfg, &lb_states, &state, service_name).await
4034 {
4035 match outcome {
4036 UsageProviderRefreshOutcome::Refreshed => {
4037 summary.refreshed += 1;
4038 summary.auto_refreshed += 1;
4039 }
4040 UsageProviderRefreshOutcome::Failed => {
4041 summary.failed += 1;
4042 summary.auto_failed += 1;
4043 }
4044 UsageProviderRefreshOutcome::MissingToken => {
4045 summary.missing_token += 1;
4046 }
4047 }
4048 }
4049 }
4050
4051 if !refreshed_provider_ids.is_empty()
4052 && let Ok(mut map) = poll_map.lock()
4053 {
4054 let now = Instant::now();
4055 for provider_id in refreshed_provider_ids {
4056 map.insert(provider_id, now);
4057 }
4058 }
4059
4060 summary
4061}
4062
4063pub fn enqueue_poll_for_codex_provider_endpoint(
4067 client: Client,
4068 cfg: Arc<ProxyConfig>,
4069 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
4070 state: Arc<ProxyState>,
4071 service_name: &str,
4072 provider_endpoint: ProviderEndpointKey,
4073) {
4074 let key = provider_endpoint.clone();
4075 let Some(initial_sleep_for) = enqueue_request_balance_refresh(key.clone()) else {
4076 return;
4077 };
4078
4079 let service_name = service_name.to_string();
4080 tokio::spawn(async move {
4081 let mut sleep_for = initial_sleep_for;
4082 loop {
4083 tokio::time::sleep(sleep_for).await;
4084 match take_request_balance_refresh_if_due(&key) {
4085 RequestBalanceQueueDue::Due => {}
4086 RequestBalanceQueueDue::NotDue(delay) => {
4087 sleep_for = delay;
4088 continue;
4089 }
4090 RequestBalanceQueueDue::Missing => return,
4091 }
4092
4093 let current_target = usage_provider_target_for_provider_endpoint(
4094 &cfg,
4095 service_name.as_str(),
4096 &provider_endpoint,
4097 );
4098 match poll_for_codex_target(
4099 client.clone(),
4100 cfg.clone(),
4101 lb_states.clone(),
4102 state.clone(),
4103 service_name.as_str(),
4104 current_target,
4105 )
4106 .await
4107 {
4108 RequestBalancePollOutcome::Attempted | RequestBalancePollOutcome::Skipped => {
4109 return;
4110 }
4111 RequestBalancePollOutcome::Deferred(delay) => {
4112 schedule_request_balance_refresh_at(key.clone(), Instant::now() + delay);
4113 sleep_for = delay;
4114 }
4115 }
4116 }
4117 });
4118}
4119
4120async fn poll_for_codex_target(
4121 client: Client,
4122 cfg: Arc<ProxyConfig>,
4123 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
4124 state: Arc<ProxyState>,
4125 service_name: &str,
4126 current_target: Option<UsageProviderTarget>,
4127) -> RequestBalancePollOutcome {
4128 if cfg!(test) {
4131 return RequestBalancePollOutcome::Skipped;
4132 }
4133
4134 let providers_file = load_providers();
4135 let Some(current_target) = current_target else {
4136 return RequestBalancePollOutcome::Skipped;
4137 };
4138
4139 let now = Instant::now();
4140 let poll_map = LAST_USAGE_POLL.get_or_init(|| Mutex::new(HashMap::new()));
4141 let mut matched_configured_provider = false;
4142 let mut configured_jobs = Vec::new();
4143 let mut next_cooldown = None::<Duration>;
4144
4145 for provider in &providers_file.providers {
4146 if !domain_matches(¤t_target.base_url, &provider.domains) {
4147 continue;
4148 }
4149 matched_configured_provider = true;
4150
4151 let Some(interval_secs) = effective_poll_interval_secs(provider) else {
4152 continue;
4153 };
4154
4155 {
4156 let mut map = match poll_map.lock() {
4157 Ok(m) => m,
4158 Err(_) => continue,
4159 };
4160 if let Some(last) = map.get(&provider.id)
4161 && let Some(cooldown) = remaining_poll_cooldown(*last, interval_secs, now)
4162 {
4163 next_cooldown =
4164 Some(next_cooldown.map_or(cooldown, |existing| existing.min(cooldown)));
4165 continue;
4166 }
4167 map.insert(provider.id.clone(), now);
4168 }
4169
4170 warn_if_provider_spans_hosts(&cfg, service_name, provider);
4171 configured_jobs.push(ConfiguredRefreshJob {
4172 provider,
4173 target: current_target.clone(),
4174 interval_secs,
4175 });
4176 }
4177
4178 if !configured_jobs.is_empty() {
4179 let _ = run_configured_refresh_jobs(
4180 &client,
4181 configured_jobs,
4182 &cfg,
4183 &lb_states,
4184 &state,
4185 service_name,
4186 )
4187 .await;
4188 return RequestBalancePollOutcome::Attempted;
4189 }
4190
4191 if matched_configured_provider {
4192 return next_cooldown
4193 .map(RequestBalancePollOutcome::Deferred)
4194 .unwrap_or(RequestBalancePollOutcome::Skipped);
4195 }
4196
4197 let auto_provider = if is_official_openai_base_url(¤t_target.base_url) {
4198 auto_openai_official_provider(¤t_target)
4199 } else {
4200 auto_usage_provider(¤t_target, first_auto_probe_kind(¤t_target))
4201 };
4202 let Some(interval_secs) = effective_poll_interval_secs(&auto_provider) else {
4203 return RequestBalancePollOutcome::Skipped;
4204 };
4205
4206 {
4207 let mut map = match poll_map.lock() {
4208 Ok(m) => m,
4209 Err(_) => return RequestBalancePollOutcome::Skipped,
4210 };
4211 if let Some(last) = map.get(&auto_provider.id)
4212 && let Some(cooldown) = remaining_poll_cooldown(*last, interval_secs, now)
4213 {
4214 return RequestBalancePollOutcome::Deferred(cooldown);
4215 }
4216 map.insert(auto_provider.id.clone(), now);
4217 }
4218
4219 let _ = auto_probe_provider_target(
4220 &client,
4221 ¤t_target,
4222 &cfg,
4223 &lb_states,
4224 &state,
4225 service_name,
4226 )
4227 .await;
4228 RequestBalancePollOutcome::Attempted
4229}
4230
4231#[cfg(test)]
4232mod tests {
4233 use super::*;
4234
4235 use crate::balance::BalanceSnapshotStatus;
4236 use crate::config::{ServiceConfig, UpstreamAuth, UpstreamConfig};
4237 use axum::routing::get;
4238 use std::net::SocketAddr;
4239 use std::sync::atomic::{AtomicUsize, Ordering};
4240 use tokio::net::TcpListener;
4241
4242 async fn spawn_axum_server(app: axum::Router) -> (SocketAddr, tokio::task::JoinHandle<()>) {
4243 let listener = TcpListener::bind("127.0.0.1:0").await.expect("bind");
4244 let addr = listener.local_addr().expect("local_addr");
4245 let handle = tokio::spawn(async move {
4246 axum::serve(listener, app).await.expect("serve");
4247 });
4248 (addr, handle)
4249 }
4250
4251 fn provider(id: &str, kind: ProviderKind) -> UsageProviderConfig {
4252 UsageProviderConfig {
4253 id: id.to_string(),
4254 kind,
4255 domains: vec!["example.com".to_string()],
4256 endpoint: "https://example.com/usage".to_string(),
4257 token_env: None,
4258 require_token_env: false,
4259 poll_interval_secs: Some(60),
4260 refresh_on_request: true,
4261 trust_exhaustion_for_routing: true,
4262 headers: BTreeMap::new(),
4263 variables: BTreeMap::new(),
4264 extract: UsageProviderExtractConfig::default(),
4265 }
4266 }
4267
4268 fn upstream() -> UpstreamRef {
4269 UpstreamRef {
4270 station_name: "right".to_string(),
4271 index: 1,
4272 provider_endpoint: None,
4273 }
4274 }
4275
4276 fn endpoint_upstream() -> UpstreamRef {
4277 UpstreamRef {
4278 station_name: "right".to_string(),
4279 index: 1,
4280 provider_endpoint: Some(ProviderEndpointKey::new("codex", "right", "default")),
4281 }
4282 }
4283
4284 fn upstream_config(base_url: &str) -> UpstreamConfig {
4285 UpstreamConfig {
4286 base_url: base_url.to_string(),
4287 auth: UpstreamAuth::default(),
4288 tags: HashMap::new(),
4289 supported_models: HashMap::new(),
4290 model_mapping: HashMap::new(),
4291 }
4292 }
4293
4294 fn endpoint_upstream_config(
4295 base_url: &str,
4296 provider_id: &str,
4297 endpoint_id: &str,
4298 ) -> UpstreamConfig {
4299 let mut upstream = upstream_config(base_url);
4300 upstream
4301 .tags
4302 .insert("provider_id".to_string(), provider_id.to_string());
4303 upstream
4304 .tags
4305 .insert("endpoint_id".to_string(), endpoint_id.to_string());
4306 upstream
4307 }
4308
4309 fn service_config(name: &str, upstreams: Vec<UpstreamConfig>) -> ServiceConfig {
4310 ServiceConfig {
4311 name: name.to_string(),
4312 alias: None,
4313 enabled: true,
4314 level: 1,
4315 upstreams,
4316 }
4317 }
4318
4319 fn proxy_config(stations: Vec<ServiceConfig>) -> ProxyConfig {
4320 let mut cfg = ProxyConfig::default();
4321 cfg.codex.configs = stations
4322 .into_iter()
4323 .map(|station| (station.name.clone(), station))
4324 .collect();
4325 cfg
4326 }
4327
4328 fn usage_provider_target(base_url: &str, provider_id: &str) -> UsageProviderTarget {
4329 UsageProviderTarget {
4330 upstream: UpstreamRef {
4331 station_name: "routing".to_string(),
4332 index: 0,
4333 provider_endpoint: Some(ProviderEndpointKey::new("codex", provider_id, "default")),
4334 },
4335 base_url: base_url.to_string(),
4336 provider_id: Some(provider_id.to_string()),
4337 }
4338 }
4339
4340 fn usage_provider_target_at(
4341 station_name: &str,
4342 upstream_index: usize,
4343 base_url: &str,
4344 provider_id: &str,
4345 ) -> UsageProviderTarget {
4346 UsageProviderTarget {
4347 upstream: UpstreamRef {
4348 station_name: station_name.to_string(),
4349 index: upstream_index,
4350 provider_endpoint: Some(ProviderEndpointKey::new("codex", provider_id, "default")),
4351 },
4352 base_url: base_url.to_string(),
4353 provider_id: Some(provider_id.to_string()),
4354 }
4355 }
4356
4357 fn clear_auto_probe_kind_state(provider_id: &str) {
4358 if let Some(hints) = AUTO_PROBE_KIND_HINTS.get()
4359 && let Ok(mut hints) = hints.lock()
4360 {
4361 hints.remove(provider_id);
4362 }
4363 if let Some(failures) = AUTO_PROBE_KIND_FAILURES.get()
4364 && let Ok(mut failures) = failures.lock()
4365 {
4366 failures.retain(|key, _| key.provider_id != provider_id);
4367 }
4368 clear_usage_provider_target_suppressions_for_provider(provider_id);
4369 }
4370
4371 #[test]
4372 fn budget_snapshot_reports_monthly_budget_and_exhaustion() {
4373 let snapshot = budget_snapshot_from_json(
4374 &provider("packycode", ProviderKind::BudgetHttpJson),
4375 &upstream(),
4376 &serde_json::json!({
4377 "monthly_budget_usd": "10.50",
4378 "monthly_spent_usd": 10.5
4379 }),
4380 100,
4381 Some(1_000),
4382 );
4383
4384 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
4385 assert_eq!(snapshot.exhausted, Some(true));
4386 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("10.5"));
4387 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("10.5"));
4388 }
4389
4390 #[test]
4391 fn budget_snapshot_keeps_missing_amounts_unknown() {
4392 let snapshot = budget_snapshot_from_json(
4393 &provider("packycode", ProviderKind::BudgetHttpJson),
4394 &upstream(),
4395 &serde_json::json!({}),
4396 100,
4397 Some(1_000),
4398 );
4399
4400 assert_eq!(snapshot.status, BalanceSnapshotStatus::Unknown);
4401 assert_eq!(snapshot.exhausted, None);
4402 }
4403
4404 #[test]
4405 fn yescode_snapshot_sums_subscription_and_paygo_balances() {
4406 let snapshot = yescode_snapshot_from_json(
4407 &provider("yescode", ProviderKind::YescodeProfile),
4408 &upstream(),
4409 &serde_json::json!({
4410 "subscription_balance": "1.25",
4411 "pay_as_you_go_balance": 2.5
4412 }),
4413 100,
4414 Some(1_000),
4415 );
4416
4417 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
4418 assert_eq!(snapshot.exhausted, Some(false));
4419 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("3.75"));
4420 assert_eq!(snapshot.subscription_balance_usd.as_deref(), Some("1.25"));
4421 assert_eq!(snapshot.paygo_balance_usd.as_deref(), Some("2.5"));
4422 }
4423
4424 #[test]
4425 fn openai_balance_endpoint_defaults_to_base_user_balance_without_v1() {
4426 let mut provider = provider("sub2api", ProviderKind::OpenAiBalanceHttpJson);
4427 provider.endpoint.clear();
4428
4429 let endpoint =
4430 resolve_endpoint(&provider, "https://relay.example.com/v1", "token").expect("endpoint");
4431
4432 assert_eq!(endpoint, "https://relay.example.com/user/balance");
4433 }
4434
4435 #[test]
4436 fn sub2api_usage_endpoint_defaults_to_upstream_usage_under_v1() {
4437 let mut provider = provider("sub2api", ProviderKind::Sub2ApiUsage);
4438 provider.endpoint.clear();
4439
4440 let endpoint =
4441 resolve_endpoint(&provider, "https://relay.example.com/v1", "token").expect("endpoint");
4442
4443 assert_eq!(endpoint, "https://relay.example.com/v1/usage");
4444 }
4445
4446 #[test]
4447 fn sub2api_auth_me_endpoint_defaults_to_dashboard_path_without_v1() {
4448 let mut provider = provider("sub2api-auth", ProviderKind::Sub2ApiAuthMe);
4449 provider.endpoint.clear();
4450
4451 let endpoint =
4452 resolve_endpoint(&provider, "https://relay.example.com/v1", "token").expect("endpoint");
4453
4454 assert_eq!(endpoint, "https://relay.example.com/api/v1/auth/me");
4455 }
4456
4457 #[test]
4458 fn provider_templates_support_variables_for_custom_headers_or_queries() {
4459 let mut provider = provider("newapi", ProviderKind::NewApiUserSelf);
4460 provider.endpoint = "{{base_url}}/api/user/self?user={{userId}}".to_string();
4461 provider
4462 .variables
4463 .insert("userId".to_string(), "42".to_string());
4464
4465 let endpoint = resolve_endpoint(&provider, "https://newapi.example.com/v1", "token")
4466 .expect("endpoint");
4467
4468 assert_eq!(endpoint, "https://newapi.example.com/api/user/self?user=42");
4469 }
4470
4471 #[test]
4472 fn new_api_token_usage_endpoint_defaults_to_model_key_usage_path() {
4473 let mut provider = provider("newapi-token", ProviderKind::NewApiTokenUsage);
4474 provider.endpoint.clear();
4475
4476 let endpoint = resolve_endpoint(&provider, "https://newapi.example.com/v1", "token")
4477 .expect("endpoint");
4478
4479 assert_eq!(endpoint, "https://newapi.example.com/api/usage/token/");
4480 }
4481
4482 #[test]
4483 fn openai_organization_costs_endpoint_defaults_to_official_v1_costs_window() {
4484 let mut provider = provider("openai", ProviderKind::OpenAiOrganizationCosts);
4485 provider.endpoint.clear();
4486
4487 let endpoint =
4488 resolve_endpoint(&provider, "https://api.openai.com/v1", "token").expect("endpoint");
4489
4490 assert!(endpoint.starts_with("https://api.openai.com/v1/organization/costs?start_time="));
4491 assert!(endpoint.ends_with("&limit=30"));
4492 let start_time = endpoint
4493 .split("start_time=")
4494 .nth(1)
4495 .and_then(|value| value.split('&').next())
4496 .and_then(|value| value.parse::<u64>().ok())
4497 .expect("numeric start_time");
4498 assert!(start_time > 0);
4499 }
4500
4501 #[test]
4502 fn require_token_env_prevents_upstream_model_key_fallback() {
4503 let mut cfg = proxy_config(vec![service_config(
4504 "right",
4505 vec![upstream_config("https://api.openai.com/v1")],
4506 )]);
4507 cfg.codex
4508 .configs
4509 .get_mut("right")
4510 .expect("station")
4511 .upstreams[0]
4512 .auth
4513 .auth_token = Some("model-key".to_string());
4514
4515 let mut provider = provider("openai", ProviderKind::OpenAiOrganizationCosts);
4516 provider.token_env = Some("__CODEX_HELPER_TEST_MISSING_TOKEN_ENV__".to_string());
4517 provider.require_token_env = true;
4518 let upstreams = [UpstreamRef {
4519 station_name: "right".to_string(),
4520 index: 0,
4521 provider_endpoint: None,
4522 }];
4523
4524 assert_eq!(resolve_token(&provider, &upstreams, &cfg, "codex"), None);
4525
4526 provider.require_token_env = false;
4527 assert_eq!(
4528 resolve_token(&provider, &upstreams, &cfg, "codex").as_deref(),
4529 Some("model-key")
4530 );
4531 }
4532
4533 #[test]
4534 fn effective_poll_interval_respects_disable_flag_zero_and_minimum() {
4535 let mut provider = provider("sub2api", ProviderKind::OpenAiBalanceHttpJson);
4536
4537 provider.poll_interval_secs = Some(0);
4538 assert_eq!(effective_poll_interval_secs(&provider), None);
4539
4540 provider.poll_interval_secs = Some(10);
4541 assert_eq!(
4542 effective_poll_interval_secs(&provider),
4543 Some(MIN_POLL_INTERVAL_SECS)
4544 );
4545
4546 provider.poll_interval_secs = None;
4547 assert_eq!(
4548 effective_poll_interval_secs(&provider),
4549 Some(DEFAULT_POLL_INTERVAL_SECS)
4550 );
4551
4552 provider.refresh_on_request = false;
4553 assert_eq!(effective_poll_interval_secs(&provider), None);
4554 }
4555
4556 #[test]
4557 fn usage_provider_http_error_detail_extracts_json_code_and_message() {
4558 let detail = usage_provider_http_error_detail(
4559 r#"{"code":"USER_INACTIVE","message":"User account is not active"}"#,
4560 )
4561 .expect("error detail");
4562
4563 assert_eq!(detail, "USER_INACTIVE: User account is not active");
4564 assert!(usage_provider_error_is_terminal(&detail));
4565 }
4566
4567 #[test]
4568 fn current_day_quota_exhaustion_blocks_followup_usage_even_when_display_only() {
4569 let mut snapshot = ProviderBalanceSnapshot {
4570 status: BalanceSnapshotStatus::Exhausted,
4571 exhausted: Some(true),
4572 exhaustion_affects_routing: false,
4573 quota_period: Some("daily".to_string()),
4574 quota_remaining_usd: Some("0".to_string()),
4575 ..ProviderBalanceSnapshot::default()
4576 };
4577 snapshot.refresh_status(0);
4578
4579 assert!(!snapshot.routing_exhausted());
4580 assert!(
4581 usage_provider_suppression_decision_from_snapshot(
4582 &snapshot,
4583 &ProxyConfig::default(),
4584 snapshot.fetched_at_ms,
4585 )
4586 .is_some()
4587 );
4588 assert_eq!(
4589 usage_provider_snapshot_suppression_reason(&snapshot).as_deref(),
4590 Some("daily package quota exhausted for current period")
4591 );
4592 }
4593
4594 #[test]
4595 fn current_day_quota_suppression_expires_after_configured_reset() {
4596 let cfg = ProxyConfig::default();
4597 let fetched_at_ms = 1_700_000_000_000;
4598 let reset_at_ms = next_reset_at_ms(
4599 fetched_at_ms,
4600 cfg.ui.usage_forecast.reset_utc_offset.as_str(),
4601 cfg.ui.usage_forecast.reset_time.as_str(),
4602 )
4603 .expect("default reset config is valid");
4604 let grace_ms = duration_millis_u64(USAGE_PROVIDER_DAILY_RESET_SUPPRESSION_GRACE);
4605 let suppress_until_ms = reset_at_ms + grace_ms;
4606
4607 let mut snapshot = ProviderBalanceSnapshot {
4608 fetched_at_ms,
4609 stale_after_ms: Some(fetched_at_ms + 60_000),
4610 status: BalanceSnapshotStatus::Exhausted,
4611 exhausted: Some(true),
4612 exhaustion_affects_routing: false,
4613 quota_period: Some("daily".to_string()),
4614 quota_remaining_usd: Some("0".to_string()),
4615 ..ProviderBalanceSnapshot::default()
4616 };
4617 snapshot.refresh_status(fetched_at_ms);
4618
4619 let decision = usage_provider_suppression_decision_from_snapshot(
4620 &snapshot,
4621 &cfg,
4622 suppress_until_ms - 1,
4623 )
4624 .expect("daily exhaustion should suppress until reset grace expires");
4625 assert_eq!(decision.ttl, Duration::from_millis(1));
4626 assert!(decision.routing_exhausted);
4627
4628 assert!(
4629 usage_provider_suppression_decision_from_snapshot(&snapshot, &cfg, suppress_until_ms)
4630 .is_none(),
4631 "a stale daily exhaustion snapshot must not suppress refresh after the reset boundary"
4632 );
4633 }
4634
4635 #[test]
4636 fn stale_non_daily_exhaustion_snapshot_does_not_renew_suppression() {
4637 let cfg = ProxyConfig::default();
4638 let mut snapshot = ProviderBalanceSnapshot {
4639 fetched_at_ms: 1_000,
4640 stale_after_ms: Some(2_000),
4641 status: BalanceSnapshotStatus::Exhausted,
4642 exhausted: Some(true),
4643 quota_period: Some("quota".to_string()),
4644 quota_remaining_usd: Some("0".to_string()),
4645 ..ProviderBalanceSnapshot::default()
4646 };
4647 snapshot.refresh_status(1_000);
4648
4649 let fresh_decision =
4650 usage_provider_suppression_decision_from_snapshot(&snapshot, &cfg, 1_500)
4651 .expect("fresh exhausted quota should suppress follow-up polling");
4652 assert_eq!(fresh_decision.ttl, USAGE_PROVIDER_EXHAUSTED_SUPPRESSION_TTL);
4653
4654 assert!(
4655 usage_provider_suppression_decision_from_snapshot(&snapshot, &cfg, 2_001).is_none(),
4656 "stale exhausted quota snapshots must not be used to renew suppression forever"
4657 );
4658 }
4659
4660 #[test]
4661 fn refreshable_weekly_window_exhaustion_suppresses_only_while_snapshot_is_fresh() {
4662 let cfg = ProxyConfig::default();
4663 let mut snapshot = ProviderBalanceSnapshot {
4664 fetched_at_ms: 1_000,
4665 stale_after_ms: Some(10_000),
4666 status: BalanceSnapshotStatus::Exhausted,
4667 exhausted: Some(true),
4668 exhaustion_affects_routing: false,
4669 quota_period: Some("weekly".to_string()),
4670 quota_remaining_usd: Some("0".to_string()),
4671 ..ProviderBalanceSnapshot::default()
4672 };
4673 snapshot.refresh_status(1_000);
4674
4675 let decision = usage_provider_suppression_decision_from_snapshot(&snapshot, &cfg, 4_000)
4676 .expect("fresh weekly window exhaustion should block follow-up usage");
4677 assert_eq!(decision.ttl, Duration::from_millis(6_000));
4678 assert!(decision.routing_exhausted);
4679
4680 assert!(
4681 usage_provider_suppression_decision_from_snapshot(&snapshot, &cfg, 10_001).is_none(),
4682 "weekly/monthly windows without explicit reset_at should be re-queried after staleness"
4683 );
4684 }
4685
4686 #[test]
4687 fn rate_limit_reset_at_drives_suppression_ttl() {
4688 let cfg = ProxyConfig::default();
4689 let reset_at_ms = 120_000;
4690 let suppress_until_ms =
4691 reset_at_ms + duration_millis_u64(USAGE_PROVIDER_DAILY_RESET_SUPPRESSION_GRACE);
4692 let mut snapshot = ProviderBalanceSnapshot {
4693 fetched_at_ms: 1_000,
4694 stale_after_ms: Some(10_000),
4695 status: BalanceSnapshotStatus::Exhausted,
4696 exhausted: Some(true),
4697 quota_period: Some("rate_limit:5h".to_string()),
4698 quota_resets_at_ms: Some(reset_at_ms),
4699 ..ProviderBalanceSnapshot::default()
4700 };
4701 snapshot.refresh_status(1_000);
4702
4703 let decision = usage_provider_suppression_decision_from_snapshot(
4704 &snapshot,
4705 &cfg,
4706 suppress_until_ms - 1,
4707 )
4708 .expect("rate limit should suppress until reset grace expires");
4709 assert_eq!(decision.ttl, Duration::from_millis(1));
4710
4711 assert!(
4712 usage_provider_suppression_decision_from_snapshot(&snapshot, &cfg, suppress_until_ms)
4713 .is_none()
4714 );
4715 }
4716
4717 #[test]
4718 fn auto_provider_uses_stable_target_id_across_probe_kinds() {
4719 let target = UsageProviderTarget {
4720 upstream: UpstreamRef {
4721 station_name: "input/sub".to_string(),
4722 index: 2,
4723 provider_endpoint: None,
4724 },
4725 base_url: "https://ai.input.im/v1".to_string(),
4726 provider_id: None,
4727 };
4728
4729 let sub2api = auto_usage_provider(&target, ProviderKind::Sub2ApiUsage);
4730 let newapi_token = auto_usage_provider(&target, ProviderKind::NewApiTokenUsage);
4731 let newapi = auto_usage_provider(&target, ProviderKind::NewApiUserSelf);
4732
4733 assert_eq!(sub2api.id, "auto:balance:input-sub:2");
4734 assert_eq!(sub2api.id, newapi_token.id);
4735 assert_eq!(sub2api.id, newapi.id);
4736 assert_eq!(sub2api.domains, vec!["ai.input.im".to_string()]);
4737 assert_eq!(
4738 resolve_endpoint(&sub2api, &target.base_url, "token").unwrap(),
4739 "https://ai.input.im/v1/usage"
4740 );
4741 assert_eq!(
4742 resolve_endpoint(&newapi_token, &target.base_url, "token").unwrap(),
4743 "https://ai.input.im/api/usage/token/"
4744 );
4745 }
4746
4747 #[test]
4748 fn auto_probe_prefers_rightcode_adapter_for_rightcode_hosts() {
4749 let target = UsageProviderTarget {
4750 upstream: UpstreamRef {
4751 station_name: "right".to_string(),
4752 index: 0,
4753 provider_endpoint: None,
4754 },
4755 base_url: "https://www.right.codes/codex/v1".to_string(),
4756 provider_id: Some("right".to_string()),
4757 };
4758
4759 assert_eq!(
4760 first_auto_probe_kind(&target),
4761 ProviderKind::RightCodeAccountSummary
4762 );
4763 assert_eq!(
4764 resolve_endpoint(
4765 &auto_usage_provider(&target, ProviderKind::RightCodeAccountSummary),
4766 &target.base_url,
4767 "token"
4768 )
4769 .unwrap(),
4770 "https://www.right.codes/account/summary"
4771 );
4772 assert_eq!(
4773 auto_usage_provider(&target, ProviderKind::RightCodeAccountSummary)
4774 .token_env
4775 .as_deref(),
4776 None
4777 );
4778 }
4779
4780 #[test]
4781 fn auto_provider_id_prefers_runtime_provider_tag() {
4782 let target = UsageProviderTarget {
4783 upstream: UpstreamRef {
4784 station_name: "routing".to_string(),
4785 index: 0,
4786 provider_endpoint: Some(ProviderEndpointKey::new("codex", "input", "default")),
4787 },
4788 base_url: "https://ai.input.im/v1".to_string(),
4789 provider_id: Some("input".to_string()),
4790 };
4791
4792 let provider = auto_usage_provider(&target, ProviderKind::Sub2ApiUsage);
4793
4794 assert_eq!(provider.id, "input");
4795 }
4796
4797 #[test]
4798 fn auto_target_provider_id_filter_matches_runtime_provider_tag() {
4799 let target = UsageProviderTarget {
4800 upstream: UpstreamRef {
4801 station_name: "routing".to_string(),
4802 index: 6,
4803 provider_endpoint: Some(ProviderEndpointKey::new("codex", "input6", "default")),
4804 },
4805 base_url: "https://input.9z1.me/v1".to_string(),
4806 provider_id: Some("input6".to_string()),
4807 };
4808
4809 assert!(auto_target_matches_provider_id_filter(&target, None));
4810 assert!(auto_target_matches_provider_id_filter(
4811 &target,
4812 Some("input6")
4813 ));
4814 assert!(!auto_target_matches_provider_id_filter(
4815 &target,
4816 Some("input5")
4817 ));
4818 }
4819
4820 #[test]
4821 fn auto_target_provider_id_filter_matches_generated_auto_id() {
4822 let target = UsageProviderTarget {
4823 upstream: UpstreamRef {
4824 station_name: "routing".to_string(),
4825 index: 6,
4826 provider_endpoint: None,
4827 },
4828 base_url: "https://input.9z1.me/v1".to_string(),
4829 provider_id: None,
4830 };
4831
4832 assert!(auto_target_matches_provider_id_filter(
4833 &target,
4834 Some("auto:balance:routing:6")
4835 ));
4836 assert!(!auto_target_matches_provider_id_filter(
4837 &target,
4838 Some("input6")
4839 ));
4840 }
4841
4842 #[test]
4843 fn provider_endpoint_target_lookup_uses_endpoint_identity() {
4844 let cfg = proxy_config(vec![service_config(
4845 "routing",
4846 vec![
4847 endpoint_upstream_config("https://input.example/v1", "input", "default"),
4848 endpoint_upstream_config("https://right.example/v1", "right", "default"),
4849 ],
4850 )]);
4851
4852 let target = usage_provider_target_for_provider_endpoint(
4853 &cfg,
4854 "codex",
4855 &ProviderEndpointKey::new("codex", "right", "default"),
4856 )
4857 .expect("provider endpoint target");
4858
4859 assert_eq!(target.upstream.station_name, "routing");
4860 assert_eq!(target.upstream.index, 1);
4861 assert_eq!(
4862 target
4863 .upstream
4864 .provider_endpoint
4865 .as_ref()
4866 .map(ProviderEndpointKey::stable_key)
4867 .as_deref(),
4868 Some("codex/right/default")
4869 );
4870 assert_eq!(target.base_url, "https://right.example/v1");
4871 assert_eq!(target.provider_id.as_deref(), Some("right"));
4872 }
4873
4874 #[test]
4875 fn request_balance_queue_deduplicates_until_due() {
4876 let key = ProviderEndpointKey::new("codex", "input", "default");
4877 let queue = REQUEST_BALANCE_QUEUE.get_or_init(|| Mutex::new(HashMap::new()));
4878 {
4879 let mut queue = queue.lock().expect("queue");
4880 queue.remove(&key);
4881 }
4882
4883 assert_eq!(
4884 enqueue_request_balance_refresh(key.clone()),
4885 Some(REQUEST_BALANCE_REFRESH_DELAY)
4886 );
4887 assert_eq!(enqueue_request_balance_refresh(key.clone()), None);
4888 assert!(matches!(
4889 take_request_balance_refresh_if_due(&key),
4890 RequestBalanceQueueDue::NotDue(_)
4891 ));
4892
4893 {
4894 let mut queue = queue.lock().expect("queue");
4895 queue.insert(key.clone(), Instant::now() - Duration::from_secs(1));
4896 }
4897
4898 assert_eq!(
4899 take_request_balance_refresh_if_due(&key),
4900 RequestBalanceQueueDue::Due
4901 );
4902 assert_eq!(
4903 take_request_balance_refresh_if_due(&key),
4904 RequestBalanceQueueDue::Missing
4905 );
4906 assert_eq!(
4907 enqueue_request_balance_refresh(key.clone()),
4908 Some(REQUEST_BALANCE_REFRESH_DELAY)
4909 );
4910
4911 queue.lock().expect("queue").remove(&key);
4912 }
4913
4914 #[test]
4915 fn request_balance_queue_does_not_extend_due_refresh() {
4916 let key = ProviderEndpointKey::new("codex", "input", "default");
4917 let queue = REQUEST_BALANCE_QUEUE.get_or_init(|| Mutex::new(HashMap::new()));
4918 {
4919 let mut queue = queue.lock().expect("queue");
4920 queue.insert(key.clone(), Instant::now() - Duration::from_secs(1));
4921 }
4922
4923 assert_eq!(
4924 enqueue_request_balance_refresh(key.clone()),
4925 Some(Duration::ZERO)
4926 );
4927 assert_eq!(
4928 take_request_balance_refresh_if_due(&key),
4929 RequestBalanceQueueDue::Due
4930 );
4931 assert_eq!(
4932 take_request_balance_refresh_if_due(&key),
4933 RequestBalanceQueueDue::Missing
4934 );
4935
4936 queue.lock().expect("queue").remove(&key);
4937 }
4938
4939 #[test]
4940 fn remaining_poll_cooldown_returns_only_unexpired_window() {
4941 let now = Instant::now();
4942 assert_eq!(
4943 remaining_poll_cooldown(now - Duration::from_secs(30), 60, now),
4944 Some(Duration::from_secs(30))
4945 );
4946 assert_eq!(
4947 remaining_poll_cooldown(now - Duration::from_secs(60), 60, now),
4948 None
4949 );
4950 assert_eq!(
4951 remaining_poll_cooldown(now - Duration::from_secs(61), 60, now),
4952 None
4953 );
4954 }
4955
4956 #[test]
4957 fn configured_target_keys_prevent_auto_probe_for_explicit_balance_domains() {
4958 let cfg = proxy_config(vec![
4959 service_config("explicit", vec![upstream_config("https://example.com/v1")]),
4960 service_config("auto", vec![upstream_config("https://ai.input.im/v1")]),
4961 ]);
4962 let configured = configured_target_keys(
4963 &cfg,
4964 "codex",
4965 &[provider("relay", ProviderKind::OpenAiBalanceHttpJson)],
4966 None,
4967 );
4968 let auto_targets = usage_provider_targets(&cfg, "codex", None)
4969 .into_iter()
4970 .filter(|target| !configured.contains(&target_key(target)))
4971 .map(|target| target.upstream.station_name)
4972 .collect::<Vec<_>>();
4973
4974 assert_eq!(auto_targets, vec!["auto".to_string()]);
4975 }
4976
4977 #[test]
4978 fn auto_probe_accepts_only_usable_balance_snapshots() {
4979 let usable = sub2api_usage_snapshot_from_json(
4980 &provider("auto", ProviderKind::Sub2ApiUsage),
4981 &upstream(),
4982 &serde_json::json!({
4983 "isValid": true,
4984 "remaining": 1
4985 }),
4986 100,
4987 Some(1_000),
4988 );
4989 let unusable = balance_http_snapshot_from_json(
4990 &provider("auto", ProviderKind::OpenAiBalanceHttpJson),
4991 &upstream(),
4992 &serde_json::json!({ "ok": true }),
4993 100,
4994 Some(1_000),
4995 );
4996
4997 assert!(auto_snapshot_is_usable(&usable));
4998 assert!(!auto_snapshot_is_usable(&unusable));
4999 }
5000
5001 #[test]
5002 fn auto_probe_error_summary_keeps_all_attempt_failures() {
5003 let errors = vec![
5004 "Sub2ApiUsage: HTTP 404".to_string(),
5005 "NewApiTokenUsage: missing quota fields".to_string(),
5006 "OpenAiBalanceHttpJson: non-JSON response".to_string(),
5007 ];
5008
5009 let summary = auto_probe_error_summary(&errors).expect("summary");
5010
5011 assert!(summary.contains("Sub2ApiUsage: HTTP 404"));
5012 assert!(summary.contains("NewApiTokenUsage: missing quota fields"));
5013 assert!(summary.contains("OpenAiBalanceHttpJson: non-JSON response"));
5014 }
5015
5016 #[test]
5017 fn auto_probe_kind_order_prioritizes_remembered_success() {
5018 let provider_id = "input-order-success";
5019 clear_auto_probe_kind_state(provider_id);
5020 let target = usage_provider_target("https://relay.example.com/v1", provider_id);
5021
5022 remember_auto_probe_kind_success(provider_id, &target, ProviderKind::NewApiUserSelf);
5023
5024 let order = auto_probe_kind_order(provider_id, &target);
5025
5026 assert_eq!(order.first(), Some(&ProviderKind::NewApiUserSelf));
5027 assert_eq!(
5028 order
5029 .iter()
5030 .filter(|kind| **kind == ProviderKind::Sub2ApiUsage)
5031 .count(),
5032 1
5033 );
5034 clear_auto_probe_kind_state(provider_id);
5035 }
5036
5037 #[test]
5038 fn auto_probe_kind_order_temporarily_skips_recent_failures() {
5039 let provider_id = "input-order-failure";
5040 clear_auto_probe_kind_state(provider_id);
5041 let target = usage_provider_target("https://relay.example.com/v1", provider_id);
5042 let now = Instant::now();
5043
5044 remember_auto_probe_kind_failure(provider_id, &target, ProviderKind::Sub2ApiUsage, now);
5045
5046 let order = auto_probe_kind_order(provider_id, &target);
5047
5048 assert!(!order.contains(&ProviderKind::Sub2ApiUsage));
5049 assert!(order.contains(&ProviderKind::NewApiTokenUsage));
5050
5051 if let Some(failures) = AUTO_PROBE_KIND_FAILURES.get()
5052 && let Ok(mut failures) = failures.lock()
5053 {
5054 failures.insert(
5055 AutoProbeKindFailureKey {
5056 provider_id: provider_id.to_string(),
5057 target: auto_probe_target_key(&target),
5058 kind: ProviderKind::Sub2ApiUsage,
5059 },
5060 now - AUTO_PROBE_KIND_FAILURE_TTL - Duration::from_secs(1),
5061 );
5062 }
5063
5064 let order = auto_probe_kind_order(provider_id, &target);
5065 assert!(order.contains(&ProviderKind::Sub2ApiUsage));
5066 clear_auto_probe_kind_state(provider_id);
5067 }
5068
5069 #[test]
5070 fn auto_probe_kind_order_can_be_empty_when_all_kinds_are_suppressed() {
5071 let provider_id = "input-order-all-suppressed";
5072 clear_auto_probe_kind_state(provider_id);
5073 let target = usage_provider_target("https://relay.example.com/v1", provider_id);
5074 let now = Instant::now();
5075
5076 for kind in AUTO_PROBE_KINDS {
5077 if kind == ProviderKind::RightCodeAccountSummary {
5078 continue;
5079 }
5080 remember_auto_probe_kind_failure(provider_id, &target, kind, now);
5081 }
5082
5083 assert!(auto_probe_kind_order(provider_id, &target).is_empty());
5084 clear_auto_probe_kind_state(provider_id);
5085 }
5086
5087 #[test]
5088 fn auto_probe_kind_failures_do_not_suppress_distinct_targets_with_same_provider_id() {
5089 let provider_id = "input-shared-provider";
5090 clear_auto_probe_kind_state(provider_id);
5091 let routing_target =
5092 usage_provider_target_at("routing", 0, "https://relay.example.com/v1", provider_id);
5093 let catalog_target =
5094 usage_provider_target_at("input", 0, "https://relay.example.com/v1", provider_id);
5095 let now = Instant::now();
5096
5097 for kind in AUTO_PROBE_KINDS {
5098 if kind == ProviderKind::RightCodeAccountSummary {
5099 continue;
5100 }
5101 remember_auto_probe_kind_failure(provider_id, &routing_target, kind, now);
5102 }
5103
5104 assert!(auto_probe_kind_order(provider_id, &routing_target).is_empty());
5105 assert!(
5106 !auto_probe_kind_order(provider_id, &catalog_target).is_empty(),
5107 "a routing target's temporary failures must not hide catalog balance probes for the same provider"
5108 );
5109 clear_auto_probe_kind_state(provider_id);
5110 }
5111
5112 #[tokio::test]
5113 async fn auto_probe_suppressed_order_records_error_snapshot() {
5114 let provider_id = "input-suppressed-snapshot";
5115 clear_auto_probe_kind_state(provider_id);
5116 let mut upstream =
5117 endpoint_upstream_config("https://relay.example.com/v1", provider_id, "default");
5118 upstream.auth.auth_token = Some("model-key".to_string());
5119 let cfg = proxy_config(vec![service_config("routing", vec![upstream])]);
5120 let lb_states = Arc::new(Mutex::new(HashMap::new()));
5121 let state = ProxyState::new();
5122 let target = usage_provider_target("https://relay.example.com/v1", provider_id);
5123 let now = Instant::now();
5124
5125 for kind in AUTO_PROBE_KINDS {
5126 if kind == ProviderKind::RightCodeAccountSummary {
5127 continue;
5128 }
5129 remember_auto_probe_kind_failure(provider_id, &target, kind, now);
5130 }
5131
5132 let outcome =
5133 auto_probe_provider_target(&Client::new(), &target, &cfg, &lb_states, &state, "codex")
5134 .await;
5135
5136 assert_eq!(outcome, UsageProviderRefreshOutcome::Failed);
5137 let view = state.get_provider_balance_view("codex").await;
5138 let snapshot = view
5139 .get("routing")
5140 .and_then(|snapshots| {
5141 snapshots
5142 .iter()
5143 .find(|snapshot| snapshot.provider_id == provider_id)
5144 })
5145 .expect("suppressed auto probe snapshot");
5146 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
5147 assert_eq!(
5148 snapshot.error.as_deref(),
5149 Some("all balance probe kinds are temporarily suppressed")
5150 );
5151 assert_eq!(snapshot.upstream_index, Some(0));
5152 let guard = lb_states.lock().expect("lb states");
5153 assert!(
5154 !guard
5155 .get("routing")
5156 .and_then(|entry| entry.usage_exhausted.first())
5157 .copied()
5158 .unwrap_or(true)
5159 );
5160 clear_auto_probe_kind_state(provider_id);
5161 }
5162
5163 #[tokio::test]
5164 async fn auto_probe_terminal_auth_failure_keeps_route_exhausted_during_suppression() {
5165 let provider_id = "input-terminal-auth";
5166 clear_auto_probe_kind_state(provider_id);
5167 let request_count = Arc::new(AtomicUsize::new(0));
5168 let counter = request_count.clone();
5169 let app = axum::Router::new().fallback(get(move || {
5170 let counter = counter.clone();
5171 async move {
5172 counter.fetch_add(1, Ordering::SeqCst);
5173 (
5174 axum::http::StatusCode::UNAUTHORIZED,
5175 axum::Json(serde_json::json!({
5176 "code": "USER_INACTIVE",
5177 "message": "User account is not active"
5178 })),
5179 )
5180 }
5181 }));
5182 let (addr, handle) = spawn_axum_server(app).await;
5183 let base_url = format!("http://{addr}/v1");
5184 let mut upstream = endpoint_upstream_config(&base_url, provider_id, "default");
5185 upstream.auth.auth_token = Some("model-key".to_string());
5186 let cfg = proxy_config(vec![service_config("routing", vec![upstream])]);
5187 let lb_states = Arc::new(Mutex::new(HashMap::new()));
5188 let state = ProxyState::new();
5189 let target = usage_provider_target(&base_url, provider_id);
5190
5191 let outcome =
5192 auto_probe_provider_target(&Client::new(), &target, &cfg, &lb_states, &state, "codex")
5193 .await;
5194
5195 assert_eq!(outcome, UsageProviderRefreshOutcome::Failed);
5196 assert!(request_count.load(Ordering::SeqCst) > 0);
5197 {
5198 let guard = lb_states.lock().expect("lb states");
5199 assert!(
5200 guard
5201 .get("routing")
5202 .and_then(|entry| entry.usage_exhausted.first())
5203 .copied()
5204 .unwrap_or(false)
5205 );
5206 }
5207 let view = state.get_provider_balance_view("codex").await;
5208 let snapshot = view
5209 .get("routing")
5210 .and_then(|snapshots| {
5211 snapshots
5212 .iter()
5213 .find(|snapshot| snapshot.provider_id == provider_id)
5214 })
5215 .expect("terminal auth failure snapshot");
5216 assert!(
5217 snapshot
5218 .error
5219 .as_deref()
5220 .unwrap_or("")
5221 .contains("User account is not active")
5222 );
5223
5224 let requests_after_first_probe = request_count.load(Ordering::SeqCst);
5225 let suppressed_outcome =
5226 auto_probe_provider_target(&Client::new(), &target, &cfg, &lb_states, &state, "codex")
5227 .await;
5228
5229 assert_eq!(suppressed_outcome, UsageProviderRefreshOutcome::Failed);
5230 assert_eq!(
5231 request_count.load(Ordering::SeqCst),
5232 requests_after_first_probe
5233 );
5234 let guard = lb_states.lock().expect("lb states");
5235 assert!(
5236 guard
5237 .get("routing")
5238 .and_then(|entry| entry.usage_exhausted.first())
5239 .copied()
5240 .unwrap_or(false)
5241 );
5242 clear_auto_probe_kind_state(provider_id);
5243 handle.abort();
5244 }
5245
5246 #[tokio::test]
5247 async fn auto_probe_daily_package_exhaustion_suppresses_followup_refresh() {
5248 let provider_id = "input-daily-exhausted";
5249 clear_auto_probe_kind_state(provider_id);
5250 let request_count = Arc::new(AtomicUsize::new(0));
5251 let counter = request_count.clone();
5252 let app = axum::Router::new().fallback(get(move || {
5253 let counter = counter.clone();
5254 async move {
5255 counter.fetch_add(1, Ordering::SeqCst);
5256 axum::Json(serde_json::json!({
5257 "isValid": true,
5258 "mode": "unrestricted",
5259 "planName": "CodeX Lite",
5260 "subscription": {
5261 "daily_usage_usd": 100,
5262 "daily_limit_usd": 100,
5263 "weekly_usage_usd": 100,
5264 "weekly_limit_usd": 0,
5265 "monthly_usage_usd": 100,
5266 "monthly_limit_usd": 0
5267 }
5268 }))
5269 }
5270 }));
5271 let (addr, handle) = spawn_axum_server(app).await;
5272 let base_url = format!("http://{addr}/v1");
5273 let mut upstream = endpoint_upstream_config(&base_url, provider_id, "default");
5274 upstream.auth.auth_token = Some("model-key".to_string());
5275 let cfg = proxy_config(vec![service_config("routing", vec![upstream])]);
5276 let lb_states = Arc::new(Mutex::new(HashMap::new()));
5277 let state = ProxyState::new();
5278 let target = usage_provider_target(&base_url, provider_id);
5279
5280 let outcome =
5281 auto_probe_provider_target(&Client::new(), &target, &cfg, &lb_states, &state, "codex")
5282 .await;
5283
5284 assert_eq!(outcome, UsageProviderRefreshOutcome::Refreshed);
5285 assert_eq!(request_count.load(Ordering::SeqCst), 1);
5286 {
5287 let guard = lb_states.lock().expect("lb states");
5288 assert!(
5289 guard
5290 .get("routing")
5291 .and_then(|entry| entry.usage_exhausted.first())
5292 .copied()
5293 .unwrap_or(false)
5294 );
5295 }
5296 let view = state.get_provider_balance_view("codex").await;
5297 let snapshot = view
5298 .get("routing")
5299 .and_then(|snapshots| {
5300 snapshots
5301 .iter()
5302 .find(|snapshot| snapshot.provider_id == provider_id)
5303 })
5304 .expect("daily exhausted snapshot");
5305 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5306 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
5307 assert!(snapshot.routing_ignored_exhaustion());
5308
5309 let suppressed_outcome =
5310 auto_probe_provider_target(&Client::new(), &target, &cfg, &lb_states, &state, "codex")
5311 .await;
5312
5313 assert_eq!(suppressed_outcome, UsageProviderRefreshOutcome::Failed);
5314 assert_eq!(request_count.load(Ordering::SeqCst), 1);
5315 let guard = lb_states.lock().expect("lb states");
5316 assert!(
5317 guard
5318 .get("routing")
5319 .and_then(|entry| entry.usage_exhausted.first())
5320 .copied()
5321 .unwrap_or(false)
5322 );
5323 clear_auto_probe_kind_state(provider_id);
5324 handle.abort();
5325 }
5326
5327 #[test]
5328 fn openai_balance_snapshot_reads_common_sub2api_balance_shape() {
5329 let snapshot = balance_http_snapshot_from_json(
5330 &provider("sub2api", ProviderKind::OpenAiBalanceHttpJson),
5331 &upstream(),
5332 &serde_json::json!({
5333 "balance": "1.25"
5334 }),
5335 100,
5336 Some(1_000),
5337 );
5338
5339 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5340 assert_eq!(snapshot.exhausted, Some(false));
5341 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("1.25"));
5342 }
5343
5344 #[test]
5345 fn json_path_supports_array_indices_for_official_balance_shapes() {
5346 let value = serde_json::json!({
5347 "balance_infos": [
5348 { "currency": "CNY", "total_balance": "3.25" }
5349 ]
5350 });
5351
5352 assert_eq!(
5353 json_value_at_path(&value, "balance_infos.0.total_balance")
5354 .and_then(|value| value.as_str()),
5355 Some("3.25")
5356 );
5357 }
5358
5359 #[test]
5360 fn openai_balance_snapshot_reads_cc_switch_official_balance_shapes() {
5361 let snapshot = balance_http_snapshot_from_json(
5362 &provider("deepseek", ProviderKind::OpenAiBalanceHttpJson),
5363 &upstream(),
5364 &serde_json::json!({
5365 "balance_infos": [
5366 { "currency": "CNY", "total_balance": "3.25" }
5367 ],
5368 "is_available": true
5369 }),
5370 100,
5371 Some(1_000),
5372 );
5373
5374 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5375 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("3.25"));
5376
5377 let snapshot = balance_http_snapshot_from_json(
5378 &provider("siliconflow", ProviderKind::OpenAiBalanceHttpJson),
5379 &upstream(),
5380 &serde_json::json!({
5381 "code": 20000,
5382 "data": {
5383 "totalBalance": "8.5",
5384 "chargeBalance": "2.5"
5385 }
5386 }),
5387 100,
5388 Some(1_000),
5389 );
5390
5391 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5392 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("8.5"));
5393 assert_eq!(snapshot.paygo_balance_usd.as_deref(), Some("2.5"));
5394 }
5395
5396 #[test]
5397 fn openai_balance_snapshot_can_derive_remaining_from_total_and_used() {
5398 let mut provider = provider("openrouter", ProviderKind::OpenAiBalanceHttpJson);
5399 provider.extract.monthly_budget_paths = vec!["data.total_credits".to_string()];
5400 provider.extract.monthly_spent_paths = vec!["data.total_usage".to_string()];
5401 provider.extract.derive_remaining_from_budget_and_spent = true;
5402
5403 let snapshot = balance_http_snapshot_from_json(
5404 &provider,
5405 &upstream(),
5406 &serde_json::json!({
5407 "data": {
5408 "total_credits": "10",
5409 "total_usage": "4"
5410 }
5411 }),
5412 100,
5413 Some(1_000),
5414 );
5415
5416 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5417 assert_eq!(snapshot.exhausted, Some(false));
5418 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("6"));
5419 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("10"));
5420 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("4"));
5421 }
5422
5423 #[test]
5424 fn openai_balance_snapshot_supports_divisor_for_minor_units() {
5425 let mut provider = provider("novita", ProviderKind::OpenAiBalanceHttpJson);
5426 provider.extract.remaining_balance_paths = vec!["availableBalance".to_string()];
5427 provider.extract.remaining_divisor = Some(10_000);
5428
5429 let snapshot = balance_http_snapshot_from_json(
5430 &provider,
5431 &upstream(),
5432 &serde_json::json!({
5433 "availableBalance": 12345
5434 }),
5435 100,
5436 Some(1_000),
5437 );
5438
5439 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5440 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("1.2345"));
5441 }
5442
5443 #[test]
5444 fn sub2api_usage_snapshot_reads_all_api_hub_usage_shape() {
5445 let snapshot = sub2api_usage_snapshot_from_json(
5446 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5447 &upstream(),
5448 &serde_json::json!({
5449 "isValid": true,
5450 "mode": "unrestricted",
5451 "planName": "CodeX Air",
5452 "remaining": 165.0877165,
5453 "usage": {
5454 "today": {
5455 "cost": 0,
5456 "requests": 0,
5457 "total_tokens": 0
5458 },
5459 "total": {
5460 "cost": 354.194748,
5461 "requests": 2691,
5462 "total_tokens": 384084697
5463 }
5464 }
5465 }),
5466 100,
5467 Some(1_000),
5468 );
5469
5470 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5471 assert_eq!(snapshot.exhausted, Some(false));
5472 assert_eq!(snapshot.plan_name.as_deref(), Some("CodeX Air"));
5473 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("165.0877165"));
5474 assert_eq!(snapshot.total_used_usd.as_deref(), Some("354.194748"));
5475 assert_eq!(snapshot.today_used_usd.as_deref(), Some("0"));
5476 assert_eq!(snapshot.total_requests, Some(2691));
5477 assert_eq!(snapshot.today_requests, Some(0));
5478 assert_eq!(snapshot.total_tokens, Some(384084697));
5479 assert_eq!(snapshot.today_tokens, Some(0));
5480 }
5481
5482 #[test]
5483 fn sub2api_usage_snapshot_reads_rates_model_stats_windows_and_alerts() {
5484 let snapshot = sub2api_usage_snapshot_from_json(
5485 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5486 &upstream(),
5487 &serde_json::json!({
5488 "isValid": true,
5489 "mode": "unrestricted",
5490 "plan_name": "CodeX Pro",
5491 "remaining": 9,
5492 "subscription": {
5493 "daily_usage_usd": 95,
5494 "daily_limit_usd": 100,
5495 "weekly_usage_usd": "120.5",
5496 "weekly_limit_usd": 0,
5497 "monthly_usage_usd": 300.25,
5498 "monthly_limit_usd": 1000,
5499 "expires_at": "2026-05-09T12:00:00.000Z"
5500 },
5501 "usage": {
5502 "today": {
5503 "request_count": "7",
5504 "input_tokens": 100,
5505 "output_tokens": 25,
5506 "total_cost_usd": "1.5"
5507 },
5508 "total": {
5509 "requests": 42,
5510 "tokens": 1234,
5511 "cost": 9.25
5512 },
5513 "average_duration_ms": "842.7",
5514 "rpm": "0.7",
5515 "tpm": 85.3
5516 },
5517 "model_stats": [
5518 {
5519 "model": "gpt-4o-mini",
5520 "request_count": "7",
5521 "prompt_tokens": 100,
5522 "completion_tokens": 25,
5523 "input_cost": "0.12",
5524 "output_cost": "0.34"
5525 }
5526 ]
5527 }),
5528 100,
5529 Some(1_000),
5530 );
5531
5532 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5533 assert_eq!(snapshot.plan_name.as_deref(), Some("CodeX Pro"));
5534 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("9"));
5535 assert_eq!(snapshot.today_requests, Some(7));
5536 assert_eq!(snapshot.today_tokens, Some(100));
5537 assert_eq!(snapshot.today_used_usd.as_deref(), Some("1.5"));
5538 assert_eq!(snapshot.total_requests, Some(42));
5539 assert_eq!(snapshot.total_tokens, Some(1234));
5540 assert_eq!(snapshot.total_used_usd.as_deref(), Some("9.25"));
5541 let rate = snapshot.usage_rate.expect("rate");
5542 assert_eq!(rate.average_duration_ms.as_deref(), Some("842.7"));
5543 assert_eq!(rate.rpm.as_deref(), Some("0.7"));
5544 assert_eq!(rate.tpm.as_deref(), Some("85.3"));
5545 assert_eq!(snapshot.usage_windows.len(), 3);
5546 assert_eq!(snapshot.usage_windows[0].period, "daily");
5547 assert_eq!(
5548 snapshot.usage_windows[0].remaining_usd.as_deref(),
5549 Some("5")
5550 );
5551 assert_eq!(snapshot.usage_windows[1].unlimited, Some(true));
5552 assert_eq!(snapshot.usage_model_stats.len(), 1);
5553 assert_eq!(snapshot.usage_model_stats[0].model, "gpt-4o-mini");
5554 assert_eq!(snapshot.usage_model_stats[0].request_count, Some(7));
5555 assert_eq!(snapshot.usage_model_stats[0].total_tokens, Some(125));
5556 assert_eq!(
5557 snapshot.usage_model_stats[0].total_cost_usd.as_deref(),
5558 Some("0.46")
5559 );
5560 assert_eq!(
5561 snapshot
5562 .usage_alerts
5563 .iter()
5564 .map(|alert| alert.kind)
5565 .collect::<Vec<_>>(),
5566 vec![
5567 ProviderUsageAlertKind::DailyUsage95,
5568 ProviderUsageAlertKind::LowBalance,
5569 ProviderUsageAlertKind::SubscriptionExpired,
5570 ]
5571 );
5572 }
5573
5574 #[test]
5575 fn sub2api_subscription_zero_remaining_is_display_only_period_capacity_exhaustion() {
5576 let snapshot = sub2api_usage_snapshot_from_json(
5577 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5578 &upstream(),
5579 &serde_json::json!({
5580 "isValid": true,
5581 "mode": "unrestricted",
5582 "planName": "CodeX Lite 年度",
5583 "remaining": 0,
5584 "subscription": {
5585 "daily_usage_usd": 100.468025,
5586 "daily_limit_usd": 100,
5587 "weekly_usage_usd": 401.441684,
5588 "weekly_limit_usd": 0,
5589 "monthly_usage_usd": 401.441684,
5590 "monthly_limit_usd": 0
5591 },
5592 "usage": {
5593 "today": { "cost": 0, "requests": 0, "total_tokens": 0 },
5594 "total": { "cost": 702.492098, "requests": 42, "total_tokens": 1234 }
5595 }
5596 }),
5597 100,
5598 Some(1_000),
5599 );
5600
5601 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5602 assert_eq!(snapshot.exhausted, Some(true));
5603 assert_eq!(snapshot.plan_name.as_deref(), Some("CodeX Lite 年度"));
5604 assert_eq!(snapshot.total_balance_usd, None);
5605 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
5606 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
5607 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("100"));
5608 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("100.468025"));
5609 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("100"));
5610 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("100.468025"));
5611 assert_eq!(snapshot.total_used_usd.as_deref(), Some("702.492098"));
5612 assert_eq!(snapshot.today_used_usd.as_deref(), Some("0"));
5613 assert!(
5614 !snapshot.routing_exhausted(),
5615 "sub2api /v1/usage skips billing checks; subscription windows are reset lazily on real requests"
5616 );
5617 }
5618
5619 #[test]
5620 fn sub2api_quota_limited_zero_remaining_still_marks_exhausted() {
5621 let snapshot = sub2api_usage_snapshot_from_json(
5622 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5623 &upstream(),
5624 &serde_json::json!({
5625 "isValid": true,
5626 "mode": "quota_limited",
5627 "quota": {
5628 "limit": 10,
5629 "used": 10,
5630 "remaining": 0,
5631 "unit": "USD"
5632 }
5633 }),
5634 100,
5635 Some(1_000),
5636 );
5637
5638 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5639 assert_eq!(snapshot.exhausted, Some(true));
5640 assert_eq!(snapshot.total_balance_usd, None);
5641 assert_eq!(snapshot.quota_period.as_deref(), Some("quota"));
5642 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
5643 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("10"));
5644 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("10"));
5645 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("10"));
5646 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("10"));
5647 assert!(snapshot.routing_exhausted());
5648 }
5649
5650 #[test]
5651 fn sub2api_quota_limited_rate_limit_exhaustion_marks_temporary_window() {
5652 let reset_at = "2026-01-02T03:04:05Z";
5653 let reset_at_ms = parse_timestamp_secs(reset_at).expect("timestamp") * 1000;
5654 let snapshot = sub2api_usage_snapshot_from_json(
5655 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5656 &upstream(),
5657 &serde_json::json!({
5658 "isValid": true,
5659 "mode": "quota_limited",
5660 "rate_limits": [
5661 {
5662 "window": "5h",
5663 "limit": 100,
5664 "used": 100,
5665 "remaining": 0,
5666 "reset_at": reset_at
5667 }
5668 ]
5669 }),
5670 100,
5671 Some(1_000),
5672 );
5673
5674 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5675 assert_eq!(snapshot.exhausted, Some(true));
5676 assert_eq!(snapshot.quota_period.as_deref(), Some("rate_limit:5h"));
5677 assert_eq!(snapshot.quota_resets_at_ms, Some(reset_at_ms));
5678 assert_eq!(snapshot.quota_remaining_usd, None);
5679 assert!(snapshot.routing_exhausted());
5680 }
5681
5682 #[test]
5683 fn sub2api_quota_limited_total_quota_exhaustion_wins_over_rate_limit_window() {
5684 let snapshot = sub2api_usage_snapshot_from_json(
5685 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5686 &upstream(),
5687 &serde_json::json!({
5688 "isValid": true,
5689 "mode": "quota_limited",
5690 "quota": {
5691 "limit": 10,
5692 "used": 10,
5693 "remaining": 0,
5694 "unit": "USD"
5695 },
5696 "rate_limits": [
5697 {
5698 "window": "5h",
5699 "limit": 100,
5700 "used": 100,
5701 "remaining": 0,
5702 "reset_at": "2026-01-02T03:04:05Z"
5703 }
5704 ]
5705 }),
5706 100,
5707 Some(1_000),
5708 );
5709
5710 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5711 assert_eq!(snapshot.quota_period.as_deref(), Some("quota"));
5712 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
5713 assert_eq!(snapshot.quota_resets_at_ms, None);
5714 }
5715
5716 #[test]
5717 fn sub2api_usage_snapshot_marks_invalid_key_as_error() {
5718 let snapshot = sub2api_usage_snapshot_from_json(
5719 &provider("sub2api", ProviderKind::Sub2ApiUsage),
5720 &upstream(),
5721 &serde_json::json!({
5722 "isValid": false
5723 }),
5724 100,
5725 Some(1_000),
5726 );
5727
5728 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
5729 assert_eq!(
5730 snapshot.error.as_deref(),
5731 Some("sub2api usage response reported invalid API key")
5732 );
5733 }
5734
5735 #[test]
5736 fn sub2api_auth_me_snapshot_reads_dashboard_balance_envelope() {
5737 let snapshot = sub2api_auth_me_snapshot_from_json(
5738 &provider("sub2api-auth", ProviderKind::Sub2ApiAuthMe),
5739 &upstream(),
5740 &serde_json::json!({
5741 "code": 0,
5742 "message": "ok",
5743 "data": {
5744 "id": 42,
5745 "username": "demo",
5746 "balance": "12.5"
5747 }
5748 }),
5749 100,
5750 Some(1_000),
5751 );
5752
5753 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5754 assert_eq!(snapshot.exhausted, Some(false));
5755 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("12.5"));
5756 }
5757
5758 #[test]
5759 fn rightcode_endpoint_defaults_to_account_summary() {
5760 let mut provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
5761 provider.endpoint.clear();
5762
5763 let endpoint = resolve_endpoint(&provider, "https://www.right.codes/codex/v1", "token")
5764 .expect("endpoint");
5765
5766 assert_eq!(endpoint, "https://www.right.codes/account/summary");
5767 }
5768
5769 #[test]
5770 fn rightcode_account_summary_reads_matching_subscription_and_balance() {
5771 let mut provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
5772 provider.trust_exhaustion_for_routing = false;
5773
5774 let snapshot = rightcode_account_summary_snapshot_from_json(
5775 &provider,
5776 &upstream(),
5777 &serde_json::json!({
5778 "balance": 3.25,
5779 "subscriptions": [
5780 {
5781 "name": "Daily",
5782 "total_quota": 20,
5783 "remaining_quota": 7.5,
5784 "reset_today": true,
5785 "available_prefixes": ["/codex"]
5786 },
5787 {
5788 "name": "Other",
5789 "total_quota": 99,
5790 "remaining_quota": 99,
5791 "reset_today": true,
5792 "available_prefixes": ["/claude"]
5793 },
5794 {
5795 "name": "Badge",
5796 "total_quota": 5,
5797 "remaining_quota": 5,
5798 "reset_today": true,
5799 "available_prefixes": ["/codex"]
5800 }
5801 ]
5802 }),
5803 "https://www.right.codes/codex/v1",
5804 100,
5805 Some(1_000),
5806 );
5807
5808 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5809 assert_eq!(snapshot.exhausted, Some(false));
5810 assert!(!snapshot.routing_exhausted());
5811 assert_eq!(snapshot.plan_name.as_deref(), Some("Daily"));
5812 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("3.25"));
5813 assert_eq!(snapshot.paygo_balance_usd, None);
5814 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
5815 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("7.5"));
5816 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("20"));
5817 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("12.5"));
5818 }
5819
5820 #[test]
5821 fn rightcode_account_summary_accounts_for_not_reset_today() {
5822 let provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
5823
5824 let snapshot = rightcode_account_summary_snapshot_from_json(
5825 &provider,
5826 &upstream(),
5827 &serde_json::json!({
5828 "subscriptions": [
5829 {
5830 "name": "Daily",
5831 "total_quota": 20,
5832 "remaining_quota": 7.5,
5833 "reset_today": false,
5834 "available_prefixes": ["/codex"]
5835 }
5836 ]
5837 }),
5838 "https://www.right.codes/codex/v1",
5839 100,
5840 Some(1_000),
5841 );
5842
5843 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
5844 assert_eq!(snapshot.exhausted, Some(false));
5845 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("27.5"));
5846 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("20"));
5847 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("0"));
5848 }
5849
5850 #[test]
5851 fn rightcode_zero_daily_quota_without_balance_is_display_only_exhaustion_by_default() {
5852 let mut provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
5853 provider.trust_exhaustion_for_routing = false;
5854
5855 let snapshot = rightcode_account_summary_snapshot_from_json(
5856 &provider,
5857 &upstream(),
5858 &serde_json::json!({
5859 "balance": 0,
5860 "subscriptions": [
5861 {
5862 "name": "Daily",
5863 "total_quota": 20,
5864 "remaining_quota": 0,
5865 "reset_today": true,
5866 "available_prefixes": ["/codex"]
5867 }
5868 ]
5869 }),
5870 "https://www.right.codes/codex/v1",
5871 100,
5872 Some(1_000),
5873 );
5874
5875 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5876 assert_eq!(snapshot.exhausted, Some(true));
5877 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
5878 assert!(!snapshot.routing_exhausted());
5879 assert!(snapshot.routing_ignored_exhaustion());
5880 }
5881
5882 #[test]
5883 fn sub2api_auth_me_snapshot_marks_business_error() {
5884 let snapshot = sub2api_auth_me_snapshot_from_json(
5885 &provider("sub2api-auth", ProviderKind::Sub2ApiAuthMe),
5886 &upstream(),
5887 &serde_json::json!({
5888 "code": 401,
5889 "message": "login required"
5890 }),
5891 100,
5892 Some(1_000),
5893 );
5894
5895 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
5896 assert_eq!(snapshot.error.as_deref(), Some("login required"));
5897 }
5898
5899 #[test]
5900 fn provider_can_disable_routing_trust_for_exhausted_balance() {
5901 let mut provider = provider("sub2api", ProviderKind::OpenAiBalanceHttpJson);
5902 provider.trust_exhaustion_for_routing = false;
5903
5904 let snapshot = balance_http_snapshot_from_json(
5905 &provider,
5906 &upstream(),
5907 &serde_json::json!({
5908 "balance": "0"
5909 }),
5910 100,
5911 Some(1_000),
5912 );
5913
5914 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
5915 assert_eq!(snapshot.exhausted, Some(true));
5916 assert!(!snapshot.exhaustion_affects_routing);
5917 assert!(!snapshot.routing_exhausted());
5918 }
5919
5920 #[test]
5921 fn provider_exhaustion_trust_defaults_to_enabled_when_omitted() {
5922 let provider: UsageProviderConfig = serde_json::from_value(serde_json::json!({
5923 "id": "sub2api",
5924 "kind": "openai_balance_http_json",
5925 "domains": ["example.com"]
5926 }))
5927 .expect("provider config");
5928
5929 assert!(provider.trust_exhaustion_for_routing);
5930 }
5931
5932 #[tokio::test]
5933 async fn provider_missing_token_clears_stale_lb_exhaustion_marker() {
5934 let cfg = proxy_config(vec![service_config(
5935 "right",
5936 vec![
5937 upstream_config("https://primary.example/v1"),
5938 upstream_config("https://backup.example/v1"),
5939 ],
5940 )]);
5941 let lb_states = Arc::new(Mutex::new(HashMap::new()));
5942 let target = UsageProviderTarget {
5943 upstream: upstream(),
5944 base_url: "https://backup.example/v1".to_string(),
5945 provider_id: Some("right".to_string()),
5946 };
5947 let upstreams = vec![target.upstream.clone()];
5948 let state = ProxyState::new();
5949 update_usage_exhausted(&lb_states, &state, &cfg, "codex", &upstreams, true).await;
5950 {
5951 let guard = lb_states.lock().expect("lb states");
5952 assert!(
5953 guard
5954 .get("right")
5955 .and_then(|entry| entry.usage_exhausted.get(1))
5956 .copied()
5957 .unwrap_or(false)
5958 );
5959 }
5960
5961 let outcome = refresh_provider_target(RefreshProviderTargetParams {
5962 client: &Client::new(),
5963 provider: &provider("sub2api", ProviderKind::OpenAiBalanceHttpJson),
5964 target: &target,
5965 cfg: &cfg,
5966 lb_states: &lb_states,
5967 state: &state,
5968 service_name: "codex",
5969 interval_secs: 60,
5970 })
5971 .await;
5972
5973 assert_eq!(outcome, UsageProviderRefreshOutcome::MissingToken);
5974 let guard = lb_states.lock().expect("lb states");
5975 assert!(
5976 !guard
5977 .get("right")
5978 .and_then(|entry| entry.usage_exhausted.get(1))
5979 .copied()
5980 .unwrap_or(true)
5981 );
5982 }
5983
5984 #[tokio::test]
5985 async fn usage_exhaustion_syncs_owned_balance_policy_action() {
5986 let cfg = proxy_config(vec![service_config(
5987 "right",
5988 vec![
5989 upstream_config("https://primary.example/v1"),
5990 upstream_config("https://backup.example/v1"),
5991 ],
5992 )]);
5993 let lb_states = Arc::new(Mutex::new(HashMap::new()));
5994 let upstreams = vec![endpoint_upstream()];
5995 let state = ProxyState::new();
5996
5997 update_usage_exhausted(&lb_states, &state, &cfg, "codex", &upstreams, true).await;
5998 let actions = state.list_policy_actions("codex").await;
5999 assert_eq!(actions.len(), 1);
6000 assert_eq!(actions[0].source_signal.kind, ProviderSignalKind::Balance);
6001 assert_eq!(actions[0].reason, "balance_exhausted");
6002
6003 update_usage_exhausted(&lb_states, &state, &cfg, "codex", &upstreams, false).await;
6004 assert!(state.list_policy_actions("codex").await.is_empty());
6005 }
6006
6007 #[test]
6008 fn openai_balance_snapshot_supports_custom_paths_and_derived_budget() {
6009 let mut provider = provider("custom", ProviderKind::OpenAiBalanceHttpJson);
6010 provider.extract.remaining_balance_paths = vec!["payload.remaining_usd".to_string()];
6011 provider.extract.monthly_spent_paths = vec!["payload.used_usd".to_string()];
6012 provider.extract.derive_budget_from_remaining_and_spent = true;
6013
6014 let snapshot = balance_http_snapshot_from_json(
6015 &provider,
6016 &upstream(),
6017 &serde_json::json!({
6018 "payload": {
6019 "remaining_usd": "2",
6020 "used_usd": "0.5"
6021 }
6022 }),
6023 100,
6024 Some(1_000),
6025 );
6026
6027 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("2"));
6028 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
6029 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("2.5"));
6030 assert_eq!(snapshot.exhausted, Some(false));
6031 }
6032
6033 #[test]
6034 fn new_api_snapshot_converts_quota_units_like_cc_switch_template() {
6035 let snapshot = new_api_snapshot_from_json(
6036 &provider("newapi", ProviderKind::NewApiUserSelf),
6037 &upstream(),
6038 &serde_json::json!({
6039 "success": true,
6040 "data": {
6041 "quota": 500000,
6042 "used_quota": 250000
6043 }
6044 }),
6045 100,
6046 Some(1_000),
6047 );
6048
6049 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
6050 assert_eq!(snapshot.exhausted, Some(false));
6051 assert_eq!(snapshot.total_balance_usd, None);
6052 assert_eq!(snapshot.quota_period.as_deref(), Some("quota"));
6053 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("1"));
6054 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("1.5"));
6055 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("0.5"));
6056 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
6057 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("1.5"));
6058 }
6059
6060 #[test]
6061 fn new_api_user_self_honors_unlimited_quota_flag() {
6062 let snapshot = new_api_snapshot_from_json(
6063 &provider("newapi", ProviderKind::NewApiUserSelf),
6064 &upstream(),
6065 &serde_json::json!({
6066 "success": true,
6067 "data": {
6068 "quota": 0,
6069 "used_quota": 250000,
6070 "unlimited_quota": true
6071 }
6072 }),
6073 100,
6074 Some(1_000),
6075 );
6076
6077 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
6078 assert_eq!(snapshot.exhausted, Some(false));
6079 assert_eq!(snapshot.total_balance_usd, None);
6080 assert_eq!(snapshot.monthly_budget_usd, None);
6081 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
6082 assert_eq!(snapshot.unlimited_quota, Some(true));
6083 }
6084
6085 #[test]
6086 fn new_api_token_usage_honors_unlimited_quota_flag() {
6087 let snapshot = new_api_token_usage_snapshot_from_json(
6088 &provider("newapi-token", ProviderKind::NewApiTokenUsage),
6089 &upstream(),
6090 &serde_json::json!({
6091 "code": true,
6092 "message": "ok",
6093 "data": {
6094 "object": "token_usage",
6095 "name": "demo-token",
6096 "total_granted": 0,
6097 "total_used": 250000,
6098 "total_available": 0,
6099 "unlimited_quota": true
6100 }
6101 }),
6102 100,
6103 Some(1_000),
6104 );
6105
6106 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
6107 assert_eq!(snapshot.exhausted, Some(false));
6108 assert_eq!(snapshot.plan_name.as_deref(), Some("demo-token"));
6109 assert_eq!(snapshot.total_balance_usd, None);
6110 assert_eq!(snapshot.monthly_budget_usd, None);
6111 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
6112 assert_eq!(snapshot.unlimited_quota, Some(true));
6113 }
6114
6115 #[test]
6116 fn openai_organization_costs_sums_official_cost_buckets_without_exhaustion() {
6117 let snapshot = openai_organization_costs_snapshot_from_json(
6118 &provider("openai", ProviderKind::OpenAiOrganizationCosts),
6119 &upstream(),
6120 &serde_json::json!({
6121 "object": "page",
6122 "data": [
6123 {
6124 "object": "bucket",
6125 "start_time": 1710000000,
6126 "end_time": 1710086400,
6127 "results": [
6128 {
6129 "object": "organization.costs.result",
6130 "amount": { "value": 1.25, "currency": "usd" }
6131 },
6132 {
6133 "object": "organization.costs.result",
6134 "amount": { "value": "2.5", "currency": "usd" }
6135 },
6136 {
6137 "object": "organization.costs.result",
6138 "amount": { "value": 99, "currency": "eur" }
6139 }
6140 ]
6141 },
6142 {
6143 "object": "bucket",
6144 "results": [
6145 {
6146 "object": "organization.costs.result",
6147 "amount": { "value": "0.25", "currency": "USD" }
6148 }
6149 ]
6150 }
6151 ],
6152 "has_more": false
6153 }),
6154 100,
6155 Some(1_000),
6156 );
6157
6158 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
6159 assert_eq!(snapshot.exhausted, None);
6160 assert!(!snapshot.exhaustion_affects_routing);
6161 assert!(!snapshot.routing_exhausted());
6162 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("4"));
6163 assert_eq!(snapshot.total_balance_usd, None);
6164 }
6165}