turbomcp_core/
state.rs

1//! State management utilities for MCP servers
2
3use crate::{Error, ErrorKind, Result};
4use serde_json::Value;
5use std::collections::HashMap;
6use std::sync::{Arc, RwLock};
7
8/// Thread-safe state manager for MCP servers
9#[derive(Debug, Clone)]
10pub struct StateManager {
11    /// Internal state storage
12    state: Arc<RwLock<HashMap<String, Value>>>,
13}
14
15impl StateManager {
16    /// Create a new state manager
17    #[must_use]
18    pub fn new() -> Self {
19        Self {
20            state: Arc::new(RwLock::new(HashMap::new())),
21        }
22    }
23
24    /// Set a value in the state
25    pub fn set(&self, key: String, value: Value) {
26        if let Ok(mut state) = self.state.write() {
27            state.insert(key, value);
28        }
29    }
30
31    /// Get a value from the state
32    #[must_use]
33    pub fn get(&self, key: &str) -> Option<Value> {
34        self.state.read().ok()?.get(key).cloned()
35    }
36
37    /// Remove a value from the state
38    #[must_use]
39    pub fn remove(&self, key: &str) -> Option<Value> {
40        self.state.write().ok()?.remove(key)
41    }
42
43    /// Check if a key exists in the state
44    #[must_use]
45    pub fn contains(&self, key: &str) -> bool {
46        self.state.read().is_ok_and(|state| state.contains_key(key))
47    }
48
49    /// Get the number of entries in the state
50    #[must_use]
51    pub fn size(&self) -> usize {
52        self.state.read().map_or(0, |state| state.len())
53    }
54
55    /// List all keys in the state
56    #[must_use]
57    pub fn list_keys(&self) -> Vec<String> {
58        self.state
59            .read()
60            .map_or_else(|_| Vec::new(), |state| state.keys().cloned().collect())
61    }
62
63    /// Clear all entries from the state
64    pub fn clear(&self) {
65        if let Ok(mut state) = self.state.write() {
66            state.clear();
67        }
68    }
69
70    /// Export the state as JSON
71    #[must_use]
72    pub fn export(&self) -> Value {
73        self.state.read().map_or_else(
74            |_| Value::Object(serde_json::Map::new()),
75            |state| Value::Object(state.iter().map(|(k, v)| (k.clone(), v.clone())).collect()),
76        )
77    }
78
79    /// Import state from JSON
80    pub fn import(&self, data: Value) -> Result<()> {
81        match data {
82            Value::Object(obj) => self.state.write().map_or_else(
83                |_| {
84                    Err(Error::new(
85                        ErrorKind::Internal,
86                        "Failed to acquire write lock",
87                    ))
88                },
89                |mut state| {
90                    for (key, value) in obj {
91                        state.insert(key, value);
92                    }
93                    Ok(())
94                },
95            ),
96            _ => Err(Error::new(
97                ErrorKind::Configuration,
98                "Import data must be a JSON object",
99            )),
100        }
101    }
102}
103
104impl Default for StateManager {
105    fn default() -> Self {
106        Self::new()
107    }
108}
109
110#[cfg(test)]
111mod tests {
112    use super::*;
113    use serde_json::json;
114
115    #[test]
116    fn test_state_operations() {
117        let state = StateManager::new();
118
119        // Test set and get
120        state.set("key1".to_string(), json!("value1"));
121        assert_eq!(state.get("key1"), Some(json!("value1")));
122
123        // Test contains
124        assert!(state.contains("key1"));
125        assert!(!state.contains("key2"));
126
127        // Test size
128        assert_eq!(state.size(), 1);
129
130        // Test remove
131        assert_eq!(state.remove("key1"), Some(json!("value1")));
132        assert!(!state.contains("key1"));
133        assert_eq!(state.size(), 0);
134    }
135
136    #[test]
137    fn test_export_import() {
138        let state1 = StateManager::new();
139        state1.set("key1".to_string(), json!("value1"));
140        state1.set("key2".to_string(), json!(42));
141
142        let exported = state1.export();
143
144        let state2 = StateManager::new();
145        assert!(state2.import(exported).is_ok());
146
147        assert_eq!(state2.get("key1"), Some(json!("value1")));
148        assert_eq!(state2.get("key2"), Some(json!(42)));
149        assert_eq!(state2.size(), 2);
150    }
151
152    #[test]
153    fn test_list_keys() {
154        let state = StateManager::new();
155        state.set("a".to_string(), json!(1));
156        state.set("b".to_string(), json!(2));
157        state.set("c".to_string(), json!(3));
158
159        let mut keys = state.list_keys();
160        keys.sort();
161        assert_eq!(keys, vec!["a", "b", "c"]);
162    }
163
164    #[test]
165    fn test_clear() {
166        let state = StateManager::new();
167        state.set("key1".to_string(), json!("value1"));
168        state.set("key2".to_string(), json!("value2"));
169
170        assert_eq!(state.size(), 2);
171        state.clear();
172        assert_eq!(state.size(), 0);
173        assert!(!state.contains("key1"));
174        assert!(!state.contains("key2"));
175    }
176}
turbomcp_core/state.rs

turbomcp_core/
state.rs
