diff --git a/astrbot/core/pipeline/preprocess_stage/stage.py b/astrbot/core/pipeline/preprocess_stage/stage.py index b7b784fb1..704c6820c 100644 --- a/astrbot/core/pipeline/preprocess_stage/stage.py +++ b/astrbot/core/pipeline/preprocess_stage/stage.py @@ -11,6 +11,7 @@ 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, @@ -47,6 +48,54 @@ class PreProcessStage(Stage): 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, @@ -129,14 +178,8 @@ class PreProcessStage(Stage): 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 + await self._normalize_image_component(event, component) 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( @@ -164,13 +207,8 @@ class PreProcessStage(Stage): ) 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 + await self._normalize_image_component(event, reply_comp) 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( diff --git a/tests/test_preprocess_stage.py b/tests/test_preprocess_stage.py index 017ec5886..1d043f97c 100644 --- a/tests/test_preprocess_stage.py +++ b/tests/test_preprocess_stage.py @@ -150,6 +150,35 @@ async def test_preprocess_image_cleanup_preserves_usable_file( 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