feat(mcp): add opt-in per-server oauth relay discovery (#39936)

Resolves LIT-7074
This commit is contained in:
tin-berri 2026-09-05 18:04:35 -07:00 committed by GitHub
parent 91ae13d07d
commit 60440ee4d3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
17 changed files with 309 additions and 8 deletions

View file

@ -0,0 +1,2 @@
-- AlterTable
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false;

View file

@ -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?

View file

@ -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

View file

@ -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]

View file

@ -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,

View file

@ -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,

View file

@ -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": [
{

View file

@ -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,

View file

@ -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

View file

@ -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?

View file

@ -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

View file

@ -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?

View file

@ -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),

View file

@ -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

View file

@ -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
)

View file

@ -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()

View file

@ -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 */