openlatch-client 0.6.9

//! Which lead session a Claude Code agent-team teammate belongs to.
//!
//! A teammate (an `Explore` agent in a tmux pane, say) is its own `claude` process with its own
//! `session_id`. Claude Code launches it with `--parent-session-id <lead uuid>` and
//! `--agent-id <name>@session-<8 hex>`, and the hook runs as a descendant of that process. On
//! Claude Code's `SessionStart` the hook forwards its own parent pid as
//! [`HOOK_PPID_HEADER`](crate::hook_output::HOOK_PPID_HEADER); the daemon walks up from it, reads
//! each ancestor's argv — never its environment — and stamps the two ids into the event as
//! [`PARENT_SESSION_KEY`] and [`TEAMMATE_KEY`], which the platform links sessions by.
//!
//! Those flags are undocumented. If Claude Code renames them, the keys are simply absent: every
//! failure here is "no link", never an error, and nothing on the verdict depends on it.

use crate::hooks::cline_app::process::{self, ProcArgv};

/// The lead session's id, as Claude Code passed it to the teammate.
pub(crate) const PARENT_SESSION_KEY: &str = "openlatch.parent_session_id";
/// The teammate's own agent id (`<name>@session-<8 hex>`).
pub(crate) const TEAMMATE_KEY: &str = "openlatch.teammate_id";

/// How many processes are inspected, starting at the hook's parent. The hook usually runs under
/// `sh -c`, so the teammate `claude` is the first or second; three leaves room for one wrapper.
const MAX_ANCESTORS: usize = 3;

/// The two ids a teammate's argv carries.
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) struct TeammateLink {
    pub(crate) parent_session_id: String,
    pub(crate) teammate_id: String,
}

/// Both flags from one argv, or None. Only `--parent-session-id` (a UUID) and `--agent-id`
/// (`^[A-Za-z0-9._-]{1,64}@session-[0-9a-f]{8}$`) are read; every other argument is dropped.
/// Both are required: a parent id without its teammate id (or the reverse) is not a teammate
/// launch we recognise, and half a link is worse than none.
pub(crate) fn parse_teammate_flags(argv: &[String]) -> Option<TeammateLink> {
    let mut parent = None;
    let mut agent = None;
    let mut it = argv.iter();
    while let Some(arg) = it.next() {
        let (slot, value) = if let Some(v) = arg.strip_prefix("--parent-session-id=") {
            (&mut parent, Some(v))
        } else if arg == "--parent-session-id" {
            (&mut parent, it.next().map(String::as_str))
        } else if let Some(v) = arg.strip_prefix("--agent-id=") {
            (&mut agent, Some(v))
        } else if arg == "--agent-id" {
            (&mut agent, it.next().map(String::as_str))
        } else {
            continue;
        };
        if let Some(v) = value {
            *slot = Some(v);
        }
    }
    let parent = parent.filter(|v| is_uuid(v))?;
    let agent = agent.filter(|v| is_agent_id(v))?;
    Some(TeammateLink {
        parent_session_id: parent.to_string(),
        teammate_id: agent.to_string(),
    })
}

/// `8-4-4-4-12` hex digits, either case.
fn is_uuid(v: &str) -> bool {
    v.len() == 36
        && v.bytes().enumerate().all(|(i, b)| match i {
            8 | 13 | 18 | 23 => b == b'-',
            _ => b.is_ascii_hexdigit(),
        })
}

/// `^[A-Za-z0-9._-]{1,64}@session-[0-9a-f]{8}$`.
fn is_agent_id(v: &str) -> bool {
    let Some((name, session)) = v.split_once("@session-") else {
        return false;
    };
    (1..=64).contains(&name.len())
        && name
            .bytes()
            .all(|b| b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-'))
        && session.len() == 8
        && session
            .bytes()
            .all(|b| b.is_ascii_digit() || (b'a'..=b'f').contains(&b))
}

/// Walk at most [`MAX_ANCESTORS`] processes up from `start`, returning the first argv that names
/// a teammate. `read` is the process table, injected for tests.
fn find_in_ancestors(start: u32, read: impl Fn(u32) -> Option<ProcArgv>) -> Option<TeammateLink> {
    let mut pid = start;
    for _ in 0..MAX_ANCESTORS {
        let proc = read(pid)?;
        if let Some(link) = parse_teammate_flags(&proc.args) {
            return Some(link);
        }
        // pid 1 (or a self-parented pid) is the top: nothing further up launched the hook.
        match proc.ppid {
            Some(ppid) if ppid > 1 && ppid != pid => pid = ppid,
            _ => return None,
        }
    }
    None
}

/// The teammate link for a hook whose parent pid is `hook_ppid`, read from the live process table.
pub(crate) fn find(hook_ppid: u32) -> Option<TeammateLink> {
    find_in_ancestors(hook_ppid, process::argv_of)
}

/// Write `link` into `data` under the two enrichment keys. A non-object `data` is left alone.
pub(crate) fn stamp(data: &mut Option<serde_json::Value>, link: TeammateLink) {
    if let Some(obj) = data.as_mut().and_then(serde_json::Value::as_object_mut) {
        obj.insert(
            PARENT_SESSION_KEY.to_string(),
            serde_json::Value::String(link.parent_session_id),
        );
        obj.insert(
            TEAMMATE_KEY.to_string(),
            serde_json::Value::String(link.teammate_id),
        );
    }
}

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

    const LEAD: &str = "0b101558-1c2d-4e5f-8a9b-0123456789ab";
    const AGENT: &str = "agent-survey@session-0b101558";

    fn argv(a: &[&str]) -> Vec<String> {
        a.iter().map(|s| (*s).to_string()).collect()
    }

    fn link() -> TeammateLink {
        TeammateLink {
            parent_session_id: LEAD.into(),
            teammate_id: AGENT.into(),
        }
    }

    /// The argv parser is the contract with an undocumented Claude Code launch: pin it.
    #[test]
    fn parse_teammate_flags_table() {
        let cases: &[(&str, &[&str], Option<TeammateLink>)] = &[
            (
                "claude 2.1.291 teammate launch, extra flags ignored",
                &[
                    "claude",
                    "--model",
                    "opus",
                    "--parent-session-id",
                    LEAD,
                    "--agent-id",
                    AGENT,
                    "--dangerously-skip-permissions",
                ],
                Some(link()),
            ),
            (
                "equals form",
                &[
                    "claude",
                    &format!("--agent-id={AGENT}"),
                    &format!("--parent-session-id={LEAD}"),
                ],
                Some(link()),
            ),
            ("no flags: a lead session", &["claude", "--resume"], None),
            (
                "only the parent id: half a link is none",
                &["claude", "--parent-session-id", LEAD],
                None,
            ),
            ("only the agent id", &["claude", "--agent-id", AGENT], None),
            (
                "parent id is not a uuid",
                &[
                    "claude",
                    "--parent-session-id",
                    "not-a-uuid",
                    "--agent-id",
                    AGENT,
                ],
                None,
            ),
            (
                "agent id lacks the session suffix",
                &[
                    "claude",
                    "--parent-session-id",
                    LEAD,
                    "--agent-id",
                    "agent-survey",
                ],
                None,
            ),
            (
                "agent id session suffix is uppercase hex",
                &[
                    "claude",
                    "--parent-session-id",
                    LEAD,
                    "--agent-id",
                    "agent-survey@session-0B101558",
                ],
                None,
            ),
            (
                "agent name carries a forbidden character",
                &[
                    "claude",
                    "--parent-session-id",
                    LEAD,
                    "--agent-id",
                    "agent/survey@session-0b101558",
                ],
                None,
            ),
            (
                "flag with no value at the end",
                &["claude", "--agent-id", AGENT, "--parent-session-id"],
                None,
            ),
        ];
        for (name, args, want) in cases {
            assert_eq!(&parse_teammate_flags(&argv(args)), want, "{name}");
        }
    }

    #[test]
    fn agent_name_is_bounded_at_64() {
        let ok = format!("{}@session-0b101558", "a".repeat(64));
        let long = format!("{}@session-0b101558", "a".repeat(65));
        assert!(is_agent_id(&ok));
        assert!(!is_agent_id(&long));
        assert!(!is_agent_id("@session-0b101558"));
    }

    fn table(procs: &[(u32, &[&str], u32)]) -> HashMap<u32, ProcArgv> {
        procs
            .iter()
            .map(|(pid, args, ppid)| {
                (
                    *pid,
                    ProcArgv {
                        args: argv(args),
                        ppid: Some(*ppid),
                    },
                )
            })
            .collect()
    }

    #[test]
    fn walks_past_the_shell_to_the_teammate() {
        let t = table(&[
            (300, &["sh", "-c", "openlatch-hook"], 200),
            (
                200,
                &["claude", "--parent-session-id", LEAD, "--agent-id", AGENT],
                100,
            ),
            (100, &["tmux"], 1),
        ]);
        assert_eq!(find_in_ancestors(300, |p| t.get(&p).cloned()), Some(link()));
    }

    #[test]
    fn stops_after_three_ancestors_and_at_the_top() {
        let deep = table(&[
            (5, &["sh"], 4),
            (4, &["sh"], 3),
            (3, &["sh"], 2),
            (
                2,
                &["claude", "--parent-session-id", LEAD, "--agent-id", AGENT],
                1,
            ),
        ]);
        assert_eq!(find_in_ancestors(5, |p| deep.get(&p).cloned()), None);
        let lead = table(&[(30, &["sh"], 20), (20, &["claude"], 1)]);
        assert_eq!(find_in_ancestors(30, |p| lead.get(&p).cloned()), None);
        assert_eq!(find_in_ancestors(99, |_| None), None);
    }

    #[test]
    fn stamp_writes_both_keys_into_an_object_only() {
        let mut data = Some(serde_json::json!({"hook_event_name": "SessionStart"}));
        stamp(&mut data, link());
        let d = data.unwrap();
        assert_eq!(d[PARENT_SESSION_KEY], LEAD);
        assert_eq!(d[TEAMMATE_KEY], AGENT);

        let mut not_obj = Some(serde_json::json!("x"));
        stamp(&mut not_obj, link());
        assert_eq!(not_obj, Some(serde_json::json!("x")));
    }
}