From dfc8e724280ab9460a48834227437ef5697e4f54 Mon Sep 17 00:00:00 2001 From: Jun <84921700+Amazingjun-j@users.noreply.github.com> Date: Sun, 13 Sep 2026 19:55:58 +0800 Subject: [PATCH] fix(agents): make read-before-write cancellation lock-safe (#5395) * test(agents): expose read-before-write cancellation races * fix(agents): make read-before-write cancellation lock-safe * test(agents): fix cancellation test import order * test(agents): apply cancellation test formatting * fix(agents): preserve first cancellation when worker task cancels * test(agents): cover cancelled worker preservation * fix(agents): preserve cancellation over worker errors * test(agents): cover worker error after cancellation --- .../read_before_write_middleware.py | 52 +++- .../test_read_before_write_cancellation.py | 278 ++++++++++++++++++ 2 files changed, 323 insertions(+), 7 deletions(-) create mode 100644 backend/tests/test_read_before_write_cancellation.py diff --git a/backend/packages/harness/deerflow/agents/middlewares/read_before_write_middleware.py b/backend/packages/harness/deerflow/agents/middlewares/read_before_write_middleware.py index 48dcabac3..322d672dc 100644 --- a/backend/packages/harness/deerflow/agents/middlewares/read_before_write_middleware.py +++ b/backend/packages/harness/deerflow/agents/middlewares/read_before_write_middleware.py @@ -115,6 +115,45 @@ def _content_hash(content: str) -> str: return hashlib.sha256(content.encode("utf-8")).hexdigest() +async def _await_off_thread(task: asyncio.Task[Any]) -> Any: + """Drain an already-dispatched worker operation before propagating cancellation.""" + first_cancel: asyncio.CancelledError | None = None + while True: + try: + result = await asyncio.shield(task) + except asyncio.CancelledError as exc: + if task.cancelled(): + if first_cancel is not None: + raise first_cancel + raise + if first_cancel is None: + first_cancel = exc + if not task.done(): + continue + except BaseException: + if first_cancel is None: + raise + else: + if first_cancel is None: + return result + + if first_cancel is not None: + if task.done() and not task.cancelled(): + task.exception() + raise first_cancel + + +async def _acquire_gate_lock(lock: threading.Lock) -> None: + """Acquire off-loop safely; threading.Lock permits cross-thread release.""" + acquire_task = asyncio.create_task(asyncio.to_thread(lock.acquire)) + try: + await _await_off_thread(acquire_task) + except asyncio.CancelledError: + if acquire_task.done() and not acquire_task.cancelled() and acquire_task.exception() is None: + lock.release() + raise + + class ReadBeforeWriteMiddleware(AgentMiddleware): """Version gate: block writes to existing files not read at their current version.""" @@ -180,13 +219,11 @@ class ReadBeforeWriteMiddleware(AgentMiddleware): return await handler(request) try: async with sandbox_authorization_scope_async(request.runtime): - # threading.Lock may be released from a different thread than the - # acquiring one, so acquiring in a worker thread and releasing on - # the event-loop thread is safe. lock = self._lock_for(request, path) - await asyncio.to_thread(lock.acquire) + await _acquire_gate_lock(lock) try: - blocked = await asyncio.to_thread(self._check_write_gate, request) + check_task = asyncio.create_task(asyncio.to_thread(self._check_write_gate, request)) + blocked = await _await_off_thread(check_task) if blocked is not None: return normalize_tool_result(blocked) return await handler(request) @@ -201,10 +238,11 @@ class ReadBeforeWriteMiddleware(AgentMiddleware): try: async with sandbox_authorization_scope_async(request.runtime): lock = self._lock_for(request, path) - await asyncio.to_thread(lock.acquire) + await _acquire_gate_lock(lock) try: result = await handler(request) - await asyncio.to_thread(self._attach_read_mark, request, result) + mark_task = asyncio.create_task(asyncio.to_thread(self._attach_read_mark, request, result)) + await _await_off_thread(mark_task) return result finally: lock.release() diff --git a/backend/tests/test_read_before_write_cancellation.py b/backend/tests/test_read_before_write_cancellation.py new file mode 100644 index 000000000..be9ae2faf --- /dev/null +++ b/backend/tests/test_read_before_write_cancellation.py @@ -0,0 +1,278 @@ +from __future__ import annotations + +import asyncio +import threading +from contextlib import suppress +from unittest.mock import MagicMock + +import pytest +from langchain_core.messages import ToolMessage +from langgraph.prebuilt.tool_node import ToolCallRequest + +from deerflow.agents.middlewares.read_before_write_middleware import ReadBeforeWriteMiddleware, _await_off_thread +from deerflow.sandbox.exceptions import SandboxAuthorizationError + +_PATH = "/mnt/user-data/outputs/report.md" + + +def _request(tool_name: str) -> ToolCallRequest: + runtime = MagicMock() + runtime.context = {"thread_id": "thread-cancel"} + args: dict[str, object] = {"description": "d", "path": _PATH} + if tool_name == "write_file": + args["content"] = "v2" + return ToolCallRequest( + tool_call={"name": tool_name, "args": args, "id": f"call-{tool_name}"}, + tool=None, + state={"messages": []}, + runtime=runtime, + ) + + +class _BlockingAcquireLock: + def __init__(self) -> None: + self.acquire_started = threading.Event() + self.allow_acquire = threading.Event() + self.acquired = threading.Event() + self.release_calls = 0 + self._guard = threading.Lock() + + def acquire(self) -> bool: + self.acquire_started.set() + assert self.allow_acquire.wait(timeout=5), "test did not unblock gate-lock acquisition" + with self._guard: + self.acquired.set() + return True + + def release(self) -> None: + with self._guard: + assert self.acquired.is_set(), "gate lock released without ownership" + self.acquired.clear() + self.release_calls += 1 + + +class _TrackingLock: + def __init__(self) -> None: + self.acquired = False + self.release_calls = 0 + self._guard = threading.Lock() + + def acquire(self) -> bool: + with self._guard: + assert not self.acquired + self.acquired = True + return True + + def release(self) -> None: + with self._guard: + assert self.acquired + self.acquired = False + self.release_calls += 1 + + +@pytest.mark.parametrize("tool_name", ["read_file", "write_file"]) +def test_async_gate_cancellation_drains_queued_lock_acquisition( + monkeypatch: pytest.MonkeyPatch, + tool_name: str, +) -> None: + async def scenario() -> None: + gate_lock = _BlockingAcquireLock() + middleware = ReadBeforeWriteMiddleware(content_reader=lambda _runtime, _path: "v1") + monkeypatch.setattr(middleware, "_lock_for", lambda _request, _path: gate_lock) + handler_called = False + + async def handler(_request: ToolCallRequest) -> ToolMessage: + nonlocal handler_called + handler_called = True + return ToolMessage(content="OK", tool_call_id=f"call-{tool_name}", name=tool_name) + + task = asyncio.create_task(middleware.awrap_tool_call(_request(tool_name), handler)) + try: + assert await asyncio.to_thread(gate_lock.acquire_started.wait, 2), "gate-lock acquisition did not start" + + task.cancel("first cancellation") + await asyncio.sleep(0) + + # threading.Lock.acquire() is already running in a worker thread and + # cannot be cancelled. The middleware must retain ownership of that + # acquisition until it lands, then release it before propagating the + # original cancellation. Returning cancellation here would orphan a + # future successful acquire with nobody left to release it. + assert not task.done() + assert gate_lock.release_calls == 0 + + task.cancel("second cancellation") + await asyncio.sleep(0) + + assert not task.done() + assert gate_lock.release_calls == 0 + assert not handler_called + + gate_lock.allow_acquire.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + + assert exc_info.value.args == ("first cancellation",) + assert gate_lock.release_calls == 1 + assert not gate_lock.acquired.is_set() + assert not handler_called + finally: + gate_lock.allow_acquire.set() + await asyncio.to_thread(gate_lock.acquired.wait, 0.1) + if gate_lock.acquired.is_set() and gate_lock.release_calls == 0: + gate_lock.release() + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("tool_name", ["read_file", "write_file"]) +def test_async_gate_cancellation_drains_sync_probe_before_unlocking( + monkeypatch: pytest.MonkeyPatch, + tool_name: str, +) -> None: + async def scenario() -> None: + probe_started = threading.Event() + allow_probe_finish = threading.Event() + gate_lock = _TrackingLock() + + def reader(_runtime: object, _path: str) -> str: + probe_started.set() + assert allow_probe_finish.wait(timeout=5), "test did not unblock gate probe" + return "v1" + + middleware = ReadBeforeWriteMiddleware(content_reader=reader) + monkeypatch.setattr(middleware, "_lock_for", lambda _request, _path: gate_lock) + handler_calls = 0 + + async def handler(_request: ToolCallRequest) -> ToolMessage: + nonlocal handler_calls + handler_calls += 1 + return ToolMessage(content="v1", tool_call_id=f"call-{tool_name}", name=tool_name) + + task = asyncio.create_task(middleware.awrap_tool_call(_request(tool_name), handler)) + try: + assert await asyncio.to_thread(probe_started.wait, 2), "gate probe did not start" + assert gate_lock.acquired + + task.cancel("first cancellation") + await asyncio.sleep(0) + + # The synchronous probe is also non-cancellable once dispatched. + # Releasing the gate lock before that worker returns lets another + # same-path tool enter the critical section concurrently with work + # that this task started under the lock. + assert not task.done() + assert gate_lock.acquired + assert gate_lock.release_calls == 0 + + task.cancel("second cancellation") + await asyncio.sleep(0) + + assert not task.done() + assert gate_lock.acquired + assert gate_lock.release_calls == 0 + + allow_probe_finish.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + + assert exc_info.value.args == ("first cancellation",) + assert not gate_lock.acquired + assert gate_lock.release_calls == 1 + assert handler_calls == (1 if tool_name == "read_file" else 0) + finally: + allow_probe_finish.set() + if gate_lock.acquired: + gate_lock.release() + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + asyncio.run(scenario()) + + +@pytest.mark.parametrize("tool_name", ["read_file", "write_file"]) +def test_async_gate_cancellation_wins_over_sync_probe_authorization_error( + monkeypatch: pytest.MonkeyPatch, + tool_name: str, +) -> None: + async def scenario() -> None: + probe_started = threading.Event() + allow_probe_finish = threading.Event() + gate_lock = _TrackingLock() + + def reader(_runtime: object, _path: str) -> str: + probe_started.set() + assert allow_probe_finish.wait(timeout=5), "test did not unblock failing gate probe" + raise SandboxAuthorizationError("probe denied") + + middleware = ReadBeforeWriteMiddleware(content_reader=reader) + monkeypatch.setattr(middleware, "_lock_for", lambda _request, _path: gate_lock) + handler_calls = 0 + + async def handler(_request: ToolCallRequest) -> ToolMessage: + nonlocal handler_calls + handler_calls += 1 + return ToolMessage(content="v1", tool_call_id=f"call-{tool_name}", name=tool_name) + + task = asyncio.create_task(middleware.awrap_tool_call(_request(tool_name), handler)) + try: + assert await asyncio.to_thread(probe_started.wait, 2), "failing gate probe did not start" + assert gate_lock.acquired + + task.cancel("first cancellation") + await asyncio.sleep(0) + + assert not task.done() + assert gate_lock.acquired + assert gate_lock.release_calls == 0 + + allow_probe_finish.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + + assert exc_info.value.args == ("first cancellation",) + assert not gate_lock.acquired + assert gate_lock.release_calls == 1 + assert handler_calls == (1 if tool_name == "read_file" else 0) + finally: + allow_probe_finish.set() + if gate_lock.acquired: + gate_lock.release() + if not task.done(): + task.cancel() + with suppress(asyncio.CancelledError): + await task + + asyncio.run(scenario()) + + +def test_await_off_thread_preserves_first_cancel_when_worker_task_is_cancelled() -> None: + async def scenario() -> None: + worker_started = asyncio.Event() + worker_can_finish = asyncio.Event() + + async def worker() -> None: + worker_started.set() + await worker_can_finish.wait() + + worker_task = asyncio.create_task(worker()) + waiter = asyncio.create_task(_await_off_thread(worker_task)) + await worker_started.wait() + + waiter.cancel("first cancellation") + await asyncio.sleep(0) + assert not waiter.done() + + worker_task.cancel("inner cancellation") + with pytest.raises(asyncio.CancelledError) as exc_info: + await waiter + + assert exc_info.value.args == ("first cancellation",) + + asyncio.run(scenario())