fastmcp 2.8.1__py3-none-any.whl → 2.9.0__py3-none-any.whl

fastmcp/server/middleware/middleware.py

@@ -0,0 +1,236 @@
+from __future__ import annotations
+
+import logging
+from collections.abc import Awaitable
+from dataclasses import dataclass, field, replace
+from datetime import datetime, timezone
+from functools import partial
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Generic,
+    Literal,
+    Protocol,
+    TypeVar,
+    runtime_checkable,
+)
+
+import mcp.types as mt
+
+from fastmcp.prompts.prompt import Prompt
+from fastmcp.resources.resource import Resource
+from fastmcp.resources.template import ResourceTemplate
+from fastmcp.tools.tool import Tool
+
+if TYPE_CHECKING:
+    from fastmcp.server.context import Context
+
+logger = logging.getLogger(__name__)
+
+
+T = TypeVar("T")
+R = TypeVar("R", covariant=True)
+
+
+@runtime_checkable
+class CallNext(Protocol[T, R]):
+    def __call__(self, context: MiddlewareContext[T]) -> Awaitable[R]: ...
+
+
+ServerResultT = TypeVar(
+    "ServerResultT",
+    bound=mt.EmptyResult
+    | mt.InitializeResult
+    | mt.CompleteResult
+    | mt.GetPromptResult
+    | mt.ListPromptsResult
+    | mt.ListResourcesResult
+    | mt.ListResourceTemplatesResult
+    | mt.ReadResourceResult
+    | mt.CallToolResult
+    | mt.ListToolsResult,
+)
+
+
+@dataclass(kw_only=True)
+class CallToolResult:
+    content: list[mt.Content]
+    isError: bool = False
+
+
+@dataclass(kw_only=True)
+class ListToolsResult:
+    tools: dict[str, Tool]
+
+
+@dataclass(kw_only=True)
+class ListResourcesResult:
+    resources: list[Resource]
+
+
+@dataclass(kw_only=True)
+class ListResourceTemplatesResult:
+    resource_templates: list[ResourceTemplate]
+
+
+@dataclass(kw_only=True)
+class ListPromptsResult:
+    prompts: list[Prompt]
+
+
+@runtime_checkable
+class ServerResultProtocol(Protocol[ServerResultT]):
+    root: ServerResultT
+
+
+@dataclass(kw_only=True, frozen=True)
+class MiddlewareContext(Generic[T]):
+    """
+    Unified context for all middleware operations.
+    """
+
+    message: T
+
+    fastmcp_context: Context | None = None
+
+    # Common metadata
+    source: Literal["client", "server"] = "client"
+    type: Literal["request", "notification"] = "request"
+    method: str | None = None
+    timestamp: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+    def copy(self, **kwargs: Any) -> MiddlewareContext[T]:
+        return replace(self, **kwargs)
+
+
+def make_middleware_wrapper(
+    middleware: Middleware, call_next: CallNext[T, R]
+) -> CallNext[T, R]:
+    """Create a wrapper that applies a single middleware to a context. The
+    closure bakes in the middleware and call_next function, so it can be
+    passed to other functions that expect a call_next function."""
+
+    async def wrapper(context: MiddlewareContext[T]) -> R:
+        return await middleware(context, call_next)
+
+    return wrapper
+
+
+class Middleware:
+    """Base class for FastMCP middleware with dispatching hooks."""
+
+    async def __call__(
+        self,
+        context: MiddlewareContext[T],
+        call_next: CallNext[T, Any],
+    ) -> Any:
+        """Main entry point that orchestrates the pipeline."""
+        handler_chain = await self._dispatch_handler(
+            context,
+            call_next=call_next,
+        )
+        return await handler_chain(context)
+
+    async def _dispatch_handler(
+        self, context: MiddlewareContext[Any], call_next: CallNext[Any, Any]
+    ) -> CallNext[Any, Any]:
+        """Builds a chain of handlers for a given message."""
+        handler = call_next
+
+        match context.method:
+            case "tools/call":
+                handler = partial(self.on_call_tool, call_next=handler)
+            case "resources/read":
+                handler = partial(self.on_read_resource, call_next=handler)
+            case "prompts/get":
+                handler = partial(self.on_get_prompt, call_next=handler)
+            case "tools/list":
+                handler = partial(self.on_list_tools, call_next=handler)
+            case "resources/list":
+                handler = partial(self.on_list_resources, call_next=handler)
+            case "resources/templates/list":
+                handler = partial(self.on_list_resource_templates, call_next=handler)
+            case "prompts/list":
+                handler = partial(self.on_list_prompts, call_next=handler)
+
+        match context.type:
+            case "request":
+                handler = partial(self.on_request, call_next=handler)
+            case "notification":
+                handler = partial(self.on_notification, call_next=handler)
+
+        handler = partial(self.on_message, call_next=handler)
+
+        return handler
+
+    async def on_message(
+        self,
+        context: MiddlewareContext[Any],
+        call_next: CallNext[Any, Any],
+    ) -> Any:
+        return await call_next(context)
+
+    async def on_request(
+        self,
+        context: MiddlewareContext[mt.Request],
+        call_next: CallNext[mt.Request, Any],
+    ) -> Any:
+        return await call_next(context)
+
+    async def on_notification(
+        self,
+        context: MiddlewareContext[mt.Notification],
+        call_next: CallNext[mt.Notification, Any],
+    ) -> Any:
+        return await call_next(context)
+
+    async def on_call_tool(
+        self,
+        context: MiddlewareContext[mt.CallToolRequestParams],
+        call_next: CallNext[mt.CallToolRequestParams, mt.CallToolResult],
+    ) -> mt.CallToolResult:
+        return await call_next(context)
+
+    async def on_read_resource(
+        self,
+        context: MiddlewareContext[mt.ReadResourceRequestParams],
+        call_next: CallNext[mt.ReadResourceRequestParams, mt.ReadResourceResult],
+    ) -> mt.ReadResourceResult:
+        return await call_next(context)
+
+    async def on_get_prompt(
+        self,
+        context: MiddlewareContext[mt.GetPromptRequestParams],
+        call_next: CallNext[mt.GetPromptRequestParams, mt.GetPromptResult],
+    ) -> mt.GetPromptResult:
+        return await call_next(context)
+
+    async def on_list_tools(
+        self,
+        context: MiddlewareContext[mt.ListToolsRequest],
+        call_next: CallNext[mt.ListToolsRequest, ListToolsResult],
+    ) -> ListToolsResult:
+        return await call_next(context)
+
+    async def on_list_resources(
+        self,
+        context: MiddlewareContext[mt.ListResourcesRequest],
+        call_next: CallNext[mt.ListResourcesRequest, ListResourcesResult],
+    ) -> ListResourcesResult:
+        return await call_next(context)
+
+    async def on_list_resource_templates(
+        self,
+        context: MiddlewareContext[mt.ListResourceTemplatesRequest],
+        call_next: CallNext[
+            mt.ListResourceTemplatesRequest, ListResourceTemplatesResult
+        ],
+    ) -> ListResourceTemplatesResult:
+        return await call_next(context)
+
+    async def on_list_prompts(
+        self,
+        context: MiddlewareContext[mt.ListPromptsRequest],
+        call_next: CallNext[mt.ListPromptsRequest, ListPromptsResult],
+    ) -> ListPromptsResult:
+        return await call_next(context)

fastmcp/server/middleware/rate_limiting.py

@@ -0,0 +1,231 @@
+"""Rate limiting middleware for protecting FastMCP servers from abuse."""
+
+import asyncio
+import time
+from collections import defaultdict, deque
+from collections.abc import Callable
+from typing import Any
+
+from mcp import McpError
+from mcp.types import ErrorData
+
+from .middleware import CallNext, Middleware, MiddlewareContext
+
+
+class RateLimitError(McpError):
+    """Error raised when rate limit is exceeded."""
+
+    def __init__(self, message: str = "Rate limit exceeded"):
+        super().__init__(ErrorData(code=-32000, message=message))
+
+
+class TokenBucketRateLimiter:
+    """Token bucket implementation for rate limiting."""
+
+    def __init__(self, capacity: int, refill_rate: float):
+        """Initialize token bucket.
+
+        Args:
+            capacity: Maximum number of tokens in the bucket
+            refill_rate: Tokens added per second
+        """
+        self.capacity = capacity
+        self.refill_rate = refill_rate
+        self.tokens = capacity
+        self.last_refill = time.time()
+        self._lock = asyncio.Lock()
+
+    async def consume(self, tokens: int = 1) -> bool:
+        """Try to consume tokens from the bucket.
+
+        Args:
+            tokens: Number of tokens to consume
+
+        Returns:
+            True if tokens were available and consumed, False otherwise
+        """
+        async with self._lock:
+            now = time.time()
+            elapsed = now - self.last_refill
+
+            # Add tokens based on elapsed time
+            self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
+            self.last_refill = now
+
+            if self.tokens >= tokens:
+                self.tokens -= tokens
+                return True
+            return False
+
+
+class SlidingWindowRateLimiter:
+    """Sliding window rate limiter implementation."""
+
+    def __init__(self, max_requests: int, window_seconds: int):
+        """Initialize sliding window rate limiter.
+
+        Args:
+            max_requests: Maximum requests allowed in the time window
+            window_seconds: Time window in seconds
+        """
+        self.max_requests = max_requests
+        self.window_seconds = window_seconds
+        self.requests = deque()
+        self._lock = asyncio.Lock()
+
+    async def is_allowed(self) -> bool:
+        """Check if a request is allowed."""
+        async with self._lock:
+            now = time.time()
+            cutoff = now - self.window_seconds
+
+            # Remove old requests outside the window
+            while self.requests and self.requests[0] < cutoff:
+                self.requests.popleft()
+
+            if len(self.requests) < self.max_requests:
+                self.requests.append(now)
+                return True
+            return False
+
+
+class RateLimitingMiddleware(Middleware):
+    """Middleware that implements rate limiting to prevent server abuse.
+
+    Uses a token bucket algorithm by default, allowing for burst traffic
+    while maintaining a sustainable long-term rate.
+
+    Example:
+        ```python
+        from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
+
+        # Allow 10 requests per second with bursts up to 20
+        rate_limiter = RateLimitingMiddleware(
+            max_requests_per_second=10,
+            burst_capacity=20
+        )
+
+        mcp = FastMCP("MyServer")
+        mcp.add_middleware(rate_limiter)
+        ```
+    """
+
+    def __init__(
+        self,
+        max_requests_per_second: float = 10.0,
+        burst_capacity: int | None = None,
+        get_client_id: Callable[[MiddlewareContext], str] | None = None,
+        global_limit: bool = False,
+    ):
+        """Initialize rate limiting middleware.
+
+        Args:
+            max_requests_per_second: Sustained requests per second allowed
+            burst_capacity: Maximum burst capacity. If None, defaults to 2x max_requests_per_second
+            get_client_id: Function to extract client ID from context. If None, uses global limiting
+            global_limit: If True, apply limit globally; if False, per-client
+        """
+        self.max_requests_per_second = max_requests_per_second
+        self.burst_capacity = burst_capacity or int(max_requests_per_second * 2)
+        self.get_client_id = get_client_id
+        self.global_limit = global_limit
+
+        # Storage for rate limiters per client
+        self.limiters: dict[str, TokenBucketRateLimiter] = defaultdict(
+            lambda: TokenBucketRateLimiter(
+                self.burst_capacity, self.max_requests_per_second
+            )
+        )
+
+        # Global rate limiter
+        if self.global_limit:
+            self.global_limiter = TokenBucketRateLimiter(
+                self.burst_capacity, self.max_requests_per_second
+            )
+
+    def _get_client_identifier(self, context: MiddlewareContext) -> str:
+        """Get client identifier for rate limiting."""
+        if self.get_client_id:
+            return self.get_client_id(context)
+        return "global"
+
+    async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
+        """Apply rate limiting to requests."""
+        if self.global_limit:
+            # Global rate limiting
+            allowed = await self.global_limiter.consume()
+            if not allowed:
+                raise RateLimitError("Global rate limit exceeded")
+        else:
+            # Per-client rate limiting
+            client_id = self._get_client_identifier(context)
+            limiter = self.limiters[client_id]
+            allowed = await limiter.consume()
+            if not allowed:
+                raise RateLimitError(f"Rate limit exceeded for client: {client_id}")
+
+        return await call_next(context)
+
+
+class SlidingWindowRateLimitingMiddleware(Middleware):
+    """Middleware that implements sliding window rate limiting.
+
+    Uses a sliding window approach which provides more precise rate limiting
+    but uses more memory to track individual request timestamps.
+
+    Example:
+        ```python
+        from fastmcp.server.middleware.rate_limiting import SlidingWindowRateLimitingMiddleware
+
+        # Allow 100 requests per minute
+        rate_limiter = SlidingWindowRateLimitingMiddleware(
+            max_requests=100,
+            window_minutes=1
+        )
+
+        mcp = FastMCP("MyServer")
+        mcp.add_middleware(rate_limiter)
+        ```
+    """
+
+    def __init__(
+        self,
+        max_requests: int,
+        window_minutes: int = 1,
+        get_client_id: Callable[[MiddlewareContext], str] | None = None,
+    ):
+        """Initialize sliding window rate limiting middleware.
+
+        Args:
+            max_requests: Maximum requests allowed in the time window
+            window_minutes: Time window in minutes
+            get_client_id: Function to extract client ID from context
+        """
+        self.max_requests = max_requests
+        self.window_seconds = window_minutes * 60
+        self.get_client_id = get_client_id
+
+        # Storage for rate limiters per client
+        self.limiters: dict[str, SlidingWindowRateLimiter] = defaultdict(
+            lambda: SlidingWindowRateLimiter(self.max_requests, self.window_seconds)
+        )
+
+    def _get_client_identifier(self, context: MiddlewareContext) -> str:
+        """Get client identifier for rate limiting."""
+        if self.get_client_id:
+            return self.get_client_id(context)
+        return "global"
+
+    async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
+        """Apply sliding window rate limiting to requests."""
+        client_id = self._get_client_identifier(context)
+        limiter = self.limiters[client_id]
+
+        allowed = await limiter.is_allowed()
+        if not allowed:
+            raise RateLimitError(
+                f"Rate limit exceeded: {self.max_requests} requests per "
+                f"{self.window_seconds // 60} minutes for client: {client_id}"
+            )
+
+        return await call_next(context)

fastmcp/server/middleware/timing.py

@@ -0,0 +1,156 @@
+"""Timing middleware for measuring and logging request performance."""
+
+import logging
+import time
+from typing import Any
+
+from .middleware import CallNext, Middleware, MiddlewareContext
+
+
+class TimingMiddleware(Middleware):
+    """Middleware that logs the execution time of requests.
+
+    Only measures and logs timing for request messages (not notifications).
+    Provides insights into performance characteristics of your MCP server.
+
+    Example:
+        ```python
+        from fastmcp.server.middleware.timing import TimingMiddleware
+
+        mcp = FastMCP("MyServer")
+        mcp.add_middleware(TimingMiddleware())
+
+        # Now all requests will be timed and logged
+        ```
+    """
+
+    def __init__(
+        self, logger: logging.Logger | None = None, log_level: int = logging.INFO
+    ):
+        """Initialize timing middleware.
+
+        Args:
+            logger: Logger instance to use. If None, creates a logger named 'fastmcp.timing'
+            log_level: Log level for timing messages (default: INFO)
+        """
+        self.logger = logger or logging.getLogger("fastmcp.timing")
+        self.log_level = log_level
+
+    async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
+        """Time request execution and log the results."""
+        method = context.method or "unknown"
+
+        start_time = time.perf_counter()
+        try:
+            result = await call_next(context)
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            self.logger.log(
+                self.log_level, f"Request {method} completed in {duration_ms:.2f}ms"
+            )
+            return result
+        except Exception as e:
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            self.logger.log(
+                self.log_level,
+                f"Request {method} failed after {duration_ms:.2f}ms: {e}",
+            )
+            raise
+
+
+class DetailedTimingMiddleware(Middleware):
+    """Enhanced timing middleware with per-operation breakdowns.
+
+    Provides detailed timing information for different types of MCP operations,
+    allowing you to identify performance bottlenecks in specific operations.
+
+    Example:
+        ```python
+        from fastmcp.server.middleware.timing import DetailedTimingMiddleware
+        import logging
+
+        # Configure logging to see the output
+        logging.basicConfig(level=logging.INFO)
+
+        mcp = FastMCP("MyServer")
+        mcp.add_middleware(DetailedTimingMiddleware())
+        ```
+    """
+
+    def __init__(
+        self, logger: logging.Logger | None = None, log_level: int = logging.INFO
+    ):
+        """Initialize detailed timing middleware.
+
+        Args:
+            logger: Logger instance to use. If None, creates a logger named 'fastmcp.timing.detailed'
+            log_level: Log level for timing messages (default: INFO)
+        """
+        self.logger = logger or logging.getLogger("fastmcp.timing.detailed")
+        self.log_level = log_level
+
+    async def _time_operation(
+        self, context: MiddlewareContext, call_next: CallNext, operation_name: str
+    ) -> Any:
+        """Helper method to time any operation."""
+        start_time = time.perf_counter()
+        try:
+            result = await call_next(context)
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            self.logger.log(
+                self.log_level, f"{operation_name} completed in {duration_ms:.2f}ms"
+            )
+            return result
+        except Exception as e:
+            duration_ms = (time.perf_counter() - start_time) * 1000
+            self.logger.log(
+                self.log_level,
+                f"{operation_name} failed after {duration_ms:.2f}ms: {e}",
+            )
+            raise
+
+    async def on_call_tool(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time tool execution."""
+        tool_name = getattr(context.message, "name", "unknown")
+        return await self._time_operation(context, call_next, f"Tool '{tool_name}'")
+
+    async def on_read_resource(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time resource reading."""
+        resource_uri = getattr(context.message, "uri", "unknown")
+        return await self._time_operation(
+            context, call_next, f"Resource '{resource_uri}'"
+        )
+
+    async def on_get_prompt(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time prompt retrieval."""
+        prompt_name = getattr(context.message, "name", "unknown")
+        return await self._time_operation(context, call_next, f"Prompt '{prompt_name}'")
+
+    async def on_list_tools(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time tool listing."""
+        return await self._time_operation(context, call_next, "List tools")
+
+    async def on_list_resources(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time resource listing."""
+        return await self._time_operation(context, call_next, "List resources")
+
+    async def on_list_resource_templates(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time resource template listing."""
+        return await self._time_operation(context, call_next, "List resource templates")
+
+    async def on_list_prompts(
+        self, context: MiddlewareContext, call_next: CallNext
+    ) -> Any:
+        """Time prompt listing."""
+        return await self._time_operation(context, call_next, "List prompts")