From 73de68e822d116f7d53c07d7baedd1fe5376fcf9 Mon Sep 17 00:00:00 2001 From: Hemang Sarkar Date: Mon, 12 May 2025 15:36:33 +0200 Subject: [PATCH 1/9] Update tests_ci.yml with a timeout at the job level. --- .github/workflows/tests_ci.yml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/workflows/tests_ci.yml b/.github/workflows/tests_ci.yml index 5291aa0..8641eb7 100644 --- a/.github/workflows/tests_ci.yml +++ b/.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 \ No newline at end of file + fi From e2e004b7b1f1048bad52444ea3e9936bc4716668 Mon Sep 17 00:00:00 2001 From: Hemang Date: Tue, 13 May 2025 09:41:09 +0200 Subject: [PATCH 2/9] Move dockerfiles inside gateway/ and update main CLI script to be able to run build, up, down and logs on a local gateway server instance. --- .env | 6 - .github/workflows/publish-images.yml | 2 +- MANIFEST.in | 2 + gateway/.env | 9 +- .../Dockerfile.gateway | 6 +- gateway/__main__.py | 209 +++++++++++++++++- .../docker-compose.local.yml | 6 +- gateway/integrations/explorer.py | 4 +- gateway/integrations/guardrails.py | 1 + run.sh | 15 +- tests/integration/docker-compose.test.yml | 2 +- .../resources/mcp/stdio/client/main.py | 9 +- 12 files changed, 241 insertions(+), 30 deletions(-) delete mode 100644 .env create mode 100644 MANIFEST.in rename Dockerfile.gateway => gateway/Dockerfile.gateway (65%) rename docker-compose.local.yml => gateway/docker-compose.local.yml (92%) diff --git a/.env b/.env deleted file mode 100644 index ae81d93..0000000 --- a/.env +++ /dev/null @@ -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://:8000 to push to the local explorer instance. -INVARIANT_API_URL=https://explorer.invariantlabs.ai -GUARDRAILS_API_URL=https://explorer.invariantlabs.ai \ No newline at end of file diff --git a/.github/workflows/publish-images.yml b/.github/workflows/publish-images.yml index b4c15dd..04e8fb9 100644 --- a/.github/workflows/publish-images.yml +++ b/.github/workflows/publish-images.yml @@ -48,7 +48,7 @@ jobs: uses: docker/build-push-action@v5 with: context: ./gateway - file: Dockerfile.gateway + file: ./gateway/Dockerfile.gateway platforms: linux/amd64 push: true tags: ${{ env.tags }} diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..2e13d6c --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,2 @@ +include gateway/docker-compose.local.yml +include gateway/Dockerfile.gateway diff --git a/gateway/.env b/gateway/.env index c84df8b..093ad91 100644 --- a/gateway/.env +++ b/gateway/.env @@ -1,4 +1,11 @@ POSTGRES_USER=postgres POSTGRES_PASSWORD=postgres POSTGRES_DB=invariantmonitor -POSTGRES_HOST=database \ No newline at end of file +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://:8000 to push to the local explorer instance. +INVARIANT_API_URL=https://explorer.invariantlabs.ai +GUARDRAILS_API_URL=https://explorer.invariantlabs.ai \ No newline at end of file diff --git a/Dockerfile.gateway b/gateway/Dockerfile.gateway similarity index 65% rename from Dockerfile.gateway rename to gateway/Dockerfile.gateway index 24e6cbd..0642c60 100644 --- a/Dockerfile.gateway +++ b/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 . diff --git a/gateway/__main__.py b/gateway/__main__.py index 942e3df..9bc9935 100644 --- a/gateway/__main__.py +++ b/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 diff --git a/docker-compose.local.yml b/gateway/docker-compose.local.yml similarity index 92% rename from docker-compose.local.yml rename to gateway/docker-compose.local.yml index f570640..0f92ca9 100644 --- a/docker-compose.local.yml +++ b/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} diff --git a/gateway/integrations/explorer.py b/gateway/integrations/explorer.py index dfd8bee..01bbb15 100644 --- a/gateway/integrations/explorer.py +++ b/gateway/integrations/explorer.py @@ -164,7 +164,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, @@ -179,7 +179,7 @@ async def fetch_guardrails_from_explorer( # Get the dataset policies. policies_response = await client.get( - f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy" + f"/api/v1/dataset/byuser/{username}/{dataset_name}/policy", timeout=5 ) if policies_response.status_code != 200: if policies_response.status_code == 404: diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index ca25016..31d4319 100644 --- a/gateway/integrations/guardrails.py +++ b/gateway/integrations/guardrails.py @@ -370,6 +370,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: diff --git a/run.sh b/run.sh index 636d634..6fff8c0 100755 --- a/run.sh +++ b/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 } @@ -143,6 +143,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" @@ -151,9 +152,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"\ @@ -185,7 +190,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 diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 42cb340..3695cc0 100644 --- a/tests/integration/docker-compose.test.yml +++ b/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 diff --git a/tests/integration/resources/mcp/stdio/client/main.py b/tests/integration/resources/mcp/stdio/client/main.py index 4238949..9feb2bb 100644 --- a/tests/integration/resources/mcp/stdio/client/main.py +++ b/tests/integration/resources/mcp/stdio/client/main.py @@ -2,7 +2,7 @@ import os import time - +from datetime import timedelta from contextlib import AsyncExitStack from typing import Any, Optional @@ -68,7 +68,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(minutes=0.5) + ) ) await self.session.initialize() @@ -130,5 +132,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: + time.sleep(2) await client.cleanup() From 8eae198eb0395c1d4864a78821e85566164a5e40 Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 14 May 2025 10:31:03 +0200 Subject: [PATCH 3/9] 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 From ed50670bef699c7c63343c25a014e11cc7157f92 Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 14 May 2025 12:07:38 +0200 Subject: [PATCH 4/9] Add MCP gateway tests with hybrid guardrails (both blocking and logging). Also refactor tests so that we can parameterize the transport type - stdio or sse. --- .../mcp/{test_mcp_sse.py => test_mcp.py} | 245 +++++++++++++---- tests/integration/mcp/test_mcp_stdio.py | 258 ------------------ tests/integration/utils.py | 2 +- 3 files changed, 197 insertions(+), 308 deletions(-) rename tests/integration/mcp/{test_mcp_sse.py => test_mcp.py} (50%) delete mode 100644 tests/integration/mcp/test_mcp_stdio.py diff --git a/tests/integration/mcp/test_mcp_sse.py b/tests/integration/mcp/test_mcp.py similarity index 50% rename from tests/integration/mcp/test_mcp_sse.py rename to tests/integration/mcp/test_mcp.py index 9c17321..07394d9 100644 --- a/tests/integration/mcp/test_mcp_sse.py +++ b/tests/integration/mcp/test_mcp.py @@ -1,9 +1,10 @@ -"""Test MCP gateway via SSE.""" +"""Test MCP gateway via SSE and stdio transports.""" import os import uuid -from resources.mcp.sse.client.main import run as mcp_client_run +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 @@ -11,26 +12,50 @@ import requests from mcp.shared.exceptions import McpError -MCP_SERVER_HOST = "mcp-messenger-sse-server" -MCP_SERVER_PORT = 8123 +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", [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""" +@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. - 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"}, - ) + 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"}, + ) assert result.isError is False assert ( @@ -73,10 +98,11 @@ async def test_mcp_sse_with_gateway(explorer_api_url, gateway_url, push_to_explo @pytest.mark.asyncio @pytest.mark.timeout(15) -async def test_mcp_sse_with_gateway_and_logging_guardrails( - explorer_api_url, gateway_url +@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 via sse and verify that logging guardrails work""" + """Test MCP gateway and verify that logging guardrails work""" project_name = "test-mcp-" + str(uuid.uuid4()) dataset_creation_response = await create_dataset( @@ -101,14 +127,24 @@ async def test_mcp_sse_with_gateway_and_logging_guardrails( ) # 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"}, - ) + 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 ( @@ -147,12 +183,8 @@ async def test_mcp_sse_with_gateway_and_logging_guardrails( {"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() + # Validate the annotations + annotations = trace["annotations"] food_annotation = None tool_call_annotation = None @@ -178,10 +210,11 @@ async def test_mcp_sse_with_gateway_and_logging_guardrails( @pytest.mark.asyncio @pytest.mark.timeout(15) -async def test_mcp_sse_with_gateway_and_blocking_guardrails( - explorer_api_url, gateway_url +@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 via sse and verify that blocking guardrails work""" + """Test MCP gateway and verify that blocking guardrails work""" project_name = "test-mcp-" + str(uuid.uuid4()) dataset_creation_response = await create_dataset( @@ -200,14 +233,24 @@ async def test_mcp_sse_with_gateway_and_blocking_guardrails( # 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 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 @@ -247,15 +290,119 @@ async def test_mcp_sse_with_gateway_and_blocking_guardrails( "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() + # 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) +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"][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"} + ] + + # Validate the annotations + annotations = trace["annotations"] + assert len(annotations) == 2 + 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" + assert ( + annotations[1]["content"] == "food in ToolOutput" + and annotations[1]["address"] == "messages.1.content.0.text:22-26" + ) + assert annotations[1]["extra_metadata"]["source"] == "guardrails-error" diff --git a/tests/integration/mcp/test_mcp_stdio.py b/tests/integration/mcp/test_mcp_stdio.py deleted file mode 100644 index f252c4e..0000000 --- a/tests/integration/mcp/test_mcp_stdio.py +++ /dev/null @@ -1,258 +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.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 -): - """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"}, - ) - - 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_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"][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_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"][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/utils.py b/tests/integration/utils.py index 6df9ce9..3c6d3bd 100644 --- a/tests/integration/utils.py +++ b/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 From e32ec74ed2b44d6cda2bfc9559e5f1ed953859f4 Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 14 May 2025 12:26:15 +0200 Subject: [PATCH 5/9] Fix paths to pyproject.toml and README in dockerfile.gateway. Also update context in publish-images.yml --- .github/workflows/publish-images.yml | 2 +- gateway/Dockerfile.gateway | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/.github/workflows/publish-images.yml b/.github/workflows/publish-images.yml index 04e8fb9..ff8e9e4 100644 --- a/.github/workflows/publish-images.yml +++ b/.github/workflows/publish-images.yml @@ -47,7 +47,7 @@ jobs: - name: build gateway image uses: docker/build-push-action@v5 with: - context: ./gateway + context: . file: ./gateway/Dockerfile.gateway platforms: linux/amd64 push: true diff --git a/gateway/Dockerfile.gateway b/gateway/Dockerfile.gateway index 0642c60..9bfd43c 100644 --- a/gateway/Dockerfile.gateway +++ b/gateway/Dockerfile.gateway @@ -2,11 +2,11 @@ FROM python:3.12 WORKDIR /srv/gateway -COPY ../pyproject.toml ./pyproject.toml +COPY pyproject.toml ./pyproject.toml COPY . . -COPY ../README.md ./README.md +COPY README.md ./README.md RUN pip install --no-cache-dir . From fefc22eea0dd4041ac13900e522f5112c3791f0e Mon Sep 17 00:00:00 2001 From: Hemang Sarkar Date: Wed, 14 May 2025 15:32:05 +0200 Subject: [PATCH 6/9] Update test_mcp.py to include the transport parameter in the hybrid guardrails test. --- tests/integration/mcp/test_mcp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/mcp/test_mcp.py b/tests/integration/mcp/test_mcp.py index 07394d9..3925585 100644 --- a/tests/integration/mcp/test_mcp.py +++ b/tests/integration/mcp/test_mcp.py @@ -302,6 +302,7 @@ async def test_mcp_with_gateway_and_blocking_guardrails( @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 ): From a214837b1e955e7432c8b26d61dc0a96de2b9d3f Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 14 May 2025 15:44:52 +0200 Subject: [PATCH 7/9] Add message to session store for MCP sse in post path before returning. --- gateway/routes/mcp_sse.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/gateway/routes/mcp_sse.py b/gateway/routes/mcp_sse.py index 677e342..2411ad7 100644 --- a/gateway/routes/mcp_sse.py +++ b/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 From 0ffde69391b0a73633e6e737e4cf0d4527d5d494 Mon Sep 17 00:00:00 2001 From: Hemang Sarkar Date: Thu, 15 May 2025 10:54:12 +0200 Subject: [PATCH 8/9] Fix test_mcp_sse_with_gateway_hybrid_guardrails so that we don't rely on order of annotations. --- tests/integration/mcp/test_mcp.py | 29 +++++++++++++++++++---------- 1 file changed, 19 insertions(+), 10 deletions(-) diff --git a/tests/integration/mcp/test_mcp.py b/tests/integration/mcp/test_mcp.py index 3925585..d977509 100644 --- a/tests/integration/mcp/test_mcp.py +++ b/tests/integration/mcp/test_mcp.py @@ -395,15 +395,24 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails( ] # 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.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 ( - 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" - assert ( - annotations[1]["content"] == "food in ToolOutput" - and annotations[1]["address"] == "messages.1.content.0.text:22-26" - ) - assert annotations[1]["extra_metadata"]["source"] == "guardrails-error" + 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" From a6c11240765d691a1086ff1d325f5ed850f61dee Mon Sep 17 00:00:00 2001 From: Hemang Sarkar Date: Thu, 15 May 2025 11:07:29 +0200 Subject: [PATCH 9/9] Update test_mcp.py --- tests/integration/mcp/test_mcp.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/integration/mcp/test_mcp.py b/tests/integration/mcp/test_mcp.py index d977509..d326992 100644 --- a/tests/integration/mcp/test_mcp.py +++ b/tests/integration/mcp/test_mcp.py @@ -395,6 +395,7 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails( ] # Validate the annotations + annotations = trace["annotations"] food_annotation = None tool_call_annotation = None