scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-30 01:52:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 70

fix(mcp): surface upstream session teardown as RuntimeError instead of cancelling callers

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-23 15:56:13 +00:00
parent 9203232ae6
commit 5fe8415f7c
2 changed files with 26 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

8
litellm/experimental_mcp_client/client.py
View file

@ -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():

21
tests/test_litellm/experimental_mcp_client/test_mcp_client.py
View file

@ -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)