hatchery-node 0.2.76

//! Git state of live sessions for the task manager: branch, head and dirty count of the
//! directory each session's process runs in. Read with `git` itself (no parsing of `.git`);
//! a session outside a repository reports nothing. Unix reads the cwd from `/proc/<pid>/cwd`;
//! elsewhere the watcher reports nothing until a portable cwd source exists.

use serde::Serialize;
use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use std::sync::{Mutex, OnceLock};

#[derive(Clone, Debug, Serialize, PartialEq, Eq)]
pub struct SessionGit {
    pub workspace_id: String,
    pub instance_id: u64,
    pub provider: String,
    pub branch: String,
    pub head: String,
    /// Changed + untracked paths (`git status --porcelain` lines), capped at 10000.
    pub dirty: u32,
    /// Commits on HEAD that the local default branch (`main`, else `master`) does not have.
    pub ahead: u32,
    pub root: String,
}

#[derive(Clone, Debug)]
struct Live {
    pid: u32,
    provider: String,
    /// Where the session was last seen working (start cwd, then the last snapshot's root):
    /// after exit `/proc/<pid>/cwd` is gone, and the final git state is read from here.
    dir: Option<PathBuf>,
}

/// A session's lifecycle as the task manager sees it (`task.session_state`).
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Lifecycle {
    pub workspace: String,
    pub instance: u64,
    /// `live` | `gone`.
    pub state: &'static str,
    /// For `gone`: where to read the session's last git state.
    pub dir: Option<PathBuf>,
    /// The session's keystore public key (hex), once it has one: core binds the session
    /// record to it, so the session's own signed task acts name it.
    pub public_key: Option<String>,
}

static LIFECYCLE: OnceLock<tokio::sync::mpsc::UnboundedSender<Lifecycle>> = OnceLock::new();

/// Install the lifecycle sink (once, at startup). Without one, nothing is sent.
pub fn set_lifecycle_sink(tx: tokio::sync::mpsc::UnboundedSender<Lifecycle>) {
    let _ = LIFECYCLE.set(tx);
}

fn emit(l: Lifecycle) {
    if let Some(tx) = LIFECYCLE.get() {
        let _ = tx.send(l);
    }
}

static LIVE: OnceLock<Mutex<BTreeMap<(String, u64), Live>>> = OnceLock::new();

fn live() -> std::sync::MutexGuard<'static, BTreeMap<(String, u64), Live>> {
    LIVE.get_or_init(Default::default).lock().unwrap_or_else(|p| p.into_inner())
}

/// The session's process started.
pub fn running(workspace: &str, instance: u64, provider: &str, pid: u32) {
    live().insert((workspace.to_owned(), instance), Live { pid, provider: provider.to_owned(), dir: cwd_of(pid) });
    emit(Lifecycle { workspace: workspace.to_owned(), instance, state: "live", dir: None, public_key: None });
}

/// The provider of a live session (from the frame publisher; `running` may not know it).
pub fn provider_hint(workspace: &str, instance: u64, provider: &str) {
    if let Some(l) = live().get_mut(&(workspace.to_owned(), instance)) {
        if l.provider != provider {
            l.provider = provider.to_owned();
        }
    }
}

/// The session's process exited.
pub fn exited(workspace: &str, instance: u64) {
    if let Some(l) = live().remove(&(workspace.to_owned(), instance)) {
        emit(Lifecycle { workspace: workspace.to_owned(), instance, state: "gone", dir: l.dir, public_key: None });
    }
}

/// The keystore issued the identity of proxy session `<incarnation>-<instance>`: report the
/// live session again with its public key.
pub fn session_key(proxy_session: &str, public_key_hex: &str) {
    let Some(instance) = proxy_session.rsplit('-').next().and_then(|i| i.parse::<u64>().ok()) else { return };
    if public_key_hex.is_empty() {
        return;
    }
    let workspace = live().keys().find(|(_, i)| *i == instance).map(|(w, _)| w.clone());
    if let Some(workspace) = workspace {
        emit(Lifecycle { workspace, instance, state: "live", dir: None, public_key: Some(public_key_hex.to_owned()) });
    }
}

/// The git state of the repository a session worked in (`dir` itself, or the `task-*`
/// repository one level below it), in the `task.session_git` shape.
pub fn read_session_repo(dir: &Path) -> Option<serde_json::Value> {
    let (branch, head, dirty, ahead, root) = read_repo_full(dir).or_else(|| task_repo_below(dir))?;
    Some(serde_json::json!({ "branch": branch, "head": head, "dirty": dirty, "ahead": ahead, "root": root }))
}

fn cwd_of(pid: u32) -> Option<PathBuf> {
    if cfg!(unix) {
        std::fs::read_link(format!("/proc/{pid}/cwd")).ok()
    } else {
        None
    }
}

fn git(dir: &Path, args: &[&str]) -> Option<String> {
    let mut cmd = std::process::Command::new("git");
    cmd.arg("-C").arg(dir).args(args).stdin(std::process::Stdio::null()).stderr(std::process::Stdio::null());
    crate::windows_process::hide_console(&mut cmd);
    let out = cmd.output().ok()?;
    out.status.success().then(|| String::from_utf8_lossy(&out.stdout).into_owned())
}

/// Branch / head / dirty of the repository containing `dir`.
pub fn read_repo(dir: &Path) -> Option<(String, String, u32, String)> {
    read_repo_full(dir).map(|(b, h, d, _, r)| (b, h, d, r))
}

/// Like [`read_repo`] plus `ahead` of the default branch.
pub fn read_repo_full(dir: &Path) -> Option<(String, String, u32, u32, String)> {
    let (branch, head, dirty, root) = read_repo_basic(dir)?;
    let base = ["main", "master"].into_iter().find(|b| git(dir, &["rev-parse", "--verify", "--quiet", &format!("refs/heads/{b}")]).is_some());
    let ahead = base
        .filter(|b| *b != branch)
        .and_then(|b| git(dir, &["rev-list", "--count", "HEAD", &format!("^{b}")]))
        .and_then(|s| s.trim().parse().ok())
        .unwrap_or(0);
    Some((branch, head, dirty, ahead, root))
}

fn read_repo_basic(dir: &Path) -> Option<(String, String, u32, String)> {
    let root = git(dir, &["rev-parse", "--show-toplevel"])?.trim().to_owned();
    let branch = git(dir, &["rev-parse", "--abbrev-ref", "HEAD"]).map(|s| s.trim().to_owned()).unwrap_or_default();
    let head = git(dir, &["rev-parse", "--short=12", "HEAD"]).map(|s| s.trim().to_owned()).unwrap_or_default();
    let dirty = git(dir, &["status", "--porcelain", "--untracked-files=normal"]).map(|s| s.lines().take(10_000).count() as u32).unwrap_or(0);
    Some((branch, head, dirty, root))
}

/// A session whose cwd is not a repository (a shared workspace root): the repository one level
/// below it that sits on a `task-*` branch (the branch a dispatched task works on).
fn task_repo_below(dir: &Path) -> Option<(String, String, u32, u32, String)> {
    let mut entries: Vec<_> = std::fs::read_dir(dir).ok()?.flatten().take(256).map(|e| e.path()).filter(|p| p.join(".git").exists()).collect();
    entries.sort();
    entries.into_iter().find_map(|p| read_repo_full(&p).filter(|r| r.0.starts_with("task-")))
}

/// One snapshot of every live session that sits in a repository.
pub fn snapshot() -> Vec<SessionGit> {
    let sessions: Vec<_> = live().iter().map(|(k, v)| (k.clone(), v.clone())).collect();
    let out: Vec<SessionGit> = sessions
        .into_iter()
        .filter_map(|((ws, inst), l)| {
            let dir = cwd_of(l.pid)?;
            let (branch, head, dirty, ahead, root) = read_repo_full(&dir).or_else(|| task_repo_below(&dir))?;
            Some(SessionGit { workspace_id: ws, instance_id: inst, provider: l.provider, branch, head, dirty, ahead, root })
        })
        .collect();
    let mut live = live();
    for s in &out {
        if let Some(l) = live.get_mut(&(s.workspace_id.clone(), s.instance_id)) {
            l.dir = Some(PathBuf::from(&s.root));
        }
    }
    out
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn reads_branch_and_dirty() {
        let d = std::env::temp_dir().join(format!("gw-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&d);
        std::fs::create_dir_all(&d).unwrap();
        let run = |a: &[&str]| assert!(std::process::Command::new("git").arg("-C").arg(&d).args(a).output().unwrap().status.success());
        run(&["init", "-q", "-b", "feat-x"]);
        run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "i"]);
        std::fs::write(d.join("a.txt"), "x").unwrap();
        let (branch, head, dirty, _) = read_repo(&d).unwrap();
        run(&["branch", "main"]);
        run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "j"]);
        assert_eq!(read_repo_full(&d).unwrap().3, 1);
        assert_eq!(branch, "feat-x");
        assert_eq!(head.len(), 12);
        assert_eq!(dirty, 1);
        let _ = std::fs::remove_dir_all(d);
    }

    #[test]
    fn finds_task_repo_below_workspace() {
        let ws = std::env::temp_dir().join(format!("gw-ws-{}", std::process::id()));
        let _ = std::fs::remove_dir_all(&ws);
        let r = ws.join("repo");
        std::fs::create_dir_all(&r).unwrap();
        let run = |a: &[&str]| assert!(std::process::Command::new("git").arg("-C").arg(&r).args(a).output().unwrap().status.success());
        run(&["init", "-q", "-b", "main"]);
        run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "i"]);
        assert!(task_repo_below(&ws).is_none());
        run(&["checkout", "-q", "-b", "task-9"]);
        run(&["-c", "core.hooksPath=/dev/null", "-c", "user.name=t", "-c", "user.email=t@t", "commit", "-q", "--allow-empty", "-m", "w"]);
        let f = task_repo_below(&ws).unwrap();
        assert_eq!((f.0.as_str(), f.2, f.3), ("task-9", 0, 1));
        let _ = std::fs::remove_dir_all(ws);
    }

    #[test]
    fn lifecycle_is_emitted_with_the_last_dir() {
        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
        set_lifecycle_sink(tx);
        running("lc", 91, "codex", std::process::id());
        session_key("inc-1-91", "ab");
        exited("lc", 91);
        exited("lc", 91);
        let mut got = Vec::new();
        while let Ok(l) = rx.try_recv() {
            if l.workspace == "lc" {
                got.push(l);
            }
        }
        assert_eq!(got.iter().map(|l| l.state).collect::<Vec<_>>(), ["live", "live", "gone"], "a second exit is not a second gone");
        assert_eq!(got[1].public_key.as_deref(), Some("ab"));
        if cfg!(target_os = "linux") {
            assert_eq!(got[2].dir, std::env::current_dir().ok());
        }
        let here = std::env::current_dir().unwrap();
        assert_eq!(read_session_repo(&here).is_some(), read_repo(&here).is_some());
    }

    #[test]
    fn live_session_snapshot() {
        running("w", 77, "claude", std::process::id());
        let here = std::env::current_dir().unwrap();
        let in_repo = read_repo(&here).is_some();
        assert_eq!(snapshot().iter().any(|s| s.instance_id == 77), in_repo && cfg!(target_os = "linux"));
        exited("w", 77);
        assert!(!snapshot().iter().any(|s| s.instance_id == 77));
    }
}