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