From 57bde56b35e8417a434b680cb3caf4d06806dd87 Mon Sep 17 00:00:00 2001 From: Alex Soffronow-Pagonidis Date: Tue, 28 Jul 2026 11:12:39 +0200 Subject: [PATCH] Graceful shutdown for supervised background tasks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Shutdown no longer cancels the work it promises to wait for. Setting a stop event and cancelling in the same breath meant the cancellation landed before the coroutine next checked the event, so the clean path never ran and an in-flight external-agents sweep was abandoned — leaving some bundles rendered from the new sources and some from the old. stop_background_task signals, waits a bounded moment, and cancels only if the task does not wind itself down, so cancellation is the timeout backstop rather than the mechanism. It lives in nerve.utils.aio rather than the gateway lifespan so a service can shut its own background task down without importing the FastAPI app. Co-Authored-By: Claude --- nerve/external_agents/sync_service.py | 34 +++++-- nerve/gateway/server.py | 7 +- nerve/utils/aio.py | 58 ++++++++++++ tests/test_external_agents_sync_service.py | 103 ++++++++++++++++++++- tests/test_utils_aio.py | 97 +++++++++++++++++++ 5 files changed, 283 insertions(+), 16 deletions(-) create mode 100644 nerve/utils/aio.py create mode 100644 tests/test_utils_aio.py diff --git a/nerve/external_agents/sync_service.py b/nerve/external_agents/sync_service.py index 92163f4b..ded81b71 100644 --- a/nerve/external_agents/sync_service.py +++ b/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(): diff --git a/nerve/gateway/server.py b/nerve/gateway/server.py index 427b3624..037b1afa 100644 --- a/nerve/gateway/server.py +++ b/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() diff --git a/nerve/utils/aio.py b/nerve/utils/aio.py new file mode 100644 index 00000000..42b32f13 --- /dev/null +++ b/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) diff --git a/tests/test_external_agents_sync_service.py b/tests/test_external_agents_sync_service.py index e29ec02e..40371455 100644 --- a/tests/test_external_agents_sync_service.py +++ b/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() diff --git a/tests/test_utils_aio.py b/tests/test_utils_aio.py new file mode 100644 index 00000000..3b69f93d --- /dev/null +++ b/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