Add an option to add extra metadata that is pushed and passed to Guardrails during an MCP session (#47)

* use select() before readline

* support for setting static metadata for MCP sessions

* nest extra mcp metadata in metadata object

* unify session metadata

* extra metadata tests

* use empty object as parameters, if None

* list_tools as tool call

* offset indices in tests

* test: adjust addresses

* mcp: make error reporting configurable

* line logging

* log version

* verbose logging + loud exception failure

* add server and client name to policy get

* append trace even if not pushing

* port tools/list message support to SSE

* use python -m build

* adjust guardrail failure address

* support for blocking tools/list in SSE

* use error-based failure response format by default

* tools/list test

* don't list_tools in stdio connect

* flaky test: handle second possible result in anthropic streaming case

---------

Co-authored-by: knielsen404 <kristian@invariantlabs.ai>

tests/integration/anthropic/test_anthropic_with_tool_call.py

@@ -248,14 +248,27 @@ async def test_streaming_response_with_tool_call(
     messages = [{"role": "user", "content": query}]
     response = weather_agent.get_streaming_response(messages)
 
-    assert response is not None
-    assert response[0][0].type == "text"
-    assert response[0][1].type == "tool_use"
-    assert response[0][1].name == "get_weather"
-    assert city in response[0][1].input["location"].lower()
+    if len(response) == 2:
+        assert response is not None
+        assert response[0][0].type == "text"
+        assert response[0][1].type == "tool_use"
+        assert response[0][1].name == "get_weather"
+        assert city in response[0][1].input["location"].lower()
 
-    assert response[1][0].type == "text"
-    assert city in response[1][0].text.lower()
+        assert response[1][0].type == "text"
+        assert city in response[1][0].text.lower()
+    elif len(response) == 1:
+        # expected output in this case is something like this:
+        # [[TextBlock(text="I'll help you check the weather in New York using the get_weather function.", type='text', citations=None), ToolUseBlock(id='toolu_019VZsmxuUhShou2EpPBxvpe', input={'location': 'New York, NY', 'unit': 'celsius'}, name='get_weather', type='tool_use')]]
+        
+        assert response is not None
+        assert response[0][0].type == "text"
+        assert response[0][1].type == "tool_use"
+        assert response[0][1].name == "get_weather"
+        assert city in response[0][1].input["location"].lower()
+
+    else:
+        assert False, "Expected response length 2 or 1, but got" + str(response)
 
     if push_to_explorer:
         # Wait for the trace to be saved

tests/integration/docker-compose.test.yml

@@ -122,6 +122,7 @@ services:
       timeout: 5s
       retries: 5
 
+
 networks:
   invariant-gateway-web-test:
     external: true

tests/integration/mcp/test_mcp.py

@@ -45,7 +45,7 @@ async def test_mcp_with_gateway(
             project_name,
             push_to_explorer=push_to_explorer,
             tool_name="get_last_message_from_user",
-            tool_args={"username": "Alice"},
+            tool_args={"username": "Alice"}
         )
     else:
         result = await mcp_stdio_client_run(
@@ -55,6 +55,7 @@ async def test_mcp_with_gateway(
             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
@@ -85,13 +86,13 @@ async def test_mcp_with_gateway(
             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"] == {
+        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"][1]["role"] == "tool"
-        assert trace["messages"][1]["content"] == [
+        assert trace["messages"][3]["role"] == "tool"
+        assert trace["messages"][3]["content"] == [
             {"type": "text", "text": "What is your favorite food?\n"}
         ]
 
@@ -173,13 +174,13 @@ async def test_mcp_with_gateway_and_logging_guardrails(
         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"] == {
+    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"][1]["role"] == "tool"
-    assert trace["messages"][1]["content"] == [
+    assert trace["messages"][3]["role"] == "tool"
+    assert trace["messages"][3]["content"] == [
         {"type": "text", "text": "What is your favorite food?\n"}
     ]
 
@@ -192,12 +193,12 @@ async def test_mcp_with_gateway_and_logging_guardrails(
     for annotation in annotations:
         if (
             annotation["content"] == "food in ToolOutput"
-            and annotation["address"] == "messages.1.content.0.text:22-26"
+            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.0.tool_calls.0"
+            and annotation["address"] == "messages.2.tool_calls.0"
         ):
             tool_call_annotation = annotation
     assert food_annotation is not None, "Missing 'food in ToolOutput' annotation"
@@ -286,8 +287,8 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
         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"] == {
+    assert trace["messages"][2]["role"] == "assistant"
+    assert trace["messages"][2]["tool_calls"][0]["function"] == {
         "name": "get_last_message_from_user",
         "arguments": {"username": "Alice"},
     }
@@ -297,7 +298,7 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
     assert len(annotations) == 1
     assert (
         annotations[0]["content"] == "get_last_message_from_user is called"
-        and annotations[0]["address"] == "messages.0.tool_calls.0"
+        and annotations[0]["address"] == "messages.2.tool_calls.0"
     )
     assert annotations[0]["extra_metadata"]["source"] == "guardrails-error"
     assert annotations[0]["extra_metadata"]["guardrail"]["action"] == "block"
@@ -387,13 +388,13 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails(
         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"] == {
+    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"][1]["role"] == "tool"
-    assert trace["messages"][1]["content"] == [
+    assert trace["messages"][3]["role"] == "tool"
+    assert trace["messages"][3]["content"] == [
         {"type": "text", "text": "What is your favorite food?\n"}
     ]
 
@@ -406,12 +407,12 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails(
     for annotation in annotations:
         if (
             annotation["content"] == "food in ToolOutput"
-            and annotation["address"] == "messages.1.content.0.text:22-26"
+            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.0.tool_calls.0"
+            and annotation["address"] == "messages.2.tool_calls.0"
         ):
             tool_call_annotation = annotation
     assert food_annotation is not None, "Missing 'food in ToolOutput' annotation"
@@ -422,3 +423,55 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails(
     assert food_annotation["extra_metadata"]["guardrail"]["action"] == "block"
     assert tool_call_annotation["extra_metadata"]["source"] == "guardrails-error"
     assert tool_call_annotation["extra_metadata"]["guardrail"]["action"] == "log"
+
+
+
+@pytest.mark.asyncio
+@pytest.mark.timeout(30)
+@pytest.mark.parametrize("transport", ["stdio", "sse"])
+async def test_mcp_tool_list_blocking(
+    explorer_api_url, invariant_gateway_package_whl_file, gateway_url, transport
+):
+    """
+    Tests that blocking guardrails work for the tools/list call.
+    
+    For those, the expected behavior is that the returned tools are all renamed to blocked_... and include an informative block notice, instead of the original tool description.
+    """
+    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_output: ToolOutput)\n   tool_call(tool_output).function.name == "tools/list"',
+        action="block",
+        invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
+    )
+
+    # Run the MCP client and make the tools/list call.
+    if transport == "sse":
+        tools_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="tools/list",
+            tool_args={},
+        )
+    else:
+        tools_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="tools/list",
+            tool_args={},
+        )
+
+    assert "blocked_get_last_message_from_user" in str(tools_result), "Expected the tool names to be renamed and blocked because of the blocking guardrail on the tools/list call. Instead got: " + str(tools_result)
+

tests/integration/resources/mcp/sse/client/main.py

@@ -97,7 +97,13 @@ async def run(
                 "PUSH-INVARIANT-EXPLORER": str(push_to_explorer),
             },
         )
-        return await client.process_query(tool_name, tool_args)
+        # list tools
+        listed_tools = await client.session.list_tools()
+        # call tool
+        if tool_name == "tools/list":
+            return listed_tools
+        else:
+            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

tests/integration/resources/mcp/stdio/client/main.py

@@ -23,6 +23,7 @@ class MCPClient:
         project_name: str,
         server_script_path: str,
         push_to_explorer: bool,
+        metadata_keys: Optional[dict[str, str]] = None,
     ):
         """
         Connect to an MCP server.
@@ -42,6 +43,12 @@ class MCPClient:
             "--project-name",
             project_name,
         ]
+        # add metadata cli args
+        if metadata_keys is not None:
+            for key, value in metadata_keys.items():
+                args.append("--metadata-" + key + "=" + value)
+
+
         if push_to_explorer:
             args.append("--push-explorer")
         args.extend(
@@ -54,6 +61,7 @@ class MCPClient:
                 os.path.basename(server_script_path),
             ]
         )
+
         server_params = StdioServerParameters(
             command="uvx",
             args=args,
@@ -73,6 +81,7 @@ class MCPClient:
             )
         )
 
+        # initialize the session
         await self.session.initialize()
 
     async def call_tool(
@@ -101,6 +110,7 @@ async def run(
     push_to_explorer: bool,
     tool_name: str,
     tool_args: dict[str, Any],
+    metadata_keys: Optional[dict[str, str]] = None,
 ) -> types.CallToolResult:
     """
     Main function to setup the MCP client and server.
@@ -123,7 +133,13 @@ async def run(
             project_name,
             server_script_path,
             push_to_explorer,
+            metadata_keys=metadata_keys
         )
-        return await client.call_tool(tool_name, tool_args)
+        listed_tools = await client.session.list_tools()
+        if tool_name == "tools/list":
+            # list tools
+            return listed_tools
+        else:
+            return await client.call_tool(tool_name, tool_args)
     finally:
         await client.cleanup()