fix(proxy): store the actual selected model in spend logs for Azure Model Router

litellm/litellm_core_utils/litellm_logging.py

@@ -5930,11 +5930,15 @@ def get_standard_logging_object_payload(
             response_model_name = final_response_obj.get("model")
 
         # For Azure Model Router, preserve the actual model in the top-level standard
-        # logging payload only when the user has opted in.
+        # logging payload.
+        from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
         requested_model: Final = kwargs.get("model")
-        if (
-            isinstance(requested_model, str)
-            and ("model_router" in requested_model.lower() or "model-router" in requested_model.lower())
+        stamped_selected_model: Final = AzureFoundryModelInfo.get_model_router_selected_model(hidden_params)
+        if stamped_selected_model is not None:
+            model_name = stamped_selected_model
+        elif (
+            AzureFoundryModelInfo.is_model_router_call(model=requested_model, hidden_params=hidden_params)
             and isinstance(response_model_name, str)
             and response_model_name
         ):

litellm/llms/azure_ai/azure_model_router/transformation.py

@@ -65,15 +65,24 @@ def transform_response(
 
         Extracts the actual model used from the Azure response (e.g., gpt-5-nano-2025-08-07)
         and returns it with the azure_ai/ prefix for proper display and cost tracking.
+
+        Also stamps that model onto ``_hidden_params`` so downstream consumers (spend logs,
+        response restamping) can read it instead of guessing the route from the model string.
         """
-        from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+        from litellm.llms.azure_ai.common_utils import (
+            AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+            AzureFoundryModelInfo,
+        )
+        from litellm.router_utils.add_retry_fallback_headers import (
+            get_hidden_params_dict,
+        )
 
         # Get base model for the parent call (strips routing prefixes for API compatibility)
         base_model: Final[str] = AzureFoundryModelInfo.get_base_model(model)
 
         # Call parent transform_response first - this will extract the actual model
         # from the raw response (e.g., "gpt-5-nano-2025-08-07")
-        model_response = super().transform_response(
+        transformed_response: Final = super().transform_response(
             model=base_model,
             raw_response=raw_response,
             model_response=model_response,
@@ -86,7 +95,15 @@ def transform_response(
             api_key=api_key,
             json_mode=json_mode,
         )
-        return model_response
+        selected_model: Final = transformed_response.model
+        if selected_model:
+            # Rebuilt rather than mutated in place: ModelResponseBase declares _hidden_params as a
+            # class-level dict, so an in-place write can bleed into unrelated responses.
+            transformed_response._hidden_params = {  # pyright: ignore[reportPrivateUsage]  # ModelResponse exposes no public hidden-params setter  # mutable-ok: ModelResponse requires _hidden_params to be a plain dict
+                **get_hidden_params_dict(transformed_response),
+                AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model,
+            }
+        return transformed_response
 
     def calculate_additional_costs(self, model: str, prompt_tokens: int, completion_tokens: int) -> dict | None:
         """

litellm/llms/azure_ai/common_utils.py

@@ -51,6 +51,9 @@ def get_azure_ai_auth_headers(
     )
 
 
+AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: Final = "azure_model_router_selected_model"
+
+
 class AzureFoundryModelInfo(BaseLLMModelInfo):
     """Model info for Azure AI / Azure Foundry models."""
 
@@ -82,6 +85,41 @@ def get_azure_ai_route(model: str) -> Literal["agents", "model_router", "default
             return "model_router"
         return "default"
 
+    @staticmethod
+    def get_model_router_selected_model(hidden_params: Mapping[str, object] | None) -> str | None:
+        """The model Azure Model Router actually served, stamped by ``AzureModelRouterConfig``.
+
+        Reading this beats re-deriving the route from a model string: the stamp is set on the
+        code path that was actually taken, so it holds no matter what the caller named the model.
+        """
+        if not hidden_params:
+            return None
+        selected: Final = hidden_params.get(AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY)
+        if isinstance(selected, str) and selected:
+            return selected
+        return None
+
+    @staticmethod
+    def is_model_router_call(
+        model: str | None = None,
+        hidden_params: Mapping[str, object] | None = None,
+    ) -> bool:
+        """Whether a request went down the Azure Model Router route.
+
+        Prefers the response stamp, then the deployment's litellm model path, and only then the
+        caller-supplied name. The last two go through ``get_azure_ai_route`` so the model-router
+        name heuristic lives in exactly one place.
+        """
+        if AzureFoundryModelInfo.get_model_router_selected_model(hidden_params) is not None:
+            return True
+        deployment_model: Final = (
+            hidden_params.get("litellm_model_name") or hidden_params.get("model") if hidden_params is not None else None
+        )
+        return any(
+            isinstance(candidate, str) and AzureFoundryModelInfo.get_azure_ai_route(candidate) == "model_router"
+            for candidate in (deployment_model, model)
+        )
+
     @staticmethod
     def get_api_base(api_base: str | None = None) -> str | None:
         return api_base or litellm.api_base or get_secret_str("AZURE_AI_API_BASE")

litellm/proxy/common_request_processing.py

@@ -1138,24 +1138,25 @@ async def keepalive_then_relay() -> AsyncGenerator[bytes, None]:
     )
 
 
-def _is_azure_model_router_request(model: str) -> bool:
+def _is_azure_model_router_request(model: str, hidden_params: Mapping[str, object] | None = None) -> bool:
     """
-    Check if the requested model is an Azure Model Router.
+    Check if a request went down the Azure Model Router route.
 
-    Azure Model Router models follow the pattern:
-    - azure_ai/model_router/<deployment-name>
-    - azure_ai/model-router
-    - model_router/<deployment-name>
-    - model-router
+    ``model`` here is what the *client* sent, a model group alias with no ``model_router/``
+    prefix, so matching on it alone only works when the operator happened to put "model-router"
+    in the alias. Where the response is in hand its stamp answers this outright, so callers
+    should pass ``hidden_params``.
 
     Args:
         model: The requested model name
+        hidden_params: ``_hidden_params`` from the response, when the caller has it
 
     Returns:
         bool: True if this is an Azure Model Router request
     """
-    model_lower: Final = model.lower()
-    return "model-router" in model_lower or "model_router" in model_lower
+    from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo
+
+    return AzureFoundryModelInfo.is_model_router_call(model=model, hidden_params=hidden_params)
 
 
 def _override_openai_response_model(
@@ -1223,7 +1224,7 @@ def _override_openai_response_model(
                 return
 
     # Check if this is an Azure Model Router request - if so, preserve the actual model used
-    if _is_azure_model_router_request(requested_model):
+    if _is_azure_model_router_request(requested_model, hidden_params):
         verbose_proxy_logger.debug(
             "%s: Azure Model Router detected - preserving actual model used from response instead of overriding to router model.",
             log_context,

litellm/proxy/spend_tracking/spend_tracking_utils.py

@@ -444,7 +444,9 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs
         or None
     )
     raw_model: Final = cast(str, kwargs.get("model") or "")
-    model_name: Final = reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
+    model_name: Final = (
+        standard_logging_payload.get("model") if standard_logging_payload is not None else None
+    ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {})
 
     try:
         payload: Final[SpendLogsPayload] = SpendLogsPayload(

ruff-strict-budget.json

@@ -168,7 +168,7 @@
     "limit": 3
   },
   "RET504": {
-    "limit": 176
+    "limit": 175
   },
   "RUF012": {
     "limit": 240

tests/test_litellm/litellm_core_utils/test_litellm_logging.py

@@ -6,7 +6,6 @@
 
 import pytest
 
-
 import time
 
 import httpx
@@ -3867,6 +3866,90 @@ def test_get_standard_logging_object_payload_includes_litellm_call_id(logging_ob
     assert payload["litellm_call_id"] == call_id
 
 
+# ── Azure Model Router selected-model attribution ────────────────────────────
+
+
+def _model_router_response(selected_model: str, stamp: bool):
+    """A ModelResponse as AzureModelRouterConfig hands it back, with or without the stamp."""
+    from litellm.llms.azure_ai.common_utils import (
+        AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+    )
+    from litellm.types.utils import ModelResponse
+
+    response = ModelResponse(model=selected_model)
+    response._hidden_params = (
+        {AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY: selected_model} if stamp else {}
+    )
+    return response
+
+
+def test_standard_logging_payload_uses_stamped_model_router_model(logging_obj):
+    """
+    The selected model must win off the stamp, not off "model-router" appearing in the
+    requested model. An operator whose model group is named anything else was invisible
+    to the name check, so their logs and spend rows named the router instead.
+    """
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import (
+        get_standard_logging_object_payload,
+    )
+
+    now = datetime.datetime.now()
+    payload = get_standard_logging_object_payload(
+        kwargs={
+            "model": "azure_ai/smart-pick",
+            "custom_llm_provider": "azure_ai",
+            "messages": [],
+            "litellm_params": {"metadata": {}},
+        },
+        init_response_obj=_model_router_response(
+            "azure_ai/grok-4-1-fast-reasoning", stamp=True
+        ),
+        start_time=now,
+        end_time=now,
+        logging_obj=logging_obj,
+        status="success",
+    )
+
+    assert payload is not None
+    assert payload["model"] == "azure_ai/grok-4-1-fast-reasoning"
+
+
+def test_standard_logging_payload_keeps_requested_model_without_router_stamp(
+    logging_obj,
+):
+    """
+    Control for the test above: an ordinary azure_ai deployment is unaffected, so the stamp
+    is what redirects attribution rather than the response model winning unconditionally.
+    """
+    import datetime
+
+    from litellm.litellm_core_utils.litellm_logging import (
+        get_standard_logging_object_payload,
+    )
+
+    now = datetime.datetime.now()
+    payload = get_standard_logging_object_payload(
+        kwargs={
+            "model": "azure_ai/smart-pick",
+            "custom_llm_provider": "azure_ai",
+            "messages": [],
+            "litellm_params": {"metadata": {}},
+        },
+        init_response_obj=_model_router_response(
+            "azure_ai/grok-4-1-fast-reasoning", stamp=False
+        ),
+        start_time=now,
+        end_time=now,
+        logging_obj=logging_obj,
+        status="success",
+    )
+
+    assert payload is not None
+    assert payload["model"] == "azure_ai/smart-pick"
+
+
 def _make_dict_logging_obj():
     """Build a Logging instance configured for a non-streaming dict result."""
     obj = LitellmLogging(

tests/test_litellm/llms/azure_ai/chat/test_azure_ai_transformation.py

@@ -201,6 +201,90 @@ def test_azure_model_router_response_shows_actual_model():
     )
 
 
+def test_azure_model_router_stamps_selected_model_on_hidden_params():
+    """
+    The selected model must be stamped on _hidden_params, not left for downstream code to
+    re-derive by looking for "model-router" in the model string. Deployments whose alias
+    does not contain that text are invisible to the string check.
+    """
+    from httpx import Response
+
+    from litellm.llms.azure_ai.common_utils import (
+        AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+        AzureFoundryModelInfo,
+    )
+    from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj
+    from litellm.types.utils import ModelResponse
+
+    raw_response_json = {
+        "id": "chatcmpl-test456",
+        "object": "chat.completion",
+        "created": 1234567890,
+        "model": "grok-4-1-fast-reasoning",
+        "choices": [
+            {
+                "index": 0,
+                "message": {"role": "assistant", "content": "pong"},
+                "finish_reason": "stop",
+            }
+        ],
+        "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15},
+    }
+
+    mock_response = MagicMock(spec=Response)
+    mock_response.json.return_value = raw_response_json
+    mock_response.text = json.dumps(raw_response_json)
+    mock_response.headers = {}
+
+    logging_obj = MagicMock(spec=LiteLLMLoggingObj)
+    logging_obj.post_call = MagicMock()
+    logging_obj.model_call_details = {}
+
+    result = AzureModelRouterConfig().transform_response(
+        model="smart-pick",
+        raw_response=mock_response,
+        model_response=ModelResponse(),
+        logging_obj=logging_obj,
+        request_data={},
+        messages=[{"role": "user", "content": "Reply with just pong"}],
+        optional_params={},
+        litellm_params={"model": "azure_ai/model_router/smart-pick"},
+        encoding=None,
+        api_key="test-key",
+        json_mode=False,
+    )
+
+    assert result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY] == result.model
+    assert (
+        result._hidden_params[AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY]
+        == "azure_ai/grok-4-1-fast-reasoning"
+    )
+    assert AzureFoundryModelInfo.get_model_router_selected_model(
+        result._hidden_params
+    ) == ("azure_ai/grok-4-1-fast-reasoning")
+    assert (
+        AzureFoundryModelInfo.is_model_router_call(
+            model="smart-pick", hidden_params=result._hidden_params
+        )
+        is True
+    )
+
+
+def test_azure_model_router_stamp_does_not_leak_across_responses():
+    """
+    ModelResponse declares _hidden_params as a class-level dict, so the stamp has to be written
+    as a fresh dict. Mutating in place would bleed the selected model into unrelated responses.
+    """
+    from litellm.llms.azure_ai.common_utils import (
+        AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY,
+    )
+    from litellm.types.utils import ModelResponse
+
+    untouched = ModelResponse()
+
+    assert AZURE_MODEL_ROUTER_SELECTED_MODEL_KEY not in (untouched._hidden_params or {})
+
+
 def test_drop_tool_level_extra_fields_strips_copilot_mcp_server_name():
     """
     Regression test: Azure AI returns 400 when tools contain copilot_mcp_server_name.