openlatch-client 0.1.18

//! Session registry + identity assurance (D-09).
//!
//! There is no active-session store in the daemon `AppState` — this builds one.
//! The Mode 1 hook already fires on `SessionStart` and every tool call; the hook
//! side (`daemon::handlers::process_envelope`) stamps a **shared** registry that
//! the boundary listener reads at request time to resolve the attribution triple
//! (`agent_id` / `source` / `agent_session_id`) plus an honest confidence signal
//! (`assurance`).
//!
//! > **B-2 / C-2b.** Nothing session-specific travels on the model request. The
//! > boundary and the hook run in the **same binary**, so the current session is
//! > correlated *in-process* against this shared registry, keyed by the stable
//! > per-install id. `agent_id` + `source` are the platform join keys the hook
//! > materializer resolves against `agents` — carried here verbatim so the
//! > economics row joins the same way the hook stream does.
//!
//! The registry is written by the hook side and read by the boundary side. Both
//! hold the **same** `Arc<SessionRegistry>` (created once in
//! `daemon::serve_with_listener`), so an upsert on the hook path is visible to
//! the very next boundary request.

use std::cmp::Reverse;
use std::time::{Duration, Instant};

use dashmap::DashMap;

/// The stable, opaque, PII-free per-install identifier. In Phase 1 it reuses the
/// existing `agent_id` (`agt_<uuid>`) — the hook writes it into
/// `ANTHROPIC_CUSTOM_HEADERS` at `init`, so it arrives back on every model
/// request as `x-openlatch-install-id`, keying the registry from both sides.
pub type InstallId = String;

/// How long an unrefreshed session stays "active". Past this quiet window an
/// entry is dropped lazily on the next read/write — a stopped agent no longer
/// counts toward the concurrency tie-break.
pub const SESSION_QUIET_WINDOW: Duration = Duration::from_secs(300);

/// Cap on distinct installs the registry tracks, so a pathological caller that
/// forges a fresh `x-openlatch-install-id` per request cannot grow it unbounded.
/// Loopback-only + one install per host makes 1024 generous headroom.
const MAX_TRACKED_INSTALLS: usize = 1024;

/// Identity-assurance — the honest confidence signal (C-7, **three** values, no
/// `asserted`). Wire field `ai.openlatch.session.assurance`; storage column
/// `identity_assurance`. Same value under two names by design.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Assurance {
    /// Exactly one agent session active for this install at request time.
    Attested,
    /// Two or more concurrent — correlation is a most-recent-active best guess.
    Inferred,
    /// No hook session active (agent without the hook, or a non-agent caller).
    Unknown,
}

impl Assurance {
    /// Frozen wire string.
    pub fn as_str(&self) -> &'static str {
        match self {
            Assurance::Attested => "attested",
            Assurance::Inferred => "inferred",
            Assurance::Unknown => "unknown",
        }
    }
}

/// One live agent session, as last seen by the hook side.
#[derive(Clone, Debug)]
pub struct SessionActivity {
    /// The attribution resolution key (C-2b) — the same `agent_id` the hook emits.
    pub agent_id: String,
    /// The agent platform (`claude-code`) — first component of the `agents`
    /// unique key. Carried verbatim from the hook envelope so the join matches.
    pub source: String,
    /// The agent session id (the hook envelope `subject`).
    pub session_id: String,
    /// Last time the hook refreshed this session.
    pub last_seen: Instant,
}

/// Shared active-session store. `DashMap` (already a dependency) gives lock-free
/// concurrent access from the hook write path and the boundary read path.
#[derive(Default)]
pub struct SessionRegistry {
    active: DashMap<InstallId, Vec<SessionActivity>>,
}

impl SessionRegistry {
    /// Hook side: record (or refresh) a session for this install. Called on
    /// `SessionStart` and every tool-call hook. Stale siblings are pruned in
    /// passing so the concurrency count stays honest.
    pub fn upsert(&self, install_id: &str, agent_id: &str, source: &str, session_id: &str) {
        // Bound growth: if we are at the install cap and this is a brand-new
        // install key, drop it rather than grow unbounded (loopback, one install
        // per host — this only trips under a forging caller).
        if !self.active.contains_key(install_id) && self.active.len() >= MAX_TRACKED_INSTALLS {
            self.evict_empty();
            if self.active.len() >= MAX_TRACKED_INSTALLS {
                return;
            }
        }

        let now = Instant::now();
        let mut entry = self.active.entry(install_id.to_string()).or_default();
        entry.retain(|a| now.duration_since(a.last_seen) < SESSION_QUIET_WINDOW);
        if let Some(existing) = entry.iter_mut().find(|a| a.session_id == session_id) {
            existing.last_seen = now;
            existing.agent_id = agent_id.to_string();
            existing.source = source.to_string();
        } else {
            entry.push(SessionActivity {
                agent_id: agent_id.to_string(),
                source: source.to_string(),
                session_id: session_id.to_string(),
                last_seen: now,
            });
        }
    }

    /// Drop install keys whose session vectors have gone fully stale/empty.
    fn evict_empty(&self) {
        let now = Instant::now();
        self.active.retain(|_, v| {
            v.retain(|a| now.duration_since(a.last_seen) < SESSION_QUIET_WINDOW);
            !v.is_empty()
        });
    }

    /// The fresh (non-stale) sessions for an install, dropping expired entries.
    fn fresh(&self, install_id: &str) -> Vec<SessionActivity> {
        let now = Instant::now();
        match self.active.get(install_id) {
            Some(v) => v
                .iter()
                .filter(|a| now.duration_since(a.last_seen) < SESSION_QUIET_WINDOW)
                .cloned()
                .collect(),
            None => Vec::new(),
        }
    }
}

/// The resolved attribution triple + assurance for one request.
#[derive(Clone, Debug)]
pub struct Resolved {
    pub agent_id: Option<String>,
    pub source: Option<String>,
    pub session_id: Option<String>,
    pub assurance: Assurance,
}

impl Resolved {
    /// The unattributed resolution — no hook session for this install: no
    /// identifiers, `unknown` assurance. Still forwarded, just Unattributed.
    pub fn unknown() -> Self {
        Resolved {
            agent_id: None,
            source: None,
            session_id: None,
            assurance: Assurance::Unknown,
        }
    }
}

/// Boundary side (D-09 / G17): resolve the attribution triple + assurance for an
/// install at request time.
///
/// - **1 active** → `attested`, that session's identifiers.
/// - **0 active** → `unknown`, no identifiers (still forwarded, Unattributed).
/// - **≥2 active** → `inferred`, the **most-recently-active** session (flagged).
pub fn resolve_session(reg: &SessionRegistry, install: &str) -> Resolved {
    let mut active = reg.fresh(install);
    match active.len() {
        1 => {
            let a = active.remove(0);
            Resolved {
                agent_id: Some(a.agent_id),
                source: Some(a.source),
                session_id: Some(a.session_id),
                assurance: Assurance::Attested,
            }
        }
        0 => Resolved::unknown(),
        _ => {
            // Most-recent-active wins, flagged inferred (the honesty mechanism).
            active.sort_by_key(|a| Reverse(a.last_seen));
            let a = active.remove(0);
            Resolved {
                agent_id: Some(a.agent_id),
                source: Some(a.source),
                session_id: Some(a.session_id),
                assurance: Assurance::Inferred,
            }
        }
    }
}

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

    #[test]
    fn zero_active_resolves_unknown() {
        // C-8b: no hook session for this install → unknown, no identifiers.
        let reg = SessionRegistry::default();
        let r = resolve_session(&reg, "agt_absent");
        assert_eq!(r.assurance, Assurance::Unknown);
        assert!(r.agent_id.is_none() && r.session_id.is_none() && r.source.is_none());
    }

    #[test]
    fn one_active_resolves_attested_with_identifiers() {
        let reg = SessionRegistry::default();
        reg.upsert("agt_1", "agt_1", "claude-code", "sess_a");
        let r = resolve_session(&reg, "agt_1");
        assert_eq!(r.assurance, Assurance::Attested);
        assert_eq!(r.session_id.as_deref(), Some("sess_a"));
        assert_eq!(r.agent_id.as_deref(), Some("agt_1"));
        assert_eq!(r.source.as_deref(), Some("claude-code"));
    }

    #[test]
    fn two_concurrent_resolve_inferred_most_recent() {
        let reg = SessionRegistry::default();
        reg.upsert("agt_1", "agt_1", "claude-code", "sess_old");
        std::thread::sleep(Duration::from_millis(5));
        reg.upsert("agt_1", "agt_1", "claude-code", "sess_new");
        let r = resolve_session(&reg, "agt_1");
        assert_eq!(r.assurance, Assurance::Inferred);
        // Most-recently-active session wins the tie-break.
        assert_eq!(r.session_id.as_deref(), Some("sess_new"));
    }

    #[test]
    fn refresh_updates_last_seen_not_count() {
        let reg = SessionRegistry::default();
        reg.upsert("agt_1", "agt_1", "claude-code", "sess_a");
        reg.upsert("agt_1", "agt_1", "claude-code", "sess_a");
        // Same session refreshed → still one active → attested, not inferred.
        let r = resolve_session(&reg, "agt_1");
        assert_eq!(r.assurance, Assurance::Attested);
    }

    #[test]
    fn assurance_wire_strings_are_the_frozen_three() {
        assert_eq!(Assurance::Attested.as_str(), "attested");
        assert_eq!(Assurance::Inferred.as_str(), "inferred");
        assert_eq!(Assurance::Unknown.as_str(), "unknown");
    }
}