1use std::collections::HashMap;
33use std::mem::replace;
34use std::sync::{Arc, Mutex, MutexGuard};
35use std::time::Duration;
36
37use rust_decimal::Decimal;
38use serde::{Deserialize, Serialize};
39use serde_json::{Value, json};
40
41use ironflow_store::entities::{RunFilter, RunStatus, StepKind, StepStatus};
42use ironflow_store::store::Store;
43
44use crate::config::StepConfig;
45use crate::error::EngineError;
46use crate::executor::{StepArtifacts, StepOutput};
47
48pub const DEFAULT_PLAN_MAX_DEPTH: u32 = 3;
50
51pub const MAX_PLANNED_STEPS: usize = 1000;
56
57pub const DEFAULT_ESTIMATE_SAMPLE_RUNS: u32 = 20;
59
60pub const DEFAULT_PLAN_TIMEOUT: Duration = Duration::from_secs(10);
76
77#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
96#[serde(tag = "state", rename_all = "snake_case")]
97pub enum ConditionResult {
98 Evaluated {
100 expression: String,
103 value: bool,
105 },
106 Skipped {
109 reason: String,
111 },
112 Unevaluable {
114 expression: String,
116 reason: String,
118 },
119}
120
121#[derive(Debug, Clone, Serialize, Deserialize)]
142pub struct PlannedStep {
143 pub name: String,
145 pub kind: StepKind,
147 pub workflow: String,
149 pub depth: u32,
151 pub depends_on: Vec<String>,
157 pub condition: Option<ConditionResult>,
163 pub parallel_group: Option<String>,
166 #[serde(with = "opt_duration_ms")]
168 pub estimated_duration: Option<Duration>,
169}
170
171#[derive(Debug, Clone, Serialize, Deserialize)]
189pub struct ExecutionPlan {
190 pub workflow: String,
192 pub steps: Vec<PlannedStep>,
194 #[serde(with = "opt_duration_ms")]
196 pub estimated_duration: Option<Duration>,
197 pub max_depth: u32,
199 pub truncated: bool,
201 pub incomplete_reason: Option<String>,
203}
204
205#[derive(Debug, Clone)]
225pub struct PlanOptions {
226 pub max_depth: u32,
228 pub estimate_durations: bool,
230 pub sample_runs: u32,
232 pub timeout: Duration,
235}
236
237impl Default for PlanOptions {
238 fn default() -> Self {
239 Self {
240 max_depth: DEFAULT_PLAN_MAX_DEPTH,
241 estimate_durations: true,
242 sample_runs: DEFAULT_ESTIMATE_SAMPLE_RUNS,
243 timeout: DEFAULT_PLAN_TIMEOUT,
244 }
245 }
246}
247
248mod opt_duration_ms {
250 use std::time::Duration;
251
252 use serde::{Deserialize, Deserializer, Serialize, Serializer};
253
254 pub(super) fn serialize<S: Serializer>(
256 value: &Option<Duration>,
257 serializer: S,
258 ) -> Result<S::Ok, S::Error> {
259 let millis = value.map(|d| d.as_millis() as u64);
260 millis.serialize(serializer)
261 }
262
263 pub(super) fn deserialize<'de, D: Deserializer<'de>>(
265 deserializer: D,
266 ) -> Result<Option<Duration>, D::Error> {
267 let millis = Option::<u64>::deserialize(deserializer)?;
268 Ok(millis.map(Duration::from_millis))
269 }
270}
271
272pub(crate) type SharedPlanRecorder = Arc<Mutex<PlanRecorder>>;
274
275pub(crate) struct PlanRecorder {
277 payload: Value,
278 workflow: String,
279 steps: Vec<PlannedStep>,
280 last_names: Vec<String>,
281 pending_condition: Option<ConditionResult>,
282 depth: u32,
283 max_depth: u32,
284 parallel_groups: u32,
285 estimates: HashMap<String, Duration>,
286 truncated: bool,
287 incomplete_reason: Option<String>,
288}
289
290impl PlanRecorder {
291 pub(crate) fn new(
293 workflow: String,
294 payload: Value,
295 max_depth: u32,
296 estimates: HashMap<String, Duration>,
297 ) -> Self {
298 Self {
299 payload,
300 workflow,
301 steps: Vec::new(),
302 last_names: Vec::new(),
303 pending_condition: None,
304 depth: 0,
305 max_depth,
306 parallel_groups: 0,
307 estimates,
308 truncated: false,
309 incomplete_reason: None,
310 }
311 }
312
313 pub(crate) fn payload(&self) -> Value {
315 self.payload.clone()
316 }
317
318 pub(crate) fn swap_payload(&mut self, next: Value) -> Value {
323 replace(&mut self.payload, next)
324 }
325
326 pub(crate) fn set_condition(&mut self, condition: ConditionResult) {
328 self.pending_condition = Some(condition);
329 }
330
331 pub(crate) fn estimate_for(&self, name: &str) -> Option<Duration> {
333 self.estimates.get(name).copied()
334 }
335
336 pub(crate) fn seed_estimate(&mut self, name: &str, duration: Duration) {
341 self.estimates.entry(name.to_string()).or_insert(duration);
342 }
343
344 pub(crate) fn record(
349 &mut self,
350 name: &str,
351 kind: StepKind,
352 workflow: &str,
353 parallel_group: Option<String>,
354 ) -> bool {
355 if self.steps.len() >= MAX_PLANNED_STEPS {
356 self.truncated = true;
357 if self.incomplete_reason.is_none() {
358 self.incomplete_reason = Some(format!("step cap of {MAX_PLANNED_STEPS} reached"));
359 }
360 return false;
361 }
362
363 let estimated_duration = self.estimates.get(name).copied();
364 self.steps.push(PlannedStep {
365 name: name.to_string(),
366 kind,
367 workflow: workflow.to_string(),
368 depth: self.depth,
369 depends_on: self.last_names.clone(),
370 condition: self.pending_condition.take(),
371 parallel_group,
372 estimated_duration,
373 });
374 true
375 }
376
377 pub(crate) fn record_or_stop(
382 &mut self,
383 name: &str,
384 kind: StepKind,
385 workflow: &str,
386 parallel_group: Option<String>,
387 ) -> Result<(), EngineError> {
388 if self.record(name, kind, workflow, parallel_group) {
389 Ok(())
390 } else {
391 Err(EngineError::InvalidWorkflow(
392 "execution plan exceeded the maximum number of steps".to_string(),
393 ))
394 }
395 }
396
397 pub(crate) fn set_last(&mut self, names: Vec<String>) {
399 self.last_names = names;
400 }
401
402 pub(crate) fn last(&self) -> &[String] {
404 &self.last_names
405 }
406
407 pub(crate) fn next_group(&mut self) -> String {
409 self.parallel_groups += 1;
410 format!("parallel-{}", self.parallel_groups)
411 }
412
413 pub(crate) fn enter_workflow(&mut self) -> bool {
415 if self.depth + 1 > self.max_depth {
416 self.truncated = true;
417 if self.incomplete_reason.is_none() {
418 self.incomplete_reason = Some(format!(
419 "sub-workflow expansion stopped at depth {}",
420 self.max_depth
421 ));
422 }
423 return false;
424 }
425 self.depth += 1;
426 true
427 }
428
429 pub(crate) fn leave_workflow(&mut self) {
431 self.depth = self.depth.saturating_sub(1);
432 }
433
434 pub(crate) fn fail(&mut self, reason: String) {
436 self.truncated = true;
437 if self.incomplete_reason.is_none() {
438 self.incomplete_reason = Some(reason);
439 }
440 }
441
442 pub(crate) fn snapshot(&self) -> ExecutionPlan {
444 ExecutionPlan {
445 workflow: self.workflow.clone(),
446 estimated_duration: total_estimate(&self.steps),
447 steps: self.steps.clone(),
448 max_depth: self.max_depth,
449 truncated: self.truncated,
450 incomplete_reason: self.incomplete_reason.clone(),
451 }
452 }
453
454 pub(crate) fn into_plan(self) -> ExecutionPlan {
456 ExecutionPlan {
457 workflow: self.workflow,
458 estimated_duration: total_estimate(&self.steps),
459 steps: self.steps,
460 max_depth: self.max_depth,
461 truncated: self.truncated,
462 incomplete_reason: self.incomplete_reason,
463 }
464 }
465}
466
467pub(crate) fn lock_plan(plan: &SharedPlanRecorder) -> MutexGuard<'_, PlanRecorder> {
472 plan.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
473}
474
475fn total_estimate(steps: &[PlannedStep]) -> Option<Duration> {
481 let mut total = Duration::ZERO;
482 let mut seen_any = false;
483 let mut current_group: Option<&str> = None;
484 let mut group_max = Duration::ZERO;
485
486 for step in steps {
487 match step.parallel_group.as_deref() {
488 Some(group) if current_group == Some(group) => {
489 if let Some(estimate) = step.estimated_duration {
490 seen_any = true;
491 group_max = group_max.max(estimate);
492 }
493 }
494 Some(group) => {
495 if current_group.is_some() {
496 total += group_max;
497 }
498 current_group = Some(group);
499 group_max = step.estimated_duration.unwrap_or(Duration::ZERO);
500 if step.estimated_duration.is_some() {
501 seen_any = true;
502 }
503 }
504 None => {
505 if current_group.is_some() {
506 total += group_max;
507 current_group = None;
508 group_max = Duration::ZERO;
509 }
510 if let Some(estimate) = step.estimated_duration {
511 seen_any = true;
512 total += estimate;
513 }
514 }
515 }
516 }
517
518 if current_group.is_some() {
519 total += group_max;
520 }
521
522 seen_any.then_some(total)
523}
524
525pub(crate) fn planned_output(config: &StepConfig, estimate: Option<Duration>) -> StepOutput {
530 let output = match config {
531 StepConfig::Shell(_) => json!({"stdout": "", "stderr": "", "exit_code": 0}),
532 StepConfig::Http(_) => json!({"status": 200, "headers": {}, "body": ""}),
533 StepConfig::Agent(_) => json!({}),
534 StepConfig::Workflow(c) => json!({
535 "run_id": Value::Null,
536 "workflow_name": c.workflow_name,
537 "status": "completed",
538 "cost_usd": 0,
539 "duration_ms": 0,
540 }),
541 StepConfig::Approval(_) | StepConfig::Decision(_) | StepConfig::Delay(_) => Value::Null,
542 };
543
544 StepOutput {
545 output,
546 duration_ms: estimate.map(|d| d.as_millis() as u64).unwrap_or(0),
547 cost_usd: Decimal::ZERO,
548 input_tokens: None,
549 cache_read_input_tokens: None,
550 cache_creation_input_tokens: None,
551 output_tokens: None,
552 model: None,
553 debug_messages: None,
554 artifacts: StepArtifacts::default(),
555 account_id: None,
556 environment_id: None,
557 }
558}
559
560pub(crate) fn planned_custom_output(estimate: Option<Duration>) -> StepOutput {
562 StepOutput {
563 output: json!({}),
564 duration_ms: estimate.map(|d| d.as_millis() as u64).unwrap_or(0),
565 cost_usd: Decimal::ZERO,
566 input_tokens: None,
567 cache_read_input_tokens: None,
568 cache_creation_input_tokens: None,
569 output_tokens: None,
570 model: None,
571 debug_messages: None,
572 artifacts: StepArtifacts::default(),
573 account_id: None,
574 environment_id: None,
575 }
576}
577
578pub async fn estimate_durations(
604 store: &Arc<dyn Store>,
605 workflow_name: &str,
606 sample_runs: u32,
607) -> Result<HashMap<String, Duration>, EngineError> {
608 let page = store
609 .list_runs(
610 RunFilter {
611 workflow_name: Some(workflow_name.to_string()),
612 status: Some(RunStatus::Completed),
613 has_steps: Some(true),
614 ..RunFilter::default()
615 },
616 1,
617 sample_runs.clamp(1, 100),
618 )
619 .await?;
620
621 let mut totals: HashMap<String, (u64, u64)> = HashMap::new();
622 for run in &page.items {
623 for step in store.list_steps(run.id).await? {
624 if step.status.state != StepStatus::Completed {
625 continue;
626 }
627 let entry = totals.entry(step.name).or_insert((0, 0));
628 entry.0 += step.duration_ms;
629 entry.1 += 1;
630 }
631 }
632
633 Ok(totals
634 .into_iter()
635 .filter(|(_, (_, count))| *count > 0)
636 .map(|(name, (sum, count))| (name, Duration::from_millis(sum / count)))
637 .collect())
638}
639
640impl ExecutionPlan {
641 pub fn estimated_duration_ms(&self) -> Option<u64> {
661 self.estimated_duration.map(|d| d.as_millis() as u64)
662 }
663}
664
665impl PlannedStep {
666 pub fn estimated_duration_ms(&self) -> Option<u64> {
689 self.estimated_duration.map(|d| d.as_millis() as u64)
690 }
691}
692
693#[cfg(test)]
694mod tests {
695 use std::collections::HashMap;
696 use std::thread::spawn;
697 use std::time::Duration;
698
699 use serde_json::{from_value, json, to_value};
700
701 use crate::config::{HttpConfig, ShellConfig};
702
703 use super::*;
704
705 fn step(name: &str, group: Option<&str>, estimate: Option<u64>) -> PlannedStep {
706 PlannedStep {
707 name: name.to_string(),
708 kind: StepKind::Shell,
709 workflow: "wf".to_string(),
710 depth: 0,
711 depends_on: Vec::new(),
712 condition: None,
713 parallel_group: group.map(str::to_string),
714 estimated_duration: estimate.map(Duration::from_millis),
715 }
716 }
717
718 #[test]
719 fn total_estimate_sums_sequential_steps() {
720 let steps = vec![
721 step("a", None, Some(100)),
722 step("b", None, Some(250)),
723 step("c", None, Some(50)),
724 ];
725 assert_eq!(total_estimate(&steps), Some(Duration::from_millis(400)));
726 }
727
728 #[test]
729 fn total_estimate_counts_a_parallel_group_once_at_its_slowest() {
730 let steps = vec![
731 step("build", None, Some(100)),
732 step("t1", Some("parallel-1"), Some(300)),
733 step("t2", Some("parallel-1"), Some(700)),
734 step("t3", Some("parallel-1"), Some(200)),
735 step("deploy", None, Some(100)),
736 ];
737 assert_eq!(total_estimate(&steps), Some(Duration::from_millis(900)));
738 }
739
740 #[test]
741 fn total_estimate_handles_a_trailing_parallel_group() {
742 let steps = vec![
743 step("build", None, Some(100)),
744 step("t1", Some("parallel-1"), Some(300)),
745 step("t2", Some("parallel-1"), Some(700)),
746 ];
747 assert_eq!(total_estimate(&steps), Some(Duration::from_millis(800)));
748 }
749
750 #[test]
751 fn total_estimate_is_none_without_any_estimate() {
752 let steps = vec![step("a", None, None), step("b", None, None)];
753 assert_eq!(total_estimate(&steps), None);
754 }
755
756 #[test]
757 fn planned_step_duration_round_trips_as_milliseconds() {
758 let original = step("a", None, Some(1234));
759 let value = to_value(&original).expect("serialize");
760 assert_eq!(value["estimated_duration"], 1234);
761
762 let back: PlannedStep = from_value(value).expect("deserialize");
763 assert_eq!(back.estimated_duration, Some(Duration::from_millis(1234)));
764 }
765
766 #[test]
767 fn planned_step_duration_round_trips_when_absent() {
768 let original = step("a", None, None);
769 let value = to_value(&original).expect("serialize");
770 assert!(value["estimated_duration"].is_null());
771
772 let back: PlannedStep = from_value(value).expect("deserialize");
773 assert_eq!(back.estimated_duration, None);
774 }
775
776 #[test]
777 fn planned_shell_and_http_outputs_look_successful() {
778 let shell = planned_output(&StepConfig::Shell(ShellConfig::new("echo hi")), None);
779 assert!(shell.is_success());
780
781 let http = planned_output(
782 &StepConfig::Http(HttpConfig::get("https://example.com")),
783 None,
784 );
785 assert!(http.is_success());
786 }
787
788 #[test]
789 fn planned_output_carries_the_estimate_as_its_duration() {
790 let output = planned_output(
791 &StepConfig::Shell(ShellConfig::new("echo hi")),
792 Some(Duration::from_millis(900)),
793 );
794 assert_eq!(output.duration_ms, 900);
795 assert_eq!(output.cost_usd, Decimal::ZERO);
796 }
797
798 #[test]
799 fn planned_custom_output_is_an_empty_object() {
800 let output = planned_custom_output(None);
801 assert_eq!(output.output, json!({}));
802 assert_eq!(output.duration_ms, 0);
803 }
804
805 #[test]
806 fn condition_result_serializes_its_state_tag() {
807 let evaluated = to_value(ConditionResult::Evaluated {
808 expression: "env == prod".to_string(),
809 value: false,
810 })
811 .expect("serialize");
812 assert_eq!(evaluated["state"], "evaluated");
813 assert_eq!(evaluated["value"], false);
814
815 let skipped = to_value(ConditionResult::Skipped {
816 reason: "not prod".to_string(),
817 })
818 .expect("serialize");
819 assert_eq!(skipped["state"], "skipped");
820
821 let unevaluable = to_value(ConditionResult::Unevaluable {
822 expression: "build succeeded".to_string(),
823 reason: "depends on a step output".to_string(),
824 })
825 .expect("serialize");
826 assert_eq!(unevaluable["state"], "unevaluable");
827 }
828
829 #[test]
830 fn recorder_records_dependencies_and_conditions() {
831 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
832 assert!(recorder.record("a", StepKind::Shell, "wf", None));
833 recorder.set_last(vec!["a".to_string()]);
834 recorder.set_condition(ConditionResult::Skipped {
835 reason: "nope".to_string(),
836 });
837 assert!(recorder.record("b", StepKind::Shell, "wf", None));
838
839 let plan = recorder.into_plan();
840 assert_eq!(plan.steps.len(), 2);
841 assert_eq!(plan.steps[1].depends_on, vec!["a".to_string()]);
842 assert!(matches!(
843 plan.steps[1].condition,
844 Some(ConditionResult::Skipped { .. })
845 ));
846 assert!(plan.steps[0].condition.is_none());
847 }
848
849 #[test]
850 fn recorder_stops_at_the_step_cap() {
851 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
852 for index in 0..MAX_PLANNED_STEPS {
853 assert!(recorder.record(&format!("s{index}"), StepKind::Shell, "wf", None));
854 }
855 assert!(!recorder.record("overflow", StepKind::Shell, "wf", None));
856
857 let plan = recorder.into_plan();
858 assert_eq!(plan.steps.len(), MAX_PLANNED_STEPS);
859 assert!(plan.truncated);
860 assert!(
861 plan.incomplete_reason
862 .expect("a reason")
863 .contains("step cap")
864 );
865 }
866
867 #[test]
868 fn plan_record_or_stop_errors_once_the_step_cap_is_reached() {
869 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
870 for index in 0..MAX_PLANNED_STEPS {
871 recorder
872 .record_or_stop(&format!("s{index}"), StepKind::Shell, "wf", None)
873 .expect("under the cap");
874 }
875 let err = recorder
876 .record_or_stop("overflow", StepKind::Shell, "wf", None)
877 .expect_err("over the cap");
878
879 assert!(matches!(err, EngineError::InvalidWorkflow(_)));
880 assert!(err.to_string().contains("maximum number of steps"));
881 let plan = recorder.into_plan();
882 assert_eq!(plan.steps.len(), MAX_PLANNED_STEPS);
883 assert!(plan.truncated);
884 }
885
886 #[test]
887 fn recorder_refuses_to_expand_past_the_depth_limit() {
888 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 1, HashMap::new());
889 assert!(recorder.enter_workflow());
890 assert!(!recorder.enter_workflow());
891
892 let plan = recorder.snapshot();
893 assert!(plan.truncated);
894 assert!(
895 plan.incomplete_reason
896 .expect("a reason")
897 .contains("depth 1")
898 );
899 }
900
901 #[test]
902 fn recorder_allocates_successive_parallel_group_names() {
903 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
904 assert_eq!(recorder.next_group(), "parallel-1");
905 assert_eq!(recorder.next_group(), "parallel-2");
906 }
907
908 #[test]
909 fn recorder_swaps_and_restores_the_payload() {
910 let mut recorder =
911 PlanRecorder::new("wf".to_string(), json!({"env": "prod"}), 3, HashMap::new());
912 let previous = recorder.swap_payload(json!({"env": "dev"}));
913 assert_eq!(previous, json!({"env": "prod"}));
914 assert_eq!(recorder.payload(), json!({"env": "dev"}));
915 recorder.swap_payload(previous);
916 assert_eq!(recorder.payload(), json!({"env": "prod"}));
917 }
918
919 #[test]
920 fn recorder_keeps_the_first_failure_reason() {
921 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
922 recorder.fail("first".to_string());
923 recorder.fail("second".to_string());
924 let plan = recorder.into_plan();
925 assert_eq!(plan.incomplete_reason.as_deref(), Some("first"));
926 assert!(plan.truncated);
927 }
928
929 #[test]
930 fn recorder_uses_the_history_estimate_for_a_known_step() {
931 let estimates = HashMap::from([("build".to_string(), Duration::from_millis(400))]);
932 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, estimates);
933 assert_eq!(
934 recorder.estimate_for("build"),
935 Some(Duration::from_millis(400))
936 );
937 assert_eq!(recorder.estimate_for("unknown"), None);
938 recorder.record("build", StepKind::Shell, "wf", None);
939 let plan = recorder.into_plan();
940 assert_eq!(plan.estimated_duration, Some(Duration::from_millis(400)));
941 }
942
943 #[test]
944 fn lock_plan_recovers_from_poisoning() {
945 let shared: SharedPlanRecorder = Arc::new(Mutex::new(PlanRecorder::new(
946 "wf".to_string(),
947 json!({}),
948 3,
949 HashMap::new(),
950 )));
951 let poisoner = Arc::clone(&shared);
952 let _ = spawn(move || {
953 let _guard = poisoner.lock().expect("lock");
954 panic!("poison the mutex");
955 })
956 .join();
957
958 assert_eq!(lock_plan(&shared).payload(), json!({}));
959 }
960}