mnemo_postgres/
storage.rs

1use mnemo_core::error::{Error, Result};
2use mnemo_core::model::acl::{Acl, Permission};
3use mnemo_core::model::agent_profile::AgentProfile;
4use mnemo_core::model::checkpoint::Checkpoint;
5use mnemo_core::model::delegation::{Delegation, DelegationScope};
6use mnemo_core::model::embedding_baseline::EmbeddingBaseline;
7use mnemo_core::model::event::AgentEvent;
8use mnemo_core::model::memory::MemoryRecord;
9use mnemo_core::model::relation::Relation;
10use mnemo_core::storage::{MemoryFilter, StorageBackend};
11use pgvector::Vector;
12use sqlx::Row;
13use uuid::Uuid;
14
15/// PostgreSQL-backed storage for Mnemo.
16///
17/// Wraps a `sqlx::PgPool` and runs schema migrations on construction.
18/// Embeddings are stored using the pgvector `vector` column type, while
19/// event embeddings are stored as `BYTEA` (serialised `Vec<f32>` in
20/// little-endian byte order), matching the DuckDB backend convention.
21pub struct PgStorage {
22    pool: sqlx::PgPool,
23    dimensions: usize,
24}
25
26impl PgStorage {
27    /// Connect to a PostgreSQL database and run migrations.
28    ///
29    /// `url` is a standard `postgres://` connection string.
30    /// `dimensions` controls the width of the pgvector `vector` column.
31    pub async fn connect(url: &str, dimensions: usize) -> Result<Self> {
32        let pool = sqlx::PgPool::connect(url)
33            .await
34            .map_err(|e| Error::Storage(e.to_string()))?;
35        let storage = Self { pool, dimensions };
36        crate::migrations::run_migrations(&storage.pool, dimensions).await?;
37        Ok(storage)
38    }
39
40    /// Build a `PgStorage` from an existing pool (useful for tests).
41    pub async fn from_pool(pool: sqlx::PgPool, dimensions: usize) -> Result<Self> {
42        crate::migrations::run_migrations(&pool, dimensions).await?;
43        Ok(Self { pool, dimensions })
44    }
45
46    /// A clone of the connection pool, so a [`crate::PgVectorIndex`] can share
47    /// the same connections for ANN search. `sqlx::PgPool` is `Arc`-backed, so
48    /// the clone is cheap and points at the same pool.
49    pub fn pool(&self) -> sqlx::PgPool {
50        self.pool.clone()
51    }
52
53    /// The pgvector `vector(dim)` column width this storage was migrated with.
54    pub fn dimensions(&self) -> usize {
55        self.dimensions
56    }
57}
58
59// ---------------------------------------------------------------------------
60// Helpers
61// ---------------------------------------------------------------------------
62
63fn map_sqlx(e: sqlx::Error) -> Error {
64    Error::Storage(e.to_string())
65}
66
67fn serialize_embedding(embedding: &Option<Vec<f32>>) -> Option<Vec<u8>> {
68    embedding
69        .as_ref()
70        .map(|v| v.iter().flat_map(|f| f.to_le_bytes()).collect())
71}
72
73fn deserialize_embedding(blob: Option<Vec<u8>>) -> Option<Vec<f32>> {
74    blob.map(|bytes| {
75        bytes
76            .chunks_exact(4)
77            .map(|chunk| f32::from_le_bytes([chunk[0], chunk[1], chunk[2], chunk[3]]))
78            .collect()
79    })
80}
81
82fn row_to_memory(row: &sqlx::postgres::PgRow) -> std::result::Result<MemoryRecord, sqlx::Error> {
83    let tags: Vec<String> = row.try_get::<Vec<String>, _>("tags").unwrap_or_default();
84    let metadata: serde_json::Value = row
85        .try_get("metadata")
86        .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
87
88    // pgvector stores the embedding as its own type; we retrieve the raw text
89    // representation and parse back to Vec<f32>. If the column is NULL we get None.
90    let embedding: Option<Vec<f32>> = {
91        let raw: Option<String> = row.try_get("embedding_text").ok().flatten();
92        raw.and_then(|s| {
93            // pgvector text output looks like "[0.1,0.2,0.3]"
94            let trimmed = s.trim_start_matches('[').trim_end_matches(']');
95            if trimmed.is_empty() {
96                None
97            } else {
98                Some(
99                    trimmed
100                        .split(',')
101                        .filter_map(|v| v.trim().parse::<f32>().ok())
102                        .collect(),
103                )
104            }
105        })
106    };
107
108    Ok(MemoryRecord {
109        id: row.get("id"),
110        agent_id: row.get("agent_id"),
111        content: row.get("content"),
112        memory_type: row
113            .get::<String, _>("memory_type")
114            .parse()
115            .unwrap_or(mnemo_core::model::memory::MemoryType::Semantic),
116        scope: row
117            .get::<String, _>("scope")
118            .parse()
119            .unwrap_or(mnemo_core::model::memory::Scope::Private),
120        importance: row.get("importance"),
121        tags,
122        metadata,
123        embedding,
124        content_hash: row.get("content_hash"),
125        prev_hash: row.get("prev_hash"),
126        source_type: row
127            .get::<String, _>("source_type")
128            .parse()
129            .unwrap_or(mnemo_core::model::memory::SourceType::Agent),
130        source_id: row.get("source_id"),
131        consolidation_state: row
132            .get::<String, _>("consolidation_state")
133            .parse()
134            .unwrap_or(mnemo_core::model::memory::ConsolidationState::Raw),
135        access_count: row.get::<i64, _>("access_count") as u64,
136        org_id: row.get("org_id"),
137        thread_id: row.get("thread_id"),
138        created_at: row.get("created_at"),
139        updated_at: row.get("updated_at"),
140        last_accessed_at: row.get("last_accessed_at"),
141        expires_at: row.get("expires_at"),
142        deleted_at: row.get("deleted_at"),
143        decay_rate: row.get("decay_rate"),
144        created_by: row.get("created_by"),
145        version: row.get::<i32, _>("version") as u32,
146        prev_version_id: row.get("prev_version_id"),
147        quarantined: row.get("quarantined"),
148        quarantine_reason: row.get("quarantine_reason"),
149        decay_function: row.get("decay_function"),
150    })
151}
152
153/// The standard SELECT column list for the memories table.
154/// We cast the pgvector `embedding` column to text so we can parse it
155/// back into `Vec<f32>` without depending on a pgvector Rust decode path.
156///
157/// NOTE (sqlx 0.9 `SqlSafeStr`): queries that interpolate this const into
158/// a `format!`ed SQL string are wrapped in `sqlx::AssertSqlSafe`. This is
159/// audited-safe: the only interpolated values are this column-list const,
160/// code-built `$N` placeholder fragments, and numeric `usize` limit/offset
161/// — all caller data is bound via `$N`, never string-interpolated.
162const MEMORY_COLUMNS: &str = r#"
163    id, agent_id, content, memory_type, scope, importance,
164    tags, metadata, embedding::text AS embedding_text,
165    content_hash, prev_hash, source_type, source_id,
166    consolidation_state, access_count, org_id, thread_id,
167    created_at, updated_at, last_accessed_at, expires_at,
168    deleted_at, decay_rate, created_by, version, prev_version_id,
169    quarantined, quarantine_reason, decay_function
170"#;
171
172fn row_to_event(row: &sqlx::postgres::PgRow) -> std::result::Result<AgentEvent, sqlx::Error> {
173    let payload: serde_json::Value = row.try_get("payload").unwrap_or(serde_json::Value::Null);
174    let embedding_blob: Option<Vec<u8>> = row.try_get("embedding").unwrap_or(None);
175
176    Ok(AgentEvent {
177        id: row.get("id"),
178        agent_id: row.get("agent_id"),
179        thread_id: row.get("thread_id"),
180        run_id: row.get("run_id"),
181        parent_event_id: row.get("parent_event_id"),
182        event_type: row
183            .get::<String, _>("event_type")
184            .parse()
185            .unwrap_or(mnemo_core::model::event::EventType::Error),
186        payload,
187        trace_id: row.get("trace_id"),
188        span_id: row.get("span_id"),
189        model: row.get("model"),
190        tokens_input: row.get("tokens_input"),
191        tokens_output: row.get("tokens_output"),
192        latency_ms: row.get("latency_ms"),
193        cost_usd: row.get("cost_usd"),
194        timestamp: row.get("timestamp"),
195        logical_clock: row.get("logical_clock"),
196        content_hash: row.get("content_hash"),
197        prev_hash: row.get("prev_hash"),
198        embedding: deserialize_embedding(embedding_blob),
199    })
200}
201
202fn row_to_relation(row: &sqlx::postgres::PgRow) -> std::result::Result<Relation, sqlx::Error> {
203    let metadata: serde_json::Value = row
204        .try_get("metadata")
205        .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
206
207    Ok(Relation {
208        id: row.get("id"),
209        source_id: row.get("source_id"),
210        target_id: row.get("target_id"),
211        relation_type: row.get("relation_type"),
212        weight: row.get("weight"),
213        metadata,
214        created_at: row.get("created_at"),
215    })
216}
217
218fn row_to_checkpoint(row: &sqlx::postgres::PgRow) -> std::result::Result<Checkpoint, sqlx::Error> {
219    let state_snapshot: serde_json::Value = row
220        .try_get("state_snapshot")
221        .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
222    let state_diff: Option<serde_json::Value> = row.try_get("state_diff").unwrap_or(None);
223
224    // memory_refs is stored as TEXT[] of UUID strings
225    let memory_refs_raw: Vec<String> = row.try_get("memory_refs").unwrap_or_default();
226    let memory_refs: Vec<Uuid> = memory_refs_raw
227        .iter()
228        .filter_map(|s| Uuid::parse_str(s).ok())
229        .collect();
230
231    let metadata: serde_json::Value = row
232        .try_get("metadata")
233        .unwrap_or(serde_json::Value::Object(serde_json::Map::new()));
234
235    Ok(Checkpoint {
236        id: row.get("id"),
237        thread_id: row.get("thread_id"),
238        agent_id: row.get("agent_id"),
239        parent_id: row.get("parent_id"),
240        branch_name: row.get("branch_name"),
241        state_snapshot,
242        state_diff,
243        memory_refs,
244        event_cursor: row.get("event_cursor"),
245        label: row.get("label"),
246        created_at: row.get("created_at"),
247        metadata,
248    })
249}
250
251fn row_to_delegation(row: &sqlx::postgres::PgRow) -> std::result::Result<Delegation, sqlx::Error> {
252    let scope_type: String = row.get("scope_type");
253    let scope_value: Option<serde_json::Value> = row.try_get("scope_value").unwrap_or(None);
254
255    let scope = match scope_type.as_str() {
256        "by_tag" => {
257            let tags: Vec<String> = scope_value
258                .and_then(|v| serde_json::from_value(v).ok())
259                .unwrap_or_default();
260            DelegationScope::ByTag(tags)
261        }
262        "by_memory_id" => {
263            let id_strs: Vec<String> = scope_value
264                .and_then(|v| serde_json::from_value(v).ok())
265                .unwrap_or_default();
266            let uuids = id_strs
267                .into_iter()
268                .filter_map(|s| Uuid::parse_str(&s).ok())
269                .collect();
270            DelegationScope::ByMemoryId(uuids)
271        }
272        _ => DelegationScope::AllMemories,
273    };
274
275    Ok(Delegation {
276        id: row.get("id"),
277        delegator_id: row.get("delegator_id"),
278        delegate_id: row.get("delegate_id"),
279        permission: row
280            .get::<String, _>("permission")
281            .parse()
282            .unwrap_or(Permission::Read),
283        scope,
284        max_depth: row.get::<i32, _>("max_depth") as u32,
285        current_depth: row.get::<i32, _>("current_depth") as u32,
286        parent_delegation_id: row.get("parent_delegation_id"),
287        created_at: row.get("created_at"),
288        expires_at: row.get("expires_at"),
289        revoked_at: row.get("revoked_at"),
290    })
291}
292
293// ---------------------------------------------------------------------------
294// StorageBackend implementation
295// ---------------------------------------------------------------------------
296
297#[async_trait::async_trait]
298impl StorageBackend for PgStorage {
299    fn backend_name(&self) -> &'static str {
300        "postgres"
301    }
302
303    // -----------------------------------------------------------------------
304    // Memory CRUD
305    // -----------------------------------------------------------------------
306
307    async fn insert_memory(&self, record: &MemoryRecord) -> Result<()> {
308        let embedding_param: Option<Vector> =
309            record.embedding.as_ref().map(|v| Vector::from(v.clone()));
310
311        let tags_slice: &[String] = &record.tags;
312
313        sqlx::query(
314            r#"
315INSERT INTO memories (
316    id, agent_id, content, memory_type, scope, importance,
317    tags, metadata, embedding,
318    content_hash, prev_hash, source_type, source_id,
319    consolidation_state, access_count, org_id, thread_id,
320    created_at, updated_at, last_accessed_at, expires_at,
321    deleted_at, decay_rate, created_by, version, prev_version_id,
322    quarantined, quarantine_reason, decay_function
323) VALUES (
324    $1, $2, $3, $4, $5, $6,
325    $7, $8, $9,
326    $10, $11, $12, $13,
327    $14, $15, $16, $17,
328    $18, $19, $20, $21,
329    $22, $23, $24, $25, $26,
330    $27, $28, $29
331)
332"#,
333        )
334        .bind(record.id)
335        .bind(&record.agent_id)
336        .bind(&record.content)
337        .bind(record.memory_type.to_string())
338        .bind(record.scope.to_string())
339        .bind(record.importance)
340        .bind(tags_slice)
341        .bind(&record.metadata)
342        .bind(&embedding_param)
343        .bind(&record.content_hash)
344        .bind(&record.prev_hash)
345        .bind(record.source_type.to_string())
346        .bind(&record.source_id)
347        .bind(record.consolidation_state.to_string())
348        .bind(record.access_count as i64)
349        .bind(&record.org_id)
350        .bind(&record.thread_id)
351        .bind(&record.created_at)
352        .bind(&record.updated_at)
353        .bind(&record.last_accessed_at)
354        .bind(&record.expires_at)
355        .bind(&record.deleted_at)
356        .bind(record.decay_rate)
357        .bind(&record.created_by)
358        .bind(record.version as i32)
359        .bind(record.prev_version_id)
360        .bind(record.quarantined)
361        .bind(&record.quarantine_reason)
362        .bind(&record.decay_function)
363        .execute(&self.pool)
364        .await
365        .map_err(map_sqlx)?;
366
367        Ok(())
368    }
369
370    async fn get_memory(&self, id: Uuid) -> Result<Option<MemoryRecord>> {
371        let sql = format!("SELECT {MEMORY_COLUMNS} FROM memories WHERE id = $1");
372        let row = sqlx::query(sqlx::AssertSqlSafe(sql.as_str()))
373            .bind(id)
374            .fetch_optional(&self.pool)
375            .await
376            .map_err(map_sqlx)?;
377
378        match row {
379            Some(r) => Ok(Some(row_to_memory(&r).map_err(map_sqlx)?)),
380            None => Ok(None),
381        }
382    }
383
384    async fn update_memory(&self, record: &MemoryRecord) -> Result<()> {
385        let embedding_param: Option<Vector> =
386            record.embedding.as_ref().map(|v| Vector::from(v.clone()));
387
388        let tags_slice: &[String] = &record.tags;
389
390        let result = sqlx::query(
391            r#"
392UPDATE memories SET
393    agent_id = $1, content = $2, memory_type = $3, scope = $4,
394    importance = $5, tags = $6, metadata = $7,
395    embedding = $8,
396    content_hash = $9, prev_hash = $10, source_type = $11,
397    source_id = $12, consolidation_state = $13, access_count = $14,
398    org_id = $15, thread_id = $16, updated_at = $17,
399    last_accessed_at = $18, expires_at = $19, deleted_at = $20,
400    decay_rate = $21, created_by = $22, version = $23,
401    prev_version_id = $24, quarantined = $25, quarantine_reason = $26,
402    decay_function = $27
403WHERE id = $28
404"#,
405        )
406        .bind(&record.agent_id)
407        .bind(&record.content)
408        .bind(record.memory_type.to_string())
409        .bind(record.scope.to_string())
410        .bind(record.importance)
411        .bind(tags_slice)
412        .bind(&record.metadata)
413        .bind(&embedding_param)
414        .bind(&record.content_hash)
415        .bind(&record.prev_hash)
416        .bind(record.source_type.to_string())
417        .bind(&record.source_id)
418        .bind(record.consolidation_state.to_string())
419        .bind(record.access_count as i64)
420        .bind(&record.org_id)
421        .bind(&record.thread_id)
422        .bind(&record.updated_at)
423        .bind(&record.last_accessed_at)
424        .bind(&record.expires_at)
425        .bind(&record.deleted_at)
426        .bind(record.decay_rate)
427        .bind(&record.created_by)
428        .bind(record.version as i32)
429        .bind(record.prev_version_id)
430        .bind(record.quarantined)
431        .bind(&record.quarantine_reason)
432        .bind(&record.decay_function)
433        .bind(record.id)
434        .execute(&self.pool)
435        .await
436        .map_err(map_sqlx)?;
437
438        if result.rows_affected() == 0 {
439            return Err(Error::NotFound(format!("memory {} not found", record.id)));
440        }
441        Ok(())
442    }
443
444    async fn soft_delete_memory(&self, id: Uuid) -> Result<()> {
445        let now = chrono::Utc::now().to_rfc3339();
446        let result = sqlx::query(
447            "UPDATE memories SET deleted_at = $1, updated_at = $2 WHERE id = $3 AND deleted_at IS NULL",
448        )
449        .bind(&now)
450        .bind(&now)
451        .bind(id)
452        .execute(&self.pool)
453        .await
454        .map_err(map_sqlx)?;
455
456        if result.rows_affected() == 0 {
457            return Err(Error::NotFound(format!(
458                "memory {id} not found or already deleted"
459            )));
460        }
461        Ok(())
462    }
463
464    async fn hard_delete_memory(&self, id: Uuid) -> Result<()> {
465        let result = sqlx::query("DELETE FROM memories WHERE id = $1")
466            .bind(id)
467            .execute(&self.pool)
468            .await
469            .map_err(map_sqlx)?;
470
471        if result.rows_affected() == 0 {
472            return Err(Error::NotFound(format!("memory {id} not found")));
473        }
474
475        // Clean up ACLs for this memory
476        sqlx::query("DELETE FROM acls WHERE memory_id = $1")
477            .bind(id)
478            .execute(&self.pool)
479            .await
480            .map_err(map_sqlx)?;
481
482        Ok(())
483    }
484
485    async fn list_memories(
486        &self,
487        filter: &MemoryFilter,
488        limit: usize,
489        offset: usize,
490    ) -> Result<Vec<MemoryRecord>> {
491        let mut conditions: Vec<String> = Vec::new();
492        // We'll track bind-parameter index manually.
493        // The MEMORY_COLUMNS select doesn't use numbered params.
494        let mut param_idx: usize = 0;
495
496        // We accumulate bind values in a specific order and push them later
497        // via a dynamic query builder. Unfortunately sqlx's dynamic queries
498        // require us to build the SQL string with numbered placeholders and
499        // bind all values in order.
500
501        // We'll collect (sql_fragment, value_type) tuples, then bind them.
502        // Use a simpler approach: build the query string, then bind
503        // parameters positionally.
504
505        if !filter.include_deleted {
506            conditions.push("deleted_at IS NULL".to_string());
507        }
508
509        // We'll use an enum-based approach below to track what to bind.
510        #[derive(Debug)]
511        enum Param {
512            Str(String),
513            F32(f32),
514        }
515        let mut params: Vec<Param> = Vec::new();
516
517        if let Some(ref agent_id) = filter.agent_id {
518            param_idx += 1;
519            conditions.push(format!("agent_id = ${param_idx}"));
520            params.push(Param::Str(agent_id.clone()));
521        }
522        if let Some(memory_type) = filter.memory_type {
523            param_idx += 1;
524            conditions.push(format!("memory_type = ${param_idx}"));
525            params.push(Param::Str(memory_type.to_string()));
526        }
527        if let Some(scope) = filter.scope {
528            param_idx += 1;
529            conditions.push(format!("scope = ${param_idx}"));
530            params.push(Param::Str(scope.to_string()));
531        }
532        if let Some(min_importance) = filter.min_importance {
533            param_idx += 1;
534            conditions.push(format!("importance >= ${param_idx}"));
535            params.push(Param::F32(min_importance));
536        }
537        if let Some(ref org_id) = filter.org_id {
538            param_idx += 1;
539            conditions.push(format!("org_id = ${param_idx}"));
540            params.push(Param::Str(org_id.clone()));
541        }
542        if let Some(ref thread_id) = filter.thread_id {
543            param_idx += 1;
544            conditions.push(format!("thread_id = ${param_idx}"));
545            params.push(Param::Str(thread_id.clone()));
546        }
547
548        let where_clause = if conditions.is_empty() {
549            String::new()
550        } else {
551            format!("WHERE {}", conditions.join(" AND "))
552        };
553
554        let sql = format!(
555            "SELECT {MEMORY_COLUMNS} FROM memories {where_clause} ORDER BY created_at DESC LIMIT {limit} OFFSET {offset}"
556        );
557
558        let mut query = sqlx::query(sqlx::AssertSqlSafe(sql.as_str()));
559        for p in &params {
560            match p {
561                Param::Str(s) => query = query.bind(s),
562                Param::F32(f) => query = query.bind(*f),
563            }
564        }
565
566        let rows = query.fetch_all(&self.pool).await.map_err(map_sqlx)?;
567        let mut results = Vec::with_capacity(rows.len());
568        for r in &rows {
569            results.push(row_to_memory(r).map_err(map_sqlx)?);
570        }
571        Ok(results)
572    }
573
574    async fn touch_memory(&self, id: Uuid) -> Result<()> {
575        let now = chrono::Utc::now().to_rfc3339();
576        sqlx::query(
577            "UPDATE memories SET access_count = access_count + 1, last_accessed_at = $1 WHERE id = $2",
578        )
579        .bind(&now)
580        .bind(id)
581        .execute(&self.pool)
582        .await
583        .map_err(map_sqlx)?;
584        Ok(())
585    }
586
587    // -----------------------------------------------------------------------
588    // ACL
589    // -----------------------------------------------------------------------
590
591    async fn insert_acl(&self, acl: &Acl) -> Result<()> {
592        sqlx::query(
593            r#"
594INSERT INTO acls (id, memory_id, principal_type, principal_id, permission, granted_by, created_at, expires_at)
595VALUES ($1, $2, $3, $4, $5, $6, $7, $8)
596"#,
597        )
598        .bind(acl.id)
599        .bind(acl.memory_id)
600        .bind(acl.principal_type.to_string())
601        .bind(&acl.principal_id)
602        .bind(acl.permission.to_string())
603        .bind(&acl.granted_by)
604        .bind(&acl.created_at)
605        .bind(&acl.expires_at)
606        .execute(&self.pool)
607        .await
608        .map_err(map_sqlx)?;
609        Ok(())
610    }
611
612    async fn check_permission(
613        &self,
614        memory_id: Uuid,
615        principal_id: &str,
616        required: Permission,
617    ) -> Result<bool> {
618        // Check if the principal is the owner
619        let owner_row = sqlx::query("SELECT agent_id FROM memories WHERE id = $1")
620            .bind(memory_id)
621            .fetch_optional(&self.pool)
622            .await
623            .map_err(map_sqlx)?;
624
625        match owner_row {
626            None => return Err(Error::NotFound(format!("memory {memory_id} not found"))),
627            Some(row) => {
628                let owner: String = row.get("agent_id");
629                if owner == principal_id {
630                    return Ok(true);
631                }
632            }
633        }
634
635        // Check ACLs (direct grants)
636        let now = chrono::Utc::now().to_rfc3339();
637        let acl_rows = sqlx::query(
638            "SELECT permission FROM acls WHERE memory_id = $1 AND principal_id = $2 AND (expires_at IS NULL OR expires_at > $3)",
639        )
640        .bind(memory_id)
641        .bind(principal_id)
642        .bind(&now)
643        .fetch_all(&self.pool)
644        .await
645        .map_err(map_sqlx)?;
646
647        for row in &acl_rows {
648            let perm_str: String = row.get("permission");
649            if let Ok(perm) = perm_str.parse::<Permission>()
650                && perm.satisfies(required)
651            {
652                return Ok(true);
653            }
654        }
655
656        // Check public ACLs
657        let public_rows = sqlx::query(
658            "SELECT permission FROM acls WHERE memory_id = $1 AND principal_type = 'public' AND (expires_at IS NULL OR expires_at > $2)",
659        )
660        .bind(memory_id)
661        .bind(&now)
662        .fetch_all(&self.pool)
663        .await
664        .map_err(map_sqlx)?;
665
666        for row in &public_rows {
667            let perm_str: String = row.get("permission");
668            if let Ok(perm) = perm_str.parse::<Permission>()
669                && perm.satisfies(required)
670            {
671                return Ok(true);
672            }
673        }
674
675        // Check delegations
676        if self
677            .check_delegation(principal_id, memory_id, required)
678            .await?
679        {
680            return Ok(true);
681        }
682
683        Ok(false)
684    }
685
686    // -----------------------------------------------------------------------
687    // Relations
688    // -----------------------------------------------------------------------
689
690    async fn insert_relation(&self, relation: &Relation) -> Result<()> {
691        sqlx::query(
692            r#"
693INSERT INTO relations (id, source_id, target_id, relation_type, weight, metadata, created_at)
694VALUES ($1, $2, $3, $4, $5, $6, $7)
695"#,
696        )
697        .bind(relation.id)
698        .bind(relation.source_id)
699        .bind(relation.target_id)
700        .bind(&relation.relation_type)
701        .bind(relation.weight)
702        .bind(&relation.metadata)
703        .bind(&relation.created_at)
704        .execute(&self.pool)
705        .await
706        .map_err(map_sqlx)?;
707        Ok(())
708    }
709
710    async fn get_relations_from(&self, source_id: Uuid) -> Result<Vec<Relation>> {
711        let rows = sqlx::query(
712            "SELECT id, source_id, target_id, relation_type, weight, metadata, created_at FROM relations WHERE source_id = $1",
713        )
714        .bind(source_id)
715        .fetch_all(&self.pool)
716        .await
717        .map_err(map_sqlx)?;
718
719        let mut results = Vec::with_capacity(rows.len());
720        for r in &rows {
721            results.push(row_to_relation(r).map_err(map_sqlx)?);
722        }
723        Ok(results)
724    }
725
726    async fn get_relations_to(&self, target_id: Uuid) -> Result<Vec<Relation>> {
727        let rows = sqlx::query(
728            "SELECT id, source_id, target_id, relation_type, weight, metadata, created_at FROM relations WHERE target_id = $1",
729        )
730        .bind(target_id)
731        .fetch_all(&self.pool)
732        .await
733        .map_err(map_sqlx)?;
734
735        let mut results = Vec::with_capacity(rows.len());
736        for r in &rows {
737            results.push(row_to_relation(r).map_err(map_sqlx)?);
738        }
739        Ok(results)
740    }
741
742    async fn delete_relation(&self, id: Uuid) -> Result<()> {
743        let result = sqlx::query("DELETE FROM relations WHERE id = $1")
744            .bind(id)
745            .execute(&self.pool)
746            .await
747            .map_err(map_sqlx)?;
748
749        if result.rows_affected() == 0 {
750            return Err(Error::NotFound(format!("relation {id} not found")));
751        }
752        Ok(())
753    }
754
755    // -----------------------------------------------------------------------
756    // Chain linking
757    // -----------------------------------------------------------------------
758
759    async fn get_latest_memory_hash(
760        &self,
761        agent_id: &str,
762        thread_id: Option<&str>,
763    ) -> Result<Option<Vec<u8>>> {
764        let row = if let Some(tid) = thread_id {
765            sqlx::query(
766                "SELECT content_hash FROM memories WHERE agent_id = $1 AND thread_id = $2 AND deleted_at IS NULL ORDER BY created_at DESC LIMIT 1",
767            )
768            .bind(agent_id)
769            .bind(tid)
770            .fetch_optional(&self.pool)
771            .await
772            .map_err(map_sqlx)?
773        } else {
774            sqlx::query(
775                "SELECT content_hash FROM memories WHERE agent_id = $1 AND thread_id IS NULL AND deleted_at IS NULL ORDER BY created_at DESC LIMIT 1",
776            )
777            .bind(agent_id)
778            .fetch_optional(&self.pool)
779            .await
780            .map_err(map_sqlx)?
781        };
782
783        Ok(row.map(|r| r.get::<Vec<u8>, _>("content_hash")))
784    }
785
786    async fn get_latest_event_hash(
787        &self,
788        agent_id: &str,
789        thread_id: Option<&str>,
790    ) -> Result<Option<Vec<u8>>> {
791        let row = if let Some(tid) = thread_id {
792            sqlx::query(
793                "SELECT content_hash FROM agent_events WHERE agent_id = $1 AND thread_id = $2 ORDER BY timestamp DESC LIMIT 1",
794            )
795            .bind(agent_id)
796            .bind(tid)
797            .fetch_optional(&self.pool)
798            .await
799            .map_err(map_sqlx)?
800        } else {
801            sqlx::query(
802                "SELECT content_hash FROM agent_events WHERE agent_id = $1 ORDER BY timestamp DESC LIMIT 1",
803            )
804            .bind(agent_id)
805            .fetch_optional(&self.pool)
806            .await
807            .map_err(map_sqlx)?
808        };
809        Ok(row.map(|r| r.get::<Vec<u8>, _>("content_hash")))
810    }
811
812    async fn get_sync_watermark(&self, key: &str) -> Result<Option<String>> {
813        let row = sqlx::query("SELECT value FROM sync_metadata WHERE key = $1")
814            .bind(key)
815            .fetch_optional(&self.pool)
816            .await
817            .map_err(map_sqlx)?;
818        Ok(row.map(|r| r.get::<String, _>("value")))
819    }
820
821    async fn set_sync_watermark(&self, key: &str, value: &str) -> Result<()> {
822        let now = chrono::Utc::now().to_rfc3339();
823        sqlx::query(
824            "INSERT INTO sync_metadata (key, value, updated_at) VALUES ($1, $2, $3) ON CONFLICT (key) DO UPDATE SET value = $2, updated_at = $3",
825        )
826        .bind(key)
827        .bind(value)
828        .bind(now)
829        .execute(&self.pool)
830        .await
831        .map_err(map_sqlx)?;
832        Ok(())
833    }
834
835    // -----------------------------------------------------------------------
836    // Permission-safe ANN
837    // -----------------------------------------------------------------------
838
839    async fn list_accessible_memory_ids(&self, agent_id: &str, limit: usize) -> Result<Vec<Uuid>> {
840        let now = chrono::Utc::now().to_rfc3339();
841        let rows = sqlx::query(
842            r#"
843SELECT id FROM memories
844WHERE (
845    agent_id = $1
846    OR scope = 'public'
847    OR id IN (
848        SELECT memory_id FROM acls
849        WHERE principal_id = $2 AND (expires_at IS NULL OR expires_at > $3)
850    )
851)
852AND deleted_at IS NULL
853LIMIT $4
854"#,
855        )
856        .bind(agent_id)
857        .bind(agent_id)
858        .bind(&now)
859        .bind(limit as i64)
860        .fetch_all(&self.pool)
861        .await
862        .map_err(map_sqlx)?;
863
864        let ids: Vec<Uuid> = rows.iter().map(|r| r.get("id")).collect();
865        Ok(ids)
866    }
867
868    // -----------------------------------------------------------------------
869    // Events
870    // -----------------------------------------------------------------------
871
872    async fn insert_event(&self, event: &AgentEvent) -> Result<()> {
873        let payload_json = &event.payload;
874        let embedding_blob = serialize_embedding(&event.embedding);
875
876        sqlx::query(
877            r#"
878INSERT INTO agent_events (
879    id, agent_id, thread_id, run_id, parent_event_id, event_type,
880    payload, trace_id, span_id, model, tokens_input, tokens_output,
881    latency_ms, cost_usd, "timestamp", logical_clock, content_hash,
882    prev_hash, embedding
883) VALUES (
884    $1, $2, $3, $4, $5, $6,
885    $7, $8, $9, $10, $11, $12,
886    $13, $14, $15, $16, $17,
887    $18, $19
888)
889"#,
890        )
891        .bind(event.id)
892        .bind(&event.agent_id)
893        .bind(&event.thread_id)
894        .bind(&event.run_id)
895        .bind(event.parent_event_id)
896        .bind(event.event_type.to_string())
897        .bind(payload_json)
898        .bind(&event.trace_id)
899        .bind(&event.span_id)
900        .bind(&event.model)
901        .bind(event.tokens_input)
902        .bind(event.tokens_output)
903        .bind(event.latency_ms)
904        .bind(event.cost_usd)
905        .bind(&event.timestamp)
906        .bind(event.logical_clock)
907        .bind(&event.content_hash)
908        .bind(&event.prev_hash)
909        .bind(&embedding_blob)
910        .execute(&self.pool)
911        .await
912        .map_err(map_sqlx)?;
913        Ok(())
914    }
915
916    async fn list_events(
917        &self,
918        agent_id: &str,
919        limit: usize,
920        offset: usize,
921    ) -> Result<Vec<AgentEvent>> {
922        let rows = sqlx::query(
923            r#"
924SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type,
925       payload, trace_id, span_id, model, tokens_input, tokens_output,
926       latency_ms, cost_usd, "timestamp", logical_clock, content_hash,
927       prev_hash, embedding
928FROM agent_events
929WHERE agent_id = $1
930ORDER BY "timestamp" DESC
931LIMIT $2 OFFSET $3
932"#,
933        )
934        .bind(agent_id)
935        .bind(limit as i64)
936        .bind(offset as i64)
937        .fetch_all(&self.pool)
938        .await
939        .map_err(map_sqlx)?;
940
941        let mut results = Vec::with_capacity(rows.len());
942        for r in &rows {
943            results.push(row_to_event(r).map_err(map_sqlx)?);
944        }
945        Ok(results)
946    }
947
948    async fn get_events_by_thread(&self, thread_id: &str, limit: usize) -> Result<Vec<AgentEvent>> {
949        let rows = sqlx::query(
950            r#"
951SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type,
952       payload, trace_id, span_id, model, tokens_input, tokens_output,
953       latency_ms, cost_usd, "timestamp", logical_clock, content_hash,
954       prev_hash, embedding
955FROM agent_events
956WHERE thread_id = $1
957ORDER BY "timestamp" ASC
958LIMIT $2
959"#,
960        )
961        .bind(thread_id)
962        .bind(limit as i64)
963        .fetch_all(&self.pool)
964        .await
965        .map_err(map_sqlx)?;
966
967        let mut results = Vec::with_capacity(rows.len());
968        for r in &rows {
969            results.push(row_to_event(r).map_err(map_sqlx)?);
970        }
971        Ok(results)
972    }
973
974    async fn get_event(&self, id: Uuid) -> Result<Option<AgentEvent>> {
975        let row = sqlx::query(
976            r#"
977SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type,
978       payload, trace_id, span_id, model, tokens_input, tokens_output,
979       latency_ms, cost_usd, "timestamp", logical_clock, content_hash,
980       prev_hash, embedding
981FROM agent_events
982WHERE id = $1
983"#,
984        )
985        .bind(id)
986        .fetch_optional(&self.pool)
987        .await
988        .map_err(map_sqlx)?;
989
990        match row {
991            Some(r) => Ok(Some(row_to_event(&r).map_err(map_sqlx)?)),
992            None => Ok(None),
993        }
994    }
995
996    async fn list_child_events(
997        &self,
998        parent_event_id: Uuid,
999        limit: usize,
1000    ) -> Result<Vec<AgentEvent>> {
1001        let rows = sqlx::query(
1002            r#"
1003SELECT id, agent_id, thread_id, run_id, parent_event_id, event_type,
1004       payload, trace_id, span_id, model, tokens_input, tokens_output,
1005       latency_ms, cost_usd, "timestamp", logical_clock, content_hash,
1006       prev_hash, embedding
1007FROM agent_events
1008WHERE parent_event_id = $1
1009ORDER BY "timestamp" ASC
1010LIMIT $2
1011"#,
1012        )
1013        .bind(parent_event_id)
1014        .bind(limit as i64)
1015        .fetch_all(&self.pool)
1016        .await
1017        .map_err(map_sqlx)?;
1018
1019        let mut results = Vec::with_capacity(rows.len());
1020        for r in &rows {
1021            results.push(row_to_event(r).map_err(map_sqlx)?);
1022        }
1023        Ok(results)
1024    }
1025
1026    // -----------------------------------------------------------------------
1027    // Ordered listing
1028    // -----------------------------------------------------------------------
1029
1030    async fn list_memories_by_agent_ordered(
1031        &self,
1032        agent_id: &str,
1033        thread_id: Option<&str>,
1034        limit: usize,
1035    ) -> Result<Vec<MemoryRecord>> {
1036        let rows = if let Some(tid) = thread_id {
1037            let sql = format!(
1038                "SELECT {MEMORY_COLUMNS} FROM memories WHERE agent_id = $1 AND thread_id = $2 AND deleted_at IS NULL ORDER BY created_at ASC LIMIT $3"
1039            );
1040            sqlx::query(sqlx::AssertSqlSafe(sql.as_str()))
1041                .bind(agent_id)
1042                .bind(tid)
1043                .bind(limit as i64)
1044                .fetch_all(&self.pool)
1045                .await
1046                .map_err(map_sqlx)?
1047        } else {
1048            let sql = format!(
1049                "SELECT {MEMORY_COLUMNS} FROM memories WHERE agent_id = $1 AND deleted_at IS NULL ORDER BY created_at ASC LIMIT $2"
1050            );
1051            sqlx::query(sqlx::AssertSqlSafe(sql.as_str()))
1052                .bind(agent_id)
1053                .bind(limit as i64)
1054                .fetch_all(&self.pool)
1055                .await
1056                .map_err(map_sqlx)?
1057        };
1058
1059        let mut results = Vec::with_capacity(rows.len());
1060        for r in &rows {
1061            results.push(row_to_memory(r).map_err(map_sqlx)?);
1062        }
1063        Ok(results)
1064    }
1065
1066    // -----------------------------------------------------------------------
1067    // Sync support
1068    // -----------------------------------------------------------------------
1069
1070    async fn list_memories_since(
1071        &self,
1072        updated_after: &str,
1073        limit: usize,
1074    ) -> Result<Vec<MemoryRecord>> {
1075        let sql = format!(
1076            "SELECT {MEMORY_COLUMNS} FROM memories WHERE updated_at > $1 ORDER BY updated_at ASC LIMIT $2"
1077        );
1078        let rows = sqlx::query(sqlx::AssertSqlSafe(sql.as_str()))
1079            .bind(updated_after)
1080            .bind(limit as i64)
1081            .fetch_all(&self.pool)
1082            .await
1083            .map_err(map_sqlx)?;
1084
1085        let mut results = Vec::with_capacity(rows.len());
1086        for r in &rows {
1087            results.push(row_to_memory(r).map_err(map_sqlx)?);
1088        }
1089        Ok(results)
1090    }
1091
1092    async fn upsert_memory(&self, record: &MemoryRecord) -> Result<()> {
1093        match self.update_memory(record).await {
1094            Ok(()) => Ok(()),
1095            Err(Error::NotFound(_)) => self.insert_memory(record).await,
1096            Err(e) => Err(e),
1097        }
1098    }
1099
1100    // -----------------------------------------------------------------------
1101    // Expired memory cleanup
1102    // -----------------------------------------------------------------------
1103
1104    async fn cleanup_expired(&self) -> Result<usize> {
1105        let now = chrono::Utc::now().to_rfc3339();
1106        let result = sqlx::query(
1107            "UPDATE memories SET deleted_at = $1 WHERE expires_at IS NOT NULL AND expires_at < $2 AND deleted_at IS NULL",
1108        )
1109        .bind(&now)
1110        .bind(&now)
1111        .execute(&self.pool)
1112        .await
1113        .map_err(map_sqlx)?;
1114
1115        Ok(result.rows_affected() as usize)
1116    }
1117
1118    // -----------------------------------------------------------------------
1119    // Delegations
1120    // -----------------------------------------------------------------------
1121
1122    async fn insert_delegation(&self, d: &Delegation) -> Result<()> {
1123        let scope_type = d.scope.to_string();
1124        let scope_value: serde_json::Value = match &d.scope {
1125            DelegationScope::AllMemories => serde_json::Value::Null,
1126            DelegationScope::ByTag(tags) => serde_json::json!(tags),
1127            DelegationScope::ByMemoryId(ids) => {
1128                serde_json::json!(ids.iter().map(|id| id.to_string()).collect::<Vec<_>>())
1129            }
1130        };
1131
1132        sqlx::query(
1133            r#"
1134INSERT INTO delegations (
1135    id, delegator_id, delegate_id, permission, scope_type, scope_value,
1136    max_depth, current_depth, parent_delegation_id,
1137    created_at, expires_at, revoked_at
1138) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
1139"#,
1140        )
1141        .bind(d.id)
1142        .bind(&d.delegator_id)
1143        .bind(&d.delegate_id)
1144        .bind(d.permission.to_string())
1145        .bind(&scope_type)
1146        .bind(&scope_value)
1147        .bind(d.max_depth as i32)
1148        .bind(d.current_depth as i32)
1149        .bind(d.parent_delegation_id)
1150        .bind(&d.created_at)
1151        .bind(&d.expires_at)
1152        .bind(&d.revoked_at)
1153        .execute(&self.pool)
1154        .await
1155        .map_err(map_sqlx)?;
1156        Ok(())
1157    }
1158
1159    async fn list_delegations_for(&self, delegate_id: &str) -> Result<Vec<Delegation>> {
1160        let now = chrono::Utc::now().to_rfc3339();
1161        let rows = sqlx::query(
1162            r#"
1163SELECT id, delegator_id, delegate_id, permission, scope_type, scope_value,
1164       max_depth, current_depth, parent_delegation_id,
1165       created_at, expires_at, revoked_at
1166FROM delegations
1167WHERE delegate_id = $1 AND revoked_at IS NULL AND (expires_at IS NULL OR expires_at > $2)
1168"#,
1169        )
1170        .bind(delegate_id)
1171        .bind(&now)
1172        .fetch_all(&self.pool)
1173        .await
1174        .map_err(map_sqlx)?;
1175
1176        let mut results = Vec::with_capacity(rows.len());
1177        for r in &rows {
1178            results.push(row_to_delegation(r).map_err(map_sqlx)?);
1179        }
1180        Ok(results)
1181    }
1182
1183    async fn revoke_delegation(&self, id: Uuid) -> Result<()> {
1184        let now = chrono::Utc::now().to_rfc3339();
1185        let result = sqlx::query(
1186            "UPDATE delegations SET revoked_at = $1 WHERE id = $2 AND revoked_at IS NULL",
1187        )
1188        .bind(&now)
1189        .bind(id)
1190        .execute(&self.pool)
1191        .await
1192        .map_err(map_sqlx)?;
1193
1194        if result.rows_affected() == 0 {
1195            return Err(Error::NotFound(format!(
1196                "delegation {id} not found or already revoked"
1197            )));
1198        }
1199        Ok(())
1200    }
1201
1202    async fn check_delegation(
1203        &self,
1204        delegate_id: &str,
1205        memory_id: Uuid,
1206        required: Permission,
1207    ) -> Result<bool> {
1208        let delegations = self.list_delegations_for(delegate_id).await?;
1209
1210        // Get the memory to inspect its tags for scope matching
1211        let memory = match self.get_memory(memory_id).await? {
1212            Some(m) => m,
1213            None => return Ok(false),
1214        };
1215
1216        for d in &delegations {
1217            if !d.permission.satisfies(required) {
1218                continue;
1219            }
1220            match &d.scope {
1221                DelegationScope::AllMemories => return Ok(true),
1222                DelegationScope::ByMemoryId(ids) => {
1223                    if ids.contains(&memory_id) {
1224                        return Ok(true);
1225                    }
1226                }
1227                DelegationScope::ByTag(tags) => {
1228                    if tags.iter().any(|t| memory.tags.contains(t)) {
1229                        return Ok(true);
1230                    }
1231                }
1232            }
1233        }
1234        Ok(false)
1235    }
1236
1237    // -----------------------------------------------------------------------
1238    // Agent Profiles
1239    // -----------------------------------------------------------------------
1240
1241    async fn insert_or_update_agent_profile(&self, profile: &AgentProfile) -> Result<()> {
1242        sqlx::query(
1243            r#"
1244INSERT INTO agent_profiles (agent_id, avg_importance, avg_content_length, total_memories, last_updated)
1245VALUES ($1, $2, $3, $4, $5)
1246ON CONFLICT (agent_id) DO UPDATE SET
1247    avg_importance = EXCLUDED.avg_importance,
1248    avg_content_length = EXCLUDED.avg_content_length,
1249    total_memories = EXCLUDED.total_memories,
1250    last_updated = EXCLUDED.last_updated
1251"#,
1252        )
1253        .bind(&profile.agent_id)
1254        .bind(profile.avg_importance)
1255        .bind(profile.avg_content_length)
1256        .bind(profile.total_memories as i64)
1257        .bind(&profile.last_updated)
1258        .execute(&self.pool)
1259        .await
1260        .map_err(map_sqlx)?;
1261        Ok(())
1262    }
1263
1264    async fn get_agent_profile(&self, agent_id: &str) -> Result<Option<AgentProfile>> {
1265        let row = sqlx::query(
1266            "SELECT agent_id, avg_importance, avg_content_length, total_memories, last_updated FROM agent_profiles WHERE agent_id = $1",
1267        )
1268        .bind(agent_id)
1269        .fetch_optional(&self.pool)
1270        .await
1271        .map_err(map_sqlx)?;
1272
1273        Ok(row.map(|r| AgentProfile {
1274            agent_id: r.get("agent_id"),
1275            avg_importance: r.get("avg_importance"),
1276            avg_content_length: r.get("avg_content_length"),
1277            total_memories: r.get::<i64, _>("total_memories") as u64,
1278            last_updated: r.get("last_updated"),
1279        }))
1280    }
1281
1282    // -----------------------------------------------------------------------
1283    // Embedding baselines (v0.3.3)
1284    // -----------------------------------------------------------------------
1285
1286    async fn insert_or_update_embedding_baseline(
1287        &self,
1288        baseline: &EmbeddingBaseline,
1289    ) -> Result<()> {
1290        let mu_json =
1291            serde_json::to_value(&baseline.mu).map_err(|e| Error::Storage(e.to_string()))?;
1292        let cov_json =
1293            serde_json::to_value(&baseline.cov_diag).map_err(|e| Error::Storage(e.to_string()))?;
1294        sqlx::query(
1295            r#"
1296INSERT INTO embedding_baseline (agent_id, mu, cov_diag, n, updated_at)
1297VALUES ($1, $2, $3, $4, $5)
1298ON CONFLICT (agent_id) DO UPDATE SET
1299    mu = EXCLUDED.mu,
1300    cov_diag = EXCLUDED.cov_diag,
1301    n = EXCLUDED.n,
1302    updated_at = EXCLUDED.updated_at
1303"#,
1304        )
1305        .bind(&baseline.agent_id)
1306        .bind(&mu_json)
1307        .bind(&cov_json)
1308        .bind(baseline.n as i64)
1309        .bind(&baseline.updated_at)
1310        .execute(&self.pool)
1311        .await
1312        .map_err(map_sqlx)?;
1313        Ok(())
1314    }
1315
1316    async fn get_embedding_baseline(&self, agent_id: &str) -> Result<Option<EmbeddingBaseline>> {
1317        let row = sqlx::query(
1318            "SELECT agent_id, mu, cov_diag, n, updated_at FROM embedding_baseline WHERE agent_id = $1",
1319        )
1320        .bind(agent_id)
1321        .fetch_optional(&self.pool)
1322        .await
1323        .map_err(map_sqlx)?;
1324
1325        match row {
1326            None => Ok(None),
1327            Some(r) => {
1328                let mu_val: serde_json::Value = r.get("mu");
1329                let cov_val: serde_json::Value = r.get("cov_diag");
1330                let mu: Vec<f32> =
1331                    serde_json::from_value(mu_val).map_err(|e| Error::Storage(e.to_string()))?;
1332                let cov_diag: Vec<f32> =
1333                    serde_json::from_value(cov_val).map_err(|e| Error::Storage(e.to_string()))?;
1334                Ok(Some(EmbeddingBaseline {
1335                    agent_id: r.get("agent_id"),
1336                    mu,
1337                    cov_diag,
1338                    n: r.get::<i64, _>("n") as u64,
1339                    updated_at: r.get("updated_at"),
1340                }))
1341            }
1342        }
1343    }
1344
1345    // -----------------------------------------------------------------------
1346    // Checkpoints
1347    // -----------------------------------------------------------------------
1348
1349    async fn insert_checkpoint(&self, cp: &Checkpoint) -> Result<()> {
1350        let memory_refs_strs: Vec<String> =
1351            cp.memory_refs.iter().map(|id| id.to_string()).collect();
1352        let refs_slice: &[String] = &memory_refs_strs;
1353
1354        sqlx::query(
1355            r#"
1356INSERT INTO checkpoints (
1357    id, thread_id, agent_id, parent_id, branch_name,
1358    state_snapshot, state_diff, memory_refs, event_cursor,
1359    label, created_at, metadata
1360) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11, $12)
1361"#,
1362        )
1363        .bind(cp.id)
1364        .bind(&cp.thread_id)
1365        .bind(&cp.agent_id)
1366        .bind(cp.parent_id)
1367        .bind(&cp.branch_name)
1368        .bind(&cp.state_snapshot)
1369        .bind(&cp.state_diff)
1370        .bind(refs_slice)
1371        .bind(cp.event_cursor)
1372        .bind(&cp.label)
1373        .bind(&cp.created_at)
1374        .bind(&cp.metadata)
1375        .execute(&self.pool)
1376        .await
1377        .map_err(map_sqlx)?;
1378        Ok(())
1379    }
1380
1381    async fn get_checkpoint(&self, id: Uuid) -> Result<Option<Checkpoint>> {
1382        let row = sqlx::query(
1383            r#"
1384SELECT id, thread_id, agent_id, parent_id, branch_name,
1385       state_snapshot, state_diff, memory_refs, event_cursor,
1386       label, created_at, metadata
1387FROM checkpoints WHERE id = $1
1388"#,
1389        )
1390        .bind(id)
1391        .fetch_optional(&self.pool)
1392        .await
1393        .map_err(map_sqlx)?;
1394
1395        match row {
1396            Some(r) => Ok(Some(row_to_checkpoint(&r).map_err(map_sqlx)?)),
1397            None => Ok(None),
1398        }
1399    }
1400
1401    async fn list_checkpoints(
1402        &self,
1403        thread_id: &str,
1404        branch: Option<&str>,
1405        limit: usize,
1406    ) -> Result<Vec<Checkpoint>> {
1407        let rows = if let Some(branch_name) = branch {
1408            sqlx::query(
1409                r#"
1410SELECT id, thread_id, agent_id, parent_id, branch_name,
1411       state_snapshot, state_diff, memory_refs, event_cursor,
1412       label, created_at, metadata
1413FROM checkpoints
1414WHERE thread_id = $1 AND branch_name = $2
1415ORDER BY created_at DESC
1416LIMIT $3
1417"#,
1418            )
1419            .bind(thread_id)
1420            .bind(branch_name)
1421            .bind(limit as i64)
1422            .fetch_all(&self.pool)
1423            .await
1424            .map_err(map_sqlx)?
1425        } else {
1426            sqlx::query(
1427                r#"
1428SELECT id, thread_id, agent_id, parent_id, branch_name,
1429       state_snapshot, state_diff, memory_refs, event_cursor,
1430       label, created_at, metadata
1431FROM checkpoints
1432WHERE thread_id = $1
1433ORDER BY created_at DESC
1434LIMIT $2
1435"#,
1436            )
1437            .bind(thread_id)
1438            .bind(limit as i64)
1439            .fetch_all(&self.pool)
1440            .await
1441            .map_err(map_sqlx)?
1442        };
1443
1444        let mut results = Vec::with_capacity(rows.len());
1445        for r in &rows {
1446            results.push(row_to_checkpoint(r).map_err(map_sqlx)?);
1447        }
1448        Ok(results)
1449    }
1450
1451    async fn get_latest_checkpoint(
1452        &self,
1453        thread_id: &str,
1454        branch: &str,
1455    ) -> Result<Option<Checkpoint>> {
1456        let row = sqlx::query(
1457            r#"
1458SELECT id, thread_id, agent_id, parent_id, branch_name,
1459       state_snapshot, state_diff, memory_refs, event_cursor,
1460       label, created_at, metadata
1461FROM checkpoints
1462WHERE thread_id = $1 AND branch_name = $2
1463ORDER BY created_at DESC
1464LIMIT 1
1465"#,
1466        )
1467        .bind(thread_id)
1468        .bind(branch)
1469        .fetch_optional(&self.pool)
1470        .await
1471        .map_err(map_sqlx)?;
1472
1473        match row {
1474            Some(r) => Ok(Some(row_to_checkpoint(&r).map_err(map_sqlx)?)),
1475            None => Ok(None),
1476        }
1477    }
1478}
mnemo_postgres/storage.rs

mnemo_postgres/
storage.rs
