wehub-resource-sync
137 lines
5.2 KiB
Python
137 lines
5.2 KiB
Python
import logging
|
|
import asyncio
|
|
|
|
from fastapi import Request, Response, HTTPException
|
|
from mcp.server.lowlevel.server import Server
|
|
from mcp.server.streamable_http_manager import StreamableHTTPSessionManager, EventStore
|
|
from mcp.server.transport_security import TransportSecuritySettings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class FastApiHttpSessionManager:
|
|
"""
|
|
FastAPI-native wrapper around StreamableHTTPSessionManager
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
mcp_server: Server,
|
|
event_store: EventStore | None = None,
|
|
json_response: bool = True, # Default to JSON for HTTP transport
|
|
security_settings: TransportSecuritySettings | None = None,
|
|
):
|
|
self.mcp_server = mcp_server
|
|
self.event_store = event_store
|
|
self.json_response = json_response
|
|
self.security_settings = security_settings
|
|
self._session_manager: StreamableHTTPSessionManager | None = None
|
|
self._manager_task: asyncio.Task | None = None
|
|
self._manager_started = False
|
|
self._startup_lock = asyncio.Lock()
|
|
|
|
async def _ensure_session_manager_started(self) -> None:
|
|
"""
|
|
Ensure the session manager is started.
|
|
|
|
This is called lazily on the first request to start the session manager
|
|
if it hasn't been started yet.
|
|
"""
|
|
if self._manager_started:
|
|
return
|
|
|
|
async with self._startup_lock:
|
|
if self._manager_started:
|
|
return
|
|
|
|
logger.debug("Starting StreamableHTTP session manager")
|
|
|
|
# Create the session manager
|
|
# Note: We don't use stateless=True because we want to support sessions
|
|
# but sessions are optional as per the MCP spec
|
|
self._session_manager = StreamableHTTPSessionManager(
|
|
app=self.mcp_server,
|
|
event_store=self.event_store,
|
|
json_response=self.json_response,
|
|
stateless=False, # Always support sessions, but they're optional
|
|
security_settings=self.security_settings,
|
|
)
|
|
|
|
# Start the session manager in a background task
|
|
async def run_session_manager():
|
|
try:
|
|
async with self._session_manager.run():
|
|
logger.info("StreamableHTTP session manager is running")
|
|
# Keep running until cancelled
|
|
await asyncio.Event().wait()
|
|
except asyncio.CancelledError:
|
|
logger.info("StreamableHTTP session manager is shutting down")
|
|
raise
|
|
except Exception:
|
|
logger.exception("Error in StreamableHTTP session manager")
|
|
raise
|
|
|
|
self._manager_task = asyncio.create_task(run_session_manager())
|
|
self._manager_started = True
|
|
|
|
# Give the session manager a moment to initialize
|
|
await asyncio.sleep(0.1)
|
|
|
|
async def handle_fastapi_request(self, request: Request) -> Response:
|
|
"""
|
|
Handle a FastAPI request by delegating to the session manager.
|
|
|
|
This converts FastAPI's Request/Response to ASGI scope/receive/send
|
|
and then converts the result back to a FastAPI Response.
|
|
"""
|
|
# Ensure session manager is started
|
|
await self._ensure_session_manager_started()
|
|
|
|
if not self._session_manager:
|
|
raise HTTPException(status_code=500, detail="Session manager not initialized")
|
|
|
|
logger.debug(f"Handling FastAPI request: {request.method} {request.url.path}")
|
|
|
|
# Capture the response from the session manager
|
|
response_started = False
|
|
response_status = 200
|
|
response_headers = []
|
|
response_body = b""
|
|
|
|
async def send_callback(message):
|
|
nonlocal response_started, response_status, response_headers, response_body
|
|
|
|
if message["type"] == "http.response.start":
|
|
response_started = True
|
|
response_status = message["status"]
|
|
response_headers = message.get("headers", [])
|
|
elif message["type"] == "http.response.body":
|
|
response_body += message.get("body", b"")
|
|
|
|
try:
|
|
# Delegate to the session manager's handle_request method
|
|
await self._session_manager.handle_request(request.scope, request.receive, send_callback)
|
|
|
|
# Convert the captured ASGI response to a FastAPI Response
|
|
headers_dict = {name.decode(): value.decode() for name, value in response_headers}
|
|
|
|
return Response(
|
|
content=response_body,
|
|
status_code=response_status,
|
|
headers=headers_dict,
|
|
)
|
|
|
|
except Exception:
|
|
logger.exception("Error in StreamableHTTPSessionManager")
|
|
raise HTTPException(status_code=500, detail="Internal server error")
|
|
|
|
async def shutdown(self) -> None:
|
|
"""Clean up the session manager and background task."""
|
|
if self._manager_task and not self._manager_task.done():
|
|
self._manager_task.cancel()
|
|
try:
|
|
await self._manager_task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
self._manager_started = False
|