fix(azure_sentinel): respect AZURE_AUTHORITY_HOST and derive the Azure Monitor audience per cloud

litellm/integrations/azure_sentinel/azure_sentinel.py

@@ -16,7 +16,10 @@
 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.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 @@ def __init__(
         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 @@ def __init__(
                 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 @@ def __init__(
         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 @@ def __init__(
         )
 
         # 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 @@ def __init__(
         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 @@ async def _get_oauth_token(self) -> str:
         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,

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"