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