Spaces:
Running
Running
Jeremiah Lowin Claude commited on
Commit ·
5800457
1
Parent(s): 62222a1
Add production-ready middleware examples with comprehensive tests
Browse filesImplements timing, logging, rate limiting, and error handling middleware to showcase
FastMCP's middleware capabilities as a headline feature. Each middleware includes:
- Production-ready implementations with full configurability
- Comprehensive unit and integration tests with real FastMCP servers
- Updated documentation with teaching examples and production usage patterns
🤖 Generated with [Claude Code](https://claude.ai/code)
Co-Authored-By: Claude <noreply@anthropic.com>
- docs/servers/middleware.mdx +213 -60
- src/fastmcp/server/middleware/__init__.py +5 -0
- src/fastmcp/server/middleware/error_handling.py +206 -0
- src/fastmcp/server/middleware/logging.py +165 -0
- src/fastmcp/server/middleware/rate_limiting.py +231 -0
- src/fastmcp/server/middleware/timing.py +156 -0
- tests/server/middleware/test_error_handling.py +601 -0
- tests/server/middleware/test_logging.py +426 -0
- tests/server/middleware/test_rate_limiting.py +448 -0
- tests/server/middleware/test_timing.py +312 -0
docs/servers/middleware.mdx
CHANGED
|
@@ -329,93 +329,246 @@ parent.mount(child, prefix="child")
|
|
| 329 |
|
| 330 |
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.
|
| 331 |
|
| 332 |
-
## Examples
|
| 333 |
|
| 334 |
-
|
| 335 |
|
| 336 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 337 |
|
| 338 |
```python
|
|
|
|
| 339 |
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
| 340 |
-
from fastmcp.exceptions import ToolError
|
| 341 |
|
| 342 |
-
class
|
| 343 |
-
def
|
| 344 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 345 |
|
| 346 |
async def on_request(self, context: MiddlewareContext, call_next):
|
| 347 |
-
|
| 348 |
-
|
| 349 |
-
|
| 350 |
-
|
| 351 |
-
|
| 352 |
-
|
| 353 |
-
|
| 354 |
-
|
| 355 |
-
|
| 356 |
-
|
| 357 |
-
|
| 358 |
-
|
| 359 |
-
except Exception:
|
| 360 |
-
pass
|
| 361 |
|
|
|
|
| 362 |
return await call_next(context)
|
|
|
|
|
|
|
|
|
|
| 363 |
|
| 364 |
-
|
| 365 |
-
|
| 366 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 367 |
```
|
| 368 |
|
| 369 |
-
|
| 370 |
|
| 371 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
| 372 |
|
| 373 |
```python
|
| 374 |
-
import time
|
| 375 |
import logging
|
|
|
|
| 376 |
|
| 377 |
-
class
|
| 378 |
def __init__(self):
|
| 379 |
-
self.logger = logging.getLogger("
|
|
|
|
| 380 |
|
| 381 |
-
async def
|
| 382 |
-
tool_name = context.message.name
|
| 383 |
-
start_time = time.time()
|
| 384 |
-
|
| 385 |
try:
|
| 386 |
-
|
| 387 |
-
|
| 388 |
-
|
| 389 |
-
|
| 390 |
-
|
| 391 |
-
)
|
| 392 |
-
|
| 393 |
-
return result
|
| 394 |
|
| 395 |
-
|
| 396 |
-
execution_time = time.time() - start_time
|
| 397 |
-
self.logger.error(
|
| 398 |
-
f"Tool {tool_name} failed after {execution_time:.3f}s: {e}"
|
| 399 |
-
)
|
| 400 |
raise
|
| 401 |
```
|
| 402 |
|
| 403 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 404 |
|
| 405 |
-
|
| 406 |
|
| 407 |
```python
|
| 408 |
-
|
| 409 |
-
|
| 410 |
-
|
| 411 |
-
|
| 412 |
-
|
| 413 |
-
|
| 414 |
-
|
| 415 |
-
|
| 416 |
-
|
| 417 |
-
|
| 418 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 419 |
|
| 420 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 421 |
```
|
|
|
|
| 329 |
|
| 330 |
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.
|
| 331 |
|
| 332 |
+
## Built-in Middleware Examples
|
| 333 |
|
| 334 |
+
FastMCP includes several production-ready middleware implementations that demonstrate best practices and provide immediately useful functionality. Let's explore how each type works by building simplified versions, then see how to use the full implementations.
|
| 335 |
|
| 336 |
+
### Timing Middleware
|
| 337 |
+
|
| 338 |
+
Performance monitoring is essential for understanding your server's behavior and identifying bottlenecks. FastMCP includes production-ready timing middleware at `fastmcp.server.middleware.timing`.
|
| 339 |
+
|
| 340 |
+
Here's an example of how it works:
|
| 341 |
|
| 342 |
```python
|
| 343 |
+
import time
|
| 344 |
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
|
|
|
| 345 |
|
| 346 |
+
class SimpleTimingMiddleware(Middleware):
|
| 347 |
+
async def on_request(self, context: MiddlewareContext, call_next):
|
| 348 |
+
start_time = time.perf_counter()
|
| 349 |
+
|
| 350 |
+
try:
|
| 351 |
+
result = await call_next(context)
|
| 352 |
+
duration_ms = (time.perf_counter() - start_time) * 1000
|
| 353 |
+
print(f"Request {context.method} completed in {duration_ms:.2f}ms")
|
| 354 |
+
return result
|
| 355 |
+
except Exception as e:
|
| 356 |
+
duration_ms = (time.perf_counter() - start_time) * 1000
|
| 357 |
+
print(f"Request {context.method} failed after {duration_ms:.2f}ms: {e}")
|
| 358 |
+
raise
|
| 359 |
+
```
|
| 360 |
+
|
| 361 |
+
To use the full production version with proper logging and configuration:
|
| 362 |
+
|
| 363 |
+
```python
|
| 364 |
+
from fastmcp.server.middleware.timing import (
|
| 365 |
+
TimingMiddleware,
|
| 366 |
+
DetailedTimingMiddleware
|
| 367 |
+
)
|
| 368 |
+
|
| 369 |
+
# Basic timing for all requests
|
| 370 |
+
mcp.add_middleware(TimingMiddleware())
|
| 371 |
+
|
| 372 |
+
# Detailed per-operation timing (tools, resources, prompts)
|
| 373 |
+
mcp.add_middleware(DetailedTimingMiddleware())
|
| 374 |
+
```
|
| 375 |
+
|
| 376 |
+
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.
|
| 377 |
+
|
| 378 |
+
### Logging Middleware
|
| 379 |
+
|
| 380 |
+
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`.
|
| 381 |
+
|
| 382 |
+
Here's an example of how it works:
|
| 383 |
+
|
| 384 |
+
```python
|
| 385 |
+
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
| 386 |
+
|
| 387 |
+
class SimpleLoggingMiddleware(Middleware):
|
| 388 |
+
async def on_message(self, context: MiddlewareContext, call_next):
|
| 389 |
+
print(f"Processing {context.method} from {context.source}")
|
| 390 |
+
|
| 391 |
+
try:
|
| 392 |
+
result = await call_next(context)
|
| 393 |
+
print(f"Completed {context.method}")
|
| 394 |
+
return result
|
| 395 |
+
except Exception as e:
|
| 396 |
+
print(f"Failed {context.method}: {e}")
|
| 397 |
+
raise
|
| 398 |
+
```
|
| 399 |
+
|
| 400 |
+
To use the full production versions with advanced features:
|
| 401 |
+
|
| 402 |
+
```python
|
| 403 |
+
from fastmcp.server.middleware.logging import (
|
| 404 |
+
LoggingMiddleware,
|
| 405 |
+
StructuredLoggingMiddleware
|
| 406 |
+
)
|
| 407 |
+
|
| 408 |
+
# Human-readable logging with payload support
|
| 409 |
+
mcp.add_middleware(LoggingMiddleware(
|
| 410 |
+
include_payloads=True,
|
| 411 |
+
max_payload_length=1000
|
| 412 |
+
))
|
| 413 |
+
|
| 414 |
+
# JSON-structured logging for log aggregation tools
|
| 415 |
+
mcp.add_middleware(StructuredLoggingMiddleware(include_payloads=True))
|
| 416 |
+
```
|
| 417 |
+
|
| 418 |
+
The production versions include payload logging, structured JSON output, custom logger support, payload size limits, and operation-specific hooks for granular control.
|
| 419 |
+
|
| 420 |
+
### Rate Limiting Middleware
|
| 421 |
+
|
| 422 |
+
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`.
|
| 423 |
+
|
| 424 |
+
Here's an example of how it works:
|
| 425 |
+
|
| 426 |
+
```python
|
| 427 |
+
import time
|
| 428 |
+
from collections import defaultdict
|
| 429 |
+
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
| 430 |
+
from mcp import McpError
|
| 431 |
+
from mcp.types import ErrorData
|
| 432 |
+
|
| 433 |
+
class SimpleRateLimitMiddleware(Middleware):
|
| 434 |
+
def __init__(self, requests_per_minute: int = 60):
|
| 435 |
+
self.requests_per_minute = requests_per_minute
|
| 436 |
+
self.client_requests = defaultdict(list)
|
| 437 |
|
| 438 |
async def on_request(self, context: MiddlewareContext, call_next):
|
| 439 |
+
current_time = time.time()
|
| 440 |
+
client_id = "default" # In practice, extract from headers or context
|
| 441 |
+
|
| 442 |
+
# Clean old requests and check limit
|
| 443 |
+
cutoff_time = current_time - 60
|
| 444 |
+
self.client_requests[client_id] = [
|
| 445 |
+
req_time for req_time in self.client_requests[client_id]
|
| 446 |
+
if req_time > cutoff_time
|
| 447 |
+
]
|
| 448 |
+
|
| 449 |
+
if len(self.client_requests[client_id]) >= self.requests_per_minute:
|
| 450 |
+
raise McpError(ErrorData(code=-32000, message="Rate limit exceeded"))
|
|
|
|
|
|
|
| 451 |
|
| 452 |
+
self.client_requests[client_id].append(current_time)
|
| 453 |
return await call_next(context)
|
| 454 |
+
```
|
| 455 |
+
|
| 456 |
+
To use the full production versions with advanced algorithms:
|
| 457 |
|
| 458 |
+
```python
|
| 459 |
+
from fastmcp.server.middleware.rate_limiting import (
|
| 460 |
+
RateLimitingMiddleware,
|
| 461 |
+
SlidingWindowRateLimitingMiddleware
|
| 462 |
+
)
|
| 463 |
+
|
| 464 |
+
# Token bucket rate limiting (allows controlled bursts)
|
| 465 |
+
mcp.add_middleware(RateLimitingMiddleware(
|
| 466 |
+
max_requests_per_second=10.0,
|
| 467 |
+
burst_capacity=20
|
| 468 |
+
))
|
| 469 |
+
|
| 470 |
+
# Sliding window rate limiting (precise time-based control)
|
| 471 |
+
mcp.add_middleware(SlidingWindowRateLimitingMiddleware(
|
| 472 |
+
max_requests=100,
|
| 473 |
+
window_minutes=1
|
| 474 |
+
))
|
| 475 |
```
|
| 476 |
|
| 477 |
+
The production versions include token bucket algorithms, per-client identification, global rate limiting, and async-safe implementations with configurable client identification functions.
|
| 478 |
|
| 479 |
+
### Error Handling Middleware
|
| 480 |
+
|
| 481 |
+
Consistent error handling and recovery is critical for robust MCP servers. FastMCP provides comprehensive error handling middleware at `fastmcp.server.middleware.error_handling`.
|
| 482 |
+
|
| 483 |
+
Here's an example of how it works:
|
| 484 |
|
| 485 |
```python
|
|
|
|
| 486 |
import logging
|
| 487 |
+
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
| 488 |
|
| 489 |
+
class SimpleErrorHandlingMiddleware(Middleware):
|
| 490 |
def __init__(self):
|
| 491 |
+
self.logger = logging.getLogger("errors")
|
| 492 |
+
self.error_counts = {}
|
| 493 |
|
| 494 |
+
async def on_message(self, context: MiddlewareContext, call_next):
|
|
|
|
|
|
|
|
|
|
| 495 |
try:
|
| 496 |
+
return await call_next(context)
|
| 497 |
+
except Exception as error:
|
| 498 |
+
# Log the error and track statistics
|
| 499 |
+
error_key = f"{type(error).__name__}:{context.method}"
|
| 500 |
+
self.error_counts[error_key] = self.error_counts.get(error_key, 0) + 1
|
|
|
|
|
|
|
|
|
|
| 501 |
|
| 502 |
+
self.logger.error(f"Error in {context.method}: {type(error).__name__}: {error}")
|
|
|
|
|
|
|
|
|
|
|
|
|
| 503 |
raise
|
| 504 |
```
|
| 505 |
|
| 506 |
+
To use the full production versions with advanced features:
|
| 507 |
+
|
| 508 |
+
```python
|
| 509 |
+
from fastmcp.server.middleware.error_handling import (
|
| 510 |
+
ErrorHandlingMiddleware,
|
| 511 |
+
RetryMiddleware
|
| 512 |
+
)
|
| 513 |
+
|
| 514 |
+
# Comprehensive error logging and transformation
|
| 515 |
+
mcp.add_middleware(ErrorHandlingMiddleware(
|
| 516 |
+
include_traceback=True,
|
| 517 |
+
transform_errors=True,
|
| 518 |
+
error_callback=my_error_callback
|
| 519 |
+
))
|
| 520 |
+
|
| 521 |
+
# Automatic retry with exponential backoff
|
| 522 |
+
mcp.add_middleware(RetryMiddleware(
|
| 523 |
+
max_retries=3,
|
| 524 |
+
retry_exceptions=(ConnectionError, TimeoutError)
|
| 525 |
+
))
|
| 526 |
+
```
|
| 527 |
+
|
| 528 |
+
The production versions include error transformation, custom callbacks, configurable retry logic, and proper MCP error formatting.
|
| 529 |
+
|
| 530 |
+
### Combining Middleware
|
| 531 |
|
| 532 |
+
These middleware work together seamlessly:
|
| 533 |
|
| 534 |
```python
|
| 535 |
+
from fastmcp import FastMCP
|
| 536 |
+
from fastmcp.server.middleware.timing import TimingMiddleware
|
| 537 |
+
from fastmcp.server.middleware.logging import LoggingMiddleware
|
| 538 |
+
from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
|
| 539 |
+
from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
|
| 540 |
+
|
| 541 |
+
mcp = FastMCP("Production Server")
|
| 542 |
+
|
| 543 |
+
# Add middleware in logical order
|
| 544 |
+
mcp.add_middleware(ErrorHandlingMiddleware()) # Handle errors first
|
| 545 |
+
mcp.add_middleware(RateLimitingMiddleware(max_requests_per_second=50))
|
| 546 |
+
mcp.add_middleware(TimingMiddleware()) # Time actual execution
|
| 547 |
+
mcp.add_middleware(LoggingMiddleware()) # Log everything
|
| 548 |
+
|
| 549 |
+
@mcp.tool
|
| 550 |
+
def my_tool(data: str) -> str:
|
| 551 |
+
return f"Processed: {data}"
|
| 552 |
+
```
|
| 553 |
+
|
| 554 |
+
This configuration provides comprehensive monitoring, protection, and observability for your MCP server.
|
| 555 |
+
|
| 556 |
+
### Custom Middleware Example
|
| 557 |
+
|
| 558 |
+
You can also create custom middleware by extending the base class:
|
| 559 |
+
|
| 560 |
+
```python
|
| 561 |
+
from fastmcp.server.middleware import Middleware, MiddlewareContext
|
| 562 |
+
|
| 563 |
+
class CustomHeaderMiddleware(Middleware):
|
| 564 |
+
async def on_request(self, context: MiddlewareContext, call_next):
|
| 565 |
+
# Add custom logic here
|
| 566 |
+
print(f"Processing {context.method}")
|
| 567 |
|
| 568 |
+
result = await call_next(context)
|
| 569 |
+
|
| 570 |
+
print(f"Completed {context.method}")
|
| 571 |
+
return result
|
| 572 |
+
|
| 573 |
+
mcp.add_middleware(CustomHeaderMiddleware())
|
| 574 |
```
|
src/fastmcp/server/middleware/__init__.py
CHANGED
|
@@ -1 +1,6 @@
|
|
| 1 |
from .middleware import Middleware, MiddlewareContext
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
from .middleware import Middleware, MiddlewareContext
|
| 2 |
+
|
| 3 |
+
__all__ = [
|
| 4 |
+
"Middleware",
|
| 5 |
+
"MiddlewareContext",
|
| 6 |
+
]
|
src/fastmcp/server/middleware/error_handling.py
ADDED
|
@@ -0,0 +1,206 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Error handling middleware for consistent error responses and tracking."""
|
| 2 |
+
|
| 3 |
+
import asyncio
|
| 4 |
+
import logging
|
| 5 |
+
import traceback
|
| 6 |
+
from collections.abc import Callable
|
| 7 |
+
from typing import Any
|
| 8 |
+
|
| 9 |
+
from mcp import McpError
|
| 10 |
+
from mcp.types import ErrorData
|
| 11 |
+
|
| 12 |
+
from .middleware import CallNext, Middleware, MiddlewareContext
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class ErrorHandlingMiddleware(Middleware):
|
| 16 |
+
"""Middleware that provides consistent error handling and logging.
|
| 17 |
+
|
| 18 |
+
Catches exceptions, logs them appropriately, and converts them to
|
| 19 |
+
proper MCP error responses. Also tracks error patterns for monitoring.
|
| 20 |
+
|
| 21 |
+
Example:
|
| 22 |
+
```python
|
| 23 |
+
from fastmcp.server.middleware.error_handling import ErrorHandlingMiddleware
|
| 24 |
+
import logging
|
| 25 |
+
|
| 26 |
+
# Configure logging to see error details
|
| 27 |
+
logging.basicConfig(level=logging.ERROR)
|
| 28 |
+
|
| 29 |
+
mcp = FastMCP("MyServer")
|
| 30 |
+
mcp.add_middleware(ErrorHandlingMiddleware())
|
| 31 |
+
```
|
| 32 |
+
"""
|
| 33 |
+
|
| 34 |
+
def __init__(
|
| 35 |
+
self,
|
| 36 |
+
logger: logging.Logger | None = None,
|
| 37 |
+
include_traceback: bool = False,
|
| 38 |
+
error_callback: Callable[[Exception, MiddlewareContext], None] | None = None,
|
| 39 |
+
transform_errors: bool = True,
|
| 40 |
+
):
|
| 41 |
+
"""Initialize error handling middleware.
|
| 42 |
+
|
| 43 |
+
Args:
|
| 44 |
+
logger: Logger instance for error logging. If None, uses 'fastmcp.errors'
|
| 45 |
+
include_traceback: Whether to include full traceback in error logs
|
| 46 |
+
error_callback: Optional callback function called for each error
|
| 47 |
+
transform_errors: Whether to transform non-MCP errors to McpError
|
| 48 |
+
"""
|
| 49 |
+
self.logger = logger or logging.getLogger("fastmcp.errors")
|
| 50 |
+
self.include_traceback = include_traceback
|
| 51 |
+
self.error_callback = error_callback
|
| 52 |
+
self.transform_errors = transform_errors
|
| 53 |
+
self.error_counts = {}
|
| 54 |
+
|
| 55 |
+
def _log_error(self, error: Exception, context: MiddlewareContext) -> None:
|
| 56 |
+
"""Log error with appropriate detail level."""
|
| 57 |
+
error_type = type(error).__name__
|
| 58 |
+
method = context.method or "unknown"
|
| 59 |
+
|
| 60 |
+
# Track error counts
|
| 61 |
+
error_key = f"{error_type}:{method}"
|
| 62 |
+
self.error_counts[error_key] = self.error_counts.get(error_key, 0) + 1
|
| 63 |
+
|
| 64 |
+
base_message = f"Error in {method}: {error_type}: {str(error)}"
|
| 65 |
+
|
| 66 |
+
if self.include_traceback:
|
| 67 |
+
self.logger.error(f"{base_message}\n{traceback.format_exc()}")
|
| 68 |
+
else:
|
| 69 |
+
self.logger.error(base_message)
|
| 70 |
+
|
| 71 |
+
# Call custom error callback if provided
|
| 72 |
+
if self.error_callback:
|
| 73 |
+
try:
|
| 74 |
+
self.error_callback(error, context)
|
| 75 |
+
except Exception as callback_error:
|
| 76 |
+
self.logger.error(f"Error in error callback: {callback_error}")
|
| 77 |
+
|
| 78 |
+
def _transform_error(self, error: Exception) -> Exception:
|
| 79 |
+
"""Transform non-MCP errors to proper MCP errors."""
|
| 80 |
+
if isinstance(error, McpError):
|
| 81 |
+
return error
|
| 82 |
+
|
| 83 |
+
if not self.transform_errors:
|
| 84 |
+
return error
|
| 85 |
+
|
| 86 |
+
# Map common exceptions to appropriate MCP error codes
|
| 87 |
+
error_type = type(error)
|
| 88 |
+
|
| 89 |
+
if error_type in (ValueError, TypeError):
|
| 90 |
+
return McpError(
|
| 91 |
+
ErrorData(code=-32602, message=f"Invalid params: {str(error)}")
|
| 92 |
+
)
|
| 93 |
+
elif error_type in (FileNotFoundError, KeyError):
|
| 94 |
+
return McpError(
|
| 95 |
+
ErrorData(code=-32001, message=f"Resource not found: {str(error)}")
|
| 96 |
+
)
|
| 97 |
+
elif error_type is PermissionError:
|
| 98 |
+
return McpError(
|
| 99 |
+
ErrorData(code=-32000, message=f"Permission denied: {str(error)}")
|
| 100 |
+
)
|
| 101 |
+
elif error_type in (TimeoutError, asyncio.TimeoutError):
|
| 102 |
+
return McpError(
|
| 103 |
+
ErrorData(code=-32000, message=f"Request timeout: {str(error)}")
|
| 104 |
+
)
|
| 105 |
+
else:
|
| 106 |
+
return McpError(
|
| 107 |
+
ErrorData(code=-32603, message=f"Internal error: {str(error)}")
|
| 108 |
+
)
|
| 109 |
+
|
| 110 |
+
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 111 |
+
"""Handle errors for all messages."""
|
| 112 |
+
try:
|
| 113 |
+
return await call_next(context)
|
| 114 |
+
except Exception as error:
|
| 115 |
+
self._log_error(error, context)
|
| 116 |
+
|
| 117 |
+
# Transform and re-raise
|
| 118 |
+
transformed_error = self._transform_error(error)
|
| 119 |
+
raise transformed_error
|
| 120 |
+
|
| 121 |
+
def get_error_stats(self) -> dict[str, int]:
|
| 122 |
+
"""Get error statistics for monitoring."""
|
| 123 |
+
return self.error_counts.copy()
|
| 124 |
+
|
| 125 |
+
|
| 126 |
+
class RetryMiddleware(Middleware):
|
| 127 |
+
"""Middleware that implements automatic retry logic for failed requests.
|
| 128 |
+
|
| 129 |
+
Retries requests that fail with transient errors, using exponential
|
| 130 |
+
backoff to avoid overwhelming the server or external dependencies.
|
| 131 |
+
|
| 132 |
+
Example:
|
| 133 |
+
```python
|
| 134 |
+
from fastmcp.server.middleware.error_handling import RetryMiddleware
|
| 135 |
+
|
| 136 |
+
# Retry up to 3 times with exponential backoff
|
| 137 |
+
retry_middleware = RetryMiddleware(
|
| 138 |
+
max_retries=3,
|
| 139 |
+
retry_exceptions=(ConnectionError, TimeoutError)
|
| 140 |
+
)
|
| 141 |
+
|
| 142 |
+
mcp = FastMCP("MyServer")
|
| 143 |
+
mcp.add_middleware(retry_middleware)
|
| 144 |
+
```
|
| 145 |
+
"""
|
| 146 |
+
|
| 147 |
+
def __init__(
|
| 148 |
+
self,
|
| 149 |
+
max_retries: int = 3,
|
| 150 |
+
base_delay: float = 1.0,
|
| 151 |
+
max_delay: float = 60.0,
|
| 152 |
+
backoff_multiplier: float = 2.0,
|
| 153 |
+
retry_exceptions: tuple[type[Exception], ...] = (ConnectionError, TimeoutError),
|
| 154 |
+
logger: logging.Logger | None = None,
|
| 155 |
+
):
|
| 156 |
+
"""Initialize retry middleware.
|
| 157 |
+
|
| 158 |
+
Args:
|
| 159 |
+
max_retries: Maximum number of retry attempts
|
| 160 |
+
base_delay: Initial delay between retries in seconds
|
| 161 |
+
max_delay: Maximum delay between retries in seconds
|
| 162 |
+
backoff_multiplier: Multiplier for exponential backoff
|
| 163 |
+
retry_exceptions: Tuple of exception types that should trigger retries
|
| 164 |
+
logger: Logger for retry attempts
|
| 165 |
+
"""
|
| 166 |
+
self.max_retries = max_retries
|
| 167 |
+
self.base_delay = base_delay
|
| 168 |
+
self.max_delay = max_delay
|
| 169 |
+
self.backoff_multiplier = backoff_multiplier
|
| 170 |
+
self.retry_exceptions = retry_exceptions
|
| 171 |
+
self.logger = logger or logging.getLogger("fastmcp.retry")
|
| 172 |
+
|
| 173 |
+
def _should_retry(self, error: Exception) -> bool:
|
| 174 |
+
"""Determine if an error should trigger a retry."""
|
| 175 |
+
return isinstance(error, self.retry_exceptions)
|
| 176 |
+
|
| 177 |
+
def _calculate_delay(self, attempt: int) -> float:
|
| 178 |
+
"""Calculate delay for the given attempt number."""
|
| 179 |
+
delay = self.base_delay * (self.backoff_multiplier**attempt)
|
| 180 |
+
return min(delay, self.max_delay)
|
| 181 |
+
|
| 182 |
+
async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 183 |
+
"""Implement retry logic for requests."""
|
| 184 |
+
last_error = None
|
| 185 |
+
|
| 186 |
+
for attempt in range(self.max_retries + 1):
|
| 187 |
+
try:
|
| 188 |
+
return await call_next(context)
|
| 189 |
+
except Exception as error:
|
| 190 |
+
last_error = error
|
| 191 |
+
|
| 192 |
+
# Don't retry on the last attempt or if it's not a retryable error
|
| 193 |
+
if attempt == self.max_retries or not self._should_retry(error):
|
| 194 |
+
break
|
| 195 |
+
|
| 196 |
+
delay = self._calculate_delay(attempt)
|
| 197 |
+
self.logger.warning(
|
| 198 |
+
f"Request {context.method} failed (attempt {attempt + 1}/{self.max_retries + 1}): "
|
| 199 |
+
f"{type(error).__name__}: {str(error)}. Retrying in {delay:.1f}s..."
|
| 200 |
+
)
|
| 201 |
+
|
| 202 |
+
await asyncio.sleep(delay)
|
| 203 |
+
|
| 204 |
+
# Re-raise the last error if all retries failed
|
| 205 |
+
if last_error:
|
| 206 |
+
raise last_error
|
src/fastmcp/server/middleware/logging.py
ADDED
|
@@ -0,0 +1,165 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Comprehensive logging middleware for FastMCP servers."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import logging
|
| 5 |
+
from typing import Any
|
| 6 |
+
|
| 7 |
+
from .middleware import CallNext, Middleware, MiddlewareContext
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class LoggingMiddleware(Middleware):
|
| 11 |
+
"""Middleware that provides comprehensive request and response logging.
|
| 12 |
+
|
| 13 |
+
Logs all MCP messages with configurable detail levels. Useful for debugging,
|
| 14 |
+
monitoring, and understanding server usage patterns.
|
| 15 |
+
|
| 16 |
+
Example:
|
| 17 |
+
```python
|
| 18 |
+
from fastmcp.server.middleware.logging import LoggingMiddleware
|
| 19 |
+
import logging
|
| 20 |
+
|
| 21 |
+
# Configure logging
|
| 22 |
+
logging.basicConfig(level=logging.INFO)
|
| 23 |
+
|
| 24 |
+
mcp = FastMCP("MyServer")
|
| 25 |
+
mcp.add_middleware(LoggingMiddleware())
|
| 26 |
+
```
|
| 27 |
+
"""
|
| 28 |
+
|
| 29 |
+
def __init__(
|
| 30 |
+
self,
|
| 31 |
+
logger: logging.Logger | None = None,
|
| 32 |
+
log_level: int = logging.INFO,
|
| 33 |
+
include_payloads: bool = False,
|
| 34 |
+
max_payload_length: int = 1000,
|
| 35 |
+
):
|
| 36 |
+
"""Initialize logging middleware.
|
| 37 |
+
|
| 38 |
+
Args:
|
| 39 |
+
logger: Logger instance to use. If None, creates a logger named 'fastmcp.requests'
|
| 40 |
+
log_level: Log level for messages (default: INFO)
|
| 41 |
+
include_payloads: Whether to include message payloads in logs
|
| 42 |
+
max_payload_length: Maximum length of payload to log (prevents huge logs)
|
| 43 |
+
"""
|
| 44 |
+
self.logger = logger or logging.getLogger("fastmcp.requests")
|
| 45 |
+
self.log_level = log_level
|
| 46 |
+
self.include_payloads = include_payloads
|
| 47 |
+
self.max_payload_length = max_payload_length
|
| 48 |
+
|
| 49 |
+
def _format_message(self, context: MiddlewareContext) -> str:
|
| 50 |
+
"""Format a message for logging."""
|
| 51 |
+
parts = [
|
| 52 |
+
f"source={context.source}",
|
| 53 |
+
f"type={context.type}",
|
| 54 |
+
f"method={context.method or 'unknown'}",
|
| 55 |
+
]
|
| 56 |
+
|
| 57 |
+
if self.include_payloads and hasattr(context.message, "__dict__"):
|
| 58 |
+
try:
|
| 59 |
+
payload = json.dumps(context.message.__dict__, default=str)
|
| 60 |
+
if len(payload) > self.max_payload_length:
|
| 61 |
+
payload = payload[: self.max_payload_length] + "..."
|
| 62 |
+
parts.append(f"payload={payload}")
|
| 63 |
+
except (TypeError, ValueError):
|
| 64 |
+
parts.append("payload=<non-serializable>")
|
| 65 |
+
|
| 66 |
+
return " ".join(parts)
|
| 67 |
+
|
| 68 |
+
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 69 |
+
"""Log all messages."""
|
| 70 |
+
message_info = self._format_message(context)
|
| 71 |
+
|
| 72 |
+
self.logger.log(self.log_level, f"Processing message: {message_info}")
|
| 73 |
+
|
| 74 |
+
try:
|
| 75 |
+
result = await call_next(context)
|
| 76 |
+
self.logger.log(
|
| 77 |
+
self.log_level, f"Completed message: {context.method or 'unknown'}"
|
| 78 |
+
)
|
| 79 |
+
return result
|
| 80 |
+
except Exception as e:
|
| 81 |
+
self.logger.log(
|
| 82 |
+
logging.ERROR, f"Failed message: {context.method or 'unknown'} - {e}"
|
| 83 |
+
)
|
| 84 |
+
raise
|
| 85 |
+
|
| 86 |
+
|
| 87 |
+
class StructuredLoggingMiddleware(Middleware):
|
| 88 |
+
"""Middleware that provides structured JSON logging for better log analysis.
|
| 89 |
+
|
| 90 |
+
Outputs structured logs that are easier to parse and analyze with log
|
| 91 |
+
aggregation tools like ELK stack, Splunk, or cloud logging services.
|
| 92 |
+
|
| 93 |
+
Example:
|
| 94 |
+
```python
|
| 95 |
+
from fastmcp.server.middleware.logging import StructuredLoggingMiddleware
|
| 96 |
+
import logging
|
| 97 |
+
|
| 98 |
+
mcp = FastMCP("MyServer")
|
| 99 |
+
mcp.add_middleware(StructuredLoggingMiddleware())
|
| 100 |
+
```
|
| 101 |
+
"""
|
| 102 |
+
|
| 103 |
+
def __init__(
|
| 104 |
+
self,
|
| 105 |
+
logger: logging.Logger | None = None,
|
| 106 |
+
log_level: int = logging.INFO,
|
| 107 |
+
include_payloads: bool = False,
|
| 108 |
+
):
|
| 109 |
+
"""Initialize structured logging middleware.
|
| 110 |
+
|
| 111 |
+
Args:
|
| 112 |
+
logger: Logger instance to use. If None, creates a logger named 'fastmcp.structured'
|
| 113 |
+
log_level: Log level for messages (default: INFO)
|
| 114 |
+
include_payloads: Whether to include message payloads in logs
|
| 115 |
+
"""
|
| 116 |
+
self.logger = logger or logging.getLogger("fastmcp.structured")
|
| 117 |
+
self.log_level = log_level
|
| 118 |
+
self.include_payloads = include_payloads
|
| 119 |
+
|
| 120 |
+
def _create_log_entry(
|
| 121 |
+
self, context: MiddlewareContext, event: str, **extra_fields
|
| 122 |
+
) -> dict:
|
| 123 |
+
"""Create a structured log entry."""
|
| 124 |
+
entry = {
|
| 125 |
+
"event": event,
|
| 126 |
+
"timestamp": context.timestamp.isoformat(),
|
| 127 |
+
"source": context.source,
|
| 128 |
+
"type": context.type,
|
| 129 |
+
"method": context.method,
|
| 130 |
+
**extra_fields,
|
| 131 |
+
}
|
| 132 |
+
|
| 133 |
+
if self.include_payloads and hasattr(context.message, "__dict__"):
|
| 134 |
+
try:
|
| 135 |
+
entry["payload"] = context.message.__dict__
|
| 136 |
+
except (TypeError, ValueError):
|
| 137 |
+
entry["payload"] = "<non-serializable>"
|
| 138 |
+
|
| 139 |
+
return entry
|
| 140 |
+
|
| 141 |
+
async def on_message(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 142 |
+
"""Log structured message information."""
|
| 143 |
+
start_entry = self._create_log_entry(context, "request_start")
|
| 144 |
+
self.logger.log(self.log_level, json.dumps(start_entry))
|
| 145 |
+
|
| 146 |
+
try:
|
| 147 |
+
result = await call_next(context)
|
| 148 |
+
|
| 149 |
+
success_entry = self._create_log_entry(
|
| 150 |
+
context,
|
| 151 |
+
"request_success",
|
| 152 |
+
result_type=type(result).__name__ if result else None,
|
| 153 |
+
)
|
| 154 |
+
self.logger.log(self.log_level, json.dumps(success_entry))
|
| 155 |
+
|
| 156 |
+
return result
|
| 157 |
+
except Exception as e:
|
| 158 |
+
error_entry = self._create_log_entry(
|
| 159 |
+
context,
|
| 160 |
+
"request_error",
|
| 161 |
+
error_type=type(e).__name__,
|
| 162 |
+
error_message=str(e),
|
| 163 |
+
)
|
| 164 |
+
self.logger.log(logging.ERROR, json.dumps(error_entry))
|
| 165 |
+
raise
|
src/fastmcp/server/middleware/rate_limiting.py
ADDED
|
@@ -0,0 +1,231 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Rate limiting middleware for protecting FastMCP servers from abuse."""
|
| 2 |
+
|
| 3 |
+
import asyncio
|
| 4 |
+
import time
|
| 5 |
+
from collections import defaultdict, deque
|
| 6 |
+
from collections.abc import Callable
|
| 7 |
+
from typing import Any
|
| 8 |
+
|
| 9 |
+
from mcp import McpError
|
| 10 |
+
from mcp.types import ErrorData
|
| 11 |
+
|
| 12 |
+
from .middleware import CallNext, Middleware, MiddlewareContext
|
| 13 |
+
|
| 14 |
+
|
| 15 |
+
class RateLimitError(McpError):
|
| 16 |
+
"""Error raised when rate limit is exceeded."""
|
| 17 |
+
|
| 18 |
+
def __init__(self, message: str = "Rate limit exceeded"):
|
| 19 |
+
super().__init__(ErrorData(code=-32000, message=message))
|
| 20 |
+
|
| 21 |
+
|
| 22 |
+
class TokenBucketRateLimiter:
|
| 23 |
+
"""Token bucket implementation for rate limiting."""
|
| 24 |
+
|
| 25 |
+
def __init__(self, capacity: int, refill_rate: float):
|
| 26 |
+
"""Initialize token bucket.
|
| 27 |
+
|
| 28 |
+
Args:
|
| 29 |
+
capacity: Maximum number of tokens in the bucket
|
| 30 |
+
refill_rate: Tokens added per second
|
| 31 |
+
"""
|
| 32 |
+
self.capacity = capacity
|
| 33 |
+
self.refill_rate = refill_rate
|
| 34 |
+
self.tokens = capacity
|
| 35 |
+
self.last_refill = time.time()
|
| 36 |
+
self._lock = asyncio.Lock()
|
| 37 |
+
|
| 38 |
+
async def consume(self, tokens: int = 1) -> bool:
|
| 39 |
+
"""Try to consume tokens from the bucket.
|
| 40 |
+
|
| 41 |
+
Args:
|
| 42 |
+
tokens: Number of tokens to consume
|
| 43 |
+
|
| 44 |
+
Returns:
|
| 45 |
+
True if tokens were available and consumed, False otherwise
|
| 46 |
+
"""
|
| 47 |
+
async with self._lock:
|
| 48 |
+
now = time.time()
|
| 49 |
+
elapsed = now - self.last_refill
|
| 50 |
+
|
| 51 |
+
# Add tokens based on elapsed time
|
| 52 |
+
self.tokens = min(self.capacity, self.tokens + elapsed * self.refill_rate)
|
| 53 |
+
self.last_refill = now
|
| 54 |
+
|
| 55 |
+
if self.tokens >= tokens:
|
| 56 |
+
self.tokens -= tokens
|
| 57 |
+
return True
|
| 58 |
+
return False
|
| 59 |
+
|
| 60 |
+
|
| 61 |
+
class SlidingWindowRateLimiter:
|
| 62 |
+
"""Sliding window rate limiter implementation."""
|
| 63 |
+
|
| 64 |
+
def __init__(self, max_requests: int, window_seconds: int):
|
| 65 |
+
"""Initialize sliding window rate limiter.
|
| 66 |
+
|
| 67 |
+
Args:
|
| 68 |
+
max_requests: Maximum requests allowed in the time window
|
| 69 |
+
window_seconds: Time window in seconds
|
| 70 |
+
"""
|
| 71 |
+
self.max_requests = max_requests
|
| 72 |
+
self.window_seconds = window_seconds
|
| 73 |
+
self.requests = deque()
|
| 74 |
+
self._lock = asyncio.Lock()
|
| 75 |
+
|
| 76 |
+
async def is_allowed(self) -> bool:
|
| 77 |
+
"""Check if a request is allowed."""
|
| 78 |
+
async with self._lock:
|
| 79 |
+
now = time.time()
|
| 80 |
+
cutoff = now - self.window_seconds
|
| 81 |
+
|
| 82 |
+
# Remove old requests outside the window
|
| 83 |
+
while self.requests and self.requests[0] < cutoff:
|
| 84 |
+
self.requests.popleft()
|
| 85 |
+
|
| 86 |
+
if len(self.requests) < self.max_requests:
|
| 87 |
+
self.requests.append(now)
|
| 88 |
+
return True
|
| 89 |
+
return False
|
| 90 |
+
|
| 91 |
+
|
| 92 |
+
class RateLimitingMiddleware(Middleware):
|
| 93 |
+
"""Middleware that implements rate limiting to prevent server abuse.
|
| 94 |
+
|
| 95 |
+
Uses a token bucket algorithm by default, allowing for burst traffic
|
| 96 |
+
while maintaining a sustainable long-term rate.
|
| 97 |
+
|
| 98 |
+
Example:
|
| 99 |
+
```python
|
| 100 |
+
from fastmcp.server.middleware.rate_limiting import RateLimitingMiddleware
|
| 101 |
+
|
| 102 |
+
# Allow 10 requests per second with bursts up to 20
|
| 103 |
+
rate_limiter = RateLimitingMiddleware(
|
| 104 |
+
max_requests_per_second=10,
|
| 105 |
+
burst_capacity=20
|
| 106 |
+
)
|
| 107 |
+
|
| 108 |
+
mcp = FastMCP("MyServer")
|
| 109 |
+
mcp.add_middleware(rate_limiter)
|
| 110 |
+
```
|
| 111 |
+
"""
|
| 112 |
+
|
| 113 |
+
def __init__(
|
| 114 |
+
self,
|
| 115 |
+
max_requests_per_second: float = 10.0,
|
| 116 |
+
burst_capacity: int | None = None,
|
| 117 |
+
get_client_id: Callable[[MiddlewareContext], str] | None = None,
|
| 118 |
+
global_limit: bool = False,
|
| 119 |
+
):
|
| 120 |
+
"""Initialize rate limiting middleware.
|
| 121 |
+
|
| 122 |
+
Args:
|
| 123 |
+
max_requests_per_second: Sustained requests per second allowed
|
| 124 |
+
burst_capacity: Maximum burst capacity. If None, defaults to 2x max_requests_per_second
|
| 125 |
+
get_client_id: Function to extract client ID from context. If None, uses global limiting
|
| 126 |
+
global_limit: If True, apply limit globally; if False, per-client
|
| 127 |
+
"""
|
| 128 |
+
self.max_requests_per_second = max_requests_per_second
|
| 129 |
+
self.burst_capacity = burst_capacity or int(max_requests_per_second * 2)
|
| 130 |
+
self.get_client_id = get_client_id
|
| 131 |
+
self.global_limit = global_limit
|
| 132 |
+
|
| 133 |
+
# Storage for rate limiters per client
|
| 134 |
+
self.limiters: dict[str, TokenBucketRateLimiter] = defaultdict(
|
| 135 |
+
lambda: TokenBucketRateLimiter(
|
| 136 |
+
self.burst_capacity, self.max_requests_per_second
|
| 137 |
+
)
|
| 138 |
+
)
|
| 139 |
+
|
| 140 |
+
# Global rate limiter
|
| 141 |
+
if self.global_limit:
|
| 142 |
+
self.global_limiter = TokenBucketRateLimiter(
|
| 143 |
+
self.burst_capacity, self.max_requests_per_second
|
| 144 |
+
)
|
| 145 |
+
|
| 146 |
+
def _get_client_identifier(self, context: MiddlewareContext) -> str:
|
| 147 |
+
"""Get client identifier for rate limiting."""
|
| 148 |
+
if self.get_client_id:
|
| 149 |
+
return self.get_client_id(context)
|
| 150 |
+
return "global"
|
| 151 |
+
|
| 152 |
+
async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 153 |
+
"""Apply rate limiting to requests."""
|
| 154 |
+
if self.global_limit:
|
| 155 |
+
# Global rate limiting
|
| 156 |
+
allowed = await self.global_limiter.consume()
|
| 157 |
+
if not allowed:
|
| 158 |
+
raise RateLimitError("Global rate limit exceeded")
|
| 159 |
+
else:
|
| 160 |
+
# Per-client rate limiting
|
| 161 |
+
client_id = self._get_client_identifier(context)
|
| 162 |
+
limiter = self.limiters[client_id]
|
| 163 |
+
allowed = await limiter.consume()
|
| 164 |
+
if not allowed:
|
| 165 |
+
raise RateLimitError(f"Rate limit exceeded for client: {client_id}")
|
| 166 |
+
|
| 167 |
+
return await call_next(context)
|
| 168 |
+
|
| 169 |
+
|
| 170 |
+
class SlidingWindowRateLimitingMiddleware(Middleware):
|
| 171 |
+
"""Middleware that implements sliding window rate limiting.
|
| 172 |
+
|
| 173 |
+
Uses a sliding window approach which provides more precise rate limiting
|
| 174 |
+
but uses more memory to track individual request timestamps.
|
| 175 |
+
|
| 176 |
+
Example:
|
| 177 |
+
```python
|
| 178 |
+
from fastmcp.server.middleware.rate_limiting import SlidingWindowRateLimitingMiddleware
|
| 179 |
+
|
| 180 |
+
# Allow 100 requests per minute
|
| 181 |
+
rate_limiter = SlidingWindowRateLimitingMiddleware(
|
| 182 |
+
max_requests=100,
|
| 183 |
+
window_minutes=1
|
| 184 |
+
)
|
| 185 |
+
|
| 186 |
+
mcp = FastMCP("MyServer")
|
| 187 |
+
mcp.add_middleware(rate_limiter)
|
| 188 |
+
```
|
| 189 |
+
"""
|
| 190 |
+
|
| 191 |
+
def __init__(
|
| 192 |
+
self,
|
| 193 |
+
max_requests: int,
|
| 194 |
+
window_minutes: int = 1,
|
| 195 |
+
get_client_id: Callable[[MiddlewareContext], str] | None = None,
|
| 196 |
+
):
|
| 197 |
+
"""Initialize sliding window rate limiting middleware.
|
| 198 |
+
|
| 199 |
+
Args:
|
| 200 |
+
max_requests: Maximum requests allowed in the time window
|
| 201 |
+
window_minutes: Time window in minutes
|
| 202 |
+
get_client_id: Function to extract client ID from context
|
| 203 |
+
"""
|
| 204 |
+
self.max_requests = max_requests
|
| 205 |
+
self.window_seconds = window_minutes * 60
|
| 206 |
+
self.get_client_id = get_client_id
|
| 207 |
+
|
| 208 |
+
# Storage for rate limiters per client
|
| 209 |
+
self.limiters: dict[str, SlidingWindowRateLimiter] = defaultdict(
|
| 210 |
+
lambda: SlidingWindowRateLimiter(self.max_requests, self.window_seconds)
|
| 211 |
+
)
|
| 212 |
+
|
| 213 |
+
def _get_client_identifier(self, context: MiddlewareContext) -> str:
|
| 214 |
+
"""Get client identifier for rate limiting."""
|
| 215 |
+
if self.get_client_id:
|
| 216 |
+
return self.get_client_id(context)
|
| 217 |
+
return "global"
|
| 218 |
+
|
| 219 |
+
async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 220 |
+
"""Apply sliding window rate limiting to requests."""
|
| 221 |
+
client_id = self._get_client_identifier(context)
|
| 222 |
+
limiter = self.limiters[client_id]
|
| 223 |
+
|
| 224 |
+
allowed = await limiter.is_allowed()
|
| 225 |
+
if not allowed:
|
| 226 |
+
raise RateLimitError(
|
| 227 |
+
f"Rate limit exceeded: {self.max_requests} requests per "
|
| 228 |
+
f"{self.window_seconds // 60} minutes for client: {client_id}"
|
| 229 |
+
)
|
| 230 |
+
|
| 231 |
+
return await call_next(context)
|
src/fastmcp/server/middleware/timing.py
ADDED
|
@@ -0,0 +1,156 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Timing middleware for measuring and logging request performance."""
|
| 2 |
+
|
| 3 |
+
import logging
|
| 4 |
+
import time
|
| 5 |
+
from typing import Any
|
| 6 |
+
|
| 7 |
+
from .middleware import CallNext, Middleware, MiddlewareContext
|
| 8 |
+
|
| 9 |
+
|
| 10 |
+
class TimingMiddleware(Middleware):
|
| 11 |
+
"""Middleware that logs the execution time of requests.
|
| 12 |
+
|
| 13 |
+
Only measures and logs timing for request messages (not notifications).
|
| 14 |
+
Provides insights into performance characteristics of your MCP server.
|
| 15 |
+
|
| 16 |
+
Example:
|
| 17 |
+
```python
|
| 18 |
+
from fastmcp.server.middleware.timing import TimingMiddleware
|
| 19 |
+
|
| 20 |
+
mcp = FastMCP("MyServer")
|
| 21 |
+
mcp.add_middleware(TimingMiddleware())
|
| 22 |
+
|
| 23 |
+
# Now all requests will be timed and logged
|
| 24 |
+
```
|
| 25 |
+
"""
|
| 26 |
+
|
| 27 |
+
def __init__(
|
| 28 |
+
self, logger: logging.Logger | None = None, log_level: int = logging.INFO
|
| 29 |
+
):
|
| 30 |
+
"""Initialize timing middleware.
|
| 31 |
+
|
| 32 |
+
Args:
|
| 33 |
+
logger: Logger instance to use. If None, creates a logger named 'fastmcp.timing'
|
| 34 |
+
log_level: Log level for timing messages (default: INFO)
|
| 35 |
+
"""
|
| 36 |
+
self.logger = logger or logging.getLogger("fastmcp.timing")
|
| 37 |
+
self.log_level = log_level
|
| 38 |
+
|
| 39 |
+
async def on_request(self, context: MiddlewareContext, call_next: CallNext) -> Any:
|
| 40 |
+
"""Time request execution and log the results."""
|
| 41 |
+
method = context.method or "unknown"
|
| 42 |
+
|
| 43 |
+
start_time = time.perf_counter()
|
| 44 |
+
try:
|
| 45 |
+
result = await call_next(context)
|
| 46 |
+
duration_ms = (time.perf_counter() - start_time) * 1000
|
| 47 |
+
self.logger.log(
|
| 48 |
+
self.log_level, f"Request {method} completed in {duration_ms:.2f}ms"
|
| 49 |
+
)
|
| 50 |
+
return result
|
| 51 |
+
except Exception as e:
|
| 52 |
+
duration_ms = (time.perf_counter() - start_time) * 1000
|
| 53 |
+
self.logger.log(
|
| 54 |
+
self.log_level,
|
| 55 |
+
f"Request {method} failed after {duration_ms:.2f}ms: {e}",
|
| 56 |
+
)
|
| 57 |
+
raise
|
| 58 |
+
|
| 59 |
+
|
| 60 |
+
class DetailedTimingMiddleware(Middleware):
|
| 61 |
+
"""Enhanced timing middleware with per-operation breakdowns.
|
| 62 |
+
|
| 63 |
+
Provides detailed timing information for different types of MCP operations,
|
| 64 |
+
allowing you to identify performance bottlenecks in specific operations.
|
| 65 |
+
|
| 66 |
+
Example:
|
| 67 |
+
```python
|
| 68 |
+
from fastmcp.server.middleware.timing import DetailedTimingMiddleware
|
| 69 |
+
import logging
|
| 70 |
+
|
| 71 |
+
# Configure logging to see the output
|
| 72 |
+
logging.basicConfig(level=logging.INFO)
|
| 73 |
+
|
| 74 |
+
mcp = FastMCP("MyServer")
|
| 75 |
+
mcp.add_middleware(DetailedTimingMiddleware())
|
| 76 |
+
```
|
| 77 |
+
"""
|
| 78 |
+
|
| 79 |
+
def __init__(
|
| 80 |
+
self, logger: logging.Logger | None = None, log_level: int = logging.INFO
|
| 81 |
+
):
|
| 82 |
+
"""Initialize detailed timing middleware.
|
| 83 |
+
|
| 84 |
+
Args:
|
| 85 |
+
logger: Logger instance to use. If None, creates a logger named 'fastmcp.timing.detailed'
|
| 86 |
+
log_level: Log level for timing messages (default: INFO)
|
| 87 |
+
"""
|
| 88 |
+
self.logger = logger or logging.getLogger("fastmcp.timing.detailed")
|
| 89 |
+
self.log_level = log_level
|
| 90 |
+
|
| 91 |
+
async def _time_operation(
|
| 92 |
+
self, context: MiddlewareContext, call_next: CallNext, operation_name: str
|
| 93 |
+
) -> Any:
|
| 94 |
+
"""Helper method to time any operation."""
|
| 95 |
+
start_time = time.perf_counter()
|
| 96 |
+
try:
|
| 97 |
+
result = await call_next(context)
|
| 98 |
+
duration_ms = (time.perf_counter() - start_time) * 1000
|
| 99 |
+
self.logger.log(
|
| 100 |
+
self.log_level, f"{operation_name} completed in {duration_ms:.2f}ms"
|
| 101 |
+
)
|
| 102 |
+
return result
|
| 103 |
+
except Exception as e:
|
| 104 |
+
duration_ms = (time.perf_counter() - start_time) * 1000
|
| 105 |
+
self.logger.log(
|
| 106 |
+
self.log_level,
|
| 107 |
+
f"{operation_name} failed after {duration_ms:.2f}ms: {e}",
|
| 108 |
+
)
|
| 109 |
+
raise
|
| 110 |
+
|
| 111 |
+
async def on_call_tool(
|
| 112 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 113 |
+
) -> Any:
|
| 114 |
+
"""Time tool execution."""
|
| 115 |
+
tool_name = getattr(context.message, "name", "unknown")
|
| 116 |
+
return await self._time_operation(context, call_next, f"Tool '{tool_name}'")
|
| 117 |
+
|
| 118 |
+
async def on_read_resource(
|
| 119 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 120 |
+
) -> Any:
|
| 121 |
+
"""Time resource reading."""
|
| 122 |
+
resource_uri = getattr(context.message, "uri", "unknown")
|
| 123 |
+
return await self._time_operation(
|
| 124 |
+
context, call_next, f"Resource '{resource_uri}'"
|
| 125 |
+
)
|
| 126 |
+
|
| 127 |
+
async def on_get_prompt(
|
| 128 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 129 |
+
) -> Any:
|
| 130 |
+
"""Time prompt retrieval."""
|
| 131 |
+
prompt_name = getattr(context.message, "name", "unknown")
|
| 132 |
+
return await self._time_operation(context, call_next, f"Prompt '{prompt_name}'")
|
| 133 |
+
|
| 134 |
+
async def on_list_tools(
|
| 135 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 136 |
+
) -> Any:
|
| 137 |
+
"""Time tool listing."""
|
| 138 |
+
return await self._time_operation(context, call_next, "List tools")
|
| 139 |
+
|
| 140 |
+
async def on_list_resources(
|
| 141 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 142 |
+
) -> Any:
|
| 143 |
+
"""Time resource listing."""
|
| 144 |
+
return await self._time_operation(context, call_next, "List resources")
|
| 145 |
+
|
| 146 |
+
async def on_list_resource_templates(
|
| 147 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 148 |
+
) -> Any:
|
| 149 |
+
"""Time resource template listing."""
|
| 150 |
+
return await self._time_operation(context, call_next, "List resource templates")
|
| 151 |
+
|
| 152 |
+
async def on_list_prompts(
|
| 153 |
+
self, context: MiddlewareContext, call_next: CallNext
|
| 154 |
+
) -> Any:
|
| 155 |
+
"""Time prompt listing."""
|
| 156 |
+
return await self._time_operation(context, call_next, "List prompts")
|
tests/server/middleware/test_error_handling.py
ADDED
|
@@ -0,0 +1,601 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for error handling middleware."""
|
| 2 |
+
|
| 3 |
+
import logging
|
| 4 |
+
from unittest.mock import AsyncMock, MagicMock
|
| 5 |
+
|
| 6 |
+
import pytest
|
| 7 |
+
from mcp import McpError
|
| 8 |
+
|
| 9 |
+
from fastmcp.server.middleware.error_handling import (
|
| 10 |
+
ErrorHandlingMiddleware,
|
| 11 |
+
RetryMiddleware,
|
| 12 |
+
)
|
| 13 |
+
from fastmcp.server.middleware.middleware import MiddlewareContext
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
@pytest.fixture
|
| 17 |
+
def mock_context():
|
| 18 |
+
"""Create a mock middleware context."""
|
| 19 |
+
context = MagicMock(spec=MiddlewareContext)
|
| 20 |
+
context.method = "test_method"
|
| 21 |
+
return context
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
@pytest.fixture
|
| 25 |
+
def mock_call_next():
|
| 26 |
+
"""Create a mock call_next function."""
|
| 27 |
+
return AsyncMock(return_value="test_result")
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class TestErrorHandlingMiddleware:
|
| 31 |
+
"""Test error handling middleware functionality."""
|
| 32 |
+
|
| 33 |
+
def test_init_default(self):
|
| 34 |
+
"""Test default initialization."""
|
| 35 |
+
middleware = ErrorHandlingMiddleware()
|
| 36 |
+
assert middleware.logger.name == "fastmcp.errors"
|
| 37 |
+
assert middleware.include_traceback is False
|
| 38 |
+
assert middleware.error_callback is None
|
| 39 |
+
assert middleware.transform_errors is True
|
| 40 |
+
assert middleware.error_counts == {}
|
| 41 |
+
|
| 42 |
+
def test_init_custom(self):
|
| 43 |
+
"""Test custom initialization."""
|
| 44 |
+
logger = logging.getLogger("custom")
|
| 45 |
+
callback = MagicMock()
|
| 46 |
+
|
| 47 |
+
middleware = ErrorHandlingMiddleware(
|
| 48 |
+
logger=logger,
|
| 49 |
+
include_traceback=True,
|
| 50 |
+
error_callback=callback,
|
| 51 |
+
transform_errors=False,
|
| 52 |
+
)
|
| 53 |
+
assert middleware.logger is logger
|
| 54 |
+
assert middleware.include_traceback is True
|
| 55 |
+
assert middleware.error_callback is callback
|
| 56 |
+
assert middleware.transform_errors is False
|
| 57 |
+
|
| 58 |
+
def test_log_error_basic(self, mock_context, caplog):
|
| 59 |
+
"""Test basic error logging."""
|
| 60 |
+
middleware = ErrorHandlingMiddleware()
|
| 61 |
+
error = ValueError("test error")
|
| 62 |
+
|
| 63 |
+
with caplog.at_level(logging.ERROR):
|
| 64 |
+
middleware._log_error(error, mock_context)
|
| 65 |
+
|
| 66 |
+
assert "Error in test_method: ValueError: test error" in caplog.text
|
| 67 |
+
assert "ValueError:test_method" in middleware.error_counts
|
| 68 |
+
assert middleware.error_counts["ValueError:test_method"] == 1
|
| 69 |
+
|
| 70 |
+
def test_log_error_with_traceback(self, mock_context, caplog):
|
| 71 |
+
"""Test error logging with traceback."""
|
| 72 |
+
middleware = ErrorHandlingMiddleware(include_traceback=True)
|
| 73 |
+
error = ValueError("test error")
|
| 74 |
+
|
| 75 |
+
with caplog.at_level(logging.ERROR):
|
| 76 |
+
middleware._log_error(error, mock_context)
|
| 77 |
+
|
| 78 |
+
assert "Error in test_method: ValueError: test error" in caplog.text
|
| 79 |
+
# The traceback is added to the log message
|
| 80 |
+
assert "Error in test_method: ValueError: test error" in caplog.text
|
| 81 |
+
|
| 82 |
+
def test_log_error_with_callback(self, mock_context):
|
| 83 |
+
"""Test error logging with callback."""
|
| 84 |
+
callback = MagicMock()
|
| 85 |
+
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
| 86 |
+
error = ValueError("test error")
|
| 87 |
+
|
| 88 |
+
middleware._log_error(error, mock_context)
|
| 89 |
+
|
| 90 |
+
callback.assert_called_once_with(error, mock_context)
|
| 91 |
+
|
| 92 |
+
def test_log_error_callback_exception(self, mock_context, caplog):
|
| 93 |
+
"""Test error logging when callback raises exception."""
|
| 94 |
+
callback = MagicMock(side_effect=RuntimeError("callback error"))
|
| 95 |
+
middleware = ErrorHandlingMiddleware(error_callback=callback)
|
| 96 |
+
error = ValueError("test error")
|
| 97 |
+
|
| 98 |
+
with caplog.at_level(logging.ERROR):
|
| 99 |
+
middleware._log_error(error, mock_context)
|
| 100 |
+
|
| 101 |
+
assert "Error in error callback: callback error" in caplog.text
|
| 102 |
+
|
| 103 |
+
def test_transform_error_mcp_error(self):
|
| 104 |
+
"""Test that MCP errors are not transformed."""
|
| 105 |
+
middleware = ErrorHandlingMiddleware()
|
| 106 |
+
from mcp.types import ErrorData
|
| 107 |
+
|
| 108 |
+
error = McpError(ErrorData(code=-32001, message="test error"))
|
| 109 |
+
|
| 110 |
+
result = middleware._transform_error(error)
|
| 111 |
+
|
| 112 |
+
assert result is error
|
| 113 |
+
|
| 114 |
+
def test_transform_error_disabled(self):
|
| 115 |
+
"""Test error transformation when disabled."""
|
| 116 |
+
middleware = ErrorHandlingMiddleware(transform_errors=False)
|
| 117 |
+
error = ValueError("test error")
|
| 118 |
+
|
| 119 |
+
result = middleware._transform_error(error)
|
| 120 |
+
|
| 121 |
+
assert result is error
|
| 122 |
+
|
| 123 |
+
def test_transform_error_value_error(self):
|
| 124 |
+
"""Test transforming ValueError."""
|
| 125 |
+
middleware = ErrorHandlingMiddleware()
|
| 126 |
+
error = ValueError("test error")
|
| 127 |
+
|
| 128 |
+
result = middleware._transform_error(error)
|
| 129 |
+
|
| 130 |
+
assert isinstance(result, McpError)
|
| 131 |
+
assert result.error.code == -32602
|
| 132 |
+
assert "Invalid params: test error" in result.error.message
|
| 133 |
+
|
| 134 |
+
def test_transform_error_file_not_found(self):
|
| 135 |
+
"""Test transforming FileNotFoundError."""
|
| 136 |
+
middleware = ErrorHandlingMiddleware()
|
| 137 |
+
error = FileNotFoundError("test error")
|
| 138 |
+
|
| 139 |
+
result = middleware._transform_error(error)
|
| 140 |
+
|
| 141 |
+
assert isinstance(result, McpError)
|
| 142 |
+
assert result.error.code == -32001
|
| 143 |
+
assert "Resource not found: test error" in result.error.message
|
| 144 |
+
|
| 145 |
+
def test_transform_error_permission_error(self):
|
| 146 |
+
"""Test transforming PermissionError."""
|
| 147 |
+
middleware = ErrorHandlingMiddleware()
|
| 148 |
+
error = PermissionError("test error")
|
| 149 |
+
|
| 150 |
+
result = middleware._transform_error(error)
|
| 151 |
+
|
| 152 |
+
assert isinstance(result, McpError)
|
| 153 |
+
assert result.error.code == -32000
|
| 154 |
+
assert "Permission denied: test error" in result.error.message
|
| 155 |
+
|
| 156 |
+
def test_transform_error_timeout_error(self):
|
| 157 |
+
"""Test transforming TimeoutError."""
|
| 158 |
+
middleware = ErrorHandlingMiddleware()
|
| 159 |
+
error = TimeoutError("test error")
|
| 160 |
+
|
| 161 |
+
result = middleware._transform_error(error)
|
| 162 |
+
|
| 163 |
+
assert isinstance(result, McpError)
|
| 164 |
+
assert result.error.code == -32000
|
| 165 |
+
assert "Request timeout: test error" in result.error.message
|
| 166 |
+
|
| 167 |
+
def test_transform_error_generic(self):
|
| 168 |
+
"""Test transforming generic error."""
|
| 169 |
+
middleware = ErrorHandlingMiddleware()
|
| 170 |
+
error = RuntimeError("test error")
|
| 171 |
+
|
| 172 |
+
result = middleware._transform_error(error)
|
| 173 |
+
|
| 174 |
+
assert isinstance(result, McpError)
|
| 175 |
+
assert result.error.code == -32603
|
| 176 |
+
assert "Internal error: test error" in result.error.message
|
| 177 |
+
|
| 178 |
+
async def test_on_message_success(self, mock_context, mock_call_next):
|
| 179 |
+
"""Test successful message handling."""
|
| 180 |
+
middleware = ErrorHandlingMiddleware()
|
| 181 |
+
|
| 182 |
+
result = await middleware.on_message(mock_context, mock_call_next)
|
| 183 |
+
|
| 184 |
+
assert result == "test_result"
|
| 185 |
+
assert mock_call_next.called
|
| 186 |
+
|
| 187 |
+
async def test_on_message_error_transform(self, mock_context, caplog):
|
| 188 |
+
"""Test error handling with transformation."""
|
| 189 |
+
middleware = ErrorHandlingMiddleware()
|
| 190 |
+
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
| 191 |
+
|
| 192 |
+
with caplog.at_level(logging.ERROR):
|
| 193 |
+
with pytest.raises(McpError) as exc_info:
|
| 194 |
+
await middleware.on_message(mock_context, mock_call_next)
|
| 195 |
+
|
| 196 |
+
assert exc_info.value.error.code == -32602
|
| 197 |
+
assert "Invalid params: test error" in exc_info.value.error.message
|
| 198 |
+
assert "Error in test_method: ValueError: test error" in caplog.text
|
| 199 |
+
|
| 200 |
+
def test_get_error_stats(self, mock_context):
|
| 201 |
+
"""Test getting error statistics."""
|
| 202 |
+
middleware = ErrorHandlingMiddleware()
|
| 203 |
+
error1 = ValueError("error1")
|
| 204 |
+
error2 = ValueError("error2")
|
| 205 |
+
error3 = RuntimeError("error3")
|
| 206 |
+
|
| 207 |
+
middleware._log_error(error1, mock_context)
|
| 208 |
+
middleware._log_error(error2, mock_context)
|
| 209 |
+
middleware._log_error(error3, mock_context)
|
| 210 |
+
|
| 211 |
+
stats = middleware.get_error_stats()
|
| 212 |
+
assert stats["ValueError:test_method"] == 2
|
| 213 |
+
assert stats["RuntimeError:test_method"] == 1
|
| 214 |
+
|
| 215 |
+
|
| 216 |
+
class TestRetryMiddleware:
|
| 217 |
+
"""Test retry middleware functionality."""
|
| 218 |
+
|
| 219 |
+
def test_init_default(self):
|
| 220 |
+
"""Test default initialization."""
|
| 221 |
+
middleware = RetryMiddleware()
|
| 222 |
+
assert middleware.max_retries == 3
|
| 223 |
+
assert middleware.base_delay == 1.0
|
| 224 |
+
assert middleware.max_delay == 60.0
|
| 225 |
+
assert middleware.backoff_multiplier == 2.0
|
| 226 |
+
assert middleware.retry_exceptions == (ConnectionError, TimeoutError)
|
| 227 |
+
assert middleware.logger.name == "fastmcp.retry"
|
| 228 |
+
|
| 229 |
+
def test_init_custom(self):
|
| 230 |
+
"""Test custom initialization."""
|
| 231 |
+
logger = logging.getLogger("custom")
|
| 232 |
+
middleware = RetryMiddleware(
|
| 233 |
+
max_retries=5,
|
| 234 |
+
base_delay=2.0,
|
| 235 |
+
max_delay=120.0,
|
| 236 |
+
backoff_multiplier=3.0,
|
| 237 |
+
retry_exceptions=(ValueError, RuntimeError),
|
| 238 |
+
logger=logger,
|
| 239 |
+
)
|
| 240 |
+
assert middleware.max_retries == 5
|
| 241 |
+
assert middleware.base_delay == 2.0
|
| 242 |
+
assert middleware.max_delay == 120.0
|
| 243 |
+
assert middleware.backoff_multiplier == 3.0
|
| 244 |
+
assert middleware.retry_exceptions == (ValueError, RuntimeError)
|
| 245 |
+
assert middleware.logger is logger
|
| 246 |
+
|
| 247 |
+
def test_should_retry_true(self):
|
| 248 |
+
"""Test retry decision for retryable errors."""
|
| 249 |
+
middleware = RetryMiddleware()
|
| 250 |
+
|
| 251 |
+
assert middleware._should_retry(ConnectionError()) is True
|
| 252 |
+
assert middleware._should_retry(TimeoutError()) is True
|
| 253 |
+
|
| 254 |
+
def test_should_retry_false(self):
|
| 255 |
+
"""Test retry decision for non-retryable errors."""
|
| 256 |
+
middleware = RetryMiddleware()
|
| 257 |
+
|
| 258 |
+
assert middleware._should_retry(ValueError()) is False
|
| 259 |
+
assert middleware._should_retry(RuntimeError()) is False
|
| 260 |
+
|
| 261 |
+
def test_calculate_delay(self):
|
| 262 |
+
"""Test delay calculation."""
|
| 263 |
+
middleware = RetryMiddleware(
|
| 264 |
+
base_delay=1.0, backoff_multiplier=2.0, max_delay=10.0
|
| 265 |
+
)
|
| 266 |
+
|
| 267 |
+
assert middleware._calculate_delay(0) == 1.0
|
| 268 |
+
assert middleware._calculate_delay(1) == 2.0
|
| 269 |
+
assert middleware._calculate_delay(2) == 4.0
|
| 270 |
+
assert middleware._calculate_delay(3) == 8.0
|
| 271 |
+
assert middleware._calculate_delay(4) == 10.0 # capped at max_delay
|
| 272 |
+
|
| 273 |
+
async def test_on_request_success_first_try(self, mock_context, mock_call_next):
|
| 274 |
+
"""Test successful request on first try."""
|
| 275 |
+
middleware = RetryMiddleware()
|
| 276 |
+
|
| 277 |
+
result = await middleware.on_request(mock_context, mock_call_next)
|
| 278 |
+
|
| 279 |
+
assert result == "test_result"
|
| 280 |
+
assert mock_call_next.call_count == 1
|
| 281 |
+
|
| 282 |
+
async def test_on_request_success_after_retries(self, mock_context, caplog):
|
| 283 |
+
"""Test successful request after retries."""
|
| 284 |
+
middleware = RetryMiddleware(base_delay=0.01) # Fast retry for testing
|
| 285 |
+
|
| 286 |
+
# Fail first two attempts, succeed on third
|
| 287 |
+
mock_call_next = AsyncMock(
|
| 288 |
+
side_effect=[
|
| 289 |
+
ConnectionError("connection failed"),
|
| 290 |
+
ConnectionError("connection failed"),
|
| 291 |
+
"test_result",
|
| 292 |
+
]
|
| 293 |
+
)
|
| 294 |
+
|
| 295 |
+
with caplog.at_level(logging.WARNING):
|
| 296 |
+
result = await middleware.on_request(mock_context, mock_call_next)
|
| 297 |
+
|
| 298 |
+
assert result == "test_result"
|
| 299 |
+
assert mock_call_next.call_count == 3
|
| 300 |
+
assert "Retrying in" in caplog.text
|
| 301 |
+
|
| 302 |
+
async def test_on_request_max_retries_exceeded(self, mock_context, caplog):
|
| 303 |
+
"""Test request failing after max retries."""
|
| 304 |
+
middleware = RetryMiddleware(max_retries=2, base_delay=0.01)
|
| 305 |
+
|
| 306 |
+
# Fail all attempts
|
| 307 |
+
mock_call_next = AsyncMock(side_effect=ConnectionError("connection failed"))
|
| 308 |
+
|
| 309 |
+
with caplog.at_level(logging.WARNING):
|
| 310 |
+
with pytest.raises(ConnectionError):
|
| 311 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 312 |
+
|
| 313 |
+
assert mock_call_next.call_count == 3 # initial + 2 retries
|
| 314 |
+
assert "Retrying in" in caplog.text
|
| 315 |
+
|
| 316 |
+
async def test_on_request_non_retryable_error(self, mock_context):
|
| 317 |
+
"""Test non-retryable error is not retried."""
|
| 318 |
+
middleware = RetryMiddleware()
|
| 319 |
+
mock_call_next = AsyncMock(side_effect=ValueError("non-retryable"))
|
| 320 |
+
|
| 321 |
+
with pytest.raises(ValueError):
|
| 322 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 323 |
+
|
| 324 |
+
assert mock_call_next.call_count == 1 # No retries
|
| 325 |
+
|
| 326 |
+
|
| 327 |
+
@pytest.fixture
|
| 328 |
+
def error_handling_server():
|
| 329 |
+
"""Create a FastMCP server specifically for error handling middleware tests."""
|
| 330 |
+
from fastmcp import FastMCP
|
| 331 |
+
|
| 332 |
+
mcp = FastMCP("ErrorHandlingTestServer")
|
| 333 |
+
|
| 334 |
+
@mcp.tool
|
| 335 |
+
def reliable_operation(data: str) -> str:
|
| 336 |
+
"""A reliable operation that always succeeds."""
|
| 337 |
+
return f"Success: {data}"
|
| 338 |
+
|
| 339 |
+
@mcp.tool
|
| 340 |
+
def failing_operation(error_type: str = "value") -> str:
|
| 341 |
+
"""An operation that fails with different error types."""
|
| 342 |
+
if error_type == "value":
|
| 343 |
+
raise ValueError("Value error occurred")
|
| 344 |
+
elif error_type == "file":
|
| 345 |
+
raise FileNotFoundError("File not found")
|
| 346 |
+
elif error_type == "permission":
|
| 347 |
+
raise PermissionError("Permission denied")
|
| 348 |
+
elif error_type == "timeout":
|
| 349 |
+
raise TimeoutError("Operation timed out")
|
| 350 |
+
elif error_type == "generic":
|
| 351 |
+
raise RuntimeError("Generic runtime error")
|
| 352 |
+
else:
|
| 353 |
+
return "Operation completed"
|
| 354 |
+
|
| 355 |
+
@mcp.tool
|
| 356 |
+
def intermittent_operation(fail_rate: float = 0.5) -> str:
|
| 357 |
+
"""An operation that fails intermittently."""
|
| 358 |
+
import random
|
| 359 |
+
|
| 360 |
+
if random.random() < fail_rate:
|
| 361 |
+
raise ConnectionError("Random connection failure")
|
| 362 |
+
return "Operation succeeded"
|
| 363 |
+
|
| 364 |
+
@mcp.tool
|
| 365 |
+
def retryable_operation(attempt_count: int = 0) -> str:
|
| 366 |
+
"""An operation that succeeds after a few attempts."""
|
| 367 |
+
# This is a simple way to simulate retry behavior
|
| 368 |
+
# In a real scenario, you might use external state
|
| 369 |
+
if attempt_count < 2:
|
| 370 |
+
raise ConnectionError("Temporary connection error")
|
| 371 |
+
return "Operation succeeded after retries"
|
| 372 |
+
|
| 373 |
+
return mcp
|
| 374 |
+
|
| 375 |
+
|
| 376 |
+
class TestErrorHandlingMiddlewareIntegration:
|
| 377 |
+
"""Integration tests for error handling middleware with real FastMCP server."""
|
| 378 |
+
|
| 379 |
+
async def test_error_handling_middleware_logs_real_errors(
|
| 380 |
+
self, error_handling_server, caplog
|
| 381 |
+
):
|
| 382 |
+
"""Test that error handling middleware logs real errors from tools."""
|
| 383 |
+
from fastmcp.client import Client
|
| 384 |
+
|
| 385 |
+
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
| 386 |
+
|
| 387 |
+
with caplog.at_level(logging.ERROR):
|
| 388 |
+
async with Client(error_handling_server) as client:
|
| 389 |
+
# Test different types of errors
|
| 390 |
+
with pytest.raises(Exception):
|
| 391 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 392 |
+
|
| 393 |
+
with pytest.raises(Exception):
|
| 394 |
+
await client.call_tool("failing_operation", {"error_type": "file"})
|
| 395 |
+
|
| 396 |
+
log_text = caplog.text
|
| 397 |
+
|
| 398 |
+
# Should have error logs for both failures
|
| 399 |
+
assert "Error in tools/call: ToolError:" in log_text
|
| 400 |
+
# Should have captured both error instances
|
| 401 |
+
error_count = log_text.count("Error in tools/call:")
|
| 402 |
+
assert error_count == 2
|
| 403 |
+
|
| 404 |
+
async def test_error_handling_middleware_tracks_error_statistics(
|
| 405 |
+
self, error_handling_server
|
| 406 |
+
):
|
| 407 |
+
"""Test that error handling middleware accurately tracks error statistics."""
|
| 408 |
+
from fastmcp.client import Client
|
| 409 |
+
|
| 410 |
+
error_middleware = ErrorHandlingMiddleware()
|
| 411 |
+
error_handling_server.add_middleware(error_middleware)
|
| 412 |
+
|
| 413 |
+
async with Client(error_handling_server) as client:
|
| 414 |
+
# Generate different types of errors
|
| 415 |
+
for _ in range(3):
|
| 416 |
+
with pytest.raises(Exception):
|
| 417 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 418 |
+
|
| 419 |
+
for _ in range(2):
|
| 420 |
+
with pytest.raises(Exception):
|
| 421 |
+
await client.call_tool("failing_operation", {"error_type": "file"})
|
| 422 |
+
|
| 423 |
+
# Try some intermittent operations (some may succeed)
|
| 424 |
+
for _ in range(5):
|
| 425 |
+
try:
|
| 426 |
+
await client.call_tool("intermittent_operation", {"fail_rate": 0.8})
|
| 427 |
+
except Exception:
|
| 428 |
+
pass # Expected failures
|
| 429 |
+
|
| 430 |
+
# Check error statistics
|
| 431 |
+
stats = error_middleware.get_error_stats()
|
| 432 |
+
|
| 433 |
+
# Should have tracked the ToolError wrapper
|
| 434 |
+
assert "ToolError:tools/call" in stats
|
| 435 |
+
assert stats["ToolError:tools/call"] >= 5 # At least the 5 deliberate failures
|
| 436 |
+
|
| 437 |
+
async def test_error_handling_middleware_with_success_and_failure(
|
| 438 |
+
self, error_handling_server, caplog
|
| 439 |
+
):
|
| 440 |
+
"""Test error handling middleware with mix of successful and failed operations."""
|
| 441 |
+
from fastmcp.client import Client
|
| 442 |
+
|
| 443 |
+
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
| 444 |
+
|
| 445 |
+
with caplog.at_level(logging.ERROR):
|
| 446 |
+
async with Client(error_handling_server) as client:
|
| 447 |
+
# Successful operation (should not generate error logs)
|
| 448 |
+
await client.call_tool("reliable_operation", {"data": "test"})
|
| 449 |
+
|
| 450 |
+
# Failed operation (should generate error log)
|
| 451 |
+
with pytest.raises(Exception):
|
| 452 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 453 |
+
|
| 454 |
+
# Another successful operation
|
| 455 |
+
await client.call_tool("reliable_operation", {"data": "test2"})
|
| 456 |
+
|
| 457 |
+
log_text = caplog.text
|
| 458 |
+
|
| 459 |
+
# Should only have one error log (for the failed operation)
|
| 460 |
+
error_count = log_text.count("Error in tools/call:")
|
| 461 |
+
assert error_count == 1
|
| 462 |
+
|
| 463 |
+
async def test_error_handling_middleware_custom_callback(
|
| 464 |
+
self, error_handling_server
|
| 465 |
+
):
|
| 466 |
+
"""Test error handling middleware with custom error callback."""
|
| 467 |
+
from fastmcp.client import Client
|
| 468 |
+
|
| 469 |
+
captured_errors = []
|
| 470 |
+
|
| 471 |
+
def error_callback(error, context):
|
| 472 |
+
captured_errors.append(
|
| 473 |
+
{
|
| 474 |
+
"error_type": type(error).__name__,
|
| 475 |
+
"message": str(error),
|
| 476 |
+
"method": context.method,
|
| 477 |
+
}
|
| 478 |
+
)
|
| 479 |
+
|
| 480 |
+
error_handling_server.add_middleware(
|
| 481 |
+
ErrorHandlingMiddleware(error_callback=error_callback)
|
| 482 |
+
)
|
| 483 |
+
|
| 484 |
+
async with Client(error_handling_server) as client:
|
| 485 |
+
# Generate some errors
|
| 486 |
+
with pytest.raises(Exception):
|
| 487 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 488 |
+
|
| 489 |
+
with pytest.raises(Exception):
|
| 490 |
+
await client.call_tool("failing_operation", {"error_type": "timeout"})
|
| 491 |
+
|
| 492 |
+
# Check that callback was called
|
| 493 |
+
assert len(captured_errors) == 2
|
| 494 |
+
assert captured_errors[0]["error_type"] == "ToolError"
|
| 495 |
+
assert captured_errors[1]["error_type"] == "ToolError"
|
| 496 |
+
assert all(error["method"] == "tools/call" for error in captured_errors)
|
| 497 |
+
|
| 498 |
+
async def test_error_handling_middleware_transform_errors(
|
| 499 |
+
self, error_handling_server
|
| 500 |
+
):
|
| 501 |
+
"""Test error transformation functionality."""
|
| 502 |
+
from fastmcp.client import Client
|
| 503 |
+
|
| 504 |
+
error_handling_server.add_middleware(
|
| 505 |
+
ErrorHandlingMiddleware(transform_errors=True)
|
| 506 |
+
)
|
| 507 |
+
|
| 508 |
+
async with Client(error_handling_server) as client:
|
| 509 |
+
# All errors should still be raised, but potentially transformed
|
| 510 |
+
with pytest.raises(Exception) as exc_info:
|
| 511 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 512 |
+
|
| 513 |
+
# Error should still exist (may be wrapped by FastMCP)
|
| 514 |
+
assert exc_info.value is not None
|
| 515 |
+
|
| 516 |
+
|
| 517 |
+
class TestRetryMiddlewareIntegration:
|
| 518 |
+
"""Integration tests for retry middleware with real FastMCP server."""
|
| 519 |
+
|
| 520 |
+
async def test_retry_middleware_with_transient_failures(
|
| 521 |
+
self, error_handling_server, caplog
|
| 522 |
+
):
|
| 523 |
+
"""Test retry middleware with operations that have transient failures."""
|
| 524 |
+
from fastmcp.client import Client
|
| 525 |
+
|
| 526 |
+
# Configure retry middleware to retry connection errors
|
| 527 |
+
error_handling_server.add_middleware(
|
| 528 |
+
RetryMiddleware(
|
| 529 |
+
max_retries=3,
|
| 530 |
+
base_delay=0.01, # Very short delay for testing
|
| 531 |
+
retry_exceptions=(ConnectionError,),
|
| 532 |
+
)
|
| 533 |
+
)
|
| 534 |
+
|
| 535 |
+
with caplog.at_level(logging.WARNING):
|
| 536 |
+
async with Client(error_handling_server) as client:
|
| 537 |
+
# This operation fails intermittently - try several times
|
| 538 |
+
success_count = 0
|
| 539 |
+
for _ in range(5):
|
| 540 |
+
try:
|
| 541 |
+
await client.call_tool(
|
| 542 |
+
"intermittent_operation", {"fail_rate": 0.7}
|
| 543 |
+
)
|
| 544 |
+
success_count += 1
|
| 545 |
+
except Exception:
|
| 546 |
+
pass # Some failures expected even with retries
|
| 547 |
+
|
| 548 |
+
# Should have some retry log messages
|
| 549 |
+
# Note: Retry logs might not appear if the underlying errors are wrapped by FastMCP
|
| 550 |
+
# The key is that some operations should succeed due to retries
|
| 551 |
+
|
| 552 |
+
async def test_retry_middleware_with_permanent_failures(
|
| 553 |
+
self, error_handling_server
|
| 554 |
+
):
|
| 555 |
+
"""Test that retry middleware doesn't retry non-retryable errors."""
|
| 556 |
+
from fastmcp.client import Client
|
| 557 |
+
|
| 558 |
+
# Configure retry middleware for connection errors only
|
| 559 |
+
error_handling_server.add_middleware(
|
| 560 |
+
RetryMiddleware(
|
| 561 |
+
max_retries=3, base_delay=0.01, retry_exceptions=(ConnectionError,)
|
| 562 |
+
)
|
| 563 |
+
)
|
| 564 |
+
|
| 565 |
+
async with Client(error_handling_server) as client:
|
| 566 |
+
# Value errors should not be retried
|
| 567 |
+
with pytest.raises(Exception):
|
| 568 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 569 |
+
|
| 570 |
+
# Should fail immediately without retries
|
| 571 |
+
|
| 572 |
+
async def test_combined_error_handling_and_retry_middleware(
|
| 573 |
+
self, error_handling_server, caplog
|
| 574 |
+
):
|
| 575 |
+
"""Test error handling and retry middleware working together."""
|
| 576 |
+
from fastmcp.client import Client
|
| 577 |
+
|
| 578 |
+
# Add both middleware
|
| 579 |
+
error_handling_server.add_middleware(ErrorHandlingMiddleware())
|
| 580 |
+
error_handling_server.add_middleware(
|
| 581 |
+
RetryMiddleware(
|
| 582 |
+
max_retries=2, base_delay=0.01, retry_exceptions=(ConnectionError,)
|
| 583 |
+
)
|
| 584 |
+
)
|
| 585 |
+
|
| 586 |
+
with caplog.at_level(logging.ERROR):
|
| 587 |
+
async with Client(error_handling_server) as client:
|
| 588 |
+
# Try intermittent operation
|
| 589 |
+
try:
|
| 590 |
+
await client.call_tool("intermittent_operation", {"fail_rate": 0.9})
|
| 591 |
+
except Exception:
|
| 592 |
+
pass # May still fail even with retries
|
| 593 |
+
|
| 594 |
+
# Try permanent failure
|
| 595 |
+
with pytest.raises(Exception):
|
| 596 |
+
await client.call_tool("failing_operation", {"error_type": "value"})
|
| 597 |
+
|
| 598 |
+
log_text = caplog.text
|
| 599 |
+
|
| 600 |
+
# Should have error logs from error handling middleware
|
| 601 |
+
assert "Error in tools/call:" in log_text
|
tests/server/middleware/test_logging.py
ADDED
|
@@ -0,0 +1,426 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for logging middleware."""
|
| 2 |
+
|
| 3 |
+
import json
|
| 4 |
+
import logging
|
| 5 |
+
from unittest.mock import AsyncMock, MagicMock
|
| 6 |
+
|
| 7 |
+
import pytest
|
| 8 |
+
|
| 9 |
+
from fastmcp.server.middleware.logging import (
|
| 10 |
+
LoggingMiddleware,
|
| 11 |
+
StructuredLoggingMiddleware,
|
| 12 |
+
)
|
| 13 |
+
from fastmcp.server.middleware.middleware import MiddlewareContext
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
@pytest.fixture
|
| 17 |
+
def mock_context():
|
| 18 |
+
"""Create a mock middleware context."""
|
| 19 |
+
context = MagicMock(spec=MiddlewareContext)
|
| 20 |
+
context.method = "test_method"
|
| 21 |
+
context.source = "client"
|
| 22 |
+
context.type = "request"
|
| 23 |
+
context.message = MagicMock()
|
| 24 |
+
context.message.__dict__ = {"param": "value"}
|
| 25 |
+
context.timestamp = MagicMock()
|
| 26 |
+
context.timestamp.isoformat.return_value = "2023-01-01T00:00:00Z"
|
| 27 |
+
return context
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
@pytest.fixture
|
| 31 |
+
def mock_call_next():
|
| 32 |
+
"""Create a mock call_next function."""
|
| 33 |
+
return AsyncMock(return_value="test_result")
|
| 34 |
+
|
| 35 |
+
|
| 36 |
+
class TestLoggingMiddleware:
|
| 37 |
+
"""Test logging middleware functionality."""
|
| 38 |
+
|
| 39 |
+
def test_init_default(self):
|
| 40 |
+
"""Test default initialization."""
|
| 41 |
+
middleware = LoggingMiddleware()
|
| 42 |
+
assert middleware.logger.name == "fastmcp.requests"
|
| 43 |
+
assert middleware.log_level == logging.INFO
|
| 44 |
+
assert middleware.include_payloads is False
|
| 45 |
+
assert middleware.max_payload_length == 1000
|
| 46 |
+
|
| 47 |
+
def test_init_custom(self):
|
| 48 |
+
"""Test custom initialization."""
|
| 49 |
+
logger = logging.getLogger("custom")
|
| 50 |
+
middleware = LoggingMiddleware(
|
| 51 |
+
logger=logger,
|
| 52 |
+
log_level=logging.DEBUG,
|
| 53 |
+
include_payloads=True,
|
| 54 |
+
max_payload_length=500,
|
| 55 |
+
)
|
| 56 |
+
assert middleware.logger is logger
|
| 57 |
+
assert middleware.log_level == logging.DEBUG
|
| 58 |
+
assert middleware.include_payloads is True
|
| 59 |
+
assert middleware.max_payload_length == 500
|
| 60 |
+
|
| 61 |
+
def test_format_message_without_payloads(self, mock_context):
|
| 62 |
+
"""Test message formatting without payloads."""
|
| 63 |
+
middleware = LoggingMiddleware()
|
| 64 |
+
formatted = middleware._format_message(mock_context)
|
| 65 |
+
|
| 66 |
+
assert "source=client" in formatted
|
| 67 |
+
assert "type=request" in formatted
|
| 68 |
+
assert "method=test_method" in formatted
|
| 69 |
+
assert "payload=" not in formatted
|
| 70 |
+
|
| 71 |
+
def test_format_message_with_payloads(self, mock_context):
|
| 72 |
+
"""Test message formatting with payloads."""
|
| 73 |
+
middleware = LoggingMiddleware(include_payloads=True)
|
| 74 |
+
formatted = middleware._format_message(mock_context)
|
| 75 |
+
|
| 76 |
+
assert "source=client" in formatted
|
| 77 |
+
assert "type=request" in formatted
|
| 78 |
+
assert "method=test_method" in formatted
|
| 79 |
+
assert 'payload={"param": "value"}' in formatted
|
| 80 |
+
|
| 81 |
+
def test_format_message_long_payload(self, mock_context):
|
| 82 |
+
"""Test message formatting with long payload truncation."""
|
| 83 |
+
middleware = LoggingMiddleware(include_payloads=True, max_payload_length=10)
|
| 84 |
+
formatted = middleware._format_message(mock_context)
|
| 85 |
+
|
| 86 |
+
assert "payload=" in formatted
|
| 87 |
+
assert "..." in formatted
|
| 88 |
+
|
| 89 |
+
async def test_on_message_success(self, mock_context, mock_call_next, caplog):
|
| 90 |
+
"""Test logging successful messages."""
|
| 91 |
+
middleware = LoggingMiddleware()
|
| 92 |
+
|
| 93 |
+
with caplog.at_level(logging.INFO):
|
| 94 |
+
result = await middleware.on_message(mock_context, mock_call_next)
|
| 95 |
+
|
| 96 |
+
assert result == "test_result"
|
| 97 |
+
assert mock_call_next.called
|
| 98 |
+
assert "Processing message:" in caplog.text
|
| 99 |
+
assert "Completed message: test_method" in caplog.text
|
| 100 |
+
|
| 101 |
+
async def test_on_message_failure(self, mock_context, caplog):
|
| 102 |
+
"""Test logging failed messages."""
|
| 103 |
+
middleware = LoggingMiddleware()
|
| 104 |
+
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
| 105 |
+
|
| 106 |
+
with caplog.at_level(logging.INFO):
|
| 107 |
+
with pytest.raises(ValueError):
|
| 108 |
+
await middleware.on_message(mock_context, mock_call_next)
|
| 109 |
+
|
| 110 |
+
assert "Processing message:" in caplog.text
|
| 111 |
+
assert "Failed message: test_method - test error" in caplog.text
|
| 112 |
+
|
| 113 |
+
|
| 114 |
+
class TestStructuredLoggingMiddleware:
|
| 115 |
+
"""Test structured logging middleware functionality."""
|
| 116 |
+
|
| 117 |
+
def test_init_default(self):
|
| 118 |
+
"""Test default initialization."""
|
| 119 |
+
middleware = StructuredLoggingMiddleware()
|
| 120 |
+
assert middleware.logger.name == "fastmcp.structured"
|
| 121 |
+
assert middleware.log_level == logging.INFO
|
| 122 |
+
assert middleware.include_payloads is False
|
| 123 |
+
|
| 124 |
+
def test_create_log_entry_basic(self, mock_context):
|
| 125 |
+
"""Test creating basic log entry."""
|
| 126 |
+
middleware = StructuredLoggingMiddleware()
|
| 127 |
+
entry = middleware._create_log_entry(mock_context, "test_event")
|
| 128 |
+
|
| 129 |
+
assert entry["event"] == "test_event"
|
| 130 |
+
assert entry["timestamp"] == "2023-01-01T00:00:00Z"
|
| 131 |
+
assert entry["source"] == "client"
|
| 132 |
+
assert entry["type"] == "request"
|
| 133 |
+
assert entry["method"] == "test_method"
|
| 134 |
+
assert "payload" not in entry
|
| 135 |
+
|
| 136 |
+
def test_create_log_entry_with_payload(self, mock_context):
|
| 137 |
+
"""Test creating log entry with payload."""
|
| 138 |
+
middleware = StructuredLoggingMiddleware(include_payloads=True)
|
| 139 |
+
entry = middleware._create_log_entry(mock_context, "test_event")
|
| 140 |
+
|
| 141 |
+
assert entry["payload"] == {"param": "value"}
|
| 142 |
+
|
| 143 |
+
def test_create_log_entry_with_extra_fields(self, mock_context):
|
| 144 |
+
"""Test creating log entry with extra fields."""
|
| 145 |
+
middleware = StructuredLoggingMiddleware()
|
| 146 |
+
entry = middleware._create_log_entry(
|
| 147 |
+
mock_context, "test_event", extra_field="extra_value"
|
| 148 |
+
)
|
| 149 |
+
|
| 150 |
+
assert entry["extra_field"] == "extra_value"
|
| 151 |
+
|
| 152 |
+
async def test_on_message_success(self, mock_context, mock_call_next, caplog):
|
| 153 |
+
"""Test structured logging of successful messages."""
|
| 154 |
+
middleware = StructuredLoggingMiddleware()
|
| 155 |
+
|
| 156 |
+
with caplog.at_level(logging.INFO):
|
| 157 |
+
result = await middleware.on_message(mock_context, mock_call_next)
|
| 158 |
+
|
| 159 |
+
assert result == "test_result"
|
| 160 |
+
|
| 161 |
+
# Check that we have structured JSON logs
|
| 162 |
+
log_lines = [record.message for record in caplog.records]
|
| 163 |
+
assert len(log_lines) == 2 # start and success entries
|
| 164 |
+
|
| 165 |
+
start_entry = json.loads(log_lines[0])
|
| 166 |
+
assert start_entry["event"] == "request_start"
|
| 167 |
+
assert start_entry["method"] == "test_method"
|
| 168 |
+
|
| 169 |
+
success_entry = json.loads(log_lines[1])
|
| 170 |
+
assert success_entry["event"] == "request_success"
|
| 171 |
+
assert success_entry["result_type"] == "str"
|
| 172 |
+
|
| 173 |
+
async def test_on_message_failure(self, mock_context, caplog):
|
| 174 |
+
"""Test structured logging of failed messages."""
|
| 175 |
+
middleware = StructuredLoggingMiddleware()
|
| 176 |
+
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
| 177 |
+
|
| 178 |
+
with caplog.at_level(logging.INFO):
|
| 179 |
+
with pytest.raises(ValueError):
|
| 180 |
+
await middleware.on_message(mock_context, mock_call_next)
|
| 181 |
+
|
| 182 |
+
# Check that we have structured JSON logs
|
| 183 |
+
log_lines = [record.message for record in caplog.records]
|
| 184 |
+
assert len(log_lines) == 2 # start and error entries
|
| 185 |
+
|
| 186 |
+
start_entry = json.loads(log_lines[0])
|
| 187 |
+
assert start_entry["event"] == "request_start"
|
| 188 |
+
|
| 189 |
+
error_entry = json.loads(log_lines[1])
|
| 190 |
+
assert error_entry["event"] == "request_error"
|
| 191 |
+
assert error_entry["error_type"] == "ValueError"
|
| 192 |
+
assert error_entry["error_message"] == "test error"
|
| 193 |
+
|
| 194 |
+
|
| 195 |
+
@pytest.fixture
|
| 196 |
+
def logging_server():
|
| 197 |
+
"""Create a FastMCP server specifically for logging middleware tests."""
|
| 198 |
+
from fastmcp import FastMCP
|
| 199 |
+
|
| 200 |
+
mcp = FastMCP("LoggingTestServer")
|
| 201 |
+
|
| 202 |
+
@mcp.tool
|
| 203 |
+
def simple_operation(data: str) -> str:
|
| 204 |
+
"""A simple operation for testing logging."""
|
| 205 |
+
return f"Processed: {data}"
|
| 206 |
+
|
| 207 |
+
@mcp.tool
|
| 208 |
+
def complex_operation(items: list[str], mode: str = "default") -> dict:
|
| 209 |
+
"""A complex operation with structured data."""
|
| 210 |
+
return {"processed_items": len(items), "mode": mode, "result": "success"}
|
| 211 |
+
|
| 212 |
+
@mcp.tool
|
| 213 |
+
def operation_with_error(should_fail: bool = False) -> str:
|
| 214 |
+
"""An operation that can be made to fail."""
|
| 215 |
+
if should_fail:
|
| 216 |
+
raise ValueError("Operation failed intentionally")
|
| 217 |
+
return "Operation completed successfully"
|
| 218 |
+
|
| 219 |
+
@mcp.resource("log://test")
|
| 220 |
+
def test_resource() -> str:
|
| 221 |
+
"""A test resource for logging."""
|
| 222 |
+
return "Test resource content"
|
| 223 |
+
|
| 224 |
+
@mcp.prompt
|
| 225 |
+
def test_prompt() -> str:
|
| 226 |
+
"""A test prompt for logging."""
|
| 227 |
+
return "Test prompt content"
|
| 228 |
+
|
| 229 |
+
return mcp
|
| 230 |
+
|
| 231 |
+
|
| 232 |
+
class TestLoggingMiddlewareIntegration:
|
| 233 |
+
"""Integration tests for logging middleware with real FastMCP server."""
|
| 234 |
+
|
| 235 |
+
async def test_logging_middleware_logs_successful_operations(
|
| 236 |
+
self, logging_server, caplog
|
| 237 |
+
):
|
| 238 |
+
"""Test that logging middleware captures successful operations."""
|
| 239 |
+
from fastmcp.client import Client
|
| 240 |
+
|
| 241 |
+
logging_server.add_middleware(LoggingMiddleware())
|
| 242 |
+
|
| 243 |
+
with caplog.at_level(logging.INFO):
|
| 244 |
+
async with Client(logging_server) as client:
|
| 245 |
+
await client.call_tool("simple_operation", {"data": "test_data"})
|
| 246 |
+
await client.call_tool(
|
| 247 |
+
"complex_operation", {"items": ["a", "b", "c"], "mode": "batch"}
|
| 248 |
+
)
|
| 249 |
+
|
| 250 |
+
log_text = caplog.text
|
| 251 |
+
|
| 252 |
+
# Should have processing and completion logs for both operations
|
| 253 |
+
assert "Processing message:" in log_text
|
| 254 |
+
assert "Completed message: tools/call" in log_text
|
| 255 |
+
|
| 256 |
+
# Should have captured both tool calls
|
| 257 |
+
processing_count = log_text.count("Processing message:")
|
| 258 |
+
completion_count = log_text.count("Completed message:")
|
| 259 |
+
assert processing_count == 2
|
| 260 |
+
assert completion_count == 2
|
| 261 |
+
|
| 262 |
+
async def test_logging_middleware_logs_failures(self, logging_server, caplog):
|
| 263 |
+
"""Test that logging middleware captures failed operations."""
|
| 264 |
+
from fastmcp.client import Client
|
| 265 |
+
|
| 266 |
+
logging_server.add_middleware(LoggingMiddleware())
|
| 267 |
+
|
| 268 |
+
with caplog.at_level(logging.INFO):
|
| 269 |
+
async with Client(logging_server) as client:
|
| 270 |
+
# This should fail and be logged
|
| 271 |
+
with pytest.raises(Exception):
|
| 272 |
+
await client.call_tool(
|
| 273 |
+
"operation_with_error", {"should_fail": True}
|
| 274 |
+
)
|
| 275 |
+
|
| 276 |
+
log_text = caplog.text
|
| 277 |
+
|
| 278 |
+
# Should have processing and failure logs
|
| 279 |
+
assert "Processing message:" in log_text
|
| 280 |
+
assert "Failed message: tools/call" in log_text
|
| 281 |
+
|
| 282 |
+
async def test_logging_middleware_with_payloads(self, logging_server, caplog):
|
| 283 |
+
"""Test logging middleware when configured to include payloads."""
|
| 284 |
+
from fastmcp.client import Client
|
| 285 |
+
|
| 286 |
+
logging_server.add_middleware(
|
| 287 |
+
LoggingMiddleware(include_payloads=True, max_payload_length=500)
|
| 288 |
+
)
|
| 289 |
+
|
| 290 |
+
with caplog.at_level(logging.INFO):
|
| 291 |
+
async with Client(logging_server) as client:
|
| 292 |
+
await client.call_tool("simple_operation", {"data": "payload_test"})
|
| 293 |
+
|
| 294 |
+
log_text = caplog.text
|
| 295 |
+
|
| 296 |
+
# Should include payload information
|
| 297 |
+
assert "Processing message:" in log_text
|
| 298 |
+
assert "payload=" in log_text
|
| 299 |
+
|
| 300 |
+
async def test_structured_logging_middleware_produces_json(
|
| 301 |
+
self, logging_server, caplog
|
| 302 |
+
):
|
| 303 |
+
"""Test that structured logging middleware produces parseable JSON logs."""
|
| 304 |
+
import json
|
| 305 |
+
|
| 306 |
+
from fastmcp.client import Client
|
| 307 |
+
|
| 308 |
+
logging_server.add_middleware(
|
| 309 |
+
StructuredLoggingMiddleware(include_payloads=True)
|
| 310 |
+
)
|
| 311 |
+
|
| 312 |
+
with caplog.at_level(logging.INFO):
|
| 313 |
+
async with Client(logging_server) as client:
|
| 314 |
+
await client.call_tool("simple_operation", {"data": "json_test"})
|
| 315 |
+
|
| 316 |
+
# Extract JSON log entries
|
| 317 |
+
log_lines = [
|
| 318 |
+
record.message
|
| 319 |
+
for record in caplog.records
|
| 320 |
+
if record.name == "fastmcp.structured"
|
| 321 |
+
]
|
| 322 |
+
|
| 323 |
+
assert len(log_lines) >= 2 # Should have start and success entries
|
| 324 |
+
|
| 325 |
+
# Each log line should be valid JSON
|
| 326 |
+
for line in log_lines:
|
| 327 |
+
log_entry = json.loads(line)
|
| 328 |
+
assert "event" in log_entry
|
| 329 |
+
assert "timestamp" in log_entry
|
| 330 |
+
assert "source" in log_entry
|
| 331 |
+
assert "type" in log_entry
|
| 332 |
+
assert "method" in log_entry
|
| 333 |
+
|
| 334 |
+
async def test_structured_logging_middleware_handles_errors(
|
| 335 |
+
self, logging_server, caplog
|
| 336 |
+
):
|
| 337 |
+
"""Test structured logging of errors with JSON format."""
|
| 338 |
+
import json
|
| 339 |
+
|
| 340 |
+
from fastmcp.client import Client
|
| 341 |
+
|
| 342 |
+
logging_server.add_middleware(StructuredLoggingMiddleware())
|
| 343 |
+
|
| 344 |
+
with caplog.at_level(logging.INFO):
|
| 345 |
+
async with Client(logging_server) as client:
|
| 346 |
+
with pytest.raises(Exception):
|
| 347 |
+
await client.call_tool(
|
| 348 |
+
"operation_with_error", {"should_fail": True}
|
| 349 |
+
)
|
| 350 |
+
|
| 351 |
+
# Extract JSON log entries
|
| 352 |
+
log_lines = [
|
| 353 |
+
record.message
|
| 354 |
+
for record in caplog.records
|
| 355 |
+
if record.name == "fastmcp.structured"
|
| 356 |
+
]
|
| 357 |
+
|
| 358 |
+
# Should have start and error entries
|
| 359 |
+
assert len(log_lines) >= 2
|
| 360 |
+
|
| 361 |
+
# Find the error entry
|
| 362 |
+
error_entries = []
|
| 363 |
+
for line in log_lines:
|
| 364 |
+
log_entry = json.loads(line)
|
| 365 |
+
if log_entry.get("event") == "request_error":
|
| 366 |
+
error_entries.append(log_entry)
|
| 367 |
+
|
| 368 |
+
assert len(error_entries) == 1
|
| 369 |
+
error_entry = error_entries[0]
|
| 370 |
+
assert "error_type" in error_entry
|
| 371 |
+
assert "error_message" in error_entry
|
| 372 |
+
|
| 373 |
+
async def test_logging_middleware_with_different_operations(
|
| 374 |
+
self, logging_server, caplog
|
| 375 |
+
):
|
| 376 |
+
"""Test logging middleware with various MCP operations."""
|
| 377 |
+
from fastmcp.client import Client
|
| 378 |
+
|
| 379 |
+
logging_server.add_middleware(LoggingMiddleware())
|
| 380 |
+
|
| 381 |
+
with caplog.at_level(logging.INFO):
|
| 382 |
+
async with Client(logging_server) as client:
|
| 383 |
+
# Test different operation types
|
| 384 |
+
await client.call_tool("simple_operation", {"data": "test"})
|
| 385 |
+
await client.read_resource("log://test")
|
| 386 |
+
await client.get_prompt("test_prompt")
|
| 387 |
+
await client.list_tools()
|
| 388 |
+
|
| 389 |
+
log_text = caplog.text
|
| 390 |
+
|
| 391 |
+
# Should have logs for all different operation types
|
| 392 |
+
# Note: Different operations may have different method names
|
| 393 |
+
processing_count = log_text.count("Processing message:")
|
| 394 |
+
completion_count = log_text.count("Completed message:")
|
| 395 |
+
|
| 396 |
+
# Should have processed all 4 operations
|
| 397 |
+
assert processing_count == 4
|
| 398 |
+
assert completion_count == 4
|
| 399 |
+
|
| 400 |
+
async def test_logging_middleware_custom_configuration(self, logging_server):
|
| 401 |
+
"""Test logging middleware with custom logger configuration."""
|
| 402 |
+
import io
|
| 403 |
+
import logging
|
| 404 |
+
|
| 405 |
+
from fastmcp.client import Client
|
| 406 |
+
|
| 407 |
+
# Create custom logger
|
| 408 |
+
log_buffer = io.StringIO()
|
| 409 |
+
handler = logging.StreamHandler(log_buffer)
|
| 410 |
+
custom_logger = logging.getLogger("custom_logging_test")
|
| 411 |
+
custom_logger.addHandler(handler)
|
| 412 |
+
custom_logger.setLevel(logging.DEBUG)
|
| 413 |
+
|
| 414 |
+
logging_server.add_middleware(
|
| 415 |
+
LoggingMiddleware(
|
| 416 |
+
logger=custom_logger, log_level=logging.DEBUG, include_payloads=True
|
| 417 |
+
)
|
| 418 |
+
)
|
| 419 |
+
|
| 420 |
+
async with Client(logging_server) as client:
|
| 421 |
+
await client.call_tool("simple_operation", {"data": "custom_test"})
|
| 422 |
+
|
| 423 |
+
# Check that our custom logger captured the logs
|
| 424 |
+
log_output = log_buffer.getvalue()
|
| 425 |
+
assert "Processing message:" in log_output
|
| 426 |
+
assert "payload=" in log_output
|
tests/server/middleware/test_rate_limiting.py
ADDED
|
@@ -0,0 +1,448 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for rate limiting middleware."""
|
| 2 |
+
|
| 3 |
+
import asyncio
|
| 4 |
+
from unittest.mock import AsyncMock, MagicMock
|
| 5 |
+
|
| 6 |
+
import pytest
|
| 7 |
+
|
| 8 |
+
from fastmcp import FastMCP
|
| 9 |
+
from fastmcp.client import Client
|
| 10 |
+
from fastmcp.exceptions import ToolError
|
| 11 |
+
from fastmcp.server.middleware.middleware import MiddlewareContext
|
| 12 |
+
from fastmcp.server.middleware.rate_limiting import (
|
| 13 |
+
RateLimitError,
|
| 14 |
+
RateLimitingMiddleware,
|
| 15 |
+
SlidingWindowRateLimiter,
|
| 16 |
+
SlidingWindowRateLimitingMiddleware,
|
| 17 |
+
TokenBucketRateLimiter,
|
| 18 |
+
)
|
| 19 |
+
|
| 20 |
+
|
| 21 |
+
@pytest.fixture
|
| 22 |
+
def mock_context():
|
| 23 |
+
"""Create a mock middleware context."""
|
| 24 |
+
context = MagicMock(spec=MiddlewareContext)
|
| 25 |
+
context.method = "test_method"
|
| 26 |
+
return context
|
| 27 |
+
|
| 28 |
+
|
| 29 |
+
@pytest.fixture
|
| 30 |
+
def mock_call_next():
|
| 31 |
+
"""Create a mock call_next function."""
|
| 32 |
+
return AsyncMock(return_value="test_result")
|
| 33 |
+
|
| 34 |
+
|
| 35 |
+
class TestTokenBucketRateLimiter:
|
| 36 |
+
"""Test token bucket rate limiter."""
|
| 37 |
+
|
| 38 |
+
def test_init(self):
|
| 39 |
+
"""Test initialization."""
|
| 40 |
+
limiter = TokenBucketRateLimiter(capacity=10, refill_rate=5.0)
|
| 41 |
+
assert limiter.capacity == 10
|
| 42 |
+
assert limiter.refill_rate == 5.0
|
| 43 |
+
assert limiter.tokens == 10
|
| 44 |
+
|
| 45 |
+
async def test_consume_success(self):
|
| 46 |
+
"""Test successful token consumption."""
|
| 47 |
+
limiter = TokenBucketRateLimiter(capacity=10, refill_rate=5.0)
|
| 48 |
+
|
| 49 |
+
# Should be able to consume tokens initially
|
| 50 |
+
assert await limiter.consume(5) is True
|
| 51 |
+
assert await limiter.consume(3) is True
|
| 52 |
+
|
| 53 |
+
async def test_consume_failure(self):
|
| 54 |
+
"""Test failed token consumption."""
|
| 55 |
+
limiter = TokenBucketRateLimiter(capacity=5, refill_rate=1.0)
|
| 56 |
+
|
| 57 |
+
# Consume all tokens
|
| 58 |
+
assert await limiter.consume(5) is True
|
| 59 |
+
|
| 60 |
+
# Should fail to consume more
|
| 61 |
+
assert await limiter.consume(1) is False
|
| 62 |
+
|
| 63 |
+
async def test_refill(self):
|
| 64 |
+
"""Test token refill over time."""
|
| 65 |
+
limiter = TokenBucketRateLimiter(
|
| 66 |
+
capacity=10, refill_rate=10.0
|
| 67 |
+
) # 10 tokens per second
|
| 68 |
+
|
| 69 |
+
# Consume all tokens
|
| 70 |
+
assert await limiter.consume(10) is True
|
| 71 |
+
assert await limiter.consume(1) is False
|
| 72 |
+
|
| 73 |
+
# Wait for refill (0.2 seconds = 2 tokens at 10/sec)
|
| 74 |
+
await asyncio.sleep(0.2)
|
| 75 |
+
assert await limiter.consume(2) is True
|
| 76 |
+
|
| 77 |
+
|
| 78 |
+
class TestSlidingWindowRateLimiter:
|
| 79 |
+
"""Test sliding window rate limiter."""
|
| 80 |
+
|
| 81 |
+
def test_init(self):
|
| 82 |
+
"""Test initialization."""
|
| 83 |
+
limiter = SlidingWindowRateLimiter(max_requests=10, window_seconds=60)
|
| 84 |
+
assert limiter.max_requests == 10
|
| 85 |
+
assert limiter.window_seconds == 60
|
| 86 |
+
assert len(limiter.requests) == 0
|
| 87 |
+
|
| 88 |
+
async def test_is_allowed_success(self):
|
| 89 |
+
"""Test allowing requests within limit."""
|
| 90 |
+
limiter = SlidingWindowRateLimiter(max_requests=3, window_seconds=60)
|
| 91 |
+
|
| 92 |
+
# Should allow requests up to the limit
|
| 93 |
+
assert await limiter.is_allowed() is True
|
| 94 |
+
assert await limiter.is_allowed() is True
|
| 95 |
+
assert await limiter.is_allowed() is True
|
| 96 |
+
|
| 97 |
+
async def test_is_allowed_failure(self):
|
| 98 |
+
"""Test rejecting requests over limit."""
|
| 99 |
+
limiter = SlidingWindowRateLimiter(max_requests=2, window_seconds=60)
|
| 100 |
+
|
| 101 |
+
# Should allow up to limit
|
| 102 |
+
assert await limiter.is_allowed() is True
|
| 103 |
+
assert await limiter.is_allowed() is True
|
| 104 |
+
|
| 105 |
+
# Should reject over limit
|
| 106 |
+
assert await limiter.is_allowed() is False
|
| 107 |
+
|
| 108 |
+
async def test_sliding_window(self):
|
| 109 |
+
"""Test sliding window behavior."""
|
| 110 |
+
limiter = SlidingWindowRateLimiter(max_requests=2, window_seconds=1)
|
| 111 |
+
|
| 112 |
+
# Use up requests
|
| 113 |
+
assert await limiter.is_allowed() is True
|
| 114 |
+
assert await limiter.is_allowed() is True
|
| 115 |
+
assert await limiter.is_allowed() is False
|
| 116 |
+
|
| 117 |
+
# Wait for window to pass
|
| 118 |
+
await asyncio.sleep(1.1)
|
| 119 |
+
|
| 120 |
+
# Should be able to make requests again
|
| 121 |
+
assert await limiter.is_allowed() is True
|
| 122 |
+
|
| 123 |
+
|
| 124 |
+
class TestRateLimitingMiddleware:
|
| 125 |
+
"""Test rate limiting middleware."""
|
| 126 |
+
|
| 127 |
+
def test_init_default(self):
|
| 128 |
+
"""Test default initialization."""
|
| 129 |
+
middleware = RateLimitingMiddleware()
|
| 130 |
+
assert middleware.max_requests_per_second == 10.0
|
| 131 |
+
assert middleware.burst_capacity == 20
|
| 132 |
+
assert middleware.get_client_id is None
|
| 133 |
+
assert middleware.global_limit is False
|
| 134 |
+
|
| 135 |
+
def test_init_custom(self):
|
| 136 |
+
"""Test custom initialization."""
|
| 137 |
+
|
| 138 |
+
def get_client_id(ctx):
|
| 139 |
+
return "test_client"
|
| 140 |
+
|
| 141 |
+
middleware = RateLimitingMiddleware(
|
| 142 |
+
max_requests_per_second=5.0,
|
| 143 |
+
burst_capacity=10,
|
| 144 |
+
get_client_id=get_client_id,
|
| 145 |
+
global_limit=True,
|
| 146 |
+
)
|
| 147 |
+
assert middleware.max_requests_per_second == 5.0
|
| 148 |
+
assert middleware.burst_capacity == 10
|
| 149 |
+
assert middleware.get_client_id is get_client_id
|
| 150 |
+
assert middleware.global_limit is True
|
| 151 |
+
|
| 152 |
+
def test_get_client_identifier_default(self, mock_context):
|
| 153 |
+
"""Test default client identifier."""
|
| 154 |
+
middleware = RateLimitingMiddleware()
|
| 155 |
+
assert middleware._get_client_identifier(mock_context) == "global"
|
| 156 |
+
|
| 157 |
+
def test_get_client_identifier_custom(self, mock_context):
|
| 158 |
+
"""Test custom client identifier."""
|
| 159 |
+
|
| 160 |
+
def get_client_id(ctx):
|
| 161 |
+
return "custom_client"
|
| 162 |
+
|
| 163 |
+
middleware = RateLimitingMiddleware(get_client_id=get_client_id)
|
| 164 |
+
assert middleware._get_client_identifier(mock_context) == "custom_client"
|
| 165 |
+
|
| 166 |
+
async def test_on_request_success(self, mock_context, mock_call_next):
|
| 167 |
+
"""Test successful request within rate limit."""
|
| 168 |
+
middleware = RateLimitingMiddleware(max_requests_per_second=100.0) # High limit
|
| 169 |
+
|
| 170 |
+
result = await middleware.on_request(mock_context, mock_call_next)
|
| 171 |
+
|
| 172 |
+
assert result == "test_result"
|
| 173 |
+
assert mock_call_next.called
|
| 174 |
+
|
| 175 |
+
async def test_on_request_rate_limited(self, mock_context, mock_call_next):
|
| 176 |
+
"""Test request rejection due to rate limiting."""
|
| 177 |
+
middleware = RateLimitingMiddleware(
|
| 178 |
+
max_requests_per_second=1.0, burst_capacity=1
|
| 179 |
+
)
|
| 180 |
+
|
| 181 |
+
# First request should succeed
|
| 182 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 183 |
+
|
| 184 |
+
# Second request should be rate limited
|
| 185 |
+
with pytest.raises(RateLimitError, match="Rate limit exceeded"):
|
| 186 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 187 |
+
|
| 188 |
+
async def test_global_rate_limiting(self, mock_context, mock_call_next):
|
| 189 |
+
"""Test global rate limiting."""
|
| 190 |
+
middleware = RateLimitingMiddleware(
|
| 191 |
+
max_requests_per_second=1.0, burst_capacity=1, global_limit=True
|
| 192 |
+
)
|
| 193 |
+
|
| 194 |
+
# First request should succeed
|
| 195 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 196 |
+
|
| 197 |
+
# Second request should be rate limited
|
| 198 |
+
with pytest.raises(RateLimitError, match="Global rate limit exceeded"):
|
| 199 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 200 |
+
|
| 201 |
+
|
| 202 |
+
class TestSlidingWindowRateLimitingMiddleware:
|
| 203 |
+
"""Test sliding window rate limiting middleware."""
|
| 204 |
+
|
| 205 |
+
def test_init_default(self):
|
| 206 |
+
"""Test default initialization."""
|
| 207 |
+
middleware = SlidingWindowRateLimitingMiddleware(max_requests=100)
|
| 208 |
+
assert middleware.max_requests == 100
|
| 209 |
+
assert middleware.window_seconds == 60
|
| 210 |
+
assert middleware.get_client_id is None
|
| 211 |
+
|
| 212 |
+
def test_init_custom(self):
|
| 213 |
+
"""Test custom initialization."""
|
| 214 |
+
|
| 215 |
+
def get_client_id(ctx):
|
| 216 |
+
return "test_client"
|
| 217 |
+
|
| 218 |
+
middleware = SlidingWindowRateLimitingMiddleware(
|
| 219 |
+
max_requests=50, window_minutes=5, get_client_id=get_client_id
|
| 220 |
+
)
|
| 221 |
+
assert middleware.max_requests == 50
|
| 222 |
+
assert middleware.window_seconds == 300 # 5 minutes
|
| 223 |
+
assert middleware.get_client_id is get_client_id
|
| 224 |
+
|
| 225 |
+
async def test_on_request_success(self, mock_context, mock_call_next):
|
| 226 |
+
"""Test successful request within rate limit."""
|
| 227 |
+
middleware = SlidingWindowRateLimitingMiddleware(max_requests=100)
|
| 228 |
+
|
| 229 |
+
result = await middleware.on_request(mock_context, mock_call_next)
|
| 230 |
+
|
| 231 |
+
assert result == "test_result"
|
| 232 |
+
assert mock_call_next.called
|
| 233 |
+
|
| 234 |
+
async def test_on_request_rate_limited(self, mock_context, mock_call_next):
|
| 235 |
+
"""Test request rejection due to rate limiting."""
|
| 236 |
+
middleware = SlidingWindowRateLimitingMiddleware(max_requests=1)
|
| 237 |
+
|
| 238 |
+
# First request should succeed
|
| 239 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 240 |
+
|
| 241 |
+
# Second request should be rate limited
|
| 242 |
+
with pytest.raises(RateLimitError, match="Rate limit exceeded"):
|
| 243 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 244 |
+
|
| 245 |
+
|
| 246 |
+
class TestRateLimitError:
|
| 247 |
+
"""Test rate limit error."""
|
| 248 |
+
|
| 249 |
+
def test_init_default(self):
|
| 250 |
+
"""Test default initialization."""
|
| 251 |
+
error = RateLimitError()
|
| 252 |
+
assert error.error.code == -32000
|
| 253 |
+
assert error.error.message == "Rate limit exceeded"
|
| 254 |
+
|
| 255 |
+
def test_init_custom(self):
|
| 256 |
+
"""Test custom initialization."""
|
| 257 |
+
error = RateLimitError("Custom message")
|
| 258 |
+
assert error.error.code == -32000
|
| 259 |
+
assert error.error.message == "Custom message"
|
| 260 |
+
|
| 261 |
+
|
| 262 |
+
@pytest.fixture
|
| 263 |
+
def rate_limit_server():
|
| 264 |
+
"""Create a FastMCP server specifically for rate limiting tests."""
|
| 265 |
+
mcp = FastMCP("RateLimitTestServer")
|
| 266 |
+
|
| 267 |
+
@mcp.tool
|
| 268 |
+
def quick_action(message: str) -> str:
|
| 269 |
+
"""A quick action for testing rate limits."""
|
| 270 |
+
return f"Processed: {message}"
|
| 271 |
+
|
| 272 |
+
@mcp.tool
|
| 273 |
+
def batch_process(items: list[str]) -> str:
|
| 274 |
+
"""Process multiple items."""
|
| 275 |
+
return f"Processed {len(items)} items"
|
| 276 |
+
|
| 277 |
+
@mcp.tool
|
| 278 |
+
def heavy_computation() -> str:
|
| 279 |
+
"""A heavy computation that might need rate limiting."""
|
| 280 |
+
# Simulate some work
|
| 281 |
+
import time
|
| 282 |
+
|
| 283 |
+
time.sleep(0.01) # Very short delay
|
| 284 |
+
return "Heavy computation complete"
|
| 285 |
+
|
| 286 |
+
return mcp
|
| 287 |
+
|
| 288 |
+
|
| 289 |
+
class TestRateLimitingMiddlewareIntegration:
|
| 290 |
+
"""Integration tests for rate limiting middleware with real FastMCP server."""
|
| 291 |
+
|
| 292 |
+
async def test_rate_limiting_allows_normal_usage(self, rate_limit_server):
|
| 293 |
+
"""Test that normal usage patterns are allowed through rate limiting."""
|
| 294 |
+
# Generous rate limit
|
| 295 |
+
rate_limit_server.add_middleware(
|
| 296 |
+
RateLimitingMiddleware(max_requests_per_second=50.0, burst_capacity=10)
|
| 297 |
+
)
|
| 298 |
+
|
| 299 |
+
async with Client(rate_limit_server) as client:
|
| 300 |
+
# Normal usage should be fine
|
| 301 |
+
for i in range(5):
|
| 302 |
+
result = await client.call_tool(
|
| 303 |
+
"quick_action", {"message": f"task_{i}"}
|
| 304 |
+
)
|
| 305 |
+
assert f"Processed: task_{i}" in str(result)
|
| 306 |
+
|
| 307 |
+
async def test_rate_limiting_blocks_rapid_requests(self, rate_limit_server):
|
| 308 |
+
"""Test that rate limiting blocks rapid successive requests."""
|
| 309 |
+
# Very restrictive rate limit
|
| 310 |
+
rate_limit_server.add_middleware(
|
| 311 |
+
RateLimitingMiddleware(max_requests_per_second=2.0, burst_capacity=3)
|
| 312 |
+
)
|
| 313 |
+
|
| 314 |
+
async with Client(rate_limit_server) as client:
|
| 315 |
+
# First few should succeed (within burst capacity)
|
| 316 |
+
await client.call_tool("quick_action", {"message": "1"})
|
| 317 |
+
await client.call_tool("quick_action", {"message": "2"})
|
| 318 |
+
await client.call_tool("quick_action", {"message": "3"})
|
| 319 |
+
|
| 320 |
+
# Next should be rate limited
|
| 321 |
+
with pytest.raises(ToolError, match="Rate limit exceeded"):
|
| 322 |
+
await client.call_tool("quick_action", {"message": "4"})
|
| 323 |
+
|
| 324 |
+
async def test_rate_limiting_with_concurrent_requests(self, rate_limit_server):
|
| 325 |
+
"""Test rate limiting behavior with concurrent requests."""
|
| 326 |
+
rate_limit_server.add_middleware(
|
| 327 |
+
RateLimitingMiddleware(max_requests_per_second=5.0, burst_capacity=3)
|
| 328 |
+
)
|
| 329 |
+
|
| 330 |
+
async with Client(rate_limit_server) as client:
|
| 331 |
+
# Fire off many concurrent requests
|
| 332 |
+
tasks = []
|
| 333 |
+
for i in range(8):
|
| 334 |
+
task = asyncio.create_task(
|
| 335 |
+
client.call_tool("quick_action", {"message": f"concurrent_{i}"})
|
| 336 |
+
)
|
| 337 |
+
tasks.append(task)
|
| 338 |
+
|
| 339 |
+
# Gather results, allowing exceptions
|
| 340 |
+
results = await asyncio.gather(*tasks, return_exceptions=True)
|
| 341 |
+
|
| 342 |
+
# Some should succeed, some should be rate limited
|
| 343 |
+
successes = [r for r in results if not isinstance(r, Exception)]
|
| 344 |
+
failures = [r for r in results if isinstance(r, ToolError)]
|
| 345 |
+
|
| 346 |
+
assert len(successes) > 0, "Some requests should succeed"
|
| 347 |
+
assert len(failures) > 0, "Some requests should be rate limited"
|
| 348 |
+
assert len(successes) + len(failures) == 8
|
| 349 |
+
|
| 350 |
+
async def test_sliding_window_rate_limiting(self, rate_limit_server):
|
| 351 |
+
"""Test sliding window rate limiting implementation."""
|
| 352 |
+
rate_limit_server.add_middleware(
|
| 353 |
+
SlidingWindowRateLimitingMiddleware(
|
| 354 |
+
max_requests=3,
|
| 355 |
+
window_minutes=1, # 1 minute window
|
| 356 |
+
)
|
| 357 |
+
)
|
| 358 |
+
|
| 359 |
+
async with Client(rate_limit_server) as client:
|
| 360 |
+
# Should allow up to the limit
|
| 361 |
+
await client.call_tool("quick_action", {"message": "1"})
|
| 362 |
+
await client.call_tool("quick_action", {"message": "2"})
|
| 363 |
+
await client.call_tool("quick_action", {"message": "3"})
|
| 364 |
+
|
| 365 |
+
# Fourth should be blocked
|
| 366 |
+
with pytest.raises(ToolError, match="Rate limit exceeded"):
|
| 367 |
+
await client.call_tool("quick_action", {"message": "4"})
|
| 368 |
+
|
| 369 |
+
async def test_rate_limiting_with_different_operations(self, rate_limit_server):
|
| 370 |
+
"""Test that rate limiting applies to all types of operations."""
|
| 371 |
+
rate_limit_server.add_middleware(
|
| 372 |
+
RateLimitingMiddleware(max_requests_per_second=3.0, burst_capacity=2)
|
| 373 |
+
)
|
| 374 |
+
|
| 375 |
+
async with Client(rate_limit_server) as client:
|
| 376 |
+
# Mix different operations
|
| 377 |
+
await client.call_tool("quick_action", {"message": "test"})
|
| 378 |
+
await client.call_tool("heavy_computation")
|
| 379 |
+
|
| 380 |
+
# Should be rate limited regardless of operation type
|
| 381 |
+
with pytest.raises(ToolError, match="Rate limit exceeded"):
|
| 382 |
+
await client.call_tool("batch_process", {"items": ["a", "b", "c"]})
|
| 383 |
+
|
| 384 |
+
async def test_custom_client_identification(self, rate_limit_server):
|
| 385 |
+
"""Test rate limiting with custom client identification."""
|
| 386 |
+
|
| 387 |
+
def get_client_id(context):
|
| 388 |
+
# In a real scenario, this might extract from headers or context
|
| 389 |
+
return "test_client_123"
|
| 390 |
+
|
| 391 |
+
rate_limit_server.add_middleware(
|
| 392 |
+
RateLimitingMiddleware(
|
| 393 |
+
max_requests_per_second=2.0,
|
| 394 |
+
burst_capacity=1,
|
| 395 |
+
get_client_id=get_client_id,
|
| 396 |
+
)
|
| 397 |
+
)
|
| 398 |
+
|
| 399 |
+
async with Client(rate_limit_server) as client:
|
| 400 |
+
# First request should succeed
|
| 401 |
+
await client.call_tool("quick_action", {"message": "first"})
|
| 402 |
+
|
| 403 |
+
# Second should be rate limited for this specific client
|
| 404 |
+
with pytest.raises(
|
| 405 |
+
ToolError, match="Rate limit exceeded for client: test_client_123"
|
| 406 |
+
):
|
| 407 |
+
await client.call_tool("quick_action", {"message": "second"})
|
| 408 |
+
|
| 409 |
+
async def test_global_rate_limiting(self, rate_limit_server):
|
| 410 |
+
"""Test global rate limiting across all clients."""
|
| 411 |
+
rate_limit_server.add_middleware(
|
| 412 |
+
RateLimitingMiddleware(
|
| 413 |
+
max_requests_per_second=2.0, burst_capacity=2, global_limit=True
|
| 414 |
+
)
|
| 415 |
+
)
|
| 416 |
+
|
| 417 |
+
async with Client(rate_limit_server) as client:
|
| 418 |
+
# Use up the global capacity
|
| 419 |
+
await client.call_tool("quick_action", {"message": "1"})
|
| 420 |
+
await client.call_tool("quick_action", {"message": "2"})
|
| 421 |
+
|
| 422 |
+
# Should be globally rate limited
|
| 423 |
+
with pytest.raises(ToolError, match="Global rate limit exceeded"):
|
| 424 |
+
await client.call_tool("quick_action", {"message": "3"})
|
| 425 |
+
|
| 426 |
+
async def test_rate_limiting_recovery_over_time(self, rate_limit_server):
|
| 427 |
+
"""Test that rate limiting allows requests again after time passes."""
|
| 428 |
+
rate_limit_server.add_middleware(
|
| 429 |
+
RateLimitingMiddleware(
|
| 430 |
+
max_requests_per_second=10.0, # 10 per second = 1 every 100ms
|
| 431 |
+
burst_capacity=1,
|
| 432 |
+
)
|
| 433 |
+
)
|
| 434 |
+
|
| 435 |
+
async with Client(rate_limit_server) as client:
|
| 436 |
+
# Use up capacity
|
| 437 |
+
await client.call_tool("quick_action", {"message": "first"})
|
| 438 |
+
|
| 439 |
+
# Should be rate limited immediately
|
| 440 |
+
with pytest.raises(ToolError):
|
| 441 |
+
await client.call_tool("quick_action", {"message": "second"})
|
| 442 |
+
|
| 443 |
+
# Wait for token bucket to refill (150ms should be enough for ~1.5 tokens)
|
| 444 |
+
await asyncio.sleep(0.15)
|
| 445 |
+
|
| 446 |
+
# Should be able to make another request
|
| 447 |
+
result = await client.call_tool("quick_action", {"message": "after_wait"})
|
| 448 |
+
assert "after_wait" in str(result)
|
tests/server/middleware/test_timing.py
ADDED
|
@@ -0,0 +1,312 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 1 |
+
"""Tests for timing middleware."""
|
| 2 |
+
|
| 3 |
+
import asyncio
|
| 4 |
+
import logging
|
| 5 |
+
import time
|
| 6 |
+
from unittest.mock import AsyncMock, MagicMock
|
| 7 |
+
|
| 8 |
+
import pytest
|
| 9 |
+
|
| 10 |
+
from fastmcp import FastMCP
|
| 11 |
+
from fastmcp.client import Client
|
| 12 |
+
from fastmcp.server.middleware.middleware import MiddlewareContext
|
| 13 |
+
from fastmcp.server.middleware.timing import DetailedTimingMiddleware, TimingMiddleware
|
| 14 |
+
|
| 15 |
+
|
| 16 |
+
@pytest.fixture
|
| 17 |
+
def mock_context():
|
| 18 |
+
"""Create a mock middleware context."""
|
| 19 |
+
context = MagicMock(spec=MiddlewareContext)
|
| 20 |
+
context.method = "test_method"
|
| 21 |
+
return context
|
| 22 |
+
|
| 23 |
+
|
| 24 |
+
@pytest.fixture
|
| 25 |
+
def mock_call_next():
|
| 26 |
+
"""Create a mock call_next function."""
|
| 27 |
+
return AsyncMock(return_value="test_result")
|
| 28 |
+
|
| 29 |
+
|
| 30 |
+
class TestTimingMiddleware:
|
| 31 |
+
"""Test timing middleware functionality."""
|
| 32 |
+
|
| 33 |
+
def test_init_default(self):
|
| 34 |
+
"""Test default initialization."""
|
| 35 |
+
middleware = TimingMiddleware()
|
| 36 |
+
assert middleware.logger.name == "fastmcp.timing"
|
| 37 |
+
assert middleware.log_level == logging.INFO
|
| 38 |
+
|
| 39 |
+
def test_init_custom(self):
|
| 40 |
+
"""Test custom initialization."""
|
| 41 |
+
logger = logging.getLogger("custom")
|
| 42 |
+
middleware = TimingMiddleware(logger=logger, log_level=logging.DEBUG)
|
| 43 |
+
assert middleware.logger is logger
|
| 44 |
+
assert middleware.log_level == logging.DEBUG
|
| 45 |
+
|
| 46 |
+
async def test_on_request_success(self, mock_context, mock_call_next, caplog):
|
| 47 |
+
"""Test timing successful requests."""
|
| 48 |
+
middleware = TimingMiddleware()
|
| 49 |
+
|
| 50 |
+
with caplog.at_level(logging.INFO):
|
| 51 |
+
result = await middleware.on_request(mock_context, mock_call_next)
|
| 52 |
+
|
| 53 |
+
assert result == "test_result"
|
| 54 |
+
assert mock_call_next.called
|
| 55 |
+
assert "Request test_method completed in" in caplog.text
|
| 56 |
+
assert "ms" in caplog.text
|
| 57 |
+
|
| 58 |
+
async def test_on_request_failure(self, mock_context, caplog):
|
| 59 |
+
"""Test timing failed requests."""
|
| 60 |
+
middleware = TimingMiddleware()
|
| 61 |
+
mock_call_next = AsyncMock(side_effect=ValueError("test error"))
|
| 62 |
+
|
| 63 |
+
with caplog.at_level(logging.INFO):
|
| 64 |
+
with pytest.raises(ValueError):
|
| 65 |
+
await middleware.on_request(mock_context, mock_call_next)
|
| 66 |
+
|
| 67 |
+
assert "Request test_method failed after" in caplog.text
|
| 68 |
+
assert "ms: test error" in caplog.text
|
| 69 |
+
|
| 70 |
+
|
| 71 |
+
class TestDetailedTimingMiddleware:
|
| 72 |
+
"""Test detailed timing middleware functionality."""
|
| 73 |
+
|
| 74 |
+
def test_init_default(self):
|
| 75 |
+
"""Test default initialization."""
|
| 76 |
+
middleware = DetailedTimingMiddleware()
|
| 77 |
+
assert middleware.logger.name == "fastmcp.timing.detailed"
|
| 78 |
+
assert middleware.log_level == logging.INFO
|
| 79 |
+
|
| 80 |
+
async def test_on_call_tool(self, caplog):
|
| 81 |
+
"""Test timing tool calls."""
|
| 82 |
+
middleware = DetailedTimingMiddleware()
|
| 83 |
+
context = MagicMock()
|
| 84 |
+
context.message.name = "test_tool"
|
| 85 |
+
mock_call_next = AsyncMock(return_value="tool_result")
|
| 86 |
+
|
| 87 |
+
with caplog.at_level(logging.INFO):
|
| 88 |
+
result = await middleware.on_call_tool(context, mock_call_next)
|
| 89 |
+
|
| 90 |
+
assert result == "tool_result"
|
| 91 |
+
assert "Tool 'test_tool' completed in" in caplog.text
|
| 92 |
+
|
| 93 |
+
async def test_on_read_resource(self, caplog):
|
| 94 |
+
"""Test timing resource reads."""
|
| 95 |
+
middleware = DetailedTimingMiddleware()
|
| 96 |
+
context = MagicMock()
|
| 97 |
+
context.message.uri = "test://resource"
|
| 98 |
+
mock_call_next = AsyncMock(return_value="resource_result")
|
| 99 |
+
|
| 100 |
+
with caplog.at_level(logging.INFO):
|
| 101 |
+
result = await middleware.on_read_resource(context, mock_call_next)
|
| 102 |
+
|
| 103 |
+
assert result == "resource_result"
|
| 104 |
+
assert "Resource 'test://resource' completed in" in caplog.text
|
| 105 |
+
|
| 106 |
+
async def test_on_get_prompt(self, caplog):
|
| 107 |
+
"""Test timing prompt retrieval."""
|
| 108 |
+
middleware = DetailedTimingMiddleware()
|
| 109 |
+
context = MagicMock()
|
| 110 |
+
context.message.name = "test_prompt"
|
| 111 |
+
mock_call_next = AsyncMock(return_value="prompt_result")
|
| 112 |
+
|
| 113 |
+
with caplog.at_level(logging.INFO):
|
| 114 |
+
result = await middleware.on_get_prompt(context, mock_call_next)
|
| 115 |
+
|
| 116 |
+
assert result == "prompt_result"
|
| 117 |
+
assert "Prompt 'test_prompt' completed in" in caplog.text
|
| 118 |
+
|
| 119 |
+
async def test_on_list_tools(self, caplog):
|
| 120 |
+
"""Test timing tool listing."""
|
| 121 |
+
middleware = DetailedTimingMiddleware()
|
| 122 |
+
context = MagicMock()
|
| 123 |
+
mock_call_next = AsyncMock(return_value="tools_result")
|
| 124 |
+
|
| 125 |
+
with caplog.at_level(logging.INFO):
|
| 126 |
+
result = await middleware.on_list_tools(context, mock_call_next)
|
| 127 |
+
|
| 128 |
+
assert result == "tools_result"
|
| 129 |
+
assert "List tools completed in" in caplog.text
|
| 130 |
+
|
| 131 |
+
async def test_operation_failure(self, caplog):
|
| 132 |
+
"""Test timing failed operations."""
|
| 133 |
+
middleware = DetailedTimingMiddleware()
|
| 134 |
+
context = MagicMock()
|
| 135 |
+
context.message.name = "failing_tool"
|
| 136 |
+
mock_call_next = AsyncMock(side_effect=RuntimeError("operation failed"))
|
| 137 |
+
|
| 138 |
+
with caplog.at_level(logging.INFO):
|
| 139 |
+
with pytest.raises(RuntimeError):
|
| 140 |
+
await middleware.on_call_tool(context, mock_call_next)
|
| 141 |
+
|
| 142 |
+
assert "Tool 'failing_tool' failed after" in caplog.text
|
| 143 |
+
assert "ms: operation failed" in caplog.text
|
| 144 |
+
|
| 145 |
+
|
| 146 |
+
@pytest.fixture
|
| 147 |
+
def timing_server():
|
| 148 |
+
"""Create a FastMCP server specifically for timing middleware tests."""
|
| 149 |
+
mcp = FastMCP("TimingTestServer")
|
| 150 |
+
|
| 151 |
+
@mcp.tool
|
| 152 |
+
def instant_task() -> str:
|
| 153 |
+
"""A task that completes instantly."""
|
| 154 |
+
return "Done instantly"
|
| 155 |
+
|
| 156 |
+
@mcp.tool
|
| 157 |
+
def short_task() -> str:
|
| 158 |
+
"""A task that takes 0.1 seconds."""
|
| 159 |
+
time.sleep(0.1)
|
| 160 |
+
return "Done after 0.1s"
|
| 161 |
+
|
| 162 |
+
@mcp.tool
|
| 163 |
+
def medium_task() -> str:
|
| 164 |
+
"""A task that takes 0.15 seconds."""
|
| 165 |
+
time.sleep(0.15)
|
| 166 |
+
return "Done after 0.15s"
|
| 167 |
+
|
| 168 |
+
@mcp.tool
|
| 169 |
+
def failing_task() -> str:
|
| 170 |
+
"""A task that always fails."""
|
| 171 |
+
raise ValueError("Task failed as expected")
|
| 172 |
+
|
| 173 |
+
@mcp.resource("timer://test")
|
| 174 |
+
def test_resource() -> str:
|
| 175 |
+
"""A resource that takes time to read."""
|
| 176 |
+
time.sleep(0.05)
|
| 177 |
+
return "Resource content after 0.05s"
|
| 178 |
+
|
| 179 |
+
@mcp.prompt
|
| 180 |
+
def test_prompt() -> str:
|
| 181 |
+
"""A prompt that takes time to generate."""
|
| 182 |
+
time.sleep(0.08)
|
| 183 |
+
return "Prompt content after 0.08s"
|
| 184 |
+
|
| 185 |
+
return mcp
|
| 186 |
+
|
| 187 |
+
|
| 188 |
+
class TestTimingMiddlewareIntegration:
|
| 189 |
+
"""Integration tests for timing middleware with real FastMCP server."""
|
| 190 |
+
|
| 191 |
+
async def test_timing_middleware_measures_tool_execution(
|
| 192 |
+
self, timing_server, caplog
|
| 193 |
+
):
|
| 194 |
+
"""Test that timing middleware accurately measures tool execution times."""
|
| 195 |
+
timing_server.add_middleware(TimingMiddleware())
|
| 196 |
+
|
| 197 |
+
with caplog.at_level(logging.INFO):
|
| 198 |
+
async with Client(timing_server) as client:
|
| 199 |
+
# Test instant task
|
| 200 |
+
await client.call_tool("instant_task")
|
| 201 |
+
|
| 202 |
+
# Test short task (0.1s)
|
| 203 |
+
await client.call_tool("short_task")
|
| 204 |
+
|
| 205 |
+
# Test medium task (0.15s)
|
| 206 |
+
await client.call_tool("medium_task")
|
| 207 |
+
|
| 208 |
+
log_text = caplog.text
|
| 209 |
+
|
| 210 |
+
# Should have timing logs for all three calls
|
| 211 |
+
timing_logs = [
|
| 212 |
+
line
|
| 213 |
+
for line in log_text.split("\n")
|
| 214 |
+
if "completed in" in line and "ms" in line
|
| 215 |
+
]
|
| 216 |
+
assert len(timing_logs) == 3
|
| 217 |
+
|
| 218 |
+
# Verify that longer tasks show longer timing (roughly)
|
| 219 |
+
assert "tools/call completed in" in log_text
|
| 220 |
+
assert "ms" in log_text
|
| 221 |
+
|
| 222 |
+
async def test_timing_middleware_handles_failures(self, timing_server, caplog):
|
| 223 |
+
"""Test that timing middleware measures time even for failed operations."""
|
| 224 |
+
timing_server.add_middleware(TimingMiddleware())
|
| 225 |
+
|
| 226 |
+
with caplog.at_level(logging.INFO):
|
| 227 |
+
async with Client(timing_server) as client:
|
| 228 |
+
# This should fail but still be timed
|
| 229 |
+
with pytest.raises(Exception):
|
| 230 |
+
await client.call_tool("failing_task")
|
| 231 |
+
|
| 232 |
+
# Should log the failure with timing
|
| 233 |
+
assert "tools/call failed after" in caplog.text
|
| 234 |
+
assert "ms:" in caplog.text
|
| 235 |
+
|
| 236 |
+
async def test_detailed_timing_middleware_per_operation(
|
| 237 |
+
self, timing_server, caplog
|
| 238 |
+
):
|
| 239 |
+
"""Test that detailed timing middleware provides operation-specific timing."""
|
| 240 |
+
timing_server.add_middleware(DetailedTimingMiddleware())
|
| 241 |
+
|
| 242 |
+
with caplog.at_level(logging.INFO):
|
| 243 |
+
async with Client(timing_server) as client:
|
| 244 |
+
# Test tool call
|
| 245 |
+
await client.call_tool("short_task")
|
| 246 |
+
|
| 247 |
+
# Test resource read
|
| 248 |
+
await client.read_resource("timer://test")
|
| 249 |
+
|
| 250 |
+
# Test prompt
|
| 251 |
+
await client.get_prompt("test_prompt")
|
| 252 |
+
|
| 253 |
+
# Test listing operations
|
| 254 |
+
await client.list_tools()
|
| 255 |
+
await client.list_resources()
|
| 256 |
+
await client.list_prompts()
|
| 257 |
+
|
| 258 |
+
log_text = caplog.text
|
| 259 |
+
|
| 260 |
+
# Should have specific timing logs for each operation type
|
| 261 |
+
assert "Tool 'short_task' completed in" in log_text
|
| 262 |
+
assert "Resource 'timer://test' completed in" in log_text
|
| 263 |
+
assert "Prompt 'test_prompt' completed in" in log_text
|
| 264 |
+
assert "List tools completed in" in log_text
|
| 265 |
+
assert "List resources completed in" in log_text
|
| 266 |
+
assert "List prompts completed in" in log_text
|
| 267 |
+
|
| 268 |
+
async def test_timing_middleware_concurrent_operations(self, timing_server, caplog):
|
| 269 |
+
"""Test timing middleware with concurrent operations."""
|
| 270 |
+
timing_server.add_middleware(TimingMiddleware())
|
| 271 |
+
|
| 272 |
+
with caplog.at_level(logging.INFO):
|
| 273 |
+
async with Client(timing_server) as client:
|
| 274 |
+
# Run multiple operations concurrently
|
| 275 |
+
tasks = [
|
| 276 |
+
client.call_tool("instant_task"),
|
| 277 |
+
client.call_tool("short_task"),
|
| 278 |
+
client.call_tool("instant_task"),
|
| 279 |
+
]
|
| 280 |
+
|
| 281 |
+
await asyncio.gather(*tasks)
|
| 282 |
+
|
| 283 |
+
log_text = caplog.text
|
| 284 |
+
|
| 285 |
+
# Should have timing logs for all concurrent operations
|
| 286 |
+
timing_logs = [line for line in log_text.split("\n") if "completed in" in line]
|
| 287 |
+
assert len(timing_logs) == 3
|
| 288 |
+
|
| 289 |
+
async def test_timing_middleware_custom_logger(self, timing_server):
|
| 290 |
+
"""Test timing middleware with custom logger configuration."""
|
| 291 |
+
import io
|
| 292 |
+
import logging
|
| 293 |
+
|
| 294 |
+
# Create a custom logger that writes to a string buffer
|
| 295 |
+
log_buffer = io.StringIO()
|
| 296 |
+
handler = logging.StreamHandler(log_buffer)
|
| 297 |
+
custom_logger = logging.getLogger("custom_timing")
|
| 298 |
+
custom_logger.addHandler(handler)
|
| 299 |
+
custom_logger.setLevel(logging.DEBUG)
|
| 300 |
+
|
| 301 |
+
# Use custom logger and log level
|
| 302 |
+
timing_server.add_middleware(
|
| 303 |
+
TimingMiddleware(logger=custom_logger, log_level=logging.DEBUG)
|
| 304 |
+
)
|
| 305 |
+
|
| 306 |
+
async with Client(timing_server) as client:
|
| 307 |
+
await client.call_tool("instant_task")
|
| 308 |
+
|
| 309 |
+
# Check that our custom logger was used
|
| 310 |
+
log_output = log_buffer.getvalue()
|
| 311 |
+
assert "tools/call completed in" in log_output
|
| 312 |
+
assert "ms" in log_output
|