mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(mcp): accept integer progressToken in host progress capture (#32402)
This commit is contained in:
parent
7cc660866a
commit
f922be32f0
2 changed files with 43 additions and 2 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue