tireless/crates/tireless-data/src/job_store.rs

//! `JobStore` over Postgres.
//!
//! Claiming is a row transition under `SELECT … FOR UPDATE SKIP LOCKED`
//! (`architecture/generic.md` §3), which makes it atomic across however many
//! runners exist. Nothing here reads a forge label to decide anything —
//! `CLAUDE.md` invariant 6.

use async_trait::async_trait;
use chrono::{DateTime, Duration, Utc};
use sqlx::{Row, postgres::PgRow};
use tireless_core::port::{DiscoveredIssue, JobStore};
use tireless_core::routing::{job_kind_for, lane_for};
use tireless_entities::{
    AgentKind, Error, Forge, IssueRef, Job, JobKind, JobState, LabelProtocol, PollSchedule,
    TrackedRepo,
};
use uuid::Uuid;

use crate::store::PgStore;

fn db(e: sqlx::Error) -> Error {
    Error::Database(e.to_string())
}

/// Parse a value written by serde back into its enum.
///
/// Round-tripping through serde rather than hand-writing a `match` is what makes
/// the schema's check constraints and the Rust types provably the same
/// vocabulary — the store tests assert the constraint lists, and this asserts
/// the parse. A hand-written match would let the two drift apart silently.
fn parse<T: serde::de::DeserializeOwned>(s: &str, what: &str) -> Result<T, Error> {
    serde_json::from_value(serde_json::Value::String(s.to_string())).map_err(|_| {
        Error::Database(format!(
            "column {what} holds {s:?}, which is not a value this build understands; \
             a migration added it without the matching enum variant"
        ))
    })
}

fn wire<T: serde::Serialize>(v: &T) -> String {
    serde_json::to_value(v)
        .ok()
        .and_then(|x| x.as_str().map(str::to_owned))
        .expect("domain enums serialise to strings")
}

fn job_from_row(row: &PgRow) -> Result<Job, Error> {
    Ok(Job {
        id: row.try_get("id").map_err(db)?,
        issue: IssueRef {
            forge: parse::<Forge>(row.try_get("forge").map_err(db)?, "job.forge")?,
            owner: row.try_get("owner").map_err(db)?,
            repo: row.try_get("repo").map_err(db)?,
            number: row.try_get("number").map_err(db)?,
        },
        kind: parse::<JobKind>(row.try_get("kind").map_err(db)?, "job.kind")?,
        state: parse::<JobState>(row.try_get("state").map_err(db)?, "job.state")?,
        claimed_by: row.try_get("claimed_by").map_err(db)?,
        claim_expires_at: row.try_get("claim_expires_at").map_err(db)?,
        parent_job_id: row.try_get("parent_job_id").map_err(db)?,
        attempts: row.try_get("attempts").map_err(db)?,
        last_error: row.try_get("last_error").map_err(db)?,
        created_at: row.try_get("created_at").map_err(db)?,
        updated_at: row.try_get("updated_at").map_err(db)?,
    })
}

fn repo_from_row(row: &PgRow) -> Result<TrackedRepo, Error> {
    Ok(TrackedRepo {
        id: row.try_get("id").map_err(db)?,
        forge: parse::<Forge>(row.try_get("forge").map_err(db)?, "tracked_repo.forge")?,
        owner: row.try_get("owner").map_err(db)?,
        repo: row.try_get("repo").map_err(db)?,
        clone_url: row.try_get("clone_url").map_err(db)?,
        default_branch: row.try_get("default_branch").map_err(db)?,
        schedule: PollSchedule {
            interval_seconds: row.try_get::<i32, _>("interval_seconds").map_err(db)? as u32,
            quiet_from: row.try_get("quiet_from").map_err(db)?,
            quiet_until: row.try_get("quiet_until").map_err(db)?,
            enabled: row.try_get("enabled").map_err(db)?,
        },
        last_etag: row.try_get("last_etag").map_err(db)?,
        last_polled_at: row.try_get("last_polled_at").map_err(db)?,
        created_at: row.try_get("created_at").map_err(db)?,
    })
}

/// How long a claim is held before the lease expires.
///
/// Longer than any single agent run's ceiling would strand a job for hours after
/// a crash; shorter than a run would let a second worker claim a job that is
/// still being worked. The runner renews while it works, so this is the *gap
/// after a worker stops renewing*, not a run budget.
const CLAIM_LEASE: Duration = Duration::minutes(10);

#[async_trait]
impl JobStore for PgStore {
    async fn tracked_repos(&self) -> Result<Vec<TrackedRepo>, Error> {
        let rows = sqlx::query(
            "select id, forge, owner, repo, clone_url, default_branch, interval_seconds, \
                    quiet_from, quiet_until, enabled, last_etag, last_polled_at, created_at \
               from tracked_repo \
              where enabled \
              order by created_at",
        )
        .fetch_all(self.pool())
        .await
        .map_err(db)?;

        rows.iter().map(repo_from_row).collect()
    }

    async fn enqueue(&self, issues: &[DiscoveredIssue]) -> Result<usize, Error> {
        let mut enqueued = 0usize;

        for issue in issues {
            // A closed issue is withdrawn work, whatever it is labelled.
            if !issue.is_open {
                continue;
            }
            // Labels decide what to do, and whether to do anything at all. An
            // issue that is not opted in never reaches here from a well-behaved
            // client, but enqueue is the last gate before work exists, so it
            // checks rather than assumes.
            let Some(kind) = job_kind_for(&issue.labels, &self.labels) else {
                continue;
            };

            // Recorded now because the labels that decide it are in hand now;
            // deriving it at claim time would cost a forge request per claim.
            let lane = lane_for(kind, false, &issue.labels, &self.labels).agent;

            // `on conflict do nothing` against the partial unique index: an
            // issue that already has a live job is skipped, and one whose jobs
            // are all terminal gets a new one. That is the same rule the index
            // encodes, so the two cannot disagree.
            let result = sqlx::query(
                "insert into job (forge, owner, repo, number, kind, state, lane) \
                 values ($1, $2, $3, $4, $5, 'pending', $6) \
                 on conflict do nothing",
            )
            .bind(wire(&issue.issue.forge))
            .bind(&issue.issue.owner)
            .bind(&issue.issue.repo)
            .bind(issue.issue.number)
            .bind(wire(&kind))
            .bind(wire(&lane))
            .execute(self.pool())
            .await
            .map_err(db)?;

            enqueued += result.rows_affected() as usize;
        }

        Ok(enqueued)
    }

    async fn claim_next(
        &self,
        worker: &str,
        allowed_lanes: &[AgentKind],
    ) -> Result<Option<Job>, Error> {
        if allowed_lanes.is_empty() {
            // Every lane is held. Asking the database would be a round trip to
            // be told nothing, and an empty `any($1)` matches nothing anyway.
            return Ok(None);
        }
        let lanes: Vec<String> = allowed_lanes.iter().map(wire).collect();

        // One statement, so the select and the update are the same transaction
        // without managing one by hand. The CTE takes the row lock with SKIP
        // LOCKED; the update writes the claim. A second claimer running
        // concurrently skips the locked row rather than blocking on it.
        let row = sqlx::query(
            "with candidate as ( \
                 select id from job \
                  where state = 'pending' and lane = any($1) \
                  order by created_at \
                  limit 1 \
                  for update skip locked \
             ) \
             update job set state = 'claimed', \
                            claimed_by = $2, \
                            claim_expires_at = now() + $3::interval, \
                            attempts = attempts + 1, \
                            updated_at = now() \
               from candidate \
              where job.id = candidate.id \
             returning job.id, job.forge, job.owner, job.repo, job.number, job.kind, \
                       job.state, job.claimed_by, job.claim_expires_at, job.parent_job_id, \
                       job.attempts, job.last_error, job.created_at, job.updated_at",
        )
        .bind(&lanes)
        .bind(worker)
        .bind(format!("{} seconds", CLAIM_LEASE.num_seconds()))
        .fetch_optional(self.pool())
        .await
        .map_err(db)?;

        row.as_ref().map(job_from_row).transpose()
    }

    async fn transition(&self, job: Uuid, to: JobState) -> Result<(), Error> {
        // A terminal job is never re-opened: re-running one is an operator
        // action that creates a *new* job, so that history is not rewritten
        // (design.md §4.3).
        let terminal: Vec<String> = [
            JobState::Delivered,
            JobState::Blocked,
            JobState::Failed,
            JobState::Abandoned,
        ]
        .iter()
        .map(wire)
        .collect();

        // Returning to pending must release the claim, or the row violates
        // `pending_holds_no_claim` — the database refuses to let this be
        // half-done, which is what makes lease expiry safe.
        let releasing = to == JobState::Pending;

        let result = sqlx::query(
            "update job \
                set state = $2, \
                    claimed_by = case when $3 then null else claimed_by end, \
                    claim_expires_at = case when $3 then null else claim_expires_at end, \
                    updated_at = now() \
              where id = $1 and state <> all($4)",
        )
        .bind(job)
        .bind(wire(&to))
        .bind(releasing)
        .bind(&terminal)
        .execute(self.pool())
        .await
        .map_err(db)?;

        if result.rows_affected() == 0 {
            // Either the job is gone or it is already terminal. Both are the
            // caller asking for something that cannot happen, and both are worth
            // surfacing rather than silently succeeding.
            return Err(Error::IllegalTransition(job));
        }
        Ok(())
    }

    async fn expire_stale_claims(&self) -> Result<usize, Error> {
        let result = sqlx::query(
            "update job \
                set state = 'pending', \
                    claimed_by = null, \
                    claim_expires_at = null, \
                    last_error = coalesce(last_error, 'claim lease expired'), \
                    updated_at = now() \
              where state in ('claimed', 'running') \
                and claim_expires_at < now()",
        )
        .execute(self.pool())
        .await
        .map_err(db)?;

        Ok(result.rows_affected() as usize)
    }
}

impl PgStore {
    /// Renew a held claim, pushing its lease out.
    ///
    /// A run can outlast `CLAIM_LEASE`, and without renewal the sweeper would
    /// hand a still-running job to a second worker — two agents on one issue,
    /// which the claim protocol exists to prevent. Guarded on `claimed_by` so a
    /// worker cannot renew a claim it lost while it was not looking.
    pub async fn renew_claim(&self, job: Uuid, worker: &str) -> Result<bool, Error> {
        let result = sqlx::query(
            "update job set claim_expires_at = now() + $3::interval, updated_at = now() \
              where id = $1 and claimed_by = $2 and state in ('claimed', 'running')",
        )
        .bind(job)
        .bind(worker)
        .bind(format!("{} seconds", CLAIM_LEASE.num_seconds()))
        .execute(self.pool())
        .await
        .map_err(db)?;

        Ok(result.rows_affected() == 1)
    }

    /// Record the labels a poll observed, refreshing the cached lane.
    ///
    /// The lane is a cache of operator intent (see `0002_job_lane.sql`), so an
    /// operator adding `tireless/agent:oc` to a queued issue has to reach the
    /// job somehow. Only live jobs are touched: re-routing history would be
    /// rewriting what happened.
    pub async fn refresh_lane(
        &self,
        issue: &IssueRef,
        labels: &[String],
        protocol: &LabelProtocol,
    ) -> Result<usize, Error> {
        let row = sqlx::query(
            "select id, kind, parent_job_id from job \
              where forge = $1 and owner = $2 and repo = $3 and number = $4 \
                and state in ('pending', 'claimed', 'running')",
        )
        .bind(wire(&issue.forge))
        .bind(&issue.owner)
        .bind(&issue.repo)
        .bind(issue.number)
        .fetch_optional(self.pool())
        .await
        .map_err(db)?;

        let Some(row) = row else { return Ok(0) };
        let kind: JobKind = parse(row.try_get("kind").map_err(db)?, "job.kind")?;
        let has_parent = row
            .try_get::<Option<Uuid>, _>("parent_job_id")
            .map_err(db)?
            .is_some();
        let lane = lane_for(kind, has_parent, labels, protocol).agent;

        let result = sqlx::query("update job set lane = $2, updated_at = now() where id = $1")
            .bind(row.try_get::<Uuid, _>("id").map_err(db)?)
            .bind(wire(&lane))
            .execute(self.pool())
            .await
            .map_err(db)?;

        Ok(result.rows_affected() as usize)
    }

    /// When the most recent finished job of a kind ran against an issue.
    ///
    /// The discovery cooldown asks this: a survey that ran this morning has
    /// nothing new to say this afternoon (design.md §2.6).
    pub async fn last_finished_at(
        &self,
        issue: &IssueRef,
        kind: JobKind,
    ) -> Result<Option<DateTime<Utc>>, Error> {
        let row = sqlx::query(
            "select max(updated_at) as last from job \
              where forge = $1 and owner = $2 and repo = $3 and number = $4 \
                and kind = $5 and state = 'delivered'",
        )
        .bind(wire(&issue.forge))
        .bind(&issue.owner)
        .bind(&issue.repo)
        .bind(issue.number)
        .bind(wire(&kind))
        .fetch_one(self.pool())
        .await
        .map_err(db)?;

        row.try_get("last").map_err(db)
    }
}