1#![warn(missing_docs)]
2mod 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);
29pub 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
39pub const SNAPSHOT_SCHEMA_VERSION: u32 = 1;
41pub const CAPABILITIES_SCHEMA_VERSION: u32 = 1;
43
44pub type AgentQuotaResult<T> = Result<T, AgentQuotaError>;
46
47#[derive(Debug)]
49pub enum AgentQuotaError {
50 ConfigRead {
52 path: PathBuf,
54 source: std::io::Error,
56 },
57 ConfigParse {
59 path: PathBuf,
61 message: String,
63 },
64 ConfigValidation {
66 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#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
100#[serde(rename_all = "snake_case")]
101pub enum ProviderId {
102 Codex,
104 Claude,
106}
107
108impl ProviderId {
109 pub fn as_str(self) -> &'static str {
111 match self {
112 Self::Codex => "codex",
113 Self::Claude => "claude",
114 }
115 }
116
117 pub fn label(self) -> &'static str {
119 match self {
120 Self::Codex => "Codex",
121 Self::Claude => "Claude Code",
122 }
123 }
124
125 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
137#[serde(rename_all = "snake_case")]
138pub enum ProviderProbeTransport {
139 LocalProcess,
141 RemoteApi,
143}
144
145#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
147#[serde(rename_all = "snake_case")]
148pub enum ProviderQuotaSource {
149 CodexAppServer,
151 AnthropicRateLimitHeaders,
153}
154
155#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
157#[serde(rename_all = "snake_case")]
158pub enum ProviderCredentialSource {
159 CodexCliSession,
161 ClaudeCodeOauthFile,
163}
164
165#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
167#[serde(rename_all = "camelCase")]
168pub struct ProviderCapability {
169 pub provider_id: ProviderId,
171 pub provider_name: String,
173 pub probe_transport: ProviderProbeTransport,
175 pub quota_source: ProviderQuotaSource,
177 pub credential_source: ProviderCredentialSource,
179 pub submits_message: bool,
181 pub may_affect_quota_or_billing: bool,
183 pub default_cache_ttl_seconds: u64,
185 pub probe_impact: String,
187}
188
189#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
191#[serde(rename_all = "camelCase")]
192pub struct ProviderCapabilitiesDocument {
193 pub schema_version: u32,
195 pub snapshot_schema_version: u32,
197 pub providers: Vec<ProviderCapability>,
199}
200
201pub fn provider_capability(provider: ProviderId) -> ProviderCapability {
203 providers::adapter(provider).capability()
204}
205
206pub 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#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
219#[serde(rename_all = "snake_case", deny_unknown_fields)]
220pub struct ProviderProfile {
221 pub id: String,
223 pub provider: ProviderId,
225 pub label: Option<String>,
227 pub command_path: Option<PathBuf>,
229 pub credentials_path: Option<PathBuf>,
231 #[serde(default)]
233 pub env: BTreeMap<String, String>,
234 #[serde(default = "default_enabled")]
236 pub enabled: bool,
237}
238
239impl ProviderProfile {
240 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 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 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#[derive(Debug, Clone, Default, PartialEq, Eq, Serialize, Deserialize)]
315#[serde(rename_all = "snake_case", deny_unknown_fields)]
316pub struct AgentQuotaConfig {
317 #[serde(default)]
319 pub profiles: Vec<ProviderProfile>,
320}
321
322impl AgentQuotaConfig {
323 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 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 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#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
372enum UsageSelection {
373 All,
374 Providers(Vec<ProviderId>),
375 Profiles(Vec<ProviderProfile>),
376}
377
378#[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 pub fn all() -> Self {
393 Self {
394 selection: UsageSelection::All,
395 }
396 }
397
398 pub fn providers(providers: impl IntoIterator<Item = ProviderId>) -> Self {
402 Self {
403 selection: UsageSelection::Providers(providers.into_iter().collect()),
404 }
405 }
406
407 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
438#[serde(rename_all = "snake_case")]
439pub enum ProbeStatus {
440 Ok,
442 AuthenticationRequired,
444 Unsupported,
446 TransientError,
448 InvalidResponse,
450 InvalidConfiguration,
452}
453
454#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
456#[serde(rename_all = "snake_case")]
457pub enum QuotaState {
458 Available,
460 Exhausted,
462 Unknown,
464}
465
466#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
468#[serde(rename_all = "snake_case")]
469pub enum ReadinessPolicy {
470 Any,
472 All,
474}
475
476#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
478#[serde(rename_all = "camelCase")]
479pub struct ReadinessSummary {
480 pub policy: ReadinessPolicy,
482 pub satisfied: bool,
484 pub total_profiles: usize,
486 pub ready_profiles: usize,
488 pub not_ready_profiles: usize,
490 pub failed_probes: usize,
492 pub exhausted_profiles: usize,
494}
495
496impl ReadinessSummary {
497 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
532#[serde(rename_all = "snake_case")]
533pub enum ProviderUsageErrorCode {
534 InvalidConfiguration,
536 ExecutableNotFound,
538 CredentialsUnavailable,
540 AuthenticationRequired,
542 Timeout,
544 Network,
546 ProviderResponse,
548 InvalidResponse,
550 Unsupported,
552}
553
554#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
556#[serde(rename_all = "camelCase")]
557pub struct ProviderUsageError {
558 pub code: ProviderUsageErrorCode,
560 pub message: String,
562 pub retryable: bool,
564}
565
566#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
568#[serde(rename_all = "snake_case")]
569pub enum ProviderUsageFreshness {
570 Live,
572 Cached,
574}
575
576#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
578#[serde(rename_all = "camelCase")]
579pub struct ProviderUsageCollectionMetadata {
580 pub freshness: ProviderUsageFreshness,
582 pub probe_duration_ms: u64,
584 pub cached_at_ms: Option<i64>,
586 pub expires_at_ms: Option<i64>,
588}
589
590#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
592#[serde(rename_all = "camelCase")]
593pub struct ProviderBillableUsage {
594 pub used: String,
596 pub limit: String,
598 pub remaining_percent: f64,
600 pub resets_at_epoch_seconds: i64,
602}
603
604#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
606#[serde(rename_all = "camelCase")]
607pub struct ProviderCredits {
608 pub balance: Option<String>,
610 pub has_credits: bool,
612 pub unlimited: bool,
614}
615
616#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
618#[serde(rename_all = "camelCase")]
619pub struct ProviderRateLimitResetCredits {
620 pub available_count: u64,
622}
623
624#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
626#[serde(rename_all = "camelCase")]
627pub struct ProviderUsageSnapshot {
628 pub schema_version: u32,
630 pub provider_id: ProviderId,
632 pub provider_name: String,
634 pub profile_id: String,
636 pub profile_name: String,
638 pub account_label: Option<String>,
640 pub source: Option<String>,
642 pub probe_status: ProbeStatus,
644 pub quota_state: QuotaState,
646 pub plan: Option<String>,
648 pub windows: Vec<ProviderUsageWindow>,
650 #[serde(default, skip_serializing_if = "Option::is_none")]
652 pub billable_usage: Option<ProviderBillableUsage>,
653 #[serde(default, skip_serializing_if = "Option::is_none")]
655 pub credits: Option<ProviderCredits>,
656 #[serde(default, skip_serializing_if = "Option::is_none")]
658 pub rate_limit_reset_credits: Option<ProviderRateLimitResetCredits>,
659 pub observed_at_ms: i64,
661 #[serde(default, skip_serializing_if = "Option::is_none")]
663 pub collection: Option<ProviderUsageCollectionMetadata>,
664 pub message: Option<String>,
666 pub error: Option<ProviderUsageError>,
668}
669
670impl ProviderUsageSnapshot {
671 pub fn is_ready(&self) -> bool {
677 self.probe_status == ProbeStatus::Ok && self.quota_state == QuotaState::Available
678 }
679}
680
681#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
683#[serde(rename_all = "camelCase")]
684pub struct ProviderUsageWindow {
685 pub kind: ProviderUsageWindowKind,
687 pub label: String,
689 pub used_percent: Option<f64>,
691 pub remaining_percent: Option<f64>,
693 pub window_minutes: Option<u64>,
695 pub resets_at_epoch_seconds: Option<i64>,
697 pub detail: Option<String>,
699}
700
701#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
703#[serde(rename_all = "snake_case")]
704pub enum ProviderUsageWindowKind {
705 FiveHour,
707 Weekly,
709 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#[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 pub fn new() -> Self {
774 Self::default()
775 }
776
777 pub fn with_http_client(mut self, http_client: reqwest::Client) -> Self {
782 self.http_client = http_client;
783 self
784 }
785
786 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 pub fn with_claude_model(mut self, model: impl Into<String>) -> Self {
797 self.claude_model = model.into();
798 self
799 }
800
801 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 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 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
1244pub 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#[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 pub fn with_ttl(ttl: Duration) -> Self {
1700 Self {
1701 ttl,
1702 ..Self::default()
1703 }
1704 }
1705
1706 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 pub async fn clear(&self) {
1772 self.inner.lock().await.clear();
1773 }
1774}
1775
1776pub async fn collect_provider_usage() -> Vec<ProviderUsageSnapshot> {
1778 collect_usage(CollectUsageOptions::all()).await
1779}
1780
1781pub async fn collect_usage(options: CollectUsageOptions) -> Vec<ProviderUsageSnapshot> {
1783 AgentQuotaClient::new().collect_usage(options).await
1784}
1785
1786pub 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
1793pub 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}