Skip to main content

talon_cli/mcp/
state.rs

1use std::collections::HashMap;
2use std::sync::{Arc, Mutex, RwLock};
3use talon_core::TalonConfig;
4
5use crate::mcp::session::ledger::TurnLedger;
6
7/// Process-local state shared across all MCP request handlers in a single
8/// `talon mcp` process lifetime.
9#[derive(Debug)]
10pub struct McpServerState {
11    pub config: Arc<ConfigState>,
12    pub sessions: Arc<RwLock<SessionStore>>,
13    pub diagnostics: Arc<DiagnosticsState>,
14}
15
16/// Resolved configuration paths and the loaded [`TalonConfig`].
17#[derive(Debug)]
18pub struct ConfigState {
19    pub config: TalonConfig,
20    pub config_path: Option<std::path::PathBuf>,
21    pub vault_path: std::path::PathBuf,
22    pub db_path: std::path::PathBuf,
23}
24
25/// Process-local session store keyed by host + session ID.
26#[derive(Debug)]
27pub struct SessionStore {
28    pub sessions: HashMap<SessionKey, SessionState>,
29}
30
31/// Composite key identifying a single agent session.
32#[derive(Debug, Clone, PartialEq, Eq, Hash)]
33pub struct SessionKey {
34    pub host: HostKind,
35    pub session_id: String,
36}
37
38/// Identifies the MCP host that opened a session.
39#[derive(Debug, Clone, PartialEq, Eq, Hash)]
40pub enum HostKind {
41    ClaudeCode,
42    Hermes,
43    Unknown(String),
44}
45
46/// Per-session runtime state.
47#[derive(Debug)]
48pub struct SessionState {
49    pub created_at_ms: i64,
50    pub last_seen_at_ms: i64,
51    /// Turn history with suppression ledger for recall deduplication.
52    pub ledger: TurnLedger,
53    /// Per-turn score decay multiplier for suppression (default `DEFAULT_DECAY`).
54    pub suppression_decay: f64,
55}
56
57/// Lightweight diagnostics visible to health-check tooling.
58#[derive(Debug)]
59pub struct DiagnosticsState {
60    pub watcher_running: std::sync::atomic::AtomicBool,
61    pub last_refresh_error: Mutex<Option<String>>,
62    pub last_embed_error: Mutex<Option<String>>,
63}
64
65impl McpServerState {
66    /// Creates a new [`McpServerState`] with empty sessions and default
67    /// diagnostics, wrapping it in an [`Arc`] for shared ownership.
68    #[must_use]
69    pub fn new(config: ConfigState) -> Arc<Self> {
70        Arc::new(Self {
71            config: Arc::new(config),
72            sessions: Arc::new(RwLock::new(SessionStore {
73                sessions: HashMap::new(),
74            })),
75            diagnostics: Arc::new(DiagnosticsState {
76                watcher_running: std::sync::atomic::AtomicBool::new(false),
77                last_refresh_error: Mutex::new(None),
78                last_embed_error: Mutex::new(None),
79            }),
80        })
81    }
82}
83
84#[cfg(test)]
85mod tests {
86    use super::{HostKind, McpServerState, SessionKey};
87    use crate::mcp::state::ConfigState;
88    use std::path::PathBuf;
89
90    fn stub_config_state() -> ConfigState {
91        let vault_path = PathBuf::from("/tmp/vault");
92        let db_path = PathBuf::from("/tmp/vault.db");
93        let config = crate::config::default_config_for_vault(vault_path.clone());
94        ConfigState {
95            config,
96            config_path: None,
97            vault_path,
98            db_path,
99        }
100    }
101
102    #[test]
103    fn mcp_server_state_new_creates_empty_session_store() {
104        let state = McpServerState::new(stub_config_state());
105        let is_empty = state
106            .sessions
107            .read()
108            .unwrap_or_else(std::sync::PoisonError::into_inner)
109            .sessions
110            .is_empty();
111        assert!(
112            is_empty,
113            "expected empty session store after McpServerState::new"
114        );
115    }
116
117    #[test]
118    fn session_key_equality() {
119        let a = SessionKey {
120            host: HostKind::ClaudeCode,
121            session_id: "abc".to_string(),
122        };
123        let b = SessionKey {
124            host: HostKind::ClaudeCode,
125            session_id: "abc".to_string(),
126        };
127        assert_eq!(a, b);
128    }
129}