Python: Propagate MCP initialization failures

python/semantic_kernel/connectors/mcp.py

@@ -7,7 +7,7 @@
 import sys
 from abc import abstractmethod
 from collections.abc import Awaitable, Callable, Sequence
-from contextlib import AbstractAsyncContextManager, AsyncExitStack, _AsyncGeneratorContextManager
+from contextlib import AbstractAsyncContextManager, AsyncExitStack, _AsyncGeneratorContextManager, suppress
 from datetime import timedelta
 from functools import partial
 from itertools import chain
@@ -300,12 +300,14 @@ async def connect(self) -> None:
         try:
             self._current_task = asyncio.create_task(self._inner_connect(ready_event))
             await ready_event.wait()
-        except KernelPluginInvalidConfigurationError:
-            ready_event.clear()
-            raise
+            if self._current_task.done():
+                await self._current_task
         except Exception as ex:
             ready_event.clear()
-            await self.close()
+            with suppress(Exception):
+                await self.close()
+            if isinstance(ex, KernelPluginInvalidConfigurationError):
+                raise
             raise FunctionExecutionException("Failed to enter context manager.") from ex
 
     async def close(self) -> None:
@@ -314,62 +316,76 @@ async def close(self) -> None:
             # Signal the stop event, which asks the _inner_connect
             # method to close the session with the exit stack
             self._stop_event.set()
-        if self._current_task:
-            # After, the signal, we wait for it to close the exit stack.
-            await self._current_task
+        try:
+            if self._current_task:
+                # After the signal, wait for the task to close the exit stack.
+                await self._current_task
+        finally:
             self._current_task = None
-        self.session = None
+            self.session = None
 
     async def _inner_connect(self, ready_event: asyncio.Event) -> None:
-        if not self.session:
-            try:
-                transport = await self._exit_stack.enter_async_context(self.get_mcp_client())
-            except Exception as ex:
-                await self._exit_stack.aclose()
-                ready_event.set()
-                raise KernelPluginInvalidConfigurationError(
-                    "Failed to connect to the MCP server. Please check your configuration."
-                ) from ex
-            try:
-                session = await self._exit_stack.enter_async_context(
-                    ClientSession(
-                        read_stream=transport[0],
-                        write_stream=transport[1],
-                        read_timeout_seconds=timedelta(seconds=self.request_timeout) if self.request_timeout else None,
-                        message_handler=self.message_handler,
-                        logging_callback=self.logging_callback,
-                        sampling_callback=self.sampling_callback,
+        try:
+            if not self.session:
+                try:
+                    transport = await self._exit_stack.enter_async_context(self.get_mcp_client())
+                except Exception as ex:
+                    raise KernelPluginInvalidConfigurationError(
+                        "Failed to connect to the MCP server. Please check your configuration."
+                    ) from ex
+                try:
+                    session = await self._exit_stack.enter_async_context(
+                        ClientSession(
+                            read_stream=transport[0],
+                            write_stream=transport[1],
+                            read_timeout_seconds=(
+                                timedelta(seconds=self.request_timeout) if self.request_timeout else None
+                            ),
+                            message_handler=self.message_handler,
+                            logging_callback=self.logging_callback,
+                            sampling_callback=self.sampling_callback,
+                        )
                     )
-                )
-            except Exception as ex:
-                await self._exit_stack.aclose()
-                raise KernelPluginInvalidConfigurationError(
-                    "Failed to create a session. Please check your configuration."
-                ) from ex
+                except Exception as ex:
+                    raise KernelPluginInvalidConfigurationError(
+                        "Failed to create a session. Please check your configuration."
+                    ) from ex
+                try:
+                    await session.initialize()
+                except Exception as ex:
+                    raise KernelPluginInvalidConfigurationError(
+                        "Failed to initialize session. Please check your configuration."
+                    ) from ex
+                self.session = session
+            elif self.session._request_id == 0:
+                # If the session is not initialized, we need to reinitialize it
+                try:
+                    await self.session.initialize()
+                except Exception as ex:
+                    raise KernelPluginInvalidConfigurationError(
+                        "Failed to initialize session. Please check your configuration."
+                    ) from ex
+            logger.debug("Connected to MCP server: %s", self.session)
+            if self.load_tools_flag:
+                await self.load_tools()
+            if self.load_prompts_flag:
+                await self.load_prompts()
+
+            if logger.level != logging.NOTSET:
+                try:
+                    await self.session.set_logging_level(
+                        next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level)
+                    )
+                except Exception:
+                    logger.warning("Failed to set log level to %s", logger.level)
+        except Exception:
             try:
-                await session.initialize()
-            except Exception as ex:
                 await self._exit_stack.aclose()
-                raise KernelPluginInvalidConfigurationError(
-                    "Failed to initialize session. Please check your configuration."
-                ) from ex
-            self.session = session
-        elif self.session._request_id == 0:
-            # If the session is not initialized, we need to reinitialize it
-            await self.session.initialize()
-        logger.debug("Connected to MCP server: %s", self.session)
-        if self.load_tools_flag:
-            await self.load_tools()
-        if self.load_prompts_flag:
-            await self.load_prompts()
-
-        if logger.level != logging.NOTSET:
-            try:
-                await self.session.set_logging_level(
-                    next(level for level, value in LOG_LEVEL_MAPPING.items() if value == logger.level)
-                )
             except Exception:
-                logger.warning("Failed to set log level to %s", logger.level)
+                logger.exception("Error closing exit stack after failed connect")
+            finally:
+                ready_event.set()
+            raise
         # Setting up is complete, will now signal the main loop that we are ready
         ready_event.set()
         # Create a stop event to signal the exit stack to close

python/tests/unit/connectors/mcp/test_mcp.py

@@ -1,5 +1,6 @@
 # Copyright (c) Microsoft. All rights reserved.
 
+import asyncio
 import logging
 import re
 from unittest.mock import AsyncMock, MagicMock, patch
@@ -17,7 +18,7 @@
     create_mcp_server_from_functions,
     create_mcp_server_from_kernel,
 )
-from semantic_kernel.exceptions import KernelPluginInvalidConfigurationError
+from semantic_kernel.exceptions import FunctionExecutionException, KernelPluginInvalidConfigurationError
 from semantic_kernel.functions import KernelFunction, KernelPlugin, kernel_function
 
 
@@ -255,6 +256,76 @@ async def test_mcp_plugin_failed_get_session():
                 pass
 
 
+@pytest.mark.parametrize("supplied_session", [False, True])
+@patch("semantic_kernel.connectors.mcp.streamablehttp_client")
+@patch("semantic_kernel.connectors.mcp.ClientSession")
+async def test_streamable_http_initialization_failure_unblocks_connect(mock_session, mock_client, supplied_session):
+    mock_read = MagicMock()
+    mock_write = MagicMock()
+    mock_callback = MagicMock()
+
+    mock_generator = MagicMock()
+    mock_generator.__aenter__.return_value = (mock_read, mock_write, mock_callback)
+    mock_generator.__aexit__.return_value = (mock_read, mock_write, mock_callback)
+    mock_client.return_value = mock_generator
+
+    session = mock_session.return_value.__aenter__.return_value
+    session._request_id = 0
+    initialization_error = RuntimeError("Unauthorized")
+    session.initialize.side_effect = initialization_error
+
+    plugin = MCPStreamableHttpPlugin(
+        name="test",
+        url="http://localhost:8080/mcp",
+        session=session if supplied_session else None,
+        load_tools=False,
+        load_prompts=False,
+    )
+
+    with pytest.raises(KernelPluginInvalidConfigurationError, match="Failed to initialize session") as exc_info:
+        await asyncio.wait_for(plugin.connect(), timeout=1)
+
+    assert exc_info.value.__cause__ is initialization_error
+    assert plugin.session is None
+
+
+@pytest.mark.parametrize("failed_loader", ["load_tools", "load_prompts"])
+@patch("semantic_kernel.connectors.mcp.streamablehttp_client")
+@patch("semantic_kernel.connectors.mcp.ClientSession")
+async def test_loader_failure_closes_contexts_and_allows_reconnect(mock_session, mock_client, failed_loader):
+    transport = MagicMock()
+    transport.__aenter__.return_value = (MagicMock(), MagicMock(), MagicMock())
+    mock_client.return_value = transport
+    session_context = mock_session.return_value
+
+    async def exit_after_yield(*args):
+        await asyncio.sleep(0)
+        return False
+
+    session_context.__aexit__.side_effect = exit_after_yield
+    plugin = MCPStreamableHttpPlugin(
+        name="test",
+        url="http://localhost:8080/mcp",
+        load_tools=failed_loader == "load_tools",
+        load_prompts=failed_loader == "load_prompts",
+    )
+    setattr(plugin, failed_loader, AsyncMock(side_effect=RuntimeError("load failed")))
+
+    with pytest.raises(FunctionExecutionException, match="Failed to enter context manager"):
+        await asyncio.wait_for(plugin.connect(), timeout=1)
+
+    session_context.__aexit__.assert_awaited_once()
+    transport.__aexit__.assert_awaited_once()
+    assert plugin.session is None
+
+    setattr(plugin, failed_loader, AsyncMock())
+    await asyncio.wait_for(plugin.connect(), timeout=1)
+    assert plugin.session is session_context.__aenter__.return_value
+    await asyncio.wait_for(plugin.close(), timeout=1)
+    assert session_context.__aexit__.await_count == 2
+    assert transport.__aexit__.await_count == 2
+
+
 @patch("semantic_kernel.connectors.mcp.stdio_client")
 @patch("semantic_kernel.connectors.mcp.ClientSession")
 async def test_with_kwargs_stdio(mock_session, mock_client, list_tool_calls, kernel: "Kernel"):