Skip to main content

max / makenotwork

13.6 KB · 463 lines History Blame Raw
1 use std::collections::{HashMap, HashSet};
2 use std::sync::Arc;
3
4 use async_trait::async_trait;
5 use serde::{Deserialize, Serialize};
6 use serde_json::Value;
7
8 use crate::error::{Error, Result};
9
10 /// A single tool exposed to an MCP client.
11 ///
12 /// Implementors must declare `kind()` — either `Read` or `Write(capability)`.
13 /// There is no default, because forgetting the classification is the sort of
14 /// hazard we can prevent at compile time.
15 #[async_trait]
16 pub trait Tool: Send + Sync + 'static {
17 fn name(&self) -> &str;
18 fn description(&self) -> &str;
19
20 /// Read or Write. Writes carry a capability that must be granted per session.
21 fn kind(&self) -> ToolKind;
22
23 /// JSON Schema for the tool's arguments. Should describe a single object.
24 fn input_schema(&self) -> Value;
25
26 async fn call(&self, args: Value) -> Result<ToolCallResult>;
27
28 /// If this tool is a hard refusal, return a stable, model-visible reason.
29 /// Registered refusals appear in `tools/list` so external clients can see
30 /// what the app has chosen not to offer, and the model gets a clear signal
31 /// instead of inventing tool names.
32 fn refusal_reason(&self) -> Option<&str> {
33 None
34 }
35
36 /// Opt-in flag for the compact projection used by small models. Defaults
37 /// to `false` — conservative: apps opt tools into the compact surface only
38 /// after verifying a small model can call them reliably.
39 fn small_model_safe(&self) -> bool {
40 false
41 }
42 }
43
44 /// Classification of a tool's side-effect profile.
45 #[derive(Debug, Clone, PartialEq, Eq)]
46 pub enum ToolKind {
47 /// No state change. Freely callable subject to app-level connection auth.
48 Read,
49 /// Mutates app state. Requires the caller to hold the named capability.
50 Write(WriteCapability),
51 }
52
53 impl ToolKind {
54 pub fn is_write(&self) -> bool {
55 matches!(self, ToolKind::Write(_))
56 }
57
58 pub fn capability_id(&self) -> Option<&str> {
59 match self {
60 ToolKind::Write(cap) => Some(&cap.id),
61 ToolKind::Read => None,
62 }
63 }
64 }
65
66 /// A named, pre-declared write capability.
67 ///
68 /// `id` is a stable public contract owned by the app — treat it like an API
69 /// surface identifier. If you rename a capability, you introduce a new one and
70 /// deprecate the old, you do not silently reassign an existing id.
71 #[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
72 pub struct WriteCapability {
73 pub id: String,
74 pub description: String,
75 }
76
77 impl WriteCapability {
78 pub fn new(id: impl Into<String>, description: impl Into<String>) -> Self {
79 Self {
80 id: id.into(),
81 description: description.into(),
82 }
83 }
84 }
85
86 /// Result of a single tool invocation.
87 ///
88 /// Shape mirrors MCP's `CallToolResult`: a list of content parts plus an
89 /// `is_error` flag. MVP carries only text content; media/resource parts land
90 /// alongside audiofiles integration.
91 #[derive(Debug, Clone, Serialize, Deserialize)]
92 pub struct ToolCallResult {
93 pub content: Vec<ContentPart>,
94 #[serde(default, rename = "isError")]
95 pub is_error: bool,
96 }
97
98 #[derive(Debug, Clone, Serialize, Deserialize)]
99 #[serde(tag = "type", rename_all = "snake_case")]
100 pub enum ContentPart {
101 Text { text: String },
102 }
103
104 impl ToolCallResult {
105 pub fn text(text: impl Into<String>) -> Self {
106 Self {
107 content: vec![ContentPart::Text { text: text.into() }],
108 is_error: false,
109 }
110 }
111
112 pub fn error(text: impl Into<String>) -> Self {
113 Self {
114 content: vec![ContentPart::Text { text: text.into() }],
115 is_error: true,
116 }
117 }
118 }
119
120 /// Serializable description of a tool for `tools/list` responses and for
121 /// translating to provider-native schemas (Ollama, OpenAI, Anthropic).
122 #[derive(Debug, Clone, Serialize, Deserialize)]
123 pub struct ToolSpec {
124 pub name: String,
125 pub description: String,
126 #[serde(rename = "inputSchema")]
127 pub input_schema: Value,
128
129 /// `"read"` or `"write"`. Extension over the MCP base spec; clients that
130 /// don't understand this can ignore it.
131 #[serde(rename = "kburgKind")]
132 pub kind: &'static str,
133
134 /// Only set when `kind == "write"`. Extension.
135 #[serde(rename = "kburgCapability", skip_serializing_if = "Option::is_none")]
136 pub capability: Option<WriteCapability>,
137
138 /// Set when this tool is a registered refusal. Extension.
139 #[serde(rename = "kburgRefusalReason", skip_serializing_if = "Option::is_none")]
140 pub refusal_reason: Option<String>,
141 }
142
143 /// Which projection of the registry to expose to a given client.
144 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
145 pub enum SurfaceProjection {
146 /// All registered tools.
147 Full,
148 /// Only tools annotated `small_model_safe`.
149 Compact,
150 }
151
152 impl SurfaceProjection {
153 fn includes(self, small_model_safe: bool) -> bool {
154 match self {
155 SurfaceProjection::Full => true,
156 SurfaceProjection::Compact => small_model_safe,
157 }
158 }
159 }
160
161 /// A registered "hard refusal" — appears in `tools/list` and returns a stable
162 /// error on invocation, preventing the model from inventing plausible names.
163 pub struct Refusal {
164 name: String,
165 description: String,
166 reason: String,
167 }
168
169 impl Refusal {
170 pub fn new(
171 name: impl Into<String>,
172 description: impl Into<String>,
173 reason: impl Into<String>,
174 ) -> Self {
175 Self {
176 name: name.into(),
177 description: description.into(),
178 reason: reason.into(),
179 }
180 }
181 }
182
183 #[async_trait]
184 impl Tool for Refusal {
185 fn name(&self) -> &str {
186 &self.name
187 }
188 fn description(&self) -> &str {
189 &self.description
190 }
191 fn kind(&self) -> ToolKind {
192 ToolKind::Read
193 }
194 fn input_schema(&self) -> Value {
195 serde_json::json!({ "type": "object", "properties": {} })
196 }
197 fn refusal_reason(&self) -> Option<&str> {
198 Some(&self.reason)
199 }
200 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
201 Err(Error::Refused {
202 tool: self.name.clone(),
203 reason: self.reason.clone(),
204 })
205 }
206 }
207
208 /// Registry of tools available to clients.
209 ///
210 /// Cheaply cloneable — the underlying map lives behind `Arc` so the server
211 /// and any agents can share the same registry without copying.
212 #[derive(Clone, Default)]
213 pub struct ToolRegistry {
214 tools: Arc<HashMap<String, Arc<dyn Tool>>>,
215 }
216
217 impl ToolRegistry {
218 pub fn new() -> Self {
219 Self::default()
220 }
221
222 /// Register a tool. Panics on duplicate names or if called after cloning
223 /// (register all tools before serving).
224 pub fn register<T: Tool>(&mut self, tool: T) {
225 let name = tool.name().to_string();
226 let map = Arc::get_mut(&mut self.tools).expect(
227 "ToolRegistry::register called after cloning; register all tools before serving",
228 );
229 assert!(
230 !map.contains_key(&name),
231 "duplicate tool registration: {name}"
232 );
233 map.insert(name, Arc::new(tool));
234 }
235
236 pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
237 self.tools.get(name).cloned()
238 }
239
240 /// All specs, unfiltered. Equivalent to `specs_projected(Full)`.
241 pub fn specs(&self) -> Vec<ToolSpec> {
242 self.specs_projected(SurfaceProjection::Full)
243 }
244
245 pub fn specs_projected(&self, projection: SurfaceProjection) -> Vec<ToolSpec> {
246 self.tools
247 .values()
248 .filter(|t| projection.includes(t.small_model_safe()))
249 .map(|t| ToolSpec {
250 name: t.name().to_string(),
251 description: t.description().to_string(),
252 input_schema: t.input_schema(),
253 kind: match t.kind() {
254 ToolKind::Read => "read",
255 ToolKind::Write(_) => "write",
256 },
257 capability: match t.kind() {
258 ToolKind::Write(cap) => Some(cap),
259 ToolKind::Read => None,
260 },
261 refusal_reason: t.refusal_reason().map(str::to_string),
262 })
263 .collect()
264 }
265
266 /// Distinct write capabilities declared by registered tools.
267 pub fn write_capabilities(&self) -> Vec<WriteCapability> {
268 let mut seen = HashSet::new();
269 let mut out = Vec::new();
270 for tool in self.tools.values() {
271 if let ToolKind::Write(cap) = tool.kind()
272 && seen.insert(cap.id.clone())
273 {
274 out.push(cap);
275 }
276 }
277 out
278 }
279
280 /// Call a tool, enforcing the granted-capabilities set. `grants` may be
281 /// `None` to bypass the check (e.g., for a fully-trusted in-process
282 /// caller); pass `Some(set)` from any surface exposed to an LLM.
283 pub async fn call(
284 &self,
285 name: &str,
286 args: Value,
287 grants: Option<&HashSet<String>>,
288 ) -> Result<ToolCallResult> {
289 let tool = self
290 .get(name)
291 .ok_or_else(|| Error::ToolNotFound(name.to_string()))?;
292
293 if let (ToolKind::Write(cap), Some(grants)) = (tool.kind(), grants)
294 && !grants.contains(&cap.id)
295 {
296 return Err(Error::CapabilityDenied {
297 tool: name.to_string(),
298 capability: cap.id,
299 });
300 }
301
302 tool.call(args).await
303 }
304 }
305
306 #[cfg(test)]
307 mod tests {
308 use super::*;
309 use serde_json::json;
310
311 struct Echo;
312 #[async_trait]
313 impl Tool for Echo {
314 fn name(&self) -> &'static str {
315 "echo"
316 }
317 fn description(&self) -> &'static str {
318 "read-only echo"
319 }
320 fn kind(&self) -> ToolKind {
321 ToolKind::Read
322 }
323 fn small_model_safe(&self) -> bool {
324 true
325 }
326 fn input_schema(&self) -> Value {
327 json!({ "type": "object", "properties": {} })
328 }
329 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
330 Ok(ToolCallResult::text("ok"))
331 }
332 }
333
334 struct Mutate;
335 #[async_trait]
336 impl Tool for Mutate {
337 fn name(&self) -> &'static str {
338 "mutate"
339 }
340 fn description(&self) -> &'static str {
341 "write, gated on `test.write`"
342 }
343 fn kind(&self) -> ToolKind {
344 ToolKind::Write(WriteCapability::new("test.write", "mutate test state"))
345 }
346 fn input_schema(&self) -> Value {
347 json!({ "type": "object", "properties": {} })
348 }
349 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
350 Ok(ToolCallResult::text("mutated"))
351 }
352 }
353
354 fn registry() -> ToolRegistry {
355 let mut r = ToolRegistry::new();
356 r.register(Echo);
357 r.register(Mutate);
358 r.register(Refusal::new(
359 "delete_everything",
360 "not offered",
361 "destructive bulk ops are never exposed",
362 ));
363 r
364 }
365
366 fn grants(ids: &[&str]) -> HashSet<String> {
367 ids.iter().map(std::string::ToString::to_string).collect()
368 }
369
370 #[tokio::test]
371 async fn read_is_callable_without_grants() {
372 let r = registry();
373 let out = r
374 .call("echo", json!({}), Some(&HashSet::new()))
375 .await
376 .unwrap();
377 assert!(!out.is_error);
378 }
379
380 #[tokio::test]
381 async fn write_denied_without_grant() {
382 let r = registry();
383 let err = r
384 .call("mutate", json!({}), Some(&HashSet::new()))
385 .await
386 .unwrap_err();
387 match err {
388 Error::CapabilityDenied { capability, .. } => assert_eq!(capability, "test.write"),
389 other => panic!("expected CapabilityDenied, got {other:?}"),
390 }
391 }
392
393 #[tokio::test]
394 async fn write_allowed_with_grant() {
395 let r = registry();
396 let out = r
397 .call("mutate", json!({}), Some(&grants(&["test.write"])))
398 .await
399 .unwrap();
400 assert_eq!(out.content.len(), 1);
401 assert!(!out.is_error);
402 }
403
404 #[tokio::test]
405 async fn none_grants_bypasses_the_check() {
406 let r = registry();
407 let out = r.call("mutate", json!({}), None).await.unwrap();
408 assert!(!out.is_error);
409 }
410
411 #[tokio::test]
412 async fn unknown_tool_is_not_found() {
413 let r = registry();
414 let err = r.call("nope", json!({}), None).await.unwrap_err();
415 assert!(matches!(err, Error::ToolNotFound(_)));
416 }
417
418 #[tokio::test]
419 async fn refusal_surfaces_in_specs_and_errors_on_call() {
420 let r = registry();
421 let spec = r
422 .specs()
423 .into_iter()
424 .find(|s| s.name == "delete_everything")
425 .expect("refusal appears in tools/list");
426 assert!(spec.refusal_reason.is_some());
427
428 let err = r
429 .call("delete_everything", json!({}), None)
430 .await
431 .unwrap_err();
432 assert!(matches!(err, Error::Refused { .. }));
433 }
434
435 #[test]
436 fn compact_projection_filters_to_small_model_safe() {
437 let r = registry();
438 let compact: Vec<_> = r
439 .specs_projected(SurfaceProjection::Compact)
440 .into_iter()
441 .map(|s| s.name)
442 .collect();
443 assert!(compact.contains(&"echo".to_string()));
444 assert!(!compact.contains(&"mutate".to_string()));
445 }
446
447 #[test]
448 fn write_capabilities_are_deduped() {
449 let r = registry();
450 let caps = r.write_capabilities();
451 assert_eq!(caps.len(), 1);
452 assert_eq!(caps[0].id, "test.write");
453 }
454
455 #[test]
456 #[should_panic(expected = "duplicate tool registration")]
457 fn duplicate_registration_panics() {
458 let mut r = ToolRegistry::new();
459 r.register(Echo);
460 r.register(Echo);
461 }
462 }
463