//! 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, 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, 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, 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 { 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::("n")) } fn map_managed(r: sqlx::postgres::PgRow) -> ManagedRow { ManagedRow { agent_id: AgentId::from(r.get::("agent_id")), node_id: r.get("node_id"), container_id: r.get("container_id"), name: r.get("name"), } }