From 464c9b69a35b90096ec97c827f2f77d60aa482d5 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Tue, 15 Sep 2026 19:06:01 +0000 Subject: [PATCH] refactor(mcp): centralize peeked-body scope key and type peek replay Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../mcp_server/auth/user_api_key_auth_mcp.py | 5 ++-- .../proxy/_experimental/mcp_server/server.py | 25 +++++++++++-------- .../mcp_server/test_mcp_server.py | 4 +-- 4 files changed, 20 insertions(+), 15 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 09442d6151e..4a00ef83e5b 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -317,6 +317,7 @@ WEBSOCKET_CLOSE_REASON_MAX_BYTES: Final = 123 BEDROCK_REALTIME_PENDING_SESSION_UPDATE_SCOPE_KEY: Final = "litellm.bedrock_realtime.pending_session_update" BEDROCK_REALTIME_SESSION_COMMITTED_SCOPE_KEY: Final = "litellm.bedrock_realtime.session_committed" BEDROCK_REALTIME_COMMITTED_FAILURE_SCOPE_KEY: Final = "litellm.bedrock_realtime.committed_failure" +MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body" REALTIME_SESSION_SUCCESS_LOGGED_KEY: Final = "realtime_session_success_logged" REALTIME_SESSION_FAILURE_LOGGED_KEY: Final = "realtime_session_failure_logged" diff --git a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py index 93e41d63c5d..b6a82b4df78 100644 --- a/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py +++ b/litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py @@ -13,6 +13,7 @@ from typing_extensions import assert_never import litellm from litellm._logging import verbose_logger +from litellm.constants import MCP_PEEKED_BODY_SCOPE_KEY from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_passthrough_resource_metadata_url, get_passthrough_www_authenticate, @@ -67,8 +68,6 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body" - _EMPTY_TOOLSET_GRANTS: Final[Mapping[str, Sequence[str]]] = MappingProxyType({}) @@ -78,7 +77,7 @@ def _admission_request(scope: Scope) -> Request: request: Final = Request(scope=scope) peeked_body: Final[bytes] = scope.get(MCP_PEEKED_BODY_SCOPE_KEY, b"{}") - async def mock_body(): + async def mock_body() -> bytes: return peeked_body request.body = mock_body diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index b71c2acf91e..3aebdd063b4 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -26,14 +26,13 @@ from starlette.types import Message, Receive, Scope, Send from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger -from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG +from litellm.constants import MAXIMUM_TRACEBACK_LINES_TO_LOG, MCP_PEEKED_BODY_SCOPE_KEY from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import ( - MCP_PEEKED_BODY_SCOPE_KEY, MCPRequestHandler, _is_mcp_admitted_user_subject, ) @@ -3798,6 +3797,14 @@ if MCP_AVAILABLE: return f"ip:{hashlib.sha256(client_ip.encode('utf-8')).hexdigest()}" return "anonymous" + def _peeked_json_object_body(body: bytes) -> bytes | None: + """The peeked bytes when they form a complete JSON-RPC object, else None so a + truncated prefix or a batch array fails closed and stays budget-enforced.""" + try: + return body if isinstance(json.loads(body), dict) else None + except (json.JSONDecodeError, TypeError): + return None + def _is_initialize_request(body: bytes) -> bool: """ Check if the request body is a JSON-RPC initialize method. @@ -4390,24 +4397,22 @@ if MCP_AVAILABLE: """Handle MCP requests through StreamableHTTP.""" try: path: Final[str] = scope.get("path", "") - consumed_messages: list[Message] = [] + consumed_messages: list[Message] = [] # mutable-ok: replay buffer for peeked ASGI messages body = b"" if scope.get("method") == "POST": consumed_messages, body = await _read_request_body_for_routing(receive) if consumed_messages: original_receive: Final = receive - async def wrapped_receive(): + async def wrapped_receive() -> Message: if consumed_messages: return consumed_messages.pop(0) return await original_receive() - receive = wrapped_receive - try: - if isinstance(json.loads(body), dict): - scope[MCP_PEEKED_BODY_SCOPE_KEY] = body - except (json.JSONDecodeError, TypeError): - pass + receive = wrapped_receive # rebind-ok: replay peeked ASGI messages to the downstream handler + peeked_object_body: Final = _peeked_json_object_body(body) + if peeked_object_body is not None: + scope[MCP_PEEKED_BODY_SCOPE_KEY] = peeked_object_body is_initialize: Final = _is_initialize_request(body) ( user_api_key_auth, 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 398f3b2ea40..20a0d97bc4f 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 @@ -2161,8 +2161,8 @@ async def test_mcp_routing_stashes_peeked_body_for_auth(): send = AsyncMock() stateless_called = [] - async def stateless_handle(s, r, se): - stateless_called.append(1) + async def stateless_handle(s, r, se) -> None: + stateless_called.append(1) # mutable-ok: records the manager the handler dispatched to with ( patch( # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam