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..ff8e9e4 100644 --- a/.github/workflows/publish-images.yml +++ b/.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 }} 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 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..9bfd43c 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 3002c02..4ab4a0c 100644 --- a/gateway/integrations/explorer.py +++ b/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: diff --git a/gateway/integrations/guardrails.py b/gateway/integrations/guardrails.py index 0f95fbe..0fed353 100644 --- a/gateway/integrations/guardrails.py +++ b/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: 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 diff --git a/run.sh b/run.sh index 0a159ee..17d3894 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 } @@ -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 diff --git a/tests/integration/docker-compose.test.yml b/tests/integration/docker-compose.test.yml index 42cb340..777c802 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 @@ -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.py b/tests/integration/mcp/test_mcp.py new file mode 100644 index 0000000..ae5baea --- /dev/null +++ b/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" 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 56a7f65..eba9429 100644 --- a/tests/integration/resources/mcp/stdio/client/main.py +++ b/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() 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 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