From b40a7b5b536722f39734798b7983ab31872832d4 Mon Sep 17 00:00:00 2001 From: Stephen Sennett Date: Thu, 10 Sep 2026 01:13:55 +1000 Subject: [PATCH 1/4] feat(guardrails): allow Entra ID auth for Azure Content Safety Azure Content Safety accepts either an API key or a Microsoft Entra ID bearer token, but the prompt shield and text moderation guardrails only ever sent Ocp-Apim-Subscription-Key and refused to start without a key. Organisations that disable local auth on the resource could not use them at all api_key is now optional. When it is omitted the guardrail authenticates with the standard credential chain (env service principal, workload identity, managed identity, az login) for the Cognitive Services scope and sends Authorization: Bearer instead. api_base stays required, since Entra auth needs the resource's custom subdomain endpoint The credential is resolved once at startup, so a deployment missing azure-identity fails at boot rather than on its first guarded request, and a key-based deployment never imports azure-identity at all Co-Authored-By: Claude Opus 5 --- .../guardrail_hooks/azure/__init__.py | 11 +- .../guardrails/guardrail_hooks/azure/base.py | 53 ++++- .../guardrail_hooks/azure/prompt_shield.py | 9 +- .../guardrail_hooks/azure/text_moderation.py | 9 +- .../guardrails/guardrail_hooks/azure/base.py | 6 +- .../guardrail_hooks/azure/conftest.py | 44 ++++ .../guardrail_hooks/azure/test_azure_base.py | 188 ++++++++++++++++++ .../azure/test_azure_prompt_shield.py | 55 +++++ .../azure/test_azure_text_moderation.py | 23 +++ .../guardrails/test_guardrail_endpoints.py | 6 +- 10 files changed, 388 insertions(+), 16 deletions(-) create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py create mode 100644 tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py index de77cd55671..4f5b15df3d6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py @@ -1,3 +1,4 @@ +from collections.abc import Callable from typing import TYPE_CHECKING, Final from litellm.types.guardrails import SupportedGuardrailIntegrations @@ -9,11 +10,13 @@ if TYPE_CHECKING: from litellm.types.guardrails import Guardrail, LitellmParams -def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail"): +def initialize_guardrail( + litellm_params: "LitellmParams", + guardrail: "Guardrail", + entra_token_provider: Callable[[], str] | None = None, +) -> AzureContentSafetyPromptShieldGuardrail | AzureContentSafetyTextModerationGuardrail: import litellm - if not litellm_params.api_key: - raise ValueError("Azure Content Safety: api_key is required") if not litellm_params.api_base: raise ValueError("Azure Content Safety: api_base is required") @@ -32,6 +35,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" **litellm_params.model_dump(exclude_none=True), "api_key": litellm_params.api_key, "api_base": litellm_params.api_base, + "entra_token_provider": entra_token_provider, "default_on": litellm_params.default_on, "event_hook": litellm_params.mode, }, @@ -43,6 +47,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" **litellm_params.model_dump(exclude_none=True), "api_key": litellm_params.api_key, "api_base": litellm_params.api_base, + "entra_token_provider": entra_token_provider, "default_on": litellm_params.default_on, "event_hook": litellm_params.mode, }, diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 42f0220cc4d..3989cc9934e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -1,4 +1,7 @@ +import asyncio import re +from collections.abc import Callable +from functools import cache from typing import TYPE_CHECKING, Any, Final from litellm._logging import verbose_proxy_logger @@ -9,6 +12,9 @@ from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, ) +from litellm.secret_managers.get_azure_ad_token_provider import ( + get_azure_ad_token_provider, +) if TYPE_CHECKING: from litellm.types.llms.openai import AllMessageValues @@ -20,6 +26,20 @@ AZURE_CONTENT_SAFETY_MAX_TEXT_LENGTH: Final = 10000 # chunk of N characters consumes ceil(N / 1000) text records. AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000 +AZURE_CONTENT_SAFETY_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default" + + +@cache +def _default_entra_token_provider() -> Callable[[], str]: + """Entra token provider for the Content Safety data plane, one credential per process.""" + try: + return get_azure_ad_token_provider(azure_scope=AZURE_CONTENT_SAFETY_ENTRA_SCOPE) + except ImportError as e: + raise ValueError( + "Azure Content Safety: api_key is not set and azure-identity is not installed. " + "Set api_key, or install azure-identity to authenticate with Microsoft Entra ID" + ) from e + class AzureGuardrailBase: """ @@ -32,8 +52,10 @@ class AzureGuardrailBase: def __init__( self, - api_key: str, + *, api_base: str, + api_key: str | None = None, + entra_token_provider: Callable[[], str] | None = None, **kwargs: Any, ): # Forward remaining kwargs to the next class in the MRO @@ -44,6 +66,30 @@ class AzureGuardrailBase: self.api_key = api_key self.api_base = api_base self.api_version: str = kwargs.get("api_version") or "2024-09-01" + self._entra_token_provider: Final = entra_token_provider or ( + None if api_key else _default_entra_token_provider() + ) + + async def _auth_header(self) -> tuple[str, str]: + """Credential header name and value for a single request. + + Azure Content Safety accepts an API key or a Microsoft Entra token and rejects each + on the other's header, so exactly one is sent. + """ + if self.api_key: + return ("Ocp-Apim-Subscription-Key", self.api_key) + + minter: Final = self._entra_token_provider or _default_entra_token_provider() + try: + token: Final = await asyncio.to_thread(minter) + except Exception as e: + verbose_proxy_logger.exception("Azure Content Safety: Entra token request failed") + raise ValueError( + "Azure Content Safety: no credential available. Set api_key, or configure an Entra " + "identity (AZURE_CLIENT_ID / AZURE_CLIENT_SECRET / AZURE_TENANT_ID, workload identity, " + "managed identity, or az login) holding the Cognitive Services User role on the resource" + ) from e + return ("Authorization", f"Bearer {token}") async def _post_to_content_safety(self, endpoint_path: str, request_body: dict[str, object]) -> dict[str, Any]: """POST to an Azure Content Safety endpoint with standard auth headers. @@ -57,8 +103,9 @@ class AzureGuardrailBase: Parsed JSON response dict. """ url: Final = f"{self.api_base}/contentsafety/{endpoint_path}?api-version={self.api_version}" - headers: Final = { - "Ocp-Apim-Subscription-Key": self.api_key, + auth_name, auth_value = await self._auth_header() + headers: Final = { # mutable-ok: AsyncHTTPHandler.post types headers as `dict | None` + auth_name: auth_value, "Content-Type": "application/json", } diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py index 6e29d44662e..60247b9c4b1 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/prompt_shield.py @@ -120,8 +120,8 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai Configuration: guardrail_name: Name of the guardrail instance - api_key: Azure Prompt Shield API key - api_base: Azure Prompt Shield API endpoint + api_base: Azure Content Safety endpoint + api_key: Azure Content Safety API key. Omit it to authenticate with Microsoft Entra ID default_on: Whether to enable by default """ @@ -129,17 +129,18 @@ class AzureContentSafetyPromptShieldGuardrail(AzureGuardrailBase, CustomGuardrai def __init__( self, + *, guardrail_name: str, - api_key: str, api_base: str, + api_key: str | None = None, **kwargs, ): """Initialize Azure Prompt Shield guardrail handler.""" # AzureGuardrailBase.__init__ stores api_key, api_base, api_version, # async_handler and forwards the rest to CustomGuardrail. super().__init__( - api_key=api_key, api_base=api_base, + api_key=api_key, guardrail_name=guardrail_name, supported_event_hooks=list(self.get_supported_event_hooks()), **kwargs, diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py index 0dca8be3307..53bfac17910 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/text_moderation.py @@ -37,8 +37,8 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr Configuration: guardrail_name: Name of the guardrail instance - api_key: Azure Text Moderation API key - api_base: Azure Text Moderation API endpoint + api_base: Azure Content Safety endpoint + api_key: Azure Content Safety API key. Omit it to authenticate with Microsoft Entra ID default_on: Whether to enable by default """ @@ -55,9 +55,10 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr def __init__( self, + *, guardrail_name: str, - api_key: str, api_base: str, + api_key: str | None = None, severity_threshold: int | None = None, severity_threshold_by_category: dict[str, int] | None = None, **kwargs, @@ -71,8 +72,8 @@ class AzureContentSafetyTextModerationGuardrail(AzureGuardrailBase, CustomGuardr # AzureGuardrailBase.__init__ stores api_key, api_base, api_version, # async_handler and forwards the rest to CustomGuardrail. super().__init__( - api_key=api_key, api_base=api_base, + api_key=api_key, guardrail_name=guardrail_name, **kwargs, ) diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py index c72888f33e1..861b1d2db02 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/azure/base.py @@ -6,7 +6,11 @@ class AzureContentSafetyConfigModel(BaseModel): api_key: str | None = Field( default=None, - description="API key for the Azure Content Safety Prompt Shield guardrail", + description=( + "API key for the Azure Content Safety resource. Optional: omit it to authenticate with " + "Microsoft Entra ID, which needs the Cognitive Services User role on the resource and a " + "custom subdomain api_base" + ), ) api_base: str | None = Field( diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py new file mode 100644 index 00000000000..8e3605137ab --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py @@ -0,0 +1,44 @@ +from typing import Final + +import httpx +import pytest + +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler +from litellm.proxy.guardrails.guardrail_hooks.azure.base import ( + _default_entra_token_provider, +) + +CLEAN_RESPONSE_FOR_BOTH_GUARDRAILS: Final = { + "userPromptAnalysis": {"attackDetected": False}, + "documentsAnalysis": [], + "categoriesAnalysis": [], + "blocklistsMatch": [], +} + + +@pytest.fixture(autouse=True) +def clear_default_entra_provider_cache(): + """A credential resolved elsewhere in the session would otherwise be reused here, masking + the dependency and failure paths these tests assert.""" + _default_entra_token_provider.cache_clear() + yield + _default_entra_token_provider.cache_clear() + + +@pytest.fixture +def api_base() -> str: + return "https://contoso.cognitiveservices.azure.com" + + +@pytest.fixture +def capturing_handler() -> tuple[AsyncHTTPHandler, list[httpx.Request]]: + """An HTTP handler answering every Content Safety call, paired with the requests it saw.""" + sent: Final[list[httpx.Request]] = [] + + def _record(request: httpx.Request) -> httpx.Response: + sent.append(request) + return httpx.Response(200, json=CLEAN_RESPONSE_FOR_BOTH_GUARDRAILS) + + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=httpx.MockTransport(_record)) + return handler, sent diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py new file mode 100644 index 00000000000..3179be458e0 --- /dev/null +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py @@ -0,0 +1,188 @@ +import sys +import threading +from typing import Final + +import pytest + +from litellm.proxy.guardrails.guardrail_hooks.azure import base +from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import ( + AzureContentSafetyPromptShieldGuardrail, +) +from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import ( + AzureContentSafetyTextModerationGuardrail, +) + +GUARDRAIL_CLASSES: Final = ( + AzureContentSafetyPromptShieldGuardrail, + AzureContentSafetyTextModerationGuardrail, +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guardrail_class", GUARDRAIL_CLASSES) +async def test_api_key_rides_the_subscription_key_header_and_mints_no_token( + guardrail_class, api_base, capturing_handler +): + handler, sent = capturing_handler + + def _must_not_mint() -> str: + raise AssertionError("Entra token minted despite api_key being set") + + guardrail: Final = guardrail_class( + guardrail_name="azure-guard", + api_base=api_base, + api_key="secret-key", + entra_token_provider=_must_not_mint, + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert len(sent) == 1 + assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" + assert "authorization" not in sent[0].headers + assert sent[0].headers["Content-Type"] == "application/json" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("guardrail_class", GUARDRAIL_CLASSES) +async def test_omitted_api_key_rides_an_entra_bearer_token(guardrail_class, api_base, capturing_handler): + handler, sent = capturing_handler + + guardrail: Final = guardrail_class( + guardrail_name="azure-guard", + api_base=api_base, + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert len(sent) == 1 + assert sent[0].headers["Authorization"] == "Bearer entra-token" + assert "ocp-apim-subscription-key" not in sent[0].headers + assert sent[0].headers["Content-Type"] == "application/json" + + +@pytest.mark.asyncio +async def test_blank_api_key_is_treated_as_absent(api_base, capturing_handler): + """An `os.environ/` api_key resolving to an empty string must reach Entra, not send a blank key.""" + handler, sent = capturing_handler + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=api_base, + api_key="", + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].headers["Authorization"] == "Bearer entra-token" + assert "ocp-apim-subscription-key" not in sent[0].headers + + +@pytest.mark.asyncio +async def test_token_minting_failure_reports_the_options_and_sends_nothing(api_base, capturing_handler): + handler, sent = capturing_handler + + def _fails() -> str: + raise RuntimeError("no managed identity endpoint found") + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=api_base, + entra_token_provider=_fails, + ) + guardrail.async_handler = handler + + with pytest.raises(ValueError, match="no credential available") as exc_info: + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent == [] + assert isinstance(exc_info.value.__cause__, RuntimeError) + + +@pytest.mark.asyncio +async def test_token_minting_failure_keeps_credential_detail_out_of_the_error(api_base, capturing_handler): + """azure-identity errors carry tenant and client ids, and this message reaches the API caller.""" + handler, _ = capturing_handler + + def _fails() -> str: + raise RuntimeError("tenant 11111111-2222-3333-4444-555555555555 rejected the request") + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=api_base, + entra_token_provider=_fails, + ) + guardrail.async_handler = handler + + with pytest.raises(ValueError, match="no credential available") as exc_info: + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert "11111111-2222-3333-4444-555555555555" not in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_token_is_minted_off_the_event_loop_thread(api_base, capturing_handler): + """Credential sources block: IMDS probes time out and the az CLI credential spawns a subprocess.""" + handler, _ = capturing_handler + minting_threads: Final[list[int]] = [] + + def _record_thread() -> str: + minting_threads.append(threading.get_ident()) + return "entra-token" + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=api_base, + entra_token_provider=_record_thread, + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert len(minting_threads) == 1 + assert minting_threads[0] != threading.get_ident() + + +def test_default_credential_is_built_once_per_process(monkeypatch): + """Rebuilding it per request costs a fresh credential and token round trip on every scan.""" + builds: Final[list[str]] = [] + + def _build(azure_scope: str): + builds.append(azure_scope) + return lambda: "entra-token" + + monkeypatch.setattr(base, "get_azure_ad_token_provider", _build) + + assert base._default_entra_token_provider() is base._default_entra_token_provider() + assert builds == [base.AZURE_CONTENT_SAFETY_ENTRA_SCOPE] + + +@pytest.mark.parametrize("guardrail_class", GUARDRAIL_CLASSES) +def test_keyless_guardrail_names_the_missing_dependency_at_startup(guardrail_class, api_base, monkeypatch): + monkeypatch.setitem(sys.modules, "azure", None) + + with pytest.raises(ValueError, match="azure-identity"): + guardrail_class(guardrail_name="azure-guard", api_base=api_base) + + +@pytest.mark.asyncio +async def test_api_key_guardrail_never_reaches_azure_identity(api_base, capturing_handler, monkeypatch): + """azure-identity ships in the proxy extra, so a key-based deployment must not depend on it.""" + handler, sent = capturing_handler + monkeypatch.setitem(sys.modules, "azure", None) + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=api_base, + api_key="secret-key", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py index 17e7222fa44..b421f52d3e0 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py @@ -4,6 +4,7 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.azure import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import ( AzureContentSafetyPromptShieldGuardrail, ) @@ -635,3 +636,57 @@ def test_update_in_memory_litellm_params_dead_env_credential_rejected_untouched( assert guardrail.api_key == "azure_prompt_shield_api_key" assert guardrail.price_per_1000_text_records == 0.38 + + +@pytest.mark.asyncio +async def test_initialize_guardrail_without_api_key_authenticates_with_entra(api_base, capturing_handler): + """A keyless config entry yields a guardrail that authenticates with Entra.""" + handler, sent = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams(guardrail="azure/prompt_shield", mode="pre_call", api_base=api_base), + {"guardrail_name": "azure-prompt-shield"}, + entra_token_provider=lambda: "entra-token", + ) + + assert isinstance(guardrail, AzureContentSafetyPromptShieldGuardrail) + assert guardrail.api_key is None + assert guardrail.api_base == api_base + + guardrail.async_handler = handler + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].headers["Authorization"] == "Bearer entra-token" + + +def test_initialize_guardrail_without_api_base_still_raises(): + """api_base carries the resource's custom subdomain, which Entra auth cannot work without.""" + with pytest.raises(ValueError, match="api_base is required"): + initialize_guardrail( + LitellmParams(guardrail="azure/prompt_shield", mode="pre_call"), + {"guardrail_name": "azure-prompt-shield"}, + entra_token_provider=lambda: "entra-token", + ) + + +@pytest.mark.asyncio +async def test_clearing_api_key_at_runtime_switches_to_entra(api_base, capturing_handler): + """A dashboard edit that removes the key must re-authenticate, not keep sending a stale header.""" + handler, sent = capturing_handler + + guardrail = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure_prompt_shield", + api_base=api_base, + api_key="secret-key", + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + guardrail.update_in_memory_litellm_params({"api_key": None}) + await guardrail.apply_guardrail(inputs={"texts": ["hello again"]}, request_data={}, input_type="request") + + assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" + assert "authorization" not in sent[0].headers + assert sent[1].headers["Authorization"] == "Bearer entra-token" + assert "ocp-apim-subscription-key" not in sent[1].headers diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index a43f95062f9..9df49db199d 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py @@ -4,9 +4,11 @@ import pytest from fastapi import HTTPException from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.guardrails.guardrail_hooks.azure import initialize_guardrail from litellm.proxy.guardrails.guardrail_hooks.azure.text_moderation import ( AzureContentSafetyTextModerationGuardrail, ) +from litellm.types.guardrails import LitellmParams from litellm.types.utils import Choices, Message, ModelResponse @@ -463,3 +465,24 @@ async def test_apply_guardrail_handles_missing_texts_key(): mock_post.assert_not_called() assert result == {"images": ["x"]} + + +@pytest.mark.asyncio +async def test_initialize_guardrail_without_api_key_authenticates_with_entra(api_base, capturing_handler): + """A keyless config entry yields a guardrail that authenticates with Entra.""" + handler, sent = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams(guardrail="azure/text_moderations", mode="pre_call", api_base=api_base), + {"guardrail_name": "azure-text-moderation"}, + entra_token_provider=lambda: "entra-token", + ) + + assert isinstance(guardrail, AzureContentSafetyTextModerationGuardrail) + assert guardrail.api_key is None + assert guardrail.api_base == api_base + + guardrail.async_handler = handler + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].headers["Authorization"] == "Bearer entra-token" diff --git a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py index 530f8ffd854..4d840394601 100644 --- a/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/test_litellm/proxy/guardrails/test_guardrail_endpoints.py @@ -632,7 +632,11 @@ def test_get_provider_specific_params(): # Check the structure of a simple field assert ( fields["api_key"]["description"] - == "API key for the Azure Content Safety Prompt Shield guardrail" + == ( + "API key for the Azure Content Safety resource. Optional: omit it to authenticate with " + "Microsoft Entra ID, which needs the Cognitive Services User role on the resource and a " + "custom subdomain api_base" + ) ) assert fields["api_key"]["required"] == False assert fields["api_key"]["type"] == "string" # Should be string, not None From 19d873799892a19627778e93bf84d4ca5c884e4f Mon Sep 17 00:00:00 2001 From: Stephen Sennett Date: Thu, 10 Sep 2026 01:13:55 +1000 Subject: [PATCH 2/4] fix(guardrails): stop sibling defaults leaking into Azure Content Safety LitellmParams mixes in every guardrail's config model, so dumping it with exclude_none handed the Azure guardrails 37 defaults belonging to other providers. One of them broke config-driven setups outright: Javelin's api_version of "v1" replaced the Content Safety default, and Azure answers api-version=v1 with 404 Resource not found Forward only the params the config actually set, the way the Prisma AIRS initializer already does. A config that omits api_version now gets the Content Safety default of 2024-09-01, and an explicit value still wins The Admin UI was unaffected, because its form submits the api_version it reads from the provider schema. Any other caller that omits api_version hits this, which is why the documented config.yaml example never worked Co-Authored-By: Claude Opus 5 --- .../guardrail_hooks/azure/__init__.py | 4 +- .../azure/test_azure_prompt_shield.py | 59 +++++++++++++++++++ .../azure/test_azure_text_moderation.py | 43 ++++++++++++++ 3 files changed, 104 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py index 4f5b15df3d6..3e278e620b6 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/__init__.py @@ -32,7 +32,7 @@ def initialize_guardrail( ) = AzureContentSafetyPromptShieldGuardrail( guardrail_name=guardrail_name, **{ - **litellm_params.model_dump(exclude_none=True), + **litellm_params.model_dump(exclude_unset=True), "api_key": litellm_params.api_key, "api_base": litellm_params.api_base, "entra_token_provider": entra_token_provider, @@ -44,7 +44,7 @@ def initialize_guardrail( azure_content_safety_guardrail = AzureContentSafetyTextModerationGuardrail( guardrail_name=guardrail_name, **{ - **litellm_params.model_dump(exclude_none=True), + **litellm_params.model_dump(exclude_unset=True), "api_key": litellm_params.api_key, "api_base": litellm_params.api_base, "entra_token_provider": entra_token_provider, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py index b421f52d3e0..c7eab7a80ba 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_prompt_shield.py @@ -690,3 +690,62 @@ async def test_clearing_api_key_at_runtime_switches_to_entra(api_base, capturing assert "authorization" not in sent[0].headers assert sent[1].headers["Authorization"] == "Bearer entra-token" assert "ocp-apim-subscription-key" not in sent[1].headers + + +@pytest.mark.asyncio +async def test_config_without_api_version_uses_the_content_safety_default(api_base, capturing_handler): + """A config that omits api_version gets the Content Safety default, not another guardrail's.""" + handler, sent = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams(guardrail="azure/prompt_shield", mode="pre_call", api_base=api_base), + {"guardrail_name": "azure-prompt-shield"}, + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].url.params["api-version"] == "2024-09-01" + + +@pytest.mark.asyncio +async def test_config_api_version_is_honoured(api_base, capturing_handler): + """Pinning an older Content Safety version stays possible.""" + handler, sent = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams(guardrail="azure/prompt_shield", mode="pre_call", api_base=api_base, api_version="2023-10-01"), + {"guardrail_name": "azure-prompt-shield"}, + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].url.params["api-version"] == "2023-10-01" + + +@pytest.mark.asyncio +async def test_config_pricing_extras_survive_the_forwarded_params(api_base, capturing_handler): + """cost_tier and price_per_1000_text_records arrive as pydantic extras rather than declared + fields, so forwarding only the params the config set must still carry them through.""" + handler, _ = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams( + guardrail="azure/prompt_shield", + mode="pre_call", + api_base=api_base, + cost_tier="paid", + price_per_1000_text_records=0.38, + ), + {"guardrail_name": "azure-prompt-shield"}, + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + request_data = {"metadata": {}} + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data=request_data, input_type="request") + + assert _recorded_guardrail_info(request_data)["guardrail_cost"] == pytest.approx(0.38 / 1000) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py index 9df49db199d..e86c53ed1e4 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_text_moderation.py @@ -1,3 +1,4 @@ +import json from unittest.mock import Mock, patch import pytest @@ -486,3 +487,45 @@ async def test_initialize_guardrail_without_api_key_authenticates_with_entra(api await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") assert sent[0].headers["Authorization"] == "Bearer entra-token" + + +@pytest.mark.asyncio +async def test_config_without_api_version_uses_the_content_safety_default(api_base, capturing_handler): + """A config that omits api_version gets the Content Safety default, not another guardrail's.""" + handler, sent = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams(guardrail="azure/text_moderations", mode="pre_call", api_base=api_base), + {"guardrail_name": "azure-text-moderation"}, + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].url.params["api-version"] == "2024-09-01" + + +@pytest.mark.asyncio +async def test_config_moderation_options_still_reach_the_request(api_base, capturing_handler): + """Forwarding only params the config set must not drop the guardrail's own options.""" + handler, sent = capturing_handler + + guardrail = initialize_guardrail( + LitellmParams( + guardrail="azure/text_moderations", + mode="pre_call", + api_base=api_base, + outputType="EightSeverityLevels", + blocklistNames=["my-blocklist"], + ), + {"guardrail_name": "azure-text-moderation"}, + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + body = json.loads(sent[0].content) + assert body["outputType"] == "EightSeverityLevels" + assert body["blocklistNames"] == ["my-blocklist"] From b5dfba190c6b380aaa4aef45c1dd26924cd0c806 Mon Sep 17 00:00:00 2001 From: Stephen Sennett Date: Thu, 10 Sep 2026 02:18:30 +1000 Subject: [PATCH 3/4] fix(guardrails): only send an Entra token to an Azure endpoint A Microsoft Entra token is scoped to every Cognitive Services resource the identity can reach, not to the one resource a key would unlock, so a typo, a stale DNS name, or a hostile api_base receives far more than a leaked resource key would give away. Check the destination before minting: the Entra path now requires HTTPS and a recognised Azure Content Safety host, and refuses anything else with an error naming api_key as the way to reach another endpoint. The check runs at startup and again per request, so clearing api_key on a running guardrail cannot quietly promote its host The api_key path is untouched. A resource key only unlocks the resource it belongs to, so gateways and test doubles stay reachable with one Co-Authored-By: Claude Opus 5 --- .../guardrails/guardrail_hooks/azure/base.py | 26 +++++++ .../guardrail_hooks/azure/test_azure_base.py | 74 +++++++++++++++++++ 2 files changed, 100 insertions(+) diff --git a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py index 3989cc9934e..c6adc05c191 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/azure/base.py +++ b/litellm/proxy/guardrails/guardrail_hooks/azure/base.py @@ -3,6 +3,7 @@ import re from collections.abc import Callable from functools import cache from typing import TYPE_CHECKING, Any, Final +from urllib.parse import urlparse from litellm._logging import verbose_proxy_logger from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -28,6 +29,28 @@ AZURE_CONTENT_SAFETY_TEXT_RECORD_LENGTH: Final = 1000 AZURE_CONTENT_SAFETY_ENTRA_SCOPE: Final = "https://cognitiveservices.azure.com/.default" +AZURE_CONTENT_SAFETY_ENTRA_HOST_SUFFIXES: Final = ( + ".cognitiveservices.azure.com", + ".cognitiveservices.azure.us", + ".cognitiveservices.azure.cn", + ".services.ai.azure.com", +) + + +def _assert_entra_destination_is_azure(api_base: str) -> None: + """An Entra token is scoped to every Cognitive Services resource the identity can reach, + not to one resource, so it is only ever sent to an Azure endpoint over TLS. Entra also + requires the resource's custom subdomain, so any other host is not a valid destination.""" + parsed: Final = urlparse(api_base) + host: Final = (parsed.hostname or "").lower() + if parsed.scheme == "https" and host.endswith(AZURE_CONTENT_SAFETY_ENTRA_HOST_SUFFIXES): + return + raise ValueError( + f"Azure Content Safety: refusing to send a Microsoft Entra token to api_base {api_base!r}. " + "Entra authentication needs the resource's HTTPS custom subdomain endpoint, for example " + "https://your-resource.cognitiveservices.azure.com. Set api_key to reach any other host" + ) + @cache def _default_entra_token_provider() -> Callable[[], str]: @@ -66,6 +89,8 @@ class AzureGuardrailBase: self.api_key = api_key self.api_base = api_base self.api_version: str = kwargs.get("api_version") or "2024-09-01" + if not api_key: + _assert_entra_destination_is_azure(api_base) self._entra_token_provider: Final = entra_token_provider or ( None if api_key else _default_entra_token_provider() ) @@ -79,6 +104,7 @@ class AzureGuardrailBase: if self.api_key: return ("Ocp-Apim-Subscription-Key", self.api_key) + _assert_entra_destination_is_azure(self.api_base) minter: Final = self._entra_token_provider or _default_entra_token_provider() try: token: Final = await asyncio.to_thread(minter) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py index 3179be458e0..bb0a1ad4c9a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py @@ -186,3 +186,77 @@ async def test_api_key_guardrail_never_reaches_azure_identity(api_base, capturin await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" + + +@pytest.mark.parametrize( + "bad_api_base", + [ + "http://contoso.cognitiveservices.azure.com", + "https://contoso.cognitiveservices.azure.com.attacker.example", + "https://attacker.example", + "https://australiaeast.api.cognitive.microsoft.com", + ], +) +def test_keyless_guardrail_refuses_a_non_azure_destination(bad_api_base): + """The Entra token covers every Cognitive Services resource the identity can reach, so a + typo or a hijacked host would receive far more than one resource's key would give away.""" + with pytest.raises(ValueError, match="refusing to send a Microsoft Entra token") as exc_info: + AzureContentSafetyPromptShieldGuardrail(guardrail_name="azure-guard", api_base=bad_api_base) + + assert "api_key" in str(exc_info.value) + + +@pytest.mark.parametrize( + "good_api_base", + [ + "https://contoso.cognitiveservices.azure.com", + "https://contoso.privatelink.cognitiveservices.azure.com", + "https://contoso.services.ai.azure.com", + "https://contoso.cognitiveservices.azure.us", + ], +) +def test_keyless_guardrail_accepts_azure_content_safety_endpoints(good_api_base): + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base=good_api_base, + entra_token_provider=lambda: "entra-token", + ) + + assert guardrail.api_base == good_api_base + + +@pytest.mark.asyncio +async def test_api_key_guardrail_may_use_any_destination(capturing_handler): + """A key is scoped to one resource, so gateways and test doubles stay reachable with one.""" + handler, sent = capturing_handler + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base="https://gateway.internal.example", + api_key="secret-key", + ) + guardrail.async_handler = handler + + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent[0].headers["Ocp-Apim-Subscription-Key"] == "secret-key" + + +@pytest.mark.asyncio +async def test_clearing_api_key_cannot_send_a_token_to_a_non_azure_destination(capturing_handler): + """A guardrail admitted on its api_key must not start minting tokens for that same host.""" + handler, sent = capturing_handler + + guardrail: Final = AzureContentSafetyPromptShieldGuardrail( + guardrail_name="azure-guard", + api_base="https://gateway.internal.example", + api_key="secret-key", + entra_token_provider=lambda: "entra-token", + ) + guardrail.async_handler = handler + guardrail.update_in_memory_litellm_params({"api_key": None}) + + with pytest.raises(ValueError, match="refusing to send a Microsoft Entra token"): + await guardrail.apply_guardrail(inputs={"texts": ["hello"]}, request_data={}, input_type="request") + + assert sent == [] From 579e0d37f3a1a7bdf4688db88ded8c206d053af2 Mon Sep 17 00:00:00 2001 From: Stephen Sennett Date: Thu, 10 Sep 2026 02:53:02 +1000 Subject: [PATCH 4/4] test(guardrails): mark the Azure capture lists as callee-filled The request, thread and scope logs the new Azure tests capture are accumulators a callback fills, which the no-mutation convention asks to be marked rather than left bare. Mark all three with a reason, and hand the captured requests back as a Sequence so a test cannot append to the fixture's own log Co-Authored-By: Claude Opus 5 --- .../proxy/guardrails/guardrail_hooks/azure/conftest.py | 5 +++-- .../guardrails/guardrail_hooks/azure/test_azure_base.py | 4 ++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py index 8e3605137ab..aad05c2017f 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/conftest.py @@ -1,3 +1,4 @@ +from collections.abc import Sequence from typing import Final import httpx @@ -31,9 +32,9 @@ def api_base() -> str: @pytest.fixture -def capturing_handler() -> tuple[AsyncHTTPHandler, list[httpx.Request]]: +def capturing_handler() -> tuple[AsyncHTTPHandler, Sequence[httpx.Request]]: """An HTTP handler answering every Content Safety call, paired with the requests it saw.""" - sent: Final[list[httpx.Request]] = [] + sent: Final[list[httpx.Request]] = [] # mutable-ok: callee-filled request log, handed back read-only def _record(request: httpx.Request) -> httpx.Response: sent.append(request) diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py index bb0a1ad4c9a..b461aca6392 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/azure/test_azure_base.py @@ -129,7 +129,7 @@ async def test_token_minting_failure_keeps_credential_detail_out_of_the_error(ap async def test_token_is_minted_off_the_event_loop_thread(api_base, capturing_handler): """Credential sources block: IMDS probes time out and the az CLI credential spawns a subprocess.""" handler, _ = capturing_handler - minting_threads: Final[list[int]] = [] + minting_threads: Final[list[int]] = [] # mutable-ok: callee-filled thread log def _record_thread() -> str: minting_threads.append(threading.get_ident()) @@ -150,7 +150,7 @@ async def test_token_is_minted_off_the_event_loop_thread(api_base, capturing_han def test_default_credential_is_built_once_per_process(monkeypatch): """Rebuilding it per request costs a fresh credential and token round trip on every scan.""" - builds: Final[list[str]] = [] + builds: Final[list[str]] = [] # mutable-ok: callee-filled scope log def _build(azure_scope: str): builds.append(azure_scope)