MCP in Production: The Practical Guide to Model Context Protocol for AI Agents That Actually Ship

import asyncio
import json
import time
import logging
from dataclasses import dataclass, field
from typing import Any, Callable
 
logger = logging.getLogger("mcp_gateway")
 
 
@dataclass
class ToolManifest:
    """A tool exposed through the gateway, with policy metadata."""
    name: str
    description: str
    input_schema: dict
    server_id: str
    risk_level: str = "low"        # low | medium | high | critical
    requires_approval: bool = False
    max_calls_per_minute: int = 60
    allowed_roles: list[str] = field(default_factory=lambda: ["*"])
 
 
@dataclass
class MCPRequest:
    """Incoming tool call from an AI client."""
    tool_name: str
    arguments: dict[str, Any]
    caller_id: str
    caller_role: str
    session_id: str
    timestamp: float = field(default_factory=time.time)
 
 
@dataclass
class MCPResponse:
    """Outgoing result to the AI client."""
    tool_name: str
    result: Any
    latency_ms: float
    server_id: str
    approved: bool = True
    error: str | None = None
 
 
class PolicyEngine:
    """Evaluates whether a tool call is allowed under the current policy."""
 
    def __init__(self, approval_callback: Callable | None = None):
        self._approval_callback = approval_callback
 
    def evaluate(self, request: MCPRequest, manifest: ToolManifest) -> tuple[bool, str]:
        # 1. Role-based access control
        if "*" not in manifest.allowed_roles and request.caller_role not in manifest.allowed_roles:
            return False, f"Role '{request.caller_role}' not authorized for tool '{manifest.name}'"
 
        # 2. Risk-level gating: critical tools always need human approval
        if manifest.risk_level == "critical":
            if self._approval_callback is None:
                return False, "Critical tool requires human approval but no approval callback is configured"
            approved = self._approval_callback(request, manifest)
            if not approved:
                return False, "Human approval denied"
 
        # 3. High-risk tools with destructive arguments
        if manifest.risk_level == "high":
            destructive_keywords = {"delete", "drop", "truncate", "remove", "destroy"}
            arg_str = json.dumps(request.arguments).lower()
            if any(kw in arg_str for kw in destructive_keywords):
                if self._approval_callback:
                    approved = self._approval_callback(request, manifest)
                    if not approved:
                        return False, "Destructive operation denied by human reviewer"
 
        return True, "OK"
 
 
class RateLimiter:
    """Token-bucket rate limiter per tool per caller."""
 
    def __init__(self):
        self._buckets: dict[str, list[float]] = {}
 
    def check(self, caller_id: str, tool_name: str, max_rpm: int) -> bool:
        key = f"{caller_id}:{tool_name}"
        now = time.time()
        window = self._buckets.setdefault(key, [])
        # Evict old entries outside the 60-second window
        self._buckets[key] = [t for t in window if now - t < 60]
        if len(self._buckets[key]) >= max_rpm:
            return False
        self._buckets[key].append(now)
        return True
 
 
class TraceCollector:
    """Collects per-call traces for observability and audit."""
 
    def __init__(self):
        self._traces: list[dict] = []
 
    def record(self, request: MCPRequest, response: MCPResponse, policy_result: str):
        self._traces.append({
            "timestamp": request.timestamp,
            "session_id": request.session_id,
            "caller_id": request.caller_id,
            "caller_role": request.caller_role,
            "tool": request.tool_name,
            "arguments": request.arguments,
            "server_id": response.server_id,
            "result_preview": str(response.result)[:200],
            "latency_ms": response.latency_ms,
            "approved": response.approved,
            "policy_result": policy_result,
            "error": response.error,
        })
 
    def export(self) -> list[dict]:
        return list(self._traces)

class MCPGateway:
    """Central gateway that mediates between AI clients and MCP servers."""
 
    def __init__(self):
        self.tools: dict[str, ToolManifest] = {}
        self.servers: dict[str, Callable] = {}  # server_id -> async call function
        self.policy = PolicyEngine()
        self.rate_limiter = RateLimiter()
        self.tracer = TraceCollector()
 
    def register_server(
        self,
        server_id: str,
        tools: list[ToolManifest],
        call_fn: Callable,
    ) -> None:
        """Register an MCP server and its tools with the gateway."""
        for tool in tools:
            tool.server_id = server_id
            self.tools[tool.name] = tool
        self.servers[server_id] = call_fn
        logger.info(f"Registered server '{server_id}' with {len(tools)} tools")
 
    def list_tools(self, caller_role: str = "*") -> list[dict]:
        """Return the composed tool manifest, filtered by caller role."""
        visible = []
        for tool in self.tools.values():
            if "*" in tool.allowed_roles or caller_role in tool.allowed_roles:
                visible.append({
                    "name": tool.name,
                    "description": tool.description,
                    "inputSchema": tool.input_schema,
                    "riskLevel": tool.risk_level,
                    "requiresApproval": tool.requires_approval,
                })
        return visible
 
    async def call_tool(self, request: MCPRequest) -> MCPResponse:
        """Route a tool call through policy, rate limiting, and the backend server."""
        t0 = time.perf_counter()
 
        # 1. Resolve the tool
        manifest = self.tools.get(request.tool_name)
        if manifest is None:
            return MCPResponse(
                tool_name=request.tool_name,
                result=None,
                latency_ms=0,
                server_id="unknown",
                approved=False,
                error=f"Tool '{request.tool_name}' not found",
            )
 
        # 2. Policy check
        allowed, reason = self.policy.evaluate(request, manifest)
        if not allowed:
            resp = MCPResponse(
                tool_name=request.tool_name,
                result=None,
                latency_ms=(time.perf_counter() - t0) * 1000,
                server_id=manifest.server_id,
                approved=False,
                error=reason,
            )
            self.tracer.record(request, resp, reason)
            return resp
 
        # 3. Rate limiting
        if not self.rate_limiter.check(
            request.caller_id, request.tool_name, manifest.max_calls_per_minute
        ):
            resp = MCPResponse(
                tool_name=request.tool_name,
                result=None,
                latency_ms=(time.perf_counter() - t0) * 1000,
                server_id=manifest.server_id,
                approved=False,
                error="Rate limit exceeded",
            )
            self.tracer.record(request, resp, "rate_limited")
            return resp
 
        # 4. Dispatch to the backend MCP server
        call_fn = self.servers[manifest.server_id]
        try:
            result = await call_fn(request.tool_name, request.arguments)
            error = None
        except Exception as e:
            result = None
            error = str(e)
            logger.error(f"Tool '{request.tool_name}' failed: {e}")
 
        resp = MCPResponse(
            tool_name=request.tool_name,
            result=result,
            latency_ms=(time.perf_counter() - t0) * 1000,
            server_id=manifest.server_id,
            approved=True,
            error=error,
        )
        self.tracer.record(request, resp, reason)
        return resp

import subprocess
 
# --- Database tool (read-only, low risk) ---
 
async def db_query_handler(tool_name: str, arguments: dict) -> dict:
    """Execute a read-only SQL query against the analytics warehouse."""
    import asyncpg
    conn = await asyncpg.connect("postgresql://readonly:***@db:5432/analytics")
    try:
        rows = await conn.fetch(arguments["query"])
        return {"rows": [dict(r) for r in rows], "count": len(rows)}
    finally:
        await conn.close()
 
 
db_tool = ToolManifest(
    name="query_analytics_db",
    description="Run a read-only SQL query against the analytics warehouse.",
    input_schema={
        "type": "object",
        "properties": {
            "query": {"type": "string", "description": "SQL SELECT query"}
        },
        "required": ["query"],
    },
    server_id="analytics_db",
    risk_level="medium",
    max_calls_per_minute=30,
    allowed_roles=["analyst", "engineer", "admin"],
)
 
 
# --- File system tool (write access, high risk) ---
 
async def file_write_handler(tool_name: str, arguments: dict) -> dict:
    """Write content to a sandboxed directory."""
    import os
 
    base = os.path.abspath("/data/agent-workspace")
    target = os.path.abspath(os.path.join(base, arguments["path"]))
 
    if not target.startswith(base):
        raise PermissionError(f"Path traversal blocked: {arguments['path']}")
 
    os.makedirs(os.path.dirname(target), exist_ok=True)
    with open(target, "w") as f:
        f.write(arguments["content"])
 
    return {"status": "written", "path": arguments["path"]}
 
 
file_tool = ToolManifest(
    name="write_workspace_file",
    description="Write a file to the agent's sandboxed workspace directory.",
    input_schema={
        "type": "object",
        "properties": {
            "path": {"type": "string", "description": "Relative path within workspace"},
            "content": {"type": "string", "description": "File content to write"},
        },
        "required": ["path", "content"],
    },
    server_id="workspace_fs",
    risk_level="high",
    requires_approval=False,
    max_calls_per_minute=10,
    allowed_roles=["engineer", "admin"],
)
 
 
# --- Compose the gateway ---
 
gateway = MCPGateway()
gateway.register_server("analytics_db", [db_tool], db_query_handler)
gateway.register_server("workspace_fs", [file_tool], file_write_handler)
 
# An agent sees only the tools its role permits:
print(gateway.list_tools(caller_role="analyst"))
# [{"name": "query_analytics_db", ...}]  — no file_write_handler visible

import hashlib
 
def verify_manifest(manifest: dict, expected_hash: str) -> bool:
    """Verify that a tool manifest has not been tampered with."""
    canonical = json.dumps(manifest, sort_keys=True, separators=(",", ":"))
    actual_hash = hashlib.sha256(canonical.encode()).hexdigest()
    return actual_hash == expected_hash

import re
 
INJECTION_PATTERNS = [
    r"ignore\s+(all\s+)?previous\s+instructions",
    r"you\s+are\s+now\s+a",
    r"system:\s*",
    r"<\|im_start\|>",
    r"<\|endoftext\|>",
]
 
def sanitize_tool_result(result: str) -> str:
    """Strip potential prompt injection patterns from tool results."""
    for pattern in INJECTION_PATTERNS:
        result = re.sub(pattern, "[REDACTED]", result, flags=re.IGNORECASE)
    return result

class ChainPolicy:
    """Detect and block risky tool call sequences within a session."""
 
    BLOCKED_CHAINS = [
        ("query_analytics_db", "write_workspace_file"),
        ("read_secrets", "send_email"),
    ]
 
    def __init__(self):
        self._session_history: dict[str, list[str]] = {}
 
    def check(self, session_id: str, tool_name: str) -> tuple[bool, str]:
        history = self._session_history.setdefault(session_id, [])
 
        for chain in self.BLOCKED_CHAINS:
            if len(history) > 0 and history[-1] == chain[0] and tool_name == chain[1]:
                return False, f"Blocked chain: {chain[0]} → {chain[1]}"
 
        history.append(tool_name)
        return True, "OK"