mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-10 03:28:53 +00:00
Add support for azure entra discovery endpoint
This commit is contained in:
Sameer Kankute
parent
2b8b614120
commit
3e4f9af955
2 changed files with 166 additions and 5 deletions
|
|
@ -1488,11 +1488,18 @@ 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 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",
|
||||
|
|
@ -1674,6 +1681,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("/"))
|
||||
|
|
@ -1713,7 +1723,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 != "login.microsoftonline.com"
|
||||
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