From 2e6f6699231283c5523490f43ed9313cb022cda7 Mon Sep 17 00:00:00 2001 From: Hemang Date: Tue, 3 Jun 2025 14:26:30 +0200 Subject: [PATCH] Create SSETransport class for MCP SSE transport. This subclasses MCPTransportBase. --- gateway/mcp/sse.py | 599 ++++++++++++++++++++++----------------------- 1 file changed, 298 insertions(+), 301 deletions(-) diff --git a/gateway/mcp/sse.py b/gateway/mcp/sse.py index aee7b44..e3668fb 100644 --- a/gateway/mcp/sse.py +++ b/gateway/mcp/sse.py @@ -3,7 +3,7 @@ import asyncio import json import re -from typing import Tuple +from typing import Any, AsyncGenerator, Optional, Tuple import httpx from httpx_sse import aconnect_sse, ServerSentEvent @@ -12,21 +12,15 @@ from fastapi.responses import StreamingResponse from gateway.common.constants import ( CLIENT_TIMEOUT, - MCP_METHOD, - MCP_TOOL_CALL, - MCP_LIST_TOOLS, UTF_8, ) from gateway.mcp.mcp_sessions_manager import ( McpSessionsManager, McpAttributes, ) +from gateway.mcp.mcp_transport_base import MCPTransportBase from gateway.mcp.utils import ( get_mcp_server_base_url, - hook_tool_call, - intercept_response, - update_mcp_server_in_session_metadata, - update_session_from_request, ) MCP_SERVER_POST_HEADERS = { @@ -43,337 +37,340 @@ MCP_SERVER_SSE_HEADERS = { MCP_SERVER_BASE_URL_HEADER = "mcp-server-base-url" gateway = APIRouter() -session_store = McpSessionsManager() +mcp_sessions_manager = McpSessionsManager() @gateway.post("/mcp/sse/messages/") -async def mcp_post_sse_gateway( - request: Request, -) -> Response: - """Proxy calls to the MCP Server tools""" - query_params = dict(request.query_params) - if not query_params.get("session_id"): - raise HTTPException( - status_code=400, - detail="Missing 'session_id' query parameter", - ) - if not session_store.session_exists(query_params.get("session_id")): - raise HTTPException( - status_code=400, - detail="Session does not exist", - ) - - session_id = query_params.get("session_id") - mcp_server_messages_endpoint = ( - get_mcp_server_base_url(request) + "/messages/?" + session_id - ) - request_body_bytes = await request.body() - request_body = json.loads(request_body_bytes) - session = session_store.get_session(session_id) - update_session_from_request(session, request_body) - - if request_body.get(MCP_METHOD) == MCP_TOOL_CALL: - # Intercept and potentially block the request - hook_tool_call_result, is_blocked = await hook_tool_call( - session_id=session_id, - session_store=session_store, - request_body=request_body, - ) - if is_blocked: - # Add the error message to the session. - # The error message is sent back to the client using the SSE stream. - await session.add_pending_error_message(hook_tool_call_result) - return Response(content="Accepted", status_code=202) - elif request_body.get(MCP_METHOD) == MCP_LIST_TOOLS: - # Intercept and potentially block the request - hook_tool_call_result, is_blocked = await hook_tool_call( - session_id=session_id, - session_store=session_store, - request_body={ - "id": request_body.get("id"), - "method": MCP_LIST_TOOLS, - "params": {"name": MCP_LIST_TOOLS, "arguments": {}}, - }, - ) - if is_blocked: - # Add the error message to the session. - # The error message is sent back to the client using the SSE stream. - await session.add_pending_error_message(hook_tool_call_result) - return Response(content="Accepted", status_code=202) - - async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: - try: - response = await client.post( - url=mcp_server_messages_endpoint, - headers={ - k: v - for k, v in request.headers.items() - if k.lower() in MCP_SERVER_POST_HEADERS - }, - json=request_body, - params=query_params, - ) - return Response( - content=response.content, - status_code=response.status_code, - headers={ - "X-Proxied-By": "mcp-gateway", - **response.headers, - }, - ) - - except httpx.RequestError as e: - print(f"[MCP POST] Request error: {str(e)}") - raise HTTPException(status_code=500, detail="Request error") from e - except Exception as e: - print(f"[MCP POST] Unexpected error: {str(e)}") - raise HTTPException(status_code=500, detail="Unexpected error") from e +async def mcp_post_sse_gateway(request: Request) -> Response: + """Proxy calls to the MCP Server using SSE transport strategy.""" + return await create_sse_transport_and_handle_post(request, mcp_sessions_manager) @gateway.get("/mcp/sse") -async def mcp_get_sse_gateway( - request: Request, -) -> StreamingResponse: - """Proxy calls to the MCP Server tools""" - mcp_server_sse_endpoint = get_mcp_server_base_url(request) + "/sse" +async def mcp_get_sse_gateway(request: Request) -> StreamingResponse: + """Proxy calls to the MCP Server using SSE transport strategy.""" + return await create_sse_transport_and_handle_stream(request, mcp_sessions_manager) + +async def create_sse_transport_and_handle_post( + request: Request, session_store: McpSessionsManager +) -> Response: + """Integration function for SSE POST route.""" query_params = dict(request.query_params) - response_headers = {} - filtered_headers = { - k: v for k, v in request.headers.items() if k.lower() in MCP_SERVER_SSE_HEADERS - } - sse_header_attributes = McpAttributes.from_request_headers(request.headers) + session_id = query_params.get("session_id") - async def event_generator(): - """ - Generate a merged stream of MCP server events and pending error messages. - The pending error messages are added in the POST messages handler. - This function runs in a loop, yielding events as they arrive. - """ - mcp_server_events_queue = asyncio.Queue() - pending_error_messages_queue = asyncio.Queue() - tasks = set() - session_id = None + if not session_id: + raise HTTPException( + status_code=400, detail="Missing 'session_id' query parameter" + ) + if not session_store.session_exists(session_id): + raise HTTPException(status_code=400, detail="Session does not exist") - try: - # MCP Server Events Processor - async def process_mcp_server_events(): - """Connect to MCP server and process its events.""" - nonlocal session_id + request_body = json.loads(await request.body()) + return await SSETransport(session_store).handle_post_request( + request, session_id, request_body + ) - async with httpx.AsyncClient( - timeout=httpx.Timeout(CLIENT_TIMEOUT) - ) as client: - try: - async with aconnect_sse( - client, - "GET", - mcp_server_sse_endpoint, - headers=filtered_headers, - params=query_params, - ) as event_source: - if event_source.response.status_code != 200: - error_content = await event_source.response.aread() - raise HTTPException( - status_code=event_source.response.status_code, - detail=error_content, - ) - response_headers.update( - dict(event_source.response.headers.items()) - ) - async for sse in event_source.aiter_sse(): - if sse.event == "endpoint": - ( - event_bytes, - extracted_id, - ) = await _handle_endpoint_event( - sse, sse_header_attributes +async def create_sse_transport_and_handle_stream( + request: Request, session_store: McpSessionsManager +) -> StreamingResponse: + """Integration function for SSE GET route.""" + return await SSETransport(session_store).handle_sse_stream(request) + + +class SSETransport(MCPTransportBase): + """ + Server-Sent Events transport implementation for MCP communication. + Handles HTTP-based SSE communication with message queuing. + """ + + async def initialize_session( + self, + *args, + **kwargs, + ) -> str: + """Initialize or get existing SSE session.""" + session_id: Optional[str] = kwargs.get("session_id", None) + session_attributes: Optional[McpAttributes] = kwargs.get( + "session_attributes", None + ) + if session_id and self.session_store.session_exists(session_id): + return session_id + + if not session_id: + raise ValueError("Session ID is required for SSE transport") + + if not self.session_store.session_exists(session_id): + if not session_attributes: + raise ValueError("Session attributes required for new session") + await self.session_store.initialize_session(session_id, session_attributes) + + return session_id + + async def handle_post_request( + self, request: Request, session_id: str, request_body: dict[str, Any] + ) -> Response: + """Handle POST request to SSE endpoint.""" + session = self.session_store.get_session(session_id) + + processed_request, is_blocked = await self.process_outgoing_request( + session_id, request_body + ) + + if is_blocked: + # Add the error message to the session for SSE delivery + await session.add_pending_error_message(processed_request) + return Response(content="Accepted", status_code=202) + + # Forward to MCP server + mcp_server_base_url = get_mcp_server_base_url(request) + mcp_server_messages_endpoint = f"{mcp_server_base_url}/messages/?{session_id}" + + # Filter headers for MCP server + mcp_headers = { + k: v + for k, v in request.headers.items() + if k.lower() in {"connection", "accept", "content-length", "content-type"} + } + + async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: + try: + response = await client.post( + url=mcp_server_messages_endpoint, + headers=mcp_headers, + json=request_body, + params=dict(request.query_params), + ) + return Response( + content=response.content, + status_code=response.status_code, + headers={"X-Proxied-By": "mcp-gateway", **response.headers}, + ) + except httpx.RequestError as e: + print(f"[MCP POST] Request error: {str(e)}") + raise HTTPException(status_code=500, detail="Request error") from e + + async def handle_sse_stream(self, request: Request) -> StreamingResponse: + """Handle SSE streaming connection.""" + mcp_server_base_url = get_mcp_server_base_url(request) + mcp_server_sse_endpoint = f"{mcp_server_base_url}/sse" + + query_params = dict(request.query_params) + response_headers = {} + + # Filter headers for SSE + filtered_headers = { + k: v + for k, v in request.headers.items() + if k.lower() in {"connection", "accept", "cache-control"} + } + + sse_header_attributes = McpAttributes.from_request_headers(request.headers) + + async def event_generator() -> AsyncGenerator[bytes, None]: + """ + Generate a merged stream of MCP server events and pending error messages. + The pending error messages are added in the POST messages handler. + This function runs in a loop, yielding events as they arrive. + """ + mcp_server_events_queue = asyncio.Queue() + pending_error_messages_queue = asyncio.Queue() + tasks = set() + session_id = None + + try: + # MCP Server Events Processor + async def process_mcp_server_events(): + nonlocal session_id + + async with httpx.AsyncClient( + timeout=httpx.Timeout(CLIENT_TIMEOUT) + ) as client: + try: + async with aconnect_sse( + client, + "GET", + mcp_server_sse_endpoint, + headers=filtered_headers, + params=query_params, + ) as event_source: + if event_source.response.status_code != 200: + error_content = await event_source.response.aread() + raise HTTPException( + status_code=event_source.response.status_code, + detail=error_content, ) - session_id = extracted_id - if ( - session_id - and "process_error_messages_task" - not in locals() - ): - process_error_messages_task = ( - asyncio.create_task( - _check_for_pending_error_messages( + response_headers.update( + dict(event_source.response.headers.items()) + ) + + async for sse in event_source.aiter_sse(): + if sse.event == "endpoint": + ( + event_bytes, + extracted_id, + ) = await self._handle_endpoint_event( + sse, sse_header_attributes + ) + session_id = extracted_id + + if ( + session_id + and "process_error_messages_task" + not in locals() + ): + process_error_messages_task = asyncio.create_task( + self._check_for_pending_error_messages( session_id, pending_error_messages_queue, ) ) + tasks.add(process_error_messages_task) + process_error_messages_task.add_done_callback( + tasks.discard + ) + + elif sse.event == "message" and session_id: + event_bytes = await self._handle_message_event( + session_id, sse ) - tasks.add(process_error_messages_task) - process_error_messages_task.add_done_callback( - tasks.discard + else: + event_bytes = f"event: {sse.event}\ndata: {sse.data}\n\n".encode( + UTF_8 ) - elif sse.event == "message" and session_id: - # Process message event - event_bytes = await _handle_message_event( - session_id, sse - ) - else: - # Pass through other event types - # pylint: disable=line-too-long - event_bytes = f"event: {sse.event}\ndata: {sse.data}\n\n".encode( - UTF_8 - ) + await mcp_server_events_queue.put(event_bytes) - # Put the processed event in the queue - await mcp_server_events_queue.put(event_bytes) + except httpx.StreamClosed as e: + print(f"Server stream closed: {e}") + except Exception as e: # pylint: disable=broad-except + print(f"Error processing server events: {e}") - except httpx.StreamClosed as e: - print(f"Server stream closed: {e}", flush=True) - except Exception as e: # pylint: disable=broad-except - print(f"Error processing server events: {e}", flush=True) - - # Start server events processor - mcp_server_events_task = asyncio.create_task(process_mcp_server_events()) - tasks.add(mcp_server_events_task) - mcp_server_events_task.add_done_callback(tasks.discard) - - # Main event loop: merge MCP server events and pending error messages - while True: - # Create futures for both queues - mcp_server_event_future = asyncio.create_task( - mcp_server_events_queue.get() - ) - pending_error_message_future = asyncio.create_task( - pending_error_messages_queue.get() + # Start server events processor + mcp_server_events_task = asyncio.create_task( + process_mcp_server_events() ) + tasks.add(mcp_server_events_task) + mcp_server_events_task.add_done_callback(tasks.discard) - # Wait for either queue to have an item, with timeout - done, pending = await asyncio.wait( - [mcp_server_event_future, pending_error_message_future], - return_when=asyncio.FIRST_COMPLETED, - timeout=0.25, - ) + # Main event loop: merge MCP server events and pending error messages + while True: + # Create futures for both queues + mcp_server_event_future = asyncio.create_task( + mcp_server_events_queue.get() + ) + pending_error_message_future = asyncio.create_task( + pending_error_messages_queue.get() + ) - for future in pending: - future.cancel() + # Wait for either queue to have an item, with timeout + done, pending = await asyncio.wait( + [mcp_server_event_future, pending_error_message_future], + return_when=asyncio.FIRST_COMPLETED, + timeout=0.25, + ) - # Timeout occurred and no future completed. - if not done: - continue + for future in pending: + future.cancel() - for future in done: - try: - event = await future - yield event - except asyncio.CancelledError: - # Future was cancelled, continue + # Timeout occurred and no future completed. + if not done: continue - finally: - # Clean up all tasks - for task in tasks: - task.cancel() + for future in done: + try: + event = await future + yield event + except asyncio.CancelledError: + # Future was cancelled, continue + continue - # Wait for all tasks to complete - if tasks: - await asyncio.wait(tasks, timeout=2) + finally: + # Clean up all tasks + for task in tasks: + task.cancel() + # Wait for all tasks to complete + if tasks: + await asyncio.wait(tasks, timeout=2) - # Return the streaming response - return StreamingResponse( - event_generator(), - media_type="text/event-stream", - headers={"X-Proxied-By": "mcp-gateway", **response_headers}, - ) - - -async def _handle_endpoint_event( - sse: ServerSentEvent, sse_header_attributes: McpAttributes -) -> Tuple[bytes, str]: - """ - Handle the endpoint event type and modify the data accordingly. - For endpoint events, we need to rewrite the endpoint to use our gateway. - - Args: - sse (ServerSentEvent): The original SSE object. - sse_header_attributes (SseHeaderAttributes): The header attributes from the request. - - Returns: - bytes: Modified SSE data as bytes. - str: session_id extracted from the data. - """ - # Extract session_id - match = re.search(r"session_id=([^&\s]+)", sse.data) - if match: - session_id = match.group(1) - # Initialize this session in our store if needed - if not session_store.session_exists(session_id): - await session_store.initialize_session(session_id, sse_header_attributes) - - # Rewrite the endpoint to use our gateway - modified_data = sse.data.replace( - "/messages/?session_id=", - "/api/v1/gateway/mcp/sse/messages/?session_id=", - ) - event_bytes = f"event: {sse.event}\ndata: {modified_data}\n\n".encode(UTF_8) - return event_bytes, session_id - - -async def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes: - """ - Handle the message event type. - - Args: - session_id (str): The session ID associated with the request. - sse (ServerSentEvent): The original SSE object. - """ - event_bytes = f"event: {sse.event}\ndata: {sse.data}\n\n".encode(UTF_8) - session = session_store.get_session(session_id) - try: - response_body = json.loads(sse.data) - update_mcp_server_in_session_metadata(session, response_body) - - intercept_response_result, is_blocked = await intercept_response( - session_id=session_id, - session_store=session_store, - response_body=response_body, + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers={"X-Proxied-By": "mcp-gateway", **response_headers}, ) - if is_blocked: - event_bytes = f"event: {sse.event}\ndata: {json.dumps(intercept_response_result)}\n\n".encode( - UTF_8 + + async def handle_communication(self, *args, **kwargs) -> StreamingResponse: + """Main communication handler for SSE transport.""" + return await self.handle_sse_stream(kwargs.get("request")) + + async def _handle_endpoint_event( + self, sse: ServerSentEvent, sse_header_attributes: McpAttributes + ) -> Tuple[bytes, str]: + """Handle endpoint event and initialize session if needed.""" + match = re.search(r"session_id=([^&\s]+)", sse.data) + session_id = match.group(1) if match else None + + if session_id: + # Initialize session if needed + await self.initialize_session( + session_id=session_id, session_attributes=sse_header_attributes ) - except json.JSONDecodeError as e: - print( - f"[MCP SSE] Error parsing message JSON: {e}", - flush=True, + + # Rewrite endpoint to use our gateway + modified_data = sse.data.replace( + "/messages/?session_id=", + "/api/v1/gateway/mcp/sse/messages/?session_id=", ) - except Exception as e: # pylint: disable=broad-except - print( - f"[MCP SSE] Error processing message: {e}", - flush=True, - ) - return event_bytes + event_bytes = f"event: {sse.event}\ndata: {modified_data}\n\n".encode(UTF_8) + return event_bytes, session_id + async def _handle_message_event( + self, session_id: str, sse: ServerSentEvent + ) -> bytes: + """Handle message event with guardrails processing.""" + try: + response_body = json.loads(sse.data) -async def _check_for_pending_error_messages( - session_id: str, pending_error_messages_queue: asyncio.Queue -): - """Periodically check for and enqueue pending error messages.""" - try: - while True: - try: - session = session_store.get_session(session_id) - error_messages = await session.get_pending_error_messages() + # Process response through guardrails + processed_response, is_blocked = await self.process_incoming_response( + session_id, response_body + ) - for error_message in error_messages: - error_bytes = ( - f"event: message\ndata: {json.dumps(error_message)}\n\n".encode( + event_bytes = f"event: {sse.event}\ndata: {sse.data}\n\n".encode(UTF_8) + if is_blocked: + event_bytes = f"event: {sse.event}\ndata: {json.dumps(processed_response)}\n\n".encode( + UTF_8 + ) + + return event_bytes + except json.JSONDecodeError as e: + print(f"[MCP SSE] Error parsing message JSON: {e}") + return f"event: {sse.event}\ndata: {sse.data}\n\n".encode(UTF_8) + except Exception as e: # pylint: disable=broad-except + print(f"[MCP SSE] Error processing message: {e}") + return f"event: {sse.event}\ndata: {sse.data}\n\n".encode(UTF_8) + + async def _check_for_pending_error_messages( + self, session_id: str, pending_error_messages_queue: asyncio.Queue + ): + """Periodically check for and enqueue pending error messages.""" + try: + while True: + try: + session = self.session_store.get_session(session_id) + error_messages = await session.get_pending_error_messages() + + for error_message in error_messages: + error_bytes = f"event: message\ndata: {json.dumps(error_message)}\n\n".encode( UTF_8 ) - ) - await pending_error_messages_queue.put(error_bytes) + await pending_error_messages_queue.put(error_bytes) - await asyncio.sleep(1) - except Exception as e: # pylint: disable=broad-except - print(f"Error checking for messages: {e}", flush=True) - await asyncio.sleep(1) - except asyncio.CancelledError: - # Task was cancelled, exit gracefully - return + await asyncio.sleep(1) + except Exception as e: # pylint: disable=broad-except + print(f"Error checking for messages: {e}") + await asyncio.sleep(1) + except asyncio.CancelledError: + return