hanzoai_agent/sdk/python/hanzo_agents/mcp_stdio_bridge.py

import asyncio
import json
import os
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Dict, Optional

import uvicorn
from fastapi import FastAPI, HTTPException

from .logger import get_logger

logger = get_logger(__name__)


@dataclass
class PendingRequest:
    """Represents a pending request waiting for response"""

    future: asyncio.Future
    timestamp: float


class StdioMCPBridge:
    """
    Bridge that converts stdio-based MCP servers to HTTP endpoints.

    This bridge starts a stdio MCP server process and provides HTTP endpoints
    that translate HTTP requests to JSON-RPC over stdio and back.
    """

    def __init__(self, server_config: dict, port: int, dev_mode: bool = False):
        self.server_config = server_config
        self.port = port
        self.dev_mode = dev_mode

        # Process management
        self.process: Optional[asyncio.subprocess.Process] = None
        self.stdin_writer: Optional[asyncio.StreamWriter] = None
        self.stdout_reader: Optional[asyncio.StreamReader] = None
        self.stderr_reader: Optional[asyncio.StreamReader] = None

        # Request correlation
        self.pending_requests: Dict[str, PendingRequest] = {}
        self.request_timeout = 30.0  # seconds

        # Server state
        self.initialized = False
        self.running = False
        self.app: Optional[FastAPI] = None
        self.server_task: Optional[asyncio.Task] = None
        self.stdio_reader_task: Optional[asyncio.Task] = None

        # Request ID counter for JSON-RPC
        self._request_id_counter = 0

    def _get_next_request_id(self) -> int:
        """Get next request ID for JSON-RPC"""
        self._request_id_counter += 1
        return self._request_id_counter

    async def start(self) -> bool:
        """Start the stdio MCP server and HTTP bridge"""
        try:
            if self.dev_mode:
                logger.debug(
                    f"Starting stdio MCP bridge for {self.server_config.get('alias', 'unknown')} "
                    f"on port {self.port}"
                )

            # Start the stdio MCP server process
            if not await self._start_stdio_process():
                return False

            # Start stdio response reader BEFORE initializing MCP session
            self.running = True
            self.stdio_reader_task = asyncio.create_task(self._read_stdio_responses())

            # Give the reader task a moment to start
            await asyncio.sleep(0.1)

            # Initialize MCP session
            if not await self._initialize_mcp_session():
                await self.stop()
                return False

            # Setup HTTP server
            self._setup_http_server()

            # Start HTTP server
            if self.app is None:
                raise RuntimeError("HTTP server not properly initialized")

            config = uvicorn.Config(
                app=self.app,
                host="localhost",
                port=self.port,
                log_level="error" if not self.dev_mode else "info",
                access_log=self.dev_mode,
            )

            server = uvicorn.Server(config)
            self.server_task = asyncio.create_task(server.serve())

            if self.dev_mode:
                logger.debug(
                    f"Stdio MCP bridge started successfully on port {self.port}"
                )

            return True

        except Exception as e:
            logger.error(f"Failed to start stdio MCP bridge: {e}")
            await self.stop()
            return False

    async def stop(self) -> None:
        """Stop the bridge and cleanup resources"""
        if self.dev_mode:
            logger.debug("Stopping stdio MCP bridge...")

        self.running = False

        # Cancel pending requests
        for request_id, pending in self.pending_requests.items():
            if not pending.future.done():
                pending.future.set_exception(Exception("Bridge shutting down"))
        self.pending_requests.clear()

        # Stop HTTP server
        if self.server_task and not self.server_task.done():
            self.server_task.cancel()
            try:
                await self.server_task
            except asyncio.CancelledError:
                pass

        # Stop stdio reader
        if self.stdio_reader_task and not self.stdio_reader_task.done():
            self.stdio_reader_task.cancel()
            try:
                await self.stdio_reader_task
            except asyncio.CancelledError:
                pass

        # Close stdio streams
        if self.stdin_writer:
            self.stdin_writer.close()
            await self.stdin_writer.wait_closed()

        # Terminate process
        if self.process:
            try:
                self.process.terminate()
                try:
                    await asyncio.wait_for(
                        asyncio.create_task(self._wait_for_process()), timeout=5.0
                    )
                except asyncio.TimeoutError:
                    self.process.kill()
                    await asyncio.create_task(self._wait_for_process())
            except Exception as e:
                logger.error(f"Error stopping process: {e}")

        if self.dev_mode:
            logger.debug("Stdio MCP bridge stopped")

    async def _wait_for_process(self):
        """Wait for process to terminate"""
        if self.process:
            await self.process.wait()

    async def health_check(self) -> bool:
        """Check if bridge and stdio process are healthy"""
        if not self.running or not self.process:
            return False

        # Check if process is still running
        if self.process.returncode is not None:
            return False

        # Try a simple tools/list request to verify communication
        try:
            await asyncio.wait_for(
                self._send_stdio_request("tools/list", {}), timeout=5.0
            )
            return True
        except Exception:
            return False

    async def _start_stdio_process(self) -> bool:
        """Start the stdio MCP server process"""
        try:
            run_command = self.server_config.get("run", "")
            if not run_command:
                raise ValueError("No run command specified in server config")

            working_dir = self.server_config.get("working_dir", ".")
            env = os.environ.copy()
            env.update(self.server_config.get("environment", {}))

            if self.dev_mode:
                logger.debug(f"Starting process: {run_command}")
                logger.debug(f"Working directory: {working_dir}")

            # Start process
            self.process = await asyncio.create_subprocess_shell(
                run_command,
                stdin=asyncio.subprocess.PIPE,
                stdout=asyncio.subprocess.PIPE,
                stderr=asyncio.subprocess.PIPE,
                cwd=working_dir,
                env=env,
            )

            if (
                self.process.stdin is None
                or self.process.stdout is None
                or self.process.stderr is None
            ):
                raise RuntimeError("Failed to create stdio pipes for process")

            self.stdin_writer = self.process.stdin
            self.stdout_reader = self.process.stdout
            self.stderr_reader = self.process.stderr

            # Give process time to start
            await asyncio.sleep(1.0)

            # Check if process started successfully
            if self.process.returncode is not None:
                stderr_output = ""
                if self.stderr_reader:
                    try:
                        stderr_data = await asyncio.wait_for(
                            self.stderr_reader.read(1024), timeout=1.0
                        )
                        stderr_output = stderr_data.decode("utf-8", errors="ignore")
                    except asyncio.TimeoutError:
                        pass

                raise RuntimeError(
                    f"Process failed to start. Exit code: {self.process.returncode}. Stderr: {stderr_output}"
                )

            return True

        except Exception as e:
            logger.error(f"Failed to start stdio process: {e}")
            return False

    async def _initialize_mcp_session(self) -> bool:
        """Initialize MCP session with handshake"""
        try:
            if self.dev_mode:
                logger.debug("Initializing MCP session...")

            # Send initialize request
            init_params = {
                "protocolVersion": "2024-11-05",
                "capabilities": {"roots": {"listChanged": True}},
                "clientInfo": {"name": "hanzo_agents-stdio-bridge", "version": "1.0.0"},
            }

            response = await self._send_stdio_request("initialize", init_params)

            if "error" in response:
                raise RuntimeError(f"Initialize failed: {response['error']}")

            # Send initialized notification (no response expected)
            await self._send_stdio_notification("notifications/initialized", {})

            self.initialized = True

            if self.dev_mode:
                logger.debug("MCP session initialized successfully")

            return True

        except Exception as e:
            logger.error(f"Failed to initialize MCP session: {e}")
            return False

    def _setup_http_server(self) -> None:
        """Setup FastAPI HTTP server with MCP endpoints"""

        @asynccontextmanager
        async def lifespan(app: FastAPI):
            # Startup
            yield
            # Shutdown
            await self.stop()

        self.app = FastAPI(
            title="MCP Stdio Bridge",
            description="HTTP bridge for stdio-based MCP servers",
            lifespan=lifespan,
        )

        @self.app.get("/health")
        async def health_endpoint():
            """Health check endpoint"""
            is_healthy = await self.health_check()
            if is_healthy:
                return {"status": "healthy", "bridge": "running", "process": "running"}
            else:
                raise HTTPException(
                    status_code=503, detail="Bridge or process not healthy"
                )

        @self.app.post("/mcp/tools/list")
        async def list_tools_endpoint():
            """List available tools from stdio MCP server"""
            try:
                response = await self._handle_list_tools({})
                return response
            except Exception as e:
                logger.error(f"Error listing tools: {e}")
                raise HTTPException(status_code=500, detail=str(e))

        @self.app.post("/mcp/tools/call")
        async def call_tool_endpoint(request: dict):
            """Call a specific tool on stdio MCP server"""
            try:
                response = await self._handle_call_tool(request)
                return response
            except Exception as e:
                logger.error(f"Error calling tool: {e}")
                raise HTTPException(status_code=500, detail=str(e))

        # Also support the standard MCP v1 endpoint format
        @self.app.post("/mcp/v1")
        async def mcp_v1_endpoint(request: dict):
            """Standard MCP v1 JSON-RPC endpoint"""
            try:
                method = request.get("method", "")
                params = request.get("params", {})

                if method == "tools/list":
                    result = await self._handle_list_tools(params)
                    return {
                        "jsonrpc": "2.0",
                        "id": request.get("id", 1),
                        "result": result,
                    }
                elif method == "tools/call":
                    result = await self._handle_call_tool(params)
                    return {
                        "jsonrpc": "2.0",
                        "id": request.get("id", 1),
                        "result": result,
                    }
                else:
                    raise HTTPException(
                        status_code=400, detail=f"Unsupported method: {method}"
                    )

            except Exception as e:
                logger.error(f"Error in MCP v1 endpoint: {e}")
                return {
                    "jsonrpc": "2.0",
                    "id": request.get("id", 1),
                    "error": {"code": -32603, "message": str(e)},
                }

    async def _handle_list_tools(self, request: dict) -> dict:
        """Handle tools/list request"""
        try:
            response = await self._send_stdio_request("tools/list", {})

            if "error" in response:
                raise RuntimeError(f"Tools list failed: {response['error']}")

            result = response.get("result", {})
            tools = result.get("tools", [])

            return {"tools": tools}

        except Exception as e:
            logger.error(f"Failed to list tools: {e}")
            raise

    async def _handle_call_tool(self, request: dict) -> dict:
        """Handle tools/call request"""
        try:
            tool_name = request.get("name")
            arguments = request.get("arguments", {})

            if not tool_name:
                raise ValueError("Tool name is required")

            params = {"name": tool_name, "arguments": arguments}

            response = await self._send_stdio_request("tools/call", params)

            if "error" in response:
                raise RuntimeError(f"Tool call failed: {response['error']}")

            return response.get("result", {})

        except Exception as e:
            logger.error(f"Failed to call tool: {e}")
            raise

    async def _send_stdio_request(self, method: str, params: dict) -> dict:
        """Send JSON-RPC request to stdio process and wait for response"""
        if not self.stdin_writer:
            raise RuntimeError("Stdio process not initialized")

        request_id = self._get_next_request_id()

        request = {
            "jsonrpc": "2.0",
            "id": request_id,
            "method": method,
            "params": params,
        }

        # Create future for response
        future = asyncio.Future()
        self.pending_requests[str(request_id)] = PendingRequest(
            future=future, timestamp=asyncio.get_event_loop().time()
        )

        try:
            # Send request
            request_json = json.dumps(request) + "\n"
            self.stdin_writer.write(request_json.encode("utf-8"))
            await self.stdin_writer.drain()

            if self.dev_mode:
                logger.debug(f"Sent request: {method} (id: {request_id})")

            # Wait for response with timeout
            response = await asyncio.wait_for(future, timeout=self.request_timeout)
            return response

        except asyncio.TimeoutError:
            # Clean up pending request
            self.pending_requests.pop(str(request_id), None)
            raise RuntimeError(f"Request timeout for {method}")
        except Exception as e:
            # Clean up pending request
            self.pending_requests.pop(str(request_id), None)
            raise RuntimeError(f"Request failed for {method}: {e}")

    async def _send_stdio_notification(self, method: str, params: dict) -> None:
        """Send JSON-RPC notification to stdio process (no response expected)"""
        if not self.stdin_writer:
            raise RuntimeError("Stdio process not initialized")

        notification = {"jsonrpc": "2.0", "method": method, "params": params}

        notification_json = json.dumps(notification) + "\n"
        self.stdin_writer.write(notification_json.encode("utf-8"))
        await self.stdin_writer.drain()

        if self.dev_mode:
            logger.debug(f"Sent notification: {method}")

    async def _read_stdio_responses(self) -> None:
        """Read responses from stdio process and correlate with pending requests"""
        if not self.stdout_reader:
            return

        try:
            while self.running:
                try:
                    # Read line from stdout
                    line = await asyncio.wait_for(
                        self.stdout_reader.readline(), timeout=1.0
                    )

                    if not line:
                        # EOF reached
                        break

                    line_str = line.decode("utf-8").strip()
                    if not line_str:
                        continue

                    # Parse JSON response
                    try:
                        response = json.loads(line_str)
                    except json.JSONDecodeError:
                        if self.dev_mode:
                            logger.warning(
                                f"Failed to parse JSON response: {line_str[:100]}..."
                            )
                        continue

                    # Handle response
                    await self._handle_stdio_response(response)

                except asyncio.TimeoutError:
                    # Check for expired requests
                    await self._cleanup_expired_requests()
                    continue
                except Exception as e:
                    if self.running:
                        logger.error(f"Error reading stdio response: {e}")
                    break

        except Exception as e:
            if self.running:
                logger.error(f"Stdio reader task failed: {e}")
        finally:
            # Cancel all pending requests
            for pending in self.pending_requests.values():
                if not pending.future.done():
                    pending.future.set_exception(Exception("Stdio reader stopped"))
            self.pending_requests.clear()

    async def _handle_stdio_response(self, response: dict) -> None:
        """Handle a response from stdio process"""
        response_id = response.get("id")

        if response_id is None:
            # This might be a notification, ignore
            return

        request_id = str(response_id)
        pending = self.pending_requests.pop(request_id, None)

        if pending and not pending.future.done():
            pending.future.set_result(response)

            if self.dev_mode:
                logger.debug(f"Received response for request {request_id}")
        elif self.dev_mode:
            logger.warning(f"Received response for unknown request {request_id}")

    async def _cleanup_expired_requests(self) -> None:
        """Clean up expired pending requests"""
        current_time = asyncio.get_event_loop().time()
        expired_ids = []

        for request_id, pending in self.pending_requests.items():
            if current_time - pending.timestamp > self.request_timeout:
                expired_ids.append(request_id)
                if not pending.future.done():
                    pending.future.set_exception(
                        asyncio.TimeoutError("Request expired")
                    )

        for request_id in expired_ids:
            self.pending_requests.pop(request_id, None)

        if expired_ids and self.dev_mode:
            logger.warning(f"Cleaned up {len(expired_ids)} expired requests")