CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-07 03:17:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Create StreamableTransport class for MCP Streamable http transport. This subclasses MCPTransportBase.

Browse Source
This commit is contained in:
Hemang
2025-06-03 15:43:25 +02:00
committed by Hemang Sarkar
parent 2e6f669923
commit 7ec13ad852
8 changed files with 456 additions and 501 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/common/constants.py
-22
View File
@@ -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"
gateway/mcp/constants.py
+23
View File
@@ -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"
gateway/mcp/mcp_sessions_manager.py
+3 -3
View File
@@ -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:
gateway/mcp/mcp_transport_base.py
+1 -8
View File
@@ -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."""
gateway/mcp/sse.py
+2 -4
View File
@@ -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,
gateway/mcp/stdio.py
+3 -5
View File
@@ -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"
gateway/mcp/streamable.py
+364 -379
View File
@@ -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
)
gateway/mcp/utils.py
+60 -80
View File
@@ -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