hatchery_node/
pty_resume.rs

1//! Node session registry: PTY sessions the user opened survive a node restart.
2//! Each live PTY session is journalled (provider, provider session id, workspace, model/effort,
3//! instance number); on node start every journalled session is spawned again with the
4//! provider's own resume flag, under the same instance number, so HQ shows the same chat.
5//! Off with `HATCHERY_SESSION_RESUME=0`.
6use serde::{Deserialize, Serialize};
7use std::path::{Path, PathBuf};
8use std::sync::Mutex;
9
10#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
11pub struct Entry {
12    pub instance_id: u64,
13    pub provider: String,
14    pub workspace_id: String,
15    pub provider_session_id: String,
16    /// m4a session key in the child's env (`M4A_SESSION_ID`): maps the runtime to its transcript.
17    #[serde(default)]
18    pub session_key: String,
19    #[serde(default)]
20    pub cwd: String,
21    /// `acp` for an ACP runtime (resumed through its managed record); PTY otherwise.
22    #[serde(default)]
23    pub acp: bool,
24    #[serde(default)]
25    pub options: Option<gate4agent_types::SessionOptionSelection>,
26    #[serde(default)]
27    pub cols: u16,
28    #[serde(default)]
29    pub rows: u16,
30}
31
32static LOCK: Mutex<()> = Mutex::new(());
33
34pub fn enabled() -> bool {
35    std::env::var("HATCHERY_SESSION_RESUME").map_or(true, |v| v != "0")
36}
37
38/// Providers whose CLI accepts a caller-chosen session id plus a resume flag.
39pub fn supported(provider: &str) -> bool {
40    provider == "claude"
41}
42
43pub fn path(state_dir: &Path) -> PathBuf {
44    state_dir.join("session-resume.json")
45}
46
47pub fn load(file: &Path) -> Vec<Entry> {
48    std::fs::read(file).ok().and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default()
49}
50
51fn save(file: &Path, entries: &[Entry]) {
52    let tmp = file.with_extension("json.tmp");
53    match serde_json::to_vec_pretty(entries) {
54        Ok(b) => {
55            if std::fs::write(&tmp, b).and_then(|_| std::fs::rename(&tmp, file)).is_err() {
56                tracing::warn!(file = %file.display(), "session resume journal not writable");
57            }
58        }
59        Err(_) => {}
60    }
61}
62
63pub fn upsert(file: &Path, entry: Entry) {
64    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
65    let mut all = load(file);
66    all.retain(|e| e.instance_id != entry.instance_id);
67    all.push(entry);
68    all.sort_by_key(|e| e.instance_id);
69    save(file, &all);
70}
71
72/// Set the provider session id of an entry that has none yet.
73pub fn set_session_id(file: &Path, instance_id: u64, sid: &str) {
74    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
75    let mut all = load(file);
76    let mut changed = false;
77    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.provider_session_id.is_empty()) {
78        e.provider_session_id = sid.to_owned();
79        changed = true;
80    }
81    if changed { save(file, &all); }
82}
83
84/// Record the m4a session key an entry's runtime got (a resume reuses the old one), so the
85/// session keeps its account and nick across node restarts.
86pub fn set_session_key(file: &Path, instance_id: u64, key: &str) {
87    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
88    let mut all = load(file);
89    let mut changed = false;
90    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.session_key != key) {
91        e.session_key = key.to_owned();
92        changed = true;
93    }
94    if changed { save(file, &all); }
95}
96
97pub fn remove(file: &Path, instance_id: u64) {
98    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
99    let mut all = load(file);
100    let n = all.len();
101    all.retain(|e| e.instance_id != instance_id);
102    if all.len() != n {
103        save(file, &all);
104    }
105}
106
107/// Fill provider session ids (and cwd) from the providers' live-session files: a claude child
108/// writes `~/.claude/sessions/<pid>.json`; its env carries the node's `M4A_SESSION_ID`.
109/// The m4a account id each node session key got (the child sees it as `M4A_SESSION_ID`).
110static ACCOUNT_REFS: std::sync::Mutex<Vec<(String, String)>> = std::sync::Mutex::new(Vec::new());
111
112/// Records that the child of `session_key` carries `account_ref` as its `M4A_SESSION_ID`.
113pub fn note_account_ref(session_key: &str, account_ref: &str) {
114    let mut m = ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner());
115    m.retain(|(k, _)| k != session_key);
116    m.push((session_key.to_owned(), account_ref.to_owned()));
117}
118
119fn account_ref_of(session_key: &str) -> Option<String> {
120    ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner()).iter().find(|(k, _)| k == session_key).map(|(_, r)| r.clone())
121}
122
123pub fn refresh(file: &Path, homes: &[PathBuf]) {
124    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
125    let mut all = load(file);
126    if all.iter().all(|e| e.session_key.is_empty()) { return; }
127    let mut found: Vec<(String, String, String)> = Vec::new();
128    for h in homes {
129        let Ok(rd) = std::fs::read_dir(h.join(".claude/sessions")) else { continue };
130        for d in rd.flatten() {
131            let Ok(b) = std::fs::read(d.path()) else { continue };
132            let Ok(v) = serde_json::from_slice::<serde_json::Value>(&b) else { continue };
133            let (Some(pid), Some(sid)) = (v["pid"].as_u64(), v["sessionId"].as_str()) else { continue };
134            let Ok(env) = std::fs::read(format!("/proc/{pid}/environ")) else { continue };
135            let key = env.split(|c| *c == 0).find_map(|kv| kv.strip_prefix(b"M4A_SESSION_ID=")).map(|k| String::from_utf8_lossy(k).into_owned());
136            if let Some(key) = key {
137                found.push((key, sid.to_owned(), v["cwd"].as_str().unwrap_or_default().to_owned()));
138            }
139        }
140    }
141    let mut changed = false;
142    for e in all.iter_mut() {
143        let alt = account_ref_of(&e.session_key);
144        if let Some((_, sid, cwd)) = found.iter().find(|(k, _, _)| !e.session_key.is_empty() && (*k == e.session_key || alt.as_deref() == Some(k.as_str()))) {
145            if e.provider_session_id != *sid || e.cwd != *cwd {
146                e.provider_session_id = sid.clone();
147                e.cwd = cwd.clone();
148                changed = true;
149            }
150        }
151    }
152    if changed { save(file, &all); }
153}
154
155/// Launch args: resume a known id, or pin a fresh one so it can be resumed later.
156pub fn launch_args(provider: &str, session_id: &str, resume: bool) -> Vec<String> {
157    match (provider, resume) {
158        ("claude", true) => vec!["--resume".into(), session_id.into()],
159        ("claude", false) => vec!["--session-id".into(), session_id.into()],
160        _ => Vec::new(),
161    }
162}
163
164pub fn new_session_id() -> String {
165    let mut b = [0u8; 16];
166    if getrandom_fill(&mut b).is_err() {
167        let t = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_nanos()).unwrap_or(0);
168        b.copy_from_slice(&(t ^ (std::process::id() as u128) << 64).to_le_bytes());
169    }
170    b[6] = (b[6] & 0x0f) | 0x40;
171    b[8] = (b[8] & 0x3f) | 0x80;
172    let h: String = b.iter().map(|x| format!("{x:02x}")).collect();
173    format!("{}-{}-{}-{}-{}", &h[0..8], &h[8..12], &h[12..16], &h[16..20], &h[20..32])
174}
175
176fn getrandom_fill(b: &mut [u8]) -> std::io::Result<()> {
177    use std::io::Read;
178    std::fs::File::open("/dev/urandom")?.read_exact(b)
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184    #[test]
185    fn journal_roundtrip_and_args() {
186        let dir = std::env::temp_dir().join(format!("pty-resume-{}", std::process::id()));
187        std::fs::create_dir_all(&dir).unwrap();
188        let f = path(&dir);
189        let id = new_session_id();
190        assert_eq!(id.len(), 36);
191        upsert(&f, Entry { instance_id: 1, provider: "claude".into(), workspace_id: "primary".into(), provider_session_id: id.clone(), session_key: String::new(), cwd: String::new(), acp: false, options: None, cols: 0, rows: 0 });
192        assert_eq!(load(&f).len(), 1);
193        remove(&f, 1);
194        assert!(load(&f).is_empty());
195        assert_eq!(launch_args("claude", &id, true)[0], "--resume");
196        assert!(launch_args("codex", &id, true).is_empty());
197    }
198}
hatchery_node/pty_resume.rs

hatchery_node/
pty_resume.rs
