Spaces:
Running
Running
File size: 5,591 Bytes
5800457 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 | """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")
|