use bytes::Bytes;
use klieo_core::{BusError, KvStore};
use serde::{Deserialize, Serialize};
use serde_json::Value;
use std::sync::Arc;
pub const RUNS_BUCKET: &str = "workflow_runs";
const IDEM_PREFIX: &str = "idem-";
const RUN_PREFIX: &str = "run-";
const ACTIVE_PREFIX: &str = "active-";
#[non_exhaustive]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum RunStatus {
Pending,
Running,
Succeeded,
Failed,
Aborted,
}
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct RunRecordView {
pub run_id: String,
pub workflow_id: String,
pub status: RunStatus,
#[serde(skip_serializing_if = "Option::is_none")]
pub result: Option<Value>,
#[serde(skip_serializing_if = "Option::is_none")]
pub error: Option<String>,
pub author: String,
pub created_at: String,
}
impl RunRecordView {
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(),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ClaimOutcome {
Created,
AlreadyExists(String),
}
#[non_exhaustive]
#[derive(Debug, thiserror::Error)]
pub enum StoreError {
#[error("kv store error: {0}")]
Kv(#[from] BusError),
#[error("run record serialisation error: {0}")]
Serde(#[from] serde_json::Error),
#[error("run record not found: {0}")]
NotFound(String),
}
#[derive(Clone)]
pub struct RunStore {
kv: Arc<dyn KvStore>,
}
impl RunStore {
pub fn new(kv: Arc<dyn KvStore>) -> Self {
Self { kv }
}
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)),
}
}
pub async fn create(&self, view: &RunRecordView) -> Result<(), StoreError> {
self.put_active_marker(&view.run_id).await?;
self.put_view(view).await
}
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
}
pub async fn set_running(&self, run_id: &str) -> Result<(), StoreError> {
self.update(run_id, |v| v.status = RunStatus::Running).await
}
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
}
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
}
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
}
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),
}
}
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)
}
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(_)));
}
#[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!()
}
}
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());
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");
}
#[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());
}
}