fix: protect long-running background runs from eviction and release them on stop

Two related bugs in the in-memory _BACKGROUND_RUNS registry that re-attach
reads from. Both surfaced in the PR #13307 review; both have to land
together so the registry's eviction policy and its cleanup path agree.

A4: _register_background_run used to pop the oldest entry by insertion
order when the dict hit _MAX_BACKGROUND_RUNS. A long-running first job
got evicted by the 101st short job mid-run; re-attach returned 404
while the still-buffering task appended into an orphaned _BackgroundRun.
The new policy prefers evicting the oldest completed entry. If every
slot is still running, evict the oldest anyway to keep the registry
bounded and log a warning so the situation is visible.

A5: stop_workflow revoked the buffer task but left the _BackgroundRun
in the registry with done=False. Re-attach readers could hang on
_cond.wait() indefinitely (the task that would have called finish() was
cancelled mid-execution), and cancelled buffers occupied memory until
LRU evicted them. New _clear_background_run helper pops the entry and
calls await bg_run.finish() so waiters wake to a clean stream end.
stop_workflow calls it after revoke_task.

Four unit tests pin both behaviors: eviction prefers completed entries,
eviction falls back to the oldest when every run is active, clear pops
and finishes the buffer, and clear is a no-op for unknown job ids. All
use real _BackgroundRun instances via monkeypatch on the module-level
dict; no mocks.

Author: ogabrielluiz

src/backend/base/langflow/api/v2/workflow.py

@@ -35,6 +35,7 @@
 from fastapi.sse import format_sse_event
 from lfx.events.event_manager import create_default_event_manager
 from lfx.graph.graph.base import Graph
+from lfx.log.logger import logger
 from lfx.schema.schema import InputValueRequest
 from lfx.schema.workflow import (
     WORKFLOW_EXECUTION_RESPONSES,
@@ -691,11 +692,41 @@ async def replay(self, start_index: int) -> AsyncIterator[bytes]:
 _BACKGROUND_RUNS: dict[str, _BackgroundRun] = {}
 
 
+async def _clear_background_run(job_id: str) -> None:
+    """Pop the background run registry entry and wake any re-attach waiters.
+
+    Called from ``stop_workflow`` after revoking the buffer task. Without
+    this, ``_cond.wait()`` inside ``_BackgroundRun.replay`` can hang
+    indefinitely (the buffer is never marked done because the task that
+    would have called ``finish()`` was cancelled mid-execution), and the
+    cancelled ``_BackgroundRun`` lingers in memory until LRU evicts it.
+    """
+    bg_run = _BACKGROUND_RUNS.pop(job_id, None)
+    if bg_run is not None:
+        await bg_run.finish()
+
+
 def _register_background_run(job_id: str, bg_run: _BackgroundRun) -> None:
-    """Register a background run, evicting the oldest entry when full."""
+    """Register a background run, evicting a completed entry when full.
+
+    Prefer evicting the oldest *completed* run so a long-running job's
+    re-attach handle survives. If every slot is still active, evict the
+    oldest one anyway to keep the registry bounded, and log a warning.
+    """
     if len(_BACKGROUND_RUNS) >= _MAX_BACKGROUND_RUNS:
-        oldest = next(iter(_BACKGROUND_RUNS))
-        _BACKGROUND_RUNS.pop(oldest, None)
+        evict_key = next(
+            (key for key, run in _BACKGROUND_RUNS.items() if run.done),
+            None,
+        )
+        if evict_key is None:
+            evict_key = next(iter(_BACKGROUND_RUNS))
+            logger.warning(
+                "Background run registry full with no completed entries; "
+                "evicting still-running job %s to make room for %s",
+                evict_key,
+                job_id,
+            )
+        _BACKGROUND_RUNS.pop(evict_key, None)
     _BACKGROUND_RUNS[job_id] = bg_run
 
 
@@ -1050,6 +1081,9 @@ async def stop_workflow(
     try:
         revoked = await task_service.revoke_task(job_id)
         await job_service.update_job_status(job_id, JobStatus.CANCELLED)
+        # Release the in-memory buffer and wake any re-attach waiters so they
+        # see a clean stream-end instead of hanging on ``_cond.wait()``.
+        await _clear_background_run(str(job_id))
 
         message = f"Job {job_id} cancelled successfully." if revoked else f"Job {job_id} is already cancelled."
         return WorkflowStopResponse(job_id=str(job_id), message=message)

src/backend/tests/unit/api/v2/test_workflow_agui.py

@@ -968,3 +968,96 @@ async def test_background_mode_with_agui_protocol_still_buffers_agui_frames(
         text = events.text
         assert "RUN_STARTED" in text
         assert "RUN_FINISHED" in text
+
+
+class TestBackgroundRunsRegistryEviction:
+    """``_register_background_run`` must not evict still-running buffers.
+
+    The registry is bounded by ``_MAX_BACKGROUND_RUNS``. The naive policy of
+    popping the oldest entry by insertion order drops the re-attach handle for
+    a long-running first job the moment the limit is hit, so the buffer task
+    keeps appending into an orphaned ``_BackgroundRun`` while re-attach
+    returns 404. Eviction must prefer completed entries first; falling back
+    to evicting the oldest only when every slot is occupied by a running run.
+    """
+
+    def test_eviction_prefers_completed_runs_over_running_ones(self, monkeypatch):
+        from langflow.api.v2 import workflow as workflow_module
+
+        monkeypatch.setattr(workflow_module, "_MAX_BACKGROUND_RUNS", 3)
+        monkeypatch.setattr(workflow_module, "_BACKGROUND_RUNS", {})
+
+        long_running = workflow_module._BackgroundRun(user_id="u", protocol_name="agui")
+        # done stays False; this is the run we must protect.
+        workflow_module._register_background_run("long", long_running)
+
+        # Fill the rest with completed runs.
+        for job_id in ("done1", "done2"):
+            done_run = workflow_module._BackgroundRun(user_id="u", protocol_name="agui")
+            done_run.done = True
+            workflow_module._register_background_run(job_id, done_run)
+
+        # Registry is now at the cap (3): [long, done1, done2]. Adding a new
+        # entry must evict a completed run, not the still-running ``long``.
+        new_run = workflow_module._BackgroundRun(user_id="u", protocol_name="agui")
+        workflow_module._register_background_run("new", new_run)
+
+        assert "long" in workflow_module._BACKGROUND_RUNS, (
+            "Still-running background run was evicted in favor of a completed one"
+        )
+        assert "new" in workflow_module._BACKGROUND_RUNS
+
+    def test_eviction_falls_back_to_oldest_when_every_run_is_active(self, monkeypatch):
+        """If every slot is occupied by a still-running run, evict the oldest anyway.
+
+        Unbounded growth would leak memory. The fallback is intentional and
+        documented; a warning log makes the situation visible.
+        """
+        from langflow.api.v2 import workflow as workflow_module
+
+        monkeypatch.setattr(workflow_module, "_MAX_BACKGROUND_RUNS", 2)
+        monkeypatch.setattr(workflow_module, "_BACKGROUND_RUNS", {})
+
+        for job_id in ("a", "b"):
+            run = workflow_module._BackgroundRun(user_id="u", protocol_name="agui")
+            workflow_module._register_background_run(job_id, run)
+
+        # All running; adding a third must evict the oldest (a).
+        third = workflow_module._BackgroundRun(user_id="u", protocol_name="agui")
+        workflow_module._register_background_run("c", third)
+
+        assert "a" not in workflow_module._BACKGROUND_RUNS
+        assert "b" in workflow_module._BACKGROUND_RUNS
+        assert "c" in workflow_module._BACKGROUND_RUNS
+
+
+class TestClearBackgroundRun:
+    """``_clear_background_run`` releases a stopped run's buffer and wakes waiters.
+
+    Without this, ``stop_workflow`` revokes the buffer task but leaves the
+    ``_BackgroundRun`` registered. Re-attach readers can hang on
+    ``_cond.wait()`` indefinitely, and up to ``_MAX_BACKGROUND_RUNS`` cancelled
+    buffers occupy memory.
+    """
+
+    async def test_clear_pops_registry_entry_and_finishes_buffer(self, monkeypatch):
+        from langflow.api.v2 import workflow as workflow_module
+
+        monkeypatch.setattr(workflow_module, "_BACKGROUND_RUNS", {})
+
+        bg_run = workflow_module._BackgroundRun(user_id="u", protocol_name="agui")
+        workflow_module._register_background_run("job-1", bg_run)
+        assert bg_run.done is False
+
+        await workflow_module._clear_background_run("job-1")
+
+        assert "job-1" not in workflow_module._BACKGROUND_RUNS
+        assert bg_run.done is True
+
+    async def test_clear_is_a_noop_for_unknown_job_id(self, monkeypatch):
+        from langflow.api.v2 import workflow as workflow_module
+
+        monkeypatch.setattr(workflow_module, "_BACKGROUND_RUNS", {})
+
+        # Must not raise even when nothing is registered.
+        await workflow_module._clear_background_run("nope")