1use super::*;
14
15#[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#[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#[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 pub fn location(&self) -> &str {
93 &self.location
94 }
95
96 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 pub(crate) fn with_schema_validation(self, validation: SchemaValidation) -> Self {
128 let schemas = Arc::new(engine::ToolChecks {
129 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 #[cfg(feature = "schema-validation")]
177 const MAX_REPORTED_ERRORS: usize = 5;
178
179 #[cfg(feature = "schema-validation")]
182 const MAX_ERROR_CHARS: usize = 300;
183
184 #[cfg(feature = "schema-validation")]
186 pub(super) struct ToolSchemas {
187 tool: String,
188 input: Option<Validator>,
189 output: Option<Validator>,
190 }
191
192 #[cfg(feature = "schema-validation")]
195 pub(super) fn options() -> jsonschema::ValidationOptions<'static> {
196 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 #[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 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 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 fn check_input(&self, arguments: &Value) -> Option<String> {
277 let validator = self.input.as_ref()?;
278 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 Some(format!(
301 "Invalid input: the arguments do not match the tool's input schema.\n{report}"
302 ))
303 }
304
305 fn check_output(&self, result: CallToolResult) -> CallToolResult {
308 let Some(validator) = self.output.as_ref() else {
309 return result;
310 };
311 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 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 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 #[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 #[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 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;