Skip to main content

oxios_memory/memory/manager/
mod.rs

1//! Memory manager — the central orchestrator for agent memory.
2//!
3//! `MemoryManager` coordinates CRUD, indexing, search, and lifecycle
4//! operations for memory entries. It wraps an abstract `MemoryStorage`
5//! backend and optional SQLite store, HNSW index, git layer, and SONA engine.
6//!
7//! All kernel-coupled types (`StateStore`, `GitLayer`, `MemoryConfig`) are
8//! accessed through traits defined in [`crate::storage`]. The kernel
9//! implements those traits and injects concrete instances.
10
11mod ops;
12mod store;
13
14use std::collections::HashMap;
15use std::sync::Arc;
16
17use anyhow::Result;
18use parking_lot::RwLock;
19
20use crate::memory::embedding::{EmbeddingProvider, EmbeddingVector, TfIdfEmbeddingProvider};
21use crate::memory::hnsw_memory_index::HnswMemoryIndex;
22use crate::memory::sona::SonaEngine;
23use crate::memory::storage::{MemoryGit, MemoryStorage};
24use crate::memory::types::{MemoryEntry, MemoryType};
25
26use super::{CurationCandidate, CurationReport, MemoryBudget};
27
28// ---------------------------------------------------------------------------
29// MemoryManager
30// ---------------------------------------------------------------------------
31
32/// Agent memory manager.
33///
34/// Stores and retrieves memory entries using a pluggable storage backend.
35/// Supports embedding-based vector search via an in-memory TF-IDF index
36/// that is rebuilt on startup.
37pub struct MemoryManager {
38    /// Storage backend (typically `StateStore` from kernel).
39    pub(crate) storage: Arc<dyn MemoryStorage>,
40    /// Maximum memories returned by recall.
41    pub(crate) max_recall: usize,
42    /// Vector index for semantic search (id → EmbeddingVector).
43    pub(crate) vector_index: RwLock<HashMap<String, EmbeddingVector>>,
44    /// Embedding provider for generating vectors.
45    pub(crate) embedding: Arc<dyn EmbeddingProvider>,
46    /// Optional git layer for version-controlled memory.
47    pub(crate) git: Option<Arc<dyn MemoryGit>>,
48    /// Optional HNSW index for fast ANN search.
49    pub(crate) hnsw_index: RwLock<Option<Arc<HnswMemoryIndex>>>,
50    /// Optional SONA learning engine (RFC-020 Phase 2).
51    pub(crate) sona_engine: Option<Arc<SonaEngine>>,
52    /// Optional SQLite-backed store (RFC-012).
53    #[cfg(feature = "sqlite-memory")]
54    pub(crate) sqlite_store: Option<Arc<crate::memory::sqlite::SqliteMemoryStore>>,
55}
56
57impl std::fmt::Debug for MemoryManager {
58    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
59        f.debug_struct("MemoryManager")
60            .field("max_recall", &self.max_recall)
61            .field("index_size", &self.vector_index.read().len())
62            .field("sona_enabled", &self.sona_engine.is_some())
63            .finish()
64    }
65}
66
67impl MemoryManager {
68    /// Create a new MemoryManager with a storage backend.
69    pub fn new(storage: Arc<dyn MemoryStorage>) -> Self {
70        Self {
71            storage,
72            max_recall: 10,
73            vector_index: RwLock::new(HashMap::new()),
74            embedding: Arc::new(TfIdfEmbeddingProvider),
75            git: None,
76            hnsw_index: RwLock::new(None),
77            sona_engine: None,
78            #[cfg(feature = "sqlite-memory")]
79            sqlite_store: None,
80        }
81    }
82
83    /// Attach a git layer for version-controlled saves.
84    pub fn set_git_layer(&mut self, gl: Arc<dyn MemoryGit>) {
85        self.git = Some(gl);
86    }
87
88    /// Attach a SQLite-backed memory store (RFC-012).
89    #[cfg(feature = "sqlite-memory")]
90    pub fn set_sqlite_store(&mut self, store: Arc<crate::memory::sqlite::SqliteMemoryStore>) {
91        self.sqlite_store = Some(store);
92    }
93
94    /// Get a reference to the SQLite store (if configured).
95    #[cfg(feature = "sqlite-memory")]
96    pub fn sqlite_store(&self) -> &Option<Arc<crate::memory::sqlite::SqliteMemoryStore>> {
97        &self.sqlite_store
98    }
99
100    /// Attach a SONA learning engine (RFC-020 Phase 2).
101    pub fn set_sona_engine(&mut self, engine: Arc<SonaEngine>) {
102        self.sona_engine = Some(engine);
103    }
104
105    /// Get a reference to the SONA engine (if configured).
106    pub fn sona_engine(&self) -> Option<&Arc<SonaEngine>> {
107        self.sona_engine.as_ref()
108    }
109
110    /// Attach an HNSW index for fast semantic search.
111    pub fn set_hnsw_index(&self, index: Arc<HnswMemoryIndex>) {
112        *self.hnsw_index.write() = Some(index);
113    }
114
115    /// Set max memories returned by recall.
116    pub fn with_max_recall(mut self, n: usize) -> Self {
117        self.max_recall = n;
118        self
119    }
120
121    /// Set max_recall in-place.
122    pub fn set_max_recall(&mut self, n: usize) {
123        self.max_recall = n;
124    }
125
126    /// Returns the number of entries in the vector index.
127    pub fn vector_index_size(&self) -> usize {
128        self.vector_index.read().len()
129    }
130
131    /// Commit a file to git if git_layer is available.
132    pub(crate) async fn git_commit(&self, rel_path: &str, message: &str) {
133        if let Some(ref gl) = self.git {
134            if gl.is_enabled() {
135                let _ = gl.commit_file(rel_path, message).await;
136            }
137        }
138    }
139
140    /// Compute effective importance of a memory entry.
141    pub fn effective_importance(entry: &MemoryEntry) -> f32 {
142        let access_boost = (1.0_f32 + entry.access_count as f32).ln();
143        entry.importance * (1.0 + access_boost)
144    }
145
146    /// Curate memories: identify candidates for removal based on budget.
147    pub async fn curate(&self, budget: &MemoryBudget) -> Result<CurationReport> {
148        let mut report = CurationReport::default();
149
150        for mt in &[
151            MemoryType::Conversation,
152            MemoryType::Session,
153            MemoryType::Fact,
154            MemoryType::Episode,
155            MemoryType::Knowledge,
156        ] {
157            let entries = self.list(*mt, budget.max_per_type * 2).await?;
158            if entries.len() <= budget.max_per_type {
159                continue;
160            }
161
162            let total_count = entries.len();
163            let mut scored: Vec<_> = entries
164                .into_iter()
165                .map(|e| (e.id.clone(), e.memory_type, Self::effective_importance(&e)))
166                .collect();
167            scored.sort_by(|a, b| a.2.partial_cmp(&b.2).unwrap_or(std::cmp::Ordering::Equal));
168
169            let to_remove = scored.len() - budget.max_per_type;
170            for (id, memory_type, score) in scored.into_iter().take(to_remove) {
171                report.candidates_for_removal.push(CurationCandidate {
172                    id,
173                    memory_type,
174                    effective_importance: score,
175                });
176            }
177            report.total_before += total_count;
178        }
179
180        for candidate in &report.candidates_for_removal {
181            if self
182                .forget(&candidate.id, candidate.memory_type)
183                .await
184                .is_ok()
185            {
186                report.removed += 1;
187            }
188        }
189
190        report.total_after = report.total_before - report.removed;
191        Ok(report)
192    }
193
194    /// Spawn a background curation task.
195    pub fn spawn_curation_task(self: &Arc<Self>, budget: MemoryBudget) {
196        let mgr = Arc::clone(self);
197        tokio::spawn(async move {
198            match mgr.curate(&budget).await {
199                Ok(report) => {
200                    if report.removed > 0 {
201                        tracing::info!(
202                            removed = report.removed,
203                            candidates = report.candidates_for_removal.len(),
204                            "Memory curation complete"
205                        );
206                    }
207                }
208                Err(e) => {
209                    tracing::warn!(error = %e, "Memory curation failed");
210                }
211            }
212        });
213    }
214}