From 72928aa7c976070ae22d3e46e9bed0525b861c2e Mon Sep 17 00:00:00 2001 From: Soulter <37870767+Soulter@users.noreply.github.com> Date: Sat, 12 Sep 2026 12:31:10 +0800 Subject: [PATCH] fix: preserve image attachment paths after event cleanup (#10042) * fix: preserve image attachment paths after event cleanup * fix: retain local image references when normalization fails * revert: undo local image reference normalization change --- astrbot/core/astr_main_agent.py | 4 +- .../core/pipeline/preprocess_stage/stage.py | 8 +-- .../sources/mattermost/mattermost_adapter.py | 8 +-- .../sources/mattermost/mattermost_event.py | 2 - tests/test_mattermost_adapter.py | 71 +++++++++++++++++++ tests/test_preprocess_stage.py | 53 +++++++++++++- tests/unit/test_astr_main_agent.py | 64 +++++++++++++++-- 7 files changed, 188 insertions(+), 22 deletions(-) diff --git a/astrbot/core/astr_main_agent.py b/astrbot/core/astr_main_agent.py index 4a4456b41..3148d638b 100644 --- a/astrbot/core/astr_main_agent.py +++ b/astrbot/core/astr_main_agent.py @@ -1471,7 +1471,7 @@ async def build_main_agent( event.track_temporary_local_file(image_path) req.image_urls.append(image_path) req.extra_user_content_parts.append( - TextPart(text=f"[Image Attachment: path {image_path}]") + TextPart(text=f"[Image Attachment: path {path}]") ) elif isinstance(comp, Record): audio_path = await comp.convert_to_file_path() @@ -1509,7 +1509,7 @@ async def build_main_agent( if _is_generated_compressed_image_path(path, image_path): event.track_temporary_local_file(image_path) req.image_urls.append(image_path) - _append_quoted_image_attachment(req, image_path) + _append_quoted_image_attachment(req, path) elif isinstance(reply_comp, Record): audio_path = await reply_comp.convert_to_file_path() req.audio_urls.append(audio_path) diff --git a/astrbot/core/pipeline/preprocess_stage/stage.py b/astrbot/core/pipeline/preprocess_stage/stage.py index d54784eda..ffebb3e35 100644 --- a/astrbot/core/pipeline/preprocess_stage/stage.py +++ b/astrbot/core/pipeline/preprocess_stage/stage.py @@ -131,9 +131,9 @@ class PreProcessStage(Stage): elif isinstance(component, Image): try: original_path = await component.convert_to_file_path() - self._track_temp_media(event, original_path) image_path = await ensure_jpeg(original_path) - self._track_temp_media(event, image_path) + if image_path != original_path: + self._track_temp_media(event, original_path) component.file = image_path component.path = image_path # Image.convert_to_file_path() prefers url, so keep it aligned. @@ -167,9 +167,9 @@ class PreProcessStage(Stage): elif isinstance(reply_comp, Image): try: original_path = await reply_comp.convert_to_file_path() - self._track_temp_media(event, original_path) image_path = await ensure_jpeg(original_path) - self._track_temp_media(event, image_path) + if image_path != original_path: + self._track_temp_media(event, original_path) reply_comp.file = image_path reply_comp.path = image_path # Image.convert_to_file_path() prefers url, so keep it aligned. diff --git a/astrbot/core/platform/sources/mattermost/mattermost_adapter.py b/astrbot/core/platform/sources/mattermost/mattermost_adapter.py index 7703c501f..1f4592d4f 100644 --- a/astrbot/core/platform/sources/mattermost/mattermost_adapter.py +++ b/astrbot/core/platform/sources/mattermost/mattermost_adapter.py @@ -230,12 +230,10 @@ class MattermostPlatformAdapter(Platform): ) if file_ids: - ( - attachment_components, - temp_paths, - ) = await self.client.parse_post_attachments(file_ids) + attachment_components, _ = await self.client.parse_post_attachments( + file_ids + ) abm.message.extend(attachment_components) - setattr(abm, "temporary_file_paths", temp_paths) abm.message_str = self._build_message_str( abm.message, diff --git a/astrbot/core/platform/sources/mattermost/mattermost_event.py b/astrbot/core/platform/sources/mattermost/mattermost_event.py index 22c445a6f..850274c2b 100644 --- a/astrbot/core/platform/sources/mattermost/mattermost_event.py +++ b/astrbot/core/platform/sources/mattermost/mattermost_event.py @@ -23,8 +23,6 @@ class MattermostMessageEvent(AstrMessageEvent): ) -> None: super().__init__(message_str, message_obj, platform_meta, session_id) self.client = client - for path in getattr(message_obj, "temporary_file_paths", []): - self.track_temporary_local_file(path) async def send(self, message: MessageChain) -> None: await self.client.send_message_chain(self.get_session_id(), message) diff --git a/tests/test_mattermost_adapter.py b/tests/test_mattermost_adapter.py index 1f414fd29..13e1e86ec 100644 --- a/tests/test_mattermost_adapter.py +++ b/tests/test_mattermost_adapter.py @@ -6,10 +6,12 @@ import pytest import astrbot.api.message_components as Comp from astrbot.api.platform import Group +from astrbot.core.pipeline.preprocess_stage import stage as preprocess_stage from astrbot.core.platform.sources.mattermost.client import MattermostClient from astrbot.core.platform.sources.mattermost.mattermost_adapter import ( MattermostPlatformAdapter, ) +from astrbot.core.utils import media_utils from tests.fixtures.helpers import make_platform_config @@ -116,6 +118,75 @@ async def test_mattermost_parse_post_attachments_maps_media_types(tmp_path): assert path.name.endswith(Path(expected_name).suffix) +@pytest.mark.asyncio +@pytest.mark.parametrize("image_kind", ["png", "jpeg", "invalid"]) +async def test_mattermost_attachments_follow_preprocessing_cleanup_rules( + tmp_path, monkeypatch, image_kind +): + import wave + + 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) + ) + image_path = tmp_path / f"image.{image_kind}" + if image_kind == "invalid": + image_path.write_bytes(b"not an image") + else: + PILImage.new("RGB", (2, 2), (255, 0, 0)).save(image_path) + audio_path = tmp_path / "voice.wav" + with wave.open(str(audio_path), "wb") as audio_file: + audio_file.setparams((1, 2, 8000, 0, "NONE", "not compressed")) + audio_file.writeframes(b"\x00\x00" * 80) + video_path = tmp_path / "clip.mp4" + video_path.write_bytes(b"video") + document_path = tmp_path / "report.pdf" + document_path.write_bytes(b"document") + downloaded_paths = [image_path, audio_path, video_path, document_path] + image = Comp.Image.fromFileSystem(str(image_path)) + adapter = _build_adapter() + adapter.client.parse_post_attachments = AsyncMock( + return_value=( + [ + image, + Comp.Record(file=str(audio_path), url=str(audio_path)), + Comp.Video.fromFileSystem(str(video_path)), + Comp.File(name="report.pdf", file=str(document_path)), + ], + [str(path) for path in downloaded_paths], + ) + ) + message = await adapter.convert_message( + post={ + "id": "post-1", + "channel_id": "channel-1", + "user_id": "user-1", + "message": "hello", + "file_ids": ["img", "audio", "video", "doc"], + }, + data={"channel_type": "D", "sender_name": "alice"}, + ) + event = adapter.create_event(message) + assert event._temporary_local_files == [] + stage = preprocess_stage.PreProcessStage() + stage.config = {} + stage.platform_settings = {} + stage.stt_settings = {"enable": False} + + await stage.process(event) + event.cleanup_temporary_local_files() + + assert image_path.exists() == (image_kind != "png") + assert not audio_path.exists() + assert video_path.exists() + assert document_path.exists() + assert Path(await image.convert_to_file_path()).exists() + if image_kind == "png": + assert image.file != str(image_path) + + @pytest.mark.asyncio async def test_mattermost_get_group_returns_members_and_channel_admins(): adapter = _build_adapter() diff --git a/tests/test_preprocess_stage.py b/tests/test_preprocess_stage.py index 1cc0b0cc9..49ec60fb5 100644 --- a/tests/test_preprocess_stage.py +++ b/tests/test_preprocess_stage.py @@ -7,6 +7,7 @@ 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 @@ -29,7 +30,7 @@ class FakeEvent: @pytest.mark.asyncio -async def test_preprocess_preserves_image_formats_and_tracks_temp_files( +async def test_preprocess_preserves_image_formats_without_tracking_temp_files( tmp_path, monkeypatch ): from PIL import Image as PILImage @@ -88,20 +89,66 @@ async def test_preprocess_preserves_image_formats_and_tracks_temp_files( 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 in event.temporary_local_files + 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 in event.temporary_local_files + 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", "invalid"]) +async def test_preprocess_image_cleanup_preserves_usable_file( + tmp_path, monkeypatch, quoted, source_kind +): + 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]) + stage = PreProcessStage() + stage.config = {} + stage.platform_settings = {} + stage.stt_settings = {"enable": False} + + await stage.process(event) + + if source_kind == "png": + assert event.temporary_local_files == [str(source_path)] + assert image.file == image.path == image.url + assert image.file != str(source_path) + with PILImage.open(image.file) as processed_img: + assert processed_img.format == "JPEG" + else: + assert event.temporary_local_files == [] + + # 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.exists() == (source_kind != "png") + 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 diff --git a/tests/unit/test_astr_main_agent.py b/tests/unit/test_astr_main_agent.py index 78284124b..e41380fcd 100644 --- a/tests/unit/test_astr_main_agent.py +++ b/tests/unit/test_astr_main_agent.py @@ -1873,14 +1873,33 @@ class TestBuildMainAgent: assert result is None @pytest.mark.asyncio + @pytest.mark.parametrize("quoted", [False, True]) + @pytest.mark.parametrize("compression_enabled", [False, True]) async def test_build_main_agent_with_images( - self, mock_event, mock_context, mock_provider + self, + mock_event, + mock_context, + mock_provider, + tmp_path, + monkeypatch, + quoted, + compression_enabled, ): - """Test building main agent with image attachments.""" + """Keep attachment paths usable after compressed visual input is cleaned up.""" + from types import SimpleNamespace + + from PIL import Image as PILImage + + from astrbot.core.utils import media_utils + module = ama - mock_image = MagicMock(spec=Image) - mock_image.convert_to_file_path = AsyncMock(return_value="/path/to/image.jpg") - mock_event.message_obj.message = [mock_image] + monkeypatch.setattr(media_utils, "get_astrbot_temp_path", lambda: str(tmp_path)) + source_path = tmp_path / "image.jpg" + PILImage.new("RGB", (8, 8), (255, 0, 0)).save(source_path) + image = Image.fromFileSystem(str(source_path)) + mock_event.message_obj.message = ( + [Reply(id="reply-1", chain=[image])] if quoted else [image] + ) mock_context.get_provider_by_id.return_value = None mock_context.get_using_provider.return_value = mock_provider @@ -1900,10 +1919,43 @@ class TestBuildMainAgent: result = await module.build_main_agent( event=mock_event, plugin_context=mock_context, - config=module.MainAgentBuildConfig(tool_call_timeout=60), + config=module.MainAgentBuildConfig( + tool_call_timeout=60, + provider_settings={ + "image_compress_enabled": compression_enabled, + "image_compress_options": {"max_size": 2}, + }, + ), ) assert result is not None + request = result.provider_request + label = "Image Attachment in quoted message" if quoted else "Image Attachment" + assert f"[{label}: path {source_path}]" in [ + part.text for part in request.extra_user_content_parts + ] + assert len(request.image_urls) == 1 + visual_path = Path(request.image_urls[0]) + if compression_enabled: + assert visual_path != source_path + with PILImage.open(visual_path) as visual_image: + assert visual_image.size == (2, 2) + mock_event.track_temporary_local_file.assert_called_once_with( + str(visual_path) + ) + else: + assert visual_path == source_path + mock_event.track_temporary_local_file.assert_not_called() + AstrMessageEvent.cleanup_temporary_local_files( + SimpleNamespace( + _temporary_local_files=[ + call.args[0] + for call in mock_event.track_temporary_local_file.call_args_list + ] + ) + ) + assert source_path.exists() + assert visual_path.exists() == (not compression_enabled) @pytest.mark.asyncio async def test_build_main_agent_skips_caption_when_main_provider_supports_images(