turbomcp-protocol 3.1.3

//! State management utilities for MCP servers

use crate::{Error, ErrorKind, Result};
use parking_lot::RwLock;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;

/// Thread-safe state manager for MCP servers.
///
/// Uses `parking_lot::RwLock` for poison-free locking — a panic in a writer
/// no longer silently disables `set` / `clear` for every subsequent caller.
#[derive(Debug, Clone)]
pub struct StateManager {
    /// Internal state storage
    state: Arc<RwLock<HashMap<String, Value>>>,
}

impl StateManager {
    /// Create a new state manager
    #[must_use]
    pub fn new() -> Self {
        Self {
            state: Arc::new(RwLock::new(HashMap::new())),
        }
    }

    /// Set a value in the state
    pub fn set(&self, key: String, value: Value) {
        self.state.write().insert(key, value);
    }

    /// Get a value from the state
    #[must_use]
    pub fn get(&self, key: &str) -> Option<Value> {
        self.state.read().get(key).cloned()
    }

    /// Remove a value from the state
    #[must_use]
    pub fn remove(&self, key: &str) -> Option<Value> {
        self.state.write().remove(key)
    }

    /// Check if a key exists in the state
    #[must_use]
    pub fn contains(&self, key: &str) -> bool {
        self.state.read().contains_key(key)
    }

    /// Get the number of entries in the state
    #[must_use]
    pub fn size(&self) -> usize {
        self.state.read().len()
    }

    /// List all keys in the state
    #[must_use]
    pub fn list_keys(&self) -> Vec<String> {
        self.state.read().keys().cloned().collect()
    }

    /// Clear all entries from the state
    pub fn clear(&self) {
        self.state.write().clear();
    }

    /// Export the state as JSON
    #[must_use]
    pub fn export(&self) -> Value {
        let state = self.state.read();
        Value::Object(state.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
    }

    /// Import state from JSON
    pub fn import(&self, data: Value) -> Result<()> {
        match data {
            Value::Object(obj) => {
                let mut state = self.state.write();
                for (key, value) in obj {
                    state.insert(key, value);
                }
                Ok(())
            }
            _ => Err(Error::new(
                ErrorKind::Configuration,
                "Import data must be a JSON object",
            )),
        }
    }
}

impl Default for StateManager {
    fn default() -> Self {
        Self::new()
    }
}

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

    #[test]
    fn test_state_operations() {
        let state = StateManager::new();

        // Test set and get
        state.set("key1".to_string(), json!("value1"));
        assert_eq!(state.get("key1"), Some(json!("value1")));

        // Test contains
        assert!(state.contains("key1"));
        assert!(!state.contains("key2"));

        // Test size
        assert_eq!(state.size(), 1);

        // Test remove
        assert_eq!(state.remove("key1"), Some(json!("value1")));
        assert!(!state.contains("key1"));
        assert_eq!(state.size(), 0);
    }

    #[test]
    fn test_export_import() {
        let state1 = StateManager::new();
        state1.set("key1".to_string(), json!("value1"));
        state1.set("key2".to_string(), json!(42));

        let exported = state1.export();

        let state2 = StateManager::new();
        assert!(state2.import(exported).is_ok());

        assert_eq!(state2.get("key1"), Some(json!("value1")));
        assert_eq!(state2.get("key2"), Some(json!(42)));
        assert_eq!(state2.size(), 2);
    }

    #[test]
    fn test_list_keys() {
        let state = StateManager::new();
        state.set("a".to_string(), json!(1));
        state.set("b".to_string(), json!(2));
        state.set("c".to_string(), json!(3));

        let mut keys = state.list_keys();
        keys.sort();
        assert_eq!(keys, vec!["a", "b", "c"]);
    }

    #[test]
    fn test_clear() {
        let state = StateManager::new();
        state.set("key1".to_string(), json!("value1"));
        state.set("key2".to_string(), json!("value2"));

        assert_eq!(state.size(), 2);
        state.clear();
        assert_eq!(state.size(), 0);
        assert!(!state.contains("key1"));
        assert!(!state.contains("key2"));
    }
}