diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index b83b6dc40c..7d33d21cc2 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -1,6 +1,6 @@ ### Stream Bridge Heartbeats -Memory/Redis bridges keep startup-only `stream_bridge.heartbeat_interval_seconds`; explicit `subscribe(..., heartbeat_interval=...)` overrides it. Provider contexts retain backend ownership through exit: drain cache/bridge `aclose()`/`close()` and SQLite/PostgreSQL checkpointer/Store `__aexit__` across cancellation before it propagates. +Memory/Redis bridges keep startup-only `stream_bridge.heartbeat_interval_seconds`; explicit `subscribe(..., heartbeat_interval=...)` overrides it. Provider contexts drain owned cache/bridge/checkpointer/Store teardown across cancellation. `close_agent_stream()` shields close to completion: preserve counts, defer host cancellation, balance repeats, keep active errors, log attached close failures at callers, map close cancellation to failure, never time out. ### Checkpoint Channel Modes (`full` / `delta`) diff --git a/backend/packages/harness/deerflow/runtime/runs/stream_cleanup.py b/backend/packages/harness/deerflow/runtime/runs/stream_cleanup.py index e3523983fd..b210a601e3 100644 --- a/backend/packages/harness/deerflow/runtime/runs/stream_cleanup.py +++ b/backend/packages/harness/deerflow/runtime/runs/stream_cleanup.py @@ -1,14 +1,111 @@ from __future__ import annotations +import asyncio import inspect +import sys from typing import Any +class AgentStreamCloseCancelledError(RuntimeError): + """Raised when the stream cancels its own asynchronous teardown.""" + + +def _stream_close_cancelled(exc: asyncio.CancelledError) -> AgentStreamCloseCancelledError: + error = AgentStreamCloseCancelledError("Agent stream cancelled its own close operation") + error.__cause__ = exc + return error + + async def close_agent_stream(stream: Any) -> None: - """Close an agent stream when its runtime exposes asynchronous cleanup.""" + """Close an agent stream before propagating caller cancellation. + + Stream teardown owns provider, graph, and tool cleanup that must finish + before the caller releases run-scoped resources. Host cancellation is + therefore deferred while the close task drains. A cancellation raised by + the close awaitable itself remains visible to the caller. + """ close = getattr(stream, "aclose", None) if close is None: return - result = close() - if inspect.isawaitable(result): - await result + caller = asyncio.current_task() + active_exception = sys.exception() + cancellation_count = caller.cancelling() if caller is not None else 0 + close_failure: Exception | None = None + result: Any = None + try: + result = close() + except asyncio.CancelledError as exc: + close_failure = _stream_close_cancelled(exc) + except Exception as exc: + close_failure = exc + if caller is None: # pragma: no cover - every running coroutine has a task + if inspect.isawaitable(result): + await result + elif close_failure is not None: + raise close_failure + return + + deferred_cancellation: asyncio.CancelledError | None = None + + def consume_new_host_cancellation(exc: asyncio.CancelledError) -> bool: + nonlocal cancellation_count, deferred_cancellation + current_count = caller.cancelling() + new_cancellations = current_count - cancellation_count + if new_cancellations <= 0: + return False + preserve_cancellation = active_exception is None and deferred_cancellation is None + cancellations_to_consume = new_cancellations - int(preserve_cancellation) + for _ in range(cancellations_to_consume): + caller.uncancel() + if preserve_cancellation: + deferred_cancellation = exc + cancellation_count = caller.cancelling() + return True + + # Deliver cancellation that was already pending before close_task exists. + # A zero count delta here is therefore host-side, not close-task cancellation. + try: + await asyncio.sleep(0) + except asyncio.CancelledError as exc: + current_count = caller.cancelling() + new_cancellations = max(0, current_count - cancellation_count) + preserve_cancellation = active_exception is None + cancellations_to_consume = new_cancellations - int(preserve_cancellation and new_cancellations > 0) + for _ in range(cancellations_to_consume): + caller.uncancel() + if preserve_cancellation: + deferred_cancellation = exc + cancellation_count = caller.cancelling() + else: + cancellation_count = caller.cancelling() + + close_task = asyncio.ensure_future(result) if inspect.isawaitable(result) else None + while close_task is not None and not close_task.done(): + try: + await asyncio.shield(close_task) + except asyncio.CancelledError as exc: + if not consume_new_host_cancellation(exc): + close_failure = _stream_close_cancelled(exc) + break + except Exception as exc: + close_failure = exc + break + + if close_task is not None and close_task.done(): + try: + close_task.result() + except asyncio.CancelledError as exc: + close_failure = _stream_close_cancelled(exc) + except Exception as exc: + close_failure = exc + + if active_exception is not None: + if close_failure is not None: + raise active_exception from close_failure + return + if deferred_cancellation is not None: + if close_failure is not None: + raise deferred_cancellation from close_failure + raise deferred_cancellation + if close_failure is not None: + raise close_failure diff --git a/backend/packages/harness/deerflow/runtime/runs/worker.py b/backend/packages/harness/deerflow/runtime/runs/worker.py index 14d4c9d023..0af012bc56 100644 --- a/backend/packages/harness/deerflow/runtime/runs/worker.py +++ b/backend/packages/harness/deerflow/runtime/runs/worker.py @@ -105,6 +105,25 @@ from .schemas import RunStatus logger = logging.getLogger(__name__) _THREAD_INCARNATION_UNSET = object() + +def _log_cancelled_stream_close_failure( + exc: asyncio.CancelledError, + *, + run_id: str, + abort_requested: bool, +) -> None: + close_failure = exc.__cause__ + if not isinstance(close_failure, Exception): + return + log = logger.warning if abort_requested else logger.debug + message = "Could not close aborted agent stream for run %s" if abort_requested else "Could not close agent stream for run %s" + log( + message, + run_id, + exc_info=(type(close_failure), close_failure, close_failure.__traceback__), + ) + + _checkpoint_locks = AsyncKeyedLockTable[str]() # Completed LangGraph runs can leave callback Contexts and AsyncPregelLoop @@ -1295,6 +1314,13 @@ async def run_agent( close_error = sys.exception() try: await close_agent_stream(stream) + except asyncio.CancelledError as exc: + _log_cancelled_stream_close_failure( + exc, + run_id=run_id, + abort_requested=broke_on_abort or record.abort_event.is_set(), + ) + raise except Exception: abort_requested = broke_on_abort or record.abort_event.is_set() if close_error is None and not abort_requested: @@ -1344,6 +1370,13 @@ async def run_agent( close_error = sys.exception() try: await close_agent_stream(stream) + except asyncio.CancelledError as exc: + _log_cancelled_stream_close_failure( + exc, + run_id=run_id, + abort_requested=broke_on_abort or record.abort_event.is_set(), + ) + raise except Exception: abort_requested = broke_on_abort or record.abort_event.is_set() if close_error is None and not abort_requested: diff --git a/backend/packages/harness/deerflow/subagents/executor.py b/backend/packages/harness/deerflow/subagents/executor.py index 5b3503060b..9b8649b3e8 100644 --- a/backend/packages/harness/deerflow/subagents/executor.py +++ b/backend/packages/harness/deerflow/subagents/executor.py @@ -1664,6 +1664,16 @@ class SubagentExecutor: ) try: await close_agent_stream(stream) + except asyncio.CancelledError as exc: + close_failure = exc.__cause__ + if isinstance(close_failure, Exception): + logger.warning( + "[trace=%s] Could not close interrupted subagent stream %s", + self.trace_id, + self.config.name, + exc_info=(type(close_failure), close_failure, close_failure.__traceback__), + ) + raise except Exception: cancel_requested = cancel_requested or result.cancel_event.is_set() if active_error is None and not cancel_requested: diff --git a/backend/tests/test_run_worker_rollback.py b/backend/tests/test_run_worker_rollback.py index 64976dfcd0..b456a197a7 100644 --- a/backend/tests/test_run_worker_rollback.py +++ b/backend/tests/test_run_worker_rollback.py @@ -1146,6 +1146,93 @@ async def test_run_agent_closes_stream_when_abort_breaks_iteration(): assert record.status == RunStatus.interrupted +@pytest.mark.parametrize("stream_modes", [["values"], ["messages-tuple", "values"]]) +@pytest.mark.parametrize("close_fails", [False, True], ids=["close-succeeds", "close-fails"]) +@pytest.mark.anyio +async def test_run_agent_repeated_cancellation_waits_for_stream_close(stream_modes, close_fails, caplog): + run_manager = RunManager() + record = await run_manager.create("thread-stream-close-cancellation") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + iteration_started = asyncio.Event() + close_started = asyncio.Event() + allow_close = asyncio.Event() + + class BlockingStream: + def __init__(self) -> None: + self.close_cancelled = False + self.closed = False + + def __aiter__(self): + return self + + async def __anext__(self): + iteration_started.set() + await asyncio.Event().wait() + + async def aclose(self) -> None: + close_started.set() + try: + await allow_close.wait() + except asyncio.CancelledError: + self.close_cancelled = True + raise + if close_fails: + raise RuntimeError("stream close failed") + self.closed = True + + stream = BlockingStream() + + class DummyAgent: + def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + return stream + + run_task = asyncio.create_task( + run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: DummyAgent(), + graph_input={}, + config={}, + stream_modes=stream_modes, + ) + ) + await iteration_started.wait() + run_task.cancel("first") + await close_started.wait() + + run_task.cancel("second") + await asyncio.sleep(0) + run_task.cancel("third") + await asyncio.sleep(0) + + assert not run_task.done() + assert not stream.close_cancelled + assert not stream.closed + bridge.publish_end.assert_not_awaited() + + with caplog.at_level(logging.DEBUG, logger="deerflow.runtime.runs.worker"): + allow_close.set() + await run_task + + assert not stream.close_cancelled + assert record.status == RunStatus.interrupted + if close_fails: + assert not stream.closed + assert "Could not close agent stream" in caplog.text + assert "stream close failed" in caplog.text + else: + assert stream.closed + assert "Could not close agent stream" not in caplog.text + bridge.publish_end.assert_awaited_once_with(record.run_id) + + @pytest.mark.parametrize("stream_modes", [["values"], ["messages-tuple", "values"]]) @pytest.mark.parametrize("abort_before_break", [True, False], ids=["early-break", "exhaustion-race"]) @pytest.mark.anyio @@ -1207,6 +1294,62 @@ async def test_run_agent_ignores_stream_close_failure_after_abort(stream_modes, bridge.publish_end.assert_awaited_once_with(record.run_id) +@pytest.mark.parametrize("stream_modes", [["values"], ["messages-tuple", "values"]]) +@pytest.mark.parametrize("synchronous_close", [False, True], ids=["async-close", "sync-close"]) +@pytest.mark.anyio +async def test_run_agent_treats_stream_originated_close_cancellation_as_error(stream_modes, synchronous_close): + run_manager = RunManager() + record = await run_manager.create("thread-stream-self-cancel") + bridge = SimpleNamespace( + publish=AsyncMock(), + publish_end=AsyncMock(), + cleanup=AsyncMock(), + ) + + class SelfCancellingCloseStream: + def __aiter__(self): + return self + + async def __anext__(self): + raise StopAsyncIteration + + def aclose(self): + if synchronous_close: + raise asyncio.CancelledError("stream cancelled its own close") + + async def cancel_close() -> None: + raise asyncio.CancelledError("stream cancelled its own close") + + return cancel_close() + + stream = SelfCancellingCloseStream() + + class DummyAgent: + def astream(self, graph_input, config=None, stream_mode=None, subgraphs=False): + del graph_input, config, stream_mode, subgraphs + return stream + + await run_agent( + bridge, + run_manager, + record, + ctx=RunContext(checkpointer=None), + agent_factory=lambda **_kwargs: DummyAgent(), + graph_input={}, + config={}, + stream_modes=stream_modes, + ) + + assert record.status == RunStatus.error + assert record.error == "Agent stream cancelled its own close operation" + error_events = [call.args for call in bridge.publish.await_args_list if call.args[1] == "error"] + assert len(error_events) == 1 + error_event = error_events[0] + assert error_event[1] == "error" + assert error_event[2]["name"] == "AgentStreamCloseCancelledError" + bridge.publish_end.assert_awaited_once_with(record.run_id) + + @pytest.mark.anyio async def test_terminal_cleanup_tasks_do_not_inherit_run_context(): marker: ContextVar[str | None] = ContextVar("run_cleanup_marker", default=None) diff --git a/backend/tests/test_stream_cleanup.py b/backend/tests/test_stream_cleanup.py new file mode 100644 index 0000000000..f70cad8194 --- /dev/null +++ b/backend/tests/test_stream_cleanup.py @@ -0,0 +1,390 @@ +from __future__ import annotations + +import asyncio + +import pytest + +from deerflow.runtime.runs.stream_cleanup import AgentStreamCloseCancelledError, close_agent_stream + + +class _BlockingStream: + def __init__(self) -> None: + self.close_started = asyncio.Event() + self.allow_close = asyncio.Event() + self.closed = asyncio.Event() + self.close_cancelled = False + + async def aclose(self) -> None: + self.close_started.set() + try: + await self.allow_close.wait() + except asyncio.CancelledError: + self.close_cancelled = True + raise + self.closed.set() + + +class _FailingCloseStream(_BlockingStream): + async def aclose(self) -> None: + self.close_started.set() + try: + await self.allow_close.wait() + except asyncio.CancelledError: + self.close_cancelled = True + raise + raise RuntimeError("close failed") + + +@pytest.mark.asyncio +async def test_repeated_cancellation_does_not_interrupt_stream_close() -> None: + stream = _BlockingStream() + running = asyncio.Event() + first_cancellation: asyncio.CancelledError | None = None + + async def run_until_cancelled() -> None: + nonlocal first_cancellation + try: + running.set() + await asyncio.Event().wait() + except asyncio.CancelledError as exc: + first_cancellation = exc + await close_agent_stream(stream) + raise + + task = asyncio.create_task(run_until_cancelled()) + await running.wait() + task.cancel("first") + await stream.close_started.wait() + + task.cancel("second") + await asyncio.sleep(0) + task.cancel("third") + await asyncio.sleep(0) + + assert not task.done() + assert not stream.close_cancelled + assert not stream.closed.is_set() + + stream.allow_close.set() + with pytest.raises(asyncio.CancelledError) as exc_info: + await task + + assert exc_info.value is first_cancellation + assert exc_info.value.args == ("first",) + assert stream.closed.is_set() + assert not stream.close_cancelled + + +@pytest.mark.asyncio +async def test_cancellation_pending_before_first_wait_is_deferred() -> None: + close_started = asyncio.Event() + allow_close = asyncio.Event() + close_cancelled = False + + class CancelBeforeAwaitStream: + def aclose(self): + caller = asyncio.current_task() + assert caller is not None + caller.cancel("pending-before-wait") + + async def drain() -> None: + nonlocal close_cancelled + close_started.set() + try: + await allow_close.wait() + except asyncio.CancelledError: + close_cancelled = True + raise + + return drain() + + task = asyncio.create_task(close_agent_stream(CancelBeforeAwaitStream())) + await close_started.wait() + await asyncio.sleep(0) + + assert not task.done() + assert not close_cancelled + + allow_close.set() + with pytest.raises(asyncio.CancelledError, match="pending-before-wait"): + await task + + assert not close_cancelled + + +@pytest.mark.asyncio +async def test_pre_wait_cancellation_does_not_replace_active_stream_error() -> None: + close_started = asyncio.Event() + allow_close = asyncio.Event() + close_cancelled = False + original_error = RuntimeError("stream failed") + + class CancelBeforeAwaitStream: + def aclose(self): + caller = asyncio.current_task() + assert caller is not None + caller.cancel("pending-before-wait") + + async def drain() -> None: + nonlocal close_cancelled + close_started.set() + try: + await allow_close.wait() + except asyncio.CancelledError: + close_cancelled = True + raise + + return drain() + + async def fail_then_close() -> None: + try: + raise original_error + finally: + await close_agent_stream(CancelBeforeAwaitStream()) + + task = asyncio.create_task(fail_then_close()) + await close_started.wait() + await asyncio.sleep(0) + + assert not task.done() + assert not close_cancelled + + allow_close.set() + with pytest.raises(RuntimeError, match="stream failed") as exc_info: + await task + + assert exc_info.value is original_error + assert not close_cancelled + assert task.cancelling() == 0 + + +@pytest.mark.asyncio +async def test_pre_entry_cancellation_does_not_replace_active_stream_error() -> None: + stream = _BlockingStream() + original_error = RuntimeError("stream failed") + + async def cancel_fail_then_close() -> None: + caller = asyncio.current_task() + assert caller is not None + caller.cancel("pending-before-entry") + try: + raise original_error + finally: + await close_agent_stream(stream) + + task = asyncio.create_task(cancel_fail_then_close()) + await stream.close_started.wait() + await asyncio.sleep(0) + + assert not task.done() + assert not stream.close_cancelled + + stream.allow_close.set() + with pytest.raises(RuntimeError, match="stream failed") as exc_info: + await task + + assert exc_info.value is original_error + assert stream.closed.is_set() + assert not stream.close_cancelled + assert task.cancelling() == 1 + + +@pytest.mark.asyncio +async def test_pre_entry_and_synchronous_close_cancellations_are_balanced() -> None: + close_started = asyncio.Event() + allow_close = asyncio.Event() + original_error = RuntimeError("stream failed") + + class CancelAgainBeforeAwaitStream: + def aclose(self): + caller = asyncio.current_task() + assert caller is not None + caller.cancel("during-close") + + async def drain() -> None: + close_started.set() + await allow_close.wait() + + return drain() + + async def cancel_fail_then_close() -> None: + caller = asyncio.current_task() + assert caller is not None + caller.cancel("before-entry") + try: + raise original_error + finally: + await close_agent_stream(CancelAgainBeforeAwaitStream()) + + task = asyncio.create_task(cancel_fail_then_close()) + await close_started.wait() + await asyncio.sleep(0) + + assert not task.done() + + allow_close.set() + with pytest.raises(RuntimeError, match="stream failed") as exc_info: + await task + + assert exc_info.value is original_error + assert task.cancelling() == 1 + + +@pytest.mark.asyncio +async def test_delivered_cancellation_count_is_preserved_while_active_error_wins() -> None: + original_error = RuntimeError("stream failed") + + class ImmediateStream: + async def aclose(self) -> None: + return None + + async def cancel_then_fail() -> None: + caller = asyncio.current_task() + assert caller is not None + caller.cancel("already delivered") + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + pass + try: + raise original_error + finally: + await close_agent_stream(ImmediateStream()) + + task = asyncio.create_task(cancel_then_fail()) + with pytest.raises(RuntimeError, match="stream failed") as exc_info: + await task + + assert exc_info.value is original_error + assert task.cancelling() == 1 + + +@pytest.mark.asyncio +async def test_delivered_cancellation_count_survives_a_new_deferred_cancellation() -> None: + stream = _BlockingStream() + old_cancellation_delivered = asyncio.Event() + + async def cancel_then_close() -> None: + caller = asyncio.current_task() + assert caller is not None + caller.cancel("already delivered") + try: + await asyncio.sleep(0) + except asyncio.CancelledError: + old_cancellation_delivered.set() + await close_agent_stream(stream) + + task = asyncio.create_task(cancel_then_close()) + await old_cancellation_delivered.wait() + await stream.close_started.wait() + + task.cancel("deferred") + await asyncio.sleep(0) + + assert not task.done() + assert task.cancelling() == 2 + assert not stream.close_cancelled + + stream.allow_close.set() + with pytest.raises(asyncio.CancelledError, match="deferred"): + await task + + assert task.cancelling() == 2 + assert stream.closed.is_set() + + +@pytest.mark.asyncio +async def test_close_failure_is_preserved_as_cancellation_cause() -> None: + stream = _FailingCloseStream() + task = asyncio.create_task(close_agent_stream(stream)) + await stream.close_started.wait() + + task.cancel("close-cancelled") + await asyncio.sleep(0) + stream.allow_close.set() + + with pytest.raises(asyncio.CancelledError, match="close-cancelled") as exc_info: + await task + + assert isinstance(exc_info.value.__cause__, RuntimeError) + assert str(exc_info.value.__cause__) == "close failed" + assert not stream.close_cancelled + + +@pytest.mark.asyncio +async def test_host_cancellation_during_close_does_not_replace_active_stream_error() -> None: + stream = _BlockingStream() + original_error = RuntimeError("stream failed") + + async def fail_then_close() -> None: + try: + raise original_error + finally: + await close_agent_stream(stream) + + task = asyncio.create_task(fail_then_close()) + await stream.close_started.wait() + task.cancel("host cancel") + await asyncio.sleep(0) + + assert not task.done() + assert not stream.close_cancelled + + stream.allow_close.set() + with pytest.raises(RuntimeError, match="stream failed") as exc_info: + await task + + assert exc_info.value is original_error + assert stream.closed.is_set() + assert not stream.close_cancelled + + +@pytest.mark.asyncio +async def test_stream_originated_cancellation_is_reported_as_close_failure() -> None: + class SelfCancellingStream: + async def aclose(self) -> None: + raise asyncio.CancelledError("stream cancelled itself") + + with pytest.raises(AgentStreamCloseCancelledError, match="cancelled its own close operation") as exc_info: + await close_agent_stream(SelfCancellingStream()) + + assert isinstance(exc_info.value.__cause__, asyncio.CancelledError) + assert exc_info.value.__cause__.args == ("stream cancelled itself",) + assert exc_info.value.__suppress_context__ is True + task = asyncio.current_task() + assert task is not None + assert task.cancelling() == 0 + + +@pytest.mark.asyncio +async def test_synchronous_stream_close_cancellation_is_reported_as_failure() -> None: + class SyncSelfCancellingStream: + def aclose(self) -> None: + raise asyncio.CancelledError("sync stream close cancellation") + + with pytest.raises(AgentStreamCloseCancelledError, match="cancelled its own close operation") as exc_info: + await close_agent_stream(SyncSelfCancellingStream()) + + assert isinstance(exc_info.value.__cause__, asyncio.CancelledError) + assert exc_info.value.__cause__.args == ("sync stream close cancellation",) + assert exc_info.value.__suppress_context__ is True + + +@pytest.mark.asyncio +async def test_cancellation_during_normal_close_is_deferred_until_close_finishes() -> None: + stream = _BlockingStream() + task = asyncio.create_task(close_agent_stream(stream)) + await stream.close_started.wait() + + task.cancel("close-cancelled") + await asyncio.sleep(0) + + assert not task.done() + assert not stream.close_cancelled + + stream.allow_close.set() + with pytest.raises(asyncio.CancelledError, match="close-cancelled"): + await task + + assert stream.closed.is_set() + assert not stream.close_cancelled diff --git a/backend/tests/test_subagent_executor.py b/backend/tests/test_subagent_executor.py index 61405f67d9..9d1e8e49b2 100644 --- a/backend/tests/test_subagent_executor.py +++ b/backend/tests/test_subagent_executor.py @@ -17,6 +17,7 @@ the real implementation in isolation. import asyncio import importlib import inspect +import logging import sys import threading import time @@ -2024,6 +2025,7 @@ class TestAsyncExecutionPath: classes, base_config, mock_agent, + caplog, ): SubagentExecutor = classes["SubagentExecutor"] started = asyncio.Event() @@ -2049,15 +2051,18 @@ class TestAsyncExecutionPath: thread_id="test-thread", ) - with patch.object(executor, "_create_agent", return_value=mock_agent): - execution = asyncio.create_task(executor._aexecute("Task")) - await started.wait() - execution.cancel("host cancellation") - with pytest.raises(asyncio.CancelledError) as raised: - await execution + with caplog.at_level(logging.WARNING, logger="deerflow.subagents.executor"): + with patch.object(executor, "_create_agent", return_value=mock_agent): + execution = asyncio.create_task(executor._aexecute("Task")) + await started.wait() + execution.cancel("host cancellation") + with pytest.raises(asyncio.CancelledError) as raised: + await execution assert raised.value.args == ("host cancellation",) assert close_attempted is True + assert "Could not close interrupted subagent stream" in caplog.text + assert "close failed" in caplog.text @pytest.mark.anyio async def test_aexecute_finally_releases_only_the_failing_subagent_lease(