diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 6aa671a5011..fe11349b62b 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -5,7 +5,7 @@ LiteLLM Proxy uses this MCP Client to connnect to other MCP servers. import asyncio import base64 from datetime import timedelta -from typing import Callable, Dict, List, Optional, Union +from typing import Awaitable, Callable, Dict, List, Optional, TypeVar, Union import httpx from mcp import ClientSession, StdioServerParameters @@ -34,6 +34,9 @@ def to_basic_auth(auth_value: str) -> str: return base64.b64encode(auth_value.encode("utf-8")).decode() +TSessionResult = TypeVar("TSessionResult") + + class MCPClient: """ MCP Client supporting: @@ -58,12 +61,6 @@ class MCPClient: self.auth_type: MCPAuthType = auth_type self.timeout: float = timeout self._mcp_auth_value: Optional[Union[str, Dict[str, str]]] = None - self._session: Optional[ClientSession] = None - self._context = None - self._transport_ctx = None - self._transport = None - self._session_ctx = None - self._task: Optional[asyncio.Task] = None self.stdio_config: Optional[MCPStdioConfig] = stdio_config self.extra_headers: Optional[Dict[str, str]] = extra_headers self.ssl_verify: Optional[VerifyTypes] = ssl_verify @@ -71,33 +68,14 @@ class MCPClient: if auth_value: self.update_auth_value(auth_value) - async def __aenter__(self): - """ - Enable async context manager support. - Initializes the transport and session. - """ - try: - await self.connect() - return self - except Exception: - await self.disconnect() - raise - - async def connect(self): - """Initialize the transport and session.""" - if self._session: - verbose_logger.debug( - f"MCP client already connected to {self.server_url or 'stdio'}" - ) - return # Already connected - - verbose_logger.info( - f"MCP client connecting to {self.server_url or 'stdio'} via {self.transport_type}" - ) + async def run_with_session( + self, operation: Callable[[ClientSession], Awaitable[TSessionResult]] + ) -> TSessionResult: + """Open a session, run the provided coroutine, and clean up.""" + transport_ctx = None try: if self.transport_type == MCPTransport.stdio: - # For stdio transport, use stdio_client with command-line parameters if not self.stdio_config: raise ValueError("stdio_config is required for stdio transport") @@ -106,117 +84,43 @@ class MCPClient: args=self.stdio_config.get("args", []), env=self.stdio_config.get("env", {}), ) - - self._transport_ctx = stdio_client(server_params) - self._transport = await self._transport_ctx.__aenter__() - self._session_ctx = ClientSession( - self._transport[0], self._transport[1] - ) - self._session = await self._session_ctx.__aenter__() - await self._session.initialize() - verbose_logger.info( - f"MCP client successfully connected via stdio: {self.stdio_config.get('command', '')}" - ) + transport_ctx = stdio_client(server_params) elif self.transport_type == MCPTransport.sse: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() - self._transport_ctx = sse_client( + transport_ctx = sse_client( url=self.server_url, timeout=self.timeout, headers=headers, httpx_client_factory=httpx_client_factory, ) - self._transport = await self._transport_ctx.__aenter__() - self._session_ctx = ClientSession( - self._transport[0], self._transport[1] - ) - self._session = await self._session_ctx.__aenter__() - await self._session.initialize() - verbose_logger.info( - f"MCP client successfully connected via SSE to {self.server_url}" - ) - else: # http + else: headers = self._get_auth_headers() httpx_client_factory = self._create_httpx_client_factory() verbose_logger.debug( "litellm headers for streamablehttp_client: %s", headers ) - self._transport_ctx = streamablehttp_client( + transport_ctx = streamablehttp_client( url=self.server_url, timeout=timedelta(seconds=self.timeout), headers=headers, httpx_client_factory=httpx_client_factory, ) - self._transport = await self._transport_ctx.__aenter__() - self._session_ctx = ClientSession( - self._transport[0], self._transport[1] - ) - self._session = await self._session_ctx.__aenter__() - await self._session.initialize() - verbose_logger.info( - f"MCP client successfully connected via HTTP to {self.server_url}" - ) - except ValueError as e: - # Re-raise ValueError exceptions (like missing stdio_config) - verbose_logger.warning(f"MCP client connection failed: {str(e)}") - await self.disconnect() + + if transport_ctx is None: + raise RuntimeError("Failed to create transport context") + + async with transport_ctx as transport: + read_stream, write_stream = transport[0], transport[1] + session_ctx = ClientSession(read_stream, write_stream) + async with session_ctx as session: + await session.initialize() + return await operation(session) + except Exception: + verbose_logger.warning( + "MCP client run_with_session failed for %s", self.server_url or "stdio" + ) raise - except Exception as e: - verbose_logger.warning(f"MCP client connection failed: {str(e)}") - await self.disconnect() - # Don't raise other exceptions, let the calling code handle it gracefully - # This allows the server manager to continue with other servers - # Instead of raising, we'll let the calling code handle the failure - pass - - async def __aexit__(self, exc_type, exc_val, exc_tb): - """Cleanup when exiting context manager.""" - await self.disconnect() - - async def disconnect(self): - """Clean up session and connections.""" - verbose_logger.info( - f"MCP client disconnecting from {self.server_url or 'stdio'}" - ) - - if self._task and not self._task.done(): - verbose_logger.debug("MCP client cancelling background task") - self._task.cancel() - try: - await self._task - except asyncio.CancelledError: - pass - - if self._session: - try: - verbose_logger.debug("MCP client closing session") - await self._session_ctx.__aexit__(None, None, None) # type: ignore - except Exception as e: - verbose_logger.debug( - f"Error closing MCP session: {type(e).__name__}: {str(e)}" - ) - pass - self._session = None - self._session_ctx = None - - if self._transport_ctx: - try: - verbose_logger.debug("MCP client closing transport") - await self._transport_ctx.__aexit__(None, None, None) - except Exception as e: - verbose_logger.debug( - f"Error closing MCP transport: {type(e).__name__}: {str(e)}" - ) - pass - self._transport_ctx = None - self._transport = None - - if self._context: - try: - await self._context.__aexit__(None, None, None) # type: ignore - except Exception: - pass - self._context = None def update_auth_value(self, mcp_auth_value: Union[str, Dict[str, str]]): """ @@ -294,24 +198,11 @@ class MCPClient: f"MCP client listing tools from {self.server_url or 'stdio'}" ) - if not self._session: - verbose_logger.debug("MCP client session not found, attempting to connect") - try: - await self.connect() - except Exception as e: - verbose_logger.error( - f"MCP client connection failed during list_tools: {type(e).__name__}: {str(e)}" - ) - return [] - - if self._session is None: - verbose_logger.error( - "MCP client session is not initialized after connection attempt" - ) - return [] + async def _list_tools_operation(session: ClientSession): + return await session.list_tools() try: - result = await self._session.list_tools() + result = await self.run_with_session(_list_tools_operation) tool_count = len(result.tools) tool_names = [tool.name for tool in result.tools] verbose_logger.info( @@ -320,7 +211,6 @@ class MCPClient: return result.tools except asyncio.CancelledError: verbose_logger.warning("MCP client list_tools was cancelled") - await self.disconnect() raise except Exception as e: error_type = type(e).__name__ @@ -339,7 +229,6 @@ class MCPClient: "the MCP server may have crashed, disconnected, or timed out" ) - await self.disconnect() # Return empty list instead of raising to allow graceful degradation return [] @@ -353,55 +242,21 @@ class MCPClient: f"MCP client calling tool '{call_tool_request_params.name}' with arguments: {call_tool_request_params.arguments}" ) - if not self._session: - verbose_logger.warning( - "MCP client session not found, attempting to connect" - ) - try: - await self.connect() - except Exception as e: - verbose_logger.error( - f"MCP client connection failed before tool call: {type(e).__name__}: {str(e)}" - ) - return MCPCallToolResult( - content=[TextContent(type="text", text=f"{str(e)}")], isError=True - ) - - if self._session is None: - verbose_logger.error( - "MCP client session is not initialized after connection attempt" - ) - return MCPCallToolResult( - content=[ - TextContent( - type="text", text="MCP client session is not initialized" - ) - ], - isError=True, - ) - - # Check session and transport state before calling tool - verbose_logger.debug( - f"MCP client state before tool call - " - f"session: {'active' if self._session else 'none'}, " - f"transport: {'active' if self._transport else 'none'}, " - f"session_ctx: {'active' if self._session_ctx else 'none'}, " - f"transport_ctx: {'active' if self._transport_ctx else 'none'}" - ) - - try: + async def _call_tool_operation(session: ClientSession): verbose_logger.debug("MCP client sending tool call to session") - tool_result = await self._session.call_tool( + return await session.call_tool( name=call_tool_request_params.name, arguments=call_tool_request_params.arguments, ) + + try: + tool_result = await self.run_with_session(_call_tool_operation) verbose_logger.info( f"MCP client tool call '{call_tool_request_params.name}' completed successfully" ) return tool_result except asyncio.CancelledError: verbose_logger.warning("MCP client tool call was cancelled") - await self.disconnect() raise except Exception as e: import traceback @@ -424,11 +279,9 @@ class MCPClient: if "BrokenResourceError" in error_type or "Broken" in error_type: verbose_logger.error( "MCP client detected broken connection/stream - " - "the MCP server may have crashed, disconnected, or timed out. " - "Session and transport will be disconnected." + "the MCP server may have crashed, disconnected, or timed out." ) - await self.disconnect() # Return a default error result instead of raising return MCPCallToolResult( content=[ diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index 5e2badb7bbe..1e301c95899 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -385,12 +385,12 @@ class MCPServerManager: ) # Update tool name to server name mapping (for both prefixed and base names) - self.tool_name_to_mcp_server_name_mapping[base_tool_name] = ( - server_prefix - ) - self.tool_name_to_mcp_server_name_mapping[prefixed_tool_name] = ( - server_prefix - ) + self.tool_name_to_mcp_server_name_mapping[ + base_tool_name + ] = server_prefix + self.tool_name_to_mcp_server_name_mapping[ + prefixed_tool_name + ] = server_prefix registered_count += 1 verbose_logger.debug( @@ -714,12 +714,6 @@ class MCPServerManager: f"Failed to get tools from server {server.name}: {str(e)}" ) return [] - finally: - if client: - try: - await client.disconnect() - except Exception: - pass async def _descovery_metadata( self, @@ -983,8 +977,6 @@ class MCPServerManager: async def _list_tools_task(): try: - await client.connect() - tools = await client.list_tools() verbose_logger.debug(f"Tools from {server_name}: {tools}") return tools @@ -1439,14 +1431,12 @@ class MCPServerManager: ) async def _call_tool_via_client(client, params): - async with client: - return await client.call_tool(params) + return await client.call_tool(params) tasks.append( asyncio.create_task(_call_tool_via_client(client, call_tool_params)) ) - # IMPORTANT: Must await tasks INSIDE the context manager to keep connection alive try: mcp_responses = await asyncio.gather(*tasks) except ( diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index e427507a4b8..4288f25740c 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -76,12 +76,12 @@ if MCP_AVAILABLE: mcp_auth_header=server_auth_header, add_prefix=False, ) - + # Filter tools based on allowed_tools configuration # Only filter if allowed_tools is explicitly configured (not None and not empty) if server.allowed_tools is not None and len(server.allowed_tools) > 0: tools = filter_tools_by_allowed_tools(tools, server) - + return _create_tool_response_objects(tools, server.mcp_info) ######################################################## @@ -212,7 +212,9 @@ if MCP_AVAILABLE: from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException from litellm.proxy.proxy_server import add_litellm_data_to_request, proxy_config - from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import MCPRequestHandler + from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( + MCPRequestHandler, + ) try: data = await request.json() @@ -222,21 +224,27 @@ if MCP_AVAILABLE: user_api_key_dict=user_api_key_dict, proxy_config=proxy_config, ) - + # FIX: Extract MCP auth headers from request # The UI sends bearer token in x-mcp-auth header and server-specific headers, # but they weren't being extracted and passed to call_mcp_tool. # This fix ensures auth headers are properly extracted from the HTTP request # and passed through to the MCP server for authentication. - mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers(request.headers) - mcp_server_auth_headers = MCPRequestHandler._get_mcp_server_auth_headers_from_headers(request.headers) - + mcp_auth_header = MCPRequestHandler._get_mcp_auth_header_from_headers( + request.headers + ) + mcp_server_auth_headers = ( + MCPRequestHandler._get_mcp_server_auth_headers_from_headers( + request.headers + ) + ) + # Add extracted headers to data dict to pass to call_mcp_tool if mcp_auth_header: data["mcp_auth_header"] = mcp_auth_header if mcp_server_auth_headers: data["mcp_server_auth_headers"] = mcp_server_auth_headers - + result = await call_mcp_tool(**data) return result except BlockedPiiEntityError as e: @@ -296,7 +304,6 @@ if MCP_AVAILABLE: Returns: Operation result or error response """ - client = None try: client = global_mcp_server_manager._create_mcp_client( server=MCPServer( @@ -315,13 +322,6 @@ if MCP_AVAILABLE: except Exception as e: verbose_logger.error(f"Error in MCP operation: {e}", exc_info=True) return {"status": "error", "message": "An internal error has occurred."} - finally: - # Ensure client is properly disconnected before response is sent - if client is not None: - try: - await client.disconnect() - except Exception as e: - verbose_logger.warning(f"Error disconnecting MCP client: {e}") @router.post("/test/connection") async def test_connection( @@ -332,7 +332,10 @@ if MCP_AVAILABLE: """ async def _test_connection_operation(client): - await client.connect() + async def _noop(session): + return "ok" + + await client.run_with_session(_noop) return {"status": "ok"} return await _execute_with_mcp_client(request, _test_connection_operation) @@ -347,7 +350,13 @@ if MCP_AVAILABLE: """ async def _list_tools_operation(client): - list_tools_result: List[MCPTool] = await client.list_tools() + async def _list_tools_session_operation(session): + return await session.list_tools() + + list_tools_response = await client.run_with_session( + _list_tools_session_operation + ) + list_tools_result: List[MCPTool] = list_tools_response.tools model_dumped_tools: List[dict] = [ tool.model_dump() for tool in list_tools_result ] diff --git a/tests/mcp_tests/test_mcp_client_unit.py b/tests/mcp_tests/test_mcp_client_unit.py index cdeee679d18..451c3774372 100644 --- a/tests/mcp_tests/test_mcp_client_unit.py +++ b/tests/mcp_tests/test_mcp_client_unit.py @@ -84,8 +84,8 @@ class TestMCPClientUnitTests: @pytest.mark.asyncio @patch("litellm.experimental_mcp_client.client.streamablehttp_client") @patch("litellm.experimental_mcp_client.client.ClientSession") - async def test_connect(self, mock_session_class, mock_transport): - """Test connecting to MCP server with authentication.""" + async def test_run_with_session(self, mock_session_class, mock_transport): + """Test run_with_session establishes session with auth headers.""" # Setup mocks mock_transport_ctx = AsyncMock() mock_transport.return_value = mock_transport_ctx @@ -102,7 +102,11 @@ class TestMCPClientUnitTests: auth_type=MCPAuth.bearer_token, auth_value="test_token", ) - await client.connect() + + async def _operation(session): + return "ok" + + await client.run_with_session(_operation) # Verify transport was created with auth headers call_args = mock_transport.call_args @@ -112,7 +116,6 @@ class TestMCPClientUnitTests: # Verify session was initialized mock_session_instance.initialize.assert_called_once() - assert client._session == mock_session_instance @pytest.mark.asyncio @patch("litellm.experimental_mcp_client.client.streamablehttp_client") diff --git a/tests/mcp_tests/test_mcp_logging.py b/tests/mcp_tests/test_mcp_logging.py index 6ee52364cc5..2b1f2699fd5 100644 --- a/tests/mcp_tests/test_mcp_logging.py +++ b/tests/mcp_tests/test_mcp_logging.py @@ -54,8 +54,6 @@ async def test_mcp_cost_tracking(): inputSchema={"type": "object", "properties": {"test": {"type": "string"}}} ) ]) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) # Mock the MCPClient constructor def mock_client_constructor(*args, **kwargs): @@ -118,7 +116,6 @@ async def test_mcp_cost_tracking(): assert response_list[0].text == "Test response" # Verify client methods were called - mock_client.__aenter__.assert_called() mock_client.call_tool.assert_called_once() ###### @@ -153,9 +150,6 @@ async def test_mcp_cost_tracking_per_tool(): inputSchema={"type": "object", "properties": {"data": {"type": "string"}}} ) ]) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - mock_client.disconnect = AsyncMock(return_value=None) # Mock the MCPClient constructor def mock_client_constructor(*args, **kwargs): @@ -297,9 +291,6 @@ async def test_mcp_tool_call_hook(): inputSchema={"type": "object", "properties": {"test": {"type": "string"}}} ) ]) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) - mock_client.disconnect = AsyncMock(return_value=None) # Mock the MCPClient constructor def mock_client_constructor(*args, **kwargs): @@ -351,4 +342,4 @@ async def test_mcp_tool_call_hook(): logged_standard_logging_payload = test_logger.standard_logging_payload print("logged_standard_logging_payload", json.dumps(logged_standard_logging_payload, indent=4)) assert logged_standard_logging_payload is not None, "Standard logging payload should not be None" - assert logged_standard_logging_payload["response_cost"] == 1.42 \ No newline at end of file + assert logged_standard_logging_payload["response_cost"] == 1.42 diff --git a/tests/mcp_tests/test_mcp_server.py b/tests/mcp_tests/test_mcp_server.py index 6e09e890ca8..c4cf67694b5 100644 --- a/tests/mcp_tests/test_mcp_server.py +++ b/tests/mcp_tests/test_mcp_server.py @@ -68,8 +68,6 @@ async def test_mcp_server_manager_https_server(): mock_client = AsyncMock() mock_client.list_tools = AsyncMock(return_value=mock_tools) mock_client.call_tool = AsyncMock(return_value=mock_result) - mock_client.__aenter__ = AsyncMock(return_value=mock_client) - mock_client.__aexit__ = AsyncMock(return_value=None) # Mock the MCPClient constructor def mock_client_constructor(*args, **kwargs): @@ -132,7 +130,6 @@ async def test_mcp_server_manager_https_server(): assert result.content[0].text == "Email sent successfully" # Verify client methods were called - mock_client.__aenter__.assert_called() mock_client.list_tools.assert_called() mock_client.call_tool.assert_called_once() @@ -301,7 +298,6 @@ async def test_mcp_http_transport_call_tool_mock(): assert result.content[0].text == "Email sent successfully to test@example.com" # Verify client methods were called - mock_client.__aenter__.assert_called() mock_client.call_tool.assert_called_once() @@ -364,7 +360,6 @@ async def test_mcp_http_transport_call_tool_error_mock(): assert "Error: Invalid email address" in result.content[0].text # Verify client methods were called - mock_client.__aenter__.assert_called() mock_client.call_tool.assert_called_once() diff --git a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py index 9ecc10cafed..14b747a5704 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -39,7 +39,10 @@ class TestMCPClient: with pytest.raises( ValueError, match="stdio_config is required for stdio transport" ): - await client.connect() + async def _noop(session): + return None + + await client.run_with_session(_noop) @pytest.mark.asyncio @patch("litellm.experimental_mcp_client.client.stdio_client") @@ -65,7 +68,10 @@ class TestMCPClient: client = MCPClient(transport_type=MCPTransport.stdio, stdio_config=stdio_config) - await client.connect() + async def _operation(session): + return "ok" + + await client.run_with_session(_operation) # Verify stdio_client was called with correct parameters mock_stdio_client.assert_called_once() @@ -109,7 +115,10 @@ class TestMCPClient: transport_type=MCPTransport.http, ) - await client.connect() + async def _operation(session): + return "ok" + + await client.run_with_session(_operation) # Verify streamablehttp_client was called mock_streamablehttp_client.assert_called_once() @@ -157,7 +166,10 @@ class TestMCPClient: ssl_verify=False, ) - await client.connect() + async def _operation(session): + return "ok" + + await client.run_with_session(_operation) # Verify sse_client was called mock_sse_client.assert_called_once() @@ -208,7 +220,10 @@ class TestMCPClient: ssl_verify=custom_ca_path, ) - await client.connect() + async def _operation(session): + return "ok" + + await client.run_with_session(_operation) # Verify streamablehttp_client was called mock_streamablehttp_client.assert_called_once() diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 2936b48c755..f929d2df107 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -540,7 +540,6 @@ async def test_oauth2_headers_passed_to_mcp_client(): ) # Return a mock client that doesn't actually connect mock_client = MagicMock() - mock_client.disconnect = AsyncMock() return mock_client # Mock _fetch_tools_with_timeout to avoid actual network calls diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index 33f013696bb..ce101dae8cd 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1347,10 +1347,10 @@ class TestMCPServerManager: @pytest.mark.asyncio async def test_call_tool_without_broken_pipe_error(self): """ - Test that call_tool properly uses async context manager to avoid broken pipe errors. - This test ensures that tasks are awaited INSIDE the context manager, keeping the connection alive. + Test that call_tool awaits the client call even without a persistent context manager. + Ensures the gathered tasks still include the MCP client call result. """ - from unittest.mock import AsyncMock, MagicMock, patch + from unittest.mock import AsyncMock, MagicMock from mcp.types import CallToolResult @@ -1369,42 +1369,17 @@ class TestMCPServerManager: manager.tool_name_to_mcp_server_name_mapping["test_tool"] = "test-server" manager.tool_name_to_mcp_server_name_mapping["test-server-test_tool"] = "test-server" - # Create mock client that tracks context manager usage - mock_client = MagicMock() - context_entered = False - context_exited = False - call_tool_called_inside_context = False - - async def mock_aenter(self): - nonlocal context_entered - context_entered = True - return self - - async def mock_aexit(self, exc_type, exc_val, exc_tb): - nonlocal context_exited - context_exited = True - # Verify that call_tool was called before context exit - assert ( - call_tool_called_inside_context - ), "call_tool must be awaited inside context manager" - return False + # Create mock client that tracks call_tool usage + mock_client = AsyncMock() async def mock_call_tool(params): - nonlocal call_tool_called_inside_context - # Verify we're inside the context when this is called - assert context_entered, "call_tool called outside context manager" - assert not context_exited, "call_tool called after context exit" - call_tool_called_inside_context = True - # Return a mock CallToolResult result = MagicMock(spec=CallToolResult) result.content = [{"type": "text", "text": "Tool executed successfully"}] result.isError = False return result - mock_client.__aenter__ = mock_aenter - mock_client.__aexit__ = mock_aexit - mock_client.call_tool = mock_call_tool + mock_client.call_tool.side_effect = mock_call_tool # Mock _create_mcp_client to return our mock client manager._create_mcp_client = MagicMock(return_value=mock_client) @@ -1437,12 +1412,8 @@ class TestMCPServerManager: assert result.isError is False assert len(result.content) > 0 - # Verify context manager was used properly - assert context_entered, "Context manager __aenter__ was not called" - assert context_exited, "Context manager __aexit__ was not called" - assert ( - call_tool_called_inside_context - ), "call_tool was not awaited inside context" + # Verify the MCP client call was awaited exactly once + assert mock_client.call_tool.await_count == 1 if __name__ == "__main__":