hatchery_node/
limit_heartbeat.rs

1//! Limit heartbeat: the only way the node reads provider limits on purpose.
2//!
3//! One light service PTY per signed-in provider account (claude, codex), outside the node's
4//! managed sessions: never shown in HQ, no record, no m4a account, no workspace. It runs the
5//! provider CLI in its own scratch folder and, every interval (env
6//! `HATCHERY_LIMIT_HEARTBEAT_SECS`, default 45, 15..=3600; `0` turns it off), types the
7//! provider's own local limit screen (`/usage` for claude, `/status` for codex), reads it off
8//! the screen into [`crate::provider_limits`] (source `heartbeat`) and closes it with Escape.
9//! Only slash commands are typed: no prompt, no model turn. A CLI that is not signed in
10//! (`none`) gets no heartbeat. Tokens, OAuth and credential files are never touched.
11//!
12//! A first start in the scratch folder may ask whether to trust it; the heartbeat answers the
13//! default ("yes") for its own folder only.
14
15use gate4agent_pty::{native_pty_system, CommandBuilder, PtySize};
16use std::io::{Read, Write};
17use std::path::PathBuf;
18use std::sync::{Arc, Mutex};
19use std::time::{Duration, Instant};
20
21pub const PROVIDERS: &[&str] = &["claude", "codex"];
22const ROWS: u16 = 50;
23const COLS: u16 = 140;
24/// How long a typed limit screen may take to show.
25const SCREEN_WAIT: Duration = Duration::from_secs(10);
26/// Wait after a start before the first command (the CLI draws its prompt).
27const BOOT_WAIT: Duration = Duration::from_secs(8);
28
29/// A CLI start may itself cost the account (start-up requests), so the PTY is kept for the
30/// life of the node and a new one is started at most this often; the readings in between are
31/// slash commands into the same CLI.
32pub const RELAUNCH_MIN: Duration = Duration::from_secs(600);
33
34/// Wait before the next start after a session that lived `lived`.
35pub fn relaunch_after(lived: Duration, last: Duration) -> Duration {
36    if lived > Duration::from_secs(3600) { RELAUNCH_MIN } else { (last * 2).clamp(RELAUNCH_MIN, Duration::from_secs(3600)) }
37}
38
39/// The heartbeat interval: `None` when off.
40pub fn interval_from(v: Option<&str>) -> Option<Duration> {
41    match v.map(str::trim).filter(|v| !v.is_empty()) {
42        None => Some(Duration::from_secs(45)),
43        Some(v) => match v.parse::<u64>() {
44            Ok(0) => None,
45            Ok(n) => Some(Duration::from_secs(n.clamp(15, 3600))),
46            Err(_) => Some(Duration::from_secs(45)),
47        },
48    }
49}
50
51/// The slash command that shows `provider`'s limits.
52pub fn command_for(provider: &str) -> Option<&'static str> {
53    match provider {
54        "claude" => Some("/usage"),
55        "codex" => Some("/status"),
56        _ => None,
57    }
58}
59
60/// Keys that answer a start-up question about trusting the heartbeat's own scratch folder
61/// with "yes": Enter when the selected option already says yes, else Down then Enter (claude
62/// selects "No, exit" first). `None` when the screen asks nothing of the kind.
63pub fn trust_keys(screen: &str) -> Option<&'static [u8]> {
64    let low = screen.to_lowercase();
65    let asks = ["i trust this folder", "one you trust", "trust the files in this folder", "allow codex to work in this folder"].iter().any(|p| low.contains(p));
66    if !asks {
67        return None;
68    }
69    let selected = low.lines().find(|l| l.trim_start().starts_with('❯') || l.trim_start().starts_with('›') || l.trim_start().starts_with("> "));
70    Some(if selected.is_some_and(|l| l.contains("yes")) { b"\r" } else { b"\x1b[B\r" })
71}
72
73/// Run the heartbeat for every signed-in provider (one thread each). Returns at once.
74/// Base of the scratch folders: `HATCHERY_LIMIT_HEARTBEAT_DIR`, else `limit-heartbeat` next to
75/// the node state file (set once by [`start`]), else the temp folder.
76static STATE_DIR: std::sync::OnceLock<PathBuf> = std::sync::OnceLock::new();
77
78pub fn scratch_base(env_dir: Option<PathBuf>, state_dir: Option<&std::path::Path>) -> PathBuf {
79    env_dir.or_else(|| state_dir.map(|d| d.join("limit-heartbeat"))).unwrap_or_else(|| std::env::temp_dir().join("hatchery-limit-heartbeat"))
80}
81
82/// `state_dir`: the folder of the node state file.
83pub fn start(state_dir: Option<PathBuf>) {
84    if let Some(d) = state_dir {
85        let _ = STATE_DIR.set(d);
86    }
87    let Some(every) = interval_from(std::env::var("HATCHERY_LIMIT_HEARTBEAT_SECS").ok().as_deref()) else {
88        tracing::info!("limit heartbeat off");
89        return;
90    };
91    for provider in PROVIDERS {
92        std::thread::Builder::new()
93            .name(format!("limit-heartbeat-{provider}"))
94            .spawn(move || run(provider, every))
95            .ok();
96    }
97}
98
99fn scratch_dir(provider: &str) -> Option<PathBuf> {
100    let base = scratch_base(std::env::var_os("HATCHERY_LIMIT_HEARTBEAT_DIR").map(PathBuf::from), STATE_DIR.get().map(PathBuf::as_path));
101    let dir = base.join(provider);
102    std::fs::create_dir_all(&dir).ok()?;
103    Some(dir)
104}
105
106/// One provider's heartbeat, for the life of the node: (re)start its PTY when it is gone.
107fn run(provider: &'static str, every: Duration) {
108    let mut backoff = RELAUNCH_MIN;
109    loop {
110        match crate::provider_inventory::account_kind(provider).as_deref() {
111            None => {
112                tracing::info!(provider, "limit heartbeat: no CLI or account kind; not started");
113                return;
114            }
115            Some("none") => {
116                tracing::info!(provider, "limit heartbeat: not signed in; checking again later");
117                std::thread::sleep(Duration::from_secs(1800));
118                continue;
119            }
120            Some(_) => {}
121        }
122        let started = Instant::now();
123        if let Err(e) = session(provider, every) {
124            tracing::warn!(provider, error = %e, "limit heartbeat session ended");
125        }
126        backoff = relaunch_after(started.elapsed(), backoff);
127        std::thread::sleep(backoff);
128    }
129}
130
131fn session(provider: &'static str, every: Duration) -> Result<(), String> {
132    let bin = crate::provider_inventory::resolve_binary(provider, &[]).ok_or("no binary")?;
133    let dir = scratch_dir(provider).ok_or("no scratch folder")?;
134    let command = command_for(provider).ok_or("no limit command")?;
135    let pair = native_pty_system().openpty(PtySize { rows: ROWS, cols: COLS, pixel_width: 0, pixel_height: 0 }).map_err(|e| e.to_string())?;
136    let mut cmd = CommandBuilder::new(&bin);
137    cmd.cwd(&dir);
138    // A plain terminal, and never a colour-less one (NO_COLOR changes nothing we read, but
139    // keep the CLI's normal screen).
140    cmd.env("TERM", "xterm-256color");
141    cmd.env_remove("NO_COLOR");
142    let mut child = pair.slave.spawn_command(cmd).map_err(|e| e.to_string())?;
143    drop(pair.slave);
144    let parser = Arc::new(Mutex::new(vt100::Parser::new(ROWS, COLS, 0)));
145    let mut reader = pair.master.try_clone_reader().map_err(|e| e.to_string())?;
146    let feed = Arc::clone(&parser);
147    std::thread::spawn(move || {
148        let mut buf = [0u8; 8192];
149        while let Ok(n) = reader.read(&mut buf) {
150            if n == 0 {
151                break;
152            }
153            feed.lock().unwrap_or_else(|p| p.into_inner()).process(&buf[..n]);
154        }
155    });
156    let mut writer = pair.master.take_writer().map_err(|e| e.to_string())?;
157    let screen = || parser.lock().unwrap_or_else(|p| p.into_inner()).screen().contents();
158    let mut send = |bytes: &[u8]| writer.write_all(bytes).and_then(|_| writer.flush()).map_err(|e| e.to_string());
159    tracing::info!(provider, "limit heartbeat started");
160    let result = (|| -> Result<(), String> {
161        std::thread::sleep(BOOT_WAIT);
162        if let Some(keys) = trust_keys(&screen()) {
163            send(&keys[..keys.len() - 1])?;
164            std::thread::sleep(Duration::from_millis(400));
165            send(b"\r")?;
166            std::thread::sleep(BOOT_WAIT);
167        }
168        loop {
169            if child.try_wait().map_err(|e| e.to_string())?.is_some() {
170                return Err("CLI exited".into());
171            }
172            // Type the command, then Enter on its own (a TUI may read a pasted line with Enter
173            // as text).
174            send(command.as_bytes())?;
175            std::thread::sleep(Duration::from_millis(400));
176            send(b"\r")?;
177            let deadline = Instant::now() + SCREEN_WAIT;
178            let mut read = false;
179            while Instant::now() < deadline {
180                std::thread::sleep(Duration::from_millis(500));
181                if crate::provider_limits::observe_screen_as(provider, &screen(), "heartbeat") {
182                    read = true;
183                    break;
184                }
185            }
186            if !read {
187                tracing::debug!(provider, "limit heartbeat: no limit screen this time");
188            }
189            // Close the panel (and any half-typed command).
190            send(b"\x1b")?;
191            std::thread::sleep(Duration::from_millis(300));
192            send(b"\x15")?;
193            std::thread::sleep(every);
194        }
195    })();
196    // Never leave the CLI behind.
197    let _ = child.kill();
198    let _ = child.wait();
199    result
200}
201
202#[cfg(test)]
203mod tests {
204    use super::*;
205
206    #[test]
207    fn scratch_lives_in_node_state_dir() {
208        let st = std::path::Path::new("/var/lib/node");
209        assert_eq!(scratch_base(None, Some(st)), st.join("limit-heartbeat"));
210        assert_eq!(scratch_base(Some("/x".into()), Some(st)), PathBuf::from("/x"));
211        assert_eq!(scratch_base(None, None), std::env::temp_dir().join("hatchery-limit-heartbeat"));
212    }
213
214    #[test]
215    fn relaunch_is_rare() {
216        let m = Duration::from_secs(60);
217        assert_eq!(relaunch_after(Duration::from_secs(5), RELAUNCH_MIN), RELAUNCH_MIN * 2);
218        assert_eq!(relaunch_after(Duration::from_secs(5), Duration::from_secs(1)), RELAUNCH_MIN);
219        assert_eq!(relaunch_after(Duration::from_secs(5), 50 * m), 60 * m);
220        assert_eq!(relaunch_after(2 * 60 * m, 60 * m), RELAUNCH_MIN);
221    }
222
223    #[test]
224    fn interval_is_configurable_and_bounded() {
225        assert_eq!(interval_from(None), Some(Duration::from_secs(45)));
226        assert_eq!(interval_from(Some("30")), Some(Duration::from_secs(30)));
227        assert_eq!(interval_from(Some("1")), Some(Duration::from_secs(15)));
228        assert_eq!(interval_from(Some("99999")), Some(Duration::from_secs(3600)));
229        assert_eq!(interval_from(Some("0")), None);
230        assert_eq!(interval_from(Some("x")), Some(Duration::from_secs(45)));
231    }
232
233    #[test]
234    fn commands_and_trust_prompt() {
235        assert_eq!((command_for("claude"), command_for("codex"), command_for("grok")), (Some("/usage"), Some("/status"), None));
236        let claude = " Quick safety check: Is this a project you created or one you trust?\n ❯ No, exit\n   Yes, I trust this folder\n Enter to confirm";
237        assert_eq!(trust_keys(claude), Some(&b"\x1b[B\r"[..]));
238        let codex = "Do you trust the contents of this directory?\n› 1. Yes, allow Codex to work in this folder without asking for approval\n  2. No, quit";
239        assert_eq!(trust_keys(codex), Some(&b"\r"[..]));
240        assert_eq!(trust_keys("> what is trust?"), None);
241    }
242}
243
244/// Live check on a machine with a signed-in claude: `cargo test -p hatchery-node --lib
245/// limit_heartbeat::live -- --ignored --nocapture`. Types only `/usage`.
246#[cfg(test)]
247mod live {
248    #[test]
249    #[ignore]
250    fn claude_heartbeat_reads_usage() {
251        std::thread::spawn(|| super::session("claude", std::time::Duration::from_secs(3600)));
252        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
253        while std::time::Instant::now() < deadline {
254            if let Some(l) = crate::provider_limits::collect().into_iter().find(|l| l.provider == "claude" && l.source == "heartbeat") {
255                println!("heartbeat: {:?}", l.windows.iter().map(|w| (&w.window, w.used_percent)).collect::<Vec<_>>());
256                return;
257            }
258            std::thread::sleep(std::time::Duration::from_secs(1));
259        }
260        panic!("no /usage reading within 60s");
261    }
262}
hatchery_node/limit_heartbeat.rs

hatchery_node/
limit_heartbeat.rs
