fix: retain upload leases through adapter sync

This commit is contained in:
hetaoBackend 2026-08-06 20:19:13 +08:00
parent cadb3a444a
commit 920dc247d2
9 changed files with 327 additions and 89 deletions

View File

@ -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.

View File

@ -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 ------------------------------------------------

View File

@ -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)

View File

@ -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)

View File

@ -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

View File

@ -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)

View File

@ -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

View File

@ -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)

View File

@ -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