Skip to main content

agent_quota_core/
lib.rs

1#![warn(missing_docs)]
2//! Local-credential quota probes for AI coding agents.
3//!
4//! Agent Quota reads locally configured accounts, performs the smallest
5//! provider-specific check available, and returns a versioned, provider-neutral
6//! snapshot. A successful probe and usable quota are represented separately so
7//! consumers do not need to infer routing decisions from transport health.
8
9mod providers;
10
11use std::collections::{BTreeMap, BTreeSet};
12use std::path::{Path, PathBuf};
13use std::sync::Arc;
14use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
15use std::{fs, process::Stdio};
16
17use reqwest::header::HeaderMap;
18use reqwest::redirect::Policy;
19use reqwest::StatusCode;
20use serde::{Deserialize, Serialize};
21use serde_json::{json, Map, Value};
22use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
23use tokio::process::{ChildStdin, Command};
24use tokio::sync::Mutex;
25use tokio::time::timeout;
26
27const DEFAULT_CODEX_TIMEOUT: Duration = Duration::from_secs(10);
28const DEFAULT_CLAUDE_TIMEOUT: Duration = Duration::from_secs(20);
29/// Default lifetime of a successful cached provider snapshot.
30pub const DEFAULT_CACHE_TTL_SECONDS: u64 = 5 * 60;
31const DEFAULT_CACHE_TTL: Duration = Duration::from_secs(DEFAULT_CACHE_TTL_SECONDS);
32const CLAUDE_USAGE_MODEL: &str = "claude-haiku-4-5-20251001";
33const CLAUDE_API_URL: &str = "https://api.anthropic.com/v1/messages";
34const CLAUDE_CREDENTIALS_RELATIVE_PATH: &str = ".claude/.credentials.json";
35
36#[cfg(target_os = "windows")]
37const CREATE_NO_WINDOW: u32 = 0x0800_0000;
38
39/// Current serialized snapshot schema.
40pub const SNAPSHOT_SCHEMA_VERSION: u32 = 1;
41/// Current serialized provider-capabilities schema.
42pub const CAPABILITIES_SCHEMA_VERSION: u32 = 1;
43
44/// Result type used for configuration operations.
45pub type AgentQuotaResult<T> = Result<T, AgentQuotaError>;
46
47/// Configuration loading or validation error.
48#[derive(Debug)]
49pub enum AgentQuotaError {
50    /// The configuration file could not be read.
51    ConfigRead {
52        /// Path that could not be read.
53        path: PathBuf,
54        /// Underlying filesystem error.
55        source: std::io::Error,
56    },
57    /// The configuration file is not valid TOML or contains unknown fields.
58    ConfigParse {
59        /// Path containing invalid configuration.
60        path: PathBuf,
61        /// Parser diagnostic.
62        message: String,
63    },
64    /// The configuration is syntactically valid but semantically unsafe.
65    ConfigValidation {
66        /// Human-readable validation diagnostic.
67        message: String,
68    },
69}
70
71impl std::fmt::Display for AgentQuotaError {
72    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
73        match self {
74            Self::ConfigRead { path, source } => {
75                write!(formatter, "could not read {}: {source}", path.display())
76            }
77            Self::ConfigParse { path, message } => {
78                write!(
79                    formatter,
80                    "invalid configuration in {}: {message}",
81                    path.display()
82                )
83            }
84            Self::ConfigValidation { message } => formatter.write_str(message),
85        }
86    }
87}
88
89impl std::error::Error for AgentQuotaError {
90    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
91        match self {
92            Self::ConfigRead { source, .. } => Some(source),
93            Self::ConfigParse { .. } | Self::ConfigValidation { .. } => None,
94        }
95    }
96}
97
98/// Supported quota provider.
99#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
100#[serde(rename_all = "snake_case")]
101pub enum ProviderId {
102    /// OpenAI Codex CLI.
103    Codex,
104    /// Anthropic Claude Code credentials.
105    Claude,
106}
107
108impl ProviderId {
109    /// Stable provider identifier.
110    pub fn as_str(self) -> &'static str {
111        match self {
112            Self::Codex => "codex",
113            Self::Claude => "claude",
114        }
115    }
116
117    /// Human-readable provider name.
118    pub fn label(self) -> &'static str {
119        match self {
120            Self::Codex => "Codex",
121            Self::Claude => "Claude Code",
122        }
123    }
124
125    /// Parse a provider name or common alias.
126    pub fn from_name(value: &str) -> Option<Self> {
127        match value.trim().to_ascii_lowercase().as_str() {
128            "codex" | "openai-codex" => Some(Self::Codex),
129            "claude" | "claude-code" | "anthropic" => Some(Self::Claude),
130            _ => None,
131        }
132    }
133}
134
135/// Transport used by a provider quota probe.
136#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
137#[serde(rename_all = "snake_case")]
138pub enum ProviderProbeTransport {
139    /// A local provider process exposes the quota interface.
140    LocalProcess,
141    /// Agent Quota calls a remote provider API directly.
142    RemoteApi,
143}
144
145/// Machine-readable source of normalized quota data.
146#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
147#[serde(rename_all = "snake_case")]
148pub enum ProviderQuotaSource {
149    /// Codex app-server account and rate-limit RPC methods.
150    CodexAppServer,
151    /// Anthropic rate-limit response headers.
152    AnthropicRateLimitHeaders,
153}
154
155/// Local credential source required by a provider probe.
156#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
157#[serde(rename_all = "snake_case")]
158pub enum ProviderCredentialSource {
159    /// The signed-in Codex CLI session.
160    CodexCliSession,
161    /// The Claude Code OAuth credential file.
162    ClaudeCodeOauthFile,
163}
164
165/// Stable capability description for one supported provider.
166#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
167#[serde(rename_all = "camelCase")]
168pub struct ProviderCapability {
169    /// Stable provider identifier.
170    pub provider_id: ProviderId,
171    /// Human-readable provider name.
172    pub provider_name: String,
173    /// Transport used to obtain quota data.
174    pub probe_transport: ProviderProbeTransport,
175    /// Machine-readable quota source.
176    pub quota_source: ProviderQuotaSource,
177    /// Local credential source read by the probe.
178    pub credential_source: ProviderCredentialSource,
179    /// Whether a quota probe submits a message to the provider.
180    pub submits_message: bool,
181    /// Whether a quota probe may affect provider quota or billing.
182    pub may_affect_quota_or_billing: bool,
183    /// Default successful-result cache lifetime.
184    pub default_cache_ttl_seconds: u64,
185    /// Concise human-readable probe impact.
186    pub probe_impact: String,
187}
188
189/// Versioned capability document for discovery by tools and user interfaces.
190#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
191#[serde(rename_all = "camelCase")]
192pub struct ProviderCapabilitiesDocument {
193    /// Capability document schema version.
194    pub schema_version: u32,
195    /// Snapshot schema produced by this Agent Quota version.
196    pub snapshot_schema_version: u32,
197    /// Capabilities for every supported provider.
198    pub providers: Vec<ProviderCapability>,
199}
200
201/// Return the stable capability description for one provider.
202pub fn provider_capability(provider: ProviderId) -> ProviderCapability {
203    providers::adapter(provider).capability()
204}
205
206/// Return a versioned document describing every supported provider.
207pub fn provider_capabilities() -> ProviderCapabilitiesDocument {
208    ProviderCapabilitiesDocument {
209        schema_version: CAPABILITIES_SCHEMA_VERSION,
210        snapshot_schema_version: SNAPSHOT_SCHEMA_VERSION,
211        providers: vec![
212            provider_capability(ProviderId::Codex),
213            provider_capability(ProviderId::Claude),
214        ],
215    }
216}
217/// One locally configured provider account.
218#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
219#[serde(rename_all = "snake_case", deny_unknown_fields)]
220pub struct ProviderProfile {
221    /// Stable, non-empty profile identifier.
222    pub id: String,
223    /// Provider used by this profile.
224    pub provider: ProviderId,
225    /// Optional human-readable label.
226    pub label: Option<String>,
227    /// Optional Codex executable path.
228    pub command_path: Option<PathBuf>,
229    /// Optional Claude Code credentials file.
230    pub credentials_path: Option<PathBuf>,
231    /// Codex process environment overrides.
232    #[serde(default)]
233    pub env: BTreeMap<String, String>,
234    /// Whether this profile participates in configured collection.
235    #[serde(default = "default_enabled")]
236    pub enabled: bool,
237}
238
239impl ProviderProfile {
240    /// Construct the conventional local profile for a provider.
241    pub fn default_for_provider(provider: ProviderId) -> Self {
242        Self {
243            id: provider.as_str().to_owned(),
244            provider,
245            label: None,
246            command_path: None,
247            credentials_path: None,
248            env: BTreeMap::new(),
249            enabled: true,
250        }
251    }
252
253    /// Human-readable profile name.
254    pub fn display_name(&self) -> String {
255        self.label
256            .as_deref()
257            .filter(|value| !value.trim().is_empty())
258            .map(ToOwned::to_owned)
259            .unwrap_or_else(|| self.provider.label().to_owned())
260    }
261
262    /// Validate provider-specific and identity fields.
263    pub fn validate(&self) -> AgentQuotaResult<()> {
264        if self.id.trim().is_empty() {
265            return Err(validation_error("profile id must not be empty"));
266        }
267        if self.id.trim() != self.id {
268            return Err(validation_error(format!(
269                "profile id `{}` must not start or end with whitespace",
270                self.id
271            )));
272        }
273        if self
274            .label
275            .as_ref()
276            .is_some_and(|label| label.trim().is_empty())
277        {
278            return Err(validation_error(format!(
279                "profile `{}` has an empty label; omit `label` to use the provider name",
280                self.id
281            )));
282        }
283
284        match self.provider {
285            ProviderId::Codex if self.credentials_path.is_some() => {
286                return Err(validation_error(format!(
287                    "profile `{}` is Codex, so `credentials_path` is not supported",
288                    self.id
289                )));
290            }
291            ProviderId::Claude if self.command_path.is_some() => {
292                return Err(validation_error(format!(
293                    "profile `{}` is Claude, so `command_path` is not supported",
294                    self.id
295                )));
296            }
297            ProviderId::Claude if !self.env.is_empty() => {
298                return Err(validation_error(format!(
299                    "profile `{}` is Claude, so process environment overrides are not supported",
300                    self.id
301                )));
302            }
303            ProviderId::Codex | ProviderId::Claude => {}
304        }
305        Ok(())
306    }
307}
308
309fn default_enabled() -> bool {
310    true
311}
312
313/// Top-level TOML configuration.
314#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
315#[serde(rename_all = "snake_case", deny_unknown_fields)]
316pub struct AgentQuotaConfig {
317    /// Configured provider accounts.
318    #[serde(default)]
319    pub profiles: Vec<ProviderProfile>,
320}
321
322impl AgentQuotaConfig {
323    /// Load and validate a TOML configuration file.
324    pub fn load(path: impl AsRef<Path>) -> AgentQuotaResult<Self> {
325        let path = path.as_ref();
326        let raw = fs::read_to_string(path).map_err(|source| AgentQuotaError::ConfigRead {
327            path: path.to_path_buf(),
328            source,
329        })?;
330        let config =
331            toml::from_str::<Self>(&raw).map_err(|error| AgentQuotaError::ConfigParse {
332                path: path.to_path_buf(),
333                message: error.to_string(),
334            })?;
335        config.validate()?;
336        Ok(config)
337    }
338
339    /// Validate profile identities and provider-specific fields.
340    pub fn validate(&self) -> AgentQuotaResult<()> {
341        let mut ids = BTreeSet::new();
342        for profile in &self.profiles {
343            profile.validate()?;
344            if !ids.insert(&profile.id) {
345                return Err(validation_error(format!(
346                    "profile id `{}` is defined more than once",
347                    profile.id
348                )));
349            }
350        }
351        Ok(())
352    }
353
354    /// Return enabled profiles in configuration order.
355    pub fn profiles(&self) -> Vec<ProviderProfile> {
356        self.profiles
357            .iter()
358            .filter(|profile| profile.enabled)
359            .cloned()
360            .collect()
361    }
362}
363
364fn validation_error(message: impl Into<String>) -> AgentQuotaError {
365    AgentQuotaError::ConfigValidation {
366        message: message.into(),
367    }
368}
369
370/// Explicit set of accounts to collect.
371#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
372enum UsageSelection {
373    All,
374    Providers(Vec<ProviderId>),
375    Profiles(Vec<ProviderProfile>),
376}
377
378/// Options identifying exactly which provider accounts to collect.
379#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
380pub struct CollectUsageOptions {
381    selection: UsageSelection,
382}
383
384impl Default for CollectUsageOptions {
385    fn default() -> Self {
386        Self::all()
387    }
388}
389
390impl CollectUsageOptions {
391    /// Collect conventional local profiles for every supported provider.
392    pub fn all() -> Self {
393        Self {
394            selection: UsageSelection::All,
395        }
396    }
397
398    /// Collect conventional profiles for exactly the supplied providers.
399    ///
400    /// An empty iterator performs no probes.
401    pub fn providers(providers: impl IntoIterator<Item = ProviderId>) -> Self {
402        Self {
403            selection: UsageSelection::Providers(providers.into_iter().collect()),
404        }
405    }
406
407    /// Collect exactly the supplied enabled profiles.
408    ///
409    /// An empty iterator performs no probes.
410    pub fn profiles(profiles: impl IntoIterator<Item = ProviderProfile>) -> Self {
411        Self {
412            selection: UsageSelection::Profiles(profiles.into_iter().collect()),
413        }
414    }
415
416    fn selected_profiles(&self) -> Vec<ProviderProfile> {
417        match &self.selection {
418            UsageSelection::All => vec![
419                ProviderProfile::default_for_provider(ProviderId::Codex),
420                ProviderProfile::default_for_provider(ProviderId::Claude),
421            ],
422            UsageSelection::Providers(providers) => providers
423                .iter()
424                .copied()
425                .map(ProviderProfile::default_for_provider)
426                .collect(),
427            UsageSelection::Profiles(profiles) => profiles
428                .iter()
429                .filter(|profile| profile.enabled)
430                .cloned()
431                .collect(),
432        }
433    }
434}
435
436/// Health of the provider probe itself.
437#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
438#[serde(rename_all = "snake_case")]
439pub enum ProbeStatus {
440    /// Provider returned a recognizable quota response.
441    Ok,
442    /// Local credentials are missing or rejected.
443    AuthenticationRequired,
444    /// Installed provider version does not expose compatible quota data.
445    Unsupported,
446    /// A retryable process, network, or provider failure occurred.
447    TransientError,
448    /// Provider returned a response that could not be interpreted safely.
449    InvalidResponse,
450    /// The supplied profile is invalid.
451    InvalidConfiguration,
452}
453
454/// Whether normalized quota permits another agent run.
455#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
456#[serde(rename_all = "snake_case")]
457pub enum QuotaState {
458    /// Every reported blocking window has remaining capacity.
459    Available,
460    /// At least one reported blocking window is exhausted.
461    Exhausted,
462    /// Capacity cannot be determined.
463    Unknown,
464}
465
466/// Policy used to evaluate whether selected profiles permit more work.
467#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
468#[serde(rename_all = "snake_case")]
469pub enum ReadinessPolicy {
470    /// At least one selected profile must be ready.
471    Any,
472    /// Every selected profile must be ready.
473    All,
474}
475
476/// Typed result of evaluating readiness across provider snapshots.
477#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
478#[serde(rename_all = "camelCase")]
479pub struct ReadinessSummary {
480    /// Policy used for this evaluation.
481    pub policy: ReadinessPolicy,
482    /// Whether the selected policy is satisfied.
483    pub satisfied: bool,
484    /// Total number of evaluated profiles.
485    pub total_profiles: usize,
486    /// Profiles with a successful probe and available quota.
487    pub ready_profiles: usize,
488    /// Profiles that are not safe positive routing signals.
489    pub not_ready_profiles: usize,
490    /// Profiles whose provider probe did not complete successfully.
491    pub failed_probes: usize,
492    /// Profiles that successfully reported exhausted quota.
493    pub exhausted_profiles: usize,
494}
495
496impl ReadinessSummary {
497    /// Evaluate provider snapshots under an `any` or `all` readiness policy.
498    pub fn evaluate(snapshots: &[ProviderUsageSnapshot], policy: ReadinessPolicy) -> Self {
499        let total_profiles = snapshots.len();
500        let ready_profiles = snapshots
501            .iter()
502            .filter(|snapshot| snapshot.is_ready())
503            .count();
504        let failed_probes = snapshots
505            .iter()
506            .filter(|snapshot| snapshot.probe_status != ProbeStatus::Ok)
507            .count();
508        let exhausted_profiles = snapshots
509            .iter()
510            .filter(|snapshot| {
511                snapshot.probe_status == ProbeStatus::Ok
512                    && snapshot.quota_state == QuotaState::Exhausted
513            })
514            .count();
515        let satisfied = match policy {
516            ReadinessPolicy::Any => ready_profiles > 0,
517            ReadinessPolicy::All => total_profiles > 0 && ready_profiles == total_profiles,
518        };
519        Self {
520            policy,
521            satisfied,
522            total_profiles,
523            ready_profiles,
524            not_ready_profiles: total_profiles.saturating_sub(ready_profiles),
525            failed_probes,
526            exhausted_profiles,
527        }
528    }
529}
530/// Stable machine-readable failure category.
531#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
532#[serde(rename_all = "snake_case")]
533pub enum ProviderUsageErrorCode {
534    /// Profile fields are invalid.
535    InvalidConfiguration,
536    /// Required executable is not installed or not discoverable.
537    ExecutableNotFound,
538    /// Credential file is unavailable.
539    CredentialsUnavailable,
540    /// Provider rejected or could not find credentials.
541    AuthenticationRequired,
542    /// Provider did not respond before the configured deadline.
543    Timeout,
544    /// Network transport failed.
545    Network,
546    /// Provider returned an unsuccessful response.
547    ProviderResponse,
548    /// Response shape could not be interpreted.
549    InvalidResponse,
550    /// Provider version does not expose the expected quota interface.
551    Unsupported,
552}
553
554/// Structured probe failure for automation and user interfaces.
555#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
556#[serde(rename_all = "camelCase")]
557pub struct ProviderUsageError {
558    /// Stable error code.
559    pub code: ProviderUsageErrorCode,
560    /// Safe human-readable diagnostic.
561    pub message: String,
562    /// Whether retrying later may succeed without user action.
563    pub retryable: bool,
564}
565
566/// Whether a returned snapshot came from a live probe or a cache hit.
567#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
568#[serde(rename_all = "snake_case")]
569pub enum ProviderUsageFreshness {
570    /// The provider was probed for this collection call.
571    Live,
572    /// A successful snapshot was returned from the profile cache.
573    Cached,
574}
575
576/// Optional collection metadata that does not change provider quota semantics.
577#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
578#[serde(rename_all = "camelCase")]
579pub struct ProviderUsageCollectionMetadata {
580    /// Whether this call performed a live probe or returned a cached snapshot.
581    pub freshness: ProviderUsageFreshness,
582    /// Time spent performing the original provider probe.
583    pub probe_duration_ms: u64,
584    /// Time at which a successful result entered the cache, when applicable.
585    pub cached_at_ms: Option<i64>,
586    /// Time at which the cached result expires, when applicable.
587    pub expires_at_ms: Option<i64>,
588}
589
590/// Optional billable or spend-controlled usage reported by a provider.
591#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
592#[serde(rename_all = "camelCase")]
593pub struct ProviderBillableUsage {
594    /// Provider-formatted amount consumed in the current billing window.
595    pub used: String,
596    /// Provider-formatted limit for the current billing window.
597    pub limit: String,
598    /// Remaining billable capacity from zero to one hundred.
599    pub remaining_percent: f64,
600    /// Billing-window reset time in Unix epoch seconds.
601    pub resets_at_epoch_seconds: i64,
602}
603
604/// Optional provider credit balance.
605#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
606#[serde(rename_all = "camelCase")]
607pub struct ProviderCredits {
608    /// Provider-formatted balance, when supplied.
609    pub balance: Option<String>,
610    /// Whether the provider reports a usable credit balance.
611    pub has_credits: bool,
612    /// Whether provider credits are unlimited.
613    pub unlimited: bool,
614}
615
616/// Optional summary of credits that can reset a provider rate limit.
617#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
618#[serde(rename_all = "camelCase")]
619pub struct ProviderRateLimitResetCredits {
620    /// Number of reset credits currently available.
621    pub available_count: u64,
622}
623
624/// Normalized provider quota snapshot.
625#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
626#[serde(rename_all = "camelCase")]
627pub struct ProviderUsageSnapshot {
628    /// Serialized schema version.
629    pub schema_version: u32,
630    /// Stable provider identifier.
631    pub provider_id: ProviderId,
632    /// Human-readable provider name.
633    pub provider_name: String,
634    /// Stable profile identifier.
635    pub profile_id: String,
636    /// Human-readable profile name.
637    pub profile_name: String,
638    /// Provider account label, when exposed.
639    pub account_label: Option<String>,
640    /// Machine-readable origin of quota data.
641    pub source: Option<String>,
642    /// Health of the quota probe.
643    pub probe_status: ProbeStatus,
644    /// Whether reported quota permits more work.
645    pub quota_state: QuotaState,
646    /// Provider subscription plan, when exposed.
647    pub plan: Option<String>,
648    /// Normalized quota windows.
649    pub windows: Vec<ProviderUsageWindow>,
650    /// Optional billable or spend-controlled usage.
651    #[serde(default, skip_serializing_if = "Option::is_none")]
652    pub billable_usage: Option<ProviderBillableUsage>,
653    /// Optional provider credit balance.
654    #[serde(default, skip_serializing_if = "Option::is_none")]
655    pub credits: Option<ProviderCredits>,
656    /// Optional available rate-limit reset credits.
657    #[serde(default, skip_serializing_if = "Option::is_none")]
658    pub rate_limit_reset_credits: Option<ProviderRateLimitResetCredits>,
659    /// Observation time in Unix epoch milliseconds.
660    pub observed_at_ms: i64,
661    /// Optional live-probe and cache freshness metadata.
662    #[serde(default, skip_serializing_if = "Option::is_none")]
663    pub collection: Option<ProviderUsageCollectionMetadata>,
664    /// Optional informational note.
665    pub message: Option<String>,
666    /// Structured error when the probe did not succeed.
667    pub error: Option<ProviderUsageError>,
668}
669
670impl ProviderUsageSnapshot {
671    /// Return whether this snapshot is a safe positive signal for new work.
672    ///
673    /// A probe must have completed successfully and reported available quota.
674    /// Keeping the two checks together prevents callers from treating an
675    /// unknown quota state or a failed probe as usable capacity.
676    pub fn is_ready(&self) -> bool {
677        self.probe_status == ProbeStatus::Ok && self.quota_state == QuotaState::Available
678    }
679}
680
681/// One provider quota window.
682#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
683#[serde(rename_all = "camelCase")]
684pub struct ProviderUsageWindow {
685    /// Normalized window category.
686    pub kind: ProviderUsageWindowKind,
687    /// Human-readable label.
688    pub label: String,
689    /// Consumed capacity from zero to one hundred.
690    pub used_percent: Option<f64>,
691    /// Remaining capacity from zero to one hundred.
692    pub remaining_percent: Option<f64>,
693    /// Window duration in minutes.
694    pub window_minutes: Option<u64>,
695    /// Reset time in Unix epoch seconds.
696    pub resets_at_epoch_seconds: Option<i64>,
697    /// Optional provider status detail.
698    pub detail: Option<String>,
699}
700
701/// Normalized quota window category.
702#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
703#[serde(rename_all = "snake_case")]
704pub enum ProviderUsageWindowKind {
705    /// Rolling five-hour window.
706    FiveHour,
707    /// Rolling or calendar weekly window.
708    Weekly,
709    /// Provider-specific window.
710    Other,
711}
712
713#[derive(Debug, Clone)]
714pub(crate) struct ProbeFailure {
715    error: ProviderUsageError,
716}
717
718impl ProbeFailure {
719    fn new(code: ProviderUsageErrorCode, message: impl Into<String>, retryable: bool) -> Self {
720        Self {
721            error: ProviderUsageError {
722                code,
723                message: message.into(),
724                retryable,
725            },
726        }
727    }
728
729    fn status(&self) -> ProbeStatus {
730        match self.error.code {
731            ProviderUsageErrorCode::InvalidConfiguration => ProbeStatus::InvalidConfiguration,
732            ProviderUsageErrorCode::CredentialsUnavailable
733            | ProviderUsageErrorCode::AuthenticationRequired => ProbeStatus::AuthenticationRequired,
734            ProviderUsageErrorCode::Unsupported => ProbeStatus::Unsupported,
735            ProviderUsageErrorCode::InvalidResponse => ProbeStatus::InvalidResponse,
736            ProviderUsageErrorCode::ExecutableNotFound
737            | ProviderUsageErrorCode::Timeout
738            | ProviderUsageErrorCode::Network
739            | ProviderUsageErrorCode::ProviderResponse => ProbeStatus::TransientError,
740        }
741    }
742}
743
744/// Reusable collector with injectable transport settings.
745#[derive(Debug, Clone)]
746pub struct AgentQuotaClient {
747    http_client: reqwest::Client,
748    claude_api_url: String,
749    claude_model: String,
750    codex_timeout: Duration,
751    claude_timeout: Duration,
752}
753
754impl Default for AgentQuotaClient {
755    fn default() -> Self {
756        let http_client = reqwest::Client::builder()
757            .user_agent(concat!("agent-quota/", env!("CARGO_PKG_VERSION")))
758            .redirect(Policy::none())
759            .build()
760            .unwrap_or_else(|_| reqwest::Client::new());
761        Self {
762            http_client,
763            claude_api_url: CLAUDE_API_URL.to_owned(),
764            claude_model: CLAUDE_USAGE_MODEL.to_owned(),
765            codex_timeout: DEFAULT_CODEX_TIMEOUT,
766            claude_timeout: DEFAULT_CLAUDE_TIMEOUT,
767        }
768    }
769}
770
771impl AgentQuotaClient {
772    /// Construct a collector with production defaults.
773    pub fn new() -> Self {
774        Self::default()
775    }
776
777    /// Override the HTTP client, useful for embedding and tests.
778    ///
779    /// The default client disables redirects. A custom client is responsible
780    /// for protecting authorization headers across redirects and proxies.
781    pub fn with_http_client(mut self, http_client: reqwest::Client) -> Self {
782        self.http_client = http_client;
783        self
784    }
785
786    /// Override the Claude Messages endpoint, useful for deterministic tests.
787    ///
788    /// The local OAuth token is sent to this URL. Never accept the URL from
789    /// untrusted input.
790    pub fn with_claude_api_url(mut self, url: impl Into<String>) -> Self {
791        self.claude_api_url = url.into();
792        self
793    }
794
795    /// Override the minimal Claude model used to obtain quota headers.
796    pub fn with_claude_model(mut self, model: impl Into<String>) -> Self {
797        self.claude_model = model.into();
798        self
799    }
800
801    /// Override provider deadlines.
802    pub fn with_timeouts(mut self, codex: Duration, claude: Duration) -> Self {
803        self.codex_timeout = codex;
804        self.claude_timeout = claude;
805        self
806    }
807
808    /// Collect normalized snapshots for an explicit selection.
809    pub async fn collect_usage(&self, options: CollectUsageOptions) -> Vec<ProviderUsageSnapshot> {
810        let probes = options
811            .selected_profiles()
812            .into_iter()
813            .map(|profile| self.collect_profile_usage(profile));
814        futures::future::join_all(probes).await
815    }
816
817    /// Collect one configured profile.
818    pub async fn collect_profile_usage(&self, profile: ProviderProfile) -> ProviderUsageSnapshot {
819        let started = Instant::now();
820        let mut snapshot = if let Err(error) = profile.validate() {
821            failed_snapshot(
822                &profile,
823                ProbeFailure::new(
824                    ProviderUsageErrorCode::InvalidConfiguration,
825                    error.to_string(),
826                    false,
827                ),
828            )
829        } else {
830            let result = providers::adapter(profile.provider)
831                .collect(self, &profile)
832                .await;
833            result.unwrap_or_else(|error| failed_snapshot(&profile, error))
834        };
835        snapshot.collection = Some(ProviderUsageCollectionMetadata {
836            freshness: ProviderUsageFreshness::Live,
837            probe_duration_ms: u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX),
838            cached_at_ms: None,
839            expires_at_ms: None,
840        });
841        snapshot
842    }
843
844    pub(crate) async fn query_codex_usage(
845        &self,
846        profile: &ProviderProfile,
847    ) -> Result<ProviderUsageSnapshot, ProbeFailure> {
848        let executable = profile
849            .command_path
850            .as_deref()
851            .unwrap_or_else(|| Path::new("codex"));
852        let mut command = Command::new(executable);
853        command.args(["app-server", "--listen", "stdio://"]);
854        command.envs(&profile.env);
855        command
856            .stdin(Stdio::piped())
857            .stdout(Stdio::piped())
858            .stderr(Stdio::null())
859            .kill_on_drop(true);
860        #[cfg(target_os = "windows")]
861        command.creation_flags(CREATE_NO_WINDOW);
862
863        let mut child = command.spawn().map_err(|error| {
864            let code = if error.kind() == std::io::ErrorKind::NotFound {
865                ProviderUsageErrorCode::ExecutableNotFound
866            } else {
867                ProviderUsageErrorCode::ProviderResponse
868            };
869            ProbeFailure::new(
870                code,
871                format!("could not start `{}`: {error}", executable.display()),
872                true,
873            )
874        })?;
875        let mut stdin = child.stdin.take().ok_or_else(|| {
876            ProbeFailure::new(
877                ProviderUsageErrorCode::InvalidResponse,
878                "Codex app-server stdin was unavailable",
879                true,
880            )
881        })?;
882        let stdout = child.stdout.take().ok_or_else(|| {
883            ProbeFailure::new(
884                ProviderUsageErrorCode::InvalidResponse,
885                "Codex app-server stdout was unavailable",
886                true,
887            )
888        })?;
889
890        write_rpc(
891            &mut stdin,
892            json!({
893                "method": "initialize",
894                "id": 1,
895                "params": {
896                    "clientInfo": {
897                        "name": "agent_quota",
898                        "title": "Agent Quota",
899                        "version": env!("CARGO_PKG_VERSION"),
900                    }
901                }
902            }),
903        )
904        .await
905        .map_err(codex_io_failure)?;
906        write_rpc(&mut stdin, json!({ "method": "initialized", "params": {} }))
907            .await
908            .map_err(codex_io_failure)?;
909        write_rpc(
910            &mut stdin,
911            json!({
912                "method": "account/read",
913                "id": 2,
914                "params": { "refreshToken": false }
915            }),
916        )
917        .await
918        .map_err(codex_io_failure)?;
919        write_rpc(
920            &mut stdin,
921            json!({ "method": "account/rateLimits/read", "id": 3 }),
922        )
923        .await
924        .map_err(codex_io_failure)?;
925
926        let mut lines = BufReader::new(stdout).lines();
927        let read = async {
928            let mut account_response = None;
929            let mut rate_limits_response = None;
930            while account_response.is_none() || rate_limits_response.is_none() {
931                let Some(line) = lines.next_line().await.map_err(codex_io_failure)? else {
932                    break;
933                };
934                let Ok(value) = serde_json::from_str::<Value>(&line) else {
935                    continue;
936                };
937                match value.get("id").and_then(Value::as_i64) {
938                    Some(2) => account_response = Some(value),
939                    Some(3) => rate_limits_response = Some(value),
940                    _ => {}
941                }
942            }
943            let account_response = account_response.ok_or_else(|| {
944                ProbeFailure::new(
945                    ProviderUsageErrorCode::InvalidResponse,
946                    "Codex account response was missing",
947                    true,
948                )
949            })?;
950            let rate_limits_response = rate_limits_response.ok_or_else(|| {
951                ProbeFailure::new(
952                    ProviderUsageErrorCode::InvalidResponse,
953                    "Codex rate-limit response was missing",
954                    true,
955                )
956            })?;
957            Ok::<_, ProbeFailure>((account_response, rate_limits_response))
958        };
959
960        let responses = timeout(self.codex_timeout, read).await.map_err(|_| {
961            ProbeFailure::new(
962                ProviderUsageErrorCode::Timeout,
963                "timed out waiting for Codex quota data",
964                true,
965            )
966        })??;
967        let _ = child.kill().await;
968        codex_snapshot_from_responses(profile, &responses.0, &responses.1)
969    }
970
971    pub(crate) async fn query_claude_usage(
972        &self,
973        profile: &ProviderProfile,
974    ) -> Result<ProviderUsageSnapshot, ProbeFailure> {
975        let token = read_claude_oauth_token(profile.credentials_path.as_deref())?;
976        let request = self
977            .http_client
978            .post(&self.claude_api_url)
979            .bearer_auth(token)
980            .header("anthropic-version", "2023-06-01")
981            .header("anthropic-beta", "oauth-2025-04-20")
982            .json(&json!({
983                "model": self.claude_model,
984                "max_tokens": 1,
985                "messages": [{ "role": "user", "content": "hi" }],
986            }));
987        let response = timeout(self.claude_timeout, request.send())
988            .await
989            .map_err(|_| {
990                ProbeFailure::new(
991                    ProviderUsageErrorCode::Timeout,
992                    "timed out waiting for Claude quota headers",
993                    true,
994                )
995            })?
996            .map_err(|error| {
997                ProbeFailure::new(
998                    ProviderUsageErrorCode::Network,
999                    format!("Claude quota request failed: {error}"),
1000                    true,
1001                )
1002            })?;
1003
1004        let status = response.status();
1005        let snapshot = claude_snapshot_from_headers(profile, response.headers(), status)?;
1006        if status.is_success() || status == StatusCode::TOO_MANY_REQUESTS {
1007            Ok(snapshot)
1008        } else {
1009            Err(http_status_failure("Claude", status))
1010        }
1011    }
1012}
1013
1014fn codex_io_failure(error: std::io::Error) -> ProbeFailure {
1015    ProbeFailure::new(
1016        ProviderUsageErrorCode::ProviderResponse,
1017        format!("Codex app-server communication failed: {error}"),
1018        true,
1019    )
1020}
1021
1022fn http_status_failure(provider: &str, status: StatusCode) -> ProbeFailure {
1023    if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) {
1024        ProbeFailure::new(
1025            ProviderUsageErrorCode::AuthenticationRequired,
1026            format!("{provider} rejected the local credentials ({status})"),
1027            false,
1028        )
1029    } else {
1030        ProbeFailure::new(
1031            ProviderUsageErrorCode::ProviderResponse,
1032            format!("{provider} returned HTTP {status}"),
1033            status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS,
1034        )
1035    }
1036}
1037
1038async fn write_rpc(stdin: &mut ChildStdin, value: Value) -> std::io::Result<()> {
1039    stdin.write_all(value.to_string().as_bytes()).await?;
1040    stdin.write_all(b"\n").await?;
1041    stdin.flush().await
1042}
1043
1044fn rpc_failure(value: &Value, operation: &str) -> Option<ProbeFailure> {
1045    let error = value.get("error")?;
1046    let message = error
1047        .get("message")
1048        .and_then(Value::as_str)
1049        .unwrap_or("unknown provider error");
1050    let lower = message.to_ascii_lowercase();
1051    let authentication = ["auth", "login", "sign in", "credential"]
1052        .iter()
1053        .any(|needle| lower.contains(needle));
1054    Some(if authentication {
1055        ProbeFailure::new(
1056            ProviderUsageErrorCode::AuthenticationRequired,
1057            format!("Codex {operation} failed: {message}"),
1058            false,
1059        )
1060    } else {
1061        ProbeFailure::new(
1062            ProviderUsageErrorCode::ProviderResponse,
1063            format!("Codex {operation} failed: {message}"),
1064            true,
1065        )
1066    })
1067}
1068
1069fn codex_snapshot_from_responses(
1070    profile: &ProviderProfile,
1071    account_response: &Value,
1072    rate_limits_response: &Value,
1073) -> Result<ProviderUsageSnapshot, ProbeFailure> {
1074    if let Some(error) = rpc_failure(account_response, "account lookup") {
1075        return Err(error);
1076    }
1077    if let Some(error) = rpc_failure(rate_limits_response, "quota lookup") {
1078        return Err(error);
1079    }
1080
1081    let account = account_response
1082        .pointer("/result/account")
1083        .and_then(Value::as_object)
1084        .ok_or_else(|| {
1085            ProbeFailure::new(
1086                ProviderUsageErrorCode::AuthenticationRequired,
1087                "Codex did not return a signed-in account",
1088                false,
1089            )
1090        })?;
1091    let rate_limits = rate_limits_response
1092        .pointer("/result/rateLimits")
1093        .ok_or_else(|| {
1094            ProbeFailure::new(
1095                ProviderUsageErrorCode::InvalidResponse,
1096                "Codex did not return rate-limit data",
1097                true,
1098            )
1099        })?;
1100
1101    let plan = account
1102        .get("planType")
1103        .and_then(Value::as_str)
1104        .map(format_plan_name);
1105    let account_label = account
1106        .get("email")
1107        .or_else(|| account.get("username"))
1108        .or_else(|| account.get("id"))
1109        .and_then(Value::as_str)
1110        .map(ToOwned::to_owned);
1111    let mut windows = Vec::new();
1112    if let Some(primary) = rate_limits.get("primary") {
1113        if let Some(window) = window_from_object("Primary window", primary) {
1114            windows.push(window);
1115        }
1116    }
1117    if let Some(secondary) = rate_limits.get("secondary") {
1118        if let Some(window) = window_from_object("Secondary window", secondary) {
1119            windows.push(window);
1120        }
1121    }
1122    sort_usage_windows(&mut windows);
1123    if windows.is_empty() {
1124        return Err(ProbeFailure::new(
1125            ProviderUsageErrorCode::InvalidResponse,
1126            "Codex rate-limit response contained no recognizable quota windows",
1127            true,
1128        ));
1129    }
1130
1131    let billable_usage = codex_billable_usage(rate_limits);
1132    let credits = codex_credits(rate_limits);
1133    let rate_limit_reset_credits = rate_limits_response
1134        .pointer("/result/rateLimitResetCredits/availableCount")
1135        .and_then(integer_value)
1136        .and_then(|value| u64::try_from(value).ok())
1137        .map(|available_count| ProviderRateLimitResetCredits { available_count });
1138    let reached = rate_limits
1139        .get("rateLimitReachedType")
1140        .and_then(Value::as_str)
1141        .filter(|value| !value.is_empty());
1142    let spend_control_reached = rate_limits
1143        .get("spendControlReached")
1144        .and_then(Value::as_bool)
1145        .unwrap_or(false)
1146        || billable_usage
1147            .as_ref()
1148            .is_some_and(|usage| usage.remaining_percent <= 0.0);
1149    let quota_state = if reached.is_some() || spend_control_reached {
1150        QuotaState::Exhausted
1151    } else {
1152        quota_state_from_windows(&windows)
1153    };
1154    let message = reached
1155        .map(|value| format!("Limit reached: {}", value.replace('_', " ")))
1156        .or_else(|| spend_control_reached.then(|| "Billable usage limit reached".to_owned()));
1157    let mut snapshot = success_snapshot(
1158        profile,
1159        account_label,
1160        Some("codex_app_server"),
1161        plan,
1162        windows,
1163        message,
1164        quota_state,
1165    );
1166    snapshot.billable_usage = billable_usage;
1167    snapshot.credits = credits;
1168    snapshot.rate_limit_reset_credits = rate_limit_reset_credits;
1169    Ok(snapshot)
1170}
1171
1172fn codex_billable_usage(rate_limits: &Value) -> Option<ProviderBillableUsage> {
1173    let value = rate_limits.get("individualLimit")?.as_object()?;
1174    Some(ProviderBillableUsage {
1175        used: scalar_string(value.get("used")?)?,
1176        limit: scalar_string(value.get("limit")?)?,
1177        remaining_percent: number_value(value.get("remainingPercent")?)?.clamp(0.0, 100.0),
1178        resets_at_epoch_seconds: normalize_epoch_seconds(integer_value(value.get("resetsAt")?)?),
1179    })
1180}
1181
1182fn codex_credits(rate_limits: &Value) -> Option<ProviderCredits> {
1183    let value = rate_limits.get("credits")?.as_object()?;
1184    Some(ProviderCredits {
1185        balance: value.get("balance").and_then(scalar_string),
1186        has_credits: value.get("hasCredits")?.as_bool()?,
1187        unlimited: value.get("unlimited")?.as_bool()?,
1188    })
1189}
1190
1191fn scalar_string(value: &Value) -> Option<String> {
1192    value
1193        .as_str()
1194        .map(ToOwned::to_owned)
1195        .or_else(|| value.as_number().map(ToString::to_string))
1196}
1197
1198fn read_claude_oauth_token(credentials_path: Option<&Path>) -> Result<String, ProbeFailure> {
1199    let path = credentials_path
1200        .map(Path::to_path_buf)
1201        .map(Ok)
1202        .unwrap_or_else(default_claude_credentials_path)
1203        .map_err(|error| {
1204            ProbeFailure::new(
1205                ProviderUsageErrorCode::CredentialsUnavailable,
1206                error.to_string(),
1207                false,
1208            )
1209        })?;
1210    let raw = fs::read_to_string(&path).map_err(|error| {
1211        ProbeFailure::new(
1212            ProviderUsageErrorCode::CredentialsUnavailable,
1213            format!(
1214                "Claude Code credentials were not readable at {}: {error}",
1215                path.display()
1216            ),
1217            false,
1218        )
1219    })?;
1220    let value = serde_json::from_str::<Value>(&raw).map_err(|error| {
1221        ProbeFailure::new(
1222            ProviderUsageErrorCode::CredentialsUnavailable,
1223            format!(
1224                "Claude Code credentials at {} are invalid: {error}",
1225                path.display()
1226            ),
1227            false,
1228        )
1229    })?;
1230    find_claude_oauth_token(&value)
1231        .map(ToOwned::to_owned)
1232        .ok_or_else(|| {
1233            ProbeFailure::new(
1234                ProviderUsageErrorCode::AuthenticationRequired,
1235                format!(
1236                    "Claude Code OAuth token was not found in {}",
1237                    path.display()
1238                ),
1239                false,
1240            )
1241        })
1242}
1243
1244/// Conventional Claude Code credentials path for the current user.
1245pub fn default_claude_credentials_path() -> AgentQuotaResult<PathBuf> {
1246    dirs::home_dir()
1247        .map(|home| home.join(CLAUDE_CREDENTIALS_RELATIVE_PATH))
1248        .ok_or_else(|| validation_error("home directory was not available"))
1249}
1250
1251fn find_claude_oauth_token(value: &Value) -> Option<&str> {
1252    [
1253        "/claudeAiOauth/accessToken",
1254        "/oauth/accessToken",
1255        "/claudeAiOauth/access_token",
1256    ]
1257    .into_iter()
1258    .find_map(|pointer| value.pointer(pointer).and_then(Value::as_str))
1259    .filter(|token| !token.trim().is_empty())
1260}
1261
1262fn claude_snapshot_from_headers(
1263    profile: &ProviderProfile,
1264    headers: &HeaderMap,
1265    status: StatusCode,
1266) -> Result<ProviderUsageSnapshot, ProbeFailure> {
1267    if matches!(status, StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN) {
1268        return Err(http_status_failure("Claude", status));
1269    }
1270
1271    let mut windows = [
1272        claude_header_window(
1273            headers,
1274            ProviderUsageWindowKind::FiveHour,
1275            "5h window",
1276            "anthropic-ratelimit-unified-5h-utilization",
1277            "anthropic-ratelimit-unified-5h-reset",
1278        ),
1279        claude_header_window(
1280            headers,
1281            ProviderUsageWindowKind::Weekly,
1282            "Weekly window",
1283            "anthropic-ratelimit-unified-7d-utilization",
1284            "anthropic-ratelimit-unified-7d-reset",
1285        ),
1286    ]
1287    .into_iter()
1288    .flatten()
1289    .collect::<Vec<_>>();
1290    sort_usage_windows(&mut windows);
1291
1292    if windows.is_empty() {
1293        let code = if status.is_success() {
1294            ProviderUsageErrorCode::Unsupported
1295        } else {
1296            ProviderUsageErrorCode::ProviderResponse
1297        };
1298        return Err(ProbeFailure::new(
1299            code,
1300            if status.is_success() {
1301                "Anthropic did not return supported unified quota headers".to_owned()
1302            } else {
1303                format!("Claude returned HTTP {status} without quota headers")
1304            },
1305            status.is_server_error() || status == StatusCode::TOO_MANY_REQUESTS,
1306        ));
1307    }
1308
1309    let provider_status = header_string(headers, "anthropic-ratelimit-unified-status");
1310    let quota_state = if status == StatusCode::TOO_MANY_REQUESTS
1311        || provider_status.as_deref().is_some_and(|value| {
1312            ["rejected", "exhausted", "blocked", "rate_limited"]
1313                .iter()
1314                .any(|candidate| value.eq_ignore_ascii_case(candidate))
1315        }) {
1316        QuotaState::Exhausted
1317    } else {
1318        quota_state_from_windows(&windows)
1319    };
1320    let message = provider_status
1321        .filter(|value| !value.eq_ignore_ascii_case("active"))
1322        .map(|value| format!("Claude quota status: {value}"))
1323        .or_else(|| Some("Live quota from Anthropic rate-limit headers.".to_owned()));
1324
1325    Ok(success_snapshot(
1326        profile,
1327        None,
1328        Some("anthropic_rate_limit_headers"),
1329        None,
1330        windows,
1331        message,
1332        quota_state,
1333    ))
1334}
1335
1336fn claude_header_window(
1337    headers: &HeaderMap,
1338    kind: ProviderUsageWindowKind,
1339    label: &str,
1340    utilization_header: &str,
1341    reset_header: &str,
1342) -> Option<ProviderUsageWindow> {
1343    let utilization = header_number(headers, utilization_header)?;
1344    if !utilization.is_finite() {
1345        return None;
1346    }
1347    let used_percent = (utilization * 100.0).clamp(0.0, 100.0);
1348    Some(ProviderUsageWindow {
1349        kind,
1350        label: label.to_owned(),
1351        used_percent: Some(used_percent),
1352        remaining_percent: Some((100.0 - used_percent).clamp(0.0, 100.0)),
1353        window_minutes: match kind {
1354            ProviderUsageWindowKind::FiveHour => Some(300),
1355            ProviderUsageWindowKind::Weekly => Some(10_080),
1356            ProviderUsageWindowKind::Other => None,
1357        },
1358        resets_at_epoch_seconds: header_integer(headers, reset_header).map(normalize_epoch_seconds),
1359        detail: header_string(
1360            headers,
1361            match kind {
1362                ProviderUsageWindowKind::FiveHour => "anthropic-ratelimit-unified-5h-status",
1363                ProviderUsageWindowKind::Weekly => "anthropic-ratelimit-unified-7d-status",
1364                ProviderUsageWindowKind::Other => "anthropic-ratelimit-unified-status",
1365            },
1366        )
1367        .map(|value| format!("Status: {value}")),
1368    })
1369}
1370
1371fn header_string(headers: &HeaderMap, name: &str) -> Option<String> {
1372    headers
1373        .get(name)
1374        .and_then(|value| value.to_str().ok())
1375        .map(str::trim)
1376        .filter(|value| !value.is_empty())
1377        .map(ToOwned::to_owned)
1378}
1379
1380fn header_number(headers: &HeaderMap, name: &str) -> Option<f64> {
1381    header_string(headers, name).and_then(|value| value.parse::<f64>().ok())
1382}
1383
1384fn header_integer(headers: &HeaderMap, name: &str) -> Option<i64> {
1385    header_string(headers, name).and_then(|value| value.parse::<i64>().ok())
1386}
1387
1388fn window_from_object(label_hint: &str, value: &Value) -> Option<ProviderUsageWindow> {
1389    value
1390        .as_object()
1391        .and_then(|map| window_from_map(label_hint, map))
1392}
1393
1394fn window_from_map(label_hint: &str, map: &Map<String, Value>) -> Option<ProviderUsageWindow> {
1395    let used_percent = number_field(
1396        map,
1397        &[
1398            "usedPercent",
1399            "used_percent",
1400            "percentUsed",
1401            "percent_used",
1402            "usagePercent",
1403            "usage_percent",
1404        ],
1405    )
1406    .filter(|value| value.is_finite())
1407    .map(|value| value.clamp(0.0, 100.0));
1408    let remaining_percent = number_field(
1409        map,
1410        &[
1411            "remainingPercent",
1412            "remaining_percent",
1413            "percentRemaining",
1414            "percent_remaining",
1415        ],
1416    )
1417    .filter(|value| value.is_finite())
1418    .map(|value| value.clamp(0.0, 100.0))
1419    .or_else(|| used_percent.map(|used| (100.0 - used).clamp(0.0, 100.0)));
1420    if used_percent.is_none() && remaining_percent.is_none() {
1421        return None;
1422    }
1423
1424    let window_minutes = number_field(
1425        map,
1426        &[
1427            "windowDurationMins",
1428            "window_duration_mins",
1429            "windowMinutes",
1430            "window_minutes",
1431        ],
1432    )
1433    .and_then(f64_to_u64)
1434    .or_else(|| {
1435        number_field(map, &["limitWindowSeconds", "limit_window_seconds"])
1436            .and_then(|seconds| f64_to_u64(seconds / 60.0))
1437    });
1438    let resets_at_epoch_seconds = integer_field(
1439        map,
1440        &["resetsAt", "resetAt", "reset_at", "resets_at", "resetTime"],
1441    )
1442    .map(normalize_epoch_seconds);
1443    let kind = kind_for_usage_window(label_hint, window_minutes);
1444    Some(ProviderUsageWindow {
1445        kind,
1446        label: label_for_usage_window(label_hint, window_minutes, kind),
1447        used_percent,
1448        remaining_percent,
1449        window_minutes,
1450        resets_at_epoch_seconds,
1451        detail: None,
1452    })
1453}
1454
1455fn number_field(map: &Map<String, Value>, keys: &[&str]) -> Option<f64> {
1456    keys.iter()
1457        .find_map(|key| map.get(*key).and_then(number_value))
1458}
1459
1460fn integer_field(map: &Map<String, Value>, keys: &[&str]) -> Option<i64> {
1461    keys.iter()
1462        .find_map(|key| map.get(*key).and_then(integer_value))
1463}
1464
1465fn number_value(value: &Value) -> Option<f64> {
1466    value
1467        .as_f64()
1468        .or_else(|| value.as_str().and_then(|raw| raw.parse::<f64>().ok()))
1469}
1470
1471fn integer_value(value: &Value) -> Option<i64> {
1472    value
1473        .as_i64()
1474        .or_else(|| value.as_str().and_then(|raw| raw.parse::<i64>().ok()))
1475}
1476
1477fn f64_to_u64(value: f64) -> Option<u64> {
1478    if value.is_finite() && value >= 0.0 {
1479        u64::try_from(value.round() as i128).ok()
1480    } else {
1481        None
1482    }
1483}
1484
1485fn normalize_epoch_seconds(value: i64) -> i64 {
1486    if value > 10_000_000_000 {
1487        value / 1000
1488    } else {
1489        value
1490    }
1491}
1492
1493fn kind_for_usage_window(label_hint: &str, window_minutes: Option<u64>) -> ProviderUsageWindowKind {
1494    let label_hint = label_hint.to_lowercase();
1495    match window_minutes {
1496        Some(300) => ProviderUsageWindowKind::FiveHour,
1497        Some(10_080) => ProviderUsageWindowKind::Weekly,
1498        _ if ["five", "5h", "5 hour", "5-hour"]
1499            .iter()
1500            .any(|needle| label_hint.contains(needle)) =>
1501        {
1502            ProviderUsageWindowKind::FiveHour
1503        }
1504        _ if ["weekly", "week", "7 day", "7-day"]
1505            .iter()
1506            .any(|needle| label_hint.contains(needle)) =>
1507        {
1508            ProviderUsageWindowKind::Weekly
1509        }
1510        _ => ProviderUsageWindowKind::Other,
1511    }
1512}
1513
1514fn label_for_usage_window(
1515    label_hint: &str,
1516    window_minutes: Option<u64>,
1517    kind: ProviderUsageWindowKind,
1518) -> String {
1519    match kind {
1520        ProviderUsageWindowKind::FiveHour => "5h window".to_owned(),
1521        ProviderUsageWindowKind::Weekly => "Weekly window".to_owned(),
1522        ProviderUsageWindowKind::Other => match window_minutes {
1523            Some(minutes) if minutes >= 60 && minutes % 60 == 0 => {
1524                format!("{}h window", minutes / 60)
1525            }
1526            Some(minutes) => format!("{minutes}m window"),
1527            None => humanize_label(label_hint),
1528        },
1529    }
1530}
1531
1532fn sort_usage_windows(windows: &mut [ProviderUsageWindow]) {
1533    windows.sort_by(|left, right| {
1534        window_order(left)
1535            .cmp(&window_order(right))
1536            .then_with(|| left.label.cmp(&right.label))
1537    });
1538}
1539
1540fn window_order(window: &ProviderUsageWindow) -> u8 {
1541    match window.kind {
1542        ProviderUsageWindowKind::FiveHour => 0,
1543        ProviderUsageWindowKind::Weekly => 1,
1544        ProviderUsageWindowKind::Other => 2,
1545    }
1546}
1547
1548fn humanize_label(value: &str) -> String {
1549    let mut label = String::new();
1550    let mut previous_lowercase = false;
1551    for ch in value.chars() {
1552        if ch == '_' || ch == '-' {
1553            label.push(' ');
1554            previous_lowercase = false;
1555        } else if ch.is_ascii_uppercase() && previous_lowercase {
1556            label.push(' ');
1557            label.push(ch);
1558            previous_lowercase = false;
1559        } else {
1560            label.push(ch);
1561            previous_lowercase = ch.is_ascii_lowercase();
1562        }
1563    }
1564    let label = label.trim();
1565    if label.eq_ignore_ascii_case("primary") {
1566        "Primary window".to_owned()
1567    } else if label.eq_ignore_ascii_case("secondary") {
1568        "Secondary window".to_owned()
1569    } else {
1570        let mut chars = label.chars();
1571        match chars.next() {
1572            Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
1573            None => "Usage window".to_owned(),
1574        }
1575    }
1576}
1577
1578fn format_plan_name(value: &str) -> String {
1579    value
1580        .replace(['_', '-'], " ")
1581        .split_whitespace()
1582        .map(|word| {
1583            let mut chars = word.chars();
1584            match chars.next() {
1585                Some(first) => first.to_uppercase().collect::<String>() + chars.as_str(),
1586                None => String::new(),
1587            }
1588        })
1589        .collect::<Vec<_>>()
1590        .join(" ")
1591}
1592
1593fn quota_state_from_windows(windows: &[ProviderUsageWindow]) -> QuotaState {
1594    if windows.is_empty() {
1595        return QuotaState::Unknown;
1596    }
1597    if windows.iter().any(|window| {
1598        window
1599            .remaining_percent
1600            .is_some_and(|remaining| remaining <= 0.0)
1601            || window.used_percent.is_some_and(|used| used >= 100.0)
1602    }) {
1603        QuotaState::Exhausted
1604    } else if windows
1605        .iter()
1606        .any(|window| window.remaining_percent.is_some() || window.used_percent.is_some())
1607    {
1608        QuotaState::Available
1609    } else {
1610        QuotaState::Unknown
1611    }
1612}
1613
1614fn success_snapshot(
1615    profile: &ProviderProfile,
1616    account_label: Option<String>,
1617    source: Option<&str>,
1618    plan: Option<String>,
1619    windows: Vec<ProviderUsageWindow>,
1620    message: Option<String>,
1621    quota_state: QuotaState,
1622) -> ProviderUsageSnapshot {
1623    ProviderUsageSnapshot {
1624        schema_version: SNAPSHOT_SCHEMA_VERSION,
1625        provider_id: profile.provider,
1626        provider_name: profile.provider.label().to_owned(),
1627        profile_id: profile.id.clone(),
1628        profile_name: profile.display_name(),
1629        account_label,
1630        source: source.map(ToOwned::to_owned),
1631        probe_status: ProbeStatus::Ok,
1632        quota_state,
1633        plan,
1634        windows,
1635        billable_usage: None,
1636        credits: None,
1637        rate_limit_reset_credits: None,
1638        observed_at_ms: now_epoch_ms(),
1639        collection: None,
1640        message,
1641        error: None,
1642    }
1643}
1644
1645fn failed_snapshot(profile: &ProviderProfile, failure: ProbeFailure) -> ProviderUsageSnapshot {
1646    ProviderUsageSnapshot {
1647        schema_version: SNAPSHOT_SCHEMA_VERSION,
1648        provider_id: profile.provider,
1649        provider_name: profile.provider.label().to_owned(),
1650        profile_id: profile.id.clone(),
1651        profile_name: profile.display_name(),
1652        account_label: None,
1653        source: None,
1654        probe_status: failure.status(),
1655        quota_state: QuotaState::Unknown,
1656        plan: None,
1657        windows: Vec::new(),
1658        billable_usage: None,
1659        credits: None,
1660        rate_limit_reset_credits: None,
1661        observed_at_ms: now_epoch_ms(),
1662        collection: None,
1663        message: None,
1664        error: Some(failure.error),
1665    }
1666}
1667
1668fn now_epoch_ms() -> i64 {
1669    SystemTime::now()
1670        .duration_since(UNIX_EPOCH)
1671        .map(|duration| i64::try_from(duration.as_millis()).unwrap_or(i64::MAX))
1672        .unwrap_or_default()
1673}
1674
1675#[derive(Debug, Clone)]
1676struct CachedProviderUsage {
1677    snapshot: ProviderUsageSnapshot,
1678    cached_at_ms: i64,
1679}
1680
1681/// Profile-keyed provider cache for repeated UI or CLI polling.
1682#[derive(Debug, Clone)]
1683pub struct ProviderUsageCache {
1684    inner: Arc<Mutex<BTreeMap<ProviderProfile, CachedProviderUsage>>>,
1685    ttl: Duration,
1686}
1687
1688impl Default for ProviderUsageCache {
1689    fn default() -> Self {
1690        Self {
1691            inner: Arc::new(Mutex::new(BTreeMap::new())),
1692            ttl: DEFAULT_CACHE_TTL,
1693        }
1694    }
1695}
1696
1697impl ProviderUsageCache {
1698    /// Construct a cache with a custom time-to-live.
1699    pub fn with_ttl(ttl: Duration) -> Self {
1700        Self {
1701            ttl,
1702            ..Self::default()
1703        }
1704    }
1705
1706    /// Collect through a cache keyed by every complete provider profile.
1707    ///
1708    /// Failed profiles are not cached, allowing immediate recovery after sign-in
1709    /// or transient failures while successful profiles keep their cached values.
1710    pub async fn collect(
1711        &self,
1712        client: &AgentQuotaClient,
1713        options: CollectUsageOptions,
1714        force_refresh: bool,
1715    ) -> Vec<ProviderUsageSnapshot> {
1716        let profiles = options.selected_profiles();
1717        let probes = profiles.into_iter().map(|profile| async move {
1718            let now = now_epoch_ms();
1719            let ttl_ms = i64::try_from(self.ttl.as_millis()).unwrap_or(i64::MAX);
1720            if !force_refresh {
1721                let guard = self.inner.lock().await;
1722                if let Some(cached) = guard.get(&profile) {
1723                    if now.saturating_sub(cached.cached_at_ms) < ttl_ms {
1724                        let mut snapshot = cached.snapshot.clone();
1725                        let metadata =
1726                            snapshot
1727                                .collection
1728                                .get_or_insert(ProviderUsageCollectionMetadata {
1729                                    freshness: ProviderUsageFreshness::Cached,
1730                                    probe_duration_ms: 0,
1731                                    cached_at_ms: Some(cached.cached_at_ms),
1732                                    expires_at_ms: Some(cached.cached_at_ms.saturating_add(ttl_ms)),
1733                                });
1734                        metadata.freshness = ProviderUsageFreshness::Cached;
1735                        return snapshot;
1736                    }
1737                }
1738            }
1739
1740            let mut snapshot = client.collect_profile_usage(profile.clone()).await;
1741            let mut guard = self.inner.lock().await;
1742            if snapshot.probe_status == ProbeStatus::Ok {
1743                let cached_at_ms = now_epoch_ms();
1744                let metadata = snapshot
1745                    .collection
1746                    .get_or_insert(ProviderUsageCollectionMetadata {
1747                        freshness: ProviderUsageFreshness::Live,
1748                        probe_duration_ms: 0,
1749                        cached_at_ms: None,
1750                        expires_at_ms: None,
1751                    });
1752                metadata.freshness = ProviderUsageFreshness::Live;
1753                metadata.cached_at_ms = Some(cached_at_ms);
1754                metadata.expires_at_ms = Some(cached_at_ms.saturating_add(ttl_ms));
1755                guard.insert(
1756                    profile,
1757                    CachedProviderUsage {
1758                        snapshot: snapshot.clone(),
1759                        cached_at_ms,
1760                    },
1761                );
1762            } else {
1763                guard.remove(&profile);
1764            }
1765            snapshot
1766        });
1767        futures::future::join_all(probes).await
1768    }
1769
1770    /// Clear every cached selection.
1771    pub async fn clear(&self) {
1772        self.inner.lock().await.clear();
1773    }
1774}
1775
1776/// Collect all conventional local provider profiles.
1777pub async fn collect_provider_usage() -> Vec<ProviderUsageSnapshot> {
1778    collect_usage(CollectUsageOptions::all()).await
1779}
1780
1781/// Collect an explicit selection with production defaults.
1782pub async fn collect_usage(options: CollectUsageOptions) -> Vec<ProviderUsageSnapshot> {
1783    AgentQuotaClient::new().collect_usage(options).await
1784}
1785
1786/// Collect one conventional provider profile.
1787pub async fn collect_single_provider_usage(provider: ProviderId) -> ProviderUsageSnapshot {
1788    AgentQuotaClient::new()
1789        .collect_profile_usage(ProviderProfile::default_for_provider(provider))
1790        .await
1791}
1792
1793/// Collect one configured provider profile.
1794pub async fn collect_profile_usage(profile: ProviderProfile) -> ProviderUsageSnapshot {
1795    AgentQuotaClient::new().collect_profile_usage(profile).await
1796}
1797
1798#[cfg(test)]
1799mod tests {
1800    use super::*;
1801    use reqwest::header::{HeaderName, HeaderValue};
1802    use std::io::{Read, Write};
1803    use std::net::{TcpListener, TcpStream};
1804    #[cfg(unix)]
1805    use std::os::unix::fs::PermissionsExt;
1806    use std::sync::atomic::{AtomicU64, Ordering};
1807    use std::thread;
1808    use std::time::Instant;
1809
1810    static TEST_FILE_SEQUENCE: AtomicU64 = AtomicU64::new(0);
1811
1812    fn header_map(entries: &[(&'static str, &'static str)]) -> HeaderMap {
1813        let mut headers = HeaderMap::new();
1814        for (name, value) in entries {
1815            headers.insert(
1816                HeaderName::from_static(name),
1817                HeaderValue::from_static(value),
1818            );
1819        }
1820        headers
1821    }
1822
1823    fn start_http_server(
1824        status: &'static str,
1825        headers: &'static str,
1826        delay: Duration,
1827    ) -> (String, thread::JoinHandle<()>) {
1828        let listener = TcpListener::bind("127.0.0.1:0").expect("test listener should bind");
1829        let address = listener.local_addr().expect("listener should have address");
1830        let server = thread::spawn(move || {
1831            let mut stream = accept_test_connection(&listener);
1832            let mut request = [0_u8; 8_192];
1833            let bytes = stream.read(&mut request).expect("request should read");
1834            assert!(bytes > 0, "request should not be empty");
1835            thread::sleep(delay);
1836            let response = format!(
1837                "HTTP/1.1 {status}\r\ncontent-length: 2\r\nconnection: close\r\n{headers}\r\n{{}}"
1838            );
1839            let _ = stream.write_all(response.as_bytes());
1840        });
1841        (format!("http://{address}/v1/messages"), server)
1842    }
1843
1844    fn accept_test_connection(listener: &TcpListener) -> TcpStream {
1845        listener
1846            .set_nonblocking(true)
1847            .expect("test listener should become nonblocking");
1848        let deadline = Instant::now() + Duration::from_secs(5);
1849
1850        loop {
1851            match listener.accept() {
1852                Ok((stream, _)) => {
1853                    stream
1854                        .set_nonblocking(false)
1855                        .expect("test stream should become blocking");
1856                    stream
1857                        .set_read_timeout(Some(Duration::from_secs(5)))
1858                        .expect("test stream should have a read timeout");
1859                    stream
1860                        .set_write_timeout(Some(Duration::from_secs(5)))
1861                        .expect("test stream should have a write timeout");
1862                    return stream;
1863                }
1864                Err(error)
1865                    if error.kind() == std::io::ErrorKind::WouldBlock
1866                        && Instant::now() < deadline =>
1867                {
1868                    thread::sleep(Duration::from_millis(5));
1869                }
1870                Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
1871                    panic!("request did not connect before the test deadline");
1872                }
1873                Err(error) => panic!("request should connect: {error}"),
1874            }
1875        }
1876    }
1877
1878    fn temporary_claude_credentials() -> PathBuf {
1879        let path = std::env::temp_dir().join(format!(
1880            "agent-quota-test-{}-{}-{}.json",
1881            std::process::id(),
1882            now_epoch_ms(),
1883            TEST_FILE_SEQUENCE.fetch_add(1, Ordering::Relaxed)
1884        ));
1885        fs::write(&path, r#"{"claudeAiOauth":{"accessToken":"test-token"}}"#)
1886            .expect("credentials fixture should write");
1887        path
1888    }
1889
1890    #[test]
1891    fn exposes_versioned_provider_capabilities() {
1892        let capabilities = provider_capabilities();
1893        assert_eq!(capabilities.schema_version, CAPABILITIES_SCHEMA_VERSION);
1894        assert_eq!(
1895            capabilities.snapshot_schema_version,
1896            SNAPSHOT_SCHEMA_VERSION
1897        );
1898        assert_eq!(capabilities.providers.len(), 2);
1899
1900        let codex = &capabilities.providers[0];
1901        assert_eq!(codex.provider_id, ProviderId::Codex);
1902        assert_eq!(codex.probe_transport, ProviderProbeTransport::LocalProcess);
1903        assert!(!codex.submits_message);
1904        assert!(!codex.may_affect_quota_or_billing);
1905
1906        let claude = &capabilities.providers[1];
1907        assert_eq!(claude.provider_id, ProviderId::Claude);
1908        assert_eq!(claude.probe_transport, ProviderProbeTransport::RemoteApi);
1909        assert!(claude.submits_message);
1910        assert!(claude.may_affect_quota_or_billing);
1911    }
1912
1913    #[test]
1914    fn evaluates_typed_readiness_counts() {
1915        let profile = ProviderProfile::default_for_provider(ProviderId::Codex);
1916        let ready = success_snapshot(
1917            &profile,
1918            None,
1919            Some("fixture"),
1920            None,
1921            Vec::new(),
1922            None,
1923            QuotaState::Available,
1924        );
1925        let exhausted = success_snapshot(
1926            &profile,
1927            None,
1928            Some("fixture"),
1929            None,
1930            Vec::new(),
1931            None,
1932            QuotaState::Exhausted,
1933        );
1934        let failed = failed_snapshot(
1935            &profile,
1936            ProbeFailure::new(ProviderUsageErrorCode::Network, "offline", true),
1937        );
1938
1939        let any = ReadinessSummary::evaluate(
1940            &[ready.clone(), exhausted.clone(), failed.clone()],
1941            ReadinessPolicy::Any,
1942        );
1943        assert!(any.satisfied);
1944        assert_eq!(any.total_profiles, 3);
1945        assert_eq!(any.ready_profiles, 1);
1946        assert_eq!(any.not_ready_profiles, 2);
1947        assert_eq!(any.failed_probes, 1);
1948        assert_eq!(any.exhausted_profiles, 1);
1949
1950        let all = ReadinessSummary::evaluate(&[ready, exhausted, failed], ReadinessPolicy::All);
1951        assert!(!all.satisfied);
1952        assert!(!ReadinessSummary::evaluate(&[], ReadinessPolicy::All).satisfied);
1953    }
1954
1955    #[tokio::test]
1956    async fn direct_collections_report_live_metadata() {
1957        let invalid_profile = ProviderProfile {
1958            id: String::new(),
1959            ..ProviderProfile::default_for_provider(ProviderId::Codex)
1960        };
1961        let snapshot = AgentQuotaClient::new()
1962            .collect_profile_usage(invalid_profile)
1963            .await;
1964        let metadata = snapshot
1965            .collection
1966            .expect("collection metadata should exist");
1967        assert_eq!(metadata.freshness, ProviderUsageFreshness::Live);
1968        assert_eq!(metadata.cached_at_ms, None);
1969        assert_eq!(metadata.expires_at_ms, None);
1970    }
1971    #[test]
1972    fn parses_and_validates_profile_config() {
1973        let config = toml::from_str::<AgentQuotaConfig>(
1974            r#"
1975                [[profiles]]
1976                id = "claude-work"
1977                provider = "claude"
1978                label = "Claude Work"
1979                credentials_path = "C:/Users/example/.claude-work/.credentials.json"
1980
1981                [[profiles]]
1982                id = "codex-private"
1983                provider = "codex"
1984                label = "Codex Private"
1985
1986                [profiles.env]
1987                CODEX_HOME = "C:/Users/example/.codex-private"
1988            "#,
1989        )
1990        .expect("profile config should parse");
1991
1992        config.validate().expect("profile config should validate");
1993        assert_eq!(config.profiles.len(), 2);
1994        assert_eq!(config.profiles[0].provider, ProviderId::Claude);
1995    }
1996
1997    #[test]
1998    fn rejects_unknown_and_duplicate_profile_fields() {
1999        let unknown = toml::from_str::<AgentQuotaConfig>(
2000            r#"
2001                [[profiles]]
2002                id = "codex"
2003                provider = "codex"
2004                comand_path = "codex"
2005            "#,
2006        );
2007        assert!(unknown.is_err());
2008
2009        let duplicate = AgentQuotaConfig {
2010            profiles: vec![
2011                ProviderProfile::default_for_provider(ProviderId::Codex),
2012                ProviderProfile::default_for_provider(ProviderId::Codex),
2013            ],
2014        };
2015        assert!(duplicate.validate().is_err());
2016    }
2017
2018    #[test]
2019    fn empty_explicit_selection_does_not_expand_to_all() {
2020        assert!(CollectUsageOptions::providers([])
2021            .selected_profiles()
2022            .is_empty());
2023        assert!(CollectUsageOptions::profiles([])
2024            .selected_profiles()
2025            .is_empty());
2026        assert_eq!(CollectUsageOptions::all().selected_profiles().len(), 2);
2027    }
2028
2029    #[test]
2030    fn maps_codex_success_and_exhaustion() {
2031        let account = json!({
2032            "id": 2,
2033            "result": {
2034                "account": { "type": "chatgpt", "email": "user@example.com", "planType": "plus" }
2035            }
2036        });
2037        let rate_limits = json!({
2038            "id": 3,
2039            "result": {
2040                "rateLimits": {
2041                    "primary": { "usedPercent": 100, "windowDurationMins": 300, "resetsAt": 1_730_947_200 },
2042                    "secondary": { "usedPercent": 46, "windowDurationMins": 10080, "resetsAt": 1_731_206_400 },
2043                    "individualLimit": { "limit": "20.00", "used": "7.50", "remainingPercent": 63, "resetsAt": 1_731_206_400 },
2044                    "credits": { "balance": "12.50", "hasCredits": true, "unlimited": false },
2045                    "spendControlReached": false,
2046                    "rateLimitReachedType": "primary"
2047                },
2048                "rateLimitResetCredits": { "availableCount": 2 }
2049            }
2050        });
2051        let profile = ProviderProfile::default_for_provider(ProviderId::Codex);
2052        let snapshot = codex_snapshot_from_responses(&profile, &account, &rate_limits)
2053            .expect("response should map");
2054
2055        assert_eq!(
2056            snapshot.probe_status,
2057            ProbeStatus::Ok,
2058            "unexpected snapshot: {snapshot:?}"
2059        );
2060        assert_eq!(snapshot.quota_state, QuotaState::Exhausted);
2061        assert!(!snapshot.is_ready());
2062        assert_eq!(snapshot.account_label.as_deref(), Some("user@example.com"));
2063        assert_eq!(snapshot.windows.len(), 2);
2064        assert_eq!(snapshot.windows[0].label, "5h window");
2065        assert_eq!(snapshot.windows[1].label, "Weekly window");
2066        let billable = snapshot
2067            .billable_usage
2068            .as_ref()
2069            .expect("billable usage should map");
2070        assert_eq!(billable.used, "7.50");
2071        assert_eq!(billable.limit, "20.00");
2072        assert_eq!(billable.remaining_percent, 63.0);
2073        let credits = snapshot.credits.as_ref().expect("credits should map");
2074        assert_eq!(credits.balance.as_deref(), Some("12.50"));
2075        assert!(credits.has_credits);
2076        assert!(!credits.unlimited);
2077        assert_eq!(
2078            snapshot
2079                .rate_limit_reset_credits
2080                .as_ref()
2081                .map(|credits| credits.available_count),
2082            Some(2)
2083        );
2084        let value = serde_json::to_value(&snapshot).expect("snapshot should serialize");
2085        assert_eq!(value["billableUsage"]["remainingPercent"], 63.0);
2086        assert_eq!(value["credits"]["balance"], "12.50");
2087        assert_eq!(value["rateLimitResetCredits"]["availableCount"], 2);
2088    }
2089
2090    #[test]
2091    fn billable_exhaustion_blocks_codex_readiness() {
2092        let account = json!({
2093            "id": 2,
2094            "result": { "account": { "planType": "plus" } }
2095        });
2096        let rate_limits = json!({
2097            "id": 3,
2098            "result": {
2099                "rateLimits": {
2100                    "primary": { "usedPercent": 25, "windowDurationMins": 300 },
2101                    "individualLimit": {
2102                        "limit": "10.00",
2103                        "used": "10.00",
2104                        "remainingPercent": 0,
2105                        "resetsAt": 1_731_206_400
2106                    },
2107                    "spendControlReached": true
2108                }
2109            }
2110        });
2111        let snapshot = codex_snapshot_from_responses(
2112            &ProviderProfile::default_for_provider(ProviderId::Codex),
2113            &account,
2114            &rate_limits,
2115        )
2116        .expect("billable limit response should map");
2117
2118        assert_eq!(snapshot.quota_state, QuotaState::Exhausted);
2119        assert!(!snapshot.is_ready());
2120        assert_eq!(
2121            snapshot.message.as_deref(),
2122            Some("Billable usage limit reached")
2123        );
2124    }
2125
2126    #[test]
2127    fn rejects_codex_rpc_errors_and_missing_windows() {
2128        let profile = ProviderProfile::default_for_provider(ProviderId::Codex);
2129        let auth_error = json!({
2130            "id": 2,
2131            "error": { "code": -32000, "message": "Please sign in to Codex" }
2132        });
2133        let rate_limits = json!({ "id": 3, "result": { "rateLimits": {} } });
2134        let error = codex_snapshot_from_responses(&profile, &auth_error, &rate_limits)
2135            .expect_err("RPC error must not become an available snapshot");
2136        assert_eq!(
2137            error.error.code,
2138            ProviderUsageErrorCode::AuthenticationRequired
2139        );
2140
2141        let account = json!({ "id": 2, "result": { "account": {} } });
2142        let error = codex_snapshot_from_responses(&profile, &account, &rate_limits)
2143            .expect_err("missing windows must be rejected");
2144        assert_eq!(error.error.code, ProviderUsageErrorCode::InvalidResponse);
2145    }
2146
2147    #[tokio::test]
2148    async fn codex_production_process_path_maps_mock_app_server() {
2149        let extension = if cfg!(windows) { "cmd" } else { "sh" };
2150        let path = std::env::temp_dir().join(format!(
2151            "agent-quota-mock-codex-{}-{}.{}",
2152            std::process::id(),
2153            now_epoch_ms(),
2154            extension
2155        ));
2156        #[cfg(windows)]
2157        let script = concat!(
2158            "@echo off\r\n",
2159            "echo {\"id\":2,\"result\":{\"account\":{\"email\":\"test@example.com\",\"planType\":\"plus\"}}}\r\n",
2160            "echo {\"id\":3,\"result\":{\"rateLimits\":{\"primary\":{\"usedPercent\":25,\"windowDurationMins\":300,\"resetsAt\":1730947200}}}}\r\n",
2161            "ping -n 3 127.0.0.1 >nul\r\n"
2162        );
2163        #[cfg(unix)]
2164        let script = concat!(
2165            "#!/bin/sh\n",
2166            "printf '%s\\n' '{\"id\":2,\"result\":{\"account\":{\"email\":\"test@example.com\",\"planType\":\"plus\"}}}'\n",
2167            "printf '%s\\n' '{\"id\":3,\"result\":{\"rateLimits\":{\"primary\":{\"usedPercent\":25,\"windowDurationMins\":300,\"resetsAt\":1730947200}}}}'\n",
2168            "sleep 2\n"
2169        );
2170        fs::write(&path, script).expect("mock Codex executable should write");
2171        #[cfg(unix)]
2172        {
2173            let mut permissions = fs::metadata(&path)
2174                .expect("mock metadata should read")
2175                .permissions();
2176            permissions.set_mode(0o755);
2177            fs::set_permissions(&path, permissions).expect("mock should become executable");
2178        }
2179
2180        let profile = ProviderProfile {
2181            command_path: Some(path.clone()),
2182            ..ProviderProfile::default_for_provider(ProviderId::Codex)
2183        };
2184        let snapshot = AgentQuotaClient::new().collect_profile_usage(profile).await;
2185        let _ = fs::remove_file(path);
2186
2187        assert_eq!(
2188            snapshot.probe_status,
2189            ProbeStatus::Ok,
2190            "unexpected snapshot: {snapshot:?}"
2191        );
2192        assert_eq!(snapshot.quota_state, QuotaState::Available);
2193        assert_eq!(snapshot.account_label.as_deref(), Some("test@example.com"));
2194        assert_eq!(snapshot.windows[0].remaining_percent, Some(75.0));
2195    }
2196
2197    #[test]
2198    fn parses_claude_headers_and_http_failures() {
2199        let headers = header_map(&[
2200            ("anthropic-ratelimit-unified-5h-utilization", "0.42"),
2201            ("anthropic-ratelimit-unified-5h-reset", "1730947200"),
2202            ("anthropic-ratelimit-unified-7d-utilization", "0.68"),
2203            ("anthropic-ratelimit-unified-7d-reset", "1731206400000"),
2204            ("anthropic-ratelimit-unified-status", "active"),
2205        ]);
2206        let profile = ProviderProfile::default_for_provider(ProviderId::Claude);
2207        let snapshot = claude_snapshot_from_headers(&profile, &headers, StatusCode::OK)
2208            .expect("headers should map");
2209        assert_eq!(snapshot.probe_status, ProbeStatus::Ok);
2210        assert_eq!(snapshot.quota_state, QuotaState::Available);
2211        assert!(snapshot.is_ready());
2212        assert_eq!(snapshot.windows[0].remaining_percent, Some(58.0));
2213        assert_eq!(
2214            snapshot.windows[1].resets_at_epoch_seconds,
2215            Some(1_731_206_400)
2216        );
2217
2218        let error =
2219            claude_snapshot_from_headers(&profile, &HeaderMap::new(), StatusCode::UNAUTHORIZED)
2220                .expect_err("401 should be authentication failure");
2221        assert_eq!(
2222            error.error.code,
2223            ProviderUsageErrorCode::AuthenticationRequired
2224        );
2225    }
2226
2227    #[test]
2228    fn reads_claude_oauth_token_from_supported_shapes() {
2229        let credentials = json!({
2230            "claudeAiOauth": {
2231                "accessToken": "oauth-token"
2232            }
2233        });
2234        assert_eq!(find_claude_oauth_token(&credentials), Some("oauth-token"));
2235    }
2236
2237    #[tokio::test]
2238    async fn cache_is_profile_keyed_and_does_not_cache_failures() {
2239        let cache = ProviderUsageCache::default();
2240        let client = AgentQuotaClient::new();
2241        let codex_profile = ProviderProfile::default_for_provider(ProviderId::Codex);
2242        cache.inner.lock().await.insert(
2243            codex_profile.clone(),
2244            CachedProviderUsage {
2245                snapshot: success_snapshot(
2246                    &codex_profile,
2247                    None,
2248                    Some("fixture"),
2249                    None,
2250                    vec![ProviderUsageWindow {
2251                        kind: ProviderUsageWindowKind::FiveHour,
2252                        label: "5h window".to_owned(),
2253                        used_percent: Some(25.0),
2254                        remaining_percent: Some(75.0),
2255                        window_minutes: Some(300),
2256                        resets_at_epoch_seconds: None,
2257                        detail: None,
2258                    }],
2259                    None,
2260                    QuotaState::Available,
2261                ),
2262                cached_at_ms: now_epoch_ms(),
2263            },
2264        );
2265        let cached = cache
2266            .collect(
2267                &client,
2268                CollectUsageOptions::profiles([codex_profile.clone()]),
2269                false,
2270            )
2271            .await;
2272        let metadata = cached[0]
2273            .collection
2274            .as_ref()
2275            .expect("cache hit should expose metadata");
2276        assert_eq!(metadata.freshness, ProviderUsageFreshness::Cached);
2277        assert!(metadata.cached_at_ms.is_some());
2278        assert!(metadata.expires_at_ms > metadata.cached_at_ms);
2279        let empty = CollectUsageOptions::profiles([]);
2280        assert!(cache
2281            .collect(&client, empty.clone(), false)
2282            .await
2283            .is_empty());
2284
2285        let invalid_profile = ProviderProfile {
2286            id: String::new(),
2287            ..ProviderProfile::default_for_provider(ProviderId::Codex)
2288        };
2289        let invalid = CollectUsageOptions::profiles([invalid_profile]);
2290        let first = cache.collect(&client, invalid.clone(), false).await;
2291        let second = cache.collect(&client, invalid, false).await;
2292        assert_eq!(first[0].probe_status, ProbeStatus::InvalidConfiguration);
2293        assert_eq!(second[0].probe_status, ProbeStatus::InvalidConfiguration);
2294        let guard = cache.inner.lock().await;
2295        assert_eq!(guard.len(), 1);
2296        assert!(guard.contains_key(&codex_profile));
2297    }
2298
2299    #[test]
2300    fn serializes_versioned_snapshot_contract() {
2301        let profile = ProviderProfile::default_for_provider(ProviderId::Codex);
2302        let snapshot = success_snapshot(
2303            &profile,
2304            Some("user@example.com".to_owned()),
2305            Some("fixture"),
2306            Some("Plus".to_owned()),
2307            vec![ProviderUsageWindow {
2308                kind: ProviderUsageWindowKind::FiveHour,
2309                label: "5h window".to_owned(),
2310                used_percent: Some(42.0),
2311                remaining_percent: Some(58.0),
2312                window_minutes: Some(300),
2313                resets_at_epoch_seconds: Some(1_730_947_200),
2314                detail: None,
2315            }],
2316            None,
2317            QuotaState::Available,
2318        );
2319        let value = serde_json::to_value(snapshot).expect("snapshot should serialize");
2320        assert_eq!(value["schemaVersion"], SNAPSHOT_SCHEMA_VERSION);
2321        assert_eq!(value["providerId"], "codex");
2322        assert_eq!(value["probeStatus"], "ok");
2323        assert_eq!(value["quotaState"], "available");
2324        assert!(value["observedAtMs"].as_i64().is_some());
2325        assert_eq!(value["windows"][0]["resetsAtEpochSeconds"], 1_730_947_200);
2326    }
2327
2328    #[test]
2329    fn shared_v1_fixture_matches_the_public_model() {
2330        let raw = include_str!("../fixtures/provider-usage-v1.json");
2331        let snapshots =
2332            serde_json::from_str::<Vec<ProviderUsageSnapshot>>(raw).expect("fixture should parse");
2333        assert_eq!(snapshots.len(), 2);
2334        assert_eq!(snapshots[0].schema_version, SNAPSHOT_SCHEMA_VERSION);
2335        assert_eq!(snapshots[0].probe_status, ProbeStatus::Ok);
2336        let collection = snapshots[0]
2337            .collection
2338            .as_ref()
2339            .expect("fixture should include optional collection metadata");
2340        assert_eq!(collection.freshness, ProviderUsageFreshness::Live);
2341        assert_eq!(collection.probe_duration_ms, 120);
2342        assert_eq!(
2343            snapshots[0]
2344                .billable_usage
2345                .as_ref()
2346                .map(|usage| usage.remaining_percent),
2347            Some(63.0)
2348        );
2349        assert_eq!(
2350            snapshots[0]
2351                .credits
2352                .as_ref()
2353                .and_then(|credits| credits.balance.as_deref()),
2354            Some("12.50")
2355        );
2356        assert_eq!(
2357            snapshots[0]
2358                .rate_limit_reset_credits
2359                .as_ref()
2360                .map(|credits| credits.available_count),
2361            Some(2)
2362        );
2363        assert!(snapshots[1].collection.is_none());
2364        assert!(snapshots[1].billable_usage.is_none());
2365        assert_eq!(
2366            snapshots[1].error.as_ref().map(|error| error.code),
2367            Some(ProviderUsageErrorCode::CredentialsUnavailable)
2368        );
2369    }
2370
2371    #[tokio::test]
2372    async fn claude_production_path_uses_injected_endpoint_and_checks_headers() {
2373        let listener = TcpListener::bind("127.0.0.1:0").expect("test listener should bind");
2374        let address = listener.local_addr().expect("listener should have address");
2375        let server = thread::spawn(move || {
2376            let mut stream = accept_test_connection(&listener);
2377            let mut request = [0_u8; 8_192];
2378            let bytes = stream.read(&mut request).expect("request should read");
2379            let request = String::from_utf8_lossy(&request[..bytes]);
2380            assert!(request.starts_with("POST /v1/messages "));
2381            assert!(request
2382                .to_ascii_lowercase()
2383                .contains("user-agent: agent-quota/"));
2384            let response = concat!(
2385                "HTTP/1.1 200 OK\r\n",
2386                "content-length: 2\r\n",
2387                "connection: close\r\n",
2388                "anthropic-ratelimit-unified-5h-utilization: 0.25\r\n",
2389                "anthropic-ratelimit-unified-5h-reset: 1730947200\r\n",
2390                "anthropic-ratelimit-unified-status: active\r\n",
2391                "\r\n{}"
2392            );
2393            stream
2394                .write_all(response.as_bytes())
2395                .expect("response should write");
2396        });
2397
2398        let credentials_path = temporary_claude_credentials();
2399        let profile = ProviderProfile {
2400            credentials_path: Some(credentials_path.clone()),
2401            ..ProviderProfile::default_for_provider(ProviderId::Claude)
2402        };
2403        let client =
2404            AgentQuotaClient::new().with_claude_api_url(format!("http://{address}/v1/messages"));
2405        let snapshot = client.collect_profile_usage(profile).await;
2406        let _ = fs::remove_file(credentials_path);
2407        server.join().expect("server should finish");
2408
2409        assert_eq!(snapshot.probe_status, ProbeStatus::Ok);
2410        assert_eq!(snapshot.quota_state, QuotaState::Available);
2411        assert_eq!(snapshot.windows[0].remaining_percent, Some(75.0));
2412    }
2413
2414    #[tokio::test]
2415    async fn claude_production_path_classifies_http_and_timeout_failures() {
2416        let credentials_path = temporary_claude_credentials();
2417        let profile = ProviderProfile {
2418            credentials_path: Some(credentials_path.clone()),
2419            ..ProviderProfile::default_for_provider(ProviderId::Claude)
2420        };
2421
2422        let (url, server) = start_http_server("401 Unauthorized", "", Duration::from_millis(0));
2423        let snapshot = AgentQuotaClient::new()
2424            .with_claude_api_url(url)
2425            .collect_profile_usage(profile.clone())
2426            .await;
2427        server.join().expect("401 server should finish");
2428        assert_eq!(snapshot.probe_status, ProbeStatus::AuthenticationRequired);
2429        assert_eq!(
2430            snapshot.error.as_ref().map(|error| error.code),
2431            Some(ProviderUsageErrorCode::AuthenticationRequired)
2432        );
2433
2434        let headers = concat!(
2435            "anthropic-ratelimit-unified-5h-utilization: 1.0\r\n",
2436            "anthropic-ratelimit-unified-5h-reset: 1730947200\r\n",
2437            "anthropic-ratelimit-unified-status: exhausted\r\n"
2438        );
2439        let (url, server) =
2440            start_http_server("429 Too Many Requests", headers, Duration::from_millis(0));
2441        let snapshot = AgentQuotaClient::new()
2442            .with_claude_api_url(url)
2443            .collect_profile_usage(profile.clone())
2444            .await;
2445        server.join().expect("429 server should finish");
2446        assert_eq!(snapshot.probe_status, ProbeStatus::Ok);
2447        assert_eq!(snapshot.quota_state, QuotaState::Exhausted);
2448
2449        let (url, server) = start_http_server("200 OK", "", Duration::from_millis(0));
2450        let snapshot = AgentQuotaClient::new()
2451            .with_claude_api_url(url)
2452            .collect_profile_usage(profile.clone())
2453            .await;
2454        server.join().expect("no-header server should finish");
2455        assert_eq!(snapshot.probe_status, ProbeStatus::Unsupported);
2456
2457        let (url, server) = start_http_server("200 OK", "", Duration::from_millis(500));
2458        let snapshot = AgentQuotaClient::new()
2459            .with_claude_api_url(url)
2460            .with_timeouts(DEFAULT_CODEX_TIMEOUT, Duration::from_millis(250))
2461            .collect_profile_usage(profile)
2462            .await;
2463        server.join().expect("timeout server should finish");
2464        let _ = fs::remove_file(credentials_path);
2465        assert_eq!(snapshot.probe_status, ProbeStatus::TransientError);
2466        assert_eq!(
2467            snapshot.error.as_ref().map(|error| error.code),
2468            Some(ProviderUsageErrorCode::Timeout)
2469        );
2470    }
2471}