fix(proxy): only pass the JSON-RPC method to MCP admission auth (#45844)

* fix(proxy): only pass the JSON-RPC method to MCP admission auth

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): expect method-only MCP admission body

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

litellm/proxy/_experimental/mcp_server/auth/user_api_key_auth_mcp.py

@@ -95,9 +95,7 @@ def _normalize_caller_admission_credential(value: str) -> str:
 
 
 def _admission_request(scope: Scope) -> Request:
-    """Request whose ``body()`` serves the routing layer's lazily peeked JSON-RPC
-    bytes (``b"{}"`` when no peek callable was stashed) instead of the ASGI
-    receive channel."""
+    """Request whose ``body()`` serves the peeked JSON-RPC method (``b"{}"`` when nothing is stashed)."""
     request: Final = Request(scope=scope)
     peeked_body: Final[Callable[[], Awaitable[bytes]] | None] = scope.get(MCP_PEEKED_BODY_SCOPE_KEY)
 

litellm/proxy/_experimental/mcp_server/server.py

@@ -1264,13 +1264,16 @@ 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."""
+    def _admission_auth_body(body: bytes) -> bytes:
+        """Return only the JSON-RPC method for admission auth."""
         try:
-            return body if isinstance(json.loads(body), dict) else None
-        except (json.JSONDecodeError, TypeError):
-            return None
+            data: Final = TypeAdapter(dict[str, object]).validate_json(body)
+        except (ValueError, TypeError):
+            return b"{}"
+        method: Final = data.get("method")
+        if not isinstance(method, str):
+            return b"{}"
+        return json.dumps({"method": method}).encode()
 
     def _is_initialize_request(body: bytes) -> bool:
         """
@@ -2193,10 +2196,10 @@ if MCP_AVAILABLE:
             if peek is not None:
                 receive = peek.receive  # rebind-ok: replay peeked ASGI messages to the downstream handler
 
-                async def peeked_json_object_body() -> bytes:
-                    return _peeked_json_object_body(await peek.admission_body()) or b"{}"
+                async def admission_auth_body() -> bytes:
+                    return _admission_auth_body(await peek.admission_body())
 
-                scope[MCP_PEEKED_BODY_SCOPE_KEY] = peeked_json_object_body
+                scope[MCP_PEEKED_BODY_SCOPE_KEY] = admission_auth_body
             (
                 user_api_key_auth,
                 mcp_auth_header,

tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py

@@ -10794,12 +10794,12 @@ async def test_admission_request_body_serves_stashed_peek_callable():
     from litellm.constants import MCP_PEEKED_BODY_SCOPE_KEY
     from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import _admission_request
 
-    jsonrpc_body = b'{"jsonrpc":"2.0","id":1,"method":"tools/list","params":{}}'
+    admission_body: Final = b'{"method":"tools/list"}'
 
     async def peek() -> bytes:
-        return jsonrpc_body
+        return admission_body
 
-    with_peek = _admission_request(
+    with_peek: Final = _admission_request(
         {
             "type": "http",
             "method": "POST",
@@ -10808,7 +10808,9 @@ async def test_admission_request_body_serves_stashed_peek_callable():
             MCP_PEEKED_BODY_SCOPE_KEY: peek,
         }
     )
-    assert await with_peek.body() == jsonrpc_body
+    assert await with_peek.body() == admission_body
 
-    without_peek = _admission_request({"type": "http", "method": "POST", "path": "/mcp", "headers": []})
+    without_peek: Final = _admission_request(
+        {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
+    )
     assert await without_peek.body() == b"{}"

tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py

@@ -1,5 +1,6 @@
 # Create server parameters for stdio connection
 import asyncio
+import json
 import os
 from typing import Final
 
@@ -7,6 +8,8 @@ import pytest
 from unittest.mock import ANY, AsyncMock, MagicMock, patch
 from contextlib import asynccontextmanager, nullcontext
 
+from starlette.types import Receive, Scope, Send
+
 
 from litellm.proxy._experimental.mcp_server.mcp_server_manager import (
     MCPServerManager,
@@ -3134,7 +3137,8 @@ async def test_mcp_admission_body_peek_times_out_without_losing_the_body(monkeyp
 
 
 @pytest.mark.asyncio
-async def test_mcp_routing_stashes_peeked_body_for_auth():
+@pytest.mark.parametrize("method", ["tools/list", "tools/call", ["tools/list"]])
+async def test_mcp_routing_stashes_only_method_for_auth_and_replays_full_body(method: str | list[str]):
     try:
         from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
             MCP_PEEKED_BODY_SCOPE_KEY,
@@ -3149,7 +3153,15 @@ async def test_mcp_routing_stashes_peeked_body_for_auth():
     except ImportError:
         pytest.skip("MCP server not available")
 
-    jsonrpc_body: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}'
+    jsonrpc_body: Final = json.dumps(
+        {
+            "jsonrpc": "2.0",
+            "id": 2,
+            "method": method,
+            "model": "zero-cost-model",
+            "params": {},
+        }
+    ).encode()
     scope = {
         "type": "http",
         "method": "POST",
@@ -3161,10 +3173,15 @@ async def test_mcp_routing_stashes_peeked_body_for_auth():
     }
     receive = AsyncMock(side_effect=[{"type": "http.request", "body": jsonrpc_body, "more_body": False}])
     send = AsyncMock()
-    stateless_called = []
 
-    async def stateless_handle(s, r, se) -> None:
-        stateless_called.append(1)  # mutable-ok: records the manager the handler dispatched to
+    async def stateless_handle(_scope: Scope, downstream_receive: Receive, _send: Send) -> None:
+        assert await downstream_receive() == {
+            "type": "http.request",
+            "body": jsonrpc_body,
+            "more_body": False,
+        }
+
+    stateless_handle_mock: Final = AsyncMock(side_effect=stateless_handle)
 
     with (
         patch(  # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam
@@ -3182,7 +3199,7 @@ async def test_mcp_routing_stashes_peeked_body_for_auth():
             True,
         ),
         patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
-            session_manager_stateless, "handle_request", side_effect=stateless_handle
+            session_manager_stateless, "handle_request", side_effect=stateless_handle_mock
         ),
         patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
             session_manager_stateless, "_server_instances", {}
@@ -3193,10 +3210,10 @@ async def test_mcp_routing_stashes_peeked_body_for_auth():
     ):
         await handle_streamable_http_mcp(scope, receive, send)
 
-    assert stateless_called, "tools/list without a session should route to the stateless manager"
-    assert await scope[MCP_PEEKED_BODY_SCOPE_KEY]() == jsonrpc_body
-    request_data: Final = await _read_request_body(_admission_request(scope))
-    assert request_data.get("method") == "tools/list"
+    stateless_handle_mock.assert_awaited_once()
+    assert await _read_request_body(_admission_request(scope)) == (
+        {"method": method} if isinstance(method, str) else {}
+    )
 
 
 @pytest.mark.asyncio

tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py

@@ -11598,7 +11598,7 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow
 
 
 @pytest.mark.asyncio
-async def test_streamable_http_admission_auth_sees_peeked_jsonrpc_body() -> None:
+async def test_streamable_http_admission_auth_sees_only_peeked_jsonrpc_method() -> None:
     from litellm.constants import MCP_PEEKED_BODY_SCOPE_KEY
     from litellm.proxy._experimental.mcp_server import server as mcp_module
 
@@ -11646,7 +11646,7 @@ async def test_streamable_http_admission_auth_sees_peeked_jsonrpc_body() -> None
     ):
         await mcp_module.handle_streamable_http_mcp(scope, receive, send)
 
-    assert auth_bodies == [_INITIALIZE]
+    assert auth_bodies == [b'{"method": "initialize"}']
     assert downstream_bodies == [_INITIALIZE]
     stateless_handle.assert_not_awaited()
     send.assert_not_awaited()