Kigi-CLI/crates/codegen/kigi-memory/src/backend.rs

//! 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<kigi_config_types::MemoryEmbeddingConfig>,
    /// Base URL for embedding API calls (CLI proxy).
    pub embed_base_url: String,
    /// API key for embedding API calls.
    pub embed_api_key: Option<String>,
    /// 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<Arc<MemoryFileWatcher>>,
    /// 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<kigi_tools::types::SharedApiKeyProvider>,
    pub auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
}

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<super::embedding::ApiEmbeddingProvider> {
        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<dyn kigi_auth::AuthCredentialProvider>>,
    api_key_provider: Option<&kigi_tools::types::SharedApiKeyProvider>,
    static_api_key: Option<&str>,
    base_url: &str,
) -> Option<super::embedding::ApiEmbeddingProvider> {
    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<kigi_config_types::MemoryEmbeddingConfig>,
    /// API base URL for embedding requests (cli-chat-proxy).
    embed_base_url: String,
    /// API key for embedding requests.
    embed_api_key: Option<String>,
    /// Search scoring config (weights, min_score, max_results).
    search_config: kigi_config_types::MemorySearchConfig,
    /// File watcher for detecting external memory edits.
    watcher: Option<Arc<MemoryFileWatcher>>,
    /// 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<std::sync::atomic::AtomicU64>,
    /// Dynamic API key provider for embedding requests.
    api_key_provider: Option<kigi_tools::types::SharedApiKeyProvider>,
    /// Refresh-capable credential provider for embedding HTTP middleware.
    auth_credentials: Option<Arc<dyn kigi_auth::AuthCredentialProvider>>,
}

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<String>,
    ) -> 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<MemoryFileWatcher>, 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<rusqlite::Connection, rusqlite::Error> {
        // 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<super::embedding::ApiEmbeddingProvider> {
        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) = &params.embed_config {
            backend = backend.with_embedding(
                ec.clone(),
                params.embed_base_url.clone(),
                params.embed_api_key.clone(),
            );
        }
        if let Some(w) = &params.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<Vec<MemorySearchResult>, Box<dyn std::error::Error + Send + Sync>> {
        // 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<dyn std::error::Error + Send + Sync> {
            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<f32>)> = 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<String> =
            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<dyn std::error::Error + Send + Sync> {
            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<usize>,
        lines: Option<usize>,
    ) -> Result<String, Box<dyn std::error::Error + Send + Sync>> {
        Ok(self.storage.read_file(Path::new(path), from, lines)?)
    }

    fn total_chunks(&self) -> Result<usize, Box<dyn std::error::Error + Send + Sync>> {
        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, &params);

        // 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, &params);

        // 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, &params);
        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(), &params);
            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, &params);
        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, &params);
        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, &params);
        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, &params);
        // 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, &params);
        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, &params);

        // 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<AtomicU32>,
            async_calls: Arc<AtomicU32>,
        }
        impl ApiKeyProvider for AsyncProbe {
            fn current_api_key(&self) -> Option<String> {
                self.sync_calls.fetch_add(1, Ordering::SeqCst);
                Some("sync-stale".into())
            }
            fn current_api_key_async(
                &self,
            ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Option<String>> + 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<T: Send + Sync>() {}
        assert_send_sync::<MemoryBackendImpl>();
    }

    #[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<String> = 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"
        );
    }
}