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

@@ -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
+                    ),
                 )
             )
 

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(

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