Merge pull request #26584 from BerriAI/litellm_mcp-oauth-azure-entra-discovery2

[Feat]Add support for azure entra discovery endpoint
This commit is contained in:
Mateo Wang 2026-04-29 14:28:41 -07:00 committed by GitHub
commit 9bc317b4d0
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 184 additions and 7 deletions

View file

@ -109,6 +109,12 @@ if not _separator_probe.is_valid:
SEP_986_URL,
)
_AZURE_ENTRA_HOSTS = {
"login.microsoftonline.com", # Global
"login.microsoftonline.us", # US Government
"login.chinacloudapi.cn", # China
}
def _warn_on_server_name_fields(
*,
@ -1503,11 +1509,28 @@ class MCPServerManager:
client = get_async_httpx_client(llm_provider=httpxSpecialProvider.MCP)
response = await client.get(server_url)
response.raise_for_status()
verbose_logger.warning(
"MCP OAuth discovery unexpectedly succeeded for %s; server did not challenge",
server_url,
(
authorization_servers,
resource_scopes,
) = await self._attempt_well_known_discovery(server_url)
metadata = await self._fetch_authorization_server_metadata(
authorization_servers
)
raise RuntimeError("OAuth discovery must not succeed without a challenge")
if (
metadata is None
and not resource_scopes
and authorization_servers
and response.status_code == 200
):
verbose_logger.warning(
"MCP OAuth discovery for %s received 200 OK without RFC 9728 challenge and no discoverable authorization metadata.",
server_url,
)
if metadata is None and resource_scopes:
return MCPOAuthMetadata(scopes=resource_scopes)
if metadata is not None and resource_scopes:
metadata.scopes = resource_scopes
return metadata
except HTTPStatusError as exc:
verbose_logger.debug(
"MCP OAuth discovery for %s received status error: %s",
@ -1525,8 +1548,8 @@ class MCPServerManager:
header_value
)
authorization_servers: List[str] = []
resource_scopes: Optional[List[str]] = None
authorization_servers = []
resource_scopes = None
if resource_metadata_url:
(
authorization_servers,
@ -1689,6 +1712,9 @@ class MCPServerManager:
f"{base}/.well-known/oauth-authorization-server/{path}"
)
candidate_urls.append(f"{base}/.well-known/openid-configuration/{path}")
candidate_urls.append(
f"{issuer_url.rstrip('/')}/.well-known/openid-configuration"
)
candidate_urls.append(f"{base}/.well-known/oauth-authorization-server")
candidate_urls.append(f"{base}/.well-known/openid-configuration")
candidate_urls.append(issuer_url.rstrip("/"))
@ -1728,7 +1754,28 @@ class MCPServerManager:
):
return metadata
return None
return self._build_azure_authorization_server_metadata(parsed)
@staticmethod
def _build_azure_authorization_server_metadata(
parsed_issuer_url: Any,
) -> Optional[MCPOAuthMetadata]:
path_parts = [
part for part in (parsed_issuer_url.path or "").split("/") if part
]
if (
parsed_issuer_url.netloc not in _AZURE_ENTRA_HOSTS
or len(path_parts) != 2
or path_parts[1] != "v2.0"
):
return None
tenant = path_parts[0]
base = f"{parsed_issuer_url.scheme}://{parsed_issuer_url.netloc}/{tenant}"
return MCPOAuthMetadata(
authorization_url=f"{base}/oauth2/v2.0/authorize",
token_url=f"{base}/oauth2/v2.0/token",
)
@staticmethod
def _decrypt_credential_field(

View file

@ -728,6 +728,136 @@ class TestMCPServerManager:
]
assert scopes == ["read", "write"]
@pytest.mark.asyncio
async def test_descovery_metadata_probes_well_known_when_server_does_not_challenge(
self,
):
manager = MCPServerManager()
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock()
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_response)
mock_metadata = MCPOAuthMetadata(
scopes=None,
authorization_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/authorize",
token_url="https://login.microsoftonline.com/tenant/oauth2/v2.0/token",
registration_url=None,
)
with (
patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_client,
),
patch.object(
manager,
"_attempt_well_known_discovery",
AsyncMock(
return_value=(
["https://login.microsoftonline.com/test-tenant-id/v2.0"],
["api://some-scope/.default"],
)
),
) as mock_well_known,
patch.object(
manager,
"_fetch_authorization_server_metadata",
AsyncMock(return_value=mock_metadata),
) as mock_fetch_auth,
):
result = await manager._descovery_metadata("http://localhost:8001/mcp")
mock_well_known.assert_awaited_once_with("http://localhost:8001/mcp")
mock_fetch_auth.assert_awaited_once_with(
["https://login.microsoftonline.com/test-tenant-id/v2.0"]
)
assert result is mock_metadata
assert result.scopes == ["api://some-scope/.default"]
@pytest.mark.asyncio
async def test_fetch_single_authorization_server_metadata_supports_azure_issuer_path(
self,
):
manager = MCPServerManager()
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
def build_response(url: str):
mock_response = MagicMock()
if url == f"{issuer}/.well-known/openid-configuration":
mock_response.json.return_value = {
"authorization_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize",
"token_endpoint": "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token",
"scopes_supported": ["api://some-scope/.default"],
}
mock_response.raise_for_status = MagicMock()
else:
request = httpx.Request("GET", url)
response_obj = httpx.Response(status_code=404, request=request)
mock_response.raise_for_status = MagicMock(
side_effect=httpx.HTTPStatusError(
"not found", request=request, response=response_obj
)
)
return mock_response
mock_client = MagicMock()
mock_client.get = AsyncMock(side_effect=build_response)
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_client,
):
result = await manager._fetch_single_authorization_server_metadata(issuer)
assert result is not None
assert (
result.authorization_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
)
assert (
result.token_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
)
assert result.scopes == ["api://some-scope/.default"]
@pytest.mark.asyncio
async def test_fetch_single_authorization_server_metadata_derives_azure_metadata(
self,
):
manager = MCPServerManager()
issuer = "https://login.microsoftonline.com/test-tenant-id/v2.0"
request = httpx.Request("GET", issuer)
response_obj = httpx.Response(status_code=404, request=request)
mock_response = MagicMock()
mock_response.raise_for_status = MagicMock(
side_effect=httpx.HTTPStatusError(
"not found", request=request, response=response_obj
)
)
mock_client = MagicMock()
mock_client.get = AsyncMock(return_value=mock_response)
with patch(
"litellm.proxy._experimental.mcp_server.mcp_server_manager.get_async_httpx_client",
return_value=mock_client,
):
result = await manager._fetch_single_authorization_server_metadata(issuer)
assert result is not None
assert (
result.authorization_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/authorize"
)
assert (
result.token_url
== "https://login.microsoftonline.com/test-tenant-id/oauth2/v2.0/token"
)
@pytest.mark.asyncio
async def test_descovery_metadata_falls_back_to_origin_when_no_auth_servers(self):
manager = MCPServerManager()