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

hatchery_node/
pty_resume.rs
