diff --git a/backend/app/channels/dingtalk.py b/backend/app/channels/dingtalk.py index 93bfc7d99..ddc405987 100644 --- a/backend/app/channels/dingtalk.py +++ b/backend/app/channels/dingtalk.py @@ -20,8 +20,9 @@ from app.channels.message_bus import InboundMessage, InboundMessageType, Message from deerflow.config.paths import get_paths from deerflow.runtime.user_context import get_effective_user_id from deerflow.sandbox.sandbox_provider import get_sandbox_provider +from deerflow.uploads.async_helpers import publish_upload_bytes_leased_async, release_published_upload_async from deerflow.uploads.layout import upload_virtual_path -from deerflow.uploads.manager import UnsafeUploadPathError, normalize_filename, publish_upload_bytes +from deerflow.uploads.manager import UnsafeUploadPathError, normalize_filename logger = logging.getLogger(__name__) @@ -627,44 +628,41 @@ class DingTalkChannel(Channel): except ValueError: safe_filename = fallback_name - def _persist() -> Path: + def _prepare_upload_dir() -> Path: # Directory prep and collision-safe publication are blocking # filesystem IO, so the whole sequence stays off the event loop. paths.ensure_thread_dirs(thread_id, user_id=effective_user_id) - uploads_dir = paths.sandbox_uploads_dir(thread_id, user_id=effective_user_id) - return publish_upload_bytes(uploads_dir, safe_filename, content) + return paths.sandbox_uploads_dir(thread_id, user_id=effective_user_id) try: - resolved_target = await asyncio.to_thread(_persist) + uploads_dir = await asyncio.to_thread(_prepare_upload_dir) + publication = await publish_upload_bytes_leased_async(uploads_dir, safe_filename, content) except (OSError, UnsafeUploadPathError): logger.exception("[DingTalk] failed to persist downloaded file: %s", safe_filename) return "" - virtual_path = upload_virtual_path(resolved_target.name) - try: - sandbox_provider = get_sandbox_provider() - # acquire_async keeps provider lifecycle work (Docker discovery, - # readiness polls) off the event loop; update_file is blocking - # transport IO on remote sandboxes, so it is offloaded too. - sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id) - if sandbox_id != "local": - sandbox = sandbox_provider.get(sandbox_id) - if sandbox is None: - # Mirror Feishu: the agent's non-local sandbox cannot see this - # file, so returning the virtual path would hand the model a - # path that reads as nothing — surface a failed-load marker. - logger.warning("[DingTalk] sandbox %s not found after acquire, dropping attachment: %s", sandbox_id, virtual_path) - return "" - await asyncio.to_thread(sandbox.update_file, virtual_path, content) - except Exception: - # Same failure mode as the sandbox-is-None branch: the bytes never - # reached the agent's sandbox, so the virtual path would read as - # nothing. Mirror Feishu and surface a failed-load marker. - logger.exception("[DingTalk] failed to sync downloaded file into non-local sandbox: %s", virtual_path) - return "" + virtual_path = upload_virtual_path(publication.path.name) - return virtual_path + try: + sandbox_provider = get_sandbox_provider() + # acquire_async keeps provider lifecycle work (Docker discovery, + # readiness polls) off the event loop; update_file is blocking + # transport IO on remote sandboxes, so it is offloaded too. + sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id) + if sandbox_id != "local": + sandbox = sandbox_provider.get(sandbox_id) + if sandbox is None: + logger.warning("[DingTalk] sandbox %s not found after acquire, dropping attachment: %s", sandbox_id, virtual_path) + return "" + await asyncio.to_thread(sandbox.update_file, virtual_path, content) + except Exception: + logger.exception("[DingTalk] failed to sync downloaded file into non-local sandbox: %s", virtual_path) + return "" + + return virtual_path + finally: + await release_published_upload_async(publication) async def _download_by_code(self, download_code: str) -> bytes | None: """Exchange a DingTalk ``downloadCode`` for the raw file bytes. diff --git a/backend/app/channels/feishu.py b/backend/app/channels/feishu.py index 5220f661d..98e8c0e8a 100644 --- a/backend/app/channels/feishu.py +++ b/backend/app/channels/feishu.py @@ -25,8 +25,8 @@ from app.channels.message_bus import ( from deerflow.config.paths import get_paths from deerflow.runtime.user_context import get_effective_user_id from deerflow.sandbox.sandbox_provider import get_sandbox_provider +from deerflow.uploads.async_helpers import publish_upload_bytes_leased_async, release_published_upload_async from deerflow.uploads.layout import upload_virtual_path -from deerflow.uploads.manager import publish_upload_bytes logger = logging.getLogger(__name__) PENDING_CLARIFICATION_TTL_SECONDS = 30 * 60 @@ -446,34 +446,37 @@ class FeishuChannel(Channel): else: filename = re.sub(r"[./\\]", "_", raw_filename) - def _persist(): + def _prepare_upload_dir(): paths.ensure_thread_dirs(thread_id, user_id=effective_user_id) - uploads_dir = paths.sandbox_uploads_dir(thread_id, user_id=effective_user_id) - return publish_upload_bytes(uploads_dir, filename, content) + return paths.sandbox_uploads_dir(thread_id, user_id=effective_user_id) try: - resolved_target = await asyncio.to_thread(_persist) + uploads_dir = await asyncio.to_thread(_prepare_upload_dir) + publication = await publish_upload_bytes_leased_async(uploads_dir, filename, content) except Exception: logger.exception("[Feishu] failed to persist downloaded resource: %s, type=%s", filename, type) return f"Failed to obtain the [{type}]" - virtual_path = upload_virtual_path(resolved_target.name) - try: - sandbox_provider = await asyncio.to_thread(get_sandbox_provider) - if not getattr(sandbox_provider, "uses_thread_data_mounts", False): - sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id) - sandbox = sandbox_provider.get(sandbox_id) - if sandbox is None: - logger.warning("[Feishu] sandbox not found for thread_id=%s", thread_id) - return f"Failed to obtain the [{type}]" - await asyncio.to_thread(sandbox.update_file, virtual_path, content) - except Exception: - logger.exception("[Feishu] failed to sync resource into non-local sandbox: %s", virtual_path) - return f"Failed to obtain the [{type}]" + virtual_path = upload_virtual_path(publication.path.name) - logger.info("[Feishu] downloaded resource mapped: file_key=%s -> %s", file_key, virtual_path) - return virtual_path + try: + sandbox_provider = await asyncio.to_thread(get_sandbox_provider) + if not getattr(sandbox_provider, "uses_thread_data_mounts", False): + sandbox_id = await sandbox_provider.acquire_async(thread_id, user_id=effective_user_id) + sandbox = sandbox_provider.get(sandbox_id) + if sandbox is None: + logger.warning("[Feishu] sandbox not found for thread_id=%s", thread_id) + return f"Failed to obtain the [{type}]" + await asyncio.to_thread(sandbox.update_file, virtual_path, content) + except Exception: + logger.exception("[Feishu] failed to sync resource into non-local sandbox: %s", virtual_path) + return f"Failed to obtain the [{type}]" + + logger.info("[Feishu] downloaded resource mapped: file_key=%s -> %s", file_key, virtual_path) + return virtual_path + finally: + await release_published_upload_async(publication) # -- message formatting ------------------------------------------------ diff --git a/backend/app/channels/wechat.py b/backend/app/channels/wechat.py index 0e6adb11a..b8b3f22f4 100644 --- a/backend/app/channels/wechat.py +++ b/backend/app/channels/wechat.py @@ -27,7 +27,7 @@ from app.channels.base import Channel 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 deerflow.uploads.manager import UnsafeUploadPathError, publish_upload_bytes +from deerflow.uploads.manager import normalize_filename, publish_upload_bytes logger = logging.getLogger(__name__) @@ -1152,7 +1152,7 @@ class WechatChannel(Channel): try: download_dir.mkdir(parents=True, exist_ok=True) return publish_upload_bytes(download_dir, filename, content) - except (OSError, UnsafeUploadPathError): + except (OSError, ValueError): logger.exception("[WeChat] failed to persist inbound media file %s", filename) return None @@ -1316,11 +1316,15 @@ class WechatChannel(Channel): @staticmethod def _normalize_inbound_filename(raw_filename: Any, *, default_prefix: str, message_id: str, index: int) -> str: + fallback = _safe_media_filename(default_prefix, ".bin", message_id=message_id, index=index) if isinstance(raw_filename, str) and raw_filename.strip(): candidate = Path(raw_filename.strip()).name if candidate: - return candidate - return _safe_media_filename(default_prefix, ".bin", message_id=message_id, index=index) + try: + return normalize_filename(candidate) + except ValueError: + pass + return fallback def _ensure_success(self, data: dict[str, Any], operation: str) -> None: ret = data.get("ret", 0) diff --git a/backend/packages/harness/deerflow/client.py b/backend/packages/harness/deerflow/client.py index 1c96730b9..a624673c2 100644 --- a/backend/packages/harness/deerflow/client.py +++ b/backend/packages/harness/deerflow/client.py @@ -64,12 +64,13 @@ from deerflow.tracing import build_tracing_callbacks, inject_langfuse_metadata from deerflow.uploads.conversion import convert_uploaded_file_to_markdown from deerflow.uploads.layout import artifact_url_for_virtual_path, conversion_virtual_path from deerflow.uploads.manager import ( + PublishedUpload, delete_file_safe, enrich_file_listing, ensure_uploads_dir, get_uploads_dir, list_files_in_dir, - publish_upload_copy, + publish_upload_copy_leased, upload_artifact_url, upload_virtual_path, ) @@ -1527,46 +1528,60 @@ class DeerFlowClient: # creating a new ThreadPoolExecutor per converted file. conversion_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1) - def _convert_in_thread(path: Path): - return asyncio.run(convert_uploaded_file_to_markdown(path)) + def _convert_in_thread(publication: PublishedUpload) -> Path | None: + return asyncio.run( + convert_uploaded_file_to_markdown( + publication.path, + publication=publication, + ) + ) try: for src_path in resolved_files: - dest = publish_upload_copy(uploads_dir, src_path.name, src_path) - dest_name = dest.name + publication = publish_upload_copy_leased(uploads_dir, src_path.name, src_path) + try: + dest = publication.path + dest_name = dest.name - info: dict[str, Any] = { - "filename": dest_name, - "size": dest.stat().st_size, - "path": str(dest), - "virtual_path": upload_virtual_path(dest_name), - "artifact_url": upload_artifact_url(thread_id, dest_name), - } - if dest_name != src_path.name: - info["original_filename"] = src_path.name + info: dict[str, Any] = { + "filename": dest_name, + "size": dest.stat().st_size, + "path": str(dest), + "virtual_path": upload_virtual_path(dest_name), + "artifact_url": upload_artifact_url(thread_id, dest_name), + } + if dest_name != src_path.name: + info["original_filename"] = src_path.name - if src_path.suffix.lower() in CONVERTIBLE_EXTENSIONS: - try: - if conversion_pool is not None: - md_path = conversion_pool.submit(_convert_in_thread, dest).result() - else: - md_path = asyncio.run(convert_uploaded_file_to_markdown(dest)) - except Exception: - logger.warning( - "Failed to convert %s to markdown", - src_path.name, - exc_info=True, - ) - md_path = None + if src_path.suffix.lower() in CONVERTIBLE_EXTENSIONS: + try: + if conversion_pool is not None: + md_path = conversion_pool.submit(_convert_in_thread, publication).result() + else: + md_path = asyncio.run( + convert_uploaded_file_to_markdown( + dest, + publication=publication, + ) + ) + except Exception: + logger.warning( + "Failed to convert %s to markdown", + src_path.name, + exc_info=True, + ) + md_path = None - if md_path is not None: - md_virtual_path = conversion_virtual_path(dest_name) - info["markdown_file"] = md_path.name - info["markdown_path"] = str(md_path) - info["markdown_virtual_path"] = md_virtual_path - info["markdown_artifact_url"] = artifact_url_for_virtual_path(thread_id, md_virtual_path) + if md_path is not None: + md_virtual_path = conversion_virtual_path(dest_name) + info["markdown_file"] = md_path.name + info["markdown_path"] = str(md_path) + info["markdown_virtual_path"] = md_virtual_path + info["markdown_artifact_url"] = artifact_url_for_virtual_path(thread_id, md_virtual_path) - uploaded_files.append(info) + uploaded_files.append(info) + finally: + publication.release() finally: if conversion_pool is not None: conversion_pool.shutdown(wait=True) diff --git a/backend/packages/harness/deerflow/uploads/async_helpers.py b/backend/packages/harness/deerflow/uploads/async_helpers.py new file mode 100644 index 000000000..dc911959b --- /dev/null +++ b/backend/packages/harness/deerflow/uploads/async_helpers.py @@ -0,0 +1,74 @@ +"""Cancellation-safe async adapters for blocking upload publication APIs.""" + +import asyncio +import logging +from pathlib import Path + +from deerflow.uploads.manager import ( + PublishedUpload, + publish_upload_bytes_leased, + rollback_published_upload, +) + +logger = logging.getLogger(__name__) + + +def _rollback_and_release(publication: PublishedUpload) -> None: + try: + rollback_published_upload(publication) + finally: + publication.release() + + +async def _drain_task(task: asyncio.Task) -> None: + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError: + continue + except BaseException: + break + + +async def publish_upload_bytes_leased_async( + base_dir: Path, + preferred_filename: str, + data: bytes, +) -> PublishedUpload: + """Publish bytes off-thread without leaking a lease when cancelled.""" + publish_task = asyncio.create_task( + asyncio.to_thread(publish_upload_bytes_leased, base_dir, preferred_filename, data), + name=f"publish-upload:{preferred_filename}", + ) + try: + return await asyncio.shield(publish_task) + except asyncio.CancelledError: + await _drain_task(publish_task) + if not publish_task.cancelled() and publish_task.exception() is None: + cleanup_task = asyncio.create_task( + asyncio.to_thread(_rollback_and_release, publish_task.result()), + name=f"rollback-cancelled-upload:{preferred_filename}", + ) + await _drain_task(cleanup_task) + try: + cleanup_task.result() + except Exception: + logger.warning("Failed to roll back cancelled upload publication", exc_info=True) + raise + + +async def release_published_upload_async(publication: PublishedUpload) -> None: + """Release a publication off-thread before propagating cancellation.""" + release_task = asyncio.create_task( + asyncio.to_thread(publication.release), + name=f"release-upload:{publication.path.name}", + ) + cancelled = False + while not release_task.done(): + try: + await asyncio.shield(release_task) + except asyncio.CancelledError: + cancelled = True + release_task.result() + if cancelled: + raise asyncio.CancelledError diff --git a/backend/tests/blocking_io/test_dingtalk_receive_file.py b/backend/tests/blocking_io/test_dingtalk_receive_file.py index 02d25801c..a182ab1f3 100644 --- a/backend/tests/blocking_io/test_dingtalk_receive_file.py +++ b/backend/tests/blocking_io/test_dingtalk_receive_file.py @@ -17,6 +17,7 @@ network leg is httpx-async and not the subject here. from __future__ import annotations import asyncio +import threading from types import SimpleNamespace from unittest.mock import AsyncMock @@ -78,3 +79,42 @@ async def test_receive_file_rejects_symlinked_upload_directory(tmp_path, monkeyp assert result == "" assert not await asyncio.to_thread((outside / "a.pdf").exists) + + +async def test_receive_file_holds_name_lease_through_remote_sync(tmp_path, monkeypatch) -> None: + from app.channels.dingtalk import DingTalkChannel + from app.channels.message_bus import MessageBus + from deerflow.config.paths import Paths + from deerflow.uploads.manager import delete_file_safe + + paths = await asyncio.to_thread(Paths, str(tmp_path)) + monkeypatch.setattr("app.channels.dingtalk.get_paths", lambda: paths) + sync_started = threading.Event() + allow_sync = threading.Event() + + class PausedSandbox: + def update_file(self, path, content): + sync_started.set() + assert allow_sync.wait(5) + + async def acquire_async(thread_id, user_id=None): + return "remote" + + monkeypatch.setattr( + "app.channels.dingtalk.get_sandbox_provider", + lambda: SimpleNamespace(acquire_async=acquire_async, get=lambda sandbox_id: PausedSandbox()), + ) + channel = DingTalkChannel(MessageBus(), config={}) + channel._download_by_code = AsyncMock(return_value=b"DATA") + + receive = asyncio.create_task(channel._receive_single_file("dc", "file", "a.pdf", "t1", user_id="default")) + assert await asyncio.to_thread(sync_started.wait, 5) + uploads = paths.sandbox_uploads_dir("t1", user_id="default") + deletion = asyncio.create_task(asyncio.to_thread(delete_file_safe, uploads, "a.pdf")) + await asyncio.sleep(0.05) + assert not deletion.done() + allow_sync.set() + + assert await receive == "/mnt/user-data/uploads/a.pdf" + await deletion + assert not await asyncio.to_thread((uploads / "a.pdf").exists) diff --git a/backend/tests/blocking_io/test_feishu_receive_file.py b/backend/tests/blocking_io/test_feishu_receive_file.py index cb6a27c1c..665b8f268 100644 --- a/backend/tests/blocking_io/test_feishu_receive_file.py +++ b/backend/tests/blocking_io/test_feishu_receive_file.py @@ -106,6 +106,45 @@ async def test_receive_file_remote_sandbox_does_not_block_event_loop(tmp_path, m assert provider.sandbox.update_thread_id != loop_thread_id +async def test_receive_file_holds_name_lease_through_remote_sync(tmp_path, monkeypatch) -> None: + from deerflow.config.paths import Paths + from deerflow.uploads.manager import delete_file_safe + + paths = await asyncio.to_thread(Paths, str(tmp_path)) + provider = _RemoteProvider() + sync_started = threading.Event() + allow_sync = threading.Event() + + def paused_update(path: str, content: bytes) -> None: + sync_started.set() + assert allow_sync.wait(5) + provider.sandbox.updates.append((path, content)) + + provider.sandbox.update_file = paused_update + monkeypatch.setattr("app.channels.feishu.get_paths", lambda: paths) + monkeypatch.setattr("app.channels.feishu.get_sandbox_provider", lambda: provider) + + receive = asyncio.create_task( + _channel_with_file()._receive_single_file( + "message-1", + "file-key", + "file", + "thread-1", + user_id="ou-user", + ) + ) + assert await asyncio.to_thread(sync_started.wait, 5) + uploads = paths.sandbox_uploads_dir("thread-1", user_id="ou-user") + deletion = asyncio.create_task(asyncio.to_thread(delete_file_safe, uploads, "report.pdf")) + await asyncio.sleep(0.05) + assert not deletion.done() + allow_sync.set() + + assert await receive == "/mnt/user-data/uploads/report.pdf" + await deletion + assert not await asyncio.to_thread((uploads / "report.pdf").exists) + + async def test_receive_file_mounted_sandbox_skips_redundant_sync(tmp_path, monkeypatch) -> None: from deerflow.config.paths import Paths diff --git a/backend/tests/test_client.py b/backend/tests/test_client.py index 939ef9164..82ae2b556 100644 --- a/backend/tests/test_client.py +++ b/backend/tests/test_client.py @@ -4,6 +4,7 @@ import asyncio import concurrent.futures import json import tempfile +import threading import zipfile from enum import Enum from pathlib import Path @@ -2122,6 +2123,44 @@ class TestMemoryManagement: class TestUploads: + def test_delete_waits_for_client_conversion_and_metadata(self, client, tmp_path): + uploads_dir = tmp_path / "user-data" / "uploads" + uploads_dir.mkdir(parents=True) + source = tmp_path / "report.pdf" + source.write_bytes(b"PDF") + conversion_started = threading.Event() + allow_conversion = threading.Event() + + async def paused_convert(path: Path, *, publication=None) -> Path: + assert publication is not None + assert publication.path == path + assert publication.is_active + conversion_started.set() + await asyncio.to_thread(allow_conversion.wait) + md_path = conversion_path_for_upload(path) + md_path.parent.mkdir(parents=True, exist_ok=True) + md_path.write_text("converted", encoding="utf-8") + return md_path + + with ( + patch("deerflow.client.ensure_uploads_dir", return_value=uploads_dir), + patch("deerflow.client.get_uploads_dir", return_value=uploads_dir), + patch("deerflow.utils.file_conversion.CONVERTIBLE_EXTENSIONS", {".pdf"}), + patch("deerflow.client.convert_uploaded_file_to_markdown", side_effect=paused_convert), + concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool, + ): + upload = pool.submit(client.upload_files, "thread-1", [source]) + assert conversion_started.wait(5) + deletion = pool.submit(client.delete_upload, "thread-1", "report.pdf") + assert not deletion.done() + allow_conversion.set() + result = upload.result(timeout=5) + deletion.result(timeout=5) + + assert result["files"][0]["markdown_file"] == "report.pdf.md" + assert not (uploads_dir / "report.pdf").exists() + assert not conversion_path_for_upload(uploads_dir / "report.pdf").exists() + def test_upload_files(self, client): with tempfile.TemporaryDirectory() as tmp: tmp_path = Path(tmp) @@ -2213,7 +2252,8 @@ class TestUploads: created_executors = [] real_executor_cls = concurrent.futures.ThreadPoolExecutor - async def fake_convert(path: Path) -> Path: + async def fake_convert(path: Path, *, publication=None) -> Path: + assert publication is not None md_path = conversion_path_for_upload(path) md_path.parent.mkdir(parents=True, exist_ok=True) md_path.write_text(f"converted {path.name}") @@ -2268,7 +2308,8 @@ class TestUploads: markdown = tmp_path / "a.md" markdown.write_bytes(b"USER") - async def fake_convert(path: Path) -> Path: + async def fake_convert(path: Path, *, publication=None) -> Path: + assert publication is not None md_path = conversion_path_for_upload(path) md_path.parent.mkdir(parents=True, exist_ok=True) md_path.write_text(f"FROM:{path.name}", encoding="utf-8") @@ -2304,7 +2345,8 @@ class TestUploads: docx.write_bytes(b"DOCX") pdf.write_bytes(b"PDF") - async def convert_failing_on_docx(path: Path) -> Path | None: + async def convert_failing_on_docx(path: Path, *, publication=None) -> Path | None: + assert publication is not None if path.suffix.lower() == ".docx": return None md_path = conversion_path_for_upload(path) diff --git a/backend/tests/test_wechat_channel.py b/backend/tests/test_wechat_channel.py index 089bc11d2..c97bad63c 100644 --- a/backend/tests/test_wechat_channel.py +++ b/backend/tests/test_wechat_channel.py @@ -1058,6 +1058,29 @@ def test_stage_downloaded_file_renames_around_planted_symlink(tmp_path: Path): assert victim.read_bytes() == b"victim" +def test_wechat_invalid_platform_filename_uses_safe_fallback(tmp_path: Path): + from app.channels.wechat import WechatChannel + + channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token", "state_dir": str(tmp_path)}) + for filename in ["a" * 256 + ".pdf", r"folder\report.pdf", ".upload-user.part"]: + safe = channel._normalize_inbound_filename( + filename, + default_prefix="wechat-file", + message_id="m1", + index=0, + ) + assert safe == "wechat-file-m1-0.bin" + assert channel._stage_downloaded_file(safe, b"payload") is not None + + +def test_wechat_plain_filename_value_error_becomes_attachment_failure(tmp_path: Path): + from app.channels.wechat import WechatChannel + + channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token", "state_dir": str(tmp_path)}) + with mock.patch("app.channels.wechat.publish_upload_bytes", side_effect=ValueError("invalid filename")): + assert channel._stage_downloaded_file("report.pdf", b"payload") is None + + def test_handle_update_downloads_inbound_file_with_media_aeskey_hex(monkeypatch, tmp_path: Path): from app.channels.wechat import WechatChannel