hatchery_node/

provider_limits.rs

//! Provider usage limits as the provider itself reports them. Honest sources only: nothing is
//! estimated, a provider with no source reports nothing.
//!
//! - claude: its own `/usage` (and `/status`) screens as rendered in a node PTY session —
//!   "Current session" = 5h window, "Current week" = 7d window, the plan from
//!   "Claude Max/Pro/Team/Enterprise". The node never types `/usage` into a session; it only
//!   reads screens the session already shows. Credential files are never opened.
//! - codex: `rate_limits` in its own rollout logs (`~/.codex/sessions/**/rollout-*.jsonl`):
//!   primary/secondary windows with `used_percent`, `window_minutes` and reset time, `plan_type`.
//! - others: implement [`LimitSource`].

use serde::Serialize;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};

#[derive(Clone, Debug, Serialize, PartialEq)]
pub struct Window {
    /// `5h`, `7d`, or `<N>m` for any other length.
    pub window: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub used_percent: Option<f64>,
    /// Reset as the provider said it: epoch ms when known, else its own text.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub resets_at_ms: Option<u64>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub resets_text: Option<String>,
}

#[derive(Clone, Debug, Serialize, PartialEq)]
pub struct ProviderLimits {
    pub provider: String,
    /// `screen` (claude /usage), `rollout` (codex logs), ...
    pub source: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tier: Option<String>,
    pub windows: Vec<Window>,
    pub observed_at_ms: u64,
}

/// A provider-specific reader. `read` must only use data the provider produced itself.
pub trait LimitSource: Send + Sync {
    fn provider(&self) -> &'static str;
    fn read(&self, now_ms: u64) -> Option<ProviderLimits>;
}

fn window_name(minutes: u64) -> String {
    match minutes {
        300 => "5h".into(),
        10080 => "7d".into(),
        m => format!("{m}m"),
    }
}

fn pct(line: &str) -> Option<f64> {
    let i = line.find('%')?;
    let head = &line[..i];
    let start = head.rfind(|c: char| !(c.is_ascii_digit() || c == '.')).map_or(0, |p| p + 1);
    head[start..].parse().ok().filter(|v: &f64| (0.0..=100.0).contains(v))
}

/// Parse a claude `/usage` (or `/status`) screen. `None` when it shows neither.
pub fn parse_claude_screen(screen: &str, now_ms: u64) -> Option<ProviderLimits> {
    let lines: Vec<&str> = screen.lines().map(str::trim).collect();
    let mut windows = Vec::new();
    let mut tier = None;
    for (i, l) in lines.iter().enumerate() {
        let low = l.to_lowercase();
        for t in ["max", "pro", "team", "enterprise"] {
            if tier.is_none() && (low.contains(&format!("claude {t}")) || low.contains(&format!("{t} plan"))) {
                tier = Some(t.to_owned());
            }
        }
        let name = if low.starts_with("current session") {
            "5h"
        } else if low.starts_with("current week") && (low.contains("all models") || !low.contains('(')) {
            "7d"
        } else {
            continue;
        };
        if windows.iter().any(|w: &Window| w.window == name) {
            continue;
        }
        let mut w = Window { window: name.into(), used_percent: None, resets_at_ms: None, resets_text: None };
        for next in lines.iter().skip(i + 1).take(3) {
            let nl = next.to_lowercase();
            if nl.starts_with("current ") {
                break;
            }
            if w.used_percent.is_none() && nl.contains("used") {
                w.used_percent = pct(next);
            }
            if let Some(p) = nl.find("resets") {
                w.resets_text = Some(next[p..].chars().take(80).collect());
            }
        }
        if w.used_percent.is_some() {
            windows.push(w);
        }
    }
    (!windows.is_empty() || tier.is_some()).then(|| ProviderLimits { provider: "claude".into(), source: "screen".into(), tier, windows, observed_at_ms: now_ms })
}

/// Parse one codex rollout line carrying `rate_limits`.
pub fn parse_codex_line(line: &str, now_ms: u64) -> Option<ProviderLimits> {
    let v: serde_json::Value = serde_json::from_str(line).ok()?;
    let rl = v.pointer("/payload/rate_limits").or_else(|| v.pointer("/payload/info/rate_limits")).or_else(|| v.get("rate_limits"))?;
    let mut windows = Vec::new();
    for key in ["primary", "secondary"] {
        let Some(w) = rl.get(key).filter(|w| w.is_object()) else { continue };
        let minutes = w.get("window_minutes").and_then(|x| x.as_u64()).unwrap_or(0);
        let resets_at_ms = w
            .get("resets_at")
            .and_then(|x| x.as_u64())
            .map(|s| s * 1000)
            .or_else(|| w.get("resets_in_seconds").and_then(|x| x.as_u64()).map(|s| now_ms + s * 1000));
        windows.push(Window { window: window_name(minutes), used_percent: w.get("used_percent").and_then(|x| x.as_f64()), resets_at_ms, resets_text: None });
    }
    let tier = rl.get("plan_type").and_then(|x| x.as_str()).map(str::to_owned);
    (!windows.is_empty()).then(|| ProviderLimits { provider: "codex".into(), source: "rollout".into(), tier, windows, observed_at_ms: now_ms })
}

/// Codex: newest rollout file, last line with `rate_limits`.
pub struct CodexRollouts {
    pub root: PathBuf,
}

impl CodexRollouts {
    pub fn from_home() -> Option<Self> {
        let home = std::env::var_os("CODEX_HOME").map(PathBuf::from).or_else(|| std::env::var_os(if cfg!(windows) { "USERPROFILE" } else { "HOME" }).map(|h| PathBuf::from(h).join(".codex")))?;
        Some(Self { root: home.join("sessions") })
    }
}

fn newest_rollout(dir: &Path, depth: u8) -> Option<PathBuf> {
    let mut best: Option<(std::time::SystemTime, PathBuf)> = None;
    for e in std::fs::read_dir(dir).ok()?.flatten() {
        let p = e.path();
        let cand = if p.is_dir() && depth > 0 {
            newest_rollout(&p, depth - 1)
        } else if p.file_name().and_then(|n| n.to_str()).is_some_and(|n| n.starts_with("rollout-") && n.ends_with(".jsonl")) {
            Some(p)
        } else {
            None
        };
        if let Some(c) = cand {
            let t = std::fs::metadata(&c).and_then(|m| m.modified()).unwrap_or(std::time::UNIX_EPOCH);
            if best.as_ref().map_or(true, |(bt, _)| t > *bt) {
                best = Some((t, c));
            }
        }
    }
    best.map(|b| b.1)
}

impl LimitSource for CodexRollouts {
    fn provider(&self) -> &'static str {
        "codex"
    }
    fn read(&self, now_ms: u64) -> Option<ProviderLimits> {
        let file = newest_rollout(&self.root, 4)?;
        let text = std::fs::read_to_string(file).ok()?;
        text.lines().rev().filter(|l| l.contains("rate_limits")).find_map(|l| parse_codex_line(l, now_ms))
    }
}

/// Latest screen-derived limits per provider (claude today).
static SCREEN: OnceLock<Mutex<BTreeMap<String, ProviderLimits>>> = OnceLock::new();

pub fn now_ms() -> u64 {
    std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(0, |d| d.as_millis() as u64)
}

/// Called from the frame publisher with every scanned screen.
pub fn observe_screen(provider: &str, screen: &str) {
    if provider != "claude" || !(screen.contains("Current session") || screen.contains("Current week") || screen.contains("Claude ")) {
        return;
    }
    let Some(new) = parse_claude_screen(screen, now_ms()) else { return };
    let mut map = SCREEN.get_or_init(Default::default).lock().unwrap_or_else(|p| p.into_inner());
    let merged = match map.remove(provider) {
        Some(mut old) => {
            if new.tier.is_some() {
                old.tier = new.tier.clone();
            }
            for w in new.windows {
                old.windows.retain(|o| o.window != w.window);
                old.windows.push(w);
            }
            old.observed_at_ms = new.observed_at_ms;
            old
        }
        None => new,
    };
    map.insert(provider.to_owned(), merged);
}

/// The 5h window usage of `provider` if a reading at most `max_age_ms` old is known.
pub fn fresh_5h_used(provider: &str, max_age_ms: u64) -> Option<f64> {
    let map = SCREEN.get()?.lock().unwrap_or_else(|p| p.into_inner());
    let l = map.get(provider)?;
    if now_ms().saturating_sub(l.observed_at_ms) > max_age_ms {
        return None;
    }
    l.windows.iter().find(|w| w.window == "5h").and_then(|w| w.used_percent)
}

/// Everything known now: screen-derived entries plus every file source.
pub fn collect(sources: &[Box<dyn LimitSource>]) -> Vec<ProviderLimits> {
    let now = now_ms();
    let mut out: Vec<ProviderLimits> = SCREEN.get().map(|m| m.lock().unwrap_or_else(|p| p.into_inner()).values().cloned().collect()).unwrap_or_default();
    out.extend(sources.iter().filter_map(|s| s.read(now)));
    out
}

pub fn default_sources() -> Vec<Box<dyn LimitSource>> {
    let mut v: Vec<Box<dyn LimitSource>> = Vec::new();
    if let Some(c) = CodexRollouts::from_home() {
        v.push(Box::new(c));
    }
    v
}

#[cfg(test)]
mod tests {
    use super::*;

    const USAGE: &str = "  Settings:  Status   Config   Usage\n\n  Current session\n  █████▌                         11% used\n  Resets 3pm (Asia/Yekaterinburg)\n\n  Current week (all models)\n  ███████████                    22% used\n  Resets Oct 23, 9am (Asia/Yekaterinburg)\n\n  Current week (Opus)\n  █                              2% used\n";

    #[test]
    fn claude_usage_screen() {
        let l = parse_claude_screen(USAGE, 7).unwrap();
        assert_eq!(l.windows.len(), 2);
        assert_eq!(l.windows[0].window, "5h");
        assert_eq!(l.windows[0].used_percent, Some(11.0));
        assert_eq!(l.windows[0].resets_text.as_deref(), Some("Resets 3pm (Asia/Yekaterinburg)"));
        assert_eq!(l.windows[1].window, "7d");
        assert_eq!(l.windows[1].used_percent, Some(22.0));
    }

    /// Real claude /usage screen captured on hatchery-box (2026-10-19).
    #[test]
    fn claude_usage_screen_live_capture() {
        let screen = "   Settings  Status   Config   Usage   Stats\n\n   Session\n\n   Total cost:            $0.0000\n   Usage:                 0 input, 0 output, 0 cache read, 0 cache write\n\n   Current session\n   ▌                                                  1% used\n   Resets 10:39am (Asia/Yekaterinburg)\n\n   Current week (all models)\n   ▌                                                  1% used\n   Resets Oct 21, 11:59pm (Asia/Yekaterinburg)\n\n   What's contributing to your limits usage?\n   31% of your usage came from the MCP server \"m4a\"\n";
        let l = parse_claude_screen(screen, 0).unwrap();
        assert_eq!(l.windows.iter().map(|w| (w.window.as_str(), w.used_percent)).collect::<Vec<_>>(), vec![("5h", Some(1.0)), ("7d", Some(1.0))]);
        assert_eq!(l.windows[1].resets_text.as_deref(), Some("Resets Oct 21, 11:59pm (Asia/Yekaterinburg)"));
    }

    #[test]
    fn claude_status_tier_and_noise() {
        assert_eq!(parse_claude_screen("Login method: Claude Max Account", 1).unwrap().tier.as_deref(), Some("max"));
        assert!(parse_claude_screen("> fix the tests\n● done", 1).is_none());
    }

    #[test]
    fn codex_rollout_line() {
        let line = r#"{"type":"event_msg","payload":{"type":"token_count","rate_limits":{"primary":{"used_percent":40.5,"window_minutes":300,"resets_in_seconds":60},"secondary":{"used_percent":9.0,"window_minutes":10080,"resets_at":1800000000},"plan_type":"plus"}}}"#;
        let l = parse_codex_line(line, 1000).unwrap();
        assert_eq!(l.tier.as_deref(), Some("plus"));
        assert_eq!(l.windows[0].window, "5h");
        assert_eq!(l.windows[0].resets_at_ms, Some(61_000));
        assert_eq!(l.windows[1].window, "7d");
        assert_eq!(l.windows[1].resets_at_ms, Some(1_800_000_000_000));
    }

    #[test]
    fn codex_source_reads_newest_file() {
        let d = std::env::temp_dir().join(format!("pl-{}", std::process::id()));
        let day = d.join("2026/10/19");
        std::fs::create_dir_all(&day).unwrap();
        std::fs::write(day.join("rollout-a.jsonl"), "{}\n{\"payload\":{\"rate_limits\":{\"primary\":{\"used_percent\":5,\"window_minutes\":300}}}}\n").unwrap();
        let l = CodexRollouts { root: d.clone() }.read(0).unwrap();
        assert_eq!(l.windows[0].used_percent, Some(5.0));
        let _ = std::fs::remove_dir_all(d);
    }
}