ironflow_engine/
guard.rs

1//! Workflow guard -- execution limits for agent workflows.
2//!
3//! Prevents runaway workflows: unbounded sub-workflow depth, uncontrolled
4//! fan-out, unlimited token consumption, and infinite execution time.
5//! Checked before every agent step invocation.
6//!
7//! # Examples
8//!
9//! ```
10//! use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState, WorkflowRejection};
11//!
12//! let config = WorkflowGuardConfig::default();
13//! assert_eq!(config.max_depth, 5);
14//!
15//! let state = WorkflowGuardState::new();
16//! assert!(state.check(&config, "agent-a").is_ok());
17//! ```
18
19use std::fmt;
20use std::sync::atomic::{AtomicU32, AtomicU64, Ordering};
21use std::sync::{Arc, Mutex};
22use std::time::Instant;
23
24/// Environment variable prefix for guard configuration.
25const ENV_PREFIX: &str = "IRONFLOW_GUARD_";
26
27/// Configurable execution limits for a workflow run.
28///
29/// Each limit has a sensible default. Limits can be set per-handler via
30/// [`WorkflowHandler::guard_config`](crate::handler::WorkflowHandler::guard_config)
31/// or globally on the [`Engine`](crate::engine::Engine).
32///
33/// # Examples
34///
35/// ```
36/// use ironflow_engine::guard::WorkflowGuardConfig;
37///
38/// let config = WorkflowGuardConfig::new()
39///     .with_max_depth(3)
40///     .with_max_fan_out(10);
41///
42/// assert_eq!(config.max_depth, 3);
43/// assert_eq!(config.max_fan_out, 10);
44/// assert_eq!(config.max_workflow_tokens, 100_000);
45/// ```
46#[derive(Debug, Clone, PartialEq, Eq)]
47pub struct WorkflowGuardConfig {
48    /// Maximum depth of nested sub-workflows (default 5).
49    pub max_depth: u32,
50    /// Maximum total number of workflow invocations (default 20).
51    pub max_fan_out: u32,
52    /// Cumulative token budget across all agent steps (default 100,000).
53    pub max_workflow_tokens: u64,
54    /// Global timeout in seconds (default 120).
55    pub workflow_timeout_secs: u64,
56}
57
58impl Default for WorkflowGuardConfig {
59    fn default() -> Self {
60        Self {
61            max_depth: 5,
62            max_fan_out: 20,
63            max_workflow_tokens: 100_000,
64            workflow_timeout_secs: 120,
65        }
66    }
67}
68
69impl WorkflowGuardConfig {
70    /// Create a configuration with default limits.
71    ///
72    /// # Examples
73    ///
74    /// ```
75    /// use ironflow_engine::guard::WorkflowGuardConfig;
76    ///
77    /// assert_eq!(WorkflowGuardConfig::new(), WorkflowGuardConfig::default());
78    /// ```
79    pub fn new() -> Self {
80        Self::default()
81    }
82
83    /// Load configuration from environment variables.
84    ///
85    /// Reads `IRONFLOW_GUARD_MAX_DEPTH`, `IRONFLOW_GUARD_MAX_FAN_OUT`,
86    /// `IRONFLOW_GUARD_MAX_WORKFLOW_TOKENS`, and `IRONFLOW_GUARD_WORKFLOW_TIMEOUT_SECS`.
87    /// Missing or unparseable variables fall back to defaults.
88    ///
89    /// # Examples
90    ///
91    /// ```
92    /// use ironflow_engine::guard::WorkflowGuardConfig;
93    ///
94    /// let config = WorkflowGuardConfig::from_env();
95    /// // Without env vars set, returns defaults.
96    /// assert_eq!(config.max_depth, 5);
97    /// ```
98    pub fn from_env() -> Self {
99        Self {
100            max_depth: parse_env("MAX_DEPTH", 5),
101            max_fan_out: parse_env("MAX_FAN_OUT", 20),
102            max_workflow_tokens: parse_env("MAX_WORKFLOW_TOKENS", 100_000),
103            workflow_timeout_secs: parse_env("WORKFLOW_TIMEOUT_SECS", 120),
104        }
105    }
106
107    /// Set the maximum sub-workflow nesting depth.
108    ///
109    /// # Examples
110    ///
111    /// ```
112    /// use ironflow_engine::guard::WorkflowGuardConfig;
113    ///
114    /// let config = WorkflowGuardConfig::new().with_max_depth(3);
115    /// assert_eq!(config.max_depth, 3);
116    /// ```
117    pub fn with_max_depth(mut self, max_depth: u32) -> Self {
118        self.max_depth = max_depth;
119        self
120    }
121
122    /// Set the maximum total workflow invocations.
123    ///
124    /// # Examples
125    ///
126    /// ```
127    /// use ironflow_engine::guard::WorkflowGuardConfig;
128    ///
129    /// let config = WorkflowGuardConfig::new().with_max_fan_out(10);
130    /// assert_eq!(config.max_fan_out, 10);
131    /// ```
132    pub fn with_max_fan_out(mut self, max_fan_out: u32) -> Self {
133        self.max_fan_out = max_fan_out;
134        self
135    }
136
137    /// Set the cumulative token budget.
138    ///
139    /// # Examples
140    ///
141    /// ```
142    /// use ironflow_engine::guard::WorkflowGuardConfig;
143    ///
144    /// let config = WorkflowGuardConfig::new().with_max_workflow_tokens(50_000);
145    /// assert_eq!(config.max_workflow_tokens, 50_000);
146    /// ```
147    pub fn with_max_workflow_tokens(mut self, max: u64) -> Self {
148        self.max_workflow_tokens = max;
149        self
150    }
151
152    /// Set the global timeout in seconds.
153    ///
154    /// # Examples
155    ///
156    /// ```
157    /// use ironflow_engine::guard::WorkflowGuardConfig;
158    ///
159    /// let config = WorkflowGuardConfig::new().with_workflow_timeout_secs(60);
160    /// assert_eq!(config.workflow_timeout_secs, 60);
161    /// ```
162    pub fn with_workflow_timeout_secs(mut self, secs: u64) -> Self {
163        self.workflow_timeout_secs = secs;
164        self
165    }
166}
167
168fn parse_env<T: std::str::FromStr>(suffix: &str, default: T) -> T {
169    std::env::var(format!("{ENV_PREFIX}{suffix}"))
170        .ok()
171        .and_then(|v| v.parse().ok())
172        .unwrap_or(default)
173}
174
175/// Reason a workflow invocation was rejected by the guard.
176///
177/// Each variant carries the observed value and the configured limit
178/// for diagnostics.
179///
180/// # Examples
181///
182/// ```
183/// use ironflow_engine::guard::WorkflowRejection;
184///
185/// let r = WorkflowRejection::MaxDepthExceeded { depth: 6, max: 5 };
186/// assert!(r.to_string().contains("6"));
187/// assert!(r.to_string().contains("5"));
188/// ```
189#[derive(Debug, Clone, PartialEq, Eq)]
190pub enum WorkflowRejection {
191    /// The sub-workflow nesting depth exceeds the configured maximum.
192    MaxDepthExceeded {
193        /// Current depth.
194        depth: u32,
195        /// Configured limit.
196        max: u32,
197    },
198    /// The target workflow is already in the call chain (direct or indirect cycle).
199    CycleDetected {
200        /// The workflow that would create a cycle.
201        target: String,
202        /// The current call chain.
203        chain: Vec<String>,
204    },
205    /// The total number of workflow invocations exceeds the configured maximum.
206    MaxFanOutExceeded {
207        /// Current invocation count.
208        invocations: u32,
209        /// Configured limit.
210        max: u32,
211    },
212    /// The cumulative token usage exceeds the configured budget.
213    TokenBudgetExhausted {
214        /// Tokens already consumed.
215        used: u64,
216        /// Configured limit.
217        max: u64,
218    },
219    /// The workflow has exceeded its global timeout.
220    WorkflowTimeout {
221        /// Seconds elapsed since the workflow started.
222        elapsed_secs: u64,
223        /// Configured limit.
224        max: u64,
225    },
226    /// The guard state is unavailable (fail closed).
227    GuardUnavailable,
228}
229
230/// Business error code for guard rejections.
231pub const WORKFLOW_GUARD_REJECTED_CODE: &str = "WORKFLOW_GUARD_REJECTED";
232
233impl fmt::Display for WorkflowRejection {
234    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235        match self {
236            Self::MaxDepthExceeded { depth, max } => {
237                write!(f, "max call depth exceeded: {depth}/{max}")
238            }
239            Self::CycleDetected { target, chain } => {
240                write!(
241                    f,
242                    "cycle detected: workflow {target:?} already in chain {chain:?}"
243                )
244            }
245            Self::MaxFanOutExceeded { invocations, max } => {
246                write!(f, "max fan-out exceeded: {invocations}/{max} invocations")
247            }
248            Self::TokenBudgetExhausted { used, max } => {
249                write!(f, "token budget exhausted: {used}/{max}")
250            }
251            Self::WorkflowTimeout { elapsed_secs, max } => {
252                write!(f, "workflow timeout: {elapsed_secs}s/{max}s")
253            }
254            Self::GuardUnavailable => {
255                write!(f, "workflow guard unavailable -- failing closed")
256            }
257        }
258    }
259}
260
261impl std::error::Error for WorkflowRejection {}
262
263/// Mutable runtime state tracked by the guard during a workflow execution.
264///
265/// Shared between parent and child workflows via `Arc<Mutex<_>>` so that
266/// limits are enforced globally across the entire run tree.
267///
268/// # Examples
269///
270/// ```
271/// use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState};
272///
273/// let mut state = WorkflowGuardState::new();
274/// let config = WorkflowGuardConfig::new().with_max_depth(2);
275///
276/// state.record_invocation("workflow-a");
277/// state.record_invocation("workflow-b");
278/// assert!(state.check(&config, "workflow-c").is_err());
279/// ```
280#[derive(Debug)]
281pub struct WorkflowGuardState {
282    /// Current sub-workflow nesting depth.
283    depth: u32,
284    /// Call chain for cycle detection (stack of workflow names).
285    call_chain: Vec<String>,
286    /// Counters of the whole run tree, shared with every branch.
287    totals: Arc<GuardTotals>,
288    /// When the root workflow started.
289    started_at: Instant,
290}
291
292/// Counters a run tree shares across the branches of its concurrent
293/// sub-workflows.
294#[derive(Debug, Default)]
295struct GuardTotals {
296    /// Total number of workflow invocations in this run tree.
297    invocations: AtomicU32,
298    /// Cumulative tokens consumed across all agent steps.
299    tokens_used: AtomicU64,
300}
301
302impl WorkflowGuardState {
303    /// Create a fresh guard state for a new workflow execution.
304    ///
305    /// # Examples
306    ///
307    /// ```
308    /// use ironflow_engine::guard::WorkflowGuardState;
309    ///
310    /// let state = WorkflowGuardState::new();
311    /// assert_eq!(state.depth(), 0);
312    /// assert_eq!(state.total_invocations(), 0);
313    /// ```
314    pub fn new() -> Self {
315        Self {
316            depth: 0,
317            call_chain: Vec::new(),
318            totals: Arc::new(GuardTotals::default()),
319            started_at: Instant::now(),
320        }
321    }
322
323    /// The state of a `target_workflow` child running concurrently with its
324    /// siblings, its invocation already claimed by
325    /// [`enter_wave`](Self::enter_wave).
326    ///
327    /// The branch is one level below this state, `target_workflow` on top of
328    /// its call chain, and keeps its own depth and chain from then on, so
329    /// siblings never push to, or pop from, each other's chain. Invocations,
330    /// tokens and the start time stay shared with the whole run tree.
331    pub(crate) fn branch_into(&self, target_workflow: &str) -> Self {
332        let mut call_chain = self.call_chain.clone();
333        call_chain.push(target_workflow.to_string());
334        Self {
335            depth: self.depth + 1,
336            call_chain,
337            totals: self.totals.clone(),
338            started_at: self.started_at,
339        }
340    }
341
342    /// Current sub-workflow nesting depth.
343    pub fn depth(&self) -> u32 {
344        self.depth
345    }
346
347    /// Total workflow invocations so far.
348    pub fn total_invocations(&self) -> u32 {
349        self.totals.invocations.load(Ordering::Relaxed)
350    }
351
352    /// Total tokens consumed so far.
353    pub fn total_tokens_used(&self) -> u64 {
354        self.totals.tokens_used.load(Ordering::Relaxed)
355    }
356
357    /// The current call chain (workflow name stack).
358    pub fn call_chain(&self) -> &[String] {
359        &self.call_chain
360    }
361
362    /// Seconds elapsed since the workflow started.
363    pub fn elapsed_secs(&self) -> u64 {
364        self.started_at.elapsed().as_secs()
365    }
366
367    /// Read-only check: would invoking `target_workflow` violate any limit?
368    ///
369    /// Does not mutate state. Call this before
370    /// [`record_invocation`](Self::record_invocation).
371    ///
372    /// # Errors
373    ///
374    /// Returns the specific [`WorkflowRejection`] variant that would be
375    /// violated.
376    ///
377    /// # Examples
378    ///
379    /// ```
380    /// use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState, WorkflowRejection};
381    ///
382    /// let state = WorkflowGuardState::new();
383    /// let config = WorkflowGuardConfig::new().with_max_depth(0);
384    ///
385    /// let err = state.check(&config, "child").unwrap_err();
386    /// assert!(matches!(err, WorkflowRejection::MaxDepthExceeded { .. }));
387    /// ```
388    pub fn check(
389        &self,
390        config: &WorkflowGuardConfig,
391        target_workflow: &str,
392    ) -> Result<(), WorkflowRejection> {
393        if self.depth >= config.max_depth {
394            return Err(WorkflowRejection::MaxDepthExceeded {
395                depth: self.depth,
396                max: config.max_depth,
397            });
398        }
399
400        if self.call_chain.iter().any(|id| id == target_workflow) {
401            return Err(WorkflowRejection::CycleDetected {
402                target: target_workflow.to_string(),
403                chain: self.call_chain.clone(),
404            });
405        }
406
407        let invocations = self.total_invocations();
408        if invocations >= config.max_fan_out {
409            return Err(WorkflowRejection::MaxFanOutExceeded {
410                invocations,
411                max: config.max_fan_out,
412            });
413        }
414
415        let used = self.total_tokens_used();
416        if used >= config.max_workflow_tokens {
417            return Err(WorkflowRejection::TokenBudgetExhausted {
418                used,
419                max: config.max_workflow_tokens,
420            });
421        }
422
423        let elapsed = self.started_at.elapsed().as_secs();
424        if elapsed >= config.workflow_timeout_secs {
425            return Err(WorkflowRejection::WorkflowTimeout {
426                elapsed_secs: elapsed,
427                max: config.workflow_timeout_secs,
428            });
429        }
430
431        Ok(())
432    }
433
434    /// [`check`](Self::check) `target_workflow`, then claim its invocation
435    /// and record it like [`record_invocation`](Self::record_invocation).
436    ///
437    /// The fan-out slot is claimed atomically: branches running concurrently
438    /// never cross `max_fan_out` together.
439    ///
440    /// # Errors
441    ///
442    /// Returns the [`WorkflowRejection`] the invocation would violate; nothing
443    /// is recorded then.
444    pub(crate) fn enter(
445        &mut self,
446        config: &WorkflowGuardConfig,
447        target_workflow: &str,
448    ) -> Result<(), WorkflowRejection> {
449        self.enter_wave(config, target_workflow, 1)?;
450        self.depth += 1;
451        self.call_chain.push(target_workflow.to_string());
452        Ok(())
453    }
454
455    /// [`check`](Self::check) a wave of `count` concurrent invocations of
456    /// `target_workflow`, then claim all of them at once. Each child then
457    /// runs on its own [`branch_into`](Self::branch_into).
458    ///
459    /// # Errors
460    ///
461    /// Returns the [`WorkflowRejection`] the wave would violate; a wave that
462    /// does not fit whole is refused whole and claims nothing.
463    pub(crate) fn enter_wave(
464        &self,
465        config: &WorkflowGuardConfig,
466        target_workflow: &str,
467        count: u32,
468    ) -> Result<(), WorkflowRejection> {
469        self.check(config, target_workflow)?;
470        self.totals
471            .invocations
472            .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |taken| {
473                taken
474                    .checked_add(count)
475                    .filter(|total| *total <= config.max_fan_out)
476            })
477            .map(|_| ())
478            .map_err(|taken| WorkflowRejection::MaxFanOutExceeded {
479                invocations: taken.saturating_add(count),
480                max: config.max_fan_out,
481            })
482    }
483
484    /// Record that a sub-workflow invocation is starting.
485    ///
486    /// Increments depth and fan-out counter, and pushes the workflow name
487    /// onto the call chain for cycle detection.
488    ///
489    /// # Examples
490    ///
491    /// ```
492    /// use ironflow_engine::guard::WorkflowGuardState;
493    ///
494    /// let mut state = WorkflowGuardState::new();
495    /// state.record_invocation("child-workflow");
496    /// assert_eq!(state.depth(), 1);
497    /// assert_eq!(state.total_invocations(), 1);
498    /// assert_eq!(state.call_chain(), &["child-workflow"]);
499    /// ```
500    pub fn record_invocation(&mut self, target_workflow: &str) {
501        self.depth += 1;
502        self.totals.invocations.fetch_add(1, Ordering::Relaxed);
503        self.call_chain.push(target_workflow.to_string());
504    }
505
506    /// Record that a sub-workflow invocation has returned.
507    ///
508    /// Decrements depth and pops the last entry from the call chain.
509    /// Safe to call even on failure paths (the guard must never leak depth).
510    ///
511    /// # Examples
512    ///
513    /// ```
514    /// use ironflow_engine::guard::WorkflowGuardState;
515    ///
516    /// let mut state = WorkflowGuardState::new();
517    /// state.record_invocation("child");
518    /// assert_eq!(state.depth(), 1);
519    ///
520    /// state.record_return();
521    /// assert_eq!(state.depth(), 0);
522    /// assert!(state.call_chain().is_empty());
523    /// ```
524    pub fn record_return(&mut self) {
525        self.depth = self.depth.saturating_sub(1);
526        self.call_chain.pop();
527    }
528
529    /// Record tokens consumed by an agent step.
530    ///
531    /// Returns the remaining token budget, or an error if the budget is
532    /// now exhausted.
533    ///
534    /// # Errors
535    ///
536    /// Returns [`WorkflowRejection::TokenBudgetExhausted`] when the
537    /// cumulative usage exceeds the configured maximum.
538    ///
539    /// # Examples
540    ///
541    /// ```
542    /// use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState, WorkflowRejection};
543    ///
544    /// let mut state = WorkflowGuardState::new();
545    /// let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
546    ///
547    /// assert!(state.record_tokens(&config, 50).is_ok());
548    /// assert!(state.record_tokens(&config, 60).is_err());
549    /// ```
550    pub fn record_tokens(
551        &mut self,
552        config: &WorkflowGuardConfig,
553        tokens_used: u64,
554    ) -> Result<u64, WorkflowRejection> {
555        let used = self
556            .totals
557            .tokens_used
558            .fetch_add(tokens_used, Ordering::Relaxed)
559            .saturating_add(tokens_used);
560        if used > config.max_workflow_tokens {
561            return Err(WorkflowRejection::TokenBudgetExhausted {
562                used,
563                max: config.max_workflow_tokens,
564            });
565        }
566        Ok(config.max_workflow_tokens - used)
567    }
568}
569
570impl Default for WorkflowGuardState {
571    fn default() -> Self {
572        Self::new()
573    }
574}
575
576/// Thread-safe handle to a shared [`WorkflowGuardState`].
577///
578/// Passed from parent to child workflows so that limits are enforced
579/// globally across the entire run tree.
580pub type SharedGuardState = Arc<Mutex<WorkflowGuardState>>;
581
582/// Create a new shared guard state for a workflow run.
583///
584/// # Examples
585///
586/// ```
587/// use ironflow_engine::guard::new_shared_guard_state;
588///
589/// let state = new_shared_guard_state();
590/// let locked = state.lock().unwrap();
591/// assert_eq!(locked.depth(), 0);
592/// ```
593pub fn new_shared_guard_state() -> SharedGuardState {
594    Arc::new(Mutex::new(WorkflowGuardState::new()))
595}
596
597#[cfg(test)]
598mod tests {
599    use super::*;
600
601    #[test]
602    fn default_config_has_documented_values() {
603        let config = WorkflowGuardConfig::default();
604        assert_eq!(config.max_depth, 5);
605        assert_eq!(config.max_fan_out, 20);
606        assert_eq!(config.max_workflow_tokens, 100_000);
607        assert_eq!(config.workflow_timeout_secs, 120);
608    }
609
610    #[test]
611    fn new_equals_default() {
612        assert_eq!(WorkflowGuardConfig::new(), WorkflowGuardConfig::default());
613    }
614
615    #[test]
616    fn config_from_env_reads_and_falls_back() {
617        // SAFETY: test-only, single-threaded access to env vars.
618        // Combined into one test to avoid races between parallel tests.
619        unsafe {
620            std::env::set_var("IRONFLOW_GUARD_MAX_DEPTH", "10");
621            std::env::set_var("IRONFLOW_GUARD_MAX_FAN_OUT", "50");
622            std::env::set_var("IRONFLOW_GUARD_MAX_WORKFLOW_TOKENS", "200000");
623            std::env::set_var("IRONFLOW_GUARD_WORKFLOW_TIMEOUT_SECS", "300");
624        }
625
626        let config = WorkflowGuardConfig::from_env();
627        assert_eq!(config.max_depth, 10);
628        assert_eq!(config.max_fan_out, 50);
629        assert_eq!(config.max_workflow_tokens, 200_000);
630        assert_eq!(config.workflow_timeout_secs, 300);
631
632        // SAFETY: test-only cleanup.
633        unsafe {
634            std::env::remove_var("IRONFLOW_GUARD_MAX_DEPTH");
635            std::env::remove_var("IRONFLOW_GUARD_MAX_FAN_OUT");
636            std::env::remove_var("IRONFLOW_GUARD_MAX_WORKFLOW_TOKENS");
637            std::env::remove_var("IRONFLOW_GUARD_WORKFLOW_TIMEOUT_SECS");
638        }
639
640        let fallback = WorkflowGuardConfig::from_env();
641        assert_eq!(fallback.max_depth, 5);
642        assert_eq!(fallback.max_fan_out, 20);
643    }
644
645    #[test]
646    fn builder_methods_set_values() {
647        let config = WorkflowGuardConfig::new()
648            .with_max_depth(3)
649            .with_max_fan_out(10)
650            .with_max_workflow_tokens(50_000)
651            .with_workflow_timeout_secs(60);
652
653        assert_eq!(config.max_depth, 3);
654        assert_eq!(config.max_fan_out, 10);
655        assert_eq!(config.max_workflow_tokens, 50_000);
656        assert_eq!(config.workflow_timeout_secs, 60);
657    }
658
659    #[test]
660    fn check_rejects_max_depth_exceeded() {
661        let config = WorkflowGuardConfig::new().with_max_depth(2);
662        let mut state = WorkflowGuardState::new();
663        state.record_invocation("a");
664        state.record_invocation("b");
665
666        let err = state.check(&config, "c").unwrap_err();
667        assert!(matches!(
668            err,
669            WorkflowRejection::MaxDepthExceeded { depth: 2, max: 2 }
670        ));
671    }
672
673    #[test]
674    fn check_allows_within_depth_limit() {
675        let config = WorkflowGuardConfig::new().with_max_depth(2);
676        let mut state = WorkflowGuardState::new();
677        state.record_invocation("a");
678
679        assert!(state.check(&config, "b").is_ok());
680    }
681
682    #[test]
683    fn check_detects_direct_cycle() {
684        let config = WorkflowGuardConfig::default();
685        let mut state = WorkflowGuardState::new();
686        state.record_invocation("workflow-a");
687
688        let err = state.check(&config, "workflow-a").unwrap_err();
689        match err {
690            WorkflowRejection::CycleDetected { target, chain } => {
691                assert_eq!(target, "workflow-a");
692                assert_eq!(chain, vec!["workflow-a"]);
693            }
694            other => panic!("expected CycleDetected, got {other:?}"),
695        }
696    }
697
698    #[test]
699    fn check_detects_indirect_cycle() {
700        let config = WorkflowGuardConfig::default();
701        let mut state = WorkflowGuardState::new();
702        state.record_invocation("a");
703        state.record_invocation("b");
704        state.record_invocation("c");
705
706        let err = state.check(&config, "a").unwrap_err();
707        match err {
708            WorkflowRejection::CycleDetected { target, chain } => {
709                assert_eq!(target, "a");
710                assert_eq!(chain, vec!["a", "b", "c"]);
711            }
712            other => panic!("expected CycleDetected, got {other:?}"),
713        }
714    }
715
716    #[test]
717    fn check_allows_no_cycle() {
718        let config = WorkflowGuardConfig::default();
719        let mut state = WorkflowGuardState::new();
720        state.record_invocation("a");
721        state.record_invocation("b");
722
723        assert!(state.check(&config, "c").is_ok());
724    }
725
726    #[test]
727    fn check_rejects_fan_out_exceeded() {
728        let config = WorkflowGuardConfig::new().with_max_fan_out(2);
729        let mut state = WorkflowGuardState::new();
730        state.record_invocation("a");
731        state.record_return();
732        state.record_invocation("b");
733        state.record_return();
734
735        let err = state.check(&config, "c").unwrap_err();
736        assert!(matches!(
737            err,
738            WorkflowRejection::MaxFanOutExceeded {
739                invocations: 2,
740                max: 2
741            }
742        ));
743    }
744
745    #[test]
746    fn check_rejects_token_budget_exhausted() {
747        let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
748        let mut state = WorkflowGuardState::new();
749        let _ = state.record_tokens(&config, 100);
750
751        let err = state.check(&config, "a").unwrap_err();
752        assert!(matches!(
753            err,
754            WorkflowRejection::TokenBudgetExhausted {
755                used: 100,
756                max: 100
757            }
758        ));
759    }
760
761    #[test]
762    fn check_rejects_workflow_timeout() {
763        // elapsed().as_secs() truncates to whole seconds, so we use a state
764        // whose started_at we can set to the past via the test-only constructor.
765        let state = WorkflowGuardState {
766            started_at: Instant::now() - std::time::Duration::from_secs(200),
767            ..WorkflowGuardState::new()
768        };
769        let config = WorkflowGuardConfig::new().with_workflow_timeout_secs(120);
770
771        let err = state.check(&config, "a").unwrap_err();
772        match err {
773            WorkflowRejection::WorkflowTimeout { elapsed_secs, max } => {
774                assert!(elapsed_secs >= 120);
775                assert_eq!(max, 120);
776            }
777            other => panic!("expected WorkflowTimeout, got {other:?}"),
778        }
779    }
780
781    #[test]
782    fn check_allows_within_timeout() {
783        let config = WorkflowGuardConfig::new().with_workflow_timeout_secs(120);
784        let state = WorkflowGuardState::new();
785        assert!(state.check(&config, "a").is_ok());
786    }
787
788    #[test]
789    fn record_invocation_updates_depth_fanout_chain() {
790        let mut state = WorkflowGuardState::new();
791
792        state.record_invocation("first");
793        assert_eq!(state.depth(), 1);
794        assert_eq!(state.total_invocations(), 1);
795        assert_eq!(state.call_chain(), &["first"]);
796
797        state.record_invocation("second");
798        assert_eq!(state.depth(), 2);
799        assert_eq!(state.total_invocations(), 2);
800        assert_eq!(state.call_chain(), &["first", "second"]);
801    }
802
803    #[test]
804    fn record_return_decrements_depth_and_pops_chain() {
805        let mut state = WorkflowGuardState::new();
806        state.record_invocation("a");
807        state.record_invocation("b");
808
809        state.record_return();
810        assert_eq!(state.depth(), 1);
811        assert_eq!(state.call_chain(), &["a"]);
812
813        state.record_return();
814        assert_eq!(state.depth(), 0);
815        assert!(state.call_chain().is_empty());
816    }
817
818    #[test]
819    fn record_return_saturates_at_zero() {
820        let mut state = WorkflowGuardState::new();
821        state.record_return();
822        assert_eq!(state.depth(), 0);
823    }
824
825    #[test]
826    fn fan_out_counts_even_after_return() {
827        let mut state = WorkflowGuardState::new();
828        state.record_invocation("a");
829        state.record_return();
830        state.record_invocation("b");
831        state.record_return();
832
833        assert_eq!(state.total_invocations(), 2);
834        assert_eq!(state.depth(), 0);
835    }
836
837    #[test]
838    fn record_tokens_accumulates_and_rejects_over_budget() {
839        let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
840        let mut state = WorkflowGuardState::new();
841
842        let remaining = state.record_tokens(&config, 40).unwrap();
843        assert_eq!(remaining, 60);
844
845        let remaining = state.record_tokens(&config, 40).unwrap();
846        assert_eq!(remaining, 20);
847
848        let err = state.record_tokens(&config, 30).unwrap_err();
849        assert!(matches!(
850            err,
851            WorkflowRejection::TokenBudgetExhausted {
852                used: 110,
853                max: 100
854            }
855        ));
856    }
857
858    #[test]
859    fn record_tokens_allows_exact_budget() {
860        let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
861        let mut state = WorkflowGuardState::new();
862
863        let remaining = state.record_tokens(&config, 100).unwrap();
864        assert_eq!(remaining, 0);
865    }
866
867    #[test]
868    fn rejection_display_max_depth() {
869        let r = WorkflowRejection::MaxDepthExceeded { depth: 6, max: 5 };
870        let msg = r.to_string();
871        assert!(msg.contains("max call depth exceeded"));
872        assert!(msg.contains("6/5"));
873    }
874
875    #[test]
876    fn rejection_display_cycle() {
877        let r = WorkflowRejection::CycleDetected {
878            target: "b".to_string(),
879            chain: vec!["a".to_string(), "b".to_string()],
880        };
881        let msg = r.to_string();
882        assert!(msg.contains("cycle detected"));
883        assert!(msg.contains("\"b\""));
884    }
885
886    #[test]
887    fn rejection_display_fan_out() {
888        let r = WorkflowRejection::MaxFanOutExceeded {
889            invocations: 21,
890            max: 20,
891        };
892        let msg = r.to_string();
893        assert!(msg.contains("max fan-out exceeded"));
894        assert!(msg.contains("21/20"));
895    }
896
897    #[test]
898    fn rejection_display_token_budget() {
899        let r = WorkflowRejection::TokenBudgetExhausted {
900            used: 100_001,
901            max: 100_000,
902        };
903        let msg = r.to_string();
904        assert!(msg.contains("token budget exhausted"));
905        assert!(msg.contains("100001/100000"));
906    }
907
908    #[test]
909    fn rejection_display_timeout() {
910        let r = WorkflowRejection::WorkflowTimeout {
911            elapsed_secs: 130,
912            max: 120,
913        };
914        let msg = r.to_string();
915        assert!(msg.contains("workflow timeout"));
916        assert!(msg.contains("130s/120s"));
917    }
918
919    #[test]
920    fn rejection_display_unavailable() {
921        let r = WorkflowRejection::GuardUnavailable;
922        assert!(r.to_string().contains("failing closed"));
923    }
924
925    #[test]
926    fn shared_guard_state_is_arc_mutex() {
927        let shared = new_shared_guard_state();
928        let mut state = shared.lock().unwrap();
929        state.record_invocation("test");
930        assert_eq!(state.depth(), 1);
931    }
932
933    #[test]
934    fn a_branch_keeps_its_own_chain_and_shares_the_totals() {
935        let mut root = WorkflowGuardState::new();
936        root.record_invocation("parent");
937        let config = WorkflowGuardConfig::new().with_max_workflow_tokens(10);
938        root.enter_wave(&config, "child", 2).expect("room for two");
939        let mut left = root.branch_into("child");
940        let mut right = root.branch_into("child");
941        left.enter(&config, "leaf").expect("no cycle on the left");
942
943        assert_eq!(left.call_chain(), ["parent", "child", "leaf"]);
944        assert_eq!(right.call_chain(), ["parent", "child"]);
945        assert_eq!(right.depth(), 2);
946        assert_eq!(root.call_chain(), ["parent"]);
947        assert_eq!(root.total_invocations(), 4);
948
949        assert_eq!(left.record_tokens(&config, 6), Ok(4));
950        assert!(
951            right.record_tokens(&config, 6).is_err(),
952            "tokens are shared"
953        );
954    }
955
956    #[test]
957    fn enter_wave_claims_room_for_every_child_or_none() {
958        let mut state = WorkflowGuardState::new();
959        state.record_invocation("earlier");
960        state.record_return();
961        let config = WorkflowGuardConfig::new().with_max_fan_out(3);
962
963        assert_eq!(
964            state.enter_wave(&config, "child", 3),
965            Err(WorkflowRejection::MaxFanOutExceeded {
966                invocations: 4,
967                max: 3
968            })
969        );
970        assert_eq!(
971            state.total_invocations(),
972            1,
973            "a refused wave claims nothing"
974        );
975        assert_eq!(state.enter_wave(&config, "child", 2), Ok(()));
976        assert_eq!(state.total_invocations(), 3);
977    }
978
979    #[test]
980    fn enter_claims_the_last_slot_once_across_branches() {
981        let root = WorkflowGuardState::new();
982        let config = WorkflowGuardConfig::new().with_max_fan_out(3);
983        root.enter_wave(&config, "child", 2).expect("room for two");
984        let mut left = root.branch_into("child");
985        let mut right = root.branch_into("child");
986
987        assert_eq!(left.enter(&config, "leaf"), Ok(()));
988        assert_eq!(
989            right.enter(&config, "leaf"),
990            Err(WorkflowRejection::MaxFanOutExceeded {
991                invocations: 3,
992                max: 3
993            })
994        );
995        assert_eq!(right.depth(), 1, "a refused invocation records nothing");
996        assert_eq!(right.call_chain(), ["child"]);
997    }
998}
ironflow_engine/guard.rs

ironflow_engine/
guard.rs
