From 79a5af071a91cc5f33bbf6c439a896388e6066dc Mon Sep 17 00:00:00 2001 From: NanPan <111261006+poijygfdyy@users.noreply.github.com> Date: Sat, 26 Sep 2026 17:59:31 +0800 Subject: [PATCH] fix(events): retain JSONL snapshot ownership across cancellation (#5795) * fix(events): retain JSONL snapshot lock through cancellation * test(events): cover cancelled JSONL snapshot ownership * docs(events): record JSONL snapshot ownership invariant * test(events): pin snapshot cancellation failure semantics * docs(runtime): name drained JSONL snapshot tasks * test(events): pin snapshot writer exclusion on disk --------- Co-authored-by: Willem Jiang --- .../harness/deerflow/runtime/AGENTS.md | 6 +- .../deerflow/runtime/events/store/jsonl.py | 46 +++++++----- .../test_jsonl_event_store_cancellation.py | 72 +++++++++++++++++++ 3 files changed, 104 insertions(+), 20 deletions(-) diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index 7d33d21cc2..80b5734b5f 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -93,7 +93,11 @@ contract. Its default implementation walks `list_messages()` backward in `before_seq` cursor, and raises when a full page has no safe progressing `seq`. Memory and database stores use that bounded path; the JSONL store overrides it with one complete thread-log read because each JSONL page would otherwise -rescan every run file. The default and JSONL paths share the public +rescan every run file. That JSONL snapshot task is named `jsonl-snapshot:{thread_id}` for +asyncio task dumps and retains the per-thread lock until its off-thread full-log +read settles even through caller cancellation. It deliberately has no drain +timeout: releasing ownership while the worker can still read files would let a +writer enter the supposedly stable snapshot. The default and JSONL paths share the public `normalize_message_ids()` and `match_ai_message_run_id()` helpers from `events/store/base.py`. Database owner filtering is inherited on every page. diff --git a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py index 374012c1dc..76c3db829f 100644 --- a/backend/packages/harness/deerflow/runtime/events/store/jsonl.py +++ b/backend/packages/harness/deerflow/runtime/events/store/jsonl.py @@ -61,6 +61,27 @@ class JsonlRunEventStore(RunEventStore): self._write_locks[thread_id] = lock return lock + @staticmethod + async def _await_owned_task[T](task: asyncio.Task[T]) -> T: + """Keep owned work attached until it settles, then propagate cancellation.""" + cancellation: asyncio.CancelledError | None = None + while not task.done(): + try: + await asyncio.shield(task) + except asyncio.CancelledError as exc: + if cancellation is None: + cancellation = exc + except Exception: + # Retrieve the failure below after the owned work has settled. + break + if cancellation is not None: + try: + task.result() + except Exception as exc: + raise cancellation from exc + raise cancellation + return task.result() + async def _run_mutation[T](self, thread_id: str, operation: Callable[[], Coroutine[Any, Any, T]]) -> T: """Drain an admitted mutation before propagating caller cancellation. @@ -71,24 +92,7 @@ class JsonlRunEventStore(RunEventStore): """ async with self._get_write_lock(thread_id): task = asyncio.create_task(operation(), name=f"jsonl-mutation:{thread_id}") - cancellation: asyncio.CancelledError | None = None - while not task.done(): - try: - await asyncio.shield(task) - except asyncio.CancelledError as exc: - if cancellation is None: - cancellation = exc - except Exception: - # Retrieve the failure below, after preserving any earlier - # cancellation. The operation has already finished rollback. - break - if cancellation is not None: - try: - task.result() - except Exception as exc: - raise cancellation from exc - raise cancellation - return task.result() + return await self._await_owned_task(task) @staticmethod def _validate_id(value: str, label: str) -> str: @@ -343,7 +347,11 @@ class JsonlRunEventStore(RunEventStore): # single-process writers. Without the write lock, reading run files one # by one can mix events from opposite sides of a concurrent append. async with self._get_write_lock(thread_id): - events = await asyncio.to_thread(self._read_thread_events, thread_id) + task = asyncio.create_task( + asyncio.to_thread(self._read_thread_events, thread_id), + name=f"jsonl-snapshot:{thread_id}", + ) + events = await self._await_owned_task(task) result: dict[str, str] = {} for event in reversed(events): match = match_ai_message_run_id(event, pending) diff --git a/backend/tests/test_jsonl_event_store_cancellation.py b/backend/tests/test_jsonl_event_store_cancellation.py index 893f188f62..c39637b2b4 100644 --- a/backend/tests/test_jsonl_event_store_cancellation.py +++ b/backend/tests/test_jsonl_event_store_cancellation.py @@ -164,6 +164,78 @@ async def test_cancellation_while_waiting_for_lock_never_starts_write(tmp_path): assert not list(tmp_path.rglob("*.jsonl")) +@pytest.mark.anyio +@pytest.mark.parametrize("cancellations", [1, 3]) +async def test_cancelled_snapshot_keeps_thread_lock_until_read_finishes(tmp_path, monkeypatch, cancellations): + store = JsonlRunEventStore(tmp_path) + await store.put(**_event()) + paused = _PausedIO(store._read_thread_events) + monkeypatch.setattr(store, "_read_thread_events", paused) + lookup = asyncio.create_task(store.find_latest_ai_message_run_ids("t1", {"missing"})) + writer = None + try: + await asyncio.wait_for(paused.entered.wait(), 5) + for index in range(cancellations): + lookup.cancel(f"cancel-{index}") + await _checkpoint() + writer = asyncio.create_task(store.put(**_event("r2", "later"))) + await _checkpoint() + assert not lookup.done() + assert not writer.done() + assert not store._run_file("t1", "r2").exists(), "writer entered before the snapshot read settled" + paused.release.set() + with pytest.raises(asyncio.CancelledError) as caught: + await lookup + assert caught.value.args == ("cancel-0",) + assert (await writer)["seq"] == 2 + assert "t1" not in store._write_locks + finally: + paused.release.set() + await asyncio.gather(lookup, *([writer] if writer is not None else []), return_exceptions=True) + await asyncio.wait_for(paused.finished.wait(), 5) + + +@pytest.mark.anyio +@pytest.mark.parametrize("cancel", [False, True]) +async def test_snapshot_read_failure_preserves_error_precedence_and_releases_lock(tmp_path, monkeypatch, cancel): + store = JsonlRunEventStore(tmp_path) + await store.put(**_event()) + + def fail_read(_thread_id): + raise UnicodeDecodeError("utf-8", b"\xff", 0, 1, "invalid start byte") + + paused = _PausedIO(fail_read) + monkeypatch.setattr(store, "_read_thread_events", paused) + lookup = asyncio.create_task(store.find_latest_ai_message_run_ids("t1", {"missing"})) + writer = None + try: + await asyncio.wait_for(paused.entered.wait(), 5) + if cancel: + lookup.cancel("snapshot-cancelled") + await _checkpoint() + writer = asyncio.create_task(store.put(**_event("r2", "later"))) + await _checkpoint() + assert not writer.done() + assert not store._run_file("t1", "r2").exists(), "writer entered before the snapshot read settled" + + paused.release.set() + if cancel: + with pytest.raises(asyncio.CancelledError) as caught: + await lookup + assert caught.value.args == ("snapshot-cancelled",) + assert isinstance(caught.value.__cause__, UnicodeDecodeError) + else: + with pytest.raises(UnicodeDecodeError): + await lookup + + assert (await writer)["seq"] == 2 + assert "t1" not in store._write_locks + finally: + paused.release.set() + await asyncio.gather(lookup, *([writer] if writer is not None else []), return_exceptions=True) + await asyncio.wait_for(paused.finished.wait(), 5) + + @pytest.mark.anyio async def test_cancelled_idempotent_write_is_visible_to_retry_and_other_threads_progress(tmp_path, monkeypatch): store = JsonlRunEventStore(tmp_path)