Rewrite MCP server to use official SDK with Streamable HTTP

orchestrator/mcp_server.py

@@ -1,17 +1,20 @@
 """
-SSE-based MCP server for coordinator integration with Claude Code.
+MCP server for coordinator integration with Claude Code.
 
 Provides an MCP-compatible server that exposes coordinator tools
-via Server-Sent Events transport. Runs as a sidecar alongside
-the orchestrator.
+via Streamable HTTP transport using the official mcp Python SDK.
+Runs as a sidecar alongside the orchestrator.
 """
 
+import functools
 import json
 import sys
 import threading
 import time
 from pathlib import Path
 
+import anyio
+
 _shared_path = Path(__file__).parent.parent / "shared"
 if _shared_path.exists() and str(_shared_path) not in sys.path:
     sys.path.insert(0, str(_shared_path))
@@ -33,35 +36,39 @@ def get_logger(name: str, **kwargs) -> logging.Logger:
 
 
 class RateLimiter:
-    """Simple token bucket rate limiter."""
+    """Simple sliding-window rate limiter (async-safe)."""
 
     def __init__(self, max_requests: int = DEFAULT_RATE_LIMIT, window_seconds: int = 60):
         self.max_requests = max_requests
         self.window_seconds = window_seconds
         self._requests: list[float] = []
-        self._lock = threading.Lock()
 
     def allow(self) -> bool:
-        """Check if a request is allowed."""
+        """Check if a request is allowed.
+
+        Safe to call from the event loop — no locks, just list operations
+        which are atomic under the GIL in CPython and fast enough that
+        contention is not a concern for single-event-loop usage.
+        """
         now = time.time()
-        with self._lock:
-            # Remove expired entries
-            cutoff = now - self.window_seconds
-            self._requests = [t for t in self._requests if t > cutoff]
-            if len(self._requests) >= self.max_requests:
-                return False
-            self._requests.append(now)
-            return True
+        cutoff = now - self.window_seconds
+        self._requests = [t for t in self._requests if t > cutoff]
+        if len(self._requests) >= self.max_requests:
+            return False
+        self._requests.append(now)
+        return True
 
 
 class MCPServer:
-    """MCP server with SSE transport for coordinator tools.
+    """MCP server with Streamable HTTP transport for coordinator tools.
+
+    Uses the official mcp Python SDK (FastMCP) to expose coordinator tools
+    over the Streamable HTTP transport protocol.
 
     Provides:
-    - Tool listing endpoint
-    - Tool execution endpoint
-    - Health check endpoint
-    - Rate limiting
+    - MCP tools via Streamable HTTP at /mcp
+    - Health check endpoint at /health
+    - Rate limiting on tool calls
 
     No authentication required — localhost-only access is enforced via
     Docker port mapping (127.0.0.1 binding in docker-compose.yml).
@@ -80,83 +87,100 @@ def __init__(
         from mcp_tools import COORDINATOR_TOOLS, CoordinatorToolHandler
 
         self.tool_handler = CoordinatorToolHandler(orchestrator_url=orchestrator_url)
-        self.tools = COORDINATOR_TOOLS
-        self._app = None
+        self.tools_config = COORDINATOR_TOOLS
+        self._mcp = None
 
     def create_app(self):
-        """Create the Flask application for the MCP server."""
-        from flask import Flask, Response, jsonify, request
-
-        app = Flask("egg-mcp-server")
-
-        @app.route("/health")
-        def health():
-            return jsonify({"status": "healthy", "service": "egg-mcp-server"})
-
-        @app.route("/mcp/v1/tools", methods=["GET"])
-        def list_tools():
-            return jsonify({"tools": self.tools})
-
-        @app.route("/mcp/v1/tools/call", methods=["POST"])
-        def call_tool():
-            if not self.rate_limiter.allow():
-                return jsonify({"error": "Rate limit exceeded"}), 429
-
-            data = request.get_json()
-            if not data:
-                return jsonify({"error": "Missing request body"}), 400
-
-            tool_name = data.get("name")
-            arguments = data.get("arguments", {})
-
-            if not tool_name:
-                return jsonify({"error": "Missing tool name"}), 400
-
-            result = self.tool_handler.handle_tool_call(tool_name, arguments)
-
-            return jsonify(
-                {
-                    "content": [{"type": "text", "text": json.dumps(result, indent=2)}],
-                    "isError": "error" in result,
-                }
-            )
-
-        @app.route("/mcp/v1/sse")
-        def sse_stream():
-            """SSE endpoint for MCP protocol events."""
-
-            def generate():
-                # Send initial tools list
-                tools_event = json.dumps({"type": "tools_list", "tools": self.tools})
-                yield f"data: {tools_event}\n\n"
-
-                # Keep connection alive with heartbeats
-                while True:
-                    time.sleep(15)
-                    yield ": heartbeat\n\n"
-
-            return Response(
-                generate(),
-                mimetype="text/event-stream",
-                headers={
-                    "Cache-Control": "no-cache",
-                    "Connection": "keep-alive",
-                    "X-Accel-Buffering": "no",
-                },
-            )
-
-        self._app = app
-        return app
-
-    def run(self, host: str = "0.0.0.0", debug: bool = False):
+        """Create the FastMCP application with coordinator tools."""
+        from mcp.server.fastmcp import FastMCP
+
+        mcp = FastMCP(
+            "egg-mcp-server",
+            host="0.0.0.0",
+            port=self.port,
+            streamable_http_path="/mcp",
+            stateless_http=True,
+            json_response=True,
+        )
+
+        rate_limiter = self.rate_limiter
+        tool_handler = self.tool_handler
+
+        # Register /health as a custom route
+        @mcp.custom_route("/health", methods=["GET"])
+        async def health(request):
+            from starlette.responses import JSONResponse
+
+            return JSONResponse({"status": "healthy", "service": "egg-mcp-server"})
+
+        # Register each coordinator tool with FastMCP.
+        # We create wrapper functions that delegate to CoordinatorToolHandler.
+        def _make_tool_fn(tool_name: str, tool_schema: dict):
+            """Build an async tool function for FastMCP from a tool schema."""
+            required = set(tool_schema.get("required", []))
+            properties = tool_schema.get("properties", {})
+
+            async def tool_fn(**kwargs) -> str:
+                if not rate_limiter.allow():
+                    return json.dumps({"error": "Rate limit exceeded"})
+                result = await anyio.to_thread.run_sync(
+                    functools.partial(tool_handler.handle_tool_call, tool_name, kwargs)
+                )
+                return json.dumps(result, indent=2)
+
+            # Build a useful signature so FastMCP can inspect parameters
+            import inspect
+
+            params = []
+            for prop_name, prop_def in properties.items():
+                default = prop_def.get("default", inspect.Parameter.empty)
+                if prop_name not in required and default is inspect.Parameter.empty:
+                    default = None
+                params.append(
+                    inspect.Parameter(
+                        prop_name,
+                        inspect.Parameter.KEYWORD_ONLY,
+                        default=default,
+                        annotation=_json_type_to_python(prop_def),
+                    )
+                )
+            tool_fn.__signature__ = inspect.Signature(params, return_annotation=str)
+            tool_fn.__name__ = tool_name
+            tool_fn.__qualname__ = tool_name
+            return tool_fn
+
+        for tool_def in self.tools_config:
+            fn = _make_tool_fn(tool_def["name"], tool_def["inputSchema"])
+            mcp.tool(
+                name=tool_def["name"],
+                description=tool_def["description"],
+            )(fn)
+
+        self._mcp = mcp
+        return mcp
+
+    def run(self):
         """Start the MCP server.
 
         Binds to 0.0.0.0 inside the container so Docker port forwarding works.
-        Localhost-only access is enforced by the docker-compose port mapping.
+        Host is set in the FastMCP constructor; localhost-only access is enforced
+        by the docker-compose port mapping.
         """
-        app = self.create_app()
-        logger.info("Starting MCP server", port=self.port, host=host)
-        app.run(host=host, port=self.port, debug=debug, threaded=True)
+        mcp = self.create_app()
+        logger.info("Starting MCP server", port=self.port)
+        mcp.run(transport="streamable-http")
+
+
+def _json_type_to_python(prop_def: dict) -> type:
+    """Map JSON Schema type to Python type annotation for FastMCP."""
+    json_type = prop_def.get("type", "string")
+    mapping = {
+        "string": str,
+        "integer": int,
+        "number": float,
+        "boolean": bool,
+    }
+    return mapping.get(json_type, str)
 
 
 def start_mcp_server(

orchestrator/requirements.txt

@@ -16,3 +16,6 @@ pyyaml>=6.0.0
 
 # Requests for simple HTTP calls
 requests>=2.31.0
+
+# MCP SDK for Streamable HTTP transport (coordinator MCP server)
+mcp[cli]>=1.20.0,<2.0.0