From 4bb1cf927c851c4a97c18745a10f50d1f78f46a5 Mon Sep 17 00:00:00 2001 From: DABH Date: Thu, 10 Sep 2026 00:31:07 -0500 Subject: [PATCH] Run google_genai MCP tests in-memory and close pooled connections The MCP tests spawned the FastMCP echo server as a stdio subprocess for every test and never shut the pooled connection down: the autouse fixture only cleared the pool dict, so each test left a ~67 MB server process (and its owner task) alive for the rest of the pytest process. On macOS CI the spawn itself also stalled: two timed-out runs show the pytest main thread inside _posixsubprocess.fork_exec for ~65s while forking the large test process, freezing the worker's event loop until the workflow deadline passed (test_mcp_side_effects, test_mcp_full_schema_propagation). Connect the pooling, schema, replay and side-effect tests to the same echo server over mcp.shared.memory streams instead; they exercise the connection pool and activity scheduling, not the transport. Keep the stdio subprocess for test_mcp_tool_discovery_and_call, and have the fixture evict (close) leftover connections so no server outlives its test. --- tests/contrib/google_genai/echo_mcp_server.py | 6 +++- tests/contrib/google_genai/test_gemini_mcp.py | 29 +++++++++++++++---- 2 files changed, 28 insertions(+), 7 deletions(-) diff --git a/tests/contrib/google_genai/echo_mcp_server.py b/tests/contrib/google_genai/echo_mcp_server.py index a69ac7b22..033810d1d 100644 --- a/tests/contrib/google_genai/echo_mcp_server.py +++ b/tests/contrib/google_genai/echo_mcp_server.py @@ -1,4 +1,8 @@ -"""A minimal stdio MCP server used by the google_genai MCP tests.""" +"""A minimal MCP echo server for the google_genai MCP tests. + +Run as a script it serves over stdio; imported, ``mcp`` can be connected to +in-memory via ``mcp.shared.memory``. +""" from mcp.server.fastmcp import FastMCP diff --git a/tests/contrib/google_genai/test_gemini_mcp.py b/tests/contrib/google_genai/test_gemini_mcp.py index 04e36bb9d..cf6225b5f 100644 --- a/tests/contrib/google_genai/test_gemini_mcp.py +++ b/tests/contrib/google_genai/test_gemini_mcp.py @@ -8,6 +8,10 @@ - full parameter-schema propagation to the model (the MCP wire-format check) - replay determinism and exact activity-scheduling counts +Only the discovery test spawns the echo server as a stdio subprocess; the +rest connect to the same server over in-memory streams, since they exercise +the pooling, activity, and replay logic rather than the transport. + Plus the server-side pass-through paths that need no shim code: - Vertex AI ``Tool(mcp_servers=[McpServer(...)])`` config serialization - Interactions API ``MCPServerToolCallStep`` / ``MCPServerToolResultStep`` rehydration @@ -56,7 +60,19 @@ @asynccontextmanager async def _echo_session() -> AsyncIterator[ClientSession]: - """Yield a connected, initialized session to the stdio echo MCP server.""" + """Yield an initialized session to the echo MCP server over in-memory streams.""" + from mcp.shared.memory import create_connected_server_and_client_session + + # Imported here so the sandbox's re-import of this module stays side-effect free. + from tests.contrib.google_genai.echo_mcp_server import mcp as echo_server + + async with create_connected_server_and_client_session(echo_server) as session: + yield session + + +@asynccontextmanager +async def _echo_stdio_session() -> AsyncIterator[ClientSession]: + """Yield an initialized session to the echo server run as a stdio subprocess.""" params = StdioServerParameters(command=sys.executable, args=[_ECHO_SERVER]) async with stdio_client(params) as (read, write): async with ClientSession(read, write) as session: @@ -85,7 +101,7 @@ def _apply_mcp_plugin( Monkeypatches ``GeminiApiCaller.activities`` (so canned generate_content responses drive the AFC loop) while leaving the plugin's MCP activities — - built from ``mcp_servers`` — to hit the real stdio echo server. + built from ``mcp_servers`` — to hit the real echo server. """ from temporalio.contrib.google_genai._gemini_activity import GeminiApiCaller @@ -128,13 +144,14 @@ async def _activity_names(handle: Any) -> list[str]: @pytest.fixture(autouse=True) -def _clear_mcp_connections(): # pyright: ignore[reportUnusedFunction] +async def _close_mcp_connections(): # pyright: ignore[reportUnusedFunction] """Isolate the module-global MCP connection pool between tests.""" from temporalio.contrib.google_genai import _mcp _mcp._CONNECTIONS.clear() yield - _mcp._CONNECTIONS.clear() + for server in list(_mcp._CONNECTIONS): + await _mcp._evict_connection(server) # --------------------------------------------------------------------------- @@ -176,7 +193,7 @@ async def run(self, server_name: str, prompt: str) -> str: async def test_mcp_tool_discovery_and_call(client: Client): - """The AFC loop discovers + calls an MCP tool through activities.""" + """The AFC loop discovers + calls a tool on a stdio MCP server via activities.""" server = "echo_basic" new_client, _ = _apply_mcp_plugin( client, @@ -184,7 +201,7 @@ async def test_mcp_tool_discovery_and_call(client: Client): make_function_call_response("echo", {"message": "hello"}), make_text_response("Done!"), ], - mcp_servers={server: _echo_session}, + mcp_servers={server: _echo_stdio_session}, ) async with new_worker(new_client, McpToolWorkflow) as worker: