mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
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:
commit
9bc317b4d0
2 changed files with 184 additions and 7 deletions
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue