Vapora/crates/vapora-rlm/src/sandbox/dispatcher.rs

// Sandbox Dispatcher - Auto-tier selection
// Routes commands to WASM (fast) or Docker (compatible) based on complexity

use std::sync::Arc;
use std::time::Instant;

use tracing::{debug, info};

use crate::sandbox::docker_pool::DockerPool;
use crate::sandbox::wasm_runtime::WasmRuntime;
use crate::sandbox::{SandboxCommand, SandboxResult};
use crate::RLMError;

/// Sandbox tier selection
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum SandboxTier {
    /// Tier 1: WASM runtime (fast, <5ms target)
    Wasm,
    /// Tier 2: Docker pool (compatible, 80-150ms target)
    Docker,
}

impl std::fmt::Display for SandboxTier {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            SandboxTier::Wasm => write!(f, "wasm"),
            SandboxTier::Docker => write!(f, "docker"),
        }
    }
}

/// Dispatcher configuration
#[derive(Debug, Clone)]
pub struct DispatcherConfig {
    /// Enable WASM tier
    pub enable_wasm: bool,
    /// Enable Docker tier
    pub enable_docker: bool,
    /// Fallback to Docker if WASM fails
    pub fallback_to_docker: bool,
}

impl Default for DispatcherConfig {
    fn default() -> Self {
        Self {
            enable_wasm: true,
            enable_docker: true,
            fallback_to_docker: true,
        }
    }
}

/// Sandbox dispatcher - routes commands to appropriate tier
pub struct SandboxDispatcher {
    wasm_runtime: Option<Arc<WasmRuntime>>,
    docker_pool: Option<Arc<DockerPool>>,
    config: DispatcherConfig,
}

impl SandboxDispatcher {
    /// Create a new dispatcher with both tiers
    pub async fn new(
        wasm_runtime: Option<Arc<WasmRuntime>>,
        docker_pool: Option<Arc<DockerPool>>,
    ) -> crate::Result<Self> {
        Ok(Self {
            wasm_runtime,
            docker_pool,
            config: DispatcherConfig::default(),
        })
    }

    /// Create with custom configuration
    pub async fn with_config(
        wasm_runtime: Option<Arc<WasmRuntime>>,
        docker_pool: Option<Arc<DockerPool>>,
        config: DispatcherConfig,
    ) -> crate::Result<Self> {
        Ok(Self {
            wasm_runtime,
            docker_pool,
            config,
        })
    }

    /// Execute a command, automatically selecting the appropriate tier
    ///
    /// # Tier Selection Logic
    /// 1. WASI-compatible commands → Tier 1 (WASM) if enabled
    /// 2. Complex commands → Tier 2 (Docker) if enabled
    /// 3. Fallback: Docker if WASM fails and fallback enabled
    ///
    /// # Returns
    /// SandboxResult with tier information
    pub async fn execute(&self, command: &SandboxCommand) -> crate::Result<SandboxResult> {
        let start = Instant::now();

        // Select tier
        let tier = self.select_tier(&command.command);

        debug!(
            "Dispatching command '{}' to {:?} tier",
            command.command, tier
        );

        // Execute in selected tier
        let result = match tier {
            SandboxTier::Wasm => {
                if let Some(ref wasm_runtime) = self.wasm_runtime {
                    match wasm_runtime.execute(command) {
                        Ok(result) => Ok(result),
                        Err(e) if self.config.fallback_to_docker => {
                            info!("WASM execution failed, falling back to Docker: {}", e);
                            self.execute_docker(command).await
                        }
                        Err(e) => Err(e),
                    }
                } else {
                    return Err(RLMError::SandboxError(
                        "WASM tier not available".to_string(),
                    ));
                }
            }
            SandboxTier::Docker => self.execute_docker(command).await,
        }?;

        let duration = start.elapsed();
        debug!(
            "Dispatched command '{}' via {} in {:?}",
            command.command, result.tier, duration
        );

        Ok(result)
    }

    /// Execute in Docker tier
    async fn execute_docker(&self, command: &SandboxCommand) -> crate::Result<SandboxResult> {
        if let Some(ref docker_pool) = self.docker_pool {
            docker_pool.execute(command).await
        } else {
            Err(RLMError::SandboxError(
                "Docker tier not available".to_string(),
            ))
        }
    }

    /// Select tier based on command complexity
    fn select_tier(&self, command: &str) -> SandboxTier {
        // WASI-compatible commands go to WASM tier (if enabled and available)
        if self.config.enable_wasm
            && self.wasm_runtime.is_some()
            && self.is_wasi_compatible(command)
        {
            return SandboxTier::Wasm;
        }

        // Non-WASI commands prefer Docker (if enabled AND available)
        if self.config.enable_docker && self.docker_pool.is_some() {
            return SandboxTier::Docker;
        }

        // Fallback: If Docker enabled but unavailable, use WASM if available
        if self.config.enable_docker && self.docker_pool.is_none() && self.wasm_runtime.is_some() {
            return SandboxTier::Wasm;
        }

        // If WASM enabled and available (for non-WASI commands when Docker not
        // preferred)
        if self.config.enable_wasm && self.wasm_runtime.is_some() {
            return SandboxTier::Wasm;
        }

        // Last resort: Docker (will error on execute if not available)
        SandboxTier::Docker
    }

    /// Check if command is WASI-compatible
    fn is_wasi_compatible(&self, command: &str) -> bool {
        matches!(command, "peek" | "grep" | "slice")
    }

    /// Get tier usage statistics
    pub fn tier_stats(&self) -> TierStats {
        // In a real implementation, would track tier usage in metrics
        // For Phase 4, return basic info
        TierStats {
            wasm_available: self.wasm_runtime.is_some(),
            docker_available: self.docker_pool.is_some(),
            docker_pool_size: self
                .docker_pool
                .as_ref()
                .map(|p| p.pool_size())
                .unwrap_or(0),
        }
    }
}

/// Tier usage statistics
#[derive(Debug, Clone)]
pub struct TierStats {
    pub wasm_available: bool,
    pub docker_available: bool,
    pub docker_pool_size: usize,
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::sandbox::wasm_runtime::WasmRuntime;

    #[tokio::test]
    async fn test_dispatcher_creation() {
        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::new(wasm, None).await.unwrap();

        let stats = dispatcher.tier_stats();
        assert!(stats.wasm_available);
        assert!(!stats.docker_available);
    }

    #[tokio::test]
    async fn test_tier_selection_wasi_compatible() {
        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::new(wasm, None).await.unwrap();

        assert_eq!(dispatcher.select_tier("peek"), SandboxTier::Wasm);
        assert_eq!(dispatcher.select_tier("grep"), SandboxTier::Wasm);
        assert_eq!(dispatcher.select_tier("slice"), SandboxTier::Wasm);
    }

    #[tokio::test]
    async fn test_tier_selection_complex_command() {
        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::new(wasm, None).await.unwrap();

        // Complex commands should prefer Docker (but WASM is selected as fallback if
        // Docker unavailable)
        assert_eq!(dispatcher.select_tier("bash"), SandboxTier::Wasm); // Fallback
    }

    #[tokio::test]
    async fn test_execute_wasm_tier() {
        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::new(wasm, None).await.unwrap();

        let command = SandboxCommand::new("peek")
            .arg("3")
            .stdin("line1\nline2\nline3\nline4");

        let result = dispatcher.execute(&command).await.unwrap();
        assert!(result.is_success());
        assert_eq!(result.tier, SandboxTier::Wasm);
        assert_eq!(result.output, "line1\nline2\nline3\n");
    }

    #[tokio::test]
    async fn test_execute_grep_wasm() {
        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::new(wasm, None).await.unwrap();

        let command = SandboxCommand::new("grep")
            .arg("error")
            .stdin("info: ok\nerror: failed\nwarn: retry");

        let result = dispatcher.execute(&command).await.unwrap();
        assert!(result.is_success());
        assert_eq!(result.tier, SandboxTier::Wasm);
        assert!(result.output.contains("error: failed"));
    }

    #[tokio::test]
    async fn test_wasm_not_available() {
        let dispatcher = SandboxDispatcher::new(None, None).await.unwrap();

        let command = SandboxCommand::new("peek").arg("5").stdin("test");

        let result = dispatcher.execute(&command).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_custom_config() {
        let config = DispatcherConfig {
            enable_wasm: false,
            enable_docker: false,
            fallback_to_docker: false,
        };

        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::with_config(wasm, None, config)
            .await
            .unwrap();

        // With WASM disabled, should select Docker (even though unavailable)
        assert_eq!(dispatcher.select_tier("peek"), SandboxTier::Docker);
    }

    #[tokio::test]
    async fn test_tier_stats() {
        let wasm = Some(Arc::new(WasmRuntime::new()));
        let dispatcher = SandboxDispatcher::new(wasm, None).await.unwrap();

        let stats = dispatcher.tier_stats();
        assert!(stats.wasm_available);
        assert!(!stats.docker_available);
        assert_eq!(stats.docker_pool_size, 0);
    }
}