blinkvm-sdk 0.4.0

//! Shared BoxLite runtime for long-lived server processes.

use std::future::Future;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::time::Duration;

use anyhow::{Context, Result, bail};
use boxlite::{
    BoxCommand, BoxliteRuntime, LiteBox, SnapshotInfo,
    runtime::options::{
        BoxArchive, BoxOptions, ExportOptions, NetworkSpec, RootfsSpec,
        SnapshotOptions, VolumeSpec,
    },
};
use tokio::runtime::Handle;
use tokio::sync::Semaphore;
use tracing::{info, warn};

use crate::boxlite_options::load_boxlite_options;
use crate::exec::exec_agent_script;
use crate::network::resolve_network_spec;
use crate::runner::run_agent_script;
use crate::AgentResult;
use blink_shared::AGENT_MEMORY_DIR;

const DEFAULT_VM_CONCURRENCY: usize = 3;
const DEFAULT_VM_THREADS: usize = 4;
const DEFAULT_CONTROL_TIMEOUT_SECS: u64 = 180;
const DEFAULT_RUN_TIMEOUT_SECS: u64 = 300;

#[derive(Clone, Debug, serde::Deserialize)]
pub struct SessionVolume {
    pub host_path: String,
    pub guest_path: String,
    #[serde(default)]
    pub read_only: bool,
}

impl From<SessionVolume> for VolumeSpec {
    fn from(volume: SessionVolume) -> Self {
        Self {
            host_path: volume.host_path,
            guest_path: volume.guest_path,
            read_only: volume.read_only,
        }
    }
}

/// VM resource limits passed through to BoxLite `BoxOptions`.
///
/// Unset fields use BoxLite defaults (4 CPUs, 4096 MiB RAM, 10 GB disk).
#[derive(Clone, Debug, Default, serde::Deserialize, serde::Serialize)]
pub struct SandboxResources {
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub cpus: Option<u8>,
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub memory_mib: Option<u32>,
    /// Sparse container rootfs virtual size in GB (at least the base image size).
    #[serde(default, skip_serializing_if = "Option::is_none")]
    pub disk_size_gb: Option<u64>,
}

impl SandboxResources {
    pub fn validate(&self) -> Result<()> {
        if let Some(0) = self.cpus {
            bail!("cpus must be >= 1");
        }
        if let Some(0) = self.memory_mib {
            bail!("memory_mib must be >= 1");
        }
        if let Some(0) = self.disk_size_gb {
            bail!("disk_size_gb must be >= 1");
        }
        Ok(())
    }
}

/// Options for creating or reusing a named session.
#[derive(Clone, Debug, Default)]
pub struct OpenSessionOptions {
    pub volumes: Vec<SessionVolume>,
    pub network: Option<boxlite::runtime::options::NetworkConfig>,
    pub resources: SandboxResources,
}

#[derive(Clone)]
pub struct BlinkContext {
    runtime: BoxliteRuntime,
    export_dir: PathBuf,
    /// When `Some`, all VM operations are isolated onto a dedicated runtime's
    /// blocking threads so that synchronous calls inside boxlite (flock, SQLite,
    /// parking_lot locks) can never block HTTP worker threads.
    vm_handle: Option<Handle>,
    /// Bounds the number of concurrently in-flight VM operations (and thus
    /// blocking threads that could be stuck on flock).
    vm_concurrency: Arc<Semaphore>,
    control_timeout: Duration,
    run_timeout: Duration,
}

#[derive(Clone, Debug, serde::Serialize)]
pub struct SessionInfo {
    pub name: Option<String>,
    pub box_id: String,
    pub status: String,
    pub running: bool,
}

impl BlinkContext {
    pub fn new() -> Result<Self> {
        let home = dirs::home_dir().context("could not resolve home directory")?;
        let export_dir = home.join(".blink").join("exports");
        std::fs::create_dir_all(&export_dir).context("failed to create export directory")?;
        Ok(Self {
            runtime: BoxliteRuntime::new(load_boxlite_options()?)
                .context("failed to initialize BoxLite runtime")?,
            export_dir,
            vm_handle: None,
            vm_concurrency: Arc::new(Semaphore::new(vm_concurrency_limit())),
            control_timeout: duration_from_env("BLINK_VM_TIMEOUT_SECS", DEFAULT_CONTROL_TIMEOUT_SECS),
            run_timeout: duration_from_env("BLINK_RUN_TIMEOUT_SECS", DEFAULT_RUN_TIMEOUT_SECS),
        })
    }

    /// Attach a dedicated VM runtime handle (server mode).
    /// When set, all boxlite operations run on blocking threads driven by this
    /// runtime, keeping HTTP workers completely free of synchronous blocking.
    pub fn with_vm_handle(mut self, handle: Handle) -> Self {
        self.vm_handle = Some(handle);
        self
    }

    pub fn export_dir(&self) -> &Path {
        &self.export_dir
    }

    // ── VM operation wrappers ──────────────────────────────────────────

    /// Run a VM operation on a blocking thread driven by the VM runtime.
    /// The semaphore limits concurrency; the timeout bounds total wall time.
    /// flock / SQLite / parking_lot inside boxlite block the blocking thread,
    /// never a tokio worker.
    async fn vm_op<F, T>(&self, timeout: Duration, future: F) -> Result<T>
    where
        F: Future<Output = Result<T>> + Send + 'static,
        T: Send + 'static,
    {
        let handle = self.vm_handle.clone();
        match handle {
            None => {
                match tokio::time::timeout(timeout, future).await {
                    Ok(result) => result,
                    Err(_) => Err(anyhow::anyhow!(
                        "timed out after {}s",
                        timeout.as_secs()
                    )),
                }
            }
            Some(handle) => {
                let semaphore = self.vm_concurrency.clone();
                let inner = async move {
                    let permit = semaphore
                        .acquire_owned()
                        .await
                        .map_err(|_| anyhow::anyhow!("VM semaphore closed"))?;
                    let join = tokio::task::spawn_blocking(move || {
                        let _permit = permit;
                        // Inner timeout: when it fires, the future is dropped,
                        // which drops boxlite's LockGuard -> releases flock.
                        // Without this, builder.build() could hang forever
                        // (some stages have no timeout) and the flock would
                        // never be released.
                        handle.block_on(async move {
                            tokio::time::timeout(timeout, future)
                                .await
                                .map_err(|_| {
                                    anyhow::anyhow!(
                                        "operation timed out after {}s",
                                        timeout.as_secs()
                                    )
                                })?
                        })
                    });
                    join.await.map_err(|e| anyhow::anyhow!("VM task panicked: {e}"))
                };
                match tokio::time::timeout(timeout, inner).await {
                    Ok(result) => result?,
                    Err(_) => Err(anyhow::anyhow!(
                        "VM operation timed out after {}s",
                        timeout.as_secs()
                    )),
                }
            }
        }
    }

    /// Lightweight variant for read-only operations that don't need the
    /// semaphore (e.g. list_checkpoints with synchronous SQLite).
    async fn vm_op_light<F, T>(&self, timeout: Duration, future: F) -> Result<T>
    where
        F: Future<Output = Result<T>> + Send + 'static,
        T: Send + 'static,
    {
        let handle = self.vm_handle.clone();
        match handle {
            None => {
                match tokio::time::timeout(timeout, future).await {
                    Ok(result) => result,
                    Err(_) => Err(anyhow::anyhow!(
                        "timed out after {}s",
                        timeout.as_secs()
                    )),
                }
            }
            Some(handle) => {
                let join = tokio::task::spawn_blocking(move || handle.block_on(future));
                match tokio::time::timeout(timeout, join).await {
                    Ok(result) => result.map_err(|e| anyhow::anyhow!("VM task panicked: {e}"))?,
                    Err(_) => Err(anyhow::anyhow!(
                        "VM operation timed out after {}s",
                        timeout.as_secs()
                    )),
                }
            }
        }
    }

    // ── Internal helpers (called from within vm_op closures) ───────────

    fn session_options(
        image: &str,
        volumes: &[VolumeSpec],
        network: NetworkSpec,
        resources: &SandboxResources,
    ) -> BoxOptions {
        let working_dir = volumes
            .first()
            .map(|volume| volume.guest_path.clone());
        BoxOptions {
            rootfs: RootfsSpec::Image(image.to_string()),
            network,
            auto_remove: false,
            detach: true,
            volumes: volumes.to_vec(),
            working_dir,
            cpus: resources.cpus,
            memory_mib: resources.memory_mib,
            disk_size_gb: resources.disk_size_gb,
            ..Default::default()
        }
    }

    async fn get_box(&self, name: &str) -> Result<LiteBox> {
        self.runtime
            .get(name)
            .await
            .context("failed to lookup session")?
            .with_context(|| format!("session '{name}' not found"))
    }

    // ── Public API ─────────────────────────────────────────────────────

    /// List sessions. Safe on HTTP runtime — boxlite wraps the DB query in
    /// `spawn_blocking` internally.
    pub async fn list_sessions(&self) -> Result<Vec<SessionInfo>> {
        let timeout = self.control_timeout;
        let infos = tokio::time::timeout(timeout, self.runtime.list_info())
            .await
            .map_err(|_| anyhow::anyhow!("list sessions timed out after {}s", timeout.as_secs()))?
            .context("list sessions")?;
        Ok(infos
            .into_iter()
            .map(|b| SessionInfo {
                name: b.name,
                box_id: b.id.to_string(),
                status: format!("{:?}", b.status),
                running: b.pid.is_some(),
            })
            .collect())
    }

    pub async fn open_session(
        &self,
        name: &str,
        image: &str,
        warm: bool,
        options: OpenSessionOptions,
    ) -> Result<(String, bool)> {
        let ctx = self.clone();
        let name = name.to_string();
        let image = image.to_string();
        self.vm_op(self.control_timeout, async move {
            options.resources.validate()?;
            info!(
                name = %name,
                image = %image,
                warm,
                volume_count = options.volumes.len(),
                cpus = ?options.resources.cpus,
                memory_mib = ?options.resources.memory_mib,
                disk_size_gb = ?options.resources.disk_size_gb,
                "opening session"
            );
            let network_spec = resolve_network_spec(options.network)?;
            let volume_specs: Vec<VolumeSpec> =
                options.volumes.into_iter().map(VolumeSpec::from).collect();
            let (litebox, created) = ctx
                .runtime
                .get_or_create(
                    Self::session_options(&image, &volume_specs, network_spec, &options.resources),
                    Some(name.clone()),
                )
                .await
                .context("open session")?;
            ensure_memory_dir(&litebox).await?;
            Ok((litebox.id().as_str().to_string(), created))
        })
        .await
    }

    pub async fn run_agent_ephemeral(
        &self,
        script_path: &Path,
        image: Option<&str>,
        resources: SandboxResources,
    ) -> Result<AgentResult> {
        let script_path = script_path.to_path_buf();
        let image = image
            .map(String::from)
            .unwrap_or_else(|| blink_shared::DEFAULT_ROOTFS_IMAGE.to_string());
        self.vm_op(self.run_timeout, async move {
            run_agent_script(&script_path, &image, resources).await
        })
        .await
    }

    pub async fn run_in_session(&self, name: &str, script_path: &Path) -> Result<AgentResult> {
        let ctx = self.clone();
        let name = name.to_string();
        let script_path = script_path.to_path_buf();
        self.vm_op(self.run_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            exec_agent_script(&litebox, &script_path).await
        })
        .await
    }

    pub async fn spawn_in_session(
        &self,
        name: &str,
        spec: crate::SpawnSpec,
    ) -> Result<crate::Execution> {
        let ctx = self.clone();
        let name = name.to_string();
        self.vm_op(self.control_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            crate::spawn_exec(&litebox, spec).await
        })
        .await
    }

    pub async fn checkpoint_session(&self, name: &str, snapshot: &str) -> Result<SnapshotInfo> {
        let ctx = self.clone();
        let name = name.to_string();
        let snapshot = snapshot.to_string();
        self.vm_op(self.run_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            litebox
                .snapshots()
                .create(SnapshotOptions::default(), &snapshot)
                .await
                .context("checkpoint")
        })
        .await
    }

    pub async fn restore_session(&self, name: &str, snapshot: &str) -> Result<()> {
        let ctx = self.clone();
        let name = name.to_string();
        let snapshot = snapshot.to_string();
        self.vm_op(self.run_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            if litebox.info().status.is_active() {
                litebox.stop().await.context("stop before restore")?;
            }
            litebox
                .snapshots()
                .restore(&snapshot)
                .await
                .context("restore")
        })
        .await
    }

    /// List checkpoints. Uses `vm_op_light` because boxlite's
    /// `snapshot_mgr.list()` does synchronous SQLite without spawn_blocking.
    pub async fn list_checkpoints(&self, name: &str) -> Result<Vec<SnapshotInfo>> {
        let ctx = self.clone();
        let name = name.to_string();
        self.vm_op_light(self.control_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            litebox
                .snapshots()
                .list()
                .await
                .context("list checkpoints")
        })
        .await
    }

    pub async fn stop_session(&self, name: &str) -> Result<()> {
        let ctx = self.clone();
        let name = name.to_string();
        self.vm_op(self.control_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            if litebox.info().status.is_active() {
                litebox.stop().await.context("stop session")?;
            }
            Ok(())
        })
        .await
    }

    pub async fn remove_session(&self, name: &str, force: bool) -> Result<()> {
        let ctx = self.clone();
        let name = name.to_string();
        self.vm_op(self.control_timeout, async move {
            if let Ok(litebox) = ctx.get_box(&name).await {
                if litebox.info().status.is_active() {
                    match litebox.stop().await {
                        Ok(()) => {}
                        Err(err) if !force => {
                            return Err(err).context("stop session before remove");
                        }
                        Err(err) => {
                            warn!(
                                name = %name,
                                error = %err,
                                "stop before remove failed; continuing with force"
                            );
                        }
                    }
                }
            }
            ctx.runtime
                .remove(&name, force)
                .await
                .context("remove session")
        })
        .await
    }

    pub async fn export_session(&self, name: &str) -> Result<PathBuf> {
        let ctx = self.clone();
        let name = name.to_string();
        self.vm_op(self.run_timeout, async move {
            let litebox = ctx.get_box(&name).await?;
            let stamp = chrono::Utc::now().format("%Y%m%d-%H%M%S");
            let dest = ctx
                .export_dir
                .join(format!("{name}-{stamp}.boxlite"));
            let archive = litebox
                .export(ExportOptions::default(), &dest)
                .await
                .context("export session")?;
            Ok(archive.path().to_path_buf())
        })
        .await
    }

    pub async fn import_session(&self, archive_path: &Path, name: Option<&str>) -> Result<String> {
        if !archive_path.exists() {
            bail!("archive not found: {}", archive_path.display());
        }
        let ctx = self.clone();
        let archive = BoxArchive::new(archive_path);
        let name_owned = name.map(String::from);
        self.vm_op(self.run_timeout, async move {
            let litebox = ctx
                .runtime
                .import_box(archive, name_owned)
                .await
                .context("import session")?;
            ensure_memory_dir(&litebox).await?;
            Ok(litebox.id().as_str().to_string())
        })
        .await
    }
}

async fn ensure_memory_dir(litebox: &LiteBox) -> Result<()> {
    let cmd = format!("mkdir -p {AGENT_MEMORY_DIR}");
    for attempt in 0..3 {
        let execution = match litebox
            .exec(BoxCommand::new("sh").arg("-c").arg(&cmd))
            .await
        {
            Ok(e) => e,
            Err(e) => {
                if attempt < 2 {
                    tokio::time::sleep(std::time::Duration::from_millis(500)).await;
                    continue;
                }
                return Err(e).context("mkdir memory dir");
            }
        };
        let status = execution.wait().await.context("mkdir wait")?;
        if status.exit_code == 0 {
            return Ok(());
        }
        if attempt < 2 {
            tokio::time::sleep(std::time::Duration::from_millis(500)).await;
        }
    }
    warn!(exit_code = "unknown", "mkdir {AGENT_MEMORY_DIR} failed after 3 attempts, agent may need to create it");
    Ok(())
}

fn vm_concurrency_limit() -> usize {
    std::env::var("BLINK_VM_CONCURRENCY")
        .ok()
        .and_then(|v| v.trim().parse::<usize>().ok())
        .filter(|v| *v > 0)
        .unwrap_or(DEFAULT_VM_CONCURRENCY)
}

pub fn vm_thread_count() -> usize {
    std::env::var("BLINK_VM_THREADS")
        .ok()
        .and_then(|v| v.trim().parse::<usize>().ok())
        .filter(|v| *v > 0)
        .unwrap_or(DEFAULT_VM_THREADS)
}

fn duration_from_env(var: &str, default_secs: u64) -> Duration {
    std::env::var(var)
        .ok()
        .and_then(|v| v.trim().parse::<u64>().ok())
        .filter(|v| *v > 0)
        .map(Duration::from_secs)
        .unwrap_or_else(|| Duration::from_secs(default_secs))
}