hatchery_node/

provider_limits.rs

//! Provider usage limits as the provider itself shows them, read ONLY from PTY screens: the
//! node's limit heartbeat ([`crate::limit_heartbeat`]: one light service PTY per signed-in
//! provider account) and any node PTY session that happens to show the screen.
//!
//! - claude: `/usage` (and `/status`) — "Current session" = 5h window, "Current week" = 7d
//!   window, the plan from "Claude Max/Pro/Team/Enterprise".
//! - codex: `/status` — its "5h limit" and "Weekly limit" lines.
//!
//! Nothing is estimated, and no other source is read: no session logs or rollouts, no
//! provider error text (that only notifies the operator), never tokens, OAuth or `auth.json`.
//! Every entry keeps its `source` and `observed_at_ms`.

use serde::Serialize;
use std::collections::BTreeMap;
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>,
    /// `100 - used`, derived here from the raw reading.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub remaining_percent: Option<f64>,
    /// What the provider showed: the line, the number and what it meant (used / left).
    #[serde(skip_serializing_if = "Option::is_none")]
    pub raw: Option<RawReading>,
    /// 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>,
}

/// A provider's own limit line, kept as it was shown (providers say `% used` or `% left`).
#[derive(Clone, Debug, Serialize, PartialEq)]
pub struct RawReading {
    /// The line, trimmed and cut to 160 characters.
    pub text: String,
    pub value: Option<f64>,
    /// `used` | `left` | `unknown`
    pub meaning: &'static str,
}

/// The percent on `line` and what it means. `left` also covers `remaining`.
pub fn reading(line: &str) -> Option<RawReading> {
    let value = pct(line)?;
    let low = line.to_lowercase();
    let after = &low[low.find('%').unwrap_or(0)..];
    let meaning = if after.contains("left") || after.contains("remaining") {
        "left"
    } else if after.contains("used") {
        "used"
    } else {
        "unknown"
    };
    let text: String = line.trim().trim_matches(|c: char| c == '│' || c == '|').trim().chars().filter(|c| !c.is_control()).take(160).collect();
    Some(RawReading { text, value: Some(value), meaning })
}

/// A window from a raw reading: used and remaining both derived from it; `unknown` keeps
/// both unset (routing never guesses which way a bare percent goes).
pub fn window_from(name: &str, raw: RawReading, resets_at_ms: Option<u64>, resets_text: Option<String>) -> Window {
    let used = match (raw.meaning, raw.value) {
        ("used", Some(v)) => Some(v),
        ("left", Some(v)) => Some(100.0 - v),
        _ => None,
    };
    Window { window: name.into(), used_percent: used, remaining_percent: used.map(|u| 100.0 - u), raw: Some(raw), resets_at_ms, resets_text }
}

#[derive(Clone, Debug, Serialize, PartialEq)]
pub struct ProviderLimits {
    pub provider: String,
    /// `screen`: a session PTY that showed it; `heartbeat`: the node's limit heartbeat.
    pub source: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub tier: Option<String>,
    pub windows: Vec<Window>,
    pub observed_at_ms: u64,
}

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 raw, mut resets_text) = (None, None);
        for next in lines.iter().skip(i + 1).take(3) {
            let nl = next.to_lowercase();
            if nl.starts_with("current ") {
                break;
            }
            if raw.is_none() && (nl.contains("used") || nl.contains("left") || nl.contains("remaining")) {
                raw = reading(next);
            }
            if let Some(p) = nl.find("resets") {
                resets_text = Some(next[p..].chars().take(80).collect());
            }
        }
        if let Some(raw) = raw {
            windows.push(window_from(name, raw, None, resets_text));
        }
    }
    (!windows.is_empty() || tier.is_some()).then(|| ProviderLimits { provider: "claude".into(), source: "screen".into(), tier, windows, observed_at_ms: now_ms })
}

/// Latest limits the node saw on screens, per (provider, source).
static OBSERVED: OnceLock<Mutex<BTreeMap<(String, String), ProviderLimits>>> = OnceLock::new();
/// Woken on every new observation, so the reporter sends it without waiting for its tick.
static CHANGED: OnceLock<tokio::sync::Notify> = OnceLock::new();

pub fn changed() -> &'static tokio::sync::Notify {
    CHANGED.get_or_init(tokio::sync::Notify::new)
}

/// Merge `new` into the observed entry of its (provider, source): its windows replace the same
/// windows, the rest stay.
fn remember(new: ProviderLimits) {
    let key = (new.provider.clone(), new.source.clone());
    let mut map = OBSERVED.get_or_init(Default::default).lock().unwrap_or_else(|p| p.into_inner());
    let before = map.remove(&key);
    let merged = match before.clone() {
        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,
    };
    // A repeat of the same reading (only its time moved) wakes nobody.
    let fresh = before.as_ref().is_none_or(|b| b.windows != merged.windows || b.tier != merged.tier);
    map.insert(key, merged);
    drop(map);
    if fresh {
        changed().notify_one();
    }
}

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

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

/// Parse a codex `/status` screen: the `5h limit` / `Weekly limit` lines (`N% used`, or
/// `N% left`), and the reset in parentheses. Nothing else on the screen is kept (it also shows
/// the account's email). `None` when it shows no limit line.
pub fn parse_codex_status(screen: &str, now_ms: u64) -> Option<ProviderLimits> {
    let mut windows: Vec<Window> = Vec::new();
    for l in screen.lines() {
        let l = l.trim().trim_matches(|c: char| c == '│' || c == '|').trim();
        let low = l.to_lowercase();
        let Some(colon) = low.find("limit:") else { continue };
        let label = &low[..colon];
        let name = if label.contains("5h") || label.contains("5-hour") || label.contains("session") {
            "5h"
        } else if label.contains("week") {
            "7d"
        } else {
            continue;
        };
        let Some(raw) = reading(l) else { continue };
        let resets_text = low.find("resets").filter(|_| low.len() == l.len()).map(|i| l[i..].trim_end_matches(|c: char| c == ')' || c.is_whitespace()).chars().take(80).collect::<String>());
        if windows.iter().all(|w| w.window != name) {
            let resets_at_ms = resets_text.as_deref().and_then(|t| codex_reset_ms(t, chrono::Local::now().fixed_offset()));
            windows.push(window_from(name, raw, resets_at_ms, resets_text));
        }
    }
    (!windows.is_empty()).then(|| ProviderLimits { provider: "codex".into(), source: "screen".into(), tier: None, windows, observed_at_ms: now_ms })
}

/// A codex reset as shown on its screen: `resets HH:MM` (the next such time) or
/// `resets HH:MM on D Mon` (that day; next year when already past). The screen gives no zone,
/// so it is the node's local time (`now`'s offset).
pub fn codex_reset_ms(text: &str, now: chrono::DateTime<chrono::FixedOffset>) -> Option<u64> {
    use chrono::{Datelike, NaiveDate, NaiveTime, TimeZone};
    let rest = text.trim().strip_prefix("resets")?.trim();
    let (hm, day) = match rest.split_once(" on ") {
        Some((hm, day)) => (hm.trim(), Some(day.trim())),
        None => (rest, None),
    };
    let time = NaiveTime::parse_from_str(hm, "%H:%M").ok()?;
    let tz = *now.offset();
    let at = match day {
        None => {
            let today = tz.from_local_datetime(&now.date_naive().and_time(time)).single()?;
            if today > now { today } else { today + chrono::Duration::days(1) }
        }
        Some(day) => {
            let mut it = day.split_whitespace();
            let d: u32 = it.next()?.parse().ok()?;
            let m = ["jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec"]
                .iter()
                .position(|x| it.clone().next().is_some_and(|w| w.to_lowercase().starts_with(x)))? as u32
                + 1;
            let mut at = tz.from_local_datetime(&NaiveDate::from_ymd_opt(now.year(), m, d)?.and_time(time)).single()?;
            if at < now - chrono::Duration::days(1) {
                at = tz.from_local_datetime(&NaiveDate::from_ymd_opt(now.year() + 1, m, d)?.and_time(time)).single()?;
            }
            at
        }
    };
    u64::try_from(at.timestamp_millis()).ok()
}

/// Parse `screen` as `provider`'s limit screen.
pub fn parse_screen(provider: &str, screen: &str, now_ms: u64) -> Option<ProviderLimits> {
    match provider {
        "claude" if screen.contains("Current session") || screen.contains("Current week") || screen.contains("Claude ") => parse_claude_screen(screen, now_ms),
        "codex" if screen.contains("limit:") => parse_codex_status(screen, now_ms),
        _ => None,
    }
}

/// Called from the frame publisher with every scanned session screen.
pub fn observe_screen(provider: &str, screen: &str) {
    observe_screen_as(provider, screen, "screen");
}

/// Remember a limit screen under `source` (`screen` or `heartbeat`). True when it was one.
pub fn observe_screen_as(provider: &str, screen: &str, source: &str) -> bool {
    match parse_screen(provider, screen, now_ms()) {
        Some(mut new) => {
            new.source = source.to_owned();
            remember(new);
            true
        }
        None => false,
    }
}

/// Everything seen on screens now.
pub fn collect() -> Vec<ProviderLimits> {
    OBSERVED.get().map(|m| m.lock().unwrap_or_else(|p| p.into_inner()).values().cloned().collect()).unwrap_or_default()
}

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

    #[test]
    fn raw_reading_keeps_meaning_and_normalizes_both_ways() {
        let codex = "│  5h limit:       [██████░░] 30% used (resets 14:05)  │\n│  Weekly limit:   [█░░░░░░░] 95% left (resets 09:00 on 23 Oct)  │";
        let l = super::parse_codex_status(codex, 0).unwrap();
        let w: Vec<_> = l.windows.iter().map(|w| (w.window.as_str(), w.used_percent, w.remaining_percent, w.raw.as_ref().map(|r| (r.value, r.meaning)))).collect();
        assert_eq!(w, vec![("5h", Some(30.0), Some(70.0), Some((Some(30.0), "used"))), ("7d", Some(5.0), Some(95.0), Some((Some(95.0), "left")))]);
        assert!(l.windows[1].raw.as_ref().unwrap().text.starts_with("Weekly limit:"));
        let sample = include_str!("../tests/fixtures/codex-status-v0.163.0.txt");
        let l = super::parse_codex_status(sample, 0).unwrap();
        assert_eq!((l.windows[0].used_percent, l.windows[0].remaining_percent, l.windows[0].raw.as_ref().unwrap().meaning), (Some(100.0), Some(0.0), "left"));
        let claude = "Current session\n███▌ 82% used\nResets 3:40pm (Asia/Yekaterinburg)\nCurrent week (all models)\n 84% left\n";
        let l = super::parse_claude_screen(claude, 0).unwrap();
        let w: Vec<_> = l.windows.iter().map(|w| (w.window.as_str(), w.used_percent, w.remaining_percent, w.raw.as_ref().unwrap().meaning)).collect();
        assert_eq!(w, vec![("5h", Some(82.0), Some(18.0), "used"), ("7d", Some(16.0), Some(84.0), "left")]);
        let bare = super::window_from("5h", super::reading("5h: 40%").unwrap(), None, None);
        assert_eq!((bare.used_percent, bare.remaining_percent, bare.raw.unwrap().meaning), (None, None, "unknown"));
    }

    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_status_screen() {
        let screen = "╭──────────────────────────────────────────────╮\n│  >_ OpenAI Codex (v0.160.0)                  │\n│  Account:        someone@example.com (Plus)  │\n│  5h limit:       [██████░░░░░░░░░░░░░░] 30% used (resets 14:05)  │\n│  Weekly limit:   [█░░░░░░░░░░░░░░░░░░░] 95% left (resets 09:00 on 23 Oct)  │\n╰──────────────────────────────────────────────╯";
        let l = parse_codex_status(screen, 5).unwrap();
        assert_eq!(l.windows.iter().map(|w| (w.window.as_str(), w.used_percent)).collect::<Vec<_>>(), vec![("5h", Some(30.0)), ("7d", Some(5.0))]);
        assert_eq!(l.windows[0].resets_text.as_deref(), Some("resets 14:05"));
        assert_eq!(l.windows[1].resets_text.as_deref(), Some("resets 09:00 on 23 Oct"));
        // The account line is never kept.
        assert!(!format!("{l:?}").contains('@'));
        assert!(parse_codex_status("› explain the rate limit: code in src/x.rs", 5).is_none());
    }

    /// Real codex v0.163.0 `/status` (op-box, 2026-10-19; account and session redacted).
    #[test]
    fn codex_status_fixture_v0_163() {
        let screen = include_str!("../tests/fixtures/codex-status-v0.163.0.txt");
        let l = parse_codex_status(screen, 5).unwrap();
        assert_eq!(l.windows.iter().map(|w| (w.window.as_str(), w.used_percent)).collect::<Vec<_>>(), vec![("5h", Some(100.0)), ("7d", Some(16.0))]);
        assert_eq!(l.windows[0].resets_text.as_deref(), Some("resets 12:11"));
        assert_eq!(l.windows[1].resets_text.as_deref(), Some("resets 06:27 on 25 Oct"));
        assert!(l.windows.iter().all(|w| w.resets_at_ms.is_some()));
        let dbg = format!("{l:?}");
        assert!(!dbg.contains('@') && !dbg.contains("00000000-"), "no account or session id kept");
    }

    #[test]
    fn codex_reset_is_node_local_time() {
        use chrono::TimeZone;
        let tz = chrono::FixedOffset::east_opt(5 * 3600).unwrap();
        let now = tz.with_ymd_and_hms(2026, 10, 19, 11, 41, 0).unwrap();
        let at = |y, mo, d, h, mi| u64::try_from(tz.with_ymd_and_hms(y, mo, d, h, mi, 0).unwrap().timestamp_millis()).unwrap();
        assert_eq!(codex_reset_ms("resets 12:11", now), Some(at(2026, 10, 19, 12, 11)));
        assert_eq!(codex_reset_ms("resets 09:00", now), Some(at(2026, 10, 20, 9, 0)), "already past today: tomorrow");
        assert_eq!(codex_reset_ms("resets 06:27 on 25 Oct", now), Some(at(2026, 10, 25, 6, 27)));
        assert_eq!(codex_reset_ms("resets 06:27 on 2 Jan", now), Some(at(2027, 1, 2, 6, 27)));
        assert_eq!(codex_reset_ms("resets soon", now), None);
    }

    #[test]
    fn screens_are_remembered_with_their_source() {
        assert!(observe_screen_as("codex", "  5h limit: [█] 88% used (resets 10:00)", "heartbeat"));
        assert!(!observe_screen_as("codex", "> hello", "heartbeat"));
        assert!(!observe_screen_as("kimi", "5h limit: 10% used", "heartbeat"));
        let got = collect().into_iter().find(|l| l.provider == "codex" && l.source == "heartbeat").unwrap();
        assert_eq!(got.windows[0].used_percent, Some(88.0));
    }
}