hatchery_node/
grok_leader.rs

1//! Per-session grok leader isolation.
2//!
3//! grok's interactive UI may attach to a shared, long-lived "leader" backend (`use_leader` in
4//! its config). The leader then runs every tool call with *its own* process environment, so
5//! the node's per-session environment (for example the m4a session identity) never reaches
6//! the tools. A node-spawned grok session therefore gets its own `--leader-socket`: the backend
7//! it starts is that session's alone and inherits that session's environment. The grok home
8//! and its login are unchanged.
9//!
10//! The leader outlives its UI by design, so the node retires it when the session ends: it
11//! stops only processes whose command line carries this session's unique socket tag. Another
12//! leader (an operator's shared one included) never carries that tag and is never touched.
13
14use std::ffi::OsString;
15use std::hash::{Hash, Hasher};
16use std::path::{Path, PathBuf};
17
18/// Set to `0` / `off` to let node-spawned grok sessions use grok's own leader configuration.
19pub const SESSION_LEADER_ENV: &str = "HATCHERY_GROK_SESSION_LEADER";
20
21/// File-name prefix of every per-session socket; the tag the retire step matches on.
22pub const SOCKET_PREFIX: &str = "hatchery-grok-leader-";
23
24pub fn disabled() -> bool {
25    std::env::var(SESSION_LEADER_ENV)
26        .map(|v| matches!(v.trim().to_ascii_lowercase().as_str(), "0" | "off" | "false" | "no"))
27        .unwrap_or(false)
28}
29
30/// Short, stable tag for one session key (keeps Unix socket paths well under their limit).
31pub fn session_tag(session_key: &str) -> String {
32    let mut h = std::collections::hash_map::DefaultHasher::new();
33    session_key.hash(&mut h);
34    format!("{SOCKET_PREFIX}{:016x}", h.finish())
35}
36
37/// `<root>/grok-leaders/<tag>.sock`; the directory is created best effort.
38pub fn socket_path(root: &Path, session_key: &str) -> PathBuf {
39    let dir = root.join("grok-leaders");
40    let _ = std::fs::create_dir_all(&dir);
41    dir.join(format!("{}.sock", session_tag(session_key)))
42}
43
44pub fn leader_socket_args(socket: &Path) -> Vec<OsString> {
45    vec![OsString::from("--leader-socket"), socket.as_os_str().to_owned()]
46}
47
48/// The tag `retire` matches on, or `None` for a path that is not a per-session socket (so a
49/// foreign path can never widen the match).
50pub fn retire_tag(socket: &Path) -> Option<String> {
51    let stem = socket.file_stem()?.to_str()?;
52    let hex = stem.strip_prefix(SOCKET_PREFIX)?;
53    (hex.len() == 16 && hex.bytes().all(|b| b.is_ascii_hexdigit())).then(|| stem.to_owned())
54}
55
56/// Stops this session's own leader (and only it), then removes its socket file.
57pub fn retire(socket: &Path) {
58    let Some(tag) = retire_tag(socket) else { return };
59    let result = stop_tagged_processes(&tag);
60    if let Err(error) = result {
61        tracing::warn!(%error, "per-session grok leader retire failed");
62    } else {
63        tracing::info!("per-session grok leader retired");
64    }
65    let _ = std::fs::remove_file(socket);
66}
67
68#[cfg(windows)]
69fn stop_tagged_processes(tag: &str) -> Result<(), String> {
70    // `tag` is validated `[A-Za-z0-9-]` only, so it is inert inside the filter string.
71    let script = format!(
72        "Get-CimInstance Win32_Process | Where-Object {{ $_.CommandLine -and $_.CommandLine.Contains('{tag}') -and $_.ProcessId -ne $PID }} | ForEach-Object {{ Stop-Process -Id $_.ProcessId -Force -ErrorAction SilentlyContinue }}"
73    );
74    let status = std::process::Command::new("powershell.exe")
75        .args(["-NoProfile", "-NonInteractive", "-Command", &script])
76        .stdin(std::process::Stdio::null())
77        .stdout(std::process::Stdio::null())
78        .stderr(std::process::Stdio::null())
79        .status()
80        .map_err(|e| e.kind().to_string())?;
81    status.success().then_some(()).ok_or_else(|| format!("powershell exit {status}"))
82}
83
84#[cfg(not(windows))]
85fn stop_tagged_processes(tag: &str) -> Result<(), String> {
86    let status = std::process::Command::new("pkill")
87        .args(["-TERM", "-f", tag])
88        .stdin(std::process::Stdio::null())
89        .stdout(std::process::Stdio::null())
90        .stderr(std::process::Stdio::null())
91        .status()
92        .map_err(|e| e.kind().to_string())?;
93    // pkill: 0 = signalled, 1 = nothing matched; both are fine.
94    matches!(status.code(), Some(0 | 1)).then_some(()).ok_or_else(|| format!("pkill exit {status}"))
95}
96
97#[cfg(test)]
98mod tests {
99    use super::*;
100
101    #[test]
102    fn socket_is_per_session_and_tagged() {
103        let root = std::env::temp_dir().join("hatchery-grok-leader-test");
104        let a = socket_path(&root, "inc-1");
105        let b = socket_path(&root, "inc-2");
106        assert_ne!(a, b);
107        assert!(a.starts_with(root.join("grok-leaders")));
108        assert!(retire_tag(&a).is_some());
109        assert_eq!(leader_socket_args(&a)[0], OsString::from("--leader-socket"));
110    }
111
112    #[test]
113    fn retire_refuses_foreign_sockets() {
114        assert!(retire_tag(Path::new("/home/user/.grok/leader.sock")).is_none());
115        assert!(retire_tag(Path::new("/tmp/hatchery-grok-leader-zz.sock")).is_none());
116        assert!(retire_tag(Path::new("/tmp/hatchery-grok-leader-0123456789abcdef'x.sock")).is_none());
117    }
118
119    #[cfg(unix)]
120    #[test]
121    fn retire_stops_only_the_tagged_process() {
122        let tag = session_tag("retire-test");
123        let mut mine = std::process::Command::new("sh")
124            .args(["-c", &format!("sleep 30; : {tag}")])
125            .spawn()
126            .unwrap();
127        let mut other = std::process::Command::new("sleep").arg("30").spawn().unwrap();
128        retire(&std::env::temp_dir().join(format!("{tag}.sock")));
129        let mut stopped = false;
130        for _ in 0..50 {
131            if mine.try_wait().unwrap().is_some() {
132                stopped = true;
133                break;
134            }
135            std::thread::sleep(std::time::Duration::from_millis(100));
136        }
137        assert!(stopped, "tagged process must be stopped");
138        assert!(other.try_wait().unwrap().is_none());
139        let _ = other.kill();
140        let _ = other.wait();
141    }
142}
hatchery_node/grok_leader.rs

hatchery_node/
grok_leader.rs
