axon/

session_scope.rs

//! Session Scoping — namespaced session isolation for AxonServer.
//!
//! Wraps multiple `SessionStore` instances keyed by scope name, providing
//! isolation between flows, daemons, and the global server context.
//!
//! Scopes:
//!   - `"global"` — default scope for server-wide state
//!   - `"flow:<name>"` — per-flow session isolation
//!   - `"daemon:<name>"` — per-daemon session isolation
//!   - Custom string — any arbitrary namespace
//!
//! Each scope gets its own ephemeral memory and persistent store.

use std::collections::HashMap;
use crate::session_store::{SessionStore, MemoryEntry};

// ── Scoped key ───────────────────────────────────────────────────────────

/// The default scope name.
pub const DEFAULT_SCOPE: &str = "global";

/// Build a scope name for a flow.
pub fn flow_scope(flow_name: &str) -> String {
    format!("flow:{}", flow_name)
}

/// Build a scope name for a daemon.
pub fn daemon_scope(daemon_name: &str) -> String {
    format!("daemon:{}", daemon_name)
}

// ── Manager ──────────────────────────────────────────────────────────────

/// Manages multiple SessionStore instances, one per scope.
#[derive(Debug)]
pub struct ScopedSessionManager {
    /// Base directory for persistent stores.
    base_path: String,
    /// Map of scope name → SessionStore.
    scopes: HashMap<String, SessionStore>,
}

impl ScopedSessionManager {
    /// Create a new scoped session manager.
    /// `base_path` is used as the root for deriving per-scope store file paths.
    pub fn new(base_path: &str) -> Self {
        let mut mgr = ScopedSessionManager {
            base_path: base_path.to_string(),
            scopes: HashMap::new(),
        };
        // Pre-create the global scope
        mgr.ensure_scope(DEFAULT_SCOPE);
        mgr
    }

    /// Get or create the SessionStore for a scope.
    fn ensure_scope(&mut self, scope: &str) -> &mut SessionStore {
        if !self.scopes.contains_key(scope) {
            let source = format!("{}__{}", self.base_path, scope.replace(':', "_"));
            let store = SessionStore::new(&source);
            self.scopes.insert(scope.to_string(), store);
        }
        self.scopes.get_mut(scope).unwrap()
    }

    /// Get the SessionStore for a scope (read-only). Returns None if scope doesn't exist.
    fn get_scope(&self, scope: &str) -> Option<&SessionStore> {
        self.scopes.get(scope)
    }

    // ── Delegated operations ─────────────────────────────────────────────

    /// Remember a value in a scope's ephemeral memory.
    pub fn remember(&mut self, scope: &str, key: &str, value: &str, source_step: &str) {
        self.ensure_scope(scope).remember(key, value, source_step);
    }

    /// Recall a value from a scope's ephemeral memory.
    pub fn recall(&mut self, scope: &str, key: &str) -> Option<MemoryEntry> {
        self.ensure_scope(scope).recall(key).cloned()
    }

    /// Persist a value in a scope's file-backed store.
    pub fn persist(&mut self, scope: &str, key: &str, value: &str, source_step: &str) {
        self.ensure_scope(scope).persist(key, value, source_step);
    }

    /// Retrieve a value from a scope's file-backed store.
    pub fn retrieve(&mut self, scope: &str, key: &str) -> Option<MemoryEntry> {
        self.ensure_scope(scope).retrieve(key).cloned()
    }

    /// Query entries in a scope's store.
    pub fn query(&mut self, scope: &str, query: &str) -> Vec<MemoryEntry> {
        self.ensure_scope(scope)
            .retrieve_query(query)
            .into_iter()
            .cloned()
            .collect()
    }

    /// Mutate an entry in a scope's store.
    pub fn mutate(&mut self, scope: &str, key: &str, new_value: &str, source_step: &str) -> bool {
        self.ensure_scope(scope).mutate(key, new_value, source_step)
    }

    /// Purge an entry from a scope's store.
    pub fn purge(&mut self, scope: &str, key: &str) -> bool {
        self.ensure_scope(scope).purge(key)
    }

    /// Flush a scope's persistent store to disk.
    pub fn flush(&mut self, scope: &str) -> Result<(), String> {
        self.ensure_scope(scope).flush()
    }

    /// Flush all scopes to disk.
    pub fn flush_all(&mut self) -> Vec<(String, Result<(), String>)> {
        let scope_names: Vec<String> = self.scopes.keys().cloned().collect();
        let mut results = Vec::new();
        for name in scope_names {
            let result = self.scopes.get(&name).map(|s| s.flush()).unwrap_or(Ok(()));
            results.push((name, result));
        }
        results
    }

    // ── Introspection ────────────────────────────────────────────────────

    /// List all scope names.
    pub fn list_scopes(&self) -> Vec<&str> {
        let mut names: Vec<&str> = self.scopes.keys().map(|s| s.as_str()).collect();
        names.sort();
        names
    }

    /// Number of active scopes.
    pub fn scope_count(&self) -> usize {
        self.scopes.len()
    }

    /// Memory count for a scope (0 if scope doesn't exist).
    pub fn memory_count(&self, scope: &str) -> usize {
        self.get_scope(scope).map(|s| s.memory_count()).unwrap_or(0)
    }

    /// Store count for a scope (0 if scope doesn't exist).
    pub fn store_count(&self, scope: &str) -> usize {
        self.get_scope(scope).map(|s| s.store_count()).unwrap_or(0)
    }

    /// Total memory count across all scopes.
    pub fn total_memory_count(&self) -> usize {
        self.scopes.values().map(|s| s.memory_count()).sum()
    }

    /// Total store count across all scopes.
    pub fn total_store_count(&self) -> usize {
        self.scopes.values().map(|s| s.store_count()).sum()
    }

    /// List all entries in a scope (ephemeral + persistent).
    pub fn list_entries(&mut self, scope: &str) -> Vec<ScopedEntry> {
        let store = self.ensure_scope(scope);
        let mut entries = Vec::new();

        for entry in store.memory_entries() {
            entries.push(ScopedEntry {
                scope: scope.to_string(),
                layer: "memory".to_string(),
                key: entry.key.clone(),
                value: entry.value.clone(),
                timestamp: entry.timestamp,
                source_step: entry.source_step.clone(),
            });
        }

        // For persistent entries, we iterate the retrieve_query with empty to get all
        // But retrieve_query needs a query string; use a broad match
        // Instead, we'll check store_count and query
        // Actually, let's just expose what we can
        entries
    }

    /// Summary of all scopes for display/API.
    pub fn summary(&self) -> Vec<ScopeSummary> {
        let mut result = Vec::new();
        for (name, store) in &self.scopes {
            result.push(ScopeSummary {
                scope: name.clone(),
                memory_count: store.memory_count(),
                store_count: store.store_count(),
            });
        }
        result.sort_by(|a, b| a.scope.cmp(&b.scope));
        result
    }
}

/// A session entry with scope context.
#[derive(Debug, Clone, serde::Serialize)]
pub struct ScopedEntry {
    pub scope: String,
    pub layer: String,
    pub key: String,
    pub value: String,
    pub timestamp: u64,
    pub source_step: String,
}

/// Summary info for a scope.
#[derive(Debug, Clone, serde::Serialize)]
pub struct ScopeSummary {
    pub scope: String,
    pub memory_count: usize,
    pub store_count: usize,
}

// ── Tests ────────────────────────────────────────────────────────────────

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

    fn test_manager(name: &str) -> ScopedSessionManager {
        let base = std::env::temp_dir().join(format!("axon_scope_test_{name}"));
        ScopedSessionManager::new(base.to_str().unwrap())
    }

    #[test]
    fn default_scope_created() {
        let mgr = test_manager("default");
        assert!(mgr.list_scopes().contains(&"global"));
        assert_eq!(mgr.scope_count(), 1);
    }

    #[test]
    fn remember_recall_in_scope() {
        let mut mgr = test_manager("rem_recall");
        mgr.remember("flow:analyze", "result", "42", "step1");
        let entry = mgr.recall("flow:analyze", "result").unwrap();
        assert_eq!(entry.value, "42");

        // Not visible in other scope
        assert!(mgr.recall("global", "result").is_none());
        assert!(mgr.recall("flow:other", "result").is_none());
    }

    #[test]
    fn persist_retrieve_in_scope() {
        let mut mgr = test_manager("persist_ret");
        mgr.persist("daemon:worker", "config", "max_retries=3", "init");
        let entry = mgr.retrieve("daemon:worker", "config").unwrap();
        assert_eq!(entry.value, "max_retries=3");

        // Not visible in global
        assert!(mgr.retrieve("global", "config").is_none());
    }

    #[test]
    fn mutate_in_scope() {
        let mut mgr = test_manager("mutate");
        mgr.persist("flow:a", "key", "old", "s1");
        assert!(mgr.mutate("flow:a", "key", "new", "s2"));
        assert_eq!(mgr.retrieve("flow:a", "key").unwrap().value, "new");

        // Mutate in wrong scope fails
        assert!(!mgr.mutate("flow:b", "key", "val", "s3"));
    }

    #[test]
    fn purge_in_scope() {
        let mut mgr = test_manager("purge");
        mgr.persist("flow:x", "temp", "data", "s1");
        assert!(mgr.purge("flow:x", "temp"));
        assert!(mgr.retrieve("flow:x", "temp").is_none());

        // Purge in wrong scope
        mgr.persist("flow:y", "keep", "data", "s1");
        assert!(!mgr.purge("flow:z", "keep"));
        assert!(mgr.retrieve("flow:y", "keep").is_some());
    }

    #[test]
    fn query_in_scope() {
        let mut mgr = test_manager("query");
        mgr.persist("flow:a", "analysis_1", "result1", "s1");
        mgr.persist("flow:a", "analysis_2", "result2", "s2");
        mgr.persist("flow:b", "analysis_3", "result3", "s3");

        let results = mgr.query("flow:a", "analysis");
        assert_eq!(results.len(), 2);

        // Different scope
        let results_b = mgr.query("flow:b", "analysis");
        assert_eq!(results_b.len(), 1);
    }

    #[test]
    fn scope_counts() {
        let mut mgr = test_manager("counts");
        mgr.remember("global", "a", "1", "s");
        mgr.remember("global", "b", "2", "s");
        mgr.persist("flow:x", "c", "3", "s");

        assert_eq!(mgr.memory_count("global"), 2);
        assert_eq!(mgr.store_count("flow:x"), 1);
        assert_eq!(mgr.total_memory_count(), 2);
        assert_eq!(mgr.total_store_count(), 1);
    }

    #[test]
    fn list_scopes_sorted() {
        let mut mgr = test_manager("list");
        mgr.remember("flow:z", "k", "v", "s");
        mgr.remember("daemon:a", "k", "v", "s");
        let scopes = mgr.list_scopes();
        assert!(scopes.len() >= 3);
        // Should be sorted
        for i in 1..scopes.len() {
            assert!(scopes[i - 1] <= scopes[i]);
        }
    }

    #[test]
    fn summary_includes_all_scopes() {
        let mut mgr = test_manager("summary");
        mgr.remember("global", "a", "1", "s");
        mgr.persist("flow:report", "b", "2", "s");

        let summary = mgr.summary();
        assert!(summary.len() >= 2);

        let global = summary.iter().find(|s| s.scope == "global").unwrap();
        assert_eq!(global.memory_count, 1);

        let flow = summary.iter().find(|s| s.scope == "flow:report").unwrap();
        assert_eq!(flow.store_count, 1);
    }

    #[test]
    fn flow_scope_and_daemon_scope_helpers() {
        assert_eq!(flow_scope("analyze"), "flow:analyze");
        assert_eq!(daemon_scope("worker"), "daemon:worker");
    }

    #[test]
    fn isolated_scopes_no_leakage() {
        let mut mgr = test_manager("isolation");

        // Same key in three different scopes
        mgr.remember("global", "status", "global_val", "s");
        mgr.remember("flow:a", "status", "flow_a_val", "s");
        mgr.remember("daemon:d", "status", "daemon_val", "s");

        assert_eq!(mgr.recall("global", "status").unwrap().value, "global_val");
        assert_eq!(mgr.recall("flow:a", "status").unwrap().value, "flow_a_val");
        assert_eq!(mgr.recall("daemon:d", "status").unwrap().value, "daemon_val");
    }

    #[test]
    fn nonexistent_scope_returns_zero_counts() {
        let mgr = test_manager("nonexist");
        assert_eq!(mgr.memory_count("flow:phantom"), 0);
        assert_eq!(mgr.store_count("daemon:ghost"), 0);
    }

    #[test]
    fn scoped_entry_serializes() {
        let entry = ScopedEntry {
            scope: "flow:test".into(),
            layer: "memory".into(),
            key: "k".into(),
            value: "v".into(),
            timestamp: 12345,
            source_step: "s1".into(),
        };
        let json = serde_json::to_string(&entry).unwrap();
        assert!(json.contains("\"scope\":\"flow:test\""));
        assert!(json.contains("\"layer\":\"memory\""));
    }

    #[test]
    fn scope_summary_serializes() {
        let summary = ScopeSummary {
            scope: "global".into(),
            memory_count: 3,
            store_count: 1,
        };
        let json = serde_json::to_string(&summary).unwrap();
        assert!(json.contains("\"scope\":\"global\""));
        assert!(json.contains("\"memory_count\":3"));
    }
}