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

mcp_proxy/stdio/interface.py

@@ -0,0 +1,580 @@
+"""Stdio interface for communication with Kiro.
+
+This module implements the stdio MCP server interface that Kiro expects,
+handling async stdin/stdout communication, subprocess lifecycle management,
+message queuing, and signal handling.
+
+Requirements satisfied:
+- 1.1: Present valid stdio MCP server interface to Kiro
+- 1.4: Maintain message ordering and correlation between stdio and HTTP protocols
+"""
+
+import asyncio
+import signal
+import sys
+from typing import Optional, Callable, Awaitable, Union
+from asyncio import Queue, StreamReader, StreamWriter
+
+from ..models import JsonRpcMessage
+from ..exceptions import ProtocolError, InvalidJsonRpcError
+from .jsonrpc import JsonRpcParser, JsonRpcSerializer
+from ..logging_config import get_logger
+
+logger = get_logger(__name__)
+
+# Default queue size limit to prevent unbounded memory growth
+DEFAULT_QUEUE_SIZE = 1000
+
+# Timeout for graceful shutdown operations
+SHUTDOWN_TIMEOUT = 5.0
+
+# Timeout for draining pending messages on shutdown
+DRAIN_TIMEOUT = 2.0
+
+
+class StdioInterface:
+    """Stdio interface for JSON-RPC communication with Kiro.
+    
+    This class manages async stdin/stdout communication, providing a clean
+    interface for sending and receiving JSON-RPC messages. It handles:
+    - Async reading from stdin with proper buffering
+    - Async writing to stdout with message queuing
+    - Subprocess lifecycle and signal management
+    - Graceful shutdown and cleanup
+    
+    The interface maintains message ordering by using FIFO queues for both
+    incoming and outgoing messages, ensuring that messages are processed
+    and sent in the order they were received/queued.
+    
+    Attributes:
+        is_running: Whether the interface is currently running
+        pending_incoming: Number of messages waiting to be processed
+        pending_outgoing: Number of messages waiting to be sent
+    """
+
+    def __init__(
+        self,
+        message_handler: Optional[Callable[[JsonRpcMessage], Awaitable[Optional[JsonRpcMessage]]]] = None,
+        queue_size: int = DEFAULT_QUEUE_SIZE
+    ):
+        """Initialize the stdio interface.
+        
+        Args:
+            message_handler: Optional async callback for handling incoming messages.
+                           Should accept a JsonRpcMessage and return a JsonRpcMessage response
+                           or None for notifications.
+            queue_size: Maximum size for message queues (default: 1000)
+        """
+        self._message_handler = message_handler
+        self._running = False
+        self._shutdown_event = asyncio.Event()
+        self._started_event = asyncio.Event()
+        
+        # Message queues for async processing with size limits
+        self._incoming_queue: Queue[JsonRpcMessage] = Queue(maxsize=queue_size)
+        self._outgoing_queue: Queue[JsonRpcMessage] = Queue(maxsize=queue_size)
+        
+        # Stream readers/writers
+        self._stdin_reader: Optional[StreamReader] = None
+        self._stdout_writer: Optional[StreamWriter] = None
+        
+        # Background tasks
+        self._tasks: list[asyncio.Task] = []
+        
+        # Parser and serializer
+        self._parser = JsonRpcParser()
+        self._serializer = JsonRpcSerializer()
+        
+        # Track if we're in the middle of shutdown
+        self._shutting_down = False
+        
+        logger.info("StdioInterface initialized", queue_size=queue_size)
+
+    @property
+    def is_running(self) -> bool:
+        """Check if the interface is currently running."""
+        return self._running
+
+    @property
+    def pending_incoming(self) -> int:
+        """Get the number of messages waiting to be processed."""
+        return self._incoming_queue.qsize()
+
+    @property
+    def pending_outgoing(self) -> int:
+        """Get the number of messages waiting to be sent."""
+        return self._outgoing_queue.qsize()
+
+    def set_message_handler(
+        self,
+        handler: Callable[[JsonRpcMessage], Awaitable[Optional[JsonRpcMessage]]]
+    ) -> None:
+        """Set or update the message handler.
+        
+        Args:
+            handler: Async callback for handling incoming messages.
+                    Should return a JsonRpcMessage for requests or None for notifications.
+        """
+        self._message_handler = handler
+        logger.debug("Message handler updated")
+
+    async def start(self) -> None:
+        """Start the stdio interface and begin processing messages.
+        
+        This method:
+        1. Sets up stdin/stdout streams
+        2. Installs signal handlers for graceful shutdown
+        3. Starts background tasks for reading, writing, and processing
+        
+        Raises:
+            RuntimeError: If interface is already running
+        """
+        if self._running:
+            raise RuntimeError("StdioInterface is already running")
+        
+        logger.info("Starting StdioInterface")
+        self._running = True
+        self._shutting_down = False
+        self._shutdown_event.clear()
+        self._started_event.clear()
+        
+        # Set up stdin/stdout streams
+        await self._setup_streams()
+        
+        # Install signal handlers for graceful shutdown
+        self._install_signal_handlers()
+        
+        # Start background tasks
+        self._tasks = [
+            asyncio.create_task(self._read_loop(), name="stdio-read"),
+            asyncio.create_task(self._write_loop(), name="stdio-write"),
+            asyncio.create_task(self._process_loop(), name="stdio-process"),
+        ]
+        
+        # Signal that we're fully started
+        self._started_event.set()
+        
+        logger.info("StdioInterface started successfully")
+
+    async def wait_until_started(self, timeout: Optional[float] = None) -> bool:
+        """Wait until the interface has fully started.
+        
+        Args:
+            timeout: Maximum time to wait in seconds (None for no timeout)
+            
+        Returns:
+            True if started successfully, False if timeout occurred
+        """
+        try:
+            if timeout is not None:
+                await asyncio.wait_for(self._started_event.wait(), timeout=timeout)
+            else:
+                await self._started_event.wait()
+            return True
+        except asyncio.TimeoutError:
+            return False
+
+    async def stop(self) -> None:
+        """Stop the stdio interface and cleanup resources.
+        
+        This method:
+        1. Signals shutdown to all background tasks
+        2. Drains pending outgoing messages
+        3. Waits for tasks to complete gracefully
+        4. Closes streams and cleans up resources
+        """
+        if not self._running:
+            logger.warning("StdioInterface is not running")
+            return
+        
+        if self._shutting_down:
+            logger.warning("StdioInterface is already shutting down")
+            return
+        
+        logger.info("Stopping StdioInterface")
+        self._shutting_down = True
+        self._running = False
+        self._shutdown_event.set()
+        
+        # Drain pending outgoing messages before stopping
+        await self._drain_outgoing_queue()
+        
+        # Cancel all background tasks
+        for task in self._tasks:
+            if not task.done():
+                task.cancel()
+        
+        # Wait for tasks to complete (with timeout)
+        if self._tasks:
+            try:
+                await asyncio.wait_for(
+                    asyncio.gather(*self._tasks, return_exceptions=True),
+                    timeout=SHUTDOWN_TIMEOUT
+                )
+            except asyncio.TimeoutError:
+                logger.warning("Timeout waiting for tasks to complete")
+        
+        # Close streams
+        await self._close_streams()
+        
+        # Clear queues
+        self._clear_queues()
+        
+        self._tasks.clear()
+        self._shutting_down = False
+        logger.info("StdioInterface stopped")
+
+    async def _drain_outgoing_queue(self) -> None:
+        """Drain and send all pending outgoing messages before shutdown.
+        
+        This ensures that any queued responses are sent before the interface
+        stops, preventing message loss.
+        """
+        if not self._stdout_writer or self._outgoing_queue.empty():
+            return
+        
+        logger.info("Draining outgoing message queue", pending=self._outgoing_queue.qsize())
+        
+        try:
+            async with asyncio.timeout(DRAIN_TIMEOUT):
+                while not self._outgoing_queue.empty():
+                    try:
+                        message = self._outgoing_queue.get_nowait()
+                        line = self._serializer.serialize(message)
+                        self._stdout_writer.write(line.encode('utf-8'))
+                        await self._stdout_writer.drain()
+                        logger.debug("Drained message", message_id=message.id)
+                    except asyncio.QueueEmpty:
+                        break
+        except asyncio.TimeoutError:
+            logger.warning("Timeout draining outgoing queue", remaining=self._outgoing_queue.qsize())
+
+    async def send_message(self, message: JsonRpcMessage) -> None:
+        """Send a JSON-RPC message to stdout.
+        
+        Messages are queued and sent asynchronously in FIFO order,
+        maintaining message ordering as required by the MCP protocol.
+        
+        Args:
+            message: JsonRpcMessage to send
+            
+        Raises:
+            RuntimeError: If interface is not running
+            asyncio.QueueFull: If the outgoing queue is full
+        """
+        if not self._running and not self._shutting_down:
+            raise RuntimeError("StdioInterface is not running")
+        
+        try:
+            # Use put_nowait to avoid blocking if queue is full
+            self._outgoing_queue.put_nowait(message)
+            logger.debug("Queued outgoing message", message_id=message.id)
+        except asyncio.QueueFull:
+            logger.error("Outgoing queue is full, message dropped", message_id=message.id)
+            raise
+
+    async def receive_message(self, timeout: Optional[float] = None) -> Optional[JsonRpcMessage]:
+        """Receive a JSON-RPC message from stdin.
+        
+        Args:
+            timeout: Maximum time to wait in seconds (None for no timeout)
+        
+        Returns:
+            JsonRpcMessage received from stdin, or None if timeout occurred
+            
+        Raises:
+            RuntimeError: If interface is not running
+        """
+        if not self._running:
+            raise RuntimeError("StdioInterface is not running")
+        
+        try:
+            if timeout is not None:
+                message = await asyncio.wait_for(
+                    self._incoming_queue.get(),
+                    timeout=timeout
+                )
+            else:
+                message = await self._incoming_queue.get()
+            logger.debug("Received message from queue", message_id=message.id)
+            return message
+        except asyncio.TimeoutError:
+            return None
+
+    async def wait_for_shutdown(self) -> None:
+        """Wait for shutdown signal.
+        
+        This method blocks until the interface is signaled to shut down,
+        useful for keeping the main process alive.
+        """
+        await self._shutdown_event.wait()
+
+    # ========================================================================
+    # Private Methods - Stream Setup
+    # ========================================================================
+
+    async def _setup_streams(self) -> None:
+        """Set up async stdin/stdout streams."""
+        loop = asyncio.get_event_loop()
+        
+        # Create stdin reader
+        self._stdin_reader = asyncio.StreamReader()
+        protocol = asyncio.StreamReaderProtocol(self._stdin_reader)
+        await loop.connect_read_pipe(lambda: protocol, sys.stdin)
+        
+        # Create stdout writer
+        transport, protocol = await loop.connect_write_pipe(
+            asyncio.streams.FlowControlMixin,
+            sys.stdout
+        )
+        self._stdout_writer = asyncio.StreamWriter(
+            transport, protocol, None, loop
+        )
+        
+        logger.debug("Stdio streams set up successfully")
+
+    async def _close_streams(self) -> None:
+        """Close stdin/stdout streams."""
+        if self._stdout_writer:
+            try:
+                self._stdout_writer.close()
+                await self._stdout_writer.wait_closed()
+            except Exception as e:
+                logger.error(f"Error closing stdout writer: {e}")
+        
+        self._stdin_reader = None
+        self._stdout_writer = None
+        logger.debug("Stdio streams closed")
+
+    # ========================================================================
+    # 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 - Background Tasks
+    # ========================================================================
+
+    async def _read_loop(self) -> None:
+        """Background task for reading messages from stdin.
+        
+        Continuously reads newline-delimited JSON-RPC messages from stdin
+        and queues them for processing. Maintains message ordering by
+        processing messages in the order they are received.
+        """
+        logger.info("Read loop started")
+        
+        try:
+            while self._running and self._stdin_reader:
+                try:
+                    # Read a line from stdin (newline-delimited)
+                    line_bytes = await self._stdin_reader.readline()
+                    
+                    if not line_bytes:
+                        # EOF reached - initiate graceful shutdown
+                        logger.info("EOF reached on stdin, initiating shutdown")
+                        asyncio.create_task(self.stop())
+                        break
+                    
+                    # Decode and parse the message
+                    line = line_bytes.decode('utf-8')
+                    
+                    # Skip empty lines
+                    if not line.strip():
+                        continue
+                    
+                    try:
+                        message = self._parser.parse(line)
+                        try:
+                            self._incoming_queue.put_nowait(message)
+                            logger.debug("Read message from stdin", message_id=message.id, method=message.method)
+                        except asyncio.QueueFull:
+                            logger.error("Incoming queue full, message dropped", message_id=message.id)
+                    except InvalidJsonRpcError as e:
+                        logger.error("Invalid JSON-RPC message", error=str(e))
+                        # Continue reading despite parse errors
+                        
+                except asyncio.CancelledError:
+                    logger.info("Read loop cancelled")
+                    break
+                except Exception as e:
+                    logger.error("Error in read loop", error=str(e), exc_info=True)
+                    # Continue reading despite errors
+                    await asyncio.sleep(0.1)
+        
+        finally:
+            logger.info("Read loop stopped")
+
+    async def _write_loop(self) -> None:
+        """Background task for writing messages to stdout.
+        
+        Continuously processes the outgoing message queue and writes
+        messages to stdout in newline-delimited format. Messages are
+        sent in FIFO order to maintain message ordering.
+        """
+        logger.info("Write loop started")
+        
+        try:
+            while self._running:
+                try:
+                    # Wait for a message with timeout to allow checking shutdown
+                    try:
+                        message = await asyncio.wait_for(
+                            self._outgoing_queue.get(),
+                            timeout=1.0
+                        )
+                    except asyncio.TimeoutError:
+                        continue
+                    
+                    if not self._stdout_writer:
+                        logger.error("Stdout writer not available")
+                        continue
+                    
+                    # Serialize and write the message
+                    line = self._serializer.serialize(message)
+                    self._stdout_writer.write(line.encode('utf-8'))
+                    await self._stdout_writer.drain()
+                    
+                    logger.debug("Wrote message to stdout", message_id=message.id)
+                    
+                except asyncio.CancelledError:
+                    logger.info("Write loop cancelled")
+                    break
+                except Exception as e:
+                    logger.error("Error in write loop", error=str(e), exc_info=True)
+                    await asyncio.sleep(0.1)
+        
+        finally:
+            logger.info("Write loop stopped")
+
+    async def _process_loop(self) -> None:
+        """Background task for processing incoming messages.
+        
+        Continuously processes incoming messages using the message handler
+        and queues responses for sending. Handles both requests (which expect
+        responses) and notifications (which don't).
+        
+        Messages are processed in FIFO order to maintain message ordering
+        as required by the MCP protocol.
+        """
+        logger.info("Process loop started")
+        
+        if not self._message_handler:
+            logger.warning("No message handler set, messages will not be processed")
+            # Still need to consume messages to prevent queue from filling up
+            while self._running:
+                try:
+                    try:
+                        message = await asyncio.wait_for(
+                            self._incoming_queue.get(),
+                            timeout=1.0
+                        )
+                        logger.warning("Message received but no handler set", message_id=message.id)
+                    except asyncio.TimeoutError:
+                        continue
+                except asyncio.CancelledError:
+                    break
+            return
+        
+        try:
+            while self._running:
+                try:
+                    # Wait for a message with timeout to allow checking shutdown
+                    try:
+                        message = await asyncio.wait_for(
+                            self._incoming_queue.get(),
+                            timeout=1.0
+                        )
+                    except asyncio.TimeoutError:
+                        continue
+                    
+                    # Process the message using the handler
+                    try:
+                        response = await self._message_handler(message)
+                        
+                        # Queue the response for sending (if any)
+                        # Notifications (messages without ID) don't get responses
+                        if response is not None:
+                            try:
+                                self._outgoing_queue.put_nowait(response)
+                                logger.debug("Processed message, queued response", 
+                                           request_id=message.id, response_id=response.id)
+                            except asyncio.QueueFull:
+                                logger.error("Outgoing queue full, response dropped",
+                                           message_id=message.id)
+                        elif message.is_notification():
+                            logger.debug("Processed notification", method=message.method)
+                        else:
+                            logger.debug("Handler returned no response", message_id=message.id)
+                            
+                    except Exception as e:
+                        logger.error("Error processing message", 
+                                   message_id=message.id, error=str(e), exc_info=True)
+                        # Create error response for requests (not notifications)
+                        if message.id is not None:
+                            from .jsonrpc import create_error_response
+                            error_response = create_error_response(
+                                request_id=message.id,
+                                code=-32603,
+                                message=f"Internal error: {str(e)}"
+                            )
+                            try:
+                                self._outgoing_queue.put_nowait(error_response)
+                            except asyncio.QueueFull:
+                                logger.error("Outgoing queue full, error response dropped",
+                                           message_id=message.id)
+                    
+                except asyncio.CancelledError:
+                    logger.info("Process loop cancelled")
+                    break
+                except Exception as e:
+                    logger.error("Error in process loop", error=str(e), exc_info=True)
+                    await asyncio.sleep(0.1)
+        
+        finally:
+            logger.info("Process loop stopped")
+
+    # ========================================================================
+    # Private Methods - Utilities
+    # ========================================================================
+
+    def _clear_queues(self) -> None:
+        """Clear all message queues."""
+        # Clear incoming queue
+        while not self._incoming_queue.empty():
+            try:
+                self._incoming_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+        
+        # Clear outgoing queue
+        while not self._outgoing_queue.empty():
+            try:
+                self._outgoing_queue.get_nowait()
+            except asyncio.QueueEmpty:
+                break
+        
+        logger.debug("Message queues cleared")