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

mcp_use/client.py

@@ -15,6 +15,7 @@ from mcp_use.types.sandbox import SandboxOptions
 
 from .config import create_connector_from_config, load_config_file
 from .logging import logger
+from .middleware import Middleware, default_logging_middleware
 from .session import MCPSession
 
 
@@ -35,6 +36,7 @@ class MCPClient:
         elicitation_callback: ElicitationFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         logging_callback: LoggingFnT | None = None,
+        middleware: list[Middleware] | None = None,
     ) -> None:
         """Initialize a new MCP client.
 
@@ -55,6 +57,12 @@ class MCPClient:
         self.elicitation_callback = elicitation_callback
         self.message_handler = message_handler
         self.logging_callback = logging_callback
+        # Add default logging middleware if no middleware provided, or prepend it to existing middleware
+        default_middleware = [default_logging_middleware]
+        if middleware:
+            self.middleware = default_middleware + middleware
+        else:
+            self.middleware = default_middleware
         # Load configuration if provided
         if config is not None:
             if isinstance(config, str):
@@ -151,6 +159,21 @@ class MCPClient:
             if name in self.active_sessions:
                 self.active_sessions.remove(name)
 
+    def add_middleware(self, middleware: Middleware) -> None:
+        """Add a middleware.
+
+        Args:
+            middleware: The middleware to add
+        """
+        if len(self.sessions) == 0 and middleware not in self.middleware:
+            self.middleware.append(middleware)
+            return
+
+        if middleware not in self.middleware:
+            self.middleware.append(middleware)
+            for session in self.sessions.values():
+                session.connector.middleware_manager.add_middleware(middleware)
+
     def get_server_names(self) -> list[str]:
         """Get the list of configured server names.
 
@@ -201,6 +224,7 @@ class MCPClient:
             elicitation_callback=self.elicitation_callback,
             message_handler=self.message_handler,
             logging_callback=self.logging_callback,
+            middleware=self.middleware,
         )
 
         # Create the session

mcp_use/config.py

@@ -13,6 +13,7 @@ from mcp_use.types.sandbox import SandboxOptions
 
 from .connectors import BaseConnector, HttpConnector, SandboxConnector, StdioConnector, WebSocketConnector
 from .connectors.utils import is_stdio_server
+from .middleware import Middleware
 
 
 def load_config_file(filepath: str) -> dict[str, Any]:
@@ -36,6 +37,7 @@ def create_connector_from_config(
     elicitation_callback: ElicitationFnT | None = None,
     message_handler: MessageHandlerFnT | None = None,
     logging_callback: LoggingFnT | None = None,
+    middleware: list[Middleware] | None = None,
 ) -> BaseConnector:
     """Create a connector based on server configuration.
     This function can be called with just the server_config parameter:
@@ -59,6 +61,7 @@ def create_connector_from_config(
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
 
     # Sandboxed connector
@@ -72,6 +75,7 @@ def create_connector_from_config(
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
 
     # HTTP connector
@@ -86,6 +90,7 @@ def create_connector_from_config(
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
 
     # WebSocket connector

mcp_use/connectors/base.py

@@ -31,6 +31,7 @@ from pydantic import AnyUrl
 import mcp_use
 
 from ..logging import logger
+from ..middleware import Middleware, MiddlewareManager
 from ..task_managers import ConnectionManager
 
 
@@ -46,6 +47,7 @@ class BaseConnector(ABC):
         elicitation_callback: ElicitationFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         logging_callback: LoggingFnT | None = None,
+        middleware: list[Middleware] | None = None,
     ):
         """Initialize base connector with common attributes."""
         self.client_session: ClientSession | None = None
@@ -62,6 +64,12 @@ class BaseConnector(ABC):
         self.logging_callback = logging_callback
         self.capabilities: ServerCapabilities | None = None
 
+        # Set up middleware manager
+        self.middleware_manager = MiddlewareManager()
+        if middleware:
+            for mw in middleware:
+                self.middleware_manager.add_middleware(mw)
+
     @property
     def client_info(self) -> Implementation:
         """Get the client info for the connector."""

mcp_use/connectors/http.py

@@ -17,6 +17,7 @@ from mcp_use.auth.oauth import OAuthClientProvider
 from ..auth import BearerAuth, OAuth
 from ..exceptions import OAuthAuthenticationError, OAuthDiscoveryError
 from ..logging import logger
+from ..middleware import CallbackClientSession, Middleware
 from ..task_managers import SseConnectionManager, StreamableHttpConnectionManager
 from .base import BaseConnector
 
@@ -39,6 +40,7 @@ class HttpConnector(BaseConnector):
         elicitation_callback: ElicitationFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         logging_callback: LoggingFnT | None = None,
+        middleware: list[Middleware] | None = None,
     ):
         """Initialize a new HTTP connector.
 
@@ -59,6 +61,7 @@ class HttpConnector(BaseConnector):
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
         self.base_url = base_url.rstrip("/")
         self.headers = headers or {}
@@ -158,7 +161,7 @@ class HttpConnector(BaseConnector):
             read_stream, write_stream = await connection_manager.start()
 
             # Test if this actually works by trying to create a client session and initialize it
-            test_client = ClientSession(
+            raw_test_client = ClientSession(
                 read_stream,
                 write_stream,
                 sampling_callback=self.sampling_callback,
@@ -167,7 +170,10 @@ class HttpConnector(BaseConnector):
                 logging_callback=self.logging_callback,
                 client_info=self.client_info,
             )
-            await test_client.__aenter__()
+            await raw_test_client.__aenter__()
+
+            # Wrap test client with middleware temporarily for testing
+            test_client = CallbackClientSession(raw_test_client, self.public_identifier, self.middleware_manager)
 
             try:
                 # Try to initialize - this is where streamable HTTP vs SSE difference should show up
@@ -209,7 +215,7 @@ class HttpConnector(BaseConnector):
                 logger.error("MCP protocol error during initialization: %s", mcp_error.error)
                 # Clean up the test client
                 try:
-                    await test_client.__aexit__(None, None, None)
+                    await raw_test_client.__aexit__(None, None, None)
                 except Exception:
                     pass
                 raise mcp_error
@@ -218,7 +224,7 @@ class HttpConnector(BaseConnector):
                 # This catches non-McpError exceptions, like a direct httpx timeout
                 # but in the most cases this won't happen. It's for safety.
                 try:
-                    await test_client.__aexit__(None, None, None)
+                    await raw_test_client.__aexit__(None, None, None)
                 except Exception:
                     pass
                 raise init_error
@@ -251,7 +257,7 @@ class HttpConnector(BaseConnector):
                     read_stream, write_stream = await connection_manager.start()
 
                     # Create the client session for SSE
-                    self.client_session = ClientSession(
+                    raw_client_session = ClientSession(
                         read_stream,
                         write_stream,
                         sampling_callback=self.sampling_callback,
@@ -260,7 +266,12 @@ class HttpConnector(BaseConnector):
                         logging_callback=self.logging_callback,
                         client_info=self.client_info,
                     )
-                    await self.client_session.__aenter__()
+                    await raw_client_session.__aenter__()
+
+                    # Wrap with middleware
+                    self.client_session = CallbackClientSession(
+                        raw_client_session, self.public_identifier, self.middleware_manager
+                    )
                     self.transport_type = "SSE"
 
                 except* Exception as sse_error:
@@ -290,4 +301,5 @@ class HttpConnector(BaseConnector):
     @property
     def public_identifier(self) -> str:
         """Get the identifier for the connector."""
-        return {"type": self.transport_type, "base_url": self.base_url}
+        transport_type = getattr(self, "transport_type", "http")
+        return f"{transport_type}:{self.base_url}"

mcp_use/connectors/sandbox.py

@@ -15,6 +15,7 @@ from mcp import ClientSession
 from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
 
 from ..logging import logger
+from ..middleware import CallbackClientSession, Middleware
 from ..task_managers import SseConnectionManager
 
 # Import E2B SDK components (optional dependency)
@@ -52,6 +53,7 @@ class SandboxConnector(BaseConnector):
         elicitation_callback: ElicitationFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         logging_callback: LoggingFnT | None = None,
+        middleware: list[Middleware] | None = None,
     ):
         """Initialize a new sandbox connector.
 
@@ -71,6 +73,7 @@ class SandboxConnector(BaseConnector):
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
         if Sandbox is None:
             raise ImportError(
@@ -226,7 +229,7 @@ class SandboxConnector(BaseConnector):
             read_stream, write_stream = await self._connection_manager.start()
 
             # Create the client session
-            self.client_session = ClientSession(
+            raw_client_session = ClientSession(
                 read_stream,
                 write_stream,
                 sampling_callback=self.sampling_callback,
@@ -235,7 +238,12 @@ class SandboxConnector(BaseConnector):
                 logging_callback=self.logging_callback,
                 client_info=self.client_info,
             )
-            await self.client_session.__aenter__()
+            await raw_client_session.__aenter__()
+
+            # Wrap with middleware
+            self.client_session = CallbackClientSession(
+                raw_client_session, self.public_identifier, self.middleware_manager
+            )
 
             # Mark as connected
             self._connected = True
@@ -299,4 +307,5 @@ class SandboxConnector(BaseConnector):
     @property
     def public_identifier(self) -> str:
         """Get the identifier for the connector."""
-        return {"type": "sandbox", "command": self.user_command, "args": self.user_args}
+        args_str = " ".join(self.user_args) if self.user_args else ""
+        return f"sandbox:{self.user_command} {args_str}".strip()

mcp_use/connectors/stdio.py

@@ -11,6 +11,7 @@ from mcp import ClientSession, StdioServerParameters
 from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
 
 from ..logging import logger
+from ..middleware import CallbackClientSession, Middleware
 from ..task_managers import StdioConnectionManager
 from .base import BaseConnector
 
@@ -33,6 +34,7 @@ class StdioConnector(BaseConnector):
         elicitation_callback: ElicitationFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
         logging_callback: LoggingFnT | None = None,
+        middleware: list[Middleware] | None = None,
     ):
         """Initialize a new stdio connector.
 
@@ -49,6 +51,7 @@ class StdioConnector(BaseConnector):
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
         self.command = command
         self.args = args or []  # Ensure args is never None
@@ -71,7 +74,7 @@ class StdioConnector(BaseConnector):
             read_stream, write_stream = await self._connection_manager.start()
 
             # Create the client session
-            self.client_session = ClientSession(
+            raw_client_session = ClientSession(
                 read_stream,
                 write_stream,
                 sampling_callback=self.sampling_callback,
@@ -80,7 +83,12 @@ class StdioConnector(BaseConnector):
                 logging_callback=self.logging_callback,
                 client_info=self.client_info,
             )
-            await self.client_session.__aenter__()
+            await raw_client_session.__aenter__()
+
+            # Wrap with middleware
+            self.client_session = CallbackClientSession(
+                raw_client_session, self.public_identifier, self.middleware_manager
+            )
 
             # Mark as connected
             self._connected = True
@@ -98,4 +106,4 @@ class StdioConnector(BaseConnector):
     @property
     def public_identifier(self) -> str:
         """Get the identifier for the connector."""
-        return {"type": "stdio", "command&args": f"{self.command} {' '.join(self.args)}"}
+        return f"stdio:{self.command} {' '.join(self.args)}"

mcp_use/connectors/websocket.py

@@ -254,4 +254,4 @@ class WebSocketConnector(BaseConnector):
     @property
     def public_identifier(self) -> str:
         """Get the identifier for the connector."""
-        return {"type": "websocket", "url": self.url}
+        return f"websocket:{self.url}"

mcp_use/middleware/__init__.py

@@ -0,0 +1,50 @@
+"""
+Middleware package for MCP request interception and processing.
+
+This package provides a flexible middleware system for intercepting MCP requests
+and responses, enabling logging, metrics, caching, and custom processing.
+
+The middleware system follows an Express.js-style pattern where middleware functions
+receive a request context and a call_next function, allowing them to process both
+incoming requests and outgoing responses.
+"""
+
+# Core middleware implementation
+# Default logging middleware
+from .logging import default_logging_middleware
+
+# Metrics middleware classes
+from .metrics import (
+    CombinedAnalyticsMiddleware,
+    ErrorTrackingMiddleware,
+    MetricsMiddleware,
+    PerformanceMetricsMiddleware,
+)
+
+# Protocol types for type-safe middleware
+from .middleware import (
+    CallbackClientSession,
+    MCPResponseContext,
+    Middleware,
+    MiddlewareContext,
+    MiddlewareManager,
+    NextFunctionT,
+)
+
+__all__ = [
+    # Core types and classes
+    "MiddlewareContext",
+    "MCPResponseContext",
+    "Middleware",
+    "MiddlewareManager",
+    "CallbackClientSession",
+    # Protocol types
+    "NextFunctionT",
+    # Default logging middleware
+    "default_logging_middleware",
+    # Metrics middleware
+    "MetricsMiddleware",
+    "PerformanceMetricsMiddleware",
+    "ErrorTrackingMiddleware",
+    "CombinedAnalyticsMiddleware",
+]

mcp_use/middleware/logging.py

@@ -0,0 +1,31 @@
+"""
+Default logging middleware for MCP requests.
+
+Simple debug logging for all MCP requests and responses.
+"""
+
+import time
+from typing import Any
+
+from ..logging import logger
+from .middleware import Middleware, MiddlewareContext, NextFunctionT
+
+
+class LoggingMiddleware(Middleware):
+    """Default logging middleware that logs all MCP requests and responses with logger.debug."""
+
+    async def on_request(self, context: MiddlewareContext[Any], call_next: NextFunctionT) -> Any:
+        """Logs all MCP requests and responses with logger.debug."""
+        logger.debug(f"[{context.id}] {context.connection_id} -> {context.method}")
+        try:
+            result = await call_next(context)
+            duration = time.time() - context.timestamp
+            logger.debug(f"[{context.id}] {context.connection_id} <- {context.method} ({duration:.3f}s)")
+            return result
+        except Exception as e:
+            duration = time.time() - context.timestamp
+            logger.debug(f"[{context.id}] {context.connection_id} <- {context.method} FAILED ({duration:.3f}s): {e}")
+            raise
+
+
+default_logging_middleware = LoggingMiddleware()