diff --git a/litellm/integrations/azure_sentinel/azure_sentinel.py b/litellm/integrations/azure_sentinel/azure_sentinel.py index f23317ae9df..563f815b582 100644 --- a/litellm/integrations/azure_sentinel/azure_sentinel.py +++ b/litellm/integrations/azure_sentinel/azure_sentinel.py @@ -16,7 +16,10 @@ import asyncio import os import time import traceback +from collections.abc import Mapping +from types import MappingProxyType from typing import Final +from urllib.parse import urlparse from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -27,6 +30,16 @@ from litellm.llms.custom_httpx.http_handler import ( ) from litellm.types.utils import StandardAuditLogPayload, StandardLoggingPayload +DEFAULT_AZURE_AUTHORITY_HOST: Final = "https://login.microsoftonline.com" +DEFAULT_AZURE_MONITOR_SCOPE: Final = "https://monitor.azure.com/.default" + +MONITOR_SCOPE_BY_AUTHORITY_HOST: Final[Mapping[str, str]] = MappingProxyType( + { + "login.microsoftonline.com": DEFAULT_AZURE_MONITOR_SCOPE, + "login.microsoftonline.us": "https://monitor.azure.us/.default", + } +) + class AzureSentinelLogger(CustomBatchLogger): """ @@ -42,6 +55,7 @@ class AzureSentinelLogger(CustomBatchLogger): client_id: str | None = None, client_secret: str | None = None, audit_stream_name: str | None = None, + authority_host: str | None = None, **kwargs, ): """ @@ -62,6 +76,10 @@ class AzureSentinelLogger(CustomBatchLogger): If not provided, will use AZURE_SENTINEL_CLIENT_SECRET or AZURE_CLIENT_SECRET env var. audit_stream_name (str, optional): Stream name from DCR for audit logs. If not provided, will use AZURE_SENTINEL_AUDIT_STREAM_NAME env var or the standard stream name. + authority_host (str, optional): Microsoft Entra authority host that issues the OAuth2 token, + e.g. "https://login.microsoftonline.us" for Azure Government. If not provided, will use + AZURE_AUTHORITY_HOST env var or default to the Azure Public Cloud authority. The Azure + Monitor audience is derived from it. """ self.async_httpx_client = get_async_httpx_client(llm_provider=httpxSpecialProvider.LoggingCallback) @@ -76,6 +94,9 @@ class AzureSentinelLogger(CustomBatchLogger): resolved_client_secret: Final = ( client_secret or os.getenv("AZURE_SENTINEL_CLIENT_SECRET") or os.getenv("AZURE_CLIENT_SECRET") ) + resolved_authority_host: Final = self._normalize_authority_host( + authority_host or os.getenv("AZURE_AUTHORITY_HOST") or DEFAULT_AZURE_AUTHORITY_HOST + ) if not resolved_dcr_immutable_id: raise ValueError( @@ -119,7 +140,8 @@ class AzureSentinelLogger(CustomBatchLogger): ) # OAuth2 scope for Azure Monitor - self.oauth_scope = "https://monitor.azure.com/.default" + self.authority_host = resolved_authority_host + self.oauth_scope = self._resolve_oauth_scope(authority_host=resolved_authority_host) self.oauth_token: str | None = None self.oauth_token_expires_at: float | None = None @@ -129,6 +151,26 @@ class AzureSentinelLogger(CustomBatchLogger): self.log_queue: list[StandardLoggingPayload] = [] self.audit_log_queue: list[StandardAuditLogPayload] = [] + @staticmethod + def _normalize_authority_host(authority_host: str) -> str: + """ + Normalize an authority host into an absolute URL with no trailing slash. + + Accepts the scheme-qualified form litellm documents ("https://login.microsoftonline.us") + and the bare-host form the azure-identity AzureAuthorityHosts constants use. + """ + stripped: Final = authority_host.strip().rstrip("/") + return stripped if "://" in stripped else f"https://{stripped}" + + @staticmethod + def _resolve_oauth_scope(authority_host: str) -> str: + """ + Map an authority host to the Azure Monitor Logs Ingestion audience for the same cloud, + falling back to the Azure Public Cloud audience for an unrecognized host. + """ + host: Final = urlparse(authority_host).hostname or "" + return MONITOR_SCOPE_BY_AUTHORITY_HOST.get(host, DEFAULT_AZURE_MONITOR_SCOPE) + @staticmethod def _build_api_endpoint(endpoint: str, dcr_immutable_id: str, stream_name: str) -> str: return f"{endpoint.rstrip('/')}/dataCollectionRules/{dcr_immutable_id}/streams/{stream_name}?api-version=2023-01-01" @@ -150,7 +192,7 @@ class AzureSentinelLogger(CustomBatchLogger): assert self.client_id is not None, "client_id is required" assert self.client_secret is not None, "client_secret is required" - token_url: Final = f"https://login.microsoftonline.com/{self.tenant_id}/oauth2/v2.0/token" + token_url: Final = f"{self.authority_host}/{self.tenant_id}/oauth2/v2.0/token" token_data: Final = { "client_id": self.client_id, diff --git a/tests/test_litellm/integrations/test_azure_sentinel.py b/tests/test_litellm/integrations/test_azure_sentinel.py index 55e462c82b8..56662eea633 100644 --- a/tests/test_litellm/integrations/test_azure_sentinel.py +++ b/tests/test_litellm/integrations/test_azure_sentinel.py @@ -296,3 +296,96 @@ async def test_azure_sentinel_audit_stream_name_from_env_var(monkeypatch): ) assert explicit_logger.audit_stream_name == "Custom-LiteLLM-Explicit" + + +def _build_logger(**overrides): + kwargs = { + "dcr_immutable_id": "dcr-test123456789", + "endpoint": "https://test-dce.eastus-1.ingest.monitor.azure.com", + "tenant_id": "test-tenant-id", + "client_id": "test-client-id", + "client_secret": "test-client-secret", + **overrides, + } + with patch("asyncio.create_task", side_effect=_close_periodic_flush_task): + return AzureSentinelLogger(**kwargs) + + +@pytest.fixture +def _no_authority_host_env(monkeypatch): + monkeypatch.delenv("AZURE_AUTHORITY_HOST", raising=False) + + +@pytest.mark.parametrize( + "authority_host, expected_authority, expected_scope", + [ + (None, "https://login.microsoftonline.com", "https://monitor.azure.com/.default"), + ("https://login.microsoftonline.us", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("https://login.microsoftonline.us/", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("login.microsoftonline.us", "https://login.microsoftonline.us", "https://monitor.azure.us/.default"), + ("https://adfs.contoso.example", "https://adfs.contoso.example", "https://monitor.azure.com/.default"), + ], +) +def test_azure_sentinel_resolves_authority_host_and_audience_together( + _no_authority_host_env, authority_host, expected_authority, expected_scope +): + """Both the Entra authority and the Azure Monitor audience must follow the configured cloud. + + Moving only the authority leaves a sovereign deployment asking sovereign Entra for the + commercial audience, which the sovereign ingestion endpoint rejects. + """ + logger = _build_logger(**({} if authority_host is None else {"authority_host": authority_host})) + + assert logger.authority_host == expected_authority + assert logger.oauth_scope == expected_scope + + +def test_azure_sentinel_authority_host_from_env_var(_no_authority_host_env, monkeypatch): + """AZURE_AUTHORITY_HOST is the documented setting and the string callback constructs the logger + with no arguments, so the env var alone has to move both values.""" + monkeypatch.setenv("AZURE_AUTHORITY_HOST", "https://login.microsoftonline.us") + + logger = _build_logger() + + assert logger.authority_host == "https://login.microsoftonline.us" + assert logger.oauth_scope == "https://monitor.azure.us/.default" + + +@pytest.mark.asyncio +async def test_azure_sentinel_token_request_uses_sovereign_authority_and_audience(_no_authority_host_env): + """The resolved values must reach the wire, not just the instance attributes.""" + logger = _build_logger(authority_host="https://login.microsoftonline.us") + logger.log_queue.append( + StandardLoggingPayload( + id="test_id", + call_type="completion", + model="gpt-3.5-turbo", + status="success", + messages=[{"role": "user", "content": "Hello"}], + response={"choices": [{"message": {"content": "Hi"}}]}, + ) + ) + + mock_token_response = MagicMock() + mock_token_response.status_code = 200 + mock_token_response.json = MagicMock(return_value={"access_token": "test-bearer-token", "expires_in": 3600}) + mock_token_response.text = "Success" + mock_api_response = MagicMock() + mock_api_response.status_code = 204 + mock_api_response.text = "Success" + + async def mock_post(*args, **kwargs): + if "oauth2/v2.0/token" in kwargs.get("url", ""): + return mock_token_response + return mock_api_response + + logger.async_httpx_client.post = AsyncMock(side_effect=mock_post) + + await logger.async_send_batch() + + token_calls = [ + call for call in logger.async_httpx_client.post.call_args_list if "oauth2/v2.0/token" in call.kwargs["url"] + ] + assert len(token_calls) == 1 + assert token_calls[0].kwargs["url"] == "https://login.microsoftonline.us/test-tenant-id/oauth2/v2.0/token" + assert token_calls[0].kwargs["data"]["scope"] == "https://monitor.azure.us/.default"