openlatch-client 0.1.5

The open-source security layer for AI agents — client forwarder
//! Identity stitching for telemetry — Phase C.
//!
//! Once authentication completes and the platform returns a `user_db_id`, we
//! emit a one-time `$create_alias` event so PostHog merges all prior
//! `agent_id`-keyed activity into the user's persistent person. From then
//! on, the same `user_db_id` is used as `distinct_id` across CLI sessions.
//!
//! Persistence lives in `~/.openlatch/identity.json`, separate from
//! `config.toml` so the wire-format config stays stable. The file is the
//! source of truth for "we already aliased this agent_id → user_db_id
//! pair" — without it we would re-alias on every login.
//!
//! This module ships in Phase A so the helper is available, but the
//! `$create_alias` payload only fires from Task 4 wiring at the auth call
//! site; everything here is a no-op when consent or the baked key is missing.

use std::path::{Path, PathBuf};

use serde::{Deserialize, Serialize};
use serde_json::{json, Map};

use super::client::QueuedEvent;
use super::events::Event;
use super::TelemetryHandle;

/// On-disk identity record. `aliased_pairs` lets us alias more than once if
/// a machine ever switches users — without re-aliasing the same pair.
#[derive(Debug, Clone, Default, Serialize, Deserialize, PartialEq, Eq)]
pub struct Identity {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub user_db_id: Option<String>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub org_id: Option<String>,
    #[serde(default)]
    pub aliased_pairs: Vec<AliasedPair>,
}

#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
pub struct AliasedPair {
    pub agent_id: String,
    pub user_db_id: String,
}

pub fn identity_path(openlatch_dir: &Path) -> PathBuf {
    openlatch_dir.join("identity.json")
}

pub fn read(path: &Path) -> Option<Identity> {
    let raw = std::fs::read_to_string(path).ok()?;
    serde_json::from_str(&raw).ok()
}

pub fn write(path: &Path, identity: &Identity) -> std::io::Result<()> {
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let serialized = serde_json::to_string_pretty(identity).map_err(std::io::Error::other)?;
    std::fs::write(path, serialized)
}

/// Was this (agent_id → user_db_id) pair already aliased? If so, skip the
/// `$create_alias` POST — PostHog tolerates duplicates but they're noise.
pub fn already_aliased(identity: &Identity, agent_id: &str, user_db_id: &str) -> bool {
    identity
        .aliased_pairs
        .iter()
        .any(|p| p.agent_id == agent_id && p.user_db_id == user_db_id)
}

/// Construct the `$create_alias` event payload.
///
/// PostHog's `$create_alias` is a regular event with two required props:
/// `distinct_id` (the canonical user) and `alias` (the prior anonymous id).
/// All super-properties get merged in by the client wrapper as usual.
pub fn create_alias_event(agent_id: &str, user_db_id: &str) -> Event {
    let mut props = Map::new();
    props.insert("alias".into(), json!(agent_id));
    props.insert("distinct_id".into(), json!(user_db_id));
    Event {
        name: "$create_alias".into(),
        properties: props,
    }
}

/// Record a successful auth: persist user_db_id + emit `$create_alias` once.
///
/// Side effects:
/// - Reads `identity.json` (or starts blank).
/// - If the (agent_id → user_db_id) pair has not been aliased yet, emits
///   `$create_alias` and records the pair.
/// - Always overwrites `user_db_id` and `org_id` to the latest values.
/// - Persists `identity.json` atomically.
///
/// All errors are swallowed: telemetry failures must not break auth.
pub fn record_auth_success(
    handle: &TelemetryHandle,
    openlatch_dir: &Path,
    agent_id: &str,
    user_db_id: &str,
    org_id: Option<&str>,
) {
    let path = identity_path(openlatch_dir);
    let mut identity = read(&path).unwrap_or_default();

    if !already_aliased(&identity, agent_id, user_db_id) {
        handle.capture(create_alias_event(agent_id, user_db_id));
        identity.aliased_pairs.push(AliasedPair {
            agent_id: agent_id.into(),
            user_db_id: user_db_id.into(),
        });
    }

    identity.user_db_id = Some(user_db_id.into());
    if let Some(o) = org_id {
        identity.org_id = Some(o.into());
    }

    let _ = write(&path, &identity);
}

/// Distinct-id resolver used at process startup: prefer `user_db_id` from
/// the persisted identity file, fall back to `agent_id`.
pub fn resolve_distinct_id(openlatch_dir: &Path, agent_id: &str) -> String {
    read(&identity_path(openlatch_dir))
        .and_then(|i| i.user_db_id)
        .unwrap_or_else(|| agent_id.to_string())
}

/// True if a `QueuedEvent` looks like a `$create_alias` — used by the
/// network layer if we ever need special handling. Currently informational.
#[allow(dead_code)]
pub fn is_alias_event(event: &QueuedEvent) -> bool {
    event.name == "$create_alias"
}

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

    #[test]
    fn test_create_alias_event_shape() {
        let e = create_alias_event("agt_x", "user_42");
        assert_eq!(e.name, "$create_alias");
        assert_eq!(e.properties["alias"], "agt_x");
        assert_eq!(e.properties["distinct_id"], "user_42");
    }

    #[test]
    fn test_already_aliased_returns_true_after_record() {
        let mut id = Identity::default();
        id.aliased_pairs.push(AliasedPair {
            agent_id: "agt_x".into(),
            user_db_id: "user_42".into(),
        });
        assert!(already_aliased(&id, "agt_x", "user_42"));
        assert!(!already_aliased(&id, "agt_x", "user_99"));
        assert!(!already_aliased(&id, "agt_y", "user_42"));
    }

    #[test]
    fn test_write_then_read_roundtrip() {
        let tmp = TempDir::new().unwrap();
        let path = identity_path(tmp.path());
        let identity = Identity {
            user_db_id: Some("user_42".into()),
            org_id: Some("org_7".into()),
            aliased_pairs: vec![AliasedPair {
                agent_id: "agt_x".into(),
                user_db_id: "user_42".into(),
            }],
        };
        write(&path, &identity).unwrap();
        let loaded = read(&path).unwrap();
        assert_eq!(loaded, identity);
    }

    #[test]
    fn test_resolve_distinct_id_prefers_user_db_id_when_present() {
        let tmp = TempDir::new().unwrap();
        let identity = Identity {
            user_db_id: Some("user_42".into()),
            ..Default::default()
        };
        write(&identity_path(tmp.path()), &identity).unwrap();
        assert_eq!(resolve_distinct_id(tmp.path(), "agt_x"), "user_42");
    }

    #[test]
    fn test_resolve_distinct_id_falls_back_to_agent_id() {
        let tmp = TempDir::new().unwrap();
        // No identity file written.
        assert_eq!(resolve_distinct_id(tmp.path(), "agt_x"), "agt_x");
    }

    #[test]
    fn test_record_auth_success_aliases_once_and_persists() {
        let tmp = TempDir::new().unwrap();
        let handle = TelemetryHandle::disabled("agt_x".into(), false);
        // disabled handle => capture is no-op, but persistence still runs.
        record_auth_success(&handle, tmp.path(), "agt_x", "user_42", Some("org_7"));
        let loaded = read(&identity_path(tmp.path())).unwrap();
        assert_eq!(loaded.user_db_id.as_deref(), Some("user_42"));
        assert_eq!(loaded.org_id.as_deref(), Some("org_7"));
        assert_eq!(loaded.aliased_pairs.len(), 1);

        // Second call with same pair must not duplicate the alias.
        record_auth_success(&handle, tmp.path(), "agt_x", "user_42", Some("org_7"));
        let loaded = read(&identity_path(tmp.path())).unwrap();
        assert_eq!(loaded.aliased_pairs.len(), 1);
    }
}