Graceful shutdown for supervised background tasks (6/13)

nerve/external_agents/sync_service.py

@@ -33,9 +33,15 @@
 from nerve.external_agents.registry import AGENT_REGISTRY, FileTarget
 from nerve.external_agents.renderers import get_renderer
 from nerve.external_agents.writer import ConfigWriter
+from nerve.utils.aio import GRACEFUL_STOP_SECONDS, stop_background_task
 
 logger = logging.getLogger(__name__)
 
+# Floor on the gap between sweeps. ``sync_interval_minutes`` is user-supplied
+# and unvalidated, and a 0 there would turn the loop into a hot spin re-reading
+# and re-hashing every target's source files.
+_MIN_SWEEP_INTERVAL_SECONDS = 60
+
 
 @dataclass
 class FileSyncStatus:
@@ -68,7 +74,8 @@ class SyncService:
     Lifecycle:
 
     - :meth:`start` spawns the background loop. Returns immediately.
-    - :meth:`stop` cancels the loop and waits for it to finish.
+    - :meth:`stop` asks the loop to finish the sweep it is in, then waits
+      for it.
     - :meth:`run_once` does one sweep synchronously — used by the
       manual ``/api/external-agents/sync`` route and by tests.
 
@@ -113,22 +120,29 @@ async def start(self) -> None:
             self._loop(), name="external-agents-sync-loop",
         )
 
-    async def stop(self) -> None:
-        """Stop the background loop and wait for the in-flight sweep."""
-        self._stop_event.set()
+    async def stop(self, timeout: float = GRACEFUL_STOP_SECONDS) -> None:
+        """Stop the background loop and wait for the in-flight sweep.
+
+        The loop only looks at ``_stop_event`` between sweeps, so cancelling
+        alongside setting it would deliver the CancelledError first and abandon
+        a sweep partway down its target list — some bundles rendered from the
+        current sources, some still from the previous ones, and no record of
+        which. Signal, then let it reach the boundary; cancellation is the
+        backstop for a sweep that wedges, not the mechanism.
+        """
         task = self._task
         if task is None:
+            # Never started, or already stopped.
+            self._stop_event.set()
             return
-        task.cancel()
-        try:
-            await task
-        except (asyncio.CancelledError, Exception):
-            pass
+        await stop_background_task(
+            task, self._stop_event, "External-agents sync", timeout,
+        )
         self._task = None
 
     async def _loop(self) -> None:
         interval = max(
-            60,
+            _MIN_SWEEP_INTERVAL_SECONDS,
             self._config.external_agents.sync_interval_minutes * 60,
         )
         while not self._stop_event.is_set():

nerve/gateway/server.py

@@ -610,9 +610,10 @@ async def _periodic_backup():
             logger.warning("Codex thread sync shutdown raised: %s", e)
         _codex_thread_sync = None
 
-    # Stop the external-agents sync service. Cheap — it just cancels
-    # the periodic loop; no per-file cleanup needed because every write
-    # is already atomic (temp + rename).
+    # Stop the external-agents sync service. It exits through its own stop event
+    # so a sweep in flight finishes the whole target set rather than stopping
+    # partway down it; cancellation is the backstop. Individual writes need no
+    # cleanup — each is atomic (temp + rename).
     if _external_agents_sync is not None:
         try:
             await _external_agents_sync.stop()

nerve/utils/aio.py

@@ -0,0 +1,58 @@
+"""Asyncio helpers for services that run supervised background tasks.
+
+Lives here rather than in ``nerve.gateway.server`` so a service can shut its
+own background task down the same way the lifespan does without importing the
+FastAPI app.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+
+logger = logging.getLogger(__name__)
+
+# How long shutdown gives a background task to stop on its own after being
+# asked, before falling back to cancellation. Sized for one external-agents
+# sweep to finish, and short enough that a wedged task cannot hold the process
+# open for a user watching it exit.
+GRACEFUL_STOP_SECONDS = 5.0
+
+
+async def stop_background_task(
+    task: asyncio.Task,
+    stop_event: asyncio.Event | None,
+    name: str,
+    timeout: float = GRACEFUL_STOP_SECONDS,
+) -> None:
+    """Stop a background task via its stop event, cancelling only if it hangs.
+
+    Setting the event and cancelling in the same breath makes the event
+    decorative: the CancelledError is delivered before the task next gets to
+    look at it, so the task dies wherever it happened to be suspended instead
+    of at the end of the cycle it was in. The external-agents sync can be
+    halfway through a sweep at that moment, leaving some bundles rendered from
+    the new sources and some from the old. So signal, wait a bounded moment for
+    the task to wind itself down, and only cancel if it doesn't.
+    """
+    if stop_event is not None:
+        stop_event.set()
+    else:
+        # No cooperative signal to give — cancellation is the only lever.
+        task.cancel()
+    try:
+        # wait_for cancels the task itself once the timeout expires and awaits
+        # the cancellation, so the backstop needs no separate cancel() and the
+        # task is never left pending on either path.
+        await asyncio.wait_for(task, timeout)
+    except asyncio.TimeoutError:
+        logger.warning(
+            "%s did not stop within %.1fs of being asked; cancelled it",
+            name, timeout,
+        )
+    except asyncio.CancelledError:
+        # The task's cancellation, not ours: either we asked for it above or
+        # something else already had. Nothing left to wait for.
+        pass
+    except Exception as e:
+        logger.warning("%s raised while shutting down: %s", name, e)

tests/test_external_agents_sync_service.py

@@ -1,12 +1,15 @@
 """Tests for the periodic external-agents sync service.
 
-Covers the core loop semantics (idempotent on second sweep, picks up
-source changes on third sweep, isolates per-agent failures) without
-spinning up a real asyncio loop — we drive ``run_once`` directly.
+Covers the core sweep semantics (idempotent on second sweep, picks up
+source changes on third sweep, isolates per-agent failures) by driving
+``run_once`` directly, plus how the background loop shuts down — the one
+part that needs a real task on a real event loop.
 """
 
 from __future__ import annotations
 
+import asyncio
+import logging
 from pathlib import Path
 from unittest.mock import patch
 
@@ -17,6 +20,7 @@
     ExternalAgentTargetConfig,
     NerveConfig,
 )
+from nerve.external_agents import sync_service as sync_service_module
 from nerve.external_agents.sync_service import SyncService
 
 
@@ -165,3 +169,96 @@ async def test_status_for_api_is_serializable(
     json.dumps(payload)
     assert "codex" in payload
     assert payload["codex"]["name"] == "codex"
+
+
+def _sweep_on_a_test_timescale(
+    monkeypatch: pytest.MonkeyPatch, config: NerveConfig, seconds: float,
+) -> None:
+    """Shrink the gap between sweeps to something a test can wait for.
+
+    Both halves of ``max(floor, configured)`` have to move, or the untouched
+    one keeps the loop parked for the best part of an hour.
+    """
+    config.external_agents.sync_interval_minutes = 0
+    monkeypatch.setattr(
+        sync_service_module, "_MIN_SWEEP_INTERVAL_SECONDS", seconds,
+    )
+
+
+class TestShutdownWaitsForTheSweep:
+    """How the gateway takes the sync service down at shutdown.
+
+    The stop event is the mechanism and cancellation is the backstop, not the
+    other way around: the loop only looks at the event between sweeps, so a
+    task cancelled where it stands can be partway down its target list, with
+    some bundles rendered from the current sources and some still from the
+    previous ones.
+    """
+
+    @pytest.mark.asyncio
+    async def test_an_in_flight_sweep_is_allowed_to_finish(
+        self, codex_config: NerveConfig, monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        _sweep_on_a_test_timescale(monkeypatch, codex_config, 0.05)
+        svc = SyncService(codex_config)
+        started, finished = asyncio.Event(), []
+
+        async def slow_sweep():
+            started.set()
+            await asyncio.sleep(0.2)  # stands in for rendering + writing
+            finished.append(1)
+            return {}
+
+        monkeypatch.setattr(svc, "run_once", slow_sweep)
+        await svc.start()  # consumes the eager first sweep
+        started.clear()
+        finished.clear()
+        # Now wait for a sweep the *loop* starts, so stop() lands mid-sweep.
+        await asyncio.wait_for(started.wait(), timeout=2)
+        task = svc._task
+
+        await svc.stop()
+
+        assert finished == [1], "stop() abandoned the sweep it waits for"
+        assert task.done() and not task.cancelled()  # left by its own exit
+        assert svc._task is None
+
+    @pytest.mark.asyncio
+    async def test_a_wedged_sweep_does_not_hold_shutdown_open(
+        self, codex_config: NerveConfig, monkeypatch: pytest.MonkeyPatch,
+        caplog,
+    ) -> None:
+        """The backstop still exists: one renderer blocking forever must not
+        keep the process alive, and giving up should say so out loud."""
+        _sweep_on_a_test_timescale(monkeypatch, codex_config, 0.05)
+        svc = SyncService(codex_config)
+        wedged = asyncio.Event()
+        sweeps = 0
+
+        async def wedges_after_the_first_sweep():
+            nonlocal sweeps
+            sweeps += 1
+            if sweeps > 1:  # let start()'s eager sweep through
+                wedged.set()
+                while True:
+                    await asyncio.sleep(0.01)
+            return {}
+
+        monkeypatch.setattr(svc, "run_once", wedges_after_the_first_sweep)
+        await svc.start()
+        await asyncio.wait_for(wedged.wait(), timeout=2)
+        task = svc._task
+
+        with caplog.at_level(logging.WARNING, logger="nerve.utils.aio"):
+            await svc.stop(timeout=0.2)
+
+        assert task.cancelled()
+        assert "did not stop within" in caplog.text
+        assert svc._task is None
+
+    @pytest.mark.asyncio
+    async def test_stop_without_start_is_harmless(
+        self, codex_config: NerveConfig,
+    ) -> None:
+        """Lifespan shutdown runs even when startup failed before start()."""
+        await SyncService(codex_config).stop()

tests/test_utils_aio.py

@@ -0,0 +1,97 @@
+"""Graceful shutdown of supervised background tasks (nerve.utils.aio)."""
+
+import asyncio
+import logging
+
+import pytest
+
+from nerve.utils.aio import stop_background_task
+
+
+async def _cycling_worker(
+    stop: asyncio.Event,
+    started: asyncio.Event,
+    finished: list[int],
+    cycle_seconds: float = 0.2,
+) -> None:
+    """A background task shaped like the real ones: a loop whose body is a
+    critical section that must not be abandoned partway through."""
+    while not stop.is_set():
+        started.set()
+        await asyncio.sleep(cycle_seconds)  # stands in for the real work
+        finished.append(1)
+
+
+class TestStopBackgroundTask:
+    """The stop event is the mechanism and cancellation is the backstop, not
+    the other way around: a task cancelled where it stands can be halfway
+    through a unit of work that has to be all-or-nothing."""
+
+    @pytest.mark.asyncio
+    async def test_the_stop_event_alone_ends_the_task(self):
+        """No cancellation involved anywhere: signalling has to be enough on
+        its own. If it isn't, the event is decorative and shutdown is really
+        just killing the task wherever it stands."""
+        stop, started, finished = asyncio.Event(), asyncio.Event(), []
+        task = asyncio.create_task(
+            _cycling_worker(stop, started, finished, cycle_seconds=0.01)
+        )
+        await asyncio.wait_for(started.wait(), timeout=2)
+
+        await stop_background_task(task, stop, "Test worker")
+        assert task.done()
+        assert not task.cancelled()  # it left through its own exit
+
+    @pytest.mark.asyncio
+    async def test_a_cycle_in_flight_is_allowed_to_finish(self):
+        """The reason the stop event has to come first. Shutdown lands while a
+        cycle is mid-flight; cancelling there abandons the work half-done, so
+        the cycle gets to complete."""
+        stop, started, finished = asyncio.Event(), asyncio.Event(), []
+        task = asyncio.create_task(_cycling_worker(stop, started, finished))
+        await asyncio.wait_for(started.wait(), timeout=2)
+
+        await stop_background_task(task, stop, "Test worker")
+        assert finished == [1]
+        assert not task.cancelled()
+
+    @pytest.mark.asyncio
+    async def test_a_task_that_will_not_stop_is_cancelled_after_the_timeout(
+        self, caplog
+    ):
+        """The backstop still exists: a wedged task must not hold shutdown open
+        indefinitely, and giving up on one should say so out loud."""
+        async def deaf_to_the_event():
+            while True:
+                await asyncio.sleep(0.01)
+
+        task = asyncio.create_task(deaf_to_the_event())
+        with caplog.at_level(logging.WARNING, logger="nerve.utils.aio"):
+            await stop_background_task(
+                task, asyncio.Event(), "Test worker", timeout=0.2
+            )
+        assert task.cancelled()
+        assert "did not stop within" in caplog.text
+
+    @pytest.mark.asyncio
+    async def test_without_a_stop_event_cancellation_is_the_only_lever(self):
+        """Setup can fail after the task exists but before the event does."""
+        async def forever():
+            while True:
+                await asyncio.sleep(0.01)
+
+        task = asyncio.create_task(forever())
+        await stop_background_task(task, None, "Test worker")
+        assert task.cancelled()
+
+    @pytest.mark.asyncio
+    async def test_a_task_that_already_died_is_logged_not_raised(self, caplog):
+        """Shutdown continues past a background task that crashed earlier —
+        there is a whole teardown sequence after this call."""
+        async def boom():
+            raise RuntimeError("worker died hours ago")
+
+        task = asyncio.create_task(boom())
+        with caplog.at_level(logging.WARNING, logger="nerve.utils.aio"):
+            await stop_background_task(task, asyncio.Event(), "Test worker")
+        assert "worker died hours ago" in caplog.text