clawmates/crates/cm-safety/src/checkpoint.rs

use sqlx::PgPool;
use uuid::Uuid;

use crate::SafetyError;

/// Suspends a run: persists the full serialized loop state and flips the
/// run to `awaiting_approval` atomically.
pub async fn suspend(
    pool: &PgPool,
    run_id: Uuid,
    state: &serde_json::Value,
) -> Result<(), SafetyError> {
    let result = sqlx::query!(
        "UPDATE agent_runs
         SET checkpoint = $2, state = 'awaiting_approval', updated_at = now()
         WHERE id = $1",
        run_id,
        state,
    )
    .execute(pool)
    .await?;
    if result.rows_affected() == 0 {
        return Err(SafetyError::NotFound);
    }
    Ok(())
}

pub async fn load(pool: &PgPool, run_id: Uuid) -> Result<serde_json::Value, SafetyError> {
    let row = sqlx::query!("SELECT checkpoint FROM agent_runs WHERE id = $1", run_id)
        .fetch_optional(pool)
        .await?
        .ok_or(SafetyError::NotFound)?;
    row.checkpoint.ok_or(SafetyError::NotFound)
}

/// Claims a suspended run for resumption. The CAS from `awaiting_approval`
/// to `running` ensures exactly one resumer proceeds.
pub async fn claim_resume(pool: &PgPool, run_id: Uuid) -> Result<bool, SafetyError> {
    let result = sqlx::query!(
        "UPDATE agent_runs SET state = 'running', updated_at = now()
         WHERE id = $1 AND state = 'awaiting_approval'",
        run_id,
    )
    .execute(pool)
    .await?;
    Ok(result.rows_affected() == 1)
}