diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index 07af4ef9d7..c33f64f960 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -4046,7 +4046,9 @@ class APIServerAdapter(OpenAICompatRoutesMixin, BasePlatformAdapter): self._activate_admitted_request() self._inflight_agent_runs += 1 try: - return await loop.run_in_executor(None, _run) + # Worker-scoped count rides along so the shutdown close gate still sees the thread + # after this handler task is cancelled (#116535); released in the worker's finally. + return await loop.run_in_executor(None, _api_runs._track_api_worker(_run)) finally: self._inflight_agent_runs -= 1 diff --git a/gateway/platforms/api_server_runs.py b/gateway/platforms/api_server_runs.py index 3a348c0cde..64b87582d5 100644 --- a/gateway/platforms/api_server_runs.py +++ b/gateway/platforms/api_server_runs.py @@ -5,6 +5,7 @@ import hashlib import json import logging import os +import threading import time import uuid from contextlib import suppress @@ -27,6 +28,43 @@ from gateway.platforms.api_server_run_idempotency import TERMINAL_STATUSES logger = logging.getLogger("gateway.platforms.api_server") + +# Executor-thread lifetime for API-server agent turns (#116535). The handler-side +# ``_inflight_agent_runs`` count in api_server.py drops in the handler's ``finally`` when the +# handler task is cancelled (disconnect/shutdown) while the executor thread behind +# ``run_in_executor`` keeps running -- so the shutdown close gate cannot rely on it. This +# module-level count is taken on the submitting thread before ``run_in_executor`` and released +# in the worker thread's own ``finally``, which also keeps it visible past ``adapters.clear()``. +_API_WORKER_LOCK = threading.Lock() +_API_WORKER_LIVE = 0 + + +def api_worker_live_count() -> int: + """Executor threads still inside an API-server agent turn (``_run_agent`` and ``/v1/runs``).""" + with _API_WORKER_LOCK: + return _API_WORKER_LIVE + + +def _track_api_worker(fn): + """Hold the worker-lifetime count for one ``run_in_executor`` submission. + + Increment on the submitting (handler) thread so the count is live before the worker can + exit; decrement in the worker thread's own ``finally`` so handler cancellation cannot drop + it early. ``fn`` still runs entirely on the worker. + """ + global _API_WORKER_LIVE + with _API_WORKER_LOCK: + _API_WORKER_LIVE += 1 + + def _counted(): + try: + return fn() + finally: + global _API_WORKER_LIVE + with _API_WORKER_LOCK: + _API_WORKER_LIVE -= 1 + + return _counted _ROOM_RETENTION_REQUEST_KEY = ( RequestKey("hermes.room_run_retention_until", float) if RequestKey is not None else "hermes.room_run_retention_until") @@ -868,7 +906,8 @@ async def _execute_run(self, run: _RunLaunch, *, _api_server) -> None: self._active_run_agents[run_id] = agent approval_notify = _make_approval_notify(self, run, _api_server=_api_server) result, usage, served_runtime = await loop.run_in_executor( - None, lambda: _run_agent_sync(self, run, agent, approval_notify, _api_server=_api_server)) + None, _track_api_worker( + lambda: _run_agent_sync(self, run, agent, approval_notify, _api_server=_api_server))) if not isinstance(result, dict): result = {} status, fields = terminal_run_status(result) diff --git a/gateway/run_shutdown.py b/gateway/run_shutdown.py index a4f8bca84f..c778aa10b9 100644 --- a/gateway/run_shutdown.py +++ b/gateway/run_shutdown.py @@ -224,6 +224,20 @@ class GatewayShutdownMixin: except Exception: return 0 + def _active_api_worker_count(self) -> int: + """API-server executor threads still inside an agent turn (#116535). + + Module-level, like the cron registry above: the handler-side adapter count is already + unreachable here (``adapters`` was cleared a phase earlier) and, worse, drops on handler + cancellation while the worker thread lives on. Read live at the close gate instead of + snapshotting. + """ + try: + from gateway.platforms.api_server_runs import api_worker_live_count + return max(0, int(api_worker_live_count())) + except Exception: + return 0 + def _interrupt_api_server_runs(self, reason: str) -> int: """Interrupt API-server agents not in ``_running_agents`` (same set ``_active_api_run_count`` counts).""" try: @@ -1977,13 +1991,18 @@ class GatewayShutdownMixin: # outlived the drain is mid-write for the same #101093 reasons; the drain already spent its # budget, so no second wait — leave the handles open (#102198). The API count is the snapshot # taken before the adapters were released; a run whose handler task was cancelled at disconnect - # has already left it, so this term under-counts rather than over-counts. - _cron_live, _api_live, _deferred_live = self._active_cron_job_count(), ctx.api_live, ctx.deferred_count() - if _cron_live or _api_live or _deferred_live: + # has already left it, so that term under-counts — the live worker-scoped count below covers + # the cancelled-handler case (#116535). + _cron_live = self._active_cron_job_count() + _api_live = ctx.api_live + _api_worker_live = self._active_api_worker_count() + _deferred_live = ctx.deferred_count() + if _cron_live or _api_live or _api_worker_live or _deferred_live: logger.warning( - "Shutdown phase: %d cron job(s) / %d API-server run(s) / %d deferred worker(s) still running " - "after the executor quiesce — skipping the SessionDB close/checkpoint, leaving state.db open " - "for the live writer (#102198)", _cron_live, _api_live, _deferred_live, + "Shutdown phase: %d cron job(s) / %d API-server run(s) / %d API-server worker(s) / " + "%d deferred worker(s) still running after the executor quiesce — skipping the SessionDB " + "close/checkpoint, leaving state.db open for the live writer (#102198, #116535)", + _cron_live, _api_live, _api_worker_live, _deferred_live, ) return _step = GatewayShutdownMixin._quiet_step diff --git a/tests/gateway/test_shutdown_cache_cleanup.py b/tests/gateway/test_shutdown_cache_cleanup.py index 4409541bc0..b9e029a690 100644 --- a/tests/gateway/test_shutdown_cache_cleanup.py +++ b/tests/gateway/test_shutdown_cache_cleanup.py @@ -69,6 +69,11 @@ class _FakeGateway: # No API server adapter -> no durable runs to stamp with the drain boundary (#115133). return 0 + def _active_api_worker_count(self): + # Worker-scoped API count the SessionDB close gate reads live (#116535). + # This fake runs no executor turns, so it is always idle. + return 0 + def _update_runtime_status(self, *_a, **_kw): pass diff --git a/tests/gateway/test_shutdown_executor_quiesce.py b/tests/gateway/test_shutdown_executor_quiesce.py index d37dc2d315..abb7fbd1d9 100644 --- a/tests/gateway/test_shutdown_executor_quiesce.py +++ b/tests/gateway/test_shutdown_executor_quiesce.py @@ -18,6 +18,7 @@ import concurrent.futures import threading import time from collections import OrderedDict +from contextlib import suppress import pytest @@ -82,6 +83,7 @@ class _FakeGateway: _api_server_hook = gw_mod.GatewayShutdownMixin._api_server_hook _mark_api_runs_shutdown_requested = gw_mod.GatewayShutdownMixin._mark_api_runs_shutdown_requested _active_api_run_count = gw_mod.GatewayShutdownMixin._active_api_run_count + _active_api_worker_count = gw_mod.GatewayShutdownMixin._active_api_worker_count _active_deferred_agent_worker_count = gw_mod.GatewayShutdownMixin._active_deferred_agent_worker_count def _update_runtime_status(self, *_a, **_kw): @@ -268,6 +270,65 @@ async def test_live_writer_outside_the_executor_skips_the_session_db_close(monke assert gw._executor_closing is True, "executor left unsealed on the outside-writer path" +@pytest.mark.asyncio +async def test_cancelled_api_handler_worker_still_blocks_session_db_close(monkeypatch): + """A cancelled request handler must not let state.db be closed under its live worker (#116535). + + The handler-side ``_inflight_agent_runs`` count drops in the handler's ``finally`` when the + handler task is cancelled, while the executor thread behind ``run_in_executor`` is still + blocked in the turn. Only the worker-scoped count still sees that thread, so the SessionDB + close gate must consult it alongside the handler snapshot -- and must not close until the + worker itself exits. + """ + import hermes_state_registry + from gateway.platforms import api_server_runs as api_runs + + events = [] + gw = _FakeGateway(events) + monkeypatch.setattr( + hermes_state_registry, "close_all", lambda: events.append("close_all") or 0 + ) + + release = threading.Event() + worker_started = threading.Event() + + def _blocked_turn(): + worker_started.set() + assert release.wait(5.0), "test tore down while the worker was still blocked" + events.append("worker_done") + + loop = asyncio.get_running_loop() + + async def _handler(): + # Same shape as the api_server call sites: the worker-scoped count is taken before + # run_in_executor and released in the worker's own finally, so cancelling this task + # drops the handler side while the thread keeps holding the worker side. + return await loop.run_in_executor(None, api_runs._track_api_worker(_blocked_turn)) + + task = asyncio.ensure_future(_handler()) + assert await loop.run_in_executor(None, worker_started.wait, 5.0), "worker never started" + task.cancel() + with suppress(asyncio.CancelledError): + await task + assert api_runs.api_worker_live_count() == 1, ( + "cancelled handler took the worker count with it" + ) + + await gw_mod.GatewayRunner.stop(gw) + + assert "close:session_db" not in events and "close_all" not in events, ( + f"SessionDB closed while the API worker was still alive: {events}" + ) + + release.set() + for _ in range(100): + if "worker_done" in events: + break + await asyncio.sleep(0.05) + assert "worker_done" in events, "worker never finished" + assert api_runs.api_worker_live_count() == 0, "worker exit did not release the count" + + def test_shutdown_executor_defaults_to_no_wait(): """The no-argument call keeps the historical fire-and-forget contract.""" gw = _FakeGateway([])