| 1 | """MCP protocol-edge normalization for unknown tools and malformed arguments.""" |
| 2 | from __future__ import annotations |
| 3 | |
| 4 | import json |
| 5 | from typing import Any, Awaitable, Callable |
| 6 | |
| 7 | from mcp.server.context import HandlerResult, ServerRequestContext |
| 8 | from mcp.types import CallToolResult, TextContent |
| 9 | from pydantic import ValidationError |
| 10 | |
| 11 | from app.audit import McpAuditSink |
| 12 | from app.contracts import McpAuthContext, McpErrorCode, McpResultEnvelope, normalize_request_id |
| 13 | from app.registry import McpToolRegistry |
| 14 | |
| 15 | |
| 16 | class McpProtocolGuard: |
| 17 | """Return the application error envelope before the SDK dispatches invalid calls.""" |
| 18 | |
| 19 | def __init__( |
| 20 | self, |
| 21 | registry: McpToolRegistry, |
| 22 | audit: McpAuditSink, |
| 23 | auth: McpAuthContext, |
| 24 | ) -> None: |
| 25 | self.registry = registry |
| 26 | self.audit = audit |
| 27 | self.auth = auth |
| 28 | |
| 29 | async def __call__( |
| 30 | self, |
| 31 | context: ServerRequestContext[Any, Any], |
| 32 | call_next: Callable[[ServerRequestContext[Any, Any]], Awaitable[HandlerResult]], |
| 33 | ) -> HandlerResult: |
| 34 | if context.method != "tools/call": |
| 35 | return await call_next(context) |
| 36 | params = context.params or {} |
| 37 | tool = str(params.get("name") or "") |
| 38 | definition = self.registry.get(tool) |
| 39 | if definition is None: |
| 40 | return self._failure(context, tool, McpErrorCode.MCP_TOOL_NOT_FOUND) |
| 41 | try: |
| 42 | definition.input_model.model_validate(params.get("arguments") or {}) |
| 43 | except ValidationError: |
| 44 | return self._failure(context, tool, McpErrorCode.INVALID_ARGUMENT, definition) |
| 45 | return await call_next(context) |
| 46 | |
| 47 | def _failure(self, context, tool, code, definition=None) -> CallToolResult: |
| 48 | request_id = normalize_request_id(context.request_id) |
| 49 | self.audit.emit( |
| 50 | "MCP_TOOL_INVOKED", |
| 51 | tool=tool, |
| 52 | request_id=request_id, |
| 53 | auth=self.auth, |
| 54 | risk_class=definition.risk_class if definition else None, |
| 55 | kind=definition.kind if definition else None, |
| 56 | ) |
| 57 | self.audit.emit( |
| 58 | "MCP_TOOL_FAILED", |
| 59 | tool=tool, |
| 60 | request_id=request_id, |
| 61 | auth=self.auth, |
| 62 | risk_class=definition.risk_class if definition else None, |
| 63 | kind=definition.kind if definition else None, |
| 64 | code=code.value, |
| 65 | ) |
| 66 | envelope = McpResultEnvelope.failure( |
| 67 | tool=tool, |
| 68 | request_id=request_id, |
| 69 | code=code, |
| 70 | ).wire() |
| 71 | return CallToolResult( |
| 72 | content=[TextContent(type="text", text=json.dumps(envelope, separators=(",", ":")))], |
| 73 | structuredContent=envelope, |
| 74 | isError=True, |
| 75 | ) |