Skip to main content

vv_agent/tools/
policy.rs

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}