diff --git a/gateway/common/constants.py b/gateway/common/constants.py index 00cfa93..a84eaea 100644 --- a/gateway/common/constants.py +++ b/gateway/common/constants.py @@ -16,25 +16,3 @@ IGNORED_HEADERS = [ CLIENT_TIMEOUT = 60.0 -# MCP related constants -INVARIANT_SESSION_ID_PREFIX = "inv-" -MCP_METHOD = "method" -MCP_TOOL_CALL = "tools/call" -MCP_LIST_TOOLS = "tools/list" -MCP_PARAMS = "params" -MCP_RESULT = "result" -MCP_SERVER_INFO = "serverInfo" -MCP_CLIENT_INFO = "clientInfo" -INVARIANT_GUARDRAILS_BLOCKED_MESSAGE = """ - [Invariant Guardrails] The MCP tool call was blocked for security reasons. - Do not attempt to circumvent this block, rather explain to the user based - on the following output what went wrong: %s - """ -INVARIANT_GUARDRAILS_BLOCKED_TOOLS_MESSAGE = """ - [Invariant Guardrails] This server was blocked from advertising its tools due to a security guardrail failure. - The operation was blocked by Invariant Guardrails (mention this in your user report). - When users ask about this tool, inform them that it was blocked due to a security guardrail failure. - %s - """ -MCP_SERVER_BASE_URL_HEADER = "mcp-server-base-url" -UTF_8 = "utf-8" diff --git a/gateway/mcp/constants.py b/gateway/mcp/constants.py new file mode 100644 index 0000000..a783d7f --- /dev/null +++ b/gateway/mcp/constants.py @@ -0,0 +1,23 @@ +"""Constants for the MCP (Model Context Protocol) Gateway.""" + +INVARIANT_SESSION_ID_PREFIX = "inv-" +MCP_METHOD = "method" +MCP_TOOL_CALL = "tools/call" +MCP_LIST_TOOLS = "tools/list" +MCP_PARAMS = "params" +MCP_RESULT = "result" +MCP_SERVER_INFO = "serverInfo" +MCP_CLIENT_INFO = "clientInfo" +INVARIANT_GUARDRAILS_BLOCKED_MESSAGE = """ + [Invariant Guardrails] The MCP tool call was blocked for security reasons. + Do not attempt to circumvent this block, rather explain to the user based + on the following output what went wrong: %s + """ +INVARIANT_GUARDRAILS_BLOCKED_TOOLS_MESSAGE = """ + [Invariant Guardrails] This server was blocked from advertising its tools due to a security guardrail failure. + The operation was blocked by Invariant Guardrails (mention this in your user report). + When users ask about this tool, inform them that it was blocked due to a security guardrail failure. + %s + """ +MCP_SERVER_BASE_URL_HEADER = "mcp-server-base-url" +UTF_8 = "utf-8" \ No newline at end of file diff --git a/gateway/mcp/mcp_sessions_manager.py b/gateway/mcp/mcp_sessions_manager.py index 0a0e0f7..2cf88ae 100644 --- a/gateway/mcp/mcp_sessions_manager.py +++ b/gateway/mcp/mcp_sessions_manager.py @@ -7,7 +7,6 @@ import getpass import os import random import socket - from typing import Any, Optional from invariant_sdk.async_client import AsyncClient @@ -16,14 +15,15 @@ from invariant_sdk.types.push_traces import PushTracesRequest from pydantic import BaseModel, Field, PrivateAttr from starlette.datastructures import Headers -from gateway.common.constants import DEFAULT_API_URL, INVARIANT_SESSION_ID_PREFIX -from gateway.common.guardrails import GuardrailRuleSet, GuardrailAction +from gateway.common.constants import DEFAULT_API_URL +from gateway.common.guardrails import GuardrailAction, GuardrailRuleSet from gateway.common.request_context import RequestContext from gateway.integrations.explorer import ( create_annotations_from_guardrails_errors, fetch_guardrails_from_explorer, ) from gateway.integrations.guardrails import check_guardrails +from gateway.mcp.constants import INVARIANT_SESSION_ID_PREFIX def user_and_host() -> str: diff --git a/gateway/mcp/mcp_transport_base.py b/gateway/mcp/mcp_transport_base.py index 3cc8d81..6c9c5ff 100644 --- a/gateway/mcp/mcp_transport_base.py +++ b/gateway/mcp/mcp_transport_base.py @@ -7,7 +7,7 @@ This module defines an abstract base class for MCP transports. from abc import ABC, abstractmethod from typing import Any, Tuple -from gateway.common.constants import ( +from gateway.mcp.constants import ( MCP_METHOD, MCP_TOOL_CALL, MCP_LIST_TOOLS, @@ -100,13 +100,6 @@ class MCPTransportBase(ABC): return interception_result, is_blocked - def _is_initialization_request(self, request_data: dict[str, Any]) -> bool: - """Check if request is an initialization request.""" - return ( - request_data.get("method") in ["initialize", "notifications/initialized"] - and "jsonrpc" in request_data - ) - @abstractmethod async def initialize_session(self, *args, **kwargs) -> str: """Initialize a session for this transport type.""" diff --git a/gateway/mcp/sse.py b/gateway/mcp/sse.py index e3668fb..3feb47a 100644 --- a/gateway/mcp/sse.py +++ b/gateway/mcp/sse.py @@ -10,10 +10,8 @@ from httpx_sse import aconnect_sse, ServerSentEvent from fastapi import APIRouter, HTTPException, Request, Response from fastapi.responses import StreamingResponse -from gateway.common.constants import ( - CLIENT_TIMEOUT, - UTF_8, -) +from gateway.common.constants import CLIENT_TIMEOUT +from gateway.mcp.constants import UTF_8 from gateway.mcp.mcp_sessions_manager import ( McpSessionsManager, McpAttributes, diff --git a/gateway/mcp/stdio.py b/gateway/mcp/stdio.py index d446048..c86b07c 100644 --- a/gateway/mcp/stdio.py +++ b/gateway/mcp/stdio.py @@ -9,18 +9,16 @@ import subprocess import sys from typing import Optional, Tuple -from gateway.common.constants import ( - UTF_8, -) -from gateway.mcp.mcp_transport_base import MCPTransportBase +from gateway.mcp.constants import UTF_8 +from gateway.mcp.log import mcp_log, MCP_LOG_FILE from gateway.mcp.mcp_sessions_manager import ( McpAttributes, McpSessionsManager, ) +from gateway.mcp.mcp_transport_base import MCPTransportBase from gateway.mcp.utils import ( generate_session_id, ) -from gateway.mcp.log import mcp_log, MCP_LOG_FILE STATUS_EOF = "eof" STATUS_DATA = "data" diff --git a/gateway/mcp/streamable.py b/gateway/mcp/streamable.py index 0716565..08d0440 100644 --- a/gateway/mcp/streamable.py +++ b/gateway/mcp/streamable.py @@ -1,36 +1,32 @@ """Gateway service to forward requests to the MCP Streamable HTTP servers""" import json +from typing import Any, Optional, Union import httpx - from httpx_sse import aconnect_sse from fastapi import APIRouter, HTTPException, Request, Response from fastapi.responses import StreamingResponse -from gateway.common.constants import ( - CLIENT_TIMEOUT, + +from gateway.common.constants import CLIENT_TIMEOUT +from gateway.mcp.constants import ( INVARIANT_SESSION_ID_PREFIX, - MCP_LIST_TOOLS, - MCP_METHOD, - MCP_TOOL_CALL, UTF_8, ) from gateway.mcp.mcp_sessions_manager import ( McpSessionsManager, McpAttributes, ) +from gateway.mcp.mcp_transport_base import MCPTransportBase from gateway.mcp.utils import ( generate_session_id, get_mcp_server_base_url, - hook_tool_call, - intercept_response, update_mcp_client_info_in_session, update_mcp_server_in_session_metadata, update_tool_call_id_in_session, ) - gateway = APIRouter() -session_store = McpSessionsManager() +mcp_sessions_manager = McpSessionsManager() CONTENT_TYPE_JSON = "application/json" CONTENT_TYPE_SSE = "text/event-stream" @@ -39,7 +35,6 @@ MCP_SESSION_ID_HEADER = "mcp-session-id" MCP_SERVER_POST_DELETE_HEADERS = { "connection", "accept", - "content-length", CONTENT_TYPE_HEADER, MCP_SESSION_ID_HEADER, } @@ -52,400 +47,390 @@ MCP_SERVER_GET_HEADERS = { @gateway.post("/mcp/streamable") -async def mcp_post_streamable_gateway(request: Request) -> StreamingResponse: - """ - Forward a POST request to the MCP Streamable server. - """ - request_body_bytes = await request.body() - request_body = json.loads(request_body_bytes) - sse_header_attributes = McpAttributes.from_request_headers(request.headers) - session_id = request.headers.get(MCP_SESSION_ID_HEADER) - is_initialization_request = _is_initialization_request(request_body) - - if session_id: - # If a session ID is provided in the request headers, it was already initialized - # in McpSessionsManager. This might be a session ID returned by the MCP server - # or a session ID generated in the gateway. - update_tool_call_id_in_session( - session_store.get_session(session_id), request_body - ) - elif is_initialization_request: - # If this is an initialization request, we generate a session ID, - # We don't call initialize_session here because we don't know - # if the MCP server is running with stateless_http set to True or False. - # If later in the response from MCP server, we don't receive a session ID then this - # will be initialized and returned back to the client else this will be - # overwritten by the session ID returned by the MCP server. - session_id = generate_session_id() - - # Intercept the request and check for guardrails. - if not is_initialization_request: - request_interception_result = await _intercept_request(session_id, request_body) - if request_interception_result: - return request_interception_result - - async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: - try: - response = await client.post( - url=_get_mcp_server_endpoint(request), - headers=_get_headers_for_mcp_post_and_delete(request), - content=request_body_bytes, - follow_redirects=True, - ) - - # Try to extract session ID from MCP server response - resp_session_id = response.headers.get(MCP_SESSION_ID_HEADER) - - # If MCP returned a session ID and we haven't seen, initialize it - if resp_session_id: - if not session_store.session_exists(resp_session_id): - await session_store.initialize_session( - resp_session_id, sse_header_attributes - ) - session_id = resp_session_id - - # If no session ID is returned, and this is an init request, initialize our own - elif is_initialization_request and not session_store.session_exists( - session_id - ): - await session_store.initialize_session( - session_id, sse_header_attributes - ) - - # Update client info if this is an initialization request - if is_initialization_request: - update_mcp_client_info_in_session( - session_store.get_session(session_id), - request_body, - ) - - # If the response is JSON type, handle it as a JSON response. - if response.headers.get(CONTENT_TYPE_HEADER) == CONTENT_TYPE_JSON: - return await _handle_mcp_json_response( - session_id=session_id, - is_initialization_request=is_initialization_request, - response=response, - ) - - # Else return SSE streaming response - return await _handle_mcp_streaming_response( - session_id=session_id, - is_initialization_request=is_initialization_request, - response=response, - ) - except httpx.RequestError as e: - print(f"[MCP POST] Request error: {str(e)}", flush=True) - raise HTTPException(status_code=500, detail="Request error") from e - except Exception as e: - print(f"[MCP POST] Unexpected error: {str(e)}", flush=True) - raise HTTPException(status_code=500, detail="Unexpected error") from e +async def mcp_post_streamable_gateway( + request: Request, +): + """Forward a POST request to the MCP Streamable server using transport strategy.""" + return await create_streamable_transport_and_handle_request( + request, "POST", mcp_sessions_manager + ) @gateway.get("/mcp/streamable") -async def mcp_get_streamable_gateway( - request: Request, -) -> Response: - """ - Forward a GET request to the MCP Streamable server. - - This allows the server to communicate to the client without the client - first sending data via HTTP POST. The server can send JSON-RPC requests - and notifications on this stream. - """ - mcp_server_endpoint = _get_mcp_server_endpoint(request) - response_headers = {} - filtered_headers = { - k: v for k, v in request.headers.items() if k.lower() in MCP_SERVER_GET_HEADERS - } - - async def event_generator(): - """Connect to MCP server and process its events.""" - - async with httpx.AsyncClient(timeout=httpx.Timeout(CLIENT_TIMEOUT)) as client: - try: - async with aconnect_sse( - client, - "GET", - mcp_server_endpoint, - headers=filtered_headers, - ) 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(): - yield sse - - 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) - - return StreamingResponse( - event_generator(), - media_type=CONTENT_TYPE_SSE, - headers={"X-Proxied-By": "mcp-gateway", **response_headers}, +async def mcp_get_streamable_gateway(request: Request) -> StreamingResponse: + """Forward a GET request to the MCP Streamable server using transport strategy.""" + return await create_streamable_transport_and_handle_request( + request, "GET", mcp_sessions_manager ) @gateway.delete("/mcp/streamable") -async def mcp_delete_streamable_gateway( - request: Request, -) -> Response: +async def mcp_delete_streamable_gateway(request: Request) -> Response: + """Forward a DELETE request to the MCP Streamable server using transport strategy.""" + return await create_streamable_transport_and_handle_request( + request, "DELETE", mcp_sessions_manager + ) + + +async def create_streamable_transport_and_handle_request( + request: Request, method: str, session_store: McpSessionsManager +) -> Union[Response, StreamingResponse]: + """Integration function for streamable routes.""" + streamable_transport = StreamableTransport(session_store) + return await streamable_transport.handle_communication( + request=request, method=method + ) + + +class StreamableTransport(MCPTransportBase): """ - Forward a DELETE request to the MCP Streamable server for explicit session termination. + Streamable HTTP transport implementation for MCP communication. + Handles HTTP POST/GET/DELETE requests with JSON and streaming responses. """ - session_id = _get_session_id(request) - if not session_store.session_exists(session_id): - raise HTTPException( - status_code=400, - detail="Session does not exist", + + async def initialize_session( + self, + *args, + **kwargs, + ) -> str: + """Initialize streamable HTTP session.""" + session_id: Optional[str] = kwargs.get("session_id", None) + session_attributes: Optional[McpAttributes] = kwargs.get( + "session_attributes", None ) - if session_id.startswith(INVARIANT_SESSION_ID_PREFIX): - return Response( - content="", - status_code=200, - headers={ - "X-Proxied-By": "mcp-gateway", - }, - ) - mcp_server_endpoint = _get_mcp_server_endpoint(request) + is_initialization_request: bool = kwargs.get("is_initialization_request", False) + if session_id and self.session_store.session_exists(session_id): + return session_id - async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: - try: - response = await client.delete( - url=mcp_server_endpoint, - headers=_get_headers_for_mcp_post_and_delete(request), - ) - await session_store.cleanup_session_lock(session_id) - return Response( - content=response.content, - status_code=response.status_code, - headers={ - "X-Proxied-By": "mcp-gateway", - **response.headers, - }, - ) + if is_initialization_request and not session_id: + session_id = generate_session_id() - except httpx.RequestError as e: - print(f"[MCP DELETE] Request error: {str(e)}") - raise HTTPException(status_code=500, detail="Request error") from e - except Exception as e: - print(f"[MCP DELETE] Unexpected error: {str(e)}") - raise HTTPException(status_code=500, detail="Unexpected error") from e - - -def _get_headers_for_mcp_post_and_delete( - request: Request, -) -> dict: - """ - Get headers for MCP server POST and DELETE requests. - This function filters out headers that are not needed for the MCP server. - If there is a session ID header, it ensures that it does not start - with INVARIANT_SESSION_ID_PREFIX since those are generated by the gateway - and not the MCP server so these should not be sent to the MCP server. - """ - return { - k: v - for k, v in request.headers.items() if ( - (k.lower() in MCP_SERVER_POST_DELETE_HEADERS) - and not ( - k.lower() == MCP_SESSION_ID_HEADER - and v.startswith(INVARIANT_SESSION_ID_PREFIX) + session_id + and not self.session_store.session_exists(session_id) + and session_attributes + ): + await self.session_store.initialize_session(session_id, session_attributes) + + return session_id + + async def handle_post_request( + self, request: Request, request_body: dict[str, Any] + ) -> Union[Response, StreamingResponse]: + """Handle POST request to streamable endpoint.""" + session_attributes = McpAttributes.from_request_headers(request.headers) + session_id = request.headers.get(MCP_SESSION_ID_HEADER) + is_initialization_request = self._is_initialization_request(request_body) + + # Handle session initialization + if session_id: + update_tool_call_id_in_session( + self.session_store.get_session(session_id), request_body ) + elif is_initialization_request: + session_id = await self.initialize_session( + session_attributes=session_attributes, is_initialization_request=True + ) + + # Process request if not initialization + if not is_initialization_request: + request_interception_result = await self._process_non_init_request( + session_id, request_body + ) + if request_interception_result: + return request_interception_result + + # Forward to MCP server + return await self._forward_to_mcp_server( + request, + request_body, + session_id, + session_attributes, + is_initialization_request, ) - } + async def handle_get_request(self, request: Request) -> StreamingResponse: + """Handle GET request for server-initiated communication.""" + mcp_server_endpoint = self._get_mcp_server_endpoint(request) + response_headers = {} -def _get_session_id(request: Request) -> str: - """Extract the session ID from request headers.""" - session_id = request.headers.get(MCP_SESSION_ID_HEADER) - if not session_id: - raise HTTPException( - status_code=400, - detail=f"Missing {MCP_SESSION_ID_HEADER} header", - ) - return session_id + filtered_headers = { + k: v + for k, v in request.headers.items() + if k.lower() in MCP_SERVER_GET_HEADERS + } + async def event_generator(): + async with httpx.AsyncClient( + timeout=httpx.Timeout(CLIENT_TIMEOUT) + ) as client: + try: + async with aconnect_sse( + client, + "GET", + mcp_server_endpoint, + headers=filtered_headers, + ) 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, + ) -def _get_mcp_server_endpoint(request: Request) -> str: - """ - Extract the MCP server endpoint from the request headers. - """ - return get_mcp_server_base_url(request) + "/mcp/" - - -def _update_mcp_response_info_in_session( - session_id: str, response_body: dict, is_json_response: bool -) -> None: - """ - Update the MCP response info in the session metadata. - """ - session = session_store.get_session(session_id) - update_mcp_server_in_session_metadata(session, response_body) - session.attributes.metadata["server_response_type"] = ( - "json" if is_json_response else "sse" - ) - - -def _is_initialization_request(request_body: dict) -> bool: - """ - Check if the request is an initialization request. - An initialization request is a JSON-RPC request with method "initialize". - Once initialization is done, the client sends a notification "notifications/initialized". - This function checks for both cases. - """ - return ( - request_body.get("method") in ["initialize", "notifications/initialized"] - and "jsonrpc" in request_body - ) - - -async def _handle_mcp_json_response( - session_id: str, is_initialization_request: bool, response: Response -) -> Response: - """ - Handle the MCP JSON response. - It checks for guardrails and returns the response accordingly. - """ - # If the response is blocked by guardrails - # return the error message else return the response as is - response_content = response.content - # The server response is empty string when client sends "notifications/initialized" - response_body = ( - json.loads(response_content.decode(UTF_8)) if response_content else {} - ) - if response_body: - _update_mcp_response_info_in_session( - session_id=session_id, response_body=response_body, is_json_response=True - ) - response_code = response.status_code - - if not is_initialization_request: - intercept_response_result, blocked = await intercept_response( - session_id=session_id, - session_store=session_store, - response_body=response_body, - ) - if blocked: - response_content = json.dumps(intercept_response_result).encode(UTF_8) - response_code = 400 - - # Build response headers, injecting gateway generated session ID if missing - response_headers = { - "X-Proxied-By": "mcp-gateway", - **response.headers, - } - if MCP_SESSION_ID_HEADER not in response.headers: - response_headers[MCP_SESSION_ID_HEADER] = session_id - return Response( - content=response_content, - status_code=response_code, - headers=response_headers, - ) - - -async def _handle_mcp_streaming_response( - session_id: str, is_initialization_request: bool, response: Response -) -> StreamingResponse: - """ - Handle the MCP streaming response. - It checks for guardrails and returns the response accordingly. - """ - - async def event_generator(): - # Events from MCP server have two parts: - # 1. event: {type} -> contains the type of event - # 2. data: {data} -> contains the actual message - # We are reading line by line so we need to buffer so that we can - # send the entire event (with both type and data) together. - # Once we receive an empty line, we end the stream. - buffer = "" - async for line in response.aiter_lines(): - stripped_line = line.strip() - if not stripped_line: - break # End of stream - if buffer: - response_body = json.loads(stripped_line.split("data: ")[1].strip()) - if not is_initialization_request: - ( - intercept_response_result, - blocked, - ) = await intercept_response( - session_id=session_id, - session_store=session_store, - response_body=response_body, - ) - if blocked: - yield ( - f"{buffer}\n" - f"data: {json.dumps(intercept_response_result)}\n\n" + response_headers.update( + dict(event_source.response.headers.items()) ) - break - else: - _update_mcp_response_info_in_session( - session_id=session_id, - response_body=response_body, - is_json_response=False, + + async for sse in event_source.aiter_sse(): + yield sse + + 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}") + + return StreamingResponse( + event_generator(), + media_type=CONTENT_TYPE_SSE, + headers={"X-Proxied-By": "mcp-gateway", **response_headers}, + ) + + async def handle_delete_request(self, request: Request) -> Response: + """Handle DELETE request for session termination.""" + session_id = self._get_session_id(request) + + if not self.session_store.session_exists(session_id): + raise HTTPException(status_code=400, detail="Session does not exist") + + if session_id.startswith(INVARIANT_SESSION_ID_PREFIX): + return Response( + content="", status_code=200, headers={"X-Proxied-By": "mcp-gateway"} + ) + + mcp_server_endpoint = self._get_mcp_server_endpoint(request) + + async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: + try: + response = await client.delete( + url=mcp_server_endpoint, + headers=self._get_headers_for_mcp_post_and_delete(request), + ) + await self.session_store.cleanup_session_lock(session_id) + 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 DELETE] Request error: {str(e)}") + raise HTTPException(status_code=500, detail="Request error") from e + + async def handle_communication( + self, *args, **kwargs + ) -> Union[Response, StreamingResponse]: + """Main communication handler for streamable transport.""" + request = kwargs.get("request") + method = kwargs.get("method", "POST") + + if method == "POST": + request_body = json.loads(await request.body()) + return await self.handle_post_request(request, request_body) + elif method == "GET": + return await self.handle_get_request(request) + elif method == "DELETE": + return await self.handle_delete_request(request) + else: + raise HTTPException(status_code=405, detail="Method not allowed") + + async def _process_non_init_request( + self, session_id: str, request_body: dict[str, Any] + ) -> Optional[Response]: + """Process non-initialization requests for guardrails.""" + processed_request, is_blocked = await self.process_outgoing_request( + session_id, request_body + ) + + if is_blocked: + return Response( + content=json.dumps(processed_request), + status_code=400, + media_type=CONTENT_TYPE_JSON, + ) + return None + + async def _forward_to_mcp_server( + self, + request: Request, + request_body: dict[str, Any], + session_id: str, + session_attributes: McpAttributes, + is_initialization_request: bool, + ) -> Union[Response, StreamingResponse]: + """Forward request to MCP server and handle response.""" + async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: + try: + response = await client.post( + url=self._get_mcp_server_endpoint(request), + headers=self._get_headers_for_mcp_post_and_delete(request), + content=json.dumps(request_body).encode(), + follow_redirects=True, + ) + + # Handle session ID from MCP server response + resp_session_id = response.headers.get(MCP_SESSION_ID_HEADER) + + if resp_session_id: + if not self.session_store.session_exists(resp_session_id): + await self.session_store.initialize_session( + resp_session_id, session_attributes + ) + session_id = resp_session_id + elif ( + is_initialization_request + and not self.session_store.session_exists(session_id) + ): + await self.session_store.initialize_session( + session_id, session_attributes ) - yield f"{buffer}\n{stripped_line}\n\n" - # Clear the buffer for the next event - buffer = "" - else: - buffer = stripped_line - # Build response headers, injecting gateway generated session ID if missing - response_headers = { - "X-Proxied-By": "mcp-gateway", - **response.headers, - } - if MCP_SESSION_ID_HEADER not in response.headers: - response_headers[MCP_SESSION_ID_HEADER] = session_id + # Update client info for initialization requests + if is_initialization_request: + update_mcp_client_info_in_session( + self.session_store.get_session(session_id), request_body + ) - return StreamingResponse( - event_generator(), - media_type=CONTENT_TYPE_SSE, - headers=response_headers, - ) + # Handle response based on content type + if response.headers.get(CONTENT_TYPE_HEADER) == CONTENT_TYPE_JSON: + return await self._handle_json_response( + session_id, is_initialization_request, response + ) + else: + return await self._handle_streaming_response( + session_id, is_initialization_request, response + ) + 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 _intercept_request(session_id: str, request_body: dict) -> Response | None: - """ - Intercept the request and check for guardrails. - This function is used to intercept requests and check for guardrails. - If the request is blocked, it returns a message indicating the block reason. - """ - if request_body.get(MCP_METHOD) == MCP_TOOL_CALL: - hook_tool_call_result, is_blocked = await hook_tool_call( - session_id=session_id, - session_store=session_store, - request_body=request_body, + async def _handle_json_response( + self, session_id: str, is_initialization_request: bool, response: httpx.Response + ) -> Response: + """Handle JSON response from MCP server.""" + response_content = response.content + response_body = ( + json.loads(response_content.decode(UTF_8)) if response_content else {} ) - if is_blocked: - return Response( - content=json.dumps(hook_tool_call_result), - status_code=400, - media_type="application/json", + + if response_body: + self._update_mcp_response_info_in_session(session_id, response_body, True) + + response_code = response.status_code + + if not is_initialization_request and response_body: + processed_response, blocked = await self.process_incoming_response( + session_id, response_body ) - elif request_body.get(MCP_METHOD) == MCP_LIST_TOOLS: - 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 blocked: + response_content = json.dumps(processed_response).encode(UTF_8) + response_code = 400 + + # Build response headers + response_headers = {"X-Proxied-By": "mcp-gateway", **response.headers} + if MCP_SESSION_ID_HEADER not in response.headers: + response_headers[MCP_SESSION_ID_HEADER] = session_id + + return Response( + content=response_content, + status_code=response_code, + headers=response_headers, ) - if is_blocked: - return Response( - content=json.dumps(hook_tool_call_result), - status_code=400, - media_type="application/json", + + async def _handle_streaming_response( + self, session_id: str, is_initialization_request: bool, response: httpx.Response + ) -> StreamingResponse: + """Handle streaming response from MCP server.""" + + async def event_generator(): + buffer = "" + async for line in response.aiter_lines(): + stripped_line = line.strip() + if not stripped_line: + break + + if buffer: + response_body = json.loads(stripped_line.split("data: ")[1].strip()) + + if not is_initialization_request: + ( + processed_response, + blocked, + ) = await self.process_incoming_response( + session_id, response_body + ) + if blocked: + yield f"{buffer}\ndata: {json.dumps(processed_response)}\n\n" + break + else: + self._update_mcp_response_info_in_session( + session_id, response_body, False + ) + + yield f"{buffer}\n{stripped_line}\n\n" + buffer = "" + else: + buffer = stripped_line + + # Build response headers + response_headers = {"X-Proxied-By": "mcp-gateway", **response.headers} + if MCP_SESSION_ID_HEADER not in response.headers: + response_headers[MCP_SESSION_ID_HEADER] = session_id + + return StreamingResponse( + event_generator(), + media_type=CONTENT_TYPE_SSE, + headers=response_headers, + ) + + def _update_mcp_response_info_in_session( + self, session_id: str, response_body: dict, is_json_response: bool + ) -> None: + """Update MCP response info in session metadata.""" + session = self.session_store.get_session(session_id) + update_mcp_server_in_session_metadata(session, response_body) + session.attributes.metadata["server_response_type"] = ( + "json" if is_json_response else "sse" + ) + + def _get_headers_for_mcp_post_and_delete(self, request: Request) -> dict: + """Get filtered headers for MCP server requests.""" + return { + k: v + for k, v in request.headers.items() + if ( + k.lower() in MCP_SERVER_POST_DELETE_HEADERS + and not ( + k.lower() == MCP_SESSION_ID_HEADER + and v.startswith(INVARIANT_SESSION_ID_PREFIX) + ) ) - return None + } + + def _get_session_id(self, request: Request) -> str: + """Extract session ID from request headers.""" + session_id = request.headers.get(MCP_SESSION_ID_HEADER) + if not session_id: + raise HTTPException(status_code=400, detail="Missing mcp-session-id header") + return session_id + + def _get_mcp_server_endpoint(self, request: Request) -> str: + """Get MCP server endpoint URL.""" + return get_mcp_server_base_url(request) + "/mcp/" + + def _is_initialization_request(self, request_data: dict[str, Any]) -> bool: + """Check if request is an initialization request.""" + return ( + request_data.get("method") in ["initialize", "notifications/initialized"] + and "jsonrpc" in request_data + ) diff --git a/gateway/mcp/utils.py b/gateway/mcp/utils.py index 2903705..831ecc3 100644 --- a/gateway/mcp/utils.py +++ b/gateway/mcp/utils.py @@ -1,54 +1,34 @@ -"""MCP utility functions.""" +"""MCP utility functions - Updated to work with transport strategy pattern.""" import asyncio import json import re import uuid - from typing import Tuple from fastapi import Request, HTTPException -from gateway.common.constants import ( + +from gateway.common.guardrails import GuardrailAction +from gateway.integrations.explorer import create_annotations_from_guardrails_errors +from gateway.mcp.constants import ( INVARIANT_GUARDRAILS_BLOCKED_MESSAGE, INVARIANT_GUARDRAILS_BLOCKED_TOOLS_MESSAGE, INVARIANT_SESSION_ID_PREFIX, MCP_CLIENT_INFO, - MCP_SERVER_BASE_URL_HEADER, MCP_LIST_TOOLS, MCP_METHOD, MCP_PARAMS, MCP_RESULT, + MCP_SERVER_BASE_URL_HEADER, MCP_SERVER_INFO, MCP_TOOL_CALL, ) -from gateway.common.guardrails import GuardrailAction -from gateway.mcp.mcp_sessions_manager import ( - McpSession, - McpSessionsManager, -) -from gateway.integrations.explorer import create_annotations_from_guardrails_errors from gateway.mcp.log import format_errors_in_response - - -def _check_if_new_errors( - session_id: str, session_store: McpSessionsManager, guardrails_result: dict -) -> bool: - """Checks if there are new errors in the guardrails result.""" - session = session_store.get_session(session_id) - annotations = create_annotations_from_guardrails_errors( - guardrails_result.get("errors", []) - ) - for annotation in annotations: - if annotation not in session.annotations: - return True - return False +from gateway.mcp.mcp_sessions_manager import McpSession, McpSessionsManager def generate_session_id() -> str: - """ - Generate a new session ID. - If the MCP server is session less then we don't have a session ID from the MCP server. - """ + """Generate a new session ID.""" return INVARIANT_SESSION_ID_PREFIX + uuid.uuid4().hex @@ -88,41 +68,8 @@ def update_session_from_request(session: McpSession, request_body: dict) -> None update_tool_call_id_in_session(session, request_body) -def _convert_localhost_to_docker_host(mcp_server_base_url: str) -> str: - """ - Convert localhost or 127.0.0.1 in an address to host.docker.internal - - Args: - mcp_server_base_url (str): The original server address from the header - - Returns: - str: Modified server address with localhost references changed to host.docker.internal - """ - if "localhost" in mcp_server_base_url or "127.0.0.1" in mcp_server_base_url: - # Replace localhost or 127.0.0.1 with host.docker.internal - modified_address = re.sub( - r"(https?://)(?:localhost|127\.0\.0\.1)(\b|:)", - r"\1host.docker.internal\2", - mcp_server_base_url, - ) - return modified_address - - return mcp_server_base_url - - def get_mcp_server_base_url(request: Request) -> str: - """ - Extract the MCP server base URL from the request headers. - - Args: - request (Request): The incoming request object. - - Returns: - str: The MCP server base URL. - - Raises: - HTTPException: If the MCP server base URL is not found in the headers. - """ + """Extract the MCP server base URL from the request headers.""" mcp_server_base_url = request.headers.get(MCP_SERVER_BASE_URL_HEADER) if not mcp_server_base_url: raise HTTPException( @@ -132,6 +79,32 @@ def get_mcp_server_base_url(request: Request) -> str: return _convert_localhost_to_docker_host(mcp_server_base_url).rstrip("/") +def _convert_localhost_to_docker_host(mcp_server_base_url: str) -> str: + """Convert localhost or 127.0.0.1 in an address to host.docker.internal.""" + if "localhost" in mcp_server_base_url or "127.0.0.1" in mcp_server_base_url: + modified_address = re.sub( + r"(https?://)(?:localhost|127\.0\.0\.1)(\b|:)", + r"\1host.docker.internal\2", + mcp_server_base_url, + ) + return modified_address + return mcp_server_base_url + + +def _check_if_new_errors( + session_id: str, session_store: McpSessionsManager, guardrails_result: dict +) -> bool: + """Checks if there are new errors in the guardrails result.""" + session = session_store.get_session(session_id) + annotations = create_annotations_from_guardrails_errors( + guardrails_result.get("errors", []) + ) + for annotation in annotations: + if annotation not in session.annotations: + return True + return False + + async def hook_tool_call( session_id: str, session_store: McpSessionsManager, request_body: dict ) -> Tuple[dict, bool]: @@ -157,14 +130,14 @@ async def hook_tool_call( }, } message = {"role": "assistant", "content": "", "tool_calls": [tool_call]} - # Check for blocking guardrails - this blocks until completion + + # Check for blocking guardrails session = session_store.get_session(session_id) guardrails_result = await session.get_guardrails_check_result( message, action=GuardrailAction.BLOCK ) - # If the request is blocked, return a message indicating the block reason. - # If there are new errors, run append_and_push_trace in background. - # If there are no new errors, just return the original request. + + # If the request is blocked, return error message if ( guardrails_result and guardrails_result.get("errors", []) @@ -187,6 +160,7 @@ async def hook_tool_call( % guardrails_result["errors"], }, }, True + # Push trace to the explorer await session_store.add_message_to_session(session_id, message, guardrails_result) return request_body, False @@ -197,7 +171,7 @@ async def hook_tool_call_response( session_store: McpSessionsManager, response_body: dict, is_tools_list=False, -) -> dict: +) -> Tuple[dict, bool]: """ Hook to process the response JSON after receiving it from the MCP server. @@ -207,17 +181,21 @@ async def hook_tool_call_response( response_body (dict): The response JSON to be processed. is_tools_list (bool): Flag to indicate if the response is from a tools/list call. Returns: - dict: The response JSON is returned if no guardrail is violated - else an error dict is returned. + Tuple[dict, bool]: A tuple containing the processed response JSON + and a boolean indicating whether the response was blocked. If the response + is blocked, the dict will contain an error message else it will contain the + original response. """ is_blocked = False result = response_body + message = { "role": "tool", "tool_call_id": f"call_{result.get('id')}", - "content": result.get(MCP_RESULT).get("content"), - "error": result.get(MCP_RESULT).get("error"), + "content": result.get(MCP_RESULT, {}).get("content"), + "error": result.get(MCP_RESULT, {}).get("error"), } + session = session_store.get_session(session_id) guardrails_result = await session.get_guardrails_check_result( message, action=GuardrailAction.BLOCK @@ -229,7 +207,7 @@ async def hook_tool_call_response( and _check_if_new_errors(session_id, session_store, guardrails_result) ): is_blocked = True - # If the request is blocked, return a message indicating the block reason + if not is_tools_list: result = { "jsonrpc": "2.0", @@ -241,7 +219,7 @@ async def hook_tool_call_response( }, } else: - # special error response for tools/list tool call + # Special error response for tools/list result = { "jsonrpc": "2.0", "id": response_body.get("id"), @@ -261,7 +239,7 @@ async def hook_tool_call_response( "title": "This tool was blocked by security guardrails.", }, } - for tool in response_body["result"]["tools"] + for tool in response_body.get("result", {}).get("tools", []) ] }, } @@ -295,6 +273,7 @@ async def intercept_response( intercept_response_result = response_body is_blocked = False + # Intercept and potentially block tool call response if method == MCP_TOOL_CALL: intercept_response_result, is_blocked = await hook_tool_call_response( @@ -304,10 +283,10 @@ async def intercept_response( ) # Intercept and potentially block list tool call response elif method == MCP_LIST_TOOLS: - # store tools in metadata - session_store.get_session(session_id).attributes.metadata["tools"] = ( - response_body.get(MCP_RESULT).get("tools") - ) + # Store tools in metadata + tools = response_body.get(MCP_RESULT, {}).get("tools", []) + session_store.get_session(session_id).attributes.metadata["tools"] = tools + intercept_response_result, is_blocked = await hook_tool_call_response( session_id=session_id, session_store=session_store, @@ -315,10 +294,11 @@ async def intercept_response( "jsonrpc": "2.0", "id": response_body.get("id"), "result": { - "content": json.dumps(response_body.get(MCP_RESULT).get("tools")), - "tools": response_body.get(MCP_RESULT).get("tools"), + "content": json.dumps(tools), + "tools": tools, }, }, is_tools_list=True, ) + return intercept_response_result, is_blocked