1use std::sync::Arc;
2
3use crate::tools::metadata::{normalize_tool_metadata_labels, utf16_cmp};
4use crate::tools::{ToolMetadata, ToolMetadataError, ToolSideEffect};
5use crate::types::{Metadata, ToolArguments};
6
7pub type CanUseToolPredicate = Arc<dyn Fn(&str, &ToolArguments) -> bool + Send + Sync + 'static>;
8
9#[derive(Clone, Default)]
10pub struct ToolPolicy {
11 pub allowed_tools: Option<Vec<String>>,
12 pub disallowed_tools: Vec<String>,
13 pub approval: ApprovalPolicy,
14 pub can_use_tool: Option<CanUseToolPredicate>,
15 pub denied_side_effects: Vec<ToolSideEffect>,
16 pub denied_capability_tags: Vec<String>,
17 pub deny_terminal_tools: bool,
18 pub denied_cost_dimensions: Vec<String>,
19}
20
21impl ToolPolicy {
22 pub fn allow_only(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
23 self.allowed_tools = Some(tools.into_iter().map(Into::into).collect());
24 self
25 }
26
27 pub fn disallow(mut self, tool: impl Into<String>) -> Self {
28 self.disallowed_tools.push(tool.into());
29 self
30 }
31
32 pub fn can_use_tool<F>(mut self, predicate: F) -> Self
33 where
34 F: Fn(&str, &ToolArguments) -> bool + Send + Sync + 'static,
35 {
36 self.can_use_tool = Some(Arc::new(predicate));
37 self
38 }
39
40 pub fn deny_side_effect(mut self, side_effect: ToolSideEffect) -> Self {
41 self.denied_side_effects.push(side_effect);
42 self.normalize_metadata_denials_in_place()
43 .expect("a side-effect enum is always a valid denial");
44 self
45 }
46
47 pub fn deny_capability_tag(
48 mut self,
49 tag: impl Into<String>,
50 ) -> Result<Self, ToolMetadataError> {
51 self.denied_capability_tags.push(tag.into());
52 self.normalize_metadata_denials_in_place()?;
53 Ok(self)
54 }
55
56 pub fn deny_terminal_tools(mut self) -> Self {
57 self.deny_terminal_tools = true;
58 self
59 }
60
61 pub fn deny_cost_dimension(
62 mut self,
63 dimension: impl Into<String>,
64 ) -> Result<Self, ToolMetadataError> {
65 self.denied_cost_dimensions.push(dimension.into());
66 self.normalize_metadata_denials_in_place()?;
67 Ok(self)
68 }
69
70 pub fn allows_arguments(&self, tool_name: &str, arguments: &ToolArguments) -> bool {
71 self.can_use_tool
72 .as_ref()
73 .is_none_or(|predicate| predicate(tool_name, arguments))
74 }
75
76 pub fn normalized(&self) -> Result<Self, ToolMetadataError> {
77 let mut normalized = self.clone();
78 normalized.normalize_metadata_denials_in_place()?;
79 Ok(normalized)
80 }
81
82 pub fn metadata_denial_source(&self, metadata: Option<&ToolMetadata>) -> Option<&'static str> {
83 let metadata = metadata?;
84 if self.denied_side_effects.contains(&metadata.side_effect) {
85 return Some("metadata.side_effect");
86 }
87 if self.deny_terminal_tools && metadata.terminal {
88 return Some("metadata.terminal");
89 }
90 if metadata
91 .capability_tags
92 .iter()
93 .any(|tag| self.denied_capability_tags.contains(tag))
94 {
95 return Some("metadata.capability_tag");
96 }
97 if metadata
98 .cost_dimensions
99 .iter()
100 .any(|dimension| self.denied_cost_dimensions.contains(dimension))
101 {
102 return Some("metadata.cost_dimension");
103 }
104 None
105 }
106
107 pub(crate) fn extend_metadata_denials(&mut self, other: &Self) {
108 self.denied_side_effects
109 .extend(other.denied_side_effects.iter().copied());
110 self.denied_capability_tags
111 .extend(other.denied_capability_tags.iter().cloned());
112 self.deny_terminal_tools |= other.deny_terminal_tools;
113 self.denied_cost_dimensions
114 .extend(other.denied_cost_dimensions.iter().cloned());
115 self.normalize_metadata_denials_in_place()
116 .expect("normalized policies remain valid when unioned");
117 }
118
119 pub(crate) fn normalize_metadata_denials_in_place(&mut self) -> Result<(), ToolMetadataError> {
120 self.denied_side_effects
121 .sort_by(|left, right| utf16_cmp(left.as_str(), right.as_str()));
122 self.denied_side_effects.dedup();
123 self.denied_capability_tags =
124 normalize_tool_metadata_labels(&self.denied_capability_tags, "denied_capability_tags")?;
125 self.denied_cost_dimensions =
126 normalize_tool_metadata_labels(&self.denied_cost_dimensions, "denied_cost_dimensions")?;
127 Ok(())
128 }
129}
130
131#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
132pub enum ApprovalPolicy {
133 #[default]
134 Default,
135 Never,
136 Always,
137 OnRequest,
138}
139
140#[derive(Debug, Clone, PartialEq, Eq)]
141pub enum ApprovalDecision {
142 Approved,
143 ApprovedForSession,
144 Denied(String),
145 TimedOut(String),
146 NeedsApproval,
147 Detailed {
148 decision: Box<ApprovalDecision>,
149 reason: Option<String>,
150 metadata: Metadata,
151 },
152}
153
154impl ApprovalDecision {
155 pub fn allow() -> Self {
156 Self::Approved
157 }
158
159 pub fn allow_session() -> Self {
160 Self::ApprovedForSession
161 }
162
163 pub fn deny(reason: impl Into<String>) -> Self {
164 Self::Denied(reason.into())
165 }
166
167 pub fn timeout(reason: impl Into<String>) -> Self {
168 Self::TimedOut(reason.into())
169 }
170
171 pub fn with_reason(self, reason: impl Into<String>) -> Self {
172 let reason = reason.into();
173 match self {
174 Self::Denied(_) => Self::Denied(reason),
175 Self::TimedOut(_) => Self::TimedOut(reason),
176 Self::Detailed {
177 decision, metadata, ..
178 } => Self::Detailed {
179 decision,
180 reason: Some(reason),
181 metadata,
182 },
183 decision => Self::Detailed {
184 decision: Box::new(decision),
185 reason: Some(reason),
186 metadata: Metadata::new(),
187 },
188 }
189 }
190
191 pub fn with_metadata(self, metadata: Metadata) -> Self {
192 match self {
193 Self::Detailed {
194 decision,
195 reason,
196 metadata: mut existing,
197 } => {
198 existing.extend(metadata);
199 Self::Detailed {
200 decision,
201 reason,
202 metadata: existing,
203 }
204 }
205 decision => Self::Detailed {
206 decision: Box::new(decision),
207 reason: None,
208 metadata,
209 },
210 }
211 }
212
213 pub fn is_approved(&self) -> bool {
214 matches!(self.action(), "allow" | "allow_session")
215 }
216
217 pub fn action(&self) -> &'static str {
218 match self {
219 Self::Approved => "allow",
220 Self::ApprovedForSession => "allow_session",
221 Self::Denied(_) => "deny",
222 Self::TimedOut(_) => "timeout",
223 Self::NeedsApproval => "needs_approval",
224 Self::Detailed { decision, .. } => decision.action(),
225 }
226 }
227
228 pub fn reason(&self) -> &str {
229 match self {
230 Self::Denied(reason) | Self::TimedOut(reason) => reason,
231 Self::Detailed {
232 decision, reason, ..
233 } => reason.as_deref().unwrap_or_else(|| decision.reason()),
234 _ => "",
235 }
236 }
237
238 pub fn metadata(&self) -> Option<&Metadata> {
239 match self {
240 Self::Detailed { metadata, .. } => Some(metadata),
241 _ => None,
242 }
243 }
244}