From ac68e4a8f77fc99a0120bc15d6ed56b7d23c40c7 Mon Sep 17 00:00:00 2001 From: Luca Beurer-Kellner Date: Fri, 16 May 2025 19:49:30 +0200 Subject: [PATCH] port tools/list message support to SSE --- gateway/mcp/mcp.py | 4 +- gateway/routes/mcp_sse.py | 37 +++ run.sh | 4 +- tests/integration/docker-compose.test.yml | 1 + tests/integration/mcp/test_mcp_stdio.py | 260 ------------------ .../resources/mcp/sse/client/main.py | 3 + .../resources/mcp/stdio/client/main.py | 3 + 7 files changed, 48 insertions(+), 264 deletions(-) delete mode 100644 tests/integration/mcp/test_mcp_stdio.py diff --git a/gateway/mcp/mcp.py b/gateway/mcp/mcp.py index 52764c5..ffb4ca7 100644 --- a/gateway/mcp/mcp.py +++ b/gateway/mcp/mcp.py @@ -306,6 +306,8 @@ async def hook_tool_call(ctx: McpContext, request: dict) -> tuple[dict, bool]: ): if ctx.push_explorer: await append_and_push_trace(ctx, message, guardrailing_result) + else: + ctx.trace.append(message) return json_rpc_error_response( request.get("id"), @@ -364,7 +366,7 @@ async def hook_tool_result(ctx: McpContext, result: dict) -> dict: message = { "role": "tool", "content": json.dumps(result.get("result").get("tools")), - "tool_call_id": "call_" + str(result.get("id")), + "tool_call_id": call_id, } # next validate it with guardrails guardrailing_result = await get_guardrails_check_result( diff --git a/gateway/routes/mcp_sse.py b/gateway/routes/mcp_sse.py index 2411ad7..d70b220 100644 --- a/gateway/routes/mcp_sse.py +++ b/gateway/routes/mcp_sse.py @@ -3,6 +3,7 @@ import asyncio import json import re +import os from typing import Tuple import httpx @@ -96,6 +97,23 @@ async def mcp_post_gateway( # The error message is sent back to the client using the SSE stream. await session.add_pending_error_message(hook_tool_call_result) return Response(content="Accepted", status_code=202) + elif request_json.get(MCP_METHOD) == MCP_LIST_TOOLS: + # Intercept and potentially block the request + hook_tool_call_result, is_blocked = await _hook_tool_call( + session_id=session_id, request_json={ + "id": request_json.get("id"), + "method": MCP_LIST_TOOLS, + "params": { + "name": request_json.get(MCP_METHOD), + "arguments": {} + }, + } + ) + if is_blocked: + # Add the error message to the session. + # The error message is sent back to the client using the SSE stream. + await session.add_pending_error_message(hook_tool_call_result) + return Response(content="Accepted", status_code=202) async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client: try: @@ -392,6 +410,10 @@ def _convert_localhost_to_docker_host(mcp_server_base_url: str) -> str: Returns: str: Modified server address with localhost references changed to host.docker.internal """ + # check if we are running in a docker container + if not os.environ.get("DOCKER_ENV"): + return mcp_server_base_url + 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( @@ -470,15 +492,30 @@ async def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes: "utf-8" ) elif method == MCP_LIST_TOOLS: + # store tools in metadata session_store.get_session(session_id).metadata["tools"] = response_json.get( MCP_RESULT ).get("tools") + # store tools/list tool call in trace + hook_tool_call_response = await _hook_tool_call_response( + session_id=session_id, + response_json={ + "id": response_json.get("id"), + "result": { + "content": json.dumps(response_json.get(MCP_RESULT).get("tools")) + } + } + ) + except json.JSONDecodeError as e: print( f"[MCP SSE] Error parsing message JSON: {e}", flush=True, ) except Exception as e: # pylint: disable=broad-except + if os.environ.get("DEBUG") == "true": + import traceback + traceback.print_exc() print( f"[MCP SSE] Error processing message: {e}", flush=True, diff --git a/run.sh b/run.sh index 17d3894..52fbdb3 100755 --- a/run.sh +++ b/run.sh @@ -145,9 +145,7 @@ integration_tests() { # Generate latest whl file for the invariant-gateway package. # This is required to run the integration tests. - pip install build - rm -rf dist - python -m build + uv build WHEEL_FILE=$(ls dist/*.whl | head -n 1) echo "WHEEL_FILE: $WHEEL_FILE" diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 777c802..169e58a 100644 --- a/tests/integration/docker-compose.test.yml +++ b/tests/integration/docker-compose.test.yml @@ -122,6 +122,7 @@ services: timeout: 5s retries: 5 + networks: invariant-gateway-web-test: external: true diff --git a/tests/integration/mcp/test_mcp_stdio.py b/tests/integration/mcp/test_mcp_stdio.py deleted file mode 100644 index 79b2b67..0000000 --- a/tests/integration/mcp/test_mcp_stdio.py +++ /dev/null @@ -1,260 +0,0 @@ -"""Test MCP gateway via stdio.""" - -import os -import uuid - -import pytest -import requests - -from mcp.shared.exceptions import McpError -from utils import create_dataset, add_guardrail_to_dataset - -from resources.mcp.stdio.client.main import run as mcp_client_run - - -@pytest.mark.asyncio -@pytest.mark.parametrize("push_to_explorer", [False, True]) -async def test_mcp_stdio_with_gateway( - explorer_api_url, invariant_gateway_package_whl_file, push_to_explorer -): - """Test MCP gateway via stdio and verify trace is pushed to explorer""" - project_name = "test-mcp-" + str(uuid.uuid4()) - - # Run the MCP client and make the tool call. - result = await mcp_client_run( - invariant_gateway_package_whl_file, - project_name, - server_script_path="resources/mcp/stdio/messenger_server/main.py", - push_to_explorer=push_to_explorer, - tool_name="get_last_message_from_user", - tool_args={"username": "Alice"}, - metadata_keys={"my-custom-key": "value1", "my-custom-key-2": "value2"}, - ) - - assert result.isError is False - assert ( - result.content[0].type == "text" - and result.content[0].text == "What is your favorite food?\n" - ) - - if push_to_explorer: - # Fetch the trace ids for the dataset - traces_response = requests.get( - f"{explorer_api_url}/api/v1/dataset/byuser/developer/{project_name}/traces", - timeout=5, - ) - traces = traces_response.json() - assert len(traces) == 1 - trace_id = traces[0]["id"] - - # Fetch the trace - trace_response = requests.get( - f"{explorer_api_url}/api/v1/trace/{trace_id}", - timeout=5, - ) - trace = trace_response.json() - metadata = trace["extra_metadata"] - assert ( - metadata["source"] == "mcp" - and metadata["mcp_client"] == "mcp" - and metadata["mcp_server"] == "messenger_server" - ) - # ensure custom keys are present - assert metadata["my-custom-key"] == "value1" - assert metadata["my-custom-key-2"] == "value2" - - assert trace["messages"][2]["role"] == "assistant" - assert trace["messages"][2]["tool_calls"][0]["function"] == { - "name": "get_last_message_from_user", - "arguments": {"username": "Alice"}, - } - assert trace["messages"][3]["role"] == "tool" - assert trace["messages"][3]["content"] == [ - {"type": "text", "text": "What is your favorite food?\n"} - ] - - -@pytest.mark.asyncio -async def test_mcp_stdio_with_gateway_and_logging_guardrails( - explorer_api_url, invariant_gateway_package_whl_file -): - """Test MCP gateway via stdio and verify that logging guardrails work""" - project_name = "test-mcp-" + str(uuid.uuid4()) - - dataset_creation_response = await create_dataset( - explorer_api_url, - invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), - dataset_name=project_name, - ) - dataset_id = dataset_creation_response["id"] - _ = await add_guardrail_to_dataset( - explorer_api_url, - dataset_id=dataset_id, - policy='raise "food in ToolOutput" if:\n (tool_output: ToolOutput)\n (chunk: str) in text(tool_output.content)\n "food" in chunk', - action="log", - invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), - ) - _ = await add_guardrail_to_dataset( - explorer_api_url, - dataset_id=dataset_id, - policy='raise "get_last_message_from_user is called" if:\n (tool_call: ToolCall)\n tool_call is tool:get_last_message_from_user', - action="log", - invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), - ) - - # Run the MCP client and make the tool call. - result = await mcp_client_run( - invariant_gateway_package_whl_file, - project_name, - server_script_path="resources/mcp/stdio/messenger_server/main.py", - push_to_explorer=True, - tool_name="get_last_message_from_user", - tool_args={"username": "Alice"}, - ) - - assert result.isError is False - assert ( - result.content[0].type == "text" - and result.content[0].text == "What is your favorite food?\n" - ) - - # Fetch the trace ids for the dataset - traces_response = requests.get( - f"{explorer_api_url}/api/v1/dataset/byuser/developer/{project_name}/traces", - timeout=5, - ) - traces = traces_response.json() - assert len(traces) == 1 - trace_id = traces[0]["id"] - - # Fetch the trace - trace_response = requests.get( - f"{explorer_api_url}/api/v1/trace/{trace_id}", - timeout=5, - ) - trace = trace_response.json() - metadata = trace["extra_metadata"] - assert ( - metadata["source"] == "mcp" - and metadata["mcp_client"] == "mcp" - and metadata["mcp_server"] == "messenger_server" - ) - assert trace["messages"][2]["role"] == "assistant" - assert trace["messages"][2]["tool_calls"][0]["function"] == { - "name": "get_last_message_from_user", - "arguments": {"username": "Alice"}, - } - assert trace["messages"][3]["role"] == "tool" - assert trace["messages"][3]["content"] == [ - {"type": "text", "text": "What is your favorite food?\n"} - ] - - # Fetch annotations - annotations_response = requests.get( - f"{explorer_api_url}/api/v1/trace/{trace_id}/annotations", - timeout=5, - ) - annotations = annotations_response.json() - food_annotation = None - tool_call_annotation = None - - assert len(annotations) == 2 - for annotation in annotations: - if ( - annotation["content"] == "food in ToolOutput" - and annotation["address"] == "messages.3.content.0.text:22-26" - ): - food_annotation = annotation - elif ( - annotation["content"] == "get_last_message_from_user is called" - and annotation["address"] == "messages.2.tool_calls.0" - ): - tool_call_annotation = annotation - assert food_annotation is not None, "Missing 'food in ToolOutput' annotation" - assert ( - tool_call_annotation is not None - ), "Missing 'get_last_message_from_user is called' annotation" - assert food_annotation["extra_metadata"]["source"] == "guardrails-error" - assert tool_call_annotation["extra_metadata"]["source"] == "guardrails-error" - - -@pytest.mark.asyncio -async def test_mcp_stdio_with_gateway_and_blocking_guardrails( - explorer_api_url, invariant_gateway_package_whl_file -): - """Test MCP gateway via stdio and verify that blocking guardrails work""" - project_name = "test-mcp-" + str(uuid.uuid4()) - - dataset_creation_response = await create_dataset( - explorer_api_url, - invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), - dataset_name=project_name, - ) - dataset_id = dataset_creation_response["id"] - _ = await add_guardrail_to_dataset( - explorer_api_url, - dataset_id=dataset_id, - policy='raise "get_last_message_from_user is called" if:\n (tool_call: ToolCall)\n tool_call is tool:get_last_message_from_user', - action="block", - invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"), - ) - - # Run the MCP client and make the tool call. - try: - _ = await mcp_client_run( - invariant_gateway_package_whl_file, - project_name, - server_script_path="resources/mcp/stdio/messenger_server/main.py", - push_to_explorer=True, - tool_name="get_last_message_from_user", - tool_args={"username": "Alice"}, - ) - # The tool call should be blocked by the guardrail - # and an error should be raised. - except McpError as e: - assert ( - "[Invariant Guardrails] The MCP tool call was blocked for security reasons" - in e.error.message - ) - assert "get_last_message_from_user is called" in e.error.message - assert e.error.code == -32600 - - # Fetch the trace ids for the dataset - traces_response = requests.get( - f"{explorer_api_url}/api/v1/dataset/byuser/developer/{project_name}/traces", - timeout=5, - ) - traces = traces_response.json() - assert len(traces) == 1 - trace_id = traces[0]["id"] - - # Fetch the trace - trace_response = requests.get( - f"{explorer_api_url}/api/v1/trace/{trace_id}", - timeout=5, - ) - trace = trace_response.json() - metadata = trace["extra_metadata"] - assert ( - metadata["source"] == "mcp" - and metadata["mcp_client"] == "mcp" - and metadata["mcp_server"] == "messenger_server" - ) - assert trace["messages"][2]["role"] == "assistant" - assert trace["messages"][2]["tool_calls"][0]["function"] == { - "name": "get_last_message_from_user", - "arguments": {"username": "Alice"}, - } - - # Fetch annotations - annotations_response = requests.get( - f"{explorer_api_url}/api/v1/trace/{trace_id}/annotations", - timeout=5, - ) - annotations = annotations_response.json() - assert len(annotations) == 1 - assert ( - annotations[0]["content"] == "get_last_message_from_user is called" - and annotations[0]["address"] == "messages.2.tool_calls.0" - ) - assert annotations[0]["extra_metadata"]["source"] == "guardrails-error" diff --git a/tests/integration/resources/mcp/sse/client/main.py b/tests/integration/resources/mcp/sse/client/main.py index 9139c14..934a94b 100644 --- a/tests/integration/resources/mcp/sse/client/main.py +++ b/tests/integration/resources/mcp/sse/client/main.py @@ -97,6 +97,9 @@ async def run( "PUSH-INVARIANT-EXPLORER": str(push_to_explorer), }, ) + # list tools + await client.session.list_tools() + # call tool return await client.process_query(tool_name, tool_args) finally: # Sleep for a while to allow the server to process the background tasks diff --git a/tests/integration/resources/mcp/stdio/client/main.py b/tests/integration/resources/mcp/stdio/client/main.py index ba710e3..3e5b35c 100644 --- a/tests/integration/resources/mcp/stdio/client/main.py +++ b/tests/integration/resources/mcp/stdio/client/main.py @@ -81,7 +81,10 @@ class MCPClient: ) ) + # initialize the session await self.session.initialize() + # list tools + await self.session.list_tools() async def call_tool( self, tool_name: str, tool_args: dict[str, Any]