1use car_ir::ToolSchema;
8use serde::{Deserialize, Serialize};
9use std::collections::HashMap;
10use tokio::sync::RwLock;
11
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
14#[serde(rename_all = "snake_case")]
15#[derive(Default)]
16pub enum ToolPermission {
17 Allow,
19 #[default]
21 AskUser,
22 Deny,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
28#[serde(rename_all = "snake_case")]
29pub enum ToolSource {
30 Builtin,
32 UserDefined,
34 Subprocess,
36 Mcp { server_name: String },
38}
39
40impl ToolSource {
41 pub const fn kind(&self) -> car_ir::ToolSourceKind {
45 match self {
46 Self::Builtin => car_ir::ToolSourceKind::Builtin,
47 Self::UserDefined => car_ir::ToolSourceKind::UserDefined,
48 Self::Subprocess => car_ir::ToolSourceKind::Subprocess,
49 Self::Mcp { server_name: _ } => car_ir::ToolSourceKind::Mcp,
50 }
51 }
52}
53
54#[derive(Debug, Clone)]
56pub struct ToolEntry {
57 pub schema: ToolSchema,
59 pub permission: ToolPermission,
61 pub source: ToolSource,
63 pub side_effects: bool,
65 pub category: Option<String>,
67}
68
69impl ToolEntry {
70 pub fn new(mut schema: ToolSchema) -> Self {
73 schema.source = car_ir::ToolSourceKind::UserDefined;
74 Self {
75 schema,
76 permission: ToolPermission::default(),
77 source: ToolSource::UserDefined,
78 side_effects: true,
79 category: None,
80 }
81 }
82
83 pub fn builtin(mut schema: ToolSchema) -> Self {
84 schema.source = car_ir::ToolSourceKind::Builtin;
85 Self {
86 permission: ToolPermission::Allow,
87 source: ToolSource::Builtin,
88 side_effects: false,
89 category: None,
90 schema,
91 }
92 }
93
94 pub fn with_permission(mut self, perm: ToolPermission) -> Self {
95 self.permission = perm;
96 self
97 }
98
99 pub fn with_source(mut self, source: ToolSource) -> Self {
100 self.schema.source = source.kind();
101 self.source = source;
102 self
103 }
104
105 pub fn with_side_effects(mut self, side_effects: bool) -> Self {
106 self.side_effects = side_effects;
107 self
108 }
109
110 pub fn with_category(mut self, category: &str) -> Self {
111 self.category = Some(category.to_string());
112 self
113 }
114}
115
116#[derive(Debug, Clone)]
118pub struct RegistryValidationError {
119 pub tool_name: String,
120 pub message: String,
121}
122
123pub struct ToolRegistry {
125 entries: RwLock<HashMap<String, ToolEntry>>,
126}
127
128impl ToolRegistry {
129 pub fn new() -> Self {
130 Self {
131 entries: RwLock::new(HashMap::new()),
132 }
133 }
134
135 pub async fn register(&self, mut entry: ToolEntry) {
137 entry.schema.source = entry.source.kind();
140 let name = entry.schema.name.clone();
141 self.entries.write().await.insert(name, entry);
142 }
143
144 pub fn try_register(&self, entry: ToolEntry) -> bool {
156 match self.entries.try_write() {
157 Ok(mut entries) => {
158 entries.insert(entry.schema.name.clone(), entry);
159 true
160 }
161 Err(_) => false,
162 }
163 }
164
165 pub async fn get(&self, name: &str) -> Option<ToolEntry> {
167 self.entries.read().await.get(name).cloned()
168 }
169
170 pub async fn contains(&self, name: &str) -> bool {
172 self.entries.read().await.contains_key(name)
173 }
174
175 pub async fn remove(&self, name: &str) -> Option<ToolEntry> {
177 self.entries.write().await.remove(name)
178 }
179
180 pub async fn names(&self) -> Vec<String> {
182 self.entries.read().await.keys().cloned().collect()
183 }
184
185 pub async fn entries(&self) -> Vec<ToolEntry> {
187 self.entries.read().await.values().cloned().collect()
188 }
189
190 pub async fn schemas(&self) -> Vec<ToolSchema> {
192 self.entries
193 .read()
194 .await
195 .values()
196 .map(|e| e.schema.clone())
197 .collect()
198 }
199
200 pub async fn allowed_schemas(&self) -> Vec<ToolSchema> {
202 self.entries
203 .read()
204 .await
205 .values()
206 .filter(|e| e.permission != ToolPermission::Deny)
207 .map(|e| e.schema.clone())
208 .collect()
209 }
210
211 pub async fn by_source(&self, source_match: &ToolSource) -> Vec<ToolEntry> {
213 self.entries
214 .read()
215 .await
216 .values()
217 .filter(|e| std::mem::discriminant(&e.source) == std::mem::discriminant(source_match))
218 .cloned()
219 .collect()
220 }
221
222 pub async fn by_category(&self, category: &str) -> Vec<ToolEntry> {
224 self.entries
225 .read()
226 .await
227 .values()
228 .filter(|e| e.category.as_deref() == Some(category))
229 .cloned()
230 .collect()
231 }
232
233 pub async fn validate(&self) -> Vec<RegistryValidationError> {
235 let entries = self.entries.read().await;
236 let mut errors = Vec::new();
237 for (name, entry) in entries.iter() {
238 if entry.schema.name != *name {
239 errors.push(RegistryValidationError {
240 tool_name: name.clone(),
241 message: format!(
242 "schema name '{}' doesn't match registry key '{}'",
243 entry.schema.name, name
244 ),
245 });
246 }
247 if entry.schema.description.is_empty() {
248 errors.push(RegistryValidationError {
249 tool_name: name.clone(),
250 message: "missing description".to_string(),
251 });
252 }
253 }
254 errors
255 }
256
257 pub async fn to_schema_map(&self) -> HashMap<String, ToolSchema> {
259 self.entries
260 .read()
261 .await
262 .iter()
263 .map(|(k, v)| (k.clone(), v.schema.clone()))
264 .collect()
265 }
266
267 pub async fn len(&self) -> usize {
269 self.entries.read().await.len()
270 }
271
272 pub async fn is_empty(&self) -> bool {
273 self.entries.read().await.is_empty()
274 }
275}
276
277impl Default for ToolRegistry {
278 fn default() -> Self {
279 Self::new()
280 }
281}
282
283#[cfg(test)]
284mod tests {
285 use super::*;
286
287 fn test_schema(name: &str) -> ToolSchema {
288 ToolSchema {
289 name: name.to_string(),
290 source: car_ir::ToolSourceKind::UserDefined,
291 description: format!("{} tool", name),
292 parameters: serde_json::json!({"type": "object"}),
293 returns: None,
294 idempotent: false,
295 cache_ttl_secs: None,
296 rate_limit: None,
297 }
298 }
299
300 #[tokio::test]
301 async fn test_register_and_get() {
302 let reg = ToolRegistry::new();
303 let entry = ToolEntry::new(test_schema("search"))
304 .with_permission(ToolPermission::Allow)
305 .with_category("network");
306 reg.register(entry).await;
307
308 let got = reg.get("search").await.unwrap();
309 assert_eq!(got.schema.name, "search");
310 assert_eq!(got.schema.source, car_ir::ToolSourceKind::UserDefined);
311 assert_eq!(got.permission, ToolPermission::Allow);
312 assert_eq!(got.category.as_deref(), Some("network"));
313 }
314
315 #[tokio::test]
316 async fn test_allowed_schemas_excludes_denied() {
317 let reg = ToolRegistry::new();
318 reg.register(ToolEntry::new(test_schema("read")).with_permission(ToolPermission::Allow))
319 .await;
320 reg.register(ToolEntry::new(test_schema("delete")).with_permission(ToolPermission::Deny))
321 .await;
322 reg.register(ToolEntry::new(test_schema("write")).with_permission(ToolPermission::AskUser))
323 .await;
324
325 let allowed = reg.allowed_schemas().await;
326 assert_eq!(allowed.len(), 2);
327 assert!(allowed.iter().all(|s| s.name != "delete"));
328 }
329
330 #[tokio::test]
331 async fn test_validation() {
332 let reg = ToolRegistry::new();
333 let mut bad_schema = test_schema("good");
334 bad_schema.description = String::new();
335 reg.register(ToolEntry::new(bad_schema)).await;
336
337 let errors = reg.validate().await;
338 assert_eq!(errors.len(), 1);
339 assert!(errors[0].message.contains("missing description"));
340 }
341
342 #[tokio::test]
343 async fn test_by_source() {
344 let reg = ToolRegistry::new();
345 let mut builtin = ToolEntry::builtin(test_schema("infer"));
346 builtin.schema.source = car_ir::ToolSourceKind::UserDefined;
349 reg.register(builtin).await;
350 reg.register(ToolEntry::new(test_schema("search"))).await;
351
352 let builtins = reg.by_source(&ToolSource::Builtin).await;
353 assert_eq!(builtins.len(), 1);
354 assert_eq!(builtins[0].schema.name, "infer");
355 assert_eq!(builtins[0].schema.source, car_ir::ToolSourceKind::Builtin);
356 }
357}