clawmates/crates/cm-runtime/tests/warm_pool.rs

//! Warm sandbox pool: pre-provisioned containers absorb the first-exec
//! latency. Real Docker — the pool fills in the background, an exec
//! takes a sandbox from it, and the warmer restores the target.

use std::process::Command;
use std::sync::Arc;
use std::time::Duration;

use cm_domain::{
    AccessPolicy, Agent, AgentId, AgentStatus, Role, User, UserId, Workspace, WorkspaceId,
};
use cm_runtime::SandboxManager;
use cm_sandbox::DockerDriver;
use sqlx::PgPool;

const IMAGE: &str = "clawmates/agent-base:dev";

fn ensure_image() {
    let exists = Command::new("docker")
        .args(["image", "inspect", IMAGE])
        .output()
        .expect("docker available")
        .status
        .success();
    if !exists {
        let root = env!("CARGO_MANIFEST_DIR");
        let status = Command::new("docker")
            .args([
                "build",
                "-t",
                IMAGE,
                "-f",
                &format!("{root}/../../images/agent-base/Dockerfile"),
                &format!("{root}/../../images/agent-base"),
            ])
            .status()
            .expect("docker build runs");
        assert!(status.success());
    }
}

async fn pool_reaches(manager: &SandboxManager, target: usize) {
    for _ in 0..120 {
        if manager.pool_size().await == target {
            return;
        }
        tokio::time::sleep(Duration::from_millis(250)).await;
    }
    panic!(
        "pool never reached {target} (now {})",
        manager.pool_size().await
    );
}

/// Inserts a workspace + owner user + agent so that
/// `agent_containers::upsert(agent, ...)` (which uses INSERT ... FROM agents
/// WHERE a.id = $1) actually persists a row. Without this seeding the upsert
/// silently inserts zero rows, the next exec's reuse lookup returns None,
/// and the "assigned" sandbox is re-provisioned from the warm pool —
/// draining it and racing the warmer's refill (this test's classic flake).
async fn seed_agent(pool: &PgPool) -> AgentId {
    let ws = Workspace {
        id: WorkspaceId::new(),
        name: "WarmPool".into(),
        plan: "team".into(),
    };
    cm_db::repo::workspaces::insert(pool, &ws).await.unwrap();
    let owner = User {
        id: UserId::new(),
        workspace_id: ws.id,
        email: format!("{}@warmpool.test", UserId::new()),
        role: Role::Owner,
        display_name: "Owner".into(),
        created_at: time::OffsetDateTime::UNIX_EPOCH,
    };
    cm_db::repo::users::insert(pool, &owner).await.unwrap();
    let agent = Agent {
        id: AgentId::new(),
        workspace_id: ws.id,
        name: "Warm".into(),
        job_title: "Analyst".into(),
        system_prompt: String::new(),
        avatar: String::new(),
        accent: String::new(),
        wallpaper: String::new(),
        managed_by: owner.id,
        status: AgentStatus::Online,
    };
    cm_db::repo::agents::insert(pool, &agent, &AccessPolicy::default())
        .await
        .unwrap();
    agent.id
}

#[tokio::test]
async fn the_pool_prefills_assigns_and_refills() {
    ensure_image();
    let driver: Arc<dyn cm_sandbox::SandboxDriver> =
        Arc::new(DockerDriver::connect().expect("docker reachable"));
    let pool = cm_testkit::test_pool().await;
    let agent = seed_agent(&pool).await;
    let manager = Arc::new(SandboxManager::new(driver, pool, "local", IMAGE)).warm(2);

    // The warmer fills the pool without any exec happening.
    pool_reaches(&manager, 2).await;

    // An exec is served from the pool — and works.
    let result = manager.exec(agent, "id -u").await.unwrap();
    assert_eq!(result.stdout.trim(), "10001");

    // The warmer restores the target while the agent keeps its sandbox.
    pool_reaches(&manager, 2).await;
    let again = manager.exec(agent, "echo still-mine").await.unwrap();
    assert_eq!(again.stdout.trim(), "still-mine");
    assert_eq!(
        manager.pool_size().await,
        2,
        "reuse must not drain the pool"
    );

    // `shutdown` tears down the POOL only — assigned sandboxes deliberately
    // survive it, so they can be reused across a redeploy. The comment here
    // used to claim it destroyed both, which is why nobody noticed that every
    // run of this test left its assigned container running: three such tests,
    // three leaked containers per `./scripts/test.sh`, 289 of them by the time
    // anyone counted.
    manager.release_agent(agent).await;
    manager.shutdown().await;
    assert_eq!(manager.pool_size().await, 0);
    assert!(
        !manager.release_agent(agent).await,
        "the agent's sandbox must be gone, not merely unpooled"
    );
}