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

Add MCP guardrailing for SSE.

Browse Source
This commit is contained in:
Hemang
2025-05-08 20:39:37 +05:30
parent 40ec6d2db2
commit 794aae0326
4 changed files with 255 additions and 34 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/common/constants.py
+14
View File
@@ -13,3 +13,17 @@ IGNORED_HEADERS = [
]
CLIENT_TIMEOUT = 60.0
# MCP related constants
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
"""
gateway/common/mcp_sessions_manager.py
+123 -7
View File
@@ -13,6 +13,14 @@ from invariant_sdk.types.push_traces import PushTracesRequest
from pydantic import BaseModel, Field, PrivateAttr
from starlette.datastructures import Headers
from gateway.common.guardrails import GuardrailRuleSet, GuardrailAction
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
DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
@@ -29,10 +37,50 @@ class McpSession(BaseModel):
push_explorer: bool
trace_id: Optional[str] = None
last_trace_length: int = 0
annotations: List[Dict[str, Any]] = Field(default_factory=list)
guardrails: GuardrailRuleSet = Field(
default_factory=lambda: GuardrailRuleSet(
blocking_guardrails=[], logging_guardrails=[]
)
)
# Lock to maintain in-order pushes to explorer
_lock: asyncio.Lock = PrivateAttr(default_factory=asyncio.Lock)
async def load_guardrails(self) -> None:
"""
Load guardrails for the session.
This method fetches guardrails from the Invariant Explorer and assigns them to the session.
"""
self.guardrails = await fetch_guardrails_from_explorer(
self.explorer_dataset,
"Bearer " + os.getenv("INVARIANT_API_KEY"),
)
def _deduplicate_annotations(self, new_annotations: list) -> list:
"""Deduplicate new_annotations using the annotations in the session."""
deduped_annotations = []
for annotation in new_annotations:
# Check if an annotation with the same content and address exists in self.annotations
# TODO: Rely on the __eq__ method of the AnnotationCreate class directly via not in
# to remove duplicates instead of using a custom logic.
# This is a temporary solution until the Invariant SDK is updated.
is_duplicate = False
for current_annotation in self.annotations:
if (
annotation.content == current_annotation.content
and annotation.address == current_annotation.address
and annotation.extra_metadata == current_annotation.extra_metadata
):
is_duplicate = True
break
if not is_duplicate:
deduped_annotations.append(annotation)
return deduped_annotations
@contextlib.asynccontextmanager
async def session_lock(self):
"""
@@ -45,7 +93,45 @@ class McpSession(BaseModel):
async with self._lock:
yield
async def add_message(self, message: Dict[str, Any]) -> None:
async def get_guardrails_check_result(
self,
message: dict,
action: GuardrailAction = GuardrailAction.BLOCK,
) -> dict:
"""
Check against guardrails of type action.
"""
# Skip if no guardrails are configured for this action
if not (
(self.guardrails.blocking_guardrails and action == GuardrailAction.BLOCK)
or (self.guardrails.logging_guardrails and action == GuardrailAction.LOG)
):
return {}
# Prepare context and select appropriate guardrails
context = RequestContext.create(
request_json={},
dataset_name=self.explorer_dataset,
invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
guardrails=self.guardrails,
)
guardrails_to_check = (
self.guardrails.blocking_guardrails
if action == GuardrailAction.BLOCK
else self.guardrails.logging_guardrails
)
result = await check_guardrails(
messages=self.messages + [message],
guardrails=guardrails_to_check,
context=context,
)
return result
async def add_message(
self, message: Dict[str, Any], guardrails_result=Dict
) -> None:
"""
Add a message to the session and optionally push to explorer.
@@ -53,13 +139,35 @@ class McpSession(BaseModel):
message: The message to add
"""
async with self.session_lock():
annotations = []
if guardrails_result and guardrails_result.get("errors", []):
annotations = create_annotations_from_guardrails_errors(
guardrails_result.get("errors")
)
if self.guardrails.logging_guardrails:
logging_guardrails_check_result = (
await self.get_guardrails_check_result(
message, action=GuardrailAction.LOG
)
)
if (
logging_guardrails_check_result
and logging_guardrails_check_result.get("errors", [])
):
annotations.extend(
create_annotations_from_guardrails_errors(
logging_guardrails_check_result["errors"]
)
)
deduplicated_annotations = self._deduplicate_annotations(annotations)
# pylint: disable=no-member
self.messages.append(message)
# If push_explorer is enabled, push the trace
if self.push_explorer:
await self._push_trace_update()
await self._push_trace_update(deduplicated_annotations)
async def _push_trace_update(self) -> None:
async def _push_trace_update(self, deduplicated_annotations: list) -> None:
"""
Push trace updates to the explorer.
@@ -86,6 +194,7 @@ class McpSession(BaseModel):
messages=[self.messages],
dataset=self.explorer_dataset,
metadata=[metadata],
annotations=[deduplicated_annotations],
)
)
self.trace_id = response.id[0]
@@ -96,8 +205,11 @@ class McpSession(BaseModel):
AppendMessagesRequest(
trace_id=self.trace_id,
messages=new_messages,
annotations=deduplicated_annotations,
)
)
# pylint: disable=no-member
self.annotations.extend(deduplicated_annotations)
self.last_trace_length = len(self.messages)
except Exception as e: # pylint: disable=broad-except
print(f"[MCP SSE] Error pushing trace for session {self.session_id}: {e}")
@@ -151,16 +263,19 @@ class McpSessionsManager:
"""Check if a session exists"""
return session_id in self._sessions
def initialize_session(
async def initialize_session(
self, session_id: str, sse_header_attributes: SseHeaderAttributes
) -> None:
"""Initialize a new session"""
if session_id not in self._sessions:
self._sessions[session_id] = McpSession(
session = McpSession(
session_id=session_id,
explorer_dataset=sse_header_attributes.explorer_dataset,
push_explorer=sse_header_attributes.push_explorer,
)
self._sessions[session_id] = session
# Load guardrails for the session from the explorer
await session.load_guardrails()
def get_session(self, session_id: str) -> McpSession:
"""Get a session by ID"""
@@ -169,7 +284,7 @@ class McpSessionsManager:
return self._sessions.get(session_id)
async def add_message_to_session(
self, session_id: str, message: Dict[str, Any]
self, session_id: str, message: Dict[str, Any], guardrails_result: dict
) -> None:
"""
Add a message to a session and push to explorer if enabled.
@@ -177,6 +292,7 @@ class McpSessionsManager:
Args:
session_id: The session ID
message: The message to add
guardrails_result: The result of the guardrails check
"""
session = self.get_session(session_id)
await session.add_message(message)
await session.add_message(message, guardrails_result)
gateway/mcp/mcp.py
+8 -8
View File
@@ -11,6 +11,12 @@ from invariant_sdk.async_client import AsyncClient
from invariant_sdk.types.append_messages import AppendMessagesRequest
from invariant_sdk.types.push_traces import PushTracesRequest
from gateway.common.constants import (
INVARIANT_GUARDRAILS_BLOCKED_MESSAGE,
MCP_METHOD,
MCP_TOOL_CALL,
MCP_LIST_TOOLS,
)
from gateway.common.guardrails import GuardrailAction
from gateway.common.request_context import RequestContext
from gateway.integrations.explorer import create_annotations_from_guardrails_errors
@@ -19,16 +25,8 @@ from gateway.mcp.log import mcp_log, MCP_LOG_FILE
from gateway.mcp.mcp_context import McpContext
from gateway.mcp.task_utils import run_task_in_background, run_task_sync
MCP_METHOD = "method"
UTF_8_ENCODING = "utf-8"
MCP_TOOL_CALL = "tools/call"
MCP_LIST_TOOLS = "tools/list"
MCP_INITIALIZE = "initialize"
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
"""
DEFAULT_API_URL = "https://explorer.invariantlabs.ai"
@@ -312,6 +310,7 @@ def stream_and_forward_stderr(
MCP_LOG_FILE.buffer.write(line)
MCP_LOG_FILE.buffer.flush()
def run_stdio_input_loop(ctx: McpContext, mcp_process: subprocess.Popen) -> None:
"""Handle standard input, intercept call and forward requests to mcp_process stdin."""
@@ -377,6 +376,7 @@ def run_stdio_input_loop(ctx: McpContext, mcp_process: subprocess.Popen) -> None
except KeyboardInterrupt:
mcp_process.terminate()
def split_args(args: list[str] = None) -> tuple[list[str], list[str]]:
"""
Splits CLI arguments into two parts:
gateway/routes/mcp_sse.py
+110 -19
View File
@@ -12,20 +12,22 @@ from fastapi.responses import StreamingResponse
from gateway.common.constants import (
CLIENT_TIMEOUT,
INVARIANT_GUARDRAILS_BLOCKED_MESSAGE,
MCP_METHOD,
MCP_TOOL_CALL,
MCP_LIST_TOOLS,
MCP_PARAMS,
MCP_RESULT,
MCP_SERVER_INFO,
MCP_CLIENT_INFO,
)
from gateway.common.guardrails import GuardrailAction
from gateway.common.mcp_sessions_manager import (
McpSessionsManager,
SseHeaderAttributes,
)
from gateway.integrations.explorer import create_annotations_from_guardrails_errors
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"
MCP_SERVER_POST_HEADERS = {
"connection",
"accept",
@@ -85,7 +87,22 @@ async def mcp_post_gateway(
)
if request_json.get(MCP_METHOD) == MCP_TOOL_CALL:
_hook_tool_call(session_id=session_id, request_json=request_json)
# Intercept and potentially block the request
hook_tool_call_result, is_blocked = await _hook_tool_call(
session_id=session_id, request_json=request_json
)
if is_blocked:
# If blocked, hook_tool_call_result contains the block message.
# Forward the block message result back to the caller.
# The original request is not passed to the MCP process.
return Response(
content=json.dumps(hook_tool_call_result),
status_code=403,
headers={
"X-Proxied-By": "mcp-gateway",
"Content-Type": "application/json",
},
)
async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client:
try:
@@ -168,7 +185,7 @@ async def mcp_get_sse_gateway(
(
event_bytes,
session_id,
) = _handle_endpoint_event(
) = await _handle_endpoint_event(
sse,
sse_header_attributes=SseHeaderAttributes.from_request_headers(
request.headers
@@ -176,7 +193,7 @@ async def mcp_get_sse_gateway(
)
case "message":
if session_id:
event_bytes = _handle_message_event(
event_bytes = await _handle_message_event(
session_id=session_id, sse=sse
)
yield event_bytes
@@ -196,7 +213,7 @@ async def mcp_get_sse_gateway(
)
def _hook_tool_call(session_id: str, request_json: dict) -> None:
async def _hook_tool_call(session_id: str, request_json: dict) -> Tuple[dict, bool]:
"""
Hook to process the request JSON before sending it to the MCP server.
@@ -213,17 +230,53 @@ def _hook_tool_call(session_id: str, request_json: dict) -> None:
},
}
message = {"role": "assistant", "content": "", "tool_calls": [tool_call]}
# Check for blocking guardrails - this blocks until completion
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 (
guardrails_result
and guardrails_result.get("errors", [])
and _check_if_new_errors(session_id, guardrails_result)
):
# Add the trace to the explorer
asyncio.create_task(
session_store.add_message_to_session(
session_id=session_id,
message=message,
guardrails_result=guardrails_result,
)
)
return {
"jsonrpc": "2.0",
"id": request_json.get("id"),
"error": {
"code": -32600,
"message": INVARIANT_GUARDRAILS_BLOCKED_MESSAGE
% guardrails_result["errors"],
},
}, True
# Push trace to the explorer - don't block on its response
asyncio.create_task(session_store.add_message_to_session(session_id, message))
asyncio.create_task(
session_store.add_message_to_session(session_id, message, guardrails_result)
)
return request_json, False
def _hook_tool_call_response(session_id: str, response_json: dict) -> None:
async def _hook_tool_call_response(session_id: str, response_json: dict) -> dict:
"""
Hook to process the response JSON after receiving it from the MCP server.
Args:
session_id (str): The session ID associated with the request.
response_json (dict): The response JSON to be processed.
Returns:
dict: The response JSON is returned if no guardrail is violated
else an error dict is returned.
"""
message = {
"role": "tool",
@@ -231,8 +284,28 @@ def _hook_tool_call_response(session_id: str, response_json: dict) -> None:
"content": response_json.get(MCP_RESULT).get("content"),
"error": response_json.get(MCP_RESULT).get("error"),
}
result = response_json
session = session_store.get_session(session_id)
guardrailing_result = await session.get_guardrails_check_result(
message, action=GuardrailAction.BLOCK
)
if guardrailing_result and guardrailing_result.get("errors", []):
# If the request is blocked, return a message indicating the block reason.
result = {
"jsonrpc": "2.0",
"id": response_json.get("id"),
"error": {
"code": -32600,
"message": INVARIANT_GUARDRAILS_BLOCKED_MESSAGE
% guardrailing_result["errors"],
},
}
# Push trace to the explorer - don't block on its response
asyncio.create_task(session_store.add_message_to_session(session_id, message))
asyncio.create_task(
session_store.add_message_to_session(session_id, message, guardrailing_result)
)
return result
def _convert_localhost_to_docker_host(mcp_server_base_url: str) -> str:
@@ -257,7 +330,7 @@ def _convert_localhost_to_docker_host(mcp_server_base_url: str) -> str:
return mcp_server_base_url
def _handle_endpoint_event(
async def _handle_endpoint_event(
sse: ServerSentEvent, sse_header_attributes: SseHeaderAttributes
) -> Tuple[bytes, str]:
"""
@@ -278,7 +351,7 @@ def _handle_endpoint_event(
session_id = match.group(1)
# Initialize this session in our store if needed
if not session_store.session_exists(session_id):
session_store.initialize_session(session_id, sse_header_attributes)
await session_store.initialize_session(session_id, sse_header_attributes)
# Rewrite the endpoint to use our gateway
modified_data = sse.data.replace(
@@ -289,7 +362,7 @@ def _handle_endpoint_event(
return event_bytes, session_id
def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes:
async def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes:
"""
Handle the message event type.
@@ -311,10 +384,16 @@ def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes:
method = session.id_to_method_mapping.get(response_json.get("id"))
if method == MCP_TOOL_CALL:
_hook_tool_call_response(
hook_tool_call_response = await _hook_tool_call_response(
session_id=session_id,
response_json=response_json,
)
# Update the event bytes with hook_tool_call_response.
# hook_tool_call_response is same as response_json if no guardrail is violated.
# If guardrail is violated, it contains the error message.
event_bytes = f"event: {sse.event}\ndata: {json.dumps(hook_tool_call_response)}\n\n".encode(
"utf-8"
)
elif method == MCP_LIST_TOOLS:
session_store.get_session(session_id).metadata["tools"] = response_json.get(
MCP_RESULT
@@ -330,3 +409,15 @@ def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes:
flush=True,
)
return event_bytes
def _check_if_new_errors(session_id: str, 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