From 482f3c37ed6de423daec8e65f36142cc5408d808 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Tue, 3 Nov 2026 00:02:21 -0800 Subject: [PATCH] fix(azure_ai): keep azure_ad_token out of the request body (#49769) * fix(azure_ai): keep azure_ad_token out of the request body azure_ad_token is not a litellm-owned param, so get_optional_params routes it into extra_body and the azure_ai chat and Claude transformations merge it into the JSON body. Azure AI Foundry then rejects the call with 400 Unknown parameter / Extra inputs are not permitted. The Authorization header already comes from litellm_params, so drop the extra_body copy once at the top of the azure_ai dispatch, as the azure branches do. Type litellm.open_ai_chat_completion_models and open_ai_text_completion_models as Set[str] so the basedpyright gate does not grow from the untyped optional_params read. * fix(azure_ai): keep Entra credentials out of every azure_ai request body * fix(azure_ai): keep the Entra token out of the image edit form * fix(azure_ai): keep Entra credentials out of image edit extra_body too --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/images/main.py | 11 +- litellm/utils.py | 14 +- .../llms/azure_ai/test_azure_ai_entra_auth.py | 164 ++++++++++++++++++ 3 files changed, 184 insertions(+), 5 deletions(-) diff --git a/litellm/images/main.py b/litellm/images/main.py index f37b252fa45..27917cf9289 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -28,7 +28,7 @@ from litellm.litellm_core_utils.mock_functions import mock_image_generation from litellm.llms.base_llm import BaseImageEditConfig, BaseImageGenerationConfig from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler -from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params +from litellm.utils import exception_type, filter_out_litellm_params, get_litellm_params, header_auth_params #################### Initialize provider clients #################### llm_http_handler: BaseLLMHTTPHandler = BaseLLMHTTPHandler() @@ -988,10 +988,15 @@ def image_edit( or custom_llm_provider == "azure" or custom_llm_provider in litellm.openai_compatible_providers ): + dropped_form_params: Final = header_auth_params(custom_llm_provider) image_edit_request_params.update( flatten_form_field_values( - non_default_params, - extra_body if isinstance(extra_body, dict) else None, + {key: value for key, value in non_default_params.items() if key not in dropped_form_params}, + ( + {key: value for key, value in extra_body.items() if key not in dropped_form_params} + if isinstance(extra_body, dict) + else None + ), ) ) diff --git a/litellm/utils.py b/litellm/utils.py index 108c6fae982..d881289deed 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6161,6 +6161,14 @@ EXTRA_BODY_ROUTING_KEYS: Final = frozenset({"model"}) LITELLM_ONLY_PARAMS: Final = frozenset(("cache_control_injection_points",)) +def header_auth_params(custom_llm_provider: str) -> frozenset[str]: + if custom_llm_provider != "azure_ai": + return frozenset() + from litellm.llms.azure_ai.common_utils import AZURE_ENTRA_CREDENTIAL_PARAM_KEYS + + return AZURE_ENTRA_CREDENTIAL_PARAM_KEYS + + def add_provider_specific_params_to_optional_params( optional_params: dict, passed_params: dict, @@ -6195,8 +6203,10 @@ def add_provider_specific_params_to_optional_params( **extra_body, } - dropped_keys: Final = EXTRA_BODY_ROUTING_KEYS | frozenset( - param for param in (additional_drop_params or ()) if isinstance(param, str) + dropped_keys: Final = ( + EXTRA_BODY_ROUTING_KEYS + | header_auth_params(custom_llm_provider) + | frozenset(param for param in (additional_drop_params or ()) if isinstance(param, str)) ) ensure_extra_body_is_safe: Final = litellm_utils.ensure_extra_body_is_safe optional_params["extra_body"] = ensure_extra_body_is_safe( diff --git a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py index 908e727bc4c..d0c1c4a6fb2 100644 --- a/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py +++ b/tests/unit/llms/azure_ai/test_azure_ai_entra_auth.py @@ -5,9 +5,15 @@ Every azure_ai route must authenticate with an Entra ID token when no API key is instead of requiring an API key. """ +import json +import queue +from collections.abc import Mapping +from typing import Final from unittest.mock import patch +import httpx import pytest +import respx import litellm from litellm.llms.azure_ai.common_utils import ( @@ -16,6 +22,8 @@ from litellm.llms.azure_ai.common_utils import ( has_azure_entra_params, resolve_azure_ai_agent_auth_header, ) +from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_httpx.http_handler import HTTPHandler ENTRA_PARAMS = {"azure_ad_token": "entra-token"} @@ -290,3 +298,159 @@ async def test_agent_auth_header_is_the_entra_bearer(): headers = await resolve_azure_ai_agent_auth_header({"azure_ad_token": "entra-token"}) assert headers == {"Authorization": "Bearer entra-token"} + + +OPENAI_CHAT_COMPLETION: Final = { + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "gpt-6.1-sol", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, +} +OPENAI_EMBEDDING: Final = { + "object": "list", + "data": [{"object": "embedding", "index": 0, "embedding": [0.1, 0.2]}], + "model": "cohere-embed-v3-english", + "usage": {"prompt_tokens": 1, "total_tokens": 1}, +} +ANTHROPIC_MESSAGE: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "OK"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} + + +class _PreCallOptionalParamsCapture(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.optional_params: tuple[dict[str, object], ...] = () + + def log_pre_api_call(self, model: str, messages: list[object], kwargs: dict[str, dict[str, object]]) -> None: + self.optional_params = (*self.optional_params, kwargs["optional_params"]) + + +@pytest.mark.parametrize( + ("model", "provider_response"), + [ + pytest.param("azure_ai/gpt-6.1-sol", OPENAI_CHAT_COMPLETION, id="openai-shaped"), + pytest.param("azure_ai/claude-sonnet-5-5", ANTHROPIC_MESSAGE, id="claude"), + ], +) +def test_completion_sends_the_entra_token_only_as_the_bearer_header( + model: str, provider_response: Mapping[str, object], monkeypatch: pytest.MonkeyPatch +): + capture: Final = _PreCallOptionalParamsCapture() + monkeypatch.setattr(litellm, "callbacks", [capture]) + sent: Final[queue.Queue[httpx.Request]] = queue.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + sent.put_nowait(request) + return httpx.Response(200, json=provider_response) + + response: Final = litellm.completion( + model=model, + api_base="https://example-resource.services.ai.azure.com", + azure_ad_token="entra-token", + messages=[{"role": "user", "content": "hi"}], + client=HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))), + ) + + request: Final = sent.get_nowait() + assert request.headers["Authorization"] == "Bearer entra-token" + assert "azure_ad_token" not in json.loads(request.content) + assert response.choices[0].message.content == "OK" + assert capture.optional_params + assert "entra-token" not in json.dumps(capture.optional_params, default=str) + + +@pytest.mark.respx(assert_all_called=True) +def test_embedding_sends_the_entra_token_only_as_the_bearer_header(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json=OPENAI_EMBEDDING) + ) + + litellm.embedding( + model="azure_ai/cohere-embed-v3-english", + api_base="https://example-resource.services.ai.azure.com", + azure_ad_token="entra-token", + input=["hello"], + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert "azure_ad_token" not in json.loads(request.content) + + +IMAGE_RESPONSE: Final = {"created": 1, "data": [{"b64_json": "aGk="}]} +PNG_BYTES: Final = b"\x89PNG\r\n\x1a\n" + b"\x00" * 32 + + +@pytest.mark.respx(assert_all_called=True) +@pytest.mark.parametrize( + "extra_kwargs", + [{}, {"extra_body": {"azure_ad_token": "entra-token", "client_secret": "entra-token"}}], + ids=["kwarg", "extra_body"], +) +def test_image_edit_sends_the_entra_token_only_as_the_bearer_header( + respx_mock: respx.MockRouter, extra_kwargs: dict[str, dict[str, str]] +): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json=IMAGE_RESPONSE) + ) + + litellm.image_edit( + model="azure_ai/FLUX.1-Kontext-pro", + api_base="https://example-resource.services.ai.azure.com", + api_version="2025-04-01-preview", + azure_ad_token="entra-token", + image=PNG_BYTES, + prompt="add a hat", + **extra_kwargs, + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert b"entra-token" not in request.content + + +@pytest.mark.respx(assert_all_called=True) +def test_image_generation_sends_the_entra_token_only_as_the_bearer_header(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json=IMAGE_RESPONSE) + ) + + litellm.image_generation( + model="azure_ai/FLUX-1.1-pro", + api_base="https://example-resource.services.ai.azure.com", + azure_ad_token="entra-token", + prompt="a red circle", + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert "azure_ad_token" not in json.loads(request.content) + + +@pytest.mark.respx(assert_all_called=True) +def test_transcription_sends_the_entra_token_only_as_the_bearer_header(respx_mock: respx.MockRouter): + route: Final = respx_mock.post(url__startswith="https://example-resource.services.ai.azure.com").mock( + return_value=httpx.Response(200, json={"text": "hi"}) + ) + + litellm.transcription( + model="azure_ai/whisper", + api_base="https://example-resource.services.ai.azure.com", + api_version="2024-06-01", + azure_ad_token="entra-token", + file=("hello.wav", b"RIFF" + b"\x00" * 32), + ) + + request: Final = route.calls.last.request + assert request.headers["Authorization"] == "Bearer entra-token" + assert b"entra-token" not in request.content