From 60440ee4d3e309a780c88ebbf550603e29e71cef Mon Sep 17 00:00:00 2001 From: tin-berri Date: Sat, 5 Sep 2026 18:04:35 -0700 Subject: [PATCH] feat(mcp): add opt-in per-server oauth relay discovery (#39936) Resolves LIT-7074 --- .../migration.sql | 2 + .../litellm_proxy_extras/schema.prisma | 1 + litellm/models/mcp_server.py | 1 + .../mcp_server/auth/user_api_key_auth_mcp.py | 2 +- .../mcp_server/discoverable_endpoints.py | 20 ++++-- .../mcp_server/mcp_server_manager.py | 34 +++++++++ litellm/proxy/_lazy_openapi_snapshot.json | 35 +++++++++ litellm/proxy/_types.py | 43 +++++++++++ .../mcp_management_endpoints.py | 22 ++++++ litellm/proxy/schema.prisma | 1 + .../types/mcp_server/mcp_server_manager.py | 11 +++ schema.prisma | 1 + .../auth/test_user_api_key_auth_mcp.py | 1 + .../mcp_server/test_db_credentials.py | 72 +++++++++++++++++++ .../mcp_server/test_discoverable_endpoints.py | 12 +++- .../mcp_server/test_mcp_server_manager.py | 44 ++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 ++++ 17 files changed, 309 insertions(+), 8 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql new file mode 100644 index 00000000000..2fa1234bbfc --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260905120000_add_per_server_oauth_discovery_to_mcp_servers/migration.sql @@ -0,0 +1,2 @@ +-- AlterTable +ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 1c43668f227..06ac177cca4 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable { delegate_auth_to_upstream Boolean @default(false) oauth_passthrough Boolean @default(false) dcr_bridge Boolean? + per_server_oauth_discovery Boolean @default(false) is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/models/mcp_server.py b/litellm/models/mcp_server.py index 6cc4a765e46..6bf21a19896 100644 --- a/litellm/models/mcp_server.py +++ b/litellm/models/mcp_server.py @@ -98,6 +98,7 @@ class LiteLLM_MCPServerTable(LiteLLMPydanticObjectBase): delegate_auth_to_upstream: bool = False oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = Field(default_factory=list) byok_api_key_help_url: str | None = None 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 b9bfb062ec7..ad9622c9cd0 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 @@ -197,7 +197,7 @@ def _gateway_dcr_challenge_target( if targets is None: return None server: Final = global_mcp_server_manager.get_mcp_server_by_name(targets[0], client_ip=client_ip) - if server is None or not server.is_gateway_managed_oauth2: + if server is None or not server.advertises_gateway_authorization_server: return None return targets[0] diff --git a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py index 3da5950ce7f..e10bfd41ed6 100644 --- a/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py +++ b/litellm/proxy/_experimental/mcp_server/discoverable_endpoints.py @@ -1481,11 +1481,19 @@ async def _persist_dcr_client_registration( ) updated_row: Final = await update_mcp_server( prisma_client=prisma_client, - data=UpdateMCPServerRequest( - server_id=mcp_server.server_id, - credentials=credentials, - oauth2_flow="authorization_code", - **({"token_url": mcp_server.token_url} if mcp_server.token_url else {}), + data=( + UpdateMCPServerRequest( + server_id=mcp_server.server_id, + credentials=credentials, + oauth2_flow="authorization_code", + token_url=mcp_server.token_url, + ) + if mcp_server.token_url + else UpdateMCPServerRequest( + server_id=mcp_server.server_id, + credentials=credentials, + oauth2_flow="authorization_code", + ) ), touched_by="mcp_oauth_dcr", ) @@ -2367,7 +2375,7 @@ async def _build_oauth_protected_resource_response( if mcp_server is None or mcp_server.auth_type != MCPAuth.oauth2_token_exchange: _raise_unless_oauth2_discovery_server(mcp_server, mcp_server_name, "not an OAuth-protected resource") - if explicitly_named and mcp_server is not None and mcp_server.is_gateway_managed_oauth2: + if explicitly_named and mcp_server is not None and mcp_server.advertises_gateway_authorization_server: return { "authorization_servers": [f"{request_base_url}/mcp"], "resource": resource_url, diff --git a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py index dc1e8db1628..612596bc803 100644 --- a/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py +++ b/litellm/proxy/_experimental/mcp_server/mcp_server_manager.py @@ -155,6 +155,7 @@ from litellm.proxy._types import ( MCPTransportType, SpecialMCPServerNames, UserAPIKeyAuth, + is_per_server_oauth_discovery_eligible, ) from litellm.proxy.auth.ip_address_utils import IPAddressUtils from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper @@ -344,6 +345,7 @@ class MCPServerConfig(TypedDict, total=False): token_endpoint_auth_method: MCPTokenEndpointAuthMethod scopes: str | Sequence[str] dcr_bridge: object + per_server_oauth_discovery: ReadOnly[object] extra_headers: _StringList allowed_tools: _StringList disallowed_tools: _StringList @@ -414,6 +416,31 @@ def _blank_to_none(value: str | None) -> str | None: return value.strip() or None +def _config_per_server_oauth_discovery( + server_config: MCPServerConfig, + server_ref: str, + auth_type: MCPAuthType | None, + oauth2_flow: object, +) -> bool: + match server_config.get("per_server_oauth_discovery", False): + case bool() as enabled: + pass + case other: + raise ValueError( + f"Invalid config for MCP server '{server_ref}': per_server_oauth_discovery must be a boolean " + f"(got {other!r})." + ) + relay_eligible: Final = is_per_server_oauth_discovery_eligible( + auth_type, oauth2_flow, server_config.get("delegate_auth_to_upstream", False) + ) + if enabled and not relay_eligible: + raise ValueError( + f"Invalid config for MCP server '{server_ref}': per_server_oauth_discovery is only supported for " + "auth_type oauth2 with oauth2_flow authorization_code and without delegate_auth_to_upstream." + ) + return enabled + + def _pinned_config_server_id(raw_server_id: object, server_name: str) -> str | None: """Return the ``server_id`` an admin pinned for this config.yaml server, or ``None`` when absent. @@ -2307,6 +2334,9 @@ class MCPServerManager: ) config_dcr_bridge = server_config.get("dcr_bridge", None) + config_per_server_oauth_discovery = _config_per_server_oauth_discovery( + server_config, server_name or server_id, auth_type, config_oauth2_flow + ) if config_dcr_bridge is not None and not isinstance(config_dcr_bridge, bool): raise ValueError( f"Invalid config for MCP server '{server_name or server_id}': dcr_bridge " @@ -2378,6 +2408,7 @@ class MCPServerManager: delegate_auth_to_upstream=bool(server_config.get("delegate_auth_to_upstream", False)), oauth_passthrough=bool(server_config.get("oauth_passthrough", False)), dcr_bridge=config_dcr_bridge, + per_server_oauth_discovery=config_per_server_oauth_discovery, # AWS SigV4 fields aws_access_key_id=server_config.get("aws_access_key_id", None), aws_secret_access_key=server_config.get("aws_secret_access_key", None), @@ -2903,6 +2934,7 @@ class MCPServerManager: delegate_auth_to_upstream=bool(getattr(mcp_server, "delegate_auth_to_upstream", False)), oauth_passthrough=bool(getattr(mcp_server, "oauth_passthrough", False)), dcr_bridge=getattr(mcp_server, "dcr_bridge", None), + per_server_oauth_discovery=bool(getattr(mcp_server, "per_server_oauth_discovery", False)), created_at=getattr(mcp_server, "created_at", None), updated_at=getattr(mcp_server, "updated_at", None), tool_name_to_display_name=_deserialize_json_dict(getattr(mcp_server, "tool_name_to_display_name", None)), @@ -6692,6 +6724,7 @@ class MCPServerManager: registration_url=server.configured_registration_url or server.registration_url, oauth2_flow=server.oauth2_flow, dcr_bridge=server.dcr_bridge, + per_server_oauth_discovery=server.per_server_oauth_discovery, token_exchange_endpoint=server.token_exchange_endpoint, audience=server.audience, subject_token_type=server.subject_token_type, @@ -6810,6 +6843,7 @@ class MCPServerManager: delegate_auth_to_upstream=server.delegate_auth_to_upstream, oauth_passthrough=getattr(server, "oauth_passthrough", False), dcr_bridge=server.dcr_bridge, + per_server_oauth_discovery=server.per_server_oauth_discovery, is_byok=server.is_byok, byok_description=server.byok_description, byok_api_key_help_url=server.byok_api_key_help_url, diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index ea757edb215..91a97ad6544 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -16355,6 +16355,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -17978,6 +17983,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -18859,6 +18869,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -20868,6 +20883,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -22342,6 +22362,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -22862,6 +22887,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { @@ -25371,6 +25401,11 @@ "title": "Oauth Passthrough", "type": "boolean" }, + "per_server_oauth_discovery": { + "default": false, + "title": "Per Server Oauth Discovery", + "type": "boolean" + }, "registration_url": { "anyOf": [ { diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f83011835fd..c28ac8848ba 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1379,6 +1379,35 @@ def _dcr_bridge_auth_type_error(auth_type: object) -> ValueError: ) +def _per_server_oauth_discovery_error() -> ValueError: + return ValueError( + "per_server_oauth_discovery is only supported for auth_type oauth2 with oauth2_flow " + "authorization_code and without delegate_auth_to_upstream." + ) + + +def is_per_server_oauth_discovery_eligible( + auth_type: object, oauth2_flow: object, delegate_auth_to_upstream: object +) -> bool: + return auth_type == MCPAuth.oauth2 and oauth2_flow == "authorization_code" and not delegate_auth_to_upstream + + +def _reject_unsupported_per_server_oauth_discovery(values: object, require_auth_type: bool) -> None: + """Partial updates may omit eligibility fields; those are checked against the stored row by the + update endpoint. Every field the payload does carry must be eligible on its own.""" + if not isinstance(values, dict) or not values.get("per_server_oauth_discovery"): + return + auth_type_ok: Final = values.get("auth_type") == MCPAuth.oauth2 or ( + not require_auth_type and "auth_type" not in values + ) + oauth2_flow_ok: Final = values.get("oauth2_flow") == "authorization_code" or ( + not require_auth_type and "oauth2_flow" not in values + ) + if auth_type_ok and oauth2_flow_ok and not values.get("delegate_auth_to_upstream"): + return + raise _per_server_oauth_discovery_error() + + class NewMCPServerRequest(LiteLLMPydanticObjectBase): server_id: str | None = None server_name: str | None = None @@ -1420,6 +1449,7 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): delegate_auth_to_upstream: bool = False oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = Field(default_factory=list) byok_api_key_help_url: str | None = None @@ -1484,6 +1514,12 @@ class NewMCPServerRequest(LiteLLMPydanticObjectBase): return values raise _dcr_bridge_auth_type_error(auth_type) + @model_validator(mode="before") + @classmethod + def validate_per_server_oauth_discovery_auth_type(cls, values: object) -> object: + _reject_unsupported_per_server_oauth_discovery(values, require_auth_type=True) + return values + class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): server_id: str @@ -1526,6 +1562,7 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): delegate_auth_to_upstream: bool = False oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = Field(default_factory=list) byok_api_key_help_url: str | None = None @@ -1570,6 +1607,12 @@ class UpdateMCPServerRequest(LiteLLMPydanticObjectBase): return values raise _dcr_bridge_auth_type_error(auth_type) + @model_validator(mode="before") + @classmethod + def validate_per_server_oauth_discovery_auth_type(cls, values: object) -> object: + _reject_unsupported_per_server_oauth_discovery(values, require_auth_type=False) + return values + from litellm.models.mcp_server import ( # noqa: E402 LiteLLM_MCPServerTable as LiteLLM_MCPServerTable, diff --git a/litellm/proxy/management_endpoints/mcp_management_endpoints.py b/litellm/proxy/management_endpoints/mcp_management_endpoints.py index ae266792391..d5c3427f29a 100644 --- a/litellm/proxy/management_endpoints/mcp_management_endpoints.py +++ b/litellm/proxy/management_endpoints/mcp_management_endpoints.py @@ -193,6 +193,7 @@ if MCP_AVAILABLE: UpdateMCPServerRequest, UserAPIKeyAuth, UserMCPManagementMode, + is_per_server_oauth_discovery_eligible, ) from litellm.proxy.auth.user_api_key_auth import ( _user_api_key_auth_builder, @@ -2714,6 +2715,27 @@ if MCP_AVAILABLE: old_server_record = None old_server_record_read_failed = True + if payload.per_server_oauth_discovery and (old_server_record is not None or old_server_record_read_failed): + relay_eligible: Final = old_server_record is not None and is_per_server_oauth_discovery_eligible( + payload.auth_type if "auth_type" in payload_fields_set else old_server_record.auth_type, + payload.oauth2_flow if "oauth2_flow" in payload_fields_set else old_server_record.oauth2_flow, + ( + payload.delegate_auth_to_upstream + if "delegate_auth_to_upstream" in payload_fields_set + else old_server_record.delegate_auth_to_upstream + ), + ) + if not relay_eligible: + raise HTTPException( + status_code=status.HTTP_400_BAD_REQUEST, + detail={ # mutable-ok: FastAPI HTTPException detail requires a plain dict + "error": ( + "per_server_oauth_discovery is only supported for auth_type oauth2 with oauth2_flow " + "authorization_code and without delegate_auth_to_upstream." + ) + }, + ) + if ( payload.dcr_bridge and payload.auth_type is None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 1c43668f227..06ac177cca4 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable { delegate_auth_to_upstream Boolean @default(false) oauth_passthrough Boolean @default(false) dcr_bridge Boolean? + per_server_oauth_discovery Boolean @default(false) is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? diff --git a/litellm/types/mcp_server/mcp_server_manager.py b/litellm/types/mcp_server/mcp_server_manager.py index 9bf3acc601c..84ffd50eea1 100644 --- a/litellm/types/mcp_server/mcp_server_manager.py +++ b/litellm/types/mcp_server/mcp_server_manager.py @@ -158,6 +158,7 @@ class MCPServer(BaseModel): # be set explicitly to avoid regressing servers that did not opt in. oauth_passthrough: bool = False dcr_bridge: bool | None = None + per_server_oauth_discovery: bool = False is_byok: bool = False byok_description: list[str] = [] byok_api_key_help_url: str | None = None @@ -241,6 +242,16 @@ class MCPServer(BaseModel): so they are excluded by construction.""" return self.auth_type == MCPAuth.oauth2 and not self.delegate_auth_to_upstream + @property + def uses_per_server_oauth_relay(self) -> bool: + """Whether named discovery should advertise the configured per-server OAuth relay.""" + return self.per_server_oauth_discovery and self.auth_type == MCPAuth.oauth2 and not self.has_client_credentials + + @property + def advertises_gateway_authorization_server(self) -> bool: + """Whether named discovery should advertise the aggregate gateway authorization server.""" + return self.is_gateway_managed_oauth2 and not self.uses_per_server_oauth_relay + @property def is_true_passthrough(self) -> bool: """True for the transparent-proxy mode: LiteLLM performs no admission auth and forwards the diff --git a/schema.prisma b/schema.prisma index 1c43668f227..06ac177cca4 100644 --- a/schema.prisma +++ b/schema.prisma @@ -343,6 +343,7 @@ model LiteLLM_MCPServerTable { delegate_auth_to_upstream Boolean @default(false) oauth_passthrough Boolean @default(false) dcr_bridge Boolean? + per_server_oauth_discovery Boolean @default(false) is_byok Boolean @default(false) byok_description String[] @default([]) byok_api_key_help_url String? 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 f1e299802fb..44eb7795659 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 @@ -7184,6 +7184,7 @@ class TestAggregateGatewayDcrChallenge: cases = [ (_server(MCPAuth.oauth2), "srv"), + (_server(MCPAuth.oauth2, per_server_oauth_discovery=True), None), (_server(MCPAuth.oauth2, oauth2_flow="client_credentials"), "srv"), (_server(MCPAuth.oauth2, delegate_auth_to_upstream=True), None), (_server(MCPAuth.oauth2_token_exchange), None), diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py index 4d9142ad4c5..e20f6646310 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_db_credentials.py @@ -1362,3 +1362,75 @@ async def test_refresh_user_oauth_token_uses_admin_entered_token_url_when_issuer assert result is not None assert captured["url"] == "https://idp.example.com/token" + + +def test_prepare_mcp_server_data_carries_per_server_oauth_discovery(): + request = NewMCPServerRequest( + server_name="relay_create", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + per_server_oauth_discovery=True, + ) + + data = _prepare_mcp_server_data(request) + + assert data["per_server_oauth_discovery"] is True + + +def test_prepare_mcp_server_data_update_carries_per_server_oauth_discovery(): + request = UpdateMCPServerRequest( + server_id="relay-update", + url="https://upstream.example.com/mcp", + transport=MCPTransport.http, + auth_type=MCPAuth.oauth2, + oauth2_flow="authorization_code", + per_server_oauth_discovery=True, + ) + + data = _prepare_mcp_server_data(request, exclude_unset=True) + + assert data["per_server_oauth_discovery"] is True + + +@pytest.mark.parametrize( + "request_cls, extra, overrides", + [ + (NewMCPServerRequest, {"server_name": "relay_create"}, {"auth_type": MCPAuth.oauth_delegate}), + (NewMCPServerRequest, {"server_name": "relay_create"}, {"oauth2_flow": "client_credentials"}), + (UpdateMCPServerRequest, {"server_id": "relay-update"}, {"delegate_auth_to_upstream": True}), + ], +) +def test_request_models_reject_unsupported_per_server_oauth_discovery(request_cls, extra, overrides): + payload = { + "url": "https://upstream.example.com/mcp", + "transport": MCPTransport.http, + "auth_type": MCPAuth.oauth2, + "oauth2_flow": "authorization_code", + "per_server_oauth_discovery": True, + **extra, + **overrides, + } + + with pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"): + request_cls(**payload) + + +@pytest.mark.parametrize( + "partial_payload", + [ + {"oauth2_flow": "client_credentials"}, + {"delegate_auth_to_upstream": True}, + {"auth_type": MCPAuth.api_key}, + ], +) +def test_partial_update_rejects_ineligible_field_alongside_per_server_oauth_discovery(partial_payload): + with pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"): + UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True, **partial_payload) + + +def test_partial_update_defers_omitted_eligibility_fields_to_the_stored_row(): + request = UpdateMCPServerRequest(server_id="relay-update", per_server_oauth_discovery=True) + + assert request.per_server_oauth_discovery is True 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 775f6e5f3b8..a9bb24bbef9 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 @@ -3342,12 +3342,13 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa mock_request.headers = {} interactive = _oauth2_server("github_mcp") + relay = _oauth2_server("relay_mcp", per_server_oauth_discovery=True) m2m = _oauth2_server("m2m_mcp", oauth2_flow="client_credentials", client_id="cid", client_secret="cs") delegated = _oauth2_server("delegated_mcp", delegate_auth_to_upstream=True) global_mcp_server_manager.registry.clear() try: - for server in (interactive, m2m, delegated): + for server in (interactive, relay, m2m, delegated): global_mcp_server_manager.registry[server.server_id] = server for name in ("github_mcp", "m2m_mcp"): @@ -3363,6 +3364,15 @@ async def test_oauth_protected_resource_gateway_managed_oauth2_advertises_gatewa assert legacy["authorization_servers"] == ["https://litellm.example.com/mcp"], name assert legacy["resource"] == f"https://litellm.example.com/{name}/mcp" + relay_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name="relay_mcp", use_standard_pattern=True + ) + assert relay_response["authorization_servers"] == ["https://litellm.example.com/relay_mcp"] + relay_legacy_response = await _build_oauth_protected_resource_response( + request=mock_request, mcp_server_name="relay_mcp", use_standard_pattern=False + ) + assert relay_legacy_response["authorization_servers"] == ["https://litellm.example.com/relay_mcp"] + delegated_response = await _build_oauth_protected_resource_response( request=mock_request, mcp_server_name="delegated_mcp", use_standard_pattern=True ) diff --git a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py index e3ed48713aa..764e2bb0e99 100644 --- a/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py +++ b/tests/test_litellm/proxy/_experimental/mcp_server/test_mcp_server_manager.py @@ -1229,6 +1229,50 @@ class TestMCPServerManager: base.update(overrides) return {"bridgeserver": base} + @pytest.mark.asyncio + async def test_load_servers_from_config_accepts_per_server_oauth_discovery_for_oauth2(self): + manager = MCPServerManager() + + with patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)): + await manager.load_servers_from_config( + self._oauth2_config(oauth2_flow="authorization_code", per_server_oauth_discovery=True) + ) + + server = next(iter(manager.config_mcp_servers.values())) + assert server.per_server_oauth_discovery is True + assert server.uses_per_server_oauth_relay is True + assert server.advertises_gateway_authorization_server is False + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "config", + [ + {"auth_type": MCPAuth.oauth_delegate}, + {"oauth2_flow": "client_credentials"}, + {"oauth2_flow": "authorization_code", "delegate_auth_to_upstream": True}, + ], + ) + async def test_load_servers_from_config_rejects_unsupported_per_server_oauth_discovery(self, config): + manager = MCPServerManager() + + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + pytest.raises(ValueError, match="per_server_oauth_discovery is only supported"), + ): + await manager.load_servers_from_config(self._oauth2_config(per_server_oauth_discovery=True, **config)) + + @pytest.mark.asyncio + async def test_load_servers_from_config_rejects_non_boolean_per_server_oauth_discovery(self): + manager = MCPServerManager() + + with ( + patch.object(manager, "_descovery_metadata", new=AsyncMock(return_value=None)), + pytest.raises(ValueError, match="per_server_oauth_discovery.*must be a boolean"), + ): + await manager.load_servers_from_config( + self._oauth2_config(oauth2_flow="authorization_code", per_server_oauth_discovery="yes") + ) + @pytest.mark.asyncio async def test_load_servers_from_config_rejects_dcr_bridge_on_gateway_managed_auth_type(self): manager = MCPServerManager() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8a850598259..b1534c19670 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -28870,6 +28870,11 @@ export interface components { * @default false */ oauth_passthrough: boolean; + /** + * Per Server Oauth Discovery + * @default false + */ + per_server_oauth_discovery: boolean; /** Registration Url */ registration_url?: string | null; /** Review Notes */ @@ -32010,6 +32015,11 @@ export interface components { * @default false */ oauth_passthrough: boolean; + /** + * Per Server Oauth Discovery + * @default false + */ + per_server_oauth_discovery: boolean; /** Registration Url */ registration_url?: string | null; /** Server Id */ @@ -37833,6 +37843,11 @@ export interface components { * @default false */ oauth_passthrough: boolean; + /** + * Per Server Oauth Discovery + * @default false + */ + per_server_oauth_discovery: boolean; /** Registration Url */ registration_url?: string | null; /** Server Id */