mirror of
https://github.com/PrefectHQ/fastmcp.git
synced 2026-08-21 21:14:17 +02:00
Merge pull request #919 from jlowin/middleware
This commit is contained in:
Jeremiah Lowin
GitHub
commit
547bceb154
11 changed files with 2764 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 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 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 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 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 built-in 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 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 built-in 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 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 built-in 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())
|
||||
```
|
||||
6
src/fastmcp/server/middleware/__init__.py
Normal file
6
src/fastmcp/server/middleware/__init__.py
Normal file
|
|
@ -0,0 +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