mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
fix(mcp): re-raise cancellation on py3.10 where Task.cancelling() is absent
This commit is contained in:
parent
ac3d072561
commit
90c8d7ba85
2 changed files with 57 additions and 9 deletions
|
|
@ -2540,16 +2540,29 @@ class MCPServerManager:
|
|||
status_code=status_code, server_name=server_name
|
||||
) from e
|
||||
|
||||
# No recoverable HTTP status. Respect a genuine *outer* cancellation
|
||||
# (shutdown / client disconnect) so cooperative cancellation still
|
||||
# works; only treat an internally-originated cancel as a failure.
|
||||
# No recoverable HTTP status. We must not launder this into an empty
|
||||
# success, but we also must not swallow a genuine *outer*
|
||||
# cancellation (shutdown / client disconnect) — re-raising those
|
||||
# keeps cooperative cancellation working.
|
||||
#
|
||||
# On Python 3.11+ `asyncio.Task.cancelling()` lets us tell an
|
||||
# internally-originated cancel-scope cancellation (the upstream
|
||||
# failure we want to surface) from a real outer cancel (which we
|
||||
# re-raise). On Python 3.10 that API does not exist, so we cannot
|
||||
# classify the cancellation; there we conservatively re-raise to
|
||||
# preserve cancellation semantics rather than risk converting a real
|
||||
# shutdown cancel into a server error. (The trade-off: an
|
||||
# upstream-induced bare cancellation on 3.10 propagates as an error
|
||||
# instead of a clean per-server failure — still never a false
|
||||
# success.)
|
||||
task = asyncio.current_task()
|
||||
is_outer_cancel = bool(
|
||||
task is not None
|
||||
and getattr(task, "cancelling", None) is not None
|
||||
and task.cancelling() > 0
|
||||
)
|
||||
if is_outer_cancel:
|
||||
cancelling = getattr(task, "cancelling", None) if task is not None else None
|
||||
if cancelling is None:
|
||||
# Python 3.10 (or no running task): cannot classify — cooperate
|
||||
# with the cancellation.
|
||||
raise
|
||||
if cancelling() > 0:
|
||||
# A real outer cancellation was requested on this task.
|
||||
raise
|
||||
verbose_logger.warning(
|
||||
f"Tool listing from {server_name} was interrupted before "
|
||||
|
|
|
|||
|
|
@ -278,3 +278,38 @@ async def test_fetch_tools_surfaces_error_on_bare_cancellation():
|
|||
await manager._fetch_tools_with_timeout(mock_client, server.name, server=server)
|
||||
assert exc_info.value.status_code is None
|
||||
mock_client.list_tools.assert_awaited_with(raise_on_error=False)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fetch_tools_reraises_genuine_outer_cancellation():
|
||||
"""A real outer cancellation (shutdown / client disconnect) must propagate,
|
||||
not be converted into MCPUpstreamError — otherwise cooperative cancellation
|
||||
breaks. On 3.11+ this is detected via Task.cancelling(); on 3.10 the
|
||||
classifier is unavailable and we re-raise unconditionally, so this contract
|
||||
holds on every supported runtime."""
|
||||
manager = MCPServerManager()
|
||||
server = MCPServer(
|
||||
server_id="o1",
|
||||
name="mock_databricks",
|
||||
url="https://upstream/mcp",
|
||||
transport=MCPTransport.http,
|
||||
auth_type=MCPAuth.oauth2,
|
||||
)
|
||||
|
||||
started = asyncio.Event()
|
||||
|
||||
async def _hang(*args, **kwargs):
|
||||
started.set()
|
||||
await asyncio.sleep(3600)
|
||||
|
||||
mock_client = MagicMock()
|
||||
mock_client.list_tools = _hang
|
||||
|
||||
task = asyncio.create_task(
|
||||
manager._fetch_tools_with_timeout(mock_client, server.name, server=server)
|
||||
)
|
||||
await started.wait()
|
||||
task.cancel()
|
||||
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue