Let the supervisor stop one flow's tasks and clear its quarantine

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01StpRc2C6au1WJ1EUU7fsfu

backend/fluksio/flow/supervision.py

@@ -44,6 +44,10 @@ class Supervisor:
     def __init__(self, events: EventBus | None = None) -> None:
         self._events = events
         self._tasks: dict[str, asyncio.Task[None]] = {}
+        #: Which flow each task belongs to. Recorded rather than read off the
+        #: task's name, because names are a flow and a node joined by a dot and
+        #: matching on that prefix would let 'hea' cancel 'heating'.
+        self._flows: dict[str, str] = {}
         self._failures: dict[str, deque[float]] = {}
         self.quarantined: set[str] = set()
 
@@ -51,6 +55,7 @@ class Supervisor:
         """Run `factory()` and keep running it until told to stop."""
         if name in self._tasks:
             return
+        self._flows[name] = flow
         self._tasks[name] = asyncio.create_task(
             self._supervise(name, flow, factory), name=f"supervised:{name}"
         )
@@ -116,6 +121,29 @@ class Supervisor:
         """Stop supervising. Idempotent, and safe to call mid-restart."""
         tasks = list(self._tasks.values())
         self._tasks.clear()
+        self._flows.clear()
+        await self._cancel(tasks)
+
+    async def cancel_flow(self, flow: str) -> None:
+        """Stop one flow's supervised tasks and give it a clean slate.
+
+        Rebuilding the whole pipeline throws the supervisor away and builds
+        another, so a flow quarantined by the last build gets another chance.
+        Rebuilding one flow has to say the same thing about that flow alone, or
+        every other flow's quarantine would go with it.
+        """
+        names = [name for name, owner in self._flows.items() if owner == flow]
+        tasks = [self._tasks.pop(name) for name in names if name in self._tasks]
+        for name in names:
+            del self._flows[name]
+        await self._cancel(tasks)
+        # The clean slate: whatever this flow spent before, the build that
+        # follows starts its budget again.
+        self.quarantined.discard(flow)
+        self._failures.pop(flow, None)
+
+    async def _cancel(self, tasks: list[asyncio.Task[None]]) -> None:
+        """Ask these tasks to stop, and wait no longer than the grace period."""
         for task in tasks:
             task.cancel()
         if not tasks:

backend/tests/flow/test_supervision.py

@@ -124,3 +124,49 @@ def test_cancel_all_is_idempotent():
         await supervisor.cancel_all()
 
     asyncio.run(scenario())
+
+
+def test_cancelling_one_flow_leaves_another_flows_tasks_running():
+    """Rebuilding one flow must not take every other flow's loops with it."""
+
+    async def loop() -> None:
+        await asyncio.sleep(3600)
+
+    async def scenario() -> None:
+        supervisor = Supervisor()
+        supervisor.spawn("gone.sub", "gone", loop)
+        supervisor.spawn("fine.sub", "fine", loop)
+        await _settle()
+
+        survivor = supervisor._tasks["fine.sub"]
+        await supervisor.cancel_flow("gone")
+
+        assert list(supervisor._tasks) == ["fine.sub"]
+        assert not survivor.done()
+        await supervisor.cancel_all()
+
+    asyncio.run(scenario())
+
+
+def test_cancelling_one_flow_lifts_only_its_quarantine():
+    async def failing() -> None:
+        raise RuntimeError("nope")
+
+    async def scenario() -> None:
+        supervisor = Supervisor()
+        supervisor.spawn("a", "broken", failing)
+        supervisor.spawn("b", "also-broken", failing)
+        await _settle()
+
+        assert supervisor.quarantined == {"broken", "also-broken"}
+
+        await supervisor.cancel_flow("broken")
+
+        assert supervisor.quarantined == {"also-broken"}
+        # And its budget, or the flow would be quarantined again on the first
+        # crash after the rebuild that was meant to fix it.
+        assert "broken" not in supervisor._failures
+        assert "also-broken" in supervisor._failures
+        await supervisor.cancel_all()
+
+    asyncio.run(scenario())