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 next_group(&mut self) -> String {
404 self.parallel_groups += 1;
405 format!("parallel-{}", self.parallel_groups)
406 }
407
408 pub(crate) fn enter_workflow(&mut self) -> bool {
410 if self.depth + 1 > self.max_depth {
411 self.truncated = true;
412 if self.incomplete_reason.is_none() {
413 self.incomplete_reason = Some(format!(
414 "sub-workflow expansion stopped at depth {}",
415 self.max_depth
416 ));
417 }
418 return false;
419 }
420 self.depth += 1;
421 true
422 }
423
424 pub(crate) fn leave_workflow(&mut self) {
426 self.depth = self.depth.saturating_sub(1);
427 }
428
429 pub(crate) fn fail(&mut self, reason: String) {
431 self.truncated = true;
432 if self.incomplete_reason.is_none() {
433 self.incomplete_reason = Some(reason);
434 }
435 }
436
437 pub(crate) fn snapshot(&self) -> ExecutionPlan {
439 ExecutionPlan {
440 workflow: self.workflow.clone(),
441 estimated_duration: total_estimate(&self.steps),
442 steps: self.steps.clone(),
443 max_depth: self.max_depth,
444 truncated: self.truncated,
445 incomplete_reason: self.incomplete_reason.clone(),
446 }
447 }
448
449 pub(crate) fn into_plan(self) -> ExecutionPlan {
451 ExecutionPlan {
452 workflow: self.workflow,
453 estimated_duration: total_estimate(&self.steps),
454 steps: self.steps,
455 max_depth: self.max_depth,
456 truncated: self.truncated,
457 incomplete_reason: self.incomplete_reason,
458 }
459 }
460}
461
462pub(crate) fn lock_plan(plan: &SharedPlanRecorder) -> MutexGuard<'_, PlanRecorder> {
467 plan.lock().unwrap_or_else(|poisoned| poisoned.into_inner())
468}
469
470fn total_estimate(steps: &[PlannedStep]) -> Option<Duration> {
476 let mut total = Duration::ZERO;
477 let mut seen_any = false;
478 let mut current_group: Option<&str> = None;
479 let mut group_max = Duration::ZERO;
480
481 for step in steps {
482 match step.parallel_group.as_deref() {
483 Some(group) if current_group == Some(group) => {
484 if let Some(estimate) = step.estimated_duration {
485 seen_any = true;
486 group_max = group_max.max(estimate);
487 }
488 }
489 Some(group) => {
490 if current_group.is_some() {
491 total += group_max;
492 }
493 current_group = Some(group);
494 group_max = step.estimated_duration.unwrap_or(Duration::ZERO);
495 if step.estimated_duration.is_some() {
496 seen_any = true;
497 }
498 }
499 None => {
500 if current_group.is_some() {
501 total += group_max;
502 current_group = None;
503 group_max = Duration::ZERO;
504 }
505 if let Some(estimate) = step.estimated_duration {
506 seen_any = true;
507 total += estimate;
508 }
509 }
510 }
511 }
512
513 if current_group.is_some() {
514 total += group_max;
515 }
516
517 seen_any.then_some(total)
518}
519
520pub(crate) fn planned_output(config: &StepConfig, estimate: Option<Duration>) -> StepOutput {
525 let output = match config {
526 StepConfig::Shell(_) => json!({"stdout": "", "stderr": "", "exit_code": 0}),
527 StepConfig::Http(_) => json!({"status": 200, "headers": {}, "body": ""}),
528 StepConfig::Agent(_) => json!({}),
529 StepConfig::Workflow(c) => json!({
530 "run_id": Value::Null,
531 "workflow_name": c.workflow_name,
532 "status": "completed",
533 "cost_usd": 0,
534 "duration_ms": 0,
535 }),
536 StepConfig::Approval(_) | StepConfig::Decision(_) | StepConfig::Delay(_) => Value::Null,
537 };
538
539 StepOutput {
540 output,
541 duration_ms: estimate.map(|d| d.as_millis() as u64).unwrap_or(0),
542 cost_usd: Decimal::ZERO,
543 input_tokens: None,
544 cache_read_input_tokens: None,
545 cache_creation_input_tokens: None,
546 output_tokens: None,
547 model: None,
548 debug_messages: None,
549 artifacts: StepArtifacts::default(),
550 account_id: None,
551 environment_id: None,
552 }
553}
554
555pub(crate) fn planned_custom_output(estimate: Option<Duration>) -> StepOutput {
557 StepOutput {
558 output: json!({}),
559 duration_ms: estimate.map(|d| d.as_millis() as u64).unwrap_or(0),
560 cost_usd: Decimal::ZERO,
561 input_tokens: None,
562 cache_read_input_tokens: None,
563 cache_creation_input_tokens: None,
564 output_tokens: None,
565 model: None,
566 debug_messages: None,
567 artifacts: StepArtifacts::default(),
568 account_id: None,
569 environment_id: None,
570 }
571}
572
573pub async fn estimate_durations(
599 store: &Arc<dyn Store>,
600 workflow_name: &str,
601 sample_runs: u32,
602) -> Result<HashMap<String, Duration>, EngineError> {
603 let page = store
604 .list_runs(
605 RunFilter {
606 workflow_name: Some(workflow_name.to_string()),
607 status: Some(RunStatus::Completed),
608 has_steps: Some(true),
609 ..RunFilter::default()
610 },
611 1,
612 sample_runs.clamp(1, 100),
613 )
614 .await?;
615
616 let mut totals: HashMap<String, (u64, u64)> = HashMap::new();
617 for run in &page.items {
618 for step in store.list_steps(run.id).await? {
619 if step.status.state != StepStatus::Completed {
620 continue;
621 }
622 let entry = totals.entry(step.name).or_insert((0, 0));
623 entry.0 += step.duration_ms;
624 entry.1 += 1;
625 }
626 }
627
628 Ok(totals
629 .into_iter()
630 .filter(|(_, (_, count))| *count > 0)
631 .map(|(name, (sum, count))| (name, Duration::from_millis(sum / count)))
632 .collect())
633}
634
635impl ExecutionPlan {
636 pub fn estimated_duration_ms(&self) -> Option<u64> {
656 self.estimated_duration.map(|d| d.as_millis() as u64)
657 }
658}
659
660impl PlannedStep {
661 pub fn estimated_duration_ms(&self) -> Option<u64> {
684 self.estimated_duration.map(|d| d.as_millis() as u64)
685 }
686}
687
688#[cfg(test)]
689mod tests {
690 use std::collections::HashMap;
691 use std::thread::spawn;
692 use std::time::Duration;
693
694 use serde_json::{from_value, json, to_value};
695
696 use crate::config::{HttpConfig, ShellConfig};
697
698 use super::*;
699
700 fn step(name: &str, group: Option<&str>, estimate: Option<u64>) -> PlannedStep {
701 PlannedStep {
702 name: name.to_string(),
703 kind: StepKind::Shell,
704 workflow: "wf".to_string(),
705 depth: 0,
706 depends_on: Vec::new(),
707 condition: None,
708 parallel_group: group.map(str::to_string),
709 estimated_duration: estimate.map(Duration::from_millis),
710 }
711 }
712
713 #[test]
714 fn total_estimate_sums_sequential_steps() {
715 let steps = vec![
716 step("a", None, Some(100)),
717 step("b", None, Some(250)),
718 step("c", None, Some(50)),
719 ];
720 assert_eq!(total_estimate(&steps), Some(Duration::from_millis(400)));
721 }
722
723 #[test]
724 fn total_estimate_counts_a_parallel_group_once_at_its_slowest() {
725 let steps = vec![
726 step("build", None, Some(100)),
727 step("t1", Some("parallel-1"), Some(300)),
728 step("t2", Some("parallel-1"), Some(700)),
729 step("t3", Some("parallel-1"), Some(200)),
730 step("deploy", None, Some(100)),
731 ];
732 assert_eq!(total_estimate(&steps), Some(Duration::from_millis(900)));
733 }
734
735 #[test]
736 fn total_estimate_handles_a_trailing_parallel_group() {
737 let steps = vec![
738 step("build", None, Some(100)),
739 step("t1", Some("parallel-1"), Some(300)),
740 step("t2", Some("parallel-1"), Some(700)),
741 ];
742 assert_eq!(total_estimate(&steps), Some(Duration::from_millis(800)));
743 }
744
745 #[test]
746 fn total_estimate_is_none_without_any_estimate() {
747 let steps = vec![step("a", None, None), step("b", None, None)];
748 assert_eq!(total_estimate(&steps), None);
749 }
750
751 #[test]
752 fn planned_step_duration_round_trips_as_milliseconds() {
753 let original = step("a", None, Some(1234));
754 let value = to_value(&original).expect("serialize");
755 assert_eq!(value["estimated_duration"], 1234);
756
757 let back: PlannedStep = from_value(value).expect("deserialize");
758 assert_eq!(back.estimated_duration, Some(Duration::from_millis(1234)));
759 }
760
761 #[test]
762 fn planned_step_duration_round_trips_when_absent() {
763 let original = step("a", None, None);
764 let value = to_value(&original).expect("serialize");
765 assert!(value["estimated_duration"].is_null());
766
767 let back: PlannedStep = from_value(value).expect("deserialize");
768 assert_eq!(back.estimated_duration, None);
769 }
770
771 #[test]
772 fn planned_shell_and_http_outputs_look_successful() {
773 let shell = planned_output(&StepConfig::Shell(ShellConfig::new("echo hi")), None);
774 assert!(shell.is_success());
775
776 let http = planned_output(
777 &StepConfig::Http(HttpConfig::get("https://example.com")),
778 None,
779 );
780 assert!(http.is_success());
781 }
782
783 #[test]
784 fn planned_output_carries_the_estimate_as_its_duration() {
785 let output = planned_output(
786 &StepConfig::Shell(ShellConfig::new("echo hi")),
787 Some(Duration::from_millis(900)),
788 );
789 assert_eq!(output.duration_ms, 900);
790 assert_eq!(output.cost_usd, Decimal::ZERO);
791 }
792
793 #[test]
794 fn planned_custom_output_is_an_empty_object() {
795 let output = planned_custom_output(None);
796 assert_eq!(output.output, json!({}));
797 assert_eq!(output.duration_ms, 0);
798 }
799
800 #[test]
801 fn condition_result_serializes_its_state_tag() {
802 let evaluated = to_value(ConditionResult::Evaluated {
803 expression: "env == prod".to_string(),
804 value: false,
805 })
806 .expect("serialize");
807 assert_eq!(evaluated["state"], "evaluated");
808 assert_eq!(evaluated["value"], false);
809
810 let skipped = to_value(ConditionResult::Skipped {
811 reason: "not prod".to_string(),
812 })
813 .expect("serialize");
814 assert_eq!(skipped["state"], "skipped");
815
816 let unevaluable = to_value(ConditionResult::Unevaluable {
817 expression: "build succeeded".to_string(),
818 reason: "depends on a step output".to_string(),
819 })
820 .expect("serialize");
821 assert_eq!(unevaluable["state"], "unevaluable");
822 }
823
824 #[test]
825 fn recorder_records_dependencies_and_conditions() {
826 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
827 assert!(recorder.record("a", StepKind::Shell, "wf", None));
828 recorder.set_last(vec!["a".to_string()]);
829 recorder.set_condition(ConditionResult::Skipped {
830 reason: "nope".to_string(),
831 });
832 assert!(recorder.record("b", StepKind::Shell, "wf", None));
833
834 let plan = recorder.into_plan();
835 assert_eq!(plan.steps.len(), 2);
836 assert_eq!(plan.steps[1].depends_on, vec!["a".to_string()]);
837 assert!(matches!(
838 plan.steps[1].condition,
839 Some(ConditionResult::Skipped { .. })
840 ));
841 assert!(plan.steps[0].condition.is_none());
842 }
843
844 #[test]
845 fn recorder_stops_at_the_step_cap() {
846 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
847 for index in 0..MAX_PLANNED_STEPS {
848 assert!(recorder.record(&format!("s{index}"), StepKind::Shell, "wf", None));
849 }
850 assert!(!recorder.record("overflow", StepKind::Shell, "wf", None));
851
852 let plan = recorder.into_plan();
853 assert_eq!(plan.steps.len(), MAX_PLANNED_STEPS);
854 assert!(plan.truncated);
855 assert!(
856 plan.incomplete_reason
857 .expect("a reason")
858 .contains("step cap")
859 );
860 }
861
862 #[test]
863 fn plan_record_or_stop_errors_once_the_step_cap_is_reached() {
864 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
865 for index in 0..MAX_PLANNED_STEPS {
866 recorder
867 .record_or_stop(&format!("s{index}"), StepKind::Shell, "wf", None)
868 .expect("under the cap");
869 }
870 let err = recorder
871 .record_or_stop("overflow", StepKind::Shell, "wf", None)
872 .expect_err("over the cap");
873
874 assert!(matches!(err, EngineError::InvalidWorkflow(_)));
875 assert!(err.to_string().contains("maximum number of steps"));
876 let plan = recorder.into_plan();
877 assert_eq!(plan.steps.len(), MAX_PLANNED_STEPS);
878 assert!(plan.truncated);
879 }
880
881 #[test]
882 fn recorder_refuses_to_expand_past_the_depth_limit() {
883 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 1, HashMap::new());
884 assert!(recorder.enter_workflow());
885 assert!(!recorder.enter_workflow());
886
887 let plan = recorder.snapshot();
888 assert!(plan.truncated);
889 assert!(
890 plan.incomplete_reason
891 .expect("a reason")
892 .contains("depth 1")
893 );
894 }
895
896 #[test]
897 fn recorder_allocates_successive_parallel_group_names() {
898 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
899 assert_eq!(recorder.next_group(), "parallel-1");
900 assert_eq!(recorder.next_group(), "parallel-2");
901 }
902
903 #[test]
904 fn recorder_swaps_and_restores_the_payload() {
905 let mut recorder =
906 PlanRecorder::new("wf".to_string(), json!({"env": "prod"}), 3, HashMap::new());
907 let previous = recorder.swap_payload(json!({"env": "dev"}));
908 assert_eq!(previous, json!({"env": "prod"}));
909 assert_eq!(recorder.payload(), json!({"env": "dev"}));
910 recorder.swap_payload(previous);
911 assert_eq!(recorder.payload(), json!({"env": "prod"}));
912 }
913
914 #[test]
915 fn recorder_keeps_the_first_failure_reason() {
916 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, HashMap::new());
917 recorder.fail("first".to_string());
918 recorder.fail("second".to_string());
919 let plan = recorder.into_plan();
920 assert_eq!(plan.incomplete_reason.as_deref(), Some("first"));
921 assert!(plan.truncated);
922 }
923
924 #[test]
925 fn recorder_uses_the_history_estimate_for_a_known_step() {
926 let estimates = HashMap::from([("build".to_string(), Duration::from_millis(400))]);
927 let mut recorder = PlanRecorder::new("wf".to_string(), json!({}), 3, estimates);
928 assert_eq!(
929 recorder.estimate_for("build"),
930 Some(Duration::from_millis(400))
931 );
932 assert_eq!(recorder.estimate_for("unknown"), None);
933 recorder.record("build", StepKind::Shell, "wf", None);
934 let plan = recorder.into_plan();
935 assert_eq!(plan.estimated_duration, Some(Duration::from_millis(400)));
936 }
937
938 #[test]
939 fn lock_plan_recovers_from_poisoning() {
940 let shared: SharedPlanRecorder = Arc::new(Mutex::new(PlanRecorder::new(
941 "wf".to_string(),
942 json!({}),
943 3,
944 HashMap::new(),
945 )));
946 let poisoner = Arc::clone(&shared);
947 let _ = spawn(move || {
948 let _guard = poisoner.lock().expect("lock");
949 panic!("poison the mutex");
950 })
951 .join();
952
953 assert_eq!(lock_plan(&shared).payload(), json!({}));
954 }
955}