klieo_workflow_api/

run_store.rs

//! Non-authoritative run-status cache over a [`KvStore`].
//!
//! This is a volatile poll cache — the authoritative, durable run outcome
//! lives in the provenance chain (wired in a later slice). Idempotency
//! claims use compare-and-set (`expected = None`) so concurrent identical
//! submits collapse to one run without a get-then-put race (CWE-367).

use bytes::Bytes;
use klieo_core::{BusError, KvStore};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;

/// KV bucket holding both idempotency claims and run records.
pub const RUNS_BUCKET: &str = "workflow_runs";

const IDEM_PREFIX: &str = "idem-";
const RUN_PREFIX: &str = "run-";
/// Marks a run as non-terminal (`Pending`/`Running`). Written by `create`,
/// deleted by every terminal transition and by `rollback_claim`, so
/// [`RunStore::list_active`] tracks the live-run count instead of a
/// backend's total history.
const ACTIVE_PREFIX: &str = "active-";

/// Lifecycle status of a submitted run.
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
    /// Accepted, not yet started.
    Pending,
    /// Executor is running the flow.
    Running,
    /// Flow completed successfully.
    Succeeded,
    /// Flow returned an error, panicked, or its record could not be finalised.
    Failed,
    /// Flow exceeded the per-run timeout and was cancelled.
    Aborted,
}

/// The client-visible snapshot of a run returned by `GET /runs/{id}`.
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RunRecordView {
    /// Server-minted run id.
    pub run_id: String,
    /// The workflow this run executes; the run's provenance `resource_ref`.
    pub workflow_id: String,
    /// Current lifecycle status.
    pub status: RunStatus,
    /// Output envelope, present once `Succeeded`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub result: Option<Value>,
    /// Failure summary, present once `Failed`/`Aborted`.
    #[serde(skip_serializing_if = "Option::is_none")]
    pub error: Option<String>,
    /// Server-stamped author (authenticated caller), never client-supplied.
    pub author: String,
    /// RFC-3339 submit time, stamped at the handler boundary.
    pub created_at: String,
}

impl RunRecordView {
    /// A freshly-accepted run, before the executor starts it.
    pub fn pending(
        run_id: impl Into<String>,
        workflow_id: impl Into<String>,
        author: impl Into<String>,
        created_at: impl Into<String>,
    ) -> Self {
        Self {
            run_id: run_id.into(),
            workflow_id: workflow_id.into(),
            status: RunStatus::Pending,
            result: None,
            error: None,
            author: author.into(),
            created_at: created_at.into(),
        }
    }
}

/// Outcome of a CAS idempotency claim.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClaimOutcome {
    /// This caller won the claim; proceed to run.
    Created,
    /// The key was already claimed; the winning run's id.
    AlreadyExists(String),
}

/// Failure modes of the run store.
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum StoreError {
    /// The underlying KV store failed.
    #[error("kv store error: {0}")]
    Kv(#[from] BusError),
    /// A stored record could not be (de)serialised.
    #[error("run record serialisation error: {0}")]
    Serde(#[from] serde_json::Error),
    /// A transition targeted a run id with no record.
    #[error("run record not found: {0}")]
    NotFound(String),
}

/// Volatile run-status cache keyed by run id, with CAS-based idempotency.
#[derive(Clone)]
pub struct RunStore {
    kv: Arc<dyn KvStore>,
}

impl RunStore {
    /// Wrap a KV store as a run-status cache.
    pub fn new(kv: Arc<dyn KvStore>) -> Self {
        Self { kv }
    }

    /// Atomically claim `idempotency_key` for `run_id` via `cas(expected =
    /// None)`. The first caller gets [`ClaimOutcome::Created`]; a concurrent
    /// or later identical submit reads the winner's id
    /// ([`ClaimOutcome::AlreadyExists`]).
    pub async fn claim_idempotent(
        &self,
        idempotency_key: &str,
        run_id: &str,
    ) -> Result<ClaimOutcome, StoreError> {
        let key = format!("{IDEM_PREFIX}{idempotency_key}");
        let value = Bytes::from(run_id.to_string());
        match self.kv.cas(RUNS_BUCKET, &key, value, None).await {
            Ok(_) => Ok(ClaimOutcome::Created),
            Err(BusError::CasConflict { .. }) => {
                let winner = self
                    .kv
                    .get(RUNS_BUCKET, &key)
                    .await?
                    .map(|e| String::from_utf8_lossy(&e.value).into_owned())
                    .ok_or_else(|| StoreError::NotFound(key.clone()))?;
                Ok(ClaimOutcome::AlreadyExists(winner))
            }
            Err(other) => Err(StoreError::Kv(other)),
        }
    }

    /// Persist a fresh run record and mark it active.
    ///
    /// Writes the active marker *before* the run record. If the process
    /// dies between the two writes, a stray marker with no record is left
    /// behind — `Self::list_active` tolerates that (skips it). The
    /// reverse order would risk the opposite, worse failure: a Pending
    /// record with no marker, which the sweep would silently never see.
    pub async fn create(&self, view: &RunRecordView) -> Result<(), StoreError> {
        self.put_active_marker(&view.run_id).await?;
        self.put_view(view).await
    }

    /// Undo a claim + its pending record + active marker after a post-claim
    /// failure (e.g. the run-start audit append failed), so a genuine retry
    /// with the same key re-runs instead of resolving to a dead run.
    /// Best-effort — a delete failure leaves the sweep as the backstop.
    pub async fn rollback_claim(
        &self,
        idempotency_key: &str,
        run_id: &str,
    ) -> Result<(), StoreError> {
        self.kv
            .delete(RUNS_BUCKET, &format!("{IDEM_PREFIX}{idempotency_key}"))
            .await?;
        self.kv
            .delete(RUNS_BUCKET, &format!("{RUN_PREFIX}{run_id}"))
            .await?;
        self.delete_active_marker(run_id).await
    }

    /// Mark the run `Running`.
    pub async fn set_running(&self, run_id: &str) -> Result<(), StoreError> {
        self.update(run_id, |v| v.status = RunStatus::Running).await
    }

    /// Mark the run `Succeeded` with its output envelope, then retire its
    /// active marker.
    pub async fn set_succeeded(&self, run_id: &str, result: Value) -> Result<(), StoreError> {
        self.update(run_id, move |v| {
            v.status = RunStatus::Succeeded;
            v.result = Some(result);
        })
        .await?;
        self.delete_active_marker(run_id).await
    }

    /// Mark the run `Failed` with a summary, then retire its active marker.
    pub async fn set_failed(&self, run_id: &str, error: String) -> Result<(), StoreError> {
        self.update(run_id, move |v| {
            v.status = RunStatus::Failed;
            v.error = Some(error);
        })
        .await?;
        self.delete_active_marker(run_id).await
    }

    /// Mark the run `Aborted` (timed out) with a summary, then retire its
    /// active marker.
    pub async fn set_aborted(&self, run_id: &str, error: String) -> Result<(), StoreError> {
        self.update(run_id, move |v| {
            v.status = RunStatus::Aborted;
            v.error = Some(error);
        })
        .await?;
        self.delete_active_marker(run_id).await
    }

    /// Read a run record, or `None` if unknown.
    pub async fn get(&self, run_id: &str) -> Result<Option<RunRecordView>, StoreError> {
        let key = format!("{RUN_PREFIX}{run_id}");
        match self.kv.get(RUNS_BUCKET, &key).await? {
            Some(entry) => Ok(Some(serde_json::from_slice(&entry.value)?)),
            None => Ok(None),
        }
    }

    /// Every non-terminal (`Pending`/`Running`) run, read through the
    /// active-run marker index. Used by the startup orphan sweep. Terminal
    /// records are never removed but their markers are, so the per-record
    /// fetches are bounded to the live-run count rather than total history.
    /// (A backend without prefix-filtered key enumeration — e.g. NATS — still
    /// lists all keys once, but skips fetching the terminal records.)
    ///
    /// On a backend that cannot enumerate keys this returns an empty list
    /// rather than erroring (an ephemeral store has nothing durable to
    /// sweep). A marker with no matching run record — the crash window
    /// documented on [`Self::create`] — is tolerated: skipped rather than
    /// surfaced as an error.
    ///
    /// Crate-internal: the host drives the sweep via [`crate::sweep_orphaned_runs`],
    /// never this method directly, so it is not part of the public API.
    pub(crate) async fn list_active(&self) -> Result<Vec<RunRecordView>, StoreError> {
        let keys = match self.kv.keys(RUNS_BUCKET).await {
            Ok(keys) => keys,
            Err(BusError::Unsupported(_)) => return Ok(Vec::new()),
            Err(other) => return Err(StoreError::Kv(other)),
        };
        let mut runs = Vec::new();
        for key in keys.iter().filter(|k| k.starts_with(ACTIVE_PREFIX)) {
            let run_id = &key[ACTIVE_PREFIX.len()..];
            if let Some(view) = self.get(run_id).await? {
                runs.push(view);
            }
        }
        Ok(runs)
    }

    /// Read-modify-write a run record. Sound because each run has a single
    /// writer: the CAS idempotency claim admits exactly one executor per
    /// `run_id`, and that executor's writes (`create` → `set_running` →
    /// terminal) are sequential. The startup `sweep_orphaned_runs` is a
    /// documented second writer, but it runs only at startup before serving
    /// traffic, so it never overlaps a live executor. A future concurrent
    /// writer on a live run (e.g. a cancel/retry endpoint) would need CAS
    /// here to avoid a lost update.
    async fn update(
        &self,
        run_id: &str,
        mutate: impl FnOnce(&mut RunRecordView),
    ) -> Result<(), StoreError> {
        let mut view = self
            .get(run_id)
            .await?
            .ok_or_else(|| StoreError::NotFound(run_id.to_string()))?;
        mutate(&mut view);
        self.put_view(&view).await
    }

    async fn put_view(&self, view: &RunRecordView) -> Result<(), StoreError> {
        let key = format!("{RUN_PREFIX}{}", view.run_id);
        let bytes = Bytes::from(serde_json::to_vec(view)?);
        self.kv.put(RUNS_BUCKET, &key, bytes).await?;
        Ok(())
    }

    async fn put_active_marker(&self, run_id: &str) -> Result<(), StoreError> {
        let key = format!("{ACTIVE_PREFIX}{run_id}");
        self.kv.put(RUNS_BUCKET, &key, Bytes::new()).await?;
        Ok(())
    }

    async fn delete_active_marker(&self, run_id: &str) -> Result<(), StoreError> {
        self.kv
            .delete(RUNS_BUCKET, &format!("{ACTIVE_PREFIX}{run_id}"))
            .await?;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use klieo_bus_memory::MemoryBus;

    fn store() -> RunStore {
        RunStore::new(MemoryBus::new().kv)
    }

    #[tokio::test]
    async fn claim_returns_created_once_then_already_exists() {
        let store = store();
        assert_eq!(
            store.claim_idempotent("k", "run-1").await.unwrap(),
            ClaimOutcome::Created
        );
        assert_eq!(
            store.claim_idempotent("k", "run-2").await.unwrap(),
            ClaimOutcome::AlreadyExists("run-1".to_string()),
        );
    }

    #[tokio::test]
    async fn status_transitions_persist() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "run-1",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        store.set_running("run-1").await.unwrap();
        assert_eq!(
            store.get("run-1").await.unwrap().unwrap().status,
            RunStatus::Running
        );
        store
            .set_succeeded("run-1", serde_json::json!({"ok": true}))
            .await
            .unwrap();
        let view = store.get("run-1").await.unwrap().unwrap();
        assert_eq!(view.status, RunStatus::Succeeded);
        assert_eq!(view.result, Some(serde_json::json!({"ok": true})));
        assert_eq!(view.author, "alice");
    }

    #[tokio::test]
    async fn missing_run_reads_none() {
        assert!(store().get("nope").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn transition_on_missing_run_errors() {
        let err = store().set_running("ghost").await.unwrap_err();
        assert!(matches!(err, StoreError::NotFound(_)));
    }

    /// A backend that cannot enumerate keys (the trait's default `keys()`
    /// returns `Unsupported`) makes `list_active()` a no-op empty rather than
    /// an error, so the startup sweep skips an ephemeral store cleanly.
    #[tokio::test]
    async fn list_active_maps_unsupported_keys_to_empty() {
        use async_trait::async_trait;
        use klieo_core::bus::{KvEntry, Lease, Revision};
        use std::time::Duration;

        struct NoKeysKv;
        #[async_trait]
        impl KvStore for NoKeysKv {
            async fn get(&self, _: &str, _: &str) -> Result<Option<KvEntry>, BusError> {
                unreachable!("list_active() must not read values when keys() is unsupported")
            }
            async fn put(&self, _: &str, _: &str, _: Bytes) -> Result<Revision, BusError> {
                unreachable!()
            }
            async fn cas(
                &self,
                _: &str,
                _: &str,
                _: Bytes,
                _: Option<Revision>,
            ) -> Result<Revision, BusError> {
                unreachable!()
            }
            async fn delete(&self, _: &str, _: &str) -> Result<(), BusError> {
                unreachable!()
            }
            async fn lease(&self, _: &str, _: &str, _: Duration) -> Result<Lease, BusError> {
                unreachable!()
            }
            // keys() uses the trait default → BusError::Unsupported.
        }

        let store = RunStore::new(std::sync::Arc::new(NoKeysKv));
        assert!(store.list_active().await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn create_writes_active_marker() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "run-1",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        let active = store.list_active().await.unwrap();
        assert_eq!(active.len(), 1);
        assert_eq!(active[0].run_id, "run-1");
    }

    #[tokio::test]
    async fn set_succeeded_deletes_active_marker() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "run-1",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        store
            .set_succeeded("run-1", serde_json::json!({"ok": true}))
            .await
            .unwrap();
        assert!(store.list_active().await.unwrap().is_empty());
        // The record itself survives — only the active marker is retired.
        assert_eq!(
            store.get("run-1").await.unwrap().unwrap().status,
            RunStatus::Succeeded
        );
    }

    #[tokio::test]
    async fn set_failed_deletes_active_marker() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "run-1",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        store.set_failed("run-1", "boom".into()).await.unwrap();
        assert!(store.list_active().await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn set_aborted_deletes_active_marker() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "run-1",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        store
            .set_aborted("run-1", "timed out".into())
            .await
            .unwrap();
        assert!(store.list_active().await.unwrap().is_empty());
    }

    #[tokio::test]
    async fn rollback_claim_deletes_active_marker() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "run-1",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        store.rollback_claim("idem-key", "run-1").await.unwrap();
        assert!(store.list_active().await.unwrap().is_empty());
        assert!(store.get("run-1").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn list_active_excludes_terminal_runs() {
        let store = store();
        store
            .create(&RunRecordView::pending(
                "done",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();
        store
            .set_succeeded("done", serde_json::json!({}))
            .await
            .unwrap();
        store
            .create(&RunRecordView::pending(
                "live",
                "wf-1",
                "alice",
                "2026-07-13T00:00:00Z",
            ))
            .await
            .unwrap();

        let active = store.list_active().await.unwrap();
        assert_eq!(active.len(), 1);
        assert_eq!(active[0].run_id, "live");
    }

    /// The crash window documented on `create`: a marker written with no
    /// matching run record is skipped, not surfaced as an error.
    #[tokio::test]
    async fn list_active_tolerates_marker_without_record() {
        let store = store();
        store.put_active_marker("ghost").await.unwrap();
        assert!(store.list_active().await.unwrap().is_empty());
    }
}