🪲 BUG-#12: Preserve real image MIME type instead of hardcoding image/png

pycodeloop/providers/_shapes.py

@@ -5,12 +5,26 @@
 from __future__ import annotations
 
 import json
+import re
 from collections.abc import Callable
 
 from pycodeloop.core.session import Message
 
 RequestBuilder = Callable[[str, "list[Message]", "list[dict]", str], dict]
 
+_DATA_URL = re.compile(r"^data:(?P<mime>[\w./+-]+);base64,(?P<data>[^\n]+)$")
+
+
+def _split_image(image: str) -> tuple[str, str]:
+    """Returns `(mime_type, base64_data)`. Accepts a full `data:<mime>;
+    base64,<data>` URL (mime taken from it) or a bare base64 string
+    (assumed `image/png`, the historical/backward-compatible behavior
+    for a caller passing raw base64 straight to `Session.add_user`)."""
+    match = _DATA_URL.match(image)
+    if match:
+        return match.group("mime"), match.group("data")
+    return "image/png", image
+
 
 def openai_tool_schema(tools: list[dict]) -> list[dict]:
     return [
@@ -38,9 +52,9 @@ def to_openai_messages(
                 content = [
                     {
                         "type": "image_url",
-                        "image_url": {"url": f"data:image/png;base64,{image}"},
+                        "image_url": {"url": f"data:{mime};base64,{data}"},
                     }
-                    for image in msg.images
+                    for mime, data in (_split_image(i) for i in msg.images)
                 ]
                 if msg.content:
                     content.append({"type": "text", "text": msg.content})
@@ -108,11 +122,11 @@ def to_anthropic_messages(messages: list[Message]) -> list[dict]:
                         "type": "image",
                         "source": {
                             "type": "base64",
-                            "media_type": "image/png",
-                            "data": image,
+                            "media_type": mime,
+                            "data": data,
                         },
                     }
-                    for image in msg.images
+                    for mime, data in (_split_image(i) for i in msg.images)
                 ]
                 if msg.content:
                     content.append({"type": "text", "text": msg.content})

tests/providers/test_shapes.py

@@ -4,13 +4,27 @@
 
 from pycodeloop.core.session import Message
 from pycodeloop.providers._shapes import (
+    _split_image,
     anthropic_tool_schema,
     request_builder_from_config,
     to_anthropic_messages,
     to_openai_messages,
 )
 
 
+class TestSplitImage(unittest.TestCase):
+    def test_data_url_rejects_embedded_newlines_in_the_base64_payload(self):
+        """A multiline/PEM-style base64 payload must not silently match
+        as a clean data URL — falls back to treating the whole string as
+        opaque `image/png` data instead of a partial, corrupted match."""
+        malformed = "data:image/png;base64,abc\ndef"
+
+        mime, data = _split_image(malformed)
+
+        self.assertEqual(mime, "image/png")
+        self.assertEqual(data, malformed)
+
+
 class TestAnthropicMessageBuilding(unittest.TestCase):
     def test_plain_text_user_message_stays_a_string(self):
         out = to_anthropic_messages([Message(role="user", content="hi")])
@@ -50,6 +64,22 @@ def test_image_only_message_omits_empty_text_block(self):
         self.assertEqual(len(out[0]["content"]), 1)
         self.assertEqual(out[0]["content"][0]["type"], "image")
 
+    def test_data_url_image_preserves_its_real_mime_type(self):
+        out = to_anthropic_messages(
+            [
+                Message(
+                    role="user",
+                    content="",
+                    images=["data:image/jpeg;base64,b64data"],
+                )
+            ]
+        )
+
+        self.assertEqual(
+            out[0]["content"][0]["source"],
+            {"type": "base64", "media_type": "image/jpeg", "data": "b64data"},
+        )
+
 
 class TestOpenAIMessageBuilding(unittest.TestCase):
     def test_plain_text_user_message_stays_a_string(self):
@@ -81,6 +111,23 @@ def test_user_message_with_images_becomes_content_blocks(self):
             },
         )
 
+    def test_data_url_image_preserves_its_real_mime_type(self):
+        out = to_openai_messages(
+            "sys",
+            [
+                Message(
+                    role="user",
+                    content="",
+                    images=["data:image/webp;base64,b64data"],
+                )
+            ],
+        )
+
+        self.assertEqual(
+            out[-1]["content"][0]["image_url"]["url"],
+            "data:image/webp;base64,b64data",
+        )
+
     def test_tool_call_extra_fields_round_trip_back_to_the_wire(self):
         """A ToolCall's vendor-specific `extra` (e.g. Gemini's
         extra_content.google.thought_signature) must be re-emitted as a