RomanNum3ral
/
hermes-agent
mirror of https://github.com/NousResearch/hermes-agent.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

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.
This commit is contained in:
briandevans 2026-08-05 20:46:12 -07:00 committed by kshitij
parent 51fa7db469
commit d9ddfb23d5
3 changed files with 399 additions and 12 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

66
gateway/platforms/api_server.py
View File

@ -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()

20
gateway/run.py
View File

@ -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)

325
tests/gateway/test_api_server_active_work_drain.py
View File

@ -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"
)