Xingyao Wang c2f46200c0
chore(lint): Apply comprehensive linting and formatting fixes (#10287)
Co-authored-by: openhands <openhands@all-hands.dev>
2025-08-13 21:13:19 +02:00

147 lines
4.8 KiB
Python

"""MCP Proxy Manager for OpenHands.
This module provides a manager class for handling FastMCP proxy instances,
including initialization, configuration, and mounting to FastAPI applications.
"""
import logging
from typing import Any, Optional
from anyio import get_cancelled_exc_class
from fastapi import FastAPI
from fastmcp import FastMCP
from fastmcp.utilities.logging import get_logger as fastmcp_get_logger
from openhands.core.config.mcp_config import MCPStdioServerConfig
logger = logging.getLogger(__name__)
fastmcp_logger = fastmcp_get_logger('fastmcp')
class MCPProxyManager:
"""Manager for FastMCP proxy instances.
This class encapsulates all the functionality related to creating, configuring,
and managing FastMCP proxy instances, including mounting them to FastAPI applications.
"""
def __init__(
self,
auth_enabled: bool = False,
api_key: Optional[str] = None,
logger_level: Optional[int] = None,
):
"""Initialize the MCP Proxy Manager.
Args:
name: Name of the proxy server
auth_enabled: Whether authentication is enabled
api_key: API key for authentication (required if auth_enabled is True)
logger_level: Logging level for the FastMCP logger
"""
self.auth_enabled = auth_enabled
self.api_key = api_key
self.proxy: Optional[FastMCP] = None
# Initialize with a valid configuration format for FastMCP
self.config: dict[str, Any] = {
'mcpServers': {},
}
# Configure FastMCP logger
if logger_level is not None:
fastmcp_logger.setLevel(logger_level)
def initialize(self) -> None:
"""Initialize the FastMCP proxy with the current configuration."""
if len(self.config['mcpServers']) == 0:
logger.info(
'No MCP servers configured for FastMCP Proxy, skipping initialization.'
)
return None
# Create a new proxy with the current configuration
self.proxy = FastMCP.as_proxy(
self.config,
auth_enabled=self.auth_enabled,
api_key=self.api_key,
)
logger.info('FastMCP Proxy initialized successfully')
async def mount_to_app(
self, app: FastAPI, allow_origins: Optional[list[str]] = None
) -> None:
"""Mount the SSE server app to a FastAPI application.
Args:
app: FastAPI application to mount to
allow_origins: List of allowed origins for CORS
"""
if len(self.config['mcpServers']) == 0:
logger.info('No MCP servers configured for FastMCP Proxy, skipping mount.')
return
if not self.proxy:
raise ValueError('FastMCP Proxy is not initialized')
def close_on_double_start(app):
async def wrapped(scope, receive, send):
start_sent = False
async def check_send(message):
nonlocal start_sent
if message['type'] == 'http.response.start':
if start_sent:
raise get_cancelled_exc_class()(
'closed because of double http.response.start (mcp issue https://github.com/modelcontextprotocol/python-sdk/issues/883)'
)
start_sent = True
await send(message)
await app(scope, receive, check_send)
return wrapped
# Get the SSE app
# mcp_app = self.proxy.http_app(path='/shttp')
mcp_app = close_on_double_start(
self.proxy.http_app(path='/sse', transport='sse')
)
app.mount('/mcp', mcp_app)
# Remove any existing mounts at root path
if '/mcp' in app.routes:
app.routes.remove('/mcp')
app.mount('/', mcp_app)
logger.info('Mounted FastMCP Proxy app at /mcp')
async def update_and_remount(
self,
app: FastAPI,
stdio_servers: list[MCPStdioServerConfig],
allow_origins: Optional[list[str]] = None,
) -> None:
"""Update the tools configuration and remount the proxy to the app.
This is a convenience method that combines updating the tools,
shutting down the existing proxy, initializing a new one, and
mounting it to the app.
Args:
app: FastAPI application to mount to
tools: List of tool configurations
allow_origins: List of allowed origins for CORS
"""
tools = {t.name: t.model_dump() for t in stdio_servers}
self.config['mcpServers'] = tools
del self.proxy
self.proxy = None
# Initialize a new proxy
self.initialize()
# Mount the new proxy to the app
await self.mount_to_app(app, allow_origins)