fastmcp 2.12.1__py3-none-any.whl → 2.13.2__py3-none-any.whl

fastmcp/server/middleware/logging.py

@@ -2,6 +2,7 @@
 
 import json
 import logging
+import time
 from collections.abc import Callable
 from logging import Logger
 from typing import Any
@@ -16,7 +17,135 @@ def default_serializer(data: Any) -> str:
     return pydantic_core.to_json(data, fallback=str).decode()
 
 
-class LoggingMiddleware(Middleware):
+class BaseLoggingMiddleware(Middleware):
+    """Base class for logging middleware."""
+
+    logger: Logger
+    log_level: int
+    include_payloads: bool
+    include_payload_length: bool
+    estimate_payload_tokens: bool
+    max_payload_length: int | None
+    methods: list[str] | None
+    structured_logging: bool
+    payload_serializer: Callable[[Any], str] | None
+
+    def _serialize_payload(self, context: MiddlewareContext[Any]) -> str:
+        payload: str
+
+        if not self.payload_serializer:
+            payload = default_serializer(context.message)
+        else:
+            try:
+                payload = self.payload_serializer(context.message)
+            except Exception as e:
+                self.logger.warning(
+                    f"Failed to serialize payload due to {e}: {context.type} {context.method} {context.source}."
+                )
+                payload = default_serializer(context.message)
+
+        return payload
+
+    def _format_message(self, message: dict[str, str | int | float]) -> str:
+        """Format a message for logging."""
+        if self.structured_logging:
+            return json.dumps(message)
+        else:
+            return " ".join([f"{k}={v}" for k, v in message.items()])
+
+    def _create_before_message(
+        self, context: MiddlewareContext[Any]
+    ) -> dict[str, str | int | float]:
+        message = {
+            "event": context.type + "_start",
+            "method": context.method or "unknown",
+            "source": context.source,
+        }
+
+        if (
+            self.include_payloads
+            or self.include_payload_length
+            or self.estimate_payload_tokens
+        ):
+            payload = self._serialize_payload(context)
+
+            if self.include_payload_length or self.estimate_payload_tokens:
+                payload_length = len(payload)
+                payload_tokens = payload_length // 4
+                if self.estimate_payload_tokens:
+                    message["payload_tokens"] = payload_tokens
+                if self.include_payload_length:
+                    message["payload_length"] = payload_length
+
+            if self.max_payload_length and len(payload) > self.max_payload_length:
+                payload = payload[: self.max_payload_length] + "..."
+
+            if self.include_payloads:
+                message["payload"] = payload
+                message["payload_type"] = type(context.message).__name__
+
+        return message
+
+    def _create_error_message(
+        self,
+        context: MiddlewareContext[Any],
+        start_time: float,
+        error: Exception,
+    ) -> dict[str, str | int | float]:
+        duration_ms: float = _get_duration_ms(start_time)
+        message = {
+            "event": context.type + "_error",
+            "method": context.method or "unknown",
+            "source": context.source,
+            "duration_ms": duration_ms,
+            "error": str(object=error),
+        }
+        return message
+
+    def _create_after_message(
+        self,
+        context: MiddlewareContext[Any],
+        start_time: float,
+    ) -> dict[str, str | int | float]:
+        duration_ms: float = _get_duration_ms(start_time)
+        message = {
+            "event": context.type + "_success",
+            "method": context.method or "unknown",
+            "source": context.source,
+            "duration_ms": duration_ms,
+        }
+        return message
+
+    def _log_message(
+        self, message: dict[str, str | int | float], log_level: int | None = None
+    ):
+        self.logger.log(log_level or self.log_level, self._format_message(message))
+
+    async def on_message(
+        self, context: MiddlewareContext[Any], call_next: CallNext[Any, Any]
+    ) -> Any:
+        """Log messages for configured methods."""
+
+        if self.methods and context.method not in self.methods:
+            return await call_next(context)
+
+        self._log_message(self._create_before_message(context))
+
+        start_time = time.perf_counter()
+        try:
+            result = await call_next(context)
+
+            self._log_message(self._create_after_message(context, start_time))
+
+            return result
+        except Exception as e:
+            self._log_message(
+                self._create_error_message(context, start_time, e), logging.ERROR
+            )
+            raise
+
+
+class LoggingMiddleware(BaseLoggingMiddleware):
     """Middleware that provides comprehensive request and response logging.
 
     Logs all MCP messages with configurable detail levels. Useful for debugging,
@@ -37,9 +166,12 @@ class LoggingMiddleware(Middleware):
 
     def __init__(
         self,
+        *,
         logger: logging.Logger | None = None,
         log_level: int = logging.INFO,
         include_payloads: bool = False,
+        include_payload_length: bool = False,
+        estimate_payload_tokens: bool = False,
         max_payload_length: int = 1000,
         methods: list[str] | None = None,
         payload_serializer: Callable[[Any], str] | None = None,
@@ -50,68 +182,25 @@ class LoggingMiddleware(Middleware):
             logger: Logger instance to use. If None, creates a logger named 'fastmcp.requests'
             log_level: Log level for messages (default: INFO)
             include_payloads: Whether to include message payloads in logs
+            include_payload_length: Whether to include response size in logs
+            estimate_payload_tokens: Whether to estimate response tokens
             max_payload_length: Maximum length of payload to log (prevents huge logs)
             methods: List of methods to log. If None, logs all methods.
+            payload_serializer: Callable that converts objects to a JSON string for the
+                payload. If not provided, uses FastMCP's default tool serializer.
         """
-        self.logger: Logger = logger or logging.getLogger("fastmcp.requests")
-        self.log_level: int = log_level
+        self.logger: Logger = logger or logging.getLogger("fastmcp.middleware.logging")
+        self.log_level = log_level
         self.include_payloads: bool = include_payloads
+        self.include_payload_length: bool = include_payload_length
+        self.estimate_payload_tokens: bool = estimate_payload_tokens
         self.max_payload_length: int = max_payload_length
         self.methods: list[str] | None = methods
         self.payload_serializer: Callable[[Any], str] | None = payload_serializer
+        self.structured_logging: bool = False
 
-    def _format_message(self, context: MiddlewareContext[Any]) -> str:
-        """Format a message for logging."""
-        parts = [
-            f"source={context.source}",
-            f"type={context.type}",
-            f"method={context.method or 'unknown'}",
-        ]
-
-        if self.include_payloads:
-            payload: str
 
-            if not self.payload_serializer:
-                payload = default_serializer(context.message)
-            else:
-                try:
-                    payload = self.payload_serializer(context.message)
-                except Exception as e:
-                    self.logger.warning(
-                        f"Failed {e} to serialize payload: {context.type} {context.method} {context.source}."
-                    )
-                    payload = default_serializer(context.message)
-
-            if len(payload) > self.max_payload_length:
-                payload = payload[: self.max_payload_length] + "..."
-
-            parts.append(f"payload={payload}")
-        return " ".join(parts)
-
-    async def on_message(
-        self, context: MiddlewareContext[Any], call_next: CallNext[Any, Any]
-    ) -> Any:
-        """Log all messages."""
-        message_info = self._format_message(context)
-        if self.methods and context.method not in self.methods:
-            return await call_next(context)
-
-        self.logger.log(self.log_level, f"Processing message: {message_info}")
-
-        try:
-            result = await call_next(context)
-            self.logger.log(
-                self.log_level, f"Completed message: {context.method or 'unknown'}"
-            )
-            return result
-        except Exception as e:
-            self.logger.log(
-                logging.ERROR, f"Failed message: {context.method or 'unknown'} - {e}"
-            )
-            raise
-
-
-class StructuredLoggingMiddleware(Middleware):
+class StructuredLoggingMiddleware(BaseLoggingMiddleware):
     """Middleware that provides structured JSON logging for better log analysis.
 
     Outputs structured logs that are easier to parse and analyze with log
@@ -129,9 +218,12 @@ class StructuredLoggingMiddleware(Middleware):
 
     def __init__(
         self,
+        *,
         logger: logging.Logger | None = None,
         log_level: int = logging.INFO,
         include_payloads: bool = False,
+        include_payload_length: bool = False,
+        estimate_payload_tokens: bool = False,
         methods: list[str] | None = None,
         payload_serializer: Callable[[Any], str] | None = None,
     ):
@@ -141,74 +233,24 @@ class StructuredLoggingMiddleware(Middleware):
             logger: Logger instance to use. If None, creates a logger named 'fastmcp.structured'
             log_level: Log level for messages (default: INFO)
             include_payloads: Whether to include message payloads in logs
+            include_payload_length: Whether to include payload size in logs
+            estimate_payload_tokens: Whether to estimate token count using length // 4
             methods: List of methods to log. If None, logs all methods.
-            serializer: Callable that converts objects to a JSON string for the
+            payload_serializer: Callable that converts objects to a JSON string for the
                 payload. If not provided, uses FastMCP's default tool serializer.
         """
-        self.logger: Logger = logger or logging.getLogger("fastmcp.structured")
+        self.logger: Logger = logger or logging.getLogger(
+            "fastmcp.middleware.structured_logging"
+        )
         self.log_level: int = log_level
         self.include_payloads: bool = include_payloads
+        self.include_payload_length: bool = include_payload_length
+        self.estimate_payload_tokens: bool = estimate_payload_tokens
         self.methods: list[str] | None = methods
         self.payload_serializer: Callable[[Any], str] | None = payload_serializer
+        self.max_payload_length: int | None = None
+        self.structured_logging: bool = True
 
-    def _create_log_entry(
-        self, context: MiddlewareContext[Any], event: str, **extra_fields: Any
-    ) -> dict[str, Any]:
-        """Create a structured log entry."""
-        entry = {
-            "event": event,
-            "timestamp": context.timestamp.isoformat(),
-            "source": context.source,
-            "type": context.type,
-            "method": context.method,
-            **extra_fields,
-        }
 
-        if self.include_payloads:
-            payload: str
-
-            if not self.payload_serializer:
-                payload = default_serializer(context.message)
-            else:
-                try:
-                    payload = self.payload_serializer(context.message)
-                except Exception as e:
-                    self.logger.warning(
-                        f"Failed {str(e)} to serialize payload: {context.type} {context.method} {context.source}."
-                    )
-                    payload = default_serializer(context.message)
-
-            entry["payload"] = payload
-
-        return entry
-
-    async def on_message(
-        self, context: MiddlewareContext[Any], call_next: CallNext[Any, Any]
-    ) -> Any:
-        """Log structured message information."""
-        start_entry = self._create_log_entry(context, "request_start")
-        if self.methods and context.method not in self.methods:
-            return await call_next(context)
-
-        self.logger.log(self.log_level, json.dumps(start_entry))
-
-        try:
-            result = await call_next(context)
-
-            success_entry = self._create_log_entry(
-                context,
-                "request_success",
-                result_type=type(result).__name__ if result else None,
-            )
-            self.logger.log(self.log_level, json.dumps(success_entry))
-
-            return result
-        except Exception as e:
-            error_entry = self._create_log_entry(
-                context,
-                "request_error",
-                error_type=type(e).__name__,
-                error_message=str(e),
-            )
-            self.logger.log(logging.ERROR, json.dumps(error_entry))
-            raise
+def _get_duration_ms(start_time: float, /) -> float:
+    return round(number=(time.perf_counter() - start_time) * 1000, ndigits=2)

fastmcp/server/middleware/middleware.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import logging
-from collections.abc import Awaitable
+from collections.abc import Awaitable, Sequence
 from dataclasses import dataclass, field, replace
 from datetime import datetime, timezone
 from functools import partial
@@ -15,6 +15,7 @@ from typing import (
 )
 
 import mcp.types as mt
+from mcp.server.lowlevel.helper_types import ReadResourceContents
 from typing_extensions import TypeVar
 
 from fastmcp.prompts.prompt import Prompt
@@ -26,9 +27,9 @@ if TYPE_CHECKING:
     from fastmcp.server.context import Context
 
 __all__ = [
+    "CallNext",
     "Middleware",
     "MiddlewareContext",
-    "CallNext",
 ]
 
 logger = logging.getLogger(__name__)
@@ -98,6 +99,8 @@ class Middleware:
         handler = call_next
 
         match context.method:
+            case "initialize":
+                handler = partial(self.on_initialize, call_next=handler)
             case "tools/call":
                 handler = partial(self.on_call_tool, call_next=handler)
             case "resources/read":
@@ -132,18 +135,25 @@ class Middleware:
 
     async def on_request(
         self,
-        context: MiddlewareContext[mt.Request],
-        call_next: CallNext[mt.Request, Any],
+        context: MiddlewareContext[mt.Request[Any, Any]],
+        call_next: CallNext[mt.Request[Any, Any], Any],
     ) -> Any:
         return await call_next(context)
 
     async def on_notification(
         self,
-        context: MiddlewareContext[mt.Notification],
-        call_next: CallNext[mt.Notification, Any],
+        context: MiddlewareContext[mt.Notification[Any, Any]],
+        call_next: CallNext[mt.Notification[Any, Any], Any],
     ) -> Any:
         return await call_next(context)
 
+    async def on_initialize(
+        self,
+        context: MiddlewareContext[mt.InitializeRequest],
+        call_next: CallNext[mt.InitializeRequest, None],
+    ) -> None:
+        return await call_next(context)
+
     async def on_call_tool(
         self,
         context: MiddlewareContext[mt.CallToolRequestParams],
@@ -154,8 +164,10 @@ class Middleware:
     async def on_read_resource(
         self,
         context: MiddlewareContext[mt.ReadResourceRequestParams],
-        call_next: CallNext[mt.ReadResourceRequestParams, mt.ReadResourceResult],
-    ) -> mt.ReadResourceResult:
+        call_next: CallNext[
+            mt.ReadResourceRequestParams, Sequence[ReadResourceContents]
+        ],
+    ) -> Sequence[ReadResourceContents]:
         return await call_next(context)
 
     async def on_get_prompt(
@@ -168,27 +180,29 @@ class Middleware:
     async def on_list_tools(
         self,
         context: MiddlewareContext[mt.ListToolsRequest],
-        call_next: CallNext[mt.ListToolsRequest, list[Tool]],
-    ) -> list[Tool]:
+        call_next: CallNext[mt.ListToolsRequest, Sequence[Tool]],
+    ) -> Sequence[Tool]:
         return await call_next(context)
 
     async def on_list_resources(
         self,
         context: MiddlewareContext[mt.ListResourcesRequest],
-        call_next: CallNext[mt.ListResourcesRequest, list[Resource]],
-    ) -> list[Resource]:
+        call_next: CallNext[mt.ListResourcesRequest, Sequence[Resource]],
+    ) -> Sequence[Resource]:
         return await call_next(context)
 
     async def on_list_resource_templates(
         self,
         context: MiddlewareContext[mt.ListResourceTemplatesRequest],
-        call_next: CallNext[mt.ListResourceTemplatesRequest, list[ResourceTemplate]],
-    ) -> list[ResourceTemplate]:
+        call_next: CallNext[
+            mt.ListResourceTemplatesRequest, Sequence[ResourceTemplate]
+        ],
+    ) -> Sequence[ResourceTemplate]:
         return await call_next(context)
 
     async def on_list_prompts(
         self,
         context: MiddlewareContext[mt.ListPromptsRequest],
-        call_next: CallNext[mt.ListPromptsRequest, list[Prompt]],
-    ) -> list[Prompt]:
+        call_next: CallNext[mt.ListPromptsRequest, Sequence[Prompt]],
+    ) -> Sequence[Prompt]:
         return await call_next(context)

fastmcp/server/middleware/rate_limiting.py

@@ -1,11 +1,11 @@
 """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
 
+import anyio
 from mcp import McpError
 from mcp.types import ErrorData
 
@@ -33,7 +33,7 @@ class TokenBucketRateLimiter:
         self.refill_rate = refill_rate
         self.tokens = capacity
         self.last_refill = time.time()
-        self._lock = asyncio.Lock()
+        self._lock = anyio.Lock()
 
     async def consume(self, tokens: int = 1) -> bool:
         """Try to consume tokens from the bucket.
@@ -71,7 +71,7 @@ class SlidingWindowRateLimiter:
         self.max_requests = max_requests
         self.window_seconds = window_seconds
         self.requests = deque()
-        self._lock = asyncio.Lock()
+        self._lock = anyio.Lock()
 
     async def is_allowed(self) -> bool:
         """Check if a request is allowed."""

fastmcp/server/middleware/tool_injection.py

@@ -0,0 +1,116 @@
+"""A middleware for injecting tools into the MCP server context."""
+
+from collections.abc import Sequence
+from logging import Logger
+from typing import Annotated, Any
+
+import mcp.types
+from mcp.server.lowlevel.helper_types import ReadResourceContents
+from mcp.types import Prompt
+from pydantic import AnyUrl
+from typing_extensions import override
+
+from fastmcp.server.context import Context
+from fastmcp.server.middleware.middleware import CallNext, Middleware, MiddlewareContext
+from fastmcp.tools.tool import Tool, ToolResult
+from fastmcp.utilities.logging import get_logger
+
+logger: Logger = get_logger(name=__name__)
+
+
+class ToolInjectionMiddleware(Middleware):
+    """A middleware for injecting tools into the context."""
+
+    def __init__(self, tools: Sequence[Tool]):
+        """Initialize the tool injection middleware."""
+        self._tools_to_inject: Sequence[Tool] = tools
+        self._tools_to_inject_by_name: dict[str, Tool] = {
+            tool.name: tool for tool in tools
+        }
+
+    @override
+    async def on_list_tools(
+        self,
+        context: MiddlewareContext[mcp.types.ListToolsRequest],
+        call_next: CallNext[mcp.types.ListToolsRequest, Sequence[Tool]],
+    ) -> Sequence[Tool]:
+        """Inject tools into the response."""
+        return [*self._tools_to_inject, *await call_next(context)]
+
+    @override
+    async def on_call_tool(
+        self,
+        context: MiddlewareContext[mcp.types.CallToolRequestParams],
+        call_next: CallNext[mcp.types.CallToolRequestParams, ToolResult],
+    ) -> ToolResult:
+        """Intercept tool calls to injected tools."""
+        if context.message.name in self._tools_to_inject_by_name:
+            tool = self._tools_to_inject_by_name[context.message.name]
+            return await tool.run(arguments=context.message.arguments or {})
+
+        return await call_next(context)
+
+
+async def list_prompts(context: Context) -> list[Prompt]:
+    """List prompts available on the server."""
+    return await context.list_prompts()
+
+
+list_prompts_tool = Tool.from_function(
+    fn=list_prompts,
+)
+
+
+async def get_prompt(
+    context: Context,
+    name: Annotated[str, "The name of the prompt to render."],
+    arguments: Annotated[
+        dict[str, Any] | None, "The arguments to pass to the prompt."
+    ] = None,
+) -> mcp.types.GetPromptResult:
+    """Render a prompt available on the server."""
+    return await context.get_prompt(name=name, arguments=arguments)
+
+
+get_prompt_tool = Tool.from_function(
+    fn=get_prompt,
+)
+
+
+class PromptToolMiddleware(ToolInjectionMiddleware):
+    """A middleware for injecting prompts as tools into the context."""
+
+    def __init__(self) -> None:
+        tools: list[Tool] = [list_prompts_tool, get_prompt_tool]
+        super().__init__(tools=tools)
+
+
+async def list_resources(context: Context) -> list[mcp.types.Resource]:
+    """List resources available on the server."""
+    return await context.list_resources()
+
+
+list_resources_tool = Tool.from_function(
+    fn=list_resources,
+)
+
+
+async def read_resource(
+    context: Context,
+    uri: Annotated[AnyUrl | str, "The URI of the resource to read."],
+) -> list[ReadResourceContents]:
+    """Read a resource available on the server."""
+    return await context.read_resource(uri=uri)
+
+
+read_resource_tool = Tool.from_function(
+    fn=read_resource,
+)
+
+
+class ResourceToolMiddleware(ToolInjectionMiddleware):
+    """A middleware for injecting resources as tools into the context."""
+
+    def __init__(self) -> None:
+        tools: list[Tool] = [list_resources_tool, read_resource_tool]
+        super().__init__(tools=tools)

fastmcp/server/openapi.py

@@ -513,11 +513,11 @@ class OpenAPITool(Tool):
                 if e.response.text:
                     error_message += f" - {e.response.text}"
 
-            raise ValueError(error_message)
+            raise ValueError(error_message) from e
 
         except httpx.RequestError as e:
             # Handle request errors (connection, timeout, etc.)
-            raise ValueError(f"Request error: {str(e)}")
+            raise ValueError(f"Request error: {e!s}") from e
 
 
 class OpenAPIResource(Resource):
@@ -531,9 +531,11 @@ class OpenAPIResource(Resource):
         name: str,
         description: str,
         mime_type: str = "application/json",
-        tags: set[str] = set(),
+        tags: set[str] | None = None,
         timeout: float | None = None,
     ):
+        if tags is None:
+            tags = set()
         super().__init__(
             uri=AnyUrl(uri),  # Convert string to AnyUrl
             name=name,
@@ -632,11 +634,11 @@ class OpenAPIResource(Resource):
                 if e.response.text:
                     error_message += f" - {e.response.text}"
 
-            raise ValueError(error_message)
+            raise ValueError(error_message) from e
 
         except httpx.RequestError as e:
             # Handle request errors (connection, timeout, etc.)
-            raise ValueError(f"Request error: {str(e)}")
+            raise ValueError(f"Request error: {e!s}") from e
 
 
 class OpenAPIResourceTemplate(ResourceTemplate):
@@ -650,9 +652,11 @@ class OpenAPIResourceTemplate(ResourceTemplate):
         name: str,
         description: str,
         parameters: dict[str, Any],
-        tags: set[str] = set(),
+        tags: set[str] | None = None,
         timeout: float | None = None,
     ):
+        if tags is None:
+            tags = set()
         super().__init__(
             uri_template=uri_template,
             name=name,
@@ -785,6 +789,7 @@ class FastMCPOpenAPI(FastMCP):
         http_routes = openapi.parse_openapi_to_http_routes(openapi_spec)
 
         # Process routes
+        num_excluded = 0
         route_maps = (route_maps or []) + DEFAULT_ROUTE_MAPPINGS
         for route in http_routes:
             # Determine route type based on mappings or default rules
@@ -823,8 +828,11 @@ class FastMCPOpenAPI(FastMCP):
                 self._create_openapi_template(route, component_name, tags=route_tags)
             elif route_type == MCPType.EXCLUDE:
                 logger.info(f"Excluding route: {route.method} {route.path}")
+                num_excluded += 1
 
-        logger.info(f"Created FastMCP OpenAPI server with {len(http_routes)} routes")
+        logger.info(
+            f"Created FastMCP OpenAPI server with {len(http_routes) - num_excluded} routes"
+        )
 
     def _generate_default_name(
         self, route: openapi.HTTPRoute, mcp_names_map: dict[str, str] | None = None