1use std::collections::HashMap;
20
21use serde::{Deserialize, Serialize};
22use sha2::{Digest, Sha256};
23
24use crate::internal::tool_runtime::runtime::{tool_risk, ToolRisk};
25use crate::tools::IDEMPOTENT_TOOLS;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Deserialize, Serialize)]
31#[serde(rename_all = "lowercase")]
32pub enum ToolPolicyMode {
33 #[default]
34 Off,
35 Warn,
36 Block,
37}
38
39#[derive(Debug, Clone, Deserialize, Serialize)]
46pub struct ToolPolicyConfig {
47 #[serde(default)]
48 pub mode: ToolPolicyMode,
49 #[serde(default)]
52 pub deny_tools: Vec<String>,
53 #[serde(default)]
58 pub max_risk: Option<ToolRisk>,
59 #[serde(default)]
65 pub allow_tools: Vec<String>,
66 #[serde(default)]
69 pub grants_mode: ToolPolicyMode,
70 #[serde(default = "default_control_chat_grant_bypass")]
71 pub control_chat_bypass: bool,
72 #[serde(default)]
73 pub grants: Vec<ToolGrantRule>,
74}
75
76impl Default for ToolPolicyConfig {
77 fn default() -> Self {
78 Self {
79 mode: ToolPolicyMode::Off,
80 deny_tools: Vec::new(),
81 max_risk: None,
82 allow_tools: Vec::new(),
83 grants_mode: ToolPolicyMode::Off,
84 control_chat_bypass: true,
85 grants: Vec::new(),
86 }
87 }
88}
89
90fn default_control_chat_grant_bypass() -> bool {
91 true
92}
93
94#[derive(Debug, Clone, Default, Deserialize, Serialize)]
95pub struct ToolGrantRule {
96 #[serde(default)]
98 pub chat_id: Option<i64>,
99 #[serde(default)]
101 pub channel: Option<String>,
102 #[serde(default)]
104 pub principal: Option<String>,
105 #[serde(default)]
107 pub allow_tools: Vec<String>,
108 #[serde(default)]
109 pub deny_tools: Vec<String>,
110 #[serde(default)]
111 pub max_risk: Option<ToolRisk>,
112}
113
114#[derive(Debug, Clone, PartialEq, Eq)]
116pub enum PolicyDecision {
117 Allow,
118 Warn(String),
120 Block(String),
122}
123
124pub fn evaluate_tool_policy(cfg: &ToolPolicyConfig, tool_name: &str) -> PolicyDecision {
127 evaluate_tool_policy_with_risk(cfg, tool_name, tool_risk(tool_name))
128}
129
130pub fn evaluate_tool_policy_with_risk(
134 cfg: &ToolPolicyConfig,
135 tool_name: &str,
136 risk: crate::internal::tool_runtime::runtime::ToolRisk,
137) -> PolicyDecision {
138 if cfg.mode == ToolPolicyMode::Off {
139 return PolicyDecision::Allow;
140 }
141 let violation = if tool_list_contains(&cfg.deny_tools, tool_name) {
144 Some(format!("tool `{tool_name}` is in tool_policy.deny_tools"))
145 } else if !cfg.allow_tools.is_empty() && !tool_list_contains(&cfg.allow_tools, tool_name) {
146 Some(format!(
147 "tool `{tool_name}` is not in tool_policy.allow_tools"
148 ))
149 } else {
150 match cfg.max_risk {
151 Some(max_risk) if risk > max_risk => Some(format!(
152 "tool `{tool_name}` risk `{}` exceeds tool_policy.max_risk `{}`",
153 risk.as_str(),
154 max_risk.as_str()
155 )),
156 _ => None,
157 }
158 };
159 match violation {
160 None => PolicyDecision::Allow,
161 Some(reason) => {
162 if cfg.mode == ToolPolicyMode::Block {
163 PolicyDecision::Block(reason)
164 } else {
165 PolicyDecision::Warn(reason)
166 }
167 }
168 }
169}
170
171pub fn global_policy_has_no_rules(cfg: &ToolPolicyConfig) -> bool {
176 cfg.mode != ToolPolicyMode::Off
177 && cfg.deny_tools.is_empty()
178 && cfg.allow_tools.is_empty()
179 && cfg.max_risk.is_none()
180}
181
182pub fn evaluate_tool_policy_for_auth(
185 cfg: &ToolPolicyConfig,
186 tool_name: &str,
187 auth: &crate::internal::tool_runtime::runtime::ToolAuthContext,
188) -> PolicyDecision {
189 evaluate_tool_policy_for_auth_with_risk(cfg, tool_name, auth, tool_risk(tool_name))
190}
191
192pub fn evaluate_tool_policy_for_auth_with_risk(
196 cfg: &ToolPolicyConfig,
197 tool_name: &str,
198 auth: &crate::internal::tool_runtime::runtime::ToolAuthContext,
199 risk: crate::internal::tool_runtime::runtime::ToolRisk,
200) -> PolicyDecision {
201 let global = evaluate_tool_policy_with_risk(cfg, tool_name, risk);
202 if matches!(global, PolicyDecision::Block(_))
203 || cfg.grants_mode == ToolPolicyMode::Off
204 || (cfg.control_chat_bypass && auth.is_control_chat())
205 {
206 return global;
207 }
208
209 let matching: Vec<&ToolGrantRule> = cfg
210 .grants
211 .iter()
212 .filter(|rule| grant_matches(rule, auth))
213 .collect();
214 let violation = if matching.is_empty() {
215 Some(format!(
216 "no tool grant matches chat {} principal `{}`",
217 auth.caller_chat_id, auth.principal
218 ))
219 } else if matching
220 .iter()
221 .any(|rule| tool_list_contains(&rule.deny_tools, tool_name))
222 {
223 Some(format!(
224 "tool `{tool_name}` is denied for principal `{}`",
225 auth.principal
226 ))
227 } else if matching
228 .iter()
229 .any(|rule| rule.max_risk.is_some_and(|max_risk| risk > max_risk))
230 {
231 Some(format!(
232 "tool `{tool_name}` risk `{}` exceeds a matching capability grant",
233 risk.as_str()
234 ))
235 } else if !matching
236 .iter()
237 .any(|rule| tool_list_contains(&rule.allow_tools, tool_name))
238 {
239 Some(format!(
240 "tool `{tool_name}` is not granted to chat {} principal `{}`",
241 auth.caller_chat_id, auth.principal
242 ))
243 } else {
244 None
245 };
246
247 match violation {
248 Some(reason) if cfg.grants_mode == ToolPolicyMode::Block => PolicyDecision::Block(reason),
249 Some(reason) => PolicyDecision::Warn(reason),
250 None => global,
251 }
252}
253
254fn grant_matches(
255 rule: &ToolGrantRule,
256 auth: &crate::internal::tool_runtime::runtime::ToolAuthContext,
257) -> bool {
258 rule.chat_id
259 .is_none_or(|chat_id| chat_id == auth.caller_chat_id)
260 && rule
261 .channel
262 .as_deref()
263 .is_none_or(|channel| channel.eq_ignore_ascii_case(&auth.caller_channel))
264 && rule
265 .principal
266 .as_deref()
267 .is_none_or(|principal| wildcard_match(principal, &auth.principal))
268}
269
270fn wildcard_match(pattern: &str, value: &str) -> bool {
271 pattern
272 .strip_suffix('*')
273 .map_or(pattern == value, |prefix| value.starts_with(prefix))
274}
275
276fn tool_list_contains(list: &[String], tool_name: &str) -> bool {
277 list.iter()
278 .any(|pattern| pattern == "*" || wildcard_match(pattern, tool_name))
279}
280
281const NO_PROGRESS_WARN_AFTER: usize = 2;
283const SAME_TOOL_FAILURE_WARN_AFTER: usize = 3;
284
285#[derive(Default)]
286pub struct GuardrailController {
287 no_progress: HashMap<(String, String), (String, usize)>,
289 same_tool_failures: HashMap<String, usize>,
291}
292
293impl GuardrailController {
294 pub fn new() -> Self {
295 Self::default()
296 }
297
298 pub fn reset_for_turn(&mut self) {
299 self.no_progress.clear();
300 self.same_tool_failures.clear();
301 }
302
303 pub fn after_call(
306 &mut self,
307 tool_name: &str,
308 args_hash: &str,
309 result: &str,
310 failed: bool,
311 ) -> Option<String> {
312 if failed {
313 self.no_progress
316 .remove(&(tool_name.to_string(), args_hash.to_string()));
317 let count = self
318 .same_tool_failures
319 .entry(tool_name.to_string())
320 .or_insert(0);
321 *count += 1;
322 if *count >= SAME_TOOL_FAILURE_WARN_AFTER {
323 return Some(format!(
324 "\n\n[Tool loop warning: `{tool_name}` has failed {count} time(s) this turn. \
325 This looks like a loop — change approach (different tool, different inputs, \
326 or summarize what you know and proceed) before retrying.]"
327 ));
328 }
329 return None;
330 }
331
332 self.same_tool_failures.remove(tool_name);
334
335 if !is_idempotent(tool_name) {
336 return None;
337 }
338
339 let result_hash = hash_str(result);
340 let key = (tool_name.to_string(), args_hash.to_string());
341 let entry = self
342 .no_progress
343 .entry(key)
344 .or_insert((result_hash.clone(), 0));
345 if entry.0 == result_hash {
346 entry.1 += 1;
347 } else {
348 *entry = (result_hash, 1);
349 }
350 let count = entry.1;
351 if count >= NO_PROGRESS_WARN_AFTER {
352 return Some(format!(
353 "\n\n[Tool loop warning: `{tool_name}` has returned the same result \
354 {count} time(s) for the same arguments. Use the result you already have \
355 — re-running it won't change the answer.]"
356 ));
357 }
358 None
359 }
360}
361
362fn is_idempotent(tool_name: &str) -> bool {
363 IDEMPOTENT_TOOLS.contains(&tool_name)
364}
365
366fn hash_str(s: &str) -> String {
367 let mut hasher = Sha256::new();
368 hasher.update(s.as_bytes());
369 let digest = hasher.finalize();
370 format!("{:x}", digest)[..16].to_string()
372}
373
374#[cfg(test)]
375mod tests {
376 use super::*;
377
378 fn cfg(mode: ToolPolicyMode) -> ToolPolicyConfig {
379 ToolPolicyConfig {
380 mode,
381 ..Default::default()
382 }
383 }
384
385 #[test]
386 fn policy_off_allows_everything() {
387 let mut c = cfg(ToolPolicyMode::Off);
388 c.deny_tools = vec!["bash".into()];
389 c.max_risk = Some(ToolRisk::Low);
390 assert_eq!(evaluate_tool_policy(&c, "bash"), PolicyDecision::Allow);
391 }
392
393 #[test]
394 fn policy_blocks_denied_tool() {
395 let mut c = cfg(ToolPolicyMode::Block);
396 c.deny_tools = vec!["bash".into()];
397 assert!(matches!(
398 evaluate_tool_policy(&c, "bash"),
399 PolicyDecision::Block(_)
400 ));
401 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
402 }
403
404 #[test]
405 fn policy_warn_mode_warns_not_blocks() {
406 let mut c = cfg(ToolPolicyMode::Warn);
407 c.deny_tools = vec!["bash".into()];
408 assert!(matches!(
409 evaluate_tool_policy(&c, "bash"),
410 PolicyDecision::Warn(_)
411 ));
412 }
413
414 #[test]
415 fn policy_max_risk_blocks_higher_tiers() {
416 let mut c = cfg(ToolPolicyMode::Block);
417 c.max_risk = Some(ToolRisk::Medium);
418 assert!(matches!(
420 evaluate_tool_policy(&c, "bash"),
421 PolicyDecision::Block(_)
422 ));
423 assert_eq!(
425 evaluate_tool_policy(&c, "write_file"),
426 PolicyDecision::Allow
427 );
428 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
430 }
431
432 #[test]
433 fn policy_max_risk_covers_mcp_tools() {
434 let mut c = cfg(ToolPolicyMode::Block);
436 c.max_risk = Some(ToolRisk::Low);
437 assert!(matches!(
438 evaluate_tool_policy(&c, "mcp_fs_delete_file"),
439 PolicyDecision::Block(_)
440 ));
441 }
442
443 #[test]
444 fn policy_with_risk_override_replaces_name_derived_risk() {
445 let mut c = cfg(ToolPolicyMode::Block);
448 c.max_risk = Some(ToolRisk::Medium);
449 assert!(matches!(
451 evaluate_tool_policy(&c, "mcp_fs_delete_file"),
452 PolicyDecision::Allow
453 ));
454 assert!(matches!(
456 evaluate_tool_policy_with_risk(&c, "mcp_fs_delete_file", ToolRisk::High),
457 PolicyDecision::Block(_)
458 ));
459 c.max_risk = Some(ToolRisk::Low);
461 assert!(matches!(
462 evaluate_tool_policy_with_risk(&c, "mcp_fs_delete_file", ToolRisk::Low),
463 PolicyDecision::Allow
464 ));
465 }
466
467 #[test]
468 fn policy_allow_list_restricts_to_listed_tools() {
469 let mut c = cfg(ToolPolicyMode::Block);
472 c.allow_tools = vec!["read_file".into(), "web_*".into()];
473 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
474 assert_eq!(
475 evaluate_tool_policy(&c, "web_search"),
476 PolicyDecision::Allow
477 );
478 match evaluate_tool_policy(&c, "bash") {
479 PolicyDecision::Block(reason) => {
480 assert!(
481 reason.contains("not in tool_policy.allow_tools"),
482 "{reason}"
483 )
484 }
485 other => panic!("expected block, got {other:?}"),
486 }
487 assert!(matches!(
488 evaluate_tool_policy(&c, "write_file"),
489 PolicyDecision::Block(_)
490 ));
491 c.mode = ToolPolicyMode::Warn;
493 assert!(matches!(
494 evaluate_tool_policy(&c, "bash"),
495 PolicyDecision::Warn(_)
496 ));
497 }
498
499 #[test]
500 fn policy_empty_allow_list_permits_everything_not_denied() {
501 let mut c = cfg(ToolPolicyMode::Block);
502 c.deny_tools = vec!["bash".into()];
503 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
504 assert_eq!(
505 evaluate_tool_policy(&c, "write_file"),
506 PolicyDecision::Allow
507 );
508 }
509
510 #[test]
511 fn policy_deny_wins_over_allow() {
512 let mut c = cfg(ToolPolicyMode::Block);
513 c.deny_tools = vec!["bash".into()];
514 c.allow_tools = vec!["bash".into(), "read_file".into()];
515 assert!(matches!(
516 evaluate_tool_policy(&c, "bash"),
517 PolicyDecision::Block(_)
518 ));
519 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
520 }
521
522 #[test]
523 fn policy_allow_list_does_not_lift_max_risk() {
524 let mut c = cfg(ToolPolicyMode::Block);
525 c.max_risk = Some(ToolRisk::Low);
526 c.allow_tools = vec!["bash".into(), "read_file".into()];
527 assert!(matches!(
528 evaluate_tool_policy(&c, "bash"),
529 PolicyDecision::Block(_)
530 ));
531 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
532 }
533
534 #[test]
535 fn policy_global_lists_support_wildcards() {
536 let mut c = cfg(ToolPolicyMode::Block);
537 c.deny_tools = vec!["mcp_*".into()];
538 assert!(matches!(
539 evaluate_tool_policy(&c, "mcp_fs_delete_file"),
540 PolicyDecision::Block(_)
541 ));
542 assert_eq!(evaluate_tool_policy(&c, "read_file"), PolicyDecision::Allow);
543
544 let mut c = cfg(ToolPolicyMode::Block);
545 c.allow_tools = vec!["*".into()];
546 assert_eq!(evaluate_tool_policy(&c, "bash"), PolicyDecision::Allow);
547 }
548
549 #[test]
550 fn policy_without_rules_is_flagged() {
551 assert!(!global_policy_has_no_rules(&cfg(ToolPolicyMode::Off)));
552 assert!(global_policy_has_no_rules(&cfg(ToolPolicyMode::Block)));
553 assert!(global_policy_has_no_rules(&cfg(ToolPolicyMode::Warn)));
554 let mut c = cfg(ToolPolicyMode::Block);
555 c.allow_tools = vec!["read_file".into()];
556 assert!(!global_policy_has_no_rules(&c));
557 let mut c = cfg(ToolPolicyMode::Block);
558 c.max_risk = Some(ToolRisk::High);
559 assert!(!global_policy_has_no_rules(&c));
560 }
561
562 #[test]
563 fn policy_invalid_max_risk_rejected_at_parse() {
564 assert!(
566 serde_yaml::from_str::<ToolPolicyConfig>("mode: block\nmax_risk: extreme").is_err()
567 );
568 let ok: ToolPolicyConfig = serde_yaml::from_str("mode: block\nmax_risk: medium").unwrap();
569 assert_eq!(ok.max_risk, Some(ToolRisk::Medium));
570 }
571
572 fn auth(
573 chat_id: i64,
574 principal: &str,
575 ) -> crate::internal::tool_runtime::runtime::ToolAuthContext {
576 crate::internal::tool_runtime::runtime::ToolAuthContext {
577 caller_channel: "telegram".into(),
578 caller_chat_id: chat_id,
579 principal: principal.into(),
580 control_chat_ids: vec![],
581 env_files: vec![],
582 }
583 }
584
585 #[test]
586 fn capability_grants_are_chat_and_principal_scoped() {
587 let mut config = cfg(ToolPolicyMode::Off);
588 config.grants_mode = ToolPolicyMode::Block;
589 config.grants = vec![ToolGrantRule {
590 chat_id: Some(7),
591 principal: Some("subagent:*".into()),
592 allow_tools: vec!["read_*".into(), "web_search".into()],
593 max_risk: Some(ToolRisk::Low),
594 ..Default::default()
595 }];
596
597 assert_eq!(
598 evaluate_tool_policy_for_auth(&config, "read_file", &auth(7, "subagent:abc")),
599 PolicyDecision::Allow
600 );
601 assert!(matches!(
602 evaluate_tool_policy_for_auth(&config, "write_file", &auth(7, "subagent:abc")),
603 PolicyDecision::Block(_)
604 ));
605 assert!(matches!(
606 evaluate_tool_policy_for_auth(&config, "read_file", &auth(8, "subagent:abc")),
607 PolicyDecision::Block(_)
608 ));
609 assert!(matches!(
610 evaluate_tool_policy_for_auth(&config, "read_file", &auth(7, "main")),
611 PolicyDecision::Block(_)
612 ));
613 }
614
615 #[test]
616 fn global_block_cannot_be_weakened_by_grant() {
617 let mut config = cfg(ToolPolicyMode::Block);
618 config.deny_tools = vec!["bash".into()];
619 config.grants_mode = ToolPolicyMode::Block;
620 config.grants = vec![ToolGrantRule {
621 allow_tools: vec!["*".into()],
622 ..Default::default()
623 }];
624 assert!(matches!(
625 evaluate_tool_policy_for_auth(&config, "bash", &auth(7, "main")),
626 PolicyDecision::Block(_)
627 ));
628 }
629
630 #[test]
631 fn idempotent_no_progress_warns_on_third_identical() {
632 let mut c = GuardrailController::new();
633 assert!(c
634 .after_call("read_file", "h1", "file contents", false)
635 .is_none());
636 let warn = c.after_call("read_file", "h1", "file contents", false);
638 assert!(warn.is_some());
639 assert!(warn.unwrap().contains("same result"));
640 }
641
642 #[test]
643 fn idempotent_changing_result_does_not_warn() {
644 let mut c = GuardrailController::new();
645 c.after_call("read_file", "h1", "v1", false);
646 c.after_call("read_file", "h1", "v2", false);
647 let warn = c.after_call("read_file", "h1", "v3", false);
648 assert!(warn.is_none());
649 }
650
651 #[test]
652 fn mutating_tool_no_warning_on_repeat() {
653 let mut c = GuardrailController::new();
654 c.after_call("write_file", "h1", "ok", false);
655 let warn = c.after_call("write_file", "h1", "ok", false);
656 assert!(warn.is_none());
657 }
658
659 #[test]
660 fn same_tool_failure_streak_warns() {
661 let mut c = GuardrailController::new();
662 c.after_call("bash", "h1", "err1", true);
663 c.after_call("bash", "h2", "err2", true);
664 let warn = c.after_call("bash", "h3", "err3", true);
666 assert!(warn.is_some());
667 assert!(warn.unwrap().contains("failed 3"));
668 }
669
670 #[test]
671 fn success_resets_failure_streak() {
672 let mut c = GuardrailController::new();
673 c.after_call("bash", "h1", "err1", true);
674 c.after_call("bash", "h2", "err2", true);
675 c.after_call("bash", "h3", "ok", false);
676 let warn = c.after_call("bash", "h4", "err4", true);
678 assert!(warn.is_none());
679 let warn = c.after_call("bash", "h5", "err5", true);
680 assert!(warn.is_none());
681 }
682
683 #[test]
684 fn failure_does_not_increment_no_progress() {
685 let mut c = GuardrailController::new();
686 c.after_call("read_file", "h1", "ok", false);
687 c.after_call("read_file", "h1", "err", true);
688 let warn = c.after_call("read_file", "h1", "ok", false);
689 assert!(warn.is_none());
691 }
692
693 #[test]
694 fn reset_for_turn_clears_state() {
695 let mut c = GuardrailController::new();
696 c.after_call("read_file", "h1", "v1", false);
697 c.after_call("read_file", "h1", "v1", false);
698 c.reset_for_turn();
699 let warn = c.after_call("read_file", "h1", "v1", false);
700 assert!(warn.is_none());
701 }
702}