mnemo/store.rs
1//! The Mnemo database — ties the storage engine and the memory model together.
2//!
3//! Layout of a populated file:
4//!
5//! ```text
6//! page 0 : Header (unencrypted)
7//! pages 1..W : write-ahead log region
8//! pages W.. : encrypted record runs + catalog runs + index +
9//! snapshot manifest (append-only)
10//! ```
11//!
12//! Durability is provided by a **write-ahead log** ([`crate::wal`]). A `flush`
13//! is one transaction: record data pages are written copy-on-write to fresh
14//! pages, then the new catalog, ANN index, and header are logged to the WAL
15//! and committed with a single fsync. A checkpoint then folds the WAL into the
16//! home pages. A crash before the commit leaves the previous state intact; a
17//! crash after it is repaired by replaying the WAL on open.
18//!
19//! Because pages are only ever appended, every past transaction's runs survive
20//! on disk. Each flush also appends an entry to a **snapshot manifest**, so
21//! [`Mnemo::restore_to`] can reinstate any past state. Stale pages — and the
22//! history along with them — are reclaimed by [`Mnemo::compact_file`].
23
24use std::collections::HashMap;
25use std::fs::OpenOptions;
26use std::io::{Read, Seek, SeekFrom, Write};
27use std::path::PathBuf;
28
29use serde::{Deserialize, Serialize};
30use ulid::Ulid;
31
32use crate::crypto::{self, KdfParams};
33use crate::error::{MnemoError, Result};
34use crate::format::{Header, FLAG_ENCRYPTED, PAGE_SIZE, PAYLOAD, VERSION, WRAPPED_DEK_LEN};
35use crate::index::{IndexConfig, IndexInfo, IvfPqIndex};
36use crate::memory::{self, Memory, MemoryType, Metric, Scope, ScoreWeights};
37use crate::pager::Pager;
38use crate::wal;
39
40/// Default initial size of the write-ahead log region, in pages (64 KiB).
41/// The region grows automatically when a transaction's control plane outgrows
42/// it, so this is a *floor* on the fresh-file footprint, not a cap on
43/// transaction size. Was 64 pages (512 KiB) in v0.1.0 — lowered because that
44/// reservation dominated small-file size (~62% of a 31-memory dogfood file).
45const DEFAULT_WAL_PAGES: u64 = 8;
46
47/// Hard floor on the WAL reservation. A single header-page flush already
48/// needs one full DATA frame plus a COMMIT frame; two pages are the bare
49/// minimum that lets the first transaction commit without immediate growth.
50const MIN_WAL_PAGES: u64 = 2;
51
52/// Default cap on retained snapshot manifest entries. Each `flush` appends
53/// one entry; without a cap, long-running processes accumulate entries
54/// forever, growing per-flush manifest-serialize cost as O(total flushes)
55/// and the manifest run size on disk in lockstep. 256 keeps roughly a
56/// week of hourly snapshots while staying tiny (256 × ~82 bytes ≈ 21 KiB
57/// of manifest). Set `MnemoConfig::max_snapshots = 0` to disable the cap.
58pub const DEFAULT_MAX_SNAPSHOTS: usize = 256;
59
60/// Upper bound for the WAL scan when recovering a torn-header file. The
61/// real WAL region size lives in the header, but a torn header is exactly
62/// the case where we can't read it — so we scan a generous fixed window
63/// from the conventional WAL start, relying on `wal::recover` stopping at
64/// the first non-magic frame. 64 covers the historical pre-0.2 default
65/// of a 64-page reservation; users who configure larger WALs and then
66/// suffer a torn header at the same time will need offline recovery.
67const TORN_HEADER_SCAN_PAGES: u64 = 64;
68
69/// Configuration for creating a new database.
70#[derive(Clone, Copy, Debug)]
71pub struct MnemoConfig {
72 /// Embedding dimensionality. Every stored vector must match this.
73 pub dimensions: usize,
74 /// Key-derivation parameters.
75 pub kdf: KdfParams,
76 /// Initial size of the WAL region in 8 KiB pages. Default is
77 /// [`DEFAULT_WAL_PAGES`] (8 pages = 64 KiB), which keeps fresh files
78 /// small; the region auto-grows when a transaction's control plane needs
79 /// more space, so raising this is only a hint for the expected steady-state
80 /// transaction size. Write-heavy databases that routinely flush large
81 /// catalogs or ANN indexes can set this higher (e.g. 64 = 512 KiB) to
82 /// avoid early grow events; tiny embedded uses can leave it at the
83 /// default. Clamped to a minimum of [`MIN_WAL_PAGES`] (2 pages).
84 pub wal_pages_initial: u64,
85 /// Maximum number of snapshot manifest entries to retain. Each `flush`
86 /// appends one entry; once the cap is hit, the *oldest* entry is dropped
87 /// from the in-memory manifest before the new one is added, so the
88 /// retained set is always the most-recent N. Defaults to
89 /// [`DEFAULT_MAX_SNAPSHOTS`] (256). Set to `0` to disable the cap and
90 /// retain every snapshot forever (the pre-v0.3 behavior).
91 ///
92 /// Pages referenced *only* by pruned snapshots stay on disk until the
93 /// next [`Mnemo::compact_file`], which reclaims them. Point-in-time
94 /// recovery via [`Mnemo::restore_to`] into a pruned `txn_id` returns
95 /// [`MnemoError::NotFound`].
96 ///
97 /// Apply to an already-open database with [`Mnemo::set_max_snapshots`].
98 pub max_snapshots: usize,
99}
100
101impl Default for MnemoConfig {
102 fn default() -> Self {
103 Self {
104 dimensions: 768,
105 kdf: KdfParams::secure(),
106 wal_pages_initial: DEFAULT_WAL_PAGES,
107 max_snapshots: DEFAULT_MAX_SNAPSHOTS,
108 }
109 }
110}
111
112/// Catalog entry: maps a memory ID to its page run on disk.
113///
114/// **Schema (v5+):** carries `accessed_at` and `access_count` so `recall`
115/// can update access stats by mutating the catalog entry rather than
116/// rewriting the entire record (Phase 2.1 of the improvement plan). The
117/// values in a memory's serialized record body become a stale snapshot
118/// after the first recall; [`Mnemo::read_memory`] re-populates them from
119/// the catalog so consumers always see the live values.
120#[derive(Serialize, Deserialize, Clone, Debug)]
121struct CatalogEntry {
122 /// ULID as a raw `u128`.
123 id: u128,
124 start_page: u64,
125 page_count: u32,
126 /// Exact serialized byte length of the record.
127 len: u32,
128 deleted: bool,
129 /// Unix seconds of the most recent recall hit (or write, for fresh entries).
130 accessed_at: i64,
131 /// How many times `recall` has surfaced this memory.
132 access_count: u32,
133}
134
135/// Frozen v4 catalog entry shape, used only by the v4→v5 migration path in
136/// [`Mnemo::open`]. Its layout matches the catalog encoding written by
137/// pre-v5 builds (5 positional fields). Do not change this struct.
138#[derive(Serialize, Deserialize, Clone, Debug)]
139struct CatalogEntryV4 {
140 id: u128,
141 start_page: u64,
142 page_count: u32,
143 len: u32,
144 deleted: bool,
145}
146
147/// Returned by [`Mnemo::prepare_for_flush`] and consumed by [`Mnemo::flush`].
148/// Carries the already-serialized control-plane bytes through the prelude
149/// so we don't re-serialize after the lease has been persisted.
150struct FlushPrelude {
151 cat_bytes: Option<Vec<u8>>,
152 idx_bytes: Option<Vec<u8>>,
153}
154
155/// One entry in the append-only snapshot manifest. Because record, catalog,
156/// and index pages are only ever *appended*, the runs a past flush wrote are
157/// still on disk; a `Snapshot` is the set of pointers needed to reconstruct
158/// the database exactly as that flush left it.
159#[derive(Serialize, Deserialize, Clone, Debug)]
160struct Snapshot {
161 txn_id: u64,
162 created_at: i64,
163 catalog_start: u64,
164 catalog_pages: u64,
165 catalog_len: u64,
166 index_start: u64,
167 index_pages: u64,
168 index_len: u64,
169 memory_count: u64,
170}
171
172/// A restorable point in a database's history — see [`Mnemo::snapshots`].
173#[derive(Clone, Copy, Debug)]
174pub struct SnapshotInfo {
175 /// Id of the transaction that produced this snapshot (monotonic from 1).
176 pub txn_id: u64,
177 /// When the snapshot was committed (unix seconds).
178 pub created_at: i64,
179 /// Live memory count captured in the snapshot.
180 pub memory_count: u64,
181}
182
183/// A retrieval request for [`Mnemo::recall`].
184#[derive(Clone, Debug)]
185pub struct RecallRequest {
186 /// Query embedding.
187 pub query: Vec<f32>,
188 /// Maximum results to return.
189 pub top_k: usize,
190 /// Restrict to these memory types (`None` = all types).
191 pub memory_types: Option<Vec<MemoryType>>,
192 /// Restrict to a single agent's view: its own memories plus shared ones.
193 pub agent_id: Option<String>,
194 /// Similarity metric.
195 pub metric: Metric,
196 /// Multi-signal score weights.
197 pub weights: ScoreWeights,
198 /// Index override: partitions to probe (`None` = index default). Ignored
199 /// when no ANN index is present.
200 pub n_probe: Option<usize>,
201 /// Index override: candidates to rerank (`None` = index default). Ignored
202 /// when no ANN index is present.
203 pub n_rerank: Option<usize>,
204 /// Whether to update each result's `accessed_at` and `access_count` in
205 /// the catalog. Defaults to `true`, matching pre-v5 behavior.
206 ///
207 /// Set to `false` for **fully read-only recall** — useful for batch
208 /// scoring, dry-run scoring, or recalls run by introspection tooling
209 /// (e.g. `mnemo about` consumers) where you don't want the score's own
210 /// observation to perturb the database. With `false`, recall does not
211 /// dirty the catalog and the next `flush` is a no-op.
212 pub track_access: bool,
213}
214
215impl RecallRequest {
216 /// A request with sensible defaults for a query embedding.
217 pub fn new(query: Vec<f32>) -> Self {
218 Self {
219 query,
220 top_k: 10,
221 memory_types: None,
222 agent_id: None,
223 metric: Metric::Cosine,
224 weights: ScoreWeights::default(),
225 n_probe: None,
226 n_rerank: None,
227 track_access: true,
228 }
229 }
230 /// Set the result cap.
231 pub fn top_k(mut self, k: usize) -> Self {
232 self.top_k = k;
233 self
234 }
235 /// Restrict to specific memory types.
236 pub fn types(mut self, t: Vec<MemoryType>) -> Self {
237 self.memory_types = Some(t);
238 self
239 }
240 /// Restrict to one agent's view.
241 pub fn agent(mut self, agent_id: impl Into<String>) -> Self {
242 self.agent_id = Some(agent_id.into());
243 self
244 }
245 /// Set the similarity metric (default: cosine).
246 pub fn metric(mut self, metric: Metric) -> Self {
247 self.metric = metric;
248 self
249 }
250 /// Replace the multi-signal score weights.
251 pub fn weights(mut self, weights: ScoreWeights) -> Self {
252 self.weights = weights;
253 self
254 }
255 /// Override the number of IVF partitions probed (accuracy/speed dial).
256 pub fn n_probe(mut self, n: usize) -> Self {
257 self.n_probe = Some(n);
258 self
259 }
260 /// Override the number of candidates reranked exactly (accuracy dial).
261 pub fn n_rerank(mut self, n: usize) -> Self {
262 self.n_rerank = Some(n);
263 self
264 }
265 /// Toggle whether recall updates `accessed_at` / `access_count` on
266 /// the returned memories' catalog entries (default `true`). Pass
267 /// `false` for a fully read-only recall — see the field doc.
268 pub fn track_access(mut self, track: bool) -> Self {
269 self.track_access = track;
270 self
271 }
272}
273
274/// One scored result from [`Mnemo::recall`].
275#[derive(Clone, Debug)]
276pub struct RecallResult {
277 /// The retrieved memory.
278 pub memory: Memory,
279 /// Combined multi-signal score.
280 pub score: f32,
281 /// The bare similarity component (before other signals).
282 pub similarity: f32,
283}
284
285/// Summary statistics for a database.
286#[derive(Clone, Debug)]
287pub struct Stats {
288 /// Live (non-deleted) memory count.
289 pub memories: usize,
290 /// Tombstoned entries awaiting compaction.
291 pub deleted: usize,
292 /// Embedding dimensionality.
293 pub dimensions: usize,
294 /// File size in bytes.
295 pub file_bytes: u64,
296 /// Distinct agent IDs present.
297 pub agents: Vec<String>,
298 /// Whether pages are encrypted (always true in v1).
299 pub encrypted: bool,
300 /// Creation time (unix seconds).
301 pub created_at: i64,
302 /// ANN index shape, if an index has been built.
303 pub index: Option<IndexInfo>,
304 /// Current size of the write-ahead log region, in 8 KiB pages.
305 pub wal_pages: u64,
306}
307
308/// Result of a [`Mnemo::compact_file`] run.
309#[derive(Clone, Copy, Debug)]
310pub struct CompactReport {
311 /// Live memories before compaction.
312 pub before: usize,
313 /// Live memories after compaction (expired ones dropped).
314 pub after: usize,
315}
316
317/// An encrypted, single-file agent memory database.
318pub struct Mnemo {
319 pager: Pager,
320 header: Header,
321 catalog: Vec<CatalogEntry>,
322 index: HashMap<u128, usize>,
323 #[allow(dead_code)]
324 path: PathBuf,
325 dimensions: usize,
326 kdf: KdfParams,
327 /// Set whenever the catalog changes; drives whether `flush` rewrites it.
328 dirty_catalog: bool,
329 /// Optional IVF+PQ approximate-nearest-neighbour index.
330 ann: Option<IvfPqIndex>,
331 /// Set whenever the ANN index changes; drives whether `flush` rewrites it.
332 dirty_index: bool,
333 /// Append-only manifest of committed snapshots, oldest first. Capped
334 /// at [`Mnemo::max_snapshots`] entries on every flush — the oldest
335 /// entries are pruned first.
336 manifest: Vec<Snapshot>,
337 /// Maximum manifest entries to retain across flushes. `0` disables
338 /// the cap (retain forever). Defaults to [`DEFAULT_MAX_SNAPSHOTS`].
339 max_snapshots: usize,
340}
341
342/// Read a run of consecutive encrypted pages and concatenate their plaintext.
343fn read_run_bytes(pager: &mut Pager, start: u64, pages: u64, len: u64) -> Result<Vec<u8>> {
344 let mut buf = Vec::with_capacity(len as usize);
345 for i in 0..pages {
346 buf.extend_from_slice(&pager.read_page(start + i)?);
347 }
348 buf.truncate(len as usize);
349 Ok(buf)
350}
351
352/// Load a v5+ catalog run (an empty `pages` yields an empty catalog).
353fn load_catalog(pager: &mut Pager, start: u64, pages: u64, len: u64) -> Result<Vec<CatalogEntry>> {
354 if pages == 0 {
355 return Ok(Vec::new());
356 }
357 let buf = read_run_bytes(pager, start, pages, len)?;
358 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
359}
360
361/// Load a **v4** catalog run using the frozen [`CatalogEntryV4`] shape.
362/// Used only by the v4→v5 migration path in [`Mnemo::open`].
363fn load_catalog_v4(
364 pager: &mut Pager,
365 start: u64,
366 pages: u64,
367 len: u64,
368) -> Result<Vec<CatalogEntryV4>> {
369 if pages == 0 {
370 return Ok(Vec::new());
371 }
372 let buf = read_run_bytes(pager, start, pages, len)?;
373 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
374}
375
376/// Walk a v4 catalog and synthesize the v5 shape. For each live entry,
377/// reads the memory's serialized body and copies its `accessed_at` and
378/// `access_count` into the catalog row (the v4 record body still carries
379/// them — they're moving *from* the body *to* the catalog). Deleted
380/// entries skip the read and zero the fields.
381fn migrate_v4_catalog(
382 pager: &mut Pager,
383 v4: Vec<CatalogEntryV4>,
384) -> Result<Vec<CatalogEntry>> {
385 let mut out = Vec::with_capacity(v4.len());
386 for e in &v4 {
387 let (accessed_at, access_count) = if e.deleted {
388 (0i64, 0u32)
389 } else {
390 let mut buf = Vec::with_capacity(e.len as usize);
391 for i in 0..e.page_count as u64 {
392 buf.extend_from_slice(&pager.read_page(e.start_page + i)?);
393 }
394 buf.truncate(e.len as usize);
395 let mem: Memory = rmp_serde::from_slice(&buf)
396 .map_err(|err| MnemoError::Serialize(err.to_string()))?;
397 (mem.accessed_at, mem.access_count)
398 };
399 out.push(CatalogEntry {
400 id: e.id,
401 start_page: e.start_page,
402 page_count: e.page_count,
403 len: e.len,
404 deleted: e.deleted,
405 accessed_at,
406 access_count,
407 });
408 }
409 Ok(out)
410}
411
412/// Load the snapshot manifest (an empty `pages` yields an empty manifest).
413fn load_manifest(pager: &mut Pager, start: u64, pages: u64, len: u64) -> Result<Vec<Snapshot>> {
414 if pages == 0 {
415 return Ok(Vec::new());
416 }
417 let buf = read_run_bytes(pager, start, pages, len)?;
418 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
419}
420
421/// Load an ANN index run, validating its dimensionality.
422fn load_index(
423 pager: &mut Pager,
424 start: u64,
425 pages: u64,
426 len: u64,
427 dims: usize,
428) -> Result<Option<IvfPqIndex>> {
429 if pages == 0 {
430 return Ok(None);
431 }
432 let buf = read_run_bytes(pager, start, pages, len)?;
433 let mut idx: IvfPqIndex =
434 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))?;
435 if idx.dims() != dims {
436 return Err(MnemoError::Invalid(
437 "ANN index dimensionality does not match the database".into(),
438 ));
439 }
440 idx.rebuild_assignment();
441 Ok(Some(idx))
442}
443
444/// True if a memory's metadata marks it as the canonical onboarding manifest
445/// (`metadata.topic == "manifest"`). Used by [`Mnemo::about`] to hoist the
446/// manifest to the top of the briefing regardless of importance ordering.
447fn is_manifest(m: &Memory) -> bool {
448 m.metadata
449 .get("topic")
450 .and_then(|v| v.as_str())
451 .map(|s| s.eq_ignore_ascii_case("manifest"))
452 .unwrap_or(false)
453}
454
455/// Build the ULID → catalog-slot lookup map.
456fn build_id_index(catalog: &[CatalogEntry]) -> HashMap<u128, usize> {
457 let mut m = HashMap::with_capacity(catalog.len());
458 for (i, e) in catalog.iter().enumerate() {
459 m.insert(e.id, i);
460 }
461 m
462}
463
464impl Mnemo {
465 /// Create a brand-new encrypted database at `path`.
466 pub fn create(path: impl Into<PathBuf>, passphrase: &str, config: MnemoConfig) -> Result<Mnemo> {
467 let path: PathBuf = path.into();
468 if config.dimensions == 0 {
469 return Err(MnemoError::Invalid("dimensions must be > 0".into()));
470 }
471 let file = OpenOptions::new()
472 .read(true)
473 .write(true)
474 .create(true)
475 .truncate(true)
476 .open(&path)?;
477
478 let salt = crypto::random_salt();
479 let dek = crypto::random_dek();
480 let kek = crypto::derive_kek(passphrase.as_bytes(), &salt, config.kdf)?;
481 let dek_nonce = crypto::random_nonce();
482 let wrapped = crypto::wrap_dek(&kek, &dek_nonce, &dek)?;
483 if wrapped.len() != WRAPPED_DEK_LEN {
484 return Err(MnemoError::Crypto("unexpected wrapped DEK length".into()));
485 }
486 let mut wrapped_dek = [0u8; WRAPPED_DEK_LEN];
487 wrapped_dek.copy_from_slice(&wrapped);
488
489 let mut header = Header {
490 version: VERSION,
491 page_size: PAGE_SIZE as u32,
492 flags: FLAG_ENCRYPTED,
493 dimensions: config.dimensions as u32,
494 created_at: memory::now_secs(),
495 write_counter: 0,
496 // Page 0 is the header; pages 1..=wal_pages are the WAL;
497 // record/catalog/index pages start after it. Initial WAL size is
498 // chosen per-config (default 8 pages = 64 KiB) and clamped to a
499 // sane floor so the first transaction can always commit.
500 next_page: 1 + config.wal_pages_initial.max(MIN_WAL_PAGES),
501 catalog_start: 0,
502 catalog_pages: 0,
503 catalog_len: 0,
504 vector_count: 0,
505 m_cost: config.kdf.m_cost,
506 t_cost: config.kdf.t_cost,
507 p_cost: config.kdf.p_cost,
508 salt,
509 dek_nonce,
510 wrapped_dek,
511 index_start: 0,
512 index_pages: 0,
513 index_len: 0,
514 wal_start: 1,
515 wal_pages: config.wal_pages_initial.max(MIN_WAL_PAGES),
516 wal_seq: 0,
517 manifest_start: 0,
518 manifest_pages: 0,
519 manifest_len: 0,
520 // Regenerated per write by `apply_seal`; `None` until we seal.
521 seal_nonce: None,
522 seal_tag: None,
523 };
524
525 let mut pager = Pager::new(file, dek, 0, VERSION);
526 let mut page = header.to_page();
527 header.apply_seal(&mut page, pager.dek())?;
528 pager.write_raw(0, &page)?;
529 pager.sync()?;
530
531 Ok(Mnemo {
532 pager,
533 header,
534 catalog: Vec::new(),
535 index: HashMap::new(),
536 path,
537 dimensions: config.dimensions,
538 kdf: config.kdf,
539 dirty_catalog: false,
540 ann: None,
541 dirty_index: false,
542 manifest: Vec::new(),
543 max_snapshots: config.max_snapshots,
544 })
545 }
546
547 /// Open an existing database. A wrong passphrase fails cleanly with
548 /// [`MnemoError::WrongPassphrase`].
549 pub fn open(path: impl Into<PathBuf>, passphrase: &str) -> Result<Mnemo> {
550 let path: PathBuf = path.into();
551 let mut file = OpenOptions::new().read(true).write(true).open(&path)?;
552
553 let mut hbuf = [0u8; PAGE_SIZE];
554 file.seek(SeekFrom::Start(0))?;
555 file.read_exact(&mut hbuf)?;
556 let mut header = match Header::from_page(&hbuf) {
557 Ok(h) => h,
558 Err(MnemoError::HeaderChecksum) => {
559 // Page 0 is torn — most likely a crash during a checkpoint's
560 // header write. Try to heal from the WAL at its default site
561 // (a never-grown WAL never moves from page 1). We don't know
562 // the WAL size (the header that holds it is what's torn), so
563 // scan a generous fixed window — `wal::recover` stops at the
564 // first non-magic frame, so reading extra is safe.
565 let healed = wal::recover(&mut file, 1, TORN_HEADER_SCAN_PAGES, 0)?;
566 let frames = healed.ok_or(MnemoError::HeaderChecksum)?;
567 for (page_no, bytes) in &frames {
568 if bytes.len() != PAGE_SIZE {
569 return Err(MnemoError::Invalid("WAL frame is not page-sized".into()));
570 }
571 file.seek(SeekFrom::Start(page_no * PAGE_SIZE as u64))?;
572 file.write_all(bytes)?;
573 }
574 file.sync_all()?;
575 file.seek(SeekFrom::Start(0))?;
576 file.read_exact(&mut hbuf)?;
577 Header::from_page(&hbuf)?
578 }
579 Err(e) => return Err(e),
580 };
581
582 // Crash recovery: replay a committed-but-uncheckpointed transaction.
583 // Each frame is a finished page image; the header frame, replayed to
584 // page 0, supersedes the header just read.
585 if let Some(frames) =
586 wal::recover(&mut file, header.wal_start, header.wal_pages, header.wal_seq)?
587 {
588 for (page_no, bytes) in &frames {
589 if bytes.len() != PAGE_SIZE {
590 return Err(MnemoError::Invalid("WAL frame is not page-sized".into()));
591 }
592 file.seek(SeekFrom::Start(page_no * PAGE_SIZE as u64))?;
593 file.write_all(bytes)?;
594 }
595 file.sync_all()?;
596 // Re-read the now-current header from the replayed page 0.
597 file.seek(SeekFrom::Start(0))?;
598 file.read_exact(&mut hbuf)?;
599 header = Header::from_page(&hbuf)?;
600 }
601
602 let kdf = KdfParams {
603 m_cost: header.m_cost,
604 t_cost: header.t_cost,
605 p_cost: header.p_cost,
606 };
607 let kek = crypto::derive_kek(passphrase.as_bytes(), &header.salt, kdf)?;
608 let dek = crypto::unwrap_dek(&kek, &header.dek_nonce, &header.wrapped_dek)?;
609
610 // Validate the v7 header seal now that we have the DEK. Returns
611 // `MnemoError::HeaderTampered` if any of the sealed mutable fields
612 // were rewritten after the last legitimate flush. No-op for v6 and
613 // older — pre-v7 files don't carry a seal.
614 header.validate_seal(&dek)?;
615
616 // Build the pager at the file's *current* version so initial reads
617 // use the matching AAD scheme — `Pager::page_aad` returns empty AAD
618 // for v4/v5 (no page binding) and `page_no.to_le_bytes()` for v6+.
619 let on_disk_version = header.version;
620 let mut pager = Pager::new(file, dek, header.write_counter, on_disk_version);
621 let dimensions = header.dimensions as usize;
622
623 // === v4→v5 catalog migration =====================================
624 // v6 reads directly; v5 reads directly; v4 replays the catalog into
625 // the v5 shape by reading each memory's body to populate the new
626 // `accessed_at` / `access_count` fields. All page reads here still
627 // use the v4/v5 AAD scheme (none).
628 let (catalog, migrated_from_v4) = if on_disk_version >= 5 {
629 (
630 load_catalog(
631 &mut pager,
632 header.catalog_start,
633 header.catalog_pages,
634 header.catalog_len,
635 )?,
636 false,
637 )
638 } else if on_disk_version == 4 {
639 let v4 = load_catalog_v4(
640 &mut pager,
641 header.catalog_start,
642 header.catalog_pages,
643 header.catalog_len,
644 )?;
645 (migrate_v4_catalog(&mut pager, v4)?, true)
646 } else {
647 return Err(MnemoError::UnsupportedVersion(on_disk_version));
648 };
649 let index = build_id_index(&catalog);
650 let ann = load_index(
651 &mut pager,
652 header.index_start,
653 header.index_pages,
654 header.index_len,
655 dimensions,
656 )?;
657 let mut manifest = load_manifest(
658 &mut pager,
659 header.manifest_start,
660 header.manifest_pages,
661 header.manifest_len,
662 )?;
663
664 // === v5→v6 page-encryption migration =============================
665 // Pre-v6 data pages were encrypted under no AAD; v6 binds the home
666 // `page_no` as AAD so a page-transplant attack fails authentication
667 // on read. Re-stage every live record page through `write_page` so
668 // the next flush re-encrypts it under v6 AAD at the same home slot.
669 // Tombstoned entries skip — they're reclaimed by `compact_file`.
670 // Old catalog / ANN / manifest pages stay as orphan ciphertext;
671 // the next flush writes the new runs at fresh slots.
672 let migrated_pages_to_v6 = on_disk_version < 6;
673 if migrated_pages_to_v6 {
674 for e in &catalog {
675 if e.deleted {
676 continue;
677 }
678 for i in 0..e.page_count as u64 {
679 let payload = pager.read_page(e.start_page + i)?;
680 pager.write_page(e.start_page + i, &payload)?;
681 }
682 }
683 pager.set_version(6);
684 }
685
686 // === v6→v7 header-seal migration =================================
687 // v7 adds an AES-GCM seal at the tail of the header page that
688 // authenticates every mutable header field. Pre-v7 files have no
689 // seal — we just bump the version in memory and mark `dirty_catalog`
690 // so the next flush writes a sealed v7 header. Pages are untouched
691 // (page crypto is unchanged from v6).
692 let migrated_header_to_v7 = on_disk_version < 7;
693 if migrated_header_to_v7 {
694 pager.set_version(VERSION);
695 }
696
697 let migrated = migrated_from_v4 || migrated_pages_to_v6 || migrated_header_to_v7;
698 if migrated {
699 // Stamp the version forward in memory; the next flush persists
700 // it via the WAL-committed header frame (and the v7 seal).
701 header.version = VERSION;
702 // Snapshots written under an older format reference page runs
703 // encrypted under the old AAD scheme; this build can't decrypt
704 // them under the new pager. Drop the manifest so the next
705 // flush records a fresh snapshot of the migrated state.
706 // Point-in-time recovery into the pre-migration past is
707 // sacrificed for migration simplicity. Old pages remain on
708 // disk until `compact_file`. (For v6→v7 alone — which doesn't
709 // change page crypto — the old snapshots would technically
710 // still be readable, but dropping them uniformly keeps the
711 // migration policy simple.)
712 manifest.clear();
713 }
714
715 // The ANN index pages on disk were sealed under the *old* AAD scheme
716 // if we migrated pages to v6; force a rewrite at fresh v6-encrypted
717 // slots on the next flush so a future reopen doesn't try to read
718 // them under the new AAD and fail. (Catalog and manifest already
719 // get rewritten via dirty_catalog and the manifest.clear() above.)
720 let dirty_index = migrated_pages_to_v6 && ann.is_some();
721
722 Ok(Mnemo {
723 pager,
724 header,
725 catalog,
726 index,
727 path,
728 dimensions,
729 kdf,
730 dirty_catalog: migrated,
731 ann,
732 dirty_index,
733 manifest,
734 // `Mnemo::open` has no [`MnemoConfig`] to read; default the cap
735 // to [`DEFAULT_MAX_SNAPSHOTS`]. Override with
736 // [`Mnemo::set_max_snapshots`] right after open if you need
737 // unlimited retention or a different cap for this handle.
738 max_snapshots: DEFAULT_MAX_SNAPSHOTS,
739 })
740 }
741
742 /// Override the manifest snapshot cap on this open handle. `0` disables
743 /// the cap; any positive value keeps the most-recent `max` snapshots and
744 /// drops the rest on the next flush. See [`MnemoConfig::max_snapshots`].
745 pub fn set_max_snapshots(&mut self, max: usize) {
746 self.max_snapshots = max;
747 }
748
749 /// Embedding dimensionality this database expects.
750 pub fn dimensions(&self) -> usize {
751 self.dimensions
752 }
753
754 /// Number of live (non-deleted) memories.
755 pub fn len(&self) -> usize {
756 self.catalog.iter().filter(|e| !e.deleted).count()
757 }
758
759 /// True if the database holds no live memories.
760 pub fn is_empty(&self) -> bool {
761 self.len() == 0
762 }
763
764 // --- internal page/record helpers -------------------------------------
765
766 fn write_record(&mut self, bytes: &[u8]) -> Result<(u64, u32)> {
767 let pc = bytes.len().div_ceil(PAYLOAD).max(1);
768 let start = self.header.next_page;
769 self.header.next_page += pc as u64;
770 for i in 0..pc {
771 let lo = i * PAYLOAD;
772 let hi = ((i + 1) * PAYLOAD).min(bytes.len());
773 self.pager.write_page(start + i as u64, &bytes[lo..hi])?;
774 }
775 Ok((start, pc as u32))
776 }
777
778 fn read_record(&mut self, e: &CatalogEntry) -> Result<Vec<u8>> {
779 let mut buf = Vec::with_capacity(e.len as usize);
780 for i in 0..e.page_count as u64 {
781 buf.extend_from_slice(&self.pager.read_page(e.start_page + i)?);
782 }
783 buf.truncate(e.len as usize);
784 Ok(buf)
785 }
786
787 fn read_memory(&mut self, e: &CatalogEntry) -> Result<Memory> {
788 let bytes = self.read_record(e)?;
789 let mut m: Memory = rmp_serde::from_slice(&bytes)
790 .map_err(|err| MnemoError::Serialize(err.to_string()))?;
791 // The catalog is the source of truth for access stats — the values
792 // in the record body are a stale snapshot from when the record was
793 // last written. Overwrite them here so consumers always see live
794 // values regardless of when the record was last serialized.
795 m.accessed_at = e.accessed_at;
796 m.access_count = e.access_count;
797 Ok(m)
798 }
799
800 /// Serialize and store a memory (insert or overwrite by ID).
801 fn put(&mut self, mut m: Memory) -> Result<Ulid> {
802 if m.vector.len() != self.dimensions {
803 return Err(MnemoError::DimensionMismatch {
804 expected: self.dimensions,
805 got: m.vector.len(),
806 });
807 }
808 if m.id == Ulid::nil() {
809 m.id = Ulid::new();
810 }
811 let id_u: u128 = m.id.0;
812 let bytes = rmp_serde::to_vec(&m).map_err(|e| MnemoError::Serialize(e.to_string()))?;
813 let (start, pc) = self.write_record(&bytes)?;
814 // For overwrites preserve the existing access stats; for inserts
815 // seed them from the memory (which carries them in its body).
816 let (prev_accessed, prev_count) = match self.index.get(&id_u).copied() {
817 Some(idx) => (self.catalog[idx].accessed_at, self.catalog[idx].access_count),
818 None => (m.accessed_at, m.access_count),
819 };
820 let entry = CatalogEntry {
821 id: id_u,
822 start_page: start,
823 page_count: pc,
824 len: bytes.len() as u32,
825 deleted: false,
826 accessed_at: prev_accessed,
827 access_count: prev_count,
828 };
829 match self.index.get(&id_u).copied() {
830 Some(idx) => self.catalog[idx] = entry,
831 None => {
832 self.index.insert(id_u, self.catalog.len());
833 self.catalog.push(entry);
834 }
835 }
836 self.dirty_catalog = true;
837
838 // Keep the ANN index complete: assign the (new or changed) vector to
839 // its nearest partition. Centroids/codebook stay fixed until rebuild.
840 if let Some(ann) = &mut self.ann {
841 ann.add(id_u, &m.vector);
842 self.dirty_index = true;
843 }
844 Ok(m.id)
845 }
846
847 // --- public CRUD ------------------------------------------------------
848
849 /// Store a memory. Returns its ULID. If the memory's ID is nil a fresh
850 /// one is assigned; an existing ID overwrites in place.
851 pub fn remember(&mut self, memory: Memory) -> Result<Ulid> {
852 self.put(memory)
853 }
854
855 /// Fetch a memory by ID.
856 pub fn get(&mut self, id: &Ulid) -> Result<Memory> {
857 let idx = *self
858 .index
859 .get(&id.0)
860 .ok_or_else(|| MnemoError::NotFound(id.to_string()))?;
861 let entry = self.catalog[idx].clone();
862 if entry.deleted {
863 return Err(MnemoError::NotFound(id.to_string()));
864 }
865 self.read_memory(&entry)
866 }
867
868 /// Soft-delete a memory (tombstoned; space reclaimed by `compact`).
869 pub fn delete(&mut self, id: &Ulid) -> Result<()> {
870 let idx = *self
871 .index
872 .get(&id.0)
873 .ok_or_else(|| MnemoError::NotFound(id.to_string()))?;
874 if !self.catalog[idx].deleted {
875 self.catalog[idx].deleted = true;
876 self.dirty_catalog = true;
877 }
878 if let Some(ann) = &mut self.ann {
879 ann.remove(id.0);
880 self.dirty_index = true;
881 }
882 Ok(())
883 }
884
885 /// Return every live memory (used by tooling and compaction).
886 pub fn memories(&mut self) -> Result<Vec<Memory>> {
887 let entries: Vec<CatalogEntry> = self
888 .catalog
889 .iter()
890 .filter(|e| !e.deleted)
891 .cloned()
892 .collect();
893 let mut out = Vec::with_capacity(entries.len());
894 for e in &entries {
895 out.push(self.read_memory(e)?);
896 }
897 Ok(out)
898 }
899
900 /// Return the database's **self-describing onboarding memories** — the
901 /// ones tagged `metadata.area = "onboarding"`. This is the engine-level
902 /// surface for the *single-file philosophy*: an agent who receives a
903 /// `.mnemo` file (and its passphrase) can call this to learn what the
904 /// file is, which embedder it expects, the recommended agent id, and any
905 /// other conventions the file's author chose to record — all without
906 /// needing any external documentation.
907 ///
908 /// Ordering: the canonical manifest (tag `metadata.topic = "manifest"`)
909 /// always comes first; everything else follows in `importance` descending,
910 /// then `created_at` ascending for deterministic results.
911 pub fn about(&mut self) -> Result<Vec<Memory>> {
912 let mut out: Vec<Memory> = self
913 .memories()?
914 .into_iter()
915 .filter(|m| {
916 m.metadata
917 .get("area")
918 .and_then(|v| v.as_str())
919 .map(|s| s.eq_ignore_ascii_case("onboarding"))
920 .unwrap_or(false)
921 })
922 .collect();
923 out.sort_by(|a, b| {
924 let a_manifest = is_manifest(a);
925 let b_manifest = is_manifest(b);
926 // Manifest topic wins first; then importance desc; then created_at asc.
927 b_manifest
928 .cmp(&a_manifest)
929 .then_with(|| b.importance.total_cmp(&a.importance))
930 .then_with(|| a.created_at.cmp(&b.created_at))
931 });
932 Ok(out)
933 }
934
935 // --- retrieval --------------------------------------------------------
936
937 /// Phase-1 brute-force search: rank live memories by raw similarity only.
938 /// Read-only — does not update access statistics.
939 pub fn search(
940 &mut self,
941 query: &[f32],
942 top_k: usize,
943 metric: Metric,
944 ) -> Result<Vec<(Memory, f32)>> {
945 if query.len() != self.dimensions {
946 return Err(MnemoError::DimensionMismatch {
947 expected: self.dimensions,
948 got: query.len(),
949 });
950 }
951 let now = memory::now_secs();
952 let entries: Vec<CatalogEntry> = self
953 .catalog
954 .iter()
955 .filter(|e| !e.deleted)
956 .cloned()
957 .collect();
958
959 let mut scored: Vec<(Memory, f32)> = Vec::new();
960 for e in &entries {
961 let m = self.read_memory(e)?;
962 if m.is_expired(now) {
963 continue;
964 }
965 let sim = memory::similarity(metric, query, &m.vector);
966 scored.push((m, sim));
967 }
968 scored.sort_by(|a, b| b.1.total_cmp(&a.1));
969 scored.truncate(top_k);
970 Ok(scored)
971 }
972
973 /// Phase-5 multi-signal recall: rank by
974 /// `α·similarity + β·recency + γ·importance + δ·ln(freq)`, with type and
975 /// agent filtering. Updates `accessed_at` / `access_count` on the
976 /// returned memories (persisted on the next `flush`).
977 ///
978 /// When an ANN index has been built ([`Mnemo::build_index`]), recall runs
979 /// the tiered IVF→PQ→rerank pipeline: it scores only the similarity-nearest
980 /// `n_rerank` candidates rather than every memory. Without an index it
981 /// scores every live memory exactly.
982 pub fn recall(&mut self, req: &RecallRequest) -> Result<Vec<RecallResult>> {
983 if req.query.len() != self.dimensions {
984 return Err(MnemoError::DimensionMismatch {
985 expected: self.dimensions,
986 got: req.query.len(),
987 });
988 }
989 let now = memory::now_secs();
990
991 // Candidate set: ANN-narrowed when an index exists, else everything.
992 let entries: Vec<CatalogEntry> = if let Some(ann) = &self.ann {
993 let ids = ann.query(&req.query, req.n_probe, req.n_rerank);
994 ids.iter()
995 .filter_map(|id| self.index.get(id).copied())
996 .map(|i| self.catalog[i].clone())
997 .filter(|e| !e.deleted)
998 .collect()
999 } else {
1000 self.catalog
1001 .iter()
1002 .filter(|e| !e.deleted)
1003 .cloned()
1004 .collect()
1005 };
1006
1007 let mut scored: Vec<RecallResult> = Vec::new();
1008 for e in &entries {
1009 let m = self.read_memory(e)?;
1010 if m.is_expired(now) {
1011 continue;
1012 }
1013 // Type filter.
1014 if let Some(types) = &req.memory_types {
1015 if !types.contains(&m.memory_type) {
1016 continue;
1017 }
1018 }
1019 // Agent-scoping: an agent sees its own memories plus shared ones.
1020 if let Some(agent) = &req.agent_id {
1021 let visible = m.agent_id == *agent || m.scope == Scope::Shared;
1022 if !visible {
1023 continue;
1024 }
1025 }
1026 let sim = memory::similarity(req.metric, &req.query, &m.vector);
1027 let age = (now - m.accessed_at) as f32;
1028 let score = req.weights.score(sim, age, m.importance, m.access_count);
1029 scored.push(RecallResult { memory: m, score, similarity: sim });
1030 }
1031 scored.sort_by(|a, b| b.score.total_cmp(&a.score));
1032 scored.truncate(req.top_k);
1033
1034 // Update access statistics for everything we surfaced — but only
1035 // mutate the catalog entry, not the full record on disk. Pre-v5
1036 // this loop called `self.put(r.memory.clone())`, which serialized
1037 // and rewrote every result (vector included) to fresh pages,
1038 // turning a top-K recall into K full record rewrites. Now the
1039 // catalog is the source of truth for these two fields and the
1040 // record body stays untouched until the next real edit.
1041 if req.track_access {
1042 for r in &mut scored {
1043 if let Some(&idx) = self.index.get(&r.memory.id.0) {
1044 let entry = &mut self.catalog[idx];
1045 entry.accessed_at = now;
1046 entry.access_count = entry.access_count.saturating_add(1);
1047 // Propagate the just-updated values onto the returned
1048 // Memory so the caller sees the post-recall state.
1049 r.memory.accessed_at = entry.accessed_at;
1050 r.memory.access_count = entry.access_count;
1051 }
1052 }
1053 // Mark the catalog dirty so the next flush persists these
1054 // bumps. The records themselves are not dirty.
1055 if !scored.is_empty() {
1056 self.dirty_catalog = true;
1057 }
1058 }
1059 Ok(scored)
1060 }
1061
1062 // --- approximate index ------------------------------------------------
1063
1064 /// Build an IVF+PQ approximate-nearest-neighbour index over every live
1065 /// memory, using default tuning. After this, [`Mnemo::recall`] runs the
1066 /// tiered pipeline instead of an exact scan. Persisted on the next
1067 /// `flush`. Returns a snapshot of the index shape.
1068 pub fn build_index(&mut self) -> Result<IndexInfo> {
1069 self.build_index_with(IndexConfig::default())
1070 }
1071
1072 /// Build the ANN index with explicit tuning.
1073 pub fn build_index_with(&mut self, cfg: IndexConfig) -> Result<IndexInfo> {
1074 let mems = self.memories()?;
1075 if mems.is_empty() {
1076 return Err(MnemoError::Invalid(
1077 "cannot build an index over an empty database".into(),
1078 ));
1079 }
1080 let items: Vec<(u128, &[f32])> =
1081 mems.iter().map(|m| (m.id.0, m.vector.as_slice())).collect();
1082 let idx = IvfPqIndex::build(self.dimensions, &items, cfg)?;
1083 let info = idx.info();
1084 self.ann = Some(idx);
1085 self.dirty_index = true;
1086 Ok(info)
1087 }
1088
1089 /// Rebuild the ANN index from scratch — re-clusters centroids and retrains
1090 /// the PQ codebook, undoing the cluster drift that accumulates as memories
1091 /// are inserted against fixed centroids.
1092 pub fn rebuild_index(&mut self) -> Result<IndexInfo> {
1093 let cfg = match &self.ann {
1094 Some(a) => IndexConfig {
1095 n_probe: a.n_probe(),
1096 n_rerank: a.n_rerank(),
1097 ..Default::default()
1098 },
1099 None => IndexConfig::default(),
1100 };
1101 self.build_index_with(cfg)
1102 }
1103
1104 /// Drop the ANN index; recall reverts to exact scans. Persisted on flush.
1105 pub fn drop_index(&mut self) {
1106 if self.ann.is_some() {
1107 self.ann = None;
1108 self.dirty_index = true;
1109 }
1110 }
1111
1112 /// Whether an ANN index is currently loaded.
1113 pub fn has_index(&self) -> bool {
1114 self.ann.is_some()
1115 }
1116
1117 // --- durability & maintenance ----------------------------------------
1118
1119 /// Seal a serialized buffer into a fresh run of encrypted home pages and
1120 /// append a WAL frame for each. Returns the run's `(start_page, pages)`.
1121 /// The pages are *not* written to their home locations here — they go to
1122 /// the WAL and are folded in later by [`Mnemo::checkpoint`].
1123 fn seal_run(&mut self, bytes: &[u8], frames: &mut Vec<wal::Frame>) -> Result<(u64, u32)> {
1124 let pc = bytes.len().div_ceil(PAYLOAD).max(1);
1125 let start = self.header.next_page;
1126 self.header.next_page += pc as u64;
1127 for i in 0..pc {
1128 let lo = i * PAYLOAD;
1129 let hi = ((i + 1) * PAYLOAD).min(bytes.len());
1130 let page_no = start + i as u64;
1131 let sealed = self.pager.seal_page(page_no, &bytes[lo..hi])?;
1132 frames.push((page_no, sealed.to_vec()));
1133 }
1134 Ok((start, pc as u32))
1135 }
1136
1137 /// Grow the WAL region if `frame_pages` worth of frames would not fit.
1138 /// Called at the top of `flush`, where the WAL is always spent (every
1139 /// flush ends with a checkpoint), so relocating it cannot strand a
1140 /// committed transaction.
1141 fn ensure_wal_capacity(&mut self, frame_pages: usize) -> Result<()> {
1142 let need = wal::txn_byte_len(frame_pages, PAGE_SIZE);
1143 if need <= self.header.wal_pages * PAGE_SIZE as u64 {
1144 return Ok(());
1145 }
1146 // Allocate a new, larger region at the file tail (~1.5x headroom).
1147 let req = need.div_ceil(PAGE_SIZE as u64);
1148 let new_pages = (req + req / 2 + 4).max(DEFAULT_WAL_PAGES);
1149 let new_start = self.header.next_page;
1150 self.header.next_page += new_pages;
1151 self.header.wal_start = new_start;
1152 self.header.wal_pages = new_pages;
1153 // Persist the relocation now: an isolated header write over an empty
1154 // WAL. A crash here leaves a consistent state with a bigger WAL.
1155 self.header.write_counter = self.pager.write_counter;
1156 let mut page = self.header.to_page();
1157 self.header.apply_seal(&mut page, self.pager.dek())?;
1158 self.pager.write_raw(0, &page)?;
1159 self.pager.sync()?;
1160 Ok(())
1161 }
1162
1163 /// Fold a committed transaction's WAL frames into their home pages.
1164 fn checkpoint(&mut self, frames: &[wal::Frame]) -> Result<()> {
1165 for (page_no, bytes) in frames {
1166 let mut img = [0u8; PAGE_SIZE];
1167 img.copy_from_slice(bytes);
1168 self.pager.write_sealed(*page_no, &img)?;
1169 }
1170 self.pager.sync()?;
1171 Ok(())
1172 }
1173
1174 /// Persist all pending changes as one **write-ahead-logged transaction**.
1175 ///
1176 /// Sequence:
1177 ///
1178 /// 1. **Prepare.** Serialize the control plane (catalog, ANN index),
1179 /// grow the WAL if needed, and **lease** counter and page slots —
1180 /// bump `header.write_counter` and `header.next_page` by upper
1181 /// bounds on this transaction's consumption and persist that
1182 /// *leased* header to disk before any encrypted page is written.
1183 /// 2. **Data pages.** [`Pager::flush`] writes dirty record pages
1184 /// copy-on-write under fresh nonces and fsyncs them. The orphan-page
1185 /// nonce-reuse window (see [Phase 1.1 of the improvement plan])
1186 /// is closed because the leased header on disk already records a
1187 /// `write_counter` past anything this step can produce.
1188 /// 3. **Control plane.** Seal the new catalog / ANN index / snapshot
1189 /// manifest into fresh home page runs.
1190 /// 4. **Commit.** Log all of step 3's frames plus the final header
1191 /// frame into the WAL and fsync — the durability point.
1192 /// 5. **Checkpoint.** Fold the WAL into the home pages.
1193 ///
1194 /// A crash before step 4 leaves the previous state intact. A crash
1195 /// after step 4 is repaired by [`Mnemo::open`] replaying the WAL.
1196 /// Safe to call repeatedly.
1197 pub fn flush(&mut self) -> Result<()> {
1198 if !self.dirty_catalog && !self.dirty_index {
1199 // No control-plane changes pending. Because every `put` that
1200 // dirties a data page also dirties the catalog, there can't
1201 // be dirty data pages either — nothing to do.
1202 return Ok(());
1203 }
1204
1205 // 1. Serialize control plane, ensure WAL capacity, lease counter +
1206 // page slots, and persist the leased header *before* any data
1207 // page hits the disk.
1208 let ctx = self.prepare_for_flush()?;
1209
1210 // 2. Record (vector) data pages: copy-on-write to fresh pages, fsynced.
1211 // Safe: the lease guarantees that even if we crash here, reopen
1212 // will see a `write_counter` past every nonce this step uses.
1213 self.pager.flush()?;
1214
1215 // 3. Seal the catalog / index control pages into fresh home runs.
1216 let mut frames: Vec<wal::Frame> = Vec::new();
1217 if let Some(bytes) = &ctx.cat_bytes {
1218 let (start, pages) = self.seal_run(bytes, &mut frames)?;
1219 self.header.catalog_start = start;
1220 self.header.catalog_pages = pages as u64;
1221 self.header.catalog_len = bytes.len() as u64;
1222 }
1223 if self.dirty_index {
1224 match &ctx.idx_bytes {
1225 Some(bytes) => {
1226 let (start, pages) = self.seal_run(bytes, &mut frames)?;
1227 self.header.index_start = start;
1228 self.header.index_pages = pages as u64;
1229 self.header.index_len = bytes.len() as u64;
1230 }
1231 None => {
1232 self.header.index_start = 0;
1233 self.header.index_pages = 0;
1234 self.header.index_len = 0;
1235 }
1236 }
1237 }
1238
1239 // 3b. Record this transaction as a restorable snapshot. The manifest
1240 // update is staged in a local and adopted only once the commit
1241 // below succeeds, so a failed flush leaves no phantom entry.
1242 //
1243 // The cap (`self.max_snapshots`) is applied to the staged local
1244 // before the new entry lands, so after appending the manifest
1245 // holds at most `max_snapshots` entries — the most-recent N.
1246 // `max_snapshots == 0` disables the cap. Pages referenced only
1247 // by pruned snapshots stay on disk until `compact_file`.
1248 let live = self.len() as u64;
1249 let txn_id = self.header.wal_seq + 1;
1250 let mut manifest = self.manifest.clone();
1251 if self.max_snapshots > 0 {
1252 // Drop the oldest entries so there's room for the new one.
1253 while manifest.len() >= self.max_snapshots {
1254 manifest.remove(0);
1255 }
1256 }
1257 manifest.push(Snapshot {
1258 txn_id,
1259 created_at: memory::now_secs(),
1260 catalog_start: self.header.catalog_start,
1261 catalog_pages: self.header.catalog_pages,
1262 catalog_len: self.header.catalog_len,
1263 index_start: self.header.index_start,
1264 index_pages: self.header.index_pages,
1265 index_len: self.header.index_len,
1266 memory_count: live,
1267 });
1268 let man_bytes =
1269 rmp_serde::to_vec(&manifest).map_err(|e| MnemoError::Serialize(e.to_string()))?;
1270 let (m_start, m_pages) = self.seal_run(&man_bytes, &mut frames)?;
1271 self.header.manifest_start = m_start;
1272 self.header.manifest_pages = m_pages as u64;
1273 self.header.manifest_len = man_bytes.len() as u64;
1274
1275 // 3c. The header is the transaction's final frame; stamp the new id
1276 // and the *actual* (post-flush) write counter, which is <= the
1277 // leased value persisted in `prepare_for_flush`. This is the
1278 // value that ultimately replaces the leased header on checkpoint.
1279 self.header.vector_count = live;
1280 self.header.write_counter = self.pager.write_counter;
1281 self.header.wal_seq = txn_id;
1282 let mut hpage = self.header.to_page();
1283 self.header.apply_seal(&mut hpage, self.pager.dek())?;
1284 frames.push((0, hpage.to_vec()));
1285
1286 // 4. COMMIT — log the transaction and fsync. Nothing at a home page
1287 // other than the leased header has changed yet; this single
1288 // fsync is what makes the rest durable.
1289 let (wal_start, wal_pages) = (self.header.wal_start, self.header.wal_pages);
1290 wal::commit(self.pager.file_mut(), wal_start, wal_pages, txn_id, &frames)?;
1291
1292 // 5. Checkpoint — fold the WAL into the home pages.
1293 self.checkpoint(&frames)?;
1294
1295 self.manifest = manifest;
1296 self.dirty_catalog = false;
1297 self.dirty_index = false;
1298 Ok(())
1299 }
1300
1301 /// Pre-flush prelude shared by [`Mnemo::flush`] and the
1302 /// `__crash_partial_flush_for_testing` hook.
1303 ///
1304 /// Serializes the control plane (catalog, ANN index), grows the WAL if
1305 /// needed, then **leases** counter and page slots — bumps a *clone* of
1306 /// the header by upper-bound amounts and persists that leased clone
1307 /// with `pager.write_raw + sync`. The in-memory `self.header` keeps
1308 /// its pre-lease `next_page` so [`Mnemo::seal_run`] continues to
1309 /// allocate from the right slot; the final committed header carries
1310 /// the actual post-flush values and overwrites the leased one on
1311 /// checkpoint.
1312 ///
1313 /// This pre-stamping is what closes the nonce-reuse window flagged in
1314 /// Phase 1.1 of the improvement plan: a crash anywhere between here
1315 /// and the WAL commit leaves a file whose on-disk header records a
1316 /// `write_counter` *past* every nonce the rest of this transaction
1317 /// could produce, so the next reopen starts from a counter that
1318 /// guarantees uniqueness against any orphan data pages still on disk.
1319 fn prepare_for_flush(&mut self) -> Result<FlushPrelude> {
1320 // Serialize.
1321 let cat_bytes = if self.dirty_catalog {
1322 Some(
1323 rmp_serde::to_vec(&self.catalog)
1324 .map_err(|e| MnemoError::Serialize(e.to_string()))?,
1325 )
1326 } else {
1327 None
1328 };
1329 let idx_bytes: Option<Vec<u8>> = if self.dirty_index {
1330 match &self.ann {
1331 Some(ann) => Some(
1332 rmp_serde::to_vec(ann).map_err(|e| MnemoError::Serialize(e.to_string()))?,
1333 ),
1334 None => None,
1335 }
1336 } else {
1337 None
1338 };
1339
1340 let cat_pc = cat_bytes.as_ref().map_or(0, |b| b.len().div_ceil(PAYLOAD).max(1));
1341 let idx_pc = idx_bytes.as_ref().map_or(0, |b| b.len().div_ceil(PAYLOAD).max(1));
1342 // Upper bound on the manifest run: one extra entry, <=82 bytes each
1343 // (rmp_serde encodes a Snapshot as a 9-element fixarray of i64/u64,
1344 // worst case 81 bytes plus a 1-byte array header).
1345 let man_upper = 9 + (self.manifest.len() + 1) * 82;
1346 let man_pc_est = man_upper.div_ceil(PAYLOAD).max(1);
1347
1348 // Grow WAL if needed (does its own header write+sync).
1349 self.ensure_wal_capacity(cat_pc + idx_pc + man_pc_est + 1)?;
1350
1351 // Lease counter and page slots. counter_lease is an upper bound on
1352 // how many `pager.write_counter` advances this transaction will
1353 // make (one per dirty data page + one per sealed control-plane
1354 // page). page_lease covers only the control-plane allocations,
1355 // because dirty data pages were already allocated past
1356 // `header.next_page` by `write_record` when `remember` ran.
1357 let data_dirty = self.pager.dirty_page_count() as u64;
1358 let counter_lease = data_dirty + (cat_pc + idx_pc + man_pc_est) as u64;
1359 let page_lease = (cat_pc + idx_pc + man_pc_est) as u64;
1360
1361 // Persist a leased *clone* of the header — in-memory `self.header`
1362 // keeps its pre-lease `next_page` for `seal_run` to consume from
1363 // the correct slot.
1364 let mut leased = self.header.clone();
1365 leased.write_counter = self.pager.write_counter + counter_lease;
1366 leased.next_page += page_lease;
1367 let mut page = leased.to_page();
1368 leased.apply_seal(&mut page, self.pager.dek())?;
1369 self.pager.write_raw(0, &page)?;
1370 self.pager.sync()?;
1371
1372 Ok(FlushPrelude { cat_bytes, idx_bytes })
1373 }
1374
1375 /// Flush and close. Equivalent to `flush()`; the file is released on drop.
1376 pub fn close(&mut self) -> Result<()> {
1377 self.flush()
1378 }
1379
1380 /// **TEST ONLY — do not use.** Reproduces a crash inside
1381 /// [`Mnemo::flush`] *after* the prelude has leased counter+page slots
1382 /// and persisted the leased header, but *before* the WAL is committed.
1383 /// Runs `prepare_for_flush` + `pager.flush` and returns — leaving
1384 /// data pages and a leased header on disk, with no commit frame.
1385 ///
1386 /// Used by `tests/integration.rs::nonce_unique_after_crashed_data_flush`
1387 /// to verify that this window (Phase 1.1 of the improvement plan) does
1388 /// not enable AES-GCM nonce reuse. Calling this in production code
1389 /// strands data pages; never expose it from a binding.
1390 #[doc(hidden)]
1391 pub fn __crash_partial_flush_for_testing(&mut self) -> Result<()> {
1392 if !self.dirty_catalog && !self.dirty_index {
1393 return Ok(());
1394 }
1395 let _prelude = self.prepare_for_flush()?;
1396 self.pager.flush()
1397 }
1398
1399 /// Bound the in-memory page cache to `pages` decrypted pages.
1400 ///
1401 /// The cache holds decrypted page payloads to speed repeated reads. By
1402 /// default it is capped at 8192 pages (~64 MiB); lower the cap to trade
1403 /// hit rate for a smaller footprint, or raise it for a hotter cache. The
1404 /// cap governs *clean* pages — pages with un-flushed writes are always
1405 /// retained until [`Mnemo::flush`], regardless of the cap.
1406 pub fn set_cache_capacity(&mut self, pages: usize) {
1407 self.pager.set_cache_capacity(pages);
1408 }
1409
1410 /// Page-cache occupancy: `(pages_cached, capacity)`.
1411 pub fn cache_stats(&self) -> (usize, usize) {
1412 self.pager.cache_stats()
1413 }
1414
1415 /// Begin a conversation [`Session`](crate::Session) for `agent_id`.
1416 ///
1417 /// The session borrows the database for its lifetime, records turns as
1418 /// working memory, and consolidates them into episodic memory when closed.
1419 pub fn session(&mut self, agent_id: impl Into<String>) -> crate::session::Session<'_> {
1420 crate::session::Session::new(self, agent_id.into())
1421 }
1422
1423 // --- snapshots & point-in-time recovery ------------------------------
1424
1425 /// Every committed transaction, oldest first — the restore points
1426 /// available to [`Mnemo::restore_to`] and [`Mnemo::restore_to_time`].
1427 ///
1428 /// Each `flush` appends one snapshot. Because the storage engine is
1429 /// append-only, the pages a past flush wrote are still on disk, so any
1430 /// listed snapshot can be reinstated exactly. The history reaches back to
1431 /// the last [`Mnemo::compact_file`], which reclaims space by collapsing
1432 /// it.
1433 pub fn snapshots(&self) -> Vec<SnapshotInfo> {
1434 let mut v: Vec<SnapshotInfo> = self
1435 .manifest
1436 .iter()
1437 .map(|s| SnapshotInfo {
1438 txn_id: s.txn_id,
1439 created_at: s.created_at,
1440 memory_count: s.memory_count,
1441 })
1442 .collect();
1443 v.sort_by_key(|s| s.txn_id);
1444 v
1445 }
1446
1447 /// Load a past snapshot's state and commit it as a new transaction.
1448 fn apply_snapshot(&mut self, snap: &Snapshot) -> Result<()> {
1449 let catalog = load_catalog(
1450 &mut self.pager,
1451 snap.catalog_start,
1452 snap.catalog_pages,
1453 snap.catalog_len,
1454 )?;
1455 let ann = load_index(
1456 &mut self.pager,
1457 snap.index_start,
1458 snap.index_pages,
1459 snap.index_len,
1460 self.dimensions,
1461 )?;
1462 self.index = build_id_index(&catalog);
1463 self.catalog = catalog;
1464 self.ann = ann;
1465 // Re-commit as a fresh transaction: crash-safe, and itself recorded
1466 // as a new snapshot so a restore can always be undone.
1467 self.dirty_catalog = true;
1468 self.dirty_index = true;
1469 self.flush()
1470 }
1471
1472 /// Restore the database to the snapshot produced by transaction `txn_id`.
1473 ///
1474 /// The restore is itself a new committed transaction (and a new
1475 /// snapshot), so it is crash-safe and reversible — restoring forward to a
1476 /// later snapshot afterwards works just as well.
1477 pub fn restore_to(&mut self, txn_id: u64) -> Result<SnapshotInfo> {
1478 let snap = self
1479 .manifest
1480 .iter()
1481 .find(|s| s.txn_id == txn_id)
1482 .cloned()
1483 .ok_or_else(|| MnemoError::NotFound(format!("snapshot for transaction {txn_id}")))?;
1484 self.apply_snapshot(&snap)?;
1485 Ok(SnapshotInfo {
1486 txn_id: snap.txn_id,
1487 created_at: snap.created_at,
1488 memory_count: snap.memory_count,
1489 })
1490 }
1491
1492 /// Restore the database to the latest snapshot committed at or before
1493 /// `unix_secs`. Returns [`MnemoError::NotFound`] if no snapshot is that
1494 /// old. Like [`Mnemo::restore_to`], the restore is a new transaction.
1495 pub fn restore_to_time(&mut self, unix_secs: i64) -> Result<SnapshotInfo> {
1496 let snap = self
1497 .manifest
1498 .iter()
1499 .filter(|s| s.created_at <= unix_secs)
1500 .max_by_key(|s| s.txn_id)
1501 .cloned()
1502 .ok_or_else(|| {
1503 MnemoError::NotFound(format!("no snapshot at or before time {unix_secs}"))
1504 })?;
1505 self.apply_snapshot(&snap)?;
1506 Ok(SnapshotInfo {
1507 txn_id: snap.txn_id,
1508 created_at: snap.created_at,
1509 memory_count: snap.memory_count,
1510 })
1511 }
1512
1513 /// Change the passphrase. Cheap: re-derives the KEK and re-wraps the DEK;
1514 /// the encrypted pages are never rewritten.
1515 pub fn rekey(&mut self, new_passphrase: &str, kdf: KdfParams) -> Result<()> {
1516 self.flush()?;
1517 let new_salt = crypto::random_salt();
1518 let new_kek = crypto::derive_kek(new_passphrase.as_bytes(), &new_salt, kdf)?;
1519 let new_nonce = crypto::random_nonce();
1520 let wrapped = crypto::wrap_dek(&new_kek, &new_nonce, self.pager.dek())?;
1521 if wrapped.len() != WRAPPED_DEK_LEN {
1522 return Err(MnemoError::Crypto("unexpected wrapped DEK length".into()));
1523 }
1524 let mut wrapped_dek = [0u8; WRAPPED_DEK_LEN];
1525 wrapped_dek.copy_from_slice(&wrapped);
1526
1527 self.header.salt = new_salt;
1528 self.header.dek_nonce = new_nonce;
1529 self.header.wrapped_dek = wrapped_dek;
1530 self.header.m_cost = kdf.m_cost;
1531 self.header.t_cost = kdf.t_cost;
1532 self.header.p_cost = kdf.p_cost;
1533 self.header.write_counter = self.pager.write_counter;
1534 self.kdf = kdf;
1535
1536 let mut page = self.header.to_page();
1537 self.header.apply_seal(&mut page, self.pager.dek())?;
1538 self.pager.write_raw(0, &page)?;
1539 self.pager.sync()?;
1540 Ok(())
1541 }
1542
1543 /// Decrypt and validate every live record. Returns the count verified.
1544 pub fn verify(&mut self) -> Result<usize> {
1545 let entries: Vec<CatalogEntry> = self
1546 .catalog
1547 .iter()
1548 .filter(|e| !e.deleted)
1549 .cloned()
1550 .collect();
1551 let n = entries.len();
1552 for e in &entries {
1553 let bytes = self.read_record(e)?;
1554 let _m: Memory = rmp_serde::from_slice(&bytes)
1555 .map_err(|err| MnemoError::Serialize(err.to_string()))?;
1556 }
1557 Ok(n)
1558 }
1559
1560 /// Summary statistics for the open database.
1561 pub fn stats(&mut self) -> Result<Stats> {
1562 let deleted = self.catalog.iter().filter(|e| e.deleted).count();
1563 let mut agents: Vec<String> = self
1564 .memories()?
1565 .into_iter()
1566 .map(|m| m.agent_id)
1567 .collect();
1568 agents.sort();
1569 agents.dedup();
1570 let file_bytes = self.header.next_page * PAGE_SIZE as u64;
1571 Ok(Stats {
1572 memories: self.len(),
1573 deleted,
1574 dimensions: self.dimensions,
1575 file_bytes,
1576 agents,
1577 encrypted: self.header.flags & FLAG_ENCRYPTED != 0,
1578 created_at: self.header.created_at,
1579 index: self.ann.as_ref().map(|a| a.info()),
1580 wal_pages: self.header.wal_pages,
1581 })
1582 }
1583
1584 /// Rewrite the file, dropping tombstoned and expired memories and
1585 /// reclaiming stale pages left by updates. Done out-of-place via a temp
1586 /// file and an atomic rename.
1587 pub fn compact_file(path: &str, passphrase: &str) -> Result<CompactReport> {
1588 let mut old = Mnemo::open(path, passphrase)?;
1589 let dims = old.dimensions;
1590 let kdf = old.kdf;
1591 let tmp = format!("{path}.compact-tmp");
1592
1593 let mut new = Mnemo::create(
1594 &tmp,
1595 passphrase,
1596 MnemoConfig { dimensions: dims, kdf, ..Default::default() },
1597 )?;
1598 let want_index = old.ann.as_ref().map(|a| (a.n_probe(), a.n_rerank()));
1599 let now = memory::now_secs();
1600 let all = old.memories()?;
1601 let before = all.len();
1602 let mut after = 0;
1603 for m in all {
1604 if m.is_expired(now) {
1605 continue;
1606 }
1607 new.remember(m)?;
1608 after += 1;
1609 }
1610 // Rebuild the index fresh (re-clustered) when the source had one.
1611 if let Some((n_probe, n_rerank)) = want_index {
1612 if after > 0 {
1613 new.build_index_with(IndexConfig {
1614 n_probe,
1615 n_rerank,
1616 ..Default::default()
1617 })?;
1618 }
1619 }
1620 new.flush()?;
1621 drop(new);
1622 drop(old);
1623
1624 std::fs::rename(&tmp, path)?;
1625 Ok(CompactReport { before, after })
1626 }
1627}