Skip to main content

codex_helper_core/
usage_providers.rs

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    /// 简单预算接口,返回 total/used,判断是否用尽
22    BudgetHttpJson,
23    /// YesCode 账户用量,基于 /api/v1/auth/profile 返回的余额信息
24    YescodeProfile,
25    /// OpenAI-compatible relay balance endpoint, defaulting to /user/balance.
26    #[serde(
27        rename = "openai_balance_http_json",
28        alias = "open_ai_balance_http_json",
29        alias = "relay_balance_http_json"
30    )]
31    OpenAiBalanceHttpJson,
32    /// Sub2API API-key telemetry endpoint, defaulting to /v1/usage.
33    #[serde(rename = "sub2api_usage", alias = "sub2api_usage_http_json")]
34    Sub2ApiUsage,
35    /// Sub2API dashboard JWT account endpoint, defaulting to /api/v1/auth/me.
36    #[serde(rename = "sub2api_auth_me", alias = "sub2api_auth_me_http_json")]
37    Sub2ApiAuthMe,
38    /// New API-style model token quota endpoint, defaulting to /api/usage/token/.
39    #[serde(
40        rename = "new_api_token_usage",
41        alias = "new_api_token_usage_http_json"
42    )]
43    NewApiTokenUsage,
44    /// New API-style user quota endpoint, defaulting to /api/user/self.
45    NewApiUserSelf,
46    /// RightCode account summary endpoint, defaulting to /account/summary.
47    #[serde(
48        rename = "rightcode_account_summary",
49        alias = "right_code_account_summary",
50        alias = "rightcode"
51    )]
52    RightCodeAccountSummary,
53    /// OpenAI official organization Costs API, defaulting to a rolling 30-day cost window.
54    #[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
232// 全局节流状态:按 provider.id 记录最近一次查询时间,避免高频请求。
233static LAST_USAGE_POLL: OnceLock<Mutex<HashMap<String, Instant>>> = OnceLock::new();
234
235const DEFAULT_POLL_INTERVAL_SECS: u64 = 60;
236// Minimal poll interval per provider to avoid hammering usage APIs.
237const 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    // RightCode subscription windows are daily capacity signals. A zero daily
355    // remainder can coexist with account balance or be reset lazily, so the
356    // built-in adapter displays it without demoting routes by default.
357    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                // Match co.yes.vg, cotest.yes.vg, and sibling subdomains.
495                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    // 写入默认配置,方便用户查看/修改。
548    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    // 优先: token_env 环境变量
743    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    // 否则: 使用绑定 upstream 的 auth_token(当前 Codex 正在使用的 token)
755    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    // Tests should be hermetic and must not depend on real user `usage_providers.json`.
2534    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
2656/// 在特定 Codex upstream 请求结束后,按需查询一次用量并更新 LB 状态。
2657/// 设计为轻量的“按需刷新”,而非后台定时轮询。
2658pub 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
2672/// Provider-endpoint keyed variant used by the route graph executor.
2673/// Station/upstream are still updated inside the usage provider as a compatibility projection
2674/// when the current runtime config can map the endpoint back to one legacy upstream.
2675pub 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    // Tests should be hermetic and should not depend on any real user `usage_providers.json` on
2697    // the machine running the suite. Disable provider polling during tests to avoid flakiness.
2698    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(&current_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(&current_target.base_url) {
2760        auto_openai_official_provider(&current_target)
2761    } else {
2762        auto_usage_provider(&current_target, first_auto_probe_kind(&current_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        &current_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}