diff --git a/backend/packages/harness/deerflow/projects/documents.py b/backend/packages/harness/deerflow/projects/documents.py index 553978e01e..246e4524be 100644 --- a/backend/packages/harness/deerflow/projects/documents.py +++ b/backend/packages/harness/deerflow/projects/documents.py @@ -593,17 +593,33 @@ async def read_file_chunks(path: Path, *, chunk_size: int = 65536) -> AsyncItera Shared by the from-thread promote path (thread file → shelf staging) and the attach path (staged shelf copy → thread-upload ingestion). + + Cancellation contract: once an open/read/close has been admitted to the + file-I/O worker, the caller's cancellation is delivered only after that + operation settles (:func:`await_drained`). The opened handle is retained + the moment the worker returns it, so a cancellation racing the result + handoff cannot leak it, and the cleanup close never races an in-flight + read. Waiting for the next chunk (the ``yield`` itself) stays immediately + cancellable. (#7491) """ def _open() -> Any: return open(path, "rb") - handle = await run_file_io(_open) + retained: list[Any] = [] + + async def _open_and_retain() -> Any: + handle = await run_file_io(_open) + retained.append(handle) + return handle + try: - while chunk := await run_file_io(handle.read, chunk_size): + handle = await await_drained(_open_and_retain()) + while chunk := await await_drained(run_file_io(handle.read, chunk_size)): yield chunk finally: - await run_file_io(handle.close) + for open_handle in retained: + await await_drained(run_file_io(open_handle.close)) def _copy_original_under_lock(paths: Paths, staging_path: Path, *, user_id: str, row: dict) -> None: diff --git a/backend/tests/blocking_io/test_read_file_chunks_cancellation.py b/backend/tests/blocking_io/test_read_file_chunks_cancellation.py new file mode 100644 index 0000000000..144ed6a28b --- /dev/null +++ b/backend/tests/blocking_io/test_read_file_chunks_cancellation.py @@ -0,0 +1,192 @@ +"""Cancellation-safety regression anchors for ``read_file_chunks`` (#7491). + +Each test gates the real ``run_file_io`` dispatch with an ``asyncio.Event`` at +the worker boundary (real executor, real temporary files) to reproduce one of +the three unsafe windows: a cancellation racing the open-result handoff, a +close racing an in-flight read, and a repeated cancellation trying to escape a +pending close. The expected contract: once an operation has been admitted to +the worker, the caller's cancellation is delivered only after it settles, the +handle always reaches the cleanup path, and the close never races a read. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable, Coroutine +from typing import Any + +import pytest + +from deerflow.projects.documents import read_file_chunks +from deerflow.utils.file_io import run_file_io as _real_run_file_io + +pytestmark = [ + pytest.mark.asyncio, + pytest.mark.allow_blocking_io, +] + + +class _Gate: + def __init__(self, mode: str) -> None: + self.mode = mode + self.admitted = asyncio.Event() + self.hold = asyncio.Event() + self.released = False + + def release(self) -> None: + self.released = True + self.hold.set() + + +def _gated_offload( + gates: dict[str, _Gate], + on_call, + on_done, +) -> tuple[list[str], Callable[..., Coroutine[Any, Any, Any]]]: + """Wrap ``documents.run_file_io`` call-through with worker-boundary gates. + + ``on_call`` fires when a dispatch enters the wrapper (admission), + ``on_done`` fires after the real worker call settles. + """ + calls: list[str] = [] + + async def gated(func, /, *args, **kwargs): + name = getattr(func, "__name__", "") + calls.append(name) + on_call(name) + gate = gates.get(name) + if gate is None: + result = await _real_run_file_io(func, *args, **kwargs) + elif gate.mode == "before": + gate.admitted.set() + await gate.hold.wait() + result = await _real_run_file_io(func, *args, **kwargs) + else: + result = await _real_run_file_io(func, *args, **kwargs) + gate.admitted.set() + await gate.hold.wait() + on_done(name) + return result + + return calls, gated + + +async def _settle() -> None: + for _ in range(10): + await asyncio.sleep(0) + + +async def test_cancel_racing_open_handoff_still_closes_handle(tmp_path, monkeypatch) -> None: + """Window 1: the open worker returns, then cancellation lands before the + result reaches the coroutine — the handle must still be closed.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"payload") + open_gate = _Gate("after") + state: dict[str, bool] = {"closed": False} + + def on_call(name: str) -> None: ... + def on_done(name: str) -> None: + if name == "close": + state["closed"] = True + + calls, gated = _gated_offload({"_open": open_gate}, on_call, on_done) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + agen = read_file_chunks(path, chunk_size=4) + task = asyncio.create_task(agen.__anext__()) + await asyncio.wait_for(open_gate.admitted.wait(), 5) + task.cancel() + await _settle() + assert "close" not in calls, "cleanup must wait for the open handoff to settle" + + open_gate.release() + with pytest.raises(asyncio.CancelledError): + await task + assert state["closed"], "handle retained at handoff must reach the cleanup close" + + +async def test_cancel_during_read_does_not_race_close(tmp_path, monkeypatch) -> None: + """Window 2: cancellation while a read is admitted must not start the + cleanup close before that read settles.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"abcdefgh") + read_gate = _Gate("before") + state: dict[str, bool] = {"close_raced": False, "read_settled": False} + + def on_call(name: str) -> None: + if name == "close" and not state["read_settled"]: + state["close_raced"] = True + + def on_done(name: str) -> None: + if name == "read" and read_gate.released: + state["read_settled"] = True + + calls, gated = _gated_offload({"read": read_gate}, on_call, on_done) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + agen = read_file_chunks(path, chunk_size=4) + task = asyncio.create_task(agen.__anext__()) + await asyncio.wait_for(read_gate.admitted.wait(), 5) + task.cancel() + await _settle() + assert "close" not in calls, "close must wait for the admitted read to settle" + + read_gate.release() + with pytest.raises(asyncio.CancelledError): + await task + assert state["read_settled"], "the admitted read must run to completion" + assert not state["close_raced"], "close must not precede or overlap the in-flight read" + + +async def test_repeated_cancellation_cannot_escape_pending_close(tmp_path, monkeypatch) -> None: + """Window 3: a cancellation arriving while the cleanup close is admitted + must be delivered only after the close settles — repeated cancellation + included.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"abc") + close_gate = _Gate("before") + state: dict[str, bool] = {"close_completed": False} + + def on_call(name: str) -> None: ... + + def on_done(name: str) -> None: + if name == "close" and close_gate.released: + state["close_completed"] = True + + calls, gated = _gated_offload({"close": close_gate}, on_call, on_done) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + async def consume() -> list[bytes]: + return [chunk async for chunk in read_file_chunks(path, chunk_size=8)] + + task = asyncio.create_task(consume()) + await asyncio.wait_for(close_gate.admitted.wait(), 5) + task.cancel() + await _settle() + task.cancel() + await _settle() + close_gate.release() + with pytest.raises(asyncio.CancelledError): + await task + assert state["close_completed"], "the caller may escape only after the close settles" + + +async def test_read_file_chunks_happy_path_unchanged(tmp_path, monkeypatch) -> None: + """Sanity: full consumption still yields the exact bytes.""" + path = tmp_path / "doc.bin" + path.write_bytes(b"abcdefgh") + calls, gated = _gated_offload({}, lambda name: None, lambda name: None) + monkeypatch.setattr("deerflow.projects.documents.run_file_io", gated) + + chunks = [chunk async for chunk in read_file_chunks(path, chunk_size=4)] + + assert chunks == [b"abcd", b"efgh"] + assert calls[0] == "_open" and calls[-1] == "close" + + +async def test_missing_path_raises_file_not_found(tmp_path) -> None: + """Pins the open-failure path: the error surfaces as-is and the cleanup + skips the close (``retained`` is empty — no handle was ever opened).""" + with pytest.raises(FileNotFoundError): + async for _ in read_file_chunks(tmp_path / "absent.bin", chunk_size=4): + pass