Merge branch 'main' into mcp-metadata

tests/integration/gemini/test_generate_content_without_tool_calls.py

@@ -103,6 +103,7 @@ async def test_generate_content(
 
 @pytest.mark.skipif(not os.getenv("GEMINI_API_KEY"), reason="No GEMINI_API_KEY set")
 @pytest.mark.parametrize("push_to_explorer", [True, False])
+@pytest.mark.skip(reason="Skipping this test: 500 error from Gemini API")
 async def test_generate_content_with_image(
     explorer_api_url, gateway_url, push_to_explorer
 ):

tests/integration/mcp/test_mcp.py

@@ -17,7 +17,7 @@ MCP_SSE_SERVER_PORT = 8123
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(15)
+@pytest.mark.timeout(30)
 @pytest.mark.parametrize(
     "push_to_explorer, transport",
     [
@@ -98,7 +98,7 @@ async def test_mcp_with_gateway(
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(15)
+@pytest.mark.timeout(30)
 @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
@@ -206,11 +206,13 @@ async def test_mcp_with_gateway_and_logging_guardrails(
         tool_call_annotation is not None
     ), "Missing 'get_last_message_from_user is called' annotation"
     assert food_annotation["extra_metadata"]["source"] == "guardrails-error"
+    assert food_annotation["extra_metadata"]["guardrail"]["action"] == "log"
     assert tool_call_annotation["extra_metadata"]["source"] == "guardrails-error"
+    assert tool_call_annotation["extra_metadata"]["guardrail"]["action"] == "log"
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(15)
+@pytest.mark.timeout(30)
 @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
@@ -299,10 +301,11 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
         and annotations[0]["address"] == "messages.0.tool_calls.0"
     )
     assert annotations[0]["extra_metadata"]["source"] == "guardrails-error"
+    assert annotations[0]["extra_metadata"]["guardrail"]["action"] == "block"
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(15)
+@pytest.mark.timeout(30)
 @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
@@ -417,4 +420,6 @@ async def test_mcp_sse_with_gateway_hybrid_guardrails(
         tool_call_annotation is not None
     ), "Missing 'get_last_message_from_user is called' annotation"
     assert food_annotation["extra_metadata"]["source"] == "guardrails-error"
+    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"

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
+
 from datetime import timedelta
 from contextlib import AsyncExitStack
 from typing import Any, Optional
@@ -77,7 +77,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(seconds=10)
+                self.stdio, self.write, read_timeout_seconds=timedelta(seconds=15)
             )
         )
 
@@ -93,10 +93,6 @@ class MCPClient:
             tool_name: Name of the tool to call
             tool_args: Arguments for the tool call
         """
-        response = await self.session.list_tools()
-        if tool_name not in [tool.name for tool in response.tools]:
-            raise ValueError(f"Tool '{tool_name}' not found in available tools")
-
         # Execute tool call
         result = await self.session.call_tool(tool_name, tool_args)
         return result
@@ -140,8 +136,4 @@ async def run(
         )
         return await client.call_tool(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()