Build a Production-Grade MCP Server in Python with Tooling

import os
from functools import lru_cache
from typing import Dict, Any

from pydantic_settings import BaseSettings, SettingsConfigDict
from fastapi import FastAPI
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

class Settings(BaseSettings):
    model_config = SettingsConfigDict(env_file=".env", extra="ignore")

    APP_NAME: str = "MCP Tool Server"
    APP_VERSION: str = "1.0.0"
    OPENAPI_URL: str = "/openapi.json"
    LOG_LEVEL: str = "INFO"

    # OAuth 2.1 settings placeholders
    OAUTH_TENANT_ID: str = ""
    OAUTH_CLIENT_ID: str = ""
    OAUTH_CLIENT_SECRET: str = ""
    OAUTH_TOKEN_URL: str = ""
    OAUTH_JWKS_URL: str = ""

    # Rate Limiting placeholders
    RATE_LIMIT_PER_MINUTE: int = 100
    RATE_LIMIT_EXPIRATION_SECONDS: int = 60

    # Redis for rate limiting (if used) placeholders
    REDIS_HOST: str = "localhost"
    REDIS_PORT: int = 6379
    REDIS_DB: int = 0

@lru_cache()
def get_settings() -> Settings:
    """Caches settings to avoid re-reading .env for every request."""
    logger.info("Loading application settings...")
    return Settings()

app = FastAPI(
    title=get_settings().APP_NAME,
    version=get_settings().APP_VERSION,
    openapi_url=get_settings().OPENAPI_URL,
    docs_url="/docs",
    redoc_url="/redoc"
)

@app.get("/health")
async def health_check() -> Dict[str, Any]:
    """
    Health check endpoint for the server.
    Returns basic status information.
    """
    logger.debug("Health check requested.")
    return {"status": "ok", "version": get_settings().APP_VERSION, "app_name": get_settings().APP_NAME}

# requirements.txt content:
# fastapi
# uvicorn
# pydantic
# pydantic-settings
# python-jose
# redis
# httpx
# opentelemetry-api
# opentelemetry-sdk
# opentelemetry-exporter-otlp
# opentelemetry-instrumentation-fastapi
# opentelemetry-instrumentation-httpx
# opentelemetry-instrumentation-redis
# opentelemetry-instrumentation-logging
# gunicorn

import os
import logging
from typing import Dict, Any, Optional
from datetime import datetime, timedelta

import httpx
from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.security import HTTPBearer, HTTPAuthorizationCredentials
from jose import jwt, jwk
from jose.exceptions import JWTError

from app.config import get_settings, Settings

logger = logging.getLogger(__name__)

# Custom HTTPBearer scheme for dependency injection
oauth2_scheme = HTTPBearer()

class JwksClient:
    """Handles fetching and caching of JSON Web Key Set (JWKS)."""
    def __init__(self, jwks_url: str):
        self.jwks_url = jwks_url
        self._keys: Dict[str, Any] = {}
        self._last_fetched: Optional[datetime] = None
        self._cache_duration = timedelta(hours=1) # Cache JWKS for 1 hour

    async def _fetch_jwks(self) -> Dict[str, Any]:
        """Fetches JWKS from the configured URL."""
        try:
            async with httpx.AsyncClient(timeout=10) as client:
                response = await client.get(self.jwks_url)
                response.raise_for_status()
                self._keys = {key['kid']: key for key in response.json().get('keys', [])}
                self._last_fetched = datetime.utcnow()
                logger.info(f"Successfully fetched JWKS from {self.jwks_url}")
                return self._keys
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error fetching JWKS from {self.jwks_url}: {e.response.status_code} - {e.response.text}")
            raise HTTPException(status_code=status.HTTP_502_BAD_GATEWAY, detail="Failed to fetch JWKS")
        except httpx.RequestError as e:
            logger.error(f"Network error fetching JWKS from {self.jwks_url}: {e}")
            raise HTTPException(status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Network error fetching JWKS")
        except Exception as e:
            logger.error(f"Unexpected error fetching JWKS from {self.jwks_url}: {e}")
            raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail="Internal server error fetching JWKS")

    async def get_key(self, kid: str) -> Optional[Dict[str, Any]]:
        """Retrieves a specific key by KID, fetching JWKS if needed or expired."""
        if not self._keys or (self._last_fetched and datetime.utcnow() - self._last_fetched > self._cache_duration):
            await self._fetch_jwks()
        return self._keys.get(kid)

jwks_client: Optional[JwksClient] = None

async def initialize_jwks_client():
    """Initializes JWKS client on application startup."""
    settings = get_settings()
    global jwks_client
    jwks_client = JwksClient(settings.OAUTH_JWKS_URL)
    await jwks_client._fetch_jwks() # Pre-fetch JWKS on startup

async def get_current_user(credentials: HTTPAuthorizationCredentials = Depends(oauth2_scheme), settings: Settings = Depends(get_settings)) -> Dict[str, Any]:
    """
    Dependency to validate the OAuth 2.1 bearer token.
    """
    if not jwks_client:
        logger.error("JWKS client not initialized.")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Authentication service not ready."
        )

    credentials_exception = HTTPException(
        status_code=status.HTTP_401_UNAUTHORIZED,
        detail="Could not validate credentials",
        headers={"WWW-Authenticate": "Bearer"},
    )
    try:
        token = credentials.credentials  # Extract token string from HTTPAuthorizationCredentials
        # Get the kid from the token header
        header = jwt.get_unverified_header(token)
        kid = header.get("kid")
        if not kid:
            logger.warning("Token header missing 'kid'.")
            raise credentials_exception

        # Retrieve the public key using the kid
        public_key_data = await jwks_client.get_key(kid)
        if not public_key_data:
            logger.warning(f"Public key with kid '{kid}' not found in JWKS.")
            raise credentials_exception

        public_key = jwk.construct(public_key_data)

        # Decode and verify the token
        payload = jwt.decode(
            token,
            public_key,
            algorithms=["RS256"], # Assuming RS256, adjust if your provider uses different
            audience=settings.OAUTH_CLIENT_ID, # The client ID of *this* server
            options={"verify_signature": True, "verify_aud": True, "verify_exp": True}
        )

        # Validate issuer and other claims if necessary
        # A more robust check might involve comparing against a list of allowed issuers
        iss = payload.get("iss", "")
        if not iss or settings.OAUTH_TENANT_ID not in iss:
            logger.warning(f"Invalid issuer: {payload.get('iss')}")
            raise credentials_exception

        logger.info(f"Token validated for client: {payload.get('aud')}")
        return payload # Return the decoded payload for further authorization checks
    except JWTError as e:
        logger.warning(f"JWT validation failed: {e}")
        raise credentials_exception
    except Exception as e:
        logger.exception(f"An unexpected error occurred during token validation: {e}")
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail="Internal authentication error."
        )

import logging
from typing import Dict, Any, Callable, List, Optional
from pydantic import BaseModel, Field
from typing import Annotated
from fastapi import Depends, HTTPException, status, APIRouter
from app.config import get_settings
from app.auth import get_current_user

logger = logging.getLogger(__name__)

class ToolMetadata(BaseModel):
    name: str = Field(..., description="Unique name of the tool.")
    description: str = Field(..., description="Description of the tool's functionality.")
    version: str = Field("1.0.0", description="Version of the tool.")
    input_schema: Dict[str, Any] = Field({}, description="JSON schema for tool input parameters.")
    output_schema: Dict[str, Any] = Field({}, description="JSON schema for tool output.")
    is_resource: bool = Field(False, description="True if this is a resource, False if it's an action tool.")
    rate_limit_per_minute: Optional[int] = Field(None, description="Rate limit for this specific tool, overrides global.")

class Tool:
    """Represents an executable tool or resource."""
    def __init__(self, metadata: ToolMetadata, func: Callable[..., Any]):
        self.metadata = metadata
        self.func = func

    async def execute(self, **kwargs: Any) -> Any:
        """Executes the tool function."""
        logger.info(f"Executing tool: {self.metadata.name} with args: {kwargs}")
        try:
            result = await self.func(**kwargs)
            logger.info(f"Tool {self.metadata.name} executed successfully.")
            return result
        except Exception as e:
            logger.error(f"Error executing tool {self.metadata.name}: {e}", exc_info=True)
            raise

class ToolRegistry:
    """Manages a collection of registered tools."""
    def __init__(self):
        self._tools: Dict[str, Tool] = {}
        logger.info("ToolRegistry initialized.")

    def register_tool(self, tool: Tool):
        """Registers a new tool."""
        if tool.metadata.name in self._tools:
            logger.warning(f"Tool '{tool.metadata.name}' already registered. Overwriting.")
        self._tools[tool.metadata.name] = tool
        logger.info(f"Tool '{tool.metadata.name}' (v{tool.metadata.version}) registered.")

    def get_tool(self, name: str) -> Optional[Tool]:
        """Retrieves a tool by its name."""
        return self._tools.get(name)

    def get_all_tools_metadata(self) -> List[ToolMetadata]:
        """Returns metadata for all registered tools."""
        return [tool.metadata for tool in self._tools.values()]

# Initialize the global tool registry
tool_registry = ToolRegistry()

# --- Example Tool Definitions ---

async def product_lookup_func(product_id: str, region: str = "us-east-1") -> Dict[str, Any]:
    """Simulates looking up product details by ID."""
    logger.info(f"Simulating product lookup for ID: {product_id} in region: {region}")
    # In a real scenario, this would call an external API or database
    if product_id == "P123":
        return {"product_id": product_id, "name": "Wireless Headphones", "price": 99.99, "region": region}
    elif product_id == "P456":
        return {"product_id": product_id, "name": "Smartwatch", "price": 199.99, "region": region}
    else:
        logger.warning(f"Product ID {product_id} not found.")
        return {"product_id": product_id, "name": "Not Found", "price": 0.0, "region": region}

product_lookup_tool = Tool(
    metadata=ToolMetadata(
        name="product_lookup",
        description="Looks up product details by a given product ID and optional region.",
        version="1.1.0",
        input_schema={
            "type": "object",
            "properties": {
                "product_id": {"type": "string", "description": "The unique identifier of the product.", "pattern": "^[P][0-9]{3}$"},
                "region": {"type": "string", "description": "The geographic region to search in (default: us-east-1).", "enum": ["us-east-1", "us-west-2", "eu-central-1"]}
            },
            "required": ["product_id"]
        },
        output_schema={
            "type": "object",
            "properties": {
                "product_id": {"type": "string"},
                "name": {"type": "string"},
                "price": {"type": "number"},
                "region": {"type": "string"}
            }
        },
        rate_limit_per_minute=50
    ),
    func=product_lookup_func
)
tool_registry.register_tool(product_lookup_tool)

async def get_customer_profile_func(customer_id: Annotated[int, Field(ge=1)], include_orders: bool = False) -> Dict[str, Any]:
    """Simulates fetching a customer's profile details."""
    logger.info(f"Simulating fetching customer profile for ID: {customer_id}, include_orders: {include_orders}")
    if customer_id == 101:
        profile = {"customer_id": 101, "name": "Alice Smith", "email": "[email protected]"}
        if include_orders:
            profile["orders"] = [{"order_id": "ORD001", "amount": 120.50}, {"order_id": "ORD002", "amount": 55.00}]
        return profile
    else:
        logger.warning(f"Customer ID {customer_id} not found.")
        return {"customer_id": customer_id, "name": "Not Found"}

customer_profile_resource = Tool(
    metadata=ToolMetadata(
        name="customer_profile_fetcher",
        description="Retrieves a customer's profile by ID. Can optionally include recent orders.",
        version="1.0.0",
        input_schema={
            "type": "object",
            "properties": {
                "customer_id": {"type": "integer", "description": "The unique identifier of the customer.", "minimum": 1},
                "include_orders": {"type": "boolean", "description": "Whether to include recent order history (default: false)."}
            },
            "required": ["customer_id"]
        },
        output_schema={
            "type": "object",
            "properties": {
                "customer_id": {"type": "integer"},
                "name": {"type": "string"},
                "email": {"type": "string"},
                "orders": {"type": "array", "items": {"type": "object"}, "description": "List of recent orders"}
            }
        },
        is_resource=True,
        rate_limit_per_minute=30
    ),
    func=get_customer_profile_func
)
tool_registry.register_tool(customer_profile_resource)

mcp_router = APIRouter(prefix="/.well-known/mcp")

@mcp_router.get("/tools", response_model=List[ToolMetadata])
async def get_mcp_tools(current_user: Dict[str, Any] = Depends(get_current_user)) -> List[ToolMetadata]:
    """
    MCP endpoint to discover available tools and resources.
    Requires authentication.
    """
    logger.info(f"MCP tool discovery requested by client: {current_user.get('aud')}")
    return tool_registry.get_all_tools_metadata()

@mcp_router.post("/invoke/{tool_name}", response_model=Dict[str, Any])
async def invoke_mcp_tool(
    tool_name: str,
    payload: Dict[str, Any],
) -> Dict[str, Any]:
    """
    MCP endpoint to invoke a specific tool.
    Requires authentication.
    """
    tool = tool_registry.get_tool(tool_name)
    if not tool:
        raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Tool '{tool_name}' not found.")

    # Basic input validation against schema (can be more robust with jsonschema library)
    # For this guide, we rely on Pydantic's internal validation where possible and tool's func errors.
    try:
        # Convert payload to Pydantic model for validation if a model was defined, or rely on func's type hints
        # For simplicity, direct pass to function and let it raise errors based on type hints
        result = await tool.execute(**payload)
        return result
    except TypeError as e:
        logger.error(f"Type error during tool invocation for '{tool_name}': {e}", exc_info=True)
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=f"Invalid input parameters for tool '{tool_name}'. Details: {e}")
    except HTTPException:
        raise # Re-raise HTTP exceptions from tool execution
    except Exception as e:
        logger.error(f"Error invoking tool '{tool_name}': {e}", exc_info=True)
        raise HTTPException(status_code=status.HTTP_500_INTERNAL_SERVER_ERROR, detail=f"Error executing tool '{tool_name}'.")

import logging
import time
from typing import Dict, Any, Optional
from datetime import datetime, timedelta

import redis.asyncio as redis
from fastapi import FastAPI, Request, Response, Depends, HTTPException, status
from fastapi.responses import JSONResponse
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.types import ASGIApp

from app.config import get_settings, Settings
from app.tools import tool_registry # Ensure tool registry is imported

logger = logging.getLogger(__name__)

# Initialize Redis client globally
redis_client: Optional[redis.Redis] = None

async def initialize_redis_client():
    """Initializes Redis client on application startup."""
    settings = get_settings()
    global redis_client
    try:
        redis_client = redis.Redis(host=settings.REDIS_HOST, port=settings.REDIS_PORT, db=settings.REDIS_DB, decode_responses=True)
        await redis_client.ping() # Test connection
        logger.info(f"Connected to Redis at {settings.REDIS_HOST}:{settings.REDIS_PORT}")
    except Exception as e:
        logger.critical(f"Failed to connect to Redis at {settings.REDIS_HOST}:{settings.REDIS_PORT}. Rate limiting will be disabled. Error: {e}")
        redis_client = None # Disable rate limiting if Redis connection fails

class RateLimitMiddleware(BaseHTTPMiddleware):
    """Middleware for global and tool-specific rate limiting."""
    def __init__(self, app: ASGIApp, settings: Settings):
        super().__init__(app)
        self.settings = settings
        logger.info("RateLimitMiddleware initialized.")

    async def dispatch(self, request: Request, call_next):
        if not redis_client:
            logger.warning("Redis client not available. Skipping rate limiting.")
            return await call_next(request)

        # Identify client for rate limiting (e.g., from validated token)
        # We'll try to get the 'user' payload from request.state if auth middleware ran before this
        client_id = "anonymous"
        current_user = getattr(request.state, 'user', None)
        if current_user:
            client_id = current_user.get("aud", "unknown_client")

        # Determine rate limit key and limits
        path = request.url.path
        rate_limit_key = f"rate_limit:{client_id}:global"
        limit_per_minute = self.settings.RATE_LIMIT_PER_MINUTE
        expiration_seconds = self.settings.RATE_LIMIT_EXPIRATION_SECONDS

        # Tool-specific rate limiting for invoke endpoint
        if path.startswith("/.well-known/mcp/invoke/"): # Check for MCP invoke path
            try:
                tool_name = path.split("/")[-1]
                tool = tool_registry.get_tool(tool_name)
                if tool and tool.metadata.rate_limit_per_minute is not None:
                    limit_per_minute = tool.metadata.rate_limit_per_minute
                    rate_limit_key = f"rate_limit:{client_id}:tool:{tool_name}"
                    logger.debug(f"Applying tool-specific rate limit for {tool_name}: {limit_per_minute}/min")
                else:
                    logger.debug(f"No specific rate limit for tool {tool_name}. Using global.")
            except IndexError:
                logger.warning(f"Malformed tool invocation path: {path}. Using global rate limit.")

        # Increment counter and check
        current_requests = await redis_client.incr(rate_limit_key)
        if current_requests == 1:
            await redis_client.expire(rate_limit_key, expiration_seconds)
        
        # Retrieve TTL to calculate remaining time
        ttl = await redis_client.ttl(rate_limit_key)
        if ttl == -1: # Key exists but no expiration set (shouldn't happen with expire above, but for safety)
            await redis_client.expire(rate_limit_key, expiration_seconds)
            ttl = expiration_seconds
        elif ttl == -2: # Key does not exist (e.g., expired between incr and ttl call, or race condition)
            ttl = expiration_seconds # Assume full duration remaining if key somehow disappeared

        if current_requests > limit_per_minute:
            logger.warning(f"Rate limit exceeded for client {client_id} on key {rate_limit_key}. Current: {current_requests}, Limit: {limit_per_minute}")
            raise HTTPException(
                status_code=status.HTTP_429_TOO_MANY_REQUESTS,
                detail=f"Rate limit exceeded. Try again in {ttl} seconds.",
                headers={
                    "Retry-After": str(ttl),
                    "X-RateLimit-Limit": str(limit_per_minute),
                    "X-RateLimit-Remaining": "0",
                    "X-RateLimit-Reset": str(int(time.time()) + ttl)
                }
            )

        response = await call_next(request)
        response.headers["X-RateLimit-Limit"] = str(limit_per_minute)
        response.headers["X-RateLimit-Remaining"] = str(max(0, limit_per_minute - current_requests))
        response.headers["X-RateLimit-Reset"] = str(int(time.time()) + ttl) # Approximate reset time
        return response

# Global exception handler for HTTPException for consistent JSON responses
async def http_exception_handler(request: Request, exc: HTTPException):
    logger.error(f"HTTP Exception caught: {exc.status_code} - {exc.detail} for path {request.url.path}")
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "detail": exc.detail,
            "error_type": exc.detail.split('.')[0] if exc.detail else "HTTP_ERROR"
        },
        headers=exc.headers
    )

# Catch-all exception handler for unexpected server errors
async def catch_all_exception_handler(request: Request, exc: Exception):
    logger.exception(f"Unhandled server error for path {request.url.path}: {exc}")
    return JSONResponse(
        status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
        content={"detail": "An unexpected server error occurred.", "error_type": "INTERNAL_SERVER_ERROR"},
    )

import os
import logging
from typing import Dict, Any

from opentelemetry import trace
from opentelemetry.sdk.resources import Resource
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import ConsoleSpanExporter, BatchSpanProcessor
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor
from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor
from opentelemetry.instrumentation.redis import RedisInstrumentor
from opentelemetry.instrumentation.logging import LoggingInstrumentor

from app.config import get_settings
from fastapi import FastAPI # Import FastAPI for instrumentation

logger = logging.getLogger(__name__)

def setup_opentelemetry(app_instance: FastAPI, settings: Settings):
    """
    Configures OpenTelemetry for the FastAPI application.
    """
    resource = Resource.create({
        "service.name": settings.APP_NAME,
        "service.version": settings.APP_VERSION,
        "deployment.environment": os.environ.get("ENV", "development")
    })

    provider = TracerProvider(resource=resource)

    # Configure exporter
    # For local development, use ConsoleSpanExporter or a local Jaeger/Zipkin collector
    # For production, use OTLPSpanExporter to send to an OTLP collector (e.g., Azure Monitor, Datadog, Honeycomb)
    otlp_endpoint = os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT")
    if otlp_endpoint:
        exporter = OTLPSpanExporter(
            endpoint=otlp_endpoint,
            headers={"Content-Type": "application/json"} # Adjust if your collector expects protobuf
        )
        logger.info(f"Using OTLP exporter for OpenTelemetry to {otlp_endpoint}")
    else:
        exporter = ConsoleSpanExporter()
        logger.info("Using Console exporter for OpenTelemetry (OTEL_EXPORTER_OTLP_ENDPOINT not set).")

    span_processor = BatchSpanProcessor(exporter)
    provider.add_span_processor(span_processor)
    trace.set_tracer_provider(provider)

    # Instrument FastAPI application
    FastAPIInstrumentor.instrument_app(app_instance, tracer_provider=provider)
    logger.info("FastAPI instrumented for OpenTelemetry.")

    # Instrument HTTPX for outgoing requests (e.g., JWKS fetching)
    HTTPXClientInstrumentor().instrument()
    logger.info("HTTPX client instrumented for OpenTelemetry.")

    # Instrument Redis client
    RedisInstrumentor().instrument()
    logger.info("Redis client instrumented for OpenTelemetry.")

    # Instrument logging to inject trace_id and span_id
    LoggingInstrumentor().instrument(set_logging_format=True) 
    logger.info("Logging instrumented for OpenTelemetry.")

# Dockerfile
FROM python:3.9-slim-buster

# Set environment variables for Python
ENV PYTHONUNBUFFERED 1
ENV APP_HOME /app

# Create app directory and set as working directory
WORKDIR $APP_HOME

# Install system dependencies (if any, e.g., for specific Python packages like cryptography)
# Uncomment and add if needed: 
# RUN apt-get update && apt-get install -y --no-install-recommends 
#     build-essential 
#     # libffi-dev 
#     # libssl-dev 
#     && rm -rf /var/lib/apt/lists/*

# Copy requirements file and install Python dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy application code
COPY app/ $APP_HOME/app/

# Expose the port FastAPI runs on
EXPOSE 8000

# Command to run the application using Gunicorn with Uvicorn workers for production
CMD ["gunicorn", "app.main:app", "--workers", "4", "--worker-class", "uvicorn.workers.UvicornWorker", "--bind", "0.0.0.0:8000", "--log-level", "info"]