fix(memory): serialize memory_tree schema init across worker pool (#2059)

Co-authored-by: Steven Enamakel <enamakel@tinyhumans.ai>

Cargo.toml

@@ -17,6 +17,10 @@ path = "src/bin/slack_backfill.rs"
 name = "gmail-backfill-3d"
 path = "src/bin/gmail_backfill_3d.rs"
 
+[[bin]]
+name = "memory-tree-init-smoke"
+path = "src/bin/memory_tree_init_smoke.rs"
+
 [[bin]]
 name = "inference-probe"
 path = "src/bin/inference_probe.rs"

src/bin/memory_tree_init_smoke.rs

@@ -0,0 +1,110 @@
+//! Manual stress smoke for the memory_tree schema-init race fix.
+//!
+//! Spins N concurrent threads racing into `memory::tree::store::with_connection`
+//! against a shared workspace. Pre-fix (without the mutex-gated init guard),
+//! cold-start runs would surface SQLite codes 14 (CANTOPEN), 1546
+//! (IOERR_TRUNCATE), or 4874 (IOERR_SHMMAP) on some threads. Post-fix,
+//! all N threads must return Ok.
+//!
+//! # Usage
+//!
+//! ```sh
+//! # Fresh workspace (forces cold-start path)
+//! rm -rf /tmp/mt-smoke
+//! OPENHUMAN_WORKSPACE=/tmp/mt-smoke \
+//!   cargo run --bin memory-tree-init-smoke -- 32
+//!
+//! # Re-run against warm DB (should also be Ok; exercises fast path)
+//! OPENHUMAN_WORKSPACE=/tmp/mt-smoke \
+//!   cargo run --bin memory-tree-init-smoke -- 32
+//! ```
+//!
+//! Arg is thread count (default 16, must be > 0). Higher = more contention.
+//! Use `RUST_LOG=debug` to see per-worker results.
+//!
+//! Exit code: 0 if all threads Ok, 1 if any failed.
+
+use std::path::PathBuf;
+use std::process::ExitCode;
+use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::Arc;
+
+use openhuman_core::openhuman::config::Config;
+use openhuman_core::openhuman::memory::tree::store::with_connection;
+
+fn main() -> ExitCode {
+    env_logger::Builder::from_env(env_logger::Env::default().default_filter_or("info"))
+        .try_init()
+        .ok();
+
+    let workspace = match std::env::var("OPENHUMAN_WORKSPACE") {
+        Ok(v) => PathBuf::from(v),
+        Err(_) => {
+            log::error!("OPENHUMAN_WORKSPACE must be set to a writable directory");
+            return ExitCode::from(2);
+        }
+    };
+    let raw = std::env::args().nth(1).unwrap_or_else(|| "16".into());
+    let n: usize = match raw.parse() {
+        Ok(v) if v > 0 => v,
+        Ok(_) => {
+            log::error!("thread count must be a positive integer (> 0), got {raw}");
+            return ExitCode::from(2);
+        }
+        Err(e) => {
+            log::error!("thread count must be a positive integer, got {raw:?}: {e}");
+            return ExitCode::from(2);
+        }
+    };
+
+    let mut cfg = Config::default();
+    cfg.workspace_dir = workspace.clone();
+
+    let db_path = workspace.join("memory_tree").join("chunks.db");
+    let cold = !db_path.exists();
+    log::info!(
+        "[smoke] workspace={} cold_start={} threads={}",
+        workspace.display(),
+        cold,
+        n
+    );
+
+    let errors = Arc::new(AtomicUsize::new(0));
+    let start = std::time::Instant::now();
+
+    let threads: Vec<_> = (0..n)
+        .map(|i| {
+            let cfg = cfg.clone();
+            let errors = errors.clone();
+            std::thread::spawn(move || match with_connection(&cfg, |_| Ok(())) {
+                Ok(_) => {
+                    log::debug!("worker {i:3} ok");
+                }
+                Err(e) => {
+                    errors.fetch_add(1, Ordering::Relaxed);
+                    log::error!("worker {i:3} FAILED: {e:#}");
+                }
+            })
+        })
+        .collect();
+
+    for t in threads {
+        t.join().expect("worker thread panicked");
+    }
+
+    let failed = errors.load(Ordering::Relaxed);
+    let elapsed = start.elapsed();
+    log::info!(
+        "[smoke] done in {:?} — {}/{} ok, {} failed",
+        elapsed,
+        n - failed,
+        n,
+        failed
+    );
+
+    if failed > 0 {
+        ExitCode::from(1)
+    } else {
+        ExitCode::SUCCESS
+    }
+}

src/openhuman/memory/tree/store.rs

@@ -26,8 +26,10 @@ use chrono::{DateTime, TimeZone, Utc};
 use parking_lot::Mutex as PMutex;
 use rusqlite::{params, Connection, OptionalExtension, Transaction};
 use std::collections::HashMap;
-use std::path::PathBuf;
+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};
 
@@ -62,14 +64,19 @@ pub const CHUNK_STATUS_SEALED: &str = "sealed";
 /// Chunk lifecycle: rejected by the admission gate (too low signal).
 pub const CHUNK_STATUS_DROPPED: &str = "dropped";
 
+// `PRAGMA foreign_keys = ON` is intentionally NOT in SCHEMA — it is
+// a connection-local pragma that resets to off on every new
+// `Connection::open`. SCHEMA only runs once per DB path (first-init);
+// applying foreign_keys here would leak FK-off into every later
+// `with_connection()` call that hits the fast path. The pragma is
+// set per-connection in `open_connection()` instead.
+
 /// `PRAGMA user_version` value once the one-shot legacy→sidecar embedding
 /// migration (#1574 §7) has run. `0` (fresh/legacy DB) triggers the copy on
 /// next open; `>= 1` skips it. Bump only for a new one-shot data migration.
 const TREE_EMBEDDING_MIGRATION_VERSION: i64 = 1;
 
 const SCHEMA: &str = "
-PRAGMA foreign_keys = ON;
-
 CREATE TABLE IF NOT EXISTS mem_tree_chunks (
     id                     TEXT PRIMARY KEY,
     source_kind            TEXT NOT NULL,
@@ -713,6 +720,76 @@ fn ms_to_utc(ms: i64) -> rusqlite::Result<DateTime<Utc>> {
     })
 }
 
+// ── Schema-apply instrumentation (test-only) ─────────────────────────────────
+//
+// Per-path counter of how many times `apply_schema` ran for each DB path,
+// gated behind `cfg(test)` so the production binary carries no overhead.
+// Used by the concurrent-init regression test to assert "exactly once per
+// path" across racing workers; it survives even when the connection cache
+// is cleared between tests because tests use distinct tempdirs.
+#[cfg(test)]
+static SCHEMA_APPLY_COUNTS: OnceLock<Mutex<HashMap<PathBuf, usize>>> = OnceLock::new();
+
+fn record_schema_apply(_path: &Path) {
+    #[cfg(test)]
+    {
+        let counts = SCHEMA_APPLY_COUNTS.get_or_init(|| Mutex::new(HashMap::new()));
+        let mut guard = counts
+            .lock()
+            .expect("memory_tree schema apply count mutex poisoned");
+        *guard.entry(_path.to_path_buf()).or_insert(0) += 1;
+    }
+}
+
+#[cfg(test)]
+#[doc(hidden)]
+pub(crate) fn schema_apply_count_for_path_for_tests(path: &Path) -> usize {
+    SCHEMA_APPLY_COUNTS
+        .get()
+        .and_then(|m| {
+            m.lock()
+                .ok()
+                .map(|guard| guard.get(path).copied().unwrap_or(0))
+        })
+        .unwrap_or(0)
+}
+
+/// SQLite extended result code `CANTOPEN` — surfaces when a cold-start
+/// caller races the lockfile/WAL creation done by another connection.
+const SQLITE_CANTOPEN: i32 = 14;
+/// SQLite extended result code `IOERR_TRUNCATE` — fires when the WAL is
+/// being truncated by another connection during bootstrap.
+const SQLITE_IOERR_TRUNCATE: i32 = 1546;
+/// SQLite extended result code `IOERR_SHMMAP` — fires when the shared
+/// memory file is resized by another connection during bootstrap.
+const SQLITE_IOERR_SHMMAP: i32 = 4874;
+
+/// True if `err` (or anything in its cause chain) is one of the three
+/// SQLite codes that fire during cold-start WAL/SHM bootstrap races:
+/// `CANTOPEN`, `IOERR_TRUNCATE`, `IOERR_SHMMAP`.
+pub(crate) fn is_transient_cold_start(err: &anyhow::Error) -> bool {
+    fn is_transient_sqlite(e: &(dyn std::error::Error + 'static)) -> bool {
+        if let Some(rusqlite::Error::SqliteFailure(ffi, _)) = e.downcast_ref::<rusqlite::Error>() {
+            return matches!(
+                ffi.extended_code,
+                SQLITE_CANTOPEN | SQLITE_IOERR_TRUNCATE | SQLITE_IOERR_SHMMAP
+            );
+        }
+        false
+    }
+    if is_transient_sqlite(err.root_cause()) {
+        return true;
+    }
+    let mut src: Option<&(dyn std::error::Error + 'static)> = Some(err.as_ref());
+    while let Some(cur) = src {
+        if is_transient_sqlite(cur) {
+            return true;
+        }
+        src = cur.source();
+    }
+    false
+}
+
 // ── Connection cache (#2206) ─────────────────────────────────────────────────
 
 /// How many consecutive init failures before the circuit breaker trips.
@@ -786,6 +863,13 @@ impl CircuitBreaker {
 struct ConnectionCache {
     connections: PMutex<HashMap<PathBuf, Arc<PMutex<Connection>>>>,
     breakers: PMutex<HashMap<PathBuf, Arc<CircuitBreaker>>>,
+    /// Per-path mutex held across the slow-path init so concurrent
+    /// workers racing into `with_connection` on a cold DB serialise on
+    /// the WAL+SHM bootstrap. Without this, N threads see "no cached
+    /// connection" simultaneously and all run `apply_schema`, which is
+    /// idempotent but reopens the cold-start race window
+    /// (OPENHUMAN-TAURI-HH / -ZM / -MB).
+    init_locks: PMutex<HashMap<PathBuf, Arc<PMutex<()>>>>,
 }
 
 static CONN_CACHE: OnceLock<ConnectionCache> = OnceLock::new();
@@ -794,6 +878,7 @@ fn conn_cache() -> &'static ConnectionCache {
     CONN_CACHE.get_or_init(|| ConnectionCache {
         connections: PMutex::new(HashMap::new()),
         breakers: PMutex::new(HashMap::new()),
+        init_locks: PMutex::new(HashMap::new()),
     })
 }
 
@@ -845,6 +930,22 @@ pub(crate) fn try_cleanup_stale_files(db_path: &std::path::Path) -> bool {
 fn init_db(conn: &Connection, config: &Config) -> Result<()> {
     conn.busy_timeout(SQLITE_BUSY_TIMEOUT)
         .context("Failed to configure memory_tree busy timeout")?;
+    // SQLite resets `foreign_keys` to off on every new connection. The
+    // ConnectionCache holds one cached `Connection` per DB path, so
+    // setting it here (alongside the rest of init) is the per-connection
+    // surface — fast-path callers reuse the cached conn with FKs already
+    // on.
+    conn.execute_batch("PRAGMA foreign_keys = ON;")
+        .context("Failed to enable memory_tree foreign_keys pragma")?;
+    apply_schema(conn)?;
+    // #1574 §7: one-shot, version-gated legacy→sidecar embedding migration.
+    migrate_legacy_embeddings_to_sidecar(conn, config)?;
+    Ok(())
+}
+
+fn apply_schema(conn: &Connection) -> Result<()> {
+    // Note: `init_db` runs the `#1574 §7` legacy→sidecar embedding migration
+    // after this returns, so the dim-equal copy step is not duplicated here.
     if let Err(wal_err) = conn.execute_batch("PRAGMA journal_mode=WAL;") {
         log::warn!(
             "[memory_tree] Failed to enable WAL mode (filesystem may not support it): {wal_err}"
@@ -891,8 +992,6 @@ fn init_db(conn: &Connection, config: &Config) -> Result<()> {
         "is_user",
         "INTEGER NOT NULL DEFAULT 0",
     )?;
-    // #1574 §7: one-shot, version-gated legacy→sidecar embedding migration.
-    migrate_legacy_embeddings_to_sidecar(conn, config)?;
     Ok(())
 }
 
@@ -936,7 +1035,26 @@ fn get_or_init_connection(config: &Config) -> Result<Arc<PMutex<Connection>>> {
         }
     }
 
-    // ── Slow path: initialise DB, then insert into cache ─────────────────
+    // ── Slow path: serialise init per-path so concurrent workers don't
+    //    all race into `open_and_init` on a cold DB.
+    let init_lock = {
+        let mut guard = conn_cache().init_locks.lock();
+        guard
+            .entry(db_path.clone())
+            .or_insert_with(|| Arc::new(PMutex::new(())))
+            .clone()
+    };
+    let _init_guard = init_lock.lock();
+
+    // Re-check the cache once we hold the init lock — another thread
+    // may have completed init while we were queued.
+    {
+        let guard = conn_cache().connections.lock();
+        if let Some(conn) = guard.get(&db_path) {
+            return Ok(Arc::clone(conn));
+        }
+    }
+
     log::debug!(
         "[memory_tree] opening and initialising DB at {}",
         db_path.display()
@@ -1017,6 +1135,7 @@ fn open_and_init(db_path: &std::path::Path, config: &Config) -> Result<Connectio
         .with_context(|| format!("Failed to open memory_tree DB: {}", db_path.display()))?;
     init_db(&conn, config)
         .with_context(|| format!("Failed to init memory_tree schema: {}", db_path.display()))?;
+    record_schema_apply(db_path);
     Ok(conn)
 }
 
@@ -1039,6 +1158,7 @@ pub(crate) fn invalidate_connection(config: &Config) {
 pub(crate) fn clear_connection_cache() {
     conn_cache().connections.lock().clear();
     conn_cache().breakers.lock().clear();
+    conn_cache().init_locks.lock().clear();
 }
 
 /// Open the memory_tree SQLite DB and run a closure against it.
@@ -1052,10 +1172,13 @@ pub(crate) fn clear_connection_cache() {
 /// reused from a process-level cache. Schema migrations run exactly once on
 /// the first call for a given `config.workspace_dir`. Subsequent calls pay
 /// only the cost of a `parking_lot::Mutex` lock and the closure itself.
-pub(crate) fn with_connection<T>(
-    config: &Config,
-    f: impl FnOnce(&Connection) -> Result<T>,
-) -> Result<T> {
+///
+/// `#[doc(hidden)] pub` (not `pub(crate)`) because the
+/// `memory-tree-init-smoke` bin in `src/bin/` is a separate crate target
+/// and must reach this entry point. It is NOT a stable API surface —
+/// downstream crates should treat it as internal.
+#[doc(hidden)]
+pub fn with_connection<T>(config: &Config, f: impl FnOnce(&Connection) -> Result<T>) -> Result<T> {
     let conn_arc = get_or_init_connection(config)?;
     let guard = conn_arc.lock();
     f(&guard)

src/openhuman/memory/tree/store_tests.rs

@@ -265,6 +265,136 @@ fn schema_has_content_path_and_content_sha256_columns() {
     .unwrap();
 }
 
+/// Regression: OPENHUMAN-TAURI-HH / -ZM / -MB.
+///
+/// Before this fix, N `tree_jobs_worker` tasks racing into
+/// `with_connection` on a cold workspace would trigger one of three
+/// SQLite cold-start codes — 14 (CANTOPEN), 1546 (IOERR_TRUNCATE),
+/// or 4874 (IOERR_SHMMAP) — surfaced as
+/// `Failed to initialize memory_tree schema`. The mutex-gated init set
+/// in `store::open_and_init_with_retry` serialises the WAL+SHM
+/// bootstrap so only one thread runs `apply_schema` per DB path.
+///
+/// Asserts:
+/// 1. All N concurrent callers return `Ok` (no races, no surfaced errors).
+/// 2. `apply_schema` runs exactly once for the shared path even though
+///    8 threads hit a cold DB simultaneously.
+#[test]
+fn with_connection_serialises_concurrent_schema_init() {
+    use std::sync::atomic::Ordering;
+
+    let (_tmp, cfg) = test_config();
+    let db_path = cfg.workspace_dir.join("memory_tree").join("chunks.db");
+    let errors = std::sync::Arc::new(std::sync::atomic::AtomicUsize::new(0));
+
+    let threads: Vec<_> = (0..8)
+        .map(|_| {
+            let cfg = cfg.clone();
+            let errors = errors.clone();
+            std::thread::spawn(move || {
+                if with_connection(&cfg, |_| Ok(())).is_err() {
+                    errors.fetch_add(1, Ordering::Relaxed);
+                }
+            })
+        })
+        .collect();
+    for t in threads {
+        t.join().expect("worker thread panicked");
+    }
+
+    assert_eq!(
+        errors.load(Ordering::Relaxed),
+        0,
+        "concurrent with_connection callers must all succeed"
+    );
+    let applied = super::schema_apply_count_for_path_for_tests(&db_path);
+    assert_eq!(
+        applied, 1,
+        "apply_schema must run exactly once per DB path under concurrent init; ran {applied} times"
+    );
+}
+
+/// Directly pins the `is_transient_cold_start` classifier — the
+/// gatekeeper for the retry loop in `open_and_init_with_retry`. The
+/// concurrent-init test above only exercises it indirectly (and only
+/// if a transient happens to fire on the dev box). A targeted test
+/// catches regressions if the match arms are edited.
+#[test]
+fn is_transient_cold_start_classifies_known_extended_codes() {
+    use rusqlite::ffi;
+    use rusqlite::ErrorCode;
+
+    // The three SHMmap/WAL bootstrap codes that fire under cold-start
+    // contention. All must classify as transient → retried.
+    for extended in [
+        14,   // CANTOPEN
+        1546, // IOERR_TRUNCATE
+        4874, // IOERR_SHMMAP
+    ] {
+        let err = anyhow::Error::from(rusqlite::Error::SqliteFailure(
+            ffi::Error {
+                code: ErrorCode::SystemIoFailure,
+                extended_code: extended,
+            },
+            None,
+        ));
+        assert!(
+            super::is_transient_cold_start(&err),
+            "extended_code {extended} must classify as transient cold-start"
+        );
+    }
+
+    // SQLITE_BUSY (extended code 5) is a real lock-contention signal,
+    // NOT a cold-start race — the caller handles it via `busy_timeout`
+    // not via this retry loop. Must NOT classify.
+    let busy = anyhow::Error::from(rusqlite::Error::SqliteFailure(
+        ffi::Error {
+            code: ErrorCode::DatabaseBusy,
+            extended_code: 5,
+        },
+        None,
+    ));
+    assert!(
+        !super::is_transient_cold_start(&busy),
+        "DatabaseBusy must not be classified as cold-start transient"
+    );
+
+    // Non-SQLite error in the chain — must not classify.
+    let other: anyhow::Error = anyhow::anyhow!("not a sqlite error");
+    assert!(
+        !super::is_transient_cold_start(&other),
+        "non-SQLite errors must not classify as transient cold-start"
+    );
+}
+
+/// Regression: `PRAGMA foreign_keys` is connection-local in SQLite and
+/// must be re-set on every `Connection::open`. After the schema-init
+/// refactor, the pragma moved out of `SCHEMA` (which only runs on
+/// first init per path) into `open_connection`. Verify both the
+/// cold-init path and the fast path return a connection with FK on.
+#[test]
+fn with_connection_keeps_foreign_keys_on_for_every_call() {
+    let (_tmp, cfg) = test_config();
+    // First call — exercises apply_schema + open_connection.
+    let fk_on_first: i64 = with_connection(&cfg, |conn| {
+        Ok(conn.query_row("PRAGMA foreign_keys;", params![], |r| r.get::<_, i64>(0))?)
+    })
+    .unwrap();
+    assert_eq!(
+        fk_on_first, 1,
+        "foreign_keys must be ON on first connection"
+    );
+    // Second call — fast path (schema init skipped); pragma must still be set.
+    let fk_on_second: i64 = with_connection(&cfg, |conn| {
+        Ok(conn.query_row("PRAGMA foreign_keys;", params![], |r| r.get::<_, i64>(0))?)
+    })
+    .unwrap();
+    assert_eq!(
+        fk_on_second, 1,
+        "foreign_keys must be ON on fast-path (post-init) connection"
+    );
+}
+
 /// #1574 §7: the one-shot, version-gated legacy→sidecar migration copies a
 /// legacy `embedding` blob whose dimensionality matches the active embedder
 /// into the per-model sidecar at the active signature, skips dim-mismatched

src/openhuman/wallet/execution.rs

@@ -1010,6 +1010,12 @@ mod tests {
 
     #[test]
     fn quote_store_round_trips_and_expires() {
+        // Must hold TEST_LOCK before clobbering the process-wide quote store,
+        // otherwise this races the async execute_prepared_* tests that store
+        // a quote and then await — `reset_quote_store_for_tests()` here can
+        // wipe their quote between store + await, surfacing as
+        // "quote 'q_retry' not found" in CI (intermittent).
+        let _guard = TEST_LOCK.lock();
         reset_quote_store_for_tests();
         let now = now_ms();
         let mut q = PreparedTransaction {