From c6b2c65b048f5a1a33c4f388040e57e4e7d6401a Mon Sep 17 00:00:00 2001 From: Weilong Liao <37870767+Soulter@users.noreply.github.com> Date: Fri, 26 Jun 2026 11:27:32 +0800 Subject: [PATCH] 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 --- .../platform/sources/mattermost/client.py | 11 +- .../platform/sources/satori/satori_event.py | 42 +++- .../platform/sources/slack/slack_event.py | 3 +- .../platform/sources/webchat/webchat_event.py | 23 +- .../platform/sources/wecom/wecom_adapter.py | 22 +- .../sources/weixin_oc/weixin_oc_adapter.py | 15 +- astrbot/core/utils/media_utils.py | 209 ++++++++++++++---- astrbot/dashboard/services/chat_service.py | 20 ++ tests/test_media_utils.py | 147 +++++++++++- ...test_platform_image_format_preservation.py | 121 ++++++++++ tests/test_preprocess_stage.py | 50 +++-- .../unit/test_webchat_upload_image_format.py | 46 ++++ 12 files changed, 617 insertions(+), 92 deletions(-) create mode 100644 tests/test_platform_image_format_preservation.py create mode 100644 tests/unit/test_webchat_upload_image_format.py diff --git a/astrbot/core/platform/sources/mattermost/client.py b/astrbot/core/platform/sources/mattermost/client.py index d47b8d497..c35b89387 100644 --- a/astrbot/core/platform/sources/mattermost/client.py +++ b/astrbot/core/platform/sources/mattermost/client.py @@ -10,7 +10,7 @@ from astrbot.api import logger from astrbot.api.event import MessageChain from astrbot.api.message_components import At, File, Image, Plain, Record, Reply, Video from astrbot.core.utils.astrbot_path import get_astrbot_temp_path -from astrbot.core.utils.media_utils import MediaResolver +from astrbot.core.utils.media_utils import MediaResolver, detect_image_mime_type_async class MattermostClient: @@ -168,12 +168,19 @@ class MattermostClient: path = await segment.convert_to_file_path() file_path = Path(path) file_bytes = await asyncio.to_thread(file_path.read_bytes) + mime_type = ( + await detect_image_mime_type_async( + file_bytes, + default_mime_type=None, + ) + or mimetypes.guess_type(file_path.name)[0] + ) file_ids.append( await self.upload_file( channel_id, file_bytes, file_path.name, - mimetypes.guess_type(file_path.name)[0] or "image/jpeg", + mime_type or "image/jpeg", ) ) elif isinstance(segment, (File, Record, Video)): diff --git a/astrbot/core/platform/sources/satori/satori_event.py b/astrbot/core/platform/sources/satori/satori_event.py index 021422283..7e2e92eb4 100644 --- a/astrbot/core/platform/sources/satori/satori_event.py +++ b/astrbot/core/platform/sources/satori/satori_event.py @@ -15,6 +15,7 @@ from astrbot.api.message_components import ( Video, ) from astrbot.api.platform import AstrBotMessage, PlatformMetadata +from astrbot.core.utils.media_utils import resolve_media_ref_to_base64_data if TYPE_CHECKING: from .satori_adapter import SatoriPlatformAdapter @@ -52,6 +53,27 @@ class SatoriPlatformEvent(AstrMessageEvent): user = login.get("user", {}) self.user_id = user.get("id") if user else None + @staticmethod + async def _image_to_data_url(component: Image) -> str | None: + """Resolve an image component to a MIME-aware data URL. + + Args: + component: Image message component to resolve. + + Returns: + A data URL preserving the detected image MIME type, or None when + the image cannot be resolved. + """ + + image_ref = component.url or component.file + if not image_ref: + return None + image_data = await resolve_media_ref_to_base64_data( + image_ref, + media_type="image", + ) + return image_data.to_data_url() if image_data else None + @classmethod async def send_with_adapter( cls, @@ -184,12 +206,14 @@ class SatoriPlatformEvent(AstrMessageEvent): await self.send(temp_chain) content_parts = [] try: - image_base64 = await component.convert_to_base64() - if image_base64: + image_data_url = await self._image_to_data_url( + component + ) + if image_data_url: img_chain = MessageChain( [ Plain( - text=f'', + text=f'', ), ], ) @@ -228,9 +252,9 @@ class SatoriPlatformEvent(AstrMessageEvent): elif isinstance(component, Image): try: - image_base64 = await component.convert_to_base64() - if image_base64: - return f'' + image_data_url = await self._image_to_data_url(component) + if image_data_url: + return f'' except Exception as e: logger.error(f"图片转换为base64失败: {e}") @@ -321,9 +345,9 @@ class SatoriPlatformEvent(AstrMessageEvent): elif isinstance(component, Image): try: - image_base64 = await component.convert_to_base64() - if image_base64: - return f'' + image_data_url = await cls._image_to_data_url(component) + if image_data_url: + return f'' except Exception as e: logger.error(f"图片转换为base64失败: {e}") diff --git a/astrbot/core/platform/sources/slack/slack_event.py b/astrbot/core/platform/sources/slack/slack_event.py index 3f62690b5..5fb26d22f 100644 --- a/astrbot/core/platform/sources/slack/slack_event.py +++ b/astrbot/core/platform/sources/slack/slack_event.py @@ -1,6 +1,7 @@ import asyncio import re from collections.abc import AsyncGenerator, Iterable +from pathlib import Path from typing import cast from slack_sdk.web.async_client import AsyncWebClient @@ -48,7 +49,7 @@ class SlackMessageEvent(AstrMessageEvent): path = await segment.convert_to_file_path() response = await web_client.files_upload_v2( file=path, - filename="image.jpg", + filename=Path(path).name, ) if not response["ok"]: logger.error(f"Slack file upload failed: {response['error']}") diff --git a/astrbot/core/platform/sources/webchat/webchat_event.py b/astrbot/core/platform/sources/webchat/webchat_event.py index bc1e1a6bc..883074046 100644 --- a/astrbot/core/platform/sources/webchat/webchat_event.py +++ b/astrbot/core/platform/sources/webchat/webchat_event.py @@ -1,13 +1,19 @@ +import asyncio import base64 import json import os import shutil import uuid +from pathlib import Path from astrbot.api import logger from astrbot.api.event import AstrMessageEvent, MessageChain from astrbot.api.message_components import File, Image, Json, Plain, Record from astrbot.core.utils.astrbot_path import get_astrbot_data_path +from astrbot.core.utils.media_utils import ( + MEDIA_MIME_EXTENSIONS, + detect_image_mime_type_async, +) from .webchat_queue_mgr import webchat_queue_mgr @@ -78,11 +84,16 @@ class WebChatMessageEvent(AstrMessageEvent): ) elif isinstance(comp, Image): # save image to local - filename = f"{str(uuid.uuid4())}.jpg" - path = os.path.join(attachments_dir, filename) image_base64 = await comp.convert_to_base64() - with open(path, "wb") as f: - f.write(base64.b64decode(image_base64)) + image_bytes = base64.b64decode(image_base64) + mime_type = await detect_image_mime_type_async( + image_bytes, + default_mime_type=None, + ) + suffix = MEDIA_MIME_EXTENSIONS.get(mime_type or "", ".jpg") + filename = f"{str(uuid.uuid4())}{suffix}" + path = os.path.join(attachments_dir, filename) + await asyncio.to_thread(Path(path).write_bytes, image_bytes) data = f"[IMAGE]{filename}" await web_chat_back_queue.put( { @@ -97,8 +108,8 @@ class WebChatMessageEvent(AstrMessageEvent): filename = f"{str(uuid.uuid4())}.wav" path = os.path.join(attachments_dir, filename) record_base64 = await comp.convert_to_base64() - with open(path, "wb") as f: - f.write(base64.b64decode(record_base64)) + record_bytes = base64.b64decode(record_base64) + await asyncio.to_thread(Path(path).write_bytes, record_bytes) data = f"[RECORD]{filename}" await web_chat_back_queue.put( { diff --git a/astrbot/core/platform/sources/wecom/wecom_adapter.py b/astrbot/core/platform/sources/wecom/wecom_adapter.py index e5d0b29f9..e80b762c2 100644 --- a/astrbot/core/platform/sources/wecom/wecom_adapter.py +++ b/astrbot/core/platform/sources/wecom/wecom_adapter.py @@ -29,7 +29,11 @@ from astrbot.core import logger from astrbot.core.platform.astr_message_event import MessageSesion from astrbot.core.platform.webhook_server import FastAPIWebhookServer from astrbot.core.utils.astrbot_path import get_astrbot_temp_path -from astrbot.core.utils.media_utils import MediaResolver +from astrbot.core.utils.media_utils import ( + MEDIA_MIME_EXTENSIONS, + MediaResolver, + detect_image_mime_type_async, +) from astrbot.core.utils.webhook_utils import log_webhook_info from .wecom_event import WecomPlatformEvent @@ -388,8 +392,7 @@ class WecomPlatformAdapter(Platform): ) temp_dir = get_astrbot_temp_path() path = os.path.join(temp_dir, f"wecom_{msg.media_id}.amr") - with open(path, "wb") as f: - f.write(resp.content) + await asyncio.to_thread(Path(path).write_bytes, resp.content) try: path_wav = await MediaResolver( @@ -453,9 +456,13 @@ class WecomPlatformAdapter(Platform): media_id, ) temp_dir = get_astrbot_temp_path() - path = os.path.join(temp_dir, f"weixinkefu_{media_id}.jpg") - with open(path, "wb") as f: - f.write(resp.content) + mime_type = await detect_image_mime_type_async( + resp.content, + default_mime_type=None, + ) + suffix = MEDIA_MIME_EXTENSIONS.get(mime_type or "", ".jpg") + path = os.path.join(temp_dir, f"weixinkefu_{media_id}{suffix}") + await asyncio.to_thread(Path(path).write_bytes, resp.content) abm.message = [Image(file=path, url=path)] elif msgtype == "voice": media_id = msg.get("voice", {}).get("media_id", "") @@ -467,8 +474,7 @@ class WecomPlatformAdapter(Platform): temp_dir = get_astrbot_temp_path() path = os.path.join(temp_dir, f"weixinkefu_{media_id}.amr") - with open(path, "wb") as f: - f.write(resp.content) + await asyncio.to_thread(Path(path).write_bytes, resp.content) try: path_wav = await MediaResolver( diff --git a/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py b/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py index 81570346a..5bbf24a72 100644 --- a/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py +++ b/astrbot/core/platform/sources/weixin_oc/weixin_oc_adapter.py @@ -29,7 +29,11 @@ from astrbot.api.platform import ( from astrbot.core import astrbot_config from astrbot.core.platform.astr_message_event import MessageSesion from astrbot.core.utils.astrbot_path import get_astrbot_temp_path -from astrbot.core.utils.media_utils import MediaResolver +from astrbot.core.utils.media_utils import ( + MEDIA_MIME_EXTENSIONS, + MediaResolver, + detect_image_mime_type_async, +) from .weixin_oc_client import WeixinOCClient from .weixin_oc_event import WeixinOCMessageEvent @@ -757,11 +761,16 @@ class WeixinOCAdapter(Platform): ) else: content = await self.client.download_cdn_bytes(encrypted_query_param) + mime_type = await detect_image_mime_type_async( + content, + default_mime_type=None, + ) + suffix = MEDIA_MIME_EXTENSIONS.get(mime_type or "", ".jpg") image_path = self._save_inbound_media( content, prefix="weixin_oc_img", - file_name="image.jpg", - fallback_suffix=".jpg", + file_name=f"image{suffix}", + fallback_suffix=suffix, ) return Image.fromFileSystem(str(image_path)) diff --git a/astrbot/core/utils/media_utils.py b/astrbot/core/utils/media_utils.py index c53d148ed..2946f0d94 100644 --- a/astrbot/core/utils/media_utils.py +++ b/astrbot/core/utils/media_utils.py @@ -358,14 +358,14 @@ def describe_media_ref(media_ref: object | None) -> str: def detect_image_mime_type( - image_bytes: bytes, + image_source: bytes | str | Path, *, default_mime_type: str | None = "image/jpeg", ) -> str | None: - """Detect an image MIME type from bytes. + """Detect an image MIME type from encoded bytes or a local path. Args: - image_bytes: Encoded image bytes to inspect. + image_source: Encoded image bytes or a local image path to inspect. default_mime_type: MIME type to return when detection fails. Returns: @@ -374,7 +374,12 @@ def detect_image_mime_type( """ try: - with PILImage.open(io.BytesIO(image_bytes)) as image: + image_file = ( + io.BytesIO(image_source) + if isinstance(image_source, bytes) + else image_source + ) + with PILImage.open(image_file) as image: image.verify() image_format = str(image.format or "").upper() except Exception: @@ -383,6 +388,29 @@ def detect_image_mime_type( return IMAGE_FORMAT_MIME_TYPES.get(image_format, default_mime_type) +async def detect_image_mime_type_async( + image_source: bytes | str | Path, + *, + default_mime_type: str | None = "image/jpeg", +) -> str | None: + """Detect an image MIME type without blocking the event loop. + + Args: + image_source: Encoded image bytes or a local image path to inspect. + default_mime_type: MIME type to return when detection fails. + + Returns: + The detected MIME type, or ``default_mime_type`` when detection fails or + the format is unknown. + """ + + return await asyncio.to_thread( + detect_image_mime_type, + image_source, + default_mime_type=default_mime_type, + ) + + def _guess_mime_type(path: Path, fallback: str | None = None) -> str | None: """Guess a MIME type from a filename, with an optional fallback.""" return mimetypes.guess_type(path.name)[0] or fallback @@ -418,18 +446,35 @@ async def _materialize_media_ref( suffix = default_suffix or DEFAULT_MEDIA_SUFFIXES.get(media_type, ".bin") if media_ref.startswith(("http://", "https://")): - parsed = urlparse(media_ref) - target_suffix = Path(parsed.path).suffix or suffix - target_path = _temp_media_path(media_type, target_suffix) + if media_type == "image": + target_path = _temp_media_path("image", ".bin") + else: + parsed = urlparse(media_ref) + target_suffix = Path(parsed.path).suffix or suffix + target_path = _temp_media_path(media_type, target_suffix) cleanup_paths.append(target_path) try: await download_file(media_ref, str(target_path)) except Exception: _cleanup_paths(cleanup_paths) raise + mime_type = _guess_mime_type(target_path) + if media_type == "image": + detected_mime_type = await detect_image_mime_type_async( + target_path, + default_mime_type=None, + ) + if detected_mime_type: + mime_type = detected_mime_type + detected_suffix = _extension_from_mime_type(detected_mime_type) + if detected_suffix and target_path.suffix.lower() != detected_suffix: + detected_path = _temp_media_path("image", detected_suffix) + await asyncio.to_thread(target_path.rename, detected_path) + cleanup_paths[-1] = detected_path + target_path = detected_path return _LocalMediaFile( path=target_path, - mime_type=_guess_mime_type(target_path), + mime_type=mime_type, cleanup_paths=cleanup_paths, ) @@ -440,10 +485,18 @@ async def _materialize_media_ref( if media_ref.startswith("data:"): mime_type, media_bytes = _parse_base64_data_uri(media_ref) target_suffix = _extension_from_mime_type(mime_type) or suffix + if media_type == "image" and target_suffix == suffix: + detected_mime_type = await detect_image_mime_type_async( + media_bytes, + default_mime_type=None, + ) + if detected_mime_type: + mime_type = detected_mime_type + target_suffix = _extension_from_mime_type(detected_mime_type) or suffix target_path = _temp_media_path(media_type, target_suffix) cleanup_paths.append(target_path) try: - target_path.write_bytes(media_bytes) + await asyncio.to_thread(target_path.write_bytes, media_bytes) except Exception: _cleanup_paths(cleanup_paths) raise @@ -458,14 +511,26 @@ async def _materialize_media_ref( media_ref.removeprefix("base64://"), error_message="invalid base64 media payload", ) - target_path = _temp_media_path(media_type, suffix) + mime_type = None + target_suffix = suffix + if media_type == "image": + mime_type = await detect_image_mime_type_async( + media_bytes, + default_mime_type=None, + ) + target_suffix = _extension_from_mime_type(mime_type) or suffix + target_path = _temp_media_path(media_type, target_suffix) cleanup_paths.append(target_path) try: - target_path.write_bytes(media_bytes) + await asyncio.to_thread(target_path.write_bytes, media_bytes) except Exception: _cleanup_paths(cleanup_paths) raise - return _LocalMediaFile(path=target_path, cleanup_paths=cleanup_paths) + return _LocalMediaFile( + path=target_path, + mime_type=mime_type, + cleanup_paths=cleanup_paths, + ) path = Path(media_ref) path_exists = False @@ -487,14 +552,26 @@ async def _materialize_media_ref( except ValueError: pass else: - target_path = _temp_media_path(media_type, suffix) + mime_type = None + target_suffix = suffix + if media_type == "image": + mime_type = await detect_image_mime_type_async( + media_bytes, + default_mime_type=None, + ) + target_suffix = _extension_from_mime_type(mime_type) or suffix + target_path = _temp_media_path(media_type, target_suffix) cleanup_paths.append(target_path) try: - target_path.write_bytes(media_bytes) + await asyncio.to_thread(target_path.write_bytes, media_bytes) except Exception: _cleanup_paths(cleanup_paths) raise - return _LocalMediaFile(path=target_path, cleanup_paths=cleanup_paths) + return _LocalMediaFile( + path=target_path, + mime_type=mime_type, + cleanup_paths=cleanup_paths, + ) return _LocalMediaFile(path=path, mime_type=_guess_mime_type(path)) @@ -707,13 +784,13 @@ class MediaResolver: if self.media_type == "image": async with self.as_path(target_format=target_format) as resolved: try: - media_bytes = resolved.read_bytes() + media_bytes = await asyncio.to_thread(resolved.read_bytes) except OSError: if strict: raise return None - mime_type = detect_image_mime_type( + mime_type = await detect_image_mime_type_async( media_bytes, default_mime_type=None, ) @@ -1191,7 +1268,7 @@ async def ensure_wav(audio_path: str, output_path: str | None = None) -> str: async def ensure_jpeg(image_path: str, output_path: str | None = None) -> str: - """Ensure the image path points to a JPEG file. + """Ensure JPEG-compatible still images point to a JPEG file. Args: image_path: Local image path to inspect and convert when needed. @@ -1200,7 +1277,8 @@ async def ensure_jpeg(image_path: str, output_path: str | None = None) -> str: Returns: The original path when the source is already a JPEG file with a jpg/jpeg - suffix or cannot be found; otherwise the converted JPEG path. + suffix, cannot be found, has alpha transparency, or is animated; otherwise + the converted JPEG path. Raises: Exception: Raised by Pillow when the source file cannot be opened or saved as @@ -1216,10 +1294,25 @@ async def ensure_jpeg(image_path: str, output_path: str | None = None) -> str: with PILImage.open(source_path) as opened_img: image_format = str(opened_img.format or "").upper() + image_has_alpha = opened_img.mode in {"RGBA", "LA"} or ( + opened_img.mode == "P" and "transparency" in opened_img.info + ) + image_is_animated = ( + getattr(opened_img, "is_animated", False) + or getattr( + opened_img, + "n_frames", + 1, + ) + > 1 + ) if image_format == "JPEG" and source_path.suffix.lower() in {".jpg", ".jpeg"}: return image_path + if image_has_alpha or image_is_animated: + return image_path + if output_path is None: temp_dir = Path(get_astrbot_temp_path()) temp_dir.mkdir(parents=True, exist_ok=True) @@ -1228,23 +1321,11 @@ async def ensure_jpeg(image_path: str, output_path: str | None = None) -> str: def convert_image_to_jpeg() -> str: converted_img: PILImage.Image | None = None - flattened_img: PILImage.Image | None = None with PILImage.open(image_path) as opened_img: try: working_img: PILImage.Image = opened_img - if opened_img.mode in {"RGBA", "LA"} or ( - opened_img.mode == "P" and "transparency" in opened_img.info - ): - flattened = PILImage.new("RGB", opened_img.size, (255, 255, 255)) - flattened_img = flattened - alpha_source = opened_img.convert("RGBA") - try: - flattened.paste(alpha_source, mask=alpha_source.getchannel("A")) - finally: - alpha_source.close() - working_img = flattened - elif opened_img.mode != "RGB": + if opened_img.mode != "RGB": converted_img = opened_img.convert("RGB") working_img = converted_img @@ -1253,8 +1334,6 @@ async def ensure_jpeg(image_path: str, output_path: str | None = None) -> str: finally: if converted_img is not None: converted_img.close() - if flattened_img is not None: - flattened_img.close() try: return await asyncio.to_thread(convert_image_to_jpeg) @@ -1381,19 +1460,46 @@ async def extract_video_cover( def _compress_image_sync( - data: bytes, + source: bytes | Path, temp_dir: Path, max_size: int, quality: int, optimize: bool, -) -> str: - """Run image compression synchronously via ``asyncio.to_thread``.""" - with PILImage.open(io.BytesIO(data)) as opened_img: +) -> str | None: + """Run image compression synchronously via ``asyncio.to_thread``. + + Args: + source: Encoded image bytes or a local path to open inside the worker. + temp_dir: Directory where the compressed image should be written. + max_size: Longest edge of the compressed image in pixels. + quality: JPEG output quality in the range 1-100. + optimize: Whether Pillow should optimize the saved image. + + Returns: + The compressed image path, or ``None`` when the image should be kept as-is. + """ + fp = io.BytesIO(source) if isinstance(source, bytes) else source + with PILImage.open(fp) as opened_img: converted_img: PILImage.Image | None = None try: + if ( + getattr(opened_img, "is_animated", False) + or getattr(opened_img, "n_frames", 1) > 1 + ): + return None + working_img = opened_img - if opened_img.mode != "RGB": + image_has_alpha = opened_img.mode in {"RGBA", "LA"} or ( + opened_img.mode == "P" and "transparency" in opened_img.info + ) + output_format = "PNG" if image_has_alpha else "JPEG" + output_suffix = ".png" if image_has_alpha else ".jpg" + + if image_has_alpha and opened_img.mode != "RGBA": + converted_img = opened_img.convert("RGBA") + working_img = converted_img + elif not image_has_alpha and opened_img.mode != "RGB": converted_img = opened_img.convert("RGB") working_img = converted_img assert working_img is not None @@ -1402,8 +1508,11 @@ def _compress_image_sync( working_img.thumbnail((max_size, max_size), PILImage.Resampling.LANCZOS) new_uuid = uuid.uuid4().hex - save_path = temp_dir / f"compressed_{new_uuid}.jpg" - working_img.save(save_path, "JPEG", quality=quality, optimize=optimize) + save_path = temp_dir / f"compressed_{new_uuid}{output_suffix}" + save_kwargs: dict[str, int | bool] = {"optimize": optimize} + if output_format == "JPEG": + save_kwargs["quality"] = quality + working_img.save(save_path, output_format, **save_kwargs) logger.debug(f"Image compressed successfully: {save_path}") return str(save_path) finally: @@ -1431,7 +1540,7 @@ async def compress_image( quality = min(max(int(quality), 1), 100) optimize = IMAGE_COMPRESS_DEFAULT_OPTIMIZE min_file_size_bytes = int(IMAGE_COMPRESS_DEFAULT_MIN_FILE_SIZE_MB * 1024 * 1024) - data = None + image_source: bytes | Path | None = None def _exceeds_max_size(source: bytes | Path) -> bool: try: @@ -1446,11 +1555,13 @@ async def compress_image( return url_or_path elif url_or_path.startswith("data:image"): _header, encoded = url_or_path.split(",", 1) - data = _decode_base64_payload( + image_source = _decode_base64_payload( encoded, error_message="invalid image data URI payload", ) - if len(data) < min_file_size_bytes and not _exceeds_max_size(data): + if len(image_source) < min_file_size_bytes and not _exceeds_max_size( + image_source + ): return url_or_path else: local_path = Path(url_or_path) @@ -1460,21 +1571,21 @@ async def compress_image( local_path ): return url_or_path - with local_path.open("rb") as f: - data = f.read() + image_source = local_path - if not data: + if image_source is None: return url_or_path temp_dir = Path(get_astrbot_temp_path()) temp_dir.mkdir(parents=True, exist_ok=True) # Offload the blocking image processing task to a thread. - return await asyncio.to_thread( + compressed_path = await asyncio.to_thread( _compress_image_sync, - data, + image_source, temp_dir, max_size, quality, optimize, ) + return compressed_path or url_or_path diff --git a/astrbot/dashboard/services/chat_service.py b/astrbot/dashboard/services/chat_service.py index 80bf6bbf3..866328ab2 100644 --- a/astrbot/dashboard/services/chat_service.py +++ b/astrbot/dashboard/services/chat_service.py @@ -26,6 +26,10 @@ from astrbot.core.platform.sources.webchat.webchat_queue_mgr import webchat_queu from astrbot.core.utils.active_event_registry import active_event_registry from astrbot.core.utils.astrbot_path import get_astrbot_data_path from astrbot.core.utils.datetime_utils import to_utc_isoformat +from astrbot.core.utils.media_utils import ( + MEDIA_MIME_EXTENSIONS, + detect_image_mime_type_async, +) SSE_HEARTBEAT = ": heartbeat\n\n" @@ -584,6 +588,22 @@ class ChatService: raise ChatServiceError("Invalid filename") await file.save(str(file_path)) + if attach_type == "image": + detected_mime_type = await detect_image_mime_type_async( + file_path, + default_mime_type=None, + ) + if detected_mime_type: + content_type = detected_mime_type + detected_suffix = MEDIA_MIME_EXTENSIONS.get(detected_mime_type) + if detected_suffix and file_path.suffix.lower() != detected_suffix: + target_path = file_path.with_suffix(detected_suffix) + if target_path.exists(): + target_path = ( + attachments_dir / f"{uuid.uuid4().hex}{detected_suffix}" + ) + await asyncio.to_thread(file_path.rename, target_path) + file_path = target_path attachment = await self.db.insert_attachment( path=str(file_path), type=attach_type, diff --git a/tests/test_media_utils.py b/tests/test_media_utils.py index c64d286c5..6edfcd4d7 100644 --- a/tests/test_media_utils.py +++ b/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 diff --git a/tests/test_platform_image_format_preservation.py b/tests/test_platform_image_format_preservation.py new file mode 100644 index 000000000..42aaa8aa0 --- /dev/null +++ b/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('