Skip to main content

tact_memory/store/
local.rs

1//! Private schema-v1 SQLite storage.
2
3use super::{MemoryError, MemoryStore, current_time_ms};
4use crate::{
5    MemoryImportReport, MemoryKey, MemoryLimits, MemoryRecord, MemoryScan,
6    model::{StoredMemory, normalize_identity},
7    secrets::contains_likely_secret,
8    server::protocol::{self, ExportCursor, SyncReport},
9};
10use rusqlite::{
11    Connection, ErrorCode, OptionalExtension, Transaction, TransactionBehavior, params,
12};
13use std::{
14    collections::{HashMap, HashSet},
15    fs,
16    future::Future,
17    path::{Path, PathBuf},
18    sync::Arc,
19    time::Duration,
20};
21use thiserror::Error;
22
23const BUSY_TIMEOUT: Duration = Duration::from_secs(5);
24const DATABASE_PAGE_SIZE_BYTES: usize = 4 * 1024;
25const SCHEMA_VERSION: i64 = 1;
26const ALLOCATOR_ROW_ID: i64 = 1;
27const INSTALL_ALLOCATOR_TRIGGER: &str = "CREATE TRIGGER IF NOT EXISTS memory_id_allocator
28     AFTER INSERT ON memories
29     BEGIN
30        SELECT CASE
31            WHEN NEW.id < (SELECT next_id FROM memory_metadata WHERE id = 1)
32            THEN RAISE(ABORT, 'memory id was already allocated')
33        END;
34        UPDATE memory_metadata SET next_id = NEW.id + 1
35        WHERE id = 1 AND next_id <= NEW.id;
36     END;";
37
38#[derive(Debug, Error)]
39enum LocalStoreError {
40    #[error("could not prepare the memory directory")]
41    Directory(#[source] std::io::Error),
42    #[error("memory storage task stopped unexpectedly")]
43    Task(#[source] tokio::task::JoinError),
44}
45
46/// Concrete private SQLite memory store.
47#[derive(Clone, Debug)]
48pub struct LocalMemoryStore {
49    pub(crate) path: Arc<PathBuf>,
50    limits: MemoryLimits,
51}
52
53impl LocalMemoryStore {
54    /// Opens or creates a private local SQLite store at `path` on first use.
55    pub fn new(path: impl Into<PathBuf>) -> Self {
56        Self {
57            path: Arc::new(path.into()),
58            limits: MemoryLimits::PRODUCTION,
59        }
60    }
61
62    #[cfg(test)]
63    pub(crate) fn with_limits(path: impl Into<PathBuf>, limits: MemoryLimits) -> Self {
64        Self {
65            path: Arc::new(path.into()),
66            limits,
67        }
68    }
69
70    /// Searches visible memories and records scan telemetry at `now_ms`.
71    async fn scan(
72        &self,
73        query: &str,
74        limit: usize,
75        now_ms: i64,
76    ) -> Result<MemoryScan, MemoryError> {
77        let store = self.clone();
78        let query = query.to_owned();
79        let limit = limit.min(self.limits.scan_results);
80        run_local(move || store.scan_local(&query, limit, now_ms)).await
81    }
82
83    pub(crate) fn scan_local(
84        &self,
85        query: &str,
86        limit: usize,
87        now_ms: i64,
88    ) -> Result<MemoryScan, MemoryError> {
89        if query.len() > self.limits.query_bytes {
90            return Err(MemoryError::QueryTooLarge {
91                maximum_bytes: self.limits.query_bytes,
92            });
93        }
94
95        let mut connection = self.open()?;
96        let transaction = connection
97            .transaction_with_behavior(TransactionBehavior::Immediate)
98            .map_err(sqlite_error)?;
99        prune_expired(&transaction, now_ms)?;
100
101        let memories = load_all(&transaction)?
102            .into_iter()
103            .filter(|memory| !contains_likely_secret(&memory.content))
104            .map(MemoryRecord::from)
105            .collect::<Vec<_>>();
106        let limit = limit.min(self.limits.scan_results);
107        let scan = MemoryScan::rank(query, &memories, limit);
108
109        for candidate in &scan.candidates {
110            transaction
111                .execute(
112                    "UPDATE memories
113                     SET last_scanned_at_ms = ?1, scan_count = scan_count + 1
114                     WHERE id = ?2",
115                    params![now_ms, candidate.key.id],
116                )
117                .map_err(sqlite_error)?;
118        }
119        transaction.commit().map_err(sqlite_error)?;
120
121        Ok(scan)
122    }
123
124    pub(crate) fn read_local(
125        &self,
126        references: &[(i64, Option<u64>)],
127        now_ms: i64,
128    ) -> Result<Vec<MemoryRecord>, MemoryError> {
129        let mut connection = self.open()?;
130        let transaction = connection
131            .transaction_with_behavior(TransactionBehavior::Immediate)
132            .map_err(sqlite_error)?;
133        prune_expired(&transaction, now_ms)?;
134
135        let mut seen = HashSet::new();
136        let mut records = Vec::with_capacity(references.len());
137        for &(id, version) in references {
138            let memory = load_one(&transaction, id)?;
139            let Some(mut memory) = memory else {
140                continue;
141            };
142            if version.is_some_and(|version| version != memory.version)
143                || contains_likely_secret(&memory.content)
144                || !seen.insert(id)
145            {
146                continue;
147            }
148
149            transaction
150                .execute(
151                    "UPDATE memories
152                     SET last_used_at_ms = ?1, use_count = use_count + 1,
153                         probation_until_ms = NULL
154                     WHERE id = ?2",
155                    params![now_ms, id],
156                )
157                .map_err(sqlite_error)?;
158            memory.last_used_at_ms = Some(now_ms);
159            memory.use_count = memory.use_count.saturating_add(1);
160            memory.probation_until_ms = None;
161            records.push(memory.into());
162        }
163        transaction.commit().map_err(sqlite_error)?;
164        Ok(records)
165    }
166
167    /// Reads records selected by unversioned IDs and versioned keys.
168    ///
169    /// IDs refer to the active local store or configured remote namespace. Keys retain their
170    /// namespace and version semantics. Missing, stale, and duplicate records are omitted.
171    /// Successful reads record use telemetry at `now_ms`.
172    async fn read(
173        &self,
174        local_ids: &[i64],
175        keys: &[MemoryKey],
176        now_ms: i64,
177    ) -> Result<Vec<MemoryRecord>, MemoryError> {
178        let store = self.clone();
179        let owned_keys = keys.to_vec();
180        let read_ids = local_ids.to_vec();
181        run_local(move || {
182            let mut references = owned_keys
183                .iter()
184                .filter(|key| key.is_local())
185                .map(|key| (key.id, Some(key.version)))
186                .collect::<Vec<_>>();
187            references.extend(distinct_ids(&read_ids).into_iter().map(|id| (id, None)));
188            store.read_local(&references, now_ms)
189        })
190        .await
191    }
192
193    /// Inserts content or compare-and-swap replaces the record identified by `replacement`.
194    async fn put(
195        &self,
196        content: &str,
197        replacement: Option<MemoryKey>,
198        now_ms: i64,
199    ) -> Result<MemoryRecord, MemoryError> {
200        if replacement.as_ref().is_some_and(|key| !key.is_local()) {
201            return Err(MemoryError::RemoteReadOnly);
202        }
203        let store = self.clone();
204        let content = content.to_owned();
205        run_local(move || store.put_local(&content, replacement, now_ms)).await
206    }
207
208    pub(crate) fn put_local(
209        &self,
210        content: &str,
211        replacement: Option<MemoryKey>,
212        now_ms: i64,
213    ) -> Result<MemoryRecord, MemoryError> {
214        validate_content(content, &self.limits)?;
215        let normalized_identity = normalize_identity(content);
216        if normalized_identity.is_empty() {
217            return Err(MemoryError::EmptyContent);
218        }
219
220        let mut connection = self.open()?;
221        let transaction = connection
222            .transaction_with_behavior(TransactionBehavior::Immediate)
223            .map_err(sqlite_error)?;
224        prune_expired(&transaction, now_ms)?;
225
226        let result = match replacement {
227            Some(key) => self.replace(&transaction, content, &normalized_identity, key, now_ms),
228            None => self.insert(&transaction, content, &normalized_identity, now_ms),
229        }?;
230        transaction.commit().map_err(sqlite_error)?;
231        Ok(result.into())
232    }
233
234    /// Compare-and-swap deletes `key` from its owning backend.
235    async fn delete(&self, key: MemoryKey) -> Result<(), MemoryError> {
236        if !key.is_local() {
237            return Err(MemoryError::RemoteReadOnly);
238        }
239        let store = self.clone();
240        run_local(move || store.delete_local(key)).await
241    }
242
243    pub(crate) fn delete_local(&self, key: MemoryKey) -> Result<(), MemoryError> {
244        let mut connection = self.open()?;
245        let transaction = connection
246            .transaction_with_behavior(TransactionBehavior::Immediate)
247            .map_err(sqlite_error)?;
248        let current_version = transaction
249            .query_row(
250                "SELECT version FROM memories WHERE id = ?1",
251                [key.id],
252                |row| row.get::<_, i64>(0),
253            )
254            .optional()
255            .map_err(sqlite_error)?;
256        let Some(current_version) = current_version else {
257            transaction.commit().map_err(sqlite_error)?;
258            return Ok(());
259        };
260        if current_version as u64 != key.version {
261            return Err(MemoryError::Conflict);
262        }
263
264        transaction
265            .execute("DELETE FROM memories WHERE id = ?1", [key.id])
266            .map_err(sqlite_error)?;
267        transaction.commit().map_err(sqlite_error)?;
268        Ok(())
269    }
270
271    /// Lists all visible records after pruning probation at `now_ms`.
272    async fn list(&self, now_ms: i64) -> Result<Vec<MemoryRecord>, MemoryError> {
273        let store = self.clone();
274        run_local(move || store.list_local(now_ms)).await
275    }
276
277    /// Imports a validated remote snapshot into a local store as new probationary records.
278    pub async fn merge_remote_export(
279        &self,
280        memories: Vec<MemoryRecord>,
281    ) -> Result<MemoryImportReport, MemoryError> {
282        if memories.is_empty() {
283            return Ok(MemoryImportReport::default());
284        }
285        let store = self.clone();
286        let now_ms = current_time_ms();
287        run_local(move || store.merge_remote_export_local(memories, now_ms)).await
288    }
289
290    fn merge_remote_export_local(
291        &self,
292        mut memories: Vec<MemoryRecord>,
293        now_ms: i64,
294    ) -> Result<MemoryImportReport, MemoryError> {
295        memories.sort_by(|left, right| {
296            left.key
297                .namespace
298                .cmp(&right.key.namespace)
299                .then_with(|| left.key.id.cmp(&right.key.id))
300        });
301        for memory in &memories {
302            let namespace = memory
303                .key
304                .namespace
305                .as_deref()
306                .ok_or(MemoryError::Conflict)?;
307            if !protocol::is_valid_namespace(namespace)
308                || memory.key.id <= 0
309                || memory.key.version == 0
310            {
311                return Err(MemoryError::Conflict);
312            }
313            validate_content(&memory.content, &self.limits)?;
314        }
315
316        let mut connection = self.open()?;
317        let transaction = connection
318            .transaction_with_behavior(TransactionBehavior::Immediate)
319            .map_err(sqlite_error)?;
320        prune_expired(&transaction, now_ms)?;
321        let totals = totals(&transaction)?;
322        let mut identities = load_all(&transaction)?
323            .into_iter()
324            .map(|memory| normalize_identity(&memory.content))
325            .collect::<HashSet<_>>();
326        let mut accepted = Vec::new();
327        let mut skipped = 0;
328        for memory in memories {
329            let identity = normalize_identity(&memory.content);
330            if !identities.insert(identity.clone()) {
331                skipped += 1;
332                continue;
333            }
334            accepted.push((memory.content, identity));
335        }
336        let resulting_records = totals.records.saturating_add(accepted.len() as u64);
337        if resulting_records > self.limits.records as u64 {
338            return Err(MemoryError::RecordCapacity {
339                maximum: self.limits.records,
340            });
341        }
342        let imported_bytes = accepted
343            .iter()
344            .map(|(content, _)| content.len() as u64)
345            .sum::<u64>();
346        if totals.content_bytes.saturating_add(imported_bytes)
347            > self.limits.total_content_bytes as u64
348        {
349            return Err(MemoryError::ContentCapacity {
350                maximum_bytes: self.limits.total_content_bytes,
351            });
352        }
353        let probation_until_ms = now_ms.saturating_add(self.limits.probation_duration_ms);
354        for (content, identity) in &accepted {
355            let id = allocate_id(&transaction)?;
356            transaction
357                .execute(
358                    "INSERT INTO memories (
359                        id, content, normalized_identity, created_at_ms, updated_at_ms,
360                        last_scanned_at_ms, scan_count, last_used_at_ms, use_count,
361                        probation_until_ms, version
362                     ) VALUES (?1, ?2, ?3, ?4, ?4, NULL, 0, NULL, 0, ?5, 1)",
363                    params![id, content, identity, now_ms, probation_until_ms],
364                )
365                .map_err(sqlite_write_error)?;
366        }
367        transaction.commit().map_err(sqlite_write_error)?;
368        Ok(MemoryImportReport {
369            inserted: accepted.len(),
370            skipped,
371        })
372    }
373
374    pub(crate) fn list_local(&self, now_ms: i64) -> Result<Vec<MemoryRecord>, MemoryError> {
375        let mut connection = self.open()?;
376        let transaction = connection
377            .transaction_with_behavior(TransactionBehavior::Immediate)
378            .map_err(sqlite_error)?;
379        prune_expired(&transaction, now_ms)?;
380        let records = load_all(&transaction)?
381            .into_iter()
382            .filter(|memory| !contains_likely_secret(&memory.content))
383            .map(MemoryRecord::from)
384            .collect();
385        transaction.commit().map_err(sqlite_error)?;
386        Ok(records)
387    }
388
389    fn insert(
390        &self,
391        transaction: &Transaction<'_>,
392        content: &str,
393        normalized_identity: &str,
394        now_ms: i64,
395    ) -> Result<StoredMemory, MemoryError> {
396        if identity_exists(transaction, normalized_identity, None)? {
397            return Err(MemoryError::Duplicate);
398        }
399        let totals = totals(transaction)?;
400        if totals.records >= self.limits.records as u64 {
401            return Err(MemoryError::RecordCapacity {
402                maximum: self.limits.records,
403            });
404        }
405        self.check_content_capacity(totals.content_bytes, 0, content.len())?;
406
407        let probation_until_ms = now_ms.saturating_add(self.limits.probation_duration_ms);
408        let id = allocate_id(transaction)?;
409        transaction
410            .execute(
411                "INSERT INTO memories (
412                    id, content, normalized_identity, created_at_ms, updated_at_ms,
413                    last_scanned_at_ms, scan_count, last_used_at_ms, use_count,
414                    probation_until_ms, version
415                 ) VALUES (?1, ?2, ?3, ?4, ?4, NULL, 0, NULL, 0, ?5, 1)",
416                params![id, content, normalized_identity, now_ms, probation_until_ms],
417            )
418            .map_err(sqlite_write_error)?;
419        load_one(transaction, id)?.ok_or(MemoryError::NotFound)
420    }
421
422    fn replace(
423        &self,
424        transaction: &Transaction<'_>,
425        content: &str,
426        normalized_identity: &str,
427        key: MemoryKey,
428        now_ms: i64,
429    ) -> Result<StoredMemory, MemoryError> {
430        let current = transaction
431            .query_row(
432                "SELECT version, length(CAST(content AS BLOB)) FROM memories WHERE id = ?1",
433                [key.id],
434                |row| Ok((row.get::<_, i64>(0)?, row.get::<_, i64>(1)?)),
435            )
436            .optional()
437            .map_err(sqlite_error)?;
438        let Some((current_version, previous_content_bytes)) = current else {
439            return Err(MemoryError::NotFound);
440        };
441        if current_version as u64 != key.version {
442            return Err(MemoryError::Conflict);
443        }
444        if identity_exists(transaction, normalized_identity, Some(key.id))? {
445            return Err(MemoryError::Duplicate);
446        }
447
448        let totals = totals(transaction)?;
449        self.check_content_capacity(
450            totals.content_bytes,
451            previous_content_bytes as usize,
452            content.len(),
453        )?;
454        let next_version = current_version
455            .checked_add(1)
456            .ok_or(MemoryError::Conflict)?;
457        let probation_until_ms = now_ms.saturating_add(self.limits.probation_duration_ms);
458        transaction
459            .execute(
460                "UPDATE memories
461                 SET content = ?1, normalized_identity = ?2, updated_at_ms = ?3,
462                     last_scanned_at_ms = NULL, scan_count = 0,
463                     last_used_at_ms = NULL, use_count = 0,
464                     probation_until_ms = ?4, version = ?5
465                 WHERE id = ?6 AND version = ?7",
466                params![
467                    content,
468                    normalized_identity,
469                    now_ms,
470                    probation_until_ms,
471                    next_version,
472                    key.id,
473                    current_version,
474                ],
475            )
476            .map_err(sqlite_write_error)?;
477        load_one(transaction, key.id)?.ok_or(MemoryError::NotFound)
478    }
479
480    fn check_content_capacity(
481        &self,
482        current_bytes: u64,
483        replaced_bytes: usize,
484        new_bytes: usize,
485    ) -> Result<(), MemoryError> {
486        let resulting_bytes = current_bytes
487            .saturating_sub(replaced_bytes as u64)
488            .saturating_add(new_bytes as u64);
489        if resulting_bytes > self.limits.total_content_bytes as u64 {
490            return Err(MemoryError::ContentCapacity {
491                maximum_bytes: self.limits.total_content_bytes,
492            });
493        }
494        Ok(())
495    }
496
497    pub(crate) fn open(&self) -> Result<Connection, MemoryError> {
498        prepare_private_parent(&self.path)?;
499        let mut connection = Connection::open(self.path.as_path()).map_err(sqlite_error)?;
500        connection
501            .busy_timeout(BUSY_TIMEOUT)
502            .map_err(sqlite_error)?;
503        let schema_version = connection
504            .query_row("PRAGMA user_version", [], |row| row.get::<_, i64>(0))
505            .map_err(sqlite_error)?;
506        if !matches!(schema_version, 0 | SCHEMA_VERSION) {
507            return Err(MemoryError::UnsupportedSchemaVersion {
508                found: schema_version,
509                supported: SCHEMA_VERSION,
510            });
511        }
512        connection
513            .pragma_update(None, "journal_mode", "DELETE")
514            .map_err(sqlite_error)?;
515        connection
516            .pragma_update(None, "page_size", DATABASE_PAGE_SIZE_BYTES as i64)
517            .map_err(sqlite_error)?;
518        let page_size = connection
519            .query_row("PRAGMA page_size", [], |row| row.get::<_, i64>(0))
520            .map_err(sqlite_error)? as usize;
521        let maximum_pages = self.limits.database_bytes.div_ceil(page_size).max(1);
522        connection
523            .pragma_update(None, "max_page_count", maximum_pages as i64)
524            .map_err(sqlite_error)?;
525        // The allocator table is a backward-compatible schema-v1 extension. Older builds ignore
526        // it; current builds retain identity history even when every memory row is deleted.
527        let transaction = connection
528            .transaction_with_behavior(TransactionBehavior::Immediate)
529            .map_err(sqlite_error)?;
530        transaction
531            .execute_batch(
532                "CREATE TABLE IF NOT EXISTS memories (
533                    id INTEGER PRIMARY KEY,
534                    content TEXT NOT NULL,
535                    normalized_identity TEXT NOT NULL UNIQUE,
536                    created_at_ms INTEGER NOT NULL,
537                    updated_at_ms INTEGER NOT NULL,
538                    last_scanned_at_ms INTEGER,
539                    scan_count INTEGER NOT NULL DEFAULT 0 CHECK (scan_count >= 0),
540                    last_used_at_ms INTEGER,
541                    use_count INTEGER NOT NULL DEFAULT 0 CHECK (use_count >= 0),
542                    probation_until_ms INTEGER,
543                    version INTEGER NOT NULL DEFAULT 1 CHECK (version > 0)
544                 );
545                 CREATE TABLE IF NOT EXISTS memory_metadata (
546                    id INTEGER PRIMARY KEY CHECK (id = 1),
547                    next_id INTEGER NOT NULL CHECK (next_id > 0)
548                 );",
549            )
550            .map_err(sqlite_write_error)?;
551        let maximum_id = transaction
552            .query_row("SELECT COALESCE(MAX(id), 0) FROM memories", [], |row| {
553                row.get::<_, i64>(0)
554            })
555            .map_err(sqlite_error)?;
556        let next_id = maximum_id
557            .checked_add(1)
558            .ok_or(MemoryError::StorageCapacity)?;
559        transaction
560            .execute(
561                "INSERT INTO memory_metadata (id, next_id) VALUES (?1, ?2)
562                 ON CONFLICT(id) DO UPDATE SET next_id = excluded.next_id
563                 WHERE memory_metadata.next_id < excluded.next_id",
564                params![ALLOCATOR_ROW_ID, next_id],
565            )
566            .map_err(sqlite_write_error)?;
567        transaction
568            .execute_batch(INSTALL_ALLOCATOR_TRIGGER)
569            .map_err(sqlite_write_error)?;
570        if schema_version == 0 {
571            transaction
572                .pragma_update(None, "user_version", SCHEMA_VERSION)
573                .map_err(sqlite_write_error)?;
574        }
575        transaction.commit().map_err(sqlite_write_error)?;
576        Ok(connection)
577    }
578
579    async fn sync_local_snapshot(
580        &self,
581        memories: Vec<MemoryRecord>,
582        now_ms: i64,
583    ) -> Result<SyncReport, MemoryError> {
584        let store = self.clone();
585        run_local(move || store.sync_local_snapshot_blocking(&memories, now_ms)).await
586    }
587
588    fn sync_local_snapshot_blocking(
589        &self,
590        memories: &[MemoryRecord],
591        now_ms: i64,
592    ) -> Result<SyncReport, MemoryError> {
593        let mut identities = HashSet::new();
594        let mut ids = HashSet::new();
595        let content_bytes = memories.iter().try_fold(0usize, |total, memory| {
596            if !memory.key.is_local()
597                || memory.key.id <= 0
598                || memory.key.version == 0
599                || !ids.insert(memory.key.id)
600            {
601                return Err(MemoryError::Conflict);
602            }
603            validate_content(&memory.content, &self.limits)?;
604            if !identities.insert(normalize_identity(&memory.content)) {
605                return Err(MemoryError::Duplicate);
606            }
607            total
608                .checked_add(memory.content.len())
609                .ok_or(MemoryError::ContentCapacity {
610                    maximum_bytes: self.limits.total_content_bytes,
611                })
612        })?;
613        if memories.len() > self.limits.records {
614            return Err(MemoryError::RecordCapacity {
615                maximum: self.limits.records,
616            });
617        }
618        if content_bytes > self.limits.total_content_bytes {
619            return Err(MemoryError::ContentCapacity {
620                maximum_bytes: self.limits.total_content_bytes,
621            });
622        }
623
624        let mut connection = self.open()?;
625        let transaction = connection
626            .transaction_with_behavior(TransactionBehavior::Immediate)
627            .map_err(sqlite_error)?;
628        prune_expired(&transaction, now_ms)?;
629        let existing = load_all(&transaction)?;
630        let previous = existing
631            .iter()
632            .cloned()
633            .map(|memory| (memory.id, MemoryRecord::from(memory)))
634            .collect::<HashMap<_, _>>();
635        let incoming_ids = memories
636            .iter()
637            .map(|memory| memory.key.id)
638            .collect::<HashSet<_>>();
639        transaction
640            .execute_batch("DROP TRIGGER memory_id_allocator")
641            .map_err(sqlite_write_error)?;
642        transaction
643            .execute("DELETE FROM memories", [])
644            .map_err(sqlite_write_error)?;
645        let mut report = SyncReport {
646            deleted: existing
647                .iter()
648                .filter(|memory| !incoming_ids.contains(&memory.id))
649                .count(),
650            ..SyncReport::default()
651        };
652        for memory in memories {
653            observe_id(&transaction, memory.key.id)?;
654            match previous.get(&memory.key.id) {
655                Some(previous) if previous == memory => report.unchanged += 1,
656                Some(_) => report.replaced += 1,
657                None => report.inserted += 1,
658            }
659            transaction.execute(
660                "INSERT INTO memories (id, content, normalized_identity, created_at_ms, updated_at_ms, last_scanned_at_ms, scan_count, last_used_at_ms, use_count, probation_until_ms, version) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11)",
661                params![memory.key.id, memory.content, normalize_identity(&memory.content), memory.created_at_ms, memory.updated_at_ms, memory.last_scanned_at_ms, memory.scan_count as i64, memory.last_used_at_ms, memory.use_count as i64, memory.probation_until_ms, memory.key.version as i64],
662            ).map_err(sqlite_write_error)?;
663        }
664        transaction
665            .execute_batch(INSTALL_ALLOCATOR_TRIGGER)
666            .map_err(sqlite_write_error)?;
667        transaction.commit().map_err(sqlite_write_error)?;
668        Ok(report)
669    }
670
671    async fn export_local_page(
672        &self,
673        cursor: Option<ExportCursor>,
674        limit: usize,
675        now_ms: i64,
676    ) -> Result<(Vec<MemoryRecord>, Option<ExportCursor>), MemoryError> {
677        let mut records = self.list(now_ms).await?;
678        let after = cursor.map_or(0, |cursor| cursor.id);
679        records.retain(|record| record.key.id > after);
680        let limit = limit.clamp(1, protocol::MAX_EXPORT_PAGE_RECORDS);
681        let has_more = records.len() > limit;
682        records.truncate(limit);
683        let next = has_more.then(|| ExportCursor {
684            namespace: String::new(),
685            id: records.last().expect("non-empty limited page").key.id,
686        });
687        Ok((records, next))
688    }
689}
690
691impl MemoryStore for LocalMemoryStore {
692    fn scan(
693        &self,
694        query: &str,
695        limit: usize,
696    ) -> impl Future<Output = Result<MemoryScan, MemoryError>> + Send {
697        LocalMemoryStore::scan(self, query, limit, current_time_ms())
698    }
699    fn read(
700        &self,
701        ids: &[i64],
702        keys: &[MemoryKey],
703    ) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
704        LocalMemoryStore::read(self, ids, keys, current_time_ms())
705    }
706    fn list(&self) -> impl Future<Output = Result<Vec<MemoryRecord>, MemoryError>> + Send {
707        LocalMemoryStore::list(self, current_time_ms())
708    }
709    fn put(
710        &self,
711        content: &str,
712        replacement: Option<MemoryKey>,
713    ) -> impl Future<Output = Result<MemoryRecord, MemoryError>> + Send {
714        LocalMemoryStore::put(self, content, replacement, current_time_ms())
715    }
716    fn delete(&self, key: MemoryKey) -> impl Future<Output = Result<(), MemoryError>> + Send {
717        LocalMemoryStore::delete(self, key)
718    }
719    fn sync(
720        &self,
721        memories: &[MemoryRecord],
722    ) -> impl Future<Output = Result<SyncReport, MemoryError>> + Send {
723        let store = self.clone();
724        let memories = memories.to_vec();
725        let now_ms = current_time_ms();
726        async move { store.sync_local_snapshot(memories, now_ms).await }
727    }
728    fn export_page(
729        &self,
730        _namespaces: Option<&[String]>,
731        cursor: Option<&ExportCursor>,
732        limit: usize,
733    ) -> impl Future<Output = Result<(Vec<MemoryRecord>, Option<ExportCursor>), MemoryError>> + Send
734    {
735        let store = self.clone();
736        let cursor = cursor.cloned();
737        let now_ms = current_time_ms();
738        async move { store.export_local_page(cursor, limit, now_ms).await }
739    }
740}
741
742async fn run_local<T>(
743    operation: impl FnOnce() -> Result<T, MemoryError> + Send + 'static,
744) -> Result<T, MemoryError>
745where
746    T: Send + 'static,
747{
748    tokio::task::spawn_blocking(operation)
749        .await
750        .map_err(|source| MemoryError::backend(LocalStoreError::Task(source)))?
751}
752
753struct Totals {
754    records: u64,
755    content_bytes: u64,
756}
757
758fn totals(transaction: &Transaction<'_>) -> Result<Totals, MemoryError> {
759    transaction
760        .query_row(
761            "SELECT COUNT(*), COALESCE(SUM(length(CAST(content AS BLOB))), 0) FROM memories",
762            [],
763            |row| {
764                Ok(Totals {
765                    records: row.get::<_, i64>(0)? as u64,
766                    content_bytes: row.get::<_, i64>(1)? as u64,
767                })
768            },
769        )
770        .map_err(sqlite_error)
771}
772
773fn identity_exists(
774    transaction: &Transaction<'_>,
775    normalized_identity: &str,
776    excluded_id: Option<i64>,
777) -> Result<bool, MemoryError> {
778    transaction
779        .query_row(
780            "SELECT EXISTS(
781                SELECT 1 FROM memories
782                WHERE normalized_identity = ?1 AND (?2 IS NULL OR id != ?2)
783             )",
784            params![normalized_identity, excluded_id],
785            |row| row.get(0),
786        )
787        .map_err(sqlite_error)
788}
789
790fn prune_expired(transaction: &Transaction<'_>, now_ms: i64) -> Result<(), MemoryError> {
791    transaction
792        .execute(
793            "DELETE FROM memories
794             WHERE probation_until_ms IS NOT NULL
795               AND probation_until_ms <= ?1
796               AND use_count = 0",
797            [now_ms],
798        )
799        .map_err(sqlite_error)?;
800    Ok(())
801}
802
803fn allocate_id(transaction: &Transaction<'_>) -> Result<i64, MemoryError> {
804    let recorded = transaction
805        .query_row(
806            "SELECT next_id FROM memory_metadata WHERE id = ?1",
807            [ALLOCATOR_ROW_ID],
808            |row| row.get::<_, i64>(0),
809        )
810        .map_err(sqlite_error)?;
811    let maximum_id = transaction
812        .query_row("SELECT COALESCE(MAX(id), 0) FROM memories", [], |row| {
813            row.get::<_, i64>(0)
814        })
815        .map_err(sqlite_error)?;
816    let id = recorded.max(
817        maximum_id
818            .checked_add(1)
819            .ok_or(MemoryError::StorageCapacity)?,
820    );
821    id.checked_add(1).ok_or(MemoryError::StorageCapacity)?;
822    Ok(id)
823}
824
825fn observe_id(transaction: &Transaction<'_>, id: i64) -> Result<(), MemoryError> {
826    let next_id = id.checked_add(1).ok_or(MemoryError::StorageCapacity)?;
827    transaction
828        .execute(
829            "UPDATE memory_metadata SET next_id = MAX(next_id, ?1) WHERE id = ?2",
830            params![next_id, ALLOCATOR_ROW_ID],
831        )
832        .map_err(sqlite_write_error)?;
833    Ok(())
834}
835
836fn load_all(transaction: &Transaction<'_>) -> Result<Vec<StoredMemory>, MemoryError> {
837    let mut statement = transaction
838        .prepare(
839            "SELECT id, content, created_at_ms, updated_at_ms,
840                    last_scanned_at_ms, scan_count, last_used_at_ms, use_count,
841                    probation_until_ms, version
842             FROM memories
843             ORDER BY id",
844        )
845        .map_err(sqlite_error)?;
846    let rows = statement
847        .query_map([], row_to_memory)
848        .map_err(sqlite_error)?;
849    rows.collect::<Result<Vec<_>, _>>().map_err(sqlite_error)
850}
851
852fn load_one(transaction: &Transaction<'_>, id: i64) -> Result<Option<StoredMemory>, MemoryError> {
853    transaction
854        .query_row(
855            "SELECT id, content, created_at_ms, updated_at_ms,
856                    last_scanned_at_ms, scan_count, last_used_at_ms, use_count,
857                    probation_until_ms, version
858             FROM memories
859             WHERE id = ?1",
860            [id],
861            row_to_memory,
862        )
863        .optional()
864        .map_err(sqlite_error)
865}
866
867fn row_to_memory(row: &rusqlite::Row<'_>) -> rusqlite::Result<StoredMemory> {
868    Ok(StoredMemory {
869        namespace: None,
870        id: row.get(0)?,
871        content: row.get(1)?,
872        created_at_ms: row.get(2)?,
873        updated_at_ms: row.get(3)?,
874        last_scanned_at_ms: row.get(4)?,
875        scan_count: row.get::<_, i64>(5)? as u64,
876        last_used_at_ms: row.get(6)?,
877        use_count: row.get::<_, i64>(7)? as u64,
878        probation_until_ms: row.get(8)?,
879        version: row.get::<_, i64>(9)? as u64,
880    })
881}
882
883fn distinct_ids(ids: &[i64]) -> Vec<i64> {
884    let mut seen = HashSet::new();
885    ids.iter().copied().filter(|id| seen.insert(*id)).collect()
886}
887
888fn sqlite_error(source: rusqlite::Error) -> MemoryError {
889    let retryable = matches!(
890        &source,
891        rusqlite::Error::SqliteFailure(error, _)
892            if matches!(error.code, ErrorCode::DatabaseBusy | ErrorCode::DatabaseLocked)
893    );
894    if retryable {
895        MemoryError::unavailable(source)
896    } else {
897        MemoryError::backend(source)
898    }
899}
900
901fn sqlite_write_error(source: rusqlite::Error) -> MemoryError {
902    match &source {
903        rusqlite::Error::SqliteFailure(error, _) if error.code == ErrorCode::DiskFull => {
904            MemoryError::StorageCapacity
905        }
906        _ => sqlite_error(source),
907    }
908}
909
910pub(crate) fn prepare_private_parent(path: &Path) -> Result<(), MemoryError> {
911    let Some(parent) = path.parent() else {
912        return Ok(());
913    };
914    if parent.as_os_str().is_empty() {
915        return Ok(());
916    }
917    fs::create_dir_all(parent)
918        .map_err(|source| MemoryError::backend(LocalStoreError::Directory(source)))?;
919
920    #[cfg(unix)]
921    {
922        use std::os::unix::fs::PermissionsExt;
923
924        fs::set_permissions(parent, fs::Permissions::from_mode(0o700))
925            .map_err(|source| MemoryError::backend(LocalStoreError::Directory(source)))?;
926    }
927    Ok(())
928}
929
930pub(crate) fn validate_content(content: &str, limits: &MemoryLimits) -> Result<(), MemoryError> {
931    if content.trim().is_empty() {
932        return Err(MemoryError::EmptyContent);
933    }
934    if content.len() > limits.content_bytes {
935        return Err(MemoryError::ContentTooLarge {
936            maximum_bytes: limits.content_bytes,
937        });
938    }
939    if contains_likely_secret(content) {
940        return Err(MemoryError::SecretRejected);
941    }
942    Ok(())
943}
944
945#[cfg(test)]
946mod allocator_tests {
947    use super::*;
948
949    #[test]
950    fn allocation_reconciles_rows_inserted_by_a_legacy_writer() {
951        let directory = tempfile::tempdir().unwrap();
952        let store = LocalMemoryStore::new(directory.path().join("memory.sqlite3"));
953        let mut connection = store.open().unwrap();
954        let transaction = connection
955            .transaction_with_behavior(TransactionBehavior::Immediate)
956            .unwrap();
957        transaction
958            .execute(
959                "INSERT INTO memories (
960                    id, content, normalized_identity, created_at_ms, updated_at_ms, version
961                 ) VALUES (1, 'legacy', 'legacy', 1, 1, 1)",
962                [],
963            )
964            .unwrap();
965
966        assert_eq!(allocate_id(&transaction).unwrap(), 2);
967    }
968
969    #[test]
970    fn legacy_writers_cannot_reuse_retired_ids() {
971        let directory = tempfile::tempdir().unwrap();
972        let path = directory.path().join("memory.sqlite3");
973        let store = LocalMemoryStore::new(&path);
974        let memory = store.put_local("retired", None, 1).unwrap();
975        store.delete_local(memory.key).unwrap();
976
977        let legacy = Connection::open(path).unwrap();
978        assert!(
979            legacy
980                .execute(
981                    "INSERT INTO memories (
982                        content, normalized_identity, created_at_ms, updated_at_ms, version
983                     ) VALUES ('legacy', 'legacy', 2, 2, 1)",
984                    [],
985                )
986                .is_err()
987        );
988    }
989}