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

mcp_use/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 .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/middleware/middleware.py

@@ -0,0 +1,262 @@
+"""
+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,
+)
+
+# 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] = []
+
+    def add_middleware(self, callback: Middleware) -> None:
+        """Add a middleware callback."""
+        self.middlewares.append(callback)
+
+    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)

mcp_use/task_managers/base.py

@@ -22,13 +22,14 @@ class ConnectionManager(Generic[T], ABC):
     used with MCP connectors.
     """
 
-    def __init__(self):
+    def __init__(self) -> None:
         """Initialize a new connection manager."""
         self._ready_event = asyncio.Event()
         self._done_event = asyncio.Event()
+        self._stop_event = asyncio.Event()
         self._exception: Exception | None = None
         self._connection: T | None = None
-        self._task: asyncio.Task | None = None
+        self._task: asyncio.Task[None] | None = None
 
     @abstractmethod
     async def _establish_connection(self) -> T:
@@ -86,20 +87,15 @@ class ConnectionManager(Generic[T], ABC):
 
     async def stop(self) -> None:
         """Stop the connection manager and close the connection."""
+        # Signal stop to the connection task instead of cancelling it, avoids
+        # propagating CancelledError to unrelated tasks.
         if self._task and not self._task.done():
-            # Cancel the task
-            logger.debug(f"Cancelling {self.__class__.__name__} task")
-            self._task.cancel()
-
-            # Wait for it to complete
-            try:
-                await self._task
-            except asyncio.CancelledError:
-                logger.debug(f"{self.__class__.__name__} task cancelled successfully")
-            except Exception as e:
-                logger.warning(f"Error stopping {self.__class__.__name__} task: {e}")
-
-        # Wait for the connection to be done
+            logger.debug(f"Signaling stop to {self.__class__.__name__} task")
+            self._stop_event.set()
+            # Wait for it to finish gracefully
+            await self._task
+
+        # Ensure cleanup completed
         await self._done_event.wait()
         logger.debug(f"{self.__class__.__name__} task completed")
 
@@ -125,14 +121,8 @@ class ConnectionManager(Generic[T], ABC):
             # Signal that the connection is ready
             self._ready_event.set()
 
-            # Wait indefinitely until cancelled
-            try:
-                # This keeps the connection open until cancelled
-                await asyncio.Event().wait()
-            except asyncio.CancelledError:
-                # Expected when stopping
-                logger.debug(f"{self.__class__.__name__} task received cancellation")
-                pass
+            # Wait until stop is requested
+            await self._stop_event.wait()
 
         except Exception as e:
             # Store the exception

mcp_use/task_managers/sse.py

@@ -7,6 +7,7 @@ that ensures proper task isolation and resource cleanup.
 
 from typing import Any
 
+import httpx
 from mcp.client.sse import sse_client
 
 from ..logging import logger
@@ -27,6 +28,7 @@ class SseConnectionManager(ConnectionManager[tuple[Any, Any]]):
         headers: dict[str, str] | None = None,
         timeout: float = 5,
         sse_read_timeout: float = 60 * 5,
+        auth: httpx.Auth | None = None,
     ):
         """Initialize a new SSE connection manager.
 
@@ -35,12 +37,14 @@ class SseConnectionManager(ConnectionManager[tuple[Any, Any]]):
             headers: Optional HTTP headers
             timeout: Timeout for HTTP operations in seconds
             sse_read_timeout: Timeout for SSE read operations in seconds
+            auth: Optional httpx.Auth instance for authentication
         """
         super().__init__()
         self.url = url
         self.headers = headers or {}
         self.timeout = timeout
         self.sse_read_timeout = sse_read_timeout
+        self.auth = auth
         self._sse_ctx = None
 
     async def _establish_connection(self) -> tuple[Any, Any]:
@@ -58,6 +62,7 @@ class SseConnectionManager(ConnectionManager[tuple[Any, Any]]):
             headers=self.headers,
             timeout=self.timeout,
             sse_read_timeout=self.sse_read_timeout,
+            auth=self.auth,
         )
 
         # Enter the context manager