hatchery-node 0.2.62

//! Node session registry: PTY sessions the user opened survive a node restart.
//! Each live PTY session is journalled (provider, provider session id, workspace, model/effort,
//! instance number); on node start every journalled session is spawned again with the
//! provider's own resume flag, under the same instance number, so HQ shows the same chat.
//! Off with `HATCHERY_SESSION_RESUME=0`.
use serde::{Deserialize, Serialize};
use std::path::{Path, PathBuf};
use std::sync::Mutex;

#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)]
pub struct Entry {
    pub instance_id: u64,
    pub provider: String,
    pub workspace_id: String,
    pub provider_session_id: String,
    /// m4a session key in the child's env (`M4A_SESSION_ID`): maps the runtime to its transcript.
    #[serde(default)]
    pub session_key: String,
    #[serde(default)]
    pub cwd: String,
    /// `acp` for an ACP runtime (resumed through its managed record); PTY otherwise.
    #[serde(default)]
    pub acp: bool,
    #[serde(default)]
    pub options: Option<gate4agent_types::SessionOptionSelection>,
    #[serde(default)]
    pub cols: u16,
    #[serde(default)]
    pub rows: u16,
}

static LOCK: Mutex<()> = Mutex::new(());

pub fn enabled() -> bool {
    std::env::var("HATCHERY_SESSION_RESUME").map_or(true, |v| v != "0")
}

/// Providers whose CLI accepts a caller-chosen session id plus a resume flag.
pub fn supported(provider: &str) -> bool {
    provider == "claude"
}

pub fn path(state_dir: &Path) -> PathBuf {
    state_dir.join("session-resume.json")
}

pub fn load(file: &Path) -> Vec<Entry> {
    std::fs::read(file).ok().and_then(|b| serde_json::from_slice(&b).ok()).unwrap_or_default()
}

fn save(file: &Path, entries: &[Entry]) {
    let tmp = file.with_extension("json.tmp");
    match serde_json::to_vec_pretty(entries) {
        Ok(b) => {
            if std::fs::write(&tmp, b).and_then(|_| std::fs::rename(&tmp, file)).is_err() {
                tracing::warn!(file = %file.display(), "session resume journal not writable");
            }
        }
        Err(_) => {}
    }
}

pub fn upsert(file: &Path, entry: Entry) {
    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let mut all = load(file);
    all.retain(|e| e.instance_id != entry.instance_id);
    all.push(entry);
    all.sort_by_key(|e| e.instance_id);
    save(file, &all);
}

/// Set the provider session id of an entry that has none yet.
pub fn set_session_id(file: &Path, instance_id: u64, sid: &str) {
    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let mut all = load(file);
    let mut changed = false;
    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.provider_session_id.is_empty()) {
        e.provider_session_id = sid.to_owned();
        changed = true;
    }
    if changed { save(file, &all); }
}

/// Record the m4a session key an entry's runtime got (a resume reuses the old one), so the
/// session keeps its account and nick across node restarts.
pub fn set_session_key(file: &Path, instance_id: u64, key: &str) {
    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let mut all = load(file);
    let mut changed = false;
    for e in all.iter_mut().filter(|e| e.instance_id == instance_id && e.session_key != key) {
        e.session_key = key.to_owned();
        changed = true;
    }
    if changed { save(file, &all); }
}

pub fn remove(file: &Path, instance_id: u64) {
    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let mut all = load(file);
    let n = all.len();
    all.retain(|e| e.instance_id != instance_id);
    if all.len() != n {
        save(file, &all);
    }
}

/// Fill provider session ids (and cwd) from the providers' live-session files: a claude child
/// writes `~/.claude/sessions/<pid>.json`; its env carries the node's `M4A_SESSION_ID`.
/// The m4a account id each node session key got (the child sees it as `M4A_SESSION_ID`).
static ACCOUNT_REFS: std::sync::Mutex<Vec<(String, String)>> = std::sync::Mutex::new(Vec::new());

/// Records that the child of `session_key` carries `account_ref` as its `M4A_SESSION_ID`.
pub fn note_account_ref(session_key: &str, account_ref: &str) {
    let mut m = ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner());
    m.retain(|(k, _)| k != session_key);
    m.push((session_key.to_owned(), account_ref.to_owned()));
}

fn account_ref_of(session_key: &str) -> Option<String> {
    ACCOUNT_REFS.lock().unwrap_or_else(|p| p.into_inner()).iter().find(|(k, _)| k == session_key).map(|(_, r)| r.clone())
}

pub fn refresh(file: &Path, homes: &[PathBuf]) {
    let _g = LOCK.lock().unwrap_or_else(|p| p.into_inner());
    let mut all = load(file);
    if all.iter().all(|e| e.session_key.is_empty()) { return; }
    let mut found: Vec<(String, String, String)> = Vec::new();
    for h in homes {
        let Ok(rd) = std::fs::read_dir(h.join(".claude/sessions")) else { continue };
        for d in rd.flatten() {
            let Ok(b) = std::fs::read(d.path()) else { continue };
            let Ok(v) = serde_json::from_slice::<serde_json::Value>(&b) else { continue };
            let (Some(pid), Some(sid)) = (v["pid"].as_u64(), v["sessionId"].as_str()) else { continue };
            let Ok(env) = std::fs::read(format!("/proc/{pid}/environ")) else { continue };
            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());
            if let Some(key) = key {
                found.push((key, sid.to_owned(), v["cwd"].as_str().unwrap_or_default().to_owned()));
            }
        }
    }
    let mut changed = false;
    for e in all.iter_mut() {
        let alt = account_ref_of(&e.session_key);
        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()))) {
            if e.provider_session_id != *sid || e.cwd != *cwd {
                e.provider_session_id = sid.clone();
                e.cwd = cwd.clone();
                changed = true;
            }
        }
    }
    if changed { save(file, &all); }
}

/// Launch args: resume a known id, or pin a fresh one so it can be resumed later.
pub fn launch_args(provider: &str, session_id: &str, resume: bool) -> Vec<String> {
    match (provider, resume) {
        ("claude", true) => vec!["--resume".into(), session_id.into()],
        ("claude", false) => vec!["--session-id".into(), session_id.into()],
        _ => Vec::new(),
    }
}

pub fn new_session_id() -> String {
    let mut b = [0u8; 16];
    if getrandom_fill(&mut b).is_err() {
        let t = std::time::SystemTime::now().duration_since(std::time::UNIX_EPOCH).map(|d| d.as_nanos()).unwrap_or(0);
        b.copy_from_slice(&(t ^ (std::process::id() as u128) << 64).to_le_bytes());
    }
    b[6] = (b[6] & 0x0f) | 0x40;
    b[8] = (b[8] & 0x3f) | 0x80;
    let h: String = b.iter().map(|x| format!("{x:02x}")).collect();
    format!("{}-{}-{}-{}-{}", &h[0..8], &h[8..12], &h[12..16], &h[16..20], &h[20..32])
}

fn getrandom_fill(b: &mut [u8]) -> std::io::Result<()> {
    use std::io::Read;
    std::fs::File::open("/dev/urandom")?.read_exact(b)
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn journal_roundtrip_and_args() {
        let dir = std::env::temp_dir().join(format!("pty-resume-{}", std::process::id()));
        std::fs::create_dir_all(&dir).unwrap();
        let f = path(&dir);
        let id = new_session_id();
        assert_eq!(id.len(), 36);
        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 });
        assert_eq!(load(&f).len(), 1);
        remove(&f, 1);
        assert!(load(&f).is_empty());
        assert_eq!(launch_args("claude", &id, true)[0], "--resume");
        assert!(launch_args("codex", &id, true).is_empty());
    }
}