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
49/// Sessions journalled for restart resume: every ACP session (restored through its managed
50/// record with ACP `session/load`, else started fresh under the same number), and PTY sessions
51/// of providers with a resume flag. The journal is written at spawn and on every session-id or
52/// key change (tmp + rename), never on a stop signal, so a hard kill (Windows) keeps it.
53pub fn journalled(provider: &str, acp: bool) -> bool {
54    acp || supported(provider)
55}
56
57pub fn path(state_dir: &Path) -> PathBuf {
58    state_dir.join("session-resume.json")
59}
60
61pub fn load(file: &Path) -> Vec<Entry> {
62    std::fs::read(file).ok().and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default()
63}
64
65fn save(file: &Path, entries: &[Entry]) {
66    let tmp = file.with_extension("json.tmp");
67    match serde_json::to_vec_pretty(entries) {
68        Ok(b) => {
69            if std::fs::write(&tmp, b).and_then(|_| std::fs::rename(&tmp, file)).is_err() {
70                tracing::warn!(file = %file.display(), "session resume journal not writable");
71            }
72        }
73        Err(_) => {}
74    }
75}
76
77pub fn upsert(file: &Path, entry: Entry) {
78    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
79    let mut all = load(file);
80    all.retain(|e| e.instance_id != entry.instance_id);
81    all.push(entry);
82    all.sort_by_key(|e| e.instance_id);
83    save(file, &all);
84}
85
86/// Set the provider session id of an entry that has none yet.
87pub fn set_session_id(file: &Path, instance_id: u64, sid: &str) {
88    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
89    let mut all = load(file);
90    let mut changed = false;
91    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.provider_session_id.is_empty()) {
92        e.provider_session_id = sid.to_owned();
93        changed = true;
94    }
95    if changed { save(file, &all); }
96}
97
98/// Record the m4a session key an entry's runtime got (a resume reuses the old one), so the
99/// session keeps its account and nick across node restarts.
100pub fn set_session_key(file: &Path, instance_id: u64, key: &str) {
101    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
102    let mut all = load(file);
103    let mut changed = false;
104    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.session_key != key) {
105        e.session_key = key.to_owned();
106        changed = true;
107    }
108    if changed { save(file, &all); }
109}
110
111pub fn remove(file: &Path, instance_id: u64) {
112    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
113    let mut all = load(file);
114    let n = all.len();
115    all.retain(|e| e.instance_id != instance_id);
116    if all.len() != n {
117        save(file, &all);
118        if let Some(dir) = file.parent() {
119            crate::session_meta::mark_gone(dir, instance_id);
120        }
121    }
122}
123
124/// Fill provider session ids (and cwd) from the providers' live-session files: a claude child
125/// writes `~/.claude/sessions/<pid>.json`; its env carries the node's `M4A_SESSION_ID`.
126/// The m4a account id each node session key got (the child sees it as `M4A_SESSION_ID`).
127static ACCOUNT_REFS: std::sync::Mutex<Vec<(String, String)>> = std::sync::Mutex::new(Vec::new());
128
129/// Records that the child of `session_key` carries `account_ref` as its `M4A_SESSION_ID`.
130pub fn note_account_ref(session_key: &str, account_ref: &str) {
131    let mut m = ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner());
132    m.retain(|(k, _)| k != session_key);
133    m.push((session_key.to_owned(), account_ref.to_owned()));
134}
135
136fn account_ref_of(session_key: &str) -> Option<String> {
137    ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner()).iter().find(|(k, _)| k == session_key).map(|(_, r)| r.clone())
138}
139
140pub fn refresh(file: &Path, homes: &[PathBuf]) {
141    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
142    let mut all = load(file);
143    if all.iter().all(|e| e.session_key.is_empty()) { return; }
144    let mut found: Vec<(String, String, String)> = Vec::new();
145    for h in homes {
146        let Ok(rd) = std::fs::read_dir(h.join(".claude/sessions")) else { continue };
147        for d in rd.flatten() {
148            let Ok(b) = std::fs::read(d.path()) else { continue };
149            let Ok(v) = serde_json::from_slice::<serde_json::Value>(&b) else { continue };
150            let (Some(pid), Some(sid)) = (v["pid"].as_u64(), v["sessionId"].as_str()) else { continue };
151            let Ok(env) = std::fs::read(format!("/proc/{pid}/environ")) else { continue };
152            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());
153            if let Some(key) = key {
154                found.push((key, sid.to_owned(), v["cwd"].as_str().unwrap_or_default().to_owned()));
155            }
156        }
157    }
158    let mut changed = false;
159    for e in all.iter_mut() {
160        let alt = account_ref_of(&e.session_key);
161        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()))) {
162            if e.provider_session_id != *sid || e.cwd != *cwd {
163                e.provider_session_id = sid.clone();
164                e.cwd = cwd.clone();
165                changed = true;
166            }
167        }
168    }
169    if changed { save(file, &all); }
170}
171
172/// Launch args: resume a known id, or pin a fresh one so it can be resumed later.
173pub fn launch_args(provider: &str, session_id: &str, resume: bool) -> Vec<String> {
174    match (provider, resume) {
175        ("claude", true) => vec!["--resume".into(), session_id.into()],
176        ("claude", false) => vec!["--session-id".into(), session_id.into()],
177        _ => Vec::new(),
178    }
179}
180
181pub fn new_session_id() -> String {
182    let mut b = [0u8; 16];
183    if getrandom_fill(&mut b).is_err() {
184        let t = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_nanos()).unwrap_or(0);
185        b.copy_from_slice(&(t ^ (std::process::id() as u128) << 64).to_le_bytes());
186    }
187    b[6] = (b[6] & 0x0f) | 0x40;
188    b[8] = (b[8] & 0x3f) | 0x80;
189    let h: String = b.iter().map(|x| format!("{x:02x}")).collect();
190    format!("{}-{}-{}-{}-{}", &h[0..8], &h[8..12], &h[12..16], &h[16..20], &h[20..32])
191}
192
193fn getrandom_fill(b: &mut [u8]) -> std::io::Result<()> {
194    use std::io::Read;
195    std::fs::File::open("/dev/urandom")?.read_exact(b)
196}
197
198#[cfg(test)]
199mod tests {
200    use super::*;
201    #[test]
202    fn journal_roundtrip_and_args() {
203        let dir = std::env::temp_dir().join(format!("pty-resume-{}", std::process::id()));
204        std::fs::create_dir_all(&dir).unwrap();
205        let f = path(&dir);
206        let id = new_session_id();
207        assert_eq!(id.len(), 36);
208        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() });
209        assert_eq!(load(&f).len(), 1);
210        remove(&f, 1);
211        assert!(load(&f).is_empty());
212        assert_eq!(launch_args("claude", &id, true)[0], "--resume");
213        assert!(launch_args("codex", &id, true).is_empty());
214    }
215}
216
217#[cfg(test)]
218mod journal_tests {
219    use super::*;
220    #[test]
221    fn acp_sessions_of_any_provider_are_journalled_at_spawn_without_a_stop() {
222        assert!(journalled("codex", true) && journalled("claude", false) && !journalled("codex", false));
223        let dir = std::env::temp_dir().join(format!("resume-journal-{}", std::process::id()));
224        std::fs::create_dir_all(&dir).unwrap();
225        let f = path(&dir);
226        upsert(&f, Entry { instance_id: 2, provider: "codex".into(), workspace_id: "primary".into(), acp: true, ..Entry::default() });
227        // No stop, no signal: the file is already on disk, and no tmp is left behind.
228        assert_eq!(load(&f).len(), 1);
229        assert!(!f.with_extension("json.tmp").exists());
230        set_session_id(&f, 2, "sess-1");
231        assert_eq!(load(&f)[0].provider_session_id, "sess-1");
232        upsert(&f, Entry { instance_id: 2, provider: "codex".into(), workspace_id: "primary".into(), acp: true, provider_session_id: "sess-1".into(), ..Entry::default() });
233        assert_eq!(load(&f).len(), 1);
234        let _ = std::fs::remove_dir_all(&dir);
235    }
236}
hatchery_node/pty_resume.rs

hatchery_node/
pty_resume.rs
