Skip to main content

awa_model/
admin.rs

1use crate::dlq::DlqMetadata;
2use crate::error::AwaError;
3use crate::job::{JobRow, JobState};
4use crate::queue_storage::QueueStorage;
5use chrono::{DateTime, Duration, Utc};
6use serde::{Deserialize, Serialize};
7use sqlx::types::Json;
8use sqlx::PgExecutor;
9use sqlx::PgPool;
10use std::cmp::max;
11use std::collections::HashMap;
12use std::fmt;
13use std::str::FromStr;
14use uuid::Uuid;
15
16#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
17pub struct JobTimelineEvent {
18    pub timestamp: DateTime<Utc>,
19    pub label: String,
20    pub state: Option<JobState>,
21    pub detail: Option<String>,
22    pub is_error: bool,
23}
24
25#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
26pub struct JobDumpSummary {
27    pub original_priority: i16,
28    pub can_retry: bool,
29    pub can_cancel: bool,
30    pub error_count: usize,
31    pub latest_error: Option<serde_json::Value>,
32}
33
34#[derive(Debug, Clone, Serialize, Deserialize)]
35pub struct JobDump {
36    pub job: JobRow,
37    pub summary: JobDumpSummary,
38    pub timeline: Vec<JobTimelineEvent>,
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub dlq: Option<DlqMetadata>,
41}
42
43#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
44#[serde(rename_all = "snake_case")]
45pub enum RunDumpSource {
46    CurrentJobRow,
47    ErrorHistory,
48}
49
50#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
51pub struct CallbackDump {
52    pub callback_id: Option<Uuid>,
53    pub callback_timeout_at: Option<DateTime<Utc>>,
54    pub callback_filter: Option<String>,
55    pub callback_on_complete: Option<String>,
56    pub callback_on_fail: Option<String>,
57    pub callback_transform: Option<String>,
58}
59
60#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
61pub struct RunDump {
62    pub job_id: i64,
63    pub kind: String,
64    pub queue: String,
65    pub selected_attempt: i16,
66    pub current_attempt: i16,
67    pub current_run_lease: i64,
68    pub selected_run_lease: Option<i64>,
69    pub source: RunDumpSource,
70    pub state: JobState,
71    pub started_at: Option<DateTime<Utc>>,
72    pub finished_at: Option<DateTime<Utc>>,
73    pub error: Option<String>,
74    pub terminal: Option<bool>,
75    pub progress: Option<serde_json::Value>,
76    pub metadata: Option<serde_json::Value>,
77    pub callback: Option<CallbackDump>,
78    pub raw_error_entry: Option<serde_json::Value>,
79    pub notes: Vec<String>,
80}
81
82#[derive(Debug, Clone)]
83struct ErrorEntry {
84    attempt: Option<i16>,
85    error: Option<String>,
86    at: Option<DateTime<Utc>>,
87    terminal: bool,
88    raw: serde_json::Value,
89}
90
91fn parse_error_entry(value: &serde_json::Value) -> Option<ErrorEntry> {
92    let obj = value.as_object()?;
93    let attempt = obj
94        .get("attempt")
95        .and_then(|v| v.as_i64())
96        .and_then(|v| i16::try_from(v).ok());
97    let error = obj
98        .get("error")
99        .and_then(|v| v.as_str())
100        .map(ToOwned::to_owned);
101    let at = obj
102        .get("at")
103        .and_then(|v| v.as_str())
104        .and_then(|v| chrono::DateTime::parse_from_rfc3339(v).ok())
105        .map(|v| v.with_timezone(&Utc));
106    let terminal = obj
107        .get("terminal")
108        .and_then(|v| v.as_bool())
109        .unwrap_or(false);
110    Some(ErrorEntry {
111        attempt,
112        error,
113        at,
114        terminal,
115        raw: value.clone(),
116    })
117}
118
119fn original_priority(job: &JobRow) -> i16 {
120    job.metadata
121        .get("_awa_original_priority")
122        .and_then(|v| v.as_i64())
123        .and_then(|v| i16::try_from(v).ok())
124        .unwrap_or(job.priority)
125}
126
127fn build_job_timeline(job: &JobRow) -> Vec<JobTimelineEvent> {
128    let mut events = Vec::new();
129    events.push(JobTimelineEvent {
130        timestamp: job.created_at,
131        label: "Created".to_string(),
132        state: None,
133        detail: None,
134        is_error: false,
135    });
136
137    if let Some(errors) = &job.errors {
138        for entry in errors.iter().filter_map(parse_error_entry) {
139            if let Some(timestamp) = entry.at {
140                let label = match entry.attempt {
141                    Some(attempt) => format!("Attempt {attempt} failed"),
142                    None => "Error".to_string(),
143                };
144                events.push(JobTimelineEvent {
145                    timestamp,
146                    label,
147                    state: Some(JobState::Failed),
148                    detail: entry.error,
149                    is_error: true,
150                });
151            }
152        }
153    }
154
155    if let Some(attempted_at) = job.attempted_at {
156        let label = if job.state == JobState::WaitingExternal {
157            format!("Attempt {} — waiting for callback", job.attempt)
158        } else if job.state.is_terminal() {
159            format!("Attempt {} started", job.attempt)
160        } else {
161            format!("Attempt {} running", job.attempt)
162        };
163        let state = if job.state.is_terminal() {
164            Some(JobState::Running)
165        } else {
166            Some(job.state)
167        };
168        events.push(JobTimelineEvent {
169            timestamp: attempted_at,
170            label,
171            state,
172            detail: None,
173            is_error: false,
174        });
175    }
176
177    if let Some(finalized_at) = job.finalized_at {
178        let label = match job.state {
179            JobState::Completed => format!("Completed (attempt {})", job.attempt),
180            JobState::Failed => format!(
181                "Failed after {} attempt{}",
182                job.attempt,
183                if job.attempt == 1 { "" } else { "s" }
184            ),
185            JobState::Cancelled => "Cancelled".to_string(),
186            _ => "Finalized".to_string(),
187        };
188        events.push(JobTimelineEvent {
189            timestamp: finalized_at,
190            label,
191            state: Some(job.state),
192            detail: None,
193            is_error: false,
194        });
195    }
196
197    events.sort_by_key(|event| event.timestamp);
198    events
199}
200
201fn build_job_dump(job: JobRow, dlq: Option<DlqMetadata>) -> JobDump {
202    let latest_error = job
203        .errors
204        .as_ref()
205        .and_then(|errors| errors.last())
206        .cloned();
207    let summary = JobDumpSummary {
208        original_priority: original_priority(&job),
209        can_retry: dlq.is_none()
210            && matches!(
211                job.state,
212                JobState::Failed | JobState::Cancelled | JobState::WaitingExternal
213            ),
214        can_cancel: dlq.is_none() && !job.state.is_terminal(),
215        error_count: job.errors.as_ref().map(|errors| errors.len()).unwrap_or(0),
216        latest_error,
217    };
218    let timeline = build_job_timeline(&job);
219    JobDump {
220        job,
221        summary,
222        timeline,
223        dlq,
224    }
225}
226
227fn callback_dump(job: &JobRow) -> Option<CallbackDump> {
228    let callback = CallbackDump {
229        callback_id: job.callback_id,
230        callback_timeout_at: job.callback_timeout_at,
231        callback_filter: job.callback_filter.clone(),
232        callback_on_complete: job.callback_on_complete.clone(),
233        callback_on_fail: job.callback_on_fail.clone(),
234        callback_transform: job.callback_transform.clone(),
235    };
236    if callback.callback_id.is_none()
237        && callback.callback_timeout_at.is_none()
238        && callback.callback_filter.is_none()
239        && callback.callback_on_complete.is_none()
240        && callback.callback_on_fail.is_none()
241        && callback.callback_transform.is_none()
242    {
243        None
244    } else {
245        Some(callback)
246    }
247}
248
249async fn active_queue_storage(pool: &PgPool) -> Result<Option<QueueStorage>, AwaError> {
250    QueueStorage::active_schema(pool)
251        .await?
252        .map(QueueStorage::from_existing_schema)
253        .transpose()
254}
255
256/// Transaction-aware variant of [`active_queue_storage`] — read the active
257/// queue-storage schema inside the caller's transaction so the engine
258/// selection and the resulting state UPDATE commit (or roll back)
259/// together. Required by the ADR-029 atomic callback resolution path so a
260/// race between this read and a routing flip can never split the chosen
261/// engine and the transition that follows.
262async fn active_queue_storage_in_tx(
263    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
264) -> Result<Option<QueueStorage>, AwaError> {
265    QueueStorage::active_schema_in_tx(tx)
266        .await?
267        .map(QueueStorage::from_existing_schema)
268        .transpose()
269}
270
271fn queue_storage_current_jobs_cte(schema: &str) -> String {
272    format!(
273        r#"
274        WITH current_available AS (
275            SELECT
276                ready.job_id,
277                ready.kind,
278                ready.queue,
279                'available'::awa.job_state AS state,
280                ready.created_at,
281                ready.run_at,
282                NULL::timestamptz AS finalized_at
283            FROM {schema}.ready_entries AS ready
284            JOIN {schema}.queue_claim_heads AS claims
285              ON claims.queue = ready.queue
286             AND claims.priority = ready.priority
287             AND claims.enqueue_shard = ready.enqueue_shard
288            WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
289              AND NOT EXISTS (
290                  SELECT 1 FROM {schema}.ready_tombstones AS tomb
291                  WHERE tomb.ready_slot = ready.ready_slot
292                    AND tomb.ready_generation = ready.ready_generation
293                    AND tomb.queue = ready.queue
294                    AND tomb.priority = ready.priority
295                    AND tomb.enqueue_shard = ready.enqueue_shard
296                    AND tomb.lane_seq = ready.lane_seq
297              )
298        ),
299        current_jobs AS (
300            SELECT job_id, kind, queue, state, created_at, run_at, finalized_at
301            FROM current_available
302            UNION ALL
303            SELECT job_id, kind, queue, state, created_at, run_at, finalized_at
304            FROM {schema}.deferred_jobs
305            UNION ALL
306            SELECT
307                leases.job_id,
308                ready.kind,
309                leases.queue,
310                leases.state,
311                ready.created_at,
312                ready.run_at,
313                NULL::timestamptz AS finalized_at
314            FROM {schema}.leases AS leases
315            JOIN {schema}.ready_entries AS ready
316              ON ready.ready_slot = leases.ready_slot
317             AND ready.ready_generation = leases.ready_generation
318             AND ready.queue = leases.queue
319             AND ready.priority = leases.priority
320             AND ready.enqueue_shard = leases.enqueue_shard
321             AND ready.lane_seq = leases.lane_seq
322            UNION ALL
323            SELECT job_id, kind, queue, state, created_at, run_at, finalized_at
324            FROM {schema}.done_entries
325            UNION ALL
326            SELECT
327                job_id,
328                kind,
329                queue,
330                'failed'::awa.job_state AS state,
331                created_at,
332                run_at,
333                finalized_at
334            FROM {schema}.dlq_entries
335        )
336        "#
337    )
338}
339
340async fn notify_cancellation_tx<'a>(
341    tx: &mut sqlx::Transaction<'a, sqlx::Postgres>,
342    job_id: i64,
343    run_lease: i64,
344) -> Result<(), AwaError> {
345    let payload = serde_json::json!({ "job_id": job_id, "run_lease": run_lease }).to_string();
346    sqlx::query("SELECT pg_notify('awa:cancel', $1)")
347        .bind(payload)
348        .execute(tx.as_mut())
349        .await?;
350    Ok(())
351}
352
353async fn list_queue_storage_jobs(
354    store: &QueueStorage,
355    pool: &PgPool,
356    filter: &ListJobsFilter,
357) -> Result<Vec<JobRow>, AwaError> {
358    let limit = filter.limit.unwrap_or(100).clamp(1, 1000);
359    let candidate_limit = if filter.tag.is_some() {
360        limit.saturating_mul(10).min(5000)
361    } else {
362        limit.saturating_mul(3).min(2000)
363    };
364
365    let sql = format!(
366        "{} \
367         SELECT job_id \
368         FROM current_jobs \
369         WHERE ($1::awa.job_state IS NULL OR state = $1) \
370           AND ($2::text IS NULL OR kind = $2) \
371           AND ($3::text IS NULL OR queue = $3) \
372           AND ($4::bigint IS NULL OR job_id < $4) \
373         ORDER BY job_id DESC \
374         LIMIT $5",
375        queue_storage_current_jobs_cte(store.schema())
376    );
377
378    let mut jobs = Vec::new();
379    let mut cursor = filter.before_id;
380
381    loop {
382        let ids: Vec<i64> = sqlx::query_scalar(&sql)
383            .bind(filter.state)
384            .bind(&filter.kind)
385            .bind(&filter.queue)
386            .bind(cursor)
387            .bind(candidate_limit)
388            .fetch_all(pool)
389            .await?;
390        if ids.is_empty() {
391            break;
392        }
393
394        for job_id in &ids {
395            let Some(job) = store.load_job(pool, *job_id).await? else {
396                continue;
397            };
398
399            if let Some(state) = filter.state {
400                if job.state != state {
401                    continue;
402                }
403            }
404            if let Some(kind) = &filter.kind {
405                if &job.kind != kind {
406                    continue;
407                }
408            }
409            if let Some(queue) = &filter.queue {
410                if &job.queue != queue {
411                    continue;
412                }
413            }
414            if let Some(tag) = &filter.tag {
415                if !job.tags.iter().any(|job_tag| job_tag == tag) {
416                    continue;
417                }
418            }
419
420            jobs.push(job);
421            if jobs.len() as i64 >= limit {
422                break;
423            }
424        }
425
426        if jobs.len() as i64 >= limit || ids.len() < candidate_limit as usize {
427            break;
428        }
429        cursor = ids.last().copied();
430    }
431
432    jobs.sort_by_key(|job| std::cmp::Reverse(job.id));
433    Ok(jobs)
434}
435
436/// Retry a single failed, cancelled, or waiting_external job.
437pub async fn retry(pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
438    if let Some(store) = active_queue_storage(pool).await? {
439        return store
440            .retry_job(pool, job_id)
441            .await?
442            .ok_or(AwaError::JobNotFound { id: job_id })
443            .map(Some);
444    }
445
446    sqlx::query_as::<_, JobRow>(
447        r#"
448        UPDATE awa.jobs
449        SET state = 'available', attempt = 0, run_at = now(),
450            finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL,
451            callback_id = NULL, callback_timeout_at = NULL,
452            callback_filter = NULL, callback_on_complete = NULL,
453            callback_on_fail = NULL, callback_transform = NULL
454        WHERE id = $1 AND state IN ('failed', 'cancelled', 'waiting_external')
455        RETURNING *
456        "#,
457    )
458    .bind(job_id)
459    .fetch_optional(pool)
460    .await?
461    .ok_or(AwaError::JobNotFound { id: job_id })
462    .map(Some)
463}
464
465/// Cancel a single non-terminal job.
466pub async fn cancel(pool: &PgPool, job_id: i64) -> Result<Option<JobRow>, AwaError> {
467    if let Some(store) = active_queue_storage(pool).await? {
468        return store
469            .cancel_job(pool, job_id)
470            .await?
471            .ok_or(AwaError::JobNotFound { id: job_id })
472            .map(Some);
473    }
474
475    // Cancel the row and capture its prior state. If we moved it out
476    // of `running` / `waiting_external`, NOTIFY listening workers so
477    // the handler currently executing it can learn about the
478    // cancellation and stop cleanly.
479    let mut tx = pool.begin().await?;
480
481    let prior_state: Option<JobState> =
482        sqlx::query_scalar::<_, JobState>("SELECT state FROM awa.jobs WHERE id = $1 FOR UPDATE")
483            .bind(job_id)
484            .fetch_optional(tx.as_mut())
485            .await?;
486
487    let Some(prior_state) = prior_state else {
488        tx.rollback().await.ok();
489        return Err(AwaError::JobNotFound { id: job_id });
490    };
491
492    let job: Option<JobRow> = sqlx::query_as::<_, JobRow>(
493        r#"
494        UPDATE awa.jobs
495        SET state = 'cancelled', finalized_at = now(),
496            callback_id = NULL, callback_timeout_at = NULL,
497            callback_filter = NULL, callback_on_complete = NULL,
498            callback_on_fail = NULL, callback_transform = NULL
499        WHERE id = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
500        RETURNING *
501        "#,
502    )
503    .bind(job_id)
504    .fetch_optional(tx.as_mut())
505    .await?;
506
507    let Some(job) = job else {
508        tx.rollback().await.ok();
509        return Err(AwaError::JobNotFound { id: job_id });
510    };
511
512    if matches!(prior_state, JobState::Running | JobState::WaitingExternal) {
513        notify_cancellation_tx(&mut tx, job.id, job.run_lease).await?;
514    }
515
516    tx.commit().await?;
517    Ok(Some(job))
518}
519
520/// Cancel a job by its unique key components.
521///
522/// Reconstructs the BLAKE3 unique key from the same inputs used at insert time
523/// (kind, optional queue, optional args, optional period bucket), then cancels
524/// the single oldest matching non-terminal job. Returns `None` if no matching
525/// job was found (already completed, already cancelled, or never existed).
526///
527/// The parameters must match what was used at insert time: pass `queue` only if
528/// the original `UniqueOpts` had `by_queue: true`, `args` only if `by_args: true`,
529/// and `period_bucket` only if `by_period` was set. Mismatched components produce
530/// a different hash and the job won't be found.
531///
532/// Only one job is cancelled per call (the oldest by `id`). This is intentional:
533/// unique key enforcement uses a state bitmask, so multiple rows with the same
534/// key can legally coexist (e.g., one `waiting_external` + one `available`).
535/// Cancelling all of them in one shot would be surprising.
536///
537/// This is useful when the caller knows the job kind and args but not the job ID —
538/// e.g., cancelling a scheduled reminder when the triggering condition is resolved.
539///
540/// # Implementation notes
541///
542/// Queries `jobs_hot` and `scheduled_jobs` directly rather than the `awa.jobs`
543/// UNION ALL view, because PostgreSQL does not support `FOR UPDATE` on UNION
544/// views. The CTE selects the oldest candidate ID without row locks, then
545/// delegates to `cancel(...)` so running jobs also emit the cooperative
546/// in-flight cancellation notification. If the selected job completes before
547/// `cancel(...)` locks it, the cancel becomes a no-op and returns `None`.
548///
549/// The lookup scans `unique_key` on both physical tables without a dedicated
550/// index. This is acceptable for low-volume use cases. For high-volume tables,
551/// consider adding a partial index on `unique_key WHERE unique_key IS NOT NULL`
552/// or routing through `job_unique_claims` (which is already indexed).
553pub async fn cancel_by_unique_key(
554    pool: &PgPool,
555    kind: &str,
556    queue: Option<&str>,
557    args: Option<&serde_json::Value>,
558    period_bucket: Option<i64>,
559) -> Result<Option<JobRow>, AwaError> {
560    let unique_key = crate::unique::compute_unique_key(kind, queue, args, period_bucket);
561
562    if let Some(store) = active_queue_storage(pool).await? {
563        let sql = format!(
564            r#"
565            WITH current_available AS (
566                SELECT ready.job_id, ready.unique_key
567                FROM {schema}.ready_entries AS ready
568                JOIN {schema}.queue_claim_heads AS claims
569                  ON claims.queue = ready.queue
570                 AND claims.priority = ready.priority
571                 AND claims.enqueue_shard = ready.enqueue_shard
572                WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
573                  AND NOT EXISTS (
574                      SELECT 1 FROM {schema}.ready_tombstones AS tomb
575                      WHERE tomb.ready_slot = ready.ready_slot
576                        AND tomb.ready_generation = ready.ready_generation
577                        AND tomb.queue = ready.queue
578                        AND tomb.priority = ready.priority
579                        AND tomb.enqueue_shard = ready.enqueue_shard
580                        AND tomb.lane_seq = ready.lane_seq
581                  )
582            ),
583            candidates AS (
584                SELECT job_id
585                FROM current_available
586                WHERE unique_key = $1
587                UNION ALL
588                SELECT job_id
589                FROM {schema}.deferred_jobs
590                WHERE unique_key = $1
591                UNION ALL
592                SELECT job_id
593                FROM {schema}.leases
594                WHERE unique_key = $1
595            )
596            SELECT job_id
597            FROM candidates
598            ORDER BY job_id ASC
599            LIMIT 1
600            "#,
601            schema = store.schema()
602        );
603
604        let candidate: Option<i64> = sqlx::query_scalar(&sql)
605            .bind(&unique_key)
606            .fetch_optional(pool)
607            .await?;
608
609        return match candidate {
610            Some(job_id) => cancel(pool, job_id).await,
611            None => Ok(None),
612        };
613    }
614
615    // Find the oldest matching job across both physical tables, then route
616    // through `cancel(...)` so running jobs emit the same cooperative
617    // cancellation notification as ID-based admin cancels.
618    let candidate: Option<i64> = sqlx::query_scalar(
619        r#"
620        WITH candidates AS (
621            SELECT id FROM awa.jobs_hot
622            WHERE unique_key = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
623            UNION ALL
624            SELECT id FROM awa.scheduled_jobs
625            WHERE unique_key = $1 AND state NOT IN ('completed', 'failed', 'cancelled')
626            ORDER BY id ASC
627            LIMIT 1
628        )
629        SELECT id
630        FROM candidates
631        "#,
632    )
633    .bind(&unique_key)
634    .fetch_optional(pool)
635    .await?;
636
637    match candidate {
638        Some(job_id) => match cancel(pool, job_id).await {
639            Ok(row) => Ok(row),
640            Err(AwaError::JobNotFound { .. }) => Ok(None),
641            Err(err) => Err(err),
642        },
643        None => Ok(None),
644    }
645}
646
647/// Retry all failed jobs of a given kind.
648pub async fn retry_failed_by_kind(pool: &PgPool, kind: &str) -> Result<Vec<JobRow>, AwaError> {
649    if let Some(store) = active_queue_storage(pool).await? {
650        let sql = format!(
651            r#"
652            SELECT job_id
653            FROM {schema}.done_entries
654            WHERE kind = $1
655              AND state = 'failed'
656            ORDER BY job_id ASC
657            "#,
658            schema = store.schema()
659        );
660        let ids: Vec<i64> = sqlx::query_scalar(&sql).bind(kind).fetch_all(pool).await?;
661        return store.retry_jobs_by_ids(pool, &ids).await;
662    }
663
664    let rows = sqlx::query_as::<_, JobRow>(
665        r#"
666        UPDATE awa.jobs
667        SET state = 'available', attempt = 0, run_at = now(),
668            finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL
669        WHERE kind = $1 AND state = 'failed'
670        RETURNING *
671        "#,
672    )
673    .bind(kind)
674    .fetch_all(pool)
675    .await?;
676
677    Ok(rows)
678}
679
680/// Retry all failed jobs in a given queue.
681pub async fn retry_failed_by_queue(pool: &PgPool, queue: &str) -> Result<Vec<JobRow>, AwaError> {
682    if let Some(store) = active_queue_storage(pool).await? {
683        let sql = format!(
684            r#"
685            SELECT job_id
686            FROM {schema}.done_entries
687            WHERE queue = $1
688              AND state = 'failed'
689            ORDER BY job_id ASC
690            "#,
691            schema = store.schema()
692        );
693        let ids: Vec<i64> = sqlx::query_scalar(&sql).bind(queue).fetch_all(pool).await?;
694        return store.retry_jobs_by_ids(pool, &ids).await;
695    }
696
697    let rows = sqlx::query_as::<_, JobRow>(
698        r#"
699        UPDATE awa.jobs
700        SET state = 'available', attempt = 0, run_at = now(),
701            finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL
702        WHERE queue = $1 AND state = 'failed'
703        RETURNING *
704        "#,
705    )
706    .bind(queue)
707    .fetch_all(pool)
708    .await?;
709
710    Ok(rows)
711}
712
713/// Discard (delete) all failed jobs of a given kind.
714pub async fn discard_failed(pool: &PgPool, kind: &str) -> Result<u64, AwaError> {
715    if let Some(store) = active_queue_storage(pool).await? {
716        return store.discard_failed_by_kind(pool, kind).await;
717    }
718
719    let result = sqlx::query("DELETE FROM awa.jobs WHERE kind = $1 AND state = 'failed'")
720        .bind(kind)
721        .execute(pool)
722        .await?;
723
724    Ok(result.rows_affected())
725}
726
727/// Pause a queue. Affects all workers immediately.
728pub async fn pause_queue<'e, E>(
729    executor: E,
730    queue: &str,
731    paused_by: Option<&str>,
732) -> Result<(), AwaError>
733where
734    E: PgExecutor<'e>,
735{
736    sqlx::query(
737        r#"
738        INSERT INTO awa.queue_meta (queue, paused, paused_at, paused_by)
739        VALUES ($1, TRUE, now(), $2)
740        ON CONFLICT (queue) DO UPDATE SET paused = TRUE, paused_at = now(), paused_by = $2
741        "#,
742    )
743    .bind(queue)
744    .bind(paused_by)
745    .execute(executor)
746    .await?;
747
748    Ok(())
749}
750
751/// Resume a paused queue.
752pub async fn resume_queue<'e, E>(executor: E, queue: &str) -> Result<(), AwaError>
753where
754    E: PgExecutor<'e>,
755{
756    sqlx::query("UPDATE awa.queue_meta SET paused = FALSE WHERE queue = $1")
757        .bind(queue)
758        .execute(executor)
759        .await?;
760
761    Ok(())
762}
763
764/// Drain a queue: cancel all non-running, non-terminal jobs.
765pub async fn drain_queue(pool: &PgPool, queue: &str) -> Result<u64, AwaError> {
766    if let Some(store) = active_queue_storage(pool).await? {
767        let sql = format!(
768            "{} \
769             SELECT job_id \
770             FROM current_jobs \
771             WHERE queue = $1 \
772               AND state IN ('available', 'scheduled', 'retryable', 'waiting_external') \
773             ORDER BY job_id ASC",
774            queue_storage_current_jobs_cte(store.schema())
775        );
776        let ids: Vec<i64> = sqlx::query_scalar(&sql).bind(queue).fetch_all(pool).await?;
777        return store
778            .cancel_jobs_by_ids(pool, &ids)
779            .await
780            .map(|rows| rows.len() as u64);
781    }
782
783    let result = sqlx::query(
784        r#"
785        UPDATE awa.jobs
786        SET state = 'cancelled', finalized_at = now(),
787            callback_id = NULL, callback_timeout_at = NULL,
788            callback_filter = NULL, callback_on_complete = NULL,
789            callback_on_fail = NULL, callback_transform = NULL
790        WHERE queue = $1 AND state IN ('available', 'scheduled', 'retryable', 'waiting_external')
791        "#,
792    )
793    .bind(queue)
794    .execute(pool)
795    .await?;
796
797    Ok(result.rows_affected())
798}
799
800/// Code-declared, operator-facing descriptor for a queue.
801#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
802pub struct QueueDescriptor {
803    pub display_name: Option<String>,
804    pub description: Option<String>,
805    pub owner: Option<String>,
806    pub docs_url: Option<String>,
807    pub tags: Vec<String>,
808    pub extra: serde_json::Value,
809}
810
811impl Default for QueueDescriptor {
812    fn default() -> Self {
813        Self {
814            display_name: None,
815            description: None,
816            owner: None,
817            docs_url: None,
818            tags: Vec::new(),
819            extra: serde_json::json!({}),
820        }
821    }
822}
823
824#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
825pub struct NamedQueueDescriptor {
826    pub queue: String,
827    #[serde(flatten)]
828    pub descriptor: QueueDescriptor,
829}
830
831#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
832pub struct JobKindDescriptor {
833    pub display_name: Option<String>,
834    pub description: Option<String>,
835    pub owner: Option<String>,
836    pub docs_url: Option<String>,
837    pub tags: Vec<String>,
838    pub extra: serde_json::Value,
839}
840
841impl Default for JobKindDescriptor {
842    fn default() -> Self {
843        Self {
844            display_name: None,
845            description: None,
846            owner: None,
847            docs_url: None,
848            tags: Vec::new(),
849            extra: serde_json::json!({}),
850        }
851    }
852}
853
854#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
855pub struct NamedJobKindDescriptor {
856    pub kind: String,
857    #[serde(flatten)]
858    pub descriptor: JobKindDescriptor,
859}
860
861/// Columns bound per descriptor row. Must match the order used in
862/// [`build_descriptor_upsert`] and the bind loop below.
863const DESCRIPTOR_PARAMS_PER_ROW: u32 = 9;
864
865/// Postgres caps bound parameters at 65535 per statement. 9 params × 7000 rows
866/// = 63k, so chunk below that with a comfortable margin.
867const DESCRIPTOR_BATCH_SIZE: usize = 5000;
868
869/// Build a batched INSERT ... VALUES (...), (...), ... ON CONFLICT upsert for
870/// a descriptor catalog. `table` is `awa.queue_descriptors` or
871/// `awa.job_kind_descriptors`; `pk` is `queue` or `kind`.
872fn build_descriptor_upsert(table: &str, pk: &str, count: usize) -> String {
873    let mut query = format!(
874        "INSERT INTO {table} (\
875             {pk}, display_name, description, owner, docs_url, tags, extra, \
876             descriptor_hash, sync_interval_ms, created_at, updated_at, last_seen_at\
877         ) VALUES "
878    );
879    let mut param_index = 1u32;
880    for i in 0..count {
881        if i > 0 {
882            query.push_str(", ");
883        }
884        query.push_str(&format!(
885            "(${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, ${}, now(), now(), now())",
886            param_index,
887            param_index + 1,
888            param_index + 2,
889            param_index + 3,
890            param_index + 4,
891            param_index + 5,
892            param_index + 6,
893            param_index + 7,
894            param_index + 8,
895        ));
896        param_index += DESCRIPTOR_PARAMS_PER_ROW;
897    }
898    // updated_at only advances when the hash changes — so an untouched
899    // descriptor keeps its original `updated_at` timestamp across ticks
900    // but still gets a fresh `last_seen_at` for liveness.
901    query.push_str(&format!(
902        " ON CONFLICT ({pk}) DO UPDATE SET \
903             display_name = EXCLUDED.display_name, \
904             description = EXCLUDED.description, \
905             owner = EXCLUDED.owner, \
906             docs_url = EXCLUDED.docs_url, \
907             tags = EXCLUDED.tags, \
908             extra = EXCLUDED.extra, \
909             descriptor_hash = EXCLUDED.descriptor_hash, \
910             sync_interval_ms = EXCLUDED.sync_interval_ms, \
911             updated_at = CASE \
912                 WHEN {table}.descriptor_hash IS DISTINCT FROM EXCLUDED.descriptor_hash \
913                 THEN now() ELSE {table}.updated_at \
914             END, \
915             last_seen_at = now()"
916    ));
917    query
918}
919
920/// Upsert queue descriptors declared by the worker runtime.
921///
922/// Descriptor rows are part of the control-plane catalog and intentionally
923/// separate from mutable queue runtime state in `awa.queue_meta`.
924///
925/// Descriptors are upserted in batched multi-row statements — one round-trip
926/// per [`DESCRIPTOR_BATCH_SIZE`] descriptors rather than one per descriptor.
927/// For typical fleets (≤100 queues / ≤500 kinds) that collapses to a single
928/// round-trip per sync call.
929pub async fn sync_queue_descriptors(
930    pool: &PgPool,
931    descriptors: &[NamedQueueDescriptor],
932    sync_interval: std::time::Duration,
933) -> Result<(), AwaError> {
934    if descriptors.is_empty() {
935        return Ok(());
936    }
937    let sync_interval_ms = sync_interval.as_millis() as i64;
938
939    for chunk in descriptors.chunks(DESCRIPTOR_BATCH_SIZE) {
940        let hashes: Vec<String> = chunk
941            .iter()
942            .map(|named| named.descriptor.descriptor_hash())
943            .collect();
944        let sql = build_descriptor_upsert("awa.queue_descriptors", "queue", chunk.len());
945        let mut query = sqlx::query(&sql);
946        for (named, hash) in chunk.iter().zip(hashes.iter()) {
947            query = query
948                .bind(&named.queue)
949                .bind(named.descriptor.display_name.as_deref())
950                .bind(named.descriptor.description.as_deref())
951                .bind(named.descriptor.owner.as_deref())
952                .bind(named.descriptor.docs_url.as_deref())
953                .bind(&named.descriptor.tags)
954                .bind(&named.descriptor.extra)
955                .bind(hash.as_str())
956                .bind(sync_interval_ms);
957        }
958        query.execute(pool).await?;
959    }
960
961    Ok(())
962}
963
964/// Upsert job-kind descriptors declared by the worker runtime. Batched the
965/// same way as [`sync_queue_descriptors`].
966pub async fn sync_job_kind_descriptors(
967    pool: &PgPool,
968    descriptors: &[NamedJobKindDescriptor],
969    sync_interval: std::time::Duration,
970) -> Result<(), AwaError> {
971    if descriptors.is_empty() {
972        return Ok(());
973    }
974    let sync_interval_ms = sync_interval.as_millis() as i64;
975
976    for chunk in descriptors.chunks(DESCRIPTOR_BATCH_SIZE) {
977        let hashes: Vec<String> = chunk
978            .iter()
979            .map(|named| named.descriptor.descriptor_hash())
980            .collect();
981        let sql = build_descriptor_upsert("awa.job_kind_descriptors", "kind", chunk.len());
982        let mut query = sqlx::query(&sql);
983        for (named, hash) in chunk.iter().zip(hashes.iter()) {
984            query = query
985                .bind(&named.kind)
986                .bind(named.descriptor.display_name.as_deref())
987                .bind(named.descriptor.description.as_deref())
988                .bind(named.descriptor.owner.as_deref())
989                .bind(named.descriptor.docs_url.as_deref())
990                .bind(&named.descriptor.tags)
991                .bind(&named.descriptor.extra)
992                .bind(hash.as_str())
993                .bind(sync_interval_ms);
994        }
995        query.execute(pool).await?;
996    }
997
998    Ok(())
999}
1000
1001impl QueueDescriptor {
1002    pub fn new() -> Self {
1003        Self::default()
1004    }
1005
1006    pub fn display_name(mut self, display_name: impl Into<String>) -> Self {
1007        self.display_name = Some(display_name.into());
1008        self
1009    }
1010
1011    pub fn description(mut self, description: impl Into<String>) -> Self {
1012        self.description = Some(description.into());
1013        self
1014    }
1015
1016    pub fn owner(mut self, owner: impl Into<String>) -> Self {
1017        self.owner = Some(owner.into());
1018        self
1019    }
1020
1021    pub fn docs_url(mut self, docs_url: impl Into<String>) -> Self {
1022        self.docs_url = Some(docs_url.into());
1023        self
1024    }
1025
1026    pub fn tags<I, S>(mut self, tags: I) -> Self
1027    where
1028        I: IntoIterator<Item = S>,
1029        S: Into<String>,
1030    {
1031        self.tags = tags.into_iter().map(Into::into).collect();
1032        self
1033    }
1034
1035    pub fn tag(mut self, tag: impl Into<String>) -> Self {
1036        self.tags.push(tag.into());
1037        self
1038    }
1039
1040    pub fn extra(mut self, extra: serde_json::Value) -> Self {
1041        self.extra = extra;
1042        self
1043    }
1044
1045    pub fn descriptor_hash(&self) -> String {
1046        let payload = serde_json::json!({
1047            "display_name": self.display_name,
1048            "description": self.description,
1049            "owner": self.owner,
1050            "docs_url": self.docs_url,
1051            "tags": self.tags,
1052            "extra": canonicalize_json(&self.extra),
1053        });
1054        let encoded = serde_json::to_vec(&payload)
1055            .expect("queue descriptor JSON serialization should not fail");
1056        hex::encode(blake3::hash(&encoded).as_bytes())
1057    }
1058}
1059
1060impl JobKindDescriptor {
1061    pub fn new() -> Self {
1062        Self::default()
1063    }
1064
1065    pub fn display_name(mut self, display_name: impl Into<String>) -> Self {
1066        self.display_name = Some(display_name.into());
1067        self
1068    }
1069
1070    pub fn description(mut self, description: impl Into<String>) -> Self {
1071        self.description = Some(description.into());
1072        self
1073    }
1074
1075    pub fn owner(mut self, owner: impl Into<String>) -> Self {
1076        self.owner = Some(owner.into());
1077        self
1078    }
1079
1080    pub fn docs_url(mut self, docs_url: impl Into<String>) -> Self {
1081        self.docs_url = Some(docs_url.into());
1082        self
1083    }
1084
1085    pub fn tags<I, S>(mut self, tags: I) -> Self
1086    where
1087        I: IntoIterator<Item = S>,
1088        S: Into<String>,
1089    {
1090        self.tags = tags.into_iter().map(Into::into).collect();
1091        self
1092    }
1093
1094    pub fn tag(mut self, tag: impl Into<String>) -> Self {
1095        self.tags.push(tag.into());
1096        self
1097    }
1098
1099    pub fn extra(mut self, extra: serde_json::Value) -> Self {
1100        self.extra = extra;
1101        self
1102    }
1103
1104    pub fn descriptor_hash(&self) -> String {
1105        let payload = serde_json::json!({
1106            "display_name": self.display_name,
1107            "description": self.description,
1108            "owner": self.owner,
1109            "docs_url": self.docs_url,
1110            "tags": self.tags,
1111            "extra": canonicalize_json(&self.extra),
1112        });
1113        let encoded = serde_json::to_vec(&payload)
1114            .expect("job kind descriptor JSON serialization should not fail");
1115        hex::encode(blake3::hash(&encoded).as_bytes())
1116    }
1117}
1118
1119fn canonicalize_json(value: &serde_json::Value) -> serde_json::Value {
1120    match value {
1121        serde_json::Value::Object(map) => {
1122            let mut keys: Vec<_> = map.keys().cloned().collect();
1123            keys.sort();
1124            let mut out = serde_json::Map::with_capacity(map.len());
1125            for key in keys {
1126                if let Some(child) = map.get(&key) {
1127                    out.insert(key, canonicalize_json(child));
1128                }
1129            }
1130            serde_json::Value::Object(out)
1131        }
1132        serde_json::Value::Array(values) => {
1133            serde_json::Value::Array(values.iter().map(canonicalize_json).collect())
1134        }
1135        _ => value.clone(),
1136    }
1137}
1138
1139#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
1140pub struct QueueOverview {
1141    pub queue: String,
1142    pub display_name: Option<String>,
1143    pub description: Option<String>,
1144    pub owner: Option<String>,
1145    pub docs_url: Option<String>,
1146    pub tags: Vec<String>,
1147    pub extra: serde_json::Value,
1148    pub descriptor_last_seen_at: Option<DateTime<Utc>>,
1149    pub descriptor_stale: bool,
1150    pub descriptor_mismatch: bool,
1151    /// All non-terminal jobs for the queue, including running and waiting_external.
1152    pub total_queued: i64,
1153    pub scheduled: i64,
1154    pub available: i64,
1155    pub retryable: i64,
1156    pub running: i64,
1157    pub failed: i64,
1158    pub waiting_external: i64,
1159    pub completed_last_hour: i64,
1160    pub lag_seconds: Option<f64>,
1161    pub paused: bool,
1162}
1163
1164#[derive(Debug, Clone, Serialize, sqlx::FromRow)]
1165pub struct JobKindOverview {
1166    pub kind: String,
1167    pub display_name: Option<String>,
1168    pub description: Option<String>,
1169    pub owner: Option<String>,
1170    pub docs_url: Option<String>,
1171    pub tags: Vec<String>,
1172    pub extra: serde_json::Value,
1173    pub descriptor_last_seen_at: Option<DateTime<Utc>>,
1174    pub descriptor_stale: bool,
1175    pub descriptor_mismatch: bool,
1176    pub job_count: i64,
1177    pub queue_count: i64,
1178    pub completed_last_hour: i64,
1179}
1180
1181/// Snapshot of a per-queue rate limit configuration.
1182#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1183pub struct RateLimitSnapshot {
1184    pub max_rate: f64,
1185    pub burst: u32,
1186}
1187
1188/// Runtime concurrency mode for a queue.
1189#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
1190#[serde(rename_all = "snake_case")]
1191pub enum QueueRuntimeMode {
1192    HardReserved,
1193    Weighted,
1194}
1195
1196/// Per-queue configuration published by a worker runtime instance.
1197#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1198pub struct QueueRuntimeConfigSnapshot {
1199    pub mode: QueueRuntimeMode,
1200    pub max_workers: Option<u32>,
1201    pub min_workers: Option<u32>,
1202    pub weight: Option<u32>,
1203    pub global_max_workers: Option<u32>,
1204    pub poll_interval_ms: u64,
1205    pub deadline_duration_secs: u64,
1206    pub priority_aging_interval_secs: u64,
1207    #[serde(default)]
1208    pub claimers: Option<u16>,
1209    #[serde(default)]
1210    pub claim_batch_size: Option<usize>,
1211    #[serde(default)]
1212    pub dlq_enabled: Option<bool>,
1213    pub rate_limit: Option<RateLimitSnapshot>,
1214}
1215
1216/// Runtime state for one queue on one worker instance.
1217#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
1218pub struct QueueRuntimeSnapshot {
1219    pub queue: String,
1220    pub in_flight: u32,
1221    pub overflow_held: Option<u32>,
1222    pub config: QueueRuntimeConfigSnapshot,
1223}
1224
1225#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash)]
1226#[serde(rename_all = "snake_case")]
1227pub enum StorageCapability {
1228    Canonical,
1229    CanonicalDrainOnly,
1230    QueueStorage,
1231}
1232
1233impl StorageCapability {
1234    pub fn as_str(self) -> &'static str {
1235        match self {
1236            Self::Canonical => "canonical",
1237            Self::CanonicalDrainOnly => "canonical_drain_only",
1238            Self::QueueStorage => "queue_storage",
1239        }
1240    }
1241}
1242
1243impl fmt::Display for StorageCapability {
1244    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1245        f.write_str(self.as_str())
1246    }
1247}
1248
1249impl FromStr for StorageCapability {
1250    type Err = String;
1251
1252    fn from_str(value: &str) -> Result<Self, Self::Err> {
1253        match value {
1254            "canonical" => Ok(Self::Canonical),
1255            "canonical_drain_only" => Ok(Self::CanonicalDrainOnly),
1256            "queue_storage" => Ok(Self::QueueStorage),
1257            _ => Err(value.to_string()),
1258        }
1259    }
1260}
1261
1262/// Operator-selected execution role used during a `0.5.x → 0.6` storage
1263/// transition. Persisted on `awa.runtime_instances` so the SQL gate for
1264/// `enter_mixed_transition` can require a runtime that will actually
1265/// execute queue-storage work after routing flips, rather than only
1266/// inspecting the more permissive `storage_capability` snapshot
1267/// (auto-role runtimes report `queue_storage` while in
1268/// canonical/prepared but downgrade to `canonical_drain_only` post-flip).
1269#[derive(Debug, Clone, Copy, Default, Serialize, Deserialize, PartialEq, Eq, Hash)]
1270#[serde(rename_all = "snake_case")]
1271pub enum TransitionRole {
1272    /// Default. Resolves storage at startup from `awa.storage_status()`
1273    /// and follows whatever the cluster is doing. Auto runtimes started
1274    /// before mixed_transition cannot satisfy the queue-storage executor
1275    /// gate; they're meant to drain canonical first and then have new
1276    /// auto runtimes spin up post-flip to take over queue_storage work.
1277    #[default]
1278    Auto,
1279    /// Stay on canonical execution regardless of transition state. Used
1280    /// for runtimes that should keep draining canonical even after
1281    /// routing flips.
1282    CanonicalDrain,
1283    /// Always execute queue-storage work, including pre-flip when state
1284    /// is `prepared`. At least one live runtime in this role is required
1285    /// before `awa.storage_enter_mixed_transition()` will succeed.
1286    QueueStorageTarget,
1287}
1288
1289impl TransitionRole {
1290    pub fn as_str(self) -> &'static str {
1291        match self {
1292            Self::Auto => "auto",
1293            Self::CanonicalDrain => "canonical_drain",
1294            Self::QueueStorageTarget => "queue_storage_target",
1295        }
1296    }
1297}
1298
1299impl fmt::Display for TransitionRole {
1300    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
1301        f.write_str(self.as_str())
1302    }
1303}
1304
1305impl FromStr for TransitionRole {
1306    type Err = String;
1307
1308    fn from_str(value: &str) -> Result<Self, Self::Err> {
1309        match value {
1310            "auto" => Ok(Self::Auto),
1311            "canonical_drain" => Ok(Self::CanonicalDrain),
1312            "queue_storage_target" => Ok(Self::QueueStorageTarget),
1313            _ => Err(value.to_string()),
1314        }
1315    }
1316}
1317
1318/// Data written by a worker runtime into the observability snapshot table.
1319#[derive(Debug, Clone, Serialize, Deserialize)]
1320pub struct RuntimeSnapshotInput {
1321    pub instance_id: Uuid,
1322    pub hostname: Option<String>,
1323    pub pid: i32,
1324    pub version: String,
1325    pub storage_capability: StorageCapability,
1326    #[serde(default)]
1327    pub transition_role: TransitionRole,
1328    pub started_at: DateTime<Utc>,
1329    pub snapshot_interval_ms: i64,
1330    pub healthy: bool,
1331    pub postgres_connected: bool,
1332    pub poll_loop_alive: bool,
1333    pub heartbeat_alive: bool,
1334    pub maintenance_alive: bool,
1335    pub shutting_down: bool,
1336    pub leader: bool,
1337    pub global_max_workers: Option<u32>,
1338    pub queues: Vec<QueueRuntimeSnapshot>,
1339    pub queue_descriptor_hashes: HashMap<String, String>,
1340    pub job_kind_descriptor_hashes: HashMap<String, String>,
1341}
1342
1343/// A worker runtime instance as exposed through the admin API.
1344#[derive(Debug, Clone, Serialize, Deserialize)]
1345pub struct RuntimeInstance {
1346    pub instance_id: Uuid,
1347    pub hostname: Option<String>,
1348    pub pid: i32,
1349    pub version: String,
1350    pub storage_capability: StorageCapability,
1351    #[serde(default)]
1352    pub transition_role: TransitionRole,
1353    pub started_at: DateTime<Utc>,
1354    pub last_seen_at: DateTime<Utc>,
1355    pub snapshot_interval_ms: i64,
1356    pub stale: bool,
1357    pub healthy: bool,
1358    pub postgres_connected: bool,
1359    pub poll_loop_alive: bool,
1360    pub heartbeat_alive: bool,
1361    pub maintenance_alive: bool,
1362    pub shutting_down: bool,
1363    pub leader: bool,
1364    pub global_max_workers: Option<u32>,
1365    pub queues: Vec<QueueRuntimeSnapshot>,
1366}
1367
1368impl RuntimeInstance {
1369    fn stale_cutoff(interval_ms: i64) -> Duration {
1370        let interval_ms = max(interval_ms, 1_000);
1371        Duration::milliseconds(max(interval_ms.saturating_mul(3), 30_000))
1372    }
1373
1374    fn from_db_row(row: RuntimeInstanceRow, now: DateTime<Utc>) -> Result<Self, AwaError> {
1375        let stale = row.last_seen_at + Self::stale_cutoff(row.snapshot_interval_ms) < now;
1376        let storage_capability =
1377            StorageCapability::from_str(&row.storage_capability).map_err(|value| {
1378                AwaError::Validation(format!(
1379                    "invalid storage capability in runtime_instances: {value}"
1380                ))
1381            })?;
1382        let transition_role = TransitionRole::from_str(&row.transition_role).map_err(|value| {
1383            AwaError::Validation(format!(
1384                "invalid transition_role in runtime_instances: {value}"
1385            ))
1386        })?;
1387        Ok(Self {
1388            instance_id: row.instance_id,
1389            hostname: row.hostname,
1390            pid: row.pid,
1391            version: row.version,
1392            storage_capability,
1393            transition_role,
1394            started_at: row.started_at,
1395            last_seen_at: row.last_seen_at,
1396            snapshot_interval_ms: row.snapshot_interval_ms,
1397            stale,
1398            healthy: row.healthy,
1399            postgres_connected: row.postgres_connected,
1400            poll_loop_alive: row.poll_loop_alive,
1401            heartbeat_alive: row.heartbeat_alive,
1402            maintenance_alive: row.maintenance_alive,
1403            shutting_down: row.shutting_down,
1404            leader: row.leader,
1405            global_max_workers: row.global_max_workers.map(|v| v as u32),
1406            queues: row.queues.0,
1407        })
1408    }
1409}
1410
1411/// Cluster-wide runtime overview.
1412#[derive(Debug, Clone, Serialize, Deserialize)]
1413pub struct RuntimeOverview {
1414    pub total_instances: usize,
1415    pub live_instances: usize,
1416    pub stale_instances: usize,
1417    pub healthy_instances: usize,
1418    pub leader_instances: usize,
1419    pub instances: Vec<RuntimeInstance>,
1420}
1421
1422/// Queue-centric runtime/config summary aggregated across worker instances.
1423#[derive(Debug, Clone, Serialize, Deserialize)]
1424pub struct QueueRuntimeSummary {
1425    pub queue: String,
1426    pub instance_count: usize,
1427    pub live_instances: usize,
1428    pub stale_instances: usize,
1429    pub healthy_instances: usize,
1430    pub total_in_flight: u64,
1431    pub overflow_held_total: Option<u64>,
1432    pub config_mismatch: bool,
1433    pub config: Option<QueueRuntimeConfigSnapshot>,
1434}
1435
1436fn queue_runtime_configs_match(
1437    left: &QueueRuntimeConfigSnapshot,
1438    right: &QueueRuntimeConfigSnapshot,
1439) -> bool {
1440    left.mode == right.mode
1441        && left.max_workers == right.max_workers
1442        && left.min_workers == right.min_workers
1443        && left.weight == right.weight
1444        && left.global_max_workers == right.global_max_workers
1445        && left.poll_interval_ms == right.poll_interval_ms
1446        && left.deadline_duration_secs == right.deadline_duration_secs
1447        && left.priority_aging_interval_secs == right.priority_aging_interval_secs
1448        && left.rate_limit == right.rate_limit
1449        && (left.dlq_enabled.is_none()
1450            || right.dlq_enabled.is_none()
1451            || left.dlq_enabled == right.dlq_enabled)
1452}
1453
1454fn representative_queue_runtime_config(
1455    configs: &[QueueRuntimeConfigSnapshot],
1456) -> Option<QueueRuntimeConfigSnapshot> {
1457    let first = configs.first()?;
1458    if first.dlq_enabled.is_some() {
1459        return Some(first.clone());
1460    }
1461
1462    configs
1463        .iter()
1464        .find(|candidate| {
1465            candidate.dlq_enabled.is_some() && queue_runtime_configs_match(first, candidate)
1466        })
1467        .cloned()
1468        .or_else(|| Some(first.clone()))
1469}
1470
1471#[derive(Debug, sqlx::FromRow)]
1472struct RuntimeInstanceRow {
1473    instance_id: Uuid,
1474    hostname: Option<String>,
1475    pid: i32,
1476    version: String,
1477    storage_capability: String,
1478    transition_role: String,
1479    started_at: DateTime<Utc>,
1480    last_seen_at: DateTime<Utc>,
1481    snapshot_interval_ms: i64,
1482    healthy: bool,
1483    postgres_connected: bool,
1484    poll_loop_alive: bool,
1485    heartbeat_alive: bool,
1486    maintenance_alive: bool,
1487    shutting_down: bool,
1488    leader: bool,
1489    global_max_workers: Option<i32>,
1490    queues: Json<Vec<QueueRuntimeSnapshot>>,
1491}
1492
1493/// Upsert a runtime observability snapshot for one worker instance.
1494pub async fn upsert_runtime_snapshot<'e, E>(
1495    executor: E,
1496    snapshot: &RuntimeSnapshotInput,
1497) -> Result<(), AwaError>
1498where
1499    E: PgExecutor<'e>,
1500{
1501    sqlx::query(
1502        r#"
1503        INSERT INTO awa.runtime_instances (
1504            instance_id,
1505            hostname,
1506            pid,
1507            version,
1508            storage_capability,
1509            transition_role,
1510            started_at,
1511            last_seen_at,
1512            snapshot_interval_ms,
1513            healthy,
1514            postgres_connected,
1515            poll_loop_alive,
1516            heartbeat_alive,
1517            maintenance_alive,
1518            shutting_down,
1519            leader,
1520            global_max_workers,
1521            queues,
1522            queue_descriptor_hashes,
1523            job_kind_descriptor_hashes
1524        )
1525        VALUES (
1526            $1, $2, $3, $4, $5, $6, $7, now(), $8, $9, $10, $11, $12, $13, $14, $15, $16, $17, $18, $19
1527        )
1528        ON CONFLICT (instance_id) DO UPDATE SET
1529            hostname = EXCLUDED.hostname,
1530            pid = EXCLUDED.pid,
1531            version = EXCLUDED.version,
1532            storage_capability = EXCLUDED.storage_capability,
1533            transition_role = EXCLUDED.transition_role,
1534            started_at = EXCLUDED.started_at,
1535            last_seen_at = now(),
1536            snapshot_interval_ms = EXCLUDED.snapshot_interval_ms,
1537            healthy = EXCLUDED.healthy,
1538            postgres_connected = EXCLUDED.postgres_connected,
1539            poll_loop_alive = EXCLUDED.poll_loop_alive,
1540            heartbeat_alive = EXCLUDED.heartbeat_alive,
1541            maintenance_alive = EXCLUDED.maintenance_alive,
1542            shutting_down = EXCLUDED.shutting_down,
1543            leader = EXCLUDED.leader,
1544            global_max_workers = EXCLUDED.global_max_workers,
1545            queues = EXCLUDED.queues,
1546            queue_descriptor_hashes = EXCLUDED.queue_descriptor_hashes,
1547            job_kind_descriptor_hashes = EXCLUDED.job_kind_descriptor_hashes
1548        "#,
1549    )
1550    .bind(snapshot.instance_id)
1551    .bind(snapshot.hostname.as_deref())
1552    .bind(snapshot.pid)
1553    .bind(&snapshot.version)
1554    .bind(snapshot.storage_capability.as_str())
1555    .bind(snapshot.transition_role.as_str())
1556    .bind(snapshot.started_at)
1557    .bind(snapshot.snapshot_interval_ms)
1558    .bind(snapshot.healthy)
1559    .bind(snapshot.postgres_connected)
1560    .bind(snapshot.poll_loop_alive)
1561    .bind(snapshot.heartbeat_alive)
1562    .bind(snapshot.maintenance_alive)
1563    .bind(snapshot.shutting_down)
1564    .bind(snapshot.leader)
1565    .bind(snapshot.global_max_workers.map(|v| v as i32))
1566    .bind(Json(&snapshot.queues))
1567    .bind(Json(&snapshot.queue_descriptor_hashes))
1568    .bind(Json(&snapshot.job_kind_descriptor_hashes))
1569    .execute(executor)
1570    .await?;
1571
1572    Ok(())
1573}
1574
1575/// Opportunistically delete long-stale runtime snapshot rows.
1576pub async fn cleanup_runtime_snapshots<'e, E>(
1577    executor: E,
1578    max_age: Duration,
1579) -> Result<u64, AwaError>
1580where
1581    E: PgExecutor<'e>,
1582{
1583    let seconds = max(max_age.num_seconds(), 1);
1584    let result = sqlx::query(
1585        "DELETE FROM awa.runtime_instances WHERE last_seen_at < now() - make_interval(secs => $1)",
1586    )
1587    .bind(seconds)
1588    .execute(executor)
1589    .await?;
1590
1591    Ok(result.rows_affected())
1592}
1593
1594/// Delete descriptor rows whose `last_seen_at` is older than `max_age`.
1595///
1596/// Intended to run on the maintenance leader's cleanup cycle — a descriptor
1597/// whose declaring code has been retired would otherwise linger in the
1598/// catalog forever, showing as permanently stale. Descriptors are
1599/// best-effort (no FK from `awa.jobs*`), so deletion is safe: if a worker
1600/// later re-declares the same queue / kind, the next sync recreates the
1601/// row from the declaration.
1602///
1603/// `table` must be `awa.queue_descriptors` or `awa.job_kind_descriptors`;
1604/// the caller is expected to dispatch both in turn.
1605pub async fn cleanup_stale_descriptors<'e, E>(
1606    executor: E,
1607    table: &str,
1608    max_age: Duration,
1609) -> Result<u64, AwaError>
1610where
1611    E: PgExecutor<'e>,
1612{
1613    if !matches!(table, "awa.queue_descriptors" | "awa.job_kind_descriptors") {
1614        return Err(AwaError::Validation(format!(
1615            "cleanup_stale_descriptors: unknown table {table:?}"
1616        )));
1617    }
1618    let seconds = max(max_age.num_seconds(), 1);
1619    // Table name is an authenticated literal from the match above — safe
1620    // to interpolate into the statement.
1621    let sql = format!("DELETE FROM {table} WHERE last_seen_at < now() - make_interval(secs => $1)");
1622    let result = sqlx::query(&sql).bind(seconds).execute(executor).await?;
1623    Ok(result.rows_affected())
1624}
1625
1626/// List all runtime instances ordered with leader/live instances first.
1627pub async fn list_runtime_instances<'e, E>(executor: E) -> Result<Vec<RuntimeInstance>, AwaError>
1628where
1629    E: PgExecutor<'e>,
1630{
1631    let rows = sqlx::query_as::<_, RuntimeInstanceRow>(
1632        r#"
1633        SELECT
1634            instance_id,
1635            hostname,
1636            pid,
1637            version,
1638            storage_capability,
1639            transition_role,
1640            started_at,
1641            last_seen_at,
1642            snapshot_interval_ms,
1643            healthy,
1644            postgres_connected,
1645            poll_loop_alive,
1646            heartbeat_alive,
1647            maintenance_alive,
1648            shutting_down,
1649            leader,
1650            global_max_workers,
1651            queues
1652        FROM awa.runtime_instances
1653        ORDER BY leader DESC, last_seen_at DESC, started_at DESC
1654        "#,
1655    )
1656    .fetch_all(executor)
1657    .await?;
1658
1659    let now = Utc::now();
1660    rows.into_iter()
1661        .map(|row| RuntimeInstance::from_db_row(row, now))
1662        .collect()
1663}
1664
1665/// Cluster runtime overview with instance list.
1666pub async fn runtime_overview<'e, E>(executor: E) -> Result<RuntimeOverview, AwaError>
1667where
1668    E: PgExecutor<'e>,
1669{
1670    let instances = list_runtime_instances(executor).await?;
1671    let total_instances = instances.len();
1672    let stale_instances = instances.iter().filter(|i| i.stale).count();
1673    let live_instances = total_instances.saturating_sub(stale_instances);
1674    let healthy_instances = instances.iter().filter(|i| !i.stale && i.healthy).count();
1675    let leader_instances = instances.iter().filter(|i| !i.stale && i.leader).count();
1676
1677    Ok(RuntimeOverview {
1678        total_instances,
1679        live_instances,
1680        stale_instances,
1681        healthy_instances,
1682        leader_instances,
1683        instances,
1684    })
1685}
1686
1687/// Queue runtime/config summary aggregated across worker snapshots.
1688pub async fn queue_runtime_summary<'e, E>(executor: E) -> Result<Vec<QueueRuntimeSummary>, AwaError>
1689where
1690    E: PgExecutor<'e>,
1691{
1692    let instances = list_runtime_instances(executor).await?;
1693    let mut by_queue: HashMap<String, Vec<(bool, bool, QueueRuntimeSnapshot)>> = HashMap::new();
1694
1695    for instance in instances {
1696        let is_live = !instance.stale;
1697        let is_healthy = is_live && instance.healthy;
1698        for queue in instance.queues {
1699            by_queue
1700                .entry(queue.queue.clone())
1701                .or_default()
1702                .push((is_live, is_healthy, queue));
1703        }
1704    }
1705
1706    let mut summaries: Vec<_> = by_queue
1707        .into_iter()
1708        .map(|(queue, entries)| {
1709            let instance_count = entries.len();
1710            let live_instances = entries.iter().filter(|(live, _, _)| *live).count();
1711            let stale_instances = instance_count.saturating_sub(live_instances);
1712            let healthy_instances = entries.iter().filter(|(_, healthy, _)| *healthy).count();
1713            let total_in_flight = entries
1714                .iter()
1715                .filter(|(live, _, _)| *live)
1716                .map(|(_, _, queue)| u64::from(queue.in_flight))
1717                .sum();
1718
1719            let overflow_total: u64 = entries
1720                .iter()
1721                .filter(|(live, _, _)| *live)
1722                .filter_map(|(_, _, queue)| queue.overflow_held.map(u64::from))
1723                .sum();
1724
1725            let live_configs: Vec<_> = entries
1726                .iter()
1727                .filter(|(live, _, _)| *live)
1728                .map(|(_, _, queue)| queue.config.clone())
1729                .collect();
1730            let config_candidates = if live_configs.is_empty() {
1731                entries
1732                    .iter()
1733                    .map(|(_, _, queue)| queue.config.clone())
1734                    .collect::<Vec<_>>()
1735            } else {
1736                live_configs
1737            };
1738            let config = representative_queue_runtime_config(&config_candidates);
1739            let config_mismatch = config_candidates.iter().skip(1).any(|candidate| {
1740                config
1741                    .as_ref()
1742                    .is_some_and(|cfg| !queue_runtime_configs_match(cfg, candidate))
1743            });
1744
1745            QueueRuntimeSummary {
1746                queue,
1747                instance_count,
1748                live_instances,
1749                stale_instances,
1750                healthy_instances,
1751                total_in_flight,
1752                overflow_held_total: config
1753                    .as_ref()
1754                    .filter(|cfg| cfg.mode == QueueRuntimeMode::Weighted)
1755                    .map(|_| overflow_total),
1756                config_mismatch,
1757                config,
1758            }
1759        })
1760        .collect();
1761
1762    summaries.sort_by(|a, b| a.queue.cmp(&b.queue));
1763    Ok(summaries)
1764}
1765
1766/// Get queue overviews for all known queues.
1767///
1768/// Hybrid read: per-state counts come from the `queue_state_counts`
1769/// cache table (eventually consistent, ~2s lag), while `lag_seconds`
1770/// and `completed_last_hour` are computed live from `jobs_hot`.
1771///
1772/// The cache is kept fresh by the maintenance leader's dirty-key
1773/// recompute (~2s) and full reconciliation (~60s). Also warmed during
1774/// `migrate()`.
1775///
1776/// For exact cached counts in tests without a running maintenance
1777/// leader, call `flush_dirty_admin_metadata()` first.
1778pub async fn queue_overviews(pool: &PgPool) -> Result<Vec<QueueOverview>, AwaError> {
1779    if let Some(store) = active_queue_storage(pool).await? {
1780        let sql = format!(
1781            r#"
1782            {current_jobs_cte},
1783            all_queues AS (
1784                SELECT queue FROM current_jobs
1785                UNION
1786                SELECT queue FROM awa.queue_descriptors
1787                UNION
1788                SELECT queue FROM awa.queue_meta
1789                UNION
1790                SELECT DISTINCT descriptor.key AS queue
1791                FROM awa.runtime_instances runtime
1792                CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
1793            ),
1794            live_queue_descriptor_variants AS (
1795                SELECT
1796                    descriptor.key AS queue,
1797                    count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
1798                FROM awa.runtime_instances runtime
1799                CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
1800                WHERE runtime.last_seen_at + make_interval(
1801                    secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
1802                ) >= now()
1803                GROUP BY descriptor.key
1804            ),
1805            queue_counts AS (
1806                SELECT
1807                    queue,
1808                    count(*) FILTER (WHERE state = 'scheduled')::bigint AS scheduled,
1809                    count(*) FILTER (WHERE state = 'available')::bigint AS available,
1810                    count(*) FILTER (WHERE state = 'retryable')::bigint AS retryable,
1811                    count(*) FILTER (WHERE state = 'running')::bigint AS running,
1812                    count(*) FILTER (WHERE state = 'failed')::bigint AS failed,
1813                    count(*) FILTER (WHERE state = 'waiting_external')::bigint AS waiting_external
1814                FROM current_jobs
1815                GROUP BY queue
1816            ),
1817            available_lag AS (
1818                SELECT
1819                    queue,
1820                    EXTRACT(EPOCH FROM (clock_timestamp() - min(run_at)))::float8 AS lag_seconds
1821                FROM current_jobs
1822                WHERE state = 'available'
1823                GROUP BY queue
1824            ),
1825            completed_recent AS (
1826                SELECT
1827                    queue,
1828                    count(*)::bigint AS completed_last_hour
1829                FROM current_jobs
1830                WHERE state = 'completed'
1831                  AND finalized_at > clock_timestamp() - interval '1 hour'
1832                GROUP BY queue
1833            )
1834            SELECT
1835                q.queue,
1836                qd.display_name,
1837                qd.description,
1838                qd.owner,
1839                qd.docs_url,
1840                COALESCE(qd.tags, ARRAY[]::text[]) AS tags,
1841                COALESCE(qd.extra, '{{}}'::jsonb) AS extra,
1842                qd.last_seen_at AS descriptor_last_seen_at,
1843                CASE
1844                    WHEN qd.last_seen_at IS NULL THEN FALSE
1845                    ELSE qd.last_seen_at + make_interval(
1846                        secs => GREATEST(((COALESCE(qd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
1847                    ) < now()
1848                END AS descriptor_stale,
1849                COALESCE(qdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
1850                COALESCE(qc.scheduled + qc.available + qc.running + qc.retryable + qc.waiting_external, 0) AS total_queued,
1851                COALESCE(qc.scheduled, 0) AS scheduled,
1852                COALESCE(qc.available, 0) AS available,
1853                COALESCE(qc.retryable, 0) AS retryable,
1854                COALESCE(qc.running, 0) AS running,
1855                COALESCE(qc.failed, 0) AS failed,
1856                COALESCE(qc.waiting_external, 0) AS waiting_external,
1857                COALESCE(cr.completed_last_hour, 0) AS completed_last_hour,
1858                al.lag_seconds,
1859                COALESCE(qm.paused, FALSE) AS paused
1860            FROM all_queues q
1861            LEFT JOIN queue_counts qc ON qc.queue = q.queue
1862            LEFT JOIN awa.queue_descriptors qd ON qd.queue = q.queue
1863            LEFT JOIN live_queue_descriptor_variants qdv ON qdv.queue = q.queue
1864            LEFT JOIN available_lag al ON al.queue = q.queue
1865            LEFT JOIN completed_recent cr ON cr.queue = q.queue
1866            LEFT JOIN awa.queue_meta qm ON qm.queue = q.queue
1867            ORDER BY q.queue
1868            "#,
1869            current_jobs_cte = queue_storage_current_jobs_cte(store.schema())
1870        );
1871
1872        let rows = sqlx::query_as::<_, QueueOverview>(&sql)
1873            .fetch_all(pool)
1874            .await?;
1875        return Ok(rows);
1876    }
1877
1878    let rows = sqlx::query_as::<_, QueueOverview>(
1879        r#"
1880        WITH all_queues AS (
1881            -- Union every source of queue-name knowledge so /queues never
1882            -- hides a queue that exists *somewhere* in the system:
1883            --   - queue_state_counts: has observed jobs
1884            --   - queue_descriptors:   declared by worker code
1885            --   - queue_meta:          pause state recorded by an operator
1886            --   - runtime_instances:   registered by a currently-running worker
1887            SELECT queue FROM awa.queue_state_counts
1888            UNION
1889            SELECT queue FROM awa.queue_descriptors
1890            UNION
1891            SELECT queue FROM awa.queue_meta
1892            UNION
1893            SELECT DISTINCT descriptor.key AS queue
1894            FROM awa.runtime_instances runtime
1895            CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
1896        ),
1897        live_queue_descriptor_variants AS (
1898            SELECT
1899                descriptor.key AS queue,
1900                count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
1901            FROM awa.runtime_instances runtime
1902            CROSS JOIN LATERAL jsonb_each_text(runtime.queue_descriptor_hashes) AS descriptor(key, value)
1903            WHERE runtime.last_seen_at + make_interval(
1904                secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
1905            ) >= now()
1906            GROUP BY descriptor.key
1907        ),
1908        available_lag AS (
1909            SELECT
1910                queue,
1911                EXTRACT(EPOCH FROM (now() - min(run_at)))::float8 AS lag_seconds
1912            FROM awa.jobs_hot
1913            WHERE state = 'available'
1914            GROUP BY queue
1915        ),
1916        completed_recent AS (
1917            SELECT
1918                queue,
1919                count(*)::bigint AS completed_last_hour
1920            FROM awa.jobs_hot
1921            WHERE state = 'completed'
1922              AND finalized_at > now() - interval '1 hour'
1923            GROUP BY queue
1924        )
1925        SELECT
1926            q.queue,
1927            qd.display_name,
1928            qd.description,
1929            qd.owner,
1930            qd.docs_url,
1931            COALESCE(qd.tags, ARRAY[]::text[]) AS tags,
1932            COALESCE(qd.extra, '{}'::jsonb) AS extra,
1933            qd.last_seen_at AS descriptor_last_seen_at,
1934            CASE
1935                WHEN qd.last_seen_at IS NULL THEN FALSE
1936                ELSE qd.last_seen_at + make_interval(
1937                    secs => GREATEST(((COALESCE(qd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
1938                ) < now()
1939            END AS descriptor_stale,
1940            COALESCE(qdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
1941            COALESCE(qs.scheduled + qs.available + qs.running + qs.retryable + qs.waiting_external, 0) AS total_queued,
1942            COALESCE(qs.scheduled, 0) AS scheduled,
1943            COALESCE(qs.available, 0) AS available,
1944            COALESCE(qs.retryable, 0) AS retryable,
1945            COALESCE(qs.running, 0) AS running,
1946            COALESCE(qs.failed, 0) AS failed,
1947            COALESCE(qs.waiting_external, 0) AS waiting_external,
1948            COALESCE(cr.completed_last_hour, 0) AS completed_last_hour,
1949            al.lag_seconds,
1950            COALESCE(qm.paused, FALSE) AS paused
1951        FROM all_queues q
1952        LEFT JOIN awa.queue_state_counts qs ON qs.queue = q.queue
1953        LEFT JOIN awa.queue_descriptors qd ON qd.queue = q.queue
1954        LEFT JOIN live_queue_descriptor_variants qdv ON qdv.queue = q.queue
1955        LEFT JOIN available_lag al ON al.queue = q.queue
1956        LEFT JOIN completed_recent cr ON cr.queue = q.queue
1957        LEFT JOIN awa.queue_meta qm ON qm.queue = q.queue
1958        ORDER BY q.queue
1959        "#,
1960    )
1961    .fetch_all(pool)
1962    .await?;
1963
1964    Ok(rows)
1965}
1966
1967/// Get one queue overview by name.
1968pub async fn queue_overview(pool: &PgPool, queue: &str) -> Result<Option<QueueOverview>, AwaError> {
1969    let rows = queue_overviews(pool).await?;
1970    Ok(rows.into_iter().find(|row| row.queue == queue))
1971}
1972
1973pub async fn queue_descriptors_for_names<'e, E>(
1974    executor: E,
1975    queues: &[String],
1976) -> Result<HashMap<String, QueueDescriptor>, AwaError>
1977where
1978    E: PgExecutor<'e>,
1979{
1980    if queues.is_empty() {
1981        return Ok(HashMap::new());
1982    }
1983
1984    let rows = sqlx::query_as::<
1985        _,
1986        (
1987            String,
1988            Option<String>,
1989            Option<String>,
1990            Option<String>,
1991            Option<String>,
1992            Vec<String>,
1993            serde_json::Value,
1994        ),
1995    >(
1996        r#"
1997        SELECT
1998            queue,
1999            display_name,
2000            description,
2001            owner,
2002            docs_url,
2003            tags,
2004            extra
2005        FROM awa.queue_descriptors
2006        WHERE queue = ANY($1)
2007        "#,
2008    )
2009    .bind(queues)
2010    .fetch_all(executor)
2011    .await?;
2012
2013    Ok(rows
2014        .into_iter()
2015        .map(
2016            |(queue, display_name, description, owner, docs_url, tags, extra)| {
2017                (
2018                    queue,
2019                    QueueDescriptor {
2020                        display_name,
2021                        description,
2022                        owner,
2023                        docs_url,
2024                        tags,
2025                        extra,
2026                    },
2027                )
2028            },
2029        )
2030        .collect())
2031}
2032
2033pub async fn job_kind_descriptors_for_names<'e, E>(
2034    executor: E,
2035    kinds: &[String],
2036) -> Result<HashMap<String, JobKindDescriptor>, AwaError>
2037where
2038    E: PgExecutor<'e>,
2039{
2040    if kinds.is_empty() {
2041        return Ok(HashMap::new());
2042    }
2043
2044    let rows = sqlx::query_as::<
2045        _,
2046        (
2047            String,
2048            Option<String>,
2049            Option<String>,
2050            Option<String>,
2051            Option<String>,
2052            Vec<String>,
2053            serde_json::Value,
2054        ),
2055    >(
2056        r#"
2057        SELECT
2058            kind,
2059            display_name,
2060            description,
2061            owner,
2062            docs_url,
2063            tags,
2064            extra
2065        FROM awa.job_kind_descriptors
2066        WHERE kind = ANY($1)
2067        "#,
2068    )
2069    .bind(kinds)
2070    .fetch_all(executor)
2071    .await?;
2072
2073    Ok(rows
2074        .into_iter()
2075        .map(
2076            |(kind, display_name, description, owner, docs_url, tags, extra)| {
2077                (
2078                    kind,
2079                    JobKindDescriptor {
2080                        display_name,
2081                        description,
2082                        owner,
2083                        docs_url,
2084                        tags,
2085                        extra,
2086                    },
2087                )
2088            },
2089        )
2090        .collect())
2091}
2092
2093/// List jobs with optional filters.
2094#[derive(Debug, Clone, Default, Serialize)]
2095pub struct ListJobsFilter {
2096    pub state: Option<JobState>,
2097    pub kind: Option<String>,
2098    pub queue: Option<String>,
2099    pub tag: Option<String>,
2100    pub before_id: Option<i64>,
2101    pub limit: Option<i64>,
2102}
2103
2104/// List jobs matching the given filter.
2105pub async fn list_jobs(pool: &PgPool, filter: &ListJobsFilter) -> Result<Vec<JobRow>, AwaError> {
2106    if let Some(store) = active_queue_storage(pool).await? {
2107        return list_queue_storage_jobs(&store, pool, filter).await;
2108    }
2109
2110    let limit = filter.limit.unwrap_or(100);
2111
2112    let rows = sqlx::query_as::<_, JobRow>(
2113        r#"
2114        SELECT * FROM awa.jobs
2115        WHERE ($1::awa.job_state IS NULL OR state = $1)
2116          AND ($2::text IS NULL OR kind = $2)
2117          AND ($3::text IS NULL OR queue = $3)
2118          AND ($4::text IS NULL OR tags @> ARRAY[$4]::text[])
2119          AND ($5::bigint IS NULL OR id < $5)
2120        ORDER BY id DESC
2121        LIMIT $6
2122        "#,
2123    )
2124    .bind(filter.state)
2125    .bind(&filter.kind)
2126    .bind(&filter.queue)
2127    .bind(&filter.tag)
2128    .bind(filter.before_id)
2129    .bind(limit)
2130    .fetch_all(pool)
2131    .await?;
2132
2133    Ok(rows)
2134}
2135
2136/// Get a single job by ID.
2137pub async fn get_job(pool: &PgPool, job_id: i64) -> Result<JobRow, AwaError> {
2138    if let Some(store) = active_queue_storage(pool).await? {
2139        let row = store.load_job(pool, job_id).await?;
2140        return row.ok_or(AwaError::JobNotFound { id: job_id });
2141    }
2142
2143    let row = sqlx::query_as::<_, JobRow>("SELECT * FROM awa.jobs WHERE id = $1")
2144        .bind(job_id)
2145        .fetch_optional(pool)
2146        .await?;
2147
2148    row.ok_or(AwaError::JobNotFound { id: job_id })
2149}
2150
2151/// Fetch a job plus optional DLQ metadata for admin surfaces that need to
2152/// distinguish a DLQ row from a normal live or terminal row.
2153pub async fn get_job_with_source(
2154    pool: &PgPool,
2155    job_id: i64,
2156) -> Result<(JobRow, Option<DlqMetadata>), AwaError> {
2157    let job = get_job(pool, job_id).await?;
2158    let dlq = if active_queue_storage(pool).await?.is_some() {
2159        crate::dlq::get_dlq_job(pool, job_id)
2160            .await?
2161            .map(|row| row.metadata())
2162    } else {
2163        None
2164    };
2165    Ok((job, dlq))
2166}
2167
2168/// Build a read-only inspection snapshot for one job.
2169pub async fn dump_job(pool: &PgPool, job_id: i64) -> Result<JobDump, AwaError> {
2170    let (job, dlq) = get_job_with_source(pool, job_id).await?;
2171    Ok(build_job_dump(job, dlq))
2172}
2173
2174/// Build a read-only inspection snapshot for one attempt.
2175///
2176/// Awa does not currently persist a standalone runs table. The current attempt
2177/// is inspected from the live job row. Historical attempts are reconstructed
2178/// from the structured `errors[]` history.
2179pub async fn dump_run(
2180    pool: &PgPool,
2181    job_id: i64,
2182    attempt: Option<i16>,
2183) -> Result<RunDump, AwaError> {
2184    let job = get_job(pool, job_id).await?;
2185    let selected_attempt = attempt.unwrap_or(job.attempt);
2186
2187    if selected_attempt < 0 {
2188        return Err(AwaError::Validation("attempt must be >= 0".to_string()));
2189    }
2190
2191    if job.attempt == 0 && selected_attempt == 0 {
2192        return Ok(RunDump {
2193            job_id: job.id,
2194            kind: job.kind.clone(),
2195            queue: job.queue.clone(),
2196            selected_attempt,
2197            current_attempt: job.attempt,
2198            current_run_lease: job.run_lease,
2199            selected_run_lease: Some(job.run_lease),
2200            source: RunDumpSource::CurrentJobRow,
2201            state: job.state,
2202            started_at: job.attempted_at,
2203            finished_at: job.finalized_at,
2204            error: None,
2205            terminal: None,
2206            progress: job.progress.clone(),
2207            metadata: Some(job.metadata.clone()),
2208            callback: callback_dump(&job),
2209            raw_error_entry: None,
2210            notes: vec!["Job has not been claimed yet; attempt 0 is the pre-run snapshot.".into()],
2211        });
2212    }
2213
2214    if selected_attempt == job.attempt {
2215        let latest_error = job
2216            .errors
2217            .as_ref()
2218            .into_iter()
2219            .flatten()
2220            .filter_map(parse_error_entry)
2221            .find(|entry| entry.attempt == Some(selected_attempt));
2222        return Ok(RunDump {
2223            job_id: job.id,
2224            kind: job.kind.clone(),
2225            queue: job.queue.clone(),
2226            selected_attempt,
2227            current_attempt: job.attempt,
2228            current_run_lease: job.run_lease,
2229            selected_run_lease: Some(job.run_lease),
2230            source: RunDumpSource::CurrentJobRow,
2231            state: job.state,
2232            started_at: job.attempted_at,
2233            finished_at: if job.state.is_terminal() {
2234                job.finalized_at
2235            } else {
2236                None
2237            },
2238            error: latest_error.as_ref().and_then(|entry| entry.error.clone()),
2239            terminal: latest_error.as_ref().map(|entry| entry.terminal),
2240            progress: job.progress.clone(),
2241            metadata: Some(job.metadata.clone()),
2242            callback: callback_dump(&job),
2243            raw_error_entry: latest_error.map(|entry| entry.raw),
2244            notes: vec![],
2245        });
2246    }
2247
2248    if selected_attempt == 0 {
2249        return Err(AwaError::Validation(format!(
2250            "attempt {selected_attempt} is not available for job {job_id}; current attempt is {}",
2251            job.attempt
2252        )));
2253    }
2254
2255    if job.attempt > 0 && selected_attempt > job.attempt {
2256        return Err(AwaError::Validation(format!(
2257            "attempt {selected_attempt} is not available for job {job_id}; current attempt is {}",
2258            job.attempt
2259        )));
2260    }
2261
2262    let historical = job
2263        .errors
2264        .as_ref()
2265        .into_iter()
2266        .flatten()
2267        .filter_map(parse_error_entry)
2268        .find(|entry| entry.attempt == Some(selected_attempt))
2269        .ok_or_else(|| {
2270            AwaError::Validation(format!(
2271                "attempt {selected_attempt} is not present in the recorded error history for job {job_id}"
2272            ))
2273        })?;
2274
2275    Ok(RunDump {
2276        job_id: job.id,
2277        kind: job.kind.clone(),
2278        queue: job.queue.clone(),
2279        selected_attempt,
2280        current_attempt: job.attempt,
2281        current_run_lease: job.run_lease,
2282        selected_run_lease: None,
2283        source: RunDumpSource::ErrorHistory,
2284        state: if historical.terminal {
2285            JobState::Failed
2286        } else {
2287            JobState::Retryable
2288        },
2289        started_at: None,
2290        finished_at: historical.at,
2291        error: historical.error,
2292        terminal: Some(historical.terminal),
2293        progress: None,
2294        metadata: None,
2295        callback: None,
2296        raw_error_entry: Some(historical.raw),
2297        notes: vec![
2298            "Historical attempts are reconstructed from errors[] because Awa does not persist a standalone runs table.".into(),
2299            "Only the current attempt has live progress, callback, and metadata fields.".into(),
2300        ],
2301    })
2302}
2303
2304/// Count jobs grouped by state.
2305///
2306/// Reads from the `queue_state_counts` cache table.
2307pub async fn state_counts(pool: &PgPool) -> Result<HashMap<JobState, i64>, AwaError> {
2308    if let Some(store) = active_queue_storage(pool).await? {
2309        let sql = format!(
2310            r#"
2311            SELECT
2312                COALESCE((SELECT count(*)::bigint FROM {schema}.deferred_jobs WHERE state = 'scheduled'), 0) AS scheduled,
2313                COALESCE((
2314                    SELECT count(*)::bigint
2315                    FROM {schema}.ready_entries AS ready
2316                    JOIN {schema}.queue_claim_heads AS claims
2317                      ON claims.queue = ready.queue
2318                     AND claims.priority = ready.priority
2319                     AND claims.enqueue_shard = ready.enqueue_shard
2320                    WHERE ready.lane_seq >= {schema}.sequence_next_value(claims.seq_name)
2321                      AND NOT EXISTS (
2322                          SELECT 1 FROM {schema}.ready_tombstones AS tomb
2323                          WHERE tomb.queue = ready.queue
2324                            AND tomb.priority = ready.priority
2325                            AND tomb.enqueue_shard = ready.enqueue_shard
2326                            AND tomb.lane_seq = ready.lane_seq
2327                            AND tomb.ready_slot = ready.ready_slot
2328                            AND tomb.ready_generation = ready.ready_generation
2329                      )
2330                ), 0) AS available,
2331                COALESCE((SELECT count(*)::bigint FROM {schema}.leases WHERE state = 'running'), 0) AS running,
2332                COALESCE((SELECT count(*)::bigint FROM {schema}.done_entries WHERE state = 'completed'), 0) AS completed,
2333                COALESCE((SELECT count(*)::bigint FROM {schema}.deferred_jobs WHERE state = 'retryable'), 0) AS retryable,
2334                COALESCE((SELECT count(*)::bigint FROM {schema}.done_entries WHERE state = 'failed'), 0)
2335                  + COALESCE((SELECT count(*)::bigint FROM {schema}.dlq_entries), 0) AS failed,
2336                COALESCE((SELECT count(*)::bigint FROM {schema}.done_entries WHERE state = 'cancelled'), 0) AS cancelled,
2337                COALESCE((SELECT count(*)::bigint FROM {schema}.leases WHERE state = 'waiting_external'), 0) AS waiting_external
2338            "#,
2339            schema = store.schema()
2340        );
2341
2342        let (
2343            scheduled,
2344            available,
2345            running,
2346            completed,
2347            retryable,
2348            failed,
2349            cancelled,
2350            waiting_external,
2351        ): (i64, i64, i64, i64, i64, i64, i64, i64) = sqlx::query_as(&sql).fetch_one(pool).await?;
2352
2353        return Ok(HashMap::from([
2354            (JobState::Scheduled, scheduled),
2355            (JobState::Available, available),
2356            (JobState::Running, running),
2357            (JobState::Completed, completed),
2358            (JobState::Retryable, retryable),
2359            (JobState::Failed, failed),
2360            (JobState::Cancelled, cancelled),
2361            (JobState::WaitingExternal, waiting_external),
2362        ]));
2363    }
2364
2365    // Single scan of queue_state_counts — sums all columns in one pass
2366    // then unpivots via VALUES join.
2367    let rows = sqlx::query_as::<_, (JobState, i64)>(
2368        r#"
2369        SELECT v.state, v.total FROM (
2370            SELECT
2371                COALESCE(sum(scheduled), 0)::bigint      AS scheduled,
2372                COALESCE(sum(available), 0)::bigint      AS available,
2373                COALESCE(sum(running), 0)::bigint        AS running,
2374                COALESCE(sum(completed), 0)::bigint      AS completed,
2375                COALESCE(sum(retryable), 0)::bigint      AS retryable,
2376                COALESCE(sum(failed), 0)::bigint         AS failed,
2377                COALESCE(sum(cancelled), 0)::bigint      AS cancelled,
2378                COALESCE(sum(waiting_external), 0)::bigint AS waiting_external
2379            FROM awa.queue_state_counts
2380        ) s,
2381        LATERAL (VALUES
2382            ('scheduled'::awa.job_state,        s.scheduled),
2383            ('available'::awa.job_state,        s.available),
2384            ('running'::awa.job_state,          s.running),
2385            ('completed'::awa.job_state,        s.completed),
2386            ('retryable'::awa.job_state,        s.retryable),
2387            ('failed'::awa.job_state,           s.failed),
2388            ('cancelled'::awa.job_state,        s.cancelled),
2389            ('waiting_external'::awa.job_state, s.waiting_external)
2390        ) AS v(state, total)
2391        "#,
2392    )
2393    .fetch_all(pool)
2394    .await?;
2395
2396    Ok(rows.into_iter().collect())
2397}
2398
2399/// Get job-kind overviews for all known kinds.
2400pub async fn job_kind_overviews(pool: &PgPool) -> Result<Vec<JobKindOverview>, AwaError> {
2401    if let Some(store) = active_queue_storage(pool).await? {
2402        let sql = format!(
2403            r#"
2404            {current_jobs_cte},
2405            all_kinds AS (
2406                SELECT kind FROM current_jobs
2407                UNION
2408                SELECT kind FROM awa.job_kind_descriptors
2409                UNION
2410                SELECT DISTINCT descriptor.key AS kind
2411                FROM awa.runtime_instances runtime
2412                CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2413            ),
2414            live_kind_descriptor_variants AS (
2415                SELECT
2416                    descriptor.key AS kind,
2417                    count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
2418                FROM awa.runtime_instances runtime
2419                CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2420                WHERE runtime.last_seen_at + make_interval(
2421                    secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
2422                ) >= now()
2423                GROUP BY descriptor.key
2424            ),
2425            kind_counts AS (
2426                SELECT
2427                    kind,
2428                    count(*)::bigint AS job_count,
2429                    count(DISTINCT queue)::bigint AS queue_count
2430                FROM current_jobs
2431                GROUP BY kind
2432            ),
2433            completed_recent AS (
2434                SELECT
2435                    kind,
2436                    count(*)::bigint AS completed_last_hour
2437                FROM current_jobs
2438                WHERE state = 'completed'
2439                  AND finalized_at > clock_timestamp() - interval '1 hour'
2440                GROUP BY kind
2441            )
2442            SELECT
2443                k.kind,
2444                kd.display_name,
2445                kd.description,
2446                kd.owner,
2447                kd.docs_url,
2448                COALESCE(kd.tags, ARRAY[]::text[]) AS tags,
2449                COALESCE(kd.extra, '{{}}'::jsonb) AS extra,
2450                kd.last_seen_at AS descriptor_last_seen_at,
2451                CASE
2452                    WHEN kd.last_seen_at IS NULL THEN FALSE
2453                    ELSE kd.last_seen_at + make_interval(
2454                        secs => GREATEST(((COALESCE(kd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
2455                    ) < now()
2456                END AS descriptor_stale,
2457                COALESCE(kdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
2458                COALESCE(kc.job_count, 0) AS job_count,
2459                COALESCE(kc.queue_count, 0) AS queue_count,
2460                COALESCE(cr.completed_last_hour, 0) AS completed_last_hour
2461            FROM all_kinds k
2462            LEFT JOIN kind_counts kc ON kc.kind = k.kind
2463            LEFT JOIN awa.job_kind_descriptors kd ON kd.kind = k.kind
2464            LEFT JOIN live_kind_descriptor_variants kdv ON kdv.kind = k.kind
2465            LEFT JOIN completed_recent cr ON cr.kind = k.kind
2466            ORDER BY k.kind
2467            "#,
2468            current_jobs_cte = queue_storage_current_jobs_cte(store.schema())
2469        );
2470
2471        let rows = sqlx::query_as::<_, JobKindOverview>(&sql)
2472            .fetch_all(pool)
2473            .await?;
2474        return Ok(rows);
2475    }
2476
2477    let rows = sqlx::query_as::<_, JobKindOverview>(
2478        r#"
2479        WITH all_kinds AS (
2480            -- Union every source of kind-name knowledge so /kinds never
2481            -- hides a kind that exists *somewhere* in the system:
2482            --   - job_kind_catalog: has observed jobs
2483            --   - job_kind_descriptors: declared by worker code
2484            --   - runtime_instances: reported by a currently-running worker
2485            SELECT kind FROM awa.job_kind_catalog WHERE ref_count > 0
2486            UNION
2487            SELECT kind FROM awa.job_kind_descriptors
2488            UNION
2489            SELECT DISTINCT descriptor.key AS kind
2490            FROM awa.runtime_instances runtime
2491            CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2492        ),
2493        live_kind_descriptor_variants AS (
2494            SELECT
2495                descriptor.key AS kind,
2496                count(DISTINCT descriptor.value)::bigint AS descriptor_variant_count
2497            FROM awa.runtime_instances runtime
2498            CROSS JOIN LATERAL jsonb_each_text(runtime.job_kind_descriptor_hashes) AS descriptor(key, value)
2499            WHERE runtime.last_seen_at + make_interval(
2500                secs => GREATEST(((GREATEST(runtime.snapshot_interval_ms, 1000) / 1000) * 3)::int, 30)
2501            ) >= now()
2502            GROUP BY descriptor.key
2503        ),
2504        completed_recent AS (
2505            SELECT
2506                kind,
2507                count(*)::bigint AS completed_last_hour
2508            FROM awa.jobs_hot
2509            WHERE state = 'completed'
2510              AND finalized_at > now() - interval '1 hour'
2511            GROUP BY kind
2512        ),
2513        queue_counts AS (
2514            -- Restrict the per-kind queue fan-out to `jobs_hot` rather than
2515            -- the full `awa.jobs` view. jobs_hot is bounded by retention
2516            -- (default 24h completed / 72h failed) so the scan cost is
2517            -- tied to in-flight volume, not historical volume. The
2518            -- semantic this produces — "queues this kind is currently
2519            -- active on" — is the one admin surfaces actually care about;
2520            -- a kind that hasn't enqueued in 3 months shouldn't be
2521            -- counted as still spanning N queues.
2522            SELECT
2523                kind,
2524                count(DISTINCT queue)::bigint AS queue_count
2525            FROM awa.jobs_hot
2526            GROUP BY kind
2527        )
2528        SELECT
2529            k.kind,
2530            kd.display_name,
2531            kd.description,
2532            kd.owner,
2533            kd.docs_url,
2534            COALESCE(kd.tags, ARRAY[]::text[]) AS tags,
2535            COALESCE(kd.extra, '{}'::jsonb) AS extra,
2536            kd.last_seen_at AS descriptor_last_seen_at,
2537            CASE
2538                WHEN kd.last_seen_at IS NULL THEN FALSE
2539                ELSE kd.last_seen_at + make_interval(
2540                    secs => GREATEST(((COALESCE(kd.sync_interval_ms, 10000) / 1000) * 3)::int, 30)
2541                ) < now()
2542            END AS descriptor_stale,
2543            COALESCE(kdv.descriptor_variant_count, 0) > 1 AS descriptor_mismatch,
2544            COALESCE(kc.ref_count, 0) AS job_count,
2545            COALESCE(qc.queue_count, 0) AS queue_count,
2546            COALESCE(cr.completed_last_hour, 0) AS completed_last_hour
2547        FROM all_kinds k
2548        LEFT JOIN awa.job_kind_catalog kc ON kc.kind = k.kind
2549        LEFT JOIN awa.job_kind_descriptors kd ON kd.kind = k.kind
2550        LEFT JOIN live_kind_descriptor_variants kdv ON kdv.kind = k.kind
2551        LEFT JOIN queue_counts qc ON qc.kind = k.kind
2552        LEFT JOIN completed_recent cr ON cr.kind = k.kind
2553        ORDER BY k.kind
2554        "#,
2555    )
2556    .fetch_all(pool)
2557    .await?;
2558
2559    Ok(rows)
2560}
2561
2562pub async fn job_kind_overview(
2563    pool: &PgPool,
2564    kind: &str,
2565) -> Result<Option<JobKindOverview>, AwaError> {
2566    let rows = job_kind_overviews(pool).await?;
2567    Ok(rows.into_iter().find(|row| row.kind == kind))
2568}
2569
2570/// Return all distinct job kinds.
2571///
2572/// Reads from the `job_kind_catalog` cache table.
2573pub async fn distinct_kinds(pool: &PgPool) -> Result<Vec<String>, AwaError> {
2574    if let Some(store) = active_queue_storage(pool).await? {
2575        let sql = format!(
2576            "{} \
2577             SELECT DISTINCT kind \
2578             FROM current_jobs \
2579             ORDER BY kind",
2580            queue_storage_current_jobs_cte(store.schema())
2581        );
2582        return sqlx::query_scalar(&sql)
2583            .fetch_all(pool)
2584            .await
2585            .map_err(AwaError::from);
2586    }
2587
2588    let rows = sqlx::query_scalar::<_, String>(
2589        "SELECT kind FROM awa.job_kind_catalog WHERE ref_count > 0 ORDER BY kind",
2590    )
2591    .fetch_all(pool)
2592    .await?;
2593
2594    Ok(rows)
2595}
2596
2597/// Return all distinct queue names.
2598///
2599/// Reads from the `job_queue_catalog` cache table.
2600pub async fn distinct_queues(pool: &PgPool) -> Result<Vec<String>, AwaError> {
2601    if let Some(store) = active_queue_storage(pool).await? {
2602        let sql = format!(
2603            "{} \
2604             SELECT DISTINCT queue \
2605             FROM current_jobs \
2606             ORDER BY queue",
2607            queue_storage_current_jobs_cte(store.schema())
2608        );
2609        return sqlx::query_scalar(&sql)
2610            .fetch_all(pool)
2611            .await
2612            .map_err(AwaError::from);
2613    }
2614
2615    let rows = sqlx::query_scalar::<_, String>(
2616        "SELECT queue FROM awa.job_queue_catalog WHERE ref_count > 0 ORDER BY queue",
2617    )
2618    .fetch_all(pool)
2619    .await?;
2620
2621    Ok(rows)
2622}
2623
2624/// Drain one batch of dirty keys and recompute exact cached rows.
2625/// Returns the number of dirty keys processed in this batch.
2626///
2627/// Called frequently by the maintenance leader (~2s). Uses per-queue
2628/// indexes for targeted recompute rather than full table scans.
2629pub async fn recompute_dirty_admin_metadata(pool: &PgPool) -> Result<i32, AwaError> {
2630    if active_queue_storage(pool).await?.is_some() {
2631        return Ok(0);
2632    }
2633    let count: i32 = sqlx::query_scalar("SELECT awa.recompute_dirty_admin_metadata(100)")
2634        .fetch_one(pool)
2635        .await?;
2636    Ok(count)
2637}
2638
2639/// Drain ALL dirty keys until the backlog is empty.
2640///
2641/// Use in tests or admin tooling where you need the cache to be fully
2642/// consistent before reading. Each call to the underlying SQL function
2643/// acquires a blocking advisory lock, so concurrent callers serialize
2644/// rather than skip.
2645pub async fn flush_dirty_admin_metadata(pool: &PgPool) -> Result<i32, AwaError> {
2646    if active_queue_storage(pool).await?.is_some() {
2647        return Ok(0);
2648    }
2649    let mut total = 0i32;
2650    loop {
2651        let count: i32 = sqlx::query_scalar("SELECT awa.recompute_dirty_admin_metadata(100)")
2652            .fetch_one(pool)
2653            .await?;
2654        total += count;
2655        if count == 0 {
2656            break;
2657        }
2658    }
2659    Ok(total)
2660}
2661
2662/// Full reconciliation of admin metadata counters from base tables.
2663///
2664/// Called infrequently by the maintenance leader (~60s) as a safety net
2665/// to correct any drift from skipped dirty keys. Also called during
2666/// migrate() to warm the cache.
2667pub async fn refresh_admin_metadata(pool: &PgPool) -> Result<(), AwaError> {
2668    if active_queue_storage(pool).await?.is_some() {
2669        return Ok(());
2670    }
2671    sqlx::query("SELECT awa.refresh_admin_metadata()")
2672        .execute(pool)
2673        .await?;
2674    Ok(())
2675}
2676
2677/// Retry multiple jobs by ID. Only retries failed, cancelled, or waiting_external jobs.
2678pub async fn bulk_retry(pool: &PgPool, ids: &[i64]) -> Result<Vec<JobRow>, AwaError> {
2679    if let Some(store) = active_queue_storage(pool).await? {
2680        return store.retry_jobs_by_ids(pool, ids).await;
2681    }
2682
2683    let rows = sqlx::query_as::<_, JobRow>(
2684        r#"
2685        UPDATE awa.jobs
2686        SET state = 'available', attempt = 0, run_at = now(),
2687            finalized_at = NULL, heartbeat_at = NULL, deadline_at = NULL,
2688            callback_id = NULL, callback_timeout_at = NULL,
2689            callback_filter = NULL, callback_on_complete = NULL,
2690            callback_on_fail = NULL, callback_transform = NULL
2691        WHERE id = ANY($1) AND state IN ('failed', 'cancelled', 'waiting_external')
2692        RETURNING *
2693        "#,
2694    )
2695    .bind(ids)
2696    .fetch_all(pool)
2697    .await?;
2698
2699    Ok(rows)
2700}
2701
2702/// Cancel multiple jobs by ID. Only cancels non-terminal jobs.
2703pub async fn bulk_cancel(pool: &PgPool, ids: &[i64]) -> Result<Vec<JobRow>, AwaError> {
2704    if let Some(store) = active_queue_storage(pool).await? {
2705        return store.cancel_jobs_by_ids(pool, ids).await;
2706    }
2707
2708    let rows = sqlx::query_as::<_, JobRow>(
2709        r#"
2710        UPDATE awa.jobs
2711        SET state = 'cancelled', finalized_at = now(),
2712            callback_id = NULL, callback_timeout_at = NULL,
2713            callback_filter = NULL, callback_on_complete = NULL,
2714            callback_on_fail = NULL, callback_transform = NULL
2715        WHERE id = ANY($1) AND state NOT IN ('completed', 'failed', 'cancelled')
2716        RETURNING *
2717        "#,
2718    )
2719    .bind(ids)
2720    .fetch_all(pool)
2721    .await?;
2722
2723    Ok(rows)
2724}
2725
2726/// A bucketed count of jobs by state over time.
2727#[derive(Debug, Clone, Serialize)]
2728pub struct StateTimeseriesBucket {
2729    pub bucket: chrono::DateTime<chrono::Utc>,
2730    pub state: JobState,
2731    pub count: i64,
2732}
2733
2734/// Return time-bucketed state counts over the last N minutes.
2735pub async fn state_timeseries(
2736    pool: &PgPool,
2737    minutes: i32,
2738) -> Result<Vec<StateTimeseriesBucket>, AwaError> {
2739    if let Some(store) = active_queue_storage(pool).await? {
2740        let sql = format!(
2741            "{} \
2742             SELECT \
2743                 date_trunc('minute', created_at) AS bucket, \
2744                 state, \
2745                 count(*) AS count \
2746             FROM current_jobs \
2747             WHERE created_at >= clock_timestamp() - make_interval(mins => $1) \
2748             GROUP BY bucket, state \
2749             ORDER BY bucket",
2750            queue_storage_current_jobs_cte(store.schema())
2751        );
2752        let rows = sqlx::query_as::<_, (chrono::DateTime<chrono::Utc>, JobState, i64)>(&sql)
2753            .bind(minutes)
2754            .fetch_all(pool)
2755            .await?;
2756
2757        return Ok(rows
2758            .into_iter()
2759            .map(|(bucket, state, count)| StateTimeseriesBucket {
2760                bucket,
2761                state,
2762                count,
2763            })
2764            .collect());
2765    }
2766
2767    let rows = sqlx::query_as::<_, (chrono::DateTime<chrono::Utc>, JobState, i64)>(
2768        r#"
2769        SELECT
2770            date_trunc('minute', created_at) AS bucket,
2771            state,
2772            count(*) AS count
2773        FROM awa.jobs
2774        WHERE created_at >= now() - make_interval(mins => $1)
2775        GROUP BY bucket, state
2776        ORDER BY bucket
2777        "#,
2778    )
2779    .bind(minutes)
2780    .fetch_all(pool)
2781    .await?;
2782
2783    Ok(rows
2784        .into_iter()
2785        .map(|(bucket, state, count)| StateTimeseriesBucket {
2786            bucket,
2787            state,
2788            count,
2789        })
2790        .collect())
2791}
2792
2793/// Register a callback for a running job, writing the callback_id and timeout
2794/// to the database immediately.
2795///
2796/// Call this BEFORE sending the callback_id to the external system to avoid
2797/// the race condition where the external system fires before the DB knows
2798/// about the callback.
2799///
2800/// Returns the generated callback UUID on success.
2801pub async fn register_callback(
2802    pool: &PgPool,
2803    job_id: i64,
2804    run_lease: i64,
2805    timeout: std::time::Duration,
2806) -> Result<Uuid, AwaError> {
2807    if let Some(store) = active_queue_storage(pool).await? {
2808        return store
2809            .register_callback(pool, job_id, run_lease, timeout)
2810            .await;
2811    }
2812
2813    let callback_id = Uuid::new_v4();
2814    let timeout_secs = timeout.as_secs_f64();
2815    let result = sqlx::query(
2816        r#"UPDATE awa.jobs
2817           SET callback_id = $2,
2818               callback_timeout_at = now() + make_interval(secs => $3),
2819               callback_filter = NULL,
2820               callback_on_complete = NULL,
2821               callback_on_fail = NULL,
2822               callback_transform = NULL
2823           WHERE id = $1 AND state = 'running' AND run_lease = $4"#,
2824    )
2825    .bind(job_id)
2826    .bind(callback_id)
2827    .bind(timeout_secs)
2828    .bind(run_lease)
2829    .execute(pool)
2830    .await?;
2831    if result.rows_affected() == 0 {
2832        return Err(AwaError::Validation("job is not in running state".into()));
2833    }
2834    Ok(callback_id)
2835}
2836
2837/// Complete a waiting job via external callback.
2838///
2839/// Accepts jobs in `waiting_external` or `running` state (race handling: the
2840/// external system may fire before the executor transitions to `waiting_external`).
2841///
2842/// When `resume` is `false` (default), the job transitions to `completed`.
2843/// When `resume` is `true`, the job transitions back to `running` with the
2844/// callback payload stored in metadata under `_awa_callback_result`. The
2845/// handler can then read the result and continue processing (sequential
2846/// callback pattern from ADR-021).
2847pub async fn complete_external(
2848    pool: &PgPool,
2849    callback_id: Uuid,
2850    payload: Option<serde_json::Value>,
2851    run_lease: Option<i64>,
2852) -> Result<JobRow, AwaError> {
2853    let mut tx = pool.begin().await?;
2854    let row = complete_external_in_tx(&mut tx, callback_id, payload, run_lease).await?;
2855    tx.commit().await?;
2856    Ok(row)
2857}
2858
2859/// Transaction-aware variant of [`complete_external`] (ADR-029). Callers
2860/// that want to dispatch follow-up specs in the same transaction as the
2861/// callback completion should drive this directly and own the
2862/// `tx.commit()`. See [`crate::admin::resolve_callback_in_tx`] for the
2863/// policy-evaluating counterpart.
2864pub async fn complete_external_in_tx(
2865    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2866    callback_id: Uuid,
2867    payload: Option<serde_json::Value>,
2868    run_lease: Option<i64>,
2869) -> Result<JobRow, AwaError> {
2870    complete_external_in_tx_inner(tx, callback_id, payload, run_lease, false).await
2871}
2872
2873/// Complete a waiting job and resume the handler with the callback payload.
2874///
2875/// Like `complete_external`, but the job transitions to `running` instead of
2876/// `completed`, allowing the handler to continue with sequential callbacks.
2877/// The payload is stored in `metadata._awa_callback_result`.
2878pub async fn resume_external(
2879    pool: &PgPool,
2880    callback_id: Uuid,
2881    payload: Option<serde_json::Value>,
2882    run_lease: Option<i64>,
2883) -> Result<JobRow, AwaError> {
2884    let mut tx = pool.begin().await?;
2885    let row = resume_external_in_tx(&mut tx, callback_id, payload, run_lease).await?;
2886    tx.commit().await?;
2887    Ok(row)
2888}
2889
2890/// Transaction-aware variant of [`resume_external`] (ADR-029). Resume is
2891/// not currently an ADR-029 outcome — the handler keeps running and will
2892/// emit its own completion event later — so callers that just need
2893/// payload delivery generally use [`resume_external`]. This `_in_tx`
2894/// shape exists for symmetry with the other callback transitions and so
2895/// applications can compose the resume `UPDATE` with their own
2896/// transaction-local writes.
2897pub async fn resume_external_in_tx(
2898    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2899    callback_id: Uuid,
2900    payload: Option<serde_json::Value>,
2901    run_lease: Option<i64>,
2902) -> Result<JobRow, AwaError> {
2903    complete_external_in_tx_inner(tx, callback_id, payload, run_lease, true).await
2904}
2905
2906async fn complete_external_in_tx_inner(
2907    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2908    callback_id: Uuid,
2909    payload: Option<serde_json::Value>,
2910    run_lease: Option<i64>,
2911    resume: bool,
2912) -> Result<JobRow, AwaError> {
2913    if let Some(store) = active_queue_storage_in_tx(tx).await? {
2914        return store
2915            .complete_external_in_tx(tx, callback_id, payload, run_lease, resume)
2916            .await;
2917    }
2918
2919    let row = if resume {
2920        // Resume: transition to running, store payload, refresh heartbeat.
2921        // The handler is still alive and polling — it will detect the state change.
2922        let payload_json = payload.unwrap_or(serde_json::Value::Null);
2923        sqlx::query_as::<_, JobRow>(
2924            r#"
2925            UPDATE awa.jobs
2926            SET state = 'running',
2927                callback_id = NULL,
2928                callback_timeout_at = NULL,
2929                callback_filter = NULL,
2930                callback_on_complete = NULL,
2931                callback_on_fail = NULL,
2932                callback_transform = NULL,
2933                heartbeat_at = now(),
2934                metadata = metadata || jsonb_build_object('_awa_callback_result', $3::jsonb)
2935            WHERE callback_id = $1 AND state IN ('waiting_external', 'running')
2936              AND ($2::bigint IS NULL OR run_lease = $2)
2937            RETURNING *
2938            "#,
2939        )
2940        .bind(callback_id)
2941        .bind(run_lease)
2942        .bind(&payload_json)
2943        .fetch_optional(tx.as_mut())
2944        .await?
2945    } else {
2946        // Complete: terminal state, clear everything.
2947        sqlx::query_as::<_, JobRow>(
2948            r#"
2949            UPDATE awa.jobs
2950            SET state = 'completed',
2951                finalized_at = now(),
2952                callback_id = NULL,
2953                callback_timeout_at = NULL,
2954                callback_filter = NULL,
2955                callback_on_complete = NULL,
2956                callback_on_fail = NULL,
2957                callback_transform = NULL,
2958                heartbeat_at = NULL,
2959                deadline_at = NULL,
2960                progress = NULL
2961            WHERE callback_id = $1 AND state IN ('waiting_external', 'running')
2962              AND ($2::bigint IS NULL OR run_lease = $2)
2963            RETURNING *
2964            "#,
2965        )
2966        .bind(callback_id)
2967        .bind(run_lease)
2968        .fetch_optional(tx.as_mut())
2969        .await?
2970    };
2971
2972    row.ok_or(AwaError::CallbackNotFound {
2973        callback_id: callback_id.to_string(),
2974    })
2975}
2976
2977/// Fail a waiting job via external callback.
2978///
2979/// Records the error and transitions to `failed`.
2980pub async fn fail_external(
2981    pool: &PgPool,
2982    callback_id: Uuid,
2983    error: &str,
2984    run_lease: Option<i64>,
2985) -> Result<JobRow, AwaError> {
2986    let mut tx = pool.begin().await?;
2987    let row = fail_external_in_tx(&mut tx, callback_id, error, run_lease).await?;
2988    tx.commit().await?;
2989    Ok(row)
2990}
2991
2992/// Transaction-aware variant of [`fail_external`] (ADR-029).
2993pub async fn fail_external_in_tx(
2994    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
2995    callback_id: Uuid,
2996    error: &str,
2997    run_lease: Option<i64>,
2998) -> Result<JobRow, AwaError> {
2999    if let Some(store) = active_queue_storage_in_tx(tx).await? {
3000        return store
3001            .fail_external_with_error_entry_in_tx(
3002                tx,
3003                callback_id,
3004                serde_json::json!({ "error": error }),
3005                run_lease,
3006            )
3007            .await;
3008    }
3009
3010    let row = sqlx::query_as::<_, JobRow>(
3011        r#"
3012        UPDATE awa.jobs
3013        SET state = 'failed',
3014            finalized_at = now(),
3015            callback_id = NULL,
3016            callback_timeout_at = NULL,
3017            callback_filter = NULL,
3018            callback_on_complete = NULL,
3019            callback_on_fail = NULL,
3020            callback_transform = NULL,
3021            heartbeat_at = NULL,
3022            deadline_at = NULL,
3023            errors = errors || jsonb_build_object(
3024                'error', $2::text,
3025                'attempt', attempt,
3026                'at', now()
3027            )::jsonb
3028        WHERE callback_id = $1 AND state IN ('waiting_external', 'running')
3029          AND ($3::bigint IS NULL OR run_lease = $3)
3030        RETURNING *
3031        "#,
3032    )
3033    .bind(callback_id)
3034    .bind(error)
3035    .bind(run_lease)
3036    .fetch_optional(tx.as_mut())
3037    .await?;
3038
3039    row.ok_or(AwaError::CallbackNotFound {
3040        callback_id: callback_id.to_string(),
3041    })
3042}
3043
3044/// Retry a waiting job via external callback.
3045///
3046/// Resets to `available` with attempt = 0. The handler must be idempotent
3047/// with respect to the external call — a retry re-executes from scratch.
3048///
3049/// Only accepts `waiting_external` state — unlike complete/fail which are
3050/// terminal transitions, retry puts the job back to `available`. Allowing
3051/// retry from `running` would risk concurrent dispatch if the original
3052/// handler hasn't finished yet.
3053pub async fn retry_external(
3054    pool: &PgPool,
3055    callback_id: Uuid,
3056    run_lease: Option<i64>,
3057) -> Result<JobRow, AwaError> {
3058    let mut tx = pool.begin().await?;
3059    let row = retry_external_in_tx(&mut tx, callback_id, run_lease).await?;
3060    tx.commit().await?;
3061    Ok(row)
3062}
3063
3064/// Transaction-aware variant of [`retry_external`] (ADR-029).
3065pub async fn retry_external_in_tx(
3066    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3067    callback_id: Uuid,
3068    run_lease: Option<i64>,
3069) -> Result<JobRow, AwaError> {
3070    if let Some(store) = active_queue_storage_in_tx(tx).await? {
3071        return store.retry_external_in_tx(tx, callback_id, run_lease).await;
3072    }
3073
3074    let row = sqlx::query_as::<_, JobRow>(
3075        r#"
3076        UPDATE awa.jobs
3077        SET state = 'available',
3078            attempt = 0,
3079            run_at = now(),
3080            finalized_at = NULL,
3081            callback_id = NULL,
3082            callback_timeout_at = NULL,
3083            callback_filter = NULL,
3084            callback_on_complete = NULL,
3085            callback_on_fail = NULL,
3086            callback_transform = NULL,
3087            heartbeat_at = NULL,
3088            deadline_at = NULL
3089        WHERE callback_id = $1 AND state = 'waiting_external'
3090          AND ($2::bigint IS NULL OR run_lease = $2)
3091        RETURNING *
3092        "#,
3093    )
3094    .bind(callback_id)
3095    .bind(run_lease)
3096    .fetch_optional(tx.as_mut())
3097    .await?;
3098
3099    row.ok_or(AwaError::CallbackNotFound {
3100        callback_id: callback_id.to_string(),
3101    })
3102}
3103
3104/// Reset the callback timeout for a long-running external operation.
3105///
3106/// External systems call this periodically to signal "still working" without
3107/// completing the job. Resets `callback_timeout_at` to `now() + timeout`.
3108/// The job stays in `waiting_external`.
3109///
3110/// Returns the updated job row, or `CallbackNotFound` if the callback ID
3111/// doesn't match a waiting job.
3112pub async fn heartbeat_callback(
3113    pool: &PgPool,
3114    callback_id: Uuid,
3115    timeout: std::time::Duration,
3116) -> Result<JobRow, AwaError> {
3117    if let Some(store) = active_queue_storage(pool).await? {
3118        return store.heartbeat_callback(pool, callback_id, timeout).await;
3119    }
3120
3121    let timeout_secs = timeout.as_secs_f64();
3122    let row = sqlx::query_as::<_, JobRow>(
3123        r#"
3124        UPDATE awa.jobs
3125        SET callback_timeout_at = now() + make_interval(secs => $2)
3126        WHERE callback_id = $1 AND state = 'waiting_external'
3127        RETURNING *
3128        "#,
3129    )
3130    .bind(callback_id)
3131    .bind(timeout_secs)
3132    .fetch_optional(pool)
3133    .await?;
3134
3135    row.ok_or(AwaError::CallbackNotFound {
3136        callback_id: callback_id.to_string(),
3137    })
3138}
3139
3140/// Cancel (clear) a registered callback for a running job.
3141///
3142/// Best-effort cleanup: returns `Ok(true)` if a row was updated,
3143/// `Ok(false)` if no match (already resolved, rescued, or wrong lease).
3144/// Callers should not treat `false` as an error.
3145pub async fn cancel_callback(pool: &PgPool, job_id: i64, run_lease: i64) -> Result<bool, AwaError> {
3146    if let Some(store) = active_queue_storage(pool).await? {
3147        return store.cancel_callback(pool, job_id, run_lease).await;
3148    }
3149
3150    let result = sqlx::query(
3151        r#"
3152        UPDATE awa.jobs
3153        SET callback_id = NULL,
3154            callback_timeout_at = NULL,
3155            callback_filter = NULL,
3156            callback_on_complete = NULL,
3157            callback_on_fail = NULL,
3158            callback_transform = NULL
3159        WHERE id = $1 AND callback_id IS NOT NULL AND state = 'running' AND run_lease = $2
3160        "#,
3161    )
3162    .bind(job_id)
3163    .bind(run_lease)
3164    .execute(pool)
3165    .await?;
3166
3167    Ok(result.rows_affected() > 0)
3168}
3169
3170// ── Sequential callback wait helpers ─────────────────────────────────
3171//
3172// These functions extract the DB-interaction logic for `wait_for_callback`
3173// so that both the Rust `JobContext` and the Python bridge call the same
3174// code paths.
3175
3176/// Result of a single poll iteration inside `wait_for_callback`.
3177#[derive(Debug)]
3178pub enum CallbackPollResult {
3179    /// The callback was resolved and the payload is ready.
3180    Resolved(serde_json::Value),
3181    /// Still waiting — caller should sleep and poll again.
3182    Pending,
3183    /// The callback token is stale (a different callback is current).
3184    Stale {
3185        token: Uuid,
3186        current: Uuid,
3187        state: JobState,
3188    },
3189    /// The job left the wait unexpectedly (rescued, cancelled, etc.).
3190    UnexpectedState { token: Uuid, state: JobState },
3191    /// The job was not found.
3192    NotFound,
3193}
3194
3195/// Transition a running job to `waiting_external` for the given callback.
3196///
3197/// Returns `Ok(true)` if the transition succeeded, `Ok(false)` if the row
3198/// did not match (the caller should check for an early-resume race).
3199pub async fn enter_callback_wait(
3200    pool: &PgPool,
3201    job_id: i64,
3202    run_lease: i64,
3203    callback_id: Uuid,
3204) -> Result<bool, AwaError> {
3205    if let Some(store) = active_queue_storage(pool).await? {
3206        return store
3207            .enter_callback_wait(pool, job_id, run_lease, callback_id)
3208            .await;
3209    }
3210
3211    let result = sqlx::query(
3212        r#"
3213        UPDATE awa.jobs
3214        SET state = 'waiting_external',
3215            heartbeat_at = NULL,
3216            deadline_at = NULL
3217        WHERE id = $1 AND state = 'running' AND run_lease = $2 AND callback_id = $3
3218        "#,
3219    )
3220    .bind(job_id)
3221    .bind(run_lease)
3222    .bind(callback_id)
3223    .execute(pool)
3224    .await?;
3225
3226    Ok(result.rows_affected() > 0)
3227}
3228
3229/// Check the current state of a job during callback wait.
3230///
3231/// Handles the early-resume race: if `resume_external` won the race before
3232/// the handler transitioned to `waiting_external`, the callback result is
3233/// already in metadata and this returns `Resolved`.
3234pub async fn check_callback_state(
3235    pool: &PgPool,
3236    job_id: i64,
3237    callback_id: Uuid,
3238) -> Result<CallbackPollResult, AwaError> {
3239    if let Some(store) = active_queue_storage(pool).await? {
3240        return store.check_callback_state(pool, job_id, callback_id).await;
3241    }
3242
3243    let row: Option<(JobState, Option<Uuid>, serde_json::Value)> =
3244        sqlx::query_as("SELECT state, callback_id, metadata FROM awa.jobs WHERE id = $1")
3245            .bind(job_id)
3246            .fetch_optional(pool)
3247            .await?;
3248
3249    match row {
3250        Some((JobState::Running, None, metadata))
3251            if metadata.get("_awa_callback_result").is_some() =>
3252        {
3253            let payload = take_callback_payload(pool, job_id, metadata).await?;
3254            Ok(CallbackPollResult::Resolved(payload))
3255        }
3256        Some((state, Some(current_callback_id), _)) if current_callback_id != callback_id => {
3257            Ok(CallbackPollResult::Stale {
3258                token: callback_id,
3259                current: current_callback_id,
3260                state,
3261            })
3262        }
3263        Some((JobState::WaitingExternal, Some(current), _)) if current == callback_id => {
3264            Ok(CallbackPollResult::Pending)
3265        }
3266        Some((state, _, _)) => Ok(CallbackPollResult::UnexpectedState {
3267            token: callback_id,
3268            state,
3269        }),
3270        None => Ok(CallbackPollResult::NotFound),
3271    }
3272}
3273
3274/// Extract the `_awa_callback_result` key from metadata and clean it up.
3275pub async fn take_callback_payload(
3276    pool: &PgPool,
3277    job_id: i64,
3278    metadata: serde_json::Value,
3279) -> Result<serde_json::Value, AwaError> {
3280    let payload = metadata
3281        .get("_awa_callback_result")
3282        .cloned()
3283        .unwrap_or(serde_json::Value::Null);
3284
3285    sqlx::query("UPDATE awa.jobs SET metadata = metadata - '_awa_callback_result' WHERE id = $1")
3286        .bind(job_id)
3287        .execute(pool)
3288        .await?;
3289
3290    Ok(payload)
3291}
3292
3293// ── CEL callback expressions ──────────────────────────────────────────
3294
3295/// Configuration for CEL callback expressions.
3296///
3297/// All fields are optional. When all are `None`, behaviour is identical to
3298/// the original `register_callback` (no expression evaluation).
3299#[derive(Debug, Clone, Default)]
3300pub struct CallbackConfig {
3301    /// Gate: should this payload be processed at all? Returns bool.
3302    pub filter: Option<String>,
3303    /// Does this payload indicate success? Returns bool.
3304    pub on_complete: Option<String>,
3305    /// Does this payload indicate failure? Returns bool. Evaluated before on_complete.
3306    pub on_fail: Option<String>,
3307    /// Reshape payload before returning to caller. Returns any Value.
3308    pub transform: Option<String>,
3309}
3310
3311impl CallbackConfig {
3312    /// Returns true if no expressions are configured.
3313    pub fn is_empty(&self) -> bool {
3314        self.filter.is_none()
3315            && self.on_complete.is_none()
3316            && self.on_fail.is_none()
3317            && self.transform.is_none()
3318    }
3319}
3320
3321/// What `resolve_callback` should do if no CEL conditions match or no
3322/// expressions are configured.
3323#[derive(Debug, Clone, Copy, PartialEq, Eq)]
3324pub enum DefaultAction {
3325    Complete,
3326    Fail,
3327    Ignore,
3328}
3329
3330/// Outcome of `resolve_callback`.
3331#[derive(Debug)]
3332pub enum ResolveOutcome {
3333    Completed {
3334        payload: Option<serde_json::Value>,
3335        job: JobRow,
3336    },
3337    Failed {
3338        job: JobRow,
3339    },
3340    Ignored {
3341        reason: String,
3342    },
3343}
3344
3345impl ResolveOutcome {
3346    pub fn is_completed(&self) -> bool {
3347        matches!(self, ResolveOutcome::Completed { .. })
3348    }
3349    pub fn is_failed(&self) -> bool {
3350        matches!(self, ResolveOutcome::Failed { .. })
3351    }
3352    pub fn is_ignored(&self) -> bool {
3353        matches!(self, ResolveOutcome::Ignored { .. })
3354    }
3355}
3356
3357/// Register a callback with optional CEL expressions.
3358///
3359/// When expressions are provided and the `cel` feature is enabled, each
3360/// expression is trial-compiled at registration time so syntax errors are
3361/// caught early.
3362///
3363/// When the `cel` feature is disabled and any expression is non-None,
3364/// returns `AwaError::Validation`.
3365pub async fn register_callback_with_config(
3366    pool: &PgPool,
3367    job_id: i64,
3368    run_lease: i64,
3369    timeout: std::time::Duration,
3370    config: &CallbackConfig,
3371) -> Result<Uuid, AwaError> {
3372    // Validate CEL expressions at registration time: compile + check references
3373    #[cfg(feature = "cel")]
3374    {
3375        for (name, expr) in [
3376            ("filter", &config.filter),
3377            ("on_complete", &config.on_complete),
3378            ("on_fail", &config.on_fail),
3379            ("transform", &config.transform),
3380        ] {
3381            if let Some(src) = expr {
3382                let program = cel::Program::compile(src).map_err(|e| {
3383                    AwaError::Validation(format!("invalid CEL expression for {name}: {e}"))
3384                })?;
3385
3386                // Reject undeclared variables — CEL only reports these at execution
3387                // time, so an expression like `missing == 1` would parse fine but
3388                // silently fall into the fail-open path at resolve time.
3389                let refs = program.references();
3390                let bad_vars: Vec<&str> = refs
3391                    .variables()
3392                    .into_iter()
3393                    .filter(|v| *v != "payload")
3394                    .collect();
3395                if !bad_vars.is_empty() {
3396                    return Err(AwaError::Validation(format!(
3397                        "CEL expression for {name} references undeclared variable(s): {}; \
3398                         only 'payload' is available",
3399                        bad_vars.join(", ")
3400                    )));
3401                }
3402            }
3403        }
3404    }
3405
3406    #[cfg(not(feature = "cel"))]
3407    {
3408        if !config.is_empty() {
3409            return Err(AwaError::Validation(
3410                "CEL expressions require the 'cel' feature".into(),
3411            ));
3412        }
3413    }
3414
3415    if let Some(store) = active_queue_storage(pool).await? {
3416        return store
3417            .register_callback_with_config(pool, job_id, run_lease, timeout, config)
3418            .await;
3419    }
3420
3421    let callback_id = Uuid::new_v4();
3422    let timeout_secs = timeout.as_secs_f64();
3423
3424    let result = sqlx::query(
3425        r#"UPDATE awa.jobs
3426           SET callback_id = $2,
3427               callback_timeout_at = now() + make_interval(secs => $3),
3428               callback_filter = $4,
3429               callback_on_complete = $5,
3430               callback_on_fail = $6,
3431               callback_transform = $7
3432           WHERE id = $1 AND state = 'running' AND run_lease = $8"#,
3433    )
3434    .bind(job_id)
3435    .bind(callback_id)
3436    .bind(timeout_secs)
3437    .bind(&config.filter)
3438    .bind(&config.on_complete)
3439    .bind(&config.on_fail)
3440    .bind(&config.transform)
3441    .bind(run_lease)
3442    .execute(pool)
3443    .await?;
3444
3445    if result.rows_affected() == 0 {
3446        return Err(AwaError::Validation("job is not in running state".into()));
3447    }
3448    Ok(callback_id)
3449}
3450
3451/// Internal action decided by CEL evaluation or default.
3452enum ResolveAction {
3453    Complete(Option<serde_json::Value>),
3454    Fail {
3455        error: String,
3456        expression: Option<String>,
3457    },
3458    Ignore(String),
3459}
3460
3461/// Resolve a callback by evaluating CEL expressions against the payload.
3462///
3463/// Uses a transaction with `SELECT ... FOR UPDATE` for atomicity.
3464/// The `default_action` determines behaviour when no CEL conditions match
3465/// or no expressions are configured.
3466pub async fn resolve_callback(
3467    pool: &PgPool,
3468    callback_id: Uuid,
3469    payload: Option<serde_json::Value>,
3470    default_action: DefaultAction,
3471    run_lease: Option<i64>,
3472) -> Result<ResolveOutcome, AwaError> {
3473    let mut tx = pool.begin().await?;
3474    let outcome =
3475        resolve_callback_in_tx(&mut tx, callback_id, payload, default_action, run_lease).await?;
3476    tx.commit().await?;
3477    Ok(outcome)
3478}
3479
3480/// Transaction-aware variant of [`resolve_callback`] (ADR-029).
3481///
3482/// Performs callback lookup, CEL evaluation, and the chosen state
3483/// transition (Complete / Fail / Ignore) inside the caller-owned
3484/// transaction. The lookup takes a row lock — `SELECT ... FOR UPDATE` on
3485/// `awa.jobs_hot` for canonical storage; `FOR UPDATE OF lease` against
3486/// the active queue-storage `leases` table — so concurrent resolves are
3487/// serialised on the same `callback_id`.
3488///
3489/// Callers that want to dispatch ADR-029 follow-up specs in the same
3490/// transaction (e.g. `Client::resolve_callback`) drive this directly,
3491/// then dispatch specs against the returned [`ResolveOutcome`] before
3492/// committing.
3493pub async fn resolve_callback_in_tx(
3494    tx: &mut sqlx::Transaction<'_, sqlx::Postgres>,
3495    callback_id: Uuid,
3496    payload: Option<serde_json::Value>,
3497    default_action: DefaultAction,
3498    run_lease: Option<i64>,
3499) -> Result<ResolveOutcome, AwaError> {
3500    if let Some(store) = active_queue_storage_in_tx(tx).await? {
3501        let job = store
3502            .callback_job_in_tx(tx, callback_id, run_lease, true)
3503            .await?
3504            .ok_or(AwaError::CallbackNotFound {
3505                callback_id: callback_id.to_string(),
3506            })?;
3507
3508        let action = evaluate_or_default(&job, &payload, default_action)?;
3509
3510        return match action {
3511            ResolveAction::Complete(transformed_payload) => {
3512                let completed_job = store
3513                    .complete_external_in_tx(tx, callback_id, None, run_lease, false)
3514                    .await?;
3515                Ok(ResolveOutcome::Completed {
3516                    payload: transformed_payload,
3517                    job: completed_job,
3518                })
3519            }
3520            ResolveAction::Fail { error, expression } => {
3521                let mut error_json = serde_json::json!({
3522                    "error": error,
3523                    "attempt": job.attempt,
3524                    "at": chrono::Utc::now().to_rfc3339(),
3525                });
3526                if let Some(expr) = expression {
3527                    error_json["expression"] = serde_json::Value::String(expr);
3528                }
3529
3530                let failed_job = store
3531                    .fail_external_with_error_entry_in_tx(tx, callback_id, error_json, run_lease)
3532                    .await?;
3533                Ok(ResolveOutcome::Failed { job: failed_job })
3534            }
3535            ResolveAction::Ignore(reason) => Ok(ResolveOutcome::Ignored { reason }),
3536        };
3537    }
3538
3539    // Query jobs_hot directly (not the awa.jobs UNION ALL view) because
3540    // FOR UPDATE is not reliably supported on UNION views. Waiting_external
3541    // and running jobs are always in jobs_hot (the check constraint on
3542    // scheduled_jobs only allows scheduled/retryable).
3543    //
3544    // Accepts both 'waiting_external' and 'running' to handle the race where
3545    // a fast callback arrives before the executor transitions running ->
3546    // waiting_external (matching complete_external/fail_external behavior).
3547    let job = sqlx::query_as::<_, JobRow>(
3548        "SELECT * FROM awa.jobs_hot WHERE callback_id = $1
3549         AND state IN ('waiting_external', 'running')
3550         AND ($2::bigint IS NULL OR run_lease = $2)
3551         FOR UPDATE",
3552    )
3553    .bind(callback_id)
3554    .bind(run_lease)
3555    .fetch_optional(tx.as_mut())
3556    .await?
3557    .ok_or(AwaError::CallbackNotFound {
3558        callback_id: callback_id.to_string(),
3559    })?;
3560
3561    let action = evaluate_or_default(&job, &payload, default_action)?;
3562
3563    match action {
3564        ResolveAction::Complete(transformed_payload) => {
3565            let completed_job = sqlx::query_as::<_, JobRow>(
3566                r#"
3567                UPDATE awa.jobs
3568                SET state = 'completed',
3569                    finalized_at = now(),
3570                    callback_id = NULL,
3571                    callback_timeout_at = NULL,
3572                    callback_filter = NULL,
3573                    callback_on_complete = NULL,
3574                    callback_on_fail = NULL,
3575                    callback_transform = NULL,
3576                    heartbeat_at = NULL,
3577                    deadline_at = NULL,
3578                    progress = NULL
3579                WHERE id = $1
3580                RETURNING *
3581                "#,
3582            )
3583            .bind(job.id)
3584            .fetch_one(tx.as_mut())
3585            .await?;
3586
3587            Ok(ResolveOutcome::Completed {
3588                payload: transformed_payload,
3589                job: completed_job,
3590            })
3591        }
3592        ResolveAction::Fail { error, expression } => {
3593            let mut error_json = serde_json::json!({
3594                "error": error,
3595                "attempt": job.attempt,
3596                "at": chrono::Utc::now().to_rfc3339(),
3597            });
3598            if let Some(expr) = expression {
3599                error_json["expression"] = serde_json::Value::String(expr);
3600            }
3601
3602            let failed_job = sqlx::query_as::<_, JobRow>(
3603                r#"
3604                UPDATE awa.jobs
3605                SET state = 'failed',
3606                    finalized_at = now(),
3607                    callback_id = NULL,
3608                    callback_timeout_at = NULL,
3609                    callback_filter = NULL,
3610                    callback_on_complete = NULL,
3611                    callback_on_fail = NULL,
3612                    callback_transform = NULL,
3613                    heartbeat_at = NULL,
3614                    deadline_at = NULL,
3615                    errors = errors || $2::jsonb
3616                WHERE id = $1
3617                RETURNING *
3618                "#,
3619            )
3620            .bind(job.id)
3621            .bind(error_json)
3622            .fetch_one(tx.as_mut())
3623            .await?;
3624
3625            Ok(ResolveOutcome::Failed { job: failed_job })
3626        }
3627        ResolveAction::Ignore(reason) => {
3628            // No state change — caller's tx will commit cleanly without
3629            // changing the FOR UPDATE row, releasing the lock.
3630            Ok(ResolveOutcome::Ignored { reason })
3631        }
3632    }
3633}
3634
3635/// Evaluate CEL expressions or fall through to default_action.
3636fn evaluate_or_default(
3637    job: &JobRow,
3638    payload: &Option<serde_json::Value>,
3639    default_action: DefaultAction,
3640) -> Result<ResolveAction, AwaError> {
3641    let has_expressions = job.callback_filter.is_some()
3642        || job.callback_on_complete.is_some()
3643        || job.callback_on_fail.is_some()
3644        || job.callback_transform.is_some();
3645
3646    if !has_expressions {
3647        return Ok(apply_default(default_action, payload));
3648    }
3649
3650    #[cfg(feature = "cel")]
3651    {
3652        Ok(evaluate_cel(job, payload, default_action))
3653    }
3654
3655    #[cfg(not(feature = "cel"))]
3656    {
3657        // Expressions are present but CEL feature is not enabled.
3658        // Return an error without mutating the job — it stays in waiting_external.
3659        let _ = (payload, default_action);
3660        Err(AwaError::Validation(
3661            "CEL expressions present but 'cel' feature is not enabled".into(),
3662        ))
3663    }
3664}
3665
3666fn apply_default(
3667    default_action: DefaultAction,
3668    payload: &Option<serde_json::Value>,
3669) -> ResolveAction {
3670    match default_action {
3671        DefaultAction::Complete => ResolveAction::Complete(payload.clone()),
3672        DefaultAction::Fail => ResolveAction::Fail {
3673            error: "callback failed: default action".to_string(),
3674            expression: None,
3675        },
3676        DefaultAction::Ignore => {
3677            ResolveAction::Ignore("no expressions configured, default is ignore".to_string())
3678        }
3679    }
3680}
3681
3682#[cfg(feature = "cel")]
3683fn evaluate_cel(
3684    job: &JobRow,
3685    payload: &Option<serde_json::Value>,
3686    default_action: DefaultAction,
3687) -> ResolveAction {
3688    let payload_value = payload.as_ref().cloned().unwrap_or(serde_json::Value::Null);
3689
3690    // 1. Evaluate filter
3691    if let Some(filter_expr) = &job.callback_filter {
3692        match eval_bool(filter_expr, &payload_value, job.id, "filter") {
3693            Ok(true) => {} // pass through
3694            Ok(false) => {
3695                return ResolveAction::Ignore("filter expression returned false".to_string());
3696            }
3697            Err(_) => {
3698                // Fail-open: treat filter error as true (pass through)
3699            }
3700        }
3701    }
3702
3703    // 2. Evaluate on_fail (before on_complete — fail takes precedence)
3704    if let Some(on_fail_expr) = &job.callback_on_fail {
3705        match eval_bool(on_fail_expr, &payload_value, job.id, "on_fail") {
3706            Ok(true) => {
3707                return ResolveAction::Fail {
3708                    error: "callback failed: on_fail expression matched".to_string(),
3709                    expression: Some(on_fail_expr.clone()),
3710                };
3711            }
3712            Ok(false) => {} // don't fail
3713            Err(_) => {
3714                // Fail-open: treat on_fail error as false (don't fail)
3715            }
3716        }
3717    }
3718
3719    // 3. Evaluate on_complete
3720    if let Some(on_complete_expr) = &job.callback_on_complete {
3721        match eval_bool(on_complete_expr, &payload_value, job.id, "on_complete") {
3722            Ok(true) => {
3723                // Complete with optional transform
3724                let transformed = apply_transform(job, &payload_value);
3725                return ResolveAction::Complete(Some(transformed));
3726            }
3727            Ok(false) => {} // don't complete
3728            Err(_) => {
3729                // Fail-open: treat on_complete error as false (don't complete)
3730            }
3731        }
3732    }
3733
3734    // 4. Neither condition matched → apply default_action
3735    apply_default(default_action, payload)
3736}
3737
3738#[cfg(feature = "cel")]
3739fn eval_bool(
3740    expression: &str,
3741    payload_value: &serde_json::Value,
3742    job_id: i64,
3743    expression_name: &str,
3744) -> Result<bool, ()> {
3745    let program = match cel::Program::compile(expression) {
3746        Ok(p) => p,
3747        Err(e) => {
3748            tracing::warn!(
3749                job_id,
3750                expression_name,
3751                expression,
3752                error = %e,
3753                "CEL compilation error during evaluation"
3754            );
3755            return Err(());
3756        }
3757    };
3758
3759    let mut context = cel::Context::default();
3760    if let Err(e) = context.add_variable("payload", payload_value.clone()) {
3761        tracing::warn!(
3762            job_id,
3763            expression_name,
3764            error = %e,
3765            "Failed to add payload variable to CEL context"
3766        );
3767        return Err(());
3768    }
3769
3770    match program.execute(&context) {
3771        Ok(cel::Value::Bool(b)) => Ok(b),
3772        Ok(other) => {
3773            tracing::warn!(
3774                job_id,
3775                expression_name,
3776                expression,
3777                result_type = ?other.type_of(),
3778                "CEL expression returned non-bool"
3779            );
3780            Err(())
3781        }
3782        Err(e) => {
3783            tracing::warn!(
3784                job_id,
3785                expression_name,
3786                expression,
3787                error = %e,
3788                "CEL execution error"
3789            );
3790            Err(())
3791        }
3792    }
3793}
3794
3795#[cfg(feature = "cel")]
3796fn apply_transform(job: &JobRow, payload_value: &serde_json::Value) -> serde_json::Value {
3797    let transform_expr = match &job.callback_transform {
3798        Some(expr) => expr,
3799        None => return payload_value.clone(),
3800    };
3801
3802    let program = match cel::Program::compile(transform_expr) {
3803        Ok(p) => p,
3804        Err(e) => {
3805            tracing::warn!(
3806                job_id = job.id,
3807                expression = transform_expr,
3808                error = %e,
3809                "CEL transform compilation error, using original payload"
3810            );
3811            return payload_value.clone();
3812        }
3813    };
3814
3815    let mut context = cel::Context::default();
3816    if let Err(e) = context.add_variable("payload", payload_value.clone()) {
3817        tracing::warn!(
3818            job_id = job.id,
3819            error = %e,
3820            "Failed to add payload variable for transform"
3821        );
3822        return payload_value.clone();
3823    }
3824
3825    match program.execute(&context) {
3826        Ok(value) => match value.json() {
3827            Ok(json) => json,
3828            Err(e) => {
3829                tracing::warn!(
3830                    job_id = job.id,
3831                    expression = transform_expr,
3832                    error = %e,
3833                    "CEL transform result could not be converted to JSON, using original payload"
3834                );
3835                payload_value.clone()
3836            }
3837        },
3838        Err(e) => {
3839            tracing::warn!(
3840                job_id = job.id,
3841                expression = transform_expr,
3842                error = %e,
3843                "CEL transform execution error, using original payload"
3844            );
3845            payload_value.clone()
3846        }
3847    }
3848}