llm_stack/tool/
depth.rs

1//! Loop depth tracking for nested tool loops.
2//!
3//! When tools spawn sub-agents (nested `tool_loop` calls), depth tracking
4//! prevents runaway recursion. Implement [`LoopDepth`] on your context type
5//! to enable automatic depth management.
6//!
7//! # Example
8//!
9//! ```rust
10//! use llm_stack::tool::LoopDepth;
11//!
12//! #[derive(Clone)]
13//! struct AgentContext {
14//!     user_id: String,
15//!     depth: u32,
16//! }
17//!
18//! impl LoopDepth for AgentContext {
19//!     fn loop_depth(&self) -> u32 {
20//!         self.depth
21//!     }
22//!
23//!     fn with_depth(&self, depth: u32) -> Self {
24//!         Self {
25//!             depth,
26//!             ..self.clone()
27//!         }
28//!     }
29//! }
30//! ```
31
32/// Trait for contexts that support automatic depth tracking in nested tool loops.
33///
34/// When `tool_loop` executes tools, it passes a context with incremented depth
35/// so that nested loops can enforce depth limits via `max_depth` in
36/// [`ToolLoopConfig`](super::ToolLoopConfig).
37///
38/// # Blanket Implementation
39///
40/// The unit type `()` has a blanket implementation that always returns depth 0.
41/// Use this for simple cases where depth tracking isn't needed:
42///
43/// ```rust
44/// use llm_stack::tool::LoopDepth;
45///
46/// // () always returns 0, ignores depth changes
47/// assert_eq!(().loop_depth(), 0);
48/// assert_eq!(().with_depth(5), ());
49/// ```
50///
51/// # Custom Implementation
52///
53/// For agent systems with nesting, implement this on your context type:
54///
55/// ```rust
56/// use llm_stack::tool::LoopDepth;
57///
58/// #[derive(Clone)]
59/// struct MyContext {
60///     session_id: String,
61///     loop_depth: u32,
62/// }
63///
64/// impl LoopDepth for MyContext {
65///     fn loop_depth(&self) -> u32 {
66///         self.loop_depth
67///     }
68///
69///     fn with_depth(&self, depth: u32) -> Self {
70///         Self {
71///             loop_depth: depth,
72///             ..self.clone()
73///         }
74///     }
75/// }
76/// ```
77pub trait LoopDepth: Clone + Send + Sync {
78    /// Returns the current nesting depth.
79    ///
80    /// A depth of 0 means this is the top-level loop (not nested).
81    fn loop_depth(&self) -> u32;
82
83    /// Returns a new context with the specified depth.
84    ///
85    /// Called by `tool_loop` when passing context to tool handlers,
86    /// incrementing depth for any nested loops.
87    #[must_use]
88    fn with_depth(&self, depth: u32) -> Self;
89}
90
91/// Blanket implementation for unit type — always depth 0, no tracking.
92///
93/// This allows simple use cases to work without implementing the trait:
94///
95/// ```rust
96/// use llm_stack::tool::{ToolLoopConfig, ToolRegistry};
97///
98/// // Works with () context, no depth tracking
99/// let registry: ToolRegistry<()> = ToolRegistry::new();
100/// ```
101impl LoopDepth for () {
102    fn loop_depth(&self) -> u32 {
103        0
104    }
105
106    fn with_depth(&self, _depth: u32) -> Self {}
107}
108
109#[cfg(test)]
110mod tests {
111    use super::*;
112
113    #[test]
114    fn test_unit_loop_depth() {
115        assert_eq!(().loop_depth(), 0);
116    }
117
118    #[test]
119    #[allow(clippy::let_unit_value)]
120    fn test_unit_with_depth_ignores_value() {
121        // with_depth on () returns (), which still has depth 0
122        let nested = ().with_depth(5);
123        // Even after "nesting", unit context still reports depth 0
124        assert_eq!(nested.loop_depth(), 0);
125    }
126
127    #[derive(Clone)]
128    struct TestContext {
129        name: String,
130        depth: u32,
131    }
132
133    impl LoopDepth for TestContext {
134        fn loop_depth(&self) -> u32 {
135            self.depth
136        }
137
138        fn with_depth(&self, depth: u32) -> Self {
139            Self {
140                depth,
141                ..self.clone()
142            }
143        }
144    }
145
146    #[test]
147    fn test_custom_context_depth() {
148        let ctx = TestContext {
149            name: "test".into(),
150            depth: 0,
151        };
152        assert_eq!(ctx.loop_depth(), 0);
153
154        let nested = ctx.with_depth(1);
155        assert_eq!(nested.loop_depth(), 1);
156        assert_eq!(nested.name, "test"); // Other fields preserved
157    }
158
159    #[test]
160    fn test_depth_increments() {
161        let ctx = TestContext {
162            name: "agent".into(),
163            depth: 0,
164        };
165
166        let level1 = ctx.with_depth(ctx.loop_depth() + 1);
167        assert_eq!(level1.loop_depth(), 1);
168
169        let level2 = level1.with_depth(level1.loop_depth() + 1);
170        assert_eq!(level2.loop_depth(), 2);
171    }
172}
llm_stack/tool/depth.rs

llm_stack/tool/
depth.rs
