[BugFix][Diffusion] Preserve image payloads with reasoning metadata (#5238)

Signed-off-by: akshatvishu <akshatnayak197@gmail.com>

tests/diffusion/test_diffusion_output_formatter.py

@@ -241,6 +241,41 @@ def test_formatter_preserves_text_envelope_metadata(monkeypatch: pytest.MonkeyPa
     }
 
 
+@pytest.mark.parametrize(
+    "text_metadata",
+    [
+        pytest.param({"think_text": "reasoning"}, id="sensenova-bagel-thinking"),
+        pytest.param({"ar_generated_text": "reasoning"}, id="hunyuan-image3-cot"),
+    ],
+)
+def test_formatter_preserves_image_with_reasoning_metadata(
+    monkeypatch: pytest.MonkeyPatch,
+    text_metadata: dict[str, str],
+) -> None:
+    monkeypatch.setattr(output_formatter, "supports_audio_output", lambda _: False)
+    metadata = {"text": text_metadata}
+    postprocess_output = normalize_diffusion_postprocess_output(
+        {
+            "payload": {"image": "image-0"},
+            "metadata": metadata,
+        }
+    )
+
+    [result] = format_diffusion_outputs(
+        request=_request("generate an image"),
+        od_config=_config(),
+        diffusion_output=DiffusionOutput(output=None),
+        output_data=None,
+        postprocess_output=postprocess_output,
+        timings=_timings(),
+    )
+
+    assert postprocess_output.primary_key == "image"
+    assert result.images == ["image-0"]
+    assert result.final_output_type == "image"
+    assert result.multimodal_output == {"metadata": metadata}
+
+
 def test_formatter_preserves_audio_output_with_model_sample_rate_fallback(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: