lc_langgraph/
persistence.rs

1// crates/lc-langgraph/src/persistence.rs
2//! Graph persistence for serialization and storage
3//!
4//! This module provides persistence capabilities for graph definitions,
5//! allowing graphs to be saved, loaded, and shared across sessions.
6
7use async_trait::async_trait;
8use chrono::{DateTime, Utc};
9use serde::{Deserialize, Serialize};
10use std::collections::HashMap;
11use std::path::PathBuf;
12use tokio::sync::Mutex;
13use uuid::Uuid;
14
15/// GraphPersistence trait for storing and loading graph definitions
16#[async_trait]
17pub trait GraphPersistence: Send + Sync {
18    /// Save a graph definition with the given ID
19    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError>;
20
21    /// Load a graph definition by ID
22    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError>;
23    /// Delete a graph definition by ID
24    async fn delete(&self, id: &str) -> Result<(), PersistenceError>;
25
26    /// Check if a graph definition exists
27    async fn exists(&self, id: &str) -> Result<bool, PersistenceError>;
28
29    /// List all stored graph IDs
30    async fn list(&self) -> Result<Vec<String>, PersistenceError>;
31}
32
33/// Persistence error types
34#[derive(Debug, thiserror::Error)]
35#[non_exhaustive]
36pub enum PersistenceError {
37    /// The requested graph was not found.
38    #[error("Graph '{0}' not found")]
39    NotFound(String),
40
41    /// Failed to serialize a graph definition.
42    #[error("Serialization error: {0}")]
43    SerializationError(String),
44
45    /// Failed to deserialize a graph definition.
46    #[error("Deserialization error: {0}")]
47    DeserializationError(String),
48
49    /// An I/O error occurred while accessing storage.
50    #[error("IO error: {0}")]
51    IoError(String),
52
53    /// The graph definition is invalid.
54    #[error("Invalid graph definition: {0}")]
55    InvalidDefinition(String),
56
57    /// A MongoDB operation failed.
58    #[error("MongoDB error: {0}")]
59    MongoError(String),
60
61    /// Failed to connect to the persistence backend.
62    #[error("Connection error: {0}")]
63    ConnectionError(String),
64}
65
66/// GraphDefinition - Serializable graph structure
67#[derive(Debug, Clone, Serialize, Deserialize)]
68pub struct GraphDefinition {
69    /// Unique identifier
70    pub id: String,
71
72    /// Human-readable name
73    pub name: Option<String>,
74
75    /// Entry point node name
76    pub entry_point: String,
77
78    /// Node definitions
79    pub nodes: Vec<NodeDefinition>,
80
81    /// Edge definitions
82    pub edges: Vec<EdgeDefinition>,
83
84    /// Router definitions
85    pub routers: Vec<RouterDefinition>,
86
87    /// Maximum recursion limit
88    pub recursion_limit: usize,
89
90    /// Creation timestamp
91    pub created_at: DateTime<Utc>,
92
93    /// Last update timestamp
94    pub updated_at: DateTime<Utc>,
95
96    /// Custom metadata
97    pub metadata: HashMap<String, serde_json::Value>,
98}
99
100impl GraphDefinition {
101    /// Create a new graph definition with the given entry point
102    pub fn new(entry_point: impl Into<String>) -> Self {
103        let now = Utc::now();
104        Self {
105            id: Uuid::new_v4().to_string(),
106            name: None,
107            entry_point: entry_point.into(),
108            nodes: Vec::new(),
109            edges: Vec::new(),
110            routers: Vec::new(),
111            recursion_limit: 25,
112            created_at: now,
113            updated_at: now,
114            metadata: HashMap::new(),
115        }
116    }
117
118    /// Set a custom ID
119    pub fn with_id(mut self, id: impl Into<String>) -> Self {
120        self.id = id.into();
121        self
122    }
123
124    /// Set a human-readable name
125    pub fn with_name(mut self, name: impl Into<String>) -> Self {
126        self.name = Some(name.into());
127        self
128    }
129
130    /// Set recursion limit
131    pub fn with_recursion_limit(mut self, limit: usize) -> Self {
132        self.recursion_limit = limit;
133        self
134    }
135
136    /// Add a node definition
137    pub fn add_node(&mut self, node: NodeDefinition) {
138        self.nodes.push(node);
139        self.updated_at = Utc::now();
140    }
141
142    /// Add an edge definition
143    pub fn add_edge(&mut self, edge: EdgeDefinition) {
144        self.edges.push(edge);
145        self.updated_at = Utc::now();
146    }
147
148    /// Add a router definition
149    pub fn add_router(&mut self, router: RouterDefinition) {
150        self.routers.push(router);
151        self.updated_at = Utc::now();
152    }
153}
154
155/// NodeDefinition - Serializable node structure
156#[derive(Debug, Clone, Serialize, Deserialize)]
157pub struct NodeDefinition {
158    /// Node name
159    pub name: String,
160
161    /// Node type
162    pub node_type: NodeType,
163
164    /// Custom configuration
165    pub config: serde_json::Value,
166}
167
168/// NodeType - Type of node execution
169#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
170pub enum NodeType {
171    /// Synchronous node.
172    Sync,
173    /// Asynchronous node.
174    Async,
175    /// A subgraph node.
176    Subgraph,
177    /// A custom node type.
178    Custom,
179}
180
181/// EdgeDefinition - Serializable edge structure
182#[derive(Debug, Clone, Serialize, Deserialize)]
183pub struct EdgeDefinition {
184    /// Edge type
185    pub edge_type: EdgeType,
186
187    /// Source node name
188    pub source: String,
189
190    /// Target node name (for fixed edges)
191    pub target: Option<String>,
192
193    /// Multiple targets (for fan-out edges)
194    pub targets: Option<Vec<String>>,
195
196    /// Router name (for conditional edges)
197    pub router_name: Option<String>,
198
199    /// Conditional targets mapping (route -> target)
200    pub conditional_targets: Option<HashMap<String, String>>,
201
202    /// Default target for conditional edges
203    pub default_target: Option<String>,
204
205    /// Source nodes for fan-in edges
206    pub sources: Option<Vec<String>>,
207}
208
209impl EdgeDefinition {
210    /// Create a fixed edge
211    pub fn fixed(source: impl Into<String>, target: impl Into<String>) -> Self {
212        Self {
213            edge_type: EdgeType::Fixed,
214            source: source.into(),
215            target: Some(target.into()),
216            targets: None,
217            router_name: None,
218            conditional_targets: None,
219            default_target: None,
220            sources: None,
221        }
222    }
223
224    /// Create a conditional edge
225    pub fn conditional(
226        source: impl Into<String>,
227        router_name: impl Into<String>,
228        targets: HashMap<String, String>,
229        default_target: Option<String>,
230    ) -> Self {
231        Self {
232            edge_type: EdgeType::Conditional,
233            source: source.into(),
234            target: None,
235            targets: None,
236            router_name: Some(router_name.into()),
237            conditional_targets: Some(targets),
238            default_target,
239            sources: None,
240        }
241    }
242
243    /// Create a fan-out edge (parallel execution)
244    pub fn fan_out(source: impl Into<String>, targets: Vec<String>) -> Self {
245        Self {
246            edge_type: EdgeType::FanOut,
247            source: source.into(),
248            target: None,
249            targets: Some(targets),
250            router_name: None,
251            conditional_targets: None,
252            default_target: None,
253            sources: None,
254        }
255    }
256
257    /// Create a fan-in edge (merge from parallel branches)
258    pub fn fan_in(sources: Vec<String>, target: impl Into<String>) -> Self {
259        Self {
260            edge_type: EdgeType::FanIn,
261            source: "__fan_in__".to_string(),
262            target: Some(target.into()),
263            targets: None,
264            router_name: None,
265            conditional_targets: None,
266            default_target: None,
267            sources: Some(sources),
268        }
269    }
270}
271
272/// EdgeType - Type of edge connection
273#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
274pub enum EdgeType {
275    /// Fixed transition
276    Fixed,
277
278    /// Conditional routing
279    Conditional,
280
281    /// Parallel fan-out
282    FanOut,
283
284    /// Merge fan-in
285    FanIn,
286}
287
288/// RouterDefinition - Serializable router structure
289#[derive(Debug, Clone, Serialize, Deserialize)]
290pub struct RouterDefinition {
291    /// Router name
292    pub name: String,
293
294    /// Router type (e.g., "function", "state_key")
295    pub router_type: String,
296
297    /// Possible routes
298    pub routes: Vec<String>,
299
300    /// Custom configuration
301    pub config: serde_json::Value,
302}
303
304/// MemoryPersistence - In-memory graph storage
305pub struct MemoryPersistence {
306    graphs: Mutex<HashMap<String, GraphDefinition>>,
307}
308
309impl MemoryPersistence {
310    /// Create a new memory persistence instance
311    pub fn new() -> Self {
312        Self {
313            graphs: Mutex::new(HashMap::new()),
314        }
315    }
316}
317
318impl Default for MemoryPersistence {
319    fn default() -> Self {
320        Self::new()
321    }
322}
323
324#[async_trait]
325impl GraphPersistence for MemoryPersistence {
326    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError> {
327        let mut graphs = self.graphs.lock().await;
328        graphs.insert(id.to_string(), definition.clone());
329        Ok(())
330    }
331
332    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
333        let graphs = self.graphs.lock().await;
334        graphs
335            .get(id)
336            .cloned()
337            .ok_or_else(|| PersistenceError::NotFound(id.to_string()))
338    }
339
340    async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
341        let mut graphs = self.graphs.lock().await;
342        graphs
343            .remove(id)
344            .map(|_| ())
345            .ok_or_else(|| PersistenceError::NotFound(id.to_string()))
346    }
347
348    async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
349        let graphs = self.graphs.lock().await;
350        Ok(graphs.contains_key(id))
351    }
352
353    async fn list(&self) -> Result<Vec<String>, PersistenceError> {
354        let graphs = self.graphs.lock().await;
355        Ok(graphs.keys().cloned().collect())
356    }
357}
358
359/// FilePersistence - File-based graph storage
360pub struct FilePersistence {
361    directory: PathBuf,
362}
363
364impl FilePersistence {
365    /// Create a new file persistence instance
366    pub fn new(directory: impl Into<PathBuf>) -> Result<Self, PersistenceError> {
367        let dir = directory.into();
368        if !dir.exists() {
369            std::fs::create_dir_all(&dir).map_err(|e| {
370                PersistenceError::IoError(format!(
371                    "Failed to create directory '{}': {}",
372                    dir.display(),
373                    e
374                ))
375            })?;
376        }
377        Ok(Self { directory: dir })
378    }
379
380    fn graph_path(&self, id: &str) -> Result<PathBuf, PersistenceError> {
381        // Sanitize id to prevent path traversal: reject ".." and absolute paths
382        if id.contains("..") || id.contains('/') || id.contains('\\') {
383            return Err(PersistenceError::InvalidDefinition(format!(
384                "Invalid graph id '{}': path traversal detected",
385                id
386            )));
387        }
388        if std::path::Path::new(id).is_absolute() {
389            return Err(PersistenceError::InvalidDefinition(format!(
390                "Invalid graph id '{}': absolute path not allowed",
391                id
392            )));
393        }
394        Ok(self.directory.join(format!("{}.json", id)))
395    }
396}
397
398// NOTE: `Default` is intentionally NOT implemented for `FilePersistence` (Q1).
399// The default constructor would have to create a directory, which is I/O that can
400// fail (read-only cwd, disk full, permissions) — `Default` cannot report that
401// failure, so it would have to panic. Use `FilePersistence::new(...)` which
402// returns a `Result` and surfaces the error instead.
403
404#[async_trait]
405impl GraphPersistence for FilePersistence {
406    async fn save(&self, id: &str, definition: &GraphDefinition) -> Result<(), PersistenceError> {
407        let path = self.graph_path(id)?;
408
409        let json = serde_json::to_string_pretty(definition)
410            .map_err(|e| PersistenceError::SerializationError(e.to_string()))?;
411
412        tokio::fs::write(&path, json)
413            .await
414            .map_err(|e| PersistenceError::IoError(e.to_string()))?;
415
416        Ok(())
417    }
418
419    async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
420        let path = self.graph_path(id)?;
421
422        if !path.exists() {
423            return Err(PersistenceError::NotFound(id.to_string()));
424        }
425
426        let json = tokio::fs::read_to_string(&path)
427            .await
428            .map_err(|e| PersistenceError::IoError(e.to_string()))?;
429
430        let definition: GraphDefinition = serde_json::from_str(&json)
431            .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
432
433        Ok(definition)
434    }
435
436    async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
437        let path = self.graph_path(id)?;
438
439        if path.exists() {
440            tokio::fs::remove_file(&path)
441                .await
442                .map_err(|e| PersistenceError::IoError(e.to_string()))?;
443        }
444
445        Ok(())
446    }
447
448    async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
449        let path = self.graph_path(id)?;
450        Ok(path.exists())
451    }
452
453    async fn list(&self) -> Result<Vec<String>, PersistenceError> {
454        let mut ids = Vec::new();
455
456        let mut entries = tokio::fs::read_dir(&self.directory)
457            .await
458            .map_err(|e| PersistenceError::IoError(e.to_string()))?;
459
460        while let Some(entry) = entries
461            .next_entry()
462            .await
463            .map_err(|e| PersistenceError::IoError(e.to_string()))?
464        {
465            let path = entry.path();
466            if path.extension().is_some_and(|ext| ext == "json") {
467                if let Some(id) = path.file_stem().and_then(|s| s.to_str()) {
468                    ids.push(id.to_string());
469                }
470            }
471        }
472
473        Ok(ids)
474    }
475}
476
477// ============================================================================
478// MongoDB persistence implementation
479// ============================================================================
480
481#[cfg(feature = "mongodb-persistence")]
482mod mongo_impl {
483    use super::*;
484    use mongodb::{
485        bson::{doc, from_document, to_document, Document},
486        options::{ClientOptions, FindOptions},
487        Client, Collection,
488    };
489
490    /// MongoDB configuration
491    pub struct MongoConfig {
492        /// MongoDB connection URI (e.g. mongodb://localhost:27017 or mongodb+srv://user:pass@cluster.mongodb.net)
493        pub uri: String,
494
495        /// Database name
496        pub database: String,
497
498        /// Collection name
499        pub collection: String,
500    }
501
502    impl MongoConfig {
503        /// Create a new MongoDB configuration
504        pub fn new(
505            uri: impl Into<String>,
506            database: impl Into<String>,
507            collection: impl Into<String>,
508        ) -> Self {
509            Self {
510                uri: uri.into(),
511                database: database.into(),
512                collection: collection.into(),
513            }
514        }
515
516        /// Create a configuration from environment variables
517        ///
518        /// Environment variables:
519        /// - MONGO_URI: MongoDB connection URI
520        /// - MONGO_DATABASE: database name (default: langgraph)
521        /// - MONGO_COLLECTION: collection name (default: graph_definitions)
522        pub fn from_env() -> Result<Self, PersistenceError> {
523            Ok(Self {
524                uri: std::env::var("MONGO_URI").map_err(|_| {
525                    PersistenceError::ConnectionError(
526                        "MONGO_URI environment variable not set".to_string(),
527                    )
528                })?,
529                database: std::env::var("MONGO_DATABASE")
530                    .unwrap_or_else(|_| "langgraph".to_string()),
531                collection: std::env::var("MONGO_COLLECTION")
532                    .unwrap_or_else(|_| "graph_definitions".to_string()),
533            })
534        }
535    }
536
537    /// MongoPersistence - MongoDB graph storage implementation
538    pub struct MongoPersistence {
539        client: Client,
540        collection: Collection<Document>,
541        database_name: String,
542        collection_name: String,
543    }
544
545    impl MongoPersistence {
546        /// Create a new MongoDB persistence instance
547        ///
548        /// # Arguments
549        /// - config: MongoDB configuration
550        ///
551        /// # Example
552        /// ```ignore
553        /// let config = MongoConfig::new(
554        ///     "mongodb://localhost:27017",
555        ///     "langgraph",
556        ///     "graph_definitions"
557        /// );
558        /// let persistence = MongoPersistence::new(config).await?;
559        /// ```
560        pub async fn new(config: MongoConfig) -> Result<Self, PersistenceError> {
561            let client_options = ClientOptions::parse(&config.uri)
562                .await
563                .map_err(|e| PersistenceError::ConnectionError(e.to_string()))?;
564
565            let client = Client::with_options(client_options)
566                .map_err(|e| PersistenceError::ConnectionError(e.to_string()))?;
567
568            let database = client.database(&config.database);
569            let collection = database.collection(&config.collection);
570
571            Ok(Self {
572                client,
573                collection,
574                database_name: config.database,
575                collection_name: config.collection,
576            })
577        }
578
579        /// Create an instance from environment variables
580        pub async fn from_env() -> Result<Self, PersistenceError> {
581            let config = MongoConfig::from_env()?;
582            Self::new(config).await
583        }
584
585        /// Get the MongoDB client
586        pub fn client(&self) -> &Client {
587            &self.client
588        }
589
590        /// Get the collection name
591        pub fn collection_name(&self) -> &str {
592            &self.collection_name
593        }
594
595        /// Get the database name
596        pub fn database_name(&self) -> &str {
597            &self.database_name
598        }
599    }
600
601    #[async_trait]
602    impl GraphPersistence for MongoPersistence {
603        async fn save(
604            &self,
605            id: &str,
606            definition: &GraphDefinition,
607        ) -> Result<(), PersistenceError> {
608            let doc = to_document(definition)
609                .map_err(|e| PersistenceError::SerializationError(e.to_string()))?;
610
611            // Use an upsert: update if it exists, insert otherwise
612            self.collection
613                .update_one(
614                    doc! { "_id": id },
615                    doc! { "$set": doc },
616                    mongodb::options::UpdateOptions::builder()
617                        .upsert(true)
618                        .build(),
619                )
620                .await
621                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
622
623            Ok(())
624        }
625
626        async fn load(&self, id: &str) -> Result<GraphDefinition, PersistenceError> {
627            let filter = doc! { "_id": id };
628
629            let result = self
630                .collection
631                .find_one(filter, None)
632                .await
633                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
634
635            match result {
636                Some(doc) => {
637                    let definition: GraphDefinition = from_document(doc)
638                        .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
639                    Ok(definition)
640                }
641                None => Err(PersistenceError::NotFound(id.to_string())),
642            }
643        }
644
645        async fn delete(&self, id: &str) -> Result<(), PersistenceError> {
646            let filter = doc! { "_id": id };
647
648            let result = self
649                .collection
650                .delete_one(filter, None)
651                .await
652                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
653
654            if result.deleted_count == 0 {
655                Err(PersistenceError::NotFound(id.to_string()))
656            } else {
657                Ok(())
658            }
659        }
660
661        async fn exists(&self, id: &str) -> Result<bool, PersistenceError> {
662            let filter = doc! { "_id": id };
663
664            let count = self
665                .collection
666                .count_documents(filter, None)
667                .await
668                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
669
670            Ok(count > 0)
671        }
672
673        async fn list(&self) -> Result<Vec<String>, PersistenceError> {
674            let filter = doc! {};
675            let options = FindOptions::builder().projection(doc! { "_id": 1 }).build();
676
677            let mut cursor = self
678                .collection
679                .find(filter, options)
680                .await
681                .map_err(|e| PersistenceError::MongoError(e.to_string()))?;
682
683            let mut ids = Vec::new();
684            while cursor
685                .advance()
686                .await
687                .map_err(|e| PersistenceError::MongoError(e.to_string()))?
688            {
689                let doc = cursor
690                    .deserialize_current()
691                    .map_err(|e| PersistenceError::DeserializationError(e.to_string()))?;
692
693                if let Ok(id) = doc.get_str("_id") {
694                    ids.push(id.to_string());
695                }
696            }
697
698            Ok(ids)
699        }
700    }
701}
702
703#[cfg(feature = "mongodb-persistence")]
704pub use mongo_impl::{MongoConfig, MongoPersistence};
705
706#[cfg(test)]
707mod tests {
708    use super::*;
709
710    #[test]
711    fn test_node_type_serialization() {
712        let types = vec![
713            NodeType::Sync,
714            NodeType::Async,
715            NodeType::Subgraph,
716            NodeType::Custom,
717        ];
718
719        for t in types {
720            let json = serde_json::to_string(&t).unwrap();
721            let parsed: NodeType = serde_json::from_str(&json).unwrap();
722            assert_eq!(parsed, t);
723        }
724    }
725
726    #[test]
727    fn test_edge_type_serialization() {
728        let types = vec![
729            EdgeType::Fixed,
730            EdgeType::Conditional,
731            EdgeType::FanOut,
732            EdgeType::FanIn,
733        ];
734
735        for t in types {
736            let json = serde_json::to_string(&t).unwrap();
737            let parsed: EdgeType = serde_json::from_str(&json).unwrap();
738            assert_eq!(parsed, t);
739        }
740    }
741
742    #[test]
743    fn test_graph_definition_builder() {
744        let def = GraphDefinition::new("entry".to_string())
745            .with_id("test-id".to_string())
746            .with_name("Test Graph".to_string())
747            .with_recursion_limit(50);
748
749        assert_eq!(def.id, "test-id");
750        assert_eq!(def.name, Some("Test Graph".to_string()));
751        assert_eq!(def.entry_point, "entry");
752        assert_eq!(def.recursion_limit, 50);
753    }
754
755    #[tokio::test]
756    async fn test_memory_persistence() {
757        let persistence = MemoryPersistence::new();
758        let def = GraphDefinition::new("entry".to_string()).with_id("test-001".to_string());
759
760        persistence.save("test-001", &def).await.unwrap();
761        assert!(persistence.exists("test-001").await.unwrap());
762
763        let loaded = persistence.load("test-001").await.unwrap();
764        assert_eq!(loaded.id, "test-001");
765
766        persistence.delete("test-001").await.unwrap();
767        assert!(!persistence.exists("test-001").await.unwrap());
768    }
769}
lc_langgraph/persistence.rs

lc_langgraph/
persistence.rs
