1use std::fmt;
20use std::sync::{Arc, Mutex};
21use std::time::Instant;
22
23const ENV_PREFIX: &str = "IRONFLOW_GUARD_";
25
26#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct WorkflowGuardConfig {
47 pub max_depth: u32,
49 pub max_fan_out: u32,
51 pub max_workflow_tokens: u64,
53 pub workflow_timeout_secs: u64,
55}
56
57impl Default for WorkflowGuardConfig {
58 fn default() -> Self {
59 Self {
60 max_depth: 5,
61 max_fan_out: 20,
62 max_workflow_tokens: 100_000,
63 workflow_timeout_secs: 120,
64 }
65 }
66}
67
68impl WorkflowGuardConfig {
69 pub fn new() -> Self {
79 Self::default()
80 }
81
82 pub fn from_env() -> Self {
98 Self {
99 max_depth: parse_env("MAX_DEPTH", 5),
100 max_fan_out: parse_env("MAX_FAN_OUT", 20),
101 max_workflow_tokens: parse_env("MAX_WORKFLOW_TOKENS", 100_000),
102 workflow_timeout_secs: parse_env("WORKFLOW_TIMEOUT_SECS", 120),
103 }
104 }
105
106 pub fn with_max_depth(mut self, max_depth: u32) -> Self {
117 self.max_depth = max_depth;
118 self
119 }
120
121 pub fn with_max_fan_out(mut self, max_fan_out: u32) -> Self {
132 self.max_fan_out = max_fan_out;
133 self
134 }
135
136 pub fn with_max_workflow_tokens(mut self, max: u64) -> Self {
147 self.max_workflow_tokens = max;
148 self
149 }
150
151 pub fn with_workflow_timeout_secs(mut self, secs: u64) -> Self {
162 self.workflow_timeout_secs = secs;
163 self
164 }
165}
166
167fn parse_env<T: std::str::FromStr>(suffix: &str, default: T) -> T {
168 std::env::var(format!("{ENV_PREFIX}{suffix}"))
169 .ok()
170 .and_then(|v| v.parse().ok())
171 .unwrap_or(default)
172}
173
174#[derive(Debug, Clone, PartialEq, Eq)]
189pub enum WorkflowRejection {
190 MaxDepthExceeded {
192 depth: u32,
194 max: u32,
196 },
197 CycleDetected {
199 target: String,
201 chain: Vec<String>,
203 },
204 MaxFanOutExceeded {
206 invocations: u32,
208 max: u32,
210 },
211 TokenBudgetExhausted {
213 used: u64,
215 max: u64,
217 },
218 WorkflowTimeout {
220 elapsed_secs: u64,
222 max: u64,
224 },
225 GuardUnavailable,
227}
228
229pub const WORKFLOW_GUARD_REJECTED_CODE: &str = "WORKFLOW_GUARD_REJECTED";
231
232impl fmt::Display for WorkflowRejection {
233 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
234 match self {
235 Self::MaxDepthExceeded { depth, max } => {
236 write!(f, "max call depth exceeded: {depth}/{max}")
237 }
238 Self::CycleDetected { target, chain } => {
239 write!(
240 f,
241 "cycle detected: workflow {target:?} already in chain {chain:?}"
242 )
243 }
244 Self::MaxFanOutExceeded { invocations, max } => {
245 write!(f, "max fan-out exceeded: {invocations}/{max} invocations")
246 }
247 Self::TokenBudgetExhausted { used, max } => {
248 write!(f, "token budget exhausted: {used}/{max}")
249 }
250 Self::WorkflowTimeout { elapsed_secs, max } => {
251 write!(f, "workflow timeout: {elapsed_secs}s/{max}s")
252 }
253 Self::GuardUnavailable => {
254 write!(f, "workflow guard unavailable -- failing closed")
255 }
256 }
257 }
258}
259
260impl std::error::Error for WorkflowRejection {}
261
262#[derive(Debug)]
280pub struct WorkflowGuardState {
281 depth: u32,
283 call_chain: Vec<String>,
285 total_invocations: u32,
287 total_tokens_used: u64,
289 started_at: Instant,
291}
292
293impl WorkflowGuardState {
294 pub fn new() -> Self {
306 Self {
307 depth: 0,
308 call_chain: Vec::new(),
309 total_invocations: 0,
310 total_tokens_used: 0,
311 started_at: Instant::now(),
312 }
313 }
314
315 pub fn depth(&self) -> u32 {
317 self.depth
318 }
319
320 pub fn total_invocations(&self) -> u32 {
322 self.total_invocations
323 }
324
325 pub fn total_tokens_used(&self) -> u64 {
327 self.total_tokens_used
328 }
329
330 pub fn call_chain(&self) -> &[String] {
332 &self.call_chain
333 }
334
335 pub fn elapsed_secs(&self) -> u64 {
337 self.started_at.elapsed().as_secs()
338 }
339
340 pub fn check(
362 &self,
363 config: &WorkflowGuardConfig,
364 target_workflow: &str,
365 ) -> Result<(), WorkflowRejection> {
366 if self.depth >= config.max_depth {
367 return Err(WorkflowRejection::MaxDepthExceeded {
368 depth: self.depth,
369 max: config.max_depth,
370 });
371 }
372
373 if self.call_chain.iter().any(|id| id == target_workflow) {
374 return Err(WorkflowRejection::CycleDetected {
375 target: target_workflow.to_string(),
376 chain: self.call_chain.clone(),
377 });
378 }
379
380 if self.total_invocations >= config.max_fan_out {
381 return Err(WorkflowRejection::MaxFanOutExceeded {
382 invocations: self.total_invocations,
383 max: config.max_fan_out,
384 });
385 }
386
387 if self.total_tokens_used >= config.max_workflow_tokens {
388 return Err(WorkflowRejection::TokenBudgetExhausted {
389 used: self.total_tokens_used,
390 max: config.max_workflow_tokens,
391 });
392 }
393
394 let elapsed = self.started_at.elapsed().as_secs();
395 if elapsed >= config.workflow_timeout_secs {
396 return Err(WorkflowRejection::WorkflowTimeout {
397 elapsed_secs: elapsed,
398 max: config.workflow_timeout_secs,
399 });
400 }
401
402 Ok(())
403 }
404
405 pub fn record_invocation(&mut self, target_workflow: &str) {
422 self.depth += 1;
423 self.total_invocations += 1;
424 self.call_chain.push(target_workflow.to_string());
425 }
426
427 pub fn record_return(&mut self) {
446 self.depth = self.depth.saturating_sub(1);
447 self.call_chain.pop();
448 }
449
450 pub fn record_tokens(
472 &mut self,
473 config: &WorkflowGuardConfig,
474 tokens_used: u64,
475 ) -> Result<u64, WorkflowRejection> {
476 self.total_tokens_used += tokens_used;
477 if self.total_tokens_used > config.max_workflow_tokens {
478 return Err(WorkflowRejection::TokenBudgetExhausted {
479 used: self.total_tokens_used,
480 max: config.max_workflow_tokens,
481 });
482 }
483 Ok(config.max_workflow_tokens - self.total_tokens_used)
484 }
485}
486
487impl Default for WorkflowGuardState {
488 fn default() -> Self {
489 Self::new()
490 }
491}
492
493pub type SharedGuardState = Arc<Mutex<WorkflowGuardState>>;
498
499pub fn new_shared_guard_state() -> SharedGuardState {
511 Arc::new(Mutex::new(WorkflowGuardState::new()))
512}
513
514#[cfg(test)]
515mod tests {
516 use super::*;
517
518 #[test]
519 fn default_config_has_documented_values() {
520 let config = WorkflowGuardConfig::default();
521 assert_eq!(config.max_depth, 5);
522 assert_eq!(config.max_fan_out, 20);
523 assert_eq!(config.max_workflow_tokens, 100_000);
524 assert_eq!(config.workflow_timeout_secs, 120);
525 }
526
527 #[test]
528 fn new_equals_default() {
529 assert_eq!(WorkflowGuardConfig::new(), WorkflowGuardConfig::default());
530 }
531
532 #[test]
533 fn config_from_env_reads_and_falls_back() {
534 unsafe {
537 std::env::set_var("IRONFLOW_GUARD_MAX_DEPTH", "10");
538 std::env::set_var("IRONFLOW_GUARD_MAX_FAN_OUT", "50");
539 std::env::set_var("IRONFLOW_GUARD_MAX_WORKFLOW_TOKENS", "200000");
540 std::env::set_var("IRONFLOW_GUARD_WORKFLOW_TIMEOUT_SECS", "300");
541 }
542
543 let config = WorkflowGuardConfig::from_env();
544 assert_eq!(config.max_depth, 10);
545 assert_eq!(config.max_fan_out, 50);
546 assert_eq!(config.max_workflow_tokens, 200_000);
547 assert_eq!(config.workflow_timeout_secs, 300);
548
549 unsafe {
551 std::env::remove_var("IRONFLOW_GUARD_MAX_DEPTH");
552 std::env::remove_var("IRONFLOW_GUARD_MAX_FAN_OUT");
553 std::env::remove_var("IRONFLOW_GUARD_MAX_WORKFLOW_TOKENS");
554 std::env::remove_var("IRONFLOW_GUARD_WORKFLOW_TIMEOUT_SECS");
555 }
556
557 let fallback = WorkflowGuardConfig::from_env();
558 assert_eq!(fallback.max_depth, 5);
559 assert_eq!(fallback.max_fan_out, 20);
560 }
561
562 #[test]
563 fn builder_methods_set_values() {
564 let config = WorkflowGuardConfig::new()
565 .with_max_depth(3)
566 .with_max_fan_out(10)
567 .with_max_workflow_tokens(50_000)
568 .with_workflow_timeout_secs(60);
569
570 assert_eq!(config.max_depth, 3);
571 assert_eq!(config.max_fan_out, 10);
572 assert_eq!(config.max_workflow_tokens, 50_000);
573 assert_eq!(config.workflow_timeout_secs, 60);
574 }
575
576 #[test]
577 fn check_rejects_max_depth_exceeded() {
578 let config = WorkflowGuardConfig::new().with_max_depth(2);
579 let mut state = WorkflowGuardState::new();
580 state.record_invocation("a");
581 state.record_invocation("b");
582
583 let err = state.check(&config, "c").unwrap_err();
584 assert!(matches!(
585 err,
586 WorkflowRejection::MaxDepthExceeded { depth: 2, max: 2 }
587 ));
588 }
589
590 #[test]
591 fn check_allows_within_depth_limit() {
592 let config = WorkflowGuardConfig::new().with_max_depth(2);
593 let mut state = WorkflowGuardState::new();
594 state.record_invocation("a");
595
596 assert!(state.check(&config, "b").is_ok());
597 }
598
599 #[test]
600 fn check_detects_direct_cycle() {
601 let config = WorkflowGuardConfig::default();
602 let mut state = WorkflowGuardState::new();
603 state.record_invocation("workflow-a");
604
605 let err = state.check(&config, "workflow-a").unwrap_err();
606 match err {
607 WorkflowRejection::CycleDetected { target, chain } => {
608 assert_eq!(target, "workflow-a");
609 assert_eq!(chain, vec!["workflow-a"]);
610 }
611 other => panic!("expected CycleDetected, got {other:?}"),
612 }
613 }
614
615 #[test]
616 fn check_detects_indirect_cycle() {
617 let config = WorkflowGuardConfig::default();
618 let mut state = WorkflowGuardState::new();
619 state.record_invocation("a");
620 state.record_invocation("b");
621 state.record_invocation("c");
622
623 let err = state.check(&config, "a").unwrap_err();
624 match err {
625 WorkflowRejection::CycleDetected { target, chain } => {
626 assert_eq!(target, "a");
627 assert_eq!(chain, vec!["a", "b", "c"]);
628 }
629 other => panic!("expected CycleDetected, got {other:?}"),
630 }
631 }
632
633 #[test]
634 fn check_allows_no_cycle() {
635 let config = WorkflowGuardConfig::default();
636 let mut state = WorkflowGuardState::new();
637 state.record_invocation("a");
638 state.record_invocation("b");
639
640 assert!(state.check(&config, "c").is_ok());
641 }
642
643 #[test]
644 fn check_rejects_fan_out_exceeded() {
645 let config = WorkflowGuardConfig::new().with_max_fan_out(2);
646 let mut state = WorkflowGuardState::new();
647 state.record_invocation("a");
648 state.record_return();
649 state.record_invocation("b");
650 state.record_return();
651
652 let err = state.check(&config, "c").unwrap_err();
653 assert!(matches!(
654 err,
655 WorkflowRejection::MaxFanOutExceeded {
656 invocations: 2,
657 max: 2
658 }
659 ));
660 }
661
662 #[test]
663 fn check_rejects_token_budget_exhausted() {
664 let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
665 let mut state = WorkflowGuardState::new();
666 let _ = state.record_tokens(&config, 100);
667
668 let err = state.check(&config, "a").unwrap_err();
669 assert!(matches!(
670 err,
671 WorkflowRejection::TokenBudgetExhausted {
672 used: 100,
673 max: 100
674 }
675 ));
676 }
677
678 #[test]
679 fn check_rejects_workflow_timeout() {
680 let state = WorkflowGuardState {
683 started_at: Instant::now() - std::time::Duration::from_secs(200),
684 ..WorkflowGuardState::new()
685 };
686 let config = WorkflowGuardConfig::new().with_workflow_timeout_secs(120);
687
688 let err = state.check(&config, "a").unwrap_err();
689 match err {
690 WorkflowRejection::WorkflowTimeout { elapsed_secs, max } => {
691 assert!(elapsed_secs >= 120);
692 assert_eq!(max, 120);
693 }
694 other => panic!("expected WorkflowTimeout, got {other:?}"),
695 }
696 }
697
698 #[test]
699 fn check_allows_within_timeout() {
700 let config = WorkflowGuardConfig::new().with_workflow_timeout_secs(120);
701 let state = WorkflowGuardState::new();
702 assert!(state.check(&config, "a").is_ok());
703 }
704
705 #[test]
706 fn record_invocation_updates_depth_fanout_chain() {
707 let mut state = WorkflowGuardState::new();
708
709 state.record_invocation("first");
710 assert_eq!(state.depth(), 1);
711 assert_eq!(state.total_invocations(), 1);
712 assert_eq!(state.call_chain(), &["first"]);
713
714 state.record_invocation("second");
715 assert_eq!(state.depth(), 2);
716 assert_eq!(state.total_invocations(), 2);
717 assert_eq!(state.call_chain(), &["first", "second"]);
718 }
719
720 #[test]
721 fn record_return_decrements_depth_and_pops_chain() {
722 let mut state = WorkflowGuardState::new();
723 state.record_invocation("a");
724 state.record_invocation("b");
725
726 state.record_return();
727 assert_eq!(state.depth(), 1);
728 assert_eq!(state.call_chain(), &["a"]);
729
730 state.record_return();
731 assert_eq!(state.depth(), 0);
732 assert!(state.call_chain().is_empty());
733 }
734
735 #[test]
736 fn record_return_saturates_at_zero() {
737 let mut state = WorkflowGuardState::new();
738 state.record_return();
739 assert_eq!(state.depth(), 0);
740 }
741
742 #[test]
743 fn fan_out_counts_even_after_return() {
744 let mut state = WorkflowGuardState::new();
745 state.record_invocation("a");
746 state.record_return();
747 state.record_invocation("b");
748 state.record_return();
749
750 assert_eq!(state.total_invocations(), 2);
751 assert_eq!(state.depth(), 0);
752 }
753
754 #[test]
755 fn record_tokens_accumulates_and_rejects_over_budget() {
756 let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
757 let mut state = WorkflowGuardState::new();
758
759 let remaining = state.record_tokens(&config, 40).unwrap();
760 assert_eq!(remaining, 60);
761
762 let remaining = state.record_tokens(&config, 40).unwrap();
763 assert_eq!(remaining, 20);
764
765 let err = state.record_tokens(&config, 30).unwrap_err();
766 assert!(matches!(
767 err,
768 WorkflowRejection::TokenBudgetExhausted {
769 used: 110,
770 max: 100
771 }
772 ));
773 }
774
775 #[test]
776 fn record_tokens_allows_exact_budget() {
777 let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
778 let mut state = WorkflowGuardState::new();
779
780 let remaining = state.record_tokens(&config, 100).unwrap();
781 assert_eq!(remaining, 0);
782 }
783
784 #[test]
785 fn rejection_display_max_depth() {
786 let r = WorkflowRejection::MaxDepthExceeded { depth: 6, max: 5 };
787 let msg = r.to_string();
788 assert!(msg.contains("max call depth exceeded"));
789 assert!(msg.contains("6/5"));
790 }
791
792 #[test]
793 fn rejection_display_cycle() {
794 let r = WorkflowRejection::CycleDetected {
795 target: "b".to_string(),
796 chain: vec!["a".to_string(), "b".to_string()],
797 };
798 let msg = r.to_string();
799 assert!(msg.contains("cycle detected"));
800 assert!(msg.contains("\"b\""));
801 }
802
803 #[test]
804 fn rejection_display_fan_out() {
805 let r = WorkflowRejection::MaxFanOutExceeded {
806 invocations: 21,
807 max: 20,
808 };
809 let msg = r.to_string();
810 assert!(msg.contains("max fan-out exceeded"));
811 assert!(msg.contains("21/20"));
812 }
813
814 #[test]
815 fn rejection_display_token_budget() {
816 let r = WorkflowRejection::TokenBudgetExhausted {
817 used: 100_001,
818 max: 100_000,
819 };
820 let msg = r.to_string();
821 assert!(msg.contains("token budget exhausted"));
822 assert!(msg.contains("100001/100000"));
823 }
824
825 #[test]
826 fn rejection_display_timeout() {
827 let r = WorkflowRejection::WorkflowTimeout {
828 elapsed_secs: 130,
829 max: 120,
830 };
831 let msg = r.to_string();
832 assert!(msg.contains("workflow timeout"));
833 assert!(msg.contains("130s/120s"));
834 }
835
836 #[test]
837 fn rejection_display_unavailable() {
838 let r = WorkflowRejection::GuardUnavailable;
839 assert!(r.to_string().contains("failing closed"));
840 }
841
842 #[test]
843 fn shared_guard_state_is_arc_mutex() {
844 let shared = new_shared_guard_state();
845 let mut state = shared.lock().unwrap();
846 state.record_invocation("test");
847 assert_eq!(state.depth(), 1);
848 }
849}