mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-23 14:04:18 +02:00
Add production-ready middleware examples with comprehensive tests
Implements timing, logging, rate limiting, and error handling middleware to showcase FastMCP's middleware capabilities as a headline feature. Each middleware includes: - Production-ready implementations with full configurability - Comprehensive unit and integration tests with real FastMCP servers - Updated documentation with teaching examples and production usage patterns 🤖 Generated with [Claude Code](https://claude.ai/code) Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
Jeremiah Lowin
parent
d76f42ceb8
commit
6d033eeaf1
10 changed files with 2763 additions and 60 deletions
|
|
@ -329,93 +329,246 @@ parent.mount(child, prefix="child")
|
|||
|
||||
When a client calls "child_tool", the request will flow through the parent's authentication middleware first, then route to the child server where it will go through the child's logging middleware.
|
||||
|
||||
## Examples
|
||||
## Built-in Middleware Examples
|
||||
|
||||
### Authentication Middleware
|
||||
FastMCP includes several production-ready middleware implementations that demonstrate best practices and provide immediately useful functionality. Let's explore how each type works by building simplified versions, then see how to use the full implementations.
|
||||
|
||||
This middleware checks for a valid authorization token on all requests:
|
||||
### Timing Middleware
|
||||
|
||||
```python
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
from fastmcp.exceptions import ToolError
|
||||
Performance monitoring is essential for understanding your server's behavior and identifying bottlenecks. FastMCP includes production-ready timing middleware at `fastmcp.server.middleware.timing`.
|
||||
|
||||
class AuthenticationMiddleware(Middleware):
|
||||
def __init__(self, required_token: str):
|
||||
self.required_token = required_token
|
||||
|
||||
async def on_request(self, context: MiddlewareContext, call_next):
|
||||
if hasattr(context, 'fastmcp_context') and context.fastmcp_context:
|
||||
try:
|
||||
request = context.fastmcp_context.get_http_request()
|
||||
auth_header = request.headers.get("Authorization")
|
||||
|
||||
if not auth_header or not auth_header.startswith("Bearer "):
|
||||
raise ToolError("Missing or invalid authorization header")
|
||||
|
||||
token = auth_header.split(" ", 1)[1]
|
||||
if token != self.required_token:
|
||||
raise ToolError("Invalid authentication token")
|
||||
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
return await call_next(context)
|
||||
|
||||
# Usage
|
||||
mcp = FastMCP("SecureServer")
|
||||
mcp.add_middleware(AuthenticationMiddleware("secret-token-123"))
|
||||
```
|
||||
|
||||
### Performance Monitoring Middleware
|
||||
|
||||
This middleware tracks how long tools take to execute:
|
||||
Here's an example of how it works:
|
||||
|
||||
```python
|
||||
import time
|
||||
import logging
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
|
||||
class PerformanceMiddleware(Middleware):
|
||||
def __init__(self):
|
||||
self.logger = logging.getLogger("performance")
|
||||
|
||||
async def on_call_tool(self, context: MiddlewareContext, call_next):
|
||||
tool_name = context.message.name
|
||||
start_time = time.time()
|
||||
class SimpleTimingMiddleware(Middleware):
|
||||
async def on_request(self, context: MiddlewareContext, call_next):
|
||||
start_time = time.perf_counter()
|
||||
|
||||
try:
|
||||
result = await call_next(context)
|
||||
execution_time = time.time() - start_time
|
||||
|
||||
self.logger.info(
|
||||
f"Tool {tool_name} completed in {execution_time:.3f}s"
|
||||
)
|
||||
|
||||
duration_ms = (time.perf_counter() - start_time) * 1000
|
||||
print(f"Request {context.method} completed in {duration_ms:.2f}ms")
|
||||
return result
|
||||
|
||||
except Exception as e:
|
||||
execution_time = time.time() - start_time
|
||||
self.logger.error(
|
||||
f"Tool {tool_name} failed after {execution_time:.3f}s: {e}"
|
||||
)
|
||||
duration_ms = (time.perf_counter() - start_time) * 1000
|
||||
print(f"Request {context.method} failed after {duration_ms:.2f}ms: {e}")
|
||||
raise
|
||||
```
|
||||
|
||||
### Request Transformation Middleware
|
||||
|
||||
This middleware adds metadata to tool calls:
|
||||
To use the full production version with proper logging and configuration:
|
||||
|
||||
```python
|
||||
class TransformationMiddleware(Middleware):
|
||||
async def on_call_tool(self, context: MiddlewareContext, call_next):
|
||||
if hasattr(context.message, 'arguments'):
|
||||
args = context.message.arguments or {}
|
||||
args['_middleware_timestamp'] = context.timestamp.isoformat()
|
||||
|
||||
modified_context = context.copy(
|
||||
message=context.message.model_copy(update={'arguments': args})
|
||||
)
|
||||
else:
|
||||
modified_context = context
|
||||
from fastmcp.server.middleware.timing import (
|
||||
TimingMiddleware,
|
||||
DetailedTimingMiddleware
|
||||
)
|
||||
|
||||
# Basic timing for all requests
|
||||
mcp.add_middleware(TimingMiddleware())
|
||||
|
||||
# Detailed per-operation timing (tools, resources, prompts)
|
||||
mcp.add_middleware(DetailedTimingMiddleware())
|
||||
```
|
||||
|
||||
The built-in versions include custom logger support, proper formatting, and **DetailedTimingMiddleware** provides operation-specific hooks like `on_call_tool` and `on_read_resource` for granular timing.
|
||||
|
||||
### Logging Middleware
|
||||
|
||||
Request and response logging is crucial for debugging, monitoring, and understanding usage patterns in your MCP server. FastMCP provides comprehensive logging middleware at `fastmcp.server.middleware.logging`.
|
||||
|
||||
Here's an example of how it works:
|
||||
|
||||
```python
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
|
||||
class SimpleLoggingMiddleware(Middleware):
|
||||
async def on_message(self, context: MiddlewareContext, call_next):
|
||||
print(f"Processing {context.method} from {context.source}")
|
||||
|
||||
return await call_next(modified_context)
|
||||
try:
|
||||
result = await call_next(context)
|
||||
print(f"Completed {context.method}")
|
||||
return result
|
||||
except Exception as e:
|
||||
print(f"Failed {context.method}: {e}")
|
||||
raise
|
||||
```
|
||||
|
||||
To use the full production versions with advanced features:
|
||||
|
||||
```python
|
||||
from fastmcp.server.middleware.logging import (
|
||||
LoggingMiddleware,
|
||||
StructuredLoggingMiddleware
|
||||
)
|
||||
|
||||
# Human-readable logging with payload support
|
||||
mcp.add_middleware(LoggingMiddleware(
|
||||
include_payloads=True,
|
||||
max_payload_length=1000
|
||||
))
|
||||
|
||||
# JSON-structured logging for log aggregation tools
|
||||
mcp.add_middleware(StructuredLoggingMiddleware(include_payloads=True))
|
||||
```
|
||||
|
||||
The production versions include payload logging, structured JSON output, custom logger support, payload size limits, and operation-specific hooks for granular control.
|
||||
|
||||
### Rate Limiting Middleware
|
||||
|
||||
Rate limiting is essential for protecting your server from abuse, ensuring fair resource usage, and maintaining performance under load. FastMCP includes sophisticated rate limiting middleware at `fastmcp.server.middleware.rate_limiting`.
|
||||
|
||||
Here's an example of how it works:
|
||||
|
||||
```python
|
||||
import time
|
||||
from collections import defaultdict
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
from mcp import McpError
|
||||
from mcp.types import ErrorData
|
||||
|
||||
class SimpleRateLimitMiddleware(Middleware):
|
||||
def __init__(self, requests_per_minute: int = 60):
|
||||
self.requests_per_minute = requests_per_minute
|
||||
self.client_requests = defaultdict(list)
|
||||
|
||||
async def on_request(self, context: MiddlewareContext, call_next):
|
||||
current_time = time.time()
|
||||
client_id = "default" # In practice, extract from headers or context
|
||||
|
||||
# Clean old requests and check limit
|
||||
cutoff_time = current_time - 60
|
||||
self.client_requests[client_id] = [
|
||||
req_time for req_time in self.client_requests[client_id]
|
||||
if req_time > cutoff_time
|
||||
]
|
||||
|
||||
if len(self.client_requests[client_id]) >= self.requests_per_minute:
|
||||
raise McpError(ErrorData(code=-32000, message="Rate limit exceeded"))
|
||||
|
||||
self.client_requests[client_id].append(current_time)
|
||||
return await call_next(context)
|
||||
```
|
||||
|
||||
To use the full production versions with advanced algorithms:
|
||||
|
||||
```python
|
||||
from fastmcp.server.middleware.rate_limiting import (
|
||||
RateLimitingMiddleware,
|
||||
SlidingWindowRateLimitingMiddleware
|
||||
)
|
||||
|
||||
# Token bucket rate limiting (allows controlled bursts)
|
||||
mcp.add_middleware(RateLimitingMiddleware(
|
||||
max_requests_per_second=10.0,
|
||||
burst_capacity=20
|
||||
))
|
||||
|
||||
# Sliding window rate limiting (precise time-based control)
|
||||
mcp.add_middleware(SlidingWindowRateLimitingMiddleware(
|
||||
max_requests=100,
|
||||
window_minutes=1
|
||||
))
|
||||
```
|
||||
|
||||
The production versions include token bucket algorithms, per-client identification, global rate limiting, and async-safe implementations with configurable client identification functions.
|
||||
|
||||
### Error Handling Middleware
|
||||
|
||||
Consistent error handling and recovery is critical for robust MCP servers. FastMCP provides comprehensive error handling middleware at `fastmcp.server.middleware.error_handling`.
|
||||
|
||||
Here's an example of how it works:
|
||||
|
||||
```python
|
||||
import logging
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
|
||||
class SimpleErrorHandlingMiddleware(Middleware):
|
||||
def __init__(self):
|
||||
self.logger = logging.getLogger("errors")
|
||||
self.error_counts = {}
|
||||
|
||||
async def on_message(self, context: MiddlewareContext, call_next):
|
||||
try:
|
||||
return await call_next(context)
|
||||
except Exception as error:
|
||||
# Log the error and track statistics
|
||||
error_key = f"{type(error).__name__}:{context.method}"
|
||||
self.error_counts[error_key] = self.error_counts.get(error_key, 0) + 1
|
||||
|
||||
self.logger.error(f"Error in {context.method}: {type(error).__name__}: {error}")
|
||||
raise
|
||||
```
|
||||
|
||||
To use the full production versions with advanced features:
|
||||
|
||||
```python
|
||||
from fastmcp.server.middleware.error_handling import (
|
||||
ErrorHandlingMiddleware,
|
||||
RetryMiddleware
|
||||
)
|
||||
|
||||
# Comprehensive error logging and transformation
|
||||
mcp.add_middleware(ErrorHandlingMiddleware(
|
||||
include_traceback=True,
|
||||
transform_errors=True,
|
||||
error_callback=my_error_callback
|
||||
))
|
||||
|
||||
# Automatic retry with exponential backoff
|
||||
mcp.add_middleware(RetryMiddleware(
|
||||
max_retries=3,
|
||||
retry_exceptions=(ConnectionError, TimeoutError)
|
||||
))
|
||||
```
|
||||
|
||||
The production versions include error transformation, custom callbacks, configurable retry logic, and proper MCP error formatting.
|
||||
|
||||
### Combining Middleware
|
||||
|
||||
These middleware work together seamlessly:
|
||||
|
||||
```python
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.server.middleware.timing import TimingMiddleware
|
||||
from fastmcp.server.middleware.logging import LoggingMiddleware
|
||||
from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
|
||||
from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
|
||||
|
||||
mcp = FastMCP("Production Server")
|
||||
|
||||
# Add middleware in logical order
|
||||
mcp.add_middleware(ErrorHandlingMiddleware()) # Handle errors first
|
||||
mcp.add_middleware(RateLimitingMiddleware(max_requests_per_second=50))
|
||||
mcp.add_middleware(TimingMiddleware()) # Time actual execution
|
||||
mcp.add_middleware(LoggingMiddleware()) # Log everything
|
||||
|
||||
@mcp.tool
|
||||
def my_tool(data: str) -> str:
|
||||
return f"Processed: {data}"
|
||||
```
|
||||
|
||||
This configuration provides comprehensive monitoring, protection, and observability for your MCP server.
|
||||
|
||||
### Custom Middleware Example
|
||||
|
||||
You can also create custom middleware by extending the base class:
|
||||
|
||||
```python
|
||||
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
||||
|
||||
class CustomHeaderMiddleware(Middleware):
|
||||
async def on_request(self, context: MiddlewareContext, call_next):
|
||||
# Add custom logic here
|
||||
print(f"Processing {context.method}")
|
||||
|
||||
result = await call_next(context)
|
||||
|
||||
print(f"Completed {context.method}")
|
||||
return result
|
||||
|
||||
mcp.add_middleware(CustomHeaderMiddleware())
|
||||
```
|
||||
|
|
@ -1 +1,6 @@
|
|||
from .middleware import Middleware, MiddlewareContext
|
||||
|
||||
__all__ = [
|
||||
"Middleware",
|
||||
"MiddlewareContext",
|
||||
]
|
||||
|
|
|
|||
206
src/fastmcp/server/middleware/error_handling.py
Normal file
206
src/fastmcp/server/middleware/error_handling.py
Normal file
|
|
@ -0,0 +1,206 @@
|
|||
"""Error handling middleware for consistent error responses and tracking."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import traceback
|
||||
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 ErrorHandlingMiddleware(Middleware):
|
||||
"""Middleware that provides consistent error handling and logging.
|
||||
|
||||
Catches exceptions, logs them appropriately, and converts them to
|
||||
proper MCP error responses. Also tracks error patterns for monitoring.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
|
||||
import logging
|
||||
|
||||
# Configure logging to see error details
|
||||
logging.basicConfig(level=logging.ERROR)
|
||||
|
||||
mcp = FastMCP("MyServer")
|
||||
mcp.add_middleware(ErrorHandlingMiddleware())
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logger: logging.Logger | None = None,
|
||||
include_traceback: bool = False,
|
||||
error_callback: Callable[[Exception, MiddlewareContext], None] | None = None,
|
||||
transform_errors: bool = True,
|
||||
):
|
||||
"""Initialize error handling middleware.
|
||||
|
||||
Args:
|
||||
logger: Logger instance for error logging. If None, uses 'fastmcp.errors'
|
||||
include_traceback: Whether to include full traceback in error logs
|
||||
error_callback: Optional callback function called for each error
|
||||
transform_errors: Whether to transform non-MCP errors to McpError
|
||||
"""
|
||||
self.logger = logger or logging.getLogger("fastmcp.errors")
|
||||
self.include_traceback = include_traceback
|
||||
self.error_callback = error_callback
|
||||
self.transform_errors = transform_errors
|
||||
self.error_counts = {}
|
||||
|
||||
def _log_error(self, error: Exception, context: MiddlewareContext) -> None:
|
||||
"""Log error with appropriate detail level."""
|
||||
error_type = type(error).__name__
|
||||
method = context.method or "unknown"
|
||||
|
||||
# Track error counts
|
||||
error_key = f"{error_type}:{method}"
|
||||
self.error_counts[error_key] = self.error_counts.get(error_key, 0) + 1
|
||||
|
||||
base_message = f"Error in {method}: {error_type}: {str(error)}"
|
||||
|
||||
if self.include_traceback:
|
||||
self.logger.error(f"{base_message}\n{traceback.format_exc()}")
|
||||
else:
|
||||
self.logger.error(base_message)
|
||||
|
||||
# Call custom error callback if provided
|
||||
if self.error_callback:
|
||||
try:
|
||||
self.error_callback(error, context)
|
||||
except Exception as callback_error:
|
||||
self.logger.error(f"Error in error callback: {callback_error}")
|
||||
|
||||
def _transform_error(self, error: Exception) -> Exception:
|
||||
"""Transform non-MCP errors to proper MCP errors."""
|
||||
if isinstance(error, McpError):
|
||||
return error
|
||||
|
||||
if not self.transform_errors:
|
||||
return error
|
||||
|
||||
# Map common exceptions to appropriate MCP error codes
|
||||
error_type = type(error)
|
||||
|
||||
if error_type in (ValueError, TypeError):
|
||||
return McpError(
|
||||
ErrorData(code=-32602, message=f"Invalid params: {str(error)}")
|
||||
)
|
||||
elif error_type in (FileNotFoundError, KeyError):
|
||||
return McpError(
|
||||
ErrorData(code=-32001, message=f"Resource not found: {str(error)}")
|
||||
)
|
||||
elif error_type is PermissionError:
|
||||
return McpError(
|
||||
ErrorData(code=-32000, message=f"Permission denied: {str(error)}")
|
||||
)
|
||||
elif error_type in (TimeoutError, asyncio.TimeoutError):
|
||||
return McpError(
|
||||
ErrorData(code=-32000, message=f"Request timeout: {str(error)}")
|
||||
)
|
||||
else:
|
||||
return McpError(
|
||||
ErrorData(code=-32603, message=f"Internal error: {str(error)}")
|
||||
)
|
||||
|
||||
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
||||
"""Handle errors for all messages."""
|
||||
try:
|
||||
return await call_next(context)
|
||||
except Exception as error:
|
||||
self._log_error(error, context)
|
||||
|
||||
# Transform and re-raise
|
||||
transformed_error = self._transform_error(error)
|
||||
raise transformed_error
|
||||
|
||||
def get_error_stats(self) -> dict[str, int]:
|
||||
"""Get error statistics for monitoring."""
|
||||
return self.error_counts.copy()
|
||||
|
||||
|
||||
class RetryMiddleware(Middleware):
|
||||
"""Middleware that implements automatic retry logic for failed requests.
|
||||
|
||||
Retries requests that fail with transient errors, using exponential
|
||||
backoff to avoid overwhelming the server or external dependencies.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastmcp.server.middleware.error_handling import RetryMiddleware
|
||||
|
||||
# Retry up to 3 times with exponential backoff
|
||||
retry_middleware = RetryMiddleware(
|
||||
max_retries=3,
|
||||
retry_exceptions=(ConnectionError, TimeoutError)
|
||||
)
|
||||
|
||||
mcp = FastMCP("MyServer")
|
||||
mcp.add_middleware(retry_middleware)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
max_retries: int = 3,
|
||||
base_delay: float = 1.0,
|
||||
max_delay: float = 60.0,
|
||||
backoff_multiplier: float = 2.0,
|
||||
retry_exceptions: tuple[type[Exception], ...] = (ConnectionError, TimeoutError),
|
||||
logger: logging.Logger | None = None,
|
||||
):
|
||||
"""Initialize retry middleware.
|
||||
|
||||
Args:
|
||||
max_retries: Maximum number of retry attempts
|
||||
base_delay: Initial delay between retries in seconds
|
||||
max_delay: Maximum delay between retries in seconds
|
||||
backoff_multiplier: Multiplier for exponential backoff
|
||||
retry_exceptions: Tuple of exception types that should trigger retries
|
||||
logger: Logger for retry attempts
|
||||
"""
|
||||
self.max_retries = max_retries
|
||||
self.base_delay = base_delay
|
||||
self.max_delay = max_delay
|
||||
self.backoff_multiplier = backoff_multiplier
|
||||
self.retry_exceptions = retry_exceptions
|
||||
self.logger = logger or logging.getLogger("fastmcp.retry")
|
||||
|
||||
def _should_retry(self, error: Exception) -> bool:
|
||||
"""Determine if an error should trigger a retry."""
|
||||
return isinstance(error, self.retry_exceptions)
|
||||
|
||||
def _calculate_delay(self, attempt: int) -> float:
|
||||
"""Calculate delay for the given attempt number."""
|
||||
delay = self.base_delay * (self.backoff_multiplier**attempt)
|
||||
return min(delay, self.max_delay)
|
||||
|
||||
async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
||||
"""Implement retry logic for requests."""
|
||||
last_error = None
|
||||
|
||||
for attempt in range(self.max_retries + 1):
|
||||
try:
|
||||
return await call_next(context)
|
||||
except Exception as error:
|
||||
last_error = error
|
||||
|
||||
# Don't retry on the last attempt or if it's not a retryable error
|
||||
if attempt == self.max_retries or not self._should_retry(error):
|
||||
break
|
||||
|
||||
delay = self._calculate_delay(attempt)
|
||||
self.logger.warning(
|
||||
f"Request {context.method} failed (attempt {attempt + 1}/{self.max_retries + 1}): "
|
||||
f"{type(error).__name__}: {str(error)}. Retrying in {delay:.1f}s..."
|
||||
)
|
||||
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Re-raise the last error if all retries failed
|
||||
if last_error:
|
||||
raise last_error
|
||||
165
src/fastmcp/server/middleware/logging.py
Normal file
165
src/fastmcp/server/middleware/logging.py
Normal file
|
|
@ -0,0 +1,165 @@
|
|||
"""Comprehensive logging middleware for FastMCP servers."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
from .middleware import CallNext, Middleware, MiddlewareContext
|
||||
|
||||
|
||||
class LoggingMiddleware(Middleware):
|
||||
"""Middleware that provides comprehensive request and response logging.
|
||||
|
||||
Logs all MCP messages with configurable detail levels. Useful for debugging,
|
||||
monitoring, and understanding server usage patterns.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastmcp.server.middleware.logging import LoggingMiddleware
|
||||
import logging
|
||||
|
||||
# Configure logging
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
|
||||
mcp = FastMCP("MyServer")
|
||||
mcp.add_middleware(LoggingMiddleware())
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logger: logging.Logger | None = None,
|
||||
log_level: int = logging.INFO,
|
||||
include_payloads: bool = False,
|
||||
max_payload_length: int = 1000,
|
||||
):
|
||||
"""Initialize logging middleware.
|
||||
|
||||
Args:
|
||||
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
|
||||
max_payload_length: Maximum length of payload to log (prevents huge logs)
|
||||
"""
|
||||
self.logger = logger or logging.getLogger("fastmcp.requests")
|
||||
self.log_level = log_level
|
||||
self.include_payloads = include_payloads
|
||||
self.max_payload_length = max_payload_length
|
||||
|
||||
def _format_message(self, context: MiddlewareContext) -> 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 and hasattr(context.message, "__dict__"):
|
||||
try:
|
||||
payload = json.dumps(context.message.__dict__, default=str)
|
||||
if len(payload) > self.max_payload_length:
|
||||
payload = payload[: self.max_payload_length] + "..."
|
||||
parts.append(f"payload={payload}")
|
||||
except (TypeError, ValueError):
|
||||
parts.append("payload=<non-serializable>")
|
||||
|
||||
return " ".join(parts)
|
||||
|
||||
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
||||
"""Log all messages."""
|
||||
message_info = self._format_message(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):
|
||||
"""Middleware that provides structured JSON logging for better log analysis.
|
||||
|
||||
Outputs structured logs that are easier to parse and analyze with log
|
||||
aggregation tools like ELK stack, Splunk, or cloud logging services.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from fastmcp.server.middleware.logging import StructuredLoggingMiddleware
|
||||
import logging
|
||||
|
||||
mcp = FastMCP("MyServer")
|
||||
mcp.add_middleware(StructuredLoggingMiddleware())
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
logger: logging.Logger | None = None,
|
||||
log_level: int = logging.INFO,
|
||||
include_payloads: bool = False,
|
||||
):
|
||||
"""Initialize structured logging middleware.
|
||||
|
||||
Args:
|
||||
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
|
||||
"""
|
||||
self.logger = logger or logging.getLogger("fastmcp.structured")
|
||||
self.log_level = log_level
|
||||
self.include_payloads = include_payloads
|
||||
|
||||
def _create_log_entry(
|
||||
self, context: MiddlewareContext, event: str, **extra_fields
|
||||
) -> dict:
|
||||
"""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 and hasattr(context.message, "__dict__"):
|
||||
try:
|
||||
entry["payload"] = context.message.__dict__
|
||||
except (TypeError, ValueError):
|
||||
entry["payload"] = "<non-serializable>"
|
||||
|
||||
return entry
|
||||
|
||||
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
||||
"""Log structured message information."""
|
||||
start_entry = self._create_log_entry(context, "request_start")
|
||||
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
|
||||
231
src/fastmcp/server/middleware/rate_limiting.py
Normal file
231
src/fastmcp/server/middleware/rate_limiting.py
Normal file
|
|
@ -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)
|
||||
156
src/fastmcp/server/middleware/timing.py
Normal file
156
src/fastmcp/server/middleware/timing.py
Normal file
|
|
@ -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")
|
||||
601
tests/server/middleware/test_error_handling.py
Normal file
601
tests/server/middleware/test_error_handling.py
Normal file
|
|
@ -0,0 +1,601 @@
|
|||
"""Tests for error handling middleware."""
|
||||
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from mcp import McpError
|
||||
|
||||
from fastmcp.server.middleware.error_handling import (
|
||||
ErrorHandlingMiddleware,
|
||||
RetryMiddleware,
|
||||
)
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock middleware context."""
|
||||
context = MagicMock(spec=MiddlewareContext)
|
||||
context.method = "test_method"
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_call_next():
|
||||
"""Create a mock call_next function."""
|
||||
return AsyncMock(return_value="test_result")
|
||||
|
||||
|
||||
class TestErrorHandlingMiddleware:
|
||||
"""Test error handling middleware functionality."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
assert middleware.logger.name == "fastmcp.errors"
|
||||
assert middleware.include_traceback is False
|
||||
assert middleware.error_callback is None
|
||||
assert middleware.transform_errors is True
|
||||
assert middleware.error_counts == {}
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
logger = logging.getLogger("custom")
|
||||
callback = MagicMock()
|
||||
|
||||
middleware = ErrorHandlingMiddleware(
|
||||
logger=logger,
|
||||
include_traceback=True,
|
||||
error_callback=callback,
|
||||
transform_errors=False,
|
||||
)
|
||||
assert middleware.logger is logger
|
||||
assert middleware.include_traceback is True
|
||||
assert middleware.error_callback is callback
|
||||
assert middleware.transform_errors is False
|
||||
|
||||
def test_log_error_basic(self, mock_context, caplog):
|
||||
"""Test basic error logging."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error = ValueError("test error")
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
assert "Error in test_method: ValueError: test error" in caplog.text
|
||||
assert "ValueError:test_method" in middleware.error_counts
|
||||
assert middleware.error_counts["ValueError:test_method"] == 1
|
||||
|
||||
def test_log_error_with_traceback(self, mock_context, caplog):
|
||||
"""Test error logging with traceback."""
|
||||
middleware = ErrorHandlingMiddleware(include_traceback=True)
|
||||
error = ValueError("test error")
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
assert "Error in test_method: ValueError: test error" in caplog.text
|
||||
# The traceback is added to the log message
|
||||
assert "Error in test_method: ValueError: test error" in caplog.text
|
||||
|
||||
def test_log_error_with_callback(self, mock_context):
|
||||
"""Test error logging with callback."""
|
||||
callback = MagicMock()
|
||||
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
||||
error = ValueError("test error")
|
||||
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
callback.assert_called_once_with(error, mock_context)
|
||||
|
||||
def test_log_error_callback_exception(self, mock_context, caplog):
|
||||
"""Test error logging when callback raises exception."""
|
||||
callback = MagicMock(side_effect=RuntimeError("callback error"))
|
||||
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
||||
error = ValueError("test error")
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
middleware._log_error(error, mock_context)
|
||||
|
||||
assert "Error in error callback: callback error" in caplog.text
|
||||
|
||||
def test_transform_error_mcp_error(self):
|
||||
"""Test that MCP errors are not transformed."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
from mcp.types import ErrorData
|
||||
|
||||
error = McpError(ErrorData(code=-32001, message="test error"))
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert result is error
|
||||
|
||||
def test_transform_error_disabled(self):
|
||||
"""Test error transformation when disabled."""
|
||||
middleware = ErrorHandlingMiddleware(transform_errors=False)
|
||||
error = ValueError("test error")
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert result is error
|
||||
|
||||
def test_transform_error_value_error(self):
|
||||
"""Test transforming ValueError."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error = ValueError("test error")
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert isinstance(result, McpError)
|
||||
assert result.error.code == -32602
|
||||
assert "Invalid params: test error" in result.error.message
|
||||
|
||||
def test_transform_error_file_not_found(self):
|
||||
"""Test transforming FileNotFoundError."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error = FileNotFoundError("test error")
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert isinstance(result, McpError)
|
||||
assert result.error.code == -32001
|
||||
assert "Resource not found: test error" in result.error.message
|
||||
|
||||
def test_transform_error_permission_error(self):
|
||||
"""Test transforming PermissionError."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error = PermissionError("test error")
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert isinstance(result, McpError)
|
||||
assert result.error.code == -32000
|
||||
assert "Permission denied: test error" in result.error.message
|
||||
|
||||
def test_transform_error_timeout_error(self):
|
||||
"""Test transforming TimeoutError."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error = TimeoutError("test error")
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert isinstance(result, McpError)
|
||||
assert result.error.code == -32000
|
||||
assert "Request timeout: test error" in result.error.message
|
||||
|
||||
def test_transform_error_generic(self):
|
||||
"""Test transforming generic error."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error = RuntimeError("test error")
|
||||
|
||||
result = middleware._transform_error(error)
|
||||
|
||||
assert isinstance(result, McpError)
|
||||
assert result.error.code == -32603
|
||||
assert "Internal error: test error" in result.error.message
|
||||
|
||||
async def test_on_message_success(self, mock_context, mock_call_next):
|
||||
"""Test successful message handling."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.called
|
||||
|
||||
async def test_on_message_error_transform(self, mock_context, caplog):
|
||||
"""Test error handling with transformation."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
with pytest.raises(McpError) as exc_info:
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert exc_info.value.error.code == -32602
|
||||
assert "Invalid params: test error" in exc_info.value.error.message
|
||||
assert "Error in test_method: ValueError: test error" in caplog.text
|
||||
|
||||
def test_get_error_stats(self, mock_context):
|
||||
"""Test getting error statistics."""
|
||||
middleware = ErrorHandlingMiddleware()
|
||||
error1 = ValueError("error1")
|
||||
error2 = ValueError("error2")
|
||||
error3 = RuntimeError("error3")
|
||||
|
||||
middleware._log_error(error1, mock_context)
|
||||
middleware._log_error(error2, mock_context)
|
||||
middleware._log_error(error3, mock_context)
|
||||
|
||||
stats = middleware.get_error_stats()
|
||||
assert stats["ValueError:test_method"] == 2
|
||||
assert stats["RuntimeError:test_method"] == 1
|
||||
|
||||
|
||||
class TestRetryMiddleware:
|
||||
"""Test retry middleware functionality."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = RetryMiddleware()
|
||||
assert middleware.max_retries == 3
|
||||
assert middleware.base_delay == 1.0
|
||||
assert middleware.max_delay == 60.0
|
||||
assert middleware.backoff_multiplier == 2.0
|
||||
assert middleware.retry_exceptions == (ConnectionError, TimeoutError)
|
||||
assert middleware.logger.name == "fastmcp.retry"
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
logger = logging.getLogger("custom")
|
||||
middleware = RetryMiddleware(
|
||||
max_retries=5,
|
||||
base_delay=2.0,
|
||||
max_delay=120.0,
|
||||
backoff_multiplier=3.0,
|
||||
retry_exceptions=(ValueError, RuntimeError),
|
||||
logger=logger,
|
||||
)
|
||||
assert middleware.max_retries == 5
|
||||
assert middleware.base_delay == 2.0
|
||||
assert middleware.max_delay == 120.0
|
||||
assert middleware.backoff_multiplier == 3.0
|
||||
assert middleware.retry_exceptions == (ValueError, RuntimeError)
|
||||
assert middleware.logger is logger
|
||||
|
||||
def test_should_retry_true(self):
|
||||
"""Test retry decision for retryable errors."""
|
||||
middleware = RetryMiddleware()
|
||||
|
||||
assert middleware._should_retry(ConnectionError()) is True
|
||||
assert middleware._should_retry(TimeoutError()) is True
|
||||
|
||||
def test_should_retry_false(self):
|
||||
"""Test retry decision for non-retryable errors."""
|
||||
middleware = RetryMiddleware()
|
||||
|
||||
assert middleware._should_retry(ValueError()) is False
|
||||
assert middleware._should_retry(RuntimeError()) is False
|
||||
|
||||
def test_calculate_delay(self):
|
||||
"""Test delay calculation."""
|
||||
middleware = RetryMiddleware(
|
||||
base_delay=1.0, backoff_multiplier=2.0, max_delay=10.0
|
||||
)
|
||||
|
||||
assert middleware._calculate_delay(0) == 1.0
|
||||
assert middleware._calculate_delay(1) == 2.0
|
||||
assert middleware._calculate_delay(2) == 4.0
|
||||
assert middleware._calculate_delay(3) == 8.0
|
||||
assert middleware._calculate_delay(4) == 10.0 # capped at max_delay
|
||||
|
||||
async def test_on_request_success_first_try(self, mock_context, mock_call_next):
|
||||
"""Test successful request on first try."""
|
||||
middleware = RetryMiddleware()
|
||||
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.call_count == 1
|
||||
|
||||
async def test_on_request_success_after_retries(self, mock_context, caplog):
|
||||
"""Test successful request after retries."""
|
||||
middleware = RetryMiddleware(base_delay=0.01) # Fast retry for testing
|
||||
|
||||
# Fail first two attempts, succeed on third
|
||||
mock_call_next = AsyncMock(
|
||||
side_effect=[
|
||||
ConnectionError("connection failed"),
|
||||
ConnectionError("connection failed"),
|
||||
"test_result",
|
||||
]
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.call_count == 3
|
||||
assert "Retrying in" in caplog.text
|
||||
|
||||
async def test_on_request_max_retries_exceeded(self, mock_context, caplog):
|
||||
"""Test request failing after max retries."""
|
||||
middleware = RetryMiddleware(max_retries=2, base_delay=0.01)
|
||||
|
||||
# Fail all attempts
|
||||
mock_call_next = AsyncMock(side_effect=ConnectionError("connection failed"))
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
with pytest.raises(ConnectionError):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert mock_call_next.call_count == 3 # initial + 2 retries
|
||||
assert "Retrying in" in caplog.text
|
||||
|
||||
async def test_on_request_non_retryable_error(self, mock_context):
|
||||
"""Test non-retryable error is not retried."""
|
||||
middleware = RetryMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("non-retryable"))
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert mock_call_next.call_count == 1 # No retries
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def error_handling_server():
|
||||
"""Create a FastMCP server specifically for error handling middleware tests."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("ErrorHandlingTestServer")
|
||||
|
||||
@mcp.tool
|
||||
def reliable_operation(data: str) -> str:
|
||||
"""A reliable operation that always succeeds."""
|
||||
return f"Success: {data}"
|
||||
|
||||
@mcp.tool
|
||||
def failing_operation(error_type: str = "value") -> str:
|
||||
"""An operation that fails with different error types."""
|
||||
if error_type == "value":
|
||||
raise ValueError("Value error occurred")
|
||||
elif error_type == "file":
|
||||
raise FileNotFoundError("File not found")
|
||||
elif error_type == "permission":
|
||||
raise PermissionError("Permission denied")
|
||||
elif error_type == "timeout":
|
||||
raise TimeoutError("Operation timed out")
|
||||
elif error_type == "generic":
|
||||
raise RuntimeError("Generic runtime error")
|
||||
else:
|
||||
return "Operation completed"
|
||||
|
||||
@mcp.tool
|
||||
def intermittent_operation(fail_rate: float = 0.5) -> str:
|
||||
"""An operation that fails intermittently."""
|
||||
import random
|
||||
|
||||
if random.random() < fail_rate:
|
||||
raise ConnectionError("Random connection failure")
|
||||
return "Operation succeeded"
|
||||
|
||||
@mcp.tool
|
||||
def retryable_operation(attempt_count: int = 0) -> str:
|
||||
"""An operation that succeeds after a few attempts."""
|
||||
# This is a simple way to simulate retry behavior
|
||||
# In a real scenario, you might use external state
|
||||
if attempt_count < 2:
|
||||
raise ConnectionError("Temporary connection error")
|
||||
return "Operation succeeded after retries"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
class TestErrorHandlingMiddlewareIntegration:
|
||||
"""Integration tests for error handling middleware with real FastMCP server."""
|
||||
|
||||
async def test_error_handling_middleware_logs_real_errors(
|
||||
self, error_handling_server, caplog
|
||||
):
|
||||
"""Test that error handling middleware logs real errors from tools."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Test different types of errors
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "file"})
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have error logs for both failures
|
||||
assert "Error in tools/call: ToolError:" in log_text
|
||||
# Should have captured both error instances
|
||||
error_count = log_text.count("Error in tools/call:")
|
||||
assert error_count == 2
|
||||
|
||||
async def test_error_handling_middleware_tracks_error_statistics(
|
||||
self, error_handling_server
|
||||
):
|
||||
"""Test that error handling middleware accurately tracks error statistics."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
error_middleware = ErrorHandlingMiddleware()
|
||||
error_handling_server.add_middleware(error_middleware)
|
||||
|
||||
async with Client(error_handling_server) as client:
|
||||
# Generate different types of errors
|
||||
for _ in range(3):
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
for _ in range(2):
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "file"})
|
||||
|
||||
# Try some intermittent operations (some may succeed)
|
||||
for _ in range(5):
|
||||
try:
|
||||
await client.call_tool("intermittent_operation", {"fail_rate": 0.8})
|
||||
except Exception:
|
||||
pass # Expected failures
|
||||
|
||||
# Check error statistics
|
||||
stats = error_middleware.get_error_stats()
|
||||
|
||||
# Should have tracked the ToolError wrapper
|
||||
assert "ToolError:tools/call" in stats
|
||||
assert stats["ToolError:tools/call"] >= 5 # At least the 5 deliberate failures
|
||||
|
||||
async def test_error_handling_middleware_with_success_and_failure(
|
||||
self, error_handling_server, caplog
|
||||
):
|
||||
"""Test error handling middleware with mix of successful and failed operations."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Successful operation (should not generate error logs)
|
||||
await client.call_tool("reliable_operation", {"data": "test"})
|
||||
|
||||
# Failed operation (should generate error log)
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
# Another successful operation
|
||||
await client.call_tool("reliable_operation", {"data": "test2"})
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should only have one error log (for the failed operation)
|
||||
error_count = log_text.count("Error in tools/call:")
|
||||
assert error_count == 1
|
||||
|
||||
async def test_error_handling_middleware_custom_callback(
|
||||
self, error_handling_server
|
||||
):
|
||||
"""Test error handling middleware with custom error callback."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
captured_errors = []
|
||||
|
||||
def error_callback(error, context):
|
||||
captured_errors.append(
|
||||
{
|
||||
"error_type": type(error).__name__,
|
||||
"message": str(error),
|
||||
"method": context.method,
|
||||
}
|
||||
)
|
||||
|
||||
error_handling_server.add_middleware(
|
||||
ErrorHandlingMiddleware(error_callback=error_callback)
|
||||
)
|
||||
|
||||
async with Client(error_handling_server) as client:
|
||||
# Generate some errors
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "timeout"})
|
||||
|
||||
# Check that callback was called
|
||||
assert len(captured_errors) == 2
|
||||
assert captured_errors[0]["error_type"] == "ToolError"
|
||||
assert captured_errors[1]["error_type"] == "ToolError"
|
||||
assert all(error["method"] == "tools/call" for error in captured_errors)
|
||||
|
||||
async def test_error_handling_middleware_transform_errors(
|
||||
self, error_handling_server
|
||||
):
|
||||
"""Test error transformation functionality."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
error_handling_server.add_middleware(
|
||||
ErrorHandlingMiddleware(transform_errors=True)
|
||||
)
|
||||
|
||||
async with Client(error_handling_server) as client:
|
||||
# All errors should still be raised, but potentially transformed
|
||||
with pytest.raises(Exception) as exc_info:
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
# Error should still exist (may be wrapped by FastMCP)
|
||||
assert exc_info.value is not None
|
||||
|
||||
|
||||
class TestRetryMiddlewareIntegration:
|
||||
"""Integration tests for retry middleware with real FastMCP server."""
|
||||
|
||||
async def test_retry_middleware_with_transient_failures(
|
||||
self, error_handling_server, caplog
|
||||
):
|
||||
"""Test retry middleware with operations that have transient failures."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
# Configure retry middleware to retry connection errors
|
||||
error_handling_server.add_middleware(
|
||||
RetryMiddleware(
|
||||
max_retries=3,
|
||||
base_delay=0.01, # Very short delay for testing
|
||||
retry_exceptions=(ConnectionError,),
|
||||
)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.WARNING):
|
||||
async with Client(error_handling_server) as client:
|
||||
# This operation fails intermittently - try several times
|
||||
success_count = 0
|
||||
for _ in range(5):
|
||||
try:
|
||||
await client.call_tool(
|
||||
"intermittent_operation", {"fail_rate": 0.7}
|
||||
)
|
||||
success_count += 1
|
||||
except Exception:
|
||||
pass # Some failures expected even with retries
|
||||
|
||||
# Should have some retry log messages
|
||||
# Note: Retry logs might not appear if the underlying errors are wrapped by FastMCP
|
||||
# The key is that some operations should succeed due to retries
|
||||
|
||||
async def test_retry_middleware_with_permanent_failures(
|
||||
self, error_handling_server
|
||||
):
|
||||
"""Test that retry middleware doesn't retry non-retryable errors."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
# Configure retry middleware for connection errors only
|
||||
error_handling_server.add_middleware(
|
||||
RetryMiddleware(
|
||||
max_retries=3, base_delay=0.01, retry_exceptions=(ConnectionError,)
|
||||
)
|
||||
)
|
||||
|
||||
async with Client(error_handling_server) as client:
|
||||
# Value errors should not be retried
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
# Should fail immediately without retries
|
||||
|
||||
async def test_combined_error_handling_and_retry_middleware(
|
||||
self, error_handling_server, caplog
|
||||
):
|
||||
"""Test error handling and retry middleware working together."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
# Add both middleware
|
||||
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
||||
error_handling_server.add_middleware(
|
||||
RetryMiddleware(
|
||||
max_retries=2, base_delay=0.01, retry_exceptions=(ConnectionError,)
|
||||
)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.ERROR):
|
||||
async with Client(error_handling_server) as client:
|
||||
# Try intermittent operation
|
||||
try:
|
||||
await client.call_tool("intermittent_operation", {"fail_rate": 0.9})
|
||||
except Exception:
|
||||
pass # May still fail even with retries
|
||||
|
||||
# Try permanent failure
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_operation", {"error_type": "value"})
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have error logs from error handling middleware
|
||||
assert "Error in tools/call:" in log_text
|
||||
426
tests/server/middleware/test_logging.py
Normal file
426
tests/server/middleware/test_logging.py
Normal file
|
|
@ -0,0 +1,426 @@
|
|||
"""Tests for logging middleware."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from fastmcp.server.middleware.logging import (
|
||||
LoggingMiddleware,
|
||||
StructuredLoggingMiddleware,
|
||||
)
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock middleware context."""
|
||||
context = MagicMock(spec=MiddlewareContext)
|
||||
context.method = "test_method"
|
||||
context.source = "client"
|
||||
context.type = "request"
|
||||
context.message = MagicMock()
|
||||
context.message.__dict__ = {"param": "value"}
|
||||
context.timestamp = MagicMock()
|
||||
context.timestamp.isoformat.return_value = "2023-01-01T00:00:00Z"
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_call_next():
|
||||
"""Create a mock call_next function."""
|
||||
return AsyncMock(return_value="test_result")
|
||||
|
||||
|
||||
class TestLoggingMiddleware:
|
||||
"""Test logging middleware functionality."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = LoggingMiddleware()
|
||||
assert middleware.logger.name == "fastmcp.requests"
|
||||
assert middleware.log_level == logging.INFO
|
||||
assert middleware.include_payloads is False
|
||||
assert middleware.max_payload_length == 1000
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
logger = logging.getLogger("custom")
|
||||
middleware = LoggingMiddleware(
|
||||
logger=logger,
|
||||
log_level=logging.DEBUG,
|
||||
include_payloads=True,
|
||||
max_payload_length=500,
|
||||
)
|
||||
assert middleware.logger is logger
|
||||
assert middleware.log_level == logging.DEBUG
|
||||
assert middleware.include_payloads is True
|
||||
assert middleware.max_payload_length == 500
|
||||
|
||||
def test_format_message_without_payloads(self, mock_context):
|
||||
"""Test message formatting without payloads."""
|
||||
middleware = LoggingMiddleware()
|
||||
formatted = middleware._format_message(mock_context)
|
||||
|
||||
assert "source=client" in formatted
|
||||
assert "type=request" in formatted
|
||||
assert "method=test_method" in formatted
|
||||
assert "payload=" not in formatted
|
||||
|
||||
def test_format_message_with_payloads(self, mock_context):
|
||||
"""Test message formatting with payloads."""
|
||||
middleware = LoggingMiddleware(include_payloads=True)
|
||||
formatted = middleware._format_message(mock_context)
|
||||
|
||||
assert "source=client" in formatted
|
||||
assert "type=request" in formatted
|
||||
assert "method=test_method" in formatted
|
||||
assert 'payload={"param": "value"}' in formatted
|
||||
|
||||
def test_format_message_long_payload(self, mock_context):
|
||||
"""Test message formatting with long payload truncation."""
|
||||
middleware = LoggingMiddleware(include_payloads=True, max_payload_length=10)
|
||||
formatted = middleware._format_message(mock_context)
|
||||
|
||||
assert "payload=" in formatted
|
||||
assert "..." in formatted
|
||||
|
||||
async def test_on_message_success(self, mock_context, mock_call_next, caplog):
|
||||
"""Test logging successful messages."""
|
||||
middleware = LoggingMiddleware()
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.called
|
||||
assert "Processing message:" in caplog.text
|
||||
assert "Completed message: test_method" in caplog.text
|
||||
|
||||
async def test_on_message_failure(self, mock_context, caplog):
|
||||
"""Test logging failed messages."""
|
||||
middleware = LoggingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert "Processing message:" in caplog.text
|
||||
assert "Failed message: test_method - test error" in caplog.text
|
||||
|
||||
|
||||
class TestStructuredLoggingMiddleware:
|
||||
"""Test structured logging middleware functionality."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
assert middleware.logger.name == "fastmcp.structured"
|
||||
assert middleware.log_level == logging.INFO
|
||||
assert middleware.include_payloads is False
|
||||
|
||||
def test_create_log_entry_basic(self, mock_context):
|
||||
"""Test creating basic log entry."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
entry = middleware._create_log_entry(mock_context, "test_event")
|
||||
|
||||
assert entry["event"] == "test_event"
|
||||
assert entry["timestamp"] == "2023-01-01T00:00:00Z"
|
||||
assert entry["source"] == "client"
|
||||
assert entry["type"] == "request"
|
||||
assert entry["method"] == "test_method"
|
||||
assert "payload" not in entry
|
||||
|
||||
def test_create_log_entry_with_payload(self, mock_context):
|
||||
"""Test creating log entry with payload."""
|
||||
middleware = StructuredLoggingMiddleware(include_payloads=True)
|
||||
entry = middleware._create_log_entry(mock_context, "test_event")
|
||||
|
||||
assert entry["payload"] == {"param": "value"}
|
||||
|
||||
def test_create_log_entry_with_extra_fields(self, mock_context):
|
||||
"""Test creating log entry with extra fields."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
entry = middleware._create_log_entry(
|
||||
mock_context, "test_event", extra_field="extra_value"
|
||||
)
|
||||
|
||||
assert entry["extra_field"] == "extra_value"
|
||||
|
||||
async def test_on_message_success(self, mock_context, mock_call_next, caplog):
|
||||
"""Test structured logging of successful messages."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
|
||||
# Check that we have structured JSON logs
|
||||
log_lines = [record.message for record in caplog.records]
|
||||
assert len(log_lines) == 2 # start and success entries
|
||||
|
||||
start_entry = json.loads(log_lines[0])
|
||||
assert start_entry["event"] == "request_start"
|
||||
assert start_entry["method"] == "test_method"
|
||||
|
||||
success_entry = json.loads(log_lines[1])
|
||||
assert success_entry["event"] == "request_success"
|
||||
assert success_entry["result_type"] == "str"
|
||||
|
||||
async def test_on_message_failure(self, mock_context, caplog):
|
||||
"""Test structured logging of failed messages."""
|
||||
middleware = StructuredLoggingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_message(mock_context, mock_call_next)
|
||||
|
||||
# Check that we have structured JSON logs
|
||||
log_lines = [record.message for record in caplog.records]
|
||||
assert len(log_lines) == 2 # start and error entries
|
||||
|
||||
start_entry = json.loads(log_lines[0])
|
||||
assert start_entry["event"] == "request_start"
|
||||
|
||||
error_entry = json.loads(log_lines[1])
|
||||
assert error_entry["event"] == "request_error"
|
||||
assert error_entry["error_type"] == "ValueError"
|
||||
assert error_entry["error_message"] == "test error"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def logging_server():
|
||||
"""Create a FastMCP server specifically for logging middleware tests."""
|
||||
from fastmcp import FastMCP
|
||||
|
||||
mcp = FastMCP("LoggingTestServer")
|
||||
|
||||
@mcp.tool
|
||||
def simple_operation(data: str) -> str:
|
||||
"""A simple operation for testing logging."""
|
||||
return f"Processed: {data}"
|
||||
|
||||
@mcp.tool
|
||||
def complex_operation(items: list[str], mode: str = "default") -> dict:
|
||||
"""A complex operation with structured data."""
|
||||
return {"processed_items": len(items), "mode": mode, "result": "success"}
|
||||
|
||||
@mcp.tool
|
||||
def operation_with_error(should_fail: bool = False) -> str:
|
||||
"""An operation that can be made to fail."""
|
||||
if should_fail:
|
||||
raise ValueError("Operation failed intentionally")
|
||||
return "Operation completed successfully"
|
||||
|
||||
@mcp.resource("log://test")
|
||||
def test_resource() -> str:
|
||||
"""A test resource for logging."""
|
||||
return "Test resource content"
|
||||
|
||||
@mcp.prompt
|
||||
def test_prompt() -> str:
|
||||
"""A test prompt for logging."""
|
||||
return "Test prompt content"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
class TestLoggingMiddlewareIntegration:
|
||||
"""Integration tests for logging middleware with real FastMCP server."""
|
||||
|
||||
async def test_logging_middleware_logs_successful_operations(
|
||||
self, logging_server, caplog
|
||||
):
|
||||
"""Test that logging middleware captures successful operations."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
logging_server.add_middleware(LoggingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool("simple_operation", {"data": "test_data"})
|
||||
await client.call_tool(
|
||||
"complex_operation", {"items": ["a", "b", "c"], "mode": "batch"}
|
||||
)
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have processing and completion logs for both operations
|
||||
assert "Processing message:" in log_text
|
||||
assert "Completed message: tools/call" in log_text
|
||||
|
||||
# Should have captured both tool calls
|
||||
processing_count = log_text.count("Processing message:")
|
||||
completion_count = log_text.count("Completed message:")
|
||||
assert processing_count == 2
|
||||
assert completion_count == 2
|
||||
|
||||
async def test_logging_middleware_logs_failures(self, logging_server, caplog):
|
||||
"""Test that logging middleware captures failed operations."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
logging_server.add_middleware(LoggingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
# This should fail and be logged
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"operation_with_error", {"should_fail": True}
|
||||
)
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have processing and failure logs
|
||||
assert "Processing message:" in log_text
|
||||
assert "Failed message: tools/call" in log_text
|
||||
|
||||
async def test_logging_middleware_with_payloads(self, logging_server, caplog):
|
||||
"""Test logging middleware when configured to include payloads."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
logging_server.add_middleware(
|
||||
LoggingMiddleware(include_payloads=True, max_payload_length=500)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool("simple_operation", {"data": "payload_test"})
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should include payload information
|
||||
assert "Processing message:" in log_text
|
||||
assert "payload=" in log_text
|
||||
|
||||
async def test_structured_logging_middleware_produces_json(
|
||||
self, logging_server, caplog
|
||||
):
|
||||
"""Test that structured logging middleware produces parseable JSON logs."""
|
||||
import json
|
||||
|
||||
from fastmcp.client import Client
|
||||
|
||||
logging_server.add_middleware(
|
||||
StructuredLoggingMiddleware(include_payloads=True)
|
||||
)
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool("simple_operation", {"data": "json_test"})
|
||||
|
||||
# Extract JSON log entries
|
||||
log_lines = [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.name == "fastmcp.structured"
|
||||
]
|
||||
|
||||
assert len(log_lines) >= 2 # Should have start and success entries
|
||||
|
||||
# Each log line should be valid JSON
|
||||
for line in log_lines:
|
||||
log_entry = json.loads(line)
|
||||
assert "event" in log_entry
|
||||
assert "timestamp" in log_entry
|
||||
assert "source" in log_entry
|
||||
assert "type" in log_entry
|
||||
assert "method" in log_entry
|
||||
|
||||
async def test_structured_logging_middleware_handles_errors(
|
||||
self, logging_server, caplog
|
||||
):
|
||||
"""Test structured logging of errors with JSON format."""
|
||||
import json
|
||||
|
||||
from fastmcp.client import Client
|
||||
|
||||
logging_server.add_middleware(StructuredLoggingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool(
|
||||
"operation_with_error", {"should_fail": True}
|
||||
)
|
||||
|
||||
# Extract JSON log entries
|
||||
log_lines = [
|
||||
record.message
|
||||
for record in caplog.records
|
||||
if record.name == "fastmcp.structured"
|
||||
]
|
||||
|
||||
# Should have start and error entries
|
||||
assert len(log_lines) >= 2
|
||||
|
||||
# Find the error entry
|
||||
error_entries = []
|
||||
for line in log_lines:
|
||||
log_entry = json.loads(line)
|
||||
if log_entry.get("event") == "request_error":
|
||||
error_entries.append(log_entry)
|
||||
|
||||
assert len(error_entries) == 1
|
||||
error_entry = error_entries[0]
|
||||
assert "error_type" in error_entry
|
||||
assert "error_message" in error_entry
|
||||
|
||||
async def test_logging_middleware_with_different_operations(
|
||||
self, logging_server, caplog
|
||||
):
|
||||
"""Test logging middleware with various MCP operations."""
|
||||
from fastmcp.client import Client
|
||||
|
||||
logging_server.add_middleware(LoggingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(logging_server) as client:
|
||||
# Test different operation types
|
||||
await client.call_tool("simple_operation", {"data": "test"})
|
||||
await client.read_resource("log://test")
|
||||
await client.get_prompt("test_prompt")
|
||||
await client.list_tools()
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have logs for all different operation types
|
||||
# Note: Different operations may have different method names
|
||||
processing_count = log_text.count("Processing message:")
|
||||
completion_count = log_text.count("Completed message:")
|
||||
|
||||
# Should have processed all 4 operations
|
||||
assert processing_count == 4
|
||||
assert completion_count == 4
|
||||
|
||||
async def test_logging_middleware_custom_configuration(self, logging_server):
|
||||
"""Test logging middleware with custom logger configuration."""
|
||||
import io
|
||||
import logging
|
||||
|
||||
from fastmcp.client import Client
|
||||
|
||||
# Create custom logger
|
||||
log_buffer = io.StringIO()
|
||||
handler = logging.StreamHandler(log_buffer)
|
||||
custom_logger = logging.getLogger("custom_logging_test")
|
||||
custom_logger.addHandler(handler)
|
||||
custom_logger.setLevel(logging.DEBUG)
|
||||
|
||||
logging_server.add_middleware(
|
||||
LoggingMiddleware(
|
||||
logger=custom_logger, log_level=logging.DEBUG, include_payloads=True
|
||||
)
|
||||
)
|
||||
|
||||
async with Client(logging_server) as client:
|
||||
await client.call_tool("simple_operation", {"data": "custom_test"})
|
||||
|
||||
# Check that our custom logger captured the logs
|
||||
log_output = log_buffer.getvalue()
|
||||
assert "Processing message:" in log_output
|
||||
assert "payload=" in log_output
|
||||
448
tests/server/middleware/test_rate_limiting.py
Normal file
448
tests/server/middleware/test_rate_limiting.py
Normal file
|
|
@ -0,0 +1,448 @@
|
|||
"""Tests for rate limiting middleware."""
|
||||
|
||||
import asyncio
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.exceptions import ToolError
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
from fastmcp.server.middleware.rate_limiting import (
|
||||
RateLimitError,
|
||||
RateLimitingMiddleware,
|
||||
SlidingWindowRateLimiter,
|
||||
SlidingWindowRateLimitingMiddleware,
|
||||
TokenBucketRateLimiter,
|
||||
)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock middleware context."""
|
||||
context = MagicMock(spec=MiddlewareContext)
|
||||
context.method = "test_method"
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_call_next():
|
||||
"""Create a mock call_next function."""
|
||||
return AsyncMock(return_value="test_result")
|
||||
|
||||
|
||||
class TestTokenBucketRateLimiter:
|
||||
"""Test token bucket rate limiter."""
|
||||
|
||||
def test_init(self):
|
||||
"""Test initialization."""
|
||||
limiter = TokenBucketRateLimiter(capacity=10, refill_rate=5.0)
|
||||
assert limiter.capacity == 10
|
||||
assert limiter.refill_rate == 5.0
|
||||
assert limiter.tokens == 10
|
||||
|
||||
async def test_consume_success(self):
|
||||
"""Test successful token consumption."""
|
||||
limiter = TokenBucketRateLimiter(capacity=10, refill_rate=5.0)
|
||||
|
||||
# Should be able to consume tokens initially
|
||||
assert await limiter.consume(5) is True
|
||||
assert await limiter.consume(3) is True
|
||||
|
||||
async def test_consume_failure(self):
|
||||
"""Test failed token consumption."""
|
||||
limiter = TokenBucketRateLimiter(capacity=5, refill_rate=1.0)
|
||||
|
||||
# Consume all tokens
|
||||
assert await limiter.consume(5) is True
|
||||
|
||||
# Should fail to consume more
|
||||
assert await limiter.consume(1) is False
|
||||
|
||||
async def test_refill(self):
|
||||
"""Test token refill over time."""
|
||||
limiter = TokenBucketRateLimiter(
|
||||
capacity=10, refill_rate=10.0
|
||||
) # 10 tokens per second
|
||||
|
||||
# Consume all tokens
|
||||
assert await limiter.consume(10) is True
|
||||
assert await limiter.consume(1) is False
|
||||
|
||||
# Wait for refill (0.2 seconds = 2 tokens at 10/sec)
|
||||
await asyncio.sleep(0.2)
|
||||
assert await limiter.consume(2) is True
|
||||
|
||||
|
||||
class TestSlidingWindowRateLimiter:
|
||||
"""Test sliding window rate limiter."""
|
||||
|
||||
def test_init(self):
|
||||
"""Test initialization."""
|
||||
limiter = SlidingWindowRateLimiter(max_requests=10, window_seconds=60)
|
||||
assert limiter.max_requests == 10
|
||||
assert limiter.window_seconds == 60
|
||||
assert len(limiter.requests) == 0
|
||||
|
||||
async def test_is_allowed_success(self):
|
||||
"""Test allowing requests within limit."""
|
||||
limiter = SlidingWindowRateLimiter(max_requests=3, window_seconds=60)
|
||||
|
||||
# Should allow requests up to the limit
|
||||
assert await limiter.is_allowed() is True
|
||||
assert await limiter.is_allowed() is True
|
||||
assert await limiter.is_allowed() is True
|
||||
|
||||
async def test_is_allowed_failure(self):
|
||||
"""Test rejecting requests over limit."""
|
||||
limiter = SlidingWindowRateLimiter(max_requests=2, window_seconds=60)
|
||||
|
||||
# Should allow up to limit
|
||||
assert await limiter.is_allowed() is True
|
||||
assert await limiter.is_allowed() is True
|
||||
|
||||
# Should reject over limit
|
||||
assert await limiter.is_allowed() is False
|
||||
|
||||
async def test_sliding_window(self):
|
||||
"""Test sliding window behavior."""
|
||||
limiter = SlidingWindowRateLimiter(max_requests=2, window_seconds=1)
|
||||
|
||||
# Use up requests
|
||||
assert await limiter.is_allowed() is True
|
||||
assert await limiter.is_allowed() is True
|
||||
assert await limiter.is_allowed() is False
|
||||
|
||||
# Wait for window to pass
|
||||
await asyncio.sleep(1.1)
|
||||
|
||||
# Should be able to make requests again
|
||||
assert await limiter.is_allowed() is True
|
||||
|
||||
|
||||
class TestRateLimitingMiddleware:
|
||||
"""Test rate limiting middleware."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = RateLimitingMiddleware()
|
||||
assert middleware.max_requests_per_second == 10.0
|
||||
assert middleware.burst_capacity == 20
|
||||
assert middleware.get_client_id is None
|
||||
assert middleware.global_limit is False
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
|
||||
def get_client_id(ctx):
|
||||
return "test_client"
|
||||
|
||||
middleware = RateLimitingMiddleware(
|
||||
max_requests_per_second=5.0,
|
||||
burst_capacity=10,
|
||||
get_client_id=get_client_id,
|
||||
global_limit=True,
|
||||
)
|
||||
assert middleware.max_requests_per_second == 5.0
|
||||
assert middleware.burst_capacity == 10
|
||||
assert middleware.get_client_id is get_client_id
|
||||
assert middleware.global_limit is True
|
||||
|
||||
def test_get_client_identifier_default(self, mock_context):
|
||||
"""Test default client identifier."""
|
||||
middleware = RateLimitingMiddleware()
|
||||
assert middleware._get_client_identifier(mock_context) == "global"
|
||||
|
||||
def test_get_client_identifier_custom(self, mock_context):
|
||||
"""Test custom client identifier."""
|
||||
|
||||
def get_client_id(ctx):
|
||||
return "custom_client"
|
||||
|
||||
middleware = RateLimitingMiddleware(get_client_id=get_client_id)
|
||||
assert middleware._get_client_identifier(mock_context) == "custom_client"
|
||||
|
||||
async def test_on_request_success(self, mock_context, mock_call_next):
|
||||
"""Test successful request within rate limit."""
|
||||
middleware = RateLimitingMiddleware(max_requests_per_second=100.0) # High limit
|
||||
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.called
|
||||
|
||||
async def test_on_request_rate_limited(self, mock_context, mock_call_next):
|
||||
"""Test request rejection due to rate limiting."""
|
||||
middleware = RateLimitingMiddleware(
|
||||
max_requests_per_second=1.0, burst_capacity=1
|
||||
)
|
||||
|
||||
# First request should succeed
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
# Second request should be rate limited
|
||||
with pytest.raises(RateLimitError, match="Rate limit exceeded"):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
async def test_global_rate_limiting(self, mock_context, mock_call_next):
|
||||
"""Test global rate limiting."""
|
||||
middleware = RateLimitingMiddleware(
|
||||
max_requests_per_second=1.0, burst_capacity=1, global_limit=True
|
||||
)
|
||||
|
||||
# First request should succeed
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
# Second request should be rate limited
|
||||
with pytest.raises(RateLimitError, match="Global rate limit exceeded"):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
|
||||
class TestSlidingWindowRateLimitingMiddleware:
|
||||
"""Test sliding window rate limiting middleware."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = SlidingWindowRateLimitingMiddleware(max_requests=100)
|
||||
assert middleware.max_requests == 100
|
||||
assert middleware.window_seconds == 60
|
||||
assert middleware.get_client_id is None
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
|
||||
def get_client_id(ctx):
|
||||
return "test_client"
|
||||
|
||||
middleware = SlidingWindowRateLimitingMiddleware(
|
||||
max_requests=50, window_minutes=5, get_client_id=get_client_id
|
||||
)
|
||||
assert middleware.max_requests == 50
|
||||
assert middleware.window_seconds == 300 # 5 minutes
|
||||
assert middleware.get_client_id is get_client_id
|
||||
|
||||
async def test_on_request_success(self, mock_context, mock_call_next):
|
||||
"""Test successful request within rate limit."""
|
||||
middleware = SlidingWindowRateLimitingMiddleware(max_requests=100)
|
||||
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.called
|
||||
|
||||
async def test_on_request_rate_limited(self, mock_context, mock_call_next):
|
||||
"""Test request rejection due to rate limiting."""
|
||||
middleware = SlidingWindowRateLimitingMiddleware(max_requests=1)
|
||||
|
||||
# First request should succeed
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
# Second request should be rate limited
|
||||
with pytest.raises(RateLimitError, match="Rate limit exceeded"):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
|
||||
class TestRateLimitError:
|
||||
"""Test rate limit error."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
error = RateLimitError()
|
||||
assert error.error.code == -32000
|
||||
assert error.error.message == "Rate limit exceeded"
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
error = RateLimitError("Custom message")
|
||||
assert error.error.code == -32000
|
||||
assert error.error.message == "Custom message"
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def rate_limit_server():
|
||||
"""Create a FastMCP server specifically for rate limiting tests."""
|
||||
mcp = FastMCP("RateLimitTestServer")
|
||||
|
||||
@mcp.tool
|
||||
def quick_action(message: str) -> str:
|
||||
"""A quick action for testing rate limits."""
|
||||
return f"Processed: {message}"
|
||||
|
||||
@mcp.tool
|
||||
def batch_process(items: list[str]) -> str:
|
||||
"""Process multiple items."""
|
||||
return f"Processed {len(items)} items"
|
||||
|
||||
@mcp.tool
|
||||
def heavy_computation() -> str:
|
||||
"""A heavy computation that might need rate limiting."""
|
||||
# Simulate some work
|
||||
import time
|
||||
|
||||
time.sleep(0.01) # Very short delay
|
||||
return "Heavy computation complete"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
class TestRateLimitingMiddlewareIntegration:
|
||||
"""Integration tests for rate limiting middleware with real FastMCP server."""
|
||||
|
||||
async def test_rate_limiting_allows_normal_usage(self, rate_limit_server):
|
||||
"""Test that normal usage patterns are allowed through rate limiting."""
|
||||
# Generous rate limit
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(max_requests_per_second=50.0, burst_capacity=10)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# Normal usage should be fine
|
||||
for i in range(5):
|
||||
result = await client.call_tool(
|
||||
"quick_action", {"message": f"task_{i}"}
|
||||
)
|
||||
assert f"Processed: task_{i}" in str(result)
|
||||
|
||||
async def test_rate_limiting_blocks_rapid_requests(self, rate_limit_server):
|
||||
"""Test that rate limiting blocks rapid successive requests."""
|
||||
# Very restrictive rate limit
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(max_requests_per_second=2.0, burst_capacity=3)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# First few should succeed (within burst capacity)
|
||||
await client.call_tool("quick_action", {"message": "1"})
|
||||
await client.call_tool("quick_action", {"message": "2"})
|
||||
await client.call_tool("quick_action", {"message": "3"})
|
||||
|
||||
# Next should be rate limited
|
||||
with pytest.raises(ToolError, match="Rate limit exceeded"):
|
||||
await client.call_tool("quick_action", {"message": "4"})
|
||||
|
||||
async def test_rate_limiting_with_concurrent_requests(self, rate_limit_server):
|
||||
"""Test rate limiting behavior with concurrent requests."""
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(max_requests_per_second=5.0, burst_capacity=3)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# Fire off many concurrent requests
|
||||
tasks = []
|
||||
for i in range(8):
|
||||
task = asyncio.create_task(
|
||||
client.call_tool("quick_action", {"message": f"concurrent_{i}"})
|
||||
)
|
||||
tasks.append(task)
|
||||
|
||||
# Gather results, allowing exceptions
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
# Some should succeed, some should be rate limited
|
||||
successes = [r for r in results if not isinstance(r, Exception)]
|
||||
failures = [r for r in results if isinstance(r, ToolError)]
|
||||
|
||||
assert len(successes) > 0, "Some requests should succeed"
|
||||
assert len(failures) > 0, "Some requests should be rate limited"
|
||||
assert len(successes) + len(failures) == 8
|
||||
|
||||
async def test_sliding_window_rate_limiting(self, rate_limit_server):
|
||||
"""Test sliding window rate limiting implementation."""
|
||||
rate_limit_server.add_middleware(
|
||||
SlidingWindowRateLimitingMiddleware(
|
||||
max_requests=3,
|
||||
window_minutes=1, # 1 minute window
|
||||
)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# Should allow up to the limit
|
||||
await client.call_tool("quick_action", {"message": "1"})
|
||||
await client.call_tool("quick_action", {"message": "2"})
|
||||
await client.call_tool("quick_action", {"message": "3"})
|
||||
|
||||
# Fourth should be blocked
|
||||
with pytest.raises(ToolError, match="Rate limit exceeded"):
|
||||
await client.call_tool("quick_action", {"message": "4"})
|
||||
|
||||
async def test_rate_limiting_with_different_operations(self, rate_limit_server):
|
||||
"""Test that rate limiting applies to all types of operations."""
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(max_requests_per_second=3.0, burst_capacity=2)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# Mix different operations
|
||||
await client.call_tool("quick_action", {"message": "test"})
|
||||
await client.call_tool("heavy_computation")
|
||||
|
||||
# Should be rate limited regardless of operation type
|
||||
with pytest.raises(ToolError, match="Rate limit exceeded"):
|
||||
await client.call_tool("batch_process", {"items": ["a", "b", "c"]})
|
||||
|
||||
async def test_custom_client_identification(self, rate_limit_server):
|
||||
"""Test rate limiting with custom client identification."""
|
||||
|
||||
def get_client_id(context):
|
||||
# In a real scenario, this might extract from headers or context
|
||||
return "test_client_123"
|
||||
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(
|
||||
max_requests_per_second=2.0,
|
||||
burst_capacity=1,
|
||||
get_client_id=get_client_id,
|
||||
)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# First request should succeed
|
||||
await client.call_tool("quick_action", {"message": "first"})
|
||||
|
||||
# Second should be rate limited for this specific client
|
||||
with pytest.raises(
|
||||
ToolError, match="Rate limit exceeded for client: test_client_123"
|
||||
):
|
||||
await client.call_tool("quick_action", {"message": "second"})
|
||||
|
||||
async def test_global_rate_limiting(self, rate_limit_server):
|
||||
"""Test global rate limiting across all clients."""
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(
|
||||
max_requests_per_second=2.0, burst_capacity=2, global_limit=True
|
||||
)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# Use up the global capacity
|
||||
await client.call_tool("quick_action", {"message": "1"})
|
||||
await client.call_tool("quick_action", {"message": "2"})
|
||||
|
||||
# Should be globally rate limited
|
||||
with pytest.raises(ToolError, match="Global rate limit exceeded"):
|
||||
await client.call_tool("quick_action", {"message": "3"})
|
||||
|
||||
async def test_rate_limiting_recovery_over_time(self, rate_limit_server):
|
||||
"""Test that rate limiting allows requests again after time passes."""
|
||||
rate_limit_server.add_middleware(
|
||||
RateLimitingMiddleware(
|
||||
max_requests_per_second=10.0, # 10 per second = 1 every 100ms
|
||||
burst_capacity=1,
|
||||
)
|
||||
)
|
||||
|
||||
async with Client(rate_limit_server) as client:
|
||||
# Use up capacity
|
||||
await client.call_tool("quick_action", {"message": "first"})
|
||||
|
||||
# Should be rate limited immediately
|
||||
with pytest.raises(ToolError):
|
||||
await client.call_tool("quick_action", {"message": "second"})
|
||||
|
||||
# Wait for token bucket to refill (150ms should be enough for ~1.5 tokens)
|
||||
await asyncio.sleep(0.15)
|
||||
|
||||
# Should be able to make another request
|
||||
result = await client.call_tool("quick_action", {"message": "after_wait"})
|
||||
assert "after_wait" in str(result)
|
||||
312
tests/server/middleware/test_timing.py
Normal file
312
tests/server/middleware/test_timing.py
Normal file
|
|
@ -0,0 +1,312 @@
|
|||
"""Tests for timing middleware."""
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import time
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
from fastmcp import FastMCP
|
||||
from fastmcp.client import Client
|
||||
from fastmcp.server.middleware.middleware import MiddlewareContext
|
||||
from fastmcp.server.middleware.timing import DetailedTimingMiddleware, TimingMiddleware
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock middleware context."""
|
||||
context = MagicMock(spec=MiddlewareContext)
|
||||
context.method = "test_method"
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_call_next():
|
||||
"""Create a mock call_next function."""
|
||||
return AsyncMock(return_value="test_result")
|
||||
|
||||
|
||||
class TestTimingMiddleware:
|
||||
"""Test timing middleware functionality."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = TimingMiddleware()
|
||||
assert middleware.logger.name == "fastmcp.timing"
|
||||
assert middleware.log_level == logging.INFO
|
||||
|
||||
def test_init_custom(self):
|
||||
"""Test custom initialization."""
|
||||
logger = logging.getLogger("custom")
|
||||
middleware = TimingMiddleware(logger=logger, log_level=logging.DEBUG)
|
||||
assert middleware.logger is logger
|
||||
assert middleware.log_level == logging.DEBUG
|
||||
|
||||
async def test_on_request_success(self, mock_context, mock_call_next, caplog):
|
||||
"""Test timing successful requests."""
|
||||
middleware = TimingMiddleware()
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert result == "test_result"
|
||||
assert mock_call_next.called
|
||||
assert "Request test_method completed in" in caplog.text
|
||||
assert "ms" in caplog.text
|
||||
|
||||
async def test_on_request_failure(self, mock_context, caplog):
|
||||
"""Test timing failed requests."""
|
||||
middleware = TimingMiddleware()
|
||||
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with pytest.raises(ValueError):
|
||||
await middleware.on_request(mock_context, mock_call_next)
|
||||
|
||||
assert "Request test_method failed after" in caplog.text
|
||||
assert "ms: test error" in caplog.text
|
||||
|
||||
|
||||
class TestDetailedTimingMiddleware:
|
||||
"""Test detailed timing middleware functionality."""
|
||||
|
||||
def test_init_default(self):
|
||||
"""Test default initialization."""
|
||||
middleware = DetailedTimingMiddleware()
|
||||
assert middleware.logger.name == "fastmcp.timing.detailed"
|
||||
assert middleware.log_level == logging.INFO
|
||||
|
||||
async def test_on_call_tool(self, caplog):
|
||||
"""Test timing tool calls."""
|
||||
middleware = DetailedTimingMiddleware()
|
||||
context = MagicMock()
|
||||
context.message.name = "test_tool"
|
||||
mock_call_next = AsyncMock(return_value="tool_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_call_tool(context, mock_call_next)
|
||||
|
||||
assert result == "tool_result"
|
||||
assert "Tool 'test_tool' completed in" in caplog.text
|
||||
|
||||
async def test_on_read_resource(self, caplog):
|
||||
"""Test timing resource reads."""
|
||||
middleware = DetailedTimingMiddleware()
|
||||
context = MagicMock()
|
||||
context.message.uri = "test://resource"
|
||||
mock_call_next = AsyncMock(return_value="resource_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_read_resource(context, mock_call_next)
|
||||
|
||||
assert result == "resource_result"
|
||||
assert "Resource 'test://resource' completed in" in caplog.text
|
||||
|
||||
async def test_on_get_prompt(self, caplog):
|
||||
"""Test timing prompt retrieval."""
|
||||
middleware = DetailedTimingMiddleware()
|
||||
context = MagicMock()
|
||||
context.message.name = "test_prompt"
|
||||
mock_call_next = AsyncMock(return_value="prompt_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_get_prompt(context, mock_call_next)
|
||||
|
||||
assert result == "prompt_result"
|
||||
assert "Prompt 'test_prompt' completed in" in caplog.text
|
||||
|
||||
async def test_on_list_tools(self, caplog):
|
||||
"""Test timing tool listing."""
|
||||
middleware = DetailedTimingMiddleware()
|
||||
context = MagicMock()
|
||||
mock_call_next = AsyncMock(return_value="tools_result")
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
result = await middleware.on_list_tools(context, mock_call_next)
|
||||
|
||||
assert result == "tools_result"
|
||||
assert "List tools completed in" in caplog.text
|
||||
|
||||
async def test_operation_failure(self, caplog):
|
||||
"""Test timing failed operations."""
|
||||
middleware = DetailedTimingMiddleware()
|
||||
context = MagicMock()
|
||||
context.message.name = "failing_tool"
|
||||
mock_call_next = AsyncMock(side_effect=RuntimeError("operation failed"))
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
with pytest.raises(RuntimeError):
|
||||
await middleware.on_call_tool(context, mock_call_next)
|
||||
|
||||
assert "Tool 'failing_tool' failed after" in caplog.text
|
||||
assert "ms: operation failed" in caplog.text
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def timing_server():
|
||||
"""Create a FastMCP server specifically for timing middleware tests."""
|
||||
mcp = FastMCP("TimingTestServer")
|
||||
|
||||
@mcp.tool
|
||||
def instant_task() -> str:
|
||||
"""A task that completes instantly."""
|
||||
return "Done instantly"
|
||||
|
||||
@mcp.tool
|
||||
def short_task() -> str:
|
||||
"""A task that takes 0.1 seconds."""
|
||||
time.sleep(0.1)
|
||||
return "Done after 0.1s"
|
||||
|
||||
@mcp.tool
|
||||
def medium_task() -> str:
|
||||
"""A task that takes 0.15 seconds."""
|
||||
time.sleep(0.15)
|
||||
return "Done after 0.15s"
|
||||
|
||||
@mcp.tool
|
||||
def failing_task() -> str:
|
||||
"""A task that always fails."""
|
||||
raise ValueError("Task failed as expected")
|
||||
|
||||
@mcp.resource("timer://test")
|
||||
def test_resource() -> str:
|
||||
"""A resource that takes time to read."""
|
||||
time.sleep(0.05)
|
||||
return "Resource content after 0.05s"
|
||||
|
||||
@mcp.prompt
|
||||
def test_prompt() -> str:
|
||||
"""A prompt that takes time to generate."""
|
||||
time.sleep(0.08)
|
||||
return "Prompt content after 0.08s"
|
||||
|
||||
return mcp
|
||||
|
||||
|
||||
class TestTimingMiddlewareIntegration:
|
||||
"""Integration tests for timing middleware with real FastMCP server."""
|
||||
|
||||
async def test_timing_middleware_measures_tool_execution(
|
||||
self, timing_server, caplog
|
||||
):
|
||||
"""Test that timing middleware accurately measures tool execution times."""
|
||||
timing_server.add_middleware(TimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(timing_server) as client:
|
||||
# Test instant task
|
||||
await client.call_tool("instant_task")
|
||||
|
||||
# Test short task (0.1s)
|
||||
await client.call_tool("short_task")
|
||||
|
||||
# Test medium task (0.15s)
|
||||
await client.call_tool("medium_task")
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have timing logs for all three calls
|
||||
timing_logs = [
|
||||
line
|
||||
for line in log_text.split("\n")
|
||||
if "completed in" in line and "ms" in line
|
||||
]
|
||||
assert len(timing_logs) == 3
|
||||
|
||||
# Verify that longer tasks show longer timing (roughly)
|
||||
assert "tools/call completed in" in log_text
|
||||
assert "ms" in log_text
|
||||
|
||||
async def test_timing_middleware_handles_failures(self, timing_server, caplog):
|
||||
"""Test that timing middleware measures time even for failed operations."""
|
||||
timing_server.add_middleware(TimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(timing_server) as client:
|
||||
# This should fail but still be timed
|
||||
with pytest.raises(Exception):
|
||||
await client.call_tool("failing_task")
|
||||
|
||||
# Should log the failure with timing
|
||||
assert "tools/call failed after" in caplog.text
|
||||
assert "ms:" in caplog.text
|
||||
|
||||
async def test_detailed_timing_middleware_per_operation(
|
||||
self, timing_server, caplog
|
||||
):
|
||||
"""Test that detailed timing middleware provides operation-specific timing."""
|
||||
timing_server.add_middleware(DetailedTimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(timing_server) as client:
|
||||
# Test tool call
|
||||
await client.call_tool("short_task")
|
||||
|
||||
# Test resource read
|
||||
await client.read_resource("timer://test")
|
||||
|
||||
# Test prompt
|
||||
await client.get_prompt("test_prompt")
|
||||
|
||||
# Test listing operations
|
||||
await client.list_tools()
|
||||
await client.list_resources()
|
||||
await client.list_prompts()
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have specific timing logs for each operation type
|
||||
assert "Tool 'short_task' completed in" in log_text
|
||||
assert "Resource 'timer://test' completed in" in log_text
|
||||
assert "Prompt 'test_prompt' completed in" in log_text
|
||||
assert "List tools completed in" in log_text
|
||||
assert "List resources completed in" in log_text
|
||||
assert "List prompts completed in" in log_text
|
||||
|
||||
async def test_timing_middleware_concurrent_operations(self, timing_server, caplog):
|
||||
"""Test timing middleware with concurrent operations."""
|
||||
timing_server.add_middleware(TimingMiddleware())
|
||||
|
||||
with caplog.at_level(logging.INFO):
|
||||
async with Client(timing_server) as client:
|
||||
# Run multiple operations concurrently
|
||||
tasks = [
|
||||
client.call_tool("instant_task"),
|
||||
client.call_tool("short_task"),
|
||||
client.call_tool("instant_task"),
|
||||
]
|
||||
|
||||
await asyncio.gather(*tasks)
|
||||
|
||||
log_text = caplog.text
|
||||
|
||||
# Should have timing logs for all concurrent operations
|
||||
timing_logs = [line for line in log_text.split("\n") if "completed in" in line]
|
||||
assert len(timing_logs) == 3
|
||||
|
||||
async def test_timing_middleware_custom_logger(self, timing_server):
|
||||
"""Test timing middleware with custom logger configuration."""
|
||||
import io
|
||||
import logging
|
||||
|
||||
# Create a custom logger that writes to a string buffer
|
||||
log_buffer = io.StringIO()
|
||||
handler = logging.StreamHandler(log_buffer)
|
||||
custom_logger = logging.getLogger("custom_timing")
|
||||
custom_logger.addHandler(handler)
|
||||
custom_logger.setLevel(logging.DEBUG)
|
||||
|
||||
# Use custom logger and log level
|
||||
timing_server.add_middleware(
|
||||
TimingMiddleware(logger=custom_logger, log_level=logging.DEBUG)
|
||||
)
|
||||
|
||||
async with Client(timing_server) as client:
|
||||
await client.call_tool("instant_task")
|
||||
|
||||
# Check that our custom logger was used
|
||||
log_output = log_buffer.getvalue()
|
||||
assert "tools/call completed in" in log_output
|
||||
assert "ms" in log_output
|
||||
Loading…
Add table
Add a link
Reference in a new issue