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

codex_cli/rate_limits/
client.rs
