Merge branch 'main' into mcp-metadata

.env

@@ -1,6 +0,0 @@
-# This specifies the Invariant Explorer instance where the gateway will push the traces 
-# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview.
-# If you want to push to a local instance of explorer, then specify the app-api docker container name like: 
-# http://<app-api-docker-container-name>:8000 to push to the local explorer instance.
-INVARIANT_API_URL=https://explorer.invariantlabs.ai
-GUARDRAILS_API_URL=https://explorer.invariantlabs.ai

.github/workflows/publish-images.yml

@@ -47,8 +47,8 @@ jobs:
       - name: build gateway image
         uses: docker/build-push-action@v5
         with:
-          context: ./gateway
-          file: Dockerfile.gateway
+          context: .
+          file: ./gateway/Dockerfile.gateway
           platforms: linux/amd64
           push: true
           tags: ${{ env.tags }}

.github/workflows/tests_ci.yml

@@ -13,6 +13,7 @@ jobs:
   test:
     name: Build & Test
     runs-on: ubuntu-latest
+    timeout-minutes: 15
     steps:
       - name: Checkout
         uses: actions/checkout@v4
@@ -57,4 +58,4 @@ jobs:
           if [[ "${{ steps.integration-tests.outcome }}" != "success" ]]; then
             echo "Integration tests failed"
             exit 1
-          fi
+          fi

MANIFEST.in

@@ -0,0 +1,2 @@
+include gateway/docker-compose.local.yml
+include gateway/Dockerfile.gateway

gateway/.env

@@ -1,4 +1,11 @@
 POSTGRES_USER=postgres
 POSTGRES_PASSWORD=postgres
 POSTGRES_DB=invariantmonitor
-POSTGRES_HOST=database
+POSTGRES_HOST=database
+
+# This specifies the Invariant Explorer instance where the gateway will push the traces 
+# Set this to https://preview-explorer.invariantlabs.ai if you want to push to Preview.
+# If you want to push to a local instance of explorer, then specify the app-api docker container name like: 
+# http://<app-api-docker-container-name>:8000 to push to the local explorer instance.
+INVARIANT_API_URL=https://explorer.invariantlabs.ai
+GUARDRAILS_API_URL=https://explorer.invariantlabs.ai

gateway/Dockerfile.gateway

@@ -2,11 +2,11 @@ FROM python:3.12
 
 WORKDIR /srv/gateway
 
-COPY pyproject.toml ./
+COPY pyproject.toml ./pyproject.toml
 
-COPY gateway/ ./
+COPY . .
 
-COPY README.md /srv/gateway/README.md
+COPY README.md ./README.md
 
 RUN pip install --no-cache-dir .
 

gateway/__main__.py

@@ -1,12 +1,20 @@
 """Script is used to run actions using the Invariant Gateway."""
 
 import asyncio
+import os
 import signal
+import subprocess
 import sys
+import time
+
+from typing import Optional
 
 from gateway.mcp import mcp
 
 
+LOCAL_COMPOSE_FILE = "gateway/docker-compose.local.yml"
+
+
 # Handle signals to ensure clean shutdown
 def signal_handler(sig, frame):
     """Handle signals for graceful shutdown."""
@@ -16,8 +24,14 @@ def signal_handler(sig, frame):
 def print_help():
     """Prints the help message."""
     actions = {
-        "mcp": "Runs the Invariant Gateway against MCP (Model Context Protocol) servers with guardrailing and push to Explorer features.",
-        "llm": "Runs the Invariant Gateway against LLM providers with guardrailing and push to Explorer features. Not implemented yet.",
+        "mcp": """
+        Runs the Invariant Gateway against MCP (Model Context Protocol) stdio servers with guardrailing and push to Explorer features.
+        """,
+        "server": """
+        Runs the Invariant Gateway server locally providing guardrailing and push to Explorer features.
+        Should be called with one of the following subcommands: build, up, down, logs.
+        A guardrails file can be passed with the flag: --guardrails-file=/path/to/guardrails/file.
+        """,
         "help": "Shows this help message.",
     }
 
@@ -25,6 +39,181 @@ def print_help():
         print(f"{verb}: {description}")
 
 
+def ensure_network_exists(network_name: str = "invariant-explorer-web") -> bool:
+    """Ensure the Docker network exists."""
+    try:
+        # Check if network exists
+        result = subprocess.run(
+            ["docker", "network", "inspect", network_name],
+            stdout=subprocess.DEVNULL,
+            stderr=subprocess.DEVNULL,
+            check=False,
+        )
+
+        if result.returncode != 0:
+            print(f"Creating Docker network: {network_name}")
+            subprocess.run(
+                ["docker", "network", "create", network_name],
+                check=True,
+            )
+
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"Error creating Docker network: {e}")
+        return False
+
+
+def setup_guardrails(guardrails_file_path: Optional[str] = None) -> bool:
+    """Configure guardrails if specified."""
+    if not guardrails_file_path:
+        return True
+
+    if not os.path.isfile(guardrails_file_path):
+        print(
+            f"Error: Specified guardrails file does not exist: {guardrails_file_path}"
+        )
+        return False
+
+    # Convert to absolute path
+    guardrails_file_path = os.path.realpath(guardrails_file_path)
+    os.environ["GUARDRAILS_FILE_PATH"] = guardrails_file_path
+
+    # Check if INVARIANT_API_KEY is set
+    if not os.environ.get("INVARIANT_API_KEY"):
+        print(
+            "Error: A guardrails file is specified, but INVARIANT_API_KEY env var is not set. "
+            "This is required to validate guardrails."
+        )
+        return False
+
+    return True
+
+
+def build():
+    """Build Docker containers using docker-compose."""
+    try:
+        print(f"Building using docker-compose file: {LOCAL_COMPOSE_FILE}")
+        subprocess.run(
+            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "build"],
+            check=True,
+        )
+        print("Build completed successfully")
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"Error building containers: {e}")
+        return False
+
+
+def up(guardrails_file_path: Optional[str] = None):
+    """Set up the local server for the Invariant Gateway."""
+    # Ensure network exists
+    if not ensure_network_exists():
+        return 1
+
+    # Set up guardrails
+    if not setup_guardrails(guardrails_file_path=guardrails_file_path):
+        return 1
+
+    # Run the server
+    try:
+        # Start containers
+        print(f"Starting containers using docker-compose file: {LOCAL_COMPOSE_FILE}")
+        subprocess.run(
+            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "up", "-d"],
+            check=True,
+        )
+
+        # Wait for containers to start
+        time.sleep(2)
+
+        # Check if gateway container is running
+        result = subprocess.run(
+            ["docker", "ps", "-qf", "name=invariant-gateway"],
+            capture_output=True,
+            text=True,
+            check=True,
+        )
+
+        if not result.stdout.strip():
+            print("The invariant-gateway container failed to start.")
+            logs = subprocess.run(
+                ["docker", "logs", "invariant-gateway"],
+                capture_output=True,
+                text=True,
+                check=False,
+            )
+            print("Last 20 lines of logs:")
+            print("\n".join(logs.stdout.strip().split("\n")[-20:]))
+            return False
+
+        print("Gateway started at http://localhost:8005/api/v1/gateway/")
+        print("See http://localhost:8005/api/v1/gateway/docs for API documentation")
+
+        if guardrails_file_path:
+            print(f"Using Guardrails File: {guardrails_file_path}")
+
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"Error starting containers: {e}")
+        return False
+
+
+def down():
+    """Stop the Docker containers."""
+    try:
+        print(f"Stopping containers using docker-compose file: {LOCAL_COMPOSE_FILE}")
+        subprocess.run(
+            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "down"],
+            check=True,
+        )
+        print("Containers stopped successfully")
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"Error stopping containers: {e}")
+        return False
+
+
+def logs():
+    """Show container logs."""
+    try:
+        subprocess.run(
+            ["docker", "compose", "-f", str(LOCAL_COMPOSE_FILE), "logs", "-f"],
+            check=True,
+        )
+        return True
+    except subprocess.CalledProcessError as e:
+        print(f"Error showing logs: {e}")
+        return False
+    except KeyboardInterrupt:
+        print("\nExiting logs view")
+        return True
+
+
+def run_server_command(command, args=None):
+    """Run a server command."""
+    if args is None:
+        args = []
+
+    if command == "build":
+        return build()
+    elif command == "up":
+        # Parse guardrails file from args
+        guardrails_file = None
+        for arg in args:
+            if arg.startswith("--guardrails-file="):
+                guardrails_file = arg.split("=", 1)[1]
+
+        return up(guardrails_file)
+    elif command == "down":
+        return down()
+    elif command == "logs":
+        return logs()
+    else:
+        print(f"Unknown server command: {command}")
+        print("Available commands: build, up, down, logs")
+        return False
+
+
 def main():
     """Entry point for the Invariant Gateway."""
     signal.signal(signal.SIGINT, signal_handler)
@@ -37,9 +226,19 @@ def main():
     verb = sys.argv[1]
     if verb == "mcp":
         return asyncio.run(mcp.execute(sys.argv[2:]))
-    if verb == "llm":
-        print("[gateway/__main__.py] 'llm' action is not implemented yet.")
-        return 1
+    if verb == "server":
+        if len(sys.argv) < 3:
+            print(
+                "Error: Missing command for server. Should be one of: build, up, down, logs"
+            )
+            print_help()
+            return 1
+
+        command = sys.argv[2]
+        args = sys.argv[3:]
+        if not run_server_command(command, args):
+            return 1
+        return 0
     if verb == "help":
         print_help()
         return 0

gateway/docker-compose.local.yml

@@ -2,8 +2,8 @@ services:
   invariant-gateway:
     container_name: invariant-gateway
     build:
-      context: .
-      dockerfile: Dockerfile.gateway
+      context: ..
+      dockerfile: gateway/Dockerfile.gateway
     working_dir: /srv/gateway
     env_file:
       - .env
@@ -13,7 +13,7 @@ services:
       - ${INVARIANT_API_KEY:+INVARIANT_API_KEY=${INVARIANT_API_KEY}}
     volumes:
       - type: bind
-        source: ./gateway
+        source: ../gateway
         target: /srv/gateway
       - type: bind
         source: ${GUARDRAILS_FILE_PATH:-/dev/null}

gateway/integrations/explorer.py

@@ -167,7 +167,7 @@ async def fetch_guardrails_from_explorer(
     )
 
     # Get the user details.
-    user_info_response = await client.get("/api/v1/user/identity")
+    user_info_response = await client.get("/api/v1/user/identity", timeout=5)
     if user_info_response.status_code == 401:
         raise HTTPException(
             status_code=401,
@@ -183,7 +183,10 @@ async def fetch_guardrails_from_explorer(
     # Get the dataset policies.
     policies_response = await client.get(
         f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy",
-        params={"client_name": client_name, "server_name": server_name},
+        params={
+            **({"client_name": client_name} if client_name else {}),
+            **({"server_name": server_name} if server_name else {}),
+        },
     )
     if policies_response.status_code != 200:
         if policies_response.status_code == 404:

gateway/integrations/guardrails.py

@@ -371,6 +371,7 @@ async def check_guardrails(
                     "Authorization": context.get_guardrailing_authorization(),
                     "Accept": "application/json"
                 },
+                timeout=5,
             )
             if not result.is_success:
                 if result.status_code == 401:

gateway/routes/mcp_sse.py

@@ -333,9 +333,7 @@ async def _hook_tool_call(session_id: str, request_json: dict) -> Tuple[dict, bo
             },
         }, True
     # Push trace to the explorer - don't block on its response
-    asyncio.create_task(
-        session_store.add_message_to_session(session_id, message, guardrails_result)
-    )
+    await session_store.add_message_to_session(session_id, message, guardrails_result)
     return request_json, False
 
 

run.sh

@@ -37,7 +37,7 @@ up() {
   fi
 
   # Start Docker Compose with the correct environment variable
-  docker compose -f docker-compose.local.yml up -d
+  docker compose -f gateway/docker-compose.local.yml up -d
 
   # Get the status of the container
   sleep 2
@@ -58,12 +58,12 @@ up() {
 
 build() {
   # Build local services
-  docker compose -f docker-compose.local.yml build
+  docker compose -f gateway/docker-compose.local.yml build
 }
 
 down() {
   # Bring down local services
-  docker compose -f docker-compose.local.yml down
+  docker compose -f gateway/docker-compose.local.yml down
   GATEWAY_ROOT_PATH=$(pwd) docker compose -f tests/integration/docker-compose.test.yml down
 }
 
@@ -146,6 +146,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
   WHEEL_FILE=$(ls dist/*.whl | head -n 1)
   echo "WHEEL_FILE: $WHEEL_FILE"
@@ -154,9 +155,13 @@ integration_tests() {
 
   docker build -t 'invariant-gateway-tests' -f ./tests/integration/Dockerfile.test ./tests
 
-  docker run \
+  docker rm -f invariant-gateway-integration-tests-container >/dev/null 2>&1 || true
+  mkdir -p /tmp/docker-logs/.invariant
+
+  docker run --name invariant-gateway-integration-tests-container \
     --mount type=bind,source=./tests/integration,target=/tests \
     --mount type=bind,source=$(realpath $WHEEL_FILE),target=/package/$(basename $WHEEL_FILE) \
+    --mount type=bind,source=/tmp/docker-logs/.invariant,target=/root/.invariant \
     --network invariant-gateway-web-test \
     -e OPENAI_API_KEY="$OPENAI_API_KEY" \
     -e ANTHROPIC_API_KEY="$ANTHROPIC_API_KEY"\
@@ -188,7 +193,7 @@ case "$1" in
     down
     ;;
   "logs")
-    docker compose -f docker-compose.local.yml logs -f
+    docker compose -f gateway/docker-compose.local.yml logs -f
     ;;
   "unit-tests")
     shift

tests/integration/docker-compose.test.yml

@@ -25,7 +25,7 @@ services:
     container_name: invariant-gateway-test
     build:
       context: ${GATEWAY_ROOT_PATH}
-      dockerfile: ${GATEWAY_ROOT_PATH}/Dockerfile.gateway
+      dockerfile: ${GATEWAY_ROOT_PATH}/gateway/Dockerfile.gateway
     depends_on:
       app-api:
         condition: service_healthy
@@ -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.py

@@ -0,0 +1,420 @@
+"""Test MCP gateway via SSE and stdio transports."""
+
+import os
+import uuid
+
+from resources.mcp.sse.client.main import run as mcp_sse_client_run
+from resources.mcp.stdio.client.main import run as mcp_stdio_client_run
+from utils import create_dataset, add_guardrail_to_dataset
+
+import pytest
+import requests
+
+from mcp.shared.exceptions import McpError
+
+MCP_SSE_SERVER_HOST = "mcp-messenger-sse-server"
+MCP_SSE_SERVER_PORT = 8123
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(15)
+@pytest.mark.parametrize(
+    "push_to_explorer, transport",
+    [
+        (False, "stdio"),
+        (False, "sse"),
+        (True, "stdio"),
+        (True, "sse"),
+    ],
+)
+async def test_mcp_with_gateway(
+    explorer_api_url,
+    invariant_gateway_package_whl_file,
+    gateway_url,
+    push_to_explorer,
+    transport,
+):
+    """Test MCP gateway and verify trace is pushed to explorer"""
+    project_name = "test-mcp-" + str(uuid.uuid4())
+
+    # Run the MCP client and make the tool call.
+    if transport == "sse":
+        result = await mcp_sse_client_run(
+            gateway_url + "/api/v1/gateway/mcp/sse",
+            f"http://{MCP_SSE_SERVER_HOST}:{MCP_SSE_SERVER_PORT}",
+            project_name,
+            push_to_explorer=push_to_explorer,
+            tool_name="get_last_message_from_user",
+            tool_args={"username": "Alice"}
+        )
+    else:
+        result = await mcp_stdio_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"
+        )
+        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
+@pytest.mark.timeout(15)
+@pytest.mark.parametrize("transport", ["stdio", "sse"])
+async def test_mcp_with_gateway_and_logging_guardrails(
+    explorer_api_url, invariant_gateway_package_whl_file, gateway_url, transport
+):
+    """Test MCP gateway 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.
+    if transport == "sse":
+        result = await mcp_sse_client_run(
+            gateway_url + "/api/v1/gateway/mcp/sse",
+            f"http://{MCP_SSE_SERVER_HOST}:{MCP_SSE_SERVER_PORT}",
+            project_name,
+            push_to_explorer=True,
+            tool_name="get_last_message_from_user",
+            tool_args={"username": "Alice"},
+        )
+    else:
+        result = await mcp_stdio_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"}
+    ]
+
+    # Validate the annotations
+    annotations = trace["annotations"]
+    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
+@pytest.mark.timeout(15)
+@pytest.mark.parametrize("transport", ["stdio", "sse"])
+async def test_mcp_with_gateway_and_blocking_guardrails(
+    explorer_api_url, invariant_gateway_package_whl_file, gateway_url, transport
+):
+    """Test MCP gateway 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:
+        if transport == "sse":
+            _ = await mcp_sse_client_run(
+                gateway_url + "/api/v1/gateway/mcp/sse",
+                f"http://{MCP_SSE_SERVER_HOST}:{MCP_SSE_SERVER_PORT}",
+                project_name,
+                push_to_explorer=True,
+                tool_name="get_last_message_from_user",
+                tool_args={"username": "Alice"},
+            )
+        else:
+            _ = await mcp_stdio_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"},
+            )
+        # 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"][2]["role"] == "assistant"
+    assert trace["messages"][2]["tool_calls"][0]["function"] == {
+        "name": "get_last_message_from_user",
+        "arguments": {"username": "Alice"},
+    }
+
+    # Validate the annotations
+    annotations = trace["annotations"]
+    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"
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(15)
+@pytest.mark.parametrize("transport", ["stdio", "sse"])
+async def test_mcp_sse_with_gateway_hybrid_guardrails(
+    explorer_api_url, invariant_gateway_package_whl_file, gateway_url, transport
+):
+    """Test MCP gateway and verify that logging and blocking guardrails work together"""
+    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="log",
+        invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
+    )
+    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="block",
+        invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
+    )
+
+    # Run the MCP client and make the tool call.
+    try:
+        if transport == "sse":
+            _ = await mcp_sse_client_run(
+                gateway_url + "/api/v1/gateway/mcp/sse",
+                f"http://{MCP_SSE_SERVER_HOST}:{MCP_SSE_SERVER_PORT}",
+                project_name,
+                push_to_explorer=True,
+                tool_name="get_last_message_from_user",
+                tool_args={"username": "Alice"},
+            )
+        else:
+            _ = await mcp_stdio_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"},
+            )
+        # If we get here, the tool call was not blocked
+        pytest.fail("Expected McpError to be raised")
+    # The tool call output 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 "food in ToolOutput" 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"},
+    }
+    assert trace["messages"][3]["role"] == "tool"
+    assert trace["messages"][3]["content"] == [
+        {"type": "text", "text": "What is your favorite food?\n"}
+    ]
+
+    # Validate the annotations
+    annotations = trace["annotations"]
+    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"

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,8 +1,8 @@
 """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
@@ -76,7 +76,9 @@ class MCPClient:
         )
         self.stdio, self.write = stdio_transport
         self.session = await self.exit_stack.enter_async_context(
-            ClientSession(self.stdio, self.write)
+            ClientSession(
+                self.stdio, self.write, read_timeout_seconds=timedelta(seconds=10)
+            )
         )
 
         await self.session.initialize()
@@ -140,5 +142,6 @@ async def run(
     finally:
         # Sleep for a while to allow the server to process the background tasks
         # like pushing traces to the explorer
-        time.sleep(2)
+        if push_to_explorer:
+            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

tests/integration/utils.py

@@ -4,7 +4,7 @@ import os
 import uuid
 from typing import Any, Dict, Literal, Optional
 
-from httpx import AsyncClient, Client
+from httpx import Client
 from openai import OpenAI
 from google import genai
 from anthropic import Anthropic