clawmates/crates/cm-db/src/repo/agent_containers.rs

//! Durable registry of the live container backing each agent, so any server
//! replica can find + reuse it (replacing the former in-process handle map).

use cm_domain::{AgentId, WorkspaceId};
use sqlx::{PgPool, Row};

use crate::DbError;

/// A registry row for one agent's container of a given kind.
#[derive(Debug, Clone)]
pub struct ContainerRow {
    pub node_id: String,
    pub container_id: String,
    pub name: String,
    pub session_count: i32,
}

/// One container to act on (idle reap / reconcile / shutdown).
#[derive(Debug, Clone)]
pub struct ManagedRow {
    pub agent_id: AgentId,
    pub node_id: String,
    pub container_id: String,
    pub name: String,
}

pub async fn get(
    pool: &PgPool,
    agent_id: AgentId,
    kind: &str,
) -> Result<Option<ContainerRow>, DbError> {
    let row = sqlx::query(
        "SELECT node_id, container_id, name, session_count
         FROM agent_containers WHERE agent_id = $1 AND kind = $2",
    )
    .bind(agent_id.as_uuid())
    .bind(kind)
    .fetch_optional(pool)
    .await?;
    Ok(row.map(|r| ContainerRow {
        node_id: r.get("node_id"),
        container_id: r.get("container_id"),
        name: r.get("name"),
        session_count: r.get("session_count"),
    }))
}

/// Record the (new) container for `(agent, kind)`; resets the session count.
/// The workspace is derived from the agent row (single source of truth).
pub async fn upsert(
    pool: &PgPool,
    agent_id: AgentId,
    kind: &str,
    node_id: &str,
    container_id: &str,
    name: &str,
) -> Result<(), DbError> {
    sqlx::query(
        "INSERT INTO agent_containers
            (agent_id, kind, node_id, container_id, name, workspace_id, session_count, last_seen)
         SELECT $1, $2, $3, $4, $5, a.workspace_id, 0, now() FROM agents a WHERE a.id = $1
         ON CONFLICT (agent_id, kind) DO UPDATE SET
            node_id = excluded.node_id, container_id = excluded.container_id,
            name = excluded.name, workspace_id = excluded.workspace_id,
            session_count = 0, last_seen = now()",
    )
    .bind(agent_id.as_uuid())
    .bind(kind)
    .bind(node_id)
    .bind(container_id)
    .bind(name)
    .execute(pool)
    .await?;
    Ok(())
}

pub async fn delete(pool: &PgPool, agent_id: AgentId, kind: &str) -> Result<(), DbError> {
    sqlx::query("DELETE FROM agent_containers WHERE agent_id = $1 AND kind = $2")
        .bind(agent_id.as_uuid())
        .bind(kind)
        .execute(pool)
        .await?;
    Ok(())
}

/// Adjust the live-session count (clamped at 0) and bump last_seen.
pub async fn add_session(
    pool: &PgPool,
    agent_id: AgentId,
    kind: &str,
    delta: i32,
) -> Result<(), DbError> {
    sqlx::query(
        "UPDATE agent_containers
         SET session_count = GREATEST(0, session_count + $3), last_seen = now()
         WHERE agent_id = $1 AND kind = $2",
    )
    .bind(agent_id.as_uuid())
    .bind(kind)
    .bind(delta)
    .execute(pool)
    .await?;
    Ok(())
}

/// Zero the live-session counts for a kind (graceful shutdown: the WS sessions
/// are ending, so the containers become idle-reapable but are not destroyed).
pub async fn reset_sessions(pool: &PgPool, kind: &str) -> Result<(), DbError> {
    sqlx::query("UPDATE agent_containers SET session_count = 0 WHERE kind = $1")
        .bind(kind)
        .execute(pool)
        .await?;
    Ok(())
}

/// Session-less containers untouched for longer than `idle_secs` (0 = all).
pub async fn idle(pool: &PgPool, kind: &str, idle_secs: i64) -> Result<Vec<ManagedRow>, DbError> {
    let rows = sqlx::query(
        "SELECT agent_id, node_id, container_id, name FROM agent_containers
         WHERE kind = $1 AND session_count <= 0
           AND last_seen < now() - ($2 * interval '1 second')",
    )
    .bind(kind)
    .bind(idle_secs)
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(map_managed).collect())
}

/// Every recorded container of a kind (for reconcile + shutdown).
pub async fn all(pool: &PgPool, kind: &str) -> Result<Vec<ManagedRow>, DbError> {
    let rows = sqlx::query(
        "SELECT agent_id, node_id, container_id, name FROM agent_containers WHERE kind = $1",
    )
    .bind(kind)
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(map_managed).collect())
}

/// Count live containers in a workspace (any kind) — for quota checks.
pub async fn count_for_workspace(pool: &PgPool, workspace_id: WorkspaceId) -> Result<i64, DbError> {
    let row = sqlx::query("SELECT count(*) AS n FROM agent_containers WHERE workspace_id = $1")
        .bind(workspace_id.as_uuid())
        .fetch_one(pool)
        .await?;
    Ok(row.get::<i64, _>("n"))
}

fn map_managed(r: sqlx::postgres::PgRow) -> ManagedRow {
    ManagedRow {
        agent_id: AgentId::from(r.get::<uuid::Uuid, _>("agent_id")),
        node_id: r.get("node_id"),
        container_id: r.get("container_id"),
        name: r.get("name"),
    }
}