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