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

//! Persistence for topology runs — both saved "compare" results and durable
//! single-topology "run" jobs (queued → running → completed/failed), the latter
//! checkpointed per step for crash-resumable long-horizon execution.

use cm_domain::WorkspaceId;
use serde_json::Value;
use sqlx::PgPool;
use time::OffsetDateTime;
use uuid::Uuid;

use crate::DbError;

/// A row summary for the recent-runs list. `iteration` and `finished_at`
/// are populated for loop iterations and for terminal runs respectively;
/// `None` for compares or still-in-flight runs.
pub struct TopologyRunSummary {
    pub id: Uuid,
    pub task: String,
    pub status: String,
    pub kind: String,
    pub created_at: OffsetDateTime,
    pub iteration: Option<i32>,
    pub finished_at: Option<OffsetDateTime>,
}

/// A full saved comparison run.
pub struct TopologyRun {
    pub id: Uuid,
    pub task: String,
    pub comparison: Value,
    pub created_at: OffsetDateTime,
}

/// A durable topology job claimed by the background worker (queued or resumed).
pub struct ClaimedTopologyRun {
    pub id: Uuid,
    pub workspace_id: Uuid,
    pub task: String,
    /// Input graph (present for `kind = 'run'` jobs).
    pub graph: Option<Value>,
    /// Completed-step outputs persisted by the last checkpoint, for resume.
    pub checkpoint: Option<Value>,
    /// Event-journal offset reached so far.
    pub last_event_id: i64,
    /// Deploy tier: `team` drives claws directly; `company`/`org` drive the
    /// recursive sub-topology executor.
    pub tier: String,
}

/// Lifecycle status + progress for a durable run (status endpoint).
pub struct TopologyRunStatus {
    pub id: Uuid,
    pub task: String,
    pub kind: String,
    pub status: String,
    pub error: Option<String>,
    pub checkpoint: Option<Value>,
    /// The final result blob (the `RunRecord`/`Comparison`) once completed.
    pub result: Option<Value>,
    pub last_event_id: i64,
    pub created_at: OffsetDateTime,
    pub updated_at: OffsetDateTime,
}

/// Save a finished comparison run for a workspace (synchronous path).
pub async fn insert(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
    task: &str,
    comparison: &Value,
) -> Result<(), DbError> {
    // status/kind default to 'completed'/'compare' for the synchronous path.
    sqlx::query!(
        "INSERT INTO topology_runs (id, workspace_id, task, comparison)
         VALUES ($1, $2, $3, $4)",
        id,
        workspace_id.as_uuid(),
        task,
        comparison,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// Enqueue a durable single-topology run job (`kind = 'run'`, `status = 'queued'`).
/// The background worker claims and executes it; the result lands in `comparison`.
/// Tier defaults to `team` (drives claws directly).
pub async fn enqueue_run(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
    task: &str,
    graph: &Value,
) -> Result<(), DbError> {
    enqueue_run_tier(pool, id, workspace_id, task, graph, "team").await
}

/// Enqueue a durable run for a specific deploy tier (`team` | `company` | `org`).
/// The worker selects the matching executor — a `company`/`org` job runs the
/// recursive sub-topology executor, which drives each child tier in turn.
pub async fn enqueue_run_tier(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
    task: &str,
    graph: &Value,
    tier: &str,
) -> Result<(), DbError> {
    sqlx::query!(
        "INSERT INTO topology_runs (id, workspace_id, task, kind, status, graph, tier)
         VALUES ($1, $2, $3, 'run', 'queued', $4, $5)",
        id,
        workspace_id.as_uuid(),
        task,
        graph,
        tier,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// Enqueue a durable run bound to a specific team. `team_id` is stored so the
/// post-terminal ephemeral-teardown hook can find the team from a completed run.
pub async fn enqueue_run_for_team(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
    task: &str,
    graph: &Value,
    team_id: Uuid,
) -> Result<(), DbError> {
    sqlx::query!(
        "INSERT INTO topology_runs
            (id, workspace_id, task, kind, status, graph, tier, team_id)
         VALUES ($1, $2, $3, 'run', 'queued', $4, 'team', $5)",
        id,
        workspace_id.as_uuid(),
        task,
        graph,
        team_id,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// Result of `check_ephemeral_teardown` when this run's terminal completion
/// should tear down its team.
pub struct EphemeralTeardown {
    pub team_id: Uuid,
    pub workspace_id: Uuid,
    pub claw_ids: Vec<Uuid>,
}

/// If this run belongs to an ephemeral team AND no siblings are still queued /
/// running, return the team's teardown context (team id + workspace + bound
/// claws). Callers then deprovision each claw and delete the team. Returns
/// `None` if there's nothing to do (permanent team, still-running siblings, or
/// no team back-ref at all).
pub async fn check_ephemeral_teardown(
    pool: &PgPool,
    run_id: Uuid,
) -> Result<Option<EphemeralTeardown>, DbError> {
    let row = sqlx::query!(
        "SELECT t.id AS team_id, t.workspace_id
           FROM topology_runs r
           JOIN teams t ON t.id = r.team_id
          WHERE r.id = $1
            AND t.lifecycle = 'ephemeral'
            AND NOT EXISTS (
                    SELECT 1 FROM topology_runs sib
                     WHERE sib.team_id = t.id
                       AND sib.id <> r.id
                       AND sib.status IN ('queued', 'running')
                )",
        run_id,
    )
    .fetch_optional(pool)
    .await?;
    let Some(row) = row else { return Ok(None) };
    let claws = sqlx::query!(
        "SELECT claw_id FROM team_members WHERE team_id = $1",
        row.team_id,
    )
    .fetch_all(pool)
    .await?;
    Ok(Some(EphemeralTeardown {
        team_id: row.team_id,
        workspace_id: row.workspace_id,
        claw_ids: claws.into_iter().map(|r| r.claw_id).collect(),
    }))
}

/// Atomically claim the oldest queued job, flipping it to `running`. Uses
/// `FOR UPDATE SKIP LOCKED` so multiple workers never claim the same job.
/// Returns `None` when the queue is empty.
pub async fn claim_next_queued(pool: &PgPool) -> Result<Option<ClaimedTopologyRun>, DbError> {
    let row = sqlx::query!(
        "UPDATE topology_runs
         SET status = 'running', started_at = COALESCE(started_at, now()), updated_at = now()
         WHERE id = (
             SELECT id FROM topology_runs
             WHERE status = 'queued'
             ORDER BY created_at
             FOR UPDATE SKIP LOCKED
             LIMIT 1
         )
         RETURNING id, workspace_id, task, graph, checkpoint, last_event_id, tier",
    )
    .fetch_optional(pool)
    .await?;
    Ok(row.map(|r| ClaimedTopologyRun {
        id: r.id,
        workspace_id: r.workspace_id,
        task: r.task,
        graph: r.graph,
        checkpoint: r.checkpoint,
        last_event_id: r.last_event_id,
        tier: r.tier,
    }))
}

/// Persist mid-run progress: the completed-step checkpoint + journal offset.
/// Touches `updated_at` so the stale-run sweeper treats the job as alive.
pub async fn checkpoint(
    pool: &PgPool,
    id: Uuid,
    checkpoint: &Value,
    last_event_id: i64,
) -> Result<(), DbError> {
    sqlx::query!(
        "UPDATE topology_runs
         SET checkpoint = $2, last_event_id = $3, updated_at = now()
         WHERE id = $1",
        id,
        checkpoint,
        last_event_id,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// Touch `updated_at` without changing the checkpoint — keeps a long-running
/// job visibly alive to the stale-run sweeper. Used by the recursive executor:
/// a parent (company/org) run can spend minutes inside one node executing a
/// child sub-topology, so every leaf turn touches the parent here to prevent
/// the 180s sweep from requeuing the parent mid-subtree.
pub async fn touch(pool: &PgPool, id: Uuid) -> Result<(), DbError> {
    sqlx::query!(
        "UPDATE topology_runs SET updated_at = now() WHERE id = $1",
        id,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// If this run belongs to a research topic AND no siblings of that topic
/// are still queued or running, transition the topic `processing → reviewing`.
/// Guarded by `status = 'processing'` so a repeat call (e.g. a retry) is a
/// no-op; a topic already reviewing/publishing/published stays put.
/// Returns `true` when the topic was transitioned.
pub async fn notify_run_completed(pool: &PgPool, id: Uuid) -> Result<bool, DbError> {
    // One statement: subquery locates the topic id, subquery counts siblings
    // still in flight (excluding *this* run — it's about to be flipped to
    // completed/failed by the caller, but ordering isn't guaranteed here).
    let row = sqlx::query!(
        "UPDATE research_topics t
            SET status = 'reviewing', updated_at = now()
          WHERE t.id = (
                    SELECT research_topic_id FROM topology_runs
                     WHERE id = $1 AND research_topic_id IS NOT NULL
                )
            AND t.status = 'processing'
            AND NOT EXISTS (
                    SELECT 1 FROM topology_runs
                     WHERE research_topic_id = t.id
                       AND id <> $1
                       AND status IN ('queued', 'running')
                )
          RETURNING t.id",
        id,
    )
    .fetch_optional(pool)
    .await?;
    Ok(row.is_some())
}

/// Mark a job completed and store its final result blob.
pub async fn complete(pool: &PgPool, id: Uuid, result: &Value) -> Result<(), DbError> {
    sqlx::query!(
        "UPDATE topology_runs
         SET status = 'completed', comparison = $2, finished_at = now(), updated_at = now()
         WHERE id = $1",
        id,
        result,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// Mark a job failed with an error message.
pub async fn fail(pool: &PgPool, id: Uuid, error: &str) -> Result<(), DbError> {
    sqlx::query!(
        "UPDATE topology_runs
         SET status = 'failed', error = $2, finished_at = now(), updated_at = now()
         WHERE id = $1",
        id,
        error,
    )
    .execute(pool)
    .await?;
    Ok(())
}

/// Cancel a run (workspace-scoped). Only `queued`/`running` jobs can be
/// cancelled; returns whether a row transitioned. The worker observes the new
/// status at its next step boundary and stops.
pub async fn cancel(pool: &PgPool, id: Uuid, workspace_id: WorkspaceId) -> Result<bool, DbError> {
    let result = sqlx::query!(
        "UPDATE topology_runs
         SET status = 'cancelled', finished_at = now(), updated_at = now()
         WHERE id = $1 AND workspace_id = $2 AND status IN ('queued', 'running')",
        id,
        workspace_id.as_uuid(),
    )
    .execute(pool)
    .await?;
    Ok(result.rows_affected() == 1)
}

/// The current status of a run (no workspace scope) — used by the worker to
/// detect cancellation mid-run without re-reading the whole row.
pub async fn current_status(pool: &PgPool, id: Uuid) -> Result<Option<String>, DbError> {
    let row = sqlx::query!("SELECT status FROM topology_runs WHERE id = $1", id)
        .fetch_optional(pool)
        .await?;
    Ok(row.map(|r| r.status))
}

/// Resume sweep: requeue `running` jobs whose worker went silent (no checkpoint
/// touch within `older_than_secs`). The next claim resumes them from checkpoint.
/// Returns how many were requeued.
pub async fn requeue_stale(pool: &PgPool, older_than_secs: f64) -> Result<u64, DbError> {
    let result = sqlx::query!(
        "UPDATE topology_runs
         SET status = 'queued', updated_at = now()
         WHERE status = 'running' AND updated_at < now() - make_interval(secs => $1)",
        older_than_secs,
    )
    .execute(pool)
    .await?;
    Ok(result.rows_affected())
}

/// How many durable runs are currently `queued` or `running` for a workspace
/// (the dashboard "running now" / status-bar counter).
pub async fn count_active(pool: &PgPool, workspace_id: WorkspaceId) -> Result<i64, DbError> {
    let row = sqlx::query!(
        "SELECT count(*) AS n FROM topology_runs
         WHERE workspace_id = $1 AND status IN ('queued', 'running')",
        workspace_id.as_uuid(),
    )
    .fetch_one(pool)
    .await?;
    Ok(row.n.unwrap_or(0))
}

/// The most recent runs for a workspace, newest first.
pub async fn list_recent(
    pool: &PgPool,
    workspace_id: WorkspaceId,
    limit: i64,
) -> Result<Vec<TopologyRunSummary>, DbError> {
    let rows = sqlx::query!(
        "SELECT id, task, status, kind, created_at, iteration, finished_at
         FROM topology_runs
         WHERE workspace_id = $1 ORDER BY created_at DESC LIMIT $2",
        workspace_id.as_uuid(),
        limit,
    )
    .fetch_all(pool)
    .await?;
    Ok(rows
        .into_iter()
        .map(|r| TopologyRunSummary {
            id: r.id,
            task: r.task,
            status: r.status,
            kind: r.kind,
            created_at: r.created_at,
            iteration: r.iteration,
            finished_at: r.finished_at,
        })
        .collect())
}

/// Iterations of a loop, newest first. Uses the partial index
/// `topology_runs_loop_idx` on `(loop_id, iteration DESC)`.
pub async fn list_by_loop(
    pool: &PgPool,
    workspace_id: WorkspaceId,
    loop_id: Uuid,
    limit: i64,
) -> Result<Vec<TopologyRunSummary>, DbError> {
    let rows = sqlx::query!(
        "SELECT id, task, status, kind, created_at, iteration, finished_at
         FROM topology_runs
         WHERE workspace_id = $1 AND loop_id = $2
         ORDER BY iteration DESC NULLS LAST, created_at DESC
         LIMIT $3",
        workspace_id.as_uuid(),
        loop_id,
        limit,
    )
    .fetch_all(pool)
    .await?;
    Ok(rows
        .into_iter()
        .map(|r| TopologyRunSummary {
            id: r.id,
            task: r.task,
            status: r.status,
            kind: r.kind,
            created_at: r.created_at,
            iteration: r.iteration,
            finished_at: r.finished_at,
        })
        .collect())
}

/// A single saved comparison/run result, scoped to its workspace.
pub async fn get(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
) -> Result<TopologyRun, DbError> {
    let row = sqlx::query!(
        "SELECT id, task, comparison, created_at FROM topology_runs
         WHERE id = $1 AND workspace_id = $2",
        id,
        workspace_id.as_uuid(),
    )
    .fetch_optional(pool)
    .await?
    .ok_or(DbError::NotFound)?;
    Ok(TopologyRun {
        id: row.id,
        task: row.task,
        // Nullable since the durable path: a queued/running run has no result yet.
        comparison: row.comparison.unwrap_or(Value::Null),
        created_at: row.created_at,
    })
}

/// Lifecycle status + progress for a durable run (status endpoint), workspace-scoped.
pub async fn status(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
) -> Result<TopologyRunStatus, DbError> {
    let row = sqlx::query!(
        "SELECT id, task, kind, status, error, checkpoint, comparison, last_event_id,
                created_at, updated_at
         FROM topology_runs WHERE id = $1 AND workspace_id = $2",
        id,
        workspace_id.as_uuid(),
    )
    .fetch_optional(pool)
    .await?
    .ok_or(DbError::NotFound)?;
    Ok(TopologyRunStatus {
        id: row.id,
        task: row.task,
        kind: row.kind,
        status: row.status,
        error: row.error,
        checkpoint: row.checkpoint,
        result: row.comparison,
        last_event_id: row.last_event_id,
        created_at: row.created_at,
        updated_at: row.updated_at,
    })
}