From 0536fb3062c148bd1ac5bceebbbadf0910a42b74 Mon Sep 17 00:00:00 2001 From: yassin Date: Fri, 18 Sep 2026 21:29:59 +0000 Subject: [PATCH] fix(mcp): fail closed on empty JWT claims and gate the REST tool routes on mcp_allowed_clients Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../mcp_server/client_allowlist.py | 2 +- .../mcp_server/rest_endpoints.py | 3 + .../proxy/_experimental/mcp_server/server.py | 12 +- .../mcp_server/test_client_allowlist.py | 6 + .../mcp_server/test_rest_endpoints.py | 124 ++++++++++++++++++ 5 files changed, 142 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/_experimental/mcp_server/client_allowlist.py b/litellm/proxy/_experimental/mcp_server/client_allowlist.py index cb708723efd..1b8bf7aa9ba 100644 --- a/litellm/proxy/_experimental/mcp_server/client_allowlist.py +++ b/litellm/proxy/_experimental/mcp_server/client_allowlist.py @@ -124,7 +124,7 @@ def resolve_mcp_client_identity( headers: Mapping[str, str], ) -> MCPClientIdentity | MCPClientRejection: """A JWT caller is identified by its configured claim alone, so a header can never override the IdP.""" - if jwt_claims and allowlist.jwt_field is not None: + if jwt_claims is not None and allowlist.jwt_field is not None: claim: Final[object] = get_nested_value(data=jwt_claims, key_path=allowlist.jwt_field) if isinstance(claim, str) and claim: return MCPClientIdentity(client_id=claim, source="jwt", source_name=allowlist.jwt_field) diff --git a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py index 6a0ab5bdec5..9d895d755dd 100644 --- a/litellm/proxy/_experimental/mcp_server/rest_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/rest_endpoints.py @@ -193,6 +193,7 @@ if MCP_AVAILABLE: filter_tools_by_allowed_tools, filter_tools_by_key_team_permissions, fire_mcp_tool_call_failure_logging, + reject_disallowed_mcp_client, ) ######################################################## @@ -875,6 +876,7 @@ if MCP_AVAILABLE: MCPRequestHandler, ) + reject_disallowed_mcp_client(request.headers, user_api_key_dict) try: mcp_server_name = _as_query_str(mcp_server_name) toolset_name = _as_query_str(toolset_name) @@ -1078,6 +1080,7 @@ if MCP_AVAILABLE: proxy_logging_obj, ) + reject_disallowed_mcp_client(request.headers, user_api_key_dict) try: user_api_key_dict = await acting_user_auth(user_api_key_dict) data = await request.json() diff --git a/litellm/proxy/_experimental/mcp_server/server.py b/litellm/proxy/_experimental/mcp_server/server.py index 244793850b2..73366b701d2 100644 --- a/litellm/proxy/_experimental/mcp_server/server.py +++ b/litellm/proxy/_experimental/mcp_server/server.py @@ -77,6 +77,7 @@ from litellm.proxy._experimental.mcp_server.oauth_utils import ( get_route_relative_request_path, well_known_root_suffix, ) +from litellm.proxy._experimental.mcp_server.ui_session_utils import is_ui_session_credential from litellm.proxy._experimental.mcp_server.utils import ( LITELLM_MCP_SERVER_DESCRIPTION, LITELLM_MCP_SERVER_NAME, @@ -3711,11 +3712,14 @@ if MCP_AVAILABLE: return load_mcp_client_allowlist(general_settings) - def _reject_disallowed_mcp_client(scope: Scope, user_api_key_auth: UserAPIKeyAuth | None) -> None: + def reject_disallowed_mcp_client(headers: Mapping[str, str], user_api_key_auth: UserAPIKeyAuth | None) -> None: + """Gate every MCP tool surface on ``mcp_allowed_clients``; the dashboard's own session is not a client app.""" + if user_api_key_auth is not None and is_ui_session_credential(user_api_key_auth): + return rejection: Final = check_mcp_client_allowed( allowlist=_load_mcp_client_allowlist(), jwt_claims=user_api_key_auth.jwt_claims if user_api_key_auth is not None else None, - headers=StarletteRequest(scope).headers, + headers=headers, ) if rejection is None: return @@ -4558,7 +4562,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) - _reject_disallowed_mcp_client(scope, user_api_key_auth) + reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control @@ -4887,7 +4891,7 @@ if MCP_AVAILABLE: oauth2_headers, raw_headers, ) = await extract_mcp_auth_context(scope, path) - _reject_disallowed_mcp_client(scope, user_api_key_auth) + reject_disallowed_mcp_client(StarletteRequest(scope).headers, user_api_key_auth) scoped_server_endpoint: Final = len(_get_mcp_servers_in_path(path) or []) == 1 # Extract client IP for MCP access control diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py index 714b044d311..d8b52b7b348 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_client_allowlist.py @@ -157,6 +157,12 @@ def test_jwt_caller_is_judged_by_its_claim_even_when_the_header_would_pass() -> assert check_mcp_client_allowed(JWT_AND_HEADER, {"azp": "antigravity-cli"}, {"x-mcp-client": "claude-code"}) is None +def test_jwt_caller_with_an_empty_claim_set_cannot_fall_back_to_the_header() -> None: + rejection: Final = check_mcp_client_allowed(JWT_AND_HEADER, {}, {"x-mcp-client": "antigravity-cli"}) + assert isinstance(rejection, MCPClientRejection) + assert "azp" in rejection.details + + def test_non_jwt_caller_falls_back_to_the_header_when_both_sources_are_configured() -> None: assert check_mcp_client_allowed(JWT_AND_HEADER, None, {"x-mcp-client": "antigravity-cli"}) is None assert check_mcp_client_allowed(JWT_AND_HEADER, None, {"x-mcp-client": "claude-code"}) is not None diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py index 4ec4ae31ca6..21de380f831 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_rest_endpoints.py @@ -4313,3 +4313,127 @@ class TestV1ResolvedOauth2Gate: assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv"]) == set() assert rest_endpoints._v1_resolved_oauth2_server_ids(["oauth2-srv", "delegate-srv"]) == {"delegate-srv"} + + +_CLIENT_ALLOWLIST_SETTINGS: Final[dict[str, object]] = { + "mcp_allowed_clients": ["antigravity-cli"], + "litellm_jwtauth": {"mcp_client_id_jwt_field": "azp"}, + "mcp_client_id_header": "x-mcp-client", +} + + +class TestClientAllowlistOnRestRoutes: + """``mcp_allowed_clients`` must gate the REST tool facade exactly like the /mcp transports, + otherwise an unlisted harness can list and call tools by switching to /mcp-rest.""" + + pytestmark = pytest.mark.asyncio + + @staticmethod + def _stub_listing(monkeypatch: pytest.MonkeyPatch) -> list[UserAPIKeyAuth]: + listed_for: list[UserAPIKeyAuth] = [] + + async def fake_contexts(user_api_key_auth): + listed_for.append(user_api_key_auth) + return [user_api_key_auth] + + async def fake_get_allowed_mcp_servers(*args, **kwargs): + return [] + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", _CLIENT_ALLOWLIST_SETTINGS, raising=False) + monkeypatch.setattr(rest_endpoints, "build_effective_auth_contexts", fake_contexts, raising=False) + monkeypatch.setattr( + rest_endpoints.global_mcp_server_manager, + "get_allowed_mcp_servers", + fake_get_allowed_mcp_servers, + raising=False, + ) + return listed_for + + @pytest.mark.parametrize( + ("caller", "headers", "expected_fragment"), + ( + (UserAPIKeyAuth(jwt_claims={"azp": "claude-code"}), {"x-mcp-client": "antigravity-cli"}, "'claude-code'"), + (UserAPIKeyAuth(jwt_claims={}), {"x-mcp-client": "antigravity-cli"}, "no 'azp' claim"), + (UserAPIKeyAuth(), {"x-mcp-client": "claude-code"}, "'claude-code'"), + (UserAPIKeyAuth(), {}, "no 'x-mcp-client' header"), + ), + ) + async def test_tools_list_rejects_unlisted_clients_before_resolving_servers( + self, + monkeypatch: pytest.MonkeyPatch, + caller: UserAPIKeyAuth, + headers: dict[str, str], + expected_fragment: str, + ) -> None: + listed_for: Final = self._stub_listing(monkeypatch) + request: Final = _build_request(headers, path="/mcp-rest/tools/list", method="GET") + + with pytest.raises(HTTPException) as denied: + await rest_endpoints.list_tool_rest_api( + request, server_id=None, mcp_server_name=None, toolset_name=None, user_api_key_dict=caller + ) + + assert denied.value.status_code == 403 + assert denied.value.detail["error"] == "Forbidden" + assert expected_fragment in denied.value.detail["details"] + assert "mcp_allowed_clients" in denied.value.detail["details"] + assert listed_for == [] + + @pytest.mark.parametrize( + ("caller", "headers"), + ( + (UserAPIKeyAuth(jwt_claims={"azp": "antigravity-cli"}), {"x-mcp-client": "claude-code"}), + (UserAPIKeyAuth(), {"x-mcp-client": "antigravity-cli"}), + ), + ) + async def test_tools_list_admits_listed_clients( + self, monkeypatch: pytest.MonkeyPatch, caller: UserAPIKeyAuth, headers: dict[str, str] + ) -> None: + listed_for: Final = self._stub_listing(monkeypatch) + request: Final = _build_request(headers, path="/mcp-rest/tools/list", method="GET") + + result: Final = await rest_endpoints.list_tool_rest_api( + request, server_id=None, mcp_server_name=None, toolset_name=None, user_api_key_dict=caller + ) + + assert result["tools"] == [] + assert listed_for == [caller] + + async def test_dashboard_session_is_not_treated_as_a_client_application( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.constants import UI_SESSION_TOKEN_TEAM_ID + + listed_for: Final = self._stub_listing(monkeypatch) + session: Final = UserAPIKeyAuth(team_id=UI_SESSION_TOKEN_TEAM_ID, user_id="admin-user", user_role="proxy_admin") + request: Final = _build_request(path="/mcp-rest/tools/list", method="GET") + + result: Final = await rest_endpoints.list_tool_rest_api( + request, server_id=None, mcp_server_name=None, toolset_name=None, user_api_key_dict=session + ) + + assert result["tools"] == [] + assert listed_for == [session] + + async def test_tools_call_rejects_unlisted_clients_before_reading_the_body( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", _CLIENT_ALLOWLIST_SETTINGS, raising=False) + acting: Final = AsyncMock() + monkeypatch.setattr(rest_endpoints, "acting_user_auth", acting, raising=False) + request: Final = _build_request( + {"x-mcp-client": "antigravity-cli"}, + path="/mcp-rest/tools/call", + method="POST", + json_body={"server_id": "server-1", "name": "demo-tool", "arguments": {}}, + ) + + with pytest.raises(HTTPException) as denied: + await rest_endpoints.call_tool_rest_api( + request, user_api_key_dict=UserAPIKeyAuth(jwt_claims={"azp": "claude-code"}) + ) + + assert denied.value.status_code == 403 + assert denied.value.detail["error"] == "Forbidden" + assert "'claude-code'" in denied.value.detail["details"] + acting.assert_not_awaited()