fix(gateway): recheck the heartbeat guard after an awaited edit

Review finding: the guard ran before `edit_message` was awaited; a restart
notice sent during that await followed by a failed edit produced a fresh
"Working" fallback bubble after the notice. Recheck before the fallback
send; the notifier ends instead.
This commit is contained in:
teknium1
2026-09-12 08:20:03 -07:00
committed by Teknium
parent 0fc3cbac59
commit ec8b4b2d6f
2 changed files with 43 additions and 0 deletions
+6
View File
@@ -3854,6 +3854,12 @@ class GatewayTurnMixin:
logger.debug("Heartbeat edit failed: %s", _ee)
_notify_res = None
if not (_notify_res and getattr(_notify_res, "success", False)):
# The edit above awaited; a drain/restart notice may have gone out meanwhile, and
# a fresh "Working" bubble after it reads as a contradiction (#10990).
if not self._should_emit_long_running_notification(
session_key, agent_holder[0], _executor_task_holder[0]
):
break
_notify_res = await _notify_adapter.send(
source.chat_id, _heartbeat_text,
metadata=_interim_metadata(_non_conversational_metadata(_status_thread_metadata, platform=source.platform)),
+37
View File
@@ -518,6 +518,43 @@ class TestLongRunningNotificationOwnership:
"sess", original_agent, executor_task=None
) is False
@pytest.mark.asyncio
async def test_restart_during_heartbeat_edit_sends_no_fallback_bubble(self, monkeypatch):
"""The guard is rechecked after the awaited edit: a restart that begins while the edit is
in flight must not be followed by a fresh "Working" send when that edit fails (#10990)."""
import asyncio
from types import SimpleNamespace
from gateway.run import GatewayRunner
from gateway.turn_context import TurnContext
monkeypatch.setenv("HERMES_AGENT_NOTIFY_INTERVAL", "0.01")
runner = object.__new__(GatewayRunner)
runner._running_agents = {}
runner._draining = runner._restart_requested = False
adapter = MagicMock()
first_send = AsyncMock(return_value=SimpleNamespace(success=True, message_id="hb-1"))
adapter.send = first_send
async def _edit_then_restart(*a, **k):
runner._restart_requested = True # restart notice goes out while the edit is awaited
return SimpleNamespace(success=False)
adapter.edit_message = AsyncMock(side_effect=_edit_then_restart)
runner._adapter_for_source = lambda source: adapter
runner._agent_activity_summary = staticmethod(lambda agent: None)
agent = MagicMock()
runner._running_agents["sess"] = agent
disp = MagicMock()
disp._display_surface_mode.return_value = "on"
disp.resolve_display_setting.return_value = False
ctx = TurnContext(source=SimpleNamespace(chat_id="c", platform="telegram"), session_key="sess")
ctx.agent_holder[0] = agent
await asyncio.wait_for(runner._run_agent_notify_long_running(disp, ctx, [None]), 5)
assert first_send.await_count == 1 # the original heartbeat only
adapter.edit_message.assert_awaited_once()
@pytest.mark.parametrize("flag", ["_draining", "_restart_requested"])
def test_notification_stops_once_shutdown_or_restart_begins(self, flag):
"""After the restart/shutdown notice a heartbeat would contradict it (#10990)."""