fix(mcp): auto cleanup on exceptions occurring in __enter__

Author: dbschmigelski

src/strands/models/bedrock.py

@@ -47,6 +47,7 @@
 
 DEFAULT_READ_TIMEOUT = 120
 
+
 class BedrockModel(Model):
     """AWS Bedrock model provider implementation.
 

src/strands/tools/mcp/mcp_client.py

@@ -83,11 +83,15 @@ def __init__(self, transport_callable: Callable[[], MCPTransport], *, startup_ti
         self._transport_callable = transport_callable
 
         self._background_thread: threading.Thread | None = None
-        self._background_thread_session: ClientSession
-        self._background_thread_event_loop: AbstractEventLoop
+        self._background_thread_session: ClientSession | None = None
+        self._background_thread_event_loop: AbstractEventLoop | None = None
 
     def __enter__(self) -> "MCPClient":
-        """Context manager entry point which initializes the MCP server connection."""
+        """Context manager entry point which initializes the MCP server connection.
+
+        TODO: Refactor to lazy initialization pattern following idiomatic Python.
+        Heavy work in __enter__ is non-idiomatic - should move connection logic to first method call instead.
+        """
         return self.start()
 
     def __exit__(self, exc_type: BaseException, exc_val: BaseException, exc_tb: TracebackType) -> None:
@@ -118,9 +122,15 @@ def start(self) -> "MCPClient":
             self._init_future.result(timeout=self._startup_timeout)
             self._log_debug_with_thread("the client initialization was successful")
         except futures.TimeoutError as e:
-            raise MCPClientInitializationError("background thread did not start in 30 seconds") from e
+            # Pass None for exc_type, exc_val, exc_tb since this isn't a context manager exit
+            self.stop(None, None, None)
+            raise MCPClientInitializationError(
+                f"background thread did not start in {self._startup_timeout} seconds"
+            ) from e
         except Exception as e:
             logger.exception("client failed to initialize")
+            # Pass None for exc_type, exc_val, exc_tb since this isn't a context manager exit
+            self.stop(None, None, None)
             raise MCPClientInitializationError("the client initialization failed") from e
         return self
 
@@ -129,21 +139,29 @@ def stop(
     ) -> None:
         """Signals the background thread to stop and waits for it to complete, ensuring proper cleanup of all resources.
 
+        This method is defensive and can handle partial initialization states that may occur
+        if start() fails partway through initialization.
+
         Args:
             exc_type: Exception type if an exception was raised in the context
             exc_val: Exception value if an exception was raised in the context
             exc_tb: Exception traceback if an exception was raised in the context
         """
         self._log_debug_with_thread("exiting MCPClient context")
 
-        async def _set_close_event() -> None:
-            self._close_event.set()
-
-        self._invoke_on_background_thread(_set_close_event()).result()
-        self._log_debug_with_thread("waiting for background thread to join")
+        # Only try to signal close event if we have a background thread
         if self._background_thread is not None:
+            # Signal close event if event loop exists
+            if self._background_thread_event_loop is not None:
+
+                async def _set_close_event() -> None:
+                    self._close_event.set()
+
+                asyncio.run_coroutine_threadsafe(_set_close_event(), self._background_thread_event_loop)
+
+            self._log_debug_with_thread("waiting for background thread to join")
             self._background_thread.join()
-        self._log_debug_with_thread("background thread joined, MCPClient context exited")
+            self._log_debug_with_thread("background thread joined, MCPClient context exited")
 
         # Reset fields to allow instance reuse
         self._init_future = futures.Future()
@@ -165,6 +183,7 @@ def list_tools_sync(self, pagination_token: Optional[str] = None) -> PaginatedLi
             raise MCPClientInitializationError(CLIENT_SESSION_NOT_RUNNING_ERROR_MESSAGE)
 
         async def _list_tools_async() -> ListToolsResult:
+            assert self._background_thread_session is not None
             return await self._background_thread_session.list_tools(cursor=pagination_token)
 
         list_tools_response: ListToolsResult = self._invoke_on_background_thread(_list_tools_async()).result()
@@ -191,6 +210,7 @@ def list_prompts_sync(self, pagination_token: Optional[str] = None) -> ListPromp
             raise MCPClientInitializationError(CLIENT_SESSION_NOT_RUNNING_ERROR_MESSAGE)
 
         async def _list_prompts_async() -> ListPromptsResult:
+            assert self._background_thread_session is not None
             return await self._background_thread_session.list_prompts(cursor=pagination_token)
 
         list_prompts_result: ListPromptsResult = self._invoke_on_background_thread(_list_prompts_async()).result()
@@ -215,6 +235,7 @@ def get_prompt_sync(self, prompt_id: str, args: dict[str, Any]) -> GetPromptResu
             raise MCPClientInitializationError(CLIENT_SESSION_NOT_RUNNING_ERROR_MESSAGE)
 
         async def _get_prompt_async() -> GetPromptResult:
+            assert self._background_thread_session is not None
             return await self._background_thread_session.get_prompt(prompt_id, arguments=args)
 
         get_prompt_result: GetPromptResult = self._invoke_on_background_thread(_get_prompt_async()).result()
@@ -250,6 +271,7 @@ def call_tool_sync(
             raise MCPClientInitializationError(CLIENT_SESSION_NOT_RUNNING_ERROR_MESSAGE)
 
         async def _call_tool_async() -> MCPCallToolResult:
+            assert self._background_thread_session is not None
             return await self._background_thread_session.call_tool(name, arguments, read_timeout_seconds)
 
         try:
@@ -285,6 +307,7 @@ async def call_tool_async(
             raise MCPClientInitializationError(CLIENT_SESSION_NOT_RUNNING_ERROR_MESSAGE)
 
         async def _call_tool_async() -> MCPCallToolResult:
+            assert self._background_thread_session is not None
             return await self._background_thread_session.call_tool(name, arguments, read_timeout_seconds)
 
         try:

tests/strands/tools/mcp/test_mcp_client.py

@@ -337,8 +337,12 @@ def test_enter_with_initialization_exception(mock_transport):
 
     client = MCPClient(mock_transport["transport_callable"])
 
-    with pytest.raises(MCPClientInitializationError, match="the client initialization failed"):
-        client.start()
+    with patch.object(client, "stop") as mock_stop:
+        with pytest.raises(MCPClientInitializationError, match="the client initialization failed"):
+            client.start()
+
+        # Verify stop() was called for cleanup
+        mock_stop.assert_called_once_with(None, None, None)
 
 
 def test_mcp_tool_result_type():
@@ -466,3 +470,18 @@ def test_get_prompt_sync_session_not_active():
 
     with pytest.raises(MCPClientInitializationError, match="client session is not running"):
         client.get_prompt_sync("test_prompt_id", {})
+
+
+def test_timeout_initialization_cleanup():
+    """Test that timeout during initialization properly cleans up."""
+
+    def slow_transport():
+        time.sleep(5)
+        return MagicMock()
+
+    client = MCPClient(slow_transport, startup_timeout=1)
+
+    with patch.object(client, "stop") as mock_stop:
+        with pytest.raises(MCPClientInitializationError, match="background thread did not start in 1 seconds"):
+            client.start()
+        mock_stop.assert_called_once_with(None, None, None)

tests_integ/test_mcp_client.py

@@ -15,6 +15,7 @@
 from strands.tools.mcp.mcp_client import MCPClient
 from strands.tools.mcp.mcp_types import MCPTransport
 from strands.types.content import Message
+from strands.types.exceptions import MCPClientInitializationError
 from strands.types.tools import ToolUse
 
 
@@ -268,3 +269,31 @@ def transport_callback() -> MCPTransport:
 
 def _messages_to_content_blocks(messages: List[Message]) -> List[ToolUse]:
     return [block["toolUse"] for message in messages for block in message["content"] if "toolUse" in block]
+
+
+def test_mcp_client_timeout_integration():
+    """Integration test for timeout scenario that caused hanging."""
+    import threading
+
+    from mcp import StdioServerParameters, stdio_client
+
+    def slow_transport():
+        time.sleep(4)  # Longer than timeout
+        return stdio_client(StdioServerParameters(command="python", args=["tests_integ/echo_server.py"]))
+
+    client = MCPClient(slow_transport, startup_timeout=2)
+    initial_threads = threading.active_count()
+
+    # First attempt should timeout
+    with pytest.raises(MCPClientInitializationError, match="background thread did not start in 2 seconds"):
+        with client:
+            pass
+
+    time.sleep(1)  # Allow cleanup
+    assert threading.active_count() == initial_threads  # No thread leak
+
+    # Should be able to recover by increasing timeout
+    client._startup_timeout = 60
+    with client:
+        tools = client.list_tools_sync()
+        assert len(tools) >= 0  # Should work now