nanny_runtime/
enforcement.rs1use nanny_core::agent::state::StopReason;
10use nanny_core::policy::{Policy, PolicyContext, PolicyDecision};
11use std::collections::HashMap;
12
13pub struct ToolPermissionPolicy {
21 allowed_tools: Vec<String>,
22}
23
24impl ToolPermissionPolicy {
25 pub fn new(allowed_tools: Vec<String>) -> Self {
26 Self { allowed_tools }
27 }
28}
29
30impl Policy for ToolPermissionPolicy {
31 fn evaluate(&self, ctx: &PolicyContext) -> PolicyDecision {
32 if let Some(tool) = &ctx.requested_tool {
33 if !self.allowed_tools.contains(tool) {
34 return PolicyDecision::Deny {
35 reason: StopReason::ToolDenied {
36 tool_name: tool.clone(),
37 },
38 };
39 }
40 }
41 PolicyDecision::Allow
42 }
43}
44
45pub struct RuleEvaluator {
54 max_calls: HashMap<String, u32>,
55}
56
57impl RuleEvaluator {
58 pub fn new(max_calls: HashMap<String, u32>) -> Self {
59 Self { max_calls }
60 }
61
62 pub fn rule_name_for(&self, tool: &str) -> Option<String> {
69 self.max_calls
70 .contains_key(tool)
71 .then(|| format!("{tool}.max_calls"))
72 }
73}
74
75impl Policy for RuleEvaluator {
76 fn evaluate(&self, ctx: &PolicyContext) -> PolicyDecision {
77 let tool = match &ctx.requested_tool {
78 Some(t) => t,
79 None => return PolicyDecision::Allow,
80 };
81 if let Some(&max) = self.max_calls.get(tool) {
82 let calls_so_far = ctx.tool_call_counts.get(tool).copied().unwrap_or(0);
83 if calls_so_far >= max {
84 return PolicyDecision::Deny {
85 reason: StopReason::RuleDenied {
86 rule_name: format!("{tool}.max_calls"),
87 },
88 };
89 }
90 }
91 PolicyDecision::Allow
92 }
93}
94
95pub struct ChainPolicy<A, B> {
99 first: A,
100 second: B,
101}
102
103impl<A, B> ChainPolicy<A, B> {
104 pub fn new(first: A, second: B) -> Self {
105 Self { first, second }
106 }
107}
108
109impl<A: Policy, B: Policy> Policy for ChainPolicy<A, B> {
110 fn evaluate(&self, ctx: &PolicyContext) -> PolicyDecision {
111 match self.first.evaluate(ctx) {
112 PolicyDecision::Allow => self.second.evaluate(ctx),
113 deny => deny,
114 }
115 }
116}
117
118#[cfg(test)]
121mod tests {
122 use super::*;
123
124 fn base_context() -> PolicyContext {
125 PolicyContext::default()
126 }
127
128 fn tool_policy() -> ToolPermissionPolicy {
129 ToolPermissionPolicy::new(vec!["http_get".to_string()])
130 }
131
132 #[test]
133 fn denies_unlisted_tool() {
134 let ctx = PolicyContext {
135 requested_tool: Some("write_file".to_string()),
136 ..base_context()
137 };
138 assert!(matches!(
139 tool_policy().evaluate(&ctx),
140 PolicyDecision::Deny {
141 reason: StopReason::ToolDenied { .. }
142 }
143 ));
144 }
145
146 #[test]
147 fn allows_listed_tool() {
148 let ctx = PolicyContext {
149 requested_tool: Some("http_get".to_string()),
150 ..base_context()
151 };
152 assert!(matches!(
153 tool_policy().evaluate(&ctx),
154 PolicyDecision::Allow
155 ));
156 }
157
158 fn rule_evaluator_with_http_get_limit(max: u32) -> RuleEvaluator {
159 let mut map = HashMap::new();
160 map.insert("http_get".to_string(), max);
161 RuleEvaluator::new(map)
162 }
163
164 #[test]
165 fn rule_evaluator_allows_when_under_limit() {
166 let re = rule_evaluator_with_http_get_limit(3);
167 let mut counts = HashMap::new();
168 counts.insert("http_get".to_string(), 2u32);
169 let ctx = PolicyContext {
170 requested_tool: Some("http_get".to_string()),
171 tool_call_counts: counts,
172 ..base_context()
173 };
174 assert!(matches!(re.evaluate(&ctx), PolicyDecision::Allow));
175 }
176
177 #[test]
178 fn rule_evaluator_denies_at_max_calls() {
179 let re = rule_evaluator_with_http_get_limit(3);
180 let mut counts = HashMap::new();
181 counts.insert("http_get".to_string(), 3u32);
182 let ctx = PolicyContext {
183 requested_tool: Some("http_get".to_string()),
184 tool_call_counts: counts,
185 ..base_context()
186 };
187 assert!(matches!(
188 re.evaluate(&ctx),
189 PolicyDecision::Deny {
190 reason: StopReason::RuleDenied { ref rule_name }
191 } if rule_name == "http_get.max_calls"
192 ));
193 }
194
195 #[test]
196 fn rule_evaluator_ignores_unconfigured_tools() {
197 let re = rule_evaluator_with_http_get_limit(1);
198 let ctx = PolicyContext {
199 requested_tool: Some("write_file".to_string()),
200 ..base_context()
201 };
202 assert!(matches!(re.evaluate(&ctx), PolicyDecision::Allow));
203 }
204
205 #[test]
206 fn rule_evaluator_allows_when_no_tool_requested() {
207 let re = rule_evaluator_with_http_get_limit(1);
208 assert!(matches!(
209 re.evaluate(&base_context()),
210 PolicyDecision::Allow
211 ));
212 }
213
214 #[test]
215 fn chain_allows_when_both_allow() {
216 let chain = ChainPolicy::new(
217 RuleEvaluator::new(HashMap::new()),
218 RuleEvaluator::new(HashMap::new()),
219 );
220 assert!(matches!(
221 chain.evaluate(&base_context()),
222 PolicyDecision::Allow
223 ));
224 }
225
226 #[test]
227 fn chain_denies_when_first_denies() {
228 let first = ToolPermissionPolicy::new(vec![]);
229 let second = RuleEvaluator::new(HashMap::new());
230 let chain = ChainPolicy::new(first, second);
231 let ctx = PolicyContext {
232 requested_tool: Some("http_get".to_string()),
233 ..base_context()
234 };
235 assert!(matches!(
236 chain.evaluate(&ctx),
237 PolicyDecision::Deny {
238 reason: StopReason::ToolDenied { .. }
239 }
240 ));
241 }
242
243 #[test]
244 fn chain_denies_when_second_denies() {
245 let first = RuleEvaluator::new(HashMap::new());
246 let re = rule_evaluator_with_http_get_limit(1);
247 let chain = ChainPolicy::new(first, re);
248 let mut counts = HashMap::new();
249 counts.insert("http_get".to_string(), 1u32);
250 let ctx = PolicyContext {
251 requested_tool: Some("http_get".to_string()),
252 tool_call_counts: counts,
253 ..base_context()
254 };
255 assert!(matches!(
256 chain.evaluate(&ctx),
257 PolicyDecision::Deny {
258 reason: StopReason::RuleDenied { .. }
259 }
260 ));
261 }
262
263 #[test]
264 fn chain_first_denial_wins_over_second() {
265 let first = ToolPermissionPolicy::new(vec![]);
268 let mut max_calls = HashMap::new();
269 max_calls.insert("http_get".to_string(), 0u32);
270 let second = RuleEvaluator::new(max_calls);
271 let chain = ChainPolicy::new(first, second);
272 let ctx = PolicyContext {
273 requested_tool: Some("http_get".to_string()),
274 ..base_context()
275 };
276 assert!(matches!(
277 chain.evaluate(&ctx),
278 PolicyDecision::Deny {
279 reason: StopReason::ToolDenied { .. }
280 }
281 ));
282 }
283}