fix(proxy): let MCP server management bodies carry their OAuth token exchange settings (#45723)

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

litellm/proxy/auth/auth_utils.py

@@ -5,6 +5,7 @@ import sys
 from collections.abc import Collection, Iterator, Mapping
 from functools import lru_cache
 from logging import Logger
+from types import MappingProxyType
 from typing import Any, Final, Protocol
 
 from fastapi import HTTPException, Request, status
@@ -37,7 +38,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import (
 )
 from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present
 from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params
-from litellm.types.utils import CustomPricingLiteLLMParams
+from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params
 
 
 def is_invalid_virtual_key_error(exception: BaseException | None) -> bool:
@@ -279,6 +280,7 @@ _DEPLOYMENT_MANAGEMENT_ROUTES: Final[frozenset[str]] = frozenset(
     ("/model/new", "/model/update", "/model/delete", "/health/test_connection")
 )
 _DEPLOYMENT_ID_UPDATE_ROUTE: Final = re.compile(r"^/model/[^/]+/update$")
+_MCP_SERVER_MANAGEMENT_ROUTE: Final = re.compile(r"^/v1/mcp/server(?:/.*)?$")
 
 
 def route_manages_deployments(route: str | None) -> bool:
@@ -297,6 +299,10 @@ def route_manages_deployments(route: str | None) -> bool:
     )
 
 
+def route_manages_mcp_servers(route: str | None) -> bool:
+    return route is not None and _MCP_SERVER_MANAGEMENT_ROUTE.fullmatch(route) is not None
+
+
 _NESTED_CONFIG_KEYS: Final[tuple[str, ...]] = ("litellm_embedding_config", "extra_body")
 
 # Metadata containers that carry per-request configuration consumed by the
@@ -456,19 +462,25 @@ _BANNED_REQUEST_BODY_PARAMS: Final[tuple[str, ...]] = (
 
 
 def _check_banned_params(
-    body: dict,
+    body: Mapping[str, object],
     general_settings: dict,
     llm_router: Router | None,
     model: str,
     *,
     manages_deployments: bool = False,
+    mcp_server_fields_allowed: bool = False,
 ) -> None:
     """Raise ``ValueError`` if ``body`` carries a banned param without admin opt-in.
 
     Shared between the root-level check and the nested-config check so a
     new banned param only needs to be added in one place.
     """
-    reject_server_owned_wif_params(body)
+    wif_body: Final[Mapping[str, object]] = (
+        MappingProxyType({key: value for key, value in body.items() if key not in oauth_token_exchange_litellm_params})
+        if mcp_server_fields_allowed
+        else body
+    )
+    reject_server_owned_wif_params(wif_body)
     if not manages_deployments:
         reject_federated_credential_reference(body)
     for param in _BANNED_REQUEST_BODY_PARAMS:
@@ -613,7 +625,14 @@ def is_request_body_safe(
     if "model_list" in request_body:
         raise ValueError("Rejected Request: model_list is not allowed in the request body.")
     manages_deployments: Final = route_manages_deployments(route)
-    _check_banned_params(request_body, general_settings, llm_router, model, manages_deployments=manages_deployments)
+    _check_banned_params(
+        request_body,
+        general_settings,
+        llm_router,
+        model,
+        manages_deployments=manages_deployments,
+        mcp_server_fields_allowed=route_manages_mcp_servers(route),
+    )
     for nested_key in _NESTED_CONFIG_KEYS:
         nested = _coerce_metadata_to_dict(request_body.get(nested_key))
         if nested is not None:

tests/unit/proxy/auth/test_auth_utils.py

@@ -64,6 +64,60 @@ def test_token_exchange_settings_in_request_body_are_rejected(field: str) -> Non
     assert field in str(error.value)
 
 
+@pytest.mark.parametrize("field", oauth_token_exchange_litellm_params)
+@pytest.mark.parametrize("route", ("/v1/mcp/server", "/v1/mcp/server/abc123"))
+def test_mcp_server_management_body_may_set_token_exchange_settings(field: str, route: str) -> None:
+    assert (
+        is_request_body_safe(
+            request_body={"server_id": "abc123", "auth_type": "oauth2_token_exchange", field: "configured-value"},
+            general_settings={},
+            llm_router=None,
+            model="",
+            route=route,
+        )
+        is True
+    )
+
+
+@pytest.mark.parametrize("field", oauth_token_exchange_litellm_params)
+@pytest.mark.parametrize("route", ("/chat/completions", "/v1/mcp/serverfoo"))
+def test_token_exchange_settings_remain_banned_outside_mcp_server_management(field: str, route: str) -> None:
+    with pytest.raises(ValueError, match=field):
+        is_request_body_safe(
+            request_body={"server_id": "abc123", "auth_type": "oauth2_token_exchange", field: "configured-value"},
+            general_settings={},
+            llm_router=None,
+            model="",
+            route=route,
+        )
+
+
+def test_mcp_server_management_body_still_rejects_anthropic_wif_settings() -> None:
+    with pytest.raises(ValueError, match="anthropic_federation_rule_id"):
+        is_request_body_safe(
+            request_body={"server_id": "abc123", "anthropic_federation_rule_id": "fdrl_admin"},
+            general_settings={},
+            llm_router=None,
+            model="",
+            route="/v1/mcp/server",
+        )
+
+
+def test_mcp_server_management_body_still_rejects_nested_token_exchange_settings() -> None:
+    with pytest.raises(ValueError, match="token_exchange_endpoint"):
+        is_request_body_safe(
+            request_body={
+                "server_id": "abc123",
+                "auth_type": "oauth2_token_exchange",
+                "metadata": {"token_exchange_endpoint": "https://identity.example.com/token"},
+            },
+            general_settings={},
+            llm_router=None,
+            model="",
+            route="/v1/mcp/server",
+        )
+
+
 def test_model_opt_in_cannot_allow_a_token_exchange_endpoint_in_a_request_body() -> None:
     router: Final = Router(
         model_list=[