diff --git a/gateway/common/constants.py b/gateway/common/constants.py index 02a8209..c8c87ff 100644 --- a/gateway/common/constants.py +++ b/gateway/common/constants.py @@ -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 + """ diff --git a/gateway/common/mcp_sessions_manager.py b/gateway/common/mcp_sessions_manager.py index e9c77be..663802a 100644 --- a/gateway/common/mcp_sessions_manager.py +++ b/gateway/common/mcp_sessions_manager.py @@ -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) diff --git a/gateway/mcp/mcp.py b/gateway/mcp/mcp.py index 2132155..b7ff75f 100644 --- a/gateway/mcp/mcp.py +++ b/gateway/mcp/mcp.py @@ -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: diff --git a/gateway/routes/mcp_sse.py b/gateway/routes/mcp_sse.py index 42cb70c..33d9415 100644 --- a/gateway/routes/mcp_sse.py +++ b/gateway/routes/mcp_sse.py @@ -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