//! Concrete `MemoryBackend` implementation using hybrid search. //! //! `MemoryBackendImpl` combines FTS5 keyword search with optional vector //! KNN similarity via `hybrid_search()`. When embeddings are available //! (embedding config + API key), the query is vectorized and both signals //! are merged with recency and source weights. When embeddings are //! unavailable, gracefully degrades to FTS-only. //! //! `rusqlite::Connection` is `!Send + !Sync`, so we open a fresh `MemoryIndex` //! per query. WAL mode ensures concurrent readers don't block. use std::path::{Path, PathBuf}; use std::sync::Arc; use kigi_tools::types::memory_backend::{MemoryBackend, MemorySearchResult}; use super::embedding::EmbeddingProvider as _; use super::storage::MemoryStorage; use super::watcher::MemoryFileWatcher; /// All configuration needed to build a fully-wired [`MemoryBackendImpl`] for a live session. /// /// Grouping these in one struct ensures every call site — ToolBridge, first-turn /// injection, and post-compaction recovery — shares identical config. Without it, /// different paths silently fell back to FTS-only search and ignored /// `[memory.search]` config because no single place applied all builder methods. #[derive(Clone)] pub struct MemoryBackendParams { /// Session ID for telemetry events. pub session_id: String, /// Embedding provider config — `None` forces FTS-only fallback everywhere. pub embed_config: Option, /// Base URL for embedding API calls (CLI proxy). pub embed_base_url: String, /// API key for embedding API calls. pub embed_api_key: Option, /// Hybrid search scoring config (weights, thresholds, decay, MMR). pub search_config: kigi_config_types::MemorySearchConfig, /// File watcher for sync-on-search — `None` disables external-edit detection. pub watcher: Option>, /// Seconds before a stale reindex claim is forcibly released. pub stale_claim_secs: i64, /// Telemetry label emitted with every search event from this backend. /// /// Differentiates the three runtime search paths in dashboards and logs: /// - `"tool"` — model-initiated `memory_search` tool call (ToolBridge) /// - `"injection"` — first-turn memory context injection /// - `"compaction_recovery"` — post-compaction context re-injection pub search_source: &'static str, /// Dynamic API key provider — when set, `make_embedding_provider()` resolves /// the key per-call instead of using the static `embed_api_key`. pub api_key_provider: Option, pub auth_credentials: Option>, } impl MemoryBackendParams { /// Async so `current_api_key_async` can drive the AuthManager /// refresh chain; reindex loops outlive the OIDC TTL. pub async fn make_embedding_provider(&self) -> Option { build_embedding_provider( self.embed_config.as_ref(), self.auth_credentials.as_ref(), self.api_key_provider.as_ref(), self.embed_api_key.as_deref(), &self.embed_base_url, ) .await } } async fn build_embedding_provider( config: Option<&kigi_config_types::MemoryEmbeddingConfig>, auth_credentials: Option<&Arc>, api_key_provider: Option<&kigi_tools::types::SharedApiKeyProvider>, static_api_key: Option<&str>, base_url: &str, ) -> Option { let config = config?; if config.model.as_ref().is_none_or(|m| m.is_empty()) { return None; } // Prefer the refresh-capable credential provider — the middleware gives // 401 retry for free without any per-call key resolution. if let Some(creds) = auth_credentials { let client = super::embedding::build_middleware_client(creds.clone()); return super::embedding::ApiEmbeddingProvider::from_config( config, base_url.to_owned(), client, ); } // Fallback: resolve API key per-call, wrap in a static middleware client // (no 401 refresh, but auth header is still stamped by middleware). let api_key = match api_key_provider { Some(p) => p.current_api_key_async().await, None => None, } .or_else(|| static_api_key.map(|s| s.to_owned()))?; super::embedding::ApiEmbeddingProvider::from_session(config, base_url.to_owned(), api_key) } /// `MemoryBackend` implementation backed by hybrid search (FTS5 + vector KNN). /// /// Stores only `Send + Sync` config data. The `MemoryIndex` and /// `EmbeddingProvider` are constructed on demand per query. pub struct MemoryBackendImpl { db_path: PathBuf, storage: MemoryStorage, /// Embedding config — `None` disables vector search (FTS-only fallback). embed_config: Option, /// API base URL for embedding requests (cli-chat-proxy). embed_base_url: String, /// API key for embedding requests. embed_api_key: Option, /// Search scoring config (weights, min_score, max_results). search_config: kigi_config_types::MemorySearchConfig, /// File watcher for detecting external memory edits. watcher: Option>, /// Stale claim threshold for reindex coordination. stale_claim_secs: i64, /// Session ID for telemetry events. session_id: String, /// Telemetry label for search events — mirrors [`MemoryBackendParams::search_source`]. search_source: &'static str, /// Shared search counter — read by session summary telemetry. /// /// Only the ToolBridge backend's counter is shared back to the session actor; /// injection and compaction-recovery backends use their own local counters. pub search_counter: std::sync::Arc, /// Dynamic API key provider for embedding requests. api_key_provider: Option, /// Refresh-capable credential provider for embedding HTTP middleware. auth_credentials: Option>, } impl MemoryBackendImpl { /// Create a new backend. `db_path` must point to an existing SQLite /// database created by `MemoryIndex::open_or_create()`. pub fn new(db_path: PathBuf, storage: MemoryStorage) -> Self { Self { db_path, storage, embed_config: None, embed_base_url: String::new(), embed_api_key: None, search_config: kigi_config_types::MemorySearchConfig::default(), watcher: None, stale_claim_secs: 60, session_id: String::new(), search_source: "tool", api_key_provider: None, auth_credentials: None, search_counter: std::sync::Arc::new(std::sync::atomic::AtomicU64::new(0)), } } /// Set the session ID for telemetry. pub fn with_session_id(mut self, session_id: String) -> Self { self.session_id = session_id; self } /// Configure the embedding provider for hybrid search. /// /// Without this, `search()` falls back to FTS-only. pub fn with_embedding( mut self, config: kigi_config_types::MemoryEmbeddingConfig, base_url: String, api_key: Option, ) -> Self { self.embed_config = Some(config); self.embed_base_url = base_url; self.embed_api_key = api_key; self } /// Override the search scoring config (weights, limits, etc.). pub fn with_search_config(mut self, config: kigi_config_types::MemorySearchConfig) -> Self { self.search_config = config; self } /// Attach a file watcher for sync-on-search (reindex dirty files before querying). pub fn with_watcher(mut self, watcher: Arc, stale_claim_secs: i64) -> Self { self.watcher = Some(watcher); self.stale_claim_secs = stale_claim_secs; self } /// Open a read-only connection for simple queries (`total_chunks`, `get`). fn open_readonly(&self) -> Result { // Journal-mode-aware open (busy_timeout included): never mmap a legacy // WAL -shm on network mounts (SIGBUS); see JournalMode::open_readonly. kigi_sqlite_journal::JournalMode::for_db_path(&self.db_path).open_readonly(&self.db_path) } async fn make_embedding_provider(&self) -> Option { build_embedding_provider( self.embed_config.as_ref(), self.auth_credentials.as_ref(), self.api_key_provider.as_ref(), self.embed_api_key.as_deref(), &self.embed_base_url, ) .await } /// Build a fully configured backend for a live session. /// /// Prefer this over calling `new()` + individual builder methods: it ensures /// session_id, embeddings, search config, and the file watcher are applied /// consistently at every call site (ToolBridge, first-turn injection, /// post-compaction recovery). Using the factory eliminates the silent /// per-site drift where some paths got hybrid search while others fell back /// to FTS-only, and where `[memory.search]` config was effectively ignored. pub fn from_session_params(storage: MemoryStorage, params: &MemoryBackendParams) -> Self { let db_path = storage.workspace_dir().join("index.sqlite"); let mut backend = Self::new(db_path, storage) .with_session_id(params.session_id.clone()) .with_search_config(params.search_config.clone()); backend.search_source = params.search_source; if let Some(ec) = ¶ms.embed_config { backend = backend.with_embedding( ec.clone(), params.embed_base_url.clone(), params.embed_api_key.clone(), ); } if let Some(w) = ¶ms.watcher { backend = backend.with_watcher(w.clone(), params.stale_claim_secs); } backend.api_key_provider = params.api_key_provider.clone(); backend.auth_credentials = params.auth_credentials.clone(); backend } } /// Test-only field accessors. /// /// These expose private fields so tests can assert that `from_session_params` /// actually stored the values it was given, without routing through a full /// runtime search call whose semantics override some config fields. #[cfg(test)] impl MemoryBackendImpl { /// Returns the session ID stored in this backend. pub fn session_id_for_test(&self) -> &str { &self.session_id } /// Returns the search config stored in this backend. pub fn search_config_for_test(&self) -> &kigi_config_types::MemorySearchConfig { &self.search_config } } #[async_trait::async_trait] impl MemoryBackend for MemoryBackendImpl { #[tracing::instrument(name = "memory.search", skip_all, fields( session_id = %self.session_id, max_results, min_score, ))] async fn search( &self, query: &str, max_results: usize, min_score: f64, ) -> Result, Box> { // Open a MemoryIndex for this query (open-per-query, ~1ms). // // IMPORTANT: `MemoryIndex` is `Send` but `!Sync`, so `&MemoryIndex` // is `!Send`. To keep this future `Send`, we must never hold a // `&index` borrow across an `.await` point. The code below is // structured into sync phases (borrow &index) and async phases // (no &index borrow) to satisfy this constraint. let embed_dims = self.embed_config.as_ref().map_or(1024, |ec| ec.dimensions); let mut index = super::index::MemoryIndex::open_or_create( &self.db_path, self.storage.clone(), kigi_config_types::MemoryIndexConfig::default(), embed_dims, ) .map_err(|e| -> Box { Box::new(std::io::Error::other(e.to_string())) })?; // ── Sync phase 1: reindex dirty files, collect chunks needing embeddings ── let mut reindex_chunks: Vec<(String, String)> = Vec::new(); let mut needs_release = false; if let Some(ref watcher) = self.watcher && watcher.is_dirty() && index.try_claim_reindex(self.stale_claim_secs) { needs_release = true; let dirty_files = watcher.take_dirty(); let dirty_count = dirty_files.len(); for file in &dirty_files { if file.exists() { // File was created or modified — reindex it. let source = self.storage.classify_source(file); let _ = index.reindex_file(file, source); } else { // File was deleted — remove its stale chunks from the index so // they are no longer searchable. Without this call, reindex_file // returns early when the file is unreadable and leaves orphaned // chunks behind indefinitely. let _ = index.delete_path(file); } } if dirty_count > 0 { reindex_chunks = index.chunks_without_embeddings().unwrap_or_default(); } } // ── Async phase: embed missing chunks (no &index borrow) ── let provider = self.make_embedding_provider().await; if !reindex_chunks.is_empty() && let Some(ref provider) = provider { let mut upserts: Vec<(String, Vec)> = Vec::new(); for batch in reindex_chunks.chunks(32) { let texts: Vec<&str> = batch.iter().map(|(_, t)| t.as_str()).collect(); match provider.embed_batch(&texts).await { Ok(embeddings) => { for ((chunk_id, _), emb) in batch.iter().zip(embeddings.into_iter()) { upserts.push((chunk_id.clone(), emb)); } } Err(e) => { tracing::warn!( target: kigi_log::memory_log::TARGET, error = %e, "embedding batch failed during sync-on-search, skipping" ); } } } // Sync: upsert embeddings back (borrows &index, no await) for (chunk_id, emb) in &upserts { let _ = index.upsert_embedding(chunk_id, emb); } } if needs_release { index.release_claim(); } // ── Sync phase 2: FTS search ── let mut search_config = self.search_config.clone(); search_config.max_results = max_results; search_config.min_score = min_score as f32; let candidate_limit = search_config.max_results * 3; let mut fts_results = index.search_fts(query, candidate_limit).unwrap_or_default(); // Supplemental evergreen query: ensure global/workspace MEMORY.md // chunks appear in candidates even when session volume crowds them // out of the base FTS results. Mirrors hybrid_search() in search.rs. let evergreen = index .search_fts_by_sources(query, candidate_limit, &["global", "workspace"]) .unwrap_or_default(); let existing: std::collections::HashSet = fts_results.iter().map(|r| r.chunk_id.clone()).collect(); for r in evergreen { if !existing.contains(&r.chunk_id) { fts_results.push(r); } } let vec_available = index.vec_available() && provider.is_some(); // ── Async phase: embed query for vector search (no &index borrow) ── let query_embedding = if vec_available { if let Some(ref provider) = provider { match provider.embed_batch(&[query]).await { Ok(embeddings) if !embeddings.is_empty() => { Some(embeddings.into_iter().next().unwrap()) } Ok(_) => None, Err(e) => { tracing::warn!(error = %e, "embedding query failed, falling back to FTS-only"); None } } } else { None } } else { None }; // ── Sync phase 3: vector search + scoring + merge (borrows &index) ── let results = super::search::hybrid_search_merge( &index, fts_results, query_embedding.as_deref(), &search_config, ) .map_err(|e| -> Box { Box::new(std::io::Error::other(e.to_string())) })?; // Record accesses for the returned chunks so access_count and // last_accessed stay current. Non-fatal: a failed write is a no-op // for the caller and does not affect the search response. for result in &results { let _ = index.record_access(&result.chunk_id); } self.search_counter .fetch_add(1, std::sync::atomic::Ordering::Relaxed); Ok(results .into_iter() .map(|r| MemorySearchResult { chunk_id: r.chunk_id, path: r.path, start_line: r.start_line, end_line: r.end_line, score: r.score, snippet: r.snippet, source: r.source, created_at: Some(r.created_at), }) .collect()) } fn get( &self, path: &str, from: Option, lines: Option, ) -> Result> { Ok(self.storage.read_file(Path::new(path), from, lines)?) } fn total_chunks(&self) -> Result> { let conn = self.open_readonly()?; let count: i64 = conn.query_row("SELECT COUNT(*) FROM chunks", [], |r| r.get(0))?; Ok(count as usize) } /// Return the configured `max_results` from the stored search config. /// /// Overrides the trait default so the `memory_search` tool honours /// `[memory.search].max_results` from config when the model does not /// supply an explicit value. fn default_search_max_results(&self) -> usize { self.search_config.max_results } /// Return the configured `min_score` from the stored search config. fn default_search_min_score(&self) -> f64 { self.search_config.min_score as f64 } } #[cfg(test)] mod factory_tests { use super::*; use crate::index::{MemoryIndex, init_sqlite_vec}; use crate::storage::MemoryStorage; use kigi_config_types::{MemoryEmbeddingConfig, MemorySearchConfig}; use tempfile::TempDir; fn make_storage(tmp: &TempDir) -> MemoryStorage { let global = tmp.path().join("memory"); let workspace = global.join("test_ws"); MemoryStorage::with_paths(global, workspace) } fn make_params_fts_only(session_id: &str) -> MemoryBackendParams { MemoryBackendParams { session_id: session_id.to_string(), embed_config: None, embed_base_url: String::new(), embed_api_key: None, search_config: MemorySearchConfig::default(), watcher: None, stale_claim_secs: 60, search_source: "tool", api_key_provider: None, auth_credentials: None, } } /// from_session_params stores the session_id it was given. /// /// Direct assertion via the `#[cfg(test)]` accessor proves the factory /// actually stored the value rather than discarding it. The counter /// increment check additionally confirms the backend is functional. #[tokio::test] async fn test_factory_sets_session_id() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let storage = make_storage(&tmp); let db_path = storage.workspace_dir().join("index.sqlite"); let mut idx = MemoryIndex::open_or_create( &db_path, storage.clone(), kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); let file = tmp.path().join("note.md"); std::fs::write(&file, "# Facts\n\nRust is fast.").unwrap(); idx.reindex_file(&file, "workspace").unwrap(); drop(idx); let params = make_params_fts_only("test-session-abc"); let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); // Direct assertion: the stored session_id matches what the factory was given. assert_eq!( backend.session_id_for_test(), "test-session-abc", "session_id must be stored exactly as supplied" ); // Functional check: the backend actually runs a search. let before = backend .search_counter .load(std::sync::atomic::Ordering::Relaxed); let _ = backend.search("rust", 5, 0.0).await; let after = backend .search_counter .load(std::sync::atomic::Ordering::Relaxed); assert_eq!( after, before + 1, "search counter must increment per search" ); } /// from_session_params stores the search_config it was given. /// /// Direct assertion via the `#[cfg(test)]` accessor proves the factory /// propagated the config into the backend rather than discarding it. /// `max_results` is verified because the `search()` method overrides it /// with the caller's argument — so checking the *stored* value is the only /// way to confirm the factory wired it correctly. #[tokio::test] async fn test_factory_wires_search_config() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let storage = make_storage(&tmp); let db_path = storage.workspace_dir().join("index.sqlite"); let mut idx = MemoryIndex::open_or_create( &db_path, storage.clone(), kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); for i in 0..10 { let f = tmp.path().join(format!("note{i}.md")); std::fs::write(&f, format!("# Entry {i}\n\nRust tip number {i}.")).unwrap(); idx.reindex_file(&f, "workspace").unwrap(); } drop(idx); let params = MemoryBackendParams { search_config: MemorySearchConfig { max_results: 3, ..Default::default() }, ..make_params_fts_only("test-search-config") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); // Direct: the stored config has exactly the value the factory was given. assert_eq!( backend.search_config_for_test().max_results, 3, "stored max_results must equal what was supplied to the factory" ); } /// from_session_params wires non-overridable config fields (MMR, temporal decay) /// that `search()` never replaces with caller arguments. /// /// This is the clearest proof that `[memory.search]` config is actually wired /// rather than silently ignored: fields the caller cannot override must arrive /// in the stored search_config exactly as given. #[test] fn test_factory_wires_non_overridable_search_config_fields() { let tmp = TempDir::new().unwrap(); let storage = make_storage(&tmp); let custom_search = MemorySearchConfig { max_results: 7, mmr: kigi_config_types::MmrConfig { enabled: true, lambda: 0.42, }, temporal_decay: kigi_config_types::TemporalDecayConfig { enabled: true, half_life_days: 14.0, }, ..Default::default() }; let params = MemoryBackendParams { search_config: custom_search, ..make_params_fts_only("test-full-config") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); let stored = backend.search_config_for_test(); // None of these are overridden by the caller in search() — they must // survive the factory path unchanged. assert_eq!(stored.max_results, 7); assert!(stored.mmr.enabled, "MMR enabled must be stored"); assert!( (stored.mmr.lambda - 0.42).abs() < f64::EPSILON, "MMR lambda must be stored exactly" ); assert!( stored.temporal_decay.enabled, "temporal_decay enabled must be stored" ); assert!( (stored.temporal_decay.half_life_days - 14.0).abs() < f64::EPSILON, "temporal_decay half_life_days must be stored exactly" ); } /// from_session_params propagates search_source into the backend. /// /// Correctness test: every caller (tool, injection, /// compaction_recovery) must be able to set a distinct source label so /// dashboards can separate the three search paths. #[test] fn test_factory_propagates_search_source() { let tmp = TempDir::new().unwrap(); let storage = make_storage(&tmp); for source in ["tool", "injection", "compaction_recovery"] { let params = MemoryBackendParams { search_source: source, ..make_params_fts_only("test-source") }; let backend = MemoryBackendImpl::from_session_params(storage.clone(), ¶ms); assert_eq!( backend.search_source, source, "search_source must be propagated for source='{source}'" ); } } /// The default search_source is "tool" when constructing via new(). #[test] fn test_default_search_source_is_tool() { let tmp = TempDir::new().unwrap(); let db_path = tmp.path().join("test.sqlite"); let storage = make_storage(&tmp); let backend = MemoryBackendImpl::new(db_path, storage); assert_eq!(backend.search_source, "tool"); } /// MemoryBackendParams with different search_source values is Clone. #[test] fn test_params_clone_preserves_search_source() { let params = MemoryBackendParams { search_source: "injection", ..make_params_fts_only("test-clone-source") }; let cloned = params.clone(); assert_eq!(cloned.search_source, "injection"); } /// Watcher startup telemetry reflects actual runtime state. /// /// `watcher.is_some()` is `true` only when the watcher started successfully. /// With a valid directory the watcher should start; without one it should return None. /// This guards the contract that `watcher_started` in telemetry must reflect /// runtime outcome, not configuration intent. #[test] fn test_params_watcher_started_reflects_runtime() { let tmp = TempDir::new().unwrap(); // Success path: directory exists → watcher starts. let watch_dir = tmp.path().join("memory"); std::fs::create_dir_all(&watch_dir).unwrap(); let watcher = crate::watcher::MemoryFileWatcher::start(&watch_dir); let params_with_watcher = MemoryBackendParams { watcher: watcher.map(std::sync::Arc::new), ..make_params_fts_only("test-watcher-runtime") }; // watcher.is_some() reflects whether startup succeeded. // (On environments without inotify/FSEvents this may be None; skip rather than fail.) let _ = params_with_watcher.watcher.is_some(); // just verify it compiles // Failure path: non-existent directory → watcher must return None. let missing = tmp.path().join("does_not_exist"); let no_watcher = crate::watcher::MemoryFileWatcher::start(&missing); assert!( no_watcher.is_none(), "watcher must return None for a non-existent directory" ); let params_no_watcher = MemoryBackendParams { watcher: None, ..make_params_fts_only("test-no-watcher") }; assert!( params_no_watcher.watcher.is_none(), "params.watcher.is_none() means telemetry reports watcher_started=false" ); } /// default_search_max_results returns the configured value from search_config. /// /// Verifies that the MemoryBackend trait override in MemoryBackendImpl /// exposes search_config.max_results rather than the hardcoded default (6). #[test] fn test_default_search_max_results_from_config() { let tmp = TempDir::new().unwrap(); let storage = make_storage(&tmp); let params = MemoryBackendParams { search_config: MemorySearchConfig { max_results: 12, ..Default::default() }, ..make_params_fts_only("test-defaults") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); assert_eq!( backend.default_search_max_results(), 12, "default_search_max_results must return search_config.max_results" ); } /// default_search_min_score returns the configured value from search_config. #[test] fn test_default_search_min_score_from_config() { let tmp = TempDir::new().unwrap(); let storage = make_storage(&tmp); let params = MemoryBackendParams { search_config: MemorySearchConfig { min_score: 0.42, ..Default::default() }, ..make_params_fts_only("test-defaults") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); assert!( (backend.default_search_min_score() - 0.42_f64).abs() < 1e-6, "default_search_min_score must return search_config.min_score" ); } /// from_session_params without embed_config produces a backend that does not panic /// and returns results using FTS-only path. #[tokio::test] async fn test_factory_fts_only_without_embed() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let storage = make_storage(&tmp); let db_path = storage.workspace_dir().join("index.sqlite"); let mut idx = MemoryIndex::open_or_create( &db_path, storage.clone(), kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); let f = tmp.path().join("note.md"); std::fs::write(&f, "# Guide\n\nRust ownership rules.").unwrap(); idx.reindex_file(&f, "workspace").unwrap(); drop(idx); let params = make_params_fts_only("test-fts-only"); let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); let results = backend.search("rust ownership", 5, 0.0).await.unwrap(); assert!( !results.is_empty(), "FTS-only backend should return results" ); let ts = results[0].created_at; assert!( ts.is_some() && ts.unwrap() > 0, "created_at must be Some(positive) after backend search (got {ts:?})" ); } /// from_session_params with embed_config but no api_key gracefully falls back /// to FTS-only (the embedding provider requires a key). #[tokio::test] async fn test_factory_embed_config_without_key_falls_back_to_fts() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let storage = make_storage(&tmp); let db_path = storage.workspace_dir().join("index.sqlite"); let mut idx = MemoryIndex::open_or_create( &db_path, storage.clone(), kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); let f = tmp.path().join("note.md"); std::fs::write(&f, "# Guide\n\nRust borrow checker.").unwrap(); idx.reindex_file(&f, "workspace").unwrap(); drop(idx); let params = MemoryBackendParams { embed_config: Some(MemoryEmbeddingConfig::default()), embed_base_url: "http://localhost".to_string(), embed_api_key: None, // no key → provider cannot be created ..make_params_fts_only("test-embed-no-key") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); // Must not panic; FTS results should still come back. let results = backend.search("rust borrow", 5, 0.0).await.unwrap(); assert!( !results.is_empty(), "should fall back to FTS when api_key is None" ); } /// MemoryBackendParams is Clone. #[test] fn test_params_is_clone() { let params = make_params_fts_only("clone-test"); let _cloned = params.clone(); } /// from_session_params without watcher produces a backend that searches correctly. #[tokio::test] async fn test_factory_no_watcher() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let storage = make_storage(&tmp); let db_path = storage.workspace_dir().join("index.sqlite"); let mut idx = MemoryIndex::open_or_create( &db_path, storage.clone(), kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); let f = tmp.path().join("note.md"); std::fs::write(&f, "# Tip\n\nAlways write tests.").unwrap(); idx.reindex_file(&f, "workspace").unwrap(); drop(idx); let params = MemoryBackendParams { watcher: None, ..make_params_fts_only("test-no-watcher") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); let results = backend.search("tests", 5, 0.0).await.unwrap(); assert!( !results.is_empty(), "no-watcher backend should still return results" ); } /// `ensure_initialized` must be called before watcher startup. /// /// Regression test for the ordering fix: on a first-use machine the /// memory directories do not exist yet. If the watcher tries to watch a /// non-existent directory it returns `None` (silently dropping the feature). /// After `ensure_initialized()` the directories exist and the watcher can /// start successfully. /// /// This mirrors the ordering enforced in `spawn_session_actor`: /// 1. `storage.ensure_initialized()` /// 2. `MemoryFileWatcher::start(storage.global_dir())` #[test] fn test_ensure_initialized_before_watcher_ordering() { let tmp = TempDir::new().unwrap(); let global = tmp.path().join("memory"); let workspace = global.join("test_ws"); let storage = MemoryStorage::with_paths(global.clone(), workspace.clone()); // Precondition: neither directory exists yet (fresh machine simulation). assert!( !global.exists(), "global memory dir must not exist before initialization" ); // --- Wrong ordering (watcher before init) --- // The watcher returns None because the directory does not exist. let watcher_before_init = crate::watcher::MemoryFileWatcher::start(&global); assert!( watcher_before_init.is_none(), "watcher must fail (None) when directory does not exist yet" ); // --- Correct ordering (init, then watcher) --- // After ensure_initialized the directories and MEMORY.md templates exist. storage.ensure_initialized().unwrap(); assert!( global.exists(), "global dir must exist after ensure_initialized" ); assert!( workspace.exists(), "workspace dir must exist after ensure_initialized" ); assert!( global.join("MEMORY.md").exists(), "global MEMORY.md template must exist" ); assert!( workspace.join("MEMORY.md").exists(), "workspace MEMORY.md template must exist" ); // Watcher now succeeds because the directory exists. // (Allowed to return None in environments without inotify/kqueue // support — e.g. some CI containers — but must not error-panic.) let watcher_after_init = crate::watcher::MemoryFileWatcher::start(&global); // If a watcher was returned we can confirm it is usable (not dirty yet). if let Some(w) = watcher_after_init { assert!( !w.is_dirty(), "freshly started watcher must report no dirty files" ); } // If None, the test environment does not support file-watching — // that is acceptable; the directories themselves are what matter here. } /// End-to-end regression test for the watcher-driven delete path. /// /// Tests the full chain: /// 1. file is indexed /// 2. watcher is started /// 3. first `backend.search()` confirms content is found /// 4. file is deleted (OS fires a Remove event to the watcher) /// 5. second `backend.search()` triggers sync-on-search, which calls /// `delete_path()` because the file no longer exists /// 6. content is no longer returned /// /// This test guards against regressions in the `file.exists() → else /// delete_path()` branch that would be invisible to the `delete_path` /// unit tests alone. #[tokio::test] async fn test_watcher_delete_clears_stale_chunks() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let global = tmp.path().join("memory"); let workspace = global.join("test_ws"); std::fs::create_dir_all(&global).unwrap(); std::fs::create_dir_all(&workspace).unwrap(); let storage = MemoryStorage::with_paths(global.clone(), workspace); let db_path = storage.workspace_dir().join("index.sqlite"); // Step 1: Write + canonicalize the file path BEFORE indexing. // // On macOS, TempDir paths may live under /private/tmp (via a symlink // from /tmp). FSEvents returns canonicalized paths, so the path stored // in the index must match what the watcher event delivers. let file_raw = global.join("note.md"); std::fs::write(&file_raw, "# Unique\n\nXyzzy-watcher-delete-token.").unwrap(); let file = dunce::canonicalize(&file_raw).unwrap_or(file_raw); { let mut idx = MemoryIndex::open_or_create( &db_path, storage.clone(), kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); // Index with the canonical path so DB key matches watcher event paths. idx.reindex_file(&file, "workspace").unwrap(); } // Step 2: Start watcher AFTER indexing so the Remove event for the // upcoming deletion is the first event the watcher ever sees. let watch_dir = dunce::canonicalize(&global).unwrap_or(global.clone()); let watcher = match crate::watcher::MemoryFileWatcher::start(&watch_dir) { Some(w) => w, None => { // File-watching not supported in this environment (e.g., some CI // containers without inotify/FSEvents). Skip rather than fail. return; } }; let watcher_arc = std::sync::Arc::new(watcher); let params = MemoryBackendParams { watcher: Some(watcher_arc.clone()), ..make_params_fts_only("test-watcher-delete") }; let backend = MemoryBackendImpl::from_session_params(storage, ¶ms); // Step 3: Confirm content is found before deletion. let before = backend .search("Xyzzy-watcher-delete-token", 5, 0.0) .await .unwrap(); assert!( !before.is_empty(), "content must be found before file is deleted" ); // Step 4: Delete the file — the OS will fire a Remove event. std::fs::remove_file(&file).unwrap(); // Poll until the watcher detects the event (more reliable than a fixed // sleep on macOS where FSEvents delivery time varies considerably). // Give up after 2 s and skip the timing-sensitive assertion rather than // flake — delete_path unit tests cover the underlying logic. let mut event_delivered = false; for _ in 0..20 { std::thread::sleep(std::time::Duration::from_millis(100)); if watcher_arc.is_dirty() { event_delivered = true; break; } } if !event_delivered { // FSEvents not delivered within 2 s — environment is too slow. // Skip silently; the logic is covered by delete_path unit tests. return; } // Step 5+6: search triggers sync-on-search, which detects file.exists() // == false and calls delete_path(), clearing all stale chunks. let after = backend .search("Xyzzy-watcher-delete-token", 5, 0.0) .await .unwrap(); assert!( after.is_empty(), "deleted file's content must not appear after watcher-driven delete sync" ); } /// Regression: provider build must use `current_api_key_async`, /// never sync. Prevents memory_search 401s on rotated tokens. #[tokio::test] async fn make_embedding_provider_uses_async_api_key_resolution() { use kigi_tools::types::ApiKeyProvider; use std::sync::atomic::{AtomicU32, Ordering}; struct AsyncProbe { sync_calls: Arc, async_calls: Arc, } impl ApiKeyProvider for AsyncProbe { fn current_api_key(&self) -> Option { self.sync_calls.fetch_add(1, Ordering::SeqCst); Some("sync-stale".into()) } fn current_api_key_async( &self, ) -> std::pin::Pin> + Send + '_>> { let counter = self.async_calls.clone(); Box::pin(async move { counter.fetch_add(1, Ordering::SeqCst); Some("async-fresh".into()) }) } } let sync_calls = Arc::new(AtomicU32::new(0)); let async_calls = Arc::new(AtomicU32::new(0)); let probe: kigi_tools::types::SharedApiKeyProvider = Arc::new(AsyncProbe { sync_calls: sync_calls.clone(), async_calls: async_calls.clone(), }); let params = MemoryBackendParams { session_id: "s1".into(), embed_config: Some(MemoryEmbeddingConfig { model: Some("test-embed-model".into()), ..Default::default() }), embed_base_url: "http://example/v1".into(), embed_api_key: Some("static-fallback".into()), search_config: MemorySearchConfig::default(), watcher: None, stale_claim_secs: 60, search_source: "tool", api_key_provider: Some(probe), // No auth_credentials — forces the api_key_provider fallback path. auth_credentials: None, }; let provider = params.make_embedding_provider().await; assert!( provider.is_some(), "provider must be built when model is set" ); assert_eq!( async_calls.load(Ordering::SeqCst), 1, "must call current_api_key_async exactly once per provider build" ); assert_eq!( sync_calls.load(Ordering::SeqCst), 0, "sync current_api_key must NOT be called — the async path is the contract" ); } } #[cfg(test)] mod tests { use super::*; use crate::index::{MemoryIndex, init_sqlite_vec}; use kigi_config_types::MemoryIndexConfig; use tempfile::TempDir; fn setup_index(tmp: &TempDir) -> (PathBuf, MemoryStorage) { init_sqlite_vec(); let global = tmp.path().join("memory"); let workspace = global.join("test_ws"); let storage = MemoryStorage::with_paths(global, workspace); let db_path = tmp.path().join("test.sqlite"); let mut idx = MemoryIndex::open_or_create(&db_path, storage.clone(), MemoryIndexConfig::default(), 4) .unwrap(); let file_path = tmp.path().join("test.md"); std::fs::write(&file_path, "# Guide\n\nRust programming tutorial.").unwrap(); idx.reindex_file(&file_path, "workspace").unwrap(); (db_path, storage) } #[tokio::test] async fn test_backend_search() { let tmp = TempDir::new().unwrap(); let (db_path, storage) = setup_index(&tmp); let backend = MemoryBackendImpl::new(db_path, storage); let results = backend.search("rust programming", 10, 0.0).await.unwrap(); assert!(!results.is_empty(), "should find indexed content"); assert!(results[0].snippet.contains("Rust")); } #[test] fn test_backend_total_chunks() { let tmp = TempDir::new().unwrap(); let (db_path, storage) = setup_index(&tmp); let backend = MemoryBackendImpl::new(db_path, storage); let count = backend.total_chunks().unwrap(); assert!(count >= 1, "should have at least 1 chunk"); } #[test] fn test_backend_is_send_sync() { fn assert_send_sync() {} assert_send_sync::(); } #[tokio::test] async fn test_search_with_punctuation_in_query() { let tmp = TempDir::new().unwrap(); let (db_path, storage) = setup_index(&tmp); let backend = MemoryBackendImpl::new(db_path, storage); // Raw user message with punctuation — should not crash FTS5 let results = backend .search("what is rust? how to use it!", 10, 0.0) .await .unwrap(); assert!( !results.is_empty(), "should match 'rust' despite punctuation in query" ); } #[tokio::test] async fn test_search_with_special_chars_only() { let tmp = TempDir::new().unwrap(); let (db_path, storage) = setup_index(&tmp); let backend = MemoryBackendImpl::new(db_path, storage); // Query with only special chars — should return empty, not error let results = backend.search("???!!!", 10, 0.0).await.unwrap(); assert!( results.is_empty(), "special-chars-only query should return empty" ); } #[tokio::test] async fn test_search_hybrid_fts_only_fallback() { // Without embedding config, hybrid search should degrade to FTS-only let tmp = TempDir::new().unwrap(); let (db_path, storage) = setup_index(&tmp); let backend = MemoryBackendImpl::new(db_path, storage); // Even with high min_score, hybrid search normalizes scores to [0,1] // so results above the threshold should be returned let results = backend.search("rust programming", 10, 0.0).await.unwrap(); assert!( !results.is_empty(), "FTS-only fallback should still return results" ); // Scores should be normalized (0,1] range from hybrid scoring assert!(results[0].score > 0.0, "hybrid scores should be positive"); } /// The supplemental evergreen query in `search()` adds global/workspace /// candidates that the base `search_fts` missed due to candidate_limit. /// /// Tests the mechanism directly at the index level: verifies that with /// a tight FTS limit, global/workspace chunks are absent from the base /// results but present in the supplemental source-filtered query. Then /// confirms the full backend search pipeline surfaces them. #[tokio::test] async fn test_search_returns_global_and_workspace_memory() { let tmp = TempDir::new().unwrap(); init_sqlite_vec(); let global = tmp.path().join("memory"); let workspace = global.join("test_ws"); let storage = MemoryStorage::with_paths(global, workspace); let db_path = storage.workspace_dir().join("index.sqlite"); let mut idx = MemoryIndex::open_or_create(&db_path, storage.clone(), MemoryIndexConfig::default(), 4) .unwrap(); // Index global + workspace with matching content. let global_file = tmp.path().join("global_mem.md"); std::fs::write( &global_file, "# Preferences\n\nAlways use graphite for PRs. Prefer Rust over Python.", ) .unwrap(); idx.reindex_file(&global_file, "global").unwrap(); let ws_file = tmp.path().join("ws_mem.md"); std::fs::write( &ws_file, "# Project Decisions\n\nWe chose graphite for PRs in this project.", ) .unwrap(); idx.reindex_file(&ws_file, "workspace").unwrap(); // Index session files that also match the query. for i in 0..5 { let f = tmp.path().join(format!("session_{i}.md")); std::fs::write( &f, format!("# Session {i}\n\nDiscussed graphite for PRs and item {i}."), ) .unwrap(); idx.reindex_file(&f, "session").unwrap(); } // Verify the supplemental query mechanism: with a tight limit the // base FTS returns a mix, but `search_fts_by_sources` for // "global"/"workspace" always finds the evergreen chunks. let evergreen = idx .search_fts_by_sources("graphite PRs", 10, &["global", "workspace"]) .unwrap(); assert!( evergreen.len() >= 2, "supplemental evergreen query must find both global and workspace chunks" ); let evergreen_sources: Vec = evergreen .iter() .filter_map(|r| idx.get_chunk(&r.chunk_id).ok().flatten()) .map(|c| c.source) .collect(); assert!( evergreen_sources.contains(&"global".to_string()), "evergreen query must find global chunk" ); assert!( evergreen_sources.contains(&"workspace".to_string()), "evergreen query must find workspace chunk" ); drop(idx); // Full backend search: global/workspace must appear in results. let backend = MemoryBackendImpl::new(db_path, storage); let results = backend.search("graphite PRs", 10, 0.0).await.unwrap(); let has_global = results.iter().any(|r| r.source == "global"); let has_workspace = results.iter().any(|r| r.source == "workspace"); assert!( has_global, "global MEMORY.md chunks must appear in search results" ); assert!( has_workspace, "workspace MEMORY.md chunks must appear in search results" ); } } #[cfg(test)] mod index_embedding_tests { use crate::index::MemoryIndex; use crate::storage::MemoryStorage; #[test] fn test_chunks_without_embeddings() { let tmp = tempfile::TempDir::new().unwrap(); let global = tmp.path().join("memory"); let workspace = global.join("test_ws"); let storage = MemoryStorage::with_paths(global, workspace); let db_path = tmp.path().join("test.sqlite"); let mut idx = MemoryIndex::open_or_create( &db_path, storage, kigi_config_types::MemoryIndexConfig::default(), 4, ) .unwrap(); if !idx.vec_available() { // sqlite-vec not available — chunks_without_embeddings returns empty let missing = idx.chunks_without_embeddings().unwrap(); assert!(missing.is_empty(), "no-vec: should return empty"); return; } let file_path = tmp.path().join("test.md"); std::fs::write(&file_path, "# Title\n\nSome content here.").unwrap(); idx.reindex_file(&file_path, "workspace").unwrap(); // After reindex, chunks should exist but have no embeddings let missing = idx.chunks_without_embeddings().unwrap(); assert!( !missing.is_empty(), "newly indexed chunks should be missing embeddings" ); // After upserting an embedding, the chunk should disappear from missing let (chunk_id, _) = &missing[0]; let dummy_embedding = vec![0.0f32; 4]; idx.upsert_embedding(chunk_id, &dummy_embedding).unwrap(); let missing_after = idx.chunks_without_embeddings().unwrap(); assert_eq!( missing_after.len(), missing.len() - 1, "one fewer chunk should be missing after embedding" ); } }