From d9ddfb23d5cb3a4346a8305f25d43e2827ed88a1 Mon Sep 17 00:00:00 2001 From: briandevans <252620095+briandevans@users.noreply.github.com> Date: Wed, 5 Aug 2026 20:46:12 -0700 Subject: [PATCH] fix(gateway): interrupt every in-flight API turn on shutdown, not just /v1/runs The shutdown drain ACCOUNTS for API-server work but never INTERRUPTS it. `_drain_active_agents()` folds `_active_api_run_count()` into both its wait loop and its `timed_out` verdict, while `_interrupt_running_agents()` iterates `self._running_agents` only -- a dict no API turn ever enters, because the API server owns its own agent lifecycle. `gateway/run.py` states the gap against itself: "API-server / desk sessions have the same structural gap (#63529)." The user-visible result is that every gateway restart with a live API or desktop turn burns the full drain timeout and then runs `_kill_tool_subprocesses("post-interrupt")`, which amputates the turn's tool subprocesses with no cooperative interrupt and no resume marker. There are seven API agent-entry points. Six funnel through `_run_agent()` (both session-chat routes, and `/v1/chat/completions` + `/v1/responses` in streaming and non-streaming form) and are counted by `_inflight_agent_runs`; the seventh, `/v1/runs`, runs its own lifecycle and is counted through `_active_run_tasks`. None of the six has a run_id, so the run_id-keyed `_active_run_agents` cannot reach them, and only two pass `agent_ref` -- which lands in a caller-local list, not a registry. So register once at the single unconditional creation site inside `_run_agent`, beside the existing `_publish_turn_process_ownership()` call, and unregister in the same `finally` that already clears it. That one symmetric pair covers all six callers. The registry is adapter-owned and keyed by object identity, kept separate from `_active_run_agents` because that dict is run_id-keyed and scoped to the public `/v1/runs` stop API. `interrupt_active_runs()` then walks both registries, deduped by identity, so the interrupt set matches the set the drain waits on. The settle window after the interrupt now polls API work as well: the interrupt is cooperative, and without this the window closes the instant `_running_agents` is empty -- which it always is for API turns -- and the tool kill lands on a turn that was asked to stop microseconds earlier. --- gateway/platforms/api_server.py | 66 +++- gateway/run.py | 20 +- .../test_api_server_active_work_drain.py | 325 ++++++++++++++++++ 3 files changed, 399 insertions(+), 12 deletions(-) diff --git a/gateway/platforms/api_server.py b/gateway/platforms/api_server.py index d24cffc05fc6b..1b8c4af9c42c1 100644 --- a/gateway/platforms/api_server.py +++ b/gateway/platforms/api_server.py @@ -1448,6 +1448,15 @@ class APIServerAdapter(BasePlatformAdapter): # (the /v1/runs path tracks its own in-flight set via # _active_run_tasks). self._inflight_agent_runs: int = 0 + # Every agent currently inside _run_agent(), i.e. exactly the turns + # counted by _inflight_agent_runs above. Shutdown needs the whole + # adapter-owned set, so this is deliberately NOT _active_run_agents: + # that one is run_id-keyed and scoped to the public /v1/runs stop API, + # and only /v1/runs has a run_id at all. Keyed by id() because the + # other six agent-entry paths have no stable identifier of their own; + # the dict holds a strong reference for the life of the turn, so an + # id() can never be recycled while it is still registered. + self._shutdown_interruptible_agents: Dict[int, Any] = {} # Back-reference to the owning GatewayRunner (set by gateway/run.py) # so /api/platforms/{platform}/events can resolve sibling adapters. # BasePlatformAdapter declares the class-level default of None. @@ -1475,22 +1484,49 @@ class APIServerAdapter(BasePlatformAdapter): return 0 def interrupt_active_runs(self, reason: str) -> int: - """Interrupt active API-run agents during gateway shutdown. + """Cooperatively interrupt every adapter-owned agent during shutdown. The gateway drain accounts for API-server work through ``active_agent_work_count()``, but those agents are owned by this - adapter rather than ``GatewayRunner._running_agents``. Expose the same - cooperative interrupt used by ``POST /v1/runs/{run_id}/stop`` so a - shutdown timeout can stop long-running API work before process teardown. + adapter rather than ``GatewayRunner._running_agents``, so + ``GatewayRunner._interrupt_running_agents()`` never reaches them: the + turn runs to the drain timeout with no cooperative interrupt and is + then amputated by the post-interrupt tool-subprocess kill. + + Cover the same set the drain waits on, so accounting and interrupt + agree: + + * ``_active_run_agents`` — the ``/v1/runs`` agents counted through + ``_active_run_tasks``. + * ``_shutdown_interruptible_agents`` — every ``_run_agent()`` turn + counted through ``_inflight_agent_runs``, i.e. both session-chat + routes, ``/v1/chat/completions`` and ``/v1/responses`` in their + streaming and non-streaming forms. + + ``_pending_agent_requests`` is intentionally not covered: it counts + admitted requests that have not constructed an agent yet, so there is + no object to interrupt. + + Returns the number of agents that accepted an interrupt. """ + agents: Dict[int, Any] = {} + for agent in list(self._active_run_agents.values()): + if agent is not None: + agents[id(agent)] = agent + for agent in list(self._shutdown_interruptible_agents.values()): + if agent is not None: + # Dedupe by object identity — the two registries are disjoint + # today (/v1/runs runs its own lifecycle, not _run_agent), but + # an agent published to both must still be interrupted once. + agents[id(agent)] = agent + interrupted = 0 - for run_id, agent in list(self._active_run_agents.items()): + for agent in agents.values(): try: - agent.interrupt(reason) - interrupted += 1 - logger.debug("[api_server] interrupted active run %s during shutdown", run_id) + if request_hard_interrupt(agent, reason): + interrupted += 1 except Exception as exc: - logger.debug("[api_server] failed interrupting active run %s: %s", run_id, exc) + logger.debug("[api_server] failed interrupting active agent: %s", exc) return interrupted @staticmethod @@ -6082,6 +6118,13 @@ class APIServerAdapter(BasePlatformAdapter): # runs its own agent lifecycle and doesn't go through # TurnRunner, so it needs its own baseline. _publish_turn_process_ownership(agent, effective_task_id) + # Shutdown interrupt coverage (#63529). Registering here, + # once, covers every _run_agent() caller — the same reason + # the _ProviderAuthResolutionError handler below lives here + # rather than in each route. Only two callers pass + # ``agent_ref``, and only /v1/runs has a run_id, so neither + # is a usable hook for the rest. + self._shutdown_interruptible_agents[id(agent)] = agent result = agent.run_conversation( user_message=user_message, conversation_history=conversation_history, @@ -6211,6 +6254,11 @@ class APIServerAdapter(BasePlatformAdapter): # in gateway/run.py's _run_sync_with_timeout_lifecycle. if agent is not None: _clear_turn_process_ownership(agent) + # Symmetric with the registration above: the turn is + # over, so it must not be interrupted by a later + # shutdown. pop() is a no-op when _create_agent + # succeeded but the turn never reached registration. + self._shutdown_interruptible_agents.pop(id(agent), None) clear_session_vars(tokens) self._activate_admitted_request() diff --git a/gateway/run.py b/gateway/run.py index 2d2934e48d4de..b2f4f2d95b819 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -7421,7 +7421,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew return 0 def _interrupt_api_server_runs(self, reason: str) -> int: - """Interrupt API-server agents that are not in ``_running_agents``.""" + """Interrupt API-server agents that are not in ``_running_agents``. + + Counterpart of ``_active_api_run_count()``: that method folds + adapter-owned API work into the shutdown drain, so this one must reach + the same agents when the drain times out. Duck-typed on the adapter so + an older adapter (or a minimal test double for this class) without the + hook is simply skipped rather than raising mid-shutdown. + """ try: adapter = getattr(self, "adapters", {}).get(Platform.API_SERVER) helper = getattr(adapter, "interrupt_active_runs", None) @@ -9259,6 +9266,9 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew logger.debug("Interrupted running agent for session %s during shutdown", session_key) except Exception as e: logger.debug("Failed interrupting agent during shutdown: %s", e) + # API-server / desk turns are adapter-owned and never enter + # _running_agents, so the loop above cannot see them even though + # _drain_active_agents() waited for them (#63529). interrupted_api = self._interrupt_api_server_runs(reason) if interrupted_api: logger.debug("Interrupted %d api_server run(s) during shutdown", interrupted_api) @@ -12929,9 +12939,13 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew _INTERRUPT_REASON_GATEWAY_RESTART if self._restart_requested else _INTERRUPT_REASON_GATEWAY_SHUTDOWN ) interrupt_deadline = asyncio.get_running_loop().time() + 5.0 + # Wait on API-server work too. The interrupt is cooperative: + # without this the settle window closes the instant + # _running_agents is empty, and an API turn that was just asked + # to stop gets its tool subprocesses killed below before it can + # unwind — the exact amputation this interrupt exists to avoid. while ( - self._running_agents - or self._active_api_run_count() + self._running_agents or self._active_api_run_count() ) and asyncio.get_running_loop().time() < interrupt_deadline: self._update_runtime_status("draining") await asyncio.sleep(0.1) diff --git a/tests/gateway/test_api_server_active_work_drain.py b/tests/gateway/test_api_server_active_work_drain.py index 45a4a467980ba..9edc39f72e0a8 100644 --- a/tests/gateway/test_api_server_active_work_drain.py +++ b/tests/gateway/test_api_server_active_work_drain.py @@ -8,6 +8,7 @@ turns once the gateway starts draining. """ import asyncio +import threading from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock, patch @@ -17,8 +18,14 @@ from aiohttp.test_utils import TestClient, TestServer from gateway.config import Platform, PlatformConfig from gateway.platforms.api_server import APIServerAdapter +from gateway.run import _INTERRUPT_REASON_GATEWAY_SHUTDOWN +from hermes_state import SessionDB from tests.gateway.restart_test_helpers import make_restart_runner +# Safety net so a regression parks the executor thread forever instead of +# hanging CI. No assertion below depends on elapsed time. +_TURN_UNBLOCK_TIMEOUT = 30.0 + class _RunTask: def __init__(self, done: bool = False): @@ -228,3 +235,321 @@ class TestDrainAdmission: assert payload["error"]["code"] == "gateway_draining" +# --------------------------------------------------------------------------- +# Shutdown interrupt coverage (#63529) +# +# The drain ACCOUNTS for every API turn (`active_agent_work_count()` sums +# `_pending_agent_requests` + `_inflight_agent_runs` + live `_active_run_tasks`) +# but `GatewayRunner._interrupt_running_agents()` only walked +# `self._running_agents`, which no API turn ever enters. So an API turn held +# the drain open for the full timeout and was then amputated by +# `_kill_tool_subprocesses("post-interrupt")` with no cooperative interrupt. +# +# `/v1/runs` is only one of seven API agent-entry points. The other six all +# funnel through `_run_agent()` — both session-chat routes and +# `/v1/chat/completions` + `/v1/responses` in streaming and non-streaming form +# — and none of them has a run_id, so `_active_run_agents` cannot reach them. +# --------------------------------------------------------------------------- + + +def _parked_agent(loop, started: asyncio.Event, release: threading.Event) -> MagicMock: + """A mock agent whose turn parks inside ``run_conversation`` until released. + + ``request_hard_interrupt`` falls back to ``agent.interrupt(reason)`` for an + unspecced ``MagicMock`` — ``inspect.getattr_static`` refuses to invent + ``hard_interrupt`` on a ``__getattr__`` proxy — which is exactly the ABI + teknium1's review asked this regression to verify. + """ + agent = MagicMock() + agent.session_id = None + agent.session_prompt_tokens = 0 + agent.session_completion_tokens = 0 + agent.session_total_tokens = 0 + agent._last_compaction_in_place = False + agent._hermes_api_runtime = {} + + def _park(user_message=None, conversation_history=None, task_id=None): + loop.call_soon_threadsafe(started.set) + release.wait(_TURN_UNBLOCK_TIMEOUT) + return {"final_response": "done", "messages": [], "api_calls": 0, "tools": []} + + agent.run_conversation.side_effect = _park + # A real agent unwinds its turn on interrupt; releasing here models that so + # the parked executor thread can finish. + agent.interrupt.side_effect = lambda *_a, **_k: release.set() + return agent + + +class _SettlingApiAdapter: + """API adapter double whose work clears a few polls AFTER it is interrupted. + + The poll count is the deterministic quantity under test: it makes "the + settle window kept polling API work" observable without timing anything. + """ + + def __init__(self, polls_to_settle: int = 3): + self._polls_to_settle = polls_to_settle + self.interrupt_reasons: list = [] + + def active_agent_work_count(self) -> int: + if not self.interrupt_reasons: + return 1 + if self._polls_to_settle > 0: + self._polls_to_settle -= 1 + return 1 + return 0 + + def interrupt_active_runs(self, reason: str) -> int: + self.interrupt_reasons.append(reason) + return 1 + + @property + def settled(self) -> bool: + """Non-consuming view of the same state, safe to read from a spy.""" + return bool(self.interrupt_reasons) and self._polls_to_settle == 0 + + +def _make_async_noop(): + async def _noop(*args, **kwargs): + return None + + return _noop + + +class TestRunAgentRegistersForShutdownInterrupt: + @pytest.mark.asyncio + async def test_run_agent_registers_and_unregisters_the_agent(self): + """One registration inside ``_run_agent`` covers all six of its callers. + + Only two callers pass ``agent_ref``, and that lands in a caller-local + list rather than any registry, so it is not a usable hook. + """ + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + agent = MagicMock() + agent.session_id = None + agent.session_prompt_tokens = 0 + agent.session_completion_tokens = 0 + agent.session_total_tokens = 0 + agent._last_compaction_in_place = False + observed = {} + + def _record(user_message=None, conversation_history=None, task_id=None): + observed["during"] = dict(adapter._shutdown_interruptible_agents) + return {"final_response": "done", "messages": [], "api_calls": 0, "tools": []} + + agent.run_conversation.side_effect = _record + + with patch.object(adapter, "_create_agent", return_value=agent): + await adapter._run_agent( + user_message="hello", + conversation_history=[], + session_id="s1", + ) + + assert list(observed["during"].values()) == [agent] + assert adapter._shutdown_interruptible_agents == {} + + @pytest.mark.asyncio + async def test_agent_is_unregistered_when_the_turn_raises(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + agent = MagicMock() + agent.run_conversation.side_effect = RuntimeError("boom") + + with patch.object(adapter, "_create_agent", return_value=agent): + with pytest.raises(RuntimeError): + await adapter._run_agent( + user_message="hello", + conversation_history=[], + session_id="s1", + ) + + assert adapter._shutdown_interruptible_agents == {} + + +class TestInterruptActiveRuns: + def test_interrupts_v1_runs_agents(self): + """The ``/v1/runs`` coverage #63963 established stays green.""" + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + agent = MagicMock() + adapter._active_run_agents = {"run-1": agent} + + assert adapter.interrupt_active_runs("gateway shutdown") == 1 + agent.interrupt.assert_called_once_with("gateway shutdown") + + def test_interrupts_each_agent_exactly_once_across_both_registries(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + shared = MagicMock() + run_only = MagicMock() + turn_only = MagicMock() + adapter._active_run_agents = {"run-1": run_only, "run-2": shared} + adapter._shutdown_interruptible_agents = { + id(shared): shared, + id(turn_only): turn_only, + } + + assert adapter.interrupt_active_runs("gateway shutdown") == 3 + shared.interrupt.assert_called_once_with("gateway shutdown") + run_only.interrupt.assert_called_once_with("gateway shutdown") + turn_only.interrupt.assert_called_once_with("gateway shutdown") + + def test_one_bad_agent_does_not_strand_the_others(self): + adapter = APIServerAdapter(PlatformConfig(enabled=True)) + exploding = MagicMock() + exploding.interrupt.side_effect = RuntimeError("already torn down") + no_abi = object() # exposes neither hard_interrupt nor interrupt + healthy = MagicMock() + adapter._shutdown_interruptible_agents = { + id(exploding): exploding, + id(no_abi): no_abi, + id(healthy): healthy, + } + + assert adapter.interrupt_active_runs("gateway shutdown") == 1 + healthy.interrupt.assert_called_once_with("gateway shutdown") + + +class TestShutdownInterruptReachesEveryApiTurn: + @pytest.mark.asyncio + async def test_chat_completions_turn_is_interrupted(self): + """A non-``/v1/runs`` API turn, end to end through the real handler. + + This is teknium1's named acceptance criterion on #63963: the drain + counts this turn, so the shutdown interrupt must reach it. + """ + runner, _adapter = make_restart_runner() + api = APIServerAdapter(PlatformConfig(enabled=True)) + runner.adapters = {Platform.API_SERVER: api} + app = _make_admission_app(api) + + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = threading.Event() + agent = _parked_agent(loop, started, release) + + try: + with patch.object(api, "_create_agent", return_value=agent): + async with TestClient(TestServer(app)) as client: + request = asyncio.ensure_future( + client.post( + "/v1/chat/completions", + json={"messages": [{"role": "user", "content": "hi"}]}, + ) + ) + await asyncio.wait_for(started.wait(), _TURN_UNBLOCK_TIMEOUT) + + # The drain sees this turn ... + assert runner._active_api_run_count() == 1 + # ... and it is not in _running_agents, so only the API + # hook can reach it. + assert runner._running_agents == {} + + runner._interrupt_running_agents(_INTERRUPT_REASON_GATEWAY_SHUTDOWN) + + agent.interrupt.assert_called_once_with( + _INTERRUPT_REASON_GATEWAY_SHUTDOWN + ) + response = await asyncio.wait_for(request, _TURN_UNBLOCK_TIMEOUT) + assert response.status == 200 + finally: + release.set() + + assert api._shutdown_interruptible_agents == {} + + @pytest.mark.asyncio + async def test_session_chat_sse_turn_is_interrupted(self, tmp_path): + """The SSE session-chat route is a second, differently shaped caller.""" + runner, _adapter = make_restart_runner() + api = APIServerAdapter(PlatformConfig(enabled=True)) + session_db = SessionDB(tmp_path / "state.db") + api._session_db = session_db + runner.adapters = {Platform.API_SERVER: api} + app = _make_admission_app(api) + session_id = session_db.create_session("sse-session", "api_server") + + loop = asyncio.get_running_loop() + started = asyncio.Event() + release = threading.Event() + agent = _parked_agent(loop, started, release) + + try: + with patch.object(api, "_create_agent", return_value=agent): + async with TestClient(TestServer(app)) as client: + request = asyncio.ensure_future( + client.post( + f"/api/sessions/{session_id}/chat/stream", + json={"message": "hi"}, + ) + ) + await asyncio.wait_for(started.wait(), _TURN_UNBLOCK_TIMEOUT) + + assert runner._active_api_run_count() == 1 + assert runner._running_agents == {} + + runner._interrupt_running_agents(_INTERRUPT_REASON_GATEWAY_SHUTDOWN) + + agent.interrupt.assert_called_once_with( + _INTERRUPT_REASON_GATEWAY_SHUTDOWN + ) + response = await asyncio.wait_for(request, _TURN_UNBLOCK_TIMEOUT) + assert response.status == 200 + await asyncio.wait_for(response.text(), _TURN_UNBLOCK_TIMEOUT) + finally: + release.set() + close = getattr(session_db, "close", None) + if callable(close): + close() + + assert api._shutdown_interruptible_agents == {} + + def test_interrupt_running_agents_is_a_noop_without_an_api_adapter(self): + """The hook is duck-typed — an adapterless runner must not raise.""" + runner, _adapter = make_restart_runner() + runner.adapters = {} + + runner._interrupt_running_agents(_INTERRUPT_REASON_GATEWAY_SHUTDOWN) + + assert runner._interrupt_api_server_runs("x") == 0 + + +class TestShutdownSettleWindow: + @pytest.mark.asyncio + async def test_settle_window_waits_for_interrupted_api_work(self, monkeypatch): + """The interrupt is cooperative, so the settle window must poll API work. + + Otherwise the window closes the instant ``_running_agents`` is empty — + which it always is for API turns — and the post-interrupt tool kill + lands on a turn that was asked to stop microseconds earlier. + """ + import tools.browser_tool as _bt + import tools.process_registry as _pr + import tools.terminal_tool as _tt + + runner, adapter = make_restart_runner() + runner._restart_drain_timeout = 0.01 # force the drain-timeout path + adapter.disconnect = _make_async_noop() + api = _SettlingApiAdapter() + runner.adapters = {Platform.TELEGRAM: adapter, Platform.API_SERVER: api} + + settled_at_kill: list = [] + + def _spy_kill_all(task_id=None): + settled_at_kill.append(api.settled) + return 0 + + monkeypatch.setattr(_pr.process_registry, "kill_all", _spy_kill_all) + monkeypatch.setattr(_tt, "cleanup_all_environments", lambda: None) + monkeypatch.setattr(_bt, "cleanup_all_browsers", lambda: None) + + with patch("gateway.status.remove_pid_file"), \ + patch("gateway.status.write_runtime_status"), \ + patch("cron.scheduler.mark_job_run"): + await runner.stop() + + assert api.interrupt_reasons == [_INTERRUPT_REASON_GATEWAY_SHUTDOWN] + assert settled_at_kill, "post-interrupt tool kill never ran" + assert settled_at_kill[0] is True, ( + "post-interrupt tool kill ran while the interrupted API turn was " + "still unwinding" + ) + +