1use std::collections::{HashMap, HashSet};
13
14use crate::base::AgentError;
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
18pub enum ToolRisk {
19 Safe,
21 Standard,
23 Dangerous,
26}
27
28impl ToolRisk {
29 pub fn name(&self) -> &'static str {
31 match self {
32 ToolRisk::Safe => "safe",
33 ToolRisk::Standard => "standard",
34 ToolRisk::Dangerous => "dangerous",
35 }
36 }
37}
38
39impl std::fmt::Display for ToolRisk {
40 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41 write!(f, "{}", self.name())
42 }
43}
44
45#[derive(Debug, Clone)]
61pub struct ToolPolicy {
62 risks: HashMap<String, ToolRisk>,
64 default_risk: ToolRisk,
66 max_permitted: ToolRisk,
68 sandboxed: HashSet<String>,
70 allow_unrestricted_dangerous: bool,
72}
73
74impl Default for ToolPolicy {
75 fn default() -> Self {
76 Self::new()
77 }
78}
79
80impl ToolPolicy {
81 pub fn new() -> Self {
84 Self {
85 risks: HashMap::new(),
86 default_risk: ToolRisk::Safe,
87 max_permitted: ToolRisk::Dangerous,
88 sandboxed: HashSet::new(),
89 allow_unrestricted_dangerous: false,
90 }
91 }
92
93 pub fn risk(mut self, name: impl Into<String>, risk: ToolRisk) -> Self {
95 self.risks.insert(name.into(), risk);
96 self
97 }
98
99 pub fn sandboxed(mut self, name: impl Into<String>) -> Self {
101 self.sandboxed.insert(name.into());
102 self
103 }
104
105 pub fn with_default_risk(mut self, risk: ToolRisk) -> Self {
107 self.default_risk = risk;
108 self
109 }
110
111 pub fn with_max_permitted(mut self, risk: ToolRisk) -> Self {
113 self.max_permitted = risk;
114 self
115 }
116
117 pub fn allow_unrestricted_dangerous(mut self, allow: bool) -> Self {
120 self.allow_unrestricted_dangerous = allow;
121 self
122 }
123
124 pub fn risk_of(&self, name: &str) -> ToolRisk {
126 self.risks.get(name).copied().unwrap_or(self.default_risk)
127 }
128
129 pub fn check(&self, name: &str) -> Result<(), AgentError> {
132 let risk = self.risk_of(name);
133
134 if risk > self.max_permitted {
136 return Err(AgentError::Other(format!(
137 "tool '{name}' requires permission tier '{risk}', max permitted is '{}'",
138 self.max_permitted
139 )));
140 }
141
142 if risk == ToolRisk::Dangerous
144 && !self.sandboxed.contains(name)
145 && !self.allow_unrestricted_dangerous
146 {
147 return Err(AgentError::Other(format!(
148 "dangerous tool '{name}' must run in a sandboxed environment \
149 (declare via ToolPolicy::sandboxed(\"{name}\"))"
150 )));
151 }
152
153 Ok(())
154 }
155}
156
157#[cfg(test)]
158mod tests {
159 use super::*;
160
161 #[test]
162 fn test_tool_risk_display_and_order() {
163 assert_eq!(ToolRisk::Safe.to_string(), "safe");
164 assert_eq!(ToolRisk::Standard.to_string(), "standard");
165 assert_eq!(ToolRisk::Dangerous.to_string(), "dangerous");
166 assert!(ToolRisk::Safe < ToolRisk::Standard);
167 assert!(ToolRisk::Standard < ToolRisk::Dangerous);
168 }
169
170 #[test]
171 fn test_policy_allows_default_safe_tools() {
172 let policy = ToolPolicy::new();
174 assert!(policy.check("any_tool").is_ok());
175 assert_eq!(policy.risk_of("any_tool"), ToolRisk::Safe);
176 }
177
178 #[test]
179 fn test_policy_permission_tier_gate() {
180 let policy = ToolPolicy::new()
181 .risk("calculator", ToolRisk::Dangerous)
182 .with_max_permitted(ToolRisk::Standard);
183 let err = policy.check("calculator").unwrap_err();
184 assert!(err.to_string().contains("permission tier"), "{}", err);
185 assert!(policy.check("other").is_ok());
187 }
188
189 #[test]
190 fn test_policy_dangerous_requires_sandbox() {
191 let policy = ToolPolicy::new().risk("code_interpreter", ToolRisk::Dangerous);
193 let err = policy.check("code_interpreter").unwrap_err();
194 assert!(err.to_string().contains("sandboxed"), "{}", err);
195
196 let policy = policy.sandboxed("code_interpreter");
198 assert!(policy.check("code_interpreter").is_ok());
199 }
200
201 #[test]
202 fn test_policy_allow_unrestricted_dangerous() {
203 let policy = ToolPolicy::new()
204 .risk("http", ToolRisk::Dangerous)
205 .allow_unrestricted_dangerous(true);
206 assert!(policy.check("http").is_ok());
207 }
208
209 #[test]
210 fn test_policy_sandboxed_but_not_dangerous_still_gated_by_tier() {
211 let policy = ToolPolicy::new()
213 .risk("calculator", ToolRisk::Dangerous)
214 .sandboxed("calculator")
215 .with_max_permitted(ToolRisk::Standard);
216 let err = policy.check("calculator").unwrap_err();
217 assert!(err.to_string().contains("permission tier"), "{}", err);
218 }
219}