fix(azure_ai): normalize MAI usage in generation response handler

Apply normalize_mai_image_usage before building ImageResponse so token-based
cost calculation works when Azure returns num_output_tokens fields.

Co-authored-by: Cursor <cursoragent@cursor.com>

litellm/llms/azure_ai/image_generation/mai_transformation.py

@@ -200,6 +200,9 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
     ) -> ImageResponse:
         response = raw_response.json()
 
+        if "usage" in response:
+            response["usage"] = self.normalize_mai_image_usage(response.get("usage"))
+
         logging_obj.post_call(
             input=request_data.get("prompt", ""),
             api_key=api_key,

tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py

@@ -1,6 +1,8 @@
 import os
 import sys
+from unittest.mock import MagicMock
 
+import httpx
 import pytest
 
 sys.path.insert(0, os.path.abspath("../../../../../.."))
@@ -209,6 +211,36 @@ class TestAzureMAIImageGeneration:
                 drop_params=False,
             )
 
+    def test_transform_image_generation_response_normalizes_mai_usage(self):
+        config = AzureFoundryMAIImageGenerationConfig()
+        raw_response = MagicMock(spec=httpx.Response)
+        raw_response.json.return_value = {
+            "created": 1780897477,
+            "data": [{"b64_json": "abc123"}],
+            "usage": {
+                "num_output_tokens": 1024,
+                "num_input_text_tokens": 22,
+                "output_image_tokens": 1024,
+            },
+        }
+
+        logging_obj = MagicMock()
+        image_response = config.transform_image_generation_response(
+            model="MAI-Image-2.5",
+            raw_response=raw_response,
+            model_response=ImageResponse(),
+            logging_obj=logging_obj,
+            request_data={"prompt": "A red fox"},
+            optional_params={"width": 1024, "height": 1024},
+            litellm_params={},
+            encoding=None,
+        )
+
+        assert image_response.data[0].b64_json == "abc123"
+        assert image_response.usage.output_tokens == 1024
+        assert image_response.usage.input_tokens == 22
+        assert image_response.usage.total_tokens == 1046
+
     def test_mai_image_cost_calculator_token_based(self):
         os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True"
         litellm.model_cost = litellm.get_model_cost_map(url="")