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

mcp_proxy/proxy.py

@@ -0,0 +1,586 @@
+"""Main MCP Proxy orchestration.
+
+This module provides the main McpProxy class that orchestrates all components:
+- Stdio interface for communication with Kiro
+- Protocol translator for stdio <-> HTTP conversion
+- Authenticated HTTP client for backend communication
+- Configuration management
+- Graceful startup and shutdown
+- Signal handling and cleanup
+
+Requirements satisfied:
+- 3.1: Perform OAuth DCR with OAuth provider on startup
+- 2.1: Establish connections to HTTP MCP server
+"""
+
+import asyncio
+import logging
+import signal
+from typing import Optional
+
+from .auth.manager import AuthenticationManagerImpl
+from .config.manager import ConfigurationManagerImpl
+from .http.authenticated_client import AuthenticatedHttpClient
+from .http.client import HttpClientImpl
+from .models import JsonRpcMessage, ProxyConfig
+from .translator.translator import ProtocolTranslatorImpl
+from .stdio.interface import StdioInterface
+from .exceptions import McpProxyError, ConfigurationError
+from .logging_config import get_logger, configure_logging
+
+logger = get_logger(__name__)
+
+
+class McpProxy:
+    """Main MCP Proxy orchestrator.
+    
+    This class brings together all components to create a complete proxy that:
+    - Presents a stdio MCP interface to Kiro
+    - Translates between stdio and HTTP MCP protocols
+    - Manages OAuth DCR authentication
+    - Handles HTTP communication with backend MCP server
+    - Provides graceful startup and shutdown
+    
+    The proxy operates as follows:
+    1. Load configuration from environment variables
+    2. Initialize authentication manager (performs DCR)
+    3. Create HTTP client with authentication
+    4. Set up protocol translator
+    5. Start stdio interface
+    6. Process messages bidirectionally
+    7. Handle shutdown signals gracefully
+    
+    Attributes:
+        config: Proxy configuration
+        is_running: Whether the proxy is currently running
+    """
+
+    def __init__(self, config: Optional[ProxyConfig] = None):
+        """Initialize the MCP Proxy.
+        
+        Args:
+            config: Optional proxy configuration. If not provided,
+                   configuration will be loaded from environment variables.
+        """
+        self._config = config
+        self._running = False
+        self._shutdown_event = asyncio.Event()
+        
+        # Components (initialized during startup)
+        self._config_manager: Optional[ConfigurationManagerImpl] = None
+        self._auth_manager: Optional[AuthenticationManagerImpl] = None
+        self._http_client: Optional[HttpClientImpl] = None
+        self._authenticated_client: Optional[AuthenticatedHttpClient] = None
+        self._translator: Optional[ProtocolTranslatorImpl] = None
+        self._stdio_interface: Optional[StdioInterface] = None
+        
+        logger.info("McpProxy instance created")
+
+    @property
+    def config(self) -> Optional[ProxyConfig]:
+        """Get the current configuration."""
+        return self._config
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the proxy is currently running."""
+        return self._running
+
+    async def start(self) -> None:
+        """Start the MCP Proxy.
+        
+        This method performs the complete startup sequence:
+        1. Load configuration (if not provided)
+        2. Set up logging
+        3. Initialize authentication manager (performs DCR)
+        4. Create and initialize HTTP client
+        5. Set up protocol translator
+        6. Start stdio interface
+        7. Install signal handlers
+        
+        Raises:
+            ConfigurationError: If configuration is invalid
+            AuthenticationError: If OAuth DCR fails
+            RuntimeError: If proxy is already running
+        """
+        if self._running:
+            raise RuntimeError("McpProxy is already running")
+        
+        logger.info("Starting MCP Proxy")
+        
+        try:
+            # Step 1: Load configuration
+            await self._load_configuration()
+            
+            # Step 2: Set up logging
+            self._setup_logging()
+            
+            # Step 3: Initialize authentication manager
+            await self._initialize_authentication()
+            
+            # Step 4: Create and initialize HTTP client
+            await self._initialize_http_client()
+            
+            # Step 5: Set up protocol translator
+            self._initialize_translator()
+            
+            # Step 6: Start stdio interface
+            await self._start_stdio_interface()
+            
+            # Step 7: Install signal handlers
+            self._install_signal_handlers()
+            
+            self._running = True
+            self._shutdown_event.clear()
+            
+            logger.info("MCP Proxy started successfully")
+            
+        except Exception as e:
+            logger.error(f"Failed to start MCP Proxy: {e}", exc_info=True)
+            # Clean up any partially initialized components
+            await self._cleanup()
+            raise
+
+    async def stop(self) -> None:
+        """Stop the MCP Proxy gracefully.
+        
+        This method performs graceful shutdown:
+        1. Signal shutdown to all components
+        2. Stop stdio interface (drains pending messages)
+        3. Close authenticated HTTP client
+        4. Close authentication manager
+        5. Clean up resources
+        """
+        if not self._running:
+            logger.warning("McpProxy is not running")
+            return
+        
+        logger.info("Stopping MCP Proxy")
+        self._running = False
+        self._shutdown_event.set()
+        
+        try:
+            # Stop stdio interface first (drains pending messages)
+            if self._stdio_interface:
+                await self._stdio_interface.stop()
+            
+            # Close authenticated HTTP client
+            if self._authenticated_client:
+                await self._authenticated_client.close()
+            
+            # Close authentication manager
+            if self._auth_manager:
+                await self._auth_manager.close()
+            
+            # Clean up remaining resources
+            await self._cleanup()
+            
+            logger.info("MCP Proxy stopped successfully")
+            
+        except Exception as e:
+            logger.error(f"Error during shutdown: {e}", exc_info=True)
+            raise
+
+    async def run(self) -> None:
+        """Run the proxy until shutdown signal is received.
+        
+        This method starts the proxy and waits for a shutdown signal.
+        It's the main entry point for running the proxy as a standalone process.
+        
+        Usage:
+            proxy = McpProxy()
+            await proxy.run()
+        """
+        await self.start()
+        
+        try:
+            # Wait for shutdown signal
+            logger.info("MCP Proxy running, waiting for shutdown signal")
+            await self._shutdown_event.wait()
+        except KeyboardInterrupt:
+            logger.info("Received keyboard interrupt")
+        finally:
+            await self.stop()
+
+    async def wait_for_shutdown(self) -> None:
+        """Wait for shutdown signal.
+        
+        This method blocks until the proxy is signaled to shut down.
+        Useful for keeping the main process alive while the proxy runs.
+        """
+        await self._shutdown_event.wait()
+
+    # ========================================================================
+    # Private Methods - Initialization
+    # ========================================================================
+
+    async def _load_configuration(self) -> None:
+        """Load configuration from environment variables if not provided."""
+        if self._config is None:
+            logger.info("Loading configuration from environment variables")
+            self._config_manager = ConfigurationManagerImpl()
+            self._config = await self._config_manager.load()
+            logger.info("Configuration loaded successfully")
+        else:
+            logger.info("Using provided configuration")
+        
+        # Log configuration (excluding sensitive data)
+        logger.info(
+            "Configuration: mcp_server_url=%s, oauth_provider_url=%s, "
+            "client_name=%s, scopes=%s, connection_timeout=%s, "
+            "retry_attempts=%s, log_level=%s",
+            self._config.mcp_server_url,
+            self._config.oauth_provider_url,
+            self._config.client_name,
+            self._config.scopes,
+            self._config.connection_timeout,
+            self._config.retry_attempts,
+            self._config.log_level,
+        )
+
+    def _setup_logging(self) -> None:
+        """Set up logging based on configuration."""
+        if self._config:
+            configure_logging(self._config.log_level)
+            logger.info(f"Logging configured with level: {self._config.log_level}")
+
+    async def _initialize_authentication(self) -> None:
+        """Initialize authentication manager and perform DCR.
+        
+        This satisfies Requirement 3.1: Perform OAuth DCR on startup.
+        """
+        if not self._config:
+            raise ConfigurationError("Configuration not loaded")
+        
+        logger.info("Initializing authentication manager")
+        self._auth_manager = AuthenticationManagerImpl(self._config)
+        
+        # Initialize performs DCR and obtains initial token
+        await self._auth_manager.initialize()
+        
+        logger.info("Authentication manager initialized successfully")
+
+    async def _initialize_http_client(self) -> None:
+        """Initialize HTTP client with authentication.
+        
+        This satisfies Requirement 2.1: Establish connections to HTTP MCP server.
+        """
+        if not self._config or not self._auth_manager:
+            raise RuntimeError("Configuration or authentication not initialized")
+        
+        logger.info("Initializing HTTP client")
+        
+        # Create retry configuration
+        from .http.client import RetryConfig
+        retry_config = RetryConfig(
+            max_retries=self._config.retry_attempts,
+            max_delay=self._config.max_backoff_seconds,
+        )
+        
+        # Create base HTTP client
+        self._http_client = HttpClientImpl(
+            base_url=str(self._config.mcp_server_url),
+            timeout=self._config.connection_timeout,
+            retry_config=retry_config,
+        )
+        
+        # Wrap with authentication
+        self._authenticated_client = AuthenticatedHttpClient(
+            http_client=self._http_client,
+            auth_manager=self._auth_manager,
+        )
+        
+        # Initialize authenticated client (ensures token is ready)
+        await self._authenticated_client.initialize()
+        
+        logger.info("HTTP client initialized successfully")
+
+    def _initialize_translator(self) -> None:
+        """Initialize protocol translator."""
+        if not self._config:
+            raise RuntimeError("Configuration not initialized")
+        
+        logger.info("Initializing protocol translator")
+        
+        self._translator = ProtocolTranslatorImpl(
+            base_url=str(self._config.mcp_server_url)
+        )
+        
+        logger.info("Protocol translator initialized successfully")
+
+    async def _start_stdio_interface(self) -> None:
+        """Start stdio interface with message handler."""
+        logger.info("Starting stdio interface")
+        
+        # Create stdio interface with message handler
+        self._stdio_interface = StdioInterface(
+            message_handler=self._handle_stdio_message
+        )
+        
+        # Start the interface
+        await self._stdio_interface.start()
+        
+        # Wait for interface to be fully started
+        started = await self._stdio_interface.wait_until_started(timeout=5.0)
+        if not started:
+            raise RuntimeError("Stdio interface failed to start within timeout")
+        
+        logger.info("Stdio interface started successfully")
+
+    # ========================================================================
+    # Private Methods - Message Handling
+    # ========================================================================
+
+    async def _handle_stdio_message(
+        self,
+        message: JsonRpcMessage
+    ) -> Optional[JsonRpcMessage]:
+        """Handle incoming stdio message from Kiro.
+        
+        This is the main message processing pipeline:
+        1. Translate stdio message to HTTP request
+        2. Send HTTP request via authenticated client
+        3. Translate HTTP response back to stdio message
+        4. Return stdio response
+        
+        Args:
+            message: JSON-RPC message from Kiro
+            
+        Returns:
+            JSON-RPC response message, or None for notifications
+        """
+        if not self._translator or not self._authenticated_client:
+            logger.error(
+                "Components not initialized",
+                has_translator=self._translator is not None,
+                has_authenticated_client=self._authenticated_client is not None,
+            )
+            return self._create_error_response(
+                message.id,
+                -32603,
+                "Internal error: proxy not fully initialized"
+            )
+        
+        try:
+            # Handle notifications (no response expected)
+            if message.is_notification():
+                logger.debug(
+                    "Received notification",
+                    method=message.method,
+                )
+                # For now, we don't forward notifications to backend
+                # This could be extended to support server-to-client notifications
+                return None
+            
+            # Translate stdio request to HTTP
+            logger.debug(
+                "Translating stdio request to HTTP",
+                method=message.method,
+                id=message.id,
+                has_params=message.params is not None,
+            )
+            http_request = await self._translator.translate_stdio_to_http(message)
+            
+            # Send HTTP request via authenticated client
+            logger.debug(
+                "Sending HTTP request",
+                http_method=http_request.method,
+                url=http_request.url,
+                session_id=http_request.session_id,
+            )
+            http_response = await self._authenticated_client.send_request(http_request)
+            
+            # Translate HTTP response back to stdio
+            logger.debug(
+                "Translating HTTP response to stdio",
+                status=http_response.status,
+                content_type=http_response.content_type,
+                body_length=len(http_response.body) if http_response.body else 0,
+            )
+            stdio_response = await self._translator.translate_http_to_stdio(
+                http_response,
+                stdio_request_id=message.id
+            )
+            
+            # Mark correlation as completed
+            if message.id is not None:
+                try:
+                    self._translator.mark_correlation_completed(message.id)
+                except Exception as e:
+                    logger.warning(
+                        "Failed to mark correlation completed",
+                        error=str(e),
+                        message_id=message.id,
+                    )
+            
+            logger.info(
+                "Message processed successfully",
+                method=message.method,
+                request_id=message.id,
+                response_has_result=stdio_response.result is not None,
+                response_has_error=stdio_response.error is not None,
+            )
+            return stdio_response
+            
+        except McpProxyError as e:
+            logger.error(
+                "Proxy error handling message",
+                error=str(e),
+                error_type=type(e).__name__,
+                message_id=message.id,
+                method=message.method,
+                has_details=hasattr(e, 'details') and e.details is not None,
+                exc_info=True,
+            )
+            
+            # Mark correlation as failed
+            if message.id is not None:
+                try:
+                    self._translator.mark_correlation_failed(message.id)
+                except Exception:
+                    pass
+            
+            # Return error response
+            from .exceptions import to_jsonrpc_error
+            error_obj = to_jsonrpc_error(e)
+            return self._create_error_response(
+                message.id,
+                error_obj["code"],
+                error_obj["message"],
+                error_obj.get("data")
+            )
+            
+        except Exception as e:
+            logger.error(
+                "Unexpected error handling message",
+                error=str(e),
+                error_type=type(e).__name__,
+                message_id=message.id,
+                method=message.method,
+                exc_info=True,
+            )
+            
+            # Mark correlation as failed
+            if message.id is not None:
+                try:
+                    self._translator.mark_correlation_failed(message.id)
+                except Exception:
+                    pass
+            
+            # Return generic error response
+            return self._create_error_response(
+                message.id,
+                -32603,
+                f"Internal error: {str(e)}"
+            )
+
+    def _create_error_response(
+        self,
+        request_id: Optional[int | str],
+        code: int,
+        message: str,
+        data: Optional[any] = None
+    ) -> JsonRpcMessage:
+        """Create a JSON-RPC error response.
+        
+        Args:
+            request_id: Request ID to respond to
+            code: JSON-RPC error code
+            message: Error message
+            data: Optional additional error data
+            
+        Returns:
+            JSON-RPC error response
+        """
+        from .models import JsonRpcError
+        
+        return JsonRpcMessage(
+            jsonrpc="2.0",
+            id=request_id,
+            error=JsonRpcError(
+                code=code,
+                message=message,
+                data=data
+            )
+        )
+
+    # ========================================================================
+    # Private Methods - Signal Handling
+    # ========================================================================
+
+    def _install_signal_handlers(self) -> None:
+        """Install signal handlers for graceful shutdown."""
+        loop = asyncio.get_event_loop()
+        
+        for sig in (signal.SIGTERM, signal.SIGINT):
+            try:
+                loop.add_signal_handler(
+                    sig,
+                    lambda s=sig: asyncio.create_task(self._handle_signal(s))
+                )
+                logger.debug(f"Installed signal handler for {sig.name}")
+            except NotImplementedError:
+                # Signal handlers not supported on this platform (e.g., Windows)
+                logger.warning(f"Signal handlers not supported for {sig.name}")
+
+    async def _handle_signal(self, sig: signal.Signals) -> None:
+        """Handle shutdown signals.
+        
+        Args:
+            sig: Signal that was received
+        """
+        logger.info(f"Received signal {sig.name}, initiating shutdown")
+        await self.stop()
+
+    # ========================================================================
+    # Private Methods - Cleanup
+    # ========================================================================
+
+    async def _cleanup(self) -> None:
+        """Clean up resources during shutdown or after failed startup."""
+        logger.debug("Cleaning up resources")
+        
+        # Close components in reverse order of initialization
+        if self._stdio_interface and self._stdio_interface.is_running:
+            try:
+                await self._stdio_interface.stop()
+            except Exception as e:
+                logger.error(f"Error stopping stdio interface: {e}")
+        
+        if self._authenticated_client:
+            try:
+                await self._authenticated_client.close()
+            except Exception as e:
+                logger.error(f"Error closing authenticated client: {e}")
+        
+        if self._auth_manager:
+            try:
+                await self._auth_manager.close()
+            except Exception as e:
+                logger.error(f"Error closing auth manager: {e}")
+        
+        # Clear references
+        self._stdio_interface = None
+        self._authenticated_client = None
+        self._http_client = None
+        self._auth_manager = None
+        self._translator = None
+        
+        logger.debug("Cleanup completed")
+
+    # ========================================================================
+    # Context Manager Support
+    # ========================================================================
+
+    async def __aenter__(self) -> "McpProxy":
+        """Async context manager entry."""
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        """Async context manager exit."""
+        await self.stop()
+
+    def __repr__(self) -> str:
+        """String representation of the proxy."""
+        return f"McpProxy(running={self._running}, config={self._config is not None})"

mcp_proxy/stdio/__init__.py

@@ -0,0 +1,31 @@
+"""Stdio interface module for MCP Proxy.
+
+This module handles stdio communication with Kiro using JSON-RPC protocol.
+"""
+
+from .interface import StdioInterface, DEFAULT_QUEUE_SIZE, SHUTDOWN_TIMEOUT, DRAIN_TIMEOUT
+from .jsonrpc import (
+    JsonRpcParser,
+    JsonRpcSerializer,
+    MessageIdGenerator,
+    MessageCorrelationTracker,
+    create_request,
+    create_response,
+    create_notification,
+    create_error_response,
+)
+
+__all__ = [
+    "StdioInterface",
+    "DEFAULT_QUEUE_SIZE",
+    "SHUTDOWN_TIMEOUT",
+    "DRAIN_TIMEOUT",
+    "JsonRpcParser",
+    "JsonRpcSerializer",
+    "MessageIdGenerator",
+    "MessageCorrelationTracker",
+    "create_request",
+    "create_response",
+    "create_notification",
+    "create_error_response",
+]