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

mcp_use/auth/oauth_callback.py

@@ -0,0 +1,214 @@
+"""OAuth callback server implementation."""
+
+import asyncio
+from dataclasses import dataclass
+
+import anyio
+import uvicorn
+from starlette.applications import Starlette
+from starlette.requests import Request
+from starlette.responses import HTMLResponse
+from starlette.routing import Route
+
+from ..logging import logger
+
+
+@dataclass
+class CallbackResponse:
+    """Response data from OAuth callback."""
+
+    code: str | None = None  # Authorization code (success)
+    state: str | None = None  # CSRF protection token
+    error: str | None = None  # Errors code (if failed)
+    error_description: str | None = None
+    error_uri: str | None = None
+
+
+class OAuthCallbackServer:
+    """Local server to handle OAuth callback."""
+
+    def __init__(self, port: int):
+        """Initialize the callback server.
+
+        Args:
+            port: Port to listen on.
+        """
+        self.port = port
+        self.redirect_uri: str | None = None
+        # Thread safe way to pass callback data to the main OAuth flow
+        self.response_queue: asyncio.Queue[CallbackResponse] = asyncio.Queue(maxsize=1)
+        self.server: uvicorn.Server | None = None
+        self._shutdown_event = anyio.Event()
+
+    async def start(self) -> str:
+        """Start the callback server and return the redirect URI."""
+        app = self._create_app()
+
+        # Create the server
+        config = uvicorn.Config(
+            app,
+            host="127.0.0.1",
+            port=self.port,
+            log_level="error",  # Suppress uvicorn logs
+        )
+        self.server = uvicorn.Server(config)
+
+        # Start server in background
+        self._server_task = asyncio.create_task(self.server.serve())
+
+        # Wait a moment for server to start
+        await asyncio.sleep(0.1)
+
+        self.redirect_uri = f"http://localhost:{self.port}/callback"
+        return self.redirect_uri
+
+    async def wait_for_code(self, timeout: float = 300) -> CallbackResponse:
+        """Wait for the OAuth callback with a timeout (default 5 minutes)."""
+        try:
+            response = await asyncio.wait_for(self.response_queue.get(), timeout=timeout)
+            return response
+        except TimeoutError:
+            raise TimeoutError(f"OAuth callback not received within {timeout} seconds") from None
+        finally:
+            await self.shutdown()
+
+    async def shutdown(self):
+        """Shutdown the callback server."""
+        self._shutdown_event.set()
+        if self.server:
+            self.server.should_exit = True
+            if hasattr(self, "_server_task"):
+                try:
+                    await asyncio.wait_for(self._server_task, timeout=5.0)
+                except TimeoutError:
+                    self._server_task.cancel()
+
+    def _create_app(self) -> Starlette:
+        """Create the Starlette application."""
+
+        async def callback(request: Request) -> HTMLResponse:
+            """Handle the OAuth callback."""
+            params = request.query_params
+
+            # Extract OAuth parameters
+            response = CallbackResponse(
+                code=params.get("code"),
+                state=params.get("state"),
+                error=params.get("error"),
+                error_description=params.get("error_description"),
+                error_uri=params.get("error_uri"),
+            )
+
+            # Log the callback response
+            logger.debug(
+                f"OAuth callback received: error={response.error}, error_description={response.error_description}"
+            )
+            if response.code:
+                logger.debug("OAuth callback received authorization code")
+            else:
+                logger.error(f"OAuth callback error: {response.error} - {response.error_description}")
+
+            # Put response in queue
+            try:
+                self.response_queue.put_nowait(response)
+            except asyncio.QueueFull:
+                pass  # Ignore if queue is already full
+
+            # Return success page
+            if response.code:
+                html = self._success_html()
+            else:
+                html = self._error_html(response.error, response.error_description)
+
+            return HTMLResponse(content=html)
+
+        routes = [Route("/callback", callback)]
+        return Starlette(routes=routes)
+
+    def _success_html(self) -> str:
+        """HTML response for successful authorization."""
+        return """
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <title>Authorization Successful</title>
+            <style>
+                body {
+                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+                    display: flex;
+                    justify-content: center;
+                    align-items: center;
+                    height: 100vh;
+                    margin: 0;
+                    background-color: #f5f5f5;
+                }
+                .container {
+                    text-align: center;
+                    padding: 2rem;
+                    background: white;
+                    border-radius: 8px;
+                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+                }
+                h1 { color: #22c55e; margin-bottom: 0.5rem; }
+                p { color: #666; margin-top: 0.5rem; }
+                .icon { font-size: 48px; margin-bottom: 1rem; }
+            </style>
+        </head>
+        <body>
+            <div class="container">
+                <div class="icon">✅</div>
+                <h1>Authorization Successful!</h1>
+                <p>You can now close this window and return to your application.</p>
+            </div>
+            <script>
+                // Auto-close after 3 seconds
+                setTimeout(() => window.close(), 3000);
+            </script>
+        </body>
+        </html>
+        """
+
+    def _error_html(self, error: str | None, description: str | None) -> str:
+        """HTML response for authorization error."""
+        error_msg = error or "Unknown error"
+        desc_msg = description or "Authorization was not completed successfully."
+
+        return f"""
+        <!DOCTYPE html>
+        <html>
+        <head>
+            <title>Authorization Error</title>
+            <style>
+                body {{
+                    font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
+                    display: flex;
+                    justify-content: center;
+                    align-items: center;
+                    height: 100vh;
+                    margin: 0;
+                    background-color: #f5f5f5;
+                }}
+                .container {{
+                    text-align: center;
+                    padding: 2rem;
+                    background: white;
+                    border-radius: 8px;
+                    box-shadow: 0 2px 4px rgba(0,0,0,0.1);
+                    max-width: 500px;
+                }}
+                h1 {{ color: #ef4444; margin-bottom: 0.5rem; }}
+                .error {{ color: #dc2626; font-weight: 600; margin: 1rem 0; }}
+                .description {{ color: #666; margin-top: 0.5rem; }}
+                .icon {{ font-size: 48px; margin-bottom: 1rem; }}
+            </style>
+        </head>
+        <body>
+            <div class="container">
+                <div class="icon">❌</div>
+                <h1>Authorization Error</h1>
+                <p class="error">{error_msg}</p>
+                <p class="description">{desc_msg}</p>
+            </div>
+        </body>
+        </html>
+        """

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.
 
@@ -192,7 +215,7 @@ class MCPClient:
 
         server_config = servers[server_name]
 
-        # Create connector with options
+        # Create connector with options and client-level auth
         connector = create_connector_from_config(
             server_config,
             sandbox=self.sandbox,
@@ -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
@@ -79,13 +83,14 @@ def create_connector_from_config(
         return HttpConnector(
             base_url=server_config["url"],
             headers=server_config.get("headers", None),
-            auth_token=server_config.get("auth_token", None),
+            auth=server_config.get("auth", {}),
             timeout=server_config.get("timeout", 5),
             sse_read_timeout=server_config.get("sse_read_timeout", 60 * 5),
             sampling_callback=sampling_callback,
             elicitation_callback=elicitation_callback,
             message_handler=message_handler,
             logging_callback=logging_callback,
+            middleware=middleware,
         )
 
     # WebSocket connector
@@ -93,7 +98,7 @@ def create_connector_from_config(
         return WebSocketConnector(
             url=server_config["ws_url"],
             headers=server_config.get("headers", None),
-            auth_token=server_config.get("auth_token", None),
+            auth=server_config.get("auth", {}),
         )
 
     raise ValueError("Cannot determine connector type from config")

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."""
@@ -111,29 +119,34 @@ class BaseConnector(ABC):
         """Clean up all resources associated with this connector."""
         errors = []
 
-        # First close the client session
-        if self.client_session:
+        # First stop the connection manager, this closes the ClientSession inside
+        # the same task where it was opened, avoiding cancel-scope mismatches.
+        if self._connection_manager:
             try:
-                logger.debug("Closing client session")
-                await self.client_session.__aexit__(None, None, None)
+                logger.debug("Stopping connection manager")
+                await self._connection_manager.stop()
             except Exception as e:
-                error_msg = f"Error closing client session: {e}"
+                error_msg = f"Error stopping connection manager: {e}"
                 logger.warning(error_msg)
                 errors.append(error_msg)
             finally:
-                self.client_session = None
+                self._connection_manager = None
 
-        # Then stop the connection manager
-        if self._connection_manager:
+        # Ensure the client_session reference is cleared (it should already be
+        # closed by the connection manager). Only attempt a direct __aexit__ if
+        # the connection manager did *not* exist, this covers edge-cases like
+        # failed connections where no manager was started.
+        if self.client_session:
             try:
-                logger.debug("Stopping connection manager")
-                await self._connection_manager.stop()
+                if not self._connection_manager:
+                    logger.debug("Closing client session (no connection manager)")
+                    await self.client_session.__aexit__(None, None, None)
             except Exception as e:
-                error_msg = f"Error stopping connection manager: {e}"
+                error_msg = f"Error closing client session: {e}"
                 logger.warning(error_msg)
                 errors.append(error_msg)
             finally:
-                self._connection_manager = None
+                self.client_session = None
 
         # Reset tools
         self._tools = None