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 bb2fe2d59b4..2c4c942ebeb 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 @@ -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) diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index a30bf30e45f..696550f2169 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/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, diff --git a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 63756ee108c..9f717f522c0 100644 --- a/tests/unit/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/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"{}" diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py index 3953d9d086c..3345683517b 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server.py +++ 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 diff --git a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py b/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py index 4a214617142..4b46406ac69 100644 --- a/tests/unit/proxy/_experimental/mcp_server/test_mcp_server_tool_calls_and_headers.py +++ b/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()