1use serde::{Deserialize, Serialize};
2use std::collections::HashMap;
3use std::time::{SystemTime, UNIX_EPOCH};
4
5use crate::config;
6
7#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
8#[serde(rename_all = "snake_case")]
9pub enum CredentialStatus {
10 Valid,
11 Expired,
12 NotFound,
13 ParseError,
14}
15
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
17#[serde(rename_all = "camelCase")]
18pub struct QuotaTier {
19 pub name: String,
20 pub utilization: f64,
21 pub resets_at: Option<String>,
22}
23
24#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
25#[serde(rename_all = "camelCase")]
26pub struct ExtraUsage {
27 pub is_enabled: bool,
28 pub monthly_limit: Option<f64>,
29 pub used_credits: Option<f64>,
30 pub utilization: Option<f64>,
31 pub currency: Option<String>,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
35#[serde(rename_all = "camelCase")]
36pub struct SubscriptionQuota {
37 pub tool: String,
38 pub credential_status: CredentialStatus,
39 pub credential_message: Option<String>,
40 pub success: bool,
41 pub tiers: Vec<QuotaTier>,
42 pub extra_usage: Option<ExtraUsage>,
43 pub error: Option<String>,
44 pub queried_at: Option<i64>,
45}
46
47impl SubscriptionQuota {
48 pub fn not_found(tool: &str) -> Self {
49 Self {
50 tool: tool.to_string(),
51 credential_status: CredentialStatus::NotFound,
52 credential_message: None,
53 success: false,
54 tiers: vec![],
55 extra_usage: None,
56 error: None,
57 queried_at: None,
58 }
59 }
60
61 pub(crate) fn error(tool: &str, status: CredentialStatus, message: String) -> Self {
62 Self {
63 tool: tool.to_string(),
64 credential_status: status,
65 credential_message: Some(message.clone()),
66 success: false,
67 tiers: vec![],
68 extra_usage: None,
69 error: Some(message),
70 queried_at: Some(now_millis()),
71 }
72 }
73}
74
75#[derive(Deserialize)]
76struct ClaudeOAuthEntry {
77 #[serde(rename = "accessToken")]
78 access_token: Option<String>,
79 #[serde(rename = "expiresAt")]
80 expires_at: Option<serde_json::Value>,
81}
82
83fn read_claude_credentials() -> (Option<String>, CredentialStatus, Option<String>) {
84 #[cfg(target_os = "macos")]
85 {
86 if let Some(result) = read_claude_credentials_from_keychain() {
87 return result;
88 }
89 }
90
91 read_claude_credentials_from_file()
92}
93
94#[cfg(target_os = "macos")]
95fn read_claude_credentials_from_keychain(
96) -> Option<(Option<String>, CredentialStatus, Option<String>)> {
97 let output = std::process::Command::new("security")
98 .args([
99 "find-generic-password",
100 "-s",
101 "Claude Code-credentials",
102 "-w",
103 ])
104 .output()
105 .ok()?;
106
107 if !output.status.success() {
108 return None;
109 }
110
111 let json_str = String::from_utf8(output.stdout).ok()?;
112 let json_str = json_str.trim();
113 if json_str.is_empty() {
114 return None;
115 }
116
117 Some(parse_claude_credentials_json(json_str))
118}
119
120fn read_claude_credentials_from_file() -> (Option<String>, CredentialStatus, Option<String>) {
121 let cred_path = config::get_claude_config_dir().join(".credentials.json");
122
123 if !cred_path.exists() {
124 return (None, CredentialStatus::NotFound, None);
125 }
126
127 let content = match std::fs::read_to_string(&cred_path) {
128 Ok(content) => content,
129 Err(error) => {
130 return (
131 None,
132 CredentialStatus::ParseError,
133 Some(format!("Failed to read credentials file: {error}")),
134 );
135 }
136 };
137
138 parse_claude_credentials_json(&content)
139}
140
141fn parse_claude_credentials_json(
142 content: &str,
143) -> (Option<String>, CredentialStatus, Option<String>) {
144 let parsed: serde_json::Value = match serde_json::from_str(content) {
145 Ok(value) => value,
146 Err(error) => {
147 return (
148 None,
149 CredentialStatus::ParseError,
150 Some(format!("Failed to parse credentials JSON: {error}")),
151 );
152 }
153 };
154
155 let entry_value = parsed
156 .get("claudeAiOauth")
157 .or_else(|| parsed.get("claude.ai_oauth"));
158
159 let Some(entry_value) = entry_value else {
160 return (
161 None,
162 CredentialStatus::ParseError,
163 Some("No OAuth entry found in credentials".to_string()),
164 );
165 };
166
167 let entry: ClaudeOAuthEntry = match serde_json::from_value(entry_value.clone()) {
168 Ok(entry) => entry,
169 Err(error) => {
170 return (
171 None,
172 CredentialStatus::ParseError,
173 Some(format!("Failed to parse OAuth entry: {error}")),
174 );
175 }
176 };
177
178 let access_token = match entry.access_token {
179 Some(token) if !token.is_empty() => token,
180 _ => {
181 return (
182 None,
183 CredentialStatus::ParseError,
184 Some("accessToken is empty or missing".to_string()),
185 );
186 }
187 };
188
189 if let Some(expires_at) = entry.expires_at {
190 if is_token_expired(&expires_at) {
191 return (
192 Some(access_token),
193 CredentialStatus::Expired,
194 Some("OAuth token has expired".to_string()),
195 );
196 }
197 }
198
199 (Some(access_token), CredentialStatus::Valid, None)
200}
201
202fn is_token_expired(expires_at: &serde_json::Value) -> bool {
203 let now_secs = SystemTime::now()
204 .duration_since(UNIX_EPOCH)
205 .unwrap_or_default()
206 .as_secs();
207
208 match expires_at {
209 serde_json::Value::Number(value) => value.as_u64().is_some_and(|timestamp| {
210 let timestamp_secs = if timestamp > 1_000_000_000_000 {
211 timestamp / 1000
212 } else {
213 timestamp
214 };
215 timestamp_secs < now_secs
216 }),
217 serde_json::Value::String(value) => {
218 if let Ok(dt) = chrono::DateTime::parse_from_rfc3339(value) {
219 (dt.timestamp() as u64) < now_secs
220 } else if let Ok(dt) =
221 chrono::NaiveDateTime::parse_from_str(value, "%Y-%m-%dT%H:%M:%S%.f")
222 {
223 (dt.and_utc().timestamp() as u64) < now_secs
224 } else {
225 false
226 }
227 }
228 _ => false,
229 }
230}
231
232#[derive(Deserialize)]
233struct ApiUsageWindow {
234 utilization: Option<f64>,
235 resets_at: Option<String>,
236}
237
238#[derive(Deserialize)]
239struct ApiExtraUsage {
240 is_enabled: Option<bool>,
241 monthly_limit: Option<f64>,
242 used_credits: Option<f64>,
243 utilization: Option<f64>,
244 currency: Option<String>,
245}
246
247const KNOWN_TIERS: &[&str] = &[
248 "five_hour",
249 "seven_day",
250 "seven_day_opus",
251 "seven_day_sonnet",
252];
253
254async fn query_claude_quota(access_token: &str) -> SubscriptionQuota {
255 let client = crate::proxy::http_client::get();
256
257 let response = match client
258 .get("https://api.anthropic.com/api/oauth/usage")
259 .header("Authorization", format!("Bearer {access_token}"))
260 .header("anthropic-beta", "oauth-2025-04-20")
261 .header("Accept", "application/json")
262 .timeout(std::time::Duration::from_secs(10))
263 .send()
264 .await
265 {
266 Ok(response) => response,
267 Err(error) => {
268 return SubscriptionQuota::error(
269 "claude",
270 CredentialStatus::Valid,
271 format!("Network error: {error}"),
272 );
273 }
274 };
275
276 let status = response.status();
277 if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
278 return SubscriptionQuota::error(
279 "claude",
280 CredentialStatus::Expired,
281 format!("Authentication failed (HTTP {status}). Please re-login with Claude CLI."),
282 );
283 }
284
285 if !status.is_success() {
286 let body = response.text().await.unwrap_or_default();
287 return SubscriptionQuota::error(
288 "claude",
289 CredentialStatus::Valid,
290 format!("API error (HTTP {status}): {body}"),
291 );
292 }
293
294 let body: serde_json::Value = match response.json().await {
295 Ok(body) => body,
296 Err(error) => {
297 return SubscriptionQuota::error(
298 "claude",
299 CredentialStatus::Valid,
300 format!("Failed to parse API response: {error}"),
301 );
302 }
303 };
304
305 let mut tiers = Vec::new();
306 for &tier_name in KNOWN_TIERS {
307 if let Some(window) = body.get(tier_name) {
308 if let Ok(window) = serde_json::from_value::<ApiUsageWindow>(window.clone()) {
309 if let Some(utilization) = window.utilization {
310 tiers.push(QuotaTier {
311 name: tier_name.to_string(),
312 utilization,
313 resets_at: window.resets_at,
314 });
315 }
316 }
317 }
318 }
319
320 if let Some(object) = body.as_object() {
321 for (key, value) in object {
322 if key == "extra_usage" || KNOWN_TIERS.contains(&key.as_str()) {
323 continue;
324 }
325 if let Ok(window) = serde_json::from_value::<ApiUsageWindow>(value.clone()) {
326 if let Some(utilization) = window.utilization {
327 tiers.push(QuotaTier {
328 name: key.clone(),
329 utilization,
330 resets_at: window.resets_at,
331 });
332 }
333 }
334 }
335 }
336
337 let extra_usage = body.get("extra_usage").and_then(|value| {
338 serde_json::from_value::<ApiExtraUsage>(value.clone())
339 .ok()
340 .map(|extra| ExtraUsage {
341 is_enabled: extra.is_enabled.unwrap_or(false),
342 monthly_limit: extra.monthly_limit,
343 used_credits: extra.used_credits,
344 utilization: extra.utilization,
345 currency: extra.currency,
346 })
347 });
348
349 SubscriptionQuota {
350 tool: "claude".to_string(),
351 credential_status: CredentialStatus::Valid,
352 credential_message: None,
353 success: true,
354 tiers,
355 extra_usage,
356 error: None,
357 queried_at: Some(now_millis()),
358 }
359}
360
361#[derive(Deserialize)]
362struct CodexAuthJson {
363 auth_mode: Option<String>,
364 tokens: Option<CodexTokens>,
365 last_refresh: Option<String>,
366}
367
368#[derive(Deserialize)]
369struct CodexTokens {
370 access_token: Option<String>,
371 account_id: Option<String>,
372}
373
374type CodexCredentials = (
375 Option<String>,
376 Option<String>,
377 CredentialStatus,
378 Option<String>,
379);
380
381fn read_codex_credentials() -> CodexCredentials {
382 #[cfg(target_os = "macos")]
383 {
384 if let Some(result) = read_codex_credentials_from_keychain() {
385 return result;
386 }
387 }
388
389 read_codex_credentials_from_file()
390}
391
392#[cfg(target_os = "macos")]
393fn read_codex_credentials_from_keychain() -> Option<CodexCredentials> {
394 let output = std::process::Command::new("security")
395 .args(["find-generic-password", "-s", "Codex Auth", "-w"])
396 .output()
397 .ok()?;
398
399 if !output.status.success() {
400 return None;
401 }
402
403 let json_str = String::from_utf8(output.stdout).ok()?;
404 let json_str = json_str.trim();
405 if json_str.is_empty() {
406 return None;
407 }
408
409 Some(parse_codex_credentials_json(json_str))
410}
411
412fn read_codex_credentials_from_file() -> CodexCredentials {
413 let auth_path = crate::codex_config::get_codex_auth_path();
414
415 if !auth_path.exists() {
416 return (None, None, CredentialStatus::NotFound, None);
417 }
418
419 let content = match std::fs::read_to_string(&auth_path) {
420 Ok(content) => content,
421 Err(error) => {
422 return (
423 None,
424 None,
425 CredentialStatus::ParseError,
426 Some(format!("Failed to read Codex auth file: {error}")),
427 );
428 }
429 };
430
431 parse_codex_credentials_json(&content)
432}
433
434fn parse_codex_credentials_json(content: &str) -> CodexCredentials {
435 let auth: CodexAuthJson = match serde_json::from_str(content) {
436 Ok(auth) => auth,
437 Err(error) => {
438 return (
439 None,
440 None,
441 CredentialStatus::ParseError,
442 Some(format!("Failed to parse Codex auth JSON: {error}")),
443 );
444 }
445 };
446
447 if auth.auth_mode.as_deref() != Some("chatgpt") {
448 return (
449 None,
450 None,
451 CredentialStatus::NotFound,
452 Some("Codex not using OAuth mode".to_string()),
453 );
454 }
455
456 let tokens = match auth.tokens {
457 Some(tokens) => tokens,
458 None => {
459 return (
460 None,
461 None,
462 CredentialStatus::ParseError,
463 Some("No tokens in Codex auth".to_string()),
464 );
465 }
466 };
467
468 let access_token = match tokens.access_token {
469 Some(token) if !token.is_empty() => token,
470 _ => {
471 return (
472 None,
473 None,
474 CredentialStatus::ParseError,
475 Some("access_token is empty or missing".to_string()),
476 );
477 }
478 };
479
480 if let Some(ref last_refresh) = auth.last_refresh {
481 if is_codex_token_stale(last_refresh) {
482 return (
483 Some(access_token),
484 tokens.account_id,
485 CredentialStatus::Expired,
486 Some("Codex token may be stale (>8 days since last refresh)".to_string()),
487 );
488 }
489 }
490
491 (
492 Some(access_token),
493 tokens.account_id,
494 CredentialStatus::Valid,
495 None,
496 )
497}
498
499fn is_codex_token_stale(last_refresh: &str) -> bool {
500 let now_secs = SystemTime::now()
501 .duration_since(UNIX_EPOCH)
502 .unwrap_or_default()
503 .as_secs();
504
505 chrono::DateTime::parse_from_rfc3339(last_refresh).is_ok_and(|dt| {
506 let age_secs = now_secs.saturating_sub(dt.timestamp() as u64);
507 age_secs > 8 * 24 * 3600
508 })
509}
510
511#[derive(Deserialize)]
512struct CodexRateLimitWindow {
513 used_percent: Option<f64>,
514 limit_window_seconds: Option<i64>,
515 reset_at: Option<i64>,
516}
517
518#[derive(Deserialize)]
519struct CodexRateLimit {
520 primary_window: Option<CodexRateLimitWindow>,
521 secondary_window: Option<CodexRateLimitWindow>,
522}
523
524#[derive(Deserialize)]
525struct CodexUsageResponse {
526 rate_limit: Option<CodexRateLimit>,
527}
528
529fn now_millis() -> i64 {
530 std::time::SystemTime::now()
531 .duration_since(std::time::UNIX_EPOCH)
532 .unwrap_or_default()
533 .as_millis() as i64
534}
535
536fn window_seconds_to_tier_name(secs: i64) -> String {
537 match secs {
538 18_000 => "five_hour".to_string(),
539 604_800 => "seven_day".to_string(),
540 secs => {
541 let hours = secs / 3600;
542 if hours >= 24 {
543 format!("{}_day", hours / 24)
544 } else {
545 format!("{}_hour", hours)
546 }
547 }
548 }
549}
550
551fn unix_ts_to_iso(ts: i64) -> Option<String> {
552 chrono::DateTime::from_timestamp(ts, 0).map(|dt| dt.to_rfc3339())
553}
554
555pub(crate) async fn query_codex_quota(
556 access_token: &str,
557 account_id: Option<&str>,
558 tool_label: &str,
559 expired_message: &str,
560) -> SubscriptionQuota {
561 let client = crate::proxy::http_client::get();
562
563 let mut request = client
564 .get("https://chatgpt.com/backend-api/wham/usage")
565 .header("Authorization", format!("Bearer {access_token}"))
566 .header("User-Agent", "codex-cli")
567 .header("Accept", "application/json");
568
569 if let Some(account_id) = account_id {
570 request = request.header("ChatGPT-Account-Id", account_id);
571 }
572
573 let response = match request
574 .timeout(std::time::Duration::from_secs(10))
575 .send()
576 .await
577 {
578 Ok(response) => response,
579 Err(error) => {
580 return SubscriptionQuota::error(
581 tool_label,
582 CredentialStatus::Valid,
583 format!("Network error: {error}"),
584 );
585 }
586 };
587
588 let status = response.status();
589 if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
590 return SubscriptionQuota::error(
591 tool_label,
592 CredentialStatus::Expired,
593 format!("{expired_message} (HTTP {status})"),
594 );
595 }
596
597 if !status.is_success() {
598 let body = response.text().await.unwrap_or_default();
599 return SubscriptionQuota::error(
600 tool_label,
601 CredentialStatus::Valid,
602 format!("API error (HTTP {status}): {body}"),
603 );
604 }
605
606 let body: CodexUsageResponse = match response.json().await {
607 Ok(body) => body,
608 Err(error) => {
609 return SubscriptionQuota::error(
610 tool_label,
611 CredentialStatus::Valid,
612 format!("Failed to parse API response: {error}"),
613 );
614 }
615 };
616
617 let mut tiers = Vec::new();
618 if let Some(rate_limit) = body.rate_limit {
619 for window in [rate_limit.primary_window, rate_limit.secondary_window]
620 .into_iter()
621 .flatten()
622 {
623 if let Some(utilization) = window.used_percent {
624 tiers.push(QuotaTier {
625 name: window
626 .limit_window_seconds
627 .map(window_seconds_to_tier_name)
628 .unwrap_or_else(|| "unknown".to_string()),
629 utilization,
630 resets_at: window.reset_at.and_then(unix_ts_to_iso),
631 });
632 }
633 }
634 }
635
636 SubscriptionQuota {
637 tool: tool_label.to_string(),
638 credential_status: CredentialStatus::Valid,
639 credential_message: None,
640 success: true,
641 tiers,
642 extra_usage: None,
643 error: None,
644 queried_at: Some(now_millis()),
645 }
646}
647
648#[derive(Deserialize)]
649struct GeminiOAuthCredsFile {
650 #[serde(alias = "accessToken")]
651 access_token: Option<String>,
652 #[serde(alias = "refreshToken")]
653 refresh_token: Option<String>,
654 #[serde(alias = "expiryDate", alias = "expiresAt")]
655 expiry_date: Option<i64>,
656 #[serde(alias = "clientId")]
657 client_id: Option<String>,
658 #[serde(alias = "clientSecret")]
659 client_secret: Option<String>,
660}
661
662#[derive(Debug, Clone)]
663struct GeminiOAuthClient {
664 client_id: String,
665 client_secret: String,
666}
667
668type GeminiCredentials = (
669 Option<String>,
670 Option<String>,
671 Option<GeminiOAuthClient>,
672 CredentialStatus,
673 Option<String>,
674);
675
676fn gemini_oauth_client_from_options(
677 client_id: Option<String>,
678 client_secret: Option<String>,
679) -> Option<GeminiOAuthClient> {
680 let client_id = client_id?.trim().to_string();
681 let client_secret = client_secret?.trim().to_string();
682 if client_id.is_empty() || client_secret.is_empty() {
683 return None;
684 }
685 Some(GeminiOAuthClient {
686 client_id,
687 client_secret,
688 })
689}
690
691fn gemini_oauth_client_from_json(value: &serde_json::Value) -> Option<GeminiOAuthClient> {
692 let client_id = value
693 .get("client_id")
694 .or_else(|| value.get("clientId"))
695 .and_then(|value| value.as_str())
696 .map(String::from);
697 let client_secret = value
698 .get("client_secret")
699 .or_else(|| value.get("clientSecret"))
700 .and_then(|value| value.as_str())
701 .map(String::from);
702 gemini_oauth_client_from_options(client_id, client_secret)
703}
704
705fn read_gemini_credentials() -> GeminiCredentials {
706 #[cfg(target_os = "macos")]
707 {
708 if let Some(result) = read_gemini_credentials_from_keychain() {
709 return result;
710 }
711 }
712
713 read_gemini_credentials_from_file()
714}
715
716#[cfg(target_os = "macos")]
717fn read_gemini_credentials_from_keychain() -> Option<GeminiCredentials> {
718 let output = std::process::Command::new("security")
719 .args([
720 "find-generic-password",
721 "-s",
722 "gemini-cli-oauth",
723 "-a",
724 "main-account",
725 "-w",
726 ])
727 .output()
728 .ok()?;
729
730 if !output.status.success() {
731 return None;
732 }
733
734 let json_str = String::from_utf8(output.stdout).ok()?;
735 let json_str = json_str.trim();
736 if json_str.is_empty() {
737 return None;
738 }
739
740 Some(parse_gemini_keychain_json(json_str))
741}
742
743#[cfg(target_os = "macos")]
744fn parse_gemini_keychain_json(content: &str) -> GeminiCredentials {
745 let parsed: serde_json::Value = match serde_json::from_str(content) {
746 Ok(value) => value,
747 Err(error) => {
748 return (
749 None,
750 None,
751 None,
752 CredentialStatus::ParseError,
753 Some(format!("Failed to parse Gemini keychain JSON: {error}")),
754 );
755 }
756 };
757
758 let Some(token) = parsed.get("token") else {
759 return parse_gemini_file_json(content);
760 };
761
762 let access_token = token
763 .get("accessToken")
764 .and_then(|value| value.as_str())
765 .map(String::from);
766 let refresh_token = token
767 .get("refreshToken")
768 .and_then(|value| value.as_str())
769 .map(String::from);
770 let expires_at = token.get("expiresAt").and_then(|value| value.as_i64());
771 let oauth_client =
772 gemini_oauth_client_from_json(token).or_else(|| gemini_oauth_client_from_json(&parsed));
773
774 match access_token {
775 Some(token) if !token.is_empty() => {
776 if expires_at.is_some_and(|expires_at| expires_at < now_millis()) {
777 return (
778 Some(token),
779 refresh_token,
780 oauth_client,
781 CredentialStatus::Expired,
782 Some("Gemini access token has expired".to_string()),
783 );
784 }
785
786 (
787 Some(token),
788 refresh_token,
789 oauth_client,
790 CredentialStatus::Valid,
791 None,
792 )
793 }
794 _ => (
795 None,
796 refresh_token,
797 oauth_client,
798 CredentialStatus::ParseError,
799 Some("accessToken is empty or missing".to_string()),
800 ),
801 }
802}
803
804fn read_gemini_credentials_from_file() -> GeminiCredentials {
805 let cred_path = crate::gemini_config::get_gemini_dir().join("oauth_creds.json");
806
807 if !cred_path.exists() {
808 return (None, None, None, CredentialStatus::NotFound, None);
809 }
810
811 let content = match std::fs::read_to_string(&cred_path) {
812 Ok(content) => content,
813 Err(error) => {
814 return (
815 None,
816 None,
817 None,
818 CredentialStatus::ParseError,
819 Some(format!("Failed to read Gemini credentials: {error}")),
820 );
821 }
822 };
823
824 parse_gemini_file_json(&content)
825}
826
827fn parse_gemini_file_json(content: &str) -> GeminiCredentials {
828 let creds: GeminiOAuthCredsFile = match serde_json::from_str(content) {
829 Ok(creds) => creds,
830 Err(error) => {
831 return (
832 None,
833 None,
834 None,
835 CredentialStatus::ParseError,
836 Some(format!("Failed to parse Gemini credentials: {error}")),
837 );
838 }
839 };
840
841 let access_token = match creds.access_token {
842 Some(token) if !token.is_empty() => token,
843 _ => {
844 return (
845 None,
846 creds.refresh_token,
847 gemini_oauth_client_from_options(creds.client_id, creds.client_secret),
848 CredentialStatus::ParseError,
849 Some("access_token is empty or missing".to_string()),
850 );
851 }
852 };
853
854 if creds
855 .expiry_date
856 .is_some_and(|expires_at| expires_at < now_millis())
857 {
858 return (
859 Some(access_token),
860 creds.refresh_token,
861 gemini_oauth_client_from_options(creds.client_id, creds.client_secret),
862 CredentialStatus::Expired,
863 Some("Gemini access token has expired".to_string()),
864 );
865 }
866
867 (
868 Some(access_token),
869 creds.refresh_token,
870 gemini_oauth_client_from_options(creds.client_id, creds.client_secret),
871 CredentialStatus::Valid,
872 None,
873 )
874}
875
876async fn refresh_gemini_token(
877 refresh_token: &str,
878 oauth_client: &GeminiOAuthClient,
879) -> Option<String> {
880 let client = crate::proxy::http_client::get();
881
882 let response = client
883 .post("https://oauth2.googleapis.com/token")
884 .form(&[
885 ("client_id", oauth_client.client_id.as_str()),
886 ("client_secret", oauth_client.client_secret.as_str()),
887 ("refresh_token", refresh_token),
888 ("grant_type", "refresh_token"),
889 ])
890 .timeout(std::time::Duration::from_secs(10))
891 .send()
892 .await
893 .ok()?;
894
895 if !response.status().is_success() {
896 return None;
897 }
898
899 let body: serde_json::Value = response.json().await.ok()?;
900 body.get("access_token")?.as_str().map(String::from)
901}
902
903#[derive(Deserialize)]
904struct GeminiLoadCodeAssistResponse {
905 #[serde(rename = "cloudaicompanionProject")]
906 cloudaicompanion_project: Option<serde_json::Value>,
907}
908
909#[derive(Deserialize)]
910struct GeminiBucketInfo {
911 #[serde(rename = "remainingFraction")]
912 remaining_fraction: Option<f64>,
913 #[serde(rename = "resetTime")]
914 reset_time: Option<String>,
915 #[serde(rename = "modelId")]
916 model_id: Option<String>,
917}
918
919#[derive(Deserialize)]
920struct GeminiQuotaResponse {
921 buckets: Option<Vec<GeminiBucketInfo>>,
922}
923
924fn extract_project_id(value: &serde_json::Value) -> Option<String> {
925 match value {
926 serde_json::Value::String(value) => Some(value.clone()),
927 serde_json::Value::Object(object) => object
928 .get("id")
929 .or_else(|| object.get("projectId"))
930 .and_then(|value| value.as_str())
931 .map(String::from),
932 _ => None,
933 }
934}
935
936fn classify_gemini_model(model_id: &str) -> &str {
937 if model_id.contains("flash-lite") {
938 "gemini_flash_lite"
939 } else if model_id.contains("flash") {
940 "gemini_flash"
941 } else if model_id.contains("pro") {
942 "gemini_pro"
943 } else {
944 model_id
945 }
946}
947
948async fn query_gemini_quota(access_token: &str) -> SubscriptionQuota {
949 let client = crate::proxy::http_client::get();
950
951 let load_response = match client
952 .post("https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist")
953 .header("Authorization", format!("Bearer {access_token}"))
954 .header("Content-Type", "application/json")
955 .json(&serde_json::json!({
956 "metadata": {
957 "ideType": "GEMINI_CLI",
958 "pluginType": "GEMINI"
959 }
960 }))
961 .timeout(std::time::Duration::from_secs(10))
962 .send()
963 .await
964 {
965 Ok(response) => response,
966 Err(error) => {
967 return SubscriptionQuota::error(
968 "gemini",
969 CredentialStatus::Valid,
970 format!("Network error (loadCodeAssist): {error}"),
971 );
972 }
973 };
974
975 let load_status = load_response.status();
976 if load_status == reqwest::StatusCode::UNAUTHORIZED
977 || load_status == reqwest::StatusCode::FORBIDDEN
978 {
979 return SubscriptionQuota::error(
980 "gemini",
981 CredentialStatus::Expired,
982 format!("Authentication failed (HTTP {load_status}). Please re-login with Gemini CLI."),
983 );
984 }
985
986 if !load_status.is_success() {
987 let body = load_response.text().await.unwrap_or_default();
988 return SubscriptionQuota::error(
989 "gemini",
990 CredentialStatus::Valid,
991 format!("loadCodeAssist failed (HTTP {load_status}): {body}"),
992 );
993 }
994
995 let load_body: GeminiLoadCodeAssistResponse = match load_response.json().await {
996 Ok(body) => body,
997 Err(error) => {
998 return SubscriptionQuota::error(
999 "gemini",
1000 CredentialStatus::Valid,
1001 format!("Failed to parse loadCodeAssist response: {error}"),
1002 );
1003 }
1004 };
1005
1006 let project_id = load_body
1007 .cloudaicompanion_project
1008 .as_ref()
1009 .and_then(extract_project_id);
1010
1011 let mut quota_body = serde_json::json!({});
1012 if let Some(project_id) = project_id {
1013 quota_body["project"] = serde_json::Value::String(project_id);
1014 }
1015
1016 let quota_response = match client
1017 .post("https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota")
1018 .header("Authorization", format!("Bearer {access_token}"))
1019 .header("Content-Type", "application/json")
1020 .json("a_body)
1021 .timeout(std::time::Duration::from_secs(10))
1022 .send()
1023 .await
1024 {
1025 Ok(response) => response,
1026 Err(error) => {
1027 return SubscriptionQuota::error(
1028 "gemini",
1029 CredentialStatus::Valid,
1030 format!("Network error (retrieveUserQuota): {error}"),
1031 );
1032 }
1033 };
1034
1035 let quota_status = quota_response.status();
1036 if quota_status == reqwest::StatusCode::UNAUTHORIZED
1037 || quota_status == reqwest::StatusCode::FORBIDDEN
1038 {
1039 return SubscriptionQuota::error(
1040 "gemini",
1041 CredentialStatus::Expired,
1042 format!("Authentication failed (HTTP {quota_status})."),
1043 );
1044 }
1045
1046 if !quota_status.is_success() {
1047 let body = quota_response.text().await.unwrap_or_default();
1048 return SubscriptionQuota::error(
1049 "gemini",
1050 CredentialStatus::Valid,
1051 format!("retrieveUserQuota failed (HTTP {quota_status}): {body}"),
1052 );
1053 }
1054
1055 let quota_data: GeminiQuotaResponse = match quota_response.json().await {
1056 Ok(body) => body,
1057 Err(error) => {
1058 return SubscriptionQuota::error(
1059 "gemini",
1060 CredentialStatus::Valid,
1061 format!("Failed to parse quota response: {error}"),
1062 );
1063 }
1064 };
1065
1066 let mut category_map: HashMap<String, (f64, Option<String>)> = HashMap::new();
1067
1068 if let Some(buckets) = quota_data.buckets {
1069 for bucket in buckets {
1070 let model_id = bucket.model_id.as_deref().unwrap_or("unknown");
1071 let category = classify_gemini_model(model_id).to_string();
1072 let remaining = bucket.remaining_fraction.unwrap_or(1.0).clamp(0.0, 1.0);
1073 let entry = category_map
1074 .entry(category)
1075 .or_insert((remaining, bucket.reset_time.clone()));
1076
1077 if remaining < entry.0 {
1078 entry.0 = remaining;
1079 if bucket.reset_time.is_some() {
1080 entry.1.clone_from(&bucket.reset_time);
1081 }
1082 }
1083 }
1084 }
1085
1086 let sort_order = |name: &str| -> usize {
1087 match name {
1088 "gemini_pro" => 0,
1089 "gemini_flash" => 1,
1090 "gemini_flash_lite" => 2,
1091 _ => 3,
1092 }
1093 };
1094
1095 let mut tiers: Vec<QuotaTier> = category_map
1096 .into_iter()
1097 .map(|(name, (remaining, reset_time))| QuotaTier {
1098 name,
1099 utilization: (1.0 - remaining) * 100.0,
1100 resets_at: reset_time,
1101 })
1102 .collect();
1103 tiers.sort_by_key(|tier| sort_order(&tier.name));
1104
1105 SubscriptionQuota {
1106 tool: "gemini".to_string(),
1107 credential_status: CredentialStatus::Valid,
1108 credential_message: None,
1109 success: true,
1110 tiers,
1111 extra_usage: None,
1112 error: None,
1113 queried_at: Some(now_millis()),
1114 }
1115}
1116
1117pub async fn get_subscription_quota(tool: &str) -> Result<SubscriptionQuota, String> {
1118 match tool {
1119 "claude" => {
1120 let (token, status, message) = read_claude_credentials();
1121 match status {
1122 CredentialStatus::NotFound => Ok(SubscriptionQuota::not_found("claude")),
1123 CredentialStatus::ParseError => Ok(SubscriptionQuota::error(
1124 "claude",
1125 CredentialStatus::ParseError,
1126 message.unwrap_or_else(|| "Failed to parse credentials".to_string()),
1127 )),
1128 CredentialStatus::Expired => {
1129 if let Some(token) = token {
1130 let result = query_claude_quota(&token).await;
1131 if result.success {
1132 return Ok(result);
1133 }
1134 }
1135 Ok(SubscriptionQuota::error(
1136 "claude",
1137 CredentialStatus::Expired,
1138 message.unwrap_or_else(|| "OAuth token has expired".to_string()),
1139 ))
1140 }
1141 CredentialStatus::Valid => {
1142 let Some(token) = token else {
1143 return Ok(SubscriptionQuota::error(
1144 "claude",
1145 CredentialStatus::ParseError,
1146 "accessToken is empty or missing".to_string(),
1147 ));
1148 };
1149 Ok(query_claude_quota(&token).await)
1150 }
1151 }
1152 }
1153 "codex" => {
1154 let (token, account_id, status, message) = read_codex_credentials();
1155 match status {
1156 CredentialStatus::NotFound => Ok(SubscriptionQuota::not_found("codex")),
1157 CredentialStatus::ParseError => Ok(SubscriptionQuota::error(
1158 "codex",
1159 CredentialStatus::ParseError,
1160 message.unwrap_or_else(|| "Failed to parse credentials".to_string()),
1161 )),
1162 CredentialStatus::Expired => {
1163 if let Some(token) = token {
1164 let result = query_codex_quota(
1165 &token,
1166 account_id.as_deref(),
1167 "codex",
1168 "Authentication failed. Please re-login with Codex CLI.",
1169 )
1170 .await;
1171 if result.success {
1172 return Ok(result);
1173 }
1174 }
1175 Ok(SubscriptionQuota::error(
1176 "codex",
1177 CredentialStatus::Expired,
1178 message.unwrap_or_else(|| "Codex OAuth token may be stale".to_string()),
1179 ))
1180 }
1181 CredentialStatus::Valid => {
1182 let Some(token) = token else {
1183 return Ok(SubscriptionQuota::error(
1184 "codex",
1185 CredentialStatus::ParseError,
1186 "access_token is empty or missing".to_string(),
1187 ));
1188 };
1189 Ok(query_codex_quota(
1190 &token,
1191 account_id.as_deref(),
1192 "codex",
1193 "Authentication failed. Please re-login with Codex CLI.",
1194 )
1195 .await)
1196 }
1197 }
1198 }
1199 "gemini" => {
1200 let (token, refresh_token, oauth_client, status, message) = read_gemini_credentials();
1201 match status {
1202 CredentialStatus::NotFound => Ok(SubscriptionQuota::not_found("gemini")),
1203 CredentialStatus::ParseError => Ok(SubscriptionQuota::error(
1204 "gemini",
1205 CredentialStatus::ParseError,
1206 message.unwrap_or_else(|| "Failed to parse credentials".to_string()),
1207 )),
1208 CredentialStatus::Expired => {
1209 if let (Some(refresh_token), Some(oauth_client)) =
1210 (refresh_token.as_deref(), oauth_client.as_ref())
1211 {
1212 if let Some(new_token) =
1213 refresh_gemini_token(refresh_token, oauth_client).await
1214 {
1215 return Ok(query_gemini_quota(&new_token).await);
1216 }
1217 }
1218 if let Some(ref token) = token {
1219 let result = query_gemini_quota(token).await;
1220 if result.success {
1221 return Ok(result);
1222 }
1223 }
1224 Ok(SubscriptionQuota::error(
1225 "gemini",
1226 CredentialStatus::Expired,
1227 message.unwrap_or_else(|| "Gemini OAuth token has expired".to_string()),
1228 ))
1229 }
1230 CredentialStatus::Valid => {
1231 let Some(token) = token else {
1232 return Ok(SubscriptionQuota::error(
1233 "gemini",
1234 CredentialStatus::ParseError,
1235 "access_token is empty or missing".to_string(),
1236 ));
1237 };
1238 Ok(query_gemini_quota(&token).await)
1239 }
1240 }
1241 }
1242 _ => Ok(SubscriptionQuota::not_found(tool)),
1243 }
1244}
1245
1246#[cfg(test)]
1247mod tests {
1248 use super::*;
1249
1250 #[test]
1251 fn parse_claude_credentials_accepts_current_and_legacy_keys() {
1252 let current =
1253 r#"{"claudeAiOauth":{"accessToken":"tok-current","expiresAt":4102444800000}}"#;
1254 let legacy = r#"{"claude.ai_oauth":{"accessToken":"tok-legacy","expiresAt":"2100-01-01T00:00:00Z"}}"#;
1255
1256 assert_eq!(
1257 parse_claude_credentials_json(current),
1258 (
1259 Some("tok-current".to_string()),
1260 CredentialStatus::Valid,
1261 None
1262 )
1263 );
1264 assert_eq!(
1265 parse_claude_credentials_json(legacy),
1266 (
1267 Some("tok-legacy".to_string()),
1268 CredentialStatus::Valid,
1269 None
1270 )
1271 );
1272 }
1273
1274 #[test]
1275 fn parse_claude_credentials_marks_expired_token() {
1276 let content = r#"{"claudeAiOauth":{"accessToken":"tok","expiresAt":1}}"#;
1277
1278 let (token, status, message) = parse_claude_credentials_json(content);
1279
1280 assert_eq!(token, Some("tok".to_string()));
1281 assert_eq!(status, CredentialStatus::Expired);
1282 assert_eq!(message, Some("OAuth token has expired".to_string()));
1283 }
1284
1285 #[test]
1286 fn parse_codex_credentials_requires_chatgpt_auth_mode() {
1287 let content =
1288 r#"{"auth_mode":"apikey","tokens":{"access_token":"tok","account_id":"acc"}}"#;
1289
1290 let (token, account_id, status, message) = parse_codex_credentials_json(content);
1291
1292 assert_eq!(token, None);
1293 assert_eq!(account_id, None);
1294 assert_eq!(status, CredentialStatus::NotFound);
1295 assert_eq!(message, Some("Codex not using OAuth mode".to_string()));
1296 }
1297
1298 #[test]
1299 fn parse_codex_credentials_returns_account_id_for_chatgpt_mode() {
1300 let content = r#"{
1301 "auth_mode":"chatgpt",
1302 "tokens":{"access_token":"tok","account_id":"acc"},
1303 "last_refresh":"2100-01-01T00:00:00Z"
1304 }"#;
1305
1306 let (token, account_id, status, message) = parse_codex_credentials_json(content);
1307
1308 assert_eq!(token, Some("tok".to_string()));
1309 assert_eq!(account_id, Some("acc".to_string()));
1310 assert_eq!(status, CredentialStatus::Valid);
1311 assert_eq!(message, None);
1312 }
1313
1314 #[test]
1315 fn parse_codex_credentials_marks_stale_token() {
1316 let content = r#"{
1317 "auth_mode":"chatgpt",
1318 "tokens":{"access_token":"tok","account_id":"acc"},
1319 "last_refresh":"2000-01-01T00:00:00Z"
1320 }"#;
1321
1322 let (token, account_id, status, message) = parse_codex_credentials_json(content);
1323
1324 assert_eq!(token, Some("tok".to_string()));
1325 assert_eq!(account_id, Some("acc".to_string()));
1326 assert_eq!(status, CredentialStatus::Expired);
1327 assert_eq!(
1328 message,
1329 Some("Codex token may be stale (>8 days since last refresh)".to_string())
1330 );
1331 }
1332
1333 #[test]
1334 fn parse_gemini_file_credentials_returns_refresh_token() {
1335 let content =
1336 r#"{"access_token":"tok","refresh_token":"refresh","expiry_date":4102444800000}"#;
1337
1338 let (token, refresh_token, oauth_client, status, message) = parse_gemini_file_json(content);
1339
1340 assert_eq!(token, Some("tok".to_string()));
1341 assert_eq!(refresh_token, Some("refresh".to_string()));
1342 assert!(oauth_client.is_none());
1343 assert_eq!(status, CredentialStatus::Valid);
1344 assert_eq!(message, None);
1345 }
1346
1347 #[test]
1348 fn parse_gemini_file_credentials_reads_refresh_client_fields() {
1349 let content = r#"{"access_token":"tok","refresh_token":"refresh","client_id":"client-id","client_secret":"client-secret","expiry_date":1}"#;
1350
1351 let (_, _, oauth_client, status, _) = parse_gemini_file_json(content);
1352
1353 let oauth_client = oauth_client.expect("oauth client fields should be parsed");
1354 assert_eq!(oauth_client.client_id, "client-id");
1355 assert_eq!(oauth_client.client_secret, "client-secret");
1356 assert_eq!(status, CredentialStatus::Expired);
1357 }
1358
1359 #[test]
1360 fn parse_gemini_file_credentials_marks_expired_token() {
1361 let content = r#"{"access_token":"tok","refresh_token":"refresh","expiry_date":1}"#;
1362
1363 let (token, refresh_token, oauth_client, status, message) = parse_gemini_file_json(content);
1364
1365 assert_eq!(token, Some("tok".to_string()));
1366 assert_eq!(refresh_token, Some("refresh".to_string()));
1367 assert!(oauth_client.is_none());
1368 assert_eq!(status, CredentialStatus::Expired);
1369 assert_eq!(message, Some("Gemini access token has expired".to_string()));
1370 }
1371
1372 #[test]
1373 fn gemini_helpers_match_upstream_shapes() {
1374 assert_eq!(
1375 extract_project_id(&serde_json::json!({"projectId": "project-1"})),
1376 Some("project-1".to_string())
1377 );
1378 assert_eq!(
1379 extract_project_id(&serde_json::json!({"id": "project-2"})),
1380 Some("project-2".to_string())
1381 );
1382 assert_eq!(
1383 extract_project_id(&serde_json::json!("project-3")),
1384 Some("project-3".to_string())
1385 );
1386
1387 assert_eq!(classify_gemini_model("gemini-3-pro"), "gemini_pro");
1388 assert_eq!(classify_gemini_model("gemini-3-flash"), "gemini_flash");
1389 assert_eq!(
1390 classify_gemini_model("gemini-3-flash-lite"),
1391 "gemini_flash_lite"
1392 );
1393 }
1394
1395 #[test]
1396 fn subscription_quota_not_found_matches_upstream_shape() {
1397 let quota = SubscriptionQuota::not_found("codex_oauth");
1398 assert_eq!(quota.tool, "codex_oauth");
1399 assert_eq!(quota.credential_status, CredentialStatus::NotFound);
1400 assert!(!quota.success);
1401 assert!(quota.tiers.is_empty());
1402 assert!(quota.queried_at.is_none());
1403 }
1404
1405 #[test]
1406 fn window_seconds_to_tier_name_matches_known_windows() {
1407 assert_eq!(window_seconds_to_tier_name(18_000), "five_hour");
1408 assert_eq!(window_seconds_to_tier_name(604_800), "seven_day");
1409 assert_eq!(window_seconds_to_tier_name(7_200), "2_hour");
1410 assert_eq!(window_seconds_to_tier_name(172_800), "2_day");
1411 }
1412}