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

@@ -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"
 

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

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,

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