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::{File, OpenOptions};
26use std::io::{Read, Seek, SeekFrom, Write};
27use std::path::{Path, PathBuf};
28
29use fs4::fs_std::FileExt;
30use serde::{Deserialize, Serialize};
31use ulid::Ulid;
32
33use crate::crypto::{self, KdfParams};
34use crate::embed::{self, Embedder};
35use crate::error::{MnemoError, Result};
36use crate::format::{Header, FLAG_ENCRYPTED, PAGE_SIZE, PAYLOAD, VERSION, WRAPPED_DEK_LEN};
37use crate::index::{IndexConfig, IndexInfo, IvfPqIndex};
38use crate::memory::{self, Memory, MemoryType, Metric, Scope, ScoreWeights};
39use crate::pager::Pager;
40use crate::result_cache::{
41 self, BatchState, CacheBudget, CacheDirectoryEntry, CacheEntry, CacheFlushPolicy, CachePutOpts,
42 CacheStats, CachedValue, CounterState, DirectoryIndex, SemanticCachePutOpts,
43};
44use crate::wal;
45
46/// Default initial size of the write-ahead log region, in pages (64 KiB).
47/// The region grows automatically when a transaction's control plane outgrows
48/// it, so this is a *floor* on the fresh-file footprint, not a cap on
49/// transaction size. Was 64 pages (512 KiB) in v0.1.0 — lowered because that
50/// reservation dominated small-file size (~62% of a 31-memory dogfood file).
51const DEFAULT_WAL_PAGES: u64 = 8;
52
53/// Hard floor on the WAL reservation. A single header-page flush already
54/// needs one full DATA frame plus a COMMIT frame; two pages are the bare
55/// minimum that lets the first transaction commit without immediate growth.
56const MIN_WAL_PAGES: u64 = 2;
57
58/// Default cap on retained snapshot manifest entries. Each `flush` appends
59/// one entry; without a cap, long-running processes accumulate entries
60/// forever, growing per-flush manifest-serialize cost as O(total flushes)
61/// and the manifest run size on disk in lockstep. 256 keeps roughly a
62/// week of hourly snapshots while staying tiny (256 × ~82 bytes ≈ 21 KiB
63/// of manifest). Set `MnemoConfig::max_snapshots = 0` to disable the cap.
64pub const DEFAULT_MAX_SNAPSHOTS: usize = 256;
65
66/// Upper bound for the WAL scan when recovering a torn-header file. The
67/// real WAL region size lives in the header, but a torn header is exactly
68/// the case where we can't read it — so we scan a generous fixed window
69/// from the conventional WAL start, relying on `wal::recover` stopping at
70/// the first non-magic frame. 64 covers the historical pre-0.2 default
71/// of a 64-page reservation; users who configure larger WALs and then
72/// suffer a torn header at the same time will need offline recovery.
73const TORN_HEADER_SCAN_PAGES: u64 = 64;
74
75/// Configuration for creating a new database.
76#[derive(Clone, Copy, Debug)]
77pub struct MnemoConfig {
78 /// Embedding dimensionality. Every stored vector must match this.
79 pub dimensions: usize,
80 /// Key-derivation parameters.
81 pub kdf: KdfParams,
82 /// Initial size of the WAL region in 8 KiB pages. Default is
83 /// [`DEFAULT_WAL_PAGES`] (8 pages = 64 KiB), which keeps fresh files
84 /// small; the region auto-grows when a transaction's control plane needs
85 /// more space, so raising this is only a hint for the expected steady-state
86 /// transaction size. Write-heavy databases that routinely flush large
87 /// catalogs or ANN indexes can set this higher (e.g. 64 = 512 KiB) to
88 /// avoid early grow events; tiny embedded uses can leave it at the
89 /// default. Clamped to a minimum of [`MIN_WAL_PAGES`] (2 pages).
90 pub wal_pages_initial: u64,
91 /// Maximum number of snapshot manifest entries to retain. Each `flush`
92 /// appends one entry; once the cap is hit, the *oldest* entry is dropped
93 /// from the in-memory manifest before the new one is added, so the
94 /// retained set is always the most-recent N. Defaults to
95 /// [`DEFAULT_MAX_SNAPSHOTS`] (256). Set to `0` to disable the cap and
96 /// retain every snapshot forever (the pre-v0.3 behavior).
97 ///
98 /// Pages referenced *only* by pruned snapshots stay on disk until the
99 /// next [`Mnemo::compact_file`], which reclaims them. Point-in-time
100 /// recovery via [`Mnemo::restore_to`] into a pruned `txn_id` returns
101 /// [`MnemoError::NotFound`].
102 ///
103 /// Apply to an already-open database with [`Mnemo::set_max_snapshots`].
104 pub max_snapshots: usize,
105}
106
107impl Default for MnemoConfig {
108 fn default() -> Self {
109 Self {
110 dimensions: 768,
111 kdf: KdfParams::secure(),
112 wal_pages_initial: DEFAULT_WAL_PAGES,
113 max_snapshots: DEFAULT_MAX_SNAPSHOTS,
114 }
115 }
116}
117
118/// Catalog entry: maps a memory ID to its page run on disk.
119///
120/// **Schema (v5+):** carries `accessed_at` and `access_count` so `recall`
121/// can update access stats by mutating the catalog entry rather than
122/// rewriting the entire record (Phase 2.1 of the improvement plan). The
123/// values in a memory's serialized record body become a stale snapshot
124/// after the first recall; [`Mnemo::read_memory`] re-populates them from
125/// the catalog so consumers always see the live values.
126#[derive(Serialize, Deserialize, Clone, Debug)]
127struct CatalogEntry {
128 /// ULID as a raw `u128`.
129 id: u128,
130 start_page: u64,
131 page_count: u32,
132 /// Exact serialized byte length of the record.
133 len: u32,
134 deleted: bool,
135 /// Unix seconds of the most recent recall hit (or write, for fresh entries).
136 accessed_at: i64,
137 /// How many times `recall` has surfaced this memory.
138 access_count: u32,
139}
140
141/// Frozen v4 catalog entry shape, used only by the v4→v5 migration path in
142/// [`Mnemo::open`]. Its layout matches the catalog encoding written by
143/// pre-v5 builds (5 positional fields). Do not change this struct.
144#[derive(Serialize, Deserialize, Clone, Debug)]
145struct CatalogEntryV4 {
146 id: u128,
147 start_page: u64,
148 page_count: u32,
149 len: u32,
150 deleted: bool,
151}
152
153/// Returned by [`Mnemo::prepare_for_flush`] and consumed by [`Mnemo::flush`].
154/// Carries the already-serialized control-plane bytes through the prelude
155/// so we don't re-serialize after the lease has been persisted.
156struct FlushPrelude {
157 cat_bytes: Option<Vec<u8>>,
158 idx_bytes: Option<Vec<u8>>,
159 /// Serialized result-cache directory bytes when `dirty_cache` is
160 /// set (Phase 10.1). `None` if the cache didn't change since the
161 /// last flush — in that case the header keeps its current
162 /// `cache_start` / `cache_pages` / `cache_len` and no fresh cache
163 /// run is written.
164 cache_bytes: Option<Vec<u8>>,
165}
166
167/// One entry in the append-only snapshot manifest. Because record, catalog,
168/// and index pages are only ever *appended*, the runs a past flush wrote are
169/// still on disk; a `Snapshot` is the set of pointers needed to reconstruct
170/// the database exactly as that flush left it.
171#[derive(Serialize, Deserialize, Clone, Debug)]
172struct Snapshot {
173 txn_id: u64,
174 created_at: i64,
175 catalog_start: u64,
176 catalog_pages: u64,
177 catalog_len: u64,
178 index_start: u64,
179 index_pages: u64,
180 index_len: u64,
181 memory_count: u64,
182}
183
184/// A restorable point in a database's history — see [`Mnemo::snapshots`].
185#[derive(Clone, Copy, Debug)]
186pub struct SnapshotInfo {
187 /// Id of the transaction that produced this snapshot (monotonic from 1).
188 pub txn_id: u64,
189 /// When the snapshot was committed (unix seconds).
190 pub created_at: i64,
191 /// Live memory count captured in the snapshot.
192 pub memory_count: u64,
193}
194
195/// A retrieval request for [`Mnemo::recall`].
196#[derive(Clone, Debug)]
197pub struct RecallRequest {
198 /// Query embedding.
199 pub query: Vec<f32>,
200 /// Maximum results to return.
201 pub top_k: usize,
202 /// Restrict to these memory types (`None` = all types).
203 pub memory_types: Option<Vec<MemoryType>>,
204 /// Restrict to a single agent's view: its own memories plus shared ones.
205 pub agent_id: Option<String>,
206 /// Similarity metric.
207 pub metric: Metric,
208 /// Multi-signal score weights.
209 pub weights: ScoreWeights,
210 /// Index override: partitions to probe (`None` = index default). Ignored
211 /// when no ANN index is present.
212 pub n_probe: Option<usize>,
213 /// Index override: candidates to rerank (`None` = index default). Ignored
214 /// when no ANN index is present.
215 pub n_rerank: Option<usize>,
216 /// Whether to update each result's `accessed_at` and `access_count` in
217 /// the catalog. Defaults to `true`, matching pre-v5 behavior.
218 ///
219 /// Set to `false` for **fully read-only recall** — useful for batch
220 /// scoring, dry-run scoring, or recalls run by introspection tooling
221 /// (e.g. `mnemo about` consumers) where you don't want the score's own
222 /// observation to perturb the database. With `false`, recall does not
223 /// dirty the catalog and the next `flush` is a no-op.
224 pub track_access: bool,
225}
226
227impl RecallRequest {
228 /// A request with sensible defaults for a query embedding.
229 pub fn new(query: Vec<f32>) -> Self {
230 Self {
231 query,
232 top_k: 10,
233 memory_types: None,
234 agent_id: None,
235 metric: Metric::Cosine,
236 weights: ScoreWeights::default(),
237 n_probe: None,
238 n_rerank: None,
239 track_access: true,
240 }
241 }
242
243 /// Build a request by embedding `query_text` with the supplied
244 /// embedder (Phase 3). Shorthand for `embedder.embed(&[text])` +
245 /// [`RecallRequest::new`]; same defaults as `new`. Fails if the
246 /// embedder returns an empty vector list or an embedding failure.
247 pub fn from_text(embedder: &dyn Embedder, query_text: &str) -> Result<Self> {
248 let mut vectors = embedder.embed(&[query_text])?;
249 let v = vectors.pop().ok_or_else(|| {
250 MnemoError::Invalid("embedder returned no vector for the query text".into())
251 })?;
252 Ok(Self::new(v))
253 }
254 /// Set the result cap.
255 pub fn top_k(mut self, k: usize) -> Self {
256 self.top_k = k;
257 self
258 }
259 /// Restrict to specific memory types.
260 pub fn types(mut self, t: Vec<MemoryType>) -> Self {
261 self.memory_types = Some(t);
262 self
263 }
264 /// Restrict to one agent's view.
265 pub fn agent(mut self, agent_id: impl Into<String>) -> Self {
266 self.agent_id = Some(agent_id.into());
267 self
268 }
269 /// Set the similarity metric (default: cosine).
270 pub fn metric(mut self, metric: Metric) -> Self {
271 self.metric = metric;
272 self
273 }
274 /// Replace the multi-signal score weights.
275 pub fn weights(mut self, weights: ScoreWeights) -> Self {
276 self.weights = weights;
277 self
278 }
279 /// Override the number of IVF partitions probed (accuracy/speed dial).
280 pub fn n_probe(mut self, n: usize) -> Self {
281 self.n_probe = Some(n);
282 self
283 }
284 /// Override the number of candidates reranked exactly (accuracy dial).
285 pub fn n_rerank(mut self, n: usize) -> Self {
286 self.n_rerank = Some(n);
287 self
288 }
289 /// Toggle whether recall updates `accessed_at` / `access_count` on
290 /// the returned memories' catalog entries (default `true`). Pass
291 /// `false` for a fully read-only recall — see the field doc.
292 pub fn track_access(mut self, track: bool) -> Self {
293 self.track_access = track;
294 self
295 }
296}
297
298/// One scored result from [`Mnemo::recall`].
299#[derive(Clone, Debug)]
300pub struct RecallResult {
301 /// The retrieved memory.
302 pub memory: Memory,
303 /// Combined multi-signal score.
304 pub score: f32,
305 /// The bare similarity component (before other signals).
306 pub similarity: f32,
307}
308
309/// Summary statistics for a database.
310#[derive(Clone, Debug)]
311pub struct Stats {
312 /// Live (non-deleted) memory count.
313 pub memories: usize,
314 /// Tombstoned entries awaiting compaction.
315 pub deleted: usize,
316 /// Embedding dimensionality.
317 pub dimensions: usize,
318 /// File size in bytes.
319 pub file_bytes: u64,
320 /// Distinct agent IDs present.
321 pub agents: Vec<String>,
322 /// Whether pages are encrypted (always true in v1).
323 pub encrypted: bool,
324 /// Creation time (unix seconds).
325 pub created_at: i64,
326 /// ANN index shape, if an index has been built.
327 pub index: Option<IndexInfo>,
328 /// Current size of the write-ahead log region, in 8 KiB pages.
329 pub wal_pages: u64,
330 /// Fsyncs issued by the most recent successful [`Mnemo::flush`] on
331 /// this handle. `0` on a fresh handle before any flush, or when
332 /// `flush()` short-circuited (nothing dirty). Phase 6 instrumentation
333 /// for the Phase 11.1 write-latency work — the planner wants this
334 /// before deciding whether cutting fsync count is a viable
335 /// durability-preserving fix.
336 pub fsyncs_last_flush: u32,
337}
338
339/// Result of a [`Mnemo::compact_file`] run.
340#[derive(Clone, Copy, Debug)]
341pub struct CompactReport {
342 /// Live memories before compaction.
343 pub before: usize,
344 /// Live memories after compaction (expired ones dropped).
345 pub after: usize,
346}
347
348/// An encrypted, single-file agent memory database.
349pub struct Mnemo {
350 pager: Pager,
351 header: Header,
352 catalog: Vec<CatalogEntry>,
353 index: HashMap<u128, usize>,
354 #[allow(dead_code)]
355 path: PathBuf,
356 dimensions: usize,
357 kdf: KdfParams,
358 /// Set whenever the catalog changes; drives whether `flush` rewrites it.
359 dirty_catalog: bool,
360 /// Optional IVF+PQ approximate-nearest-neighbour index.
361 ann: Option<IvfPqIndex>,
362 /// Set whenever the ANN index changes; drives whether `flush` rewrites it.
363 dirty_index: bool,
364 /// Append-only manifest of committed snapshots, oldest first. Capped
365 /// at [`Mnemo::max_snapshots`] entries on every flush — the oldest
366 /// entries are pruned first.
367 manifest: Vec<Snapshot>,
368 /// Maximum manifest entries to retain across flushes. `0` disables
369 /// the cap (retain forever). Defaults to [`DEFAULT_MAX_SNAPSHOTS`].
370 max_snapshots: usize,
371 /// If `true`, this handle was opened via [`Mnemo::open_read_only`] and
372 /// every mutating API returns [`MnemoError::ReadOnly`]. The underlying
373 /// file was opened without write permission and holds an OS shared
374 /// lock; the OS would reject a write() anyway, but gating in Rust
375 /// gives callers a clean, typed error.
376 read_only: bool,
377 /// Result-cache directory (Phase 10.1). See [`crate::result_cache`].
378 /// Loaded from `header.cache_start` on open; persisted on `flush`
379 /// via the same seal_run pipeline as catalog/index/manifest.
380 cache_dir: Vec<CacheDirectoryEntry>,
381 /// In-memory index over `cache_dir` for O(1) `(namespace, key_hash)`
382 /// lookups. Rebuilt on open + kept in sync on `cache_put`/`delete`.
383 cache_idx: DirectoryIndex,
384 /// Set whenever `cache_dir` changes; drives whether `flush` rewrites
385 /// the cache directory run.
386 dirty_cache: bool,
387 /// Per-namespace eviction budgets. Missing entries fall back to
388 /// [`CacheBudget::default`]. Not persisted (per-handle policy).
389 cache_budgets: std::collections::HashMap<String, CacheBudget>,
390 /// Flush policy for cache mutations. Memory writes ignore this and
391 /// always behave as `Strict`.
392 cache_flush_policy: CacheFlushPolicy,
393 /// Bookkeeping for `CacheFlushPolicy::Batched` — pending mutation
394 /// count and first-dirty timestamp.
395 cache_batch: BatchState,
396 /// In-memory hit/miss/eviction counters (reset on reopen).
397 cache_counters: CounterState,
398 /// Optional handle-scoped embedder (Phase 3, PR 10a). Registered
399 /// via [`Mnemo::set_embedder`]; consulted by [`Mnemo::remember_text`]
400 /// and [`RecallRequest::from_text`]. `None` means text-side APIs
401 /// return [`MnemoError::Invalid`] with a "no embedder" message.
402 /// Not persisted — the handle is re-registered after every open.
403 embedder: Option<std::sync::Arc<dyn Embedder>>,
404 /// Fsyncs issued by the most recent `flush()` call. Snapshotted
405 /// from the thread-local fsync counter in `lib.rs::fsync_counter`:
406 /// `flush()` reads it on entry, does its work (each `sync_all` the
407 /// pager / wal / recovery path makes bumps the counter), and
408 /// stores the delta here. `0` on a fresh handle before any flush,
409 /// or on a read-only handle where `flush()` short-circuits.
410 /// Exposed through [`Stats::fsyncs_last_flush`] for diagnostics —
411 /// the Phase 6 planner wants this before deciding whether cutting
412 /// fsync count is a viable write-latency win.
413 fsyncs_last_flush: u32,
414}
415
416/// Open a file and take an OS advisory lock, returning `Locked { path }`
417/// if the requested lock cannot be acquired. On drop of the returned
418/// [`File`], the OS releases the lock. `flock`/`LockFileEx` semantics
419/// (per file-descriptor on Unix, per-open-file-handle on Windows) mean
420/// the lock is scoped to this handle: cloning the file (or the returned
421/// `File`) shares the lock, but a fresh `open()` from any process gets a
422/// distinct handle and will collide.
423fn open_with_lock(path: &Path, write: bool, create: bool) -> Result<File> {
424 let mut opts = OpenOptions::new();
425 opts.read(true);
426 if write {
427 opts.write(true);
428 }
429 if create {
430 opts.create(true).truncate(true);
431 }
432 let file = opts.open(path)?;
433 // On success the returned `Result<()>` is `Ok(())`. On conflict:
434 // * Unix `flock(2)` → `ErrorKind::WouldBlock` (EWOULDBLOCK / EAGAIN)
435 // * Windows `LockFileEx` → raw OS code 33 (ERROR_LOCK_VIOLATION)
436 // or 32 (ERROR_SHARING_VIOLATION), both surfaced by Rust as
437 // `ErrorKind::Uncategorized` because `std::io::ErrorKind`
438 // doesn't carry a portable "locked" variant.
439 // Anything else is a real I/O failure (permissions, invalid handle,
440 // ...) and gets bubbled up as `MnemoError::Io`.
441 let result = if write {
442 FileExt::try_lock_exclusive(&file)
443 } else {
444 FileExt::try_lock_shared(&file)
445 };
446 match result {
447 Ok(()) => Ok(file),
448 Err(e) if is_lock_conflict(&e) => Err(MnemoError::Locked {
449 path: path.to_path_buf(),
450 }),
451 Err(e) => Err(MnemoError::Io(e)),
452 }
453}
454
455/// Was this I/O error the OS reporting a conflicting file lock, or a
456/// genuine I/O failure? Platform-conditional because
457/// [`std::io::ErrorKind`] carries no portable "resource-locked" variant
458/// as of Rust 1.75 — Windows surfaces `ERROR_LOCK_VIOLATION` /
459/// `ERROR_SHARING_VIOLATION` as `Uncategorized`.
460fn is_lock_conflict(e: &std::io::Error) -> bool {
461 if e.kind() == std::io::ErrorKind::WouldBlock {
462 return true;
463 }
464 #[cfg(windows)]
465 if let Some(code) = e.raw_os_error() {
466 // 33 = ERROR_LOCK_VIOLATION, 32 = ERROR_SHARING_VIOLATION —
467 // both mean "someone else holds an incompatible lock on this
468 // region of the file". See
469 // https://learn.microsoft.com/en-us/windows/win32/debug/system-error-codes--0-499-
470 if code == 33 || code == 32 {
471 return true;
472 }
473 }
474 false
475}
476
477/// Read a run of consecutive encrypted pages and concatenate their plaintext.
478fn read_run_bytes(pager: &mut Pager, start: u64, pages: u64, len: u64) -> Result<Vec<u8>> {
479 let mut buf = Vec::with_capacity(len as usize);
480 for i in 0..pages {
481 buf.extend_from_slice(&pager.read_page(start + i)?);
482 }
483 buf.truncate(len as usize);
484 Ok(buf)
485}
486
487/// Load a v5+ catalog run (an empty `pages` yields an empty catalog).
488fn load_catalog(pager: &mut Pager, start: u64, pages: u64, len: u64) -> Result<Vec<CatalogEntry>> {
489 if pages == 0 {
490 return Ok(Vec::new());
491 }
492 let buf = read_run_bytes(pager, start, pages, len)?;
493 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
494}
495
496/// Load a **v4** catalog run using the frozen [`CatalogEntryV4`] shape.
497/// Used only by the v4→v5 migration path in [`Mnemo::open`].
498fn load_catalog_v4(
499 pager: &mut Pager,
500 start: u64,
501 pages: u64,
502 len: u64,
503) -> Result<Vec<CatalogEntryV4>> {
504 if pages == 0 {
505 return Ok(Vec::new());
506 }
507 let buf = read_run_bytes(pager, start, pages, len)?;
508 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
509}
510
511/// Walk a v4 catalog and synthesize the v5 shape. For each live entry,
512/// reads the memory's serialized body and copies its `accessed_at` and
513/// `access_count` into the catalog row (the v4 record body still carries
514/// them — they're moving *from* the body *to* the catalog). Deleted
515/// entries skip the read and zero the fields.
516fn migrate_v4_catalog(pager: &mut Pager, v4: Vec<CatalogEntryV4>) -> Result<Vec<CatalogEntry>> {
517 let mut out = Vec::with_capacity(v4.len());
518 for e in &v4 {
519 let (accessed_at, access_count) = if e.deleted {
520 (0i64, 0u32)
521 } else {
522 let mut buf = Vec::with_capacity(e.len as usize);
523 for i in 0..e.page_count as u64 {
524 buf.extend_from_slice(&pager.read_page(e.start_page + i)?);
525 }
526 buf.truncate(e.len as usize);
527 let mem: Memory = rmp_serde::from_slice(&buf)
528 .map_err(|err| MnemoError::Serialize(err.to_string()))?;
529 (mem.accessed_at, mem.access_count)
530 };
531 out.push(CatalogEntry {
532 id: e.id,
533 start_page: e.start_page,
534 page_count: e.page_count,
535 len: e.len,
536 deleted: e.deleted,
537 accessed_at,
538 access_count,
539 });
540 }
541 Ok(out)
542}
543
544/// Load the snapshot manifest (an empty `pages` yields an empty manifest).
545fn load_manifest(pager: &mut Pager, start: u64, pages: u64, len: u64) -> Result<Vec<Snapshot>> {
546 if pages == 0 {
547 return Ok(Vec::new());
548 }
549 let buf = read_run_bytes(pager, start, pages, len)?;
550 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
551}
552
553/// Load the result-cache directory (v8+). An empty `pages` yields an
554/// empty directory, which is the natural state for v7 files migrated
555/// to v8 (no cache entries yet) and for freshly created v8 files.
556fn load_cache_directory(
557 pager: &mut Pager,
558 start: u64,
559 pages: u64,
560 len: u64,
561) -> Result<Vec<CacheDirectoryEntry>> {
562 if pages == 0 {
563 return Ok(Vec::new());
564 }
565 let buf = read_run_bytes(pager, start, pages, len)?;
566 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))
567}
568
569/// Load an ANN index run, validating its dimensionality.
570fn load_index(
571 pager: &mut Pager,
572 start: u64,
573 pages: u64,
574 len: u64,
575 dims: usize,
576) -> Result<Option<IvfPqIndex>> {
577 if pages == 0 {
578 return Ok(None);
579 }
580 let buf = read_run_bytes(pager, start, pages, len)?;
581 let mut idx: IvfPqIndex =
582 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))?;
583 if idx.dims() != dims {
584 return Err(MnemoError::Invalid(
585 "ANN index dimensionality does not match the database".into(),
586 ));
587 }
588 idx.rebuild_assignment();
589 Ok(Some(idx))
590}
591
592/// True if a memory's metadata marks it as the canonical onboarding manifest
593/// (`metadata.topic == "manifest"`). Used by [`Mnemo::about`] to hoist the
594/// manifest to the top of the briefing regardless of importance ordering.
595fn is_manifest(m: &Memory) -> bool {
596 m.metadata
597 .get("topic")
598 .and_then(|v| v.as_str())
599 .map(|s| s.eq_ignore_ascii_case("manifest"))
600 .unwrap_or(false)
601}
602
603/// Build the ULID → catalog-slot lookup map.
604fn build_id_index(catalog: &[CatalogEntry]) -> HashMap<u128, usize> {
605 let mut m = HashMap::with_capacity(catalog.len());
606 for (i, e) in catalog.iter().enumerate() {
607 m.insert(e.id, i);
608 }
609 m
610}
611
612impl Mnemo {
613 /// Create a brand-new encrypted database at `path`.
614 pub fn create(
615 path: impl Into<PathBuf>,
616 passphrase: &str,
617 config: MnemoConfig,
618 ) -> Result<Mnemo> {
619 let path: PathBuf = path.into();
620 if config.dimensions == 0 {
621 return Err(MnemoError::Invalid("dimensions must be > 0".into()));
622 }
623 // Exclusive lock covers the write path; second writer sees `Locked`.
624 let file = open_with_lock(&path, /*write=*/ true, /*create=*/ true)?;
625
626 let salt = crypto::random_salt();
627 let dek = crypto::random_dek();
628 let kek = crypto::derive_kek(passphrase.as_bytes(), &salt, config.kdf)?;
629 let dek_nonce = crypto::random_nonce();
630 let wrapped = crypto::wrap_dek(&kek, &dek_nonce, &dek)?;
631 if wrapped.len() != WRAPPED_DEK_LEN {
632 return Err(MnemoError::Crypto("unexpected wrapped DEK length".into()));
633 }
634 let mut wrapped_dek = [0u8; WRAPPED_DEK_LEN];
635 wrapped_dek.copy_from_slice(&wrapped);
636
637 let mut header = Header {
638 version: VERSION,
639 page_size: PAGE_SIZE as u32,
640 flags: FLAG_ENCRYPTED,
641 dimensions: config.dimensions as u32,
642 created_at: memory::now_secs(),
643 write_counter: 0,
644 // Page 0 is the header; pages 1..=wal_pages are the WAL;
645 // record/catalog/index pages start after it. Initial WAL size is
646 // chosen per-config (default 8 pages = 64 KiB) and clamped to a
647 // sane floor so the first transaction can always commit.
648 next_page: 1 + config.wal_pages_initial.max(MIN_WAL_PAGES),
649 catalog_start: 0,
650 catalog_pages: 0,
651 catalog_len: 0,
652 vector_count: 0,
653 m_cost: config.kdf.m_cost,
654 t_cost: config.kdf.t_cost,
655 p_cost: config.kdf.p_cost,
656 salt,
657 dek_nonce,
658 wrapped_dek,
659 index_start: 0,
660 index_pages: 0,
661 index_len: 0,
662 wal_start: 1,
663 wal_pages: config.wal_pages_initial.max(MIN_WAL_PAGES),
664 wal_seq: 0,
665 manifest_start: 0,
666 manifest_pages: 0,
667 manifest_len: 0,
668 // Regenerated per write by `apply_seal`; `None` until we seal.
669 seal_nonce: None,
670 seal_tag: None,
671 // v8+ result-cache pointer — zero-initialized on create,
672 // populated only when `cache_put` runs at least once and
673 // the next `flush` seals a directory run.
674 cache_start: 0,
675 cache_pages: 0,
676 cache_len: 0,
677 };
678
679 let mut pager = Pager::new(file, dek, 0, VERSION);
680 let mut page = header.to_page();
681 header.apply_seal(&mut page, pager.dek())?;
682 pager.write_raw(0, &page)?;
683 pager.sync()?;
684
685 Ok(Mnemo {
686 pager,
687 header,
688 catalog: Vec::new(),
689 index: HashMap::new(),
690 path,
691 dimensions: config.dimensions,
692 kdf: config.kdf,
693 dirty_catalog: false,
694 ann: None,
695 dirty_index: false,
696 manifest: Vec::new(),
697 max_snapshots: config.max_snapshots,
698 read_only: false,
699 cache_dir: Vec::new(),
700 cache_idx: DirectoryIndex::new(),
701 dirty_cache: false,
702 cache_budgets: std::collections::HashMap::new(),
703 cache_flush_policy: CacheFlushPolicy::default(),
704 cache_batch: BatchState::default(),
705 cache_counters: CounterState::default(),
706 embedder: None,
707 fsyncs_last_flush: 0,
708 })
709 }
710
711 /// Open an existing database. A wrong passphrase fails cleanly with
712 /// [`MnemoError::WrongPassphrase`]. Takes an OS exclusive lock for
713 /// the lifetime of the returned handle; a second concurrent open
714 /// (from any process) returns [`MnemoError::Locked`].
715 ///
716 /// For a caller that only needs to read (and can tolerate the
717 /// pending-recovery / pending-migration cases documented on
718 /// [`Mnemo::open_read_only`]), prefer that entry point — it takes a
719 /// shared lock and coexists with other readers.
720 pub fn open(path: impl Into<PathBuf>, passphrase: &str) -> Result<Mnemo> {
721 let path: PathBuf = path.into();
722 let mut file = open_with_lock(&path, /*write=*/ true, /*create=*/ false)?;
723
724 let mut hbuf = [0u8; PAGE_SIZE];
725 file.seek(SeekFrom::Start(0))?;
726 file.read_exact(&mut hbuf)?;
727 let mut header = match Header::from_page(&hbuf) {
728 Ok(h) => h,
729 Err(MnemoError::HeaderChecksum) => {
730 // Page 0 is torn — most likely a crash during a checkpoint's
731 // header write. Try to heal from the WAL at its default site
732 // (a never-grown WAL never moves from page 1). We don't know
733 // the WAL size (the header that holds it is what's torn), so
734 // scan a generous fixed window — `wal::recover` stops at the
735 // first non-magic frame, so reading extra is safe.
736 let healed = wal::recover(&mut file, 1, TORN_HEADER_SCAN_PAGES, 0)?;
737 let frames = healed.ok_or(MnemoError::HeaderChecksum)?;
738 for (page_no, bytes) in &frames {
739 if bytes.len() != PAGE_SIZE {
740 return Err(MnemoError::Invalid("WAL frame is not page-sized".into()));
741 }
742 file.seek(SeekFrom::Start(page_no * PAGE_SIZE as u64))?;
743 file.write_all(bytes)?;
744 }
745 file.sync_all()?;
746 crate::bump_fsync();
747 file.seek(SeekFrom::Start(0))?;
748 file.read_exact(&mut hbuf)?;
749 Header::from_page(&hbuf)?
750 }
751 Err(e) => return Err(e),
752 };
753
754 // Crash recovery: replay a committed-but-uncheckpointed transaction.
755 // Each frame is a finished page image; the header frame, replayed to
756 // page 0, supersedes the header just read.
757 if let Some(frames) = wal::recover(
758 &mut file,
759 header.wal_start,
760 header.wal_pages,
761 header.wal_seq,
762 )? {
763 for (page_no, bytes) in &frames {
764 if bytes.len() != PAGE_SIZE {
765 return Err(MnemoError::Invalid("WAL frame is not page-sized".into()));
766 }
767 file.seek(SeekFrom::Start(page_no * PAGE_SIZE as u64))?;
768 file.write_all(bytes)?;
769 }
770 file.sync_all()?;
771 crate::bump_fsync();
772 // Re-read the now-current header from the replayed page 0.
773 file.seek(SeekFrom::Start(0))?;
774 file.read_exact(&mut hbuf)?;
775 header = Header::from_page(&hbuf)?;
776 }
777
778 let kdf = KdfParams {
779 m_cost: header.m_cost,
780 t_cost: header.t_cost,
781 p_cost: header.p_cost,
782 };
783 let kek = crypto::derive_kek(passphrase.as_bytes(), &header.salt, kdf)?;
784 let dek = crypto::unwrap_dek(&kek, &header.dek_nonce, &header.wrapped_dek)?;
785
786 // Validate the v7 header seal now that we have the DEK. Returns
787 // `MnemoError::HeaderTampered` if any of the sealed mutable fields
788 // were rewritten after the last legitimate flush. No-op for v6 and
789 // older — pre-v7 files don't carry a seal.
790 header.validate_seal(&dek)?;
791
792 // Build the pager at the file's *current* version so initial reads
793 // use the matching AAD scheme — `Pager::page_aad` returns empty AAD
794 // for v4/v5 (no page binding) and `page_no.to_le_bytes()` for v6+.
795 let on_disk_version = header.version;
796 let mut pager = Pager::new(file, dek, header.write_counter, on_disk_version);
797 let dimensions = header.dimensions as usize;
798
799 // === v4→v5 catalog migration =====================================
800 // v6 reads directly; v5 reads directly; v4 replays the catalog into
801 // the v5 shape by reading each memory's body to populate the new
802 // `accessed_at` / `access_count` fields. All page reads here still
803 // use the v4/v5 AAD scheme (none).
804 let (catalog, migrated_from_v4) = if on_disk_version >= 5 {
805 (
806 load_catalog(
807 &mut pager,
808 header.catalog_start,
809 header.catalog_pages,
810 header.catalog_len,
811 )?,
812 false,
813 )
814 } else if on_disk_version == 4 {
815 let v4 = load_catalog_v4(
816 &mut pager,
817 header.catalog_start,
818 header.catalog_pages,
819 header.catalog_len,
820 )?;
821 (migrate_v4_catalog(&mut pager, v4)?, true)
822 } else {
823 return Err(MnemoError::UnsupportedVersion(on_disk_version));
824 };
825 let index = build_id_index(&catalog);
826 let ann = load_index(
827 &mut pager,
828 header.index_start,
829 header.index_pages,
830 header.index_len,
831 dimensions,
832 )?;
833 let mut manifest = load_manifest(
834 &mut pager,
835 header.manifest_start,
836 header.manifest_pages,
837 header.manifest_len,
838 )?;
839
840 // === v5→v6 page-encryption migration =============================
841 // Pre-v6 data pages were encrypted under no AAD; v6 binds the home
842 // `page_no` as AAD so a page-transplant attack fails authentication
843 // on read. Re-stage every live record page through `write_page` so
844 // the next flush re-encrypts it under v6 AAD at the same home slot.
845 // Tombstoned entries skip — they're reclaimed by `compact_file`.
846 // Old catalog / ANN / manifest pages stay as orphan ciphertext;
847 // the next flush writes the new runs at fresh slots.
848 let migrated_pages_to_v6 = on_disk_version < 6;
849 if migrated_pages_to_v6 {
850 for e in &catalog {
851 if e.deleted {
852 continue;
853 }
854 for i in 0..e.page_count as u64 {
855 let payload = pager.read_page(e.start_page + i)?;
856 pager.write_page(e.start_page + i, &payload)?;
857 }
858 }
859 pager.set_version(6);
860 }
861
862 // === v6→v7 header-seal migration =================================
863 // v7 adds an AES-GCM seal at the tail of the header page that
864 // authenticates every mutable header field. Pre-v7 files have no
865 // seal — we just bump the version in memory and mark `dirty_catalog`
866 // so the next flush writes a sealed v7 header. Pages are untouched
867 // (page crypto is unchanged from v6).
868 let migrated_header_to_v7 = on_disk_version < 7;
869 if migrated_header_to_v7 {
870 pager.set_version(7);
871 }
872
873 // === v7→v8 cache-directory migration =============================
874 // v8 adds three u64 header fields (cache_start/pages/len) for the
875 // result-cache directory (Phase 10.1). Pre-v8 files have those
876 // three bytes as zero (the on-disk region beyond the v7 seal is
877 // zero-initialized), which is the legal "empty cache directory"
878 // state — no data migration required. We bump the version in
879 // memory and mark dirty_catalog so the next flush writes a
880 // v8-sealed header with the new AAD shape covering the cache
881 // pointer. Pages are untouched.
882 let migrated_header_to_v8 = on_disk_version < 8;
883 if migrated_header_to_v8 {
884 pager.set_version(VERSION);
885 }
886
887 let migrated = migrated_from_v4
888 || migrated_pages_to_v6
889 || migrated_header_to_v7
890 || migrated_header_to_v8;
891 if migrated {
892 // Stamp the version forward in memory; the next flush persists
893 // it via the WAL-committed header frame (and the v7 seal).
894 header.version = VERSION;
895 // Snapshots written under an older format reference page runs
896 // encrypted under the old AAD scheme; this build can't decrypt
897 // them under the new pager. Drop the manifest so the next
898 // flush records a fresh snapshot of the migrated state.
899 // Point-in-time recovery into the pre-migration past is
900 // sacrificed for migration simplicity. Old pages remain on
901 // disk until `compact_file`. (For v6→v7 alone — which doesn't
902 // change page crypto — the old snapshots would technically
903 // still be readable, but dropping them uniformly keeps the
904 // migration policy simple.)
905 manifest.clear();
906 }
907
908 // The ANN index pages on disk were sealed under the *old* AAD scheme
909 // if we migrated pages to v6; force a rewrite at fresh v6-encrypted
910 // slots on the next flush so a future reopen doesn't try to read
911 // them under the new AAD and fail. (Catalog and manifest already
912 // get rewritten via dirty_catalog and the manifest.clear() above.)
913 let dirty_index = migrated_pages_to_v6 && ann.is_some();
914
915 // Load the cache directory (v8+ only). Pre-v8 files have
916 // cache_pages == 0 by header layout, so this trivially yields
917 // an empty directory without touching the disk.
918 let cache_dir = load_cache_directory(
919 &mut pager,
920 header.cache_start,
921 header.cache_pages,
922 header.cache_len,
923 )?;
924 let cache_idx = DirectoryIndex::build(&cache_dir);
925
926 Ok(Mnemo {
927 pager,
928 header,
929 catalog,
930 index,
931 path,
932 dimensions,
933 kdf,
934 dirty_catalog: migrated,
935 ann,
936 dirty_index,
937 manifest,
938 // `Mnemo::open` has no [`MnemoConfig`] to read; default the cap
939 // to [`DEFAULT_MAX_SNAPSHOTS`]. Override with
940 // [`Mnemo::set_max_snapshots`] right after open if you need
941 // unlimited retention or a different cap for this handle.
942 max_snapshots: DEFAULT_MAX_SNAPSHOTS,
943 read_only: false,
944 cache_dir,
945 cache_idx,
946 dirty_cache: false,
947 cache_budgets: std::collections::HashMap::new(),
948 cache_flush_policy: CacheFlushPolicy::default(),
949 cache_batch: BatchState::default(),
950 cache_counters: CounterState::default(),
951 embedder: None,
952 fsyncs_last_flush: 0,
953 })
954 }
955
956 /// Open an existing database in **read-only** mode. Takes an OS shared
957 /// lock; multiple read-only handles can coexist, but any attempt to
958 /// open the same file read-write elsewhere returns
959 /// [`MnemoError::Locked`]. All mutating methods on the returned handle
960 /// return [`MnemoError::ReadOnly`], including `recall` with
961 /// `track_access(true)`.
962 ///
963 /// Two conditions block a read-only open and return
964 /// [`MnemoError::NeedsWriteOpen`]:
965 ///
966 /// 1. The file has a committed-but-uncheckpointed WAL transaction that
967 /// would ordinarily be replayed on open. A read-only handle
968 /// cannot write the recovered pages back.
969 /// 2. The on-disk format version is older than the current
970 /// `format::VERSION`, requiring an on-disk migration.
971 ///
972 /// In either case, open the file read-write once (which recovers /
973 /// migrates in place), close the handle cleanly, and reopen
974 /// read-only. This is documented in the README's Durability section.
975 pub fn open_read_only(path: impl Into<PathBuf>, passphrase: &str) -> Result<Mnemo> {
976 let path: PathBuf = path.into();
977 let mut file = open_with_lock(&path, /*write=*/ false, /*create=*/ false)?;
978
979 let mut hbuf = [0u8; PAGE_SIZE];
980 file.seek(SeekFrom::Start(0))?;
981 file.read_exact(&mut hbuf)?;
982 // A torn header on a read-only open cannot be healed (that requires
983 // writing recovered frames back). Surface as `NeedsWriteOpen` so the
984 // caller knows to reopen read-write once.
985 let header = Header::from_page(&hbuf).map_err(|e| match e {
986 MnemoError::HeaderChecksum => MnemoError::NeedsWriteOpen {
987 reason: "torn header page needs recovery from WAL",
988 },
989 other => other,
990 })?;
991
992 // Refuse to migrate on read-only opens. v6→v7 is header-only
993 // (no page writes required), but even that bumps `version` on
994 // disk and re-seals the header — which needs write access. Keep
995 // the policy uniform across all sub-version bumps.
996 if header.version < VERSION {
997 return Err(MnemoError::NeedsWriteOpen {
998 reason: "pending format-version migration",
999 });
1000 }
1001
1002 // Refuse if there is a committed-but-uncheckpointed transaction
1003 // waiting for `wal::recover` to fold it into home pages.
1004 if wal::recover(
1005 &mut file,
1006 header.wal_start,
1007 header.wal_pages,
1008 header.wal_seq,
1009 )?
1010 .is_some()
1011 {
1012 return Err(MnemoError::NeedsWriteOpen {
1013 reason: "pending WAL recovery",
1014 });
1015 }
1016
1017 let kdf = KdfParams {
1018 m_cost: header.m_cost,
1019 t_cost: header.t_cost,
1020 p_cost: header.p_cost,
1021 };
1022 let kek = crypto::derive_kek(passphrase.as_bytes(), &header.salt, kdf)?;
1023 let dek = crypto::unwrap_dek(&kek, &header.dek_nonce, &header.wrapped_dek)?;
1024 header.validate_seal(&dek)?;
1025
1026 let mut pager = Pager::new(file, dek, header.write_counter, header.version);
1027 let dimensions = header.dimensions as usize;
1028
1029 let catalog = load_catalog(
1030 &mut pager,
1031 header.catalog_start,
1032 header.catalog_pages,
1033 header.catalog_len,
1034 )?;
1035 let index = build_id_index(&catalog);
1036 let ann = load_index(
1037 &mut pager,
1038 header.index_start,
1039 header.index_pages,
1040 header.index_len,
1041 dimensions,
1042 )?;
1043 let manifest = load_manifest(
1044 &mut pager,
1045 header.manifest_start,
1046 header.manifest_pages,
1047 header.manifest_len,
1048 )?;
1049 let cache_dir = load_cache_directory(
1050 &mut pager,
1051 header.cache_start,
1052 header.cache_pages,
1053 header.cache_len,
1054 )?;
1055 let cache_idx = DirectoryIndex::build(&cache_dir);
1056
1057 Ok(Mnemo {
1058 pager,
1059 header,
1060 catalog,
1061 index,
1062 path,
1063 dimensions,
1064 kdf,
1065 // No mutation is possible; leave both dirty flags false so
1066 // `close()` is a genuine no-op even if we later add an
1067 // auto-flush on drop.
1068 dirty_catalog: false,
1069 ann,
1070 dirty_index: false,
1071 manifest,
1072 max_snapshots: DEFAULT_MAX_SNAPSHOTS,
1073 read_only: true,
1074 cache_dir,
1075 cache_idx,
1076 dirty_cache: false,
1077 cache_budgets: std::collections::HashMap::new(),
1078 cache_flush_policy: CacheFlushPolicy::default(),
1079 cache_batch: BatchState::default(),
1080 cache_counters: CounterState::default(),
1081 embedder: None,
1082 fsyncs_last_flush: 0,
1083 })
1084 }
1085
1086 /// Override the manifest snapshot cap on this open handle. `0` disables
1087 /// the cap; any positive value keeps the most-recent `max` snapshots and
1088 /// drops the rest on the next flush. See [`MnemoConfig::max_snapshots`].
1089 pub fn set_max_snapshots(&mut self, max: usize) {
1090 self.max_snapshots = max;
1091 }
1092
1093 /// Embedding dimensionality this database expects.
1094 pub fn dimensions(&self) -> usize {
1095 self.dimensions
1096 }
1097
1098 /// Number of live (non-deleted) memories.
1099 pub fn len(&self) -> usize {
1100 self.catalog.iter().filter(|e| !e.deleted).count()
1101 }
1102
1103 /// True if the database holds no live memories.
1104 pub fn is_empty(&self) -> bool {
1105 self.len() == 0
1106 }
1107
1108 // --- internal page/record helpers -------------------------------------
1109
1110 fn write_record(&mut self, bytes: &[u8]) -> Result<(u64, u32)> {
1111 let pc = bytes.len().div_ceil(PAYLOAD).max(1);
1112 let start = self.header.next_page;
1113 self.header.next_page += pc as u64;
1114 for i in 0..pc {
1115 let lo = i * PAYLOAD;
1116 let hi = ((i + 1) * PAYLOAD).min(bytes.len());
1117 self.pager.write_page(start + i as u64, &bytes[lo..hi])?;
1118 }
1119 Ok((start, pc as u32))
1120 }
1121
1122 fn read_record(&mut self, e: &CatalogEntry) -> Result<Vec<u8>> {
1123 let mut buf = Vec::with_capacity(e.len as usize);
1124 for i in 0..e.page_count as u64 {
1125 buf.extend_from_slice(&self.pager.read_page(e.start_page + i)?);
1126 }
1127 buf.truncate(e.len as usize);
1128 Ok(buf)
1129 }
1130
1131 fn read_memory(&mut self, e: &CatalogEntry) -> Result<Memory> {
1132 let bytes = self.read_record(e)?;
1133 let mut m: Memory =
1134 rmp_serde::from_slice(&bytes).map_err(|err| MnemoError::Serialize(err.to_string()))?;
1135 // The catalog is the source of truth for access stats — the values
1136 // in the record body are a stale snapshot from when the record was
1137 // last written. Overwrite them here so consumers always see live
1138 // values regardless of when the record was last serialized.
1139 m.accessed_at = e.accessed_at;
1140 m.access_count = e.access_count;
1141 Ok(m)
1142 }
1143
1144 /// Serialize and store a memory (insert or overwrite by ID).
1145 fn put(&mut self, mut m: Memory) -> Result<Ulid> {
1146 if m.vector.len() != self.dimensions {
1147 return Err(MnemoError::DimensionMismatch {
1148 expected: self.dimensions,
1149 got: m.vector.len(),
1150 });
1151 }
1152 // Reject all-zero vectors (Phase 3 PR 10a; breaking vs 0.4.x).
1153 // A zero vector scores exactly 0.0 on every cosine query — it
1154 // can't contribute signal, but recency / importance can still
1155 // pull it into top-k as noise. We'd rather fail loudly at
1156 // `remember` than let a mis-embedded entry pollute recall
1157 // quality for the life of the database. Scaffold manifests now
1158 // use a sentinel unit vector so this check doesn't trip them.
1159 if embed::is_zero_vector(&m.vector) {
1160 return Err(MnemoError::Invalid(
1161 "all-zero vector refused: it scores 0.0 on every cosine query and \
1162 pollutes recall; check your embedder returned a real embedding"
1163 .into(),
1164 ));
1165 }
1166 if m.id == Ulid::nil() {
1167 m.id = Ulid::new();
1168 }
1169 let id_u: u128 = m.id.0;
1170 let bytes = rmp_serde::to_vec(&m).map_err(|e| MnemoError::Serialize(e.to_string()))?;
1171 let (start, pc) = self.write_record(&bytes)?;
1172 // For overwrites preserve the existing access stats; for inserts
1173 // seed them from the memory (which carries them in its body).
1174 let (prev_accessed, prev_count) = match self.index.get(&id_u).copied() {
1175 Some(idx) => (
1176 self.catalog[idx].accessed_at,
1177 self.catalog[idx].access_count,
1178 ),
1179 None => (m.accessed_at, m.access_count),
1180 };
1181 let entry = CatalogEntry {
1182 id: id_u,
1183 start_page: start,
1184 page_count: pc,
1185 len: bytes.len() as u32,
1186 deleted: false,
1187 accessed_at: prev_accessed,
1188 access_count: prev_count,
1189 };
1190 match self.index.get(&id_u).copied() {
1191 Some(idx) => self.catalog[idx] = entry,
1192 None => {
1193 self.index.insert(id_u, self.catalog.len());
1194 self.catalog.push(entry);
1195 }
1196 }
1197 self.dirty_catalog = true;
1198
1199 // Keep the ANN index complete: assign the (new or changed) vector to
1200 // its nearest partition. Centroids/codebook stay fixed until rebuild.
1201 if let Some(ann) = &mut self.ann {
1202 ann.add(id_u, &m.vector);
1203 self.dirty_index = true;
1204 }
1205 Ok(m.id)
1206 }
1207
1208 // --- public CRUD ------------------------------------------------------
1209
1210 /// Register a text embedder for this handle. Enables
1211 /// [`Mnemo::remember_text`] and the text-side `RecallRequest`
1212 /// convenience path. Does two checks before storing the embedder:
1213 ///
1214 /// 1. **Dimension check.** The embedder's `identity().dimensions`
1215 /// must equal [`Mnemo::dimensions`] (the width this database
1216 /// was created with). Mismatch returns
1217 /// [`MnemoError::DimensionMismatch`].
1218 /// 2. **Manifest identity check.** If the scaffold / onboarding
1219 /// manifest already records an embedder — the canonical shape
1220 /// is `metadata.embedder = { name, dimensions, normalize }`,
1221 /// with the legacy flat `embedder_name` / `embedder_dimensions`
1222 /// / `embedder_normalize` keys accepted as an alias — the new
1223 /// embedder's identity must match exactly. A different name —
1224 /// even at the same dimensions — returns
1225 /// [`MnemoError::Invalid`] with
1226 /// [`EmbedderIdentity::MISMATCH_MSG`] in the message. If the
1227 /// manifest has no embedder metadata yet, this call upserts
1228 /// the nested object into the first onboarding/manifest memory
1229 /// (write + flush so a subsequent reopen sees the identity).
1230 ///
1231 /// **Read-only handles.** `set_embedder` on an `open_read_only`
1232 /// handle performs both checks and adopts the embedder in memory;
1233 /// it never writes. Readers (CLI `--text`, agent-side `recall_text`)
1234 /// can therefore use the text path next to a writer that holds
1235 /// the exclusive lock.
1236 ///
1237 /// On both failure paths the handle's embedder slot is left
1238 /// unchanged. On success the embedder is wrapped in an `Arc<dyn
1239 /// Embedder>` and consulted by every text-side call made through
1240 /// this handle. The slot is per-handle and not persisted — a
1241 /// fresh open must re-register (same identity, so the manifest
1242 /// check round-trips cleanly).
1243 pub fn set_embedder<E: Embedder + 'static>(&mut self, embedder: E) -> Result<()> {
1244 let identity = embedder.identity();
1245 if identity.dimensions != self.dimensions {
1246 return Err(MnemoError::DimensionMismatch {
1247 expected: self.dimensions,
1248 got: identity.dimensions,
1249 });
1250 }
1251
1252 // Find a scaffold/onboarding manifest to inspect (or upsert
1253 // into). `about()` already filters + sorts the onboarding set.
1254 let onboarding = self.about()?;
1255 let mut upsert_target: Option<usize> = None;
1256 for (i, m) in onboarding.iter().enumerate() {
1257 if let Some(saw) = embed::read_identity_from_metadata(&m.metadata) {
1258 if saw != identity {
1259 return Err(embed::identity_mismatch_err(&saw, &identity));
1260 }
1261 // Match — nothing to persist; adopt + done. Works the
1262 // same on read-only and read-write handles.
1263 self.embedder = Some(std::sync::Arc::new(embedder));
1264 return Ok(());
1265 }
1266 if upsert_target.is_none() {
1267 upsert_target = Some(i);
1268 }
1269 }
1270
1271 // No existing identity on the manifest. On a read-write handle
1272 // upsert into the first onboarding memory so later opens see
1273 // it; on a read-only handle skip the write (readers must not
1274 // mutate the file), but still adopt the embedder in memory.
1275 if !self.read_only {
1276 if let Some(idx) = upsert_target {
1277 let mut m = onboarding[idx].clone();
1278 embed::write_identity_into_metadata(&mut m.metadata, &identity);
1279 // Reuse the memory's existing id so this is an in-place
1280 // update of the scaffold manifest, not a new entry.
1281 self.put(m)?;
1282 self.flush()?;
1283 }
1284 }
1285
1286 self.embedder = Some(std::sync::Arc::new(embedder));
1287 Ok(())
1288 }
1289
1290 /// Borrow the currently-registered embedder, if any. Returns
1291 /// `None` when no embedder has been set via
1292 /// [`Mnemo::set_embedder`]. Bindings use this to drive their own
1293 /// text-side convenience methods without duplicating the identity
1294 /// check.
1295 pub fn embedder(&self) -> Option<&dyn Embedder> {
1296 self.embedder.as_deref()
1297 }
1298
1299 /// Embed `content` with the registered embedder and `remember` the
1300 /// result. Returns the new memory's ULID. Convenience path for
1301 /// callers who have a string, want it stored, and don't need any
1302 /// of the per-memory options [`Memory`] exposes (agent, importance,
1303 /// metadata, ttl, scope) — build a [`Memory`] by hand and call
1304 /// [`Mnemo::remember`] for those.
1305 ///
1306 /// Fails with [`MnemoError::Invalid`] if no embedder has been
1307 /// registered on this handle (call [`Mnemo::set_embedder`] first);
1308 /// with [`MnemoError::DimensionMismatch`] if the embedder returned
1309 /// the wrong width for this database.
1310 pub fn remember_text(&mut self, content: &str, memory_type: MemoryType) -> Result<Ulid> {
1311 if self.read_only {
1312 return Err(MnemoError::ReadOnly);
1313 }
1314 let embedder = self.embedder.as_ref().ok_or_else(|| {
1315 MnemoError::Invalid(
1316 "no embedder registered on this handle; call set_embedder first".into(),
1317 )
1318 })?;
1319 let mut vectors = embedder.embed(&[content])?;
1320 let vector = vectors.pop().ok_or_else(|| {
1321 MnemoError::Invalid("embedder returned no vector for the input text".into())
1322 })?;
1323 if vector.len() != self.dimensions {
1324 return Err(MnemoError::DimensionMismatch {
1325 expected: self.dimensions,
1326 got: vector.len(),
1327 });
1328 }
1329 let memory = Memory::new(content.to_string(), memory_type, vector);
1330 self.remember(memory)
1331 }
1332
1333 /// Store a memory. Returns its ULID. If the memory's ID is nil a fresh
1334 /// one is assigned; an existing ID overwrites in place.
1335 pub fn remember(&mut self, memory: Memory) -> Result<Ulid> {
1336 if self.read_only {
1337 return Err(MnemoError::ReadOnly);
1338 }
1339 self.put(memory)
1340 }
1341
1342 /// Fetch a memory by ID.
1343 pub fn get(&mut self, id: &Ulid) -> Result<Memory> {
1344 let idx = *self
1345 .index
1346 .get(&id.0)
1347 .ok_or_else(|| MnemoError::NotFound(id.to_string()))?;
1348 let entry = self.catalog[idx].clone();
1349 if entry.deleted {
1350 return Err(MnemoError::NotFound(id.to_string()));
1351 }
1352 self.read_memory(&entry)
1353 }
1354
1355 /// Soft-delete a memory (tombstoned; space reclaimed by `compact`).
1356 pub fn delete(&mut self, id: &Ulid) -> Result<()> {
1357 if self.read_only {
1358 return Err(MnemoError::ReadOnly);
1359 }
1360 let idx = *self
1361 .index
1362 .get(&id.0)
1363 .ok_or_else(|| MnemoError::NotFound(id.to_string()))?;
1364 if !self.catalog[idx].deleted {
1365 self.catalog[idx].deleted = true;
1366 self.dirty_catalog = true;
1367 }
1368 if let Some(ann) = &mut self.ann {
1369 ann.remove(id.0);
1370 self.dirty_index = true;
1371 }
1372 Ok(())
1373 }
1374
1375 /// Return every live memory (used by tooling and compaction).
1376 pub fn memories(&mut self) -> Result<Vec<Memory>> {
1377 let entries: Vec<CatalogEntry> = self
1378 .catalog
1379 .iter()
1380 .filter(|e| !e.deleted)
1381 .cloned()
1382 .collect();
1383 let mut out = Vec::with_capacity(entries.len());
1384 for e in &entries {
1385 out.push(self.read_memory(e)?);
1386 }
1387 Ok(out)
1388 }
1389
1390 /// Return the database's **self-describing onboarding memories** — the
1391 /// ones tagged `metadata.area = "onboarding"`. This is the engine-level
1392 /// surface for the *single-file philosophy*: an agent who receives a
1393 /// `.mnemo` file (and its passphrase) can call this to learn what the
1394 /// file is, which embedder it expects, the recommended agent id, and any
1395 /// other conventions the file's author chose to record — all without
1396 /// needing any external documentation.
1397 ///
1398 /// Ordering: the canonical manifest (tag `metadata.topic = "manifest"`)
1399 /// always comes first; everything else follows in `importance` descending,
1400 /// then `created_at` ascending for deterministic results.
1401 pub fn about(&mut self) -> Result<Vec<Memory>> {
1402 let mut out: Vec<Memory> = self
1403 .memories()?
1404 .into_iter()
1405 .filter(|m| {
1406 m.metadata
1407 .get("area")
1408 .and_then(|v| v.as_str())
1409 .map(|s| s.eq_ignore_ascii_case("onboarding"))
1410 .unwrap_or(false)
1411 })
1412 .collect();
1413 out.sort_by(|a, b| {
1414 let a_manifest = is_manifest(a);
1415 let b_manifest = is_manifest(b);
1416 // Manifest topic wins first; then importance desc; then created_at asc.
1417 b_manifest
1418 .cmp(&a_manifest)
1419 .then_with(|| b.importance.total_cmp(&a.importance))
1420 .then_with(|| a.created_at.cmp(&b.created_at))
1421 });
1422 Ok(out)
1423 }
1424
1425 // --- retrieval --------------------------------------------------------
1426
1427 /// Phase-1 brute-force search: rank live memories by raw similarity only.
1428 /// Read-only — does not update access statistics.
1429 pub fn search(
1430 &mut self,
1431 query: &[f32],
1432 top_k: usize,
1433 metric: Metric,
1434 ) -> Result<Vec<(Memory, f32)>> {
1435 if query.len() != self.dimensions {
1436 return Err(MnemoError::DimensionMismatch {
1437 expected: self.dimensions,
1438 got: query.len(),
1439 });
1440 }
1441 let now = memory::now_secs();
1442 let entries: Vec<CatalogEntry> = self
1443 .catalog
1444 .iter()
1445 .filter(|e| !e.deleted)
1446 .cloned()
1447 .collect();
1448
1449 let mut scored: Vec<(Memory, f32)> = Vec::new();
1450 for e in &entries {
1451 let m = self.read_memory(e)?;
1452 if m.is_expired(now) {
1453 continue;
1454 }
1455 let sim = memory::similarity(metric, query, &m.vector);
1456 scored.push((m, sim));
1457 }
1458 scored.sort_by(|a, b| b.1.total_cmp(&a.1));
1459 scored.truncate(top_k);
1460 Ok(scored)
1461 }
1462
1463 /// Phase-5 multi-signal recall: rank by
1464 /// `α·similarity + β·recency + γ·importance + δ·ln(freq)`, with type and
1465 /// agent filtering. Updates `accessed_at` / `access_count` on the
1466 /// returned memories (persisted on the next `flush`).
1467 ///
1468 /// When an ANN index has been built ([`Mnemo::build_index`]), recall runs
1469 /// the tiered IVF→PQ→rerank pipeline: it scores only the similarity-nearest
1470 /// `n_rerank` candidates rather than every memory. Without an index it
1471 /// scores every live memory exactly.
1472 pub fn recall(&mut self, req: &RecallRequest) -> Result<Vec<RecallResult>> {
1473 if req.query.len() != self.dimensions {
1474 return Err(MnemoError::DimensionMismatch {
1475 expected: self.dimensions,
1476 got: req.query.len(),
1477 });
1478 }
1479 // Access-tracking bumps `accessed_at` / `access_count` in the
1480 // catalog, which is a mutation. On a read-only handle the caller
1481 // must pass `track_access(false)` (via CLI `--no-track`).
1482 if self.read_only && req.track_access {
1483 return Err(MnemoError::ReadOnly);
1484 }
1485 let now = memory::now_secs();
1486
1487 // Candidate set: ANN-narrowed when an index exists, else everything.
1488 let entries: Vec<CatalogEntry> = if let Some(ann) = &self.ann {
1489 let ids = ann.query(&req.query, req.n_probe, req.n_rerank);
1490 ids.iter()
1491 .filter_map(|id| self.index.get(id).copied())
1492 .map(|i| self.catalog[i].clone())
1493 .filter(|e| !e.deleted)
1494 .collect()
1495 } else {
1496 self.catalog
1497 .iter()
1498 .filter(|e| !e.deleted)
1499 .cloned()
1500 .collect()
1501 };
1502
1503 let mut scored: Vec<RecallResult> = Vec::new();
1504 for e in &entries {
1505 let m = self.read_memory(e)?;
1506 if m.is_expired(now) {
1507 continue;
1508 }
1509 // Type filter.
1510 if let Some(types) = &req.memory_types {
1511 if !types.contains(&m.memory_type) {
1512 continue;
1513 }
1514 }
1515 // Agent-scoping: an agent sees its own memories plus shared ones.
1516 if let Some(agent) = &req.agent_id {
1517 let visible = m.agent_id == *agent || m.scope == Scope::Shared;
1518 if !visible {
1519 continue;
1520 }
1521 }
1522 let sim = memory::similarity(req.metric, &req.query, &m.vector);
1523 let age = (now - m.accessed_at) as f32;
1524 let score = req.weights.score(sim, age, m.importance, m.access_count);
1525 scored.push(RecallResult {
1526 memory: m,
1527 score,
1528 similarity: sim,
1529 });
1530 }
1531 scored.sort_by(|a, b| b.score.total_cmp(&a.score));
1532 scored.truncate(req.top_k);
1533
1534 // Update access statistics for everything we surfaced — but only
1535 // mutate the catalog entry, not the full record on disk. Pre-v5
1536 // this loop called `self.put(r.memory.clone())`, which serialized
1537 // and rewrote every result (vector included) to fresh pages,
1538 // turning a top-K recall into K full record rewrites. Now the
1539 // catalog is the source of truth for these two fields and the
1540 // record body stays untouched until the next real edit.
1541 if req.track_access {
1542 for r in &mut scored {
1543 if let Some(&idx) = self.index.get(&r.memory.id.0) {
1544 let entry = &mut self.catalog[idx];
1545 entry.accessed_at = now;
1546 entry.access_count = entry.access_count.saturating_add(1);
1547 // Propagate the just-updated values onto the returned
1548 // Memory so the caller sees the post-recall state.
1549 r.memory.accessed_at = entry.accessed_at;
1550 r.memory.access_count = entry.access_count;
1551 }
1552 }
1553 // Mark the catalog dirty so the next flush persists these
1554 // bumps. The records themselves are not dirty.
1555 if !scored.is_empty() {
1556 self.dirty_catalog = true;
1557 }
1558 }
1559 Ok(scored)
1560 }
1561
1562 // --- approximate index ------------------------------------------------
1563
1564 /// Build an IVF+PQ approximate-nearest-neighbour index over every live
1565 /// memory, using default tuning. After this, [`Mnemo::recall`] runs the
1566 /// tiered pipeline instead of an exact scan. Persisted on the next
1567 /// `flush`. Returns a snapshot of the index shape.
1568 pub fn build_index(&mut self) -> Result<IndexInfo> {
1569 self.build_index_with(IndexConfig::default())
1570 }
1571
1572 /// Build the ANN index with explicit tuning.
1573 pub fn build_index_with(&mut self, cfg: IndexConfig) -> Result<IndexInfo> {
1574 if self.read_only {
1575 return Err(MnemoError::ReadOnly);
1576 }
1577 let mems = self.memories()?;
1578 if mems.is_empty() {
1579 return Err(MnemoError::Invalid(
1580 "cannot build an index over an empty database".into(),
1581 ));
1582 }
1583 let items: Vec<(u128, &[f32])> =
1584 mems.iter().map(|m| (m.id.0, m.vector.as_slice())).collect();
1585 let idx = IvfPqIndex::build(self.dimensions, &items, cfg)?;
1586 let info = idx.info();
1587 self.ann = Some(idx);
1588 self.dirty_index = true;
1589 Ok(info)
1590 }
1591
1592 /// Rebuild the ANN index from scratch — re-clusters centroids and retrains
1593 /// the PQ codebook, undoing the cluster drift that accumulates as memories
1594 /// are inserted against fixed centroids.
1595 pub fn rebuild_index(&mut self) -> Result<IndexInfo> {
1596 let cfg = match &self.ann {
1597 Some(a) => IndexConfig {
1598 n_probe: a.n_probe(),
1599 n_rerank: a.n_rerank(),
1600 ..Default::default()
1601 },
1602 None => IndexConfig::default(),
1603 };
1604 self.build_index_with(cfg)
1605 }
1606
1607 /// Drop the ANN index; recall reverts to exact scans. Persisted on flush.
1608 /// No-op on a read-only handle (signature is infallible — callers of
1609 /// `build_index` / `flush` on a read-only handle already get a typed
1610 /// `ReadOnly` error, which is the enforceable boundary).
1611 pub fn drop_index(&mut self) {
1612 if self.read_only {
1613 return;
1614 }
1615 if self.ann.is_some() {
1616 self.ann = None;
1617 self.dirty_index = true;
1618 }
1619 }
1620
1621 /// Whether an ANN index is currently loaded.
1622 pub fn has_index(&self) -> bool {
1623 self.ann.is_some()
1624 }
1625
1626 // --- durability & maintenance ----------------------------------------
1627
1628 /// Seal a serialized buffer into a fresh run of encrypted home pages and
1629 /// append a WAL frame for each. Returns the run's `(start_page, pages)`.
1630 /// The pages are *not* written to their home locations here — they go to
1631 /// the WAL and are folded in later by [`Mnemo::checkpoint`].
1632 fn seal_run(&mut self, bytes: &[u8], frames: &mut Vec<wal::Frame>) -> Result<(u64, u32)> {
1633 let pc = bytes.len().div_ceil(PAYLOAD).max(1);
1634 let start = self.header.next_page;
1635 self.header.next_page += pc as u64;
1636 for i in 0..pc {
1637 let lo = i * PAYLOAD;
1638 let hi = ((i + 1) * PAYLOAD).min(bytes.len());
1639 let page_no = start + i as u64;
1640 let sealed = self.pager.seal_page(page_no, &bytes[lo..hi])?;
1641 frames.push((page_no, sealed.to_vec()));
1642 }
1643 Ok((start, pc as u32))
1644 }
1645
1646 /// Grow the WAL region if `frame_pages` worth of frames would not fit.
1647 /// Called at the top of `flush`, where the WAL is always spent (every
1648 /// flush ends with a checkpoint), so relocating it cannot strand a
1649 /// committed transaction.
1650 fn ensure_wal_capacity(&mut self, frame_pages: usize) -> Result<()> {
1651 let need = wal::txn_byte_len(frame_pages, PAGE_SIZE);
1652 if need <= self.header.wal_pages * PAGE_SIZE as u64 {
1653 return Ok(());
1654 }
1655 // Allocate a new, larger region at the file tail (~1.5x headroom).
1656 let req = need.div_ceil(PAGE_SIZE as u64);
1657 let new_pages = (req + req / 2 + 4).max(DEFAULT_WAL_PAGES);
1658 let new_start = self.header.next_page;
1659 self.header.next_page += new_pages;
1660 self.header.wal_start = new_start;
1661 self.header.wal_pages = new_pages;
1662 // Persist the relocation now: an isolated header write over an empty
1663 // WAL. A crash here leaves a consistent state with a bigger WAL.
1664 self.header.write_counter = self.pager.write_counter;
1665 let mut page = self.header.to_page();
1666 self.header.apply_seal(&mut page, self.pager.dek())?;
1667 self.pager.write_raw(0, &page)?;
1668 self.pager.sync()?;
1669 Ok(())
1670 }
1671
1672 /// Fold a committed transaction's WAL frames into their home pages.
1673 fn checkpoint(&mut self, frames: &[wal::Frame]) -> Result<()> {
1674 for (page_no, bytes) in frames {
1675 let mut img = [0u8; PAGE_SIZE];
1676 img.copy_from_slice(bytes);
1677 self.pager.write_sealed(*page_no, &img)?;
1678 }
1679 self.pager.sync()?;
1680 Ok(())
1681 }
1682
1683 /// Persist all pending changes as one **write-ahead-logged transaction**.
1684 ///
1685 /// Sequence:
1686 ///
1687 /// 1. **Prepare.** Serialize the control plane (catalog, ANN index),
1688 /// grow the WAL if needed, and **lease** counter and page slots —
1689 /// bump `header.write_counter` and `header.next_page` by upper
1690 /// bounds on this transaction's consumption and persist that
1691 /// *leased* header to disk before any encrypted page is written.
1692 /// 2. **Data pages.** [`Pager::flush`] writes dirty record pages
1693 /// copy-on-write under fresh nonces and fsyncs them. The orphan-page
1694 /// nonce-reuse window (see [Phase 1.1 of the improvement plan])
1695 /// is closed because the leased header on disk already records a
1696 /// `write_counter` past anything this step can produce.
1697 /// 3. **Control plane.** Seal the new catalog / ANN index / snapshot
1698 /// manifest into fresh home page runs.
1699 /// 4. **Commit.** Log all of step 3's frames plus the final header
1700 /// frame into the WAL and fsync — the durability point.
1701 /// 5. **Checkpoint.** Fold the WAL into the home pages.
1702 ///
1703 /// A crash before step 4 leaves the previous state intact. A crash
1704 /// after step 4 is repaired by [`Mnemo::open`] replaying the WAL.
1705 /// Safe to call repeatedly.
1706 pub fn flush(&mut self) -> Result<()> {
1707 if self.read_only {
1708 return Err(MnemoError::ReadOnly);
1709 }
1710 // Snapshot the thread-local fsync counter. Every `sync_all` the
1711 // pager / wal / checkpoint path makes between here and the end
1712 // of this function bumps the same counter; `delta` lands on
1713 // `self.fsyncs_last_flush` so `Stats::fsyncs_last_flush` can
1714 // report it. The planner wants this to decide whether cutting
1715 // fsync count is a viable write-latency win before we touch
1716 // the write path's durability model.
1717 let fsync_before = crate::fsync_count();
1718 if !self.dirty_catalog && !self.dirty_index && !self.dirty_cache {
1719 // No control-plane changes pending. Because every `put` that
1720 // dirties a data page also dirties the catalog, and every
1721 // `cache_put` that dirties a data page also dirties
1722 // `dirty_cache`, there can't be dirty data pages either —
1723 // nothing to do.
1724 self.fsyncs_last_flush = 0;
1725 return Ok(());
1726 }
1727
1728 // 1. Serialize control plane, ensure WAL capacity, lease counter +
1729 // page slots, and persist the leased header *before* any data
1730 // page hits the disk.
1731 let ctx = self.prepare_for_flush()?;
1732
1733 // 2. Record (vector) data pages: copy-on-write to fresh pages, fsynced.
1734 // Safe: the lease guarantees that even if we crash here, reopen
1735 // will see a `write_counter` past every nonce this step uses.
1736 self.pager.flush()?;
1737
1738 // 3. Seal the catalog / index control pages into fresh home runs.
1739 let mut frames: Vec<wal::Frame> = Vec::new();
1740 if let Some(bytes) = &ctx.cat_bytes {
1741 let (start, pages) = self.seal_run(bytes, &mut frames)?;
1742 self.header.catalog_start = start;
1743 self.header.catalog_pages = pages as u64;
1744 self.header.catalog_len = bytes.len() as u64;
1745 }
1746 if self.dirty_index {
1747 match &ctx.idx_bytes {
1748 Some(bytes) => {
1749 let (start, pages) = self.seal_run(bytes, &mut frames)?;
1750 self.header.index_start = start;
1751 self.header.index_pages = pages as u64;
1752 self.header.index_len = bytes.len() as u64;
1753 }
1754 None => {
1755 self.header.index_start = 0;
1756 self.header.index_pages = 0;
1757 self.header.index_len = 0;
1758 }
1759 }
1760 }
1761 // v8+ result-cache directory (Phase 10.1). Seal the serialized
1762 // directory into a fresh page run and update the pointer in the
1763 // header. The seal AAD already covers cache_start/pages/len for
1764 // v8+ (see `header_seal_aad`), so tampering with the pointer
1765 // trips the header seal on the next open.
1766 if self.dirty_cache {
1767 match &ctx.cache_bytes {
1768 Some(bytes) if !self.cache_dir.is_empty() => {
1769 let (start, pages) = self.seal_run(bytes, &mut frames)?;
1770 self.header.cache_start = start;
1771 self.header.cache_pages = pages as u64;
1772 self.header.cache_len = bytes.len() as u64;
1773 }
1774 _ => {
1775 // Either the directory serialized to nothing (all
1776 // entries deleted after tombstone-only migration) or
1777 // it's empty — collapse the pointer to zero.
1778 self.header.cache_start = 0;
1779 self.header.cache_pages = 0;
1780 self.header.cache_len = 0;
1781 }
1782 }
1783 }
1784
1785 // 3b. Record this transaction as a restorable snapshot. The manifest
1786 // update is staged in a local and adopted only once the commit
1787 // below succeeds, so a failed flush leaves no phantom entry.
1788 //
1789 // The cap (`self.max_snapshots`) is applied to the staged local
1790 // before the new entry lands, so after appending the manifest
1791 // holds at most `max_snapshots` entries — the most-recent N.
1792 // `max_snapshots == 0` disables the cap. Pages referenced only
1793 // by pruned snapshots stay on disk until `compact_file`.
1794 let live = self.len() as u64;
1795 let txn_id = self.header.wal_seq + 1;
1796 let mut manifest = self.manifest.clone();
1797 if self.max_snapshots > 0 {
1798 // Drop the oldest entries so there's room for the new one.
1799 while manifest.len() >= self.max_snapshots {
1800 manifest.remove(0);
1801 }
1802 }
1803 manifest.push(Snapshot {
1804 txn_id,
1805 created_at: memory::now_secs(),
1806 catalog_start: self.header.catalog_start,
1807 catalog_pages: self.header.catalog_pages,
1808 catalog_len: self.header.catalog_len,
1809 index_start: self.header.index_start,
1810 index_pages: self.header.index_pages,
1811 index_len: self.header.index_len,
1812 memory_count: live,
1813 });
1814 let man_bytes =
1815 rmp_serde::to_vec(&manifest).map_err(|e| MnemoError::Serialize(e.to_string()))?;
1816 let (m_start, m_pages) = self.seal_run(&man_bytes, &mut frames)?;
1817 self.header.manifest_start = m_start;
1818 self.header.manifest_pages = m_pages as u64;
1819 self.header.manifest_len = man_bytes.len() as u64;
1820
1821 // 3c. The header is the transaction's final frame; stamp the new id
1822 // and the *actual* (post-flush) write counter, which is <= the
1823 // leased value persisted in `prepare_for_flush`. This is the
1824 // value that ultimately replaces the leased header on checkpoint.
1825 self.header.vector_count = live;
1826 self.header.write_counter = self.pager.write_counter;
1827 self.header.wal_seq = txn_id;
1828 let mut hpage = self.header.to_page();
1829 self.header.apply_seal(&mut hpage, self.pager.dek())?;
1830 frames.push((0, hpage.to_vec()));
1831
1832 // 4. COMMIT — log the transaction and fsync. Nothing at a home page
1833 // other than the leased header has changed yet; this single
1834 // fsync is what makes the rest durable.
1835 let (wal_start, wal_pages) = (self.header.wal_start, self.header.wal_pages);
1836 wal::commit(self.pager.file_mut(), wal_start, wal_pages, txn_id, &frames)?;
1837
1838 // 5. Checkpoint — fold the WAL into the home pages.
1839 self.checkpoint(&frames)?;
1840
1841 self.manifest = manifest;
1842 self.dirty_catalog = false;
1843 self.dirty_index = false;
1844 self.dirty_cache = false;
1845 // Batched policy: successful flush consumes the pending batch.
1846 self.cache_batch.reset();
1847 // Record how many fsyncs this flush issued, for Stats reporting.
1848 // Thread-local counter guarantees we see only this handle's own
1849 // sync_all calls between the snapshot above and this read.
1850 self.fsyncs_last_flush =
1851 (crate::fsync_count().wrapping_sub(fsync_before)).min(u32::MAX as u64) as u32;
1852 Ok(())
1853 }
1854
1855 /// Pre-flush prelude shared by [`Mnemo::flush`] and the
1856 /// `__crash_partial_flush_for_testing` hook.
1857 ///
1858 /// Serializes the control plane (catalog, ANN index), grows the WAL if
1859 /// needed, then **leases** counter and page slots — bumps a *clone* of
1860 /// the header by upper-bound amounts and persists that leased clone
1861 /// with `pager.write_raw + sync`. The in-memory `self.header` keeps
1862 /// its pre-lease `next_page` so [`Mnemo::seal_run`] continues to
1863 /// allocate from the right slot; the final committed header carries
1864 /// the actual post-flush values and overwrites the leased one on
1865 /// checkpoint.
1866 ///
1867 /// This pre-stamping is what closes the nonce-reuse window flagged in
1868 /// Phase 1.1 of the improvement plan: a crash anywhere between here
1869 /// and the WAL commit leaves a file whose on-disk header records a
1870 /// `write_counter` *past* every nonce the rest of this transaction
1871 /// could produce, so the next reopen starts from a counter that
1872 /// guarantees uniqueness against any orphan data pages still on disk.
1873 fn prepare_for_flush(&mut self) -> Result<FlushPrelude> {
1874 // Serialize.
1875 let cat_bytes = if self.dirty_catalog {
1876 Some(
1877 rmp_serde::to_vec(&self.catalog)
1878 .map_err(|e| MnemoError::Serialize(e.to_string()))?,
1879 )
1880 } else {
1881 None
1882 };
1883 let idx_bytes: Option<Vec<u8>> = if self.dirty_index {
1884 match &self.ann {
1885 Some(ann) => {
1886 Some(rmp_serde::to_vec(ann).map_err(|e| MnemoError::Serialize(e.to_string()))?)
1887 }
1888 None => None,
1889 }
1890 } else {
1891 None
1892 };
1893 // Serialize the result-cache directory when it's dirty. An
1894 // empty directory serializes to a small header + 0-element
1895 // array; still cheap enough that we don't special-case it.
1896 let cache_bytes: Option<Vec<u8>> = if self.dirty_cache {
1897 Some(
1898 rmp_serde::to_vec(&self.cache_dir)
1899 .map_err(|e| MnemoError::Serialize(e.to_string()))?,
1900 )
1901 } else {
1902 None
1903 };
1904
1905 let cat_pc = cat_bytes
1906 .as_ref()
1907 .map_or(0, |b| b.len().div_ceil(PAYLOAD).max(1));
1908 let idx_pc = idx_bytes
1909 .as_ref()
1910 .map_or(0, |b| b.len().div_ceil(PAYLOAD).max(1));
1911 let cache_pc = cache_bytes
1912 .as_ref()
1913 .map_or(0, |b| b.len().div_ceil(PAYLOAD).max(1));
1914 // Upper bound on the manifest run: one extra entry, <=82 bytes each
1915 // (rmp_serde encodes a Snapshot as a 9-element fixarray of i64/u64,
1916 // worst case 81 bytes plus a 1-byte array header).
1917 let man_upper = 9 + (self.manifest.len() + 1) * 82;
1918 let man_pc_est = man_upper.div_ceil(PAYLOAD).max(1);
1919
1920 // Grow WAL if needed (does its own header write+sync). The
1921 // cache directory rides in the same transaction so its page
1922 // count contributes to the reservation.
1923 self.ensure_wal_capacity(cat_pc + idx_pc + cache_pc + man_pc_est + 1)?;
1924
1925 // Lease counter and page slots. counter_lease is an upper bound on
1926 // how many `pager.write_counter` advances this transaction will
1927 // make (one per dirty data page + one per sealed control-plane
1928 // page). page_lease covers only the control-plane allocations,
1929 // because dirty data pages were already allocated past
1930 // `header.next_page` by `write_record` when `remember` ran or
1931 // `cache_put` ran.
1932 let data_dirty = self.pager.dirty_page_count() as u64;
1933 let counter_lease = data_dirty + (cat_pc + idx_pc + cache_pc + man_pc_est) as u64;
1934 let page_lease = (cat_pc + idx_pc + cache_pc + man_pc_est) as u64;
1935
1936 // Persist a leased *clone* of the header — in-memory `self.header`
1937 // keeps its pre-lease `next_page` for `seal_run` to consume from
1938 // the correct slot.
1939 let mut leased = self.header.clone();
1940 leased.write_counter = self.pager.write_counter + counter_lease;
1941 leased.next_page += page_lease;
1942 let mut page = leased.to_page();
1943 leased.apply_seal(&mut page, self.pager.dek())?;
1944 self.pager.write_raw(0, &page)?;
1945 self.pager.sync()?;
1946
1947 Ok(FlushPrelude {
1948 cat_bytes,
1949 idx_bytes,
1950 cache_bytes,
1951 })
1952 }
1953
1954 /// Flush and close. Equivalent to `flush()` on a read-write handle; on
1955 /// a read-only handle it's a no-op (no dirty state to persist). The
1956 /// file is released on drop, which also drops the OS advisory lock.
1957 pub fn close(&mut self) -> Result<()> {
1958 if self.read_only {
1959 return Ok(());
1960 }
1961 self.flush()
1962 }
1963
1964 /// **TEST ONLY — do not use.** Reproduces a crash inside
1965 /// [`Mnemo::flush`] *after* the prelude has leased counter+page slots
1966 /// and persisted the leased header, but *before* the WAL is committed.
1967 /// Runs `prepare_for_flush` + `pager.flush` and returns — leaving
1968 /// data pages and a leased header on disk, with no commit frame.
1969 ///
1970 /// Used by `tests/integration.rs::nonce_unique_after_crashed_data_flush`
1971 /// to verify that this window (Phase 1.1 of the improvement plan) does
1972 /// not enable AES-GCM nonce reuse. Calling this in production code
1973 /// strands data pages; never expose it from a binding.
1974 #[doc(hidden)]
1975 pub fn __crash_partial_flush_for_testing(&mut self) -> Result<()> {
1976 if self.read_only {
1977 return Err(MnemoError::ReadOnly);
1978 }
1979 if !self.dirty_catalog && !self.dirty_index {
1980 return Ok(());
1981 }
1982 let _prelude = self.prepare_for_flush()?;
1983 self.pager.flush()
1984 }
1985
1986 /// Bound the in-memory page cache to `pages` decrypted pages.
1987 ///
1988 /// The cache holds decrypted page payloads to speed repeated reads. By
1989 /// default it is capped at 8192 pages (~64 MiB); lower the cap to trade
1990 /// hit rate for a smaller footprint, or raise it for a hotter cache. The
1991 /// cap governs *clean* pages — pages with un-flushed writes are always
1992 /// retained until [`Mnemo::flush`], regardless of the cap.
1993 pub fn set_cache_capacity(&mut self, pages: usize) {
1994 self.pager.set_cache_capacity(pages);
1995 }
1996
1997 /// Page-cache occupancy: `(pages_cached, capacity)`. Renamed from
1998 /// `cache_stats` in v0.4.0: the result-cache API (Phase 10.1) took
1999 /// that name for its per-namespace [`CacheStats`] summary, and the
2000 /// page-cache pair is the older/less-user-facing of the two, so it
2001 /// gets the more-specific name.
2002 pub fn page_cache_stats(&self) -> (usize, usize) {
2003 self.pager.cache_stats()
2004 }
2005
2006 // === Result cache (Phase 10.1 + 10.3) ============================
2007
2008 /// Insert or overwrite a cache entry at `(namespace, key)`.
2009 ///
2010 /// The engine SHA-256-hashes `key` — callers pass the raw input
2011 /// string. Persists to disk on the next `flush()` under the current
2012 /// flush policy (see [`Mnemo::set_cache_flush_policy`]); in
2013 /// [`CacheFlushPolicy::Batched`] the engine may call `flush()`
2014 /// internally when a threshold trips.
2015 ///
2016 /// On overflow of the namespace's [`CacheBudget`] (default 10,000
2017 /// entries / 64 MiB) the LRU entry is tombstoned before the new
2018 /// insert. Tombstones are reclaimed by
2019 /// [`Mnemo::compact_file`], not the runtime flush.
2020 ///
2021 /// Returns [`MnemoError::ReadOnly`] on a read-only handle.
2022 pub fn cache_put(
2023 &mut self,
2024 namespace: &str,
2025 key: &str,
2026 value: &[u8],
2027 opts: CachePutOpts,
2028 ) -> Result<()> {
2029 if self.read_only {
2030 return Err(MnemoError::ReadOnly);
2031 }
2032 let key_hash = result_cache::hash_key(key);
2033 let key_preview = result_cache::key_preview(key);
2034 let now = memory::now_secs();
2035 let size_bytes = value.len() as u64;
2036
2037 let entry = CacheEntry {
2038 namespace: namespace.to_string(),
2039 key_hash,
2040 key_preview: key_preview.clone(),
2041 value: value.to_vec(),
2042 content_type: opts.content_type.clone(),
2043 created_at: now,
2044 accessed_at: now,
2045 access_count: 0,
2046 ttl_secs: opts.ttl_secs,
2047 cost_hint_ms: opts.cost_hint_ms,
2048 size_bytes,
2049 vector: None,
2050 model: None,
2051 };
2052 let bytes = rmp_serde::to_vec(&entry).map_err(|e| MnemoError::Serialize(e.to_string()))?;
2053 let (start_page, page_count) = self.write_record(&bytes)?;
2054
2055 // Tombstone any prior entry for this key. Its pages remain on
2056 // disk until `compact_file`, exactly like memory-record
2057 // overwrites — the tombstone just hides it from lookups.
2058 if let Some(prev_pos) = self.cache_idx.get(namespace, &key_hash) {
2059 self.cache_dir[prev_pos].deleted = true;
2060 self.cache_idx.remove(namespace, &key_hash);
2061 }
2062
2063 let dir_entry = CacheDirectoryEntry {
2064 namespace: namespace.to_string(),
2065 key_hash,
2066 start_page,
2067 page_count,
2068 len: bytes.len() as u32,
2069 deleted: false,
2070 accessed_at: now,
2071 access_count: 0,
2072 size_bytes,
2073 ttl_secs: opts.ttl_secs,
2074 created_at: now,
2075 vector: None,
2076 model: None,
2077 };
2078 let pos = self.cache_dir.len();
2079 self.cache_dir.push(dir_entry);
2080 self.cache_idx.insert(namespace, key_hash, pos);
2081
2082 // Enforce budget — evict LRU tombstones until the namespace
2083 // fits under both caps.
2084 self.evict_to_budget(namespace)?;
2085
2086 self.dirty_cache = true;
2087 self.cache_batch.record_dirty();
2088
2089 // Batched policy: auto-flush if a threshold tripped. Successful
2090 // `flush()` calls `cache_batch.reset()` for us.
2091 if self.cache_batch.should_auto_flush(&self.cache_flush_policy) {
2092 self.flush()?;
2093 }
2094 Ok(())
2095 }
2096
2097 /// Look up `(namespace, key)`. Returns `None` on miss (either the
2098 /// key was never inserted or the entry is TTL-expired). On hit
2099 /// updates catalog-side access stats (mirrors the recall v5-trick);
2100 /// on a read-only handle the stats update is silently skipped.
2101 ///
2102 /// A miss is recorded as a miss in this handle's counters even if
2103 /// the key was TTL-expired — from the caller's perspective the two
2104 /// are indistinguishable.
2105 pub fn cache_get(&mut self, namespace: &str, key: &str) -> Result<Option<CachedValue>> {
2106 let key_hash = result_cache::hash_key(key);
2107 let pos = match self.cache_idx.get(namespace, &key_hash) {
2108 Some(p) => p,
2109 None => {
2110 self.cache_counters.misses += 1;
2111 return Ok(None);
2112 }
2113 };
2114 let now = memory::now_secs();
2115 // Snapshot the directory entry fields we need before any
2116 // mutation — the borrow checker gets grumpy otherwise.
2117 let (start_page, page_count, len, ttl_secs, expired, created_at) = {
2118 let e = &self.cache_dir[pos];
2119 (
2120 e.start_page,
2121 e.page_count,
2122 e.len,
2123 e.ttl_secs,
2124 e.is_expired(now),
2125 e.created_at,
2126 )
2127 };
2128 if expired {
2129 // Treat expired as a miss for the caller. Don't tombstone
2130 // here — cache_purge does that. Just skip.
2131 self.cache_counters.misses += 1;
2132 return Ok(None);
2133 }
2134 // Decrypt and decode the record body.
2135 let mut buf = Vec::with_capacity(len as usize);
2136 for i in 0..page_count as u64 {
2137 buf.extend_from_slice(&self.pager.read_page(start_page + i)?);
2138 }
2139 buf.truncate(len as usize);
2140 let entry: CacheEntry =
2141 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))?;
2142
2143 // Bump access stats catalog-only (v5 trick — no full record
2144 // rewrite on a hit). Skipped on read-only handles so a shared-lock
2145 // reader can't dirty the catalog.
2146 if !self.read_only {
2147 self.cache_dir[pos].accessed_at = now;
2148 self.cache_dir[pos].access_count = self.cache_dir[pos].access_count.saturating_add(1);
2149 self.dirty_cache = true;
2150 }
2151 self.cache_counters.hits += 1;
2152 // On a mutating handle: access_count reflects the just-bumped
2153 // value. On a read-only handle: the on-disk value is unchanged,
2154 // so returning it verbatim is honest (the caller's reads did
2155 // not, and cannot, alter the persisted counter).
2156 let access_count = self.cache_dir[pos].access_count;
2157 Ok(Some(CachedValue {
2158 value: entry.value,
2159 content_type: entry.content_type,
2160 created_at,
2161 accessed_at: if self.read_only {
2162 entry.accessed_at
2163 } else {
2164 now
2165 },
2166 access_count,
2167 ttl_secs,
2168 }))
2169 }
2170
2171 /// Tombstone the entry at `(namespace, key)`. Silently succeeds if
2172 /// no such entry exists. Space is reclaimed by `compact_file`.
2173 pub fn cache_delete(&mut self, namespace: &str, key: &str) -> Result<()> {
2174 if self.read_only {
2175 return Err(MnemoError::ReadOnly);
2176 }
2177 let key_hash = result_cache::hash_key(key);
2178 if let Some(pos) = self.cache_idx.get(namespace, &key_hash) {
2179 self.cache_dir[pos].deleted = true;
2180 self.cache_idx.remove(namespace, &key_hash);
2181 self.dirty_cache = true;
2182 self.cache_batch.record_dirty();
2183 if self.cache_batch.should_auto_flush(&self.cache_flush_policy) {
2184 self.flush()?;
2185 }
2186 }
2187 Ok(())
2188 }
2189
2190 /// Purge cache entries. If `namespace` is `Some`, restrict to that
2191 /// namespace; otherwise purge across the whole cache. If
2192 /// `expired_only` is `true`, only TTL-expired entries are removed;
2193 /// otherwise every live entry in the target scope is tombstoned.
2194 /// Returns the number of entries newly tombstoned.
2195 pub fn cache_purge(&mut self, namespace: Option<&str>, expired_only: bool) -> Result<usize> {
2196 if self.read_only {
2197 return Err(MnemoError::ReadOnly);
2198 }
2199 let now = memory::now_secs();
2200 let mut purged = 0usize;
2201 // Iterate the directory by index; mark matching entries as
2202 // deleted and pull them out of the in-memory index.
2203 for pos in 0..self.cache_dir.len() {
2204 if self.cache_dir[pos].deleted {
2205 continue;
2206 }
2207 if let Some(ns) = namespace {
2208 if self.cache_dir[pos].namespace != ns {
2209 continue;
2210 }
2211 }
2212 if expired_only && !self.cache_dir[pos].is_expired(now) {
2213 continue;
2214 }
2215 let key_hash = self.cache_dir[pos].key_hash;
2216 let ns = self.cache_dir[pos].namespace.clone();
2217 self.cache_dir[pos].deleted = true;
2218 self.cache_idx.remove(&ns, &key_hash);
2219 purged += 1;
2220 }
2221 if purged > 0 {
2222 self.dirty_cache = true;
2223 self.cache_batch.record_dirty();
2224 if self.cache_batch.should_auto_flush(&self.cache_flush_policy) {
2225 self.flush()?;
2226 }
2227 }
2228 Ok(purged)
2229 }
2230
2231 /// Statistics for `namespace` (or the whole cache if `None`).
2232 /// `hits`/`misses`/`evictions` are per-handle in-memory counters
2233 /// that reset on reopen; `entries`/`bytes` reflect the durable
2234 /// on-disk state (live entries only, tombstones excluded).
2235 pub fn cache_stats(&self, namespace: Option<&str>) -> CacheStats {
2236 let mut entries = 0usize;
2237 let mut bytes = 0u64;
2238 for e in &self.cache_dir {
2239 if e.deleted {
2240 continue;
2241 }
2242 if let Some(ns) = namespace {
2243 if e.namespace != ns {
2244 continue;
2245 }
2246 }
2247 entries += 1;
2248 bytes += e.size_bytes;
2249 }
2250 let hits = self.cache_counters.hits;
2251 let misses = self.cache_counters.misses;
2252 let total = hits + misses;
2253 let hit_rate = if total > 0 {
2254 hits as f64 / total as f64
2255 } else {
2256 0.0
2257 };
2258 CacheStats {
2259 entries,
2260 bytes,
2261 hits,
2262 misses,
2263 hit_rate,
2264 evictions: self.cache_counters.evictions,
2265 }
2266 }
2267
2268 /// Set (or update) the eviction budget for `namespace`. Applies to
2269 /// subsequent `cache_put` calls; does not retroactively evict
2270 /// entries that already sit under a looser cap.
2271 pub fn set_cache_budget(&mut self, namespace: &str, budget: CacheBudget) {
2272 self.cache_budgets.insert(namespace.to_string(), budget);
2273 }
2274
2275 /// Set the per-handle cache flush policy. Defaults to
2276 /// [`CacheFlushPolicy::Strict`]. Memory writes are never batched;
2277 /// this affects only cache mutations.
2278 pub fn set_cache_flush_policy(&mut self, policy: CacheFlushPolicy) {
2279 self.cache_flush_policy = policy;
2280 // Switching mid-run: if we're moving to Strict and the batch is
2281 // non-empty, the pending count still counts toward dirty_cache
2282 // which the next flush() will consume.
2283 }
2284
2285 /// Evict LRU live entries in `namespace` until it fits under both
2286 /// budget caps. Called from `cache_put` after the insert; the new
2287 /// entry is included in the count so we always leave room.
2288 fn evict_to_budget(&mut self, namespace: &str) -> Result<()> {
2289 let budget = self
2290 .cache_budgets
2291 .get(namespace)
2292 .copied()
2293 .unwrap_or_default();
2294 loop {
2295 // Sum current live count + bytes in the namespace.
2296 let mut live_count = 0usize;
2297 let mut live_bytes = 0u64;
2298 for e in &self.cache_dir {
2299 if !e.deleted && e.namespace == namespace {
2300 live_count += 1;
2301 live_bytes += e.size_bytes;
2302 }
2303 }
2304 if live_count <= budget.max_entries && live_bytes <= budget.max_bytes {
2305 return Ok(());
2306 }
2307 // Find the LRU (oldest accessed_at) live entry.
2308 let victim_pos = self
2309 .cache_dir
2310 .iter()
2311 .enumerate()
2312 .filter(|(_, e)| !e.deleted && e.namespace == namespace)
2313 .min_by_key(|(_, e)| e.accessed_at)
2314 .map(|(i, _)| i);
2315 let Some(pos) = victim_pos else {
2316 // Nothing to evict but still over budget — the single
2317 // just-inserted entry itself is bigger than
2318 // max_bytes. Nothing we can safely do here without
2319 // rejecting the put; leave it and return.
2320 return Ok(());
2321 };
2322 let key_hash = self.cache_dir[pos].key_hash;
2323 let ns = self.cache_dir[pos].namespace.clone();
2324 self.cache_dir[pos].deleted = true;
2325 self.cache_idx.remove(&ns, &key_hash);
2326 self.cache_counters.evictions += 1;
2327 }
2328 }
2329
2330 // === Semantic cache (Phase 10.2) ==================================
2331
2332 /// Insert or overwrite a **semantic** cache entry — same as
2333 /// [`Mnemo::cache_put`] but with a required embedding vector and
2334 /// model identifier. The vector is stored on the directory entry
2335 /// so [`Mnemo::cache_get_semantic`] can score without decrypting
2336 /// record bodies.
2337 ///
2338 /// Vector dimensionality is validated against the database's
2339 /// configured `dimensions` — a mismatch returns
2340 /// [`MnemoError::DimensionMismatch`], the same signal `remember`
2341 /// gives for the memory catalog.
2342 ///
2343 /// The `model` string is matched **exactly** on lookup — a
2344 /// different model's cache is invisible to a given query. Pick
2345 /// canonical identifiers (`"text-embedding-3-small"`,
2346 /// `"bge-large-en-v1.5"`, ...) and stick to them.
2347 pub fn cache_put_semantic(
2348 &mut self,
2349 namespace: &str,
2350 key: &str,
2351 vector: Vec<f32>,
2352 value: &[u8],
2353 opts: SemanticCachePutOpts,
2354 ) -> Result<()> {
2355 if self.read_only {
2356 return Err(MnemoError::ReadOnly);
2357 }
2358 if vector.len() != self.dimensions {
2359 return Err(MnemoError::DimensionMismatch {
2360 expected: self.dimensions,
2361 got: vector.len(),
2362 });
2363 }
2364 let key_hash = result_cache::hash_key(key);
2365 let key_preview = result_cache::key_preview(key);
2366 let now = memory::now_secs();
2367 let size_bytes = value.len() as u64;
2368
2369 let entry = CacheEntry {
2370 namespace: namespace.to_string(),
2371 key_hash,
2372 key_preview: key_preview.clone(),
2373 value: value.to_vec(),
2374 content_type: opts.content_type.clone(),
2375 created_at: now,
2376 accessed_at: now,
2377 access_count: 0,
2378 ttl_secs: opts.ttl_secs,
2379 cost_hint_ms: opts.cost_hint_ms,
2380 size_bytes,
2381 vector: Some(vector.clone()),
2382 model: Some(opts.model.clone()),
2383 };
2384 let bytes = rmp_serde::to_vec(&entry).map_err(|e| MnemoError::Serialize(e.to_string()))?;
2385 let (start_page, page_count) = self.write_record(&bytes)?;
2386
2387 if let Some(prev_pos) = self.cache_idx.get(namespace, &key_hash) {
2388 self.cache_dir[prev_pos].deleted = true;
2389 self.cache_idx.remove(namespace, &key_hash);
2390 }
2391
2392 let dir_entry = CacheDirectoryEntry {
2393 namespace: namespace.to_string(),
2394 key_hash,
2395 start_page,
2396 page_count,
2397 len: bytes.len() as u32,
2398 deleted: false,
2399 accessed_at: now,
2400 access_count: 0,
2401 size_bytes,
2402 ttl_secs: opts.ttl_secs,
2403 created_at: now,
2404 vector: Some(vector),
2405 model: Some(opts.model),
2406 };
2407 let pos = self.cache_dir.len();
2408 self.cache_dir.push(dir_entry);
2409 self.cache_idx.insert(namespace, key_hash, pos);
2410
2411 self.evict_to_budget(namespace)?;
2412
2413 self.dirty_cache = true;
2414 self.cache_batch.record_dirty();
2415 if self.cache_batch.should_auto_flush(&self.cache_flush_policy) {
2416 self.flush()?;
2417 }
2418 Ok(())
2419 }
2420
2421 /// Semantic recall over the cache: score every live entry in
2422 /// `namespace` whose `model` matches against `query_vector` using
2423 /// cosine similarity; return the top-1 hit and its similarity iff
2424 /// `sim >= threshold`. `None` otherwise.
2425 ///
2426 /// Uses [`crate::result_cache::DEFAULT_SEMANTIC_THRESHOLD`] if you
2427 /// want a sensible default — pass it explicitly (`0.97` today).
2428 ///
2429 /// Exact-key entries (those without a vector) are transparently
2430 /// skipped, as are entries whose model doesn't match — a
2431 /// different-model hit would be a semantic bug, not a win.
2432 /// TTL-expired entries are also skipped.
2433 ///
2434 /// Access-stat updates on hit follow the same policy as
2435 /// [`Mnemo::cache_get`]: bumped on read-write handles, silently
2436 /// skipped on read-only. Hits and misses are recorded in the
2437 /// per-handle counters just like exact-key lookups.
2438 pub fn cache_get_semantic(
2439 &mut self,
2440 namespace: &str,
2441 query_vector: &[f32],
2442 threshold: f32,
2443 model: &str,
2444 ) -> Result<Option<(CachedValue, f32)>> {
2445 if query_vector.len() != self.dimensions {
2446 return Err(MnemoError::DimensionMismatch {
2447 expected: self.dimensions,
2448 got: query_vector.len(),
2449 });
2450 }
2451 let now = memory::now_secs();
2452
2453 // Scan namespace live entries, filter by model, compute
2454 // cosine, keep the best. Exact-key entries have `vector: None`
2455 // and are skipped by the filter.
2456 let mut best: Option<(usize, f32)> = None;
2457 for (i, e) in self.cache_dir.iter().enumerate() {
2458 if e.deleted || e.namespace != namespace || e.is_expired(now) {
2459 continue;
2460 }
2461 let Some(ref v) = e.vector else { continue };
2462 let Some(ref m) = e.model else { continue };
2463 if m != model {
2464 continue;
2465 }
2466 let sim = memory::cosine(query_vector, v);
2467 if best.map(|(_, s)| sim > s).unwrap_or(true) {
2468 best = Some((i, sim));
2469 }
2470 }
2471
2472 let Some((pos, sim)) = best else {
2473 self.cache_counters.misses += 1;
2474 return Ok(None);
2475 };
2476 if sim < threshold {
2477 self.cache_counters.misses += 1;
2478 return Ok(None);
2479 }
2480
2481 // Extract fields before any mutation to keep the borrow checker
2482 // happy.
2483 let (start_page, page_count, len, ttl_secs, created_at) = {
2484 let e = &self.cache_dir[pos];
2485 (e.start_page, e.page_count, e.len, e.ttl_secs, e.created_at)
2486 };
2487
2488 // Decrypt and decode the record body.
2489 let mut buf = Vec::with_capacity(len as usize);
2490 for i in 0..page_count as u64 {
2491 buf.extend_from_slice(&self.pager.read_page(start_page + i)?);
2492 }
2493 buf.truncate(len as usize);
2494 let entry: CacheEntry =
2495 rmp_serde::from_slice(&buf).map_err(|e| MnemoError::Serialize(e.to_string()))?;
2496
2497 if !self.read_only {
2498 self.cache_dir[pos].accessed_at = now;
2499 self.cache_dir[pos].access_count = self.cache_dir[pos].access_count.saturating_add(1);
2500 self.dirty_cache = true;
2501 }
2502 self.cache_counters.hits += 1;
2503 let access_count = self.cache_dir[pos].access_count;
2504 Ok(Some((
2505 CachedValue {
2506 value: entry.value,
2507 content_type: entry.content_type,
2508 created_at,
2509 accessed_at: if self.read_only {
2510 entry.accessed_at
2511 } else {
2512 now
2513 },
2514 access_count,
2515 ttl_secs,
2516 },
2517 sim,
2518 )))
2519 }
2520
2521 /// Begin a conversation [`Session`](crate::Session) for `agent_id`.
2522 ///
2523 /// The session borrows the database for its lifetime, records turns as
2524 /// working memory, and consolidates them into episodic memory when closed.
2525 pub fn session(&mut self, agent_id: impl Into<String>) -> crate::session::Session<'_> {
2526 crate::session::Session::new(self, agent_id.into())
2527 }
2528
2529 // --- snapshots & point-in-time recovery ------------------------------
2530
2531 /// Every committed transaction, oldest first — the restore points
2532 /// available to [`Mnemo::restore_to`] and [`Mnemo::restore_to_time`].
2533 ///
2534 /// Each `flush` appends one snapshot. Because the storage engine is
2535 /// append-only, the pages a past flush wrote are still on disk, so any
2536 /// listed snapshot can be reinstated exactly. The history reaches back to
2537 /// the last [`Mnemo::compact_file`], which reclaims space by collapsing
2538 /// it.
2539 pub fn snapshots(&self) -> Vec<SnapshotInfo> {
2540 let mut v: Vec<SnapshotInfo> = self
2541 .manifest
2542 .iter()
2543 .map(|s| SnapshotInfo {
2544 txn_id: s.txn_id,
2545 created_at: s.created_at,
2546 memory_count: s.memory_count,
2547 })
2548 .collect();
2549 v.sort_by_key(|s| s.txn_id);
2550 v
2551 }
2552
2553 /// Load a past snapshot's state and commit it as a new transaction.
2554 fn apply_snapshot(&mut self, snap: &Snapshot) -> Result<()> {
2555 let catalog = load_catalog(
2556 &mut self.pager,
2557 snap.catalog_start,
2558 snap.catalog_pages,
2559 snap.catalog_len,
2560 )?;
2561 let ann = load_index(
2562 &mut self.pager,
2563 snap.index_start,
2564 snap.index_pages,
2565 snap.index_len,
2566 self.dimensions,
2567 )?;
2568 self.index = build_id_index(&catalog);
2569 self.catalog = catalog;
2570 self.ann = ann;
2571 // Re-commit as a fresh transaction: crash-safe, and itself recorded
2572 // as a new snapshot so a restore can always be undone.
2573 self.dirty_catalog = true;
2574 self.dirty_index = true;
2575 self.flush()
2576 }
2577
2578 /// Restore the database to the snapshot produced by transaction `txn_id`.
2579 ///
2580 /// The restore is itself a new committed transaction (and a new
2581 /// snapshot), so it is crash-safe and reversible — restoring forward to a
2582 /// later snapshot afterwards works just as well.
2583 pub fn restore_to(&mut self, txn_id: u64) -> Result<SnapshotInfo> {
2584 if self.read_only {
2585 return Err(MnemoError::ReadOnly);
2586 }
2587 let snap = self
2588 .manifest
2589 .iter()
2590 .find(|s| s.txn_id == txn_id)
2591 .cloned()
2592 .ok_or_else(|| MnemoError::NotFound(format!("snapshot for transaction {txn_id}")))?;
2593 self.apply_snapshot(&snap)?;
2594 Ok(SnapshotInfo {
2595 txn_id: snap.txn_id,
2596 created_at: snap.created_at,
2597 memory_count: snap.memory_count,
2598 })
2599 }
2600
2601 /// Restore the database to the latest snapshot committed at or before
2602 /// `unix_secs`. Returns [`MnemoError::NotFound`] if no snapshot is that
2603 /// old. Like [`Mnemo::restore_to`], the restore is a new transaction.
2604 pub fn restore_to_time(&mut self, unix_secs: i64) -> Result<SnapshotInfo> {
2605 if self.read_only {
2606 return Err(MnemoError::ReadOnly);
2607 }
2608 let snap = self
2609 .manifest
2610 .iter()
2611 .filter(|s| s.created_at <= unix_secs)
2612 .max_by_key(|s| s.txn_id)
2613 .cloned()
2614 .ok_or_else(|| {
2615 MnemoError::NotFound(format!("no snapshot at or before time {unix_secs}"))
2616 })?;
2617 self.apply_snapshot(&snap)?;
2618 Ok(SnapshotInfo {
2619 txn_id: snap.txn_id,
2620 created_at: snap.created_at,
2621 memory_count: snap.memory_count,
2622 })
2623 }
2624
2625 /// Change the passphrase. Cheap: re-derives the KEK and re-wraps the DEK;
2626 /// the encrypted pages are never rewritten.
2627 pub fn rekey(&mut self, new_passphrase: &str, kdf: KdfParams) -> Result<()> {
2628 if self.read_only {
2629 return Err(MnemoError::ReadOnly);
2630 }
2631 self.flush()?;
2632 let new_salt = crypto::random_salt();
2633 let new_kek = crypto::derive_kek(new_passphrase.as_bytes(), &new_salt, kdf)?;
2634 let new_nonce = crypto::random_nonce();
2635 let wrapped = crypto::wrap_dek(&new_kek, &new_nonce, self.pager.dek())?;
2636 if wrapped.len() != WRAPPED_DEK_LEN {
2637 return Err(MnemoError::Crypto("unexpected wrapped DEK length".into()));
2638 }
2639 let mut wrapped_dek = [0u8; WRAPPED_DEK_LEN];
2640 wrapped_dek.copy_from_slice(&wrapped);
2641
2642 self.header.salt = new_salt;
2643 self.header.dek_nonce = new_nonce;
2644 self.header.wrapped_dek = wrapped_dek;
2645 self.header.m_cost = kdf.m_cost;
2646 self.header.t_cost = kdf.t_cost;
2647 self.header.p_cost = kdf.p_cost;
2648 self.header.write_counter = self.pager.write_counter;
2649 self.kdf = kdf;
2650
2651 let mut page = self.header.to_page();
2652 self.header.apply_seal(&mut page, self.pager.dek())?;
2653 self.pager.write_raw(0, &page)?;
2654 self.pager.sync()?;
2655 Ok(())
2656 }
2657
2658 /// Decrypt and validate every live record. Returns the count verified.
2659 pub fn verify(&mut self) -> Result<usize> {
2660 let entries: Vec<CatalogEntry> = self
2661 .catalog
2662 .iter()
2663 .filter(|e| !e.deleted)
2664 .cloned()
2665 .collect();
2666 let n = entries.len();
2667 for e in &entries {
2668 let bytes = self.read_record(e)?;
2669 let _m: Memory = rmp_serde::from_slice(&bytes)
2670 .map_err(|err| MnemoError::Serialize(err.to_string()))?;
2671 }
2672 Ok(n)
2673 }
2674
2675 /// Summary statistics for the open database.
2676 pub fn stats(&mut self) -> Result<Stats> {
2677 let deleted = self.catalog.iter().filter(|e| e.deleted).count();
2678 let mut agents: Vec<String> = self.memories()?.into_iter().map(|m| m.agent_id).collect();
2679 agents.sort();
2680 agents.dedup();
2681 let file_bytes = self.header.next_page * PAGE_SIZE as u64;
2682 Ok(Stats {
2683 memories: self.len(),
2684 deleted,
2685 dimensions: self.dimensions,
2686 file_bytes,
2687 agents,
2688 encrypted: self.header.flags & FLAG_ENCRYPTED != 0,
2689 created_at: self.header.created_at,
2690 index: self.ann.as_ref().map(|a| a.info()),
2691 wal_pages: self.header.wal_pages,
2692 fsyncs_last_flush: self.fsyncs_last_flush,
2693 })
2694 }
2695
2696 /// Rewrite the file, dropping tombstoned and expired memories and
2697 /// reclaiming stale pages left by updates. Done out-of-place via a temp
2698 /// file and an atomic rename.
2699 pub fn compact_file(path: &str, passphrase: &str) -> Result<CompactReport> {
2700 let mut old = Mnemo::open(path, passphrase)?;
2701 let dims = old.dimensions;
2702 let kdf = old.kdf;
2703 let tmp = format!("{path}.compact-tmp");
2704
2705 let mut new = Mnemo::create(
2706 &tmp,
2707 passphrase,
2708 MnemoConfig {
2709 dimensions: dims,
2710 kdf,
2711 ..Default::default()
2712 },
2713 )?;
2714 let want_index = old.ann.as_ref().map(|a| (a.n_probe(), a.n_rerank()));
2715 let now = memory::now_secs();
2716 let all = old.memories()?;
2717 let before = all.len();
2718 let mut after = 0;
2719 for m in all {
2720 if m.is_expired(now) {
2721 continue;
2722 }
2723 new.remember(m)?;
2724 after += 1;
2725 }
2726 // Rebuild the index fresh (re-clustered) when the source had one.
2727 if let Some((n_probe, n_rerank)) = want_index {
2728 if after > 0 {
2729 new.build_index_with(IndexConfig {
2730 n_probe,
2731 n_rerank,
2732 ..Default::default()
2733 })?;
2734 }
2735 }
2736 new.flush()?;
2737 drop(new);
2738 drop(old);
2739
2740 std::fs::rename(&tmp, path)?;
2741 Ok(CompactReport { before, after })
2742 }
2743}