From 6bbcf1dbbf3d533ee1948b9b3183d8962ace4263 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 28 Jul 2026 12:08:08 +0000 Subject: [PATCH 1/3] fix(mcp): keep the streamable-HTTP routing peek on a UTF-8 boundary Fixes https://github.com/BerriAI/litellm/issues/34917 --- .../proxy/_experimental/mcp_server/server.py | 24 +++- .../mcp_server/test_mcp_server.py | 104 ++++++++++++++++++ 2 files changed, 125 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 14673cf12c1..af5da275961 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -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 diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index 1753b0d92a8..7f79e5aebda 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -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(): """ From a835e75620b1e3fb8c5e11d39c70fc73f17feb04 Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:24:19 -0700 Subject: [PATCH 2/3] fix(mcp): handle split UTF-8 routing previews in place --- .../proxy/_experimental/mcp_server/server.py | 20 +- .../mcp_server/test_mcp_server.py | 173 ++++++++---------- 2 files changed, 78 insertions(+), 115 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index e0fb7cede03..397a82cfa45 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -277,24 +277,6 @@ 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``. @@ -4057,7 +4039,7 @@ if MCP_AVAILABLE: # directly from the original `receive` via wrapped_receive. break - return consumed_messages, _utf8_boundary_prefix(b"".join(body_chunks)) + return consumed_messages, b"".join(body_chunks) async def _handle_stale_mcp_session( scope: Scope, diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py index aff3a2f0483..b2eded67430 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py @@ -1992,6 +1992,7 @@ async def test_streamable_http_session_manager_is_stateless(): ( ("POST", b'{"jsonrpc":"2.0","id":1,"method":"initialize","params":{}}', True), ("POST", b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}', False), + ("POST", b"", False), ("GET", b"", False), ("DELETE", b"", False), ), @@ -2466,106 +2467,65 @@ async def test_mcp_routing_caps_body_peek_for_oversized_chunked_body(): @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. +@pytest.mark.parametrize("method", ("initialize", "tools/call")) +@pytest.mark.parametrize("chunked", (False, True)) +@pytest.mark.parametrize( + ("character", "bytes_before_cap"), + (("é", 0), ("é", 1), ("中", 1), ("中", 2), ("😀", 1), ("😀", 2), ("😀", 3)), +) +async def test_mcp_routing_peek_survives_multibyte_char_split_at_cap( + method: str, chunked: bool, character: str, bytes_before_cap: int +) -> None: + from litellm.proxy._experimental.mcp_server import server as mcp_module - 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") + params: Final = ( + { + "protocolVersion": LATEST_HANDSHAKE_VERSION, + "capabilities": {}, + "clientInfo": {"name": "<>", "version": "1"}, + } + if method == "initialize" + else {"name": "update_full_document", "arguments": {"markdown": "<>"}} + ) + template: Final = json.dumps({"jsonrpc": "2.0", "id": 1, "method": method, "params": params}).encode() + prefix, suffix = template.split(b"<>") + cap: Final = mcp_module._MCP_ROUTING_PEEK_MAX_BYTES + body: Final = prefix + b"x" * (cap - bytes_before_cap - len(prefix)) + character.encode() + b"tail" + suffix + chunks: Final = (body[: cap - 1], body[cap - 1 : cap], body[cap:]) if chunked else (body,) + messages: Final[tuple[Message, ...]] = tuple( + {"type": "http.request", "body": chunk, "more_body": index < len(chunks) - 1} + for index, chunk in enumerate(chunks) + ) + receive: Final = AsyncMock(side_effect=messages) + send: Final = AsyncMock() + received: Final[asyncio.Future[bytes]] = asyncio.get_running_loop().create_future() - 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") + async def handle_request(_: Scope, downstream_receive: Receive, outgoing: Send) -> None: + assert receive.await_count == (2 if chunked else 1) + received.set_result(await _drain_body(downstream_receive)) + await outgoing({"type": "http.response.start", "status": 200, "headers": []}) + await outgoing({"type": "http.response.body", "body": b"{}"}) + stateless_handle: Final = AsyncMock(side_effect=handle_request) + stateful_handle: Final = AsyncMock() + scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []} with ( + _client_allowlist_patches({}, None), patch( - "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context", - new_callable=AsyncMock, - return_value=(MagicMock(), None, ["progress_test"], None, None, None), + "litellm.proxy._experimental.mcp_server.server.session_manager_stateless", + SimpleNamespace(handle_request=stateless_handle), ), - patch("litellm.proxy._experimental.mcp_server.server.set_auth_context"), patch( - "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", - True, + "litellm.proxy._experimental.mcp_server.server.session_manager_stateful", + SimpleNamespace(handle_request=stateful_handle), ), - 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) + await mcp_module.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 + assert send.call_args_list[0].args[0]["status"] == 200 + assert received.result() == body + stateless_handle.assert_awaited_once() + stateful_handle.assert_not_awaited() @pytest.mark.asyncio @@ -4119,7 +4079,12 @@ def test_jsonrpc_text_has_top_level_method_ignores_nested_method(): @pytest.mark.asyncio -async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): +@pytest.mark.parametrize("response_field", ("result", "error")) +@pytest.mark.parametrize(("character", "bytes_before_cap"), (("", 0), ("x", 0), ("é", 1), ("中", 2), ("😀", 3))) +@pytest.mark.parametrize("cancel_request", (False, True)) +async def test_truncated_jsonrpc_response_with_nested_method_skips_lock( + response_field: str, character: str, bytes_before_cap: int, cancel_request: bool +) -> None: """Regression: a large JSON-RPC *response* POST whose ``result`` payload nests a ``method`` key must skip the per-session lock so it does not deadlock behind the in-flight request POST that is holding the lock while @@ -4147,7 +4112,7 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): async def handle(s, r, se): msg = await r() body = msg.get("body", b"") or b"" - if b'"result"' in body: + if body == response_body: response_handled.set() else: request_in_handle.set() @@ -4174,9 +4139,16 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): # A JSON-RPC response larger than the routing peek cap so it can't be fully # parsed, with a nested "method" key in the first bytes to trip a flat # substring heuristic. - response_body = ( - '{"jsonrpc":"2.0","id":99,"result":{"toolResult":{"method":"GET","payload":"' + ("x" * 5000) + '"}}}' + response_prefix: Final = ( + '{"jsonrpc":"2.0","id":99,"' + response_field + + '":{"code":-32000,"message":"test","data":{"method":"GET","payload":"' ).encode() + response_body: Final = ( + response_prefix + + b"x" * (mcp_server._MCP_ROUTING_PEEK_MAX_BYTES - bytes_before_cap - len(response_prefix) if character else 0) + + character.encode() + + b'tail"}}}' + ) try: with ( @@ -4204,8 +4176,17 @@ async def test_truncated_jsonrpc_response_with_nested_method_skips_lock(): # lock held by req_task and this wait would time out (deadlock). await asyncio.wait_for(response_handled.wait(), timeout=1.0) - gate.set() - await asyncio.gather(req_task, resp_task) + await resp_task + assert not req_task.done() + if cancel_request: + req_task.cancel() + with pytest.raises(asyncio.CancelledError): + await req_task + else: + gate.set() + await req_task + assert not mcp_server._stateful_session_locks[session_id].locked() + assert session_id not in mcp_server._stateful_session_active_request_counts finally: gate.set() mcp_server._stateful_session_auth_contexts.pop(session_id, None) From 2da7e6dfe20febb05e7c00302f9fd29ccb24dc8f Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:31:18 -0700 Subject: [PATCH 3/3] ci: trigger MCP fix checks against main