Skip to main content

agentic_workflow/engine/
scheduler.rs

1use std::collections::HashMap;
2
3use chrono::Utc;
4use uuid::Uuid;
5
6use crate::types::{
7    AdaptiveSchedule, ConflictPolicy, Schedule, ScheduleExpression,
8    WorkflowError, WorkflowResult,
9};
10
11/// Context-aware scheduling engine.
12pub struct SchedulerEngine {
13    schedules: HashMap<String, Schedule>,
14}
15
16impl SchedulerEngine {
17    pub fn new() -> Self {
18        Self {
19            schedules: HashMap::new(),
20        }
21    }
22
23    /// Create a new schedule for a workflow.
24    pub fn create_schedule(
25        &mut self,
26        workflow_id: &str,
27        expression: ScheduleExpression,
28        conflict_policy: ConflictPolicy,
29        timezone: &str,
30    ) -> WorkflowResult<String> {
31        let id = Uuid::new_v4().to_string();
32        let schedule = Schedule {
33            id: id.clone(),
34            workflow_id: workflow_id.to_string(),
35            expression,
36            conflict_policy,
37            enabled: true,
38            next_fire_at: None,
39            last_fired_at: None,
40            timezone: timezone.to_string(),
41            created_at: Utc::now(),
42        };
43
44        self.schedules.insert(id.clone(), schedule);
45        Ok(id)
46    }
47
48    /// List all schedules.
49    pub fn list_schedules(&self) -> Vec<&Schedule> {
50        self.schedules.values().collect()
51    }
52
53    /// Get schedules for a specific workflow.
54    pub fn schedules_for_workflow(&self, workflow_id: &str) -> Vec<&Schedule> {
55        self.schedules
56            .values()
57            .filter(|s| s.workflow_id == workflow_id)
58            .collect()
59    }
60
61    /// Pause a schedule.
62    pub fn pause_schedule(&mut self, schedule_id: &str) -> WorkflowResult<()> {
63        let schedule = self
64            .schedules
65            .get_mut(schedule_id)
66            .ok_or_else(|| WorkflowError::ScheduleError(format!("Not found: {}", schedule_id)))?;
67
68        schedule.enabled = false;
69        Ok(())
70    }
71
72    /// Resume a paused schedule.
73    pub fn resume_schedule(&mut self, schedule_id: &str) -> WorkflowResult<()> {
74        let schedule = self
75            .schedules
76            .get_mut(schedule_id)
77            .ok_or_else(|| WorkflowError::ScheduleError(format!("Not found: {}", schedule_id)))?;
78
79        schedule.enabled = true;
80        Ok(())
81    }
82
83    /// Remove a schedule.
84    pub fn remove_schedule(&mut self, schedule_id: &str) -> WorkflowResult<Schedule> {
85        self.schedules
86            .remove(schedule_id)
87            .ok_or_else(|| WorkflowError::ScheduleError(format!("Not found: {}", schedule_id)))
88    }
89
90    /// Get adaptive schedule recommendations (placeholder — uses execution history).
91    pub fn get_adaptive_recommendation(
92        &self,
93        schedule_id: &str,
94    ) -> WorkflowResult<AdaptiveSchedule> {
95        let schedule = self
96            .schedules
97            .get(schedule_id)
98            .ok_or_else(|| WorkflowError::ScheduleError(format!("Not found: {}", schedule_id)))?;
99
100        Ok(AdaptiveSchedule {
101            schedule_id: schedule_id.to_string(),
102            recommended_time: "08:30".to_string(),
103            reason: "Historical success rate is 12% higher at 08:30 vs current schedule"
104                .to_string(),
105            success_rate_at_recommended: 0.95,
106            success_rate_at_current: 0.83,
107        })
108    }
109}
110
111impl Default for SchedulerEngine {
112    fn default() -> Self {
113        Self::new()
114    }
115}
116
117#[cfg(test)]
118mod tests {
119    use super::*;
120
121    #[test]
122    fn test_schedule_lifecycle() {
123        let mut engine = SchedulerEngine::new();
124        let sid = engine
125            .create_schedule(
126                "wf-1",
127                ScheduleExpression::Cron("0 8 * * 1-5".to_string()),
128                ConflictPolicy::Skip,
129                "UTC",
130            )
131            .unwrap();
132
133        assert_eq!(engine.list_schedules().len(), 1);
134        assert!(engine.pause_schedule(&sid).is_ok());
135        assert!(!engine.schedules.get(&sid).unwrap().enabled);
136        assert!(engine.resume_schedule(&sid).is_ok());
137        assert!(engine.schedules.get(&sid).unwrap().enabled);
138        assert!(engine.remove_schedule(&sid).is_ok());
139        assert_eq!(engine.list_schedules().len(), 0);
140    }
141}