agentic_memory_mcp/session/
manager.rs

1//! Graph lifecycle management, file I/O, and session tracking.
2
3use std::path::PathBuf;
4use std::time::{Duration, Instant};
5
6use agentic_memory::{
7    AmemReader, AmemWriter, CognitiveEventBuilder, Edge, EdgeType, EventType, MemoryGraph,
8    QueryEngine, WriteEngine,
9};
10
11use crate::types::{McpError, McpResult};
12
13/// Default auto-save interval.
14const DEFAULT_AUTO_SAVE_SECS: u64 = 30;
15
16/// Manages the memory graph lifecycle, file I/O, and session state.
17pub struct SessionManager {
18    graph: MemoryGraph,
19    query_engine: QueryEngine,
20    write_engine: WriteEngine,
21    file_path: PathBuf,
22    current_session: u32,
23    dirty: bool,
24    last_save: Instant,
25    auto_save_interval: Duration,
26}
27
28impl SessionManager {
29    /// Open or create a memory file at the given path.
30    pub fn open(path: &str) -> McpResult<Self> {
31        let file_path = PathBuf::from(path);
32        let dimension = agentic_memory::DEFAULT_DIMENSION;
33
34        let graph = if file_path.exists() {
35            tracing::info!("Opening existing memory file: {}", file_path.display());
36            AmemReader::read_from_file(&file_path)
37                .map_err(|e| McpError::AgenticMemory(format!("Failed to read memory file: {e}")))?
38        } else {
39            tracing::info!("Creating new memory file: {}", file_path.display());
40            // Ensure parent directory exists
41            if let Some(parent) = file_path.parent() {
42                std::fs::create_dir_all(parent).map_err(|e| {
43                    McpError::Io(std::io::Error::other(format!(
44                        "Failed to create directory {}: {e}",
45                        parent.display()
46                    )))
47                })?;
48            }
49            MemoryGraph::new(dimension)
50        };
51
52        // Determine the next session ID from existing sessions
53        let session_ids = graph.session_index().session_ids();
54        let current_session = session_ids.iter().copied().max().unwrap_or(0) + 1;
55
56        tracing::info!(
57            "Session {} started. Graph has {} nodes, {} edges.",
58            current_session,
59            graph.node_count(),
60            graph.edge_count()
61        );
62
63        Ok(Self {
64            graph,
65            query_engine: QueryEngine::new(),
66            write_engine: WriteEngine::new(dimension),
67            file_path,
68            current_session,
69            dirty: false,
70            last_save: Instant::now(),
71            auto_save_interval: Duration::from_secs(DEFAULT_AUTO_SAVE_SECS),
72        })
73    }
74
75    /// Get an immutable reference to the graph.
76    pub fn graph(&self) -> &MemoryGraph {
77        &self.graph
78    }
79
80    /// Get a mutable reference to the graph and mark as dirty.
81    pub fn graph_mut(&mut self) -> &mut MemoryGraph {
82        self.dirty = true;
83        &mut self.graph
84    }
85
86    /// Get the query engine.
87    pub fn query_engine(&self) -> &QueryEngine {
88        &self.query_engine
89    }
90
91    /// Get the write engine.
92    pub fn write_engine(&self) -> &WriteEngine {
93        &self.write_engine
94    }
95
96    /// Current session ID.
97    pub fn current_session_id(&self) -> u32 {
98        self.current_session
99    }
100
101    /// Start a new session, optionally with an explicit ID.
102    pub fn start_session(&mut self, explicit_id: Option<u32>) -> McpResult<u32> {
103        let session_id = explicit_id.unwrap_or_else(|| {
104            let ids = self.graph.session_index().session_ids();
105            ids.iter().copied().max().unwrap_or(0) + 1
106        });
107
108        self.current_session = session_id;
109        tracing::info!("Started session {session_id}");
110        Ok(session_id)
111    }
112
113    /// End a session and optionally create an episode summary.
114    pub fn end_session_with_episode(&mut self, session_id: u32, summary: &str) -> McpResult<u64> {
115        let episode_id = self
116            .write_engine
117            .compress_session(&mut self.graph, session_id, summary)
118            .map_err(|e| McpError::AgenticMemory(format!("Failed to compress session: {e}")))?;
119
120        self.dirty = true;
121        self.save()?;
122
123        tracing::info!("Ended session {session_id}, created episode node {episode_id}");
124
125        Ok(episode_id)
126    }
127
128    /// Save the graph to file.
129    pub fn save(&mut self) -> McpResult<()> {
130        if !self.dirty {
131            return Ok(());
132        }
133
134        let writer = AmemWriter::new(self.graph.dimension());
135        writer
136            .write_to_file(&self.graph, &self.file_path)
137            .map_err(|e| McpError::AgenticMemory(format!("Failed to write memory file: {e}")))?;
138
139        self.dirty = false;
140        self.last_save = Instant::now();
141        tracing::debug!("Saved memory file: {}", self.file_path.display());
142        Ok(())
143    }
144
145    /// Check if auto-save is needed and save if so.
146    pub fn maybe_auto_save(&mut self) -> McpResult<()> {
147        if self.dirty && self.last_save.elapsed() >= self.auto_save_interval {
148            self.save()?;
149        }
150        Ok(())
151    }
152
153    /// Mark the graph as dirty (needs saving).
154    pub fn mark_dirty(&mut self) {
155        self.dirty = true;
156    }
157
158    /// Get the file path.
159    pub fn file_path(&self) -> &PathBuf {
160        &self.file_path
161    }
162
163    /// Add a cognitive event to the graph.
164    pub fn add_event(
165        &mut self,
166        event_type: EventType,
167        content: &str,
168        confidence: f32,
169        edges: Vec<(u64, EdgeType, f32)>,
170    ) -> McpResult<(u64, usize)> {
171        let event = CognitiveEventBuilder::new(event_type, content.to_string())
172            .session_id(self.current_session)
173            .confidence(confidence)
174            .build();
175
176        // First, add the node to get its assigned ID
177        let result = self
178            .write_engine
179            .ingest(&mut self.graph, vec![event], vec![])
180            .map_err(|e| McpError::AgenticMemory(format!("Failed to add event: {e}")))?;
181
182        let node_id = result.new_node_ids.first().copied().ok_or_else(|| {
183            McpError::InternalError("No node ID returned from ingest".to_string())
184        })?;
185
186        // Then add edges with the correct source_id
187        let mut edge_count = 0;
188        for (target_id, edge_type, weight) in &edges {
189            let edge = Edge::new(node_id, *target_id, *edge_type, *weight);
190            self.graph
191                .add_edge(edge)
192                .map_err(|e| McpError::AgenticMemory(format!("Failed to add edge: {e}")))?;
193            edge_count += 1;
194        }
195
196        self.dirty = true;
197        self.maybe_auto_save()?;
198
199        Ok((node_id, edge_count))
200    }
201
202    /// Correct a previous belief.
203    pub fn correct_node(&mut self, old_node_id: u64, new_content: &str) -> McpResult<u64> {
204        let new_id = self
205            .write_engine
206            .correct(
207                &mut self.graph,
208                old_node_id,
209                new_content,
210                self.current_session,
211            )
212            .map_err(|e| McpError::AgenticMemory(format!("Failed to correct node: {e}")))?;
213
214        self.dirty = true;
215        self.maybe_auto_save()?;
216
217        Ok(new_id)
218    }
219}
220
221impl Drop for SessionManager {
222    fn drop(&mut self) {
223        if self.dirty {
224            if let Err(e) = self.save() {
225                tracing::error!("Failed to save on drop: {e}");
226            }
227        }
228    }
229}
agentic_memory_mcp/session/manager.rs

agentic_memory_mcp/session/
manager.rs
