use crate::auth;
use std::time::Duration;
use reqwest::blocking::{Client, Response};
use reqwest::{StatusCode, redirect::Policy};
use serde::Deserialize;
use crate::auth::ProviderCredential;
use crate::config::McPaths;
use crate::http_body::read_bounded_response_text;
const CODEX_USAGE_URL: &str = "https://chatgpt.com/backend-api/wham/usage";
const REQUEST_TIMEOUT: Duration = Duration::from_secs(15);
const MAX_RESPONSE_BYTES: u64 = 256 * 1024;
const AUTH_ERROR: &str = "Codex usage authentication failed; run /login openai-codex";
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct CodexUsage {
pub(crate) plan: Option<String>,
pub(crate) windows: Vec<CodexUsageWindow>,
}
#[derive(Clone, Debug, PartialEq)]
pub(crate) struct CodexUsageWindow {
pub(crate) label: String,
pub(crate) limit_window_seconds: Option<u64>,
pub(crate) used_percent: f64,
pub(crate) reset_at: Option<i64>,
}
#[derive(Clone, Copy, PartialEq, Eq)]
pub(crate) struct CodexAccountIdentity([u8; 32]);
impl std::fmt::Debug for CodexAccountIdentity {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("CodexAccountIdentity(<redacted>)")
}
}
impl CodexAccountIdentity {
pub(crate) fn from_oauth(access: &str, account_id: Option<&str>) -> Option<Self> {
use sha2::{Digest, Sha256};
let account_id = account_id
.filter(|id| !id.trim().is_empty())
.map(str::to_owned)
.or_else(|| auth::extract_chatgpt_account_id_from_jwt(access).ok())?;
Some(Self(Sha256::digest(account_id.as_bytes()).into()))
}
}
#[derive(Debug)]
pub(crate) struct AccountCodexUsage {
pub(crate) account: CodexAccountIdentity,
pub(crate) usage: CodexUsage,
}
pub(crate) fn load_codex_usage(paths: &McPaths) -> Result<AccountCodexUsage, String> {
let credential = auth::codex_credential_from_store(paths).map_err(|_| AUTH_ERROR.to_owned())?;
let client = usage_client(REQUEST_TIMEOUT)?;
fetch_codex_usage(&client, CODEX_USAGE_URL, credential, || {
auth::force_refresh_codex_credential_from_store(paths).map_err(|_| AUTH_ERROR.to_owned())
})
}
fn usage_client(timeout: Duration) -> Result<Client, String> {
Client::builder()
.timeout(timeout)
.connect_timeout(timeout)
.redirect(Policy::none())
.build()
.map_err(|_| "Could not create Codex usage HTTP client".to_owned())
}
fn fetch_codex_usage(
client: &Client,
url: &str,
mut credential: ProviderCredential,
refresh: impl FnOnce() -> Result<ProviderCredential, String>,
) -> Result<AccountCodexUsage, String> {
let response = request_usage(client, url, &credential)?;
let response = if response.status() == StatusCode::UNAUTHORIZED {
drop(response);
credential = refresh().map_err(|_| AUTH_ERROR.to_owned())?;
request_usage(client, url, &credential)?
} else {
response
};
let ProviderCredential::OAuth { access, account_id } = &credential else {
return Err(AUTH_ERROR.to_owned());
};
let account = CodexAccountIdentity::from_oauth(access, account_id.as_deref())
.ok_or_else(|| AUTH_ERROR.to_owned())?;
let status = response.status();
if status == StatusCode::UNAUTHORIZED || status == StatusCode::FORBIDDEN {
return Err(AUTH_ERROR.to_owned());
}
if !status.is_success() {
return Err(format!(
"Codex usage request failed (HTTP {})",
status.as_u16()
));
}
let body = read_bounded_response_text(response, MAX_RESPONSE_BYTES).map_err(|_| {
"Could not read Codex usage response within size and time limits".to_owned()
})?;
Ok(AccountCodexUsage {
account,
usage: parse_codex_usage(&body)?,
})
}
fn request_usage(
client: &Client,
url: &str,
credential: &ProviderCredential,
) -> Result<Response, String> {
let ProviderCredential::OAuth { access, account_id } = credential else {
return Err(AUTH_ERROR.to_owned());
};
if access.trim().is_empty() {
return Err(AUTH_ERROR.to_owned());
}
let mut request = client
.get(url)
.bearer_auth(access)
.header("Accept", "application/json")
.header(
"User-Agent",
concat!("magi-code/", env!("CARGO_PKG_VERSION")),
)
.header("Cache-Control", "no-cache")
.header("Pragma", "no-cache");
if let Some(account_id) = account_id
.as_deref()
.filter(|value| !value.trim().is_empty())
{
request = request.header("ChatGPT-Account-Id", account_id);
}
request.send().map_err(|error| {
if error.is_timeout() {
"Codex usage request timed out".to_owned()
} else {
"Codex usage request failed".to_owned()
}
})
}
#[derive(Deserialize)]
struct UsageResponse {
plan_type: Option<String>,
rate_limit: Option<RateLimit>,
}
#[derive(Deserialize)]
struct RateLimit {
primary_window: Option<UsageWindow>,
secondary_window: Option<UsageWindow>,
}
#[derive(Deserialize)]
struct UsageWindow {
used_percent: f64,
limit_window_seconds: Option<u64>,
reset_at: Option<i64>,
}
fn parse_codex_usage(body: &str) -> Result<CodexUsage, String> {
let response: UsageResponse =
serde_json::from_str(body).map_err(|_| "Invalid Codex usage response".to_owned())?;
let limits = response
.rate_limit
.ok_or_else(|| "Codex usage windows are missing".to_owned())?;
let mut windows = Vec::with_capacity(2);
for (fallback_label, window) in [
("Primary", limits.primary_window),
("Secondary", limits.secondary_window),
] {
let Some(window) = window else { continue };
if !window.used_percent.is_finite()
|| window.used_percent < 0.0
|| window.limit_window_seconds == Some(0)
{
return Err("Invalid Codex usage window".to_owned());
}
let label = match window.limit_window_seconds {
Some(18_000) => "5h",
Some(604_800) => "Weekly",
_ => fallback_label,
};
windows.push(CodexUsageWindow {
label: label.to_owned(),
limit_window_seconds: window.limit_window_seconds,
used_percent: window.used_percent,
reset_at: window.reset_at,
});
}
if windows.is_empty() {
return Err("Codex usage windows are missing".to_owned());
}
Ok(CodexUsage {
plan: response.plan_type,
windows,
})
}