Skip to main content

tower_mcp/tool/
validation.rs

1//! Protocol shape checks and opt-in JSON Schema validation for tools.
2//!
3//! With the `schema-validation` feature, every tool checks `arguments` against
4//! its advertised `inputSchema` before the handler runs, and a result's
5//! `structuredContent` against its `outputSchema` before the result is
6//! returned. Both schemas are compiled once, when the tool is built.
7//!
8//! The check wraps the tool's service, MRTR handler, and live handler, the same
9//! places [`Tool::with_guard`] wraps, so it covers every handler kind and both
10//! synchronous and task-backed calls. The structured-content shape check
11//! runs independently of the feature and the declared schemas.
12
13use super::*;
14
15/// Per-tool switches for schema validation, set on [`ToolBuilder`].
16///
17/// Both default to on. Output also controls the always-on shape check.
18#[derive(Debug, Clone, Copy)]
19#[cfg_attr(not(feature = "schema-validation"), allow(dead_code))]
20pub(crate) struct SchemaValidation {
21    pub(crate) input: bool,
22    pub(crate) output: bool,
23}
24
25impl Default for SchemaValidation {
26    fn default() -> Self {
27        Self {
28            input: true,
29            output: true,
30        }
31    }
32}
33
34/// Check that a JSON Schema compiles, with the options tool schemas are
35/// validated with.
36///
37/// A tool whose `inputSchema` or `outputSchema` does not compile is still
38/// built: the failure is logged on `mcp::tools` and that schema is left
39/// unchecked, because a schema taken from a backend or a proxy must not take
40/// the process down. Code that builds tools from definitions it loads at run
41/// time (a plugin, a script, a config file) can call this first and report a
42/// bad schema against its source instead.
43///
44/// The options are the ones every tool schema is compiled with: the draft is
45/// detected from `$schema` (2020-12 when absent), `format` is enforced, and no
46/// reference is resolved from the network or the filesystem, so a remote or
47/// `file:` `$ref` does not compile. Building a tool does not call this and
48/// its behavior does not change.
49///
50/// ```rust
51/// use serde_json::json;
52/// use tower_mcp::tool::check_schema;
53///
54/// let schema = json!({
55///     "type": "object",
56///     "properties": { "path": { "type": "string", "minLength": 1 } },
57///     "required": ["path"]
58/// });
59/// assert!(check_schema(&schema).is_ok());
60///
61/// // `required` must be a list of names.
62/// let error = check_schema(&json!({ "type": "object", "required": "path" })).unwrap_err();
63/// assert_eq!(error.location(), "/required");
64/// assert!(error.to_string().starts_with("at /required: "), "{error}");
65/// ```
66#[cfg(feature = "schema-validation")]
67pub fn check_schema(schema: &Value) -> std::result::Result<(), SchemaError> {
68    engine::options()
69        .build(schema)
70        .map(drop)
71        .map_err(|error| SchemaError {
72            location: error.instance_path().as_str().to_string(),
73            message: error.to_string(),
74        })
75}
76
77/// Why a JSON Schema does not compile, returned by [`check_schema`].
78///
79/// Displays as `at <location>: <message>`, with `(root)` for the schema
80/// itself, the same form a failed argument check uses.
81#[cfg(feature = "schema-validation")]
82#[derive(Debug, Clone, PartialEq, Eq)]
83pub struct SchemaError {
84    location: String,
85    message: String,
86}
87
88#[cfg(feature = "schema-validation")]
89impl SchemaError {
90    /// A JSON Pointer to the part of the schema at fault; empty for the
91    /// schema itself.
92    pub fn location(&self) -> &str {
93        &self.location
94    }
95
96    /// What is wrong at [`location`](Self::location).
97    pub fn message(&self) -> &str {
98        &self.message
99    }
100}
101
102#[cfg(feature = "schema-validation")]
103impl std::fmt::Display for SchemaError {
104    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105        let location = match self.location.as_str() {
106            "" => "(root)",
107            path => path,
108        };
109        write!(f, "at {location}: {}", self.message)
110    }
111}
112
113#[cfg(feature = "schema-validation")]
114impl std::error::Error for SchemaError {}
115
116impl Tool {
117    /// Enforce this tool's schemas around every handler it has.
118    ///
119    /// Compiles the input schema and, when one is declared, the output schema.
120    /// A schema that does not compile is logged and left unchecked; it does
121    /// not panic, because a tool built from a backend's schema (a proxy, a
122    /// dynamic registration) must not take the process down. Callers that
123    /// want the failure instead use [`check_schema`] before building.
124    /// Always wraps because the structured-content shape check applies to
125    /// every tool regardless of the feature or declared schemas (#1540), so
126    /// no tool may bypass the wrapper.
127    pub(crate) fn with_schema_validation(self, validation: SchemaValidation) -> Self {
128        let schemas = Arc::new(engine::ToolChecks {
129            // Like the schema check, diagnostics retain the build-time name;
130            // later with_name_prefix calls (including nesting) are not reflected.
131            tool: self.name.clone(),
132            check_shape: validation.output,
133            #[cfg(feature = "schema-validation")]
134            schemas: engine::ToolSchemas::compile(
135                &self.name,
136                validation.input.then_some(&self.input_schema),
137                validation
138                    .output
139                    .then_some(self.output_schema.as_ref())
140                    .flatten(),
141            ),
142        });
143        Tool {
144            service: self.service.map(|inner| {
145                BoxCloneService::new(engine::ValidateService {
146                    inner,
147                    schemas: schemas.clone(),
148                })
149            }),
150            #[cfg(feature = "stateless")]
151            mrtr_handler: self.mrtr_handler.map(|inner| {
152                Arc::new(engine::ValidatedMrtrToolHandler {
153                    inner,
154                    schemas: schemas.clone(),
155                }) as Arc<dyn MrtrToolHandler>
156            }),
157            live_handler: self.live_handler.map(|inner| {
158                Arc::new(engine::ValidatedLiveToolHandler {
159                    inner,
160                    schemas: schemas.clone(),
161                }) as Arc<dyn LiveToolHandler>
162            }),
163            ..self
164        }
165    }
166}
167
168mod engine {
169    use super::*;
170
171    #[cfg(feature = "schema-validation")]
172    use jsonschema::Validator;
173
174    /// Most violations named in one error result. A model needs the first few
175    /// to correct a call; the rest would only bloat the message.
176    #[cfg(feature = "schema-validation")]
177    const MAX_REPORTED_ERRORS: usize = 5;
178
179    /// Longest a single violation is allowed to run, in characters. A
180    /// violation echoes the offending value, which can be arbitrarily large.
181    #[cfg(feature = "schema-validation")]
182    const MAX_ERROR_CHARS: usize = 300;
183
184    /// The compiled schemas of one tool.
185    #[cfg(feature = "schema-validation")]
186    pub(super) struct ToolSchemas {
187        tool: String,
188        input: Option<Validator>,
189        output: Option<Validator>,
190    }
191
192    /// The options every tool schema is compiled with, shared with
193    /// [`check_schema`](super::check_schema) so the two cannot disagree.
194    #[cfg(feature = "schema-validation")]
195    pub(super) fn options() -> jsonschema::ValidationOptions<'static> {
196        // Formats are annotations by default in draft 2019-09 and later. The
197        // schemas here are contracts with a caller, so `format` is enforced.
198        jsonschema::options().should_validate_formats(true)
199    }
200
201    #[cfg(feature = "schema-validation")]
202    fn compile(tool: &str, which: &'static str, schema: &Value) -> Option<Validator> {
203        match options().build(schema) {
204            Ok(validator) => Some(validator),
205            Err(error) => {
206                tracing::error!(
207                    target: "mcp::tools",
208                    tool,
209                    schema = which,
210                    %error,
211                    "tool schema does not compile; it will not be validated"
212                );
213                None
214            }
215        }
216    }
217
218    /// One line per violation, up to [`MAX_REPORTED_ERRORS`], or `None` when
219    /// the instance is valid.
220    #[cfg(feature = "schema-validation")]
221    fn violations(validator: &Validator, instance: &Value) -> Option<String> {
222        let mut errors = validator.iter_errors(instance);
223        let lines: Vec<String> = errors
224            .by_ref()
225            .take(MAX_REPORTED_ERRORS)
226            .map(|error| {
227                let location = match error.instance_path().as_str() {
228                    "" => "(root)",
229                    path => path,
230                };
231                let mut message = error.to_string();
232                if let Some((end, _)) = message.char_indices().nth(MAX_ERROR_CHARS) {
233                    message.truncate(end);
234                    message.push_str("...");
235                }
236                // The last segment of the schema path is the keyword that failed.
237                match error.schema_path().as_str().rsplit('/').next() {
238                    Some(rule) if !rule.is_empty() => {
239                        format!("- at {location}: {message} ({rule})")
240                    }
241                    _ => format!("- at {location}: {message}"),
242                }
243            })
244            .collect();
245        if lines.is_empty() {
246            return None;
247        }
248        let mut report = lines.join("\n");
249        if errors.next().is_some() {
250            report.push_str("\n- and more violations");
251        }
252        Some(report)
253    }
254
255    #[cfg(feature = "schema-validation")]
256    impl ToolSchemas {
257        /// `None` when neither schema is to be checked, or neither compiled.
258        pub(super) fn compile(
259            tool: &str,
260            input: Option<&Value>,
261            output: Option<&Value>,
262        ) -> Option<Self> {
263            let input = input.and_then(|schema| compile(tool, "inputSchema", schema));
264            let output = output.and_then(|schema| compile(tool, "outputSchema", schema));
265            if input.is_none() && output.is_none() {
266                return None;
267            }
268            Some(Self {
269                tool: tool.to_string(),
270                input,
271                output,
272            })
273        }
274
275        /// The error message for arguments that violate the input schema.
276        fn check_input(&self, arguments: &Value) -> Option<String> {
277            let validator = self.input.as_ref()?;
278            // `arguments` is optional in `tools/call`, and an omitted one
279            // arrives as null. An input schema describes an object, so an
280            // omitted value is checked as an empty one, unless the schema
281            // takes null itself (a handler declared over `()`).
282            let empty = Value::Object(Map::new());
283            let instance = if arguments.is_null() {
284                if validator.is_valid(arguments) {
285                    return None;
286                }
287                &empty
288            } else {
289                arguments
290            };
291            let report = violations(validator, instance)?;
292            tracing::debug!(
293                target: "mcp::tools",
294                tool = %self.tool,
295                %report,
296                "tool arguments do not match the input schema"
297            );
298            // Same prefix as the serde failure a typed handler reports, so a
299            // caller matching on it sees one shape with or without the feature.
300            Some(format!(
301                "Invalid input: the arguments do not match the tool's input schema.\n{report}"
302            ))
303        }
304
305        /// The result itself, or an error result when its `structuredContent`
306        /// violates the output schema.
307        fn check_output(&self, result: CallToolResult) -> CallToolResult {
308            let Some(validator) = self.output.as_ref() else {
309                return result;
310            };
311            // An error result is not a conforming structured result and
312            // carries no structured content of its own to check.
313            if result.is_error {
314                return result;
315            }
316            let Some(structured) = result.structured_content.as_ref() else {
317                return result;
318            };
319            let Some(report) = violations(validator, structured) else {
320                return result;
321            };
322            // The caller cannot fix this: the handler produced a result that
323            // contradicts the schema it advertises.
324            tracing::warn!(
325                target: "mcp::tools",
326                tool = %self.tool,
327                %report,
328                "tool structuredContent does not match the output schema"
329            );
330            CallToolResult::error(format!(
331                "Tool returned structured content that does not match its output schema.\n{report}"
332            ))
333        }
334    }
335
336    pub(super) struct ToolChecks {
337        pub(super) tool: String,
338        pub(super) check_shape: bool,
339        #[cfg(feature = "schema-validation")]
340        pub(super) schemas: Option<ToolSchemas>,
341    }
342
343    impl ToolChecks {
344        fn check_input(&self, _arguments: &Value) -> Option<String> {
345            #[cfg(feature = "schema-validation")]
346            return self.schemas.as_ref()?.check_input(_arguments);
347            #[cfg(not(feature = "schema-validation"))]
348            None
349        }
350
351        fn check_output(&self, result: CallToolResult, final_protocol: bool) -> CallToolResult {
352            // #1540: constructors cannot know the request revision. Older MCP
353            // revisions require objects, while 2026-07-28 permits any JSON value.
354            if self.check_shape
355                && !final_protocol
356                && !result.is_error
357                && let Some(structured) = result.structured_content.as_ref()
358            {
359                let kind = match structured {
360                    Value::Object(_) => None,
361                    Value::Array(_) => Some("array"),
362                    Value::String(_) => Some("string"),
363                    Value::Number(_) => Some("number"),
364                    Value::Bool(_) => Some("boolean"),
365                    Value::Null => Some("null"),
366                };
367                if let Some(kind) = kind {
368                    let message = format!(
369                        "Tool '{}' returned structuredContent of type {kind}; MCP 2025-11-25 and earlier require a JSON object. Wrap a list with CallToolResult::from_list or return an object.",
370                        self.tool,
371                    );
372                    tracing::warn!(
373                        target: "mcp::tools",
374                        tool = %self.tool,
375                        %message,
376                        "tool structuredContent has an invalid shape for the protocol revision"
377                    );
378                    return CallToolResult::error(message);
379                }
380            }
381            #[cfg(feature = "schema-validation")]
382            if let Some(schemas) = &self.schemas {
383                return schemas.check_output(result);
384            }
385            result
386        }
387    }
388
389    /// Validates a tool service's arguments before it runs and its result
390    /// after.
391    #[derive(Clone)]
392    pub(super) struct ValidateService<S> {
393        pub(super) inner: S,
394        pub(super) schemas: Arc<ToolChecks>,
395    }
396
397    impl<S> Service<ToolRequest> for ValidateService<S>
398    where
399        S: Service<ToolRequest, Response = CallToolResult, Error = Infallible>,
400        S::Future: Send + 'static,
401    {
402        type Response = CallToolResult;
403        type Error = Infallible;
404        type Future = BoxFuture<'static, std::result::Result<CallToolResult, Infallible>>;
405
406        fn poll_ready(
407            &mut self,
408            cx: &mut Context<'_>,
409        ) -> Poll<std::result::Result<(), Self::Error>> {
410            self.inner.poll_ready(cx)
411        }
412
413        fn call(&mut self, req: ToolRequest) -> Self::Future {
414            if let Some(message) = self.schemas.check_input(&req.args) {
415                return Box::pin(std::future::ready(Ok(CallToolResult::error(message))));
416            }
417            let final_protocol = crate::router::is_final_protocol_request(req.ctx.extensions());
418            let future = self.inner.call(req);
419            let schemas = self.schemas.clone();
420            Box::pin(async move {
421                future
422                    .await
423                    .map(|result| schemas.check_output(result, final_protocol))
424            })
425        }
426    }
427
428    /// The same check around an MRTR handler. Only a completed outcome has a
429    /// result to check; an input-required outcome passes through.
430    #[cfg(feature = "stateless")]
431    pub(super) struct ValidatedMrtrToolHandler {
432        pub(super) inner: Arc<dyn MrtrToolHandler>,
433        pub(super) schemas: Arc<ToolChecks>,
434    }
435
436    #[cfg(feature = "stateless")]
437    impl MrtrToolHandler for ValidatedMrtrToolHandler {
438        fn call(
439            &self,
440            ctx: RequestContext,
441            args: Value,
442        ) -> BoxFuture<'_, Result<RequestOutcome<CallToolResult>>> {
443            if let Some(message) = self.schemas.check_input(&args) {
444                return Box::pin(async move {
445                    Ok(RequestOutcome::Complete(CallToolResult::error(message)))
446                });
447            }
448            let final_protocol = crate::router::is_final_protocol_request(ctx.extensions());
449            let future = self.inner.call(ctx, args);
450            Box::pin(async move {
451                Ok(match future.await? {
452                    RequestOutcome::Complete(result) => {
453                        RequestOutcome::Complete(self.schemas.check_output(result, final_protocol))
454                    }
455                    other => other,
456                })
457            })
458        }
459
460        fn input_schema(&self) -> Value {
461            self.inner.input_schema()
462        }
463    }
464
465    /// The same check around a live handler. A rejected call completes the
466    /// task with an error result, as a rejected guard does.
467    pub(super) struct ValidatedLiveToolHandler {
468        pub(super) inner: Arc<dyn LiveToolHandler>,
469        pub(super) schemas: Arc<ToolChecks>,
470    }
471
472    #[async_trait::async_trait]
473    impl LiveToolHandler for ValidatedLiveToolHandler {
474        async fn call(
475            &self,
476            ctx: RequestContext,
477            task: TaskContext,
478            arguments: Value,
479        ) -> Result<TaskOutcome> {
480            if let Some(message) = self.schemas.check_input(&arguments) {
481                return Ok(TaskOutcome::Completed(CallToolResult::error(message)));
482            }
483            let final_protocol = crate::router::is_final_protocol_request(ctx.extensions());
484            Ok(match self.inner.call(ctx, task, arguments).await? {
485                TaskOutcome::Completed(result) => {
486                    TaskOutcome::Completed(self.schemas.check_output(result, final_protocol))
487                }
488                other => other,
489            })
490        }
491    }
492}
493
494#[cfg(test)]
495mod tests;