bamboo_engine/

session_repository.rs

//! Canonical session coordinator owned by the framework.
//!
//! [`SessionRepository`] bundles the three tiers a Bamboo session lives in — the
//! in-memory [`SessionCache`], the durable [`Storage`], and the
//! merge-on-write [`LockedSessionStore`] — and provides the one canonical
//! load/save coordination (cache → storage → backfill, and dual-write).
//!
//! This is a *framework* capability, not a server one: previously the
//! coordination lived only as inherent methods on the server's `AppState`,
//! which meant anything outside the HTTP server (the SDK, in-process embedders)
//! could not load or persist sessions consistently. `SessionRepository` lets any
//! caller that holds the three tiers share the exact same behaviour; the
//! server's `AppState` now delegates to it.

use std::sync::Arc;

use bamboo_agent_core::storage::Storage;
use bamboo_agent_core::Session;
use bamboo_storage::LockedSessionStore;

use crate::{read_cached_session, SessionCache};

/// Framework-owned coordinator over a session's cache / storage / persistence
/// tiers. Cheap to clone (all fields are `Arc`).
#[derive(Clone)]
pub struct SessionRepository {
    cache: SessionCache,
    storage: Arc<dyn Storage>,
    persistence: Arc<LockedSessionStore>,
}

impl SessionRepository {
    pub fn new(
        cache: SessionCache,
        storage: Arc<dyn Storage>,
        persistence: Arc<LockedSessionStore>,
    ) -> Self {
        Self {
            cache,
            storage,
            persistence,
        }
    }

    pub fn cache(&self) -> &SessionCache {
        &self.cache
    }

    pub fn storage(&self) -> &Arc<dyn Storage> {
        &self.storage
    }

    pub fn persistence(&self) -> &Arc<LockedSessionStore> {
        &self.persistence
    }

    /// Load a session from the memory cache, falling back to durable storage
    /// (and back-filling the cache on a storage hit). `None` if absent in both.
    pub async fn load(&self, session_id: &str) -> Option<Session> {
        if let Some(session) = read_cached_session(&self.cache, session_id) {
            return Some(session);
        }

        match self.storage.load_session(session_id).await {
            Ok(Some(session)) => {
                self.cache.insert(
                    session_id.to_string(),
                    Arc::new(parking_lot::RwLock::new(session.clone())),
                );
                Some(session)
            }
            _ => None,
        }
    }

    /// Like [`load`](Self::load), but surfaces storage errors instead of
    /// swallowing them to `None`. Cache hit short-circuits; a storage hit
    /// back-fills the cache.
    pub async fn try_load(&self, session_id: &str) -> std::io::Result<Option<Session>> {
        if let Some(session) = read_cached_session(&self.cache, session_id) {
            return Ok(Some(session));
        }
        let loaded = self.storage.load_session(session_id).await?;
        if let Some(ref session) = loaded {
            self.cache.insert(
                session_id.to_string(),
                Arc::new(parking_lot::RwLock::new(session.clone())),
            );
        }
        Ok(loaded)
    }

    /// Persist the session (merge-on-write) and refresh the cache, surfacing
    /// storage errors. Use [`save_and_cache`](Self::save_and_cache) for the
    /// fire-and-forget variant that logs and continues on failure.
    pub async fn save(&self, session: &mut Session) -> std::io::Result<()> {
        self.persistence.merge_save_runtime(session).await?;
        self.cache.insert(
            session.id.clone(),
            Arc::new(parking_lot::RwLock::new(session.clone())),
        );
        Ok(())
    }

    /// Atomically mutate the latest durable runtime session and refresh the
    /// cache with the saved value. This is the safe path for narrow metadata
    /// indexes that can be updated concurrently with runner message writes.
    pub async fn update_runtime_session<F>(
        &self,
        session_id: &str,
        metadata_keys: &[&str],
        mutate: F,
    ) -> std::io::Result<Option<Session>>
    where
        F: FnOnce(&mut Session),
    {
        let saved = self
            .persistence
            .update_runtime_config(session_id, mutate)
            .await?;
        if let (Some(saved), Some(cached)) = (saved.as_ref(), self.cache.get(session_id)) {
            let mut cached = cached.write();
            for key in metadata_keys {
                if let Some(value) = saved.metadata.get(*key) {
                    cached.metadata.insert((*key).to_string(), value.clone());
                } else {
                    cached.metadata.remove(*key);
                }
            }
        }
        Ok(saved)
    }

    /// Load a session, creating a fresh `Session::new(id, model)` if absent.
    pub async fn load_or_create(&self, session_id: &str, model: &str) -> Session {
        if let Some(session) = self.load(session_id).await {
            return session;
        }
        Session::new(session_id.to_string(), model.to_string())
    }

    /// Load a session, reconciling the memory and storage copies via a
    /// preference heuristic: storage wins when it is strictly newer, or when it
    /// is the same age but still carries a pending question memory lost. Storage
    /// is **never** preferred when it is strictly older than memory.
    ///
    /// The cache is refreshed cache-aside but with a no-regression guarantee:
    /// `load_merged` never overwrites a newer cached session with an older
    /// storage copy, so it is safe to call from hot read paths.
    pub async fn load_merged(&self, session_id: &str) -> Option<Session> {
        let memory_session = read_cached_session(&self.cache, session_id);
        let storage_session = self
            .storage
            .load_session(session_id)
            .await
            .unwrap_or_default();

        match (memory_session, storage_session) {
            (Some(memory), Some(storage)) => {
                let prefer_storage = should_prefer_storage(&memory, &storage);
                let diverged = prefer_storage || memory.messages.len() != storage.messages.len();
                let chosen_len = if prefer_storage {
                    storage.messages.len()
                } else {
                    memory.messages.len()
                };
                macro_rules! merged_log {
                    ($level:ident) => {
                        tracing::$level!(
                            "[{}] load_session_merged: memory={} msgs (updated_at={}), storage={} msgs (updated_at={}), prefer_storage={} -> chose {} msgs",
                            session_id,
                            memory.messages.len(),
                            memory.updated_at,
                            storage.messages.len(),
                            storage.updated_at,
                            prefer_storage,
                            chosen_len,
                        )
                    };
                }
                if diverged {
                    merged_log!(debug);
                } else {
                    merged_log!(trace);
                }
                let memory_updated_at = memory.updated_at;
                let chosen = if prefer_storage { storage } else { memory };
                // Cache-aside refresh with a hard no-regression invariant: only
                // write back when we actually reconciled *to storage* (a memory
                // win is already the cached copy; re-inserting it would needlessly
                // replace a possibly-live Arc) AND the reconciled copy is not
                // older than what memory already holds. This is what makes
                // `load_merged` safe on hot read paths — it can never clobber a
                // freshly-updated session with a stale storage copy.
                if prefer_storage && chosen.updated_at >= memory_updated_at {
                    self.cache.insert(
                        session_id.to_string(),
                        Arc::new(parking_lot::RwLock::new(chosen.clone())),
                    );
                }
                Some(chosen)
            }
            (Some(memory), None) => Some(memory),
            (None, Some(storage)) => {
                self.cache.insert(
                    session_id.to_string(),
                    Arc::new(parking_lot::RwLock::new(storage.clone())),
                );
                Some(storage)
            }
            (None, None) => None,
        }
    }

    /// Persist the session (merge-on-write, preserving concurrent UI edits to
    /// the authoritative metadata group) and refresh the in-memory cache.
    pub async fn save_and_cache(&self, session: &mut Session) {
        if let Err(error) = self.persistence.merge_save_runtime(session).await {
            tracing::warn!("[{}] Failed to save session: {}", session.id, error);
        }
        self.cache.insert(
            session.id.clone(),
            Arc::new(parking_lot::RwLock::new(session.clone())),
        );
    }
}

fn should_prefer_storage(memory_session: &Session, storage_session: &Session) -> bool {
    // Never reconcile *backwards* to a strictly-older storage copy: if memory is
    // newer it is authoritative (e.g. it just answered and cleared a pending
    // question while storage still holds the stale one). Respecting `updated_at`
    // here is what stops `load_merged` from returning — and caching — stale data.
    if storage_session.updated_at < memory_session.updated_at {
        return false;
    }
    // Storage is same-age or newer: prefer it when strictly newer, or when it
    // still carries a pending question that the (same-age) memory copy lost, so
    // a genuine clarification is never dropped.
    storage_session.updated_at > memory_session.updated_at
        || (memory_session.pending_question.is_none() && storage_session.pending_question.is_some())
}

/// `SessionRepository` is the canonical `RuntimeSessionPersistence`: the runtime
/// can persist a session through the same coordinator (merge-on-write + cache
/// refresh) instead of a bespoke adapter.
#[async_trait::async_trait]
impl bamboo_domain::RuntimeSessionPersistence for SessionRepository {
    async fn save_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
        // Runtime authorization reads through this same cache. Refresh it even
        // when durable storage fails so a current activation can never observe
        // a previous run's skill allowlist. The error is still returned to the
        // caller and durable state remains unchanged.
        let result = self.persistence.merge_save_runtime(session).await;
        self.cache.insert(
            session.id.clone(),
            Arc::new(parking_lot::RwLock::new(session.clone())),
        );
        result
    }

    async fn checkpoint_runtime_session(&self, session: &mut Session) -> std::io::Result<()> {
        // The execute-boundary checkpoint uses LockedSessionStore's atomic
        // append-safe transcript merge, then publishes that reconciled snapshot
        // to the runtime cache.  On failure, leave the existing cache alone: the
        // checkpoint may have failed to load the latest durable transcript, and
        // replacing a fresher cache entry with the stale runner snapshot would
        // set up a later SHRINK write.
        let result = self.persistence.checkpoint_runtime_session(session).await;
        if result.is_ok() {
            self.cache.insert(
                session.id.clone(),
                Arc::new(parking_lot::RwLock::new(session.clone())),
            );
        }
        result
    }

    async fn load_runtime_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
        self.try_load(session_id).await
    }

    async fn clear_legacy_pending_messages(
        &self,
        session_id: &str,
        expected: &[serde_json::Value],
    ) -> std::io::Result<bool> {
        // Do not use the trait default here: `load_runtime_session` is
        // cache-first, so a stale cache could erase a message concurrently
        // appended to the durable legacy queue. Delegate the compare-and-clear
        // to LockedSessionStore's single per-session critical section.
        let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
            self.persistence.as_ref(),
            session_id,
            expected,
        )
        .await?;
        if !cleared {
            return Ok(false);
        }

        // Refresh only after the durable CAS succeeds. A failed comparison or
        // persistence error leaves the live cache untouched.
        let latest = self
            .storage
            .load_session(session_id)
            .await?
            .ok_or_else(|| {
                std::io::Error::new(
                    std::io::ErrorKind::NotFound,
                    format!("session disappeared after clearing legacy inbox: {session_id}"),
                )
            })?;
        self.cache.insert(
            session_id.to_string(),
            Arc::new(parking_lot::RwLock::new(latest)),
        );
        Ok(true)
    }

    async fn append_token_usage_record(
        &self,
        session_id: &str,
        json_line: &str,
    ) -> std::io::Result<()> {
        self.storage
            .append_token_usage_record(session_id, json_line)
            .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use bamboo_agent_core::storage::Storage;
    use chrono::Utc;
    use std::collections::HashMap;
    use std::sync::Mutex;

    #[derive(Default)]
    struct MapStorage {
        sessions: Mutex<HashMap<String, Session>>,
    }

    struct FailingSaveStorage {
        persisted: Mutex<Option<Session>>,
    }

    #[async_trait::async_trait]
    impl Storage for MapStorage {
        async fn save_session(&self, session: &Session) -> std::io::Result<()> {
            self.sessions
                .lock()
                .unwrap()
                .insert(session.id.clone(), session.clone());
            Ok(())
        }
        async fn load_session(&self, session_id: &str) -> std::io::Result<Option<Session>> {
            Ok(self.sessions.lock().unwrap().get(session_id).cloned())
        }
        async fn delete_session(&self, session_id: &str) -> std::io::Result<bool> {
            Ok(self.sessions.lock().unwrap().remove(session_id).is_some())
        }
    }

    #[async_trait::async_trait]
    impl Storage for FailingSaveStorage {
        async fn save_session(&self, _session: &Session) -> std::io::Result<()> {
            Err(std::io::Error::other("injected save failure"))
        }

        async fn load_session(&self, _session_id: &str) -> std::io::Result<Option<Session>> {
            Ok(self.persisted.lock().unwrap().clone())
        }

        async fn delete_session(&self, _session_id: &str) -> std::io::Result<bool> {
            Ok(false)
        }
    }

    fn test_repo(storage: Arc<dyn Storage>) -> SessionRepository {
        let cache: SessionCache = Arc::new(dashmap::DashMap::new());
        let persistence = Arc::new(LockedSessionStore::new(storage.clone()));
        SessionRepository::new(cache, storage, persistence)
    }

    fn cache_put(repo: &SessionRepository, session: &Session) {
        repo.cache().insert(
            session.id.clone(),
            Arc::new(parking_lot::RwLock::new(session.clone())),
        );
    }

    #[tokio::test]
    async fn narrow_runtime_metadata_transaction_preserves_live_and_durable_non_owned_state() {
        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
        let repo = test_repo(storage.clone());
        let id = "narrow-metadata";
        let mut durable = Session::new(id, "durable-model");
        durable.add_message(bamboo_agent_core::Message::user("durable user turn"));
        durable
            .metadata
            .insert("external.durable".to_string(), "keep".to_string());
        storage.save_session(&durable).await.expect("seed durable");

        let mut live = durable.clone();
        live.add_message(bamboo_agent_core::Message::assistant(
            "in-flight assistant tool call",
            None,
        ));
        live.model = "live-model".to_string();
        live.metadata
            .insert("external.live".to_string(), "keep".to_string());
        cache_put(&repo, &live);

        repo.update_runtime_session(id, &["workflow.owned"], |latest| {
            latest
                .metadata
                .insert("workflow.owned".to_string(), "active".to_string());
        })
        .await
        .expect("transaction")
        .expect("session exists");

        let saved = storage
            .load_session(id)
            .await
            .expect("load durable")
            .expect("durable exists");
        assert_eq!(
            saved.messages.len(),
            1,
            "transaction never writes stale live messages"
        );
        assert_eq!(
            saved.metadata.get("external.durable").map(String::as_str),
            Some("keep")
        );
        assert_eq!(
            saved.metadata.get("workflow.owned").map(String::as_str),
            Some("active")
        );

        let cached = read_cached_session(repo.cache(), id).expect("live cache");
        assert_eq!(
            cached.messages.len(),
            2,
            "cache live tool call is not replaced"
        );
        assert_eq!(cached.model, "live-model");
        assert_eq!(
            cached.metadata.get("external.live").map(String::as_str),
            Some("keep")
        );
        assert_eq!(
            cached.metadata.get("workflow.owned").map(String::as_str),
            Some("active")
        );
    }

    /// Regression guard: a strictly-newer in-memory session (e.g. one that just
    /// answered and cleared its pending question) must win over a strictly-older
    /// storage copy that still carries the pending question — both in the value
    /// returned AND in the cache (no clobber).
    #[tokio::test]
    async fn load_merged_does_not_regress_to_older_storage() {
        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
        let repo = test_repo(storage.clone());
        let id = "s1";

        let mut stale = Session::new(id.to_string(), "m");
        stale.set_pending_question(
            "tc1".into(),
            "kind".into(),
            "q?".into(),
            vec!["OK".into()],
            true,
        );
        stale.updated_at = Utc::now() - chrono::Duration::seconds(10);
        storage.save_session(&stale).await.unwrap();

        let mut fresh = Session::new(id.to_string(), "m");
        fresh.updated_at = Utc::now();
        cache_put(&repo, &fresh);

        let merged = repo.load_merged(id).await.expect("session exists");
        assert!(
            merged.pending_question.is_none(),
            "must return the newer answered memory copy, not the stale storage one"
        );
        let cached = read_cached_session(repo.cache(), id).expect("cached");
        assert!(
            cached.pending_question.is_none(),
            "load_merged must never regress the cache to a stale storage copy"
        );
    }

    /// The pending-question recovery still works when storage is the same age:
    /// if memory lost a pending question that same-age storage retains, prefer
    /// storage so a genuine clarification is not dropped.
    #[tokio::test]
    async fn load_merged_recovers_pending_question_from_same_age_storage() {
        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
        let repo = test_repo(storage.clone());
        let id = "s2";
        let ts = Utc::now();

        let mut with_pending = Session::new(id.to_string(), "m");
        with_pending.set_pending_question(
            "tc".into(),
            "k".into(),
            "q".into(),
            vec!["OK".into()],
            true,
        );
        with_pending.updated_at = ts;
        storage.save_session(&with_pending).await.unwrap();

        let mut lost = with_pending.clone();
        lost.clear_pending_question();
        lost.updated_at = ts;
        cache_put(&repo, &lost);

        let merged = repo.load_merged(id).await.expect("session exists");
        assert!(
            merged.pending_question.is_some(),
            "same-age storage carrying a pending question must still be recovered"
        );
    }

    #[tokio::test]
    async fn runtime_publish_refreshes_cache_even_when_storage_fails() {
        let id = "runtime-selection";
        let mut previous = Session::new(id.to_string(), "m");
        previous.metadata.insert(
            "skill_runtime_selected_skill_ids".to_string(),
            "[\"plan\"]".to_string(),
        );
        let storage: Arc<dyn Storage> = Arc::new(FailingSaveStorage {
            persisted: Mutex::new(Some(previous.clone())),
        });
        let repo = test_repo(storage.clone());
        cache_put(&repo, &previous);

        let mut current = previous.clone();
        current.metadata.insert(
            "skill_runtime_selected_skill_ids".to_string(),
            "[\"review\"]".to_string(),
        );
        current.updated_at = Utc::now();

        let result =
            bamboo_domain::RuntimeSessionPersistence::save_runtime_session(&repo, &mut current)
                .await;
        assert!(result.is_err(), "durable failure must still be surfaced");

        let cached = repo.load(id).await.expect("cached current session");
        assert_eq!(
            cached
                .metadata
                .get("skill_runtime_selected_skill_ids")
                .map(String::as_str),
            Some("[\"review\"]")
        );
        let allowlist = bamboo_skills::access_control::extract_skill_allowlist(&cached.metadata)
            .expect("runtime authorization allowlist");
        assert!(allowlist.contains("review"));
        assert!(!allowlist.contains("plan"));
        let durable = storage
            .load_session(id)
            .await
            .expect("load durable state")
            .expect("previous durable session");
        assert_eq!(
            durable
                .metadata
                .get("skill_runtime_selected_skill_ids")
                .map(String::as_str),
            Some("[\"plan\"]")
        );
    }

    #[tokio::test]
    async fn legacy_clear_uses_durable_cas_and_never_erases_concurrent_append_from_stale_cache() {
        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
        let repo = test_repo(storage.clone());
        let id = "legacy-cas-race";
        let expected = vec![serde_json::json!({"content": "first"})];

        let mut stale_cache = Session::new(id, "m");
        stale_cache.set_pending_injected_messages(expected.clone());
        cache_put(&repo, &stale_cache);

        let mut durable = stale_cache.clone();
        durable.set_pending_injected_messages(vec![
            serde_json::json!({"content": "first"}),
            serde_json::json!({"content": "concurrent"}),
        ]);
        storage.save_session(&durable).await.unwrap();

        let cleared = bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
            &repo, id, &expected,
        )
        .await
        .unwrap();
        assert!(!cleared, "the durable compare-and-clear must reject drift");
        assert_eq!(
            storage
                .load_session(id)
                .await
                .unwrap()
                .unwrap()
                .pending_injected_messages()
                .unwrap(),
            durable.pending_injected_messages().unwrap(),
            "the concurrent durable append must remain intact"
        );
        assert_eq!(
            read_cached_session(repo.cache(), id)
                .unwrap()
                .pending_injected_messages()
                .unwrap(),
            expected,
            "a failed CAS must not mutate the existing cache"
        );
    }

    #[tokio::test]
    async fn successful_legacy_clear_refreshes_stale_cache_from_durable_state() {
        let storage: Arc<dyn Storage> = Arc::new(MapStorage::default());
        let repo = test_repo(storage.clone());
        let id = "legacy-cas-success";
        let expected = vec![serde_json::json!({"content": "first"})];

        let mut stale_cache = Session::new(id, "stale-model");
        stale_cache.set_pending_injected_messages(expected.clone());
        cache_put(&repo, &stale_cache);

        let mut durable = Session::new(id, "durable-model");
        durable.set_pending_injected_messages(expected.clone());
        durable
            .metadata
            .insert("durable-only".to_string(), "keep".to_string());
        storage.save_session(&durable).await.unwrap();

        assert!(
            bamboo_domain::RuntimeSessionPersistence::clear_legacy_pending_messages(
                &repo, id, &expected,
            )
            .await
            .unwrap()
        );
        let cached = read_cached_session(repo.cache(), id).unwrap();
        assert!(!cached.has_pending_injected_messages());
        assert_eq!(cached.model, "durable-model");
        assert_eq!(
            cached.metadata.get("durable-only").map(String::as_str),
            Some("keep")
        );
    }
}