From bba880b6763bab774675bb027b6fc1f185633162 Mon Sep 17 00:00:00 2001 From: Sun Shihao <75820689+sun-sh@users.noreply.github.com> Date: Sat, 24 Oct 2026 10:47:40 +0800 Subject: [PATCH] fix(projects): drain shelf staging on cancellation (#7360) --- README.md | 2 + .../harness/deerflow/projects/documents.py | 19 +- .../packages/harness/deerflow/utils/AGENTS.md | 4 + ...t_project_document_staging_cancellation.py | 244 ++++++++++++++++++ .../tests/test_project_documents_promotion.py | 2 +- 5 files changed, 265 insertions(+), 6 deletions(-) create mode 100644 backend/tests/test_project_document_staging_cancellation.py diff --git a/README.md b/README.md index 51bac9e8c1..c449574250 100644 --- a/README.md +++ b/README.md @@ -3378,6 +3378,8 @@ from the project page's Documents section: Oversized uploads return `413`; deduplicated uploads return `200`. Staging cleanup failures are logged without replacing these outcomes or a missing, foreign, or archived project's `404` response. + Cancelling a shelf upload discards its partial file after any active disk + operation finishes; waiting for further upload data remains cancellable. - **List** entries with name, size, modified time, and provenance (uploaded vs. saved from a conversation), and preview or download any entry. - **Save to project** from a thread file: the read-only conversation-files diff --git a/backend/packages/harness/deerflow/projects/documents.py b/backend/packages/harness/deerflow/projects/documents.py index ab29684ed8..553978e01e 100644 --- a/backend/packages/harness/deerflow/projects/documents.py +++ b/backend/packages/harness/deerflow/projects/documents.py @@ -165,10 +165,14 @@ async def stage_document_bytes(paths: Paths, *, user_id: str, project_id: str, c callers reject a missing, foreign, or archived project first. """ staging_path = _new_staging_path(paths, user_id=user_id, project_id=project_id) - handle = await run_file_io(_open_staging, staging_path.parent, staging_path) + handle = None digest = hashlib.sha256() size = 0 + async def prepare() -> None: + nonlocal handle + handle = await run_file_io(_open_staging, staging_path.parent, staging_path) + async def _aiter() -> AsyncIterator[bytes]: if isinstance(chunks, list): for chunk in chunks: @@ -178,15 +182,20 @@ async def stage_document_bytes(paths: Paths, *, user_id: str, project_id: str, c yield chunk try: + # Retain the opened handle even when cancellation arrives during its + # worker-to-caller handoff. Only file operations are drained; waiting + # for the next incoming chunk must remain cancellable. + await await_drained(prepare()) async for chunk in _aiter(): size += len(chunk) if size > max_bytes: raise ShelfUploadTooLargeError(f"File exceeds the {max_bytes}-byte limit") digest.update(chunk) - await run_file_io(handle.write, chunk) - await run_file_io(handle.close) - except Exception: - await run_file_io(_close_and_unlink, handle, staging_path) + await await_drained(run_file_io(handle.write, chunk)) + await await_drained(run_file_io(handle.close)) + except BaseException: + if handle is not None: + await await_drained(run_file_io(_close_and_unlink, handle, staging_path)) raise return StagedDocument(staging_path=staging_path, sha256=digest.hexdigest(), size_bytes=size) diff --git a/backend/packages/harness/deerflow/utils/AGENTS.md b/backend/packages/harness/deerflow/utils/AGENTS.md index 326ba0e701..3d06189f1f 100644 --- a/backend/packages/harness/deerflow/utils/AGENTS.md +++ b/backend/packages/harness/deerflow/utils/AGENTS.md @@ -18,6 +18,10 @@ temporary; pass a raw descriptor only until `fdopen` succeeds, or an explicitly transferred duplicate descriptor owned by the cleanup caller. Shelf discards must preserve both a rejected project's 404 and a dedup hit's 200 response. Callers retain their own serialization, permissions and publication policy. +Shelf `stage_document_bytes` drains preparation, writes, close and failure +cleanup before propagating cancellation. Retain the opened handle across its +handoff, but leave waiting for incoming chunks cancellable; regressions in +`tests/test_project_document_staging_cancellation.py` cover both shelf entry points. Local Linux/macOS text revisions opt into a short `temporary_prefix` and `preserve_metadata_from` (a caller-owned descriptor). Staging starts at `0600`; `file_revision_metadata` transfers ownership, mode, ACLs and xattrs before data. diff --git a/backend/tests/test_project_document_staging_cancellation.py b/backend/tests/test_project_document_staging_cancellation.py new file mode 100644 index 0000000000..91e7a2f036 --- /dev/null +++ b/backend/tests/test_project_document_staging_cancellation.py @@ -0,0 +1,244 @@ +"""Cancelled shelf staging retains filesystem ownership until cleanup ends.""" + +from __future__ import annotations + +import asyncio +import hashlib +import threading +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from _router_auth_helpers import call_unwrapped +from fastapi import Response + +from deerflow.config.paths import Paths +from deerflow.projects import documents + + +@pytest.mark.asyncio +async def test_cancelled_chunk_stream_closes_and_removes_staging(tmp_path, monkeypatch): + paths = Paths(tmp_path) + waiting = asyncio.Event() + never = asyncio.Event() + handles = [] + real_open = documents._open_staging + + def observe_open(*args): + handle = real_open(*args) + handles.append(handle) + return handle + + async def chunks(): + yield b"first chunk\n" + waiting.set() + await never.wait() + + monkeypatch.setattr(documents, "_open_staging", observe_open) + task = asyncio.create_task(documents.stage_document_bytes(paths, user_id="u1", project_id="p1", chunks=chunks(), max_bytes=1024)) + try: + await asyncio.wait_for(waiting.wait(), 10) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 10) + assert handles[0].closed, "cancelled staging retained its file descriptor" + assert not list((paths.project_documents_dir("u1", "p1") / ".staging").iterdir()) + assert not never.is_set(), "cancellation must not drain a stalled incoming stream" + finally: + task.cancel() + await asyncio.gather(task, return_exceptions=True) + for handle in handles: + handle.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["open", "write", "close", "cleanup"]) +@pytest.mark.parametrize("worker_fails", [False, True], ids=["success", "failure"]) +async def test_cancelled_staging_drains_file_operation(tmp_path, monkeypatch, operation, worker_fails): + paths = Paths(tmp_path) + loop = asyncio.get_running_loop() + loop_thread = threading.get_ident() + started = asyncio.Event() + finished = threading.Event() + release = threading.Event() + handles = [] + worker_threads = [] + cleanup_before_worker_finish = [] + real_open = documents._open_staging + real_cleanup = documents._close_and_unlink + real_file_io = documents.run_file_io + + def gated(call): + worker_threads.append(threading.get_ident()) + loop.call_soon_threadsafe(started.set) + try: + if threading.get_ident() != loop_thread: + assert release.wait(10), "file operation was not released" + if worker_fails and operation != "cleanup": + raise OSError("synthetic staging worker failure") + return call() + finally: + finished.set() + + def observe_open(*args): + def create(): + handle = real_open(*args) + handles.append(handle) + return handle + + handle = gated(create) if operation == "open" else create() + if operation in {"write", "close"}: + real_operation = getattr(handle, operation) + first = True + + def observed_operation(*args): + nonlocal first + if first: + first = False + return gated(lambda: real_operation(*args)) + return real_operation(*args) + + setattr(handle, operation, observed_operation) + elif operation == "cleanup" and worker_fails: + real_close = handle.close + + def close_with_error(): + real_close() + raise OSError("staging close failed after closing") + + handle.close = close_with_error + return handle + + def observe_cleanup(*args): + if operation == "cleanup": + # Hold cleanup before the real close/unlink. The failure case + # raises from close only after releasing its real descriptor. + gated(lambda: None) + return real_cleanup(*args) + + async def observe_file_io(func, *args, **kwargs): + if func is observe_cleanup and operation in {"write", "close"}: + cleanup_before_worker_finish.append(not finished.is_set()) + return await real_file_io(func, *args, **kwargs) + + monkeypatch.setattr(documents, "_open_staging", observe_open) + monkeypatch.setattr(documents, "_close_and_unlink", observe_cleanup) + monkeypatch.setattr(documents, "run_file_io", observe_file_io) + limit = 1 if operation == "cleanup" else 1024 + task = asyncio.create_task(documents.stage_document_bytes(paths, user_id="u1", project_id="p1", chunks=[b"payload\n"], max_bytes=limit)) + try: + await asyncio.wait_for(started.wait(), 10) + assert len(worker_threads) == 1 + assert worker_threads[0] != loop_thread, "staging blocked the event loop" + task.cancel() + await asyncio.sleep(0) + task.cancel() + for _ in range(10): + await asyncio.sleep(0) + assert not task.done(), "cancellation detached an admitted file operation" + assert not finished.is_set() + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 10) + assert finished.is_set() + assert not any(cleanup_before_worker_finish), "cleanup started before its file operation finished" + assert all(handle.closed for handle in handles) + staging = paths.project_documents_dir("u1", "p1") / ".staging" + assert not staging.exists() or not list(staging.iterdir()) + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + await asyncio.to_thread(finished.wait, 10) + for handle in handles: + if not handle.closed: + handle.close() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("streamed", [False, True], ids=["list", "async-stream"]) +async def test_completed_staging_keeps_exact_bytes_and_hash(tmp_path, streamed): + paths = Paths(tmp_path) + chunks = [b"UTF-8\n", "\u4f60\u597d\r\n".encode("utf-8"), b"\x00\xff"] + + async def stream(): + for chunk in chunks: + yield chunk + + staged = await documents.stage_document_bytes(paths, user_id="u1", project_id="p1", chunks=stream() if streamed else chunks, max_bytes=1024) + content = b"".join(chunks) + assert staged.staging_path.read_bytes() == content + assert staged.size_bytes == len(content) + assert staged.sha256 == hashlib.sha256(content).hexdigest() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("entry", ["upload", "promotion"]) +async def test_shelf_entry_cancellation_reclaims_staging_before_insert(tmp_path, monkeypatch, entry): + """Both route consumers retain the real staging writer and worker pool.""" + from app.gateway.routers import project_documents as router + + paths = Paths(tmp_path) + started = asyncio.Event() + release = threading.Event() + finished = threading.Event() + loop = asyncio.get_running_loop() + handles = [] + real_open = documents._open_staging + insert = AsyncMock() + source = paths.sandbox_outputs_dir("thread-1", user_id="u1") / "source.txt" + source.parent.mkdir(parents=True) + source.write_bytes(b"source bytes\n") + + def observe_open(*args): + handle = real_open(*args) + handles.append(handle) + real_write = handle.write + + def paused_write(chunk): + loop.call_soon_threadsafe(started.set) + try: + assert release.wait(10), "route staging writer was not released" + return real_write(chunk) + finally: + finished.set() + + handle.write = paused_write + return handle + + monkeypatch.setattr(documents, "_open_staging", observe_open) + monkeypatch.setattr(router, "get_paths", lambda: paths) + monkeypatch.setattr(router, "get_effective_user_id", lambda: "u1") + monkeypatch.setattr(router, "_require_active_project", AsyncMock(return_value={"id": "p1", "status": "active"})) + monkeypatch.setattr(router, "get_project_document_repo", lambda request: SimpleNamespace()) + monkeypatch.setattr(router, "get_thread_store", lambda request: SimpleNamespace(get=AsyncMock(return_value={"thread_id": "thread-1"}))) + monkeypatch.setattr(router, "add_staged_document", insert) + + if entry == "upload": + file = SimpleNamespace(filename="source.txt", read=AsyncMock(side_effect=[b"source bytes\n", b""])) + call = call_unwrapped(router.upload_project_document, "p1", request=SimpleNamespace(), response=Response(), file=file, name=None, config=SimpleNamespace(uploads={})) + else: + body = router.PromoteThreadFileRequest(thread_id="thread-1", kind="output", name="source.txt") + call = call_unwrapped(router.promote_thread_file_to_shelf, "p1", body=body, request=SimpleNamespace(), response=Response(), config=SimpleNamespace(uploads={})) + task = asyncio.create_task(call) + try: + await asyncio.wait_for(started.wait(), 10) + task.cancel() + for _ in range(10): + await asyncio.sleep(0) + assert not task.done(), "route abandoned its real staging writer" + release.set() + with pytest.raises(asyncio.CancelledError): + await asyncio.wait_for(task, 10) + assert finished.is_set() + assert handles[0].closed + assert not list((paths.project_documents_dir("u1", "p1") / ".staging").iterdir()) + insert.assert_not_awaited() + assert source.read_bytes() == b"source bytes\n" + finally: + release.set() + task.cancel() + await asyncio.gather(task, return_exceptions=True) + for handle in handles: + if not handle.closed: + handle.close() diff --git a/backend/tests/test_project_documents_promotion.py b/backend/tests/test_project_documents_promotion.py index fc929efddb..0d1a776b4c 100644 --- a/backend/tests/test_project_documents_promotion.py +++ b/backend/tests/test_project_documents_promotion.py @@ -566,7 +566,7 @@ class TestAttach: assert response.status_code == 500 assert response.json() == {"detail": "Failed to upload report.txt."} - assert str(staged_path) in caplog.text + assert str(disk_full) in caplog.text assert not (_thread_uploads("thread-1") / "report.txt").exists() def test_attach_sync_failure_retains_host_file(self, tmp_path):