From 7421d8b707a6092c623901159c5632173a3fbe53 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 07:16:58 +0000 Subject: [PATCH 1/7] fix: harden memory persistence boundaries --- src/memory/chunks/connection.rs | 180 ++-- src/memory/chunks/connection_breaker.rs | 50 + src/memory/chunks/embeddings.rs | 4 +- src/memory/chunks/migrations.rs | 17 +- src/memory/chunks/mod.rs | 52 +- src/memory/chunks/produce.rs | 93 +- src/memory/chunks/produce_tests.rs | 60 +- src/memory/chunks/raw_refs.rs | 29 +- src/memory/chunks/recovery.rs | 94 +- src/memory/chunks/recovery_tests.rs | 26 +- src/memory/chunks/store.rs | 136 +-- src/memory/chunks/store_conn_tests.rs | 12 +- src/memory/chunks/store_delete.rs | 282 +++--- src/memory/chunks/store_embed_tests.rs | 7 +- src/memory/chunks/store_list.rs | 121 +++ src/memory/chunks/store_sources.rs | 13 +- src/memory/chunks/store_tests.rs | 275 +----- src/memory/chunks/store_tests_more.rs | 397 ++++++++ src/memory/chunks/types_tests.rs | 77 ++ src/memory/conversations/bus.rs | 38 +- src/memory/conversations/bus_tests.rs | 39 + src/memory/conversations/inverted_index.rs | 44 +- .../conversations/inverted_index_tests.rs | 20 + src/memory/conversations/mod.rs | 10 +- src/memory/conversations/store_index.rs | 79 +- src/memory/conversations/store_ops.rs | 17 +- src/memory/conversations/store_tests.rs | 854 +----------------- src/memory/conversations/store_tests_late.rs | 478 ++++++++++ src/memory/conversations/store_tests_more.rs | 427 +++++++++ src/memory/diff/ledger.rs | 231 +---- src/memory/diff/ledger_helpers.rs | 215 +++++ src/memory/diff/ledger_tests.rs | 76 ++ src/memory/diff/mod.rs | 29 +- src/memory/diff/snapshot.rs | 2 +- src/memory/entities/frontmatter.rs | 2 +- src/memory/fsutil.rs | 2 +- src/memory/goals/mod.rs | 2 +- src/memory/goals/reflect.rs | 25 +- src/memory/goals/reflect_tests.rs | 18 + src/memory/goals/types.rs | 2 +- src/memory/sources/mod.rs | 6 +- src/memory/sources/readers/conversation.rs | 8 +- .../sources/readers/conversation_tests.rs | 22 + src/memory/sources/readers/folder.rs | 20 +- src/memory/sources/readers/folder_tests.rs | 21 + src/memory/sources/registry.rs | 176 +--- src/memory/sources/registry_tests.rs | 67 +- src/memory/sources/types.rs | 176 ++++ src/memory/sources/validation.rs | 11 +- src/memory/store/content/atomic.rs | 24 +- .../store/content/compose/compose_tests.rs | 129 +-- .../content/compose/compose_tests_more.rs | 129 +++ src/memory/store/content/compose/yaml.rs | 2 +- src/memory/store/content/mod.rs | 12 +- src/memory/store/content/tags.rs | 17 +- src/memory/store/entity_index/mod.rs | 23 +- src/memory/store/entity_index/store.rs | 140 +-- src/memory/store/entity_index/store_tests.rs | 37 + src/memory/store/entity_index/transaction.rs | 120 +++ src/memory/store/kv.rs | 49 +- src/memory/store/kv_tests.rs | 37 + src/memory/store/memory_trait.rs | 266 ++++++ src/memory/store/memory_trait_tests.rs | 69 ++ src/memory/store/mod.rs | 15 +- src/memory/store/safety/mod.rs | 16 +- src/memory/store/safety/pii.rs | 630 +------------ src/memory/store/safety/pii/checks.rs | 232 +++++ src/memory/store/safety/pii_tests.rs | 388 ++++++++ src/memory/store/store.rs | 92 +- src/memory/store/store_tests.rs | 60 ++ src/memory/store/types.rs | 12 +- src/memory/store/vectors/store.rs | 53 +- src/memory/store/vectors/store_tests.rs | 171 ++-- src/memory/store/vectors/store_tests_more.rs | 114 +++ src/memory/tool_memory/render.rs | 37 +- src/memory/tool_memory/render_tests.rs | 24 + src/memory/tool_memory/store.rs | 26 +- src/memory/tool_memory/store_tests.rs | 39 + src/memory/tool_memory/types.rs | 8 +- 79 files changed, 4644 insertions(+), 3369 deletions(-) create mode 100644 src/memory/chunks/connection_breaker.rs create mode 100644 src/memory/chunks/store_list.rs create mode 100644 src/memory/chunks/store_tests_more.rs create mode 100644 src/memory/conversations/store_tests_late.rs create mode 100644 src/memory/conversations/store_tests_more.rs create mode 100644 src/memory/diff/ledger_helpers.rs create mode 100644 src/memory/store/content/compose/compose_tests_more.rs create mode 100644 src/memory/store/entity_index/transaction.rs create mode 100644 src/memory/store/memory_trait.rs create mode 100644 src/memory/store/memory_trait_tests.rs create mode 100644 src/memory/store/safety/pii/checks.rs create mode 100644 src/memory/store/safety/pii_tests.rs create mode 100644 src/memory/store/store_tests.rs create mode 100644 src/memory/store/vectors/store_tests_more.rs diff --git a/src/memory/chunks/connection.rs b/src/memory/chunks/connection.rs index d0a7770..8661553 100644 --- a/src/memory/chunks/connection.rs +++ b/src/memory/chunks/connection.rs @@ -5,22 +5,26 @@ //! `ConnectionCache` keyed by DB path: each entry holds one //! `parking_lot::Mutex` that is initialised once (schema + //! migrations) and reused for all subsequent calls. A per-entry -//! [`CircuitBreaker`] stops retrying after [`CB_THRESHOLD`] consecutive init -//! failures for [`CB_COOLDOWN`] so a broken install does not busy-loop. +//! [`CircuitBreaker`] stops retrying after the configured failure threshold +//! for a cooldown period so a broken install does not busy-loop. use anyhow::{Context, Result}; use parking_lot::Mutex as PMutex; use rusqlite::Connection; use std::collections::HashMap; use std::path::{Path, PathBuf}; -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; #[cfg(test)] use std::sync::Mutex; use std::sync::{Arc, OnceLock}; -use std::time::{Duration, Instant}; + +use super::connection_breaker::CircuitBreaker; +#[cfg(test)] +pub(crate) use super::connection_breaker::CB_THRESHOLD; use super::migrations::{migrate_legacy_embeddings_to_sidecar, purge_global_topic_trees}; -use super::recovery::{is_io_open_error, try_cleanup_stale_files}; +use super::recovery::{ + is_corrupt_error, is_io_open_error, quarantine_corrupt_files, try_cleanup_stale_files, +}; use super::schema::SCHEMA; use super::{db_path_for, SQLITE_BUSY_TIMEOUT}; use crate::memory::config::MemoryConfig; @@ -60,72 +64,6 @@ pub(crate) fn schema_apply_count_for_path_for_tests(path: &Path) -> usize { .unwrap_or(0) } -// ── Circuit breaker ────────────────────────────────────────────────────────── - -/// How many consecutive init failures before the circuit breaker trips. -pub(crate) const CB_THRESHOLD: u32 = 3; -/// How long the circuit breaker holds the DB closed after tripping. -pub(crate) const CB_COOLDOWN: Duration = Duration::from_secs(30); - -/// Per-path circuit breaker: after [`CB_THRESHOLD`] consecutive init failures -/// the breaker trips and `get_or_init_connection` returns an error immediately -/// until [`CB_COOLDOWN`] elapses. On the first success it resets to zero. -struct CircuitBreaker { - /// Count of consecutive init failures since the last success. Reset to 0 - /// on any success. - consecutive_failures: AtomicU32, - /// Whether the breaker is currently open (rejecting new attempts). - tripped: AtomicBool, - /// Timestamp of the most recent trip/re-trip, used by [`Self::is_open`] to - /// compute whether [`CB_COOLDOWN`] has elapsed. `None` when never tripped - /// or after a reset. - last_trip: PMutex>, -} - -impl CircuitBreaker { - fn new() -> Self { - Self { - consecutive_failures: AtomicU32::new(0), - tripped: AtomicBool::new(false), - last_trip: PMutex::new(None), - } - } - - /// Records a successful init. Returns `true` if this call cleared a - /// previously-tripped breaker (a transition back to healthy). - fn record_success(&self) -> bool { - self.consecutive_failures.store(0, Ordering::Relaxed); - *self.last_trip.lock() = None; - self.tripped.swap(false, Ordering::Relaxed) - } - - /// Records one more failure. Returns `true` if this call just tripped the - /// breaker (i.e. the threshold was crossed right now). - fn record_failure(&self) -> bool { - let prev = self.consecutive_failures.fetch_add(1, Ordering::Relaxed); - let count = prev + 1; - if count >= CB_THRESHOLD && !self.tripped.swap(true, Ordering::Relaxed) { - *self.last_trip.lock() = Some(Instant::now()); - return true; - } - // Re-arm the cooldown on each subsequent failure while already tripped. - if self.tripped.load(Ordering::Relaxed) { - *self.last_trip.lock() = Some(Instant::now()); - } - false - } - - /// Returns `true` when the breaker is open AND the cooldown has not yet - /// elapsed. Returns `false` (allowing a retry) once the cooldown passes. - fn is_open(&self) -> bool { - if !self.tripped.load(Ordering::Relaxed) { - return false; - } - let guard = self.last_trip.lock(); - matches!(*guard, Some(t) if t.elapsed() < CB_COOLDOWN) - } -} - // ── Connection cache ───────────────────────────────────────────────────────── /// Process-global, per-DB-path connection state. @@ -135,7 +73,7 @@ impl CircuitBreaker { /// circuit breaker tracking recent init failures, and a per-path init lock /// used to serialise cold-start initialisation across concurrent callers. /// Entries are never evicted except by [`invalidate_connection`], -/// [`drop_cached_connection`], or [`clear_connection_cache`] (test-only) — one +/// [`drop_cached_connection`], or test-only path-local reset — one /// entry accumulates per distinct workspace path for the life of the process. struct ConnectionCache { /// One live, already-initialised `Connection` per DB path, wrapped so @@ -213,13 +151,10 @@ fn apply_schema(conn: &Connection) -> Result<()> { let journal_mode: String = conn .query_row("PRAGMA journal_mode=TRUNCATE", [], |row| row.get(0)) .context("Failed to set chunk DB journal_mode=TRUNCATE")?; - // NOTE: SQLite can refuse a journal-mode change (e.g. an open transaction - // elsewhere, or WAL mode held by another connection to the same file) by - // simply returning the *previous* mode instead of erroring. This check is - // currently a no-op — a refusal is silently accepted here, and the - // synchronous=FULL crash-safety assumption in `init_db` is only actually - // valid when the mode really did become TRUNCATE. See audit finding SC-9. - let _ = journal_mode.eq_ignore_ascii_case("truncate"); + anyhow::ensure!( + journal_mode.eq_ignore_ascii_case("truncate"), + "SQLite refused journal_mode=TRUNCATE (active mode: {journal_mode})" + ); conn.execute_batch(SCHEMA) .context("Failed to initialize chunk DB schema")?; // Additive, idempotent migrations. @@ -300,8 +235,8 @@ pub(super) fn add_column_if_missing( /// per-path circuit breaker is open (see [`CircuitBreaker`]). Otherwise /// returns `Err` if opening the SQLite file or running [`init_db`] fails /// (after the single stale-file-cleanup retry) — each failure is recorded -/// against the breaker, and three consecutive failures trip it for -/// [`CB_COOLDOWN`]. +/// against the breaker, and three consecutive failures trip it for the +/// configured cooldown. pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result>> { let db_path = db_path_for(config); @@ -346,14 +281,23 @@ pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result { @@ -369,11 +313,12 @@ pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result { @@ -448,15 +393,48 @@ pub(super) fn drop_cached_connection(config: &MemoryConfig) { conn_cache().breakers.lock().remove(&db_path); } -/// Clear cached connections and init locks for test isolation. +/// Quarantine and rebuild a corrupt database while holding the same per-path +/// initialization lock used by cold opens. +pub(super) fn recover_corrupt_connection(config: &MemoryConfig) -> Result { + let db_path = db_path_for(config); + let init_lock = { + let mut locks = conn_cache().init_locks.lock(); + locks + .entry(db_path.clone()) + .or_insert_with(|| Arc::new(PMutex::new(()))) + .clone() + }; + let _init_guard = init_lock.lock(); + + drop_cached_connection(config); + let quarantined = quarantine_corrupt_files(config)?; + let conn = open_and_init(&db_path, config) + .context("failed to rebuild chunk DB schema after quarantining corrupt DB")?; + let connection = Arc::new(PMutex::new(conn)); + conn_cache() + .connections + .lock() + .insert(db_path.clone(), connection); + let breaker = { + let mut breakers = conn_cache().breakers.lock(); + breakers + .entry(db_path) + .or_insert_with(|| Arc::new(CircuitBreaker::new())) + .clone() + }; + breaker.record_success(); + Ok(quarantined) +} + +/// Remove only one workspace from the test connection cache. /// -/// Breakers are deliberately retained: tests use unique temporary workspace -/// paths, and clearing the process-wide breaker map races with parallel tests -/// that are proving threshold behavior for another path. +/// Keeping cache-reset scope local prevents parallel tests with unrelated +/// tempdirs from invalidating one another's live `Arc`s. #[cfg(test)] -pub(crate) fn clear_connection_cache() { - conn_cache().connections.lock().clear(); - conn_cache().init_locks.lock().clear(); +pub(crate) fn clear_connection_cache_for(config: &MemoryConfig) { + let path = db_path_for(config); + conn_cache().connections.lock().remove(&path); + conn_cache().init_locks.lock().remove(&path); } /// Open the chunk SQLite DB and run a closure against it. diff --git a/src/memory/chunks/connection_breaker.rs b/src/memory/chunks/connection_breaker.rs new file mode 100644 index 0000000..6cbe483 --- /dev/null +++ b/src/memory/chunks/connection_breaker.rs @@ -0,0 +1,50 @@ +//! Per-database circuit breaker for repeated connection-init failures. + +use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; +use std::time::{Duration, Instant}; + +use parking_lot::Mutex; + +pub(crate) const CB_THRESHOLD: u32 = 3; +pub(crate) const CB_COOLDOWN: Duration = Duration::from_secs(30); + +pub(super) struct CircuitBreaker { + consecutive_failures: AtomicU32, + tripped: AtomicBool, + last_trip: Mutex>, +} + +impl CircuitBreaker { + pub fn new() -> Self { + Self { + consecutive_failures: AtomicU32::new(0), + tripped: AtomicBool::new(false), + last_trip: Mutex::new(None), + } + } + + pub fn record_success(&self) -> bool { + self.consecutive_failures.store(0, Ordering::Relaxed); + *self.last_trip.lock() = None; + self.tripped.swap(false, Ordering::Relaxed) + } + + pub fn record_failure(&self) -> bool { + let count = self.consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1; + if count >= CB_THRESHOLD && !self.tripped.swap(true, Ordering::Relaxed) { + *self.last_trip.lock() = Some(Instant::now()); + return true; + } + if self.tripped.load(Ordering::Relaxed) { + *self.last_trip.lock() = Some(Instant::now()); + } + false + } + + pub fn is_open(&self) -> bool { + if !self.tripped.load(Ordering::Relaxed) { + return false; + } + matches!(*self.last_trip.lock(), Some(tripped) if tripped.elapsed() < CB_COOLDOWN) + } +} diff --git a/src/memory/chunks/embeddings.rs b/src/memory/chunks/embeddings.rs index 24575b5..0e07e25 100644 --- a/src/memory/chunks/embeddings.rs +++ b/src/memory/chunks/embeddings.rs @@ -119,7 +119,7 @@ pub fn set_chunk_embedding_for_signature( /// caller commits (or rolls back) `tx`. /// /// # Errors -/// See [`upsert_chunk_embedding_conn`]. +/// Uses the same validation as the connection-scoped upsert helper. pub fn set_chunk_embedding_for_signature_tx( tx: &rusqlite::Transaction<'_>, chunk_id: &str, @@ -132,7 +132,7 @@ pub fn set_chunk_embedding_for_signature_tx( /// Transaction-scoped summary embedding upsert (used by the legacy migration). /// /// # Errors -/// See [`upsert_summary_embedding_conn`]. +/// Uses the same validation as the connection-scoped upsert helper. pub fn set_summary_embedding_for_signature_tx( tx: &rusqlite::Transaction<'_>, summary_id: &str, diff --git a/src/memory/chunks/migrations.rs b/src/memory/chunks/migrations.rs index 3457a64..80c6db6 100644 --- a/src/memory/chunks/migrations.rs +++ b/src/memory/chunks/migrations.rs @@ -6,12 +6,9 @@ //! ## Contract //! Both migrations here follow the same shape: read `user_version`, bail out //! if it is already at/past the migration's target version, otherwise do the -//! work inside one `unchecked_transaction`, commit, then bump -//! `user_version` in a *separate* statement after the commit. That last step -//! is not itself transactional with the migration body — a crash between -//! commit and the `pragma_update` re-runs the (idempotent, `DELETE`/copy-only) -//! migration body on next open, which is safe. The reverse is not safe: see -//! the NOTE on [`migrate_legacy_embeddings_to_sidecar`] (audit finding SC-10). +//! work and bump `user_version` inside one `unchecked_transaction`. A crash can +//! therefore expose neither the migrated rows nor the version marker, or both, +//! but never a split state. //! //! Neither migration is exercised by an automated test in this module (audit //! test-coverage gap) — both are asserted only via the module's behavior at @@ -97,9 +94,9 @@ pub(super) fn migrate_legacy_embeddings_to_sidecar( } } + tx.pragma_update(None, "user_version", TREE_EMBEDDING_MIGRATION_VERSION) + .context("set PRAGMA user_version during embedding migration")?; tx.commit()?; - conn.pragma_update(None, "user_version", TREE_EMBEDDING_MIGRATION_VERSION) - .context("set PRAGMA user_version after embedding migration")?; Ok(()) } @@ -166,6 +163,8 @@ pub(super) fn purge_global_topic_trees(conn: &Connection, config: &MemoryConfig) "DELETE FROM mem_tree_jobs WHERE kind IN ('topic_route','digest_daily')", [], )?; + tx.pragma_update(None, "user_version", GLOBAL_TOPIC_PURGE_MIGRATION_VERSION) + .context("set PRAGMA user_version during global/topic purge")?; tx.commit()?; // On-disk: drop `wiki/summaries/global*` and `topic-*` summary folders. @@ -181,7 +180,5 @@ pub(super) fn purge_global_topic_trees(conn: &Connection, config: &MemoryConfig) } } - conn.pragma_update(None, "user_version", GLOBAL_TOPIC_PURGE_MIGRATION_VERSION) - .context("set PRAGMA user_version after global/topic purge")?; Ok(()) } diff --git a/src/memory/chunks/mod.rs b/src/memory/chunks/mod.rs index e8851e1..1704079 100644 --- a/src/memory/chunks/mod.rs +++ b/src/memory/chunks/mod.rs @@ -3,13 +3,13 @@ //! One module for the full chunk lifecycle, ported faithfully from //! OpenHuman's `memory_store/chunks`: //! -//! - [`types`] — [`Chunk`], [`Metadata`], [`SourceKind`], [`DataSource`], +//! - `types` — [`Chunk`], [`Metadata`], [`SourceKind`], [`DataSource`], //! [`SourceRef`], [`StagedChunk`] and the deterministic //! [`chunk_id`] / token-estimate functions. The persisted shape. -//! - [`produce`] — source-kind-dispatch chunker (chat / email / document). +//! - `produce` — source-kind-dispatch chunker (chat / email / document). //! Splits canonical Markdown into bounded chunks with stable per-source //! sequence numbers. -//! - [`semantic`] — heading- and paragraph-aware chunker used to split large +//! - `semantic` — heading- and paragraph-aware chunker used to split large //! documents into LLM-context-sized pieces while preserving heading context. //! Exported as [`chunk_semantic`]. //! - `store` / `connection` / `migrations` / `raw_refs` / `embeddings` — @@ -26,7 +26,6 @@ //! one-shot SQLite migrations; everything else is left to the modules that //! own those tables. -use std::collections::HashSet; use std::path::PathBuf; use std::time::Duration; @@ -36,6 +35,7 @@ mod schema; #[path = "connection.rs"] mod connection; +mod connection_breaker; #[path = "recovery.rs"] mod recovery; @@ -55,6 +55,7 @@ mod semantic; mod store; #[path = "store_delete.rs"] mod store_delete; +mod store_list; #[path = "store_sources.rs"] mod store_sources; #[path = "types.rs"] @@ -101,16 +102,17 @@ pub use store::{ claim_source_ingest_tx, count_chunks, count_chunks_by_lifecycle_status, count_raw_paths_ingested_with_prefix, delete_source_ingest, extraction_coverage, filter_raw_paths_not_ingested, get_chunk, get_chunk_lifecycle_status, get_chunks_batch, - is_source_ingested, list_chunks, list_source_ids_with_prefix, mark_raw_paths_ingested, + is_source_ingested, list_source_ids_with_prefix, mark_raw_paths_ingested, set_chunk_lifecycle_status, update_chunk_content_sha256, update_summary_content_sha256, - upsert_chunks, upsert_chunks_tx, upsert_staged_chunks_tx, ListChunksQuery, - CHUNK_STATUS_ADMITTED, CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, - CHUNK_STATUS_PENDING_EXTRACTION, CHUNK_STATUS_SEALED, RAW_FILE_GATE_KIND, + upsert_chunks, upsert_chunks_tx, upsert_staged_chunks_tx, CHUNK_STATUS_ADMITTED, + CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, CHUNK_STATUS_PENDING_EXTRACTION, + CHUNK_STATUS_SEALED, RAW_FILE_GATE_KIND, }; pub use store_delete::{ delete_chunks_by_owner, delete_chunks_by_source, delete_chunks_by_source_prefix, delete_orphaned_source_tree, }; +pub use store_list::{list_chunks, ListChunksQuery}; pub use store_sources::{get_chunk_lifecycle_status_tx, set_chunk_lifecycle_status_tx}; // ── Shared internal constants / helpers ───────────────────────────────────── @@ -154,37 +156,3 @@ pub(crate) fn redact(s: &str) -> String { let d = h.finalize(); format!("{:08x}", u32::from_be_bytes([d[0], d[1], d[2], d[3]])) } - -/// Tag marking a chunk as belonging to a configurable "memory source". -const MEMORY_SOURCE_TAG: &str = "memory_sources"; - -/// Extract the memory-source id from a composite `mem_src::` source -/// id, or `None` when the id is not in that format. Mirrors OpenHuman's -/// `memory::sync::extract_mem_src_id`. -fn extract_mem_src_id(composite_source_id: &str) -> Option<&str> { - let rest = composite_source_id.strip_prefix("mem_src:")?; - let colon_pos = rest.find(':')?; - let source_id = &rest[..colon_pos]; - if colon_pos + 1 >= rest.len() { - return None; - } - Some(source_id) -} - -/// Whether a chunk is allowed under a per-profile memory-source allowlist. -/// Non-memory-source chunks always pass; memory-source chunks pass only when -/// their (possibly composite) source id resolves into `set`. -pub(crate) fn chunk_source_allowed_in( - set: &HashSet, - tags: &[String], - source_id: &str, -) -> bool { - let is_memory_source = tags.iter().any(|t| t == MEMORY_SOURCE_TAG); - if !is_memory_source { - return true; - } - if set.contains(source_id) { - return true; - } - extract_mem_src_id(source_id).is_some_and(|id| set.contains(id)) -} diff --git a/src/memory/chunks/produce.rs b/src/memory/chunks/produce.rs index 0c84985..30399c4 100644 --- a/src/memory/chunks/produce.rs +++ b/src/memory/chunks/produce.rs @@ -67,26 +67,15 @@ pub struct ChunkerInput { /// /// ## Dispatch by source kind /// -/// - **Chat / Email**: split at message/email boundaries, then greedy-pack -/// consecutive units into a single chunk until adding the next unit would -/// exceed `max_tokens`. Oversize units (a single message > `max_tokens`) -/// fall back to [`split_by_token_budget`] and emit each piece with -/// `partial_message = true`. -/// - **Document**: split by [`split_by_token_budget`] — sized by the +/// - **Chat / Email**: one chunk per message/email boundary. Oversize units +/// fall back to the bounded splitter and emit each piece with +/// `partial_message = true`. IDs derive from the complete logical unit plus +/// the piece index, so overlapping deliveries are stable regardless of batch +/// position. +/// - **Document**: split by the bounded token-budget helper — sized by the /// conservative token estimate (paragraph → sentence → whitespace → /// hard-char) with ~12% overlap between adjacent chunks. /// -/// # NOTE — re-ingest of a grown source can leave stale rows (audit SC-8) -/// Chunk ids are deterministic in `(source_kind, source_id, seq, content)`, so -/// re-chunking byte-identical input reproduces the exact same ids — an -/// idempotent no-op against the store. But greedy packing means *appending* -/// new messages/emails to a source changes the content of what used to be the -/// last packed chunk at a given `seq`, so re-chunking a grown source produces -/// a **new** id at that `seq` rather than replacing the old row. This -/// function itself has no store access and cannot dedupe across calls; the -/// caller ([`super::store::upsert_chunks`]) only adds/replaces by id, so the -/// old row from the pre-growth chunking is never removed unless the caller -/// explicitly reconciles by source. pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec { let now = chrono::Utc::now(); let max_tokens = opts.max_tokens.max(1); @@ -121,47 +110,18 @@ pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec .collect(); } - // For Chat and Email: greedy-pack consecutive units into chunks. - let unit_separator = "\n\n"; - let sep_chars = unit_separator.chars().count(); - + // Chat/email IDs are tied to each complete logical unit, not its batch-local + // sequence. This makes overlapping redelivery idempotent. let mut out: Vec = Vec::new(); - let mut acc: Vec = Vec::new(); - let mut acc_chars = 0usize; - - // Flush accumulated units as one packed chunk. - let flush = |acc: &mut Vec, acc_chars: &mut usize, out: &mut Vec| { - if acc.is_empty() { - return; - } - let content = acc.join(unit_separator); - let seq = out.len() as u32; - let tc = approx_token_count(&content); - let id = chunk_id(input.source_kind, &input.source_id, seq, &content); - out.push(Chunk { - id, - content, - metadata: input.metadata.clone(), - token_count: tc, - seq_in_source: seq, - created_at: now, - partial_message: false, - }); - acc.clear(); - *acc_chars = 0; - }; - for unit in units { let unit_chars = unit.chars().count(); if unit_chars > max_chars { - // Oversize: flush any pending accumulator first, then sub-split. - flush(&mut acc, &mut acc_chars, &mut out); let sub_pieces = split_by_token_budget(&unit, max_tokens); - for piece in sub_pieces { + for (part, piece) in sub_pieces.into_iter().enumerate() { let seq = out.len() as u32; let tc = approx_token_count(&piece); - let id = chunk_id(input.source_kind, &input.source_id, seq, &piece); + let id = chunk_id(input.source_kind, &input.source_id, part as u32, &unit); out.push(Chunk { id, content: piece, @@ -174,29 +134,20 @@ pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec } continue; } - - // Compute projected size if we add this unit to the accumulator. - let projected = if acc.is_empty() { - unit_chars - } else { - acc_chars + sep_chars + unit_chars - }; - - if projected > max_chars { - // Adding this unit would overflow — flush the accumulator first. - flush(&mut acc, &mut acc_chars, &mut out); - } - - if !acc.is_empty() { - acc_chars += sep_chars; - } - acc_chars += unit_chars; - acc.push(unit); + let seq = out.len() as u32; + let token_count = approx_token_count(&unit); + let id = chunk_id(input.source_kind, &input.source_id, 0, &unit); + out.push(Chunk { + id, + content: unit, + metadata: input.metadata.clone(), + token_count, + seq_in_source: seq, + created_at: now, + partial_message: false, + }); } - // Flush any remaining accumulated units. - flush(&mut acc, &mut acc_chars, &mut out); - if out.is_empty() { // Degenerate: empty input → one empty chunk, matching original behaviour. let id = chunk_id(input.source_kind, &input.source_id, 0, ""); diff --git a/src/memory/chunks/produce_tests.rs b/src/memory/chunks/produce_tests.rs index 64ebc22..72339f7 100644 --- a/src/memory/chunks/produce_tests.rs +++ b/src/memory/chunks/produce_tests.rs @@ -46,7 +46,7 @@ fn empty_chat_input_produces_one_empty_chunk() { } #[test] -fn chat_messages_pack_into_one_chunk_when_small() { +fn chat_messages_keep_one_chunk_per_message_when_small() { let md = "## 2026-01-01T00:00:00Z — alice\nHello world\n\n## 2026-01-01T00:01:00Z — bob\nParagraph one.\n\nParagraph two.".to_string(); let input = ChunkerInput { source_kind: SourceKind::Chat, @@ -57,14 +57,14 @@ fn chat_messages_pack_into_one_chunk_when_small() { let chunks = chunk_markdown(&input, &ChunkerOptions::default()); assert_eq!( chunks.len(), - 1, - "small messages should be packed into one chunk; got {chunks:?}" + 2, + "each message must keep a stable replay boundary; got {chunks:?}" ); assert!(chunks[0].content.contains("alice")); - assert!(chunks[0].content.contains("bob")); - assert!(chunks[0].content.contains("Paragraph one.")); - assert!(chunks[0].content.contains("Paragraph two.")); - assert!(!chunks[0].partial_message); + assert!(chunks[1].content.contains("bob")); + assert!(chunks[1].content.contains("Paragraph one.")); + assert!(chunks[1].content.contains("Paragraph two.")); + assert!(chunks.iter().all(|chunk| !chunk.partial_message)); } #[test] @@ -93,7 +93,7 @@ fn chat_messages_split_at_boundary_when_large() { } #[test] -fn email_threads_pack_into_one_chunk_when_small() { +fn email_threads_keep_one_chunk_per_message_when_small() { let md = "---\nFrom: alice@example.com\nSubject: Hello\nDate: 2026-01-01T00:00:00Z\n\nFirst body.\n---\nFrom: bob@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:01:00Z\n\nSecond body.\n---\nFrom: carol@example.com\nSubject: Re: Hello\nDate: 2026-01-01T00:02:00Z\n\nThird body.".to_string(); let input = ChunkerInput { source_kind: SourceKind::Email, @@ -104,13 +104,13 @@ fn email_threads_pack_into_one_chunk_when_small() { let chunks = chunk_markdown(&input, &ChunkerOptions::default()); assert_eq!( chunks.len(), - 1, - "three small emails should pack into one chunk; got {chunks:?}" + 3, + "each email must keep a stable replay boundary; got {chunks:?}" ); assert!(chunks[0].content.contains("First body.")); - assert!(chunks[0].content.contains("Second body.")); - assert!(chunks[0].content.contains("Third body.")); - assert!(!chunks[0].partial_message); + assert!(chunks[1].content.contains("Second body.")); + assert!(chunks[2].content.contains("Third body.")); + assert!(chunks.iter().all(|chunk| !chunk.partial_message)); } #[test] @@ -159,7 +159,7 @@ fn oversize_single_email_splits_with_partial_flag() { } #[test] -fn packed_units_joined_by_double_newline() { +fn chat_units_keep_independent_chunks() { let md = "## 2026-01-01T00:00:00Z — alice\nfoo\n\n## 2026-01-01T00:01:00Z — bob\nbar".to_string(); let input = ChunkerInput { @@ -169,12 +169,32 @@ fn packed_units_joined_by_double_newline() { metadata: meta(), }; let chunks = chunk_markdown(&input, &ChunkerOptions::default()); - assert_eq!(chunks.len(), 1); - assert!( - chunks[0].content.contains("\n\n"), - "packed units must be joined by \\n\\n; content={:?}", - chunks[0].content - ); + assert_eq!(chunks.len(), 2); + assert!(chunks[0].content.contains("alice")); + assert!(chunks[1].content.contains("bob")); +} + +#[test] +fn overlapping_chat_batches_reuse_message_chunk_ids() { + let first = ChunkerInput { + source_kind: SourceKind::Chat, + source_id: "slack:eng".into(), + markdown: "## 2026-01-01T00:00:00Z — alice\none\n\n## 2026-01-01T00:01:00Z — bob\ntwo" + .into(), + metadata: meta(), + }; + let second = ChunkerInput { + source_kind: SourceKind::Chat, + source_id: "slack:eng".into(), + markdown: "## 2026-01-01T00:01:00Z — bob\ntwo\n\n## 2026-01-01T00:02:00Z — carol\nthree" + .into(), + metadata: meta(), + }; + let first_chunks = chunk_markdown(&first, &ChunkerOptions::default()); + let second_chunks = chunk_markdown(&second, &ChunkerOptions::default()); + + assert_eq!(first_chunks[1].content, second_chunks[0].content); + assert_eq!(first_chunks[1].id, second_chunks[0].id); } #[test] diff --git a/src/memory/chunks/raw_refs.rs b/src/memory/chunks/raw_refs.rs index 50f32dd..14d02d5 100644 --- a/src/memory/chunks/raw_refs.rs +++ b/src/memory/chunks/raw_refs.rs @@ -9,9 +9,8 @@ //! *only* copy of the message body — the SQLite `content` column holds just //! the preview. This module only reads/writes the pointers; it does not //! delete the archive files themselves. Chunk deletion -//! ([`super::store_delete`]) currently never parses `raw_refs_json` to remove -//! the files it points at, so deleting an email chunk leaves its raw body on -//! disk indefinitely (audit finding SC-7). +//! ([`super::store_delete`]) owns reachability-based cleanup of unreferenced +//! raw files and their ingest-gate rows. use anyhow::{Context, Result}; use rusqlite::{params, OptionalExtension, Transaction}; @@ -142,14 +141,10 @@ pub fn list_chunk_raw_ref_paths_with_prefix( /// Return both `content_path` and `content_sha256` stored in SQLite for /// `chunk_id`. /// -/// # Gotcha (audit finding SC-18) -/// Returns `None` only when the row is missing or either column is SQL NULL. -/// Some chunk kinds (email — see the module doc) store `content_path` / -/// `content_sha256` as an **empty string**, not NULL, because their body -/// lives entirely in the raw archive rather than an MD content file. For -/// those rows this function returns `Some(("", ""))`, not `None` — callers -/// that treat `Some(_)` as "there is a readable content file at this path" -/// will mis-handle that case; check for a non-empty path before using it. +/// Returns `None` when the row is missing, either column is SQL NULL, or the +/// path is empty. Email chunks intentionally use empty pointer strings because +/// their bodies live in the raw archive; those are not readable content-file +/// pointers. /// /// # Errors /// Returns `Err` only if the underlying query fails. @@ -169,15 +164,14 @@ pub fn get_chunk_content_pointers( }, ) .optional()?; - Ok(row.and_then(|(p, s)| p.zip(s))) + Ok(row.and_then(|(p, s)| p.zip(s).filter(|(path, _)| !path.is_empty()))) }) } /// Return the `content_path` stored in SQLite for `chunk_id`, if any. /// -/// Same empty-string-vs-NULL gotcha as [`get_chunk_content_pointers`]: a -/// `Some(String::new())` result means the column is set to `""`, not that a -/// content file exists at that path. +/// Empty strings are normalised to `None`, matching +/// [`get_chunk_content_pointers`]. /// /// # Errors /// Returns `Err` only if the underlying query fails. @@ -191,14 +185,13 @@ pub fn get_chunk_content_path(config: &MemoryConfig, chunk_id: &str) -> Result bool { || msg.contains("truncate file") } +/// Whether an error chain reports a corrupt or non-database SQLite image. +pub fn is_corrupt_error(err: &anyhow::Error) -> bool { + err.chain().any(|cause| { + cause + .downcast_ref::() + .is_some_and(|error| { + matches!( + error, + rusqlite::Error::SqliteFailure(code, _) + if matches!( + code.code, + rusqlite::ErrorCode::DatabaseCorrupt + | rusqlite::ErrorCode::NotADatabase + ) + ) + }) + }) +} + /// Clean up WAL/SHM side-files that can block a clean DB open after a crash, /// **without ever discarding committed data**. /// @@ -227,51 +244,16 @@ fn quick_check_ok(db_path: &Path) -> Result { Ok(result.eq_ignore_ascii_case("ok")) } -/// Recover from a `SQLITE_CORRUPT` (malformed image) on the chunk DB. -/// -/// Quarantines the damaged file (and its WAL/SHM side-files) to a timestamped -/// `.corrupt-` copy — preserved, not deleted — then rebuilds an empty -/// schema so the store resumes instead of wedging indefinitely. -/// -/// Returns `Ok(true)` when a quarantine + rebuild happened, `Ok(false)` when a -/// fresh `PRAGMA quick_check` now passes (the earlier failure was transient), -/// and `Err` when the quarantine rename or the schema rebuild failed. -/// -/// # NOTE — not wired into any production error path (audit finding SC-5) -/// `#[allow(dead_code)]` here is not incidental: nothing in `connection.rs` -/// currently calls this on a `SQLITE_CORRUPT` failure, so a real corruption -/// event wedges `get_or_init_connection` (via the circuit breaker) rather -/// than triggering this recovery. Wiring it in requires care beyond adding a -/// call site: the two-step "drop cached connection, then rename" sequence -/// below does not hold any lock across the gap, so a concurrent -/// `with_connection` call for the same path can reopen and re-cache the -/// about-to-be-quarantined file between step 1 and step 3 — its writes then -/// land in the file this function is about to rename out from under it, and -/// step 4's fresh `get_or_init_connection` call returns that same stale -/// cached `Arc` instead of a connection to the newly rebuilt schema. A safe -/// wiring needs the per-path init lock held across the whole -/// quarantine-and-rebuild sequence. +/// Reconfirm corruption and quarantine the main database plus side files. /// -/// # Errors -/// Returns `Err` if quarantining any of the main/`-wal`/`-shm` files fails -/// (e.g. permissions, or a lingering file handle keeping the rename from -/// succeeding on platforms with mandatory file locking), or if rebuilding the -/// schema via [`get_or_init_connection`] fails. -#[allow(dead_code)] -pub fn recover_corrupt_db(config: &MemoryConfig) -> Result { +/// The caller must hold the connection cache's per-path initialization lock +/// and must have removed the cached connection before calling this helper. +pub(super) fn quarantine_corrupt_files(config: &MemoryConfig) -> Result { let db_path = db_path_for(config); - - // 1. Drop any cached (corrupt) connection + breaker so the OS file handle - // is closed before we rename. - drop_cached_connection(config); - - // 2. Re-confirm corruption against the on-disk file. If `quick_check` now - // passes, the image is actually healthy — don't destroy good data. if db_path.exists() && matches!(quick_check_ok(&db_path), Ok(true)) { return Ok(false); } - // 3. Quarantine the main DB + WAL/SHM side-files to `.corrupt-`. let ts = Utc::now().format("%Y%m%dT%H%M%SZ").to_string(); for suffix in &["", "-wal", "-shm"] { let src = with_name_suffix(&db_path, suffix); @@ -287,14 +269,28 @@ pub fn recover_corrupt_db(config: &MemoryConfig) -> Result { ) })?; } - - // 4. Rebuild an empty schema by forcing a fresh open. - get_or_init_connection(config) - .context("failed to rebuild chunk DB schema after quarantining corrupt DB")?; - Ok(true) } +/// Recover from a `SQLITE_CORRUPT` (malformed image) on the chunk DB. +/// +/// Quarantines the damaged file (and its WAL/SHM side-files) to a timestamped +/// `.corrupt-` copy — preserved, not deleted — then rebuilds an empty +/// schema so the store resumes instead of wedging indefinitely. +/// +/// Returns `Ok(true)` when a quarantine + rebuild happened, `Ok(false)` when a +/// fresh `PRAGMA quick_check` now passes (the earlier failure was transient), +/// and `Err` when the quarantine rename or the schema rebuild failed. +/// +/// # Errors +/// Returns `Err` if quarantining any of the main/`-wal`/`-shm` files fails +/// (e.g. permissions, or a lingering file handle keeping the rename from +/// succeeding on platforms with mandatory file locking), or if rebuilding the +/// schema initialization fails. +pub fn recover_corrupt_db(config: &MemoryConfig) -> Result { + recover_corrupt_connection(config) +} + #[cfg(test)] #[path = "recovery_tests.rs"] mod tests; diff --git a/src/memory/chunks/recovery_tests.rs b/src/memory/chunks/recovery_tests.rs index ffceeb5..6cfd33c 100644 --- a/src/memory/chunks/recovery_tests.rs +++ b/src/memory/chunks/recovery_tests.rs @@ -1,6 +1,6 @@ //! Unit tests for corrupt-DB recovery (`super`). -use super::super::connection::{clear_connection_cache, with_connection}; +use super::super::connection::{clear_connection_cache_for, with_connection}; use super::super::db_path_for; use super::recover_corrupt_db; use crate::memory::config::MemoryConfig; @@ -17,7 +17,6 @@ fn corrupt_test_config() -> (TempDir, MemoryConfig) { /// a fresh, queryable schema so the store resumes. #[test] fn recover_corrupt_db_quarantines_and_rebuilds() { - clear_connection_cache(); let (_tmp, cfg) = corrupt_test_config(); let db_path = db_path_for(&cfg); std::fs::create_dir_all(db_path.parent().unwrap()).unwrap(); @@ -41,7 +40,7 @@ fn recover_corrupt_db_quarantines_and_rebuilds() { "exactly one quarantined copy should exist" ); - clear_connection_cache(); + clear_connection_cache_for(&cfg); let count: i64 = with_connection(&cfg, |conn| { conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| r.get(0)) .context("count jobs") @@ -50,6 +49,26 @@ fn recover_corrupt_db_quarantines_and_rebuilds() { assert_eq!(count, 0, "rebuilt jobs table starts empty"); } +#[test] +fn cold_open_automatically_quarantines_corrupt_database() { + let (_tmp, cfg) = corrupt_test_config(); + let db_path = db_path_for(&cfg); + std::fs::create_dir_all(db_path.parent().unwrap()).unwrap(); + std::fs::write(&db_path, b"not a database").unwrap(); + + let count: i64 = with_connection(&cfg, |conn| { + conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |row| row.get(0)) + .map_err(anyhow::Error::from) + }) + .expect("cold open should recover and initialize schema"); + + assert_eq!(count, 0); + assert!(std::fs::read_dir(db_path.parent().unwrap()) + .unwrap() + .filter_map(Result::ok) + .any(|entry| entry.file_name().to_string_lossy().contains(".corrupt-"))); +} + /// A legacy DB still in WAL mode can hold committed-but-uncheckpointed /// transactions *only* in its `-wal` side-file. The stale-file cleanup path /// (run on a cold-start I/O open error) must fold that data back into the main @@ -109,7 +128,6 @@ fn cleanup_preserves_committed_wal_data() { /// preserved and recovery is a no-op returning `Ok(false)`. #[test] fn recover_corrupt_db_is_noop_on_healthy_db() { - clear_connection_cache(); let (_tmp, cfg) = corrupt_test_config(); with_connection(&cfg, |conn| { conn.query_row("SELECT COUNT(*) FROM mem_tree_jobs", [], |r| { diff --git a/src/memory/chunks/store.rs b/src/memory/chunks/store.rs index f9a8713..3e1b2b4 100644 --- a/src/memory/chunks/store.rs +++ b/src/memory/chunks/store.rs @@ -16,9 +16,6 @@ use super::connection::with_connection; use super::types::{Chunk, Metadata, SourceKind, SourceRef, StagedChunk}; use crate::memory::config::MemoryConfig; -pub(super) const DEFAULT_LIST_LIMIT: usize = 100; -pub(super) const MAX_LIST_LIMIT: usize = 10_000; - /// Chunk lifecycle: freshly persisted, awaiting the async extract job. pub const CHUNK_STATUS_PENDING_EXTRACTION: &str = "pending_extraction"; /// Chunk lifecycle: extract ran and the chunk passed admission. @@ -41,16 +38,9 @@ pub const CHUNK_STATUS_DROPPED: &str = "dropped"; /// /// Returns `Ok(0)` immediately (no DB access) for an empty `chunks` slice. /// -/// # Gotcha (audit finding SC-17) -/// This `ON CONFLICT` clause only overwrites the plain-content columns; it -/// never touches `content_path`, `content_sha256`, or `lifecycle_status`. If -/// a chunk id was previously staged via `upsert_staged_chunks_tx` (so it has -/// a `content_path`) or was marked `dropped` by the admission gate, calling -/// this function again for the same id leaves those columns exactly as they -/// were — it does not clear a stale `content_path` pointing at now-orphaned -/// content, and it cannot resurrect a `dropped` chunk back to an admitted -/// state. Use the staged path / [`super::store_sources::set_chunk_lifecycle_status`] -/// explicitly when either of those needs to change. +/// Replacing a previously staged row clears its file pointer and restores the +/// lifecycle to `admitted`, so the row's storage representation and lifecycle +/// match this plain-content write. /// /// # Errors /// Returns `Err` if beginning the transaction, preparing/executing @@ -98,7 +88,10 @@ const UPSERT_SQL: &str = "INSERT INTO mem_tree_chunks ( content = excluded.content, token_count = excluded.token_count, seq_in_source = excluded.seq_in_source, - created_at_ms = excluded.created_at_ms"; + created_at_ms = excluded.created_at_ms, + content_path = NULL, + content_sha256 = NULL, + lifecycle_status = 'admitted'"; /// Bind and execute `UPSERT_SQL` once per chunk against an already-prepared /// statement. Split out from [`upsert_chunks`] so the statement is prepared @@ -257,7 +250,7 @@ pub fn list_source_ids_with_prefix( }) } -const SELECT_COLUMNS: &str = "id, source_kind, source_id, path_scope, source_ref, owner, +pub(super) const SELECT_COLUMNS: &str = "id, source_kind, source_id, path_scope, source_ref, owner, timestamp_ms, time_range_start_ms, time_range_end_ms, tags_json, content, token_count, seq_in_source, created_at_ms"; @@ -326,109 +319,6 @@ pub fn get_chunks_batch( }) } -/// Query parameters for [`list_chunks`]. All fields are optional filters — -/// callers pass `ListChunksQuery::default()` to get recent-across-everything. -#[derive(Debug, Default, Clone)] -pub struct ListChunksQuery { - /// Restrict to one source kind. - pub source_kind: Option, - /// Restrict to one logical source id. - pub source_id: Option, - /// Restrict to one owner. - pub owner: Option, - /// Inclusive lower bound on `timestamp` (milliseconds since epoch). - pub since_ms: Option, - /// Inclusive upper bound on `timestamp` (milliseconds since epoch). - pub until_ms: Option, - /// Max rows to return (default 100 when `None`). - pub limit: Option, - /// Per-profile memory-source allowlist. When `Some`, memory-source chunks - /// whose source identifier is not in the set are dropped *before* the row - /// limit is applied. Non-source chunks always pass. - pub source_scope: Option>, - /// When `true`, rows the admission gate rejected (`lifecycle_status = - /// 'dropped'`) are excluded. - pub exclude_dropped: bool, -} - -/// List chunks matching the provided filters, ordered by `timestamp_ms` DESC, -/// then `seq_in_source` ASC as a tiebreaker. -/// -/// # Gotcha (audit finding SC-15) -/// When `query.source_scope` is `Some`, the allowlist filter has to run in -/// Rust *after* the SQL fetch (SQLite doesn't know about the memory-source -/// tag semantics), so this fetches up to `MAX_LIST_LIMIT` (10,000) -/// candidate rows from SQL — ordered newest-first — before filtering and -/// truncating to `query.limit` in Rust. If a workspace has more than 10,000 -/// chunks newer than the allowed ones, valid, in-scope rows past that SQL -/// cutoff are silently dropped and never reach the Rust-side filter, even -/// though they would have passed it. This only affects scoped queries; an -/// unscoped `list_chunks` call applies `query.limit` directly in SQL and has -/// no such gap. -/// -/// # Errors -/// Returns `Err` if the query fails to prepare/execute or any row fails to -/// decode (see `row_to_chunk`). -pub fn list_chunks(config: &MemoryConfig, query: &ListChunksQuery) -> Result> { - with_connection(config, |conn| { - let mut sql = format!("SELECT {SELECT_COLUMNS} FROM mem_tree_chunks WHERE 1=1"); - let mut bound: Vec> = Vec::new(); - - if let Some(kind) = query.source_kind { - sql.push_str(" AND source_kind = ?"); - bound.push(Box::new(kind.as_str().to_string())); - } - if let Some(ref source_id) = query.source_id { - sql.push_str(" AND source_id = ?"); - bound.push(Box::new(source_id.clone())); - } - if let Some(ref owner) = query.owner { - sql.push_str(" AND owner = ?"); - bound.push(Box::new(owner.clone())); - } - if let Some(since_ms) = query.since_ms { - sql.push_str(" AND timestamp_ms >= ?"); - bound.push(Box::new(since_ms)); - } - if let Some(until_ms) = query.until_ms { - sql.push_str(" AND timestamp_ms <= ?"); - bound.push(Box::new(until_ms)); - } - if query.exclude_dropped { - sql.push_str(" AND lifecycle_status != ?"); - bound.push(Box::new(CHUNK_STATUS_DROPPED.to_string())); - } - let requested_limit = normalized_limit(query.limit); - // When a profile source-scope is active, fetch a wider candidate set and - // apply the gate in Rust *before* truncating, so a disallowed-source - // prefix can't push permitted rows past the requested limit. - let sql_limit = if query.source_scope.is_some() { - MAX_LIST_LIMIT as i64 - } else { - requested_limit - }; - sql.push_str(" ORDER BY timestamp_ms DESC, seq_in_source ASC LIMIT ?"); - bound.push(Box::new(sql_limit)); - - let mut stmt = conn.prepare(&sql)?; - let param_refs: Vec<&dyn rusqlite::ToSql> = bound - .iter() - .map(|b| b.as_ref() as &dyn rusqlite::ToSql) - .collect(); - let mut rows = stmt - .query_map(param_refs.as_slice(), row_to_chunk)? - .collect::>>() - .context("Failed to collect chunks")?; - if let Some(ref allowed) = query.source_scope { - rows.retain(|c| { - super::chunk_source_allowed_in(allowed, &c.metadata.tags, &c.metadata.source_id) - }); - rows.truncate(requested_limit as usize); - } - Ok(rows) - }) -} - /// Count total chunks in the store (no filters — every row in /// `mem_tree_chunks`, regardless of lifecycle status). /// @@ -475,16 +365,6 @@ pub fn extraction_coverage(config: &MemoryConfig) -> Result { }) } -/// Clamp a caller-requested row limit to `[1, MAX_LIST_LIMIT]`, defaulting to -/// [`DEFAULT_LIST_LIMIT`] when `requested` is `None`. Never returns 0 (a -/// caller-requested limit of 0 is treated as 1, not "no rows"). -pub(super) fn normalized_limit(requested: Option) -> i64 { - let clamped = requested - .unwrap_or(DEFAULT_LIST_LIMIT) - .clamp(1, MAX_LIST_LIMIT); - i64::try_from(clamped).unwrap_or(MAX_LIST_LIMIT as i64) -} - /// Decode one `mem_tree_chunks` row (columns in exactly [`SELECT_COLUMNS`] /// order) into a [`Chunk`]. Shared by every plain-chunk query in this module. /// diff --git a/src/memory/chunks/store_conn_tests.rs b/src/memory/chunks/store_conn_tests.rs index e06bf70..c6a4af9 100644 --- a/src/memory/chunks/store_conn_tests.rs +++ b/src/memory/chunks/store_conn_tests.rs @@ -3,12 +3,11 @@ //! embedding / delete / migration accessors against a tempdir-backed SQLite //! store. //! -//! Because the connection cache is a process-level singleton, tests that -//! exercise cache behaviour call `clear_connection_cache()` at the start, or -//! use unique tempdirs that cannot collide with other tests. +//! Tests use unique tempdirs so parallel cache checks cannot collide. Tests +//! that must force a reopen reset only their own workspace path. use super::connection::{ - clear_connection_cache, get_or_init_connection, invalidate_connection, + clear_connection_cache_for, get_or_init_connection, invalidate_connection, schema_apply_count_for_path_for_tests, with_connection, CB_THRESHOLD, }; use super::embeddings::{active_embedding_dims, embedding_to_blob}; @@ -227,7 +226,6 @@ fn legacy_embeddings_migrate_to_sidecar_once() { #[test] fn connection_cache_returns_same_arc_for_same_workspace() { - clear_connection_cache(); let (_tmp, cfg) = test_config(); let arc1 = get_or_init_connection(&cfg).expect("first get_or_init"); @@ -240,7 +238,6 @@ fn connection_cache_returns_same_arc_for_same_workspace() { #[test] fn connection_cache_uses_separate_connections_for_different_workspaces() { - clear_connection_cache(); let (_tmp1, cfg1) = test_config(); let (_tmp2, cfg2) = test_config(); @@ -259,7 +256,6 @@ fn connection_cache_uses_separate_connections_for_different_workspaces() { #[test] fn circuit_breaker_trips_after_threshold() { - clear_connection_cache(); let tmp = TempDir::new().expect("tempdir"); // Create a regular file where the memory_tree *directory* would be. @@ -344,7 +340,7 @@ fn existing_wal_db_migrates_to_truncate() { .expect("seed"); } - clear_connection_cache(); + clear_connection_cache_for(&cfg); with_connection(&cfg, |conn| { let mode: String = conn.query_row("PRAGMA journal_mode", [], |r| r.get(0))?; assert!( diff --git a/src/memory/chunks/store_delete.rs b/src/memory/chunks/store_delete.rs index 959071e..4f75d00 100644 --- a/src/memory/chunks/store_delete.rs +++ b/src/memory/chunks/store_delete.rs @@ -6,21 +6,14 @@ //! subsystem is not ported here); it deletes only the chunk-owned rows and //! files. //! -//! ## Deletion completeness gaps (audit findings SC-7, SC-20) -//! - Only files at `content_path` are removed. A chunk's `raw_refs_json` -//! pointers (the raw-archive mirror — the *only* copy of the body for -//! email chunks, see [`super::raw_refs`]) are never parsed here, so their -//! files are never deleted, and no reachability check runs to see whether -//! another surviving chunk still references the same raw file. Deleting an -//! email account's chunks leaves every message body on disk. -//! - Matching `RAW_FILE_GATE_KIND` rows in `mem_tree_ingested_sources` (the -//! raw-archive ingest gate — see [`super::store::RAW_FILE_GATE_KIND`]) are -//! never cleared alongside the deleted chunks. -//! - Every public entry point loads every chunk row for the source kind into -//! memory and filters in Rust, then issues five `DELETE` statements per -//! matched chunk — `O(N·M)` round-trips for `N` chunks touched across `M` -//! dependent tables, rather than a single set-based `DELETE ... WHERE id IN -//! (subquery)` per table. +//! Raw archives are cascade-cleaned by reachability: a file and its raw-file +//! ingest gate are removed only after the transaction confirms no surviving +//! chunk references that path. A malformed surviving pointer row makes cleanup +//! conservative and preserves all candidate raw files. +//! +//! Selection and dependent-row cleanup are set-based in SQLite. Rust only +//! materialises the matched rows whose filesystem pointers and source-tree +//! scopes must be processed after the database transaction commits. use anyhow::{Context, Result}; use rusqlite::params; @@ -28,6 +21,8 @@ use std::collections::HashSet; use super::connection::with_connection; use super::content_root; +use super::raw_refs::RawRef; +use super::store_sources::RAW_FILE_GATE_KIND; use super::types::SourceKind; use crate::memory::config::MemoryConfig; @@ -46,12 +41,7 @@ pub fn delete_chunks_by_source( source_kind: SourceKind, source_id: &str, ) -> Result { - delete_chunks_by_source_filter( - config, - source_kind, - |candidate, _owner| candidate == source_id, - |candidate| candidate == source_id, - ) + delete_chunks_by_source_filter(config, source_kind, DeleteFilter::ExactSource(source_id)) } /// Delete all chunk rows whose source id starts with `source_id_prefix`. @@ -69,8 +59,7 @@ pub fn delete_chunks_by_source_prefix( delete_chunks_by_source_filter( config, source_kind, - |candidate, _owner| candidate.starts_with(source_id_prefix), - |candidate| candidate.starts_with(source_id_prefix), + DeleteFilter::SourcePrefix(source_id_prefix), ) } @@ -90,25 +79,40 @@ pub fn delete_chunks_by_owner( source_kind: SourceKind, owner: &str, ) -> Result { - delete_chunks_by_source_filter( - config, - source_kind, - |_source_id, candidate_owner| candidate_owner == owner, - |_source_id| false, - ) + delete_chunks_by_source_filter(config, source_kind, DeleteFilter::Owner(owner)) +} + +#[derive(Clone, Copy)] +enum DeleteFilter<'a> { + ExactSource(&'a str), + SourcePrefix(&'a str), + Owner(&'a str), +} + +impl<'a> DeleteFilter<'a> { + fn value(self) -> &'a str { + match self { + Self::ExactSource(value) | Self::SourcePrefix(value) | Self::Owner(value) => value, + } + } + + fn chunk_predicate(self) -> &'static str { + match self { + Self::ExactSource(_) => "source_id = ?2", + Self::SourcePrefix(_) => "substr(source_id, 1, length(?2)) = ?2", + Self::Owner(_) => "owner = ?2", + } + } } /// Shared implementation behind [`delete_chunks_by_source`], /// [`delete_chunks_by_source_prefix`], and [`delete_chunks_by_owner`]. /// -/// Loads every chunk row for `source_kind`, keeps those where `matches_chunk` -/// (given `(source_id, owner)`) returns `true`, then — inside one -/// transaction — deletes each matched chunk's score/entity-index/embedding/ -/// reembed-skip rows before the chunk row itself, and finally removes any -/// `mem_tree_ingested_sources` row whose `source_id` either satisfies -/// `matches_ingested_source` directly or became fully orphaned (zero -/// remaining chunks) as a result of this delete. On-disk `content_path` -/// files are collected during the transaction but removed only after it +/// Selects matching rows into a temporary SQLite table, then deletes each +/// dependent table with one set-based statement before removing the chunk +/// rows. Ingest gates selected directly or made fully orphaned are removed in +/// the same transaction. On-disk `content_path` and unreferenced +/// `raw_refs_json` files are collected during the transaction but removed only after it /// commits, via [`remove_chunk_content_files`] (best-effort; a filesystem /// failure there does not roll back or fail the DB-side delete, and does not /// surface as an `Err` from this function). @@ -121,121 +125,159 @@ pub fn delete_chunks_by_owner( fn delete_chunks_by_source_filter( config: &MemoryConfig, source_kind: SourceKind, - matches_chunk: impl Fn(&str, &str) -> bool, - matches_ingested_source: impl Fn(&str) -> bool, + filter: DeleteFilter<'_>, ) -> Result { let mut content_paths = Vec::new(); let deleted = with_connection(config, |conn| { let tx = conn.unchecked_transaction()?; + tx.execute_batch( + "CREATE TEMP TABLE IF NOT EXISTS mem_tree_delete_selection ( + id TEXT PRIMARY KEY, + source_id TEXT NOT NULL, + path_scope TEXT, + content_path TEXT, + raw_refs_json TEXT + ); + DELETE FROM mem_tree_delete_selection;", + )?; + let selection_sql = format!( + "INSERT INTO mem_tree_delete_selection + (id, source_id, path_scope, content_path, raw_refs_json) + SELECT id, source_id, path_scope, content_path, raw_refs_json + FROM mem_tree_chunks + WHERE source_kind = ?1 AND {}", + filter.chunk_predicate() + ); + tx.execute( + &selection_sql, + params![source_kind.as_str(), filter.value()], + )?; + let chunks = { let mut stmt = tx.prepare( - "SELECT id, source_id, owner, content_path, path_scope - FROM mem_tree_chunks - WHERE source_kind = ?1", + "SELECT id, source_id, content_path, path_scope, raw_refs_json + FROM mem_tree_delete_selection", )?; - let rows = stmt.query_map(params![source_kind.as_str()], |row| { + let rows = stmt.query_map([], |row| { Ok(( row.get::<_, String>(0)?, row.get::<_, String>(1)?, - row.get::<_, String>(2)?, + row.get::<_, Option>(2)?, row.get::<_, Option>(3)?, row.get::<_, Option>(4)?, )) })?; - rows.filter_map(|row| match row { - Ok((id, source_id, owner, content_path, path_scope)) - if matches_chunk(&source_id, &owner) => - { - Some(Ok((id, source_id, content_path, path_scope))) - } - Ok(_) => None, - Err(error) => Some(Err(error)), - }) - .collect::>>() - .context("Failed to collect chunks by source")? + rows.collect::>>() + .context("Failed to collect chunks by source")? }; - let deleted_source_ids: HashSet = chunks - .iter() - .map(|(_, source_id, _, _)| source_id.clone()) - .collect(); let deleted_tree_scopes: HashSet = chunks .iter() - .map(|(_, source_id, _, path_scope)| { + .map(|(_, source_id, _, path_scope, _)| { path_scope.clone().unwrap_or_else(|| source_id.clone()) }) .collect(); - for (chunk_id, _source_id, content_path, _path_scope) in &chunks { - tx.execute( - "DELETE FROM mem_tree_score WHERE chunk_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_chunk_embeddings WHERE chunk_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_chunk_reembed_skipped WHERE chunk_id = ?1", - params![chunk_id], - )?; - tx.execute( - "DELETE FROM mem_tree_chunks WHERE id = ?1", - params![chunk_id], - )?; + let raw_path_candidates: HashSet = chunks + .iter() + .filter_map(|(_, _, _, _, json)| json.as_deref()) + .filter_map(|json| serde_json::from_str::>(json).ok()) + .flatten() + .map(|raw_ref| raw_ref.path) + .collect(); + + for (_chunk_id, _source_id, content_path, _path_scope, _raw_refs_json) in &chunks { if let Some(path) = content_path.as_ref().filter(|path| !path.is_empty()) { content_paths.push(path.clone()); } } + tx.execute( + "DELETE FROM mem_tree_score + WHERE chunk_id IN (SELECT id FROM mem_tree_delete_selection)", + [], + )?; + tx.execute( + "DELETE FROM mem_tree_entity_index + WHERE node_id IN (SELECT id FROM mem_tree_delete_selection)", + [], + )?; + tx.execute( + "DELETE FROM mem_tree_chunk_embeddings + WHERE chunk_id IN (SELECT id FROM mem_tree_delete_selection)", + [], + )?; + tx.execute( + "DELETE FROM mem_tree_chunk_reembed_skipped + WHERE chunk_id IN (SELECT id FROM mem_tree_delete_selection)", + [], + )?; + tx.execute( + "DELETE FROM mem_tree_chunks + WHERE id IN (SELECT id FROM mem_tree_delete_selection)", + [], + )?; - // A fully-orphaned source (no chunks left) has its ingest gate removed. - let mut orphaned_deleted_sources = HashSet::new(); - for source_id in &deleted_source_ids { - let remaining: i64 = tx.query_row( - "SELECT COUNT(*) - FROM mem_tree_chunks - WHERE source_kind = ?1 AND source_id = ?2", - params![source_kind.as_str(), source_id], - |row| row.get(0), - )?; - if remaining == 0 { - orphaned_deleted_sources.insert(source_id.clone()); - } - } - - let ingested_sources = { + // Clean raw files only when every surviving pointer row is readable. + // Otherwise an unknown reference could make deleting a candidate file + // destructive, so fail closed and leave the archive untouched. + let mut surviving_raw_paths = HashSet::new(); + let mut has_corrupt_surviving_refs = false; + { let mut stmt = tx.prepare( - "SELECT source_id - FROM mem_tree_ingested_sources - WHERE source_kind = ?1", + "SELECT raw_refs_json FROM mem_tree_chunks + WHERE raw_refs_json IS NOT NULL AND raw_refs_json != ''", )?; - let rows = - stmt.query_map(params![source_kind.as_str()], |row| row.get::<_, String>(0))?; - rows.filter_map(|row| match row { - Ok(source_id) - if matches_ingested_source(&source_id) - || orphaned_deleted_sources.contains(&source_id) => - { - Some(Ok(source_id)) + let rows = stmt.query_map([], |row| row.get::<_, String>(0))?; + for row in rows { + match serde_json::from_str::>(&row?) { + Ok(refs) => surviving_raw_paths.extend(refs.into_iter().map(|r| r.path)), + Err(_) => has_corrupt_surviving_refs = true, } - Ok(_) => None, - Err(error) => Some(Err(error)), - }) - .collect::>>() - .context("Failed to collect ingested sources")? - }; + } + } + if !has_corrupt_surviving_refs { + for path in raw_path_candidates.difference(&surviving_raw_paths) { + tx.execute( + "DELETE FROM mem_tree_ingested_sources + WHERE source_kind = ?1 AND source_id = ?2", + params![RAW_FILE_GATE_KIND, path], + )?; + content_paths.push(path.clone()); + } + } - for source_id in &ingested_sources { - tx.execute( - "DELETE FROM mem_tree_ingested_sources - WHERE source_kind = ?1 AND source_id = ?2", - params![source_kind.as_str(), source_id], - )?; + // A fully-orphaned source (no chunks left) has its ingest gate removed. + tx.execute( + "DELETE FROM mem_tree_ingested_sources AS gate + WHERE gate.source_kind = ?1 + AND gate.source_id IN ( + SELECT DISTINCT source_id FROM mem_tree_delete_selection + ) + AND NOT EXISTS ( + SELECT 1 FROM mem_tree_chunks AS chunk + WHERE chunk.source_kind = gate.source_kind + AND chunk.source_id = gate.source_id + )", + params![source_kind.as_str()], + )?; + match filter { + DeleteFilter::ExactSource(source_id) => { + tx.execute( + "DELETE FROM mem_tree_ingested_sources + WHERE source_kind = ?1 AND source_id = ?2", + params![source_kind.as_str(), source_id], + )?; + } + DeleteFilter::SourcePrefix(prefix) => { + tx.execute( + "DELETE FROM mem_tree_ingested_sources + WHERE source_kind = ?1 + AND substr(source_id, 1, length(?2)) = ?2", + params![source_kind.as_str(), prefix], + )?; + } + DeleteFilter::Owner(_) => {} } for scope in &deleted_tree_scopes { diff --git a/src/memory/chunks/store_embed_tests.rs b/src/memory/chunks/store_embed_tests.rs index 8984b97..4d364a8 100644 --- a/src/memory/chunks/store_embed_tests.rs +++ b/src/memory/chunks/store_embed_tests.rs @@ -3,12 +3,11 @@ //! embedding / delete / migration accessors against a tempdir-backed SQLite //! store. //! -//! Because the connection cache is a process-level singleton, tests that -//! exercise cache behaviour call `clear_connection_cache()` at the start, or -//! use unique tempdirs that cannot collide with other tests. +//! Tests use unique tempdirs so parallel cache checks cannot collide. Tests +//! that must force a reopen reset only their own workspace path. use super::connection::{ - clear_connection_cache, get_or_init_connection, invalidate_connection, + clear_connection_cache_for, get_or_init_connection, invalidate_connection, schema_apply_count_for_path_for_tests, with_connection, CB_THRESHOLD, }; use super::embeddings::{active_embedding_dims, embedding_to_blob}; diff --git a/src/memory/chunks/store_list.rs b/src/memory/chunks/store_list.rs new file mode 100644 index 0000000..5539c43 --- /dev/null +++ b/src/memory/chunks/store_list.rs @@ -0,0 +1,121 @@ +//! Filtered and paginated chunk listing. + +use anyhow::{Context, Result}; + +use super::connection::with_connection; +use super::store::{row_to_chunk, CHUNK_STATUS_DROPPED, SELECT_COLUMNS}; +use super::types::{Chunk, SourceKind}; +use crate::memory::config::MemoryConfig; + +const DEFAULT_LIST_LIMIT: usize = 100; +const MAX_LIST_LIMIT: usize = 10_000; + +/// Optional filters and pagination for [`list_chunks`]. +#[derive(Debug, Default, Clone)] +pub struct ListChunksQuery { + /// Exact source-kind filter. + pub source_kind: Option, + /// Exact logical source-id filter. + pub source_id: Option, + /// Exact owner filter. + pub owner: Option, + /// Inclusive lower source-time bound in epoch milliseconds. + pub since_ms: Option, + /// Inclusive upper source-time bound in epoch milliseconds. + pub until_ms: Option, + /// Maximum rows, clamped to the store safety cap. + pub limit: Option, + /// Ordered rows to skip for internal pagination. + pub offset: Option, + /// Allowed memory-source identifiers. + pub source_scope: Option>, + /// Exclude lifecycle-dropped chunks. + pub exclude_dropped: bool, +} + +/// List chunks matching all supplied filters in deterministic newest-first +/// order. Source allowlisting is applied in SQL before the row limit. +/// +/// # Errors +/// Returns an error when SQLite preparation, execution, or row decoding fails. +pub fn list_chunks(config: &MemoryConfig, query: &ListChunksQuery) -> Result> { + with_connection(config, |conn| { + let mut sql = format!("SELECT {SELECT_COLUMNS} FROM mem_tree_chunks WHERE 1=1"); + let mut bound: Vec> = Vec::new(); + if let Some(kind) = query.source_kind { + sql.push_str(" AND source_kind = ?"); + bound.push(Box::new(kind.as_str().to_string())); + } + for (clause, value) in [ + (" AND source_id = ?", query.source_id.as_ref()), + (" AND owner = ?", query.owner.as_ref()), + ] { + if let Some(value) = value { + sql.push_str(clause); + bound.push(Box::new(value.clone())); + } + } + if let Some(value) = query.since_ms { + sql.push_str(" AND timestamp_ms >= ?"); + bound.push(Box::new(value)); + } + if let Some(value) = query.until_ms { + sql.push_str(" AND timestamp_ms <= ?"); + bound.push(Box::new(value)); + } + if query.exclude_dropped { + sql.push_str(" AND lifecycle_status != ?"); + bound.push(Box::new(CHUNK_STATUS_DROPPED.to_string())); + } + append_source_scope(&mut sql, &mut bound, query.source_scope.as_ref()); + sql.push_str(" ORDER BY timestamp_ms DESC, seq_in_source ASC, id ASC LIMIT ? OFFSET ?"); + bound.push(Box::new(normalized_limit(query.limit))); + bound.push(Box::new( + i64::try_from(query.offset.unwrap_or(0)).unwrap_or(i64::MAX), + )); + let params = bound + .iter() + .map(|value| value.as_ref() as &dyn rusqlite::ToSql) + .collect::>(); + conn.prepare(&sql)? + .query_map(params.as_slice(), row_to_chunk)? + .collect::>>() + .context("Failed to collect chunks") + }) +} + +fn append_source_scope( + sql: &mut String, + bound: &mut Vec>, + allowed: Option<&std::collections::HashSet>, +) { + let Some(allowed) = allowed else { return }; + sql.push_str( + " AND (NOT EXISTS (SELECT 1 FROM json_each(mem_tree_chunks.tags_json) + WHERE value = 'memory_sources')", + ); + if !allowed.is_empty() { + sql.push_str(" OR ("); + for (index, source_id) in allowed.iter().enumerate() { + if index > 0 { + sql.push_str(" OR "); + } + sql.push_str("source_id = ? OR source_id LIKE ? ESCAPE '\\'"); + bound.push(Box::new(source_id.clone())); + let escaped = source_id + .replace('\\', "\\\\") + .replace('%', "\\%") + .replace('_', "\\_"); + bound.push(Box::new(format!("mem_src:{escaped}:%"))); + } + sql.push(')'); + } + sql.push(')'); +} + +fn normalized_limit(requested: Option) -> i64 { + let limit = requested + .unwrap_or(DEFAULT_LIST_LIMIT) + .clamp(1, MAX_LIST_LIMIT); + i64::try_from(limit).unwrap_or(MAX_LIST_LIMIT as i64) +} diff --git a/src/memory/chunks/store_sources.rs b/src/memory/chunks/store_sources.rs index a0f7f0e..a0f1566 100644 --- a/src/memory/chunks/store_sources.rs +++ b/src/memory/chunks/store_sources.rs @@ -18,9 +18,7 @@ use crate::memory::config::MemoryConfig; /// stores it as-is. /// /// # Errors -/// Returns `Err` only if the underlying `UPDATE` fails. Silently affects zero -/// rows (no error, no signal) if `chunk_id` does not exist — see the NOTE on -/// `set_chunk_lifecycle_status_conn`. +/// Returns `Err` if the update fails or `chunk_id` does not exist. pub fn set_chunk_lifecycle_status( config: &MemoryConfig, chunk_id: &str, @@ -47,15 +45,8 @@ pub fn set_chunk_lifecycle_status_tx( /// Core `UPDATE ... SET lifecycle_status` over an arbitrary `&Connection`. /// -/// # NOTE -/// `changed` (the affected-row count) is computed but never inspected — a -/// call for a nonexistent `chunk_id` silently succeeds with zero rows -/// touched rather than surfacing that as an error or a return value the -/// caller could check. Callers cannot currently distinguish "status set" from -/// "chunk_id did not exist" without a separate existence check. -/// /// # Errors -/// Returns `Err` only if the `UPDATE` statement itself fails. +/// Returns `Err` if the statement fails or no chunk row matches. fn set_chunk_lifecycle_status_conn(conn: &Connection, chunk_id: &str, status: &str) -> Result<()> { let changed = conn.execute( "UPDATE mem_tree_chunks SET lifecycle_status = ?1 WHERE id = ?2", diff --git a/src/memory/chunks/store_tests.rs b/src/memory/chunks/store_tests.rs index 1f30fde..9f9a835 100644 --- a/src/memory/chunks/store_tests.rs +++ b/src/memory/chunks/store_tests.rs @@ -2,7 +2,7 @@ //! Unit tests for the chunk store (`super`): upsert / list / delete. use super::connection::{ - clear_connection_cache, get_or_init_connection, invalidate_connection, + clear_connection_cache_for, get_or_init_connection, invalidate_connection, schema_apply_count_for_path_for_tests, with_connection, CB_THRESHOLD, }; use super::embeddings::{active_embedding_dims, embedding_to_blob}; @@ -115,6 +115,27 @@ fn list_chunks_source_scope_filters_before_limit() { assert_eq!(rows[0].metadata.source_id, "slack:#secret"); } +#[test] +fn list_chunks_source_scope_treats_composite_ids_as_literals() { + let (_tmp, cfg) = test_config(); + let mut allowed_chunk = sample_chunk("mem_src:a_b:item-1", 0, 2_000); + allowed_chunk.metadata.tags = vec!["memory_sources".into()]; + let mut wildcard_lookalike = sample_chunk("mem_src:axb:item-2", 0, 3_000); + wildcard_lookalike.metadata.tags = vec!["memory_sources".into()]; + upsert_chunks(&cfg, &[wildcard_lookalike, allowed_chunk]).unwrap(); + + let rows = list_chunks( + &cfg, + &ListChunksQuery { + source_scope: Some(std::collections::HashSet::from(["a_b".into()])), + ..Default::default() + }, + ) + .unwrap(); + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].metadata.source_id, "mem_src:a_b:item-1"); +} + #[test] fn upsert_is_idempotent() { let (_tmp, cfg) = test_config(); @@ -429,253 +450,5 @@ fn delete_chunks_by_owner_preserves_other_owners_for_same_source() { assert!(!is_source_ingested(&cfg, SourceKind::Chat, "slack:c-1-only").unwrap()); } -#[test] -fn delete_chunks_by_source_removes_safe_content_files_but_rejects_escape_paths() { - let (_tmp, cfg) = test_config(); - let safe = sample_chunk("slack:c-1", 0, 1_700_000_000_000); - let unsafe_chunk = sample_chunk("slack:c-1", 1, 1_700_000_001_000); - upsert_chunks(&cfg, &[safe.clone(), unsafe_chunk.clone()]).unwrap(); - - let root = content_root(&cfg); - let safe_rel = "chunks/safe.md"; - let safe_path = root.join(safe_rel); - std::fs::create_dir_all(safe_path.parent().unwrap()).unwrap(); - std::fs::write(&safe_path, "safe").unwrap(); - - let outside_path = root.parent().unwrap().join("outside.md"); - std::fs::write(&outside_path, "outside").unwrap(); - - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", - params![safe_rel, safe.id], - )?; - conn.execute( - "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", - params!["../outside.md", unsafe_chunk.id], - )?; - Ok(()) - }) - .unwrap(); - - let deleted = delete_chunks_by_source(&cfg, SourceKind::Chat, "slack:c-1").unwrap(); - - assert_eq!(deleted, 2); - assert!(!safe_path.exists()); - assert!(outside_path.exists()); -} - -#[test] -fn raw_refs_round_trip_and_prefix_listing_tolerates_corrupt_rows() { - let (_tmp, cfg) = test_config(); - let c1 = sample_chunk("gmail:thread-1", 0, 1_700_000_000_000); - let c2 = sample_chunk("gmail:thread-2", 0, 1_700_000_001_000); - let corrupt = sample_chunk("gmail:thread-3", 0, 1_700_000_002_000); - upsert_chunks(&cfg, &[c1.clone(), c2.clone(), corrupt.clone()]).unwrap(); - - let refs = vec![ - RawRef { - path: "raw/gmail/thread_1/message_1.md".into(), - start: 10, - end: Some(42), - }, - RawRef { - path: "raw/gmail/thread_1/message_2.md".into(), - start: 0, - end: None, - }, - ]; - set_chunk_raw_refs(&cfg, &c1.id, &refs).unwrap(); - set_chunk_raw_refs( - &cfg, - &c2.id, - &[RawRef { - path: "raw/slack/channel/message.md".into(), - start: 0, - end: None, - }], - ) - .unwrap(); - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_chunks SET raw_refs_json = 'not valid json' WHERE id = ?1", - params![corrupt.id], - )?; - Ok(()) - }) - .unwrap(); - - let got = get_chunk_raw_refs(&cfg, &c1.id).unwrap().unwrap(); - assert_eq!(got.len(), 2); - assert_eq!(got[0].path, "raw/gmail/thread_1/message_1.md"); - assert_eq!(got[0].start, 10); - assert_eq!(got[0].end, Some(42)); - assert!(get_chunk_raw_refs(&cfg, "missing").unwrap().is_none()); - - let paths = list_chunk_raw_ref_paths_with_prefix(&cfg, "raw/gmail/").unwrap(); - assert_eq!(paths.len(), 2); - assert!(paths.contains("raw/gmail/thread_1/message_1.md")); - assert!(paths.contains("raw/gmail/thread_1/message_2.md")); -} - -#[test] -fn content_pointer_accessors_return_only_complete_non_deleted_rows() { - let (_tmp, cfg) = test_config(); - let chunk = sample_chunk("notion:page-1", 0, 1_700_000_000_000); - upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); - with_connection(&cfg, |conn| { - conn.execute( - "UPDATE mem_tree_chunks SET content_path = ?1, content_sha256 = ?2 WHERE id = ?3", - params!["chunks/notion/page-1.md", "abc123", chunk.id], - )?; - conn.execute( - "INSERT OR IGNORE INTO mem_tree_trees (id, kind, scope, created_at_ms) - VALUES ('tree-content-pointers', 'source', 'notion:page-1', 0)", - [], - )?; - conn.execute( - "INSERT INTO mem_tree_summaries ( - id, tree_id, tree_kind, level, child_ids_json, content, token_count, - entities_json, topics_json, time_range_start_ms, time_range_end_ms, - score, sealed_at_ms, deleted, content_path, content_sha256 - ) VALUES - ('summary-live', 'tree-content-pointers', 'source', 1, '[]', 'live', 1, - '[]', '[]', 0, 0, 1.0, 0, 0, 'summaries/live.md', 'sum123'), - ('summary-deleted', 'tree-content-pointers', 'source', 1, '[]', 'deleted', 1, - '[]', '[]', 0, 0, 1.0, 0, 1, 'summaries/deleted.md', 'sum456'), - ('summary-incomplete', 'tree-content-pointers', 'source', 1, '[]', 'incomplete', 1, - '[]', '[]', 0, 0, 1.0, 0, 0, 'summaries/incomplete.md', NULL)", - [], - )?; - Ok(()) - }) - .unwrap(); - - assert_eq!( - get_chunk_content_path(&cfg, &chunk.id).unwrap().as_deref(), - Some("chunks/notion/page-1.md") - ); - assert_eq!( - get_chunk_content_pointers(&cfg, &chunk.id).unwrap(), - Some(("chunks/notion/page-1.md".into(), "abc123".into())) - ); - assert!(get_chunk_content_pointers(&cfg, "missing") - .unwrap() - .is_none()); - assert_eq!( - get_summary_content_pointers(&cfg, "summary-live").unwrap(), - Some(("summaries/live.md".into(), "sum123".into())) - ); - - let summaries = list_summaries_with_content_path(&cfg).unwrap(); - assert_eq!( - summaries, - vec![( - "summary-live".into(), - "summaries/live.md".into(), - "sum123".into() - )] - ); -} - -#[test] -fn raw_file_ingest_gate_is_order_preserving_and_prefix_scoped() { - let (_tmp, cfg) = test_config(); - let paths = vec![ - "raw/gmail/a.md".to_string(), - "raw/gmail/b.md".to_string(), - "raw/slack/c.md".to_string(), - ]; - - assert_eq!(filter_raw_paths_not_ingested(&cfg, &paths).unwrap(), paths); - assert_eq!(mark_raw_paths_ingested(&cfg, &paths[..2]).unwrap(), 2); - assert_eq!(mark_raw_paths_ingested(&cfg, &paths[..2]).unwrap(), 0); - assert_eq!( - filter_raw_paths_not_ingested(&cfg, &paths).unwrap(), - vec!["raw/slack/c.md".to_string()] - ); - assert_eq!( - count_raw_paths_ingested_with_prefix(&cfg, "raw/gmail/").unwrap(), - 2 - ); - assert_eq!( - count_raw_paths_ingested_with_prefix(&cfg, "raw/slack/").unwrap(), - 0 - ); -} - -#[test] -fn staged_chunk_upsert_persists_preview_and_content_pointers() { - let (_tmp, cfg) = test_config(); - let mut chunk = sample_chunk("notion:staged", 0, 1_700_000_000_000); - chunk.metadata.source_kind = SourceKind::Document; - chunk.metadata.path_scope = Some("notion:workspace".into()); - chunk.content = "x".repeat(620); - chunk.token_count = 155; - let staged = StagedChunk { - chunk: chunk.clone(), - content_path: "chunks/notion/workspace/page.md".into(), - content_sha256: "sha-staged".into(), - }; - - let inserted = with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - let inserted = upsert_staged_chunks_tx(&tx, std::slice::from_ref(&staged))?; - tx.commit()?; - Ok(inserted) - }) - .unwrap(); - assert_eq!(inserted, 1); - - let got = get_chunk(&cfg, &chunk.id).unwrap().unwrap(); - assert_eq!(got.content.len(), 500); - assert!(got.content.chars().all(|ch| ch == 'x')); - assert_eq!(got.metadata.path_scope.as_deref(), Some("notion:workspace")); - assert_eq!(got.token_count, 155); - assert_eq!( - get_chunk_content_pointers(&cfg, &chunk.id).unwrap(), - Some(( - "chunks/notion/workspace/page.md".into(), - "sha-staged".into() - )) - ); -} - -#[test] -fn missing_chunk_returns_none() { - let (_tmp, cfg) = test_config(); - assert!(get_chunk(&cfg, "nonexistent").unwrap().is_none()); -} - -#[test] -fn empty_batch_is_noop() { - let (_tmp, cfg) = test_config(); - assert_eq!(upsert_chunks(&cfg, &[]).unwrap(), 0); - assert_eq!(count_chunks(&cfg).unwrap(), 0); -} - -#[test] -fn schema_has_content_path_and_content_sha256_columns() { - let (_tmp, cfg) = test_config(); - with_connection(&cfg, |conn| { - let mut stmt = conn.prepare("PRAGMA table_info(mem_tree_chunks)")?; - let names: Vec = stmt - .query_map(params![], |row| row.get::<_, String>(1))? - .filter_map(|r| r.ok()) - .collect(); - assert!( - names.iter().any(|n| n == "path_scope"), - "missing path_scope; found: {names:?}" - ); - assert!( - names.iter().any(|n| n == "content_path"), - "missing content_path; found: {names:?}" - ); - assert!( - names.iter().any(|n| n == "content_sha256"), - "missing content_sha256; found: {names:?}" - ); - Ok(()) - }) - .unwrap(); -} +#[path = "store_tests_more.rs"] +mod more; diff --git a/src/memory/chunks/store_tests_more.rs b/src/memory/chunks/store_tests_more.rs new file mode 100644 index 0000000..6b7b378 --- /dev/null +++ b/src/memory/chunks/store_tests_more.rs @@ -0,0 +1,397 @@ +use super::*; + +#[test] +fn delete_chunks_by_source_removes_safe_content_files_but_rejects_escape_paths() { + let (_tmp, cfg) = test_config(); + let safe = sample_chunk("slack:c-1", 0, 1_700_000_000_000); + let unsafe_chunk = sample_chunk("slack:c-1", 1, 1_700_000_001_000); + upsert_chunks(&cfg, &[safe.clone(), unsafe_chunk.clone()]).unwrap(); + + let root = content_root(&cfg); + let safe_rel = "chunks/safe.md"; + let safe_path = root.join(safe_rel); + std::fs::create_dir_all(safe_path.parent().unwrap()).unwrap(); + std::fs::write(&safe_path, "safe").unwrap(); + + let outside_path = root.parent().unwrap().join("outside.md"); + std::fs::write(&outside_path, "outside").unwrap(); + + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", + params![safe_rel, safe.id], + )?; + conn.execute( + "UPDATE mem_tree_chunks SET content_path = ?1 WHERE id = ?2", + params!["../outside.md", unsafe_chunk.id], + )?; + Ok(()) + }) + .unwrap(); + + let deleted = delete_chunks_by_source(&cfg, SourceKind::Chat, "slack:c-1").unwrap(); + + assert_eq!(deleted, 2); + assert!(!safe_path.exists()); + assert!(outside_path.exists()); +} + +#[test] +fn delete_chunks_by_source_removes_only_unreferenced_raw_files_and_gates() { + let (_tmp, cfg) = test_config(); + let first = sample_chunk("gmail:thread-1", 0, 1_700_000_000_000); + let second = sample_chunk("gmail:thread-2", 0, 1_700_000_001_000); + upsert_chunks(&cfg, &[first.clone(), second.clone()]).unwrap(); + + let shared = "raw/gmail/shared.md".to_string(); + let unique = "raw/gmail/unique.md".to_string(); + set_chunk_raw_refs( + &cfg, + &first.id, + &[ + RawRef { + path: shared.clone(), + start: 0, + end: None, + }, + RawRef { + path: unique.clone(), + start: 0, + end: None, + }, + ], + ) + .unwrap(); + set_chunk_raw_refs( + &cfg, + &second.id, + &[RawRef { + path: shared.clone(), + start: 0, + end: None, + }], + ) + .unwrap(); + let root = content_root(&cfg); + for path in [&shared, &unique] { + let absolute = root.join(path); + std::fs::create_dir_all(absolute.parent().unwrap()).unwrap(); + std::fs::write(absolute, path).unwrap(); + } + mark_raw_paths_ingested(&cfg, &[shared.clone(), unique.clone()]).unwrap(); + + assert_eq!( + delete_chunks_by_source(&cfg, SourceKind::Chat, "gmail:thread-1").unwrap(), + 1 + ); + assert!(root.join(&shared).exists()); + assert!(!root.join(&unique).exists()); + assert_eq!( + filter_raw_paths_not_ingested(&cfg, &[shared.clone(), unique.clone()]).unwrap(), + vec![unique.clone()] + ); + + assert_eq!( + delete_chunks_by_source(&cfg, SourceKind::Chat, "gmail:thread-2").unwrap(), + 1 + ); + assert!(!root.join(&shared).exists()); + assert_eq!( + filter_raw_paths_not_ingested(&cfg, std::slice::from_ref(&shared)).unwrap(), + vec![shared] + ); +} + +#[test] +fn corrupt_surviving_raw_refs_make_deletion_cleanup_fail_closed() { + let (_tmp, cfg) = test_config(); + let removed = sample_chunk("gmail:removed", 0, 1_700_000_000_000); + let corrupt = sample_chunk("gmail:corrupt", 0, 1_700_000_001_000); + upsert_chunks(&cfg, &[removed.clone(), corrupt.clone()]).unwrap(); + let raw_path = "raw/gmail/maybe-shared.md".to_string(); + set_chunk_raw_refs( + &cfg, + &removed.id, + &[RawRef { + path: raw_path.clone(), + start: 0, + end: None, + }], + ) + .unwrap(); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_chunks SET raw_refs_json='corrupt' WHERE id=?1", + params![corrupt.id], + )?; + Ok(()) + }) + .unwrap(); + let absolute = content_root(&cfg).join(&raw_path); + std::fs::create_dir_all(absolute.parent().unwrap()).unwrap(); + std::fs::write(&absolute, "preserve").unwrap(); + mark_raw_paths_ingested(&cfg, std::slice::from_ref(&raw_path)).unwrap(); + + delete_chunks_by_source(&cfg, SourceKind::Chat, "gmail:removed").unwrap(); + + assert!(absolute.exists()); + assert!(filter_raw_paths_not_ingested(&cfg, &[raw_path]) + .unwrap() + .is_empty()); +} + +#[test] +fn raw_refs_round_trip_and_prefix_listing_tolerates_corrupt_rows() { + let (_tmp, cfg) = test_config(); + let c1 = sample_chunk("gmail:thread-1", 0, 1_700_000_000_000); + let c2 = sample_chunk("gmail:thread-2", 0, 1_700_000_001_000); + let corrupt = sample_chunk("gmail:thread-3", 0, 1_700_000_002_000); + upsert_chunks(&cfg, &[c1.clone(), c2.clone(), corrupt.clone()]).unwrap(); + + let refs = vec![ + RawRef { + path: "raw/gmail/thread_1/message_1.md".into(), + start: 10, + end: Some(42), + }, + RawRef { + path: "raw/gmail/thread_1/message_2.md".into(), + start: 0, + end: None, + }, + ]; + set_chunk_raw_refs(&cfg, &c1.id, &refs).unwrap(); + set_chunk_raw_refs( + &cfg, + &c2.id, + &[RawRef { + path: "raw/slack/channel/message.md".into(), + start: 0, + end: None, + }], + ) + .unwrap(); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_chunks SET raw_refs_json = 'not valid json' WHERE id = ?1", + params![corrupt.id], + )?; + Ok(()) + }) + .unwrap(); + + let got = get_chunk_raw_refs(&cfg, &c1.id).unwrap().unwrap(); + assert_eq!(got.len(), 2); + assert_eq!(got[0].path, "raw/gmail/thread_1/message_1.md"); + assert_eq!(got[0].start, 10); + assert_eq!(got[0].end, Some(42)); + assert!(get_chunk_raw_refs(&cfg, "missing").unwrap().is_none()); + + let paths = list_chunk_raw_ref_paths_with_prefix(&cfg, "raw/gmail/").unwrap(); + assert_eq!(paths.len(), 2); + assert!(paths.contains("raw/gmail/thread_1/message_1.md")); + assert!(paths.contains("raw/gmail/thread_1/message_2.md")); +} + +#[test] +fn content_pointer_accessors_return_only_complete_non_deleted_rows() { + let (_tmp, cfg) = test_config(); + let chunk = sample_chunk("notion:page-1", 0, 1_700_000_000_000); + upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_chunks SET content_path = ?1, content_sha256 = ?2 WHERE id = ?3", + params!["chunks/notion/page-1.md", "abc123", chunk.id], + )?; + conn.execute( + "INSERT OR IGNORE INTO mem_tree_trees (id, kind, scope, created_at_ms) + VALUES ('tree-content-pointers', 'source', 'notion:page-1', 0)", + [], + )?; + conn.execute( + "INSERT INTO mem_tree_summaries ( + id, tree_id, tree_kind, level, child_ids_json, content, token_count, + entities_json, topics_json, time_range_start_ms, time_range_end_ms, + score, sealed_at_ms, deleted, content_path, content_sha256 + ) VALUES + ('summary-live', 'tree-content-pointers', 'source', 1, '[]', 'live', 1, + '[]', '[]', 0, 0, 1.0, 0, 0, 'summaries/live.md', 'sum123'), + ('summary-deleted', 'tree-content-pointers', 'source', 1, '[]', 'deleted', 1, + '[]', '[]', 0, 0, 1.0, 0, 1, 'summaries/deleted.md', 'sum456'), + ('summary-incomplete', 'tree-content-pointers', 'source', 1, '[]', 'incomplete', 1, + '[]', '[]', 0, 0, 1.0, 0, 0, 'summaries/incomplete.md', NULL)", + [], + )?; + Ok(()) + }) + .unwrap(); + + assert_eq!( + get_chunk_content_path(&cfg, &chunk.id).unwrap().as_deref(), + Some("chunks/notion/page-1.md") + ); + assert_eq!( + get_chunk_content_pointers(&cfg, &chunk.id).unwrap(), + Some(("chunks/notion/page-1.md".into(), "abc123".into())) + ); + assert!(get_chunk_content_pointers(&cfg, "missing") + .unwrap() + .is_none()); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_chunks SET content_path = '', content_sha256 = '' WHERE id = ?1", + rusqlite::params![chunk.id], + )?; + Ok(()) + }) + .unwrap(); + assert!(get_chunk_content_path(&cfg, &chunk.id).unwrap().is_none()); + assert!(get_chunk_content_pointers(&cfg, &chunk.id) + .unwrap() + .is_none()); + assert_eq!( + get_summary_content_pointers(&cfg, "summary-live").unwrap(), + Some(("summaries/live.md".into(), "sum123".into())) + ); + + let summaries = list_summaries_with_content_path(&cfg).unwrap(); + assert_eq!( + summaries, + vec![( + "summary-live".into(), + "summaries/live.md".into(), + "sum123".into() + )] + ); +} + +#[test] +fn raw_file_ingest_gate_is_order_preserving_and_prefix_scoped() { + let (_tmp, cfg) = test_config(); + let paths = vec![ + "raw/gmail/a.md".to_string(), + "raw/gmail/b.md".to_string(), + "raw/slack/c.md".to_string(), + ]; + + assert_eq!(filter_raw_paths_not_ingested(&cfg, &paths).unwrap(), paths); + assert_eq!(mark_raw_paths_ingested(&cfg, &paths[..2]).unwrap(), 2); + assert_eq!(mark_raw_paths_ingested(&cfg, &paths[..2]).unwrap(), 0); + assert_eq!( + filter_raw_paths_not_ingested(&cfg, &paths).unwrap(), + vec!["raw/slack/c.md".to_string()] + ); + assert_eq!( + count_raw_paths_ingested_with_prefix(&cfg, "raw/gmail/").unwrap(), + 2 + ); + assert_eq!( + count_raw_paths_ingested_with_prefix(&cfg, "raw/slack/").unwrap(), + 0 + ); +} + +#[test] +fn staged_chunk_upsert_persists_preview_and_content_pointers() { + let (_tmp, cfg) = test_config(); + let mut chunk = sample_chunk("notion:staged", 0, 1_700_000_000_000); + chunk.metadata.source_kind = SourceKind::Document; + chunk.metadata.path_scope = Some("notion:workspace".into()); + chunk.content = "x".repeat(620); + chunk.token_count = 155; + let staged = StagedChunk { + chunk: chunk.clone(), + content_path: "chunks/notion/workspace/page.md".into(), + content_sha256: "sha-staged".into(), + }; + + let inserted = with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + let inserted = upsert_staged_chunks_tx(&tx, std::slice::from_ref(&staged))?; + tx.commit()?; + Ok(inserted) + }) + .unwrap(); + assert_eq!(inserted, 1); + + let got = get_chunk(&cfg, &chunk.id).unwrap().unwrap(); + assert_eq!(got.content.len(), 500); + assert!(got.content.chars().all(|ch| ch == 'x')); + assert_eq!(got.metadata.path_scope.as_deref(), Some("notion:workspace")); + assert_eq!(got.token_count, 155); + assert_eq!( + get_chunk_content_pointers(&cfg, &chunk.id).unwrap(), + Some(( + "chunks/notion/workspace/page.md".into(), + "sha-staged".into() + )) + ); +} + +#[test] +fn missing_chunk_returns_none() { + let (_tmp, cfg) = test_config(); + assert!(get_chunk(&cfg, "nonexistent").unwrap().is_none()); +} + +#[test] +fn empty_batch_is_noop() { + let (_tmp, cfg) = test_config(); + assert_eq!(upsert_chunks(&cfg, &[]).unwrap(), 0); + assert_eq!(count_chunks(&cfg).unwrap(), 0); +} + +#[test] +fn schema_has_content_path_and_content_sha256_columns() { + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + let mut stmt = conn.prepare("PRAGMA table_info(mem_tree_chunks)")?; + let names: Vec = stmt + .query_map(params![], |row| row.get::<_, String>(1))? + .filter_map(|r| r.ok()) + .collect(); + assert!( + names.iter().any(|n| n == "path_scope"), + "missing path_scope; found: {names:?}" + ); + assert!( + names.iter().any(|n| n == "content_path"), + "missing content_path; found: {names:?}" + ); + assert!( + names.iter().any(|n| n == "content_sha256"), + "missing content_sha256; found: {names:?}" + ); + Ok(()) + }) + .unwrap(); +} + +#[test] +fn plain_upsert_clears_staged_pointer_and_readmits_dropped_row() { + let (_tmp, cfg) = test_config(); + let chunk = sample_chunk("notion:plain-replace", 0, 1_700_000_000_000); + upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_chunks + SET content_path = 'stale.md', content_sha256 = 'stale', + lifecycle_status = 'dropped' + WHERE id = ?1", + rusqlite::params![chunk.id], + )?; + Ok(()) + }) + .unwrap(); + + upsert_chunks(&cfg, std::slice::from_ref(&chunk)).unwrap(); + assert!(get_chunk_content_pointers(&cfg, &chunk.id) + .unwrap() + .is_none()); + assert_eq!( + super::super::get_chunk_lifecycle_status(&cfg, &chunk.id) + .unwrap() + .as_deref(), + Some("admitted") + ); +} diff --git a/src/memory/chunks/types_tests.rs b/src/memory/chunks/types_tests.rs index 62797bf..e7db31f 100644 --- a/src/memory/chunks/types_tests.rs +++ b/src/memory/chunks/types_tests.rs @@ -1,6 +1,7 @@ //! Unit tests for the chunk model (`super`). use super::*; +use chrono::TimeZone; #[test] fn chunk_id_is_deterministic() { @@ -129,3 +130,79 @@ fn approx_token_count_scales_linearly() { assert_eq!(approx_token_count("abcde"), 2); // 5→2 assert_eq!(approx_token_count(&"x".repeat(400)), 100); } + +#[test] +fn source_kind_parse_rejects_unknown_wire_values() { + assert_eq!( + SourceKind::parse("video").unwrap_err(), + "unknown source kind: video" + ); +} + +#[test] +fn metadata_constructor_and_source_ref_fill_documented_defaults() { + let timestamp = Utc.timestamp_millis_opt(1_700_000_000_123).unwrap(); + let mut metadata = Metadata::point_in_time(SourceKind::Document, "doc-1", "alice", timestamp); + metadata.source_ref = Some(SourceRef::new("notion://doc-1")); + + assert_eq!(metadata.source_id, "doc-1"); + assert_eq!(metadata.owner, "alice"); + assert_eq!(metadata.time_range, (timestamp, timestamp)); + assert!(metadata.tags.is_empty()); + assert_eq!(metadata.source_ref.unwrap().value, "notion://doc-1"); +} + +#[test] +fn chunk_json_round_trips_millisecond_time_range_and_partial_default() { + let timestamp = Utc.timestamp_millis_opt(1_700_000_000_123).unwrap(); + let chunk = Chunk { + id: "chunk".into(), + content: "body".into(), + metadata: Metadata::point_in_time(SourceKind::Chat, "channel", "alice", timestamp), + token_count: 1, + seq_in_source: 0, + created_at: timestamp, + partial_message: true, + }; + let encoded = serde_json::to_value(&chunk).unwrap(); + assert_eq!( + encoded["metadata"]["time_range"]["start_ms"], + timestamp.timestamp_millis() + ); + assert_eq!(serde_json::from_value::(encoded).unwrap(), chunk); + + let mut legacy = serde_json::to_value(&chunk).unwrap(); + legacy.as_object_mut().unwrap().remove("partial_message"); + assert!( + !serde_json::from_value::(legacy) + .unwrap() + .partial_message + ); +} + +#[test] +fn chunk_json_rejects_out_of_range_time_range_endpoints() { + let timestamp = Utc.timestamp_millis_opt(1_700_000_000_123).unwrap(); + let chunk = Chunk { + id: "chunk".into(), + content: "body".into(), + metadata: Metadata::point_in_time(SourceKind::Chat, "channel", "alice", timestamp), + token_count: 1, + seq_in_source: 0, + created_at: timestamp, + partial_message: false, + }; + let mut encoded = serde_json::to_value(chunk).unwrap(); + encoded["metadata"]["time_range"]["start_ms"] = serde_json::json!(i64::MAX); + assert!(serde_json::from_value::(encoded.clone()) + .unwrap_err() + .to_string() + .contains("invalid start_ms")); + + encoded["metadata"]["time_range"]["start_ms"] = serde_json::json!(0); + encoded["metadata"]["time_range"]["end_ms"] = serde_json::json!(i64::MAX); + assert!(serde_json::from_value::(encoded) + .unwrap_err() + .to_string() + .contains("invalid end_ms")); +} diff --git a/src/memory/conversations/bus.rs b/src/memory/conversations/bus.rs index 151188a..ee8ca2b 100644 --- a/src/memory/conversations/bus.rs +++ b/src/memory/conversations/bus.rs @@ -279,13 +279,6 @@ struct ChannelTurnDescriptor<'a> { /// short-lived channel threads, O(n) per turn (O(n²) over a thread's /// lifetime) for long-running ones. /// -/// NOTE (audit TR-8): every call passes `labels: Some(vec!["general"])` to -/// `ensure_thread`, and the thread-log fold (`ThreadLogEntry::Upsert` -/// handling in `store_index::thread_index_unlocked`) lets a `Some` label list -/// on a later `Upsert` override the thread's current labels. So a user-set -/// label (e.g. via `update_thread_labels`) is silently reset to `["general"]` -/// the next time a message lands on this thread — this call site should pass -/// `labels: None` once thread creation and per-turn touch are distinguished. fn persist_channel_turn( workspace_dir: &Path, descriptor: ChannelTurnDescriptor<'_>, @@ -311,7 +304,9 @@ fn persist_channel_turn( title, created_at: created_at.clone(), parent_thread_id: None, - labels: Some(vec!["general".to_string()]), + // The store infers `general` when creating a channel thread. On + // later touches, `None` preserves any labels the user assigned. + labels: None, personality_id: None, }, )?; @@ -354,21 +349,30 @@ fn persist_channel_turn( /// channels append a `_thread:` suffix when a non-blank `thread_ts` is /// present) and prefixes it with `channel:`. /// -/// NOTE (audit TR-15): `base_key` is a plain `_`-joined concatenation of -/// `channel`, `sender`, and `reply_target`, none of which are guaranteed -/// underscore-free. Two distinct triples can collide onto the same thread id -/// — e.g. `("slack", "a_b", "c")` and `("slack", "a", "b_c")` both yield -/// `slack_a_b_c`. A length-prefixed or hex-encoded join (as -/// `ConversationStore::thread_messages_path`'s `hex_encode` already does for -/// filenames) would be collision-free; this preserves the existing wire -/// format for already-persisted thread ids, so it is not changed here. fn persisted_channel_thread_id( channel: &str, sender: &str, reply_target: &str, thread_ts: Option<&str>, ) -> String { - let base_key = format!("{channel}_{sender}_{reply_target}"); + let mut base_key = format!("{channel}_{sender}_{reply_target}"); + if [channel, sender, reply_target] + .iter() + .any(|component| component.contains('_')) + { + use sha2::{Digest, Sha256}; + let mut hasher = Sha256::new(); + for component in [channel, sender, reply_target] { + hasher.update(component.len().to_be_bytes()); + hasher.update(component.as_bytes()); + } + let suffix = hasher.finalize()[..6] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + base_key.push_str("__"); + base_key.push_str(&suffix); + } let key = if channel == "telegram" { base_key } else { diff --git a/src/memory/conversations/bus_tests.rs b/src/memory/conversations/bus_tests.rs index f43f688..591352f 100644 --- a/src/memory/conversations/bus_tests.rs +++ b/src/memory/conversations/bus_tests.rs @@ -64,6 +64,38 @@ async fn persists_inbound_and_processed_turns_into_workspace_thread() { assert_eq!(messages[1].extra_metadata["success"], true); } +#[tokio::test] +async fn later_channel_turn_preserves_user_assigned_labels() { + let temp = TempDir::new().expect("tempdir"); + let subscriber = ConversationPersistenceSubscriber::new(temp.path().to_path_buf()); + let event = |message_id: &str| ChannelEvent::Received { + channel: "slack".into(), + message_id: message_id.into(), + sender: "alice".into(), + reply_target: "general".into(), + content: "hello".into(), + thread_ts: None, + workspace_dir: temp.path().to_path_buf(), + }; + + subscriber.handle(&event("m1")).await; + super::super::update_thread_labels( + temp.path().to_path_buf(), + "channel:slack_alice_general", + vec!["tasks".into()], + "2026-07-14T00:00:00Z", + ) + .unwrap(); + subscriber.handle(&event("m2")).await; + + let thread = super::super::list_threads(temp.path().to_path_buf()) + .unwrap() + .into_iter() + .find(|thread| thread.id == "channel:slack_alice_general") + .unwrap(); + assert_eq!(thread.labels, vec!["tasks"]); +} + #[tokio::test] async fn telegram_thread_ts_does_not_split_persisted_thread() { let temp = TempDir::new().expect("tempdir"); @@ -129,6 +161,13 @@ fn persisted_channel_thread_id_ignores_blank_thread_ts() { assert_eq!(without, with_blank); } +#[test] +fn persisted_channel_thread_ids_do_not_alias_underscore_components() { + let left = persisted_channel_thread_id("slack", "a_b", "c", None); + let right = persisted_channel_thread_id("slack", "a", "b_c", None); + assert_ne!(left, right); +} + #[test] fn channel_thread_title_uses_thread_suffix_only_for_non_telegram_threads() { assert_eq!( diff --git a/src/memory/conversations/inverted_index.rs b/src/memory/conversations/inverted_index.rs index 6f12055..00efe75 100644 --- a/src/memory/conversations/inverted_index.rs +++ b/src/memory/conversations/inverted_index.rs @@ -86,6 +86,7 @@ struct DocEntry { content: String, // original, returned verbatim in hits content_normalized: String, // for Phase 2 substring verification created_at: String, + created_at_ms: i64, } /// In-memory trigram/bigram inverted index over conversation messages. @@ -140,6 +141,9 @@ impl InvertedIndex { return; } let normalized = normalize(&content); + let created_at_ms = chrono::DateTime::parse_from_rfc3339(&created_at) + .map(|value| value.timestamp_millis()) + .unwrap_or(i64::MIN); let doc_id = self.docs.len() as u32; for ngram in ngrams(&normalized) { if let Some(posting) = self.postings.get_mut(ngram) { @@ -159,6 +163,7 @@ impl InvertedIndex { content, content_normalized: normalized, created_at, + created_at_ms, })); self.by_message.insert(key, doc_id); } @@ -243,7 +248,16 @@ impl InvertedIndex { let mut per_term: Vec> = Vec::with_capacity(terms.len()); for term in &terms { let candidates = match self.candidates_for_term(term) { - Some(v) => v, + Some(v) => { + if v.len() > LARGE_CANDIDATE_LIMIT { + return self.recency_fallback(exclude_thread_id, limit); + } + v + } + // Terms too short for an ngram (notably one CJK character) + // require exact verification over the corpus. Returning the + // recency fallback here would manufacture unrelated score-0 + // hits instead of answering the query. None => self .docs .iter() @@ -251,9 +265,6 @@ impl InvertedIndex { .filter_map(|(i, slot)| slot.as_ref().map(|_| i as u32)) .collect::>(), }; - if candidates.len() > LARGE_CANDIDATE_LIMIT { - return self.recency_fallback(exclude_thread_id, limit); - } per_term.push(candidates); } @@ -283,21 +294,16 @@ impl InvertedIndex { // Ranking by `matched` (usize) is order-equivalent to ranking by // `score = matched / total_terms` since `total_terms` is a positive // constant, so the returned order is unchanged. - let mut ranked: Vec<(u32, usize, &str)> = hit_counts + let mut ranked: Vec<(u32, usize, i64)> = hit_counts .into_iter() .map(|(doc_id, matched)| { let entry = self.docs[doc_id as usize] .as_ref() .expect("doc_id from hit_counts must be live"); - (doc_id, matched, entry.created_at.as_str()) + (doc_id, matched, entry.created_at_ms) }) .collect(); - // NOTE: `created_at` is compared as a raw string (lexicographic, not - // parsed-to-instant). Values are RFC3339 with a consistent zero-padded - // width in practice, so this tiebreak orders correctly for same-offset - // timestamps, but mixed offsets (`+00:00` vs `Z` vs a non-UTC offset) - // can misorder — see audit TR-16. - ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| b.2.cmp(a.2))); + ranked.sort_by(|a, b| b.1.cmp(&a.1).then_with(|| b.2.cmp(&a.2))); ranked.truncate(limit); ranked @@ -385,11 +391,16 @@ impl InvertedIndex { exclude_thread_id: Option<&str>, limit: usize, ) -> Vec { - let mut hits: Vec = self + let mut entries: Vec<&DocEntry> = self .docs .iter() .filter_map(|slot| slot.as_ref()) .filter(|entry| exclude_thread_id != Some(entry.thread_id.as_ref())) + .collect(); + entries.sort_by_key(|entry| std::cmp::Reverse(entry.created_at_ms)); + entries.truncate(limit); + entries + .into_iter() .map(|entry| CrossThreadHit { thread_id: entry.thread_id.to_string(), message_id: entry.message_id.clone(), @@ -402,12 +413,7 @@ impl InvertedIndex { // the newest entries first. score: 0.0, }) - .collect(); - // NOTE: lexicographic string comparison, not a parsed-timestamp - // comparison — see the TR-16 note on the Phase 2 ranking above. - hits.sort_by(|a, b| b.created_at.cmp(&a.created_at)); - hits.truncate(limit); - hits + .collect() } } diff --git a/src/memory/conversations/inverted_index_tests.rs b/src/memory/conversations/inverted_index_tests.rs index 3bbe43d..7f577f4 100644 --- a/src/memory/conversations/inverted_index_tests.rs +++ b/src/memory/conversations/inverted_index_tests.rs @@ -162,6 +162,12 @@ fn pathological_query_short_circuits_to_recency() { assert_eq!(hits.len(), 5, "fallback must still respect limit"); // Score 0.0 is the recency-fallback marker. assert!(hits.iter().all(|h| h.score == 0.0)); + + idx.insert("bulk", msg("cjk", "猫だけ", "2026-04-11T10:00:00Z")); + let exact = idx.search("猫", 5, None); + assert_eq!(exact.len(), 1); + assert_eq!(exact[0].message_id, "cjk"); + assert_eq!(exact[0].score, 1.0); } #[test] @@ -232,3 +238,17 @@ fn ranks_by_score_then_recency_before_truncating() { hits[1].score ); } + +#[test] +fn recency_tiebreak_compares_instants_across_rfc3339_offsets() { + let mut idx = InvertedIndex::new(); + idx.insert( + "t1", + msg("older", "alpha match", "2026-04-10T12:30:00+02:00"), + ); + idx.insert("t1", msg("newer", "alpha match", "2026-04-10T11:00:00Z")); + + let hits = idx.search("alpha", 2, None); + assert_eq!(hits[0].message_id, "newer"); + assert_eq!(hits[1].message_id, "older"); +} diff --git a/src/memory/conversations/mod.rs b/src/memory/conversations/mod.rs index 739d7d5..de220cc 100644 --- a/src/memory/conversations/mod.rs +++ b/src/memory/conversations/mod.rs @@ -14,12 +14,12 @@ //! //! ## Layout //! -//! - [`types`] — the on-disk wire types (threads, messages, patches, hits). -//! - [`tokenize`] — multilingual normalization + character n-gram tokenizer. -//! - [`inverted_index`] — in-memory trigram/bigram index over message content. -//! - [`store`] — the JSONL [`ConversationStore`] (append/read/update/delete, +//! - `types` — the on-disk wire types (threads, messages, patches, hits). +//! - `tokenize` — multilingual normalization + character n-gram tokenizer. +//! - `inverted_index` — in-memory trigram/bigram index over message content. +//! - `store` — the JSONL [`ConversationStore`] (append/read/update/delete, //! process-wide write serialization, warm-index cache, cross-thread search). -//! - [`bus`] — a channel-persistence subscriber that mirrors inbound/processed +//! - `bus` — a channel-persistence subscriber that mirrors inbound/processed //! channel turns into the store. //! //! ## Ported-from-OpenHuman notes diff --git a/src/memory/conversations/store_index.rs b/src/memory/conversations/store_index.rs index 3277f98..106be6c 100644 --- a/src/memory/conversations/store_index.rs +++ b/src/memory/conversations/store_index.rs @@ -157,23 +157,14 @@ impl ConversationStore { pub(super) fn list_threads_unlocked(&self) -> Result, String> { let mut index = self.thread_index_unlocked()?; - // Backfill stats for any thread with no MessageAppended/Stats history - // yet (legacy data). The slow per-thread file read happens at most once - // per thread — we persist a `Stats` snapshot so subsequent - // list_threads calls hit the fast path. - let needs_backfill: Vec = index - .iter() - .filter_map(|(id, entry)| { - if entry.message_count.is_none() { - Some(id.clone()) - } else { - None - } - }) - .collect(); - if !needs_backfill.is_empty() { + // Reconcile the derived stat trail against the authoritative message + // files. Appending a message and its compact stat event spans two files + // and cannot be one filesystem transaction; this check repairs either + // crash window instead of trusting a possibly-short stat history. + let thread_ids = index.keys().cloned().collect::>(); + if !thread_ids.is_empty() { let threads_path = self.ensure_root()?.join(THREADS_FILENAME); - for thread_id in &needs_backfill { + for thread_id in &thread_ids { let (count, last_message_at) = self.measure_messages_unlocked(thread_id)?; // Treat created_at as last_message_at when there are no // messages — keeps the sort key meaningful and matches the @@ -184,14 +175,20 @@ impl ConversationStore { .map(|e| e.created_at.clone()) .unwrap_or_default() }); - append_jsonl( - &threads_path, - &ThreadLogEntry::Stats { - thread_id: thread_id.clone(), - message_count: count, - last_message_at: resolved_last.clone(), - }, - )?; + let differs = index.get(thread_id).is_none_or(|entry| { + entry.message_count != Some(count) + || entry.last_message_at.as_deref() != Some(resolved_last.as_str()) + }); + if differs { + append_jsonl( + &threads_path, + &ThreadLogEntry::Stats { + thread_id: thread_id.clone(), + message_count: count, + last_message_at: resolved_last.clone(), + }, + )?; + } if let Some(entry) = index.get_mut(thread_id) { entry.message_count = Some(count); entry.last_message_at = Some(resolved_last); @@ -221,22 +218,17 @@ impl ConversationStore { } }) .collect(); - // NOTE: both fields are compared as raw strings (lexicographic), not - // parsed timestamps. RFC3339 sorts correctly this way only when every - // value shares the same offset representation; mixed `+00:00` / `Z` / - // non-UTC offsets across threads can misorder the listing — see audit - // TR-16. threads.sort_by(|a, b| { - b.last_message_at - .cmp(&a.last_message_at) - .then_with(|| b.created_at.cmp(&a.created_at)) + timestamp_millis(&b.last_message_at) + .cmp(×tamp_millis(&a.last_message_at)) + .then_with(|| timestamp_millis(&b.created_at).cmp(×tamp_millis(&a.created_at))) }); Ok(threads) } /// Count messages and find the newest timestamp by reading the per-thread - /// JSONL file. Slow path — used only when the threads-log stat history is - /// missing (legacy data) so we can write a one-time `Stats` snapshot. + /// JSONL file. This is the authoritative source used to reconcile the + /// compact thread stat trail after either side of a two-file append crash. pub(super) fn measure_messages_unlocked( &self, thread_id: &str, @@ -260,17 +252,8 @@ impl ConversationStore { Some(entry) => entry, None => return Ok(None), }; - // Prefer the index-tracked stats (cheap). Fall back to a single - // per-thread file read for legacy threads with no stat history — - // list_threads is responsible for permanently backfilling those. - let (message_count, last_message_at) = - match (entry.message_count, entry.last_message_at.as_ref()) { - (Some(count), Some(last_at)) => (count, last_at.clone()), - _ => { - let (count, last_at) = self.measure_messages_unlocked(thread_id)?; - (count, last_at.unwrap_or_else(|| entry.created_at.clone())) - } - }; + let (message_count, last_at) = self.measure_messages_unlocked(thread_id)?; + let last_message_at = last_at.unwrap_or_else(|| entry.created_at.clone()); Ok(Some(ConversationThread { id: thread_id.to_string(), title: entry.title.clone(), @@ -397,3 +380,9 @@ impl ConversationStore { }) } } + +fn timestamp_millis(value: &str) -> i64 { + chrono::DateTime::parse_from_rfc3339(value) + .map(|timestamp| timestamp.timestamp_millis()) + .unwrap_or(i64::MIN) +} diff --git a/src/memory/conversations/store_ops.rs b/src/memory/conversations/store_ops.rs index 395dfde..1189cb1 100644 --- a/src/memory/conversations/store_ops.rs +++ b/src/memory/conversations/store_ops.rs @@ -111,19 +111,10 @@ impl ConversationStore { /// Append a message to the thread's JSONL file. Errors if the thread is missing. /// - /// Persists via two separate fsync'd appends — the message row itself, - /// then a `MessageAppended` stat entry in `threads.jsonl` — plus, if the - /// in-memory search index is already warm for this workspace, an - /// in-process index update. None of the three are transactional with - /// each other. - /// - /// NOTE (audit TR-17): a crash between the message append and the - /// `MessageAppended` append leaves the message durably persisted but the - /// thread's cached `message_count`/`last_message_at` permanently one - /// short — `list_threads_unlocked`'s legacy backfill only re-derives - /// those stats from the raw message file when there is *no* prior - /// `MessageAppended`/`Stats` history at all (`message_count.is_none()`), - /// so a thread that already has stat history does not get corrected. + /// Persists via two separate fsync'd appends — the authoritative message + /// row, then a compact `MessageAppended` stat entry. Thread reads reconcile + /// that stat trail against the message file, repairing a crash between the + /// two appends. pub fn append_message( &self, thread_id: &str, diff --git a/src/memory/conversations/store_tests.rs b/src/memory/conversations/store_tests.rs index ad55a20..a47d879 100644 --- a/src/memory/conversations/store_tests.rs +++ b/src/memory/conversations/store_tests.rs @@ -468,855 +468,5 @@ fn store_handles_labels_and_inference() { } } -#[test] -fn conversation_store_new() { - let tmp = TempDir::new().unwrap(); - let store = ConversationStore::new(tmp.path().to_path_buf()); - let threads = store.list_threads().unwrap(); - assert!(threads.is_empty()); -} - -#[test] -fn conversation_purge_stats_default() { - let stats = ConversationPurgeStats::default(); - assert_eq!(stats.thread_count, 0); - assert_eq!(stats.message_count, 0); -} - -#[test] -fn list_threads_does_not_read_per_thread_files_after_first_call() { - // After the first list_threads (which may backfill), deleting every - // per-thread messages file must leave count + last_message_at intact — - // proving the slow path is no longer on the hot loop. - let (temp, store) = make_store(); - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "t1".to_string(), - title: "T1".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - for i in 0..3 { - store - .append_message( - "t1", - ConversationMessage { - id: format!("m{i}"), - content: format!("hi {i}"), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: format!("2026-04-10T12:0{}:00Z", i + 1), - }, - ) - .unwrap(); - } - // Warm-up: list_threads folds the MessageAppended entries. - let _ = store.list_threads().unwrap(); - - // Now blow away the per-thread JSONL. If list_threads still reads it, - // the count would drop to 0. If our index-only path works, the cached - // (3, latest_ts) survives. - let messages_dir = temp - .path() - .join("memory") - .join("conversations") - .join("threads"); - let entries: Vec<_> = std::fs::read_dir(&messages_dir) - .unwrap() - .filter_map(Result::ok) - .collect(); - for entry in entries { - std::fs::remove_file(entry.path()).unwrap(); - } - - let threads = store.list_threads().unwrap(); - assert_eq!(threads.len(), 1); - assert_eq!(threads[0].message_count, 3); - assert_eq!(threads[0].last_message_at, "2026-04-10T12:03:00Z"); -} - -#[test] -fn backfill_writes_stats_snapshot_for_legacy_threads() { - // Simulate legacy data: write only an Upsert entry (no MessageAppended) - // plus a per-thread messages file. The first list_threads must backfill. - let (temp, store) = make_store(); - let conversations_dir = temp.path().join("memory").join("conversations"); - std::fs::create_dir_all(conversations_dir.join("threads")).unwrap(); - - let threads_log = conversations_dir.join("threads.jsonl"); - let upsert = serde_json::json!({ - "op": "upsert", - "thread_id": "legacy-1", - "title": "Legacy", - "created_at": "2026-04-10T08:00:00Z", - "updated_at": "2026-04-10T08:00:00Z", - }); - std::fs::write(&threads_log, format!("{}\n", upsert)).unwrap(); - - // Write 2 messages directly to the per-thread file (no MessageAppended - // entries — this is what pre-upgrade data looks like). - let messages_file = conversations_dir - .join("threads") - .join(format!("{}.jsonl", hex_encode("legacy-1".as_bytes()))); - let m1 = serde_json::json!({ - "id": "m1", "content": "a", "type": "text", - "extraMetadata": {}, "sender": "user", - "createdAt": "2026-04-10T09:00:00Z", - }); - let m2 = serde_json::json!({ - "id": "m2", "content": "b", "type": "text", - "extraMetadata": {}, "sender": "user", - "createdAt": "2026-04-10T09:05:00Z", - }); - std::fs::write(&messages_file, format!("{m1}\n{m2}\n")).unwrap(); - - let threads = store.list_threads().unwrap(); - assert_eq!(threads.len(), 1); - assert_eq!(threads[0].message_count, 2); - assert_eq!(threads[0].last_message_at, "2026-04-10T09:05:00Z"); - - // The backfill should have appended a Stats entry — check the log - // contents now contain "op":"stats" for legacy-1. - let log = std::fs::read_to_string(&threads_log).unwrap(); - assert!( - log.contains("\"op\":\"stats\"") && log.contains("legacy-1"), - "expected backfilled Stats entry in threads.jsonl, got:\n{log}", - ); - - // Second call: blow away the messages file. Stats from the log keep - // count + last_message_at correct without re-reading. - std::fs::remove_file(&messages_file).unwrap(); - let threads2 = store.list_threads().unwrap(); - assert_eq!(threads2[0].message_count, 2); - assert_eq!(threads2[0].last_message_at, "2026-04-10T09:05:00Z"); -} - -#[test] -fn legacy_log_without_stats_still_parses() { - // Old on-disk format (only Upsert + Delete variants) must still load - // without errors after the enum gained MessageAppended + Stats. - let (temp, store) = make_store(); - let conversations_dir = temp.path().join("memory").join("conversations"); - std::fs::create_dir_all(conversations_dir.join("threads")).unwrap(); - let threads_log = conversations_dir.join("threads.jsonl"); - let upsert = serde_json::json!({ - "op": "upsert", - "thread_id": "old", - "title": "Old", - "created_at": "2026-04-10T08:00:00Z", - "updated_at": "2026-04-10T08:00:00Z", - }); - std::fs::write(&threads_log, format!("{}\n", upsert)).unwrap(); - - let threads = store.list_threads().unwrap(); - assert_eq!(threads.len(), 1); - assert_eq!(threads[0].id, "old"); - assert_eq!(threads[0].message_count, 0); - // No messages → last_message_at falls back to created_at. - assert_eq!(threads[0].last_message_at, "2026-04-10T08:00:00Z"); -} - -#[test] -fn delete_thread_clears_stats_from_index() { - let (_temp, store) = make_store(); - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "doomed".to_string(), - title: "Doomed".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "doomed", - ConversationMessage { - id: "m1".to_string(), - content: "x".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - assert_eq!(store.list_threads().unwrap().len(), 1); - - store - .delete_thread("doomed", "2026-04-10T12:02:00Z") - .unwrap(); - assert!(store.list_threads().unwrap().is_empty()); -} - -#[test] -fn search_cross_thread_messages_finds_hits_outside_excluded_thread() { - let (_temp, store) = make_store(); - - // Chat A — durable fact lives here. - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "thread-a".to_string(), - title: "Chat A".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "thread-a", - ConversationMessage { - id: "m-a-1".to_string(), - content: "Remember: my project is called Phoenix and uses Go and PostgreSQL." - .to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - - // Chat B — active chat, asking dependent question. Should be excluded - // so its own text doesn't echo back into [Cross-chat context]. - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "thread-b".to_string(), - title: "Chat B".to_string(), - created_at: "2026-04-10T13:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "thread-b", - ConversationMessage { - id: "m-b-1".to_string(), - content: "What database does my project use?".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T13:01:00Z".to_string(), - }, - ) - .unwrap(); - - let hits = store - .search_cross_thread_messages("What database does my project use", 10, Some("thread-b")) - .expect("cross-thread search"); - - assert_eq!(hits.len(), 1, "exactly one cross-thread hit"); - let hit = &hits[0]; - assert_eq!(hit.thread_id, "thread-a"); - assert!(hit.content.contains("PostgreSQL")); - assert!(hit.score > 0.0); -} - -#[test] -fn search_cross_thread_messages_excludes_active_thread() { - let (_temp, store) = make_store(); - - // Single thread — the only matching message lives in the thread we're - // about to exclude. Expect zero hits (don't echo same-chat history). - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "thread-only".to_string(), - title: "Only".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "thread-only", - ConversationMessage { - id: "m-1".to_string(), - content: "PostgreSQL deployment running on staging".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - - let hits = store - .search_cross_thread_messages("PostgreSQL deployment staging", 10, Some("thread-only")) - .expect("cross-thread search"); - assert!( - hits.is_empty(), - "active thread must not echo into cross-chat" - ); - - // Sanity: without exclude, the hit is returned. - let hits_no_exclude = store - .search_cross_thread_messages("PostgreSQL deployment staging", 10, None) - .expect("cross-thread search"); - assert_eq!(hits_no_exclude.len(), 1); -} - -#[test] -fn search_cross_thread_messages_skips_short_terms_and_empty_queries() { - let (_temp, store) = make_store(); - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "t".to_string(), - title: "T".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "t", - ConversationMessage { - id: "m".to_string(), - content: "Postgres".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - - // All terms < 3 chars → empty - assert!(store - .search_cross_thread_messages("a is on", 10, None) - .unwrap() - .is_empty()); - // Empty query → empty - assert!(store - .search_cross_thread_messages("", 10, None) - .unwrap() - .is_empty()); -} - -#[test] -fn search_cross_thread_messages_finds_polish_substring_without_diacritics() { - let (_temp, store) = make_store(); - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "thread-pl".to_string(), - title: "PL".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "thread-pl", - ConversationMessage { - id: "m1".to_string(), - content: "Lecę w piątek do Łodzi a potem Krakowa".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - - // Query without diacritics should still find content with them. - let hits = store - .search_cross_thread_messages("Lodzi", 10, None) - .expect("cross-thread search"); - assert_eq!(hits.len(), 1, "ł-fold should match Łodzi via lodzi"); - - let hits = store - .search_cross_thread_messages("krakow", 10, None) - .expect("cross-thread search"); - assert_eq!(hits.len(), 1, "diacritic strip should match Krakowa"); -} - -#[test] -fn search_cross_thread_messages_finds_japanese_bigram_match() { - let (_temp, store) = make_store(); - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "thread-jp".to_string(), - title: "JP".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "thread-jp", - ConversationMessage { - id: "m1".to_string(), - content: "明日東京に行きます".to_string(), // "Tomorrow I'm going to Tokyo" - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - - let hits = store - .search_cross_thread_messages("東京", 10, None) - .expect("cross-thread search"); - assert_eq!(hits.len(), 1, "CJK bigram lookup should find 東京"); - assert_eq!(hits[0].message_id, "m1"); -} - -#[test] -fn search_cross_thread_messages_rebuilds_index_from_jsonl_after_reopen() { - // First store handle writes messages, second handle (simulating - // process restart on the same workspace dir) must lazy-rebuild the - // index from JSONL and still answer search queries. - let temp = TempDir::new().expect("tempdir"); - let workspace = temp.path().to_path_buf(); - { - let store = ConversationStore::new(workspace.clone()); - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "thread-x".to_string(), - title: "X".to_string(), - created_at: "2026-04-10T12:00:00Z".to_string(), - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "thread-x", - ConversationMessage { - id: "m1".to_string(), - content: "persisted across reopen — checksum kitten".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-04-10T12:01:00Z".to_string(), - }, - ) - .unwrap(); - } - // The cache key is per-workspace path; this TempDir was never seen - // before, so a fresh store handle will trigger a lazy rebuild. - let reopened = ConversationStore::new(workspace); - let hits = reopened - .search_cross_thread_messages("kitten", 10, None) - .expect("cross-thread search"); - assert_eq!(hits.len(), 1, "reopened store must rebuild index from disk"); -} - -#[test] -fn update_thread_labels_missing_thread_returns_error() { - let (_temp, store) = make_store(); - let err = store - .update_thread_labels("missing", vec!["work".into()], "2026-04-10T12:05:00Z") - .unwrap_err(); - assert!(err.contains("thread missing not found")); -} - -#[test] -fn cold_search_does_not_serialize_on_outer_lock() { - // Issue #2849: verify that a cold-cache search releases the store - // lock before the JSONL rebuild, so concurrent writes aren't blocked. - let (_temp, store) = make_store(); - - // Seed a thread with a message so the search has something to find. - store - .ensure_thread(CreateConversationThread { - id: "t1".to_string(), - title: "test thread".to_string(), - created_at: "2026-01-01T00:00:00Z".to_string(), - parent_thread_id: None, - labels: None, - personality_id: None, - }) - .unwrap(); - store - .append_message( - "t1", - ConversationMessage { - id: "m1".to_string(), - content: "hello world".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "user".to_string(), - created_at: "2026-01-01T00:00:00Z".to_string(), - }, - ) - .unwrap(); - - // Evict any warm cache so the next search triggers a cold rebuild. - { - let mut cache = CONVERSATION_INDEX_CACHE.lock(); - cache.remove(&store.root_dir()); - } - - // Spawn a thread that tries to append a message while a cold search - // is (conceptually) running. In the old code this would deadlock or - // serialize behind the full rebuild; in the fixed code the store lock - // is released after the thread-list snapshot and the append succeeds - // concurrently. - let store2 = store.clone(); - let writer = std::thread::spawn(move || { - store2 - .append_message( - "t1", - ConversationMessage { - id: "m2".to_string(), - content: "concurrent write".to_string(), - message_type: "text".to_string(), - extra_metadata: json!({}), - sender: "assistant".to_string(), - created_at: "2026-01-01T00:00:01Z".to_string(), - }, - ) - .unwrap(); - }); - - // Run the cold search — should not deadlock. - let results = store - .search_cross_thread_messages("hello", 10, None) - .unwrap(); - assert!(!results.is_empty(), "search should find seeded message"); - - // The concurrent write must also succeed. - writer.join().expect("concurrent write must not deadlock"); -} - -#[test] -fn read_jsonl_skips_invalid_lines_but_keeps_valid_ones() { - let tmp = TempDir::new().unwrap(); - let path = tmp.path().join("messages.jsonl"); - std::fs::write( - &path, - concat!( - "{\"id\":\"m1\",\"content\":\"ok\",\"type\":\"text\",\"extraMetadata\":{},\"sender\":\"user\",\"createdAt\":\"2026-04-10T12:00:00Z\"}\n", - "{not valid json}\n", - "{\"id\":\"m2\",\"content\":\"ok2\",\"type\":\"text\",\"extraMetadata\":{},\"sender\":\"agent\",\"createdAt\":\"2026-04-10T12:01:00Z\"}\n" - ), - ) - .unwrap(); - - let messages: Vec = read_jsonl(&path).expect("read jsonl"); - assert_eq!(messages.len(), 2); - assert_eq!(messages[0].id, "m1"); - assert_eq!(messages[1].id, "m2"); -} - -// ── concurrency: search cold rebuild must not block concurrent append ──────── - -/// Regression test for issue #2849. -/// -/// Before the fix, `search_cross_thread_messages` held `CONVERSATION_STORE_LOCK` -/// for the entire cold index rebuild, stalling every concurrent -/// `append_message` call for as long as the rebuild took. The fix -/// moves the rebuild outside the outer lock (`prime_index_if_cold`), so -/// an append in flight during a cold rebuild acquires the outer lock -/// independently and completes promptly. -/// -/// The test seeds a fresh workspace (cold cache), races a search against -/// an append using a barrier, and asserts the append finishes within a -/// generous timeout that would be violated if the two operations were -/// serialised through the outer lock. -#[test] -fn search_cold_rebuild_does_not_block_concurrent_append() { - use std::sync::{mpsc, Arc, Barrier}; - use std::thread; - use std::time::Duration; - - let ts = "2026-04-10T12:00:00Z".to_string(); - - // Fresh TempDir → path never seen by the process-level cache → cold. - // Note: append_message only updates an *existing* cache entry; it never - // inserts one, so the cache stays cold until the first search call. - let temp = TempDir::new().unwrap(); - let store = ConversationStore::new(temp.path().to_path_buf()); - - store - .ensure_thread(CreateConversationThread { - parent_thread_id: None, - id: "t1".to_string(), - title: "Rebuild thread".to_string(), - created_at: ts.clone(), - labels: None, - personality_id: None, - }) - .unwrap(); - - // Seed enough messages to give the rebuild real work. - for i in 0..200_usize { - store - .append_message( - "t1", - ConversationMessage { - id: format!("seed-{i}"), - content: format!("seed message {i} for cold rebuild test"), - message_type: "text".to_string(), - extra_metadata: serde_json::json!({}), - sender: "user".to_string(), - created_at: ts.clone(), - }, - ) - .unwrap(); - } - - let store_search = store.clone(); - let store_append = store.clone(); - - // Both threads start at the same time. - let barrier = Arc::new(Barrier::new(2)); - let b_search = Arc::clone(&barrier); - let b_append = Arc::clone(&barrier); - - let search_handle = thread::spawn(move || { - b_search.wait(); - store_search.search_cross_thread_messages("seed message", 5, None) - }); - - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - b_append.wait(); - let result = store_append.append_message( - "t1", - ConversationMessage { - id: "concurrent-append".to_string(), - content: "written during cold rebuild".to_string(), - message_type: "text".to_string(), - extra_metadata: serde_json::json!({}), - sender: "user".to_string(), - created_at: ts, - }, - ); - let _ = tx.send(result); - }); - - // append_message must complete even if the rebuild is in progress. On the - // old code this blocked for the full rebuild duration; on fixed code the - // two operations proceed concurrently. The 30 s budget tolerates a slow CI - // runner — a genuine deadlock never completes, so a regression still fails. - let append_result = rx - .recv_timeout(Duration::from_secs(30)) - .expect("append_message did not complete within 30 s — likely blocked by cold rebuild"); - assert!( - append_result.is_ok(), - "append failed: {:?}", - append_result.err() - ); - - let search_result = search_handle.join().expect("search thread panicked"); - assert!( - search_result.is_ok(), - "search failed: {:?}", - search_result.err() - ); -} - -// ── legacy workspace (pre-Stats backfill path) ─────────────────────────────── - -/// Regression test for issue #2849 (backfill path). -/// -/// A workspace where `threads.jsonl` contains only `Upsert` entries with no -/// `MessageAppended` / `Stats` history is a "pre-Stats" workspace — common for -/// data written before the Stats log was introduced. When -/// `list_threads_unlocked` encounters such threads it calls -/// `measure_messages_unlocked` per thread and appends a `Stats` entry to -/// `threads.jsonl`, all while holding `CONVERSATION_STORE_LOCK`. -/// -/// `prime_index_if_cold` must NOT call `list_threads_unlocked`. It uses -/// `thread_index_unlocked` (header-only, no per-thread I/O) to snapshot -/// thread IDs under the lock, then reads per-thread JSONL content outside -/// the lock. This test verifies that a cold search on such a workspace -/// still finds the correct messages, and that the former blocking code path -/// is no longer reachable from `prime_index_if_cold`. -#[test] -fn prime_index_cold_build_works_on_legacy_workspace_without_stats() { - let temp = TempDir::new().unwrap(); - let store = ConversationStore::new(temp.path().to_path_buf()); - let ts = "2023-06-01T00:00:00Z".to_string(); - - // Bootstrap a legacy workspace by writing directly to the JSONL files — - // bypassing ensure_thread / append_message so no MessageAppended or Stats - // entries end up in threads.jsonl. This is exactly the shape produced by - // versions of the code predating the Stats log. - let root = store.root_dir(); - std::fs::create_dir_all(root.join(THREAD_MESSAGES_DIR)).unwrap(); - - // Write an Upsert-only threads.jsonl (no MessageAppended / Stats entries). - append_jsonl( - &root.join(THREADS_FILENAME), - &ThreadLogEntry::Upsert { - thread_id: "legacy-t1".to_string(), - title: "Legacy Thread".to_string(), - created_at: ts.clone(), - updated_at: ts.clone(), - parent_thread_id: None, - labels: None, - personality_id: None, - }, - ) - .unwrap(); - - // Write messages directly to the per-thread JSONL file, bypassing - // append_message so message_count stays None in the index. - let msg_path = store.thread_messages_path("legacy-t1"); - for i in 0..3_usize { - append_jsonl( - &msg_path, - &ConversationMessage { - id: format!("lm{i}"), - content: format!("legacy kitten message {i}"), - message_type: "text".to_string(), - extra_metadata: serde_json::json!({}), - sender: "user".to_string(), - created_at: ts.clone(), - }, - ) - .unwrap(); - } - - // Cold build on a pre-Stats workspace must index all messages without - // triggering measure_messages_unlocked under CONVERSATION_STORE_LOCK. - let hits = store - .search_cross_thread_messages("kitten", 10, None) - .expect("search on legacy workspace"); - assert_eq!( - hits.len(), - 3, - "all three legacy messages must be found via cold build" - ); - assert!( - hits.iter().any(|h| h.message_id == "lm0"), - "lm0 must be in results" - ); -} - -/// Extends the concurrent-append test to the legacy (no-Stats) workspace shape. -/// -/// Before the fix, `prime_index_if_cold` called `list_threads_unlocked` under -/// the outer lock; for pre-Stats workspaces this triggered a slow -/// `measure_messages_unlocked` + `Stats` append per thread — stalling any -/// concurrent `append_message`. After the fix, `thread_index_unlocked` is -/// used instead (header-only) so the append proceeds concurrently. -#[test] -fn legacy_workspace_cold_rebuild_does_not_block_concurrent_append() { - use std::sync::{mpsc, Arc, Barrier}; - use std::thread; - use std::time::Duration; - - let temp = TempDir::new().unwrap(); - let store = ConversationStore::new(temp.path().to_path_buf()); - let ts = "2023-06-01T00:00:00Z".to_string(); - - // Build a pre-Stats workspace with many threads to make the rebuild - // measurable (each thread has a per-thread JSONL file but no Stats entry). - let root = store.root_dir(); - std::fs::create_dir_all(root.join(THREAD_MESSAGES_DIR)).unwrap(); - - for t in 0..20_usize { - let tid = format!("legacy-t{t}"); - append_jsonl( - &root.join(THREADS_FILENAME), - &ThreadLogEntry::Upsert { - thread_id: tid.clone(), - title: format!("Legacy {t}"), - created_at: ts.clone(), - updated_at: ts.clone(), - parent_thread_id: None, - labels: None, - personality_id: None, - }, - ) - .unwrap(); - let msg_path = store.thread_messages_path(&tid); - for m in 0..50_usize { - append_jsonl( - &msg_path, - &ConversationMessage { - id: format!("lm-{t}-{m}"), - content: format!("legacy content thread {t} message {m}"), - message_type: "text".to_string(), - extra_metadata: serde_json::json!({}), - sender: "user".to_string(), - created_at: ts.clone(), - }, - ) - .unwrap(); - } - } - - // Also need a thread that append_message can target — create it properly - // so it exists in threads.jsonl (still Upsert-only, no Stats). - append_jsonl( - &root.join(THREADS_FILENAME), - &ThreadLogEntry::Upsert { - thread_id: "append-target".to_string(), - title: "Append Target".to_string(), - created_at: ts.clone(), - updated_at: ts.clone(), - parent_thread_id: None, - labels: None, - personality_id: None, - }, - ) - .unwrap(); - - let store_search = store.clone(); - let store_append = store.clone(); - - let barrier = Arc::new(Barrier::new(2)); - let b_search = Arc::clone(&barrier); - let b_append = Arc::clone(&barrier); - - let search_handle = thread::spawn(move || { - b_search.wait(); - store_search.search_cross_thread_messages("legacy content", 5, None) - }); - - let (tx, rx) = mpsc::channel(); - thread::spawn(move || { - b_append.wait(); - let result = store_append.append_message( - "append-target", - ConversationMessage { - id: "concurrent-legacy-append".to_string(), - content: "written during legacy cold rebuild".to_string(), - message_type: "text".to_string(), - extra_metadata: serde_json::json!({}), - sender: "user".to_string(), - created_at: ts, - }, - ); - let _ = tx.send(result); - }); - - let append_result = rx - .recv_timeout(Duration::from_secs(30)) - .expect("append_message blocked — legacy workspace cold rebuild held STORE_LOCK too long"); - assert!( - append_result.is_ok(), - "append failed: {:?}", - append_result.err() - ); - - let search_result = search_handle.join().expect("search thread panicked"); - assert!( - search_result.is_ok(), - "search failed: {:?}", - search_result.err() - ); -} +#[path = "store_tests_more.rs"] +mod more; diff --git a/src/memory/conversations/store_tests_late.rs b/src/memory/conversations/store_tests_late.rs new file mode 100644 index 0000000..4927ae2 --- /dev/null +++ b/src/memory/conversations/store_tests_late.rs @@ -0,0 +1,478 @@ +use super::*; + +#[test] +fn search_cross_thread_messages_finds_japanese_bigram_match() { + let (_temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-jp".to_string(), + title: "JP".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "thread-jp", + ConversationMessage { + id: "m1".to_string(), + content: "明日東京に行きます".to_string(), // "Tomorrow I'm going to Tokyo" + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + let hits = store + .search_cross_thread_messages("東京", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "CJK bigram lookup should find 東京"); + assert_eq!(hits[0].message_id, "m1"); +} + +#[test] +fn search_cross_thread_messages_rebuilds_index_from_jsonl_after_reopen() { + // First store handle writes messages, second handle (simulating + // process restart on the same workspace dir) must lazy-rebuild the + // index from JSONL and still answer search queries. + let temp = TempDir::new().expect("tempdir"); + let workspace = temp.path().to_path_buf(); + { + let store = ConversationStore::new(workspace.clone()); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-x".to_string(), + title: "X".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "thread-x", + ConversationMessage { + id: "m1".to_string(), + content: "persisted across reopen — checksum kitten".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + } + // The cache key is per-workspace path; this TempDir was never seen + // before, so a fresh store handle will trigger a lazy rebuild. + let reopened = ConversationStore::new(workspace); + let hits = reopened + .search_cross_thread_messages("kitten", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "reopened store must rebuild index from disk"); +} + +#[test] +fn update_thread_labels_missing_thread_returns_error() { + let (_temp, store) = make_store(); + let err = store + .update_thread_labels("missing", vec!["work".into()], "2026-04-10T12:05:00Z") + .unwrap_err(); + assert!(err.contains("thread missing not found")); +} + +#[test] +fn cold_search_does_not_serialize_on_outer_lock() { + // Issue #2849: verify that a cold-cache search releases the store + // lock before the JSONL rebuild, so concurrent writes aren't blocked. + let (_temp, store) = make_store(); + + // Seed a thread with a message so the search has something to find. + store + .ensure_thread(CreateConversationThread { + id: "t1".to_string(), + title: "test thread".to_string(), + created_at: "2026-01-01T00:00:00Z".to_string(), + parent_thread_id: None, + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "t1", + ConversationMessage { + id: "m1".to_string(), + content: "hello world".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-01-01T00:00:00Z".to_string(), + }, + ) + .unwrap(); + + // Evict any warm cache so the next search triggers a cold rebuild. + { + let mut cache = CONVERSATION_INDEX_CACHE.lock(); + cache.remove(&store.root_dir()); + } + + // Spawn a thread that tries to append a message while a cold search + // is (conceptually) running. In the old code this would deadlock or + // serialize behind the full rebuild; in the fixed code the store lock + // is released after the thread-list snapshot and the append succeeds + // concurrently. + let store2 = store.clone(); + let writer = std::thread::spawn(move || { + store2 + .append_message( + "t1", + ConversationMessage { + id: "m2".to_string(), + content: "concurrent write".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "assistant".to_string(), + created_at: "2026-01-01T00:00:01Z".to_string(), + }, + ) + .unwrap(); + }); + + // Run the cold search — should not deadlock. + let results = store + .search_cross_thread_messages("hello", 10, None) + .unwrap(); + assert!(!results.is_empty(), "search should find seeded message"); + + // The concurrent write must also succeed. + writer.join().expect("concurrent write must not deadlock"); +} + +#[test] +fn read_jsonl_skips_invalid_lines_but_keeps_valid_ones() { + let tmp = TempDir::new().unwrap(); + let path = tmp.path().join("messages.jsonl"); + std::fs::write( + &path, + concat!( + "{\"id\":\"m1\",\"content\":\"ok\",\"type\":\"text\",\"extraMetadata\":{},\"sender\":\"user\",\"createdAt\":\"2026-04-10T12:00:00Z\"}\n", + "{not valid json}\n", + "{\"id\":\"m2\",\"content\":\"ok2\",\"type\":\"text\",\"extraMetadata\":{},\"sender\":\"agent\",\"createdAt\":\"2026-04-10T12:01:00Z\"}\n" + ), + ) + .unwrap(); + + let messages: Vec = read_jsonl(&path).expect("read jsonl"); + assert_eq!(messages.len(), 2); + assert_eq!(messages[0].id, "m1"); + assert_eq!(messages[1].id, "m2"); +} + +// ── concurrency: search cold rebuild must not block concurrent append ──────── + +/// Regression test for issue #2849. +/// +/// Before the fix, `search_cross_thread_messages` held `CONVERSATION_STORE_LOCK` +/// for the entire cold index rebuild, stalling every concurrent +/// `append_message` call for as long as the rebuild took. The fix +/// moves the rebuild outside the outer lock (`prime_index_if_cold`), so +/// an append in flight during a cold rebuild acquires the outer lock +/// independently and completes promptly. +/// +/// The test seeds a fresh workspace (cold cache), races a search against +/// an append using a barrier, and asserts the append finishes within a +/// generous timeout that would be violated if the two operations were +/// serialised through the outer lock. +#[test] +fn search_cold_rebuild_does_not_block_concurrent_append() { + use std::sync::{mpsc, Arc, Barrier}; + use std::thread; + use std::time::Duration; + + let ts = "2026-04-10T12:00:00Z".to_string(); + + // Fresh TempDir → path never seen by the process-level cache → cold. + // Note: append_message only updates an *existing* cache entry; it never + // inserts one, so the cache stays cold until the first search call. + let temp = TempDir::new().unwrap(); + let store = ConversationStore::new(temp.path().to_path_buf()); + + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "t1".to_string(), + title: "Rebuild thread".to_string(), + created_at: ts.clone(), + labels: None, + personality_id: None, + }) + .unwrap(); + + // Seed enough messages to give the rebuild real work. + for i in 0..200_usize { + store + .append_message( + "t1", + ConversationMessage { + id: format!("seed-{i}"), + content: format!("seed message {i} for cold rebuild test"), + message_type: "text".to_string(), + extra_metadata: serde_json::json!({}), + sender: "user".to_string(), + created_at: ts.clone(), + }, + ) + .unwrap(); + } + + let store_search = store.clone(); + let store_append = store.clone(); + + // Both threads start at the same time. + let barrier = Arc::new(Barrier::new(2)); + let b_search = Arc::clone(&barrier); + let b_append = Arc::clone(&barrier); + + let search_handle = thread::spawn(move || { + b_search.wait(); + store_search.search_cross_thread_messages("seed message", 5, None) + }); + + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + b_append.wait(); + let result = store_append.append_message( + "t1", + ConversationMessage { + id: "concurrent-append".to_string(), + content: "written during cold rebuild".to_string(), + message_type: "text".to_string(), + extra_metadata: serde_json::json!({}), + sender: "user".to_string(), + created_at: ts, + }, + ); + let _ = tx.send(result); + }); + + // append_message must complete even if the rebuild is in progress. On the + // old code this blocked for the full rebuild duration; on fixed code the + // two operations proceed concurrently. The 30 s budget tolerates a slow CI + // runner — a genuine deadlock never completes, so a regression still fails. + let append_result = rx + .recv_timeout(Duration::from_secs(30)) + .expect("append_message did not complete within 30 s — likely blocked by cold rebuild"); + assert!( + append_result.is_ok(), + "append failed: {:?}", + append_result.err() + ); + + let search_result = search_handle.join().expect("search thread panicked"); + assert!( + search_result.is_ok(), + "search failed: {:?}", + search_result.err() + ); +} + +// ── legacy workspace (pre-Stats backfill path) ─────────────────────────────── + +/// Regression test for issue #2849 (backfill path). +/// +/// A workspace where `threads.jsonl` contains only `Upsert` entries with no +/// `MessageAppended` / `Stats` history is a "pre-Stats" workspace — common for +/// data written before the Stats log was introduced. When +/// `list_threads_unlocked` encounters such threads it calls +/// `measure_messages_unlocked` per thread and appends a `Stats` entry to +/// `threads.jsonl`, all while holding `CONVERSATION_STORE_LOCK`. +/// +/// `prime_index_if_cold` must NOT call `list_threads_unlocked`. It uses +/// `thread_index_unlocked` (header-only, no per-thread I/O) to snapshot +/// thread IDs under the lock, then reads per-thread JSONL content outside +/// the lock. This test verifies that a cold search on such a workspace +/// still finds the correct messages, and that the former blocking code path +/// is no longer reachable from `prime_index_if_cold`. +#[test] +fn prime_index_cold_build_works_on_legacy_workspace_without_stats() { + let temp = TempDir::new().unwrap(); + let store = ConversationStore::new(temp.path().to_path_buf()); + let ts = "2023-06-01T00:00:00Z".to_string(); + + // Bootstrap a legacy workspace by writing directly to the JSONL files — + // bypassing ensure_thread / append_message so no MessageAppended or Stats + // entries end up in threads.jsonl. This is exactly the shape produced by + // versions of the code predating the Stats log. + let root = store.root_dir(); + std::fs::create_dir_all(root.join(THREAD_MESSAGES_DIR)).unwrap(); + + // Write an Upsert-only threads.jsonl (no MessageAppended / Stats entries). + append_jsonl( + &root.join(THREADS_FILENAME), + &ThreadLogEntry::Upsert { + thread_id: "legacy-t1".to_string(), + title: "Legacy Thread".to_string(), + created_at: ts.clone(), + updated_at: ts.clone(), + parent_thread_id: None, + labels: None, + personality_id: None, + }, + ) + .unwrap(); + + // Write messages directly to the per-thread JSONL file, bypassing + // append_message so message_count stays None in the index. + let msg_path = store.thread_messages_path("legacy-t1"); + for i in 0..3_usize { + append_jsonl( + &msg_path, + &ConversationMessage { + id: format!("lm{i}"), + content: format!("legacy kitten message {i}"), + message_type: "text".to_string(), + extra_metadata: serde_json::json!({}), + sender: "user".to_string(), + created_at: ts.clone(), + }, + ) + .unwrap(); + } + + // Cold build on a pre-Stats workspace must index all messages without + // triggering measure_messages_unlocked under CONVERSATION_STORE_LOCK. + let hits = store + .search_cross_thread_messages("kitten", 10, None) + .expect("search on legacy workspace"); + assert_eq!( + hits.len(), + 3, + "all three legacy messages must be found via cold build" + ); + assert!( + hits.iter().any(|h| h.message_id == "lm0"), + "lm0 must be in results" + ); +} + +/// Extends the concurrent-append test to the legacy (no-Stats) workspace shape. +/// +/// Before the fix, `prime_index_if_cold` called `list_threads_unlocked` under +/// the outer lock; for pre-Stats workspaces this triggered a slow +/// `measure_messages_unlocked` + `Stats` append per thread — stalling any +/// concurrent `append_message`. After the fix, `thread_index_unlocked` is +/// used instead (header-only) so the append proceeds concurrently. +#[test] +fn legacy_workspace_cold_rebuild_does_not_block_concurrent_append() { + use std::sync::{mpsc, Arc, Barrier}; + use std::thread; + use std::time::Duration; + + let temp = TempDir::new().unwrap(); + let store = ConversationStore::new(temp.path().to_path_buf()); + let ts = "2023-06-01T00:00:00Z".to_string(); + + // Build a pre-Stats workspace with many threads to make the rebuild + // measurable (each thread has a per-thread JSONL file but no Stats entry). + let root = store.root_dir(); + std::fs::create_dir_all(root.join(THREAD_MESSAGES_DIR)).unwrap(); + + for t in 0..20_usize { + let tid = format!("legacy-t{t}"); + append_jsonl( + &root.join(THREADS_FILENAME), + &ThreadLogEntry::Upsert { + thread_id: tid.clone(), + title: format!("Legacy {t}"), + created_at: ts.clone(), + updated_at: ts.clone(), + parent_thread_id: None, + labels: None, + personality_id: None, + }, + ) + .unwrap(); + let msg_path = store.thread_messages_path(&tid); + for m in 0..50_usize { + append_jsonl( + &msg_path, + &ConversationMessage { + id: format!("lm-{t}-{m}"), + content: format!("legacy content thread {t} message {m}"), + message_type: "text".to_string(), + extra_metadata: serde_json::json!({}), + sender: "user".to_string(), + created_at: ts.clone(), + }, + ) + .unwrap(); + } + } + + // Also need a thread that append_message can target — create it properly + // so it exists in threads.jsonl (still Upsert-only, no Stats). + append_jsonl( + &root.join(THREADS_FILENAME), + &ThreadLogEntry::Upsert { + thread_id: "append-target".to_string(), + title: "Append Target".to_string(), + created_at: ts.clone(), + updated_at: ts.clone(), + parent_thread_id: None, + labels: None, + personality_id: None, + }, + ) + .unwrap(); + + let store_search = store.clone(); + let store_append = store.clone(); + + let barrier = Arc::new(Barrier::new(2)); + let b_search = Arc::clone(&barrier); + let b_append = Arc::clone(&barrier); + + let search_handle = thread::spawn(move || { + b_search.wait(); + store_search.search_cross_thread_messages("legacy content", 5, None) + }); + + let (tx, rx) = mpsc::channel(); + thread::spawn(move || { + b_append.wait(); + let result = store_append.append_message( + "append-target", + ConversationMessage { + id: "concurrent-legacy-append".to_string(), + content: "written during legacy cold rebuild".to_string(), + message_type: "text".to_string(), + extra_metadata: serde_json::json!({}), + sender: "user".to_string(), + created_at: ts, + }, + ); + let _ = tx.send(result); + }); + + let append_result = rx + .recv_timeout(Duration::from_secs(30)) + .expect("append_message blocked — legacy workspace cold rebuild held STORE_LOCK too long"); + assert!( + append_result.is_ok(), + "append failed: {:?}", + append_result.err() + ); + + let search_result = search_handle.join().expect("search thread panicked"); + assert!( + search_result.is_ok(), + "search failed: {:?}", + search_result.err() + ); +} diff --git a/src/memory/conversations/store_tests_more.rs b/src/memory/conversations/store_tests_more.rs new file mode 100644 index 0000000..0dcbe2d --- /dev/null +++ b/src/memory/conversations/store_tests_more.rs @@ -0,0 +1,427 @@ +use super::*; + +#[test] +fn conversation_store_new() { + let tmp = TempDir::new().unwrap(); + let store = ConversationStore::new(tmp.path().to_path_buf()); + let threads = store.list_threads().unwrap(); + assert!(threads.is_empty()); +} + +#[test] +fn conversation_purge_stats_default() { + let stats = ConversationPurgeStats::default(); + assert_eq!(stats.thread_count, 0); + assert_eq!(stats.message_count, 0); +} + +#[test] +fn list_threads_reconciles_stats_with_authoritative_message_files() { + let (temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "t1".to_string(), + title: "T1".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + for i in 0..3 { + store + .append_message( + "t1", + ConversationMessage { + id: format!("m{i}"), + content: format!("hi {i}"), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: format!("2026-04-10T12:0{}:00Z", i + 1), + }, + ) + .unwrap(); + } + // Warm-up: list_threads folds the MessageAppended entries. + let _ = store.list_threads().unwrap(); + + // Removing the authoritative transcript must not leave stale derived + // counts behind. + let messages_dir = temp + .path() + .join("memory") + .join("conversations") + .join("threads"); + let entries: Vec<_> = std::fs::read_dir(&messages_dir) + .unwrap() + .filter_map(Result::ok) + .collect(); + for entry in entries { + std::fs::remove_file(entry.path()).unwrap(); + } + + let threads = store.list_threads().unwrap(); + assert_eq!(threads.len(), 1); + assert_eq!(threads[0].message_count, 0); + assert_eq!(threads[0].last_message_at, "2026-04-10T12:00:00Z"); +} + +#[test] +fn backfill_writes_stats_snapshot_for_legacy_threads() { + // Simulate legacy data: write only an Upsert entry (no MessageAppended) + // plus a per-thread messages file. The first list_threads must backfill. + let (temp, store) = make_store(); + let conversations_dir = temp.path().join("memory").join("conversations"); + std::fs::create_dir_all(conversations_dir.join("threads")).unwrap(); + + let threads_log = conversations_dir.join("threads.jsonl"); + let upsert = serde_json::json!({ + "op": "upsert", + "thread_id": "legacy-1", + "title": "Legacy", + "created_at": "2026-04-10T08:00:00Z", + "updated_at": "2026-04-10T08:00:00Z", + }); + std::fs::write(&threads_log, format!("{}\n", upsert)).unwrap(); + + // Write 2 messages directly to the per-thread file (no MessageAppended + // entries — this is what pre-upgrade data looks like). + let messages_file = conversations_dir + .join("threads") + .join(format!("{}.jsonl", hex_encode("legacy-1".as_bytes()))); + let m1 = serde_json::json!({ + "id": "m1", "content": "a", "type": "text", + "extraMetadata": {}, "sender": "user", + "createdAt": "2026-04-10T09:00:00Z", + }); + let m2 = serde_json::json!({ + "id": "m2", "content": "b", "type": "text", + "extraMetadata": {}, "sender": "user", + "createdAt": "2026-04-10T09:05:00Z", + }); + std::fs::write(&messages_file, format!("{m1}\n{m2}\n")).unwrap(); + + let threads = store.list_threads().unwrap(); + assert_eq!(threads.len(), 1); + assert_eq!(threads[0].message_count, 2); + assert_eq!(threads[0].last_message_at, "2026-04-10T09:05:00Z"); + + // The backfill should have appended a Stats entry — check the log + // contents now contain "op":"stats" for legacy-1. + let log = std::fs::read_to_string(&threads_log).unwrap(); + assert!( + log.contains("\"op\":\"stats\"") && log.contains("legacy-1"), + "expected backfilled Stats entry in threads.jsonl, got:\n{log}", + ); + + // The message file remains authoritative even after a Stats snapshot. + std::fs::remove_file(&messages_file).unwrap(); + let threads2 = store.list_threads().unwrap(); + assert_eq!(threads2[0].message_count, 0); + assert_eq!(threads2[0].last_message_at, "2026-04-10T08:00:00Z"); +} + +#[test] +fn list_threads_repairs_message_append_without_matching_stat_event() { + let (temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "crash-window".to_string(), + title: "Crash window".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "crash-window", + ConversationMessage { + id: "m1".into(), + content: "first".into(), + message_type: "text".into(), + extra_metadata: json!({}), + sender: "user".into(), + created_at: "2026-04-10T12:01:00Z".into(), + }, + ) + .unwrap(); + + let path = temp + .path() + .join("memory/conversations/threads") + .join(format!("{}.jsonl", hex_encode("crash-window".as_bytes()))); + let second = ConversationMessage { + id: "m2".into(), + content: "persisted before crash".into(), + message_type: "text".into(), + extra_metadata: json!({}), + sender: "assistant".into(), + created_at: "2026-04-10T12:02:00Z".into(), + }; + use std::io::Write; + writeln!( + std::fs::OpenOptions::new().append(true).open(path).unwrap(), + "{}", + serde_json::to_string(&second).unwrap() + ) + .unwrap(); + + let thread = store.list_threads().unwrap().remove(0); + assert_eq!(thread.message_count, 2); + assert_eq!(thread.last_message_at, "2026-04-10T12:02:00Z"); +} + +#[test] +fn legacy_log_without_stats_still_parses() { + // Old on-disk format (only Upsert + Delete variants) must still load + // without errors after the enum gained MessageAppended + Stats. + let (temp, store) = make_store(); + let conversations_dir = temp.path().join("memory").join("conversations"); + std::fs::create_dir_all(conversations_dir.join("threads")).unwrap(); + let threads_log = conversations_dir.join("threads.jsonl"); + let upsert = serde_json::json!({ + "op": "upsert", + "thread_id": "old", + "title": "Old", + "created_at": "2026-04-10T08:00:00Z", + "updated_at": "2026-04-10T08:00:00Z", + }); + std::fs::write(&threads_log, format!("{}\n", upsert)).unwrap(); + + let threads = store.list_threads().unwrap(); + assert_eq!(threads.len(), 1); + assert_eq!(threads[0].id, "old"); + assert_eq!(threads[0].message_count, 0); + // No messages → last_message_at falls back to created_at. + assert_eq!(threads[0].last_message_at, "2026-04-10T08:00:00Z"); +} + +#[test] +fn delete_thread_clears_stats_from_index() { + let (_temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "doomed".to_string(), + title: "Doomed".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "doomed", + ConversationMessage { + id: "m1".to_string(), + content: "x".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + assert_eq!(store.list_threads().unwrap().len(), 1); + + store + .delete_thread("doomed", "2026-04-10T12:02:00Z") + .unwrap(); + assert!(store.list_threads().unwrap().is_empty()); +} + +#[test] +fn search_cross_thread_messages_finds_hits_outside_excluded_thread() { + let (_temp, store) = make_store(); + + // Chat A — durable fact lives here. + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-a".to_string(), + title: "Chat A".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "thread-a", + ConversationMessage { + id: "m-a-1".to_string(), + content: "Remember: my project is called Phoenix and uses Go and PostgreSQL." + .to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + // Chat B — active chat, asking dependent question. Should be excluded + // so its own text doesn't echo back into [Cross-chat context]. + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-b".to_string(), + title: "Chat B".to_string(), + created_at: "2026-04-10T13:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "thread-b", + ConversationMessage { + id: "m-b-1".to_string(), + content: "What database does my project use?".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T13:01:00Z".to_string(), + }, + ) + .unwrap(); + + let hits = store + .search_cross_thread_messages("What database does my project use", 10, Some("thread-b")) + .expect("cross-thread search"); + + assert_eq!(hits.len(), 1, "exactly one cross-thread hit"); + let hit = &hits[0]; + assert_eq!(hit.thread_id, "thread-a"); + assert!(hit.content.contains("PostgreSQL")); + assert!(hit.score > 0.0); +} + +#[test] +fn search_cross_thread_messages_excludes_active_thread() { + let (_temp, store) = make_store(); + + // Single thread — the only matching message lives in the thread we're + // about to exclude. Expect zero hits (don't echo same-chat history). + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-only".to_string(), + title: "Only".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "thread-only", + ConversationMessage { + id: "m-1".to_string(), + content: "PostgreSQL deployment running on staging".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + let hits = store + .search_cross_thread_messages("PostgreSQL deployment staging", 10, Some("thread-only")) + .expect("cross-thread search"); + assert!( + hits.is_empty(), + "active thread must not echo into cross-chat" + ); + + // Sanity: without exclude, the hit is returned. + let hits_no_exclude = store + .search_cross_thread_messages("PostgreSQL deployment staging", 10, None) + .expect("cross-thread search"); + assert_eq!(hits_no_exclude.len(), 1); +} + +#[test] +fn search_cross_thread_messages_skips_short_terms_and_empty_queries() { + let (_temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "t".to_string(), + title: "T".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "t", + ConversationMessage { + id: "m".to_string(), + content: "Postgres".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + // All terms < 3 chars → empty + assert!(store + .search_cross_thread_messages("a is on", 10, None) + .unwrap() + .is_empty()); + // Empty query → empty + assert!(store + .search_cross_thread_messages("", 10, None) + .unwrap() + .is_empty()); +} + +#[test] +fn search_cross_thread_messages_finds_polish_substring_without_diacritics() { + let (_temp, store) = make_store(); + store + .ensure_thread(CreateConversationThread { + parent_thread_id: None, + id: "thread-pl".to_string(), + title: "PL".to_string(), + created_at: "2026-04-10T12:00:00Z".to_string(), + labels: None, + personality_id: None, + }) + .unwrap(); + store + .append_message( + "thread-pl", + ConversationMessage { + id: "m1".to_string(), + content: "Lecę w piątek do Łodzi a potem Krakowa".to_string(), + message_type: "text".to_string(), + extra_metadata: json!({}), + sender: "user".to_string(), + created_at: "2026-04-10T12:01:00Z".to_string(), + }, + ) + .unwrap(); + + // Query without diacritics should still find content with them. + let hits = store + .search_cross_thread_messages("Lodzi", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "ł-fold should match Łodzi via lodzi"); + + let hits = store + .search_cross_thread_messages("krakow", 10, None) + .expect("cross-thread search"); + assert_eq!(hits.len(), 1, "diacritic strip should match Krakowa"); +} + +#[path = "store_tests_late.rs"] +mod late; diff --git a/src/memory/diff/ledger.rs b/src/memory/diff/ledger.rs index dc399e1..f538a7b 100644 --- a/src/memory/diff/ledger.rs +++ b/src/memory/diff/ledger.rs @@ -15,21 +15,22 @@ //! - **Diff** → `git diff ..` scoped to the source path //! //! Item identity is the file name: each item is one flat blob whose name is the -//! item id encoded into a git-safe path component ([`encode_item_id`]). A +//! item id encoded into a git-safe path component. A //! content change keeps the same name → `Modified`; renaming the item id is //! `Removed` + `Added`, matching the id-keyed semantics. //! //! Snapshot metadata that has no natural git home (source kind/label, trigger, //! item count, millisecond timestamp) rides in the commit message as trailers. -//! All mutations serialise through a process-global [`WRITE_LOCK`] because the +//! All mutations serialise through a process-global write lock because the //! repository's parent/HEAD bookkeeping is not safe to interleave. -use std::collections::HashMap; use std::path::Path; use std::sync::Mutex; use anyhow::{Context, Result}; -use git2::{Delta, DiffOptions, Object, ObjectType, Oid, Repository, Signature, Time}; +use git2::{Delta, DiffOptions, Object, ObjectType, Oid, Repository}; + +use super::ledger_helpers::*; use super::types::{ChangeKind, Checkpoint, DiffSummary, ItemChange, Snapshot, SnapshotTrigger}; @@ -40,13 +41,13 @@ static WRITE_LOCK: Mutex<()> = Mutex::new(()); const BLOB_MODE: i32 = 0o100644; const TREE_MODE: i32 = 0o040000; -const SIG_NAME: &str = "TinyCortex Memory"; -const SIG_EMAIL: &str = "memory-diff@tinycortex.local"; -const READ_MARKER_PREFIX: &str = "refs/openhuman/read/"; +pub(super) const SIG_NAME: &str = "TinyCortex Memory"; +pub(super) const SIG_EMAIL: &str = "memory-diff@tinycortex.local"; +pub(super) const READ_MARKER_PREFIX: &str = "refs/openhuman/read/"; const CHECKPOINT_PREFIX: &str = "ckpt_"; /// Upper bound on a single modified-item unified diff embedded in `text_diff`. -const MAX_TEXT_DIFF_CHARS: usize = 2000; +pub(super) const MAX_TEXT_DIFF_CHARS: usize = 2000; // ── Repository handle ────────────────────────────────────────────────── @@ -57,7 +58,7 @@ pub struct Ledger { } /// Metadata describing the snapshot a commit represents. Persisted as commit -/// trailers and reconstructed by [`Ledger::snapshot_from_commit`]. +/// trailers and reconstructed by the ledger's commit parser. pub struct SnapshotMeta { /// Logical source id. pub source_id: String, @@ -72,18 +73,18 @@ pub struct SnapshotMeta { impl Ledger { /// Open the ledger, initialising the repository on first use. /// - /// NOTE: any `Repository::open` failure — not only "repo doesn't exist - /// yet" — falls through to `Repository::init` on the same path. A - /// transiently-locked or genuinely corrupt repository is therefore - /// silently re-initialised rather than surfaced as an error, which can - /// discard ledger history instead of reporting the real problem. pub fn open(workspace_dir: &Path) -> Result { let repo_path = workspace_dir.join("memory_diff").join("repo"); std::fs::create_dir_all(&repo_path) .with_context(|| format!("create memory_diff repo dir: {}", repo_path.display()))?; + let git_marker = repo_path.join(".git"); let repo = match Repository::open(&repo_path) { Ok(repo) => repo, + Err(err) if git_marker.exists() => { + return Err(err) + .with_context(|| format!("open memory_diff repo: {}", repo_path.display())); + } Err(_) => Repository::init(&repo_path) .with_context(|| format!("init memory_diff repo: {}", repo_path.display()))?, }; @@ -102,6 +103,7 @@ impl Ledger { taken_at_ms: i64, ) -> Result { let _guard = WRITE_LOCK.lock().expect("memory_diff write lock poisoned"); + validate_source_id(&meta.source_id)?; // Build the source subtree from scratch: one blob per item. let source_tree_oid = { @@ -328,6 +330,28 @@ impl Ledger { let oid = Oid::from_str(snapshot_id) .with_context(|| format!("bad read-marker snapshot id: {snapshot_id}"))?; let name = read_marker_ref(source_id); + let commit = self + .repo + .find_commit(oid) + .with_context(|| format!("read-marker snapshot not found: {snapshot_id}"))?; + let snapshot = self.snapshot_from_commit(&commit); + anyhow::ensure!( + snapshot.source_id == source_id, + "snapshot {snapshot_id} belongs to source '{}', not '{source_id}'", + snapshot.source_id + ); + + if let Ok(current) = self.repo.find_reference(&name) { + if let Some(current_oid) = current.target() { + if current_oid == oid { + return Ok(()); + } + anyhow::ensure!( + self.repo.graph_descendant_of(oid, current_oid)?, + "refusing to move read marker for '{source_id}' backwards" + ); + } + } self.repo .reference(&name, oid, true, "advance memory_diff read marker") .with_context(|| format!("set read marker ref: {name}"))?; @@ -375,7 +399,7 @@ impl Ledger { // git2 0.21: Tag::message() is Result, _>; a non-UTF8 // or missing message degrades to an empty checkpoint body. tag.message().ok().flatten().unwrap_or(""), - ))) + )?)) } /// List checkpoints newest-first, up to `limit`. @@ -458,183 +482,6 @@ impl Ledger { } } -// ── Free helpers ─────────────────────────────────────────────────────── - -fn signature(at_ms: i64) -> Result> { - let time = Time::new(at_ms / 1000, 0); - Signature::new(SIG_NAME, SIG_EMAIL, &time).context("build git signature") -} - -fn read_marker_ref(source_id: &str) -> String { - format!("{READ_MARKER_PREFIX}{}", encode_source_id(source_id)) -} - -fn build_commit_message(meta: &SnapshotMeta, item_count: u32, taken_at_ms: i64) -> String { - format!( - "snapshot: {source} ({count} item(s))\n\n\ - Source-Id: {source}\n\ - Source-Kind: {kind}\n\ - Source-Label: {label}\n\ - Trigger: {trigger}\n\ - Item-Count: {count}\n\ - Taken-At-Ms: {taken}\n", - source = meta.source_id, - kind = meta.source_kind, - label = sanitize_trailer(&meta.label), - trigger = meta.trigger.as_str(), - count = item_count, - taken = taken_at_ms, - ) -} - -/// Trailer values are single-line; collapse newlines so a multi-line label -/// can't corrupt the trailer block. -fn sanitize_trailer(s: &str) -> String { - s.replace(['\n', '\r'], " ") -} - -/// Parse `Key: value` trailer lines from a commit message into a -/// lowercase-keyed map. -fn parse_trailers(message: &str) -> HashMap { - let mut map = HashMap::new(); - for line in message.lines() { - if let Some((k, v)) = line.split_once(':') { - let key = k.trim().to_ascii_lowercase(); - if !key.is_empty() && !key.contains(' ') { - map.insert(key, v.trim().to_string()); - } - } - } - map -} - -fn checkpoint_message(label: &str, snapshot_ids: &[String], created_at_ms: i64) -> String { - let payload = serde_json::json!({ - "label": label, - "snapshot_ids": snapshot_ids, - "created_at_ms": created_at_ms, - }); - payload.to_string() -} - -fn checkpoint_from_message(id: &str, message: &str) -> Checkpoint { - let value: serde_json::Value = serde_json::from_str(message.trim()).unwrap_or_default(); - Checkpoint { - id: id.to_string(), - label: value - .get("label") - .and_then(|v| v.as_str()) - .unwrap_or("") - .to_string(), - created_at_ms: value - .get("created_at_ms") - .and_then(|v| v.as_i64()) - .unwrap_or(0), - snapshot_ids: value - .get("snapshot_ids") - .and_then(|v| v.as_array()) - .map(|arr| { - arr.iter() - .filter_map(|v| v.as_str().map(str::to_string)) - .collect() - }) - .unwrap_or_default(), - } -} - -/// A git blob oid as a content hash, or `None` for the zero oid (absent side). -fn oid_hash(oid: Oid) -> Option { - if oid.is_zero() { - None - } else { - Some(oid.to_string()) - } -} - -/// Render a single delta's unified patch, truncated to [`MAX_TEXT_DIFF_CHARS`]. -fn patch_text(diff: &git2::Diff, delta_idx: usize) -> Option { - let mut patch = git2::Patch::from_diff(diff, delta_idx).ok().flatten()?; - let buf = patch.to_buf().ok()?; - // git2 0.21: Buf::as_str() returns Result<&str, Utf8Error>. - let text = buf.as_str().ok()?; - if text.trim().is_empty() { - None - } else { - Some(truncate(text, MAX_TEXT_DIFF_CHARS)) - } -} - -/// Encode an item id into a single git-safe path component. Bytes outside -/// `[A-Za-z0-9._-]` become `%XX`; an `i_` prefix keeps the result clear of the -/// reserved names `.`/`..`/empty. Reversible via [`decode_item_id`]. -pub(crate) fn encode_item_id(item_id: &str) -> String { - let mut out = String::with_capacity(item_id.len() + 2); - out.push_str("i_"); - for &b in item_id.as_bytes() { - if b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-') { - out.push(b as char); - } else { - out.push('%'); - out.push_str(&format!("{b:02X}")); - } - } - out -} - -/// Encode a source id for use as a top-level git tree entry and read-marker -/// ref component. Same reversible encoding as item ids; kept as a named helper -/// so call sites make the source-vs-item boundary explicit. -pub(crate) fn encode_source_id(source_id: &str) -> String { - encode_item_id(source_id) -} - -/// Inverse of [`encode_item_id`]. -pub(crate) fn decode_item_id(encoded: &str) -> String { - let body = encoded.strip_prefix("i_").unwrap_or(encoded); - let bytes = body.as_bytes(); - let mut out = Vec::with_capacity(bytes.len()); - let mut i = 0; - while i < bytes.len() { - if bytes[i] == b'%' && i + 2 < bytes.len() { - let hex = std::str::from_utf8(&bytes[i + 1..i + 3]).unwrap_or(""); - if let Ok(byte) = u8::from_str_radix(hex, 16) { - out.push(byte); - i += 3; - continue; - } - } - out.push(bytes[i]); - i += 1; - } - String::from_utf8_lossy(&out).into_owned() -} - -fn truncate(s: &str, max_chars: usize) -> String { - if s.len() <= max_chars { - s.to_string() - } else { - let mut end = max_chars; - while !s.is_char_boundary(end) && end > 0 { - end -= 1; - } - format!("{}…(truncated)", &s[..end]) - } -} - -/// Derive a human-readable title from item content: the first non-empty line -/// (Markdown heading markers stripped), bounded. Falls back to the item id. -fn derive_title(item_id: &str, content: &str) -> String { - let first_line = content - .lines() - .map(str::trim) - .find(|l| !l.is_empty()) - .map(|l| l.trim_start_matches('#').trim()); - match first_line { - Some(l) if !l.is_empty() => truncate(l, 120), - _ => item_id.to_string(), - } -} - #[cfg(test)] #[path = "ledger_tests.rs"] mod tests; diff --git a/src/memory/diff/ledger_helpers.rs b/src/memory/diff/ledger_helpers.rs new file mode 100644 index 0000000..92e1865 --- /dev/null +++ b/src/memory/diff/ledger_helpers.rs @@ -0,0 +1,215 @@ +//! Serialization, identity, and patch helpers for the git ledger. + +use std::collections::HashMap; + +use anyhow::{Context, Result}; +use git2::{Oid, Signature, Time}; + +use super::ledger::{SnapshotMeta, MAX_TEXT_DIFF_CHARS, READ_MARKER_PREFIX, SIG_EMAIL, SIG_NAME}; +use super::types::Checkpoint; + +// ── Free helpers ─────────────────────────────────────────────────────── + +pub(super) fn signature(at_ms: i64) -> Result> { + let time = Time::new(at_ms / 1000, 0); + Signature::new(SIG_NAME, SIG_EMAIL, &time).context("build git signature") +} + +pub(super) fn read_marker_ref(source_id: &str) -> String { + format!("{READ_MARKER_PREFIX}{}", encode_source_id(source_id)) +} + +pub(super) fn build_commit_message( + meta: &SnapshotMeta, + item_count: u32, + taken_at_ms: i64, +) -> String { + format!( + "snapshot: {source} ({count} item(s))\n\n\ + Source-Id: {source}\n\ + Source-Kind: {kind}\n\ + Source-Label: {label}\n\ + Trigger: {trigger}\n\ + Item-Count: {count}\n\ + Taken-At-Ms: {taken}\n", + source = meta.source_id, + kind = meta.source_kind, + label = sanitize_trailer(&meta.label), + trigger = meta.trigger.as_str(), + count = item_count, + taken = taken_at_ms, + ) +} + +/// Trailer values are single-line; collapse newlines so a multi-line label +/// can't corrupt the trailer block. +pub(super) fn sanitize_trailer(s: &str) -> String { + s.replace(['\n', '\r'], " ") +} + +/// Parse `Key: value` trailer lines from a commit message into a +/// lowercase-keyed map. +pub(super) fn parse_trailers(message: &str) -> HashMap { + let mut map = HashMap::new(); + let trailer_block = message + .trim_end() + .rsplit_once("\n\n") + .map_or(message, |(_, block)| block); + for line in trailer_block.lines() { + if let Some((k, v)) = line.split_once(':') { + let key = k.trim().to_ascii_lowercase(); + if !key.is_empty() && !key.contains(' ') { + map.insert(key, v.trim().to_string()); + } + } + } + map +} + +pub(super) fn validate_source_id(source_id: &str) -> Result<()> { + anyhow::ensure!(!source_id.trim().is_empty(), "source id must not be blank"); + anyhow::ensure!( + !source_id.chars().any(char::is_control), + "source id must not contain control characters" + ); + Ok(()) +} + +pub(super) fn checkpoint_message( + label: &str, + snapshot_ids: &[String], + created_at_ms: i64, +) -> String { + let payload = serde_json::json!({ + "label": label, + "snapshot_ids": snapshot_ids, + "created_at_ms": created_at_ms, + }); + payload.to_string() +} + +pub(super) fn checkpoint_from_message(id: &str, message: &str) -> Result { + let value: serde_json::Value = serde_json::from_str(message.trim()) + .with_context(|| format!("checkpoint '{id}' has invalid JSON metadata"))?; + let label = value + .get("label") + .and_then(|v| v.as_str()) + .with_context(|| format!("checkpoint '{id}' is missing string label"))?; + let created_at_ms = value + .get("created_at_ms") + .and_then(|v| v.as_i64()) + .with_context(|| format!("checkpoint '{id}' is missing integer created_at_ms"))?; + let snapshot_values = value + .get("snapshot_ids") + .and_then(|v| v.as_array()) + .with_context(|| format!("checkpoint '{id}' is missing snapshot_ids array"))?; + let snapshot_ids = snapshot_values + .iter() + .map(|value| { + value + .as_str() + .map(str::to_string) + .with_context(|| format!("checkpoint '{id}' contains a non-string snapshot id")) + }) + .collect::>>()?; + + Ok(Checkpoint { + id: id.to_string(), + label: label.to_string(), + created_at_ms, + snapshot_ids, + }) +} + +/// A git blob oid as a content hash, or `None` for the zero oid (absent side). +pub(super) fn oid_hash(oid: Oid) -> Option { + if oid.is_zero() { + None + } else { + Some(oid.to_string()) + } +} + +/// Render a single delta's unified patch, truncated to [`MAX_TEXT_DIFF_CHARS`]. +pub(super) fn patch_text(diff: &git2::Diff, delta_idx: usize) -> Option { + let mut patch = git2::Patch::from_diff(diff, delta_idx).ok().flatten()?; + let buf = patch.to_buf().ok()?; + // git2 0.21: Buf::as_str() returns Result<&str, Utf8Error>. + let text = buf.as_str().ok()?; + if text.trim().is_empty() { + None + } else { + Some(truncate(text, MAX_TEXT_DIFF_CHARS)) + } +} + +/// Encode an item id into a single git-safe path component. Bytes outside +/// `[A-Za-z0-9._-]` become `%XX`; an `i_` prefix keeps the result clear of the +/// reserved names `.`/`..`/empty. Reversible via [`decode_item_id`]. +pub(crate) fn encode_item_id(item_id: &str) -> String { + let mut out = String::with_capacity(item_id.len() + 2); + out.push_str("i_"); + for &b in item_id.as_bytes() { + if b.is_ascii_alphanumeric() || matches!(b, b'.' | b'_' | b'-') { + out.push(b as char); + } else { + out.push('%'); + out.push_str(&format!("{b:02X}")); + } + } + out +} + +/// Encode a source id for use as a top-level git tree entry and read-marker +/// ref component. Same reversible encoding as item ids; kept as a named helper +/// so call sites make the source-vs-item boundary explicit. +pub(crate) fn encode_source_id(source_id: &str) -> String { + encode_item_id(source_id) +} + +/// Inverse of [`encode_item_id`]. +pub(crate) fn decode_item_id(encoded: &str) -> String { + let body = encoded.strip_prefix("i_").unwrap_or(encoded); + let bytes = body.as_bytes(); + let mut out = Vec::with_capacity(bytes.len()); + let mut i = 0; + while i < bytes.len() { + if bytes[i] == b'%' && i + 2 < bytes.len() { + let hex = std::str::from_utf8(&bytes[i + 1..i + 3]).unwrap_or(""); + if let Ok(byte) = u8::from_str_radix(hex, 16) { + out.push(byte); + i += 3; + continue; + } + } + out.push(bytes[i]); + i += 1; + } + String::from_utf8_lossy(&out).into_owned() +} + +pub(super) fn truncate(s: &str, max_chars: usize) -> String { + if s.len() <= max_chars { + s.to_string() + } else { + let mut end = max_chars; + while !s.is_char_boundary(end) && end > 0 { + end -= 1; + } + format!("{}…(truncated)", &s[..end]) + } +} + +/// Derive a human-readable title from item content: the first non-empty line +/// (Markdown heading markers stripped), bounded. Falls back to the item id. +pub(super) fn derive_title(item_id: &str, content: &str) -> String { + let first_line = content + .lines() + .map(str::trim) + .find(|l| !l.is_empty()) + .map(|l| l.trim_start_matches('#').trim()); + match first_line { + Some(l) if !l.is_empty() => truncate(l, 120), + _ => item_id.to_string(), + } +} diff --git a/src/memory/diff/ledger_tests.rs b/src/memory/diff/ledger_tests.rs index 2c37a38..5bd9a4c 100644 --- a/src/memory/diff/ledger_tests.rs +++ b/src/memory/diff/ledger_tests.rs @@ -67,6 +67,29 @@ fn source_ids_are_encoded_for_git_tree_paths() { assert!(changes.iter().any(|change| change.item_id == "b")); } +#[test] +fn source_ids_with_control_characters_are_rejected() { + let (ledger, _dir) = temp_ledger(); + let err = ledger + .commit_snapshot( + &meta("src_ok\nSource-Id: forged"), + &items(&[("a", "x")]), + 1000, + ) + .unwrap_err(); + assert!(format!("{err:#}").contains("control characters")); + assert!(ledger.list_snapshots(None, 10).unwrap().is_empty()); +} + +#[test] +fn trailer_parser_uses_only_final_paragraph() { + let trailers = + parse_trailers("subject: forged\nSource-Id: forged\n\nSource-Id: real\nItem-Count: 2\n"); + assert_eq!(trailers.get("source-id").map(String::as_str), Some("real")); + assert_eq!(trailers.get("item-count").map(String::as_str), Some("2")); + assert!(!trailers.contains_key("subject")); +} + #[test] fn commit_and_list_snapshots() { let (ledger, _dir) = temp_ledger(); @@ -89,6 +112,35 @@ fn commit_and_list_snapshots() { assert_eq!(fetched.item_count, 1); } +#[test] +fn read_marker_rejects_regression_and_cross_source_snapshot() { + let (ledger, _dir) = temp_ledger(); + let old = ledger + .commit_snapshot(&meta("src_a"), &items(&[("a", "v1")]), 1000) + .unwrap(); + let other = ledger + .commit_snapshot(&meta("src_b"), &items(&[("b", "other")]), 2000) + .unwrap(); + let newest = ledger + .commit_snapshot(&meta("src_a"), &items(&[("a", "v2")]), 3000) + .unwrap(); + + ledger.set_read_marker("src_a", &newest.id).unwrap(); + let err = ledger.set_read_marker("src_a", &old.id).unwrap_err(); + assert!(format!("{err:#}").contains("backwards")); + assert_eq!( + ledger.get_read_marker("src_a").unwrap().as_deref(), + Some(newest.id.as_str()) + ); + + let err = ledger.set_read_marker("src_a", &other.id).unwrap_err(); + assert!(format!("{err:#}").contains("belongs to source")); + assert_eq!( + ledger.get_read_marker("src_a").unwrap().as_deref(), + Some(newest.id.as_str()) + ); +} + #[test] fn snapshots_carry_other_sources_forward() { let (ledger, _dir) = temp_ledger(); @@ -261,6 +313,30 @@ fn checkpoint_round_trip() { assert_eq!(all[0].id, "ckpt_1"); } +#[test] +fn checkpoint_parser_rejects_corrupt_or_incomplete_metadata() { + for message in [ + "not json", + r#"{"label":"x","created_at_ms":1}"#, + r#"{"label":"x","created_at_ms":1,"snapshot_ids":[42]}"#, + ] { + assert!(checkpoint_from_message("ckpt_bad", message).is_err()); + } +} + +#[test] +fn opening_an_existing_corrupt_repository_surfaces_the_error() { + let dir = tempfile::TempDir::new().unwrap(); + let repo_dir = dir.path().join("memory_diff/repo"); + std::fs::create_dir_all(&repo_dir).unwrap(); + std::fs::write(repo_dir.join(".git"), b"not a git directory").unwrap(); + + let err = Ledger::open(dir.path()) + .err() + .expect("corruption must fail"); + assert!(format!("{err:#}").contains("open memory_diff repo")); +} + #[test] fn cleanup_checkpoints_removes_old_tags() { let (ledger, _dir) = temp_ledger(); diff --git a/src/memory/diff/mod.rs b/src/memory/diff/mod.rs index 718d959..04c4a37 100644 --- a/src/memory/diff/mod.rs +++ b/src/memory/diff/mod.rs @@ -6,7 +6,7 @@ //! over time. //! //! Snapshots are built from **already-ingested** data (supplied by an injected -//! [`SnapshotItemSource`], not by re-calling source readers), so diffs cost no +//! `SnapshotItemSource`, not by re-calling source readers), so diffs cost no //! upstream API calls. The authoritative store remains whatever backs the item //! source; this module's storage is a *derived* git ledger. //! @@ -20,31 +20,34 @@ //! - Read marker → ref `refs/openhuman/read/`. //! - Diff → git tree diff scoped to a source path. //! -//! See [`ledger`] for the mapping and [`DiffEngine`] for the operations. +//! See the ledger module for the mapping and `DiffEngine` for the operations. //! //! ## Decoupling from `chunks` //! //! The chunk store is ported separately, so the engine takes a -//! [`source::SnapshotItemSource`] by injection rather than hard-depending on -//! `chunks`. [`source::InMemoryItemSource`] is a reference/test backend. +//! `SnapshotItemSource` by injection rather than hard-depending on +//! `chunks`. `InMemoryItemSource` is a reference/test backend. //! //! ## Operations //! -//! - [`DiffEngine::take_snapshot`] / [`DiffEngine::auto_snapshot_after_sync`] -//! - [`DiffEngine::list_snapshots`] -//! - [`DiffEngine::compute_diff`] (explicit pair) -//! - [`DiffEngine::diff_since_last`] (latest vs previous) -//! - [`DiffEngine::diff_since_read`] (latest vs read marker, optional advance) -//! - [`DiffEngine::mark_read`] -//! - [`DiffEngine::create_checkpoint`] / [`DiffEngine::list_checkpoints`] -//! - [`DiffEngine::diff_since_checkpoint`] (cross-source) -//! - [`DiffEngine::cleanup`] +//! - `DiffEngine::take_snapshot` / `DiffEngine::auto_snapshot_after_sync` +//! - `DiffEngine::list_snapshots` +//! - `DiffEngine::compute_diff` (explicit pair) +//! - `DiffEngine::diff_since_last` (latest vs previous) +//! - `DiffEngine::diff_since_read` (latest vs read marker, optional advance) +//! - `DiffEngine::mark_read` +//! - `DiffEngine::create_checkpoint` / `DiffEngine::list_checkpoints` +//! - `DiffEngine::diff_since_checkpoint` (cross-source) +//! - `DiffEngine::cleanup` use std::path::PathBuf; pub mod checkpoint; +// Keep the established `memory::diff::diff` path for downstream callers. +#[allow(clippy::module_inception)] pub mod diff; pub mod ledger; +mod ledger_helpers; pub mod snapshot; pub mod source; pub mod types; diff --git a/src/memory/diff/snapshot.rs b/src/memory/diff/snapshot.rs index 2d4ce3d..f851223 100644 --- a/src/memory/diff/snapshot.rs +++ b/src/memory/diff/snapshot.rs @@ -1,7 +1,7 @@ //! Snapshot capture for the diff engine. //! //! Snapshots are built from already-ingested data supplied by the injected -//! [`SnapshotItemSource`](super::source::SnapshotItemSource) — never by +//! [`SnapshotItemSource`] — never by //! re-calling source readers — and committed to the git ledger. The chunk //! store (whatever backs the item source) remains authoritative; the ledger is //! a derived, rebuildable view used purely for change tracking. diff --git a/src/memory/entities/frontmatter.rs b/src/memory/entities/frontmatter.rs index 810979a..d9b26cd 100644 --- a/src/memory/entities/frontmatter.rs +++ b/src/memory/entities/frontmatter.rs @@ -26,7 +26,7 @@ //! Only the well-known scalar/list shapes the [`Entity`] uses are emitted and //! parsed; this is deliberately not a general YAML implementation. The free //! text after the closing `---` is the notes body and is never interpreted — -//! [`notes_body`] hands it back verbatim so upserts can round-trip it. +//! The notes-body helper hands it back verbatim so upserts can round-trip it. use chrono::{DateTime, Utc}; diff --git a/src/memory/fsutil.rs b/src/memory/fsutil.rs index 0a5f971..c2f4939 100644 --- a/src/memory/fsutil.rs +++ b/src/memory/fsutil.rs @@ -1,6 +1,6 @@ //! Shared filesystem primitives for the memory engine. //! -//! [`atomic_write`] is the single crash-safe write path used by the on-disk +//! `atomic_write` is the single crash-safe write path used by the on-disk //! stores (goals list, time-tree nodes, staged summaries). It writes to a //! hidden same-directory temp file, fsyncs it, then renames it over the //! destination. Because POSIX `rename(2)` is atomic, a crash or a concurrent diff --git a/src/memory/goals/mod.rs b/src/memory/goals/mod.rs index 62d5d2e..fed1869 100644 --- a/src/memory/goals/mod.rs +++ b/src/memory/goals/mod.rs @@ -15,7 +15,7 @@ //! //! - **Explicitly** — via the [`store`] operations (`list` / `add` / `edit` / //! `delete`), which back both RPC handlers and agent tools in OpenHuman. -//! - **By reflection** — a turn-based pass ([`reflect`]) that reads the current +//! - **By reflection** — a turn-based pass ([`reflect()`]) that reads the current //! list plus a context nudge and applies add/edit/delete. On an empty list //! (first run) it populates an initial set; otherwise it makes minimal //! changes. The LLM step is abstracted behind [`reflect::GoalsGenerator`]; diff --git a/src/memory/goals/reflect.rs b/src/memory/goals/reflect.rs index c2946b9..a50aa1a 100644 --- a/src/memory/goals/reflect.rs +++ b/src/memory/goals/reflect.rs @@ -122,11 +122,6 @@ fn normalise(text: &str) -> String { /// Apply `mutations` to `doc` deterministically, de-duplicating additions by /// normalised text. Returns `(applied, skipped)` counts. Does not persist. /// -/// NOTE: the normalised-duplicate check only guards `Add`. An `Edit` that -/// rewrites one goal's text to match another existing goal's (normalised) -/// text is still applied as-is, so a generator can converge the list to -/// several items with identical text. Callers that need strict uniqueness -/// must check for this themselves. fn apply_mutations(doc: &mut GoalsDoc, mutations: &[GoalMutation]) -> (usize, usize) { let mut applied = 0usize; let mut skipped = 0usize; @@ -145,10 +140,22 @@ fn apply_mutations(doc: &mut GoalsDoc, mutations: &[GoalMutation]) -> (usize, us Err(_) => skipped += 1, } } - GoalMutation::Edit { id, text } => match doc.edit(id, text) { - Ok(()) => applied += 1, - Err(_) => skipped += 1, - }, + GoalMutation::Edit { id, text } => { + let norm = normalise(text); + let duplicate = !norm.is_empty() + && doc + .items + .iter() + .any(|item| item.id != *id && normalise(&item.text) == norm); + if duplicate { + skipped += 1; + continue; + } + match doc.edit(id, text) { + Ok(()) => applied += 1, + Err(_) => skipped += 1, + } + } GoalMutation::Delete { id } => match doc.delete(id) { Ok(()) => applied += 1, Err(_) => skipped += 1, diff --git a/src/memory/goals/reflect_tests.rs b/src/memory/goals/reflect_tests.rs index c781fcd..135fa93 100644 --- a/src/memory/goals/reflect_tests.rs +++ b/src/memory/goals/reflect_tests.rs @@ -130,3 +130,21 @@ fn reflect_skips_secret_or_pii_goal_proposals() { assert_eq!(outcome.skipped, 2); assert_eq!(outcome.goals.items[0].text, "ship the memory engine"); } + +#[test] +fn edit_to_normalized_duplicate_is_skipped() { + let mut doc = GoalsDoc::default(); + let first = doc.add("Ship the engine").unwrap(); + let second = doc.add("Write documentation").unwrap(); + let (applied, skipped) = apply_mutations( + &mut doc, + &[GoalMutation::Edit { + id: second, + text: " SHIP the ENGINE ".into(), + }], + ); + assert_eq!((applied, skipped), (0, 1)); + assert_eq!(doc.items.len(), 2); + assert_eq!(doc.items[0].id, first); + assert_eq!(doc.items[1].text, "Write documentation"); +} diff --git a/src/memory/goals/types.rs b/src/memory/goals/types.rs index 9c9c380..4a60a02 100644 --- a/src/memory/goals/types.rs +++ b/src/memory/goals/types.rs @@ -9,7 +9,7 @@ //! This module is pure: it owns parse/render and the in-memory mutation logic //! (`add` / `edit` / `delete`) plus their validation rules. The cap-enforcing //! persistence layer lives in [`super::store`]; the reflection apply/dedupe -//! logic lives in [`super::reflect`]. +//! logic lives in [`super::reflect()`]. use serde::{Deserialize, Serialize}; diff --git a/src/memory/sources/mod.rs b/src/memory/sources/mod.rs index 95450c9..f013a1d 100644 --- a/src/memory/sources/mod.rs +++ b/src/memory/sources/mod.rs @@ -29,8 +29,8 @@ pub mod types; pub mod validation; pub use readers::{is_locally_readable, reader_for, SourceReader}; -pub use registry::{ - memory_sync_defaults_for_toolkit, ComposioUpsertTarget, MemorySourcePatch, SourceRegistry, +pub use registry::{memory_sync_defaults_for_toolkit, ComposioUpsertTarget, SourceRegistry}; +pub use types::{ + ContentType, MemorySourceEntry, MemorySourcePatch, SourceContent, SourceItem, SourceKind, }; -pub use types::{ContentType, MemorySourceEntry, SourceContent, SourceItem, SourceKind}; pub use validation::{ensure_within_base, validate_entry}; diff --git a/src/memory/sources/readers/conversation.rs b/src/memory/sources/readers/conversation.rs index 7a2c5b5..949ac2e 100644 --- a/src/memory/sources/readers/conversation.rs +++ b/src/memory/sources/readers/conversation.rs @@ -74,12 +74,6 @@ impl SourceReader for ConversationReader { /// Read one thread's content by its `item_id` (the thread's file stem, as /// produced by [`list_items`](Self::list_items)). /// - /// NOTE: the traversal guard below rejects any `item_id` containing the - /// substring `".."`, not just path-separator-adjacent `..` segments. A - /// thread file whose stem legitimately contains `".."` (e.g. - /// `standup..2026.json`, which [`list_items`](Self::list_items) itself - /// would happily produce as an id) is therefore unreadable even though it - /// poses no traversal risk once separators are excluded. async fn read_item( &self, _source: &MemorySourceEntry, @@ -87,7 +81,7 @@ impl SourceReader for ConversationReader { config: &MemoryConfig, ) -> MemoryEngineResult { // Validate item_id to prevent path traversal before touching the FS. - if item_id.contains("..") || item_id.contains('/') || item_id.contains('\\') { + if matches!(item_id, "." | "..") || item_id.contains('/') || item_id.contains('\\') { return Err(MemoryError::Invalid( "invalid item_id: path traversal denied".to_string(), )); diff --git a/src/memory/sources/readers/conversation_tests.rs b/src/memory/sources/readers/conversation_tests.rs index 57dcf24..e7427d7 100644 --- a/src/memory/sources/readers/conversation_tests.rs +++ b/src/memory/sources/readers/conversation_tests.rs @@ -154,6 +154,28 @@ async fn read_item_returns_formatted_content() { assert!(content.body.contains("**assistant**: 4")); } +#[tokio::test] +async fn read_item_accepts_legitimate_double_dot_in_stem() { + let tmp = tempdir().unwrap(); + let threads_dir = tmp.path().join("threads"); + fs::create_dir_all(&threads_dir).unwrap(); + fs::write( + threads_dir.join("standup..2026.json"), + r#"{"title":"Standup","messages":[]}"#, + ) + .unwrap(); + let reader = ConversationReader; + let content = reader + .read_item( + &conversation_source(), + "standup..2026", + &MemoryConfig::new(tmp.path()), + ) + .await + .unwrap(); + assert_eq!(content.id, "standup..2026"); +} + #[tokio::test] async fn read_item_returns_error_for_missing_thread() { let tmp = tempdir().unwrap(); diff --git a/src/memory/sources/readers/folder.rs b/src/memory/sources/readers/folder.rs index c13329f..d88f238 100644 --- a/src/memory/sources/readers/folder.rs +++ b/src/memory/sources/readers/folder.rs @@ -4,9 +4,9 @@ //! `**/*.md`), and reads their content as markdown, HTML, or plaintext. //! //! Safety: file sizes are capped at -//! [`FOLDER_FILE_SIZE_CAP_BYTES`](crate::memory::config::FOLDER_FILE_SIZE_CAP_BYTES) +//! [`FOLDER_FILE_SIZE_CAP_BYTES`] //! (10 MB) on both list and read, and `read_item` is guarded against path -//! traversal via [`ensure_within_base`](crate::memory::sources::validation::ensure_within_base). +//! traversal via [`ensure_within_base`]. //! //! The directory walk uses `walkdir`; glob patterns are compiled to a `regex` //! (matched against the slash-normalised path relative to the folder root). @@ -103,13 +103,6 @@ impl SourceReader for FolderReader { /// relative to the source's `path`, as produced by /// [`list_items`](Self::list_items)). /// - /// NOTE: only path traversal is guarded here (via [`ensure_within_base`]); - /// the source's configured glob is *not* re-applied. Any `item_id` that - /// resolves to a real file under `path` is readable, even if it would not - /// have been matched — and thus would not have been listed — by - /// [`list_items`](Self::list_items)'s glob filter. A source scoped to - /// `docs/**/*.md` can still be made to read `docs/.env` by passing that - /// item id directly. async fn read_item( &self, source: &MemorySourceEntry, @@ -121,6 +114,15 @@ impl SourceReader for FolderReader { .as_deref() .ok_or_else(|| MemoryError::Invalid("folder source requires a path".to_string()))?; + let pattern = source.glob.as_deref().unwrap_or("**/*.md"); + let matcher = glob_to_regex(pattern)?; + let normalized_id = normalize_rel(Path::new(item_id)); + if !matcher.is_match(&normalized_id) { + return Err(MemoryError::Invalid(format!( + "item '{item_id}' is outside source glob '{pattern}'" + ))); + } + let file_path = Path::new(base_path).join(item_id); if !file_path.exists() { return Err(MemoryError::NotFound(format!( diff --git a/src/memory/sources/readers/folder_tests.rs b/src/memory/sources/readers/folder_tests.rs index ea403d7..831de23 100644 --- a/src/memory/sources/readers/folder_tests.rs +++ b/src/memory/sources/readers/folder_tests.rs @@ -97,6 +97,27 @@ async fn read_item_returns_file_content() { assert_eq!(content.content_type, ContentType::Markdown); } +#[tokio::test] +async fn read_item_enforces_configured_glob() { + let tmp = TempDir::new().unwrap(); + fs::create_dir_all(tmp.path().join("docs")).unwrap(); + fs::write(tmp.path().join("docs/allowed.md"), "allowed").unwrap(); + fs::write(tmp.path().join("docs/secret.env"), "secret").unwrap(); + let mut source = folder_source(&tmp.path().to_string_lossy()); + source.glob = Some("docs/**/*.md".into()); + let reader = FolderReader; + + assert!(reader + .read_item(&source, "docs/allowed.md", &config()) + .await + .is_ok()); + let err = reader + .read_item(&source, "docs/secret.env", &config()) + .await + .unwrap_err(); + assert!(err.to_string().contains("outside source glob")); +} + #[tokio::test] async fn read_item_prevents_path_traversal() { let tmp = TempDir::new().unwrap(); diff --git a/src/memory/sources/registry.rs b/src/memory/sources/registry.rs index d564b43..8705932 100644 --- a/src/memory/sources/registry.rs +++ b/src/memory/sources/registry.rs @@ -12,18 +12,24 @@ //! Each on-disk write (`SourceRegistry::atomic_write`) is atomic (temp file + //! rename), so a crash mid-write cannot leave a truncated `config.toml`. //! -//! NOTE: the load-modify-save cycle itself is *not* locked — there is no mutex -//! or file lock guarding the read-mutate-write window (contrast the diff -//! ledger's `WRITE_LOCK` or goals' mutation lock). Two concurrent mutations on -//! the same registry race: the second writer's `list()` snapshot does not see -//! the first writer's change, so its `write_all` silently overwrites it — -//! last writer wins and one mutation is lost. +//! The complete load-modify-save cycle is guarded by a process-wide mutation +//! lock, so separate [`SourceRegistry`] handles cannot overwrite one another's +//! in-process updates. Atomic rename protects each individual disk write. use std::path::{Path, PathBuf}; +use std::sync::{LazyLock, Mutex}; use anyhow::{anyhow, bail, Context, Result}; -use super::types::{MemorySourceEntry, SourceKind}; +use super::types::{MemorySourceEntry, MemorySourcePatch, SourceKind}; + +/// Serializes each registry load-modify-save transaction in this process. +/// +/// A single lock deliberately covers every path: registry mutation is rare, +/// and correctness is more important than allowing unrelated config files to +/// race through their atomic renames. The on-disk rename remains the crash- +/// safety boundary; this mutex closes the in-process lost-update window. +static REGISTRY_MUTATION_LOCK: LazyLock> = LazyLock::new(|| Mutex::new(())); /// Conservative default sync caps for a Composio toolkit, keyed by toolkit slug. /// @@ -112,12 +118,9 @@ impl SourceRegistry { /// leaving a truncated `config.toml`, matching the OpenHuman source /// registry contract. /// - /// NOTE: this re-reads the file via [`SourceRegistry::read_table`] to - /// obtain the "other top-level keys" to preserve. That read is a second, - /// unsynchronized file access separate from whatever `list()` call the - /// caller used to build `entries` — under concurrent writers the two reads - /// can observe different file versions, so `entries` and the preserved - /// keys in the final file may not have come from the same snapshot. + /// Mutation callers hold [`REGISTRY_MUTATION_LOCK`] across their initial + /// read and this preserving re-read, keeping the two snapshots ordered with + /// respect to every other in-process writer. fn write_all(&self, entries: &[MemorySourceEntry]) -> Result<()> { let mut table = self.read_table()?; let value = toml::Value::try_from(entries).context("failed to encode memory_sources")?; @@ -177,6 +180,9 @@ impl SourceRegistry { /// Validate and add a new source. Fails if the id already exists. pub fn add(&self, entry: MemorySourceEntry) -> Result { + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); entry.validate().map_err(|e| anyhow!(e))?; let mut sources = self.list()?; if sources.iter().any(|s| s.id == entry.id) { @@ -190,12 +196,16 @@ impl SourceRegistry { /// Apply a [`MemorySourcePatch`] to an existing source, then re-validate and /// save. Fails if no source has the given id. pub fn update(&self, id: &str, patch: MemorySourcePatch) -> Result { + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); let mut sources = self.list()?; let entry = sources .iter_mut() .find(|s| s.id == id) .ok_or_else(|| anyhow!("source '{id}' not found"))?; + patch.validate_for_kind(entry.kind.clone())?; patch.apply_to(entry); entry.validate().map_err(|e| anyhow!(e))?; let updated = entry.clone(); @@ -205,6 +215,9 @@ impl SourceRegistry { /// Remove a source by id. Returns `true` if an entry was removed. pub fn remove(&self, id: &str) -> Result { + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); let mut sources = self.list()?; let before = sources.len(); sources.retain(|s| s.id != id); @@ -219,6 +232,9 @@ impl SourceRegistry { /// removed. Mirrors [`SourceRegistry::upsert_composio_source`], which keys /// composio sources on `connection_id` rather than the `src_*` id. pub fn remove_composio_source_by_connection_id(&self, connection_id: &str) -> Result { + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); let mut sources = self.list()?; let before = sources.len(); sources.retain(|s| { @@ -242,6 +258,9 @@ impl SourceRegistry { connection_id: &str, label: &str, ) -> Result { + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); let mut sources = self.list()?; let (entry, _was_insert) = upsert_composio_entry_in_place(&mut sources, toolkit, connection_id, label); @@ -254,6 +273,9 @@ impl SourceRegistry { if targets.is_empty() { return Ok(0); } + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); let mut sources = self.list()?; for (toolkit, connection_id, label) in targets { upsert_composio_entry_in_place(&mut sources, toolkit, connection_id, label); @@ -264,6 +286,9 @@ impl SourceRegistry { /// Enable every source and clear all per-source caps ("All In" mode). pub fn apply_all_in(&self) -> Result> { + let _guard = REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned"); let mut sources = self.list()?; for source in &mut sources { source.enabled = true; @@ -328,131 +353,6 @@ pub(crate) fn upsert_composio_entry_in_place( (entry, true) } -/// Partial update payload for a source entry. Absent fields are left unchanged. -#[derive(Debug, Default, serde::Deserialize)] -pub struct MemorySourcePatch { - /// New human-readable label for the source. - #[serde(default)] - pub label: Option, - /// Toggle whether the source participates in sync. - #[serde(default)] - pub enabled: Option, - /// Composio toolkit slug (e.g. `gmail`, `slack`). - #[serde(default)] - pub toolkit: Option, - /// Composio connection id this source binds to. - #[serde(default)] - pub connection_id: Option, - /// Filesystem root for a local-files source. - #[serde(default)] - pub path: Option, - /// Glob filter applied under [`MemorySourcePatch::path`]. - #[serde(default)] - pub glob: Option, - /// Remote URL for a git/web source. - #[serde(default)] - pub url: Option, - /// Git branch to track. - #[serde(default)] - pub branch: Option, - /// Explicit path allowlist within a repo source. - #[serde(default)] - pub paths: Option>, - /// Search/filter query string for query-driven sources. - #[serde(default)] - pub query: Option, - /// Lookback window in days for items to ingest. - #[serde(default)] - pub since_days: Option, - /// Cap on the number of items pulled per sync. - #[serde(default)] - pub max_items: Option, - /// Source-specific selector (e.g. a Notion database or Slack channel). - #[serde(default)] - pub selector: Option, - /// Token budget per sync run. - #[serde(default)] - pub max_tokens_per_sync: Option, - /// Cost budget per sync run, in USD. - #[serde(default)] - pub max_cost_per_sync_usd: Option, - /// History depth in days for tree/summary backfill. - #[serde(default)] - pub sync_depth_days: Option, - /// Cap on commits ingested from a git source. - #[serde(default)] - pub max_commits: Option, - /// Cap on issues ingested from a repo source. - #[serde(default)] - pub max_issues: Option, - /// Cap on pull requests ingested from a repo source. - #[serde(default)] - pub max_prs: Option, -} - -impl MemorySourcePatch { - /// Apply each present field of this patch onto `entry` in place. - fn apply_to(self, entry: &mut MemorySourceEntry) { - if let Some(label) = self.label { - entry.label = label; - } - if let Some(enabled) = self.enabled { - entry.enabled = enabled; - } - if let Some(toolkit) = self.toolkit { - entry.toolkit = Some(toolkit); - } - if let Some(connection_id) = self.connection_id { - entry.connection_id = Some(connection_id); - } - if let Some(path) = self.path { - entry.path = Some(path); - } - if let Some(glob) = self.glob { - entry.glob = Some(glob); - } - if let Some(url) = self.url { - entry.url = Some(url); - } - if let Some(branch) = self.branch { - entry.branch = Some(branch); - } - if let Some(paths) = self.paths { - entry.paths = paths; - } - if let Some(query) = self.query { - entry.query = Some(query); - } - if let Some(since_days) = self.since_days { - entry.since_days = Some(since_days); - } - if let Some(max_items) = self.max_items { - entry.max_items = Some(max_items); - } - if let Some(selector) = self.selector { - entry.selector = Some(selector); - } - if let Some(v) = self.max_tokens_per_sync { - entry.max_tokens_per_sync = Some(v); - } - if let Some(v) = self.max_cost_per_sync_usd { - entry.max_cost_per_sync_usd = Some(v); - } - if let Some(v) = self.sync_depth_days { - entry.sync_depth_days = Some(v); - } - if let Some(v) = self.max_commits { - entry.max_commits = Some(v); - } - if let Some(v) = self.max_issues { - entry.max_issues = Some(v); - } - if let Some(v) = self.max_prs { - entry.max_prs = Some(v); - } - } -} - #[cfg(test)] #[path = "registry_tests.rs"] mod tests; diff --git a/src/memory/sources/registry_tests.rs b/src/memory/sources/registry_tests.rs index 6e4a015..61e4958 100644 --- a/src/memory/sources/registry_tests.rs +++ b/src/memory/sources/registry_tests.rs @@ -75,6 +75,38 @@ fn add_rejects_invalid_entry() { assert!(reg.list().unwrap().is_empty()); } +#[test] +fn concurrent_registry_adds_preserve_every_source() { + let (_tmp, reg) = registry(); + let writers = 24; + let barrier = std::sync::Arc::new(std::sync::Barrier::new(writers)); + let mut threads = Vec::new(); + for index in 0..writers { + let reg = reg.clone(); + let barrier = barrier.clone(); + threads.push(std::thread::spawn(move || { + barrier.wait(); + reg.add(folder_entry(&format!("src_concurrent_{index}"))) + .unwrap(); + })); + } + for thread in threads { + thread.join().unwrap(); + } + + let mut ids: Vec<_> = reg + .list() + .unwrap() + .into_iter() + .map(|entry| entry.id) + .collect(); + ids.sort(); + assert_eq!(ids.len(), writers); + for index in 0..writers { + assert!(ids.contains(&format!("src_concurrent_{index}"))); + } +} + #[test] fn update_applies_patch_and_persists() { let (_tmp, reg) = registry(); @@ -243,8 +275,37 @@ fn memory_source_patch_deserializes_partial_and_github_fields() { let patch: MemorySourcePatch = serde_json::from_value(json).unwrap(); assert_eq!(patch.label.as_deref(), Some("New label")); assert_eq!(patch.enabled, Some(false)); - assert_eq!(patch.max_commits, Some(100)); - assert_eq!(patch.max_issues, Some(50)); - assert_eq!(patch.max_prs, Some(25)); + assert_eq!(patch.max_commits, Some(Some(100))); + assert_eq!(patch.max_issues, Some(Some(50))); + assert_eq!(patch.max_prs, Some(Some(25))); assert!(patch.toolkit.is_none()); } + +#[test] +fn memory_source_patch_can_clear_optional_fields_with_null() { + let (_tmp, reg) = registry(); + let mut entry = folder_entry("src_clear"); + entry.glob = Some("**/*.md".into()); + entry.max_items = Some(10); + reg.add(entry).unwrap(); + + let patch: MemorySourcePatch = serde_json::from_value(serde_json::json!({ + "glob": null, + "max_items": null + })) + .unwrap(); + let updated = reg.update("src_clear", patch).unwrap(); + assert!(updated.glob.is_none()); + assert!(updated.max_items.is_none()); +} + +#[test] +fn update_rejects_fields_that_do_not_apply_to_source_kind() { + let (_tmp, reg) = registry(); + reg.add(folder_entry("src_kind")).unwrap(); + let patch: MemorySourcePatch = serde_json::from_value(serde_json::json!({ + "url": "https://example.com/repo" + })) + .unwrap(); + assert!(reg.update("src_kind", patch).is_err()); +} diff --git a/src/memory/sources/types.rs b/src/memory/sources/types.rs index 5f3b756..6f4d138 100644 --- a/src/memory/sources/types.rs +++ b/src/memory/sources/types.rs @@ -156,6 +156,182 @@ impl MemorySourceEntry { } } +fn deserialize_double_option<'de, D, T>(deserializer: D) -> Result>, D::Error> +where + D: serde::Deserializer<'de>, + T: serde::Deserialize<'de>, +{ + as serde::Deserialize>::deserialize(deserializer).map(Some) +} + +/// Partial update payload for a source entry. +/// +/// An absent field leaves the current value unchanged. For optional source +/// properties, an explicit JSON `null` clears the value while a concrete value +/// replaces it. +#[derive(Debug, Default, Deserialize)] +pub struct MemorySourcePatch { + /// New human-readable label for the source. + #[serde(default)] + pub label: Option, + /// Toggle whether the source participates in sync. + #[serde(default)] + pub enabled: Option, + /// Composio toolkit slug (e.g. `gmail`, `slack`). + #[serde(default, deserialize_with = "deserialize_double_option")] + pub toolkit: Option>, + /// Composio connection id this source binds to. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub connection_id: Option>, + /// Filesystem root for a local-files source. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub path: Option>, + /// Glob filter applied under [`MemorySourcePatch::path`]. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub glob: Option>, + /// Remote URL for a git/web source. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub url: Option>, + /// Git branch to track. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub branch: Option>, + /// Explicit path allowlist within a repo source. + #[serde(default)] + pub paths: Option>, + /// Search/filter query string for query-driven sources. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub query: Option>, + /// Lookback window in days for items to ingest. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub since_days: Option>, + /// Cap on the number of items pulled per sync. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub max_items: Option>, + /// Source-specific selector (e.g. a CSS selector for a web page). + #[serde(default, deserialize_with = "deserialize_double_option")] + pub selector: Option>, + /// Token budget per sync run. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub max_tokens_per_sync: Option>, + /// Cost budget per sync run, in USD. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub max_cost_per_sync_usd: Option>, + /// History depth in days for tree/summary backfill. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub sync_depth_days: Option>, + /// Cap on commits ingested from a git source. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub max_commits: Option>, + /// Cap on issues ingested from a repo source. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub max_issues: Option>, + /// Cap on pull requests ingested from a repo source. + #[serde(default, deserialize_with = "deserialize_double_option")] + pub max_prs: Option>, +} + +impl MemorySourcePatch { + pub(crate) fn validate_for_kind(&self, kind: SourceKind) -> anyhow::Result<()> { + let reject = |field: &str| { + Err(anyhow::anyhow!( + "field '{field}' is not applicable to source kind '{}'", + kind.as_str() + )) + }; + if (self.toolkit.is_some() || self.connection_id.is_some()) && kind != SourceKind::Composio + { + return reject("toolkit/connection_id"); + } + if (self.path.is_some() || self.glob.is_some()) && kind != SourceKind::Folder { + return reject("path/glob"); + } + if (self.branch.is_some() + || self.paths.is_some() + || self.max_commits.is_some() + || self.max_issues.is_some() + || self.max_prs.is_some()) + && kind != SourceKind::GithubRepo + { + return reject("github repository fields"); + } + if self.query.is_some() && kind != SourceKind::TwitterQuery { + return reject("query"); + } + if self.selector.is_some() && kind != SourceKind::WebPage { + return reject("selector"); + } + if self.url.is_some() + && kind != SourceKind::GithubRepo + && kind != SourceKind::RssFeed + && kind != SourceKind::WebPage + { + return reject("url"); + } + Ok(()) + } + + /// Apply each present field of this patch onto `entry` in place. + pub(crate) fn apply_to(self, entry: &mut MemorySourceEntry) { + if let Some(value) = self.label { + entry.label = value; + } + if let Some(value) = self.enabled { + entry.enabled = value; + } + if let Some(value) = self.toolkit { + entry.toolkit = value; + } + if let Some(value) = self.connection_id { + entry.connection_id = value; + } + if let Some(value) = self.path { + entry.path = value; + } + if let Some(value) = self.glob { + entry.glob = value; + } + if let Some(value) = self.url { + entry.url = value; + } + if let Some(value) = self.branch { + entry.branch = value; + } + if let Some(value) = self.paths { + entry.paths = value; + } + if let Some(value) = self.query { + entry.query = value; + } + if let Some(value) = self.since_days { + entry.since_days = value; + } + if let Some(value) = self.max_items { + entry.max_items = value; + } + if let Some(value) = self.selector { + entry.selector = value; + } + if let Some(value) = self.max_tokens_per_sync { + entry.max_tokens_per_sync = value; + } + if let Some(value) = self.max_cost_per_sync_usd { + entry.max_cost_per_sync_usd = value; + } + if let Some(value) = self.sync_depth_days { + entry.sync_depth_days = value; + } + if let Some(value) = self.max_commits { + entry.max_commits = value; + } + if let Some(value) = self.max_issues { + entry.max_issues = value; + } + if let Some(value) = self.max_prs { + entry.max_prs = value; + } + } +} + /// One item listed from a source reader. /// /// `id` is reader-scoped (e.g. a folder-relative path or a thread id) and is diff --git a/src/memory/sources/validation.rs b/src/memory/sources/validation.rs index d4d9dc0..55f47b0 100644 --- a/src/memory/sources/validation.rs +++ b/src/memory/sources/validation.rs @@ -21,16 +21,13 @@ use super::types::{MemorySourceEntry, SourceKind}; /// Returns a human-readable error message describing the first failing rule. /// `id` and `label` are required for every kind; kind-specific fields follow. /// -/// NOTE: `id` is only checked for non-emptiness — it is not constrained to a -/// safe charset. An id containing `\n` or `:` is accepted here but can -/// corrupt downstream consumers that assume single-line, colon-free ids: -/// the diff ledger's commit-message trailers (a `\n` injects extra trailer -/// lines) and `extract_item_id`'s first-`:` split (a `:` produces spurious -/// added/removed churn instead of a stable item id). pub fn validate_entry(entry: &MemorySourceEntry) -> Result<(), String> { - if entry.id.is_empty() { + if entry.id.trim().is_empty() { return Err("id is required".to_string()); } + if entry.id.contains(':') || entry.id.chars().any(char::is_control) { + return Err("id must not contain ':' or control characters".to_string()); + } if entry.label.is_empty() { return Err("label is required".to_string()); } diff --git a/src/memory/store/content/atomic.rs b/src/memory/store/content/atomic.rs index 00237ca..d31043f 100644 --- a/src/memory/store/content/atomic.rs +++ b/src/memory/store/content/atomic.rs @@ -37,8 +37,13 @@ pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result { .map_err(|e| anyhow::anyhow!("fsync tempfile {:?}: {e}", tmp_path))?; } - match std::fs::rename(&tmp_path, abs_path) { + // Publish without replacement. A sibling hard link is atomic and fails + // with AlreadyExists when another writer won the same destination; plain + // `rename` cannot provide this contract because it replaces the target on + // Unix. + match std::fs::hard_link(&tmp_path, abs_path) { Ok(()) => { + let _ = std::fs::remove_file(&tmp_path); // fsync the parent directory so the rename is durable across a crash. #[cfg(unix)] if let Some(parent) = abs_path.parent() { @@ -55,7 +60,7 @@ pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result { Ok(false) } else { Err(anyhow::anyhow!( - "rename {:?} -> {:?}: {e}", + "publish {:?} -> {:?}: {e}", tmp_path, abs_path )) @@ -173,20 +178,9 @@ pub fn sha256_hex(bytes: &[u8]) -> String { out } -/// Tiny deterministic-ish hex string for temp file names. +/// Collision-resistant random suffix for sibling temp files. fn uuid_v4_hex() -> String { - use std::time::{SystemTime, UNIX_EPOCH}; - let t = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .subsec_nanos(); - static COUNTER: std::sync::atomic::AtomicU64 = std::sync::atomic::AtomicU64::new(0); - let n = COUNTER.fetch_add(1, std::sync::atomic::Ordering::Relaxed); - format!( - "{:08x}{:016x}", - t, - n.wrapping_mul(0x9e37_79b9_7f4a_7c15).wrapping_add(t as u64) - ) + uuid::Uuid::new_v4().simple().to_string() } #[cfg(test)] diff --git a/src/memory/store/content/compose/compose_tests.rs b/src/memory/store/content/compose/compose_tests.rs index f8974a2..eac4ffc 100644 --- a/src/memory/store/content/compose/compose_tests.rs +++ b/src/memory/store/content/compose/compose_tests.rs @@ -424,130 +424,5 @@ fn compose_global_summary_alias_format() { assert!(composed.front_matter.contains("global digest")); } -#[test] -fn compose_topic_summary_alias_format() { - let input = sample_summary_input(SummaryTreeKind::Topic, "person:alex-johnson", 1); - let composed = compose_summary_md(&input); - assert!(composed.front_matter.contains("tree_kind: topic")); - assert!(composed.front_matter.contains("topic")); -} - -#[test] -fn compose_summary_with_zero_children() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "summary:L0:empty", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 0, - child_ids: &[], - child_basenames: None, - child_count: 0, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "empty", - }; - let composed = compose_summary_md(&input); - assert!(composed.front_matter.contains("children: []")); - assert!(composed.front_matter.contains("child_count: 0")); -} - -#[test] -fn compose_summary_same_start_end_date_single_date_alias() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "summary:L1:sameday", - tree_kind: SummaryTreeKind::Global, - tree_id: "t1", - tree_scope: "global", - level: 1, - child_ids: &["child-a".to_string()], - child_basenames: None, - child_count: 1, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "day recap", - }; - let composed = compose_summary_md(&input); - let alias_line = composed - .front_matter - .lines() - .find(|l| l.contains("L1") && l.contains("global digest")) - .expect("alias line must be present"); - let date_str = ts.format("%Y-%m-%d").to_string(); - assert!(alias_line.contains(&date_str), "got: {alias_line}"); - assert!(!alias_line.contains('\u{2013}'), "got: {alias_line}"); -} - -#[test] -fn scope_short_label_two_participants() { - let label = scope_short_label("gmail:alice@x.com|bob@y.com"); - assert_eq!(label, "alice@x.com \u{2194} bob@y.com"); -} - -#[test] -fn scope_short_label_many_participants() { - let label = scope_short_label("gmail:alice@x.com|bob@y.com|carol@z.com"); - assert_eq!(label, "alice@x.com + 2 others"); -} - -#[test] -fn scope_short_label_non_gmail_returns_raw() { - let label = scope_short_label("slack:#general"); - assert_eq!(label, "slack:#general"); -} - -#[test] -fn rewrite_summary_tags_delegates_to_rewrite_tags() { - let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - let input = SummaryComposeInput { - summary_id: "sum:L1:rwttest", - tree_kind: SummaryTreeKind::Source, - tree_id: "t1", - tree_scope: "gmail:alice@x.com", - level: 1, - child_ids: &["c1".to_string()], - child_basenames: None, - child_count: 1, - time_range_start: ts, - time_range_end: ts, - sealed_at: ts, - body: "summary body text", - }; - let composed = compose_summary_md(&input); - let file_bytes = composed.full.as_bytes(); - let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()]; - let rewritten = rewrite_summary_tags(file_bytes, &new_tags).unwrap(); - let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); - assert!(rewritten_str.contains(" - person/Alice-Smith")); - assert!(rewritten_str.contains(" - topic/Memory")); - assert!(!rewritten_str.contains("tags: []")); - assert!(rewritten_str.contains(&format!( - "openhuman_core_version: {}", - OPENHUMAN_CORE_VERSION - ))); - assert!(rewritten_str.contains(&format!( - "memory_artifact_format: {}", - MEMORY_ARTIFACT_FORMAT - ))); - assert!(rewritten_str.ends_with("summary body text")); -} - -#[test] -fn rewrite_summary_tags_backfills_missing_provenance() { - let file = b"---\nid: legacy\nkind: summary\ntags: []\naliases:\n - legacy\n---\nlegacy body"; - let rewritten = rewrite_summary_tags(file, &["person/Alice".to_string()]).unwrap(); - let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); - assert!(rewritten_str.contains(&format!( - "openhuman_core_version: {}", - OPENHUMAN_CORE_VERSION - ))); - assert!(rewritten_str.contains(&format!( - "memory_artifact_format: {}", - MEMORY_ARTIFACT_FORMAT - ))); - assert!(rewritten_str.ends_with("legacy body")); -} +#[path = "compose_tests_more.rs"] +mod more; diff --git a/src/memory/store/content/compose/compose_tests_more.rs b/src/memory/store/content/compose/compose_tests_more.rs new file mode 100644 index 0000000..97ea48c --- /dev/null +++ b/src/memory/store/content/compose/compose_tests_more.rs @@ -0,0 +1,129 @@ +use super::*; + +#[test] +fn compose_topic_summary_alias_format() { + let input = sample_summary_input(SummaryTreeKind::Topic, "person:alex-johnson", 1); + let composed = compose_summary_md(&input); + assert!(composed.front_matter.contains("tree_kind: topic")); + assert!(composed.front_matter.contains("topic")); +} + +#[test] +fn compose_summary_with_zero_children() { + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let input = SummaryComposeInput { + summary_id: "summary:L0:empty", + tree_kind: SummaryTreeKind::Source, + tree_id: "t1", + tree_scope: "gmail:alice@x.com", + level: 0, + child_ids: &[], + child_basenames: None, + child_count: 0, + time_range_start: ts, + time_range_end: ts, + sealed_at: ts, + body: "empty", + }; + let composed = compose_summary_md(&input); + assert!(composed.front_matter.contains("children: []")); + assert!(composed.front_matter.contains("child_count: 0")); +} + +#[test] +fn compose_summary_same_start_end_date_single_date_alias() { + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let input = SummaryComposeInput { + summary_id: "summary:L1:sameday", + tree_kind: SummaryTreeKind::Global, + tree_id: "t1", + tree_scope: "global", + level: 1, + child_ids: &["child-a".to_string()], + child_basenames: None, + child_count: 1, + time_range_start: ts, + time_range_end: ts, + sealed_at: ts, + body: "day recap", + }; + let composed = compose_summary_md(&input); + let alias_line = composed + .front_matter + .lines() + .find(|l| l.contains("L1") && l.contains("global digest")) + .expect("alias line must be present"); + let date_str = ts.format("%Y-%m-%d").to_string(); + assert!(alias_line.contains(&date_str), "got: {alias_line}"); + assert!(!alias_line.contains('\u{2013}'), "got: {alias_line}"); +} + +#[test] +fn scope_short_label_two_participants() { + let label = scope_short_label("gmail:alice@x.com|bob@y.com"); + assert_eq!(label, "alice@x.com \u{2194} bob@y.com"); +} + +#[test] +fn scope_short_label_many_participants() { + let label = scope_short_label("gmail:alice@x.com|bob@y.com|carol@z.com"); + assert_eq!(label, "alice@x.com + 2 others"); +} + +#[test] +fn scope_short_label_non_gmail_returns_raw() { + let label = scope_short_label("slack:#general"); + assert_eq!(label, "slack:#general"); +} + +#[test] +fn rewrite_summary_tags_delegates_to_rewrite_tags() { + let ts = chrono::Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let input = SummaryComposeInput { + summary_id: "sum:L1:rwttest", + tree_kind: SummaryTreeKind::Source, + tree_id: "t1", + tree_scope: "gmail:alice@x.com", + level: 1, + child_ids: &["c1".to_string()], + child_basenames: None, + child_count: 1, + time_range_start: ts, + time_range_end: ts, + sealed_at: ts, + body: "summary body text", + }; + let composed = compose_summary_md(&input); + let file_bytes = composed.full.as_bytes(); + let new_tags = vec!["person/Alice-Smith".to_string(), "topic/Memory".to_string()]; + let rewritten = rewrite_summary_tags(file_bytes, &new_tags).unwrap(); + let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); + assert!(rewritten_str.contains(" - person/Alice-Smith")); + assert!(rewritten_str.contains(" - topic/Memory")); + assert!(!rewritten_str.contains("tags: []")); + assert!(rewritten_str.contains(&format!( + "openhuman_core_version: {}", + OPENHUMAN_CORE_VERSION + ))); + assert!(rewritten_str.contains(&format!( + "memory_artifact_format: {}", + MEMORY_ARTIFACT_FORMAT + ))); + assert!(rewritten_str.ends_with("summary body text")); +} + +#[test] +fn rewrite_summary_tags_backfills_missing_provenance() { + let file = b"---\nid: legacy\nkind: summary\ntags: []\naliases:\n - legacy\n---\nlegacy body"; + let rewritten = rewrite_summary_tags(file, &["person/Alice".to_string()]).unwrap(); + let rewritten_str = std::str::from_utf8(&rewritten).unwrap(); + assert!(rewritten_str.contains(&format!( + "openhuman_core_version: {}", + OPENHUMAN_CORE_VERSION + ))); + assert!(rewritten_str.contains(&format!( + "memory_artifact_format: {}", + MEMORY_ARTIFACT_FORMAT + ))); + assert!(rewritten_str.ends_with("legacy body")); +} diff --git a/src/memory/store/content/compose/yaml.rs b/src/memory/store/content/compose/yaml.rs index d2d5fc2..dc8159a 100644 --- a/src/memory/store/content/compose/yaml.rs +++ b/src/memory/store/content/compose/yaml.rs @@ -116,7 +116,7 @@ pub fn split_front_matter(content: &str) -> Option<(&str, &str)> { /// provider-controlled values (`source_id`, `owner`, `source_ref`, tags) must /// not be able to inject additional front-matter lines or terminate the block /// early with an embedded `\n---\n`. Escapes are decoded by -/// [`unescape_double_quoted`] via [`scan_fm_field`], so values round-trip. +/// the matching scanner unescape path, so values round-trip. pub fn yaml_scalar(s: &str) -> String { let needs_quoting = s.is_empty() || s.trim() != s diff --git a/src/memory/store/content/mod.rs b/src/memory/store/content/mod.rs index 5905894..17f3e60 100644 --- a/src/memory/store/content/mod.rs +++ b/src/memory/store/content/mod.rs @@ -7,12 +7,12 @@ //! //! ## Module layout //! -//! - [`paths`] — path generation + `slugify_source_id` + summary path builders -//! - [`compose`] — YAML front-matter + body composition; tag rewriting -//! - [`atomic`] — tempfile + fsync + rename writes; SHA-256; `stage_summary` -//! - [`read`] — read + SHA-256 verification + `split_front_matter` -//! - [`tags`] — `update_chunk_tags` + Obsidian tag slugifiers -//! - [`raw`] — verbatim per-item raw archive (`raw///…`) +//! - `paths` — path generation, slugification, and summary path builders +//! - `compose` — YAML front matter, body composition, and tag rewriting +//! - `atomic` — tempfile, fsync, rename, SHA-256, and summary staging +//! - `read` — reads, SHA-256 verification, and front-matter splitting +//! - `tags` — chunk-tag updates and Obsidian tag slugifiers +//! - `raw` — verbatim per-item raw archive (`raw///…`) //! //! ## Deferred //! diff --git a/src/memory/store/content/tags.rs b/src/memory/store/content/tags.rs index f64437c..0a171dc 100644 --- a/src/memory/store/content/tags.rs +++ b/src/memory/store/content/tags.rs @@ -156,22 +156,9 @@ fn augment_with_source_tag_for_chunk(file_bytes: &[u8], tags: &[String]) -> Vec< out } -/// Generate a temp-file name suffix from the current time's sub-second -/// nanoseconds. -/// -/// NOTE: unlike [`super::atomic::write_if_new`]'s temp-name generator (which -/// mixes in a per-process atomic counter), this is `subsec_nanos()` alone — -/// two rewrites of the same chunk file landing in the same directory close -/// enough in time can produce the same `.tmp_tags_.md` name and clobber -/// each other's staging file before the rename. This is a known gap (tracked -/// as SC-21 in the storage-primitives audit), not a documentation error. +/// Generate a collision-resistant temp-file suffix. fn crate_temp_id() -> String { - use std::time::{SystemTime, UNIX_EPOCH}; - let ns = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or_default() - .subsec_nanos(); - format!("{ns:08x}") + uuid::Uuid::new_v4().simple().to_string() } #[cfg(test)] diff --git a/src/memory/store/entity_index/mod.rs b/src/memory/store/entity_index/mod.rs index 5518b06..adddfa1 100644 --- a/src/memory/store/entity_index/mod.rs +++ b/src/memory/store/entity_index/mod.rs @@ -9,24 +9,19 @@ //! //! ## Concurrency / atomicity contract //! -//! [`EntityIndex::open`] runs `PRAGMA journal_mode = WAL` unconditionally on -//! whatever `db_path` it is given. The chunk store (`crate::memory::chunks`) -//! declares this same `mem_tree_entity_index` table in its own schema and -//! deliberately enforces `TRUNCATE` journal mode on its database file. Callers -//! must therefore either (a) point `EntityIndex` at a dedicated database file -//! distinct from the chunk store's — accepting that the two copies of the -//! table are then independent and cascade-deletes / coverage counts on one -//! side won't see the other's rows — or (b) share the chunk store's -//! connection/path deliberately and accept that the WAL pragma here will flip -//! that database out of `TRUNCATE` mode. Neither option is handled -//! automatically by this module; picking one and documenting it is the -//! caller's responsibility. +//! Use [`EntityIndex::for_memory_config`] for engine data. It wraps the chunk +//! store's shared connection, so scoring, retrieval, cascades, and this typed +//! facade all observe the same table without changing owner-managed pragmas. +//! [`EntityIndex::open`] remains available for a deliberately independent +//! standalone index. pub mod store; +mod transaction; pub mod types; -pub use store::{ +pub use store::{EntityIndex, NoSelfIdentity, SelfIdentity}; +pub use transaction::{ clear_entity_index_for_node_tx, index_entities_tx, index_entities_tx_with_identity, - index_summary_entity_ids_tx_with_identity, EntityIndex, NoSelfIdentity, SelfIdentity, + index_summary_entity_ids_tx_with_identity, }; pub use types::{CanonicalEntity, EntityHit, EntityKind}; diff --git a/src/memory/store/entity_index/store.rs b/src/memory/store/entity_index/store.rs index 8e59a3f..f6119db 100644 --- a/src/memory/store/entity_index/store.rs +++ b/src/memory/store/entity_index/store.rs @@ -61,13 +61,27 @@ const SCHEMA_SQL: &str = " const OWNED_CONNECTION_PRAGMAS: &str = " PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL; + PRAGMA busy_timeout = 15000; "; -const UPSERT_SQL: &str = "INSERT OR REPLACE INTO mem_tree_entity_index ( +pub(super) const UPSERT_SQL: &str = "INSERT OR REPLACE INTO mem_tree_entity_index ( entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms, tree_id, is_user ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)"; +pub(super) const UPSERT_PRESERVE_USER_SQL: &str = "INSERT INTO mem_tree_entity_index ( + entity_id, node_id, node_kind, entity_kind, surface, + score, timestamp_ms, tree_id, is_user + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9) + ON CONFLICT(entity_id, node_id) DO UPDATE SET + node_kind = excluded.node_kind, + entity_kind = excluded.entity_kind, + surface = excluded.surface, + score = excluded.score, + timestamp_ms = excluded.timestamp_ms, + tree_id = excluded.tree_id, + is_user = MAX(mem_tree_entity_index.is_user, excluded.is_user)"; + /// SQLite-backed entity occurrence index. /// /// Thread-safe: the connection is behind a `parking_lot::Mutex`. @@ -77,6 +91,21 @@ pub struct EntityIndex { } impl EntityIndex { + /// Wrap the engine's shared chunk-database connection with the default + /// no-op identity resolver. + pub fn for_memory_config(config: &crate::memory::config::MemoryConfig) -> anyhow::Result { + Self::for_memory_config_with_identity(config, Arc::new(NoSelfIdentity)) + } + + /// Wrap the engine's shared chunk-database connection with a host-provided + /// identity resolver. + pub fn for_memory_config_with_identity( + config: &crate::memory::config::MemoryConfig, + identity: Arc, + ) -> anyhow::Result { + Self::from_shared_connection(crate::memory::chunks::shared_connection(config)?, identity) + } + /// Open (or create) an index at `db_path` with the default no-op identity. pub fn open(db_path: &Path) -> anyhow::Result { Self::open_with_identity(db_path, Arc::new(NoSelfIdentity)) @@ -332,7 +361,8 @@ impl EntityIndex { /// Run a closure with the raw transaction, for callers that need to fold /// entity indexing into a larger atomic write. The closure receives a - /// [`Transaction`] and may call [`index_entities_tx`]. + /// [`Transaction`] and may call + /// [`crate::memory::store::entity_index::index_entities_tx`]. pub fn with_transaction( &self, f: impl FnOnce(&Transaction<'_>) -> anyhow::Result, @@ -347,7 +377,7 @@ impl EntityIndex { /// Resolve `is_user` for a summary canonical id (`":"`), returning /// `false` for malformed ids or non-matchable kinds. -fn canonical_id_is_user(identity: &dyn SelfIdentity, canonical_id: &str) -> bool { +pub(super) fn canonical_id_is_user(identity: &dyn SelfIdentity, canonical_id: &str) -> bool { let Some((kind_str, value)) = canonical_id.split_once(':') else { return false; }; @@ -357,110 +387,6 @@ fn canonical_id_is_user(identity: &dyn SelfIdentity, canonical_id: &str) -> bool identity.is_self(kind, value) } -/// Transaction-scoped batch index, for folding into a larger atomic write via -/// [`EntityIndex::with_transaction`]. `is_user` defaults to `false` here since -/// the identity resolver is not in scope on a bare transaction. -/// -/// NOTE: because the upsert is keyed on `(entity_id, node_id)` and always -/// writes `is_user = 0`, re-indexing a node through this path clobbers a row -/// that a prior [`EntityIndex::index_entity`]/[`EntityIndex::index_entities`] -/// call had correctly marked `is_user = 1` — there is no read-before-write to -/// preserve the existing flag. Callers that need `is_user` preserved across -/// re-indexing should not mix this entry point with the identity-aware ones -/// for the same node. -pub fn index_entities_tx( - tx: &Transaction<'_>, - entities: &[CanonicalEntity], - node_id: &str, - node_kind: &str, - timestamp_ms: i64, - tree_id: Option<&str>, -) -> anyhow::Result { - index_entities_tx_with_identity( - tx, - entities, - node_id, - node_kind, - timestamp_ms, - tree_id, - &NoSelfIdentity, - ) -} - -/// Identity-aware transaction-scoped batch index. -pub fn index_entities_tx_with_identity( - tx: &Transaction<'_>, - entities: &[CanonicalEntity], - node_id: &str, - node_kind: &str, - timestamp_ms: i64, - tree_id: Option<&str>, - identity: &dyn SelfIdentity, -) -> anyhow::Result { - if entities.is_empty() { - return Ok(0); - } - let mut stmt = tx.prepare(UPSERT_SQL)?; - for e in entities { - stmt.execute(params![ - e.canonical_id, - node_id, - node_kind, - e.kind.as_str(), - e.surface, - e.score, - timestamp_ms, - tree_id, - identity.is_self(e.kind, &e.surface) as i32, - ])?; - } - Ok(entities.len()) -} - -/// Remove a node's entity rows inside a caller-owned transaction. -pub fn clear_entity_index_for_node_tx( - tx: &Transaction<'_>, - node_id: &str, -) -> anyhow::Result { - Ok(tx.execute( - "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", - params![node_id], - )?) -} - -/// Index summary entity ids inside a caller-owned transaction. -pub fn index_summary_entity_ids_tx_with_identity( - tx: &Transaction<'_>, - entity_ids: &[String], - node_id: &str, - score: f32, - timestamp_ms: i64, - tree_id: Option<&str>, - identity: &dyn SelfIdentity, -) -> anyhow::Result { - if entity_ids.is_empty() { - return Ok(0); - } - let mut stmt = tx.prepare(UPSERT_SQL)?; - for canonical_id in entity_ids { - let entity_kind = canonical_id - .split_once(':') - .map_or(canonical_id.as_str(), |(kind, _)| kind); - stmt.execute(params![ - canonical_id, - node_id, - "summary", - entity_kind, - canonical_id, - score, - timestamp_ms, - tree_id, - canonical_id_is_user(identity, canonical_id) as i32, - ])?; - } - Ok(entity_ids.len()) -} - #[cfg(test)] #[path = "store_tests.rs"] mod tests; diff --git a/src/memory/store/entity_index/store_tests.rs b/src/memory/store/entity_index/store_tests.rs index a5bf7ac..42d35a2 100644 --- a/src/memory/store/entity_index/store_tests.rs +++ b/src/memory/store/entity_index/store_tests.rs @@ -1,3 +1,6 @@ +use super::super::transaction::{ + index_entities_tx, index_entities_tx_with_identity, index_summary_entity_ids_tx_with_identity, +}; use super::*; use crate::memory::store::entity_index::types::{CanonicalEntity, EntityKind}; use std::sync::Arc; @@ -45,6 +48,21 @@ fn shared_connection_preserves_owner_pragmas_and_visibility() { assert_eq!(count, 1); } +#[test] +fn memory_config_constructor_shares_rows_with_score_and_retrieval_store() { + let temp = tempfile::tempdir().unwrap(); + let config = crate::memory::config::MemoryConfig::new(temp.path()); + let index = EntityIndex::for_memory_config(&config).unwrap(); + index + .index_entity(&sample_entity("alice"), "chunk-1", "leaf", 1000, None) + .unwrap(); + + let hits = crate::memory::score::store::lookup_entity(&config, "email:alice", None).unwrap(); + + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].node_id, "chunk-1"); +} + #[derive(Debug)] struct EmailSelf; @@ -321,3 +339,22 @@ fn custom_identity_marks_self_rows() { "summary self-identity should resolve from canonical id" ); } + +#[test] +fn identity_free_transaction_reindex_preserves_existing_user_flag() { + struct OnlyAlice; + impl SelfIdentity for OnlyAlice { + fn is_self(&self, kind: EntityKind, surface: &str) -> bool { + kind == EntityKind::Email && surface == "alice@example.com" + } + } + + let idx = EntityIndex::open_in_memory_with_identity(std::sync::Arc::new(OnlyAlice)).unwrap(); + let alice = sample_entity("alice"); + idx.index_entity(&alice, "chunk-1", "leaf", 1000, None) + .unwrap(); + idx.with_transaction(|tx| index_entities_tx(tx, &[alice], "chunk-1", "leaf", 2000, None)) + .unwrap(); + + assert!(idx.lookup_entity("email:alice", None).unwrap()[0].is_user); +} diff --git a/src/memory/store/entity_index/transaction.rs b/src/memory/store/entity_index/transaction.rs new file mode 100644 index 0000000..1d092ec --- /dev/null +++ b/src/memory/store/entity_index/transaction.rs @@ -0,0 +1,120 @@ +//! Transaction-scoped entity-index writes. + +use rusqlite::{params, Transaction}; + +use super::store::{ + canonical_id_is_user, NoSelfIdentity, SelfIdentity, UPSERT_PRESERVE_USER_SQL, UPSERT_SQL, +}; +use super::types::CanonicalEntity; + +/// Index canonical entities inside the caller's transaction. +pub fn index_entities_tx( + tx: &Transaction<'_>, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, +) -> anyhow::Result { + index_entities_tx_inner( + tx, + entities, + node_id, + node_kind, + timestamp_ms, + tree_id, + &NoSelfIdentity, + UPSERT_PRESERVE_USER_SQL, + ) +} + +/// Index canonical entities with self-identity classification in one transaction. +pub fn index_entities_tx_with_identity( + tx: &Transaction<'_>, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, + identity: &dyn SelfIdentity, +) -> anyhow::Result { + index_entities_tx_inner( + tx, + entities, + node_id, + node_kind, + timestamp_ms, + tree_id, + identity, + UPSERT_SQL, + ) +} + +#[allow(clippy::too_many_arguments)] +fn index_entities_tx_inner( + tx: &Transaction<'_>, + entities: &[CanonicalEntity], + node_id: &str, + node_kind: &str, + timestamp_ms: i64, + tree_id: Option<&str>, + identity: &dyn SelfIdentity, + sql: &str, +) -> anyhow::Result { + let mut statement = tx.prepare(sql)?; + for entity in entities { + statement.execute(params![ + entity.canonical_id, + node_id, + node_kind, + entity.kind.as_str(), + entity.surface, + entity.score, + timestamp_ms, + tree_id, + identity.is_self(entity.kind, &entity.surface) as i32, + ])?; + } + Ok(entities.len()) +} + +/// Delete all entity-index rows for a node inside the caller's transaction. +pub fn clear_entity_index_for_node_tx( + tx: &Transaction<'_>, + node_id: &str, +) -> anyhow::Result { + Ok(tx.execute( + "DELETE FROM mem_tree_entity_index WHERE node_id = ?1", + params![node_id], + )?) +} + +/// Index canonical summary entity ids with self-identity classification. +pub fn index_summary_entity_ids_tx_with_identity( + tx: &Transaction<'_>, + entity_ids: &[String], + node_id: &str, + score: f32, + timestamp_ms: i64, + tree_id: Option<&str>, + identity: &dyn SelfIdentity, +) -> anyhow::Result { + let mut statement = tx.prepare(UPSERT_SQL)?; + for canonical_id in entity_ids { + let entity_kind = canonical_id + .split_once(':') + .map_or(canonical_id.as_str(), |(kind, _)| kind); + statement.execute(params![ + canonical_id, + node_id, + "summary", + entity_kind, + canonical_id, + score, + timestamp_ms, + tree_id, + canonical_id_is_user(identity, canonical_id) as i32, + ])?; + } + Ok(entity_ids.len()) +} diff --git a/src/memory/store/kv.rs b/src/memory/store/kv.rs index a71775c..7e7e2a3 100644 --- a/src/memory/store/kv.rs +++ b/src/memory/store/kv.rs @@ -38,6 +38,7 @@ const SCHEMA_SQL: &str = " const OWNED_CONNECTION_PRAGMAS: &str = " PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL; + PRAGMA busy_timeout = 15000; "; /// SQLite-backed global + namespace JSON key-value store. @@ -136,10 +137,8 @@ impl KvStore { /// Read a global key, returning `None` if absent. /// - /// NOTE: a present row whose `value_json` fails to parse (corrupt data) - /// also returns `Ok(None)` — `serde_json::from_str(..).ok()` collapses - /// "absent" and "present but corrupt" into the same result, so callers - /// cannot distinguish the two from this return value alone. + /// A present row with corrupt JSON returns an error rather than masquerading + /// as a missing key. pub fn get_global(&self, key: &str) -> Result, String> { let conn = self.conn.lock(); let value: Option = conn @@ -150,7 +149,9 @@ impl KvStore { ) .optional() .map_err(|e| format!("get_global: {e}"))?; - Ok(value.and_then(|v| serde_json::from_str(&v).ok())) + value + .map(|raw| serde_json::from_str(&raw).map_err(|e| format!("get_global JSON: {e}"))) + .transpose() } /// Insert or update a namespace-scoped key-value pair. @@ -185,8 +186,7 @@ impl KvStore { /// Read a namespace-scoped key, returning `None` if absent. /// - /// Same caveat as [`Self::get_global`]: a corrupt stored `value_json` - /// also reads back as `Ok(None)`, indistinguishable from a missing key. + /// A corrupt stored value returns an error. pub fn get_namespace(&self, namespace: &str, key: &str) -> Result, String> { let conn = self.conn.lock(); let value: Option = conn @@ -197,7 +197,9 @@ impl KvStore { ) .optional() .map_err(|e| format!("get_namespace: {e}"))?; - Ok(value.and_then(|v| serde_json::from_str(&v).ok())) + value + .map(|raw| serde_json::from_str(&raw).map_err(|e| format!("get_namespace JSON: {e}"))) + .transpose() } /// Delete a global key. Returns `true` if a row was removed. @@ -224,9 +226,7 @@ impl KvStore { /// List all keys in a namespace, most recently updated first, as a JSON /// array of `{key, value, updatedAt}` objects. /// - /// A row whose `value_json` fails to parse contributes `"value": null` - /// rather than being skipped or erroring — corrupt rows still appear in - /// the listing, just with a null payload. + /// A corrupt stored value fails the listing so corruption is observable. pub fn list_namespace(&self, namespace: &str) -> Result, String> { let conn = self.conn.lock(); let mut stmt = conn @@ -244,9 +244,12 @@ impl KvStore { .map_err(|e| format!("list_namespace row: {e}"))? { let value_raw: String = row.get(1).map_err(|e| e.to_string())?; + let key = row.get::<_, String>(0).map_err(|e| e.to_string())?; + let value = serde_json::from_str::(&value_raw) + .map_err(|e| format!("list_namespace JSON for key '{key}': {e}"))?; out.push(json!({ - "key": row.get::<_, String>(0).map_err(|e| e.to_string())?, - "value": serde_json::from_str::(&value_raw).unwrap_or(Value::Null), + "key": key, + "value": value, "updatedAt": row.get::<_, f64>(2).map_err(|e| e.to_string())?, })); } @@ -268,8 +271,7 @@ impl KvStore { /// All records in a namespace as typed [`MemoryKvRecord`]s, newest first. /// - /// Same corrupt-row handling as [`Self::list_namespace`]: an unparsable - /// `value_json` surfaces as `Value::Null` rather than an error. + /// An unparsable stored value returns an error. pub fn records_namespace(&self, namespace: &str) -> Result, String> { let ns = Self::sanitize_namespace(namespace); let conn = self.conn.lock(); @@ -288,10 +290,13 @@ impl KvStore { .map_err(|e| format!("row records_namespace: {e}"))? { let value_raw: String = row.get(1).map_err(|e| e.to_string())?; + let key: String = row.get(0).map_err(|e| e.to_string())?; + let value = serde_json::from_str(&value_raw) + .map_err(|e| format!("records_namespace JSON for key '{key}': {e}"))?; out.push(MemoryKvRecord { namespace: Some(ns.clone()), - key: row.get(0).map_err(|e| e.to_string())?, - value: serde_json::from_str(&value_raw).unwrap_or(Value::Null), + key, + value, updated_at: row.get(2).map_err(|e| e.to_string())?, }); } @@ -300,8 +305,7 @@ impl KvStore { /// All global records as typed [`MemoryKvRecord`]s, newest first. /// - /// Same corrupt-row handling as [`Self::list_namespace`]: an unparsable - /// `value_json` surfaces as `Value::Null` rather than an error. + /// An unparsable stored value returns an error. pub fn records_global(&self) -> Result, String> { let conn = self.conn.lock(); let mut stmt = conn @@ -316,10 +320,13 @@ impl KvStore { .map_err(|e| format!("row records_global: {e}"))? { let value_raw: String = row.get(1).map_err(|e| e.to_string())?; + let key: String = row.get(0).map_err(|e| e.to_string())?; + let value = serde_json::from_str(&value_raw) + .map_err(|e| format!("records_global JSON for key '{key}': {e}"))?; out.push(MemoryKvRecord { namespace: None, - key: row.get(0).map_err(|e| e.to_string())?, - value: serde_json::from_str(&value_raw).unwrap_or(Value::Null), + key, + value, updated_at: row.get(2).map_err(|e| e.to_string())?, }); } diff --git a/src/memory/store/kv_tests.rs b/src/memory/store/kv_tests.rs index 57e4ff1..46a3357 100644 --- a/src/memory/store/kv_tests.rs +++ b/src/memory/store/kv_tests.rs @@ -127,3 +127,40 @@ fn missing_keys_return_none() { assert!(!kv.delete_global("nope").unwrap()); assert!(!kv.delete_namespace("ns", "nope").unwrap()); } + +#[test] +fn corrupt_json_is_reported_by_every_read_shape() { + let kv = store(); + { + let conn = kv.conn.lock(); + conn.execute( + "INSERT INTO kv_global (key, value_json, updated_at) VALUES ('bad', '{', 0)", + [], + ) + .unwrap(); + conn.execute( + "INSERT INTO kv_namespace (namespace, key, value_json, updated_at) + VALUES ('ns', 'bad', '{', 0)", + [], + ) + .unwrap(); + } + + assert!(kv.get_global("bad").is_err()); + assert!(kv.get_namespace("ns", "bad").is_err()); + assert!(kv.list_namespace("ns").is_err()); + assert!(kv.records_namespace("ns").is_err()); + assert!(kv.records_global().is_err()); + assert!(kv.records_for_scope("ns").is_err()); +} + +#[test] +fn owned_connection_configures_busy_timeout() { + let kv = store(); + let timeout: i64 = kv + .conn + .lock() + .query_row("PRAGMA busy_timeout", [], |row| row.get(0)) + .unwrap(); + assert_eq!(timeout, 15_000); +} diff --git a/src/memory/store/memory_trait.rs b/src/memory/store/memory_trait.rs new file mode 100644 index 0000000..5aca824 --- /dev/null +++ b/src/memory/store/memory_trait.rs @@ -0,0 +1,266 @@ +//! High-level [`Memory`](crate::memory::Memory) implementation for the +//! in-process reference backend. + +use std::collections::HashMap; + +use anyhow::{anyhow, Result}; +use async_trait::async_trait; +use serde_json::Value; + +use super::store::InMemoryMemoryStore; +use super::types::{MemoryInput, MemoryRecord}; +use crate::memory::traits::Memory; +use crate::memory::types::{ + MemoryCategory, MemoryEntry, MemoryTaint, NamespaceSummary, RecallOpts, GLOBAL_NAMESPACE, +}; + +const KEY: &str = "tinycortex.key"; +const CATEGORY: &str = "tinycortex.category"; +const SESSION: &str = "tinycortex.session_id"; +const TAINT: &str = "tinycortex.taint"; + +fn text_field<'a>(record: &'a MemoryRecord, name: &str) -> Option<&'a str> { + record.metadata.get(name).and_then(Value::as_str) +} + +fn entry_from_record(record: &MemoryRecord) -> MemoryEntry { + let category = text_field(record, CATEGORY) + .and_then(|value| value.parse().ok()) + .unwrap_or(MemoryCategory::Core); + MemoryEntry { + id: record.id.to_string(), + key: text_field(record, KEY) + .map(str::to_owned) + .unwrap_or_else(|| record.id.to_string()), + content: record.content.clone(), + namespace: Some(record.namespace.clone()), + category, + timestamp: record.updated_at.to_rfc3339(), + session_id: text_field(record, SESSION).map(str::to_owned), + score: None, + taint: MemoryTaint::from_db_str(text_field(record, TAINT).unwrap_or("external_sync")), + } +} + +fn query_score(content: &str, query: &str) -> Option { + let terms = query + .split_whitespace() + .map(str::to_lowercase) + .collect::>(); + if terms.is_empty() { + return Some(1.0); + } + let content = content.to_lowercase(); + let matched = terms + .iter() + .filter(|term| content.contains(term.as_str())) + .count(); + (matched != 0).then_some(matched as f64 / terms.len() as f64) +} + +#[async_trait] +impl Memory for InMemoryMemoryStore { + fn name(&self) -> &str { + "in_memory" + } + + async fn store( + &self, + namespace: &str, + key: &str, + content: &str, + category: MemoryCategory, + session_id: Option<&str>, + ) -> Result<()> { + self.store_with_taint( + namespace, + key, + content, + category, + session_id, + MemoryTaint::Internal, + ) + .await + } + + async fn store_with_taint( + &self, + namespace: &str, + key: &str, + content: &str, + category: MemoryCategory, + session_id: Option<&str>, + taint: MemoryTaint, + ) -> Result<()> { + if key.trim().is_empty() { + return Err(anyhow!("memory key cannot be empty")); + } + let content = content.trim(); + if content.is_empty() { + return Err(anyhow!("memory content cannot be empty")); + } + let mut records = self + .records + .write() + .map_err(|_| anyhow!("memory store lock poisoned"))?; + let existing_id = records + .values() + .find(|record| record.namespace == namespace && text_field(record, KEY) == Some(key)) + .map(|record| record.id); + let mut metadata = serde_json::Map::new(); + metadata.insert(KEY.into(), Value::String(key.to_string())); + metadata.insert(CATEGORY.into(), Value::String(category.to_string())); + if let Some(session_id) = session_id { + metadata.insert(SESSION.into(), Value::String(session_id.to_string())); + } + metadata.insert(TAINT.into(), Value::String(taint.as_db_str().to_string())); + + if let Some(id) = existing_id { + let record = records.get_mut(&id).expect("id came from the same map"); + record.content = content.to_string(); + record.metadata = metadata; + record.updated_at = chrono::Utc::now(); + } else { + let record = MemoryRecord::from_input(MemoryInput { + namespace: namespace.to_string(), + content: content.to_string(), + metadata, + })?; + records.insert(record.id, record); + } + Ok(()) + } + + async fn recall( + &self, + query: &str, + limit: usize, + opts: RecallOpts<'_>, + ) -> Result> { + let namespace = opts.namespace.unwrap_or(GLOBAL_NAMESPACE); + let records = self + .records + .read() + .map_err(|_| anyhow!("memory store lock poisoned"))?; + let mut entries = records + .values() + .filter(|record| record.namespace == namespace) + .map(entry_from_record) + .filter(|entry| { + opts.category + .as_ref() + .is_none_or(|category| entry.category == *category) + }) + .filter(|entry| { + opts.session_id.is_none_or(|session| { + entry.session_id.as_deref() == Some(session) + || (opts.cross_session && entry.category == MemoryCategory::Conversation) + }) + }) + .filter_map(|mut entry| { + let score = query_score(&entry.content, query)?; + if opts.min_score.is_some_and(|minimum| score < minimum) { + return None; + } + entry.score = Some(score); + Some(entry) + }) + .collect::>(); + entries.sort_by(|left, right| { + right + .score + .partial_cmp(&left.score) + .unwrap_or(std::cmp::Ordering::Equal) + .then_with(|| right.timestamp.cmp(&left.timestamp)) + .then_with(|| left.key.cmp(&right.key)) + }); + entries.truncate(limit); + Ok(entries) + } + + async fn get(&self, namespace: &str, key: &str) -> Result> { + let records = self + .records + .read() + .map_err(|_| anyhow!("memory store lock poisoned"))?; + Ok(records + .values() + .find(|record| record.namespace == namespace && text_field(record, KEY) == Some(key)) + .map(entry_from_record)) + } + + async fn list( + &self, + namespace: Option<&str>, + category: Option<&MemoryCategory>, + session_id: Option<&str>, + ) -> Result> { + let records = self + .records + .read() + .map_err(|_| anyhow!("memory store lock poisoned"))?; + let mut entries = records + .values() + .filter(|record| namespace.is_none_or(|value| record.namespace == value)) + .map(entry_from_record) + .filter(|entry| category.is_none_or(|value| entry.category == *value)) + .filter(|entry| { + session_id.is_none_or(|value| entry.session_id.as_deref() == Some(value)) + }) + .collect::>(); + entries.sort_by(|left, right| left.key.cmp(&right.key)); + Ok(entries) + } + + async fn forget(&self, namespace: &str, key: &str) -> Result { + let mut records = self + .records + .write() + .map_err(|_| anyhow!("memory store lock poisoned"))?; + let id = records + .values() + .find(|record| record.namespace == namespace && text_field(record, KEY) == Some(key)) + .map(|record| record.id); + Ok(id.and_then(|id| records.remove(&id)).is_some()) + } + + async fn namespace_summaries(&self) -> Result> { + let records = self + .records + .read() + .map_err(|_| anyhow!("memory store lock poisoned"))?; + let mut summaries: HashMap)> = HashMap::new(); + for record in records.values() { + let summary = summaries + .entry(record.namespace.clone()) + .or_insert((0, record.updated_at)); + summary.0 += 1; + summary.1 = summary.1.max(record.updated_at); + } + let mut summaries = summaries + .into_iter() + .map(|(namespace, (count, updated))| NamespaceSummary { + namespace, + count, + last_updated: Some(updated.to_rfc3339()), + }) + .collect::>(); + summaries.sort_by(|left, right| left.namespace.cmp(&right.namespace)); + Ok(summaries) + } + + async fn count(&self) -> Result { + self.records + .read() + .map(|records| records.len()) + .map_err(|_| anyhow!("memory store lock poisoned")) + } + + async fn health_check(&self) -> bool { + self.records.read().is_ok() + } +} + +#[cfg(test)] +#[path = "memory_trait_tests.rs"] +mod tests; diff --git a/src/memory/store/memory_trait_tests.rs b/src/memory/store/memory_trait_tests.rs new file mode 100644 index 0000000..d244ece --- /dev/null +++ b/src/memory/store/memory_trait_tests.rs @@ -0,0 +1,69 @@ +use super::*; + +#[tokio::test] +async fn headline_memory_contract_supports_upsert_recall_and_taint() { + let store = InMemoryMemoryStore::new(); + store + .store_with_taint( + "agent", + "preference", + "Prefers dark mode", + MemoryCategory::Core, + Some("s1"), + MemoryTaint::ExternalSync, + ) + .await + .unwrap(); + store + .store_with_taint( + "agent", + "preference", + "Prefers dark themes", + MemoryCategory::Core, + Some("s1"), + MemoryTaint::ExternalSync, + ) + .await + .unwrap(); + + assert_eq!(store.count().await.unwrap(), 1); + let hits = store + .recall( + "dark themes", + 10, + RecallOpts { + namespace: Some("agent"), + ..Default::default() + }, + ) + .await + .unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].key, "preference"); + assert_eq!(hits[0].taint, MemoryTaint::ExternalSync); + assert_eq!(hits[0].score, Some(1.0)); +} + +#[tokio::test] +async fn memory_contract_filters_lists_summarises_and_forgets() { + let store = InMemoryMemoryStore::new(); + store + .store("a", "one", "alpha", MemoryCategory::Daily, None) + .await + .unwrap(); + store + .store("b", "two", "beta", MemoryCategory::Core, Some("s2")) + .await + .unwrap(); + + assert_eq!( + store.list(Some("b"), None, Some("s2")).await.unwrap().len(), + 1 + ); + assert!(store.get("b", "two").await.unwrap().is_some()); + let summaries = store.namespace_summaries().await.unwrap(); + assert_eq!(summaries.iter().map(|row| row.count).sum::(), 2); + assert!(store.health_check().await); + assert!(store.forget("a", "one").await.unwrap()); + assert!(!store.forget("a", "one").await.unwrap()); +} diff --git a/src/memory/store/mod.rs b/src/memory/store/mod.rs index 42ad17b..f2e8912 100644 --- a/src/memory/store/mod.rs +++ b/src/memory/store/mod.rs @@ -15,13 +15,9 @@ //! //! [`KvStore`], [`VectorStore`], and [`EntityIndex`] each hold their SQLite //! connection behind a `parking_lot::Mutex`, so calls into one instance -//! serialize cleanly. None of the three sets `PRAGMA busy_timeout`, so a -//! second process (or a second in-process connection) opening the same -//! database file gets an immediate `SQLITE_BUSY` on lock contention instead -//! of blocking and retrying — unlike the chunk store's connection pool, -//! which configures a 15s busy-timeout. Two independent handles to the same -//! path from this module should therefore not be assumed to compose safely -//! under write contention. +//! serialize cleanly. Their SQLite connections use a 15-second busy timeout, +//! allowing independent handles to the same database file to tolerate brief +//! write contention instead of failing immediately with `SQLITE_BUSY`. /// Markdown content store: source-of-truth files with YAML front matter, /// atomic writes, and `content_path`/`content_sha256` provenance pointers. @@ -30,19 +26,20 @@ pub mod content; pub mod entity_index; /// Global + namespace-scoped JSON key-value store (writes pass the [`safety`] guard). pub mod kv; +mod memory_trait; /// Secret / PII detection guard applied before KV writes are persisted. pub mod safety; /// Starter in-memory reference backend used by the smoke test. #[allow(clippy::module_inception)] pub mod store; -/// Core memory contract types ([`MemoryError`], [`MemoryRecord`], etc.). +/// Core memory contract types ([`StoreError`], [`MemoryRecord`], etc.). pub mod types; /// SQLite-backed packed-`f32` cosine vector DB and embedding backends. pub mod vectors; pub use store::{InMemoryMemoryStore, MemoryStore}; pub use types::{ - MemoryError, MemoryId, MemoryInput, MemoryQuery, MemoryRecord, MemoryResult, SearchHit, + MemoryId, MemoryInput, MemoryQuery, MemoryRecord, MemoryResult, SearchHit, StoreError, }; // ── Ported storage-primitive re-exports ───────────────────────────────────── diff --git a/src/memory/store/safety/mod.rs b/src/memory/store/safety/mod.rs index c466fc5..b1ea1b3 100644 --- a/src/memory/store/safety/mod.rs +++ b/src/memory/store/safety/mod.rs @@ -5,7 +5,7 @@ //! //! The exhaustive multilingual national-ID PII module (`safety::pii`, ~1k lines //! of checksum logic) is ported from OpenHuman and runs as part of -//! [`sanitize_text`]. The write-rejection boundary stays stricter than content +//! `sanitize_text`. The write-rejection boundary stays stricter than content //! scrubbing: formatted national IDs are rejected, while phone/email-like text is //! scrubbed from content without rejecting every write that mentions them. @@ -15,7 +15,7 @@ use regex::Regex; use serde_json::Value; /// Exhaustive checksum-gated multilingual national-ID PII module (ported from -/// OpenHuman). Content scrubbing runs from [`sanitize_text`]; the boundary +/// OpenHuman). Content scrubbing runs from `sanitize_text`; the boundary /// check is re-exported as [`has_likely_pii`]. pub mod pii; @@ -29,18 +29,18 @@ const MAX_JSON_SANITIZE_DEPTH: usize = 128; #[derive(Debug, Clone, Copy, Default, PartialEq, Eq)] pub struct SanitizationReport { /// Count of secret/token pattern matches rewritten in string text by the - /// [`REDACTION_PATTERNS`] pass. + /// text-pattern redaction pass. pub text_redactions: usize, /// Count of JSON object entries dropped wholesale because their key was - /// classified as sensitive by [`is_sensitive_key`]. + /// classified as sensitive by the key classifier. pub key_redactions: usize, - /// Count of full private-key blocks ([`BLOCK_PATTERNS`]) replaced; these are + /// Count of full private-key blocks replaced; these are /// the most severe hits since the entire block is removed. pub blocked_secret_hits: usize, /// Count of nodes collapsed because JSON nesting reached - /// [`MAX_JSON_SANITIZE_DEPTH`]; the subtree is replaced rather than walked. + /// the JSON traversal depth cap; the subtree is replaced rather than walked. pub depth_redactions: usize, - /// Count of personal-identifier matches ([`PII_PATTERNS`]) replaced by the + /// Count of personal-identifier matches replaced by the /// lightweight PII screen. pub pii_redactions: usize, } @@ -219,7 +219,7 @@ pub fn sanitize_text(value: &str) -> Sanitized { } /// Recursively scrub a JSON value: sensitive keys are replaced wholesale and -/// every string value runs through [`sanitize_text`]. +/// every string value runs through `sanitize_text`. pub fn sanitize_json(value: &Value) -> Sanitized { sanitize_json_inner(value, 0) } diff --git a/src/memory/store/safety/pii.rs b/src/memory/store/safety/pii.rs index 5f9b756..2f2a631 100644 --- a/src/memory/store/safety/pii.rs +++ b/src/memory/store/safety/pii.rs @@ -9,7 +9,7 @@ //! identifiers. The false-positive rate from format alone is too high; the //! checksums bring it back to acceptable. //! -//! 2. **Bypass-resistant.** Inputs are run through [`normalize`] before +//! 2. **Bypass-resistant.** Inputs are normalized before //! matching, which: //! - strips zero-width characters (U+200B/200C/200D/FEFF/2060/180E), //! - folds fullwidth digits (`０-９` → `0-9`) and fullwidth `．－／：` @@ -32,6 +32,9 @@ use std::sync::LazyLock; use super::{SanitizationReport, Sanitized}; +mod checks; +use checks::*; + // ---------- Replacement tokens ---------- const PII_RFC: &str = "[REDACTED_PII_RFC]"; @@ -487,627 +490,6 @@ fn fold_char(c: char) -> char { // ---------- Checksum helpers ---------- -fn digits(s: &str) -> Vec { - s.chars() - .filter(|c| c.is_ascii_digit()) - .map(|c| c.to_digit(10).expect("ascii digit")) - .collect() -} - -fn valid_cpf(d: &[u32]) -> bool { - if d.len() != 11 || d.iter().all(|x| *x == d[0]) { - return false; - } - let s1: u32 = (0..9).map(|i| d[i] * (10 - i as u32)).sum(); - let dv1 = (s1 * 10) % 11 % 10; - if dv1 != d[9] { - return false; - } - let s2: u32 = (0..10).map(|i| d[i] * (11 - i as u32)).sum(); - let dv2 = (s2 * 10) % 11 % 10; - dv2 == d[10] -} - -fn valid_cnpj(d: &[u32]) -> bool { - if d.len() != 14 || d.iter().all(|x| *x == d[0]) { - return false; - } - let w1: [u32; 12] = [5, 4, 3, 2, 9, 8, 7, 6, 5, 4, 3, 2]; - let s1: u32 = (0..12).map(|i| d[i] * w1[i]).sum(); - let r1 = s1 % 11; - let dv1 = if r1 < 2 { 0 } else { 11 - r1 }; - if dv1 != d[12] { - return false; - } - let w2: [u32; 13] = [6, 5, 4, 3, 2, 9, 8, 7, 6, 5, 4, 3, 2]; - let s2: u32 = (0..13).map(|i| d[i] * w2[i]).sum(); - let r2 = s2 % 11; - let dv2 = if r2 < 2 { 0 } else { 11 - r2 }; - dv2 == d[13] -} - -fn valid_cuit(d: &[u32]) -> bool { - if d.len() != 11 { - return false; - } - let w: [u32; 10] = [5, 4, 3, 2, 7, 6, 5, 4, 3, 2]; - let s: u32 = (0..10).map(|i| d[i] * w[i]).sum(); - let r = s % 11; - let dv = match r { - 0 => 0, - 1 => return false, - _ => 11 - r, - }; - dv == d[10] -} - -// Luhn — used for credit-card validation. -fn valid_luhn(s: &str) -> bool { - let d = digits(s); - if d.len() < 13 || d.len() > 19 { - return false; - } - let mut sum = 0u32; - let mut alt = false; - for x in d.iter().rev() { - let v = if alt { - let doubled = x * 2; - if doubled > 9 { - doubled - 9 - } else { - doubled - } - } else { - *x - }; - sum += v; - alt = !alt; - } - sum.is_multiple_of(10) -} - -// IBAN mod-97. Steps: strip spaces, move first 4 chars to end, expand letters -// (A=10..Z=35), divide as a big-integer mod 97, require remainder == 1. -fn valid_iban(s: &str) -> bool { - let cleaned: String = s.chars().filter(|c| !c.is_whitespace()).collect(); - if cleaned.len() < 15 || cleaned.len() > 34 { - return false; - } - if !cleaned.chars().take(2).all(|c| c.is_ascii_alphabetic()) { - return false; - } - if !cleaned[2..4].chars().all(|c| c.is_ascii_digit()) { - return false; - } - let rotated: String = cleaned[4..].chars().chain(cleaned[..4].chars()).collect(); - let mut remainder: u64 = 0; - for c in rotated.chars() { - let chunk = if let Some(d) = c.to_digit(10) { - d as u64 - } else if c.is_ascii_alphabetic() { - (c.to_ascii_uppercase() as u64) - ('A' as u64) + 10 - } else { - return false; - }; - // Expand into the running remainder digit-by-digit so we never need - // u128. Each letter contributes 2 decimal digits. - if chunk >= 10 { - remainder = (remainder * 100 + chunk) % 97; - } else { - remainder = (remainder * 10 + chunk) % 97; - } - } - remainder == 1 -} - -// Verhoeff — used for Aadhaar. -const VERHOEFF_D: [[u8; 10]; 10] = [ - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], - [1, 2, 3, 4, 0, 6, 7, 8, 9, 5], - [2, 3, 4, 0, 1, 7, 8, 9, 5, 6], - [3, 4, 0, 1, 2, 8, 9, 5, 6, 7], - [4, 0, 1, 2, 3, 9, 5, 6, 7, 8], - [5, 9, 8, 7, 6, 0, 4, 3, 2, 1], - [6, 5, 9, 8, 7, 1, 0, 4, 3, 2], - [7, 6, 5, 9, 8, 2, 1, 0, 4, 3], - [8, 7, 6, 5, 9, 3, 2, 1, 0, 4], - [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], -]; -const VERHOEFF_P: [[u8; 10]; 8] = [ - [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], - [1, 5, 7, 6, 2, 8, 3, 0, 9, 4], - [5, 8, 0, 3, 7, 9, 6, 1, 4, 2], - [8, 9, 1, 6, 0, 4, 3, 5, 2, 7], - [9, 4, 5, 3, 1, 2, 6, 8, 7, 0], - [4, 2, 8, 6, 5, 7, 3, 9, 0, 1], - [2, 7, 9, 3, 8, 0, 6, 4, 1, 5], - [7, 0, 4, 6, 9, 1, 3, 2, 5, 8], -]; - -fn valid_verhoeff(d: &[u32]) -> bool { - if d.len() != 12 { - return false; - } - // Aadhaar can't start with 0 or 1. - if d[0] < 2 { - return false; - } - let mut c: u8 = 0; - for (i, digit) in d.iter().rev().enumerate() { - c = VERHOEFF_D[c as usize][VERHOEFF_P[i % 8][*digit as usize] as usize]; - } - c == 0 -} - -// US SSN reserved/invalid ranges per SSA. -fn valid_ssn(s: &str) -> bool { - let d = digits(s); - if d.len() != 9 { - return false; - } - let area = d[0] * 100 + d[1] * 10 + d[2]; - let group = d[3] * 10 + d[4]; - let serial = d[5] * 1000 + d[6] * 100 + d[7] * 10 + d[8]; - if area == 0 || area == 666 || area >= 900 { - return false; - } - if group == 0 || serial == 0 { - return false; - } - true -} - -// Spain DNI check letter — 8 digits mod 23 indexes into a fixed letter table. -const DNI_LETTERS: &[u8; 23] = b"TRWAGMYFPDXBNJZSQVHLCKE"; - -fn valid_dni_es(s: &str) -> bool { - let upper = s.to_ascii_uppercase(); - let bytes = upper.as_bytes(); - if bytes.len() != 9 { - return false; - } - let num_str = &upper[..8]; - let letter = bytes[8]; - let Ok(num) = num_str.parse::() else { - return false; - }; - DNI_LETTERS[(num % 23) as usize] == letter -} - -fn valid_nie_es(s: &str) -> bool { - let upper = s.to_ascii_uppercase(); - let bytes = upper.as_bytes(); - if bytes.len() != 9 { - return false; - } - let prefix = match bytes[0] { - b'X' => 0u32, - b'Y' => 1, - b'Z' => 2, - _ => return false, - }; - let Ok(rest) = std::str::from_utf8(&bytes[1..8]) else { - return false; - }; - let Ok(num) = rest.parse::() else { - return false; - }; - let composed = prefix * 10_000_000 + num; - DNI_LETTERS[(composed % 23) as usize] == bytes[8] -} - -// UK NINO reserved-prefix blacklist. -fn valid_nino(s: &str) -> bool { - let upper = s.to_ascii_uppercase(); - let bytes = upper.as_bytes(); - if bytes.len() != 9 { - return false; - } - // First char cannot be D F I Q U V; second cannot be D F I O Q U V. - let bad_first = b"DFIQUV"; - let bad_second = b"DFIOQUV"; - if bad_first.contains(&bytes[0]) || bad_second.contains(&bytes[1]) { - return false; - } - // Reserved two-letter prefixes. - let reserved = ["BG", "GB", "KN", "NK", "NT", "TN", "ZZ"]; - let prefix = &upper[..2]; - if reserved.contains(&prefix) { - return false; - } - true -} - -// ---------- Tests ---------- - #[cfg(test)] -mod tests { - use super::*; - - fn redacts(input: &str, token: &str) { - let out = redact_pii(input); - assert!( - out.value.contains(token), - "expected {token} in output. input={input:?} output={out:?}" - ); - } - - fn unchanged(input: &str) { - let out = redact_pii(input); - assert_eq!( - out.value, input, - "expected no change; report={:?}", - out.report - ); - assert_eq!(out.report.pii_redactions, 0); - } - - // --- CPF --- - #[test] - fn cpf_formatted_valid_redacted() { - redacts("CPF: 111.444.777-35.", PII_CPF); - } - #[test] - fn cpf_formatted_invalid_kept() { - unchanged("CPF 111.444.777-99 nope"); - } - #[test] - fn cpf_all_same_digits_rejected() { - unchanged("Test 111.111.111-11"); - } - #[test] - fn cpf_bare_valid_redacted() { - redacts("Sem mascara 11144477735 ok", PII_CPF); - } - - // --- CNPJ --- - #[test] - fn cnpj_formatted_valid_redacted() { - redacts("CNPJ 11.222.333/0001-81", PII_CNPJ); - } - #[test] - fn cnpj_bare_valid_redacted() { - redacts("contract 11222333000181 yes", PII_CNPJ); - } - - // --- CUIT --- - #[test] - fn cuit_valid_redacted() { - redacts("CUIT 20-11111111-2", PII_CUIT); - } - #[test] - fn cuit_invalid_kept() { - unchanged("noise 20-12345678-0 noise"); - } - - // --- RFC --- - #[test] - fn rfc_redacted() { - redacts("Mi RFC VECJ880326XK4 .", PII_RFC); - } - #[test] - fn rfc_lowercase_redacted() { - redacts("rfc vecj880326xk4", PII_RFC); - } - - // --- My Number --- - #[test] - fn my_number_redacted_with_keyword() { - redacts("マイナンバー: 123456789012", PII_MYNUM); - } - #[test] - fn bare_12_digits_without_keyword_kept() { - unchanged("Order 123456789012 shipped today."); - } - - // --- E.164 + NANP phone --- - #[test] - fn e164_redacted() { - redacts("phone +15551234567", PII_PHONE); - } - #[test] - fn nanp_formatted_redacted() { - redacts("call 415-555-0123 thanks", PII_PHONE); - } - #[test] - fn nanp_with_country_code_redacted() { - redacts("+1 (212) 555-7890", PII_PHONE); - } - #[test] - fn nanp_invalid_area_code_kept() { - unchanged("score 115-555-0123 ish"); - } - - // --- SSN --- - #[test] - fn ssn_valid_redacted() { - redacts("ssn 123-45-6789", PII_SSN); - } - #[test] - fn ssn_reserved_area_kept() { - unchanged("test 666-12-3456"); - } - #[test] - fn ssn_zero_serial_kept() { - unchanged("test 123-45-0000"); - } - - // --- Credit card / Luhn --- - #[test] - fn credit_card_visa_redacted() { - // Visa test number with valid Luhn. - redacts("card 4111 1111 1111 1111 thanks", PII_CC); - } - #[test] - fn credit_card_amex_redacted() { - redacts("card 378282246310005 used", PII_CC); - } - #[test] - fn credit_card_invalid_luhn_kept() { - unchanged("invoice 4111 1111 1111 1112"); - } - - // --- IBAN --- - #[test] - fn iban_de_redacted() { - // Known test IBAN with valid mod-97. - redacts("IBAN DE89370400440532013000 ok", PII_IBAN); - } - #[test] - fn iban_invalid_kept() { - unchanged("noise DE89370400440532013001 noise"); - } - - // --- Aadhaar --- - #[test] - fn aadhaar_formatted_verhoeff_valid_redacted() { - // 234123412346 is a known Verhoeff-valid Aadhaar test number. - redacts("Aadhaar 2341 2341 2346", PII_AADHAAR); - } - #[test] - fn aadhaar_keyword_bare_redacted() { - redacts("Aadhaar: 234123412346", PII_AADHAAR); - } - #[test] - fn aadhaar_invalid_verhoeff_kept() { - unchanged("Random 2341 2341 2345 nope"); - } - - // --- PAN-IN --- - #[test] - fn pan_in_redacted() { - redacts("PAN: ABCDE1234F", PII_PAN_IN); - } - - // --- NINO --- - #[test] - fn nino_redacted() { - redacts("NI no AB123456C", PII_NINO); - } - #[test] - fn nino_reserved_prefix_kept() { - unchanged("BG123456A"); - } - - // --- DNI / NIE --- - #[test] - fn dni_es_redacted() { - redacts("DNI 12345678Z", PII_DNI); - } - #[test] - fn dni_es_bad_letter_kept() { - unchanged("ID 12345678A code"); - } - #[test] - fn nie_es_redacted() { - redacts("NIE X1234567L", PII_DNI); - } - - // --- RRN Korea --- - #[test] - fn rrn_kr_redacted() { - redacts("주민번호 900101-1234567", PII_RRN); - } - #[test] - fn rrn_kr_bad_gender_digit_kept() { - unchanged("ref 900101-5234567 nope"); - } - - // --- Bypass resistance --- - #[test] - fn fullwidth_digits_cannot_bypass_cpf() { - // 111.444.777-35 with fullwidth digits and punctuation. - let input = "CPF: １１１．４４４．７７７－３５ done"; - let out = redact_pii(input); - assert!(out.value.contains(PII_CPF), "got {out:?}"); - } - - #[test] - fn zero_width_chars_cannot_bypass_ssn() { - // U+200B inserted between digits. - let input = "ssn 1\u{200B}23-4\u{200B}5-6789 done"; - let out = redact_pii(input); - assert!(out.value.contains(PII_SSN), "got {out:?}"); - } - - #[test] - fn arabic_indic_digits_normalize_for_phone() { - let input = "phone +١٥٥٥١٢٣٤٥٦٧"; - let out = redact_pii(input); - assert!(out.value.contains(PII_PHONE), "got {out:?}"); - } - - // --- Aggressive mix end-to-end --- - #[test] - fn aggressive_mixed_document() { - let input = "\ -Cliente RFC VECJ880326XK4. \ -Empresa CNPJ 11.222.333/0001-81. \ -Argentino CUIT 20-11111111-2. \ -Brasileiro CPF 111.444.777-35. \ -マイナンバー: 123456789012. \ -SSN 123-45-6789. \ -Card 4111 1111 1111 1111. \ -IBAN DE89370400440532013000. \ -PAN ABCDE1234F. \ -NI AB123456C. \ -DNI 12345678Z. \ -RRN 900101-1234567. \ -Phone +15551234567."; - let out = redact_pii(input); - for token in [ - PII_RFC, PII_CNPJ, PII_CUIT, PII_CPF, PII_MYNUM, PII_SSN, PII_CC, PII_IBAN, PII_PAN_IN, - PII_NINO, PII_DNI, PII_RRN, PII_PHONE, - ] { - assert!( - out.value.contains(token), - "missing {token} in: {}", - out.value - ); - } - assert!(out.report.pii_redactions >= 13); - } - - // --- has_likely_pii --- - #[test] - fn has_likely_pii_detects_cpf() { - assert!(has_likely_pii("user/111.444.777-35")); - } - - #[test] - fn has_likely_email_detects_email_without_changing_boundary_pii() { - assert!(has_likely_email("user/alice@example.com")); - assert!(!has_likely_pii("user/alice@example.com")); - } - #[test] - fn has_likely_pii_quiet_on_normal_text() { - assert!(!has_likely_pii("memory/global/preferences")); - } - - /// Regression: zero-padded millisecond-timestamp keys must NOT be - /// flagged as PII even when the digit run happens to satisfy Luhn. - /// `redact_pii` content scrubbing may still flag the same string — - /// `has_likely_pii` (used for boundary rejection of internal keys) - /// must stay strict to formatted/keyword PII only. - #[test] - fn has_likely_pii_ignores_bare_luhn_timestamp_keys() { - // 18-digit padded timestamps where the digit total mod 10 == 0 - // (the Luhn-passing case that previously rejected autocomplete - // KV writes and screen-intelligence document writes). - for key in [ - "accepted:000001747729035001", - "completion:000001747729035011", - "screen_intelligence_vision-1747729035001-VSCode", - ] { - assert!( - !has_likely_pii(key), - "internal key {key:?} must not be rejected as PII" - ); - } - } - - /// Strict boundary check should still reject formatted PII even though - /// it skips bare-numeric checksum patterns. - #[test] - fn has_likely_pii_still_blocks_formatted_secrets() { - assert!(has_likely_pii("ssn-123-45-6789")); - assert!(has_likely_pii("cliente-RFC-VECJ880326XK4")); - assert!(has_likely_pii("cuit-20-11111111-2")); - } - - /// Regression for Sentry TAURI-RUST-54T / GH #2848: scanner-built - /// `namespace` and `key` values containing bare-numeric phone-shaped - /// digit runs (WhatsApp group JID `-@g.us`, WhatsApp - /// broadcast `@broadcast`, US-prefixed WhatsApp 1:1 JID, - /// telegram numeric peer ID) must NOT be rejected by the boundary - /// PII check. NANP matches `\d{10,11}` with optional separators — - /// strict mode must skip it. Content scrubbing via `redact_pii` - /// continues to redact these substrings (see - /// `redact_pii_still_blurs_bare_phone_in_content` below). - #[test] - fn has_likely_pii_ignores_scanner_bare_phone_keys() { - for key in [ - // WhatsApp group JID — chat_id = "-@g.us" - "12025551234-1543890267@g.us:2026-05-30", - // WhatsApp broadcast list - "12025551234@broadcast:2026-05-30", - // WhatsApp 1:1 JID, country-coded US number (`1` + 10 digits) - "12025551234@c.us:2026-05-30", - // Same shape carried in the namespace - "whatsapp-web:12025551234@c.us", - "whatsapp-web:12025551234-1543890267@g.us", - // Telegram numeric peer_id key - "4123456789:2026-05-30", - ] { - assert!( - !has_likely_pii(key), - "scanner-built key {key:?} must not be rejected as PII" - ); - } - } - - /// Same regression but for the E.164 (`+`-prefixed) shape — iMessage - /// posts `key = format!("{chat_id}:{day}")` where `chat_id` can be - /// `+12025551234`. Strict mode must skip; content redaction stays. - #[test] - fn has_likely_pii_ignores_bare_e164_phone_keys() { - for key in [ - "+12025551234:2026-05-30", - "imessage:+12025551234", - "imessage:+12025551234:2026-05-30", - ] { - assert!( - !has_likely_pii(key), - "E.164-shaped key {key:?} must not be rejected as PII" - ); - } - } - - /// `redact_pii` (content scrubbing path — NOT the boundary check) - /// must still redact formatted NANP and E.164 phone numbers found - /// inside document bodies. False positives in the content path only - /// blur substring bytes; they do not reject the write — which is the - /// asymmetry this PR preserves vs. the boundary check. - /// - /// Note: bare 10-digit NANP runs (`2025551234` with no separators) - /// are NOT reached by `redact_pii` at all — the SCREEN fast-path - /// requires either `\d{11,}`, a separator, or `+`, so a bare 10-digit - /// run short-circuits as "no candidate". That pre-existed this PR; a - /// pinning sentinel for it lives below. - #[test] - fn redact_pii_still_blurs_formatted_and_e164_phone_in_content() { - let out = redact_pii("call me at 202-555-1234 or +12025551234"); - let n_phone = out.value.matches(PII_PHONE).count(); - assert!( - n_phone >= 2, - "redact_pii must still blur both formatted NANP and E.164 phones in content, \ - got {n_phone} PII_PHONE token(s) in: {}", - out.value - ); - assert!(out.report.pii_redactions >= 2); - } - - /// Sentinel pinning a pre-existing SCREEN limitation: a bare 10-digit - /// NANP run (`2025551234` with no separators) is short-circuited by - /// the `SCREEN` fast-path because no `SCREEN` regex matches a 10-digit - /// bare run (`\d{11,}` is the closest, but it needs 11+). This is the - /// status quo on `main` — this PR does not change it. The test exists - /// so any future widening of `SCREEN` (e.g. to catch bare NANP) trips - /// here as a deliberate review checkpoint, NOT a regression. - #[test] - fn redact_pii_does_not_reach_bare_10_digit_nanp_today() { - let out = redact_pii("call me at 2025551234 thanks"); - assert!( - !out.value.contains(PII_PHONE), - "SCREEN fast-path historically skips bare 10-digit NANP — \ - if this test fails, SCREEN was widened; revisit the boundary-check \ - behavior in `has_likely_pii` before adjusting. Got: {}", - out.value - ); - } - - #[test] - fn empty_text_is_noop() { - unchanged(""); - } -} +#[path = "pii_tests.rs"] +mod tests; diff --git a/src/memory/store/safety/pii/checks.rs b/src/memory/store/safety/pii/checks.rs new file mode 100644 index 0000000..0b64470 --- /dev/null +++ b/src/memory/store/safety/pii/checks.rs @@ -0,0 +1,232 @@ +//! Checksum and structural validators for PII candidates. + +pub(super) fn digits(s: &str) -> Vec { + s.chars() + .filter(|c| c.is_ascii_digit()) + .map(|c| c.to_digit(10).expect("ascii digit")) + .collect() +} + +pub(super) fn valid_cpf(d: &[u32]) -> bool { + if d.len() != 11 || d.iter().all(|x| *x == d[0]) { + return false; + } + let s1: u32 = (0..9).map(|i| d[i] * (10 - i as u32)).sum(); + let dv1 = (s1 * 10) % 11 % 10; + if dv1 != d[9] { + return false; + } + let s2: u32 = (0..10).map(|i| d[i] * (11 - i as u32)).sum(); + let dv2 = (s2 * 10) % 11 % 10; + dv2 == d[10] +} + +pub(super) fn valid_cnpj(d: &[u32]) -> bool { + if d.len() != 14 || d.iter().all(|x| *x == d[0]) { + return false; + } + let w1: [u32; 12] = [5, 4, 3, 2, 9, 8, 7, 6, 5, 4, 3, 2]; + let s1: u32 = (0..12).map(|i| d[i] * w1[i]).sum(); + let r1 = s1 % 11; + let dv1 = if r1 < 2 { 0 } else { 11 - r1 }; + if dv1 != d[12] { + return false; + } + let w2: [u32; 13] = [6, 5, 4, 3, 2, 9, 8, 7, 6, 5, 4, 3, 2]; + let s2: u32 = (0..13).map(|i| d[i] * w2[i]).sum(); + let r2 = s2 % 11; + let dv2 = if r2 < 2 { 0 } else { 11 - r2 }; + dv2 == d[13] +} + +pub(super) fn valid_cuit(d: &[u32]) -> bool { + if d.len() != 11 { + return false; + } + let w: [u32; 10] = [5, 4, 3, 2, 7, 6, 5, 4, 3, 2]; + let s: u32 = (0..10).map(|i| d[i] * w[i]).sum(); + let r = s % 11; + let dv = match r { + 0 => 0, + 1 => return false, + _ => 11 - r, + }; + dv == d[10] +} + +// Luhn — used for credit-card validation. +pub(super) fn valid_luhn(s: &str) -> bool { + let d = digits(s); + if d.len() < 13 || d.len() > 19 { + return false; + } + let mut sum = 0u32; + let mut alt = false; + for x in d.iter().rev() { + let v = if alt { + let doubled = x * 2; + if doubled > 9 { + doubled - 9 + } else { + doubled + } + } else { + *x + }; + sum += v; + alt = !alt; + } + sum.is_multiple_of(10) +} + +// IBAN mod-97. Steps: strip spaces, move first 4 chars to end, expand letters +// (A=10..Z=35), divide as a big-integer mod 97, require remainder == 1. +pub(super) fn valid_iban(s: &str) -> bool { + let cleaned: String = s.chars().filter(|c| !c.is_whitespace()).collect(); + if cleaned.len() < 15 || cleaned.len() > 34 { + return false; + } + if !cleaned.chars().take(2).all(|c| c.is_ascii_alphabetic()) { + return false; + } + if !cleaned[2..4].chars().all(|c| c.is_ascii_digit()) { + return false; + } + let rotated: String = cleaned[4..].chars().chain(cleaned[..4].chars()).collect(); + let mut remainder: u64 = 0; + for c in rotated.chars() { + let chunk = if let Some(d) = c.to_digit(10) { + d as u64 + } else if c.is_ascii_alphabetic() { + (c.to_ascii_uppercase() as u64) - ('A' as u64) + 10 + } else { + return false; + }; + // Expand into the running remainder digit-by-digit so we never need + // u128. Each letter contributes 2 decimal digits. + if chunk >= 10 { + remainder = (remainder * 100 + chunk) % 97; + } else { + remainder = (remainder * 10 + chunk) % 97; + } + } + remainder == 1 +} + +// Verhoeff — used for Aadhaar. +const VERHOEFF_D: [[u8; 10]; 10] = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 2, 3, 4, 0, 6, 7, 8, 9, 5], + [2, 3, 4, 0, 1, 7, 8, 9, 5, 6], + [3, 4, 0, 1, 2, 8, 9, 5, 6, 7], + [4, 0, 1, 2, 3, 9, 5, 6, 7, 8], + [5, 9, 8, 7, 6, 0, 4, 3, 2, 1], + [6, 5, 9, 8, 7, 1, 0, 4, 3, 2], + [7, 6, 5, 9, 8, 2, 1, 0, 4, 3], + [8, 7, 6, 5, 9, 3, 2, 1, 0, 4], + [9, 8, 7, 6, 5, 4, 3, 2, 1, 0], +]; +const VERHOEFF_P: [[u8; 10]; 8] = [ + [0, 1, 2, 3, 4, 5, 6, 7, 8, 9], + [1, 5, 7, 6, 2, 8, 3, 0, 9, 4], + [5, 8, 0, 3, 7, 9, 6, 1, 4, 2], + [8, 9, 1, 6, 0, 4, 3, 5, 2, 7], + [9, 4, 5, 3, 1, 2, 6, 8, 7, 0], + [4, 2, 8, 6, 5, 7, 3, 9, 0, 1], + [2, 7, 9, 3, 8, 0, 6, 4, 1, 5], + [7, 0, 4, 6, 9, 1, 3, 2, 5, 8], +]; + +pub(super) fn valid_verhoeff(d: &[u32]) -> bool { + if d.len() != 12 { + return false; + } + // Aadhaar can't start with 0 or 1. + if d[0] < 2 { + return false; + } + let mut c: u8 = 0; + for (i, digit) in d.iter().rev().enumerate() { + c = VERHOEFF_D[c as usize][VERHOEFF_P[i % 8][*digit as usize] as usize]; + } + c == 0 +} + +// US SSN reserved/invalid ranges per SSA. +pub(super) fn valid_ssn(s: &str) -> bool { + let d = digits(s); + if d.len() != 9 { + return false; + } + let area = d[0] * 100 + d[1] * 10 + d[2]; + let group = d[3] * 10 + d[4]; + let serial = d[5] * 1000 + d[6] * 100 + d[7] * 10 + d[8]; + if area == 0 || area == 666 || area >= 900 { + return false; + } + if group == 0 || serial == 0 { + return false; + } + true +} + +// Spain DNI check letter — 8 digits mod 23 indexes into a fixed letter table. +const DNI_LETTERS: &[u8; 23] = b"TRWAGMYFPDXBNJZSQVHLCKE"; + +pub(super) fn valid_dni_es(s: &str) -> bool { + let upper = s.to_ascii_uppercase(); + let bytes = upper.as_bytes(); + if bytes.len() != 9 { + return false; + } + let num_str = &upper[..8]; + let letter = bytes[8]; + let Ok(num) = num_str.parse::() else { + return false; + }; + DNI_LETTERS[(num % 23) as usize] == letter +} + +pub(super) fn valid_nie_es(s: &str) -> bool { + let upper = s.to_ascii_uppercase(); + let bytes = upper.as_bytes(); + if bytes.len() != 9 { + return false; + } + let prefix = match bytes[0] { + b'X' => 0u32, + b'Y' => 1, + b'Z' => 2, + _ => return false, + }; + let Ok(rest) = std::str::from_utf8(&bytes[1..8]) else { + return false; + }; + let Ok(num) = rest.parse::() else { + return false; + }; + let composed = prefix * 10_000_000 + num; + DNI_LETTERS[(composed % 23) as usize] == bytes[8] +} + +// UK NINO reserved-prefix blacklist. +pub(super) fn valid_nino(s: &str) -> bool { + let upper = s.to_ascii_uppercase(); + let bytes = upper.as_bytes(); + if bytes.len() != 9 { + return false; + } + // First char cannot be D F I Q U V; second cannot be D F I O Q U V. + let bad_first = b"DFIQUV"; + let bad_second = b"DFIOQUV"; + if bad_first.contains(&bytes[0]) || bad_second.contains(&bytes[1]) { + return false; + } + // Reserved two-letter prefixes. + let reserved = ["BG", "GB", "KN", "NK", "NT", "TN", "ZZ"]; + let prefix = &upper[..2]; + if reserved.contains(&prefix) { + return false; + } + true +} diff --git a/src/memory/store/safety/pii_tests.rs b/src/memory/store/safety/pii_tests.rs new file mode 100644 index 0000000..f97a292 --- /dev/null +++ b/src/memory/store/safety/pii_tests.rs @@ -0,0 +1,388 @@ +use super::*; + +fn redacts(input: &str, token: &str) { + let out = redact_pii(input); + assert!( + out.value.contains(token), + "expected {token} in output. input={input:?} output={out:?}" + ); +} + +fn unchanged(input: &str) { + let out = redact_pii(input); + assert_eq!( + out.value, input, + "expected no change; report={:?}", + out.report + ); + assert_eq!(out.report.pii_redactions, 0); +} + +// --- CPF --- +#[test] +fn cpf_formatted_valid_redacted() { + redacts("CPF: 111.444.777-35.", PII_CPF); +} +#[test] +fn cpf_formatted_invalid_kept() { + unchanged("CPF 111.444.777-99 nope"); +} +#[test] +fn cpf_all_same_digits_rejected() { + unchanged("Test 111.111.111-11"); +} +#[test] +fn cpf_bare_valid_redacted() { + redacts("Sem mascara 11144477735 ok", PII_CPF); +} + +// --- CNPJ --- +#[test] +fn cnpj_formatted_valid_redacted() { + redacts("CNPJ 11.222.333/0001-81", PII_CNPJ); +} +#[test] +fn cnpj_bare_valid_redacted() { + redacts("contract 11222333000181 yes", PII_CNPJ); +} + +// --- CUIT --- +#[test] +fn cuit_valid_redacted() { + redacts("CUIT 20-11111111-2", PII_CUIT); +} +#[test] +fn cuit_invalid_kept() { + unchanged("noise 20-12345678-0 noise"); +} + +// --- RFC --- +#[test] +fn rfc_redacted() { + redacts("Mi RFC VECJ880326XK4 .", PII_RFC); +} +#[test] +fn rfc_lowercase_redacted() { + redacts("rfc vecj880326xk4", PII_RFC); +} + +// --- My Number --- +#[test] +fn my_number_redacted_with_keyword() { + redacts("マイナンバー: 123456789012", PII_MYNUM); +} +#[test] +fn bare_12_digits_without_keyword_kept() { + unchanged("Order 123456789012 shipped today."); +} + +// --- E.164 + NANP phone --- +#[test] +fn e164_redacted() { + redacts("phone +15551234567", PII_PHONE); +} +#[test] +fn nanp_formatted_redacted() { + redacts("call 415-555-0123 thanks", PII_PHONE); +} +#[test] +fn nanp_with_country_code_redacted() { + redacts("+1 (212) 555-7890", PII_PHONE); +} +#[test] +fn nanp_invalid_area_code_kept() { + unchanged("score 115-555-0123 ish"); +} + +// --- SSN --- +#[test] +fn ssn_valid_redacted() { + redacts("ssn 123-45-6789", PII_SSN); +} +#[test] +fn ssn_reserved_area_kept() { + unchanged("test 666-12-3456"); +} +#[test] +fn ssn_zero_serial_kept() { + unchanged("test 123-45-0000"); +} + +// --- Credit card / Luhn --- +#[test] +fn credit_card_visa_redacted() { + // Visa test number with valid Luhn. + redacts("card 4111 1111 1111 1111 thanks", PII_CC); +} +#[test] +fn credit_card_amex_redacted() { + redacts("card 378282246310005 used", PII_CC); +} +#[test] +fn credit_card_invalid_luhn_kept() { + unchanged("invoice 4111 1111 1111 1112"); +} + +// --- IBAN --- +#[test] +fn iban_de_redacted() { + // Known test IBAN with valid mod-97. + redacts("IBAN DE89370400440532013000 ok", PII_IBAN); +} +#[test] +fn iban_invalid_kept() { + unchanged("noise DE89370400440532013001 noise"); +} + +// --- Aadhaar --- +#[test] +fn aadhaar_formatted_verhoeff_valid_redacted() { + // 234123412346 is a known Verhoeff-valid Aadhaar test number. + redacts("Aadhaar 2341 2341 2346", PII_AADHAAR); +} +#[test] +fn aadhaar_keyword_bare_redacted() { + redacts("Aadhaar: 234123412346", PII_AADHAAR); +} +#[test] +fn aadhaar_invalid_verhoeff_kept() { + unchanged("Random 2341 2341 2345 nope"); +} + +// --- PAN-IN --- +#[test] +fn pan_in_redacted() { + redacts("PAN: ABCDE1234F", PII_PAN_IN); +} + +// --- NINO --- +#[test] +fn nino_redacted() { + redacts("NI no AB123456C", PII_NINO); +} +#[test] +fn nino_reserved_prefix_kept() { + unchanged("BG123456A"); +} + +// --- DNI / NIE --- +#[test] +fn dni_es_redacted() { + redacts("DNI 12345678Z", PII_DNI); +} +#[test] +fn dni_es_bad_letter_kept() { + unchanged("ID 12345678A code"); +} +#[test] +fn nie_es_redacted() { + redacts("NIE X1234567L", PII_DNI); +} + +// --- RRN Korea --- +#[test] +fn rrn_kr_redacted() { + redacts("주민번호 900101-1234567", PII_RRN); +} +#[test] +fn rrn_kr_bad_gender_digit_kept() { + unchanged("ref 900101-5234567 nope"); +} + +// --- Bypass resistance --- +#[test] +fn fullwidth_digits_cannot_bypass_cpf() { + // 111.444.777-35 with fullwidth digits and punctuation. + let input = "CPF: １１１．４４４．７７７－３５ done"; + let out = redact_pii(input); + assert!(out.value.contains(PII_CPF), "got {out:?}"); +} + +#[test] +fn zero_width_chars_cannot_bypass_ssn() { + // U+200B inserted between digits. + let input = "ssn 1\u{200B}23-4\u{200B}5-6789 done"; + let out = redact_pii(input); + assert!(out.value.contains(PII_SSN), "got {out:?}"); +} + +#[test] +fn arabic_indic_digits_normalize_for_phone() { + let input = "phone +١٥٥٥١٢٣٤٥٦٧"; + let out = redact_pii(input); + assert!(out.value.contains(PII_PHONE), "got {out:?}"); +} + +// --- Aggressive mix end-to-end --- +#[test] +fn aggressive_mixed_document() { + let input = "\ +Cliente RFC VECJ880326XK4. \ +Empresa CNPJ 11.222.333/0001-81. \ +Argentino CUIT 20-11111111-2. \ +Brasileiro CPF 111.444.777-35. \ +マイナンバー: 123456789012. \ +SSN 123-45-6789. \ +Card 4111 1111 1111 1111. \ +IBAN DE89370400440532013000. \ +PAN ABCDE1234F. \ +NI AB123456C. \ +DNI 12345678Z. \ +RRN 900101-1234567. \ +Phone +15551234567."; + let out = redact_pii(input); + for token in [ + PII_RFC, PII_CNPJ, PII_CUIT, PII_CPF, PII_MYNUM, PII_SSN, PII_CC, PII_IBAN, PII_PAN_IN, + PII_NINO, PII_DNI, PII_RRN, PII_PHONE, + ] { + assert!( + out.value.contains(token), + "missing {token} in: {}", + out.value + ); + } + assert!(out.report.pii_redactions >= 13); +} + +// --- has_likely_pii --- +#[test] +fn has_likely_pii_detects_cpf() { + assert!(has_likely_pii("user/111.444.777-35")); +} + +#[test] +fn has_likely_email_detects_email_without_changing_boundary_pii() { + assert!(has_likely_email("user/alice@example.com")); + assert!(!has_likely_pii("user/alice@example.com")); +} +#[test] +fn has_likely_pii_quiet_on_normal_text() { + assert!(!has_likely_pii("memory/global/preferences")); +} + +/// Regression: zero-padded millisecond-timestamp keys must NOT be +/// flagged as PII even when the digit run happens to satisfy Luhn. +/// `redact_pii` content scrubbing may still flag the same string — +/// `has_likely_pii` (used for boundary rejection of internal keys) +/// must stay strict to formatted/keyword PII only. +#[test] +fn has_likely_pii_ignores_bare_luhn_timestamp_keys() { + // 18-digit padded timestamps where the digit total mod 10 == 0 + // (the Luhn-passing case that previously rejected autocomplete + // KV writes and screen-intelligence document writes). + for key in [ + "accepted:000001747729035001", + "completion:000001747729035011", + "screen_intelligence_vision-1747729035001-VSCode", + ] { + assert!( + !has_likely_pii(key), + "internal key {key:?} must not be rejected as PII" + ); + } +} + +/// Strict boundary check should still reject formatted PII even though +/// it skips bare-numeric checksum patterns. +#[test] +fn has_likely_pii_still_blocks_formatted_secrets() { + assert!(has_likely_pii("ssn-123-45-6789")); + assert!(has_likely_pii("cliente-RFC-VECJ880326XK4")); + assert!(has_likely_pii("cuit-20-11111111-2")); +} + +/// Regression for Sentry TAURI-RUST-54T / GH #2848: scanner-built +/// `namespace` and `key` values containing bare-numeric phone-shaped +/// digit runs (WhatsApp group JID `-@g.us`, WhatsApp +/// broadcast `@broadcast`, US-prefixed WhatsApp 1:1 JID, +/// telegram numeric peer ID) must NOT be rejected by the boundary +/// PII check. NANP matches `\d{10,11}` with optional separators — +/// strict mode must skip it. Content scrubbing via `redact_pii` +/// continues to redact these substrings (see +/// `redact_pii_still_blurs_bare_phone_in_content` below). +#[test] +fn has_likely_pii_ignores_scanner_bare_phone_keys() { + for key in [ + // WhatsApp group JID — chat_id = "-@g.us" + "12025551234-1543890267@g.us:2026-05-30", + // WhatsApp broadcast list + "12025551234@broadcast:2026-05-30", + // WhatsApp 1:1 JID, country-coded US number (`1` + 10 digits) + "12025551234@c.us:2026-05-30", + // Same shape carried in the namespace + "whatsapp-web:12025551234@c.us", + "whatsapp-web:12025551234-1543890267@g.us", + // Telegram numeric peer_id key + "4123456789:2026-05-30", + ] { + assert!( + !has_likely_pii(key), + "scanner-built key {key:?} must not be rejected as PII" + ); + } +} + +/// Same regression but for the E.164 (`+`-prefixed) shape — iMessage +/// posts `key = format!("{chat_id}:{day}")` where `chat_id` can be +/// `+12025551234`. Strict mode must skip; content redaction stays. +#[test] +fn has_likely_pii_ignores_bare_e164_phone_keys() { + for key in [ + "+12025551234:2026-05-30", + "imessage:+12025551234", + "imessage:+12025551234:2026-05-30", + ] { + assert!( + !has_likely_pii(key), + "E.164-shaped key {key:?} must not be rejected as PII" + ); + } +} + +/// `redact_pii` (content scrubbing path — NOT the boundary check) +/// must still redact formatted NANP and E.164 phone numbers found +/// inside document bodies. False positives in the content path only +/// blur substring bytes; they do not reject the write — which is the +/// asymmetry this PR preserves vs. the boundary check. +/// +/// Note: bare 10-digit NANP runs (`2025551234` with no separators) +/// are NOT reached by `redact_pii` at all — the SCREEN fast-path +/// requires either `\d{11,}`, a separator, or `+`, so a bare 10-digit +/// run short-circuits as "no candidate". That pre-existed this PR; a +/// pinning sentinel for it lives below. +#[test] +fn redact_pii_still_blurs_formatted_and_e164_phone_in_content() { + let out = redact_pii("call me at 202-555-1234 or +12025551234"); + let n_phone = out.value.matches(PII_PHONE).count(); + assert!( + n_phone >= 2, + "redact_pii must still blur both formatted NANP and E.164 phones in content, \ + got {n_phone} PII_PHONE token(s) in: {}", + out.value + ); + assert!(out.report.pii_redactions >= 2); +} + +/// Sentinel pinning a pre-existing SCREEN limitation: a bare 10-digit +/// NANP run (`2025551234` with no separators) is short-circuited by +/// the `SCREEN` fast-path because no `SCREEN` regex matches a 10-digit +/// bare run (`\d{11,}` is the closest, but it needs 11+). This is the +/// status quo on `main` — this PR does not change it. The test exists +/// so any future widening of `SCREEN` (e.g. to catch bare NANP) trips +/// here as a deliberate review checkpoint, NOT a regression. +#[test] +fn redact_pii_does_not_reach_bare_10_digit_nanp_today() { + let out = redact_pii("call me at 2025551234 thanks"); + assert!( + !out.value.contains(PII_PHONE), + "SCREEN fast-path historically skips bare 10-digit NANP — \ + if this test fails, SCREEN was widened; revisit the boundary-check \ + behavior in `has_likely_pii` before adjusting. Got: {}", + out.value + ); +} + +#[test] +fn empty_text_is_noop() { + unchanged(""); +} diff --git a/src/memory/store/store.rs b/src/memory/store/store.rs index dd83adb..64f3160 100644 --- a/src/memory/store/store.rs +++ b/src/memory/store/store.rs @@ -12,7 +12,7 @@ use std::sync::{Arc, RwLock}; use async_trait::async_trait; use super::types::{ - MemoryError, MemoryId, MemoryInput, MemoryQuery, MemoryRecord, MemoryResult, SearchHit, + MemoryId, MemoryInput, MemoryQuery, MemoryRecord, MemoryResult, SearchHit, StoreError, }; /// Storage backend contract for memory records. @@ -25,9 +25,9 @@ pub trait MemoryStore: Send + Sync { /// Materialize a record from `input` and persist it, returning the stored /// record (with its assigned [`MemoryId`] and timestamps). async fn insert(&self, input: MemoryInput) -> MemoryResult; - /// Fetch the record for `id`, or [`MemoryError::NotFound`] if absent. + /// Fetch the record for `id`, or [`StoreError::NotFound`] if absent. async fn get(&self, id: MemoryId) -> MemoryResult; - /// Remove and return the record for `id`, or [`MemoryError::NotFound`] if absent. + /// Remove and return the record for `id`, or [`StoreError::NotFound`] if absent. async fn delete(&self, id: MemoryId) -> MemoryResult; /// Return scored [`SearchHit`]s matching `query`, ordered most relevant first. async fn search(&self, query: MemoryQuery) -> MemoryResult>; @@ -42,7 +42,7 @@ pub trait MemoryStore: Send + Sync { #[derive(Clone, Debug, Default)] pub struct InMemoryMemoryStore { /// Records indexed by [`MemoryId`]; the `BTreeMap` keeps a stable key order. - records: Arc>>, + pub(super) records: Arc>>, } impl InMemoryMemoryStore { @@ -69,7 +69,7 @@ impl MemoryStore for InMemoryMemoryStore { .expect("memory store lock poisoned") .get(&id) .cloned() - .ok_or(MemoryError::NotFound(id)) + .ok_or(StoreError::NotFound(id)) } async fn delete(&self, id: MemoryId) -> MemoryResult { @@ -77,11 +77,20 @@ impl MemoryStore for InMemoryMemoryStore { .write() .expect("memory store lock poisoned") .remove(&id) - .ok_or(MemoryError::NotFound(id)) + .ok_or(StoreError::NotFound(id)) } async fn search(&self, query: MemoryQuery) -> MemoryResult> { - let needle = query.text.as_deref().map(str::to_lowercase); + let terms = query + .text + .as_deref() + .map(str::to_lowercase) + .map(|text| { + text.split_whitespace() + .map(str::to_string) + .collect::>() + }) + .unwrap_or_default(); let limit = query.limit.unwrap_or(20); let mut hits = self @@ -96,18 +105,18 @@ impl MemoryStore for InMemoryMemoryStore { .is_none_or(|namespace| record.namespace == namespace) }) .filter_map(|record| { - let score = match needle.as_deref() { - Some(text) if !text.is_empty() => { - let content = record.content.to_lowercase(); - if !content.contains(text) { - return None; - } - text.split_whitespace() - .filter(|term| content.contains(term)) - .count() - .max(1) as f32 + let score = if terms.is_empty() { + 1.0 + } else { + let content = record.content.to_lowercase(); + let matched = terms + .iter() + .filter(|term| content.contains(term.as_str())) + .count(); + if matched == 0 { + return None; } - _ => 1.0, + matched as f32 / terms.len() as f32 }; Some(SearchHit { @@ -130,48 +139,5 @@ impl MemoryStore for InMemoryMemoryStore { } #[cfg(test)] -mod tests { - use super::*; - - #[tokio::test] - async fn inserts_and_retrieves_memory() { - let store = InMemoryMemoryStore::new(); - let inserted = store - .insert(MemoryInput::new("profile", "User prefers dark mode")) - .await - .expect("insert memory"); - - let fetched = store.get(inserted.id).await.expect("get memory"); - - assert_eq!(fetched.namespace, "profile"); - assert_eq!(fetched.content, "User prefers dark mode"); - } - - #[tokio::test] - async fn searches_by_namespace_and_text() { - let store = InMemoryMemoryStore::new(); - store - .insert(MemoryInput::new("profile", "User prefers dark mode")) - .await - .expect("insert profile memory"); - store - .insert(MemoryInput::new( - "project", - "TinyCortex stores durable memories", - )) - .await - .expect("insert project memory"); - - let hits = store - .search(MemoryQuery { - namespace: Some("project".to_owned()), - text: Some("durable".to_owned()), - limit: Some(10), - }) - .await - .expect("search memory"); - - assert_eq!(hits.len(), 1); - assert_eq!(hits[0].record.namespace, "project"); - } -} +#[path = "store_tests.rs"] +mod tests; diff --git a/src/memory/store/store_tests.rs b/src/memory/store/store_tests.rs new file mode 100644 index 0000000..5662413 --- /dev/null +++ b/src/memory/store/store_tests.rs @@ -0,0 +1,60 @@ +use super::*; + +#[tokio::test] +async fn inserts_and_retrieves_memory() { + let store = InMemoryMemoryStore::new(); + let inserted = store + .insert(MemoryInput::new("profile", "User prefers dark mode")) + .await + .expect("insert memory"); + + let fetched = store.get(inserted.id).await.expect("get memory"); + + assert_eq!(fetched.namespace, "profile"); + assert_eq!(fetched.content, "User prefers dark mode"); +} + +#[tokio::test] +async fn searches_by_namespace_and_text() { + let store = InMemoryMemoryStore::new(); + store + .insert(MemoryInput::new("profile", "User prefers dark mode")) + .await + .expect("insert profile memory"); + store + .insert(MemoryInput::new( + "project", + "TinyCortex stores durable memories", + )) + .await + .expect("insert project memory"); + + let hits = store + .search(MemoryQuery { + namespace: Some("project".to_owned()), + text: Some("durable".to_owned()), + limit: Some(10), + }) + .await + .expect("search memory"); + + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].record.namespace, "project"); +} + +#[tokio::test] +async fn search_matches_individual_terms_and_ranks_by_coverage() { + let store = InMemoryMemoryStore::new(); + store + .insert(MemoryInput::new("profile", "dark theme")) + .await + .unwrap(); + store + .insert(MemoryInput::new("profile", "dark mode theme")) + .await + .unwrap(); + let hits = store.search(MemoryQuery::text("dark mode")).await.unwrap(); + assert_eq!(hits.len(), 2); + assert_eq!(hits[0].score, 1.0); + assert_eq!(hits[1].score, 0.5); +} diff --git a/src/memory/store/types.rs b/src/memory/store/types.rs index 31da414..51a543a 100644 --- a/src/memory/store/types.rs +++ b/src/memory/store/types.rs @@ -3,7 +3,7 @@ //! Defines the records that flow through the store: an untrusted [`MemoryInput`] //! supplied by callers, the persisted [`MemoryRecord`] minted from it, the //! [`MemoryQuery`] filter shape, the [`SearchHit`] returned by retrieval, and -//! the [`MemoryError`] failure modes. These are pure data contracts; persistence +//! the [`StoreError`] failure modes. These are pure data contracts; persistence //! and retrieval behavior live in sibling store modules. use chrono::{DateTime, Utc}; @@ -14,12 +14,12 @@ use uuid::Uuid; /// Stable, machine-readable identity for a [`MemoryRecord`]; a v4 UUID. pub type MemoryId = Uuid; -/// Result alias for fallible store operations, carrying [`MemoryError`]. -pub type MemoryResult = Result; +/// Result alias for fallible store operations, carrying [`StoreError`]. +pub type MemoryResult = Result; /// Failure modes for memory store operations. #[derive(Debug, Error, PartialEq, Eq)] -pub enum MemoryError { +pub enum StoreError { /// No record exists for the requested [`MemoryId`]. #[error("memory record not found: {0}")] NotFound(MemoryId), @@ -76,12 +76,12 @@ impl MemoryRecord { /// Promotes a [`MemoryInput`] into a persisted record. /// /// Trims `content` and rejects empty results with - /// [`MemoryError::EmptyContent`]. Mints a fresh [`MemoryId`] and stamps both + /// [`StoreError::EmptyContent`]. Mints a fresh [`MemoryId`] and stamps both /// timestamps with the current UTC instant. pub fn from_input(input: MemoryInput) -> MemoryResult { let content = input.content.trim().to_owned(); if content.is_empty() { - return Err(MemoryError::EmptyContent); + return Err(StoreError::EmptyContent); } let now = Utc::now(); diff --git a/src/memory/store/vectors/store.rs b/src/memory/store/vectors/store.rs index d237943..4b8e8e2 100644 --- a/src/memory/store/vectors/store.rs +++ b/src/memory/store/vectors/store.rs @@ -24,7 +24,7 @@ use std::sync::Arc; use anyhow::Context; use parking_lot::Mutex; -use rusqlite::Connection; +use rusqlite::{Connection, OptionalExtension}; use super::embedding::EmbeddingBackend; @@ -32,6 +32,7 @@ use super::embedding::EmbeddingBackend; const INIT_SQL: &str = " PRAGMA journal_mode = WAL; PRAGMA synchronous = NORMAL; + PRAGMA busy_timeout = 15000; CREATE TABLE IF NOT EXISTS vectors ( id TEXT NOT NULL, @@ -143,7 +144,7 @@ impl VectorStore { [], |row| row.get(0), ) - .ok(); + .optional()?; match stored_dims { None => { @@ -160,9 +161,11 @@ impl VectorStore { } } Some(dims_str) => { - let stored: usize = dims_str.parse().unwrap_or(0); + let stored: usize = dims_str.parse().with_context(|| { + format!("invalid vector store embed_dims metadata: {dims_str}") + })?; let runtime = backend.dimensions(); - if stored != 0 && runtime != 0 && stored != runtime { + if stored != runtime { anyhow::bail!( "vector store dimension mismatch: database was created with \ {stored}-dim embeddings but the current backend ({}) uses \ @@ -207,6 +210,12 @@ impl VectorStore { embedding: &[f32], metadata: serde_json::Value, ) -> anyhow::Result<()> { + anyhow::ensure!( + embedding.len() == self.backend.dimensions(), + "embedding dimension mismatch: expected {}, got {}", + self.backend.dimensions(), + embedding.len() + ); let blob = vec_to_bytes(embedding); let meta_str = serde_json::to_string(&metadata)?; let now = now_ts(); @@ -321,17 +330,17 @@ impl VectorStore { let mut scored: Vec = rows .into_iter() .map(|(id, namespace, text, blob, meta_str)| { - let stored_vec = bytes_to_vec(&blob); + let stored_vec = bytes_to_vec(&blob)?; let score = cosine_similarity(query_vec, &stored_vec); - ScoredRow { + Ok(ScoredRow { score, id, namespace, text, meta_str, - } + }) }) - .collect(); + .collect::>>()?; // Sort descending by score. scored.sort_by(|a, b| { @@ -428,30 +437,28 @@ pub fn vec_to_bytes(v: &[f32]) -> Vec { /// Deserializes little-endian bytes back to a float vector. /// -/// Uses `chunks_exact(4)`, so any trailing 1-3 bytes that don't form a full -/// `f32` are silently dropped rather than erroring — a truncated or corrupt -/// blob yields a shorter vector instead of a decode failure. (Contrast with -/// `crate::memory::chunks::embedding_from_blob`, the chunk-store's own blob -/// decoder, which errors instead of truncating — the two decoders currently -/// have different contracts for malformed input.) -pub fn bytes_to_vec(bytes: &[u8]) -> Vec { - bytes +/// Returns an error when the blob length is not divisible by four rather than +/// silently truncating corrupt trailing bytes. +pub fn bytes_to_vec(bytes: &[u8]) -> anyhow::Result> { + anyhow::ensure!( + bytes.len().is_multiple_of(4), + "invalid embedding blob byte length: {}", + bytes.len() + ); + Ok(bytes .chunks_exact(4) .map(|chunk| { let arr: [u8; 4] = chunk.try_into().unwrap_or([0; 4]); f32::from_le_bytes(arr) }) - .collect() + .collect()) } /// Computes cosine similarity between two vectors. Returns 0.0 for /// mismatched lengths, empty vectors, or zero-magnitude vectors. /// -/// The raw cosine value (mathematically in `[-1.0, 1.0]`) is clamped to -/// `[0.0, 1.0]` before being returned, so two vectors pointing in opposite -/// directions score identically to two unrelated (orthogonal) vectors — both -/// report `0.0`. Callers that need to distinguish "opposite" from "unrelated" -/// cannot do so from this score alone. +/// The raw cosine value is clamped only for floating-point drift, preserving +/// its mathematical `[-1.0, 1.0]` range. pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 { if a.len() != b.len() || a.is_empty() { return 0.0; @@ -470,7 +477,7 @@ pub fn cosine_similarity(a: &[f32], b: &[f32]) -> f64 { if denom <= f64::EPSILON { return 0.0; } - (dot / denom).clamp(0.0, 1.0) + (dot / denom).clamp(-1.0, 1.0) } fn now_ts() -> f64 { diff --git a/src/memory/store/vectors/store_tests.rs b/src/memory/store/vectors/store_tests.rs index 5c3416b..e204151 100644 --- a/src/memory/store/vectors/store_tests.rs +++ b/src/memory/store/vectors/store_tests.rs @@ -69,17 +69,17 @@ fn roundtrip_vec_bytes() { let original = vec![1.0_f32, -2.5, 3.125, 0.0, f32::MAX, f32::MIN]; let bytes = vec_to_bytes(&original); assert_eq!(bytes.len(), original.len() * 4); - assert_eq!(original, bytes_to_vec(&bytes)); + assert_eq!(original, bytes_to_vec(&bytes).unwrap()); } #[test] fn empty_vec_roundtrip() { - assert!(bytes_to_vec(&vec_to_bytes(&[])).is_empty()); + assert!(bytes_to_vec(&vec_to_bytes(&[])).unwrap().is_empty()); } #[test] -fn bytes_to_vec_truncates_partial_bytes() { - assert_eq!(bytes_to_vec(&[0u8; 5]).len(), 1); +fn bytes_to_vec_rejects_partial_bytes() { + assert!(bytes_to_vec(&[0u8; 5]).is_err()); } // ── cosine_similarity ─────────────────────────────────── @@ -97,7 +97,7 @@ fn cosine_orthogonal() { #[test] fn cosine_opposite() { - assert!(cosine_similarity(&[1.0, 0.0], &[-1.0, 0.0]).abs() < 1e-6); + assert!((cosine_similarity(&[1.0, 0.0], &[-1.0, 0.0]) + 1.0).abs() < 1e-6); } #[test] @@ -158,6 +158,53 @@ fn open_reopen_different_dims_errors() { assert!(msg.contains("8"), "should mention runtime dims: {msg}"); } +#[test] +fn open_rejects_corrupt_dimension_metadata() { + let dir = tempfile::tempdir().unwrap(); + let db_path = dir.path().join("v.db"); + { + let store = VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 4 })).unwrap(); + store + .conn + .lock() + .execute( + "UPDATE store_meta SET value = 'not-a-number' WHERE key = 'embed_dims'", + [], + ) + .unwrap(); + } + let err = VectorStore::open(&db_path, Arc::new(FakeEmbedding { dims: 4 })) + .err() + .expect("corrupt metadata must fail closed"); + assert!(err.to_string().contains("invalid vector store embed_dims")); +} + +#[test] +fn insert_with_vector_rejects_wrong_dimensions() { + let store = fake_store(3); + let err = store + .insert_with_vector("id", "ns", "text", &[1.0, 2.0], json!({})) + .unwrap_err(); + assert!(err.to_string().contains("expected 3, got 2")); + assert_eq!(store.count(None).unwrap(), 0); +} + +#[test] +fn search_rejects_corrupt_embedding_blob() { + let store = fake_store(2); + store + .conn + .lock() + .execute( + "INSERT INTO vectors + (id, namespace, text, embedding, metadata, created_at, updated_at) + VALUES ('bad', 'ns', 'text', X'0000000000', '{}', 0, 0)", + [], + ) + .unwrap(); + assert!(store.search_by_vector("ns", &[1.0, 0.0], 5).is_err()); +} + #[test] fn backend_accessor() { let store = fake_store(3); @@ -392,115 +439,5 @@ fn search_by_vector_falls_back_to_null_on_invalid_metadata() { ); } -#[test] -fn search_by_vector_scores_correct() { - let store = fake_store(3); - store - .insert_with_vector("x", "ns", "x", &[1.0, 0.0, 0.0], json!({})) - .unwrap(); - store - .insert_with_vector("y", "ns", "y", &[0.0, 1.0, 0.0], json!({})) - .unwrap(); - let results = store.search_by_vector("ns", &[1.0, 0.0, 0.0], 2).unwrap(); - assert_eq!(results[0].id, "x"); - assert!((results[0].score - 1.0).abs() < 1e-6); - assert!(results[1].score < 1e-6); -} - -#[test] -fn search_by_vector_preserves_metadata() { - let store = fake_store(2); - store - .insert_with_vector("a", "ns", "t", &[1.0, 0.0], json!({"key": "value"})) - .unwrap(); - assert_eq!( - store.search_by_vector("ns", &[1.0, 0.0], 1).unwrap()[0].metadata["key"], - "value" - ); -} - -#[test] -fn search_handles_invalid_metadata_json() { - let store = fake_store(2); - { - let conn = store.conn.lock(); - conn.execute( - "INSERT INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) - VALUES ('bad', 'ns', 'text', ?1, 'not-json', 0.0, 0.0)", - rusqlite::params![vec_to_bytes(&[1.0, 0.0])], - ) - .unwrap(); - } - let results = store.search_by_vector("ns", &[1.0, 0.0], 1).unwrap(); - assert_eq!(results[0].id, "bad"); - assert!(results[0].metadata.is_null()); -} - -// ── delete ────────────────────────────────────────────── - -#[tokio::test] -async fn delete_existing() { - let store = fake_store(4); - store.insert("a", "ns", "text", json!({})).await.unwrap(); - assert!(store.delete("ns", "a").unwrap()); - assert_eq!(store.count(Some("ns")).unwrap(), 0); -} - -#[test] -fn delete_nonexistent() { - assert!(!fake_store(3).delete("ns", "no-such-id").unwrap()); -} - -#[tokio::test] -async fn delete_wrong_namespace() { - let store = fake_store(4); - store.insert("a", "ns1", "text", json!({})).await.unwrap(); - assert!(!store.delete("ns2", "a").unwrap()); - assert_eq!(store.count(Some("ns1")).unwrap(), 1); -} - -// ── clear_namespace ───────────────────────────────────── - -#[tokio::test] -async fn clear_namespace_removes_all() { - let store = fake_store(4); - store.insert("a", "ns", "one", json!({})).await.unwrap(); - store.insert("b", "ns", "two", json!({})).await.unwrap(); - store - .insert("c", "other", "three", json!({})) - .await - .unwrap(); - assert_eq!(store.clear_namespace("ns").unwrap(), 2); - assert_eq!(store.count(Some("ns")).unwrap(), 0); - assert_eq!(store.count(Some("other")).unwrap(), 1); -} - -#[test] -fn clear_empty_namespace() { - assert_eq!(fake_store(3).clear_namespace("empty").unwrap(), 0); -} - -// ── list_namespaces ───────────────────────────────────── - -#[tokio::test] -async fn list_namespaces_empty() { - assert!(fake_store(3).list_namespaces().unwrap().is_empty()); -} - -#[tokio::test] -async fn list_namespaces_populated() { - let store = fake_store(4); - store.insert("a", "beta", "t", json!({})).await.unwrap(); - store.insert("b", "alpha", "t", json!({})).await.unwrap(); - store.insert("c", "beta", "t", json!({})).await.unwrap(); - assert_eq!(store.list_namespaces().unwrap(), vec!["alpha", "beta"]); -} - -// ── count ─────────────────────────────────────────────── - -#[test] -fn count_empty() { - let store = fake_store(3); - assert_eq!(store.count(None).unwrap(), 0); - assert_eq!(store.count(Some("ns")).unwrap(), 0); -} +#[path = "store_tests_more.rs"] +mod more; diff --git a/src/memory/store/vectors/store_tests_more.rs b/src/memory/store/vectors/store_tests_more.rs new file mode 100644 index 0000000..2216b38 --- /dev/null +++ b/src/memory/store/vectors/store_tests_more.rs @@ -0,0 +1,114 @@ +use super::*; + +#[test] +fn search_by_vector_scores_correct() { + let store = fake_store(3); + store + .insert_with_vector("x", "ns", "x", &[1.0, 0.0, 0.0], json!({})) + .unwrap(); + store + .insert_with_vector("y", "ns", "y", &[0.0, 1.0, 0.0], json!({})) + .unwrap(); + let results = store.search_by_vector("ns", &[1.0, 0.0, 0.0], 2).unwrap(); + assert_eq!(results[0].id, "x"); + assert!((results[0].score - 1.0).abs() < 1e-6); + assert!(results[1].score < 1e-6); +} + +#[test] +fn search_by_vector_preserves_metadata() { + let store = fake_store(2); + store + .insert_with_vector("a", "ns", "t", &[1.0, 0.0], json!({"key": "value"})) + .unwrap(); + assert_eq!( + store.search_by_vector("ns", &[1.0, 0.0], 1).unwrap()[0].metadata["key"], + "value" + ); +} + +#[test] +fn search_handles_invalid_metadata_json() { + let store = fake_store(2); + { + let conn = store.conn.lock(); + conn.execute( + "INSERT INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) + VALUES ('bad', 'ns', 'text', ?1, 'not-json', 0.0, 0.0)", + rusqlite::params![vec_to_bytes(&[1.0, 0.0])], + ) + .unwrap(); + } + let results = store.search_by_vector("ns", &[1.0, 0.0], 1).unwrap(); + assert_eq!(results[0].id, "bad"); + assert!(results[0].metadata.is_null()); +} + +// ── delete ────────────────────────────────────────────── + +#[tokio::test] +async fn delete_existing() { + let store = fake_store(4); + store.insert("a", "ns", "text", json!({})).await.unwrap(); + assert!(store.delete("ns", "a").unwrap()); + assert_eq!(store.count(Some("ns")).unwrap(), 0); +} + +#[test] +fn delete_nonexistent() { + assert!(!fake_store(3).delete("ns", "no-such-id").unwrap()); +} + +#[tokio::test] +async fn delete_wrong_namespace() { + let store = fake_store(4); + store.insert("a", "ns1", "text", json!({})).await.unwrap(); + assert!(!store.delete("ns2", "a").unwrap()); + assert_eq!(store.count(Some("ns1")).unwrap(), 1); +} + +// ── clear_namespace ───────────────────────────────────── + +#[tokio::test] +async fn clear_namespace_removes_all() { + let store = fake_store(4); + store.insert("a", "ns", "one", json!({})).await.unwrap(); + store.insert("b", "ns", "two", json!({})).await.unwrap(); + store + .insert("c", "other", "three", json!({})) + .await + .unwrap(); + assert_eq!(store.clear_namespace("ns").unwrap(), 2); + assert_eq!(store.count(Some("ns")).unwrap(), 0); + assert_eq!(store.count(Some("other")).unwrap(), 1); +} + +#[test] +fn clear_empty_namespace() { + assert_eq!(fake_store(3).clear_namespace("empty").unwrap(), 0); +} + +// ── list_namespaces ───────────────────────────────────── + +#[tokio::test] +async fn list_namespaces_empty() { + assert!(fake_store(3).list_namespaces().unwrap().is_empty()); +} + +#[tokio::test] +async fn list_namespaces_populated() { + let store = fake_store(4); + store.insert("a", "beta", "t", json!({})).await.unwrap(); + store.insert("b", "alpha", "t", json!({})).await.unwrap(); + store.insert("c", "beta", "t", json!({})).await.unwrap(); + assert_eq!(store.list_namespaces().unwrap(), vec!["alpha", "beta"]); +} + +// ── count ─────────────────────────────────────────────── + +#[test] +fn count_empty() { + let store = fake_store(3); + assert_eq!(store.count(None).unwrap(), 0); + assert_eq!(store.count(Some("ns")).unwrap(), 0); +} diff --git a/src/memory/tool_memory/render.rs b/src/memory/tool_memory/render.rs index 99004eb..a06a108 100644 --- a/src/memory/tool_memory/render.rs +++ b/src/memory/tool_memory/render.rs @@ -76,37 +76,22 @@ impl ToolMemoryRulesSection { /// Pure rendering helper — public so callers that pre-render the block /// (e.g. tests, dynamic prompt sources) can share the same logic. /// -/// NOTE: `rule.rule` and `rule.tool_name` are concatenated into the output -/// verbatim (only `trim()`ed) — neither is escaped. A rule body containing -/// its own `\n### \`tool\`` heading-shaped text renders as a fake extra -/// tool section inside a block the prompt frames as a hard constraint, and -/// a backtick in `tool_name` breaks out of the `` ` `` code span in the -/// `### \`tool_name\`` heading. Callers that accept rule content from -/// untrusted sources (e.g. auto-captured tool-failure text) must sanitize -/// before storing, since this function does not. -/// -/// NOTE: the sort above orders by priority first, tool name second, so a -/// tool with rules at two different priorities (e.g. one Critical, one -/// High) is split across the Critical block and the High block by -/// construction. Heading emission tracks only the *immediately previous* -/// rule's `tool_name`, so that tool gets two separate `### \`tool\`` -/// headings — one per block — instead of one grouped section. This is a -/// real, reachable rendering artifact whenever a tool has both a Critical -/// and a High rule. +/// Tool names and bodies are normalized to single-line prompt text. This keeps +/// stored newlines/backticks from forging headings or escaping code spans. pub fn render_tool_memory_rules(rules: &[ToolMemoryRule]) -> String { if rules.is_empty() { return String::new(); } - // Stable order: Critical first, then High; within a priority, by tool - // name, then by rule body, then by id. Callers may pass an + // Stable order: group by normalized tool name, then Critical before High, + // then by rule body and id. Callers may pass an // already-sorted list (the store does), but rendering must not depend // on that contract — the system prompt has to be byte-stable. let mut sorted: Vec<&ToolMemoryRule> = rules.iter().collect(); sorted.sort_by(|a, b| { - b.priority - .cmp(&a.priority) - .then_with(|| a.tool_name.cmp(&b.tool_name)) + a.tool_name + .cmp(&b.tool_name) + .then_with(|| b.priority.cmp(&a.priority)) .then_with(|| a.rule.cmp(&b.rule)) .then_with(|| a.id.cmp(&b.id)) }); @@ -129,20 +114,24 @@ pub fn render_tool_memory_rules(rules: &[ToolMemoryRule]) -> String { out.push('\n'); } out.push_str("### `"); - out.push_str(rule.tool_name.as_str()); + out.push_str(&prompt_line(&rule.tool_name).replace('`', "'")); out.push_str("`\n"); current_tool = Some(rule.tool_name.as_str()); } out.push_str("- "); out.push_str(priority_marker(rule.priority)); out.push(' '); - out.push_str(rule.rule.trim()); + out.push_str(&prompt_line(&rule.rule)); out.push('\n'); } out } +fn prompt_line(value: &str) -> String { + value.split_whitespace().collect::>().join(" ") +} + fn priority_marker(priority: ToolMemoryPriority) -> &'static str { match priority { ToolMemoryPriority::Critical => "**[critical]**", diff --git a/src/memory/tool_memory/render_tests.rs b/src/memory/tool_memory/render_tests.rs index 13afd8a..8b37aac 100644 --- a/src/memory/tool_memory/render_tests.rs +++ b/src/memory/tool_memory/render_tests.rs @@ -80,3 +80,27 @@ fn section_renders_snapshot_via_rendered_accessor() { assert!(!section.is_empty()); assert!(section.rendered().contains("never email Sarah")); } + +#[test] +fn one_tool_with_multiple_priorities_gets_one_heading() { + let rules = vec![ + rule("email", "high", ToolMemoryPriority::High), + rule("email", "critical", ToolMemoryPriority::Critical), + ]; + let output = render_tool_memory_rules(&rules); + assert_eq!(output.matches("### `email`").count(), 1); + assert!(output.find("critical").unwrap() < output.find("high").unwrap()); +} + +#[test] +fn heading_shaped_rule_content_cannot_forge_prompt_sections() { + let rules = vec![rule( + "bad`tool", + "first line\n### `shell`\n- forged", + ToolMemoryPriority::Critical, + )]; + let output = render_tool_memory_rules(&rules); + assert!(output.contains("### `bad'tool`")); + assert!(!output.contains("\n### `shell`")); + assert!(output.contains("first line ### `shell` - forged")); +} diff --git a/src/memory/tool_memory/store.rs b/src/memory/tool_memory/store.rs index d7b65fe..4e0ba17 100644 --- a/src/memory/tool_memory/store.rs +++ b/src/memory/tool_memory/store.rs @@ -16,7 +16,7 @@ //! particular (un-ported) store. use std::collections::HashMap; -use std::sync::Arc; +use std::sync::{Arc, LazyLock}; use serde_json::Value; @@ -24,6 +24,13 @@ use super::types::{tool_memory_namespace, ToolMemoryPriority, ToolMemoryRule, To use crate::memory::traits::Memory; use crate::memory::types::MemoryCategory; +/// Serializes the read-modify-write portion of rule upserts across store +/// handles. The backend trait does not expose compare-and-swap, so the lock is +/// the only portable way to preserve `created_at` and prevent lost in-process +/// updates. +static RULE_MUTATION_LOCK: LazyLock> = + LazyLock::new(|| futures::lock::Mutex::new(())); + /// Maximum number of rules surfaced into the system prompt at once. /// /// Keeps the cache-friendly prefix bounded even when callers stash a long @@ -79,11 +86,9 @@ impl ToolMemoryStore { /// per-tool heading. Normalize `tool_name` before calling if case /// consistency matters to the caller. /// - /// NOTE: this is a racy read (`fetch_rule`) → write (`Memory::store`) - /// with no mutation lock (unlike the process-wide lock the `goals` - /// module uses around its own load→mutate→save sequences). Concurrent - /// upserts of the same `(tool_name, id)` can interleave and lose an - /// update. + /// The read (`fetch_rule`) → write (`Memory::store`) transaction is + /// serialized across all in-process store handles because the generic + /// backend contract does not expose compare-and-swap. pub async fn put_rule(&self, mut rule: ToolMemoryRule) -> Result { if rule.tool_name.trim().is_empty() { return Err("tool_name is required".to_string()); @@ -94,6 +99,9 @@ impl ToolMemoryStore { if rule.id.trim().is_empty() { rule.id = ToolMemoryRule::generate_id(); } + rule.tool_name = rule.tool_name.trim().to_lowercase(); + + let _guard = RULE_MUTATION_LOCK.lock().await; let namespace = tool_memory_namespace(&rule.tool_name); let key = ToolMemoryRule::storage_key(&rule.id); @@ -210,7 +218,11 @@ impl ToolMemoryStore { .cmp(&a.priority) .then_with(|| b.updated_at.cmp(&a.updated_at)) }); - collected.truncate(TOOL_MEMORY_PROMPT_CAP); + let critical_count = collected + .iter() + .take_while(|rule| rule.priority == ToolMemoryPriority::Critical) + .count(); + collected.truncate(TOOL_MEMORY_PROMPT_CAP.max(critical_count)); let mut out: HashMap> = HashMap::new(); for rule in collected { diff --git a/src/memory/tool_memory/store_tests.rs b/src/memory/tool_memory/store_tests.rs index 0e8f812..24b1318 100644 --- a/src/memory/tool_memory/store_tests.rs +++ b/src/memory/tool_memory/store_tests.rs @@ -81,6 +81,23 @@ async fn put_rule_preserves_created_at_on_upsert() { assert_eq!(updated.rule, "updated rule body"); } +#[tokio::test] +async fn put_rule_normalizes_tool_name_with_namespace() { + let store = fresh_store(); + let stored = store + .record( + " Email ", + "normalized", + ToolMemoryPriority::Critical, + ToolMemorySource::Programmatic, + vec![], + ) + .await + .unwrap(); + assert_eq!(stored.tool_name, "email"); + assert_eq!(store.list_rules("email").await.unwrap().len(), 1); +} + #[tokio::test] async fn list_rules_sorts_critical_first_then_freshest() { let store = fresh_store(); @@ -303,6 +320,28 @@ async fn rules_for_prompt_caps_results() { assert_eq!(count, TOOL_MEMORY_PROMPT_CAP); } +#[tokio::test] +async fn rules_for_prompt_never_truncates_critical_rules() { + let store = fresh_store(); + for idx in 0..(TOOL_MEMORY_PROMPT_CAP + 5) { + store + .record( + "shell", + &format!("critical rule {idx}"), + ToolMemoryPriority::Critical, + ToolMemorySource::UserExplicit, + vec![], + ) + .await + .unwrap(); + } + let rendered = store + .rules_for_prompt(&["shell".to_string()]) + .await + .unwrap(); + assert_eq!(rendered["shell"].len(), TOOL_MEMORY_PROMPT_CAP + 5); +} + #[tokio::test] async fn list_rules_skips_malformed_entries() { let memory: Arc = Arc::new(MockMemory::default()); diff --git a/src/memory/tool_memory/types.rs b/src/memory/tool_memory/types.rs index 0079e2d..095d2e5 100644 --- a/src/memory/tool_memory/types.rs +++ b/src/memory/tool_memory/types.rs @@ -150,12 +150,8 @@ impl ToolMemoryRule { /// reason about the namespace without ambiguity. Always build the /// namespace through this helper — never hard-code the `tool-` format. /// -/// NOTE: only the namespace is normalised here. [`ToolMemoryRule::tool_name`] -/// itself is stored verbatim by [`ToolMemoryStore::put_rule`], so -/// `tool_name = "Email"` and `tool_name = "email"` land in the same -/// namespace but are still distinct strings for grouping/display purposes -/// downstream (list output, prompt-section headings). Normalize the input -/// before constructing a [`ToolMemoryRule`] if case consistency matters. +/// [`ToolMemoryStore::put_rule`] applies the same normalization to the stored +/// rule so namespace and display/grouping identity cannot diverge. /// /// [`ToolMemoryStore::put_rule`]: super::store::ToolMemoryStore::put_rule pub fn tool_memory_namespace(tool_name: &str) -> String { From 7da649cbb43b91fc23ce34a716dd3b1388305dd0 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 07:17:09 +0000 Subject: [PATCH 2/7] fix: make retrieval and config contracts operative --- src/lib.rs | 9 +- src/memory/config.rs | 195 ++++++++++++++-- src/memory/config/policy.rs | 166 +++++++++++++ src/memory/config_tests.rs | 120 +++++++++- src/memory/graph/query.rs | 24 +- src/memory/graph/types.rs | 13 ++ src/memory/mod.rs | 18 +- src/memory/retrieval/cover.rs | 58 ++--- src/memory/retrieval/cover_tests.rs | 26 +++ src/memory/retrieval/drill_down.rs | 18 +- src/memory/retrieval/drill_down_tests.rs | 56 +++++ src/memory/retrieval/fast.rs | 97 +------- src/memory/retrieval/fast_tests.rs | 226 ++++++++++++++++++ src/memory/retrieval/fetch.rs | 11 +- src/memory/retrieval/global.rs | 91 ++++---- src/memory/retrieval/global_tests.rs | 102 +++++++- src/memory/retrieval/graph_adapter.rs | 68 ++++-- src/memory/retrieval/graph_adapter_tests.rs | 41 +++- src/memory/retrieval/mod.rs | 2 +- src/memory/retrieval/rerank.rs | 36 +-- src/memory/retrieval/rerank_tests.rs | 75 ++++++ src/memory/retrieval/search.rs | 40 ++-- src/memory/retrieval/search_tests.rs | 47 +++- src/memory/retrieval/source.rs | 61 ++--- src/memory/retrieval/source_tests.rs | 2 +- src/memory/score/embed.rs | 8 +- src/memory/score/extract/llm_tests.rs | 243 +------------------- src/memory/score/extract/llm_tests_more.rs | 243 ++++++++++++++++++++ src/memory/score/extract/mod.rs | 4 +- src/memory/score/extract/regex.rs | 11 +- src/memory/score/extract/regex_tests.rs | 12 + src/memory/score/mod.rs | 81 +++---- src/memory/score/mod_tests.rs | 26 +++ src/memory/score/signals/mod.rs | 2 +- src/memory/score/signals/types.rs | 3 +- src/memory/score/store.rs | 120 +++------- src/memory/score/store_query.rs | 92 ++++++++ src/memory/score/store_tests.rs | 85 +++++++ src/memory/types.rs | 75 ++++-- src/memory/types_tests.rs | 36 ++- 40 files changed, 1922 insertions(+), 721 deletions(-) create mode 100644 src/memory/config/policy.rs create mode 100644 src/memory/retrieval/fast_tests.rs create mode 100644 src/memory/retrieval/rerank_tests.rs create mode 100644 src/memory/score/extract/llm_tests_more.rs create mode 100644 src/memory/score/store_query.rs diff --git a/src/lib.rs b/src/lib.rs index 16ad507..e814a1e 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,14 +19,11 @@ //! (`memory::queue::runtime`) instead of driving `run_once` by hand. //! - `git-diff`: compiles in `memory::diff` (git-backed source snapshots, //! diffs, checkpoints, read markers) and its native `git2`/libgit2 dependency. -//! - `providers-http`: compiles in `memory::providers` (reqwest-based -//! embedding / LLM HTTP providers). Implies `tokio`. -//! - `rpc`: compiles in `memory::rpc` (the serde wire-envelope surface for -//! exposing the engine over an RPC boundary). +//! - `sync`: compiles in the live Composio and workspace synchronization +//! pipelines and implies `tokio`. //! //! With every feature off, `cargo check` / `cargo test` still exercise the full //! synchronous engine (storage, ingest, retrieval, tree, graph, goals, …); the -//! feature-gated modules only reserve a seam until their concrete -//! implementations land. +//! feature-gated modules contain concrete optional implementations. pub mod memory; diff --git a/src/memory/config.rs b/src/memory/config.rs index a3c80c9..b46134a 100644 --- a/src/memory/config.rs +++ b/src/memory/config.rs @@ -7,12 +7,18 @@ use std::path::PathBuf; +use anyhow::Context; use serde::{Deserialize, Serialize}; +mod policy; +pub use policy::{QueueConfig, RetrievalLimits, ScoringPolicyConfig}; + /// OpenHuman tree-summarisation input budget (tokens). pub const INPUT_TOKEN_BUDGET: u32 = 50_000; /// OpenHuman tree-summarisation output budget (tokens). pub const OUTPUT_TOKEN_BUDGET: u32 = 5_000; +/// Prompt/instruction headroom reserved inside the summarisation input budget. +pub const SUMMARY_OVERHEAD_RESERVE_TOKENS: u32 = 2_048; /// Number of summary siblings before a bucket seals. pub const SUMMARY_FANOUT: u32 = 10; /// Default flush age for stale buffers (7 days, in seconds). @@ -28,12 +34,11 @@ pub const FOLDER_FILE_SIZE_CAP_BYTES: u64 = 10 * 1024 * 1024; /// Top-level configuration for a memory engine instance. /// -/// This struct performs no validation or filesystem I/O on its own — building -/// one (via [`Self::new`], `Default::default()` on the nested configs, or -/// deserializing from JSON/TOML) never touches disk or fails. Path sandboxing -/// (rejecting traversal / symlink escapes out of `workspace`) and any other -/// invariant enforcement happen where a config is actually consumed (see -/// `memory::sources::validation`), not here. +/// Direct construction and serde deserialization do not touch disk; call +/// [`Self::validate`] after constructing manually. [`Self::from_toml_file`] +/// is the safe file-loading entry point and validates before returning. Path +/// sandboxing for individual source paths remains enforced where those paths +/// are consumed (see `memory::sources::validation`). #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemoryConfig { /// Workspace root. Markdown content, SQLite indexes, and ledgers live under @@ -54,6 +59,16 @@ pub struct MemoryConfig { /// Default hybrid retrieval weighting. #[serde(default)] pub retrieval: RetrievalConfig, + /// Admission-scoring policy. Runtime extractor implementations remain + /// injected separately and are combined with this serializable policy. + #[serde(default)] + pub scoring: ScoringPolicyConfig, + /// Deterministic document-extraction policy. + #[serde(default)] + pub ingestion: crate::memory::ingest::MemoryIngestionConfig, + /// Queue locking, retry, and concurrency policy. + #[serde(default)] + pub queue: QueueConfig, /// Per-source sync budget ceilings (enforced when a host invokes ingest). #[serde(default)] pub sync_budget: SyncBudgetConfig, @@ -82,14 +97,144 @@ impl MemoryConfig { embedding: EmbeddingConfig::default(), tree: TreeConfig::default(), retrieval: RetrievalConfig::default(), + scoring: ScoringPolicyConfig::default(), + ingestion: crate::memory::ingest::MemoryIngestionConfig::default(), + queue: QueueConfig::default(), sync_budget: SyncBudgetConfig::default(), sync: SyncConfig::default(), } } + + /// Load a TOML configuration file and validate its runtime invariants. + /// + /// Relative workspace paths remain relative to the host's current working + /// directory; the loader does not silently reinterpret them relative to the + /// config file. + pub fn from_toml_file(path: impl AsRef) -> anyhow::Result { + let path = path.as_ref(); + let text = std::fs::read_to_string(path) + .map_err(anyhow::Error::from) + .with_context(|| format!("failed to read memory config {}", path.display()))?; + let config: Self = toml::from_str(&text) + .map_err(anyhow::Error::from) + .with_context(|| format!("failed to parse memory config {}", path.display()))?; + config.validate()?; + Ok(config) + } + + /// Validate configuration values that would otherwise produce degenerate + /// stores, budgets, or ranking behavior. + pub fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + !self.workspace.as_os_str().is_empty(), + "workspace must not be empty" + ); + self.scoring.validate()?; + self.queue.validate()?; + self.retrieval.limits.validate()?; + anyhow::ensure!( + (0.0..=1.0).contains(&self.ingestion.entity_threshold) + && (0.0..=1.0).contains(&self.ingestion.relation_threshold) + && (0.0..=1.0).contains(&self.ingestion.adjacency_threshold), + "ingestion thresholds must be between zero and one" + ); + anyhow::ensure!( + self.ingestion.batch_size > 0, + "ingestion.batch_size must be positive" + ); + anyhow::ensure!( + self.embedding.dim > 0, + "embedding.dim must be greater than zero" + ); + anyhow::ensure!( + self.tree.input_token_budget > 0, + "tree.input_token_budget must be greater than zero" + ); + anyhow::ensure!( + self.tree.output_token_budget > 0, + "tree.output_token_budget must be greater than zero" + ); + anyhow::ensure!( + self.tree.summary_overhead_reserve_tokens < self.tree.input_token_budget, + "tree.summary_overhead_reserve_tokens must be smaller than tree.input_token_budget" + ); + anyhow::ensure!( + self.tree.summary_fanout > 0, + "tree.summary_fanout must be greater than zero" + ); + anyhow::ensure!( + self.tree.flush_age_secs > 0, + "tree.flush_age_secs must be greater than zero" + ); + anyhow::ensure!( + self.tree.flavour_root_token_budget > 0, + "tree.flavour_root_token_budget must be greater than zero" + ); + let weights = [ + self.retrieval.default_profile.graph, + self.retrieval.default_profile.vector, + self.retrieval.default_profile.keyword, + self.retrieval.default_profile.freshness, + ]; + anyhow::ensure!( + weights + .iter() + .all(|weight| weight.is_finite() && *weight >= 0.0), + "retrieval weights must be finite and non-negative" + ); + anyhow::ensure!( + weights.iter().sum::() > 0.0, + "at least one retrieval weight must be positive" + ); + validate_budget("sync_budget", &self.sync_budget)?; + validate_budget("sync.budget", &self.sync.budget)?; + if let Some(composio) = &self.sync.composio { + let base_url = composio.base_url.trim(); + anyhow::ensure!( + base_url.starts_with("https://") || base_url.starts_with("http://"), + "sync.composio.base_url must be an http(s) URL" + ); + anyhow::ensure!( + composio.api_key.as_ref().is_none_or(|key| !key.is_empty()), + "sync.composio.api_key must not be blank" + ); + anyhow::ensure!( + composio + .bearer_token + .as_ref() + .is_none_or(|token| !token.is_empty()), + "sync.composio.bearer_token must not be blank" + ); + } + Ok(()) + } +} + +fn validate_budget(name: &str, budget: &SyncBudgetConfig) -> anyhow::Result<()> { + anyhow::ensure!( + budget.max_items.is_none_or(|value| value > 0), + "{name}.max_items must be greater than zero" + ); + anyhow::ensure!( + budget.max_tokens_per_sync.is_none_or(|value| value > 0), + "{name}.max_tokens_per_sync must be greater than zero" + ); + anyhow::ensure!( + budget + .max_cost_per_sync_usd + .is_none_or(|value| value.is_finite() && value >= 0.0), + "{name}.max_cost_per_sync_usd must be finite and non-negative" + ); + anyhow::ensure!( + budget.sync_depth_days.is_none_or(|value| value > 0), + "{name}.sync_depth_days must be greater than zero" + ); + Ok(()) } /// Embedding backend configuration. #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] pub struct EmbeddingConfig { /// Vector dimension. OpenHuman fixes this at 768. pub dim: usize, @@ -112,11 +257,16 @@ impl Default for EmbeddingConfig { /// Summary-tree budgets and sealing behavior. #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] pub struct TreeConfig { /// Max input tokens fed into one summarisation pass (see [`INPUT_TOKEN_BUDGET`]). pub input_token_budget: u32, /// Max tokens a summary may emit (see [`OUTPUT_TOKEN_BUDGET`]). pub output_token_budget: u32, + /// Tokens reserved for provider instructions and formatting rather than + /// source inputs. + #[serde(default = "default_summary_overhead_reserve_tokens")] + pub summary_overhead_reserve_tokens: u32, /// Number of summary siblings accumulated before a bucket seals into a parent /// (see [`SUMMARY_FANOUT`]). pub summary_fanout: u32, @@ -137,11 +287,16 @@ fn default_flavour_root_token_budget() -> u32 { FLAVOUR_ROOT_TOKEN_BUDGET } +fn default_summary_overhead_reserve_tokens() -> u32 { + SUMMARY_OVERHEAD_RESERVE_TOKENS +} + impl Default for TreeConfig { fn default() -> Self { Self { input_token_budget: INPUT_TOKEN_BUDGET, output_token_budget: OUTPUT_TOKEN_BUDGET, + summary_overhead_reserve_tokens: SUMMARY_OVERHEAD_RESERVE_TOKENS, summary_fanout: SUMMARY_FANOUT, flush_age_secs: DEFAULT_FLUSH_AGE_SECS, flavour_root_token_budget: FLAVOUR_ROOT_TOKEN_BUDGET, @@ -202,38 +357,42 @@ impl WeightProfile { freshness: 0.0, }; - /// Resolve a profile by its wire name. Unknown names fall back to balanced. + /// Resolve a profile by its wire name, returning `None` for unknown names. /// /// # Examples /// /// ``` /// use tinycortex::memory::config::WeightProfile; /// - /// assert_eq!(WeightProfile::by_name("semantic"), WeightProfile::SEMANTIC); - /// // Unrecognised names fail closed to the balanced default, never an error. - /// assert_eq!(WeightProfile::by_name("nonexistent"), WeightProfile::BALANCED); + /// assert_eq!(WeightProfile::by_name("semantic"), Some(WeightProfile::SEMANTIC)); + /// assert_eq!(WeightProfile::by_name("nonexistent"), None); /// ``` - pub fn by_name(name: &str) -> Self { + pub fn by_name(name: &str) -> Option { match name { - "semantic" => Self::SEMANTIC, - "lexical" => Self::LEXICAL, - "graph_first" => Self::GRAPH_FIRST, - _ => Self::BALANCED, + "balanced" => Some(Self::BALANCED), + "semantic" => Some(Self::SEMANTIC), + "lexical" => Some(Self::LEXICAL), + "graph_first" => Some(Self::GRAPH_FIRST), + _ => None, } } } /// Default retrieval configuration. #[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] pub struct RetrievalConfig { /// Default weight profile applied when a query does not specify one. pub default_profile: WeightProfile, + /// Operative result, candidate, graph, and paging bounds. + pub limits: RetrievalLimits, } impl Default for RetrievalConfig { fn default() -> Self { Self { default_profile: WeightProfile::BALANCED, + limits: RetrievalLimits::default(), } } } @@ -247,6 +406,7 @@ impl Default for RetrievalConfig { /// fallback a host applies when a source has not set its own override, rather /// than an already-enforced global cap. #[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] pub struct SyncBudgetConfig { /// Maximum records persisted by one sync tick. pub max_items: Option, @@ -260,6 +420,7 @@ pub struct SyncBudgetConfig { /// Live synchronization configuration. #[derive(Debug, Clone, Default, Serialize, Deserialize)] +#[serde(default)] pub struct SyncConfig { /// Request and ingest ceilings. #[serde(default)] @@ -267,8 +428,8 @@ pub struct SyncConfig { /// Composio transport configuration; absent disables Composio sync. #[serde(default, skip_serializing_if = "Option::is_none")] pub composio: Option, - /// Global periodic cadence. `Some(0)` means manual-only; shorter non-zero - /// values are clamped to the 24-hour default. + /// Global periodic cadence. `Some(0)` means manual-only; every other + /// configured value is honoured exactly. #[serde(default, skip_serializing_if = "Option::is_none")] pub interval_secs: Option, } diff --git a/src/memory/config/policy.rs b/src/memory/config/policy.rs new file mode 100644 index 0000000..cd60574 --- /dev/null +++ b/src/memory/config/policy.rs @@ -0,0 +1,166 @@ +//! Declarative scoring, retrieval, and queue policies. + +use serde::{Deserialize, Serialize}; + +/// Runtime bounds shared by retrieval entry points. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct RetrievalLimits { + /// Default number of returned hits. + pub default_limit: usize, + /// Default number of entity-search matches. + pub search_default_limit: usize, + /// Default number of structural time-window cover nodes. + pub cover_default_limit: usize, + /// Hard cap on caller-supplied result limits. + pub max_limit: usize, + /// Candidate occurrences loaded for graph routing. + pub occurrence_lookup_limit: usize, + /// Maximum graph traversal depth. + pub max_graph_hops: u32, + /// Default graph traversal depth. + pub default_graph_hops: u32, + /// Entity-index candidates considered by topic retrieval. + pub topic_lookup_limit: usize, + /// Page size for complete time-window chunk scans. + pub window_chunk_page_size: usize, + /// Maximum raw leaves hydrated in one public fetch. + pub fetch_batch_limit: usize, + /// Freshness signal half-life in days. + pub freshness_half_life_days: f64, +} + +impl Default for RetrievalLimits { + fn default() -> Self { + Self { + default_limit: 10, + search_default_limit: 5, + cover_default_limit: 200, + max_limit: 100, + occurrence_lookup_limit: 500, + max_graph_hops: 4, + default_graph_hops: 2, + topic_lookup_limit: 200, + window_chunk_page_size: 5_000, + fetch_batch_limit: 20, + freshness_half_life_days: 7.0, + } + } +} + +impl RetrievalLimits { + pub(super) fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.default_limit > 0 && self.search_default_limit > 0 && self.cover_default_limit > 0, + "retrieval default limits must be positive" + ); + anyhow::ensure!( + self.max_limit >= self.default_limit, + "retrieval.limits.max_limit must cover the default" + ); + anyhow::ensure!( + self.occurrence_lookup_limit > 0 + && self.topic_lookup_limit > 0 + && self.window_chunk_page_size > 0 + && self.fetch_batch_limit > 0, + "retrieval candidate and page limits must be positive" + ); + anyhow::ensure!( + self.default_graph_hops > 0 && self.default_graph_hops <= self.max_graph_hops, + "retrieval graph hop bounds are invalid" + ); + anyhow::ensure!( + self.freshness_half_life_days.is_finite() && self.freshness_half_life_days > 0.0, + "retrieval freshness half-life must be positive" + ); + Ok(()) + } +} + +/// Serializable scoring policy; extractor implementations are host-injected. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct ScoringPolicyConfig { + /// Per-signal weights. + pub weights: crate::memory::score::signals::SignalWeights, + /// Final admission threshold. + pub drop_threshold: f32, + /// Cheap-score threshold for unconditional admission. + pub definite_keep_threshold: f32, + /// Cheap-score threshold for unconditional rejection. + pub definite_drop_threshold: f32, +} + +impl Default for ScoringPolicyConfig { + fn default() -> Self { + Self { + weights: crate::memory::score::signals::SignalWeights::default(), + drop_threshold: crate::memory::score::DEFAULT_DROP_THRESHOLD, + definite_keep_threshold: crate::memory::score::DEFAULT_DEFINITE_KEEP, + definite_drop_threshold: crate::memory::score::DEFAULT_DEFINITE_DROP, + } + } +} + +impl ScoringPolicyConfig { + pub(super) fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.drop_threshold.is_finite() && (0.0..=1.0).contains(&self.drop_threshold), + "scoring.drop_threshold must be between zero and one" + ); + anyhow::ensure!( + self.definite_drop_threshold.is_finite() + && self.definite_keep_threshold.is_finite() + && self.definite_drop_threshold <= self.definite_keep_threshold, + "scoring definite thresholds are invalid" + ); + Ok(()) + } +} + +/// Queue locking, retry backoff, and LLM concurrency policy. +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(default)] +pub struct QueueConfig { + /// Lease duration for claimed jobs. + pub lock_duration_ms: i64, + /// Initial retry delay. + pub retry_base_ms: i64, + /// Maximum retry delay. + pub retry_cap_ms: i64, + /// Default number of attempts per job. + pub max_attempts: u32, + /// Concurrent LLM jobs permitted per process. + pub llm_permits: usize, + /// Maximum age of a repeatedly deferred job. + pub max_defer_age_ms: i64, +} + +impl Default for QueueConfig { + fn default() -> Self { + Self { + lock_duration_ms: 5 * 60 * 1_000, + retry_base_ms: 60 * 1_000, + retry_cap_ms: 60 * 60 * 1_000, + max_attempts: 5, + llm_permits: 1, + max_defer_age_ms: 7 * 24 * 60 * 60 * 1_000, + } + } +} + +impl QueueConfig { + pub(super) fn validate(&self) -> anyhow::Result<()> { + anyhow::ensure!( + self.lock_duration_ms > 0 + && self.retry_base_ms > 0 + && self.retry_cap_ms >= self.retry_base_ms, + "queue timing policy is invalid" + ); + anyhow::ensure!( + self.max_attempts > 0 && self.llm_permits > 0 && self.max_defer_age_ms > 0, + "queue limits must be positive" + ); + Ok(()) + } +} diff --git a/src/memory/config_tests.rs b/src/memory/config_tests.rs index 85f78e1..fd4480b 100644 --- a/src/memory/config_tests.rs +++ b/src/memory/config_tests.rs @@ -35,15 +35,24 @@ fn default_config_uses_openhuman_constants() { #[test] fn weight_profiles_match_spec() { - assert_eq!(WeightProfile::by_name("balanced"), WeightProfile::BALANCED); - assert_eq!(WeightProfile::by_name("semantic"), WeightProfile::SEMANTIC); - assert_eq!(WeightProfile::by_name("lexical"), WeightProfile::LEXICAL); + assert_eq!( + WeightProfile::by_name("balanced"), + Some(WeightProfile::BALANCED) + ); + assert_eq!( + WeightProfile::by_name("semantic"), + Some(WeightProfile::SEMANTIC) + ); + assert_eq!( + WeightProfile::by_name("lexical"), + Some(WeightProfile::LEXICAL) + ); assert_eq!( WeightProfile::by_name("graph_first"), - WeightProfile::GRAPH_FIRST + Some(WeightProfile::GRAPH_FIRST) ); - // Unknown falls back to balanced. - assert_eq!(WeightProfile::by_name("nope"), WeightProfile::BALANCED); + // Unknown names are rejected rather than silently changing behavior. + assert_eq!(WeightProfile::by_name("nope"), None); let b = WeightProfile::BALANCED; assert!((b.graph + b.vector + b.keyword + b.freshness - 1.0).abs() < 1e-9); @@ -64,5 +73,104 @@ fn config_fills_defaults_for_absent_sections() { let parsed: MemoryConfig = serde_json::from_str(r#"{"workspace":"/tmp/ws"}"#).unwrap(); assert_eq!(parsed.embedding.dim, 768); assert_eq!(parsed.tree.summary_fanout, 10); + assert_eq!(parsed.retrieval.limits.search_default_limit, 5); + assert_eq!(parsed.queue.max_attempts, 5); + assert_eq!(parsed.ingestion.batch_size, 16); + assert_eq!( + parsed.scoring.drop_threshold, + crate::memory::score::DEFAULT_DROP_THRESHOLD + ); assert!(parsed.sync_budget.max_tokens_per_sync.is_none()); } + +#[test] +fn config_validation_rejects_degenerate_values() { + let mut config = MemoryConfig::new("/tmp/ws"); + assert!(config.validate().is_ok()); + + config.embedding.dim = 0; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("embedding.dim")); + config.embedding.dim = 3; + config.tree.summary_overhead_reserve_tokens = config.tree.input_token_budget; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("summary_overhead")); + config.tree.summary_overhead_reserve_tokens = 10; + config.retrieval.default_profile.graph = -1.0; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("weights")); + config.retrieval.default_profile = WeightProfile { + graph: 0.0, + vector: 0.0, + keyword: 0.0, + freshness: 0.0, + }; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("at least one")); + + config.retrieval.default_profile = WeightProfile::BALANCED; + config.queue.retry_cap_ms = config.queue.retry_base_ms - 1; + assert!(config.validate().unwrap_err().to_string().contains("queue")); + config.queue = QueueConfig::default(); + config.retrieval.limits.window_chunk_page_size = 0; + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("retrieval")); +} + +#[test] +fn from_toml_file_applies_partial_defaults_and_validates() { + let temp = tempfile::tempdir().unwrap(); + let path = temp.path().join("memory.toml"); + std::fs::write( + &path, + "workspace = '/tmp/ws'\n[embedding]\nmodel = 'custom'\n", + ) + .unwrap(); + + let config = MemoryConfig::from_toml_file(&path).unwrap(); + + assert_eq!(config.embedding.model, "custom"); + assert_eq!(config.embedding.dim, DEFAULT_EMBEDDING_DIM); + + std::fs::write(&path, "workspace = '/tmp/ws'\n[embedding]\ndim = 0\n").unwrap(); + assert!(MemoryConfig::from_toml_file(&path).is_err()); +} + +#[test] +fn composio_validation_rejects_bad_url_and_blank_injected_secret() { + let mut config = MemoryConfig::new("/tmp/ws"); + config.sync.composio = Some(ComposioSyncConfig { + mode: ComposioMode::Direct, + base_url: "backend.example".into(), + api_key: None, + bearer_token: None, + entity_id: None, + }); + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("base_url")); + config.sync.composio.as_mut().unwrap().base_url = "https://backend.example".into(); + config.sync.composio.as_mut().unwrap().api_key = Some(SecretString::new(" ")); + assert!(config + .validate() + .unwrap_err() + .to_string() + .contains("api_key")); +} diff --git a/src/memory/graph/query.rs b/src/memory/graph/query.rs index b7dae16..bcf5afb 100644 --- a/src/memory/graph/query.rs +++ b/src/memory/graph/query.rs @@ -13,10 +13,9 @@ //! LIMIT ?2 //! ``` //! -//! Here the same derivation runs in Rust over the injected -//! [`EntityOccurrenceIndex`] so the graph stays decoupled from storage: gather -//! the subject's nodes, fan out to the entities sharing each node, and count -//! distinct shared nodes per neighbour. +//! Persistent adapters can execute that self-join through the injected +//! [`EntityOccurrenceIndex`] fast path. Lightweight indexes used by tests fall +//! back to gathering the subject's nodes and counting in Rust. use std::collections::{HashMap, HashSet}; @@ -34,14 +33,8 @@ const DEFAULT_LIMIT: usize = 100; /// output regardless of the index's iteration order. Self-edges are excluded. /// `limit` caps the result set; `None` defaults to `DEFAULT_LIMIT` (100). /// -/// # Cost -/// -/// This issues one [`EntityOccurrenceIndex::nodes_for_entity`] call plus one -/// [`EntityOccurrenceIndex::entities_on_node`] call per node the subject -/// occurs on — `1 + subject_nodes.len()` index round-trips in total. `limit` -/// only truncates the *output*; it does not bound the number of nodes -/// gathered or the fan-out cost, so a subject with many occurrences is -/// expensive to query regardless of how small `limit` is. +/// Backends with a native co-occurrence implementation execute one set-based +/// query. Other indexes use `1 + subject_nodes.len()` portable trait calls. /// /// # Errors /// @@ -54,6 +47,13 @@ pub fn co_occurring_entities( ) -> Result> { let cap = limit.unwrap_or(DEFAULT_LIMIT); + if let Some(edges) = index + .co_occurring_entities(subject_entity, cap) + .with_context(|| format!("co_occurring_entities({subject_entity})"))? + { + return Ok(edges); + } + // For each neighbour, the set of distinct nodes shared with the subject. // Using a set (rather than a bare counter) makes the distinct-node count // robust even if an index implementation returns a node more than once. diff --git a/src/memory/graph/types.rs b/src/memory/graph/types.rs index 7435d7d..b1c3e51 100644 --- a/src/memory/graph/types.rs +++ b/src/memory/graph/types.rs @@ -42,6 +42,19 @@ pub struct GraphEdge { /// [`nodes_for_entity`]: EntityOccurrenceIndex::nodes_for_entity /// [`entities_on_node`]: EntityOccurrenceIndex::entities_on_node pub trait EntityOccurrenceIndex { + /// Optional backend-native co-occurrence query. + /// + /// Persistent adapters should implement this as one set-based query; + /// lightweight test indexes may return `Ok(None)` to use the portable + /// `nodes_for_entity` + `entities_on_node` derivation. + fn co_occurring_entities( + &self, + _subject_entity: &str, + _limit: usize, + ) -> Result>> { + Ok(None) + } + /// Distinct node ids on which `entity_id` has been indexed. /// /// An entity with no occurrences yields an empty vector (not an error). diff --git a/src/memory/mod.rs b/src/memory/mod.rs index 2892210..4c5396c 100644 --- a/src/memory/mod.rs +++ b/src/memory/mod.rs @@ -81,22 +81,6 @@ pub mod sync; pub mod tool_memory; pub mod tree; -// ── Feature-gated boundary surfaces ───────────────────────────────────────── -/// reqwest-based embedding / LLM HTTP providers. -/// -/// Gated behind the `providers-http` feature (implies `tokio`). Reserves the -/// HTTP provider seam and gates the reqwest dependency; the concrete providers -/// land with goals C3/M3. -#[cfg(feature = "providers-http")] -pub mod providers; - -/// serde schema / envelope surface for the RPC boundary. -/// -/// Gated behind the `rpc` feature. Reserves the wire-facing surface for goal -/// C5 without adding heavy dependencies. -#[cfg(feature = "rpc")] -pub mod rpc; - // ── Re-exports ────────────────────────────────────────────────────────────── pub use config::{MemoryConfig, WeightProfile}; pub use error::{MemoryEngineResult, MemoryError as MemoryEngineError}; @@ -110,6 +94,6 @@ pub use types::{ // Starter in-memory store API (kept stable for the smoke test and as a simple // reference backend while richer backends are ported under `store`). pub use store::types::{ - MemoryError, MemoryId, MemoryInput, MemoryQuery, MemoryRecord, MemoryResult, SearchHit, + MemoryId, MemoryInput, MemoryQuery, MemoryRecord, MemoryResult, SearchHit, StoreError, }; pub use store::{InMemoryMemoryStore, MemoryStore}; diff --git a/src/memory/retrieval/cover.rs b/src/memory/retrieval/cover.rs index 1f6b937..6ff0dca 100644 --- a/src/memory/retrieval/cover.rs +++ b/src/memory/retrieval/cover.rs @@ -28,20 +28,8 @@ use crate::memory::tree::{SummaryNode, TreeKind}; use super::types::{hydrated_chunk_hit, hydrated_summary_hit, QueryResponse, RetrievalHit}; -/// Default cap on returned cover items when the caller passes `limit = 0`. -const DEFAULT_LIMIT: usize = 200; - -/// Upper bound on in-window chunks scanned across all sources. -/// -/// NOTE: this is a silent truncation with no signal distinct from the -/// `limit`-based one. [`QueryResponse::truncated`] is derived only from -/// `total > hits.len()` after the [`cover_window`] result-count cut. If the -/// window contains more in-window chunks than this cap, `list_chunks` -/// truncates before `cover_window` ever sees the rest, and the caller has no -/// way to distinguish "there were more chunks than the scan cap allows" from -/// "there were exactly this many, all covered by summaries". -const MAX_WINDOW_CHUNKS: usize = 5_000; - +/// Page size used while scanning the complete in-window chunk set. Pagination +/// bounds each SQLite materialisation without silently dropping later sources. /// Derive the tree scope a chunk seals under: `path_scope` overrides /// `source_id` for shared-directory document sources, mirroring the append /// path's scope derivation. @@ -101,7 +89,11 @@ pub fn cover_window_scoped( source_scope: Option>, limit: usize, ) -> Result { - let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; + let limit = if limit == 0 { + config.retrieval.limits.cover_default_limit + } else { + limit + }; if until_ms < since_ms { return Err(anyhow::anyhow!( "cover_window: until_ms ({until_ms}) precedes since_ms ({since_ms})" @@ -140,19 +132,29 @@ fn collect_cover( source_scope: Option>, ) -> Result> { let restrict_summaries = source_id.is_some() || source_scope.is_some(); - let chunks = list_chunks( - config, - &ListChunksQuery { - source_id: source_id.map(|s| s.to_string()), - source_kind, - since_ms: Some(since_ms), - until_ms: Some(until_ms), - limit: Some(MAX_WINDOW_CHUNKS), - exclude_dropped: true, - source_scope, - ..Default::default() - }, - )?; + let page_size = config.retrieval.limits.window_chunk_page_size; + let mut chunks = Vec::new(); + loop { + let page = list_chunks( + config, + &ListChunksQuery { + source_id: source_id.map(str::to_owned), + source_kind, + since_ms: Some(since_ms), + until_ms: Some(until_ms), + limit: Some(page_size), + offset: Some(chunks.len()), + exclude_dropped: true, + source_scope: source_scope.clone(), + ..Default::default() + }, + )?; + let page_len = page.len(); + chunks.extend(page); + if page_len < page_size { + break; + } + } let mut by_source: HashMap> = HashMap::new(); for chunk in chunks { diff --git a/src/memory/retrieval/cover_tests.rs b/src/memory/retrieval/cover_tests.rs index 8b0c021..8262dc5 100644 --- a/src/memory/retrieval/cover_tests.rs +++ b/src/memory/retrieval/cover_tests.rs @@ -170,3 +170,29 @@ fn scoped_cover_filters_memory_sources_before_result_limit() { assert_eq!(response.hits.len(), 1); assert_eq!(response.hits[0].tree_scope, "slack:#allowed"); } + +#[test] +fn cover_paginates_past_scan_page_without_starving_later_source() { + let (_tmp, cfg) = test_config(); + let ts = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + let page_size = cfg.retrieval.limits.window_chunk_page_size; + let mut chunks = (0..page_size) + .map(|seq| sample_chunk_at("alpha", seq as u32, &format!("alpha-{seq}"), ts)) + .collect::>(); + chunks.push(sample_chunk_at("zulu", 0, "must survive page boundary", ts)); + insert_chunks(&cfg, &chunks); + + let response = cover_window( + &cfg, + ts.timestamp_millis() - 1, + ts.timestamp_millis() + 1, + None, + None, + chunks.len(), + ) + .unwrap(); + + assert_eq!(response.total, chunks.len()); + assert!(response.hits.iter().any(|hit| hit.tree_scope == "zulu")); + assert!(!response.truncated); +} diff --git a/src/memory/retrieval/drill_down.rs b/src/memory/retrieval/drill_down.rs index 4c4f133..408c8c8 100644 --- a/src/memory/retrieval/drill_down.rs +++ b/src/memory/retrieval/drill_down.rs @@ -120,7 +120,7 @@ fn walk_with_embeddings( // Update per-document latest version with any doc-roots on THIS level // before walking, so two side-by-side revisions resolve to the newer. for id in ¤t_level { - if let Some(s) = summary_by_id.get(id) { + if let Some(s) = summary_by_id.get(id).filter(|summary| !summary.deleted) { if let Some(doc_id) = s.doc_id.as_deref() { let v = s.version_ms.unwrap_or(i64::MIN); max_version_by_doc @@ -167,28 +167,20 @@ fn walk_with_embeddings( for id in ¤t_level { if let Some(summary) = summary_by_id.remove(id) { + if summary.deleted { + continue; + } // Latest-wins: skip a doc-root superseded by a newer revision. if let Some(doc_id) = summary.doc_id.as_deref() { let v = summary.version_ms.unwrap_or(i64::MIN); if max_version_by_doc.get(doc_id).is_some_and(|&max| v < max) { continue; } - // Dedup duplicates at the winning version. - // - // NOTE (see module-level doc): this insert claims - // `emitted_docs` for `doc_id` even if the soft-delete - // check below then skips this node. A soft-deleted - // newest revision therefore suppresses every older - // sibling revision too — the whole document vanishes - // from this level instead of falling back to the - // latest surviving one. + // Dedup duplicates at the winning surviving version. if !emitted_docs.insert(doc_id.to_string()) { continue; } } - if summary.deleted { - continue; - } let scope = tree_by_id .get(&summary.tree_id) .map(|t| t.scope.clone()) diff --git a/src/memory/retrieval/drill_down_tests.rs b/src/memory/retrieval/drill_down_tests.rs index 19fadd0..2706b71 100644 --- a/src/memory/retrieval/drill_down_tests.rs +++ b/src/memory/retrieval/drill_down_tests.rs @@ -277,3 +277,59 @@ async fn drill_down_surfaces_only_latest_doc_version() { "superseded chunk filtered" ); } + +#[tokio::test] +async fn drill_down_falls_back_when_latest_doc_revision_is_deleted() { + let (_tmp, cfg) = test_config(); + let ts = fixed_ts(); + let old_chunk = leaf_chunk("notion:pageA", 0, "surviving old body", ts); + let deleted_chunk = leaf_chunk("notion:pageA", 1, "deleted new body", ts); + insert_chunks(&cfg, &[old_chunk.clone(), deleted_chunk.clone()]); + + let tree = source_tree("tree:notion", "notion:conn1", Some("s:merge"), 1000); + insert_tree_row(&cfg, &tree); + let mut old = summary_node( + "s:v1", + "tree:notion", + 1, + Some("s:merge"), + &[&old_chunk.id], + "old", + ts, + ); + old.doc_id = Some("notion:conn1:pageA".into()); + old.version_ms = Some(100); + let mut deleted = summary_node( + "s:v2", + "tree:notion", + 1, + Some("s:merge"), + &[&deleted_chunk.id], + "new", + ts, + ); + deleted.doc_id = old.doc_id.clone(); + deleted.version_ms = Some(200); + deleted.deleted = true; + let merge = summary_node( + "s:merge", + "tree:notion", + 1000, + None, + &["s:v1", "s:v2"], + "merge", + ts, + ); + insert_summary(&cfg, &old); + insert_summary(&cfg, &deleted); + insert_summary(&cfg, &merge); + + let out = drill_down(&cfg, "s:merge", 3, None, &inert(), None) + .await + .unwrap(); + let ids: Vec<_> = out.iter().map(|hit| hit.node_id.as_str()).collect(); + assert!(ids.contains(&"s:v1")); + assert!(ids.contains(&old_chunk.id.as_str())); + assert!(!ids.contains(&"s:v2")); + assert!(!ids.contains(&deleted_chunk.id.as_str())); +} diff --git a/src/memory/retrieval/fast.rs b/src/memory/retrieval/fast.rs index 70a09f6..2322b9a 100644 --- a/src/memory/retrieval/fast.rs +++ b/src/memory/retrieval/fast.rs @@ -11,15 +11,12 @@ use crate::memory::score::embed::Embedder; use crate::memory::score::store::lookup_entity; use crate::memory::tree::store::{get_summaries_batch, get_tree}; -use super::fetch::{fetch_leaves, MAX_BATCH}; +use super::fetch::fetch_leaves; use super::source::query_source; use super::types::{hydrated_summary_hit, QueryResponse}; -const LOOKUP_LIMIT: usize = 500; const DEFAULT_LIMIT: usize = 10; -const MAX_RETRIEVE_LIMIT: usize = 100; const DEFAULT_MAX_HOPS: u32 = 2; -const MAX_GRAPH_HOPS: u32 = 4; #[derive(Clone, Debug)] pub struct FastRetrieveOptions { @@ -46,15 +43,16 @@ pub async fn fast_retrieve( source_scope: Option<&HashSet>, options: FastRetrieveOptions, ) -> Result { + let limits = &config.retrieval.limits; let limit = if options.limit == 0 { - DEFAULT_LIMIT + limits.default_limit } else { - options.limit.min(MAX_RETRIEVE_LIMIT) + options.limit.min(limits.max_limit) }; let max_hops = if options.max_hops == 0 { - DEFAULT_MAX_HOPS + limits.default_graph_hops } else { - options.max_hops.min(MAX_GRAPH_HOPS) + options.max_hops.min(limits.max_graph_hops) }; let query = query.trim(); if query.is_empty() { @@ -134,8 +132,9 @@ fn local_candidates( ) -> Result> { let mut output = HashMap::new(); for pair in pairs { - let left = lookup_entity(config, &pair.a, Some(LOOKUP_LIMIT))?; - let right = lookup_entity(config, &pair.b, Some(LOOKUP_LIMIT))?; + let occurrence_limit = config.retrieval.limits.occurrence_lookup_limit; + let left = lookup_entity(config, &pair.a, Some(occurrence_limit))?; + let right = lookup_entity(config, &pair.b, Some(occurrence_limit))?; let right_nodes: HashMap<_, _> = right .iter() .map(|hit| (hit.node_id.as_str(), hit.timestamp_ms)) @@ -186,7 +185,7 @@ fn resolve_local( .map(|(id, _)| id.clone()) .collect(); let mut by_id = HashMap::new(); - for ids in leaf_ids.chunks(MAX_BATCH) { + for ids in leaf_ids.chunks(config.retrieval.limits.fetch_batch_limit) { for hit in fetch_leaves(config, ids)? { by_id.insert(hit.node_id.clone(), hit); } @@ -284,77 +283,5 @@ fn dedup_ids(ids: impl Iterator) -> Vec { } #[cfg(test)] -mod tests { - use super::*; - use crate::memory::graph::{pairs_from_entities, upsert_edges}; - use crate::memory::retrieval::test_support::{ - fixed_ts, index_entity_occurrence, insert_summary, insert_tree_row, source_tree, - summary_node, test_config, - }; - use crate::memory::score::embed::InertEmbedder; - use crate::memory::score::extract::EntityKind; - - #[test] - fn options_and_ids_are_bounded_deterministically() { - assert_eq!(FastRetrieveOptions::default().limit, 10); - assert_eq!( - dedup_ids(vec!["a".into(), "a".into(), "b".into()].into_iter()), - vec!["a", "b"] - ); - } - - #[tokio::test] - async fn local_branch_intersects_entities_and_applies_scope_before_limit() { - let (_temp, config) = test_config(); - let allowed_tree = source_tree("allowed-tree", "slack:#allowed", Some("allowed"), 1); - let denied_tree = source_tree("denied-tree", "slack:#denied", Some("denied"), 1); - insert_tree_row(&config, &allowed_tree); - insert_tree_row(&config, &denied_tree); - for (id, tree_id) in [("allowed", "allowed-tree"), ("denied", "denied-tree")] { - insert_summary( - &config, - &summary_node(id, tree_id, 1, None, &[], id, fixed_ts()), - ); - for (entity, kind) in [ - ("person:alice", EntityKind::Person), - ("topic:launch", EntityKind::Topic), - ] { - index_entity_occurrence( - &config, - entity, - kind, - entity, - id, - "summary", - fixed_ts().timestamp_millis(), - Some(tree_id), - ); - } - } - upsert_edges( - &config, - &pairs_from_entities(&["person:alice".into(), "topic:launch".into()]), - fixed_ts().timestamp_millis(), - ) - .unwrap(); - let scope = HashSet::from(["slack:#allowed".to_string()]); - - let response = fast_retrieve( - &config, - "alice launch", - &["person:alice".into(), "topic:launch".into()], - &InertEmbedder, - Some(&scope), - FastRetrieveOptions { - limit: 1, - ..Default::default() - }, - ) - .await - .unwrap(); - - assert_eq!(response.total, 1); - assert_eq!(response.hits[0].node_id, "allowed"); - assert_eq!(response.hits[0].score, 2.0); - } -} +#[path = "fast_tests.rs"] +mod tests; diff --git a/src/memory/retrieval/fast_tests.rs b/src/memory/retrieval/fast_tests.rs new file mode 100644 index 0000000..11a6406 --- /dev/null +++ b/src/memory/retrieval/fast_tests.rs @@ -0,0 +1,226 @@ +use super::*; +use crate::memory::graph::{pairs_from_entities, upsert_edges, PairDistance}; +use crate::memory::retrieval::test_support::{ + fixed_ts, index_entity_occurrence, insert_chunks, insert_summary, insert_tree_row, + sample_chunk, source_tree, summary_node, test_config, +}; +use crate::memory::score::embed::InertEmbedder; +use crate::memory::score::extract::EntityKind; + +#[test] +fn options_and_ids_are_bounded_deterministically() { + assert_eq!(FastRetrieveOptions::default().limit, 10); + assert_eq!(FastRetrieveOptions::default().max_hops, 2); + assert_eq!( + dedup_ids(vec!["a".into(), "a".into(), "b".into()].into_iter()), + vec!["a", "b"] + ); +} + +#[tokio::test] +async fn blank_query_is_empty_without_opening_storage() { + let (_temp, config) = test_config(); + let response = fast_retrieve( + &config, + " ", + &[], + &InertEmbedder, + None, + FastRetrieveOptions::default(), + ) + .await + .unwrap(); + assert!(response.hits.is_empty()); +} + +#[tokio::test] +async fn dense_fallback_filters_scope_before_limit() { + let (_temp, config) = test_config(); + for (id, scope) in [("allowed", "slack:#allowed"), ("denied", "slack:#denied")] { + let tree_id = format!("tree-{id}"); + insert_tree_row(&config, &source_tree(&tree_id, scope, Some(id), 1)); + insert_summary( + &config, + &summary_node(id, &tree_id, 1, None, &[], id, fixed_ts()), + ); + } + let scope = HashSet::from(["slack:#allowed".to_string()]); + + let response = fast_retrieve( + &config, + "anything", + &[], + &InertEmbedder, + Some(&scope), + FastRetrieveOptions { + limit: 1, + max_hops: 99, + time_window_days: None, + }, + ) + .await + .unwrap(); + + assert_eq!(response.total, 1); + assert_eq!(response.hits[0].node_id, "allowed"); +} + +#[tokio::test] +async fn occurrence_fallback_prioritizes_hits_containing_query_entities() { + let (_temp, config) = test_config(); + let tree = source_tree("tree", "slack:#eng", Some("matching"), 1); + insert_tree_row(&config, &tree); + for (id, entities) in [ + ("unrelated", Vec::new()), + ("matching", vec!["person:alice".to_string()]), + ] { + let mut node = summary_node(id, "tree", 1, None, &[], id, fixed_ts()); + node.entities = entities; + insert_summary(&config, &node); + } + + let response = fast_retrieve( + &config, + "alice", + &["person:alice".into()], + &InertEmbedder, + None, + FastRetrieveOptions { + limit: 1, + max_hops: 0, + time_window_days: None, + }, + ) + .await + .unwrap(); + + assert_eq!(response.total, 2); + assert_eq!(response.hits[0].node_id, "matching"); + assert!(response.truncated); +} + +#[test] +fn local_candidates_intersect_occurrences_and_keep_latest_timestamp() { + let (_temp, config) = test_config(); + for (entity, kind, timestamp) in [ + ("person:alice", EntityKind::Person, 10), + ("topic:launch", EntityKind::Topic, 20), + ] { + index_entity_occurrence( + &config, + entity, + kind, + entity, + "shared-node", + "summary", + timestamp, + Some("tree"), + ); + } + let pairs = vec![PairDistance { + a: "person:alice".into(), + b: "topic:launch".into(), + dist: 1, + }]; + + let candidates = local_candidates(&config, &pairs).unwrap(); + + assert_eq!(candidates["shared-node"].matched.len(), 2); + assert_eq!(candidates["shared-node"].latest_ts, 20); +} + +#[test] +fn resolve_local_hydrates_leaf_and_summary_and_skips_missing_or_out_of_scope() { + let (_temp, config) = test_config(); + let leaf = sample_chunk("slack:#allowed", 0, "leaf"); + insert_chunks(&config, std::slice::from_ref(&leaf)); + insert_tree_row( + &config, + &source_tree("summary-tree", "slack:#allowed", Some("summary"), 1), + ); + insert_summary( + &config, + &summary_node( + "summary", + "summary-tree", + 1, + None, + &[], + "summary", + fixed_ts(), + ), + ); + let candidate = |kind: &str, matched: &[&str], latest_ts| Candidate { + node_kind: kind.into(), + matched: matched.iter().map(|value| (*value).to_string()).collect(), + latest_ts, + }; + let candidates = HashMap::from([ + (leaf.id.clone(), candidate("leaf", &["a"], 3)), + ("summary".into(), candidate("summary", &["a", "b"], 2)), + ("missing".into(), candidate("summary", &["a", "b", "c"], 1)), + ]); + let scope = HashSet::from(["slack:#allowed".to_string()]); + + let response = resolve_local(&config, candidates, Some(&scope), 10).unwrap(); + + assert_eq!(response.total, 2); + assert_eq!(response.hits[0].node_id, "summary"); + assert_eq!(response.hits[0].score, 2.0); + assert_eq!(response.hits[1].node_id, leaf.id); +} + +#[tokio::test] +async fn local_branch_intersects_entities_and_applies_scope_before_limit() { + let (_temp, config) = test_config(); + let allowed_tree = source_tree("allowed-tree", "slack:#allowed", Some("allowed"), 1); + let denied_tree = source_tree("denied-tree", "slack:#denied", Some("denied"), 1); + insert_tree_row(&config, &allowed_tree); + insert_tree_row(&config, &denied_tree); + for (id, tree_id) in [("allowed", "allowed-tree"), ("denied", "denied-tree")] { + insert_summary( + &config, + &summary_node(id, tree_id, 1, None, &[], id, fixed_ts()), + ); + for (entity, kind) in [ + ("person:alice", EntityKind::Person), + ("topic:launch", EntityKind::Topic), + ] { + index_entity_occurrence( + &config, + entity, + kind, + entity, + id, + "summary", + fixed_ts().timestamp_millis(), + Some(tree_id), + ); + } + } + upsert_edges( + &config, + &pairs_from_entities(&["person:alice".into(), "topic:launch".into()]), + fixed_ts().timestamp_millis(), + ) + .unwrap(); + let scope = HashSet::from(["slack:#allowed".to_string()]); + + let response = fast_retrieve( + &config, + "alice launch", + &["person:alice".into(), "topic:launch".into()], + &InertEmbedder, + Some(&scope), + FastRetrieveOptions { + limit: 1, + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(response.total, 1); + assert_eq!(response.hits[0].node_id, "allowed"); + assert_eq!(response.hits[0].score, 2.0); +} diff --git a/src/memory/retrieval/fetch.rs b/src/memory/retrieval/fetch.rs index e92fd2b..807d721 100644 --- a/src/memory/retrieval/fetch.rs +++ b/src/memory/retrieval/fetch.rs @@ -3,7 +3,7 @@ //! //! The contract, ported from OpenHuman's `memory_tree::retrieval::fetch`: //! "given these chunk ids, give me the full content + metadata so I can cite." -//! The batch is capped at [`MAX_BATCH`] to keep the round-trip bounded. Missing +//! The batch is capped by retrieval config to keep the round-trip bounded. Missing //! ids are silently skipped — the return is best-effort, so partial failures //! are visible via `hits.len() < ids.len()`. //! @@ -19,8 +19,8 @@ use crate::memory::score::store::get_scores_batch; use super::types::{hydrated_chunk_hit, RetrievalHit}; -/// Max batch size. Callers that pass more than this get truncated (no error so -/// the caller still sees a partial result). +/// Default fetch batch retained for source compatibility. Runtime behavior is +/// controlled by [`RetrievalLimits::fetch_batch_limit`](crate::memory::config::RetrievalLimits::fetch_batch_limit). pub const MAX_BATCH: usize = 20; /// Fetch chunk rows by id in the provided order. Missing ids are dropped from @@ -30,8 +30,9 @@ pub fn fetch_leaves(config: &MemoryConfig, chunk_ids: &[String]) -> Result MAX_BATCH { - &chunk_ids[..MAX_BATCH] + let max_batch = config.retrieval.limits.fetch_batch_limit; + let ids: &[String] = if chunk_ids.len() > max_batch { + &chunk_ids[..max_batch] } else { chunk_ids }; diff --git a/src/memory/retrieval/global.rs b/src/memory/retrieval/global.rs index fbb3f13..ff49cd0 100644 --- a/src/memory/retrieval/global.rs +++ b/src/memory/retrieval/global.rs @@ -19,7 +19,7 @@ use chrono::{DateTime, TimeZone, Utc}; use crate::memory::chunks::{get_chunk_embeddings_batch, get_chunks_batch, SourceKind}; use crate::memory::config::MemoryConfig; use crate::memory::score::embed::Embedder; -use crate::memory::score::store::lookup_entity; +use crate::memory::score::store::{get_dropped_chunk_ids_batch, lookup_entity_in_window}; use crate::memory::tree::store::{ get_summaries_batch, get_summary_embeddings_batch, get_trees_batch, }; @@ -29,14 +29,6 @@ use super::types::{hit_from_chunk, hit_from_summary, QueryResponse}; /// Default result cap applied by both primitives when the caller passes /// `limit = 0`. -const DEFAULT_LIMIT: usize = 10; - -/// Per-entity node cap for the topic projection: `resolve_topic_hits` reads -/// at most this many newest entity-index rows for a given `entity_id` before -/// any window filter is applied (see the gotcha documented on -/// [`query_topic`]). -const TOPIC_LOOKUP_CAP: usize = 200; - /// Cross-source digest over `[since_ms, until_ms]`. /// /// Gathers summaries from every source tree (optionally narrowed by @@ -46,11 +38,8 @@ const TOPIC_LOOKUP_CAP: usize = 200; /// /// Returns an error if `until_ms < since_ms`. /// -/// Cost note: the window filter is applied in Rust, after -/// `collect_source_hits` has materialized every non-deleted summary across -/// every selected source tree (embeddings included) — see -/// `collect_source_hits`'s docs for why this does not scale with store -/// size. +/// The window predicate is applied in SQLite per selected tree before summary +/// embeddings are hydrated. pub async fn query_global( config: &MemoryConfig, since_ms: i64, @@ -60,17 +49,18 @@ pub async fn query_global( embedder: &dyn Embedder, limit: usize, ) -> Result { - let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; + let limits = &config.retrieval.limits; + let limit = if limit == 0 { + limits.default_limit + } else { + limit.min(limits.max_limit) + }; if until_ms < since_ms { return Err(anyhow::anyhow!( "query_global: until_ms ({until_ms}) precedes since_ms ({since_ms})" )); } - let since = ms_to_utc(since_ms); - let until = ms_to_utc(until_ms); - - let mut scored = collect_source_hits(config, None, source_kind)?; - scored.retain(|(h, _)| h.time_range_end >= since && h.time_range_start <= until); + let scored = collect_source_hits(config, None, source_kind, Some((since_ms, until_ms)))?; let total = scored.len(); let mut hits = order_hits(scored, query, embedder).await; @@ -85,15 +75,10 @@ pub async fn query_global( /// hit (summaries hydrated with their tree scope, leaves with their source id), /// optionally restricted to `[since_ms, until_ms]` and reranked by `query`. /// -/// # Gotcha: the window filter runs after a hard row cap -/// -/// `resolve_topic_hits` caps the entity-index lookup at `TOPIC_LOOKUP_CAP` -/// newest rows *before* the `since_ms`/`until_ms` window is applied here. For -/// an entity with more mentions than the cap, a window reaching further back -/// than the cap's newest rows can return zero or partial hits even though -/// matching rows exist further back in the index, and the returned `total` -/// under-reports the true match count for that window. Callers should not -/// assume completeness for a high-mention entity queried with an old window. +/// The entity-index timestamp window is pushed into SQLite before the +/// the configured topic lookup cap, so historical windows are not crowded +/// out by newer mentions. A window containing more than the cap remains +/// intentionally bounded. pub async fn query_topic( config: &MemoryConfig, entity_id: &str, @@ -103,13 +88,18 @@ pub async fn query_topic( embedder: &dyn Embedder, limit: usize, ) -> Result { - let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; + let limits = &config.retrieval.limits; + let limit = if limit == 0 { + limits.default_limit + } else { + limit.min(limits.max_limit) + }; let entity_id = entity_id.trim(); if entity_id.is_empty() { return Ok(QueryResponse::empty()); } - let mut scored = resolve_topic_hits(config, entity_id)?; + let mut scored = resolve_topic_hits(config, entity_id, since_ms, until_ms)?; if since_ms.is_some() || until_ms.is_some() { let since = since_ms.map(ms_to_utc); @@ -130,14 +120,25 @@ pub async fn query_topic( /// preserving the entity index's newest-first order and hydrating embeddings /// (summary sidecar / chunk sidecar) for the optional rerank pass. /// -/// Caps the lookup at `TOPIC_LOOKUP_CAP` rows (see the caller-facing gotcha -/// on [`query_topic`]). Soft-deleted summaries are excluded (`node.deleted` +/// Caps the already-windowed lookup at the configured topic limit. Soft-deleted +/// summaries are excluded (`node.deleted` /// check below); leaf chunks have no equivalent tombstone check here, so a /// dropped chunk that is still indexed can still surface in topic results. /// An entity-index row whose `node_id` resolves to neither a summary nor a /// chunk (a stale index entry) is silently skipped. -fn resolve_topic_hits(config: &MemoryConfig, entity_id: &str) -> Result> { - let entity_hits = lookup_entity(config, entity_id, Some(TOPIC_LOOKUP_CAP))?; +fn resolve_topic_hits( + config: &MemoryConfig, + entity_id: &str, + since_ms: Option, + until_ms: Option, +) -> Result> { + let entity_hits = lookup_entity_in_window( + config, + entity_id, + since_ms, + until_ms, + Some(config.retrieval.limits.topic_lookup_limit), + )?; if entity_hits.is_empty() { return Ok(Vec::new()); } @@ -155,6 +156,7 @@ fn resolve_topic_hits(config: &MemoryConfig, entity_id: &str) -> Result Result Result::MIN_UTC` -/// / `MAX_UTC` — it falls back to `Utc::now()`. A caller passing a sentinel -/// meant to mean "no lower/upper bound" for `since_ms`/`until_ms` will -/// instead get "now", silently narrowing the intended window. fn ms_to_utc(ms: i64) -> DateTime { - Utc.timestamp_millis_opt(ms) - .single() - .unwrap_or_else(Utc::now) + Utc.timestamp_millis_opt(ms).single().unwrap_or(if ms < 0 { + DateTime::::MIN_UTC + } else { + DateTime::::MAX_UTC + }) } #[cfg(test)] diff --git a/src/memory/retrieval/global_tests.rs b/src/memory/retrieval/global_tests.rs index 91db0c7..0f9835a 100644 --- a/src/memory/retrieval/global_tests.rs +++ b/src/memory/retrieval/global_tests.rs @@ -3,8 +3,8 @@ use super::*; use crate::memory::config::MemoryConfig; use crate::memory::retrieval::test_support::{ - fixed_ts, index_entity_occurrence, insert_chunks, insert_summary, insert_tree_row, - sample_chunk, source_tree, summary_node, test_config, + fixed_ts, index_entity_occurrence, insert_chunks, insert_score, insert_summary, + insert_tree_row, sample_chunk, sample_chunk_at, source_tree, summary_node, test_config, }; use crate::memory::score::embed::InertEmbedder; use crate::memory::score::extract::EntityKind; @@ -14,6 +14,16 @@ fn inert() -> InertEmbedder { InertEmbedder::new() } +#[test] +fn millisecond_conversion_saturates_out_of_range_values() { + assert_eq!(ms_to_utc(i64::MIN), DateTime::::MIN_UTC); + assert_eq!(ms_to_utc(i64::MAX), DateTime::::MAX_UTC); + assert_eq!( + ms_to_utc(1_700_000_000_000), + Utc.timestamp_millis_opt(1_700_000_000_000).unwrap() + ); +} + /// Seed a source tree at `scope` with one L1 summary `summary_id` at `ts`. fn seed_tree(cfg: &MemoryConfig, scope: &str, summary_id: &str, ts: DateTime) { let tree_id = format!("tree:{scope}"); @@ -189,3 +199,91 @@ async fn query_topic_window_filters_old_nodes() { assert_eq!(resp.hits.len(), 1); assert_eq!(resp.hits[0].node_id, "s:new"); } + +#[tokio::test] +async fn query_topic_applies_historical_window_before_lookup_cap() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + let old = Utc.timestamp_millis_opt(1_000_000_000_000).unwrap(); + let chunk = sample_chunk_at("slack:#history", 0, "historic phoenix", old); + insert_chunks(&cfg, std::slice::from_ref(&chunk)); + index_entity_occurrence( + &cfg, + "topic:phoenix", + EntityKind::Topic, + "phoenix", + &chunk.id, + "leaf", + old.timestamp_millis(), + None, + ); + // More than the configured topic cap recent stale index rows used to hide the old + // real row because LIMIT ran before the historical-window filter. + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + for index in 0..=cfg.retrieval.limits.topic_lookup_limit { + tx.execute( + "INSERT INTO mem_tree_entity_index + (entity_id,node_id,node_kind,entity_kind,surface,score,timestamp_ms,tree_id,is_user) + VALUES ('topic:phoenix',?1,'leaf','topic','phoenix',1.0,?2,NULL,0)", + rusqlite::params![ + format!("stale-{index}"), + fixed_ts().timestamp_millis() + index as i64 + ], + )?; + } + tx.commit()?; + Ok(()) + }) + .unwrap(); + + let response = query_topic( + &cfg, + "topic:phoenix", + Some(old.timestamp_millis() - 1), + Some(old.timestamp_millis() + 1), + None, + &inert(), + 10, + ) + .await + .unwrap(); + + assert_eq!(response.hits.len(), 1); + assert_eq!(response.hits[0].node_id, chunk.id); +} + +#[tokio::test] +async fn query_topic_excludes_chunks_with_dropped_score_rows() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + let chunk = sample_chunk("slack:#eng", 0, "phoenix"); + insert_chunks(&cfg, std::slice::from_ref(&chunk)); + insert_score(&cfg, &chunk.id, 0.1); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_score SET dropped=1 WHERE chunk_id=?1", + rusqlite::params![chunk.id], + )?; + Ok(()) + }) + .unwrap(); + index_entity_occurrence( + &cfg, + "topic:phoenix", + EntityKind::Topic, + "phoenix", + &chunk.id, + "leaf", + fixed_ts().timestamp_millis(), + None, + ); + + let response = query_topic(&cfg, "topic:phoenix", None, None, None, &inert(), 10) + .await + .unwrap(); + + assert!(response.hits.is_empty()); +} diff --git a/src/memory/retrieval/graph_adapter.rs b/src/memory/retrieval/graph_adapter.rs index feb429c..9632b59 100644 --- a/src/memory/retrieval/graph_adapter.rs +++ b/src/memory/retrieval/graph_adapter.rs @@ -8,15 +8,9 @@ //! binds a [`MemoryConfig`] so retrieval can compute graph relevance for an //! entity without the graph layer depending on SQLite. //! -//! NOTE: the co-occurrence derivation in `crate::memory::graph` calls -//! [`ConfigEntityIndex::entities_on_node`] once per node returned by -//! [`ConfigEntityIndex::nodes_for_entity`] — i.e. 1 + up to -//! `OCCURRENCE_LOOKUP_LIMIT` separate SQL queries per derivation, rather -//! than a single self-join. Combined with the newest-first truncation at -//! `OCCURRENCE_LOOKUP_LIMIT`, edge weights for popular entities (more -//! occurrences than the cap) are systematically undercounted and -//! "strongest neighbor" ordering can be wrong. See the improvement plan for -//! the planned SQL self-join fast path. +//! The adapter implements the graph trait's backend-native co-occurrence fast +//! path as one SQL self-join, without the per-entity occurrence cap. Direct +//! calls to `nodes_for_entity` remain capped for bounded diagnostics. //! //! Also recall the crate-level NOTE in [`crate::memory::retrieval`]: nothing //! in this crate currently calls into `crate::memory::graph` through this @@ -25,13 +19,11 @@ use anyhow::Result; use crate::memory::config::MemoryConfig; -use crate::memory::graph::EntityOccurrenceIndex; +use crate::memory::graph::{EntityOccurrenceIndex, GraphEdge}; use crate::memory::score::store::{list_entity_ids_for_node, lookup_entity}; -/// Per-entity lookup cap for the occurrence reads. Popular entities can touch -/// many nodes; this bounds the fan-out while staying well above any realistic -/// `k`. -const OCCURRENCE_LOOKUP_LIMIT: usize = 500; +#[cfg(test)] +pub(crate) const OCCURRENCE_LOOKUP_LIMIT: usize = 500; /// SQLite-backed [`EntityOccurrenceIndex`] over `mem_tree_entity_index`. pub struct ConfigEntityIndex<'a> { @@ -46,11 +38,49 @@ impl<'a> ConfigEntityIndex<'a> { } impl EntityOccurrenceIndex for ConfigEntityIndex<'_> { + fn co_occurring_entities( + &self, + subject_entity: &str, + limit: usize, + ) -> Result>> { + crate::memory::chunks::with_connection(self.config, |conn| { + let limit = limit.min(i64::MAX as usize) as i64; + let mut stmt = conn.prepare( + "SELECT b.entity_id, COUNT(DISTINCT a.node_id) AS weight + FROM mem_tree_entity_index a + JOIN mem_tree_entity_index b ON a.node_id = b.node_id + LEFT JOIN mem_tree_score score ON score.chunk_id = a.node_id + LEFT JOIN mem_tree_summaries summary ON summary.id = a.node_id + WHERE a.entity_id = ?1 AND b.entity_id <> ?1 + AND ((a.node_kind = 'summary' AND summary.id IS NOT NULL AND summary.deleted = 0) + OR (a.node_kind <> 'summary' AND COALESCE(score.dropped, 0) = 0)) + GROUP BY b.entity_id + ORDER BY weight DESC, b.entity_id ASC + LIMIT ?2", + )?; + let rows = stmt + .query_map(rusqlite::params![subject_entity, limit], |row| { + let weight: i64 = row.get(1)?; + Ok(GraphEdge { + subject: subject_entity.to_string(), + object: row.get(0)?, + weight: weight.max(0).min(i64::from(u32::MAX)) as u32, + }) + })? + .collect::>>()?; + Ok(Some(rows)) + }) + } + /// Distinct node ids indexed against `entity_id`, newest-first, - /// deduplicated, capped at `OCCURRENCE_LOOKUP_LIMIT`. Returns `Err` on + /// deduplicated, capped by retrieval config. Returns `Err` on /// a SQLite read failure; an unknown `entity_id` yields `Ok(vec![])`. fn nodes_for_entity(&self, entity_id: &str) -> Result> { - let hits = lookup_entity(self.config, entity_id, Some(OCCURRENCE_LOOKUP_LIMIT))?; + let hits = lookup_entity( + self.config, + entity_id, + Some(self.config.retrieval.limits.occurrence_lookup_limit), + )?; let mut out: Vec = Vec::with_capacity(hits.len()); // `lookup_entity` already returns distinct (entity_id, node_id) rows // ordered newest-first; collect node ids preserving first-seen order. @@ -63,9 +93,9 @@ impl EntityOccurrenceIndex for ConfigEntityIndex<'_> { Ok(out) } - /// Canonical entity ids indexed against `node_id`. Called once per node - /// returned by [`nodes_for_entity`](Self::nodes_for_entity) — see the - /// module-level NOTE on the resulting query-count cost. + /// Canonical entity ids indexed against `node_id`. This supports the + /// portable graph fallback and direct diagnostics; production + /// co-occurrence uses the set-based method above. fn entities_on_node(&self, node_id: &str) -> Result> { list_entity_ids_for_node(self.config, node_id) } diff --git a/src/memory/retrieval/graph_adapter_tests.rs b/src/memory/retrieval/graph_adapter_tests.rs index 63949d5..e172d53 100644 --- a/src/memory/retrieval/graph_adapter_tests.rs +++ b/src/memory/retrieval/graph_adapter_tests.rs @@ -1,7 +1,8 @@ //! Tests for the SQLite-backed graph occurrence adapter. +use crate::memory::graph::co_occurring_entities; use crate::memory::graph::EntityOccurrenceIndex; -use crate::memory::retrieval::graph_adapter::ConfigEntityIndex; +use crate::memory::retrieval::graph_adapter::{ConfigEntityIndex, OCCURRENCE_LOOKUP_LIMIT}; use crate::memory::retrieval::test_support::{fixed_ts, index_entity_occurrence, test_config}; use crate::memory::score::extract::EntityKind; @@ -52,3 +53,41 @@ fn config_entity_index_reads_nodes_and_entities_from_sqlite_index() { assert!(idx.nodes_for_entity("topic:missing").unwrap().is_empty()); assert!(idx.entities_on_node("node-missing").unwrap().is_empty()); } + +#[test] +fn sqlite_cooccurrence_fast_path_counts_beyond_occurrence_cap_and_filters_dropped_nodes() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + for index in 0..=OCCURRENCE_LOOKUP_LIMIT { + let node = format!("node-{index}"); + for entity in ["topic:phoenix", "person:alice"] { + tx.execute( + "INSERT INTO mem_tree_entity_index + (entity_id,node_id,node_kind,entity_kind,surface,score,timestamp_ms,tree_id,is_user) + VALUES (?1,?2,'leaf','topic',?1,1.0,?3,NULL,0)", + rusqlite::params![entity, node, index as i64], + )?; + } + } + tx.execute( + "INSERT INTO mem_tree_score + (chunk_id,total,token_count_signal,unique_words_signal,metadata_weight, + source_weight,interaction_weight,entity_density,dropped,computed_at_ms) + VALUES ('node-0',0,0,0,0,0,0,0,1,0)", + [], + )?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + + let edges = + co_occurring_entities(&ConfigEntityIndex::new(&cfg), "topic:phoenix", Some(10)).unwrap(); + + assert_eq!(edges.len(), 1); + assert_eq!(edges[0].object, "person:alice"); + assert_eq!(edges[0].weight, OCCURRENCE_LOOKUP_LIMIT as u32); +} diff --git a/src/memory/retrieval/mod.rs b/src/memory/retrieval/mod.rs index be860c8..a8a2f40 100644 --- a/src/memory/retrieval/mod.rs +++ b/src/memory/retrieval/mod.rs @@ -16,7 +16,7 @@ //! - [`query_topic`] — entity/topic-scoped retrieval reconstructed from the //! entity index plus hydrated source-tree nodes. //! - [`search_entities`] — fuzzy `LIKE` lookup over the entity index. -//! - [`drill_down`] — descend a summary's `child_ids` (BFS, optional rerank). +//! - [`drill_down()`] — descend a summary's `child_ids` (BFS, optional rerank). //! - [`cover_window`] — minimum-node cover of a `[since, until]` window. //! - [`fetch_leaves`] — batch-hydrate raw chunk leaves by id, capped. //! diff --git a/src/memory/retrieval/rerank.rs b/src/memory/retrieval/rerank.rs index 02c036b..783be8b 100644 --- a/src/memory/retrieval/rerank.rs +++ b/src/memory/retrieval/rerank.rs @@ -5,11 +5,9 @@ //! and the hit's stored embedding. Hits with no embedding (legacy rows, or //! leaves whose chunk was never embedded) sort to the bottom. //! -//! NOTE: the un-embedded tail is NOT kept in incoming order — every -//! un-embedded hit shares the `NEG_INFINITY` similarity sentinel, so the tie -//! break (`time_range_end` DESC, see [`rerank_by_semantic_similarity`]) reorders -//! it newest-first. The same tie break applies among embedded hits with equal -//! similarity. +//! Un-embedded hits sort after embedded hits while preserving their incoming +//! order. Equal-similarity embedded hits use recency as a deterministic tie +//! break. //! //! Embedding failures (e.g. a local model being unavailable) never surface as //! an error to the caller: the helper logs nothing (per repo rules) and falls @@ -28,10 +26,9 @@ use super::types::RetrievalHit; /// shorter of the two rather than panicking, since [`Iterator::zip`] stops at /// the shorter side). /// -/// Ordering: embedded hits sort before un-embedded hits; within each group, -/// ties break on `time_range_end` DESC (see the module-level NOTE — this -/// applies to the *entire* un-embedded group, not just genuine similarity -/// ties). On any embed failure (e.g. `embedder.embed` erroring) `hits` is +/// Ordering: embedded hits sort before un-embedded hits; similarity ties among +/// embedded hits break on `time_range_end` DESC, while un-embedded hits retain +/// input order. On any embed failure (e.g. `embedder.embed` erroring) `hits` is /// returned as-is, in its original incoming order, with no decoration or /// sorting attempted. pub(crate) async fn rerank_by_semantic_similarity( @@ -48,28 +45,33 @@ pub(crate) async fn rerank_by_semantic_similarity( // Decorate each hit with (similarity, has_embedding). Un-embedded rows get // `NEG_INFINITY` so they sort last but keep their incoming relative order. - let mut decorated: Vec<(f32, bool, RetrievalHit)> = hits + let mut decorated: Vec<(f32, bool, usize, RetrievalHit)> = hits .into_iter() .zip(embeddings) - .map(|(h, emb)| match emb { + .enumerate() + .map(|(index, (h, emb))| match emb { Some(v) if v.len() == query_vec.len() => { let sim = cosine_similarity(&query_vec, &v); - (sim, true, h) + (sim, true, index, h) } - _ => (f32::NEG_INFINITY, false, h), + _ => (f32::NEG_INFINITY, false, index, h), }) .collect(); decorated.sort_by(|a, b| match (a.1, b.1) { (true, false) => std::cmp::Ordering::Less, (false, true) => std::cmp::Ordering::Greater, - // Both ranked (or both unranked): similarity DESC, then recency DESC. - _ => { + (false, false) => a.2.cmp(&b.2), + (true, true) => { b.0.partial_cmp(&a.0) .unwrap_or(std::cmp::Ordering::Equal) - .then_with(|| b.2.time_range_end.cmp(&a.2.time_range_end)) + .then_with(|| b.3.time_range_end.cmp(&a.3.time_range_end)) } }); - decorated.into_iter().map(|(_, _, h)| h).collect() + decorated.into_iter().map(|(_, _, _, h)| h).collect() } + +#[cfg(test)] +#[path = "rerank_tests.rs"] +mod tests; diff --git a/src/memory/retrieval/rerank_tests.rs b/src/memory/retrieval/rerank_tests.rs new file mode 100644 index 0000000..761e389 --- /dev/null +++ b/src/memory/retrieval/rerank_tests.rs @@ -0,0 +1,75 @@ +use anyhow::Result; +use async_trait::async_trait; +use chrono::{TimeZone, Utc}; + +use super::*; +use crate::memory::retrieval::types::NodeKind; +use crate::memory::tree::TreeKind; + +struct FixedEmbedder; + +#[async_trait] +impl Embedder for FixedEmbedder { + fn name(&self) -> &'static str { + "fixed" + } + + async fn embed(&self, _text: &str) -> Result> { + Ok(vec![1.0, 0.0]) + } +} + +fn hit(id: &str, timestamp_ms: i64) -> RetrievalHit { + let timestamp = Utc.timestamp_millis_opt(timestamp_ms).unwrap(); + RetrievalHit { + node_id: id.into(), + node_kind: NodeKind::Summary, + tree_id: "tree".into(), + tree_kind: TreeKind::Source, + tree_scope: "scope".into(), + level: 1, + content: id.into(), + entities: vec![], + topics: vec![], + time_range_start: timestamp, + time_range_end: timestamp, + score: 0.0, + child_ids: vec![], + source_ref: None, + } +} + +#[tokio::test] +async fn rerank_orders_embedded_hits_by_similarity_descending() { + let hits = vec![hit("opposite", 3), hit("same", 1), hit("orthogonal", 2)]; + let embeddings = vec![ + Some(vec![-1.0, 0.0]), + Some(vec![1.0, 0.0]), + Some(vec![0.0, 1.0]), + ]; + let ranked = rerank_by_semantic_similarity(&FixedEmbedder, "q", hits, embeddings).await; + let ids: Vec<_> = ranked.iter().map(|hit| hit.node_id.as_str()).collect(); + assert_eq!(ids, vec!["same", "orthogonal", "opposite"]); +} + +#[tokio::test] +async fn rerank_preserves_incoming_order_for_unembedded_tail() { + let hits = vec![ + hit("old-unembedded", 1), + hit("ranked", 2), + hit("new-unembedded", 3), + ]; + let embeddings = vec![None, Some(vec![1.0, 0.0]), None]; + let ranked = rerank_by_semantic_similarity(&FixedEmbedder, "q", hits, embeddings).await; + let ids: Vec<_> = ranked.iter().map(|hit| hit.node_id.as_str()).collect(); + assert_eq!(ids, vec!["ranked", "old-unembedded", "new-unembedded"]); +} + +#[tokio::test] +async fn rerank_treats_dimension_mismatch_as_unembedded() { + let hits = vec![hit("mismatch", 3), hit("ranked", 1)]; + let embeddings = vec![Some(vec![1.0]), Some(vec![1.0, 0.0])]; + let ranked = rerank_by_semantic_similarity(&FixedEmbedder, "q", hits, embeddings).await; + assert_eq!(ranked[0].node_id, "ranked"); + assert_eq!(ranked[1].node_id, "mismatch"); +} diff --git a/src/memory/retrieval/search.rs b/src/memory/retrieval/search.rs index 0cf75d4..310afb7 100644 --- a/src/memory/retrieval/search.rs +++ b/src/memory/retrieval/search.rs @@ -8,18 +8,12 @@ //! //! Matching rules (ported from OpenHuman's `memory_tree::retrieval::search`): //! - The query is lowercased before binding into the `LIKE` parameter. -//! - A row matches when `entity_id LIKE '%q%'` OR `surface LIKE '%q%'`. +//! - A row matches when `entity_id LIKE '%q%'` OR `surface LIKE '%q%'`, with +//! SQLite wildcard characters in `q` treated literally. //! - `kinds` narrows the match by `entity_kind IN (...)` when non-empty. //! - Output is ordered by mention count DESC (then recency) so the strongest //! matches surface first. //! -//! NOTE: `query`'s SQLite `LIKE` metacharacters (`%`, `_`) are NOT escaped -//! before being embedded in the `%…%` pattern — there is no `ESCAPE` clause. -//! A literal query of `"100%"` or `"_"` is interpreted as a wildcard, not a -//! literal string, and can match far more broadly than the caller intended -//! (up to "everything, capped at `limit`"). Escape `%`/`_`/the escape -//! character itself in `query` before calling if literal matching is -//! required. use anyhow::{Context, Result}; use rusqlite::params_from_iter; @@ -30,11 +24,6 @@ use crate::memory::score::extract::EntityKind; use super::types::EntityMatch; -/// Default result cap applied when the caller passes `limit = 0`. -const DEFAULT_LIMIT: usize = 5; -/// Hard ceiling on `limit` regardless of what the caller requests. -const MAX_LIMIT: usize = 100; - /// Search the entity index for canonical ids matching `query`. /// /// Returns at most `limit` matches (default 5, clamped to 100). Each match is @@ -43,9 +32,6 @@ const MAX_LIMIT: usize = 100; /// whitespace-only query returns no matches (rather than dumping the whole /// index via `LIKE '%%'`). /// -/// See the module-level NOTE: `query` is matched as a raw `LIKE` pattern with -/// no metacharacter escaping. -/// /// # Errors /// /// Returns `Err` on a SQLite statement-prepare or row-collection failure @@ -57,13 +43,13 @@ pub fn search_entities( kinds: Option<&[EntityKind]>, limit: usize, ) -> Result> { - let limit = normalise_limit(limit); + let limit = normalise_limit(config, limit); let query = query.trim(); if query.is_empty() { return Ok(Vec::new()); } - let q_lower = query.to_lowercase(); + let q_lower = escape_like_literal(&query.to_lowercase()); with_connection(config, |conn| { let pattern = format!("%{q_lower}%"); let (sql, params) = build_sql_and_params(&pattern, kinds, limit); @@ -78,12 +64,19 @@ pub fn search_entities( }) } -/// Clamp `limit` into `[1, MAX_LIMIT]`, substituting [`DEFAULT_LIMIT`] for `0`. -fn normalise_limit(limit: usize) -> usize { +fn escape_like_literal(query: &str) -> String { + query + .replace('\\', "\\\\") + .replace('%', "\\%") + .replace('_', "\\_") +} + +/// Apply the configured default and hard cap. +fn normalise_limit(config: &MemoryConfig, limit: usize) -> usize { if limit == 0 { - DEFAULT_LIMIT + config.retrieval.limits.search_default_limit } else { - limit.min(MAX_LIMIT) + limit.min(config.retrieval.limits.max_limit) } } @@ -103,7 +96,8 @@ fn build_sql_and_params( COUNT(*) AS mention_count, MAX(timestamp_ms) AS last_seen_ms FROM mem_tree_entity_index - WHERE (LOWER(entity_id) LIKE ?1 OR LOWER(surface) LIKE ?1)", + WHERE (LOWER(entity_id) LIKE ?1 ESCAPE '\\' + OR LOWER(surface) LIKE ?1 ESCAPE '\\')", ); let mut params: Vec = vec![Value::Text(pattern.to_string())]; diff --git a/src/memory/retrieval/search_tests.rs b/src/memory/retrieval/search_tests.rs index ed57f96..e8b5828 100644 --- a/src/memory/retrieval/search_tests.rs +++ b/src/memory/retrieval/search_tests.rs @@ -81,6 +81,40 @@ fn matches_on_surface_substring() { ); } +#[test] +fn like_metacharacters_are_matched_literally() { + let (_tmp, cfg) = test_config(); + index_entity_occurrence( + &cfg, + "topic:100%_ready", + EntityKind::Topic, + "100%_ready", + "leaf-literal", + "leaf", + 1_700_000_000_000, + None, + ); + index_entity_occurrence( + &cfg, + "topic:100x-ready", + EntityKind::Topic, + "100x-ready", + "leaf-wildcard-decoy", + "leaf", + 1_700_000_000_001, + None, + ); + + let matches = search_entities(&cfg, "100%_", None, 10).unwrap(); + assert_eq!(matches.len(), 1); + assert_eq!(matches[0].canonical_id, "topic:100%_ready"); +} + +#[test] +fn escape_like_literal_escapes_escape_and_wildcard_characters() { + assert_eq!(escape_like_literal(r"a\b%c_d"), r"a\\b\%c\_d"); +} + #[test] fn kind_filter_narrows_results() { let (_tmp, cfg) = test_config(); @@ -154,6 +188,7 @@ fn limit_truncates_results() { fn build_sql_without_kinds_has_no_in_clause() { let (sql, _params) = build_sql_and_params("%a%", None, 5); assert!(sql.contains("LOWER(entity_id) LIKE")); + assert!(sql.contains("ESCAPE '\\'")); assert!(!sql.contains("entity_kind IN")); } @@ -168,10 +203,18 @@ fn build_sql_with_kinds_adds_in_clause() { #[test] fn zero_limit_defaults_to_five() { - assert_eq!(normalise_limit(0), DEFAULT_LIMIT); + let config = MemoryConfig::new("/tmp/search-limit-test"); + assert_eq!( + normalise_limit(&config, 0), + config.retrieval.limits.search_default_limit + ); } #[test] fn huge_limit_is_clamped() { - assert_eq!(normalise_limit(10_000), MAX_LIMIT); + let config = MemoryConfig::new("/tmp/search-limit-test"); + assert_eq!( + normalise_limit(&config, 10_000), + config.retrieval.limits.max_limit + ); } diff --git a/src/memory/retrieval/source.rs b/src/memory/retrieval/source.rs index eb87301..32d1662 100644 --- a/src/memory/retrieval/source.rs +++ b/src/memory/retrieval/source.rs @@ -7,23 +7,18 @@ //! kind (chat / email / document). //! 3. Neither → every source tree. //! -//! For each tree we pull all level ≥ 1 summaries. With `time_window_days` we -//! keep only summaries whose `[time_range_start, time_range_end]` overlaps -//! `[now − window, now]`. When `query` is `Some`, hits are reranked by cosine +//! For each tree we pull level ≥ 1 summaries. With `time_window_days`, an SQL +//! predicate keeps only summaries whose `[time_range_start, time_range_end]` +//! overlaps `[now − window, now]` before embeddings are hydrated. When `query` +//! is `Some`, hits are reranked by cosine //! similarity between the query embedding and each summary's stored embedding; //! otherwise they are ordered newest-first by `time_range_end`. //! //! This is a thin read-only view over `mem_tree_trees` and `mem_tree_summaries` //! — no new indexes or tables. //! -//! NOTE: the `time_window_days` filter above is applied in Rust *after* every -//! summary (across every selected tree, at every level ≥ 1) has been loaded -//! and embedding-hydrated — see `collect_source_hits`. There is no -//! SQL-level window push-down, so a query like "last 7 days, limit 10" still -//! pays the full O(total summaries in scope) read + hydrate cost on a store -//! with years of history. -//! [`crate::memory::tree::store::list_summaries_in_window`] exists and does -//! push the window into SQL (used by [`super::cover`]) but is not used here. +//! Unwindowed queries still materialize all selected summaries; callers should +//! supply a window for bounded historical retrieval. use anyhow::Result; use chrono::{Duration, Utc}; @@ -32,16 +27,14 @@ use crate::memory::chunks::SourceKind; use crate::memory::config::MemoryConfig; use crate::memory::score::embed::Embedder; use crate::memory::tree::store::{ - get_summary_embeddings_batch, get_tree_by_scope, list_summaries_at_level, list_trees_by_kind, + get_summary_embeddings_batch, get_tree_by_scope, list_summaries_at_level, + list_summaries_overlapping_window, list_trees_by_kind, }; use crate::memory::tree::{Tree, TreeKind}; use super::rerank::rerank_by_semantic_similarity; use super::types::{hydrated_summary_hit, QueryResponse, RetrievalHit}; -/// Default result cap applied when the caller passes `limit = 0`. -const DEFAULT_LIMIT: usize = 10; - /// A summary hit paired with its (possibly sidecar-hydrated) embedding, kept /// together so window-filtering and reranking stay aligned. pub(crate) type ScoredHit = (RetrievalHit, Option>); @@ -61,14 +54,19 @@ pub async fn query_source( embedder: &dyn Embedder, limit: usize, ) -> Result { - let limit = if limit == 0 { DEFAULT_LIMIT } else { limit }; - - let mut scored = collect_source_hits(config, source_id, source_kind)?; - if let Some(days) = time_window_days { + let limits = &config.retrieval.limits; + let limit = if limit == 0 { + limits.default_limit + } else { + limit.min(limits.max_limit) + }; + + let window = time_window_days.map(|days| { let now = Utc::now(); let start = now - Duration::days(days as i64); - scored.retain(|(h, _)| h.time_range_end >= start && h.time_range_start <= now); - } + (start.timestamp_millis(), now.timestamp_millis()) + }); + let scored = collect_source_hits(config, source_id, source_kind, window)?; let total = scored.len(); let sorted = order_hits(scored, query, embedder).await; @@ -101,14 +99,14 @@ pub(crate) async fn order_hits( /// the selected source trees, hydrating each summary's embedding from the /// per-model sidecar when the legacy in-row column is empty. /// -/// No time or count filtering happens here — callers ([`query_source`], -/// [`super::global::query_global`]) apply their window filter afterwards, in -/// Rust, over the full result. This is O(number of non-deleted summaries +/// When `window_ms` is present, filtering happens in SQLite before embedding +/// hydration. Without a window this remains O(number of non-deleted summaries /// across the selected trees), unbounded by any `limit` argument. pub(crate) fn collect_source_hits( config: &MemoryConfig, source_id: Option<&str>, source_kind: Option, + window_ms: Option<(i64, i64)>, ) -> Result> { let trees = select_trees(config, source_id, source_kind)?; @@ -121,8 +119,8 @@ pub(crate) fn collect_source_hits( if tree.max_level == 0 && tree.root_id.is_none() { continue; } - for level in 1..=tree.max_level { - for node in list_summaries_at_level(config, &tree.id, level)? { + if let Some((since_ms, until_ms)) = window_ms { + for node in list_summaries_overlapping_window(config, &tree.id, since_ms, until_ms)? { if node.deleted { continue; } @@ -130,6 +128,17 @@ pub(crate) fn collect_source_hits( embeddings.push(node.embedding.clone()); hits.push(hydrated_summary_hit(config, &node, &tree.scope)); } + } else { + for level in 1..=tree.max_level { + for node in list_summaries_at_level(config, &tree.id, level)? { + if node.deleted { + continue; + } + node_ids.push(node.id.clone()); + embeddings.push(node.embedding.clone()); + hits.push(hydrated_summary_hit(config, &node, &tree.scope)); + } + } } } diff --git a/src/memory/retrieval/source_tests.rs b/src/memory/retrieval/source_tests.rs index 52fc3d0..343aded 100644 --- a/src/memory/retrieval/source_tests.rs +++ b/src/memory/retrieval/source_tests.rs @@ -184,7 +184,7 @@ async fn sidecar_embeddings_hydrate_for_rerank() { // collect_source_hits must surface the sidecar vector even though the // in-row embedding column is NULL. - let scored = collect_source_hits(&cfg, None, Some(SourceKind::Chat)).unwrap(); + let scored = collect_source_hits(&cfg, None, Some(SourceKind::Chat), None).unwrap(); assert_eq!(scored.len(), 1); assert_eq!(scored[0].1.as_deref(), Some(unit_vec(0).as_slice())); } diff --git a/src/memory/score/embed.rs b/src/memory/score/embed.rs index bb5e24b..02a9e60 100644 --- a/src/memory/score/embed.rs +++ b/src/memory/score/embed.rs @@ -2,12 +2,12 @@ //! //! Produces a fixed-dimension vector per chunk / summary so retrieval can //! rerank candidates by semantic similarity. The backend is abstracted behind -//! the [`Embedder`] trait; the crate ships only the deterministic -//! [`InertEmbedder`] (zero vectors) used by tests. Real network-backed backends +//! the `Embedder` trait; the crate ships only the deterministic +//! `InertEmbedder` (zero vectors) used by tests. Real network-backed backends //! (Ollama / OpenAI-compatible / cloud) are wired in by a host adapter that -//! implements [`Embedder`] — TinyCortex never makes a network call. +//! implements `Embedder` — TinyCortex never makes a network call. //! -//! Dimension is fixed at [`EMBEDDING_DIM`] (768, from +//! Dimension is fixed at `EMBEDDING_DIM` (768, from //! [`crate::memory::config::DEFAULT_EMBEDDING_DIM`]) — mixing dimensions //! mid-run would corrupt cosine comparisons, so we catch that at the trait //! level rather than deferring to retrieval-time diagnostics. diff --git a/src/memory/score/extract/llm_tests.rs b/src/memory/score/extract/llm_tests.rs index b61cc28..323f817 100644 --- a/src/memory/score/extract/llm_tests.rs +++ b/src/memory/score/extract/llm_tests.rs @@ -434,244 +434,5 @@ fn build_prompt_carries_user_text_and_kind_tag() { assert!(prompt.system.contains("\"importance\"")); } -#[test] -fn truncate_for_log_short_input_unchanged() { - assert_eq!(truncate_for_log("hi", 10), "hi"); -} - -#[test] -fn truncate_for_log_long_input_appends_ellipsis() { - let long = "x".repeat(500); - let out = truncate_for_log(&long, 10); - assert_eq!(out.chars().count(), 11); // 10 + "…" - assert!(out.ends_with('…')); -} - -#[tokio::test] -async fn extract_retries_on_truncated_response() { - // First response is truncated mid-JSON (serde EOF) — a stream cutoff, - // not a wrong-shape body. It must be treated as retryable rather than - // silently dropped; the second (complete) response then recovers the - // entities. - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct TruncatedThenCompleteProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for TruncatedThenCompleteProvider { - fn name(&self) -> &str { - "test:truncated" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - let n = self.calls.fetch_add(1, Ordering::SeqCst); - if n == 0 { - // Array never closes → serde reports EOF (is_eof()). - Ok(r#"{"entities":[{"kind":"person","text":"Alice"}"#.to_string()) - } else { - Ok(r#"{"entities":[{"kind":"person","text":"Alice"}],"importance":0.5,"importance_reason":"r"}"#.to_string()) - } - } - } - - let mock = Arc::new(TruncatedThenCompleteProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("Alice met Bob.").await.unwrap(); - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 2, - "truncation should trigger a retry, not a silent drop" - ); - assert_eq!(out.entities.len(), 1); - assert_eq!(out.entities[0].text, "Alice"); -} - -#[tokio::test] -async fn extract_does_not_retry_on_wrong_shape_response() { - // A *complete* but wrong-shape body is a serde error that is NOT EOF. - // Unlike a mid-JSON cutoff it's deterministic and won't fix itself on - // retry, so it must return immediately (`calls == 1`) with an empty - // extraction. - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct WrongShapeProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for WrongShapeProvider { - fn name(&self) -> &str { - "test:wrong-shape" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - // `entities` must be an array; a scalar is a complete-input type - // error (not a truncation), so serde reports a non-EOF error. - Ok(r#"{"entities":123}"#.to_string()) - } - } - - let mock = Arc::new(WrongShapeProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("Alice met Bob.").await.unwrap(); - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 1, - "a non-EOF wrong-shape response must not trigger a retry" - ); - assert!( - out.entities.is_empty(), - "wrong-shape response should yield an empty extraction" - ); -} - -#[test] -fn build_prompt_sets_extraction_max_tokens_cap() { - // Extraction must cap output tokens so a credit-metered provider prices - // the request against a realistic budget. build_prompt is the single - // source of that cap. - use async_trait::async_trait; - use std::sync::Arc; - - struct NoopProvider; - #[async_trait] - impl ChatProvider for NoopProvider { - fn name(&self) -> &str { - "test:noop" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - Ok("{}".into()) - } - } - - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), Arc::new(NoopProvider)); - let prompt = ex.build_prompt("hello"); - assert_eq!(prompt.max_tokens, Some(EXTRACTION_MAX_OUTPUT_TOKENS)); - assert_eq!(EXTRACTION_MAX_OUTPUT_TOKENS, 8192); -} - -#[tokio::test] -async fn extract_does_not_retry_on_permanent_402() { - // A 402 (the BYO provider account is out of credits) is a permanent client - // error: retrying reproduces it. extract() must call the provider exactly - // once and return empty. - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct InsufficientCreditsProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for InsufficientCreditsProvider { - fn name(&self) -> &str { - "test:402" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Err(anyhow::anyhow!( - "myopenrouter API error (402 Payment Required): This request requires more \ - credits, or fewer max_tokens. You requested up to 65536 tokens, but can only \ - afford 49732." - )) - } - } - - let mock = Arc::new(InsufficientCreditsProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("some text").await.unwrap(); - - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 1, - "a permanent 402 must not be retried" - ); - assert!(out.entities.is_empty()); -} - -#[tokio::test] -async fn extract_does_not_retry_on_500_wrapped_monthly_quota() { - // A 500-envelope wrapping an inner 402 monthly-quota refusal is still a - // permanent error (MONTHLY_REQUEST_COUNT) — it must call the provider - // exactly once. - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct MonthlyQuotaProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for MonthlyQuotaProvider { - fn name(&self) -> &str { - "test:kiro" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Err(anyhow::anyhow!( - "kiro API error (500 Internal Server Error): HTTP 402 from Kiro IDE: \ - You have reached the limit. reason: MONTHLY_REQUEST_COUNT" - )) - } - } - - let mock = Arc::new(MonthlyQuotaProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("some text").await.unwrap(); - - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 1, - "a 500-wrapped monthly-quota refusal must not be retried" - ); - assert!(out.entities.is_empty()); -} - -#[tokio::test] -async fn extract_retries_transient_provider_error() { - // A transport/transient failure (no 4xx, no auth marker) must still exhaust - // the retry budget before falling back. - use async_trait::async_trait; - use std::sync::atomic::{AtomicUsize, Ordering}; - use std::sync::Arc; - - struct TransientProvider { - calls: AtomicUsize, - } - #[async_trait] - impl ChatProvider for TransientProvider { - fn name(&self) -> &str { - "test:transient" - } - async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { - self.calls.fetch_add(1, Ordering::SeqCst); - Err(anyhow::anyhow!( - "error sending request for url (https://api): connection refused" - )) - } - } - - let mock = Arc::new(TransientProvider { - calls: AtomicUsize::new(0), - }); - let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); - let out = ex.extract("some text").await.unwrap(); - - assert_eq!( - mock.calls.load(Ordering::SeqCst), - 3, - "a transient error must still exhaust the retry budget" - ); - assert!(out.entities.is_empty()); -} +#[path = "llm_tests_more.rs"] +mod more; diff --git a/src/memory/score/extract/llm_tests_more.rs b/src/memory/score/extract/llm_tests_more.rs new file mode 100644 index 0000000..23b8b78 --- /dev/null +++ b/src/memory/score/extract/llm_tests_more.rs @@ -0,0 +1,243 @@ +use super::*; + +#[test] +fn truncate_for_log_short_input_unchanged() { + assert_eq!(truncate_for_log("hi", 10), "hi"); +} + +#[test] +fn truncate_for_log_long_input_appends_ellipsis() { + let long = "x".repeat(500); + let out = truncate_for_log(&long, 10); + assert_eq!(out.chars().count(), 11); // 10 + "…" + assert!(out.ends_with('…')); +} + +#[tokio::test] +async fn extract_retries_on_truncated_response() { + // First response is truncated mid-JSON (serde EOF) — a stream cutoff, + // not a wrong-shape body. It must be treated as retryable rather than + // silently dropped; the second (complete) response then recovers the + // entities. + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct TruncatedThenCompleteProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for TruncatedThenCompleteProvider { + fn name(&self) -> &str { + "test:truncated" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + let n = self.calls.fetch_add(1, Ordering::SeqCst); + if n == 0 { + // Array never closes → serde reports EOF (is_eof()). + Ok(r#"{"entities":[{"kind":"person","text":"Alice"}"#.to_string()) + } else { + Ok(r#"{"entities":[{"kind":"person","text":"Alice"}],"importance":0.5,"importance_reason":"r"}"#.to_string()) + } + } + } + + let mock = Arc::new(TruncatedThenCompleteProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("Alice met Bob.").await.unwrap(); + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 2, + "truncation should trigger a retry, not a silent drop" + ); + assert_eq!(out.entities.len(), 1); + assert_eq!(out.entities[0].text, "Alice"); +} + +#[tokio::test] +async fn extract_does_not_retry_on_wrong_shape_response() { + // A *complete* but wrong-shape body is a serde error that is NOT EOF. + // Unlike a mid-JSON cutoff it's deterministic and won't fix itself on + // retry, so it must return immediately (`calls == 1`) with an empty + // extraction. + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct WrongShapeProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for WrongShapeProvider { + fn name(&self) -> &str { + "test:wrong-shape" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + // `entities` must be an array; a scalar is a complete-input type + // error (not a truncation), so serde reports a non-EOF error. + Ok(r#"{"entities":123}"#.to_string()) + } + } + + let mock = Arc::new(WrongShapeProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("Alice met Bob.").await.unwrap(); + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 1, + "a non-EOF wrong-shape response must not trigger a retry" + ); + assert!( + out.entities.is_empty(), + "wrong-shape response should yield an empty extraction" + ); +} + +#[test] +fn build_prompt_sets_extraction_max_tokens_cap() { + // Extraction must cap output tokens so a credit-metered provider prices + // the request against a realistic budget. build_prompt is the single + // source of that cap. + use async_trait::async_trait; + use std::sync::Arc; + + struct NoopProvider; + #[async_trait] + impl ChatProvider for NoopProvider { + fn name(&self) -> &str { + "test:noop" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + Ok("{}".into()) + } + } + + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), Arc::new(NoopProvider)); + let prompt = ex.build_prompt("hello"); + assert_eq!(prompt.max_tokens, Some(EXTRACTION_MAX_OUTPUT_TOKENS)); + assert_eq!(EXTRACTION_MAX_OUTPUT_TOKENS, 8192); +} + +#[tokio::test] +async fn extract_does_not_retry_on_permanent_402() { + // A 402 (the BYO provider account is out of credits) is a permanent client + // error: retrying reproduces it. extract() must call the provider exactly + // once and return empty. + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct InsufficientCreditsProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for InsufficientCreditsProvider { + fn name(&self) -> &str { + "test:402" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(anyhow::anyhow!( + "myopenrouter API error (402 Payment Required): This request requires more \ + credits, or fewer max_tokens. You requested up to 65536 tokens, but can only \ + afford 49732." + )) + } + } + + let mock = Arc::new(InsufficientCreditsProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("some text").await.unwrap(); + + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 1, + "a permanent 402 must not be retried" + ); + assert!(out.entities.is_empty()); +} + +#[tokio::test] +async fn extract_does_not_retry_on_500_wrapped_monthly_quota() { + // A 500-envelope wrapping an inner 402 monthly-quota refusal is still a + // permanent error (MONTHLY_REQUEST_COUNT) — it must call the provider + // exactly once. + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct MonthlyQuotaProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for MonthlyQuotaProvider { + fn name(&self) -> &str { + "test:kiro" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(anyhow::anyhow!( + "kiro API error (500 Internal Server Error): HTTP 402 from Kiro IDE: \ + You have reached the limit. reason: MONTHLY_REQUEST_COUNT" + )) + } + } + + let mock = Arc::new(MonthlyQuotaProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("some text").await.unwrap(); + + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 1, + "a 500-wrapped monthly-quota refusal must not be retried" + ); + assert!(out.entities.is_empty()); +} + +#[tokio::test] +async fn extract_retries_transient_provider_error() { + // A transport/transient failure (no 4xx, no auth marker) must still exhaust + // the retry budget before falling back. + use async_trait::async_trait; + use std::sync::atomic::{AtomicUsize, Ordering}; + use std::sync::Arc; + + struct TransientProvider { + calls: AtomicUsize, + } + #[async_trait] + impl ChatProvider for TransientProvider { + fn name(&self) -> &str { + "test:transient" + } + async fn chat_for_json(&self, _p: &ChatPrompt) -> anyhow::Result { + self.calls.fetch_add(1, Ordering::SeqCst); + Err(anyhow::anyhow!( + "error sending request for url (https://api): connection refused" + )) + } + } + + let mock = Arc::new(TransientProvider { + calls: AtomicUsize::new(0), + }); + let ex = LlmEntityExtractor::new(LlmExtractorConfig::default(), mock.clone()); + let out = ex.extract("some text").await.unwrap(); + + assert_eq!( + mock.calls.load(Ordering::SeqCst), + 3, + "a transient error must still exhaust the retry budget" + ); + assert!(out.entities.is_empty()); +} diff --git a/src/memory/score/extract/mod.rs b/src/memory/score/extract/mod.rs index 8eb225a..f7da29b 100644 --- a/src/memory/score/extract/mod.rs +++ b/src/memory/score/extract/mod.rs @@ -1,9 +1,9 @@ //! Entity extraction (Phase 2 / #708). //! //! Exposes [`EntityExtractor`] as a pluggable interface and a default -//! [`CompositeExtractor`] that runs a chain of extractors and merges their +//! `CompositeExtractor` that runs a chain of extractors and merges their //! output. The mechanical regex extractor is always available; semantic NER -//! ([`LlmEntityExtractor`]) plugs in behind the [`llm::ChatProvider`] trait +//! (`LlmEntityExtractor`) plugs in behind the `ChatProvider` trait //! without changing any call sites — and the crate never calls a real model. #[path = "composite.rs"] diff --git a/src/memory/score/extract/regex.rs b/src/memory/score/extract/regex.rs index 88931e0..60d726f 100644 --- a/src/memory/score/extract/regex.rs +++ b/src/memory/score/extract/regex.rs @@ -28,8 +28,9 @@ static RE_URL: LazyLock = LazyLock::new(|| { Regex::new(r"https?://[^\s<>\]\[()]+[^\s<>\]\[()\.\,;:\!\?]").unwrap() }); -static RE_HANDLE: LazyLock = - LazyLock::new(|| Regex::new(r"(?:^|[\s(])@([A-Za-z0-9_][A-Za-z0-9_.\-]{1,})").unwrap()); +static RE_HANDLE: LazyLock = LazyLock::new(|| { + Regex::new(r"(?:^|[\s(])@([A-Za-z0-9_](?:[A-Za-z0-9_.\-]*[A-Za-z0-9_])?)").unwrap() +}); static RE_DISCRIM: LazyLock = LazyLock::new(|| Regex::new(r"\b([A-Za-z0-9_.\-]{2,32})#\d{4}\b").unwrap()); @@ -45,12 +46,6 @@ static RE_HASHTAG: LazyLock = /// Hashtag matches are additionally surfaced as [`ExtractedTopic`] rows so /// they can be promoted into the canonical entity stream by the resolver. /// -/// NOTE: `RE_HANDLE` allows `.`/`-` anywhere in the captured group, including -/// at the end, so trailing sentence punctuation folded into a handle (e.g. -/// "ping @alice." at end-of-sentence) is captured as part of the surface -/// form. `handle:alice.` and `handle:alice` then canonicalise to distinct -/// entities, splitting co-occurrence weight for what is really one person -/// (see audit finding RS-12). pub fn extract(text: &str) -> ExtractedEntities { let mut entities: Vec = Vec::new(); let mut topics: Vec = Vec::new(); diff --git a/src/memory/score/extract/regex_tests.rs b/src/memory/score/extract/regex_tests.rs index f386769..3665450 100644 --- a/src/memory/score/extract/regex_tests.rs +++ b/src/memory/score/extract/regex_tests.rs @@ -45,6 +45,18 @@ fn handle_vs_email_boundary() { assert_eq!(emails, vec!["alice@example.com"]); } +#[test] +fn handle_stops_before_trailing_sentence_punctuation() { + let output = extract("ping @alice. then ask @bob- and @carol_b"); + let handles: Vec<_> = output + .entities + .iter() + .filter(|entity| entity.kind == EntityKind::Handle) + .map(|entity| entity.text.as_str()) + .collect(); + assert_eq!(handles, vec!["alice", "bob", "carol_b"]); +} + #[test] fn discord_style_handle() { let o = extract("ping alice#1234"); diff --git a/src/memory/score/mod.rs b/src/memory/score/mod.rs index 2a2ec86..976b5a9 100644 --- a/src/memory/score/mod.rs +++ b/src/memory/score/mod.rs @@ -47,6 +47,7 @@ pub mod resolver; pub mod signals; /// Persistence for score rows and the entity occurrence index. pub mod store; +mod store_query; use std::sync::Arc; @@ -136,6 +137,19 @@ pub struct ScoringConfig { } impl ScoringConfig { + /// Construct the default regex extractor with policy values from the + /// engine's declarative configuration. + pub fn from_memory_config(config: &MemoryConfig) -> Self { + Self { + extractor: Arc::new(extract::CompositeExtractor::regex_only()), + weights: config.scoring.weights.clone(), + drop_threshold: config.scoring.drop_threshold, + llm_extractor: None, + definite_keep_threshold: config.scoring.definite_keep_threshold, + definite_drop_threshold: config.scoring.definite_drop_threshold, + } + } + /// Default: regex-only extractor, default weights, default threshold. pub fn default_regex_only() -> Self { Self { @@ -354,14 +368,8 @@ pub async fn score_chunks_fast(chunks: &[Chunk], cfg: &ScoringConfig) -> Result< /// written here: TinyCortex's `memory::graph` module is not yet ported, so this /// helper persists the score row and the entity index only. /// -/// NOTE (RS-13): the entity index is only cleared/re-indexed when -/// `result.kept` is `true`. If a chunk is scored a second time (e.g. a -/// re-score after content edit) and flips from kept to dropped, this function -/// takes the `!result.kept` branch and never calls -/// [`store::clear_entity_index_for_node`] — the entity-index rows from the -/// earlier, kept scoring are left behind as phantom rows pointing at a -/// now-dropped chunk. Callers that need drop-time cleanup must clear the -/// entity index themselves before calling this function. +/// Existing entity rows are always cleared before the new result is applied; +/// a kept-to-dropped re-score therefore cannot leave phantom topic/graph hits. pub fn persist_score( config: &MemoryConfig, result: &ScoreResult, @@ -371,22 +379,19 @@ pub fn persist_score( let row = score_row(result); store::upsert_score(config, &row)?; - if result.kept { - // Clear any stale entity-index rows for this chunk before re-indexing. - // INSERT OR REPLACE never deletes rows whose entity_id is no longer - // present in the new extraction — so a re-score that drops an entity - // would otherwise leave a phantom index row. - store::clear_entity_index_for_node(config, &result.chunk_id)?; - if !result.canonical_entities.is_empty() { - store::index_entities( - config, - &result.canonical_entities, - &result.chunk_id, - "leaf", - timestamp_ms, - tree_id, - )?; - } + // INSERT OR REPLACE never deletes rows whose entity_id is absent from the + // new extraction. Clear for both kept and dropped results so stale rows do + // not survive a re-score. + store::clear_entity_index_for_node(config, &result.chunk_id)?; + if result.kept && !result.canonical_entities.is_empty() { + store::index_entities( + config, + &result.canonical_entities, + &result.chunk_id, + "leaf", + timestamp_ms, + tree_id, + )?; } Ok(()) @@ -394,9 +399,8 @@ pub fn persist_score( /// Transactional variant of [`persist_score`] — writes the score row and /// entity-index rows on the caller's open [`Transaction`] so the persistence -/// is atomic with the surrounding ingest write. Same kept/clear-before-reindex -/// semantics as [`persist_score`], including the same phantom-row gotcha on a -/// kept→dropped re-score (see the RS-13 note there). +/// is atomic with the surrounding ingest write. It has the same unconditional +/// clear-before-optional-reindex semantics as [`persist_score`]. pub fn persist_score_tx( tx: &Transaction<'_>, result: &ScoreResult, @@ -406,19 +410,16 @@ pub fn persist_score_tx( let row = score_row(result); store::upsert_score_tx(tx, &row)?; - if result.kept { - // See persist_score for why we clear before re-indexing. - store::clear_entity_index_for_node_tx(tx, &result.chunk_id)?; - if !result.canonical_entities.is_empty() { - store::index_entities_tx( - tx, - &result.canonical_entities, - &result.chunk_id, - "leaf", - timestamp_ms, - tree_id, - )?; - } + store::clear_entity_index_for_node_tx(tx, &result.chunk_id)?; + if result.kept && !result.canonical_entities.is_empty() { + store::index_entities_tx( + tx, + &result.canonical_entities, + &result.chunk_id, + "leaf", + timestamp_ms, + tree_id, + )?; } Ok(()) diff --git a/src/memory/score/mod_tests.rs b/src/memory/score/mod_tests.rs index 8d8ec2e..eefb891 100644 --- a/src/memory/score/mod_tests.rs +++ b/src/memory/score/mod_tests.rs @@ -300,3 +300,29 @@ async fn llm_consulted_reports_full_total() { expected ); } + +#[tokio::test] +async fn kept_to_dropped_rescore_clears_entity_index() { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + let chunk = test_chunk("Contact alice@example.com about the Phoenix launch decision."); + let scoring = ScoringConfig::default_regex_only(); + let mut result = score_chunk(&chunk, &scoring).await.unwrap(); + assert!(!result.canonical_entities.is_empty()); + result.kept = true; + result.drop_reason = None; + persist_score(&config, &result, 1, None).unwrap(); + assert!(store::count_entity_index(&config).unwrap() > 0); + + result.kept = false; + result.drop_reason = Some("rescored below threshold".into()); + persist_score(&config, &result, 2, None).unwrap(); + + assert_eq!(store::count_entity_index(&config).unwrap(), 0); + assert!( + store::get_score(&config, &chunk.id) + .unwrap() + .unwrap() + .dropped + ); +} diff --git a/src/memory/score/signals/mod.rs b/src/memory/score/signals/mod.rs index 78ace04..07ec740 100644 --- a/src/memory/score/signals/mod.rs +++ b/src/memory/score/signals/mod.rs @@ -1,6 +1,6 @@ //! Score signals + weighted combine (Phase 2 / #708). //! -//! Each submodule computes one scoring signal in `[0.0, 1.0]`. [`combine`] +//! Each submodule computes one scoring signal in `[0.0, 1.0]`. `combine` //! aggregates them into a total score using per-signal weights. The output //! is still `[0.0, 1.0]` after normalisation by total weight. //! diff --git a/src/memory/score/signals/types.rs b/src/memory/score/signals/types.rs index 8a957db..d7bd7df 100644 --- a/src/memory/score/signals/types.rs +++ b/src/memory/score/signals/types.rs @@ -38,7 +38,8 @@ pub struct ScoreSignals { /// `llm_importance` defaults to `0.0` (disabled). Callers who configure an /// LLM extractor should bump it (typical: 2.0 — comparable to the /// metadata/source weights, well below the interaction-direct signal). -#[derive(Clone, Debug)] +#[derive(Clone, Debug, serde::Serialize, serde::Deserialize)] +#[serde(default)] pub struct SignalWeights { /// Multiplier for [`ScoreSignals::token_count`]. Default `1.0`. pub token_count: f32, diff --git a/src/memory/score/store.rs b/src/memory/score/store.rs index e45feb0..07d7ba9 100644 --- a/src/memory/score/store.rs +++ b/src/memory/score/store.rs @@ -30,7 +30,7 @@ //! The `mem_tree_score` table in TinyCortex does not carry `llm_importance` / //! `llm_importance_reason` columns, so those fields are admission-time signals //! only: the persisted row records the resulting `total`. They are kept on -//! [`ScoreRow`] / [`ScoreSignals`] for API compatibility and read back as their +//! `ScoreRow` / [`ScoreSignals`] for API compatibility and read back as their //! defaults (`0.0` / `None`). TinyCortex also has no identity registry, so the //! `is_user` column is always written as `0`. @@ -47,6 +47,11 @@ use crate::memory::score::extract::EntityKind; use crate::memory::score::resolver::CanonicalEntity; use crate::memory::score::signals::ScoreSignals; +pub use super::store_query::{ + count_entity_index, count_scores, list_entity_ids_for_node, lookup_entity, + lookup_entity_in_window, +}; + /// Resolve `is_user` for one canonical entity. /// /// TinyCortex has no Composio-style identity registry, so this is always @@ -229,6 +234,37 @@ pub fn get_scores_batch( }) } +/// Return the subset of `chunk_ids` whose latest persisted admission result is +/// dropped. Missing score rows are not considered dropped. +pub fn get_dropped_chunk_ids_batch( + config: &MemoryConfig, + chunk_ids: &[String], +) -> Result> { + if chunk_ids.is_empty() { + return Ok(std::collections::HashSet::new()); + } + with_connection(config, |conn| { + let mut out = std::collections::HashSet::new(); + for window in chunk_ids.chunks(MAX_FETCH_BATCH) { + let placeholders = (1..=window.len()) + .map(|index| format!("?{index}")) + .collect::>() + .join(","); + let sql = format!( + "SELECT chunk_id FROM mem_tree_score + WHERE dropped != 0 AND chunk_id IN ({placeholders})" + ); + let mut stmt = conn.prepare(&sql)?; + let params: Vec<&dyn rusqlite::ToSql> = + window.iter().map(|id| id as &dyn rusqlite::ToSql).collect(); + for row in stmt.query_map(params.as_slice(), |row| row.get::<_, String>(0))? { + out.insert(row?); + } + } + Ok(out) + }) +} + const ENTITY_INDEX_UPSERT_SQL: &str = "INSERT OR REPLACE INTO mem_tree_entity_index ( entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms, tree_id, is_user @@ -441,88 +477,6 @@ pub struct EntityHit { pub is_user: bool, } -/// Find all nodes indexed against `entity_id`, newest first. -pub fn lookup_entity( - config: &MemoryConfig, - entity_id: &str, - limit: Option, -) -> Result> { - // Clamp to i64::MAX before casting so callers can't wrap a large usize into - // a negative LIMIT and bypass it. - let limit = limit.unwrap_or(100).min(i64::MAX as usize) as i64; - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT entity_id, node_id, node_kind, entity_kind, surface, - score, timestamp_ms, tree_id, is_user - FROM mem_tree_entity_index - WHERE entity_id = ?1 - ORDER BY timestamp_ms DESC - LIMIT ?2", - )?; - let rows = stmt - .query_map(params![entity_id, limit], |row| { - let kind_s: String = row.get(3)?; - let entity_kind = EntityKind::parse(&kind_s).map_err(|e| { - rusqlite::Error::FromSqlConversionFailure( - 3, - rusqlite::types::Type::Text, - e.into(), - ) - })?; - let is_user_int: i32 = row.get(8)?; - Ok(EntityHit { - entity_id: row.get(0)?, - node_id: row.get(1)?, - node_kind: row.get(2)?, - entity_kind, - surface: row.get(4)?, - score: row.get(5)?, - timestamp_ms: row.get(6)?, - tree_id: row.get(7)?, - is_user: is_user_int != 0, - }) - })? - .collect::>>()?; - Ok(rows) - }) -} - -/// All distinct canonical entity ids associated with `node_id`, ordered by -/// score (desc) then recency. Used by topic-routing to pick which topic trees a -/// node should fan into. -pub fn list_entity_ids_for_node(config: &MemoryConfig, node_id: &str) -> Result> { - with_connection(config, |conn| { - let mut stmt = conn.prepare( - "SELECT DISTINCT entity_id - FROM mem_tree_entity_index - WHERE node_id = ?1 - ORDER BY score DESC, timestamp_ms DESC, entity_id ASC", - )?; - let rows = stmt - .query_map(params![node_id], |row| row.get::<_, String>(0))? - .collect::>>()?; - Ok(rows) - }) -} - -/// Count rows in the entity index (for tests / diagnostics). -pub fn count_entity_index(config: &MemoryConfig) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_entity_index", [], |r| { - r.get(0) - })?; - Ok(n.max(0) as u64) - }) -} - -/// Count score rows (for tests / diagnostics). -pub fn count_scores(config: &MemoryConfig) -> Result { - with_connection(config, |conn| { - let n: i64 = conn.query_row("SELECT COUNT(*) FROM mem_tree_score", [], |r| r.get(0))?; - Ok(n.max(0) as u64) - }) -} - #[cfg(test)] #[path = "store_tests.rs"] mod tests; diff --git a/src/memory/score/store_query.rs b/src/memory/score/store_query.rs new file mode 100644 index 0000000..3abab9d --- /dev/null +++ b/src/memory/score/store_query.rs @@ -0,0 +1,92 @@ +//! Read-only score and entity-index diagnostics. + +use anyhow::Result; +use rusqlite::params; + +use super::extract::EntityKind; +use super::store::EntityHit; +use crate::memory::chunks::with_connection; +use crate::memory::config::MemoryConfig; + +pub fn lookup_entity( + config: &MemoryConfig, + entity_id: &str, + limit: Option, +) -> Result> { + lookup_entity_in_window(config, entity_id, None, None, limit) +} + +pub fn lookup_entity_in_window( + config: &MemoryConfig, + entity_id: &str, + since_ms: Option, + until_ms: Option, + limit: Option, +) -> Result> { + let limit = limit.unwrap_or(100).min(i64::MAX as usize) as i64; + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT entity_id, node_id, node_kind, entity_kind, surface, + score, timestamp_ms, tree_id, is_user + FROM mem_tree_entity_index + WHERE entity_id = ?1 + AND (?2 IS NULL OR timestamp_ms >= ?2) + AND (?3 IS NULL OR timestamp_ms <= ?3) + ORDER BY timestamp_ms DESC LIMIT ?4", + )?; + let rows = stmt + .query_map(params![entity_id, since_ms, until_ms, limit], |row| { + let raw: String = row.get(3)?; + let entity_kind = EntityKind::parse(&raw).map_err(|error| { + rusqlite::Error::FromSqlConversionFailure( + 3, + rusqlite::types::Type::Text, + error.into(), + ) + })?; + Ok(EntityHit { + entity_id: row.get(0)?, + node_id: row.get(1)?, + node_kind: row.get(2)?, + entity_kind, + surface: row.get(4)?, + score: row.get(5)?, + timestamp_ms: row.get(6)?, + tree_id: row.get(7)?, + is_user: row.get::<_, i32>(8)? != 0, + }) + })? + .collect::>>()?; + Ok(rows) + }) +} + +pub fn list_entity_ids_for_node(config: &MemoryConfig, node_id: &str) -> Result> { + with_connection(config, |conn| { + let mut stmt = conn.prepare( + "SELECT DISTINCT entity_id FROM mem_tree_entity_index + WHERE node_id = ?1 + ORDER BY score DESC, timestamp_ms DESC, entity_id ASC", + )?; + let rows = stmt + .query_map(params![node_id], |row| row.get(0))? + .collect::>>()?; + Ok(rows) + }) +} + +pub fn count_entity_index(config: &MemoryConfig) -> Result { + count_table(config, "mem_tree_entity_index") +} + +pub fn count_scores(config: &MemoryConfig) -> Result { + count_table(config, "mem_tree_score") +} + +fn count_table(config: &MemoryConfig, table: &str) -> Result { + with_connection(config, |conn| { + let sql = format!("SELECT COUNT(*) FROM {table}"); + let count: i64 = conn.query_row(&sql, [], |row| row.get(0))?; + Ok(count.max(0) as u64) + }) +} diff --git a/src/memory/score/store_tests.rs b/src/memory/score/store_tests.rs index 2d40c70..e2dbb31 100644 --- a/src/memory/score/store_tests.rs +++ b/src/memory/score/store_tests.rs @@ -237,3 +237,88 @@ fn get_scores_batch_empty_input_and_missing_chunk_ids() { assert!((map.get("c1").copied().unwrap() - 0.7).abs() < 1e-6); assert!(!map.contains_key("ghost:no-such")); } + +#[test] +fn transactional_score_and_entity_helpers_commit_together() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + let row = sample_row("tx-chunk", false); + let entities = vec![sample_entity("alice"), sample_entity("bob")]; + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + upsert_score_tx(&tx, &row)?; + assert_eq!( + index_entities_tx(&tx, &entities, "tx-chunk", "leaf", 42, Some("tree"))?, + 2 + ); + assert_eq!(clear_entity_index_for_node_tx(&tx, "tx-chunk")?, 2); + assert_eq!( + index_entities_tx(&tx, &entities[..1], "tx-chunk", "leaf", 43, None)?, + 1 + ); + tx.commit()?; + Ok(()) + }) + .unwrap(); + + assert_eq!(get_score(&cfg, "tx-chunk").unwrap().unwrap().total, 0.7); + assert_eq!( + list_entity_ids_for_node(&cfg, "tx-chunk").unwrap(), + vec!["email:alice"] + ); +} + +#[test] +fn empty_transactional_entity_batches_are_noops() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + assert_eq!(index_entities_tx(&tx, &[], "node", "leaf", 0, None)?, 0); + assert_eq!( + index_summary_entity_ids_tx(&tx, &[], "node", 0.0, 0, None)?, + 0 + ); + tx.commit()?; + Ok(()) + }) + .unwrap(); + assert_eq!( + index_entities(&cfg, &[], "node", "leaf", 0, None).unwrap(), + 0 + ); +} + +#[test] +fn score_batch_reads_across_sql_parameter_windows() { + let (_tmp, cfg) = test_config(); + let ids: Vec = (0..=MAX_FETCH_BATCH) + .map(|index| format!("chunk-{index}")) + .collect(); + for id in &ids { + upsert_score(&cfg, &sample_row(id, false)).unwrap(); + } + + let scores = get_scores_batch(&cfg, &ids).unwrap(); + + assert_eq!(scores.len(), MAX_FETCH_BATCH + 1); + assert!(ids.iter().all(|id| scores.contains_key(id))); +} + +#[test] +fn summary_entity_without_colon_uses_whole_id_as_kind_and_surfaces_parse_error() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + index_summary_entity_ids_tx(&tx, &["unknown".into()], "summary", 1.0, 1, None)?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + + assert!(lookup_entity(&cfg, "unknown", None).is_err()); +} diff --git a/src/memory/types.rs b/src/memory/types.rs index 1c1870f..900db69 100644 --- a/src/memory/types.rs +++ b/src/memory/types.rs @@ -39,8 +39,7 @@ pub const GLOBAL_NAMESPACE: &str = "global"; /// Sync paths that ingest text from third-party services (Gmail / Slack / /// Notion / Composio / MCP / …) MUST set this to [`MemoryTaint::ExternalSync`] /// at write time so callers can refuse external-effect tools on tainted context. -#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)] -#[serde(rename_all = "snake_case")] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)] pub enum MemoryTaint { /// User-driven memory (chat, manual remember, internal heuristics). #[default] @@ -49,6 +48,25 @@ pub enum MemoryTaint { ExternalSync, } +impl Serialize for MemoryTaint { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_str(self.as_db_str()) + } +} + +impl<'de> Deserialize<'de> for MemoryTaint { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let raw = String::deserialize(deserializer)?; + Ok(Self::from_db_str(&raw)) + } +} + impl MemoryTaint { /// Serialised form used by the SQLite `memory_docs.taint` column. /// @@ -96,8 +114,7 @@ impl MemoryTaint { } /// Categories used to organize and filter memories by nature and lifecycle. -#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] -#[serde(rename_all = "snake_case")] +#[derive(Debug, Clone, PartialEq, Eq)] pub enum MemoryCategory { /// Long-term foundational facts, user preferences, permanent decisions. Core, @@ -109,26 +126,56 @@ pub enum MemoryCategory { Custom(String), } -/// Renders the wire label for the built-in variants (`"core"`, `"daily"`, -/// `"conversation"`) and the raw inner name for [`MemoryCategory::Custom`]. -/// -/// This is a human-readable rendering only and is **not** the same as the -/// `serde` wire format: `Custom(name)` serializes to JSON as -/// `{"custom": name}` (serde's default externally-tagged representation for a -/// newtype variant), whereas [`Display`](std::fmt::Display) renders it as bare -/// `name` with no `"custom"` wrapper. Do not round-trip through `Display` for -/// persistence or RPC — use `serde_json` (de)serialization instead. +/// The stable wire/display representation uses the built-in labels directly +/// and prefixes custom values with `custom:`. The prefix keeps +/// `Custom("core")` distinct from [`MemoryCategory::Core`] and makes Display, +/// serde, and [`std::str::FromStr`] true inverses. impl std::fmt::Display for MemoryCategory { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { match self { Self::Core => write!(f, "core"), Self::Daily => write!(f, "daily"), Self::Conversation => write!(f, "conversation"), - Self::Custom(name) => write!(f, "{name}"), + Self::Custom(name) => write!(f, "custom:{name}"), + } + } +} + +impl std::str::FromStr for MemoryCategory { + type Err = String; + + fn from_str(value: &str) -> Result { + match value { + "core" => Ok(Self::Core), + "daily" => Ok(Self::Daily), + "conversation" => Ok(Self::Conversation), + value if value.starts_with("custom:") && value.len() > "custom:".len() => { + Ok(Self::Custom(value["custom:".len()..].to_string())) + } + _ => Err(format!("unknown memory category: {value}")), } } } +impl Serialize for MemoryCategory { + fn serialize(&self, serializer: S) -> Result + where + S: serde::Serializer, + { + serializer.serialize_str(&self.to_string()) + } +} + +impl<'de> Deserialize<'de> for MemoryCategory { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + let value = String::deserialize(deserializer)?; + value.parse().map_err(serde::de::Error::custom) + } +} + /// A single stored memory entry with associated metadata. #[derive(Debug, Clone, Serialize, Deserialize)] pub struct MemoryEntry { diff --git a/src/memory/types_tests.rs b/src/memory/types_tests.rs index 3854239..1fe1a52 100644 --- a/src/memory/types_tests.rs +++ b/src/memory/types_tests.rs @@ -15,7 +15,7 @@ fn memory_category_display_outputs_expected_values() { assert_eq!(MemoryCategory::Conversation.to_string(), "conversation"); assert_eq!( MemoryCategory::Custom("project_notes".into()).to_string(), - "project_notes" + "custom:project_notes" ); } @@ -33,6 +33,28 @@ fn memory_category_serde_uses_snake_case() { serde_json::to_string(&MemoryCategory::Conversation).unwrap(), "\"conversation\"" ); + assert_eq!( + serde_json::to_string(&MemoryCategory::Custom("core".into())).unwrap(), + "\"custom:core\"" + ); + for category in [ + MemoryCategory::Core, + MemoryCategory::Daily, + MemoryCategory::Conversation, + MemoryCategory::Custom("core".into()), + MemoryCategory::Custom("tool_memory".into()), + ] { + assert_eq!( + category.to_string().parse::().unwrap(), + category + ); + let json = serde_json::to_string(&category).unwrap(); + assert_eq!( + serde_json::from_str::(&json).unwrap(), + category + ); + } + assert!("project_notes".parse::().is_err()); } #[test] @@ -90,6 +112,18 @@ fn memory_taint_db_str_roundtrip_and_fails_closed() { ); } +#[test] +fn memory_taint_serde_unknown_values_fail_closed() { + assert_eq!( + serde_json::from_str::("\"unexpected\"").unwrap(), + MemoryTaint::ExternalSync + ); + assert_eq!( + serde_json::to_string(&MemoryTaint::ExternalSync).unwrap(), + "\"external_sync\"" + ); +} + #[test] fn memory_item_kind_serde_uses_snake_case() { assert_eq!( From f474971e29c5e5010447bc6b126a68f68ad0fcb2 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 07:17:18 +0000 Subject: [PATCH 3/7] fix: make ingest queue and tree transitions durable --- src/memory/archivist/store.rs | 68 +-- src/memory/archivist/store_tests.rs | 53 +++ src/memory/ingest/canonicalize/chat.rs | 28 +- src/memory/ingest/canonicalize/chat_tests.rs | 41 ++ src/memory/ingest/canonicalize/email.rs | 48 +- src/memory/ingest/canonicalize/email_clean.rs | 4 +- src/memory/ingest/canonicalize/email_tests.rs | 28 ++ src/memory/ingest/canonicalize/mod.rs | 41 +- src/memory/ingest/extract/extract_tests.rs | 32 +- src/memory/ingest/extract/mod.rs | 14 +- src/memory/ingest/mod.rs | 2 +- src/memory/ingest/pipeline.rs | 198 ++++---- src/memory/ingest/pipeline_tests.rs | 55 ++- src/memory/ingest/types.rs | 33 +- src/memory/queue/gate.rs | 5 +- src/memory/queue/handlers.rs | 124 +++-- src/memory/queue/mod.rs | 1 + src/memory/queue/ops.rs | 20 +- src/memory/queue/payloads.rs | 148 ++++++ src/memory/queue/runtime/mod.rs | 39 +- src/memory/queue/runtime/tests.rs | 64 +++ src/memory/queue/store.rs | 15 +- src/memory/queue/store_settle.rs | 83 ++-- src/memory/queue/store_settle_tests.rs | 59 +++ src/memory/queue/types.rs | 194 +------- src/memory/queue/worker.rs | 93 +++- src/memory/queue/worker_tests.rs | 71 +++ src/memory/sync/composio/client.rs | 1 - src/memory/sync/composio/connect_tests.rs | 157 +++++++ src/memory/sync/composio/providers/mod.rs | 1 + src/memory/sync/composio/providers/slack.rs | 76 +-- .../sync/composio/providers/slack_parse.rs | 81 ++++ src/memory/sync/dispatcher.rs | 78 +++- src/memory/sync/periodic.rs | 7 +- src/memory/sync/rebuild.rs | 2 + src/memory/tree/bucket_seal.rs | 431 +----------------- src/memory/tree/bucket_seal_tests.rs | 52 +++ src/memory/tree/document_seal.rs | 304 ++++++++++++ src/memory/tree/flavoured.rs | 7 +- src/memory/tree/flush.rs | 15 +- src/memory/tree/flush_tests.rs | 69 +++ src/memory/tree/label_resolver.rs | 51 +++ src/memory/tree/mod.rs | 10 +- src/memory/tree/runtime/engine.rs | 244 ++++------ src/memory/tree/runtime/engine_tests.rs | 53 +++ src/memory/tree/runtime/fold.rs | 77 ++++ src/memory/tree/runtime/mod.rs | 2 + src/memory/tree/runtime/rebuild_fs.rs | 162 +++++++ src/memory/tree/runtime/rebuild_fs_tests.rs | 48 ++ src/memory/tree/runtime/store/buffer.rs | 26 +- src/memory/tree/runtime/store/mod.rs | 9 +- src/memory/tree/runtime/store/nodes.rs | 20 +- src/memory/tree/runtime/store/paths.rs | 31 +- src/memory/tree/runtime/store/store_tests.rs | 33 ++ src/memory/tree/store/hotness_tests.rs | 65 +++ src/memory/tree/store/mod.rs | 2 +- src/memory/tree/store/store_tests.rs | 93 +--- src/memory/tree/store/store_tests_more.rs | 93 ++++ src/memory/tree/store/summaries.rs | 25 + src/memory/tree/store/types.rs | 14 +- src/memory/tree/summarise.rs | 16 +- src/memory/tree/summarise_tests.rs | 6 +- src/memory/tree/types.rs | 82 ++++ 63 files changed, 2679 insertions(+), 1325 deletions(-) create mode 100644 src/memory/queue/payloads.rs create mode 100644 src/memory/sync/composio/providers/slack_parse.rs create mode 100644 src/memory/tree/document_seal.rs create mode 100644 src/memory/tree/label_resolver.rs create mode 100644 src/memory/tree/runtime/fold.rs create mode 100644 src/memory/tree/runtime/rebuild_fs.rs create mode 100644 src/memory/tree/runtime/rebuild_fs_tests.rs create mode 100644 src/memory/tree/store/store_tests_more.rs create mode 100644 src/memory/tree/types.rs diff --git a/src/memory/archivist/store.rs b/src/memory/archivist/store.rs index 7373c98..4c89eb1 100644 --- a/src/memory/archivist/store.rs +++ b/src/memory/archivist/store.rs @@ -6,7 +6,7 @@ //! ``` //! //! Writes use the same atomic tempfile + rename contract as -//! [`write_if_new`](crate::memory::store::content::atomic::write_if_new), with +//! [`write_if_new`], with //! one important difference: we want to *append* turns to a session, so the seq //! is computed from the existing directory contents on each call. //! @@ -19,6 +19,7 @@ use std::fs; use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; +use sha2::{Digest, Sha256}; use crate::memory::archivist::types::ArchivedTurn; use crate::memory::config::MemoryConfig; @@ -41,7 +42,8 @@ fn session_dir(config: &MemoryConfig, session_id: &str) -> PathBuf { /// Map any non-`[A-Za-z0-9_-]` character to `_` so a session id is always a /// safe single path component. fn sanitize_session(s: &str) -> String { - s.chars() + let sanitized: String = s + .chars() .map(|c| { if c.is_alphanumeric() || c == '-' || c == '_' { c @@ -49,7 +51,19 @@ fn sanitize_session(s: &str) -> String { '_' } }) - .collect() + .collect(); + if sanitized == s && !sanitized.is_empty() { + return sanitized; + } + + // Replacement alone is collision-prone (`a/b` and `a?b`). Keep a short + // digest of the exact machine id whenever sanitisation changed it. + let digest = Sha256::digest(s.as_bytes()); + let suffix = digest[..6] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + format!("{sanitized}-{suffix}") } /// Next free sequence number for a session: one past the highest `NNNNNN.md` @@ -75,10 +89,10 @@ fn next_seq(dir: &Path) -> u32 { /// Render a turn as a YAML front-matter block followed by its body. fn compose_turn(turn: &ArchivedTurn) -> String { let mut yaml = String::from("---\n"); - yaml.push_str(&format!("session_id: {}\n", turn.session_id)); + yaml.push_str(&format!("session_id: {}\n", yaml_escape(&turn.session_id))); yaml.push_str(&format!("seq: {}\n", turn.seq)); yaml.push_str(&format!("timestamp_ms: {}\n", turn.timestamp_ms)); - yaml.push_str(&format!("role: {}\n", turn.role)); + yaml.push_str(&format!("role: {}\n", yaml_escape(&turn.role))); yaml.push_str(&format!("cost_microdollars: {}\n", turn.cost_microdollars)); if let Some(lesson) = turn.lesson.as_ref() { yaml.push_str(&format!("lesson: {}\n", yaml_escape(lesson))); @@ -94,18 +108,11 @@ fn compose_turn(turn: &ArchivedTurn) -> String { yaml } -/// Quote any string that contains characters with YAML semantic meaning. Simple -/// double-quote escaping is good enough for these single-line front-matter -/// fields. +/// Encode a string as a JSON string literal. JSON quoted scalars are valid +/// YAML, and serde's encoder correctly escapes newlines and every control +/// character that must not appear literally in front matter. fn yaml_escape(s: &str) -> String { - let needs_quote = s - .chars() - .any(|c| matches!(c, ':' | '#' | '\n' | '"' | '\'' | '[' | ']' | '{' | '}')); - if needs_quote { - format!("\"{}\"", s.replace('\\', "\\\\").replace('"', "\\\"")) - } else { - s.to_string() - } + serde_json::to_string(s).expect("serializing a string cannot fail") } /// Append a turn to its session's archive. Returns the assigned `seq`. @@ -115,12 +122,23 @@ fn yaml_escape(s: &str) -> String { pub fn record_turn(config: &MemoryConfig, mut turn: ArchivedTurn) -> Result { let dir = session_dir(config, &turn.session_id); fs::create_dir_all(&dir).with_context(|| format!("failed to mkdir -p {}", dir.display()))?; - turn.seq = next_seq(&dir); - let path = dir.join(format!("{:06}.md", turn.seq)); - let bytes = compose_turn(&turn).into_bytes(); - write_if_new(&path, &bytes) - .with_context(|| format!("failed to write episodic turn {}", path.display()))?; - Ok(turn) + loop { + turn.seq = next_seq(&dir); + let path = dir.join(format!("{:06}.md", turn.seq)); + let bytes = compose_turn(&turn).into_bytes(); + match write_if_new(&path, &bytes) { + Ok(true) => return Ok(turn), + // Another writer may have claimed the sequence after `next_seq`. + // `write_if_new` uses create-new semantics, so retrying computes a + // fresh sequence without ever overwriting the winning turn. + Ok(false) => continue, + Err(_) if path.exists() => continue, + Err(err) => { + return Err(err) + .with_context(|| format!("failed to write episodic turn {}", path.display())) + } + } + } } /// Read every turn for `session_id`, sorted by seq ascending. A missing session @@ -167,11 +185,7 @@ fn parse_turn(text: &str) -> Option { }; let k = k.trim(); let v = v.trim(); - let v_unquoted = v - .strip_prefix('"') - .and_then(|s| s.strip_suffix('"')) - .map(|s| s.replace("\\\"", "\"").replace("\\\\", "\\")) - .unwrap_or_else(|| v.to_string()); + let v_unquoted = serde_json::from_str::(v).unwrap_or_else(|_| v.to_string()); match k { "session_id" => turn.session_id = v_unquoted, "seq" => turn.seq = v_unquoted.parse().unwrap_or(0), diff --git a/src/memory/archivist/store_tests.rs b/src/memory/archivist/store_tests.rs index c7ffeb3..4c34d5f 100644 --- a/src/memory/archivist/store_tests.rs +++ b/src/memory/archivist/store_tests.rs @@ -53,6 +53,34 @@ fn append_increments_seq() { assert_eq!(read[2].content, "three"); } +#[test] +fn concurrent_record_turn_retries_sequence_collisions_without_loss() { + let (_tmp, cfg) = test_config(); + let writers = 24; + let barrier = std::sync::Arc::new(std::sync::Barrier::new(writers)); + let mut threads = Vec::new(); + for index in 0..writers { + let cfg = cfg.clone(); + let barrier = barrier.clone(); + threads.push(std::thread::spawn(move || { + barrier.wait(); + record_turn(&cfg, turn("shared", "user", &format!("turn-{index}"))).unwrap() + })); + } + let mut assigned = Vec::new(); + for thread in threads { + assigned.push(thread.join().unwrap().seq); + } + assigned.sort_unstable(); + assert_eq!(assigned, (0..writers as u32).collect::>()); + + let entries = session_entries(&cfg, "shared").unwrap(); + assert_eq!(entries.len(), writers); + let contents: std::collections::HashSet<_> = + entries.into_iter().map(|entry| entry.content).collect(); + assert_eq!(contents.len(), writers); +} + #[test] fn missing_session_returns_empty() { let (_tmp, cfg) = test_config(); @@ -79,6 +107,31 @@ fn preserves_lesson_and_tool_calls() { assert_eq!(read[0].cost_microdollars, 1234); } +#[test] +fn front_matter_round_trips_multiline_and_delimiter_like_scalars() { + let (_tmp, cfg) = test_config(); + let mut t = turn("session:one", "assistant\nadmin", "body stays separate"); + t.lesson = Some("first line\n---\nsecond: line\\tail".into()); + record_turn(&cfg, t.clone()).unwrap(); + + let read = session_entries(&cfg, &t.session_id).unwrap(); + assert_eq!(read.len(), 1); + assert_eq!(read[0].session_id, t.session_id); + assert_eq!(read[0].role, t.role); + assert_eq!(read[0].lesson, t.lesson); + assert_eq!(read[0].content, t.content); +} + +#[test] +fn unsafe_session_ids_have_collision_resistant_directories() { + let (_tmp, cfg) = test_config(); + record_turn(&cfg, turn("a/b", "user", "slash")).unwrap(); + record_turn(&cfg, turn("a?b", "user", "question")).unwrap(); + + assert_eq!(session_entries(&cfg, "a/b").unwrap()[0].content, "slash"); + assert_eq!(session_entries(&cfg, "a?b").unwrap()[0].content, "question"); +} + #[test] fn distinct_sessions_dont_mix() { let (_tmp, cfg) = test_config(); diff --git a/src/memory/ingest/canonicalize/chat.rs b/src/memory/ingest/canonicalize/chat.rs index a59f889..c6dce28 100644 --- a/src/memory/ingest/canonicalize/chat.rs +++ b/src/memory/ingest/canonicalize/chat.rs @@ -14,14 +14,8 @@ //! Reply body here. //! ``` //! -//! NOTE: `author` and `text` are rendered verbatim into the boundary line -//! (`## — `) with no escaping (known gap, see audit finding -//! QI-14 in `docs/spec/audit/04-queue-ingest.md`). A message body containing a -//! line that itself starts with `## ` is structurally indistinguishable from a -//! new message boundary and will be split into a bogus extra chunk by the -//! downstream Markdown chunker. Callers that accept untrusted message text -//! should sanitise or escape leading `## ` sequences before calling -//! [`canonicalise`] if boundary-splitting must be prevented. +//! Header newlines are collapsed, and body lines beginning with `## ` are +//! escaped so untrusted content cannot forge the chunker's message boundary. use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -82,11 +76,25 @@ pub fn canonicalise( // belongs in the MD front-matter. The chunker splits this output at `## ` // boundaries so each message becomes one chunk. for msg in &messages { + let author = msg.author.replace(['\n', '\r'], " "); + let body = msg + .text + .trim() + .lines() + .map(|line| { + if line.starts_with("## ") { + format!("\\{line}") + } else { + line.to_string() + } + }) + .collect::>() + .join("\n"); md.push_str(&format!( "## {} — {}\n{}\n\n", msg.timestamp.to_rfc3339(), - msg.author, - msg.text.trim() + author, + body )); } diff --git a/src/memory/ingest/canonicalize/chat_tests.rs b/src/memory/ingest/canonicalize/chat_tests.rs index a95adc0..1b55ff2 100644 --- a/src/memory/ingest/canonicalize/chat_tests.rs +++ b/src/memory/ingest/canonicalize/chat_tests.rs @@ -152,3 +152,44 @@ fn timestamp_numeric_string_accepted() { let msg: ChatMessage = serde_json::from_str(json).expect("numeric string should parse"); assert_eq!(msg.timestamp.timestamp_millis(), 1_700_000_000_000); } + +#[test] +fn timestamp_epoch_seconds_are_rejected_instead_of_treated_as_millis() { + for timestamp in [ + serde_json::json!(1_700_000_000), + serde_json::json!("1700000000"), + ] { + let value = serde_json::json!({ + "author": "alice", + "text": "hello", + "timestamp": timestamp, + }); + let err = serde_json::from_value::(value).unwrap_err(); + assert!(err.to_string().contains("milliseconds, not seconds")); + } +} + +#[test] +fn message_content_cannot_inject_chat_boundaries() { + let batch = ChatBatch { + platform: "slack".into(), + channel_label: "eng".into(), + messages: vec![ChatMessage { + author: "alice\n## forged-author".into(), + timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + text: "hello\n## forged message\nbye".into(), + source_ref: None, + }], + }; + let out = canonicalise("channel", "owner", &[], batch) + .unwrap() + .unwrap(); + assert_eq!( + out.markdown + .lines() + .filter(|line| line.starts_with("## ")) + .count(), + 1 + ); + assert!(out.markdown.contains("\\## forged message")); +} diff --git a/src/memory/ingest/canonicalize/email.rs b/src/memory/ingest/canonicalize/email.rs index b370b16..4ab2028 100644 --- a/src/memory/ingest/canonicalize/email.rs +++ b/src/memory/ingest/canonicalize/email.rs @@ -7,16 +7,9 @@ //! rendering to strip reply chains, marketing footers, legal disclaimers, and //! other boilerplate. //! -//! NOTE: `from`/`to`/`cc`/`subject` headers are rendered verbatim (no -//! escaping — [`email_clean::md_escape`] exists but is not applied here) into -//! the `---\nFrom: ...` boundary block the downstream chunker splits on -//! (known gap, see audit finding QI-14 in -//! `docs/spec/audit/04-queue-ingest.md`). A body or header value containing a -//! line of the form `---` immediately followed by a `From:` line can inject a -//! bogus message boundary and split one email into multiple spurious chunks. -//! Callers that accept untrusted header/body content should sanitise these -//! fields before calling [`canonicalise`] if boundary injection must be -//! prevented. +//! Header values collapse newlines through [`email_clean::md_escape`], and +//! body lines equal to `---` are escaped so untrusted content cannot forge the +//! downstream email boundary grammar. use chrono::{DateTime, Utc}; use serde::{Deserialize, Serialize}; @@ -89,25 +82,48 @@ pub fn canonicalise( // boundaries so each message becomes one chunk. for msg in &messages { md.push_str("---\n"); - md.push_str(&format!("From: {}\n", msg.from)); + md.push_str(&format!("From: {}\n", email_clean::md_escape(&msg.from))); if !msg.to.is_empty() { - md.push_str(&format!("To: {}\n", msg.to.join(", "))); + md.push_str(&format!( + "To: {}\n", + email_clean::md_escape(&msg.to.join(", ")) + )); } if !msg.cc.is_empty() { - md.push_str(&format!("Cc: {}\n", msg.cc.join(", "))); + md.push_str(&format!( + "Cc: {}\n", + email_clean::md_escape(&msg.cc.join(", ")) + )); } - md.push_str(&format!("Subject: {}\n", msg.subject)); + md.push_str(&format!( + "Subject: {}\n", + email_clean::md_escape(&msg.subject) + )); md.push_str(&format!("Date: {}\n", msg.sent_at.to_rfc3339())); if let Some(unsub) = &msg.list_unsubscribe { - md.push_str(&format!("List-Unsubscribe: {}\n", unsub)); + md.push_str(&format!( + "List-Unsubscribe: {}\n", + email_clean::md_escape(unsub) + )); } md.push('\n'); let cleaned = email_clean::clean_body(msg.body.trim()); if cleaned.is_empty() { md.push('\n'); } else { - md.push_str(&cleaned); + let safe_body = cleaned + .lines() + .map(|line| { + if line.trim_end() == "---" { + format!("\\{line}") + } else { + line.to_string() + } + }) + .collect::>() + .join("\n"); + md.push_str(&safe_body); } md.push_str("\n\n"); } diff --git a/src/memory/ingest/canonicalize/email_clean.rs b/src/memory/ingest/canonicalize/email_clean.rs index 2a6ea3a..c48f84a 100644 --- a/src/memory/ingest/canonicalize/email_clean.rs +++ b/src/memory/ingest/canonicalize/email_clean.rs @@ -19,7 +19,7 @@ use serde_json::Value; /// message we already render directly above. /// 2. **Drop footer noise** — `Unsubscribe`, `View in browser`, copyright /// lines, legal disclaimers, and address blocks. We cut at the first line -/// containing any of [`FOOTER_TRIGGERS`]. +/// containing a known footer trigger. /// /// The two passes run in order so a quoted-chain preamble below a /// "view in browser" line still gets stripped on its own merits even if the @@ -108,7 +108,7 @@ pub fn drop_reply_chain(s: &str) -> String { } /// Strip everything from the first line containing a footer trigger onward. -/// See [`FOOTER_TRIGGERS`] for the matched list. +/// Uses the module's known footer-trigger list. pub fn drop_footer_noise(s: &str) -> String { let mut offset = 0usize; for line in s.split_inclusive('\n') { diff --git a/src/memory/ingest/canonicalize/email_tests.rs b/src/memory/ingest/canonicalize/email_tests.rs index 061cbd0..a7daffd 100644 --- a/src/memory/ingest/canonicalize/email_tests.rs +++ b/src/memory/ingest/canonicalize/email_tests.rs @@ -176,3 +176,31 @@ fn sent_at_numeric_string_accepted() { let msg: EmailMessage = serde_json::from_str(json).expect("numeric string should parse"); assert_eq!(msg.sent_at.timestamp_millis(), 1_700_000_000_000); } + +#[test] +fn headers_and_body_cannot_inject_email_boundaries() { + let thread = EmailThread { + provider: "gmail".into(), + thread_subject: "thread".into(), + messages: vec![EmailMessage { + from: "alice@example.com\n---\nFrom: attacker@example.com".into(), + to: vec!["bob@example.com\nFrom: forged@example.com".into()], + cc: vec![], + subject: "hello\n---\nFrom: forged@example.com".into(), + sent_at: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + body: "real body\n---\nFrom: body-forgery@example.com\nmore".into(), + source_ref: None, + list_unsubscribe: None, + }], + }; + let out = canonicalise("thread", "owner", &[], thread) + .unwrap() + .unwrap(); + assert_eq!( + out.markdown.lines().filter(|line| *line == "---").count(), + 1 + ); + assert!(out + .markdown + .contains("\\---\nFrom: body-forgery@example.com")); +} diff --git a/src/memory/ingest/canonicalize/mod.rs b/src/memory/ingest/canonicalize/mod.rs index 026aeb3..e2d7a7b 100644 --- a/src/memory/ingest/canonicalize/mod.rs +++ b/src/memory/ingest/canonicalize/mod.rs @@ -3,7 +3,7 @@ //! //! Each source kind has its own adapter. They all return the same shape: //! a [`CanonicalisedSource`] containing the markdown blob plus a seed -//! [`Metadata`](crate::memory::chunks::Metadata) that the chunker will clone +//! [`Metadata`] that the chunker will clone //! onto each produced chunk. //! //! Adapters do not interpret content semantically — they only normalise @@ -28,20 +28,6 @@ use crate::memory::chunks::{Metadata, SourceRef}; /// On an unparseable string a serde error is returned (no silent default). /// Shared across chat, email, and document canonicalisers. /// -/// NOTE: known gap (audit finding QI-15 in -/// `docs/spec/audit/04-queue-ingest.md`) — both the `Millis` branch and the -/// decimal-string fallback accept **any** parseable `i64` as epoch -/// milliseconds with no range check. A caller that accidentally sends -/// epoch-**seconds** (~10 orders of magnitude smaller) has that value silently -/// accepted and interpreted as epoch-milliseconds, producing a timestamp near -/// the Unix epoch (1970) instead of a serde error. The resulting timestamp -/// feeds `timestamp` / `time_range` on the [`Metadata`] each canonicaliser -/// seeds, which in turn drive tree ordering and flush-staleness math — a -/// poisoned timestamp here can misorder chunks or make a fresh chunk look -/// arbitrarily stale. Callers -/// should range-check values below roughly `1e11` (below year ~1973 in -/// milliseconds) before they reach this deserializer if they cannot guarantee -/// the unit. pub(crate) fn deserialize_flexible_timestamp<'de, D>( deserializer: D, ) -> Result, D::Error> @@ -55,21 +41,30 @@ where Text(String), } + fn epoch_millis(ms: i64) -> Result, E> { + // Contemporary epoch seconds are ten digits while epoch milliseconds + // are thirteen. Reject the ambiguous near-epoch range so a seconds + // value cannot silently poison ordering and staleness calculations. + const MIN_PLAUSIBLE_EPOCH_MILLIS: u64 = 100_000_000_000; + if ms.unsigned_abs() < MIN_PLAUSIBLE_EPOCH_MILLIS { + return Err(E::custom(format!( + "epoch-ms value {ms} is too small; pass milliseconds, not seconds" + ))); + } + chrono::TimeZone::timestamp_millis_opt(&Utc, ms) + .single() + .ok_or_else(|| E::custom(format!("invalid epoch-ms: {ms}"))) + } + let raw = RawTs::deserialize(deserializer)?; match raw { - RawTs::Millis(ms) => chrono::TimeZone::timestamp_millis_opt(&Utc, ms) - .single() - .ok_or_else(|| serde::de::Error::custom(format!("invalid epoch-ms: {ms}"))), + RawTs::Millis(ms) => epoch_millis(ms), RawTs::Text(s) => { if let Ok(dt) = DateTime::parse_from_rfc3339(&s) { return Ok(dt.with_timezone(&Utc)); } if let Ok(ms) = s.parse::() { - return chrono::TimeZone::timestamp_millis_opt(&Utc, ms) - .single() - .ok_or_else(|| { - serde::de::Error::custom(format!("invalid epoch-ms string: {s}")) - }); + return epoch_millis(ms); } Err(serde::de::Error::custom(format!( "cannot parse '{s}' as RFC 3339 or epoch-ms" diff --git a/src/memory/ingest/extract/extract_tests.rs b/src/memory/ingest/extract/extract_tests.rs index fa2e747..7b7fd34 100644 --- a/src/memory/ingest/extract/extract_tests.rs +++ b/src/memory/ingest/extract/extract_tests.rs @@ -6,7 +6,8 @@ //! ownership boundary. Here we assert directly on the deterministic extractor's //! recovered entities, relations, preferences, and decisions. -use super::{extract_document, ExtractionMode, MemoryIngestionConfig}; +use super::{extract_document, extract_enriched_document, ExtractionMode, MemoryIngestionConfig}; +use crate::memory::types::{MemoryTaint, NamespaceDocumentInput}; fn fixture(path: &str) -> String { let base = std::path::Path::new(env!("CARGO_MANIFEST_DIR")); @@ -123,3 +124,32 @@ fn sentence_mode_extraction_recovers_east_and_west_spatial_relations() { r.subject == "KITCHEN" && r.predicate == "EAST_OF" && r.object == "WEST ROOM" })); } + +#[test] +fn enriched_document_merges_header_metadata_and_reports_counts() { + let input = NamespaceDocumentInput { + namespace: "inbox".into(), + key: "message-1".into(), + title: "Project handoff".into(), + content: "From: Alice Smith \nTags: launch, urgent\n\nAlice Smith owns Project Atlas.".into(), + source_type: "email".into(), + priority: "normal".into(), + tags: vec!["existing".into()], + metadata: serde_json::json!({}), + category: "episodic".into(), + session_id: None, + document_id: None, + taint: MemoryTaint::ExternalSync, + }; + + let (enriched, result) = extract_enriched_document(&input, &MemoryIngestionConfig::default()); + + assert_eq!(enriched.title, input.title); + assert!(enriched.tags.contains(&"existing".to_string())); + assert_eq!(result.document_id, ""); + assert_eq!(result.namespace, ""); + assert!(result.entity_count > 0); + assert_eq!(result.entity_count, result.entities.len()); + assert_eq!(result.relation_count, result.relations.len()); + assert_eq!(result.tags, enriched.tags); +} diff --git a/src/memory/ingest/extract/mod.rs b/src/memory/ingest/extract/mod.rs index 17417cd..f7fc375 100644 --- a/src/memory/ingest/extract/mod.rs +++ b/src/memory/ingest/extract/mod.rs @@ -4,15 +4,15 @@ //! `memory/ingestion` pipeline. It takes raw unstructured text and recovers //! structured knowledge with no model calls: //! -//! 1. **Chunking** — split the document into manageable pieces ([`chunking`], -//! [`text`]). +//! 1. **Chunking** — split the document into manageable pieces (`chunking`, +//! `text`). //! 2. **Structured extraction** — regex rules for email headers, prefixed -//! fields, and explicit graph facts ([`regex`], [`parse_lines`], -//! [`parse_relations`]). +//! fields, and explicit graph facts (`regex`, `parse_lines`, +//! `parse_relations`). //! 3. **Heuristic extraction** — recipient / spatial relations over extraction -//! units ([`parse_units`]). -//! 4. **Aggregation** — alias resolution, dedup, thresholding ([`alias`], -//! [`aggregate`], [`rules`]). +//! units (`parse_units`). +//! 4. **Aggregation** — alias resolution, dedup, thresholding (`alias`, +//! `aggregate`, `rules`). //! //! ## Ownership boundary //! diff --git a/src/memory/ingest/mod.rs b/src/memory/ingest/mod.rs index 101350f..0daca0a 100644 --- a/src/memory/ingest/mod.rs +++ b/src/memory/ingest/mod.rs @@ -49,4 +49,4 @@ pub use pipeline::{ ingest_canonical, ingest_chat, ingest_document, ingest_document_versioned, ingest_document_with_scope, ingest_email, ingest_email_with_raw_refs, }; -pub use types::{IngestOptions, IngestSummary, NullJobSink, TreeJobSink}; +pub use types::{IngestOptions, IngestSummary, NullJobSink, QueueJobSink, TreeJobSink}; diff --git a/src/memory/ingest/pipeline.rs b/src/memory/ingest/pipeline.rs index 5163cf0..2a303b7 100644 --- a/src/memory/ingest/pipeline.rs +++ b/src/memory/ingest/pipeline.rs @@ -16,12 +16,13 @@ use anyhow::{anyhow, bail, Result}; use chrono::Utc; use crate::memory::chunks::{ - self, chunk_markdown, claim_source_ingest_tx, delete_source_ingest, get_chunk_lifecycle_status, - is_source_ingested, set_chunk_lifecycle_status, set_chunk_raw_refs, upsert_staged_chunks_tx, - with_connection, ChunkerInput, ChunkerOptions, SourceKind, CHUNK_STATUS_PENDING_EXTRACTION, + self, chunk_markdown, claim_source_ingest_tx, get_chunk_lifecycle_status_tx, + is_source_ingested, set_chunk_lifecycle_status_tx, set_chunk_raw_refs_tx, + upsert_staged_chunks_tx, with_connection, ChunkerInput, ChunkerOptions, SourceKind, + CHUNK_STATUS_PENDING_EXTRACTION, }; use crate::memory::config::MemoryConfig; -use crate::memory::score::{persist_score, score_chunks_fast, ScoringConfig}; +use crate::memory::score::{persist_score_tx, score_chunks_fast, ScoringConfig}; use crate::memory::store::content; use super::canonicalize::chat::{self, ChatBatch}; @@ -37,32 +38,19 @@ use super::types::{IngestOptions, IngestSummary, TreeJobSink}; /// `source_id` is a stream identifier under which many batches accumulate) and /// rely on deterministic chunk ids for replay idempotency. /// -/// # Atomicity / failure-mode caveats (see `docs/spec/audit/04-queue-ingest.md`) +/// # Atomicity (see `docs/spec/audit/04-queue-ingest.md`) /// -/// - **The gate is a compensated reservation (QI-1).** The document gate is -/// claimed before asynchronous scoring so concurrent ingests cannot both -/// proceed. Any ordinary error in a later stage releases that reservation, -/// allowing a retry. A hard process kill in the narrow interval between the -/// claim and persistence still requires stale-reservation recovery by the -/// host; SQLite transactions cannot span the asynchronous scoring call. -/// - **Step 7 (persist score → set lifecycle → enqueue) is a non-atomic, -/// per-chunk 3-write sequence with a TOCTOU on the `prior` snapshot taken in -/// step 4 (QI-12).** A crash between the lifecycle write and the enqueue -/// strands the chunk at `pending_extraction` (unrecoverable for documents, -/// given the gate caveat above); a concurrent extract worker admitting the -/// chunk between the snapshot read and this reset can cause already-buffered -/// or already-sealed content to flow through the summary tree a second time -/// — the exact double-processing hazard the step-4 snapshot exists to -/// prevent, not fully closed by it. -/// - **Chat/email replay idempotency is weaker than the doc comment above -/// implies (QI-13).** Chunk ids are a function of `(kind, source_id, seq, -/// content)` where `seq` is assigned per-batch and chunk boundaries depend on -/// greedy token packing. Any re-delivery whose batching or boundaries differ -/// even slightly from the original (not just a byte-identical replay) -/// produces a fresh set of chunk ids and re-flows the same messages through -/// the tree. Callers that need exactly-once semantics under redelivery must -/// dedupe upstream (e.g. by message content hash) rather than relying on this -/// function. +/// - The document gate is claimed in the same SQLite transaction as chunk, +/// score, lifecycle, raw-reference, and extract-job persistence. Scoring and +/// content staging happen first, so no await or filesystem operation splits +/// the authoritative database commit. +/// - Chunk upsert, score persistence, lifecycle transition, raw pointers, and +/// extract-job enqueue commit in one SQLite transaction. The prior lifecycle +/// is read inside that transaction, preventing re-ingest from resetting a +/// chunk a worker has already admitted (QI-12). +/// - Chat/email logical-unit ids derive from complete message content rather +/// than batch-local sequence, so overlapping deliveries reuse persisted rows +/// and active queue dedupe keys. pub async fn ingest_canonical( config: &MemoryConfig, source_id: &str, @@ -85,60 +73,15 @@ pub async fn ingest_canonical( return Ok(IngestSummary::empty(source_id)); } - // 2. Authoritative source gate (documents only). Claimed before any write - // so two concurrent ingests of the same document can't both proceed. - // - // The gate row is committed in its own transaction (staging/scoring is - // async and can't hold a SQLite transaction across `.await`). To keep the - // gate honest we treat it as a *reservation*: if any later stage fails we - // release it (see the compensation below) so the document is never left - // permanently marked ingested with zero persisted chunks — a retry can - // then re-claim and finish the job. - let gate_key: Option = if source_kind == SourceKind::Document { - let gate_key = match opts.gate_version_ms { - Some(v) => format!("{source_id}@{v}"), - None => source_id.to_string(), - }; - let claimed = with_connection(config, |conn| { - let tx = conn.unchecked_transaction()?; - let claimed = - claim_source_ingest_tx(&tx, source_kind, &gate_key, Utc::now().timestamp_millis())?; - tx.commit()?; - Ok(claimed) - })?; - if !claimed { - return Ok(IngestSummary::already_ingested(source_id)); - } - Some(gate_key) - } else { - None - }; - - // Run the remaining persist/score/enqueue stages. On any failure after the - // gate was claimed, release the reservation so a retry is not short-circuited. - let result = persist_score_enqueue(config, source_id, chunks, sink, scoring, opts).await; - if result.is_err() { - if let Some(key) = gate_key.as_deref() { - if let Err(cleanup_err) = delete_source_ingest(config, source_kind, key) { - // Best-effort: the original error is the one worth surfacing, but - // a failed release would leave the source wrongly gated, so make - // the operator aware of it. - eprintln!( - "ingest: failed to release source gate {key} after ingest error: {cleanup_err}" - ); - } - } - } - result + persist_score_enqueue(config, source_id, source_kind, chunks, sink, scoring, opts).await } /// Persist chunk bodies + rows, fast-score, and enqueue extract jobs. /// -/// Split out from [`ingest_canonical`] so the source-gate reservation can wrap -/// exactly this fallible tail and compensate (release the gate) on any error. async fn persist_score_enqueue( config: &MemoryConfig, source_id: &str, + source_kind: SourceKind, chunks: Vec, sink: &dyn TreeJobSink, scoring: &ScoringConfig, @@ -149,31 +92,7 @@ async fn persist_score_enqueue( let staged = content::stage_chunks(&content_root, &chunks) .map_err(|e| anyhow!("stage_chunks failed: {e}"))?; - // 4. Snapshot each chunk's CURRENT lifecycle BEFORE the upsert. A chunk that - // already progressed past `pending_extraction` on a prior ingest must not - // be re-scheduled, or already-buffered/sealed content would flow through - // the tree twice. - let mut prior: Vec> = Vec::with_capacity(chunks.len()); - for chunk in &chunks { - prior.push(get_chunk_lifecycle_status(config, &chunk.id)?); - } - - // 5. Persist chunk rows (idempotent on deterministic chunk id). - let chunks_written = with_connection(config, |connection| { - let transaction = connection.unchecked_transaction()?; - let count = upsert_staged_chunks_tx(&transaction, &staged)?; - transaction.commit()?; - Ok(count) - })?; - - // 5b. Raw-archive-backed bodies: attach refs so a worker can resolve them. - if let Some(refs) = opts.raw_refs.as_ref() { - for chunk in &chunks { - set_chunk_raw_refs(config, &chunk.id, refs)?; - } - } - - // 6. Cheap fast-score (no LLM on the ingest hot path). + // 4. Cheap fast-score (no LLM on the ingest hot path). let scores = score_chunks_fast(&chunks, scoring).await?; if scores.len() != chunks.len() { bail!( @@ -183,30 +102,67 @@ async fn persist_score_enqueue( ); } - // 7. Persist scores, mark newly-scheduled chunks pending, and enqueue an - // extract job per scheduled chunk. The fast-score `kept` flag only feeds - // the dropped count for reporting — final admission happens in the worker. - let mut chunks_dropped = 0usize; - let mut extract_jobs_enqueued = 0usize; - let mut chunk_ids = Vec::with_capacity(chunks.len()); - for ((chunk, result), pre) in chunks.iter().zip(scores.iter()).zip(prior.iter()) { - chunk_ids.push(chunk.id.clone()); - if !result.kept { - chunks_dropped += 1; + // 5. Persist the whole database tail atomically. Snapshot lifecycle before + // the upsert while holding this transaction; the worker cannot advance a + // row between the check and its guarded re-schedule. + let persisted = with_connection(config, |connection| { + let transaction = connection.unchecked_transaction()?; + if source_kind == SourceKind::Document { + let gate_key = match opts.gate_version_ms { + Some(version) => format!("{source_id}@{version}"), + None => source_id.to_string(), + }; + if !claim_source_ingest_tx( + &transaction, + source_kind, + &gate_key, + Utc::now().timestamp_millis(), + )? { + return Ok(None); + } } + let prior = chunks + .iter() + .map(|chunk| get_chunk_lifecycle_status_tx(&transaction, &chunk.id)) + .collect::>>()?; + let chunks_written = upsert_staged_chunks_tx(&transaction, &staged)?; - let needs_processing = - matches!(pre.as_deref(), None | Some(CHUNK_STATUS_PENDING_EXTRACTION)); - if !needs_processing { - continue; + if let Some(refs) = opts.raw_refs.as_ref() { + for chunk in &chunks { + set_chunk_raw_refs_tx(&transaction, &chunk.id, refs)?; + } } - let ts_ms = chunk.metadata.timestamp.timestamp_millis(); - persist_score(config, result, ts_ms, None)?; - set_chunk_lifecycle_status(config, &chunk.id, CHUNK_STATUS_PENDING_EXTRACTION)?; - sink.enqueue_extract(&chunk.id)?; - extract_jobs_enqueued += 1; - } + let mut jobs = 0usize; + for ((chunk, result), pre) in chunks.iter().zip(scores.iter()).zip(prior.iter()) { + let needs_processing = + matches!(pre.as_deref(), None | Some(CHUNK_STATUS_PENDING_EXTRACTION)); + if !needs_processing { + continue; + } + persist_score_tx( + &transaction, + result, + chunk.metadata.timestamp.timestamp_millis(), + None, + )?; + set_chunk_lifecycle_status_tx( + &transaction, + &chunk.id, + CHUNK_STATUS_PENDING_EXTRACTION, + )?; + jobs += usize::from(sink.enqueue_extract_tx(&transaction, &chunk.id)?); + } + transaction.commit()?; + Ok(Some((chunks_written, jobs))) + })?; + + let Some((chunks_written, extract_jobs_enqueued)) = persisted else { + return Ok(IngestSummary::already_ingested(source_id)); + }; + + let chunks_dropped = scores.iter().filter(|result| !result.kept).count(); + let chunk_ids = chunks.iter().map(|chunk| chunk.id.clone()).collect(); Ok(IngestSummary { source_id: source_id.to_string(), diff --git a/src/memory/ingest/pipeline_tests.rs b/src/memory/ingest/pipeline_tests.rs index 5af58b1..af241c4 100644 --- a/src/memory/ingest/pipeline_tests.rs +++ b/src/memory/ingest/pipeline_tests.rs @@ -14,7 +14,8 @@ use crate::memory::config::MemoryConfig; use crate::memory::ingest::canonicalize::chat::{ChatBatch, ChatMessage}; use crate::memory::ingest::canonicalize::document::DocumentInput; use crate::memory::ingest::canonicalize::email::{EmailMessage, EmailThread}; -use crate::memory::ingest::types::{NullJobSink, TreeJobSink}; +use crate::memory::ingest::types::{NullJobSink, QueueJobSink, TreeJobSink}; +use crate::memory::queue::{count_by_status, JobStatus}; use crate::memory::score::ScoringConfig; /// Records every enqueued extract job id for assertions. @@ -24,9 +25,13 @@ struct RecordingJobSink { } impl TreeJobSink for RecordingJobSink { - fn enqueue_extract(&self, chunk_id: &str) -> anyhow::Result<()> { + fn enqueue_extract_tx( + &self, + _tx: &rusqlite::Transaction<'_>, + chunk_id: &str, + ) -> anyhow::Result { self.ids.lock().unwrap().push(chunk_id.to_string()); - Ok(()) + Ok(true) } } @@ -39,12 +44,16 @@ struct TogglingJobSink { } impl TreeJobSink for TogglingJobSink { - fn enqueue_extract(&self, chunk_id: &str) -> anyhow::Result<()> { + fn enqueue_extract_tx( + &self, + _tx: &rusqlite::Transaction<'_>, + chunk_id: &str, + ) -> anyhow::Result { if self.fail.load(Ordering::SeqCst) { anyhow::bail!("simulated enqueue failure"); } self.ids.lock().unwrap().push(chunk_id.to_string()); - Ok(()) + Ok(true) } } @@ -114,6 +123,36 @@ async fn ingest_chat_writes_chunks_and_enqueues_extract_jobs() { ); } +#[tokio::test] +async fn queue_sink_commits_chunk_lifecycle_and_extract_job_together() { + let (_tmp, cfg) = test_config(); + let out = ingest_chat( + &cfg, + "slack:#atomic", + "alice", + vec![], + substantive_batch(), + &QueueJobSink, + &ScoringConfig::default_regex_only(), + ) + .await + .unwrap(); + + assert_eq!(out.extract_jobs_enqueued, out.chunk_ids.len()); + assert_eq!( + count_by_status(&cfg, JobStatus::Ready).unwrap(), + out.extract_jobs_enqueued as u64 + ); + for chunk_id in out.chunk_ids { + assert_eq!( + get_chunk_lifecycle_status(&cfg, &chunk_id) + .unwrap() + .as_deref(), + Some(CHUNK_STATUS_PENDING_EXTRACTION) + ); + } +} + #[tokio::test] async fn ingest_chat_empty_batch_is_noop() { let (_tmp, cfg) = test_config(); @@ -185,7 +224,7 @@ async fn second_document_ingest_with_same_source_id_is_short_circuited() { } #[tokio::test] -async fn failed_document_ingest_releases_gate_so_retry_ingests() { +async fn failed_document_ingest_rolls_back_gate_so_retry_ingests() { let (_tmp, cfg) = test_config(); let sink = TogglingJobSink::default(); let scoring = ScoringConfig::default_regex_only(); @@ -198,8 +237,8 @@ async fn failed_document_ingest_releases_gate_so_retry_ingests() { source_ref: Some("notion://page/abc".into()), }; - // First attempt: a post-gate stage (enqueue) fails. The whole ingest must - // error, and the source gate must NOT be left claimed. + // First attempt: enqueue fails inside the database tail. The whole + // transaction, including the source gate, must roll back. sink.fail.store(true, Ordering::SeqCst); let err = ingest_document( &cfg, diff --git a/src/memory/ingest/types.rs b/src/memory/ingest/types.rs index 11f018e..3c31483 100644 --- a/src/memory/ingest/types.rs +++ b/src/memory/ingest/types.rs @@ -2,21 +2,19 @@ //! options, and the result summary. use anyhow::Result; +use rusqlite::Transaction; use crate::memory::chunks::RawRef; +use crate::memory::queue::types::{ExtractChunkPayload, NewJob}; /// Sink for the tree jobs the ingest path produces. /// -/// The async job queue (`crate::memory::queue`) is being ported concurrently, -/// so the ingest orchestrator does **not** hard-depend on it: every kept chunk -/// is handed to a [`TreeJobSink`] for downstream extraction/admission/tree -/// append/seal. A host wires the real queue behind this trait; tests use an -/// in-memory recorder. +/// Implementations enqueue against the caller-owned chunk-database transaction +/// so chunk lifecycle and queue visibility commit atomically. pub trait TreeJobSink: Send + Sync { - /// Enqueue an `extract_chunk` job for `chunk_id`. The downstream worker runs - /// full extraction, the admission gate, buffer append, and sealing — none of - /// which happen on this hot path. - fn enqueue_extract(&self, chunk_id: &str) -> Result<()>; + /// Enqueue an `extract_chunk` job for `chunk_id` inside `tx`. Returns true + /// when a new job was created and false for a deliberate no-op or dedupe. + fn enqueue_extract_tx(&self, tx: &Transaction<'_>, chunk_id: &str) -> Result; } /// A no-op [`TreeJobSink`] that drops every job. Useful when a caller only wants @@ -25,8 +23,21 @@ pub trait TreeJobSink: Send + Sync { pub struct NullJobSink; impl TreeJobSink for NullJobSink { - fn enqueue_extract(&self, _chunk_id: &str) -> Result<()> { - Ok(()) + fn enqueue_extract_tx(&self, _tx: &Transaction<'_>, _chunk_id: &str) -> Result { + Ok(false) + } +} + +/// Production sink backed by the durable SQLite queue. +#[derive(Debug, Default, Clone, Copy)] +pub struct QueueJobSink; + +impl TreeJobSink for QueueJobSink { + fn enqueue_extract_tx(&self, tx: &Transaction<'_>, chunk_id: &str) -> Result { + let job = NewJob::extract_chunk(&ExtractChunkPayload { + chunk_id: chunk_id.to_string(), + })?; + Ok(crate::memory::queue::store::enqueue_tx(tx, &job)?.is_some()) } } diff --git a/src/memory/queue/gate.rs b/src/memory/queue/gate.rs index b3f1678..9b89737 100644 --- a/src/memory/queue/gate.rs +++ b/src/memory/queue/gate.rs @@ -6,8 +6,9 @@ //! the async runtime (`tokio`) is a dev-only dependency here, so this is a //! self-contained, runtime-agnostic re-implementation built on `parking_lot`. //! -//! [`acquire`] blocks until a permit is free and returns an RAII [`Permit`] -//! that returns the slot on drop. The deterministic [`try_acquire`] never +//! [`LlmGate::acquire`] blocks until a permit is free and returns an RAII +//! [`Permit`] that returns the slot on drop. The deterministic +//! [`LlmGate::try_acquire`] never //! blocks — it is the seam tests use to assert the gate actually limits //! concurrency. The worker holds a permit for the duration of an LLM-bound //! handler (see [`JobKind::is_llm_bound`](crate::memory::queue::types::JobKind::is_llm_bound)) diff --git a/src/memory/queue/handlers.rs b/src/memory/queue/handlers.rs index ad64dc3..8c0780f 100644 --- a/src/memory/queue/handlers.rs +++ b/src/memory/queue/handlers.rs @@ -117,15 +117,14 @@ pub enum ReembedProgress { /// ## Idempotency contract /// /// The queue is at-least-once: a crash or `SQLITE_BUSY` after a delegate -/// method has run its side effect but before the caller's follow-up enqueue -/// (or before [`store_settle::mark_done`](crate::memory::queue::store_settle::mark_done) -/// commits) leaves the row `running` past its lease, so +/// method has run its side effect but before parent settlement commits leaves +/// the row `running` past its lease, so /// [`recover_stale_locks`](crate::memory::queue::store_settle::recover_stale_locks) /// puts it back to `ready` and the whole handler — including this delegate -/// call — runs again. `extract_chunk` and `seal_document` in particular are -/// not wrapped in a single transaction with their follow-up enqueue (see -/// `handle_extract`, `handle_seal_document`), so implementations MUST make -/// their side effects safe to repeat: upsert rather than insert, treat +/// call — runs again. Follow-up jobs are committed atomically with settlement, +/// but external/LLM delegate work cannot share that SQLite transaction. +/// Implementations MUST therefore make their side effects safe to repeat: +/// upsert rather than insert, treat /// "already applied" as success, and avoid effects that are observably wrong /// the second time (e.g. re-running a paid LLM call is wasteful but not /// incorrect; double-appending a chunk into a tree would be). @@ -184,12 +183,53 @@ pub trait QueueDelegates: Send + Sync { fn has_uncovered_reembed_work(&self, config: &MemoryConfig, signature: &str) -> Result; } +/// Handler outcome plus durable follow-up jobs. The worker commits all +/// follow-ups in the same transaction that marks the parent job done. +pub(crate) struct JobPlan { + pub outcome: JobOutcome, + pub follow_ups: Vec, +} + +impl JobPlan { + fn done(follow_ups: Vec) -> Self { + Self { + outcome: JobOutcome::Done, + follow_ups, + } + } + + fn outcome(outcome: JobOutcome) -> Self { + Self { + outcome, + follow_ups: Vec::new(), + } + } +} + /// Dispatch a claimed job to the matching per-kind handler. pub async fn handle_job( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, ) -> Result { + let plan = plan_job(config, job, delegates).await?; + for follow_up in &plan.follow_ups { + if store::enqueue(config, follow_up)?.is_some() + && follow_up.kind == JobKind::ReembedBackfill + { + set_backfill_in_progress(true); + } + } + Ok(plan.outcome) +} + +/// Plan a handler result without persisting follow-ups. Used by the worker so +/// parent settlement and child enqueue cannot be separated by a crash. +pub(crate) async fn plan_job( + config: &MemoryConfig, + job: &Job, + delegates: &dyn QueueDelegates, +) -> Result { match job.kind { JobKind::ExtractChunk => handle_extract(config, job, delegates).await, JobKind::AppendBuffer => handle_append_buffer(config, job, delegates).await, @@ -218,13 +258,15 @@ async fn handle_extract( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, -) -> Result { +) -> Result { let payload: ExtractChunkPayload = parse_payload(job, "ExtractChunk")?; let Some(decision) = delegates.extract_chunk(config, &payload.chunk_id).await? else { // Chunk row vanished between enqueue and claim — nothing to do. - return Ok(JobOutcome::Done); + return Ok(JobPlan::done(Vec::new())); }; + let mut follow_ups = Vec::new(); + // Admitted, flat-buffer source: enqueue the append-buffer follow-up. The // per-document-versioned sources (Notion) skip the flat L0 buffer — their // tree is built by a SealDocument job enqueued at ingest. @@ -237,41 +279,44 @@ async fn handle_extract( source_id: decision.tree_scope.clone(), }, })?; - store::enqueue(config, &follow_up)?; + follow_ups.push(follow_up); } // Anything admitted arms the re-embed backfill so the embedding pass starts // promptly (extract no longer embeds inline). if decision.kept { - crate::memory::queue::ops::ensure_reembed_backfill(config, delegates)?; + if let Some(job) = crate::memory::queue::ops::planned_reembed_backfill(config, delegates)? { + follow_ups.push(job); + } } - Ok(JobOutcome::Done) + Ok(JobPlan::done(follow_ups)) } async fn handle_append_buffer( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, -) -> Result { +) -> Result { let payload: AppendBufferPayload = parse_payload(job, "AppendBuffer")?; let Some(decision) = delegates .append_node(config, &payload.node, &payload.target) .await? else { // Missing chunk/summary, or the target topic tree was archived — drop. - return Ok(JobOutcome::Done); + return Ok(JobPlan::done(Vec::new())); }; + let mut follow_ups = Vec::new(); if decision.should_seal { let seal = SealPayload { tree_id: decision.tree_id, level: 0, force_now_ms: None, }; - store::enqueue(config, &NewJob::seal(&seal)?)?; + follow_ups.push(NewJob::seal(&seal)?); } - Ok(JobOutcome::Done) + Ok(JobPlan::done(follow_ups)) } /// Run the `Seal` handler for exactly one buffer level. @@ -279,60 +324,59 @@ async fn handle_append_buffer( /// A cascading seal returns the parent payload so each level stays its own /// crash-recovery checkpoint (see [`SealPayload::dedupe_key`]). /// -/// Edge-triggered gap: [`SealPayload::dedupe_key`] is scoped to -/// `(tree_id, level)` and only suppresses duplicates while a seal for that key -/// is `ready`/`running` (see the partial unique index documented on -/// [`store::enqueue`]). If a new enqueue attempt for the same level arrives -/// while a seal is still `running` and the buffer crosses its gate again -/// afterward, the suppressed attempt is simply dropped — nothing re-checks the -/// gate when the in-flight seal finishes, so the newly-buffered content waits -/// for the next `flush_stale` tick (which can be hours away) instead of being -/// sealed promptly. +/// Settlement re-checks the live same-level buffer after releasing this job's +/// active dedupe key. Content appended while the seal was running therefore +/// restores a suppressed seal edge immediately when it crosses the gate. async fn handle_seal( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, -) -> Result { +) -> Result { let payload: SealPayload = parse_payload(job, "Seal")?; // Seal exactly one level. A cascading seal returns the parent payload so // each level stays its own crash-recovery checkpoint. - if let Some(parent) = delegates.seal_level(config, &payload).await? { - store::enqueue(config, &NewJob::seal(&parent)?)?; - } - Ok(JobOutcome::Done) + let follow_ups = delegates + .seal_level(config, &payload) + .await? + .map(|parent| NewJob::seal(&parent)) + .transpose()? + .into_iter() + .collect(); + Ok(JobPlan::done(follow_ups)) } async fn handle_flush_stale( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, -) -> Result { +) -> Result { let payload: FlushStalePayload = parse_payload(job, "FlushStale")?; let age_secs = payload.max_age_secs.unwrap_or(L0_DEFAULT_FLUSH_AGE_SECS); let now_ms = chrono::Utc::now().timestamp_millis(); + let mut follow_ups = Vec::new(); for buf in delegates.list_stale_buffers(config, age_secs).await? { let seal = SealPayload { tree_id: buf.tree_id, level: buf.level, force_now_ms: Some(now_ms), }; - store::enqueue(config, &NewJob::seal(&seal)?)?; + follow_ups.push(NewJob::seal(&seal)?); } - Ok(JobOutcome::Done) + Ok(JobPlan::done(follow_ups)) } async fn handle_seal_document( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, -) -> Result { +) -> Result { let payload: SealDocumentPayload = parse_payload(job, "SealDocument")?; if payload.chunk_ids.is_empty() { // Empty version set — nothing to seal. - return Ok(JobOutcome::Done); + return Ok(JobPlan::done(Vec::new())); } delegates.seal_document(config, &payload).await?; - Ok(JobOutcome::Done) + Ok(JobPlan::done(Vec::new())) } /// Run one step of the `ReembedBackfill` chain. @@ -350,7 +394,7 @@ async fn handle_reembed_backfill( config: &MemoryConfig, job: &Job, delegates: &dyn QueueDelegates, -) -> Result { +) -> Result { let payload: ReembedBackfillPayload = parse_payload(job, "ReembedBackfill")?; match delegates.reembed_batch(config, &payload.signature).await? { @@ -360,10 +404,10 @@ async fn handle_reembed_backfill( set_backfill_in_progress(true); // More rows may remain — reschedule THIS row (no re-enqueue, so the // per-signature dedupe key stays valid). - Ok(JobOutcome::Defer { + Ok(JobPlan::outcome(JobOutcome::Defer { until_ms: chrono::Utc::now().timestamp_millis() + REEMBED_BACKFILL_REVISIT_MS, reason: "re-embed backfill: batch done, more pending".to_string(), - }) + })) } ReembedProgress::Wrote { more_pending: false, @@ -372,7 +416,7 @@ async fn handle_reembed_backfill( | ReembedProgress::NoProvider | ReembedProgress::StaleSignature => { set_backfill_in_progress(false); - Ok(JobOutcome::Done) + Ok(JobPlan::done(Vec::new())) } } } diff --git a/src/memory/queue/mod.rs b/src/memory/queue/mod.rs index 4982572..14b9702 100644 --- a/src/memory/queue/mod.rs +++ b/src/memory/queue/mod.rs @@ -41,6 +41,7 @@ pub mod gate; mod handlers; mod ops; +mod payloads; mod redact; /// Always-on async background worker + scheduler loops (feature `tokio`). /// diff --git a/src/memory/queue/ops.rs b/src/memory/queue/ops.rs index b354448..e752846 100644 --- a/src/memory/queue/ops.rs +++ b/src/memory/queue/ops.rs @@ -43,11 +43,7 @@ pub fn ensure_reembed_backfill( config: &MemoryConfig, delegates: &dyn QueueDelegates, ) -> Result<()> { - let sig = delegates.active_signature(config); - if delegates.has_uncovered_reembed_work(config, &sig)? { - let job = NewJob::reembed_backfill(&ReembedBackfillPayload { - signature: sig.clone(), - })?; + if let Some(job) = planned_reembed_backfill(config, delegates)? { if store::enqueue(config, &job)?.is_some() { set_backfill_in_progress(true); } @@ -55,6 +51,20 @@ pub fn ensure_reembed_backfill( Ok(()) } +/// Build the deduplicated re-embed follow-up without persisting it. Workers +/// use this to commit follow-up creation atomically with parent settlement. +pub(crate) fn planned_reembed_backfill( + config: &MemoryConfig, + delegates: &dyn QueueDelegates, +) -> Result> { + let sig = delegates.active_signature(config); + if delegates.has_uncovered_reembed_work(config, &sig)? { + let job = NewJob::reembed_backfill(&ReembedBackfillPayload { signature: sig })?; + return Ok(Some(job)); + } + Ok(None) +} + #[cfg(test)] #[path = "ops_tests.rs"] mod tests; diff --git a/src/memory/queue/payloads.rs b/src/memory/queue/payloads.rs new file mode 100644 index 0000000..79c0d92 --- /dev/null +++ b/src/memory/queue/payloads.rs @@ -0,0 +1,148 @@ +//! Strongly typed payloads persisted in queue job rows. + +use serde::{Deserialize, Serialize}; + +/// Reference to either a leaf chunk or a sealed summary node. +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum NodeRef { + /// Leaf chunk id. + Leaf { + /// Chunk row id. + chunk_id: String, + }, + /// Sealed summary id. + Summary { + /// Summary row id. + summary_id: String, + }, +} + +impl NodeRef { + /// Stable kind-prefixed identity for queue deduplication. + pub fn dedupe_fragment(&self) -> String { + match self { + Self::Leaf { chunk_id } => format!("leaf:{chunk_id}"), + Self::Summary { summary_id } => format!("summary:{summary_id}"), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +/// Extract/admit one chunk. +pub struct ExtractChunkPayload { + /// Chunk row id. + pub chunk_id: String, +} + +impl ExtractChunkPayload { + /// Stable queue deduplication key. + pub fn dedupe_key(&self) -> String { + format!("extract:{}", self.chunk_id) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +#[serde(tag = "kind", rename_all = "snake_case")] +/// Destination for an append-buffer job. +pub enum AppendTarget { + /// Source tree selected by source id. + Source { + /// Logical source id. + source_id: String, + }, + /// Topic tree selected by physical tree id. + Topic { + /// Physical topic tree id. + tree_id: String, + }, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +/// Node and destination for one buffer append. +pub struct AppendBufferPayload { + /// Leaf or summary to append. + pub node: NodeRef, + /// Destination buffer. + pub target: AppendTarget, +} + +impl AppendBufferPayload { + /// Stable queue deduplication key. + pub fn dedupe_key(&self) -> String { + let node = self.node.dedupe_fragment(); + match &self.target { + AppendTarget::Source { source_id } => format!("append:source:{source_id}:{node}"), + AppendTarget::Topic { tree_id } => format!("append:topic:{tree_id}:{node}"), + } + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +/// Seal one buffer level. +pub struct SealPayload { + /// Physical tree id. + pub tree_id: String, + /// Zero-based buffer level. + pub level: u32, + /// Presence forces sealing; value supplies the scheduling instant. + pub force_now_ms: Option, +} + +impl SealPayload { + /// Stable per-tree-and-level deduplication key. + pub fn dedupe_key(&self) -> String { + format!("seal:{}:{}", self.tree_id, self.level) + } +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +/// Scan and seal stale buffers. +pub struct FlushStalePayload { + /// Optional maximum buffer age override. + pub max_age_secs: Option, +} + +impl FlushStalePayload { + /// Stable key scoped to one three-hour scheduler block. + pub fn dedupe_key(&self, date_iso: &str, hour_block: u32) -> String { + format!("flush_stale:{date_iso}-h{hour_block}") + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +/// Re-embed rows missing the named vector signature. +pub struct ReembedBackfillPayload { + /// Embedding provider/model/dimension signature. + pub signature: String, +} + +impl ReembedBackfillPayload { + /// Stable per-signature deduplication key. + pub fn dedupe_key(&self) -> String { + format!("reembed_backfill:{}", self.signature) + } +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +/// Build and merge one versioned document subtree. +pub struct SealDocumentPayload { + /// Connection-level tree scope. + pub tree_scope: String, + /// Stable document identity. + pub doc_id: String, + /// Source version in epoch milliseconds. + pub version_ms: Option, + /// Ordered leaf chunk ids for this version. + pub chunk_ids: Vec, +} + +impl SealDocumentPayload { + /// Stable per-document-version deduplication key. + pub fn dedupe_key(&self) -> String { + match self.version_ms { + Some(version) => format!("seal_doc:{}@{version}", self.doc_id), + None => format!("seal_doc:{}", self.doc_id), + } + } +} diff --git a/src/memory/queue/runtime/mod.rs b/src/memory/queue/runtime/mod.rs index d20a86b..96e1f69 100644 --- a/src/memory/queue/runtime/mod.rs +++ b/src/memory/queue/runtime/mod.rs @@ -8,16 +8,16 @@ //! //! ## Loops //! -//! - [`run_worker`]: repeatedly [`run_once`], sleeping between polls. Idle polls -//! back off by [`WorkerLoopConfig::idle_backoff`]; transient SQLite errors +//! - `run_worker`: repeatedly runs one job, sleeping between polls. Idle polls +//! use the configured backoff; transient SQLite errors //! back off by kind (busy / I/O / disk-full) and re-poll; corruption is fatal //! and returned to the host (mirroring OpenHuman's quarantine policy). -//! - [`run_scheduler`]: on a fixed cadence, recover expired leases, enqueue +//! - `run_scheduler`: on a fixed cadence, recover expired leases, enqueue //! `flush_stale` (dedupe-safe), and `self_heal` transiently-failed jobs. -//! - [`run`]: [`bootstrap`] once, then drive both loops concurrently until +//! - `run`: bootstrap once, then drive both loops concurrently until //! shutdown or a fatal error. //! -//! All loops observe a shared [`Shutdown`] flag between jobs and during backoff, +//! All loops observe a shared shutdown flag between jobs and during backoff, //! so a triggered shutdown stops them promptly without aborting an in-flight //! handler. //! @@ -29,7 +29,7 @@ pub mod types; use std::time::Duration; -use anyhow::Result; +use anyhow::{Context, Result}; use tokio::time::sleep; use crate::memory::config::MemoryConfig; @@ -65,10 +65,17 @@ pub async fn run_worker( tokio::task::yield_now().await; } Ok(false) => interruptible_sleep(opts.idle_backoff, shutdown).await, - Err(err) => match backoff_for(&err, opts) { - Some(backoff) => interruptible_sleep(backoff, shutdown).await, - None => return Err(err), - }, + Err(err) => { + if worker::is_sqlite_corrupt(&err) { + crate::memory::chunks::recover_corrupt_db(config) + .context("recover corrupt queue database")?; + continue; + } + match backoff_for(&err, opts) { + Some(backoff) => interruptible_sleep(backoff, shutdown).await, + None => return Err(err), + } + } } } Ok(()) @@ -88,17 +95,23 @@ pub async fn run_scheduler( while !shutdown.is_triggered() { if let Err(err) = crate::memory::queue::store_settle::recover_stale_locks(config) { if worker::is_sqlite_corrupt(&err) { - return Err(err); + crate::memory::chunks::recover_corrupt_db(config) + .context("recover corrupt queue database")?; + continue; } } if let Err(err) = scheduler::enqueue_flush_stale(config) { if worker::is_sqlite_corrupt(&err) { - return Err(err); + crate::memory::chunks::recover_corrupt_db(config) + .context("recover corrupt queue database")?; + continue; } } if let Err(err) = scheduler::self_heal(config) { if worker::is_sqlite_corrupt(&err) { - return Err(err); + crate::memory::chunks::recover_corrupt_db(config) + .context("recover corrupt queue database")?; + continue; } } interruptible_sleep(opts.tick, shutdown).await; diff --git a/src/memory/queue/runtime/tests.rs b/src/memory/queue/runtime/tests.rs index edb74bf..e9f0140 100644 --- a/src/memory/queue/runtime/tests.rs +++ b/src/memory/queue/runtime/tests.rs @@ -169,3 +169,67 @@ fn backoff_classifies_busy_as_transient() { )); assert_eq!(backoff_for(&busy, &opts), Some(opts.busy_backoff)); } + +#[test] +fn backoff_classifies_disk_full_io_and_generic_errors() { + let opts = WorkerLoopConfig::default(); + let sqlite = |code, extended_code, message: &str| { + anyhow::Error::from(rusqlite::Error::SqliteFailure( + rusqlite::ffi::Error { + code, + extended_code, + }, + Some(message.into()), + )) + }; + assert_eq!( + backoff_for( + &sqlite( + rusqlite::ErrorCode::DiskFull, + 13, + "database or disk is full" + ), + &opts + ), + Some(opts.disk_full_backoff) + ); + assert_eq!( + backoff_for( + &sqlite(rusqlite::ErrorCode::SystemIoFailure, 10, "disk I/O error"), + &opts + ), + Some(opts.io_backoff) + ); + assert_eq!( + backoff_for(&anyhow::anyhow!("other"), &opts), + Some(opts.error_backoff) + ); +} + +#[tokio::test] +async fn run_bootstraps_and_returns_when_shutdown_is_pretriggered() { + let (_tmp, cfg) = test_config(); + let delegates = RecordingDelegates::admitting(); + let shutdown = Shutdown::new(); + shutdown.trigger(); + + run( + &cfg, + &delegates, + &fast_worker_opts(), + &SchedulerLoopConfig { + tick: Duration::from_millis(1), + }, + &shutdown, + ) + .await + .unwrap(); +} + +#[tokio::test] +async fn interruptible_sleep_handles_zero_and_pretriggered_shutdown() { + let shutdown = Shutdown::new(); + interruptible_sleep(Duration::ZERO, &shutdown).await; + shutdown.trigger(); + interruptible_sleep(Duration::from_secs(1), &shutdown).await; +} diff --git a/src/memory/queue/store.rs b/src/memory/queue/store.rs index 38c449d..5ec2f1a 100644 --- a/src/memory/queue/store.rs +++ b/src/memory/queue/store.rs @@ -36,7 +36,9 @@ use crate::memory::queue::types::{Job, JobKind, JobStatus, NewJob, RETIRED_JOB_K pub const DEFAULT_LOCK_DURATION_MS: i64 = 5 * 60 * 1_000; /// Backoff math for retry. Returns `now + min(base * 2^attempts, cap)`. +#[cfg(test)] pub(crate) const RETRY_BASE_MS: i64 = 60 * 1_000; +#[cfg(test)] pub(crate) const RETRY_CAP_MS: i64 = 60 * 60 * 1_000; pub(crate) const DEFAULT_MAX_ATTEMPTS: u32 = 5; @@ -257,13 +259,22 @@ pub(crate) fn row_to_job(row: &rusqlite::Row<'_>) -> rusqlite::Result { } /// Exponential backoff: attempt 1 → 60s, 2 → 120s, 3 → 240s, capped at 1h. +#[cfg(test)] pub(crate) fn backoff_ms(attempts_so_far: u32) -> i64 { + backoff_ms_with_policy(attempts_so_far, RETRY_BASE_MS, RETRY_CAP_MS) +} + +pub(crate) fn backoff_ms_with_policy( + attempts_so_far: u32, + retry_base_ms: i64, + retry_cap_ms: i64, +) -> i64 { // `attempts_so_far` is the count BEFORE the next retry's attempt — so the // first retry uses `attempts_so_far == 1`, giving base * 2^0 = 60s. let exp = attempts_so_far.saturating_sub(1).min(20); let mult = 1i64 << exp; - let raw = RETRY_BASE_MS.saturating_mul(mult); - raw.min(RETRY_CAP_MS) + let raw = retry_base_ms.saturating_mul(mult); + raw.min(retry_cap_ms) } /// True if `kind` is a retired wire string (used by tests / callers that want diff --git a/src/memory/queue/store_settle.rs b/src/memory/queue/store_settle.rs index ec4558f..b736e60 100644 --- a/src/memory/queue/store_settle.rs +++ b/src/memory/queue/store_settle.rs @@ -13,17 +13,34 @@ use rusqlite::params; use crate::memory::chunks::with_connection; use crate::memory::config::MemoryConfig; -use crate::memory::queue::store::backoff_ms; -use crate::memory::queue::types::{Job, JobFailure}; +use crate::memory::queue::store::backoff_ms_with_policy; +use crate::memory::queue::types::{Job, JobFailure, NewJob}; + +/// Maximum lifetime of a repeatedly deferred job. Deferral is intentionally +/// free of the handler failure budget, but it must still have a terminal bound. +#[cfg(test)] +pub(crate) const MAX_DEFER_AGE_MS: i64 = 7 * 24 * 60 * 60 * 1_000; /// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`. pub fn mark_done(config: &MemoryConfig, job: &Job) -> Result<()> { + mark_done_with_followups(config, job, &[]) +} + +/// Mark a claimed job done and enqueue all of its follow-up jobs in the same +/// SQLite transaction. A crash therefore exposes either neither transition or +/// both, never a completed parent with a missing child edge. +pub(crate) fn mark_done_with_followups( + config: &MemoryConfig, + job: &Job, + follow_ups: &[NewJob], +) -> Result<()> { let job_id = &job.id; let claim_attempts = job.attempts as i64; let claim_started_at = job.started_at_ms; with_connection(config, |conn| { let now_ms = Utc::now().timestamp_millis(); - conn.execute( + let tx = conn.unchecked_transaction()?; + let updated = tx.execute( "UPDATE mem_tree_jobs SET status = 'done', completed_at_ms = ?1, @@ -34,6 +51,12 @@ pub fn mark_done(config: &MemoryConfig, job: &Job) -> Result<()> { AND started_at_ms IS ?4", params![now_ms, job_id, claim_attempts, claim_started_at], )?; + if updated != 0 { + for follow_up in follow_ups { + crate::memory::queue::store::enqueue_tx(&tx, follow_up)?; + } + } + tx.commit()?; Ok(()) }) } @@ -93,7 +116,11 @@ pub fn mark_failed_typed( ], )?; } else { - let next_at = now_ms.saturating_add(backoff_ms(attempts as u32)); + let next_at = now_ms.saturating_add(backoff_ms_with_policy( + attempts as u32, + config.queue.retry_base_ms, + config.queue.retry_cap_ms, + )); conn.execute( "UPDATE mem_tree_jobs SET status = 'ready', @@ -119,32 +146,35 @@ pub fn mark_failed_typed( /// `reason` is recorded in `last_error` for visibility and `started_at_ms` is /// cleared so the next claim stamps a fresh start time. /// -/// There is no cap on how many times a row may be deferred, and deferring -/// never advances toward `max_attempts`. A handler that repeatedly makes no -/// progress and always returns `JobOutcome::Defer` (rather than eventually -/// erroring) re-defers the same row forever at zero budget cost — this -/// function has no way to distinguish that from legitimate rate-limit -/// backoff. Callers that need a bound should track defer count or job age in -/// the payload and translate it to an `Err` (burning the retry budget) once a -/// cap is hit. +/// Deferral remains free of the attempt budget, but a row that has spent seven +/// days in the defer cycle is parked as an unrecoverable failure. This bounds +/// permanently non-progressing handlers without penalising ordinary +/// rate-limit or multi-batch backfill deferrals. pub fn mark_deferred(config: &MemoryConfig, job: &Job, until_ms: i64, reason: &str) -> Result<()> { let job_id = &job.id; let claim_attempts = job.attempts as i64; let pre_claim_attempts = claim_attempts.saturating_sub(1); let claim_started_at = job.started_at_ms; with_connection(config, |conn| { + let now_ms = Utc::now().timestamp_millis(); + let oldest_allowed_ms = now_ms.saturating_sub(config.queue.max_defer_age_ms); conn.execute( "UPDATE mem_tree_jobs - SET status = 'ready', - attempts = ?1, - available_at_ms = ?2, + SET status = CASE WHEN created_at_ms <= ?2 THEN 'failed' ELSE 'ready' END, + attempts = ?3, + available_at_ms = ?4, locked_until_ms = NULL, started_at_ms = NULL, - last_error = ?3 - WHERE id = ?4 - AND attempts = ?5 - AND started_at_ms IS ?6", + completed_at_ms = CASE WHEN created_at_ms <= ?2 THEN ?1 ELSE NULL END, + last_error = ?5, + failure_reason = CASE WHEN created_at_ms <= ?2 THEN 'defer_age_exceeded' ELSE NULL END, + failure_class = CASE WHEN created_at_ms <= ?2 THEN 'unrecoverable' ELSE NULL END + WHERE id = ?6 + AND attempts = ?7 + AND started_at_ms IS ?8", params![ + now_ms, + oldest_allowed_ms, pre_claim_attempts, until_ms, reason, @@ -182,20 +212,17 @@ pub fn recover_stale_locks(config: &MemoryConfig) -> Result { /// worker pool, so any `running` row at clean-shutdown time was claimed by us. /// Returns the number of rows released. /// -/// Unlike [`mark_deferred`], this does **not** revert the `attempts` bump -/// [`claim_next`](super::store::claim_next) applied when the row was claimed -/// — released rows keep the attempt they were charged for. A process that is -/// restarted repeatedly mid-job (e.g. five graceful shutdowns while the same -/// job is running) burns through `max_attempts` purely from the -/// release/re-claim cycle, with no actual handler failure involved, and the -/// row can end up settled `failed` without ever having been given a real -/// chance to complete. +/// This reverts the `attempts` bump applied by +/// [`claim_next`](super::store::claim_next): a graceful interruption is not a +/// handler failure and must not consume the job's retry budget. pub fn release_running_locks(config: &MemoryConfig) -> Result { with_connection(config, |conn| { let n = conn.execute( "UPDATE mem_tree_jobs SET status = 'ready', - locked_until_ms = NULL + attempts = CASE WHEN attempts > 0 THEN attempts - 1 ELSE 0 END, + locked_until_ms = NULL, + started_at_ms = NULL WHERE status = 'running'", [], )?; diff --git a/src/memory/queue/store_settle_tests.rs b/src/memory/queue/store_settle_tests.rs index bcfe904..02085f5 100644 --- a/src/memory/queue/store_settle_tests.rs +++ b/src/memory/queue/store_settle_tests.rs @@ -39,6 +39,33 @@ fn mark_done_settles_current_lessee() { assert!(row.locked_until_ms.is_none()); } +#[test] +fn mark_done_commits_followup_in_parent_settlement_transaction() { + let (_tmp, cfg) = test_config(); + let parent_id = enqueue(&cfg, &extract_job("c-parent", 5)) + .unwrap() + .expect("inserted"); + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + let follow_up = NewJob::append_buffer(&AppendBufferPayload { + node: NodeRef::Leaf { + chunk_id: "c-parent".into(), + }, + target: AppendTarget::Source { + source_id: "slack:#atomic".into(), + }, + }) + .unwrap(); + + mark_done_with_followups(&cfg, &claimed, &[follow_up]).unwrap(); + + assert_eq!( + get_job(&cfg, &parent_id).unwrap().unwrap().status, + JobStatus::Done + ); + assert_eq!(count_by_status(&cfg, JobStatus::Ready).unwrap(), 1); + assert_eq!(count_total(&cfg).unwrap(), 2); +} + #[test] fn mark_failed_typed_unrecoverable_terminates_immediately() { let (_tmp, cfg) = test_config(); @@ -245,7 +272,9 @@ fn release_running_locks_resets_running_rows_regardless_of_lease() { assert_eq!(released, 1); let row = get_job(&cfg, &id).unwrap().unwrap(); assert_eq!(row.status, JobStatus::Ready); + assert_eq!(row.attempts, 0, "shutdown must not spend a retry"); assert!(row.locked_until_ms.is_none()); + assert!(row.started_at_ms.is_none()); } #[test] @@ -312,6 +341,36 @@ fn mark_deferred_does_not_increment_attempts() { assert_eq!(claim2.attempts, 1, "Defer didn't count toward the budget"); } +#[test] +fn mark_deferred_parks_jobs_that_exceed_the_defer_age_bound() { + let (_tmp, cfg) = test_config(); + let id = enqueue(&cfg, &extract_job("c-defer-expired", 5)) + .unwrap() + .expect("inserted"); + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + with_connection(&cfg, |conn| { + conn.execute( + "UPDATE mem_tree_jobs SET created_at_ms = ?1 WHERE id = ?2", + rusqlite::params![ + Utc::now() + .timestamp_millis() + .saturating_sub(MAX_DEFER_AGE_MS + 1), + id, + ], + )?; + Ok(()) + }) + .unwrap(); + + mark_deferred(&cfg, &claimed, i64::MAX, "still blocked").unwrap(); + + let row = get_job(&cfg, &id).unwrap().unwrap(); + assert_eq!(row.status, JobStatus::Failed); + assert_eq!(row.failure_reason.as_deref(), Some("defer_age_exceeded")); + assert_eq!(row.failure_class.as_deref(), Some("unrecoverable")); + assert!(row.completed_at_ms.is_some()); +} + #[test] fn deferred_row_not_claimable_until_until_ms() { let (_tmp, cfg) = test_config(); diff --git a/src/memory/queue/types.rs b/src/memory/queue/types.rs index da6ac13..ca07356 100644 --- a/src/memory/queue/types.rs +++ b/src/memory/queue/types.rs @@ -13,7 +13,11 @@ //! (`transient` / `unrecoverable`), and [`JobFailure::is_unrecoverable`]. use anyhow::{anyhow, Result}; -use serde::{Deserialize, Serialize}; + +pub use super::payloads::{ + AppendBufferPayload, AppendTarget, ExtractChunkPayload, FlushStalePayload, NodeRef, + ReembedBackfillPayload, SealDocumentPayload, SealPayload, +}; /// Discriminator persisted in `mem_tree_jobs.kind`. #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] @@ -218,194 +222,6 @@ impl std::fmt::Display for JobFailure { impl std::error::Error for JobFailure {} -// ── Payloads ─────────────────────────────────────────────────────────────── - -/// Reference to either a leaf chunk or a sealed summary node. Used by payloads -/// that route content through the pipeline regardless of which kind of source -/// produced it. -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum NodeRef { - /// A leaf chunk (wire tag `leaf`), identified by its chunk id. - Leaf { - /// Chunk row id of the leaf. - chunk_id: String, - }, - /// A sealed summary node (wire tag `summary`), identified by its summary id. - Summary { - /// Summary-tree row id of the node. - summary_id: String, - }, -} - -impl NodeRef { - /// Stringified id with kind prefix (`leaf:` or `summary:`), suitable for - /// dedupe-key composition. - pub fn dedupe_fragment(&self) -> String { - match self { - NodeRef::Leaf { chunk_id } => format!("leaf:{chunk_id}"), - NodeRef::Summary { summary_id } => format!("summary:{summary_id}"), - } - } -} - -/// Payload for [`JobKind::ExtractChunk`]: the single chunk to run extraction -/// and admission over. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ExtractChunkPayload { - /// Chunk row id to extract. - pub chunk_id: String, -} - -impl ExtractChunkPayload { - /// Stable dedupe key written to `mem_tree_jobs.dedupe_key` so the partial - /// unique index can suppress in-flight duplicates. - pub fn dedupe_key(&self) -> String { - format!("extract:{}", self.chunk_id) - } -} - -/// Where an `AppendBuffer` job should land its node. Source-tree appends are -/// keyed by `source_id`; topic-tree appends are keyed by `tree_id` because -/// there can be many topic trees per node. -#[derive(Clone, Debug, Serialize, Deserialize)] -#[serde(tag = "kind", rename_all = "snake_case")] -pub enum AppendTarget { - /// Append into the source tree keyed by `source_id` (wire tag `source`). - Source { - /// Source id whose tree receives the node. - source_id: String, - }, - /// Append into a specific topic tree keyed by `tree_id` (wire tag `topic`), - /// since a node may fan out to many topic trees. - Topic { - /// Topic-tree id that receives the node. - tree_id: String, - }, -} - -/// Payload for [`JobKind::AppendBuffer`]: which node to append and into which -/// tree's L0 buffer. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct AppendBufferPayload { - /// The leaf or summary node to push into the buffer. - pub node: NodeRef, - /// Destination tree for the append. - pub target: AppendTarget, -} - -impl AppendBufferPayload { - /// Stable dedupe key written to `mem_tree_jobs.dedupe_key`. - pub fn dedupe_key(&self) -> String { - let node_part = self.node.dedupe_fragment(); - match &self.target { - AppendTarget::Source { source_id } => { - format!("append:source:{source_id}:{node_part}") - } - AppendTarget::Topic { tree_id } => { - format!("append:topic:{tree_id}:{node_part}") - } - } - } -} - -/// Payload for [`JobKind::Seal`]: seal exactly one buffer level of one tree. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct SealPayload { - /// Tree whose buffer level is sealed. - pub tree_id: String, - /// Buffer level to seal (0-based; L0 is the leaf buffer). - pub level: u32, - /// When `Some`, the seal handler bypasses the buffer-budget check and - /// force-seals — used by the time-based flush path. The wall-clock is - /// passed through so the seal stamps a deterministic `sealed_at`. - pub force_now_ms: Option, -} - -impl SealPayload { - /// Stable dedupe key written to `mem_tree_jobs.dedupe_key`. - /// - /// Active seal-job uniqueness is enforced per `(tree, level)`: a seal - /// already in flight suppresses duplicate enqueues. Once the job completes - /// the partial index releases the key, so the next time the buffer crosses - /// its gate a fresh seal can be enqueued. - pub fn dedupe_key(&self) -> String { - format!("seal:{}:{}", self.tree_id, self.level) - } -} - -/// Payload for [`JobKind::FlushStale`]: scan stale buffers and enqueue seals -/// for any past the age cap. -#[derive(Clone, Debug, Serialize, Deserialize, Default)] -pub struct FlushStalePayload { - /// Override the configured default flush age. Optional so the scheduler can - /// enqueue with `None` and let the handler use the configured default. - pub max_age_secs: Option, -} - -impl FlushStalePayload { - /// Dedupe key scoped to a 3-hour UTC block (`hour_block = hour / 3`, - /// `0..=7`) so flush_stale runs up to 8× per day. Without this, low-volume - /// sources wait a full day between seal opportunities. - /// - /// Pure: both `date_iso` and `hour_block` are supplied by the caller from a - /// single `Utc::now()` reading, which keeps the key deterministic in tests - /// and avoids a 3-hour-boundary race. - pub fn dedupe_key(&self, date_iso: &str, hour_block: u32) -> String { - format!("flush_stale:{date_iso}-h{hour_block}") - } -} - -/// Re-embed backfill. One chain per embedding signature: the `dedupe_key` is -/// the signature, so re-triggering while a chain is in-flight is correctly -/// suppressed (exactly one backfill per space). The handler self-continues via -/// [`JobOutcome::Defer`] (reschedules this same row) rather than re-enqueuing, -/// so the fixed dedupe key is safe. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct ReembedBackfillPayload { - /// The embedding signature this chain re-embeds under. If the active - /// signature has since changed, the handler treats this job as stale and - /// finishes (a fresh chain for the new signature takes over). - pub signature: String, -} - -impl ReembedBackfillPayload { - /// Stable dedupe key — one in-flight backfill chain per signature. - pub fn dedupe_key(&self) -> String { - format!("reembed_backfill:{}", self.signature) - } -} - -/// Build (or re-build for a new version) one document's per-doc subtree and -/// merge its doc-root into the connection tree. Carries the full leaf chunk set -/// for the version so the seal handler can run the per-document side-cascade in -/// one shot. -#[derive(Clone, Debug, Serialize, Deserialize)] -pub struct SealDocumentPayload { - /// Connection-level source-tree scope, e.g. `notion:{connection_id}`. All - /// of a connection's documents share this one tree. - pub tree_scope: String, - /// Document identity (the chunk `source_id`, e.g. - /// `notion:{connection_id}:{page_id}`). - pub doc_id: String, - /// Document version as epoch-ms (`last_edited_time`). `None` for sources - /// that don't carry a version. - pub version_ms: Option, - /// Leaf chunk ids for this version, in document order. - pub chunk_ids: Vec, -} - -impl SealDocumentPayload { - /// Stable dedupe key — one in-flight seal per `(doc, version)`. A new - /// version gets a distinct key so it isn't suppressed by an older one. - pub fn dedupe_key(&self) -> String { - match self.version_ms { - Some(v) => format!("seal_doc:{}@{}", self.doc_id, v), - None => format!("seal_doc:{}", self.doc_id), - } - } -} - /// One row in `mem_tree_jobs`. `payload_json` is left as a raw string so /// callers parse it lazily based on `kind`. #[derive(Clone, Debug)] diff --git a/src/memory/queue/worker.rs b/src/memory/queue/worker.rs index 1ba3ef9..65076b1 100644 --- a/src/memory/queue/worker.rs +++ b/src/memory/queue/worker.rs @@ -27,11 +27,14 @@ use anyhow::Result; use crate::memory::config::MemoryConfig; use crate::memory::queue::gate::{LlmGate, Permit}; use crate::memory::queue::handlers::{self, QueueDelegates}; -use crate::memory::queue::store::{claim_next, purge_retired_jobs, DEFAULT_LOCK_DURATION_MS}; +use crate::memory::queue::ops::set_backfill_in_progress; +use crate::memory::queue::store::{claim_next, purge_retired_jobs}; use crate::memory::queue::store_settle::{ - mark_deferred, mark_done, mark_failed_typed, recover_stale_locks, + mark_deferred, mark_done_with_followups, mark_failed_typed, recover_stale_locks, +}; +use crate::memory::queue::types::{ + Job, JobFailure, JobKind, JobOutcome, JobStatus, NewJob, SealPayload, }; -use crate::memory::queue::types::{Job, JobFailure, JobOutcome}; /// Process-wide LLM concurrency gate. Single-slot by default (mirrors the /// upstream single-permit semaphore); LLM-bound jobs hold a permit for the @@ -60,7 +63,7 @@ pub fn bootstrap(config: &MemoryConfig) -> Result<(usize, usize)> { /// Claim and run a single job. Returns `true` when work was processed, `false` /// when no eligible row was available. /// -/// The job's lease ([`DEFAULT_LOCK_DURATION_MS`]) starts counting at claim +/// The job's configured lease starts counting at claim /// time, before the LLM gate is acquired below — a job that waits a long time /// for a free permit eats into its own lease window before its handler even /// starts. @@ -80,7 +83,7 @@ async fn run_once_with_gate( delegates: &dyn QueueDelegates, gate: &LlmGate, ) -> Result { - let Some(job) = claim_next(config, DEFAULT_LOCK_DURATION_MS)? else { + let Some(job) = claim_next(config, config.queue.lock_duration_ms)? else { return Ok(false); }; @@ -101,22 +104,74 @@ async fn run_once_with_gate( None }; - let result = handlers::handle_job(config, &job, delegates).await; + let result = handlers::plan_job(config, &job, delegates).await; drop(permit); - settle_job(config, &job, result)?; + settle_planned_job(config, &job, result)?; Ok(true) } -/// Translate a handler's [`Result`] into the matching store -/// settlement call (`mark_done` / `mark_deferred` / `mark_failed_typed`). +/// Translate a handler plan into the matching store settlement call. Successful +/// follow-up enqueues commit in the same transaction as `mark_done`. /// Errors from the settlement call itself (e.g. a busy database) propagate to /// the caller — the job's `running` row is then left for lease-expiry /// recovery rather than being retried inline here. +#[cfg(test)] fn settle_job(config: &MemoryConfig, job: &Job, result: Result) -> Result<()> { + settle_planned_job( + config, + job, + result.map(|outcome| handlers::JobPlan { + outcome, + follow_ups: Vec::new(), + }), + ) +} + +fn settle_planned_job( + config: &MemoryConfig, + job: &Job, + result: Result, +) -> Result<()> { match result { - Ok(JobOutcome::Done) => mark_done(config, job), - Ok(JobOutcome::Defer { until_ms, reason }) => { + Ok(handlers::JobPlan { + outcome: JobOutcome::Done, + follow_ups, + }) => { + let arms_reembed = follow_ups + .iter() + .any(|follow_up| follow_up.kind == JobKind::ReembedBackfill); + mark_done_with_followups(config, job, &follow_ups)?; + if arms_reembed { + set_backfill_in_progress(true); + } + + // A same-level seal enqueue can be suppressed while this row is + // running. Once its dedupe key is released, re-check the live + // buffer and restore the edge if content appended during the seal + // crossed the gate again. + if job.kind == JobKind::Seal + && matches!( + crate::memory::queue::store::get_job(config, &job.id)?, + Some(current) if current.status == JobStatus::Done + ) + { + let payload: SealPayload = serde_json::from_str(&job.payload_json)?; + let buffer = crate::memory::tree::store::get_buffer( + config, + &payload.tree_id, + payload.level, + )?; + if crate::memory::tree::should_seal(config, &buffer) { + crate::memory::queue::store::enqueue(config, &NewJob::seal(&payload)?)?; + } + } + Ok(()) + } + Ok(handlers::JobPlan { + outcome: JobOutcome::Defer { until_ms, reason }, + .. + }) => { // Defer is normal operation (transient blocker) — does NOT burn the // failure budget; `mark_deferred` reverts the claim's attempts bump. mark_deferred(config, job, until_ms, &reason) @@ -128,7 +183,21 @@ fn settle_job(config: &MemoryConfig, job: &Job, result: Result) -> R // burning the retry budget. let message = format!("{err:#}"); let typed = err.downcast_ref::(); - mark_failed_typed(config, job, &message, typed) + mark_failed_typed(config, job, &message, typed)?; + + // The handler clears this flag on every successful terminal + // outcome. Mirror that cleanup when a re-embed row exhausts its + // budget (or fails fast), otherwise retrieval remains in the + // process-wide "backfill pending" mode indefinitely. + if job.kind == JobKind::ReembedBackfill + && matches!( + crate::memory::queue::store::get_job(config, &job.id)?, + Some(current) if current.status == JobStatus::Failed + ) + { + set_backfill_in_progress(false); + } + Ok(()) } } } diff --git a/src/memory/queue/worker_tests.rs b/src/memory/queue/worker_tests.rs index 6ad6ef0..da48b49 100644 --- a/src/memory/queue/worker_tests.rs +++ b/src/memory/queue/worker_tests.rs @@ -2,6 +2,7 @@ use super::*; use crate::memory::config::MemoryConfig; use crate::memory::queue::gate::DEFAULT_LLM_PERMITS; use crate::memory::queue::handlers::ReembedProgress; +use crate::memory::queue::store::DEFAULT_LOCK_DURATION_MS; use crate::memory::queue::store::{count_by_status, enqueue, get_job}; use crate::memory::queue::test_support::RecordingDelegates; use crate::memory::queue::types::{FlushStalePayload, JobStatus, NewJob, ReembedBackfillPayload}; @@ -110,6 +111,76 @@ async fn run_once_reschedules_reembed_jobs_that_defer() { assert!(job.available_at_ms > chrono::Utc::now().timestamp_millis()); } +#[test] +fn terminal_reembed_failure_clears_process_backfill_flag() { + let (_tmp, cfg) = test_config(); + let mut new_job = NewJob::reembed_backfill(&ReembedBackfillPayload { + signature: "provider=test;model=terminal;dims=3".into(), + }) + .unwrap(); + new_job.max_attempts = Some(1); + enqueue(&cfg, &new_job).unwrap().unwrap(); + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + + crate::memory::queue::ops::set_backfill_in_progress(true); + settle_job(&cfg, &claimed, Err(anyhow::anyhow!("provider unavailable"))).unwrap(); + + assert_eq!( + get_job(&cfg, &claimed.id).unwrap().unwrap().status, + JobStatus::Failed + ); + assert!(!crate::memory::queue::ops::backfill_in_progress()); +} + +#[test] +fn completed_seal_rearms_same_level_when_live_buffer_still_crosses_gate() { + use crate::memory::chunks::with_connection; + use crate::memory::queue::types::SealPayload; + + let (_tmp, mut cfg) = test_config(); + cfg.tree.input_token_budget = 10; + let tree = crate::memory::tree::get_or_create_tree( + &cfg, + crate::memory::tree::TreeKind::Source, + "queue:edge", + ) + .unwrap(); + let payload = SealPayload { + tree_id: tree.id, + level: 0, + force_now_ms: None, + }; + let original_id = enqueue(&cfg, &NewJob::seal(&payload).unwrap()) + .unwrap() + .unwrap(); + let claimed = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS).unwrap().unwrap(); + + // Models content appended after the handler took and consumed its own + // snapshot, while this job's active dedupe key still suppressed producers. + with_connection(&cfg, |conn| { + conn.execute( + "INSERT INTO mem_tree_buffers + (tree_id, level, item_ids_json, token_sum, oldest_at_ms, updated_at_ms) + VALUES (?1, 0, '[\"late-chunk\"]', 10, 1, 1)", + rusqlite::params![payload.tree_id], + )?; + Ok(()) + }) + .unwrap(); + + settle_job(&cfg, &claimed, Ok(JobOutcome::Done)).unwrap(); + + assert_eq!( + get_job(&cfg, &original_id).unwrap().unwrap().status, + JobStatus::Done + ); + let replacement = claim_next(&cfg, DEFAULT_LOCK_DURATION_MS) + .unwrap() + .expect("same-level seal edge restored"); + assert_eq!(replacement.kind, crate::memory::queue::types::JobKind::Seal); + assert_ne!(replacement.id, original_id); +} + #[tokio::test] async fn run_once_holds_an_llm_permit_for_llm_bound_jobs() { // While a single worker is sequential, the gate must still be acquired and diff --git a/src/memory/sync/composio/client.rs b/src/memory/sync/composio/client.rs index 2039d2b..d2c4ce5 100644 --- a/src/memory/sync/composio/client.rs +++ b/src/memory/sync/composio/client.rs @@ -135,7 +135,6 @@ impl ComposioClient { .as_ref() .filter(|key| !key.is_empty()) .map(|key| key.expose().to_owned()) - .or_else(|| std::env::var("COMPOSIO_API_KEY").ok()) .filter(|key| !key.trim().is_empty()) .ok_or_else(|| anyhow::anyhow!("Composio direct API key is not configured"))?; let url = format!( diff --git a/src/memory/sync/composio/connect_tests.rs b/src/memory/sync/composio/connect_tests.rs index e24fdc8..f370303 100644 --- a/src/memory/sync/composio/connect_tests.rs +++ b/src/memory/sync/composio/connect_tests.rs @@ -3,6 +3,8 @@ use super::*; use serde_json::json; +use wiremock::matchers::{body_partial_json, header, method, path, query_param}; +use wiremock::{Mock, MockServer, ResponseTemplate}; #[test] fn generated_entity_id_is_prefixed_and_unique() { @@ -146,3 +148,158 @@ fn entity_store_load_tolerates_missing_and_corrupt_files() { std::fs::write(&corrupt, "{ not json ").unwrap(); assert_eq!(EntityStore::load(&corrupt).get("gmail"), None); } + +#[tokio::test] +async fn list_auth_configs_sends_filter_and_decodes_success() { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/api/v3/auth_configs")) + .and(query_param("toolkit_slug", "gmail")) + .and(header("x-api-key", "secret")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "items": [{"id": "ac_gmail", "toolkit": {"slug": "gmail"}}] + }))) + .mount(&server) + .await; + + let value = list_auth_configs( + &reqwest::Client::new(), + &format!("{}/api/v3/", server.uri()), + "secret", + Some("gmail"), + ) + .await + .unwrap(); + + assert_eq!( + resolve_auth_config_id(&value, "gmail").as_deref(), + Some("ac_gmail") + ); +} + +#[tokio::test] +async fn list_auth_configs_redacts_error_body_and_rejects_invalid_json() { + let failure = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/auth_configs")) + .respond_with(ResponseTemplate::new(401).set_body_string("echoed-secret")) + .mount(&failure) + .await; + let error = list_auth_configs( + &reqwest::Client::new(), + &failure.uri(), + "echoed-secret", + None, + ) + .await + .unwrap_err() + .to_string(); + assert!(error.contains("HTTP 401")); + assert!(!error.contains("echoed-secret")); + + let malformed = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/auth_configs")) + .respond_with(ResponseTemplate::new(200).set_body_string("not-json")) + .mount(&malformed) + .await; + assert!( + list_auth_configs(&reqwest::Client::new(), &malformed.uri(), "key", None) + .await + .unwrap_err() + .to_string() + .contains("decode failed") + ); +} + +#[tokio::test] +async fn create_connection_link_sends_trimmed_callback_and_extracts_fields() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/connected_accounts/link")) + .and(header("x-api-key", "secret")) + .and(body_partial_json(json!({ + "auth_config_id": "ac_1", + "user_id": "user_1", + "callback_url": "https://callback" + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "connected_account_id": "ca_1", + "redirect_url": "https://oauth" + }))) + .mount(&server) + .await; + + let link = create_connection_link( + &reqwest::Client::new(), + &server.uri(), + "secret", + "ac_1", + "user_1", + Some(" https://callback "), + ) + .await + .unwrap(); + assert_eq!(link.connected_account_id, "ca_1"); + assert_eq!(link.redirect_url.as_deref(), Some("https://oauth")); +} + +#[tokio::test] +async fn create_connection_link_rejects_http_decode_and_missing_id_failures() { + for (status, body, expected) in [ + (500, "server error", "HTTP 500"), + (200, "not-json", "decode failed"), + (200, "{}", "missing a connected account id"), + ] { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/connected_accounts/link")) + .respond_with(ResponseTemplate::new(status).set_body_string(body)) + .mount(&server) + .await; + let error = create_connection_link( + &reqwest::Client::new(), + &server.uri(), + "key", + "ac", + "user", + Some(" "), + ) + .await + .unwrap_err() + .to_string(); + assert!(error.contains(expected), "unexpected error: {error}"); + } +} + +#[tokio::test] +async fn get_connection_status_handles_success_and_safe_failures() { + let active = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/connected_accounts/ca_1")) + .and(header("x-api-key", "key")) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({"status": "ACTIVE"}))) + .mount(&active) + .await; + assert_eq!( + get_connection_status(&reqwest::Client::new(), &active.uri(), "key", "ca_1") + .await + .unwrap() + .as_deref(), + Some("ACTIVE") + ); + + for (status, body, expected) in [(403, "key", "HTTP 403"), (200, "invalid", "decode failed")] { + let server = MockServer::start().await; + Mock::given(method("GET")) + .and(path("/connected_accounts/ca_2")) + .respond_with(ResponseTemplate::new(status).set_body_string(body)) + .mount(&server) + .await; + let error = get_connection_status(&reqwest::Client::new(), &server.uri(), "key", "ca_2") + .await + .unwrap_err() + .to_string(); + assert!(error.contains(expected)); + } +} diff --git a/src/memory/sync/composio/providers/mod.rs b/src/memory/sync/composio/providers/mod.rs index 3cad8d1..2cbe5d1 100644 --- a/src/memory/sync/composio/providers/mod.rs +++ b/src/memory/sync/composio/providers/mod.rs @@ -6,6 +6,7 @@ mod github; mod linear; mod notion; mod slack; +mod slack_parse; pub use clickup::ClickUpSyncPipeline; pub use github::GitHubSyncPipeline; diff --git a/src/memory/sync/composio/providers/slack.rs b/src/memory/sync/composio/providers/slack.rs index 5dbea46..58cab7f 100644 --- a/src/memory/sync/composio/providers/slack.rs +++ b/src/memory/sync/composio/providers/slack.rs @@ -1,11 +1,13 @@ -use std::collections::{BTreeMap, HashMap}; -use std::sync::OnceLock; +use std::collections::HashMap; use async_trait::async_trait; use chrono::Utc; use serde_json::Value; use super::common::{checked_execute, document, first_array, pick_str}; +use super::slack_parse::{ + decode_cursors, next_cursor, parse_ts, replace_mentions, search_matches, search_total_pages, +}; use crate::memory::config::MemoryConfig; use crate::memory::sync::composio::{ run_incremental_sync, ActionExecutor, ComposioClient, IncrementalSource, PageFetch, SyncItem, @@ -448,73 +450,3 @@ async fn fetch_users( } users } - -fn mention_regex() -> &'static regex::Regex { - static REGEX: OnceLock = OnceLock::new(); - REGEX.get_or_init(|| regex::Regex::new(r"<@(U[A-Z0-9]+)>").expect("Slack mention regex")) -} - -fn replace_mentions(text: &str, users: Option<&serde_json::Map>) -> String { - mention_regex() - .replace_all(text, |captures: ®ex::Captures<'_>| { - let id = &captures[1]; - let resolved = users - .and_then(|users| users.get(id)) - .and_then(Value::as_str) - .unwrap_or(id); - format!("@{resolved}") - }) - .into_owned() -} - -fn next_cursor(data: &Value) -> Option { - [ - "/data/response_metadata/next_cursor", - "/response_metadata/next_cursor", - "/data/next_cursor", - "/next_cursor", - "/data/data/response_metadata/next_cursor", - ] - .iter() - .find_map(|path| data.pointer(path).and_then(Value::as_str)) - .map(str::trim) - .filter(|cursor| !cursor.is_empty()) - .map(str::to_owned) -} - -fn search_matches(data: &Value) -> Vec { - first_array( - data, - &[ - "/data/messages/matches", - "/messages/matches", - "/data/data/messages/matches", - "/messages", - ], - ) -} - -fn search_total_pages(data: &Value) -> u32 { - [ - "/data/messages/paging/pages", - "/messages/paging/pages", - "/data/data/messages/paging/pages", - "/pages", - ] - .iter() - .find_map(|path| data.pointer(path).and_then(Value::as_u64)) - .unwrap_or(1) as u32 -} - -fn decode_cursors(raw: Option<&str>) -> BTreeMap { - raw.and_then(|raw| serde_json::from_str(raw).ok()) - .unwrap_or_default() -} - -fn parse_ts(ts: &str) -> Option<(i64, u64)> { - let mut parts = ts.splitn(2, '.'); - Some(( - parts.next()?.parse().ok()?, - parts.next().unwrap_or("0").parse().ok()?, - )) -} diff --git a/src/memory/sync/composio/providers/slack_parse.rs b/src/memory/sync/composio/providers/slack_parse.rs new file mode 100644 index 0000000..0b9e0ef --- /dev/null +++ b/src/memory/sync/composio/providers/slack_parse.rs @@ -0,0 +1,81 @@ +//! Slack response cursor, mention, and timestamp parsing. + +use std::collections::BTreeMap; +use std::sync::OnceLock; + +use serde_json::Value; + +use super::common::first_array; + +pub(super) fn mention_regex() -> &'static regex::Regex { + static REGEX: OnceLock = OnceLock::new(); + REGEX.get_or_init(|| regex::Regex::new(r"<@(U[A-Z0-9]+)>").expect("Slack mention regex")) +} + +pub(super) fn replace_mentions( + text: &str, + users: Option<&serde_json::Map>, +) -> String { + mention_regex() + .replace_all(text, |captures: ®ex::Captures<'_>| { + let id = &captures[1]; + let resolved = users + .and_then(|users| users.get(id)) + .and_then(Value::as_str) + .unwrap_or(id); + format!("@{resolved}") + }) + .into_owned() +} + +pub(super) fn next_cursor(data: &Value) -> Option { + [ + "/data/response_metadata/next_cursor", + "/response_metadata/next_cursor", + "/data/next_cursor", + "/next_cursor", + "/data/data/response_metadata/next_cursor", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_str)) + .map(str::trim) + .filter(|cursor| !cursor.is_empty()) + .map(str::to_owned) +} + +pub(super) fn search_matches(data: &Value) -> Vec { + first_array( + data, + &[ + "/data/messages/matches", + "/messages/matches", + "/data/data/messages/matches", + "/messages", + ], + ) +} + +pub(super) fn search_total_pages(data: &Value) -> u32 { + [ + "/data/messages/paging/pages", + "/messages/paging/pages", + "/data/data/messages/paging/pages", + "/pages", + ] + .iter() + .find_map(|path| data.pointer(path).and_then(Value::as_u64)) + .unwrap_or(1) as u32 +} + +pub(super) fn decode_cursors(raw: Option<&str>) -> BTreeMap { + raw.and_then(|raw| serde_json::from_str(raw).ok()) + .unwrap_or_default() +} + +pub(super) fn parse_ts(ts: &str) -> Option<(i64, u64)> { + let mut parts = ts.splitn(2, '.'); + Some(( + parts.next()?.parse().ok()?, + parts.next().unwrap_or("0").parse().ok()?, + )) +} diff --git a/src/memory/sync/dispatcher.rs b/src/memory/sync/dispatcher.rs index 8f8e1f6..befcb34 100644 --- a/src/memory/sync/dispatcher.rs +++ b/src/memory/sync/dispatcher.rs @@ -132,6 +132,7 @@ mod tests { struct FakePipeline { id: &'static str, fail: bool, + init_fail: bool, } #[async_trait] @@ -143,6 +144,9 @@ mod tests { SyncPipelineKind::Workspace } async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + if self.init_fail { + anyhow::bail!("expected init failure") + } Ok(()) } async fn tick(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result { @@ -223,19 +227,22 @@ mod tests { .register(Arc::new(FakePipeline { id: "z-fail", fail: true, + init_fail: false, })) .unwrap(); dispatcher .register(Arc::new(FakePipeline { id: "a-ok", fail: false, + init_fail: false, })) .unwrap(); assert_eq!(dispatcher.ids(), vec!["a-ok", "z-fail"]); assert!(dispatcher .register(Arc::new(FakePipeline { id: "a-ok", - fail: false + fail: false, + init_fail: false, })) .is_err()); let results = dispatcher @@ -249,4 +256,73 @@ mod tests { .unwrap() .contains("expected failure")); } + + #[tokio::test] + async fn register_rejects_blank_ids_and_tick_reports_unknown_pipeline() { + let mut dispatcher = SyncDispatcher::new(); + assert!(dispatcher + .register(Arc::new(FakePipeline { + id: " ", + fail: false, + init_fail: false, + })) + .is_err()); + let error = dispatcher + .tick("missing", &MemoryConfig::new("/tmp/unused"), &context()) + .await + .unwrap_err(); + assert!(error.to_string().contains("unknown sync pipeline")); + } + + #[tokio::test] + async fn init_all_and_individual_tick_preserve_success_and_failure_details() { + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(Arc::new(FakePipeline { + id: "a-init-fails", + fail: false, + init_fail: true, + })) + .unwrap(); + dispatcher + .register(Arc::new(FakePipeline { + id: "b-ok", + fail: false, + init_fail: false, + })) + .unwrap(); + dispatcher + .register(Arc::new(FakePipeline { + id: "c-tick-fails", + fail: true, + init_fail: false, + })) + .unwrap(); + let config = MemoryConfig::new("/tmp/unused"); + let context = context(); + + let initialized = dispatcher.init_all(&config, &context).await; + assert!(initialized[0] + .error + .as_deref() + .unwrap() + .contains("expected init failure")); + assert!(initialized[1].outcome.is_some()); + assert_eq!( + dispatcher + .tick("b-ok", &config, &context) + .await + .unwrap() + .records_ingested, + 3 + ); + assert!(dispatcher + .tick("c-tick-fails", &config, &context) + .await + .is_err()); + + let encoded = serde_json::to_value(&initialized[0]).unwrap(); + assert_eq!(encoded["pipeline_id"], "a-init-fails"); + assert!(encoded.get("outcome").is_none()); + } } diff --git a/src/memory/sync/periodic.rs b/src/memory/sync/periodic.rs index 23dd899..7bd5210 100644 --- a/src/memory/sync/periodic.rs +++ b/src/memory/sync/periodic.rs @@ -13,7 +13,7 @@ pub const DEFAULT_SYNC_INTERVAL_SECS: u64 = 24 * 60 * 60; pub fn effective_interval_secs(configured: Option) -> Option { match configured { Some(0) => None, - Some(seconds) => Some(seconds.max(DEFAULT_SYNC_INTERVAL_SECS)), + Some(seconds) => Some(seconds), None => Some(DEFAULT_SYNC_INTERVAL_SECS), } } @@ -121,10 +121,7 @@ mod tests { #[test] fn cadence_handles_manual_minimum_and_persisted_success() { assert_eq!(effective_interval_secs(Some(0)), None); - assert_eq!( - effective_interval_secs(Some(60)), - Some(DEFAULT_SYNC_INTERVAL_SECS) - ); + assert_eq!(effective_interval_secs(Some(60)), Some(60)); let now = Utc::now(); let sources = vec![ source("new", SourceKind::Folder, true), diff --git a/src/memory/sync/rebuild.rs b/src/memory/sync/rebuild.rs index 73e5e8c..52245d3 100644 --- a/src/memory/sync/rebuild.rs +++ b/src/memory/sync/rebuild.rs @@ -200,6 +200,8 @@ pub async fn rebuild_tree_from_raw_with_audit( tree_kind: TreeKind::Source, target_level: 1, token_budget: config.tree.output_token_budget, + input_token_budget: config.tree.input_token_budget, + overhead_reserve_tokens: config.tree.summary_overhead_reserve_tokens, ask: tree.ask.as_deref(), }; let call = match summariser diff --git a/src/memory/tree/bucket_seal.rs b/src/memory/tree/bucket_seal.rs index 29f2eaf..efd00a6 100644 --- a/src/memory/tree/bucket_seal.rs +++ b/src/memory/tree/bucket_seal.rs @@ -21,148 +21,28 @@ //! progress events, seal-time embedding, and per-document subtree paths are not //! ported (deferred). Summary bodies are stored inline in `mem_tree_summaries`. -use std::collections::BTreeSet; -use std::sync::Arc; - use anyhow::{Context, Result}; use chrono::{DateTime, Utc}; -use futures::stream::{StreamExt, TryStreamExt}; use crate::memory::chunks::with_connection; use crate::memory::config::MemoryConfig; -use crate::memory::score::embed::Embedder; -use crate::memory::score::extract::EntityExtractor; -use crate::memory::score::resolver::canonicalise; use crate::memory::score::store::index_summary_entity_ids_tx; use crate::memory::store::content::{ slugify_source_id, stage_summary_with_layout, SummaryComposeInput, SummaryDiskLayout, SummaryTreeKind, }; use crate::memory::tree::hydrate::hydrate_inputs; +use crate::memory::tree::label_resolver::resolve_labels; use crate::memory::tree::registry::new_summary_id; use crate::memory::tree::store::{self, Buffer, SummaryNode, Tree}; -use crate::memory::tree::summarise::{fallback_summary, Summariser, SummaryContext, SummaryInput}; +use crate::memory::tree::summarise::{fallback_summary, Summariser, SummaryContext}; +pub(crate) use crate::memory::tree::types::NoopSealObserver; +pub use crate::memory::tree::types::{LabelStrategy, LeafRef, SealObserver, SealServices}; /// Hard cap on cascade depth — guards against runaway loops if token accounting /// ever slips. const MAX_CASCADE_DEPTH: u32 = 32; -/// Product callbacks around a seal. Engine state is already durable when -/// `summary_committed` runs; hosts use it for mirrors such as wiki-git. -pub trait SealObserver: Send + Sync { - fn progress(&self, _tree: &Tree, _step: &str, _level: u32, _item_count: Option) {} - fn summary_committed( - &self, - _tree: &Tree, - _node: &SummaryNode, - _content_path: &str, - _reason: &str, - ) -> Result<()> { - Ok(()) - } -} - -pub(crate) struct NoopSealObserver; -impl SealObserver for NoopSealObserver {} - -/// Injected compute and product notifications used by the crate-owned seal -/// pipeline. `embedder = None` deliberately persists a re-embeddable summary. -pub struct SealServices<'a> { - pub summariser: &'a dyn Summariser, - pub embedder: Option<&'a dyn Embedder>, - pub observer: &'a dyn SealObserver, -} - -/// How a sealed summary node's `entities` and `topics` fields get populated. -#[derive(Clone)] -pub enum LabelStrategy { - /// Run the extractor on the new summary's content; canonicalise the result - /// into `entities` (canonical_ids) and `topics` (labels). - ExtractFromContent(Arc), - /// Dedup-merge each input's `entities` and `topics` into the parent. - UnionFromChildren, - /// Leave both fields empty regardless of inputs. - Empty, -} - -impl std::fmt::Debug for LabelStrategy { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::ExtractFromContent(ex) => write!(f, "ExtractFromContent({})", ex.name()), - Self::UnionFromChildren => f.write_str("UnionFromChildren"), - Self::Empty => f.write_str("Empty"), - } - } -} - -/// Resolve `entities` and `topics` for a freshly-summarised node. -async fn resolve_labels( - strategy: &LabelStrategy, - inputs: &[SummaryInput], - summary_content: &str, -) -> Result<(Vec, Vec)> { - match strategy { - LabelStrategy::ExtractFromContent(extractor) => { - let extracted = extractor - .extract(summary_content) - .await - .context("seal-time extractor failed")?; - let canonical = canonicalise(&extracted); - let mut entities: Vec = canonical - .into_iter() - .map(|c| c.canonical_id) - .collect::>() - .into_iter() - .collect(); - entities.sort(); - let mut topics: Vec = extracted - .topics - .into_iter() - .map(|t| t.label) - .collect::>() - .into_iter() - .collect(); - topics.sort(); - Ok((entities, topics)) - } - LabelStrategy::UnionFromChildren => { - let mut entities: BTreeSet = BTreeSet::new(); - let mut topics: BTreeSet = BTreeSet::new(); - for inp in inputs { - entities.extend(inp.entities.iter().cloned()); - topics.extend(inp.topics.iter().cloned()); - } - Ok((entities.into_iter().collect(), topics.into_iter().collect())) - } - LabelStrategy::Empty => Ok((Vec::new(), Vec::new())), - } -} - -/// A single leaf being appended to an L0 buffer. -#[derive(Clone, Debug)] -pub struct LeafRef { - /// Persisted chunk id this leaf points at; used as the buffer `item_id` and - /// deduped on append. - pub chunk_id: String, - /// Chunk token count; added to the L0 buffer's `token_sum` to drive the - /// token-budget seal gate. - pub token_count: u32, - /// Chunk timestamp; folded into the buffer's `oldest_at` and the sealed - /// summary's time range. - pub timestamp: DateTime, - /// Raw chunk text, hydrated as a summariser input at seal time. - pub content: String, - /// Canonical entity ids carried up to the parent under - /// [`LabelStrategy::UnionFromChildren`]. - pub entities: Vec, - /// Topic labels carried up to the parent under - /// [`LabelStrategy::UnionFromChildren`]. - pub topics: Vec, - /// Chunk relevance score; the sealed summary takes the max over its inputs - /// (clamped to `>= 0.0`). - pub score: f32, -} - /// Append a leaf to `tree`, sealing buffers as they fill. Returns the ids of /// any summaries that sealed during this call. pub async fn append_leaf( @@ -210,6 +90,14 @@ pub fn append_to_buffer( ) -> Result<()> { with_connection(config, |conn| { let tx = conn.unchecked_transaction()?; + let status: String = tx + .query_row( + "SELECT status FROM mem_tree_trees WHERE id = ?1", + rusqlite::params![tree_id], + |row| row.get(0), + ) + .context("Failed to read tree status before append")?; + anyhow::ensure!(status == "active", "tree '{tree_id}' is archived"); let mut buf = store::get_buffer_conn(&tx, tree_id, level)?; if buf.item_ids.iter().any(|existing| existing == item_id) { return Ok(()); // retry after a failed cascade — no double count @@ -374,6 +262,8 @@ pub async fn seal_one_level_with_services( tree_kind: tree.kind, target_level, token_budget: budget, + input_token_budget: config.tree.input_token_budget, + overhead_reserve_tokens: config.tree.summary_overhead_reserve_tokens, ask: tree.ask.as_deref(), }; // Treat a blank summary the same as a hard error — fall back to the @@ -490,14 +380,14 @@ pub async fn seal_one_level_with_services( with_connection(config, move |conn| { let tx = conn.unchecked_transaction()?; - let current_max: u32 = tx + let (current_max, status): (u32, String) = tx .query_row( - "SELECT max_level FROM mem_tree_trees WHERE id = ?1", + "SELECT max_level, status FROM mem_tree_trees WHERE id = ?1", rusqlite::params![&tree_id], - |r| r.get::<_, i64>(0), + |r| Ok((r.get::<_, i64>(0)?.max(0) as u32, r.get(1)?)), ) - .map(|n| n.max(0) as u32) - .context("Failed to read current max_level for tree")?; + .context("Failed to read current state for tree")?; + anyhow::ensure!(status == "active", "tree '{tree_id}' is archived"); store::insert_staged_summary_tx(&tx, &node_for_tx, Some(&staged_for_tx), &signature)?; index_summary_entity_ids_tx( @@ -565,288 +455,7 @@ pub async fn seal_one_level_with_services( Ok(summary_id) } - -/// Level offset reserved for cross-document merge nodes. -pub const MERGE_LEVEL_BASE: u32 = 1_000; -const DOC_SUBTREE_MAX_FANIN: usize = 32; -const DOC_SUBTREE_SEAL_CONCURRENCY: usize = 8; - -/// Build one immutable document-version subtree and feed its root into the -/// shared cross-document merge tier. -pub async fn seal_document_subtree_with_services( - config: &MemoryConfig, - tree: &Tree, - doc_id: &str, - version_ms: Option, - chunk_ids: &[String], - services: &SealServices<'_>, - strategy: &LabelStrategy, -) -> Result { - if chunk_ids.is_empty() { - anyhow::bail!("seal_document_subtree: empty chunk set"); - } - log::debug!( - "[memory_tree:seal_document] enter chunks={} has_version={}", - chunk_ids.len(), - version_ms.is_some() - ); - let mut level = 0; - let mut current_ids = chunk_ids.to_vec(); - let doc_root = loop { - let batches = if level == 0 { - batch_leaves_by_token_budget(config, ¤t_ids)? - } else { - batch_by_count(¤t_ids, DOC_SUBTREE_MAX_FANIN) - }; - let batch_futures: Vec<_> = batches - .iter() - .map(|batch| { - seal_explicit_children( - config, tree, level, batch, doc_id, version_ms, services, strategy, - ) - }) - .collect(); - let nodes: Vec = futures::stream::iter(batch_futures) - .buffered(DOC_SUBTREE_SEAL_CONCURRENCY) - .try_collect() - .await?; - current_ids = nodes.iter().map(|node| node.id.clone()).collect(); - level += 1; - if current_ids.len() <= 1 { - break nodes - .into_iter() - .next() - .context("document seal produced no root")?; - } - }; - - append_to_buffer( - config, - &tree.id, - MERGE_LEVEL_BASE, - &doc_root.id, - doc_root.token_count as i64, - doc_root.time_range_start, - )?; - let root_id = doc_root.id.clone(); - let root_level = doc_root.level; - with_connection(config, move |connection| { - let transaction = connection.unchecked_transaction()?; - let (current_root, current_max): (Option, u32) = transaction.query_row( - "SELECT root_id, max_level FROM mem_tree_trees WHERE id = ?1", - [&tree.id], - |row| Ok((row.get(0)?, row.get(1)?)), - )?; - if current_root.as_deref() != Some(root_id.as_str()) || root_level > current_max { - store::update_tree_after_seal_tx( - &transaction, - &tree.id, - &root_id, - root_level, - Utc::now(), - )?; - } - transaction.commit()?; - Ok(()) - })?; - cascade_all_from_with_services( - config, - tree, - MERGE_LEVEL_BASE, - false, - services, - strategy, - false, - ) - .await?; - log::debug!("[memory_tree:seal_document] complete levels={level}"); - Ok(doc_root.id) -} - -fn batch_leaves_by_token_budget( - config: &MemoryConfig, - chunk_ids: &[String], -) -> Result>> { - let chunks = crate::memory::chunks::get_chunks_batch(config, chunk_ids)?; - let mut batches = Vec::new(); - let mut current = Vec::new(); - let mut tokens = 0_i64; - for id in chunk_ids { - let Some(chunk) = chunks.get(id) else { - continue; - }; - let next = chunk.token_count as i64; - if !current.is_empty() - && (tokens + next > config.tree.input_token_budget as i64 - || current.len() >= DOC_SUBTREE_MAX_FANIN) - { - batches.push(std::mem::take(&mut current)); - tokens = 0; - } - current.push(id.clone()); - tokens += next; - } - if !current.is_empty() { - batches.push(current); - } - if batches.is_empty() { - anyhow::bail!("seal_document_subtree: no resolvable chunks"); - } - Ok(batches) -} - -fn batch_by_count(ids: &[String], max: usize) -> Vec> { - ids.chunks(max.max(1)).map(<[String]>::to_vec).collect() -} - -#[allow(clippy::too_many_arguments)] -async fn seal_explicit_children( - config: &MemoryConfig, - tree: &Tree, - level: u32, - child_ids: &[String], - doc_id: &str, - version_ms: Option, - services: &SealServices<'_>, - strategy: &LabelStrategy, -) -> Result { - let target_level = level + 1; - let inputs = hydrate_inputs(config, level, child_ids)?; - if inputs.is_empty() { - anyhow::bail!("document seal has no hydrated inputs at level {level}"); - } - let time_range_start = inputs.iter().map(|i| i.time_range_start).min().unwrap(); - let time_range_end = inputs.iter().map(|i| i.time_range_end).max().unwrap(); - let score = inputs - .iter() - .map(|i| i.score) - .fold(f32::NEG_INFINITY, f32::max) - .max(0.0); - let output = if inputs.len() == 1 && inputs[0].token_count <= config.tree.output_token_budget { - crate::memory::tree::SummaryOutput { - content: inputs[0].content.clone(), - token_count: inputs[0].token_count, - ..Default::default() - } - } else { - let context = SummaryContext { - tree_id: &tree.id, - tree_kind: tree.kind, - target_level, - token_budget: config.tree.output_token_budget, - ask: tree.ask.as_deref(), - }; - match services.summariser.summarise(&inputs, &context).await { - Ok(output) if !output.content.trim().is_empty() => output, - _ => fallback_summary(&inputs, context.token_budget), - } - }; - let (entities, topics) = resolve_labels(strategy, &inputs, &output.content).await?; - let embedding = match services.embedder { - Some(embedder) if !output.content.trim().is_empty() => { - let input: String = output.content.chars().take(4_000).collect(); - Some( - embedder - .embed(&input) - .await - .context("embed document summary")?, - ) - } - _ => None, - }; - let now = Utc::now(); - let node = SummaryNode { - id: new_summary_id(target_level), - tree_id: tree.id.clone(), - tree_kind: tree.kind, - level: target_level, - parent_id: None, - child_ids: child_ids.to_vec(), - content: output.content, - token_count: output.token_count, - entities, - topics, - time_range_start, - time_range_end, - score, - sealed_at: now, - deleted: false, - embedding, - doc_id: Some(doc_id.to_string()), - version_ms, - }; - let summary_kind = match tree.kind { - crate::memory::tree::TreeKind::Source => SummaryTreeKind::Source, - crate::memory::tree::TreeKind::Topic => SummaryTreeKind::Topic, - crate::memory::tree::TreeKind::Global => SummaryTreeKind::Global, - crate::memory::tree::TreeKind::Flavoured => SummaryTreeKind::Flavoured, - }; - let doc_slug = slugify_source_id(doc_id); - let staged = stage_summary_with_layout( - &crate::memory::chunks::content_root(config), - &SummaryComposeInput { - summary_id: &node.id, - tree_kind: summary_kind, - tree_id: &node.tree_id, - tree_scope: &tree.scope, - level: node.level, - child_ids: &node.child_ids, - child_basenames: None, - child_count: node.child_ids.len(), - time_range_start: node.time_range_start, - time_range_end: node.time_range_end, - sealed_at: node.sealed_at, - body: &node.content, - }, - &slugify_source_id(&tree.scope), - SummaryDiskLayout::DocSubtree { - doc_slug: &doc_slug, - version_ms, - }, - )?; - let signature = crate::memory::chunks::tree_active_signature(config); - let node_for_tx = node.clone(); - let staged_for_tx = staged.clone(); - with_connection(config, move |connection| { - let transaction = connection.unchecked_transaction()?; - store::insert_staged_summary_tx( - &transaction, - &node_for_tx, - Some(&staged_for_tx), - &signature, - )?; - index_summary_entity_ids_tx( - &transaction, - &node_for_tx.entities, - &node_for_tx.id, - node_for_tx.score, - now.timestamp_millis(), - Some(&node_for_tx.tree_id), - )?; - for child_id in &node_for_tx.child_ids { - if level == 0 { - transaction.execute( - "UPDATE mem_tree_chunks SET parent_summary_id = ?1 WHERE id = ?2", - rusqlite::params![&node_for_tx.id, child_id], - )?; - } else { - transaction.execute( - "UPDATE mem_tree_summaries SET parent_id = ?1 WHERE id = ?2 AND parent_id IS NULL", - rusqlite::params![&node_for_tx.id, child_id], - )?; - } - } - transaction.commit()?; - Ok(()) - })?; - services.observer.summary_committed( - tree, - &node, - &staged.content_path, - "document_subtree_seal", - )?; - Ok(node) -} +pub use super::document_seal::{seal_document_subtree_with_services, MERGE_LEVEL_BASE}; #[cfg(test)] #[path = "bucket_seal_label_tests.rs"] diff --git a/src/memory/tree/bucket_seal_tests.rs b/src/memory/tree/bucket_seal_tests.rs index 1ecf98f..9a3dbd6 100644 --- a/src/memory/tree/bucket_seal_tests.rs +++ b/src/memory/tree/bucket_seal_tests.rs @@ -5,6 +5,7 @@ //! embedding assertions are dropped (those features are deferred). use super::*; +use crate::memory::tree::SummaryInput; use tempfile::TempDir; use crate::memory::chunks::upsert_chunks; @@ -103,6 +104,57 @@ async fn append_below_budget_does_not_seal() { assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); } +#[tokio::test] +async fn archived_tree_rejects_new_leaf() { + let (_tmp, cfg) = test_config(); + let tree = get_or_create_tree(&cfg, TreeKind::Source, "slack:#archived").unwrap(); + store::archive_tree(&cfg, &tree.id).unwrap(); + let summariser = ConcatSummariser::new(); + + let err = append_leaf( + &cfg, + &tree, + &mk_leaf("leaf-archived", 100, 1_700_000_000_000), + &summariser, + &LabelStrategy::Empty, + ) + .await + .unwrap_err(); + + assert!(format!("{err:#}").contains("archived")); + assert!(store::get_buffer(&cfg, &tree.id, 0).unwrap().is_empty()); +} + +#[tokio::test] +async fn archived_tree_rejects_seal_of_existing_buffer() { + let (_tmp, cfg) = test_config(); + let tree = get_or_create_tree(&cfg, TreeKind::Source, "slack:#archived").unwrap(); + let chunk = seed_chunk(&cfg, 0, "buffered before archive", 100, vec![]); + append_to_buffer( + &cfg, + &tree.id, + 0, + &chunk.id, + chunk.token_count as i64, + chunk.created_at, + ) + .unwrap(); + store::archive_tree(&cfg, &tree.id).unwrap(); + + let err = cascade_all_from( + &cfg, + &tree, + 0, + true, + &ConcatSummariser::new(), + &LabelStrategy::Empty, + ) + .await + .unwrap_err(); + assert!(format!("{err:#}").contains("archived")); + assert_eq!(store::count_summaries(&cfg, &tree.id).unwrap(), 0); +} + #[tokio::test] async fn service_seal_stages_body_and_enqueues_parent_atomically() { let (_tmp, mut cfg) = test_config(); diff --git a/src/memory/tree/document_seal.rs b/src/memory/tree/document_seal.rs new file mode 100644 index 0000000..9bb5982 --- /dev/null +++ b/src/memory/tree/document_seal.rs @@ -0,0 +1,304 @@ +//! Per-document immutable subtree sealing and merge-tier folding. + +use anyhow::{Context, Result}; +use chrono::Utc; +use futures::stream::{StreamExt, TryStreamExt}; + +use super::bucket_seal::{append_to_buffer, cascade_all_from_with_services}; +use super::hydrate::hydrate_inputs; +use super::label_resolver::resolve_labels; +use super::registry::new_summary_id; +use super::store::{self, SummaryNode, Tree}; +use super::summarise::{fallback_summary, SummaryContext}; +use super::types::{LabelStrategy, SealServices}; +use crate::memory::chunks::with_connection; +use crate::memory::config::MemoryConfig; +use crate::memory::score::store::index_summary_entity_ids_tx; +use crate::memory::store::content::{ + slugify_source_id, stage_summary_with_layout, SummaryComposeInput, SummaryDiskLayout, + SummaryTreeKind, +}; + +/// Level offset reserved for cross-document merge nodes. +pub const MERGE_LEVEL_BASE: u32 = 1_000; +const DOC_SUBTREE_MAX_FANIN: usize = 32; +const DOC_SUBTREE_SEAL_CONCURRENCY: usize = 8; + +/// Build one immutable document-version subtree and feed its root into the +/// shared cross-document merge tier. +pub async fn seal_document_subtree_with_services( + config: &MemoryConfig, + tree: &Tree, + doc_id: &str, + version_ms: Option, + chunk_ids: &[String], + services: &SealServices<'_>, + strategy: &LabelStrategy, +) -> Result { + if chunk_ids.is_empty() { + anyhow::bail!("seal_document_subtree: empty chunk set"); + } + log::debug!( + "[memory_tree:seal_document] enter chunks={} has_version={}", + chunk_ids.len(), + version_ms.is_some() + ); + let mut level = 0; + let mut current_ids = chunk_ids.to_vec(); + let doc_root = loop { + let batches = if level == 0 { + batch_leaves_by_token_budget(config, ¤t_ids)? + } else { + batch_by_count(¤t_ids, DOC_SUBTREE_MAX_FANIN) + }; + let batch_futures: Vec<_> = batches + .iter() + .map(|batch| { + seal_explicit_children( + config, tree, level, batch, doc_id, version_ms, services, strategy, + ) + }) + .collect(); + let nodes: Vec = futures::stream::iter(batch_futures) + .buffered(DOC_SUBTREE_SEAL_CONCURRENCY) + .try_collect() + .await?; + current_ids = nodes.iter().map(|node| node.id.clone()).collect(); + level += 1; + if current_ids.len() <= 1 { + break nodes + .into_iter() + .next() + .context("document seal produced no root")?; + } + }; + + append_to_buffer( + config, + &tree.id, + MERGE_LEVEL_BASE, + &doc_root.id, + doc_root.token_count as i64, + doc_root.time_range_start, + )?; + let root_id = doc_root.id.clone(); + let root_level = doc_root.level; + with_connection(config, move |connection| { + let transaction = connection.unchecked_transaction()?; + let (current_root, current_max): (Option, u32) = transaction.query_row( + "SELECT root_id, max_level FROM mem_tree_trees WHERE id = ?1", + [&tree.id], + |row| Ok((row.get(0)?, row.get(1)?)), + )?; + if current_root.as_deref() != Some(root_id.as_str()) || root_level > current_max { + store::update_tree_after_seal_tx( + &transaction, + &tree.id, + &root_id, + root_level, + Utc::now(), + )?; + } + transaction.commit()?; + Ok(()) + })?; + cascade_all_from_with_services( + config, + tree, + MERGE_LEVEL_BASE, + false, + services, + strategy, + false, + ) + .await?; + log::debug!("[memory_tree:seal_document] complete levels={level}"); + Ok(doc_root.id) +} + +fn batch_leaves_by_token_budget( + config: &MemoryConfig, + chunk_ids: &[String], +) -> Result>> { + let chunks = crate::memory::chunks::get_chunks_batch(config, chunk_ids)?; + let mut batches = Vec::new(); + let mut current = Vec::new(); + let mut tokens = 0_i64; + for id in chunk_ids { + let Some(chunk) = chunks.get(id) else { + continue; + }; + let next = chunk.token_count as i64; + if !current.is_empty() + && (tokens + next > config.tree.input_token_budget as i64 + || current.len() >= DOC_SUBTREE_MAX_FANIN) + { + batches.push(std::mem::take(&mut current)); + tokens = 0; + } + current.push(id.clone()); + tokens += next; + } + if !current.is_empty() { + batches.push(current); + } + if batches.is_empty() { + anyhow::bail!("seal_document_subtree: no resolvable chunks"); + } + Ok(batches) +} + +fn batch_by_count(ids: &[String], max: usize) -> Vec> { + ids.chunks(max.max(1)).map(<[String]>::to_vec).collect() +} + +#[allow(clippy::too_many_arguments)] +async fn seal_explicit_children( + config: &MemoryConfig, + tree: &Tree, + level: u32, + child_ids: &[String], + doc_id: &str, + version_ms: Option, + services: &SealServices<'_>, + strategy: &LabelStrategy, +) -> Result { + let target_level = level + 1; + let inputs = hydrate_inputs(config, level, child_ids)?; + if inputs.is_empty() { + anyhow::bail!("document seal has no hydrated inputs at level {level}"); + } + let time_range_start = inputs.iter().map(|i| i.time_range_start).min().unwrap(); + let time_range_end = inputs.iter().map(|i| i.time_range_end).max().unwrap(); + let score = inputs + .iter() + .map(|i| i.score) + .fold(f32::NEG_INFINITY, f32::max) + .max(0.0); + let output = if inputs.len() == 1 && inputs[0].token_count <= config.tree.output_token_budget { + crate::memory::tree::SummaryOutput { + content: inputs[0].content.clone(), + token_count: inputs[0].token_count, + ..Default::default() + } + } else { + let context = SummaryContext { + tree_id: &tree.id, + tree_kind: tree.kind, + target_level, + token_budget: config.tree.output_token_budget, + input_token_budget: config.tree.input_token_budget, + overhead_reserve_tokens: config.tree.summary_overhead_reserve_tokens, + ask: tree.ask.as_deref(), + }; + match services.summariser.summarise(&inputs, &context).await { + Ok(output) if !output.content.trim().is_empty() => output, + _ => fallback_summary(&inputs, context.token_budget), + } + }; + let (entities, topics) = resolve_labels(strategy, &inputs, &output.content).await?; + let embedding = match services.embedder { + Some(embedder) if !output.content.trim().is_empty() => { + let input: String = output.content.chars().take(4_000).collect(); + Some( + embedder + .embed(&input) + .await + .context("embed document summary")?, + ) + } + _ => None, + }; + let now = Utc::now(); + let node = SummaryNode { + id: new_summary_id(target_level), + tree_id: tree.id.clone(), + tree_kind: tree.kind, + level: target_level, + parent_id: None, + child_ids: child_ids.to_vec(), + content: output.content, + token_count: output.token_count, + entities, + topics, + time_range_start, + time_range_end, + score, + sealed_at: now, + deleted: false, + embedding, + doc_id: Some(doc_id.to_string()), + version_ms, + }; + let summary_kind = match tree.kind { + crate::memory::tree::TreeKind::Source => SummaryTreeKind::Source, + crate::memory::tree::TreeKind::Topic => SummaryTreeKind::Topic, + crate::memory::tree::TreeKind::Global => SummaryTreeKind::Global, + crate::memory::tree::TreeKind::Flavoured => SummaryTreeKind::Flavoured, + }; + let doc_slug = slugify_source_id(doc_id); + let staged = stage_summary_with_layout( + &crate::memory::chunks::content_root(config), + &SummaryComposeInput { + summary_id: &node.id, + tree_kind: summary_kind, + tree_id: &node.tree_id, + tree_scope: &tree.scope, + level: node.level, + child_ids: &node.child_ids, + child_basenames: None, + child_count: node.child_ids.len(), + time_range_start: node.time_range_start, + time_range_end: node.time_range_end, + sealed_at: node.sealed_at, + body: &node.content, + }, + &slugify_source_id(&tree.scope), + SummaryDiskLayout::DocSubtree { + doc_slug: &doc_slug, + version_ms, + }, + )?; + let signature = crate::memory::chunks::tree_active_signature(config); + let node_for_tx = node.clone(); + let staged_for_tx = staged.clone(); + with_connection(config, move |connection| { + let transaction = connection.unchecked_transaction()?; + store::insert_staged_summary_tx( + &transaction, + &node_for_tx, + Some(&staged_for_tx), + &signature, + )?; + index_summary_entity_ids_tx( + &transaction, + &node_for_tx.entities, + &node_for_tx.id, + node_for_tx.score, + now.timestamp_millis(), + Some(&node_for_tx.tree_id), + )?; + for child_id in &node_for_tx.child_ids { + if level == 0 { + transaction.execute( + "UPDATE mem_tree_chunks SET parent_summary_id = ?1 WHERE id = ?2", + rusqlite::params![&node_for_tx.id, child_id], + )?; + } else { + transaction.execute( + "UPDATE mem_tree_summaries SET parent_id = ?1 WHERE id = ?2 AND parent_id IS NULL", + rusqlite::params![&node_for_tx.id, child_id], + )?; + } + } + transaction.commit()?; + Ok(()) + })?; + services.observer.summary_committed( + tree, + &node, + &staged.content_path, + "document_subtree_seal", + )?; + Ok(node) +} diff --git a/src/memory/tree/flavoured.rs b/src/memory/tree/flavoured.rs index c43d762..3658122 100644 --- a/src/memory/tree/flavoured.rs +++ b/src/memory/tree/flavoured.rs @@ -6,8 +6,9 @@ //! prompt-ready markdown file: a style guide / preference profile a host can //! inject verbatim into a system prompt. //! -//! [`compile_flavoured_root`] fetches the tree's current root [`SummaryNode`], -//! clamps its body to [`TreeConfig::flavour_root_token_budget`] tokens, wraps it +//! [`compile_flavoured_root`] fetches the tree's current root +//! [`crate::memory::tree::SummaryNode`], clamps its body to +//! [`crate::memory::config::TreeConfig::flavour_root_token_budget`] tokens, wraps it //! in light front-matter (ask, tree id, scope, sealed-at, evidence changelog), //! and stages it at a stable, overwritten-in-place path //! (`flavoured/.md`) so hosts can read a fixed location. The engine @@ -51,7 +52,7 @@ pub fn flavoured_root_abs_path(config: &MemoryConfig, scope: &str) -> PathBuf { /// profile, stage it at [`flavoured_root_rel_path`] (overwriting any prior /// version in place), and return the full markdown (front-matter + body). /// -/// The body is the tree's root [`SummaryNode`] content clamped to +/// The body is the tree's root [`crate::memory::tree::SummaryNode`] content clamped to /// [`TreeConfig::flavour_root_token_budget`](crate::memory::config::TreeConfig::flavour_root_token_budget) /// tokens. Before the first seal (`root_id == None`) the body is empty; the /// artifact is still written so hosts always find the fixed path. diff --git a/src/memory/tree/flush.rs b/src/memory/tree/flush.rs index 02b6c6e..ba8117e 100644 --- a/src/memory/tree/flush.rs +++ b/src/memory/tree/flush.rs @@ -62,6 +62,7 @@ pub async fn flush_stale_buffers_with_services( let tree_by_id = store::get_trees_batch(config, &distinct_tree_ids)?; let mut seals = 0; + let mut failures = Vec::new(); for buf in stale { let Some(tree) = tree_by_id.get(&buf.tree_id) else { continue; // orphan buffer — tree row gone @@ -69,8 +70,18 @@ pub async fn flush_stale_buffers_with_services( let sealed = cascade_all_from_with_services( config, tree, buf.level, true, services, strategy, false, ) - .await?; - seals += sealed.len(); + .await; + match sealed { + Ok(ids) => seals += ids.len(), + Err(err) => failures.push(format!("{} level {}: {err:#}", buf.tree_id, buf.level)), + } + } + if !failures.is_empty() { + anyhow::bail!( + "failed to flush {} stale buffer(s) after sealing {seals}: {}", + failures.len(), + failures.join("; ") + ); } Ok(seals) } diff --git a/src/memory/tree/flush_tests.rs b/src/memory/tree/flush_tests.rs index 47eabb0..5d206bb 100644 --- a/src/memory/tree/flush_tests.rs +++ b/src/memory/tree/flush_tests.rs @@ -13,6 +13,25 @@ use crate::memory::tree::registry::get_or_create_tree; use crate::memory::tree::store::{upsert_buffer_tx, Buffer, TreeKind}; use crate::memory::tree::summarise::ConcatSummariser; +struct RejectTreeObserver { + tree_id: String, +} + +impl crate::memory::tree::bucket_seal::SealObserver for RejectTreeObserver { + fn summary_committed( + &self, + tree: &crate::memory::tree::store::Tree, + _node: &crate::memory::tree::store::SummaryNode, + _content_path: &str, + _reason: &str, + ) -> anyhow::Result<()> { + if tree.id == self.tree_id { + anyhow::bail!("injected observer failure") + } + Ok(()) + } +} + fn test_config() -> (TempDir, MemoryConfig) { let tmp = TempDir::new().unwrap(); let cfg = MemoryConfig::new(tmp.path()); @@ -170,6 +189,56 @@ async fn flush_seals_multiple_distinct_trees_via_batched_lookup() { assert_eq!(store::count_summaries(&cfg, &tree_b.id).unwrap(), 1); } +#[tokio::test] +async fn flush_continues_after_one_tree_fails() { + let (_tmp, cfg) = test_config(); + let tree_a = get_or_create_tree(&cfg, TreeKind::Source, "slack:#poison").unwrap(); + let tree_b = get_or_create_tree(&cfg, TreeKind::Source, "slack:#healthy").unwrap(); + let summariser = ConcatSummariser::new(); + let old_ts = Utc::now() - Duration::days(10); + + for (src, tree) in [("slack:#poison", &tree_a), ("slack:#healthy", &tree_b)] { + let chunk = seed_chunk(&cfg, src, 0, src, old_ts); + append_leaf( + &cfg, + tree, + &LeafRef { + chunk_id: chunk.id, + token_count: 100, + timestamp: old_ts, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 0.5, + }, + &summariser, + &LabelStrategy::Empty, + ) + .await + .unwrap(); + } + + let observer = RejectTreeObserver { + tree_id: tree_a.id.clone(), + }; + let services = SealServices { + summariser: &summariser, + embedder: None, + observer: &observer, + }; + let err = flush_stale_buffers_with_services( + &cfg, + Duration::days(7), + &services, + &LabelStrategy::Empty, + ) + .await + .expect_err("one tree should report the injected failure"); + + assert!(format!("{err:#}").contains("injected observer failure")); + assert_eq!(store::count_summaries(&cfg, &tree_b.id).unwrap(), 1); +} + #[tokio::test] async fn flush_default_noops_when_no_stale_buffers_exist() { let (_tmp, cfg) = test_config(); diff --git a/src/memory/tree/label_resolver.rs b/src/memory/tree/label_resolver.rs new file mode 100644 index 0000000..4f55de6 --- /dev/null +++ b/src/memory/tree/label_resolver.rs @@ -0,0 +1,51 @@ +//! Entity/topic label resolution for newly sealed summaries. + +use anyhow::{Context, Result}; +use std::collections::BTreeSet; + +use super::summarise::SummaryInput; +use super::types::LabelStrategy; +use crate::memory::score::resolver::canonicalise; + +/// Resolve `entities` and `topics` for a freshly-summarised node. +pub(super) async fn resolve_labels( + strategy: &LabelStrategy, + inputs: &[SummaryInput], + summary_content: &str, +) -> Result<(Vec, Vec)> { + match strategy { + LabelStrategy::ExtractFromContent(extractor) => { + let extracted = extractor + .extract(summary_content) + .await + .context("seal-time extractor failed")?; + let canonical = canonicalise(&extracted); + let mut entities: Vec = canonical + .into_iter() + .map(|c| c.canonical_id) + .collect::>() + .into_iter() + .collect(); + entities.sort(); + let mut topics: Vec = extracted + .topics + .into_iter() + .map(|t| t.label) + .collect::>() + .into_iter() + .collect(); + topics.sort(); + Ok((entities, topics)) + } + LabelStrategy::UnionFromChildren => { + let mut entities: BTreeSet = BTreeSet::new(); + let mut topics: BTreeSet = BTreeSet::new(); + for inp in inputs { + entities.extend(inp.entities.iter().cloned()); + topics.extend(inp.topics.iter().cloned()); + } + Ok((entities.into_iter().collect(), topics.into_iter().collect())) + } + LabelStrategy::Empty => Ok((Vec::new(), Vec::new())), + } +} diff --git a/src/memory/tree/mod.rs b/src/memory/tree/mod.rs index 56bc937..bdff43e 100644 --- a/src/memory/tree/mod.rs +++ b/src/memory/tree/mod.rs @@ -4,7 +4,8 @@ //! Two engines live here: //! //! - **Bucket-seal SQLite trees** ([`store`], [`bucket_seal`], [`flush`], -//! [`registry`], [`factory`], [`io`], [`read`], [`summarise`], [`hydrate`]). +//! [`registry`], [`factory`], [`io`], [`read`], [`summarise`], and internal +//! hydration helpers). //! Source/topic/global trees keyed by `(kind, scope)` in `mem_tree_trees`. //! Leaves accumulate in per-`(tree, level)` buffers; when a buffer crosses its //! token (L0) or fan-in (L≥1) gate it seals into an immutable @@ -26,23 +27,25 @@ pub mod store; pub mod bucket_seal; mod direct_ingest; +mod document_seal; pub mod factory; pub mod flavoured; pub mod flush; mod hydrate; pub mod io; +mod label_resolver; pub mod read; pub mod registry; pub mod runtime; pub mod summarise; +pub mod types; // ── Public API surface ────────────────────────────────────────────────────── pub use bucket_seal::{ append_leaf, append_leaf_deferred, append_to_buffer, cascade_all_from, cascade_all_from_with_services, seal_document_subtree_with_services, - seal_one_level_with_services, should_seal, LabelStrategy, LeafRef, SealObserver, SealServices, - MERGE_LEVEL_BASE, + seal_one_level_with_services, should_seal, MERGE_LEVEL_BASE, }; pub use direct_ingest::{ingest_summary, SummaryIngestInput, SummaryIngestOutcome}; pub use factory::{TreeFactory, TreeProfile, GLOBAL_SCOPE}; @@ -69,3 +72,4 @@ pub use summarise::{ fallback_summary, finish_provider_summary, prepare_summary_prompt, ConcatSummariser, PreparedSummaryPrompt, Summariser, SummaryCall, SummaryContext, SummaryInput, SummaryOutput, }; +pub use types::{LabelStrategy, LeafRef, SealObserver, SealServices}; diff --git a/src/memory/tree/runtime/engine.rs b/src/memory/tree/runtime/engine.rs index 7204fe2..11ad862 100644 --- a/src/memory/tree/runtime/engine.rs +++ b/src/memory/tree/runtime/engine.rs @@ -7,17 +7,15 @@ //! through); event-bus progress events and the background hourly loop are not //! ported. -use std::collections::BTreeMap; - use anyhow::{Context, Result}; use async_trait::async_trait; use chrono::{DateTime, Utc}; -use super::store; -use super::types::{ - derive_node_ids, derive_parent_id, estimate_tokens, level_from_node_id, NodeLevel, TreeNode, - TreeStatus, +use super::fold::{ + clear_pending_fold_receipts, group_by_hour, pending_fold_receipt, PendingFoldReceipt, }; +use super::store; +use super::types::{derive_parent_id, estimate_tokens, NodeLevel, TreeNode, TreeStatus}; use crate::memory::config::MemoryConfig; /// Maximum characters for a summary response (hard cap after the summariser call). @@ -58,13 +56,9 @@ impl RuntimeObserver for NoopObserver {} /// step succeeds, so a transient failure leaves the raw entries in place for /// the next run to retry. /// -/// # NOTE: retry after a partial failure can double-fold entries (`TR-11`) -/// Hour leaves are always written unconditionally before propagation runs; if -/// only an upper-level propagation fails (`_e` below is discarded, so the -/// specific failure is not surfaced to the caller), the next call re-reads the -/// buffer (not yet deleted), re-appends its content onto the *already updated* -/// hour leaf (`to_summarize` prepends `existing_summary`), and folds it a -/// second time. See `docs/spec/audit/03-tree-archivist-conversations.md`. +/// Hour leaves carry an internal receipt for the buffer filenames already +/// folded into them. If upper-level propagation fails, a retry skips those +/// entries while still incorporating entries appended after the failed run. /// /// `_ts` is currently unused — hour bucketing is derived from each buffer /// entry's own filename timestamp, not from this parameter. @@ -84,6 +78,7 @@ pub async fn run_summarization_observed( _ts: DateTime, observer: &dyn RuntimeObserver, ) -> Result> { + super::rebuild_fs::recover_interrupted_swap(config, namespace)?; let buffered = store::buffer_read(config, namespace)?; if buffered.is_empty() { return Ok(None); @@ -93,42 +88,78 @@ pub async fn run_summarization_observed( let mut all_propagation_ids: Vec<(String, NodeLevel)> = Vec::new(); let mut last_hour_node: Option = None; - - for (hour_id, entries) in &hour_groups { - let combined = entries.join("\n\n---\n\n"); - let (existing_summary, existing_created_at) = - match store::read_node(config, namespace, hour_id)? { - Some(existing) => (Some(existing.summary), Some(existing.created_at)), - None => (None, None), + let mut pending_hour_ids = Vec::new(); + + for (hour_id, group) in &hour_groups { + let existing = store::read_node(config, namespace, hour_id)?; + let receipt = existing + .as_ref() + .and_then(|node| pending_fold_receipt(node.metadata.as_deref())); + let already_applied: std::collections::HashSet<&str> = receipt + .as_ref() + .map(|receipt| { + receipt + .buffer_filenames + .iter() + .map(String::as_str) + .collect() + }) + .unwrap_or_default(); + let new_entries = group + .entries + .iter() + .filter(|(filename, _)| !already_applied.contains(filename.as_str())) + .collect::>(); + let new_content = new_entries + .iter() + .map(|(_, content)| content.as_str()) + .collect::>() + .join("\n\n---\n\n"); + + let hour_node = if new_entries.is_empty() { + existing.context("pending hour receipt exists without an hour node")? + } else { + let to_summarize = match existing.as_ref() { + Some(previous) => format!("{}\n\n---\n\n{new_content}", previous.summary), + None => new_content, }; - let to_summarize = match existing_summary { - Some(prev) => format!("{prev}\n\n---\n\n{combined}"), - None => combined, - }; - let hour_summary = summarize_to_limit( - summariser, - &to_summarize, - NodeLevel::Hour.max_tokens(), - "hour", - hour_id, - ) - .await - .context("summarize hour leaf")?; - - let now = Utc::now(); - let hour_node = TreeNode { - node_id: hour_id.clone(), - namespace: namespace.to_string(), - level: NodeLevel::Hour, - parent_id: derive_parent_id(hour_id), - summary: hour_summary.clone(), - token_count: estimate_tokens(&hour_summary), - child_count: 0, - created_at: existing_created_at.unwrap_or(now), - updated_at: now, - metadata: None, + let hour_summary = summarize_to_limit( + summariser, + &to_summarize, + NodeLevel::Hour.max_tokens(), + "hour", + hour_id, + ) + .await + .context("summarize hour leaf")?; + let now = Utc::now(); + let previous_metadata = receipt + .as_ref() + .and_then(|receipt| receipt.previous_metadata.clone()) + .or_else(|| existing.as_ref().and_then(|node| node.metadata.clone())); + let metadata = serde_json::to_string(&PendingFoldReceipt { + buffer_filenames: group + .entries + .iter() + .map(|(filename, _)| filename.clone()) + .collect(), + previous_metadata, + })?; + let node = TreeNode { + node_id: hour_id.clone(), + namespace: namespace.to_string(), + level: NodeLevel::Hour, + parent_id: derive_parent_id(hour_id), + summary: hour_summary.clone(), + token_count: estimate_tokens(&hour_summary), + child_count: 0, + created_at: existing.as_ref().map(|node| node.created_at).unwrap_or(now), + updated_at: now, + metadata: Some(metadata), + }; + store::write_node(config, &node)?; + node }; - store::write_node(config, &hour_node)?; observer.hour_completed(namespace, hour_id, hour_node.token_count); let (_, day_id, month_id, year_id, root_id) = derive_node_ids_from_hour_id(hour_id); @@ -137,6 +168,7 @@ pub async fn run_summarization_observed( all_propagation_ids.push((year_id, NodeLevel::Year)); all_propagation_ids.push((root_id, NodeLevel::Root)); last_hour_node = Some(hour_node); + pending_hour_ids.push(hour_id.clone()); } // Propagate bottom-up; a single node's failure does not void the whole run. @@ -164,23 +196,15 @@ pub async fn run_summarization_observed( if failed.is_empty() { store::buffer_delete(config, namespace, &buffer_filenames) .context("delete buffer entries after successful summarization")?; + clear_pending_fold_receipts(config, namespace, &pending_hour_ids)?; } Ok(last_hour_node) } /// Rebuild the entire tree from hour leaves upward, preserving unsummarised -/// buffer content. -/// -/// # NOTE: not crash-safe (`TR-2`) -/// This deletes the whole on-disk tree directory ([`store::delete_tree`]) and -/// then rewrites every hour leaf from the in-memory `hour_leaves` vec. A crash -/// between the delete and the full rewrite permanently loses every summary in -/// the namespace — there is no atomic swap (rebuild-into-temp + rename). The -/// buffer-preservation dance around it has the same gap: if the process -/// crashes after the rename to `tree_buffer_backup` but before the restore -/// rename back, the backup directory is left orphaned — no code path on a -/// later run adopts it, so buffered (unsummarised) content is stranded outside -/// the active buffer dir. See `docs/spec/audit/03-tree-archivist-conversations.md`. +/// buffer content. The replacement is built in a staging workspace and then +/// published with a recoverable directory swap, so a process crash cannot +/// expose a partially rebuilt tree or strand the live ingestion buffer. /// /// # Errors /// Propagates any filesystem error from the delete/rename/rewrite steps. @@ -201,6 +225,7 @@ pub async fn rebuild_tree_observed( namespace: &str, observer: &dyn RuntimeObserver, ) -> Result { + super::rebuild_fs::recover_interrupted_swap(config, namespace)?; let status = store::get_tree_status(config, namespace)?; if status.total_nodes == 0 { return Ok(status); @@ -208,36 +233,15 @@ pub async fn rebuild_tree_observed( let base = store::tree_dir(config, namespace); let mut hour_leaves: Vec = Vec::new(); - collect_hour_leaves_recursive(&base, namespace, "", &mut hour_leaves)?; + super::rebuild_fs::collect_hour_leaves(&base, namespace, "", &mut hour_leaves)?; if hour_leaves.is_empty() { return store::get_tree_status(config, namespace); } - // Preserve the buffer outside the tree dir while we delete + rebuild. - let buffer_path = store::buffer_dir(config, namespace); - let tree_base = store::tree_dir(config, namespace); - let buffer_backup = tree_base - .parent() - .unwrap_or(&tree_base) - .join("tree_buffer_backup"); - let buffer_existed = buffer_path.exists(); - if buffer_existed { - if buffer_backup.exists() { - std::fs::remove_dir_all(&buffer_backup)?; - } - std::fs::rename(&buffer_path, &buffer_backup).context("backup buffer before rebuild")?; - } - store::delete_tree(config, namespace)?; - if buffer_existed && buffer_backup.exists() { - let restored = store::buffer_dir(config, namespace); - if let Some(parent) = restored.parent() { - std::fs::create_dir_all(parent)?; - } - std::fs::rename(&buffer_backup, &restored).context("restore buffer after rebuild")?; - } + let staged = super::rebuild_fs::prepare(config, namespace)?; for leaf in &hour_leaves { - store::write_node(config, leaf)?; + store::write_node(&staged.config, leaf)?; } let mut day_ids = std::collections::BTreeSet::new(); @@ -259,7 +263,7 @@ pub async fn rebuild_tree_observed( // rebuild (the hour leaves are already re-written above). for day_id in &day_ids { let _ = propagate_node_observed( - config, + &staged.config, summariser, namespace, day_id, @@ -270,7 +274,7 @@ pub async fn rebuild_tree_observed( } for month_id in &month_ids { let _ = propagate_node_observed( - config, + &staged.config, summariser, namespace, month_id, @@ -281,7 +285,7 @@ pub async fn rebuild_tree_observed( } for year_id in &year_ids { let _ = propagate_node_observed( - config, + &staged.config, summariser, namespace, year_id, @@ -291,7 +295,7 @@ pub async fn rebuild_tree_observed( .await; } let _ = propagate_node_observed( - config, + &staged.config, summariser, namespace, "root", @@ -300,6 +304,7 @@ pub async fn rebuild_tree_observed( ) .await; + super::rebuild_fs::publish(staged)?; let status = store::get_tree_status(config, namespace)?; observer.rebuild_completed(namespace, status.total_nodes); Ok(status) @@ -399,28 +404,6 @@ async fn summarize_to_limit( Ok(response) } -/// Group buffer entries by hour from their filename timestamps. -pub(crate) fn group_by_hour(entries: &[(String, String)]) -> BTreeMap> { - let mut groups: BTreeMap> = BTreeMap::new(); - for (filename, content) in entries { - let hour_id = hour_id_from_buffer_filename(filename).unwrap_or_else(|| { - let (hour, _, _, _, _) = derive_node_ids(&Utc::now()); - hour - }); - groups.entry(hour_id).or_default().push(content.clone()); - } - groups -} - -/// Extract the hour node ID from a buffer filename like `1711972800000_abc.md`. -fn hour_id_from_buffer_filename(filename: &str) -> Option { - let ts_str = filename.split('_').next()?; - let millis: i64 = ts_str.parse().ok()?; - let dt = DateTime::from_timestamp_millis(millis)?; - let (hour_id, _, _, _, _) = derive_node_ids(&dt); - Some(hour_id) -} - /// Derive propagation IDs from an hour node_id like "2024/03/15/14". fn derive_node_ids_from_hour_id(hour_id: &str) -> (String, String, String, String, String) { let parts: Vec<&str> = hour_id.split('/').collect(); @@ -440,49 +423,6 @@ fn derive_node_ids_from_hour_id(hour_id: &str) -> (String, String, String, Strin } } -/// Recursively collect all hour leaf nodes from the tree directory. -fn collect_hour_leaves_recursive( - dir: &std::path::Path, - namespace: &str, - prefix: &str, - leaves: &mut Vec, -) -> Result<()> { - if !dir.exists() { - return Ok(()); - } - for entry in std::fs::read_dir(dir)? { - let entry = entry?; - let name = entry.file_name().to_string_lossy().to_string(); - let ft = entry.file_type()?; - if ft.is_dir() { - if name == "buffer" || name == "buffer_backup" { - continue; - } - let child_prefix = if prefix.is_empty() { - name.clone() - } else { - format!("{prefix}/{name}") - }; - collect_hour_leaves_recursive(&entry.path(), namespace, &child_prefix, leaves)?; - } else if ft.is_file() && name.ends_with(".md") && name != "summary.md" && name != "root.md" - { - let hour_part = name.trim_end_matches(".md"); - let node_id = if prefix.is_empty() { - hour_part.to_string() - } else { - format!("{prefix}/{hour_part}") - }; - if level_from_node_id(&node_id) == NodeLevel::Hour { - let raw = std::fs::read_to_string(entry.path())?; - let node = store::parse_node_markdown_pub(&raw, namespace, &node_id) - .with_context(|| format!("failed to parse hour leaf '{node_id}'"))?; - leaves.push(node); - } - } - } - Ok(()) -} - /// Discover namespaces that have pending buffer data. pub fn discover_active_namespaces(config: &MemoryConfig) -> Vec { let namespaces_dir = config.workspace.join("memory").join("namespaces"); diff --git a/src/memory/tree/runtime/engine_tests.rs b/src/memory/tree/runtime/engine_tests.rs index 41b78fe..9a32db3 100644 --- a/src/memory/tree/runtime/engine_tests.rs +++ b/src/memory/tree/runtime/engine_tests.rs @@ -65,6 +65,26 @@ struct FailAtLevelSummariser { fail_level: &'static str, reply: String, } + +#[derive(Default)] +struct ReceiptRetrySummariser { + hour_inputs: Mutex>, +} + +#[async_trait] +impl Summariser for ReceiptRetrySummariser { + async fn summarise(&self, system: Option<&str>, content: &str) -> anyhow::Result { + let system = system.unwrap_or(""); + if system.contains("at the hour level") { + self.hour_inputs.lock().unwrap().push(content.to_string()); + return Ok("word ".repeat(5_000)); + } + if system.contains("at the day level") { + anyhow::bail!("simulated day failure"); + } + Ok("ok".to_string()) + } +} #[async_trait] impl Summariser for FailAtLevelSummariser { async fn summarise(&self, system: Option<&str>, _content: &str) -> anyhow::Result { @@ -289,6 +309,39 @@ async fn run_summarization_multi_hour_groups_produce_multiple_leaves() { assert!(store::buffer_read(&cfg, ns).unwrap().is_empty()); } +#[tokio::test] +async fn retry_after_propagation_failure_does_not_double_fold_buffer() { + let tmp = TempDir::new().unwrap(); + let cfg = test_config(&tmp); + let ns = "receipt-retry"; + let ts = Utc.with_ymd_and_hms(2024, 3, 15, 8, 0, 0).unwrap(); + store::buffer_write(&cfg, ns, "only once", &ts, None).unwrap(); + store::buffer_write( + &cfg, + ns, + "also once", + &(ts + chrono::Duration::hours(1)), + None, + ) + .unwrap(); + let failing = ReceiptRetrySummariser::default(); + + run_summarization(&cfg, &failing, ns, ts).await.unwrap(); + run_summarization(&cfg, &failing, ns, ts).await.unwrap(); + + assert_eq!(failing.hour_inputs.lock().unwrap().len(), 2); + assert_eq!(store::buffer_read(&cfg, ns).unwrap().len(), 2); + + run_summarization(&cfg, &StubSummariser::with_reply("recovered"), ns, ts) + .await + .unwrap(); + assert!(store::buffer_read(&cfg, ns).unwrap().is_empty()); + let hour = store::read_node(&cfg, ns, "2024/03/15/08") + .unwrap() + .unwrap(); + assert!(hour.metadata.is_none(), "internal receipt must be cleared"); +} + #[tokio::test] async fn rebuild_tree_on_empty_namespace_is_noop() { let tmp = TempDir::new().unwrap(); diff --git a/src/memory/tree/runtime/fold.rs b/src/memory/tree/runtime/fold.rs new file mode 100644 index 0000000..0b4574e --- /dev/null +++ b/src/memory/tree/runtime/fold.rs @@ -0,0 +1,77 @@ +//! Retry receipts and buffer grouping for hour-leaf summarisation. + +use std::collections::BTreeMap; + +use anyhow::Result; +use chrono::{DateTime, Utc}; +use serde::{Deserialize, Serialize}; + +use super::store; +use super::types::derive_node_ids; +use crate::memory::config::MemoryConfig; + +#[derive(Debug, Default)] +pub(super) struct BufferedHour { + pub entries: Vec<(String, String)>, +} + +impl BufferedHour { + #[cfg(test)] + pub fn len(&self) -> usize { + self.entries.len() + } +} + +#[derive(Debug, Serialize, Deserialize)] +pub(super) struct PendingFoldReceipt { + pub buffer_filenames: Vec, + pub previous_metadata: Option, +} + +pub(super) fn pending_fold_receipt(metadata: Option<&str>) -> Option { + serde_json::from_str(metadata?).ok() +} + +pub(super) fn clear_pending_fold_receipts( + config: &MemoryConfig, + namespace: &str, + hour_ids: &[String], +) -> Result<()> { + for hour_id in hour_ids { + let Some(mut node) = store::read_node(config, namespace, hour_id)? else { + continue; + }; + let Some(receipt) = pending_fold_receipt(node.metadata.as_deref()) else { + continue; + }; + node.metadata = receipt.previous_metadata; + store::write_node(config, &node)?; + } + Ok(()) +} + +/// Group buffer entries by hour from their filename timestamps. +pub(super) fn group_by_hour(entries: &[(String, String)]) -> BTreeMap { + let mut groups: BTreeMap = BTreeMap::new(); + for (filename, content) in entries { + let hour_id = hour_id_from_buffer_filename(filename).unwrap_or_else(|| { + let (hour, _, _, _, _) = derive_node_ids(&Utc::now()); + hour + }); + groups + .entry(hour_id) + .or_default() + .entries + .push((filename.clone(), content.clone())); + } + groups +} + +/// Extract the hour node ID from a buffer filename like `1711972800000_abc.md`. +fn hour_id_from_buffer_filename(filename: &str) -> Option { + let ts_str = filename.split('_').next()?; + let millis: i64 = ts_str.parse().ok()?; + let dt = DateTime::from_timestamp_millis(millis)?; + let (hour_id, _, _, _, _) = derive_node_ids(&dt); + Some(hour_id) +} diff --git a/src/memory/tree/runtime/mod.rs b/src/memory/tree/runtime/mod.rs index 95c93ea..bdd87d6 100644 --- a/src/memory/tree/runtime/mod.rs +++ b/src/memory/tree/runtime/mod.rs @@ -9,6 +9,8 @@ //! behind [`engine::Summariser`]; the RPC/CLI/bus/schema surfaces are not ported. pub mod engine; +mod fold; +mod rebuild_fs; pub mod store; pub mod types; diff --git a/src/memory/tree/runtime/rebuild_fs.rs b/src/memory/tree/runtime/rebuild_fs.rs new file mode 100644 index 0000000..a62811c --- /dev/null +++ b/src/memory/tree/runtime/rebuild_fs.rs @@ -0,0 +1,162 @@ +//! Crash-safe filesystem staging and swap helpers for time-tree rebuilds. + +use std::path::{Path, PathBuf}; + +use anyhow::{Context, Result}; + +use super::store; +use super::types::{level_from_node_id, NodeLevel, TreeNode}; +use crate::memory::config::MemoryConfig; + +const STAGING_WORKSPACE: &str = ".memory-rebuild-staging"; +const BACKUP_DIR: &str = "tree.rebuild-backup"; + +pub(super) struct StagedTree { + pub config: MemoryConfig, + active: PathBuf, + staged: PathBuf, + backup: PathBuf, +} + +/// Restore an interrupted directory swap and adopt any buffer left in backup. +pub(super) fn recover_interrupted_swap(config: &MemoryConfig, namespace: &str) -> Result<()> { + let active = store::tree_dir(config, namespace); + let backup = active.with_file_name(BACKUP_DIR); + if !backup.exists() { + return Ok(()); + } + if !active.exists() { + std::fs::rename(&backup, &active).context("restore interrupted tree rebuild")?; + return Ok(()); + } + + let backup_buffer = backup.join("buffer"); + if backup_buffer.exists() { + let active_buffer = active.join("buffer"); + std::fs::create_dir_all(&active_buffer)?; + for entry in std::fs::read_dir(&backup_buffer)? { + let entry = entry?; + let destination = active_buffer.join(entry.file_name()); + if !destination.exists() { + std::fs::rename(entry.path(), destination)?; + } + } + } + std::fs::remove_dir_all(&backup).context("remove adopted tree rebuild backup")?; + Ok(()) +} + +/// Create an empty sibling workspace where a replacement tree can be built. +pub(super) fn prepare(config: &MemoryConfig, namespace: &str) -> Result { + recover_interrupted_swap(config, namespace)?; + let active = store::tree_dir(config, namespace); + let backup = active.with_file_name(BACKUP_DIR); + let mut staged_config = config.clone(); + staged_config.workspace = config.workspace.join(STAGING_WORKSPACE); + let staged = store::tree_dir(&staged_config, namespace); + if staged.exists() { + std::fs::remove_dir_all(&staged).context("remove stale tree rebuild staging directory")?; + } + Ok(StagedTree { + config: staged_config, + active, + staged, + backup, + }) +} + +/// Atomically publish a fully-built staged tree while preserving the latest +/// live ingestion buffer. Interrupted swaps are recoverable by +/// [`recover_interrupted_swap`]. +pub(super) fn publish(staged: StagedTree) -> Result<()> { + if let Some(parent) = staged.active.parent() { + std::fs::create_dir_all(parent)?; + } + if staged.backup.exists() { + std::fs::remove_dir_all(&staged.backup)?; + } + if staged.active.exists() { + std::fs::rename(&staged.active, &staged.backup) + .context("move active tree to rebuild backup")?; + } + if let Err(error) = std::fs::rename(&staged.staged, &staged.active) { + if !staged.active.exists() && staged.backup.exists() { + let _ = std::fs::rename(&staged.backup, &staged.active); + } + return Err(error).context("publish rebuilt tree"); + } + + // The staging tree never contains a buffer. Move the latest live buffer + // from the backup only after the new summary tree is visible. + recover_interrupted_swap_for_paths(&staged.active, &staged.backup) +} + +fn recover_interrupted_swap_for_paths(active: &Path, backup: &Path) -> Result<()> { + if !backup.exists() { + return Ok(()); + } + let backup_buffer = backup.join("buffer"); + if backup_buffer.exists() { + let active_buffer = active.join("buffer"); + std::fs::create_dir_all(&active_buffer)?; + for entry in std::fs::read_dir(&backup_buffer)? { + let entry = entry?; + let destination = active_buffer.join(entry.file_name()); + if !destination.exists() { + std::fs::rename(entry.path(), destination)?; + } + } + } + std::fs::remove_dir_all(backup).context("remove tree rebuild backup")?; + Ok(()) +} + +/// Recursively collect all hour leaf nodes from a tree directory. +pub(super) fn collect_hour_leaves( + dir: &Path, + namespace: &str, + prefix: &str, + leaves: &mut Vec, +) -> Result<()> { + if !dir.exists() { + return Ok(()); + } + for entry in std::fs::read_dir(dir)? { + let entry = entry?; + let name = entry.file_name().to_string_lossy().to_string(); + let file_type = entry.file_type()?; + if file_type.is_dir() { + if name == "buffer" || name == "buffer_backup" { + continue; + } + let child_prefix = if prefix.is_empty() { + name.clone() + } else { + format!("{prefix}/{name}") + }; + collect_hour_leaves(&entry.path(), namespace, &child_prefix, leaves)?; + } else if file_type.is_file() + && name.ends_with(".md") + && name != "summary.md" + && name != "root.md" + { + let hour = name.trim_end_matches(".md"); + let node_id = if prefix.is_empty() { + hour.to_string() + } else { + format!("{prefix}/{hour}") + }; + if level_from_node_id(&node_id) == NodeLevel::Hour { + let raw = std::fs::read_to_string(entry.path())?; + let node = store::parse_node_markdown_pub(&raw, namespace, &node_id) + .with_context(|| format!("failed to parse hour leaf '{node_id}'"))?; + leaves.push(node); + } + } + } + Ok(()) +} + +#[cfg(test)] +#[path = "rebuild_fs_tests.rs"] +mod tests; diff --git a/src/memory/tree/runtime/rebuild_fs_tests.rs b/src/memory/tree/runtime/rebuild_fs_tests.rs new file mode 100644 index 0000000..c2f247e --- /dev/null +++ b/src/memory/tree/runtime/rebuild_fs_tests.rs @@ -0,0 +1,48 @@ +use chrono::{TimeZone, Utc}; +use tempfile::TempDir; + +use super::*; + +fn config(tmp: &TempDir) -> MemoryConfig { + MemoryConfig::new(tmp.path().join("workspace")) +} + +#[test] +fn recovery_restores_backup_when_swap_stopped_before_publish() { + let tmp = TempDir::new().unwrap(); + let config = config(&tmp); + let namespace = "restore-backup"; + let timestamp = Utc.with_ymd_and_hms(2026, 7, 14, 5, 0, 0).unwrap(); + store::buffer_write(&config, namespace, "pending", ×tamp, None).unwrap(); + + let active = store::tree_dir(&config, namespace); + let backup = active.with_file_name(BACKUP_DIR); + std::fs::rename(&active, &backup).unwrap(); + + recover_interrupted_swap(&config, namespace).unwrap(); + + assert!(active.exists()); + assert!(!backup.exists()); + assert_eq!(store::buffer_read(&config, namespace).unwrap().len(), 1); +} + +#[test] +fn recovery_adopts_backup_buffer_after_new_tree_is_visible() { + let tmp = TempDir::new().unwrap(); + let config = config(&tmp); + let namespace = "adopt-buffer"; + let timestamp = Utc.with_ymd_and_hms(2026, 7, 14, 5, 0, 0).unwrap(); + store::buffer_write(&config, namespace, "pending", ×tamp, None).unwrap(); + + let active = store::tree_dir(&config, namespace); + let backup = active.with_file_name(BACKUP_DIR); + std::fs::rename(&active, &backup).unwrap(); + std::fs::create_dir_all(active.join("2026/07/14")).unwrap(); + std::fs::write(active.join("2026/07/14/05.md"), "replacement").unwrap(); + + recover_interrupted_swap(&config, namespace).unwrap(); + + assert!(active.join("2026/07/14/05.md").exists()); + assert!(!backup.exists()); + assert_eq!(store::buffer_read(&config, namespace).unwrap().len(), 1); +} diff --git a/src/memory/tree/runtime/store/buffer.rs b/src/memory/tree/runtime/store/buffer.rs index a7c1a61..1bc956e 100644 --- a/src/memory/tree/runtime/store/buffer.rs +++ b/src/memory/tree/runtime/store/buffer.rs @@ -94,25 +94,15 @@ pub fn buffer_drain(config: &MemoryConfig, namespace: &str) -> Result String { - let trimmed = raw.trim_start(); - if !trimmed.starts_with("---") { + let Some(after_open) = raw.strip_prefix("---\nmetadata: ") else { + return raw.to_string(); + }; + let Some((metadata, body)) = after_open.split_once("\n---\n") else { + return raw.to_string(); + }; + if serde_json::from_str::(metadata).is_err() { return raw.to_string(); } - let after_open = &trimmed[3..]; - if let Some(close_pos) = after_open.find("\n---") { - after_open[close_pos + 4..] - .trim_start_matches('\n') - .to_string() - } else { - raw.to_string() - } + body.strip_prefix('\n').unwrap_or(body).to_string() } diff --git a/src/memory/tree/runtime/store/mod.rs b/src/memory/tree/runtime/store/mod.rs index 53f1def..d40872c 100644 --- a/src/memory/tree/runtime/store/mod.rs +++ b/src/memory/tree/runtime/store/mod.rs @@ -4,10 +4,11 @@ //! `{workspace}/memory/namespaces/{namespace}/tree/`. The folder hierarchy //! mirrors the time hierarchy (`root.md`, `2024/summary.md`, //! `2024/03/15/14.md`). Ported from OpenHuman's `tree_runtime/store.rs` with -//! logging stripped and `Config` → [`MemoryConfig`]. Split across files to keep -//! each under the repo's 500-line cap: [`paths`] (path/validation), [`nodes`] -//! (node read/write + markdown parsing), [`scan`] (counts / status / collection), -//! and [`buffer`] (ingestion buffer). +//! logging stripped and `Config` → +//! [`crate::memory::config::MemoryConfig`]. Split across files to keep each +//! under the repo's 500-line cap: `paths` (path/validation), `nodes` (node +//! read/write + markdown parsing), `scan` (counts/status/collection), and +//! `buffer` (ingestion buffer). mod buffer; mod nodes; diff --git a/src/memory/tree/runtime/store/nodes.rs b/src/memory/tree/runtime/store/nodes.rs index 52f285d..e2b93db 100644 --- a/src/memory/tree/runtime/store/nodes.rs +++ b/src/memory/tree/runtime/store/nodes.rs @@ -24,13 +24,13 @@ pub fn write_node(config: &MemoryConfig, node: &TreeNode) -> Result<()> { .with_context(|| format!("create dirs for {}", parent.display()))?; } let metadata_line = match &node.metadata { - Some(m) => format!("metadata: {m}\n"), + Some(m) => format!("metadata: {}\n", yaml_string(m)), None => String::new(), }; let frontmatter = format!( "---\n\ - node_id: \"{}\"\n\ - namespace: \"{}\"\n\ + node_id: {}\n\ + namespace: {}\n\ level: {}\n\ parent_id: {}\n\ token_count: {}\n\ @@ -39,11 +39,11 @@ pub fn write_node(config: &MemoryConfig, node: &TreeNode) -> Result<()> { updated_at: {}\n\ {}\ ---\n\n", - node.node_id, - node.namespace, + yaml_string(&node.node_id), + yaml_string(&node.namespace), node.level.as_str(), match &node.parent_id { - Some(pid) => format!("\"{pid}\""), + Some(pid) => yaml_string(pid), None => "~".to_string(), }, node.token_count, @@ -60,6 +60,10 @@ pub fn write_node(config: &MemoryConfig, node: &TreeNode) -> Result<()> { Ok(()) } +fn yaml_string(value: &str) -> String { + serde_json::to_string(value).expect("serializing a string cannot fail") +} + /// Read a single tree node from its markdown file. `None` if it does not exist. pub fn read_node( config: &MemoryConfig, @@ -275,7 +279,9 @@ pub(crate) fn split_frontmatter(raw: &str) -> (HashMap, String) } if let Some(colon_pos) = line.find(':') { let key = line[..colon_pos].trim().to_string(); - let value = line[colon_pos + 1..].trim().trim_matches('"').to_string(); + let raw_value = line[colon_pos + 1..].trim(); + let value = serde_json::from_str::(raw_value) + .unwrap_or_else(|_| raw_value.trim_matches('"').to_string()); map.insert(key, value); } } diff --git a/src/memory/tree/runtime/store/paths.rs b/src/memory/tree/runtime/store/paths.rs index 234cc03..05eeaaa 100644 --- a/src/memory/tree/runtime/store/paths.rs +++ b/src/memory/tree/runtime/store/paths.rs @@ -25,23 +25,22 @@ pub fn node_file_path(config: &MemoryConfig, namespace: &str, node_id: &str) -> tree_dir(config, namespace).join(node_id_to_path(node_id)) } -/// Sanitise a namespace string for use as a directory name: trims whitespace, -/// maps each of `/ \ : * ? " < > | .` to `_`, then collapses `__` runs to `_`. -/// -/// # NOTE: not collision-free (`TR-15`) -/// This maps distinct namespaces onto the same directory name whenever they -/// differ only in which sanitised character produced a given `_`, e.g. -/// `"a/b"` and `"a.b"` both sanitise to `"a_b"`. The single-pass `replace("__", -/// "_")` also does not fully collapse triple-or-more underscore runs -/// consistently across inputs that already contained literal underscores. -/// Prefer length-prefixing or hex-encoding raw bytes for a collision-free -/// mapping (as `thread_messages_path` in `conversations` already does). See -/// `docs/spec/audit/03-tree-archivist-conversations.md`. +/// Sanitise a namespace string for use as a directory name. Safe names retain +/// their legacy spelling; transformed names carry a digest of the exact raw +/// namespace so distinct machine ids cannot alias the same directory. fn sanitize(namespace: &str) -> String { - namespace - .trim() - .replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|', '.'], "_") - .replace("__", "_") + let raw = namespace.trim(); + let sanitized = raw.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|', '.'], "_"); + if sanitized == raw { + return sanitized; + } + use sha2::{Digest, Sha256}; + let digest = Sha256::digest(raw.as_bytes()); + let suffix = digest[..6] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + format!("{sanitized}-{suffix}") } /// Validate a namespace string, erroring on empty / dangerous input. diff --git a/src/memory/tree/runtime/store/store_tests.rs b/src/memory/tree/runtime/store/store_tests.rs index 5ff2061..949b2e1 100644 --- a/src/memory/tree/runtime/store/store_tests.rs +++ b/src/memory/tree/runtime/store/store_tests.rs @@ -41,6 +41,22 @@ fn write_and_read_node_roundtrip() { assert!(read_back.parent_id.is_none()); } +#[test] +fn node_frontmatter_strings_escape_newlines_and_delimiters() { + let tmp = TempDir::new().unwrap(); + let config = test_config(&tmp); + let mut node = make_node("namespace\n---\nforged: true", "2024", "summary"); + node.metadata = Some("line one\n---\nnode_id: forged\\tail".into()); + write_node(&config, &node).unwrap(); + + let read = read_node(&config, &node.namespace, &node.node_id) + .unwrap() + .unwrap(); + assert_eq!(read.namespace, node.namespace); + assert_eq!(read.metadata, node.metadata); + assert_eq!(read.summary, node.summary); +} + #[test] fn write_and_read_hour_leaf() { let tmp = TempDir::new().unwrap(); @@ -166,6 +182,16 @@ fn buffer_write_with_metadata() { assert_eq!(drained[0].1, "entry with meta"); } +#[test] +fn buffer_content_starting_with_horizontal_rule_is_not_truncated() { + let tmp = TempDir::new().unwrap(); + let config = test_config(&tmp); + let content = "---\nuser-authored heading\n---\nbody"; + buffer_write(&config, "test-ns", content, &Utc::now(), None).unwrap(); + let drained = buffer_drain(&config, "test-ns").unwrap(); + assert_eq!(drained[0].1, content); +} + #[test] fn ancestors_walk_to_root() { let tmp = TempDir::new().unwrap(); @@ -239,6 +265,13 @@ fn validate_namespace_accepts_and_rejects() { } } +#[test] +fn transformed_namespace_paths_are_collision_resistant() { + let tmp = TempDir::new().unwrap(); + let cfg = test_config(&tmp); + assert_ne!(tree_dir(&cfg, "a/b"), tree_dir(&cfg, "a.b")); +} + #[test] fn list_namespaces_with_root_returns_only_summarised() { let tmp = TempDir::new().unwrap(); diff --git a/src/memory/tree/store/hotness_tests.rs b/src/memory/tree/store/hotness_tests.rs index 9479c6e..5d5125c 100644 --- a/src/memory/tree/store/hotness_tests.rs +++ b/src/memory/tree/store/hotness_tests.rs @@ -47,3 +47,68 @@ fn upsert_round_trip() { assert_eq!(get(&cfg, &c.entity_id).unwrap().unwrap(), c); assert_eq!(count(&cfg).unwrap(), 1); } + +#[test] +fn upsert_updates_existing_row_and_get_or_fresh_returns_persisted_value() { + let (_tmp, cfg) = test_config(); + let mut counters = HotnessCounters::fresh("topic:rust", 100); + upsert(&cfg, &counters).unwrap(); + counters.mention_count_30d = 7; + counters.query_hits_30d = 3; + counters.last_hotness = Some(4.5); + counters.last_updated_ms = 200; + upsert(&cfg, &counters).unwrap(); + + assert_eq!(count(&cfg).unwrap(), 1); + assert_eq!(get_or_fresh(&cfg, "topic:rust").unwrap(), counters); +} + +#[test] +fn distinct_sources_counts_unique_non_null_tree_ids() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + for (node, tree) in [ + ("one", Some("tree-a")), + ("two", Some("tree-a")), + ("three", Some("tree-b")), + ("four", None), + ] { + conn.execute( + "INSERT INTO mem_tree_entity_index + (entity_id,node_id,node_kind,entity_kind,surface,score,timestamp_ms,tree_id,is_user) + VALUES ('topic:rust',?1,'leaf','topic','rust',1.0,1,?2,0)", + rusqlite::params![node, tree], + )?; + } + Ok(()) + }) + .unwrap(); + + assert_eq!(distinct_sources_for(&cfg, "topic:rust").unwrap(), 2); + assert_eq!(distinct_sources_for(&cfg, "topic:missing").unwrap(), 0); +} + +#[test] +fn negative_legacy_counters_decode_fail_closed_to_zero() { + use crate::memory::chunks::with_connection; + + let (_tmp, cfg) = test_config(); + with_connection(&cfg, |conn| { + conn.execute( + "INSERT INTO mem_tree_entity_hotness + (entity_id,mention_count_30d,distinct_sources,last_seen_ms,query_hits_30d, + graph_centrality,ingests_since_check,last_hotness,last_updated_ms) + VALUES ('legacy',-1,-2,NULL,-3,NULL,-4,NULL,0)", + [], + )?; + Ok(()) + }) + .unwrap(); + let counters = get(&cfg, "legacy").unwrap().unwrap(); + assert_eq!(counters.mention_count_30d, 0); + assert_eq!(counters.distinct_sources, 0); + assert_eq!(counters.query_hits_30d, 0); + assert_eq!(counters.ingests_since_check, 0); +} diff --git a/src/memory/tree/store/mod.rs b/src/memory/tree/store/mod.rs index ae9d885..cfdcae5 100644 --- a/src/memory/tree/store/mod.rs +++ b/src/memory/tree/store/mod.rs @@ -37,7 +37,7 @@ pub use summaries::{ get_summary_embedding_for_signature, get_summary_embeddings_batch, get_summary_embeddings_for_signature_batch, insert_staged_summary_tx, insert_summary_tx, list_children_of_summary, list_summaries_at_level, list_summaries_in_window, - set_summary_embedding, set_summary_embedding_for_signature, + list_summaries_overlapping_window, set_summary_embedding, set_summary_embedding_for_signature, }; // ── Buffers ───────────────────────────────────────────────────────────────── diff --git a/src/memory/tree/store/store_tests.rs b/src/memory/tree/store/store_tests.rs index ab3f7ea..06d5bcc 100644 --- a/src/memory/tree/store/store_tests.rs +++ b/src/memory/tree/store/store_tests.rs @@ -433,94 +433,5 @@ fn get_trees_batch_returns_present_ids_and_skips_missing() { assert!(!map.contains_key("ghost")); } -#[test] -fn get_summaries_batch_returns_present_ids_and_skips_missing() { - let (_tmp, cfg) = test_config(); - insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap(); - let a = sample_summary("sum-a", "tree-1", 1); - let b = sample_summary("sum-b", "tree-1", 1); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &a, "test")?; - insert_summary_tx(&tx, &b, "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - assert!(get_summaries_batch(&cfg, &[]).unwrap().is_empty()); - let ids = vec![ - "sum-a".to_string(), - "sum-b".to_string(), - "ghost".to_string(), - ]; - let map = get_summaries_batch(&cfg, &ids).unwrap(); - assert_eq!(map.len(), 2); - assert_eq!(map.get("sum-a").unwrap(), &a); - assert_eq!(map.get("sum-b").unwrap(), &b); -} - -#[test] -fn summary_batch_embedding_lookup_returns_only_signature_scoped_rows() { - let (_tmp, cfg) = test_config(); - insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap(); - for sid in ["sum-1", "sum-2", "sum-3"] { - let node = sample_summary(sid, "tree-1", 1); - with_connection(&cfg, |conn| { - let tx = conn.unchecked_transaction()?; - insert_summary_tx(&tx, &node, "test")?; - tx.commit()?; - Ok(()) - }) - .unwrap(); - } - let sig_a = "openai/text-embedding-3-small@1536"; - let sig_b = "local/bge-small@384"; - set_summary_embedding_for_signature(&cfg, "sum-1", sig_a, &[0.1, 0.2]).unwrap(); - set_summary_embedding_for_signature(&cfg, "sum-2", sig_a, &[0.3, 0.4]).unwrap(); - set_summary_embedding_for_signature(&cfg, "sum-3", sig_b, &[0.5, 0.6, 0.7]).unwrap(); - - let ids = vec!["sum-1".into(), "sum-2".into(), "sum-3".into()]; - let map_a = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_a).unwrap(); - assert_eq!(map_a.len(), 2); - assert_eq!(map_a.get("sum-1").cloned(), Some(vec![0.1, 0.2])); - assert!(!map_a.contains_key("sum-3")); - - let map_b = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); - assert_eq!(map_b.len(), 1); - assert_eq!(map_b.get("sum-3").cloned(), Some(vec![0.5, 0.6, 0.7])); - - assert!(get_summary_embeddings_for_signature_batch(&cfg, &[], sig_a) - .unwrap() - .is_empty()); - assert!(get_summary_embeddings_batch(&cfg, &ids).unwrap().is_empty()); -} - -#[test] -fn list_trees_by_kind_and_archive() { - let (_tmp, cfg) = test_config(); - // Distinct created_at so the `ORDER BY created_at_ms ASC` is deterministic. - let mut s1 = sample_tree("source-1", "chat:slack:#eng"); - s1.created_at = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); - insert_tree(&cfg, &s1).unwrap(); - let mut topic = sample_tree("topic-1", "person:alice"); - topic.kind = TreeKind::Topic; - insert_tree(&cfg, &topic).unwrap(); - let mut s2 = sample_tree("source-2", "chat:discord:#ops"); - s2.created_at = Utc.timestamp_millis_opt(1_700_000_001_000).unwrap(); - insert_tree(&cfg, &s2).unwrap(); - - let source_ids: Vec = list_trees_by_kind(&cfg, TreeKind::Source) - .unwrap() - .into_iter() - .map(|t| t.id) - .collect(); - assert_eq!(source_ids, vec!["source-1", "source-2"]); - assert_eq!( - list_trees_by_kind(&cfg, TreeKind::Topic).unwrap()[0].id, - "topic-1" - ); - - archive_tree(&cfg, "source-1").unwrap(); - let archived = get_tree(&cfg, "source-1").unwrap().unwrap(); - assert_eq!(archived.status, TreeStatus::Archived); -} +#[path = "store_tests_more.rs"] +mod more; diff --git a/src/memory/tree/store/store_tests_more.rs b/src/memory/tree/store/store_tests_more.rs new file mode 100644 index 0000000..f77ad6a --- /dev/null +++ b/src/memory/tree/store/store_tests_more.rs @@ -0,0 +1,93 @@ +use super::*; + +#[test] +fn get_summaries_batch_returns_present_ids_and_skips_missing() { + let (_tmp, cfg) = test_config(); + insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap(); + let a = sample_summary("sum-a", "tree-1", 1); + let b = sample_summary("sum-b", "tree-1", 1); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + insert_summary_tx(&tx, &a, "test")?; + insert_summary_tx(&tx, &b, "test")?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + assert!(get_summaries_batch(&cfg, &[]).unwrap().is_empty()); + let ids = vec![ + "sum-a".to_string(), + "sum-b".to_string(), + "ghost".to_string(), + ]; + let map = get_summaries_batch(&cfg, &ids).unwrap(); + assert_eq!(map.len(), 2); + assert_eq!(map.get("sum-a").unwrap(), &a); + assert_eq!(map.get("sum-b").unwrap(), &b); +} + +#[test] +fn summary_batch_embedding_lookup_returns_only_signature_scoped_rows() { + let (_tmp, cfg) = test_config(); + insert_tree(&cfg, &sample_tree("tree-1", "slack:#eng")).unwrap(); + for sid in ["sum-1", "sum-2", "sum-3"] { + let node = sample_summary(sid, "tree-1", 1); + with_connection(&cfg, |conn| { + let tx = conn.unchecked_transaction()?; + insert_summary_tx(&tx, &node, "test")?; + tx.commit()?; + Ok(()) + }) + .unwrap(); + } + let sig_a = "openai/text-embedding-3-small@1536"; + let sig_b = "local/bge-small@384"; + set_summary_embedding_for_signature(&cfg, "sum-1", sig_a, &[0.1, 0.2]).unwrap(); + set_summary_embedding_for_signature(&cfg, "sum-2", sig_a, &[0.3, 0.4]).unwrap(); + set_summary_embedding_for_signature(&cfg, "sum-3", sig_b, &[0.5, 0.6, 0.7]).unwrap(); + + let ids = vec!["sum-1".into(), "sum-2".into(), "sum-3".into()]; + let map_a = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_a).unwrap(); + assert_eq!(map_a.len(), 2); + assert_eq!(map_a.get("sum-1").cloned(), Some(vec![0.1, 0.2])); + assert!(!map_a.contains_key("sum-3")); + + let map_b = get_summary_embeddings_for_signature_batch(&cfg, &ids, sig_b).unwrap(); + assert_eq!(map_b.len(), 1); + assert_eq!(map_b.get("sum-3").cloned(), Some(vec![0.5, 0.6, 0.7])); + + assert!(get_summary_embeddings_for_signature_batch(&cfg, &[], sig_a) + .unwrap() + .is_empty()); + assert!(get_summary_embeddings_batch(&cfg, &ids).unwrap().is_empty()); +} + +#[test] +fn list_trees_by_kind_and_archive() { + let (_tmp, cfg) = test_config(); + // Distinct created_at so the `ORDER BY created_at_ms ASC` is deterministic. + let mut s1 = sample_tree("source-1", "chat:slack:#eng"); + s1.created_at = Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(); + insert_tree(&cfg, &s1).unwrap(); + let mut topic = sample_tree("topic-1", "person:alice"); + topic.kind = TreeKind::Topic; + insert_tree(&cfg, &topic).unwrap(); + let mut s2 = sample_tree("source-2", "chat:discord:#ops"); + s2.created_at = Utc.timestamp_millis_opt(1_700_000_001_000).unwrap(); + insert_tree(&cfg, &s2).unwrap(); + + let source_ids: Vec = list_trees_by_kind(&cfg, TreeKind::Source) + .unwrap() + .into_iter() + .map(|t| t.id) + .collect(); + assert_eq!(source_ids, vec!["source-1", "source-2"]); + assert_eq!( + list_trees_by_kind(&cfg, TreeKind::Topic).unwrap()[0].id, + "topic-1" + ); + + archive_tree(&cfg, "source-1").unwrap(); + let archived = get_tree(&cfg, "source-1").unwrap().unwrap(); + assert_eq!(archived.status, TreeStatus::Archived); +} diff --git a/src/memory/tree/store/summaries.rs b/src/memory/tree/store/summaries.rs index 85d4869..57c1acd 100644 --- a/src/memory/tree/store/summaries.rs +++ b/src/memory/tree/store/summaries.rs @@ -178,6 +178,31 @@ pub fn list_summaries_in_window( }) } +/// List non-deleted summaries whose time envelope overlaps the inclusive +/// `[since_ms, until_ms]` window, across every summary level in one tree. +/// +/// Unlike [`list_summaries_in_window`], which requires full containment for +/// cover construction, this overlap query is intended for retrieval. +pub fn list_summaries_overlapping_window( + config: &MemoryConfig, + tree_id: &str, + since_ms: i64, + until_ms: i64, +) -> Result> { + with_connection(config, |conn| { + let mut stmt = conn.prepare(&format!( + "{SELECT_SUMMARY_BASE} WHERE tree_id = ?1 AND deleted = 0 AND level >= 1 \ + AND time_range_end_ms >= ?2 AND time_range_start_ms <= ?3 \ + ORDER BY level ASC, time_range_start_ms ASC" + ))?; + let rows = stmt + .query_map(params![tree_id, since_ms, until_ms], row_to_summary)? + .collect::>>() + .context("Failed to collect overlapping summaries")?; + Ok(rows) + }) +} + /// List non-deleted summaries whose `parent_id` is `parent` (its direct /// children). Ordered by `sealed_at` ASC. Used by the tree-walk read path. pub fn list_children_of_summary( diff --git a/src/memory/tree/store/types.rs b/src/memory/tree/store/types.rs index 0368beb..1b35d6b 100644 --- a/src/memory/tree/store/types.rs +++ b/src/memory/tree/store/types.rs @@ -64,22 +64,14 @@ impl TreeKind { } } -/// Activity state of a tree. Archived trees are *intended* to stay queryable -/// but reject new leaves. -/// -/// # NOTE: the archived invariant is not enforced today -/// Nothing in [`crate::memory::tree::bucket_seal::append_leaf`], -/// [`crate::memory::tree::registry::get_or_create_tree`], or -/// [`crate::memory::tree::bucket_seal::cascade_all_from`] checks this field — -/// an archived tree currently accepts new leaves and seals exactly like an -/// active one. See `TR-9` in `docs/spec/audit/03-tree-archivist-conversations.md`. +/// Activity state of a tree. Archived trees stay queryable but transactional +/// append and seal commits reject further writes. #[derive(Clone, Copy, Debug, Eq, PartialEq, Hash, Serialize, Deserialize)] #[serde(rename_all = "snake_case")] pub enum TreeStatus { /// Tree accepts new leaves and seals normally. Active, - /// Tree is meant to be frozen (no new leaves or seals) but the write paths - /// do not currently check this status — see the type-level NOTE above. + /// Tree is frozen: new buffer items and seal commits are rejected. Archived, } diff --git a/src/memory/tree/summarise.rs b/src/memory/tree/summarise.rs index 79a53fe..aca071d 100644 --- a/src/memory/tree/summarise.rs +++ b/src/memory/tree/summarise.rs @@ -14,10 +14,6 @@ use chrono::{DateTime, Utc}; use crate::memory::chunks::approx_token_count; use crate::memory::tree::store::TreeKind; -const MAX_SUMMARY_OUTPUT_TOKENS: u32 = 5_000; -const NUM_CTX_TOKENS: u32 = 60_000; -const OVERHEAD_RESERVE_TOKENS: u32 = 2_048; - /// One contribution being folded — a raw leaf at L0→L1, or a lower-level /// summary at L_n→L_{n+1}. #[derive(Clone, Debug)] @@ -52,8 +48,12 @@ pub struct SummaryContext<'a> { pub target_level: u32, /// Maximum approximate tokens the produced summary may occupy. pub token_budget: u32, + /// Total input/context budget available to this fold. + pub input_token_budget: u32, + /// Prompt and formatting headroom withheld from source inputs. + pub overhead_reserve_tokens: u32, /// Natural-language ask that steers the fold, for - /// [`TreeKind::Flavoured`](crate::memory::tree::TreeKind::Flavoured) trees. + /// [`TreeKind::Flavoured`] trees. /// When present, [`prepare_summary_prompt`] emits a flavour-directed system /// prompt instead of the generic folding prompt. `None` for every other /// tree kind. @@ -96,13 +96,13 @@ pub fn prepare_summary_prompt( ctx: &SummaryContext<'_>, output_language: Option<&str>, ) -> Option { - let effective_budget = ctx.token_budget.min(MAX_SUMMARY_OUTPUT_TOKENS); + let effective_budget = ctx.token_budget; let per_input_cap = if inputs.is_empty() { 0 } else { - NUM_CTX_TOKENS + ctx.input_token_budget .saturating_sub(effective_budget) - .saturating_sub(OVERHEAD_RESERVE_TOKENS) + .saturating_sub(ctx.overhead_reserve_tokens) / inputs.len() as u32 }; let mut ordered: Vec<_> = inputs.iter().collect(); diff --git a/src/memory/tree/summarise_tests.rs b/src/memory/tree/summarise_tests.rs index 3f36062..7317f1b 100644 --- a/src/memory/tree/summarise_tests.rs +++ b/src/memory/tree/summarise_tests.rs @@ -50,6 +50,8 @@ async fn concat_summariser_matches_fallback() { tree_kind: TreeKind::Source, target_level: 1, token_budget: 10_000, + input_token_budget: 50_000, + overhead_reserve_tokens: 2_048, ask: None, }; let out = ConcatSummariser::new() @@ -72,11 +74,13 @@ fn provider_prompt_is_priority_ordered_language_aware_and_budgeted() { tree_kind: TreeKind::Source, target_level: 1, token_budget: 9_000, + input_token_budget: 50_000, + overhead_reserve_tokens: 2_048, ask: None, }; let prompt = prepare_summary_prompt(&[low, high], &context, Some("French")).unwrap(); assert!(prompt.user.starts_with("[high]")); assert!(prompt.system.contains("Write the summary in French")); - assert_eq!(prompt.effective_budget, 5_000); + assert_eq!(prompt.effective_budget, 9_000); assert!(prepare_summary_prompt(&[], &context, None).is_none()); } diff --git a/src/memory/tree/types.rs b/src/memory/tree/types.rs new file mode 100644 index 0000000..54388b6 --- /dev/null +++ b/src/memory/tree/types.rs @@ -0,0 +1,82 @@ +//! Public service and input types for bucket-seal trees. + +use std::sync::Arc; + +use anyhow::Result; +use chrono::{DateTime, Utc}; + +use super::store::{SummaryNode, Tree}; +use super::summarise::Summariser; +use crate::memory::score::embed::Embedder; +use crate::memory::score::extract::EntityExtractor; + +/// Product callbacks around a seal after durable state transitions. +pub trait SealObserver: Send + Sync { + /// Report structural progress without summary content. + fn progress(&self, _tree: &Tree, _step: &str, _level: u32, _item_count: Option) {} + /// Notify a host mirror after a summary commits. + fn summary_committed( + &self, + _tree: &Tree, + _node: &SummaryNode, + _content_path: &str, + _reason: &str, + ) -> Result<()> { + Ok(()) + } +} + +pub(crate) struct NoopSealObserver; +impl SealObserver for NoopSealObserver {} + +/// Injected compute and product services for sealing. +pub struct SealServices<'a> { + /// Service that produces a summary from hydrated child inputs. + pub summariser: &'a dyn Summariser, + /// Optional embedding service used to index committed summaries. + pub embedder: Option<&'a dyn Embedder>, + /// Observer notified around durable seal transitions. + pub observer: &'a dyn SealObserver, +} + +/// Strategy for populating a summary's entity and topic labels. +#[derive(Clone)] +pub enum LabelStrategy { + /// Extract labels from the new summary content. + ExtractFromContent(Arc), + /// Union labels already carried by children. + UnionFromChildren, + /// Leave labels empty. + Empty, +} + +impl std::fmt::Debug for LabelStrategy { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::ExtractFromContent(extractor) => { + write!(f, "ExtractFromContent({})", extractor.name()) + } + Self::UnionFromChildren => f.write_str("UnionFromChildren"), + Self::Empty => f.write_str("Empty"), + } + } +} + +/// One persisted leaf being appended to an L0 buffer. +#[derive(Clone, Debug)] +pub struct LeafRef { + /// Stable identifier of the persisted chunk. + pub chunk_id: String, + /// Token contribution used by the L0 seal gate. + pub token_count: u32, + /// Source timestamp used for buffer age and summary bounds. + pub timestamp: DateTime, + /// Chunk content available to callers constructing leaf inputs. + pub content: String, + /// Canonical entity identifiers attached to the chunk. + pub entities: Vec, + /// Topic labels attached to the chunk. + pub topics: Vec, + /// Importance score carried by the leaf. + pub score: f32, +} From 95454a69959482fff161d8697884974f167f3045 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 07:17:37 +0000 Subject: [PATCH 4/7] test: enforce memory remediation quality gates --- .github/workflows/ci.yml | 13 +-- .github/workflows/release.yml | 7 +- Cargo.toml | 14 +-- README.md | 28 ++++-- docs/spec/README.md | 55 +++++++++++- docs/spec/improvement-plan.md | 9 ++ scripts/setup.sh | 14 +++ scripts/test.sh | 15 ++++ src/memory/providers/mod.rs | 24 ----- src/memory/rpc/mod.rs | 15 ---- tests/composio_sync_mock.rs | 13 +-- tests/smoke.rs | 164 ++++++++++++++++++++++++++++++++++ 12 files changed, 294 insertions(+), 77 deletions(-) create mode 100755 scripts/setup.sh create mode 100755 scripts/test.sh delete mode 100644 src/memory/providers/mod.rs delete mode 100644 src/memory/rpc/mod.rs diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2e8bd3d..f59b866 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -45,6 +45,9 @@ jobs: - name: Test run: cargo test --all-features + - name: Documentation + run: RUSTDOCFLAGS="-D warnings" cargo doc --all-features --no-deps + features: name: Features (${{ matrix.features.name }}) runs-on: ubuntu-latest @@ -61,10 +64,8 @@ jobs: flags: --no-default-features --features tokio - name: git-diff flags: --no-default-features --features git-diff - - name: providers-http - flags: --no-default-features --features providers-http - - name: rpc - flags: --no-default-features --features rpc + - name: sync + flags: --no-default-features --features sync - name: all features flags: --all-features steps: @@ -80,5 +81,5 @@ jobs: workspaces: . key: features-${{ matrix.features.name }} - - name: Check - run: cargo check --all-targets ${{ matrix.features.flags }} + - name: Test + run: cargo test --all-targets ${{ matrix.features.flags }} diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index febccc8..237eeb9 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -45,10 +45,13 @@ jobs: run: cargo fmt --all -- --check - name: Run clippy - run: cargo clippy --all-targets -- -D warnings + run: cargo clippy --all-targets --all-features -- -D warnings - name: Run tests - run: cargo test + run: cargo test --all-features + + - name: Build documentation + run: RUSTDOCFLAGS="-D warnings" cargo doc --all-features --no-deps - name: Compute next version id: version diff --git a/Cargo.toml b/Cargo.toml index f33f9a2..718633b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,6 @@ exclude = [ ".vscode/", "benchmarks/", "docs/", - "examples/", "gitbooks/", "paper/", "requirements.txt", @@ -34,21 +33,10 @@ tokio = ["dep:tokio"] # module (and libgit2) compiles out. git-diff = ["dep:git2"] -# reqwest-based embedding / LLM HTTP providers (`memory::providers`). Pulls in -# the reqwest stack and needs an async runtime, so it implies `tokio`. The -# concrete providers land with goals C3/M3; this feature reserves the seam and -# gates the dependency. -providers-http = ["dep:reqwest", "tokio"] - # Live source synchronization (Composio HTTP pipelines and workspace scans). # Host schedulers, credentials, RPC, and event buses remain outside the crate. sync = ["dep:reqwest", "dep:tracing", "tokio"] -# serde schema / envelope surface for the RPC boundary (`memory::rpc`). Reserved -# for goal C5; gates the wire-facing surface without adding heavy dependencies -# (serde is already a core dependency). -rpc = [] - [dependencies] anyhow = "1" log = "0.4" @@ -79,7 +67,7 @@ walkdir = "2" git2 = { version = "0.21", default-features = false, features = [ "vendored-libgit2", ], optional = true } -# reqwest pulls a large async HTTP stack — only for the `providers-http` feature. +# reqwest pulls a large async HTTP stack — only for live synchronization. reqwest = { version = "0.12", default-features = false, features = [ "json", "rustls-tls", diff --git a/README.md b/README.md index 665414b..62729fd 100644 --- a/README.md +++ b/README.md @@ -16,7 +16,7 @@ The human brain is a master at compression. It doesn't try to remember every passing detail; instead it aggressively prunes noise to keep a sharp, focused, easily accessible recall of what truly matters. Traditional AI memory systems do the opposite — they try to remember _everything_ and retrieve whatever is _similar_. But similar doesn't mean important. The result? Your AI drowns in stale, irrelevant context that degrades every response. -**TinyCortex** takes the brain's approach: it **intelligently forgets noise**. A multi-signal admission gate drops low-value content at ingest, and retrieval applies exponential time-decay so stale memories fade from recall while fresh, high-signal knowledge ranks first. The result is a memory engine that stays lean and focused instead of drowning in stale context. +**TinyCortex** takes the brain's approach: it **intelligently forgets noise**. A multi-signal admission gate drops low-value content at ingest. Retrieval provides semantic, keyword, graph, summary-tree, and freshness-scoring building blocks so hosts can choose an explicit ranking policy instead of treating every stored item as equally useful. The hosted TinyCortex platform has been evaluated on [RAGAS](https://www.ragas.io/), TemporalBench, [BABILong](https://github.com/booydar/babilong/), and [Vending-Bench](https://andonlabs.com/evals/vending-bench-2) — see [Benchmarks](#-benchmarks) below for the reported results and what you can reproduce from this repo. @@ -26,7 +26,7 @@ The hosted TinyCortex platform has been evaluated on [RAGAS](https://www.ragas.i ## Intelligent Noise Filtering -Every chunk passes a multi-signal admission gate at ingest — low-value content is dropped before it ever pollutes the store — and retrieval ranking applies exponential time-decay (7-day half-life by default) so stale memories fade from recall. The store stays lean on its own — no manual cleanup. +Every chunk passes a multi-signal admission gate at ingest, so low-value content is dropped before it pollutes the store. The retrieval module exposes exponential freshness decay (7-day half-life by default) as a composable signal; built-in tree queries use semantic similarity or recency unless the host applies the hybrid composer. ![Interaction graph highlighting important knowledge](docs/images/gif/AppleEmailGraph.gif) @@ -42,7 +42,7 @@ Markdown files are the source of truth. SQLite chunk rows, summary trees, vector ## Recency-Weighted Recall -Retrieval blends keyword relevance, vector similarity, graph proximity, and freshness into a single explainable score, so what comes back is a focused slice of long-term history rather than a noisy dump. Summary-tree hotness tracking promotes frequently written topics into their own trees. (The fully proactive "conscious recall" experience — surfacing memories without an explicit query — is part of the hosted TinyCortex platform, built on these primitives.) +TinyCortex exposes keyword relevance, vector similarity, graph proximity, freshness, and an explainable hybrid-score composer. These are policy primitives rather than a hidden global ranking rule: hosts can apply the supplied `WeightProfile`, while the built-in source/topic/global queries use their documented semantic or recency ordering. Summary-tree hotness tracking promotes frequently written topics into their own trees. (The fully proactive "conscious recall" experience — surfacing memories without an explicit query — is part of the hosted TinyCortex platform, built on these primitives.) # ⚡ Getting Started @@ -68,7 +68,7 @@ async fn main() -> anyhow::Result<()> { .await?; // Recall it with a keyword query. - let hits = store.search(MemoryQuery::text("theme preference")).await?; + let hits = store.search(MemoryQuery::text("dark mode")).await?; for hit in hits { println!("{:.3} {}", hit.score, hit.record.content); } @@ -76,7 +76,15 @@ async fn main() -> anyhow::Result<()> { } ``` -The `InMemoryMemoryStore` is the simple reference backend. The full engine — content store, chunking, scoring, summary trees, vector/keyword/graph/hybrid retrieval, the diff ledger, and the async job queue — lives under the [`memory`](https://docs.rs/tinycortex/latest/tinycortex/memory/) module. See the **[documentation](https://tinyhumans.gitbook.io/tinycortex/)** for the architecture, concepts, and end-to-end ingest walkthroughs. +The `InMemoryMemoryStore` is the simple reference backend. The full engine — content store, chunking, scoring, summary trees, and vector/keyword/graph retrieval primitives — lives under the [`memory`](https://docs.rs/tinycortex/latest/tinycortex/memory/) module. Optional surfaces require Cargo features: + +| Feature | Enables | +| --- | --- | +| `tokio` | Always-on queue worker and scheduler loops | +| `git-diff` | Git-backed snapshots, checkpoints, and read markers | +| `sync` | Composio and workspace synchronization pipelines | + +See the **[documentation](https://tinyhumans.gitbook.io/tinycortex/)** for the architecture, concepts, and end-to-end ingest walkthroughs. # 🧩 How It Works @@ -96,12 +104,12 @@ source payload | ----------------------------- | -------------------------------------------------------------------------------------------- | | **Storage primitives** | Markdown content store, SQLite chunks, summary trees, vector DB, KV, entity index | | **Ingest** | Canonicalize → chunk → score → embed → tree | -| **Retrieval** | Vector, keyword, graph, tree drill-down, and hybrid search with explainable score breakdowns | -| **Diff** | Git-backed source snapshots, checkpoints, and read-markers for change awareness | +| **Retrieval** | Vector, keyword, graph, tree drill-down, and composable hybrid-score primitives | +| **Diff** (`git-diff`) | Git-backed source snapshots, checkpoints, and read-markers for change awareness | | **Entities & Graph** | Entity markdown files + a co-occurrence graph derived from the entity index | | **Goals / Tool Memory** | Compact long-term goal list and durable tool-scoped rules | | **Conversations / Archivist** | Transcript storage and conversion of turns into summary-tree leaves | -| **Queue** | Async jobs: extract, append, seal, flush-stale, re-embed, seal-document | +| **Queue** | Durable jobs; optional `tokio` feature adds always-on worker/scheduler loops | Full details live in the **[documentation](https://tinyhumans.gitbook.io/tinycortex/)**. @@ -138,7 +146,9 @@ See [`benchmarks/`](./benchmarks/README.md) for the full reported tables and for # 🤝 Contributing -TinyCortex is built in Rust (2021 edition). Clone the repo and: +TinyCortex is built in Rust (2021 edition). Clone the repo, then run +`./scripts/setup.sh` once and `./scripts/test.sh` for the full local CI suite. +The underlying commands are: ```bash cargo test # run unit + integration tests diff --git a/docs/spec/README.md b/docs/spec/README.md index 7698445..ea27dc0 100644 --- a/docs/spec/README.md +++ b/docs/spec/README.md @@ -1,6 +1,6 @@ # Memory Engine Audit & Improvement Spec -_Audit date: 2026-07-11 · Baseline: `main` @ `c9d1afd` · All 1025 unit tests + 1 integration test green at time of audit._ +_Initial audit: 2026-07-11 · Architecture follow-up: 2026-07-14 · Remediation implementation: 2026-07-14._ This folder holds the results of a full-codebase audit of the TinyCortex memory engine and the improvement plan derived from it. The audit swept every @@ -25,7 +25,58 @@ carries `file:line` evidence. | [audit/09-verification-infrastructure.md](audit/09-verification-infrastructure.md) | Per-block verifiability: test coverage map, CI feature matrix, setup/test scripts | | [audit/10-simplification-dead-weight.md](audit/10-simplification-dead-weight.md) | Duplication, dead/speculative code, dependency weight, error/async story | -## Executive summary +## Remediation status + +The correctness remediation described by audits 01–06 is implemented. The +work preserves wire ids while closing the verified data-loss, concurrency, +parser, queue-settlement, retrieval, and contract defects. Fixes that were +already present were retained and covered rather than reimplemented. + +The resulting trust baseline includes: + +- atomic or compensated persistence at every audited crash boundary, including + document ingest, queue follow-ups, staged summaries, time-tree rebuilds, + goals, and source/tool-memory registries; +- concurrency-safe buffer sealing, conversation sequencing, source mutation, + read markers, and shared SQLite/entity-index ownership; +- SQL-bounded retrieval and graph operations without the audited 200/500/5000 + silent truncation cliffs; +- fail-closed taint/category contracts, a real `Memory` implementation for the + reference store, validated partial config loading, and configurable scoring, + ingestion, retrieval-limit, queue, sync, and tree policy roots; +- a real ingest → durable queue drain → tree seal → retrieval functional test, + feature-matrix tests (including `sync`), setup/test scripts, corruption and + crash-window tests, and source files split below the 500-line repository cap; +- removal of the empty `rpc` and `providers-http` feature surfaces. Optional + features now correspond to concrete code rather than reserved placeholders. + +Audits 07–10 intentionally also contain longer-horizon design alternatives, +such as replacing the shared transactional SQLite database with separately +owned persistence services or adding a remote store backend. Those are not +correctness defects and are not silently presented as completed here. The +shared database remains the documented atomicity boundary; the public +in-memory backend now implements the consolidated `Memory` contract, while +the separate [configurable-store specification](configurable-store.md) remains +the design for a future network backend. + +Verification is recorded from clean commands run at the end of remediation; +the exact test and coverage totals below are updated only from command output. + +| Gate | Final result | +| --- | --- | +| Formatting | `cargo fmt --all -- --check` passes | +| Static validation | all-feature check and clippy with warnings denied pass | +| Unit tests | 1300 passed | +| Functional tests | 15 passed; 1 credentialed live smoke intentionally ignored | +| Doctests | 4 passed; 1 example intentionally ignored | +| Documentation | strict all-feature rustdoc passes with no warnings | +| Line coverage | **90.60%** (`cargo llvm-cov --all-features --workspace`) | +| Packaging | `cargo package --no-verify` succeeds without warnings | + +## Original audit executive summary + +The following describes the pre-remediation baseline and is retained as the +rationale for the fixes, not as a statement about the current tree. The engine is well-layered and richly documented, and the full test suite passes — but the audit surfaced **4 critical** and **~20 major** verified diff --git a/docs/spec/improvement-plan.md b/docs/spec/improvement-plan.md index 98fb6d1..1f108ce 100644 --- a/docs/spec/improvement-plan.md +++ b/docs/spec/improvement-plan.md @@ -1,5 +1,14 @@ # Memory Engine Improvement Plan +> Remediation record (2026-07-14): phases 0, 1, 2, and 4 have been +> implemented with regression coverage. Phase 3's contract consolidation is +> complete for the local reference backend (`InMemoryMemoryStore` implements +> `Memory` and the conflicting store error was renamed); the network backend +> remains a separate future product capability described in +> [configurable-store.md](configurable-store.md), not an implemented feature. +> The architecture alternatives in audits 07–10 remain decision records where +> they propose replacing the shared transactional SQLite boundary. + Derived from the [audit findings](README.md). Phases are ordered by risk: each phase is independently shippable, and within a phase the workstreams are parallelizable (use separate worktrees per the repo guidelines). Finding IDs diff --git a/scripts/setup.sh b/scripts/setup.sh new file mode 100755 index 0000000..2b472b1 --- /dev/null +++ b/scripts/setup.sh @@ -0,0 +1,14 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$repo_root" + +git submodule update --init --recursive + +if command -v rustup >/dev/null 2>&1; then + rustup toolchain install stable --profile minimal --component rustfmt,clippy +fi + +cargo fetch --locked +echo "TinyCortex development dependencies are ready." diff --git a/scripts/test.sh b/scripts/test.sh new file mode 100755 index 0000000..6ae9929 --- /dev/null +++ b/scripts/test.sh @@ -0,0 +1,15 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" +cd "$repo_root" + +cargo fmt --all -- --check +cargo clippy --all-targets --all-features -- -D warnings +cargo test --all-targets --all-features +RUSTDOCFLAGS="-D warnings" cargo doc --all-features --no-deps + +for feature in tokio git-diff sync; do + cargo test --all-targets --no-default-features --features "$feature" +done +cargo test --all-targets --no-default-features diff --git a/src/memory/providers/mod.rs b/src/memory/providers/mod.rs deleted file mode 100644 index 6bf4f6a..0000000 --- a/src/memory/providers/mod.rs +++ /dev/null @@ -1,24 +0,0 @@ -//! HTTP-backed embedding / LLM providers (feature `providers-http`). -//! -//! This module is the seam for reqwest-based provider implementations that back -//! the scoring and embedding signals with real network services. The concrete -//! providers (and their wire contracts) land with goals **C3** (embedding -//! providers) and **M3** (LLM providers); this module reserves the boundary and -//! gates the heavy [`reqwest`] dependency behind the `providers-http` feature so -//! the dependency-light core never links it. -//! -//! ## Embedding signature invariant -//! -//! Any embedding provider added here MUST report its identity as the canonical -//! signature string `provider=;model=;dims=` so that -//! vectors stay comparable across the store, retrieval, and re-embed backfill -//! paths. See `docs/plan/01-migration-library.md` §1.3. - -/// Shared async HTTP client type for provider implementations. -/// -/// Aliased here so the reqwest dependency stays confined to this feature-gated -/// module and provider code has a single, stable client type to build against. -/// As of this writing no provider implementation exists yet under this -/// module; the alias exists so future providers (goals **C3**/**M3**) share -/// one client type rather than each pulling `reqwest` in independently. -pub type HttpClient = reqwest::Client; diff --git a/src/memory/rpc/mod.rs b/src/memory/rpc/mod.rs deleted file mode 100644 index f4c3685..0000000 --- a/src/memory/rpc/mod.rs +++ /dev/null @@ -1,15 +0,0 @@ -//! serde schema / envelope surface for the RPC boundary (feature `rpc`). -//! -//! This module is the seam for the wire-facing request/response envelope and -//! schema types that expose the memory engine over an RPC boundary. The -//! concrete envelope contract lands with goal **C5**; this module reserves the -//! surface and gates it behind the `rpc` feature so hosts that only embed the -//! engine in-process never compile the wire layer. -//! -//! No heavy dependencies are pulled in: the RPC surface is built on the core -//! `serde` / `serde_json` stack that the engine already depends on. -//! -//! Wire-format invariants that this surface MUST preserve when populated (see -//! `docs/plan/01-migration-library.md` §1.3): embedding signatures, archivist -//! leaf ids, and the `MemoryTaint` `internal` / `external_sync` strings that -//! fail **closed** to `external_sync`. diff --git a/tests/composio_sync_mock.rs b/tests/composio_sync_mock.rs index ec3362b..8c93d6f 100644 --- a/tests/composio_sync_mock.rs +++ b/tests/composio_sync_mock.rs @@ -155,12 +155,13 @@ async fn gmail_sync_paginates_persists_cursor_taint_and_is_idempotent() { let first = pipeline.tick(&config, &context).await.unwrap(); assert_eq!(first.records_ingested, 2); assert_eq!(first.actions_called, 2); - let docs = captures.documents.lock().unwrap(); - assert_eq!(docs.len(), 2); - assert!(docs - .iter() - .all(|doc| doc.metadata["taint"] == "external_sync")); - drop(docs); + { + let docs = captures.documents.lock().unwrap(); + assert_eq!(docs.len(), 2); + assert!(docs + .iter() + .all(|doc| doc.metadata["taint"] == "external_sync")); + } let state = SyncState::load(captures.as_ref(), "gmail", "conn-1") .await diff --git a/tests/smoke.rs b/tests/smoke.rs index b2542e2..9d5e9d8 100644 --- a/tests/smoke.rs +++ b/tests/smoke.rs @@ -1,5 +1,123 @@ +use async_trait::async_trait; +use chrono::{TimeZone, Utc}; +use tempfile::TempDir; +use tinycortex::memory::chunks::get_chunk; +use tinycortex::memory::config::MemoryConfig; +use tinycortex::memory::ingest::canonicalize::chat::{ChatBatch, ChatMessage}; +use tinycortex::memory::ingest::{ingest_chat, QueueJobSink}; +use tinycortex::memory::queue::{ + drain_until_idle, AppendDecision, AppendTarget, ExtractDecision, NodeRef, QueueDelegates, + ReembedProgress, SealDocumentPayload, SealPayload, StaleBuffer, +}; +use tinycortex::memory::retrieval::query_source; +use tinycortex::memory::score::embed::InertEmbedder; +use tinycortex::memory::score::ScoringConfig; +use tinycortex::memory::tree::store::get_tree; +use tinycortex::memory::tree::{append_leaf_deferred, ConcatSummariser, LeafRef, TreeFactory}; use tinycortex::memory::{InMemoryMemoryStore, MemoryInput, MemoryQuery, MemoryStore}; +struct EngineDelegates; + +#[async_trait] +impl QueueDelegates for EngineDelegates { + async fn extract_chunk( + &self, + config: &MemoryConfig, + chunk_id: &str, + ) -> anyhow::Result> { + Ok(get_chunk(config, chunk_id)?.map(|chunk| ExtractDecision { + kept: true, + uses_document_subtree: false, + tree_scope: chunk.metadata.source_id, + })) + } + + async fn append_node( + &self, + config: &MemoryConfig, + node: &NodeRef, + target: &AppendTarget, + ) -> anyhow::Result> { + let NodeRef::Leaf { chunk_id } = node else { + return Ok(None); + }; + let Some(chunk) = get_chunk(config, chunk_id)? else { + return Ok(None); + }; + let source_id = match target { + AppendTarget::Source { source_id } => source_id, + AppendTarget::Topic { .. } => return Ok(None), + }; + let factory = TreeFactory::source(source_id.as_str()); + let tree = factory.get_or_create(config)?; + append_leaf_deferred( + config, + &tree, + &LeafRef { + chunk_id: chunk.id, + token_count: chunk.token_count, + timestamp: chunk.metadata.timestamp, + content: chunk.content, + entities: vec![], + topics: vec![], + score: 1.0, + }, + )?; + Ok(Some(AppendDecision { + tree_id: tree.id, + should_seal: true, + })) + } + + async fn seal_level( + &self, + config: &MemoryConfig, + payload: &SealPayload, + ) -> anyhow::Result> { + let tree = get_tree(config, &payload.tree_id)?.expect("queued tree exists"); + TreeFactory::from_tree(&tree) + .seal_now(config, &ConcatSummariser::new()) + .await?; + Ok(None) + } + + async fn list_stale_buffers( + &self, + _config: &MemoryConfig, + _max_age_secs: i64, + ) -> anyhow::Result> { + Ok(vec![]) + } + + async fn seal_document( + &self, + _config: &MemoryConfig, + _payload: &SealDocumentPayload, + ) -> anyhow::Result<()> { + Ok(()) + } + + async fn reembed_batch( + &self, + _config: &MemoryConfig, + _signature: &str, + ) -> anyhow::Result { + Ok(ReembedProgress::Covered) + } + + fn active_signature(&self, _config: &MemoryConfig) -> String { + "inert".into() + } + + fn has_uncovered_reembed_work( + &self, + _config: &MemoryConfig, + _signature: &str, + ) -> anyhow::Result { + Ok(false) + } +} + #[tokio::test] async fn stores_and_finds_memory() { let store = InMemoryMemoryStore::new(); @@ -19,3 +137,49 @@ async fn stores_and_finds_memory() { assert_eq!(hits.len(), 1); assert_eq!(hits[0].record.namespace, "default"); } + +#[tokio::test] +async fn ingests_drains_seals_and_retrieves_real_engine_data() { + let workspace = TempDir::new().unwrap(); + let config = MemoryConfig::new(workspace.path()); + let batch = ChatBatch { + platform: "slack".into(), + channel_label: "#launch".into(), + messages: vec![ChatMessage { + author: "alice".into(), + timestamp: Utc.timestamp_millis_opt(1_700_000_000_000).unwrap(), + text: "The Phoenix launch is Friday after the staging review and runbook sign-off." + .into(), + source_ref: Some("slack://launch/1".into()), + }], + }; + + let ingested = ingest_chat( + &config, + "slack:#launch", + "alice", + vec![], + batch, + &QueueJobSink, + &ScoringConfig::from_memory_config(&config), + ) + .await + .unwrap(); + assert_eq!(ingested.chunks_written, 1); + assert_eq!(ingested.extract_jobs_enqueued, 1); + + drain_until_idle(&config, &EngineDelegates).await.unwrap(); + let response = query_source( + &config, + Some("slack:#launch"), + None, + None, + None, + &InertEmbedder, + 10, + ) + .await + .unwrap(); + assert_eq!(response.hits.len(), 1); + assert!(response.hits[0].content.contains("Phoenix launch")); +} From fc8d569c9aa2cd6d2fb67ed0c1c4012dfc6a3fc6 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 09:37:15 +0000 Subject: [PATCH 5/7] fix: address pull request review feedback --- src/memory/archivist/store.rs | 26 +-- src/memory/chunks/connection.rs | 35 +-- src/memory/chunks/connection_breaker.rs | 61 +++-- src/memory/chunks/migrations.rs | 62 ++++-- src/memory/chunks/produce.rs | 5 +- src/memory/chunks/store.rs | 30 ++- src/memory/chunks/store_delete.rs | 2 +- src/memory/chunks/store_list.rs | 10 +- src/memory/config.rs | 7 + src/memory/config/policy.rs | 8 +- src/memory/config/policy_tests.rs | 47 ++++ src/memory/conversations/bus.rs | 41 +++- src/memory/conversations/inverted_index.rs | 13 +- src/memory/conversations/store_index.rs | 35 ++- src/memory/conversations/store_tests_more.rs | 16 +- src/memory/diff/ledger_helpers.rs | 11 + src/memory/fsutil.rs | 21 ++ src/memory/ingest/pipeline.rs | 6 +- src/memory/ingest/pipeline_tests.rs | 2 + src/memory/ingest/types.rs | 26 ++- src/memory/queue/ops_tests.rs | 1 + src/memory/queue/runtime/mod.rs | 6 +- src/memory/queue/store.rs | 24 +- src/memory/queue/store_settle.rs | 17 +- src/memory/queue/store_settle_tests.rs | 3 +- src/memory/queue/test_support.rs | 3 + src/memory/queue/worker.rs | 19 +- src/memory/queue/worker_tests.rs | 1 + src/memory/retrieval/global.rs | 11 +- src/memory/retrieval/source.rs | 16 +- src/memory/score/mod.rs | 25 +-- src/memory/score/store_query.rs | 4 + src/memory/score/store_query_tests.rs | 45 ++++ src/memory/sources/readers/folder.rs | 2 +- src/memory/sources/registry.rs | 34 ++- src/memory/sources/types.rs | 6 + src/memory/store/content/atomic.rs | 2 +- src/memory/store/entity_index/transaction.rs | 5 + src/memory/store/kv.rs | 13 +- src/memory/store/memory_trait.rs | 14 +- src/memory/store/mod.rs | 19 ++ src/memory/store/safety/pii/checks.rs | 4 + src/memory/store/safety/pii/checks_tests.rs | 43 ++++ src/memory/store/store.rs | 25 +-- src/memory/store/vectors/store.rs | 29 +-- src/memory/store/vectors/store_tests_more.rs | 8 +- .../sync/composio/providers/slack_parse.rs | 10 + src/memory/sync/dispatcher.rs | 209 +----------------- src/memory/sync/dispatcher_tests.rs | 209 ++++++++++++++++++ src/memory/tool_memory/store.rs | 36 +-- src/memory/tree/bucket_seal.rs | 6 +- src/memory/tree/document_seal.rs | 22 +- src/memory/tree/runtime/store/paths.rs | 17 +- src/memory/types.rs | 1 + 54 files changed, 832 insertions(+), 521 deletions(-) create mode 100644 src/memory/config/policy_tests.rs create mode 100644 src/memory/score/store_query_tests.rs create mode 100644 src/memory/store/safety/pii/checks_tests.rs create mode 100644 src/memory/sync/dispatcher_tests.rs diff --git a/src/memory/archivist/store.rs b/src/memory/archivist/store.rs index 4c89eb1..288fcce 100644 --- a/src/memory/archivist/store.rs +++ b/src/memory/archivist/store.rs @@ -19,7 +19,6 @@ use std::fs; use std::path::{Path, PathBuf}; use anyhow::{Context, Result}; -use sha2::{Digest, Sha256}; use crate::memory::archivist::types::ArchivedTurn; use crate::memory::config::MemoryConfig; @@ -42,28 +41,9 @@ fn session_dir(config: &MemoryConfig, session_id: &str) -> PathBuf { /// Map any non-`[A-Za-z0-9_-]` character to `_` so a session id is always a /// safe single path component. fn sanitize_session(s: &str) -> String { - let sanitized: String = s - .chars() - .map(|c| { - if c.is_alphanumeric() || c == '-' || c == '_' { - c - } else { - '_' - } - }) - .collect(); - if sanitized == s && !sanitized.is_empty() { - return sanitized; - } - - // Replacement alone is collision-prone (`a/b` and `a?b`). Keep a short - // digest of the exact machine id whenever sanitisation changed it. - let digest = Sha256::digest(s.as_bytes()); - let suffix = digest[..6] - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::(); - format!("{sanitized}-{suffix}") + crate::memory::fsutil::sanitize_component_with_digest(s, |character| { + character.is_alphanumeric() || character == '-' || character == '_' + }) } /// Next free sequence number for a session: one past the highest `NNNNNN.md` diff --git a/src/memory/chunks/connection.rs b/src/memory/chunks/connection.rs index 8661553..e1786a7 100644 --- a/src/memory/chunks/connection.rs +++ b/src/memory/chunks/connection.rs @@ -271,6 +271,20 @@ pub(crate) fn get_or_init_connection(config: &MemoryConfig) -> Result Result { @@ -406,7 +406,12 @@ pub(super) fn recover_corrupt_connection(config: &MemoryConfig) -> Result }; let _init_guard = init_lock.lock(); - drop_cached_connection(config); + // Remove the cache entry, then quiesce every caller that already cloned + // its Arc before renaming the database. The guard remains held through + // quarantine and rebuild, so no write can land in the quarantined file. + let cached = conn_cache().connections.lock().remove(&db_path); + let _cached_guard = cached.as_ref().map(|connection| connection.lock()); + conn_cache().breakers.lock().remove(&db_path); let quarantined = quarantine_corrupt_files(config)?; let conn = open_and_init(&db_path, config) .context("failed to rebuild chunk DB schema after quarantining corrupt DB")?; diff --git a/src/memory/chunks/connection_breaker.rs b/src/memory/chunks/connection_breaker.rs index 6cbe483..27c2550 100644 --- a/src/memory/chunks/connection_breaker.rs +++ b/src/memory/chunks/connection_breaker.rs @@ -1,6 +1,5 @@ //! Per-database circuit breaker for repeated connection-init failures. -use std::sync::atomic::{AtomicBool, AtomicU32, Ordering}; use std::time::{Duration, Instant}; use parking_lot::Mutex; @@ -8,43 +7,65 @@ use parking_lot::Mutex; pub(crate) const CB_THRESHOLD: u32 = 3; pub(crate) const CB_COOLDOWN: Duration = Duration::from_secs(30); +struct BreakerState { + consecutive_failures: u32, + tripped: bool, + last_trip: Option, +} + pub(super) struct CircuitBreaker { - consecutive_failures: AtomicU32, - tripped: AtomicBool, - last_trip: Mutex>, + state: Mutex, } impl CircuitBreaker { + /// Create a closed breaker with no recorded initialization failures. pub fn new() -> Self { Self { - consecutive_failures: AtomicU32::new(0), - tripped: AtomicBool::new(false), - last_trip: Mutex::new(None), + state: Mutex::new(BreakerState { + consecutive_failures: 0, + tripped: false, + last_trip: None, + }), } } + /// Reset failure state after a successful initialization. + /// + /// Returns `true` when this call recovered a previously tripped breaker. pub fn record_success(&self) -> bool { - self.consecutive_failures.store(0, Ordering::Relaxed); - *self.last_trip.lock() = None; - self.tripped.swap(false, Ordering::Relaxed) + let mut state = self.state.lock(); + let was_tripped = state.tripped; + state.consecutive_failures = 0; + state.tripped = false; + state.last_trip = None; + was_tripped } + /// Record one initialization failure and trip at [`CB_THRESHOLD`]. + /// + /// Returns `true` only for the failure that transitions the breaker from + /// closed to tripped. Further failures refresh the cooldown timestamp. pub fn record_failure(&self) -> bool { - let count = self.consecutive_failures.fetch_add(1, Ordering::Relaxed) + 1; - if count >= CB_THRESHOLD && !self.tripped.swap(true, Ordering::Relaxed) { - *self.last_trip.lock() = Some(Instant::now()); - return true; - } - if self.tripped.load(Ordering::Relaxed) { - *self.last_trip.lock() = Some(Instant::now()); + let mut state = self.state.lock(); + state.consecutive_failures = state.consecutive_failures.saturating_add(1); + let just_tripped = state.consecutive_failures >= CB_THRESHOLD && !state.tripped; + if state.consecutive_failures >= CB_THRESHOLD { + state.tripped = true; + state.last_trip = Some(Instant::now()); } - false + just_tripped } + /// Return whether the breaker is tripped and still within its cooldown. + /// + /// Once [`CB_COOLDOWN`] elapses this returns `false`, allowing one + /// serialized initialization attempt; success resets the breaker and + /// failure refreshes the cooldown. pub fn is_open(&self) -> bool { - if !self.tripped.load(Ordering::Relaxed) { + let state = self.state.lock(); + if !state.tripped { return false; } - matches!(*self.last_trip.lock(), Some(tripped) if tripped.elapsed() < CB_COOLDOWN) + matches!(state.last_trip, Some(tripped) if tripped.elapsed() < CB_COOLDOWN) } } diff --git a/src/memory/chunks/migrations.rs b/src/memory/chunks/migrations.rs index 80c6db6..25fe6da 100644 --- a/src/memory/chunks/migrations.rs +++ b/src/memory/chunks/migrations.rs @@ -4,11 +4,9 @@ //! once per vault. Called from [`super::connection`] during DB initialisation. //! //! ## Contract -//! Both migrations here follow the same shape: read `user_version`, bail out -//! if it is already at/past the migration's target version, otherwise do the -//! work and bump `user_version` inside one `unchecked_transaction`. A crash can -//! therefore expose neither the migrated rows nor the version marker, or both, -//! but never a split state. +//! Both migrations are version-gated. Database-only work and its version marker +//! share a transaction. Migrations with filesystem cleanup complete that +//! cleanup before recording the version so an interruption remains retryable. //! //! Neither migration is exercised by an automated test in this module (audit //! test-coverage gap) — both are asserted only via the module's behavior at @@ -112,16 +110,13 @@ pub(super) fn migrate_legacy_embeddings_to_sidecar( /// parent `mem_tree_summaries` / `mem_tree_trees` rows, so this works whether /// or not `ON DELETE CASCADE` is declared for a given foreign key. /// -/// The on-disk `wiki/summaries/global*` / `topic-*` folder cleanup after the -/// DB transaction is best-effort: a filesystem error there is swallowed -/// (`let _ =`) rather than propagated, so a failed directory removal does not -/// block the `user_version` bump — those folders may survive as harmless -/// orphans if cleanup fails, but the DB-side purge still completes. +/// The on-disk `wiki/summaries/global*` / `topic-*` folder cleanup happens +/// before the version marker is committed. A cleanup failure therefore leaves +/// the migration retryable on the next startup. /// /// # Errors /// Returns `Err` if any `DELETE` statement or the transaction commit fails, or -/// if bumping `user_version` afterward fails. Does not return `Err` for -/// filesystem cleanup failures (see above). +/// if filesystem cleanup, bumping `user_version`, or commit fails. pub(super) fn purge_global_topic_trees(conn: &Connection, config: &MemoryConfig) -> Result<()> { let version: i64 = conn .query_row("PRAGMA user_version", [], |r| r.get(0)) @@ -163,22 +158,43 @@ pub(super) fn purge_global_topic_trees(conn: &Connection, config: &MemoryConfig) "DELETE FROM mem_tree_jobs WHERE kind IN ('topic_route','digest_daily')", [], )?; - tx.pragma_update(None, "user_version", GLOBAL_TOPIC_PURGE_MIGRATION_VERSION) - .context("set PRAGMA user_version during global/topic purge")?; - tx.commit()?; - // On-disk: drop `wiki/summaries/global*` and `topic-*` summary folders. - // Best-effort — a filesystem error must not abort the version bump. let summaries_root = content_root(config).join("wiki").join("summaries"); - if let Ok(entries) = std::fs::read_dir(&summaries_root) { - for entry in entries.flatten() { - let name = entry.file_name(); - let name = name.to_string_lossy(); - if name.starts_with("global") || name.starts_with("topic-") { - let _ = std::fs::remove_dir_all(entry.path()); + match std::fs::read_dir(&summaries_root) { + Ok(entries) => { + for entry in entries { + let entry = entry.with_context(|| { + format!( + "read summary directory entry in {}", + summaries_root.display() + ) + })?; + let name = entry.file_name(); + let name = name.to_string_lossy(); + if name.starts_with("global") || name.starts_with("topic-") { + std::fs::remove_dir_all(entry.path()).with_context(|| { + format!( + "remove retired summary directory {}", + entry.path().display() + ) + })?; + } } } + Err(error) if error.kind() == std::io::ErrorKind::NotFound => {} + Err(error) => { + return Err(error).with_context(|| { + format!( + "read retired summary directories in {}", + summaries_root.display() + ) + }); + } } + tx.pragma_update(None, "user_version", GLOBAL_TOPIC_PURGE_MIGRATION_VERSION) + .context("set PRAGMA user_version during global/topic purge")?; + tx.commit()?; + Ok(()) } diff --git a/src/memory/chunks/produce.rs b/src/memory/chunks/produce.rs index 30399c4..6aff6d1 100644 --- a/src/memory/chunks/produce.rs +++ b/src/memory/chunks/produce.rs @@ -121,7 +121,10 @@ pub fn chunk_markdown(input: &ChunkerInput, opts: &ChunkerOptions) -> Vec for (part, piece) in sub_pieces.into_iter().enumerate() { let seq = out.len() as u32; let tc = approx_token_count(&piece); - let id = chunk_id(input.source_kind, &input.source_id, part as u32, &unit); + // Include the emitted piece in the identity. If the split + // boundary changes, a body must not silently reuse the old + // embedding/score identity for the same part number. + let id = chunk_id(input.source_kind, &input.source_id, part as u32, &piece); out.push(Chunk { id, content: piece, diff --git a/src/memory/chunks/store.rs b/src/memory/chunks/store.rs index 3e1b2b4..681b8d3 100644 --- a/src/memory/chunks/store.rs +++ b/src/memory/chunks/store.rs @@ -50,12 +50,29 @@ pub fn upsert_chunks(config: &MemoryConfig, chunks: &[Chunk]) -> Result { if chunks.is_empty() { return Ok(0); } - with_connection(config, |conn| { + let (count, replaced_content_paths) = with_connection(config, |conn| { let tx = conn.unchecked_transaction()?; + let mut replaced_content_paths = Vec::new(); + for chunk in chunks { + if let Some(path) = tx + .query_row( + "SELECT content_path FROM mem_tree_chunks WHERE id = ?1", + [&chunk.id], + |row| row.get::<_, Option>(0), + ) + .optional()? + .flatten() + .filter(|path| !path.is_empty()) + { + replaced_content_paths.push(path); + } + } let count = upsert_chunks_tx(&tx, chunks)?; tx.commit()?; - Ok(count) - }) + Ok((count, replaced_content_paths)) + })?; + super::store_delete::remove_chunk_content_files(config, &replaced_content_paths); + Ok(count) } pub fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result { @@ -73,8 +90,8 @@ pub fn upsert_chunks_tx(tx: &Transaction<'_>, chunks: &[Chunk]) -> Result const UPSERT_SQL: &str = "INSERT INTO mem_tree_chunks ( id, source_kind, source_id, path_scope, source_ref, owner, timestamp_ms, time_range_start_ms, time_range_end_ms, - tags_json, content, token_count, seq_in_source, created_at_ms - ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14) + tags_json, content, token_count, seq_in_source, created_at_ms, lifecycle_status + ) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15) ON CONFLICT(id) DO UPDATE SET source_kind = excluded.source_kind, source_id = excluded.source_id, @@ -91,7 +108,7 @@ const UPSERT_SQL: &str = "INSERT INTO mem_tree_chunks ( created_at_ms = excluded.created_at_ms, content_path = NULL, content_sha256 = NULL, - lifecycle_status = 'admitted'"; + lifecycle_status = excluded.lifecycle_status"; /// Bind and execute `UPSERT_SQL` once per chunk against an already-prepared /// statement. Split out from [`upsert_chunks`] so the statement is prepared @@ -121,6 +138,7 @@ fn upsert_chunks_with_statement( chunk.token_count, chunk.seq_in_source, chunk.created_at.timestamp_millis(), + CHUNK_STATUS_ADMITTED, ])?; } Ok(()) diff --git a/src/memory/chunks/store_delete.rs b/src/memory/chunks/store_delete.rs index 4f75d00..450928f 100644 --- a/src/memory/chunks/store_delete.rs +++ b/src/memory/chunks/store_delete.rs @@ -366,7 +366,7 @@ pub fn delete_orphaned_source_tree( /// Best-effort removal of on-disk chunk content files, with strict sandboxing: /// a `content_path` that escapes the content root (via `..`, an absolute path, /// or a symlink pointing outside) is refused rather than followed. -fn remove_chunk_content_files(config: &MemoryConfig, content_paths: &[String]) { +pub(super) fn remove_chunk_content_files(config: &MemoryConfig, content_paths: &[String]) { use std::path::{Component, Path}; let root = content_root(config); diff --git a/src/memory/chunks/store_list.rs b/src/memory/chunks/store_list.rs index 5539c43..78447fa 100644 --- a/src/memory/chunks/store_list.rs +++ b/src/memory/chunks/store_list.rs @@ -100,13 +100,11 @@ fn append_source_scope( if index > 0 { sql.push_str(" OR "); } - sql.push_str("source_id = ? OR source_id LIKE ? ESCAPE '\\'"); + sql.push_str("source_id = ? OR substr(source_id, 1, length(?)) = ?"); bound.push(Box::new(source_id.clone())); - let escaped = source_id - .replace('\\', "\\\\") - .replace('%', "\\%") - .replace('_', "\\_"); - bound.push(Box::new(format!("mem_src:{escaped}:%"))); + let prefix = format!("mem_src:{source_id}:"); + bound.push(Box::new(prefix.clone())); + bound.push(Box::new(prefix)); } sql.push(')'); } diff --git a/src/memory/config.rs b/src/memory/config.rs index b46134a..74f2545 100644 --- a/src/memory/config.rs +++ b/src/memory/config.rs @@ -158,6 +158,13 @@ impl MemoryConfig { self.tree.summary_overhead_reserve_tokens < self.tree.input_token_budget, "tree.summary_overhead_reserve_tokens must be smaller than tree.input_token_budget" ); + anyhow::ensure!( + self.tree + .output_token_budget + .checked_add(self.tree.summary_overhead_reserve_tokens) + .is_some_and(|reserved| reserved < self.tree.input_token_budget), + "tree output and summary overhead budgets must leave positive input headroom" + ); anyhow::ensure!( self.tree.summary_fanout > 0, "tree.summary_fanout must be greater than zero" diff --git a/src/memory/config/policy.rs b/src/memory/config/policy.rs index cd60574..bf6c5f3 100644 --- a/src/memory/config/policy.rs +++ b/src/memory/config/policy.rs @@ -2,6 +2,8 @@ use serde::{Deserialize, Serialize}; +pub(crate) const DEFAULT_MAX_DEFER_AGE_MS: i64 = 7 * 24 * 60 * 60 * 1_000; + /// Runtime bounds shared by retrieval entry points. #[derive(Debug, Clone, Serialize, Deserialize)] #[serde(default)] @@ -144,7 +146,7 @@ impl Default for QueueConfig { retry_cap_ms: 60 * 60 * 1_000, max_attempts: 5, llm_permits: 1, - max_defer_age_ms: 7 * 24 * 60 * 60 * 1_000, + max_defer_age_ms: DEFAULT_MAX_DEFER_AGE_MS, } } } @@ -164,3 +166,7 @@ impl QueueConfig { Ok(()) } } + +#[cfg(test)] +#[path = "policy_tests.rs"] +mod tests; diff --git a/src/memory/config/policy_tests.rs b/src/memory/config/policy_tests.rs new file mode 100644 index 0000000..e4c1507 --- /dev/null +++ b/src/memory/config/policy_tests.rs @@ -0,0 +1,47 @@ +use super::*; + +#[test] +fn retrieval_limits_validate_defaults_and_reject_invalid_bounds() { + let mut limits = RetrievalLimits::default(); + limits.validate().unwrap(); + limits.default_limit = 0; + assert!(limits.validate().is_err()); + limits = RetrievalLimits::default(); + limits.max_limit = limits.default_limit - 1; + assert!(limits.validate().is_err()); + limits = RetrievalLimits::default(); + limits.default_graph_hops = limits.max_graph_hops + 1; + assert!(limits.validate().is_err()); + limits = RetrievalLimits::default(); + limits.freshness_half_life_days = f64::NAN; + assert!(limits.validate().is_err()); +} + +#[test] +fn scoring_policy_validates_threshold_ranges_and_order() { + let mut policy = ScoringPolicyConfig::default(); + policy.validate().unwrap(); + policy.drop_threshold = 1.1; + assert!(policy.validate().is_err()); + policy = ScoringPolicyConfig::default(); + policy.definite_drop_threshold = 0.9; + policy.definite_keep_threshold = 0.1; + assert!(policy.validate().is_err()); +} + +#[test] +fn queue_policy_validates_timings_and_limits() { + let mut queue = QueueConfig::default(); + queue.validate().unwrap(); + queue.retry_cap_ms = queue.retry_base_ms - 1; + assert!(queue.validate().is_err()); + queue = QueueConfig::default(); + queue.max_attempts = 0; + assert!(queue.validate().is_err()); + queue = QueueConfig::default(); + queue.llm_permits = 0; + assert!(queue.validate().is_err()); + queue = QueueConfig::default(); + queue.max_defer_age_ms = 0; + assert!(queue.validate().is_err()); +} diff --git a/src/memory/conversations/bus.rs b/src/memory/conversations/bus.rs index ee8ca2b..2b13d0e 100644 --- a/src/memory/conversations/bus.rs +++ b/src/memory/conversations/bus.rs @@ -28,7 +28,8 @@ use chrono::Utc; use serde_json::json; use super::{ - append_message, ensure_thread, get_messages, ConversationMessage, CreateConversationThread, + append_message, ensure_thread, get_messages, list_threads, ConversationMessage, + CreateConversationThread, }; static CONVERSATION_PERSISTENCE_WORKSPACE: OnceLock>> = OnceLock::new(); @@ -283,12 +284,27 @@ fn persist_channel_turn( workspace_dir: &Path, descriptor: ChannelTurnDescriptor<'_>, ) -> Result<(), String> { - let thread_id = persisted_channel_thread_id( + let collision_safe_id = persisted_channel_thread_id( descriptor.channel, descriptor.sender, descriptor.reply_target, descriptor.thread_ts, ); + let legacy_id = legacy_channel_thread_id( + descriptor.channel, + descriptor.sender, + descriptor.reply_target, + descriptor.thread_ts, + ); + let thread_id = if legacy_id != collision_safe_id + && list_threads(workspace_dir.to_path_buf())? + .iter() + .any(|thread| thread.id == legacy_id) + { + legacy_id + } else { + collision_safe_id + }; let title = channel_thread_title( descriptor.channel, descriptor.sender, @@ -384,6 +400,27 @@ fn persisted_channel_thread_id( format!("channel:{key}") } +/// Derive the pre-collision-fix id so existing underscore-containing channel +/// threads can continue receiving turns. New threads use +/// [`persisted_channel_thread_id`]. +fn legacy_channel_thread_id( + channel: &str, + sender: &str, + reply_target: &str, + thread_ts: Option<&str>, +) -> String { + let base_key = format!("{channel}_{sender}_{reply_target}"); + let key = if channel == "telegram" { + base_key + } else { + match thread_ts.and_then(non_empty_trimmed) { + Some(thread_ts) => format!("{base_key}_thread:{thread_ts}"), + None => base_key, + } + }; + format!("channel:{key}") +} + /// Build the human-readable thread title shown for a channel thread. /// Cosmetic only — never parsed back — so it does not need the collision /// safety `persisted_channel_thread_id` requires. diff --git a/src/memory/conversations/inverted_index.rs b/src/memory/conversations/inverted_index.rs index 00efe75..e29ca14 100644 --- a/src/memory/conversations/inverted_index.rs +++ b/src/memory/conversations/inverted_index.rs @@ -335,15 +335,10 @@ impl InvertedIndex { /// posting lists: `acc` is rewritten in place once per remaining /// ngram, allocating zero intermediate sets. /// - /// NOTE: the `None` case in `search`'s caller treats "no ngrams" as "every - /// live doc is a candidate" (a full-corpus scan), not as "no candidates". - /// For a large corpus this routinely exceeds `LARGE_CANDIDATE_LIMIT` and - /// trips the recency-only fallback (see `search`), returning score-`0.0` - /// hits that are visually indistinguishable from genuine substring - /// matches. Query terms that are too short to ngram-index (sub-3-byte - /// non-CJK terms already fail `MIN_TERM_BYTES`, so this mostly affects - /// single CJK characters) degrade silently rather than returning no - /// results. + /// The `None` case intentionally performs an uncapped full-corpus scan and + /// does not invoke the recency fallback. This avoids fabricating score-0.0 + /// hits for short or single-CJK terms, at the cost of scanning every live + /// document for those queries. fn candidates_for_term(&self, term: &str) -> Option> { let term_ngrams = ngrams(term); if term_ngrams.is_empty() { diff --git a/src/memory/conversations/store_index.rs b/src/memory/conversations/store_index.rs index 106be6c..ce44322 100644 --- a/src/memory/conversations/store_index.rs +++ b/src/memory/conversations/store_index.rs @@ -157,15 +157,23 @@ impl ConversationStore { pub(super) fn list_threads_unlocked(&self) -> Result, String> { let mut index = self.thread_index_unlocked()?; - // Reconcile the derived stat trail against the authoritative message - // files. Appending a message and its compact stat event spans two files - // and cannot be one filesystem transaction; this check repairs either - // crash window instead of trusting a possibly-short stat history. - let thread_ids = index.keys().cloned().collect::>(); + // Reconcile only cold/recovery entries whose derived stat trail is + // absent. Once stats exist, routine list calls stay header-only rather + // than rescanning every message file. + let thread_ids = index + .iter() + .filter(|(_, entry)| entry.message_count.is_none() || entry.last_message_at.is_none()) + .map(|(thread_id, _)| thread_id.clone()) + .collect::>(); if !thread_ids.is_empty() { let threads_path = self.ensure_root()?.join(THREADS_FILENAME); for thread_id in &thread_ids { - let (count, last_message_at) = self.measure_messages_unlocked(thread_id)?; + let Ok((count, last_message_at)) = self.measure_messages_unlocked(thread_id) else { + // One unreadable transcript must not make thread + // navigation unavailable. Leave it unreconciled so a + // later call can retry after repair. + continue; + }; // Treat created_at as last_message_at when there are no // messages — keeps the sort key meaningful and matches the // pre-refactor semantics. @@ -252,8 +260,19 @@ impl ConversationStore { Some(entry) => entry, None => return Ok(None), }; - let (message_count, last_at) = self.measure_messages_unlocked(thread_id)?; - let last_message_at = last_at.unwrap_or_else(|| entry.created_at.clone()); + let (message_count, last_message_at) = match (entry.message_count, &entry.last_message_at) { + (Some(count), Some(last)) => (count, last.clone()), + _ => match self.measure_messages_unlocked(thread_id) { + Ok((count, last)) => (count, last.unwrap_or_else(|| entry.created_at.clone())), + Err(_) => ( + entry.message_count.unwrap_or(0), + entry + .last_message_at + .clone() + .unwrap_or_else(|| entry.created_at.clone()), + ), + }, + }; Ok(Some(ConversationThread { id: thread_id.to_string(), title: entry.title.clone(), diff --git a/src/memory/conversations/store_tests_more.rs b/src/memory/conversations/store_tests_more.rs index 0dcbe2d..dd19b07 100644 --- a/src/memory/conversations/store_tests_more.rs +++ b/src/memory/conversations/store_tests_more.rs @@ -46,8 +46,9 @@ fn list_threads_reconciles_stats_with_authoritative_message_files() { // Warm-up: list_threads folds the MessageAppended entries. let _ = store.list_threads().unwrap(); - // Removing the authoritative transcript must not leave stale derived - // counts behind. + // Removing the transcript after a stats snapshot does not force routine + // list calls to rescan every message file. Cached navigation metadata + // remains available; explicit recovery can rebuild it when needed. let messages_dir = temp .path() .join("memory") @@ -63,8 +64,8 @@ fn list_threads_reconciles_stats_with_authoritative_message_files() { let threads = store.list_threads().unwrap(); assert_eq!(threads.len(), 1); - assert_eq!(threads[0].message_count, 0); - assert_eq!(threads[0].last_message_at, "2026-04-10T12:00:00Z"); + assert_eq!(threads[0].message_count, 3); + assert_eq!(threads[0].last_message_at, "2026-04-10T12:03:00Z"); } #[test] @@ -115,11 +116,12 @@ fn backfill_writes_stats_snapshot_for_legacy_threads() { "expected backfilled Stats entry in threads.jsonl, got:\n{log}", ); - // The message file remains authoritative even after a Stats snapshot. + // Once a Stats snapshot exists, routine reads remain header-only even if + // the transcript later becomes unavailable. std::fs::remove_file(&messages_file).unwrap(); let threads2 = store.list_threads().unwrap(); - assert_eq!(threads2[0].message_count, 0); - assert_eq!(threads2[0].last_message_at, "2026-04-10T08:00:00Z"); + assert_eq!(threads2[0].message_count, 2); + assert_eq!(threads2[0].last_message_at, "2026-04-10T09:05:00Z"); } #[test] diff --git a/src/memory/diff/ledger_helpers.rs b/src/memory/diff/ledger_helpers.rs index 92e1865..15ab6b4 100644 --- a/src/memory/diff/ledger_helpers.rs +++ b/src/memory/diff/ledger_helpers.rs @@ -19,6 +19,10 @@ pub(super) fn read_marker_ref(source_id: &str) -> String { format!("{READ_MARKER_PREFIX}{}", encode_source_id(source_id)) } +/// Build the snapshot commit subject and stable trailer block. +/// +/// Labels are reduced to one safe trailer line; source identity, kind, +/// trigger, item count, and capture time keep their machine-readable layout. pub(super) fn build_commit_message( meta: &SnapshotMeta, item_count: u32, @@ -66,6 +70,9 @@ pub(super) fn parse_trailers(message: &str) -> HashMap { map } +/// Validate a source id before using it in ledger refs and trailers. +/// +/// Empty ids and ids containing control characters are rejected. pub(super) fn validate_source_id(source_id: &str) -> Result<()> { anyhow::ensure!(!source_id.trim().is_empty(), "source id must not be blank"); anyhow::ensure!( @@ -75,6 +82,8 @@ pub(super) fn validate_source_id(source_id: &str) -> Result<()> { Ok(()) } +/// Encode one checkpoint using the stable human-readable subject and trailer +/// schema consumed by [`checkpoint_from_message`]. pub(super) fn checkpoint_message( label: &str, snapshot_ids: &[String], @@ -88,6 +97,8 @@ pub(super) fn checkpoint_message( payload.to_string() } +/// Decode a checkpoint from commit trailers, rejecting missing, malformed, or +/// invalid source identity and numeric fields. pub(super) fn checkpoint_from_message(id: &str, message: &str) -> Result { let value: serde_json::Value = serde_json::from_str(message.trim()) .with_context(|| format!("checkpoint '{id}' has invalid JSON metadata"))?; diff --git a/src/memory/fsutil.rs b/src/memory/fsutil.rs index c2f4939..4f1451a 100644 --- a/src/memory/fsutil.rs +++ b/src/memory/fsutil.rs @@ -15,6 +15,27 @@ use std::io; use std::path::Path; +/// Sanitize one machine identifier into a path component while preserving +/// collision resistance. Safe, non-empty inputs retain their spelling; +/// transformed or empty inputs receive a short digest of the exact raw value. +pub(crate) fn sanitize_component_with_digest(raw: &str, allowed: impl Fn(char) -> bool) -> String { + use sha2::{Digest, Sha256}; + + let sanitized = raw + .chars() + .map(|character| if allowed(character) { character } else { '_' }) + .collect::(); + if sanitized == raw && !sanitized.is_empty() { + return sanitized; + } + let digest = Sha256::digest(raw.as_bytes()); + let suffix = digest[..6] + .iter() + .map(|byte| format!("{byte:02x}")) + .collect::(); + format!("{sanitized}-{suffix}") +} + /// Atomically write `bytes` to `path`, replacing any existing file. /// /// The bytes are written to a hidden same-directory temp file and fsynced, then diff --git a/src/memory/ingest/pipeline.rs b/src/memory/ingest/pipeline.rs index 2a303b7..23a8ebc 100644 --- a/src/memory/ingest/pipeline.rs +++ b/src/memory/ingest/pipeline.rs @@ -151,7 +151,11 @@ async fn persist_score_enqueue( &chunk.id, CHUNK_STATUS_PENDING_EXTRACTION, )?; - jobs += usize::from(sink.enqueue_extract_tx(&transaction, &chunk.id)?); + jobs += usize::from(sink.enqueue_extract_tx( + &transaction, + &chunk.id, + config.queue.max_attempts, + )?); } transaction.commit()?; Ok(Some((chunks_written, jobs))) diff --git a/src/memory/ingest/pipeline_tests.rs b/src/memory/ingest/pipeline_tests.rs index af241c4..10a9d91 100644 --- a/src/memory/ingest/pipeline_tests.rs +++ b/src/memory/ingest/pipeline_tests.rs @@ -29,6 +29,7 @@ impl TreeJobSink for RecordingJobSink { &self, _tx: &rusqlite::Transaction<'_>, chunk_id: &str, + _default_max_attempts: u32, ) -> anyhow::Result { self.ids.lock().unwrap().push(chunk_id.to_string()); Ok(true) @@ -48,6 +49,7 @@ impl TreeJobSink for TogglingJobSink { &self, _tx: &rusqlite::Transaction<'_>, chunk_id: &str, + _default_max_attempts: u32, ) -> anyhow::Result { if self.fail.load(Ordering::SeqCst) { anyhow::bail!("simulated enqueue failure"); diff --git a/src/memory/ingest/types.rs b/src/memory/ingest/types.rs index 3c31483..621ee73 100644 --- a/src/memory/ingest/types.rs +++ b/src/memory/ingest/types.rs @@ -14,7 +14,12 @@ use crate::memory::queue::types::{ExtractChunkPayload, NewJob}; pub trait TreeJobSink: Send + Sync { /// Enqueue an `extract_chunk` job for `chunk_id` inside `tx`. Returns true /// when a new job was created and false for a deliberate no-op or dedupe. - fn enqueue_extract_tx(&self, tx: &Transaction<'_>, chunk_id: &str) -> Result; + fn enqueue_extract_tx( + &self, + tx: &Transaction<'_>, + chunk_id: &str, + default_max_attempts: u32, + ) -> Result; } /// A no-op [`TreeJobSink`] that drops every job. Useful when a caller only wants @@ -23,7 +28,12 @@ pub trait TreeJobSink: Send + Sync { pub struct NullJobSink; impl TreeJobSink for NullJobSink { - fn enqueue_extract_tx(&self, _tx: &Transaction<'_>, _chunk_id: &str) -> Result { + fn enqueue_extract_tx( + &self, + _tx: &Transaction<'_>, + _chunk_id: &str, + _default_max_attempts: u32, + ) -> Result { Ok(false) } } @@ -33,11 +43,19 @@ impl TreeJobSink for NullJobSink { pub struct QueueJobSink; impl TreeJobSink for QueueJobSink { - fn enqueue_extract_tx(&self, tx: &Transaction<'_>, chunk_id: &str) -> Result { + fn enqueue_extract_tx( + &self, + tx: &Transaction<'_>, + chunk_id: &str, + default_max_attempts: u32, + ) -> Result { let job = NewJob::extract_chunk(&ExtractChunkPayload { chunk_id: chunk_id.to_string(), })?; - Ok(crate::memory::queue::store::enqueue_tx(tx, &job)?.is_some()) + Ok( + crate::memory::queue::store::enqueue_tx_with_default(tx, &job, default_max_attempts)? + .is_some(), + ) } } diff --git a/src/memory/queue/ops_tests.rs b/src/memory/queue/ops_tests.rs index 976fc7a..2308e6c 100644 --- a/src/memory/queue/ops_tests.rs +++ b/src/memory/queue/ops_tests.rs @@ -25,6 +25,7 @@ fn count_reembed_jobs(cfg: &MemoryConfig) -> u64 { #[test] fn backfill_flag_setters_are_wired() { + let _flag_guard = crate::memory::queue::test_support::BACKFILL_FLAG_TEST_LOCK.lock(); // The flag is process-global and shared across parallel tests, so this only // exercises the setters compile/run path (it does not assert the value, // which other tests may concurrently toggle). diff --git a/src/memory/queue/runtime/mod.rs b/src/memory/queue/runtime/mod.rs index 96e1f69..8716a6d 100644 --- a/src/memory/queue/runtime/mod.rs +++ b/src/memory/queue/runtime/mod.rs @@ -9,9 +9,9 @@ //! ## Loops //! //! - `run_worker`: repeatedly runs one job, sleeping between polls. Idle polls -//! use the configured backoff; transient SQLite errors -//! back off by kind (busy / I/O / disk-full) and re-poll; corruption is fatal -//! and returned to the host (mirroring OpenHuman's quarantine policy). +//! use the configured backoff; transient SQLite errors back off by kind +//! (busy / I/O / disk-full) and re-poll; corruption triggers database +//! recovery and the loop continues polling. //! - `run_scheduler`: on a fixed cadence, recover expired leases, enqueue //! `flush_stale` (dedupe-safe), and `self_heal` transiently-failed jobs. //! - `run`: bootstrap once, then drive both loops concurrently until diff --git a/src/memory/queue/store.rs b/src/memory/queue/store.rs index 5ec2f1a..567c211 100644 --- a/src/memory/queue/store.rs +++ b/src/memory/queue/store.rs @@ -46,7 +46,9 @@ pub(crate) const DEFAULT_MAX_ATTEMPTS: u32 = 5; /// `ready`/`running`) shares it. Returns `Some(id)` if inserted, `None` if a /// duplicate was suppressed. pub fn enqueue(config: &MemoryConfig, job: &NewJob) -> Result> { - with_connection(config, |conn| enqueue_conn(conn, job)) + with_connection(config, |conn| { + enqueue_conn(conn, job, config.queue.max_attempts) + }) } /// Enqueue inside a caller-owned transaction. Use this when the producer is @@ -54,14 +56,28 @@ pub fn enqueue(config: &MemoryConfig, job: &NewJob) -> Result> { /// insert lands atomically with the side-effect. `Transaction` derefs to /// `Connection`, so callers just pass `&tx`. pub fn enqueue_tx(tx: &Transaction<'_>, job: &NewJob) -> Result> { - enqueue_conn(tx, job) + enqueue_conn(tx, job, DEFAULT_MAX_ATTEMPTS) +} + +/// Enqueue inside a transaction using the caller's configured default retry +/// budget when the job does not override it. +pub(crate) fn enqueue_tx_with_default( + tx: &Transaction<'_>, + job: &NewJob, + default_max_attempts: u32, +) -> Result> { + enqueue_conn(tx, job, default_max_attempts) } -pub(crate) fn enqueue_conn(conn: &Connection, job: &NewJob) -> Result> { +pub(crate) fn enqueue_conn( + conn: &Connection, + job: &NewJob, + default_max_attempts: u32, +) -> Result> { let id = format!("job:{}", Uuid::new_v4()); let now_ms = Utc::now().timestamp_millis(); let available_at = job.available_at_ms.unwrap_or(now_ms); - let max_attempts = job.max_attempts.unwrap_or(DEFAULT_MAX_ATTEMPTS) as i64; + let max_attempts = job.max_attempts.unwrap_or(default_max_attempts) as i64; let inserted = conn.execute( "INSERT OR IGNORE INTO mem_tree_jobs ( diff --git a/src/memory/queue/store_settle.rs b/src/memory/queue/store_settle.rs index b736e60..896d8fe 100644 --- a/src/memory/queue/store_settle.rs +++ b/src/memory/queue/store_settle.rs @@ -16,14 +16,9 @@ use crate::memory::config::MemoryConfig; use crate::memory::queue::store::backoff_ms_with_policy; use crate::memory::queue::types::{Job, JobFailure, NewJob}; -/// Maximum lifetime of a repeatedly deferred job. Deferral is intentionally -/// free of the handler failure budget, but it must still have a terminal bound. -#[cfg(test)] -pub(crate) const MAX_DEFER_AGE_MS: i64 = 7 * 24 * 60 * 60 * 1_000; - /// Mark a claimed job as `done`. Clears the lock and stamps `completed_at_ms`. pub fn mark_done(config: &MemoryConfig, job: &Job) -> Result<()> { - mark_done_with_followups(config, job, &[]) + mark_done_with_followups(config, job, &[]).map(|_| ()) } /// Mark a claimed job done and enqueue all of its follow-up jobs in the same @@ -33,7 +28,7 @@ pub(crate) fn mark_done_with_followups( config: &MemoryConfig, job: &Job, follow_ups: &[NewJob], -) -> Result<()> { +) -> Result { let job_id = &job.id; let claim_attempts = job.attempts as i64; let claim_started_at = job.started_at_ms; @@ -53,11 +48,15 @@ pub(crate) fn mark_done_with_followups( )?; if updated != 0 { for follow_up in follow_ups { - crate::memory::queue::store::enqueue_tx(&tx, follow_up)?; + crate::memory::queue::store::enqueue_tx_with_default( + &tx, + follow_up, + config.queue.max_attempts, + )?; } } tx.commit()?; - Ok(()) + Ok(updated != 0) }) } diff --git a/src/memory/queue/store_settle_tests.rs b/src/memory/queue/store_settle_tests.rs index 02085f5..a022923 100644 --- a/src/memory/queue/store_settle_tests.rs +++ b/src/memory/queue/store_settle_tests.rs @@ -344,6 +344,7 @@ fn mark_deferred_does_not_increment_attempts() { #[test] fn mark_deferred_parks_jobs_that_exceed_the_defer_age_bound() { let (_tmp, cfg) = test_config(); + let max_defer_age_ms = cfg.queue.max_defer_age_ms; let id = enqueue(&cfg, &extract_job("c-defer-expired", 5)) .unwrap() .expect("inserted"); @@ -354,7 +355,7 @@ fn mark_deferred_parks_jobs_that_exceed_the_defer_age_bound() { rusqlite::params![ Utc::now() .timestamp_millis() - .saturating_sub(MAX_DEFER_AGE_MS + 1), + .saturating_sub(max_defer_age_ms + 1), id, ], )?; diff --git a/src/memory/queue/test_support.rs b/src/memory/queue/test_support.rs index 60880a7..1c819e2 100644 --- a/src/memory/queue/test_support.rs +++ b/src/memory/queue/test_support.rs @@ -19,6 +19,9 @@ use crate::memory::queue::handlers::{ }; use crate::memory::queue::types::{AppendTarget, NodeRef, SealDocumentPayload, SealPayload}; +/// Serializes tests that mutate the process-wide re-embed backfill flag. +pub(crate) static BACKFILL_FLAG_TEST_LOCK: Mutex<()> = parking_lot::const_mutex(()); + /// Observable call counters, cloneable so a test can read them after `drain`. #[derive(Default)] pub(crate) struct Counts { diff --git a/src/memory/queue/worker.rs b/src/memory/queue/worker.rs index 65076b1..6c1291a 100644 --- a/src/memory/queue/worker.rs +++ b/src/memory/queue/worker.rs @@ -20,7 +20,8 @@ //! The Sentry-once emission and the storage-degraded flag stay host-owned //! regardless. -use std::sync::LazyLock; +use std::collections::HashMap; +use std::sync::{Arc, LazyLock}; use anyhow::Result; @@ -40,6 +41,8 @@ use crate::memory::queue::types::{ /// upstream single-permit semaphore); LLM-bound jobs hold a permit for the /// duration of their handler. static LLM_GATE: LazyLock = LazyLock::new(LlmGate::default); +static CONFIGURED_LLM_GATES: LazyLock>>> = + LazyLock::new(|| parking_lot::Mutex::new(HashMap::new())); /// Short delay used when an LLM job is claimed while the process-wide gate is /// full. Deferring releases the job lease and attempt immediately, avoiding a @@ -75,7 +78,15 @@ pub fn bootstrap(config: &MemoryConfig) -> Result<(usize, usize)> { /// handler and drop it; the other blocks the only executor thread inside /// `acquire`). pub async fn run_once(config: &MemoryConfig, delegates: &dyn QueueDelegates) -> Result { - run_once_with_gate(config, delegates, &LLM_GATE).await + let gate = { + let mut gates = CONFIGURED_LLM_GATES.lock(); + Arc::clone( + gates + .entry(config.queue.llm_permits) + .or_insert_with(|| Arc::new(LlmGate::new(config.queue.llm_permits))), + ) + }; + run_once_with_gate(config, delegates, &gate).await } async fn run_once_with_gate( @@ -141,8 +152,8 @@ fn settle_planned_job( let arms_reembed = follow_ups .iter() .any(|follow_up| follow_up.kind == JobKind::ReembedBackfill); - mark_done_with_followups(config, job, &follow_ups)?; - if arms_reembed { + let committed = mark_done_with_followups(config, job, &follow_ups)?; + if committed && arms_reembed { set_backfill_in_progress(true); } diff --git a/src/memory/queue/worker_tests.rs b/src/memory/queue/worker_tests.rs index da48b49..ea7d328 100644 --- a/src/memory/queue/worker_tests.rs +++ b/src/memory/queue/worker_tests.rs @@ -113,6 +113,7 @@ async fn run_once_reschedules_reembed_jobs_that_defer() { #[test] fn terminal_reembed_failure_clears_process_backfill_flag() { + let _flag_guard = crate::memory::queue::test_support::BACKFILL_FLAG_TEST_LOCK.lock(); let (_tmp, cfg) = test_config(); let mut new_job = NewJob::reembed_backfill(&ReembedBackfillPayload { signature: "provider=test;model=terminal;dims=3".into(), diff --git a/src/memory/retrieval/global.rs b/src/memory/retrieval/global.rs index ff49cd0..b760a4d 100644 --- a/src/memory/retrieval/global.rs +++ b/src/memory/retrieval/global.rs @@ -76,7 +76,7 @@ pub async fn query_global( /// optionally restricted to `[since_ms, until_ms]` and reranked by `query`. /// /// The entity-index timestamp window is pushed into SQLite before the -/// the configured topic lookup cap, so historical windows are not crowded +/// configured topic lookup cap, so historical windows are not crowded /// out by newer mentions. A window containing more than the cap remains /// intentionally bounded. pub async fn query_topic( @@ -121,11 +121,10 @@ pub async fn query_topic( /// (summary sidecar / chunk sidecar) for the optional rerank pass. /// /// Caps the already-windowed lookup at the configured topic limit. Soft-deleted -/// summaries are excluded (`node.deleted` -/// check below); leaf chunks have no equivalent tombstone check here, so a -/// dropped chunk that is still indexed can still surface in topic results. -/// An entity-index row whose `node_id` resolves to neither a summary nor a -/// chunk (a stale index entry) is silently skipped. +/// summaries are excluded via `node.deleted`; leaf chunks are filtered against +/// the fetched `dropped_chunk_ids` tombstone set. An entity-index row whose +/// `node_id` resolves to neither a summary nor a chunk (a stale index entry) is +/// silently skipped. fn resolve_topic_hits( config: &MemoryConfig, entity_id: &str, diff --git a/src/memory/retrieval/source.rs b/src/memory/retrieval/source.rs index 32d1662..24fb841 100644 --- a/src/memory/retrieval/source.rs +++ b/src/memory/retrieval/source.rs @@ -20,7 +20,7 @@ //! Unwindowed queries still materialize all selected summaries; callers should //! supply a window for bounded historical retrieval. -use anyhow::Result; +use anyhow::{Context, Result}; use chrono::{Duration, Utc}; use crate::memory::chunks::SourceKind; @@ -61,11 +61,15 @@ pub async fn query_source( limit.min(limits.max_limit) }; - let window = time_window_days.map(|days| { - let now = Utc::now(); - let start = now - Duration::days(days as i64); - (start.timestamp_millis(), now.timestamp_millis()) - }); + let window = time_window_days + .map(|days| -> Result<_> { + let now = Utc::now(); + let start = now + .checked_sub_signed(Duration::days(days as i64)) + .context("retrieval time window exceeds supported timestamp range")?; + Ok((start.timestamp_millis(), now.timestamp_millis())) + }) + .transpose()?; let scored = collect_source_hits(config, source_id, source_kind, window)?; let total = scored.len(); diff --git a/src/memory/score/mod.rs b/src/memory/score/mod.rs index 976b5a9..0a60632 100644 --- a/src/memory/score/mod.rs +++ b/src/memory/score/mod.rs @@ -376,25 +376,12 @@ pub fn persist_score( timestamp_ms: i64, tree_id: Option<&str>, ) -> Result<()> { - let row = score_row(result); - store::upsert_score(config, &row)?; - - // INSERT OR REPLACE never deletes rows whose entity_id is absent from the - // new extraction. Clear for both kept and dropped results so stale rows do - // not survive a re-score. - store::clear_entity_index_for_node(config, &result.chunk_id)?; - if result.kept && !result.canonical_entities.is_empty() { - store::index_entities( - config, - &result.canonical_entities, - &result.chunk_id, - "leaf", - timestamp_ms, - tree_id, - )?; - } - - Ok(()) + crate::memory::chunks::with_connection(config, |connection| { + let transaction = connection.unchecked_transaction()?; + persist_score_tx(&transaction, result, timestamp_ms, tree_id)?; + transaction.commit()?; + Ok(()) + }) } /// Transactional variant of [`persist_score`] — writes the score row and diff --git a/src/memory/score/store_query.rs b/src/memory/score/store_query.rs index 3abab9d..a978bf8 100644 --- a/src/memory/score/store_query.rs +++ b/src/memory/score/store_query.rs @@ -90,3 +90,7 @@ fn count_table(config: &MemoryConfig, table: &str) -> Result { Ok(count.max(0) as u64) }) } + +#[cfg(test)] +#[path = "store_query_tests.rs"] +mod tests; diff --git a/src/memory/score/store_query_tests.rs b/src/memory/score/store_query_tests.rs new file mode 100644 index 0000000..ed1798e --- /dev/null +++ b/src/memory/score/store_query_tests.rs @@ -0,0 +1,45 @@ +use super::*; + +fn config() -> (tempfile::TempDir, MemoryConfig) { + let temp = tempfile::tempdir().unwrap(); + let config = MemoryConfig::new(temp.path()); + (temp, config) +} + +fn insert_entity(config: &MemoryConfig, entity: &str, node: &str, timestamp_ms: i64, score: f64) { + with_connection(config, |connection| { + connection.execute( + "INSERT INTO mem_tree_entity_index + (entity_id, node_id, node_kind, entity_kind, surface, score, timestamp_ms, tree_id, is_user) + VALUES (?1, ?2, 'leaf', 'topic', ?1, ?3, ?4, NULL, 0)", + rusqlite::params![entity, node, score, timestamp_ms], + )?; + Ok(()) + }) + .unwrap(); +} + +#[test] +fn lookup_entity_in_window_filters_orders_and_limits() { + let (_temp, config) = config(); + insert_entity(&config, "topic:x", "old", 10, 0.2); + insert_entity(&config, "topic:x", "mid", 20, 0.4); + insert_entity(&config, "topic:x", "new", 30, 0.6); + insert_entity(&config, "topic:y", "other", 25, 0.9); + + let hits = lookup_entity_in_window(&config, "topic:x", Some(15), Some(35), Some(1)).unwrap(); + assert_eq!(hits.len(), 1); + assert_eq!(hits[0].node_id, "new"); +} + +#[test] +fn list_entity_ids_for_node_is_distinct_and_deterministic() { + let (_temp, config) = config(); + insert_entity(&config, "topic:a", "node", 10, 0.4); + insert_entity(&config, "topic:b", "node", 20, 0.9); + + assert_eq!( + list_entity_ids_for_node(&config, "node").unwrap(), + vec!["topic:b", "topic:a"] + ); +} diff --git a/src/memory/sources/readers/folder.rs b/src/memory/sources/readers/folder.rs index d88f238..4dbf688 100644 --- a/src/memory/sources/readers/folder.rs +++ b/src/memory/sources/readers/folder.rs @@ -114,7 +114,7 @@ impl SourceReader for FolderReader { .as_deref() .ok_or_else(|| MemoryError::Invalid("folder source requires a path".to_string()))?; - let pattern = source.glob.as_deref().unwrap_or("**/*.md"); + let pattern = source.glob.as_deref().unwrap_or(DEFAULT_GLOB); let matcher = glob_to_regex(pattern)?; let normalized_id = normalize_rel(Path::new(item_id)); if !matcher.is_match(&normalized_id) { diff --git a/src/memory/sources/registry.rs b/src/memory/sources/registry.rs index 8705932..377abbb 100644 --- a/src/memory/sources/registry.rs +++ b/src/memory/sources/registry.rs @@ -31,6 +31,12 @@ use super::types::{MemorySourceEntry, MemorySourcePatch, SourceKind}; /// safety boundary; this mutex closes the in-process lost-update window. static REGISTRY_MUTATION_LOCK: LazyLock> = LazyLock::new(|| Mutex::new(())); +fn mutation_guard() -> std::sync::MutexGuard<'static, ()> { + REGISTRY_MUTATION_LOCK + .lock() + .expect("source registry mutation lock poisoned") +} + /// Conservative default sync caps for a Composio toolkit, keyed by toolkit slug. /// /// Single source of truth for the cheap out-of-the-box sync volume. Applied to a @@ -180,9 +186,7 @@ impl SourceRegistry { /// Validate and add a new source. Fails if the id already exists. pub fn add(&self, entry: MemorySourceEntry) -> Result { - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); entry.validate().map_err(|e| anyhow!(e))?; let mut sources = self.list()?; if sources.iter().any(|s| s.id == entry.id) { @@ -196,9 +200,7 @@ impl SourceRegistry { /// Apply a [`MemorySourcePatch`] to an existing source, then re-validate and /// save. Fails if no source has the given id. pub fn update(&self, id: &str, patch: MemorySourcePatch) -> Result { - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); let mut sources = self.list()?; let entry = sources .iter_mut() @@ -215,9 +217,7 @@ impl SourceRegistry { /// Remove a source by id. Returns `true` if an entry was removed. pub fn remove(&self, id: &str) -> Result { - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); let mut sources = self.list()?; let before = sources.len(); sources.retain(|s| s.id != id); @@ -232,9 +232,7 @@ impl SourceRegistry { /// removed. Mirrors [`SourceRegistry::upsert_composio_source`], which keys /// composio sources on `connection_id` rather than the `src_*` id. pub fn remove_composio_source_by_connection_id(&self, connection_id: &str) -> Result { - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); let mut sources = self.list()?; let before = sources.len(); sources.retain(|s| { @@ -258,9 +256,7 @@ impl SourceRegistry { connection_id: &str, label: &str, ) -> Result { - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); let mut sources = self.list()?; let (entry, _was_insert) = upsert_composio_entry_in_place(&mut sources, toolkit, connection_id, label); @@ -273,9 +269,7 @@ impl SourceRegistry { if targets.is_empty() { return Ok(0); } - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); let mut sources = self.list()?; for (toolkit, connection_id, label) in targets { upsert_composio_entry_in_place(&mut sources, toolkit, connection_id, label); @@ -286,9 +280,7 @@ impl SourceRegistry { /// Enable every source and clear all per-source caps ("All In" mode). pub fn apply_all_in(&self) -> Result> { - let _guard = REGISTRY_MUTATION_LOCK - .lock() - .expect("source registry mutation lock poisoned"); + let _guard = mutation_guard(); let mut sources = self.list()?; for source in &mut sources { source.enabled = true; diff --git a/src/memory/sources/types.rs b/src/memory/sources/types.rs index 6f4d138..9519190 100644 --- a/src/memory/sources/types.rs +++ b/src/memory/sources/types.rs @@ -257,9 +257,15 @@ impl MemorySourcePatch { if self.query.is_some() && kind != SourceKind::TwitterQuery { return reject("query"); } + if self.since_days.is_some() && kind != SourceKind::TwitterQuery { + return reject("since_days"); + } if self.selector.is_some() && kind != SourceKind::WebPage { return reject("selector"); } + if self.max_items.is_some() && kind != SourceKind::RssFeed { + return reject("max_items"); + } if self.url.is_some() && kind != SourceKind::GithubRepo && kind != SourceKind::RssFeed diff --git a/src/memory/store/content/atomic.rs b/src/memory/store/content/atomic.rs index d31043f..db22409 100644 --- a/src/memory/store/content/atomic.rs +++ b/src/memory/store/content/atomic.rs @@ -60,7 +60,7 @@ pub fn write_if_new(abs_path: &Path, bytes: &[u8]) -> anyhow::Result { Ok(false) } else { Err(anyhow::anyhow!( - "publish {:?} -> {:?}: {e}", + "filesystem does not support required atomic hard-link publish {:?} -> {:?}: {e}", tmp_path, abs_path )) diff --git a/src/memory/store/entity_index/transaction.rs b/src/memory/store/entity_index/transaction.rs index 1d092ec..3e7ffb8 100644 --- a/src/memory/store/entity_index/transaction.rs +++ b/src/memory/store/entity_index/transaction.rs @@ -8,6 +8,11 @@ use super::store::{ use super::types::CanonicalEntity; /// Index canonical entities inside the caller's transaction. +/// +/// New rows are inserted with `is_user = false` because [`NoSelfIdentity`] +/// performs no identity classification. On conflict, +/// [`UPSERT_PRESERVE_USER_SQL`] deliberately preserves the existing `is_user` +/// value rather than resetting a prior identity match. pub fn index_entities_tx( tx: &Transaction<'_>, entities: &[CanonicalEntity], diff --git a/src/memory/store/kv.rs b/src/memory/store/kv.rs index 7e7e2a3..f95a8dd 100644 --- a/src/memory/store/kv.rs +++ b/src/memory/store/kv.rs @@ -16,6 +16,10 @@ use std::path::Path; use std::sync::Arc; use crate::memory::store::safety; + +fn parse_value_json(raw: &str, key: &str, context: &str) -> Result { + serde_json::from_str(raw).map_err(|error| format!("{context} JSON for key '{key}': {error}")) +} use crate::memory::types::MemoryKvRecord; const SCHEMA_SQL: &str = " @@ -245,8 +249,7 @@ impl KvStore { { let value_raw: String = row.get(1).map_err(|e| e.to_string())?; let key = row.get::<_, String>(0).map_err(|e| e.to_string())?; - let value = serde_json::from_str::(&value_raw) - .map_err(|e| format!("list_namespace JSON for key '{key}': {e}"))?; + let value = parse_value_json(&value_raw, &key, "list_namespace")?; out.push(json!({ "key": key, "value": value, @@ -291,8 +294,7 @@ impl KvStore { { let value_raw: String = row.get(1).map_err(|e| e.to_string())?; let key: String = row.get(0).map_err(|e| e.to_string())?; - let value = serde_json::from_str(&value_raw) - .map_err(|e| format!("records_namespace JSON for key '{key}': {e}"))?; + let value = parse_value_json(&value_raw, &key, "records_namespace")?; out.push(MemoryKvRecord { namespace: Some(ns.clone()), key, @@ -321,8 +323,7 @@ impl KvStore { { let value_raw: String = row.get(1).map_err(|e| e.to_string())?; let key: String = row.get(0).map_err(|e| e.to_string())?; - let value = serde_json::from_str(&value_raw) - .map_err(|e| format!("records_global JSON for key '{key}': {e}"))?; + let value = parse_value_json(&value_raw, &key, "records_global")?; out.push(MemoryKvRecord { namespace: None, key, diff --git a/src/memory/store/memory_trait.rs b/src/memory/store/memory_trait.rs index 5aca824..c8b66a8 100644 --- a/src/memory/store/memory_trait.rs +++ b/src/memory/store/memory_trait.rs @@ -43,19 +43,7 @@ fn entry_from_record(record: &MemoryRecord) -> MemoryEntry { } fn query_score(content: &str, query: &str) -> Option { - let terms = query - .split_whitespace() - .map(str::to_lowercase) - .collect::>(); - if terms.is_empty() { - return Some(1.0); - } - let content = content.to_lowercase(); - let matched = terms - .iter() - .filter(|term| content.contains(term.as_str())) - .count(); - (matched != 0).then_some(matched as f64 / terms.len() as f64) + super::query_match_score(content, query).map(f64::from) } #[async_trait] diff --git a/src/memory/store/mod.rs b/src/memory/store/mod.rs index f2e8912..32f65c2 100644 --- a/src/memory/store/mod.rs +++ b/src/memory/store/mod.rs @@ -49,3 +49,22 @@ pub use vectors::{ bytes_to_vec, cosine_similarity, vec_to_bytes, EmbeddingBackend, InertEmbedding, SearchResult, VectorStore, }; + +/// Score a case-insensitive whitespace-term query against content. +/// +/// Empty queries match fully; content with no matching term is excluded. +pub(super) fn query_match_score(content: &str, query: &str) -> Option { + let terms = query + .split_whitespace() + .map(str::to_lowercase) + .collect::>(); + if terms.is_empty() { + return Some(1.0); + } + let content = content.to_lowercase(); + let matched = terms + .iter() + .filter(|term| content.contains(term.as_str())) + .count(); + (matched != 0).then_some(matched as f32 / terms.len() as f32) +} diff --git a/src/memory/store/safety/pii/checks.rs b/src/memory/store/safety/pii/checks.rs index 0b64470..04fe09d 100644 --- a/src/memory/store/safety/pii/checks.rs +++ b/src/memory/store/safety/pii/checks.rs @@ -230,3 +230,7 @@ pub(super) fn valid_nino(s: &str) -> bool { } true } + +#[cfg(test)] +#[path = "checks_tests.rs"] +mod tests; diff --git a/src/memory/store/safety/pii/checks_tests.rs b/src/memory/store/safety/pii/checks_tests.rs new file mode 100644 index 0000000..77aeea5 --- /dev/null +++ b/src/memory/store/safety/pii/checks_tests.rs @@ -0,0 +1,43 @@ +use super::*; + +#[test] +fn tax_ids_enforce_lengths_checksums_and_repetition_rules() { + assert!(valid_cpf(&digits("529.982.247-25"))); + assert!(!valid_cpf(&digits("111.111.111-11"))); + assert!(!valid_cpf(&digits("5299822472"))); + assert!(valid_cnpj(&digits("11.222.333/0001-81"))); + assert!(!valid_cnpj(&digits("11.222.333/0001-82"))); + assert!(!valid_cnpj(&digits("00000000000000"))); + assert!(valid_cuit(&digits("20-12345678-6"))); + assert!(!valid_cuit(&digits("20-12345678-7"))); + assert!(!valid_cuit(&digits("2012345678"))); +} + +#[test] +fn payment_checksums_reject_bad_bounds_and_checksums() { + assert!(valid_luhn("4111 1111 1111 1111")); + assert!(!valid_luhn("4111 1111 1111 1112")); + assert!(!valid_luhn("7992739871")); + assert!(valid_iban("GB82 WEST 1234 5698 7654 32")); + assert!(!valid_iban("GB82 WEST 1234 5698 7654 33")); + assert!(!valid_iban("GB00")); +} + +#[test] +fn identity_validators_cover_checksums_reserved_values_and_prefixes() { + assert!(valid_verhoeff(&digits("234567890124"))); + assert!(!valid_verhoeff(&digits("134567890124"))); + assert!(!valid_verhoeff(&digits("234567890125"))); + assert!(valid_ssn("123-45-6789")); + assert!(!valid_ssn("666-45-6789")); + assert!(!valid_ssn("123-00-6789")); + assert!(!valid_ssn("123-45-0000")); + assert!(valid_dni_es("12345678Z")); + assert!(!valid_dni_es("12345678A")); + assert!(valid_nie_es("X1234567L")); + assert!(!valid_nie_es("A1234567L")); + assert!(valid_nino("AA123456A")); + assert!(!valid_nino("BG123456A")); + assert!(!valid_nino("DA123456A")); + assert!(!valid_nino("AA12345A")); +} diff --git a/src/memory/store/store.rs b/src/memory/store/store.rs index 64f3160..c2f7b31 100644 --- a/src/memory/store/store.rs +++ b/src/memory/store/store.rs @@ -81,16 +81,7 @@ impl MemoryStore for InMemoryMemoryStore { } async fn search(&self, query: MemoryQuery) -> MemoryResult> { - let terms = query - .text - .as_deref() - .map(str::to_lowercase) - .map(|text| { - text.split_whitespace() - .map(str::to_string) - .collect::>() - }) - .unwrap_or_default(); + let query_text = query.text.as_deref().unwrap_or_default(); let limit = query.limit.unwrap_or(20); let mut hits = self @@ -105,19 +96,7 @@ impl MemoryStore for InMemoryMemoryStore { .is_none_or(|namespace| record.namespace == namespace) }) .filter_map(|record| { - let score = if terms.is_empty() { - 1.0 - } else { - let content = record.content.to_lowercase(); - let matched = terms - .iter() - .filter(|term| content.contains(term.as_str())) - .count(); - if matched == 0 { - return None; - } - matched as f32 / terms.len() as f32 - }; + let score = super::query_match_score(&record.content, query_text)?; Some(SearchHit { record: record.clone(), diff --git a/src/memory/store/vectors/store.rs b/src/memory/store/vectors/store.rs index 4b8e8e2..db55065 100644 --- a/src/memory/store/vectors/store.rs +++ b/src/memory/store/vectors/store.rs @@ -84,11 +84,8 @@ impl VectorStore { /// On first open the backend name and dimensions are persisted to a /// `store_meta` table. On subsequent opens the stored dimensions are /// compared against the runtime backend and an error is returned if they - /// mismatch (prevents silent cosine-similarity corruption from - /// mixed-dimension vectors) — see the private `check_or_store_meta` helper - /// below for known gaps in that guard (an unreadable `store_meta` row is - /// indistinguishable from first-open, and a corrupt stored-dims string - /// silently disables the check rather than failing closed). + /// mismatch. Metadata read errors and malformed stored dimensions fail + /// closed rather than being treated as a first open. pub fn open(db_path: &Path, backend: Arc) -> anyhow::Result { if let Some(parent) = db_path.parent() { std::fs::create_dir_all(parent)?; @@ -123,16 +120,9 @@ impl VectorStore { /// Persist or validate the embedding configuration in `store_meta`. /// - /// NOTE: two related gaps, both currently present: - /// - A failed read of the `embed_dims` row (e.g. a transient I/O error) - /// collapses to `None` via `.ok()` and is treated identically to a - /// genuinely first-ever open, so it silently *overwrites* - /// `embed_provider`/`embed_dims` with the current runtime backend - /// instead of surfacing the read failure. - /// - The stored dims string is parsed with `.parse().unwrap_or(0)`; a - /// corrupted value collapses to `0`, and the mismatch check is skipped - /// whenever either side is `0` — so a corrupt row disables the guard - /// entirely instead of failing closed. + /// Missing metadata initializes the store. Query failures, malformed + /// stored dimensions, and runtime/stored dimension mismatches are returned + /// as errors so the compatibility guard fails closed. fn check_or_store_meta( conn: &Connection, backend: &dyn EmbeddingBackend, @@ -197,11 +187,8 @@ impl VectorStore { /// Inserts with a pre-computed embedding vector (skips the embed call). /// - /// NOTE: the vector's length is never checked against the store's - /// configured dimensions (`store_meta.embed_dims`). A wrong-dimension - /// vector is stored as-is; [`cosine_similarity`] then returns `0.0` for - /// any query against it (length mismatch), so the row silently never - /// surfaces in search results rather than erroring at insert time. + /// The vector length is checked against the active backend dimensions and + /// a mismatch is rejected before any row is written. pub fn insert_with_vector( &self, id: &str, @@ -441,7 +428,7 @@ pub fn vec_to_bytes(v: &[f32]) -> Vec { /// silently truncating corrupt trailing bytes. pub fn bytes_to_vec(bytes: &[u8]) -> anyhow::Result> { anyhow::ensure!( - bytes.len().is_multiple_of(4), + bytes.len() % 4 == 0, "invalid embedding blob byte length: {}", bytes.len() ); diff --git a/src/memory/store/vectors/store_tests_more.rs b/src/memory/store/vectors/store_tests_more.rs index 2216b38..87a26e7 100644 --- a/src/memory/store/vectors/store_tests_more.rs +++ b/src/memory/store/vectors/store_tests_more.rs @@ -30,12 +30,14 @@ fn search_by_vector_preserves_metadata() { #[test] fn search_handles_invalid_metadata_json() { let store = fake_store(2); + store + .insert_with_vector("bad", "ns", "text", &[1.0, 0.0], json!({})) + .unwrap(); { let conn = store.conn.lock(); conn.execute( - "INSERT INTO vectors (id, namespace, text, embedding, metadata, created_at, updated_at) - VALUES ('bad', 'ns', 'text', ?1, 'not-json', 0.0, 0.0)", - rusqlite::params![vec_to_bytes(&[1.0, 0.0])], + "UPDATE vectors SET metadata = 'not-json' WHERE id = 'bad' AND namespace = 'ns'", + [], ) .unwrap(); } diff --git a/src/memory/sync/composio/providers/slack_parse.rs b/src/memory/sync/composio/providers/slack_parse.rs index 0b9e0ef..2419f60 100644 --- a/src/memory/sync/composio/providers/slack_parse.rs +++ b/src/memory/sync/composio/providers/slack_parse.rs @@ -7,11 +7,14 @@ use serde_json::Value; use super::common::first_array; +/// Return the cached matcher for Slack `<@USERID>` mentions. pub(super) fn mention_regex() -> &'static regex::Regex { static REGEX: OnceLock = OnceLock::new(); REGEX.get_or_init(|| regex::Regex::new(r"<@(U[A-Z0-9]+)>").expect("Slack mention regex")) } +/// Replace Slack mention tokens with resolved display names, falling back to +/// the raw user id when the optional user map has no match. pub(super) fn replace_mentions( text: &str, users: Option<&serde_json::Map>, @@ -28,6 +31,7 @@ pub(super) fn replace_mentions( .into_owned() } +/// Read the first non-blank next cursor across supported response envelopes. pub(super) fn next_cursor(data: &Value) -> Option { [ "/data/response_metadata/next_cursor", @@ -43,6 +47,7 @@ pub(super) fn next_cursor(data: &Value) -> Option { .map(str::to_owned) } +/// Extract Slack search matches across legacy and nested response envelopes. pub(super) fn search_matches(data: &Value) -> Vec { first_array( data, @@ -55,6 +60,7 @@ pub(super) fn search_matches(data: &Value) -> Vec { ) } +/// Extract the search page count, defaulting to one when paging is absent. pub(super) fn search_total_pages(data: &Value) -> u32 { [ "/data/messages/paging/pages", @@ -67,11 +73,15 @@ pub(super) fn search_total_pages(data: &Value) -> u32 { .unwrap_or(1) as u32 } +/// Decode persisted per-scope cursors, returning an empty map for absent or +/// malformed JSON so synchronization can restart safely. pub(super) fn decode_cursors(raw: Option<&str>) -> BTreeMap { raw.and_then(|raw| serde_json::from_str(raw).ok()) .unwrap_or_default() } +/// Parse Slack's `seconds.fraction` timestamp into numeric components. +/// Missing fractions become zero; malformed numeric components return `None`. pub(super) fn parse_ts(ts: &str) -> Option<(i64, u64)> { let mut parts = ts.splitn(2, '.'); Some(( diff --git a/src/memory/sync/dispatcher.rs b/src/memory/sync/dispatcher.rs index befcb34..c2f2805 100644 --- a/src/memory/sync/dispatcher.rs +++ b/src/memory/sync/dispatcher.rs @@ -119,210 +119,5 @@ impl SyncDispatcher { } #[cfg(test)] -mod tests { - use std::collections::HashMap; - use std::sync::Mutex; - - use async_trait::async_trait; - - use super::*; - use crate::memory::sync::state::SyncStateStore; - use crate::memory::sync::traits::{SkillDocSink, SkillDocument, SyncEvent, SyncEventSink}; - - struct FakePipeline { - id: &'static str, - fail: bool, - init_fail: bool, - } - - #[async_trait] - impl SyncPipeline for FakePipeline { - fn id(&self) -> &str { - self.id - } - fn kind(&self) -> SyncPipelineKind { - SyncPipelineKind::Workspace - } - async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { - if self.init_fail { - anyhow::bail!("expected init failure") - } - Ok(()) - } - async fn tick(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result { - if self.fail { - anyhow::bail!("expected failure") - } - Ok(SyncOutcome { - records_ingested: 3, - more_pending: false, - actions_called: 0, - provider_cost_usd: 0.0, - note: None, - }) - } - } - - #[derive(Default)] - struct NoopHost(Mutex>); - #[async_trait] - impl SkillDocSink for NoopHost { - async fn store(&self, _: SkillDocument) -> anyhow::Result<()> { - Ok(()) - } - async fn delete(&self, _: &str, _: &str) -> anyhow::Result<()> { - Ok(()) - } - } - #[async_trait] - impl SyncEventSink for NoopHost { - async fn emit(&self, _: SyncEvent) -> anyhow::Result<()> { - Ok(()) - } - } - #[async_trait] - impl SyncStateStore for NoopHost { - async fn get( - &self, - namespace: &str, - key: &str, - ) -> anyhow::Result> { - Ok(self - .0 - .lock() - .unwrap() - .get(&format!("{namespace}:{key}")) - .cloned()) - } - async fn set( - &self, - namespace: &str, - key: &str, - value: &serde_json::Value, - ) -> anyhow::Result<()> { - self.0 - .lock() - .unwrap() - .insert(format!("{namespace}:{key}"), value.clone()); - Ok(()) - } - } - - fn context() -> SyncContext { - let host = Arc::new(NoopHost::default()); - SyncContext { - events: host.clone(), - documents: host.clone(), - state: host, - local_documents: None, - external_sources: None, - summariser: None, - } - } - - #[tokio::test] - async fn tick_all_is_deterministic_and_isolates_failures() { - let mut dispatcher = SyncDispatcher::new(); - dispatcher - .register(Arc::new(FakePipeline { - id: "z-fail", - fail: true, - init_fail: false, - })) - .unwrap(); - dispatcher - .register(Arc::new(FakePipeline { - id: "a-ok", - fail: false, - init_fail: false, - })) - .unwrap(); - assert_eq!(dispatcher.ids(), vec!["a-ok", "z-fail"]); - assert!(dispatcher - .register(Arc::new(FakePipeline { - id: "a-ok", - fail: false, - init_fail: false, - })) - .is_err()); - let results = dispatcher - .tick_all(&MemoryConfig::new("/tmp/unused"), &context()) - .await; - assert_eq!(results.len(), 2); - assert_eq!(results[0].outcome.as_ref().unwrap().records_ingested, 3); - assert!(results[1] - .error - .as_deref() - .unwrap() - .contains("expected failure")); - } - - #[tokio::test] - async fn register_rejects_blank_ids_and_tick_reports_unknown_pipeline() { - let mut dispatcher = SyncDispatcher::new(); - assert!(dispatcher - .register(Arc::new(FakePipeline { - id: " ", - fail: false, - init_fail: false, - })) - .is_err()); - let error = dispatcher - .tick("missing", &MemoryConfig::new("/tmp/unused"), &context()) - .await - .unwrap_err(); - assert!(error.to_string().contains("unknown sync pipeline")); - } - - #[tokio::test] - async fn init_all_and_individual_tick_preserve_success_and_failure_details() { - let mut dispatcher = SyncDispatcher::new(); - dispatcher - .register(Arc::new(FakePipeline { - id: "a-init-fails", - fail: false, - init_fail: true, - })) - .unwrap(); - dispatcher - .register(Arc::new(FakePipeline { - id: "b-ok", - fail: false, - init_fail: false, - })) - .unwrap(); - dispatcher - .register(Arc::new(FakePipeline { - id: "c-tick-fails", - fail: true, - init_fail: false, - })) - .unwrap(); - let config = MemoryConfig::new("/tmp/unused"); - let context = context(); - - let initialized = dispatcher.init_all(&config, &context).await; - assert!(initialized[0] - .error - .as_deref() - .unwrap() - .contains("expected init failure")); - assert!(initialized[1].outcome.is_some()); - assert_eq!( - dispatcher - .tick("b-ok", &config, &context) - .await - .unwrap() - .records_ingested, - 3 - ); - assert!(dispatcher - .tick("c-tick-fails", &config, &context) - .await - .is_err()); - - let encoded = serde_json::to_value(&initialized[0]).unwrap(); - assert_eq!(encoded["pipeline_id"], "a-init-fails"); - assert!(encoded.get("outcome").is_none()); - } -} +#[path = "dispatcher_tests.rs"] +mod tests; diff --git a/src/memory/sync/dispatcher_tests.rs b/src/memory/sync/dispatcher_tests.rs new file mode 100644 index 0000000..a2483b5 --- /dev/null +++ b/src/memory/sync/dispatcher_tests.rs @@ -0,0 +1,209 @@ +use std::collections::HashMap; +use std::sync::{Arc, Mutex}; + +use async_trait::async_trait; + +use super::*; +use crate::memory::sync::state::SyncStateStore; +use crate::memory::sync::traits::{SkillDocSink, SkillDocument, SyncEvent, SyncEventSink}; + +struct FakePipeline { + id: &'static str, + fail: bool, + init_fail: bool, +} + +#[async_trait] +impl SyncPipeline for FakePipeline { + fn id(&self) -> &str { + self.id + } + + fn kind(&self) -> SyncPipelineKind { + SyncPipelineKind::Workspace + } + + async fn init(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result<()> { + if self.init_fail { + anyhow::bail!("expected init failure") + } + Ok(()) + } + + async fn tick(&self, _: &MemoryConfig, _: &SyncContext) -> anyhow::Result { + if self.fail { + anyhow::bail!("expected failure") + } + Ok(SyncOutcome { + records_ingested: 3, + more_pending: false, + actions_called: 0, + provider_cost_usd: 0.0, + note: None, + }) + } +} + +#[derive(Default)] +struct NoopHost(Mutex>); + +#[async_trait] +impl SkillDocSink for NoopHost { + async fn store(&self, _: SkillDocument) -> anyhow::Result<()> { + Ok(()) + } + + async fn delete(&self, _: &str, _: &str) -> anyhow::Result<()> { + Ok(()) + } +} + +#[async_trait] +impl SyncEventSink for NoopHost { + async fn emit(&self, _: SyncEvent) -> anyhow::Result<()> { + Ok(()) + } +} + +#[async_trait] +impl SyncStateStore for NoopHost { + async fn get(&self, namespace: &str, key: &str) -> anyhow::Result> { + Ok(self + .0 + .lock() + .unwrap() + .get(&format!("{namespace}:{key}")) + .cloned()) + } + + async fn set( + &self, + namespace: &str, + key: &str, + value: &serde_json::Value, + ) -> anyhow::Result<()> { + self.0 + .lock() + .unwrap() + .insert(format!("{namespace}:{key}"), value.clone()); + Ok(()) + } +} + +fn context() -> SyncContext { + let host = Arc::new(NoopHost::default()); + SyncContext { + events: host.clone(), + documents: host.clone(), + state: host, + local_documents: None, + external_sources: None, + summariser: None, + } +} + +#[tokio::test] +async fn tick_all_is_deterministic_and_isolates_failures() { + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(Arc::new(FakePipeline { + id: "z-fail", + fail: true, + init_fail: false, + })) + .unwrap(); + dispatcher + .register(Arc::new(FakePipeline { + id: "a-ok", + fail: false, + init_fail: false, + })) + .unwrap(); + assert_eq!(dispatcher.ids(), vec!["a-ok", "z-fail"]); + assert!(dispatcher + .register(Arc::new(FakePipeline { + id: "a-ok", + fail: false, + init_fail: false, + })) + .is_err()); + let results = dispatcher + .tick_all(&MemoryConfig::new("/tmp/unused"), &context()) + .await; + assert_eq!(results.len(), 2); + assert_eq!(results[0].outcome.as_ref().unwrap().records_ingested, 3); + assert!(results[1] + .error + .as_deref() + .unwrap() + .contains("expected failure")); +} + +#[tokio::test] +async fn register_rejects_blank_ids_and_tick_reports_unknown_pipeline() { + let mut dispatcher = SyncDispatcher::new(); + assert!(dispatcher + .register(Arc::new(FakePipeline { + id: " ", + fail: false, + init_fail: false, + })) + .is_err()); + let error = dispatcher + .tick("missing", &MemoryConfig::new("/tmp/unused"), &context()) + .await + .unwrap_err(); + assert!(error.to_string().contains("unknown sync pipeline")); +} + +#[tokio::test] +async fn init_all_and_individual_tick_preserve_success_and_failure_details() { + let mut dispatcher = SyncDispatcher::new(); + dispatcher + .register(Arc::new(FakePipeline { + id: "a-init-fails", + fail: false, + init_fail: true, + })) + .unwrap(); + dispatcher + .register(Arc::new(FakePipeline { + id: "b-ok", + fail: false, + init_fail: false, + })) + .unwrap(); + dispatcher + .register(Arc::new(FakePipeline { + id: "c-tick-fails", + fail: true, + init_fail: false, + })) + .unwrap(); + let config = MemoryConfig::new("/tmp/unused"); + let context = context(); + + let initialized = dispatcher.init_all(&config, &context).await; + assert!(initialized[0] + .error + .as_deref() + .unwrap() + .contains("expected init failure")); + assert!(initialized[1].outcome.is_some()); + assert_eq!( + dispatcher + .tick("b-ok", &config, &context) + .await + .unwrap() + .records_ingested, + 3 + ); + assert!(dispatcher + .tick("c-tick-fails", &config, &context) + .await + .is_err()); + + let encoded = serde_json::to_value(&initialized[0]).unwrap(); + assert_eq!(encoded["pipeline_id"], "a-init-fails"); + assert!(encoded.get("outcome").is_none()); +} diff --git a/src/memory/tool_memory/store.rs b/src/memory/tool_memory/store.rs index 4e0ba17..b505744 100644 --- a/src/memory/tool_memory/store.rs +++ b/src/memory/tool_memory/store.rs @@ -65,26 +65,10 @@ impl ToolMemoryStore { /// `created_at` is preserved. `tool_name` is sourced from the rule /// itself to avoid storage/namespace skew. /// - /// Validation is limited to "non-empty `tool_name`" and "non-empty - /// `rule` body" — callers should be aware of the following gaps: - /// - /// NOTE: `rule.rule` is not rejected or sanitized for embedded newlines. - /// [`render_tool_memory_rules`](super::render::render_tool_memory_rules) - /// concatenates the rule body verbatim into the pinned system-prompt - /// block, so a stored rule containing `"...\n### \`shell\`\n- ..."` can - /// forge what looks like a second tool section inside a block the - /// prompt frames as a hard constraint. This matters because rules can - /// arrive from [`ToolMemorySource::PostTurn`] auto-capture of untrusted - /// tool-failure text. Sanitize/reject `\n`/`\r` before calling this if - /// the rule body may be attacker-influenced. - /// - /// NOTE: `tool_name` is stored verbatim (only the derived namespace is - /// lower-cased via [`tool_memory_namespace`]), so `"Email"` and - /// `"email"` land in the same namespace but are treated as two distinct - /// tools by [`list_rules`](Self::list_rules) grouping and by - /// [`render_tool_memory_rules`](super::render::render_tool_memory_rules)'s - /// per-tool heading. Normalize `tool_name` before calling if case - /// consistency matters to the caller. + /// `tool_name` is trimmed and lower-cased before new writes. Legacy + /// mixed-case rows remain unchanged until they are rewritten. Rule bodies + /// may contain newlines; the renderer treats them as body text rather than + /// allowing them to create additional tool sections. /// /// The read (`fetch_rule`) → write (`Memory::store`) transaction is /// serialized across all in-process store handles because the generic @@ -179,15 +163,9 @@ impl ToolMemoryStore { /// Returns the set of rules whose [`ToolMemoryPriority`] indicates /// they must be eagerly surfaced (Critical + High), grouped by tool - /// name. Result is bounded by [`TOOL_MEMORY_PROMPT_CAP`] entries - /// total — Critical rules are always preferred over High when the - /// cap is reached (the truncate happens after a Critical-first sort). - /// - /// NOTE: the cap is a hard `truncate`, not a per-priority reservation. - /// If more than [`TOOL_MEMORY_PROMPT_CAP`] Critical rules exist across - /// the scanned tools, the ones sorted past the cap (oldest `updated_at` - /// within Critical) are silently dropped from the prompt injection — - /// there is no overflow signal to the caller. + /// name. Every Critical rule is retained; [`TOOL_MEMORY_PROMPT_CAP`] + /// limits only the High-priority remainder. The result can therefore + /// exceed the cap when there are more Critical rules than the cap. /// /// `tools` constrains which tool namespaces to inspect; passing an /// empty slice scans every known tool namespace via diff --git a/src/memory/tree/bucket_seal.rs b/src/memory/tree/bucket_seal.rs index efd00a6..4863dc2 100644 --- a/src/memory/tree/bucket_seal.rs +++ b/src/memory/tree/bucket_seal.rs @@ -436,7 +436,11 @@ pub async fn seal_one_level_with_services( level: target_level, force_now_ms: None, }; - crate::memory::queue::enqueue_tx(&tx, &crate::memory::queue::NewJob::seal(&payload)?)?; + crate::memory::queue::store::enqueue_tx_with_default( + &tx, + &crate::memory::queue::NewJob::seal(&payload)?, + config.queue.max_attempts, + )?; } if target_level > current_max { diff --git a/src/memory/tree/document_seal.rs b/src/memory/tree/document_seal.rs index 9bb5982..34d28fe 100644 --- a/src/memory/tree/document_seal.rs +++ b/src/memory/tree/document_seal.rs @@ -85,11 +85,13 @@ pub async fn seal_document_subtree_with_services( let root_level = doc_root.level; with_connection(config, move |connection| { let transaction = connection.unchecked_transaction()?; - let (current_root, current_max): (Option, u32) = transaction.query_row( - "SELECT root_id, max_level FROM mem_tree_trees WHERE id = ?1", - [&tree.id], - |row| Ok((row.get(0)?, row.get(1)?)), - )?; + let (current_root, current_max, status): (Option, u32, String) = transaction + .query_row( + "SELECT root_id, max_level, status FROM mem_tree_trees WHERE id = ?1", + [&tree.id], + |row| Ok((row.get(0)?, row.get(1)?, row.get(2)?)), + )?; + anyhow::ensure!(status == "active", "tree '{}' is archived", tree.id); if current_root.as_deref() != Some(root_id.as_str()) || root_level > current_max { store::update_tree_after_seal_tx( &transaction, @@ -264,6 +266,16 @@ async fn seal_explicit_children( let staged_for_tx = staged.clone(); with_connection(config, move |connection| { let transaction = connection.unchecked_transaction()?; + let status: String = transaction.query_row( + "SELECT status FROM mem_tree_trees WHERE id = ?1", + [&node_for_tx.tree_id], + |row| row.get(0), + )?; + anyhow::ensure!( + status == "active", + "tree '{}' is archived", + node_for_tx.tree_id + ); store::insert_staged_summary_tx( &transaction, &node_for_tx, diff --git a/src/memory/tree/runtime/store/paths.rs b/src/memory/tree/runtime/store/paths.rs index 05eeaaa..cc34971 100644 --- a/src/memory/tree/runtime/store/paths.rs +++ b/src/memory/tree/runtime/store/paths.rs @@ -30,17 +30,12 @@ pub fn node_file_path(config: &MemoryConfig, namespace: &str, node_id: &str) -> /// namespace so distinct machine ids cannot alias the same directory. fn sanitize(namespace: &str) -> String { let raw = namespace.trim(); - let sanitized = raw.replace(['/', '\\', ':', '*', '?', '"', '<', '>', '|', '.'], "_"); - if sanitized == raw { - return sanitized; - } - use sha2::{Digest, Sha256}; - let digest = Sha256::digest(raw.as_bytes()); - let suffix = digest[..6] - .iter() - .map(|byte| format!("{byte:02x}")) - .collect::(); - format!("{sanitized}-{suffix}") + crate::memory::fsutil::sanitize_component_with_digest(raw, |character| { + !matches!( + character, + '/' | '\\' | ':' | '*' | '?' | '"' | '<' | '>' | '|' | '.' + ) + }) } /// Validate a namespace string, erroring on empty / dangerous input. diff --git a/src/memory/types.rs b/src/memory/types.rs index 900db69..5943cae 100644 --- a/src/memory/types.rs +++ b/src/memory/types.rs @@ -152,6 +152,7 @@ impl std::str::FromStr for MemoryCategory { value if value.starts_with("custom:") && value.len() > "custom:".len() => { Ok(Self::Custom(value["custom:".len()..].to_string())) } + value if !value.is_empty() => Ok(Self::Custom(value.to_string())), _ => Err(format!("unknown memory category: {value}")), } } From e8ee978680deeeb1360ed0705a429c56c12e92b0 Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 09:50:31 +0000 Subject: [PATCH 6/7] fix: reconcile upstream feature surfaces --- src/memory/persona/mod.rs | 17 +++++++++++------ src/memory/persona/types.rs | 2 +- src/memory/providers/openrouter.rs | 12 ++++++++---- src/memory/providers/openrouter_tests.rs | 2 ++ src/memory/sources/types.rs | 4 ++-- src/memory/store/entity_index/transaction.rs | 2 +- src/memory/store/vectors/store.rs | 1 + src/memory/types_tests.rs | 5 ++++- 8 files changed, 30 insertions(+), 15 deletions(-) diff --git a/src/memory/persona/mod.rs b/src/memory/persona/mod.rs index 0668a2e..5d191db 100644 --- a/src/memory/persona/mod.rs +++ b/src/memory/persona/mod.rs @@ -3,13 +3,18 @@ //! memory layer** — personality, communication style, coding style, and //! tool/stack preferences compiled into a small, prompt-ready context pack. //! -//! The surface is organised around the canonical evidence model in [`types`]: -//! readers ([`readers`]) emit redacted [`PersonaEvidence`], the map step -//! ([`distill`]) turns batches of evidence into [`SessionDigest`]s via a +//! The surface is organised around the canonical evidence model in +//! [`types`](crate::memory::persona::types): readers +//! ([`readers`](crate::memory::persona::readers)) emit redacted +//! [`PersonaEvidence`](crate::memory::persona::types::PersonaEvidence), the map +//! step ([`distill`](crate::memory::persona::distill)) turns batches of evidence +//! into [`SessionDigest`](crate::memory::persona::types::SessionDigest)s via a //! [`ChatProvider`](crate::memory::score::extract::ChatProvider), the reduce -//! step folds digests into seven facet flavoured trees, and the [`compile`] -//! step assembles `persona/PERSONA.md`. [`state`] makes runs incremental and -//! resumable; [`pipeline`] wires it all together for the CLI harness. +//! step folds digests into seven facet flavoured trees, and the +//! [`compile`](crate::memory::persona::compile) step assembles +//! `persona/PERSONA.md`. [`state`](crate::memory::persona::state) makes runs +//! incremental and resumable; [`pipeline`](crate::memory::persona::pipeline) +//! wires it all together for the CLI harness. //! //! Everything is local-first and depends only on the crate's `ChatProvider` / //! `Summariser` / `EmbeddingBackend` trait seams — nothing here names a diff --git a/src/memory/persona/types.rs b/src/memory/persona/types.rs index 9bdfeae..16055bb 100644 --- a/src/memory/persona/types.rs +++ b/src/memory/persona/types.rs @@ -10,7 +10,7 @@ //! //! [`PersonaEvidence`] can only be constructed through [`PersonaEvidence::new`], //! which runs the raw excerpt through -//! [`sanitize_text`](crate::memory::store::safety::sanitize_text) *before* the +//! [`sanitize_text`] *before* the //! value is stored on the struct. `sanitize_text` is the composite redactor: it //! scrubs secrets/tokens/keys (OpenAI `sk-…`, GitHub `gh*_…`, OAuth/bearer //! credentials) *and* runs the formatted-PII pass diff --git a/src/memory/providers/openrouter.rs b/src/memory/providers/openrouter.rs index d793601..9ec8a0c 100644 --- a/src/memory/providers/openrouter.rs +++ b/src/memory/providers/openrouter.rs @@ -4,13 +4,17 @@ //! not a hard dependency. [`OpenRouterProvider`] implements all three provider //! seams against OpenRouter's OpenAI-compatible endpoints: //! -//! - [`ChatProvider`] — `POST /chat/completions` in JSON mode (persona digests). -//! - [`Summariser`] — the same endpoint in plain-text mode (flavoured-tree folds). -//! - [`EmbeddingBackend`] — `POST /embeddings` (vector retrieval). +//! - [`ChatProvider`](crate::memory::score::extract::ChatProvider) — +//! `POST /chat/completions` in JSON mode (persona digests). +//! - [`Summariser`](crate::memory::tree::Summariser) — the same endpoint in +//! plain-text mode (flavoured-tree folds). +//! - [`EmbeddingBackend`](crate::memory::store::vectors::EmbeddingBackend) — +//! `POST /embeddings` (vector retrieval). //! //! Nothing under `memory::persona` names this type: the pipeline depends only on //! the traits, and OpenHuman injects its own routes. Secrets are held as -//! [`SecretString`]; token usage is accumulated per-run and the provider aborts +//! [`SecretString`](crate::memory::config::SecretString); token usage is +//! accumulated per-run and the provider aborts //! cleanly (returns a non-retryable error) once a configured cost/call budget is //! hit, mirroring the [`DailyBudget`](crate::memory::sync::state::DailyBudget) //! pattern. Transport and `429`/`5xx` failures retry with backoff; `4xx` client diff --git a/src/memory/providers/openrouter_tests.rs b/src/memory/providers/openrouter_tests.rs index fc7b45a..7f13b0e 100644 --- a/src/memory/providers/openrouter_tests.rs +++ b/src/memory/providers/openrouter_tests.rs @@ -183,6 +183,8 @@ async fn summariser_folds_via_chat_and_reports_usage() { tree_kind: TreeKind::Flavoured, target_level: 1, token_budget: 200, + input_token_budget: 4_096, + overhead_reserve_tokens: 256, ask: Some("Distill workflow habits."), }; let call = Summariser::summarise_with_usage(&provider, &inputs, &ctx) diff --git a/src/memory/sources/types.rs b/src/memory/sources/types.rs index 9519190..18b4da5 100644 --- a/src/memory/sources/types.rs +++ b/src/memory/sources/types.rs @@ -257,13 +257,13 @@ impl MemorySourcePatch { if self.query.is_some() && kind != SourceKind::TwitterQuery { return reject("query"); } - if self.since_days.is_some() && kind != SourceKind::TwitterQuery { + if matches!(self.since_days, Some(Some(_))) && kind != SourceKind::TwitterQuery { return reject("since_days"); } if self.selector.is_some() && kind != SourceKind::WebPage { return reject("selector"); } - if self.max_items.is_some() && kind != SourceKind::RssFeed { + if matches!(self.max_items, Some(Some(_))) && kind != SourceKind::RssFeed { return reject("max_items"); } if self.url.is_some() diff --git a/src/memory/store/entity_index/transaction.rs b/src/memory/store/entity_index/transaction.rs index 3e7ffb8..78f0de7 100644 --- a/src/memory/store/entity_index/transaction.rs +++ b/src/memory/store/entity_index/transaction.rs @@ -11,7 +11,7 @@ use super::types::CanonicalEntity; /// /// New rows are inserted with `is_user = false` because [`NoSelfIdentity`] /// performs no identity classification. On conflict, -/// [`UPSERT_PRESERVE_USER_SQL`] deliberately preserves the existing `is_user` +/// `UPSERT_PRESERVE_USER_SQL` deliberately preserves the existing `is_user` /// value rather than resetting a prior identity match. pub fn index_entities_tx( tx: &Transaction<'_>, diff --git a/src/memory/store/vectors/store.rs b/src/memory/store/vectors/store.rs index db55065..8500461 100644 --- a/src/memory/store/vectors/store.rs +++ b/src/memory/store/vectors/store.rs @@ -426,6 +426,7 @@ pub fn vec_to_bytes(v: &[f32]) -> Vec { /// /// Returns an error when the blob length is not divisible by four rather than /// silently truncating corrupt trailing bytes. +#[allow(clippy::manual_is_multiple_of)] // Keep compatibility below Rust 1.87. pub fn bytes_to_vec(bytes: &[u8]) -> anyhow::Result> { anyhow::ensure!( bytes.len() % 4 == 0, diff --git a/src/memory/types_tests.rs b/src/memory/types_tests.rs index 1fe1a52..5ee61b5 100644 --- a/src/memory/types_tests.rs +++ b/src/memory/types_tests.rs @@ -54,7 +54,10 @@ fn memory_category_serde_uses_snake_case() { category ); } - assert!("project_notes".parse::().is_err()); + assert_eq!( + "project_notes".parse::().unwrap(), + MemoryCategory::Custom("project_notes".into()) + ); } #[test] From daf9a0ebf392603aee022cc6083ad8bf8bacf39c Mon Sep 17 00:00:00 2001 From: Steven Enamakel Date: Tue, 14 Jul 2026 10:08:48 +0000 Subject: [PATCH 7/7] fix: address final review edge cases --- src/memory/chunks/mod.rs | 1 + src/memory/chunks/store_delete.rs | 37 ++++++++++++++++++++++++++ src/memory/ingest/pipeline.rs | 5 ++++ src/memory/ingest/pipeline_tests.rs | 14 ++++++++++ src/memory/retrieval/fast_tests.rs | 14 +++++++--- src/memory/retrieval/source.rs | 14 ++++++---- src/memory/store/memory_trait.rs | 3 +++ src/memory/store/memory_trait_tests.rs | 30 +++++++++++++++++++++ 8 files changed, 110 insertions(+), 8 deletions(-) diff --git a/src/memory/chunks/mod.rs b/src/memory/chunks/mod.rs index 1704079..167e4d9 100644 --- a/src/memory/chunks/mod.rs +++ b/src/memory/chunks/mod.rs @@ -108,6 +108,7 @@ pub use store::{ CHUNK_STATUS_BUFFERED, CHUNK_STATUS_DROPPED, CHUNK_STATUS_PENDING_EXTRACTION, CHUNK_STATUS_SEALED, RAW_FILE_GATE_KIND, }; +pub(crate) use store_delete::remove_unreferenced_content_files; pub use store_delete::{ delete_chunks_by_owner, delete_chunks_by_source, delete_chunks_by_source_prefix, delete_orphaned_source_tree, diff --git a/src/memory/chunks/store_delete.rs b/src/memory/chunks/store_delete.rs index 450928f..d89ed74 100644 --- a/src/memory/chunks/store_delete.rs +++ b/src/memory/chunks/store_delete.rs @@ -401,3 +401,40 @@ pub(super) fn remove_chunk_content_files(config: &MemoryConfig, content_paths: & let _ = std::fs::remove_file(&path); } } + +/// Remove staged chunk files that have no committed chunk or summary pointer. +/// +/// Used after a document-ingest gate race: the losing writer may have staged +/// bodies before discovering the committed winner. Paths referenced by the +/// winner are preserved. +pub(crate) fn remove_unreferenced_content_files( + config: &MemoryConfig, + content_paths: &[String], +) -> Result<()> { + let mut unique = content_paths + .iter() + .filter(|path| !path.is_empty()) + .cloned() + .collect::>(); + if unique.is_empty() { + return Ok(()); + } + with_connection(config, |connection| { + unique.retain(|path| { + connection + .query_row( + "SELECT NOT EXISTS ( + SELECT 1 FROM mem_tree_chunks WHERE content_path = ?1 + UNION ALL + SELECT 1 FROM mem_tree_summaries WHERE content_path = ?1 + )", + [path], + |row| row.get::<_, bool>(0), + ) + .unwrap_or(false) + }); + Ok(()) + })?; + remove_chunk_content_files(config, &unique.into_iter().collect::>()); + Ok(()) +} diff --git a/src/memory/ingest/pipeline.rs b/src/memory/ingest/pipeline.rs index 23a8ebc..7c58386 100644 --- a/src/memory/ingest/pipeline.rs +++ b/src/memory/ingest/pipeline.rs @@ -162,6 +162,11 @@ async fn persist_score_enqueue( })?; let Some((chunks_written, extract_jobs_enqueued)) = persisted else { + let staged_paths = staged + .iter() + .map(|chunk| chunk.content_path.clone()) + .collect::>(); + chunks::remove_unreferenced_content_files(config, &staged_paths)?; return Ok(IngestSummary::already_ingested(source_id)); }; diff --git a/src/memory/ingest/pipeline_tests.rs b/src/memory/ingest/pipeline_tests.rs index 10a9d91..bdbb757 100644 --- a/src/memory/ingest/pipeline_tests.rs +++ b/src/memory/ingest/pipeline_tests.rs @@ -199,6 +199,11 @@ async fn second_document_ingest_with_same_source_id_is_short_circuited() { .unwrap(); assert!(!first.already_ingested); assert!(first.chunks_written >= 1); + let staged_files_before = walkdir::WalkDir::new(crate::memory::chunks::content_root(&cfg)) + .into_iter() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_file()) + .count(); // Even with completely different content under the same source_id the second // ingest must write nothing: documents are append-only and source_id is the @@ -223,6 +228,15 @@ async fn second_document_ingest_with_same_source_id_is_short_circuited() { assert!(second.chunk_ids.is_empty()); assert_eq!(count_chunks(&cfg).unwrap(), first.chunks_written as u64); + let staged_files_after = walkdir::WalkDir::new(crate::memory::chunks::content_root(&cfg)) + .into_iter() + .filter_map(Result::ok) + .filter(|entry| entry.file_type().is_file()) + .count(); + assert_eq!( + staged_files_after, staged_files_before, + "a losing gate attempt must reclaim its unreferenced staged bodies" + ); } #[tokio::test] diff --git a/src/memory/retrieval/fast_tests.rs b/src/memory/retrieval/fast_tests.rs index 11a6406..2b38a43 100644 --- a/src/memory/retrieval/fast_tests.rs +++ b/src/memory/retrieval/fast_tests.rs @@ -35,13 +35,21 @@ async fn blank_query_is_empty_without_opening_storage() { #[tokio::test] async fn dense_fallback_filters_scope_before_limit() { - let (_temp, config) = test_config(); - for (id, scope) in [("allowed", "slack:#allowed"), ("denied", "slack:#denied")] { + let (_temp, mut config) = test_config(); + config.retrieval.limits.max_limit = 1; + for (id, scope, timestamp) in [ + ("allowed", "slack:#allowed", fixed_ts()), + ( + "denied", + "slack:#denied", + fixed_ts() + chrono::Duration::seconds(1), + ), + ] { let tree_id = format!("tree-{id}"); insert_tree_row(&config, &source_tree(&tree_id, scope, Some(id), 1)); insert_summary( &config, - &summary_node(id, &tree_id, 1, None, &[], id, fixed_ts()), + &summary_node(id, &tree_id, 1, None, &[], id, timestamp), ); } let scope = HashSet::from(["slack:#allowed".to_string()]); diff --git a/src/memory/retrieval/source.rs b/src/memory/retrieval/source.rs index 24fb841..fffeb78 100644 --- a/src/memory/retrieval/source.rs +++ b/src/memory/retrieval/source.rs @@ -41,10 +41,12 @@ pub(crate) type ScoredHit = (RetrievalHit, Option>); /// Retrieve summary hits from the selected source trees. /// -/// `limit` defaults to 10 when 0. When `query` is `Some`, the (inert-in-tests) -/// `embedder` is used to semantically rerank; otherwise ordering is -/// newest-first. See the module-level NOTE for the cost profile of -/// `time_window_days`. +/// `limit` defaults to 10 when 0 and is capped to the configured public limit. +/// The internal `usize::MAX` sentinel bypasses that cap so callers that apply a +/// narrower scope afterward can collect the complete candidate set first. +/// When `query` is `Some`, the (inert-in-tests) `embedder` is used to +/// semantically rerank; otherwise ordering is newest-first. See the +/// module-level NOTE for the cost profile of `time_window_days`. pub async fn query_source( config: &MemoryConfig, source_id: Option<&str>, @@ -55,7 +57,9 @@ pub async fn query_source( limit: usize, ) -> Result { let limits = &config.retrieval.limits; - let limit = if limit == 0 { + let limit = if limit == usize::MAX { + usize::MAX + } else if limit == 0 { limits.default_limit } else { limit.min(limits.max_limit) diff --git a/src/memory/store/memory_trait.rs b/src/memory/store/memory_trait.rs index c8b66a8..6898b37 100644 --- a/src/memory/store/memory_trait.rs +++ b/src/memory/store/memory_trait.rs @@ -125,6 +125,9 @@ impl Memory for InMemoryMemoryStore { limit: usize, opts: RecallOpts<'_>, ) -> Result> { + if query.trim().is_empty() { + return Ok(Vec::new()); + } let namespace = opts.namespace.unwrap_or(GLOBAL_NAMESPACE); let records = self .records diff --git a/src/memory/store/memory_trait_tests.rs b/src/memory/store/memory_trait_tests.rs index d244ece..a90717d 100644 --- a/src/memory/store/memory_trait_tests.rs +++ b/src/memory/store/memory_trait_tests.rs @@ -44,6 +44,36 @@ async fn headline_memory_contract_supports_upsert_recall_and_taint() { assert_eq!(hits[0].score, Some(1.0)); } +#[tokio::test] +async fn blank_recall_query_returns_no_unrelated_memories() { + let store = InMemoryMemoryStore::new(); + store + .store( + "agent", + "secret", + "unrelated memory", + MemoryCategory::Core, + None, + ) + .await + .unwrap(); + + for query in ["", " \n\t"] { + assert!(store + .recall( + query, + 10, + RecallOpts { + namespace: Some("agent"), + ..Default::default() + }, + ) + .await + .unwrap() + .is_empty()); + } +} + #[tokio::test] async fn memory_contract_filters_lists_summarises_and_forgets() { let store = InMemoryMemoryStore::new();