Nicer error handling

Author: amithkk

src/mcpadapt/auth/__init__.py

@@ -1,40 +1,128 @@
-"""Authentication module for MCPAdapt."""
+"""Authentication module for MCPAdapt.
+
+This module provides OAuth, API Key, and Bearer token authentication support
+for MCP servers.
+
+Example usage with OAuth:
+
+```python
+from mcp.client.auth import OAuthClientProvider
+from mcp.shared.auth import OAuthClientMetadata
+from pydantic import HttpUrl
+
+from mcpadapt.auth import (
+    InMemoryTokenStorage, 
+    LocalBrowserOAuthHandler
+)
+from mcpadapt.core import MCPAdapt
+from mcpadapt.smolagents_adapter import SmolAgentsAdapter
+
+# Create OAuth provider directly
+client_metadata = OAuthClientMetadata(
+    client_name="My App",
+    redirect_uris=[HttpUrl("http://localhost:3030/callback")],
+    grant_types=["authorization_code", "refresh_token"],
+    response_types=["code"],
+    token_endpoint_auth_method="client_secret_post",
+)
+
+oauth_handler = LocalBrowserOAuthHandler(callback_port=3030)
+token_storage = InMemoryTokenStorage()
+
+oauth_provider = OAuthClientProvider(
+    server_url="https://example.com",
+    client_metadata=client_metadata,
+    storage=token_storage,
+    redirect_handler=oauth_handler.handle_redirect,
+    callback_handler=oauth_handler.handle_callback,
+)
+
+# Use with MCPAdapt
+with MCPAdapt(
+    serverparams={"url": "https://example.com/mcp", "transport": "streamable-http"},
+    adapter=SmolAgentsAdapter(),
+    auth_provider=oauth_provider,
+) as tools:
+    print(f"Connected with {len(tools)} tools")
+```
+
+Example usage with API Key:
+
+```python
+from mcpadapt.auth import ApiKeyAuthProvider
+from mcpadapt.core import MCPAdapt
+from mcpadapt.smolagents_adapter import SmolAgentsAdapter
+
+# Create API Key provider
+api_key_provider = ApiKeyAuthProvider(
+    header_name="X-API-Key",
+    header_value="your-api-key-here"
+)
+
+with MCPAdapt(
+    serverparams={"url": "https://example.com/mcp", "transport": "streamable-http"},
+    adapter=SmolAgentsAdapter(),
+    auth_provider=api_key_provider,
+) as tools:
+    print(f"Connected with {len(tools)} tools")
+```
+
+For custom implementations, extend BaseOAuthHandler:
+
+```python
+from mcpadapt.auth import BaseOAuthHandler
+
+class CustomOAuthHandler(BaseOAuthHandler):
+    async def handle_redirect(self, authorization_url: str) -> None:
+        # Custom redirect logic (e.g., print URL for headless environments)
+        print(f"Please open: {authorization_url}")
+    
+    async def handle_callback(self) -> tuple[str, str | None]:
+        # Custom callback logic (e.g., manual code input)
+        auth_code = input("Enter authorization code: ")
+        return auth_code, None
+```
+"""
 
-from .handlers import default_callback_handler, default_redirect_handler
 from .oauth import InMemoryTokenStorage
+from .handlers import (
+    BaseOAuthHandler,
+    LocalBrowserOAuthHandler,
+    LocalCallbackServer,
+)
 from .providers import (
     ApiKeyAuthProvider,
     BearerAuthProvider,
-    create_auth_provider,
     get_auth_headers,
 )
-from .models import (
-    ApiKeyConfig,
-    AuthConfig,
-    AuthConfigBase,
-    BearerAuthConfig,
-    CallbackHandler,
-    OAuthConfig,
-    RedirectHandler,
+from .exceptions import (
+    OAuthError,
+    OAuthTimeoutError,
+    OAuthCancellationError,
+    OAuthNetworkError,
+    OAuthConfigurationError,
+    OAuthServerError,
+    OAuthCallbackError,
 )
 
 __all__ = [
-    # Types
-    "AuthConfig",
-    "AuthConfigBase",
-    "OAuthConfig",
-    "ApiKeyConfig",
-    "BearerAuthConfig",
-    "CallbackHandler",
-    "RedirectHandler",
-    # OAuth utilities
-    "InMemoryTokenStorage",
-    # Handlers
-    "default_callback_handler",
-    "default_redirect_handler",
-    # Providers
+    # Handler classes
+    "BaseOAuthHandler",
+    "LocalBrowserOAuthHandler",
+    "LocalCallbackServer",
+    # Provider classes
     "ApiKeyAuthProvider",
     "BearerAuthProvider",
-    "create_auth_provider",
+    # Default implementations
+    "InMemoryTokenStorage",
+    # Provider functions
     "get_auth_headers",
+    # Exception classes
+    "OAuthError",
+    "OAuthTimeoutError",
+    "OAuthCancellationError",
+    "OAuthNetworkError",
+    "OAuthConfigurationError",
+    "OAuthServerError",
+    "OAuthCallbackError",
 ]

src/mcpadapt/auth/exceptions.py

@@ -0,0 +1,85 @@
+"""Custom exceptions for OAuth authentication errors."""
+
+
+class OAuthError(Exception):
+    """Base class for all OAuth authentication errors."""
+    
+    def __init__(self, message: str, error_code: str | None = None, context: dict | None = None):
+        """Initialize OAuth error.
+        
+        Args:
+            message: Human-readable error message
+            error_code: Machine-readable error code (optional)
+            context: Additional context about the error (optional)
+        """
+        super().__init__(message)
+        self.error_code = error_code
+        self.context = context or {}
+
+
+class OAuthTimeoutError(OAuthError):
+    """Raised when OAuth callback doesn't arrive within the specified timeout."""
+    
+    def __init__(self, timeout_seconds: int, context: dict | None = None):
+        message = (
+            f"OAuth authentication timed out after {timeout_seconds} seconds. "
+            f"The user may have closed the browser window or the OAuth server may be unreachable. "
+            f"Try refreshing the browser or check your network connection."
+        )
+        super().__init__(message, "oauth_timeout", context)
+        self.timeout_seconds = timeout_seconds
+
+
+class OAuthCancellationError(OAuthError):
+    """Raised when the user cancels or denies the OAuth authorization."""
+    
+    def __init__(self, error_details: str | None = None, context: dict | None = None):
+        base_message = "OAuth authorization was cancelled or denied by the user."
+        if error_details:
+            message = f"{base_message} Details: {error_details}"
+        else:
+            message = base_message
+        super().__init__(message, "oauth_cancelled", context)
+        self.error_details = error_details
+
+
+class OAuthNetworkError(OAuthError):
+    """Raised when network-related issues prevent OAuth completion."""
+    
+    def __init__(self, original_error: Exception, context: dict | None = None):
+        message = (
+            f"OAuth authentication failed due to network error: {str(original_error)}. "
+            f"Check your internet connection and try again."
+        )
+        super().__init__(message, "oauth_network_error", context)
+        self.original_error = original_error
+
+
+class OAuthConfigurationError(OAuthError):
+    """Raised when OAuth configuration is invalid or incomplete."""
+    
+    def __init__(self, config_issue: str, context: dict | None = None):
+        message = f"OAuth configuration error: {config_issue}"
+        super().__init__(message, "oauth_config_error", context)
+        self.config_issue = config_issue
+
+
+class OAuthServerError(OAuthError):
+    """Raised when the OAuth server returns an error response."""
+    
+    def __init__(self, server_error: str, error_description: str | None = None, context: dict | None = None):
+        message = f"OAuth server error: {server_error}"
+        if error_description:
+            message += f" - {error_description}"
+        super().__init__(message, "oauth_server_error", context)
+        self.server_error = server_error
+        self.error_description = error_description
+
+
+class OAuthCallbackError(OAuthError):
+    """Raised when there's an issue with the OAuth callback handling."""
+    
+    def __init__(self, callback_issue: str, context: dict | None = None):
+        message = f"OAuth callback error: {callback_issue}"
+        super().__init__(message, "oauth_callback_error", context)
+        self.callback_issue = callback_issue