leindex 1.9.5

LeIndex MCP and semantic code search engine for AI tools and large codebases
// Model path resolution for the worker process
//
// VAL-CPHASE-010: Worker model resolution uses the documented precedence:
// 1. Explicit env override (LEINDEX_MODEL_PATH)
// 2. Bundled models near the binary
// 3. User cache fallback (~/.leindex/models/ or $LEINDEX_HOME/models/)
//
// The resolver is used by the worker runtime to locate model and tokenizer
// files without requiring the main daemon to pass paths explicitly.

use std::path::{Path, PathBuf};

/// Error during model path resolution.
#[derive(Debug, Clone)]
pub struct ModelResolutionError {
    pub message: String,
}

impl std::fmt::Display for ModelResolutionError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "model resolution error: {}", self.message)
    }
}

impl std::error::Error for ModelResolutionError {}

/// Resolves model file paths using the documented precedence chain.
pub struct ModelResolver;

impl ModelResolver {
    fn bundled_model_dirs() -> Vec<PathBuf> {
        let mut dirs = Vec::new();
        if let Ok(exe_path) = std::env::current_exe() {
            if let Some(parent) = exe_path.parent() {
                // 1. exe_parent/models (e.g., target/release/models)
                dirs.push(parent.join("models"));
                if let Some(grandparent) = parent.parent() {
                    // 2. exe_parent/../models (e.g., target/models)
                    dirs.push(grandparent.join("models"));
                    // VAL-ONNX-004: 3. exe_parent/../../models (e.g., workspace root models)
                    // When running from target/release/, models are at ../../models/
                    if let Some(great_grandparent) = grandparent.parent() {
                        dirs.push(great_grandparent.join("models"));
                    }
                }
            }
        }
        dirs
    }

    fn user_model_dirs() -> Vec<PathBuf> {
        let mut dirs = Vec::new();
        if let Some(configured) = crate::config::model_dir_path() {
            dirs.push(configured);
        }
        if let Some(home) = dirs::home_dir() {
            let default_home = home.join(".leindex").join("models");
            if !dirs.iter().any(|dir| dir == &default_home) {
                dirs.push(default_home);
            }
        }
        dirs
    }

    /// Resolve the ONNX model file path for the given model name.
    ///
    /// VAL-CPHASE-010: Uses the precedence chain:
    /// 1. LEINDEX_MODEL_PATH env override
    /// 2. Bundled models directory (relative to the worker binary)
    /// 3. User cache directory (~/.leindex/models/ or $LEINDEX_HOME/models/)
    pub fn resolve(model_name: &str) -> Result<PathBuf, ModelResolutionError> {
        let model_filename = format!("{}.onnx", model_name);

        // 1. Explicit env override
        if let Ok(path) = std::env::var("LEINDEX_MODEL_PATH") {
            let model_path = PathBuf::from(path).join(&model_filename);
            if model_path.exists() {
                tracing::debug!("model resolved via env override: {}", model_path.display());
                return Ok(model_path);
            }
        }

        // 2. Bundled models (relative to the running binary)
        for bundled_dir in Self::bundled_model_dirs() {
            let model_path = bundled_dir.join(&model_filename);
            if model_path.exists() {
                tracing::debug!("model resolved via bundled path: {}", model_path.display());
                return Ok(model_path);
            }
        }

        // 3. User cache fallback
        for user_models in Self::user_model_dirs() {
            let model_path = user_models.join(&model_filename);
            if model_path.exists() {
                tracing::debug!("model resolved via user cache: {}", model_path.display());
                return Ok(model_path);
            }
        }

        Err(ModelResolutionError {
            message: format!(
                "model '{}' not found in any standard location (env, bundled, user cache)",
                model_name
            ),
        })
    }

    /// Resolve the tokenizer file path for the given model name.
    ///
    /// Uses the same precedence chain as model resolution.
    pub fn resolve_tokenizer(model_name: &str) -> Result<PathBuf, ModelResolutionError> {
        // Tokenizer is typically shared across model variants
        let _ = model_name; // Model name may be used for variant-specific tokenizers in future

        // 1. Explicit env override
        if let Ok(path) = std::env::var("LEINDEX_MODEL_PATH") {
            let tokenizer_path = PathBuf::from(path).join("tokenizer.json");
            if tokenizer_path.exists() {
                return Ok(tokenizer_path);
            }
        }

        // 2. Bundled models
        for bundled_dir in Self::bundled_model_dirs() {
            let tokenizer_path = bundled_dir.join("tokenizer.json");
            if tokenizer_path.exists() {
                return Ok(tokenizer_path);
            }
        }

        // 3. User cache fallback
        for user_models in Self::user_model_dirs() {
            let tokenizer_path = user_models.join("tokenizer.json");
            if tokenizer_path.exists() {
                return Ok(tokenizer_path);
            }
        }

        Err(ModelResolutionError {
            message: "tokenizer not found in any standard location".to_string(),
        })
    }

    /// Determine the source of a resolved path for reporting.
    ///
    /// Returns one of: "env_override", "bundled", "user_cache".
    pub fn source_for_path(path: &Path) -> &'static str {
        // Check env override first — if the env var is set and the path is
        // rooted under it, report as env_override regardless of whether the
        // file also happens to live near the binary.
        if let Ok(env_path) = std::env::var("LEINDEX_MODEL_PATH") {
            let env_dir = PathBuf::from(env_path);
            if path.starts_with(&env_dir) {
                return "env_override";
            }
        }

        // Check if it's near the binary
        if let Ok(exe_path) = std::env::current_exe() {
            if let Some(parent) = exe_path.parent() {
                if path.starts_with(parent) {
                    return "bundled";
                }
            }
        }

        "user_cache"
    }
}

/// All tests in this module mutate the `LEINDEX_MODEL_PATH` env var and must
/// not run concurrently.  We use a single test-serialising attribute so that
/// `cargo test -- --test-threads=N` still works correctly.
#[cfg(test)]
mod tests {
    use super::*;

    // Use the crate-level shared lock for env-var-mutating test serialization.
    use crate::embed::test_util::ENV_TEST_LOCK;

    #[test]
    fn test_resolve_model_not_found() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // Clear any env override
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var(crate::config::LEINDEX_HOME_ENV) };

        let result = ModelResolver::resolve("nonexistent-model-xyz");
        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(err.message.contains("nonexistent-model-xyz"));
        assert!(err.message.contains("not found"));
    }

    #[test]
    fn test_resolve_tokenizer_not_found() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var(crate::config::LEINDEX_HOME_ENV) };

        // Use a model name guaranteed not to correspond to a real model. The
        // user-cache fallback only triggers when `tokenizer.json` actually
        // exists in `~/.leindex/models/`, so this test correctly fails on dev
        // machines where that file is present. We therefore only assert the
        // error path when the user cache does NOT have a tokenizer.
        let user_has_tokenizer = dirs::home_dir()
            .map(|h| {
                h.join(".leindex")
                    .join("models")
                    .join("tokenizer.json")
                    .exists()
            })
            .unwrap_or(false);
        if !user_has_tokenizer {
            let result = ModelResolver::resolve_tokenizer("nonexistent");
            assert!(result.is_err());
            assert!(result.unwrap_err().message.contains("tokenizer not found"));
        } else {
            // On dev machines with `~/.leindex/models/tokenizer.json`, the
            // user-cache fallback legitimately resolves, so skip the assertion.
            // This is a pre-existing environment coupling, not a regression.
        }
    }

    #[test]
    fn test_resolve_with_env_override_missing_file() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // Set env to a temp dir that doesn't have the model
        let temp_dir = tempfile::tempdir().unwrap();
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::set_var("LEINDEX_MODEL_PATH", temp_dir.path()) };

        let result = ModelResolver::resolve("test-model");
        // Should still fail because the file doesn't exist in the temp dir
        assert!(result.is_err());

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
    }

    #[test]
    fn test_resolve_with_env_override_existing_file() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        let temp_dir = tempfile::tempdir().unwrap();
        let model_file = temp_dir.path().join("test-model.onnx");
        std::fs::write(&model_file, b"fake model").unwrap();

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::set_var("LEINDEX_MODEL_PATH", temp_dir.path()) };

        let result = ModelResolver::resolve("test-model");
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), model_file);

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
    }

    #[test]
    fn test_resolve_with_leindex_home_models() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
        let temp_dir = tempfile::tempdir().unwrap();
        let models_dir = temp_dir.path().join("models");
        std::fs::create_dir_all(&models_dir).unwrap();
        let model_file = models_dir.join("home-model.onnx");
        std::fs::write(&model_file, b"fake model").unwrap();

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::set_var(crate::config::LEINDEX_HOME_ENV, temp_dir.path()) };

        let result = ModelResolver::resolve("home-model");
        assert_eq!(result.unwrap(), model_file);

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var(crate::config::LEINDEX_HOME_ENV) };
    }

    #[test]
    fn test_resolve_tokenizer_with_env_override() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        let temp_dir = tempfile::tempdir().unwrap();
        let tokenizer_file = temp_dir.path().join("tokenizer.json");
        std::fs::write(&tokenizer_file, b"{}").unwrap();

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::set_var("LEINDEX_MODEL_PATH", temp_dir.path()) };

        let result = ModelResolver::resolve_tokenizer("test");
        assert!(result.is_ok());
        assert_eq!(result.unwrap(), tokenizer_file);

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var(crate::config::LEINDEX_HOME_ENV) };
    }

    #[test]
    fn test_source_for_path_env_override() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        let temp_dir = tempfile::tempdir().unwrap();
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::set_var("LEINDEX_MODEL_PATH", temp_dir.path()) };

        let path = temp_dir.path().join("model.onnx");
        assert_eq!(ModelResolver::source_for_path(&path), "env_override");

        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
    }

    #[test]
    fn test_source_for_path_user_cache() {
        let _guard = ENV_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
        // FIXME: Audit that the environment access only happens in single-threaded code.
        unsafe { std::env::remove_var("LEINDEX_MODEL_PATH") };
        let path = PathBuf::from("/some/random/path/model.onnx");
        assert_eq!(ModelResolver::source_for_path(&path), "user_cache");
    }

    #[test]
    fn test_model_resolution_error_display() {
        let err = ModelResolutionError {
            message: "test error".to_string(),
        };
        assert!(err.to_string().contains("test error"));
    }
}