diff --git a/README.md b/README.md index cd6bb04c1f..5641643634 100644 --- a/README.md +++ b/README.md @@ -582,6 +582,13 @@ DeerFlow still uses `Forwarded` / `X-Forwarded-*` headers to recover the browser > > The local two-Gateway harness disables IM providers under both `channels` and `channel_connections` by default. `DEERFLOW_MI_KEEP_CHANNELS=1` preserves them; `channels.leader_lease` stays as configured. > +> Cleaning up an in-memory run stream, or closing its bridge, releases retained +> payloads and signals a replay gap to subscribers with unfinished replay. They +> reload durable state instead of waiting on a detached buffer or consuming stale +> history. Established readers that fully consumed a terminal stream retain their +> normal `end` signal. A new reader whose initial attachment loses the race with +> cleanup receives a gap so it can recover the saved result. +> > During SSE replay-gap recovery, the Web UI preserves a reconnect pointer claimed by another run while durable state is loading. An older recovery cannot overwrite that pointer when it rejoins its own stream. Even a stale foreign pointer is preserved, so the next reload may first check that run before active-run recovery rejoins, adding a reconnect cycle. > > Store-only SSE and `/wait` consumers can finish after a heartbeat when run-manager diff --git a/backend/packages/harness/deerflow/runtime/AGENTS.md b/backend/packages/harness/deerflow/runtime/AGENTS.md index e532474066..142ceadab4 100644 --- a/backend/packages/harness/deerflow/runtime/AGENTS.md +++ b/backend/packages/harness/deerflow/runtime/AGENTS.md @@ -16,6 +16,15 @@ to sibling namespaces. Keep the cursor loop and outage/TTL fallback unchanged. Observed Redis streams removed during replay yield `StreamGap` with empty bounds; initially empty streams still wait. After this gap, `/wait` refreshes store-only status and completes only for a terminal run with no retained stream. +Memory stream cleanup retires the captured stream generation, clears its retained +payloads and wakes its Condition. Established readers emit one empty-bound +`StreamGap` and stop if replay is unfinished, including on bridge close. Fully +consumed terminal streams keep END, using the retired generation's tail offset; +readers never attach to a recreated stream. An unpublished stream without explicit +cleanup still heartbeats. +New readers that capture a generation but acquire its Condition only after +retirement always receive a gap, including cursor-less terminal reattachments. + 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. Reconnect cursor parsing requires ASCII digits and must treat integer-conversion failures as malformed input. Redis cursors additionally require unsigned 64-bit components; use the same parser before sending a cursor to `XREAD` and when detecting replay gaps. Preserve malformed-cursor policy: memory replays retained events, Redis tails active streams and replays ended streams through `end`. `tests/test_stream_bridge.py` directly tests both parsers offline, including oversized sequences, Redis component overflow, non-ASCII digits, and valid uint64 boundaries; integration tests also exercise malformed cursors against real Redis. diff --git a/backend/packages/harness/deerflow/runtime/stream_bridge/memory.py b/backend/packages/harness/deerflow/runtime/stream_bridge/memory.py index b3ecef3a0c..4fb5019565 100644 --- a/backend/packages/harness/deerflow/runtime/stream_bridge/memory.py +++ b/backend/packages/harness/deerflow/runtime/stream_bridge/memory.py @@ -22,6 +22,7 @@ class _RunStream: condition: asyncio.Condition = field(default_factory=asyncio.Condition) ended: bool = False start_offset: int = 0 + retired: bool = False class MemoryStreamBridge(StreamBridge): @@ -77,6 +78,9 @@ class MemoryStreamBridge(StreamBridge): ) def _resolve_start_offset(self, stream: _RunStream, last_event_id: str | None) -> int | StreamGap: + # Cleanup can win the Condition after a new reader captures this generation. + if stream.retired: + return StreamGap(last_event_id, None, None) if last_event_id is None: return stream.start_offset @@ -152,13 +156,16 @@ class MemoryStreamBridge(StreamBridge): while True: async with stream.condition: - if next_offset < stream.start_offset: + if stream.retired: + entry: StreamItem = END_SENTINEL if stream.ended and next_offset >= stream.start_offset else StreamGap(cursor_event_id, None, None) + should_stop = True + elif next_offset < stream.start_offset: logger.warning( "subscriber for run %s fell behind retained buffer at offset %s", run_id, next_offset, ) - entry: StreamItem = self._make_gap(stream, cursor_event_id) + entry = self._make_gap(stream, cursor_event_id) should_stop = True else: should_stop = False @@ -188,9 +195,21 @@ class MemoryStreamBridge(StreamBridge): async def cleanup(self, run_id: str, *, delay: float = 0) -> None: if delay > 0: await asyncio.sleep(delay) - self._streams.pop(run_id, None) - self._counters.pop(run_id, None) + stream = self._streams.get(run_id) + if stream is None: + return + async with stream.condition: + # Established generators retain this object after registry removal. + # Retire that generation, release its payloads, and wake readers so + # they can reload durable state rather than waiting on an orphan. + stream.retired = True + stream.start_offset += len(stream.events) + stream.events.clear() + if self._streams.get(run_id) is stream: + self._streams.pop(run_id, None) + self._counters.pop(run_id, None) + stream.condition.notify_all() async def close(self) -> None: - self._streams.clear() - self._counters.clear() + for run_id in tuple(self._streams): + await self.cleanup(run_id) diff --git a/backend/tests/test_memory_stream_retirement.py b/backend/tests/test_memory_stream_retirement.py new file mode 100644 index 0000000000..c31224d3f6 --- /dev/null +++ b/backend/tests/test_memory_stream_retirement.py @@ -0,0 +1,348 @@ +"""Retiring a memory stream must release its buffer and wake established readers.""" + +import asyncio +import json +import weakref +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest + +from deerflow.runtime import MemoryStreamBridge, RunManager, RunStatus, StreamGap +from deerflow.runtime.runs.store.memory import MemoryRunStore +from deerflow.runtime.stream_bridge.base import END_SENTINEL, HEARTBEAT_SENTINEL + + +async def _finish(bridge, streams, pending=()): + for task in pending: + task.cancel() + await asyncio.gather(*pending, return_exceptions=True) + for stream in streams: + await stream.aclose() + await bridge.close() + + +async def _attach_behind_retirement(bridge, run_id, stream, operation, monkeypatch): + """Queue retirement before attachment on the real Condition's fair lock.""" + condition = bridge._streams[run_id].condition + real_acquire = condition.acquire + retirement_queued = asyncio.Event() + attachment_queued = asyncio.Event() + tasks = [] + + async def observe_acquire(): + if asyncio.current_task() is tasks[0]: + retirement_queued.set() + else: + attachment_queued.set() + return await real_acquire() + + try: + async with condition: + monkeypatch.setattr(condition, "acquire", observe_acquire) + retirement = bridge.cleanup(run_id) if operation == "cleanup" else bridge.close() + tasks.append(asyncio.create_task(retirement)) + await asyncio.wait_for(retirement_queued.wait(), 1) + tasks.append(asyncio.create_task(anext(stream))) + await asyncio.wait_for(attachment_queued.wait(), 1) + await asyncio.wait_for(asyncio.shield(tasks[0]), 1) + return await asyncio.wait_for(asyncio.shield(tasks[1]), 1) + finally: + for task in tasks: + task.cancel() + await asyncio.gather(*tasks, return_exceptions=True) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["cleanup", "close"]) +@pytest.mark.parametrize("ended", [False, True]) +@pytest.mark.parametrize("with_cursor", [False, True]) +async def test_attachment_queued_behind_retirement_reports_gap(monkeypatch, operation, ended, with_cursor): + bridge = MemoryStreamBridge() + await bridge.publish("attach-window", "messages", "earlier") + await bridge.publish("attach-window", "messages", "unread") + cursor = bridge._streams["attach-window"].events[0].id if with_cursor else None + if ended: + await bridge.publish_end("attach-window") + stream = bridge.subscribe("attach-window", last_event_id=cursor) + try: + assert await _attach_behind_retirement(bridge, "attach-window", stream, operation, monkeypatch) == StreamGap(cursor, None, None) + with pytest.raises(StopAsyncIteration): + await anext(stream) + finally: + await _finish(bridge, [stream]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("observer", [False, True]) +@pytest.mark.parametrize("emit_gap_on_missing_stream", [False, True]) +async def test_gateway_late_attachment_reloads_instead_of_reporting_an_empty_run(monkeypatch, observer, emit_gap_on_missing_stream): + from app.gateway.services import sse_consumer + + manager = RunManager() + bridge = MemoryStreamBridge() + record = await manager.create("retirement-attach") + await manager.set_status(record.run_id, RunStatus.success) + monkeypatch.setattr(manager, "cancel", AsyncMock()) + request = SimpleNamespace(headers={}, is_disconnected=AsyncMock(return_value=False)) + await bridge.publish(record.run_id, "messages", "saved answer") + await bridge.publish_end(record.run_id) + stream = sse_consumer(bridge, record, request, manager, apply_on_disconnect=not observer, emit_gap_on_missing_stream=emit_gap_on_missing_stream) + try: + frame = await _attach_behind_retirement(bridge, record.run_id, stream, "cleanup", monkeypatch) + assert "event: gap" in frame and "event: end" not in frame + payload = json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + assert payload["recovery"] == "reload_durable_state" + assert payload["requested_event_id"] is None + assert payload["earliest_available_event_id"] is None + assert payload["latest_available_event_id"] is None + with pytest.raises(StopAsyncIteration): + await anext(stream) + manager.cancel.assert_not_awaited() + finally: + await _finish(bridge, [stream]) + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["cleanup", "close"]) +@pytest.mark.parametrize("published", [False, True]) +async def test_retirement_wakes_a_waiting_subscriber(monkeypatch, operation, published): + bridge = MemoryStreamBridge(heartbeat_interval=30) + run = "retire-waiting" + stream = bridge.subscribe(run) + cursor = None + if published: + await bridge.publish(run, "message", {"text": "partial"}) + cursor = (await anext(stream)).id + else: + # Establish an empty subscription, then wait on its real Condition. + initial = bridge.subscribe(run, heartbeat_interval=0.001) + try: + assert await anext(initial) is HEARTBEAT_SENTINEL + finally: + await initial.aclose() + condition = bridge._streams[run].condition + entered = asyncio.Event() + real_wait = condition.wait + + async def observe_wait(): + entered.set() + return await real_wait() + + monkeypatch.setattr(condition, "wait", observe_wait) + pending = asyncio.create_task(anext(stream)) + try: + await asyncio.wait_for(entered.wait(), 1) + if operation == "cleanup": + await bridge.cleanup(run) + else: + await bridge.close() + assert await asyncio.wait_for(asyncio.shield(pending), 0.5) == StreamGap(cursor, None, None) + with pytest.raises(StopAsyncIteration): + await anext(stream) + assert not await bridge.stream_exists(run) + finally: + await _finish(bridge, [stream], [pending]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["cleanup", "close"]) +@pytest.mark.parametrize("ended", [False, True]) +async def test_retired_buffer_is_released_and_does_not_replay_a_stale_tail(operation, ended): + class Payload: + pass + + bridge = MemoryStreamBridge() + await bridge.publish("buffer", "metadata", {}) + payload = Payload() + reference = weakref.ref(payload) + await bridge.publish("buffer", "messages", payload) + del payload + if ended: + await bridge.publish_end("buffer") + stream = bridge.subscribe("buffer") + first = await anext(stream) + try: + if operation == "cleanup": + await bridge.cleanup("buffer") + else: + await bridge.close() + assert reference() is None, "the established generator must not retain the retired payload" + assert await anext(stream) == StreamGap(first.id, None, None) + with pytest.raises(StopAsyncIteration): + await anext(stream) + finally: + await _finish(bridge, [stream]) + + +@pytest.mark.asyncio +async def test_cleanup_is_scoped_and_old_readers_cannot_consume_a_recreated_stream(): + bridge = MemoryStreamBridge() + for run in ("retired", "untouched"): + await bridge.publish(run, "messages", {"generation": "old", "run": run}) + old = bridge.subscribe("retired") + cursor = (await anext(old)).id + try: + await bridge.cleanup("retired") + await bridge.cleanup("missing") + await bridge.cleanup("retired") + await bridge.publish("retired", "messages", {"generation": "new"}) + await bridge.publish_end("retired") + await bridge.publish_end("untouched") + assert await asyncio.wait_for(anext(old), 0.5) == StreamGap(cursor, None, None) + replacement = [entry async for entry in bridge.subscribe("retired")] + unaffected = [entry async for entry in bridge.subscribe("untouched")] + assert replacement[0].data == {"generation": "new"} + assert replacement[-1] is END_SENTINEL + assert unaffected[0].data["run"] == "untouched" + assert unaffected[-1] is END_SENTINEL + finally: + await _finish(bridge, [old]) + + +@pytest.mark.asyncio +async def test_close_retires_all_established_readers_and_is_idempotent(): + bridge = MemoryStreamBridge() + streams = [] + cursors = [] + try: + for index in range(3): + await bridge.publish(str(index), "messages", index) + stream = bridge.subscribe(str(index)) + cursors.append((await anext(stream)).id) + streams.append(stream) + await bridge.close() + await bridge.close() + for stream, cursor in zip(streams, cursors, strict=True): + assert await asyncio.wait_for(anext(stream), 0.5) == StreamGap(cursor, None, None) + assert not bridge._streams and not bridge._counters + finally: + await _finish(bridge, streams) + + +@pytest.mark.asyncio +async def test_delayed_cleanup_preserves_replay_until_its_grace_expires(): + bridge = MemoryStreamBridge() + await bridge.publish("grace", "messages", "first") + stream = bridge.subscribe("grace") + cleanup = asyncio.create_task(bridge.cleanup("grace", delay=0.05)) + try: + first = await anext(stream) + assert first.data == "first" + assert await bridge.stream_exists("grace") + await cleanup + assert await asyncio.wait_for(anext(stream), 0.5) == StreamGap(first.id, None, None) + finally: + await _finish(bridge, [stream], [cleanup]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ["cleanup", "close"]) +@pytest.mark.parametrize("count", [0, 1, 5]) +async def test_fully_consumed_terminal_generation_keeps_end_even_after_recreation(operation, count): + bridge = MemoryStreamBridge(queue_maxsize=2) + for index in range(count): + await bridge.publish("complete", "messages", index) + stream = bridge.subscribe("complete", heartbeat_interval=0.001) + try: + if count: + for _ in range(min(count, 2)): + await anext(stream) + else: + assert await anext(stream) is HEARTBEAT_SENTINEL + await bridge.publish_end("complete") + if operation == "cleanup": + await bridge.cleanup("complete") + else: + await bridge.close() + await bridge.publish("complete", "messages", "replacement") + assert await asyncio.wait_for(anext(stream), 0.5) is END_SENTINEL + with pytest.raises(StopAsyncIteration): + await anext(stream) + finally: + await _finish(bridge, [stream]) + + +@pytest.mark.asyncio +async def test_unretired_empty_stream_heartbeats_then_delivers_normally(): + bridge = MemoryStreamBridge(heartbeat_interval=0.001) + stream = bridge.subscribe("live") + try: + assert await anext(stream) is HEARTBEAT_SENTINEL + await bridge.publish("live", "messages", "answer") + await bridge.publish_end("live") + entries = [entry async for entry in stream] + assert entries[0].data == "answer" + assert entries[-1] is END_SENTINEL + finally: + await _finish(bridge, [stream]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("observer", [False, True]) +@pytest.mark.parametrize("terminal", [False, True]) +async def test_gateway_reports_recovery_gap_without_cancelling_a_running_run(monkeypatch, observer, terminal): + from app.gateway.services import sse_consumer + + manager = RunManager() + bridge = MemoryStreamBridge() + record = await manager.create("retirement-sse") + await manager.set_status(record.run_id, RunStatus.running) + monkeypatch.setattr(manager, "cancel", AsyncMock()) + request = SimpleNamespace(headers={}, is_disconnected=AsyncMock(return_value=False)) + await bridge.publish(record.run_id, "messages", {"text": "partial"}) + stream = sse_consumer(bridge, record, request, manager, apply_on_disconnect=not observer) + try: + assert "event: messages" in await anext(stream) + if terminal: + await bridge.publish_end(record.run_id) + await manager.set_status(record.run_id, RunStatus.success) + await bridge.cleanup(record.run_id) + frame = await asyncio.wait_for(anext(stream), 0.5) + if terminal: + assert "event: end" in frame and "event: gap" not in frame + else: + assert "event: gap" in frame and "event: end" not in frame + payload = json.loads(next(line[6:] for line in frame.splitlines() if line.startswith("data: "))) + assert payload["recovery"] == "reload_durable_state" + assert payload["earliest_available_event_id"] is None + assert payload["latest_available_event_id"] is None + with pytest.raises(StopAsyncIteration): + await anext(stream) + manager.cancel.assert_not_awaited() + finally: + await _finish(bridge, [stream]) + await manager.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("store_only", [False, True]) +async def test_gateway_wait_finishes_from_terminal_state_after_retirement(store_only): + from app.gateway.services import wait_for_run_completion + + store = MemoryRunStore() if store_only else None + manager = RunManager(store=store) + bridge = MemoryStreamBridge(heartbeat_interval=30) + if store_only: + await store.put("retirement-wait", thread_id="retirement", status="running") + record = await manager.get("retirement-wait") + else: + record = await manager.create("retirement") + await manager.set_status(record.run_id, RunStatus.running) + + async def retire_after_first_delivery(): + if store_only: + await store.update_status(record.run_id, "success") + else: + await manager.set_status(record.run_id, RunStatus.success) + await bridge.cleanup(record.run_id) + return False + + request = SimpleNamespace(headers={}, is_disconnected=retire_after_first_delivery) + await bridge.publish(record.run_id, "messages", "last update") + try: + assert await asyncio.wait_for(wait_for_run_completion(bridge, record, request, manager), 0.5) + finally: + await bridge.close() + await manager.shutdown()