Skip to main content

faucet_cli/serve/
runner.rs

1//! The run lifecycle: validate + queue a submission (`submit`), then run it under
2//! a permit. A cancel / timeout / shutdown trigger cooperatively cancels the
3//! pipeline (so a buffered sink flushes at its next page boundary, #146 H16) and
4//! grants a bounded flush grace before hard-dropping it; the task then finalizes
5//! an authoritative terminal status. See spec §7 + §20.
6
7use crate::auth_catalog::build_auth_catalog;
8use crate::executor::{ExecuteOptions, RunSummary, run_expanded};
9use crate::registry::build_source;
10use crate::serve::error::ServeError;
11use crate::serve::history::{Claim, InvocationRecord, RunRecord, RunStatus};
12use crate::serve::history::{ClaimedShard, ShardInsert};
13use crate::serve::load::{ConfigFormat, LoadedSubmission, load_submission};
14use crate::serve::rbac::AuthContext;
15use crate::serve::state::ServerState;
16use crate::serve::{idempotency, metrics};
17use chrono::{DateTime, FixedOffset, Utc};
18use serde::{Deserialize, Serialize};
19use std::collections::BTreeMap;
20use std::time::Duration;
21use tokio_util::sync::CancellationToken;
22use tracing::Instrument;
23
24/// `Retry-After` advertised when the queue is full.
25const QUEUE_FULL_RETRY_AFTER_SECS: u64 = 5;
26
27/// Grace granted to a cancelled / timed-out / shutting-down run to flush
28/// buffered sink output cooperatively before its future is hard-dropped (which
29/// aborts the pipeline's task set, the backstop for a run stuck mid-write so a
30/// hung run can't wedge shutdown). Generous enough for an S3 multipart
31/// completion (#146 H16).
32pub(crate) const RUN_FLUSH_GRACE: Duration = Duration::from_secs(30);
33
34/// `POST /v1/runs` request body.
35#[derive(Debug, Deserialize)]
36pub struct SubmitRequest {
37    pub config: String,
38    #[serde(default)]
39    pub config_format: ConfigFormatWire,
40    pub name: Option<String>,
41    #[serde(default)]
42    pub labels: BTreeMap<String, String>,
43    pub timeout_secs: Option<u64>,
44    #[serde(default)]
45    pub doctor_first: bool,
46    pub idempotency_key: Option<String>,
47    pub clock: Option<String>,
48    /// Optional completion callback fired when this run reaches a terminal
49    /// state (#481). Validated at submit time so a bad destination is a 422 on
50    /// the submission rather than a silent no-op when the run finishes.
51    #[serde(default)]
52    pub callback: Option<crate::serve::callback::CallbackSpec>,
53}
54
55/// Wire enum mirroring `load::ConfigFormat` with serde rename.
56#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize)]
57#[serde(rename_all = "lowercase")]
58pub enum ConfigFormatWire {
59    #[default]
60    Yaml,
61    Json,
62}
63
64impl From<ConfigFormatWire> for ConfigFormat {
65    fn from(w: ConfigFormatWire) -> Self {
66        match w {
67            ConfigFormatWire::Yaml => ConfigFormat::Yaml,
68            ConfigFormatWire::Json => ConfigFormat::Json,
69        }
70    }
71}
72
73/// `POST /v1/runs` success body (202).
74#[derive(Debug, Serialize)]
75pub struct SubmitResponse {
76    pub run_id: String,
77    pub status: RunStatus,
78    pub submitted_at: DateTime<Utc>,
79}
80
81/// Re-run a claimed Pending run on this instance (cluster mode). Reconstructs the
82/// execution inputs from the persisted record (re-resolving config with this
83/// instance's own env/credentials), acquires a permit, and runs the shared tail.
84pub fn resume_claimed_run(state: ServerState, rec: RunRecord) {
85    tokio::spawn(async move {
86        let run_id = rec.run_id.clone();
87        let Some(body) = rec.config_body.as_deref() else {
88            tracing::error!(run_id, "claimed run has no stored config; failing it");
89            finalize(
90                &state,
91                &run_id,
92                rec.submitted_at,
93                Terminal::Failed {
94                    reason: "claimed run record missing config_body".into(),
95                    records: 0,
96                    invs: Vec::new(),
97                },
98            )
99            .await;
100            return;
101        };
102        let format = rec.config_format.unwrap_or_default();
103        let loaded = match load_submission(body, format, state.default_base().as_ref()).await {
104            Ok(l) => l,
105            Err(e) => {
106                finalize(
107                    &state,
108                    &run_id,
109                    rec.submitted_at,
110                    Terminal::Failed {
111                        reason: format!(
112                            "re-loading claimed config: {}",
113                            e.api_error().error.message
114                        ),
115                        records: 0,
116                        invs: Vec::new(),
117                    },
118                )
119                .await;
120                return;
121            }
122        };
123
124        // Mode B (#230): a sharded run is expanded into shard rows here — the
125        // claiming instance acts as the (ephemeral) coordinator — and is NOT
126        // executed as a whole. Enumeration + insert is idempotent, so two
127        // instances both claiming + coordinating converge on the same shard set.
128        if let Some(sh) = loaded.cfg.shard.clone()
129            && sh.count >= 2
130        {
131            match coordinate_sharded_run(&state, &run_id, &loaded, sh.count).await {
132                Ok(true) => return, // expanded into shards — shard loop runs them
133                Ok(false) => {}     // not shardable → fall through, run the whole run
134                Err(e) => {
135                    finalize(
136                        &state,
137                        &run_id,
138                        rec.submitted_at,
139                        Terminal::Failed {
140                            reason: format!("sharding: {e}"),
141                            records: 0,
142                            invs: Vec::new(),
143                        },
144                    )
145                    .await;
146                    return;
147                }
148            }
149        }
150
151        // The claim loop only claims up to available_permits and is the sole
152        // permit consumer, so this acquire returns immediately.
153        let _permit = state
154            .semaphore()
155            .acquire_owned()
156            .await
157            .expect("semaphore not closed");
158        // Register a local cancel token so a cross-instance cancel (the claim loop
159        // calling registry.cancel) reaches this run.
160        let run_token = CancellationToken::new();
161        state.registry().register(run_id.clone(), run_token.clone());
162        execute_run(
163            state.clone(),
164            loaded,
165            run_id,
166            run_token,
167            rec.submitted_at,
168            rec.timeout_secs,
169            rec.clock.clone(),
170            false,
171        )
172        .await;
173    });
174}
175
176/// Coordinator step (Mode B): expand a sharded run into `faucet_serve_shards`
177/// rows. Returns `Ok(true)` when the run was sharded (caller must not execute it
178/// as a whole), `Ok(false)` when it isn't shardable (caller runs it whole).
179///
180/// Idempotent: enumeration is deterministic and the insert is
181/// `ON CONFLICT DO NOTHING`, so a re-coordinated run (e.g. after the coordinator
182/// crashed and the Pending run was requeued) converges on the same shard set.
183async fn coordinate_sharded_run(
184    state: &ServerState,
185    run_id: &str,
186    loaded: &LoadedSubmission,
187    count: usize,
188) -> crate::error::CliResult<bool> {
189    use crate::error::CliError;
190
191    // Sharding applies to a single-source pipeline; a matrix fan-out is not
192    // shardable (each row is already an independent unit — use Mode A).
193    if loaded.nodes.len() != 1 {
194        tracing::warn!(
195            run_id,
196            nodes = loaded.nodes.len(),
197            "shard requested but the run is not a single-node pipeline; running it whole"
198        );
199        return Ok(false);
200    }
201    let node = &loaded.nodes[0];
202    let auth = build_auth_catalog(loaded.cfg.auth.as_ref())
203        .map_err(|e| CliError::Internal(format!("auth catalog: {e}")))?;
204    let source = build_source(&node.source.kind, node.source.config.clone(), &auth, None).await?;
205    if !source.is_shardable() {
206        tracing::warn!(
207            run_id,
208            kind = %node.source.kind,
209            "source is not shardable; running the run whole"
210        );
211        return Ok(false);
212    }
213
214    // Mark the parent run Sharded BEFORE inserting any shard rows (F27).
215    // Orphan recovery (`recover_orphans` / `reclaim_orphans`) only reclaims
216    // `running` runs; once the parent is `Sharded` it is immune to a false
217    // fail. Doing this first guarantees the invariant that **no shard row ever
218    // exists while the parent is still `running`** — otherwise a coordinator
219    // crash between `insert_shards` and the status flip would let orphan
220    // recovery mark the run Failed while its already-inserted shards are
221    // claimed, execute, and write data (a run the user believes failed and may
222    // resubmit → duplicate writes). If the flip itself fails, the run stays
223    // `running` + owned by us with no shards inserted, so a crash here is
224    // safely *requeued* by reclaim, not falsely failed. (A crash after the flip
225    // but before `insert_shards` leaves a 0-shard `Sharded` run, which
226    // `all_terminal()` never finalizes — an availability leak, not a
227    // correctness/duplication bug; the desired trade.)
228    {
229        let mut r = state
230            .history()
231            .get(run_id)
232            .await
233            .map_err(|e| CliError::Internal(e.to_string()))?
234            .ok_or_else(|| CliError::Internal(format!("run {run_id} vanished before sharding")))?;
235        r.status = RunStatus::Sharded;
236        state
237            .history()
238            .upsert(&r)
239            .await
240            .map_err(|e| CliError::Internal(e.to_string()))?;
241    }
242
243    let shards = source
244        .enumerate_shards(count)
245        .await
246        .map_err(|e| CliError::Internal(format!("enumerate_shards: {e}")))?;
247    let inserts: Vec<ShardInsert> = shards
248        .iter()
249        .map(|s| ShardInsert {
250            shard_id: s.id.clone(),
251            descriptor: s.descriptor.clone(),
252            size_estimate: s.size_estimate,
253        })
254        .collect();
255    let inserted = state
256        .history()
257        .insert_shards(run_id, &inserts)
258        .await
259        .map_err(|e| CliError::Internal(e.to_string()))?;
260    tracing::info!(
261        run_id,
262        shards = inserts.len(),
263        inserted,
264        "expanded run into shards (Mode B)"
265    );
266
267    // Wake the local claim loop so it picks up the freshly-inserted shards.
268    state.cluster().kick();
269    Ok(true)
270}
271
272/// Execute one claimed shard (Mode B): rebuild + narrow the source to the shard,
273/// run it under a permit, owner-fenced-finalize the shard, then finalize the
274/// parent run once every shard is terminal.
275pub fn resume_claimed_shard(state: ServerState, claimed: ClaimedShard) {
276    tokio::spawn(async move {
277        let ClaimedShard {
278            run_id,
279            shard_id,
280            descriptor,
281            run,
282        } = claimed;
283
284        let Some(body) = run.config_body.clone() else {
285            tracing::error!(run_id, shard_id, "claimed shard's run has no stored config");
286            let _ = state
287                .history()
288                .finalize_shard(&run_id, &shard_id, false)
289                .await;
290            maybe_finalize_parent(&state, &run_id).await;
291            return;
292        };
293        let format = run.config_format.unwrap_or_default();
294        let loaded = match load_submission(&body, format, state.default_base().as_ref()).await {
295            Ok(l) => l,
296            Err(e) => {
297                tracing::error!(
298                    run_id,
299                    shard_id,
300                    error = %e.api_error().error.message,
301                    "re-loading shard config failed"
302                );
303                let _ = state
304                    .history()
305                    .finalize_shard(&run_id, &shard_id, false)
306                    .await;
307                maybe_finalize_parent(&state, &run_id).await;
308                return;
309            }
310        };
311
312        let _permit = state
313            .semaphore()
314            .acquire_owned()
315            .await
316            .expect("semaphore not closed");
317
318        let shard = faucet_core::ShardSpec {
319            id: shard_id.clone(),
320            descriptor,
321            size_estimate: None,
322        };
323        // Register a per-shard cooperative-cancel token BEFORE running so a
324        // cross-instance cancel reaches this shard: `cancel_run` flags the parent
325        // → `pending_shard_cancellations` → the claim loop calls
326        // `registry().cancel_run_shards(run_id)`, which fires this token. Removed
327        // on return so a terminated shard never leaks a token (shard accounting is
328        // separate from the run's `in_flight`, so a plain `deregister` is used).
329        let coop = CancellationToken::new();
330        state
331            .registry()
332            .register_shard(&run_id, &shard_id, coop.clone());
333        // Count the shard in `in_flight` so the shutdown drain waits for it and
334        // fires the cooperative-flush cancel (audit #321 H5). The guard releases
335        // the slot + token on EVERY return path (including panic), so a shard can
336        // never leak the counter and wedge graceful shutdown.
337        state.registry().mark_shard_running();
338        let _shard_guard = ShardInFlightGuard {
339            state: state.clone(),
340            run_id: run_id.clone(),
341            shard_id: shard_id.clone(),
342        };
343        let success = execute_shard(
344            &state,
345            loaded,
346            &run_id,
347            &shard_id,
348            shard,
349            coop,
350            run.timeout_secs,
351            run.clock.clone(),
352            run.submitted_at,
353        )
354        .await;
355
356        match state
357            .history()
358            .finalize_shard(&run_id, &shard_id, success)
359            .await
360        {
361            Ok(true) => {}
362            Ok(false) => tracing::warn!(
363                run_id,
364                shard_id,
365                "shard was reclaimed by another instance; discarding result"
366            ),
367            Err(e) => tracing::error!(run_id, shard_id, error = %e, "finalize_shard failed"),
368        }
369        maybe_finalize_parent(&state, &run_id).await;
370    });
371}
372
373/// Run one shard's pipeline (single node, source narrowed via `opts.shard`).
374/// Returns `true` on clean completion. Does not touch the parent run record —
375/// the caller finalizes the shard and the parent.
376#[allow(clippy::too_many_arguments)]
377async fn execute_shard(
378    state: &ServerState,
379    loaded: LoadedSubmission,
380    run_id: &str,
381    shard_id: &str,
382    shard: faucet_core::ShardSpec,
383    coop: CancellationToken,
384    timeout_secs: Option<u64>,
385    clock_flag: Option<String>,
386    submitted_at: DateTime<Utc>,
387) -> bool {
388    let LoadedSubmission { cfg, nodes } = loaded;
389    let pipeline_name = cfg.name.clone().unwrap_or_else(|| "serve".to_string());
390
391    let auth = match build_auth_catalog(cfg.auth.as_ref()) {
392        Ok(a) => a,
393        Err(e) => {
394            tracing::error!(run_id, shard_id, "shard auth catalog: {e}");
395            return false;
396        }
397    };
398    let clock = match resolve_clock(clock_flag.as_deref(), submitted_at) {
399        Ok(c) => c,
400        Err(e) => {
401            tracing::error!(
402                run_id,
403                shard_id,
404                "shard clock: {}",
405                e.api_error().error.message
406            );
407            return false;
408        }
409    };
410    let resilience = match &cfg.resilience {
411        Some(spec) => match spec.to_policy() {
412            Ok(p) => Some(p),
413            Err(e) => {
414                tracing::error!(run_id, shard_id, "shard resilience: {e}");
415                return false;
416            }
417        },
418        None => None,
419    };
420    #[cfg(feature = "lineage")]
421    let lineage = match crate::lineage_glue::build_emitter(cfg.lineage.as_ref()) {
422        Ok(l) => l,
423        Err(e) => {
424            tracing::error!(run_id, shard_id, "shard lineage: {e}");
425            return false;
426        }
427    };
428
429    // `coop` is the per-shard cancel token registered by `resume_claimed_shard`;
430    // a cross-instance cancel (F10), a server shutdown, or a timeout all fire it
431    // so the shard's pipeline flushes at its next page boundary.
432    let opts = ExecuteOptions {
433        pipeline_name,
434        // Correlate notifications to the submitted run (#480). Every shard of a
435        // sharded run reports the same `run_id`; they differ by `invocation_id`.
436        run_id: Some(run_id.to_string()),
437        execution: cfg.execution.clone(),
438        dry_run: false,
439        limit: None,
440        state_path_override: None,
441        shard: Some(shard),
442        auth,
443        clock,
444        cancel: Some(coop.clone()),
445        resilience,
446        // Inert under `shard: Some(..)` (a shard's volume is not the row's) —
447        // carried for uniformity with the whole-run path below.
448        sla: cfg.sla.clone(),
449        #[cfg(feature = "lineage")]
450        lineage,
451        #[cfg(feature = "lineage")]
452        lineage_cfg: cfg.lineage.clone(),
453        #[cfg(feature = "notify")]
454        notifier: None,
455        // Inert under `shard: Some(..)` — a shard's records/URIs are a slice
456        // of the row's; the catalog records whole runs only.
457        #[cfg(feature = "catalog")]
458        catalog: None,
459    };
460
461    let server_shutdown = state.shutdown_token();
462    let span = tracing::info_span!("faucet.serve.shard", serve_run_id = %run_id, shard = %shard_id);
463    let work = async move { classify_run(run_expanded(nodes, opts).await) }.instrument(span);
464    tokio::pin!(work);
465    let timeout_fut = async {
466        match timeout_secs {
467            Some(s) => tokio::time::sleep(Duration::from_secs(s)).await,
468            None => std::future::pending::<()>().await,
469        }
470    };
471    tokio::pin!(timeout_fut);
472
473    // Cancel triggers (shutdown / timeout) cooperatively cancel + flush, like
474    // execute_run. A failed shard simply returns false → its lease eventually
475    // reassigns it (or it poisons after max_attempts).
476    // A cross-instance cancel (F10) flows in via the shard's registered coop
477    // token: `cancel_run` flags the parent → `pending_cancellations` →
478    // `claim_loop` fires `registry().cancel(run_id)`. The token is registered
479    // under the run id by `resume_claimed_shard` before this runs.
480    let terminal = tokio::select! {
481        biased;
482        t = &mut work => t,
483        _ = coop.cancelled() => {
484            // Fired by the claim loop for a remote cancel. Flush within the grace
485            // window; a cooperative cancel returns Ok(partial), so the shard is
486            // Cancelled — but a flush that FAILS must surface, not be masked.
487            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
488                Ok(failed @ Terminal::Failed { .. }) => failed,
489                Ok(_) | Err(_) => Terminal::Cancelled,
490            }
491        }
492        _ = server_shutdown.cancelled() => {
493            coop.cancel();
494            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
495                Ok(failed @ Terminal::Failed { .. }) => failed,
496                Ok(_) | Err(_) => Terminal::ShutdownFailed,
497            }
498        }
499        _ = &mut timeout_fut => {
500            coop.cancel();
501            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
502                Ok(failed @ Terminal::Failed { .. }) => failed,
503                Ok(_) | Err(_) => Terminal::Timeout { secs: timeout_secs.unwrap_or(0) },
504            }
505        }
506    };
507    matches!(terminal, Terminal::Completed { .. })
508}
509
510/// Finalize a `Sharded` parent run once all its shards are terminal. The last
511/// shard to finish always observes `all_terminal` (its own `finalize_shard`
512/// committed first), so the run never lingers `Sharded`. A benign double-finalize
513/// (two shards finishing simultaneously) writes the same terminal status twice.
514async fn maybe_finalize_parent(state: &ServerState, run_id: &str) {
515    let progress = match state.history().shard_progress(run_id).await {
516        Ok(p) => p,
517        Err(e) => {
518            tracing::warn!(run_id, error = %e, "shard_progress failed");
519            return;
520        }
521    };
522    if !progress.all_terminal() {
523        return;
524    }
525    let success = progress.failed == 0;
526    let status = if success {
527        RunStatus::Completed
528    } else {
529        RunStatus::Failed
530    };
531    let error =
532        (!success).then(|| format!("{}/{} shard(s) failed", progress.failed, progress.total));
533    // Status-fenced finalize: transitions the parent only while it is still
534    // `Sharded`, and does NOT re-stamp owner/lease on the terminal record, so a
535    // near-simultaneous double-finalize from two instances has a single winner
536    // (F45). The loser (and any non-`Sharded` parent) is a no-op.
537    match state
538        .history()
539        .finalize_sharded_parent(run_id, status, Utc::now(), error)
540        .await
541    {
542        Ok(true) => {
543            metrics::record_run_finished(status, if success { "ok" } else { "error" });
544            tracing::info!(
545                run_id,
546                shards = progress.total,
547                failed = progress.failed,
548                "sharded run finalized"
549            );
550            // Fire the parent's completion callback exactly once — `Ok(true)` is
551            // the status-fenced winner, so a near-simultaneous double-finalize
552            // from two instances delivers a single callback (#481). Re-read the
553            // record: `finalize_sharded_parent` takes fields, not the record, so
554            // this is the only way to get the persisted terminal state.
555            match state.history().get(run_id).await {
556                Ok(Some(rec)) => crate::serve::callback::fire(&rec).await,
557                Ok(None) => tracing::warn!(
558                    run_id,
559                    "sharded parent vanished before its callback could fire"
560                ),
561                Err(e) => tracing::warn!(
562                    run_id,
563                    error = %e,
564                    "could not read sharded parent for its completion callback"
565                ),
566            }
567        }
568        Ok(false) => {} // already finalized by another shard/instance — nothing to do
569        Err(e) => {
570            tracing::error!(run_id, error = %e, "finalizing sharded parent run failed");
571        }
572    }
573}
574
575/// Warn at submit when a clustered or source-sharded run is at-least-once with
576/// a non-idempotent destination (F26/F39). Cluster failover and shard reclaim
577/// are at-least-once by construction: an owner whose lease lapses while it is
578/// still alive (slow, paused), or a reassigned shard whose source re-reads from
579/// the start (postgres PK-range sharding yields no resumable bookmark), can
580/// re-execute work and write **duplicate** rows to an append-mode sink. The
581/// safe configuration is `write_mode: upsert`/`delete` (keyed, so re-writes are
582/// idempotent) or `delivery: exactly_once`. We surface the risk loudly rather
583/// than silently duplicating — the repo's #1 worst bug class.
584/// Pure decision for [`warn_if_cluster_at_least_once`]: the sink kinds that
585/// would write non-idempotently under cluster failover / shard reclaim. Empty
586/// when the run is neither clustered nor sharded, is `exactly_once`, or every
587/// sink is keyed (`upsert`/`delete`).
588fn at_least_once_risky_sinks(
589    loaded: &LoadedSubmission,
590    clustered: bool,
591    sharded: bool,
592) -> Vec<&str> {
593    if !(clustered || sharded) || loaded.cfg.delivery == faucet_core::DeliveryMode::ExactlyOnce {
594        return Vec::new();
595    }
596    loaded
597        .nodes
598        .iter()
599        .filter(|n| {
600            !matches!(
601                n.sink
602                    .config
603                    .get("write_mode")
604                    .and_then(|v| v.as_str())
605                    .unwrap_or("append"),
606                "upsert" | "delete"
607            )
608        })
609        .map(|n| n.sink.kind.as_str())
610        .collect()
611}
612
613fn warn_if_cluster_at_least_once(loaded: &LoadedSubmission, clustered: bool, sharded: bool) {
614    let risky = at_least_once_risky_sinks(loaded, clustered, sharded);
615    if !risky.is_empty() {
616        let scope = if sharded {
617            "source-sharded (Mode B)"
618        } else {
619            "clustered"
620        };
621        tracing::warn!(
622            sinks = ?risky,
623            "{scope} execution is at-least-once: a failover or shard reclaim can re-run work \
624             and write duplicate rows to an append-mode sink. Set `write_mode: upsert` (or \
625             `delivery: exactly_once`) on the destination to make re-execution idempotent (F26/F39)."
626        );
627    }
628}
629
630/// Best-effort release of an idempotency claim whose paired run-record upsert
631/// failed (F21). No-op when the request carried no key; the release itself is
632/// scoped by `run_id` (a newer run that re-claimed the key keeps its claim).
633async fn release_orphaned_claim(state: &ServerState, req: &SubmitRequest, run_id: &str) {
634    if req.idempotency_key.is_some()
635        && let Err(e) = state.history().release_idempotency(run_id).await
636    {
637        tracing::warn!(
638            run_id,
639            error = %e,
640            "failed to release idempotency claim after a run-record write error; \
641             a replay of the key may 404 until the claim self-expires"
642        );
643    }
644}
645
646/// Validate, idempotency-claim, queue, and spawn a submission.
647pub async fn submit(
648    state: ServerState,
649    req: SubmitRequest,
650    actor: AuthContext,
651) -> Result<SubmitResponse, ServeError> {
652    let format: ConfigFormat = req.config_format.into();
653    let loaded = load_submission(&req.config, format, state.default_base().as_ref()).await?;
654
655    // At-least-once duplicate-write warning for clustered / source-sharded runs
656    // with an append-mode destination (F26/F39).
657    let sharded = loaded.cfg.shard.as_ref().is_some_and(|s| s.count >= 2);
658    warn_if_cluster_at_least_once(&loaded, state.cluster().enabled(), sharded);
659
660    // Completion-callback validation (#481) — before the queue reservation, so a
661    // malformed or refused destination costs nothing and surfaces as a 422 on
662    // the submission rather than as a silent no-op when the run finishes.
663    if let Some(cb) = req.callback.as_ref() {
664        cb.validate(state.callback_allow_hosts())
665            .map_err(|message| ServeError::Unprocessable {
666                message,
667                details: None,
668            })?;
669        cb.reject_secrets_in_cluster(state.cluster().enabled())
670            .map_err(|message| ServeError::Unprocessable {
671                message,
672                details: None,
673            })?;
674    }
675
676    // Reserve a queue slot first, so a Fresh idempotency claim is always followed
677    // by a spawned run (no orphaned claims — spec §20.2).
678    if !state.registry().try_reserve() {
679        return Err(ServeError::QueueFull {
680            retry_after_secs: QUEUE_FULL_RETRY_AFTER_SECS,
681        });
682    }
683    // Releases the reservation on ANY early return below (doctor_first 422 /
684    // replay / conflict / claim or upsert error). Defused just before spawn.
685    let reservation = ReservationGuard::new(state.clone());
686
687    // doctor_first preflight — run BEHIND the reservation so concurrent preflight
688    // probing is bounded by `max_queued_runs` rather than running unthrottled
689    // before any limit applies (#146 R). On failure the guard releases the slot
690    // via the early `?`. The (redacted) report is stored on the run record below
691    // so `GET /v1/runs/{id}` exposes it (#146 R: doctor_report was never set).
692    let doctor_report = if req.doctor_first {
693        Some(run_doctor_first(&state, &loaded).await?)
694    } else {
695        None
696    };
697
698    let run_id = uuid::Uuid::now_v7().to_string();
699
700    // Config fingerprint (sha256) — the idempotency identity AND the value
701    // recorded on the `run.submit` audit entry (#205).
702    let merged = serde_json::to_value(&loaded.cfg).unwrap_or(serde_json::Value::Null);
703    let fp_config = idempotency::fingerprint(&merged, loaded.cfg.name.as_deref());
704
705    // Idempotency claim (if a key was supplied).
706    if let Some(key) = &req.idempotency_key {
707        // Fold the run-affecting request fields (clock / timeout_secs / labels)
708        // into the fingerprint, not just the config — so a key replayed with a
709        // different backfill `clock` is a 409, not a replay of the original
710        // run's window (#146 M7).
711        let fp = idempotency::request_fingerprint(
712            &fp_config,
713            req.clock.as_deref(),
714            req.timeout_secs,
715            &req.labels,
716        );
717        match state
718            .history()
719            .claim_idempotency(key, &fp, &run_id, state.idempotency_retention())
720            .await
721            .map_err(|e| match e {
722                // Degraded backend can't safely honor idempotency → 503, retry.
723                crate::serve::history::HistoryError::Degraded(m) => ServeError::Unavailable(m),
724                other => ServeError::Internal(other.to_string()),
725            })? {
726            Claim::Fresh => {}
727            Claim::Replay(existing) => {
728                metrics::record_idempotency_hit();
729                return replay_response(&state, &existing).await;
730            }
731            Claim::Conflict => {
732                return Err(ServeError::Conflict(
733                    "idempotency key reused with a different payload".into(),
734                ));
735            }
736        }
737        // A `Fresh` claim is recorded BEFORE the record upsert below. The memory
738        // backend's `upsert` is infallible, but a SQL backend's can fail — so if
739        // the upsert fails, `release_orphaned_claim` drops the claim (F21) rather
740        // than leaving a replay 404-ing until the claim self-expires.
741    }
742
743    let submitted_at = Utc::now();
744    let mut rec = RunRecord::queued(
745        run_id.clone(),
746        req.name.clone(),
747        req.labels.clone(),
748        req.idempotency_key.clone(),
749        submitted_at,
750    );
751    rec.doctor_report = doctor_report;
752    rec.callback = req.callback.clone();
753
754    if state.cluster().enabled() {
755        // A degraded (DB-unreachable) backend can't coordinate a cluster: the
756        // claim loop's claim_pending is a no-op on the in-memory fallback, so a
757        // Pending run would never be claimed. Fail closed with a retryable 503
758        // rather than silently orphaning the run (#197 spec §9).
759        if state.history().degraded() {
760            return Err(ServeError::Unavailable(
761                "clustered run-history backend is degraded; runs cannot be claimed \
762                 by any instance — retry once it recovers"
763                    .into(),
764            ));
765        }
766        // Cluster mode: persist the RAW config so any instance can re-resolve +
767        // run it, mark the run Pending, and wake the local claim loop. No local
768        // queue slot / spawn — the claim loop owns execution.
769        rec.status = RunStatus::Pending;
770        rec.config_body = Some(req.config.clone());
771        rec.config_format = Some(req.config_format.into());
772        rec.timeout_secs = req.timeout_secs;
773        rec.clock = req.clock.clone();
774        if let Err(e) = state.history().upsert(&rec).await {
775            // The record write that should follow a `Fresh` claim failed — release
776            // the orphaned claim so a replay starts fresh instead of 404-ing for
777            // the whole retention window (F21). Best-effort.
778            release_orphaned_claim(&state, &req, &run_id).await;
779            return Err(ServeError::Internal(e.to_string()));
780        }
781        // Release the local queue reservation (cluster runs are bounded by the
782        // claim loop + semaphore, not the submit-side queue).
783        drop(reservation);
784        state.cluster().kick();
785        crate::serve::audit::write(
786            &state,
787            &actor,
788            "run.submit",
789            Some(run_id.clone()),
790            Some(fp_config.clone()),
791            "ok",
792        )
793        .await;
794        return Ok(SubmitResponse {
795            run_id,
796            status: RunStatus::Pending,
797            submitted_at,
798        });
799    }
800
801    if let Err(e) = state.history().upsert(&rec).await {
802        // See the cluster path above: release the orphaned claim (F21).
803        release_orphaned_claim(&state, &req, &run_id).await;
804        return Err(ServeError::Internal(e.to_string()));
805    }
806
807    let run_token = CancellationToken::new();
808    state.registry().register(run_id.clone(), run_token.clone());
809    metrics::set_run_gauges(&state);
810
811    // The spawned task now owns the queued→running→finished lifecycle.
812    reservation.defuse();
813    spawn_run(
814        state.clone(),
815        loaded,
816        req,
817        run_id.clone(),
818        run_token,
819        submitted_at,
820    );
821
822    crate::serve::audit::write(
823        &state,
824        &actor,
825        "run.submit",
826        Some(run_id.clone()),
827        Some(fp_config.clone()),
828        "ok",
829    )
830    .await;
831
832    Ok(SubmitResponse {
833        run_id,
834        status: RunStatus::Queued,
835        submitted_at,
836    })
837}
838
839/// Run the `doctor_first` probes; on any failure return 422 with the report.
840/// Run the `doctor_first` probes. On success returns the (redacted) report so
841/// the caller can store it on the run record (`doctor_report`); on any probe
842/// failure returns 422 with the same redacted report as `details`.
843pub(crate) async fn run_doctor_first(
844    state: &ServerState,
845    loaded: &LoadedSubmission,
846) -> Result<serde_json::Value, ServeError> {
847    use faucet_core::check::CheckContext;
848    let auth =
849        build_auth_catalog(loaded.cfg.auth.as_ref()).map_err(|e| ServeError::Unprocessable {
850            message: e.to_string(),
851            details: None,
852        })?;
853    let ctx = CheckContext {
854        timeout: state.probe_timeout(),
855    };
856    // Same pipeline-name derivation as the run path above, so SLA probes read
857    // the state keys the executor writes.
858    let pipeline_name = loaded
859        .cfg
860        .name
861        .clone()
862        .unwrap_or_else(|| "serve".to_string());
863    let mut invs = crate::commands::doctor::probe_roots(
864        &loaded.nodes,
865        &auth,
866        &ctx,
867        loaded.cfg.sla.as_ref(),
868        &pipeline_name,
869    )
870    .await;
871    let failed = crate::commands::doctor::count_failures(&invs);
872    // Redact regardless of outcome — the report is surfaced either way (as the
873    // 422 `details` on failure, or stored on the run record on success).
874    crate::commands::doctor::redact_invocations(&mut invs);
875    let report = serde_json::json!({ "invocations": invs });
876    if failed > 0 {
877        return Err(ServeError::Unprocessable {
878            message: format!("doctor_first preflight failed: {failed} probe(s) failed"),
879            details: Some(report),
880        });
881    }
882    Ok(report)
883}
884
885/// Build the replay response for an idempotency hit (the existing run's status).
886async fn replay_response(state: &ServerState, run_id: &str) -> Result<SubmitResponse, ServeError> {
887    let rec = state
888        .history()
889        .get(run_id)
890        .await
891        .map_err(|e| ServeError::Internal(e.to_string()))?
892        .ok_or(ServeError::NotFound)?;
893    Ok(SubmitResponse {
894        run_id: rec.run_id,
895        status: rec.status,
896        submitted_at: rec.submitted_at,
897    })
898}
899
900/// Releases a queue reservation on drop unless [`Self::defuse`]d. Guarantees the
901/// `queued` counter is balanced on every early-return path (replay / conflict /
902/// claim-or-upsert error) without a manual `release_reservation` at each site.
903/// Defused once the run is handed to the spawned task, which then owns the
904/// queued→running transition via `mark_running`.
905struct ReservationGuard {
906    state: Option<ServerState>,
907}
908
909impl ReservationGuard {
910    fn new(state: ServerState) -> Self {
911        Self { state: Some(state) }
912    }
913
914    /// Hand the reservation off to the spawned task (no release on drop).
915    fn defuse(mut self) {
916        self.state = None;
917    }
918}
919
920impl Drop for ReservationGuard {
921    fn drop(&mut self) {
922        if let Some(state) = self.state.take() {
923            state.registry().release_reservation();
924            metrics::set_run_gauges(&state);
925        }
926    }
927}
928
929/// Releases the in-flight slot (decrement `in_flight`, drop the cancel token, wake
930/// the shutdown drain) on drop — on EVERY path including panic. Without this, a
931/// panic between `mark_running` and a manual `mark_finished` would leak the
932/// counter and hang graceful shutdown forever.
933struct InFlightGuard {
934    state: ServerState,
935    run_id: String,
936}
937
938impl Drop for InFlightGuard {
939    fn drop(&mut self) {
940        self.state.registry().mark_finished(&self.run_id);
941        metrics::set_run_gauges(&self.state);
942    }
943}
944
945/// Releases a claimed shard's in-flight slot (decrement `in_flight`, drop the
946/// shard cancel token, wake the shutdown drain) on drop — on EVERY path
947/// including panic. Mirrors [`InFlightGuard`] but keys on the per-shard token so
948/// a Mode-B shard participates in the graceful-shutdown drain (audit #321 H5).
949struct ShardInFlightGuard {
950    state: ServerState,
951    run_id: String,
952    shard_id: String,
953}
954
955impl Drop for ShardInFlightGuard {
956    fn drop(&mut self) {
957        self.state
958            .registry()
959            .mark_shard_finished(&self.run_id, &self.shard_id);
960        metrics::set_run_gauges(&self.state);
961    }
962}
963
964/// Terminal classification of a run task.
965enum Terminal {
966    Completed {
967        records: u64,
968        invs: Vec<InvocationRecord>,
969    },
970    Failed {
971        reason: String,
972        records: u64,
973        invs: Vec<InvocationRecord>,
974    },
975    Timeout {
976        secs: u64,
977    },
978    Cancelled,
979    ShutdownFailed,
980}
981
982impl Terminal {
983    /// (status, metric reason label, records, invocations, error message)
984    fn into_parts(
985        self,
986    ) -> (
987        RunStatus,
988        &'static str,
989        u64,
990        Vec<InvocationRecord>,
991        Option<String>,
992    ) {
993        match self {
994            Terminal::Completed { records, invs } => {
995                (RunStatus::Completed, "ok", records, invs, None)
996            }
997            Terminal::Failed {
998                reason,
999                records,
1000                invs,
1001            } => (RunStatus::Failed, "error", records, invs, Some(reason)),
1002            Terminal::Timeout { secs } => (
1003                RunStatus::Failed,
1004                "timeout",
1005                0,
1006                Vec::new(),
1007                Some(format!("run exceeded timeout_secs ({secs}s)")),
1008            ),
1009            Terminal::Cancelled => (RunStatus::Cancelled, "cancelled", 0, Vec::new(), None),
1010            Terminal::ShutdownFailed => (
1011                RunStatus::Failed,
1012                "server_shutdown",
1013                0,
1014                Vec::new(),
1015                Some("server shutdown before the run finished".into()),
1016            ),
1017        }
1018    }
1019}
1020
1021/// Classify a finished `run_expanded` result into a `Terminal`.
1022fn classify_run(result: crate::error::CliResult<RunSummary>) -> Terminal {
1023    match result {
1024        Ok(summary) => {
1025            let records: u64 = summary
1026                .invocations
1027                .iter()
1028                .map(|i| i.records_written as u64)
1029                .sum();
1030            let invs: Vec<InvocationRecord> = summary
1031                .invocations
1032                .iter()
1033                .map(InvocationRecord::from)
1034                .collect();
1035            if summary.had_failures() {
1036                Terminal::Failed {
1037                    reason: format!("{} invocation(s) failed", summary.failure_count()),
1038                    records,
1039                    invs,
1040                }
1041            } else {
1042                Terminal::Completed { records, invs }
1043            }
1044        }
1045        Err(e) => Terminal::Failed {
1046            reason: e.to_string(),
1047            records: 0,
1048            invs: Vec::new(),
1049        },
1050    }
1051}
1052
1053/// Parse the optional request `clock` (RFC3339), defaulting to `submitted_at`.
1054fn resolve_clock(
1055    flag: Option<&str>,
1056    default: DateTime<Utc>,
1057) -> Result<DateTime<FixedOffset>, ServeError> {
1058    match flag {
1059        None => Ok(default.fixed_offset()),
1060        Some(s) => DateTime::parse_from_rfc3339(s)
1061            .map_err(|_| ServeError::BadConfig(format!("clock '{s}' is not RFC3339"))),
1062    }
1063}
1064
1065/// Spawn the detached run task: acquire a permit, run under the 3-arm select,
1066/// finalize the terminal status.
1067fn spawn_run(
1068    state: ServerState,
1069    loaded: LoadedSubmission,
1070    req: SubmitRequest,
1071    run_id: String,
1072    run_token: CancellationToken,
1073    submitted_at: DateTime<Utc>,
1074) {
1075    let server_shutdown = state.shutdown_token();
1076    tokio::spawn(async move {
1077        // Race the permit acquisition against cancel / shutdown so a run
1078        // cancelled while STILL QUEUED (before any permit frees) is finalized
1079        // immediately, instead of only after it eventually acquires a permit
1080        // (#146 R). `biased` prefers the cancel/shutdown signals over a
1081        // simultaneously-available permit.
1082        let _permit = tokio::select! {
1083            biased;
1084            _ = run_token.cancelled() => {
1085                finalize_queued_cancel(&state, &run_id, submitted_at, Terminal::Cancelled).await;
1086                return;
1087            }
1088            _ = server_shutdown.cancelled() => {
1089                finalize_queued_cancel(&state, &run_id, submitted_at, Terminal::ShutdownFailed).await;
1090                return;
1091            }
1092            permit = state.semaphore().acquire_owned() => permit.expect("semaphore not closed"),
1093        };
1094        execute_run(
1095            state,
1096            loaded,
1097            run_id,
1098            run_token,
1099            submitted_at,
1100            req.timeout_secs,
1101            req.clock,
1102            true,
1103        )
1104        .await;
1105        // `_permit` drops here.
1106    });
1107}
1108
1109/// The queued→running→finalize execution tail, shared by the submit path
1110/// (`spawn_run`) and the cluster claim path (`resume_claimed_run`). Assumes the
1111/// caller already holds an execution permit and has registered `run_token`.
1112/// `from_queue` is `true` when the run consumed a local queue slot (submit path)
1113/// and `false` for a cluster claim-path run that never reserved one (#228).
1114#[allow(clippy::too_many_arguments)]
1115async fn execute_run(
1116    state: ServerState,
1117    loaded: LoadedSubmission,
1118    run_id: String,
1119    run_token: CancellationToken,
1120    submitted_at: DateTime<Utc>,
1121    timeout_secs: Option<u64>,
1122    clock_flag: Option<String>,
1123    from_queue: bool,
1124) {
1125    let server_shutdown = state.shutdown_token();
1126    let LoadedSubmission { cfg, nodes } = loaded;
1127
1128    // Queued → running. From here the guard guarantees `mark_finished` (and a
1129    // gauge refresh) on EVERY exit, including early returns and panics.
1130    // A submit-path run consumed a local queue slot (Queued→Running); a cluster
1131    // claim-path run never did (#228) — only bump in_flight for it.
1132    if from_queue {
1133        state.registry().mark_running();
1134    } else {
1135        state.registry().mark_running_unqueued();
1136    }
1137    let _guard = InFlightGuard {
1138        state: state.clone(),
1139        run_id: run_id.clone(),
1140    };
1141    let started = Utc::now();
1142    if let Ok(Some(mut rec)) = state.history().get(&run_id).await {
1143        rec.status = RunStatus::Running;
1144        rec.started_at = Some(started);
1145        let _ = state.history().upsert(&rec).await;
1146    }
1147    metrics::set_run_gauges(&state);
1148
1149    // Build execution options (auth/clock failures finalize as Failed).
1150    let pipeline_name = cfg.name.clone().unwrap_or_else(|| "serve".to_string());
1151    let auth = match build_auth_catalog(cfg.auth.as_ref()) {
1152        Ok(a) => a,
1153        Err(e) => {
1154            finalize(
1155                &state,
1156                &run_id,
1157                started,
1158                Terminal::Failed {
1159                    reason: format!("auth catalog: {e}"),
1160                    records: 0,
1161                    invs: Vec::new(),
1162                },
1163            )
1164            .await;
1165            return;
1166        }
1167    };
1168    let clock = match resolve_clock(clock_flag.as_deref(), submitted_at) {
1169        Ok(c) => c,
1170        Err(e) => {
1171            finalize(
1172                &state,
1173                &run_id,
1174                started,
1175                Terminal::Failed {
1176                    reason: e.api_error().error.message,
1177                    records: 0,
1178                    invs: Vec::new(),
1179                },
1180            )
1181            .await;
1182            return;
1183        }
1184    };
1185
1186    // Cooperative-cancel token the pipeline observes so it flushes buffered
1187    // output (e.g. a Parquet footer, an S3 multipart upload) at its next
1188    // page boundary on cancel / timeout / shutdown — instead of having its
1189    // future hard-dropped, which flushes nothing (#146 H16).
1190    let coop = CancellationToken::new();
1191    // Resilience policy from the (merged) submitted config. A malformed
1192    // `resilience:` block finalizes the run as Failed, mirroring the
1193    // auth/clock failure handling above.
1194    let resilience = match &cfg.resilience {
1195        Some(spec) => match spec.to_policy() {
1196            Ok(p) => Some(p),
1197            Err(e) => {
1198                finalize(
1199                    &state,
1200                    &run_id,
1201                    started,
1202                    Terminal::Failed {
1203                        reason: format!("resilience: {e}"),
1204                        records: 0,
1205                        invs: Vec::new(),
1206                    },
1207                )
1208                .await;
1209                return;
1210            }
1211        },
1212        None => None,
1213    };
1214    // Build the per-run OpenLineage emitter from the (merged) submitted
1215    // config. A malformed `lineage:` block finalizes the run as Failed,
1216    // mirroring the auth/clock failure handling above.
1217    #[cfg(feature = "lineage")]
1218    let lineage = match crate::lineage_glue::build_emitter(cfg.lineage.as_ref()) {
1219        Ok(l) => l,
1220        Err(e) => {
1221            finalize(
1222                &state,
1223                &run_id,
1224                started,
1225                Terminal::Failed {
1226                    reason: format!("lineage: {e}"),
1227                    records: 0,
1228                    invs: Vec::new(),
1229                },
1230            )
1231            .await;
1232            return;
1233        }
1234    };
1235    // Notifier from the submitted config's `notifications:` block. A malformed
1236    // block disables notifications for this run (logged) rather than failing an
1237    // already-accepted run.
1238    #[cfg(feature = "notify")]
1239    let notifier = crate::notify::Notifier::from_specs(&cfg.notifications).unwrap_or_else(|e| {
1240        tracing::error!(%run_id, "notifications config invalid, disabling: {e}");
1241        None
1242    });
1243    let opts = ExecuteOptions {
1244        pipeline_name,
1245        // The id returned by `POST /v1/runs`, so a completion notification can be
1246        // matched back to the submission (#480).
1247        run_id: Some(run_id.clone()),
1248        execution: cfg.execution.clone(),
1249        dry_run: false,
1250        limit: None,
1251        state_path_override: None,
1252        shard: None,
1253        auth,
1254        clock,
1255        cancel: Some(coop.clone()),
1256        resilience,
1257        sla: cfg.sla.clone(),
1258        #[cfg(feature = "lineage")]
1259        lineage,
1260        #[cfg(feature = "lineage")]
1261        lineage_cfg: cfg.lineage.clone(),
1262        #[cfg(feature = "notify")]
1263        notifier,
1264        // Record every serve run into the Data Movement Catalog (#279),
1265        // stored alongside the run history (`--history` backend) and
1266        // attributed to this serve run id.
1267        #[cfg(feature = "catalog")]
1268        catalog: Some(crate::catalog::CatalogHandle {
1269            store: state.history(),
1270            run_id: Some(run_id.clone()),
1271            sample_records: crate::catalog::DEFAULT_SAMPLE_RECORDS,
1272        }),
1273    };
1274
1275    let span = tracing::info_span!("faucet.serve.run", serve_run_id = %run_id);
1276    let work = async move {
1277        // Emitted inside the run span so it is captured by the SSE log layer
1278        // (and gives every `/logs` reader at least one line to anchor on).
1279        tracing::info!("pipeline run starting");
1280        classify_run(run_expanded(nodes, opts).await)
1281    }
1282    .instrument(span);
1283    tokio::pin!(work);
1284
1285    // The run timeout is modelled as a cancel trigger (not a hard
1286    // `tokio::time::timeout` drop) so a timed-out run still flushes.
1287    let timeout_fut = async {
1288        match timeout_secs {
1289            Some(s) => tokio::time::sleep(Duration::from_secs(s)).await,
1290            None => std::future::pending::<()>().await,
1291        }
1292    };
1293    tokio::pin!(timeout_fut);
1294
1295    enum Trigger {
1296        Done(Terminal),
1297        Cancel,
1298        Shutdown,
1299        Timeout(u64),
1300    }
1301
1302    // Phase 1: run to natural completion, or until a cancel trigger fires.
1303    // `biased` prefers a just-completed run over a simultaneous trigger.
1304    let trigger = tokio::select! {
1305        biased;
1306        t = &mut work => Trigger::Done(t),
1307        _ = run_token.cancelled() => Trigger::Cancel,
1308        _ = server_shutdown.cancelled() => Trigger::Shutdown,
1309        _ = &mut timeout_fut => Trigger::Timeout(timeout_secs.unwrap_or(0)),
1310    };
1311
1312    let terminal = match trigger {
1313        Trigger::Done(t) => t,
1314        triggered => {
1315            // Phase 2: a trigger fired. Cancel cooperatively and give the
1316            // pipeline a bounded grace to flush at its next page boundary,
1317            // then hard-drop it (drops the JoinSet, aborting any pipeline
1318            // genuinely stuck mid-write) so a hung run can't wedge shutdown.
1319            coop.cancel();
1320            let trigger_terminal = match triggered {
1321                Trigger::Cancel => Terminal::Cancelled,
1322                Trigger::Shutdown => Terminal::ShutdownFailed,
1323                Trigger::Timeout(secs) => Terminal::Timeout { secs },
1324                Trigger::Done(_) => unreachable!("matched in the outer arm"),
1325            };
1326            // A cooperative cancel makes `run_stream` flush and return Ok(partial),
1327            // so the trigger label (Cancelled/Timeout/ShutdownFailed) is the
1328            // correct status for that path. But if the flush itself FAILS within
1329            // the grace window, surface that real failure — never mask it behind
1330            // the trigger label, which would hide a data error / partial write.
1331            match tokio::time::timeout(RUN_FLUSH_GRACE, &mut work).await {
1332                Ok(failed @ Terminal::Failed { .. }) => failed,
1333                Ok(_) | Err(_) => trigger_terminal,
1334            }
1335        }
1336    };
1337
1338    finalize(&state, &run_id, started, terminal).await;
1339    // Signal `/logs` readers the run is done, then drop the buffer after a
1340    // drain window so a late fetcher can still replay it (spec §12).
1341    state.log_hub().finish(&run_id);
1342    schedule_log_drop(state.clone(), run_id.clone());
1343    // `_guard` drops here → mark_finished + gauge refresh.
1344}
1345
1346/// Write the authoritative terminal record + the run-finished metric.
1347async fn finalize(state: &ServerState, run_id: &str, started: DateTime<Utc>, term: Terminal) {
1348    let finished = Utc::now();
1349    let elapsed = (finished - started).to_std().ok().map(|d| d.as_secs_f64());
1350    let (status, reason, records, invs, error) = term.into_parts();
1351    // Read-modify-write the existing record to preserve its metadata
1352    // (name / labels / idempotency_key / submitted_at). If it can't be read —
1353    // the backend errored, or the record was purged / landed in another store
1354    // under degraded fallback — DON'T silently drop the terminal status (#146
1355    // M6): reconstruct a minimal terminal record and upsert it, so the run
1356    // never lingers non-terminal while `record_run_finished` has already fired.
1357    let mut rec = match state.history().get(run_id).await {
1358        Ok(Some(rec)) => rec,
1359        Ok(None) => {
1360            tracing::warn!(
1361                run_id,
1362                "finalize: run record not found; writing a fresh terminal record"
1363            );
1364            RunRecord::queued(run_id.to_string(), None, BTreeMap::new(), None, started)
1365        }
1366        Err(e) => {
1367            tracing::warn!(
1368                run_id,
1369                error = %e,
1370                "finalize: failed to read run record; writing a fresh terminal record"
1371            );
1372            RunRecord::queued(run_id.to_string(), None, BTreeMap::new(), None, started)
1373        }
1374    };
1375    rec.status = status;
1376    rec.started_at.get_or_insert(started);
1377    rec.finished_at = Some(finished);
1378    rec.elapsed_secs = elapsed;
1379    rec.records_written = records;
1380    rec.invocations = invs;
1381    rec.error = error;
1382    if state.cluster().enabled() {
1383        // Owner-fenced: if another instance reclaimed this run (lease expired),
1384        // our write is a no-op and we discard our result — the reclaimer is now
1385        // authoritative (#197).
1386        match state.history().finalize_owned(&rec).await {
1387            Ok(true) => {
1388                metrics::record_run_finished(status, reason);
1389                // Only the instance that won the owner-fenced write fires the
1390                // callback, so a reclaimed run cannot deliver twice (#481).
1391                crate::serve::callback::fire(&rec).await;
1392            }
1393            Ok(false) => tracing::warn!(
1394                run_id,
1395                "finalize: run was reclaimed by another instance; discarding result"
1396            ),
1397            Err(e) => {
1398                tracing::error!(run_id, error = %e, "finalize: owner-fenced write failed")
1399            }
1400        }
1401    } else {
1402        if let Err(e) = state.history().upsert(&rec).await {
1403            tracing::error!(
1404                run_id,
1405                error = %e,
1406                "finalize: failed to persist terminal run record"
1407            );
1408        }
1409        metrics::record_run_finished(status, reason);
1410        crate::serve::callback::fire(&rec).await;
1411    }
1412}
1413
1414/// Finalize a run that was cancelled / hit shutdown while still QUEUED (before
1415/// it acquired an execution permit, so it never became in-flight and has no
1416/// `InFlightGuard`). Releases the queue slot, writes the terminal record, closes
1417/// the log buffer, and refreshes the gauges — the queued-path analogue of the
1418/// normal `finalize` + `InFlightGuard`-drop cleanup.
1419async fn finalize_queued_cancel(
1420    state: &ServerState,
1421    run_id: &str,
1422    submitted_at: DateTime<Utc>,
1423    term: Terminal,
1424) {
1425    state.registry().mark_queued_cancelled(run_id);
1426    finalize(state, run_id, submitted_at, term).await;
1427    state.log_hub().finish(run_id);
1428    schedule_log_drop(state.clone(), run_id.to_string());
1429    metrics::set_run_gauges(state);
1430}
1431
1432/// Spawn a detached timer that drops a finished run's log buffer after the drain
1433/// window, freeing its ring once late `/logs` fetchers have had a chance to read.
1434fn schedule_log_drop(state: ServerState, run_id: String) {
1435    tokio::spawn(async move {
1436        tokio::time::sleep(crate::serve::logs::LOG_DRAIN).await;
1437        state.log_hub().drop_run(&run_id);
1438    });
1439}
1440
1441#[cfg(test)]
1442mod tests {
1443    use super::*;
1444
1445    /// A stand-in admin actor for the submit() tests.
1446    fn admin_actor() -> AuthContext {
1447        AuthContext {
1448            principal: "test".into(),
1449            role: crate::serve::rbac::Role::Admin,
1450            source_ip: None,
1451        }
1452    }
1453
1454    #[test]
1455    fn classify_ok_no_failures_is_completed() {
1456        let summary = RunSummary {
1457            invocations: vec![crate::executor::InvocationOutcome {
1458                row_id: "r".into(),
1459                parent_record_key: None,
1460                records_written: 3,
1461                error: None,
1462                metrics: None,
1463            }],
1464        };
1465        let (status, reason, records, _, error) = classify_run(Ok(summary)).into_parts();
1466        assert_eq!(status, RunStatus::Completed);
1467        assert_eq!(reason, "ok");
1468        assert_eq!(records, 3);
1469        assert!(error.is_none());
1470    }
1471
1472    #[test]
1473    fn classify_ok_with_failures_is_failed() {
1474        let summary = RunSummary {
1475            invocations: vec![crate::executor::InvocationOutcome {
1476                row_id: "r".into(),
1477                parent_record_key: None,
1478                records_written: 0,
1479                error: Some("boom".into()),
1480                metrics: None,
1481            }],
1482        };
1483        let (status, reason, _, _, error) = classify_run(Ok(summary)).into_parts();
1484        assert_eq!(status, RunStatus::Failed);
1485        assert_eq!(reason, "error");
1486        assert!(error.unwrap().contains("invocation(s) failed"));
1487    }
1488
1489    #[test]
1490    fn timeout_maps_to_failed_with_timeout_reason() {
1491        let (status, reason, _, _, error) = Terminal::Timeout { secs: 30 }.into_parts();
1492        assert_eq!(status, RunStatus::Failed);
1493        assert_eq!(reason, "timeout");
1494        assert!(error.unwrap().contains("30s"));
1495    }
1496
1497    #[test]
1498    fn resolve_clock_defaults_and_parses() {
1499        let default = Utc::now();
1500        assert_eq!(
1501            resolve_clock(None, default).unwrap(),
1502            default.fixed_offset()
1503        );
1504        assert!(resolve_clock(Some("2026-01-31T00:00:00Z"), default).is_ok());
1505        assert!(resolve_clock(Some("not-a-time"), default).is_err());
1506    }
1507
1508    #[tokio::test]
1509    async fn conflict_releases_reservation() {
1510        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1511        use crate::serve::history::RunHistory;
1512        use crate::serve::history::memory::MemoryHistory;
1513        use crate::serve::state::ServerState;
1514        use std::sync::Arc;
1515        use tokio_util::sync::CancellationToken;
1516
1517        let cfg = ServeConfig {
1518            listen: "127.0.0.1:0".parse().unwrap(),
1519            auth: AuthMode::None,
1520            max_concurrent_runs: 4,
1521            max_queued_runs: 4,
1522            default_config_path: None,
1523            history: HistoryBackendSpec::Memory,
1524            cors_origins: vec![],
1525            body_limit_bytes: 1_048_576,
1526            shutdown_grace: Duration::from_secs(60),
1527            retain_terminal_runs: Duration::from_secs(60),
1528            idempotency_retention: Duration::from_secs(60),
1529            lease_ttl: Duration::from_secs(30),
1530            probe_timeout: Duration::from_secs(10),
1531            env_file: None,
1532            no_env_file: false,
1533            log_level: "info".into(),
1534            ui_enabled: true,
1535            cluster: crate::serve::cluster::ClusterConfig::disabled(),
1536            triggers_path: None,
1537            callback_allow_hosts: Vec::new(),
1538        };
1539        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1540        let state = ServerState::new(
1541            &cfg,
1542            None,
1543            CancellationToken::new(),
1544            history,
1545            crate::serve::logs::LogHub::new(),
1546            None,
1547            #[cfg(feature = "triggers")]
1548            crate::serve::triggers::health::TriggersHandle::empty(),
1549        );
1550
1551        // Pre-claim the key with a DIFFERENT fingerprint so submit() hits Conflict.
1552        state
1553            .history()
1554            .claim_idempotency("k", "different-fp", "prior", Duration::from_secs(60))
1555            .await
1556            .unwrap();
1557
1558        let req = SubmitRequest {
1559            config: "version: 1\npipeline:\n  source: { type: csv, config: { path: x.csv } }\n  sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1560            config_format: ConfigFormatWire::Yaml,
1561            name: None,
1562            labels: BTreeMap::new(),
1563            timeout_secs: None,
1564            doctor_first: false,
1565            callback: None,
1566            idempotency_key: Some("k".into()),
1567            clock: None,
1568        };
1569
1570        let err = submit(state.clone(), req, admin_actor()).await.unwrap_err();
1571        assert!(
1572            matches!(err, ServeError::Conflict(_)),
1573            "expected Conflict, got {err:?}"
1574        );
1575        // The reservation taken before the claim must have been released by the guard.
1576        assert_eq!(state.registry().queued(), 0);
1577    }
1578
1579    #[tokio::test]
1580    async fn cluster_submit_writes_pending_with_config_and_does_not_spawn() {
1581        use crate::serve::cluster::ClusterConfig;
1582        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1583        use crate::serve::history::RunHistory;
1584        use crate::serve::history::memory::MemoryHistory;
1585        use crate::serve::state::ServerState;
1586        use std::sync::Arc;
1587        use tokio_util::sync::CancellationToken;
1588
1589        let mut cluster = ClusterConfig::disabled();
1590        cluster.enabled = true;
1591        let cfg = ServeConfig {
1592            listen: "127.0.0.1:0".parse().unwrap(),
1593            auth: AuthMode::None,
1594            max_concurrent_runs: 4,
1595            max_queued_runs: 4,
1596            default_config_path: None,
1597            history: HistoryBackendSpec::Memory,
1598            cors_origins: vec![],
1599            body_limit_bytes: 1_048_576,
1600            shutdown_grace: Duration::from_secs(60),
1601            retain_terminal_runs: Duration::from_secs(60),
1602            idempotency_retention: Duration::from_secs(60),
1603            lease_ttl: Duration::from_secs(30),
1604            probe_timeout: Duration::from_secs(10),
1605            env_file: None,
1606            no_env_file: false,
1607            log_level: "info".into(),
1608            ui_enabled: true,
1609            cluster,
1610            triggers_path: None,
1611            callback_allow_hosts: Vec::new(),
1612        };
1613        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1614        let state = ServerState::new(
1615            &cfg,
1616            None,
1617            CancellationToken::new(),
1618            history,
1619            crate::serve::logs::LogHub::new(),
1620            None,
1621            #[cfg(feature = "triggers")]
1622            crate::serve::triggers::health::TriggersHandle::empty(),
1623        );
1624
1625        let req = SubmitRequest {
1626            config: "version: 1\npipeline:\n  source: { type: csv, config: { path: x.csv } }\n  sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1627            config_format: ConfigFormatWire::Yaml,
1628            name: Some("n".into()),
1629            labels: BTreeMap::new(),
1630            timeout_secs: Some(99),
1631            doctor_first: false,
1632            callback: None,
1633            idempotency_key: None,
1634            clock: None,
1635        };
1636        let resp = submit(state.clone(), req, admin_actor()).await.unwrap();
1637        assert_eq!(resp.status, RunStatus::Pending);
1638        // No local queue slot was consumed (cluster runs don't queue locally).
1639        assert_eq!(state.registry().queued(), 0);
1640        let rec = state.history().get(&resp.run_id).await.unwrap().unwrap();
1641        assert_eq!(rec.status, RunStatus::Pending);
1642        assert!(rec.config_body.as_deref().unwrap().contains("version: 1"));
1643        assert_eq!(rec.timeout_secs, Some(99));
1644    }
1645
1646    /// A `ServerState` backed by an in-memory history, for finalize tests.
1647    fn memory_state() -> crate::serve::state::ServerState {
1648        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1649        use crate::serve::history::RunHistory;
1650        use crate::serve::history::memory::MemoryHistory;
1651        use crate::serve::state::ServerState;
1652        use std::sync::Arc;
1653        use tokio_util::sync::CancellationToken;
1654
1655        let cfg = ServeConfig {
1656            listen: "127.0.0.1:0".parse().unwrap(),
1657            auth: AuthMode::None,
1658            max_concurrent_runs: 4,
1659            max_queued_runs: 4,
1660            default_config_path: None,
1661            history: HistoryBackendSpec::Memory,
1662            cors_origins: vec![],
1663            body_limit_bytes: 1_048_576,
1664            shutdown_grace: Duration::from_secs(60),
1665            retain_terminal_runs: Duration::from_secs(60),
1666            idempotency_retention: Duration::from_secs(60),
1667            lease_ttl: Duration::from_secs(30),
1668            probe_timeout: Duration::from_secs(10),
1669            env_file: None,
1670            no_env_file: false,
1671            log_level: "info".into(),
1672            ui_enabled: true,
1673            cluster: crate::serve::cluster::ClusterConfig::disabled(),
1674            triggers_path: None,
1675            callback_allow_hosts: Vec::new(),
1676        };
1677        let history = Arc::new(MemoryHistory::new(Duration::from_secs(60))) as Arc<dyn RunHistory>;
1678        ServerState::new(
1679            &cfg,
1680            None,
1681            CancellationToken::new(),
1682            history,
1683            crate::serve::logs::LogHub::new(),
1684            None,
1685            #[cfg(feature = "triggers")]
1686            crate::serve::triggers::health::TriggersHandle::empty(),
1687        )
1688    }
1689
1690    #[tokio::test]
1691    async fn finalize_writes_terminal_record_when_record_is_missing() {
1692        // M6 (#146): if the run record can't be read at finalize time (purged,
1693        // or split to another store under degraded fallback), the terminal
1694        // status must NOT be silently dropped — a fresh terminal record is
1695        // written so the run never lingers non-terminal while the run-finished
1696        // metric has already fired.
1697        let state = memory_state();
1698        let started = Utc::now();
1699        finalize(
1700            &state,
1701            "ghost",
1702            started,
1703            Terminal::Failed {
1704                reason: "boom".into(),
1705                records: 0,
1706                invs: Vec::new(),
1707            },
1708        )
1709        .await;
1710        let rec = state
1711            .history()
1712            .get("ghost")
1713            .await
1714            .unwrap()
1715            .expect("finalize must create a terminal record even when none existed");
1716        assert_eq!(rec.status, RunStatus::Failed);
1717        assert!(rec.finished_at.is_some());
1718        assert!(rec.started_at.is_some());
1719        assert_eq!(rec.error.as_deref(), Some("boom"));
1720    }
1721
1722    #[tokio::test]
1723    async fn finalize_preserves_metadata_of_existing_record() {
1724        // The happy path still read-modify-writes, preserving name/labels/key.
1725        let state = memory_state();
1726        let started = Utc::now();
1727        let mut rec = RunRecord::queued(
1728            "r1".into(),
1729            Some("nightly".into()),
1730            BTreeMap::new(),
1731            Some("idem-k".into()),
1732            started,
1733        );
1734        rec.status = RunStatus::Running;
1735        rec.started_at = Some(started);
1736        state.history().upsert(&rec).await.unwrap();
1737
1738        finalize(
1739            &state,
1740            "r1",
1741            started,
1742            Terminal::Completed {
1743                records: 5,
1744                invs: Vec::new(),
1745            },
1746        )
1747        .await;
1748        let got = state.history().get("r1").await.unwrap().unwrap();
1749        assert_eq!(got.status, RunStatus::Completed);
1750        assert_eq!(got.records_written, 5);
1751        assert_eq!(got.name.as_deref(), Some("nightly"));
1752        assert_eq!(got.idempotency_key.as_deref(), Some("idem-k"));
1753    }
1754
1755    #[cfg(any(feature = "serve-history-sqlite", feature = "serve-history-postgres"))]
1756    #[tokio::test]
1757    async fn cluster_submit_503s_when_history_degraded() {
1758        // #197 spec §9: a degraded backend can't coordinate a cluster, so submit
1759        // must fail closed with 503 rather than orphan a never-claimable Pending run.
1760        use crate::serve::cluster::ClusterConfig;
1761        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1762        use crate::serve::history::RunHistory;
1763        use crate::serve::history::fallback::FallbackHistory;
1764        use crate::serve::state::ServerState;
1765        use std::sync::Arc;
1766        use tokio_util::sync::CancellationToken;
1767
1768        let mut cluster = ClusterConfig::disabled();
1769        cluster.enabled = true;
1770        let cfg = ServeConfig {
1771            listen: "127.0.0.1:0".parse().unwrap(),
1772            auth: AuthMode::None,
1773            max_concurrent_runs: 4,
1774            max_queued_runs: 4,
1775            default_config_path: None,
1776            history: HistoryBackendSpec::Memory,
1777            cors_origins: vec![],
1778            body_limit_bytes: 1_048_576,
1779            shutdown_grace: Duration::from_secs(60),
1780            retain_terminal_runs: Duration::from_secs(60),
1781            idempotency_retention: Duration::from_secs(60),
1782            lease_ttl: Duration::from_secs(30),
1783            probe_timeout: Duration::from_secs(10),
1784            env_file: None,
1785            no_env_file: false,
1786            log_level: "info".into(),
1787            ui_enabled: true,
1788            cluster,
1789            triggers_path: None,
1790            callback_allow_hosts: Vec::new(),
1791        };
1792        // A backend that is degraded from startup (primary unreachable).
1793        let history = Arc::new(FallbackHistory::degraded_at_startup(
1794            Duration::from_secs(60),
1795            "test",
1796        )) as Arc<dyn RunHistory>;
1797        assert!(history.degraded());
1798        let state = ServerState::new(
1799            &cfg,
1800            None,
1801            CancellationToken::new(),
1802            history,
1803            crate::serve::logs::LogHub::new(),
1804            None,
1805            #[cfg(feature = "triggers")]
1806            crate::serve::triggers::health::TriggersHandle::empty(),
1807        );
1808        let req = SubmitRequest {
1809            config: "version: 1\npipeline:\n  source: { type: csv, config: { path: x.csv } }\n  sink: { type: jsonl, config: { path: out.jsonl } }\n".into(),
1810            config_format: ConfigFormatWire::Yaml,
1811            name: None,
1812            labels: BTreeMap::new(),
1813            timeout_secs: None,
1814            doctor_first: false,
1815            callback: None,
1816            idempotency_key: None,
1817            clock: None,
1818        };
1819        let err = submit(state.clone(), req, admin_actor()).await.unwrap_err();
1820        assert!(
1821            matches!(err, ServeError::Unavailable(_)),
1822            "expected 503 Unavailable, got {err:?}"
1823        );
1824        // The queue reservation must have been released (no leak).
1825        assert_eq!(state.registry().queued(), 0);
1826    }
1827
1828    // ── Mode B coverage: coordinator / parent finalize / shard execution ─────
1829    // SQLite-backed so the shard RunHistory methods are live (the memory backend
1830    // is inert for shards). No Docker: the S3 source builds offline and its
1831    // enumerate_shards is pure (hash-modulo); csv→jsonl runs entirely on temp
1832    // files.
1833    #[cfg(feature = "serve-history-sqlite")]
1834    mod shards {
1835        use super::*;
1836        use crate::serve::config::{AuthMode, HistoryBackendSpec, ServeConfig};
1837        use crate::serve::history::RunHistory;
1838        use crate::serve::history::sqlite::SqliteHistory;
1839        use crate::serve::load::{ConfigFormat, load_submission};
1840        use crate::serve::state::ServerState;
1841        use faucet_core::ShardSpec;
1842        use std::collections::BTreeMap;
1843        use std::sync::Arc;
1844        use tokio_util::sync::CancellationToken;
1845
1846        async fn sqlite_state(dir: &std::path::Path) -> ServerState {
1847            let url = format!("sqlite://{}/h.db", dir.display());
1848            let history = Arc::new(
1849                SqliteHistory::connect(
1850                    &url,
1851                    Duration::from_secs(300),
1852                    Duration::from_secs(300),
1853                    "inst-test".into(),
1854                )
1855                .await
1856                .expect("sqlite history"),
1857            ) as Arc<dyn RunHistory>;
1858            let cfg = ServeConfig {
1859                listen: "127.0.0.1:0".parse().unwrap(),
1860                auth: AuthMode::None,
1861                max_concurrent_runs: 4,
1862                max_queued_runs: 4,
1863                default_config_path: None,
1864                history: HistoryBackendSpec::Memory,
1865                cors_origins: vec![],
1866                body_limit_bytes: 1_048_576,
1867                shutdown_grace: Duration::from_secs(60),
1868                retain_terminal_runs: Duration::from_secs(60),
1869                idempotency_retention: Duration::from_secs(60),
1870                lease_ttl: Duration::from_secs(30),
1871                probe_timeout: Duration::from_secs(10),
1872                env_file: None,
1873                no_env_file: false,
1874                log_level: "info".into(),
1875                ui_enabled: true,
1876                cluster: crate::serve::cluster::ClusterConfig::disabled(),
1877                triggers_path: None,
1878                callback_allow_hosts: Vec::new(),
1879            };
1880            ServerState::new(
1881                &cfg,
1882                None,
1883                CancellationToken::new(),
1884                history,
1885                crate::serve::logs::LogHub::new(),
1886                None,
1887                #[cfg(feature = "triggers")]
1888                crate::serve::triggers::health::TriggersHandle::empty(),
1889            )
1890        }
1891
1892        async fn loaded(yaml: &str) -> LoadedSubmission {
1893            load_submission(yaml, ConfigFormat::Yaml, None)
1894                .await
1895                .expect("load submission")
1896        }
1897
1898        async fn seed_run(state: &ServerState, run_id: &str, status: RunStatus) {
1899            let mut rec = RunRecord::queued(run_id.into(), None, BTreeMap::new(), None, Utc::now());
1900            rec.status = status;
1901            rec.config_body = Some("version: 1".into());
1902            state.history().upsert(&rec).await.expect("seed run");
1903        }
1904
1905        #[tokio::test]
1906        async fn coordinate_matrix_run_is_not_shardable() {
1907            // A matrix expands to >1 node → not shardable → Ok(false), no build.
1908            let dir = tempfile::tempdir().unwrap();
1909            let state = sqlite_state(dir.path()).await;
1910            let l = loaded(
1911                "version: 1\nname: m\nmatrix:\n  - id: a\n  - id: b\npipeline:\n  \
1912                 source: { type: rest, config: { url: \"http://localhost/x\" } }\n  \
1913                 sink: { type: stdout, config: {} }\n",
1914            )
1915            .await;
1916            assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1917        }
1918
1919        #[tokio::test]
1920        async fn coordinate_non_shardable_source_runs_whole() {
1921            // A csv source is not shardable → Ok(false) (built offline).
1922            let dir = tempfile::tempdir().unwrap();
1923            let state = sqlite_state(dir.path()).await;
1924            let input = dir.path().join("in.csv");
1925            std::fs::write(&input, "id\n1\n").unwrap();
1926            let l = loaded(&format!(
1927                "version: 1\npipeline:\n  \
1928                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
1929                 sink: {{ type: stdout, config: {{}} }}\n",
1930                input.display()
1931            ))
1932            .await;
1933            assert!(!coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1934        }
1935
1936        #[tokio::test]
1937        async fn coordinate_s3_source_inserts_shards_and_marks_sharded() {
1938            let dir = tempfile::tempdir().unwrap();
1939            let state = sqlite_state(dir.path()).await;
1940            seed_run(&state, "r", RunStatus::Running).await;
1941            let l = loaded(
1942                "version: 1\npipeline:\n  \
1943                 source: { type: s3, config: { bucket: my-bucket, prefix: null, \
1944                 region: null, endpoint_url: null, file_format: json_lines, \
1945                 max_objects: null, concurrency: 10 } }\n  \
1946                 sink: { type: stdout, config: {} }\n",
1947            )
1948            .await;
1949            assert!(coordinate_sharded_run(&state, "r", &l, 4).await.unwrap());
1950            // 4 shards inserted; parent flipped to Sharded.
1951            let prog = state.history().shard_progress("r").await.unwrap();
1952            assert_eq!(prog.total, 4);
1953            assert_eq!(prog.pending, 4);
1954            assert_eq!(
1955                state.history().get("r").await.unwrap().unwrap().status,
1956                RunStatus::Sharded
1957            );
1958        }
1959
1960        #[tokio::test]
1961        async fn at_least_once_risky_sinks_flags_append_cluster_and_shard() {
1962            // F26/F39: a clustered or sharded run with an append-mode sink is
1963            // at-least-once → flagged. Neither flag → not flagged.
1964            let append = loaded(
1965                "version: 1\npipeline:\n  \
1966                 source: { type: rest, config: { url: \"http://localhost/x\" } }\n  \
1967                 sink: { type: stdout, config: {} }\n",
1968            )
1969            .await;
1970            // Not clustered, not sharded → no warning.
1971            assert!(at_least_once_risky_sinks(&append, false, false).is_empty());
1972            // Clustered append → flagged.
1973            assert_eq!(
1974                at_least_once_risky_sinks(&append, true, false),
1975                vec!["stdout"]
1976            );
1977            // Sharded append → flagged.
1978            assert_eq!(
1979                at_least_once_risky_sinks(&append, false, true),
1980                vec!["stdout"]
1981            );
1982        }
1983
1984        #[tokio::test]
1985        async fn at_least_once_risky_sinks_safe_for_upsert_and_exactly_once() {
1986            // Upsert sink → keyed re-writes are idempotent → not flagged.
1987            let upsert = loaded(
1988                "version: 1\npipeline:\n  \
1989                 source: { type: postgres, config: { connection_url: \"postgres://x\", \
1990                 query: \"select 1\" } }\n  \
1991                 sink: { type: postgres, config: { connection_url: \"postgres://y\", \
1992                 table_name: t, column_mapping: auto_map, write_mode: upsert, key: [id] } }\n",
1993            )
1994            .await;
1995            assert!(at_least_once_risky_sinks(&upsert, true, true).is_empty());
1996
1997            // exactly_once delivery → not flagged even with an append sink.
1998            let eo = loaded(
1999                "version: 1\ndelivery: exactly_once\npipeline:\n  \
2000                 source: { type: postgres-cdc, config: {} }\n  \
2001                 sink: { type: sqlite, config: {} }\n  \
2002                 state: { type: file, config: { path: \"/tmp/x.json\" } }\n",
2003            )
2004            .await;
2005            assert!(at_least_once_risky_sinks(&eo, true, false).is_empty());
2006        }
2007
2008        async fn seed_sharded_with_shards(state: &ServerState, run_id: &str, n: usize) {
2009            use crate::serve::history::ShardInsert;
2010            seed_run(state, run_id, RunStatus::Sharded).await;
2011            let shards: Vec<ShardInsert> = (0..n)
2012                .map(|i| ShardInsert {
2013                    shard_id: i.to_string(),
2014                    descriptor: serde_json::json!({ "i": i }),
2015                    size_estimate: None,
2016                })
2017                .collect();
2018            state
2019                .history()
2020                .insert_shards(run_id, &shards)
2021                .await
2022                .unwrap();
2023            // Claim them so they are 'running' and finalizable by this instance.
2024            let claimed = state.history().claim_shards(n).await.unwrap();
2025            assert_eq!(claimed.len(), n);
2026        }
2027
2028        #[tokio::test]
2029        async fn maybe_finalize_parent_completes_when_all_shards_succeed() {
2030            let dir = tempfile::tempdir().unwrap();
2031            let state = sqlite_state(dir.path()).await;
2032            seed_sharded_with_shards(&state, "r", 3).await;
2033            for i in 0..3 {
2034                state
2035                    .history()
2036                    .finalize_shard("r", &i.to_string(), true)
2037                    .await
2038                    .unwrap();
2039            }
2040            maybe_finalize_parent(&state, "r").await;
2041            assert_eq!(
2042                state.history().get("r").await.unwrap().unwrap().status,
2043                RunStatus::Completed
2044            );
2045        }
2046
2047        #[tokio::test]
2048        async fn maybe_finalize_parent_fails_when_a_shard_fails() {
2049            let dir = tempfile::tempdir().unwrap();
2050            let state = sqlite_state(dir.path()).await;
2051            seed_sharded_with_shards(&state, "r", 2).await;
2052            state
2053                .history()
2054                .finalize_shard("r", "0", true)
2055                .await
2056                .unwrap();
2057            state
2058                .history()
2059                .finalize_shard("r", "1", false)
2060                .await
2061                .unwrap();
2062            maybe_finalize_parent(&state, "r").await;
2063            assert_eq!(
2064                state.history().get("r").await.unwrap().unwrap().status,
2065                RunStatus::Failed
2066            );
2067        }
2068
2069        #[tokio::test]
2070        async fn maybe_finalize_parent_keeps_sharded_until_all_terminal() {
2071            let dir = tempfile::tempdir().unwrap();
2072            let state = sqlite_state(dir.path()).await;
2073            seed_sharded_with_shards(&state, "r", 2).await;
2074            // Only one shard finalized → run stays Sharded.
2075            state
2076                .history()
2077                .finalize_shard("r", "0", true)
2078                .await
2079                .unwrap();
2080            maybe_finalize_parent(&state, "r").await;
2081            assert_eq!(
2082                state.history().get("r").await.unwrap().unwrap().status,
2083                RunStatus::Sharded
2084            );
2085        }
2086
2087        #[tokio::test]
2088        async fn finalize_sharded_parent_is_status_fenced_and_idempotent() {
2089            // F45: the status-fenced finalize transitions the parent exactly
2090            // once. A concurrent second finalize (e.g. two shards finishing on
2091            // two instances) is a no-op, and a non-Sharded run is never touched.
2092            let dir = tempfile::tempdir().unwrap();
2093            let state = sqlite_state(dir.path()).await;
2094            seed_sharded_with_shards(&state, "r", 2).await;
2095
2096            let first = state
2097                .history()
2098                .finalize_sharded_parent("r", RunStatus::Completed, Utc::now(), None)
2099                .await
2100                .unwrap();
2101            assert!(first, "first finalize wins");
2102            assert_eq!(
2103                state.history().get("r").await.unwrap().unwrap().status,
2104                RunStatus::Completed
2105            );
2106
2107            // Second finalize sees a non-Sharded parent → no-op, status unchanged.
2108            let second = state
2109                .history()
2110                .finalize_sharded_parent("r", RunStatus::Failed, Utc::now(), Some("late".into()))
2111                .await
2112                .unwrap();
2113            assert!(!second, "second finalize is a no-op");
2114            let r = state.history().get("r").await.unwrap().unwrap();
2115            assert_eq!(r.status, RunStatus::Completed, "status not overwritten");
2116            assert!(r.error.is_none(), "late error must not be stamped");
2117
2118            // An unknown / never-sharded run id is not finalized.
2119            let missing = state
2120                .history()
2121                .finalize_sharded_parent("does-not-exist", RunStatus::Completed, Utc::now(), None)
2122                .await
2123                .unwrap();
2124            assert!(!missing);
2125        }
2126
2127        #[tokio::test]
2128        async fn execute_shard_runs_a_csv_to_jsonl_shard() {
2129            let dir = tempfile::tempdir().unwrap();
2130            let state = sqlite_state(dir.path()).await;
2131            let input = dir.path().join("in.csv");
2132            std::fs::write(&input, "id,name\n1,alice\n2,bob\n").unwrap();
2133            let output = dir.path().join("out.jsonl");
2134            let yaml = format!(
2135                "version: 1\npipeline:\n  \
2136                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
2137                 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2138                input.display(),
2139                output.display()
2140            );
2141            let l = loaded(&yaml).await;
2142            // The whole-dataset shard is a no-op for the (non-shardable) csv source,
2143            // exercising the apply_shard call + per-shard state-key path end-to-end.
2144            let ok = execute_shard(
2145                &state,
2146                l,
2147                "r",
2148                "0",
2149                ShardSpec::whole(),
2150                CancellationToken::new(),
2151                None,
2152                None,
2153                Utc::now(),
2154            )
2155            .await;
2156            assert!(ok, "csv→jsonl shard should complete");
2157            let written = std::fs::read_to_string(&output).unwrap();
2158            assert_eq!(written.lines().count(), 2, "both rows written");
2159            assert!(written.contains("alice") && written.contains("bob"));
2160        }
2161
2162        #[tokio::test]
2163        async fn resume_claimed_shard_executes_and_finalizes_parent() {
2164            // End-to-end per-shard entry point: claim a shard whose run is a
2165            // csv→jsonl pipeline, dispatch it, and confirm the shard runs, is
2166            // finalized, and the parent run flips to Completed.
2167            let dir = tempfile::tempdir().unwrap();
2168            let state = sqlite_state(dir.path()).await;
2169            let input = dir.path().join("in.csv");
2170            std::fs::write(&input, "id,name\n1,alice\n").unwrap();
2171            let output = dir.path().join("out.jsonl");
2172            let yaml = format!(
2173                "version: 1\npipeline:\n  \
2174                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
2175                 sink: {{ type: jsonl, config: {{ path: \"{}\" }} }}\n",
2176                input.display(),
2177                output.display()
2178            );
2179            // Seed the parent Sharded run carrying the pipeline config, + one shard.
2180            let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2181            rec.status = RunStatus::Sharded;
2182            rec.config_body = Some(yaml);
2183            state.history().upsert(&rec).await.unwrap();
2184            use crate::serve::history::ShardInsert;
2185            state
2186                .history()
2187                .insert_shards(
2188                    "r",
2189                    &[ShardInsert {
2190                        shard_id: "0".into(),
2191                        descriptor: serde_json::Value::Null,
2192                        size_estimate: None,
2193                    }],
2194                )
2195                .await
2196                .unwrap();
2197            let claimed = state.history().claim_shards(1).await.unwrap();
2198            assert_eq!(claimed.len(), 1);
2199
2200            resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2201
2202            // Poll until the parent run reaches a terminal state (the spawned
2203            // task runs the shard, finalizes it, then finalizes the parent).
2204            let mut status = RunStatus::Sharded;
2205            for _ in 0..100 {
2206                tokio::time::sleep(Duration::from_millis(50)).await;
2207                status = state.history().get("r").await.unwrap().unwrap().status;
2208                if status.is_terminal() {
2209                    break;
2210                }
2211            }
2212            assert_eq!(status, RunStatus::Completed, "shard ran → parent completed");
2213            assert!(output.exists(), "shard wrote its output");
2214        }
2215
2216        #[tokio::test]
2217        async fn resume_claimed_shard_with_no_config_fails_the_shard() {
2218            // A claimed shard whose parent run has no config_body can't run →
2219            // the shard is finalized failed and the parent run fails.
2220            let dir = tempfile::tempdir().unwrap();
2221            let state = sqlite_state(dir.path()).await;
2222            let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2223            rec.status = RunStatus::Sharded; // config_body intentionally None
2224            state.history().upsert(&rec).await.unwrap();
2225            use crate::serve::history::ShardInsert;
2226            state
2227                .history()
2228                .insert_shards(
2229                    "r",
2230                    &[ShardInsert {
2231                        shard_id: "0".into(),
2232                        descriptor: serde_json::Value::Null,
2233                        size_estimate: None,
2234                    }],
2235                )
2236                .await
2237                .unwrap();
2238            let claimed = state.history().claim_shards(1).await.unwrap();
2239            resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2240
2241            let mut status = RunStatus::Sharded;
2242            for _ in 0..100 {
2243                tokio::time::sleep(Duration::from_millis(50)).await;
2244                status = state.history().get("r").await.unwrap().unwrap().status;
2245                if status.is_terminal() {
2246                    break;
2247                }
2248            }
2249            assert_eq!(status, RunStatus::Failed, "no-config shard → parent failed");
2250        }
2251
2252        #[tokio::test]
2253        async fn coordinate_returns_err_when_source_build_fails() {
2254            // A shardable-looking source with an invalid config fails to build →
2255            // coordinate_sharded_run surfaces the error (the caller fails the run).
2256            let dir = tempfile::tempdir().unwrap();
2257            let state = sqlite_state(dir.path()).await;
2258            // s3 missing the required `file_format` field → build_source errors.
2259            let l = loaded(
2260                "version: 1\npipeline:\n  \
2261                 source: { type: s3, config: { bucket: b } }\n  \
2262                 sink: { type: stdout, config: {} }\n",
2263            )
2264            .await;
2265            assert!(coordinate_sharded_run(&state, "r", &l, 4).await.is_err());
2266        }
2267
2268        #[tokio::test]
2269        async fn execute_shard_returns_false_on_malformed_resilience() {
2270            // A resilience block that parses but fails to compile makes
2271            // execute_shard fail fast (false) before running the pipeline.
2272            let dir = tempfile::tempdir().unwrap();
2273            let state = sqlite_state(dir.path()).await;
2274            let input = dir.path().join("in.csv");
2275            std::fs::write(&input, "id\n1\n").unwrap();
2276            let yaml = format!(
2277                "version: 1\nresilience:\n  retry:\n    max_attempts: 0\npipeline:\n  \
2278                 source: {{ type: csv, config: {{ path: \"{}\" }} }}\n  \
2279                 sink: {{ type: stdout, config: {{}} }}\n",
2280                input.display()
2281            );
2282            let l = loaded(&yaml).await;
2283            let ok = execute_shard(
2284                &state,
2285                l,
2286                "r",
2287                "0",
2288                ShardSpec::whole(),
2289                CancellationToken::new(),
2290                None,
2291                None,
2292                Utc::now(),
2293            )
2294            .await;
2295            assert!(!ok, "malformed resilience → shard fails fast");
2296        }
2297
2298        #[tokio::test]
2299        async fn resume_claimed_shard_with_unloadable_config_fails_the_shard() {
2300            // A claimed shard whose run config fails to re-load (malformed body)
2301            // exercises resume_claimed_shard's load-error branch → shard failed.
2302            let dir = tempfile::tempdir().unwrap();
2303            let state = sqlite_state(dir.path()).await;
2304            let mut rec = RunRecord::queued("r".into(), None, BTreeMap::new(), None, Utc::now());
2305            rec.status = RunStatus::Sharded;
2306            rec.config_body = Some("this: is: not: valid: yaml: [".into());
2307            state.history().upsert(&rec).await.unwrap();
2308            use crate::serve::history::ShardInsert;
2309            state
2310                .history()
2311                .insert_shards(
2312                    "r",
2313                    &[ShardInsert {
2314                        shard_id: "0".into(),
2315                        descriptor: serde_json::Value::Null,
2316                        size_estimate: None,
2317                    }],
2318                )
2319                .await
2320                .unwrap();
2321            let claimed = state.history().claim_shards(1).await.unwrap();
2322            resume_claimed_shard(state.clone(), claimed.into_iter().next().unwrap());
2323
2324            let mut status = RunStatus::Sharded;
2325            for _ in 0..100 {
2326                tokio::time::sleep(Duration::from_millis(50)).await;
2327                status = state.history().get("r").await.unwrap().unwrap().status;
2328                if status.is_terminal() {
2329                    break;
2330                }
2331            }
2332            assert_eq!(
2333                status,
2334                RunStatus::Failed,
2335                "unloadable config → parent failed"
2336            );
2337        }
2338
2339        // ── F10: cross-instance cancel of a Sharded run ──────────────────────
2340
2341        #[tokio::test]
2342        async fn request_cancel_flags_a_sharded_parent() {
2343            // F10: a Sharded parent (status 'sharded') must accept request_cancel
2344            // (the guard was broadened from 'running' to ('running','sharded')).
2345            let dir = tempfile::tempdir().unwrap();
2346            let state = sqlite_state(dir.path()).await;
2347            seed_run(&state, "r", RunStatus::Sharded).await;
2348            // request_cancel is fire-and-forget; verify it took effect by reading
2349            // back via pending_shard_cancellations after a shard is claimed below.
2350            state.history().request_cancel("r").await.unwrap();
2351            // Insert + claim a shard so this instance owns a running shard under r.
2352            use crate::serve::history::ShardInsert;
2353            state
2354                .history()
2355                .insert_shards(
2356                    "r",
2357                    &[ShardInsert {
2358                        shard_id: "0".into(),
2359                        descriptor: serde_json::Value::Null,
2360                        size_estimate: None,
2361                    }],
2362                )
2363                .await
2364                .unwrap();
2365            let claimed = state.history().claim_shards(1).await.unwrap();
2366            assert_eq!(claimed.len(), 1, "shard claimed (running, owned)");
2367
2368            let flagged = state.history().pending_shard_cancellations().await.unwrap();
2369            assert_eq!(
2370                flagged,
2371                vec!["r".to_string()],
2372                "the flagged sharded parent's run id is returned for its running shard"
2373            );
2374        }
2375
2376        #[tokio::test]
2377        async fn pending_shard_cancellations_filters_unflagged_and_pending_shards() {
2378            let dir = tempfile::tempdir().unwrap();
2379            let state = sqlite_state(dir.path()).await;
2380            use crate::serve::history::ShardInsert;
2381            let one = |id: &str| {
2382                vec![ShardInsert {
2383                    shard_id: id.into(),
2384                    descriptor: serde_json::Value::Null,
2385                    size_estimate: None,
2386                }]
2387            };
2388
2389            // Runs A (flagged) and C (NOT flagged) each get one shard; claim both
2390            // so they are running+owned here.
2391            seed_run(&state, "A", RunStatus::Sharded).await;
2392            state.history().request_cancel("A").await.unwrap();
2393            state.history().insert_shards("A", &one("0")).await.unwrap();
2394            seed_run(&state, "C", RunStatus::Sharded).await;
2395            state.history().insert_shards("C", &one("0")).await.unwrap();
2396            let claimed = state.history().claim_shards(8).await.unwrap();
2397            assert_eq!(claimed.len(), 2, "A and C shards claimed (running)");
2398
2399            // Run B: flagged, but its shard stays PENDING (inserted after the
2400            // claim, never claimed) → must be excluded (the join requires a
2401            // 'running' shard owned by this instance).
2402            seed_run(&state, "B", RunStatus::Sharded).await;
2403            state.history().request_cancel("B").await.unwrap();
2404            state.history().insert_shards("B", &one("0")).await.unwrap();
2405
2406            let flagged = state.history().pending_shard_cancellations().await.unwrap();
2407            assert_eq!(
2408                flagged,
2409                vec!["A".to_string()],
2410                "only A (flagged + a running owned shard); B pending-shard, C unflagged"
2411            );
2412        }
2413
2414        // ── F11: orphaned-Sharded-parent sweep ───────────────────────────────
2415
2416        #[tokio::test]
2417        async fn finalize_sweep_completes_an_all_success_sharded_parent() {
2418            let dir = tempfile::tempdir().unwrap();
2419            let state = sqlite_state(dir.path()).await;
2420            seed_sharded_with_shards(&state, "r", 3).await;
2421            for i in 0..3 {
2422                state
2423                    .history()
2424                    .finalize_shard("r", &i.to_string(), true)
2425                    .await
2426                    .unwrap();
2427            }
2428            // No inline maybe_finalize_parent — the sweep must finalize it.
2429            let n = state
2430                .history()
2431                .finalize_completed_sharded_parents()
2432                .await
2433                .unwrap();
2434            assert_eq!(n, 1, "one sharded parent finalized");
2435            let rec = state.history().get("r").await.unwrap().unwrap();
2436            assert_eq!(rec.status, RunStatus::Completed);
2437            assert!(rec.finished_at.is_some());
2438            assert!(rec.error.is_none());
2439
2440            // Idempotent: a second sweep finalizes nothing.
2441            assert_eq!(
2442                state
2443                    .history()
2444                    .finalize_completed_sharded_parents()
2445                    .await
2446                    .unwrap(),
2447                0,
2448                "already-terminal parent is not re-finalized"
2449            );
2450        }
2451
2452        #[tokio::test]
2453        async fn finalize_sweep_fails_a_parent_with_a_failed_shard() {
2454            let dir = tempfile::tempdir().unwrap();
2455            let state = sqlite_state(dir.path()).await;
2456            seed_sharded_with_shards(&state, "r", 3).await;
2457            state
2458                .history()
2459                .finalize_shard("r", "0", true)
2460                .await
2461                .unwrap();
2462            state
2463                .history()
2464                .finalize_shard("r", "1", false)
2465                .await
2466                .unwrap();
2467            state
2468                .history()
2469                .finalize_shard("r", "2", true)
2470                .await
2471                .unwrap();
2472            let n = state
2473                .history()
2474                .finalize_completed_sharded_parents()
2475                .await
2476                .unwrap();
2477            assert_eq!(n, 1);
2478            let rec = state.history().get("r").await.unwrap().unwrap();
2479            assert_eq!(rec.status, RunStatus::Failed);
2480            assert!(rec.finished_at.is_some());
2481            assert_eq!(rec.error.as_deref(), Some("1/3 shard(s) failed"));
2482        }
2483
2484        #[tokio::test]
2485        async fn finalize_sweep_leaves_a_not_all_terminal_parent_sharded() {
2486            let dir = tempfile::tempdir().unwrap();
2487            let state = sqlite_state(dir.path()).await;
2488            seed_sharded_with_shards(&state, "r", 2).await;
2489            // Only one shard terminal → the parent stays sharded.
2490            state
2491                .history()
2492                .finalize_shard("r", "0", true)
2493                .await
2494                .unwrap();
2495            let n = state
2496                .history()
2497                .finalize_completed_sharded_parents()
2498                .await
2499                .unwrap();
2500            assert_eq!(n, 0, "parent with a still-running shard is not finalized");
2501            assert_eq!(
2502                state.history().get("r").await.unwrap().unwrap().status,
2503                RunStatus::Sharded
2504            );
2505        }
2506    }
2507}