openlatch-client 0.5.4

//! Bounded, caller-owned session state for live schema-2 evaluation.

use std::collections::{HashMap, VecDeque};
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Mutex;
use std::time::{Duration, Instant};

use crate::zone_eval::{SessionState, StateLayout};

const DEFAULT_CAPACITY: usize = 1_024;
const DEFAULT_IDLE_TTL: Duration = Duration::from_secs(60 * 60);

#[derive(Clone)]
struct Entry {
    state: SessionState,
    layout: StateLayout,
    touched: Instant,
}

#[derive(Default)]
struct Inner {
    entries: HashMap<String, Entry>,
    order: VecDeque<String>,
}

/// Per-daemon bounded state store. Missing and incompatible state both return
/// `None`, preserving the evaluator's `on_evict` semantics.
pub struct SessionStateManager {
    inner: Mutex<Inner>,
    capacity: usize,
    idle_ttl: Duration,
    evicted: AtomicU64,
}

/// Bounded daemon-owned history for tool-dispatch interventions.
///
/// This state is deliberately separate from the evaluator register file: it is
/// adapter delivery state, not an input to the pure policy engine. Entries keep
/// hashes and paths only; raw tool input and output never survive the request.
#[derive(Clone, Debug, Default)]
pub struct DispatchSessionState {
    pub last_completed_key: Option<String>,
    pub last_action_key: Option<String>,
    pub last_result_hash: Option<String>,
    pub last_result_was_error: bool,
    pub identical_completions: u64,
    pub identical_error_completions: u64,
    pub pending_output_paths: HashMap<String, String>,
    pub seen_completion_ids: VecDeque<String>,
    pub last_steer_key: Option<String>,
    pub last_steer_at: Option<Instant>,
    pub pending_stop_reason: Option<String>,
    pub remaining_stop_blocks: u8,
}

#[derive(Default)]
struct DispatchInner {
    entries: HashMap<String, (DispatchSessionState, Instant)>,
    order: VecDeque<String>,
}

/// Shares the evaluator store's capacity/TTL policy without coupling adapter
/// history to a bundle's register layout.
pub struct DispatchStateManager {
    inner: Mutex<DispatchInner>,
    capacity: usize,
    idle_ttl: Duration,
}

impl Default for DispatchStateManager {
    fn default() -> Self {
        Self::new(DEFAULT_CAPACITY, DEFAULT_IDLE_TTL)
    }
}

impl DispatchStateManager {
    pub fn new(capacity: usize, idle_ttl: Duration) -> Self {
        Self {
            inner: Mutex::new(DispatchInner::default()),
            capacity: capacity.max(1),
            idle_ttl,
        }
    }

    /// Mutate one session atomically and evict stale/old entries on the same
    /// bounded schedule as evaluator state.
    pub fn with_session<R>(
        &self,
        session_id: &str,
        now: Instant,
        f: impl FnOnce(&mut DispatchSessionState) -> R,
    ) -> R {
        let mut inner = self.inner.lock().unwrap_or_else(|error| error.into_inner());
        if inner
            .entries
            .get(session_id)
            .is_some_and(|(_, touched)| now.saturating_duration_since(*touched) > self.idle_ttl)
        {
            inner.entries.remove(session_id);
            inner.order.retain(|key| key != session_id);
        }
        inner.order.retain(|key| key != session_id);
        inner.order.push_back(session_id.to_string());
        let (state, touched) = inner
            .entries
            .entry(session_id.to_string())
            .or_insert_with(|| (DispatchSessionState::default(), now));
        *touched = now;
        let result = f(state);
        while inner.entries.len() > self.capacity {
            let Some(oldest) = inner.order.pop_front() else {
                break;
            };
            inner.entries.remove(&oldest);
        }
        result
    }

    pub fn remove(&self, session_id: &str) {
        let Ok(mut inner) = self.inner.lock() else {
            return;
        };
        inner.entries.remove(session_id);
        inner.order.retain(|key| key != session_id);
    }
}

impl Default for SessionStateManager {
    fn default() -> Self {
        Self::new(DEFAULT_CAPACITY, DEFAULT_IDLE_TTL)
    }
}

impl SessionStateManager {
    pub fn new(capacity: usize, idle_ttl: Duration) -> Self {
        Self {
            inner: Mutex::new(Inner::default()),
            capacity: capacity.max(1),
            idle_ttl,
            evicted: AtomicU64::new(0),
        }
    }

    /// Take a compatible snapshot without manufacturing a blank state.
    pub fn get(
        &self,
        session_id: &str,
        layout: &StateLayout,
        now: Instant,
    ) -> Option<SessionState> {
        let mut inner = self.inner.lock().ok()?;
        let expired_or_changed = inner.entries.get(session_id).is_some_and(|entry| {
            now.saturating_duration_since(entry.touched) > self.idle_ttl
                || entry.layout != *layout
                || !entry.state.matches(layout)
        });
        if expired_or_changed {
            inner.entries.remove(session_id);
            inner.order.retain(|key| key != session_id);
            self.note_eviction(session_id, "idle_or_layout_change");
            return None;
        }
        let state = inner.entries.get_mut(session_id)?;
        state.touched = now;
        Some(state.state.clone())
    }

    /// Commit only a successfully evaluated state.
    pub fn commit(
        &self,
        session_id: String,
        layout: StateLayout,
        state: SessionState,
        now: Instant,
    ) {
        if !state.matches(&layout) {
            return;
        }
        let Ok(mut inner) = self.inner.lock() else {
            return;
        };
        inner.order.retain(|key| key != &session_id);
        inner.order.push_back(session_id.clone());
        inner.entries.insert(
            session_id,
            Entry {
                state,
                layout,
                touched: now,
            },
        );
        while inner.entries.len() > self.capacity {
            let Some(oldest) = inner.order.pop_front() else {
                break;
            };
            if inner.entries.remove(&oldest).is_some() {
                self.note_eviction(&oldest, "capacity");
            }
        }
    }

    /// Evaluate and commit while holding the session store lock, so two
    /// concurrent hooks for one session cannot both advance from the same
    /// register snapshot and lose an update.
    pub fn evaluate<R>(
        &self,
        session_id: &str,
        layout: StateLayout,
        now: Instant,
        evaluate: impl FnOnce(Option<&SessionState>) -> (R, SessionState),
    ) -> R {
        let mut inner = self.inner.lock().unwrap_or_else(|error| error.into_inner());
        let incompatible = inner.entries.get(session_id).is_some_and(|entry| {
            now.saturating_duration_since(entry.touched) > self.idle_ttl
                || entry.layout != layout
                || !entry.state.matches(&layout)
        });
        if incompatible {
            inner.entries.remove(session_id);
            inner.order.retain(|key| key != session_id);
            self.note_eviction(session_id, "idle_or_layout_change");
        }
        let prior = inner.entries.get(session_id).map(|entry| &entry.state);
        let (result, state) = evaluate(prior);
        if state.matches(&layout) {
            inner.order.retain(|key| key != session_id);
            inner.order.push_back(session_id.to_string());
            inner.entries.insert(
                session_id.to_string(),
                Entry {
                    state,
                    layout,
                    touched: now,
                },
            );
            while inner.entries.len() > self.capacity {
                let Some(oldest) = inner.order.pop_front() else {
                    break;
                };
                if inner.entries.remove(&oldest).is_some() {
                    self.note_eviction(&oldest, "capacity");
                }
            }
        }
        result
    }

    /**
     * Read one immutable prior snapshot without touching its recency or evicting
     * it. A caller that ultimately delivers the evaluated action follows with
     * exactly one [`Self::commit`]; failed previews leave the store unchanged.
     */
    pub fn preview(
        &self,
        session_id: &str,
        layout: &StateLayout,
        now: Instant,
    ) -> Option<SessionState> {
        let inner = self.inner.lock().ok()?;
        let entry = inner.entries.get(session_id)?;
        (now.saturating_duration_since(entry.touched) <= self.idle_ttl
            && entry.layout == *layout
            && entry.state.matches(layout))
        .then(|| entry.state.clone())
    }

    pub fn remove(&self, session_id: &str) {
        let Ok(mut inner) = self.inner.lock() else {
            return;
        };
        inner.entries.remove(session_id);
        inner.order.retain(|key| key != session_id);
    }

    pub fn evicted_count(&self) -> u64 {
        self.evicted.load(Ordering::Relaxed)
    }

    fn note_eviction(&self, session_id: &str, reason: &str) {
        self.evicted.fetch_add(1, Ordering::Relaxed);
        tracing::warn!(
            target: "policy",
            code = "OL-1214",
            %session_id,
            %reason,
            "session policy state evicted; on_evict applies on the next evaluation"
        );
    }
}

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

    fn layout(c: usize) -> StateLayout {
        StateLayout {
            c,
            ..StateLayout::default()
        }
    }

    #[test]
    fn missing_and_layout_changed_state_are_unavailable() {
        let manager = SessionStateManager::new(2, Duration::from_secs(60));
        let now = Instant::now();
        assert_eq!(manager.get("s", &layout(1), now), None);
        manager.commit("s".into(), layout(1), SessionState::blank(&layout(1)), now);
        assert!(manager.get("s", &layout(1), now).is_some());
        assert_eq!(manager.get("s", &layout(2), now), None);
        assert_eq!(manager.evicted_count(), 1);
    }

    #[test]
    fn capacity_evicts_oldest() {
        let manager = SessionStateManager::new(1, Duration::from_secs(60));
        let now = Instant::now();
        manager.commit("a".into(), layout(0), SessionState::default(), now);
        manager.commit("b".into(), layout(0), SessionState::default(), now);
        assert_eq!(manager.get("a", &layout(0), now), None);
        assert!(manager.get("b", &layout(0), now).is_some());
        assert_eq!(manager.evicted_count(), 1);
    }

    #[test]
    fn concurrent_session_transactions_retain_every_update() {
        let manager = std::sync::Arc::new(SessionStateManager::default());
        let start = std::sync::Arc::new(std::sync::Barrier::new(16));
        let mut threads = Vec::new();
        for _ in 0..16 {
            let manager = manager.clone();
            let start = start.clone();
            threads.push(std::thread::spawn(move || {
                start.wait();
                manager.evaluate("shared", layout(1), Instant::now(), |prior| {
                    let mut next = prior
                        .cloned()
                        .unwrap_or_else(|| SessionState::blank(&layout(1)));
                    next.c[0] += 1;
                    ((), next)
                });
            }));
        }
        for thread in threads {
            thread.join().expect("transaction thread");
        }
        assert_eq!(
            manager.get("shared", &layout(1), Instant::now()).unwrap().c,
            vec![16]
        );
    }
}