From 8eae198eb0395c1d4864a78821e85566164a5e40 Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 14 May 2025 10:31:03 +0200 Subject: [PATCH] Add integration tests for MCP SSE via gateway with guardrails. --- tests/integration/docker-compose.test.yml | 16 ++ tests/integration/mcp/test_mcp_sse.py | 261 ++++++++++++++++++ tests/integration/mcp/test_mcp_stdio.py | 3 + tests/integration/requirements.txt | 1 + .../integration/resources/mcp/sse/__init__.py | 0 .../resources/mcp/sse/client/__init__.py | 0 .../resources/mcp/sse/client/main.py | 106 +++++++ .../messenger_server/Dockerfile.mcp-server | 11 + .../mcp/sse/messenger_server/__init__.py | 0 .../mcp/sse/messenger_server/main.py | 84 ++++++ .../resources/mcp/stdio/client/main.py | 8 +- .../mcp/stdio/messenger_server/main.py | 1 - 12 files changed, 486 insertions(+), 5 deletions(-) create mode 100644 tests/integration/mcp/test_mcp_sse.py create mode 100644 tests/integration/resources/mcp/sse/__init__.py create mode 100644 tests/integration/resources/mcp/sse/client/__init__.py create mode 100644 tests/integration/resources/mcp/sse/client/main.py create mode 100644 tests/integration/resources/mcp/sse/messenger_server/Dockerfile.mcp-server create mode 100644 tests/integration/resources/mcp/sse/messenger_server/__init__.py create mode 100644 tests/integration/resources/mcp/sse/messenger_server/main.py diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 3695cc0..777c802 100644 --- a/tests/integration/docker-compose.test.yml +++ b/tests/integration/docker-compose.test.yml @@ -106,6 +106,22 @@ services: timeout: 5s retries: 5 + mcp-messenger-sse-server: + # MCP SSE server used in integration tests + build: + context: ${GATEWAY_ROOT_PATH} + dockerfile: ${GATEWAY_ROOT_PATH}/tests/integration/resources/mcp/sse/messenger_server/Dockerfile.mcp-server + container_name: invariant-gateway-test-mcp-server + networks: + - invariant-gateway-web-test + ports: + - "8123:8123" + healthcheck: + test: ["CMD", "curl", "-f", "http://localhost:8123/sse"] + interval: 3s + timeout: 5s + retries: 5 + networks: invariant-gateway-web-test: external: true diff --git a/tests/integration/mcp/test_mcp_sse.py b/tests/integration/mcp/test_mcp_sse.py new file mode 100644 index 0000000..9c17321 --- /dev/null +++ b/tests/integration/mcp/test_mcp_sse.py @@ -0,0 +1,261 @@ +"""Test MCP gateway via SSE.""" + +import os +import uuid + +from resources.mcp.sse.client.main import run as mcp_client_run +from utils import create_dataset, add_guardrail_to_dataset + +import pytest +import requests + +from mcp.shared.exceptions import McpError + +MCP_SERVER_HOST = "mcp-messenger-sse-server" +MCP_SERVER_PORT = 8123 + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +@pytest.mark.parametrize("push_to_explorer", [False, True]) +async def test_mcp_sse_with_gateway(explorer_api_url, gateway_url, push_to_explorer): + """Test MCP gateway via sse 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( + gateway_url + "/api/v1/gateway/mcp/sse", + f"http://{MCP_SERVER_HOST}:{MCP_SERVER_PORT}", + project_name, + push_to_explorer=push_to_explorer, + 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" + ) + + 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" + ) + assert trace["messages"][0]["role"] == "assistant" + assert trace["messages"][0]["tool_calls"][0]["function"] == { + "name": "get_last_message_from_user", + "arguments": {"username": "Alice"}, + } + assert trace["messages"][1]["role"] == "tool" + assert trace["messages"][1]["content"] == [ + {"type": "text", "text": "What is your favorite food?\n"} + ] + + +@pytest.mark.asyncio +@pytest.mark.timeout(15) +async def test_mcp_sse_with_gateway_and_logging_guardrails( + explorer_api_url, gateway_url +): + """Test MCP gateway via sse 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( + gateway_url + "/api/v1/gateway/mcp/sse", + f"http://{MCP_SERVER_HOST}:{MCP_SERVER_PORT}", + project_name, + 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"][0]["role"] == "assistant" + assert trace["messages"][0]["tool_calls"][0]["function"] == { + "name": "get_last_message_from_user", + "arguments": {"username": "Alice"}, + } + assert trace["messages"][1]["role"] == "tool" + assert trace["messages"][1]["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.1.content.0.text:22-26" + ): + food_annotation = annotation + elif ( + annotation["content"] == "get_last_message_from_user is called" + and annotation["address"] == "messages.0.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 +@pytest.mark.timeout(15) +async def test_mcp_sse_with_gateway_and_blocking_guardrails( + explorer_api_url, gateway_url +): + """Test MCP gateway via sse 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( + gateway_url + "/api/v1/gateway/mcp/sse", + f"http://{MCP_SERVER_HOST}:{MCP_SERVER_PORT}", + project_name, + push_to_explorer=True, + tool_name="get_last_message_from_user", + tool_args={"username": "Alice"}, + ) + # If we get here, the tool call was not blocked + pytest.fail("Expected McpError to be raised") + # 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"][0]["role"] == "assistant" + assert trace["messages"][0]["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.0.tool_calls.0" + ) + assert annotations[0]["extra_metadata"]["source"] == "guardrails-error" diff --git a/tests/integration/mcp/test_mcp_stdio.py b/tests/integration/mcp/test_mcp_stdio.py index 8946cad..f252c4e 100644 --- a/tests/integration/mcp/test_mcp_stdio.py +++ b/tests/integration/mcp/test_mcp_stdio.py @@ -13,6 +13,7 @@ from resources.mcp.stdio.client.main import run as mcp_client_run @pytest.mark.asyncio +@pytest.mark.timeout(15) @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 @@ -70,6 +71,7 @@ async def test_mcp_stdio_with_gateway( @pytest.mark.asyncio +@pytest.mark.timeout(15) async def test_mcp_stdio_with_gateway_and_logging_guardrails( explorer_api_url, invariant_gateway_package_whl_file ): @@ -174,6 +176,7 @@ async def test_mcp_stdio_with_gateway_and_logging_guardrails( @pytest.mark.asyncio +@pytest.mark.timeout(15) async def test_mcp_stdio_with_gateway_and_blocking_guardrails( explorer_api_url, invariant_gateway_package_whl_file ): diff --git a/tests/integration/requirements.txt b/tests/integration/requirements.txt index 248001a..8c878b8 100644 --- a/tests/integration/requirements.txt +++ b/tests/integration/requirements.txt @@ -6,5 +6,6 @@ openai pillow pytest pytest-asyncio +pytest-timeout tavily-python uv \ No newline at end of file diff --git a/tests/integration/resources/mcp/sse/__init__.py b/tests/integration/resources/mcp/sse/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/resources/mcp/sse/client/__init__.py b/tests/integration/resources/mcp/sse/client/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/resources/mcp/sse/client/main.py b/tests/integration/resources/mcp/sse/client/main.py new file mode 100644 index 0000000..9139c14 --- /dev/null +++ b/tests/integration/resources/mcp/sse/client/main.py @@ -0,0 +1,106 @@ +"""This is a simple example of how to use the MCP client with SSE transport.""" + +# pylint: disable=E1101 +# pylint: disable=W0201 + +import asyncio +from datetime import timedelta + +from typing import Any, Optional +from contextlib import AsyncExitStack + +from mcp import ClientSession +from mcp.client.sse import sse_client + + +class MCPClient: + """MCP Client for interacting with a MCP SSE server and processing queries""" + + def __init__(self): + # Initialize session and client objects + self.session: Optional[ClientSession] = None + self.exit_stack = AsyncExitStack() + self._streams_context = None # Initialize these to None + self._session_context = None # so they always exist + + async def connect_to_sse_server( + self, server_url: str, headers: Optional[dict] = None + ): + """ + Connect to an MCP server running with SSE transport + + Args: + server_url: URL of the MCP server + headers: Optional headers to include in the request + """ + # Store the context managers so they stay alive + self._streams_context = sse_client( + url=server_url, + timeout=5, + headers=headers or {}, + sse_read_timeout=10, + ) + streams = await self._streams_context.__aenter__() + + self._session_context = ClientSession(*streams) + # pylint: disable=C2801 + self.session: ClientSession = await self._session_context.__aenter__() + + # Initialize + await self.session.initialize() + + async def cleanup(self): + """Clean up the session and streams""" + # Check if the session context exists before trying to exit it + if hasattr(self, "_session_context") and self._session_context is not None: + await self._session_context.__aexit__(None, None, None) + + # Check if the streams context exists before trying to exit it + if hasattr(self, "_streams_context") and self._streams_context is not None: + await self._streams_context.__aexit__(None, None, None) + + async def process_query(self, tool_name: str, tool_args: dict) -> str: + """Process a query using MCP server""" + result = await self.session.call_tool( + tool_name, tool_args, read_timeout_seconds=timedelta(seconds=10) + ) + return result + + +async def run( + gateway_url: str, + mcp_server_base_url: str, + project_name: str, + push_to_explorer: bool, + tool_name: str, + tool_args: dict[str, Any], +): + """ + Run the MCP client with the given parameters. + + Args: + gateway_url: URL of the Invariant Gateway + mcp_server_base_url: Base URL of the MCP server + project_name: Name of the project in Invariant Explorer + push_to_explorer: Whether to push traces to the Invariant Explorer + tool_name: Name of the tool to call + tool_args: Arguments for the tool call + + """ + client = MCPClient() + try: + await client.connect_to_sse_server( + server_url=gateway_url, + headers={ + "MCP-SERVER-BASE-URL": mcp_server_base_url, + "INVARIANT-PROJECT-NAME": project_name, + "PUSH-INVARIANT-EXPLORER": str(push_to_explorer), + }, + ) + return await client.process_query(tool_name, tool_args) + finally: + # Sleep for a while to allow the server to process the background tasks + # like pushing traces to the explorer + if push_to_explorer: + await asyncio.sleep(2) + await client.cleanup() diff --git a/tests/integration/resources/mcp/sse/messenger_server/Dockerfile.mcp-server b/tests/integration/resources/mcp/sse/messenger_server/Dockerfile.mcp-server new file mode 100644 index 0000000..df21625 --- /dev/null +++ b/tests/integration/resources/mcp/sse/messenger_server/Dockerfile.mcp-server @@ -0,0 +1,11 @@ +FROM python:3.12-slim + +WORKDIR /app + +# Copy the messenger server code +COPY tests/integration/resources/mcp/sse/messenger_server /app/messenger_server + +# Install dependencies +RUN pip install --no-cache-dir "uvicorn[standard]" "httpx" "mcp[cli]" "starlette" + +CMD ["python", "messenger_server/main.py", "--host", "0.0.0.0", "--port", "8123"] \ No newline at end of file diff --git a/tests/integration/resources/mcp/sse/messenger_server/__init__.py b/tests/integration/resources/mcp/sse/messenger_server/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/integration/resources/mcp/sse/messenger_server/main.py b/tests/integration/resources/mcp/sse/messenger_server/main.py new file mode 100644 index 0000000..39658a2 --- /dev/null +++ b/tests/integration/resources/mcp/sse/messenger_server/main.py @@ -0,0 +1,84 @@ +"""This is a messenger server implementation that returns a few messages based on the username.""" + +import argparse +import hashlib + +import uvicorn + +from mcp.server.fastmcp import FastMCP +from mcp.server import Server +from mcp.server.sse import SseServerTransport +from starlette.applications import Starlette +from starlette.requests import Request +from starlette.routing import Mount, Route + +# Initialize FastMCP server +mcp = FastMCP("messenger_server") + + +MESSAGES = [ + "What about you?", + "What are you doing?", + "What is your name?", + "What is your favorite color?", + "What is your favorite food?", + "What is your favorite movie?", + "What is your favorite book?", +] + + +def _deterministic_index_from_username(username: str, limit: int) -> int: + """Deterministically calculate the index of messages to return based on the username.""" + hash_val = int(hashlib.sha256(username.encode()).hexdigest(), 16) + return hash_val % limit + 1 + + +@mcp.tool() +async def get_last_message_from_user(username: str) -> str: + """Get the last message sent by the username.""" + return MESSAGES[_deterministic_index_from_username(username, len(MESSAGES))] + "\n" + + +@mcp.tool() +async def send_message(username: str, message: str) -> str: + """Send a message to the username.""" + return f"Message '{message}' sent to {username}." + + +def create_starlette_app(server: Server, *, debug: bool = False) -> Starlette: + """Create a Starlette application that can server the provied mcp server with SSE.""" + sse = SseServerTransport("/messages/") + + async def handle_sse(request: Request) -> None: + async with sse.connect_sse( + request.scope, + request.receive, + request._send, # pylint: disable=W0212 + ) as (read_stream, write_stream): + await server.run( + read_stream, + write_stream, + server.create_initialization_options(), + ) + + return Starlette( + debug=debug, + routes=[ + Route("/sse", endpoint=handle_sse), + Mount("/messages/", app=sse.handle_post_message), + ], + ) + + +if __name__ == "__main__": + mcp_server = mcp._mcp_server # pylint: disable=W0212 + + parser = argparse.ArgumentParser(description="Run MCP SSE-based server") + parser.add_argument("--host", help="Host to bind to", required=True) + parser.add_argument("--port", help="Port to listen on", required=True, type=int) + args = parser.parse_args() + + # Bind SSE request handling to MCP server + starlette_app = create_starlette_app(mcp_server, debug=True) + + uvicorn.run(starlette_app, host=args.host, port=args.port) diff --git a/tests/integration/resources/mcp/stdio/client/main.py b/tests/integration/resources/mcp/stdio/client/main.py index 9feb2bb..020db99 100644 --- a/tests/integration/resources/mcp/stdio/client/main.py +++ b/tests/integration/resources/mcp/stdio/client/main.py @@ -1,7 +1,7 @@ """A MCP client implementation that interacts with MCP server to make tool calls.""" +import asyncio import os -import time from datetime import timedelta from contextlib import AsyncExitStack from typing import Any, Optional @@ -11,7 +11,7 @@ from mcp.client.stdio import stdio_client class MCPClient: - """MCP Client for interacting with a MCP server and processing queries""" + """MCP Client for interacting with a MCP stdio server and processing queries""" def __init__(self): self.session: Optional[ClientSession] = None @@ -69,7 +69,7 @@ class MCPClient: self.stdio, self.write = stdio_transport self.session = await self.exit_stack.enter_async_context( ClientSession( - self.stdio, self.write, read_timeout_seconds=timedelta(minutes=0.5) + self.stdio, self.write, read_timeout_seconds=timedelta(seconds=10) ) ) @@ -133,5 +133,5 @@ async def run( # Sleep for a while to allow the server to process the background tasks # like pushing traces to the explorer if push_to_explorer: - time.sleep(2) + await asyncio.sleep(2) await client.cleanup() diff --git a/tests/integration/resources/mcp/stdio/messenger_server/main.py b/tests/integration/resources/mcp/stdio/messenger_server/main.py index 823a989..7f29a8a 100644 --- a/tests/integration/resources/mcp/stdio/messenger_server/main.py +++ b/tests/integration/resources/mcp/stdio/messenger_server/main.py @@ -1,6 +1,5 @@ """This is a messenger server implementation that returns a few messages based on the username.""" -import random import hashlib from mcp.server.fastmcp import FastMCP