Skip to main content

mcp_utils/client/
tool_filter.rs

1use llm::ToolDefinition;
2use utils::matches_name_pattern;
3
4#[derive(Debug, Clone, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
5#[serde(untagged)]
6pub enum ToolMatcher {
7    Name(String),
8    Annotations(ToolAnnotationMatcher),
9}
10
11impl ToolMatcher {
12    pub fn name(pattern: impl Into<String>) -> Self {
13        Self::Name(pattern.into())
14    }
15
16    pub fn read_only() -> Self {
17        Self::Annotations(ToolAnnotationMatcher { read_only: Some(true), ..ToolAnnotationMatcher::default() })
18    }
19
20    pub fn annotations(matcher: ToolAnnotationMatcher) -> Self {
21        Self::Annotations(matcher)
22    }
23
24    pub fn matches(&self, tool: &ToolDefinition) -> bool {
25        match self {
26            Self::Name(pattern) => matches_name_pattern(pattern, &tool.name),
27            Self::Annotations(matcher) => matcher.matches(tool),
28        }
29    }
30}
31
32#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
33#[serde(rename_all = "camelCase", deny_unknown_fields)]
34pub struct ToolAnnotationMatcher {
35    #[serde(default, skip_serializing_if = "Option::is_none")]
36    pub read_only: Option<bool>,
37    #[serde(default, skip_serializing_if = "Option::is_none")]
38    pub destructive: Option<bool>,
39    #[serde(default, skip_serializing_if = "Option::is_none")]
40    pub idempotent: Option<bool>,
41    #[serde(default, skip_serializing_if = "Option::is_none")]
42    pub open_world: Option<bool>,
43}
44
45impl ToolAnnotationMatcher {
46    pub fn matches(&self, tool: &ToolDefinition) -> bool {
47        let Some(annotations) = tool.annotations.as_ref() else { return false };
48        let pairs = [
49            (self.read_only, annotations.read_only_hint),
50            (self.destructive, annotations.destructive_hint),
51            (self.idempotent, annotations.idempotent_hint),
52            (self.open_world, annotations.open_world_hint),
53        ];
54        if pairs.iter().all(|(field, _)| field.is_none()) {
55            return false;
56        }
57        pairs.iter().all(|(field, hint)| field.is_none_or(|value| *hint == Some(value)))
58    }
59}
60
61/// Filter for restricting which MCP tools an agent may discover and execute.
62#[derive(Debug, Clone, Default, PartialEq, Eq, serde::Deserialize, serde::Serialize, schemars::JsonSchema)]
63#[serde(rename_all = "camelCase", deny_unknown_fields)]
64pub struct ToolFilter {
65    #[serde(default, skip_serializing_if = "Vec::is_empty")]
66    pub allow: Vec<ToolMatcher>,
67    #[serde(default, skip_serializing_if = "Vec::is_empty")]
68    pub deny: Vec<ToolMatcher>,
69}
70
71impl ToolFilter {
72    pub fn is_empty(&self) -> bool {
73        self.allow.is_empty() && self.deny.is_empty()
74    }
75
76    pub fn apply(&self, tools: Vec<ToolDefinition>) -> Vec<ToolDefinition> {
77        tools.into_iter().filter(|tool| self.is_tool_allowed(tool)).collect()
78    }
79
80    pub fn is_tool_allowed(&self, tool: &ToolDefinition) -> bool {
81        let allowed = self.allow.is_empty() || self.allow.iter().any(|matcher| matcher.matches(tool));
82        let denied = self.deny.iter().any(|matcher| matcher.matches(tool));
83        allowed && !denied
84    }
85}