Skip to main content

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::{Arc, Mutex};
21use std::time::Instant;
22
23/// Environment variable prefix for guard configuration.
24const ENV_PREFIX: &str = "IRONFLOW_GUARD_";
25
26/// Configurable execution limits for a workflow run.
27///
28/// Each limit has a sensible default. Limits can be set per-handler via
29/// [`WorkflowHandler::guard_config`](crate::handler::WorkflowHandler::guard_config)
30/// or globally on the [`Engine`](crate::engine::Engine).
31///
32/// # Examples
33///
34/// ```
35/// use ironflow_engine::guard::WorkflowGuardConfig;
36///
37/// let config = WorkflowGuardConfig::new()
38///     .with_max_depth(3)
39///     .with_max_fan_out(10);
40///
41/// assert_eq!(config.max_depth, 3);
42/// assert_eq!(config.max_fan_out, 10);
43/// assert_eq!(config.max_workflow_tokens, 100_000);
44/// ```
45#[derive(Debug, Clone, PartialEq, Eq)]
46pub struct WorkflowGuardConfig {
47    /// Maximum depth of nested sub-workflows (default 5).
48    pub max_depth: u32,
49    /// Maximum total number of workflow invocations (default 20).
50    pub max_fan_out: u32,
51    /// Cumulative token budget across all agent steps (default 100,000).
52    pub max_workflow_tokens: u64,
53    /// Global timeout in seconds (default 120).
54    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    /// Create a configuration with default limits.
70    ///
71    /// # Examples
72    ///
73    /// ```
74    /// use ironflow_engine::guard::WorkflowGuardConfig;
75    ///
76    /// assert_eq!(WorkflowGuardConfig::new(), WorkflowGuardConfig::default());
77    /// ```
78    pub fn new() -> Self {
79        Self::default()
80    }
81
82    /// Load configuration from environment variables.
83    ///
84    /// Reads `IRONFLOW_GUARD_MAX_DEPTH`, `IRONFLOW_GUARD_MAX_FAN_OUT`,
85    /// `IRONFLOW_GUARD_MAX_WORKFLOW_TOKENS`, and `IRONFLOW_GUARD_WORKFLOW_TIMEOUT_SECS`.
86    /// Missing or unparseable variables fall back to defaults.
87    ///
88    /// # Examples
89    ///
90    /// ```
91    /// use ironflow_engine::guard::WorkflowGuardConfig;
92    ///
93    /// let config = WorkflowGuardConfig::from_env();
94    /// // Without env vars set, returns defaults.
95    /// assert_eq!(config.max_depth, 5);
96    /// ```
97    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    /// Set the maximum sub-workflow nesting depth.
107    ///
108    /// # Examples
109    ///
110    /// ```
111    /// use ironflow_engine::guard::WorkflowGuardConfig;
112    ///
113    /// let config = WorkflowGuardConfig::new().with_max_depth(3);
114    /// assert_eq!(config.max_depth, 3);
115    /// ```
116    pub fn with_max_depth(mut self, max_depth: u32) -> Self {
117        self.max_depth = max_depth;
118        self
119    }
120
121    /// Set the maximum total workflow invocations.
122    ///
123    /// # Examples
124    ///
125    /// ```
126    /// use ironflow_engine::guard::WorkflowGuardConfig;
127    ///
128    /// let config = WorkflowGuardConfig::new().with_max_fan_out(10);
129    /// assert_eq!(config.max_fan_out, 10);
130    /// ```
131    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    /// Set the cumulative token budget.
137    ///
138    /// # Examples
139    ///
140    /// ```
141    /// use ironflow_engine::guard::WorkflowGuardConfig;
142    ///
143    /// let config = WorkflowGuardConfig::new().with_max_workflow_tokens(50_000);
144    /// assert_eq!(config.max_workflow_tokens, 50_000);
145    /// ```
146    pub fn with_max_workflow_tokens(mut self, max: u64) -> Self {
147        self.max_workflow_tokens = max;
148        self
149    }
150
151    /// Set the global timeout in seconds.
152    ///
153    /// # Examples
154    ///
155    /// ```
156    /// use ironflow_engine::guard::WorkflowGuardConfig;
157    ///
158    /// let config = WorkflowGuardConfig::new().with_workflow_timeout_secs(60);
159    /// assert_eq!(config.workflow_timeout_secs, 60);
160    /// ```
161    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/// Reason a workflow invocation was rejected by the guard.
175///
176/// Each variant carries the observed value and the configured limit
177/// for diagnostics.
178///
179/// # Examples
180///
181/// ```
182/// use ironflow_engine::guard::WorkflowRejection;
183///
184/// let r = WorkflowRejection::MaxDepthExceeded { depth: 6, max: 5 };
185/// assert!(r.to_string().contains("6"));
186/// assert!(r.to_string().contains("5"));
187/// ```
188#[derive(Debug, Clone, PartialEq, Eq)]
189pub enum WorkflowRejection {
190    /// The sub-workflow nesting depth exceeds the configured maximum.
191    MaxDepthExceeded {
192        /// Current depth.
193        depth: u32,
194        /// Configured limit.
195        max: u32,
196    },
197    /// The target workflow is already in the call chain (direct or indirect cycle).
198    CycleDetected {
199        /// The workflow that would create a cycle.
200        target: String,
201        /// The current call chain.
202        chain: Vec<String>,
203    },
204    /// The total number of workflow invocations exceeds the configured maximum.
205    MaxFanOutExceeded {
206        /// Current invocation count.
207        invocations: u32,
208        /// Configured limit.
209        max: u32,
210    },
211    /// The cumulative token usage exceeds the configured budget.
212    TokenBudgetExhausted {
213        /// Tokens already consumed.
214        used: u64,
215        /// Configured limit.
216        max: u64,
217    },
218    /// The workflow has exceeded its global timeout.
219    WorkflowTimeout {
220        /// Seconds elapsed since the workflow started.
221        elapsed_secs: u64,
222        /// Configured limit.
223        max: u64,
224    },
225    /// The guard state is unavailable (fail closed).
226    GuardUnavailable,
227}
228
229/// Business error code for guard rejections.
230pub 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/// Mutable runtime state tracked by the guard during a workflow execution.
263///
264/// Shared between parent and child workflows via `Arc<Mutex<_>>` so that
265/// limits are enforced globally across the entire run tree.
266///
267/// # Examples
268///
269/// ```
270/// use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState};
271///
272/// let mut state = WorkflowGuardState::new();
273/// let config = WorkflowGuardConfig::new().with_max_depth(2);
274///
275/// state.record_invocation("workflow-a");
276/// state.record_invocation("workflow-b");
277/// assert!(state.check(&config, "workflow-c").is_err());
278/// ```
279#[derive(Debug)]
280pub struct WorkflowGuardState {
281    /// Current sub-workflow nesting depth.
282    depth: u32,
283    /// Call chain for cycle detection (stack of workflow names).
284    call_chain: Vec<String>,
285    /// Total number of workflow invocations in this run tree.
286    total_invocations: u32,
287    /// Cumulative tokens consumed across all agent steps.
288    total_tokens_used: u64,
289    /// When the root workflow started.
290    started_at: Instant,
291}
292
293impl WorkflowGuardState {
294    /// Create a fresh guard state for a new workflow execution.
295    ///
296    /// # Examples
297    ///
298    /// ```
299    /// use ironflow_engine::guard::WorkflowGuardState;
300    ///
301    /// let state = WorkflowGuardState::new();
302    /// assert_eq!(state.depth(), 0);
303    /// assert_eq!(state.total_invocations(), 0);
304    /// ```
305    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    /// Current sub-workflow nesting depth.
316    pub fn depth(&self) -> u32 {
317        self.depth
318    }
319
320    /// Total workflow invocations so far.
321    pub fn total_invocations(&self) -> u32 {
322        self.total_invocations
323    }
324
325    /// Total tokens consumed so far.
326    pub fn total_tokens_used(&self) -> u64 {
327        self.total_tokens_used
328    }
329
330    /// The current call chain (workflow name stack).
331    pub fn call_chain(&self) -> &[String] {
332        &self.call_chain
333    }
334
335    /// Seconds elapsed since the workflow started.
336    pub fn elapsed_secs(&self) -> u64 {
337        self.started_at.elapsed().as_secs()
338    }
339
340    /// Read-only check: would invoking `target_workflow` violate any limit?
341    ///
342    /// Does not mutate state. Call this before
343    /// [`record_invocation`](Self::record_invocation).
344    ///
345    /// # Errors
346    ///
347    /// Returns the specific [`WorkflowRejection`] variant that would be
348    /// violated.
349    ///
350    /// # Examples
351    ///
352    /// ```
353    /// use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState, WorkflowRejection};
354    ///
355    /// let state = WorkflowGuardState::new();
356    /// let config = WorkflowGuardConfig::new().with_max_depth(0);
357    ///
358    /// let err = state.check(&config, "child").unwrap_err();
359    /// assert!(matches!(err, WorkflowRejection::MaxDepthExceeded { .. }));
360    /// ```
361    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    /// Record that a sub-workflow invocation is starting.
406    ///
407    /// Increments depth and fan-out counter, and pushes the workflow name
408    /// onto the call chain for cycle detection.
409    ///
410    /// # Examples
411    ///
412    /// ```
413    /// use ironflow_engine::guard::WorkflowGuardState;
414    ///
415    /// let mut state = WorkflowGuardState::new();
416    /// state.record_invocation("child-workflow");
417    /// assert_eq!(state.depth(), 1);
418    /// assert_eq!(state.total_invocations(), 1);
419    /// assert_eq!(state.call_chain(), &["child-workflow"]);
420    /// ```
421    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    /// Record that a sub-workflow invocation has returned.
428    ///
429    /// Decrements depth and pops the last entry from the call chain.
430    /// Safe to call even on failure paths (the guard must never leak depth).
431    ///
432    /// # Examples
433    ///
434    /// ```
435    /// use ironflow_engine::guard::WorkflowGuardState;
436    ///
437    /// let mut state = WorkflowGuardState::new();
438    /// state.record_invocation("child");
439    /// assert_eq!(state.depth(), 1);
440    ///
441    /// state.record_return();
442    /// assert_eq!(state.depth(), 0);
443    /// assert!(state.call_chain().is_empty());
444    /// ```
445    pub fn record_return(&mut self) {
446        self.depth = self.depth.saturating_sub(1);
447        self.call_chain.pop();
448    }
449
450    /// Record tokens consumed by an agent step.
451    ///
452    /// Returns the remaining token budget, or an error if the budget is
453    /// now exhausted.
454    ///
455    /// # Errors
456    ///
457    /// Returns [`WorkflowRejection::TokenBudgetExhausted`] when the
458    /// cumulative usage exceeds the configured maximum.
459    ///
460    /// # Examples
461    ///
462    /// ```
463    /// use ironflow_engine::guard::{WorkflowGuardConfig, WorkflowGuardState, WorkflowRejection};
464    ///
465    /// let mut state = WorkflowGuardState::new();
466    /// let config = WorkflowGuardConfig::new().with_max_workflow_tokens(100);
467    ///
468    /// assert!(state.record_tokens(&config, 50).is_ok());
469    /// assert!(state.record_tokens(&config, 60).is_err());
470    /// ```
471    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
493/// Thread-safe handle to a shared [`WorkflowGuardState`].
494///
495/// Passed from parent to child workflows so that limits are enforced
496/// globally across the entire run tree.
497pub type SharedGuardState = Arc<Mutex<WorkflowGuardState>>;
498
499/// Create a new shared guard state for a workflow run.
500///
501/// # Examples
502///
503/// ```
504/// use ironflow_engine::guard::new_shared_guard_state;
505///
506/// let state = new_shared_guard_state();
507/// let locked = state.lock().unwrap();
508/// assert_eq!(locked.depth(), 0);
509/// ```
510pub 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        // SAFETY: test-only, single-threaded access to env vars.
535        // Combined into one test to avoid races between parallel tests.
536        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        // SAFETY: test-only cleanup.
550        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        // elapsed().as_secs() truncates to whole seconds, so we use a state
681        // whose started_at we can set to the past via the test-only constructor.
682        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}