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

fix(api-server): epoch-gate abandoned-run reaps and cover the /v1/runs sibling 

The API server intentionally lets concurrent runs share a client-provided
session_id (= process task_id), so the SSE-disconnect reap could kill a
process a still-live concurrent run spawned after the disconnecting run's
baseline — the same stale-reaper bug class the gateway path gates via
run_generation.

- Per-task-id run epochs (monotonic counter): each run claims the epoch at
  publish; a reaper holding a superseded epoch declines to kill. A missing
  entry (the run's own clear pruned it) still reaps, so the leak fix isn't
  silently disabled.
- _publish_turn_process_ownership / _clear_turn_process_ownership helpers
  replace the copy-pasted marker set/clear blocks, so attribute names and
  epoch bookkeeping can't drift between surfaces.
- /v1/runs — the third own-lifecycle surface — now records ownership and
  reaps on POST /v1/runs/{id}/stop and on server-side SSE cancellation,
  closing the remaining sibling paths of #76115.
This commit is contained in:
kshitijk4poor 2026-08-02 14:09:56 +05:30 committed by kshitij
parent eb4772ec2f
commit 8f0f55eaac
2 changed files with 250 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

121
gateway/platforms/api_server.py
View File

@ -43,6 +43,7 @@ import asyncio
import errno
import hashlib
import hmac
import itertools
import json
from contextlib import contextmanager, nullcontext
from contextvars import ContextVar
@ -52,6 +53,7 @@ import os
import re
import sqlite3
import sys
import threading
import time
import uuid
from pathlib import Path
@ -617,32 +619,108 @@ def _multimodal_validation_error(exc: ValueError, *, param: str) -> "web.Respons
)
def _reap_disconnected_agent_processes(agent: Any) -> None:
"""Reap background processes an SSE-disconnected agent's turn created.
def _reap_disconnected_agent_processes(
agent: Any, *, source: str = "api_server_sse_disconnect"
) -> None:
"""Reap background processes an abandoned API-server turn created.
Mirrors the gateway-turn cleanup in ``gateway/run.py`` (#76115) for this
API-server surface, which runs its own agent lifecycle via ``_run_agent``
and never passes through ``TurnRunner`` — so it needs its own trigger for
the same baseline-diff reap. Fire-and-forget on a daemon thread so the
SSE handler's own cleanup isn't blocked on process-tree teardown.
Reaping is epoch-gated: client-provided session IDs are conversation
scopes, and multiple concurrent runs can intentionally share one (see
``_handle_runs``). Without the gate, run A disconnecting could kill a
process a still-live run B (same task_id) spawned after A's baseline
snapshot — the same stale-reaper bug class the gateway path gates via
``run_generation``. The epoch closure skips the reap when a newer run
has since claimed the task_id; that newer run's own baseline covers its
eventual cleanup.
"""
process_task_id = getattr(agent, "_gateway_turn_process_task_id", "")
process_baseline = getattr(agent, "_gateway_turn_process_baseline", None)
if not process_task_id or process_baseline is None:
return
import threading
epoch = getattr(agent, "_gateway_turn_process_epoch", None)
is_still_current: Optional[Any] = None
if epoch is not None:
def _epoch_still_current(_task_id=process_task_id, _epoch=epoch):
# Skip only when a NEWER run has claimed this task_id. A missing
# entry means the abandoned run's own clear pruned it (worker
# returned after the interrupt) — no newer claimant exists, so
# the reap must still proceed or the leak survives. This matches
# the gateway gate's semantics: worker completion does not bump
# run_generation either.
with _TURN_PROCESS_EPOCH_LOCK:
current = _TURN_PROCESS_EPOCHS.get(_task_id)
return current is None or current == _epoch
is_still_current = _epoch_still_current
from gateway.run import _reap_gateway_turn_processes
threading.Thread(
target=_reap_gateway_turn_processes,
args=(process_task_id, process_baseline),
kwargs={"source": "api_server_sse_disconnect"},
kwargs={"source": source, "is_still_current": is_still_current},
name=f"api-turn-reaper-{process_task_id[:12]}",
daemon=True,
).start()
# Per-task-id run epochs for the reap gate above. task_id is a conversation
# scope shared by concurrent API runs, so each run that claims it bumps the
# epoch; a reaper holding a stale epoch declines to kill. Epochs come from a
# single monotonic counter (never reused), so pruning an entry and later
# re-claiming the task_id can never resurrect a stale reaper's claim.
# Entries are pruned on clear when still current, bounding the dict to
# in-flight runs.
_TURN_PROCESS_EPOCHS: Dict[str, int] = {}
_TURN_PROCESS_EPOCH_LOCK = threading.Lock()
_TURN_PROCESS_EPOCH_COUNTER = itertools.count(1)
def _publish_turn_process_ownership(agent: Any, task_id: str) -> None:
"""Snapshot the process baseline and claim the task_id's current epoch.
Single place all API-server agent lifecycles (chat/responses ``_run_agent``
and ``/v1/runs``) record turn ownership, so the marker attribute names and
epoch bookkeeping cannot drift between surfaces.
"""
from tools.process_registry import process_registry
with _TURN_PROCESS_EPOCH_LOCK:
epoch = next(_TURN_PROCESS_EPOCH_COUNTER)
_TURN_PROCESS_EPOCHS[task_id] = epoch
agent._gateway_turn_process_task_id = task_id
agent._gateway_turn_process_baseline = process_registry.snapshot_running_ids(
task_id
)
agent._gateway_turn_process_epoch = epoch
def _clear_turn_process_ownership(agent: Any) -> None:
"""Clear turn ownership the moment the turn finishes (success or crash).
A disconnect/cancel landing after this point must not reap background
work the turn deliberately left running — mirrors the same race-window
guard in ``gateway/run.py``'s ``_run_sync_with_timeout_lifecycle``.
"""
task_id = getattr(agent, "_gateway_turn_process_task_id", "")
epoch = getattr(agent, "_gateway_turn_process_epoch", None)
if task_id and epoch is not None:
with _TURN_PROCESS_EPOCH_LOCK:
# Prune only when this run is still the current claimant; a
# newer concurrent run owns the entry otherwise.
if _TURN_PROCESS_EPOCHS.get(task_id) == epoch:
del _TURN_PROCESS_EPOCHS[task_id]
agent._gateway_turn_process_task_id = ""
agent._gateway_turn_process_baseline = frozenset()
agent._gateway_turn_process_epoch = None
def _session_chat_user_message(body: Dict[str, Any], *, param: str = "message") -> tuple[Any, Optional["web.Response"]]:
"""Parse and normalize session chat ``message`` / ``input`` like chat completions."""
user_message = body.get("message") or body.get("input")
@ -4860,6 +4938,13 @@ class APIServerAdapter(BasePlatformAdapter):
agent.interrupt("SSE task cancelled")
except Exception:
pass
# Same abandonment as a client disconnect: the run will never
# be resumed, so reap the background processes it created
# (#76115). Epoch-gated; no-op when the turn already
# finished and cleared its markers.
_reap_disconnected_agent_processes(
agent, source="api_server_sse_cancelled"
)
if not agent_task.done():
agent_task.cancel()
logger.info("SSE task cancelled; persisted incomplete snapshot for %s", response_id)
@ -5868,11 +5953,7 @@ class APIServerAdapter(BasePlatformAdapter):
# gateway-turn cleanup (#76115); this API-server surface
# runs its own agent lifecycle and doesn't go through
# TurnRunner, so it needs its own baseline.
from tools.process_registry import process_registry
agent._gateway_turn_process_task_id = effective_task_id
agent._gateway_turn_process_baseline = (
process_registry.snapshot_running_ids(effective_task_id)
)
_publish_turn_process_ownership(agent, effective_task_id)
result = agent.run_conversation(
user_message=user_message,
conversation_history=conversation_history,
@ -6001,8 +6082,7 @@ class APIServerAdapter(BasePlatformAdapter):
# running on purpose. Mirrors the same race-window guard
# in gateway/run.py's _run_sync_with_timeout_lifecycle.
if agent is not None:
agent._gateway_turn_process_task_id = ""
agent._gateway_turn_process_baseline = frozenset()
_clear_turn_process_ownership(agent)
clear_session_vars(tokens)
self._activate_admitted_request()
@ -6355,12 +6435,23 @@ class APIServerAdapter(BasePlatformAdapter):
session_id=session_id or "",
)
register_gateway_notify(approval_session_key, _approval_notify)
# /v1/runs runs its own agent lifecycle (no
# TurnRunner, no _run_agent) — record turn process
# ownership so stop/cancel can reap only the
# background processes this run created (#76115).
_publish_turn_process_ownership(agent, effective_task_id)
r = agent.run_conversation(
user_message=user_message,
conversation_history=conversation_history,
task_id=effective_task_id,
)
finally:
# Worker finished (interrupted or complete) —
# clear turn ownership immediately so a later
# stop/cancel can't reap background work this
# run deliberately left running (same race-window
# guard as gateway/run.py and _run_agent above).
_clear_turn_process_ownership(agent)
try:
unregister_gateway_notify(approval_session_key)
finally:
@ -6700,6 +6791,14 @@ class APIServerAdapter(BasePlatformAdapter):
agent.interrupt("Stop requested via API")
except Exception:
pass
# The stopped run is abandoned — reap only the background
# processes it created (#76115). Epoch-gated inside, so a
# concurrent run sharing the same session_id keeps its own
# processes; no-op if the run already finished and cleared
# its ownership markers.
_reap_disconnected_agent_processes(
agent, source="api_server_run_stop"
)
return web.json_response({"run_id": run_id, "status": "stopping"})

140
tests/gateway/test_api_server.py
View File

@ -481,6 +481,146 @@ class TestDisconnectedAgentReap:
time.sleep(0.1)
assert calls == []
def test_stale_epoch_skips_reap_when_newer_run_claimed_task_id(self, monkeypatch):
"""#76188 follow-up: concurrent API runs can share a client-provided
session_id (same task_id). A disconnecting run whose epoch has been
superseded must NOT kill the newer run's processes."""
from gateway.platforms.api_server import (
_clear_turn_process_ownership,
_publish_turn_process_ownership,
_reap_disconnected_agent_processes,
)
from tools.process_registry import process_registry
calls = []
monkeypatch.setattr(
process_registry,
"kill_started_since",
lambda *a, **k: calls.append(True) or 1,
)
monkeypatch.setattr(
process_registry, "snapshot_running_ids", lambda _tid: frozenset()
)
run_a = types.SimpleNamespace()
run_b = types.SimpleNamespace()
_publish_turn_process_ownership(run_a, "shared-session")
# Run B claims the same session_id — supersedes A's epoch.
_publish_turn_process_ownership(run_b, "shared-session")
_reap_disconnected_agent_processes(run_a)
time.sleep(0.2)
assert calls == [], "stale run A must not reap run B's processes"
# Run B disconnecting IS current — its reap proceeds.
_reap_disconnected_agent_processes(run_b)
deadline = time.time() + 1.0
while not calls and time.time() < deadline:
time.sleep(0.01)
assert calls == [True]
_clear_turn_process_ownership(run_b)
def test_reap_proceeds_when_own_clear_pruned_the_epoch_entry(self, monkeypatch):
"""A missing epoch entry (the abandoned run's own finally already
cleared it) means no newer claimant — the reap must proceed using a
pre-captured marker snapshot, or the leak survives."""
from gateway.platforms.api_server import (
_clear_turn_process_ownership,
_publish_turn_process_ownership,
_reap_disconnected_agent_processes,
)
from tools.process_registry import process_registry
calls = []
monkeypatch.setattr(
process_registry,
"kill_started_since",
lambda *a, **k: calls.append(True) or 1,
)
monkeypatch.setattr(
process_registry, "snapshot_running_ids", lambda _tid: frozenset()
)
run = types.SimpleNamespace()
_publish_turn_process_ownership(run, "solo-session")
# Simulate the disconnect handler capturing the agent while the
# worker's finally clears ownership: snapshot markers, then clear.
stale_view = types.SimpleNamespace(
_gateway_turn_process_task_id=run._gateway_turn_process_task_id,
_gateway_turn_process_baseline=run._gateway_turn_process_baseline,
_gateway_turn_process_epoch=run._gateway_turn_process_epoch,
)
_clear_turn_process_ownership(run)
_reap_disconnected_agent_processes(stale_view)
deadline = time.time() + 1.0
while not calls and time.time() < deadline:
time.sleep(0.01)
assert calls == [True]
def test_publish_and_clear_ownership_roundtrip(self, monkeypatch):
from gateway.platforms.api_server import (
_TURN_PROCESS_EPOCHS,
_clear_turn_process_ownership,
_publish_turn_process_ownership,
)
from tools.process_registry import process_registry
monkeypatch.setattr(
process_registry,
"snapshot_running_ids",
lambda tid: frozenset({f"pre-{tid}"}),
)
agent = types.SimpleNamespace()
_publish_turn_process_ownership(agent, "sess-rt")
assert agent._gateway_turn_process_task_id == "sess-rt"
assert agent._gateway_turn_process_baseline == frozenset({"pre-sess-rt"})
assert isinstance(agent._gateway_turn_process_epoch, int)
assert "sess-rt" in _TURN_PROCESS_EPOCHS
_clear_turn_process_ownership(agent)
assert agent._gateway_turn_process_task_id == ""
assert agent._gateway_turn_process_baseline == frozenset()
assert agent._gateway_turn_process_epoch is None
# Entry pruned — dict stays bounded to in-flight runs.
assert "sess-rt" not in _TURN_PROCESS_EPOCHS
@pytest.mark.asyncio
async def test_stop_run_reaps_owned_processes(self, adapter, monkeypatch):
"""POST /v1/runs/{id}/stop abandons the run — it must reap the
background processes that run created (#76115 sibling surface)."""
from gateway.platforms.api_server import _publish_turn_process_ownership
from tools.process_registry import process_registry
calls = []
monkeypatch.setattr(
process_registry,
"kill_started_since",
lambda task_id, baseline, *, source: calls.append(
(task_id, baseline, source)
)
or 1,
)
monkeypatch.setattr(
process_registry, "snapshot_running_ids", lambda _tid: frozenset()
)
agent = MagicMock()
_publish_turn_process_ownership(agent, "run-stop-sess")
adapter._active_run_agents["run_x"] = agent
request = MagicMock()
request.match_info = {"run_id": "run_x"}
resp = await adapter._handle_stop_run(request)
assert resp.status == 200
deadline = time.time() + 1.0
while not calls and time.time() < deadline:
time.sleep(0.01)
assert calls == [("run-stop-sess", frozenset(), "api_server_run_stop")]
agent.interrupt.assert_called_once()
class TestRunEventCallback: