diff --git a/headroom/proxy/server.py b/headroom/proxy/server.py index 1b58db04e..c26160a4a 100644 --- a/headroom/proxy/server.py +++ b/headroom/proxy/server.py @@ -2553,6 +2553,7 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: app.state.started_at = time.time() app.state.ready = False app.state.startup_error = None + app.state.periodic_toin_stats_task = None try: try: @@ -2560,7 +2561,9 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: # Startup await proxy.startup() if config.periodic_toin_stats_enabled: - asyncio.create_task(_log_toin_stats_periodically()) + app.state.periodic_toin_stats_task = asyncio.create_task( + _log_toin_stats_periodically() + ) if proxy.usage_reporter: await proxy.usage_reporter.start(proxy) if proxy.traffic_learner: @@ -2610,6 +2613,16 @@ def create_app(config: ProxyConfig | None = None) -> FastAPI: exc, ) + periodic_toin_stats_task = app.state.periodic_toin_stats_task + if periodic_toin_stats_task is not None: + periodic_toin_stats_task.cancel() + await _timed( + asyncio.gather(periodic_toin_stats_task, return_exceptions=True), + label="periodic_toin_stats.stop", + timeout=3.0, + ) + app.state.periodic_toin_stats_task = None + if _cc_reconciler is not None: await _timed(_cc_reconciler.stop(), label="cc_reconciler.stop", timeout=3.0) if _beacon_is_owner[0]: diff --git a/tests/test_proxy_telemetry_env.py b/tests/test_proxy_telemetry_env.py index e3ffd65a6..0dc51c117 100644 --- a/tests/test_proxy_telemetry_env.py +++ b/tests/test_proxy_telemetry_env.py @@ -98,3 +98,33 @@ class TestProxyPeriodicTOINStatsEnv: pass assert requested is True + + def test_lifespan_cancels_periodic_toin_stats_on_shutdown(self, monkeypatch): + """Shutdown cancels and awaits the periodic TOIN stats task.""" + monkeypatch.setenv("HEADROOM_SKIP_UPSTREAM_CHECK", "1") + + async def hold_periodic_stats_task(): + await asyncio.Event().wait() + + monkeypatch.setattr( + "headroom.proxy.server._log_toin_stats_periodically", + hold_periodic_stats_task, + ) + + app = create_app( + ProxyConfig( + optimize=False, + cache_enabled=False, + rate_limit_enabled=False, + cost_tracking_enabled=False, + periodic_toin_stats_enabled=True, + ) + ) + + with TestClient(app): + task = app.state.periodic_toin_stats_task + assert task is not None + assert not task.done() + + assert task.cancelled() + assert app.state.periodic_toin_stats_task is None