Skip to main content

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/// Node start: stops per-session leaders a previous run of this node left behind (its sessions
69/// died with it). Only processes carrying [`SOCKET_PREFIX`] are matched; an operator's own
70/// leader never carries it.
71pub fn sweep_orphans() {
72    if disabled() {
73        return;
74    }
75    match stop_tagged_processes(SOCKET_PREFIX) {
76        Ok(()) => tracing::info!("orphaned per-session grok leaders swept"),
77        Err(error) => tracing::warn!(%error, "orphaned grok leader sweep failed"),
78    }
79}
80
81#[cfg(windows)]
82fn stop_tagged_processes(tag: &str) -> Result<(), String> {
83    // `tag` is validated `[A-Za-z0-9-]` only, so it is inert inside the filter string.
84    let script = format!(
85        "Get-CimInstance Win32_Process | Where-Object {{ $_.CommandLine -and $_.CommandLine.Contains('{tag}') -and $_.ProcessId -ne $PID }} | ForEach-Object {{ Stop-Process -Id $_.ProcessId -Force -ErrorAction SilentlyContinue }}"
86    );
87    let mut cmd = std::process::Command::new("powershell.exe");
88    cmd.args(["-NoProfile", "-NonInteractive", "-Command", &script])
89        .stdin(std::process::Stdio::null())
90        .stdout(std::process::Stdio::null())
91        .stderr(std::process::Stdio::null());
92    #[cfg(windows)]
93    {
94        use std::os::windows::process::CommandExt;
95        cmd.creation_flags(0x0800_0000);
96    }
97    let status = cmd.status()
98        .map_err(|e| e.kind().to_string())?;
99    status.success().then_some(()).ok_or_else(|| format!("powershell exit {status}"))
100}
101
102#[cfg(not(windows))]
103fn stop_tagged_processes(tag: &str) -> Result<(), String> {
104    let status = std::process::Command::new("pkill")
105        .args(["-TERM", "-f", tag])
106        .stdin(std::process::Stdio::null())
107        .stdout(std::process::Stdio::null())
108        .stderr(std::process::Stdio::null())
109        .status()
110        .map_err(|e| e.kind().to_string())?;
111    // pkill: 0 = signalled, 1 = nothing matched; both are fine.
112    matches!(status.code(), Some(0 | 1)).then_some(()).ok_or_else(|| format!("pkill exit {status}"))
113}
114
115#[cfg(test)]
116mod tests {
117    use super::*;
118
119    #[test]
120    fn socket_is_per_session_and_tagged() {
121        let root = std::env::temp_dir().join("hatchery-grok-leader-test");
122        let a = socket_path(&root, "inc-1");
123        let b = socket_path(&root, "inc-2");
124        assert_ne!(a, b);
125        assert!(a.starts_with(root.join("grok-leaders")));
126        assert!(retire_tag(&a).is_some());
127        assert_eq!(leader_socket_args(&a)[0], OsString::from("--leader-socket"));
128    }
129
130    #[test]
131    fn retire_refuses_foreign_sockets() {
132        assert!(retire_tag(Path::new("/home/user/.grok/leader.sock")).is_none());
133        assert!(retire_tag(Path::new("/tmp/hatchery-grok-leader-zz.sock")).is_none());
134        assert!(retire_tag(Path::new("/tmp/hatchery-grok-leader-0123456789abcdef'x.sock")).is_none());
135    }
136
137    #[cfg(unix)]
138    #[test]
139    fn retire_stops_only_the_tagged_process() {
140        let tag = session_tag("retire-test");
141        let mut mine = std::process::Command::new("sh")
142            .args(["-c", &format!("sleep 30; : {tag}")])
143            .spawn()
144            .unwrap();
145        let mut other = std::process::Command::new("sleep").arg("30").spawn().unwrap();
146        retire(&std::env::temp_dir().join(format!("{tag}.sock")));
147        let mut stopped = false;
148        for _ in 0..50 {
149            if mine.try_wait().unwrap().is_some() {
150                stopped = true;
151                break;
152            }
153            std::thread::sleep(std::time::Duration::from_millis(100));
154        }
155        assert!(stopped, "tagged process must be stopped");
156        assert!(other.try_wait().unwrap().is_none());
157        let _ = other.kill();
158        let _ = other.wait();
159    }
160}