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 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 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 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 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 let acl_result = {
503 let conn = self.conn.lock().await;
504
505 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 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 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 }; 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 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 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 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 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 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 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 let fetched = storage.get_memory(record.id).await.unwrap().unwrap();
1496 assert!(fetched.deleted_at.is_some());
1497
1498 let filter = MemoryFilter::default();
1500 let list = storage.list_memories(&filter, 100, 0).await.unwrap();
1501 assert!(list.is_empty());
1502
1503 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 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 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 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 assert!(
1591 storage
1592 .check_permission(record.id, "agent-1", Permission::Admin)
1593 .await
1594 .unwrap()
1595 );
1596
1597 assert!(
1599 !storage
1600 .check_permission(record.id, "agent-2", Permission::Read)
1601 .await
1602 .unwrap()
1603 );
1604
1605 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 assert!(
1620 storage
1621 .check_permission(record.id, "agent-2", Permission::Read)
1622 .await
1623 .unwrap()
1624 );
1625 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 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 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 let all = storage
1819 .list_checkpoints("thread-1", None, 10)
1820 .await
1821 .unwrap();
1822 assert_eq!(all.len(), 3);
1823
1824 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 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 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 let none = storage
1855 .get_latest_checkpoint("thread-1", "nonexistent")
1856 .await
1857 .unwrap();
1858 assert!(none.is_none());
1859 }
1860
1861 #[test]
1870 fn embeddings_round_trip_through_the_blob_encoding() {
1871 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 #[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 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}