agentic_workflow/resilience/
rollback.rs

1use std::collections::HashMap;
2
3use chrono::Utc;
4
5use crate::types::{
6    RollbackAction, RollbackReceipt, RollbackScope, RollbackStepResult,
7    WorkflowError, WorkflowResult,
8};
9
10/// Rollback architecture — per-step undo with verification.
11pub struct RollbackEngine {
12    actions: HashMap<String, RollbackAction>,
13    receipts: Vec<RollbackReceipt>,
14}
15
16impl RollbackEngine {
17    pub fn new() -> Self {
18        Self {
19            actions: HashMap::new(),
20            receipts: Vec::new(),
21        }
22    }
23
24    /// Define a rollback action for a step.
25    pub fn define_action(&mut self, action: RollbackAction) -> WorkflowResult<()> {
26        self.actions.insert(action.step_id.clone(), action);
27        Ok(())
28    }
29
30    /// Get rollback action for a step.
31    pub fn get_action(&self, step_id: &str) -> Option<&RollbackAction> {
32        self.actions.get(step_id)
33    }
34
35    /// Preview what a rollback would do.
36    pub fn preview(
37        &self,
38        scope: &RollbackScope,
39        completed_step_ids: &[String],
40    ) -> Vec<String> {
41        match scope {
42            RollbackScope::Full => completed_step_ids
43                .iter()
44                .rev()
45                .filter(|id| self.actions.contains_key(id.as_str()))
46                .cloned()
47                .collect(),
48            RollbackScope::FromStep { step_id } => {
49                let start_idx = completed_step_ids
50                    .iter()
51                    .position(|id| id == step_id)
52                    .unwrap_or(0);
53                completed_step_ids[start_idx..]
54                    .iter()
55                    .rev()
56                    .filter(|id| self.actions.contains_key(id.as_str()))
57                    .cloned()
58                    .collect()
59            }
60            RollbackScope::Selective { step_ids } => step_ids
61                .iter()
62                .rev()
63                .filter(|id| self.actions.contains_key(id.as_str()))
64                .cloned()
65                .collect(),
66        }
67    }
68
69    /// Execute rollback and produce a receipt.
70    pub fn execute_rollback(
71        &mut self,
72        execution_id: &str,
73        scope: RollbackScope,
74        steps_to_rollback: &[String],
75    ) -> WorkflowResult<RollbackReceipt> {
76        let started_at = Utc::now();
77        let mut results = Vec::new();
78        let mut overall_success = true;
79
80        for step_id in steps_to_rollback {
81            match self.actions.get(step_id) {
82                Some(action) => {
83                    // In a real implementation, this would execute the rollback action
84                    let result = RollbackStepResult {
85                        step_id: step_id.clone(),
86                        success: true,
87                        error: None,
88                        verification_passed: action.verification.as_ref().map(|_| true),
89                    };
90                    results.push(result);
91                }
92                None => {
93                    results.push(RollbackStepResult {
94                        step_id: step_id.clone(),
95                        success: false,
96                        error: Some("No rollback action defined".to_string()),
97                        verification_passed: None,
98                    });
99                    overall_success = false;
100                }
101            }
102        }
103
104        let receipt = RollbackReceipt {
105            execution_id: execution_id.to_string(),
106            scope,
107            rolled_back_steps: results,
108            started_at,
109            completed_at: Utc::now(),
110            overall_success,
111        };
112
113        self.receipts.push(receipt.clone());
114        Ok(receipt)
115    }
116
117    /// Get rollback receipts for an execution.
118    pub fn get_receipts(&self, execution_id: &str) -> Vec<&RollbackReceipt> {
119        self.receipts
120            .iter()
121            .filter(|r| r.execution_id == execution_id)
122            .collect()
123    }
124
125    /// List all defined rollback actions.
126    pub fn list_actions(&self) -> Vec<&RollbackAction> {
127        self.actions.values().collect()
128    }
129}
130
131impl Default for RollbackEngine {
132    fn default() -> Self {
133        Self::new()
134    }
135}
136
137#[cfg(test)]
138mod tests {
139    use super::*;
140    use crate::types::RollbackType;
141
142    #[test]
143    fn test_rollback_preview() {
144        let mut engine = RollbackEngine::new();
145
146        engine
147            .define_action(RollbackAction {
148                id: "r1".into(),
149                step_id: "step-1".into(),
150                action_type: RollbackType::NotPossible {
151                    reason: "Email already sent".into(),
152                },
153                description: "Cannot undo email".into(),
154                verification: None,
155            })
156            .unwrap();
157
158        engine
159            .define_action(RollbackAction {
160                id: "r2".into(),
161                step_id: "step-2".into(),
162                action_type: RollbackType::Command {
163                    command: "undo.sh".into(),
164                    args: vec![],
165                },
166                description: "Undo step 2".into(),
167                verification: None,
168            })
169            .unwrap();
170
171        let completed = vec!["step-1".to_string(), "step-2".to_string()];
172        let preview = engine.preview(&RollbackScope::Full, &completed);
173        assert_eq!(preview.len(), 2);
174        assert_eq!(preview[0], "step-2"); // Reverse order
175    }
176}
agentic_workflow/resilience/rollback.rs

agentic_workflow/resilience/
rollback.rs
