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

parrot/mcp/server.py

@@ -0,0 +1,225 @@
+"""
+MCP Server Implementation - Expose AI-Parrot Tools via MCP Protocol
+=================================================================
+This creates an MCP server that exposes your existing AbstractTool instances
+as MCP tools that can be consumed by any MCP client.
+"""
+import logging
+import asyncio
+import argparse
+from typing import List, Optional
+from aiohttp import web
+
+# AI-Parrot imports
+from parrot.tools.abstract import AbstractTool
+
+# Config
+from parrot.mcp.config import MCPServerConfig
+
+# Transports
+from parrot.mcp.transports.stdio import StdioMCPServer
+from parrot.mcp.transports.http import HttpMCPServer
+from parrot.mcp.transports.sse import SseMCPServer
+from parrot.mcp.transports.unix import UnixMCPServer
+from parrot.mcp.transports.quic import QuicMCPServer
+
+# Suppress noisy loggers
+logging.getLogger('matplotlib').setLevel(logging.ERROR)
+logging.getLogger('PIL').setLevel(logging.ERROR)
+logging.getLogger('urllib3').setLevel(logging.ERROR)
+logging.getLogger('requests').setLevel(logging.ERROR)
+
+
+class MCPServer:
+    """Main MCP server class that chooses transport."""
+
+    def __init__(self, config: MCPServerConfig, parent_app: Optional[web.Application] = None):
+        self.config = config
+
+        if config.transport == "stdio":
+            self.server = StdioMCPServer(config)
+        elif config.transport == "http":
+            self.server = HttpMCPServer(config, parent_app=parent_app)
+        elif config.transport == "sse":
+            self.server = SseMCPServer(config, parent_app=parent_app)
+        elif config.transport == "unix":
+            self.server = UnixMCPServer(config)
+        elif config.transport == "quic":
+            self.server = QuicMCPServer(config)
+        else:
+            raise ValueError(
+                f"Unsupported transport: {config.transport}"
+            )
+
+    def register_tool(self, tool: AbstractTool):
+        """Register a tool."""
+        self.server.register_tool(tool)
+
+    def register_tools(self, tools: List[AbstractTool]):
+        """Register multiple tools."""
+        self.server.register_tools(tools)
+
+    async def start(self):
+        """Start the server."""
+        await self.server.start()
+
+    async def stop(self):
+        """Start the server."""
+        await self.server.stop()
+
+    async def __aenter__(self):
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.stop()
+
+
+# Convenience functions
+
+def create_stdio_mcp_server(
+    name: str = "ai-parrot-tools",
+    tools: Optional[List[AbstractTool]] = None,
+    **kwargs
+) -> MCPServer:
+    """Create a stdio MCP server."""
+    config = MCPServerConfig(
+        name=name,
+        transport="stdio",
+        **kwargs
+    )
+
+    server = MCPServer(config)
+
+    if tools:
+        server.register_tools(tools)
+
+    return server
+
+
+def create_http_mcp_server(
+    name: str = "ai-parrot-tools",
+    host: str = "localhost",
+    port: int = 8080,
+    tools: Optional[List[AbstractTool]] = None,
+    parent_app: Optional[web.Application] = None,
+    **kwargs
+) -> MCPServer:
+    """Create an HTTP MCP server."""
+    config = MCPServerConfig(
+        name=name,
+        transport="http",
+        host=host,
+        port=port,
+        **kwargs
+    )
+
+    server = HttpMCPServer(
+        config,
+        parent_app=parent_app
+    )
+
+    if tools:
+        server.register_tools(tools)
+
+    return server
+
+def create_sse_mcp_server(
+    name: str = "ai-parrot-tools",
+    host: str = "localhost",
+    port: int = 8080,
+    tools: Optional[List[AbstractTool]] = None,
+    parent_app: Optional[web.Application] = None,
+    **kwargs,
+) -> MCPServer:
+    """Create an SSE MCP server."""
+    config = MCPServerConfig(
+        name=name,
+        transport="sse",
+        host=host,
+        port=port,
+        **kwargs,
+    )
+
+    server = MCPServer(config, parent_app=parent_app)
+    if tools:
+        server.register_tools(tools)
+
+    return server
+
+def create_unix_mcp_server(
+    name: str = "ai-parrot-tools",
+    socket_path: Optional[str] = None,
+    tools: Optional[List[AbstractTool]] = None,
+    **kwargs,
+) -> MCPServer:
+    """Create an Unix MCP server."""
+    config = MCPServerConfig(
+        name=name,
+        transport="unix",
+        socket_path=socket_path,
+        **kwargs,
+    )
+
+    server = MCPServer(config)
+    if tools:
+        server.register_tools(tools)
+
+    return server
+
+# CLI support
+async def main():
+    """Main CLI entry point."""
+    parser = argparse.ArgumentParser(
+        description="AI-Parrot MCP Server"
+    )
+    parser.add_argument("--transport", choices=["stdio", "http", "sse"], default="stdio",
+                        help="Transport type")
+    parser.add_argument("--host", default="localhost",
+                        help="Host for HTTP server")
+    parser.add_argument("--port", type=int, default=8080,
+                        help="Port for HTTP server")
+    parser.add_argument("--name", default="ai-parrot-tools",
+                        help="Server name")
+    parser.add_argument("--log-level", default="INFO",
+                        help="Log level")
+
+    args = parser.parse_args()
+
+    # Create server config
+    config = MCPServerConfig(
+        name=args.name,
+        transport=args.transport,
+        host=args.host,
+        port=args.port,
+        log_level=args.log_level
+    )
+
+    # Create server
+    server = MCPServer(config)
+
+    # Register example tools:
+    # server.register_tool(YourOpenWeatherTool())
+    # server.register_tool(YourDatabaseQueryTool())
+
+    try:
+        if args.transport in {"http", "sse"}:
+            await server.start()
+            print(f"Server running at http://{args.host}:{args.port}")
+            print("Press Ctrl+C to stop")
+
+            # Keep running
+            while True:
+                await asyncio.sleep(1)
+        else:
+            # For stdio, just start and let it handle stdin
+            await server.start()
+
+    except KeyboardInterrupt:
+        print("\nShutting down...")
+    finally:
+        await server.stop()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

parrot/mcp/transports/__init__.py

@@ -0,0 +1,3 @@
+from .websocket import WebSocketMCPServer, WebSocketMCPSession
+
+__all__ = ['WebSocketMCPServer', 'WebSocketMCPSession']

parrot/mcp/transports/base.py

@@ -0,0 +1,279 @@
+from abc import ABC, abstractmethod
+import logging
+import time
+from typing import Dict, List, Optional, Any
+from aiohttp import web
+
+from parrot.tools.abstract import AbstractTool
+from parrot.mcp.config import MCPServerConfig, AuthMethod
+from parrot.mcp.adapter import MCPToolAdapter
+from parrot.mcp.oauth import OAuthAuthorizationServer, APIKeyStore, ExternalOAuthValidator
+
+
+class MCPServerBase(ABC):
+    """Base class for MCP servers."""
+
+    def __init__(self, config: MCPServerConfig):
+        self.config = config
+        self.tools: Dict[str, MCPToolAdapter] = {}
+        self.logger = logging.getLogger(f"MCPServer.{config.name}")
+        log_level = getattr(logging, config.log_level.upper(), logging.WARNING)
+        self.logger.setLevel(log_level)
+
+        # Authentication components
+        self.oauth_server: Optional[OAuthAuthorizationServer] = None
+        self.api_key_store: Optional[APIKeyStore] = None
+        self.external_oauth: Optional[ExternalOAuthValidator] = None
+
+        # Initialize authentication based on method
+        self._init_authentication()
+
+    def _init_authentication(self) -> None:
+        """Initialize authentication based on config.auth_method."""
+        auth_method = self.config.auth_method
+
+        # Backward compatibility: enable_oauth maps to OAUTH2_INTERNAL
+        if self.config.enable_oauth and auth_method == AuthMethod.NONE:
+            auth_method = AuthMethod.OAUTH2_INTERNAL
+            self.config.auth_method = auth_method
+
+        if auth_method == AuthMethod.API_KEY:
+            self.api_key_store = self.config.api_key_store or APIKeyStore()
+            self.logger.info("Authentication: API Key enabled")
+
+        elif auth_method == AuthMethod.OAUTH2_INTERNAL:
+            self.oauth_server = OAuthAuthorizationServer(
+                default_scopes=self.config.oauth_scopes,
+                allow_dynamic_registration=self.config.oauth_allow_dynamic_registration,
+                token_ttl=self.config.oauth_token_ttl,
+                code_ttl=self.config.oauth_code_ttl,
+            )
+            self.logger.info("Authentication: OAuth2 (internal) enabled")
+
+        elif auth_method == AuthMethod.OAUTH2_EXTERNAL:
+            if not self.config.oauth2_introspection_endpoint:
+                raise ValueError(
+                    "oauth2_introspection_endpoint required for OAUTH2_EXTERNAL"
+                )
+            self.external_oauth = ExternalOAuthValidator(
+                introspection_endpoint=self.config.oauth2_introspection_endpoint,
+                client_id=self.config.oauth2_client_id or "",
+                client_secret=self.config.oauth2_client_secret or "",
+                resource_server_url=self.config.oauth2_resource_server_url,
+            )
+            self.logger.info(
+                f"Authentication: OAuth2 (external) enabled - {self.config.oauth2_issuer_url}"
+            )
+
+        elif auth_method == AuthMethod.BEARER:
+            self.logger.info("Authentication: Bearer (navigator-auth) enabled")
+
+        else:
+            self.logger.debug("Authentication: None (open access)")
+
+    def register_tool(self, tool: AbstractTool):
+        """Register an AI-Parrot tool with the MCP server."""
+        tool_name = tool.name
+
+        # Apply filtering
+        if self.config.allowed_tools and tool_name not in self.config.allowed_tools:
+            self.logger.info(f"Skipping tool {tool_name} (not in allowed_tools)")
+            return
+
+        if self.config.blocked_tools and tool_name in self.config.blocked_tools:
+            self.logger.info(f"Skipping tool {tool_name} (in blocked_tools)")
+            return
+
+        adapter = MCPToolAdapter(tool)
+        self.tools[tool_name] = adapter
+        self.logger.info(f"Registered tool: {tool_name}")
+
+    def register_tools(self, tools: List[AbstractTool]):
+        """Register multiple tools."""
+        for tool in tools:
+            self.register_tool(tool)
+
+
+    async def _authenticate_request(self, request: web.Request) -> Optional[web.Response]:
+        """
+        Authenticate request based on configured auth method.
+
+        Returns None if authenticated, or a web.Response with error if not.
+        """
+        auth_method = self.config.auth_method
+
+        if auth_method == AuthMethod.NONE:
+            return None
+
+        elif auth_method == AuthMethod.API_KEY:
+            return await self._authenticate_api_key(request)
+
+        elif auth_method == AuthMethod.OAUTH2_INTERNAL:
+            return self._authenticate_oauth_internal(request)
+
+        elif auth_method == AuthMethod.OAUTH2_EXTERNAL:
+            return await self._authenticate_oauth_external(request)
+
+        elif auth_method == AuthMethod.BEARER:
+            return await self._authenticate_bearer(request)
+
+        return None
+
+    async def _authenticate_api_key(self, request: web.Request) -> Optional[web.Response]:
+        """Validate API key from header."""
+        api_key = request.headers.get(self.config.api_key_header)
+        if not api_key:
+            return self._unauthorized_response(
+                "API key required",
+                f'X-API-Key realm="mcp"'
+            )
+
+        record = self.api_key_store.validate_key(api_key)
+        if not record:
+            return self._unauthorized_response("Invalid or expired API key")
+
+        # Log session start
+        self.api_key_store.log_session_start(api_key, record.user_id, time.time())
+        self.logger.debug(f"API key authenticated for user: {record.user_id}")
+
+        # Store user info in request for downstream use
+        request["mcp_user"] = {"user_id": record.user_id, "scopes": record.scopes}
+        return None
+
+    def _authenticate_oauth_internal(self, request: web.Request) -> Optional[web.Response]:
+        """Validate OAuth access token from internal OAuth server."""
+        if not self.oauth_server:
+            return None
+
+        token = self.oauth_server.bearer_token_from_header(
+            request.headers.get("Authorization")
+        )
+        if not self.oauth_server.is_token_valid(token):
+            return self._unauthorized_response("Valid Bearer token is required")
+
+        return None
+
+    async def _authenticate_oauth_external(self, request: web.Request) -> Optional[web.Response]:
+        """Validate OAuth access token via external introspection."""
+        if not self.external_oauth:
+            return None
+
+        token = self._extract_bearer_token(request.headers.get("Authorization"))
+        if not token:
+            return self._unauthorized_response("Bearer token required")
+
+        token_info = await self.external_oauth.validate_token(token)
+        if not token_info:
+            return self._unauthorized_response("Invalid or expired token")
+
+        # Store token info in request
+        request["mcp_user"] = {
+            "user_id": token_info.get("sub") or token_info.get("client_id"),
+            "scopes": token_info.get("scope", "").split() if token_info.get("scope") else [],
+            "token_info": token_info,
+        }
+        return None
+
+    async def _authenticate_bearer(self, request: web.Request) -> Optional[web.Response]:
+        """Validate bearer token via navigator-auth."""
+        auth = request.app.get('auth')
+        if not auth:
+            self.logger.warning("navigator-auth not configured in app['auth']")
+            # Fall through if auth not configured (development mode)
+            return None
+
+        try:
+            userdata = await auth.get_session(request)
+            if not userdata:
+                return web.json_response(
+                    {"error": "unauthorized", "error_description": "Session required"},
+                    status=401,
+                    headers={"WWW-Authenticate": 'Bearer realm="mcp"'}
+                )
+
+            # Store user info in request
+            request["mcp_user"] = userdata
+            return None
+
+        except Exception as e:
+            self.logger.error(f"navigator-auth error: {e}")
+            return web.json_response(
+                {"error": "unauthorized", "error_description": "Authentication failed"},
+                status=401,
+                headers={"WWW-Authenticate": 'Bearer realm="mcp"'}
+            )
+
+    def _extract_bearer_token(self, auth_header: Optional[str]) -> Optional[str]:
+        """Extract bearer token from Authorization header."""
+        if not auth_header:
+            return None
+        if not auth_header.lower().startswith("bearer "):
+            return None
+        return auth_header.split(" ", 1)[1].strip()
+
+    def _unauthorized_response(
+        self,
+        message: str,
+        www_authenticate: str = 'Bearer realm="mcp"'
+    ) -> web.Response:
+        """Create a 401 unauthorized response."""
+        return web.json_response(
+            {"error": "unauthorized", "error_description": message},
+            status=401,
+            headers={"WWW-Authenticate": www_authenticate}
+        )
+
+
+    async def handle_initialize(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Handle MCP initialize request."""
+        self.logger.info("Initializing MCP server...")
+
+        return {
+            "protocolVersion": "2024-11-05",
+            "capabilities": {
+                "tools": {
+                    "listChanged": False
+                }
+            },
+            "serverInfo": {
+                "name": self.config.name,
+                "version": self.config.version,
+                "description": self.config.description
+            }
+        }
+
+    async def handle_tools_list(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Handle tools/list request."""
+        self.logger.info(f"Listing {len(self.tools)} available tools")
+
+        tools = []
+        tools.extend(
+            adapter.to_mcp_tool_definition() for adapter in self.tools.values()
+        )
+
+        return {"tools": tools}
+
+    async def handle_tools_call(self, params: Dict[str, Any]) -> Dict[str, Any]:
+        """Handle tools/call request."""
+        tool_name = params.get("name")
+        arguments = params.get("arguments", {})
+
+        self.logger.info(f"Calling tool: {tool_name} with args: {arguments}")
+
+        if tool_name not in self.tools:
+            raise RuntimeError(
+                f"Tool not found: {tool_name}"
+            )
+
+        adapter = self.tools[tool_name]
+        return await adapter.execute(arguments)
+
+    @abstractmethod
+    async def start(self):
+        """Start the MCP server."""
+        pass
+
+    @abstractmethod
+    async def stop(self):
+        """Stop the MCP server."""
+        pass

parrot/mcp/transports/grpc_session.py

@@ -0,0 +1,163 @@
+import grpc
+import asyncio
+from typing import Dict, Any, Optional
+from dataclasses import dataclass
+# from .proto import mcp_pb2, mcp_pb2_grpc
+
+
+@dataclass
+class GrpcMCPConfig:
+    """Configuration for gRPC MCP transport."""
+    host: str = "localhost"
+    port: int = 50051
+    use_tls: bool = True
+    cert_path: Optional[str] = None
+    # For protobuf message format instead of JSON-RPC
+    use_protobuf_messages: bool = False
+
+
+class GrpcMCPSession:
+    """MCP session for gRPC transport with optional protobuf messages."""
+    
+    def __init__(self, config: 'MCPServerConfig', logger):
+        self.config = config
+        self.logger = logger
+        self._request_id = 0
+        self._channel: Optional[grpc.aio.Channel] = None
+        self._stub = None
+        self._initialized = False
+        self._response_queue: asyncio.Queue = asyncio.Queue()
+        self._stream_task: Optional[asyncio.Task] = None
+        
+    async def connect(self):
+        """Connect to MCP server via gRPC."""
+        try:
+            grpc_config = self.config.grpc_config or GrpcMCPConfig()
+            
+            # Create channel
+            target = f"{grpc_config.host}:{grpc_config.port}"
+            
+            if grpc_config.use_tls:
+                if grpc_config.cert_path:
+                    with open(grpc_config.cert_path, 'rb') as f:
+                        creds = grpc.ssl_channel_credentials(f.read())
+                else:
+                    creds = grpc.ssl_channel_credentials()
+                self._channel = grpc.aio.secure_channel(target, creds)
+            else:
+                self._channel = grpc.aio.insecure_channel(target)
+            
+            # Create stub
+            self._stub = mcp_pb2_grpc.MCPServiceStub(self._channel)
+            
+            # Start bidirectional stream
+            self._stream = self._stub.BiDirectionalStream()
+            self._stream_task = asyncio.create_task(self._read_responses())
+            
+            # Initialize MCP session
+            await self._initialize_session()
+            self._initialized = True
+            
+            self.logger.info(f"gRPC connection established to {self.config.name}")
+            
+        except Exception as e:
+            await self.disconnect()
+            raise MCPConnectionError(f"gRPC connection failed: {e}") from e
+    
+    async def _read_responses(self):
+        """Background task to read responses from gRPC stream."""
+        try:
+            async for response in self._stream:
+                # Convert protobuf to dict if needed
+                if hasattr(response, 'json_rpc_message'):
+                    msg = json.loads(response.json_rpc_message)
+                else:
+                    msg = self._protobuf_to_dict(response)
+                
+                await self._response_queue.put(msg)
+        except grpc.aio.AioRpcError as e:
+            self.logger.error(f"gRPC stream error: {e}")
+        except asyncio.CancelledError:
+            pass
+    
+    async def send_request(self, method: str, params: Dict = None) -> Dict[str, Any]:
+        """Send JSON-RPC request over gRPC."""
+        self._request_id += 1
+        
+        request = {
+            "jsonrpc": "2.0",
+            "id": self._request_id,
+            "method": method,
+            "params": params or {}
+        }
+        
+        # Option 1: Send JSON-RPC as string in protobuf wrapper
+        grpc_request = mcp_pb2.MCPRequest(
+            json_rpc_message=json.dumps(request)
+        )
+        
+        # Option 2: Use native protobuf messages (more efficient)
+        # grpc_request = self._dict_to_protobuf(request)
+        
+        await self._stream.write(grpc_request)
+        
+        # Wait for response with matching ID
+        while True:
+            response = await asyncio.wait_for(
+                self._response_queue.get(),
+                timeout=self.config.timeout
+            )
+            if response.get("id") == self._request_id:
+                return response
+            # Put back if not matching (notifications, etc.)
+            await self._response_queue.put(response)
+    
+    async def _initialize_session(self):
+        """Initialize MCP session."""
+        response = await self.send_request("initialize", {
+            "protocolVersion": "2024-11-05",
+            "capabilities": {},
+            "clientInfo": {
+                "name": "ai-parrot",
+                "version": "1.0.0"
+            }
+        })
+        
+        if "error" in response:
+            raise MCPConnectionError(f"Initialize failed: {response['error']}")
+        
+        # Send initialized notification
+        await self._stream.write(mcp_pb2.MCPRequest(
+            json_rpc_message=json.dumps({
+                "jsonrpc": "2.0",
+                "method": "notifications/initialized"
+            })
+        ))
+    
+    async def list_tools(self) -> list:
+        """List available tools."""
+        response = await self.send_request("tools/list")
+        return response.get("result", {}).get("tools", [])
+    
+    async def call_tool(self, name: str, arguments: Dict) -> Dict:
+        """Call a tool."""
+        response = await self.send_request("tools/call", {
+            "name": name,
+            "arguments": arguments
+        })
+        return response.get("result", {})
+    
+    async def disconnect(self):
+        """Disconnect from gRPC server."""
+        self._initialized = False
+        
+        if self._stream_task and not self._stream_task.done():
+            self._stream_task.cancel()
+            with contextlib.suppress(asyncio.CancelledError):
+                await self._stream_task
+        
+        if self._channel:
+            await self._channel.close()
+            self._channel = None
+        
+        self.logger.info("gRPC disconnected")