1use 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#[derive(Clone, Debug)]
48pub struct LocalMemoryStore {
49 pub(crate) path: Arc<PathBuf>,
50 limits: MemoryLimits,
51}
52
53impl LocalMemoryStore {
54 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 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 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 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 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 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 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 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}