diff --git a/api/core/mcp/mcp_client.py b/api/core/mcp/mcp_client.py index c6d33292dae..387bb40d270 100644 --- a/api/core/mcp/mcp_client.py +++ b/api/core/mcp/mcp_client.py @@ -76,6 +76,12 @@ class MCPClient: self.connect_server(sse_client, "sse") except (MCPConnectionError, ValueError): logger.debug("MCP connection failed with 'sse', falling back to 'mcp' method.") + # The SSE transport and ClientSession can already be entered when + # the handshake fails. Close that partial connection before + # opening the fallback so both transports are never live at once. + self._exit_stack.close() + self._exit_stack = ExitStack() + self._session = None self.connect_server(streamablehttp_client, "mcp") def connect_server(self, client_factory: Callable[..., AbstractContextManager[Any]], method_name: str) -> None: diff --git a/api/tests/unit_tests/core/mcp/test_mcp_client.py b/api/tests/unit_tests/core/mcp/test_mcp_client.py index c6407648caa..4aac002eec0 100644 --- a/api/tests/unit_tests/core/mcp/test_mcp_client.py +++ b/api/tests/unit_tests/core/mcp/test_mcp_client.py @@ -214,6 +214,43 @@ class TestMCPClient: # Verify session was created with MCP assert client._session == mock_session + @patch("core.mcp.mcp_client.sse_client") + @patch("core.mcp.mcp_client.streamablehttp_client") + @patch("core.mcp.mcp_client.ClientSession") + def test_initialize_fallback_closes_partial_sse_connection( + self, mock_client_session, mock_streamable_client, mock_sse_client + ): + """A failed SSE handshake is closed before streamable HTTP fallback starts.""" + events: list[str] = [] + + sse_context = MagicMock() + sse_context.__enter__.return_value = (Mock(), Mock()) + sse_context.__exit__.side_effect = lambda *_args: events.append("sse_closed") + mock_sse_client.return_value = sse_context + + mcp_context = MagicMock() + mcp_context.__enter__.side_effect = lambda *_args: events.append("mcp_opened") or (Mock(), Mock(), Mock()) + mock_streamable_client.return_value = mcp_context + + failed_session_context = MagicMock() + failed_session = Mock() + failed_session.initialize.side_effect = MCPConnectionError("SSE handshake failed") + failed_session_context.__enter__.return_value = failed_session + failed_session_context.__exit__.side_effect = lambda *_args: events.append("session_closed") + + live_session_context = MagicMock() + live_session = Mock() + live_session_context.__enter__.return_value = live_session + mock_client_session.side_effect = [failed_session_context, live_session_context] + + client = MCPClient(server_url="http://test.example.com/unknown") + client._initialize() + + assert events[:3] == ["session_closed", "sse_closed", "mcp_opened"] + failed_session_context.__exit__.assert_called_once() + sse_context.__exit__.assert_called_once() + assert client._session is live_session + @patch("core.mcp.client.sse_client.ssrf_proxy_sse_connect") @patch("core.mcp.client.sse_client.create_ssrf_proxy_mcp_http_client") @patch("core.mcp.mcp_client.streamablehttp_client")