Skip to main content

mnemo_core/storage/
duckdb.rs

1use std::path::Path;
2use std::sync::Arc;
3use tokio::sync::Mutex;
4
5use crate::error::{Error, Result};
6use crate::model::acl::{Acl, Permission};
7use crate::model::agent_profile::AgentProfile;
8use crate::model::checkpoint::Checkpoint;
9use crate::model::delegation::{Delegation, DelegationScope};
10use crate::model::embedding_baseline::EmbeddingBaseline;
11use crate::model::event::AgentEvent;
12use crate::model::memory::MemoryRecord;
13use crate::model::relation::Relation;
14use crate::model::write_provenance::{
15    WriteOp, WriteProvenance, flags_from_storage, flags_to_storage,
16};
17use crate::storage::{MemoryFilter, StorageBackend};
18use uuid::Uuid;
19
20pub struct DuckDbStorage {
21    conn: Arc<Mutex<duckdb::Connection>>,
22}
23
24impl DuckDbStorage {
25    pub fn open(path: &Path) -> Result<Self> {
26        let conn = duckdb::Connection::open(path)?;
27        super::migrations::run_migrations(&conn)?;
28        Ok(Self {
29            conn: Arc::new(Mutex::new(conn)),
30        })
31    }
32
33    pub fn open_in_memory() -> Result<Self> {
34        let conn = duckdb::Connection::open_in_memory()?;
35        super::migrations::run_migrations(&conn)?;
36        Ok(Self {
37            conn: Arc::new(Mutex::new(conn)),
38        })
39    }
40}
41
42fn serialize_embedding(embedding: &Option<Vec<f32>>) -> Option<Vec<u8>> {
43    embedding
44        .as_ref()
45        .map(|v| v.iter().flat_map(|f| f.to_le_bytes()).collect())
46}
47
48fn deserialize_embedding(blob: Option<Vec<u8>>) -> Option<Vec<f32>> {
49    blob.map(|bytes| {
50        // `as_chunks::<4>()` rather than `chunks_exact(4)`: it yields `&[u8; 4]`
51        // directly, so `from_le_bytes` takes the array instead of four indexed
52        // reads. Required by clippy 1.98's `chunks_exact_to_as_chunks`, and the
53        // result is the better code anyway - the indexing form carried four
54        // bounds checks the type system can prove unnecessary.
55        bytes
56            .as_chunks::<4>()
57            .0
58            .iter()
59            .map(|chunk| f32::from_le_bytes(*chunk))
60            .collect()
61    })
62}
63
64fn row_to_memory(row: &duckdb::Row<'_>) -> duckdb::Result<MemoryRecord> {
65    let id_str: String = row.get(0)?;
66    let tags_json: Option<String> = row.get(6)?;
67    let metadata_json: Option<String> = row.get(7)?;
68    let embedding_blob: Option<Vec<u8>> = row.get(8)?;
69    let content_hash: Vec<u8> = row.get(9)?;
70    let prev_hash: Option<Vec<u8>> = row.get(10)?;
71
72    let memory_type_str: String = row.get(3)?;
73    let scope_str: String = row.get(4)?;
74    let source_type_str: String = row.get(11)?;
75    let consolidation_state_str: String = row.get(13)?;
76
77    Ok(MemoryRecord {
78        id: Uuid::parse_str(&id_str)
79            .map_err(|e| duckdb::Error::FromSqlConversionFailure(0, duckdb::types::Type::Text, Box::new(e)))?,
80        agent_id: row.get(1)?,
81        content: row.get(2)?,
82        memory_type: memory_type_str.parse()
83            .map_err(|e: Error| duckdb::Error::FromSqlConversionFailure(3, duckdb::types::Type::Text, e.to_string().into()))?,
84        scope: scope_str.parse()
85            .map_err(|e: Error| duckdb::Error::FromSqlConversionFailure(4, duckdb::types::Type::Text, e.to_string().into()))?,
86        importance: row.get(5)?,
87        tags: match tags_json {
88            Some(ref s) => serde_json::from_str(s).unwrap_or_else(|e| {
89                tracing::warn!(id = %id_str, error = %e, raw = %s, "corrupted tags JSON, defaulting to empty");
90                vec![]
91            }),
92            None => vec![],
93        },
94        metadata: match metadata_json {
95            Some(ref s) => serde_json::from_str(s).unwrap_or_else(|e| {
96                tracing::warn!(id = %id_str, error = %e, "corrupted metadata JSON, defaulting to empty");
97                serde_json::Value::Object(serde_json::Map::new())
98            }),
99            None => serde_json::Value::Object(serde_json::Map::new()),
100        },
101        embedding: deserialize_embedding(embedding_blob),
102        content_hash,
103        prev_hash,
104        source_type: source_type_str.parse()
105            .map_err(|e: Error| duckdb::Error::FromSqlConversionFailure(11, duckdb::types::Type::Text, e.to_string().into()))?,
106        source_id: row.get(12)?,
107        consolidation_state: consolidation_state_str.parse()
108            .map_err(|e: Error| duckdb::Error::FromSqlConversionFailure(13, duckdb::types::Type::Text, e.to_string().into()))?,
109        access_count: u64::try_from(row.get::<_, i64>(14)?).unwrap_or(0),
110        org_id: row.get(15)?,
111        thread_id: row.get(16)?,
112        created_at: row.get(17)?,
113        updated_at: row.get(18)?,
114        last_accessed_at: row.get(19)?,
115        expires_at: row.get(20)?,
116        deleted_at: row.get(21)?,
117        decay_rate: row.get(22)?,
118        created_by: row.get(23)?,
119        version: u32::try_from(row.get::<_, i32>(24)?).unwrap_or(1),
120        prev_version_id: match row.get::<_, Option<String>>(25)? {
121            Some(s) => Uuid::parse_str(&s).map_err(|e| {
122                tracing::warn!(memory_id = %id_str, error = %e, "corrupted prev_version_id UUID");
123                e
124            }).ok(),
125            None => None,
126        },
127        quarantined: row.get::<_, bool>(26)?,
128        quarantine_reason: row.get(27)?,
129        decay_function: row.get(28).unwrap_or(None),
130    })
131}
132
133#[async_trait::async_trait]
134impl StorageBackend for DuckDbStorage {
135    fn backend_name(&self) -> &'static str {
136        "duckdb"
137    }
138
139    fn records_write_provenance(&self) -> bool {
140        true
141    }
142
143    async fn insert_write_provenance(&self, prov: &WriteProvenance) -> Result<()> {
144        let conn = self.conn.lock().await;
145        conn.execute(
146            "INSERT INTO write_provenance (id, memory_id, principal, capability_id, session_id, op, authored_at, flags, prev_hash, content_hash) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
147            duckdb::params![
148                prov.id.to_string(),
149                prov.memory_id.to_string(),
150                prov.principal,
151                prov.capability_id.map(|id| id.to_string()),
152                prov.session_id,
153                prov.op.as_str(),
154                prov.authored_at.to_rfc3339(),
155                flags_to_storage(&prov.flags),
156                prov.prev_hash,
157                prov.content_hash,
158            ],
159        )?;
160        Ok(())
161    }
162
163    async fn get_write_provenance(&self, memory_id: Uuid) -> Result<Option<WriteProvenance>> {
164        let conn = self.conn.lock().await;
165        let mut stmt = conn.prepare(
166            "SELECT id, memory_id, principal, capability_id, session_id, op, authored_at, flags, prev_hash, content_hash \
167             FROM write_provenance WHERE memory_id = ? ORDER BY id DESC LIMIT 1",
168        )?;
169        let mut rows = stmt.query_map(
170            duckdb::params![memory_id.to_string()],
171            row_to_write_provenance,
172        )?;
173        match rows.next() {
174            Some(r) => Ok(Some(r.map_err(|e| Error::Storage(e.to_string()))?)),
175            None => Ok(None),
176        }
177    }
178
179    async fn get_latest_provenance_hash(&self) -> Result<Option<Vec<u8>>> {
180        let conn = self.conn.lock().await;
181        // UUID v7 ids are time-ordered, so `id DESC` is the append order — the
182        // chain head is the most recently inserted record.
183        let mut stmt =
184            conn.prepare("SELECT content_hash FROM write_provenance ORDER BY id DESC LIMIT 1")?;
185        let mut rows = stmt.query_map([], |row| row.get::<_, Vec<u8>>(0))?;
186        match rows.next() {
187            Some(r) => Ok(Some(r.map_err(|e| Error::Storage(e.to_string()))?)),
188            None => Ok(None),
189        }
190    }
191
192    async fn list_provenance_by_principal(
193        &self,
194        principal: &str,
195        limit: usize,
196    ) -> Result<Vec<WriteProvenance>> {
197        let conn = self.conn.lock().await;
198        let mut stmt = conn.prepare(
199            "SELECT id, memory_id, principal, capability_id, session_id, op, authored_at, flags, prev_hash, content_hash \
200             FROM write_provenance WHERE principal = ? ORDER BY id DESC LIMIT ?",
201        )?;
202        let rows = stmt.query_map(
203            duckdb::params![principal, limit as i64],
204            row_to_write_provenance,
205        )?;
206        let mut out = Vec::new();
207        for r in rows {
208            out.push(r.map_err(|e| Error::Storage(e.to_string()))?);
209        }
210        Ok(out)
211    }
212
213    async fn list_provenance_by_session(
214        &self,
215        session_id: &str,
216        limit: usize,
217    ) -> Result<Vec<WriteProvenance>> {
218        let conn = self.conn.lock().await;
219        let mut stmt = conn.prepare(
220            "SELECT id, memory_id, principal, capability_id, session_id, op, authored_at, flags, prev_hash, content_hash \
221             FROM write_provenance WHERE session_id = ? ORDER BY id DESC LIMIT ?",
222        )?;
223        let rows = stmt.query_map(
224            duckdb::params![session_id, limit as i64],
225            row_to_write_provenance,
226        )?;
227        let mut out = Vec::new();
228        for r in rows {
229            out.push(r.map_err(|e| Error::Storage(e.to_string()))?);
230        }
231        Ok(out)
232    }
233
234    async fn list_memory_ids_by_principal(&self, principal: &str) -> Result<Vec<Uuid>> {
235        let conn = self.conn.lock().await;
236        // Only `remember` writes create a memory; `share` grants access to an
237        // existing one. FORGET BY PROVENANCE by principal revokes what the
238        // principal actually authored.
239        let mut stmt = conn.prepare(
240            "SELECT DISTINCT memory_id FROM write_provenance WHERE principal = ? AND op = 'remember'",
241        )?;
242        let rows = stmt.query_map(duckdb::params![principal], |row| row.get::<_, String>(0))?;
243        collect_memory_ids(rows)
244    }
245
246    async fn list_memory_ids_by_session(&self, session_id: &str) -> Result<Vec<Uuid>> {
247        let conn = self.conn.lock().await;
248        let mut stmt = conn.prepare(
249            "SELECT DISTINCT memory_id FROM write_provenance WHERE session_id = ? AND op = 'remember'",
250        )?;
251        let rows = stmt.query_map(duckdb::params![session_id], |row| row.get::<_, String>(0))?;
252        collect_memory_ids(rows)
253    }
254
255    async fn list_all_provenance(&self, limit: usize) -> Result<Vec<WriteProvenance>> {
256        let conn = self.conn.lock().await;
257        let mut stmt = conn.prepare(
258            "SELECT id, memory_id, principal, capability_id, session_id, op, authored_at, flags, prev_hash, content_hash \
259             FROM write_provenance ORDER BY id ASC LIMIT ?",
260        )?;
261        let rows = stmt.query_map(duckdb::params![limit as i64], row_to_write_provenance)?;
262        let mut out = Vec::new();
263        for r in rows {
264            out.push(r.map_err(|e| Error::Storage(e.to_string()))?);
265        }
266        Ok(out)
267    }
268
269    async fn insert_memory(&self, record: &MemoryRecord) -> Result<()> {
270        let conn = self.conn.lock().await;
271        let tags_json = serde_json::to_string(&record.tags)?;
272        let metadata_json = serde_json::to_string(&record.metadata)?;
273        let embedding_blob = serialize_embedding(&record.embedding);
274
275        conn.execute(
276            "INSERT INTO memories (id, agent_id, content, memory_type, scope, importance, tags, metadata, embedding, content_hash, prev_hash, source_type, source_id, consolidation_state, access_count, org_id, thread_id, created_at, updated_at, last_accessed_at, expires_at, deleted_at, decay_rate, created_by, version, prev_version_id, quarantined, quarantine_reason, decay_function) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
277            duckdb::params![
278                record.id.to_string(),
279                record.agent_id,
280                record.content,
281                record.memory_type.to_string(),
282                record.scope.to_string(),
283                record.importance,
284                tags_json,
285                metadata_json,
286                embedding_blob,
287                record.content_hash,
288                record.prev_hash,
289                record.source_type.to_string(),
290                record.source_id,
291                record.consolidation_state.to_string(),
292                record.access_count as i64,
293                record.org_id,
294                record.thread_id,
295                record.created_at,
296                record.updated_at,
297                record.last_accessed_at,
298                record.expires_at,
299                record.deleted_at,
300                record.decay_rate,
301                record.created_by,
302                record.version as i32,
303                record.prev_version_id.map(|id| id.to_string()),
304                record.quarantined,
305                record.quarantine_reason,
306                record.decay_function,
307            ],
308        )?;
309        Ok(())
310    }
311
312    async fn get_memory(&self, id: Uuid) -> Result<Option<MemoryRecord>> {
313        let conn = self.conn.lock().await;
314        let mut stmt = conn.prepare(
315            "SELECT id, agent_id, content, memory_type, scope, importance, tags, metadata, embedding, content_hash, prev_hash, source_type, source_id, consolidation_state, access_count, org_id, thread_id, created_at, updated_at, last_accessed_at, expires_at, deleted_at, decay_rate, created_by, version, prev_version_id, quarantined, quarantine_reason, decay_function FROM memories WHERE id = ?",
316        )?;
317        let result = stmt.query_row([id.to_string()], row_to_memory);
318        match result {
319            Ok(record) => Ok(Some(record)),
320            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
321            Err(e) => Err(Error::Storage(e.to_string())),
322        }
323    }
324
325    async fn update_memory(&self, record: &MemoryRecord) -> Result<()> {
326        let conn = self.conn.lock().await;
327        let tags_json = serde_json::to_string(&record.tags)?;
328        let metadata_json = serde_json::to_string(&record.metadata)?;
329        let embedding_blob = serialize_embedding(&record.embedding);
330
331        let affected = conn.execute(
332            "UPDATE memories SET agent_id=?, content=?, memory_type=?, scope=?, importance=?, tags=?, metadata=?, embedding=?, content_hash=?, prev_hash=?, source_type=?, source_id=?, consolidation_state=?, access_count=?, org_id=?, thread_id=?, updated_at=?, last_accessed_at=?, expires_at=?, deleted_at=?, decay_rate=?, created_by=?, version=?, prev_version_id=?, quarantined=?, quarantine_reason=?, decay_function=? WHERE id=?",
333            duckdb::params![
334                record.agent_id,
335                record.content,
336                record.memory_type.to_string(),
337                record.scope.to_string(),
338                record.importance,
339                tags_json,
340                metadata_json,
341                embedding_blob,
342                record.content_hash,
343                record.prev_hash,
344                record.source_type.to_string(),
345                record.source_id,
346                record.consolidation_state.to_string(),
347                record.access_count as i64,
348                record.org_id,
349                record.thread_id,
350                record.updated_at,
351                record.last_accessed_at,
352                record.expires_at,
353                record.deleted_at,
354                record.decay_rate,
355                record.created_by,
356                record.version as i32,
357                record.prev_version_id.map(|id| id.to_string()),
358                record.quarantined,
359                record.quarantine_reason,
360                record.decay_function,
361                record.id.to_string(),
362            ],
363        )?;
364        if affected == 0 {
365            return Err(Error::NotFound(format!("memory {} not found", record.id)));
366        }
367        Ok(())
368    }
369
370    async fn soft_delete_memory(&self, id: Uuid) -> Result<()> {
371        let conn = self.conn.lock().await;
372        let now = chrono::Utc::now().to_rfc3339();
373        let affected = conn.execute(
374            "UPDATE memories SET deleted_at = ?, updated_at = ? WHERE id = ? AND deleted_at IS NULL",
375            duckdb::params![now, now, id.to_string()],
376        )?;
377        if affected == 0 {
378            return Err(Error::NotFound(format!(
379                "memory {id} not found or already deleted"
380            )));
381        }
382        Ok(())
383    }
384
385    async fn hard_delete_memory(&self, id: Uuid) -> Result<()> {
386        let conn = self.conn.lock().await;
387        let affected = conn.execute(
388            "DELETE FROM memories WHERE id = ?",
389            duckdb::params![id.to_string()],
390        )?;
391        if affected == 0 {
392            return Err(Error::NotFound(format!("memory {id} not found")));
393        }
394        // Also clean up ACLs
395        conn.execute(
396            "DELETE FROM acls WHERE memory_id = ?",
397            duckdb::params![id.to_string()],
398        )?;
399        Ok(())
400    }
401
402    async fn list_memories(
403        &self,
404        filter: &MemoryFilter,
405        limit: usize,
406        offset: usize,
407    ) -> Result<Vec<MemoryRecord>> {
408        let conn = self.conn.lock().await;
409        let mut conditions = Vec::new();
410        let mut params: Vec<Box<dyn duckdb::ToSql>> = Vec::new();
411
412        if !filter.include_deleted {
413            conditions.push("deleted_at IS NULL".to_string());
414        }
415
416        if let Some(ref agent_id) = filter.agent_id {
417            conditions.push(format!("agent_id = ${}", params.len() + 1));
418            params.push(Box::new(agent_id.clone()));
419        }
420
421        if let Some(memory_type) = filter.memory_type {
422            conditions.push(format!("memory_type = ${}", params.len() + 1));
423            params.push(Box::new(memory_type.to_string()));
424        }
425
426        if let Some(scope) = filter.scope {
427            conditions.push(format!("scope = ${}", params.len() + 1));
428            params.push(Box::new(scope.to_string()));
429        }
430
431        if let Some(min_importance) = filter.min_importance {
432            conditions.push(format!("importance >= ${}", params.len() + 1));
433            params.push(Box::new(min_importance));
434        }
435
436        if let Some(ref org_id) = filter.org_id {
437            conditions.push(format!("org_id = ${}", params.len() + 1));
438            params.push(Box::new(org_id.clone()));
439        }
440
441        if let Some(ref thread_id) = filter.thread_id {
442            conditions.push(format!("thread_id = ${}", params.len() + 1));
443            params.push(Box::new(thread_id.clone()));
444        }
445
446        let where_clause = if conditions.is_empty() {
447            String::new()
448        } else {
449            format!("WHERE {}", conditions.join(" AND "))
450        };
451
452        let sql = format!(
453            "SELECT id, agent_id, content, memory_type, scope, importance, tags, metadata, embedding, content_hash, prev_hash, source_type, source_id, consolidation_state, access_count, org_id, thread_id, created_at, updated_at, last_accessed_at, expires_at, deleted_at, decay_rate, created_by, version, prev_version_id, quarantined, quarantine_reason, decay_function FROM memories {where_clause} ORDER BY created_at DESC LIMIT {limit} OFFSET {offset}"
454        );
455
456        let mut stmt = conn.prepare(&sql)?;
457        let param_refs: Vec<&dyn duckdb::ToSql> = params.iter().map(|p| p.as_ref()).collect();
458        let rows = stmt.query_map(param_refs.as_slice(), row_to_memory)?;
459
460        let mut results = Vec::new();
461        for row in rows {
462            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
463        }
464        Ok(results)
465    }
466
467    async fn touch_memory(&self, id: Uuid) -> Result<()> {
468        let conn = self.conn.lock().await;
469        let now = chrono::Utc::now().to_rfc3339();
470        conn.execute(
471            "UPDATE memories SET access_count = access_count + 1, last_accessed_at = ? WHERE id = ?",
472            duckdb::params![now, id.to_string()],
473        )?;
474        Ok(())
475    }
476
477    async fn insert_acl(&self, acl: &Acl) -> Result<()> {
478        let conn = self.conn.lock().await;
479        conn.execute(
480            "INSERT INTO acls (id, memory_id, principal_type, principal_id, permission, granted_by, created_at, expires_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
481            duckdb::params![
482                acl.id.to_string(),
483                acl.memory_id.to_string(),
484                acl.principal_type.to_string(),
485                acl.principal_id,
486                acl.permission.to_string(),
487                acl.granted_by,
488                acl.created_at,
489                acl.expires_at,
490            ],
491        )?;
492        Ok(())
493    }
494
495    async fn check_permission(
496        &self,
497        memory_id: Uuid,
498        principal_id: &str,
499        required: Permission,
500    ) -> Result<bool> {
501        // Do all DuckDB work in one block, then release the lock before delegation check
502        let acl_result = {
503            let conn = self.conn.lock().await;
504
505            // Check if the principal is the owner (agent_id matches)
506            let mut stmt = conn.prepare("SELECT agent_id FROM memories WHERE id = ?")?;
507            let owner_result =
508                stmt.query_row([memory_id.to_string()], |row| row.get::<_, String>(0));
509            match owner_result {
510                Ok(owner) if owner == principal_id => return Ok(true),
511                Err(duckdb::Error::QueryReturnedNoRows) => {
512                    return Err(Error::NotFound(format!("memory {memory_id} not found")));
513                }
514                _ => {}
515            }
516
517            // Check ACLs
518            let now = chrono::Utc::now().to_rfc3339();
519            let mut stmt = conn.prepare(
520                "SELECT permission FROM acls WHERE memory_id = ? AND principal_id = ? AND (expires_at IS NULL OR expires_at > ?)",
521            )?;
522            let rows = stmt.query_map(
523                duckdb::params![memory_id.to_string(), principal_id, now.clone()],
524                |row| row.get::<_, String>(0),
525            )?;
526
527            let mut perms: Vec<String> = Vec::new();
528            for row in rows {
529                perms.push(row.map_err(|e| Error::Storage(e.to_string()))?);
530            }
531
532            // Check public ACLs
533            let mut stmt = conn.prepare(
534                "SELECT permission FROM acls WHERE memory_id = ? AND principal_type = 'public' AND (expires_at IS NULL OR expires_at > ?)",
535            )?;
536            let rows = stmt.query_map(duckdb::params![memory_id.to_string(), now], |row| {
537                row.get::<_, String>(0)
538            })?;
539
540            for row in rows {
541                perms.push(row.map_err(|e| Error::Storage(e.to_string()))?);
542            }
543
544            perms
545        }; // conn lock dropped here
546
547        for perm_str in &acl_result {
548            if let Ok(perm) = perm_str.parse::<Permission>()
549                && perm.satisfies(required)
550            {
551                return Ok(true);
552            }
553        }
554
555        // Check delegations (conn lock is released)
556        if self
557            .check_delegation(principal_id, memory_id, required)
558            .await?
559        {
560            return Ok(true);
561        }
562
563        Ok(false)
564    }
565
566    async fn insert_relation(&self, relation: &Relation) -> Result<()> {
567        let conn = self.conn.lock().await;
568        conn.execute(
569            "INSERT INTO relations (id, source_id, target_id, relation_type, weight, metadata, created_at) VALUES (?, ?, ?, ?, ?, ?, ?)",
570            duckdb::params![
571                relation.id.to_string(),
572                relation.source_id.to_string(),
573                relation.target_id.to_string(),
574                relation.relation_type,
575                relation.weight,
576                serde_json::to_string(&relation.metadata)?,
577                relation.created_at,
578            ],
579        )?;
580        Ok(())
581    }
582
583    async fn get_relations_from(&self, source_id: Uuid) -> Result<Vec<Relation>> {
584        let conn = self.conn.lock().await;
585        let mut stmt = conn.prepare(
586            "SELECT id, source_id, target_id, relation_type, weight, metadata, created_at FROM relations WHERE source_id = ?",
587        )?;
588        let rows = stmt.query_map([source_id.to_string()], row_to_relation)?;
589        let mut results = Vec::new();
590        for row in rows {
591            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
592        }
593        Ok(results)
594    }
595
596    async fn get_relations_to(&self, target_id: Uuid) -> Result<Vec<Relation>> {
597        let conn = self.conn.lock().await;
598        let mut stmt = conn.prepare(
599            "SELECT id, source_id, target_id, relation_type, weight, metadata, created_at FROM relations WHERE target_id = ?",
600        )?;
601        let rows = stmt.query_map([target_id.to_string()], row_to_relation)?;
602        let mut results = Vec::new();
603        for row in rows {
604            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
605        }
606        Ok(results)
607    }
608
609    async fn delete_relation(&self, id: Uuid) -> Result<()> {
610        let conn = self.conn.lock().await;
611        let affected = conn.execute(
612            "DELETE FROM relations WHERE id = ?",
613            duckdb::params![id.to_string()],
614        )?;
615        if affected == 0 {
616            return Err(Error::NotFound(format!("relation {id} not found")));
617        }
618        Ok(())
619    }
620
621    async fn get_latest_memory_hash(
622        &self,
623        agent_id: &str,
624        thread_id: Option<&str>,
625    ) -> Result<Option<Vec<u8>>> {
626        let conn = self.conn.lock().await;
627        let (sql, result) = if let Some(tid) = thread_id {
628            let mut stmt = conn.prepare(
629                "SELECT content_hash FROM memories WHERE agent_id = ? AND thread_id = ? AND deleted_at IS NULL ORDER BY created_at DESC LIMIT 1",
630            )?;
631            let r = stmt.query_row(duckdb::params![agent_id, tid], |row| {
632                row.get::<_, Vec<u8>>(0)
633            });
634            ((), r)
635        } else {
636            let mut stmt = conn.prepare(
637                "SELECT content_hash FROM memories WHERE agent_id = ? AND thread_id IS NULL AND deleted_at IS NULL ORDER BY created_at DESC LIMIT 1",
638            )?;
639            let r = stmt.query_row(duckdb::params![agent_id], |row| row.get::<_, Vec<u8>>(0));
640            ((), r)
641        };
642        let _ = sql;
643        match result {
644            Ok(hash) => Ok(Some(hash)),
645            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
646            Err(e) => Err(Error::Storage(e.to_string())),
647        }
648    }
649
650    async fn get_latest_event_hash(
651        &self,
652        agent_id: &str,
653        thread_id: Option<&str>,
654    ) -> Result<Option<Vec<u8>>> {
655        let conn = self.conn.lock().await;
656        let result = if let Some(tid) = thread_id {
657            let mut stmt = conn.prepare(
658                "SELECT content_hash FROM agent_events WHERE agent_id = ? AND thread_id = ? ORDER BY timestamp DESC LIMIT 1",
659            )?;
660            stmt.query_row(duckdb::params![agent_id, tid], |row| {
661                row.get::<_, Vec<u8>>(0)
662            })
663        } else {
664            let mut stmt = conn.prepare(
665                "SELECT content_hash FROM agent_events WHERE agent_id = ? ORDER BY timestamp DESC LIMIT 1",
666            )?;
667            stmt.query_row(duckdb::params![agent_id], |row| row.get::<_, Vec<u8>>(0))
668        };
669        match result {
670            Ok(hash) => Ok(Some(hash)),
671            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
672            Err(e) => Err(Error::Storage(e.to_string())),
673        }
674    }
675
676    async fn get_sync_watermark(&self, key: &str) -> Result<Option<String>> {
677        let conn = self.conn.lock().await;
678        let mut stmt = conn.prepare("SELECT value FROM sync_metadata WHERE key = ?")?;
679        let result = stmt.query_row(duckdb::params![key], |row| row.get::<_, String>(0));
680        match result {
681            Ok(value) => Ok(Some(value)),
682            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
683            Err(e) => Err(Error::Storage(e.to_string())),
684        }
685    }
686
687    async fn set_sync_watermark(&self, key: &str, value: &str) -> Result<()> {
688        let conn = self.conn.lock().await;
689        let now = chrono::Utc::now().to_rfc3339();
690        // Try update first, then insert
691        let affected = conn.execute(
692            "UPDATE sync_metadata SET value = ?, updated_at = ? WHERE key = ?",
693            duckdb::params![value, now, key],
694        )?;
695        if affected == 0 {
696            conn.execute(
697                "INSERT INTO sync_metadata (key, value, updated_at) VALUES (?, ?, ?)",
698                duckdb::params![key, value, now],
699            )?;
700        }
701        Ok(())
702    }
703
704    async fn list_accessible_memory_ids(&self, agent_id: &str, limit: usize) -> Result<Vec<Uuid>> {
705        let conn = self.conn.lock().await;
706        let now = chrono::Utc::now().to_rfc3339();
707        let mut stmt = conn.prepare(
708            "SELECT id FROM memories WHERE (agent_id = ? OR scope = 'public' OR id IN (SELECT memory_id FROM acls WHERE principal_id = ? AND (expires_at IS NULL OR expires_at > ?))) AND deleted_at IS NULL LIMIT ?",
709        )?;
710        let rows = stmt.query_map(
711            duckdb::params![agent_id, agent_id, now, limit as i64],
712            |row| row.get::<_, String>(0),
713        )?;
714        let mut ids = Vec::new();
715        for row in rows {
716            let id_str = row.map_err(|e| Error::Storage(e.to_string()))?;
717            ids.push(Uuid::parse_str(&id_str).map_err(|e| Error::Storage(e.to_string()))?);
718        }
719        Ok(ids)
720    }
721
722    async fn insert_event(&self, event: &AgentEvent) -> Result<()> {
723        let conn = self.conn.lock().await;
724        let payload_json = serde_json::to_string(&event.payload)?;
725        let embedding_blob = serialize_embedding(&event.embedding);
726        conn.execute(
727            "INSERT INTO agent_events (id, agent_id, thread_id, run_id, parent_event_id, event_type, payload, trace_id, span_id, model, tokens_input, tokens_output, latency_ms, cost_usd, timestamp, logical_clock, content_hash, prev_hash, embedding) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
728            duckdb::params![
729                event.id.to_string(),
730                event.agent_id,
731                event.thread_id,
732                event.run_id,
733                event.parent_event_id.map(|id| id.to_string()),
734                event.event_type.to_string(),
735                payload_json,
736                event.trace_id,
737                event.span_id,
738                event.model,
739                event.tokens_input,
740                event.tokens_output,
741                event.latency_ms,
742                event.cost_usd,
743                event.timestamp,
744                event.logical_clock,
745                event.content_hash,
746                event.prev_hash,
747                embedding_blob,
748            ],
749        )?;
750        Ok(())
751    }
752
753    async fn list_events(
754        &self,
755        agent_id: &str,
756        limit: usize,
757        offset: usize,
758    ) -> Result<Vec<AgentEvent>> {
759        let conn = self.conn.lock().await;
760        let mut stmt = conn.prepare(
761            "SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type, payload, trace_id, span_id, model, tokens_input, tokens_output, latency_ms, cost_usd, timestamp, logical_clock, content_hash, prev_hash, embedding FROM agent_events WHERE agent_id = ? ORDER BY timestamp DESC LIMIT ? OFFSET ?",
762        )?;
763        let rows = stmt.query_map(
764            duckdb::params![agent_id, limit as i64, offset as i64],
765            row_to_event,
766        )?;
767        let mut results = Vec::new();
768        for row in rows {
769            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
770        }
771        Ok(results)
772    }
773
774    async fn get_events_by_thread(&self, thread_id: &str, limit: usize) -> Result<Vec<AgentEvent>> {
775        let conn = self.conn.lock().await;
776        let mut stmt = conn.prepare(
777            "SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type, payload, trace_id, span_id, model, tokens_input, tokens_output, latency_ms, cost_usd, timestamp, logical_clock, content_hash, prev_hash, embedding FROM agent_events WHERE thread_id = ? ORDER BY timestamp ASC LIMIT ?",
778        )?;
779        let rows = stmt.query_map(duckdb::params![thread_id, limit as i64], row_to_event)?;
780        let mut results = Vec::new();
781        for row in rows {
782            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
783        }
784        Ok(results)
785    }
786
787    async fn get_event(&self, id: Uuid) -> Result<Option<AgentEvent>> {
788        let conn = self.conn.lock().await;
789        let mut stmt = conn.prepare(
790            "SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type, payload, trace_id, span_id, model, tokens_input, tokens_output, latency_ms, cost_usd, timestamp, logical_clock, content_hash, prev_hash, embedding FROM agent_events WHERE id = ?",
791        )?;
792        let result = stmt.query_row([id.to_string()], row_to_event);
793        match result {
794            Ok(event) => Ok(Some(event)),
795            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
796            Err(e) => Err(Error::Storage(e.to_string())),
797        }
798    }
799
800    async fn list_child_events(
801        &self,
802        parent_event_id: Uuid,
803        limit: usize,
804    ) -> Result<Vec<AgentEvent>> {
805        let conn = self.conn.lock().await;
806        let mut stmt = conn.prepare(
807            "SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type, payload, trace_id, span_id, model, tokens_input, tokens_output, latency_ms, cost_usd, timestamp, logical_clock, content_hash, prev_hash, embedding FROM agent_events WHERE parent_event_id = ? ORDER BY timestamp ASC LIMIT ?",
808        )?;
809        let rows = stmt.query_map(
810            duckdb::params![parent_event_id.to_string(), limit as i64],
811            row_to_event,
812        )?;
813        let mut results = Vec::new();
814        for row in rows {
815            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
816        }
817        Ok(results)
818    }
819
820    async fn list_memories_by_agent_ordered(
821        &self,
822        agent_id: &str,
823        thread_id: Option<&str>,
824        limit: usize,
825    ) -> Result<Vec<MemoryRecord>> {
826        let conn = self.conn.lock().await;
827        let (result,) = if let Some(tid) = thread_id {
828            let mut stmt = conn.prepare(
829                "SELECT id, agent_id, content, memory_type, scope, importance, tags, metadata, embedding, content_hash, prev_hash, source_type, source_id, consolidation_state, access_count, org_id, thread_id, created_at, updated_at, last_accessed_at, expires_at, deleted_at, decay_rate, created_by, version, prev_version_id, quarantined, quarantine_reason, decay_function FROM memories WHERE agent_id = ? AND thread_id = ? AND deleted_at IS NULL ORDER BY created_at ASC LIMIT ?",
830            )?;
831            let rows =
832                stmt.query_map(duckdb::params![agent_id, tid, limit as i64], row_to_memory)?;
833            let mut results = Vec::new();
834            for row in rows {
835                results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
836            }
837            (results,)
838        } else {
839            let mut stmt = conn.prepare(
840                "SELECT id, agent_id, content, memory_type, scope, importance, tags, metadata, embedding, content_hash, prev_hash, source_type, source_id, consolidation_state, access_count, org_id, thread_id, created_at, updated_at, last_accessed_at, expires_at, deleted_at, decay_rate, created_by, version, prev_version_id, quarantined, quarantine_reason, decay_function FROM memories WHERE agent_id = ? AND deleted_at IS NULL ORDER BY created_at ASC LIMIT ?",
841            )?;
842            let rows = stmt.query_map(duckdb::params![agent_id, limit as i64], row_to_memory)?;
843            let mut results = Vec::new();
844            for row in rows {
845                results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
846            }
847            (results,)
848        };
849        Ok(result)
850    }
851
852    async fn list_memories_since(
853        &self,
854        updated_after: &str,
855        limit: usize,
856    ) -> Result<Vec<MemoryRecord>> {
857        let conn = self.conn.lock().await;
858        let mut stmt = conn.prepare(
859            "SELECT id, agent_id, content, memory_type, scope, importance, tags, metadata, embedding, content_hash, prev_hash, source_type, source_id, consolidation_state, access_count, org_id, thread_id, created_at, updated_at, last_accessed_at, expires_at, deleted_at, decay_rate, created_by, version, prev_version_id, quarantined, quarantine_reason, decay_function FROM memories WHERE updated_at > ? ORDER BY updated_at ASC LIMIT ?",
860        )?;
861        let rows = stmt.query_map(duckdb::params![updated_after, limit as i64], row_to_memory)?;
862        let mut results = Vec::new();
863        for row in rows {
864            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
865        }
866        Ok(results)
867    }
868
869    async fn upsert_memory(&self, record: &MemoryRecord) -> Result<()> {
870        // Try update first; if no rows affected, insert
871        match self.update_memory(record).await {
872            Ok(()) => Ok(()),
873            Err(Error::NotFound(_)) => self.insert_memory(record).await,
874            Err(e) => Err(e),
875        }
876    }
877
878    async fn cleanup_expired(&self) -> Result<usize> {
879        let conn = self.conn.lock().await;
880        let now = chrono::Utc::now().to_rfc3339();
881        let affected = conn.execute(
882            "UPDATE memories SET deleted_at = ? WHERE expires_at IS NOT NULL AND expires_at < ? AND deleted_at IS NULL",
883            duckdb::params![now.clone(), now],
884        )?;
885        Ok(affected)
886    }
887
888    async fn insert_delegation(&self, d: &Delegation) -> Result<()> {
889        let conn = self.conn.lock().await;
890        let scope_type = d.scope.to_string();
891        let scope_value = match &d.scope {
892            DelegationScope::AllMemories => serde_json::Value::Null,
893            DelegationScope::ByTag(tags) => serde_json::json!(tags),
894            DelegationScope::ByMemoryId(ids) => {
895                serde_json::json!(ids.iter().map(|id| id.to_string()).collect::<Vec<_>>())
896            }
897        };
898        let scope_value_json = serde_json::to_string(&scope_value)?;
899
900        conn.execute(
901            "INSERT INTO delegations (id, delegator_id, delegate_id, permission, scope_type, scope_value, max_depth, current_depth, parent_delegation_id, created_at, expires_at, revoked_at) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
902            duckdb::params![
903                d.id.to_string(),
904                d.delegator_id,
905                d.delegate_id,
906                d.permission.to_string(),
907                scope_type,
908                scope_value_json,
909                d.max_depth as i32,
910                d.current_depth as i32,
911                d.parent_delegation_id.map(|id| id.to_string()),
912                d.created_at,
913                d.expires_at,
914                d.revoked_at,
915            ],
916        )?;
917        Ok(())
918    }
919
920    async fn list_delegations_for(&self, delegate_id: &str) -> Result<Vec<Delegation>> {
921        let conn = self.conn.lock().await;
922        let now = chrono::Utc::now().to_rfc3339();
923        let mut stmt = conn.prepare(
924            "SELECT id, delegator_id, delegate_id, permission, scope_type, scope_value, max_depth, current_depth, parent_delegation_id, created_at, expires_at, revoked_at FROM delegations WHERE delegate_id = ? AND revoked_at IS NULL AND (expires_at IS NULL OR expires_at > ?)",
925        )?;
926        let rows = stmt.query_map(duckdb::params![delegate_id, now], row_to_delegation)?;
927        let mut results = Vec::new();
928        for row in rows {
929            results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
930        }
931        Ok(results)
932    }
933
934    async fn revoke_delegation(&self, id: Uuid) -> Result<()> {
935        let conn = self.conn.lock().await;
936        let now = chrono::Utc::now().to_rfc3339();
937        let affected = conn.execute(
938            "UPDATE delegations SET revoked_at = ? WHERE id = ? AND revoked_at IS NULL",
939            duckdb::params![now, id.to_string()],
940        )?;
941        if affected == 0 {
942            return Err(Error::NotFound(format!(
943                "delegation {id} not found or already revoked"
944            )));
945        }
946        Ok(())
947    }
948
949    async fn check_delegation(
950        &self,
951        delegate_id: &str,
952        memory_id: Uuid,
953        required: Permission,
954    ) -> Result<bool> {
955        let delegations = self.list_delegations_for(delegate_id).await?;
956        // Get the memory to check scope
957        let memory = match self.get_memory(memory_id).await? {
958            Some(m) => m,
959            None => return Ok(false),
960        };
961
962        for d in &delegations {
963            if !d.permission.satisfies(required) {
964                continue;
965            }
966            match &d.scope {
967                DelegationScope::AllMemories => return Ok(true),
968                DelegationScope::ByMemoryId(ids) => {
969                    if ids.contains(&memory_id) {
970                        return Ok(true);
971                    }
972                }
973                DelegationScope::ByTag(tags) => {
974                    if tags.iter().any(|t| memory.tags.contains(t)) {
975                        return Ok(true);
976                    }
977                }
978            }
979        }
980        Ok(false)
981    }
982
983    async fn insert_or_update_agent_profile(&self, profile: &AgentProfile) -> Result<()> {
984        let conn = self.conn.lock().await;
985        // Try update first, then insert
986        let affected = conn.execute(
987            "UPDATE agent_profiles SET avg_importance = ?, avg_content_length = ?, total_memories = ?, last_updated = ? WHERE agent_id = ?",
988            duckdb::params![
989                profile.avg_importance,
990                profile.avg_content_length,
991                profile.total_memories as i64,
992                profile.last_updated,
993                profile.agent_id,
994            ],
995        )?;
996        if affected == 0 {
997            conn.execute(
998                "INSERT INTO agent_profiles (agent_id, avg_importance, avg_content_length, total_memories, last_updated) VALUES (?, ?, ?, ?, ?)",
999                duckdb::params![
1000                    profile.agent_id,
1001                    profile.avg_importance,
1002                    profile.avg_content_length,
1003                    profile.total_memories as i64,
1004                    profile.last_updated,
1005                ],
1006            )?;
1007        }
1008        Ok(())
1009    }
1010
1011    async fn get_agent_profile(&self, agent_id: &str) -> Result<Option<AgentProfile>> {
1012        let conn = self.conn.lock().await;
1013        let mut stmt = conn.prepare(
1014            "SELECT agent_id, avg_importance, avg_content_length, total_memories, last_updated FROM agent_profiles WHERE agent_id = ?",
1015        )?;
1016        let result = stmt.query_row([agent_id], |row| {
1017            Ok(AgentProfile {
1018                agent_id: row.get(0)?,
1019                avg_importance: row.get(1)?,
1020                avg_content_length: row.get(2)?,
1021                total_memories: row.get::<_, i64>(3)? as u64,
1022                last_updated: row.get(4)?,
1023            })
1024        });
1025        match result {
1026            Ok(profile) => Ok(Some(profile)),
1027            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
1028            Err(e) => Err(Error::Storage(e.to_string())),
1029        }
1030    }
1031
1032    async fn insert_or_update_embedding_baseline(
1033        &self,
1034        baseline: &EmbeddingBaseline,
1035    ) -> Result<()> {
1036        let conn = self.conn.lock().await;
1037        let mu_json = serde_json::to_string(&baseline.mu)?;
1038        let cov_json = serde_json::to_string(&baseline.cov_diag)?;
1039        let affected = conn.execute(
1040            "UPDATE embedding_baseline SET mu = ?, cov_diag = ?, n = ?, updated_at = ? WHERE agent_id = ?",
1041            duckdb::params![
1042                mu_json,
1043                cov_json,
1044                baseline.n as i64,
1045                baseline.updated_at,
1046                baseline.agent_id,
1047            ],
1048        )?;
1049        if affected == 0 {
1050            let mu_json = serde_json::to_string(&baseline.mu)?;
1051            let cov_json = serde_json::to_string(&baseline.cov_diag)?;
1052            conn.execute(
1053                "INSERT INTO embedding_baseline (agent_id, mu, cov_diag, n, updated_at) VALUES (?, ?, ?, ?, ?)",
1054                duckdb::params![
1055                    baseline.agent_id,
1056                    mu_json,
1057                    cov_json,
1058                    baseline.n as i64,
1059                    baseline.updated_at,
1060                ],
1061            )?;
1062        }
1063        Ok(())
1064    }
1065
1066    async fn get_embedding_baseline(&self, agent_id: &str) -> Result<Option<EmbeddingBaseline>> {
1067        let conn = self.conn.lock().await;
1068        let mut stmt = conn.prepare(
1069            "SELECT agent_id, mu, cov_diag, n, updated_at FROM embedding_baseline WHERE agent_id = ?",
1070        )?;
1071        let result: duckdb::Result<(String, String, String, i64, String)> =
1072            stmt.query_row([agent_id], |row| {
1073                Ok((
1074                    row.get(0)?,
1075                    row.get(1)?,
1076                    row.get(2)?,
1077                    row.get(3)?,
1078                    row.get(4)?,
1079                ))
1080            });
1081        match result {
1082            Ok((agent_id, mu_json, cov_json, n, updated_at)) => {
1083                let mu: Vec<f32> = serde_json::from_str(&mu_json)?;
1084                let cov_diag: Vec<f32> = serde_json::from_str(&cov_json)?;
1085                Ok(Some(EmbeddingBaseline {
1086                    agent_id,
1087                    mu,
1088                    cov_diag,
1089                    n: n as u64,
1090                    updated_at,
1091                }))
1092            }
1093            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
1094            Err(e) => Err(Error::Storage(e.to_string())),
1095        }
1096    }
1097
1098    async fn insert_checkpoint(&self, cp: &Checkpoint) -> Result<()> {
1099        let conn = self.conn.lock().await;
1100        let state_snapshot_json = serde_json::to_string(&cp.state_snapshot)?;
1101        let state_diff_json = cp
1102            .state_diff
1103            .as_ref()
1104            .map(serde_json::to_string)
1105            .transpose()?;
1106        let memory_refs_json = serde_json::to_string(
1107            &cp.memory_refs
1108                .iter()
1109                .map(|id| id.to_string())
1110                .collect::<Vec<_>>(),
1111        )?;
1112        let metadata_json = serde_json::to_string(&cp.metadata)?;
1113
1114        conn.execute(
1115            "INSERT INTO checkpoints (id, thread_id, agent_id, parent_id, branch_name, state_snapshot, state_diff, memory_refs, event_cursor, label, created_at, metadata) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
1116            duckdb::params![
1117                cp.id.to_string(),
1118                cp.thread_id,
1119                cp.agent_id,
1120                cp.parent_id.map(|id| id.to_string()),
1121                cp.branch_name,
1122                state_snapshot_json,
1123                state_diff_json,
1124                memory_refs_json,
1125                cp.event_cursor.map(|id| id.to_string()),
1126                cp.label,
1127                cp.created_at,
1128                metadata_json,
1129            ],
1130        )?;
1131        Ok(())
1132    }
1133
1134    async fn get_checkpoint(&self, id: Uuid) -> Result<Option<Checkpoint>> {
1135        let conn = self.conn.lock().await;
1136        let mut stmt = conn.prepare(
1137            "SELECT id, thread_id, agent_id, parent_id, branch_name, state_snapshot, state_diff, memory_refs, event_cursor, label, created_at, metadata FROM checkpoints WHERE id = ?",
1138        )?;
1139        let result = stmt.query_row([id.to_string()], row_to_checkpoint);
1140        match result {
1141            Ok(cp) => Ok(Some(cp)),
1142            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
1143            Err(e) => Err(Error::Storage(e.to_string())),
1144        }
1145    }
1146
1147    async fn list_checkpoints(
1148        &self,
1149        thread_id: &str,
1150        branch: Option<&str>,
1151        limit: usize,
1152    ) -> Result<Vec<Checkpoint>> {
1153        let conn = self.conn.lock().await;
1154        let (sql, rows_result) = if let Some(branch_name) = branch {
1155            let mut stmt = conn.prepare(
1156                "SELECT id, thread_id, agent_id, parent_id, branch_name, state_snapshot, state_diff, memory_refs, event_cursor, label, created_at, metadata FROM checkpoints WHERE thread_id = ? AND branch_name = ? ORDER BY created_at DESC LIMIT ?",
1157            )?;
1158            let rows = stmt.query_map(
1159                duckdb::params![thread_id, branch_name, limit as i64],
1160                row_to_checkpoint,
1161            )?;
1162            let mut results = Vec::new();
1163            for row in rows {
1164                results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
1165            }
1166            ((), Ok(results))
1167        } else {
1168            let mut stmt = conn.prepare(
1169                "SELECT id, thread_id, agent_id, parent_id, branch_name, state_snapshot, state_diff, memory_refs, event_cursor, label, created_at, metadata FROM checkpoints WHERE thread_id = ? ORDER BY created_at DESC LIMIT ?",
1170            )?;
1171            let rows =
1172                stmt.query_map(duckdb::params![thread_id, limit as i64], row_to_checkpoint)?;
1173            let mut results = Vec::new();
1174            for row in rows {
1175                results.push(row.map_err(|e| Error::Storage(e.to_string()))?);
1176            }
1177            ((), Ok(results))
1178        };
1179        let _ = sql;
1180        rows_result
1181    }
1182
1183    async fn get_latest_checkpoint(
1184        &self,
1185        thread_id: &str,
1186        branch: &str,
1187    ) -> Result<Option<Checkpoint>> {
1188        let conn = self.conn.lock().await;
1189        let mut stmt = conn.prepare(
1190            "SELECT id, thread_id, agent_id, parent_id, branch_name, state_snapshot, state_diff, memory_refs, event_cursor, label, created_at, metadata FROM checkpoints WHERE thread_id = ? AND branch_name = ? ORDER BY created_at DESC LIMIT 1",
1191        )?;
1192        let result = stmt.query_row(duckdb::params![thread_id, branch], row_to_checkpoint);
1193        match result {
1194            Ok(cp) => Ok(Some(cp)),
1195            Err(duckdb::Error::QueryReturnedNoRows) => Ok(None),
1196            Err(e) => Err(Error::Storage(e.to_string())),
1197        }
1198    }
1199}
1200
1201fn row_to_event(row: &duckdb::Row<'_>) -> duckdb::Result<AgentEvent> {
1202    let id_str: String = row.get(0)?;
1203    let parent_id_str: Option<String> = row.get(4)?;
1204    let payload_json: Option<String> = row.get(6)?;
1205    let event_type_str: String = row.get(5)?;
1206    let content_hash: Vec<u8> = row.get(16)?;
1207    let prev_hash: Option<Vec<u8>> = row.get(17)?;
1208    let embedding_blob: Option<Vec<u8>> = row.get(18).unwrap_or(None);
1209
1210    Ok(AgentEvent {
1211        id: Uuid::parse_str(&id_str).map_err(|e| {
1212            duckdb::Error::FromSqlConversionFailure(0, duckdb::types::Type::Text, Box::new(e))
1213        })?,
1214        agent_id: row.get(1)?,
1215        thread_id: row.get(2)?,
1216        run_id: row.get(3)?,
1217        parent_event_id: parent_id_str.and_then(|s| Uuid::parse_str(&s).ok()),
1218        event_type: event_type_str.parse().map_err(|e: Error| {
1219            duckdb::Error::FromSqlConversionFailure(
1220                5,
1221                duckdb::types::Type::Text,
1222                e.to_string().into(),
1223            )
1224        })?,
1225        payload: payload_json
1226            .and_then(|s| serde_json::from_str(&s).ok())
1227            .unwrap_or(serde_json::Value::Null),
1228        trace_id: row.get(7)?,
1229        span_id: row.get(8)?,
1230        model: row.get(9)?,
1231        tokens_input: row.get(10)?,
1232        tokens_output: row.get(11)?,
1233        latency_ms: row.get(12)?,
1234        cost_usd: row.get(13)?,
1235        timestamp: row.get(14)?,
1236        logical_clock: row.get(15)?,
1237        content_hash,
1238        prev_hash,
1239        embedding: deserialize_embedding(embedding_blob),
1240    })
1241}
1242
1243fn row_to_checkpoint(row: &duckdb::Row<'_>) -> duckdb::Result<Checkpoint> {
1244    let id_str: String = row.get(0)?;
1245    let parent_id_str: Option<String> = row.get(3)?;
1246    let state_snapshot_json: Option<String> = row.get(5)?;
1247    let state_diff_json: Option<String> = row.get(6)?;
1248    let memory_refs_json: Option<String> = row.get(7)?;
1249    let event_cursor_str: Option<String> = row.get(8)?;
1250    let metadata_json: Option<String> = row.get(11)?;
1251
1252    Ok(Checkpoint {
1253        id: Uuid::parse_str(&id_str).map_err(|e| {
1254            duckdb::Error::FromSqlConversionFailure(0, duckdb::types::Type::Text, Box::new(e))
1255        })?,
1256        thread_id: row.get(1)?,
1257        agent_id: row.get(2)?,
1258        parent_id: parent_id_str.and_then(|s| Uuid::parse_str(&s).ok()),
1259        branch_name: row.get(4)?,
1260        state_snapshot: state_snapshot_json
1261            .and_then(|s| serde_json::from_str(&s).ok())
1262            .unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
1263        state_diff: state_diff_json.and_then(|s| serde_json::from_str(&s).ok()),
1264        memory_refs: memory_refs_json
1265            .and_then(|s| serde_json::from_str::<Vec<String>>(&s).ok())
1266            .map(|v| {
1267                v.into_iter()
1268                    .filter_map(|s| Uuid::parse_str(&s).ok())
1269                    .collect()
1270            })
1271            .unwrap_or_default(),
1272        event_cursor: event_cursor_str.and_then(|s| Uuid::parse_str(&s).ok()),
1273        label: row.get(9)?,
1274        created_at: row.get(10)?,
1275        metadata: metadata_json
1276            .and_then(|s| serde_json::from_str(&s).ok())
1277            .unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
1278    })
1279}
1280
1281fn collect_memory_ids(rows: impl Iterator<Item = duckdb::Result<String>>) -> Result<Vec<Uuid>> {
1282    let mut out = Vec::new();
1283    for r in rows {
1284        let s = r.map_err(|e| Error::Storage(e.to_string()))?;
1285        out.push(Uuid::parse_str(&s).map_err(|e| Error::Storage(e.to_string()))?);
1286    }
1287    Ok(out)
1288}
1289
1290fn row_to_write_provenance(row: &duckdb::Row<'_>) -> duckdb::Result<WriteProvenance> {
1291    let id_str: String = row.get(0)?;
1292    let memory_id_str: String = row.get(1)?;
1293    let principal: String = row.get(2)?;
1294    let capability_id_str: Option<String> = row.get(3)?;
1295    let session_id: Option<String> = row.get(4)?;
1296    let op_str: String = row.get(5)?;
1297    let authored_at_str: String = row.get(6)?;
1298    // `flags` was added after the table's first release; an older row stores NULL
1299    // (read as None → empty flag set). Column order: ..., authored_at(6),
1300    // flags(7), prev_hash(8), content_hash(9).
1301    let flags_str: Option<String> = row.get(7)?;
1302    let prev_hash: Option<Vec<u8>> = row.get(8)?;
1303    let content_hash: Vec<u8> = row.get(9)?;
1304    let flags = flags_from_storage(flags_str.as_deref().unwrap_or(""));
1305
1306    let uuid_err = |col: usize, e: uuid::Error| {
1307        duckdb::Error::FromSqlConversionFailure(col, duckdb::types::Type::Text, Box::new(e))
1308    };
1309    let op = match op_str.as_str() {
1310        "remember" => WriteOp::Remember,
1311        "share" => WriteOp::Share,
1312        other => {
1313            return Err(duckdb::Error::FromSqlConversionFailure(
1314                5,
1315                duckdb::types::Type::Text,
1316                format!("unknown write op `{other}`").into(),
1317            ));
1318        }
1319    };
1320    let authored_at = chrono::DateTime::parse_from_rfc3339(&authored_at_str)
1321        .map(|dt| dt.with_timezone(&chrono::Utc))
1322        .map_err(|e| {
1323            duckdb::Error::FromSqlConversionFailure(6, duckdb::types::Type::Text, Box::new(e))
1324        })?;
1325    let capability_id = match capability_id_str {
1326        Some(s) => Some(Uuid::parse_str(&s).map_err(|e| uuid_err(3, e))?),
1327        None => None,
1328    };
1329    Ok(WriteProvenance {
1330        id: Uuid::parse_str(&id_str).map_err(|e| uuid_err(0, e))?,
1331        memory_id: Uuid::parse_str(&memory_id_str).map_err(|e| uuid_err(1, e))?,
1332        principal,
1333        capability_id,
1334        session_id,
1335        op,
1336        authored_at,
1337        flags,
1338        prev_hash,
1339        content_hash,
1340    })
1341}
1342
1343fn row_to_delegation(row: &duckdb::Row<'_>) -> duckdb::Result<Delegation> {
1344    let id_str: String = row.get(0)?;
1345    let scope_type: String = row.get(4)?;
1346    let scope_value_json: Option<String> = row.get(5)?;
1347    let parent_id_str: Option<String> = row.get(8)?;
1348
1349    let scope = match scope_type.as_str() {
1350        "by_tag" => {
1351            let tags: Vec<String> = scope_value_json
1352                .and_then(|s| serde_json::from_str(&s).ok())
1353                .unwrap_or_default();
1354            DelegationScope::ByTag(tags)
1355        }
1356        "by_memory_id" => {
1357            let ids: Vec<String> = scope_value_json
1358                .and_then(|s| serde_json::from_str(&s).ok())
1359                .unwrap_or_default();
1360            let uuids = ids
1361                .into_iter()
1362                .filter_map(|s| Uuid::parse_str(&s).ok())
1363                .collect();
1364            DelegationScope::ByMemoryId(uuids)
1365        }
1366        _ => DelegationScope::AllMemories,
1367    };
1368
1369    let permission_str: String = row.get(3)?;
1370
1371    Ok(Delegation {
1372        id: Uuid::parse_str(&id_str).map_err(|e| {
1373            duckdb::Error::FromSqlConversionFailure(0, duckdb::types::Type::Text, Box::new(e))
1374        })?,
1375        delegator_id: row.get(1)?,
1376        delegate_id: row.get(2)?,
1377        permission: permission_str.parse().map_err(|e: Error| {
1378            duckdb::Error::FromSqlConversionFailure(
1379                3,
1380                duckdb::types::Type::Text,
1381                e.to_string().into(),
1382            )
1383        })?,
1384        scope,
1385        max_depth: row.get::<_, i32>(6)? as u32,
1386        current_depth: row.get::<_, i32>(7)? as u32,
1387        parent_delegation_id: parent_id_str.and_then(|s| Uuid::parse_str(&s).ok()),
1388        created_at: row.get(9)?,
1389        expires_at: row.get(10)?,
1390        revoked_at: row.get(11)?,
1391    })
1392}
1393
1394fn row_to_relation(row: &duckdb::Row<'_>) -> duckdb::Result<Relation> {
1395    let id_str: String = row.get(0)?;
1396    let source_str: String = row.get(1)?;
1397    let target_str: String = row.get(2)?;
1398    let metadata_json: Option<String> = row.get(5)?;
1399
1400    Ok(Relation {
1401        id: Uuid::parse_str(&id_str).map_err(|e| {
1402            duckdb::Error::FromSqlConversionFailure(0, duckdb::types::Type::Text, Box::new(e))
1403        })?,
1404        source_id: Uuid::parse_str(&source_str).map_err(|e| {
1405            duckdb::Error::FromSqlConversionFailure(1, duckdb::types::Type::Text, Box::new(e))
1406        })?,
1407        target_id: Uuid::parse_str(&target_str).map_err(|e| {
1408            duckdb::Error::FromSqlConversionFailure(2, duckdb::types::Type::Text, Box::new(e))
1409        })?,
1410        relation_type: row.get(3)?,
1411        weight: row.get(4)?,
1412        metadata: metadata_json
1413            .and_then(|s| serde_json::from_str(&s).ok())
1414            .unwrap_or(serde_json::Value::Object(serde_json::Map::new())),
1415        created_at: row.get(6)?,
1416    })
1417}
1418
1419#[cfg(test)]
1420mod tests {
1421    use super::*;
1422    use crate::hash::compute_content_hash;
1423    use crate::model::acl::PrincipalType;
1424    use crate::model::checkpoint::Checkpoint;
1425    use crate::model::event::{AgentEvent, EventType};
1426    use crate::model::memory::{ConsolidationState, MemoryType, Scope, SourceType};
1427
1428    fn make_record(agent_id: &str) -> MemoryRecord {
1429        let now = chrono::Utc::now().to_rfc3339();
1430        let content = "test memory content";
1431        MemoryRecord {
1432            id: Uuid::now_v7(),
1433            agent_id: agent_id.to_string(),
1434            content: content.to_string(),
1435            memory_type: MemoryType::Semantic,
1436            scope: Scope::Private,
1437            importance: 0.7,
1438            tags: vec!["test".to_string()],
1439            metadata: serde_json::json!({"key": "value"}),
1440            embedding: Some(vec![0.1, 0.2, 0.3]),
1441            content_hash: compute_content_hash(content, agent_id, &now),
1442            prev_hash: None,
1443            source_type: SourceType::Agent,
1444            source_id: None,
1445            consolidation_state: ConsolidationState::Raw,
1446            access_count: 0,
1447            org_id: None,
1448            thread_id: None,
1449            created_at: now.clone(),
1450            updated_at: now,
1451            last_accessed_at: None,
1452            expires_at: None,
1453            deleted_at: None,
1454            decay_rate: None,
1455            created_by: None,
1456            version: 1,
1457            prev_version_id: None,
1458            quarantined: false,
1459            quarantine_reason: None,
1460            decay_function: None,
1461        }
1462    }
1463
1464    #[tokio::test]
1465    async fn test_insert_and_get() {
1466        let storage = DuckDbStorage::open_in_memory().unwrap();
1467        let record = make_record("agent-1");
1468        storage.insert_memory(&record).await.unwrap();
1469
1470        let fetched = storage.get_memory(record.id).await.unwrap().unwrap();
1471        assert_eq!(fetched.id, record.id);
1472        assert_eq!(fetched.content, record.content);
1473        assert_eq!(fetched.agent_id, record.agent_id);
1474        assert_eq!(fetched.memory_type, record.memory_type);
1475        assert_eq!(fetched.tags, record.tags);
1476        assert_eq!(fetched.embedding, record.embedding);
1477    }
1478
1479    #[tokio::test]
1480    async fn test_get_nonexistent() {
1481        let storage = DuckDbStorage::open_in_memory().unwrap();
1482        let result = storage.get_memory(Uuid::now_v7()).await.unwrap();
1483        assert!(result.is_none());
1484    }
1485
1486    #[tokio::test]
1487    async fn test_soft_delete() {
1488        let storage = DuckDbStorage::open_in_memory().unwrap();
1489        let record = make_record("agent-1");
1490        storage.insert_memory(&record).await.unwrap();
1491
1492        storage.soft_delete_memory(record.id).await.unwrap();
1493
1494        // Should still exist in DB but with deleted_at set
1495        let fetched = storage.get_memory(record.id).await.unwrap().unwrap();
1496        assert!(fetched.deleted_at.is_some());
1497
1498        // Should not appear in list by default
1499        let filter = MemoryFilter::default();
1500        let list = storage.list_memories(&filter, 100, 0).await.unwrap();
1501        assert!(list.is_empty());
1502
1503        // Should appear with include_deleted
1504        let filter_with_deleted = MemoryFilter {
1505            include_deleted: true,
1506            ..Default::default()
1507        };
1508        let list = storage
1509            .list_memories(&filter_with_deleted, 100, 0)
1510            .await
1511            .unwrap();
1512        assert_eq!(list.len(), 1);
1513    }
1514
1515    #[tokio::test]
1516    async fn test_hard_delete() {
1517        let storage = DuckDbStorage::open_in_memory().unwrap();
1518        let record = make_record("agent-1");
1519        storage.insert_memory(&record).await.unwrap();
1520
1521        storage.hard_delete_memory(record.id).await.unwrap();
1522
1523        let result = storage.get_memory(record.id).await.unwrap();
1524        assert!(result.is_none());
1525    }
1526
1527    #[tokio::test]
1528    async fn test_list_with_filters() {
1529        let storage = DuckDbStorage::open_in_memory().unwrap();
1530
1531        let mut r1 = make_record("agent-1");
1532        r1.memory_type = MemoryType::Episodic;
1533        storage.insert_memory(&r1).await.unwrap();
1534
1535        let mut r2 = make_record("agent-1");
1536        r2.memory_type = MemoryType::Semantic;
1537        storage.insert_memory(&r2).await.unwrap();
1538
1539        let mut r3 = make_record("agent-2");
1540        r3.memory_type = MemoryType::Semantic;
1541        storage.insert_memory(&r3).await.unwrap();
1542
1543        // Filter by agent
1544        let filter = MemoryFilter {
1545            agent_id: Some("agent-1".to_string()),
1546            ..Default::default()
1547        };
1548        let list = storage.list_memories(&filter, 100, 0).await.unwrap();
1549        assert_eq!(list.len(), 2);
1550
1551        // Filter by type
1552        let filter = MemoryFilter {
1553            memory_type: Some(MemoryType::Semantic),
1554            ..Default::default()
1555        };
1556        let list = storage.list_memories(&filter, 100, 0).await.unwrap();
1557        assert_eq!(list.len(), 2);
1558
1559        // Filter by agent + type
1560        let filter = MemoryFilter {
1561            agent_id: Some("agent-1".to_string()),
1562            memory_type: Some(MemoryType::Episodic),
1563            ..Default::default()
1564        };
1565        let list = storage.list_memories(&filter, 100, 0).await.unwrap();
1566        assert_eq!(list.len(), 1);
1567    }
1568
1569    #[tokio::test]
1570    async fn test_touch_memory() {
1571        let storage = DuckDbStorage::open_in_memory().unwrap();
1572        let record = make_record("agent-1");
1573        storage.insert_memory(&record).await.unwrap();
1574
1575        storage.touch_memory(record.id).await.unwrap();
1576        storage.touch_memory(record.id).await.unwrap();
1577
1578        let fetched = storage.get_memory(record.id).await.unwrap().unwrap();
1579        assert_eq!(fetched.access_count, 2);
1580        assert!(fetched.last_accessed_at.is_some());
1581    }
1582
1583    #[tokio::test]
1584    async fn test_acl_and_permission_check() {
1585        let storage = DuckDbStorage::open_in_memory().unwrap();
1586        let record = make_record("agent-1");
1587        storage.insert_memory(&record).await.unwrap();
1588
1589        // Owner always has permission
1590        assert!(
1591            storage
1592                .check_permission(record.id, "agent-1", Permission::Admin)
1593                .await
1594                .unwrap()
1595        );
1596
1597        // Non-owner has no permission by default
1598        assert!(
1599            !storage
1600                .check_permission(record.id, "agent-2", Permission::Read)
1601                .await
1602                .unwrap()
1603        );
1604
1605        // Grant read to agent-2
1606        let acl = Acl {
1607            id: Uuid::now_v7(),
1608            memory_id: record.id,
1609            principal_type: PrincipalType::Agent,
1610            principal_id: "agent-2".to_string(),
1611            permission: Permission::Read,
1612            granted_by: "agent-1".to_string(),
1613            created_at: chrono::Utc::now().to_rfc3339(),
1614            expires_at: None,
1615        };
1616        storage.insert_acl(&acl).await.unwrap();
1617
1618        // Now agent-2 can read
1619        assert!(
1620            storage
1621                .check_permission(record.id, "agent-2", Permission::Read)
1622                .await
1623                .unwrap()
1624        );
1625        // But not write
1626        assert!(
1627            !storage
1628                .check_permission(record.id, "agent-2", Permission::Write)
1629                .await
1630                .unwrap()
1631        );
1632    }
1633
1634    #[tokio::test]
1635    async fn test_event_insert_and_list() {
1636        let storage = DuckDbStorage::open_in_memory().unwrap();
1637        let now = chrono::Utc::now().to_rfc3339();
1638        let event = AgentEvent {
1639            id: Uuid::now_v7(),
1640            agent_id: "agent-1".to_string(),
1641            thread_id: Some("thread-1".to_string()),
1642            run_id: None,
1643            parent_event_id: None,
1644            event_type: EventType::MemoryWrite,
1645            payload: serde_json::json!({"memory_id": "abc"}),
1646            trace_id: None,
1647            span_id: None,
1648            model: None,
1649            tokens_input: None,
1650            tokens_output: None,
1651            latency_ms: None,
1652            cost_usd: None,
1653            timestamp: now.clone(),
1654            logical_clock: 1,
1655            content_hash: vec![1, 2, 3],
1656            prev_hash: None,
1657            embedding: None,
1658        };
1659
1660        storage.insert_event(&event).await.unwrap();
1661
1662        let events = storage.list_events("agent-1", 10, 0).await.unwrap();
1663        assert_eq!(events.len(), 1);
1664        assert_eq!(events[0].id, event.id);
1665        assert_eq!(events[0].event_type, EventType::MemoryWrite);
1666        assert_eq!(events[0].agent_id, "agent-1");
1667
1668        // Get single event
1669        let fetched = storage.get_event(event.id).await.unwrap().unwrap();
1670        assert_eq!(fetched.id, event.id);
1671        assert_eq!(fetched.content_hash, vec![1, 2, 3]);
1672    }
1673
1674    #[tokio::test]
1675    async fn test_events_by_thread() {
1676        let storage = DuckDbStorage::open_in_memory().unwrap();
1677        let now = chrono::Utc::now().to_rfc3339();
1678
1679        for i in 0..3 {
1680            let event = AgentEvent {
1681                id: Uuid::now_v7(),
1682                agent_id: "agent-1".to_string(),
1683                thread_id: Some("thread-A".to_string()),
1684                run_id: None,
1685                parent_event_id: None,
1686                event_type: EventType::MemoryWrite,
1687                payload: serde_json::json!({"i": i}),
1688                trace_id: None,
1689                span_id: None,
1690                model: None,
1691                tokens_input: None,
1692                tokens_output: None,
1693                latency_ms: None,
1694                cost_usd: None,
1695                timestamp: now.clone(),
1696                logical_clock: i,
1697                content_hash: vec![i as u8],
1698                prev_hash: None,
1699                embedding: None,
1700            };
1701            storage.insert_event(&event).await.unwrap();
1702        }
1703
1704        // Different thread
1705        let event = AgentEvent {
1706            id: Uuid::now_v7(),
1707            agent_id: "agent-1".to_string(),
1708            thread_id: Some("thread-B".to_string()),
1709            run_id: None,
1710            parent_event_id: None,
1711            event_type: EventType::MemoryRead,
1712            payload: serde_json::json!({}),
1713            trace_id: None,
1714            span_id: None,
1715            model: None,
1716            tokens_input: None,
1717            tokens_output: None,
1718            latency_ms: None,
1719            cost_usd: None,
1720            timestamp: now.clone(),
1721            logical_clock: 0,
1722            content_hash: vec![99],
1723            prev_hash: None,
1724            embedding: None,
1725        };
1726        storage.insert_event(&event).await.unwrap();
1727
1728        let thread_a = storage.get_events_by_thread("thread-A", 10).await.unwrap();
1729        assert_eq!(thread_a.len(), 3);
1730
1731        let thread_b = storage.get_events_by_thread("thread-B", 10).await.unwrap();
1732        assert_eq!(thread_b.len(), 1);
1733        assert_eq!(thread_b[0].event_type, EventType::MemoryRead);
1734    }
1735
1736    #[tokio::test]
1737    async fn test_checkpoint_insert_and_get() {
1738        let storage = DuckDbStorage::open_in_memory().unwrap();
1739        let mem_id = Uuid::now_v7();
1740        let cp = Checkpoint {
1741            id: Uuid::now_v7(),
1742            thread_id: "thread-1".to_string(),
1743            agent_id: "agent-1".to_string(),
1744            parent_id: None,
1745            branch_name: "main".to_string(),
1746            state_snapshot: serde_json::json!({"step": 1}),
1747            state_diff: None,
1748            memory_refs: vec![mem_id],
1749            event_cursor: None,
1750            label: Some("initial".to_string()),
1751            created_at: chrono::Utc::now().to_rfc3339(),
1752            metadata: serde_json::json!({}),
1753        };
1754
1755        storage.insert_checkpoint(&cp).await.unwrap();
1756
1757        let fetched = storage.get_checkpoint(cp.id).await.unwrap().unwrap();
1758        assert_eq!(fetched.id, cp.id);
1759        assert_eq!(fetched.thread_id, "thread-1");
1760        assert_eq!(fetched.branch_name, "main");
1761        assert_eq!(fetched.memory_refs, vec![mem_id]);
1762        assert_eq!(fetched.label, Some("initial".to_string()));
1763    }
1764
1765    #[tokio::test]
1766    async fn test_checkpoint_list_and_latest() {
1767        let storage = DuckDbStorage::open_in_memory().unwrap();
1768
1769        let cp1 = Checkpoint {
1770            id: Uuid::now_v7(),
1771            thread_id: "thread-1".to_string(),
1772            agent_id: "agent-1".to_string(),
1773            parent_id: None,
1774            branch_name: "main".to_string(),
1775            state_snapshot: serde_json::json!({"step": 1}),
1776            state_diff: None,
1777            memory_refs: vec![],
1778            event_cursor: None,
1779            label: Some("first".to_string()),
1780            created_at: "2025-01-01T00:00:00Z".to_string(),
1781            metadata: serde_json::json!({}),
1782        };
1783        storage.insert_checkpoint(&cp1).await.unwrap();
1784
1785        let cp2 = Checkpoint {
1786            id: Uuid::now_v7(),
1787            thread_id: "thread-1".to_string(),
1788            agent_id: "agent-1".to_string(),
1789            parent_id: Some(cp1.id),
1790            branch_name: "main".to_string(),
1791            state_snapshot: serde_json::json!({"step": 2}),
1792            state_diff: Some(serde_json::json!({"step": [1, 2]})),
1793            memory_refs: vec![],
1794            event_cursor: None,
1795            label: Some("second".to_string()),
1796            created_at: "2025-01-02T00:00:00Z".to_string(),
1797            metadata: serde_json::json!({}),
1798        };
1799        storage.insert_checkpoint(&cp2).await.unwrap();
1800
1801        let cp3 = Checkpoint {
1802            id: Uuid::now_v7(),
1803            thread_id: "thread-1".to_string(),
1804            agent_id: "agent-1".to_string(),
1805            parent_id: Some(cp1.id),
1806            branch_name: "experiment".to_string(),
1807            state_snapshot: serde_json::json!({"step": "alt"}),
1808            state_diff: None,
1809            memory_refs: vec![],
1810            event_cursor: None,
1811            label: None,
1812            created_at: "2025-01-03T00:00:00Z".to_string(),
1813            metadata: serde_json::json!({}),
1814        };
1815        storage.insert_checkpoint(&cp3).await.unwrap();
1816
1817        // List all for thread
1818        let all = storage
1819            .list_checkpoints("thread-1", None, 10)
1820            .await
1821            .unwrap();
1822        assert_eq!(all.len(), 3);
1823
1824        // List by branch
1825        let main_cps = storage
1826            .list_checkpoints("thread-1", Some("main"), 10)
1827            .await
1828            .unwrap();
1829        assert_eq!(main_cps.len(), 2);
1830
1831        let exp_cps = storage
1832            .list_checkpoints("thread-1", Some("experiment"), 10)
1833            .await
1834            .unwrap();
1835        assert_eq!(exp_cps.len(), 1);
1836
1837        // Latest on main
1838        let latest = storage
1839            .get_latest_checkpoint("thread-1", "main")
1840            .await
1841            .unwrap()
1842            .unwrap();
1843        assert_eq!(latest.id, cp2.id);
1844
1845        // Latest on experiment
1846        let latest_exp = storage
1847            .get_latest_checkpoint("thread-1", "experiment")
1848            .await
1849            .unwrap()
1850            .unwrap();
1851        assert_eq!(latest_exp.id, cp3.id);
1852
1853        // No checkpoints for nonexistent branch
1854        let none = storage
1855            .get_latest_checkpoint("thread-1", "nonexistent")
1856            .await
1857            .unwrap();
1858        assert!(none.is_none());
1859    }
1860
1861    /// Embeddings must survive the byte round-trip exactly.
1862    ///
1863    /// Added when `deserialize_embedding` moved from `chunks_exact(4)` to
1864    /// `as_chunks::<4>()` for clippy 1.98. The two have identical semantics,
1865    /// including discarding a trailing partial chunk, but nothing in the suite
1866    /// asserted that: this pair is the only thing standing between a stored blob
1867    /// and a vector the index searches, and a silent corruption here would
1868    /// surface as bad recall rather than as an error.
1869    #[test]
1870    fn embeddings_round_trip_through_the_blob_encoding() {
1871        // Values chosen to catch byte-order and sign mistakes, which a
1872        // symmetric bug in both directions would otherwise hide.
1873        let original = vec![0.0_f32, 1.0, -1.0, 0.5, -0.25, f32::MIN, f32::MAX, 1e-30];
1874        let blob = serialize_embedding(&Some(original.clone())).expect("serialises");
1875        assert_eq!(
1876            blob.len(),
1877            original.len() * 4,
1878            "each f32 must occupy exactly 4 bytes"
1879        );
1880        let restored = deserialize_embedding(Some(blob)).expect("deserialises");
1881        assert_eq!(
1882            restored, original,
1883            "embedding did not survive the round-trip; stored vectors would no longer \
1884             match what was written and recall would degrade silently"
1885        );
1886    }
1887
1888    /// `None` stays `None`, and a trailing partial chunk is discarded rather
1889    /// than panicking on a short read.
1890    #[test]
1891    fn embedding_encoding_handles_absent_and_ragged_input() {
1892        assert!(serialize_embedding(&None).is_none());
1893        assert!(deserialize_embedding(None).is_none());
1894        assert_eq!(
1895            deserialize_embedding(Some(vec![])),
1896            Some(vec![]),
1897            "an empty blob is an empty embedding, not an error"
1898        );
1899        // 6 bytes is one whole f32 plus a 2-byte remainder.
1900        let ragged = deserialize_embedding(Some(vec![0, 0, 128, 63, 1, 2])).expect("some");
1901        assert_eq!(
1902            ragged,
1903            vec![1.0_f32],
1904            "a trailing partial chunk must be dropped, matching the previous \
1905             `chunks_exact` behaviour, rather than panicking or producing a value"
1906        );
1907    }
1908}