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

use cm_domain::{
    AccessPolicy, Agent, AgentId, AgentScope, AgentStatus, HumanScope, UserId, WorkspaceId,
};
use sqlx::PgPool;
use uuid::Uuid;

use crate::DbError;

/// Count of non-deleted agents in a workspace (for per-workspace quotas).
pub async fn count_active(pool: &PgPool, workspace_id: WorkspaceId) -> Result<i64, DbError> {
    let n: i64 = sqlx::query_scalar(
        "SELECT count(*) FROM agents WHERE workspace_id = $1 AND deleted_at IS NULL",
    )
    .bind(workspace_id.as_uuid())
    .fetch_one(pool)
    .await?;
    Ok(n)
}

/// Inserts an agent together with its access policy in one transaction —
/// an agent without a policy must never be observable (§7.7).
pub async fn insert(pool: &PgPool, agent: &Agent, policy: &AccessPolicy) -> Result<(), DbError> {
    let (humans_mode, human_ids): (&str, Vec<Uuid>) = match &policy.humans {
        HumanScope::EntireTeam => ("entire_team", Vec::new()),
        HumanScope::Specific(ids) => ("specific", ids.iter().map(UserId::as_uuid).collect()),
    };
    let (agents_mode, agent_ids): (&str, Vec<Uuid>) = match &policy.agents {
        AgentScope::Any => ("any", Vec::new()),
        AgentScope::Specific(ids) => ("specific", ids.iter().map(AgentId::as_uuid).collect()),
    };

    let mut tx = pool.begin().await.map_err(DbError::from)?;
    sqlx::query!(
        "INSERT INTO agents
            (id, workspace_id, name, job_title, system_prompt, avatar, accent,
             wallpaper, managed_by, status)
         VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10)",
        agent.id.as_uuid(),
        agent.workspace_id.as_uuid(),
        agent.name,
        agent.job_title,
        agent.system_prompt,
        agent.avatar,
        agent.accent,
        agent.wallpaper,
        agent.managed_by.as_uuid(),
        agent.status.as_str(),
    )
    .execute(&mut *tx)
    .await?;
    sqlx::query!(
        "INSERT INTO access_policies
            (agent_id, humans_mode, human_ids, agents_mode, agent_ids)
         VALUES ($1, $2, $3, $4, $5)",
        agent.id.as_uuid(),
        humans_mode,
        &human_ids,
        agents_mode,
        &agent_ids,
    )
    .execute(&mut *tx)
    .await?;
    tx.commit().await.map_err(DbError::from)?;
    Ok(())
}

/// Fetch an agent even if it has been soft-deleted.
///
/// For the PURGE path only. `get` hides soft-deleted rows, which is right for
/// every read — but it also meant the hard purge could not see the rows it
/// exists to remove: a soft-deleted agent was unreachable from every route and
/// accumulated forever with no way out of the application. Four of them dated
/// from June before anyone noticed, because the UI correctly never showed them.
pub async fn get_any(pool: &PgPool, agent_id: AgentId) -> Result<Agent, DbError> {
    // `sqlx::query_as` rather than the checked macro: this is the same columns
    // as `get` minus one predicate, and adding a second compile-time query for
    // that would mean regenerating the offline cache on every machine that
    // builds this.
    let row: Option<(
        uuid::Uuid,
        uuid::Uuid,
        String,
        String,
        String,
        String,
        String,
        String,
        uuid::Uuid,
        String,
    )> = sqlx::query_as(
        "SELECT id, workspace_id, name, job_title, system_prompt, avatar,
                accent, wallpaper, managed_by, status
         FROM agents WHERE id = $1",
    )
    .bind(agent_id.as_uuid())
    .fetch_optional(pool)
    .await?;
    let row = row.ok_or(DbError::NotFound)?;
    Ok(Agent {
        id: AgentId::from(row.0),
        workspace_id: WorkspaceId::from(row.1),
        name: row.2,
        job_title: row.3,
        system_prompt: row.4,
        avatar: row.5,
        accent: row.6,
        wallpaper: row.7,
        managed_by: UserId::from(row.8),
        status: row.9.parse().expect("status CHECK constraint"),
    })
}

pub async fn get(pool: &PgPool, agent_id: AgentId) -> Result<Agent, DbError> {
    let row = sqlx::query!(
        "SELECT id, workspace_id, name, job_title, system_prompt, avatar,
                accent, wallpaper, managed_by, status
         FROM agents WHERE id = $1 AND deleted_at IS NULL",
        agent_id.as_uuid(),
    )
    .fetch_one(pool)
    .await?;
    Ok(Agent {
        id: AgentId::from(row.id),
        workspace_id: WorkspaceId::from(row.workspace_id),
        name: row.name,
        job_title: row.job_title,
        system_prompt: row.system_prompt,
        avatar: row.avatar,
        accent: row.accent,
        wallpaper: row.wallpaper,
        managed_by: UserId::from(row.managed_by),
        status: row.status.parse().expect("status CHECK constraint"),
    })
}

/// Persist the model a claw was deployed with (e.g. "claude", "glm",
/// "glm-5.2") — the runtime config is otherwise the only record of it.
pub async fn set_model_binding(
    pool: &PgPool,
    agent_id: AgentId,
    model: &str,
) -> Result<(), DbError> {
    sqlx::query!(
        "UPDATE agents SET model_binding = $2 WHERE id = $1 AND deleted_at IS NULL",
        agent_id.as_uuid(),
        model,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// The persisted model binding for a claw, if any (NULL for pre-team claws).
pub async fn model_binding(pool: &PgPool, agent_id: AgentId) -> Result<Option<String>, DbError> {
    let row = sqlx::query!(
        "SELECT model_binding FROM agents WHERE id = $1 AND deleted_at IS NULL",
        agent_id.as_uuid(),
    )
    .fetch_one(pool)
    .await?;
    Ok(row.model_binding)
}

/// Patch-style profile update (§7.7 Edit profile): only provided fields
/// change; the system prompt is the Job Description textarea verbatim.
#[allow(clippy::too_many_arguments)]
pub async fn update_profile(
    pool: &PgPool,
    agent_id: AgentId,
    name: Option<&str>,
    job_title: Option<&str>,
    system_prompt: Option<&str>,
    avatar: Option<&str>,
    accent: Option<&str>,
    wallpaper: Option<&str>,
) -> Result<Agent, DbError> {
    let result = sqlx::query!(
        "UPDATE agents SET
            name = COALESCE($2, name),
            job_title = COALESCE($3, job_title),
            system_prompt = COALESCE($4, system_prompt),
            avatar = COALESCE($5, avatar),
            accent = COALESCE($6, accent),
            wallpaper = COALESCE($7, wallpaper)
         WHERE id = $1 AND deleted_at IS NULL",
        agent_id.as_uuid(),
        name,
        job_title,
        system_prompt,
        avatar,
        accent,
        wallpaper,
    )
    .execute(pool)
    .await?;
    if result.rows_affected() == 0 {
        return Err(DbError::NotFound);
    }
    get(pool, agent_id).await
}

/// The left-rail roster (§4): live agents of a workspace, oldest first.
pub async fn roster(pool: &PgPool, workspace_id: WorkspaceId) -> Result<Vec<Agent>, DbError> {
    let rows = sqlx::query!(
        "SELECT id, workspace_id, name, job_title, system_prompt, avatar,
                accent, wallpaper, managed_by, status
         FROM agents
         WHERE workspace_id = $1 AND deleted_at IS NULL
         ORDER BY created_at, id",
        workspace_id.as_uuid(),
    )
    .fetch_all(pool)
    .await?;
    Ok(rows
        .into_iter()
        .map(|row| Agent {
            id: AgentId::from(row.id),
            workspace_id: WorkspaceId::from(row.workspace_id),
            name: row.name,
            job_title: row.job_title,
            system_prompt: row.system_prompt,
            avatar: row.avatar,
            accent: row.accent,
            wallpaper: row.wallpaper,
            managed_by: UserId::from(row.managed_by),
            status: row.status.parse().expect("status CHECK constraint"),
        })
        .collect())
}

/// Replaces an agent's access policy (§7.7 access toggles).
pub async fn set_access_policy(
    pool: &PgPool,
    agent_id: AgentId,
    policy: &AccessPolicy,
) -> Result<(), DbError> {
    let (humans_mode, human_ids): (&str, Vec<Uuid>) = match &policy.humans {
        HumanScope::EntireTeam => ("entire_team", Vec::new()),
        HumanScope::Specific(ids) => ("specific", ids.iter().map(UserId::as_uuid).collect()),
    };
    let (agents_mode, agent_ids): (&str, Vec<Uuid>) = match &policy.agents {
        AgentScope::Any => ("any", Vec::new()),
        AgentScope::Specific(ids) => ("specific", ids.iter().map(AgentId::as_uuid).collect()),
    };
    let result = sqlx::query!(
        "UPDATE access_policies SET humans_mode = $2, human_ids = $3,
                agents_mode = $4, agent_ids = $5
         WHERE agent_id = $1",
        agent_id.as_uuid(),
        humans_mode,
        &human_ids,
        agents_mode,
        &agent_ids,
    )
    .execute(pool)
    .await?;
    if result.rows_affected() == 0 {
        return Err(DbError::NotFound);
    }
    Ok(())
}

pub async fn access_policy(pool: &PgPool, agent_id: AgentId) -> Result<AccessPolicy, DbError> {
    let row = sqlx::query!(
        "SELECT humans_mode, human_ids, agents_mode, agent_ids
         FROM access_policies WHERE agent_id = $1",
        agent_id.as_uuid(),
    )
    .fetch_one(pool)
    .await?;
    let humans = if row.humans_mode == "entire_team" {
        HumanScope::EntireTeam
    } else {
        HumanScope::Specific(row.human_ids.into_iter().map(UserId::from).collect())
    };
    let agents = if row.agents_mode == "any" {
        AgentScope::Any
    } else {
        AgentScope::Specific(row.agent_ids.into_iter().map(AgentId::from).collect())
    };
    Ok(AccessPolicy { humans, agents })
}

pub async fn set_status(
    pool: &PgPool,
    agent_id: AgentId,
    status: AgentStatus,
) -> Result<(), DbError> {
    let result = sqlx::query!(
        "UPDATE agents SET status = $2 WHERE id = $1 AND deleted_at IS NULL",
        agent_id.as_uuid(),
        status.as_str(),
    )
    .execute(pool)
    .await?;
    if result.rows_affected() == 0 {
        return Err(DbError::NotFound);
    }
    Ok(())
}

/// Deleting a claw is destructive and gated (§7.7); rows are kept for audit.
pub async fn soft_delete(pool: &PgPool, agent_id: AgentId) -> Result<(), DbError> {
    let result = sqlx::query!(
        "UPDATE agents SET deleted_at = now(), status = 'offline'
         WHERE id = $1 AND deleted_at IS NULL",
        agent_id.as_uuid(),
    )
    .execute(pool)
    .await?;
    if result.rows_affected() == 0 {
        return Err(DbError::NotFound);
    }
    Ok(())
}

/// Counts of the rows reaped by [`hard_purge`], for the progress summary.
#[derive(Debug, Default, Clone, Copy)]
pub struct PurgeCounts {
    pub sessions: u64,
    pub approvals: u64,
    pub connections: u64,
    pub files: u64,
}

/// Hard-delete an agent and everything that references it, in FK-dependency
/// order, in one transaction. Auto-CASCADE handles access_policies,
/// installed_skills, routines(+routine_runs) and team_members; the non-cascade
/// references (chat history, approvals, threads, connections, queued mail, file
/// drives, usage) are cleared first so the final `DELETE FROM agents` succeeds.
/// Only the immutable `audit_log` survives. Returns NotFound if the agent is gone.
pub async fn hard_purge(pool: &PgPool, agent_id: AgentId) -> Result<PurgeCounts, DbError> {
    let aid = agent_id.as_uuid();
    let mut tx = pool.begin().await?;
    let mut c = PurgeCounts::default();

    // Chat history: steps → messages → agent_runs → sessions (none cascade).
    sqlx::query(
        "DELETE FROM steps WHERE message_id IN \
         (SELECT m.id FROM messages m JOIN sessions s ON m.session_id = s.id WHERE s.agent_id = $1)",
    )
    .bind(aid)
    .execute(&mut *tx)
    .await?;
    sqlx::query(
        "DELETE FROM messages WHERE session_id IN (SELECT id FROM sessions WHERE agent_id = $1)",
    )
    .bind(aid)
    .execute(&mut *tx)
    .await?;
    sqlx::query(
        "DELETE FROM agent_runs WHERE session_id IN (SELECT id FROM sessions WHERE agent_id = $1)",
    )
    .bind(aid)
    .execute(&mut *tx)
    .await?;
    c.sessions = sqlx::query("DELETE FROM sessions WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?
        .rows_affected();

    // Approvals: execution_grants → approvals.
    sqlx::query(
        "DELETE FROM execution_grants WHERE approval_id IN \
         (SELECT id FROM approvals WHERE requested_by_agent = $1)",
    )
    .bind(aid)
    .execute(&mut *tx)
    .await?;
    c.approvals = sqlx::query("DELETE FROM approvals WHERE requested_by_agent = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?
        .rows_affected();

    // Inter-agent threads, queued mail, oauth flows.
    sqlx::query("DELETE FROM thread_messages WHERE from_agent = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?;
    sqlx::query("DELETE FROM thread_participants WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?;
    sqlx::query("DELETE FROM outbox WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?;
    sqlx::query("DELETE FROM oauth_states WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?;
    c.connections = sqlx::query("DELETE FROM app_connections WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?
        .rows_affected();
    c.files = sqlx::query("DELETE FROM file_nodes WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?
        .rows_affected();
    sqlx::query("DELETE FROM usage_events WHERE agent_id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?;

    // Finally the agent itself (cascades the rest).
    let n = sqlx::query("DELETE FROM agents WHERE id = $1")
        .bind(aid)
        .execute(&mut *tx)
        .await?
        .rows_affected();
    if n == 0 {
        return Err(DbError::NotFound);
    }
    tx.commit().await?;
    Ok(c)
}