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, #[serde(with = "time::serde::rfc3339::option")] pub last_run_at: Option, } pub async fn create( pool: &PgPool, agent_id: AgentId, name: &str, schedule_cron: &str, action: serde_json::Value, next_run_at: OffsetDateTime, ) -> Result { 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, 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, 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, ) -> 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 { 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(()) }