fix(gateway): respect routed profile busy modes

This commit is contained in:
Trevin Chow
2026-08-10 17:36:23 -07:00
committed by Teknium
parent 9d6c5a920c
commit a31be48030
2 changed files with 484 additions and 12 deletions

View File

@@ -6020,6 +6020,11 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
self._show_reasoning = self._load_show_reasoning()
self._busy_input_mode = self._load_busy_input_mode()
self._busy_text_mode = self._load_busy_text_mode()
# Secondary-profile busy modes are snapshotted during multiplex
# startup. Busy-message handlers consult these maps by routed source
# without rereading config or mutating process-global environment.
self._busy_input_modes_by_profile: Dict[str, str] = {}
self._busy_text_modes_by_profile: Dict[str, str] = {}
self._restart_drain_timeout = self._load_restart_drain_timeout()
self._restart_after_turn_timeout = self._load_restart_after_turn_timeout()
self._provider_routing = self._load_provider_routing()
@@ -7883,11 +7888,14 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
def _status_action_gerund(self) -> str:
return "restarting" if self._restart_requested else "shutting down"
def _queue_during_drain_enabled(self) -> bool:
def _queue_during_drain_enabled(
self, busy_input_mode: Optional[str] = None
) -> bool:
# Both "queue" and "steer" modes imply the user doesn't want messages
# to be lost during restart — queue them for the newly-spawned gateway
# process to pick up. "interrupt" mode drops them (current behaviour).
return self._restart_requested and self._busy_input_mode in {"queue", "steer"}
mode = busy_input_mode or self._busy_input_mode
return self._restart_requested and mode in {"queue", "steer"}
# -------- /queue FIFO helpers --------------------------------------
# /queue must produce one full agent turn per invocation, in FIFO
@@ -8525,6 +8533,76 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
input_mode = GatewayRunner._load_busy_input_mode()
return "queue" if input_mode == "queue" else "interrupt"
@staticmethod
def _busy_modes_from_config(
config: dict,
*,
fallback_input: str,
fallback_text: str,
) -> tuple[str, str]:
"""Resolve one profile's busy modes without consulting process env."""
raw_input = str(
cfg_get(config, "display", "busy_input_mode", default="") or ""
).strip().lower()
input_mode = (
raw_input
if raw_input in {"interrupt", "queue", "steer"}
else fallback_input
)
raw_text = str(
cfg_get(config, "display", "busy_text_mode", default="") or ""
).strip().lower()
if raw_text in {"interrupt", "queue"}:
text_mode = raw_text
elif raw_input in {"interrupt", "queue", "steer"}:
text_mode = "queue" if input_mode == "queue" else "interrupt"
else:
text_mode = fallback_text
return input_mode, text_mode
def _snapshot_profile_busy_modes(self, profile_name: str, config: dict) -> None:
"""Cache a routed profile's busy policy for this gateway lifetime."""
input_mode, text_mode = self._busy_modes_from_config(
config,
fallback_input=getattr(self, "_busy_input_mode", "interrupt"),
fallback_text=getattr(self, "_busy_text_mode", "interrupt"),
)
input_modes = self.__dict__.setdefault("_busy_input_modes_by_profile", {})
text_modes = self.__dict__.setdefault("_busy_text_modes_by_profile", {})
input_modes[profile_name] = input_mode
text_modes[profile_name] = text_mode
def _busy_profile_name_for_source(self, source: SessionSource) -> Optional[str]:
"""Return the routed profile whose busy policy applies, if any."""
if not getattr(getattr(self, "config", None), "multiplex_profiles", False):
return None
name = str(getattr(source, "profile", "") or "").strip()
if not name:
try:
name = str(self._profile_name_for_source(source) or "").strip()
except Exception:
name = ""
return name or None
def _effective_busy_input_mode(self, source: SessionSource) -> str:
"""Resolve busy input mode from the routed profile startup snapshot."""
fallback = getattr(self, "_busy_input_mode", "interrupt")
profile_name = self._busy_profile_name_for_source(source)
if not profile_name:
return fallback
modes = getattr(self, "_busy_input_modes_by_profile", None)
return modes.get(profile_name, fallback) if isinstance(modes, dict) else fallback
def _effective_busy_text_mode(self, source: SessionSource) -> str:
"""Resolve legacy busy text mode from the routed profile snapshot."""
fallback = getattr(self, "_busy_text_mode", "interrupt")
profile_name = self._busy_profile_name_for_source(source)
if not profile_name:
return fallback
modes = getattr(self, "_busy_text_modes_by_profile", None)
return modes.get(profile_name, fallback) if isinstance(modes, dict) else fallback
@staticmethod
def _load_restart_drain_timeout() -> float:
"""Load graceful gateway restart/stop drain timeout in seconds."""
@@ -8970,6 +9048,8 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
return True # handled (silently dropped); do not fall through
effective_mode = self._effective_busy_input_mode(event.source)
# --- Draining case (gateway restarting/stopping) ---
if self._draining:
adapter = self._adapter_for_source(event.source)
@@ -8978,7 +9058,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
reply_anchor = self._reply_anchor_for_event(event)
thread_meta = self._thread_metadata_for_source(event.source, reply_anchor)
if self._queue_during_drain_enabled():
if self._queue_during_drain_enabled(effective_mode):
self._queue_or_replace_pending_event(session_key, event)
message = f"⏳ Gateway {self._status_action_gerund()} — queued for the next turn after it comes back."
else:
@@ -9095,8 +9175,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
_busy_state = self._peek_session_state(session_key)
running_agent = _busy_state.turn.agent if _busy_state else None
effective_mode = self._busy_input_mode
busy_text_mode = getattr(self, "_busy_text_mode", "interrupt")
busy_text_mode = self._effective_busy_text_mode(event.source)
if (
event.message_type == MessageType.TEXT
and busy_text_mode == "queue"
@@ -13577,8 +13656,10 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
from gateway.config import load_gateway_config
with _profile_runtime_scope(profile_home):
profile_runtime_cfg = _load_gateway_runtime_config()
profile_cfg = load_gateway_config()
violation = _own_policy_open_startup_violation(profile_cfg)
self._snapshot_profile_busy_modes(profile_name, profile_runtime_cfg)
if violation:
raise MultiplexConfigError(
f"Profile '{profile_name}' enables {violation}. "
@@ -13712,7 +13793,9 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
self._make_profile_fatal_error_handler(profile_name, platform)
)
adapter.set_session_store(self.session_store)
adapter.set_busy_session_handler(self._handle_active_session_busy_message)
adapter.set_busy_session_handler(
self._make_profile_busy_session_handler(profile_name)
)
_set_reaction = getattr(adapter, "set_reaction_handler", None)
if callable(_set_reaction):
_set_reaction(self._handle_reaction_event)
@@ -13720,7 +13803,12 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
adapter.set_authorization_check(
self._make_adapter_auth_check(platform, profile_name=profile_name)
)
adapter._busy_text_mode = self._busy_text_mode
text_modes = getattr(self, "_busy_text_modes_by_profile", None)
adapter._busy_text_mode = (
text_modes.get(profile_name, self._busy_text_mode)
if isinstance(text_modes, dict)
else self._busy_text_mode
)
async def _run_secondary_profile_reconnect(
self, profile_name: str, platform: Platform
@@ -13917,6 +14005,21 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
return _handler
def _make_profile_busy_session_handler(self, profile_name: str):
"""Stamp an owning adapter's profile before resolving busy policy."""
async def _handler(event, _session_key):
try:
if getattr(event, "source", None) is not None and not event.source.profile:
event.source.profile = profile_name
except Exception:
pass
routed_session_key = self._session_key_for_source(event.source)
return await self._handle_active_session_busy_message(
event, routed_session_key
)
return _handler
def _make_default_profile_message_handler(self):
"""Scope a multiplexed default-profile message from ingress onward."""
profile_home = Path(get_hermes_home())
@@ -15225,6 +15328,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
merge_pending_message_event(adapter._pending_messages, _quick_key, event)
return None
effective_busy_input_mode = self._effective_busy_input_mode(source)
_telegram_followup_grace = float(
os.getenv("HERMES_TELEGRAM_FOLLOWUP_GRACE_SECONDS", "3.0")
)
@@ -15244,7 +15348,7 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
adapter = self._adapter_for_source(source)
if adapter:
if self._busy_input_mode == "queue":
if effective_busy_input_mode == "queue":
self._enqueue_fifo(_quick_key, event, adapter)
else:
merge_pending_message_event(
@@ -15276,18 +15380,21 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew
)
return None
if self._draining:
if self._queue_during_drain_enabled():
queue_during_drain = self._queue_during_drain_enabled(
effective_busy_input_mode
)
if queue_during_drain:
self._queue_or_replace_pending_event(_quick_key, event)
return (
f"⏳ Gateway {self._status_action_gerund()} — queued for the next turn after it comes back."
if self._queue_during_drain_enabled()
if queue_during_drain
else f"⏳ Gateway is {self._status_action_gerund()} and is not accepting another turn right now."
)
if self._busy_input_mode == "queue":
if effective_busy_input_mode == "queue":
logger.debug("PRIORITY queue follow-up for session %s", _quick_key)
self._queue_or_replace_pending_event(_quick_key, event)
return None
if self._busy_input_mode == "steer":
if effective_busy_input_mode == "steer":
# Steer mode: inject text into the running agent mid-run via
# agent.steer(). Falls back to queue semantics if the payload
# is empty, the agent lacks steer(), or steer() rejects.

View File

@@ -0,0 +1,365 @@
"""Profile-specific busy-input behavior for multiplexed gateways."""
import asyncio
from unittest.mock import AsyncMock, MagicMock
import pytest
from gateway.config import GatewayConfig, Platform, PlatformConfig
from gateway.platforms.base import (
BasePlatformAdapter,
MessageEvent,
MessageType,
SendResult,
SessionSource,
build_session_key,
)
from gateway.profile_routing import ProfileRoute
from gateway.run import GatewayRunner
class _ProfileAdapter(BasePlatformAdapter):
async def connect(self, *, is_reconnect: bool = False) -> bool:
return True
async def disconnect(self):
pass
async def send(self, chat_id, content, reply_to=None, metadata=None) -> SendResult:
return SendResult(success=True)
async def get_chat_info(self, chat_id):
return {}
def _runner(*, default_mode: str = "interrupt") -> GatewayRunner:
runner = GatewayRunner.__new__(GatewayRunner)
runner.config = GatewayConfig(multiplex_profiles=True)
runner._busy_input_mode = default_mode
runner._busy_text_mode = "queue" if default_mode == "queue" else "interrupt"
runner._profile_adapters = {}
runner.adapters = {}
runner._sessions = {}
runner._draining = False
runner._restart_requested = False
runner.session_store = None
runner.hooks = MagicMock()
runner.hooks.emit = AsyncMock()
runner.pairing_store = MagicMock()
runner.pairing_store.is_approved.return_value = True
runner._is_user_authorized = lambda source: True
runner._session_has_compression_in_flight = AsyncMock(return_value=False)
return runner
def _event(*, profile: str | None) -> MessageEvent:
return MessageEvent(
text="follow up",
message_type=MessageType.TEXT,
source=SessionSource(
platform=Platform.TELEGRAM,
chat_id="chat-1",
chat_type="dm",
user_id="user-1",
profile=profile,
),
message_id="message-1",
)
def _adapter() -> _ProfileAdapter:
adapter = _ProfileAdapter(
PlatformConfig(enabled=True, token="test-token"),
Platform.TELEGRAM,
)
return adapter
async def _load_profile_snapshot(
runner: GatewayRunner,
profile_home,
mode: str | None,
*,
legacy_text_mode: str | None = None,
) -> _ProfileAdapter:
display = "display:\n"
if mode is not None:
display += f" busy_input_mode: {mode}\n"
if legacy_text_mode is not None:
display += f" busy_text_mode: {legacy_text_mode}\n"
profile_home.mkdir()
(profile_home / "config.yaml").write_text(display, encoding="utf-8")
assert await runner._start_one_profile_adapters("research", profile_home, {}) == 0
adapter = _adapter()
runner._profile_adapters["research"][Platform.TELEGRAM] = adapter
runner._configure_profile_adapter(adapter, "research", Platform.TELEGRAM)
return adapter
@pytest.mark.asyncio
@pytest.mark.parametrize(
("secondary_mode", "handled", "expected_action", "expected_text_mode"),
[
("queue", False, "queue", "queue"),
("steer", True, "steer", "interrupt"),
("interrupt", True, "interrupt", "interrupt"),
],
)
async def test_secondary_profile_busy_mode_controls_live_busy_behavior(
tmp_path,
monkeypatch,
secondary_mode,
handled,
expected_action,
expected_text_mode,
):
"""A routed profile chooses queue/steer/interrupt independently."""
monkeypatch.setenv("HERMES_GATEWAY_BUSY_ACK_ENABLED", "false")
runner = _runner(default_mode="interrupt")
adapter = await _load_profile_snapshot(
runner,
tmp_path / "research",
secondary_mode,
)
event = _event(profile="research")
session_key = runner._session_key_for_source(event.source)
agent = MagicMock()
agent._active_children = []
agent.steer.return_value = True
runner._running_agents[session_key] = agent
result = await runner._handle_active_session_busy_message(event, session_key)
assert result is handled
assert adapter._busy_text_mode == expected_text_mode
if expected_action == "queue":
agent.steer.assert_not_called()
agent.interrupt.assert_not_called()
elif expected_action == "steer":
agent.steer.assert_called_once_with("follow up")
agent.interrupt.assert_not_called()
else:
agent.steer.assert_not_called()
agent.interrupt.assert_called_once_with("follow up")
@pytest.mark.asyncio
@pytest.mark.parametrize("secondary_mode", ["queue", "steer"])
async def test_secondary_profile_busy_mode_controls_priority_path(
tmp_path,
monkeypatch,
secondary_mode,
):
"""The runner's early active-agent path uses the same routed policy."""
monkeypatch.setenv("HERMES_TELEGRAM_FOLLOWUP_GRACE_SECONDS", "0")
runner = _runner(default_mode="interrupt")
adapter = await _load_profile_snapshot(
runner,
tmp_path / "research",
secondary_mode,
)
event = _event(profile="research")
session_key = runner._session_key_for_source(event.source)
agent = MagicMock()
agent._active_children = []
agent.steer.return_value = True
runner._running_agents[session_key] = agent
assert await runner._handle_message(event) is None
agent.interrupt.assert_not_called()
if secondary_mode == "queue":
agent.steer.assert_not_called()
assert adapter._pending_messages[session_key] is event
else:
agent.steer.assert_called_once_with("follow up")
assert session_key not in adapter._pending_messages
@pytest.mark.asyncio
@pytest.mark.parametrize(
("default_mode", "secondary_mode", "queued"),
[
("interrupt", "queue", True),
("queue", "interrupt", False),
],
)
async def test_secondary_profile_busy_mode_controls_busy_handler_restart_drain(
tmp_path,
default_mode,
secondary_mode,
queued,
):
runner = _runner(default_mode=default_mode)
adapter = await _load_profile_snapshot(
runner,
tmp_path / "research",
secondary_mode,
)
runner._draining = True
runner._restart_requested = True
event = _event(profile="research")
session_key = runner._session_key_for_source(event.source)
assert await runner._handle_active_session_busy_message(event, session_key) is True
assert (session_key in adapter._pending_messages) is queued
@pytest.mark.asyncio
async def test_secondary_profile_busy_mode_controls_priority_restart_drain(
tmp_path,
monkeypatch,
):
monkeypatch.setenv("HERMES_TELEGRAM_FOLLOWUP_GRACE_SECONDS", "0")
runner = _runner(default_mode="interrupt")
adapter = await _load_profile_snapshot(
runner,
tmp_path / "research",
"queue",
)
runner._draining = True
runner._restart_requested = True
event = _event(profile="research")
session_key = runner._session_key_for_source(event.source)
agent = MagicMock()
agent._active_children = []
runner._running_agents[session_key] = agent
response = await runner._handle_message(event)
assert isinstance(response, str)
assert "queued" in response
assert adapter._pending_messages[session_key] is event
agent.interrupt.assert_not_called()
@pytest.mark.asyncio
async def test_secondary_adapter_busy_guard_stamps_profile_before_resolving_mode(
tmp_path,
monkeypatch,
):
"""Per-profile adapters route busy events before the message wrapper runs."""
monkeypatch.setenv("HERMES_GATEWAY_BUSY_ACK_ENABLED", "false")
runner = _runner(default_mode="interrupt")
adapter = await _load_profile_snapshot(
runner,
tmp_path / "research",
"steer",
)
event = _event(profile=None)
adapter_session_key = build_session_key(event.source)
adapter._active_sessions[adapter_session_key] = asyncio.Event()
routed_source = _event(profile="research").source
routed_session_key = runner._session_key_for_source(routed_source)
agent = MagicMock()
agent._active_children = []
agent.steer.return_value = True
runner._running_agents[routed_session_key] = agent
await adapter.handle_message(event)
assert event.source.profile == "research"
agent.steer.assert_called_once_with("follow up")
agent.interrupt.assert_not_called()
@pytest.mark.asyncio
async def test_secondary_legacy_busy_text_mode_is_profile_specific(tmp_path):
runner = _runner(default_mode="interrupt")
adapter = await _load_profile_snapshot(
runner,
tmp_path / "research",
"interrupt",
legacy_text_mode="queue",
)
event = _event(profile="research")
session_key = runner._session_key_for_source(event.source)
agent = MagicMock()
agent._active_children = []
runner._running_agents[session_key] = agent
assert await runner._handle_active_session_busy_message(event, session_key) is False
assert adapter._busy_text_mode == "queue"
agent.interrupt.assert_not_called()
@pytest.mark.asyncio
async def test_default_busy_mode_is_unchanged_by_secondary_profile(tmp_path, monkeypatch):
monkeypatch.setenv("HERMES_GATEWAY_BUSY_ACK_ENABLED", "false")
runner = _runner(default_mode="interrupt")
await _load_profile_snapshot(runner, tmp_path / "research", "steer")
adapter = _adapter()
runner.adapters[Platform.TELEGRAM] = adapter
event = _event(profile=None)
session_key = runner._session_key_for_source(event.source)
agent = MagicMock()
agent._active_children = []
runner._running_agents[session_key] = agent
assert await runner._handle_active_session_busy_message(event, session_key) is True
agent.interrupt.assert_called_once_with("follow up")
agent.steer.assert_not_called()
assert runner._busy_input_mode == "interrupt"
@pytest.mark.asyncio
@pytest.mark.parametrize("secondary_mode", [None, "not-a-mode"])
async def test_missing_or_invalid_secondary_mode_falls_back_to_gateway_default(
tmp_path,
secondary_mode,
):
runner = _runner(default_mode="queue")
await _load_profile_snapshot(runner, tmp_path / "research", secondary_mode)
source = _event(profile="research").source
assert runner._effective_busy_input_mode(source) == "queue"
assert runner._effective_busy_text_mode(source) == "queue"
assert runner._busy_input_mode == "queue"
assert runner._busy_text_mode == "queue"
def test_profile_route_and_nonmultiplexed_resolution_preserve_boundaries():
runner = _runner(default_mode="interrupt")
runner._snapshot_profile_busy_modes(
"research",
{"display": {"busy_input_mode": "steer"}},
)
runner.config.profile_routes = [
ProfileRoute(
name="research-chat",
platform="telegram",
profile="research",
chat_id="chat-1",
)
]
source = _event(profile=None).source
assert runner._effective_busy_input_mode(source) == "steer"
runner.config.multiplex_profiles = False
source.profile = "research"
assert runner._effective_busy_input_mode(source) == "interrupt"
@pytest.mark.asyncio
async def test_effective_mode_uses_startup_snapshot_without_rereading_config(
tmp_path,
monkeypatch,
):
import gateway.run as gateway_run
runner = _runner(default_mode="interrupt")
await _load_profile_snapshot(runner, tmp_path / "research", "steer")
source = _event(profile="research").source
def fail_config_read():
raise AssertionError("busy-mode lookup reread config after startup")
monkeypatch.setattr(gateway_run, "_load_gateway_runtime_config", fail_config_read)
monkeypatch.setattr(gateway_run, "_load_gateway_config", fail_config_read)
assert runner._effective_busy_input_mode(source) == "steer"
assert runner._effective_busy_input_mode(source) == "steer"
assert runner._effective_busy_text_mode(source) == "interrupt"