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    pub fn new_from_task(sid: String, task_name: &str) -> Self {
21        Self {
22            id: sid,
23            graph_id: "default".to_string(),
24            current_task_id: task_name.to_string(),
25            status_message: None,
26            context: Context::new(),
27        }
28    }
29}
30
31/// Trait for storing and retrieving graphs
32#[async_trait]
33pub trait GraphStorage: Send + Sync {
34    async fn save(&self, id: String, graph: Arc<Graph>) -> Result<()>;
35    async fn get(&self, id: &str) -> Result<Option<Arc<Graph>>>;
36    async fn delete(&self, id: &str) -> Result<()>;
37}
38
39/// Trait for storing and retrieving sessions
40#[async_trait]
41pub trait SessionStorage: Send + Sync {
42    async fn save(&self, session: Session) -> Result<()>;
43    async fn get(&self, id: &str) -> Result<Option<Session>>;
44    async fn delete(&self, id: &str) -> Result<()>;
45}
46
47/// In-memory implementation of GraphStorage
48pub struct InMemoryGraphStorage {
49    graphs: Arc<DashMap<String, Arc<Graph>>>,
50}
51
52impl Default for InMemoryGraphStorage {
53    fn default() -> Self {
54        Self::new()
55    }
56}
57
58impl InMemoryGraphStorage {
59    pub fn new() -> Self {
60        Self {
61            graphs: Arc::new(DashMap::new()),
62        }
63    }
64}
65
66#[async_trait]
67impl GraphStorage for InMemoryGraphStorage {
68    async fn save(&self, id: String, graph: Arc<Graph>) -> Result<()> {
69        self.graphs.insert(id, graph);
70        Ok(())
71    }
72
73    async fn get(&self, id: &str) -> Result<Option<Arc<Graph>>> {
74        Ok(self.graphs.get(id).map(|entry| entry.clone()))
75    }
76
77    async fn delete(&self, id: &str) -> Result<()> {
78        self.graphs.remove(id);
79        Ok(())
80    }
81}
82
83/// In-memory implementation of SessionStorage
84pub struct InMemorySessionStorage {
85    sessions: Arc<DashMap<String, Session>>,
86}
87
88impl Default for InMemorySessionStorage {
89    fn default() -> Self {
90        Self::new()
91    }
92}
93
94impl InMemorySessionStorage {
95    pub fn new() -> Self {
96        Self {
97            sessions: Arc::new(DashMap::new()),
98        }
99    }
100}
101
102#[async_trait]
103impl SessionStorage for InMemorySessionStorage {
104    async fn save(&self, session: Session) -> Result<()> {
105        self.sessions.insert(session.id.clone(), session);
106        Ok(())
107    }
108
109    async fn get(&self, id: &str) -> Result<Option<Session>> {
110        Ok(self.sessions.get(id).map(|entry| entry.clone()))
111    }
112
113    async fn delete(&self, id: &str) -> Result<()> {
114        self.sessions.remove(id);
115        Ok(())
116    }
117}
graph_flow/storage.rs

graph_flow/
storage.rs
