1use std::sync::Arc;
2
3use crate::types::{Metadata, ToolArguments};
4
5pub type CanUseToolPredicate = Arc<dyn Fn(&str, &ToolArguments) -> bool + Send + Sync + 'static>;
6
7#[derive(Clone, Default)]
8pub struct ToolPolicy {
9 pub allowed_tools: Option<Vec<String>>,
10 pub disallowed_tools: Vec<String>,
11 pub approval: ApprovalPolicy,
12 pub can_use_tool: Option<CanUseToolPredicate>,
13}
14
15impl ToolPolicy {
16 pub fn allow_only(mut self, tools: impl IntoIterator<Item = impl Into<String>>) -> Self {
17 self.allowed_tools = Some(tools.into_iter().map(Into::into).collect());
18 self
19 }
20
21 pub fn disallow(mut self, tool: impl Into<String>) -> Self {
22 self.disallowed_tools.push(tool.into());
23 self
24 }
25
26 pub fn can_use_tool<F>(mut self, predicate: F) -> Self
27 where
28 F: Fn(&str, &ToolArguments) -> bool + Send + Sync + 'static,
29 {
30 self.can_use_tool = Some(Arc::new(predicate));
31 self
32 }
33
34 pub fn allows_arguments(&self, tool_name: &str, arguments: &ToolArguments) -> bool {
35 self.can_use_tool
36 .as_ref()
37 .is_none_or(|predicate| predicate(tool_name, arguments))
38 }
39}
40
41#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
42pub enum ApprovalPolicy {
43 #[default]
44 Default,
45 Never,
46 Always,
47 OnRequest,
48}
49
50#[derive(Debug, Clone, PartialEq, Eq)]
51pub enum ApprovalDecision {
52 Approved,
53 ApprovedForSession,
54 Denied(String),
55 TimedOut(String),
56 NeedsApproval,
57 Detailed {
58 decision: Box<ApprovalDecision>,
59 reason: Option<String>,
60 metadata: Metadata,
61 },
62}
63
64impl ApprovalDecision {
65 pub fn allow() -> Self {
66 Self::Approved
67 }
68
69 pub fn allow_session() -> Self {
70 Self::ApprovedForSession
71 }
72
73 pub fn deny(reason: impl Into<String>) -> Self {
74 Self::Denied(reason.into())
75 }
76
77 pub fn timeout(reason: impl Into<String>) -> Self {
78 Self::TimedOut(reason.into())
79 }
80
81 pub fn with_reason(self, reason: impl Into<String>) -> Self {
82 let reason = reason.into();
83 match self {
84 Self::Denied(_) => Self::Denied(reason),
85 Self::TimedOut(_) => Self::TimedOut(reason),
86 Self::Detailed {
87 decision, metadata, ..
88 } => Self::Detailed {
89 decision,
90 reason: Some(reason),
91 metadata,
92 },
93 decision => Self::Detailed {
94 decision: Box::new(decision),
95 reason: Some(reason),
96 metadata: Metadata::new(),
97 },
98 }
99 }
100
101 pub fn with_metadata(self, metadata: Metadata) -> Self {
102 match self {
103 Self::Detailed {
104 decision,
105 reason,
106 metadata: mut existing,
107 } => {
108 existing.extend(metadata);
109 Self::Detailed {
110 decision,
111 reason,
112 metadata: existing,
113 }
114 }
115 decision => Self::Detailed {
116 decision: Box::new(decision),
117 reason: None,
118 metadata,
119 },
120 }
121 }
122
123 pub fn is_approved(&self) -> bool {
124 matches!(self.action(), "allow" | "allow_session")
125 }
126
127 pub fn action(&self) -> &'static str {
128 match self {
129 Self::Approved => "allow",
130 Self::ApprovedForSession => "allow_session",
131 Self::Denied(_) => "deny",
132 Self::TimedOut(_) => "timeout",
133 Self::NeedsApproval => "needs_approval",
134 Self::Detailed { decision, .. } => decision.action(),
135 }
136 }
137
138 pub fn reason(&self) -> &str {
139 match self {
140 Self::Denied(reason) | Self::TimedOut(reason) => reason,
141 Self::Detailed {
142 decision, reason, ..
143 } => reason.as_deref().unwrap_or_else(|| decision.reason()),
144 _ => "",
145 }
146 }
147
148 pub fn metadata(&self) -> Option<&Metadata> {
149 match self {
150 Self::Detailed { metadata, .. } => Some(metadata),
151 _ => None,
152 }
153 }
154}