From cc3e96c20a6141db4286f4008751d27a1f8d4c7a Mon Sep 17 00:00:00 2001 From: Hemang Date: Wed, 4 Jun 2025 10:07:22 +0200 Subject: [PATCH] Clean up MCP tests and clients. --- gateway/mcp/mcp_transport_base.py | 10 +- tests/integration/mcp/test_mcp.py | 97 ++++------- .../resources/mcp/sse/client/main.py | 105 +++--------- .../resources/mcp/stdio/client/main.py | 157 ++++++------------ .../resources/mcp/streamable/client/main.py | 48 +++--- 5 files changed, 134 insertions(+), 283 deletions(-) diff --git a/gateway/mcp/mcp_transport_base.py b/gateway/mcp/mcp_transport_base.py index 29d728c..7557472 100644 --- a/gateway/mcp/mcp_transport_base.py +++ b/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", diff --git a/tests/integration/mcp/test_mcp.py b/tests/integration/mcp/test_mcp.py index c892127..2606c2b 100644 --- a/tests/integration/mcp/test_mcp.py +++ b/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={ diff --git a/tests/integration/resources/mcp/sse/client/main.py b/tests/integration/resources/mcp/sse/client/main.py index 098b67b..d74b592 100644 --- a/tests/integration/resources/mcp/sse/client/main.py +++ b/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) + ) diff --git a/tests/integration/resources/mcp/stdio/client/main.py b/tests/integration/resources/mcp/stdio/client/main.py index ebc55b5..3aeaaf3 100644 --- a/tests/integration/resources/mcp/stdio/client/main.py +++ b/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) diff --git a/tests/integration/resources/mcp/streamable/client/main.py b/tests/integration/resources/mcp/streamable/client/main.py index a498e5b..da7fc23 100644 --- a/tests/integration/resources/mcp/streamable/client/main.py +++ b/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) + )