Production Patterns for MCP Servers

"""
RobustToolExecutor: maps exceptions to MCP error types,
retries transient errors, and implements a circuit breaker.
"""
import asyncio
import time
import logging
from enum import Enum
from dataclasses import dataclass, field
from typing import Callable, Any
from mcp.server.fastmcp import FastMCP
from mcp.types import (
    ToolExecutionError,
    InvalidParamsError,
    InternalError,
)

logger = logging.getLogger(__name__)

# ── Circuit Breaker ──────────────────────────────────────────
class CBState(Enum):
    CLOSED = "closed"
    OPEN = "open"
    HALF_OPEN = "half_open"

@dataclass
class CircuitBreaker:
    name: str
    failure_threshold: int = 5
    reset_timeout: float = 60.0       # seconds before trying half-open
    _state: CBState = field(default=CBState.CLOSED, init=False)
    _failures: int = field(default=0, init=False)
    _opened_at: float = field(default=0.0, init=False)

    def record_success(self) -> None:
        self._failures = 0
        self._state = CBState.CLOSED

    def record_failure(self) -> None:
        self._failures += 1
        if self._failures >= self.failure_threshold:
            self._state = CBState.OPEN
            self._opened_at = time.monotonic()
            logger.warning(
                "circuit_breaker_opened",
                extra={"cb_name": self.name, "failures": self._failures}
            )

    def allow_request(self) -> bool:
        if self._state == CBState.CLOSED:
            return True
        if self._state == CBState.OPEN:
            elapsed = time.monotonic() - self._opened_at
            if elapsed >= self.reset_timeout:
                self._state = CBState.HALF_OPEN
                logger.info("circuit_breaker_half_open", extra={"cb_name": self.name})
                return True
            return False
        # HALF_OPEN: allow exactly one probe
        return True

# ── Retry helper ─────────────────────────────────────────────
TRANSIENT_EXCEPTIONS = (ConnectionError, TimeoutError, asyncio.TimeoutError)

async def retry_with_backoff(
    fn: Callable,
    args: tuple,
    kwargs: dict,
    max_retries: int = 3,
    base_delay: float = 0.5,
) -> Any:
    last_exc: Exception | None = None
    for attempt in range(max_retries):
        try:
            return await fn(*args, **kwargs)
        except TRANSIENT_EXCEPTIONS as exc:
            last_exc = exc
            delay = base_delay * (2 ** attempt)
            logger.warning(
                "tool_retry",
                extra={"attempt": attempt + 1, "delay": delay, "error": str(exc)}
            )
            await asyncio.sleep(delay)
    raise last_exc  # all retries exhausted

# ── Executor ─────────────────────────────────────────────────
class RobustToolExecutor:
    def __init__(self, cb: CircuitBreaker | None = None):
        self.cb = cb or CircuitBreaker(name="default")

    async def execute(
        self,
        tool_name: str,
        fn: Callable,
        args: tuple = (),
        kwargs: dict | None = None,
    ) -> Any:
        kwargs = kwargs or {}
        if not self.cb.allow_request():
            raise ToolExecutionError(
                message=f"Circuit breaker OPEN for '{self.cb.name}'. "
                        f"Retry after {self.cb.reset_timeout:.0f}s.",
                tool_name=tool_name,
            )
        try:
            result = await retry_with_backoff(fn, args, kwargs)
            self.cb.record_success()
            return result
        except InvalidParamsError:
            raise  # deterministic, do not retry or count as failure
        except TRANSIENT_EXCEPTIONS as exc:
            self.cb.record_failure()
            raise ToolExecutionError(
                message=f"Transient error after retries: {exc}",
                tool_name=tool_name,
            ) from exc
        except Exception as exc:
            self.cb.record_failure()
            logger.exception("tool_internal_error", extra={"tool": tool_name})
            raise InternalError(message=f"Unexpected error in {tool_name}: {exc}") from exc

# ── Usage with FastMCP ────────────────────────────────────────
mcp = FastMCP("robust-server")
executor = RobustToolExecutor(CircuitBreaker("database", failure_threshold=5, reset_timeout=60))

@mcp.tool()
async def query_database(sql: str) -> str:
    """Run a read-only SQL query."""
    if not sql.strip().upper().startswith("SELECT"):
        raise InvalidParamsError("Only SELECT statements are allowed.")

    async def _run():
        # Replace with real DB call
        await asyncio.sleep(0.01)
        return f"Results for: {sql}"

    return await executor.execute("query_database", _run)

"""
Redis-backed token bucket rate limiter for MCP servers.
Requires: pip install redis[hiredis]
"""
import time
import redis.asyncio as aioredis
from dataclasses import dataclass
from mcp.types import ToolExecutionError

# ── Tier configuration ────────────────────────────────────────
TIERS = {
    "standard": {"capacity": 60,  "refill_rate": 1.0},   # 60/min
    "premium":  {"capacity": 600, "refill_rate": 10.0},  # 600/min
}

# ── Atomic Lua script ─────────────────────────────────────────
# Returns [allowed (0/1), tokens_remaining, retry_after_ms]
_LUA = """
local key       = KEYS[1]
local capacity  = tonumber(ARGV[1])
local rate      = tonumber(ARGV[2])   -- tokens per second
local now       = tonumber(ARGV[3])   -- current time (ms)
local ttl       = tonumber(ARGV[4])   -- key TTL (seconds)

local bucket = redis.call("HMGET", key, "tokens", "last_refill")
local tokens     = tonumber(bucket[1]) or capacity
local last_refill = tonumber(bucket[2]) or now

-- Refill tokens based on elapsed time
local elapsed = math.max(0, (now - last_refill) / 1000.0)
tokens = math.min(capacity, tokens + elapsed * rate)

if tokens >= 1 then
    tokens = tokens - 1
    redis.call("HMSET", key, "tokens", tokens, "last_refill", now)
    redis.call("EXPIRE", key, ttl)
    return {1, math.floor(tokens), 0}
else
    -- Calculate when the next token will be available
    local wait_ms = math.ceil((1 - tokens) / rate * 1000)
    redis.call("HMSET", key, "tokens", tokens, "last_refill", now)
    redis.call("EXPIRE", key, ttl)
    return {0, 0, wait_ms}
end
"""

@dataclass
class RateLimitResult:
    allowed: bool
    tokens_remaining: int
    retry_after_ms: int

class RedisRateLimiter:
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self._redis = aioredis.from_url(redis_url, decode_responses=True)
        self._script = self._redis.register_script(_LUA)

    async def check(self, client_id: str, tier: str = "standard") -> RateLimitResult:
        cfg = TIERS.get(tier, TIERS["standard"])
        now_ms = int(time.time() * 1000)
        key = f"ratelimit:{client_id}"

        result = await self._script(
            keys=[key],
            args=[cfg["capacity"], cfg["refill_rate"], now_ms, 120],
        )
        return RateLimitResult(
            allowed=bool(result[0]),
            tokens_remaining=int(result[1]),
            retry_after_ms=int(result[2]),
        )

# ── Middleware wrapper for FastMCP ────────────────────────────
class RateLimitedServer:
    def __init__(self, mcp_app, limiter: RedisRateLimiter):
        self.app = mcp_app
        self.limiter = limiter

    async def handle_tool_call(self, client_id: str, tier: str, tool_fn, *args, **kwargs):
        result = await self.limiter.check(client_id, tier)
        if not result.allowed:
            retry_after_s = result.retry_after_ms / 1000
            raise ToolExecutionError(
                message=(
                    f"Rate limit exceeded. "
                    f"X-RateLimit-Remaining: 0. "
                    f"Retry-After: {retry_after_s:.1f}s"
                ),
                tool_name=tool_fn.__name__,
            )
        return await tool_fn(*args, **kwargs)

"""
Health check HTTP server for MCP servers.
Runs on port 8090 alongside the MCP transport.
Requires: pip install aiohttp
"""
import asyncio
import time
from aiohttp import web
import redis.asyncio as aioredis

# ── Metrics store (replace with prometheus_client in production) ─
_metrics = {
    "tool_calls_total": 0,
    "tool_errors_total": 0,
    "tool_latency_sum_ms": 0.0,
}

def record_call(duration_ms: float, error: bool = False):
    _metrics["tool_calls_total"] += 1
    _metrics["tool_latency_sum_ms"] += duration_ms
    if error:
        _metrics["tool_errors_total"] += 1

# ── Dependency checkers ──────────────────────────────────────
async def check_redis(redis_url: str) -> tuple[bool, str]:
    try:
        r = aioredis.from_url(redis_url)
        await r.ping()
        await r.aclose()
        return True, "ok"
    except Exception as exc:
        return False, str(exc)

# ── Handlers ─────────────────────────────────────────────────
async def handle_health(request: web.Request) -> web.Response:
    """Liveness: is the process alive?"""
    return web.json_response({"status": "ok", "ts": time.time()})

async def handle_ready(request: web.Request) -> web.Response:
    """Readiness: can we serve requests right now?"""
    redis_ok, redis_msg = await check_redis(
        request.app["redis_url"]
    )
    checks = {
        "redis": {"ok": redis_ok, "detail": redis_msg},
    }
    all_ok = all(c["ok"] for c in checks.values())
    status = 200 if all_ok else 503
    return web.json_response(
        {"status": "ready" if all_ok else "degraded", "checks": checks},
        status=status,
    )

async def handle_metrics(request: web.Request) -> web.Response:
    """Prometheus text format."""
    m = _metrics
    body = "\n".join([
        "# HELP mcp_tool_calls_total Total tool calls",
        "# TYPE mcp_tool_calls_total counter",
        f'mcp_tool_calls_total {m["tool_calls_total"]}',
        "# HELP mcp_tool_errors_total Total tool errors",
        "# TYPE mcp_tool_errors_total counter",
        f'mcp_tool_errors_total {m["tool_errors_total"]}',
        "# HELP mcp_tool_latency_ms_sum Cumulative latency ms",
        "# TYPE mcp_tool_latency_ms_sum counter",
        f'mcp_tool_latency_ms_sum {m["tool_latency_sum_ms"]}',
    ])
    return web.Response(text=body, content_type="text/plain")

# ── App factory ───────────────────────────────────────────────
def create_health_app(redis_url: str = "redis://localhost:6379") -> web.Application:
    app = web.Application()
    app["redis_url"] = redis_url
    app.router.add_get("/health", handle_health)
    app.router.add_get("/ready", handle_ready)
    app.router.add_get("/metrics", handle_metrics)
    return app

async def run_health_server(port: int = 8090):
    app = create_health_app()
    runner = web.AppRunner(app)
    await runner.setup()
    site = web.TCPSite(runner, "0.0.0.0", port)
    await site.start()
    print(f"Health server running on :{port}")

"""
Graceful shutdown for asyncio-based MCP servers.
Handles SIGTERM from Kubernetes with a 30s drain window.
"""
import asyncio
import signal
import logging

logger = logging.getLogger(__name__)

class ShutdownManager:
    def __init__(self, drain_timeout: float = 30.0):
        self.drain_timeout = drain_timeout
        self._shutting_down = False
        self._active_requests = 0
        self._drain_event = asyncio.Event()

    @property
    def is_shutting_down(self) -> bool:
        return self._shutting_down

    def enter_request(self) -> bool:
        """Call at the start of each tool handler. Returns False if shutting down."""
        if self._shutting_down:
            return False
        self._active_requests += 1
        return True

    def exit_request(self) -> None:
        """Call at the end of each tool handler (use try/finally)."""
        self._active_requests -= 1
        if self._shutting_down and self._active_requests == 0:
            self._drain_event.set()

    async def _drain(self) -> None:
        """Wait for active requests to finish."""
        if self._active_requests == 0:
            return
        logger.info("graceful_shutdown_draining", extra={
            "active_requests": self._active_requests,
            "timeout": self.drain_timeout,
        })
        try:
            await asyncio.wait_for(self._drain_event.wait(), timeout=self.drain_timeout)
            logger.info("graceful_shutdown_drained")
        except asyncio.TimeoutError:
            logger.warning("graceful_shutdown_timeout", extra={
                "remaining_requests": self._active_requests,
            })

    async def shutdown(self, cleanup_fn=None) -> None:
        """Perform the full shutdown sequence."""
        logger.info("graceful_shutdown_started")
        self._shutting_down = True

        await self._drain()

        if cleanup_fn:
            try:
                await cleanup_fn()
            except Exception:
                logger.exception("graceful_shutdown_cleanup_error")

        logger.info("graceful_shutdown_complete")

# ── Wire into the event loop ──────────────────────────────────
def install_sigterm_handler(
    loop: asyncio.AbstractEventLoop,
    shutdown_mgr: ShutdownManager,
    cleanup_fn=None,
) -> None:
    async def _handle():
        await shutdown_mgr.shutdown(cleanup_fn)
        loop.stop()

    loop.add_signal_handler(
        signal.SIGTERM,
        lambda: asyncio.ensure_future(_handle()),
    )
    # Also handle SIGINT (Ctrl+C in dev)
    loop.add_signal_handler(
        signal.SIGINT,
        lambda: asyncio.ensure_future(_handle()),
    )

# ── Usage in a FastMCP server ────────────────────────────────
async def main():
    from mcp.server.fastmcp import FastMCP
    mcp = FastMCP("my-server")
    shutdown = ShutdownManager(drain_timeout=30)

    @mcp.tool()
    async def slow_tool(n: int) -> str:
        if not shutdown.enter_request():
            raise RuntimeError("Server is shutting down")
        try:
            await asyncio.sleep(n)  # simulate work
            return f"done after {n}s"
        finally:
            shutdown.exit_request()

    loop = asyncio.get_event_loop()
    install_sigterm_handler(loop, shutdown)
    await mcp.run_async()

if __name__ == "__main__":
    asyncio.run(main())

"""
Structured JSON logging for MCP servers using structlog.
Requires: pip install structlog
"""
import time
import uuid
import structlog
import logging

def configure_logging(server_name: str, level: str = "INFO") -> None:
    """Call once at server startup."""
    structlog.configure(
        processors=[
            structlog.contextvars.merge_contextvars,
            structlog.processors.add_log_level,
            structlog.processors.TimeStamper(fmt="iso"),
            structlog.processors.JSONRenderer(),
        ],
        context_class=dict,
        logger_factory=structlog.PrintLoggerFactory(),
        cache_logger_on_first_use=True,
    )
    # Bind server-level context (present in every log line)
    structlog.contextvars.bind_contextvars(server_name=server_name)
    logging.basicConfig(level=getattr(logging, level.upper()))

log = structlog.get_logger()

# ── Tool call logging decorator ───────────────────────────────
def logged_tool(tool_name: str):
    """Decorator that adds structured logging to any async tool function."""
    def decorator(fn):
        async def wrapper(*args, client_id: str = "unknown", **kwargs):
            trace_id = str(uuid.uuid4())[:8]
            start = time.monotonic()
            # Bind per-request context
            structlog.contextvars.bind_contextvars(
                tool_name=tool_name,
                client_id=client_id,
                trace_id=trace_id,
            )
            log.debug("tool_call_start", args_count=len(args))
            try:
                result = await fn(*args, **kwargs)
                duration = (time.monotonic() - start) * 1000
                log.info("tool_call_success", duration_ms=round(duration, 1))
                return result
            except Exception as exc:
                duration = (time.monotonic() - start) * 1000
                log.error(
                    "tool_call_error",
                    duration_ms=round(duration, 1),
                    error=str(exc),
                    error_type=type(exc).__name__,
                )
                raise
            finally:
                structlog.contextvars.clear_contextvars()
        return wrapper
    return decorator

# ── Usage ─────────────────────────────────────────────────────
configure_logging("my-mcp-server")

@logged_tool("search_documents")
async def search_documents(query: str) -> list[str]:
    # ... real implementation
    return [f"result for {query}"]

# Sample output (one line, formatted here for readability):
# {
#   "server_name": "my-mcp-server",
#   "tool_name": "search_documents",
#   "client_id": "cl_abc",
#   "trace_id": "a1b2c3d4",
#   "event": "tool_call_success",
#   "level": "info",
#   "timestamp": "2024-01-15T14:23:01.442Z",
#   "duration_ms": 18.3
# }

# ── Stage 1: Build ──────────────────────────────────────────
FROM python:3.12-slim AS builder
WORKDIR /build

# Install build dependencies
RUN apt-get update && apt-get install -y --no-install-recommends \
    gcc libpq-dev \
    && rm -rf /var/lib/apt/lists/*

COPY requirements.txt .
RUN pip install --prefix=/install --no-cache-dir -r requirements.txt

# ── Stage 2: Runtime ─────────────────────────────────────────
FROM python:3.12-slim AS runtime
WORKDIR /app

# Non-root user
RUN groupadd -r mcp && useradd -r -g mcp -u 1001 mcp

# Copy only runtime dependencies from builder
COPY --from=builder /install /usr/local
COPY --chown=mcp:mcp . .

USER 1001

ENV PYTHONDONTWRITEBYTECODE=1 \
    PYTHONUNBUFFERED=1

EXPOSE 8090

HEALTHCHECK --interval=10s --timeout=3s --retries=3 \
  CMD python -c "import urllib.request; urllib.request.urlopen('http://localhost:8090/health')"

CMD ["python", "-m", "mcp_server"]