lean-ctx-engine 3.12.3

// SPDX-License-Identifier: Apache-2.0
//! Per-session value snapshot: the small file every display channel reads.
//!
//! Written by the MCP server after tool calls (throttled, atomic) to
//! `<data_dir>/value/sessions/<id>.json`, `value/current.json` and, keyed by
//! the session's project root, `value/projects/<hash>.json`. Readers (status
//! line, prompt segment, hooks, IDE) never open the ledger on their hot path;
//! `lean-ctx value` re-derives every number from the chains instead.

use std::path::{Path, PathBuf};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use sha2::{Digest, Sha256};

use crate::core::security_events::SecurityCounts;

pub const SCHEMA: u32 = 1;

/// Minimum interval between two snapshot writes of one process.
const WRITE_INTERVAL: Duration = Duration::from_secs(1);

#[derive(Debug, Clone, Default, PartialEq, Serialize, Deserialize)]
#[serde(default)]
pub struct ValueSnapshot {
    pub schema: u32,
    pub session_id: String,
    pub project_root: Option<String>,
    pub started_at: Option<DateTime<Utc>>,
    pub updated_at: Option<DateTime<Utc>>,
    pub tool_calls: u64,
    /// Tokens lean-ctx's tools would have sent without compression.
    pub tokens_input: u64,
    /// Tokens kept out of the model's context (measured, per tool result).
    pub tokens_saved: u64,
    pub cache_hits: u64,
    pub files_read: u64,
    pub commands_run: u64,
    pub security: SecurityCounts,
    /// The agent process that started the writing MCP server (see [`super::host`]).
    pub host_pid: Option<u32>,
}

impl ValueSnapshot {
    pub fn from_session(session: &crate::core::session::SessionState) -> Self {
        let s = &session.stats;
        Self {
            schema: SCHEMA,
            session_id: session.id.clone(),
            project_root: session.project_root.clone(),
            started_at: Some(session.started_at),
            updated_at: Some(Utc::now()),
            tool_calls: u64::from(s.total_tool_calls),
            tokens_input: s.total_tokens_input,
            tokens_saved: s.total_tokens_saved,
            cache_hits: u64::from(s.cache_hits),
            files_read: u64::from(s.files_read),
            commands_run: u64::from(s.commands_run),
            security: s.security,
            host_pid: super::host::host_pid(),
        }
    }

    /// Share of tool input kept out of context, in percent (derived).
    pub fn saved_pct(&self) -> Option<f64> {
        (self.tokens_input > 0).then(|| self.tokens_saved as f64 * 100.0 / self.tokens_input as f64)
    }

    /// True when the snapshot was updated within `max_age`. Display channels
    /// show nothing for a stale snapshot — never a misleading old number.
    pub fn is_fresh(&self, max_age: Duration) -> bool {
        let Some(updated) = self.updated_at else {
            return false;
        };
        let age = Utc::now().signed_duration_since(updated);
        age.to_std().map_or(true, |age| age <= max_age)
    }

    /// Nothing worth showing yet.
    pub fn is_empty(&self) -> bool {
        self.tokens_saved == 0 && self.security.is_empty()
    }
}

pub fn value_dir() -> Option<PathBuf> {
    crate::core::data_dir::lean_ctx_data_dir()
        .ok()
        .map(|dir| dir.join("value"))
}

fn sanitize_id(id: &str) -> Option<&str> {
    let ok = !id.is_empty()
        && id.len() <= 128
        && id
            .chars()
            .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '_' || c == '.')
        && !id.starts_with('.');
    ok.then_some(id)
}

pub fn session_path(dir: &Path, id: &str) -> Option<PathBuf> {
    sanitize_id(id).map(|id| dir.join("sessions").join(format!("{id}.json")))
}

pub fn current_path(dir: &Path) -> PathBuf {
    dir.join("current.json")
}

/// Where the latest snapshot for project `root` lives. Hosts (a Claude status
/// line, a Stop hook) know their working directory but not lean-ctx's session
/// id; this pointer is how they find the numbers of the session serving them.
pub fn project_path(dir: &Path, root: &str) -> PathBuf {
    let digest = crate::core::agent_identity::hex_encode(&Sha256::digest(root.as_bytes()));
    dir.join("projects").join(format!("{}.json", &digest[..16]))
}

/// Writes `snap` as the session's snapshot, as `current.json` and, when the
/// session has a project root, as that project's snapshot.
pub fn write_to(dir: &Path, snap: &ValueSnapshot) -> Result<(), String> {
    let path = session_path(dir, &snap.session_id).ok_or("invalid session id")?;
    let json = serde_json::to_vec(snap).map_err(|e| e.to_string())?;
    let project = snap
        .project_root
        .as_deref()
        .filter(|root| !root.is_empty())
        .map(|root| project_path(dir, root));
    for target in [Some(&path), project.as_ref()].into_iter().flatten() {
        if let Some(parent) = target.parent() {
            std::fs::create_dir_all(parent).map_err(|e| e.to_string())?;
        }
        crate::core::atomic_fs::try_atomic_write(target, &json, None).map_err(|e| e.to_string())?;
    }
    crate::core::atomic_fs::try_atomic_write(&current_path(dir), &json, None)
        .map_err(|e| e.to_string())
}

/// Throttled write used by the MCP server: at most once per
/// `WRITE_INTERVAL` (1 s), except that a changed security tally always writes so a
/// blocked command or a kept-out secret shows up immediately.
pub fn write_throttled(snap: &ValueSnapshot) {
    static LAST: Mutex<Option<(Instant, SecurityCounts)>> = Mutex::new(None);
    {
        let Ok(mut last) = LAST.lock() else { return };
        if let Some((at, security)) = *last
            && at.elapsed() < WRITE_INTERVAL
            && security == snap.security
        {
            return;
        }
        *last = Some((Instant::now(), snap.security));
    }
    let Some(dir) = value_dir() else { return };
    if let Err(error) = write_to(&dir, snap) {
        tracing::debug!("lean-ctx: value snapshot write failed: {error}");
    }
}

pub fn read_path(path: &Path) -> Option<ValueSnapshot> {
    let bytes = std::fs::read(path).ok()?;
    serde_json::from_slice(&bytes).ok()
}

/// The snapshot of session `id`, or of the most recently active session.
pub fn load(id: Option<&str>) -> Option<ValueSnapshot> {
    let dir = value_dir()?;
    match id {
        Some(id) => read_path(&session_path(&dir, id)?),
        None => read_path(&current_path(&dir)),
    }
}

/// The latest snapshot of the project containing `cwd`: the nearest ancestor
/// (canonical first, then as given) that has one.
pub fn load_for_dir(cwd: &Path) -> Option<ValueSnapshot> {
    load_for_dir_in(&value_dir()?, cwd)
}

pub fn load_for_dir_in(dir: &Path, cwd: &Path) -> Option<ValueSnapshot> {
    let canonical = crate::core::pathutil::safe_canonicalize_or_self(cwd);
    [canonical.as_path(), cwd]
        .into_iter()
        .flat_map(Path::ancestors)
        .filter_map(Path::to_str)
        .find_map(|root| {
            read_path(&project_path(dir, root))
                .filter(|snap| snap.project_root.as_deref() == Some(root))
        })
}

/// The snapshot of the session serving the calling agent: a status line or hook
/// is started by the same agent process that started its lean-ctx server, so
/// that server's snapshot names one of the caller's ancestors as `host_pid`.
/// Falls back to the project snapshot (whichever session wrote last) only while
/// no server in the project records its host yet (all from older builds): once
/// one does, a caller whose own session is not found yet (a server just
/// restarted) shows nothing rather than a neighbour's numbers.
pub fn load_for_caller_in(dir: &Path, cwd: &Path, max_age: Duration) -> Option<ValueSnapshot> {
    load_own_or_legacy_in(dir, cwd, &super::host::ancestors(), max_age)
}

/// [`load_for_caller_in`] for an explicit ancestor chain.
pub fn load_own_or_legacy_in(
    dir: &Path,
    cwd: &Path,
    hosts: &[u32],
    max_age: Duration,
) -> Option<ValueSnapshot> {
    let sessions = fresh_project_sessions(dir, cwd, max_age);
    if let Some(own) = pick_by_host(&sessions, hosts) {
        return Some(own);
    }
    if sessions.iter().any(|s| s.host_pid.is_some()) {
        return None;
    }
    load_for_dir_in(dir, cwd).filter(|s| s.is_fresh(max_age))
}

/// Among fresh sessions of the project containing `cwd`, the one whose host is
/// the nearest of `hosts` (nearest first); the newest when one host has several.
pub fn load_for_hosts_in(
    dir: &Path,
    cwd: &Path,
    hosts: &[u32],
    max_age: Duration,
) -> Option<ValueSnapshot> {
    pick_by_host(&fresh_project_sessions(dir, cwd, max_age), hosts)
}

/// Whether any fresh session of the project records its host: from then on a
/// caller that finds no session of its own must not borrow another's.
pub fn host_recorded_in(dir: &Path, cwd: &Path, max_age: Duration) -> bool {
    fresh_project_sessions(dir, cwd, max_age)
        .iter()
        .any(|s| s.host_pid.is_some())
}

fn pick_by_host(sessions: &[ValueSnapshot], hosts: &[u32]) -> Option<ValueSnapshot> {
    sessions
        .iter()
        .filter_map(|s| Some((hosts.iter().position(|&h| Some(h) == s.host_pid)?, s)))
        .min_by(|(ra, a), (rb, b)| ra.cmp(rb).then(b.updated_at.cmp(&a.updated_at)))
        .map(|(_, s)| s.clone())
}

/// Every session of the project containing `cwd` written within `max_age`.
fn fresh_project_sessions(dir: &Path, cwd: &Path, max_age: Duration) -> Vec<ValueSnapshot> {
    let canonical = crate::core::pathutil::safe_canonicalize_or_self(cwd);
    let in_project = |root: &str| canonical.starts_with(root) || cwd.starts_with(root);
    let (Some(cutoff), Ok(entries)) = (
        std::time::SystemTime::now().checked_sub(max_age),
        std::fs::read_dir(dir.join("sessions")),
    ) else {
        return Vec::new();
    };
    entries
        .filter_map(Result::ok)
        // Cheap mtime filter first: the directory keeps every past session.
        .filter(|e| {
            e.metadata()
                .and_then(|m| m.modified())
                .is_ok_and(|t| t >= cutoff)
        })
        .filter_map(|e| read_path(&e.path()))
        .filter(|s| s.is_fresh(max_age) && s.project_root.as_deref().is_some_and(in_project))
        .collect()
}

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

    fn snap(id: &str) -> ValueSnapshot {
        ValueSnapshot {
            schema: SCHEMA,
            session_id: id.into(),
            updated_at: Some(Utc::now()),
            tokens_input: 1_000,
            tokens_saved: 400,
            ..ValueSnapshot::default()
        }
    }

    #[test]
    fn project_snapshot_is_found_from_any_subdirectory() {
        let data = tempfile::tempdir().unwrap();
        let project = tempfile::tempdir().unwrap();
        let root = crate::core::pathutil::safe_canonicalize_or_self(project.path());
        let nested = root.join("src/deep");
        std::fs::create_dir_all(&nested).unwrap();
        let mut s = snap("sess-p");
        s.project_root = Some(root.to_string_lossy().into_owned());
        write_to(data.path(), &s).unwrap();

        assert_eq!(load_for_dir_in(data.path(), &nested), Some(s.clone()));
        assert_eq!(load_for_dir_in(data.path(), &root), Some(s));
        let elsewhere = tempfile::tempdir().unwrap();
        assert_eq!(load_for_dir_in(data.path(), elsewhere.path()), None);
    }

    #[test]
    fn each_agent_reads_its_own_session_not_the_last_writer() {
        let data = tempfile::tempdir().unwrap();
        let project = tempfile::tempdir().unwrap();
        let root = crate::core::pathutil::safe_canonicalize_or_self(project.path());
        let root_str = root.to_string_lossy().into_owned();
        let mut ours = snap("ours");
        ours.project_root = Some(root_str.clone());
        ours.host_pid = Some(4_100);
        ours.tokens_saved = 4_200;
        let mut theirs = snap("theirs");
        theirs.project_root = Some(root_str);
        theirs.host_pid = Some(9_900);
        theirs.tokens_saved = 357_500;
        write_to(data.path(), &ours).unwrap();
        write_to(data.path(), &theirs).unwrap(); // last writer owns the project file

        let max = Duration::from_hours(12);
        let found = load_for_hosts_in(data.path(), &root, &[77, 4_100, 1_000], max).unwrap();
        assert_eq!(found.session_id, "ours");
        // No ancestor matches (our server just restarted): nothing, never the
        // neighbour the project file names, because servers record their host.
        assert_eq!(load_for_hosts_in(data.path(), &root, &[5, 6], max), None);
        assert_eq!(
            load_for_dir_in(data.path(), &root).unwrap().session_id,
            "theirs"
        );
        assert_eq!(
            load_own_or_legacy_in(data.path(), &root, &[5, 6], max),
            None
        );
        assert!(host_recorded_in(data.path(), &root, max));
        let own = load_own_or_legacy_in(data.path(), &root, &[4_100], max).unwrap();
        assert_eq!(own.session_id, "ours");
        let elsewhere = tempfile::tempdir().unwrap();
        assert_eq!(
            load_for_hosts_in(data.path(), elsewhere.path(), &[4_100], max),
            None
        );
    }

    #[test]
    fn servers_without_a_host_record_keep_the_project_fallback() {
        let data = tempfile::tempdir().unwrap();
        let project = tempfile::tempdir().unwrap();
        let root = crate::core::pathutil::safe_canonicalize_or_self(project.path());
        let mut legacy = snap("legacy");
        legacy.project_root = Some(root.to_string_lossy().into_owned());
        write_to(data.path(), &legacy).unwrap();
        let max = Duration::from_hours(12);
        assert!(!host_recorded_in(data.path(), &root, max));
        let found = load_own_or_legacy_in(data.path(), &root, &[1_234], max).unwrap();
        assert_eq!(found.session_id, "legacy");
    }

    #[test]
    fn roundtrip_through_session_and_current_file() {
        let dir = tempfile::tempdir().unwrap();
        let s = snap("sess-1");
        write_to(dir.path(), &s).unwrap();
        assert_eq!(
            read_path(&session_path(dir.path(), "sess-1").unwrap()),
            Some(s.clone())
        );
        assert_eq!(read_path(&current_path(dir.path())), Some(s));
    }

    #[test]
    fn rejects_path_traversal_ids() {
        let dir = tempfile::tempdir().unwrap();
        assert!(session_path(dir.path(), "../evil").is_none());
        assert!(session_path(dir.path(), "a/b").is_none());
        assert!(session_path(dir.path(), "").is_none());
        assert!(write_to(dir.path(), &snap("../x")).is_err());
    }

    #[test]
    fn staleness_is_detected() {
        let mut s = snap("s");
        assert!(s.is_fresh(Duration::from_mins(1)));
        s.updated_at = Some(Utc::now() - chrono::Duration::hours(2));
        assert!(!s.is_fresh(Duration::from_mins(1)));
        s.updated_at = None;
        assert!(!s.is_fresh(Duration::from_mins(1)));
    }

    #[test]
    fn saved_pct_is_derived_from_measured_counts() {
        assert_eq!(snap("s").saved_pct(), Some(40.0));
        assert_eq!(ValueSnapshot::default().saved_pct(), None);
    }

    #[test]
    fn unknown_fields_and_missing_fields_are_tolerated() {
        let s: ValueSnapshot =
            serde_json::from_str(r#"{"session_id":"x","future_field":1}"#).unwrap();
        assert_eq!(s.session_id, "x");
        assert!(s.is_empty());
    }
}