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

use cm_domain::AgentId;
use sqlx::PgPool;
use time::OffsetDateTime;
use uuid::Uuid;

use crate::DbError;

/// A scheduled task an agent runs on a cron cadence (§7.6).
#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
pub struct Routine {
    pub id: Uuid,
    pub agent_id: Uuid,
    pub name: String,
    pub schedule_cron: String,
    /// `{"message": "..."}` — the prompt sent into the routine's session.
    pub action: serde_json::Value,
    pub status: String,
    #[serde(with = "time::serde::rfc3339::option")]
    pub next_run_at: Option<OffsetDateTime>,
    #[serde(with = "time::serde::rfc3339::option")]
    pub last_run_at: Option<OffsetDateTime>,
}

pub async fn create(
    pool: &PgPool,
    agent_id: AgentId,
    name: &str,
    schedule_cron: &str,
    action: serde_json::Value,
    next_run_at: OffsetDateTime,
) -> Result<Routine, DbError> {
    let id = Uuid::now_v7();
    sqlx::query!(
        "INSERT INTO routines (id, agent_id, name, schedule_cron, action, next_run_at)
         VALUES ($1, $2, $3, $4, $5, $6)",
        id,
        agent_id.as_uuid(),
        name,
        schedule_cron,
        action,
        next_run_at,
    )
    .execute(pool)
    .await?;
    Ok(Routine {
        id,
        agent_id: agent_id.as_uuid(),
        name: name.to_owned(),
        schedule_cron: schedule_cron.to_owned(),
        action,
        status: "active".into(),
        next_run_at: Some(next_run_at),
        last_run_at: None,
    })
}

pub async fn list_by_agent(pool: &PgPool, agent_id: AgentId) -> Result<Vec<Routine>, DbError> {
    let rows = sqlx::query_as!(
        Routine,
        r#"SELECT id, agent_id, name, schedule_cron, action, status,
                  next_run_at, last_run_at
           FROM routines WHERE agent_id = $1 ORDER BY created_at"#,
        agent_id.as_uuid(),
    )
    .fetch_all(pool)
    .await?;
    Ok(rows)
}

/// Claims every due routine atomically (SKIP LOCKED: one firing per
/// routine even with multiple scheduler replicas) and advances its clock.
/// `next_runs` computes the following occurrence per claimed routine.
pub async fn claim_due(pool: &PgPool, now: OffsetDateTime) -> Result<Vec<Routine>, DbError> {
    let rows = sqlx::query_as!(
        Routine,
        r#"UPDATE routines SET last_run_at = $1
           WHERE id IN (
               SELECT id FROM routines
               WHERE status = 'active' AND next_run_at IS NOT NULL
                 AND next_run_at <= $1
               FOR UPDATE SKIP LOCKED
           )
           RETURNING id, agent_id, name, schedule_cron, action, status,
                     next_run_at, last_run_at"#,
        now,
    )
    .fetch_all(pool)
    .await?;
    Ok(rows)
}

/// Schedules the next firing after a claim.
pub async fn set_next_run(
    pool: &PgPool,
    id: Uuid,
    next_run_at: Option<OffsetDateTime>,
) -> Result<(), DbError> {
    sqlx::query!(
        "UPDATE routines SET next_run_at = $2 WHERE id = $1",
        id,
        next_run_at,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// What claiming an occurrence found.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum FireClaim {
    /// Nobody has taken this occurrence. Fire it.
    Fresh,
    /// A previous attempt took it and never recorded an outcome — a crash
    /// between claim and dispatch. Safe to fire again: no completion was ever
    /// written, so nothing downstream saw a result.
    Retry,
    /// Already dispatched (or already failed). Do not fire; just advance the
    /// clock. This is the branch that makes a scheduled mission cost one
    /// container instead of one per restart.
    Settled,
}

/// Take ownership of one occurrence before dispatching it.
///
/// `scheduled_at` is the occurrence's own timestamp — the `next_run_at` that
/// came due — not the wall clock at claim time. That is what makes the claim
/// idempotent across restarts: the same occurrence always maps to the same
/// row.
pub async fn claim_fire(
    pool: &PgPool,
    routine_id: Uuid,
    scheduled_at: OffsetDateTime,
) -> Result<FireClaim, DbError> {
    use sqlx::Row;
    // Insert-or-look-at-what's-there in one statement, so two schedulers
    // racing the same occurrence cannot both see "fresh".
    let row = sqlx::query(
        "INSERT INTO routine_fires (routine_id, scheduled_at)
         VALUES ($1, $2)
         ON CONFLICT (routine_id, scheduled_at) DO UPDATE
             SET routine_id = routine_fires.routine_id
         RETURNING status, (xmax = 0) AS inserted",
    )
    .bind(routine_id)
    .bind(scheduled_at)
    .fetch_one(pool)
    .await?;

    // `xmax = 0` distinguishes a genuine insert from a no-op update — the
    // usual Postgres trick, and the reason for the otherwise pointless
    // self-assignment in DO UPDATE (a bare DO NOTHING returns no row at all).
    let inserted: bool = row.try_get("inserted").unwrap_or(false);
    if inserted {
        return Ok(FireClaim::Fresh);
    }
    let status: String = row.try_get("status").unwrap_or_default();
    Ok(match status.as_str() {
        "claimed" => FireClaim::Retry,
        _ => FireClaim::Settled,
    })
}

/// Record how a dispatched occurrence ended. Called after the work is handed
/// off, so a crash before this leaves the row `claimed` and retryable.
pub async fn complete_fire(
    pool: &PgPool,
    routine_id: Uuid,
    scheduled_at: OffsetDateTime,
    error: Option<&str>,
) -> Result<(), DbError> {
    sqlx::query(
        "UPDATE routine_fires
            SET status = $3, completed_at = now(), error = $4
          WHERE routine_id = $1 AND scheduled_at = $2",
    )
    .bind(routine_id)
    .bind(scheduled_at)
    .bind(if error.is_some() { "failed" } else { "fired" })
    .bind(error)
    .execute(pool)
    .await?;
    Ok(())
}