diff --git a/gateway/run.py b/gateway/run.py index c3aae82cf53a..d936f1bae976 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -19211,6 +19211,9 @@ def _register_heartbeat_watch(self, quick_key: str, source: Any, session_id: str The registry maps ``quick_key`` → ``(source, session_id)`` so the poller can rebuild a MessageEvent and enqueue via the adapter FIFO. + The stored session_id is only a starting point — the route can rotate + onto a new session underneath it, so each poll re-resolves it through + :meth:`_live_heartbeat_session_id`. In-memory by design: heartbeat STATE survives restarts in SessionDB, but firing resumes when the user touches /heartbeat again in the new gateway process (documented; durable schedules belong to cron). @@ -19227,6 +19230,30 @@ def _unregister_heartbeat_watch(self, quick_key: str) -> None: if watch: watch.pop(quick_key, None) + def _live_heartbeat_session_id(self, quick_key: str, fallback: str) -> str: + """Session id currently bound to a watched route. + + The registry caches the id the route had when /heartbeat was set, but + a context compression rotates the route onto a fresh session and + carries the heartbeat over with ``migrate_heartbeat_to_session``, which + marks the parent row cleared. Polling the captured id would then read + that cleared row, conclude the heartbeat was removed, and drop the + watch — silencing the heartbeat for the rest of the process while + ``/heartbeat status`` (which resolves the live session) still reports + it as active and counting down. + + ``peek_session_id`` is the read-only, lock-held accessor for the + key→session_id mapping that ``advance_compression_session`` repairs on + rotation. It never creates a session, so a route the user has since + ended stays absent and falls through to the caller's drop path. + """ + try: + sid = self.session_store.peek_session_id(quick_key) + except Exception as exc: + logger.debug("heartbeat: session lookup for %s failed: %s", quick_key, exc) + return fallback + return sid or fallback + def _start_heartbeat_poller(self) -> None: """Start the single gateway-wide heartbeat poll task (idempotent).""" existing = getattr(self, "_heartbeat_poll_task", None) @@ -19248,6 +19275,14 @@ async def _poll_loop(): continue from hermes_cli.heartbeat import HeartbeatManager + # Follow the route if it rotated onto a new session + # (compression) since the watch was registered, and + # cache the result so the next tick starts from it. + live_sid = self._live_heartbeat_session_id(quick_key, session_id) + if live_sid != session_id: + session_id = live_sid + watch[quick_key] = (source, session_id) + mgr = HeartbeatManager(session_id=session_id) if not mgr.has_heartbeat(): watch.pop(quick_key, None) diff --git a/tests/gateway/test_heartbeat_poller_session_rotation.py b/tests/gateway/test_heartbeat_poller_session_rotation.py new file mode 100644 index 000000000000..a92bdebf9222 --- /dev/null +++ b/tests/gateway/test_heartbeat_poller_session_rotation.py @@ -0,0 +1,157 @@ +"""Regression tests: the gateway heartbeat poller must follow session rotation. + +``/heartbeat`` registers a watch as ``quick_key -> (source, session_id)``, and +before this fix the poller kept polling that captured ``session_id`` forever. + +A context compression rotates the route onto a fresh session and carries the +heartbeat across with ``migrate_heartbeat_to_session``, which marks the parent +row ``cleared``. The poller, still holding the parent id, read that cleared row, +concluded the user had removed the heartbeat, and dropped the watch — so the +heartbeat went silent for the rest of the gateway process. Nothing surfaced it: +``/heartbeat status`` resolves the *live* session, so it kept reporting the +heartbeat as active with a ticking countdown. + +The poller now re-resolves the route's current session id every tick through +``_live_heartbeat_session_id`` (``SessionStore.peek_session_id``, the same +mapping ``advance_compression_session`` repairs on rotation). +""" +import asyncio +import time +from types import SimpleNamespace + +import pytest + +from hermes_cli.heartbeat import ( + HeartbeatState, + migrate_heartbeat_to_session, + save_heartbeat, +) + + +def _make_runner(routes, running=()): + """Bare GatewayRunner with just the collaborators the poller touches. + + ``routes`` is the persisted session-key → session-id mapping that + ``SessionStore.peek_session_id`` reads; mutating it mid-test is exactly what + a compression rotation does to the real store. + """ + from gateway.run import GatewayRunner + + runner = object.__new__(GatewayRunner) + runner.session_store = SimpleNamespace( + peek_session_id=lambda key: routes.get(key) + ) + # _running_agents is a mapping-backed session field view, not a plain set. + runner._running_agents = {key: object() for key in running} + runner._adapter_for_source = lambda source: object() + runner._background_tasks = set() + + enqueued = [] + runner._enqueue_fifo = lambda key, event, adapter: enqueued.append((key, event)) + return runner, enqueued + + +async def _let_poller_tick(runner, ticks=6): + """Yield to the poll loop long enough for a few iterations.""" + for _ in range(ticks): + await asyncio.sleep(0) + await asyncio.sleep(0.01) + task = getattr(runner, "_heartbeat_poll_task", None) + if task is not None: + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + +def _due_state(prompt="Check the deploy", interval=600): + """A heartbeat whose next tick is already due.""" + return HeartbeatState( + prompt=prompt, + interval_seconds=interval, + created_at=time.time() - (interval + 100), + ) + + +@pytest.mark.asyncio +async def test_poller_follows_compression_rotation(monkeypatch): + """The heartbeat keeps firing after the route rotates onto a new session.""" + monkeypatch.setattr("hermes_cli.heartbeat.POLL_SECONDS", 0.01) + + parent, child = "rot-parent-sid", "rot-child-sid" + save_heartbeat(parent, _due_state()) + + routes = {"route-rot": parent} + runner, enqueued = _make_runner(routes) + runner._register_heartbeat_watch("route-rot", "source-obj", parent) + + # Compression rotates the route and carries the heartbeat to the child. + assert migrate_heartbeat_to_session(parent, child) is True + routes["route-rot"] = child + + await _let_poller_tick(runner) + + assert enqueued, "heartbeat never fired after the session rotated" + key, event = enqueued[0] + assert key == "route-rot" + assert "Check the deploy" in event.text + + # The watch survived and now caches the child id, so the next tick starts + # from the live session instead of re-resolving from a stale one. + assert runner._heartbeat_watch["route-rot"] == ("source-obj", child) + + +@pytest.mark.asyncio +async def test_poller_drops_watch_when_heartbeat_is_gone(monkeypatch): + """A route with no heartbeat on its live session still unregisters.""" + monkeypatch.setattr("hermes_cli.heartbeat.POLL_SECONDS", 0.01) + + routes = {"route-gone": "gone-sid"} + runner, enqueued = _make_runner(routes) + runner._register_heartbeat_watch("route-gone", "source-obj", "gone-sid") + + await _let_poller_tick(runner) + + assert enqueued == [] + assert "route-gone" not in runner._heartbeat_watch + + +@pytest.mark.asyncio +async def test_busy_route_coalesces_instead_of_firing(monkeypatch): + """An in-flight turn defers the tick — unchanged by the rotation fix.""" + monkeypatch.setattr("hermes_cli.heartbeat.POLL_SECONDS", 0.01) + + save_heartbeat("busy-sid", _due_state()) + routes = {"route-busy": "busy-sid"} + runner, enqueued = _make_runner(routes, running=("route-busy",)) + runner._register_heartbeat_watch("route-busy", "source-obj", "busy-sid") + + await _let_poller_tick(runner) + + assert enqueued == [] + # Deferred, not dropped. + assert "route-busy" in runner._heartbeat_watch + + +def test_live_session_id_prefers_the_route_mapping(): + runner, _ = _make_runner({"route-a": "live-sid"}) + assert runner._live_heartbeat_session_id("route-a", "stale-sid") == "live-sid" + + +def test_live_session_id_falls_back_for_unknown_route(): + """An unmapped route keeps the captured id rather than resolving to None.""" + runner, _ = _make_runner({}) + assert runner._live_heartbeat_session_id("route-missing", "stale-sid") == "stale-sid" + + +def test_live_session_id_falls_back_when_the_store_raises(): + """A failing lookup must not take the poller down.""" + from gateway.run import GatewayRunner + + def _boom(_key): + raise RuntimeError("store unavailable") + + runner = object.__new__(GatewayRunner) + runner.session_store = SimpleNamespace(peek_session_id=_boom) + assert runner._live_heartbeat_session_id("route-x", "stale-sid") == "stale-sid"