mirror of
https://github.com/NousResearch/hermes-agent.git
synced 2026-04-25 00:51:20 +00:00
476 lines
18 KiB
Python
476 lines
18 KiB
Python
"""Tests for the session race guard that prevents concurrent agent runs.
|
|
|
|
The sentinel-based guard ensures that when _handle_message passes the
|
|
"is an agent already running?" check and proceeds to the slow async
|
|
setup path (vision enrichment, STT, hooks, session hygiene), a second
|
|
message for the same session is correctly recognized as "already running"
|
|
and routed through the interrupt/queue path instead of spawning a
|
|
duplicate agent.
|
|
"""
|
|
|
|
import asyncio
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from gateway.config import GatewayConfig, Platform, PlatformConfig
|
|
from gateway.platforms.base import MessageEvent, MessageType, merge_pending_message_event
|
|
from gateway.run import GatewayRunner, _AGENT_PENDING_SENTINEL
|
|
from gateway.session import SessionSource, build_session_key
|
|
|
|
|
|
class _FakeAdapter:
|
|
"""Minimal adapter stub for testing."""
|
|
|
|
def __init__(self):
|
|
self._pending_messages = {}
|
|
self._active_sessions = {}
|
|
self.interrupted_sessions = []
|
|
|
|
async def send(self, chat_id, text, **kwargs):
|
|
pass
|
|
|
|
async def interrupt_session_activity(self, session_key, chat_id):
|
|
self.interrupted_sessions.append((session_key, chat_id))
|
|
event = self._active_sessions.get(session_key)
|
|
if event is not None:
|
|
event.set()
|
|
|
|
|
|
def _make_runner():
|
|
runner = object.__new__(GatewayRunner)
|
|
runner.config = GatewayConfig(
|
|
platforms={Platform.TELEGRAM: PlatformConfig(enabled=True, token="***")}
|
|
)
|
|
runner.adapters = {Platform.TELEGRAM: _FakeAdapter()}
|
|
runner._running_agents = {}
|
|
runner._running_agents_ts = {}
|
|
runner._session_run_generation = {}
|
|
runner._pending_messages = {}
|
|
runner._pending_approvals = {}
|
|
runner._voice_mode = {}
|
|
runner._background_tasks = set()
|
|
runner._draining = False
|
|
runner._restart_requested = False
|
|
runner._restart_task_started = False
|
|
runner._restart_detached = False
|
|
runner._restart_via_service = False
|
|
runner._restart_drain_timeout = 0.0
|
|
runner._stop_task = None
|
|
runner._exit_code = None
|
|
runner._update_runtime_status = MagicMock()
|
|
runner._is_user_authorized = lambda _source: True
|
|
runner.hooks = MagicMock()
|
|
runner.hooks.emit = AsyncMock()
|
|
runner.session_store = MagicMock()
|
|
runner.delivery_router = MagicMock()
|
|
return runner
|
|
|
|
|
|
def _make_event(text="hello", chat_id="12345"):
|
|
source = SessionSource(
|
|
platform=Platform.TELEGRAM, chat_id=chat_id, chat_type="dm",
|
|
user_id="u1",
|
|
)
|
|
return MessageEvent(text=text, message_type=MessageType.TEXT, source=source)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 1: Sentinel is placed before _handle_message_with_agent runs
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_sentinel_placed_before_agent_setup():
|
|
"""After passing the 'not running' guard, the sentinel must be
|
|
written into _running_agents *before* any await, so that a
|
|
concurrent message sees the session as occupied."""
|
|
runner = _make_runner()
|
|
event = _make_event()
|
|
session_key = build_session_key(event.source)
|
|
|
|
# Patch _handle_message_with_agent to capture state at entry
|
|
sentinel_was_set = False
|
|
|
|
async def mock_inner(self_inner, ev, src, qk, generation):
|
|
nonlocal sentinel_was_set
|
|
sentinel_was_set = runner._running_agents.get(qk) is _AGENT_PENDING_SENTINEL
|
|
return "ok"
|
|
|
|
with patch.object(GatewayRunner, "_handle_message_with_agent", mock_inner):
|
|
await runner._handle_message(event)
|
|
|
|
assert sentinel_was_set, (
|
|
"Sentinel must be in _running_agents when _handle_message_with_agent starts"
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 2: Sentinel is cleaned up after _handle_message_with_agent
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_sentinel_cleaned_up_after_handler_returns():
|
|
"""If _handle_message_with_agent returns normally, the sentinel
|
|
must be removed so the session is not permanently locked."""
|
|
runner = _make_runner()
|
|
event = _make_event()
|
|
session_key = build_session_key(event.source)
|
|
|
|
async def mock_inner(self_inner, ev, src, qk, generation):
|
|
return "ok"
|
|
|
|
with patch.object(GatewayRunner, "_handle_message_with_agent", mock_inner):
|
|
await runner._handle_message(event)
|
|
|
|
assert session_key not in runner._running_agents, (
|
|
"Sentinel must be removed after handler completes"
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 3: Sentinel cleaned up on exception
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_sentinel_cleaned_up_on_exception():
|
|
"""If _handle_message_with_agent raises, the sentinel must still
|
|
be cleaned up so the session is not permanently locked."""
|
|
runner = _make_runner()
|
|
event = _make_event()
|
|
session_key = build_session_key(event.source)
|
|
|
|
async def mock_inner(self_inner, ev, src, qk, generation):
|
|
raise RuntimeError("boom")
|
|
|
|
with patch.object(GatewayRunner, "_handle_message_with_agent", mock_inner):
|
|
with pytest.raises(RuntimeError, match="boom"):
|
|
await runner._handle_message(event)
|
|
|
|
assert session_key not in runner._running_agents, (
|
|
"Sentinel must be removed even if handler raises"
|
|
)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 4: Second message during sentinel sees "already running"
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_second_message_during_sentinel_queued_not_duplicate():
|
|
"""While the sentinel is set (agent setup in progress), a second
|
|
message for the same session must hit the 'already running' branch
|
|
and be queued — not start a second agent."""
|
|
runner = _make_runner()
|
|
event1 = _make_event(text="first message")
|
|
event2 = _make_event(text="second message")
|
|
session_key = build_session_key(event1.source)
|
|
|
|
barrier = asyncio.Event()
|
|
|
|
async def slow_inner(self_inner, ev, src, qk, generation):
|
|
# Simulate slow setup — wait until test tells us to proceed
|
|
await barrier.wait()
|
|
return "ok"
|
|
|
|
with patch.object(GatewayRunner, "_handle_message_with_agent", slow_inner):
|
|
# Start first message (will block at barrier)
|
|
task1 = asyncio.create_task(runner._handle_message(event1))
|
|
# Yield so task1 enters slow_inner and sentinel is set
|
|
await asyncio.sleep(0)
|
|
|
|
# Verify sentinel is set
|
|
assert runner._running_agents.get(session_key) is _AGENT_PENDING_SENTINEL
|
|
|
|
# Second message should see "already running" and be queued
|
|
result2 = await runner._handle_message(event2)
|
|
assert result2 is None, "Second message should return None (queued)"
|
|
|
|
# The second message should have been queued in adapter pending
|
|
adapter = runner.adapters[Platform.TELEGRAM]
|
|
assert session_key in adapter._pending_messages, (
|
|
"Second message should be queued as pending"
|
|
)
|
|
assert adapter._pending_messages[session_key] is event2
|
|
|
|
# Let first message complete
|
|
barrier.set()
|
|
await task1
|
|
|
|
|
|
def test_merge_pending_message_event_merges_text_and_photo_followups():
|
|
pending = {}
|
|
source = SessionSource(
|
|
platform=Platform.TELEGRAM,
|
|
chat_id="12345",
|
|
chat_type="dm",
|
|
user_id="u1",
|
|
)
|
|
session_key = build_session_key(source)
|
|
|
|
text_event = MessageEvent(
|
|
text="first follow-up",
|
|
message_type=MessageType.TEXT,
|
|
source=source,
|
|
)
|
|
photo_event = MessageEvent(
|
|
text="see screenshot",
|
|
message_type=MessageType.PHOTO,
|
|
source=source,
|
|
media_urls=["/tmp/test.png"],
|
|
media_types=["image/png"],
|
|
)
|
|
|
|
merge_pending_message_event(pending, session_key, text_event, merge_text=True)
|
|
merge_pending_message_event(pending, session_key, photo_event, merge_text=True)
|
|
|
|
merged = pending[session_key]
|
|
assert merged.message_type == MessageType.PHOTO
|
|
assert merged.text == "first follow-up\n\nsee screenshot"
|
|
assert merged.media_urls == ["/tmp/test.png"]
|
|
assert merged.media_types == ["image/png"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recent_telegram_text_followup_is_queued_without_interrupt():
|
|
runner = _make_runner()
|
|
event = _make_event(text="follow-up")
|
|
session_key = build_session_key(event.source)
|
|
|
|
fake_agent = MagicMock()
|
|
fake_agent.get_activity_summary.return_value = {"seconds_since_activity": 0}
|
|
runner._running_agents[session_key] = fake_agent
|
|
import time as _time
|
|
runner._running_agents_ts[session_key] = _time.time()
|
|
|
|
result = await runner._handle_message(event)
|
|
|
|
assert result is None
|
|
fake_agent.interrupt.assert_not_called()
|
|
adapter = runner.adapters[Platform.TELEGRAM]
|
|
assert adapter._pending_messages[session_key].text == "follow-up"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_recent_telegram_followups_append_in_pending_queue():
|
|
runner = _make_runner()
|
|
first = _make_event(text="part one")
|
|
second = _make_event(text="part two")
|
|
session_key = build_session_key(first.source)
|
|
|
|
fake_agent = MagicMock()
|
|
fake_agent.get_activity_summary.return_value = {"seconds_since_activity": 0}
|
|
runner._running_agents[session_key] = fake_agent
|
|
import time as _time
|
|
runner._running_agents_ts[session_key] = _time.time()
|
|
|
|
await runner._handle_message(first)
|
|
await runner._handle_message(second)
|
|
|
|
fake_agent.interrupt.assert_not_called()
|
|
adapter = runner.adapters[Platform.TELEGRAM]
|
|
assert adapter._pending_messages[session_key].text == "part one\npart two"
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 5: Sentinel not placed for command messages
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_command_messages_do_not_leave_sentinel():
|
|
"""Slash commands (/help, /status, etc.) return early from
|
|
_handle_message. They must NOT leave a sentinel behind."""
|
|
runner = _make_runner()
|
|
source = SessionSource(
|
|
platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm",
|
|
user_id="u1",
|
|
)
|
|
event = MessageEvent(
|
|
text="/help", message_type=MessageType.TEXT, source=source
|
|
)
|
|
session_key = build_session_key(source)
|
|
|
|
# Mock the help handler to avoid needing full runner setup
|
|
runner._handle_help_command = AsyncMock(return_value="Help text")
|
|
# Need hooks for command emission
|
|
runner.hooks = MagicMock()
|
|
runner.hooks.emit = AsyncMock()
|
|
|
|
await runner._handle_message(event)
|
|
|
|
assert session_key not in runner._running_agents, (
|
|
"Command handlers must not leave sentinel in _running_agents"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
@pytest.mark.parametrize(
|
|
("command_text", "handler_attr", "handler_result"),
|
|
[
|
|
("/help", "_handle_help_command", "Help text"),
|
|
("/commands", "_handle_commands_command", "Commands text"),
|
|
("/update", "_handle_update_command", "Update text"),
|
|
("/profile", "_handle_profile_command", "Profile text"),
|
|
],
|
|
)
|
|
async def test_active_session_bypass_commands_dispatch_without_interrupt(
|
|
command_text,
|
|
handler_attr,
|
|
handler_result,
|
|
):
|
|
"""Gateway-handled bypass commands must return directly while an agent runs."""
|
|
runner = _make_runner()
|
|
event = _make_event(text=command_text)
|
|
session_key = build_session_key(event.source)
|
|
|
|
fake_agent = MagicMock()
|
|
fake_agent.get_activity_summary.return_value = {"seconds_since_activity": 0}
|
|
runner._running_agents[session_key] = fake_agent
|
|
setattr(runner, handler_attr, AsyncMock(return_value=handler_result))
|
|
|
|
result = await runner._handle_message(event)
|
|
|
|
assert result == handler_result
|
|
fake_agent.interrupt.assert_not_called()
|
|
assert session_key not in runner.adapters[Platform.TELEGRAM]._pending_messages
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 6: /stop during sentinel force-cleans and unlocks session
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_stop_during_sentinel_force_cleans_session():
|
|
"""If /stop arrives while the sentinel is set (agent still starting),
|
|
it should force-clean the sentinel and unlock the session."""
|
|
runner = _make_runner()
|
|
event1 = _make_event(text="hello")
|
|
session_key = build_session_key(event1.source)
|
|
|
|
barrier = asyncio.Event()
|
|
|
|
async def slow_inner(self_inner, ev, src, qk, generation):
|
|
await barrier.wait()
|
|
return "ok"
|
|
|
|
with patch.object(GatewayRunner, "_handle_message_with_agent", slow_inner):
|
|
task1 = asyncio.create_task(runner._handle_message(event1))
|
|
await asyncio.sleep(0)
|
|
|
|
# Sentinel should be set
|
|
assert runner._running_agents.get(session_key) is _AGENT_PENDING_SENTINEL
|
|
|
|
# Send /stop — should force-clean the sentinel
|
|
stop_event = _make_event(text="/stop")
|
|
result = await runner._handle_message(stop_event)
|
|
assert result is not None, "/stop during sentinel should return a message"
|
|
assert "stopped" in result.lower()
|
|
assert session_key not in runner._running_agents, (
|
|
"/stop must remove sentinel so the session is unlocked"
|
|
)
|
|
|
|
# Should NOT be queued as pending
|
|
adapter = runner.adapters[Platform.TELEGRAM]
|
|
assert session_key not in adapter._pending_messages
|
|
|
|
barrier.set()
|
|
await task1
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 6b: /stop hard-kills a running agent and unlocks session
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_stop_hard_kills_running_agent():
|
|
"""When /stop arrives while a real agent is running, it must:
|
|
1. Call interrupt() on the agent
|
|
2. Force-clean _running_agents to unlock the session
|
|
3. Return a confirmation message
|
|
This fixes the bug where a hung agent kept the session locked
|
|
forever — showing 'writing...' but never producing output."""
|
|
runner = _make_runner()
|
|
session_key = build_session_key(
|
|
SessionSource(platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm", user_id="u1")
|
|
)
|
|
|
|
# Simulate a running (possibly hung) agent
|
|
fake_agent = MagicMock()
|
|
fake_agent.get_activity_summary.return_value = {"seconds_since_activity": 0}
|
|
runner._running_agents[session_key] = fake_agent
|
|
runner.adapters[Platform.TELEGRAM]._active_sessions[session_key] = asyncio.Event()
|
|
|
|
# Send /stop
|
|
stop_event = _make_event(text="/stop")
|
|
result = await runner._handle_message(stop_event)
|
|
|
|
# Agent must have been interrupted
|
|
fake_agent.interrupt.assert_called_once_with("Stop requested")
|
|
|
|
# Session must be unlocked
|
|
assert session_key not in runner._running_agents, (
|
|
"/stop must remove the agent from _running_agents so the session is unlocked"
|
|
)
|
|
assert runner.adapters[Platform.TELEGRAM].interrupted_sessions == [
|
|
(session_key, "12345")
|
|
]
|
|
assert runner.adapters[Platform.TELEGRAM]._active_sessions[session_key].is_set()
|
|
|
|
# Must return a confirmation
|
|
assert result is not None
|
|
assert "stopped" in result.lower()
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 6c: /stop clears pending messages to prevent stale replays
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_stop_clears_pending_messages():
|
|
"""When /stop hard-kills a running agent, any pending messages
|
|
queued during the run must be discarded."""
|
|
runner = _make_runner()
|
|
session_key = build_session_key(
|
|
SessionSource(platform=Platform.TELEGRAM, chat_id="12345", chat_type="dm", user_id="u1")
|
|
)
|
|
|
|
fake_agent = MagicMock()
|
|
fake_agent.get_activity_summary.return_value = {"seconds_since_activity": 0}
|
|
runner._running_agents[session_key] = fake_agent
|
|
runner._pending_messages[session_key] = "some queued text"
|
|
|
|
# Queue a pending message in the adapter too
|
|
adapter = runner.adapters[Platform.TELEGRAM]
|
|
adapter._pending_messages[session_key] = _make_event(text="queued")
|
|
adapter.get_pending_message = MagicMock(return_value=_make_event(text="queued"))
|
|
adapter.has_pending_interrupt = MagicMock(return_value=False)
|
|
|
|
stop_event = _make_event(text="/stop")
|
|
await runner._handle_message(stop_event)
|
|
|
|
# Pending messages must be cleared
|
|
assert session_key not in runner._pending_messages
|
|
adapter.get_pending_message.assert_called_once_with(session_key)
|
|
|
|
|
|
# ------------------------------------------------------------------
|
|
# Test 7: Shutdown skips sentinel entries
|
|
# ------------------------------------------------------------------
|
|
@pytest.mark.asyncio
|
|
async def test_shutdown_skips_sentinel():
|
|
"""During gateway shutdown, sentinel entries in _running_agents
|
|
should be skipped without raising AttributeError."""
|
|
runner = _make_runner()
|
|
session_key = "telegram:dm:99999"
|
|
|
|
# Simulate a sentinel in _running_agents
|
|
runner._running_agents[session_key] = _AGENT_PENDING_SENTINEL
|
|
|
|
# Also add a real agent mock to verify it still gets interrupted
|
|
real_agent = MagicMock()
|
|
runner._running_agents["telegram:dm:88888"] = real_agent
|
|
|
|
runner.adapters = {} # No adapters to disconnect
|
|
runner._running = True
|
|
runner._shutdown_event = asyncio.Event()
|
|
runner._exit_reason = None
|
|
runner._shutdown_all_gateway_honcho = lambda: None
|
|
|
|
with patch("gateway.status.remove_pid_file"), \
|
|
patch("gateway.status.write_runtime_status"):
|
|
await runner.stop()
|
|
|
|
# Real agent should have been interrupted
|
|
real_agent.interrupt.assert_called_once()
|
|
# Should not have raised on the sentinel
|