mcp-use 1.3.12__py3-none-any.whl → 1.3.13__py3-none-any.whl

mcp_use/client/middleware/metrics.py

@@ -0,0 +1,314 @@
+"""
+Metrics middleware for MCP contexts.
+
+Classes for collecting comprehensive metrics about MCP context patterns,
+performance, and errors with simple instantiation.
+"""
+
+import asyncio
+import time
+from collections import Counter, defaultdict
+from typing import Any
+
+from mcp_use.client.middleware.middleware import Middleware, MiddlewareContext, NextFunctionT
+
+# Constants for performance thresholds
+SLOW_THRESHOLD_MS = 1000  # 1 second
+FAST_THRESHOLD_MS = 50  # 50ms
+MAX_SLOW_CONTEXTS = 100
+MAX_FAST_CONTEXTS = 100
+MAX_ERROR_TIMESTAMPS = 1000
+MAX_RECENT_ERRORS = 50
+SECONDS_PER_HOUR = 3600
+SECONDS_PER_MINUTE = 60
+
+
+class MetricsMiddleware(Middleware):
+    """Collects basic metrics about MCP contexts including counts, durations, and errors."""
+
+    def __init__(self):
+        self.metrics = {
+            "total_contexts": 0,
+            "total_errors": 0,
+            "method_counts": defaultdict(int),
+            "method_durations": defaultdict(list),
+            "active_contexts": 0,
+            "start_time": time.time(),
+        }
+        self.lock = asyncio.Lock()
+
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        async with self.lock:
+            self.metrics["total_contexts"] += 1
+            self.metrics["active_contexts"] += 1
+            self.metrics["method_counts"][context.method] = self.metrics["method_counts"].get(context.method, 0) + 1
+
+        try:
+            result = await call_next(context)
+            duration = time.time() - context.timestamp
+
+            async with self.lock:
+                self.metrics["active_contexts"] -= 1
+                self._record_duration(context.method, duration)
+
+            return result
+        except Exception:
+            duration = time.time() - context.timestamp
+            async with self.lock:
+                self.metrics["total_errors"] += 1
+                self.metrics["active_contexts"] -= 1
+                self._record_duration(context.method, duration)
+            raise
+
+    def _record_duration(self, method: str, duration: float) -> None:
+        """Record duration for a method in a thread-safe manner."""
+        self.metrics["method_durations"][method].append(duration)
+
+    def get_metrics(self) -> dict[str, Any]:
+        """Get current metrics snapshot."""
+        uptime = time.time() - self.metrics["start_time"]
+
+        return {
+            **self.metrics,
+            "uptime_seconds": uptime,
+            "contexts_per_second": self.metrics["total_contexts"] / uptime if uptime > 0 else 0,
+            "error_rate": self.metrics["total_errors"] / self.metrics["total_contexts"]
+            if self.metrics["total_contexts"] > 0
+            else 0,
+            "method_avg_duration": {
+                method: sum(durations) / len(durations) if durations else 0
+                for method, durations in self.metrics["method_durations"].items()
+            },
+            "method_min_duration": {
+                method: min(durations) if durations else 0
+                for method, durations in self.metrics["method_durations"].items()
+            },
+            "method_max_duration": {
+                method: max(durations) if durations else 0
+                for method, durations in self.metrics["method_durations"].items()
+            },
+        }
+
+
+class PerformanceMetricsMiddleware(Middleware):
+    """Advanced performance metrics including percentiles, throughput, and performance trends."""
+
+    def __init__(self):
+        self.performance_data = {
+            "context_times": defaultdict(list),
+            "hourly_counts": defaultdict(int),
+            "connector_performance": defaultdict(list),
+            "slow_contexts": [],  # contexts over threshold
+            "fast_contexts": [],  # Fastest contexts
+            "slow_threshold_ms": SLOW_THRESHOLD_MS,
+            "fast_threshold_ms": FAST_THRESHOLD_MS,
+        }
+        self.lock = asyncio.Lock()
+
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        start_time = time.time()
+
+        try:
+            result = await call_next(context)
+            duration_ms = (time.time() - start_time) * 1000
+
+            async with self.lock:
+                # Track performance by method and connector
+                self.performance_data["context_times"][context.method].append(duration_ms)
+                self.performance_data["connector_performance"][context.connection_id].append(duration_ms)
+
+                # Track hourly patterns
+                hour = int(time.time() // SECONDS_PER_HOUR)
+                self.performance_data["hourly_counts"][hour] += 1
+
+                # Identify slow/fast contexts
+                if duration_ms > self.performance_data["slow_threshold_ms"]:
+                    self.performance_data["slow_contexts"].append(
+                        {
+                            "method": context.method,
+                            "connector": context.connection_id,
+                            "duration_ms": duration_ms,
+                            "timestamp": context.timestamp,
+                        }
+                    )
+                    # Keep only last MAX_SLOW_CONTEXTS slow contexts
+                    if len(self.performance_data["slow_contexts"]) > MAX_SLOW_CONTEXTS:
+                        self.performance_data["slow_contexts"] = self.performance_data["slow_contexts"][
+                            -MAX_SLOW_CONTEXTS:
+                        ]
+
+                if duration_ms < self.performance_data["fast_threshold_ms"]:
+                    self.performance_data["fast_contexts"].append(
+                        {
+                            "method": context.method,
+                            "connector": context.connection_id,
+                            "duration_ms": duration_ms,
+                            "timestamp": context.timestamp,
+                        }
+                    )
+                    # Keep only last MAX_FAST_CONTEXTS fast contexts
+                    if len(self.performance_data["fast_contexts"]) > MAX_FAST_CONTEXTS:
+                        self.performance_data["fast_contexts"] = self.performance_data["fast_contexts"][
+                            -MAX_FAST_CONTEXTS:
+                        ]
+
+            return result
+
+        except Exception:
+            # Still track duration even for failed contexts
+            duration_ms = (time.time() - start_time) * 1000
+            async with self.lock:
+                self.performance_data["context_times"][context.method].append(duration_ms)
+                self.performance_data["connector_performance"][context.connection_id].append(duration_ms)
+            raise
+
+    def get_performance_metrics(self) -> dict[str, Any]:
+        """Get detailed performance statistics."""
+
+        def calculate_percentiles(values):
+            if not values:
+                return {"p50": 0, "p90": 0, "p95": 0, "p99": 0}
+            sorted_values = sorted(values)
+            n = len(sorted_values)
+            return {
+                "p50": sorted_values[int(n * 0.5)],
+                "p90": sorted_values[int(n * 0.9)],
+                "p95": sorted_values[int(n * 0.95)],
+                "p99": sorted_values[int(n * 0.99)],
+            }
+
+        method_stats = {}
+        for method, times in self.performance_data["context_times"].items():
+            if times:
+                method_stats[method] = {
+                    "count": len(times),
+                    "avg_ms": sum(times) / len(times),
+                    "min_ms": min(times),
+                    "max_ms": max(times),
+                    **calculate_percentiles(times),
+                }
+
+        connector_stats = {}
+        for connector, times in self.performance_data["connector_performance"].items():
+            if times:
+                connector_stats[connector] = {
+                    "count": len(times),
+                    "avg_ms": sum(times) / len(times),
+                    "min_ms": min(times),
+                    "max_ms": max(times),
+                    **calculate_percentiles(times),
+                }
+
+        return {
+            "method_performance": method_stats,
+            "connector_performance": connector_stats,
+            "slow_contexts": self.performance_data["slow_contexts"][-10:],  # Last 10 slow contexts
+            "fast_contexts": self.performance_data["fast_contexts"][-10:],  # Last 10 fast contexts
+            "hourly_distribution": dict(self.performance_data["hourly_counts"]),
+            "thresholds": {
+                "slow_ms": self.performance_data["slow_threshold_ms"],
+                "fast_ms": self.performance_data["fast_threshold_ms"],
+            },
+        }
+
+
+class ErrorTrackingMiddleware(Middleware):
+    """Error tracking and analysis middleware for detailed error analytics."""
+
+    def __init__(self):
+        self.error_data = {
+            "error_counts": Counter(),
+            "error_by_method": defaultdict(Counter),
+            "error_by_connector": defaultdict(Counter),
+            "recent_errors": [],
+            "error_timestamps": [],
+        }
+        self.lock = asyncio.Lock()
+
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        try:
+            return await call_next(context)
+
+        except Exception as e:
+            async with self.lock:
+                error_type = type(e).__name__
+                error_msg = str(e)
+
+                # Track error patterns
+                self.error_data["error_counts"][error_type] += 1
+                self.error_data["error_by_method"][context.method][error_type] += 1
+                self.error_data["error_by_connector"][context.connection_id][error_type] += 1
+                self.error_data["error_timestamps"].append(time.time())
+
+                # Keep recent errors for analysis
+                error_info = {
+                    "timestamp": context.timestamp,
+                    "method": context.method,
+                    "connector": context.connection_id,
+                    "error_type": error_type,
+                    "error_message": error_msg,
+                    "context_id": context.id,
+                }
+                self.error_data["recent_errors"].append(error_info)
+
+                # Keep only last MAX_RECENT_ERRORS errors
+                if len(self.error_data["recent_errors"]) > MAX_RECENT_ERRORS:
+                    self.error_data["recent_errors"] = self.error_data["recent_errors"][-MAX_RECENT_ERRORS:]
+
+                # Keep only last MAX_ERROR_TIMESTAMPS timestamps
+                if len(self.error_data["error_timestamps"]) > MAX_ERROR_TIMESTAMPS:
+                    self.error_data["error_timestamps"] = self.error_data["error_timestamps"][-MAX_ERROR_TIMESTAMPS:]
+
+            raise  # Re-raise the error
+
+    def get_error_analytics(self) -> dict[str, Any]:
+        """Get detailed error analytics."""
+
+        # Calculate error rate over time windows
+        now = time.time()
+        recent_errors = [t for t in self.error_data["error_timestamps"] if now - t < SECONDS_PER_HOUR]  # Last hour
+        very_recent_errors = [
+            t for t in self.error_data["error_timestamps"] if now - t < 5 * SECONDS_PER_MINUTE
+        ]  # Last 5 min
+
+        return {
+            "total_errors": sum(self.error_data["error_counts"].values()),
+            "error_types": dict(self.error_data["error_counts"]),
+            "errors_by_method": {method: dict(errors) for method, errors in self.error_data["error_by_method"].items()},
+            "errors_by_connector": {
+                connector: dict(errors) for connector, errors in self.error_data["error_by_connector"].items()
+            },
+            "recent_errors": self.error_data["recent_errors"][-10:],  # Last 10 errors
+            "error_rate_last_hour": len(recent_errors),
+            "error_rate_last_5min": len(very_recent_errors),
+            "most_common_errors": self.error_data["error_counts"].most_common(5),
+        }
+
+
+class CombinedAnalyticsMiddleware(Middleware):
+    """Comprehensive middleware combining metrics, performance, and error tracking."""
+
+    def __init__(self):
+        self.metrics_mw = MetricsMiddleware()
+        self.perf_mw = PerformanceMetricsMiddleware()
+        self.error_mw = ErrorTrackingMiddleware()
+
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        # Chain the middleware in the desired order: Metrics -> Errors -> Performance
+        async def chain(ctx):
+            # The final call in the chain is the original `call_next`
+            return await self.perf_mw.on_request(ctx, call_next)
+
+        async def error_chain(ctx):
+            return await self.error_mw.on_request(ctx, chain)
+
+        return await self.metrics_mw.on_request(context, error_chain)
+
+    def get_combined_analytics(self) -> dict[str, Any]:
+        """Get all analytics data in one comprehensive report."""
+        return {
+            "metrics": self.metrics_mw.get_metrics(),
+            "performance": self.perf_mw.get_performance_metrics(),
+            "errors": self.error_mw.get_error_analytics(),
+            "generated_at": time.time(),
+        }

mcp_use/client/middleware/middleware.py

@@ -0,0 +1,266 @@
+"""
+Core middleware system for MCP requests.
+
+This module provides a robust and extensible middleware architecture:
+- A typed MiddlewareContext to carry request data.
+- A Middleware base class with a dispatcher that routes to strongly-typed hooks.
+- A MiddlewareManager to build and execute the processing chain.
+- A CallbackClientSession that acts as an adapter, creating the initial context
+  without requiring changes to upstream callers like HttpConnector.
+"""
+
+import time
+import uuid
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from functools import partial
+from typing import Any, Generic, Protocol, TypeVar
+
+from mcp import ClientSession
+from mcp.types import (
+    CallToolRequestParams,
+    CallToolResult,
+    GetPromptRequestParams,
+    GetPromptResult,
+    InitializeRequestParams,
+    InitializeResult,
+    JSONRPCResponse,
+    ListPromptsRequest,
+    ListPromptsResult,
+    ListResourcesRequest,
+    ListResourcesResult,
+    ListToolsRequest,
+    ListToolsResult,
+    ReadResourceRequestParams,
+    ReadResourceResult,
+)
+
+from mcp_use.telemetry.telemetry import telemetry
+
+# Generig TypeVars for context and results
+T = TypeVar("T")
+R = TypeVar("R", covariant=True)
+
+
+@dataclass
+class MiddlewareContext(Generic[T]):
+    """Unified, typed context for all middleware operations."""
+
+    id: str
+    method: str  # The JSON-RPC method name, e.g., "tools/call"
+    params: T  # The typed parameters for the method
+    connection_id: str
+    timestamp: float
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class MCPResponseContext:
+    """Extended context for MCP responses with middleware metadata."""
+
+    request_id: str
+    result: Any
+    error: Exception | None
+    duration: float
+    metadata: dict[str, Any]
+    jsonrpc_response: JSONRPCResponse | None = None
+
+    @classmethod
+    def create(cls, request_id: str, result: Any = None, error: Exception = None) -> "MCPResponseContext":
+        return cls(request_id=request_id, result=result, error=error, duration=0.0, metadata={})
+
+
+# Protocol definition for middleware
+class NextFunctionT(Protocol[T, R]):
+    """Protocol for the `call_next` function passed to middleware."""
+
+    async def __call__(self, context: MiddlewareContext[T]) -> R: ...
+
+
+class Middleware:
+    """Base class for middlewares with hooks."""
+
+    async def __call__(self, context: MiddlewareContext[T], call_next: NextFunctionT[T, Any]) -> Any:
+        """Main entry point that orchestrates the chain"""
+        handler_chain = await self._dispatch_handler(context, call_next)
+        return await handler_chain(context)
+
+    async def _dispatch_handler(
+        self, context: MiddlewareContext[Any], call_next: NextFunctionT[Any, Any]
+    ) -> NextFunctionT[Any, Any]:
+        """Build a chain of handlers"""
+        handler = call_next
+
+        method_map = {
+            "initialize": self.on_initialize,
+            "tools/call": self.on_call_tool,
+            "tools/list": self.on_list_tools,
+            "resources/list": self.on_list_resources,
+            "resources/read": self.on_read_resource,
+            "prompts/list": self.on_list_prompts,
+            "prompts/get": self.on_get_prompt,
+        }
+
+        if hook := method_map.get(context.method):
+            handler = partial(hook, call_next=handler)
+
+        # We can assume that all intercepted calls are requests
+        handler = partial(self.on_request, call_next=handler)
+
+        return handler
+
+    # Default implementations for all hooks
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        return await call_next(context)
+
+    async def on_initialize(
+        self, context: MiddlewareContext[InitializeRequestParams], call_next: NextFunctionT
+    ) -> InitializeResult:
+        return await call_next(context)
+
+    async def on_call_tool(
+        self, context: MiddlewareContext[CallToolRequestParams], call_next: NextFunctionT
+    ) -> CallToolResult:
+        return await call_next(context)
+
+    async def on_read_resource(
+        self, context: MiddlewareContext[ReadResourceRequestParams], call_next: NextFunctionT
+    ) -> ReadResourceResult:
+        return await call_next(context)
+
+    async def on_get_prompt(
+        self, context: MiddlewareContext[GetPromptRequestParams], call_next: NextFunctionT
+    ) -> GetPromptResult:
+        return await call_next(context)
+
+    async def on_list_tools(
+        self, context: MiddlewareContext[ListToolsRequest], call_next: NextFunctionT
+    ) -> ListToolsResult:
+        return await call_next(context)
+
+    async def on_list_resources(
+        self, context: MiddlewareContext[ListResourcesRequest], call_next: NextFunctionT
+    ) -> ListResourcesResult:
+        return await call_next(context)
+
+    async def on_list_prompts(
+        self, context: MiddlewareContext[ListPromptsRequest], call_next: NextFunctionT
+    ) -> ListPromptsResult:
+        return await call_next(context)
+
+
+class MiddlewareManager:
+    """Manages middleware callbacks for MCP requests."""
+
+    def __init__(self):
+        self.middlewares: list[Middleware] = []
+
+    @telemetry("middleware_add")
+    def add_middleware(self, callback: Middleware) -> None:
+        """Add a middleware callback."""
+        self.middlewares.append(callback)
+
+    @telemetry("middleware_process_request")
+    async def process_request(self, context: MiddlewareContext, original_call: Callable) -> MCPResponseContext:
+        """
+        Runs the full middleware chain, captures timing and errors,
+        and returns a structured MCPResponseContext.
+        """
+
+        try:
+            # Chain middleware callbacks
+            async def execute_call(_: MiddlewareContext) -> Any:
+                return await original_call()
+
+            call_chain = execute_call
+            for middleware in reversed(self.middlewares):
+                call_chain = partial(middleware, call_next=call_chain)
+
+            # Execute the chain
+            start_time = time.time()
+
+            # The result of the chain is the reaw result (e.g., CallToolResult)
+            raw_result = await call_chain(context)
+
+            # Success, now wrap the result in response context
+            duration = time.time() - start_time
+
+            response = MCPResponseContext.create(request_id=context.id, result=raw_result)
+            response.duration = duration
+            return response
+
+        except Exception as error:
+            duration = time.time() - context.timestamp
+            response = MCPResponseContext.create(request_id=context.id, error=error)
+            response.duration = duration
+            return response
+
+
+class CallbackClientSession:
+    """ClientSession wrapper that uses callback-based middleware."""
+
+    def __init__(self, client_session: ClientSession, connector_id: str, middleware_manager: MiddlewareManager):
+        self._client_session = client_session
+        self.connector_id = connector_id
+        self.middleware_manager = middleware_manager
+
+    async def _intercept_call(self, method_name: str, params: Any, original_call: Callable) -> Any:
+        """
+        Creates the context, runs it through the manager, and unwraps the final response.
+        """
+        context = MiddlewareContext(
+            id=str(uuid.uuid4()),
+            method=method_name,
+            params=params,
+            connection_id=self.connector_id,
+            timestamp=time.time(),
+        )
+
+        # This now returns a rich MCPResponseContext
+        response_context = await self.middleware_manager.process_request(context, original_call)
+
+        # If there is an error, return it
+        if response_context.error:
+            raise response_context.error
+
+        return response_context.result
+
+    # Wrap all MCP methods with specific params
+    async def initialize(self, *args, **kwargs) -> InitializeResult:
+        return await self._intercept_call("initialize", None, lambda: self._client_session.initialize(*args, **kwargs))
+
+    # List requests usually don't have parameters
+    async def list_tools(self, *args, **kwargs) -> ListToolsResult:
+        return await self._intercept_call("tools/list", None, lambda: self._client_session.list_tools(*args, **kwargs))
+
+    async def call_tool(self, name: str, arguments: dict[str, Any], *args, **kwargs) -> CallToolResult:
+        params = CallToolRequestParams(name=name, arguments=arguments)
+        return await self._intercept_call(
+            "tools/call", params, lambda: self._client_session.call_tool(name, arguments, *args, **kwargs)
+        )
+
+    async def list_resources(self, *args, **kwargs) -> ListResourcesResult:
+        return await self._intercept_call(
+            "resources/list", None, lambda: self._client_session.list_resources(*args, **kwargs)
+        )
+
+    async def read_resource(self, uri: str, *args, **kwargs) -> ReadResourceResult:
+        params = ReadResourceRequestParams(uri=uri)
+        return await self._intercept_call(
+            "resources/read", params, lambda: self._client_session.read_resource(uri, *args, **kwargs)
+        )
+
+    async def list_prompts(self, *args, **kwargs) -> ListPromptsResult:
+        return await self._intercept_call(
+            "prompts/list", None, lambda: self._client_session.list_prompts(*args, **kwargs)
+        )
+
+    async def get_prompt(self, name: str, *args, **kwargs) -> GetPromptResult:
+        params = GetPromptRequestParams(name=name, **kwargs)
+        return await self._intercept_call(
+            "prompts/get", params, lambda: self._client_session.get_prompt(name, *args, **kwargs)
+        )
+
+    def __getattr__(self, name: str) -> Any:
+        """Delegate other attributes to the wrapped session."""
+        return getattr(self._client_session, name)