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(())
}

/// Count `queued` + `running` runs whose `research_topic_id` matches. The
/// research canvas polls this so it can show a spinner "the pipeline is
/// running" and suppress the manual "Submit for review" escape hatch
/// while any run is still in flight.
pub async fn active_runs_for_research_topic(
    pool: &PgPool,
    research_topic_id: Uuid,
) -> Result<i64, DbError> {
    let row = sqlx::query!(
        "SELECT count(*) AS n
           FROM topology_runs
          WHERE research_topic_id = $1
            AND status IN ('queued', 'running')",
        research_topic_id,
    )
    .fetch_one(pool)
    .await?;
    Ok(row.n.unwrap_or(0))
}

/// Live-run panel companion to `active_runs_for_research_topic`: return
/// the actual run ids (queued + running) so the UI can subscribe to
/// their SSE event streams. Ordered newest first — the freshest run is
/// the one the user just kicked off.
/// Batch run-count feeder for the research topic list. Returns a
/// (topic_id, in_flight, failed) tuple per topic in `topic_ids`,
/// omitting topics with zero runs. Used to render the errored-state
/// icon + "rerun" affordance on cards in the left sidebar.
///
/// `failed` counts runs that terminated in `failed` since the topic's
/// most recent successful run (or all-time if none have succeeded).
/// That way an old failure on a topic that later succeeded doesn't
/// keep the card flagged as broken.
pub async fn run_counts_by_research_topic(
    pool: &PgPool,
    topic_ids: &[Uuid],
) -> Result<Vec<(Uuid, i64, i64)>, DbError> {
    use sqlx::Row;
    if topic_ids.is_empty() {
        return Ok(Vec::new());
    }
    let rows: Vec<sqlx::postgres::PgRow> = sqlx::query(
        "WITH last_success AS (
           SELECT research_topic_id, max(created_at) AS ts
             FROM topology_runs
            WHERE research_topic_id = ANY($1)
              AND status = 'completed'
            GROUP BY research_topic_id
         )
         SELECT r.research_topic_id AS topic_id,
                count(*) FILTER (WHERE r.status IN ('queued','running')) AS in_flight,
                count(*) FILTER (
                  WHERE r.status = 'failed'
                    AND r.created_at > coalesce(ls.ts, 'epoch'::timestamptz)
                ) AS failed
           FROM topology_runs r
           LEFT JOIN last_success ls
             ON ls.research_topic_id = r.research_topic_id
          WHERE r.research_topic_id = ANY($1)
          GROUP BY r.research_topic_id",
    )
    .bind(topic_ids)
    .fetch_all(pool)
    .await?;
    Ok(rows
        .into_iter()
        .map(|r| {
            (
                r.get::<Uuid, _>("topic_id"),
                r.get::<i64, _>("in_flight"),
                r.get::<i64, _>("failed"),
            )
        })
        .collect())
}

pub async fn active_run_ids_for_research_topic(
    pool: &PgPool,
    research_topic_id: Uuid,
) -> Result<Vec<Uuid>, DbError> {
    use sqlx::Row;
    // Dynamic query (not `sqlx::query!`) so cm-db builds air-gapped
    // without a fresh `cargo sqlx prepare` round-trip. Schema shape is
    // identical to `active_runs_for_research_topic` above.
    let rows: Vec<sqlx::postgres::PgRow> = sqlx::query(
        "SELECT id
           FROM topology_runs
          WHERE research_topic_id = $1
            AND status IN ('queued', 'running')
          ORDER BY created_at DESC",
    )
    .bind(research_topic_id)
    .fetch_all(pool)
    .await?;
    Ok(rows.into_iter().map(|r| r.get::<Uuid, _>("id")).collect())
}

/// The research topic this run belongs to, if any. Used by the topology
/// worker's `freeze_research_outcome` post-hook to snapshot the run's
/// final synthesis into `research_outcomes`.
pub async fn research_topic_id(pool: &PgPool, id: Uuid) -> Result<Option<Uuid>, DbError> {
    let row = sqlx::query!(
        "SELECT research_topic_id FROM topology_runs WHERE id = $1",
        id,
    )
    .fetch_optional(pool)
    .await?;
    Ok(row.and_then(|r| r.research_topic_id))
}

/// The loop this run belongs to, if any. Mirror of `research_topic_id`.
/// Used by `topology_worker` to look up the per-loop gateway URL so a
/// loop's runs land on its isolated daemon (P2). Non-loop runs return
/// None.
pub async fn loop_id_for_run(pool: &PgPool, id: Uuid) -> Result<Option<Uuid>, DbError> {
    use sqlx::Row;
    let row: Option<sqlx::postgres::PgRow> =
        sqlx::query("SELECT loop_id FROM topology_runs WHERE id = $1")
            .bind(id)
            .fetch_optional(pool)
            .await?;
    Ok(row.and_then(|r| r.try_get::<Option<Uuid>, _>("loop_id").ok().flatten()))
}

/// The iteration counter for a loop-bound run. Returns None for chat /
/// research runs (iteration column is nullable). Used by the reorder
/// rationale hook so the mini-timeline can order events by iteration.
pub async fn iteration_for_run(pool: &PgPool, id: Uuid) -> Result<Option<i32>, DbError> {
    use sqlx::Row;
    let row: Option<sqlx::postgres::PgRow> =
        sqlx::query("SELECT iteration FROM topology_runs WHERE id = $1")
            .bind(id)
            .fetch_optional(pool)
            .await?;
    Ok(row.and_then(|r| r.try_get::<Option<i32>, _>("iteration").ok().flatten()))
}

/// Enqueue a durable run bound to a research topic. `research_topic_id` is
/// stored so `notify_run_completed` can flip the owning topic
/// `processing → reviewing` when its last run terminates (see
/// `topology_worker::maybe_transition_research_topic`).
pub async fn enqueue_run_for_research_topic(
    pool: &PgPool,
    id: Uuid,
    workspace_id: WorkspaceId,
    task: &str,
    graph: &Value,
    research_topic_id: Uuid,
) -> Result<(), DbError> {
    sqlx::query!(
        "INSERT INTO topology_runs
            (id, workspace_id, task, kind, status, graph, tier, research_topic_id)
         VALUES ($1, $2, $3, 'run', 'queued', $4, 'team', $5)",
        id,
        workspace_id.as_uuid(),
        task,
        graph,
        research_topic_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).
    //
    // Only advance the topic when it has AT LEAST ONE outcome — otherwise a
    // failed run with no synthesis would push the topic into `reviewing`,
    // the UI would offer "Request publish", the user would click Approve, and
    // decide_publish would 409 on the "no outcome" guard. Stays in
    // `processing` when zero outcomes exist so the loop's next iteration
    // still has a chance to produce one.
    // Dynamic query — the added EXISTS clause on research_outcomes
    // doesn't have an entry in the offline sqlx cache, so we bind
    // values by hand instead of using the `query!` macro.
    use sqlx::Row;
    let row: Option<sqlx::postgres::PgRow> = 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 EXISTS (
                    SELECT 1 FROM research_outcomes
                     WHERE topic_id = t.id
                )
            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",
    )
    .bind(id)
    .fetch_optional(pool)
    .await?;
    Ok(row
        .map(|r| r.try_get::<Uuid, _>("id").is_ok())
        .unwrap_or(false))
}

/// 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,
    })
}