magi-code 0.78.0

use std::time::Duration;

use reqwest::blocking::{Client, Response};
use reqwest::{StatusCode, redirect::Policy};
use serde::Deserialize;

use crate::config::{self, McPaths, ProviderCredential};
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>,
}

/// Only used for equality checks; never display account identifiers.
#[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(|| config::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,
}

/// Fetches a fresh snapshot. Blocking: callers must run this outside the UI thread.
pub(crate) fn load_codex_usage(paths: &McPaths) -> Result<AccountCodexUsage, String> {
    let credential =
        config::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, || {
        config::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,
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::io::{Read, Write};
    use std::net::TcpListener;
    use std::thread;

    const SNAPSHOT: &str = r#"{"plan_type":"plus","rate_limit":{"primary_window":{"used_percent":12.5,"limit_window_seconds":18000,"reset_at":1700000000},"secondary_window":{"used_percent":105,"limit_window_seconds":604800}}}"#;

    fn credential() -> ProviderCredential {
        ProviderCredential::OAuth {
            access: "test-token".to_owned(),
            account_id: Some("test-account".to_owned()),
        }
    }

    fn server(responses: Vec<String>) -> (String, thread::JoinHandle<Vec<String>>) {
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        listener.set_nonblocking(true).unwrap();
        let url = format!("http://{}/usage", listener.local_addr().unwrap());
        let handle = thread::spawn(move || {
            let mut requests = Vec::new();
            for response in responses {
                let deadline = std::time::Instant::now() + Duration::from_secs(5);
                let mut stream = loop {
                    match listener.accept() {
                        Ok((stream, _)) => break stream,
                        Err(error) if error.kind() == std::io::ErrorKind::WouldBlock => {
                            assert!(std::time::Instant::now() < deadline, "missing request");
                            thread::sleep(Duration::from_millis(5));
                        }
                        Err(error) => panic!("accept failed: {error}"),
                    }
                };
                stream
                    .set_read_timeout(Some(Duration::from_secs(2)))
                    .unwrap();
                let mut bytes = Vec::new();
                while !bytes.ends_with(b"\r\n\r\n") {
                    let mut byte = [0];
                    stream.read_exact(&mut byte).unwrap();
                    bytes.push(byte[0]);
                    assert!(bytes.len() < 16384);
                }
                requests.push(String::from_utf8(bytes).unwrap());
                // A client rejecting a large response may close before all bytes are sent.
                let _ = stream.write_all(response.as_bytes());
            }
            requests
        });
        (url, handle)
    }

    fn response(status: &str, body: &str) -> String {
        format!(
            "HTTP/1.1 {status}\r\nContent-Length: {}\r\nConnection: close\r\n\r\n{body}",
            body.len()
        )
    }

    #[test]
    fn parses_plan_windows_fractional_usage_and_optional_reset() {
        let usage = parse_codex_usage(SNAPSHOT).unwrap();
        assert_eq!(usage.plan.as_deref(), Some("plus"));
        assert_eq!(usage.windows.len(), 2);
        assert_eq!(usage.windows[0].label, "5h");
        assert_eq!(usage.windows[0].used_percent, 12.5);
        assert_eq!(usage.windows[0].reset_at, Some(1700000000));
        assert_eq!(usage.windows[1].label, "Weekly");
        assert_eq!(usage.windows[1].used_percent, 105.0);
        assert_eq!(usage.windows[1].reset_at, None);
    }

    #[test]
    fn rejects_missing_or_invalid_windows_without_exposing_body() {
        for body in [
            "private-invalid-body",
            "{}",
            r#"{"rate_limit":{}}"#,
            r#"{"rate_limit":{"primary_window":{"used_percent":-1,"limit_window_seconds":18000}}}"#,
            r#"{"rate_limit":{"primary_window":{"used_percent":1,"limit_window_seconds":0}}}"#,
        ] {
            let error = parse_codex_usage(body).unwrap_err();
            assert!(!error.contains(body));
        }
        let usage =
            parse_codex_usage(r#"{"rate_limit":{"secondary_window":{"used_percent":1}}}"#).unwrap();
        assert_eq!(usage.plan, None);
        assert_eq!(usage.windows[0].label, "Secondary");
    }

    #[test]
    fn successive_requests_show_changed_quota_then_error_without_cached_fallback() {
        let updated = SNAPSHOT.replace("12.5", "50");
        let (url, server) = server(vec![
            response("200 OK", SNAPSHOT),
            response("200 OK", &updated),
            response("503 Service Unavailable", "private-body"),
        ]);
        let client = usage_client(REQUEST_TIMEOUT).unwrap();
        let fetch =
            || fetch_codex_usage(&client, &url, credential(), || panic!("unexpected refresh"));
        assert_eq!(fetch().unwrap().usage.windows[0].used_percent, 12.5);
        assert_eq!(fetch().unwrap().usage.windows[0].used_percent, 50.0);
        assert_eq!(
            fetch().unwrap_err(),
            "Codex usage request failed (HTTP 503)"
        );
        assert_eq!(server.join().unwrap().len(), 3);
    }

    #[test]
    fn requests_fresh_usage_with_oauth_headers_and_refreshes_once() {
        let (url, server) = server(vec![
            response("401 Unauthorized", "private-body"),
            response("200 OK", SNAPSHOT),
        ]);
        let usage = fetch_codex_usage(
            &usage_client(REQUEST_TIMEOUT).unwrap(),
            &url,
            credential(),
            || {
                Ok(ProviderCredential::OAuth {
                    access: "refreshed-test-token".to_owned(),
                    account_id: Some("refreshed-test-account".to_owned()),
                })
            },
        )
        .unwrap();
        assert_eq!(usage.usage.windows[0].used_percent, 12.5);
        assert_eq!(
            usage.account,
            CodexAccountIdentity::from_oauth(
                "refreshed-test-token",
                Some("refreshed-test-account")
            )
            .unwrap()
        );
        let requests = server.join().unwrap();
        assert_eq!(requests.len(), 2);
        let first = requests[0].to_ascii_lowercase();
        assert!(first.starts_with("get /usage http/1.1"));
        assert!(first.contains("authorization: bearer test-token\r\n"));
        assert!(first.contains("chatgpt-account-id: test-account\r\n"));
        assert!(first.contains("cache-control: no-cache\r\n"));
        assert!(requests[1].contains("Bearer refreshed-test-token"));
        assert!(requests[1].contains("refreshed-test-account"));
    }

    #[test]
    fn repeated_unauthorized_stops_after_one_refresh() {
        let (url, server) = server(vec![response("401 Unauthorized", "private-body"); 2]);
        let error = fetch_codex_usage(
            &usage_client(REQUEST_TIMEOUT).unwrap(),
            &url,
            credential(),
            || Ok(credential()),
        )
        .unwrap_err();
        assert_eq!(error, AUTH_ERROR);
        assert_eq!(server.join().unwrap().len(), 2);
    }

    #[test]
    fn rejects_status_errors_redirects_and_oversized_bodies_without_refresh() {
        let redirect = "HTTP/1.1 302 Found\r\nLocation: http://127.0.0.1:1/private\r\nContent-Length: 0\r\nConnection: close\r\n\r\n".to_owned();
        for reply in [
            response("403 Forbidden", "private-body"),
            response("500 Internal Server Error", "private-body"),
            redirect,
            response("200 OK", &"x".repeat(MAX_RESPONSE_BYTES as usize + 1)),
        ] {
            let (url, server) = server(vec![reply]);
            let error = fetch_codex_usage(
                &usage_client(REQUEST_TIMEOUT).unwrap(),
                &url,
                credential(),
                || panic!("unexpected refresh"),
            )
            .unwrap_err();
            assert!(!error.contains("private"));
            assert_eq!(server.join().unwrap().len(), 1);
        }
    }

    #[test]
    fn missing_credentials_return_login_guidance_without_network() {
        let directory = tempfile::TempDir::new().unwrap();
        let paths = McPaths::from_root(directory.path().join("mc"));
        assert_eq!(load_codex_usage(&paths).unwrap_err(), AUTH_ERROR);
    }

    #[test]
    fn stalled_server_times_out_without_refresh() {
        // Keep the listening socket alive but never send response headers.
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let url = format!("http://{}/usage", listener.local_addr().unwrap());
        let error = fetch_codex_usage(
            &usage_client(Duration::from_millis(100)).unwrap(),
            &url,
            credential(),
            || panic!("unexpected refresh"),
        )
        .unwrap_err();
        assert_eq!(error, "Codex usage request timed out");
    }

    #[test]
    fn failed_refresh_returns_safe_auth_error() {
        let (url, server) = server(vec![response("401 Unauthorized", "private-body")]);
        let error = fetch_codex_usage(
            &usage_client(REQUEST_TIMEOUT).unwrap(),
            &url,
            credential(),
            || Err("private-auth-detail".to_owned()),
        )
        .unwrap_err();
        assert_eq!(error, AUTH_ERROR);
        assert_eq!(server.join().unwrap().len(), 1);
    }
}