ai-parrot 0.17.2__cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl

parrot/mcp/transports/sse.py

@@ -0,0 +1,330 @@
+import asyncio
+import json
+import logging
+import uuid
+import contextlib
+from typing import Dict, Any, Optional
+from aiohttp import web
+import aiohttp
+
+from parrot.mcp.config import MCPServerConfig
+from parrot.mcp.transports.base import MCPServerBase
+from parrot.mcp.oauth import OAuthRoutesMixin
+from parrot.mcp.client import MCPClientConfig, MCPConnectionError
+from parrot.mcp.transports.http import HttpMCPSession
+from aiohttp_sse_client import client as sse_client
+
+class SseMCPServer(OAuthRoutesMixin, MCPServerBase):
+    """MCP server using SSE transport compatible with ChatGPT and OpenAI MCP clients."""
+
+    def __init__(self, config: MCPServerConfig, parent_app: Optional[web.Application] = None):
+        super().__init__(config)
+        self.app = parent_app or web.Application()
+        self.base_path = config.base_path or "/mcp"
+        self.events_path = config.events_path or f"{self.base_path.rstrip('/')}/events"
+        self._init_oauth_support()
+        self.runner = None
+        self.site = None
+        self.sessions: Dict[str, asyncio.Queue] = {}
+        self._external_setup = parent_app is not None
+        self.app.router.add_post(self.base_path, self._handle_http_request)
+        # Alias GET on base path to the SSE stream for clients that connect at /mcp
+        self.app.router.add_get(self.base_path, self._handle_sse, allow_head=True)
+        self.app.router.add_get(self.events_path, self._handle_sse, allow_head=True)
+        self.app.router.add_get("/", self._handle_info, allow_head=True)
+        if self.oauth_server:
+            self.oauth_server.register_routes(self.app)
+
+    async def start(self):
+        """Start the SSE MCP server."""
+        if self._external_setup:
+            self.logger.info("SSE MCP server using existing aiohttp application")
+            return
+
+        self.logger.info(
+            f"Starting SSE MCP server on {self.config.host}:{self.config.port}"
+        )
+
+        self.runner = web.AppRunner(self.app)
+        await self.runner.setup()
+
+        self.site = web.TCPSite(
+            self.runner,
+            self.config.host,
+            self.config.port
+        )
+        await self.site.start()
+
+        self.logger.info(f"SSE MCP server started at http://{self.config.host}:{self.config.port}")
+        self.logger.info(
+            "MCP endpoints: "
+            f"events at http://{self.config.host}:{self.config.port}{self.events_path}, "
+            f"requests at http://{self.config.host}:{self.config.port}{self.base_path}"
+        )
+
+    async def stop(self):
+        """Stop the SSE server."""
+        if not self._external_setup:
+            if self.site:
+                await self.site.stop()
+            if self.runner:
+                await self.runner.cleanup()
+
+        # Clear any pending sessions
+        for session_id, queue in list(self.sessions.items()):
+            with contextlib.suppress(Exception):
+                queue.put_nowait(None)
+            self.sessions.pop(session_id, None)
+
+        self.logger.info("SSE MCP server stopped")
+
+    async def _handle_info(self, request: web.Request) -> web.Response:
+        auth_response = self._authenticate_request(request)
+        if auth_response:
+            return auth_response
+
+        info = {
+            "name": self.config.name,
+            "version": self.config.version,
+            "description": self.config.description,
+            "transport": "sse",
+            "endpoint": self.events_path,
+            "tools": list(self.tools.keys()),
+            "tool_count": len(self.tools)
+        }
+        return web.json_response(info)
+
+    def _get_session_id(self, request: web.Request) -> str:
+        return request.headers.get("X-Session-Id") or request.query.get("session") or str(uuid.uuid4())
+
+    async def _handle_sse(self, request: web.Request) -> web.StreamResponse:
+        auth_response = self._authenticate_request(request)
+        if auth_response:
+            return auth_response
+
+        session_id = self._get_session_id(request)
+        queue: asyncio.Queue = asyncio.Queue()
+        self.sessions[session_id] = queue
+        self.logger.info(f"SSE client connected: {session_id}")
+
+        response = web.StreamResponse(
+            status=200,
+            headers={
+                "Content-Type": "text/event-stream",
+                "Cache-Control": "no-cache",
+                "Connection": "keep-alive",
+                "X-Session-Id": session_id,
+            },
+        )
+
+        await response.prepare(request)
+        await response.write(self._format_sse_event({"type": "connected", "session": session_id}, event="connection"))
+
+        try:
+            while True:
+                try:
+                    message = await asyncio.wait_for(queue.get(), timeout=15)
+                    if message is None:
+                        break
+                    await response.write(self._format_sse_event(message))
+                except asyncio.TimeoutError:
+                    await response.write(b": keep-alive\n\n")
+        except asyncio.CancelledError:
+            self.logger.info(f"SSE client disconnected: {session_id}")
+        finally:
+            self.sessions.pop(session_id, None)
+            with contextlib.suppress(Exception):
+                await response.write_eof()
+
+        return response
+
+    def _format_sse_event(self, payload: Any, event: str = "message") -> bytes:
+        data = json.dumps(payload)
+        return f"event: {event}\ndata: {data}\n\n".encode("utf-8")
+
+    async def _push_to_session(self, session_id: Optional[str], message: Dict[str, Any]):
+        if not session_id:
+            return
+        if queue := self.sessions.get(session_id):
+            try:
+                queue.put_nowait(message)
+            except Exception as e:
+                self.logger.warning(f"Failed to enqueue SSE message for {session_id}: {e}")
+
+    async def _handle_http_request(self, request: web.Request) -> web.Response:
+        try:
+            auth_response = self._authenticate_request(request)
+            if auth_response:
+                return auth_response
+
+            data = await request.json()
+            method = data.get("method")
+            params = data.get("params", {})
+            request_id = data.get("id")
+            session_id = request.headers.get("X-Session-Id") or request.query.get("session")
+
+            self.logger.info(f"SSE HTTP request: {method} (session {session_id or 'none'})")
+
+            try:
+                if method == "initialize":
+                    result = await self.handle_initialize(params)
+                elif method == "tools/list":
+                    result = await self.handle_tools_list(params)
+                elif method == "notifications/initialized":
+                    # This is a notification, no response needed
+                    self.logger.info("Client initialization complete")
+                    return web.Response(status=204)  # No Content
+                elif method == "tools/call":
+                    result = await self.handle_tools_call(params)
+                else:
+                    raise RuntimeError(
+                        f"Unknown method: {method}"
+                    )
+
+                response = {
+                    "jsonrpc": "2.0",
+                    "id": request_id,
+                    "result": result
+                }
+
+            except Exception as e:
+                self.logger.error(f"Error handling {method}: {e}")
+                response = {
+                    "jsonrpc": "2.0",
+                    "id": request_id,
+                    "error": {
+                        "code": -32603,
+                        "message": str(e)
+                    }
+                }
+
+            await self._push_to_session(session_id, response)
+            return web.json_response(response)
+
+        except Exception as e:
+            self.logger.error(f"SSE HTTP request error: {e}")
+            return web.json_response(
+                {
+                    "jsonrpc": "2.0",
+                    "id": None,
+                    "error": {
+                        "code": -32700,
+                        "message": "Parse error",
+                    }
+                },
+            )
+
+
+class SseMCPSession(HttpMCPSession):
+    """MCP session using SSE (Server-Sent Events) for transport."""
+
+    def __init__(self, config: MCPClientConfig, logger):
+        super().__init__(config, logger)
+        self._sse_task = None
+        self._session_id = str(uuid.uuid4())
+        self._sse_ready = asyncio.Event()
+
+    async def connect(self):
+        """Connect to MCP server via SSE + HTTP."""
+        try:
+            # 1. Setup headers (auth etc) - reusing logic from HttpMCPSession
+            # We need to manually trigger auth setup if we want headers for SSE connection
+            if self.config.auth_type:
+                from parrot.mcp.client import MCPAuthHandler
+                self._auth_handler = MCPAuthHandler(
+                    self.config.auth_type,
+                    self.config.auth_config
+                )
+                auth_headers = await self._auth_handler.get_auth_headers()
+                self._base_headers.update(auth_headers)
+
+            self._base_headers.update(self.config.headers)
+            
+            # Add Session ID to headers for correlation
+            self._base_headers["X-Session-Id"] = self._session_id
+
+            # 2. Start SSE Listener
+            self._sse_task = asyncio.create_task(self._listen_sse())
+            
+            # Wait for connection to be ready
+            try:
+                await asyncio.wait_for(self._sse_ready.wait(), timeout=self.config.timeout)
+            except asyncio.TimeoutError:
+                self.logger.warning("Timeout waiting for SSE connection ready, proceeding anyway")
+
+            # 3. Create HTTP session for POST requests
+            timeout = aiohttp.ClientTimeout(total=self.config.timeout)
+            self._session = aiohttp.ClientSession(
+                timeout=timeout,
+                headers=self._base_headers
+            )
+
+            # 4. Initialize MCP
+            await self._initialize_session()
+            self._initialized = True
+            self.logger.info(f"SSE connection established to {self.config.name} (session {self._session_id})")
+
+        except Exception as e:
+            await self.disconnect()
+            raise MCPConnectionError(f"SSE connection failed: {e}") from e
+
+    async def _listen_sse(self):
+        """Listen for SSE events."""
+        events_url = self.config.events_path
+        if not events_url:
+            events_url = self.config.url
+        
+        self.logger.debug(f"Connecting to SSE endpoint: {events_url}")
+
+        try:
+            async with sse_client.EventSource(
+                events_url,
+                headers=self._base_headers
+            ) as event_source:
+                self._sse_ready.set()
+                
+                async for event in event_source:
+                    try:
+                        if event.type == 'error':
+                            self.logger.error(f"SSE Error: {event.data}")
+                            continue
+                            
+                        if event.type == 'connection':
+                            self.logger.debug(f"SSE Connection event: {event.data}")
+                            continue
+
+                        if event.type == 'message':
+                            await self._handle_sse_message(event.data)
+                            
+                    except Exception as e:
+                        self.logger.error(f"Error processing SSE event: {e}")
+
+        except asyncio.CancelledError:
+            self.logger.debug("SSE listener cancelled")
+        except Exception as e:
+            self.logger.error(f"SSE connection error: {e}")
+
+    async def _handle_sse_message(self, data: str):
+        """Handle incoming JSON-RPC message from SSE."""
+        try:
+            message = json.loads(data)
+            
+            if "method" in message and "id" not in message:
+                self.logger.info(f"Received notification: {message['method']}")
+            elif "method" in message and "id" in message:
+                 self.logger.info(f"Received server request: {message['method']}")
+            
+        except json.JSONDecodeError:
+            self.logger.error("Failed to decode SSE message")
+
+    async def disconnect(self):
+        """Disconnect SSE session."""
+        if self._sse_task:
+            self._sse_task.cancel()
+            try:
+                await self._sse_task
+            except asyncio.CancelledError:
+                pass
+            self._sse_task = None
+            
+        await super().disconnect()

parrot/mcp/transports/stdio.py

@@ -0,0 +1,309 @@
+import sys
+import os
+import asyncio
+import json
+import logging
+from typing import Dict, Any, Optional
+
+from parrot.mcp.config import MCPServerConfig
+from parrot.mcp.transports.base import MCPServerBase
+from parrot.mcp.client import MCPClientConfig, MCPConnectionError
+
+class StdioMCPServer(MCPServerBase):
+    """MCP server using stdio transport."""
+
+    def __init__(self, config: MCPServerConfig):
+        super().__init__(config)
+        self._request_id = 0
+        self._running = False
+
+    async def start(self):
+        """Start the stdio MCP server."""
+        self.logger.info(f"Starting stdio MCP server with {len(self.tools)} tools...")
+        self._running = True
+
+        while self._running:
+            try:
+                # Read line from stdin
+                line = sys.stdin.readline()
+                if not line:
+                    break
+
+                line = line.strip()
+                if not line:
+                    continue
+
+                # Parse JSON-RPC request
+                try:
+                    request = json.loads(line)
+                    response = await self._handle_request(request)
+
+                    if response:
+                        print(json.dumps(response), flush=True)
+
+                except json.JSONDecodeError as e:
+                    self.logger.warning(f"Invalid JSON received: {e}")
+                    continue
+
+            except KeyboardInterrupt:
+                break
+            except Exception as e:
+                self.logger.error(f"Error in main loop: {e}")
+                continue
+
+        self.logger.info("Stdio MCP server stopped")
+
+    async def stop(self):
+        """Stop the stdio server."""
+        self._running = False
+
+    async def _handle_request(self, request: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+        """Handle a JSON-RPC request."""
+        method = request.get("method")
+        params = request.get("params", {})
+        request_id = request.get("id")
+
+        try:
+            if method == "initialize":
+                result = await self.handle_initialize(params)
+            elif method == "tools/list":
+                result = await self.handle_tools_list(params)
+            elif method == "tools/call":
+                result = await self.handle_tools_call(params)
+            elif method == "notifications/initialized":
+                # This is a notification, no response needed
+                self.logger.info("Client initialization complete")
+                return None
+            else:
+                raise RuntimeError(f"Unknown method: {method}")
+
+            # Return success response
+            return {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "result": result
+            }
+
+        except Exception as e:
+            self.logger.error(f"Error handling {method}: {e}")
+            return {
+                "jsonrpc": "2.0",
+                "id": request_id,
+                "error": {
+                    "code": -32603,
+                    "message": str(e)
+                }
+            }
+
+
+class StdioMCPSession:
+    """MCP session for stdio transport."""
+
+    def __init__(self, config: MCPClientConfig, logger):
+        self.config = config
+        self.logger = logger
+        self._request_id = 0
+        self._process = None
+        self._stdin = None
+        self._stdout = None
+        self._stderr = None
+        self._initialized = False
+
+    async def connect(self):
+        """Connect to MCP server via stdio."""
+        if self._process:
+            await self.disconnect()
+
+        try:
+            await self._start_process()
+            await self._initialize_session()
+            self._initialized = True
+            self.logger.info(f"Stdio connection established to {self.config.name}")
+
+        except Exception as e:
+            await self.disconnect()
+            raise MCPConnectionError(f"Stdio connection failed: {e}") from e
+
+    async def _start_process(self):
+        """Start the MCP server process."""
+        if not self.config.command:
+            raise ValueError("Command required for stdio transport")
+
+        args = self.config.args or []
+        env = dict(os.environ)
+        if self.config.env:
+            env.update(self.config.env)
+
+        self._process = await asyncio.create_subprocess_exec(
+            self.config.command,
+            *args,
+            stdin=asyncio.subprocess.PIPE,
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+            env=env
+        )
+
+        self._stdin = self._process.stdin
+        self._stdout = self._process.stdout
+        self._stderr = self._process.stderr
+
+        await asyncio.sleep(self.config.startup_delay)
+
+        if self._process.returncode is not None:
+            stderr_output = ""
+            if self._stderr:
+                try:
+                    stderr_data = await asyncio.wait_for(self._stderr.read(1024), timeout=2.0)
+                    stderr_output = stderr_data.decode('utf-8', errors='replace')
+                except asyncio.TimeoutError:
+                    stderr_output = "No error output available"
+
+            raise RuntimeError(f"Process failed to start: {stderr_output}")
+
+    async def _initialize_session(self):
+        """Initialize the MCP session."""
+        try:
+            await self._send_request("initialize", {
+                "protocolVersion": "2024-11-05",
+                "capabilities": {"tools": {}},
+                "clientInfo": {"name": "ai-parrot-mcp-client", "version": "1.0.0"}
+            })
+            await self._send_notification("notifications/initialized")
+        except Exception as e:
+            raise MCPConnectionError(f"Session initialization failed: {e}") from e
+
+    async def _send_request(self, method: str, params: dict = None) -> dict:
+        if not self._process or self._process.returncode is not None:
+            raise MCPConnectionError("Process is not running")
+
+        request_id = self._get_next_id()
+        request = {"jsonrpc": "2.0", "id": request_id, "method": method}
+        if params:
+            request["params"] = params
+
+        loop = asyncio.get_running_loop()
+        deadline = loop.time() + self.config.timeout
+
+        try:
+            line = (json.dumps(request) + "\n").encode("utf-8")
+            self.logger.debug(f"Stdio sending: {line.decode().strip()}")
+            self._stdin.write(line)
+            await self._stdin.drain()
+
+            response = None
+            while True:
+                timeout = max(0.1, deadline - loop.time())
+                response_line = await asyncio.wait_for(self._stdout.readline(), timeout=timeout)
+                if not response_line:
+                    raise MCPConnectionError("Empty response - connection closed")
+
+                response_str = response_line.decode("utf-8", errors="replace").strip()
+                if not response_str:
+                    continue
+                self.logger.debug(f"Stdio received: {response_str}")
+
+                # Skip non-JSON garbage
+                try:
+                    candidate = json.loads(response_str)
+                except json.JSONDecodeError:
+                    self.logger.debug(f"Ignoring non-JSON stdout: {response_str!r}")
+                    continue
+
+                # Only accept responses with our request id; ignore notifications/others
+                if candidate.get("id") != request_id:
+                    # could be a notification or another message; ignore
+                    continue
+
+                response = candidate
+                break
+
+            if "error" in response:
+                raise MCPConnectionError(f"Server error: {response['error']}")
+
+            return response.get("result", {})
+
+        except asyncio.TimeoutError:
+            raise MCPConnectionError(f"Request timeout after {self.config.timeout} seconds")
+        except Exception as e:
+            if isinstance(e, MCPConnectionError):
+                raise
+            raise MCPConnectionError(f"Request failed: {e}") from e
+
+    async def _send_notification(self, method: str, params: dict = None):
+        """Send JSON-RPC notification."""
+        notification = {"jsonrpc": "2.0", "method": method}
+        if params:
+            notification["params"] = params
+
+        notification_line = json.dumps(notification) + "\n"
+        self.logger.debug(f"Stdio notification: {notification_line.strip()}")
+
+        self._stdin.write(notification_line.encode('utf-8'))
+        await self._stdin.drain()
+
+    def _get_next_id(self):
+        self._request_id += 1
+        return self._request_id
+
+    async def list_tools(self):
+        """List available tools."""
+        if not self._initialized:
+            raise MCPConnectionError("Session not initialized")
+
+        result = await self._send_request("tools/list")
+        tools = result.get("tools", [])
+
+        tool_objects = []
+        for tool_dict in tools:
+            tool_obj = type('MCPTool', (), tool_dict)()
+            tool_objects.append(tool_obj)
+
+        return tool_objects
+
+    async def call_tool(self, tool_name: str, arguments: dict):
+        """Call a tool."""
+        if not self._initialized:
+            raise MCPConnectionError("Session not initialized")
+
+        result = await self._send_request("tools/call", {
+            "name": tool_name,
+            "arguments": arguments
+        })
+
+        content_items = []
+        if "content" in result:
+            for item in result["content"]:
+                content_obj = type('ContentItem', (), item)()
+                content_items.append(content_obj)
+
+        result_obj = type('ToolCallResult', (), {"content": content_items})()
+        return result_obj
+
+    async def disconnect(self):
+        """Disconnect stdio session."""
+        self._initialized = False
+
+        if self._process:
+            try:
+                if self._stdin and not self._stdin.is_closing():
+                    self._stdin.close()
+                    await self._stdin.wait_closed()
+
+                if self._process.returncode is None:
+                    try:
+                        await asyncio.wait_for(
+                            self._process.wait(),
+                            timeout=self.config.kill_timeout
+                        )
+                    except asyncio.TimeoutError:
+                        self.logger.warning("Process didn't terminate, force killing")
+                        self._process.kill()
+                        await self._process.wait()
+
+            except Exception as e:
+                self.logger.debug(f"Error during disconnect: {e}")
+            finally:
+                self._process = None
+                self._stdin = None
+                self._stdout = None
+                self._stderr = None