agent_memo/
lib.rs

1//! `agent-memo` — the **on-device / local** context store, backed by
2//! **aria memo**.
3//!
4//! The contract ([`ContextStore`], [`ContextFragment`], [`RecallQuery`], the
5//! local [`embed`] module) lives in `aria-agent-core::context`; this crate is
6//! the local/embedded backend used by the native SDK (Swift / Kotlin via
7//! UniFFI) and by the `local` side of the `both` backend.
8//!
9//! Storage is **aria memo compatible**: the same `memories` table, the same
10//! `memo_type` strings, the same embedding BLOB encoding and the same
11//! `metadata` JSON shape as the `aria-memo` product / CLI, so a `memo.db`
12//! written here can be read and edited with `aria-memo list --json`
13//! (and vice versa) — see the `cli_interop_*` tests below.
14//!
15//! The previous sled implementation has been removed; there is no second
16//! on-device format to keep in sync. See `docs/adr/0011-memory-backend-switch.md`.
17
18use agent_core::context::embed::Embedder as _;
19use agent_core::context::{
20    embed, ensure_embedding, rank, ContextError, ContextFragment, ContextStore, FragmentKind,
21    RecallQuery, EMBED_DIM,
22};
23use async_trait::async_trait;
24use rusqlite::Connection;
25use std::path::Path;
26use std::sync::{Arc, Mutex};
27
28pub use agent_core::context::{now_ms, ContextError as MemoError};
29
30/// aria memo's `memories` schema (byte-for-byte the DDL used by the product),
31/// so its CLI can open the same database.
32const SCHEMA: &str = "
33CREATE TABLE IF NOT EXISTS memories (
34    id TEXT PRIMARY KEY,
35    memo_type TEXT NOT NULL,
36    content TEXT NOT NULL,
37    embedding BLOB,
38    metadata TEXT NOT NULL,
39    importance REAL NOT NULL,
40    version INTEGER NOT NULL,
41    created_at INTEGER NOT NULL,
42    updated_at INTEGER NOT NULL,
43    deleted INTEGER NOT NULL DEFAULT 0
44);
45CREATE INDEX IF NOT EXISTS idx_memories_type ON memories(memo_type);
46CREATE INDEX IF NOT EXISTS idx_memories_updated_at ON memories(updated_at);
47CREATE INDEX IF NOT EXISTS idx_memories_deleted ON memories(deleted);
48";
49
50/// `aria-memo` writes `long_term:*` for durable memories; the mapping below
51/// keeps agent fragment kinds and memo types reversible.
52fn memo_type_of(kind: FragmentKind) -> &'static str {
53    match kind {
54        FragmentKind::Message => "working",
55        FragmentKind::ToolResult => "short_term",
56        FragmentKind::LongTerm => "long_term:semantic",
57        FragmentKind::Note => "long_term:episodic",
58    }
59}
60
61fn kind_of(memo_type: &str) -> Option<FragmentKind> {
62    match memo_type {
63        "working" => Some(FragmentKind::Message),
64        "short_term" => Some(FragmentKind::ToolResult),
65        "long_term:semantic" | "long_term:entity" | "long_term:graph" => {
66            Some(FragmentKind::LongTerm)
67        }
68        "long_term:episodic" => Some(FragmentKind::Note),
69        _ => None,
70    }
71}
72
73fn importance_of(kind: FragmentKind) -> f32 {
74    match kind {
75        FragmentKind::LongTerm => 0.8,
76        FragmentKind::Message => 0.5,
77        FragmentKind::ToolResult => 0.4,
78        FragmentKind::Note => 0.6,
79    }
80}
81
82/// aria memo's embedding BLOB: empty = `None`, otherwise `u32 LE` length
83/// followed by little-endian `f32`s.
84fn serialize_embedding(emb: Option<&[f32]>) -> Vec<u8> {
85    match emb {
86        None => Vec::new(),
87        Some(v) => {
88            let mut buf = Vec::with_capacity(4 + v.len() * 4);
89            buf.extend_from_slice(&(v.len() as u32).to_le_bytes());
90            for f in v {
91                buf.extend_from_slice(&f.to_le_bytes());
92            }
93            buf
94        }
95    }
96}
97
98fn deserialize_embedding(buf: &[u8]) -> Option<Vec<f32>> {
99    if buf.is_empty() {
100        return None;
101    }
102    if buf.len() < 4 {
103        return None;
104    }
105    let n = u32::from_le_bytes([buf[0], buf[1], buf[2], buf[3]]) as usize;
106    if buf.len() != 4 + n * 4 {
107        return None;
108    }
109    let mut v = Vec::with_capacity(n);
110    for i in 0..n {
111        let s = 4 + i * 4;
112        v.push(f32::from_le_bytes([
113            buf[s],
114            buf[s + 1],
115            buf[s + 2],
116            buf[s + 3],
117        ]));
118    }
119    Some(v)
120}
121
122fn metadata_json(session: &str, key: Option<&str>) -> String {
123    let mut map = std::collections::HashMap::new();
124    map.insert("session".to_string(), session.to_string());
125    if let Some(k) = key {
126        map.insert("key".to_string(), k.to_string());
127    }
128    serde_json::to_string(&map).unwrap_or_else(|_| "{}".to_string())
129}
130
131/// aria memo backed context store (SQLite, local / on-device).
132pub struct MemoContextStore {
133    conn: Mutex<Connection>,
134}
135
136impl MemoContextStore {
137    /// Open (or create) the aria memo database at `path`.
138    pub fn open(path: &Path) -> Result<Arc<Self>, ContextError> {
139        if let Some(parent) = path.parent() {
140            if !parent.as_os_str().is_empty() {
141                std::fs::create_dir_all(parent)
142                    .map_err(|e| ContextError::Storage(format!("create db dir: {e}")))?;
143            }
144        }
145        let conn = Connection::open(path)
146            .map_err(|e| ContextError::Storage(format!("open memo db: {e}")))?;
147        conn.execute_batch(SCHEMA)
148            .map_err(|e| ContextError::Storage(format!("migrate memo db: {e}")))?;
149        Ok(Arc::new(Self {
150            conn: Mutex::new(conn),
151        }))
152    }
153
154    /// Throwaway database (tests, and the SDK default when no path is set).
155    pub fn memory() -> Result<Arc<Self>, ContextError> {
156        let dir = std::env::temp_dir().join(format!("aria-memo-{}", uuid::Uuid::new_v4()));
157        std::fs::create_dir_all(&dir)
158            .map_err(|e| ContextError::Storage(format!("create temp db dir: {e}")))?;
159        let path = dir.join("memo.db");
160        let store = Self::open(&path)?;
161        // Keep the directory around for the process lifetime (removed on reboot);
162        // tests do not rely on cleanup.
163        Ok(store)
164    }
165
166    fn all(&self) -> Result<Vec<ContextFragment>, ContextError> {
167        let conn = self
168            .conn
169            .lock()
170            .map_err(|e| ContextError::Storage(format!("memo db lock poisoned: {e}")))?;
171        let mut stmt = conn
172            .prepare(
173                "SELECT id, memo_type, content, embedding, metadata, created_at \
174                 FROM memories WHERE deleted = 0 ORDER BY created_at",
175            )
176            .map_err(|e| ContextError::Storage(e.to_string()))?;
177        let rows = stmt
178            .query_map([], |row| {
179                let emb: Vec<u8> = row.get(3)?;
180                let meta: String = row.get(4)?;
181                Ok((
182                    row.get::<_, String>(0)?,
183                    row.get::<_, String>(1)?,
184                    row.get::<_, String>(2)?,
185                    deserialize_embedding(&emb),
186                    meta,
187                    row.get::<_, i64>(5)?,
188                ))
189            })
190            .map_err(|e| ContextError::Storage(e.to_string()))?;
191        let mut out = Vec::new();
192        for r in rows {
193            let (id, memo_type, content, embedding, meta, created_at) =
194                r.map_err(|e| ContextError::Storage(e.to_string()))?;
195            let Some(kind) = kind_of(&memo_type) else {
196                continue;
197            };
198            let parsed: std::collections::HashMap<String, String> =
199                serde_json::from_str(&meta).unwrap_or_default();
200            let session = parsed.get("session").cloned().unwrap_or_default();
201            let key = parsed.get("key").cloned();
202            out.push(ContextFragment {
203                id,
204                session,
205                key,
206                kind,
207                content,
208                created_at,
209                embedding,
210            });
211        }
212        Ok(out)
213    }
214}
215
216#[async_trait]
217impl ContextStore for MemoContextStore {
218    async fn memorize(&self, mut frag: ContextFragment) -> Result<(), ContextError> {
219        ensure_embedding(&mut frag);
220        let embedding = serialize_embedding(frag.embedding.as_deref());
221        let now_sec = now_ms() / 1000;
222        let meta = metadata_json(&frag.session, frag.key.as_deref());
223        let conn = self
224            .conn
225            .lock()
226            .map_err(|e| ContextError::Storage(format!("memo db lock poisoned: {e}")))?;
227        conn.execute(
228            "INSERT INTO memories \
229             (id, memo_type, content, embedding, metadata, importance, version, \
230              created_at, updated_at, deleted) \
231             VALUES (?1, ?2, ?3, ?4, ?5, ?6, 1, ?7, ?8, 0) \
232             ON CONFLICT(id) DO UPDATE SET \
233                memo_type = excluded.memo_type, \
234                content = excluded.content, \
235                embedding = excluded.embedding, \
236                metadata = excluded.metadata, \
237                importance = excluded.importance, \
238                updated_at = excluded.updated_at, \
239                deleted = 0",
240            rusqlite::params![
241                frag.id,
242                memo_type_of(frag.kind),
243                frag.content,
244                embedding,
245                meta,
246                importance_of(frag.kind) as f64,
247                now_sec,
248                now_sec,
249            ],
250        )
251        .map_err(|e| ContextError::Storage(e.to_string()))?;
252        Ok(())
253    }
254
255    async fn recall(&self, query: &RecallQuery) -> Result<Vec<ContextFragment>, ContextError> {
256        if query.text.trim().is_empty() {
257            return Ok(Vec::new());
258        }
259        let candidates: Vec<ContextFragment> = self
260            .all()?
261            .into_iter()
262            .filter(|f| f.session == query.session)
263            .filter(|f| query.kind.map(|k| f.kind == k).unwrap_or(true))
264            .collect();
265        let q_emb = embed::LocalEmbedder::new(EMBED_DIM).embed(&query.text).ok();
266        Ok(rank(&query.text, q_emb.as_deref(), candidates, query.top_k))
267    }
268
269    async fn compact(&self, session: &str) -> Result<ContextFragment, ContextError> {
270        let parts: Vec<String> = self
271            .all()?
272            .into_iter()
273            .filter(|f| f.session == session)
274            .map(|f| format!("[{}] {}", f.kind.as_str(), f.content))
275            .collect();
276        if parts.is_empty() {
277            return Err(ContextError::NotFound(session.to_string()));
278        }
279        let merged = ContextFragment::new(session, FragmentKind::Note, parts.join("\n---\n"))
280            .with_key(format!("__compact__{session}"));
281        self.memorize(merged.clone()).await?;
282        Ok(merged)
283    }
284
285    async fn get_by_key(
286        &self,
287        session: &str,
288        key: &str,
289    ) -> Result<Option<ContextFragment>, ContextError> {
290        Ok(self
291            .all()?
292            .into_iter()
293            .filter(|f| f.session == session && f.key.as_deref() == Some(key))
294            .max_by_key(|f| f.created_at))
295    }
296
297    async fn list_session(
298        &self,
299        session: &str,
300        top_k: usize,
301    ) -> Result<Vec<ContextFragment>, ContextError> {
302        let mut out: Vec<ContextFragment> = self
303            .all()?
304            .into_iter()
305            .filter(|f| f.session == session)
306            .collect();
307        out.sort_by_key(|f| f.created_at);
308        out.truncate(top_k);
309        Ok(out)
310    }
311}
312
313#[cfg(test)]
314mod tests {
315    use super::*;
316
317    fn store() -> Arc<MemoContextStore> {
318        MemoContextStore::memory().unwrap()
319    }
320
321    #[tokio::test]
322    async fn memorize_recall_roundtrip() {
323        let store = store();
324        let f =
325            ContextFragment::new("s1", FragmentKind::Message, "the sky is blue").with_key("fact1");
326        store.memorize(f).await.unwrap();
327        let out = store.recall(&RecallQuery::new("s1", "sky")).await.unwrap();
328        assert!(out.iter().any(|f| f.content.contains("sky")));
329        assert!(store.get_by_key("s1", "fact1").await.unwrap().is_some());
330    }
331
332    #[tokio::test]
333    async fn compact_merges_session() {
334        let store = store();
335        store
336            .memorize(ContextFragment::new("s2", FragmentKind::Message, "a"))
337            .await
338            .unwrap();
339        store
340            .memorize(ContextFragment::new("s2", FragmentKind::Message, "b"))
341            .await
342            .unwrap();
343        let c = store.compact("s2").await.unwrap();
344        assert!(c.content.contains("a") && c.content.contains("b"));
345        assert!(matches!(
346            store.compact("ghost").await,
347            Err(ContextError::NotFound(_))
348        ));
349    }
350
351    #[tokio::test]
352    async fn vector_recall_prefers_similar() {
353        let store = store();
354        store
355            .memorize(ContextFragment::new(
356                "s3",
357                FragmentKind::Message,
358                "user prefers rust for systems programming",
359            ))
360            .await
361            .unwrap();
362        store
363            .memorize(ContextFragment::new(
364                "s3",
365                FragmentKind::Message,
366                "banana smoothie recipe with ice",
367            ))
368            .await
369            .unwrap();
370        let frags = store
371            .recall(&RecallQuery::new("s3", "rust programming language"))
372            .await
373            .unwrap();
374        assert!(!frags.is_empty());
375        assert!(frags[0].content.contains("rust"));
376        assert!(frags[0].embedding.is_some());
377    }
378
379    #[tokio::test]
380    async fn recall_filters_kind_and_session() {
381        let store = store();
382        store
383            .memorize(ContextFragment::new(
384                "a",
385                FragmentKind::Message,
386                "shared text",
387            ))
388            .await
389            .unwrap();
390        store
391            .memorize(ContextFragment::new(
392                "b",
393                FragmentKind::Message,
394                "shared text",
395            ))
396            .await
397            .unwrap();
398        let out = store
399            .recall(&RecallQuery::new("a", "shared"))
400            .await
401            .unwrap();
402        assert!(out.iter().all(|f| f.session == "a"));
403        assert_eq!(out.len(), 1);
404
405        store
406            .memorize(ContextFragment::new("a", FragmentKind::Note, "shared text"))
407            .await
408            .unwrap();
409        let notes = store
410            .recall(&RecallQuery::new("a", "shared").with_kind(FragmentKind::Note))
411            .await
412            .unwrap();
413        assert_eq!(notes.len(), 1);
414        assert_eq!(notes[0].kind, FragmentKind::Note);
415    }
416
417    #[tokio::test]
418    async fn empty_query_recalls_nothing() {
419        let store = store();
420        store
421            .memorize(ContextFragment::new("s", FragmentKind::Message, "hello"))
422            .await
423            .unwrap();
424        assert!(store
425            .recall(&RecallQuery::new("s", ""))
426            .await
427            .unwrap()
428            .is_empty());
429    }
430
431    #[tokio::test]
432    async fn memorize_populates_embedding_of_fixed_dim() {
433        let store = store();
434        store
435            .memorize(
436                ContextFragment::new("s", FragmentKind::Message, "rust programming").with_key("ke"),
437            )
438            .await
439            .unwrap();
440        let got = store.get_by_key("s", "ke").await.unwrap().unwrap();
441        assert_eq!(got.embedding.as_ref().unwrap().len(), EMBED_DIM);
442    }
443
444    #[tokio::test]
445    async fn upsert_replaces_existing_id() {
446        let store = store();
447        let mut frag = ContextFragment::new("s", FragmentKind::LongTerm, "v1");
448        frag.id = "fixed".into();
449        store.memorize(frag.clone()).await.unwrap();
450        let mut updated = ContextFragment::new("s", FragmentKind::LongTerm, "v2");
451        updated.id = "fixed".into();
452        store.memorize(updated).await.unwrap();
453        let all = store.all().unwrap();
454        assert_eq!(all.len(), 1);
455        assert_eq!(all[0].content, "v2");
456    }
457
458    #[test]
459    fn embedding_blob_roundtrip_matches_aria_memo_format() {
460        let v = vec![0.5f32, -0.25, 1.0];
461        let blob = serialize_embedding(Some(&v));
462        // 4-byte LE length prefix + LE floats.
463        assert_eq!(&blob[..4], &3u32.to_le_bytes());
464        assert_eq!(deserialize_embedding(&blob), Some(v));
465        assert_eq!(deserialize_embedding(&[]), None);
466    }
467
468    #[test]
469    fn memo_type_mapping_is_reversible() {
470        for kind in [
471            FragmentKind::Message,
472            FragmentKind::ToolResult,
473            FragmentKind::LongTerm,
474            FragmentKind::Note,
475        ] {
476            assert_eq!(kind_of(memo_type_of(kind)), Some(kind));
477        }
478        assert_eq!(kind_of("unknown-type"), None);
479    }
480
481    /// Skip helper: the interop tests need the real `aria-memo` binary.
482    fn aria_memo_bin() -> Option<String> {
483        std::env::var("ARIA_MEMO_BIN")
484            .ok()
485            .filter(|s| !s.trim().is_empty())
486            .or_else(|| Some("aria-memo".to_string()))
487            .filter(|bin| {
488                std::process::Command::new(bin)
489                    .arg("version")
490                    .output()
491                    .map(|o| o.status.success())
492                    .unwrap_or(false)
493            })
494    }
495
496    #[tokio::test]
497    async fn cli_interop_rust_write_is_readable_by_aria_memo() {
498        let Some(bin) = aria_memo_bin() else {
499            eprintln!("skip: aria-memo CLI not installed");
500            return;
501        };
502        let dir = std::env::temp_dir().join(format!("aria-memo-interop-{}", uuid::Uuid::new_v4()));
503        std::fs::create_dir_all(&dir).unwrap();
504        let db = dir.join("memo.db");
505        let store = MemoContextStore::open(&db).unwrap();
506        store
507            .memorize(ContextFragment::new(
508                "s",
509                FragmentKind::LongTerm,
510                "rust interop fact",
511            ))
512            .await
513            .unwrap();
514
515        let out = std::process::Command::new(bin)
516            .args(["--db"]) // keep the flag and value separate
517            .arg(&db)
518            .args(["list", "--json"])
519            .output()
520            .expect("run aria-memo list");
521        assert!(
522            out.status.success(),
523            "aria-memo list failed: {:?}",
524            out.stderr
525        );
526        let listed = String::from_utf8_lossy(&out.stdout);
527        assert!(listed.contains("rust interop fact"), "got: {listed}");
528        let _ = std::fs::remove_dir_all(&dir);
529    }
530
531    #[tokio::test]
532    async fn cli_interop_aria_memo_write_is_readable_by_rust() {
533        let Some(bin) = aria_memo_bin() else {
534            eprintln!("skip: aria-memo CLI not installed");
535            return;
536        };
537        let dir = std::env::temp_dir().join(format!("aria-memo-interop2-{}", uuid::Uuid::new_v4()));
538        std::fs::create_dir_all(&dir).unwrap();
539        let db = dir.join("memo.db");
540        let out = std::process::Command::new(bin)
541            .arg("--db")
542            .arg(&db)
543            .args([
544                "add",
545                "--type",
546                "long_term:semantic",
547                "--content",
548                "cli wrote this fact",
549                "--importance",
550                "0.8",
551            ])
552            .output()
553            .expect("run aria-memo add");
554        assert!(
555            out.status.success(),
556            "aria-memo add failed: {:?}",
557            out.stderr
558        );
559
560        let store = MemoContextStore::open(&db).unwrap();
561        let all = store.all().unwrap();
562        assert!(
563            all.iter().any(|f| f.content == "cli wrote this fact"),
564            "aria memo rows must be visible to the Rust store: {all:?}"
565        );
566        let _ = std::fs::remove_dir_all(&dir);
567    }
568}
agent_memo/lib.rs

agent_memo/
lib.rs
