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    /// Changes on each upsert so delayed cleanup cannot erase a retry or replacement.
13    #[serde(default)]
14    pub revision: u64,
15    pub instance_id: u64,
16    pub provider: String,
17    pub workspace_id: String,
18    pub provider_session_id: String,
19    /// m4a session key in the child's env (`M4A_SESSION_ID`): maps the runtime to its transcript.
20    #[serde(default)]
21    pub session_key: String,
22    #[serde(default)]
23    pub cwd: String,
24    /// `acp` for an ACP runtime (resumed through its managed record); PTY otherwise.
25    #[serde(default)]
26    pub acp: bool,
27    #[serde(default)]
28    pub options: Option<gate4agent_types::SessionOptionSelection>,
29    #[serde(default)]
30    pub cols: u16,
31    #[serde(default)]
32    pub rows: u16,
33    /// `human` | `agent` — stored on the node so origin survives HQ restarts.
34    #[serde(default)]
35    pub origin: Option<String>,
36    /// Parent chat signature, e.g. `claude#3`, when this session was spawned by another.
37    #[serde(default)]
38    pub parent_session_id: Option<String>,
39}
40
41static LOCK: Mutex<()> = Mutex::new(());
42
43pub fn enabled() -> bool {
44    std::env::var("HATCHERY_SESSION_RESUME").map_or(true, |v| v != "0")
45}
46
47/// Providers whose CLI can resume a known session id.
48pub fn supported(provider: &str) -> bool {
49    matches!(provider, "claude" | "codex")
50}
51
52/// Sessions journalled for restart resume: every ACP session (restored through its managed
53/// record with ACP `session/load`, else started fresh under the same number), and PTY sessions
54/// of providers with a resume flag. The journal is written at spawn and on every session-id or
55/// key change (tmp + rename), never on a stop signal, so a hard kill (Windows) keeps it.
56pub fn journalled(provider: &str, acp: bool) -> bool {
57    acp || supported(provider)
58}
59
60pub fn path(state_dir: &Path) -> PathBuf {
61    state_dir.join("session-resume.json")
62}
63
64pub fn load(file: &Path) -> Vec<Entry> {
65    std::fs::read(file).ok().and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default()
66}
67
68fn save(file: &Path, entries: &[Entry]) {
69    let tmp = file.with_extension("json.tmp");
70    match serde_json::to_vec_pretty(entries) {
71        Ok(b) => {
72            if std::fs::write(&tmp, b).and_then(|_| std::fs::rename(&tmp, file)).is_err() {
73                tracing::warn!(file = %file.display(), "session resume journal not writable");
74            }
75        }
76        Err(_) => {}
77    }
78}
79
80pub fn upsert(file: &Path, mut entry: Entry) {
81    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
82    let mut all = load(file);
83    entry.revision = all.iter().find(|e| e.instance_id == entry.instance_id)
84        .map_or(1, |e| e.revision.saturating_add(1));
85    all.retain(|e| e.instance_id != entry.instance_id);
86    all.push(entry);
87    all.sort_by_key(|e| e.instance_id);
88    save(file, &all);
89}
90
91/// Keep the journal aligned with the runtime's authoritative session identity.
92pub fn set_session_id(file: &Path, instance_id: u64, sid: &str) {
93    if sid.is_empty() { return; }
94    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
95    let mut all = load(file);
96    let mut changed = false;
97    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.provider_session_id != sid) {
98        e.provider_session_id = sid.to_owned();
99        changed = true;
100    }
101    if changed { save(file, &all); }
102}
103
104/// Record the m4a session key an entry's runtime got (a resume reuses the old one), so the
105/// session keeps its account and nick across node restarts.
106pub fn set_session_key(file: &Path, instance_id: u64, key: &str) {
107    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
108    let mut all = load(file);
109    let mut changed = false;
110    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.session_key != key) {
111        e.session_key = key.to_owned();
112        changed = true;
113    }
114    if changed { save(file, &all); }
115}
116
117pub fn remove(file: &Path, instance_id: u64) {
118    remove_matching(file, instance_id, None);
119}
120
121pub fn remove_revision(file: &Path, instance_id: u64, revision: u64) {
122    remove_matching(file, instance_id, Some(revision));
123}
124
125fn remove_matching(file: &Path, instance_id: u64, revision: Option<u64>) {
126    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
127    let mut all = load(file);
128    let n = all.len();
129    all.retain(|e| e.instance_id != instance_id || revision.is_some_and(|r| e.revision != r));
130    if all.len() != n {
131        save(file, &all);
132        if let Some(dir) = file.parent() {
133            crate::session_meta::mark_gone(dir, instance_id);
134        }
135    }
136}
137
138/// Fill provider session ids (and cwd) from the providers' live-session files: a claude child
139/// writes `~/.claude/sessions/<pid>.json`; its env carries the node's `M4A_SESSION_ID`.
140/// The m4a account id each node session key got (the child sees it as `M4A_SESSION_ID`).
141static ACCOUNT_REFS: std::sync::Mutex<Vec<(String, String)>> = std::sync::Mutex::new(Vec::new());
142
143/// Records that the child of `session_key` carries `account_ref` as its `M4A_SESSION_ID`.
144pub fn note_account_ref(session_key: &str, account_ref: &str) {
145    let mut m = ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner());
146    m.retain(|(k, _)| k != session_key);
147    m.push((session_key.to_owned(), account_ref.to_owned()));
148}
149
150fn account_ref_of(session_key: &str) -> Option<String> {
151    ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner()).iter().find(|(k, _)| k == session_key).map(|(_, r)| r.clone())
152}
153
154pub fn refresh(file: &Path, homes: &[PathBuf]) {
155    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
156    let mut all = load(file);
157    if all.iter().all(|e| e.session_key.is_empty()) { return; }
158    let mut found: Vec<(String, String, String)> = Vec::new();
159    for h in homes {
160        let Ok(rd) = std::fs::read_dir(h.join(".claude/sessions")) else { continue };
161        for d in rd.flatten() {
162            let Ok(b) = std::fs::read(d.path()) else { continue };
163            let Ok(v) = serde_json::from_slice::<serde_json::Value>(&b) else { continue };
164            let (Some(pid), Some(sid)) = (v["pid"].as_u64(), v["sessionId"].as_str()) else { continue };
165            let Ok(env) = std::fs::read(format!("/proc/{pid}/environ")) else { continue };
166            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());
167            if let Some(key) = key {
168                found.push((key, sid.to_owned(), v["cwd"].as_str().unwrap_or_default().to_owned()));
169            }
170        }
171    }
172    let mut changed = false;
173    for e in all.iter_mut().filter(|e| e.provider == "claude") {
174        let alt = account_ref_of(&e.session_key);
175        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()))) {
176            if e.provider_session_id != *sid || e.cwd != *cwd {
177                e.provider_session_id = sid.clone();
178                e.cwd = cwd.clone();
179                changed = true;
180            }
181        }
182    }
183    if changed { save(file, &all); }
184}
185
186/// Launch args: resume a known id, or pin a fresh one so it can be resumed later.
187pub fn launch_args(provider: &str, session_id: &str, resume: bool) -> Vec<String> {
188    match (provider, resume) {
189        ("claude", true) => vec!["--resume".into(), session_id.into()],
190        ("claude", false) => vec!["--session-id".into(), session_id.into()],
191        ("codex", true) => vec!["resume".into(), session_id.into()],
192        _ => Vec::new(),
193    }
194}
195
196pub fn new_session_id() -> String {
197    let mut b = [0u8; 16];
198    if getrandom_fill(&mut b).is_err() {
199        let t = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_nanos()).unwrap_or(0);
200        b.copy_from_slice(&(t ^ (std::process::id() as u128) << 64).to_le_bytes());
201    }
202    b[6] = (b[6] & 0x0f) | 0x40;
203    b[8] = (b[8] & 0x3f) | 0x80;
204    let h: String = b.iter().map(|x| format!("{x:02x}")).collect();
205    format!("{}-{}-{}-{}-{}", &h[0..8], &h[8..12], &h[12..16], &h[16..20], &h[20..32])
206}
207
208fn getrandom_fill(b: &mut [u8]) -> std::io::Result<()> {
209    use std::io::Read;
210    std::fs::File::open("/dev/urandom")?.read_exact(b)
211}
212
213#[cfg(test)]
214mod tests {
215    use super::*;
216    #[test]
217    fn journal_roundtrip_and_args() {
218        let dir = std::env::temp_dir().join(format!("pty-resume-{}", std::process::id()));
219        std::fs::create_dir_all(&dir).unwrap();
220        let f = path(&dir);
221        let id = new_session_id();
222        assert_eq!(id.len(), 36);
223        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() });
224        assert_eq!(load(&f).len(), 1);
225        remove(&f, 1);
226        assert!(load(&f).is_empty());
227        assert_eq!(launch_args("claude", &id, true)[0], "--resume");
228        assert_eq!(launch_args("codex", &id, true), ["resume", &id]);
229        assert!(launch_args("codex", &id, false).is_empty());
230    }
231}
232
233#[cfg(test)]
234mod journal_tests {
235    use super::*;
236    #[test]
237    fn codex_identity_changes_and_delayed_cleanup_preserve_the_retry() {
238        let dir = std::env::temp_dir().join(format!("resume-retry-{}", new_session_id()));
239        std::fs::create_dir_all(&dir).unwrap();
240        let file = path(&dir);
241        upsert(&file, Entry { instance_id: 7, provider: "codex".into(), ..Entry::default() });
242        set_session_id(&file, 7, "thread-first");
243        set_session_id(&file, 7, "thread-current");
244        set_session_id(&file, 7, "");
245        let previous = load(&file).pop().unwrap();
246        assert_eq!(previous.provider_session_id, "thread-current");
247        upsert(&file, previous.clone());
248        remove_revision(&file, 7, previous.revision);
249        let retry = load(&file).pop().unwrap();
250        assert_eq!(retry.provider_session_id, "thread-current");
251        assert!(retry.revision > previous.revision);
252        remove_revision(&file, 7, retry.revision);
253        assert!(load(&file).is_empty());
254        std::fs::remove_dir_all(dir).unwrap();
255    }
256
257    #[test]
258    fn acp_sessions_of_any_provider_are_journalled_at_spawn_without_a_stop() {
259        assert!(journalled("codex", true) && journalled("claude", false) && journalled("codex", false));
260        let dir = std::env::temp_dir().join(format!("resume-journal-{}", std::process::id()));
261        std::fs::create_dir_all(&dir).unwrap();
262        let f = path(&dir);
263        upsert(&f, Entry { instance_id: 2, provider: "codex".into(), workspace_id: "primary".into(), acp: true, ..Entry::default() });
264        // No stop, no signal: the file is already on disk, and no tmp is left behind.
265        assert_eq!(load(&f).len(), 1);
266        assert!(!f.with_extension("json.tmp").exists());
267        set_session_id(&f, 2, "sess-1");
268        assert_eq!(load(&f)[0].provider_session_id, "sess-1");
269        upsert(&f, Entry { instance_id: 2, provider: "codex".into(), workspace_id: "primary".into(), acp: true, provider_session_id: "sess-1".into(), ..Entry::default() });
270        assert_eq!(load(&f).len(), 1);
271        let _ = std::fs::remove_dir_all(&dir);
272    }
273}
hatchery_node/pty_resume.rs

hatchery_node/
pty_resume.rs
