optimclaw/src/tools/wasm/loader.rs

//! Generic WASM tool loader for loading tools from files or directories.
//!
//! This module provides a way to load WASM tools dynamically at runtime from:
//! - A directory containing `<name>.wasm` and `<name>.capabilities.json`
//! - Build artifacts in `tools-src/` (dev mode, auto-detected)
//! - Database storage (via [`WasmToolStore`])
//!
//! # Example: Loading from Directory
//!
//! ```text
//! ~/.ironclaw/tools/
//! ├── slack.wasm
//! ├── slack.capabilities.json
//! ├── github.wasm
//! └── github.capabilities.json
//! ```
//!
//! ```ignore
//! let loader = WasmToolLoader::new(runtime, registry);
//! loader.load_from_dir(Path::new("~/.ironclaw/tools/")).await?;
//! ```
//!
//! # Dev Mode
//!
//! When `load_dev_tools()` is called, the loader scans `tools-src/*/` for build
//! artifacts. Tools found there are loaded directly from the build output,
//! skipping the install directory. This means during development you just
//! rebuild the WASM and restart the host, no manual copy step needed.
//!
//! # Security
//!
//! Tools loaded from files are assigned `TrustLevel::User` by default, meaning
//! they run with the most restrictive permissions. Only tools explicitly marked
//! as `verified` or `system` in the database get elevated trust.

use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::Arc;

use tokio::fs;

use crate::secrets::SecretsStore;
use crate::tools::registry::{ToolRegistry, WasmRegistrationError, WasmToolRegistration};
use crate::tools::wasm::capabilities_schema::CapabilitiesFile;
use crate::tools::wasm::{
    Capabilities, OAuthRefreshConfig, WasmError, WasmStorageError, WasmToolRuntime, WasmToolStore,
};

/// Error during WASM tool loading.
#[derive(Debug, thiserror::Error)]
pub enum WasmLoadError {
    #[error("IO error: {0}")]
    Io(#[from] std::io::Error),

    #[error("WASM file not found: {0}")]
    WasmNotFound(PathBuf),

    #[error("Capabilities file not found: {0}")]
    CapabilitiesNotFound(PathBuf),

    #[error("Invalid capabilities JSON: {0}")]
    InvalidCapabilities(String),

    #[error("WASM compilation error: {0}")]
    Compilation(#[from] WasmError),

    #[error("Storage error: {0}")]
    Storage(#[from] WasmStorageError),

    #[error("Registration error: {0}")]
    Registration(#[from] WasmRegistrationError),

    #[error("Invalid tool name: {0}")]
    InvalidName(String),

    #[error("WIT version mismatch: {0}")]
    WitVersionMismatch(String),
}

/// Loads WASM tools from files or storage into the registry.
pub struct WasmToolLoader {
    runtime: Arc<WasmToolRuntime>,
    registry: Arc<ToolRegistry>,
    secrets_store: Option<Arc<dyn SecretsStore + Send + Sync>>,
}

impl WasmToolLoader {
    /// Create a new loader with the given runtime and registry.
    pub fn new(runtime: Arc<WasmToolRuntime>, registry: Arc<ToolRegistry>) -> Self {
        Self {
            runtime,
            registry,
            secrets_store: None,
        }
    }

    /// Set the secrets store for credential injection in WASM tools.
    pub fn with_secrets_store(mut self, store: Arc<dyn SecretsStore + Send + Sync>) -> Self {
        self.secrets_store = Some(store);
        self
    }

    /// Load a single WASM tool from a file pair.
    ///
    /// Expects:
    /// - `wasm_path`: Path to the `.wasm` file
    /// - `capabilities_path`: Path to the `.capabilities.json` file (optional)
    ///
    /// If no capabilities file is provided, the tool gets no capabilities (default deny).
    pub async fn load_from_files(
        &self,
        name: &str,
        wasm_path: &Path,
        capabilities_path: Option<&Path>,
    ) -> Result<(), WasmLoadError> {
        if name.is_empty() || name.contains('/') || name.contains('\\') {
            return Err(WasmLoadError::InvalidName(name.to_string()));
        }

        // Read WASM bytes
        if !wasm_path.exists() {
            return Err(WasmLoadError::WasmNotFound(wasm_path.to_path_buf()));
        }
        let wasm_bytes = fs::read(wasm_path).await?;

        // Read capabilities (optional) and extract OAuth refresh config
        let (capabilities, oauth_refresh) = if let Some(cap_path) = capabilities_path {
            if cap_path.exists() {
                let cap_bytes = fs::read(cap_path).await?;
                let cap_file = CapabilitiesFile::from_bytes(&cap_bytes)
                    .map_err(|e| WasmLoadError::InvalidCapabilities(e.to_string()))?;
                cap_file.validate(name);

                // Check WIT version compatibility
                check_wit_version_compat(
                    name,
                    cap_file.wit_version.as_deref(),
                    crate::tools::wasm::WIT_TOOL_VERSION,
                )?;

                let caps = cap_file.to_capabilities();
                let oauth = resolve_oauth_refresh_config(&cap_file);
                (caps, oauth)
            } else {
                tracing::warn!(
                    path = %cap_path.display(),
                    "Capabilities file not found, using default (no permissions)"
                );
                (Capabilities::default(), None)
            }
        } else {
            (Capabilities::default(), None)
        };

        // Register the tool
        self.registry
            .register_wasm(WasmToolRegistration {
                name,
                wasm_bytes: &wasm_bytes,
                runtime: &self.runtime,
                capabilities,
                limits: None,
                description: None,
                schema: None,
                secrets_store: self.secrets_store.clone(),
                oauth_refresh,
            })
            .await?;

        tracing::info!(
            name = name,
            wasm_path = %wasm_path.display(),
            "Loaded WASM tool from file"
        );

        Ok(())
    }

    /// Load all WASM tools from a directory.
    ///
    /// Scans the directory for `*.wasm` files and loads each one, looking for
    /// a matching `*.capabilities.json` sidecar file.
    ///
    /// # Directory Layout
    ///
    /// ```text
    /// tools/
    /// ├── slack.wasm                  <- Tool WASM component
    /// ├── slack.capabilities.json     <- Capabilities (optional)
    /// ├── github.wasm
    /// └── github.capabilities.json
    /// ```
    ///
    /// Tools without a capabilities file get no permissions (default deny).
    pub async fn load_from_dir(&self, dir: &Path) -> Result<LoadResults, WasmLoadError> {
        if !dir.is_dir() {
            return Err(WasmLoadError::Io(std::io::Error::new(
                std::io::ErrorKind::NotADirectory,
                format!("{} is not a directory", dir.display()),
            )));
        }

        let mut results = LoadResults::default();

        // Collect all .wasm entries first, then load in parallel
        let mut tool_entries = Vec::new();
        let mut entries = fs::read_dir(dir).await?;

        while let Some(entry) = entries.next_entry().await? {
            let path = entry.path();

            if path.extension().and_then(|e| e.to_str()) != Some("wasm") {
                continue;
            }

            let name = match path.file_stem().and_then(|s| s.to_str()) {
                Some(n) => n.to_string(),
                None => {
                    results.errors.push((
                        path.clone(),
                        WasmLoadError::InvalidName("invalid filename".to_string()),
                    ));
                    continue;
                }
            };

            let cap_path = path.with_extension("capabilities.json");
            let has_cap = cap_path.exists();
            tool_entries.push((name, path, if has_cap { Some(cap_path) } else { None }));
        }

        // Load all tools in parallel (file I/O + WASM compilation + registration)
        let load_futures = tool_entries
            .iter()
            .map(|(name, path, cap_path)| self.load_from_files(name, path, cap_path.as_deref()));

        let load_results = futures::future::join_all(load_futures).await;

        for ((name, path, _), result) in tool_entries.into_iter().zip(load_results) {
            match result {
                Ok(()) => {
                    results.loaded.push(name);
                }
                Err(e) => {
                    tracing::error!(
                        name = name,
                        path = %path.display(),
                        error = %e,
                        "Failed to load WASM tool"
                    );
                    results.errors.push((path, e));
                }
            }
        }

        if !results.loaded.is_empty() {
            tracing::info!(
                count = results.loaded.len(),
                tools = ?results.loaded,
                "Loaded WASM tools from directory"
            );
        }

        Ok(results)
    }

    /// Load a WASM tool from database storage.
    ///
    /// This is a convenience wrapper around [`ToolRegistry::register_wasm_from_storage`].
    pub async fn load_from_storage(
        &self,
        store: &dyn WasmToolStore,
        user_id: &str,
        tool_name: &str,
    ) -> Result<(), WasmLoadError> {
        self.registry
            .register_wasm_from_storage(store, &self.runtime, user_id, tool_name)
            .await?;

        tracing::info!(
            user_id = user_id,
            name = tool_name,
            "Loaded WASM tool from storage"
        );

        Ok(())
    }

    /// Load all active WASM tools for a user from storage.
    pub async fn load_all_from_storage(
        &self,
        store: &dyn WasmToolStore,
        user_id: &str,
    ) -> Result<LoadResults, WasmLoadError> {
        let tools = store.list(user_id).await?;
        let mut results = LoadResults::default();

        for tool in tools {
            // Skip non-active tools
            if tool.status != crate::tools::wasm::ToolStatus::Active {
                continue;
            }

            match self.load_from_storage(store, user_id, &tool.name).await {
                Ok(()) => {
                    results.loaded.push(tool.name);
                }
                Err(e) => {
                    tracing::error!(
                        name = tool.name,
                        user_id = user_id,
                        error = %e,
                        "Failed to load WASM tool from storage"
                    );
                    results.errors.push((PathBuf::from(&tool.name), e));
                }
            }
        }

        Ok(results)
    }
}

/// Check that a declared WIT version is compatible with the host WIT version.
///
/// Compatibility rules (semver):
/// - Same major version required (0.x is special: same minor required)
/// - Extension WIT version must not be greater than host version
///
/// If `declared` is `None`, the check is skipped (pre-versioning extension).
pub fn check_wit_version_compat(
    name: &str,
    declared: Option<&str>,
    host_version: &str,
) -> Result<(), WasmLoadError> {
    let Some(declared_str) = declared else {
        return Ok(());
    };

    let declared = semver::Version::parse(declared_str).map_err(|e| {
        WasmLoadError::WitVersionMismatch(format!(
            "Extension '{name}' has invalid wit_version '{declared_str}': {e}"
        ))
    })?;

    let host = semver::Version::parse(host_version).map_err(|e| {
        WasmLoadError::WitVersionMismatch(format!(
            "Host WIT version '{host_version}' is invalid: {e}"
        ))
    })?;

    // Major version must match
    if declared.major != host.major {
        return Err(WasmLoadError::WitVersionMismatch(format!(
            "Extension '{name}' compiled against WIT {declared}, but host supports WIT {host}. \
             Major version mismatch — rebuild the extension."
        )));
    }

    // For 0.x versions, minor must also match (semver: 0.x.y has no compatibility guarantees)
    if declared.major == 0 && declared.minor != host.minor {
        return Err(WasmLoadError::WitVersionMismatch(format!(
            "Extension '{name}' compiled against WIT {declared}, but host supports WIT {host}. \
             Rebuild the extension against the current WIT."
        )));
    }

    // Extension cannot be newer than host
    if declared > host {
        return Err(WasmLoadError::WitVersionMismatch(format!(
            "Extension '{name}' compiled against WIT {declared}, but host only supports WIT {host}. \
             Update the host or rebuild with an older WIT."
        )));
    }

    Ok(())
}

/// Extract OAuth refresh configuration from a parsed capabilities file.
///
/// Returns `None` if there's no `auth.oauth` section or if the client_id
/// can't be resolved from any source (inline, env var, or built-in defaults).
///
/// Fallback chain for client_id:
///   `oauth.client_id` > env var (`oauth.client_id_env`) > `builtin_credentials()`
fn resolve_oauth_refresh_config(cap_file: &CapabilitiesFile) -> Option<OAuthRefreshConfig> {
    let auth = cap_file.auth.as_ref()?;
    let oauth = auth.oauth.as_ref()?;

    let builtin = crate::cli::oauth_defaults::builtin_credentials(&auth.secret_name);

    let client_id = oauth
        .client_id
        .clone()
        .or_else(|| {
            oauth
                .client_id_env
                .as_ref()
                .and_then(|env| std::env::var(env).ok())
        })
        .or_else(|| builtin.as_ref().map(|c| c.client_id.to_string()))?;

    let client_secret = oauth
        .client_secret
        .clone()
        .or_else(|| {
            oauth
                .client_secret_env
                .as_ref()
                .and_then(|env| std::env::var(env).ok())
        })
        .or_else(|| builtin.as_ref().map(|c| c.client_secret.to_string()));

    Some(OAuthRefreshConfig {
        token_url: oauth.token_url.clone(),
        client_id,
        client_secret,
        secret_name: auth.secret_name.clone(),
        provider: auth.provider.clone(),
    })
}

/// Results from loading multiple tools.
#[derive(Debug, Default)]
pub struct LoadResults {
    /// Names of successfully loaded tools.
    pub loaded: Vec<String>,

    /// Errors encountered (path/name, error).
    pub errors: Vec<(PathBuf, WasmLoadError)>,
}

impl LoadResults {
    /// Check if all tools loaded successfully.
    pub fn all_succeeded(&self) -> bool {
        self.errors.is_empty()
    }

    /// Get the count of successfully loaded tools.
    pub fn success_count(&self) -> usize {
        self.loaded.len()
    }

    /// Get the count of failed tools.
    pub fn error_count(&self) -> usize {
        self.errors.len()
    }
}

/// Compile-time project root, used to locate tools-src/ in dev builds.
const CARGO_MANIFEST_DIR: &str = env!("CARGO_MANIFEST_DIR");

/// Resolve the WASM target directory for a given crate directory.
///
/// Checks (in order):
/// 1. `CARGO_TARGET_DIR` env var (shared target dir)
/// 2. `<crate_dir>/target/` (default per-crate layout)
pub fn resolve_wasm_target_dir(crate_dir: &Path) -> PathBuf {
    crate::registry::artifacts::resolve_target_dir(crate_dir)
}

/// Return the expected path to a compiled WASM artifact for a given crate.
///
/// Combines [`resolve_wasm_target_dir`] with the `wasm32-wasip2/release/` subdirectory
/// and the binary name without extension (e.g. `slack_tool`).
///
/// `binary_name` should not include the `.wasm` extension; it is appended automatically.
///
/// This is a convenience function for callers that know the exact triple (wasip2)
/// and binary name. For multi-triple search, use
/// [`crate::registry::artifacts::find_wasm_artifact`] instead.
pub fn wasm_artifact_path(crate_dir: &Path, binary_name: &str) -> PathBuf {
    resolve_wasm_target_dir(crate_dir)
        .join("wasm32-wasip2/release")
        .join(format!("{}.wasm", binary_name))
}

/// Resolve the tools source directory.
///
/// Checks (in order):
/// 1. `IRONCLAW_TOOLS_SRC` env var
/// 2. `<CARGO_MANIFEST_DIR>/tools-src/` (dev builds)
fn tools_src_dir() -> PathBuf {
    if let Ok(dir) = std::env::var("IRONCLAW_TOOLS_SRC") {
        return PathBuf::from(dir);
    }
    PathBuf::from(CARGO_MANIFEST_DIR).join("tools-src")
}

/// Discover WASM tools available as build artifacts in `tools-src/`.
///
/// Scans each subdirectory for:
/// - `tools-src/<name>/target/wasm32-wasip2/release/<crate_name>_tool.wasm`
/// - `tools-src/<name>/<name>-tool.capabilities.json`
///
/// Returns a map of install-name (e.g. "gmail-tool") to paths.
pub async fn discover_dev_tools() -> Result<HashMap<String, DiscoveredTool>, std::io::Error> {
    let src_dir = tools_src_dir();
    let mut tools = HashMap::new();

    if !src_dir.is_dir() {
        return Ok(tools);
    }

    let mut entries = fs::read_dir(&src_dir).await?;
    while let Some(entry) = entries.next_entry().await? {
        let path = entry.path();
        if !path.is_dir() {
            continue;
        }

        let dir_name = match path.file_name().and_then(|n| n.to_str()) {
            Some(n) => n.to_string(),
            None => continue,
        };

        // Convention: crate name uses underscores, directory uses hyphens
        let crate_name = dir_name.replace('-', "_");
        let install_name = format!("{}-tool", dir_name);

        let wasm_path = wasm_artifact_path(&path, &format!("{}_tool", crate_name));

        if !wasm_path.exists() {
            continue;
        }

        let caps_path = path.join(format!("{}-tool.capabilities.json", dir_name));

        tools.insert(
            install_name,
            DiscoveredTool {
                wasm_path,
                capabilities_path: if caps_path.exists() {
                    Some(caps_path)
                } else {
                    None
                },
            },
        );
    }

    Ok(tools)
}

/// Load WASM tools from build artifacts in `tools-src/`.
///
/// In dev mode, tools can be loaded directly from their build output without
/// needing to install them to `~/.ironclaw/tools/` first. Build artifacts
/// that are newer than installed copies take priority.
///
/// Set `IRONCLAW_TOOLS_SRC` env var to override the source directory.
pub async fn load_dev_tools(
    loader: &WasmToolLoader,
    install_dir: &Path,
) -> Result<LoadResults, WasmLoadError> {
    let dev_tools = discover_dev_tools().await?;
    let mut results = LoadResults::default();

    if dev_tools.is_empty() {
        return Ok(results);
    }

    for (name, discovered) in &dev_tools {
        // Check if the build artifact is newer than the installed copy
        let installed_path = install_dir.join(format!("{}.wasm", name));
        let should_load = if installed_path.exists() {
            // Compare modification times: prefer fresher build artifact
            match (
                fs::metadata(&discovered.wasm_path).await,
                fs::metadata(&installed_path).await,
            ) {
                (Ok(dev_meta), Ok(inst_meta)) => {
                    let dev_modified = dev_meta.modified().unwrap_or(std::time::UNIX_EPOCH);
                    let inst_modified = inst_meta.modified().unwrap_or(std::time::UNIX_EPOCH);
                    dev_modified > inst_modified
                }
                _ => true,
            }
        } else {
            true
        };

        if !should_load {
            continue;
        }

        tracing::info!(
            name = name,
            wasm_path = %discovered.wasm_path.display(),
            "Loading dev tool from build artifacts (newer than installed)"
        );

        match loader
            .load_from_files(
                name,
                &discovered.wasm_path,
                discovered.capabilities_path.as_deref(),
            )
            .await
        {
            Ok(()) => {
                results.loaded.push(name.clone());
            }
            Err(e) => {
                tracing::error!(
                    name = name,
                    error = %e,
                    "Failed to load dev tool"
                );
                results.errors.push((discovered.wasm_path.clone(), e));
            }
        }
    }

    if !results.loaded.is_empty() {
        tracing::info!(
            count = results.loaded.len(),
            tools = ?results.loaded,
            "Loaded dev tools from build artifacts"
        );
    }

    Ok(results)
}

/// Discover WASM tool files in a directory without loading them.
///
/// Returns a map of tool name -> (wasm_path, capabilities_path).
pub async fn discover_tools(dir: &Path) -> Result<HashMap<String, DiscoveredTool>, std::io::Error> {
    let mut tools = HashMap::new();

    if !dir.is_dir() {
        return Ok(tools);
    }

    let mut entries = fs::read_dir(dir).await?;

    while let Some(entry) = entries.next_entry().await? {
        let path = entry.path();

        if path.extension().and_then(|e| e.to_str()) != Some("wasm") {
            continue;
        }

        let name = match path.file_stem().and_then(|s| s.to_str()) {
            Some(n) => n.to_string(),
            None => continue,
        };

        let cap_path = path.with_extension("capabilities.json");

        tools.insert(
            name,
            DiscoveredTool {
                wasm_path: path,
                capabilities_path: if cap_path.exists() {
                    Some(cap_path)
                } else {
                    None
                },
            },
        );
    }

    Ok(tools)
}

/// A discovered WASM tool (not yet loaded).
#[derive(Debug)]
pub struct DiscoveredTool {
    /// Path to the WASM file.
    pub wasm_path: PathBuf,

    /// Path to the capabilities file (if present).
    pub capabilities_path: Option<PathBuf>,
}

#[cfg(test)]
mod tests {
    use std::io::Write;

    use tempfile::TempDir;

    use crate::tools::wasm::loader::{WasmLoadError, check_wit_version_compat, discover_tools};

    #[test]
    fn wit_version_compat_none_is_ok() {
        // Pre-versioning extensions (no wit_version declared) should always pass
        assert!(check_wit_version_compat("test", None, "0.2.0").is_ok());
    }

    #[test]
    fn wit_version_compat_exact_match() {
        assert!(check_wit_version_compat("test", Some("0.2.0"), "0.2.0").is_ok());
    }

    #[test]
    fn wit_version_compat_patch_older_ok() {
        // Extension on older patch of same minor is compatible
        assert!(check_wit_version_compat("test", Some("0.2.0"), "0.2.1").is_ok());
    }

    #[test]
    fn wit_version_compat_minor_mismatch_0x() {
        // For 0.x, different minor is breaking
        assert!(check_wit_version_compat("test", Some("0.1.0"), "0.2.0").is_err());
        assert!(check_wit_version_compat("test", Some("0.3.0"), "0.2.0").is_err());
    }

    #[test]
    fn wit_version_compat_major_mismatch() {
        assert!(check_wit_version_compat("test", Some("1.0.0"), "2.0.0").is_err());
    }

    #[test]
    fn wit_version_compat_extension_newer_than_host() {
        assert!(check_wit_version_compat("test", Some("0.2.1"), "0.2.0").is_err());
    }

    #[test]
    fn wit_version_compat_invalid_version() {
        assert!(check_wit_version_compat("test", Some("not-a-version"), "0.2.0").is_err());
    }

    #[tokio::test]
    async fn test_discover_tools_empty_dir() {
        let dir = TempDir::new().unwrap();
        let tools = discover_tools(dir.path()).await.unwrap();
        assert!(tools.is_empty());
    }

    #[tokio::test]
    async fn test_discover_tools_with_wasm() {
        let dir = TempDir::new().unwrap();

        // Create a fake .wasm file
        let wasm_path = dir.path().join("test_tool.wasm");
        std::fs::File::create(&wasm_path).unwrap();

        let tools = discover_tools(dir.path()).await.unwrap();
        assert_eq!(tools.len(), 1);
        assert!(tools.contains_key("test_tool"));
        assert!(tools["test_tool"].capabilities_path.is_none());
    }

    #[tokio::test]
    async fn test_discover_tools_with_capabilities() {
        let dir = TempDir::new().unwrap();

        // Create wasm and capabilities files
        std::fs::File::create(dir.path().join("slack.wasm")).unwrap();
        let mut cap_file =
            std::fs::File::create(dir.path().join("slack.capabilities.json")).unwrap();
        cap_file.write_all(b"{}").unwrap();

        let tools = discover_tools(dir.path()).await.unwrap();
        assert_eq!(tools.len(), 1);
        assert!(tools["slack"].capabilities_path.is_some());
    }

    #[tokio::test]
    async fn test_discover_tools_ignores_non_wasm() {
        let dir = TempDir::new().unwrap();

        // Create non-wasm files
        std::fs::File::create(dir.path().join("readme.md")).unwrap();
        std::fs::File::create(dir.path().join("config.json")).unwrap();
        std::fs::File::create(dir.path().join("tool.wasm")).unwrap();

        let tools = discover_tools(dir.path()).await.unwrap();
        assert_eq!(tools.len(), 1);
        assert!(tools.contains_key("tool"));
    }

    #[test]
    fn test_load_error_display() {
        let err = WasmLoadError::InvalidName("bad/name".to_string());
        assert!(err.to_string().contains("bad/name"));

        let err = WasmLoadError::WasmNotFound(std::path::PathBuf::from("/foo/bar.wasm"));
        assert!(err.to_string().contains("/foo/bar.wasm"));
    }

    #[test]
    fn test_tools_src_dir_default() {
        let dir = super::tools_src_dir();
        assert!(dir.ends_with("tools-src"));
    }

    #[tokio::test]
    async fn test_discover_dev_tools_finds_build_artifacts() {
        // This test relies on the actual tools-src/ directory in the repo.
        // If build artifacts exist, they should be discovered.
        let tools = super::discover_dev_tools().await.unwrap();

        // If any tools have been built, they should appear with "-tool" suffix
        for (name, discovered) in &tools {
            assert!(
                name.ends_with("-tool"),
                "Dev tool name should end with -tool: {}",
                name
            );
            assert!(
                discovered.wasm_path.exists(),
                "WASM should exist: {:?}",
                discovered.wasm_path
            );
        }
    }

    #[test]
    fn test_resolve_oauth_refresh_config_with_oauth() {
        use crate::tools::wasm::capabilities_schema::{
            AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
        };

        let caps = CapabilitiesFile {
            auth: Some(AuthCapabilitySchema {
                secret_name: "google_oauth_token".to_string(),
                provider: Some("google".to_string()),
                oauth: Some(OAuthConfigSchema {
                    authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
                    token_url: "https://oauth2.googleapis.com/token".to_string(),
                    client_id: Some("test-client-id".to_string()),
                    client_secret: Some("test-client-secret".to_string()),
                    ..Default::default()
                }),
                ..Default::default()
            }),
            ..Default::default()
        };

        let config = super::resolve_oauth_refresh_config(&caps);
        assert!(config.is_some());

        let config = config.unwrap();
        assert_eq!(config.token_url, "https://oauth2.googleapis.com/token");
        assert_eq!(config.client_id, "test-client-id");
        assert_eq!(config.client_secret, Some("test-client-secret".to_string()));
        assert_eq!(config.secret_name, "google_oauth_token");
        assert_eq!(config.provider, Some("google".to_string()));
    }

    #[test]
    fn test_resolve_oauth_refresh_config_no_auth() {
        use crate::tools::wasm::capabilities_schema::CapabilitiesFile;

        let caps = CapabilitiesFile::default();
        let config = super::resolve_oauth_refresh_config(&caps);
        assert!(config.is_none());
    }

    #[test]
    fn test_resolve_oauth_refresh_config_no_oauth() {
        use crate::tools::wasm::capabilities_schema::{AuthCapabilitySchema, CapabilitiesFile};

        let caps = CapabilitiesFile {
            auth: Some(AuthCapabilitySchema {
                secret_name: "manual_token".to_string(),
                ..Default::default()
            }),
            ..Default::default()
        };

        let config = super::resolve_oauth_refresh_config(&caps);
        assert!(config.is_none());
    }

    #[test]
    fn test_resolve_oauth_refresh_config_no_client_id() {
        use crate::tools::wasm::capabilities_schema::{
            AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
        };

        // A non-Google provider with no client_id anywhere should return None
        let caps = CapabilitiesFile {
            auth: Some(AuthCapabilitySchema {
                secret_name: "unknown_provider_token".to_string(),
                oauth: Some(OAuthConfigSchema {
                    authorization_url: "https://example.com/auth".to_string(),
                    token_url: "https://example.com/token".to_string(),
                    // No client_id, no client_id_env, no builtin
                    ..Default::default()
                }),
                ..Default::default()
            }),
            ..Default::default()
        };

        let config = super::resolve_oauth_refresh_config(&caps);
        assert!(config.is_none());
    }

    #[test]
    fn test_resolve_oauth_refresh_config_builtin_google() {
        use crate::tools::wasm::capabilities_schema::{
            AuthCapabilitySchema, CapabilitiesFile, OAuthConfigSchema,
        };

        // google_oauth_token should fall back to built-in credentials
        let caps = CapabilitiesFile {
            auth: Some(AuthCapabilitySchema {
                secret_name: "google_oauth_token".to_string(),
                provider: Some("google".to_string()),
                oauth: Some(OAuthConfigSchema {
                    authorization_url: "https://accounts.google.com/o/oauth2/v2/auth".to_string(),
                    token_url: "https://oauth2.googleapis.com/token".to_string(),
                    // No inline client_id, should fall back to builtin
                    ..Default::default()
                }),
                ..Default::default()
            }),
            ..Default::default()
        };

        let config = super::resolve_oauth_refresh_config(&caps);
        assert!(config.is_some());
        let config = config.unwrap();
        assert!(!config.client_id.is_empty());
        assert!(config.client_secret.is_some());
    }

    // ---------------------------------------------------------------
    // Security regression tests
    // ---------------------------------------------------------------

    use std::sync::Arc;

    use crate::tools::registry::ToolRegistry;
    use crate::tools::wasm::{WasmRuntimeConfig, WasmToolRuntime};

    /// Helper: create a WasmToolLoader backed by a real runtime + registry.
    fn make_loader() -> super::WasmToolLoader {
        let runtime = Arc::new(
            WasmToolRuntime::new(WasmRuntimeConfig::for_testing())
                .expect("failed to create WASM runtime for test"),
        );
        let registry = Arc::new(ToolRegistry::new());
        super::WasmToolLoader::new(runtime, registry)
    }

    #[tokio::test]
    async fn test_tool_name_rejects_path_separators() {
        let dir = TempDir::new().unwrap();
        // Create a valid wasm file so the name check is the only failure path
        let wasm_path = dir.path().join("dummy.wasm");
        std::fs::File::create(&wasm_path).unwrap();

        let loader = make_loader();

        for bad_name in &["../evil", "foo/bar", "foo\\bar"] {
            let result = loader.load_from_files(bad_name, &wasm_path, None).await;
            assert!(
                result.is_err(),
                "Expected error for name {:?}, got Ok",
                bad_name
            );
            let err = result.unwrap_err();
            assert!(
                matches!(err, WasmLoadError::InvalidName(_)),
                "Expected InvalidName for {:?}, got: {}",
                bad_name,
                err
            );
        }
    }

    #[tokio::test]
    async fn test_tool_name_rejects_empty() {
        let dir = TempDir::new().unwrap();
        let wasm_path = dir.path().join("dummy.wasm");
        std::fs::File::create(&wasm_path).unwrap();

        let loader = make_loader();
        let result = loader.load_from_files("", &wasm_path, None).await;

        assert!(result.is_err(), "Expected error for empty name, got Ok");
        let err = result.unwrap_err();
        assert!(
            matches!(err, WasmLoadError::InvalidName(_)),
            "Expected InvalidName for empty string, got: {}",
            err
        );
    }

    #[tokio::test]
    async fn test_load_nonexistent_wasm_file() {
        let loader = make_loader();
        let bogus_path = std::path::PathBuf::from("/tmp/nonexistent_tool_12345.wasm");

        let result = loader.load_from_files("bogus", &bogus_path, None).await;
        assert!(
            result.is_err(),
            "Expected error for nonexistent file, got Ok"
        );
        let err = result.unwrap_err();
        assert!(
            matches!(err, WasmLoadError::WasmNotFound(_)),
            "Expected WasmNotFound, got: {}",
            err
        );
    }

    #[tokio::test]
    async fn test_load_invalid_wasm_bytes() {
        let dir = TempDir::new().unwrap();
        let wasm_path = dir.path().join("invalid.wasm");

        // Write random invalid bytes (not a valid WASM module)
        let mut f = std::fs::File::create(&wasm_path).unwrap();
        f.write_all(b"this is not a valid wasm module at all")
            .unwrap();

        let loader = make_loader();
        let result = loader.load_from_files("invalid", &wasm_path, None).await;

        assert!(
            result.is_err(),
            "Expected error for invalid WASM bytes, got Ok"
        );
        // The error should come from WASM compilation or registration, not name validation
        let err = result.unwrap_err();
        assert!(
            !matches!(err, WasmLoadError::InvalidName(_)),
            "Got InvalidName instead of a compilation/registration error: {}",
            err
        );
    }

    #[tokio::test]
    async fn test_discover_skips_dotfiles() {
        let dir = TempDir::new().unwrap();

        // Create a dotfile .wasm and a normal .wasm
        std::fs::File::create(dir.path().join(".hidden.wasm")).unwrap();
        std::fs::File::create(dir.path().join("visible.wasm")).unwrap();

        let tools = discover_tools(dir.path()).await.unwrap();

        // The current implementation discovers ALL .wasm files including dotfiles.
        // This test documents the current behavior: .hidden.wasm IS discovered
        // with the stem ".hidden". A future hardening pass could add dotfile
        // filtering, at which point this assertion should be updated.
        assert!(
            tools.contains_key("visible"),
            "visible.wasm should be discovered"
        );
        assert!(
            tools.contains_key(".hidden"),
            "dotfile .hidden.wasm is currently discovered (no dotfile filter yet)"
        );
        assert_eq!(tools.len(), 2);
    }

    #[tokio::test]
    async fn test_discover_tools_ignores_subdirectories() {
        let dir = TempDir::new().unwrap();

        // Create a top-level wasm file
        std::fs::File::create(dir.path().join("top_level.wasm")).unwrap();

        // Create a subdirectory with a wasm file inside
        let sub_dir = dir.path().join("subdir");
        std::fs::create_dir(&sub_dir).unwrap();
        std::fs::File::create(sub_dir.join("nested.wasm")).unwrap();

        let tools = discover_tools(dir.path()).await.unwrap();

        // Only top-level files should be discovered (read_dir is not recursive)
        assert_eq!(tools.len(), 1, "Only top-level .wasm files should be found");
        assert!(
            tools.contains_key("top_level"),
            "top_level.wasm should be discovered"
        );
        assert!(
            !tools.contains_key("nested"),
            "nested.wasm inside subdir should NOT be discovered"
        );
    }
}