Skip to main content

hatchery_node/
provider_limits.rs

1//! Provider usage limits as the provider itself reports them. Honest sources only: nothing is
2//! estimated, a provider with no source reports nothing.
3//!
4//! - claude: its own `/usage` (and `/status`) screens as rendered in a node PTY session —
5//!   "Current session" = 5h window, "Current week" = 7d window, the plan from
6//!   "Claude Max/Pro/Team/Enterprise". The node never types `/usage` into a session; it only
7//!   reads screens the session already shows. Credential files are never opened.
8//! - codex: `rate_limits` in its own rollout logs (`~/.codex/sessions/**/rollout-*.jsonl`):
9//!   primary/secondary windows with `used_percent`, `window_minutes` and reset time, `plan_type`.
10//! - others: implement [`LimitSource`].
11
12use serde::Serialize;
13use std::collections::BTreeMap;
14use std::path::{Path, PathBuf};
15use std::sync::{Mutex, OnceLock};
16
17#[derive(Clone, Debug, Serialize, PartialEq)]
18pub struct Window {
19    /// `5h`, `7d`, or `<N>m` for any other length.
20    pub window: String,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub used_percent: Option<f64>,
23    /// Reset as the provider said it: epoch ms when known, else its own text.
24    #[serde(skip_serializing_if = "Option::is_none")]
25    pub resets_at_ms: Option<u64>,
26    #[serde(skip_serializing_if = "Option::is_none")]
27    pub resets_text: Option<String>,
28}
29
30#[derive(Clone, Debug, Serialize, PartialEq)]
31pub struct ProviderLimits {
32    pub provider: String,
33    /// `screen` (claude /usage), `rollout` (codex logs), ...
34    pub source: String,
35    #[serde(skip_serializing_if = "Option::is_none")]
36    pub tier: Option<String>,
37    pub windows: Vec<Window>,
38    pub observed_at_ms: u64,
39}
40
41/// A provider-specific reader. `read` must only use data the provider produced itself.
42pub trait LimitSource: Send + Sync {
43    fn provider(&self) -> &'static str;
44    fn read(&self, now_ms: u64) -> Option<ProviderLimits>;
45}
46
47fn window_name(minutes: u64) -> String {
48    match minutes {
49        300 => "5h".into(),
50        10080 => "7d".into(),
51        m => format!("{m}m"),
52    }
53}
54
55fn pct(line: &str) -> Option<f64> {
56    let i = line.find('%')?;
57    let head = &line[..i];
58    let start = head.rfind(|c: char| !(c.is_ascii_digit() || c == '.')).map_or(0, |p| p + 1);
59    head[start..].parse().ok().filter(|v: &f64| (0.0..=100.0).contains(v))
60}
61
62/// Parse a claude `/usage` (or `/status`) screen. `None` when it shows neither.
63pub fn parse_claude_screen(screen: &str, now_ms: u64) -> Option<ProviderLimits> {
64    let lines: Vec<&str> = screen.lines().map(str::trim).collect();
65    let mut windows = Vec::new();
66    let mut tier = None;
67    for (i, l) in lines.iter().enumerate() {
68        let low = l.to_lowercase();
69        for t in ["max", "pro", "team", "enterprise"] {
70            if tier.is_none() && (low.contains(&format!("claude {t}")) || low.contains(&format!("{t} plan"))) {
71                tier = Some(t.to_owned());
72            }
73        }
74        let name = if low.starts_with("current session") {
75            "5h"
76        } else if low.starts_with("current week") && (low.contains("all models") || !low.contains('(')) {
77            "7d"
78        } else {
79            continue;
80        };
81        if windows.iter().any(|w: &Window| w.window == name) {
82            continue;
83        }
84        let mut w = Window { window: name.into(), used_percent: None, resets_at_ms: None, resets_text: None };
85        for next in lines.iter().skip(i + 1).take(3) {
86            let nl = next.to_lowercase();
87            if nl.starts_with("current ") {
88                break;
89            }
90            if w.used_percent.is_none() && nl.contains("used") {
91                w.used_percent = pct(next);
92            }
93            if let Some(p) = nl.find("resets") {
94                w.resets_text = Some(next[p..].chars().take(80).collect());
95            }
96        }
97        if w.used_percent.is_some() {
98            windows.push(w);
99        }
100    }
101    (!windows.is_empty() || tier.is_some()).then(|| ProviderLimits { provider: "claude".into(), source: "screen".into(), tier, windows, observed_at_ms: now_ms })
102}
103
104/// Parse one codex rollout line carrying `rate_limits`.
105pub fn parse_codex_line(line: &str, now_ms: u64) -> Option<ProviderLimits> {
106    let v: serde_json::Value = serde_json::from_str(line).ok()?;
107    let rl = v.pointer("/payload/rate_limits").or_else(|| v.pointer("/payload/info/rate_limits")).or_else(|| v.get("rate_limits"))?;
108    let mut windows = Vec::new();
109    for key in ["primary", "secondary"] {
110        let Some(w) = rl.get(key).filter(|w| w.is_object()) else { continue };
111        let minutes = w.get("window_minutes").and_then(|x| x.as_u64()).unwrap_or(0);
112        let resets_at_ms = w
113            .get("resets_at")
114            .and_then(|x| x.as_u64())
115            .map(|s| s * 1000)
116            .or_else(|| w.get("resets_in_seconds").and_then(|x| x.as_u64()).map(|s| now_ms + s * 1000));
117        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 });
118    }
119    let tier = rl.get("plan_type").and_then(|x| x.as_str()).map(str::to_owned);
120    (!windows.is_empty()).then(|| ProviderLimits { provider: "codex".into(), source: "rollout".into(), tier, windows, observed_at_ms: now_ms })
121}
122
123/// Codex: newest rollout file, last line with `rate_limits`.
124pub struct CodexRollouts {
125    pub root: PathBuf,
126}
127
128impl CodexRollouts {
129    pub fn from_home() -> Option<Self> {
130        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")))?;
131        Some(Self { root: home.join("sessions") })
132    }
133}
134
135fn newest_rollout(dir: &Path, depth: u8) -> Option<PathBuf> {
136    let mut best: Option<(std::time::SystemTime, PathBuf)> = None;
137    for e in std::fs::read_dir(dir).ok()?.flatten() {
138        let p = e.path();
139        let cand = if p.is_dir() && depth > 0 {
140            newest_rollout(&p, depth - 1)
141        } else if p.file_name().and_then(|n| n.to_str()).is_some_and(|n| n.starts_with("rollout-") && n.ends_with(".jsonl")) {
142            Some(p)
143        } else {
144            None
145        };
146        if let Some(c) = cand {
147            let t = std::fs::metadata(&c).and_then(|m| m.modified()).unwrap_or(std::time::UNIX_EPOCH);
148            if best.as_ref().map_or(true, |(bt, _)| t > *bt) {
149                best = Some((t, c));
150            }
151        }
152    }
153    best.map(|b| b.1)
154}
155
156impl LimitSource for CodexRollouts {
157    fn provider(&self) -> &'static str {
158        "codex"
159    }
160    fn read(&self, now_ms: u64) -> Option<ProviderLimits> {
161        let file = newest_rollout(&self.root, 4)?;
162        let text = std::fs::read_to_string(file).ok()?;
163        text.lines().rev().filter(|l| l.contains("rate_limits")).find_map(|l| parse_codex_line(l, now_ms))
164    }
165}
166
167/// Latest screen-derived limits per provider (claude today).
168static SCREEN: OnceLock<Mutex<BTreeMap<String, ProviderLimits>>> = OnceLock::new();
169
170pub fn now_ms() -> u64 {
171    std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map_or(0, |d| d.as_millis() as u64)
172}
173
174/// Called from the frame publisher with every scanned screen.
175pub fn observe_screen(provider: &str, screen: &str) {
176    if provider != "claude" || !(screen.contains("Current session") || screen.contains("Current week") || screen.contains("Claude ")) {
177        return;
178    }
179    let Some(new) = parse_claude_screen(screen, now_ms()) else { return };
180    let mut map = SCREEN.get_or_init(Default::default).lock().unwrap_or_else(|p| p.into_inner());
181    let merged = match map.remove(provider) {
182        Some(mut old) => {
183            if new.tier.is_some() {
184                old.tier = new.tier.clone();
185            }
186            for w in new.windows {
187                old.windows.retain(|o| o.window != w.window);
188                old.windows.push(w);
189            }
190            old.observed_at_ms = new.observed_at_ms;
191            old
192        }
193        None => new,
194    };
195    map.insert(provider.to_owned(), merged);
196}
197
198/// Everything known now: screen-derived entries plus every file source.
199pub fn collect(sources: &[Box<dyn LimitSource>]) -> Vec<ProviderLimits> {
200    let now = now_ms();
201    let mut out: Vec<ProviderLimits> = SCREEN.get().map(|m| m.lock().unwrap_or_else(|p| p.into_inner()).values().cloned().collect()).unwrap_or_default();
202    out.extend(sources.iter().filter_map(|s| s.read(now)));
203    out
204}
205
206pub fn default_sources() -> Vec<Box<dyn LimitSource>> {
207    let mut v: Vec<Box<dyn LimitSource>> = Vec::new();
208    if let Some(c) = CodexRollouts::from_home() {
209        v.push(Box::new(c));
210    }
211    v
212}
213
214#[cfg(test)]
215mod tests {
216    use super::*;
217
218    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";
219
220    #[test]
221    fn claude_usage_screen() {
222        let l = parse_claude_screen(USAGE, 7).unwrap();
223        assert_eq!(l.windows.len(), 2);
224        assert_eq!(l.windows[0].window, "5h");
225        assert_eq!(l.windows[0].used_percent, Some(11.0));
226        assert_eq!(l.windows[0].resets_text.as_deref(), Some("Resets 3pm (Asia/Yekaterinburg)"));
227        assert_eq!(l.windows[1].window, "7d");
228        assert_eq!(l.windows[1].used_percent, Some(22.0));
229    }
230
231    /// Real claude /usage screen captured on hatchery-box (2026-10-19).
232    #[test]
233    fn claude_usage_screen_live_capture() {
234        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";
235        let l = parse_claude_screen(screen, 0).unwrap();
236        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))]);
237        assert_eq!(l.windows[1].resets_text.as_deref(), Some("Resets Oct 21, 11:59pm (Asia/Yekaterinburg)"));
238    }
239
240    #[test]
241    fn claude_status_tier_and_noise() {
242        assert_eq!(parse_claude_screen("Login method: Claude Max Account", 1).unwrap().tier.as_deref(), Some("max"));
243        assert!(parse_claude_screen("> fix the tests\n● done", 1).is_none());
244    }
245
246    #[test]
247    fn codex_rollout_line() {
248        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"}}}"#;
249        let l = parse_codex_line(line, 1000).unwrap();
250        assert_eq!(l.tier.as_deref(), Some("plus"));
251        assert_eq!(l.windows[0].window, "5h");
252        assert_eq!(l.windows[0].resets_at_ms, Some(61_000));
253        assert_eq!(l.windows[1].window, "7d");
254        assert_eq!(l.windows[1].resets_at_ms, Some(1_800_000_000_000));
255    }
256
257    #[test]
258    fn codex_source_reads_newest_file() {
259        let d = std::env::temp_dir().join(format!("pl-{}", std::process::id()));
260        let day = d.join("2026/10/19");
261        std::fs::create_dir_all(&day).unwrap();
262        std::fs::write(day.join("rollout-a.jsonl"), "{}\n{\"payload\":{\"rate_limits\":{\"primary\":{\"used_percent\":5,\"window_minutes\":300}}}}\n").unwrap();
263        let l = CodexRollouts { root: d.clone() }.read(0).unwrap();
264        assert_eq!(l.windows[0].used_percent, Some(5.0));
265        let _ = std::fs::remove_dir_all(d);
266    }
267}