From 6266daa12e8a72d30e6a303a2dce25d6c2d9a548 Mon Sep 17 00:00:00 2001
From: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
Date: Mon, 8 Jun 2026 20:47:23 +0000
Subject: [PATCH] fix(azure_ai): wrap MAI generation response JSON parsing in
 error handling

---
 .../image_generation/mai_transformation.py    |  8 ++++++-
 .../test_mai_image_generation.py              | 24 +++++++++++++++++++
 2 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/litellm/llms/azure_ai/image_generation/mai_transformation.py b/litellm/llms/azure_ai/image_generation/mai_transformation.py
index fd271ea3a17..c08d44f9031 100644
--- a/litellm/llms/azure_ai/image_generation/mai_transformation.py
+++ b/litellm/llms/azure_ai/image_generation/mai_transformation.py
@@ -5,6 +5,7 @@ import httpx
 from litellm.llms.base_llm.image_generation.transformation import (
     BaseImageGenerationConfig,
 )
+from litellm.llms.openai.common_utils import OpenAIError
 from litellm.types.llms.openai import OpenAIImageGenerationOptionalParams
 from litellm.types.utils import ImageResponse
 from litellm.utils import convert_to_model_response_object
@@ -199,7 +200,12 @@ class AzureFoundryMAIImageGenerationConfig(BaseImageGenerationConfig):
         api_key: Optional[str] = None,
         json_mode: Optional[bool] = None,
     ) -> ImageResponse:
-        response = raw_response.json()
+        try:
+            response = raw_response.json()
+        except Exception:
+            raise OpenAIError(
+                message=raw_response.text, status_code=raw_response.status_code
+            )
 
         if "usage" in response:
             response["usage"] = self.normalize_mai_image_usage(response.get("usage"))
diff --git a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
index f7ca99f9716..7afdf1f8499 100644
--- a/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
+++ b/tests/test_litellm/llms/azure_ai/image_generation/test_mai_image_generation.py
@@ -246,6 +246,30 @@ class TestAzureMAIImageGeneration:
         assert image_response.usage.input_tokens == 22
         assert image_response.usage.total_tokens == 1046
 
+    def test_transform_image_generation_response_non_json_raises_openai_error(self):
+        from litellm.llms.openai.common_utils import OpenAIError
+
+        config = AzureFoundryMAIImageGenerationConfig()
+        raw_response = MagicMock(spec=httpx.Response)
+        raw_response.json.side_effect = ValueError("not json")
+        raw_response.text = "upstream gateway error"
+        raw_response.status_code = 502
+
+        with pytest.raises(OpenAIError) as exc_info:
+            config.transform_image_generation_response(
+                model="MAI-Image-2.5",
+                raw_response=raw_response,
+                model_response=ImageResponse(),
+                logging_obj=MagicMock(),
+                request_data={"prompt": "A red fox"},
+                optional_params={"width": 1024, "height": 1024},
+                litellm_params={},
+                encoding=None,
+            )
+
+        assert exc_info.value.status_code == 502
+        assert exc_info.value.message == "upstream gateway error"
+
     def test_normalize_mai_usage_preserves_zero_output_tokens(self):
         config = AzureFoundryMAIImageGenerationConfig()
         normalized = config.normalize_mai_image_usage(