From cade3e312f3f324997ac68676f8d63c8054dff1e Mon Sep 17 00:00:00 2001 From: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Date: Mon, 21 Sep 2026 12:22:18 -0700 Subject: [PATCH] fix(mcp): reject connection grants after delegated mode changes --- .../mcp_server/auth/user_api_key_auth_mcp.py | 1 + .../mcp_server/gateway_dcr_flow.py | 6 +++++- .../auth/test_user_api_key_auth_mcp.py | 21 ++++++++++++++++--- .../mcp_server/test_discoverable_endpoints.py | 9 +++++--- 4 files changed, 30 insertions(+), 7 deletions(-) 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 0b2e42e8e45..8ce7ef2b20c 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 @@ -568,6 +568,7 @@ class MCPRequestHandler: if ( target is None or target.server_id not in allowed + or not target.is_gateway_managed_oauth2 or not target.needs_user_oauth_token or target.oauth_identity_binding is not None ): diff --git a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py index 2e53e8c76b3..f02ab09ebbd 100644 --- a/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py +++ b/litellm/proxy/_experimental/mcp_server/gateway_dcr_flow.py @@ -1638,7 +1638,11 @@ async def validate_connection_binding(request: Request, binding: ConnectionBindi ) if server is None or server.server_id not in allowed: raise HTTPException(status_code=403, detail="Key is not allowed to access the selected MCP server") - if not server.needs_user_oauth_token or server.oauth_identity_binding is not None: + if ( + not server.is_gateway_managed_oauth2 + or not server.needs_user_oauth_token + or server.oauth_identity_binding is not None + ): raise HTTPException(status_code=400, detail="Server does not support a keyed connection grant") return server diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py index 92be8d411ed..d7ab4628cc6 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/auth/test_user_api_key_auth_mcp.py @@ -9450,7 +9450,10 @@ class TestScopedSessionAdmission: ], ) @pytest.mark.parametrize("grant_state", ["valid", "expired", "refresh", "oversized"]) -async def test_connection_credential_requires_exact_key_and_server(monkeypatch, key, selector, expected, grant_state): +@pytest.mark.parametrize("server_mode", ["managed", "delegated"]) +async def test_connection_credential_requires_exact_key_and_server( + monkeypatch, key, selector, expected, grant_state, server_mode +): import json from litellm.proxy._experimental.mcp_server import gateway_dcr_flow as flow from litellm.proxy._experimental.mcp_server.auth import user_api_key_auth_mcp as admission @@ -9470,6 +9473,7 @@ async def test_connection_credential_requires_exact_key_and_server(monkeypatch, transport="http", auth_type=MCPAuth.oauth2, oauth2_flow="authorization_code", + delegate_auth_to_upstream=server_mode == "delegated", ) monkeypatch.setitem(global_mcp_server_manager.registry, server.server_id, server) auth = UserAPIKeyAuth( @@ -9499,7 +9503,13 @@ async def test_connection_credential_requires_exact_key_and_server(monkeypatch, if grant_state == "refresh" else issued["access_token"] ) - expected_status = 401 if expected == 200 and grant_state != "valid" else expected + expected_status = ( + 403 + if server_mode == "delegated" and key is not None + else 401 + if expected == 200 and grant_state != "valid" + else expected + ) scope = { "type": "http", "method": "POST", @@ -9516,7 +9526,12 @@ async def test_connection_credential_requires_exact_key_and_server(monkeypatch, with pytest.raises(HTTPException) as exc: await MCPRequestHandler.process_mcp_request(scope) assert exc.value.status_code == expected_status - if key == "sk-original" and selector == "connection-target" and grant_state != "valid": + if ( + server_mode == "managed" + and key == "sk-original" + and selector == "connection-target" + and grant_state != "valid" + ): assert "resource_metadata=" in exc.value.headers["www-authenticate"] assert flow.CONNECTION_SCOPE_KEY not in scope return diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py index e25932b1a8c..d07bfc4d31e 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_discoverable_endpoints.py @@ -12991,7 +12991,7 @@ def test_keyed_connection_rejects_wrong_endpoint_and_grant_type(keyed_oauth_clie @pytest.mark.asyncio -@pytest.mark.parametrize("mode", ["wrong_resource", "m2m"]) +@pytest.mark.parametrize("mode", ["wrong_resource", "m2m", "delegated"]) async def test_keyed_connection_binding_rejects_resource_or_server_mode(keyed_oauth_client, monkeypatch, mode): from starlette.requests import Request from litellm.proxy import proxy_server @@ -13016,7 +13016,9 @@ async def test_keyed_connection_binding_rejects_resource_or_server_mode(keyed_oa monkeypatch.setitem( global_mcp_server_manager.registry, harness.server.server_id, - harness.server.model_copy(update={"oauth2_flow": "client_credentials"}), + harness.server.model_copy( + update={"delegate_auth_to_upstream": True} if mode == "delegated" else {"oauth2_flow": "client_credentials"} + ), ) request = Request( { @@ -13025,7 +13027,8 @@ async def test_keyed_connection_binding_rejects_resource_or_server_mode(keyed_oa "method": "GET", "path": "/authorize", "headers": [(b"host", b"wrong.example" if mode == "wrong_resource" else b"gateway.example")], - } + }, + receive=AsyncMock(return_value={"type": "http.request", "body": b"", "more_body": False}), ) with pytest.raises(HTTPException) as exc: await harness.real_validate(request, harness.binding)