mcp-use 1.3.10__py3-none-any.whl → 1.3.11__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

@@ -192,7 +192,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,

mcp_use/config.py

@@ -79,7 +79,7 @@ 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,
@@ -93,7 +93,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

@@ -111,29 +111,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

mcp_use/connectors/http.py

@@ -5,10 +5,17 @@ This module provides a connector for communicating with MCP implementations
 through HTTP APIs with SSE or Streamable HTTP for transport.
 """
 
+from typing import Any
+
 import httpx
 from mcp import ClientSession
 from mcp.client.session import ElicitationFnT, LoggingFnT, MessageHandlerFnT, SamplingFnT
+from mcp.shared.exceptions import McpError
+
+from mcp_use.auth.oauth import OAuthClientProvider
 
+from ..auth import BearerAuth, OAuth
+from ..exceptions import OAuthAuthenticationError, OAuthDiscoveryError
 from ..logging import logger
 from ..task_managers import SseConnectionManager, StreamableHttpConnectionManager
 from .base import BaseConnector
@@ -24,10 +31,10 @@ class HttpConnector(BaseConnector):
     def __init__(
         self,
         base_url: str,
-        auth_token: str | None = None,
         headers: dict[str, str] | None = None,
         timeout: float = 5,
         sse_read_timeout: float = 60 * 5,
+        auth: str | dict[str, Any] | httpx.Auth | None = None,
         sampling_callback: SamplingFnT | None = None,
         elicitation_callback: ElicitationFnT | None = None,
         message_handler: MessageHandlerFnT | None = None,
@@ -37,10 +44,13 @@ class HttpConnector(BaseConnector):
 
         Args:
             base_url: The base URL of the MCP HTTP API.
-            auth_token: Optional authentication token.
             headers: Optional additional headers.
             timeout: Timeout for HTTP operations in seconds.
             sse_read_timeout: Timeout for SSE read operations in seconds.
+            auth: Authentication method - can be:
+                - A string token: Use Bearer token authentication
+                - A dict with OAuth config: {"client_id": "...", "client_secret": "...", "scope": "..."}
+                - An httpx.Auth object: Use custom authentication
             sampling_callback: Optional sampling callback.
             elicitation_callback: Optional elicitation callback.
         """
@@ -51,12 +61,57 @@ class HttpConnector(BaseConnector):
             logging_callback=logging_callback,
         )
         self.base_url = base_url.rstrip("/")
-        self.auth_token = auth_token
         self.headers = headers or {}
-        if auth_token:
-            self.headers["Authorization"] = f"Bearer {auth_token}"
         self.timeout = timeout
         self.sse_read_timeout = sse_read_timeout
+        self._auth: httpx.Auth | None = None
+        self._oauth: OAuth | None = None
+
+        # Handle authentication
+        if auth is not None:
+            self._set_auth(auth)
+
+    def _set_auth(self, auth: str | dict[str, Any] | httpx.Auth) -> None:
+        """Set authentication method.
+
+        Args:
+            auth: Authentication method - can be:
+                - A string token: Use Bearer token authentication
+                - A dict with OAuth config: {"client_id": "...", "client_secret": "...", "scope": "..."}
+                - An httpx.Auth object: Use custom authentication
+        """
+        if isinstance(auth, str):
+            # Treat as bearer token
+            self._auth = BearerAuth(token=auth)
+            self.headers["Authorization"] = f"Bearer {auth}"
+        elif isinstance(auth, dict):
+            # Check if this is an OAuth provider configuration
+            if "oauth_provider" in auth:
+                oauth_provider = auth["oauth_provider"]
+                if isinstance(oauth_provider, dict):
+                    oauth_provider = OAuthClientProvider(**oauth_provider)
+                self._oauth = OAuth(
+                    self.base_url,
+                    scope=auth.get("scope"),
+                    client_id=auth.get("client_id"),
+                    client_secret=auth.get("client_secret"),
+                    callback_port=auth.get("callback_port"),
+                    oauth_provider=oauth_provider,
+                )
+                self._oauth_config = auth
+            else:
+                self._oauth = OAuth(
+                    self.base_url,
+                    scope=auth.get("scope"),
+                    client_id=auth.get("client_id"),
+                    client_secret=auth.get("client_secret"),
+                    callback_port=auth.get("callback_port"),
+                )
+                self._oauth_config = auth
+        elif isinstance(auth, httpx.Auth):
+            self._auth = auth
+        else:
+            raise ValueError(f"Invalid auth type: {type(auth)}")
 
     async def connect(self) -> None:
         """Establish a connection to the MCP implementation."""
@@ -64,6 +119,29 @@ class HttpConnector(BaseConnector):
             logger.debug("Already connected to MCP implementation")
             return
 
+        # Handle OAuth if needed
+        if self._oauth:
+            try:
+                # Create a temporary client for OAuth metadata discovery
+                async with httpx.AsyncClient() as client:
+                    bearer_auth = await self._oauth.initialize(client)
+                    if not bearer_auth:
+                        # Need to perform OAuth flow
+                        logger.info("OAuth authentication required")
+                        bearer_auth = await self._oauth.authenticate()
+
+                    # Update auth and headers
+                    self._auth = bearer_auth
+                    self.headers["Authorization"] = f"Bearer {bearer_auth.token.get_secret_value()}"
+            except OAuthDiscoveryError:
+                # OAuth discovery failed - it means server doesn't support OAuth default urls
+                logger.debug("OAuth discovery failed, continuing without initialization.")
+                self._oauth = None
+                self._auth = None
+            except OAuthAuthenticationError as e:
+                logger.error(f"OAuth initialization failed: {e}")
+                raise
+
         # Try streamable HTTP first (new transport), fall back to SSE (old transport)
         # This implements backwards compatibility per MCP specification
         self.transport_type = None
@@ -73,7 +151,7 @@ class HttpConnector(BaseConnector):
             # First, try the new streamable HTTP transport
             logger.debug(f"Attempting streamable HTTP connection to: {self.base_url}")
             connection_manager = StreamableHttpConnectionManager(
-                self.base_url, self.headers, self.timeout, self.sse_read_timeout
+                self.base_url, self.headers, self.timeout, self.sse_read_timeout, auth=self._auth
             )
 
             # Test if this is a streamable HTTP server by attempting initialization
@@ -94,9 +172,9 @@ class HttpConnector(BaseConnector):
             try:
                 # Try to initialize - this is where streamable HTTP vs SSE difference should show up
                 result = await test_client.initialize()
+                logger.debug(f"Streamable HTTP initialization result: {result}")
 
                 # If we get here, streamable HTTP works
-
                 self.client_session = test_client
                 self.transport_type = "streamable HTTP"
                 self._initialized = True  # Mark as initialized since we just called initialize()
@@ -125,14 +203,28 @@ class HttpConnector(BaseConnector):
                 else:
                     self._prompts = []
 
-            except Exception as init_error:
+            # Only McpError is raised from client's initialization because
+            # exceptions are handled internally.
+            except McpError as mcp_error:
+                logger.error("MCP protocol error during initialization: %s", mcp_error.error)
                 # Clean up the test client
+                try:
+                    await test_client.__aexit__(None, None, None)
+                except Exception:
+                    pass
+                raise mcp_error
+
+            except Exception as init_error:
+                # 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)
                 except Exception:
                     pass
                 raise init_error
 
+        # Exception from the inner try is propagated here and in
+        # the most cases is an McpError, so checking instances is useless
         except Exception as streamable_error:
             logger.debug(f"Streamable HTTP failed: {streamable_error}")
 
@@ -143,24 +235,17 @@ class HttpConnector(BaseConnector):
                 except Exception:
                     pass
 
-            # Check if this is a 4xx error that indicates we should try SSE fallback
-            should_fallback = False
-            if isinstance(streamable_error, httpx.HTTPStatusError):
-                if streamable_error.response.status_code in [404, 405]:
-                    should_fallback = True
-            elif "405 Method Not Allowed" in str(streamable_error) or "404 Not Found" in str(streamable_error):
-                should_fallback = True
-            else:
-                # For other errors, still try fallback but they might indicate
-                # real connectivity issues
-                should_fallback = True
+            # It doesn't make sense to check error types. Because client
+            # always return a McpError, if he can't reach the server
+            # because it's offline, or if it has an auth problem.
+            should_fallback = True
 
             if should_fallback:
                 try:
                     # Fall back to the old SSE transport
                     logger.debug(f"Attempting SSE fallback connection to: {self.base_url}")
                     connection_manager = SseConnectionManager(
-                        self.base_url, self.headers, self.timeout, self.sse_read_timeout
+                        self.base_url, self.headers, self.timeout, self.sse_read_timeout, auth=self._auth
                     )
 
                     read_stream, write_stream = await connection_manager.start()
@@ -178,7 +263,18 @@ class HttpConnector(BaseConnector):
                     await self.client_session.__aenter__()
                     self.transport_type = "SSE"
 
-                except Exception as sse_error:
+                except* Exception as sse_error:
+                    # Get the exception from the ExceptionGroup, and here we will get the correct type.
+                    sse_error = sse_error.exceptions[0]
+                    if isinstance(sse_error, httpx.HTTPStatusError) and sse_error.response.status_code in [
+                        401,
+                        403,
+                        407,
+                    ]:
+                        raise OAuthAuthenticationError(
+                            f"Server requires authentication (HTTP {sse_error.response.status_code}) "
+                            "but auth failed. Please provide auth configuration manually."
+                        ) from sse_error
                     logger.error(
                         f"Both transport methods failed. Streamable HTTP: {streamable_error}, SSE: {sse_error}"
                     )

mcp_use/connectors/websocket.py

@@ -10,6 +10,7 @@ import json
 import uuid
 from typing import Any
 
+import httpx
 from mcp.types import Tool
 from websockets import ClientConnection
 
@@ -28,21 +29,29 @@ class WebSocketConnector(BaseConnector):
     def __init__(
         self,
         url: str,
-        auth_token: str | None = None,
         headers: dict[str, str] | None = None,
+        auth: str | dict[str, Any] | httpx.Auth | None = None,
     ):
         """Initialize a new WebSocket connector.
 
         Args:
             url: The WebSocket URL to connect to.
-            auth_token: Optional authentication token.
             headers: Optional additional headers.
+            auth: Authentication method - can be:
+                - A string token: Use Bearer token authentication
+                - A dict: Not supported for WebSocket (will log warning)
+                - An httpx.Auth object: Not supported for WebSocket (will log warning)
         """
         self.url = url
-        self.auth_token = auth_token
         self.headers = headers or {}
-        if auth_token:
-            self.headers["Authorization"] = f"Bearer {auth_token}"
+
+        # Handle authentication - WebSocket only supports bearer tokens
+        # An auth field it's not needed
+        if auth is not None:
+            if isinstance(auth, str):
+                self.headers["Authorization"] = f"Bearer {auth}"
+            else:
+                logger.warning("WebSocket connector only supports bearer token authentication")
 
         self.ws: ClientConnection | None = None
         self._connection_manager: ConnectionManager | None = None

mcp_use/exceptions.py

@@ -0,0 +1,31 @@
+"""MCP-use exceptions."""
+
+
+class MCPError(Exception):
+    """Base exception for MCP-use."""
+
+    pass
+
+
+class OAuthDiscoveryError(MCPError):
+    """OAuth discovery auth metadata error"""
+
+    pass
+
+
+class OAuthAuthenticationError(MCPError):
+    """OAuth authentication-related errors"""
+
+    pass
+
+
+class ConnectionError(MCPError):
+    """Connection-related errors."""
+
+    pass
+
+
+class ConfigurationError(MCPError):
+    """Configuration-related errors."""
+
+    pass

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