Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion src/strands/tools/mcp/mcp_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,18 @@ def stop(
This method is defensive and can handle partial initialization states that may occur
if start() fails partway through initialization.

Resources to cleanup:
- _background_thread: Thread running the async event loop
- _background_thread_session: MCP ClientSession (auto-closed by context manager)
- _background_thread_event_loop: AsyncIO event loop in background thread
- _close_event: AsyncIO event to signal thread shutdown
- _init_future: Future for initialization synchronization

Cleanup order:
1. Signal close event to background thread (if session initialized)
2. Wait for background thread to complete
3. Reset all state for reuse

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
Expand All @@ -158,7 +170,9 @@ def stop(
async def _set_close_event() -> None:
self._close_event.set()

self._invoke_on_background_thread(_set_close_event()).result()
# Not calling _invoke_on_background_thread since the session does not need to exist
# we only need the thread and event loop to exist.
asyncio.run_coroutine_threadsafe(coro=_set_close_event(), loop=self._background_thread_event_loop)

self._log_debug_with_thread("waiting for background thread to join")
self._background_thread.join()
Expand All @@ -168,6 +182,8 @@ async def _set_close_event() -> None:
self._init_future = futures.Future()
self._close_event = asyncio.Event()
self._background_thread = None
self._background_thread_session = None
self._background_thread_event_loop = None
self._session_id = uuid.uuid4()

def list_tools_sync(self, pagination_token: Optional[str] = None) -> PaginatedList[MCPAgentTool]:
Expand Down
19 changes: 19 additions & 0 deletions tests/strands/tools/mcp/test_mcp_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -521,3 +521,22 @@ def test_stop_with_background_thread_but_no_event_loop():

# Verify cleanup occurred
assert client._background_thread is None


def test_mcp_client_state_reset_after_timeout():
"""Test that all client state is properly reset after timeout."""
def slow_transport():
time.sleep(4) # Longer than timeout
return MagicMock()

client = MCPClient(slow_transport, startup_timeout=2)

# First attempt should timeout
with pytest.raises(MCPClientInitializationError, match="background thread did not start in 2 seconds"):
client.start()

# Verify all state is reset
assert client._background_thread is None
assert client._background_thread_session is None
assert client._background_thread_event_loop is None
assert not client._init_future.done() # New future created
1 change: 1 addition & 0 deletions tests_integ/test_mcp_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -297,3 +297,4 @@ def slow_transport():
with client:
tools = client.list_tools_sync()
assert len(tools) >= 0 # Should work now

Loading