fix: preserve image formats in media handling (#9019)

* fix: preserve image formats in media handling

* fix: address image format review feedback

* fix: avoid tainted image temp path rename

tests/test_media_utils.py

@@ -73,6 +73,54 @@ async def test_media_resolver_to_path_detaches_for_component_lifetimes(
         Path(image_path).unlink(missing_ok=True)
 
 
+@pytest.mark.asyncio
+async def test_image_from_base64_uses_detected_image_suffix(tmp_path, monkeypatch):
+    from PIL import Image as PILImage
+
+    monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path))
+    image_buffer = BytesIO()
+    PILImage.new("RGBA", (1, 1), (255, 0, 0, 128)).save(image_buffer, format="PNG")
+    image_base64 = base64.b64encode(image_buffer.getvalue()).decode()
+
+    image_path = await Image.fromBase64(image_base64).convert_to_file_path()
+
+    try:
+        assert Path(image_path).suffix == ".png"
+        with PILImage.open(image_path) as resolved_img:
+            assert resolved_img.format == "PNG"
+    finally:
+        Path(image_path).unlink(missing_ok=True)
+
+
+@pytest.mark.asyncio
+async def test_http_image_without_suffix_uses_detected_image_suffix(
+    tmp_path,
+    monkeypatch,
+):
+    from PIL import Image as PILImage
+
+    monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path))
+    image_buffer = BytesIO()
+    PILImage.new("RGBA", (1, 1), (255, 0, 0, 128)).save(image_buffer, format="PNG")
+
+    async def fake_download_file(_url: str, target_path: str) -> None:
+        Path(target_path).write_bytes(image_buffer.getvalue())
+
+    monkeypatch.setattr(media_utils, "download_file", fake_download_file)
+
+    image_path = await media_utils.MediaResolver(
+        "https://example.com/image?id=123",
+        media_type="image",
+    ).to_path()
+
+    try:
+        assert Path(image_path).suffix == ".png"
+        with PILImage.open(image_path) as resolved_img:
+            assert resolved_img.format == "PNG"
+    finally:
+        Path(image_path).unlink(missing_ok=True)
+
+
 @pytest.mark.asyncio
 async def test_resolve_audio_ref_to_base64_data_decodes_base64_scheme(
     tmp_path, monkeypatch
@@ -146,6 +194,18 @@ async def test_resolve_image_ref_to_base64_data_detects_png(tmp_path):
     assert resolved.to_data_url().startswith("data:image/png;base64,")
 
 
+def test_detect_image_mime_type_accepts_path(tmp_path):
+    from PIL import Image as PILImage
+
+    image_path = tmp_path / "image.png"
+    PILImage.new("RGBA", (1, 1), (255, 0, 0, 255)).save(image_path)
+
+    assert (
+        media_utils.detect_image_mime_type(image_path, default_mime_type=None)
+        == "image/png"
+    )
+
+
 @pytest.mark.asyncio
 async def test_resolve_image_ref_to_base64_data_decodes_data_uri(tmp_path, monkeypatch):
     from PIL import Image as PILImage
@@ -174,7 +234,7 @@ async def test_ensure_jpeg_converts_png_to_temp_jpg(tmp_path, monkeypatch):
     temp_dir = tmp_path / "temp"
     monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir))
     image_path = tmp_path / "image.png"
-    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(image_path)
+    PILImage.new("RGB", (2, 2), (255, 0, 0)).save(image_path)
 
     converted_path = Path(await media_utils.ensure_jpeg(str(image_path)))
 
@@ -185,6 +245,43 @@ async def test_ensure_jpeg_converts_png_to_temp_jpg(tmp_path, monkeypatch):
         assert converted_img.format == "JPEG"
 
 
+@pytest.mark.asyncio
+async def test_ensure_jpeg_keeps_alpha_png(tmp_path, monkeypatch):
+    from PIL import Image as PILImage
+
+    temp_dir = tmp_path / "temp"
+    monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir))
+    image_path = tmp_path / "transparent.png"
+    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(image_path)
+
+    converted_path = await media_utils.ensure_jpeg(str(image_path))
+
+    assert converted_path == str(image_path)
+    assert not temp_dir.exists()
+
+
+@pytest.mark.asyncio
+async def test_ensure_jpeg_keeps_animated_gif(tmp_path, monkeypatch):
+    from PIL import Image as PILImage
+
+    temp_dir = tmp_path / "temp"
+    monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir))
+    image_path = tmp_path / "animated.gif"
+    PILImage.new("RGB", (2, 2), (255, 0, 0)).save(
+        image_path,
+        format="GIF",
+        save_all=True,
+        append_images=[PILImage.new("RGB", (2, 2), (0, 0, 255))],
+        duration=100,
+        loop=0,
+    )
+
+    converted_path = await media_utils.ensure_jpeg(str(image_path))
+
+    assert converted_path == str(image_path)
+    assert not temp_dir.exists()
+
+
 @pytest.mark.asyncio
 async def test_ensure_jpeg_keeps_existing_jpg(tmp_path, monkeypatch):
     from PIL import Image as PILImage
@@ -200,6 +297,54 @@ async def test_ensure_jpeg_keeps_existing_jpg(tmp_path, monkeypatch):
     assert not temp_dir.exists()
 
 
+@pytest.mark.asyncio
+async def test_compress_image_preserves_alpha_png(tmp_path, monkeypatch):
+    from PIL import Image as PILImage
+
+    temp_dir = tmp_path / "temp"
+    monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir))
+    image_path = tmp_path / "transparent.png"
+    PILImage.new("RGBA", (8, 8), (255, 0, 0, 128)).save(image_path)
+
+    compressed_path = Path(
+        await media_utils.compress_image(str(image_path), max_size=2)
+    )
+
+    try:
+        assert compressed_path != image_path
+        assert compressed_path.suffix == ".png"
+        assert compressed_path.parent == temp_dir
+        with PILImage.open(compressed_path) as compressed_img:
+            assert compressed_img.format == "PNG"
+            assert compressed_img.mode == "RGBA"
+            assert max(compressed_img.size) <= 2
+            assert compressed_img.getpixel((0, 0))[3] == 128
+    finally:
+        compressed_path.unlink(missing_ok=True)
+
+
+@pytest.mark.asyncio
+async def test_compress_image_keeps_animated_gif(tmp_path, monkeypatch):
+    from PIL import Image as PILImage
+
+    temp_dir = tmp_path / "temp"
+    monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(temp_dir))
+    image_path = tmp_path / "animated.gif"
+    PILImage.new("RGB", (8, 8), (255, 0, 0)).save(
+        image_path,
+        format="GIF",
+        save_all=True,
+        append_images=[PILImage.new("RGB", (8, 8), (0, 0, 255))],
+        duration=100,
+        loop=0,
+    )
+
+    compressed_path = await media_utils.compress_image(str(image_path), max_size=2)
+
+    assert compressed_path == str(image_path)
+    assert not list(temp_dir.iterdir())
+
+
 @pytest.mark.asyncio
 async def test_resolve_image_ref_to_base64_data_keeps_base64_scheme_fallback(
     tmp_path, monkeypatch

tests/test_platform_image_format_preservation.py

@@ -0,0 +1,121 @@
+import asyncio
+import base64
+from io import BytesIO
+from pathlib import Path
+from unittest.mock import AsyncMock
+
+import pytest
+from PIL import Image as PILImage
+
+from astrbot.api.event import MessageChain
+from astrbot.api.message_components import Image
+from astrbot.core.platform.sources.mattermost.client import MattermostClient
+from astrbot.core.platform.sources.satori.satori_event import SatoriPlatformEvent
+from astrbot.core.platform.sources.slack.slack_event import SlackMessageEvent
+from astrbot.core.platform.sources.webchat import webchat_event
+
+
+@pytest.mark.asyncio
+async def test_satori_image_data_url_preserves_png_mime_type():
+    image_buffer = BytesIO()
+    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(
+        image_buffer,
+        format="PNG",
+    )
+    image_ref = (
+        "data:image/png;base64," + base64.b64encode(image_buffer.getvalue()).decode()
+    )
+
+    result = await SatoriPlatformEvent._convert_component_to_satori_static(
+        Image(file=image_ref),
+    )
+
+    assert result.startswith('<img src="data:image/png;base64,')
+
+
+@pytest.mark.asyncio
+async def test_satori_image_data_url_preserves_jpeg_mime_type():
+    image_buffer = BytesIO()
+    PILImage.new("RGB", (2, 2), (0, 255, 0)).save(
+        image_buffer,
+        format="JPEG",
+    )
+    image_ref = (
+        "data:image/jpeg;base64," + base64.b64encode(image_buffer.getvalue()).decode()
+    )
+
+    result = await SatoriPlatformEvent._convert_component_to_satori_static(
+        Image(file=image_ref),
+    )
+
+    assert result.startswith('<img src="data:image/jpeg;base64,')
+
+
+@pytest.mark.asyncio
+async def test_webchat_image_attachment_uses_detected_extension(tmp_path, monkeypatch):
+    image_buffer = BytesIO()
+    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(
+        image_buffer,
+        format="PNG",
+    )
+    image_ref = (
+        "data:image/png;base64," + base64.b64encode(image_buffer.getvalue()).decode()
+    )
+    queue = asyncio.Queue()
+    monkeypatch.setattr(webchat_event, "attachments_dir", str(tmp_path))
+    monkeypatch.setattr(
+        webchat_event.webchat_queue_mgr,
+        "get_or_create_back_queue",
+        lambda *_args: queue,
+    )
+
+    await webchat_event.WebChatMessageEvent._send(
+        "message-1",
+        MessageChain([Image(file=image_ref)]),
+        "webchat!user!conversation-1",
+    )
+
+    payload = await queue.get()
+    filename = payload["data"].removeprefix("[IMAGE]")
+    assert filename.endswith(".png")
+    assert (tmp_path / filename).exists()
+
+
+@pytest.mark.asyncio
+async def test_slack_image_upload_uses_resolved_filename(tmp_path):
+    image_path = tmp_path / "transparent.png"
+    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(image_path)
+    web_client = AsyncMock()
+    web_client.files_upload_v2.return_value = {
+        "ok": True,
+        "files": [{"url_private": "https://slack.example/files/transparent.png"}],
+    }
+
+    result = await SlackMessageEvent._from_segment_to_slack_block(
+        Image(file=str(image_path)),
+        web_client,
+    )
+
+    assert result is not None
+    assert web_client.files_upload_v2.await_args.kwargs["filename"] == "transparent.png"
+
+
+@pytest.mark.asyncio
+async def test_mattermost_image_upload_uses_detected_mime_type(tmp_path):
+    image_path = tmp_path / "transparent.bin"
+    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(image_path, format="PNG")
+    client = MattermostClient("https://chat.example.com", "test_token")
+    client.upload_file = AsyncMock(return_value="file-id")
+    client.create_post = AsyncMock(return_value={"id": "post-1"})
+
+    await client.send_message_chain(
+        "channel-1",
+        MessageChain([Image(file=str(image_path))]),
+    )
+
+    assert client.upload_file.await_args.args == (
+        "channel-1",
+        Path(image_path).read_bytes(),
+        "transparent.bin",
+        "image/png",
+    )

tests/test_preprocess_stage.py

@@ -1,3 +1,5 @@
+import base64
+from io import BytesIO
 from types import SimpleNamespace
 
 import pytest
@@ -27,7 +29,7 @@ class FakeEvent:
 
 
 @pytest.mark.asyncio
-async def test_preprocess_converts_images_to_jpeg_and_tracks_temp_files(
+async def test_preprocess_preserves_image_formats_and_tracks_temp_files(
     tmp_path, monkeypatch
 ):
     from PIL import Image as PILImage
@@ -39,15 +41,34 @@ async def test_preprocess_converts_images_to_jpeg_and_tracks_temp_files(
         "get_astrbot_temp_path",
         lambda: str(temp_dir),
     )
-    main_image_path = tmp_path / "main.png"
-    reply_image_path = tmp_path / "reply.webp"
-    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(main_image_path)
-    PILImage.new("RGB", (2, 2), (0, 255, 0)).save(reply_image_path)
+    main_image_buffer = BytesIO()
+    PILImage.new("RGBA", (2, 2), (255, 0, 0, 128)).save(
+        main_image_buffer,
+        format="PNG",
+    )
+    main_image_ref = (
+        "data:image/png;base64,"
+        + base64.b64encode(main_image_buffer.getvalue()).decode()
+    )
 
-    reply_image = Image(file=str(reply_image_path))
+    reply_image_buffer = BytesIO()
+    PILImage.new("RGB", (2, 2), (0, 255, 0)).save(
+        reply_image_buffer,
+        format="GIF",
+        save_all=True,
+        append_images=[PILImage.new("RGB", (2, 2), (0, 0, 255))],
+        duration=100,
+        loop=0,
+    )
+    reply_image_ref = (
+        "data:image/gif;base64,"
+        + base64.b64encode(reply_image_buffer.getvalue()).decode()
+    )
+
+    reply_image = Image(file=reply_image_ref)
     event = FakeEvent(
         [
-            Image(file=str(main_image_path)),
+            Image(file=main_image_ref),
             Reply(
                 id="reply-1",
                 chain=[Plain(text="quoted"), reply_image],
@@ -66,16 +87,19 @@ async def test_preprocess_converts_images_to_jpeg_and_tracks_temp_files(
     main_image = event.get_messages()[0]
     assert isinstance(main_image, Image)
     assert main_image.file == main_image.path == main_image.url
-    assert main_image.file.endswith(".jpg")
+    assert main_image.file.endswith(".png")
     assert main_image.file in event.temporary_local_files
-    with PILImage.open(main_image.file) as converted_img:
-        assert converted_img.format == "JPEG"
+    with PILImage.open(main_image.file) as processed_img:
+        assert processed_img.format == "PNG"
+        assert processed_img.getpixel((0, 0))[3] == 128
 
     assert reply_image.file == reply_image.path == reply_image.url
-    assert reply_image.file.endswith(".jpg")
+    assert reply_image.file.endswith(".gif")
     assert reply_image.file in event.temporary_local_files
-    with PILImage.open(reply_image.file) as converted_img:
-        assert converted_img.format == "JPEG"
+    with PILImage.open(reply_image.file) as processed_img:
+        assert processed_img.format == "GIF"
+        assert processed_img.is_animated
+        assert processed_img.n_frames == 2
 
 
 @pytest.mark.asyncio

tests/unit/test_webchat_upload_image_format.py

@@ -0,0 +1,46 @@
+from io import BytesIO
+from types import SimpleNamespace
+
+import pytest
+from PIL import Image as PILImage
+
+from astrbot.dashboard.services.chat_service import ChatService
+
+
+@pytest.mark.asyncio
+async def test_webchat_upload_uses_detected_image_type(tmp_path):
+    image_buffer = BytesIO()
+    PILImage.new("RGB", (2, 2), (255, 0, 0)).save(image_buffer, format="JPEG")
+
+    class FakeUploadFile:
+        filename = "pasted.png"
+        content_type = "image/png"
+
+        async def save(self, destination):
+            with open(destination, "wb") as output:
+                output.write(image_buffer.getvalue())
+
+    class FakeDatabase:
+        def __init__(self):
+            self.inserted = None
+
+        async def insert_attachment(self, path, type, mime_type):
+            self.inserted = {
+                "path": path,
+                "type": type,
+                "mime_type": mime_type,
+            }
+            return SimpleNamespace(attachment_id="attachment-1", path=path)
+
+    fake_db = FakeDatabase()
+    service = ChatService.__new__(ChatService)
+    service.db = fake_db
+    service.attachments_dir = str(tmp_path)
+
+    result = await service.save_uploaded_file(FakeUploadFile())
+
+    assert result["filename"].endswith(".jpg")
+    assert fake_db.inserted["mime_type"] == "image/jpeg"
+    assert fake_db.inserted["type"] == "image"
+    assert (tmp_path / result["filename"]).exists()
+    assert not (tmp_path / "pasted.png").exists()