fix(proxy): exempt zero-spend MCP discovery routes from budget checks (#41276)

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

litellm/constants.py

@@ -387,6 +387,8 @@ BEDROCK_REALTIME_SDK_DISTRIBUTION: Final = "aws-sdk-bedrock-runtime"
 BEDROCK_REALTIME_SDK_SUPPORTED_RANGE: Final = ">=0.10.0,<0.12.0"
 CLIENT_REQUESTED_MODEL_SCOPE_KEY: Final = "litellm.client_requested_model"
 MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY: Final = "litellm.model_group_alias_resolved"
+MCP_PEEKED_BODY_SCOPE_KEY: Final = "litellm_mcp_peeked_body"
+MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS: Final = 5.0
 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

@@ -1,6 +1,6 @@
 import re
 import secrets
-from collections.abc import Mapping, Sequence
+from collections.abc import Awaitable, Callable, Mapping, Sequence
 from collections.abc import Set as AbstractSet
 from dataclasses import dataclass
 from datetime import datetime, timezone
@@ -17,7 +17,7 @@ from typing_extensions import assert_never
 import litellm
 from litellm._internal_context import with_service_target
 from litellm._logging import verbose_logger
-from litellm.constants import MCP_ALL_TOOLS_WILDCARD
+from litellm.constants import MCP_ALL_TOOLS_WILDCARD, MCP_PEEKED_BODY_SCOPE_KEY
 from litellm.experimental_mcp_client.client import strip_auth_scheme
 from litellm.proxy._experimental.mcp_server.catalog import global_manager
 from litellm.proxy._experimental.mcp_server.oauth_utils import (
@@ -94,6 +94,22 @@ def _normalize_caller_admission_credential(value: str) -> str:
     return _get_bearer_token_or_received_api_key(strip_auth_scheme(value, "Bearer")).strip()
 
 
+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: Final = Request(scope=scope)
+    peeked_body: Final[Callable[[], Awaitable[bytes]] | None] = scope.get(MCP_PEEKED_BODY_SCOPE_KEY)
+
+    async def mock_body() -> bytes:
+        if peeked_body is None:
+            return b"{}"
+        return await peeked_body()
+
+    request.body = mock_body
+    return request
+
+
 def _as_list(values: Sequence[str] | None) -> list[str] | None:  # mutable-ok: resolver returns a list
     """Widen a read-only allowlist back to the mutable list the resolver's own contract returns,
     preserving the ``None`` that means "no restriction"."""
@@ -505,12 +521,7 @@ class MCPRequestHandler:
                 if mcp_servers_header == "" or (mcp_servers is not None and len(mcp_servers) == 0):
                     mcp_servers = []
             # Create a proper Request object with mock body method to avoid ASGI receive channel issues
-            request: Final = Request(scope=scope)
-
-            async def mock_body():
-                return b"{}"
-
-            request.body = mock_body
+            request: Final = _admission_request(scope)
             # Inline import — auth_utils participates in a proxy import cycle.
             from litellm.proxy.auth.auth_utils import (  # noqa: PLC0415
                 get_request_route,

litellm/proxy/_experimental/mcp_server/server.py

@@ -26,7 +26,9 @@ from starlette.types import Message, Receive, Scope, Send
 
 from litellm._logging import verbose_logger
 from litellm.constants import (
+    MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS,
     MCP_GATEWAY_SESSION_ID_PREFIX_LENGTH,
+    MCP_PEEKED_BODY_SCOPE_KEY,
 )
 from litellm.llms.custom_httpx.http_handler import (
     get_async_httpx_client,
@@ -1262,6 +1264,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.
@@ -1402,6 +1412,7 @@ if MCP_AVAILABLE:
         receive: Receive,
         *,
         full_body: bool = False,
+        already_read: int = 0,
     ) -> tuple[list[Message], bytes]:
         """
         Read just enough of the request body to decide whether this is a
@@ -1412,14 +1423,14 @@ if MCP_AVAILABLE:
         Stops reading from the wire as soon as either (a) we have peeked
         ``_MCP_ROUTING_PEEK_MAX_BYTES`` of body, or (b) the body is complete.
         The remainder of an oversized body is streamed lazily through
-        ``wrapped_receive`` in the caller — so an authenticated client cannot
+        ``_LazyPeekedBody.receive`` in the caller — so an authenticated client cannot
         force the proxy to buffer an arbitrarily large payload just to make a
         routing decision. Modern interaction preflight requests the full body,
         which the SDK's single-exchange transport also requires.
         """
         consumed_messages: Final[list[Message]] = []
         body_chunks: Final[list[bytes]] = []
-        peeked_bytes = 0
+        peeked_bytes = already_read
 
         while True:
             message = await receive()
@@ -1457,6 +1468,60 @@ if MCP_AVAILABLE:
 
         return consumed_messages, b"".join(body_chunks)
 
+    _PeekedRead = tuple[list[Message], bytes]
+
+    class _LazyPeekedBody:
+        """Reads the POST body for routing and auth only when first asked, so an
+        allowlist rejection never touches the receive channel. Whatever it
+        consumed is replayed to the downstream handler through ``receive``,
+        which falls through to the wire once the replay buffer is drained."""
+
+        __slots__ = ("_consumed_messages", "_peeked_body", "_read", "_receive")
+
+        def __init__(self, receive: Receive) -> None:
+            self._receive = receive
+            self._read: asyncio.Future[_PeekedRead] | None = None
+            self._consumed_messages: list[Message] | None = None
+            self._peeked_body: bytes = b""
+
+        def _start_read(self) -> asyncio.Future[_PeekedRead]:
+            if self._read is None:
+                self._read = asyncio.ensure_future(_read_request_body_for_routing(self._receive))
+            return self._read
+
+        async def body(self) -> bytes:
+            if self._consumed_messages is None:
+                self._consumed_messages, self._peeked_body = await self._start_read()
+            return self._peeked_body
+
+        async def admission_body(self) -> bytes:
+            try:
+                await asyncio.wait_for(
+                    asyncio.shield(self._start_read()), timeout=MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS
+                )
+            except asyncio.TimeoutError:
+                return b""
+            return await self.body()
+
+        async def full_body(self) -> bytes:
+            await self.body()
+            consumed: Final = self._consumed_messages or []
+            consumed_body: Final = b"".join(message.get("body", b"") or b"" for message in consumed)
+            if len(consumed_body) > session_manager_stateless.max_request_body_size:
+                raise HTTPException(status_code=413, detail="Request body too large")
+            if not consumed or consumed[-1].get("type") != "http.request" or not consumed[-1].get("more_body", False):
+                return consumed_body
+            rest_messages, rest_body = await _read_request_body_for_routing(
+                self._receive, full_body=True, already_read=len(consumed_body)
+            )
+            self._consumed_messages = [*consumed, *rest_messages]
+            return consumed_body + rest_body
+
+        async def receive(self) -> Message:
+            if self._consumed_messages:
+                return self._consumed_messages.pop(0)
+            return await self._receive()
+
     async def _handle_stale_mcp_session(
         scope: Scope,
         receive: Receive,
@@ -2124,6 +2189,14 @@ if MCP_AVAILABLE:
                 )(scope, receive, send)
                 return
             path: Final[str] = scope.get("path", "")
+            peek: Final = _LazyPeekedBody(receive) if scope.get("method") == "POST" else None
+            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"{}"
+
+                scope[MCP_PEEKED_BODY_SCOPE_KEY] = peeked_json_object_body
             (
                 user_api_key_auth,
                 mcp_auth_header,
@@ -2133,6 +2206,8 @@ if MCP_AVAILABLE:
                 raw_headers,
             ) = await extract_mcp_auth_context(scope, path)
             reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth)
+            body: Final = await peek.body() if peek is not None else b""
+            is_initialize: Final = _is_initialize_request(body)
             scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1
             request_headers: Final = StarletteRequest(scope).headers
             defer_upstream_probes: Final = request_headers.get(
@@ -2217,8 +2292,6 @@ if MCP_AVAILABLE:
             # - No session ID + initialize → stateful (so client gets mcp-session-id)
             # - No session ID + other → stateless (curl, Inspector, Notion)
             session_id = _get_session_id_from_scope(scope)
-            is_initialize = False
-            consumed_messages: list[Message] = []
 
             # Owner-binding: a live stateful session may only be driven by the
             # caller that created it. Reject mismatches with 403 so a leaked
@@ -2226,8 +2299,7 @@ if MCP_AVAILABLE:
             #
             # Run before ``_handle_stale_mcp_session`` so a non-owner cannot
             # force-clean another caller's residual tracking entries via a
-            # stale DELETE, and before peeking the request body so the 403
-            # response sees a pristine ``receive`` channel.
+            # stale DELETE.
             if session_id:
                 expected_owner: Final = _stateful_session_owners.get(session_id)
                 request_owner = _owner_fingerprint_for(user_api_key_auth, oauth2_headers, _client_ip)
@@ -2256,27 +2328,23 @@ if MCP_AVAILABLE:
                     return
                 session_id = _get_session_id_from_scope(scope)
 
-            body = b""
-            if scope.get("method") == "POST":
-                consumed_messages, body = await _read_request_body_for_routing(receive, full_body=defer_upstream_probes)
-                if defer_upstream_probes:
-                    rejection: Final = await _preflight_modern_interaction(
-                        scope,
-                        body,
-                        OperationContext(
-                            _caller=user_api_key_auth,
-                            mcp_auth_header=mcp_auth_header,
-                            mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None,
-                            mcp_server_auth_headers=mcp_server_auth_headers,
-                            oauth2_headers=oauth2_headers,
-                            raw_headers=raw_headers,
-                            client_ip=_client_ip,
-                        ),
-                    )
-                    if rejection is not None:
-                        await rejection(scope, receive, send)
-                        return
-                is_initialize = _is_initialize_request(body)
+            if peek is not None and defer_upstream_probes:
+                rejection: Final = await _preflight_modern_interaction(
+                    scope,
+                    await peek.full_body(),
+                    OperationContext(
+                        _caller=user_api_key_auth,
+                        mcp_auth_header=mcp_auth_header,
+                        mcp_servers=tuple(mcp_servers) if mcp_servers is not None else None,
+                        mcp_server_auth_headers=mcp_server_auth_headers,
+                        oauth2_headers=oauth2_headers,
+                        raw_headers=raw_headers,
+                        client_ip=_client_ip,
+                    ),
+                )
+                if rejection is not None:
+                    await rejection(scope, receive, send)
+                    return
 
             use_stateful: Final = bool(session_id or is_initialize)
             target_manager: Final = session_manager_stateful if use_stateful else session_manager_stateless
@@ -2306,17 +2374,6 @@ if MCP_AVAILABLE:
                     await too_many_response(scope, receive, send)
                     return
 
-            # Replay body messages if we consumed them for peeking
-            original_receive: Final = receive
-            if consumed_messages:
-
-                async def wrapped_receive():
-                    if consumed_messages:
-                        return consumed_messages.pop(0)
-                    return await original_receive()
-
-                receive = wrapped_receive
-
             # Serialize requests on the same stateful session so concurrent
             # callers don't clobber each other's auth context mid-flight.
             #

litellm/proxy/auth/auth_checks.py

@@ -971,6 +971,21 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
     }
 )
 
+MCP_DISCOVERY_ROUTES: Final = frozenset({"/mcp-rest/tools/list", "/v1/mcp/tools"})
+
+MCP_ZERO_SPEND_JSONRPC_METHODS: Final = frozenset({"initialize", "notifications/initialized", "ping", "tools/list"})
+
+MCP_TOOL_CALL_ROUTES: Final = frozenset({"/mcp/tools/call", "/mcp-rest/tools/call"})
+
+
+def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) -> bool:
+    if route in MCP_DISCOVERY_ROUTES:
+        return True
+    if route in MCP_TOOL_CALL_ROUTES or not (route == "/mcp" or route.startswith("/mcp/")):
+        return False
+    method: Final = request_body.get("method")
+    return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS
+
 
 def route_skips_budget_checks(route: str) -> bool:
     return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and (
@@ -1029,7 +1044,11 @@ async def common_checks(
         team_id=valid_token.team_id if valid_token is not None else None,
     )
 
-    skip_all_budget_checks: Final = skip_budget_checks or route_skips_budget_checks(route=route)
+    skip_all_budget_checks: Final = (
+        skip_budget_checks
+        or route_skips_budget_checks(route=route)
+        or is_mcp_discovery_request(route=route, request_body=_typed_request_body(request_body))
+    )
 
     membership_user_id: Final = (
         valid_token.user_id if valid_token is not None and (bool(_model) or not skip_all_budget_checks) else None

litellm/proxy/auth/user_api_key_auth.py

@@ -71,6 +71,7 @@ from litellm.proxy.auth.auth_checks import (  # noqa: F401  # legacy module expo
     get_team_object,
     get_user_object,
     get_user_role,
+    is_mcp_discovery_request,
     is_model_cost_zero,
     is_user_proxy_admin,
     is_valid_fallback_model,
@@ -1549,7 +1550,7 @@ async def user_api_key_auth_builder(
     anthropic_api_key_header: str | None,
     google_ai_studio_api_key_header: str | None,
     azure_apim_header: str | None,
-    request_data: dict,
+    request_data: dict[str, object],
     custom_litellm_key_header: str | None = None,
 ) -> UserAPIKeyAuth:
     from litellm.proxy.proxy_server import (
@@ -1888,8 +1889,8 @@ async def user_api_key_auth_builder(
                             llm_router=llm_router,
                             team_id=valid_token.team_id,
                         )
-                        skip_budget_checks = False
-                        if model is not None and llm_router is not None:
+                        skip_budget_checks = is_mcp_discovery_request(route=route, request_body=request_data)
+                        if not skip_budget_checks and model is not None and llm_router is not None:
                             from litellm.proxy.auth.auth_checks import is_model_cost_zero
 
                             skip_budget_checks = (  # rebind-ok: pre-existing rebinding on a rename-only line
@@ -2263,9 +2264,7 @@ async def user_api_key_auth_builder(
 
         return await validate_resolved_virtual_key(
             request=request,
-            request_data=cast(  # cast-ok: model-alias checks must mutate the original request
-                dict[str, object], request_data
-            ),
+            request_data=request_data,
             valid_token=valid_token,
             api_key=api_key,
             route=route,
@@ -2390,8 +2389,8 @@ async def validate_resolved_virtual_key(  # noqa: C901  # Preserve ordering of e
             llm_router=llm_router,
             team_id=valid_token.team_id,
         )
-        skip_budget_checks = False
-        if model is not None and llm_router is not None:
+        skip_budget_checks = is_mcp_discovery_request(route=route, request_body=request_data)
+        if not skip_budget_checks and model is not None and llm_router is not None:
             from litellm.proxy.auth.auth_checks import is_model_cost_zero
 
             skip_budget_checks = is_model_cost_zero(  # rebind-ok: pre-existing rebinding on a rename-only line
@@ -3295,7 +3294,7 @@ async def _reserve_budget_after_common_checks(
 
 
 def _should_skip_budget_checks(
-    request_data: dict,
+    request_data: dict[str, object],
     route: str,
     request: Request | None,
     llm_router: litellm.Router | None,
@@ -3308,6 +3307,8 @@ def _should_skip_budget_checks(
         llm_router=llm_router,
         team_id=team_id,
     )
+    if is_mcp_discovery_request(route=route, request_body=request_data):
+        return True
     if model is not None and llm_router is not None:
         return is_model_cost_zero(model=model, llm_router=llm_router)
     return False

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

@@ -10791,3 +10791,28 @@ async def test_catalog_refresh_uses_current_virtual_key_policy_and_keeps_session
     assert caller.team_id == "old-team" and not caller.requires_fresh_policy
     assert refreshed.access_group_ids == current_groups
     assert caller.access_group_ids == ["original-group"]
+
+
+@pytest.mark.asyncio
+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":{}}'
+
+    async def peek() -> bytes:
+        return jsonrpc_body
+
+    with_peek = _admission_request(
+        {
+            "type": "http",
+            "method": "POST",
+            "path": "/mcp",
+            "headers": [],
+            MCP_PEEKED_BODY_SCOPE_KEY: peek,
+        }
+    )
+    assert await with_peek.body() == jsonrpc_body
+
+    without_peek = _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,8 @@
 # Create server parameters for stdio connection
+import asyncio
 import os
+from typing import Final
+
 import pytest
 from unittest.mock import ANY, AsyncMock, MagicMock, patch
 from contextlib import asynccontextmanager, nullcontext
@@ -3105,3 +3108,204 @@ async def test_call_mcp_tool_resolves_unprefixed_tool_name_and_checks_permission
     assert mock_get_server.call_args_list[0][0][0] == "gmail_send_email"
     # Permissions check should be invoked with the resolved server name
     mock_is_allowed.assert_called_once()
+
+
+@pytest.mark.asyncio
+async def test_mcp_admission_body_peek_times_out_without_losing_the_body(monkeypatch: pytest.MonkeyPatch) -> None:
+    try:
+        from litellm.proxy._experimental.mcp_server import server
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    monkeypatch.setattr(server, "MCP_ADMISSION_BODY_PEEK_TIMEOUT_SECONDS", 0.01)
+    jsonrpc_body: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}'
+    message: Final = {"type": "http.request", "body": jsonrpc_body, "more_body": False}
+    body_arrived: Final = asyncio.Event()
+
+    async def stalled_receive() -> dict[str, object]:
+        await body_arrived.wait()
+        return message
+
+    peek: Final = server._LazyPeekedBody(stalled_receive)
+    assert await asyncio.wait_for(peek.admission_body(), timeout=1) == b""
+    body_arrived.set()
+    assert await asyncio.wait_for(peek.body(), timeout=1) == jsonrpc_body
+    assert await peek.receive() == message
+
+
+@pytest.mark.asyncio
+async def test_mcp_routing_stashes_peeked_body_for_auth():
+    try:
+        from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+            MCP_PEEKED_BODY_SCOPE_KEY,
+            _admission_request,
+        )
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager_stateful,
+            session_manager_stateless,
+        )
+        from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    jsonrpc_body: Final = b'{"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}'
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/mcp",
+        "headers": [
+            (b"content-type", b"application/json"),
+            (b"authorization", b"Bearer test-key"),
+        ],
+    }
+    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
+
+    with (
+        patch(  # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            new_callable=AsyncMock,
+            return_value=(MagicMock(), None, None, None, None, None),
+        ),
+        patch(  # test-quality-ok: registry is empty in unit tests; key owns one server
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            new_callable=AsyncMock,
+            return_value=[MagicMock()],
+        ),
+        patch(  # test-quality-ok: init flag is a module global; no injection seam
+            "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+            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
+        ),
+        patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
+            session_manager_stateless, "_server_instances", {}
+        ),
+        patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
+            session_manager_stateful, "_server_instances", {}
+        ),
+    ):
+        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"
+
+
+@pytest.mark.asyncio
+async def test_mcp_routing_batch_body_is_not_stashed_for_auth():
+    try:
+        from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+            MCP_PEEKED_BODY_SCOPE_KEY,
+            _admission_request,
+        )
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager_stateless,
+        )
+        from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    batch_body: Final = b'[{"jsonrpc":"2.0","id":1,"method":"tools/list"},{"jsonrpc":"2.0","id":2,"method":"tools/call"}]'
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/mcp",
+        "headers": [
+            (b"content-type", b"application/json"),
+            (b"authorization", b"Bearer test-key"),
+        ],
+    }
+    receive = AsyncMock(side_effect=[{"type": "http.request", "body": batch_body, "more_body": False}])
+    send = AsyncMock()
+
+    with (
+        patch(  # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            new_callable=AsyncMock,
+            return_value=(MagicMock(), None, None, None, None, None),
+        ),
+        patch(  # test-quality-ok: registry is empty in unit tests; key owns one server
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            new_callable=AsyncMock,
+            return_value=[MagicMock()],
+        ),
+        patch(  # test-quality-ok: init flag is a module global; no injection seam
+            "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+            True,
+        ),
+        patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
+            session_manager_stateless, "handle_request", new=AsyncMock()
+        ),
+        patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
+            session_manager_stateless, "_server_instances", {}
+        ),
+    ):
+        await handle_streamable_http_mcp(scope, receive, send)
+
+    assert await scope[MCP_PEEKED_BODY_SCOPE_KEY]() == b"{}"
+    assert await _read_request_body(_admission_request(scope)) == {}
+
+
+@pytest.mark.asyncio
+async def test_mcp_routing_truncated_json_body_is_not_stashed_for_auth():
+    try:
+        from litellm.proxy._experimental.mcp_server.auth.user_api_key_auth_mcp import (
+            MCP_PEEKED_BODY_SCOPE_KEY,
+            _admission_request,
+        )
+        from litellm.proxy._experimental.mcp_server.server import (
+            handle_streamable_http_mcp,
+            session_manager_stateless,
+        )
+        from litellm.proxy.common_utils.http_parsing_utils import _read_request_body
+    except ImportError:
+        pytest.skip("MCP server not available")
+
+    truncated_body: Final = b'{"jsonrpc":"2.0","id":1,"method":"tools'
+    scope = {
+        "type": "http",
+        "method": "POST",
+        "path": "/mcp",
+        "headers": [
+            (b"content-type", b"application/json"),
+            (b"authorization", b"Bearer test-key"),
+        ],
+    }
+    receive = AsyncMock(side_effect=[{"type": "http.request", "body": truncated_body, "more_body": False}])
+    send = AsyncMock()
+
+    with (
+        patch(  # test-quality-ok: the ASGI handler reads auth from a module-level helper; the suite's only seam
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            new_callable=AsyncMock,
+            return_value=(MagicMock(), None, None, None, None, None),
+        ),
+        patch(  # test-quality-ok: registry is empty in unit tests; key owns one server
+            "litellm.proxy._experimental.mcp_server.server._get_allowed_mcp_servers",
+            new_callable=AsyncMock,
+            return_value=[MagicMock()],
+        ),
+        patch(  # test-quality-ok: init flag is a module global; no injection seam
+            "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED",
+            True,
+        ),
+        patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
+            session_manager_stateless, "handle_request", new=AsyncMock()
+        ),
+        patch.object(  # test-quality-ok: session managers are module-level singletons; the suite's only seam
+            session_manager_stateless, "_server_instances", {}
+        ),
+    ):
+        await handle_streamable_http_mcp(scope, receive, send)
+
+    assert await scope[MCP_PEEKED_BODY_SCOPE_KEY]() == b"{}"
+    assert await _read_request_body(_admission_request(scope)) == {}

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

@@ -11595,3 +11595,58 @@ async def test_passthrough_probes_bind_each_token_to_its_authorized_server(allow
         if allowed and header_shape != "non_authorization"
         else ()
     )
+
+
+@pytest.mark.asyncio
+async def test_streamable_http_admission_auth_sees_peeked_jsonrpc_body() -> None:
+    from litellm.constants import MCP_PEEKED_BODY_SCOPE_KEY
+    from litellm.proxy._experimental.mcp_server import server as mcp_module
+
+    scope: Final[Scope] = {"type": "http", "method": "POST", "path": "/mcp", "headers": []}
+    receive: Final = AsyncMock(
+        side_effect=[
+            {"type": "http.request", "body": _INITIALIZE[:20], "more_body": True},
+            {"type": "http.request", "body": _INITIALIZE[20:], "more_body": False},
+        ]
+    )
+    send: Final = AsyncMock()
+    downstream_bodies: Final[list[bytes]] = []
+    auth_bodies: Final[list[bytes]] = []
+
+    async def extract(auth_scope: Scope, _: str):
+        auth_bodies.append(await auth_scope[MCP_PEEKED_BODY_SCOPE_KEY]())
+        return (UserAPIKeyAuth(user_id="allowlist-user", jwt_claims=_LISTED_JWT), None, None, None, None, {})
+
+    async def handle_request(_: Scope, downstream_receive: Receive, __: Send) -> None:
+        downstream_bodies.append(await _drain_body(downstream_receive))
+
+    stateful_handle: Final = AsyncMock(side_effect=handle_request)
+    stateless_handle: Final = AsyncMock()
+
+    with (
+        patch(  # test-quality-ok: the ASGI handler resolves auth through a module-level function; no injection seam
+            "litellm.proxy._experimental.mcp_server.server.extract_mcp_auth_context",
+            new_callable=AsyncMock,
+            side_effect=extract,
+        ),
+        patch(  # test-quality-ok: module flag guarding lazy session-manager startup; no injection seam
+            "litellm.proxy._experimental.mcp_server.server._SESSION_MANAGERS_INITIALIZED", True
+        ),
+        patch(  # test-quality-ok: the allowlist is read off this module global; no injection seam
+            "litellm.proxy.proxy_server.general_settings", _ALLOWLIST_SETTINGS
+        ),
+        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=stateful_handle),
+        ),
+        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=stateless_handle),
+        ),
+    ):
+        await mcp_module.handle_streamable_http_mcp(scope, receive, send)
+
+    assert auth_bodies == [_INITIALIZE]
+    assert downstream_bodies == [_INITIALIZE]
+    stateless_handle.assert_not_awaited()
+    send.assert_not_awaited()

tests/unit/proxy/auth/test_auth_checks.py

@@ -1649,3 +1649,122 @@ def test_is_model_cost_zero_judges_an_alias_chain_by_the_deployment_its_entry_ro
 
     assert verdicts == expected
     assert {name: is_model_cost_zero(model=name, llm_router=router) for name in order} == expected
+
+
+MCP_DISCOVERY_REST_ROUTES = [
+    "/mcp-rest/tools/list",
+    "/v1/mcp/tools",
+]
+
+
+@pytest.mark.parametrize("route", MCP_DISCOVERY_REST_ROUTES)
+@pytest.mark.asyncio
+async def test_mcp_tool_discovery_route_bypasses_team_budget(route):
+    from litellm.proxy.auth.auth_checks import common_checks
+
+    team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
+
+    result = await common_checks(
+        request_body={},
+        team_object=team_object,
+        user_object=None,
+        end_user_object=None,
+        global_proxy_spend=None,
+        general_settings={},
+        route=route,
+        llm_router=None,
+        proxy_logging_obj=AsyncMock(),
+        valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
+        request=MagicMock(),
+    )
+
+    assert result is True
+
+
+@pytest.mark.parametrize("method", ["initialize", "notifications/initialized", "ping", "tools/list"])
+@pytest.mark.parametrize("route", ["/mcp", "/mcp/some-server"])
+@pytest.mark.asyncio
+async def test_mcp_jsonrpc_zero_spend_method_bypasses_team_budget(route, method):
+    from litellm.proxy.auth.auth_checks import common_checks
+
+    team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
+
+    result = await common_checks(
+        request_body={"jsonrpc": "2.0", "id": 1, "method": method, "params": {}},
+        team_object=team_object,
+        user_object=None,
+        end_user_object=None,
+        global_proxy_spend=None,
+        general_settings={},
+        route=route,
+        llm_router=None,
+        proxy_logging_obj=AsyncMock(),
+        valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
+        request=MagicMock(),
+    )
+
+    assert result is True
+
+
+@pytest.mark.parametrize(
+    "route,request_body",
+    [
+        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
+        ("/mcp/some-server", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
+        ("/mcp-rest/tools/call", {"name": "add", "arguments": {}}),
+        ("/mcp/tools/call", {"jsonrpc": "2.0", "id": 1, "method": "initialize", "params": {}}),
+        ("/mcp/tools", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
+        ("/mcp/tools/list", {"jsonrpc": "2.0", "id": 1, "method": "tools/call", "params": {"name": "add"}}),
+    ],
+)
+@pytest.mark.asyncio
+async def test_mcp_tool_call_route_still_enforces_team_budget(route, request_body):
+    from litellm.proxy.auth.auth_checks import common_checks
+
+    team_object = LiteLLM_TeamTable(team_id="test-team", spend=150.0, max_budget=100.0)
+
+    with pytest.raises(litellm.BudgetExceededError):
+        await common_checks(
+            request_body=request_body,
+            team_object=team_object,
+            user_object=None,
+            end_user_object=None,
+            global_proxy_spend=None,
+            general_settings={},
+            route=route,
+            llm_router=None,
+            proxy_logging_obj=AsyncMock(),
+            valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
+            request=MagicMock(),
+        )
+
+
+@pytest.mark.parametrize(
+    "route,request_body,expected",
+    [
+        ("/mcp-rest/tools/list", {}, True),
+        ("/v1/mcp/tools", {}, True),
+        ("/mcp/tools", {}, False),
+        ("/mcp/tools/list", {}, False),
+        ("/mcp/tools", {"method": "tools/call"}, False),
+        ("/mcp/tools", {"method": "tools/list"}, True),
+        ("/mcp", {"method": "initialize"}, True),
+        ("/mcp", {"method": "notifications/initialized"}, True),
+        ("/mcp", {"method": "ping"}, True),
+        ("/mcp", {"method": "tools/list"}, True),
+        ("/mcp/some-server", {"method": "tools/list"}, True),
+        ("/mcp", {"method": "tools/call"}, False),
+        ("/mcp/some-server", {"method": "tools/call"}, False),
+        ("/mcp", {}, False),
+        ("/mcp", {"method": ["tools/list"]}, False),
+        ("/mcp", {"method": {"name": "ping"}}, False),
+        ("/mcp-rest/tools/call", {"method": "initialize"}, False),
+        ("/mcp/tools/call", {"method": "initialize"}, False),
+        ("/v1/chat/completions", {"method": "initialize"}, False),
+        ("/mcp-rest/other", {"method": "tools/list"}, False),
+    ],
+)
+def test_is_mcp_discovery_request(route, request_body, expected):
+    from litellm.proxy.auth.auth_checks import is_mcp_discovery_request
+
+    assert is_mcp_discovery_request(route=route, request_body=request_body) is expected

tests/unit/proxy/auth/test_user_api_key_auth.py

@@ -1900,3 +1900,71 @@ async def test_user_api_key_auth_websocket_rejection_adds_no_traceback_for_other
     proxy_records: Final = [record for record in caplog.records if record.name == verbose_proxy_logger.name]
     assert [record.getMessage() for record in proxy_records if record.exc_info is not None] == []
     assert [record.levelno for record in proxy_records if record.levelno >= logging.WARNING] == [logging.ERROR]
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "request_data, expect_user_model_budget_check",
+    [
+        ({"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, False),
+        ({"jsonrpc": "2.0", "id": 1, "method": "initialize"}, False),
+        ({"jsonrpc": "2.0", "id": 1, "method": "tools/call"}, True),
+    ],
+)
+async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods(
+    monkeypatch: pytest.MonkeyPatch,
+    request_data: dict[str, object],
+    expect_user_model_budget_check: bool,
+) -> None:
+    from litellm.caching.dual_cache import DualCache
+    from litellm.proxy._types import LiteLLM_JWTAuth
+    from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder
+
+    jwt_response: Final = {
+        "is_proxy_admin": False,
+        "jwt_claims": {},
+        "team_id": None,
+        "team_object": None,
+        "user_id": "jwt-user",
+        "user_email": None,
+        "user_object": LiteLLM_UserTable(
+            user_id="jwt-user",
+            model_max_budget={"gpt-4o": {"budget_limit": 1.0, "time_period": "1d"}},
+        ),
+        "org_id": None,
+        "org_object": None,
+        "end_user_id": None,
+        "end_user_object": None,
+        "token": "fake.jwt.token",
+    }
+    mock_request: Final = MagicMock()
+    mock_request.url.path = "/mcp"
+    mock_request.method = "POST"
+    mock_request.headers = {"authorization": "Bearer fake.jwt.token"}
+    mock_request.query_params = {}
+
+    monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True})
+    litellm.proxy.proxy_server.jwt_handler.update_environment(
+        prisma_client=None,
+        user_api_key_cache=DualCache(),
+        litellm_jwtauth=LiteLLM_JWTAuth(),
+    )
+    user_model_budget_check: Final = AsyncMock()
+
+    with (
+        patch("litellm.proxy.proxy_server.premium_user", True),
+        patch("litellm.proxy.auth.handle_jwt.JWTAuthManager.auth_builder", return_value=jwt_response),
+        patch("litellm.proxy.auth.user_api_key_auth._check_user_model_budget", user_model_budget_check),
+    ):
+        result: Final = await user_api_key_auth_builder(
+            request=mock_request,
+            api_key="Bearer fake.jwt.token",
+            azure_api_key_header="",
+            anthropic_api_key_header=None,
+            google_ai_studio_api_key_header=None,
+            azure_apim_header=None,
+            request_data=request_data,
+        )
+
+    assert result.user_id == "jwt-user"
+    assert user_model_budget_check.await_count == (1 if expect_user_model_budget_check else 0)

tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py

@@ -10501,3 +10501,93 @@ async def test_lens_oauth2_dispatch_preserves_opt_in_premium_and_route_boundary(
     else:
         assert error.value.code == "400"
         assert "No connected db" in error.value.message
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize(
+    "route, request_data, expect_blocked",
+    [
+        ("/mcp-rest/tools/list", {}, False),
+        ("/v1/mcp/tools", {}, False),
+        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "initialize"}, False),
+        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}, False),
+        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/call"}, True),
+        ("/mcp", {}, True),
+        ("/mcp-rest/tools/call", {}, True),
+    ],
+)
+async def test_exhausted_key_budget_spares_only_zero_spend_mcp_discovery(
+    route: str, request_data: dict[str, object], expect_blocked: bool
+) -> None:
+    from litellm.proxy.utils import hash_token
+
+    api_key: Final = "sk-exhausted-mcp-discovery"
+    hashed_token: Final = hash_token(api_key)
+    user_api_key_cache: Final = DualCache()
+    await cache_key_object(
+        hashed_token=hashed_token,
+        user_api_key_obj=UserAPIKeyAuth(token=hashed_token, max_budget=2.0, spend=5.0),
+        user_api_key_cache=user_api_key_cache,
+        proxy_logging_obj=None,
+    )
+
+    mock_request: Final = MagicMock()
+    mock_request.url.path = route
+    mock_request.method = "POST"
+    mock_request.headers = {"authorization": f"Bearer {api_key}"}
+    mock_request.query_params = {}
+    mock_request.state = SimpleNamespace()
+
+    proxy_logging_obj: Final = MagicMock()
+    proxy_logging_obj.budget_alerts = AsyncMock()
+    proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
+
+    async def authenticate() -> UserAPIKeyAuth:
+        return await user_api_key_auth_builder(
+            request=mock_request,
+            api_key=f"Bearer {api_key}",
+            azure_api_key_header="",
+            anthropic_api_key_header=None,
+            google_ai_studio_api_key_header=None,
+            azure_apim_header=None,
+            request_data=request_data,
+        )
+
+    with (
+        patch("litellm.proxy.proxy_server.general_settings", {}),
+        patch("litellm.proxy.proxy_server.master_key", "sk-master"),
+        patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
+        patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),
+        patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj),
+        patch(
+            "litellm.proxy.auth.user_api_key_auth.virtual_key_max_budget_alert_check",
+            new_callable=AsyncMock,
+        ),
+    ):
+        if expect_blocked:
+            with pytest.raises(ProxyException) as exc_info:
+                await authenticate()
+            assert "Budget has been exceeded! Key=key Current cost: 5.0, Max budget: 2.0" in exc_info.value.message
+        else:
+            result: Final = await authenticate()
+            assert result.token == hashed_token
+            assert result.spend == 5.0
+
+
+@pytest.mark.parametrize(
+    "route, request_data, expected",
+    [
+        ("/mcp-rest/tools/list", {}, True),
+        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "ping"}, True),
+        ("/mcp", {"jsonrpc": "2.0", "id": 1, "method": "tools/call"}, False),
+        ("/mcp-rest/tools/call", {}, False),
+    ],
+)
+def test_should_skip_budget_checks_spares_only_zero_spend_mcp_discovery(
+    route: str, request_data: dict[str, object], expected: bool
+) -> None:
+    from litellm.proxy.auth.user_api_key_auth import _should_skip_budget_checks
+
+    assert (
+        _should_skip_budget_checks(request_data=request_data, route=route, request=None, llm_router=None) is expected
+    )