Skip to main content

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}