agentic_codebase/mcp/
tenant.rs

1//! Multi-tenant session registry — lazy-loads per-user graph files.
2
3use std::collections::HashMap;
4use std::path::{Path, PathBuf};
5use std::sync::Arc;
6
7use super::server::McpServer;
8use crate::AcbResult;
9
10/// Mutex type selected by feature: `tokio::sync::Mutex` for SSE,
11/// `std::sync::Mutex` when tokio is unavailable.
12#[cfg(feature = "sse")]
13type MutexType<T> = tokio::sync::Mutex<T>;
14#[cfg(not(feature = "sse"))]
15type MutexType<T> = std::sync::Mutex<T>;
16
17/// Registry of per-user MCP server instances for multi-tenant mode.
18pub struct TenantRegistry {
19    data_dir: PathBuf,
20    servers: HashMap<String, Arc<MutexType<McpServer>>>,
21}
22
23impl TenantRegistry {
24    /// Create a new tenant registry backed by the given data directory.
25    pub fn new(data_dir: &Path) -> Self {
26        Self {
27            data_dir: data_dir.to_path_buf(),
28            servers: HashMap::new(),
29        }
30    }
31
32    /// Get or create an MCP server for the given user ID.
33    ///
34    /// On first access, loads `{data_dir}/{user_id}.acb` and creates
35    /// a server with the graph pre-loaded.
36    pub fn get_or_create(&mut self, user_id: &str) -> AcbResult<Arc<MutexType<McpServer>>> {
37        if let Some(server) = self.servers.get(user_id) {
38            return Ok(server.clone());
39        }
40
41        // Ensure data directory exists
42        std::fs::create_dir_all(&self.data_dir).map_err(|e| {
43            crate::AcbError::Io(std::io::Error::other(format!(
44                "Failed to create data dir {}: {e}",
45                self.data_dir.display()
46            )))
47        })?;
48
49        let graph_path = self.data_dir.join(format!("{user_id}.acb"));
50
51        tracing::info!(
52            "Opening graph for user '{user_id}': {}",
53            graph_path.display()
54        );
55
56        let mut server = McpServer::new();
57
58        // If the graph file exists, load it
59        if graph_path.is_file() {
60            match crate::AcbReader::read_from_file(&graph_path) {
61                Ok(graph) => {
62                    server.load_graph(user_id.to_string(), graph);
63                }
64                Err(e) => {
65                    tracing::warn!(
66                        "Failed to load graph for user '{user_id}': {e} — starting empty"
67                    );
68                }
69            }
70        }
71
72        let server = Arc::new(MutexType::new(server));
73        self.servers.insert(user_id.to_string(), server.clone());
74
75        Ok(server)
76    }
77
78    /// Number of active tenant sessions.
79    pub fn count(&self) -> usize {
80        self.servers.len()
81    }
82}
agentic_codebase/mcp/tenant.rs

agentic_codebase/mcp/
tenant.rs
