mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-15 00:19:14 +00:00
fix: close remaining upload lifecycle races
This commit is contained in:
parent
96ba5dccd0
commit
a8f00f30fd
@ -1461,11 +1461,11 @@ Multi-file upload with automatic document conversion:
|
||||
- Reuses one conversion worker per request when called from an active event loop
|
||||
- Files stored in thread-isolated directories under the resolving user's bucket (`users/{user_id}/threads/{thread_id}/user-data/uploads`). For IM channels the owner is threaded explicitly via the `user_id=` kwarg (see IM Channels → Owner-scoped file storage); HTTP/embedded callers resolve it from `get_effective_user_id()`
|
||||
- Every ingress stages a complete payload and atomically publishes it without replacing an existing entry. Collisions across requests, processes, HTTP, embedded, and IM adapters use `name.ext`, `name_1.ext`, `name_2.ext`; storage that cannot provide atomic no-replace publication fails explicitly.
|
||||
- Exact-name generation leases use a portable NFC-plus-casefold coordination key, so case and Unicode-normalization aliases cannot bypass an active generation on case-insensitive filesystems. The original filename remains the published name. Final lease release is the commit point: cancellation newly arriving during release is delayed and swallowed so a committed upload is returned as success rather than an indeterminate cancelled result.
|
||||
- Exact-name generation leases use a portable NFC-plus-casefold coordination key, so case and Unicode-normalization aliases cannot bypass an active generation on case-insensitive filesystems. The original filename remains the published name. Multi-file Gateway and embedded-client requests reserve those portable keys while retaining their publication leases, so a later alias chooses an `_N` candidate instead of waiting on its own batch. Alias-based deletion resolves the primary's actual directory entry by inode before deriving a long-name conversion path. Final lease release is the commit point: cancellation newly arriving during release is delayed and swallowed so a committed upload is returned as success rather than an indeterminate cancelled result.
|
||||
- Filenames containing NUL, `<`, `>`, or reserved model-context boundary markers are rejected before staging so accepted filenames and exact virtual paths remain lossless in model-visible upload context. Legacy files discovered on disk are still neutralized when listed.
|
||||
- Gateway HTTP uploads use same-directory `.upload-*.part` staging files. Each active stage holds a cross-process liveness lock under `.upload-conversions/.locks/stages/`; startup cleanup skips held stages and sweeps only crash-orphaned files. Cancellation during staging creation drains the worker and aborts the returned stage before propagating. Staging files are hidden from upload listings, agent upload context, and sandbox listing/search tools.
|
||||
- Generated Markdown is owned by `user-data/.upload-conversions/<actual-primary-filename>.md` and is omitted from primary upload listings. Deletion removes only the selected primary and that exact generated asset; it never guesses or deletes a legacy/user-owned `uploads/<stem>.md` sibling.
|
||||
- Gateway HTTP upload/list/delete handlers offload filesystem work through `deerflow.utils.file_io.run_file_io`, a dedicated ContextVar-preserving file IO executor; potentially blocking name-lease acquisition uses a separate pool so waiters cannot starve release, and cold sandbox-provider construction is also offloaded. Gateway, embedded-client, and IM ingresses share provider-aware publication: mounted providers make the exact host paths sandbox-readable; non-mounted providers acquire the sandbox and synchronize the primary plus generated conversion to their exact virtual paths. Each ingress records attempted remote paths before the write can commit and, on later failure or cancellation, calls the provider-neutral `Sandbox.remove_file()` for those paths before host rollback and lease release. Embedded multi-file calls retain every publication and receipt until the whole response is built, then roll back the complete batch on failure.
|
||||
- Gateway HTTP upload/list/delete handlers offload filesystem work through `deerflow.utils.file_io.run_file_io`, a dedicated ContextVar-preserving file IO executor; potentially blocking publication and deletion name-lease acquisition uses a separate pool so waiters cannot starve general-pool release work, and cold sandbox-provider construction is also offloaded. Gateway, embedded-client, and IM ingresses share provider-aware publication: mounted providers make the exact host paths sandbox-readable; non-mounted providers acquire the sandbox and synchronize the primary plus generated conversion to their exact virtual paths. Each ingress records attempted remote paths before the write can commit and, on later failure or cancellation, calls the provider-neutral `Sandbox.remove_file()` for those paths before host rollback and lease release. Embedded multi-file calls retain every publication and receipt until the whole response is built, then roll back the complete batch on failure. WeChat download publication uses the cancellation-safe async lease adapter, so cancellation drains and rolls back a publication worker that completes late.
|
||||
- Mounted upload paths skip both sandbox acquisition and per-file synchronization. For AIO remote/provisioner deployments this requires an explicit, accurate `sandbox.thread_data_mounts: true`; omission preserves backend auto-detection.
|
||||
- Agent receives uploaded file list via `UploadsMiddleware`
|
||||
|
||||
|
||||
@ -27,7 +27,8 @@ 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 normalize_filename, publish_upload_bytes
|
||||
from deerflow.uploads.async_helpers import publish_upload_bytes_leased_async, release_published_upload_async
|
||||
from deerflow.uploads.manager import normalize_filename
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
@ -1079,7 +1080,7 @@ 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)
|
||||
stored_path = await self._stage_downloaded_file(filename, decrypted)
|
||||
if stored_path is None:
|
||||
return None
|
||||
|
||||
@ -1130,7 +1131,7 @@ 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)
|
||||
stored_path = await self._stage_downloaded_file(filename, decrypted)
|
||||
if stored_path is None:
|
||||
return None
|
||||
|
||||
@ -1145,16 +1146,20 @@ class WechatChannel(Channel):
|
||||
"full_url": full_url,
|
||||
}
|
||||
|
||||
def _stage_downloaded_file(self, filename: str, content: bytes) -> Path | None:
|
||||
async def _stage_downloaded_file(self, filename: str, content: bytes) -> Path | None:
|
||||
download_dir = self._download_dir()
|
||||
if download_dir is None:
|
||||
return None
|
||||
try:
|
||||
download_dir.mkdir(parents=True, exist_ok=True)
|
||||
return publish_upload_bytes(download_dir, filename, content)
|
||||
await asyncio.to_thread(download_dir.mkdir, parents=True, exist_ok=True)
|
||||
publication = await publish_upload_bytes_leased_async(download_dir, filename, content)
|
||||
except (OSError, ValueError):
|
||||
logger.exception("[WeChat] failed to persist inbound media file %s", filename)
|
||||
return None
|
||||
try:
|
||||
return publication.path
|
||||
finally:
|
||||
await release_published_upload_async(publication)
|
||||
|
||||
@staticmethod
|
||||
def _decode_base64_aes_key(value: str) -> bytes | None:
|
||||
|
||||
@ -235,8 +235,28 @@ async def _run_file_io_commit(function, *args):
|
||||
return task.result()
|
||||
|
||||
|
||||
async def _publish_staged_upload_cancellation_safe(staged: StagedUpload, filename: str) -> PublishedUpload:
|
||||
publish_task = asyncio.create_task(run_upload_lease_io(publish_staged_upload_leased, staged, filename))
|
||||
async def _run_upload_lease_io_cancellation_safe(function, *args, **kwargs):
|
||||
task = asyncio.create_task(run_upload_lease_io(function, *args, **kwargs))
|
||||
cancelled = await wait_for_task_completion(task)
|
||||
result = task.result()
|
||||
if cancelled:
|
||||
raise asyncio.CancelledError
|
||||
return result
|
||||
|
||||
|
||||
async def _publish_staged_upload_cancellation_safe(
|
||||
staged: StagedUpload,
|
||||
filename: str,
|
||||
reserved_coordination_keys: set[str] | None = None,
|
||||
) -> PublishedUpload:
|
||||
publish_task = asyncio.create_task(
|
||||
run_upload_lease_io(
|
||||
publish_staged_upload_leased,
|
||||
staged,
|
||||
filename,
|
||||
reserved_coordination_keys=reserved_coordination_keys,
|
||||
)
|
||||
)
|
||||
try:
|
||||
return await asyncio.shield(publish_task)
|
||||
except asyncio.CancelledError:
|
||||
@ -309,6 +329,7 @@ async def _write_upload_file_with_limits(
|
||||
max_single_file_size: int,
|
||||
max_total_size: int,
|
||||
total_size: int,
|
||||
reserved_coordination_keys: set[str] | None = None,
|
||||
) -> tuple[PublishedUpload, int, int]:
|
||||
file_size = 0
|
||||
upload_temp: StagedUpload | None = None
|
||||
@ -323,7 +344,11 @@ async def _write_upload_file_with_limits(
|
||||
raise HTTPException(status_code=413, detail="Total upload size too large")
|
||||
await run_file_io(upload_temp.handle.write, chunk)
|
||||
|
||||
publication = await _publish_staged_upload_cancellation_safe(upload_temp, display_filename)
|
||||
publication = await _publish_staged_upload_cancellation_safe(
|
||||
upload_temp,
|
||||
display_filename,
|
||||
reserved_coordination_keys,
|
||||
)
|
||||
upload_temp = None
|
||||
except BaseException:
|
||||
if upload_temp is not None:
|
||||
@ -375,6 +400,7 @@ async def upload_files(
|
||||
sandbox_sync_targets = []
|
||||
attempted_sandbox_paths: list[str] = []
|
||||
skipped_files = []
|
||||
reserved_coordination_keys: set[str] = set()
|
||||
total_size = 0
|
||||
sandbox_provider = await asyncio.to_thread(get_sandbox_provider)
|
||||
sync_to_sandbox = not _uses_thread_data_mounts(sandbox_provider)
|
||||
@ -406,6 +432,7 @@ async def upload_files(
|
||||
max_single_file_size=limits.max_file_size,
|
||||
max_total_size=limits.max_total_size,
|
||||
total_size=total_size,
|
||||
reserved_coordination_keys=reserved_coordination_keys,
|
||||
)
|
||||
publications.append(publication)
|
||||
file_path = publication.path
|
||||
@ -533,7 +560,12 @@ async def list_uploaded_files(thread_id: ThreadId, request: Request) -> UploadLi
|
||||
async def delete_uploaded_file(thread_id: ThreadId, filename: str, request: Request) -> dict:
|
||||
"""Delete a file from a thread's uploads directory."""
|
||||
try:
|
||||
return await run_file_io(_delete_uploaded_file_for_thread, thread_id, filename, get_effective_user_id())
|
||||
return await _run_upload_lease_io_cancellation_safe(
|
||||
_delete_uploaded_file_for_thread,
|
||||
thread_id,
|
||||
filename,
|
||||
get_effective_user_id(),
|
||||
)
|
||||
except FileNotFoundError:
|
||||
raise HTTPException(status_code=404, detail=f"File not found: {filename}")
|
||||
except PathTraversalError:
|
||||
|
||||
@ -1521,6 +1521,7 @@ class DeerFlowClient:
|
||||
uploaded_files: list[dict] = []
|
||||
publications: list[PublishedUpload] = []
|
||||
sandbox_receipts: list[SandboxSyncReceipt] = []
|
||||
reserved_coordination_keys: set[str] = set()
|
||||
|
||||
conversion_pool = None
|
||||
if has_convertible_file:
|
||||
@ -1545,7 +1546,12 @@ class DeerFlowClient:
|
||||
|
||||
try:
|
||||
for src_path in resolved_files:
|
||||
publication = publish_upload_copy_leased(uploads_dir, src_path.name, src_path)
|
||||
publication = publish_upload_copy_leased(
|
||||
uploads_dir,
|
||||
src_path.name,
|
||||
src_path,
|
||||
reserved_coordination_keys=reserved_coordination_keys,
|
||||
)
|
||||
publications.append(publication)
|
||||
dest = publication.path
|
||||
dest_name = dest.name
|
||||
|
||||
@ -13,7 +13,7 @@ from .layout import (
|
||||
existing_conversion_path_for_upload,
|
||||
validate_conversion_dir,
|
||||
)
|
||||
from .lease import UploadIdentity, UploadNameLease
|
||||
from .lease import UploadIdentity, UploadNameLease, portable_name_coordination_key
|
||||
from .manager import (
|
||||
UPLOAD_STAGING_PREFIX,
|
||||
UPLOAD_STAGING_SUFFIX,
|
||||
@ -75,6 +75,7 @@ __all__ = [
|
||||
"AtomicUploadPublishError",
|
||||
"UploadIdentity",
|
||||
"UploadNameLease",
|
||||
"portable_name_coordination_key",
|
||||
"StagedUpload",
|
||||
"PublishedUpload",
|
||||
"UPLOAD_STAGING_PREFIX",
|
||||
|
||||
@ -36,7 +36,7 @@ class _ThreadLockEntry:
|
||||
_THREAD_LOCKS: dict[tuple[int, int, str], _ThreadLockEntry] = {}
|
||||
|
||||
|
||||
def _portable_name_coordination_key(filename: str) -> str:
|
||||
def portable_name_coordination_key(filename: str) -> str:
|
||||
"""Collapse portable filesystem case and Unicode aliases for lease locking."""
|
||||
return unicodedata.normalize("NFC", filename).casefold()
|
||||
|
||||
@ -273,7 +273,7 @@ class UploadNameLease:
|
||||
raise UnsafeUploadPathError("Upload lease filename is too long")
|
||||
|
||||
uploads_dir = Path(uploads_dir)
|
||||
coordination_key = _portable_name_coordination_key(filename)
|
||||
coordination_key = portable_name_coordination_key(filename)
|
||||
digest = hashlib.sha256(coordination_key.encode("utf-8")).hexdigest()
|
||||
thread_lock_key, thread_lock_entry = _acquire_thread_lock(uploads_dir, coordination_key)
|
||||
lock_file: BinaryIO | None = None
|
||||
|
||||
@ -25,7 +25,7 @@ from deerflow.uploads.layout import (
|
||||
existing_conversion_path_for_upload,
|
||||
upload_virtual_path,
|
||||
)
|
||||
from deerflow.uploads.lease import UploadIdentity, UploadNameLease, UploadStageLease
|
||||
from deerflow.uploads.lease import UploadIdentity, UploadNameLease, UploadStageLease, portable_name_coordination_key
|
||||
from deerflow.utils.thread_id import validate_thread_id
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
@ -260,7 +260,12 @@ def _release_lease_without_masking(lease: UploadNameLease) -> None:
|
||||
logger.warning("Failed to release upload name lease: %s", lease.lock_path, exc_info=True)
|
||||
|
||||
|
||||
def publish_staged_upload_leased(staged: StagedUpload, preferred_filename: str) -> PublishedUpload:
|
||||
def publish_staged_upload_leased(
|
||||
staged: StagedUpload,
|
||||
preferred_filename: str,
|
||||
*,
|
||||
reserved_coordination_keys: set[str] | None = None,
|
||||
) -> PublishedUpload:
|
||||
"""Atomically publish a staging file and retain its actual-name lease."""
|
||||
safe_name = normalize_filename(preferred_filename)
|
||||
if not staged.handle.closed:
|
||||
@ -268,6 +273,9 @@ def publish_staged_upload_leased(staged: StagedUpload, preferred_filename: str)
|
||||
_validate_staged_upload(staged)
|
||||
staged_identity = UploadIdentity.from_path(staged.path)
|
||||
for candidate_name in _filename_candidates(safe_name):
|
||||
coordination_key = portable_name_coordination_key(candidate_name)
|
||||
if reserved_coordination_keys is not None and coordination_key in reserved_coordination_keys:
|
||||
continue
|
||||
candidate = staged.base_dir / candidate_name
|
||||
try:
|
||||
os.lstat(candidate)
|
||||
@ -311,7 +319,10 @@ def publish_staged_upload_leased(staged: StagedUpload, preferred_filename: str)
|
||||
if not stat.S_ISREG(candidate_stat.st_mode) or candidate_stat.st_nlink != 1 or not staged_identity.matches(candidate):
|
||||
_rollback_link_without_masking(candidate, staged_identity)
|
||||
raise AtomicUploadPublishError("Published upload did not become an exclusive regular file")
|
||||
return PublishedUpload(path=candidate, identity=staged_identity, lease=lease)
|
||||
publication = PublishedUpload(path=candidate, identity=staged_identity, lease=lease)
|
||||
if reserved_coordination_keys is not None:
|
||||
reserved_coordination_keys.add(coordination_key)
|
||||
return publication
|
||||
except BaseException:
|
||||
if linked:
|
||||
_rollback_link_without_masking(candidate, staged_identity)
|
||||
@ -356,14 +367,24 @@ def publish_upload_bytes(base_dir: Path, preferred_filename: str, data: bytes) -
|
||||
publication.release()
|
||||
|
||||
|
||||
def publish_upload_copy_leased(base_dir: Path, preferred_filename: str, source_path: Path) -> PublishedUpload:
|
||||
def publish_upload_copy_leased(
|
||||
base_dir: Path,
|
||||
preferred_filename: str,
|
||||
source_path: Path,
|
||||
*,
|
||||
reserved_coordination_keys: set[str] | None = None,
|
||||
) -> PublishedUpload:
|
||||
"""Copy a source into staging, publish it, and retain the actual-name lease."""
|
||||
safe_name = normalize_filename(preferred_filename)
|
||||
staged = create_upload_staging_file(base_dir)
|
||||
try:
|
||||
with Path(source_path).open("rb") as source:
|
||||
shutil.copyfileobj(source, staged.handle)
|
||||
return publish_staged_upload_leased(staged, safe_name)
|
||||
return publish_staged_upload_leased(
|
||||
staged,
|
||||
safe_name,
|
||||
reserved_coordination_keys=reserved_coordination_keys,
|
||||
)
|
||||
except BaseException:
|
||||
_abort_staged_upload_without_masking(staged)
|
||||
raise
|
||||
@ -530,6 +551,22 @@ def list_files_in_dir(directory: Path) -> dict:
|
||||
return {"files": files, "count": len(files)}
|
||||
|
||||
|
||||
def _find_upload_path_by_identity(base_dir: Path, identity: UploadIdentity) -> Path:
|
||||
"""Return the directory entry that actually names *identity*."""
|
||||
with os.scandir(base_dir) as entries:
|
||||
for entry in entries:
|
||||
try:
|
||||
entry_stat = entry.stat(follow_symlinks=False)
|
||||
except FileNotFoundError:
|
||||
continue
|
||||
if stat.S_ISREG(entry_stat.st_mode) and (entry_stat.st_dev, entry_stat.st_ino) == (
|
||||
identity.device,
|
||||
identity.inode,
|
||||
):
|
||||
return Path(entry.path)
|
||||
raise UnsafeUploadPathError("Upload directory entry changed during deletion")
|
||||
|
||||
|
||||
def delete_file_safe(base_dir: Path, filename: str) -> dict:
|
||||
"""Delete a primary upload and only its exact owned conversion.
|
||||
|
||||
@ -562,10 +599,14 @@ def delete_file_safe(base_dir: Path, filename: str) -> dict:
|
||||
if not stat.S_ISREG(file_stat.st_mode) or file_stat.st_nlink != 1:
|
||||
raise UnsafeUploadPathError(f"Unsafe upload file: {safe_name}")
|
||||
|
||||
owned_conversion = existing_conversion_path_for_upload(file_path)
|
||||
identity = UploadIdentity(device=file_stat.st_dev, inode=file_stat.st_ino)
|
||||
actual_file_path = _find_upload_path_by_identity(base_dir, identity)
|
||||
owned_conversion = existing_conversion_path_for_upload(actual_file_path)
|
||||
if owned_conversion is not None:
|
||||
owned_conversion.unlink(missing_ok=True)
|
||||
file_path.unlink()
|
||||
if not identity.matches(actual_file_path):
|
||||
raise UnsafeUploadPathError("Upload changed during deletion")
|
||||
actual_file_path.unlink()
|
||||
finally:
|
||||
lease.release()
|
||||
|
||||
|
||||
@ -2294,6 +2294,25 @@ class TestUploads:
|
||||
assert (uploads_dir / "same.txt").read_bytes() == b"first"
|
||||
assert (uploads_dir / "same_1.txt").read_bytes() == b"second"
|
||||
|
||||
def test_upload_files_renames_portable_aliases_within_one_batch(self, client, tmp_path):
|
||||
uploads_dir = tmp_path / "user-data" / "uploads"
|
||||
uploads_dir.mkdir(parents=True)
|
||||
first_dir = tmp_path / "first"
|
||||
second_dir = tmp_path / "second"
|
||||
first_dir.mkdir()
|
||||
second_dir.mkdir()
|
||||
first = first_dir / "Report.txt"
|
||||
second = second_dir / "report.txt"
|
||||
first.write_bytes(b"first")
|
||||
second.write_bytes(b"second")
|
||||
|
||||
with patch("deerflow.client.ensure_uploads_dir", return_value=uploads_dir):
|
||||
result = client.upload_files("thread-aliases", [first, second])
|
||||
|
||||
assert [file["filename"] for file in result["files"]] == ["Report.txt", "report_1.txt"]
|
||||
assert (uploads_dir / "Report.txt").read_bytes() == b"first"
|
||||
assert (uploads_dir / "report_1.txt").read_bytes() == b"second"
|
||||
|
||||
def test_concurrent_client_uploads_preserve_all_payloads(self, client, tmp_path):
|
||||
uploads_dir = tmp_path / "user-data" / "uploads"
|
||||
uploads_dir.mkdir(parents=True)
|
||||
|
||||
@ -214,6 +214,42 @@ def test_multibyte_long_conversion_name_is_utf8_safe():
|
||||
assert converted.endswith(".md")
|
||||
|
||||
|
||||
def test_delete_long_upload_through_case_alias_removes_actual_owned_conversion(tmp_path):
|
||||
uploads = tmp_path / "user-data" / "uploads"
|
||||
uploads.mkdir(parents=True)
|
||||
filename = "A" * 251 + ".pdf"
|
||||
alias = filename.lower()
|
||||
upload = publish_upload_bytes(uploads, filename, b"PDF")
|
||||
if not (uploads / alias).exists():
|
||||
pytest.skip("filesystem is case-sensitive")
|
||||
conversion = conversion_path_for_upload(upload)
|
||||
conversion.parent.mkdir(parents=True, exist_ok=True)
|
||||
conversion.write_text("converted", encoding="utf-8")
|
||||
|
||||
delete_file_safe(uploads, alias)
|
||||
|
||||
assert not upload.exists()
|
||||
assert not conversion.exists()
|
||||
|
||||
|
||||
def test_delete_long_upload_through_unicode_alias_removes_actual_owned_conversion(tmp_path):
|
||||
uploads = tmp_path / "user-data" / "uploads"
|
||||
uploads.mkdir(parents=True)
|
||||
filename = "e\u0301" * 83 + ".pdf"
|
||||
alias = "\u00e9" * 83 + ".pdf"
|
||||
upload = publish_upload_bytes(uploads, filename, b"PDF")
|
||||
if not (uploads / alias).exists():
|
||||
pytest.skip("filesystem does not resolve Unicode normalization aliases")
|
||||
conversion = conversion_path_for_upload(upload)
|
||||
conversion.parent.mkdir(parents=True, exist_ok=True)
|
||||
conversion.write_text("converted", encoding="utf-8")
|
||||
|
||||
delete_file_safe(uploads, alias)
|
||||
|
||||
assert not upload.exists()
|
||||
assert not conversion.exists()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_conversion_uses_owned_full_filename_target(tmp_path):
|
||||
uploads = tmp_path / "user-data" / "uploads"
|
||||
|
||||
@ -32,6 +32,7 @@ from deerflow.uploads.manager import (
|
||||
list_files_in_dir,
|
||||
normalize_filename,
|
||||
publish_staged_upload,
|
||||
publish_staged_upload_leased,
|
||||
publish_upload_bytes,
|
||||
publish_upload_bytes_leased,
|
||||
publish_upload_copy,
|
||||
@ -153,6 +154,40 @@ def test_portable_filesystem_aliases_share_one_generation_lease(tmp_path, first_
|
||||
alias.release()
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("first_name", "alias_name", "expected_alias_name"),
|
||||
[
|
||||
("Report.pdf", "report.pdf", "report_1.pdf"),
|
||||
("caf\u00e9.pdf", "cafe\u0301.pdf", "cafe\u0301_1.pdf"),
|
||||
],
|
||||
)
|
||||
def test_batch_portable_alias_reservation_chooses_distinct_name(
|
||||
tmp_path,
|
||||
first_name,
|
||||
alias_name,
|
||||
expected_alias_name,
|
||||
):
|
||||
reserved_keys: set[str] = set()
|
||||
publications = []
|
||||
try:
|
||||
for filename, payload in [(first_name, b"first"), (alias_name, b"second")]:
|
||||
staged = create_upload_staging_file(tmp_path)
|
||||
staged.handle.write(payload)
|
||||
publications.append(
|
||||
publish_staged_upload_leased(
|
||||
staged,
|
||||
filename,
|
||||
reserved_coordination_keys=reserved_keys,
|
||||
)
|
||||
)
|
||||
finally:
|
||||
for publication in reversed(publications):
|
||||
publication.release()
|
||||
|
||||
assert [publication.path.name for publication in publications] == [first_name, expected_alias_name]
|
||||
assert {publication.path.read_bytes() for publication in publications} == {b"first", b"second"}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_release_commit_delays_and_swallows_new_cancellation(tmp_path, monkeypatch):
|
||||
publication = publish_upload_bytes_leased(tmp_path, "report.pdf", b"payload")
|
||||
|
||||
@ -420,6 +420,31 @@ async def test_waiting_publication_cannot_starve_lease_release(tmp_path, monkeyp
|
||||
assert second.path.name == "report.pdf"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_waiting_delete_cannot_starve_general_io_lease_release(tmp_path, monkeypatch):
|
||||
import deerflow.utils.file_io as file_io_module
|
||||
|
||||
publication = publish_upload_bytes_leased(tmp_path, "notes.txt", b"payload")
|
||||
single_worker = ThreadPoolExecutor(max_workers=1)
|
||||
monkeypatch.setattr(file_io_module, "_FILE_IO_EXECUTOR", single_worker)
|
||||
|
||||
with patch.object(uploads, "get_uploads_dir", return_value=tmp_path):
|
||||
deletion = asyncio.create_task(call_unwrapped(uploads.delete_uploaded_file, "thread-delete", "notes.txt", request=MagicMock()))
|
||||
await asyncio.sleep(0.05)
|
||||
release = asyncio.create_task(uploads._run_file_io_commit(publication.release))
|
||||
try:
|
||||
await asyncio.wait_for(asyncio.shield(release), timeout=0.2)
|
||||
await asyncio.wait_for(deletion, timeout=2)
|
||||
finally:
|
||||
if publication.is_active:
|
||||
publication.release()
|
||||
await release
|
||||
await deletion
|
||||
single_worker.shutdown(wait=True)
|
||||
|
||||
assert not (tmp_path / "notes.txt").exists()
|
||||
|
||||
|
||||
def test_upload_files_writes_thread_storage_and_skips_local_sandbox_sync(tmp_path):
|
||||
thread_uploads_dir = tmp_path / "uploads"
|
||||
thread_uploads_dir.mkdir(parents=True)
|
||||
@ -985,6 +1010,33 @@ def test_upload_files_adjusts_read_permissions_for_mounted_non_local_sandbox(tmp
|
||||
assert called_path.name == "notes.txt"
|
||||
|
||||
|
||||
def test_upload_files_renames_portable_aliases_within_one_batch(tmp_path):
|
||||
thread_uploads_dir = tmp_path / "uploads"
|
||||
thread_uploads_dir.mkdir(parents=True)
|
||||
provider = _mounted_provider()
|
||||
|
||||
with (
|
||||
patch.object(uploads, "ensure_uploads_dir", return_value=thread_uploads_dir),
|
||||
patch.object(uploads, "get_sandbox_provider", return_value=provider),
|
||||
):
|
||||
result = asyncio.run(
|
||||
call_unwrapped(
|
||||
uploads.upload_files,
|
||||
"thread-aliases",
|
||||
request=MagicMock(),
|
||||
files=[
|
||||
UploadFile(filename="Report.txt", file=BytesIO(b"first")),
|
||||
UploadFile(filename="report.txt", file=BytesIO(b"second")),
|
||||
],
|
||||
config=SimpleNamespace(),
|
||||
)
|
||||
)
|
||||
|
||||
assert [file.filename for file in result.files] == ["Report.txt", "report_1.txt"]
|
||||
assert (thread_uploads_dir / "Report.txt").read_bytes() == b"first"
|
||||
assert (thread_uploads_dir / "report_1.txt").read_bytes() == b"second"
|
||||
|
||||
|
||||
def test_upload_files_rejects_dotdot_and_dot_filenames(tmp_path):
|
||||
thread_uploads_dir = tmp_path / "uploads"
|
||||
thread_uploads_dir.mkdir(parents=True)
|
||||
|
||||
@ -6,12 +6,14 @@ import asyncio
|
||||
import base64
|
||||
import json
|
||||
import logging
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
import threading
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
from unittest import mock
|
||||
from unittest.mock import AsyncMock
|
||||
|
||||
import pytest
|
||||
|
||||
from app.channels.message_bus import InboundMessageType, MessageBus, OutboundMessage
|
||||
|
||||
|
||||
@ -1024,8 +1026,10 @@ def test_stage_downloaded_file_preserves_concurrent_same_name_payloads(tmp_path:
|
||||
channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token", "state_dir": str(tmp_path)})
|
||||
payloads = [f"payload-{index}".encode() for index in range(8)]
|
||||
|
||||
with ThreadPoolExecutor(max_workers=len(payloads)) as pool:
|
||||
paths = list(pool.map(lambda payload: channel._stage_downloaded_file("report.pdf", payload), payloads))
|
||||
async def stage_all():
|
||||
return await asyncio.gather(*(channel._stage_downloaded_file("report.pdf", payload) for payload in payloads))
|
||||
|
||||
paths = _run(stage_all())
|
||||
|
||||
assert all(path is not None for path in paths)
|
||||
assert {path.name for path in paths if path is not None} == {
|
||||
@ -1051,7 +1055,7 @@ def test_stage_downloaded_file_renames_around_planted_symlink(tmp_path: Path):
|
||||
victim.write_bytes(b"victim")
|
||||
(download_dir / "report.pdf").symlink_to(victim)
|
||||
|
||||
stored = channel._stage_downloaded_file("report.pdf", b"attachment")
|
||||
stored = _run(channel._stage_downloaded_file("report.pdf", b"attachment"))
|
||||
|
||||
assert stored == download_dir / "report_1.pdf"
|
||||
assert stored.read_bytes() == b"attachment"
|
||||
@ -1070,15 +1074,49 @@ def test_wechat_invalid_platform_filename_uses_safe_fallback(tmp_path: Path):
|
||||
index=0,
|
||||
)
|
||||
assert safe == "wechat-file-m1-0.bin"
|
||||
assert channel._stage_downloaded_file(safe, b"payload") is not None
|
||||
assert _run(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
|
||||
with mock.patch(
|
||||
"app.channels.wechat.publish_upload_bytes_leased_async",
|
||||
new=AsyncMock(side_effect=ValueError("invalid filename")),
|
||||
):
|
||||
assert _run(channel._stage_downloaded_file("report.pdf", b"payload")) is None
|
||||
|
||||
|
||||
def test_cancelled_stage_downloaded_file_rolls_back_late_publication(tmp_path: Path):
|
||||
import deerflow.uploads.async_helpers as async_helpers
|
||||
from app.channels.wechat import WechatChannel
|
||||
|
||||
channel = WechatChannel(bus=MessageBus(), config={"bot_token": "test-token", "state_dir": str(tmp_path)})
|
||||
publish_started = threading.Event()
|
||||
allow_publish = threading.Event()
|
||||
real_publish = async_helpers.publish_upload_bytes_leased
|
||||
|
||||
def paused_publish(*args, **kwargs):
|
||||
publish_started.set()
|
||||
assert allow_publish.wait(5)
|
||||
return real_publish(*args, **kwargs)
|
||||
|
||||
async def exercise():
|
||||
with mock.patch("deerflow.uploads.async_helpers.publish_upload_bytes_leased", side_effect=paused_publish):
|
||||
task = asyncio.create_task(channel._stage_downloaded_file("report.pdf", b"payload"))
|
||||
assert await asyncio.to_thread(publish_started.wait, 5)
|
||||
task.cancel()
|
||||
await asyncio.sleep(0.05)
|
||||
assert not task.done()
|
||||
allow_publish.set()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
_run(exercise())
|
||||
download_dir = tmp_path / channel.DEFAULT_IMAGE_DOWNLOAD_DIRNAME
|
||||
assert not (download_dir / "report.pdf").exists()
|
||||
assert not list(download_dir.glob("report*.pdf"))
|
||||
|
||||
|
||||
def test_handle_update_downloads_inbound_file_with_media_aeskey_hex(monkeypatch, tmp_path: Path):
|
||||
|
||||
Loading…
x
Reference in New Issue
Block a user