lc_langgraph/

persistence.rs

// crates/lc-langgraph/src/persistence.rs
//! Graph persistence for serialization and storage
//!
//! This module provides persistence capabilities for graph definitions,
//! allowing graphs to be saved, loaded, and shared across sessions.

use async_trait::async_trait;
use chrono::{DateTime, Utc};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use tokio::sync::Mutex;
use uuid::Uuid;

/// GraphPersistence trait for storing and loading graph definitions
#[async_trait]
pub trait GraphPersistence: Send + Sync {
    /// Save a graph definition with the given ID
    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError>;

    /// Load a graph definition by ID
    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError>;
    /// Delete a graph definition by ID
    async fn delete(&self, id: &str) -> Result<(), PersistenceError>;

    /// Check if a graph definition exists
    async fn exists(&self, id: &str) -> Result<bool, PersistenceError>;

    /// List all stored graph IDs
    async fn list(&self) -> Result<Vec<String>, PersistenceError>;
}

/// Persistence error types
#[derive(Debug, thiserror::Error)]
#[non_exhaustive]
pub enum PersistenceError {
    /// The requested graph was not found.
    #[error("Graph '{0}' not found")]
    NotFound(String),

    /// Failed to serialize a graph definition.
    #[error("Serialization error: {0}")]
    SerializationError(String),

    /// Failed to deserialize a graph definition.
    #[error("Deserialization error: {0}")]
    DeserializationError(String),

    /// An I/O error occurred while accessing storage.
    #[error("IO error: {0}")]
    IoError(String),

    /// The graph definition is invalid.
    #[error("Invalid graph definition: {0}")]
    InvalidDefinition(String),

    /// A MongoDB operation failed.
    #[error("MongoDB error: {0}")]
    MongoError(String),

    /// Failed to connect to the persistence backend.
    #[error("Connection error: {0}")]
    ConnectionError(String),
}

/// GraphDefinition - Serializable graph structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct GraphDefinition {
    /// Unique identifier
    pub id: String,

    /// Human-readable name
    pub name: Option<String>,

    /// Entry point node name
    pub entry_point: String,

    /// Node definitions
    pub nodes: Vec<NodeDefinition>,

    /// Edge definitions
    pub edges: Vec<EdgeDefinition>,

    /// Router definitions
    pub routers: Vec<RouterDefinition>,

    /// Maximum recursion limit
    pub recursion_limit: usize,

    /// Creation timestamp
    pub created_at: DateTime<Utc>,

    /// Last update timestamp
    pub updated_at: DateTime<Utc>,

    /// Custom metadata
    pub metadata: HashMap<String, serde_json::Value>,
}

impl GraphDefinition {
    /// Create a new graph definition with the given entry point
    pub fn new(entry_point: impl Into<String>) -> Self {
        let now = Utc::now();
        Self {
            id: Uuid::new_v4().to_string(),
            name: None,
            entry_point: entry_point.into(),
            nodes: Vec::new(),
            edges: Vec::new(),
            routers: Vec::new(),
            recursion_limit: 25,
            created_at: now,
            updated_at: now,
            metadata: HashMap::new(),
        }
    }

    /// Set a custom ID
    pub fn with_id(mut self, id: impl Into<String>) -> Self {
        self.id = id.into();
        self
    }

    /// Set a human-readable name
    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = Some(name.into());
        self
    }

    /// Set recursion limit
    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
        self.recursion_limit = limit;
        self
    }

    /// Add a node definition
    pub fn add_node(&mut self, node: NodeDefinition) {
        self.nodes.push(node);
        self.updated_at = Utc::now();
    }

    /// Add an edge definition
    pub fn add_edge(&mut self, edge: EdgeDefinition) {
        self.edges.push(edge);
        self.updated_at = Utc::now();
    }

    /// Add a router definition
    pub fn add_router(&mut self, router: RouterDefinition) {
        self.routers.push(router);
        self.updated_at = Utc::now();
    }
}

/// NodeDefinition - Serializable node structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct NodeDefinition {
    /// Node name
    pub name: String,

    /// Node type
    pub node_type: NodeType,

    /// Custom configuration
    pub config: serde_json::Value,
}

/// NodeType - Type of node execution
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum NodeType {
    /// Synchronous node.
    Sync,
    /// Asynchronous node.
    Async,
    /// A subgraph node.
    Subgraph,
    /// A custom node type.
    Custom,
}

/// EdgeDefinition - Serializable edge structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct EdgeDefinition {
    /// Edge type
    pub edge_type: EdgeType,

    /// Source node name
    pub source: String,

    /// Target node name (for fixed edges)
    pub target: Option<String>,

    /// Multiple targets (for fan-out edges)
    pub targets: Option<Vec<String>>,

    /// Router name (for conditional edges)
    pub router_name: Option<String>,

    /// Conditional targets mapping (route -> target)
    pub conditional_targets: Option<HashMap<String, String>>,

    /// Default target for conditional edges
    pub default_target: Option<String>,

    /// Source nodes for fan-in edges
    pub sources: Option<Vec<String>>,
}

impl EdgeDefinition {
    /// Create a fixed edge
    pub fn fixed(source: impl Into<String>, target: impl Into<String>) -> Self {
        Self {
            edge_type: EdgeType::Fixed,
            source: source.into(),
            target: Some(target.into()),
            targets: None,
            router_name: None,
            conditional_targets: None,
            default_target: None,
            sources: None,
        }
    }

    /// Create a conditional edge
    pub fn conditional(
        source: impl Into<String>,
        router_name: impl Into<String>,
        targets: HashMap<String, String>,
        default_target: Option<String>,
    ) -> Self {
        Self {
            edge_type: EdgeType::Conditional,
            source: source.into(),
            target: None,
            targets: None,
            router_name: Some(router_name.into()),
            conditional_targets: Some(targets),
            default_target,
            sources: None,
        }
    }

    /// Create a fan-out edge (parallel execution)
    pub fn fan_out(source: impl Into<String>, targets: Vec<String>) -> Self {
        Self {
            edge_type: EdgeType::FanOut,
            source: source.into(),
            target: None,
            targets: Some(targets),
            router_name: None,
            conditional_targets: None,
            default_target: None,
            sources: None,
        }
    }

    /// Create a fan-in edge (merge from parallel branches)
    pub fn fan_in(sources: Vec<String>, target: impl Into<String>) -> Self {
        Self {
            edge_type: EdgeType::FanIn,
            source: "__fan_in__".to_string(),
            target: Some(target.into()),
            targets: None,
            router_name: None,
            conditional_targets: None,
            default_target: None,
            sources: Some(sources),
        }
    }
}

/// EdgeType - Type of edge connection
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub enum EdgeType {
    /// Fixed transition
    Fixed,

    /// Conditional routing
    Conditional,

    /// Parallel fan-out
    FanOut,

    /// Merge fan-in
    FanIn,
}

/// RouterDefinition - Serializable router structure
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RouterDefinition {
    /// Router name
    pub name: String,

    /// Router type (e.g., "function", "state_key")
    pub router_type: String,

    /// Possible routes
    pub routes: Vec<String>,

    /// Custom configuration
    pub config: serde_json::Value,
}

/// MemoryPersistence - In-memory graph storage
pub struct MemoryPersistence {
    graphs: Mutex<HashMap<String, GraphDefinition>>,
}

impl MemoryPersistence {
    /// Create a new memory persistence instance
    pub fn new() -> Self {
        Self {
            graphs: Mutex::new(HashMap::new()),
        }
    }
}

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

#[async_trait]
impl GraphPersistence for MemoryPersistence {
    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError> {
        let mut graphs = self.graphs.lock().await;
        graphs.insert(id.to_string(), definition.clone());
        Ok(())
    }

    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
        let graphs = self.graphs.lock().await;
        graphs
            .get(id)
            .cloned()
            .ok_or_else(|| PersistenceError::NotFound(id.to_string()))
    }

    async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
        let mut graphs = self.graphs.lock().await;
        graphs
            .remove(id)
            .map(|_| ())
            .ok_or_else(|| PersistenceError::NotFound(id.to_string()))
    }

    async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
        let graphs = self.graphs.lock().await;
        Ok(graphs.contains_key(id))
    }

    async fn list(&self) -> Result<Vec<String>, PersistenceError> {
        let graphs = self.graphs.lock().await;
        Ok(graphs.keys().cloned().collect())
    }
}

/// FilePersistence - File-based graph storage
pub struct FilePersistence {
    directory: PathBuf,
}

impl FilePersistence {
    /// Create a new file persistence instance
    pub fn new(directory: impl Into<PathBuf>) -> Result<Self, PersistenceError> {
        let dir = directory.into();
        if !dir.exists() {
            std::fs::create_dir_all(&dir).map_err(|e| {
                PersistenceError::IoError(format!(
                    "Failed to create directory '{}': {}",
                    dir.display(),
                    e
                ))
            })?;
        }
        Ok(Self { directory: dir })
    }

    fn graph_path(&self, id: &str) -> Result<PathBuf, PersistenceError> {
        // Sanitize id to prevent path traversal: reject ".." and absolute paths
        if id.contains("..") || id.contains('/') || id.contains('\\') {
            return Err(PersistenceError::InvalidDefinition(format!(
                "Invalid graph id '{}': path traversal detected",
                id
            )));
        }
        if std::path::Path::new(id).is_absolute() {
            return Err(PersistenceError::InvalidDefinition(format!(
                "Invalid graph id '{}': absolute path not allowed",
                id
            )));
        }
        Ok(self.directory.join(format!("{}.json", id)))
    }
}

// NOTE: `Default` is intentionally NOT implemented for `FilePersistence` (Q1).
// The default constructor would have to create a directory, which is I/O that can
// fail (read-only cwd, disk full, permissions) — `Default` cannot report that
// failure, so it would have to panic. Use `FilePersistence::new(...)` which
// returns a `Result` and surfaces the error instead.

#[async_trait]
impl GraphPersistence for FilePersistence {
    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError> {
        let path = self.graph_path(id)?;

        let json = serde_json::to_string_pretty(definition)
            .map_err(|e| PersistenceError::SerializationError(e.to_string()))?;

        tokio::fs::write(&path, json)
            .await
            .map_err(|e| PersistenceError::IoError(e.to_string()))?;

        Ok(())
    }

    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
        let path = self.graph_path(id)?;

        if !path.exists() {
            return Err(PersistenceError::NotFound(id.to_string()));
        }

        let json = tokio::fs::read_to_string(&path)
            .await
            .map_err(|e| PersistenceError::IoError(e.to_string()))?;

        let definition: GraphDefinition = serde_json::from_str(&json)
            .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;

        Ok(definition)
    }

    async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
        let path = self.graph_path(id)?;

        if path.exists() {
            tokio::fs::remove_file(&path)
                .await
                .map_err(|e| PersistenceError::IoError(e.to_string()))?;
        }

        Ok(())
    }

    async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
        let path = self.graph_path(id)?;
        Ok(path.exists())
    }

    async fn list(&self) -> Result<Vec<String>, PersistenceError> {
        let mut ids = Vec::new();

        let mut entries = tokio::fs::read_dir(&self.directory)
            .await
            .map_err(|e| PersistenceError::IoError(e.to_string()))?;

        while let Some(entry) = entries
            .next_entry()
            .await
            .map_err(|e| PersistenceError::IoError(e.to_string()))?
        {
            let path = entry.path();
            if path.extension().is_some_and(|ext| ext == "json") {
                if let Some(id) = path.file_stem().and_then(|s| s.to_str()) {
                    ids.push(id.to_string());
                }
            }
        }

        Ok(ids)
    }
}

// ============================================================================
// MongoDB 持久化实现
// ============================================================================

#[cfg(feature = "mongodb-persistence")]
mod mongo_impl {
    use super::*;
    use mongodb::{
        bson::{doc, from_document, to_document, Document},
        options::{ClientOptions, FindOptions},
        Client, Collection,
    };

    /// MongoDB配置
    pub struct MongoConfig {
        /// MongoDB连接URI (例如: mongodb://localhost:27017 或 mongodb+srv://user:pass@cluster.mongodb.net)
        pub uri: String,

        /// 数据库名称
        pub database: String,

        /// 集合名称
        pub collection: String,
    }

    impl MongoConfig {
        /// 创建新的MongoDB配置
        pub fn new(
            uri: impl Into<String>,
            database: impl Into<String>,
            collection: impl Into<String>,
        ) -> Self {
            Self {
                uri: uri.into(),
                database: database.into(),
                collection: collection.into(),
            }
        }

        /// 从环境变量创建配置
        ///
        /// 环境变量:
        /// - MONGO_URI: MongoDB连接URI
        /// - MONGO_DATABASE: 数据库名称 (默认: langgraph)
        /// - MONGO_COLLECTION: 集合名称 (默认: graph_definitions)
        pub fn from_env() -> Result<Self, PersistenceError> {
            Ok(Self {
                uri: std::env::var("MONGO_URI").map_err(|_| {
                    PersistenceError::ConnectionError(
                        "MONGO_URI environment variable not set".to_string(),
                    )
                })?,
                database: std::env::var("MONGO_DATABASE")
                    .unwrap_or_else(|_| "langgraph".to_string()),
                collection: std::env::var("MONGO_COLLECTION")
                    .unwrap_or_else(|_| "graph_definitions".to_string()),
            })
        }
    }

    /// MongoPersistence - MongoDB图存储实现
    pub struct MongoPersistence {
        client: Client,
        collection: Collection<Document>,
        database_name: String,
        collection_name: String,
    }

    impl MongoPersistence {
        /// 创建新的MongoDB持久化实例
        ///
        /// # 参数
        /// - config: MongoDB配置
        ///
        /// # 示例
        /// ```ignore
        /// let config = MongoConfig::new(
        ///     "mongodb://localhost:27017",
        ///     "langgraph",
        ///     "graph_definitions"
        /// );
        /// let persistence = MongoPersistence::new(config).await?;
        /// ```
        pub async fn new(config: MongoConfig) -> Result<Self, PersistenceError> {
            let client_options = ClientOptions::parse(&config.uri)
                .await
                .map_err(|e| PersistenceError::ConnectionError(e.to_string()))?;

            let client = Client::with_options(client_options)
                .map_err(|e| PersistenceError::ConnectionError(e.to_string()))?;

            let database = client.database(&config.database);
            let collection = database.collection(&config.collection);

            Ok(Self {
                client,
                collection,
                database_name: config.database,
                collection_name: config.collection,
            })
        }

        /// 从环境变量创建实例
        pub async fn from_env() -> Result<Self, PersistenceError> {
            let config = MongoConfig::from_env()?;
            Self::new(config).await
        }

        /// 获取MongoDB客户端
        pub fn client(&self) -> &Client {
            &self.client
        }

        /// 获取集合名称
        pub fn collection_name(&self) -> &str {
            &self.collection_name
        }

        /// 获取数据库名称
        pub fn database_name(&self) -> &str {
            &self.database_name
        }
    }

    #[async_trait]
    impl GraphPersistence for MongoPersistence {
        async fn save(
            &self,
            id: &str,
            definition: &GraphDefinition,
        ) -> Result<(), PersistenceError> {
            let doc = to_document(definition)
                .map_err(|e| PersistenceError::SerializationError(e.to_string()))?;

            // 使用 upsert 操作：如果存在则更新，不存在则插入
            self.collection
                .update_one(
                    doc! { "_id": id },
                    doc! { "$set": doc },
                    mongodb::options::UpdateOptions::builder()
                        .upsert(true)
                        .build(),
                )
                .await
                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;

            Ok(())
        }

        async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
            let filter = doc! { "_id": id };

            let result = self
                .collection
                .find_one(filter, None)
                .await
                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;

            match result {
                Some(doc) => {
                    let definition: GraphDefinition = from_document(doc)
                        .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
                    Ok(definition)
                }
                None => Err(PersistenceError::NotFound(id.to_string())),
            }
        }

        async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
            let filter = doc! { "_id": id };

            let result = self
                .collection
                .delete_one(filter, None)
                .await
                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;

            if result.deleted_count == 0 {
                Err(PersistenceError::NotFound(id.to_string()))
            } else {
                Ok(())
            }
        }

        async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
            let filter = doc! { "_id": id };

            let count = self
                .collection
                .count_documents(filter, None)
                .await
                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;

            Ok(count > 0)
        }

        async fn list(&self) -> Result<Vec<String>, PersistenceError> {
            let filter = doc! {};
            let options = FindOptions::builder().projection(doc! { "_id": 1 }).build();

            let mut cursor = self
                .collection
                .find(filter, options)
                .await
                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;

            let mut ids = Vec::new();
            while cursor
                .advance()
                .await
                .map_err(|e| PersistenceError::MongoError(e.to_string()))?
            {
                let doc = cursor
                    .deserialize_current()
                    .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;

                if let Ok(id) = doc.get_str("_id") {
                    ids.push(id.to_string());
                }
            }

            Ok(ids)
        }
    }
}

#[cfg(feature = "mongodb-persistence")]
pub use mongo_impl::{MongoConfig, MongoPersistence};

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

    #[test]
    fn test_node_type_serialization() {
        let types = vec![
            NodeType::Sync,
            NodeType::Async,
            NodeType::Subgraph,
            NodeType::Custom,
        ];

        for t in types {
            let json = serde_json::to_string(&t).unwrap();
            let parsed: NodeType = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, t);
        }
    }

    #[test]
    fn test_edge_type_serialization() {
        let types = vec![
            EdgeType::Fixed,
            EdgeType::Conditional,
            EdgeType::FanOut,
            EdgeType::FanIn,
        ];

        for t in types {
            let json = serde_json::to_string(&t).unwrap();
            let parsed: EdgeType = serde_json::from_str(&json).unwrap();
            assert_eq!(parsed, t);
        }
    }

    #[test]
    fn test_graph_definition_builder() {
        let def = GraphDefinition::new("entry".to_string())
            .with_id("test-id".to_string())
            .with_name("Test Graph".to_string())
            .with_recursion_limit(50);

        assert_eq!(def.id, "test-id");
        assert_eq!(def.name, Some("Test Graph".to_string()));
        assert_eq!(def.entry_point, "entry");
        assert_eq!(def.recursion_limit, 50);
    }

    #[tokio::test]
    async fn test_memory_persistence() {
        let persistence = MemoryPersistence::new();
        let def = GraphDefinition::new("entry".to_string()).with_id("test-001".to_string());

        persistence.save("test-001", &def).await.unwrap();
        assert!(persistence.exists("test-001").await.unwrap());

        let loaded = persistence.load("test-001").await.unwrap();
        assert_eq!(loaded.id, "test-001");

        persistence.delete("test-001").await.unwrap();
        assert!(!persistence.exists("test-001").await.unwrap());
    }
}