scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-23 00:41:40 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 81

fix(mcp): admit an allowlisted initialize that fills the peek cap exactly and type the allowlist test helpers

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
yassin 2026-09-17 22:22:19 +00:00
parent fb99e3dde3
commit 5f6702ee47
2 changed files with 62 additions and 15 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/_experimental/mcp_server/server.py
View file

@ -3881,6 +3881,7 @@ if MCP_AVAILABLE:
async def _read_request_body_for_routing(
receive: Receive,
peek_max_bytes: int = _MCP_ROUTING_PEEK_MAX_BYTES,
settle_truncation: bool = False,
) -> tuple[list[Message], bytes, bool]:
"""
Read just enough of the request body to decide whether this is a
@ -3889,8 +3890,11 @@ if MCP_AVAILABLE:
the peeked body bytes (capped at ``peek_max_bytes``), plus whether
the body was cut off at that cap.
Stops reading from the wire as soon as either (a) we have peeked
``peek_max_bytes`` of body, or (b) the body is complete.
Stops reading from the wire as soon as either (a) the peek budget is
spent, or (b) the body is complete. A body that fills the budget
exactly in a frame with ``more_body`` is reported as cut off unless
``settle_truncation`` is set, in which case one more frame is read to
find out whether anything actually follows.
The remainder of an oversized body is streamed lazily through
``wrapped_receive`` in the caller — so an authenticated client cannot
force the proxy to buffer an arbitrarily large payload just to make a
@ -3922,12 +3926,9 @@ if MCP_AVAILABLE:
peeked_bytes += min(len(body), remaining)
truncated = truncated or len(body) > remaining
if not message.get("more_body", False):
if truncated or not message.get("more_body", False):
break
if peeked_bytes >= peek_max_bytes:
# Stop draining; downstream replay will pull remaining chunks
# directly from the original `receive` via wrapped_receive.
if not settle_truncation and peeked_bytes >= peek_max_bytes:
truncated = True
break
@ -4583,7 +4584,7 @@ if MCP_AVAILABLE:
if scope.get("method") == "POST":
allowed_clients: Final = _load_allowed_mcp_clients()
consumed_messages, body, body_truncated = await _read_request_body_for_routing(
receive, _routing_peek_limit(allowed_clients)
receive, _routing_peek_limit(allowed_clients), settle_truncation=allowed_clients is not None
)
if allowed_clients is not None and body_truncated and not session_id:
await _reject_oversized_unidentified_request(scope, receive, send, _client_ip)
@ -4863,11 +4864,11 @@ if MCP_AVAILABLE:
sse_allowed_clients: Final = _load_allowed_mcp_clients()
sse_consumed_messages, sse_body, sse_truncated = (
await _read_request_body_for_routing(receive, _routing_peek_limit(sse_allowed_clients))
if scope.get("method") == "POST"
await _read_request_body_for_routing(receive, MCP_ALLOWLIST_PEEK_MAX_BYTES, settle_truncation=True)
if sse_allowed_clients is not None and scope.get("method") == "POST"
else ((), b"", False)
)
if sse_allowed_clients is not None and sse_truncated:
if sse_truncated:
await _reject_oversized_unidentified_request(scope, receive, send, _sse_client_ip)
return
if _is_initialize_request(sse_body) and await _reject_initialize_from_disallowed_client(

54
tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server.py
View file

@ -3,6 +3,7 @@ import contextlib
import contextvars
import json
import os
from collections.abc import Iterator
from datetime import datetime, timedelta
from types import SimpleNamespace
from typing import Final
@ -19,6 +20,7 @@ from mcp.types import (
TextContent,
TextResourceContents,
)
from starlette.types import Receive, Scope, Send
from litellm.constants import MCP_ALLOWLIST_PEEK_MAX_BYTES
from litellm.proxy._types import (
@ -2050,7 +2052,7 @@ _ANONYMOUS_INITIALIZE: Final = b'{"jsonrpc":"2.0","id":0,"method":"initialize","
_TOOLS_LIST: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}'
async def _drain_body(receive) -> bytes:
async def _drain_body(receive: Receive) -> bytes:
chunks: list[bytes] = []
while True:
message = await receive()
@ -2066,7 +2068,7 @@ def _forbidden_client_response(send: AsyncMock) -> tuple[int, dict[str, str]]:
@contextlib.contextmanager
def _client_allowlist_patches(allowed_clients: object):
def _client_allowlist_patches(allowed_clients: list[str] | list[dict[str, str]] | str | None) -> Iterator[None]:
settings: Final = {} if allowed_clients is None else {"mcp_allowed_clients": allowed_clients}
with (
patch( # test-quality-ok: the ASGI handler resolves auth through a module-level function; no injection seam
@ -2188,7 +2190,9 @@ async def test_streamable_http_admits_listed_or_unrestricted_initialize_and_repl
@pytest.mark.asyncio
@pytest.mark.parametrize("allowed_clients", ([], "claude-code", [{"name": "claude-code"}]))
async def test_streamable_http_empty_or_malformed_allowlist_admits_nobody(allowed_clients: object) -> None:
async def test_streamable_http_empty_or_malformed_allowlist_admits_nobody(
allowed_clients: list[str] | list[dict[str, str]] | str,
) -> None:
from starlette.types import Scope
from litellm.proxy._experimental.mcp_server import server as mcp_module
@ -2243,6 +2247,10 @@ async def test_streamable_http_allowlist_only_inspects_initialize_requests() ->
def _oversized_initialize(client_name: str, peek_cap: int) -> bytes:
return _padded_initialize(client_name, peek_cap * 2)
def _padded_initialize(client_name: str, padding: int) -> bytes:
return json.dumps(
{
"jsonrpc": "2.0",
@ -2250,7 +2258,7 @@ def _oversized_initialize(client_name: str, peek_cap: int) -> bytes:
"method": "initialize",
"params": {
"protocolVersion": "2025-06-18",
"capabilities": {"experimental": {"padding": "x" * (peek_cap * 2)}},
"capabilities": {"experimental": {"padding": "x" * padding}},
"clientInfo": {"name": client_name, "version": "1.0.0"},
},
}
@ -2349,6 +2357,44 @@ async def test_streamable_http_allowlist_bounds_the_body_it_buffers_for_unidenti
stateless_handle.assert_not_awaited()
@pytest.mark.asyncio
async def test_streamable_http_allowlist_admits_an_initialize_of_exactly_the_peek_cap() -> None:
from litellm.proxy._experimental.mcp_server import server as mcp_module
request_body: Final = _padded_initialize(
"antigravity-cli", MCP_ALLOWLIST_PEEK_MAX_BYTES - len(_padded_initialize("antigravity-cli", 0))
)
assert len(request_body) == MCP_ALLOWLIST_PEEK_MAX_BYTES
scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
receive: Final = AsyncMock(
side_effect=[
{"type": "http.request", "body": request_body, "more_body": True},
{"type": "http.request", "body": b"", "more_body": False},
]
)
send: Final = AsyncMock()
downstream_bodies: Final[list[bytes]] = []
async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None:
downstream_bodies.append(await _drain_body(downstream_receive))
with (
_client_allowlist_patches(["antigravity-cli"]),
patch( # test-quality-ok: session managers are module singletons; the downstream call is the observable
"litellm.proxy._experimental.mcp_server.server.session_manager_stateful",
SimpleNamespace(handle_request=AsyncMock(side_effect=handle_request)),
),
patch( # test-quality-ok: session managers are module singletons; the downstream call is the observable
"litellm.proxy._experimental.mcp_server.server.session_manager_stateless",
SimpleNamespace(handle_request=AsyncMock(side_effect=handle_request)),
),
):
await mcp_module.handle_streamable_http_mcp(scope, receive, send)
assert downstream_bodies == [request_body]
send.assert_not_awaited()
@pytest.mark.asyncio
async def test_streamable_http_allowlist_streams_large_posts_on_an_admitted_session() -> None:
from starlette.types import Receive, Scope, Send