Skip to main content

max / makenotwork

17.2 KB · 566 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 = "kbergKind")]
132 pub kind: &'static str,
133
134 /// Only set when `kind == "write"`. Extension.
135 #[serde(rename = "kbergCapability", 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 = "kbergRefusalReason", 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 and the tool's own
281 /// argument schema. `grants` may be `None` to bypass the capability check
282 /// (e.g., for a fully-trusted in-process caller); pass `Some(set)` from any
283 /// surface exposed to an LLM.
284 ///
285 /// Order matters: capability before schema. A caller that may not invoke a
286 /// tool at all should learn that, not receive a critique of arguments it
287 /// was never entitled to submit.
288 ///
289 /// Schema validation is not optional and has no per-tool opt-out. A tool
290 /// whose real contract its schema cannot express still validates what the
291 /// schema does say, then raises [`Error::InvalidArgs`] itself for the rest.
292 pub async fn call(
293 &self,
294 name: &str,
295 args: Value,
296 grants: Option<&HashSet<String>>,
297 ) -> Result<ToolCallResult> {
298 let tool = self
299 .get(name)
300 .ok_or_else(|| Error::ToolNotFound(name.to_string()))?;
301
302 if let (ToolKind::Write(cap), Some(grants)) = (tool.kind(), grants)
303 && !grants.contains(&cap.id)
304 {
305 return Err(Error::CapabilityDenied {
306 tool: name.to_string(),
307 capability: cap.id,
308 });
309 }
310
311 if let Err(slots) = crate::schema::validate(&tool.input_schema(), &args) {
312 return Err(Error::InvalidArguments {
313 tool: name.to_string(),
314 slots,
315 });
316 }
317
318 tool.call(args).await
319 }
320 }
321
322 #[cfg(test)]
323 mod tests {
324 use super::*;
325 use serde_json::json;
326
327 struct Echo;
328 #[async_trait]
329 impl Tool for Echo {
330 fn name(&self) -> &'static str {
331 "echo"
332 }
333 fn description(&self) -> &'static str {
334 "read-only echo"
335 }
336 fn kind(&self) -> ToolKind {
337 ToolKind::Read
338 }
339 fn small_model_safe(&self) -> bool {
340 true
341 }
342 fn input_schema(&self) -> Value {
343 json!({ "type": "object", "properties": {} })
344 }
345 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
346 Ok(ToolCallResult::text("ok"))
347 }
348 }
349
350 struct Mutate;
351 #[async_trait]
352 impl Tool for Mutate {
353 fn name(&self) -> &'static str {
354 "mutate"
355 }
356 fn description(&self) -> &'static str {
357 "write, gated on `test.write`"
358 }
359 fn kind(&self) -> ToolKind {
360 ToolKind::Write(WriteCapability::new("test.write", "mutate test state"))
361 }
362 fn input_schema(&self) -> Value {
363 json!({ "type": "object", "properties": {} })
364 }
365 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
366 Ok(ToolCallResult::text("mutated"))
367 }
368 }
369
370 fn registry() -> ToolRegistry {
371 let mut r = ToolRegistry::new();
372 r.register(Echo);
373 r.register(Mutate);
374 r.register(Refusal::new(
375 "delete_everything",
376 "not offered",
377 "destructive bulk ops are never exposed",
378 ));
379 r
380 }
381
382 fn grants(ids: &[&str]) -> HashSet<String> {
383 ids.iter().map(std::string::ToString::to_string).collect()
384 }
385
386 #[tokio::test]
387 async fn read_is_callable_without_grants() {
388 let r = registry();
389 let out = r
390 .call("echo", json!({}), Some(&HashSet::new()))
391 .await
392 .unwrap();
393 assert!(!out.is_error);
394 }
395
396 #[tokio::test]
397 async fn write_denied_without_grant() {
398 let r = registry();
399 let err = r
400 .call("mutate", json!({}), Some(&HashSet::new()))
401 .await
402 .unwrap_err();
403 match err {
404 Error::CapabilityDenied { capability, .. } => assert_eq!(capability, "test.write"),
405 other => panic!("expected CapabilityDenied, got {other:?}"),
406 }
407 }
408
409 #[tokio::test]
410 async fn write_allowed_with_grant() {
411 let r = registry();
412 let out = r
413 .call("mutate", json!({}), Some(&grants(&["test.write"])))
414 .await
415 .unwrap();
416 assert_eq!(out.content.len(), 1);
417 assert!(!out.is_error);
418 }
419
420 #[tokio::test]
421 async fn none_grants_bypasses_the_check() {
422 let r = registry();
423 let out = r.call("mutate", json!({}), None).await.unwrap();
424 assert!(!out.is_error);
425 }
426
427 #[tokio::test]
428 async fn unknown_tool_is_not_found() {
429 let r = registry();
430 let err = r.call("nope", json!({}), None).await.unwrap_err();
431 assert!(matches!(err, Error::ToolNotFound(_)));
432 }
433
434 #[tokio::test]
435 async fn refusal_surfaces_in_specs_and_errors_on_call() {
436 let r = registry();
437 let spec = r
438 .specs()
439 .into_iter()
440 .find(|s| s.name == "delete_everything")
441 .expect("refusal appears in tools/list");
442 assert!(spec.refusal_reason.is_some());
443
444 let err = r
445 .call("delete_everything", json!({}), None)
446 .await
447 .unwrap_err();
448 assert!(matches!(err, Error::Refused { .. }));
449 }
450
451 #[test]
452 fn compact_projection_filters_to_small_model_safe() {
453 let r = registry();
454 let compact: Vec<_> = r
455 .specs_projected(SurfaceProjection::Compact)
456 .into_iter()
457 .map(|s| s.name)
458 .collect();
459 assert!(compact.contains(&"echo".to_string()));
460 assert!(!compact.contains(&"mutate".to_string()));
461 }
462
463 #[test]
464 fn write_capabilities_are_deduped() {
465 let r = registry();
466 let caps = r.write_capabilities();
467 assert_eq!(caps.len(), 1);
468 assert_eq!(caps[0].id, "test.write");
469 }
470
471 #[test]
472 #[should_panic(expected = "duplicate tool registration")]
473 fn duplicate_registration_panics() {
474 let mut r = ToolRegistry::new();
475 r.register(Echo);
476 r.register(Echo);
477 }
478
479 /// Records whether its body ran, so a test can prove validation rejected
480 /// the call *before* dispatch rather than after.
481 struct Strict(Arc<std::sync::atomic::AtomicBool>);
482 #[async_trait]
483 impl Tool for Strict {
484 fn name(&self) -> &'static str {
485 "strict"
486 }
487 fn description(&self) -> &'static str {
488 "requires a string `slug`"
489 }
490 fn kind(&self) -> ToolKind {
491 ToolKind::Read
492 }
493 fn input_schema(&self) -> Value {
494 json!({
495 "type": "object",
496 "properties": { "slug": { "type": "string" } },
497 "required": ["slug"]
498 })
499 }
500 async fn call(&self, _args: Value) -> Result<ToolCallResult> {
501 self.0.store(true, std::sync::atomic::Ordering::SeqCst);
502 Ok(ToolCallResult::text("ran"))
503 }
504 }
505
506 fn strict_registry() -> (ToolRegistry, Arc<std::sync::atomic::AtomicBool>) {
507 let ran = Arc::new(std::sync::atomic::AtomicBool::new(false));
508 let mut r = ToolRegistry::new();
509 r.register(Strict(ran.clone()));
510 (r, ran)
511 }
512
513 #[tokio::test]
514 async fn bad_args_are_rejected_before_the_tool_runs() {
515 let (r, ran) = strict_registry();
516 let err = r.call("strict", json!({}), None).await.unwrap_err();
517 match err {
518 Error::InvalidArguments { tool, slots } => {
519 assert_eq!(tool, "strict");
520 assert_eq!(slots.len(), 1);
521 assert_eq!(slots[0].field, "slug");
522 }
523 other => panic!("expected InvalidArguments, got {other:?}"),
524 }
525 assert!(
526 !ran.load(std::sync::atomic::Ordering::SeqCst),
527 "tool body must not run when validation fails"
528 );
529 }
530
531 #[tokio::test]
532 async fn good_args_reach_the_tool() {
533 let (r, ran) = strict_registry();
534 let out = r
535 .call("strict", json!({ "slug": "x" }), None)
536 .await
537 .unwrap();
538 assert!(!out.is_error);
539 assert!(ran.load(std::sync::atomic::Ordering::SeqCst));
540 }
541
542 #[tokio::test]
543 async fn capability_is_checked_before_the_schema() {
544 // An ungranted caller learns it may not call the tool at all, rather
545 // than receiving a critique of arguments it was never entitled to send.
546 let r = registry();
547 let err = r
548 .call("mutate", json!({ "bogus": 1 }), Some(&HashSet::new()))
549 .await
550 .unwrap_err();
551 assert!(matches!(err, Error::CapabilityDenied { .. }));
552 }
553
554 #[tokio::test]
555 async fn argument_less_tools_still_accept_anything() {
556 // `Echo` publishes an empty property set; validation must not start
557 // rejecting the extra keys some clients send.
558 let r = registry();
559 let out = r
560 .call("echo", json!({ "stray": true }), None)
561 .await
562 .unwrap();
563 assert!(!out.is_error);
564 }
565 }
566