diff --git a/astrbot/core/pipeline/preprocess_stage/stage.py b/astrbot/core/pipeline/preprocess_stage/stage.py index b7b784fb11..27347751ea 100644 --- a/astrbot/core/pipeline/preprocess_stage/stage.py +++ b/astrbot/core/pipeline/preprocess_stage/stage.py @@ -1,246 +1,284 @@ -import asyncio -import random -import re -import traceback -from collections.abc import AsyncGenerator -from pathlib import Path - -from astrbot.core import logger -from astrbot.core.message.components import Image, Plain, Record, Reply -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.utils.astrbot_path import get_astrbot_temp_path -from astrbot.core.utils.media_utils import ( - describe_media_ref, - ensure_wav, - file_uri_to_path, - is_file_uri, -) - -from ..context import PipelineContext -from ..stage import Stage, register_stage - - -@register_stage -class PreProcessStage(Stage): - async def initialize(self, ctx: PipelineContext) -> None: - self.ctx = ctx - self.config = ctx.astrbot_config - self.plugin_manager = ctx.plugin_manager - - self.stt_settings: dict = self.config.get("provider_stt_settings", {}) - self.platform_settings: dict = self.config.get("platform_settings", {}) - - @staticmethod - def _track_temp_media(event: AstrMessageEvent, media_path: str) -> None: - """Track a media file owned by the current event. - - Args: - event: Message event whose lifecycle owns the temporary file. - media_path: Local media path to track when it lives under AstrBot temp. - """ - - try: - path = Path(media_path).resolve() - temp_dir = Path(get_astrbot_temp_path()).resolve() - path.relative_to(temp_dir) - except (OSError, ValueError): - return - event.track_temporary_local_file(str(path)) - - async def process( - self, - event: AstrMessageEvent, - ) -> None | AsyncGenerator[None, None]: - """在处理事件之前的预处理""" - # 平台特异配置:platform_specific..pre_ack_emoji - supported = {"telegram", "lark", "discord"} - platform = event.get_platform_name() - cfg = ( - self.config.get("platform_specific", {}) - .get(platform, {}) - .get("pre_ack_emoji", {}) - ) or {} - emojis = cfg.get("emojis") or [] - if ( - cfg.get("enable", False) - and platform in supported - and emojis - and event.is_at_or_wake_command - ): - try: - await event.react(random.choice(emojis)) - except Exception as e: - logger.warning( - f"Failed to send a pre-response reaction on {platform}: {e}" - ) - - # 路径映射 - if mappings := self.platform_settings.get("path_mapping", []): - # 支持 Record,Image 消息段的路径映射。 - message_chain = event.get_messages() - - for idx, component in enumerate(message_chain): - if isinstance(component, Record | Image) and component.url: - for mapping in mappings: - # ":" is ambiguous for Windows absolute paths. Parse - # the separator after a drive-lettered source first, - # then handle a drive-lettered target. - drive_source_mapping = re.fullmatch( - r"([A-Za-z]:[^:]+):(.+)", mapping - ) - drive_target_mapping = re.fullmatch( - r"(.+):([A-Za-z]:.+)", mapping - ) - if drive_source_mapping: - from_, to_ = drive_source_mapping.groups() - elif drive_target_mapping: - from_, to_ = drive_target_mapping.groups() - else: - from_, separator, to_ = mapping.partition(":") - if not separator: - logger.warning(f"Invalid path mapping: {mapping}") - continue - from_ = from_.removesuffix("/").removesuffix("\\") - to_ = to_.removesuffix("/").removesuffix("\\") - - url = ( - file_uri_to_path(component.url) - if is_file_uri(component.url) - else component.url - ) - if url.startswith(from_): - component.url = url.replace(from_, to_, 1) - logger.debug(f"Path mapping: {url} -> {component.url}") - message_chain[idx] = component - - # Localize source images and normalize audio for downstream processing. - message_chain = event.get_messages() - for idx, component in enumerate(message_chain): - if isinstance(component, Record): - try: - original_path = await component.convert_to_file_path() - self._track_temp_media(event, original_path) - record_path = await ensure_wav(original_path) - self._track_temp_media(event, record_path) - component.file = record_path - component.path = record_path - message_chain[idx] = component - except Exception as e: - logger.warning(f"Voice processing failed: {e}") - elif isinstance(component, Image): - try: - image_path = await component.convert_to_file_path() - component.file = image_path - component.path = image_path - # Image.convert_to_file_path() prefers url, so keep it aligned. - component.url = image_path - message_chain[idx] = component - # Attachment references outlive the event; model copies do not. - event.untrack_temporary_local_file(image_path) - except Exception as e: - media_ref = component.url or component.file - logger.warning( - "Image processing failed for %s: %s", - describe_media_ref(media_ref), - e, - ) - - # Also normalize media components inside Reply chains. - for component in event.get_messages(): - if isinstance(component, Reply) and component.chain: - for idx, reply_comp in enumerate(component.chain): - if isinstance(reply_comp, Record): - try: - original_path = await reply_comp.convert_to_file_path() - self._track_temp_media(event, original_path) - record_path = await ensure_wav(original_path) - self._track_temp_media(event, record_path) - reply_comp.file = record_path - reply_comp.path = record_path - component.chain[idx] = reply_comp - except Exception as e: - logger.warning( - f"Voice processing in reply chain failed: {e}" - ) - elif isinstance(reply_comp, Image): - try: - image_path = await reply_comp.convert_to_file_path() - reply_comp.file = image_path - reply_comp.path = image_path - # Image.convert_to_file_path() prefers url, so keep it aligned. - reply_comp.url = image_path - component.chain[idx] = reply_comp - event.untrack_temporary_local_file(image_path) - except Exception as e: - media_ref = reply_comp.url or reply_comp.file - logger.warning( - "Image processing in reply chain failed for %s: %s", - describe_media_ref(media_ref), - e, - ) - - # STT - if self.stt_settings.get("enable", False): - # TODO: 独立 - ctx = self.plugin_manager.context - stt_provider = await ctx.get_using_stt_provider_async( - event.unified_msg_origin - ) - if not stt_provider: - logger.warning( - f"Session {event.unified_msg_origin} has no speech-to-text " - "provider configured.", - ) - return - - async def _stt_record(record_comp: Record, is_reply: bool = False): - """对单个 Record 组件执行语音转文本,成功返回 Plain,失败返回 None。""" - prefix = "referenced " if is_reply else "" - try: - path = await record_comp.convert_to_file_path() - except Exception as e: - logger.warning(f"Failed to resolve the {prefix}voice path: {e}") - return None - - retry = 5 - for i in range(retry): - try: - result = await stt_provider.get_text(audio_url=path) - if result: - suffix = " (referenced message)" if is_reply else "" - logger.info(f"Speech-to-text{suffix} result: " + result) - return Plain(result) - break - except FileNotFoundError: - # napcat workaround: file may not be ready immediately - logger.debug( - f"File is not ready ({path}); retrying {i + 1}/{retry}." - ) - await asyncio.sleep(0.5) - continue - except BaseException as e: - logger.error(traceback.format_exc()) - suffix = " (referenced message)" if is_reply else "" - logger.error(f"Speech-to-text{suffix} failed: {e}") - break - return None - - message_chain = event.get_messages() - for idx, component in enumerate(message_chain): - if isinstance(component, Record): - plain_comp = await _stt_record(component) - if plain_comp: - message_chain[idx] = plain_comp - event.message_str += plain_comp.text - event.message_obj.message_str += plain_comp.text - - # Also STT for Record components inside Reply chains - for component in event.get_messages(): - if isinstance(component, Reply) and component.chain: - for idx, reply_comp in enumerate(component.chain): - if isinstance(reply_comp, Record): - plain_comp = await _stt_record(reply_comp, is_reply=True) - if plain_comp: - component.chain[idx] = plain_comp - event.message_str += plain_comp.text - event.message_obj.message_str += plain_comp.text +import asyncio +import random +import re +import traceback +from collections.abc import AsyncGenerator +from pathlib import Path + +from astrbot.core import logger +from astrbot.core.message.components import Image, Plain, Record, Reply +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.utils.astrbot_path import get_astrbot_temp_path +from astrbot.core.utils.media_utils import ( + describe_media_ref, + detect_image_mime_type_async, + ensure_wav, + file_uri_to_path, + is_file_uri, +) + +from ..context import PipelineContext +from ..stage import Stage, register_stage + + +@register_stage +class PreProcessStage(Stage): + async def initialize(self, ctx: PipelineContext) -> None: + self.ctx = ctx + self.config = ctx.astrbot_config + self.plugin_manager = ctx.plugin_manager + + self.stt_settings: dict = self.config.get("provider_stt_settings", {}) + self.platform_settings: dict = self.config.get("platform_settings", {}) + + @staticmethod + def _track_temp_media(event: AstrMessageEvent, media_path: str) -> None: + """Track a media file owned by the current event. + + Args: + event: Message event whose lifecycle owns the temporary file. + media_path: Local media path to track when it lives under AstrBot temp. + """ + + try: + path = Path(media_path).resolve() + temp_dir = Path(get_astrbot_temp_path()).resolve() + path.relative_to(temp_dir) + except (OSError, ValueError): + return + event.track_temporary_local_file(str(path)) + + @staticmethod + def _is_existing_local_image_ref(media_ref: str | None) -> bool: + """Return whether an image reference already points at a local file.""" + if not media_ref: + return False + if is_file_uri(media_ref): + return True + if media_ref.startswith(("http://", "https://", "data:", "base64://")): + return False + try: + return Path(media_ref).exists() + except OSError: + return False + + async def _normalize_image_component( + self, + event: AstrMessageEvent, + component: Image, + ) -> None: + """Resolve and validate an image while preserving cleanup ownership.""" + media_ref = component.url or component.file + image_path: str | None = None + materialized = False + try: + image_path = await component.convert_to_file_path() + materialized = ( + not self._is_existing_local_image_ref(media_ref) + and Path(image_path).is_file() + ) + if materialized: + self._track_temp_media(event, image_path) + detected_mime_type = await detect_image_mime_type_async( + image_path, + default_mime_type=None, + ) + if detected_mime_type is None: + raise ValueError("image content could not be identified") + except Exception: + if image_path and not materialized: + event.untrack_temporary_local_file(image_path) + raise + + component.file = image_path + component.path = image_path + # Image.convert_to_file_path() prefers url, so keep it aligned. + component.url = image_path + event.untrack_temporary_local_file(image_path) + + async def process( + self, + event: AstrMessageEvent, + ) -> None | AsyncGenerator[None, None]: + """在处理事件之前的预处理""" + # 平台特异配置:platform_specific..pre_ack_emoji + supported = {"telegram", "lark", "discord"} + platform = event.get_platform_name() + cfg = ( + self.config.get("platform_specific", {}) + .get(platform, {}) + .get("pre_ack_emoji", {}) + ) or {} + emojis = cfg.get("emojis") or [] + if ( + cfg.get("enable", False) + and platform in supported + and emojis + and event.is_at_or_wake_command + ): + try: + await event.react(random.choice(emojis)) + except Exception as e: + logger.warning( + f"Failed to send a pre-response reaction on {platform}: {e}" + ) + + # 路径映射 + if mappings := self.platform_settings.get("path_mapping", []): + # 支持 Record,Image 消息段的路径映射。 + message_chain = event.get_messages() + + for idx, component in enumerate(message_chain): + if isinstance(component, Record | Image) and component.url: + for mapping in mappings: + # ":" is ambiguous for Windows absolute paths. Parse + # the separator after a drive-lettered source first, + # then handle a drive-lettered target. + drive_source_mapping = re.fullmatch( + r"([A-Za-z]:[^:]+):(.+)", mapping + ) + drive_target_mapping = re.fullmatch( + r"(.+):([A-Za-z]:.+)", mapping + ) + if drive_source_mapping: + from_, to_ = drive_source_mapping.groups() + elif drive_target_mapping: + from_, to_ = drive_target_mapping.groups() + else: + from_, separator, to_ = mapping.partition(":") + if not separator: + logger.warning(f"Invalid path mapping: {mapping}") + continue + from_ = from_.removesuffix("/").removesuffix("\\") + to_ = to_.removesuffix("/").removesuffix("\\") + + url = ( + file_uri_to_path(component.url) + if is_file_uri(component.url) + else component.url + ) + if url.startswith(from_): + component.url = url.replace(from_, to_, 1) + logger.debug(f"Path mapping: {url} -> {component.url}") + message_chain[idx] = component + + # Localize source images and normalize audio for downstream processing. + message_chain = event.get_messages() + for idx, component in enumerate(message_chain): + if isinstance(component, Record): + try: + original_path = await component.convert_to_file_path() + self._track_temp_media(event, original_path) + record_path = await ensure_wav(original_path) + self._track_temp_media(event, record_path) + component.file = record_path + component.path = record_path + message_chain[idx] = component + except Exception as e: + logger.warning(f"Voice processing failed: {e}") + elif isinstance(component, Image): + try: + await self._normalize_image_component(event, component) + message_chain[idx] = component + except Exception as e: + media_ref = component.url or component.file + logger.warning( + "Image processing failed for %s: %s", + describe_media_ref(media_ref), + e, + ) + + # Also normalize media components inside Reply chains. + for component in event.get_messages(): + if isinstance(component, Reply) and component.chain: + for idx, reply_comp in enumerate(component.chain): + if isinstance(reply_comp, Record): + try: + original_path = await reply_comp.convert_to_file_path() + self._track_temp_media(event, original_path) + record_path = await ensure_wav(original_path) + self._track_temp_media(event, record_path) + reply_comp.file = record_path + reply_comp.path = record_path + component.chain[idx] = reply_comp + except Exception as e: + logger.warning( + f"Voice processing in reply chain failed: {e}" + ) + elif isinstance(reply_comp, Image): + try: + await self._normalize_image_component(event, reply_comp) + component.chain[idx] = reply_comp + except Exception as e: + media_ref = reply_comp.url or reply_comp.file + logger.warning( + "Image processing in reply chain failed for %s: %s", + describe_media_ref(media_ref), + e, + ) + + # STT + if self.stt_settings.get("enable", False): + # TODO: 独立 + ctx = self.plugin_manager.context + stt_provider = await ctx.get_using_stt_provider_async( + event.unified_msg_origin + ) + if not stt_provider: + logger.warning( + f"Session {event.unified_msg_origin} has no speech-to-text " + "provider configured.", + ) + return + + async def _stt_record(record_comp: Record, is_reply: bool = False): + """对单个 Record 组件执行语音转文本,成功返回 Plain,失败返回 None。""" + prefix = "referenced " if is_reply else "" + try: + path = await record_comp.convert_to_file_path() + except Exception as e: + logger.warning(f"Failed to resolve the {prefix}voice path: {e}") + return None + + retry = 5 + for i in range(retry): + try: + result = await stt_provider.get_text(audio_url=path) + if result: + suffix = " (referenced message)" if is_reply else "" + logger.info(f"Speech-to-text{suffix} result: " + result) + return Plain(result) + break + except FileNotFoundError: + # napcat workaround: file may not be ready immediately + logger.debug( + f"File is not ready ({path}); retrying {i + 1}/{retry}." + ) + await asyncio.sleep(0.5) + continue + except BaseException as e: + logger.error(traceback.format_exc()) + suffix = " (referenced message)" if is_reply else "" + logger.error(f"Speech-to-text{suffix} failed: {e}") + break + return None + + message_chain = event.get_messages() + for idx, component in enumerate(message_chain): + if isinstance(component, Record): + plain_comp = await _stt_record(component) + if plain_comp: + message_chain[idx] = plain_comp + event.message_str += plain_comp.text + event.message_obj.message_str += plain_comp.text + + # Also STT for Record components inside Reply chains + for component in event.get_messages(): + if isinstance(component, Reply) and component.chain: + for idx, reply_comp in enumerate(component.chain): + if isinstance(reply_comp, Record): + plain_comp = await _stt_record(reply_comp, is_reply=True) + if plain_comp: + component.chain[idx] = plain_comp + event.message_str += plain_comp.text + event.message_obj.message_str += plain_comp.text diff --git a/tests/test_preprocess_stage.py b/tests/test_preprocess_stage.py index 017ec5886b..e271c79900 100644 --- a/tests/test_preprocess_stage.py +++ b/tests/test_preprocess_stage.py @@ -1,199 +1,228 @@ -import base64 -from io import BytesIO -from types import SimpleNamespace - -import pytest - -from astrbot.core.message.components import Image, Plain, Reply -from astrbot.core.pipeline.preprocess_stage import stage as preprocess_stage -from astrbot.core.pipeline.preprocess_stage.stage import PreProcessStage -from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.utils import media_utils - - -class FakeEvent: - def __init__(self, message): - self.message_obj = SimpleNamespace(message=message, message_str="") - self.message_str = "" - self.is_at_or_wake_command = False - self.temporary_local_files: list[str] = [] - self._temporary_local_files = self.temporary_local_files - - def get_platform_name(self): - return "test" - - def get_messages(self): - return self.message_obj.message - - track_temporary_local_file = AstrMessageEvent.track_temporary_local_file - untrack_temporary_local_file = AstrMessageEvent.untrack_temporary_local_file - - -@pytest.mark.asyncio -async def test_preprocess_preserves_image_formats_without_tracking_temp_files( - 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)) - monkeypatch.setattr( - preprocess_stage, - "get_astrbot_temp_path", - lambda: str(temp_dir), - ) - 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_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=main_image_ref), - Reply( - id="reply-1", - chain=[Plain(text="quoted"), reply_image], - sender_nickname="Alice", - message_str="quoted", - ), - ] - ) - stage = PreProcessStage() - stage.config = {} - stage.platform_settings = {} - stage.stt_settings = {"enable": False} - - await stage.process(event) - - 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(".png") - assert main_image.file not in event.temporary_local_files - 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(".gif") - assert reply_image.file not in event.temporary_local_files - 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 -@pytest.mark.parametrize("quoted", [False, True]) -@pytest.mark.parametrize( - "source_kind", ["png", "jpeg", "gif", "webp", "bmp", "invalid"] -) -@pytest.mark.parametrize("pretracked", [False, True]) -async def test_preprocess_image_cleanup_preserves_usable_file( - tmp_path, monkeypatch, quoted, source_kind, pretracked -): - from pathlib import Path - - from PIL import Image as PILImage - - monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path)) - monkeypatch.setattr( - preprocess_stage, "get_astrbot_temp_path", lambda: str(tmp_path) - ) - source_path = tmp_path / f"source.{source_kind}" - if source_kind == "invalid": - source_path.write_bytes(b"not an image") - else: - PILImage.new("RGB", (2, 2), (255, 0, 0)).save(source_path) - - image = Image.fromFileSystem(str(source_path)) - event = FakeEvent([Reply(id="reply-1", chain=[image])] if quoted else [image]) - original = source_path.read_bytes() - if pretracked: - event.track_temporary_local_file(str(source_path)) - stage = PreProcessStage() - stage.config = {} - stage.platform_settings = {} - stage.stt_settings = {"enable": False} - - await stage.process(event) - - assert event.temporary_local_files == [] - assert image.file == image.path == image.url == str(source_path) - assert source_path.read_bytes() == original - - # Exercise event cleanup to verify that the usable image survives. - AstrMessageEvent.cleanup_temporary_local_files( - SimpleNamespace(_temporary_local_files=event.temporary_local_files) - ) - assert source_path.read_bytes() == original - assert Path(await image.convert_to_file_path()).exists() - - -@pytest.mark.asyncio -async def test_preprocess_path_mapping_accepts_file_uri(tmp_path): - from PIL import Image as PILImage - - source_root = tmp_path / "source" - target_root = tmp_path / "target" - source_root.mkdir() - target_root.mkdir() - source_image = source_root / "photo.jpg" - target_image = target_root / "photo.jpg" - PILImage.new("RGB", (2, 2), (255, 0, 0)).save(target_image) - event = FakeEvent([Image(file="", url=source_image.as_uri())]) - stage = PreProcessStage() - stage.config = {} - stage.platform_settings = {"path_mapping": [f"{source_root}:{target_root}"]} - stage.stt_settings = {"enable": False} - - await stage.process(event) - - image = event.get_messages()[0] - assert isinstance(image, Image) - assert image.file == image.path == image.url == str(target_image) - - -@pytest.mark.asyncio -async def test_preprocess_path_mapping_accepts_windows_source_to_posix_target( - tmp_path, -): - from PIL import Image as PILImage - - target_root = tmp_path / "target" - target_root.mkdir() - target_image = target_root / "photo.jpg" - PILImage.new("RGB", (2, 2), (255, 0, 0)).save(target_image) - - source_prefix = r"C:\remote\media" - event = FakeEvent([Image(file="", url=f"{source_prefix}/photo.jpg")]) - stage = PreProcessStage() - stage.config = {} - stage.platform_settings = {"path_mapping": [f"{source_prefix}:{target_root}"]} - stage.stt_settings = {"enable": False} - - await stage.process(event) - - image = event.get_messages()[0] - assert isinstance(image, Image) - assert image.file == image.path == image.url == str(target_image) +import base64 +from io import BytesIO +from types import SimpleNamespace + +import pytest + +from astrbot.core.message.components import Image, Plain, Reply +from astrbot.core.pipeline.preprocess_stage import stage as preprocess_stage +from astrbot.core.pipeline.preprocess_stage.stage import PreProcessStage +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.utils import media_utils + + +class FakeEvent: + def __init__(self, message): + self.message_obj = SimpleNamespace(message=message, message_str="") + self.message_str = "" + self.is_at_or_wake_command = False + self.temporary_local_files: list[str] = [] + self._temporary_local_files = self.temporary_local_files + + def get_platform_name(self): + return "test" + + def get_messages(self): + return self.message_obj.message + + track_temporary_local_file = AstrMessageEvent.track_temporary_local_file + untrack_temporary_local_file = AstrMessageEvent.untrack_temporary_local_file + + +@pytest.mark.asyncio +async def test_preprocess_preserves_image_formats_without_tracking_temp_files( + 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)) + monkeypatch.setattr( + preprocess_stage, + "get_astrbot_temp_path", + lambda: str(temp_dir), + ) + 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_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=main_image_ref), + Reply( + id="reply-1", + chain=[Plain(text="quoted"), reply_image], + sender_nickname="Alice", + message_str="quoted", + ), + ] + ) + stage = PreProcessStage() + stage.config = {} + stage.platform_settings = {} + stage.stt_settings = {"enable": False} + + await stage.process(event) + + 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(".png") + assert main_image.file not in event.temporary_local_files + 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(".gif") + assert reply_image.file not in event.temporary_local_files + 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 +@pytest.mark.parametrize("quoted", [False, True]) +@pytest.mark.parametrize( + "source_kind", ["png", "jpeg", "gif", "webp", "bmp", "invalid"] +) +@pytest.mark.parametrize("pretracked", [False, True]) +async def test_preprocess_image_cleanup_preserves_usable_file( + tmp_path, monkeypatch, quoted, source_kind, pretracked +): + from pathlib import Path + + from PIL import Image as PILImage + + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path)) + monkeypatch.setattr( + preprocess_stage, "get_astrbot_temp_path", lambda: str(tmp_path) + ) + source_path = tmp_path / f"source.{source_kind}" + if source_kind == "invalid": + source_path.write_bytes(b"not an image") + else: + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(source_path) + + image = Image.fromFileSystem(str(source_path)) + event = FakeEvent([Reply(id="reply-1", chain=[image])] if quoted else [image]) + original = source_path.read_bytes() + if pretracked: + event.track_temporary_local_file(str(source_path)) + stage = PreProcessStage() + stage.config = {} + stage.platform_settings = {} + stage.stt_settings = {"enable": False} + + await stage.process(event) + + assert event.temporary_local_files == [] + assert image.file == image.path == image.url == str(source_path) + assert source_path.read_bytes() == original + + # Exercise event cleanup to verify that the usable image survives. + AstrMessageEvent.cleanup_temporary_local_files( + SimpleNamespace(_temporary_local_files=event.temporary_local_files) + ) + assert source_path.read_bytes() == original + assert Path(await image.convert_to_file_path()).exists() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("quoted", [False, True]) +async def test_preprocess_image_cleanup_removes_invalid_materialized_file( + tmp_path, monkeypatch, quoted +): + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path)) + monkeypatch.setattr( + preprocess_stage, "get_astrbot_temp_path", lambda: str(tmp_path) + ) + reference = "data:image/png;base64," + base64.b64encode(b"not an image").decode() + image = Image(file=reference) + event = FakeEvent([Reply(id="reply-1", chain=[image])] if quoted else [image]) + stage = PreProcessStage() + stage.config = {} + stage.platform_settings = {} + stage.stt_settings = {"enable": False} + + await stage.process(event) + + materialized = list(tmp_path.glob("media_image_*")) + assert materialized + assert image.file == reference + assert len(event.temporary_local_files) == 1 + AstrMessageEvent.cleanup_temporary_local_files( + SimpleNamespace(_temporary_local_files=event.temporary_local_files) + ) + assert not [path for path in materialized if path.exists()] + + +@pytest.mark.asyncio +async def test_preprocess_path_mapping_accepts_file_uri(tmp_path): + from PIL import Image as PILImage + + source_root = tmp_path / "source" + target_root = tmp_path / "target" + source_root.mkdir() + target_root.mkdir() + source_image = source_root / "photo.jpg" + target_image = target_root / "photo.jpg" + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(target_image) + event = FakeEvent([Image(file="", url=source_image.as_uri())]) + stage = PreProcessStage() + stage.config = {} + stage.platform_settings = {"path_mapping": [f"{source_root}:{target_root}"]} + stage.stt_settings = {"enable": False} + + await stage.process(event) + + image = event.get_messages()[0] + assert isinstance(image, Image) + assert image.file == image.path == image.url == str(target_image) + + +@pytest.mark.asyncio +async def test_preprocess_path_mapping_accepts_windows_source_to_posix_target( + tmp_path, +): + from PIL import Image as PILImage + + target_root = tmp_path / "target" + target_root.mkdir() + target_image = target_root / "photo.jpg" + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(target_image) + + source_prefix = r"C:\remote\media" + event = FakeEvent([Image(file="", url=f"{source_prefix}/photo.jpg")]) + stage = PreProcessStage() + stage.config = {} + stage.platform_settings = {"path_mapping": [f"{source_prefix}:{target_root}"]} + stage.stt_settings = {"enable": False} + + await stage.process(event) + + image = event.get_messages()[0] + assert isinstance(image, Image) + assert image.file == image.path == image.url == str(target_image)