From 5fe8415f7ccdda13e97243508664b290a8745b88 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 23 Sep 2026 15:56:13 +0000 Subject: [PATCH] fix(mcp): surface upstream session teardown as RuntimeError instead of cancelling callers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/experimental_mcp_client/client.py | 8 ++++--- .../test_mcp_client.py | 21 +++++++++++++++++++ 2 files changed, 26 insertions(+), 3 deletions(-) diff --git a/litellm/experimental_mcp_client/client.py b/litellm/experimental_mcp_client/client.py index 6eb00292649..e16d1599d0a 100644 --- a/litellm/experimental_mcp_client/client.py +++ b/litellm/experimental_mcp_client/client.py @@ -1222,8 +1222,10 @@ class PersistentMCPSession: self._active = future (outcome,) = await asyncio.gather(operation(session), return_exceptions=True) if not future.done(): - if isinstance(outcome, BaseException): + if isinstance(outcome, Exception): future.set_exception(outcome) + elif isinstance(outcome, BaseException): + future.set_exception(RuntimeError("upstream MCP operation was cancelled")) else: future.set_result(outcome) if isinstance(outcome, (ValueError, httpx2.HTTPError, OSError, MCPError)): @@ -1237,9 +1239,9 @@ class PersistentMCPSession: except asyncio.CancelledError: self._fail_waiters(None) raise - self._fail_waiters(ended if isinstance(ended, BaseException) else None) + self._fail_waiters(ended if isinstance(ended, Exception) else None) - def _fail_waiters(self, cause: BaseException | None) -> None: + def _fail_waiters(self, cause: Exception | None) -> None: pending: Final = (self._ready, self._active, *(future for _, future in self._drained())) for future in pending: if future is not None and not future.done(): 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 eb806763abd..5c3c24e2b07 100644 --- a/tests/test_litellm/experimental_mcp_client/test_mcp_client.py +++ b/tests/test_litellm/experimental_mcp_client/test_mcp_client.py @@ -3013,3 +3013,24 @@ async def test_closing_persistent_session_mid_operation_fails_the_waiter_instead with pytest.raises(RuntimeError, match="upstream MCP session closed"): await asyncio.wait_for(waiter, 5) await asyncio.wait_for(session.wait_closed(), 5) + + +@pytest.mark.asyncio +async def test_persistent_session_reports_an_upstream_cancellation_as_a_runtime_error_not_a_cancelled_caller(): + app: Final = _stateful_upstream() + async with app.router.lifespan_context(app): + client, session = _client_with_session(app) + + async def cancelled_upstream(_: object) -> str: + raise asyncio.CancelledError() + + try: + with pytest.raises(RuntimeError, match="cancelled"): + await asyncio.wait_for(session.run(cancelled_upstream), 5) + selected: Final = await client.call_tool( + CallToolRequestParams(name="select_project", arguments={"name": "e"}), persistent_session=session + ) + finally: + session.close() + assert selected.is_error is False, "an upstream cancellation must not be mistaken for a cancelled caller" + await asyncio.wait_for(session.wait_closed(), 5)