Add integration tests for MCP SSE via gateway with guardrails.

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

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"

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
 ):

tests/integration/requirements.txt

@@ -6,5 +6,6 @@ openai
 pillow
 pytest
 pytest-asyncio
+pytest-timeout
 tavily-python
 uv

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()

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"]

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)

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()

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