From 019950560d43f1058d0966b95480d73ed8daf034 Mon Sep 17 00:00:00 2001
From: kshitijk4poor <82637225+kshitijk4poor@users.noreply.github.com>
Date: Thu, 2 Jul 2026 16:57:01 +0530
Subject: [PATCH] refactor(image-gen): reuse shared image sniffer + raster
 allowlist in codex backend

Replace the plugin-local _IMAGE_MAGIC_MIME table + _sniff_image_mime
body with a delegation to agent.image_routing._sniff_mime_from_bytes,
the canonical magic-byte sniffer already used across the codebase, then
gate its result to the raster formats gpt-image-2's Responses
input_image actually accepts (png/jpeg/gif/webp).

The shared sniffer also recognizes SVG/TIFF/ICO; without the allowlist
those would pass local validation and be rejected server-side with an
opaque HTTP 400. Gating locally fails them cleanly as invalid_image_input.
Adds a regression test for SVG rejection.

Follow-up on top of @CrazyBoyM's #55828.
---
 plugins/image_gen/openai-codex/__init__.py    | 32 ++++++++++++-------
 .../image_gen/test_openai_codex_provider.py   | 14 ++++++++
 2 files changed, 34 insertions(+), 12 deletions(-)

diff --git a/plugins/image_gen/openai-codex/__init__.py b/plugins/image_gen/openai-codex/__init__.py
index 50fb3b7e8..6790028bb 100644
--- a/plugins/image_gen/openai-codex/__init__.py
+++ b/plugins/image_gen/openai-codex/__init__.py
@@ -87,12 +87,12 @@ _CODEX_INSTRUCTIONS = (
 
 _MAX_REFERENCE_IMAGES = 16
 _MAX_INPUT_IMAGE_BYTES = 25 * 1024 * 1024
-_IMAGE_MAGIC_MIME = (
-    (b"\x89PNG\r\n\x1a\n", "image/png"),
-    (b"\xff\xd8\xff", "image/jpeg"),
-    (b"RIFF", "image/webp"),
-    (b"GIF87a", "image/gif"),
-    (b"GIF89a", "image/gif"),
+# gpt-image-2's Responses ``input_image`` accepts raster formats only. The
+# shared magic-byte sniffer also recognizes SVG/TIFF/ICO, which the API
+# rejects server-side — gate to this allowlist so unsupported inputs fail
+# locally with a clear error instead of an opaque HTTP 400.
+_ACCEPTED_INPUT_MIME = frozenset(
+    {"image/png", "image/jpeg", "image/gif", "image/webp"}
 )
 
 
@@ -159,12 +159,20 @@ def _read_codex_access_token() -> Optional[str]:
 
 
 def _sniff_image_mime(raw: bytes) -> Optional[str]:
-    """Return a safe image MIME type based on magic bytes, not filename labels."""
-    if raw.startswith(b"RIFF") and len(raw) >= 12 and raw[8:12] == b"WEBP":
-        return "image/webp"
-    for magic, mime in _IMAGE_MAGIC_MIME:
-        if raw.startswith(magic):
-            return mime
+    """Return a safe raster image MIME from magic bytes (not filename labels).
+
+    Delegates magic-byte detection to the shared sniffer in
+    ``agent.image_routing`` (single source of truth), then gates the result
+    to :data:`_ACCEPTED_INPUT_MIME` — the raster formats gpt-image-2's
+    ``input_image`` actually accepts. SVG/TIFF/ICO (which the shared sniffer
+    also recognizes) are rejected here so they fail locally with a clear
+    error instead of an opaque server-side HTTP 400.
+    """
+    from agent.image_routing import _sniff_mime_from_bytes
+
+    mime = _sniff_mime_from_bytes(raw)
+    if mime in _ACCEPTED_INPUT_MIME:
+        return mime
     return None
 
 
diff --git a/tests/plugins/image_gen/test_openai_codex_provider.py b/tests/plugins/image_gen/test_openai_codex_provider.py
index fd6d59306..dc8560c8b 100644
--- a/tests/plugins/image_gen/test_openai_codex_provider.py
+++ b/tests/plugins/image_gen/test_openai_codex_provider.py
@@ -234,6 +234,20 @@ class TestGenerate:
         assert result["error_type"] == "invalid_image_input"
         assert "not a supported image" in result["error"]
 
+    def test_rejects_svg_local_source(self, provider, monkeypatch, tmp_path):
+        # The shared magic-byte sniffer recognizes SVG, but gpt-image-2's
+        # input_image accepts raster only — SVG must fail locally with a clear
+        # error, not get embedded and rejected server-side with an opaque 400.
+        monkeypatch.setattr(codex_plugin, "_read_codex_access_token", lambda: "codex-token")
+        svg_path = tmp_path / "vector.svg"
+        svg_path.write_text('<svg xmlns="http://www.w3.org/2000/svg"></svg>')
+
+        result = provider.generate("edit this", image_url=str(svg_path))
+
+        assert result["success"] is False
+        assert result["error_type"] == "invalid_image_input"
+        assert "not a supported image" in result["error"]
+
     def test_partial_image_event_used_when_done_missing(self):
         """If output_item.done is missing, partial_image_b64 is accepted."""
         payload = {