CalvinBackup/invariant-gateway
1
0
Fork 0
mirror of https://github.com/invariantlabs-ai/invariant-gateway.git synced 2026-07-06 10:57:50 +02:00
Code Issues Packages Projects Releases Wiki Activity

port tools/list message support to SSE

Browse Source
This commit is contained in:
Luca Beurer-Kellner
2025-05-16 19:49:30 +02:00
parent f05731b55f
commit ac68e4a8f7
7 changed files with 48 additions and 264 deletions
Download Patch File Download Diff File Expand all files Collapse all files
gateway/mcp/mcp.py
+3 -1
View File
@@ -306,6 +306,8 @@ async def hook_tool_call(ctx: McpContext, request: dict) -> tuple[dict, bool]:
):
if ctx.push_explorer:
await append_and_push_trace(ctx, message, guardrailing_result)
else:
ctx.trace.append(message)
return json_rpc_error_response(
request.get("id"),
@@ -364,7 +366,7 @@ async def hook_tool_result(ctx: McpContext, result: dict) -> dict:
message = {
"role": "tool",
"content": json.dumps(result.get("result").get("tools")),
"tool_call_id": "call_" + str(result.get("id")),
"tool_call_id": call_id,
}
# next validate it with guardrails
guardrailing_result = await get_guardrails_check_result(
gateway/routes/mcp_sse.py
+37
View File
@@ -3,6 +3,7 @@
import asyncio
import json
import re
import os
from typing import Tuple
import httpx
@@ -96,6 +97,23 @@ async def mcp_post_gateway(
# The error message is sent back to the client using the SSE stream.
await session.add_pending_error_message(hook_tool_call_result)
return Response(content="Accepted", status_code=202)
elif request_json.get(MCP_METHOD) == MCP_LIST_TOOLS:
# Intercept and potentially block the request
hook_tool_call_result, is_blocked = await _hook_tool_call(
session_id=session_id, request_json={
"id": request_json.get("id"),
"method": MCP_LIST_TOOLS,
"params": {
"name": request_json.get(MCP_METHOD),
"arguments": {}
},
}
)
if is_blocked:
# Add the error message to the session.
# The error message is sent back to the client using the SSE stream.
await session.add_pending_error_message(hook_tool_call_result)
return Response(content="Accepted", status_code=202)
async with httpx.AsyncClient(timeout=CLIENT_TIMEOUT) as client:
try:
@@ -392,6 +410,10 @@ def _convert_localhost_to_docker_host(mcp_server_base_url: str) -> str:
Returns:
str: Modified server address with localhost references changed to host.docker.internal
"""
# check if we are running in a docker container
if not os.environ.get("DOCKER_ENV"):
return mcp_server_base_url
if "localhost" in mcp_server_base_url or "127.0.0.1" in mcp_server_base_url:
# Replace localhost or 127.0.0.1 with host.docker.internal
modified_address = re.sub(
@@ -470,15 +492,30 @@ async def _handle_message_event(session_id: str, sse: ServerSentEvent) -> bytes:
"utf-8"
)
elif method == MCP_LIST_TOOLS:
# store tools in metadata
session_store.get_session(session_id).metadata["tools"] = response_json.get(
MCP_RESULT
).get("tools")
# store tools/list tool call in trace
hook_tool_call_response = await _hook_tool_call_response(
session_id=session_id,
response_json={
"id": response_json.get("id"),
"result": {
"content": json.dumps(response_json.get(MCP_RESULT).get("tools"))
}
}
)
except json.JSONDecodeError as e:
print(
f"[MCP SSE] Error parsing message JSON: {e}",
flush=True,
)
except Exception as e: # pylint: disable=broad-except
if os.environ.get("DEBUG") == "true":
import traceback
traceback.print_exc()
print(
f"[MCP SSE] Error processing message: {e}",
flush=True,
run.sh
+1 -3
View File
@@ -145,9 +145,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
uv build
WHEEL_FILE=$(ls dist/*.whl | head -n 1)
echo "WHEEL_FILE: $WHEEL_FILE"
tests/integration/docker-compose.test.yml
+1
View File
@@ -122,6 +122,7 @@ services:
timeout: 5s
retries: 5
networks:
invariant-gateway-web-test:
external: true
tests/integration/mcp/test_mcp_stdio.py
-260
View File
@@ -1,260 +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.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"},
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"
)
# ensure custom keys are present
assert metadata["my-custom-key"] == "value1"
assert metadata["my-custom-key-2"] == "value2"
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
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"][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"}
]
# 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.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
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"][2]["role"] == "assistant"
assert trace["messages"][2]["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.2.tool_calls.0"
)
assert annotations[0]["extra_metadata"]["source"] == "guardrails-error"
tests/integration/resources/mcp/sse/client/main.py
+3
View File
@@ -97,6 +97,9 @@ async def run(
"PUSH-INVARIANT-EXPLORER": str(push_to_explorer),
},
)
# list tools
await client.session.list_tools()
# call tool
return await client.process_query(tool_name, tool_args)
finally:
# Sleep for a while to allow the server to process the background tasks
tests/integration/resources/mcp/stdio/client/main.py
+3
View File
@@ -81,7 +81,10 @@ class MCPClient:
)
)
# initialize the session
await self.session.initialize()
# list tools
await self.session.list_tools()
async def call_tool(
self, tool_name: str, tool_args: dict[str, Any]