mirrors/hermes-agent
1
0
Fork 0
mirror of https://github.com/NousResearch/hermes-agent.git synced 2026-07-20 15:33:54 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 6

fix(gateway): complete API-server shutdown drain

Browse source
This commit is contained in:
kshitijk4poor 2026-07-13 16:27:18 +05:30 • committed by kshitij
parent 021ee34546
commit 104ffeae23
4 changed files with 198 additions and 110 deletions
Download patch file Download diff file Expand all files Collapse all files

59
gateway/platforms/api_server.py
View file

@ -914,18 +914,50 @@ class APIServerAdapter(BasePlatformAdapter):
self._inflight_agent_runs: int = 0
def active_agent_work_count(self) -> int:
"""Count of in-flight agent work served by this API server adapter.
"""Return all live agent work owned by this API adapter.
Covers the non-streaming chat/responses paths (``_inflight_agent_runs``)
plus live ``/v1/runs`` agents (``_active_run_agents``). Used by the
gateway shutdown drain so desk/API sessions get the same drain window
as messaging sessions and cron jobs (#63529).
``/v1/runs`` registers an asyncio task before it constructs and stores
its agent, so ``_active_run_agents`` has a real queued-before-agent gap.
Reuse the task-based accounting used by the concurrent-run limit: it
covers that gap and excludes completed tasks retained until cleanup.
"""
try:
return int(self._inflight_agent_runs) + len(self._active_run_agents)
return int(self._inflight_agent_runs) + sum(
not task.done() for task in self._active_run_tasks.values()
)
except Exception:
return 0
@staticmethod
def _gateway_is_draining() -> bool:
"""Whether the owning gateway currently refuses new agent turns."""
try:
from gateway.run import _gateway_runner_ref
runner = _gateway_runner_ref()
return bool(
runner
and (
getattr(runner, "_draining", False)
or getattr(runner, "_external_drain_active", False)
)
)
except Exception:
return False
def _draining_response(self) -> Optional["web.Response"]:
"""Return a retryable response while the gateway drains existing work."""
if not self._gateway_is_draining():
return None
return web.json_response(
_openai_error(
"Gateway is draining existing work; retry shortly.",
code="gateway_draining",
),
status=503,
headers={"Retry-After": "1"},
)
def _readiness_work_counts(self) -> tuple[int, int, int]:
"""Return bounded work counts from each subsystem's public state."""
active_api_runs = sum(
@ -1937,6 +1969,9 @@ class APIServerAdapter(BasePlatformAdapter):
auth_err = self._check_auth(request)
if auth_err:
return auth_err
draining = self._draining_response()
if draining is not None:
return draining
gateway_session_key, key_err = self._parse_session_key_header(request)
if key_err is not None:
return key_err
@ -1981,6 +2016,9 @@ class APIServerAdapter(BasePlatformAdapter):
auth_err = self._check_auth(request)
if auth_err:
return auth_err
draining = self._draining_response()
if draining is not None:
return draining
gateway_session_key, key_err = self._parse_session_key_header(request)
if key_err is not None:
return key_err
@ -2122,6 +2160,9 @@ class APIServerAdapter(BasePlatformAdapter):
auth_err = self._check_auth(request)
if auth_err:
return auth_err
draining = self._draining_response()
if draining is not None:
return draining
# Bound total in-flight agent runs (configurable; #7483).
limited = self._concurrency_limited_response()
@ -3258,6 +3299,9 @@ class APIServerAdapter(BasePlatformAdapter):
auth_err = self._check_auth(request)
if auth_err:
return auth_err
draining = self._draining_response()
if draining is not None:
return draining
# Bound total in-flight agent runs (configurable; #7483).
limited = self._concurrency_limited_response()
@ -4233,6 +4277,9 @@ class APIServerAdapter(BasePlatformAdapter):
auth_err = self._check_auth(request)
if auth_err:
return auth_err
draining = self._draining_response()
if draining is not None:
return draining
# Long-term memory scope header (see chat_completions for details).
gateway_session_key, key_err = self._parse_session_key_header(request)

50
gateway/run.py
View file

@ -4143,55 +4143,19 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
return 0
def _active_api_run_count(self) -> int:
"""Count of in-flight api_server agent runs across all adapters.
"""Count API-server work that is outside ``_running_agents``.
Chat/responses/desktop API sessions live entirely inside
``APIServerAdapter`` (``_inflight_agent_runs`` + ``_active_run_agents``)
and are invisible to ``self._running_agents``. Without this fold-in the
shutdown drain can report ``active_at_start=0`` and exit while desk/API
tools are still mid-flight (#63529) — the same structural failure cron
had before ``_active_cron_job_count`` (#60432). Best-effort: skips
adapters that don't expose the helper.
The primary API server owns the sole HTTP listener. Secondary multiplex
profiles cannot create an ``api_server`` adapter because it binds a port,
so only the primary registry is a supported source of this work.
"""
total = 0
try:
from gateway.config import Platform
adapter = getattr(self, "adapters", {}).get(Platform.API_SERVER)
helper = getattr(adapter, "active_agent_work_count", None)
return max(0, int(helper())) if callable(helper) else 0
except Exception:
return 0
def _count_from(adapters) -> int:
n = 0
if not adapters:
return 0
try:
values = adapters.values() if hasattr(adapters, "values") else adapters
except Exception:
return 0
for adapter in values:
try:
if getattr(adapter, "platform", None) != Platform.API_SERVER:
continue
helper = getattr(adapter, "active_agent_work_count", None)
if callable(helper):
n += int(helper())
else:
inflight = int(getattr(adapter, "_inflight_agent_runs", 0) or 0)
active = getattr(adapter, "_active_run_agents", None) or {}
n += inflight + len(active)
except Exception:
continue
return n
try:
total += _count_from(getattr(self, "adapters", None))
except Exception:
pass
try:
total += _count_from(getattr(self, "_profile_adapters", None))
except Exception:
pass
return total
# ── scale-to-zero idle detection / dormant-quiesce (Phase 0) ──────────────
# The gateway-side BEHAVIOUR that consumes the relay scale-to-zero primitives
# (gateway-gateway Phase 5). Pure logic lives in gateway/scale_to_zero.py; the

194
tests/gateway/test_api_server_active_work_drain.py
View file

@ -1,11 +1,10 @@
"""Tests for #63529: the gateway shutdown drain was structurally blind to
in-flight api_server (desk/API) agent runs.
"""Regression coverage for #63529 API-server shutdown draining.
API-server runs are tracked only inside ``APIServerAdapter``
(``_inflight_agent_runs`` + ``_active_run_agents``) and never enter
``GatewayRunner._running_agents``. Without folding them into the drain,
stop/restart reported ``active_at_start=0`` and let systemd SIGKILL mid-tool.
Mirrors tests/gateway/test_cron_active_work_drain.py for cron (#60432).
API-server work is adapter-owned rather than tracked by
``GatewayRunner._running_agents``. The shutdown drain must account for the
same live state as the API concurrency limiter, including a ``/v1/runs`` task
that exists before its agent has been constructed, and it must refuse new API
turns once the gateway starts draining.
"""
import asyncio
@ -13,81 +12,106 @@ from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from aiohttp import web
from aiohttp.test_utils import TestClient, TestServer
from gateway.config import Platform, PlatformConfig
from gateway.platforms.api_server import APIServerAdapter
from tests.gateway.restart_test_helpers import make_restart_runner
def _make_api_adapter(*, inflight: int = 0, active_ids=None):
from gateway.config import Platform
class _RunTask:
def __init__(self, done: bool = False):
self._done = done
active = {rid: MagicMock() for rid in (active_ids or [])}
def done(self) -> bool:
return self._done
def _make_api_adapter(*, inflight: int = 0, queued_ids=()):
tasks = {run_id: _RunTask() for run_id in queued_ids}
adapter = SimpleNamespace(
platform=Platform.API_SERVER,
_inflight_agent_runs=inflight,
_active_run_agents=active,
_active_run_tasks=tasks,
)
def active_agent_work_count() -> int:
return int(adapter._inflight_agent_runs) + len(adapter._active_run_agents)
return int(adapter._inflight_agent_runs) + sum(
not task.done() for task in adapter._active_run_tasks.values()
)
adapter.active_agent_work_count = active_agent_work_count
return adapter
def _make_admission_app(adapter: APIServerAdapter) -> web.Application:
app = web.Application()
app.router.add_post("/api/sessions/{session_id}/chat", adapter._handle_session_chat)
app.router.add_post(
"/api/sessions/{session_id}/chat/stream", adapter._handle_session_chat_stream
)
app.router.add_post("/v1/chat/completions", adapter._handle_chat_completions)
app.router.add_post("/v1/responses", adapter._handle_responses)
app.router.add_post("/v1/runs", adapter._handle_runs)
return app
class TestActiveApiRunCount:
def test_zero_when_no_api_adapters(self):
def test_zero_when_no_api_adapter(self):
runner, _adapter = make_restart_runner()
runner.adapters = {}
runner._profile_adapters = {}
assert runner._active_api_run_count() == 0
def test_sums_inflight_and_active_run_agents(self):
def test_delegates_to_primary_api_adapter(self):
runner, _adapter = make_restart_runner()
adapter = _make_api_adapter(inflight=2, active_ids=["r1"])
runner.adapters = {"api": adapter}
runner._profile_adapters = {}
assert runner._active_api_run_count() == 3
def test_includes_profile_adapters(self):
runner, _adapter = make_restart_runner()
runner.adapters = {"api": _make_api_adapter(inflight=1)}
runner._profile_adapters = {"p1": _make_api_adapter(active_ids=["a", "b"])}
runner.adapters = {
Platform.API_SERVER: _make_api_adapter(inflight=2, queued_ids=["r1"])
}
assert runner._active_api_run_count() == 3
def test_ignores_non_api_platforms(self):
from gateway.config import Platform
runner, _adapter = make_restart_runner()
other = SimpleNamespace(
platform=Platform.DISCORD,
_inflight_agent_runs=99,
_active_run_agents={"x": MagicMock()},
active_agent_work_count=lambda: 99,
)
runner.adapters = {"discord": other}
runner._profile_adapters = {}
runner.adapters = {Platform.DISCORD: other}
assert runner._active_api_run_count() == 0
def test_never_raises_on_broken_adapters(self):
def test_never_raises_on_broken_adapter(self):
runner, _adapter = make_restart_runner()
class Bad:
platform = property(lambda self: (_ for _ in ()).throw(RuntimeError("boom")))
platform = Platform.API_SERVER
runner.adapters = {"bad": Bad()}
runner._profile_adapters = None
@staticmethod
def active_agent_work_count() -> int:
raise RuntimeError("boom")
runner.adapters = {Platform.API_SERVER: Bad()}
assert runner._active_api_run_count() == 0
class TestAPIServerAdapterWorkCount:
def test_active_agent_work_count_on_real_class_method(self):
from gateway.platforms.api_server import APIServerAdapter
# Instantiate unbound helpers via object.__new__ to avoid full connect.
adapter = object.__new__(APIServerAdapter)
def test_counts_live_run_task_before_agent_creation(self):
adapter = APIServerAdapter(PlatformConfig(enabled=True))
adapter._inflight_agent_runs = 2
adapter._active_run_agents = {"r1": object(), "r2": object()}
assert adapter.active_agent_work_count() == 4
adapter._active_run_tasks = {
"queued": _RunTask(),
"finished": _RunTask(done=True),
}
adapter._active_run_agents = {}
assert adapter.active_agent_work_count() == 3
def test_does_not_double_count_started_run_agent(self):
adapter = APIServerAdapter(PlatformConfig(enabled=True))
adapter._inflight_agent_runs = 0
adapter._active_run_tasks = {"run-1": _RunTask()}
adapter._active_run_agents = {"run-1": object()}
assert adapter.active_agent_work_count() == 1
class TestDrainWaitsForApiWork:
@ -95,56 +119,79 @@ class TestDrainWaitsForApiWork:
async def test_drain_returns_immediately_when_nothing_active(self):
runner, _adapter = make_restart_runner()
runner.adapters = {}
runner._profile_adapters = {}
_snapshot, timed_out = await runner._drain_active_agents(5.0)
assert timed_out is False
@pytest.mark.asyncio
async def test_drain_waits_for_in_flight_api_run(self):
async def test_drain_waits_for_real_queued_run_before_agent_creation(self):
"""A live /v1/runs task must block drain before it has an agent."""
runner, _adapter = make_restart_runner()
api = _make_api_adapter(inflight=1)
runner.adapters = {"api": api}
runner._profile_adapters = {}
api = APIServerAdapter(PlatformConfig(enabled=True))
runner.adapters = {Platform.API_SERVER: api}
app = _make_admission_app(api)
original_create_task = asyncio.create_task
task_started = asyncio.Event()
allow_task = asyncio.Event()
async def finish_run():
await asyncio.sleep(0.12)
api._inflight_agent_runs = 0
def delayed_create_task(coro):
async def delayed():
task_started.set()
await allow_task.wait()
return await coro
task = asyncio.create_task(finish_run())
_snapshot, timed_out = await runner._drain_active_agents(2.0)
await task
return original_create_task(delayed())
assert timed_out is False, (
"drain must wait for api_server work, not report active_at_start=0"
)
mock_agent = MagicMock()
mock_agent.run_conversation.return_value = {"final_response": "done"}
mock_agent.session_prompt_tokens = 0
mock_agent.session_completion_tokens = 0
mock_agent.session_total_tokens = 0
with patch(
"gateway.platforms.api_server.asyncio.create_task",
side_effect=delayed_create_task,
), patch.object(api, "_create_agent", return_value=mock_agent):
async with TestClient(TestServer(app)) as client:
response = await client.post("/v1/runs", json={"input": "hello"})
assert response.status == 202
await task_started.wait()
assert api._active_run_agents == {}
assert runner._active_api_run_count() == 1
drain_task = original_create_task(runner._drain_active_agents(2.0))
await asyncio.sleep(0.1)
assert not drain_task.done()
allow_task.set()
_snapshot, timed_out = await drain_task
assert timed_out is False
@pytest.mark.asyncio
async def test_drain_times_out_if_api_run_outlives_the_window(self):
runner, _adapter = make_restart_runner()
runner.adapters = {"api": _make_api_adapter(inflight=1)}
runner._profile_adapters = {}
runner.adapters = {Platform.API_SERVER: _make_api_adapter(queued_ids=["run-1"])}
_snapshot, timed_out = await runner._drain_active_agents(0.1)
assert timed_out is True
@pytest.mark.asyncio
async def test_drain_still_waits_for_chat_and_cron(self):
async def test_drain_still_waits_for_chat_cron_and_api_work(self):
import cron.scheduler as sched
runner, _adapter = make_restart_runner()
runner._running_agents = {"session-1": MagicMock()}
sched._running_job_ids.add("job-1")
runner.adapters = {"api": _make_api_adapter(inflight=1)}
runner._profile_adapters = {}
runner.adapters = {Platform.API_SERVER: _make_api_adapter(queued_ids=["run-1"])}
async def finish_all():
await asyncio.sleep(0.12)
runner._running_agents.clear()
sched._running_job_ids.discard("job-1")
runner.adapters["api"]._inflight_agent_runs = 0
runner.adapters[Platform.API_SERVER]._active_run_tasks.clear()
task = asyncio.create_task(finish_all())
try:
@ -154,3 +201,28 @@ class TestDrainWaitsForApiWork:
sched._running_job_ids.discard("job-1")
assert timed_out is False
class TestDrainAdmission:
@pytest.mark.asyncio
async def test_drain_refuses_every_agent_start_endpoint(self):
adapter = APIServerAdapter(PlatformConfig(enabled=True))
runner = SimpleNamespace(_draining=True, _external_drain_active=False)
app = _make_admission_app(adapter)
paths = (
"/api/sessions/missing/chat",
"/api/sessions/missing/chat/stream",
"/v1/chat/completions",
"/v1/responses",
"/v1/runs",
)
with patch("gateway.run._gateway_runner_ref", lambda: runner):
async with TestClient(TestServer(app)) as client:
for path in paths:
response = await client.post(path, json={})
payload = await response.json()
assert response.status == 503
assert response.headers["Retry-After"] == "1"
assert payload["error"]["code"] == "gateway_draining"

5
tests/gateway/test_shutdown_cache_cleanup.py
View file

@ -60,6 +60,11 @@ class _FakeGateway:
# there's never in-flight cron work to report.
return 0
def _active_api_run_count(self):
# The shutdown log also reports adapter-owned API work (#63529).
# This fake has no API server adapter, so it is always idle.
return 0
def _update_runtime_status(self, *_a, **_kw):
pass