diff --git a/astrbot/dashboard/services/chat_service.py b/astrbot/dashboard/services/chat_service.py index 36b0700b92..a1bd86883d 100644 --- a/astrbot/dashboard/services/chat_service.py +++ b/astrbot/dashboard/services/chat_service.py @@ -36,6 +36,13 @@ SSE_HEARTBEAT = ": heartbeat\n\n" CHAT_RUN_SUBSCRIBER_QUEUE_SIZE = 256 +WEBCHAT_IMAGE_MIME_TYPES = { + ".jpg": "image/jpeg", + ".jpeg": "image/jpeg", + ".png": "image/png", + ".gif": "image/gif", + ".webp": "image/webp", +} def sanitize_upload_filename(filename: str | None) -> str: @@ -504,7 +511,6 @@ def __init__( self.webchat_img_dir = os.path.join(get_astrbot_data_path(), "webchat", "imgs") os.makedirs(self.attachments_dir, exist_ok=True) - self.supported_imgs = ["jpg", "jpeg", "png", "gif", "webp"] self.conv_mgr = core_lifecycle.conversation_manager self.platform_history_mgr = core_lifecycle.platform_message_history_manager self.umop_config_router = core_lifecycle.umop_config_router @@ -557,8 +563,8 @@ async def resolve_webchat_file( filename_ext = file_path.suffix.lower() if filename_ext == ".wav": return str(file_path), "audio/wav" - if filename_ext[1:] in self.supported_imgs: - return str(file_path), "image/jpeg" + if filename_ext in WEBCHAT_IMAGE_MIME_TYPES: + return str(file_path), WEBCHAT_IMAGE_MIME_TYPES[filename_ext] return str(file_path), None async def resolve_webchat_file_from_dashboard_query( diff --git a/tests/unit/test_webchat_upload_image_format.py b/tests/unit/test_webchat_upload_image_format.py index 14cce65501..7125c6961f 100644 --- a/tests/unit/test_webchat_upload_image_format.py +++ b/tests/unit/test_webchat_upload_image_format.py @@ -4,7 +4,10 @@ import pytest from PIL import Image as PILImage -from astrbot.dashboard.services.chat_service import ChatService +from astrbot.dashboard.services.chat_service import ( + WEBCHAT_IMAGE_MIME_TYPES, + ChatService, +) @pytest.mark.asyncio @@ -44,3 +47,35 @@ async def insert_attachment(self, path, type, mime_type): assert fake_db.inserted["type"] == "image" assert (tmp_path / result["filename"]).exists() assert not (tmp_path / "pasted.png").exists() + + +@pytest.mark.parametrize( + ("filename", "expected_mime_type"), + [(f"photo{ext}", mime_type) for ext, mime_type in WEBCHAT_IMAGE_MIME_TYPES.items()], +) +@pytest.mark.asyncio +async def test_resolve_webchat_file_uses_image_extension_mime_type( + tmp_path, + monkeypatch, + filename, + expected_mime_type, +): + monkeypatch.setattr( + "astrbot.dashboard.services.chat_service.get_astrbot_data_path", + lambda: str(tmp_path), + ) + service = ChatService( + SimpleNamespace(), + SimpleNamespace( + conversation_manager=None, + platform_message_history_manager=None, + umop_config_router=None, + ), + ) + file_path = tmp_path / "attachments" / filename + file_path.write_bytes(b"image-bytes") + + resolved_path, mime_type = await service.resolve_webchat_file(filename) + + assert resolved_path == str(file_path.resolve(strict=False)) + assert mime_type == expected_mime_type