Clean up MCP tests and clients.

gateway/mcp/mcp_transport_base.py

@@ -238,12 +238,10 @@ class MCPTransportBase(ABC):
             )
         ):
             # Add the trace to the explorer
-            asyncio.create_task(
-                session_store.add_message_to_session(
-                    session_id=session_id,
-                    message=message,
-                    guardrails_result=guardrails_result,
-                )
+            await session_store.add_message_to_session(
+                session_id=session_id,
+                message=message,
+                guardrails_result=guardrails_result,
             )
             return {
                 "jsonrpc": "2.0",

tests/integration/mcp/test_mcp.py

@@ -82,14 +82,12 @@ async def _invoke_mcp_tool(
     elif transport == "sse":
         return await mcp_sse_client_run(
             f"{gateway_url}/api/v1/gateway/mcp/sse",
-            push,
             tool_name,
             tool_args,
             headers=_get_headers(_get_server_base_url(transport), project_name, push),
         )
     return await mcp_streamable_client_run(
         f"{gateway_url}/api/v1/gateway/mcp/streamable",
-        push,
         tool_name,
         tool_args,
         headers=_get_headers(_get_server_base_url(transport), project_name, push),
@@ -97,7 +95,7 @@ async def _invoke_mcp_tool(
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(30)
+@pytest.mark.timeout(20)
 @pytest.mark.parametrize(
     "transport",
     [
@@ -182,7 +180,7 @@ async def test_mcp_with_gateway(
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(30)
+@pytest.mark.timeout(20)
 @pytest.mark.parametrize(
     "transport",
     [
@@ -313,7 +311,7 @@ async def test_mcp_with_gateway_and_logging_guardrails(
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(30)
+@pytest.mark.timeout(20)
 @pytest.mark.parametrize(
     "transport",
     [
@@ -345,12 +343,10 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
         invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
     )
 
-    # Run the MCP client and make the tool call.
-    try:
+    with pytest.raises(ExceptionGroup) as exc_group:
         if transport == "sse":
             _ = await mcp_sse_client_run(
                 gateway_url + "/api/v1/gateway/mcp/sse",
-                push_to_explorer=True,
                 tool_name="get_last_message_from_user",
                 tool_args={"username": "Alice"},
                 headers=_get_headers(
@@ -366,30 +362,16 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
                 tool_name="get_last_message_from_user",
                 tool_args={"username": "Alice"},
             )
-        if not transport.startswith("streamable-"):
-            # 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
-
-    if transport.startswith("streamable-"):
-        with pytest.raises(ExceptionGroup) as exc_group:
+        else:
             _ = await mcp_streamable_client_run(
                 gateway_url + "/api/v1/gateway/mcp/streamable",
-                push_to_explorer=True,
                 tool_name="get_last_message_from_user",
                 tool_args={"username": "Alice"},
                 headers=_get_headers(
                     _get_streamable_server_base_url(transport), project_name, True
                 ),
             )
+    if transport.startswith("streamable-"):
         # Extract the actual HTTPStatusError
         http_errors = [
             e
@@ -397,6 +379,14 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
             if isinstance(e, httpx.HTTPStatusError)
         ]
         assert http_errors[0].response.status_code == 400
+    else:
+        mcp_error = [e for e in exc_group.value.exceptions][0].exceptions[0]
+        assert (
+            "[Invariant Guardrails] The MCP tool call was blocked for security reasons"
+            in mcp_error.error.message
+        )
+        assert "get_last_message_from_user is called" in mcp_error.error.message
+        assert -32600 == mcp_error.error.code
 
     # Fetch the trace ids for the dataset
     traces_response = requests.get(
@@ -439,7 +429,7 @@ async def test_mcp_with_gateway_and_blocking_guardrails(
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(30)
+@pytest.mark.timeout(20)
 @pytest.mark.parametrize(
     "transport",
     [
@@ -479,12 +469,10 @@ async def test_mcp_with_gateway_hybrid_guardrails(
         invariant_authorization="Bearer " + os.getenv("INVARIANT_API_KEY"),
     )
 
-    # Run the MCP client and make the tool call.
-    try:
+    with pytest.raises(ExceptionGroup) as exc_group:
         if transport == "sse":
             _ = await mcp_sse_client_run(
                 gateway_url + "/api/v1/gateway/mcp/sse",
-                push_to_explorer=True,
                 tool_name="get_last_message_from_user",
                 tool_args={"username": "Alice"},
                 headers=_get_headers(
@@ -500,46 +488,31 @@ async def test_mcp_with_gateway_hybrid_guardrails(
                 tool_name="get_last_message_from_user",
                 tool_args={"username": "Alice"},
             )
-        if not transport.startswith("streamable-"):
-            # 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
-
-    if transport.startswith("streamable-"):
-        with pytest.raises(ExceptionGroup) as exc_group:
+        else:
             _ = await mcp_streamable_client_run(
                 gateway_url + "/api/v1/gateway/mcp/streamable",
-                push_to_explorer=True,
                 tool_name="get_last_message_from_user",
                 tool_args={"username": "Alice"},
                 headers=_get_headers(
                     _get_streamable_server_base_url(transport), project_name, True
                 ),
             )
-        if transport.startswith("streamable-json"):
-            # Extract the actual HTTPStatusError
-            http_errors = [
-                e
-                for e in exc_group.value.exceptions
-                if isinstance(e, httpx.HTTPStatusError)
-            ]
-            assert http_errors[0].response.status_code == 400
-        else:
-            mcp_error = [e for e in exc_group.value.exceptions][0].exceptions[0]
-            assert (
-                "[Invariant Guardrails] The MCP tool call was blocked for security reasons"
-                in mcp_error.error.message
-            )
-            assert "food in ToolOutput" in mcp_error.error.message
-            assert -32600 == mcp_error.error.code
+    if transport.startswith("streamable-json"):
+        # Extract the actual HTTPStatusError
+        http_errors = [
+            e
+            for e in exc_group.value.exceptions
+            if isinstance(e, httpx.HTTPStatusError)
+        ]
+        assert http_errors[0].response.status_code == 400
+    else:
+        mcp_error = [e for e in exc_group.value.exceptions][0].exceptions[0]
+        assert (
+            "[Invariant Guardrails] The MCP tool call was blocked for security reasons"
+            in mcp_error.error.message
+        )
+        assert "food in ToolOutput" in mcp_error.error.message
+        assert -32600 == mcp_error.error.code
 
     # Fetch the trace ids for the dataset
     traces_response = requests.get(
@@ -602,7 +575,7 @@ async def test_mcp_with_gateway_hybrid_guardrails(
 
 
 @pytest.mark.asyncio
-@pytest.mark.timeout(30)
+@pytest.mark.timeout(20)
 @pytest.mark.parametrize(
     "transport",
     [
@@ -642,7 +615,6 @@ async def test_mcp_tool_list_blocking(
         with pytest.raises(ExceptionGroup) as exc_group:
             _ = await mcp_streamable_client_run(
                 gateway_url + "/api/v1/gateway/mcp/streamable",
-                push_to_explorer=True,
                 tool_name="tools/list",
                 tool_args={},
                 headers=_get_headers(
@@ -700,7 +672,6 @@ async def test_mcp_sse_post_endpoint_exceptions(gateway_url):
     with pytest.raises(ExceptionGroup) as exc_group:
         await mcp_sse_client_run(
             gateway_url + "/api/v1/gateway/mcp/sse",
-            push_to_explorer=True,
             tool_name="get_last_message_from_user",
             tool_args={"username": "Alice"},
             headers={

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

@@ -1,79 +1,18 @@
 """This is a simple example of how to use the MCP client with SSE transport."""
 
-# pylint: disable=E1101
-# pylint: disable=W0201
-
-import asyncio
+from typing import Any
 from datetime import timedelta
 
-from typing import Any, Optional
-from contextlib import AsyncExitStack
-
-from mcp import ClientSession
+from mcp import ClientSession, types
 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,
-    push_to_explorer: bool,
     tool_name: str,
     tool_args: dict[str, Any],
     headers: dict[str, str] = None,
-):
+) -> types.CallToolResult | types.ListToolsResult:
     """
     Run the MCP client with the given parameters.
 
@@ -82,23 +21,23 @@ async def run(
         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
-
+        headers: Optional headers to include in the request
     """
-    client = MCPClient()
-    try:
-        await client.connect_to_sse_server(
-            server_url=gateway_url, headers=headers or {}
-        )
-        # 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
-        if push_to_explorer:
-            await asyncio.sleep(2)
-        await client.cleanup()
+    client = sse_client(
+        url=gateway_url,
+        timeout=5,
+        headers=headers or {},
+        sse_read_timeout=10,
+    )
+    async with client as streams:
+        async with ClientSession(*streams) as session:
+            await session.initialize()
+            # list tools
+            listed_tools = await session.list_tools()
+            # call tool
+            if tool_name == "tools/list":
+                return listed_tools
+            else:
+                return await session.call_tool(
+                    tool_name, tool_args, read_timeout_seconds=timedelta(seconds=10)
+                )

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

@@ -1,106 +1,55 @@
 """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
 
 from mcp import ClientSession, StdioServerParameters, types
 from mcp.client.stdio import stdio_client
 
 
-class MCPClient:
-    """MCP Client for interacting with a MCP stdio server and processing queries"""
+def _get_server_params(
+    invariant_gateway_package_whl_file: str,
+    project_name: str,
+    server_script_path: str,
+    push_to_explorer: bool,
+    metadata_keys: Optional[dict[str, str]] = None,
+) -> StdioServerParameters:
+    args = [
+        "--from",
+        invariant_gateway_package_whl_file,
+        "invariant-gateway",
+        "mcp",
+        "--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)
 
-    def __init__(self):
-        self.session: Optional[ClientSession] = None
-        self.exit_stack = AsyncExitStack()
-
-    async def connect_to_server(
-        self,
-        invariant_gateway_package_whl_file: str,
-        project_name: str,
-        server_script_path: str,
-        push_to_explorer: bool,
-        metadata_keys: Optional[dict[str, str]] = None,
-    ):
-        """
-        Connect to an MCP server.
-
-        Args:
-            invariant_gateway_package_whl_file: Path to the Invariant Gateway package
-                .whl file
-            project_name: Name of the project in Invariant Explorer
-            server_script_path: Path to the server script
-            push_to_explorer: Whether to push traces to the Invariant Explorer
-        """
-        args = [
-            "--from",
-            invariant_gateway_package_whl_file,
-            "invariant-gateway",
-            "mcp",
-            "--project-name",
-            project_name,
+    if push_to_explorer:
+        args.append("--push-explorer")
+    args.extend(
+        [
+            "--exec",
+            "uv",
+            "--directory",
+            os.path.abspath(os.path.dirname(server_script_path)),
+            "run",
+            os.path.basename(server_script_path),
         ]
-        # 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(
-            [
-                "--exec",
-                "uv",
-                "--directory",
-                os.path.abspath(os.path.dirname(server_script_path)),
-                "run",
-                os.path.basename(server_script_path),
-            ]
-        )
-
-        server_params = StdioServerParameters(
-            command="uvx",
-            args=args,
-            env={
-                "INVARIANT_API_KEY": os.environ.get("INVARIANT_API_KEY"),
-                "INVARIANT_API_URL": "http://invariant-gateway-test-explorer-app-api:8000",
-            },
-        )
-
-        stdio_transport = await self.exit_stack.enter_async_context(
-            stdio_client(server_params)
-        )
-        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=15)
-            )
-        )
-
-        # initialize the session
-        await self.session.initialize()
-
-    async def call_tool(
-        self, tool_name: str, tool_args: dict[str, Any]
-    ) -> types.CallToolResult:
-        """
-        Make a tool call on the MCP server.
-
-        Args:
-            tool_name: Name of the tool to call
-            tool_args: Arguments for the tool call
-        """
-        # Execute tool call
-        result = await self.session.call_tool(tool_name, tool_args)
-        return result
-
-    async def cleanup(self):
-        """Clean up resources"""
-        await self.exit_stack.aclose()
+    return StdioServerParameters(
+        command="uvx",
+        args=args,
+        env={
+            "INVARIANT_API_KEY": os.environ.get("INVARIANT_API_KEY"),
+            "INVARIANT_API_URL": os.environ.get("INVARIANT_API_URL"),
+        },
+    )
 
 
 async def run(
@@ -111,7 +60,7 @@ async def run(
     tool_name: str,
     tool_args: dict[str, Any],
     metadata_keys: Optional[dict[str, str]] = None,
-) -> types.CallToolResult:
+) -> types.CallToolResult | types.ListToolsResult:
     """
     Main function to setup the MCP client and server.
     It calls a tool on the server with the given args.
@@ -124,24 +73,26 @@ async def run(
         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
+        metadata_keys: Optional metadata keys to include in the request
     """
-
-    client = MCPClient()
-    try:
-        await client.connect_to_server(
+    client = stdio_client(
+        _get_server_params(
             invariant_gateway_package_whl_file,
             project_name,
             server_script_path,
             push_to_explorer,
             metadata_keys=metadata_keys,
         )
-        listed_tools = await client.session.list_tools()
-        if tool_name == "tools/list":
+    )
+    async with client as (stdio, write):
+        async with ClientSession(
+            stdio, write, read_timeout_seconds=timedelta(seconds=10)
+        ) as session:
+            await session.initialize()
             # list tools
-            return listed_tools
-        else:
-            return await client.call_tool(tool_name, tool_args)
-    finally:
-        if push_to_explorer:
-            await asyncio.sleep(2)
-        await client.cleanup()
+            listed_tools = await session.list_tools()
+            # call tool
+            if tool_name == "tools/list":
+                return listed_tools
+            else:
+                return await session.call_tool(tool_name, tool_args)

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

@@ -1,21 +1,19 @@
 """This is a simple example of how to use the MCP client with Streamable HTTP transport."""
 
-import asyncio
 from datetime import timedelta
 
 from typing import Any
 
-from mcp import ClientSession
+from mcp import ClientSession, types
 from mcp.client.streamable_http import streamablehttp_client
 
 
 async def run(
     gateway_url: str,
-    push_to_explorer: bool,
     tool_name: str,
     tool_args: dict[str, Any],
     headers: dict[str, str] = None,
-):
+) -> types.CallToolResult | types.ListToolsResult:
     """
     Run the MCP client with the given parameters.
 
@@ -26,27 +24,21 @@ async def run(
         tool_args: Arguments for the tool call
 
     """
-    try:
-        streams_context = streamablehttp_client(
-            url=gateway_url,
-            headers=headers or {},
-            timeout=timedelta(seconds=5),
-            sse_read_timeout=timedelta(seconds=10),
-        )
-        async with streams_context as (read_stream, write_stream, _):
-            async with ClientSession(read_stream, write_stream) as session:
-                await session.initialize()
-                # list tools
-                listed_tools = await session.list_tools()
-                # call tool
-                if tool_name == "tools/list":
-                    return listed_tools
-                else:
-                    return await session.call_tool(
-                        tool_name, tool_args, read_timeout_seconds=timedelta(seconds=10)
-                    )
-    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)
+    client = streamablehttp_client(
+        url=gateway_url,
+        headers=headers or {},
+        timeout=timedelta(seconds=5),
+        sse_read_timeout=timedelta(seconds=10),
+    )
+    async with client as (read_stream, write_stream, _):
+        async with ClientSession(read_stream, write_stream) as session:
+            await session.initialize()
+            # list tools
+            listed_tools = await session.list_tools()
+            # call tool
+            if tool_name == "tools/list":
+                return listed_tools
+            else:
+                return await session.call_tool(
+                    tool_name, tool_args, read_timeout_seconds=timedelta(seconds=10)
+                )