fix(mcp): accept integer progressToken in host progress capture (#32402)

This commit is contained in:
tin-berri 2026-07-07 21:46:36 -07:00 committed by GitHub
parent 7cc660866a
commit f922be32f0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 43 additions and 2 deletions

View file

@ -716,7 +716,7 @@ if MCP_AVAILABLE:
if not (host_ctx and hasattr(host_ctx, "meta") and host_ctx.meta):
return None
host_token = getattr(host_ctx.meta, "progressToken", None)
if not (host_token and hasattr(host_ctx, "session") and host_ctx.session):
if host_token is None or not (hasattr(host_ctx, "session") and host_ctx.session):
return None
host_session = host_ctx.session
@ -732,7 +732,7 @@ if MCP_AVAILABLE:
except Exception as e:
verbose_logger.error(f"Failed to forward progress to Host: {e}")
verbose_logger.debug(f"Host progressToken captured: {host_token[:8]}...")
verbose_logger.debug(f"Host progressToken captured: {str(host_token)[:8]}...")
return forward_progress
async def _build_virtual_call_logging_obj(

View file

@ -789,6 +789,47 @@ class TestCaptureHostProgressCallback:
host.request_context.session = MagicMock()
assert callable(_capture_host_progress_callback(host))
def test_returns_callable_when_token_is_integer(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
host = MagicMock()
host.request_context.meta.progressToken = 12345
host.request_context.session = MagicMock()
assert callable(_capture_host_progress_callback(host))
def test_returns_callable_when_token_is_zero(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
host = MagicMock()
host.request_context.meta.progressToken = 0
host.request_context.session = MagicMock()
assert callable(_capture_host_progress_callback(host))
@pytest.mark.asyncio
async def test_forwarded_progress_token_preserves_integer_value(self) -> None:
from litellm.proxy._experimental.mcp_server.server import (
_capture_host_progress_callback,
)
host = MagicMock()
host.request_context.meta.progressToken = 12345
session = AsyncMock()
host.request_context.session = session
callback = _capture_host_progress_callback(host)
assert callback is not None
await callback(0.5, 1.0)
session.send_progress_notification.assert_awaited_once_with(
progress_token=12345,
progress=0.5,
total=1.0,
)
class TestHandleListToolsVirtual:
"""Covers the protocol list_tools early-return when the flag is enabled."""