From be46df481e4edd331883d517a3ffb74dd5cbd72b Mon Sep 17 00:00:00 2001 From: ogabrielluiz Date: Thu, 4 Jun 2026 00:18:32 -0300 Subject: [PATCH] feat(workflows): add in-memory live bus with durable replay reattach --- .../services/background_execution/live_bus.py | 147 ++++++++++++++++++ .../background_execution/test_live_bus.py | 79 ++++++++++ 2 files changed, 226 insertions(+) create mode 100644 src/backend/base/langflow/services/background_execution/live_bus.py create mode 100644 src/backend/tests/unit/services/background_execution/test_live_bus.py diff --git a/src/backend/base/langflow/services/background_execution/live_bus.py b/src/backend/base/langflow/services/background_execution/live_bus.py new file mode 100644 index 0000000000..84997ea6a6 --- /dev/null +++ b/src/backend/base/langflow/services/background_execution/live_bus.py @@ -0,0 +1,147 @@ +"""In-process live event bus for background workflow runs. + +Publishers (the runner) push ``LiveFrame``s for a job; live subscribers receive +them through a bounded per-subscriber queue. ``reattach`` first replays the +durable log (via the injected ``read_durable``) from the caller's last seen seq, +then switches to the live tail — deduplicating any seq already replayed so a +frame straddling the replay/tail boundary is emitted exactly once. + +This bus is process-local; the redis backend (Phase 3) swaps it for redis +Streams behind the same facade method (``events``). Nothing else changes. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from collections.abc import AsyncIterator, Awaitable, Callable +from dataclasses import dataclass + +# Per-subscriber buffer cap. A slow client falling behind a fast producer drops +# the oldest live frame rather than growing without bound; the durable log is +# the source of truth for anything the client missed, reachable via reattach. +_SUBSCRIBER_MAXSIZE = 1000 + + +@dataclass(frozen=True) +class LiveFrame: + """One framed event on the live bus. + + ``seq`` is the durable cursor used for ``Last-Event-ID``. For a durable + milestone it is the row's ``job_events.seq``; for an ephemeral frame it is + the seq of the most recent durable milestone (so ordering stays monotonic). + ``data`` is the already-formatted SSE frame bytes. + """ + + seq: int + data: bytes + + +ReadDurable = Callable[[int], Awaitable[list[LiveFrame]]] + +# Sentinel pushed into a subscriber queue to signal end-of-stream. +_CLOSED = object() + + +class InMemoryLiveBus: + def __init__(self) -> None: + # job_id -> set of subscriber queues. + self._subscribers: dict[str, set[asyncio.Queue]] = {} + # job_id -> True once closed; late subscribers end immediately. + self._closed: dict[str, bool] = {} + + def _new_queue(self, job_id: str) -> asyncio.Queue: + queue: asyncio.Queue = asyncio.Queue(maxsize=_SUBSCRIBER_MAXSIZE) + self._subscribers.setdefault(job_id, set()).add(queue) + return queue + + def _drop_queue(self, job_id: str, queue: asyncio.Queue) -> None: + subs = self._subscribers.get(job_id) + if subs is not None: + subs.discard(queue) + if not subs: + self._subscribers.pop(job_id, None) + + async def publish(self, job_id: str, frame: LiveFrame) -> None: + """Push ``frame`` to every live subscriber for ``job_id``.""" + for queue in list(self._subscribers.get(job_id, set())): + if queue.full(): + with contextlib.suppress(asyncio.QueueEmpty): + queue.get_nowait() + with contextlib.suppress(asyncio.QueueFull): + queue.put_nowait(frame) + + async def close(self, job_id: str) -> None: + """Mark the job's stream ended and wake every subscriber.""" + self._closed[job_id] = True + for queue in list(self._subscribers.get(job_id, set())): + with contextlib.suppress(asyncio.QueueFull): + queue.put_nowait(_CLOSED) + + def subscribe(self, job_id: str) -> AsyncIterator[LiveFrame]: + """Live-only subscription: frames published from now until close. + + The subscriber queue is registered eagerly here (not inside the + generator) so a frame published between ``subscribe(...)`` and the + first ``__anext__`` is not lost — the queue already exists to catch it. + """ + if self._closed.get(job_id): + + async def _empty() -> AsyncIterator[LiveFrame]: + return + yield # pragma: no cover - makes this an async generator + + return _empty() + + queue = self._new_queue(job_id) + + async def _gen() -> AsyncIterator[LiveFrame]: + try: + while True: + item = await queue.get() + if item is _CLOSED: + return + yield item + finally: + self._drop_queue(job_id, queue) + + return _gen() + + def reattach( + self, + job_id: str, + last_seq: int, + read_durable: ReadDurable, + ) -> AsyncIterator[LiveFrame]: + """Replay the durable log after ``last_seq``, then tail the live bus. + + A subscriber queue is registered BEFORE the durable read so no live + frame published during replay is lost. Frames whose seq was already + replayed (or already <= last_seq) are skipped so the boundary emits + each seq exactly once. + """ + # Register eagerly (not inside the generator) so frames published + # between ``reattach(...)`` and the first ``__anext__`` are captured. + queue = self._new_queue(job_id) + + async def _gen() -> AsyncIterator[LiveFrame]: + highest = last_seq + try: + for frame in await read_durable(last_seq): + highest = max(highest, frame.seq) + yield frame + # If the job closed before/while we replayed, drain nothing more. + if self._closed.get(job_id) and queue.empty(): + return + while True: + item = await queue.get() + if item is _CLOSED: + return + if item.seq <= highest: + continue + highest = item.seq + yield item + finally: + self._drop_queue(job_id, queue) + + return _gen() diff --git a/src/backend/tests/unit/services/background_execution/test_live_bus.py b/src/backend/tests/unit/services/background_execution/test_live_bus.py new file mode 100644 index 0000000000..8ce3de26b5 --- /dev/null +++ b/src/backend/tests/unit/services/background_execution/test_live_bus.py @@ -0,0 +1,79 @@ +"""In-memory live bus: publish to live subscribers, reattach replays then tails.""" + +from __future__ import annotations + +import asyncio + +import pytest +from langflow.services.background_execution.live_bus import InMemoryLiveBus, LiveFrame + + +async def test_subscriber_receives_published_frame(): + bus = InMemoryLiveBus() + sub = bus.subscribe("job-1") + await bus.publish("job-1", LiveFrame(seq=1, data=b"hello")) + frame = await asyncio.wait_for(sub.__anext__(), timeout=2) + assert frame.seq == 1 + assert frame.data == b"hello" + await bus.close("job-1") + + +async def test_close_ends_subscription(): + bus = InMemoryLiveBus() + sub = bus.subscribe("job-1") + await bus.close("job-1") + with pytest.raises(StopAsyncIteration): + await asyncio.wait_for(sub.__anext__(), timeout=2) + + +async def test_multiple_subscribers_each_get_frame(): + bus = InMemoryLiveBus() + sub_a = bus.subscribe("job-1") + sub_b = bus.subscribe("job-1") + await bus.publish("job-1", LiveFrame(seq=1, data=b"x")) + fa = await asyncio.wait_for(sub_a.__anext__(), timeout=2) + fb = await asyncio.wait_for(sub_b.__anext__(), timeout=2) + assert fa.data == fb.data == b"x" + await bus.close("job-1") + + +async def test_reattach_replays_durable_then_tails_live(): + bus = InMemoryLiveBus() + + # Fake durable store: seq->data already persisted before reattach. + persisted = [LiveFrame(seq=1, data=b"a"), LiveFrame(seq=2, data=b"b")] + + async def read_durable(after_seq: int) -> list[LiveFrame]: + return [f for f in persisted if f.seq > after_seq] + + # Reattach from seq 0: should replay seq 1,2 from durable, then live seq 3. + stream = bus.reattach("job-1", last_seq=0, read_durable=read_durable) + + f1 = await asyncio.wait_for(stream.__anext__(), timeout=2) + f2 = await asyncio.wait_for(stream.__anext__(), timeout=2) + assert [f1.seq, f2.seq] == [1, 2] + + # Now a live frame arrives; the tail picks it up. + await bus.publish("job-1", LiveFrame(seq=3, data=b"c")) + f3 = await asyncio.wait_for(stream.__anext__(), timeout=2) + assert f3.seq == 3 + await bus.close("job-1") + + +async def test_reattach_dedupes_overlap_between_durable_and_live(): + """A frame that lands in live while we replay durable must not double-emit.""" + bus = InMemoryLiveBus() + persisted = [LiveFrame(seq=1, data=b"a")] + + async def read_durable(after_seq: int) -> list[LiveFrame]: + return [f for f in persisted if f.seq > after_seq] + + stream = bus.reattach("job-1", last_seq=0, read_durable=read_durable) + f1 = await asyncio.wait_for(stream.__anext__(), timeout=2) + assert f1.seq == 1 + # Publish seq 1 again (already replayed) and seq 2 (new). Only seq 2 should pass. + await bus.publish("job-1", LiveFrame(seq=1, data=b"a")) + await bus.publish("job-1", LiveFrame(seq=2, data=b"b")) + f = await asyncio.wait_for(stream.__anext__(), timeout=2) + assert f.seq == 2 + await bus.close("job-1")