mcp-proxy-oauth-dcr 0.1.0__py3-none-any.whl

mcp_proxy/http/authenticated_client.py

@@ -0,0 +1,388 @@
+"""Authenticated HTTP client that integrates OAuth authentication with HTTP requests.
+
+This module provides an HTTP client wrapper that:
+- Automatically attaches Bearer tokens to all HTTP MCP requests
+- Handles 401 Unauthorized responses with automatic token refresh
+- Implements fallback to DCR when token refresh fails
+- Provides transparent authentication for all HTTP operations
+- Comprehensive error logging and diagnostics
+"""
+
+import logging
+from typing import Optional, AsyncIterator
+
+from ..auth.manager import AuthenticationManagerImpl
+from ..models import HttpMcpRequest, HttpMcpResponse
+from ..exceptions import (
+    AuthenticationError,
+    TokenError,
+    HttpError,
+)
+from .client import HttpClientImpl, SSEEvent
+from ..logging_config import get_logger
+
+logger = get_logger(__name__)
+
+
+class AuthenticatedHttpClient:
+    """HTTP client with integrated OAuth authentication.
+    
+    This client wraps HttpClientImpl and AuthenticationManagerImpl to provide
+    transparent authentication for all HTTP requests. It automatically:
+    - Attaches valid Bearer tokens to requests
+    - Refreshes tokens when they expire
+    - Handles 401 responses by refreshing tokens and retrying
+    - Falls back to DCR if token refresh fails
+    
+    Attributes:
+        http_client: Underlying HTTP client for making requests
+        auth_manager: Authentication manager for token management
+    """
+
+    def __init__(
+        self,
+        http_client: HttpClientImpl,
+        auth_manager: AuthenticationManagerImpl,
+    ):
+        """Initialize the authenticated HTTP client.
+        
+        Args:
+            http_client: HTTP client instance for making requests
+            auth_manager: Authentication manager for token management
+        """
+        self._http_client = http_client
+        self._auth_manager = auth_manager
+        self._max_auth_retries = 2  # Maximum retries for 401 errors
+        
+        logger.debug("AuthenticatedHttpClient initialized")
+
+    @property
+    def http_client(self) -> HttpClientImpl:
+        """Get the underlying HTTP client."""
+        return self._http_client
+
+    @property
+    def auth_manager(self) -> AuthenticationManagerImpl:
+        """Get the authentication manager."""
+        return self._auth_manager
+
+    @property
+    def session_id(self) -> Optional[str]:
+        """Get the current session ID from HTTP client."""
+        return self._http_client.session_id
+
+    @session_id.setter
+    def session_id(self, value: Optional[str]) -> None:
+        """Set the session ID on HTTP client."""
+        self._http_client.session_id = value
+
+    async def initialize(self) -> None:
+        """Initialize the authenticated client.
+        
+        This method:
+        1. Initializes the authentication manager (performs DCR if needed)
+        2. Obtains an initial access token
+        3. Sets the token on the HTTP client
+        
+        Raises:
+            AuthenticationError: If initialization fails
+        """
+        logger.info("Initializing authenticated HTTP client")
+        
+        # Initialize authentication manager (performs DCR and gets token)
+        await self._auth_manager.initialize()
+        
+        # Set initial token on HTTP client
+        await self._update_http_client_token()
+        
+        logger.info("Authenticated HTTP client initialized successfully")
+
+    async def _update_http_client_token(self) -> None:
+        """Update the HTTP client with the current access token.
+        
+        Raises:
+            TokenError: If unable to obtain a valid token
+        """
+        token = await self._auth_manager.get_access_token()
+        self._http_client.set_auth_token(token)
+        logger.debug("Updated HTTP client with fresh access token")
+
+    async def send_request(self, request: HttpMcpRequest) -> HttpMcpResponse:
+        """Send an authenticated HTTP request with automatic token refresh.
+        
+        This method:
+        1. Ensures a valid token is attached to the request
+        2. Sends the request via the HTTP client
+        3. Handles 401 responses by refreshing the token and retrying
+        4. Falls back to DCR if token refresh fails
+        
+        Args:
+            request: HTTP MCP request to send
+            
+        Returns:
+            HTTP MCP response
+            
+        Raises:
+            AuthenticationError: If authentication fails after all retries
+            HttpError: If HTTP error occurs (non-401)
+            ConnectionError: If connection fails
+        """
+        # Ensure we have a valid token before sending
+        await self._update_http_client_token()
+        
+        # Try sending the request with automatic 401 handling
+        for attempt in range(self._max_auth_retries + 1):
+            try:
+                response = await self._http_client.send_request(request)
+                
+                # Check for 401 Unauthorized
+                if response.is_auth_error():
+                    logger.warning(
+                        "Received 401 Unauthorized response",
+                        attempt=attempt + 1,
+                        max_attempts=self._max_auth_retries + 1,
+                        url=request.url,
+                        method=request.method,
+                        session_id=self._http_client.session_id,
+                        response_body=response.body[:500] if response.body else None,  # Log response for debugging
+                        response_headers=response.headers,
+                    )
+                    
+                    # If this is not the last attempt, try to refresh and retry
+                    if attempt < self._max_auth_retries:
+                        await self._handle_auth_error(attempt)
+                        continue
+                    else:
+                        # Last attempt failed, raise error with diagnostics
+                        auth_diagnostics = self._auth_manager.get_diagnostic_info()
+                        logger.error(
+                            "Authentication failed after all retry attempts",
+                            attempts=attempt + 1,
+                            url=request.url,
+                            method=request.method,
+                            auth_diagnostics=auth_diagnostics,
+                        )
+                        raise AuthenticationError(
+                            "Authentication failed after token refresh attempts",
+                            details={
+                                "status": 401,
+                                "body": response.body,
+                                "diagnostics": auth_diagnostics,
+                            }
+                        )
+                
+                # Success - return response
+                logger.debug(
+                    "Request completed successfully",
+                    url=request.url,
+                    method=request.method,
+                    status=response.status,
+                )
+                return response
+                
+            except HttpError as e:
+                # If it's a 401 error from the HTTP client, handle it
+                if e.status_code == 401 and attempt < self._max_auth_retries:
+                    logger.warning(
+                        "HTTP client raised 401 error",
+                        attempt=attempt + 1,
+                        max_attempts=self._max_auth_retries + 1,
+                        url=request.url,
+                        method=request.method,
+                    )
+                    await self._handle_auth_error(attempt)
+                    continue
+                else:
+                    # Not a 401 or last attempt, re-raise
+                    logger.error(
+                        "HTTP error during request",
+                        status_code=e.status_code,
+                        error=str(e),
+                        url=request.url,
+                        method=request.method,
+                    )
+                    raise
+        
+        # Should not reach here, but just in case
+        raise AuthenticationError("Authentication failed with unknown error")
+
+    async def _handle_auth_error(self, attempt: int) -> None:
+        """Handle 401 authentication error by refreshing token.
+        
+        This method implements the authentication recovery strategy:
+        1. First attempt: Try to refresh the token
+        2. Second attempt: Fall back to DCR if refresh fails
+        
+        Args:
+            attempt: Current attempt number (0-indexed)
+            
+        Raises:
+            AuthenticationError: If unable to recover from auth error
+        """
+        try:
+            if attempt == 0:
+                # First 401: Try token refresh
+                logger.info(
+                    "Attempting token refresh after 401 response",
+                    attempt=attempt + 1,
+                )
+                new_token = await self._auth_manager.refresh_token()
+                self._http_client.set_auth_token(new_token)
+                logger.info("Token refreshed successfully, retrying request")
+            else:
+                # Second 401: Fall back to DCR
+                logger.warning(
+                    "Token refresh failed to resolve 401, falling back to DCR",
+                    attempt=attempt + 1,
+                )
+                await self._auth_manager.perform_dcr()
+                new_token = await self._auth_manager.refresh_token()
+                self._http_client.set_auth_token(new_token)
+                logger.info("DCR and token refresh completed, retrying request")
+                
+        except TokenError as e:
+            logger.error(
+                "Failed to refresh token",
+                error=str(e),
+                error_type=type(e).__name__,
+                attempt=attempt + 1,
+                auth_diagnostics=self._auth_manager.get_diagnostic_info(),
+            )
+            # If first attempt, we'll try DCR on next iteration
+            # If second attempt, we'll fail on next iteration
+            if attempt >= 1:
+                raise AuthenticationError(
+                    "Failed to recover from authentication error",
+                    details={
+                        "original_error": str(e),
+                        "diagnostics": self._auth_manager.get_diagnostic_info(),
+                    }
+                )
+        except Exception as e:
+            logger.error(
+                "Unexpected error during auth recovery",
+                error=str(e),
+                error_type=type(e).__name__,
+                attempt=attempt + 1,
+                exc_info=True,
+            )
+            raise AuthenticationError(
+                "Unexpected error during authentication recovery",
+                details={
+                    "error": str(e),
+                    "error_type": type(e).__name__,
+                }
+            )
+
+    async def post(
+        self,
+        path: str,
+        body: str,
+        headers: Optional[dict] = None,
+    ) -> HttpMcpResponse:
+        """Send an authenticated POST request.
+        
+        Convenience method that wraps send_request for POST requests.
+        
+        Args:
+            path: URL path (will be appended to base URL)
+            body: Request body (JSON string)
+            headers: Additional headers
+            
+        Returns:
+            HTTP MCP response
+        """
+        request = HttpMcpRequest(
+            method="POST",
+            url=path,
+            headers=headers or {},
+            body=body,
+            session_id=self._http_client.session_id,
+        )
+        return await self.send_request(request)
+
+    async def get(
+        self,
+        path: str,
+        headers: Optional[dict] = None,
+    ) -> HttpMcpResponse:
+        """Send an authenticated GET request.
+        
+        Convenience method that wraps send_request for GET requests.
+        
+        Args:
+            path: URL path (will be appended to base URL)
+            headers: Additional headers
+            
+        Returns:
+            HTTP MCP response
+        """
+        request = HttpMcpRequest(
+            method="GET",
+            url=path,
+            headers=headers or {},
+            session_id=self._http_client.session_id,
+        )
+        return await self.send_request(request)
+
+    async def open_sse_stream(
+        self,
+        path: str = "",
+        headers: Optional[dict] = None,
+    ) -> AsyncIterator[SSEEvent]:
+        """Open an authenticated Server-Sent Events stream.
+        
+        This method ensures a valid token is attached before opening the stream.
+        Note: SSE streams are long-lived, so token expiration during the stream
+        is not automatically handled. The caller should handle reconnection if needed.
+        
+        Args:
+            path: URL path for SSE endpoint (default: base URL)
+            headers: Additional headers
+            
+        Yields:
+            SSEEvent objects as they arrive from the server
+            
+        Raises:
+            AuthenticationError: If authentication fails
+            ConnectionError: If connection fails
+            StreamError: If stream encounters an error
+        """
+        # Ensure we have a valid token before opening stream
+        await self._update_http_client_token()
+        
+        # Open the SSE stream (delegates to HTTP client)
+        async for event in self._http_client.open_sse_stream(path, headers):
+            yield event
+
+    async def close(self) -> None:
+        """Close the authenticated client and release resources.
+        
+        This closes both the HTTP client and authentication manager.
+        """
+        logger.info("Closing authenticated HTTP client")
+        
+        # Close authentication manager first (stops background tasks)
+        await self._auth_manager.close()
+        
+        # Close HTTP client (releases connections)
+        await self._http_client.close()
+        
+        logger.info("Authenticated HTTP client closed")
+
+    async def __aenter__(self) -> "AuthenticatedHttpClient":
+        """Async context manager entry."""
+        await self.initialize()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Async context manager exit."""
+        await self.close()
+
+    def __repr__(self) -> str:
+        """String representation of the authenticated client."""
+        return (
+            f"AuthenticatedHttpClient("
+            f"authenticated={self._auth_manager.is_token_valid()}, "
+            f"session_id={self.session_id})"
+        )