Skip to main content

agent_orchestrator/dynamic_orchestration/
step_pool.rs

1use std::collections::HashMap;
2
3use crate::config::StepPrehookContext;
4use serde::{Deserialize, Serialize};
5
6pub use orchestrator_config::dynamic_step::DynamicStepConfig;
7
8/// Extension trait adding CEL trigger evaluation to `DynamicStepConfig`.
9pub trait DynamicStepConfigExt {
10    /// Check if this step matches the current context
11    fn matches(&self, context: &StepPrehookContext) -> bool;
12}
13
14impl DynamicStepConfigExt for DynamicStepConfig {
15    fn matches(&self, context: &StepPrehookContext) -> bool {
16        if let Some(ref trigger) = self.trigger {
17            return evaluate_trigger_condition(trigger, context).unwrap_or(false);
18        }
19        false
20    }
21}
22
23/// Pool of dynamic steps available for runtime selection
24#[derive(Debug, Clone, Serialize, Deserialize, Default)]
25pub struct DynamicStepPool {
26    /// Map of dynamic step ID -> config
27    #[serde(default)]
28    pub steps: HashMap<String, DynamicStepConfig>,
29}
30
31impl DynamicStepPool {
32    /// Create a new empty pool
33    pub fn new() -> Self {
34        Self {
35            steps: HashMap::new(),
36        }
37    }
38
39    /// Add a step to the pool
40    pub fn add_step(&mut self, step: DynamicStepConfig) {
41        self.steps.insert(step.id.clone(), step);
42    }
43
44    /// Find steps that match the current context
45    pub fn find_matching_steps(&self, context: &StepPrehookContext) -> Vec<&DynamicStepConfig> {
46        let mut matches: Vec<_> = self
47            .steps
48            .values()
49            .filter(|step| step.matches(context))
50            .collect();
51
52        // Sort by priority (descending)
53        matches.sort_by(|a, b| b.priority.cmp(&a.priority));
54        matches
55    }
56
57    /// Get a step by ID
58    pub fn get(&self, id: &str) -> Option<&DynamicStepConfig> {
59        self.steps.get(id)
60    }
61}
62
63pub(crate) fn evaluate_trigger_condition(
64    condition: &str,
65    context: &StepPrehookContext,
66) -> anyhow::Result<bool> {
67    crate::prehook::evaluate_step_prehook_expression(condition, context)
68}
69
70#[cfg(test)]
71mod tests {
72    use super::*;
73
74    #[test]
75    fn test_dynamic_step_pool() {
76        let mut pool = DynamicStepPool::new();
77        pool.add_step(DynamicStepConfig {
78            id: "step1".to_string(),
79            description: None,
80            step_type: "fix".to_string(),
81            agent_id: Some("fixer".to_string()),
82            template: None,
83            trigger: Some("active_ticket_count > 0".to_string()),
84            priority: 10,
85            max_runs: None,
86        });
87
88        let context = StepPrehookContext {
89            active_ticket_count: 5,
90            upstream_artifacts: Vec::new(),
91            build_error_count: 0,
92            test_failure_count: 0,
93            build_exit_code: None,
94            test_exit_code: None,
95            ..Default::default()
96        };
97
98        let matches = pool.find_matching_steps(&context);
99        assert!(!matches.is_empty());
100    }
101
102    #[test]
103    fn test_dynamic_step_pool_priority() {
104        let mut pool = DynamicStepPool::new();
105
106        pool.add_step(DynamicStepConfig {
107            id: "low_priority".to_string(),
108            description: None,
109            step_type: "fix".to_string(),
110            agent_id: None,
111            template: None,
112            trigger: Some("active_ticket_count > 0".to_string()),
113            priority: 1,
114            max_runs: None,
115        });
116
117        pool.add_step(DynamicStepConfig {
118            id: "high_priority".to_string(),
119            description: None,
120            step_type: "fix".to_string(),
121            agent_id: None,
122            template: None,
123            trigger: Some("active_ticket_count > 0".to_string()),
124            priority: 100,
125            max_runs: None,
126        });
127
128        let context = StepPrehookContext {
129            active_ticket_count: 5,
130            upstream_artifacts: Vec::new(),
131            build_error_count: 0,
132            test_failure_count: 0,
133            build_exit_code: None,
134            test_exit_code: None,
135            ..Default::default()
136        };
137
138        let matches = pool.find_matching_steps(&context);
139        assert_eq!(matches.len(), 2);
140        assert_eq!(matches[0].id, "high_priority");
141    }
142
143    #[test]
144    fn test_dynamic_step_pool_empty() {
145        let pool = DynamicStepPool::new();
146        assert!(pool.steps.is_empty());
147        let ctx = StepPrehookContext::default();
148        let matches = pool.find_matching_steps(&ctx);
149        assert!(matches.is_empty());
150    }
151
152    #[test]
153    fn test_dynamic_step_pool_get() {
154        let mut pool = DynamicStepPool::new();
155        pool.add_step(DynamicStepConfig {
156            id: "s1".to_string(),
157            description: Some("desc".to_string()),
158            step_type: "fix".to_string(),
159            agent_id: None,
160            template: None,
161            trigger: None,
162            priority: 0,
163            max_runs: Some(3),
164        });
165        assert!(pool.get("s1").is_some());
166        assert_eq!(pool.get("s1").expect("s1 should exist").max_runs, Some(3));
167        assert!(pool.get("nonexistent").is_none());
168    }
169
170    #[test]
171    fn test_dynamic_step_pool_overwrite() {
172        let mut pool = DynamicStepPool::new();
173        pool.add_step(DynamicStepConfig {
174            id: "s1".to_string(),
175            description: None,
176            step_type: "fix".to_string(),
177            agent_id: None,
178            template: None,
179            trigger: None,
180            priority: 1,
181            max_runs: None,
182        });
183        pool.add_step(DynamicStepConfig {
184            id: "s1".to_string(),
185            description: None,
186            step_type: "qa".to_string(),
187            agent_id: None,
188            template: None,
189            trigger: None,
190            priority: 99,
191            max_runs: None,
192        });
193        assert_eq!(pool.steps.len(), 1);
194        assert_eq!(pool.get("s1").expect("s1 should exist").step_type, "qa");
195        assert_eq!(pool.get("s1").expect("s1 should exist").priority, 99);
196    }
197
198    #[test]
199    fn test_dynamic_step_no_trigger_does_not_match() {
200        let step = DynamicStepConfig {
201            id: "s1".to_string(),
202            description: None,
203            step_type: "fix".to_string(),
204            agent_id: None,
205            template: None,
206            trigger: None,
207            priority: 0,
208            max_runs: None,
209        };
210        let ctx = StepPrehookContext {
211            active_ticket_count: 5,
212            ..Default::default()
213        };
214        assert!(!step.matches(&ctx));
215    }
216
217    #[test]
218    fn test_dynamic_step_matches_qa_exit_code_nonzero() {
219        let step = DynamicStepConfig {
220            id: "s1".to_string(),
221            description: None,
222            step_type: "fix".to_string(),
223            agent_id: None,
224            template: None,
225            trigger: Some("qa_exit_code != 0".to_string()),
226            priority: 0,
227            max_runs: None,
228        };
229        let ctx = StepPrehookContext {
230            qa_exit_code: Some(1),
231            ..Default::default()
232        };
233        assert!(step.matches(&ctx));
234    }
235
236    #[test]
237    fn test_dynamic_step_matches_qa_exit_code_zero() {
238        let step = DynamicStepConfig {
239            id: "s1".to_string(),
240            description: None,
241            step_type: "done".to_string(),
242            agent_id: None,
243            template: None,
244            trigger: Some("qa_exit_code == 0".to_string()),
245            priority: 0,
246            max_runs: None,
247        };
248        let ctx = StepPrehookContext {
249            qa_exit_code: Some(0),
250            ..Default::default()
251        };
252        assert!(step.matches(&ctx));
253    }
254
255    #[test]
256    fn test_dynamic_step_matches_qa_confidence_high() {
257        let step = DynamicStepConfig {
258            id: "s1".to_string(),
259            description: None,
260            step_type: "fix".to_string(),
261            agent_id: None,
262            template: None,
263            trigger: Some("qa_confidence > 0.8".to_string()),
264            priority: 0,
265            max_runs: None,
266        };
267        let ctx = StepPrehookContext {
268            qa_confidence: Some(0.9),
269            ..Default::default()
270        };
271        assert!(step.matches(&ctx));
272    }
273
274    #[test]
275    fn test_dynamic_step_matches_qa_confidence_low() {
276        let step = DynamicStepConfig {
277            id: "s1".to_string(),
278            description: None,
279            step_type: "fix".to_string(),
280            agent_id: None,
281            template: None,
282            trigger: Some("qa_confidence > 0.8".to_string()),
283            priority: 0,
284            max_runs: None,
285        };
286        let ctx = StepPrehookContext {
287            qa_confidence: Some(0.3),
288            ..Default::default()
289        };
290        assert!(!step.matches(&ctx));
291    }
292
293    #[test]
294    fn test_dynamic_step_matches_qa_confidence_medium() {
295        let step = DynamicStepConfig {
296            id: "s1".to_string(),
297            description: None,
298            step_type: "fix".to_string(),
299            agent_id: None,
300            template: None,
301            trigger: Some("qa_confidence > 0.5".to_string()),
302            priority: 0,
303            max_runs: None,
304        };
305        let ctx = StepPrehookContext {
306            qa_confidence: Some(0.6),
307            ..Default::default()
308        };
309        assert!(step.matches(&ctx));
310    }
311
312    #[test]
313    fn test_dynamic_step_matches_cycle_gt_2() {
314        let step = DynamicStepConfig {
315            id: "s1".to_string(),
316            description: None,
317            step_type: "fix".to_string(),
318            agent_id: None,
319            template: None,
320            trigger: Some("cycle > 2".to_string()),
321            priority: 0,
322            max_runs: None,
323        };
324        let ctx = StepPrehookContext {
325            cycle: 3,
326            ..Default::default()
327        };
328        assert!(step.matches(&ctx));
329    }
330
331    #[test]
332    fn test_dynamic_step_matches_cycle_gt_0() {
333        let step = DynamicStepConfig {
334            id: "s1".to_string(),
335            description: None,
336            step_type: "fix".to_string(),
337            agent_id: None,
338            template: None,
339            trigger: Some("cycle > 0".to_string()),
340            priority: 0,
341            max_runs: None,
342        };
343        let ctx = StepPrehookContext {
344            cycle: 1,
345            ..Default::default()
346        };
347        assert!(step.matches(&ctx));
348    }
349
350    #[test]
351    fn test_dynamic_step_does_not_match_cycle_zero() {
352        let step = DynamicStepConfig {
353            id: "s1".to_string(),
354            description: None,
355            step_type: "fix".to_string(),
356            agent_id: None,
357            template: None,
358            trigger: Some("cycle > 0".to_string()),
359            priority: 0,
360            max_runs: None,
361        };
362        let ctx = StepPrehookContext {
363            cycle: 0,
364            ..Default::default()
365        };
366        assert!(!step.matches(&ctx));
367    }
368
369    #[test]
370    fn test_dynamic_step_matches_active_tickets_zero() {
371        let step = DynamicStepConfig {
372            id: "s1".to_string(),
373            description: None,
374            step_type: "done".to_string(),
375            agent_id: None,
376            template: None,
377            trigger: Some("active_ticket_count == 0".to_string()),
378            priority: 0,
379            max_runs: None,
380        };
381        let ctx = StepPrehookContext {
382            active_ticket_count: 0,
383            ..Default::default()
384        };
385        assert!(step.matches(&ctx));
386    }
387
388    #[test]
389    fn test_dynamic_step_unknown_condition_returns_false() {
390        let step = DynamicStepConfig {
391            id: "s1".to_string(),
392            description: None,
393            step_type: "fix".to_string(),
394            agent_id: None,
395            template: None,
396            trigger: Some("some_unknown_field == true".to_string()),
397            priority: 0,
398            max_runs: None,
399        };
400        let ctx = StepPrehookContext::default();
401        assert!(!step.matches(&ctx));
402    }
403
404    #[test]
405    fn test_dynamic_step_pool_priority_three_steps() {
406        let mut pool = DynamicStepPool::new();
407        pool.add_step(DynamicStepConfig {
408            id: "low".to_string(),
409            description: None,
410            step_type: "fix".to_string(),
411            agent_id: None,
412            template: None,
413            trigger: Some("active_ticket_count > 0".to_string()),
414            priority: -5,
415            max_runs: None,
416        });
417        pool.add_step(DynamicStepConfig {
418            id: "mid".to_string(),
419            description: None,
420            step_type: "fix".to_string(),
421            agent_id: None,
422            template: None,
423            trigger: Some("active_ticket_count > 0".to_string()),
424            priority: 0,
425            max_runs: None,
426        });
427        pool.add_step(DynamicStepConfig {
428            id: "high".to_string(),
429            description: None,
430            step_type: "fix".to_string(),
431            agent_id: None,
432            template: None,
433            trigger: Some("active_ticket_count > 0".to_string()),
434            priority: 50,
435            max_runs: None,
436        });
437        let ctx = StepPrehookContext {
438            active_ticket_count: 1,
439            ..Default::default()
440        };
441        let matches = pool.find_matching_steps(&ctx);
442        assert_eq!(matches.len(), 3);
443        assert_eq!(matches[0].id, "high");
444        assert_eq!(matches[1].id, "mid");
445        assert_eq!(matches[2].id, "low");
446    }
447
448    #[test]
449    fn test_dynamic_step_pool_no_matches() {
450        let mut pool = DynamicStepPool::new();
451        pool.add_step(DynamicStepConfig {
452            id: "s1".to_string(),
453            description: None,
454            step_type: "fix".to_string(),
455            agent_id: None,
456            template: None,
457            trigger: Some("active_ticket_count > 0".to_string()),
458            priority: 10,
459            max_runs: None,
460        });
461        let ctx = StepPrehookContext {
462            active_ticket_count: 0,
463            ..Default::default()
464        };
465        let matches = pool.find_matching_steps(&ctx);
466        assert!(matches.is_empty());
467    }
468
469    #[test]
470    fn test_dynamic_step_pool_serde_round_trip() {
471        let mut pool = DynamicStepPool::new();
472        pool.add_step(DynamicStepConfig {
473            id: "s1".to_string(),
474            description: Some("my step".to_string()),
475            step_type: "fix".to_string(),
476            agent_id: Some("agent1".to_string()),
477            template: Some("tpl".to_string()),
478            trigger: Some("active_ticket_count > 0".to_string()),
479            priority: 42,
480            max_runs: Some(3),
481        });
482        let json = serde_json::to_string(&pool).expect("serialize dynamic step pool");
483        let pool2: DynamicStepPool =
484            serde_json::from_str(&json).expect("deserialize dynamic step pool");
485        assert_eq!(pool2.steps.len(), 1);
486        let s = pool2.get("s1").expect("s1 should exist after round-trip");
487        assert_eq!(s.priority, 42);
488        assert_eq!(s.max_runs, Some(3));
489    }
490
491    #[test]
492    fn test_step_prehook_context_default() {
493        let ctx = StepPrehookContext::default();
494        assert_eq!(ctx.cycle, 0);
495        assert_eq!(ctx.active_ticket_count, 0);
496        assert!(!ctx.qa_failed);
497        assert!(!ctx.fix_required);
498        assert!(ctx.qa_exit_code.is_none());
499        assert!(!ctx.self_test_passed);
500        assert!(!ctx.is_last_cycle);
501        assert_eq!(ctx.max_cycles, 0);
502    }
503}