graph_flow/
storage.rs

1use async_trait::async_trait;
2use dashmap::DashMap;
3use serde::{Deserialize, Serialize};
4use std::sync::Arc;
5
6use crate::{Context, error::Result, graph::Graph};
7
8/// Session information
9#[derive(Debug, Clone, Serialize, Deserialize)]
10pub struct Session {
11    pub id: String,
12    pub graph_id: String,
13    pub current_task_id: String,
14    /// Optional status message from the last executed task
15    pub status_message: Option<String>,
16    pub context: crate::context::Context,
17}
18
19impl Session {
20    /// Create a new session positioned at the given task.
21    ///
22    /// Note: `graph_id` is set to `"default"`; assign it explicitly if you
23    /// store multiple graphs.
24    pub fn new_from_task(sid: String, task_name: &str) -> Self {
25        Self {
26            id: sid,
27            graph_id: "default".to_string(),
28            current_task_id: task_name.to_string(),
29            status_message: None,
30            context: Context::new(),
31        }
32    }
33}
34
35/// Trait for storing and retrieving graphs
36#[async_trait]
37pub trait GraphStorage: Send + Sync {
38    async fn save(&self, id: String, graph: Arc<Graph>) -> Result<()>;
39    async fn get(&self, id: &str) -> Result<Option<Arc<Graph>>>;
40    async fn delete(&self, id: &str) -> Result<()>;
41}
42
43/// Trait for storing and retrieving sessions
44#[async_trait]
45pub trait SessionStorage: Send + Sync {
46    async fn save(&self, session: Session) -> Result<()>;
47    async fn get(&self, id: &str) -> Result<Option<Session>>;
48    async fn delete(&self, id: &str) -> Result<()>;
49}
50
51/// In-memory implementation of GraphStorage
52pub struct InMemoryGraphStorage {
53    graphs: DashMap<String, Arc<Graph>>,
54}
55
56impl Default for InMemoryGraphStorage {
57    fn default() -> Self {
58        Self::new()
59    }
60}
61
62impl InMemoryGraphStorage {
63    pub fn new() -> Self {
64        Self {
65            graphs: DashMap::new(),
66        }
67    }
68}
69
70#[async_trait]
71impl GraphStorage for InMemoryGraphStorage {
72    async fn save(&self, id: String, graph: Arc<Graph>) -> Result<()> {
73        self.graphs.insert(id, graph);
74        Ok(())
75    }
76
77    async fn get(&self, id: &str) -> Result<Option<Arc<Graph>>> {
78        Ok(self.graphs.get(id).map(|entry| entry.clone()))
79    }
80
81    async fn delete(&self, id: &str) -> Result<()> {
82        self.graphs.remove(id);
83        Ok(())
84    }
85}
86
87/// In-memory implementation of SessionStorage
88pub struct InMemorySessionStorage {
89    sessions: DashMap<String, Session>,
90}
91
92impl Default for InMemorySessionStorage {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98impl InMemorySessionStorage {
99    pub fn new() -> Self {
100        Self {
101            sessions: DashMap::new(),
102        }
103    }
104}
105
106#[async_trait]
107impl SessionStorage for InMemorySessionStorage {
108    async fn save(&self, session: Session) -> Result<()> {
109        self.sessions.insert(session.id.clone(), session);
110        Ok(())
111    }
112
113    async fn get(&self, id: &str) -> Result<Option<Session>> {
114        Ok(self.sessions.get(id).map(|entry| entry.clone()))
115    }
116
117    async fn delete(&self, id: &str) -> Result<()> {
118        self.sessions.remove(id);
119        Ok(())
120    }
121}
graph_flow/storage.rs

graph_flow/
storage.rs
