hatchery_node/

limit_heartbeat.rs

//! Limit heartbeat: the only way the node reads provider limits on purpose.
//!
//! One light service PTY per signed-in provider account (claude, codex, grok, kimi), outside the node's
//! managed sessions: never shown in HQ, no record, no m4a account, no workspace. It runs the
//! provider CLI in its own scratch folder and, every interval (env
//! `HATCHERY_LIMIT_HEARTBEAT_SECS`, default 45, 15..=3600; `0` turns it off), types the
//! provider's own local limit screen (`/usage` for claude, grok and kimi, `/status` for codex), reads it off
//! the screen into [`crate::provider_limits`] (source `heartbeat`) and closes it with Escape.
//! Only slash commands are typed: no prompt, no model turn. A CLI that is not signed in
//! (`none`, or a screen that asks to log in) gets no heartbeat until a later check. grok and
//! kimi have no `auth status`: their limit screen is the sign-in test. Tokens, OAuth and
//! credential files are never touched.
//!
//! A first start in the scratch folder may ask whether to trust it; the heartbeat answers the
//! default ("yes") for its own folder only.

use gate4agent_pty::{native_pty_system, CommandBuilder, PtySize};
use std::io::{Read, Write};
use std::path::PathBuf;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};

pub const PROVIDERS: &[&str] = &["claude", "codex", "grok", "kimi"];
const ROWS: u16 = 50;
const COLS: u16 = 140;
/// How long a typed limit screen may take to show.
const SCREEN_WAIT: Duration = Duration::from_secs(10);
/// Wait after a start before the first command (the CLI draws its prompt).
const BOOT_WAIT: Duration = Duration::from_secs(8);

/// A CLI start may itself cost the account (start-up requests), so the PTY is kept for the
/// life of the node and a new one is started at most this often; the readings in between are
/// slash commands into the same CLI.
pub const RELAUNCH_MIN: Duration = Duration::from_secs(600);

/// Wait before the next start after a session that lived `lived`.
pub fn relaunch_after(lived: Duration, last: Duration) -> Duration {
    if lived > Duration::from_secs(3600) { RELAUNCH_MIN } else { (last * 2).clamp(RELAUNCH_MIN, Duration::from_secs(3600)) }
}

/// The heartbeat interval: `None` when off.
pub fn interval_from(v: Option<&str>) -> Option<Duration> {
    match v.map(str::trim).filter(|v| !v.is_empty()) {
        None => Some(Duration::from_secs(45)),
        Some(v) => match v.parse::<u64>() {
            Ok(0) => None,
            Ok(n) => Some(Duration::from_secs(n.clamp(15, 3600))),
            Err(_) => Some(Duration::from_secs(45)),
        },
    }
}

/// The slash command that shows `provider`'s limits.
pub fn command_for(provider: &str) -> Option<&'static str> {
    match provider {
        "claude" | "grok" | "kimi" => Some("/usage"),
        "codex" => Some("/status"),
        _ => None,
    }
}

/// Keys that answer a start-up question about trusting the heartbeat's own scratch folder
/// with "yes": Enter when the selected option already says yes, else Down then Enter (claude
/// selects "No, exit" first). `None` when the screen asks nothing of the kind.
pub fn trust_keys(screen: &str) -> Option<&'static [u8]> {
    let low = screen.to_lowercase();
    // grok: "Grok Build may run or modify contents in this directory" with `y` / `n` keys.
    if low.contains("grok build may run or modify") {
        return Some(b"y\r");
    }
    let asks = ["i trust this folder", "one you trust", "trust the files in this folder", "allow codex to work in this folder", "trust this folder?"].iter().any(|p| low.contains(p));
    if !asks {
        return None;
    }
    let selected = low.lines().find(|l| l.trim_start().starts_with('❯') || l.trim_start().starts_with('›') || l.trim_start().starts_with("> "));
    Some(if selected.is_some_and(|l| l.contains("yes") || l.contains("trust and continue")) { b"\r" } else { b"\x1b[B\r" })
}

/// How a provider's sign-in is judged before its heartbeat starts: claude and codex by their
/// `auth`/`login status`; grok and kimi have none, so only an installed CLI is needed and
/// the limit screen itself tells.
fn sign_in_known(provider: &str) -> bool {
    matches!(provider, "claude" | "codex")
}

/// Run the heartbeat for every signed-in provider (one thread each). Returns at once.
/// Base of the scratch folders: `HATCHERY_LIMIT_HEARTBEAT_DIR`, else `limit-heartbeat` next to
/// the node state file (set once by [`start`]), else the temp folder.
static STATE_DIR: std::sync::OnceLock<PathBuf> = std::sync::OnceLock::new();

pub fn scratch_base(env_dir: Option<PathBuf>, state_dir: Option<&std::path::Path>) -> PathBuf {
    env_dir.or_else(|| state_dir.map(|d| d.join("limit-heartbeat"))).unwrap_or_else(|| std::env::temp_dir().join("hatchery-limit-heartbeat"))
}

/// `state_dir`: the folder of the node state file.
pub fn start(state_dir: Option<PathBuf>) {
    if let Some(d) = state_dir {
        let _ = STATE_DIR.set(d);
    }
    let Some(every) = interval_from(std::env::var("HATCHERY_LIMIT_HEARTBEAT_SECS").ok().as_deref()) else {
        tracing::info!("limit heartbeat off");
        return;
    };
    for provider in PROVIDERS {
        std::thread::Builder::new()
            .name(format!("limit-heartbeat-{provider}"))
            .spawn(move || run(provider, every))
            .ok();
    }
}

fn scratch_dir(provider: &str) -> Option<PathBuf> {
    let base = scratch_base(std::env::var_os("HATCHERY_LIMIT_HEARTBEAT_DIR").map(PathBuf::from), STATE_DIR.get().map(PathBuf::as_path));
    let dir = base.join(provider);
    std::fs::create_dir_all(&dir).ok()?;
    Some(dir)
}

/// One provider's heartbeat, for the life of the node: (re)start its PTY when it is gone.
fn run(provider: &'static str, every: Duration) {
    let mut backoff = RELAUNCH_MIN;
    loop {
        let kind = if sign_in_known(provider) {
            crate::provider_inventory::account_kind(provider)
        } else {
            crate::provider_inventory::resolve_binary(provider, &[]).map(|_| "screen".to_owned())
        };
        match kind.as_deref() {
            None => {
                tracing::info!(provider, "limit heartbeat: no CLI or account kind; not started");
                return;
            }
            Some("none") => {
                tracing::info!(provider, "limit heartbeat: not signed in; checking again later");
                std::thread::sleep(Duration::from_secs(1800));
                continue;
            }
            Some(_) => {}
        }
        let started = Instant::now();
        match session(provider, every) {
            Err(e) if e == NOT_SIGNED_IN => {
                tracing::info!(provider, "limit heartbeat: the CLI asks to sign in; checking again in an hour");
                std::thread::sleep(Duration::from_secs(3600));
                continue;
            }
            Err(e) => tracing::warn!(provider, error = %e, "limit heartbeat session ended"),
            Ok(()) => {}
        }
        backoff = relaunch_after(started.elapsed(), backoff);
        std::thread::sleep(backoff);
    }
}

const NOT_SIGNED_IN: &str = "not signed in";

fn session(provider: &'static str, every: Duration) -> Result<(), String> {
    let bin = crate::provider_inventory::resolve_binary(provider, &[]).ok_or("no binary")?;
    let dir = scratch_dir(provider).ok_or("no scratch folder")?;
    let command = command_for(provider).ok_or("no limit command")?;
    let pair = native_pty_system().openpty(PtySize { rows: ROWS, cols: COLS, pixel_width: 0, pixel_height: 0 }).map_err(|e| e.to_string())?;
    let mut cmd = CommandBuilder::new(&bin);
    cmd.cwd(&dir);
    // A plain terminal, and never a colour-less one (NO_COLOR changes nothing we read, but
    // keep the CLI's normal screen).
    cmd.env("TERM", "xterm-256color");
    cmd.env_remove("NO_COLOR");
    let mut child = pair.slave.spawn_command(cmd).map_err(|e| e.to_string())?;
    drop(pair.slave);
    let parser = Arc::new(Mutex::new(vt100::Parser::new(ROWS, COLS, 0)));
    let mut reader = pair.master.try_clone_reader().map_err(|e| e.to_string())?;
    let feed = Arc::clone(&parser);
    std::thread::spawn(move || {
        let mut buf = [0u8; 8192];
        while let Ok(n) = reader.read(&mut buf) {
            if n == 0 {
                break;
            }
            feed.lock().unwrap_or_else(|p| p.into_inner()).process(&buf[..n]);
        }
    });
    let mut writer = pair.master.take_writer().map_err(|e| e.to_string())?;
    let screen = || parser.lock().unwrap_or_else(|p| p.into_inner()).screen().contents();
    let mut send = |bytes: &[u8]| writer.write_all(bytes).and_then(|_| writer.flush()).map_err(|e| e.to_string());
    tracing::info!(provider, "limit heartbeat started");
    let result = (|| -> Result<(), String> {
        std::thread::sleep(BOOT_WAIT);
        // A CLI that opens on a sign-in screen (grok starts a device login by itself): leave
        // at once, type nothing.
        if crate::provider_limits::says_not_signed_in(provider, &screen()) {
            return Err(NOT_SIGNED_IN.into());
        }
        if let Some(keys) = trust_keys(&screen()) {
            send(&keys[..keys.len() - 1])?;
            std::thread::sleep(Duration::from_millis(400));
            send(b"\r")?;
            std::thread::sleep(BOOT_WAIT);
        }
        loop {
            if child.try_wait().map_err(|e| e.to_string())?.is_some() {
                return Err("CLI exited".into());
            }
            // Type the command, then Enter on its own (a TUI may read a pasted line with Enter
            // as text).
            send(command.as_bytes())?;
            std::thread::sleep(Duration::from_millis(400));
            send(b"\r")?;
            let deadline = Instant::now() + SCREEN_WAIT;
            let mut read = false;
            while Instant::now() < deadline {
                std::thread::sleep(Duration::from_millis(500));
                if crate::provider_limits::observe_screen_as(provider, &screen(), "heartbeat") {
                    read = true;
                    break;
                }
            }
            if !read {
                if crate::provider_limits::says_not_signed_in(provider, &screen()) {
                    return Err(NOT_SIGNED_IN.into());
                }
                tracing::debug!(provider, "limit heartbeat: no limit screen this time");
            }
            // Close the panel (and any half-typed command).
            send(b"\x1b")?;
            std::thread::sleep(Duration::from_millis(300));
            send(b"\x15")?;
            std::thread::sleep(every);
        }
    })();
    // Never leave the CLI behind.
    let _ = child.kill();
    let _ = child.wait();
    result
}

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

    #[test]
    fn scratch_lives_in_node_state_dir() {
        let st = std::path::Path::new("/var/lib/node");
        assert_eq!(scratch_base(None, Some(st)), st.join("limit-heartbeat"));
        assert_eq!(scratch_base(Some("/x".into()), Some(st)), PathBuf::from("/x"));
        assert_eq!(scratch_base(None, None), std::env::temp_dir().join("hatchery-limit-heartbeat"));
    }

    #[test]
    fn relaunch_is_rare() {
        let m = Duration::from_secs(60);
        assert_eq!(relaunch_after(Duration::from_secs(5), RELAUNCH_MIN), RELAUNCH_MIN * 2);
        assert_eq!(relaunch_after(Duration::from_secs(5), Duration::from_secs(1)), RELAUNCH_MIN);
        assert_eq!(relaunch_after(Duration::from_secs(5), 50 * m), 60 * m);
        assert_eq!(relaunch_after(2 * 60 * m, 60 * m), RELAUNCH_MIN);
    }

    #[test]
    fn interval_is_configurable_and_bounded() {
        assert_eq!(interval_from(None), Some(Duration::from_secs(45)));
        assert_eq!(interval_from(Some("30")), Some(Duration::from_secs(30)));
        assert_eq!(interval_from(Some("1")), Some(Duration::from_secs(15)));
        assert_eq!(interval_from(Some("99999")), Some(Duration::from_secs(3600)));
        assert_eq!(interval_from(Some("0")), None);
        assert_eq!(interval_from(Some("x")), Some(Duration::from_secs(45)));
    }

    #[test]
    fn commands_and_trust_prompt() {
        assert_eq!((command_for("claude"), command_for("codex"), command_for("grok"), command_for("kimi"), command_for("bash")), (Some("/usage"), Some("/status"), Some("/usage"), Some("/usage"), None));
        let grok = "Do you trust the contents of this directory?\n Grok Build may run or modify contents in this directory,\n posing security risks.\n   Yes, proceed                 y\n   No, quit                     n";
        assert_eq!(trust_keys(grok), Some(&b"y\r"[..]));
        let kimi = "  Trust this folder?\n  ↑↓ navigate · Enter select · Esc exit\n   ❯ Trust and continue\n     Exit";
        assert_eq!(trust_keys(kimi), Some(&b"\r"[..]));
        assert!(sign_in_known("claude") && sign_in_known("codex") && !sign_in_known("grok") && !sign_in_known("kimi"));
        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";
        assert_eq!(trust_keys(claude), Some(&b"\x1b[B\r"[..]));
        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";
        assert_eq!(trust_keys(codex), Some(&b"\r"[..]));
        assert_eq!(trust_keys("> what is trust?"), None);
    }
}

/// Live check on a machine with a signed-in claude: `cargo test -p hatchery-node --lib
/// limit_heartbeat::live -- --ignored --nocapture`. Types only `/usage`.
#[cfg(test)]
mod live {
    #[test]
    #[ignore]
    fn claude_heartbeat_reads_usage() {
        std::thread::spawn(|| super::session("claude", std::time::Duration::from_secs(3600)));
        let deadline = std::time::Instant::now() + std::time::Duration::from_secs(60);
        while std::time::Instant::now() < deadline {
            if let Some(l) = crate::provider_limits::collect().into_iter().find(|l| l.provider == "claude" && l.source == "heartbeat") {
                println!("heartbeat: {:?}", l.windows.iter().map(|w| (&w.window, w.used_percent)).collect::<Vec<_>>());
                return;
            }
            std::thread::sleep(std::time::Duration::from_secs(1));
        }
        panic!("no /usage reading within 60s");
    }
}