codex_cli/rate_limits/
client.rs

1use anyhow::{Context, Result};
2use reqwest::blocking::Client;
3use serde_json::Value;
4use std::path::{Path, PathBuf};
5use std::time::Duration;
6
7use crate::auth::remote;
8use crate::json;
9use crate::paths;
10
11pub struct UsageResponse {
12    pub body: String,
13    pub json: Value,
14}
15
16pub struct UsageRequest {
17    pub target_file: PathBuf,
18    pub refresh_on_401: bool,
19    pub suppress_auth_refresh_output: bool,
20    pub base_url: String,
21    pub connect_timeout_seconds: u64,
22    pub max_time_seconds: u64,
23}
24
25pub fn fetch_usage(request: &UsageRequest) -> Result<UsageResponse> {
26    let (access_token, account_id) = read_tokens(&request.target_file)?;
27    let mut response = send_request(request, &access_token, account_id.as_deref())?;
28
29    if response.status == 401 && request.refresh_on_401 {
30        let refreshed_tokens =
31            refresh_target(&request.target_file, request.suppress_auth_refresh_output)
32                .or_else(|| read_tokens(&request.target_file).ok());
33        if let Some((access_token, account_id)) = refreshed_tokens {
34            response = send_request(request, &access_token, account_id.as_deref())?;
35        }
36    }
37
38    if response.status != 200 {
39        let preview = response
40            .body
41            .chars()
42            .take(200)
43            .collect::<String>()
44            .replace(['\n', '\r'], " ");
45        if preview.is_empty() {
46            anyhow::bail!(
47                "codex-rate-limits: GET {} failed (HTTP {})",
48                response.url,
49                response.status
50            );
51        }
52        anyhow::bail!(
53            "codex-rate-limits: GET {} failed (HTTP {})\ncodex-rate-limits: body: {}",
54            response.url,
55            response.status,
56            preview
57        );
58    }
59
60    let json: Value =
61        serde_json::from_str(&response.body).context("invalid JSON from usage endpoint")?;
62
63    Ok(UsageResponse {
64        body: response.body,
65        json,
66    })
67}
68
69pub fn read_tokens(target_file: &Path) -> Result<(String, Option<String>)> {
70    let value = json::read_json(target_file)?;
71    read_tokens_from_value(&value)
72}
73
74fn read_tokens_from_value(value: &Value) -> Result<(String, Option<String>)> {
75    let access_token = json::string_at(value, &["tokens", "access_token"])
76        .or_else(|| json::string_at(value, &["access_token"]))
77        .ok_or_else(|| anyhow::anyhow!("missing access_token"))?;
78    let account_id = json::string_at(value, &["tokens", "account_id"])
79        .or_else(|| json::string_at(value, &["account_id"]));
80    Ok((access_token, account_id))
81}
82
83fn send_request(
84    request: &UsageRequest,
85    access_token: &str,
86    account_id: Option<&str>,
87) -> Result<HttpResponse> {
88    let client = Client::builder()
89        .connect_timeout(Duration::from_secs(request.connect_timeout_seconds))
90        .timeout(Duration::from_secs(request.max_time_seconds))
91        .build()?;
92
93    let url = format!("{}/wham/usage", request.base_url.trim_end_matches('/'));
94    let mut req = client
95        .get(&url)
96        .header("Authorization", format!("Bearer {}", access_token))
97        .header("Accept", "application/json")
98        .header("User-Agent", "codex-cli");
99    if let Some(account_id) = account_id {
100        req = req.header("ChatGPT-Account-Id", account_id);
101    }
102
103    let resp = req.send();
104    let resp = match resp {
105        Ok(value) => value,
106        Err(_) => {
107            anyhow::bail!("codex-rate-limits: request failed: {}", url);
108        }
109    };
110
111    let status = resp.status().as_u16();
112    let body = resp.text().unwrap_or_default();
113    Ok(HttpResponse { status, body, url })
114}
115
116fn refresh_target(target_file: &Path, suppress_output: bool) -> Option<(String, Option<String>)> {
117    if let Some(auth_file) = paths::resolve_auth_file()
118        && auth_file == target_file
119    {
120        let _ = if suppress_output {
121            crate::auth::refresh::run_silent(&[])
122        } else {
123            crate::auth::refresh::run(&[])
124        };
125        return read_tokens(target_file).ok();
126    }
127
128    if let Some(secret_dir) = paths::resolve_secret_dir()
129        && let Some(file_name) = target_file.file_name().and_then(|n| n.to_str())
130    {
131        let path = secret_dir.join(file_name);
132        if path == target_file {
133            if let Ok(Some(payload)) = remote::export_access_only_for_target_from_env(target_file)
134                && let Ok(tokens) = read_tokens_from_value(&payload.auth)
135            {
136                return Some(tokens);
137            }
138            let args = [file_name.to_string()];
139            let _ = if suppress_output {
140                crate::auth::refresh::run_silent(&args)
141            } else {
142                crate::auth::refresh::run(&args)
143            };
144            return read_tokens(target_file).ok();
145        }
146    }
147
148    None
149}
150
151struct HttpResponse {
152    status: u16,
153    body: String,
154    url: String,
155}
codex_cli/rate_limits/client.rs

codex_cli/rate_limits/
client.rs
