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::{BalanceSnapshotStatus, ProviderBalanceSnapshot};
12use crate::config::{ProxyConfig, ServiceConfigManager, proxy_home_dir};
13use crate::lb::LbState;
14use crate::pricing::UsdAmount;
15use crate::runtime_identity::ProviderEndpointKey;
16use crate::state::ProxyState;
17
18#[derive(Debug, Clone, Copy, Deserialize, Serialize, PartialEq, Eq)]
19#[serde(rename_all = "snake_case")]
20enum ProviderKind {
21 BudgetHttpJson,
23 YescodeProfile,
25 #[serde(
27 rename = "openai_balance_http_json",
28 alias = "open_ai_balance_http_json",
29 alias = "relay_balance_http_json"
30 )]
31 OpenAiBalanceHttpJson,
32 #[serde(rename = "sub2api_usage", alias = "sub2api_usage_http_json")]
34 Sub2ApiUsage,
35 #[serde(rename = "sub2api_auth_me", alias = "sub2api_auth_me_http_json")]
37 Sub2ApiAuthMe,
38 #[serde(
40 rename = "new_api_token_usage",
41 alias = "new_api_token_usage_http_json"
42 )]
43 NewApiTokenUsage,
44 NewApiUserSelf,
46 #[serde(
48 rename = "rightcode_account_summary",
49 alias = "right_code_account_summary",
50 alias = "rightcode"
51 )]
52 RightCodeAccountSummary,
53 #[serde(
55 rename = "openai_organization_costs",
56 alias = "openai_org_costs",
57 alias = "openai_costs"
58 )]
59 OpenAiOrganizationCosts,
60}
61
62impl ProviderKind {
63 fn source_name(&self) -> &'static str {
64 match self {
65 ProviderKind::BudgetHttpJson => "usage_provider:budget_http_json",
66 ProviderKind::YescodeProfile => "usage_provider:yescode_profile",
67 ProviderKind::OpenAiBalanceHttpJson => "usage_provider:openai_balance_http_json",
68 ProviderKind::Sub2ApiUsage => "usage_provider:sub2api_usage",
69 ProviderKind::Sub2ApiAuthMe => "usage_provider:sub2api_auth_me",
70 ProviderKind::NewApiTokenUsage => "usage_provider:new_api_token_usage",
71 ProviderKind::NewApiUserSelf => "usage_provider:new_api_user_self",
72 ProviderKind::RightCodeAccountSummary => "usage_provider:rightcode_account_summary",
73 ProviderKind::OpenAiOrganizationCosts => "usage_provider:openai_organization_costs",
74 }
75 }
76
77 fn default_endpoint(&self) -> Option<&'static str> {
78 match self {
79 ProviderKind::OpenAiBalanceHttpJson => Some("{{base_url}}/user/balance"),
80 ProviderKind::Sub2ApiUsage => Some("{{base_url}}/v1/usage"),
81 ProviderKind::Sub2ApiAuthMe => Some("{{base_url}}/api/v1/auth/me"),
82 ProviderKind::NewApiTokenUsage => Some("{{base_url}}/api/usage/token/"),
83 ProviderKind::NewApiUserSelf => Some("{{base_url}}/api/user/self"),
84 ProviderKind::RightCodeAccountSummary => {
85 Some("https://www.right.codes/account/summary")
86 }
87 ProviderKind::OpenAiOrganizationCosts => {
88 Some("{{base_url}}/v1/organization/costs?start_time={{unix_days_ago:30}}&limit=30")
89 }
90 _ => None,
91 }
92 }
93}
94
95#[derive(Debug, Deserialize, Serialize, Default, Clone)]
96#[serde(default)]
97struct UsageProviderExtractConfig {
98 #[serde(skip_serializing_if = "Vec::is_empty")]
99 remaining_balance_paths: Vec<String>,
100 #[serde(skip_serializing_if = "Vec::is_empty")]
101 subscription_balance_paths: Vec<String>,
102 #[serde(skip_serializing_if = "Vec::is_empty")]
103 paygo_balance_paths: Vec<String>,
104 #[serde(skip_serializing_if = "Vec::is_empty")]
105 monthly_budget_paths: Vec<String>,
106 #[serde(skip_serializing_if = "Vec::is_empty")]
107 monthly_spent_paths: Vec<String>,
108 #[serde(skip_serializing_if = "Vec::is_empty")]
109 exhausted_paths: Vec<String>,
110 #[serde(skip_serializing_if = "Option::is_none")]
111 remaining_divisor: Option<u64>,
112 #[serde(skip_serializing_if = "Option::is_none")]
113 monthly_budget_divisor: Option<u64>,
114 #[serde(skip_serializing_if = "Option::is_none")]
115 monthly_spent_divisor: Option<u64>,
116 #[serde(skip_serializing_if = "bool_is_false")]
117 derive_budget_from_remaining_and_spent: bool,
118 #[serde(skip_serializing_if = "bool_is_false")]
119 derive_remaining_from_budget_and_spent: bool,
120}
121
122impl UsageProviderExtractConfig {
123 fn is_empty(&self) -> bool {
124 self.remaining_balance_paths.is_empty()
125 && self.subscription_balance_paths.is_empty()
126 && self.paygo_balance_paths.is_empty()
127 && self.monthly_budget_paths.is_empty()
128 && self.monthly_spent_paths.is_empty()
129 && self.exhausted_paths.is_empty()
130 && self.remaining_divisor.is_none()
131 && self.monthly_budget_divisor.is_none()
132 && self.monthly_spent_divisor.is_none()
133 && !self.derive_budget_from_remaining_and_spent
134 && !self.derive_remaining_from_budget_and_spent
135 }
136}
137
138#[derive(Debug, Deserialize, Serialize)]
139struct UsageProviderConfig {
140 id: String,
141 kind: ProviderKind,
142 domains: Vec<String>,
143 #[serde(default)]
144 endpoint: String,
145 #[serde(default)]
146 token_env: Option<String>,
147 #[serde(default, skip_serializing_if = "bool_is_false")]
148 require_token_env: bool,
149 #[serde(default)]
150 poll_interval_secs: Option<u64>,
151 #[serde(
152 default = "default_refresh_on_request",
153 skip_serializing_if = "bool_is_true"
154 )]
155 refresh_on_request: bool,
156 #[serde(
157 default = "default_trust_exhaustion_for_routing",
158 skip_serializing_if = "bool_is_true"
159 )]
160 trust_exhaustion_for_routing: bool,
161 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
162 headers: BTreeMap<String, String>,
163 #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
164 variables: BTreeMap<String, String>,
165 #[serde(default, skip_serializing_if = "UsageProviderExtractConfig::is_empty")]
166 extract: UsageProviderExtractConfig,
167}
168
169#[derive(Debug, Deserialize, Serialize, Default)]
170struct UsageProvidersFile {
171 #[serde(default)]
172 providers: Vec<UsageProviderConfig>,
173}
174
175#[derive(Debug, Clone)]
176struct UpstreamRef {
177 station_name: String,
178 index: usize,
179 provider_endpoint: Option<ProviderEndpointKey>,
180}
181
182#[derive(Debug, Clone)]
183struct UsageProviderTarget {
184 upstream: UpstreamRef,
185 base_url: String,
186 provider_id: Option<String>,
187}
188
189#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord)]
190struct UsageProviderTargetKey {
191 station_name: String,
192 upstream_index: usize,
193}
194
195#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
196pub struct UsageProviderRefreshSummary {
197 pub providers_configured: usize,
198 pub providers_matched: usize,
199 pub upstreams_matched: usize,
200 pub attempted: usize,
201 pub refreshed: usize,
202 pub failed: usize,
203 pub missing_token: usize,
204 #[serde(skip_serializing_if = "usize_is_zero")]
205 pub auto_attempted: usize,
206 #[serde(skip_serializing_if = "usize_is_zero")]
207 pub auto_refreshed: usize,
208 #[serde(skip_serializing_if = "usize_is_zero")]
209 pub auto_failed: usize,
210 #[serde(skip_serializing_if = "usize_is_zero")]
211 pub deduplicated: usize,
212}
213
214#[derive(Debug, Clone, Copy, PartialEq, Eq)]
215enum UsageProviderRefreshOutcome {
216 Refreshed,
217 Failed,
218 MissingToken,
219}
220
221struct RefreshProviderTargetParams<'a> {
222 client: &'a Client,
223 provider: &'a UsageProviderConfig,
224 target: &'a UsageProviderTarget,
225 cfg: &'a ProxyConfig,
226 lb_states: &'a Arc<Mutex<HashMap<String, LbState>>>,
227 state: &'a Arc<ProxyState>,
228 service_name: &'a str,
229 interval_secs: u64,
230}
231
232static LAST_USAGE_POLL: OnceLock<Mutex<HashMap<String, Instant>>> = OnceLock::new();
234
235const DEFAULT_POLL_INTERVAL_SECS: u64 = 60;
236const MIN_POLL_INTERVAL_SECS: u64 = 20;
238const BALANCE_REFRESH_CONCURRENCY: usize = 6;
239const BALANCE_HTTP_REQUEST_TIMEOUT: Duration = Duration::from_secs(6);
240const AUTO_PROVIDER_ID_PREFIX: &str = "auto:balance:";
241const AUTO_PROBE_KINDS: [ProviderKind; 5] = [
242 ProviderKind::RightCodeAccountSummary,
243 ProviderKind::Sub2ApiUsage,
244 ProviderKind::NewApiTokenUsage,
245 ProviderKind::NewApiUserSelf,
246 ProviderKind::OpenAiBalanceHttpJson,
247];
248
249fn bool_is_false(value: &bool) -> bool {
250 !*value
251}
252
253fn bool_is_true(value: &bool) -> bool {
254 *value
255}
256
257fn usize_is_zero(value: &usize) -> bool {
258 *value == 0
259}
260
261fn default_refresh_on_request() -> bool {
262 true
263}
264
265fn default_trust_exhaustion_for_routing() -> bool {
266 true
267}
268
269fn unix_now_ms() -> u64 {
270 std::time::SystemTime::now()
271 .duration_since(std::time::UNIX_EPOCH)
272 .map(|d| d.as_millis() as u64)
273 .unwrap_or(0)
274}
275
276fn unix_now_secs() -> u64 {
277 std::time::SystemTime::now()
278 .duration_since(std::time::UNIX_EPOCH)
279 .map(|d| d.as_secs())
280 .unwrap_or(0)
281}
282
283fn stale_after_ms(fetched_at_ms: u64, interval_secs: u64) -> Option<u64> {
284 fetched_at_ms.checked_add(interval_secs.saturating_mul(3).saturating_mul(1_000))
285}
286
287fn snapshot_refresh_interval_secs(provider: &UsageProviderConfig) -> u64 {
288 let interval_secs = provider
289 .poll_interval_secs
290 .unwrap_or(DEFAULT_POLL_INTERVAL_SECS);
291 if interval_secs == 0 {
292 DEFAULT_POLL_INTERVAL_SECS
293 } else {
294 interval_secs.max(MIN_POLL_INTERVAL_SECS)
295 }
296}
297
298fn effective_poll_interval_secs(provider: &UsageProviderConfig) -> Option<u64> {
299 if !provider.refresh_on_request {
300 return None;
301 }
302
303 let interval_secs = provider
304 .poll_interval_secs
305 .unwrap_or(DEFAULT_POLL_INTERVAL_SECS);
306 if interval_secs == 0 {
307 return None;
308 }
309 Some(interval_secs.max(MIN_POLL_INTERVAL_SECS))
310}
311
312fn usage_providers_path() -> std::path::PathBuf {
313 proxy_home_dir().join("usage_providers.json")
314}
315
316fn service_manager<'a>(cfg: &'a ProxyConfig, service_name: &str) -> &'a ServiceConfigManager {
317 match service_name {
318 "claude" => &cfg.claude,
319 _ => &cfg.codex,
320 }
321}
322
323fn default_provider_config(
324 id: &str,
325 kind: ProviderKind,
326 domains: Vec<&str>,
327 endpoint: &str,
328 extract: UsageProviderExtractConfig,
329) -> UsageProviderConfig {
330 UsageProviderConfig {
331 id: id.to_string(),
332 kind,
333 domains: domains.into_iter().map(str::to_string).collect(),
334 endpoint: endpoint.to_string(),
335 token_env: None,
336 require_token_env: false,
337 poll_interval_secs: Some(60),
338 refresh_on_request: true,
339 trust_exhaustion_for_routing: true,
340 headers: BTreeMap::new(),
341 variables: BTreeMap::new(),
342 extract,
343 }
344}
345
346fn default_rightcode_provider_config(id: &str) -> UsageProviderConfig {
347 let mut provider = default_provider_config(
348 id,
349 ProviderKind::RightCodeAccountSummary,
350 vec!["www.right.codes", "right.codes"],
351 "https://www.right.codes/account/summary",
352 UsageProviderExtractConfig::default(),
353 );
354 provider.trust_exhaustion_for_routing = false;
358 provider
359}
360
361fn host_from_base_url(base_url: &str) -> Option<String> {
362 reqwest::Url::parse(base_url)
363 .ok()
364 .and_then(|url| url.host_str().map(|host| host.to_ascii_lowercase()))
365}
366
367fn is_official_openai_base_url(base_url: &str) -> bool {
368 host_from_base_url(base_url).as_deref() == Some("api.openai.com")
369}
370
371fn is_rightcode_base_url(base_url: &str) -> bool {
372 matches!(
373 host_from_base_url(base_url).as_deref(),
374 Some("www.right.codes" | "right.codes")
375 )
376}
377
378fn provider_id_component(value: &str) -> String {
379 let component = value
380 .chars()
381 .map(|ch| {
382 if ch.is_ascii_alphanumeric() || ch == '-' || ch == '_' || ch == '.' {
383 ch
384 } else {
385 '-'
386 }
387 })
388 .collect::<String>()
389 .trim_matches('-')
390 .to_string();
391 if component.is_empty() {
392 "station".to_string()
393 } else {
394 component
395 }
396}
397
398fn auto_provider_id(target: &UsageProviderTarget) -> String {
399 if let Some(provider_id) = target
400 .provider_id
401 .as_deref()
402 .map(str::trim)
403 .filter(|value| !value.is_empty())
404 {
405 return provider_id.to_string();
406 }
407 format!(
408 "{}{}:{}",
409 AUTO_PROVIDER_ID_PREFIX,
410 provider_id_component(&target.upstream.station_name),
411 target.upstream.index
412 )
413}
414
415fn auto_usage_provider(target: &UsageProviderTarget, kind: ProviderKind) -> UsageProviderConfig {
416 let mut provider = UsageProviderConfig {
417 id: auto_provider_id(target),
418 kind,
419 domains: host_from_base_url(&target.base_url)
420 .into_iter()
421 .collect::<Vec<_>>(),
422 endpoint: String::new(),
423 token_env: None,
424 require_token_env: false,
425 poll_interval_secs: Some(DEFAULT_POLL_INTERVAL_SECS),
426 refresh_on_request: true,
427 trust_exhaustion_for_routing: true,
428 headers: BTreeMap::new(),
429 variables: BTreeMap::new(),
430 extract: UsageProviderExtractConfig::default(),
431 };
432 if matches!(kind, ProviderKind::RightCodeAccountSummary) {
433 provider.trust_exhaustion_for_routing = false;
434 }
435 provider
436}
437
438fn first_auto_probe_kind(target: &UsageProviderTarget) -> ProviderKind {
439 if is_rightcode_base_url(&target.base_url) {
440 ProviderKind::RightCodeAccountSummary
441 } else {
442 ProviderKind::Sub2ApiUsage
443 }
444}
445
446fn auto_openai_official_provider(target: &UsageProviderTarget) -> UsageProviderConfig {
447 let mut provider = auto_usage_provider(target, ProviderKind::OpenAiOrganizationCosts);
448 provider.token_env = Some("OPENAI_ADMIN_KEY".to_string());
449 provider.require_token_env = true;
450 provider.refresh_on_request = false;
451 provider.trust_exhaustion_for_routing = false;
452 provider
453}
454
455fn default_providers() -> UsageProvidersFile {
456 let openrouter_extract = UsageProviderExtractConfig {
457 monthly_budget_paths: vec!["data.total_credits".to_string()],
458 monthly_spent_paths: vec!["data.total_usage".to_string()],
459 derive_remaining_from_budget_and_spent: true,
460 ..Default::default()
461 };
462
463 let novita_extract = UsageProviderExtractConfig {
464 remaining_balance_paths: vec!["availableBalance".to_string()],
465 remaining_divisor: Some(10_000),
466 ..Default::default()
467 };
468
469 let mut openai_official = default_provider_config(
470 "openai-official-costs",
471 ProviderKind::OpenAiOrganizationCosts,
472 vec!["api.openai.com"],
473 "https://api.openai.com/v1/organization/costs?start_time={{unix_days_ago:30}}&limit=30",
474 UsageProviderExtractConfig::default(),
475 );
476 openai_official.token_env = Some("OPENAI_ADMIN_KEY".to_string());
477 openai_official.require_token_env = true;
478 openai_official.refresh_on_request = false;
479 openai_official.trust_exhaustion_for_routing = false;
480
481 UsageProvidersFile {
482 providers: vec![
483 default_rightcode_provider_config("rightcode"),
484 default_provider_config(
485 "packycode",
486 ProviderKind::BudgetHttpJson,
487 vec!["packycode.com"],
488 "https://www.packycode.com/api/backend/users/info",
489 UsageProviderExtractConfig::default(),
490 ),
491 default_provider_config(
492 "yescode",
493 ProviderKind::YescodeProfile,
494 vec!["yes.vg"],
496 "https://co.yes.vg/api/v1/auth/profile",
497 UsageProviderExtractConfig::default(),
498 ),
499 default_provider_config(
500 "deepseek",
501 ProviderKind::OpenAiBalanceHttpJson,
502 vec!["api.deepseek.com"],
503 "https://api.deepseek.com/user/balance",
504 UsageProviderExtractConfig::default(),
505 ),
506 default_provider_config(
507 "stepfun",
508 ProviderKind::OpenAiBalanceHttpJson,
509 vec!["api.stepfun.ai", "api.stepfun.com"],
510 "https://api.stepfun.com/v1/accounts",
511 UsageProviderExtractConfig::default(),
512 ),
513 default_provider_config(
514 "siliconflow",
515 ProviderKind::OpenAiBalanceHttpJson,
516 vec!["api.siliconflow.cn", "api.siliconflow.com"],
517 "{{base_url}}/v1/user/info",
518 UsageProviderExtractConfig::default(),
519 ),
520 default_provider_config(
521 "openrouter",
522 ProviderKind::OpenAiBalanceHttpJson,
523 vec!["openrouter.ai"],
524 "https://openrouter.ai/api/v1/credits",
525 openrouter_extract,
526 ),
527 default_provider_config(
528 "novita",
529 ProviderKind::OpenAiBalanceHttpJson,
530 vec!["api.novita.ai"],
531 "https://api.novita.ai/v3/user/balance",
532 novita_extract,
533 ),
534 openai_official,
535 ],
536 }
537}
538
539fn load_providers() -> UsageProvidersFile {
540 let path = usage_providers_path();
541 if let Ok(text) = std::fs::read_to_string(&path)
542 && let Ok(file) = serde_json::from_str::<UsageProvidersFile>(&text)
543 {
544 return file;
545 }
546
547 let default = default_providers();
549 if let Ok(text) = serde_json::to_string_pretty(&default) {
550 if let Some(parent) = path.parent() {
551 let _ = std::fs::create_dir_all(parent);
552 }
553 let _ = std::fs::write(&path, text);
554 }
555 default
556}
557
558fn domain_matches(base_url: &str, domains: &[String]) -> bool {
559 let url = match reqwest::Url::parse(base_url) {
560 Ok(u) => u,
561 Err(_) => return false,
562 };
563 let host = match url.host_str() {
564 Some(h) => h,
565 None => return false,
566 };
567 let host = host.to_ascii_lowercase();
568 for d in domains {
569 let domain = d.trim().to_ascii_lowercase();
570 if host == domain || host.ends_with(&format!(".{}", domain)) {
571 return true;
572 }
573 }
574 false
575}
576
577fn matching_provider_targets(
578 cfg: &ProxyConfig,
579 service_name: &str,
580 provider: &UsageProviderConfig,
581 station_name_filter: Option<&str>,
582) -> Vec<UsageProviderTarget> {
583 let mut stations: Vec<_> = service_manager(cfg, service_name)
584 .stations()
585 .iter()
586 .collect();
587 stations.sort_by_key(|(name, _)| name.as_str());
588
589 let mut targets = Vec::new();
590 for (station_name, service) in stations {
591 if station_name_filter.is_some_and(|filter| filter != station_name.as_str()) {
592 continue;
593 }
594 for (index, upstream) in service.upstreams.iter().enumerate() {
595 if domain_matches(&upstream.base_url, &provider.domains) {
596 targets.push(UsageProviderTarget {
597 upstream: UpstreamRef {
598 station_name: station_name.clone(),
599 index,
600 provider_endpoint: upstream.provider_endpoint_key(service_name),
601 },
602 base_url: upstream.base_url.clone(),
603 provider_id: upstream.tags.get("provider_id").cloned(),
604 });
605 }
606 }
607 }
608
609 targets
610}
611
612fn usage_provider_targets(
613 cfg: &ProxyConfig,
614 service_name: &str,
615 station_name_filter: Option<&str>,
616) -> Vec<UsageProviderTarget> {
617 let mut stations: Vec<_> = service_manager(cfg, service_name)
618 .stations()
619 .iter()
620 .collect();
621 stations.sort_by_key(|(name, _)| name.as_str());
622
623 let mut targets = Vec::new();
624 for (station_name, service) in stations {
625 if station_name_filter.is_some_and(|filter| filter != station_name.as_str()) {
626 continue;
627 }
628 for (index, upstream) in service.upstreams.iter().enumerate() {
629 targets.push(UsageProviderTarget {
630 upstream: UpstreamRef {
631 station_name: station_name.clone(),
632 index,
633 provider_endpoint: upstream.provider_endpoint_key(service_name),
634 },
635 base_url: upstream.base_url.clone(),
636 provider_id: upstream.tags.get("provider_id").cloned(),
637 });
638 }
639 }
640
641 targets
642}
643
644fn target_key(target: &UsageProviderTarget) -> UsageProviderTargetKey {
645 UsageProviderTargetKey {
646 station_name: target.upstream.station_name.clone(),
647 upstream_index: target.upstream.index,
648 }
649}
650
651fn usage_provider_target_for_legacy_upstream(
652 cfg: &ProxyConfig,
653 service_name: &str,
654 station_name: &str,
655 upstream_index: usize,
656) -> Option<UsageProviderTarget> {
657 let current_service = service_manager(cfg, service_name).station(station_name)?;
658 let current_upstream = current_service.upstreams.get(upstream_index)?;
659 Some(UsageProviderTarget {
660 upstream: UpstreamRef {
661 station_name: station_name.to_string(),
662 index: upstream_index,
663 provider_endpoint: current_upstream.provider_endpoint_key(service_name),
664 },
665 base_url: current_upstream.base_url.clone(),
666 provider_id: current_upstream.tags.get("provider_id").cloned(),
667 })
668}
669
670fn usage_provider_target_for_provider_endpoint(
671 cfg: &ProxyConfig,
672 service_name: &str,
673 provider_endpoint: &ProviderEndpointKey,
674) -> Option<UsageProviderTarget> {
675 service_manager(cfg, service_name)
676 .stations()
677 .iter()
678 .filter_map(|(station_name, service)| {
679 service
680 .upstreams
681 .iter()
682 .enumerate()
683 .find_map(|(index, upstream)| {
684 let upstream_endpoint = upstream.provider_endpoint_key(service_name)?;
685 if upstream_endpoint != *provider_endpoint {
686 return None;
687 }
688 Some(UsageProviderTarget {
689 upstream: UpstreamRef {
690 station_name: station_name.clone(),
691 index,
692 provider_endpoint: Some(upstream_endpoint),
693 },
694 base_url: upstream.base_url.clone(),
695 provider_id: upstream.tags.get("provider_id").cloned(),
696 })
697 })
698 })
699 .next()
700}
701
702trait UsageProviderUpstreamIdentityExt {
703 fn provider_endpoint_key(&self, service_name: &str) -> Option<ProviderEndpointKey>;
704}
705
706impl UsageProviderUpstreamIdentityExt for crate::config::UpstreamConfig {
707 fn provider_endpoint_key(&self, service_name: &str) -> Option<ProviderEndpointKey> {
708 let provider_id = self.tags.get("provider_id")?.trim();
709 let endpoint_id = self.tags.get("endpoint_id")?.trim();
710 if provider_id.is_empty() || endpoint_id.is_empty() {
711 return None;
712 }
713 Some(ProviderEndpointKey::new(
714 service_name.to_string(),
715 provider_id.to_string(),
716 endpoint_id.to_string(),
717 ))
718 }
719}
720
721fn configured_target_keys(
722 cfg: &ProxyConfig,
723 service_name: &str,
724 providers: &[UsageProviderConfig],
725 station_name_filter: Option<&str>,
726) -> HashSet<UsageProviderTargetKey> {
727 providers
728 .iter()
729 .flat_map(|provider| {
730 matching_provider_targets(cfg, service_name, provider, station_name_filter)
731 })
732 .map(|target| target_key(&target))
733 .collect()
734}
735
736fn resolve_token(
737 provider: &UsageProviderConfig,
738 upstreams: &[UpstreamRef],
739 cfg: &ProxyConfig,
740 service_name: &str,
741) -> Option<String> {
742 if let Some(env_name) = &provider.token_env
744 && let Ok(v) = std::env::var(env_name)
745 && !v.trim().is_empty()
746 {
747 return Some(v);
748 }
749
750 if provider.require_token_env {
751 return None;
752 }
753
754 for uref in upstreams {
756 if let Some(service) = service_manager(cfg, service_name).station(&uref.station_name)
757 && let Some(up) = service.upstreams.get(uref.index)
758 {
759 if let Some(token) = up.auth.resolve_auth_token() {
760 return Some(token);
761 }
762 if let Some(token) = up.auth.resolve_api_key() {
763 return Some(token);
764 }
765 }
766 }
767 None
768}
769
770fn normalized_balance_base_url(base_url: &str) -> Option<String> {
771 let mut url = reqwest::Url::parse(base_url).ok()?;
772 url.set_query(None);
773 url.set_fragment(None);
774 let path = url.path().trim_end_matches('/').to_string();
775 if path.eq_ignore_ascii_case("/v1") {
776 url.set_path("");
777 } else if path.to_ascii_lowercase().ends_with("/v1") {
778 let new_path = &path[..path.len().saturating_sub(3)];
779 url.set_path(if new_path.is_empty() { "/" } else { new_path });
780 }
781 Some(url.as_str().trim_end_matches('/').to_string())
782}
783
784fn base_path_prefixes(base_url: &str) -> Vec<String> {
785 let Some(normalized) = normalized_balance_base_url(base_url) else {
786 return Vec::new();
787 };
788 let Ok(url) = reqwest::Url::parse(&normalized) else {
789 return Vec::new();
790 };
791 let parts = url
792 .path_segments()
793 .map(|segments| {
794 segments
795 .filter(|segment| !segment.is_empty())
796 .collect::<Vec<_>>()
797 })
798 .unwrap_or_default();
799 let mut prefixes = Vec::new();
800 for len in (1..=parts.len()).rev() {
801 prefixes.push(format!("/{}", parts[..len].join("/")));
802 }
803 if prefixes.is_empty() {
804 prefixes.push("/".to_string());
805 }
806 prefixes
807}
808
809fn path_prefixes_match(provider_prefixes: &[String], available_prefixes: &[String]) -> bool {
810 if provider_prefixes.is_empty() || available_prefixes.is_empty() {
811 return false;
812 }
813 provider_prefixes.iter().any(|provider_prefix| {
814 available_prefixes.iter().any(|available_prefix| {
815 provider_prefix == available_prefix
816 || provider_prefix
817 .strip_prefix(available_prefix)
818 .is_some_and(|suffix| suffix.starts_with('/'))
819 })
820 })
821}
822
823fn render_provider_template(
824 template: &str,
825 base_url: &str,
826 upstream_base_url: &str,
827 token: &str,
828 variables: &BTreeMap<String, String>,
829) -> String {
830 let mut out = template
831 .replace("{{baseUrl}}", base_url)
832 .replace("{{base_url}}", base_url)
833 .replace("{{upstreamBaseUrl}}", upstream_base_url)
834 .replace("{{upstream_base_url}}", upstream_base_url)
835 .replace("{{apiKey}}", token)
836 .replace("{{accessToken}}", token)
837 .replace("{{token}}", token);
838
839 out = out
840 .replace("{{unix_now}}", &unix_now_secs().to_string())
841 .replace("{{unix_now_ms}}", &unix_now_ms().to_string());
842
843 while let Some(start) = out.find("{{unix_days_ago:") {
844 let Some(end_offset) = out[start..].find("}}") else {
845 break;
846 };
847 let end = start + end_offset + 2;
848 let days_str = out[start + "{{unix_days_ago:".len()..end - 2].trim();
849 let replacement = days_str
850 .parse::<u64>()
851 .ok()
852 .map(|days| unix_now_secs().saturating_sub(days.saturating_mul(24 * 60 * 60)))
853 .map(|secs| secs.to_string())
854 .unwrap_or_default();
855 out.replace_range(start..end, &replacement);
856 }
857
858 while let Some(start) = out.find("{{env:") {
859 let Some(end_offset) = out[start..].find("}}") else {
860 break;
861 };
862 let end = start + end_offset + 2;
863 let env_name = out[start + 6..end - 2].trim();
864 let value = std::env::var(env_name).unwrap_or_default();
865 out.replace_range(start..end, &value);
866 }
867
868 for (name, value_template) in variables {
869 let value = render_provider_template(
870 value_template,
871 base_url,
872 upstream_base_url,
873 token,
874 &BTreeMap::new(),
875 );
876 out = out.replace(&format!("{{{{{name}}}}}"), &value);
877 }
878
879 out
880}
881
882fn resolve_endpoint(
883 provider: &UsageProviderConfig,
884 upstream_base_url: &str,
885 token: &str,
886) -> Result<String> {
887 let base_url = normalized_balance_base_url(upstream_base_url)
888 .ok_or_else(|| anyhow::anyhow!("invalid upstream base_url for balance endpoint"))?;
889 let endpoint = if provider.endpoint.trim().is_empty() {
890 provider
891 .kind
892 .default_endpoint()
893 .unwrap_or_default()
894 .to_string()
895 } else {
896 provider.endpoint.trim().to_string()
897 };
898 if endpoint.is_empty() {
899 anyhow::bail!(
900 "usage provider '{}' has no endpoint and kind {:?} has no default endpoint",
901 provider.id,
902 provider.kind
903 );
904 }
905
906 let rendered = render_provider_template(
907 &endpoint,
908 &base_url,
909 upstream_base_url,
910 token,
911 &provider.variables,
912 );
913 if rendered.starts_with("http://") || rendered.starts_with("https://") {
914 return Ok(rendered);
915 }
916
917 let path = if rendered.starts_with('/') {
918 rendered
919 } else {
920 format!("/{rendered}")
921 };
922 Ok(format!("{base_url}{path}"))
923}
924
925fn endpoint_origin(endpoint: &str) -> String {
926 reqwest::Url::parse(endpoint)
927 .ok()
928 .and_then(|url| {
929 let host = url.host_str()?;
930 let origin = match url.port() {
931 Some(port) => format!("{}://{}:{}", url.scheme(), host, port),
932 None => format!("{}://{}", url.scheme(), host),
933 };
934 Some(origin)
935 })
936 .unwrap_or_else(|| "unknown-origin".to_string())
937}
938
939async fn poll_provider_http_json(
940 client: &Client,
941 provider: &UsageProviderConfig,
942 upstream_base_url: &str,
943 token: &str,
944) -> Result<serde_json::Value> {
945 let endpoint = resolve_endpoint(provider, upstream_base_url, token)?;
946 let origin = endpoint_origin(&endpoint);
947 let base_url = normalized_balance_base_url(upstream_base_url).unwrap_or_default();
948 let mut req = client
949 .get(endpoint)
950 .timeout(BALANCE_HTTP_REQUEST_TIMEOUT)
951 .header("Accept", "application/json")
952 .header(
953 "User-Agent",
954 concat!("codex-helper/", env!("CARGO_PKG_VERSION")),
955 );
956
957 match provider.kind {
958 ProviderKind::YescodeProfile => {
959 req = req.header("X-API-Key", token);
960 }
961 _ => {
962 req = req.header("Authorization", format!("Bearer {}", token));
963 }
964 }
965
966 for (name, template) in &provider.headers {
967 let value = render_provider_template(
968 template,
969 &base_url,
970 upstream_base_url,
971 token,
972 &provider.variables,
973 );
974 if !value.trim().is_empty() {
975 req = req.header(name.as_str(), value);
976 }
977 }
978
979 let resp = req.send().await.with_context(|| {
980 format!(
981 "usage provider request failed for {} via {:?}",
982 origin, provider.kind
983 )
984 })?;
985
986 if !resp.status().is_success() {
987 anyhow::bail!(
988 "usage provider HTTP {} from {} via {:?}",
989 resp.status(),
990 origin,
991 provider.kind
992 );
993 }
994 let content_type = resp
995 .headers()
996 .get(reqwest::header::CONTENT_TYPE)
997 .and_then(|value| value.to_str().ok())
998 .map(str::to_string)
999 .unwrap_or_else(|| "unknown".to_string());
1000 let text = resp.text().await.with_context(|| {
1001 format!(
1002 "usage provider response read failed from {} via {:?}",
1003 origin, provider.kind
1004 )
1005 })?;
1006 serde_json::from_str(&text).with_context(|| {
1007 format!(
1008 "usage provider returned non-JSON response from {} via {:?} (content-type {}, {} bytes)",
1009 origin,
1010 provider.kind,
1011 content_type,
1012 text.len()
1013 )
1014 })
1015}
1016
1017fn amount_from_json(value: &serde_json::Value) -> Option<UsdAmount> {
1018 let raw = match value {
1019 serde_json::Value::Number(number) => number.to_string(),
1020 serde_json::Value::String(text) => text.trim().to_string(),
1021 _ => return None,
1022 };
1023 UsdAmount::from_decimal_str(raw.as_str())
1024}
1025
1026fn amount_from_json_with_divisor(
1027 value: &serde_json::Value,
1028 divisor: Option<u64>,
1029) -> Option<UsdAmount> {
1030 let amount = amount_from_json(value)?;
1031 match divisor {
1032 Some(divisor) => amount.checked_div_u64(divisor),
1033 None => Some(amount),
1034 }
1035}
1036
1037fn json_value_at_path<'a>(
1038 value: &'a serde_json::Value,
1039 path: &str,
1040) -> Option<&'a serde_json::Value> {
1041 let mut current = value;
1042 for segment in path
1043 .split('.')
1044 .map(str::trim)
1045 .filter(|segment| !segment.is_empty())
1046 {
1047 current = match current {
1048 serde_json::Value::Array(items) => {
1049 let index = segment.parse::<usize>().ok()?;
1050 items.get(index)?
1051 }
1052 _ => current.get(segment)?,
1053 };
1054 }
1055 Some(current)
1056}
1057
1058fn first_amount_from_paths(
1059 value: &serde_json::Value,
1060 custom_paths: &[String],
1061 default_paths: &[&str],
1062 divisor: Option<u64>,
1063) -> Option<UsdAmount> {
1064 custom_paths
1065 .iter()
1066 .map(String::as_str)
1067 .chain(default_paths.iter().copied())
1068 .find_map(|path| {
1069 json_value_at_path(value, path)
1070 .and_then(|value| amount_from_json_with_divisor(value, divisor))
1071 })
1072}
1073
1074fn bool_from_json(value: &serde_json::Value) -> Option<bool> {
1075 match value {
1076 serde_json::Value::Bool(value) => Some(*value),
1077 serde_json::Value::Number(number) => number.as_i64().map(|value| value != 0),
1078 serde_json::Value::String(text) => match text.trim().to_ascii_lowercase().as_str() {
1079 "true" | "yes" | "1" | "exhausted" => Some(true),
1080 "false" | "no" | "0" | "ok" => Some(false),
1081 _ => None,
1082 },
1083 _ => None,
1084 }
1085}
1086
1087fn first_bool_from_paths(
1088 value: &serde_json::Value,
1089 custom_paths: &[String],
1090 default_paths: &[&str],
1091) -> Option<bool> {
1092 custom_paths
1093 .iter()
1094 .map(String::as_str)
1095 .chain(default_paths.iter().copied())
1096 .find_map(|path| json_value_at_path(value, path).and_then(bool_from_json))
1097}
1098
1099fn string_from_json(value: &serde_json::Value) -> Option<String> {
1100 match value {
1101 serde_json::Value::String(text) => {
1102 let text = text.trim();
1103 if text.is_empty() {
1104 None
1105 } else {
1106 Some(text.to_string())
1107 }
1108 }
1109 _ => None,
1110 }
1111}
1112
1113fn first_string_from_paths(value: &serde_json::Value, default_paths: &[&str]) -> Option<String> {
1114 default_paths
1115 .iter()
1116 .copied()
1117 .find_map(|path| json_value_at_path(value, path).and_then(string_from_json))
1118}
1119
1120fn u64_from_json(value: &serde_json::Value) -> Option<u64> {
1121 match value {
1122 serde_json::Value::Number(number) => number.as_u64(),
1123 serde_json::Value::String(text) => {
1124 let text = text.trim();
1125 if text.is_empty() {
1126 None
1127 } else {
1128 text.parse::<u64>().ok()
1129 }
1130 }
1131 _ => None,
1132 }
1133}
1134
1135fn first_u64_from_paths(value: &serde_json::Value, default_paths: &[&str]) -> Option<u64> {
1136 default_paths
1137 .iter()
1138 .copied()
1139 .find_map(|path| json_value_at_path(value, path).and_then(u64_from_json))
1140}
1141
1142fn array_from_json_path<'a>(
1143 value: &'a serde_json::Value,
1144 path: &str,
1145) -> Option<&'a Vec<serde_json::Value>> {
1146 json_value_at_path(value, path).and_then(|value| value.as_array())
1147}
1148
1149fn amount_to_string(amount: UsdAmount) -> String {
1150 amount.format_usd()
1151}
1152
1153#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1154struct QuotaWindowSnapshot {
1155 period: &'static str,
1156 remaining: UsdAmount,
1157 used: UsdAmount,
1158 limit: UsdAmount,
1159}
1160
1161fn base_snapshot(
1162 provider: &UsageProviderConfig,
1163 upstream: &UpstreamRef,
1164 fetched_at_ms: u64,
1165 stale_after_ms: Option<u64>,
1166) -> ProviderBalanceSnapshot {
1167 let mut snapshot = ProviderBalanceSnapshot::new(
1168 provider.id.clone(),
1169 upstream.station_name.clone(),
1170 upstream.index,
1171 provider.kind.source_name(),
1172 fetched_at_ms,
1173 stale_after_ms,
1174 );
1175 snapshot.exhaustion_affects_routing = provider.trust_exhaustion_for_routing;
1176 snapshot
1177}
1178
1179fn snapshot_error(
1180 provider: &UsageProviderConfig,
1181 upstream: &UpstreamRef,
1182 fetched_at_ms: u64,
1183 stale_after_ms: Option<u64>,
1184 message: impl Into<String>,
1185) -> ProviderBalanceSnapshot {
1186 base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms).with_error(message)
1187}
1188
1189fn budget_snapshot_from_json(
1190 provider: &UsageProviderConfig,
1191 upstream: &UpstreamRef,
1192 value: &serde_json::Value,
1193 fetched_at_ms: u64,
1194 stale_after_ms: Option<u64>,
1195) -> ProviderBalanceSnapshot {
1196 let monthly_budget = first_amount_from_paths(
1197 value,
1198 &provider.extract.monthly_budget_paths,
1199 &["monthly_budget_usd", "data.monthly_budget_usd"],
1200 provider.extract.monthly_budget_divisor,
1201 );
1202 let monthly_spent = first_amount_from_paths(
1203 value,
1204 &provider.extract.monthly_spent_paths,
1205 &["monthly_spent_usd", "data.monthly_spent_usd"],
1206 provider.extract.monthly_spent_divisor,
1207 );
1208 let exhausted = match (monthly_budget, monthly_spent) {
1209 (Some(budget), Some(spent)) if !budget.is_zero() => Some(spent >= budget),
1210 (Some(_), Some(_)) => Some(false),
1211 _ => None,
1212 };
1213
1214 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1215 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
1216 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1217 snapshot.exhausted = exhausted;
1218 snapshot.refresh_status(fetched_at_ms);
1219 snapshot
1220}
1221
1222fn yescode_snapshot_from_json(
1223 provider: &UsageProviderConfig,
1224 upstream: &UpstreamRef,
1225 value: &serde_json::Value,
1226 fetched_at_ms: u64,
1227 stale_after_ms: Option<u64>,
1228) -> ProviderBalanceSnapshot {
1229 let subscription_balance = first_amount_from_paths(
1230 value,
1231 &provider.extract.subscription_balance_paths,
1232 &["subscription_balance", "data.subscription_balance"],
1233 provider.extract.remaining_divisor,
1234 );
1235 let paygo_balance = first_amount_from_paths(
1236 value,
1237 &provider.extract.paygo_balance_paths,
1238 &[
1239 "pay_as_you_go_balance",
1240 "paygo_balance",
1241 "data.pay_as_you_go_balance",
1242 "data.paygo_balance",
1243 ],
1244 provider.extract.remaining_divisor,
1245 );
1246 let total_balance = match (subscription_balance, paygo_balance) {
1247 (Some(subscription), Some(paygo)) => Some(subscription.saturating_add(paygo)),
1248 (Some(subscription), None) => Some(subscription),
1249 (None, Some(paygo)) => Some(paygo),
1250 (None, None) => None,
1251 };
1252
1253 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1254 snapshot.total_balance_usd = total_balance.map(amount_to_string);
1255 snapshot.subscription_balance_usd = subscription_balance.map(amount_to_string);
1256 snapshot.paygo_balance_usd = paygo_balance.map(amount_to_string);
1257 snapshot.exhausted = total_balance.map(UsdAmount::is_zero);
1258 snapshot.refresh_status(fetched_at_ms);
1259 snapshot
1260}
1261
1262fn balance_http_snapshot_from_json(
1263 provider: &UsageProviderConfig,
1264 upstream: &UpstreamRef,
1265 value: &serde_json::Value,
1266 fetched_at_ms: u64,
1267 stale_after_ms: Option<u64>,
1268) -> ProviderBalanceSnapshot {
1269 let remaining_balance = first_amount_from_paths(
1270 value,
1271 &provider.extract.remaining_balance_paths,
1272 &[
1273 "balance",
1274 "remaining",
1275 "remain",
1276 "available",
1277 "available_balance",
1278 "credit",
1279 "credits",
1280 "total_balance",
1281 "total_balance_usd",
1282 "totalBalance",
1283 "availableBalance",
1284 "available_balance_usd",
1285 "balance_infos.0.total_balance",
1286 "data.balance",
1287 "data.remaining",
1288 "data.available",
1289 "data.available_balance",
1290 "data.credit",
1291 "data.credits",
1292 "data.total_balance",
1293 "data.totalBalance",
1294 "data.availableBalance",
1295 ],
1296 provider.extract.remaining_divisor,
1297 );
1298 let subscription_balance = first_amount_from_paths(
1299 value,
1300 &provider.extract.subscription_balance_paths,
1301 &[
1302 "subscription_balance",
1303 "subscription_balance_usd",
1304 "subscriptionBalance",
1305 "data.subscription_balance",
1306 "data.subscription_balance_usd",
1307 "data.subscriptionBalance",
1308 ],
1309 provider.extract.remaining_divisor,
1310 );
1311 let paygo_balance = first_amount_from_paths(
1312 value,
1313 &provider.extract.paygo_balance_paths,
1314 &[
1315 "pay_as_you_go_balance",
1316 "paygo_balance",
1317 "paygo",
1318 "paygoBalance",
1319 "chargeBalance",
1320 "voucherBalance",
1321 "data.pay_as_you_go_balance",
1322 "data.paygo_balance",
1323 "data.paygo",
1324 "data.paygoBalance",
1325 "data.chargeBalance",
1326 "data.voucherBalance",
1327 ],
1328 provider.extract.remaining_divisor,
1329 );
1330 let component_remaining = match (subscription_balance, paygo_balance) {
1331 (Some(subscription), Some(paygo)) => Some(subscription.saturating_add(paygo)),
1332 (Some(subscription), None) => Some(subscription),
1333 (None, Some(paygo)) => Some(paygo),
1334 (None, None) => None,
1335 };
1336 let monthly_spent = first_amount_from_paths(
1337 value,
1338 &provider.extract.monthly_spent_paths,
1339 &[
1340 "monthly_spent_usd",
1341 "spent",
1342 "used",
1343 "used_balance",
1344 "usedBalance",
1345 "total_usage",
1346 "data.monthly_spent_usd",
1347 "data.spent",
1348 "data.used",
1349 "data.used_balance",
1350 "data.usedBalance",
1351 "data.total_usage",
1352 ],
1353 provider.extract.monthly_spent_divisor,
1354 );
1355 let monthly_budget = first_amount_from_paths(
1356 value,
1357 &provider.extract.monthly_budget_paths,
1358 &[
1359 "monthly_budget_usd",
1360 "budget",
1361 "limit",
1362 "quota_total",
1363 "creditLimit",
1364 "total_credits",
1365 "data.monthly_budget_usd",
1366 "data.budget",
1367 "data.limit",
1368 "data.quota_total",
1369 "data.creditLimit",
1370 "data.total_credits",
1371 ],
1372 provider.extract.monthly_budget_divisor,
1373 )
1374 .or_else(|| {
1375 if provider.extract.derive_budget_from_remaining_and_spent {
1376 match (remaining_balance.or(component_remaining), monthly_spent) {
1377 (Some(remaining), Some(spent)) => Some(remaining.saturating_add(spent)),
1378 _ => None,
1379 }
1380 } else {
1381 None
1382 }
1383 });
1384 let total_balance = remaining_balance.or(component_remaining).or_else(|| {
1385 match (
1386 provider.extract.derive_remaining_from_budget_and_spent,
1387 monthly_budget,
1388 monthly_spent,
1389 ) {
1390 (true, Some(budget), Some(spent)) => Some(budget.saturating_sub(spent)),
1391 _ => None,
1392 }
1393 });
1394 let exhausted = first_bool_from_paths(
1395 value,
1396 &provider.extract.exhausted_paths,
1397 &[
1398 "exhausted",
1399 "quota_exhausted",
1400 "balance_exhausted",
1401 "data.exhausted",
1402 "data.quota_exhausted",
1403 "data.balance_exhausted",
1404 ],
1405 )
1406 .or_else(|| total_balance.map(UsdAmount::is_zero));
1407
1408 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1409 snapshot.total_balance_usd = total_balance.map(amount_to_string);
1410 snapshot.subscription_balance_usd = subscription_balance.map(amount_to_string);
1411 snapshot.paygo_balance_usd = paygo_balance.map(amount_to_string);
1412 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
1413 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1414 snapshot.exhausted = exhausted;
1415 snapshot.refresh_status(fetched_at_ms);
1416 snapshot
1417}
1418
1419fn has_any_json_path(value: &serde_json::Value, paths: &[&str]) -> bool {
1420 paths
1421 .iter()
1422 .any(|path| json_value_at_path(value, path).is_some())
1423}
1424
1425fn populate_sub2api_usage_fields(
1426 snapshot: &mut ProviderBalanceSnapshot,
1427 value: &serde_json::Value,
1428) {
1429 snapshot.plan_name = first_string_from_paths(value, &["planName", "data.planName"]);
1430 snapshot.total_used_usd = first_amount_from_paths(
1431 value,
1432 &[],
1433 &["usage.total.cost", "data.usage.total.cost"],
1434 None,
1435 )
1436 .map(amount_to_string);
1437 snapshot.today_used_usd = first_amount_from_paths(
1438 value,
1439 &[],
1440 &["usage.today.cost", "data.usage.today.cost"],
1441 None,
1442 )
1443 .map(amount_to_string);
1444 snapshot.total_requests = first_u64_from_paths(
1445 value,
1446 &["usage.total.requests", "data.usage.total.requests"],
1447 );
1448 snapshot.today_requests = first_u64_from_paths(
1449 value,
1450 &["usage.today.requests", "data.usage.today.requests"],
1451 );
1452 snapshot.total_tokens = first_u64_from_paths(
1453 value,
1454 &[
1455 "usage.total.total_tokens",
1456 "usage.total.tokens",
1457 "data.usage.total.total_tokens",
1458 "data.usage.total.tokens",
1459 ],
1460 );
1461 snapshot.today_tokens = first_u64_from_paths(
1462 value,
1463 &[
1464 "usage.today.total_tokens",
1465 "usage.today.tokens",
1466 "data.usage.today.total_tokens",
1467 "data.usage.today.tokens",
1468 ],
1469 );
1470}
1471
1472fn sub2api_subscription_limit_snapshot(
1473 value: &serde_json::Value,
1474 period: &'static str,
1475 limit_paths: &[&str],
1476 usage_paths: &[&str],
1477) -> Option<QuotaWindowSnapshot> {
1478 let budget = first_amount_from_paths(value, &[], limit_paths, None)?;
1479 if budget.is_zero() {
1480 return None;
1481 }
1482 let spent = first_amount_from_paths(value, &[], usage_paths, None).unwrap_or(UsdAmount::ZERO);
1483 let remaining = budget.saturating_sub(spent);
1484 Some(QuotaWindowSnapshot {
1485 period,
1486 remaining,
1487 used: spent,
1488 limit: budget,
1489 })
1490}
1491
1492fn sub2api_limiting_subscription_window(value: &serde_json::Value) -> Option<QuotaWindowSnapshot> {
1493 let windows = [
1494 sub2api_subscription_limit_snapshot(
1495 value,
1496 "daily",
1497 &[
1498 "subscription.daily_limit_usd",
1499 "data.subscription.daily_limit_usd",
1500 ],
1501 &[
1502 "subscription.daily_usage_usd",
1503 "data.subscription.daily_usage_usd",
1504 ],
1505 ),
1506 sub2api_subscription_limit_snapshot(
1507 value,
1508 "weekly",
1509 &[
1510 "subscription.weekly_limit_usd",
1511 "data.subscription.weekly_limit_usd",
1512 ],
1513 &[
1514 "subscription.weekly_usage_usd",
1515 "data.subscription.weekly_usage_usd",
1516 ],
1517 ),
1518 sub2api_subscription_limit_snapshot(
1519 value,
1520 "monthly",
1521 &[
1522 "subscription.monthly_limit_usd",
1523 "data.subscription.monthly_limit_usd",
1524 ],
1525 &[
1526 "subscription.monthly_usage_usd",
1527 "data.subscription.monthly_usage_usd",
1528 ],
1529 ),
1530 ];
1531
1532 windows
1533 .into_iter()
1534 .flatten()
1535 .min_by_key(|window| window.remaining)
1536}
1537
1538fn sub2api_usage_snapshot_from_json(
1539 provider: &UsageProviderConfig,
1540 upstream: &UpstreamRef,
1541 value: &serde_json::Value,
1542 fetched_at_ms: u64,
1543 stale_after_ms: Option<u64>,
1544) -> ProviderBalanceSnapshot {
1545 if json_value_at_path(value, "isValid").and_then(bool_from_json) == Some(false) {
1546 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
1547 .with_error("sub2api usage response reported invalid API key");
1548 }
1549
1550 let mode = first_string_from_paths(value, &["mode", "data.mode"]);
1551 let has_subscription = has_any_json_path(value, &["subscription", "data.subscription"]);
1552
1553 if mode.as_deref() == Some("quota_limited") {
1554 let quota_remaining = first_amount_from_paths(
1555 value,
1556 &provider.extract.remaining_balance_paths,
1557 &[
1558 "quota.remaining",
1559 "data.quota.remaining",
1560 "remaining",
1561 "data.remaining",
1562 ],
1563 provider.extract.remaining_divisor,
1564 );
1565 let quota_limit = first_amount_from_paths(
1566 value,
1567 &provider.extract.monthly_budget_paths,
1568 &["quota.limit", "data.quota.limit"],
1569 provider.extract.monthly_budget_divisor,
1570 );
1571 let quota_used = first_amount_from_paths(
1572 value,
1573 &provider.extract.monthly_spent_paths,
1574 &["quota.used", "data.quota.used"],
1575 provider.extract.monthly_spent_divisor,
1576 );
1577 let exhausted = first_bool_from_paths(
1578 value,
1579 &provider.extract.exhausted_paths,
1580 &[
1581 "exhausted",
1582 "data.exhausted",
1583 "quota_exhausted",
1584 "data.quota_exhausted",
1585 ],
1586 )
1587 .or_else(|| quota_remaining.map(UsdAmount::is_zero));
1588
1589 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1590 snapshot.quota_period = Some("quota".to_string());
1591 snapshot.quota_remaining_usd = quota_remaining.map(amount_to_string);
1592 snapshot.quota_limit_usd = quota_limit.map(amount_to_string);
1593 snapshot.quota_used_usd = quota_used.map(amount_to_string);
1594 snapshot.monthly_budget_usd = quota_limit.map(amount_to_string);
1595 snapshot.monthly_spent_usd = quota_used.map(amount_to_string);
1596 snapshot.exhausted = exhausted;
1597 populate_sub2api_usage_fields(&mut snapshot, value);
1598 snapshot.refresh_status(fetched_at_ms);
1599 return snapshot;
1600 }
1601
1602 if mode.as_deref() == Some("unrestricted") && has_subscription {
1603 let limiting_window = sub2api_limiting_subscription_window(value);
1604 let exhausted = first_bool_from_paths(
1605 value,
1606 &provider.extract.exhausted_paths,
1607 &[
1608 "exhausted",
1609 "data.exhausted",
1610 "quota_exhausted",
1611 "data.quota_exhausted",
1612 ],
1613 )
1614 .or_else(|| limiting_window.map(|window| window.remaining.is_zero()))
1615 .or(Some(false));
1616
1617 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1618 if let Some(window) = limiting_window {
1619 snapshot.quota_period = Some(window.period.to_string());
1620 snapshot.quota_remaining_usd = Some(amount_to_string(window.remaining));
1621 snapshot.quota_limit_usd = Some(amount_to_string(window.limit));
1622 snapshot.quota_used_usd = Some(amount_to_string(window.used));
1623 snapshot.monthly_budget_usd = Some(amount_to_string(window.limit));
1624 snapshot.monthly_spent_usd = Some(amount_to_string(window.used));
1625 }
1626 snapshot.exhaustion_affects_routing = false;
1627 snapshot.exhausted = exhausted;
1628 populate_sub2api_usage_fields(&mut snapshot, value);
1629 snapshot.refresh_status(fetched_at_ms);
1630 return snapshot;
1631 }
1632
1633 let mut snapshot =
1634 balance_http_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms);
1635 populate_sub2api_usage_fields(&mut snapshot, value);
1636 snapshot.refresh_status(fetched_at_ms);
1637 snapshot
1638}
1639
1640fn sub2api_auth_me_snapshot_from_json(
1641 provider: &UsageProviderConfig,
1642 upstream: &UpstreamRef,
1643 value: &serde_json::Value,
1644 fetched_at_ms: u64,
1645 stale_after_ms: Option<u64>,
1646) -> ProviderBalanceSnapshot {
1647 if json_value_at_path(value, "code")
1648 .and_then(|value| value.as_i64())
1649 .is_some_and(|code| code != 0)
1650 {
1651 let message = json_value_at_path(value, "message")
1652 .and_then(|value| value.as_str())
1653 .unwrap_or("sub2api auth/me response reported failure");
1654 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
1655 .with_error(message.to_string());
1656 }
1657
1658 balance_http_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
1659}
1660
1661fn rightcode_available_prefixes(value: &serde_json::Value) -> Vec<String> {
1662 array_from_json_path(value, "available_prefixes")
1663 .into_iter()
1664 .flatten()
1665 .filter_map(|value| value.as_str())
1666 .map(str::trim)
1667 .filter(|value| !value.is_empty())
1668 .map(ToOwned::to_owned)
1669 .collect()
1670}
1671
1672fn rightcode_subscription_window(value: &serde_json::Value) -> Option<QuotaWindowSnapshot> {
1673 let limit = json_value_at_path(value, "total_quota").and_then(amount_from_json)?;
1674 if limit < UsdAmount::from_decimal_str("10").unwrap_or(UsdAmount::ZERO) {
1675 return None;
1676 }
1677 let raw_remaining = json_value_at_path(value, "remaining_quota").and_then(amount_from_json)?;
1678 let reset_today = json_value_at_path(value, "reset_today").and_then(bool_from_json);
1679 let remaining = if reset_today == Some(true) {
1680 raw_remaining
1681 } else {
1682 raw_remaining.saturating_add(limit)
1683 };
1684 let used = limit.saturating_sub(remaining);
1685 Some(QuotaWindowSnapshot {
1686 period: "daily",
1687 remaining,
1688 used,
1689 limit,
1690 })
1691}
1692
1693fn rightcode_account_summary_snapshot_from_json(
1694 provider: &UsageProviderConfig,
1695 upstream: &UpstreamRef,
1696 value: &serde_json::Value,
1697 upstream_base_url: &str,
1698 fetched_at_ms: u64,
1699 stale_after_ms: Option<u64>,
1700) -> ProviderBalanceSnapshot {
1701 let balance = json_value_at_path(value, "balance").and_then(amount_from_json);
1702 let provider_prefixes = base_path_prefixes(upstream_base_url);
1703 let mut matched_windows = Vec::new();
1704 let mut matched_plan_names = Vec::new();
1705
1706 if let Some(subscriptions) = array_from_json_path(value, "subscriptions") {
1707 for subscription in subscriptions {
1708 let available_prefixes = rightcode_available_prefixes(subscription);
1709 if !path_prefixes_match(&provider_prefixes, &available_prefixes) {
1710 continue;
1711 }
1712 let Some(window) = rightcode_subscription_window(subscription) else {
1713 continue;
1714 };
1715 matched_windows.push(window);
1716 if let Some(name) = json_value_at_path(subscription, "name").and_then(string_from_json)
1717 {
1718 matched_plan_names.push(name);
1719 }
1720 }
1721 }
1722
1723 if balance.is_none() && matched_windows.is_empty() {
1724 return snapshot_error(
1725 provider,
1726 upstream,
1727 fetched_at_ms,
1728 stale_after_ms,
1729 "rightcode account summary missing balance and matching subscription quota fields",
1730 );
1731 }
1732
1733 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1734 snapshot.total_balance_usd = balance.map(amount_to_string);
1735
1736 if !matched_windows.is_empty() {
1737 let mut remaining = UsdAmount::ZERO;
1738 let mut used = UsdAmount::ZERO;
1739 let mut limit = UsdAmount::ZERO;
1740 for window in matched_windows {
1741 remaining = remaining.saturating_add(window.remaining);
1742 used = used.saturating_add(window.used);
1743 limit = limit.saturating_add(window.limit);
1744 }
1745 snapshot.quota_period = Some("daily".to_string());
1746 snapshot.quota_remaining_usd = Some(amount_to_string(remaining));
1747 snapshot.quota_used_usd = Some(amount_to_string(used));
1748 snapshot.quota_limit_usd = Some(amount_to_string(limit));
1749 if !matched_plan_names.is_empty() {
1750 matched_plan_names.sort();
1751 matched_plan_names.dedup();
1752 snapshot.plan_name = Some(matched_plan_names.join(", "));
1753 }
1754 snapshot.exhausted = Some(remaining.is_zero() && balance.is_none_or(UsdAmount::is_zero));
1755 } else {
1756 snapshot.exhausted = balance.map(UsdAmount::is_zero);
1757 }
1758
1759 snapshot.refresh_status(fetched_at_ms);
1760 snapshot
1761}
1762
1763fn new_api_token_usage_snapshot_from_json(
1764 provider: &UsageProviderConfig,
1765 upstream: &UpstreamRef,
1766 value: &serde_json::Value,
1767 fetched_at_ms: u64,
1768 stale_after_ms: Option<u64>,
1769) -> ProviderBalanceSnapshot {
1770 if json_value_at_path(value, "success").and_then(bool_from_json) == Some(false)
1771 || json_value_at_path(value, "code").and_then(bool_from_json) == Some(false)
1772 {
1773 let message = json_value_at_path(value, "message")
1774 .and_then(|value| value.as_str())
1775 .unwrap_or("new api token usage response reported failure");
1776 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
1777 .with_error(message.to_string());
1778 }
1779
1780 let mut effective = provider.extract.clone();
1781 if effective.remaining_balance_paths.is_empty() {
1782 effective.remaining_balance_paths = vec![
1783 "data.total_available".to_string(),
1784 "data.remain_quota".to_string(),
1785 "total_available".to_string(),
1786 "remain_quota".to_string(),
1787 ];
1788 }
1789 if effective.monthly_spent_paths.is_empty() {
1790 effective.monthly_spent_paths = vec![
1791 "data.total_used".to_string(),
1792 "data.used_quota".to_string(),
1793 "total_used".to_string(),
1794 "used_quota".to_string(),
1795 ];
1796 }
1797 if effective.monthly_budget_paths.is_empty() {
1798 effective.monthly_budget_paths = vec![
1799 "data.total_granted".to_string(),
1800 "total_granted".to_string(),
1801 ];
1802 }
1803 effective.remaining_divisor = effective.remaining_divisor.or(Some(500_000));
1804 effective.monthly_spent_divisor = effective.monthly_spent_divisor.or(Some(500_000));
1805 effective.monthly_budget_divisor = effective.monthly_budget_divisor.or(Some(500_000));
1806
1807 let unlimited_quota =
1808 first_bool_from_paths(value, &[], &["data.unlimited_quota", "unlimited_quota"])
1809 == Some(true);
1810 let remaining_balance = first_amount_from_paths(
1811 value,
1812 &effective.remaining_balance_paths,
1813 &[],
1814 effective.remaining_divisor,
1815 );
1816 let monthly_spent = first_amount_from_paths(
1817 value,
1818 &effective.monthly_spent_paths,
1819 &[],
1820 effective.monthly_spent_divisor,
1821 );
1822 let monthly_budget = first_amount_from_paths(
1823 value,
1824 &effective.monthly_budget_paths,
1825 &[],
1826 effective.monthly_budget_divisor,
1827 )
1828 .or_else(|| match (remaining_balance, monthly_spent) {
1829 (Some(remaining), Some(spent)) => Some(remaining.saturating_add(spent)),
1830 _ => None,
1831 });
1832 let exhausted = if unlimited_quota {
1833 Some(false)
1834 } else {
1835 first_bool_from_paths(
1836 value,
1837 &effective.exhausted_paths,
1838 &["data.exhausted", "exhausted"],
1839 )
1840 .or_else(|| remaining_balance.map(UsdAmount::is_zero))
1841 };
1842
1843 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1844 snapshot.plan_name = first_string_from_paths(value, &["data.name", "name"]);
1845 snapshot.unlimited_quota = Some(unlimited_quota);
1846 if !unlimited_quota {
1847 snapshot.quota_period = Some("token".to_string());
1848 snapshot.quota_remaining_usd = remaining_balance.map(amount_to_string);
1849 snapshot.quota_limit_usd = monthly_budget.map(amount_to_string);
1850 snapshot.quota_used_usd = monthly_spent.map(amount_to_string);
1851 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
1852 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1853 } else {
1854 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1855 }
1856 snapshot.exhausted = exhausted;
1857 snapshot.refresh_status(fetched_at_ms);
1858 snapshot
1859}
1860
1861fn new_api_snapshot_from_json(
1862 provider: &UsageProviderConfig,
1863 upstream: &UpstreamRef,
1864 value: &serde_json::Value,
1865 fetched_at_ms: u64,
1866 stale_after_ms: Option<u64>,
1867) -> ProviderBalanceSnapshot {
1868 if json_value_at_path(value, "success").and_then(bool_from_json) == Some(false) {
1869 let message = json_value_at_path(value, "message")
1870 .and_then(|value| value.as_str())
1871 .unwrap_or("new api balance response reported failure");
1872 return base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms)
1873 .with_error(message.to_string());
1874 }
1875
1876 let mut effective = provider.extract.clone();
1877 if effective.remaining_balance_paths.is_empty() {
1878 effective.remaining_balance_paths = vec!["data.quota".to_string(), "quota".to_string()];
1879 }
1880 if effective.monthly_spent_paths.is_empty() {
1881 effective.monthly_spent_paths =
1882 vec!["data.used_quota".to_string(), "used_quota".to_string()];
1883 }
1884 effective.remaining_divisor = effective.remaining_divisor.or(Some(500_000));
1885 effective.monthly_spent_divisor = effective.monthly_spent_divisor.or(Some(500_000));
1886 effective.monthly_budget_divisor = effective.monthly_budget_divisor.or(Some(500_000));
1887
1888 let remaining_balance = first_amount_from_paths(
1889 value,
1890 &effective.remaining_balance_paths,
1891 &[],
1892 effective.remaining_divisor,
1893 );
1894 let monthly_spent = first_amount_from_paths(
1895 value,
1896 &effective.monthly_spent_paths,
1897 &[],
1898 effective.monthly_spent_divisor,
1899 );
1900 let monthly_budget = first_amount_from_paths(
1901 value,
1902 &effective.monthly_budget_paths,
1903 &["data.total_quota", "total_quota"],
1904 effective.monthly_budget_divisor,
1905 )
1906 .or_else(|| match (remaining_balance, monthly_spent) {
1907 (Some(remaining), Some(spent)) => Some(remaining.saturating_add(spent)),
1908 _ => None,
1909 });
1910 let unlimited_quota =
1911 first_bool_from_paths(value, &[], &["data.unlimited_quota", "unlimited_quota"])
1912 == Some(true);
1913 let exhausted = if unlimited_quota {
1914 Some(false)
1915 } else {
1916 first_bool_from_paths(
1917 value,
1918 &effective.exhausted_paths,
1919 &["data.exhausted", "exhausted"],
1920 )
1921 .or_else(|| remaining_balance.map(UsdAmount::is_zero))
1922 };
1923
1924 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1925 snapshot.unlimited_quota = Some(unlimited_quota);
1926 if !unlimited_quota {
1927 snapshot.quota_period = Some("quota".to_string());
1928 snapshot.quota_remaining_usd = remaining_balance.map(amount_to_string);
1929 snapshot.quota_limit_usd = monthly_budget.map(amount_to_string);
1930 snapshot.quota_used_usd = monthly_spent.map(amount_to_string);
1931 snapshot.monthly_budget_usd = monthly_budget.map(amount_to_string);
1932 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1933 } else {
1934 snapshot.monthly_spent_usd = monthly_spent.map(amount_to_string);
1935 }
1936 snapshot.exhausted = exhausted;
1937 snapshot.refresh_status(fetched_at_ms);
1938 snapshot
1939}
1940
1941fn openai_cost_result_usd_amount(result: &serde_json::Value) -> Option<UsdAmount> {
1942 let amount = json_value_at_path(result, "amount.value").and_then(amount_from_json)?;
1943 let currency = json_value_at_path(result, "amount.currency").and_then(|value| value.as_str());
1944 match currency {
1945 Some(currency) if currency.eq_ignore_ascii_case("usd") => Some(amount),
1946 None => Some(amount),
1947 _ => None,
1948 }
1949}
1950
1951fn openai_organization_costs_total(value: &serde_json::Value) -> Option<UsdAmount> {
1952 let buckets = json_value_at_path(value, "data")?.as_array()?;
1953 let mut total = UsdAmount::ZERO;
1954
1955 for bucket in buckets {
1956 let Some(results) =
1957 json_value_at_path(bucket, "results").and_then(|value| value.as_array())
1958 else {
1959 continue;
1960 };
1961 for result in results {
1962 if let Some(amount) = openai_cost_result_usd_amount(result) {
1963 total = total.saturating_add(amount);
1964 }
1965 }
1966 }
1967
1968 Some(total)
1969}
1970
1971fn openai_organization_costs_snapshot_from_json(
1972 provider: &UsageProviderConfig,
1973 upstream: &UpstreamRef,
1974 value: &serde_json::Value,
1975 fetched_at_ms: u64,
1976 stale_after_ms: Option<u64>,
1977) -> ProviderBalanceSnapshot {
1978 let spent = first_amount_from_paths(
1979 value,
1980 &provider.extract.monthly_spent_paths,
1981 &[],
1982 provider.extract.monthly_spent_divisor,
1983 )
1984 .or_else(|| openai_organization_costs_total(value));
1985
1986 let mut snapshot = base_snapshot(provider, upstream, fetched_at_ms, stale_after_ms);
1987 snapshot.monthly_spent_usd = spent.map(amount_to_string);
1988 snapshot.exhausted = None;
1989 snapshot.exhaustion_affects_routing = false;
1990 snapshot.refresh_status(fetched_at_ms);
1991 snapshot
1992}
1993
1994async fn update_usage_exhausted(
1995 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
1996 state: &Arc<ProxyState>,
1997 cfg: &ProxyConfig,
1998 service_name: &str,
1999 upstreams: &[UpstreamRef],
2000 exhausted: bool,
2001) {
2002 if let Ok(mut map) = lb_states.lock() {
2003 for uref in upstreams {
2004 let service = match service_manager(cfg, service_name).station(&uref.station_name) {
2005 Some(s) => s,
2006 None => continue,
2007 };
2008
2009 let entry = map
2010 .entry(uref.station_name.clone())
2011 .or_insert_with(LbState::default);
2012 entry.ensure_layout(service.name.as_str(), &service.upstreams);
2013 if uref.index < entry.usage_exhausted.len() {
2014 entry.usage_exhausted[uref.index] = exhausted;
2015 }
2016 }
2017 }
2018
2019 for uref in upstreams {
2020 if let Some(endpoint_key) = uref.provider_endpoint.clone() {
2021 state
2022 .set_provider_endpoint_usage_exhausted(service_name, endpoint_key, exhausted)
2023 .await;
2024 }
2025 }
2026}
2027
2028fn provider_hosts_for_diagnostics(
2029 cfg: &ProxyConfig,
2030 service_name: &str,
2031 provider: &UsageProviderConfig,
2032) -> Vec<String> {
2033 let mut hosts: Vec<String> = Vec::new();
2034 for service in service_manager(cfg, service_name).stations().values() {
2035 for upstream in &service.upstreams {
2036 if domain_matches(&upstream.base_url, &provider.domains)
2037 && let Ok(url) = reqwest::Url::parse(&upstream.base_url)
2038 && let Some(host) = url.host_str()
2039 {
2040 hosts.push(host.to_string());
2041 }
2042 }
2043 }
2044 hosts.sort();
2045 hosts.dedup();
2046 hosts
2047}
2048
2049fn warn_if_provider_spans_hosts(
2050 cfg: &ProxyConfig,
2051 service_name: &str,
2052 provider: &UsageProviderConfig,
2053) {
2054 let hosts = provider_hosts_for_diagnostics(cfg, service_name, provider);
2055 if hosts.len() > 1 {
2056 warn!(
2057 "usage provider '{}' is associated with multiple hosts: {:?}; \
2058将按统一额度处理这些 upstream,如需区分配额请拆分为多个 provider 配置",
2059 provider.id, hosts
2060 );
2061 }
2062}
2063
2064fn snapshot_from_provider_json(
2065 provider: &UsageProviderConfig,
2066 upstream: &UpstreamRef,
2067 value: &serde_json::Value,
2068 upstream_base_url: &str,
2069 fetched_at_ms: u64,
2070 stale_after_ms: Option<u64>,
2071) -> ProviderBalanceSnapshot {
2072 match provider.kind {
2073 ProviderKind::BudgetHttpJson => {
2074 budget_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
2075 }
2076 ProviderKind::YescodeProfile => {
2077 yescode_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
2078 }
2079 ProviderKind::OpenAiBalanceHttpJson => balance_http_snapshot_from_json(
2080 provider,
2081 upstream,
2082 value,
2083 fetched_at_ms,
2084 stale_after_ms,
2085 ),
2086 ProviderKind::Sub2ApiUsage => sub2api_usage_snapshot_from_json(
2087 provider,
2088 upstream,
2089 value,
2090 fetched_at_ms,
2091 stale_after_ms,
2092 ),
2093 ProviderKind::Sub2ApiAuthMe => sub2api_auth_me_snapshot_from_json(
2094 provider,
2095 upstream,
2096 value,
2097 fetched_at_ms,
2098 stale_after_ms,
2099 ),
2100 ProviderKind::NewApiTokenUsage => new_api_token_usage_snapshot_from_json(
2101 provider,
2102 upstream,
2103 value,
2104 fetched_at_ms,
2105 stale_after_ms,
2106 ),
2107 ProviderKind::NewApiUserSelf => {
2108 new_api_snapshot_from_json(provider, upstream, value, fetched_at_ms, stale_after_ms)
2109 }
2110 ProviderKind::RightCodeAccountSummary => rightcode_account_summary_snapshot_from_json(
2111 provider,
2112 upstream,
2113 value,
2114 upstream_base_url,
2115 fetched_at_ms,
2116 stale_after_ms,
2117 ),
2118 ProviderKind::OpenAiOrganizationCosts => openai_organization_costs_snapshot_from_json(
2119 provider,
2120 upstream,
2121 value,
2122 fetched_at_ms,
2123 stale_after_ms,
2124 ),
2125 }
2126}
2127
2128async fn refresh_provider_target(
2129 params: RefreshProviderTargetParams<'_>,
2130) -> UsageProviderRefreshOutcome {
2131 let RefreshProviderTargetParams {
2132 client,
2133 provider,
2134 target,
2135 cfg,
2136 lb_states,
2137 state,
2138 service_name,
2139 interval_secs,
2140 } = params;
2141
2142 let upstreams = vec![target.upstream.clone()];
2143 let fetched_at_ms = unix_now_ms();
2144 let stale_after_ms = stale_after_ms(fetched_at_ms, interval_secs);
2145
2146 let Some(token) = resolve_token(provider, &upstreams, cfg, service_name) else {
2147 let snapshot = if provider.kind == ProviderKind::OpenAiOrganizationCosts {
2148 base_snapshot(provider, &upstreams[0], fetched_at_ms, stale_after_ms)
2149 } else {
2150 base_snapshot(provider, &upstreams[0], fetched_at_ms, stale_after_ms)
2151 .with_error("no usable token; checked provider token_env and upstream auth")
2152 };
2153 state
2154 .record_provider_balance_snapshot(service_name, snapshot)
2155 .await;
2156 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
2157 if provider.kind == ProviderKind::OpenAiOrganizationCosts {
2158 warn!(
2159 "usage provider '{}' is missing OPENAI_ADMIN_KEY; OpenAI official costs stay unknown",
2160 provider.id
2161 );
2162 } else {
2163 warn!(
2164 "usage provider '{}' has no usable token (checked token_env and associated upstream auth_token); \
2165跳过本次用量查询,请检查 usage_providers.json 和 ~/.codex-helper/config.json",
2166 provider.id
2167 );
2168 }
2169 return UsageProviderRefreshOutcome::MissingToken;
2170 };
2171
2172 match poll_provider_http_json(client, provider, &target.base_url, &token).await {
2173 Ok(value) => {
2174 let snapshot = snapshot_from_provider_json(
2175 provider,
2176 &upstreams[0],
2177 &value,
2178 &target.base_url,
2179 fetched_at_ms,
2180 stale_after_ms,
2181 );
2182 let exhausted_for_lb = snapshot.routing_exhausted();
2183 update_usage_exhausted(
2184 lb_states,
2185 state,
2186 cfg,
2187 service_name,
2188 &upstreams,
2189 exhausted_for_lb,
2190 )
2191 .await;
2192 state
2193 .record_provider_balance_snapshot(service_name, snapshot)
2194 .await;
2195 info!(
2196 "usage provider '{}' refreshed {}[{}], exhausted = {}, routing_trusted = {}",
2197 provider.id,
2198 target.upstream.station_name,
2199 target.upstream.index,
2200 exhausted_for_lb,
2201 provider.trust_exhaustion_for_routing
2202 );
2203 UsageProviderRefreshOutcome::Refreshed
2204 }
2205 Err(err) => {
2206 state
2207 .record_provider_balance_snapshot(
2208 service_name,
2209 base_snapshot(provider, &upstreams[0], fetched_at_ms, stale_after_ms)
2210 .with_error(err.to_string()),
2211 )
2212 .await;
2213 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
2214 warn!(
2215 "usage provider '{}' poll failed for {}[{}]: {}",
2216 provider.id, target.upstream.station_name, target.upstream.index, err
2217 );
2218 UsageProviderRefreshOutcome::Failed
2219 }
2220 }
2221}
2222
2223struct ConfiguredRefreshJob<'a> {
2224 provider: &'a UsageProviderConfig,
2225 target: UsageProviderTarget,
2226 interval_secs: u64,
2227}
2228
2229struct AutoRefreshJob {
2230 target: UsageProviderTarget,
2231}
2232
2233async fn run_configured_refresh_job<'a>(
2234 client: &'a Client,
2235 job: ConfiguredRefreshJob<'a>,
2236 cfg: &'a ProxyConfig,
2237 lb_states: &'a Arc<Mutex<HashMap<String, LbState>>>,
2238 state: &'a Arc<ProxyState>,
2239 service_name: &'a str,
2240) -> (String, UsageProviderRefreshOutcome) {
2241 let provider_id = job.provider.id.clone();
2242 let outcome = refresh_provider_target(RefreshProviderTargetParams {
2243 client,
2244 provider: job.provider,
2245 target: &job.target,
2246 cfg,
2247 lb_states,
2248 state,
2249 service_name,
2250 interval_secs: job.interval_secs,
2251 })
2252 .await;
2253 (provider_id, outcome)
2254}
2255
2256async fn run_auto_refresh_job(
2257 client: &Client,
2258 job: AutoRefreshJob,
2259 cfg: &ProxyConfig,
2260 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
2261 state: &Arc<ProxyState>,
2262 service_name: &str,
2263) -> UsageProviderRefreshOutcome {
2264 auto_probe_provider_target(client, &job.target, cfg, lb_states, state, service_name).await
2265}
2266
2267async fn run_configured_refresh_jobs<'a>(
2268 client: &'a Client,
2269 jobs: Vec<ConfiguredRefreshJob<'a>>,
2270 cfg: &'a ProxyConfig,
2271 lb_states: &'a Arc<Mutex<HashMap<String, LbState>>>,
2272 state: &'a Arc<ProxyState>,
2273 service_name: &'a str,
2274) -> Vec<(String, UsageProviderRefreshOutcome)> {
2275 let mut pending = jobs.into_iter();
2276 let mut running = FuturesUnordered::new();
2277 let mut results = Vec::new();
2278 let concurrency = BALANCE_REFRESH_CONCURRENCY.max(1);
2279
2280 for _ in 0..concurrency {
2281 let Some(job) = pending.next() else {
2282 break;
2283 };
2284 running.push(run_configured_refresh_job(
2285 client,
2286 job,
2287 cfg,
2288 lb_states,
2289 state,
2290 service_name,
2291 ));
2292 }
2293
2294 while let Some(result) = running.next().await {
2295 results.push(result);
2296 if let Some(job) = pending.next() {
2297 running.push(run_configured_refresh_job(
2298 client,
2299 job,
2300 cfg,
2301 lb_states,
2302 state,
2303 service_name,
2304 ));
2305 }
2306 }
2307
2308 results
2309}
2310
2311async fn run_auto_refresh_jobs(
2312 client: &Client,
2313 jobs: Vec<AutoRefreshJob>,
2314 cfg: &ProxyConfig,
2315 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
2316 state: &Arc<ProxyState>,
2317 service_name: &str,
2318) -> Vec<UsageProviderRefreshOutcome> {
2319 let mut pending = jobs.into_iter();
2320 let mut running = FuturesUnordered::new();
2321 let mut results = Vec::new();
2322 let concurrency = BALANCE_REFRESH_CONCURRENCY.max(1);
2323
2324 for _ in 0..concurrency {
2325 let Some(job) = pending.next() else {
2326 break;
2327 };
2328 running.push(run_auto_refresh_job(
2329 client,
2330 job,
2331 cfg,
2332 lb_states,
2333 state,
2334 service_name,
2335 ));
2336 }
2337
2338 while let Some(result) = running.next().await {
2339 results.push(result);
2340 if let Some(job) = pending.next() {
2341 running.push(run_auto_refresh_job(
2342 client,
2343 job,
2344 cfg,
2345 lb_states,
2346 state,
2347 service_name,
2348 ));
2349 }
2350 }
2351
2352 results
2353}
2354
2355fn auto_snapshot_is_usable(snapshot: &ProviderBalanceSnapshot) -> bool {
2356 snapshot.error.is_none()
2357 && matches!(
2358 snapshot.status,
2359 BalanceSnapshotStatus::Ok | BalanceSnapshotStatus::Exhausted
2360 )
2361}
2362
2363async fn auto_probe_provider_target(
2364 client: &Client,
2365 target: &UsageProviderTarget,
2366 cfg: &ProxyConfig,
2367 lb_states: &Arc<Mutex<HashMap<String, LbState>>>,
2368 state: &Arc<ProxyState>,
2369 service_name: &str,
2370) -> UsageProviderRefreshOutcome {
2371 let upstreams = vec![target.upstream.clone()];
2372 let fetched_at_ms = unix_now_ms();
2373 let interval_secs = DEFAULT_POLL_INTERVAL_SECS;
2374 let stale_after_ms = stale_after_ms(fetched_at_ms, interval_secs);
2375
2376 if is_official_openai_base_url(&target.base_url) {
2377 let provider = auto_openai_official_provider(target);
2378 let Some(token) = resolve_token(&provider, &upstreams, cfg, service_name) else {
2379 state
2380 .record_provider_balance_snapshot(
2381 service_name,
2382 base_snapshot(&provider, &target.upstream, fetched_at_ms, stale_after_ms),
2383 )
2384 .await;
2385 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
2386 warn!(
2387 "OpenAI organization costs require OPENAI_ADMIN_KEY; balance stays unknown for {}[{}]",
2388 target.upstream.station_name, target.upstream.index
2389 );
2390 return UsageProviderRefreshOutcome::MissingToken;
2391 };
2392
2393 return match poll_provider_http_json(client, &provider, &target.base_url, &token).await {
2394 Ok(value) => {
2395 let snapshot = snapshot_from_provider_json(
2396 &provider,
2397 &upstreams[0],
2398 &value,
2399 &target.base_url,
2400 fetched_at_ms,
2401 stale_after_ms,
2402 );
2403 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false)
2404 .await;
2405 state
2406 .record_provider_balance_snapshot(service_name, snapshot)
2407 .await;
2408 UsageProviderRefreshOutcome::Refreshed
2409 }
2410 Err(err) => {
2411 state
2412 .record_provider_balance_snapshot(
2413 service_name,
2414 base_snapshot(&provider, &upstreams[0], fetched_at_ms, stale_after_ms)
2415 .with_error(err.to_string()),
2416 )
2417 .await;
2418 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false)
2419 .await;
2420 warn!(
2421 "OpenAI organization costs poll failed for {}[{}]: {}",
2422 target.upstream.station_name, target.upstream.index, err
2423 );
2424 UsageProviderRefreshOutcome::Failed
2425 }
2426 };
2427 }
2428
2429 let first_provider = auto_usage_provider(target, first_auto_probe_kind(target));
2430
2431 let Some(token) = resolve_token(&first_provider, &upstreams, cfg, service_name) else {
2432 state
2433 .record_provider_balance_snapshot(
2434 service_name,
2435 base_snapshot(
2436 &first_provider,
2437 &target.upstream,
2438 fetched_at_ms,
2439 stale_after_ms,
2440 )
2441 .with_error("no usable token; checked upstream auth"),
2442 )
2443 .await;
2444 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
2445 return UsageProviderRefreshOutcome::MissingToken;
2446 };
2447
2448 let mut last_error: Option<String> = None;
2449 for kind in AUTO_PROBE_KINDS {
2450 if kind == ProviderKind::RightCodeAccountSummary && !is_rightcode_base_url(&target.base_url)
2451 {
2452 continue;
2453 }
2454 let provider = auto_usage_provider(target, kind);
2455 match poll_provider_http_json(client, &provider, &target.base_url, &token).await {
2456 Ok(value) => {
2457 let snapshot = snapshot_from_provider_json(
2458 &provider,
2459 &upstreams[0],
2460 &value,
2461 &target.base_url,
2462 fetched_at_ms,
2463 stale_after_ms,
2464 );
2465 if auto_snapshot_is_usable(&snapshot) {
2466 let exhausted_for_lb = snapshot.routing_exhausted();
2467 update_usage_exhausted(
2468 lb_states,
2469 state,
2470 cfg,
2471 service_name,
2472 &upstreams,
2473 exhausted_for_lb,
2474 )
2475 .await;
2476 state
2477 .record_provider_balance_snapshot(service_name, snapshot)
2478 .await;
2479 info!(
2480 "auto usage provider '{}' refreshed {}[{}] via {:?}, exhausted = {}",
2481 provider.id,
2482 target.upstream.station_name,
2483 target.upstream.index,
2484 kind,
2485 exhausted_for_lb
2486 );
2487 return UsageProviderRefreshOutcome::Refreshed;
2488 }
2489 last_error = snapshot.error.or_else(|| {
2490 Some(format!(
2491 "auto probe {:?} returned no usable balance fields",
2492 kind
2493 ))
2494 });
2495 }
2496 Err(err) => {
2497 last_error = Some(err.to_string());
2498 }
2499 }
2500 }
2501
2502 if let Some(error) = last_error {
2503 warn!(
2504 "auto usage provider '{}' found no usable balance endpoint for {}[{}]: {}",
2505 first_provider.id, target.upstream.station_name, target.upstream.index, error
2506 );
2507 state
2508 .record_provider_balance_snapshot(
2509 service_name,
2510 base_snapshot(
2511 &first_provider,
2512 &target.upstream,
2513 fetched_at_ms,
2514 stale_after_ms,
2515 )
2516 .with_error(error),
2517 )
2518 .await;
2519 update_usage_exhausted(lb_states, state, cfg, service_name, &upstreams, false).await;
2520 }
2521 UsageProviderRefreshOutcome::Failed
2522}
2523
2524pub async fn refresh_balances_for_service(
2525 client: &Client,
2526 cfg: Arc<ProxyConfig>,
2527 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
2528 state: Arc<ProxyState>,
2529 service_name: &str,
2530 station_name_filter: Option<&str>,
2531 provider_id_filter: Option<&str>,
2532) -> UsageProviderRefreshSummary {
2533 if cfg!(test) {
2535 return UsageProviderRefreshSummary::default();
2536 }
2537
2538 let station_name_filter = station_name_filter
2539 .map(str::trim)
2540 .filter(|value| !value.is_empty());
2541 let provider_id_filter = provider_id_filter
2542 .map(str::trim)
2543 .filter(|value| !value.is_empty());
2544 let providers_file = load_providers();
2545 let mut summary = UsageProviderRefreshSummary {
2546 providers_configured: providers_file.providers.len(),
2547 ..UsageProviderRefreshSummary::default()
2548 };
2549 let configured_keys = if provider_id_filter.is_none() {
2550 configured_target_keys(
2551 &cfg,
2552 service_name,
2553 &providers_file.providers,
2554 station_name_filter,
2555 )
2556 } else {
2557 HashSet::new()
2558 };
2559
2560 let poll_map = LAST_USAGE_POLL.get_or_init(|| Mutex::new(HashMap::new()));
2561 let mut configured_jobs = Vec::new();
2562 for provider in &providers_file.providers {
2563 if provider_id_filter.is_some_and(|filter| filter != provider.id.as_str()) {
2564 continue;
2565 }
2566
2567 let targets = matching_provider_targets(&cfg, service_name, provider, station_name_filter);
2568 if targets.is_empty() {
2569 continue;
2570 }
2571
2572 summary.providers_matched += 1;
2573 summary.upstreams_matched += targets.len();
2574 warn_if_provider_spans_hosts(&cfg, service_name, provider);
2575
2576 let interval_secs = snapshot_refresh_interval_secs(provider);
2577 for target in targets {
2578 summary.attempted += 1;
2579 configured_jobs.push(ConfiguredRefreshJob {
2580 provider,
2581 target,
2582 interval_secs,
2583 });
2584 }
2585 }
2586
2587 let mut refreshed_provider_ids = HashSet::new();
2588 if !configured_jobs.is_empty() {
2589 for (provider_id, outcome) in run_configured_refresh_jobs(
2590 client,
2591 configured_jobs,
2592 &cfg,
2593 &lb_states,
2594 &state,
2595 service_name,
2596 )
2597 .await
2598 {
2599 match outcome {
2600 UsageProviderRefreshOutcome::Refreshed => {
2601 summary.refreshed += 1;
2602 refreshed_provider_ids.insert(provider_id);
2603 }
2604 UsageProviderRefreshOutcome::Failed => summary.failed += 1,
2605 UsageProviderRefreshOutcome::MissingToken => summary.missing_token += 1,
2606 }
2607 }
2608 }
2609
2610 if provider_id_filter.is_none() {
2611 let mut auto_jobs = Vec::new();
2612 for target in usage_provider_targets(&cfg, service_name, station_name_filter) {
2613 if configured_keys.contains(&target_key(&target)) {
2614 continue;
2615 }
2616
2617 summary.attempted += 1;
2618 summary.auto_attempted += 1;
2619 auto_jobs.push(AutoRefreshJob { target });
2620 }
2621
2622 if !auto_jobs.is_empty() {
2623 for outcome in
2624 run_auto_refresh_jobs(client, auto_jobs, &cfg, &lb_states, &state, service_name)
2625 .await
2626 {
2627 match outcome {
2628 UsageProviderRefreshOutcome::Refreshed => {
2629 summary.refreshed += 1;
2630 summary.auto_refreshed += 1;
2631 }
2632 UsageProviderRefreshOutcome::Failed => {
2633 summary.failed += 1;
2634 summary.auto_failed += 1;
2635 }
2636 UsageProviderRefreshOutcome::MissingToken => {
2637 summary.missing_token += 1;
2638 }
2639 }
2640 }
2641 }
2642 }
2643
2644 if !refreshed_provider_ids.is_empty()
2645 && let Ok(mut map) = poll_map.lock()
2646 {
2647 let now = Instant::now();
2648 for provider_id in refreshed_provider_ids {
2649 map.insert(provider_id, now);
2650 }
2651 }
2652
2653 summary
2654}
2655
2656pub async fn poll_for_codex_upstream(
2659 client: Client,
2660 cfg: Arc<ProxyConfig>,
2661 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
2662 state: Arc<ProxyState>,
2663 service_name: &str,
2664 station_name: &str,
2665 upstream_index: usize,
2666) {
2667 let current_target =
2668 usage_provider_target_for_legacy_upstream(&cfg, service_name, station_name, upstream_index);
2669 poll_for_codex_target(client, cfg, lb_states, state, service_name, current_target).await;
2670}
2671
2672pub async fn poll_for_codex_provider_endpoint(
2676 client: Client,
2677 cfg: Arc<ProxyConfig>,
2678 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
2679 state: Arc<ProxyState>,
2680 service_name: &str,
2681 provider_endpoint: ProviderEndpointKey,
2682) {
2683 let current_target =
2684 usage_provider_target_for_provider_endpoint(&cfg, service_name, &provider_endpoint);
2685 poll_for_codex_target(client, cfg, lb_states, state, service_name, current_target).await;
2686}
2687
2688async fn poll_for_codex_target(
2689 client: Client,
2690 cfg: Arc<ProxyConfig>,
2691 lb_states: Arc<Mutex<HashMap<String, LbState>>>,
2692 state: Arc<ProxyState>,
2693 service_name: &str,
2694 current_target: Option<UsageProviderTarget>,
2695) {
2696 if cfg!(test) {
2699 return;
2700 }
2701
2702 let providers_file = load_providers();
2703 let Some(current_target) = current_target else {
2704 return;
2705 };
2706
2707 let now = Instant::now();
2708 let poll_map = LAST_USAGE_POLL.get_or_init(|| Mutex::new(HashMap::new()));
2709 let mut matched_configured_provider = false;
2710 let mut configured_jobs = Vec::new();
2711
2712 for provider in &providers_file.providers {
2713 if !domain_matches(¤t_target.base_url, &provider.domains) {
2714 continue;
2715 }
2716 matched_configured_provider = true;
2717
2718 let Some(interval_secs) = effective_poll_interval_secs(provider) else {
2719 continue;
2720 };
2721
2722 {
2723 let mut map = match poll_map.lock() {
2724 Ok(m) => m,
2725 Err(_) => continue,
2726 };
2727 if let Some(last) = map.get(&provider.id)
2728 && now.duration_since(*last) < Duration::from_secs(interval_secs)
2729 {
2730 continue;
2731 }
2732 map.insert(provider.id.clone(), now);
2733 }
2734
2735 warn_if_provider_spans_hosts(&cfg, service_name, provider);
2736 configured_jobs.push(ConfiguredRefreshJob {
2737 provider,
2738 target: current_target.clone(),
2739 interval_secs,
2740 });
2741 }
2742
2743 if !configured_jobs.is_empty() {
2744 let _ = run_configured_refresh_jobs(
2745 &client,
2746 configured_jobs,
2747 &cfg,
2748 &lb_states,
2749 &state,
2750 service_name,
2751 )
2752 .await;
2753 }
2754
2755 if matched_configured_provider {
2756 return;
2757 }
2758
2759 let auto_provider = if is_official_openai_base_url(¤t_target.base_url) {
2760 auto_openai_official_provider(¤t_target)
2761 } else {
2762 auto_usage_provider(¤t_target, first_auto_probe_kind(¤t_target))
2763 };
2764 let Some(interval_secs) = effective_poll_interval_secs(&auto_provider) else {
2765 return;
2766 };
2767
2768 {
2769 let mut map = match poll_map.lock() {
2770 Ok(m) => m,
2771 Err(_) => return,
2772 };
2773 if let Some(last) = map.get(&auto_provider.id)
2774 && now.duration_since(*last) < Duration::from_secs(interval_secs)
2775 {
2776 return;
2777 }
2778 map.insert(auto_provider.id.clone(), now);
2779 }
2780
2781 let _ = auto_probe_provider_target(
2782 &client,
2783 ¤t_target,
2784 &cfg,
2785 &lb_states,
2786 &state,
2787 service_name,
2788 )
2789 .await;
2790}
2791
2792#[cfg(test)]
2793mod tests {
2794 use super::*;
2795
2796 use crate::balance::BalanceSnapshotStatus;
2797 use crate::config::{ServiceConfig, UpstreamAuth, UpstreamConfig};
2798
2799 fn provider(id: &str, kind: ProviderKind) -> UsageProviderConfig {
2800 UsageProviderConfig {
2801 id: id.to_string(),
2802 kind,
2803 domains: vec!["example.com".to_string()],
2804 endpoint: "https://example.com/usage".to_string(),
2805 token_env: None,
2806 require_token_env: false,
2807 poll_interval_secs: Some(60),
2808 refresh_on_request: true,
2809 trust_exhaustion_for_routing: true,
2810 headers: BTreeMap::new(),
2811 variables: BTreeMap::new(),
2812 extract: UsageProviderExtractConfig::default(),
2813 }
2814 }
2815
2816 fn upstream() -> UpstreamRef {
2817 UpstreamRef {
2818 station_name: "right".to_string(),
2819 index: 1,
2820 provider_endpoint: None,
2821 }
2822 }
2823
2824 fn upstream_config(base_url: &str) -> UpstreamConfig {
2825 UpstreamConfig {
2826 base_url: base_url.to_string(),
2827 auth: UpstreamAuth::default(),
2828 tags: HashMap::new(),
2829 supported_models: HashMap::new(),
2830 model_mapping: HashMap::new(),
2831 }
2832 }
2833
2834 fn endpoint_upstream_config(
2835 base_url: &str,
2836 provider_id: &str,
2837 endpoint_id: &str,
2838 ) -> UpstreamConfig {
2839 let mut upstream = upstream_config(base_url);
2840 upstream
2841 .tags
2842 .insert("provider_id".to_string(), provider_id.to_string());
2843 upstream
2844 .tags
2845 .insert("endpoint_id".to_string(), endpoint_id.to_string());
2846 upstream
2847 }
2848
2849 fn service_config(name: &str, upstreams: Vec<UpstreamConfig>) -> ServiceConfig {
2850 ServiceConfig {
2851 name: name.to_string(),
2852 alias: None,
2853 enabled: true,
2854 level: 1,
2855 upstreams,
2856 }
2857 }
2858
2859 fn proxy_config(stations: Vec<ServiceConfig>) -> ProxyConfig {
2860 let mut cfg = ProxyConfig::default();
2861 cfg.codex.configs = stations
2862 .into_iter()
2863 .map(|station| (station.name.clone(), station))
2864 .collect();
2865 cfg
2866 }
2867
2868 #[test]
2869 fn budget_snapshot_reports_monthly_budget_and_exhaustion() {
2870 let snapshot = budget_snapshot_from_json(
2871 &provider("packycode", ProviderKind::BudgetHttpJson),
2872 &upstream(),
2873 &serde_json::json!({
2874 "monthly_budget_usd": "10.50",
2875 "monthly_spent_usd": 10.5
2876 }),
2877 100,
2878 Some(1_000),
2879 );
2880
2881 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
2882 assert_eq!(snapshot.exhausted, Some(true));
2883 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("10.5"));
2884 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("10.5"));
2885 }
2886
2887 #[test]
2888 fn budget_snapshot_keeps_missing_amounts_unknown() {
2889 let snapshot = budget_snapshot_from_json(
2890 &provider("packycode", ProviderKind::BudgetHttpJson),
2891 &upstream(),
2892 &serde_json::json!({}),
2893 100,
2894 Some(1_000),
2895 );
2896
2897 assert_eq!(snapshot.status, BalanceSnapshotStatus::Unknown);
2898 assert_eq!(snapshot.exhausted, None);
2899 }
2900
2901 #[test]
2902 fn yescode_snapshot_sums_subscription_and_paygo_balances() {
2903 let snapshot = yescode_snapshot_from_json(
2904 &provider("yescode", ProviderKind::YescodeProfile),
2905 &upstream(),
2906 &serde_json::json!({
2907 "subscription_balance": "1.25",
2908 "pay_as_you_go_balance": 2.5
2909 }),
2910 100,
2911 Some(1_000),
2912 );
2913
2914 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
2915 assert_eq!(snapshot.exhausted, Some(false));
2916 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("3.75"));
2917 assert_eq!(snapshot.subscription_balance_usd.as_deref(), Some("1.25"));
2918 assert_eq!(snapshot.paygo_balance_usd.as_deref(), Some("2.5"));
2919 }
2920
2921 #[test]
2922 fn openai_balance_endpoint_defaults_to_base_user_balance_without_v1() {
2923 let mut provider = provider("sub2api", ProviderKind::OpenAiBalanceHttpJson);
2924 provider.endpoint.clear();
2925
2926 let endpoint =
2927 resolve_endpoint(&provider, "https://relay.example.com/v1", "token").expect("endpoint");
2928
2929 assert_eq!(endpoint, "https://relay.example.com/user/balance");
2930 }
2931
2932 #[test]
2933 fn sub2api_usage_endpoint_defaults_to_upstream_usage_under_v1() {
2934 let mut provider = provider("sub2api", ProviderKind::Sub2ApiUsage);
2935 provider.endpoint.clear();
2936
2937 let endpoint =
2938 resolve_endpoint(&provider, "https://relay.example.com/v1", "token").expect("endpoint");
2939
2940 assert_eq!(endpoint, "https://relay.example.com/v1/usage");
2941 }
2942
2943 #[test]
2944 fn sub2api_auth_me_endpoint_defaults_to_dashboard_path_without_v1() {
2945 let mut provider = provider("sub2api-auth", ProviderKind::Sub2ApiAuthMe);
2946 provider.endpoint.clear();
2947
2948 let endpoint =
2949 resolve_endpoint(&provider, "https://relay.example.com/v1", "token").expect("endpoint");
2950
2951 assert_eq!(endpoint, "https://relay.example.com/api/v1/auth/me");
2952 }
2953
2954 #[test]
2955 fn provider_templates_support_variables_for_custom_headers_or_queries() {
2956 let mut provider = provider("newapi", ProviderKind::NewApiUserSelf);
2957 provider.endpoint = "{{base_url}}/api/user/self?user={{userId}}".to_string();
2958 provider
2959 .variables
2960 .insert("userId".to_string(), "42".to_string());
2961
2962 let endpoint = resolve_endpoint(&provider, "https://newapi.example.com/v1", "token")
2963 .expect("endpoint");
2964
2965 assert_eq!(endpoint, "https://newapi.example.com/api/user/self?user=42");
2966 }
2967
2968 #[test]
2969 fn new_api_token_usage_endpoint_defaults_to_model_key_usage_path() {
2970 let mut provider = provider("newapi-token", ProviderKind::NewApiTokenUsage);
2971 provider.endpoint.clear();
2972
2973 let endpoint = resolve_endpoint(&provider, "https://newapi.example.com/v1", "token")
2974 .expect("endpoint");
2975
2976 assert_eq!(endpoint, "https://newapi.example.com/api/usage/token/");
2977 }
2978
2979 #[test]
2980 fn openai_organization_costs_endpoint_defaults_to_official_v1_costs_window() {
2981 let mut provider = provider("openai", ProviderKind::OpenAiOrganizationCosts);
2982 provider.endpoint.clear();
2983
2984 let endpoint =
2985 resolve_endpoint(&provider, "https://api.openai.com/v1", "token").expect("endpoint");
2986
2987 assert!(endpoint.starts_with("https://api.openai.com/v1/organization/costs?start_time="));
2988 assert!(endpoint.ends_with("&limit=30"));
2989 let start_time = endpoint
2990 .split("start_time=")
2991 .nth(1)
2992 .and_then(|value| value.split('&').next())
2993 .and_then(|value| value.parse::<u64>().ok())
2994 .expect("numeric start_time");
2995 assert!(start_time > 0);
2996 }
2997
2998 #[test]
2999 fn require_token_env_prevents_upstream_model_key_fallback() {
3000 let mut cfg = proxy_config(vec![service_config(
3001 "right",
3002 vec![upstream_config("https://api.openai.com/v1")],
3003 )]);
3004 cfg.codex
3005 .configs
3006 .get_mut("right")
3007 .expect("station")
3008 .upstreams[0]
3009 .auth
3010 .auth_token = Some("model-key".to_string());
3011
3012 let mut provider = provider("openai", ProviderKind::OpenAiOrganizationCosts);
3013 provider.token_env = Some("__CODEX_HELPER_TEST_MISSING_TOKEN_ENV__".to_string());
3014 provider.require_token_env = true;
3015 let upstreams = [UpstreamRef {
3016 station_name: "right".to_string(),
3017 index: 0,
3018 provider_endpoint: None,
3019 }];
3020
3021 assert_eq!(resolve_token(&provider, &upstreams, &cfg, "codex"), None);
3022
3023 provider.require_token_env = false;
3024 assert_eq!(
3025 resolve_token(&provider, &upstreams, &cfg, "codex").as_deref(),
3026 Some("model-key")
3027 );
3028 }
3029
3030 #[test]
3031 fn effective_poll_interval_respects_disable_flag_zero_and_minimum() {
3032 let mut provider = provider("sub2api", ProviderKind::OpenAiBalanceHttpJson);
3033
3034 provider.poll_interval_secs = Some(0);
3035 assert_eq!(effective_poll_interval_secs(&provider), None);
3036
3037 provider.poll_interval_secs = Some(10);
3038 assert_eq!(
3039 effective_poll_interval_secs(&provider),
3040 Some(MIN_POLL_INTERVAL_SECS)
3041 );
3042
3043 provider.poll_interval_secs = None;
3044 assert_eq!(
3045 effective_poll_interval_secs(&provider),
3046 Some(DEFAULT_POLL_INTERVAL_SECS)
3047 );
3048
3049 provider.refresh_on_request = false;
3050 assert_eq!(effective_poll_interval_secs(&provider), None);
3051 }
3052
3053 #[test]
3054 fn auto_provider_uses_stable_target_id_across_probe_kinds() {
3055 let target = UsageProviderTarget {
3056 upstream: UpstreamRef {
3057 station_name: "input/sub".to_string(),
3058 index: 2,
3059 provider_endpoint: None,
3060 },
3061 base_url: "https://ai.input.im/v1".to_string(),
3062 provider_id: None,
3063 };
3064
3065 let sub2api = auto_usage_provider(&target, ProviderKind::Sub2ApiUsage);
3066 let newapi_token = auto_usage_provider(&target, ProviderKind::NewApiTokenUsage);
3067 let newapi = auto_usage_provider(&target, ProviderKind::NewApiUserSelf);
3068
3069 assert_eq!(sub2api.id, "auto:balance:input-sub:2");
3070 assert_eq!(sub2api.id, newapi_token.id);
3071 assert_eq!(sub2api.id, newapi.id);
3072 assert_eq!(sub2api.domains, vec!["ai.input.im".to_string()]);
3073 assert_eq!(
3074 resolve_endpoint(&sub2api, &target.base_url, "token").unwrap(),
3075 "https://ai.input.im/v1/usage"
3076 );
3077 assert_eq!(
3078 resolve_endpoint(&newapi_token, &target.base_url, "token").unwrap(),
3079 "https://ai.input.im/api/usage/token/"
3080 );
3081 }
3082
3083 #[test]
3084 fn auto_probe_prefers_rightcode_adapter_for_rightcode_hosts() {
3085 let target = UsageProviderTarget {
3086 upstream: UpstreamRef {
3087 station_name: "right".to_string(),
3088 index: 0,
3089 provider_endpoint: None,
3090 },
3091 base_url: "https://www.right.codes/codex/v1".to_string(),
3092 provider_id: Some("right".to_string()),
3093 };
3094
3095 assert_eq!(
3096 first_auto_probe_kind(&target),
3097 ProviderKind::RightCodeAccountSummary
3098 );
3099 assert_eq!(
3100 resolve_endpoint(
3101 &auto_usage_provider(&target, ProviderKind::RightCodeAccountSummary),
3102 &target.base_url,
3103 "token"
3104 )
3105 .unwrap(),
3106 "https://www.right.codes/account/summary"
3107 );
3108 assert_eq!(
3109 auto_usage_provider(&target, ProviderKind::RightCodeAccountSummary)
3110 .token_env
3111 .as_deref(),
3112 None
3113 );
3114 }
3115
3116 #[test]
3117 fn auto_provider_id_prefers_runtime_provider_tag() {
3118 let target = UsageProviderTarget {
3119 upstream: UpstreamRef {
3120 station_name: "routing".to_string(),
3121 index: 0,
3122 provider_endpoint: Some(ProviderEndpointKey::new("codex", "input", "default")),
3123 },
3124 base_url: "https://ai.input.im/v1".to_string(),
3125 provider_id: Some("input".to_string()),
3126 };
3127
3128 let provider = auto_usage_provider(&target, ProviderKind::Sub2ApiUsage);
3129
3130 assert_eq!(provider.id, "input");
3131 }
3132
3133 #[test]
3134 fn provider_endpoint_target_lookup_uses_endpoint_identity() {
3135 let cfg = proxy_config(vec![service_config(
3136 "routing",
3137 vec![
3138 endpoint_upstream_config("https://input.example/v1", "input", "default"),
3139 endpoint_upstream_config("https://right.example/v1", "right", "default"),
3140 ],
3141 )]);
3142
3143 let target = usage_provider_target_for_provider_endpoint(
3144 &cfg,
3145 "codex",
3146 &ProviderEndpointKey::new("codex", "right", "default"),
3147 )
3148 .expect("provider endpoint target");
3149
3150 assert_eq!(target.upstream.station_name, "routing");
3151 assert_eq!(target.upstream.index, 1);
3152 assert_eq!(
3153 target
3154 .upstream
3155 .provider_endpoint
3156 .as_ref()
3157 .map(ProviderEndpointKey::stable_key)
3158 .as_deref(),
3159 Some("codex/right/default")
3160 );
3161 assert_eq!(target.base_url, "https://right.example/v1");
3162 assert_eq!(target.provider_id.as_deref(), Some("right"));
3163 }
3164
3165 #[test]
3166 fn configured_target_keys_prevent_auto_probe_for_explicit_balance_domains() {
3167 let cfg = proxy_config(vec![
3168 service_config("explicit", vec![upstream_config("https://example.com/v1")]),
3169 service_config("auto", vec![upstream_config("https://ai.input.im/v1")]),
3170 ]);
3171 let configured = configured_target_keys(
3172 &cfg,
3173 "codex",
3174 &[provider("relay", ProviderKind::OpenAiBalanceHttpJson)],
3175 None,
3176 );
3177 let auto_targets = usage_provider_targets(&cfg, "codex", None)
3178 .into_iter()
3179 .filter(|target| !configured.contains(&target_key(target)))
3180 .map(|target| target.upstream.station_name)
3181 .collect::<Vec<_>>();
3182
3183 assert_eq!(auto_targets, vec!["auto".to_string()]);
3184 }
3185
3186 #[test]
3187 fn auto_probe_accepts_only_usable_balance_snapshots() {
3188 let usable = sub2api_usage_snapshot_from_json(
3189 &provider("auto", ProviderKind::Sub2ApiUsage),
3190 &upstream(),
3191 &serde_json::json!({
3192 "isValid": true,
3193 "remaining": 1
3194 }),
3195 100,
3196 Some(1_000),
3197 );
3198 let unusable = balance_http_snapshot_from_json(
3199 &provider("auto", ProviderKind::OpenAiBalanceHttpJson),
3200 &upstream(),
3201 &serde_json::json!({ "ok": true }),
3202 100,
3203 Some(1_000),
3204 );
3205
3206 assert!(auto_snapshot_is_usable(&usable));
3207 assert!(!auto_snapshot_is_usable(&unusable));
3208 }
3209
3210 #[test]
3211 fn openai_balance_snapshot_reads_common_sub2api_balance_shape() {
3212 let snapshot = balance_http_snapshot_from_json(
3213 &provider("sub2api", ProviderKind::OpenAiBalanceHttpJson),
3214 &upstream(),
3215 &serde_json::json!({
3216 "balance": "1.25"
3217 }),
3218 100,
3219 Some(1_000),
3220 );
3221
3222 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3223 assert_eq!(snapshot.exhausted, Some(false));
3224 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("1.25"));
3225 }
3226
3227 #[test]
3228 fn json_path_supports_array_indices_for_official_balance_shapes() {
3229 let value = serde_json::json!({
3230 "balance_infos": [
3231 { "currency": "CNY", "total_balance": "3.25" }
3232 ]
3233 });
3234
3235 assert_eq!(
3236 json_value_at_path(&value, "balance_infos.0.total_balance")
3237 .and_then(|value| value.as_str()),
3238 Some("3.25")
3239 );
3240 }
3241
3242 #[test]
3243 fn openai_balance_snapshot_reads_cc_switch_official_balance_shapes() {
3244 let snapshot = balance_http_snapshot_from_json(
3245 &provider("deepseek", ProviderKind::OpenAiBalanceHttpJson),
3246 &upstream(),
3247 &serde_json::json!({
3248 "balance_infos": [
3249 { "currency": "CNY", "total_balance": "3.25" }
3250 ],
3251 "is_available": true
3252 }),
3253 100,
3254 Some(1_000),
3255 );
3256
3257 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3258 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("3.25"));
3259
3260 let snapshot = balance_http_snapshot_from_json(
3261 &provider("siliconflow", ProviderKind::OpenAiBalanceHttpJson),
3262 &upstream(),
3263 &serde_json::json!({
3264 "code": 20000,
3265 "data": {
3266 "totalBalance": "8.5",
3267 "chargeBalance": "2.5"
3268 }
3269 }),
3270 100,
3271 Some(1_000),
3272 );
3273
3274 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3275 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("8.5"));
3276 assert_eq!(snapshot.paygo_balance_usd.as_deref(), Some("2.5"));
3277 }
3278
3279 #[test]
3280 fn openai_balance_snapshot_can_derive_remaining_from_total_and_used() {
3281 let mut provider = provider("openrouter", ProviderKind::OpenAiBalanceHttpJson);
3282 provider.extract.monthly_budget_paths = vec!["data.total_credits".to_string()];
3283 provider.extract.monthly_spent_paths = vec!["data.total_usage".to_string()];
3284 provider.extract.derive_remaining_from_budget_and_spent = true;
3285
3286 let snapshot = balance_http_snapshot_from_json(
3287 &provider,
3288 &upstream(),
3289 &serde_json::json!({
3290 "data": {
3291 "total_credits": "10",
3292 "total_usage": "4"
3293 }
3294 }),
3295 100,
3296 Some(1_000),
3297 );
3298
3299 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3300 assert_eq!(snapshot.exhausted, Some(false));
3301 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("6"));
3302 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("10"));
3303 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("4"));
3304 }
3305
3306 #[test]
3307 fn openai_balance_snapshot_supports_divisor_for_minor_units() {
3308 let mut provider = provider("novita", ProviderKind::OpenAiBalanceHttpJson);
3309 provider.extract.remaining_balance_paths = vec!["availableBalance".to_string()];
3310 provider.extract.remaining_divisor = Some(10_000);
3311
3312 let snapshot = balance_http_snapshot_from_json(
3313 &provider,
3314 &upstream(),
3315 &serde_json::json!({
3316 "availableBalance": 12345
3317 }),
3318 100,
3319 Some(1_000),
3320 );
3321
3322 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3323 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("1.2345"));
3324 }
3325
3326 #[test]
3327 fn sub2api_usage_snapshot_reads_all_api_hub_usage_shape() {
3328 let snapshot = sub2api_usage_snapshot_from_json(
3329 &provider("sub2api", ProviderKind::Sub2ApiUsage),
3330 &upstream(),
3331 &serde_json::json!({
3332 "isValid": true,
3333 "mode": "unrestricted",
3334 "planName": "CodeX Air",
3335 "remaining": 165.0877165,
3336 "usage": {
3337 "today": {
3338 "cost": 0,
3339 "requests": 0,
3340 "total_tokens": 0
3341 },
3342 "total": {
3343 "cost": 354.194748,
3344 "requests": 2691,
3345 "total_tokens": 384084697
3346 }
3347 }
3348 }),
3349 100,
3350 Some(1_000),
3351 );
3352
3353 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3354 assert_eq!(snapshot.exhausted, Some(false));
3355 assert_eq!(snapshot.plan_name.as_deref(), Some("CodeX Air"));
3356 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("165.0877165"));
3357 assert_eq!(snapshot.total_used_usd.as_deref(), Some("354.194748"));
3358 assert_eq!(snapshot.today_used_usd.as_deref(), Some("0"));
3359 assert_eq!(snapshot.total_requests, Some(2691));
3360 assert_eq!(snapshot.today_requests, Some(0));
3361 assert_eq!(snapshot.total_tokens, Some(384084697));
3362 assert_eq!(snapshot.today_tokens, Some(0));
3363 }
3364
3365 #[test]
3366 fn sub2api_subscription_zero_remaining_is_display_only_period_capacity_exhaustion() {
3367 let snapshot = sub2api_usage_snapshot_from_json(
3368 &provider("sub2api", ProviderKind::Sub2ApiUsage),
3369 &upstream(),
3370 &serde_json::json!({
3371 "isValid": true,
3372 "mode": "unrestricted",
3373 "planName": "CodeX Lite 年度",
3374 "remaining": 0,
3375 "subscription": {
3376 "daily_usage_usd": 100.468025,
3377 "daily_limit_usd": 100,
3378 "weekly_usage_usd": 401.441684,
3379 "weekly_limit_usd": 0,
3380 "monthly_usage_usd": 401.441684,
3381 "monthly_limit_usd": 0
3382 },
3383 "usage": {
3384 "today": { "cost": 0, "requests": 0, "total_tokens": 0 },
3385 "total": { "cost": 702.492098, "requests": 42, "total_tokens": 1234 }
3386 }
3387 }),
3388 100,
3389 Some(1_000),
3390 );
3391
3392 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
3393 assert_eq!(snapshot.exhausted, Some(true));
3394 assert_eq!(snapshot.plan_name.as_deref(), Some("CodeX Lite 年度"));
3395 assert_eq!(snapshot.total_balance_usd, None);
3396 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
3397 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
3398 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("100"));
3399 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("100.468025"));
3400 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("100"));
3401 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("100.468025"));
3402 assert_eq!(snapshot.total_used_usd.as_deref(), Some("702.492098"));
3403 assert_eq!(snapshot.today_used_usd.as_deref(), Some("0"));
3404 assert!(
3405 !snapshot.routing_exhausted(),
3406 "sub2api /v1/usage skips billing checks; subscription windows are reset lazily on real requests"
3407 );
3408 }
3409
3410 #[test]
3411 fn sub2api_quota_limited_zero_remaining_still_marks_exhausted() {
3412 let snapshot = sub2api_usage_snapshot_from_json(
3413 &provider("sub2api", ProviderKind::Sub2ApiUsage),
3414 &upstream(),
3415 &serde_json::json!({
3416 "isValid": true,
3417 "mode": "quota_limited",
3418 "quota": {
3419 "limit": 10,
3420 "used": 10,
3421 "remaining": 0,
3422 "unit": "USD"
3423 }
3424 }),
3425 100,
3426 Some(1_000),
3427 );
3428
3429 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
3430 assert_eq!(snapshot.exhausted, Some(true));
3431 assert_eq!(snapshot.total_balance_usd, None);
3432 assert_eq!(snapshot.quota_period.as_deref(), Some("quota"));
3433 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
3434 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("10"));
3435 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("10"));
3436 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("10"));
3437 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("10"));
3438 assert!(snapshot.routing_exhausted());
3439 }
3440
3441 #[test]
3442 fn sub2api_usage_snapshot_marks_invalid_key_as_error() {
3443 let snapshot = sub2api_usage_snapshot_from_json(
3444 &provider("sub2api", ProviderKind::Sub2ApiUsage),
3445 &upstream(),
3446 &serde_json::json!({
3447 "isValid": false
3448 }),
3449 100,
3450 Some(1_000),
3451 );
3452
3453 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
3454 assert_eq!(
3455 snapshot.error.as_deref(),
3456 Some("sub2api usage response reported invalid API key")
3457 );
3458 }
3459
3460 #[test]
3461 fn sub2api_auth_me_snapshot_reads_dashboard_balance_envelope() {
3462 let snapshot = sub2api_auth_me_snapshot_from_json(
3463 &provider("sub2api-auth", ProviderKind::Sub2ApiAuthMe),
3464 &upstream(),
3465 &serde_json::json!({
3466 "code": 0,
3467 "message": "ok",
3468 "data": {
3469 "id": 42,
3470 "username": "demo",
3471 "balance": "12.5"
3472 }
3473 }),
3474 100,
3475 Some(1_000),
3476 );
3477
3478 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3479 assert_eq!(snapshot.exhausted, Some(false));
3480 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("12.5"));
3481 }
3482
3483 #[test]
3484 fn rightcode_endpoint_defaults_to_account_summary() {
3485 let mut provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
3486 provider.endpoint.clear();
3487
3488 let endpoint = resolve_endpoint(&provider, "https://www.right.codes/codex/v1", "token")
3489 .expect("endpoint");
3490
3491 assert_eq!(endpoint, "https://www.right.codes/account/summary");
3492 }
3493
3494 #[test]
3495 fn rightcode_account_summary_reads_matching_subscription_and_balance() {
3496 let mut provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
3497 provider.trust_exhaustion_for_routing = false;
3498
3499 let snapshot = rightcode_account_summary_snapshot_from_json(
3500 &provider,
3501 &upstream(),
3502 &serde_json::json!({
3503 "balance": 3.25,
3504 "subscriptions": [
3505 {
3506 "name": "Daily",
3507 "total_quota": 20,
3508 "remaining_quota": 7.5,
3509 "reset_today": true,
3510 "available_prefixes": ["/codex"]
3511 },
3512 {
3513 "name": "Other",
3514 "total_quota": 99,
3515 "remaining_quota": 99,
3516 "reset_today": true,
3517 "available_prefixes": ["/claude"]
3518 },
3519 {
3520 "name": "Badge",
3521 "total_quota": 5,
3522 "remaining_quota": 5,
3523 "reset_today": true,
3524 "available_prefixes": ["/codex"]
3525 }
3526 ]
3527 }),
3528 "https://www.right.codes/codex/v1",
3529 100,
3530 Some(1_000),
3531 );
3532
3533 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3534 assert_eq!(snapshot.exhausted, Some(false));
3535 assert!(!snapshot.routing_exhausted());
3536 assert_eq!(snapshot.plan_name.as_deref(), Some("Daily"));
3537 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("3.25"));
3538 assert_eq!(snapshot.paygo_balance_usd, None);
3539 assert_eq!(snapshot.quota_period.as_deref(), Some("daily"));
3540 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("7.5"));
3541 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("20"));
3542 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("12.5"));
3543 }
3544
3545 #[test]
3546 fn rightcode_account_summary_accounts_for_not_reset_today() {
3547 let provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
3548
3549 let snapshot = rightcode_account_summary_snapshot_from_json(
3550 &provider,
3551 &upstream(),
3552 &serde_json::json!({
3553 "subscriptions": [
3554 {
3555 "name": "Daily",
3556 "total_quota": 20,
3557 "remaining_quota": 7.5,
3558 "reset_today": false,
3559 "available_prefixes": ["/codex"]
3560 }
3561 ]
3562 }),
3563 "https://www.right.codes/codex/v1",
3564 100,
3565 Some(1_000),
3566 );
3567
3568 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3569 assert_eq!(snapshot.exhausted, Some(false));
3570 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("27.5"));
3571 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("20"));
3572 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("0"));
3573 }
3574
3575 #[test]
3576 fn rightcode_zero_daily_quota_without_balance_is_display_only_exhaustion_by_default() {
3577 let mut provider = provider("rightcode", ProviderKind::RightCodeAccountSummary);
3578 provider.trust_exhaustion_for_routing = false;
3579
3580 let snapshot = rightcode_account_summary_snapshot_from_json(
3581 &provider,
3582 &upstream(),
3583 &serde_json::json!({
3584 "balance": 0,
3585 "subscriptions": [
3586 {
3587 "name": "Daily",
3588 "total_quota": 20,
3589 "remaining_quota": 0,
3590 "reset_today": true,
3591 "available_prefixes": ["/codex"]
3592 }
3593 ]
3594 }),
3595 "https://www.right.codes/codex/v1",
3596 100,
3597 Some(1_000),
3598 );
3599
3600 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
3601 assert_eq!(snapshot.exhausted, Some(true));
3602 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("0"));
3603 assert!(!snapshot.routing_exhausted());
3604 assert!(snapshot.routing_ignored_exhaustion());
3605 }
3606
3607 #[test]
3608 fn sub2api_auth_me_snapshot_marks_business_error() {
3609 let snapshot = sub2api_auth_me_snapshot_from_json(
3610 &provider("sub2api-auth", ProviderKind::Sub2ApiAuthMe),
3611 &upstream(),
3612 &serde_json::json!({
3613 "code": 401,
3614 "message": "login required"
3615 }),
3616 100,
3617 Some(1_000),
3618 );
3619
3620 assert_eq!(snapshot.status, BalanceSnapshotStatus::Error);
3621 assert_eq!(snapshot.error.as_deref(), Some("login required"));
3622 }
3623
3624 #[test]
3625 fn provider_can_disable_routing_trust_for_exhausted_balance() {
3626 let mut provider = provider("sub2api", ProviderKind::OpenAiBalanceHttpJson);
3627 provider.trust_exhaustion_for_routing = false;
3628
3629 let snapshot = balance_http_snapshot_from_json(
3630 &provider,
3631 &upstream(),
3632 &serde_json::json!({
3633 "balance": "0"
3634 }),
3635 100,
3636 Some(1_000),
3637 );
3638
3639 assert_eq!(snapshot.status, BalanceSnapshotStatus::Exhausted);
3640 assert_eq!(snapshot.exhausted, Some(true));
3641 assert!(!snapshot.exhaustion_affects_routing);
3642 assert!(!snapshot.routing_exhausted());
3643 }
3644
3645 #[test]
3646 fn provider_exhaustion_trust_defaults_to_enabled_when_omitted() {
3647 let provider: UsageProviderConfig = serde_json::from_value(serde_json::json!({
3648 "id": "sub2api",
3649 "kind": "openai_balance_http_json",
3650 "domains": ["example.com"]
3651 }))
3652 .expect("provider config");
3653
3654 assert!(provider.trust_exhaustion_for_routing);
3655 }
3656
3657 #[tokio::test]
3658 async fn provider_missing_token_clears_stale_lb_exhaustion_marker() {
3659 let cfg = proxy_config(vec![service_config(
3660 "right",
3661 vec![
3662 upstream_config("https://primary.example/v1"),
3663 upstream_config("https://backup.example/v1"),
3664 ],
3665 )]);
3666 let lb_states = Arc::new(Mutex::new(HashMap::new()));
3667 let target = UsageProviderTarget {
3668 upstream: upstream(),
3669 base_url: "https://backup.example/v1".to_string(),
3670 provider_id: Some("right".to_string()),
3671 };
3672 let upstreams = vec![target.upstream.clone()];
3673 let state = ProxyState::new();
3674 update_usage_exhausted(&lb_states, &state, &cfg, "codex", &upstreams, true).await;
3675 {
3676 let guard = lb_states.lock().expect("lb states");
3677 assert!(
3678 guard
3679 .get("right")
3680 .and_then(|entry| entry.usage_exhausted.get(1))
3681 .copied()
3682 .unwrap_or(false)
3683 );
3684 }
3685
3686 let outcome = refresh_provider_target(RefreshProviderTargetParams {
3687 client: &Client::new(),
3688 provider: &provider("sub2api", ProviderKind::OpenAiBalanceHttpJson),
3689 target: &target,
3690 cfg: &cfg,
3691 lb_states: &lb_states,
3692 state: &state,
3693 service_name: "codex",
3694 interval_secs: 60,
3695 })
3696 .await;
3697
3698 assert_eq!(outcome, UsageProviderRefreshOutcome::MissingToken);
3699 let guard = lb_states.lock().expect("lb states");
3700 assert!(
3701 !guard
3702 .get("right")
3703 .and_then(|entry| entry.usage_exhausted.get(1))
3704 .copied()
3705 .unwrap_or(true)
3706 );
3707 }
3708
3709 #[test]
3710 fn openai_balance_snapshot_supports_custom_paths_and_derived_budget() {
3711 let mut provider = provider("custom", ProviderKind::OpenAiBalanceHttpJson);
3712 provider.extract.remaining_balance_paths = vec!["payload.remaining_usd".to_string()];
3713 provider.extract.monthly_spent_paths = vec!["payload.used_usd".to_string()];
3714 provider.extract.derive_budget_from_remaining_and_spent = true;
3715
3716 let snapshot = balance_http_snapshot_from_json(
3717 &provider,
3718 &upstream(),
3719 &serde_json::json!({
3720 "payload": {
3721 "remaining_usd": "2",
3722 "used_usd": "0.5"
3723 }
3724 }),
3725 100,
3726 Some(1_000),
3727 );
3728
3729 assert_eq!(snapshot.total_balance_usd.as_deref(), Some("2"));
3730 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
3731 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("2.5"));
3732 assert_eq!(snapshot.exhausted, Some(false));
3733 }
3734
3735 #[test]
3736 fn new_api_snapshot_converts_quota_units_like_cc_switch_template() {
3737 let snapshot = new_api_snapshot_from_json(
3738 &provider("newapi", ProviderKind::NewApiUserSelf),
3739 &upstream(),
3740 &serde_json::json!({
3741 "success": true,
3742 "data": {
3743 "quota": 500000,
3744 "used_quota": 250000
3745 }
3746 }),
3747 100,
3748 Some(1_000),
3749 );
3750
3751 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3752 assert_eq!(snapshot.exhausted, Some(false));
3753 assert_eq!(snapshot.total_balance_usd, None);
3754 assert_eq!(snapshot.quota_period.as_deref(), Some("quota"));
3755 assert_eq!(snapshot.quota_remaining_usd.as_deref(), Some("1"));
3756 assert_eq!(snapshot.quota_limit_usd.as_deref(), Some("1.5"));
3757 assert_eq!(snapshot.quota_used_usd.as_deref(), Some("0.5"));
3758 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
3759 assert_eq!(snapshot.monthly_budget_usd.as_deref(), Some("1.5"));
3760 }
3761
3762 #[test]
3763 fn new_api_user_self_honors_unlimited_quota_flag() {
3764 let snapshot = new_api_snapshot_from_json(
3765 &provider("newapi", ProviderKind::NewApiUserSelf),
3766 &upstream(),
3767 &serde_json::json!({
3768 "success": true,
3769 "data": {
3770 "quota": 0,
3771 "used_quota": 250000,
3772 "unlimited_quota": true
3773 }
3774 }),
3775 100,
3776 Some(1_000),
3777 );
3778
3779 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3780 assert_eq!(snapshot.exhausted, Some(false));
3781 assert_eq!(snapshot.total_balance_usd, None);
3782 assert_eq!(snapshot.monthly_budget_usd, None);
3783 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
3784 assert_eq!(snapshot.unlimited_quota, Some(true));
3785 }
3786
3787 #[test]
3788 fn new_api_token_usage_honors_unlimited_quota_flag() {
3789 let snapshot = new_api_token_usage_snapshot_from_json(
3790 &provider("newapi-token", ProviderKind::NewApiTokenUsage),
3791 &upstream(),
3792 &serde_json::json!({
3793 "code": true,
3794 "message": "ok",
3795 "data": {
3796 "object": "token_usage",
3797 "name": "demo-token",
3798 "total_granted": 0,
3799 "total_used": 250000,
3800 "total_available": 0,
3801 "unlimited_quota": true
3802 }
3803 }),
3804 100,
3805 Some(1_000),
3806 );
3807
3808 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3809 assert_eq!(snapshot.exhausted, Some(false));
3810 assert_eq!(snapshot.plan_name.as_deref(), Some("demo-token"));
3811 assert_eq!(snapshot.total_balance_usd, None);
3812 assert_eq!(snapshot.monthly_budget_usd, None);
3813 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("0.5"));
3814 assert_eq!(snapshot.unlimited_quota, Some(true));
3815 }
3816
3817 #[test]
3818 fn openai_organization_costs_sums_official_cost_buckets_without_exhaustion() {
3819 let snapshot = openai_organization_costs_snapshot_from_json(
3820 &provider("openai", ProviderKind::OpenAiOrganizationCosts),
3821 &upstream(),
3822 &serde_json::json!({
3823 "object": "page",
3824 "data": [
3825 {
3826 "object": "bucket",
3827 "start_time": 1710000000,
3828 "end_time": 1710086400,
3829 "results": [
3830 {
3831 "object": "organization.costs.result",
3832 "amount": { "value": 1.25, "currency": "usd" }
3833 },
3834 {
3835 "object": "organization.costs.result",
3836 "amount": { "value": "2.5", "currency": "usd" }
3837 },
3838 {
3839 "object": "organization.costs.result",
3840 "amount": { "value": 99, "currency": "eur" }
3841 }
3842 ]
3843 },
3844 {
3845 "object": "bucket",
3846 "results": [
3847 {
3848 "object": "organization.costs.result",
3849 "amount": { "value": "0.25", "currency": "USD" }
3850 }
3851 ]
3852 }
3853 ],
3854 "has_more": false
3855 }),
3856 100,
3857 Some(1_000),
3858 );
3859
3860 assert_eq!(snapshot.status, BalanceSnapshotStatus::Ok);
3861 assert_eq!(snapshot.exhausted, None);
3862 assert!(!snapshot.exhaustion_affects_routing);
3863 assert!(!snapshot.routing_exhausted());
3864 assert_eq!(snapshot.monthly_spent_usd.as_deref(), Some("4"));
3865 assert_eq!(snapshot.total_balance_usd, None);
3866 }
3867}