agent_diva_core/memory/
manager.rs

1//! Memory manager for handling long-term memory
2
3use super::provider::{
4    MemoryProvider, PrefetchRequest, PrefetchResponse, PrefetchStatus, SessionEndRequest,
5    SessionEndResponse, SessionEndStatus, StartupInjectionShape, SyncTurnRequest,
6    SyncTurnResponse, SyncTurnStatus, SystemPromptBlock, SystemPromptRequest,
7    SystemPromptResponse,
8};
9use super::storage::{DailyNote, Memory};
10use parking_lot::Mutex;
11use std::collections::HashSet;
12use std::path::{Path, PathBuf};
13
14/// Manages long-term memory storage
15#[derive(Debug)]
16pub struct MemoryManager {
17    /// Workspace directory
18    _workspace: PathBuf,
19    /// Memory file path
20    memory_path: PathBuf,
21    /// Daily notes directory
22    notes_dir: PathBuf,
23    /// History file path
24    history_path: PathBuf,
25    /// Session IDs whose shutdown hook has already been handled.
26    handled_session_end_ids: Mutex<HashSet<String>>,
27}
28
29impl MemoryManager {
30    /// Create a new memory manager
31    pub fn new<P: AsRef<Path>>(workspace: P) -> Self {
32        let workspace = workspace.as_ref().to_path_buf();
33        let memory_path = workspace.join("memory").join("MEMORY.md");
34        let history_path = workspace.join("memory").join("HISTORY.md");
35        let notes_dir = workspace.join("memory");
36
37        Self {
38            _workspace: workspace,
39            memory_path,
40            notes_dir,
41            history_path,
42            handled_session_end_ids: Mutex::new(HashSet::new()),
43        }
44    }
45
46    /// Load the long-term memory
47    pub fn load_memory(&self) -> Memory {
48        if self.memory_path.exists() {
49            match std::fs::read_to_string(&self.memory_path) {
50                Ok(content) => Memory::with_content(content),
51                Err(_) => Memory::new(),
52            }
53        } else {
54            Memory::new()
55        }
56    }
57
58    /// Save the long-term memory
59    pub fn save_memory(&self, memory: &Memory) -> crate::Result<()> {
60        if let Some(parent) = self.memory_path.parent() {
61            std::fs::create_dir_all(parent)?;
62        }
63        std::fs::write(&self.memory_path, &memory.content)?;
64        Ok(())
65    }
66
67    /// Load history entries from `HISTORY.md`
68    pub fn load_history(&self) -> String {
69        if self.history_path.exists() {
70            std::fs::read_to_string(&self.history_path).unwrap_or_default()
71        } else {
72            String::new()
73        }
74    }
75
76    /// Append an entry to `HISTORY.md`
77    pub fn append_history(&self, entry: &str) -> crate::Result<()> {
78        if entry.trim().is_empty() {
79            return Ok(());
80        }
81        if let Some(parent) = self.history_path.parent() {
82            std::fs::create_dir_all(parent)?;
83        }
84        let mut content = self.load_history();
85        if !content.is_empty() && !content.ends_with('\n') {
86            content.push('\n');
87        }
88        content.push_str(entry.trim_end());
89        content.push_str("\n\n");
90        std::fs::write(&self.history_path, content)?;
91        Ok(())
92    }
93
94    /// Load a daily note
95    pub fn load_daily_note(&self, date: impl AsRef<str>) -> DailyNote {
96        let date = date.as_ref();
97        let path = self.notes_dir.join(format!("{}.md", date));
98
99        if path.exists() {
100            match std::fs::read_to_string(&path) {
101                Ok(content) => {
102                    let mut note = DailyNote::for_date(date);
103                    note.content = content;
104                    note
105                }
106                Err(_) => DailyNote::for_date(date),
107            }
108        } else {
109            DailyNote::for_date(date)
110        }
111    }
112
113    /// Load today's note
114    pub fn load_today_note(&self) -> DailyNote {
115        let today = chrono::Local::now().format("%Y-%m-%d").to_string();
116        self.load_daily_note(&today)
117    }
118
119    /// Save a daily note
120    pub fn save_daily_note(&self, note: &DailyNote) -> crate::Result<()> {
121        std::fs::create_dir_all(&self.notes_dir)?;
122        let path = self.notes_dir.join(note.filename());
123        std::fs::write(&path, &note.content)?;
124        Ok(())
125    }
126
127    /// List all daily notes
128    pub fn list_notes(&self) -> Vec<String> {
129        let mut notes = Vec::new();
130
131        if let Ok(entries) = std::fs::read_dir(&self.notes_dir) {
132            for entry in entries.flatten() {
133                if let Some(name) = entry.file_name().to_str() {
134                    if name.ends_with(".md") && name != "MEMORY.md" {
135                        let date = name.trim_end_matches(".md").to_string();
136                        notes.push(date);
137                    }
138                }
139            }
140        }
141
142        notes.sort_by(|a, b| b.cmp(a)); // Newest first
143        notes
144    }
145
146    /// Get the memory directory path
147    pub fn memory_dir(&self) -> &Path {
148        &self.notes_dir
149    }
150
151    /// Append content to today's daily note
152    pub fn append_today(&self, content: &str) -> crate::Result<()> {
153        let mut note = self.load_today_note();
154
155        if note.content.is_empty() {
156            // Add header for new day
157            let today = chrono::Local::now().format("%Y-%m-%d").to_string();
158            note.content = format!("# {}\n\n{}", today, content);
159        } else {
160            // Append to existing content
161            note.content.push('\n');
162            note.content.push_str(content);
163        }
164
165        self.save_daily_note(&note)
166    }
167
168    /// Get memories from the last N days
169    pub fn get_recent_memories(&self, days: usize) -> String {
170        use chrono::Duration;
171
172        let mut memories = Vec::new();
173        let today = chrono::Local::now().date_naive();
174
175        for i in 0..days {
176            let date = today - Duration::days(i as i64);
177            let date_str = date.format("%Y-%m-%d").to_string();
178            let note = self.load_daily_note(&date_str);
179
180            if !note.content.is_empty() {
181                memories.push(note.content);
182            }
183        }
184
185        memories.join("\n\n---\n\n")
186    }
187
188    /// List all memory files sorted by date (newest first)
189    pub fn list_memory_files(&self) -> Vec<PathBuf> {
190        let mut files = Vec::new();
191
192        if let Ok(entries) = std::fs::read_dir(&self.notes_dir) {
193            for entry in entries.flatten() {
194                let path = entry.path();
195                if let Some(name) = path.file_name().and_then(|n| n.to_str()) {
196                    // Match pattern YYYY-MM-DD.md
197                    if name.len() == 13 && name.ends_with(".md") && name != "MEMORY.md" {
198                        let date_part = &name[..10];
199                        // Basic validation: check if it looks like a date
200                        if date_part.chars().filter(|c| *c == '-').count() == 2 {
201                            files.push(path);
202                        }
203                    }
204                }
205            }
206        }
207
208        // Sort by filename (which is the date) in reverse order
209        files.sort_by(|a, b| b.cmp(a));
210        files
211    }
212
213    /// Get memory context for the agent.
214    /// The redesigned memory model injects only long-term memory into prompts.
215    pub fn get_memory_context(&self) -> String {
216        let memory = self.load_memory();
217        if memory.content.is_empty() {
218            String::new()
219        } else {
220            format!("## Long-term Memory\n{}", memory.content)
221        }
222    }
223
224}
225
226#[async_trait::async_trait]
227impl MemoryProvider for MemoryManager {
228    fn system_prompt_block(
229        &self,
230        _request: &SystemPromptRequest,
231    ) -> crate::Result<SystemPromptResponse> {
232        let context = self.get_memory_context();
233        if context.is_empty() {
234            Ok(SystemPromptResponse::degraded(
235                "startup continuity unavailable; no long-term memory available",
236            ))
237        } else {
238            Ok(SystemPromptResponse::ready(SystemPromptBlock {
239                shape: StartupInjectionShape::CompactRenderedMarkdown,
240                markdown: context,
241            }))
242        }
243    }
244
245    async fn prefetch(&self, request: PrefetchRequest) -> crate::Result<PrefetchResponse> {
246        if request.intent.trim().is_empty() {
247            return Ok(PrefetchResponse {
248                status: PrefetchStatus::SkippedNoIntent,
249                prompt_block: None,
250            });
251        }
252
253        Ok(PrefetchResponse {
254            status: PrefetchStatus::Failed {
255                reason: format!(
256                    "prefetch recall is unavailable in the default MemoryManager for intent '{}'",
257                    request.intent.trim()
258                ),
259            },
260            prompt_block: None,
261        })
262    }
263
264    async fn sync_turn(&self, request: SyncTurnRequest) -> crate::Result<SyncTurnResponse> {
265        let mut persisted = false;
266
267        if let Some(memory_update) = request.memory_update_markdown.as_deref() {
268            if !memory_update.trim().is_empty() {
269                let memory = Memory::with_content(memory_update);
270                if let Err(err) = self.save_memory(&memory) {
271                    return Ok(SyncTurnResponse {
272                        status: SyncTurnStatus::Failed {
273                            reason: format!("failed to persist MEMORY.md: {err}"),
274                        },
275                    });
276                }
277                persisted = true;
278            }
279        }
280
281        if let Some(history_entry) = request.history_entry.as_deref() {
282            if !history_entry.trim().is_empty() {
283                if let Err(err) = self.append_history(history_entry) {
284                    return Ok(SyncTurnResponse {
285                        status: SyncTurnStatus::Failed {
286                            reason: format!("failed to append HISTORY.md: {err}"),
287                        },
288                    });
289                }
290                persisted = true;
291            }
292        }
293
294        Ok(SyncTurnResponse {
295            status: if persisted {
296                SyncTurnStatus::Persisted
297            } else {
298                SyncTurnStatus::Noop
299            },
300        })
301    }
302
303    async fn on_session_end(
304        &self,
305        request: SessionEndRequest,
306    ) -> crate::Result<SessionEndResponse> {
307        if let Some(session_id) = request.session_id {
308            let mut handled = self.handled_session_end_ids.lock();
309            if !handled.insert(session_id) {
310                return Ok(SessionEndResponse {
311                    status: SessionEndStatus::AlreadyHandled,
312                });
313            }
314        }
315
316        Ok(SessionEndResponse {
317            status: SessionEndStatus::Noop,
318        })
319    }
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325    use crate::memory::{
326        MemoryProvider, PrefetchRequest, PrefetchStatus, SessionEndRequest, SessionEndStatus,
327        StartupStatus, SyncTurnRequest, SyncTurnStatus, SystemPromptRequest,
328    };
329    use tempfile::TempDir;
330
331    #[test]
332    fn test_append_history() {
333        let temp_dir = TempDir::new().unwrap();
334        let manager = MemoryManager::new(temp_dir.path());
335        manager
336            .append_history("[2026-02-12 09:00] Added memory event")
337            .unwrap();
338
339        let history = manager.load_history();
340        assert!(history.contains("Added memory event"));
341    }
342
343    #[tokio::test]
344    async fn test_memory_provider_system_prompt_block_reads_memory() {
345        let temp_dir = TempDir::new().unwrap();
346        let manager = MemoryManager::new(temp_dir.path());
347        manager
348            .save_memory(&Memory::with_content("Long term provider context"))
349            .unwrap();
350
351        let response = manager
352            .system_prompt_block(&SystemPromptRequest {
353                workspace_root: temp_dir.path().to_path_buf(),
354            })
355            .unwrap();
356
357        assert_eq!(response.status, StartupStatus::Ready);
358
359        let block = response
360            .prompt_block
361            .expect("memory-backed provider should emit a prompt block");
362
363        assert!(block.markdown.contains("## Long-term Memory"));
364        assert!(block.markdown.contains("Long term provider context"));
365        assert_eq!(
366            block.shape,
367            crate::memory::StartupInjectionShape::CompactRenderedMarkdown
368        );
369    }
370
371    #[tokio::test]
372    async fn test_memory_provider_sync_turn_persists_memory_and_history() {
373        let temp_dir = TempDir::new().unwrap();
374        let manager = MemoryManager::new(temp_dir.path());
375
376        let result = manager
377            .sync_turn(SyncTurnRequest {
378                workspace_root: temp_dir.path().to_path_buf(),
379                memory_update_markdown: Some("Updated memory from sync_turn".to_string()),
380                history_entry: Some("[2026-05-08 10:00 UTC] synchronized turn".to_string()),
381            })
382            .await
383            .unwrap();
384
385        assert_eq!(result.status, SyncTurnStatus::Persisted);
386        assert_eq!(manager.load_memory().content, "Updated memory from sync_turn");
387        assert!(manager.load_history().contains("synchronized turn"));
388    }
389
390    #[tokio::test]
391    async fn test_memory_provider_session_end_is_noop_by_default() {
392        let temp_dir = TempDir::new().unwrap();
393        let manager = MemoryManager::new(temp_dir.path());
394
395        let response = manager
396            .on_session_end(SessionEndRequest {
397                workspace_root: temp_dir.path().to_path_buf(),
398                session_id: Some("session-1".to_string()),
399            })
400            .await
401            .unwrap();
402
403        assert_eq!(response.status, SessionEndStatus::Noop);
404    }
405
406    #[tokio::test]
407    async fn test_memory_provider_returns_degraded_startup_when_no_context_is_available() {
408        let temp_dir = TempDir::new().unwrap();
409        let manager = MemoryManager::new(temp_dir.path());
410
411        let response = manager
412            .system_prompt_block(&SystemPromptRequest {
413                workspace_root: temp_dir.path().to_path_buf(),
414            })
415            .unwrap();
416
417        match response.status {
418            StartupStatus::Degraded {
419                reason,
420                last_usable_wakeup,
421            } => {
422                assert!(reason.contains("startup continuity unavailable"));
423                assert!(last_usable_wakeup.is_none());
424            }
425            other => panic!("expected degraded startup, got {other:?}"),
426        }
427
428        let block = response
429            .prompt_block
430            .expect("degraded startup should still provide explicit startup text");
431        assert!(block.markdown.contains("status: degraded"));
432        assert!(block.markdown.contains("last_usable_wakeup: omitted"));
433    }
434
435    #[tokio::test]
436    async fn test_memory_provider_prefetch_distinguishes_no_intent_from_failed_recall() {
437        let temp_dir = TempDir::new().unwrap();
438        let manager = MemoryManager::new(temp_dir.path());
439
440        let skipped = manager
441            .prefetch(PrefetchRequest {
442                workspace_root: temp_dir.path().to_path_buf(),
443                intent: "   ".to_string(),
444                current_room: None,
445                user_message: Some("help".to_string()),
446            })
447            .await
448            .unwrap();
449        assert_eq!(skipped.status, PrefetchStatus::SkippedNoIntent);
450        assert!(skipped.prompt_block.is_none());
451
452        let failed = manager
453            .prefetch(PrefetchRequest {
454                workspace_root: temp_dir.path().to_path_buf(),
455                intent: "recall-project-status".to_string(),
456                current_room: Some("roadmap".to_string()),
457                user_message: Some("what changed?".to_string()),
458            })
459            .await
460            .unwrap();
461        match failed.status {
462            PrefetchStatus::Failed { reason } => {
463                assert!(reason.contains("prefetch recall is unavailable"));
464            }
465            other => panic!("expected failed recall, got {other:?}"),
466        }
467        assert!(failed.prompt_block.is_none());
468    }
469
470    #[tokio::test]
471    async fn test_memory_provider_sync_turn_failure_is_explicit() {
472        let temp_dir = TempDir::new().unwrap();
473        let workspace = temp_dir.path().to_path_buf();
474        std::fs::create_dir_all(workspace.join("memory")).unwrap();
475        std::fs::write(workspace.join("memory").join("MEMORY.md"), "locked").unwrap();
476        std::fs::remove_file(workspace.join("memory").join("MEMORY.md")).unwrap();
477        std::fs::create_dir_all(workspace.join("memory").join("MEMORY.md")).unwrap();
478
479        let manager = MemoryManager::new(&workspace);
480        let result = manager
481            .sync_turn(SyncTurnRequest {
482                workspace_root: workspace,
483                memory_update_markdown: Some("cannot persist".to_string()),
484                history_entry: None,
485            })
486            .await
487            .unwrap();
488
489        match result.status {
490            SyncTurnStatus::Failed { reason } => {
491                assert!(reason.contains("failed to persist MEMORY.md"));
492            }
493            other => panic!("expected sync failure, got {other:?}"),
494        }
495    }
496
497    #[tokio::test]
498    async fn test_memory_provider_session_end_is_idempotent_for_duplicates() {
499        let temp_dir = TempDir::new().unwrap();
500        let manager = MemoryManager::new(temp_dir.path());
501
502        let first = manager
503            .on_session_end(SessionEndRequest {
504                workspace_root: temp_dir.path().to_path_buf(),
505                session_id: Some("session-dup".to_string()),
506            })
507            .await
508            .unwrap();
509        assert_eq!(first.status, SessionEndStatus::Noop);
510
511        let duplicate = manager
512            .on_session_end(SessionEndRequest {
513                workspace_root: temp_dir.path().to_path_buf(),
514                session_id: Some("session-dup".to_string()),
515            })
516            .await
517            .unwrap();
518        assert_eq!(duplicate.status, SessionEndStatus::AlreadyHandled);
519    }
520}
agent_diva_core/memory/manager.rs

agent_diva_core/memory/
manager.rs
