use crate::error::{Result, WorkflowError};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;
use std::time::{SystemTime, UNIX_EPOCH};
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
pub struct WorkflowCheckpoint {
pub workflow_id: String,
pub completed_steps: Vec<String>,
pub step_outputs: HashMap<String, serde_json::Value>,
pub created_at: u64,
pub workflow_version: u32,
}
impl WorkflowCheckpoint {
#[must_use]
pub fn new(workflow_id: impl Into<String>, workflow_version: u32) -> Self {
let created_at = SystemTime::now()
.duration_since(UNIX_EPOCH)
.map(|d| d.as_secs())
.unwrap_or(0);
Self {
workflow_id: workflow_id.into(),
completed_steps: Vec::new(),
step_outputs: HashMap::new(),
created_at,
workflow_version,
}
}
pub fn mark_step_completed(
&mut self,
step_name: impl Into<String>,
output: Option<serde_json::Value>,
) {
let name = step_name.into();
if !self.completed_steps.contains(&name) {
self.completed_steps.push(name.clone());
}
if let Some(v) = output {
self.step_outputs.insert(name, v);
}
}
#[must_use]
pub fn is_step_completed(&self, step_name: &str) -> bool {
self.completed_steps.iter().any(|s| s == step_name)
}
#[must_use]
pub fn step_output(&self, step_name: &str) -> Option<&serde_json::Value> {
self.step_outputs.get(step_name)
}
#[must_use]
pub fn completed_count(&self) -> usize {
self.completed_steps.len()
}
}
#[derive(Debug, Clone)]
pub struct CheckpointManager {
storage_dir: PathBuf,
}
impl CheckpointManager {
#[must_use]
pub fn new(storage_dir: PathBuf) -> Self {
Self { storage_dir }
}
#[must_use]
pub fn storage_dir(&self) -> &PathBuf {
&self.storage_dir
}
fn checkpoint_path(&self, workflow_id: &str) -> PathBuf {
self.storage_dir
.join(format!("{}.checkpoint.json", workflow_id))
}
fn ensure_dir(&self) -> Result<()> {
std::fs::create_dir_all(&self.storage_dir).map_err(WorkflowError::Io)
}
pub fn save(&self, checkpoint: &WorkflowCheckpoint) -> Result<()> {
self.ensure_dir()?;
let path = self.checkpoint_path(&checkpoint.workflow_id);
let json = serde_json::to_vec_pretty(checkpoint)?;
std::fs::write(&path, &json).map_err(WorkflowError::Io)
}
pub fn load(&self, workflow_id: &str) -> Result<WorkflowCheckpoint> {
let path = self.checkpoint_path(workflow_id);
if !path.exists() {
return Err(WorkflowError::FileNotFound(path));
}
let bytes = std::fs::read(&path).map_err(WorkflowError::Io)?;
let cp: WorkflowCheckpoint = serde_json::from_slice(&bytes)?;
Ok(cp)
}
pub fn delete(&self, workflow_id: &str) -> Result<()> {
let path = self.checkpoint_path(workflow_id);
if path.exists() {
std::fs::remove_file(&path).map_err(WorkflowError::Io)?;
}
Ok(())
}
#[must_use]
pub fn exists(&self, workflow_id: &str) -> bool {
self.checkpoint_path(workflow_id).exists()
}
pub fn list_checkpoints(&self) -> Result<Vec<String>> {
if !self.storage_dir.exists() {
return Ok(Vec::new());
}
let mut ids = Vec::new();
let entries = std::fs::read_dir(&self.storage_dir).map_err(WorkflowError::Io)?;
for entry in entries {
let entry = entry.map_err(WorkflowError::Io)?;
let file_name = entry.file_name();
let name = file_name.to_string_lossy();
if let Some(id) = name.strip_suffix(".checkpoint.json") {
if self.exists(id) {
ids.push(id.to_string());
}
}
}
ids.sort();
Ok(ids)
}
}
#[must_use]
pub fn should_skip_step(checkpoint: &Option<WorkflowCheckpoint>, step_name: &str) -> bool {
checkpoint
.as_ref()
.map(|cp| cp.is_step_completed(step_name))
.unwrap_or(false)
}
#[cfg(test)]
mod tests {
use super::*;
fn temp_dir(suffix: &str) -> PathBuf {
let base = std::env::temp_dir().join(format!("oximedia_cp_test_{}", suffix));
let _ = std::fs::remove_dir_all(&base);
base
}
#[test]
fn test_checkpoint_new_fields() {
let cp = WorkflowCheckpoint::new("wf-999", 2);
assert_eq!(cp.workflow_id, "wf-999");
assert_eq!(cp.workflow_version, 2);
assert!(cp.completed_steps.is_empty());
assert!(cp.step_outputs.is_empty());
assert!(cp.created_at > 0);
}
#[test]
fn test_mark_step_completed_no_output() {
let mut cp = WorkflowCheckpoint::new("wf-001", 1);
cp.mark_step_completed("ingest", None);
assert!(cp.is_step_completed("ingest"));
assert!(cp.step_output("ingest").is_none());
assert_eq!(cp.completed_count(), 1);
}
#[test]
fn test_mark_step_completed_with_output() {
let mut cp = WorkflowCheckpoint::new("wf-001", 1);
let val = serde_json::json!({"frames": 1000});
cp.mark_step_completed("analyse", Some(val.clone()));
assert!(cp.is_step_completed("analyse"));
assert_eq!(cp.step_output("analyse"), Some(&val));
}
#[test]
fn test_mark_step_completed_dedup() {
let mut cp = WorkflowCheckpoint::new("wf-001", 1);
cp.mark_step_completed("transcode", None);
cp.mark_step_completed("transcode", None);
assert_eq!(cp.completed_count(), 1);
}
#[test]
fn test_is_step_completed_false() {
let cp = WorkflowCheckpoint::new("wf-001", 1);
assert!(!cp.is_step_completed("missing_step"));
}
#[test]
fn test_save_and_load() {
let dir = temp_dir("save_load");
let mgr = CheckpointManager::new(dir.clone());
let mut cp = WorkflowCheckpoint::new("wf-save", 1);
cp.mark_step_completed("step-a", Some(serde_json::json!({"ok": true})));
mgr.save(&cp).expect("save should succeed");
let loaded = mgr.load("wf-save").expect("load should succeed");
assert_eq!(loaded.workflow_id, "wf-save");
assert_eq!(loaded.completed_steps, vec!["step-a"]);
assert_eq!(
loaded.step_output("step-a"),
Some(&serde_json::json!({"ok": true}))
);
}
#[test]
fn test_load_nonexistent_returns_error() {
let dir = temp_dir("load_missing");
let mgr = CheckpointManager::new(dir);
let result = mgr.load("no-such-workflow");
assert!(
result.is_err(),
"loading a missing checkpoint must return Err"
);
}
#[test]
fn test_save_overwrites_previous() {
let dir = temp_dir("overwrite");
let mgr = CheckpointManager::new(dir);
let mut cp1 = WorkflowCheckpoint::new("wf-ow", 1);
cp1.mark_step_completed("step-a", None);
mgr.save(&cp1).expect("first save");
let mut cp2 = WorkflowCheckpoint::new("wf-ow", 1);
cp2.mark_step_completed("step-a", None);
cp2.mark_step_completed("step-b", None);
mgr.save(&cp2).expect("second save");
let loaded = mgr.load("wf-ow").expect("load after overwrite");
assert_eq!(loaded.completed_count(), 2);
}
#[test]
fn test_exists_true_after_save() {
let dir = temp_dir("exists_true");
let mgr = CheckpointManager::new(dir);
let cp = WorkflowCheckpoint::new("wf-ex", 1);
mgr.save(&cp).expect("save");
assert!(mgr.exists("wf-ex"));
}
#[test]
fn test_exists_false_before_save() {
let dir = temp_dir("exists_false");
let mgr = CheckpointManager::new(dir);
assert!(!mgr.exists("wf-not-saved"));
}
#[test]
fn test_exists_false_after_delete() {
let dir = temp_dir("exists_after_delete");
let mgr = CheckpointManager::new(dir);
let cp = WorkflowCheckpoint::new("wf-del-ex", 1);
mgr.save(&cp).expect("save");
mgr.delete("wf-del-ex").expect("delete");
assert!(!mgr.exists("wf-del-ex"));
}
#[test]
fn test_delete_existing() {
let dir = temp_dir("delete_existing");
let mgr = CheckpointManager::new(dir);
let cp = WorkflowCheckpoint::new("wf-del", 1);
mgr.save(&cp).expect("save");
assert!(mgr.exists("wf-del"));
mgr.delete("wf-del").expect("delete should succeed");
assert!(!mgr.exists("wf-del"));
}
#[test]
fn test_delete_nonexistent_ok() {
let dir = temp_dir("delete_nonexistent");
let mgr = CheckpointManager::new(dir);
assert!(mgr.delete("ghost-workflow").is_ok());
}
#[test]
fn test_list_empty_when_no_dir() {
let dir = temp_dir("list_nodir");
let mgr = CheckpointManager::new(dir);
let ids = mgr.list_checkpoints().expect("list should succeed");
assert!(ids.is_empty());
}
#[test]
fn test_list_checkpoints_after_multiple_saves() {
let dir = temp_dir("list_multiple");
let mgr = CheckpointManager::new(dir);
for id in &["wf-alpha", "wf-beta", "wf-gamma"] {
let cp = WorkflowCheckpoint::new(*id, 1);
mgr.save(&cp).expect("save");
}
let mut ids = mgr.list_checkpoints().expect("list");
ids.sort();
assert_eq!(ids, vec!["wf-alpha", "wf-beta", "wf-gamma"]);
}
#[test]
fn test_list_excludes_deleted() {
let dir = temp_dir("list_after_delete");
let mgr = CheckpointManager::new(dir);
mgr.save(&WorkflowCheckpoint::new("wf-keep", 1))
.expect("save keep");
mgr.save(&WorkflowCheckpoint::new("wf-gone", 1))
.expect("save gone");
mgr.delete("wf-gone").expect("delete gone");
let ids = mgr.list_checkpoints().expect("list");
assert!(ids.contains(&"wf-keep".to_string()));
assert!(!ids.contains(&"wf-gone".to_string()));
}
#[test]
fn test_should_skip_step_with_completed_step() {
let mut cp = WorkflowCheckpoint::new("wf-skip", 1);
cp.mark_step_completed("encode", None);
assert!(should_skip_step(&Some(cp), "encode"));
}
#[test]
fn test_should_skip_step_with_missing_step() {
let cp = WorkflowCheckpoint::new("wf-skip", 1);
assert!(!should_skip_step(&Some(cp), "not-yet-done"));
}
#[test]
fn test_should_skip_step_with_no_checkpoint() {
assert!(!should_skip_step(&None, "any-step"));
}
#[test]
fn test_should_skip_step_multiple_steps() {
let mut cp = WorkflowCheckpoint::new("wf-multi", 1);
cp.mark_step_completed("ingest", None);
cp.mark_step_completed("transcode", None);
assert!(should_skip_step(&Some(cp.clone()), "ingest"));
assert!(should_skip_step(&Some(cp.clone()), "transcode"));
assert!(!should_skip_step(&Some(cp), "upload"));
}
}