diff --git a/README.md b/README.md index 68b591908d..b2abd8319e 100644 --- a/README.md +++ b/README.md @@ -1217,6 +1217,7 @@ Scheduled-task result and lifecycle notices can also be pushed to a bound Feishu 6. Outbound images/files enforce `max_outbound_image_bytes` / `max_outbound_file_bytes` (20 MiB / 50 MiB defaults) while reading, including files that grow after resolution. Oversize reads are rejected before encryption/upload instead of sending a truncated prefix. Non-positive limits disable the corresponding cap. 7. `allowed_users` takes iLink user IDs. Unset, `null`, `[]`, or a blank string allows everyone. A single ID is one entry, not a sequence of characters, and an unquoted integer-valued number is stored as that integer's text. A scalar string containing commas or interior whitespace logs a warning but remains one literal ID; use a YAML list for multiple IDs. Invalid entries are ignored with a warning; any other configured value that yields no valid ID denies every user and logs an error. `/connect` is still accepted before that check, and a denied sender is dropped before inbound media is downloaded. 8. Shutdown waits for in-flight cursor writes. On token expiry, DeerFlow persists the cursor reset and removes the saved token before completing poller cancellation. +9. Inbound attachments keep separate staged copies even when filenames match within one message or across messages. Their original filenames are preserved for upload ingestion. Private staged copies are removed after ingestion or when a message is dropped or fails. **WeCom Setup** diff --git a/backend/app/channels/AGENTS.md b/backend/app/channels/AGENTS.md index 3eebff65cf..891b10145a 100644 --- a/backend/app/channels/AGENTS.md +++ b/backend/app/channels/AGENTS.md @@ -1,9 +1,9 @@ ### IM Channels System (`app/channels/`) -IM channels reach DeerFlow through `langgraph-sdk` HTTP; Gateway owns thread lifecycle. - -WeChat `_read_outbound_bytes`: cap+1 off-loop; overflow returns `None` before encryption/upload; nonpositive caps are unlimited. -Cursor writes drain before cancellation; token expiry drains cursor/auth clearing under `_auth_lock`. +IM channels use `langgraph-sdk` HTTP; Gateway owns thread lifecycle. +WeChat outbound: off-loop cap+1; overflow `None` before encryption/upload; <=0 unlimited. +Cancel drains cursor writes; expiry drains cursor/auth reset under `_auth_lock`. +Inbound: ingest/drop drains off-loop cleanup of owned dirs, never metadata paths. Telegram retries share bootstrap/Bot cleanup. Dead loops warn once; clear matching futures under `_shutdown_resources_lock` (closure unconfirmed). diff --git a/backend/app/channels/manager.py b/backend/app/channels/manager.py index 4c306b4555..14bdcd7289 100644 --- a/backend/app/channels/manager.py +++ b/backend/app/channels/manager.py @@ -40,6 +40,7 @@ from app.channels.message_bus import ( MessageBus, OutboundMessage, ResolvedAttachment, + cleanup_inbound_files, ) from app.channels.run_policy import CHANNEL_RUN_POLICY, ChannelRunPolicy from app.channels.store import ChannelStore @@ -63,6 +64,7 @@ from deerflow.skills.storage import get_or_new_skill_storage from deerflow.skills.storage.skill_storage import SkillStorage from deerflow.trace_context import ensure_trace_context from deerflow.uploads.manager import apply_upload_sandbox_permits +from deerflow.utils.file_io import await_drained from deerflow.utils.messages import ORIGINAL_USER_CONTENT_KEY logger = logging.getLogger(__name__) @@ -458,7 +460,7 @@ async def _read_wechat_inbound_file(file_info: dict[str, Any], client: httpx.Asy raw_path = file_info.get("path") if isinstance(raw_path, str) and raw_path.strip(): try: - return await asyncio.to_thread(Path(raw_path).read_bytes) + return await await_drained(asyncio.to_thread(Path(raw_path).read_bytes)) except OSError: logger.exception("[Manager] failed to read WeChat inbound file from local path: %s", raw_path) return None @@ -1357,6 +1359,13 @@ def _make_inbound_file_sandbox_readable(file_path: Path) -> None: async def _ingest_inbound_files(thread_id: str, msg: InboundMessage, *, user_id: str | None = None) -> list[dict[str, Any]]: + try: + return await _ingest_inbound_files_impl(thread_id, msg, user_id=user_id) + finally: + await cleanup_inbound_files(msg.files) + + +async def _ingest_inbound_files_impl(thread_id: str, msg: InboundMessage, *, user_id: str | None = None) -> list[dict[str, Any]]: if not msg.files: return [] @@ -2266,7 +2275,7 @@ class ChannelManager: except asyncio.CancelledError: for task in worker_tasks: task.cancel() - discarded_messages = self.bus.discard_pending_inbound() + discarded_messages = await self._discard_pending_inbound() if not join_task.done(): join_task.cancel() if invalidated_reservations or discarded_messages: @@ -2288,7 +2297,7 @@ class ChannelManager: # Anything still queued never started and must not keep Queue.join() # pending after its workers have been cancelled. - discarded_messages = self.bus.discard_pending_inbound() + discarded_messages = await self._discard_pending_inbound() if not join_task.done(): join_task.cancel() @@ -2312,6 +2321,20 @@ class ChannelManager: logger.info("ChannelManager stopped") + async def _discard_pending_inbound(self) -> int: + """Release attachments belonging to messages shutdown never dispatched.""" + discarded = 0 + while True: + try: + msg = self.bus.get_inbound_nowait() + except asyncio.QueueEmpty: + return discarded + try: + await cleanup_inbound_files(msg.files) + finally: + self.bus.inbound_task_done() + discarded += 1 + # -- worker pool ------------------------------------------------------- async def _worker_loop(self, worker_index: int) -> None: @@ -2376,7 +2399,10 @@ class ChannelManager: # worker pool by letting cleanup escape this loop. logger.exception("[Manager] failed to release inbound dedupe key after worker error") finally: - self.bus.inbound_task_done() + try: + await cleanup_inbound_files(msg.files) + finally: + self.bus.inbound_task_done() @staticmethod def _inbound_dedupe_key(msg: InboundMessage) -> tuple[str, str, str, str] | None: @@ -2486,8 +2512,11 @@ class ChannelManager: logger.error("[Manager] unhandled error in message task: %s", exc, exc_info=exc) async def _handle_message(self, msg: InboundMessage) -> None: - with _run_admission_scope(): - await self._dispatch_message(msg) + try: + with _run_admission_scope(): + await self._dispatch_message(msg) + finally: + await cleanup_inbound_files(msg.files) async def _dispatch_message(self, msg: InboundMessage) -> None: msg = _apply_effective_owner(msg) diff --git a/backend/app/channels/message_bus.py b/backend/app/channels/message_bus.py index 2a476ca18f..09d2c83112 100644 --- a/backend/app/channels/message_bus.py +++ b/backend/app/channels/message_bus.py @@ -7,11 +7,14 @@ import logging import threading import time from collections.abc import Callable, Coroutine +from contextlib import suppress from dataclasses import dataclass, field from enum import StrEnum from pathlib import Path from typing import Any +from deerflow.utils.file_io import await_drained + logger = logging.getLogger(__name__) DEFAULT_INBOUND_QUEUE_MAXSIZE = 1000 @@ -29,6 +32,29 @@ RESOLVED_FROM_PENDING_CLARIFICATION_METADATA_KEY = "resolved_from_pending_clarif # Adapter-owned bytes may use this transient key while crossing the channel # boundary. ChannelManager consumes and removes it before persisting metadata. INBOUND_FILE_CONTENT_KEY = "_content" +# Only adapters can supply a callable cleanup capability; a path in provider +# metadata never grants deletion authority. Consume it before persisting files. +INBOUND_FILE_CLEANUP_KEY = "_cleanup" + + +async def cleanup_inbound_files(files: list[dict[str, Any]]) -> None: + """Release adapter-owned staging once, off-loop, including on cancellation.""" + callbacks = [f.pop(INBOUND_FILE_CLEANUP_KEY, None) for f in files if isinstance(f, dict)] + callbacks = [callback for callback in callbacks if callable(callback)] + if not callbacks: + return + + def cleanup() -> None: + for callback in callbacks: + try: + callback() + except Exception as exc: + # A diagnostic failure must not hide the primary error or + # prevent the remaining owned attachments from being released. + with suppress(Exception): + logger.warning("Failed to clean up inbound attachment staging: %s", type(exc).__name__) + + await await_drained(asyncio.to_thread(cleanup)) # --------------------------------------------------------------------------- diff --git a/backend/app/channels/wechat.py b/backend/app/channels/wechat.py index 9cba7864a8..37512acaba 100644 --- a/backend/app/channels/wechat.py +++ b/backend/app/channels/wechat.py @@ -11,10 +11,12 @@ import logging import math import mimetypes import secrets +import shutil import tempfile import time from collections.abc import Mapping from enum import IntEnum +from functools import partial from pathlib import Path from typing import Any from urllib.parse import quote, urlparse @@ -27,7 +29,7 @@ from app.channels.allowed_users import parse_allowed_users from app.channels.base import Channel, ChannelUnavailable from app.channels.commands import is_known_channel_command from app.channels.connection_identity import attach_connection_identity -from app.channels.message_bus import InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment +from app.channels.message_bus import INBOUND_FILE_CLEANUP_KEY, InboundMessage, InboundMessageType, MessageBus, OutboundMessage, ResolvedAttachment, cleanup_inbound_files from deerflow.utils.file_io import await_drained logger = logging.getLogger(__name__) @@ -815,38 +817,43 @@ class WechatChannel(Channel): return files = await self._extract_inbound_files(raw_message) - if not text and not files: - return + published = False + try: + if not text and not files: + return - thread_ts = context_token or str(raw_message.get("client_id") or raw_message.get("msg_id") or "").strip() or None + thread_ts = context_token or str(raw_message.get("client_id") or raw_message.get("msg_id") or "").strip() or None - if context_token: - self._context_tokens_by_chat[chat_id] = context_token - if thread_ts: - self._context_tokens_by_thread[thread_ts] = context_token + if context_token: + self._context_tokens_by_chat[chat_id] = context_token + if thread_ts: + self._context_tokens_by_thread[thread_ts] = context_token - inbound = self._make_inbound( - chat_id=chat_id, - user_id=chat_id, - text=text, - msg_type=InboundMessageType.COMMAND if is_known_channel_command(text) else InboundMessageType.CHAT, - thread_ts=thread_ts, - files=files, - metadata={ - "context_token": context_token, - "ilink_user_id": chat_id, - "message_id": str(raw_message.get("message_id") or raw_message.get("msg_id") or "").strip(), - "ref_msg": self._extract_ref_message(raw_message), - "raw_message": raw_message, - }, - ) - inbound.topic_id = None - inbound = await self._attach_connection_identity(inbound) - if context_token and not raw_message.get("group_id") and inbound.connection_id and inbound.owner_user_id: - await self._remember_notification_context(connection_id=inbound.connection_id, owner_user_id=inbound.owner_user_id, chat_id=chat_id, context_token=context_token) - # The iLink poll loop processes updates sequentially on the Gateway - # loop, so no provider-side task needs a pre-handoff reservation. - await self._publish_inbound_or_drop(inbound) + inbound = self._make_inbound( + chat_id=chat_id, + user_id=chat_id, + text=text, + msg_type=InboundMessageType.COMMAND if is_known_channel_command(text) else InboundMessageType.CHAT, + thread_ts=thread_ts, + files=files, + metadata={ + "context_token": context_token, + "ilink_user_id": chat_id, + "message_id": str(raw_message.get("message_id") or raw_message.get("msg_id") or "").strip(), + "ref_msg": self._extract_ref_message(raw_message), + "raw_message": raw_message, + }, + ) + inbound.topic_id = None + inbound = await self._attach_connection_identity(inbound) + if context_token and not raw_message.get("group_id") and inbound.connection_id and inbound.owner_user_id: + await self._remember_notification_context(connection_id=inbound.connection_id, owner_user_id=inbound.owner_user_id, chat_id=chat_id, context_token=context_token) + # The iLink poll loop processes updates sequentially on the Gateway + # loop, so no provider-side task needs a pre-handoff reservation. + published = await self._publish_inbound_or_drop(inbound) + finally: + if not published: + await cleanup_inbound_files(files) async def _attach_connection_identity(self, inbound: InboundMessage) -> InboundMessage: return await attach_connection_identity( @@ -1300,17 +1307,21 @@ class WechatChannel(Channel): message_id = str(raw_message.get("message_id") or raw_message.get("msg_id") or raw_message.get("client_id") or "msg") - for index, item in enumerate(item_list): - if not isinstance(item, Mapping): - continue - if item.get("type") == int(MessageItemType.IMAGE): - image_file = await self._extract_image_file(item, message_id=message_id, index=index) - if image_file: - files.append(image_file) - elif item.get("type") == int(MessageItemType.FILE): - file_info = await self._extract_file_item(item, message_id=message_id, index=index) - if file_info: - files.append(file_info) + try: + for index, item in enumerate(item_list): + if not isinstance(item, Mapping): + continue + if item.get("type") == int(MessageItemType.IMAGE): + image_file = await self._extract_image_file(item, message_id=message_id, index=index) + if image_file: + files.append(image_file) + elif item.get("type") == int(MessageItemType.FILE): + file_info = await self._extract_file_item(item, message_id=message_id, index=index) + if file_info: + files.append(file_info) + except (Exception, asyncio.CancelledError): + await cleanup_inbound_files(files) + raise return files async def _extract_image_file(self, item: Mapping[str, Any], *, message_id: str, index: int) -> dict[str, Any] | None: @@ -1372,16 +1383,15 @@ class WechatChannel(Channel): detected_image = _detect_image_extension_and_mime(decrypted) image_extension = detected_image[0] if detected_image else ".jpg" filename = _safe_media_filename("wechat-image", image_extension, message_id=message_id, index=index) - stored_path = await asyncio.to_thread(self._stage_downloaded_file, filename, decrypted) - if stored_path is None: + staged_file = await self._stage_inbound_file(filename, decrypted) + if staged_file is None: return None mime_type = detected_image[1] if detected_image else mimetypes.guess_type(filename)[0] or "image/jpeg" return { "type": "image", - "filename": stored_path.name, + **staged_file, "size": len(decrypted), - "path": str(stored_path), "mime_type": mime_type, "source": "wechat", "message_item_type": int(MessageItemType.IMAGE), @@ -1444,35 +1454,63 @@ class WechatChannel(Channel): logger.warning("[WeChat] inbound file exceeds size limit (%d bytes), skipping message_id=%s", len(decrypted), message_id) return None - stored_path = await asyncio.to_thread(self._stage_downloaded_file, filename, decrypted) - if stored_path is None: + staged_file = await self._stage_inbound_file(filename, decrypted) + if staged_file is None: return None return { "type": "file", - "filename": stored_path.name, + **staged_file, "size": len(decrypted), - "path": str(stored_path), "mime_type": mime_type, "source": "wechat", "message_item_type": int(MessageItemType.FILE), "full_url": full_url, } + async def _stage_inbound_file(self, filename: str, content: bytes) -> dict[str, Any] | None: + staged_file: dict[str, Any] = {} + + def stage() -> None: + path = self._stage_downloaded_file(filename, content) + if path is not None: + staged_file.update({"filename": path.name, "path": str(path), INBOUND_FILE_CLEANUP_KEY: partial(shutil.rmtree, path.parent)}) + + try: + # Drain the write before releasing its directory, even if cancelled + # before the worker has returned the new path to the event loop. + await await_drained(asyncio.to_thread(stage)) + except (Exception, asyncio.CancelledError): + await cleanup_inbound_files([staged_file]) + raise + return staged_file or None + def _stage_downloaded_file(self, filename: str, content: bytes) -> Path | None: + if not filename or filename in {".", ".."} or Path(filename).name != filename or "/" in filename or "\\" in filename or "\x00" in filename: + logger.warning("[WeChat] unsafe inbound staging filename, dropping media file") + return None download_dir = self._download_dir() if download_dir is None: # Silent None here made an attachment vanish with no log line — # the same observability gap as a mislabeled skip reason. logger.warning("[WeChat] no state directory configured, dropping staged inbound media file %s", filename) return None + staging_dir = None try: download_dir.mkdir(parents=True, exist_ok=True) - path = download_dir / filename + # Keep the display name while giving each attachment exclusive + # storage until the manager has consumed the published message. + staging_dir = Path(tempfile.mkdtemp(prefix="inbound-", dir=download_dir)) + path = staging_dir / filename path.write_bytes(content) return path except OSError: logger.exception("[WeChat] failed to persist inbound media file %s", filename) + if staging_dir is not None: + try: + shutil.rmtree(staging_dir) + except OSError: + logger.warning("[WeChat] failed to clean up staged inbound media file %s", filename, exc_info=True) return None @staticmethod diff --git a/backend/tests/blocking_io/test_wechat_channel_state.py b/backend/tests/blocking_io/test_wechat_channel_state.py index ff8236a3fd..af76dec8be 100644 --- a/backend/tests/blocking_io/test_wechat_channel_state.py +++ b/backend/tests/blocking_io/test_wechat_channel_state.py @@ -33,7 +33,7 @@ from pathlib import Path import pytest -from app.channels.message_bus import MessageBus +from app.channels.message_bus import MessageBus, cleanup_inbound_files from app.channels.wechat import WechatChannel, _encrypt_aes_128_ecb pytestmark = pytest.mark.asyncio @@ -96,6 +96,8 @@ async def test_wechat_inbound_file_staging_does_not_block_event_loop(tmp_path: P staged = Path(published[0].files[0]["path"]) staged_exists = await asyncio.to_thread(staged.exists) assert staged_exists, "inbound image should be staged under the tmp state dir" + await cleanup_inbound_files(published[0].files) + assert not await asyncio.to_thread(staged.parent.exists), "staging cleanup must also stay off-loop" async def test_wechat_auth_state_load_does_not_block_event_loop(tmp_path: Path) -> None: diff --git a/backend/tests/test_wechat_channel.py b/backend/tests/test_wechat_channel.py index e21567b0f8..662b986a9c 100644 --- a/backend/tests/test_wechat_channel.py +++ b/backend/tests/test_wechat_channel.py @@ -1330,6 +1330,115 @@ def test_handle_update_downloads_inbound_file(monkeypatch, tmp_path: Path): _run(go()) +@pytest.mark.parametrize("separate_updates", [False, True]) +@pytest.mark.parametrize("filename", ["report.pdf", "报告.pdf"]) +def test_handle_update_preserves_same_name_attachments(tmp_path: Path, monkeypatch, separate_updates, filename): + from app.channels.manager import _ingest_inbound_files + from app.channels.wechat import WechatChannel, _encrypt_aes_128_ecb + from deerflow.config.paths import Paths + + upload_paths = Paths(tmp_path / "owner-data") + monkeypatch.setattr("deerflow.uploads.manager.get_paths", lambda: upload_paths) + + async def go(): + bus = MessageBus() + published = [] + + async def capture(msg): + published.append(msg) + + bus.publish_inbound = capture # type: ignore[method-assign] + channel = WechatChannel(bus=bus, config={"bot_token": "test-token", "state_dir": str(tmp_path)}) + aes_key = b"1234567890abcdef" + payloads = [b"FIRST ATTACHMENT", b"SECOND ATTACHMENT"] + encrypted = {f"https://cdn.weixin.qq.com/file-{index}.bin": _encrypt_aes_128_ecb(payload, aes_key) for index, payload in enumerate(payloads)} + + async def download(url, **kwargs): + return encrypted[url] + + channel._download_cdn_bytes = download # type: ignore[method-assign] + items = [{"type": 4, "file_item": {"file_name": filename, "aeskey": aes_key.hex(), "media": {"full_url": url}}} for url in encrypted] + batches = [[item] for item in items] if separate_updates else [items] + for index, batch in enumerate(batches): + await channel._handle_update({"message_type": 1, "message_id": 303 + index, "from_user_id": "wx-user-1", "context_token": "ctx-file", "item_list": batch}) + + assert len(published) == len(batches) + files = [file for msg in published for file in msg.files] + assert len(files) == 2 + paths = [Path(file["path"]) for file in files] + assert len(set(paths)) == 2 + assert [path.read_bytes() for path in paths] == payloads + assert [file["filename"] for file in files] == [filename, filename] + assert [file["size"] for file in files] == [len(payload) for payload in payloads] + + created = [] + for msg in published: + created.extend(await _ingest_inbound_files("wechat-thread", msg, user_id="owner-1")) + saved_names = [filename, f"{Path(filename).stem}_1.pdf"] + assert [file["filename"] for file in created] == saved_names + target = upload_paths.sandbox_uploads_dir("wechat-thread", user_id="owner-1") + assert [(target / name).read_bytes() for name in saved_names] == payloads + assert not upload_paths.sandbox_uploads_dir("wechat-thread").exists() + assert all(not path.parent.exists() for path in paths) + + _run(go()) + + +def test_concurrent_download_staging_preserves_each_payload(tmp_path: Path): + from concurrent.futures import ThreadPoolExecutor + + from app.channels.wechat import WechatChannel + + channel = WechatChannel(bus=MessageBus(), config={"state_dir": str(tmp_path)}) + payloads = [f"ATTACHMENT {index}".encode() for index in range(8)] + with ThreadPoolExecutor(max_workers=8) as pool: + paths = list(pool.map(lambda payload: channel._stage_downloaded_file("report.pdf", payload), payloads)) + + assert all(path is not None for path in paths) + assert len(set(paths)) == len(payloads) + assert [path.name for path in paths] == ["report.pdf"] * len(payloads) + assert [path.read_bytes() for path in paths] == payloads + + +def test_failed_download_staging_cleans_only_its_own_files(tmp_path: Path, monkeypatch, caplog): + from app.channels.wechat import WechatChannel + + channel = WechatChannel(bus=MessageBus(), config={"state_dir": str(tmp_path)}) + first = channel._stage_downloaded_file("report.pdf", b"FIRST ATTACHMENT") + before = set(tmp_path.rglob("*")) + real_write = Path.write_bytes + + def fail_write(path, content): + real_write(path, content[:3]) + (path.parent / "partial").mkdir() + real_write(path.parent / "partial" / "chunk", b"partial write") + raise OSError("simulated disk failure") + + monkeypatch.setattr(Path, "write_bytes", fail_write) + with caplog.at_level(logging.ERROR): + second = channel._stage_downloaded_file("report.pdf", b"SECOND ATTACHMENT") + + assert second is None + assert first.read_bytes() == b"FIRST ATTACHMENT" + assert set(tmp_path.rglob("*")) == before + assert "failed to persist inbound media file" in caplog.text + + +def test_download_staging_allocation_failure_is_reported(tmp_path: Path, monkeypatch, caplog): + from app.channels.wechat import WechatChannel + + channel = WechatChannel(bus=MessageBus(), config={"state_dir": str(tmp_path)}) + + def fail_allocate(**kwargs): + raise OSError("simulated directory allocation failure") + + monkeypatch.setattr("app.channels.wechat.tempfile.mkdtemp", fail_allocate) + with caplog.at_level(logging.ERROR): + assert channel._stage_downloaded_file("report.pdf", b"ATTACHMENT") is None + assert list((tmp_path / "downloads").iterdir()) == [] + assert "failed to persist inbound media file" in caplog.text + + def test_handle_update_downloads_inbound_file_with_media_aeskey_hex(monkeypatch, tmp_path: Path): from app.channels.wechat import WechatChannel diff --git a/backend/tests/test_wechat_inbound_staging_cleanup.py b/backend/tests/test_wechat_inbound_staging_cleanup.py new file mode 100644 index 0000000000..2589170047 --- /dev/null +++ b/backend/tests/test_wechat_inbound_staging_cleanup.py @@ -0,0 +1,246 @@ +"""WeChat staging ownership ends on ingestion, rejection, failure or cancellation.""" + +from __future__ import annotations + +import asyncio +import threading +from pathlib import Path +from unittest.mock import AsyncMock, Mock + +import pytest + + +@pytest.fixture +def staging(tmp_path, monkeypatch): + from app.channels.message_bus import MessageBus + from app.channels.wechat import WechatChannel, _encrypt_aes_128_ecb + from deerflow.config.paths import Paths + + paths = Paths(tmp_path / "uploads") + monkeypatch.setattr("deerflow.uploads.manager.get_paths", lambda: paths) + channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test", "state_dir": str(tmp_path / "wechat")}) + key = b"1234567890abcdef" + channel._download_cdn_bytes = AsyncMock(return_value=_encrypt_aes_128_ecb(b"attachment", key)) + raw = { + "message_type": 1, + "message_id": "message-1", + "from_user_id": "wx-user", + "item_list": [{"type": 4, "file_item": {"file_name": "report.pdf", "aeskey": key.hex(), "media": {"full_url": "https://cdn.weixin.qq.com/file"}}}], + } + return channel, raw + + +async def _queued_message(staging): + channel, raw = staging + await channel._handle_update(raw) + msg = channel.bus.get_inbound_nowait() + channel.bus.inbound_task_done() + return msg + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["prepare", "read", "write", "unsafe_name"]) +async def test_ingestion_failure_releases_owned_staging(staging, monkeypatch, failure): + from app.channels.manager import INBOUND_FILE_READERS, _ingest_inbound_files + + msg = await _queued_message(staging) + path = Path(msg.files[0]["path"]) + if failure == "read": + monkeypatch.setitem(INBOUND_FILE_READERS, "wechat", AsyncMock(side_effect=OSError("read failed"))) + elif failure == "unsafe_name": + msg.files[0]["filename"] = ".." + else: + target = "ensure_uploads_dir" if failure == "prepare" else "write_upload_file_no_symlink" + monkeypatch.setattr(f"deerflow.uploads.manager.{target}", Mock(side_effect=OSError("upload failed"))) + if failure == "prepare": + with pytest.raises(OSError, match="upload failed"): + await _ingest_inbound_files("thread", msg, user_id="owner") + else: + assert await _ingest_inbound_files("thread", msg, user_id="owner") == [] + assert not path.parent.exists() + + +@pytest.mark.asyncio +async def test_cleanup_uses_allocated_directory_not_mutable_file_path(staging, tmp_path): + from app.channels.manager import _ingest_inbound_files + + msg = await _queued_message(staging) + owned = Path(msg.files[0]["path"]) + foreign = tmp_path / "foreign" / "report.pdf" + foreign.parent.mkdir() + foreign.write_bytes(b"unrelated") + msg.files[0]["path"] = str(foreign) + assert len(await _ingest_inbound_files("thread", msg, user_id="owner")) == 1 + assert not owned.parent.exists() + assert foreign.read_bytes() == b"unrelated" + + +@pytest.mark.asyncio +async def test_unowned_local_path_is_not_removed(staging, tmp_path): + from app.channels.manager import _ingest_inbound_files + from app.channels.message_bus import InboundMessage + + foreign = tmp_path / "inbound-foreign" / "report.pdf" + foreign.parent.mkdir() + foreign.write_bytes(b"unrelated") + msg = InboundMessage(channel_name="wechat", chat_id="chat", user_id="user", text="", files=[{"path": str(foreign), "filename": foreign.name, "_cleanup": str(foreign.parent)}]) + assert len(await _ingest_inbound_files("thread", msg, user_id="owner")) == 1 + assert foreign.read_bytes() == b"unrelated" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("disposition", ["duplicate", "rejected", "handler_error", "queued_shutdown"]) +async def test_manager_releases_staging_without_ingestion(staging, disposition): + from app.channels.manager import ChannelManager + + channel, raw = staging + await channel._handle_update(raw) + paths = list(channel._download_dir().glob("inbound-*")) + assert len(paths) == 1 + manager = ChannelManager(bus=channel.bus, store=Mock()) + manager._is_duplicate_inbound = AsyncMock(return_value=disposition == "duplicate") + manager._dispatch_message = AsyncMock(side_effect=RuntimeError("handler failed") if disposition == "handler_error" else None) + if disposition == "queued_shutdown": + manager._shutdown_grace_period_seconds = 0 + await manager.stop() + else: + await manager._worker_loop(0) + assert not paths[0].exists() + await asyncio.wait_for(channel.bus.join_inbound(), 1) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("failure", ["queue_closed", "queue_full", "identity", "publication"]) +async def test_adapter_retains_cleanup_until_publication_succeeds(staging, failure): + from app.channels.message_bus import InboundMessage, MessageBus + + channel, raw = staging + if failure == "queue_closed": + channel.bus.close_inbound() + elif failure == "queue_full": + channel.bus = MessageBus(inbound_queue_maxsize=1) + await channel.bus.publish_inbound(InboundMessage(channel_name="test", chat_id="chat", user_id="user", text="full")) + elif failure == "identity": + channel._attach_connection_identity = AsyncMock(side_effect=RuntimeError("identity failed")) + else: + channel.bus.publish_inbound = AsyncMock(side_effect=RuntimeError("publication failed")) + if failure in {"identity", "publication"}: + with pytest.raises(RuntimeError, match="failed"): + await channel._handle_update(raw) + else: + await channel._handle_update(raw) + assert list(channel._download_dir().iterdir()) == [] + + +@pytest.mark.asyncio +async def test_partial_extraction_failure_releases_earlier_attachments(staging): + channel, raw = staging + extract = channel._extract_file_item + calls = 0 + + async def fail_second(*args, **kwargs): + nonlocal calls + calls += 1 + if calls == 2: + raise RuntimeError("second attachment failed") + return await extract(*args, **kwargs) + + channel._extract_file_item = fail_second + raw["item_list"] *= 2 + with pytest.raises(RuntimeError, match="second attachment failed"): + await channel._handle_update(raw) + assert list(channel._download_dir().iterdir()) == [] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("phase", ["staging", "reading", "cleanup"]) +async def test_cancellation_drains_file_io_before_releasing_staging(staging, monkeypatch, phase): + from app.channels.manager import _ingest_inbound_files + + channel, raw = staging + started = threading.Event() + release = threading.Event() + if phase == "staging": + stage = channel._stage_downloaded_file + + def blocked_stage(*args): + path = stage(*args) + started.set() + assert release.wait(5) + return path + + channel._stage_downloaded_file = blocked_stage + operation = channel._handle_update(raw) + else: + if phase == "cleanup": + import shutil + + remove = shutil.rmtree + + def blocked_remove(*args, **kwargs): + started.set() + assert release.wait(5) + return remove(*args, **kwargs) + + monkeypatch.setattr("app.channels.wechat.shutil.rmtree", blocked_remove) + msg = await _queued_message(staging) + source = Path(msg.files[0]["path"]) + read = Path.read_bytes + + def blocked_read(path): + if phase == "reading" and path == source: + started.set() + assert release.wait(5) + return read(path) + + monkeypatch.setattr(Path, "read_bytes", blocked_read) + operation = _ingest_inbound_files("thread", msg, user_id="owner") + task = asyncio.create_task(operation) + try: + assert await asyncio.to_thread(started.wait, 5) + task.cancel() + await asyncio.sleep(0) + task.cancel() + await asyncio.sleep(0) + assert not task.done(), "cancellation must drain the worker before cleanup" + finally: + release.set() + results = await asyncio.gather(task, return_exceptions=True) + assert isinstance(results[0], asyncio.CancelledError) + assert list(channel._download_dir().iterdir()) == [] + + +@pytest.mark.parametrize("filename", ["../keep.txt", "sub/report.pdf", "sub\\report.pdf", "", ".", "..", "bad\x00name.pdf"]) +def test_staging_rejects_non_basename_paths(staging, filename): + channel, _ = staging + downloads = channel._download_dir() + downloads.mkdir(parents=True) + sibling = downloads / "keep.txt" + sibling.write_bytes(b"keep") + assert channel._stage_downloaded_file(filename, b"overwrite") is None + assert sibling.read_bytes() == b"keep" + assert list(downloads.iterdir()) == [sibling] + + +@pytest.mark.asyncio +async def test_cleanup_failure_preserves_primary_error_and_releases_other_files(staging, monkeypatch): + from app.channels.manager import _ingest_inbound_files + from app.channels.message_bus import INBOUND_FILE_CLEANUP_KEY + + channel, raw = staging + raw["item_list"] *= 2 + msg = await _queued_message(staging) + paths = [Path(file["path"]) for file in msg.files] + assert len(paths) == 2 + cleanup = msg.files[0][INBOUND_FILE_CLEANUP_KEY] + msg.files[0][INBOUND_FILE_CLEANUP_KEY] = Mock(side_effect=OSError("cleanup failed")) + monkeypatch.setattr("deerflow.uploads.manager.ensure_uploads_dir", Mock(side_effect=OSError("prepare failed"))) + monkeypatch.setattr("app.channels.message_bus.logger.warning", Mock(side_effect=RuntimeError("logging failed"))) + try: + with pytest.raises(OSError, match="prepare failed"): + await _ingest_inbound_files("thread", msg, user_id="owner") + assert paths[0].exists() + assert not paths[1].parent.exists() + assert all(INBOUND_FILE_CLEANUP_KEY not in file for file in msg.files) + finally: + cleanup()