mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
fix(mcp): keep the streamable-HTTP routing peek on a UTF-8 boundary
Fixes https://github.com/BerriAI/litellm/issues/34917
This commit is contained in:
parent
daf22ec871
commit
6bbcf1dbbf
2 changed files with 125 additions and 3 deletions
|
|
@ -237,6 +237,24 @@ def _jsonrpc_text_has_top_level_method(text: str) -> bool:
|
|||
return False
|
||||
|
||||
|
||||
def _utf8_boundary_prefix(data: bytes) -> bytes:
|
||||
"""``data`` with any trailing incomplete UTF-8 sequence removed.
|
||||
|
||||
Cutting a body at a fixed byte budget can land in the middle of a multibyte
|
||||
character, and ``json.loads`` on such bytes raises ``UnicodeDecodeError``
|
||||
rather than ``JSONDecodeError``. Trimming to a character boundary keeps the
|
||||
truncated peek decodable so callers only have to handle malformed JSON.
|
||||
"""
|
||||
for trailing in range(0, min(3, len(data)) + 1):
|
||||
candidate = data[: len(data) - trailing]
|
||||
try:
|
||||
candidate.decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
continue
|
||||
return candidate
|
||||
return data
|
||||
|
||||
|
||||
def _mcp_meta_trace_carrier(req_ctx: object) -> dict[str, str] | None:
|
||||
"""The W3C trace context (``traceparent``/``tracestate``) the MCP client
|
||||
propagated in the request's ``params._meta`` (SEP-414), or ``None``.
|
||||
|
|
@ -3411,7 +3429,7 @@ if MCP_AVAILABLE:
|
|||
try:
|
||||
data = json.loads(body)
|
||||
return isinstance(data, dict) and data.get("method") == "initialize"
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
|
||||
return False
|
||||
|
||||
async def _read_request_body_for_routing(
|
||||
|
|
@ -3462,7 +3480,7 @@ if MCP_AVAILABLE:
|
|||
# directly from the original `receive` via wrapped_receive.
|
||||
break
|
||||
|
||||
return consumed_messages, b"".join(body_chunks)
|
||||
return consumed_messages, _utf8_boundary_prefix(b"".join(body_chunks))
|
||||
|
||||
async def _handle_stale_mcp_session(
|
||||
scope: Scope,
|
||||
|
|
@ -4227,7 +4245,7 @@ if MCP_AVAILABLE:
|
|||
"MCP: detected JSON-RPC response POST (id=%s), skipping session lock to avoid deadlock",
|
||||
_peeked.get("id"),
|
||||
)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
except (json.JSONDecodeError, UnicodeDecodeError, TypeError):
|
||||
# Peek cap truncated the body, so it can't be fully parsed.
|
||||
# Scan the top-level keys (depth-aware) instead of a flat
|
||||
# substring search: a response's result payload may nest a
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
import asyncio
|
||||
import contextvars
|
||||
import json
|
||||
from datetime import datetime, timedelta
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -1689,6 +1690,109 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body():
|
|||
assert total_streamed == len(first_chunk) + sum(len(b) for b in oversized_tail)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap():
|
||||
"""
|
||||
A tool-call POST whose UTF-8 body is larger than the routing peek cap, with a
|
||||
multibyte character straddling the cap boundary, must still be forwarded
|
||||
intact instead of blowing up with a UnicodeDecodeError 500.
|
||||
|
||||
Regression test for https://github.com/BerriAI/litellm/issues/34917
|
||||
"""
|
||||
try:
|
||||
from litellm.proxy._experimental.mcp_server import server as mcp_server
|
||||
from litellm.proxy._experimental.mcp_server.server import (
|
||||
handle_streamable_http_mcp,
|
||||
session_manager_stateful,
|
||||
session_manager_stateless,
|
||||
)
|
||||
except ImportError:
|
||||
pytest.skip("MCP server not available")
|
||||
|
||||
peek_cap = mcp_server._MCP_ROUTING_PEEK_MAX_BYTES
|
||||
|
||||
def _splits_multibyte_at_cap(candidate: bytes) -> bool:
|
||||
try:
|
||||
candidate[:peek_cap].decode("utf-8")
|
||||
except UnicodeDecodeError:
|
||||
return True
|
||||
return False
|
||||
|
||||
def _build_body() -> bytes:
|
||||
for pad in range(4):
|
||||
candidate = json.dumps(
|
||||
{
|
||||
"jsonrpc": "2.0",
|
||||
"id": 1,
|
||||
"method": "tools/call",
|
||||
"params": {
|
||||
"name": "update_full_document" + "x" * pad,
|
||||
"arguments": {"markdown": "щ" * 3000},
|
||||
},
|
||||
},
|
||||
ensure_ascii=False,
|
||||
).encode("utf-8")
|
||||
if len(candidate) > peek_cap and _splits_multibyte_at_cap(candidate):
|
||||
return candidate
|
||||
raise AssertionError("could not build a body splitting a multibyte char at the peek cap")
|
||||
|
||||
body = _build_body()
|
||||
|
||||
messages = [{"type": "http.request", "body": body, "more_body": False}]
|
||||
receive_calls = {"count": 0}
|
||||
|
||||
async def receive():
|
||||
idx = receive_calls["count"]
|
||||
receive_calls["count"] += 1
|
||||
return messages[idx]
|
||||
|
||||
scope = {
|
||||
"type": "http",
|
||||
"method": "POST",
|
||||
"path": "/mcp/progress_test",
|
||||
"headers": [
|
||||
(b"content-type", b"application/json"),
|
||||
(b"authorization", b"Bearer test-key"),
|
||||
],
|
||||
}
|
||||
send = AsyncMock()
|
||||
|
||||
streamed_chunks = []
|
||||
|
||||
async def stateless_handle(s, r, se):
|
||||
while True:
|
||||
msg = await r()
|
||||
if msg.get("type") != "http.request":
|
||||
break
|
||||
streamed_chunks.append(msg.get("body", b"") or b"")
|
||||
if not msg.get("more_body", False):
|
||||
break
|
||||
|
||||
async def stateful_handle(s, r, se):
|
||||
raise AssertionError("non-initialize POST should not reach stateful manager")
|
||||
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
|
||||
new_callable=AsyncMock,
|
||||
return_value=(MagicMock(), None, ["progress_test"], None, None, None),
|
||||
),
|
||||
patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"),
|
||||
patch(
|
||||
"litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
|
||||
True,
|
||||
),
|
||||
patch.object(session_manager_stateless, "handle_request", side_effect=stateless_handle),
|
||||
patch.object(session_manager_stateful, "handle_request", side_effect=stateful_handle),
|
||||
patch.object(session_manager_stateless, "_server_instances", {}),
|
||||
patch.object(session_manager_stateful, "_server_instances", {}),
|
||||
):
|
||||
await handle_streamable_http_mcp(scope, receive, send)
|
||||
|
||||
assert send.await_count == 0, f"unexpected response emitted by the proxy: {send.await_args_list}"
|
||||
assert b"".join(streamed_chunks) == body
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_enforce_stateful_session_cap_evicts_oldest_idle_then_rejects():
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue