diff --git a/band-bridge/bridge_core/bridge.py b/band-bridge/bridge_core/bridge.py index 29f1968fd..84b02407b 100644 --- a/band-bridge/bridge_core/bridge.py +++ b/band-bridge/bridge_core/bridge.py @@ -43,6 +43,7 @@ from band.platform.link import BandLink from .config import AgentConfig, BridgeConfig, ReconnectConfig +from .control import ControlSignalHandler from .forwarder import Forwarder, build_forwarder from .health import HealthServer @@ -105,6 +106,23 @@ def __init__( self._shutdown_event = shutdown_event self._connected_event = asyncio.Event() self._processed_message_ids: OrderedDict[str, None] = OrderedDict() + # In-flight forward tasks per room, so a control signal (interrupt/stop) + # can cancel whatever forwards are currently live for a room. Tracked as + # a set because same-room events serialize on the per-room lock: while + # one task forwards, others wait on that lock, and all of them must be + # cancellable — otherwise interrupt cancels a queued waiter while the + # running forward keeps going. Populated by _safe_handle_event at the + # top of processing, discarded in its finally. + self._active_room_tasks: dict[str, set[asyncio.Task[None]]] = {} + # Control-signal handling (interrupt/stop/play) lives in its own type; + # it delegates cancel/nudge/room-list back to this runner but owns the + # correlation-id dedup and routing (see control.py). + self._control = ControlSignalHandler(self) + # Background tasks spawned by control signals (play nudges) so the WS + # receive task isn't blocked awaiting them. Tracked here to be + # cancelled on close() and to keep a strong ref (create_task alone + # doesn't). + self._control_tasks: set[asyncio.Task[None]] = set() # Per-room locks: events for the same room are forwarded sequentially # so the container always sees a settled history (its own prior reply # is already posted before the next invocation reads room context). @@ -214,6 +232,8 @@ def _backoff_sleep_seconds(self, delay: float) -> float: async def close(self) -> None: """Disconnect link and close forwarder. Idempotent.""" + for task in tuple(self._control_tasks): + task.cancel() try: await self._link.disconnect() except Exception: @@ -236,6 +256,11 @@ async def _connect_and_consume(self) -> None: await self._link.connect() self._connected_event.set() + # Route control signals (interrupt/stop/play) to the control handler. + # Set as soon as connected so the hook is live for the agent_control + # channel BandLink.connect() already joined. + self._link.on_control = self._control.handle + await self._link.subscribe_agent_rooms(self._config.agent_id) existing_rooms = await self._fetch_existing_rooms() @@ -331,31 +356,36 @@ async def _rehydrate_backlog(self, room_ids: list[str]) -> None: marking, so ``/next`` would return the same message until the container marks it processed. The container's own drain loop pulls the rest. """ + # Rooms are independent; nudge concurrently. Per-room locks and dedup + # still serialize each nudge against any live event for that room. + await asyncio.gather(*(self._nudge_room(rid) for rid in room_ids)) - async def nudge(room_id: str) -> None: - try: - msg = await self._link.get_next_message(room_id) - except Exception: - logger.warning( - "Agent %s: rehydration /next failed for room %s", - self._config.agent_id, - room_id, - exc_info=True, - ) - return - if msg is None: - return - logger.info( - "Agent %s: rehydrating room %s with backlog message %s", + async def _nudge_room(self, room_id: str) -> None: + """Forward one backlog message for a room, if ``/next`` has one. + + Shared by startup rehydration (:meth:`_rehydrate_backlog`) and a live + ``play`` control signal (``ControlSignalHandler.handle``) — both cases + are the same "catch up this room via /next" operation. + """ + try: + msg = await self._link.get_next_message(room_id) + except Exception: + logger.warning( + "Agent %s: rehydration /next failed for room %s", self._config.agent_id, room_id, - msg.id, + exc_info=True, ) - await self._safe_handle_event(self._backlog_event(msg)) - - # Rooms are independent; nudge concurrently. Per-room locks and dedup - # still serialize each nudge against any live event for that room. - await asyncio.gather(*(nudge(rid) for rid in room_ids)) + return + if msg is None: + return + logger.info( + "Agent %s: rehydrating room %s with backlog message %s", + self._config.agent_id, + room_id, + msg.id, + ) + await self._safe_handle_event(self._backlog_event(msg)) def _backlog_event(self, msg: PlatformMessage) -> MessageEvent: """Wrap a ``/next`` PlatformMessage as a synthetic message_created event. @@ -382,20 +412,36 @@ def _backlog_event(self, msg: PlatformMessage) -> MessageEvent: ) async def _safe_handle_event(self, event: object) -> None: - # Semaphore caps concurrent in-flight forwards per agent. Holding it - # across ``_handle_event`` (which includes the per-room lock + the - # forwarder call) is the point — bursts wait here instead of stacking - # up inside the forwarder. - async with self._forward_semaphore: - try: + # Track this task among the in-flight forwards for its room so a control + # signal (interrupt/stop) can cancel it — whether it's the one holding + # the per-room lock and forwarding, or a same-room task still waiting on + # that lock. Room-less events (e.g. contact events) aren't tracked — + # there's nothing room-scoped to cancel for them. + room_id = getattr(event, "room_id", None) + task = asyncio.current_task() + if room_id and task is not None: + self._active_room_tasks.setdefault(room_id, set()).add(task) + try: + # Semaphore caps concurrent in-flight forwards per agent. Holding + # it across ``_handle_event`` (which includes the per-room lock + + # the forwarder call) is the point — bursts wait here instead of + # stacking up inside the forwarder. + async with self._forward_semaphore: await self._handle_event(event) - except Exception: - logger.warning( - "Agent %s: error handling event %s", - self._config.agent_id, - type(event).__name__, - exc_info=True, - ) + except Exception: + logger.warning( + "Agent %s: error handling event %s", + self._config.agent_id, + type(event).__name__, + exc_info=True, + ) + finally: + if room_id: + tasks = self._active_room_tasks.get(room_id) + if tasks is not None: + tasks.discard(task) + if not tasks: + self._active_room_tasks.pop(room_id, None) async def _handle_event(self, event: object) -> None: """Manage room subscriptions; forward forwardable events. @@ -533,6 +579,56 @@ def _remember_processed_message(self, message_id: str) -> None: if len(self._processed_message_ids) > _DEDUP_MAX_SIZE: self._processed_message_ids.popitem(last=False) + # --- Control signals (interrupt/stop/play) --- + # + # Signal routing/dedup lives in ControlSignalHandler (control.py); it calls + # back into the operations below, which touch this runner's forwarding + # state (the per-room task registry) and the /next nudge path. + + def _spawn_play_nudge(self, room_ids: list[str]) -> None: + """Fire ``play`` /next nudges as tracked background tasks. + + Runs off the WS receive task so a following interrupt/stop isn't queued + behind the nudge's /next + forward. Each spawned forward still lands in + ``_active_room_tasks`` (via ``_safe_handle_event``), so a later + interrupt/stop can cancel it. Tasks are tracked in ``_control_tasks`` + and cancelled on ``close()``. + """ + for room_id in room_ids: + task = asyncio.create_task(self._nudge_room(room_id)) + self._control_tasks.add(task) + task.add_done_callback(self._control_tasks.discard) + + def _cancel_active_forward(self, room_id: str) -> None: + """Cancel all in-flight forward tasks for a room, if any. + + Cancels both the task currently forwarding and any same-room tasks + still queued on the per-room lock, so a signal can't leave a waiter + behind to run the moment the cancelled forward releases the lock. + + Best-effort: this stops the bridge from waiting and releases the room + lock immediately. It does not guarantee the remote invocation itself + stops — ``AgentCoreForwarder``'s underlying boto3 call runs in a + thread and isn't killed by cancelling the awaiting task (see + ``forwarder.py``); that is a pre-existing, documented limitation of + that transport, not something this signal can fix. + """ + tasks = self._active_room_tasks.get(room_id) + if not tasks: + return + cancelled = 0 + for task in tuple(tasks): + if not task.done(): + task.cancel() + cancelled += 1 + if cancelled: + logger.info( + "Agent %s: cancelled %d in-flight forward(s) for room %s", + self._config.agent_id, + cancelled, + room_id, + ) + class BandBridge: """Bridge orchestrator: N AgentRunners + signal handling + health server.""" diff --git a/band-bridge/bridge_core/control.py b/band-bridge/bridge_core/control.py new file mode 100644 index 000000000..8fa6b9900 --- /dev/null +++ b/band-bridge/bridge_core/control.py @@ -0,0 +1,135 @@ +"""Agent control-signal handling (interrupt / stop / play) for the bridge. + +Split out of ``AgentRunner``: the runner owns forwarding, rehydration, and the +per-room in-flight task registry; this module owns only the control-signal +concern — correlation-id dedup, target-room resolution, and mode dispatch — +delegating the actual cancel / nudge / room-list operations back to the runner. +""" + +from __future__ import annotations + +import logging +from collections import OrderedDict +from typing import TYPE_CHECKING + +from band.client.streaming import AgentControlPayload + +if TYPE_CHECKING: + from .bridge import AgentRunner + +logger = logging.getLogger(__name__) + +_CONTROL_DEDUP_MAX_SIZE = 256 + + +class ControlSignalHandler: + """Applies ``agent.control`` signals to a runner's in-flight forwards. + + The bridge holds no Band lifecycle logic, so ``interrupt`` and ``stop`` are + handled identically: cancel whatever forward task is currently in flight + for the target room(s), if any. There is nothing to do if no forward is in + flight — the interrupt-vs-stop message-lifecycle distinction (consume vs. + leave-for-replay) is already handled downstream by the container via its + own ``/next`` claim check. ``play`` proactively nudges the room(s) via + ``/next`` so a queued message is picked up without waiting for the next + natural bridge event. + """ + + def __init__(self, runner: AgentRunner) -> None: + self._runner = runner + # Control-signal dedup. The server does not deduplicate agent.control + # pushes, so a stale duplicate could otherwise reach out and cancel + # whatever unrelated new task now occupies a room slot. Bounded LRU; + # only touched from the WebSocket receive task (same as AgentRuntime's + # _seen_control_ids on the long-running SDK path). + self._seen_control_ids: OrderedDict[str, bool] = OrderedDict() + + async def handle(self, payload: AgentControlPayload) -> None: + """Apply an ``agent.control`` signal to in-flight forwards. + + Routing mirrors ``AgentRuntime.handle_control``: a ``room_id`` targets + that room only; ``scope == "agent"`` with no ``room_id`` fans out to + all of this agent's rooms; any other combination is a no-op. Dedupes on + ``correlation_id`` (the server does not). + """ + if self._is_duplicate(payload): + return + + room_ids = await self._resolve_rooms(payload) + if room_ids is None: + return + + agent_id = self._runner.agent_id + logger.info( + "Agent %s: applying control mode=%s scope=%s to %d room(s) " + "(correlation_id=%s)", + agent_id, + payload.mode, + payload.scope, + len(room_ids), + payload.correlation_id, + ) + + match payload.mode: + case "interrupt" | "stop": + for room_id in room_ids: + self._runner._cancel_active_forward(room_id) + case "play": + # Nudge in the background so a following stop/interrupt isn't + # blocked behind this /next + forward on the WS receive task. + self._runner._spawn_play_nudge(room_ids) + case _: + logger.warning( + "Agent %s: ignoring control signal with unknown mode=%s", + agent_id, + payload.mode, + ) + + def _is_duplicate(self, payload: AgentControlPayload) -> bool: + """Return True if this signal's ``correlation_id`` was already handled. + + Records unseen ids in a bounded LRU. Signals without a + ``correlation_id`` are never treated as duplicates (the server omits + it on some pushes) but are logged. + """ + agent_id = self._runner.agent_id + cid = payload.correlation_id + if cid is None: + logger.debug( + "Agent %s: control signal mode=%s has no correlation_id; not deduped", + agent_id, + payload.mode, + ) + return False + + if cid in self._seen_control_ids: + logger.debug( + "Agent %s: duplicate control signal %s ignored", + agent_id, + cid, + ) + return True + + self._seen_control_ids[cid] = True + self._seen_control_ids.move_to_end(cid) + if len(self._seen_control_ids) > _CONTROL_DEDUP_MAX_SIZE: + self._seen_control_ids.popitem(last=False) + return False + + async def _resolve_rooms(self, payload: AgentControlPayload) -> list[str] | None: + """Resolve the signal's target room ids, or None for a no-op. + + A ``room_id`` targets that room; ``scope == "agent"`` with no + ``room_id`` fans out to all of the agent's rooms (possibly empty); any + other combination returns None so the caller skips dispatch. + """ + if payload.room_id is not None: + return [payload.room_id] + if payload.scope == "agent": + return await self._runner._fetch_existing_rooms() + logger.warning( + "Agent %s: control signal scope=%s with no room_id; no-op", + self._runner.agent_id, + payload.scope, + ) + return None diff --git a/band-bridge/bridge_core/forwarder.py b/band-bridge/bridge_core/forwarder.py index 91b8d9733..7b6df50ba 100644 --- a/band-bridge/bridge_core/forwarder.py +++ b/band-bridge/bridge_core/forwarder.py @@ -157,6 +157,10 @@ def _call() -> None: body.close() try: + # interrupt/stop are best-effort for this transport: cancelling the + # await frees the bridge but can't kill the in-flight boto3 thread — + # the remote invocation runs to completion (or its own timeout) + # regardless. await asyncio.wait_for( asyncio.to_thread(_call), timeout=self._target.timeout, diff --git a/src/band/client/streaming/__init__.py b/src/band/client/streaming/__init__.py index 31f64fedf..65bcaefd1 100644 --- a/src/band/client/streaming/__init__.py +++ b/src/band/client/streaming/__init__.py @@ -23,6 +23,7 @@ ContactAddedPayload, ContactRemovedPayload, SupersedePayload, + AgentControlPayload, ) from band.client.streaming.errors import WebSocketUpgradeError @@ -44,4 +45,5 @@ "ContactAddedPayload", "ContactRemovedPayload", "SupersedePayload", + "AgentControlPayload", ] diff --git a/src/band/client/streaming/client.py b/src/band/client/streaming/client.py index 6032319b2..7bbac17a2 100644 --- a/src/band/client/streaming/client.py +++ b/src/band/client/streaming/client.py @@ -6,7 +6,7 @@ from enum import StrEnum import logging import random -from typing import Any +from typing import Any, Literal from phoenix_channels_python_client.client import ( PHXChannelsClient, @@ -230,6 +230,27 @@ class ContactRemovedPayload(BaseModel): id: str +class AgentControlPayload(BaseModel): + """Payload for ``agent.control`` events on the agent_control channel. + + Pushed by the platform to interrupt, stop, or resume (play) an agent. + ``room_id`` is null for agent-scoped fan-out (all of the agent's rooms); + set for a single (agent, room) target. The server does not deduplicate, so + consumers should dedup on ``correlation_id``. + """ + + model_config = ConfigDict(extra="allow") + + mode: Literal["interrupt", "stop", "play"] + scope: Literal["agent", "room"] + agent_id: str + type: str | None = None + execution_id: str | None = None + room_id: str | None = None + reason: str | None = None + correlation_id: str | None = None + + class SupersedePayload(BaseModel): """Payload for terminal agent_control supersede events.""" @@ -268,6 +289,7 @@ def to_disconnect_reason(self) -> WebSocketDisconnectReason: "contact_added": ContactAddedPayload, "contact_removed": ContactRemovedPayload, "supersede": SupersedePayload, + "agent.control": AgentControlPayload, } @@ -438,13 +460,24 @@ async def join_agent_control_channel( self, agent_id: str, on_supersede: Callable[[SupersedePayload], Awaitable[None]], + on_control: Callable[[AgentControlPayload], Awaitable[None]] | None = None, ): - """Subscribe to terminal agent-control events for this agent.""" + """Subscribe to agent-control events for this agent. + + Handles terminal ``supersede`` events and, when ``on_control`` is + provided, ``agent.control`` interrupt/stop/play signals. + """ topic = f"agent_control:{agent_id}" logger.info("[WebSocket] Subscribing to topic: %s", topic) + handlers: dict[str, Callable[..., Awaitable[None]]] = { + "supersede": on_supersede + } + if on_control is not None: + handlers["agent.control"] = on_control + async def message_handler(message): - await self._handle_events(message, {"supersede": on_supersede}) + await self._handle_events(message, handlers) result = await self._require_client().subscribe_to_topic(topic, message_handler) logger.info("[WebSocket] Subscribed to topic: %s", topic) diff --git a/src/band/platform/link.py b/src/band/platform/link.py index 9590b6376..6c2067ff2 100644 --- a/src/band/platform/link.py +++ b/src/band/platform/link.py @@ -9,6 +9,7 @@ import asyncio import logging +from collections.abc import Awaitable, Callable from datetime import datetime, timezone from typing import TYPE_CHECKING @@ -46,6 +47,7 @@ ContactAddedPayload, ContactRemovedPayload, SupersedePayload, + AgentControlPayload, ) logger = logging.getLogger(__name__) @@ -102,6 +104,12 @@ def __init__( # Durable terminal disconnect reason for the current connection lifecycle. self._last_disconnect_reason: WebSocketDisconnectReason | None = None + # Preemptive control-signal hook (interrupt/stop/play). Set by the + # runtime. Invoked DIRECTLY from the WebSocket receive task — never via + # the serialized _event_queue — so a control signal can act on a cycle + # already in flight instead of queuing behind it. + self.on_control: Callable[[AgentControlPayload], Awaitable[None]] | None = None + # Debounce flag for activity-report failures: keep-alive runs at a few # seconds per room, so a down endpoint would otherwise flood the log on # every refresh. Log the first failure and the recovery, suppress repeats. @@ -151,6 +159,7 @@ async def connect(self) -> None: await self._ws.join_agent_control_channel( self.agent_id, on_supersede=self._on_supersede, + on_control=self._on_control, ) except Exception: await self._ws.__aexit__(None, None, None) @@ -332,6 +341,22 @@ async def _on_supersede(self, payload: "SupersedePayload") -> None: ) self._queue_event(WebSocketDisconnectedEvent(payload=reason)) + async def _on_control(self, payload: "AgentControlPayload") -> None: + """Handle an ``agent.control`` push (interrupt/stop/play). + + Invoked directly from the WebSocket receive task. Forwards to the + registered ``on_control`` hook WITHOUT touching the serialized event + queue, so the signal can preempt a cycle already in flight. If no hook + is registered, the push is a safe no-op. + """ + if self.on_control is None: + logger.debug( + "agent.control received (mode=%s) but no on_control hook registered", + payload.mode, + ) + return + await self.on_control(payload) + async def _on_disconnected(self, error: Exception | None) -> None: """Handle PHX client disconnection.""" if self.last_disconnect_reason: diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index 7e97c5d70..e308c9af2 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -15,6 +15,7 @@ from __future__ import annotations import asyncio +import contextlib import logging from datetime import datetime, timezone from enum import Enum, StrEnum @@ -158,6 +159,30 @@ async def request_resync(self) -> None: """ ... + def interrupt(self, *, kind: str = "interrupt") -> bool: + """Abort the in-flight reasoning cycle for this room. + + Called preemptively from the WebSocket receive task on an ``interrupt`` + or ``stop`` control signal. ``AgentRuntime`` ``hasattr``-guards this, so + custom ``Execution`` implementations that omit it degrade to a no-op. + """ + ... + + def stop_room(self) -> None: + """Durably stop this room until a play signal. + + Aborts the in-flight cycle and goes quiet. Trigger suppression is + platform-authoritative. ``AgentRuntime`` ``hasattr``-guards this. + """ + ... + + async def resume_room(self) -> None: + """Resume a stopped room (play): catch up rehydration-style via /next. + + ``AgentRuntime`` ``hasattr``-guards this. + """ + ... + # Type for execution callback ExecutionHandler = Callable[["ExecutionContext", PlatformEvent], Awaitable[None]] @@ -275,6 +300,36 @@ def __init__( self._pending_system_messages: list[str] = [] self._reconnect_sync_requested = False + # Per-cycle interrupt. The reasoning cycle runs as a child + # task so a control signal can abort just this turn without killing the + # room loop. ``_interrupt_kind`` is set by interrupt() on the receive + # task BEFORE cancelling the child, then read-and-cleared in the loop + # coroutine's cancel handler so it can't leak across cycles. + self._active_cycle_task: asyncio.Task[None] | None = None + self._interrupt_kind: str | None = None # "interrupt" | "stop" | None + + # Signal that landed in the claim->cycle window, where a message is + # claimed (mark_processing) and hydrating but the cancellable cycle task + # doesn't exist yet, so interrupt() has nothing to cancel. ``_cycle_armed`` + # marks that window open; interrupt() records ``_pending_interrupt`` while + # it is, and _run_cycle honors it before invoking the handler. Both are + # cleared as the cycle starts and in the per-message ``finally``, so a + # signal can never leak onto a later cycle. + self._cycle_armed: bool = False + self._pending_interrupt: str | None = None + + # Durable stop (play to resume). Trigger suppression is + # platform-authoritative (dispatch gated server-side, persists across + # reconnect); this flag is a PURE LOCAL EFFICIENCY CACHE — it pauses + # idle /next polling and short-circuits WS triggers while stopped to + # avoid /next->204 and mark->204/reply->403 churn. Not persisted. + self._stopped: bool = False + + # Optional seam for clearing user-visible activity state ("reasoning…") + # when a cycle is interrupted/stopped. Filled by the activity-signal + # work; a no-op (None) here. + self._on_activity_clear: Callable[[], Awaitable[None]] | None = None + @property def thread_id(self) -> str: """LangGraph thread_id = room_id.""" @@ -453,6 +508,14 @@ async def stop(self, timeout: float | None = None) -> bool: self.room_id, ) + # Cancel any in-flight cycle child task BEFORE cancelling the loop so it + # is not orphaned (the loop's await on it is not auto-cancelled when the + # loop task is cancelled). Capture the reference first because the loop's + # finally clears _active_cycle_task as it unwinds. + cycle_task = self._active_cycle_task + if cycle_task is not None and not cycle_task.done(): + cycle_task.cancel() + # Signal stop and cancel the task self._is_running = False self._process_loop_task.cancel() @@ -462,6 +525,12 @@ async def stop(self, timeout: float | None = None) -> bool: pass self._process_loop_task = None + # Drain the (now cancelled) cycle task so it does not leak as pending. + if cycle_task is not None: + with contextlib.suppress(asyncio.CancelledError): + await cycle_task + self._active_cycle_task = None + # Defensively clear any lingering working-state keep-alive so a removed # room can't leak its refresh task. Idempotent: a no-op if not active # (the per-cycle finally normally clears it already). @@ -534,6 +603,84 @@ async def request_resync(self) -> None: self.queue.put_nowait(_ResyncRequest()) # type: ignore[arg-type] # Sentinel is intentionally not a PlatformEvent. logger.debug("ExecutionContext %s: Resync sentinel enqueued", self.room_id) + def interrupt(self, *, kind: str = "interrupt") -> bool: + """Abort the in-flight reasoning cycle, if any. + + Called from the WebSocket receive task. The receive-side surface is + deliberately minimal — set a flag and cancel the child cycle task. All + message-status / dedupe bookkeeping happens in the loop coroutine's + cancel handler (``_run_cycle``), never here, so the two coroutines do + not race on shared state. + + Args: + kind: ``"interrupt"`` (consume the message) or ``"stop"`` (leave it + actionable for replay on play). Distinguishes the two unwind + paths in ``_run_cycle``. + + Returns: + True if the signal took effect — either a running cycle was + cancelled, or a claimed-but-not-yet-started cycle was armed to abort + (the claim->cycle window). Between cycles this is a clean no-op and + does NOT set ``_interrupt_kind``/``_pending_interrupt`` (which would + otherwise mis-flag the next cycle). + """ + task = self._active_cycle_task + if task is not None and not task.done(): + self._interrupt_kind = kind + task.cancel() + logger.info( + "ExecutionContext %s: %s requested, cancelling in-flight cycle", + self.room_id, + kind, + ) + return True + if self._cycle_armed: + # A message is claimed and its cycle is imminent but the cancellable + # task doesn't exist yet; record the request so _run_cycle aborts + # before invoking the handler instead of losing the signal. + self._pending_interrupt = kind + logger.info( + "ExecutionContext %s: %s requested during claim->cycle window", + self.room_id, + kind, + ) + return True + return False + + def stop_room(self) -> None: + """Durable stop for this room: abort the in-flight cycle and + go quiet until a play signal. + + The platform is authoritative on trigger suppression; ``_stopped`` is a + local efficiency cache only (pause idle /next polling, short-circuit WS + triggers). Called from the receive task — surface stays flag + cancel. + Leaves any in-flight message in 'processing' so the platform replays it + via /next on play. + """ + self._stopped = True + self.interrupt(kind="stop") + + async def resume_room(self) -> None: + """Resume a stopped room (play): clear the local stop flag and catch up + rehydration-style via /next, so callouts made while stopped are seen. + + Clears ``_stopped`` BEFORE enqueuing the resync sentinel so the loop + does not skip the catch-up it just requested. + """ + self._stopped = False + await self.request_resync() + + async def _clear_activity(self) -> None: + """Invoke the optional activity-clear seam after an aborted cycle.""" + if self._on_activity_clear is None: + return + try: + await self._on_activity_clear() + except Exception: + logger.exception( + "ExecutionContext %s: activity-clear hook failed", self.room_id + ) + # --- Participant management --- def add_participant(self, participant: dict) -> bool: @@ -897,6 +1044,13 @@ async def _process_loop(self) -> None: async with asyncio.timeout(self.config.idle_resync_seconds): event = await self.queue.get() except asyncio.TimeoutError: + if self._stopped: + # Efficiency: a stopped room would only get /next->204. + logger.debug( + "ExecutionContext %s: stopped, skipping idle /next poll", + self.room_id, + ) + continue logger.debug( "ExecutionContext %s: Idle for %ss, re-polling /next", self.room_id, @@ -906,6 +1060,15 @@ async def _process_loop(self) -> None: continue if isinstance(event, _ResyncRequest): + if self._stopped: + # resume_room() clears _stopped before enqueuing its + # sentinel, so a sentinel seen while stopped is a stale + # reconnect resync — platform gate keeps us quiet anyway. + logger.debug( + "ExecutionContext %s: stopped, ignoring resync sentinel", + self.room_id, + ) + continue logger.debug( "ExecutionContext %s: Resync requested (post-reconnect)", self.room_id, @@ -1009,6 +1172,20 @@ async def _synchronize_with_next(self) -> bool: if result == _BacklogProcessResult.RETRY_LATER: return False + if self._stopped: + # A stop control signal landed mid-cycle: the message was + # deliberately left in 'processing' for replay on play, but + # /next excludes only 'processed' messages, so the next + # iteration would just re-fetch and fully re-run the very + # cycle stop just aborted. Pause here instead; play's + # resync will pick the message back up. + logger.debug( + "ExecutionContext %s: stopped mid-backlog-sync, " + "pausing /next polling", + self.room_id, + ) + break + if self._retry_tracker.is_permanently_failed(next_msg.id): logger.warning( "ExecutionContext %s: Message %s permanently failed", @@ -1033,7 +1210,20 @@ async def _recover_stale_processing_messages(self) -> bool: state on the server. The /next endpoint skips these messages, so the agent would never pick them up again. This method finds such messages and re-processes them by calling mark_processing (creates a new attempt). + + Skipped while stopped: the stop path deliberately leaves the interrupted + message in 'processing', and a reconnect must not resurrect it through + the recovery sweep. The platform replays it via /next on play instead. + This keeps stop-survives-reconnect correct in the SDK without relying on + the platform gating the mark endpoint for this path. """ + if self._stopped: + logger.debug( + "ExecutionContext %s: stopped, skipping stale-processing recovery", + self.room_id, + ) + return True + stale_messages = await self.link.get_stale_processing_messages(self.room_id) if not stale_messages: return True @@ -1118,6 +1308,17 @@ async def _resync_pending_messages(self) -> bool: caught_up, ) + if self._stopped: + # See the matching guard in _synchronize_with_next: a stop + # mid-cycle leaves the message 'processing' for replay, and + # /next would just hand it straight back next iteration. + logger.debug( + "ExecutionContext %s: stopped mid-resync, pausing " + "/next polling", + self.room_id, + ) + break + if self._retry_tracker.is_permanently_failed(next_msg.id): break @@ -1218,6 +1419,11 @@ async def _process_claimed_backlog_message( ) return _BacklogProcessResult.ADVANCED + # Open the claim->cycle window: until _run_cycle creates the + # cancellable task, an interrupt/stop has no task to cancel, so + # interrupt() records it as pending instead. + self._cycle_armed = True + # Mark as processing on server BEFORE we start. If this fails, do not # invoke the adapter; otherwise the platform will keep returning the # same message and the agent may replay side effects. @@ -1279,9 +1485,14 @@ async def _process_claimed_backlog_message( ), ) - # Call execution handler (backlog messages are always reasoning - # cycles, so always report the working signal). - await self._execute_message_cycle(event) + # Call execution handler as a cancellable cycle (backlog messages + # are always reasoning cycles, so the working signal is always + # reported via _execute_message_cycle inside _invoke_handler). A + # control signal can abort just this turn; when it does, status is + # handled inside _run_cycle and we advance without sending + # anything. + if not await self._run_cycle(event, msg_id): + return _BacklogProcessResult.ADVANCED # SUCCESS: Mark as processed on server durable_processed = await self.link.mark_processed(self.room_id, msg_id) @@ -1314,6 +1525,10 @@ async def _process_claimed_backlog_message( return _BacklogProcessResult.ADVANCED finally: + # Close the claim->cycle window on every exit so a pending signal + # can't leak onto the next backlog message. + self._cycle_armed = False + self._pending_interrupt = None self._set_state(ExecutionState.IDLE) def _drain_duplicate_from_queue(self, msg_id: str) -> None: @@ -1342,6 +1557,19 @@ def _drain_duplicate_from_queue(self, msg_id: str) -> None: for item in items: self.queue.put_nowait(item) + async def _invoke_handler(self, event: PlatformEvent) -> None: + """Coroutine wrapper around the execution handler so it can run as a + cancellable ``asyncio.Task`` (the handler is typed ``Awaitable``). + + Message-driven cycles are bracketed by the working-state signal via + ``_execute_message_cycle``; participant add/remove events are + housekeeping and skip it. + """ + if isinstance(event, MessageEvent): + await self._execute_message_cycle(event) + else: + await self._on_execute(self, event) + async def _report_working_state(self, working: bool) -> bool: """Report the room's boolean working state (wired into the reporter).""" return await self.link.report_activity( @@ -1365,6 +1593,105 @@ async def _execute_message_cycle(self, event: PlatformEvent) -> None: finally: await self._working_reporter.stop() + async def _run_cycle(self, event: PlatformEvent, msg_id: str | None) -> bool: + """Run the execution handler as a cancellable child task. + + Wrapping the cycle in its own task lets a control signal abort just this + turn (via ``interrupt()``) without cancelling the room's process loop. + + Returns: + True if the cycle ran to completion (caller proceeds to mark the + message processed as usual). False if a control signal aborted the + cycle — message status has already been handled here and the caller + must send nothing further. + + Raises: + asyncio.CancelledError: when the cancel was a genuine shutdown of + the loop task (``_interrupt_kind`` unset), so the loop's own handler + can exit. ``CancelledError`` is a ``BaseException`` subclass, so it + bypasses the callers' ``except Exception`` and reaches the loop's + ``except asyncio.CancelledError``; we only swallow it for an + interrupt/stop, never for shutdown. + """ + # Honor a signal that landed in the claim->cycle window (interrupt()/ + # stop_room() with no cycle task to cancel yet). Reading/clearing the + # flags and creating the task below all run without an intervening + # await, so interrupt() on the receive task can't interleave here. + pending = self._pending_interrupt + self._pending_interrupt = None + self._cycle_armed = False + if pending is not None: + return await self._abort_cycle(pending, msg_id) + + self._active_cycle_task = asyncio.create_task(self._invoke_handler(event)) + try: + await self._active_cycle_task + return True + except asyncio.CancelledError: + # Read-and-clear is atomic here (no await between the two lines). + # If two control signals raced before this ran, last-writer-wins on + # _interrupt_kind — benign, since re-cancelling a cancelling task is + # a no-op and both signals wanted the cycle dead. + kind = self._interrupt_kind + self._interrupt_kind = None + if kind is None: + # Shutdown cancel of the loop task propagating through the child + # await — let it propagate so the loop exits. + raise + return await self._abort_cycle(kind, msg_id) + finally: + self._active_cycle_task = None + + async def _abort_cycle(self, kind: str, msg_id: str | None) -> bool: + """Unwind an aborted cycle (interrupt/stop): drop work, send nothing. + + Shared by the in-flight cancel path (``_run_cycle``'s ``CancelledError`` + handler) and the claim->cycle window where interrupt()/stop_room() + landed before the cycle task existed. Returns False so the caller sends + nothing further. + """ + # The handler never ran to completion, so uncharge the attempt + # `record_attempt` already billed before this cycle started — otherwise + # a message that's merely stopped/interrupted a couple of times gets + # poisoned into permanently_failed before it's ever actually attempted. + if msg_id: + self._retry_tracker.discard_attempt(msg_id) + await self._clear_activity() + if kind == "interrupt" and msg_id: + # Consume the message so the idle /next resync does not re-return it + # (excludes-only-processed) and re-fire the cycle the user just + # interrupted. Mirror the success-path bookkeeping. + if await self.link.mark_processed(self.room_id, msg_id): + self.claims.remember_completed(self.room_id, msg_id) + else: + # Durable ack failed. Mark the message locally consumed and + # queue the ack for background retry, exactly as the success + # path does — otherwise the interrupted message stays locally + # replayable and the idle /next resync re-fires the cycle the + # user just interrupted. + self.claims.remember_ack_pending(self.room_id, msg_id) + logger.warning( + "ExecutionContext %s: durable mark_processed failed for " + "interrupted message %s; retrying ack in background", + self.room_id, + msg_id, + ) + # For "stop" we deliberately leave the message in 'processing' so the + # platform replays it via /next on play; do not mark or remember. + # + # CROSS-SYSTEM INVARIANT: this replay depends on the platform's /next + # (Chat.get_next_actionable_message) excluding ONLY 'processed' — a + # 'processing' message must still be returned. If the platform ever also + # excludes 'processing', stopped messages are silently dropped on play. + # Covered by the stop->play replay test. + logger.info( + "ExecutionContext %s: cycle %s (message %s) — nothing sent", + self.room_id, + "interrupted" if kind == "interrupt" else "stopped", + msg_id, + ) + return False + async def _process_event(self, event: PlatformEvent) -> bool: """ Process single event through execution callback. @@ -1382,14 +1709,37 @@ async def _process_event(self, event: PlatformEvent) -> bool: try: if self._reconnect_sync_requested: self._reconnect_sync_requested = False - while not await self._synchronize_with_next(): - self._set_state(ExecutionState.IDLE) - await asyncio.sleep(self.config.idle_resync_seconds) + if self._stopped: + # Efficiency: a stopped room's /next is guaranteed 204 + # (platform-authoritative gate) — skip locally instead + # of making a call known to come back empty, same as + # the idle-timeout and resync-sentinel paths. + logger.debug( + "ExecutionContext %s: stopped, skipping reconnect /next sync", + self.room_id, + ) + else: + while not await self._synchronize_with_next(): + self._set_state(ExecutionState.IDLE) + await asyncio.sleep(self.config.idle_resync_seconds) logger.debug("Event %s processed successfully", event.type) finally: self._set_state(ExecutionState.IDLE) return True + # While stopped, leave message triggers actionable for replay on play. + # Suppression is platform-authoritative; skipping here is a local + # efficiency short-circuit that avoids claiming/marking (mark->204) and + # never reaches the adapter (reply->403). The message is left untouched + # so /next replays it on play. + if self._stopped and isinstance(event, MessageEvent): + logger.debug( + "ExecutionContext %s: stopped, skipping message %s (left for replay)", + self.room_id, + event.payload.id if event.payload else None, + ) + return True + payload = event.payload if isinstance(event, MessageEvent) else None msg_id = payload.id if payload else None @@ -1488,6 +1838,11 @@ async def _process_event_body( ) return True + # Open the claim->cycle window: from here until _run_cycle + # creates the cancellable task, an interrupt/stop has no task to + # cancel, so interrupt() records it as pending instead. + self._cycle_armed = True + # For messages: mark as processing on server if not await self.link.mark_processing(self.room_id, msg_id): logger.warning( @@ -1509,12 +1864,13 @@ async def _process_event_body( self.remove_participant(event.payload.id) await self._notify_participant_removed(event) - # Call execution handler. Only message-driven cycles report the - # working signal; participant add/remove events are housekeeping. - if isinstance(event, MessageEvent): - await self._execute_message_cycle(event) - else: - await self._on_execute(self, event) + # Call execution handler as a cancellable cycle. A control signal + # can abort just this turn; when it does, status is handled inside + # _run_cycle and we send nothing further. Only message-driven + # cycles report the working signal (see _invoke_handler); + # participant add/remove events are housekeeping and skip it. + if not await self._run_cycle(event, msg_id): + return True # For messages: mark as processed on server if isinstance(event, MessageEvent) and msg_id: @@ -1549,4 +1905,9 @@ async def _process_event_body( return True finally: + # Close the claim->cycle window on every exit (e.g. hydration raised + # before _run_cycle consumed the flags) so a pending signal can't + # leak onto the next message. + self._cycle_armed = False + self._pending_interrupt = None self._set_state(ExecutionState.IDLE) diff --git a/src/band/runtime/platform_runtime.py b/src/band/runtime/platform_runtime.py index 03f12d4d3..63c55bc12 100644 --- a/src/band/runtime/platform_runtime.py +++ b/src/band/runtime/platform_runtime.py @@ -177,6 +177,11 @@ async def start( on_participant_removed=self._on_participant_removed, ) + # Route preemptive control signals (interrupt/stop/play) to the + # runtime BEFORE starting (which connects the WebSocket), so the + # hook is live as soon as the agent_control channel is joined. + self._link.on_control = self._runtime.handle_control + await self._runtime.start() # Set up contact event handling after WebSocket is connected diff --git a/src/band/runtime/retry_tracker.py b/src/band/runtime/retry_tracker.py index d949d8031..fc257a20f 100644 --- a/src/band/runtime/retry_tracker.py +++ b/src/band/runtime/retry_tracker.py @@ -55,6 +55,21 @@ def mark_success(self, msg_id: str) -> None: """Clear tracking for successfully processed message.""" self._attempts.pop(msg_id, None) + def discard_attempt(self, msg_id: str) -> None: + """Uncharge the attempt aborted by our own control signal (interrupt/stop). + + A cycle cancelled by interrupt/stop never actually ran the handler, so + it must not count against the message's retry budget. Decrement by one + rather than clearing the counter: earlier genuine failures on the same + message must stay charged, otherwise interrupting one retry silently + resets the whole budget. + """ + remaining = self._attempts.get(msg_id, 0) - 1 + if remaining > 0: + self._attempts[msg_id] = remaining + else: + self._attempts.pop(msg_id, None) + def mark_permanently_failed(self, msg_id: str) -> None: """Explicitly mark message as permanently failed.""" self._failed.add(msg_id) diff --git a/src/band/runtime/runtime.py b/src/band/runtime/runtime.py index ab5262862..1dd696510 100644 --- a/src/band/runtime/runtime.py +++ b/src/band/runtime/runtime.py @@ -8,6 +8,7 @@ from __future__ import annotations import logging +from collections import OrderedDict from typing import TYPE_CHECKING, Awaitable, Callable, Protocol from band.platform.event import PlatformEvent @@ -22,6 +23,7 @@ ) if TYPE_CHECKING: + from band.client.streaming import AgentControlPayload from band.platform.link import BandLink logger = logging.getLogger(__name__) @@ -132,6 +134,12 @@ def __init__( # Per-room executions self.executions: dict[str, Execution] = {} + # Control-signal dedup. The server does not deduplicate + # agent.control pushes, so we drop repeats by correlation_id. Bounded + # LRU; only touched from the WebSocket receive task. + self._seen_control_ids: OrderedDict[str, bool] = OrderedDict() + self._max_seen_control_ids: int = 256 + # Shared by default contexts so a room/message pair executes at most # once per runtime, including across context recreation. self._claim_registry = MessageClaimRegistry() @@ -248,6 +256,106 @@ async def _on_reconnected(self) -> None: except Exception as e: logger.warning("Failed to request resync for room %s: %s", room_id, e) + # --- Control signals --- + + async def handle_control(self, payload: "AgentControlPayload") -> None: + """Apply an ``agent.control`` signal (interrupt/stop/play) to executions. + + Invoked directly from the WebSocket receive task (via + ``BandLink.on_control``) so it can preempt a cycle already in flight. + + Routing: + - ``scope == "agent"`` with no ``room_id`` → fan out to all executions. + - any signal with a ``room_id`` → that room only. + - unknown room → no-op. + + Dedupes on ``correlation_id`` (the server does not). Note: exceptions + raised here are swallowed-and-logged by the channel handler, so this + method must not rely on propagating errors to signal anything. + """ + cid = payload.correlation_id + if cid is not None: + if cid in self._seen_control_ids: + logger.debug("Duplicate control signal %s ignored", cid) + return + self._seen_control_ids[cid] = True + self._seen_control_ids.move_to_end(cid) + if len(self._seen_control_ids) > self._max_seen_control_ids: + self._seen_control_ids.popitem(last=False) + else: + logger.debug( + "Control signal mode=%s has no correlation_id; not deduped", + payload.mode, + ) + + # Resolve target executions. + if payload.room_id is not None: + execution = self.executions.get(payload.room_id) + if execution is None: + logger.info( + "Control signal (mode=%s) for unknown room %s; no-op", + payload.mode, + payload.room_id, + ) + return + targets = [execution] + elif payload.scope == "agent": + targets = list(self.executions.values()) + else: + logger.warning( + "Control signal scope=%s with no room_id; no-op", payload.scope + ) + return + + logger.info( + "Applying control mode=%s scope=%s to %d room(s) " + "(correlation_id=%s execution_id=%s)", + payload.mode, + payload.scope, + len(targets), + cid, + payload.execution_id, + ) + + for execution in targets: + await self._apply_control(execution, payload.mode) + + async def _apply_control(self, execution: Execution, mode: str) -> None: + """Dispatch one control mode to one execution, degrading gracefully. + + Custom ``Execution`` implementations that omit the control methods are + skipped with a log (mirrors how ``request_resync`` degrades). + """ + if mode == "interrupt": + fn = getattr(execution, "interrupt", None) + if fn is None: + logger.debug( + "Execution for room %s has no interrupt(); skipping", + getattr(execution, "room_id", "?"), + ) + return + fn() + elif mode == "stop": + fn = getattr(execution, "stop_room", None) + if fn is None: + logger.debug( + "Execution for room %s has no stop_room(); skipping", + getattr(execution, "room_id", "?"), + ) + return + fn() + elif mode == "play": + fn = getattr(execution, "resume_room", None) + if fn is None: + logger.debug( + "Execution for room %s has no resume_room(); skipping", + getattr(execution, "room_id", "?"), + ) + return + await fn() + else: + logger.warning("Unknown control mode %r; ignoring", mode) + # --- Execution management --- async def _create_execution(self, room_id: str) -> Execution: diff --git a/tests/bridge/test_bridge_control.py b/tests/bridge/test_bridge_control.py new file mode 100644 index 000000000..8818a9f65 --- /dev/null +++ b/tests/bridge/test_bridge_control.py @@ -0,0 +1,424 @@ +"""Tests for AgentRunner control-signal handling (interrupt/stop/play). + +The bridge holds no Band lifecycle logic, so interrupt and stop are handled +identically here: cancel whatever forward task is currently in flight for the +target room(s), if any. The interrupt-vs-stop distinction (consume vs. +leave-for-replay) is handled downstream by the container via ``/next``, not +by the bridge. ``play`` proactively nudges the room(s) via ``/next`` so a +queued message is picked up without waiting for the next natural bridge +event. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from typing import Any +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from band.client.streaming import AgentControlPayload + +from .conftest import FakeForwarder +from .test_bridge import _build_runner, _make_message_event, _make_platform_message + +pytestmark = pytest.mark.asyncio + + +def _control( + mode: str, *, scope: str, room_id: str | None, **kw: Any +) -> AgentControlPayload: + return AgentControlPayload( + mode=mode, scope=scope, agent_id="agent-1", room_id=room_id, **kw + ) + + +class _HangingForwarder: + """Hangs on the first forward() until cancelled; forwards normally after. + + ``started`` lets a test wait until the forward call has actually begun + before issuing a control signal, avoiding a race between task creation + and the signal arriving before there's anything in flight to cancel. + """ + + def __init__(self) -> None: + self.started = asyncio.Event() + self.hang = True + self.forwarded: list[dict[str, Any]] = [] + + async def forward(self, payload: dict[str, Any]) -> None: + if self.hang: + self.started.set() + await asyncio.sleep(120) + self.forwarded.append(payload) + + async def close(self) -> None: + return + + +class _MultiRoomHangingForwarder: + """Hangs forever, independently, for each of a fixed set of rooms.""" + + def __init__(self, rooms: list[str]) -> None: + self.started: dict[str, asyncio.Event] = {r: asyncio.Event() for r in rooms} + self.forwarded: list[dict[str, Any]] = [] + + async def forward(self, payload: dict[str, Any]) -> None: + room_id = payload.get("room_id") + self.started[room_id].set() + await asyncio.sleep(120) + self.forwarded.append(payload) + + async def close(self) -> None: + return + + +async def _await_cancelled(task: asyncio.Task[None]) -> None: + with contextlib.suppress(asyncio.CancelledError): + await task + assert task.cancelled() + + +async def _drain_control_tasks(runner: Any) -> None: + """Await the background tasks a ``play`` signal spawns (it's fire-and-forget + off the receive path, so tests must let the nudges finish before asserting). + """ + tasks = tuple(runner._control_tasks) + if tasks: + await asyncio.gather(*tasks) + + +# --------------------------------------------------------------------------- +# Wiring +# --------------------------------------------------------------------------- + + +class TestControlWiring: + async def test_connect_wires_on_control_to_handler(self) -> None: + runner, _, link = _build_runner() + link.__anext__ = AsyncMock(side_effect=StopAsyncIteration()) + + await runner._connect_and_consume() + + assert link.on_control == runner._control.handle + + +# --------------------------------------------------------------------------- +# Interrupt / stop — cancel an in-flight forward +# --------------------------------------------------------------------------- + + +class TestControlCancelsInFlightForward: + async def test_interrupt_cancels_in_flight_forward(self) -> None: + fwd = _HangingForwarder() + runner, _, _ = _build_runner(forwarder=fwd) + task = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m1", room_id="r1") + ) + ) + await fwd.started.wait() + + await runner._control.handle(_control("interrupt", scope="room", room_id="r1")) + + await _await_cancelled(task) + assert fwd.forwarded == [] + + async def test_stop_cancels_in_flight_forward_identically(self) -> None: + fwd = _HangingForwarder() + runner, _, _ = _build_runner(forwarder=fwd) + task = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m1", room_id="r1") + ) + ) + await fwd.started.wait() + + await runner._control.handle(_control("stop", scope="room", room_id="r1")) + + await _await_cancelled(task) + assert fwd.forwarded == [] + + async def test_cancelled_message_is_not_remembered_as_forwarded(self) -> None: + """A cancelled forward must stay retryable — never marked processed.""" + fwd = _HangingForwarder() + runner, _, _ = _build_runner(forwarder=fwd) + task = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m-cancel", room_id="r1") + ) + ) + await fwd.started.wait() + + await runner._control.handle(_control("interrupt", scope="room", room_id="r1")) + await _await_cancelled(task) + + assert "m-cancel" not in runner._processed_message_ids + + async def test_room_lock_releases_promptly_after_interrupt(self) -> None: + """A second event for the same room must not wait on the cancelled task.""" + fwd = _HangingForwarder() + runner, _, _ = _build_runner(forwarder=fwd) + task1 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m1", room_id="r1") + ) + ) + await fwd.started.wait() + + await runner._control.handle(_control("interrupt", scope="room", room_id="r1")) + await _await_cancelled(task1) + + fwd.hang = False + await asyncio.wait_for( + runner._safe_handle_event( + _make_message_event(message_id="m2", room_id="r1") + ), + timeout=1.0, + ) + + ids = {(p.get("payload") or {}).get("id") for p in fwd.forwarded} + assert ids == {"m2"} + + async def test_interrupt_with_no_active_forward_is_noop(self) -> None: + runner, fwd, _ = _build_runner() + + await runner._control.handle(_control("interrupt", scope="room", room_id="r1")) + + assert isinstance(fwd, FakeForwarder) + assert fwd.forwarded == [] + assert "r1" not in runner._active_room_tasks + + +# --------------------------------------------------------------------------- +# Routing +# --------------------------------------------------------------------------- + + +class TestControlRouting: + async def test_room_scope_only_cancels_target_room(self) -> None: + fwd = _MultiRoomHangingForwarder(rooms=["r1", "r2"]) + runner, _, _ = _build_runner(forwarder=fwd) + t1 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m1", room_id="r1") + ) + ) + t2 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m2", room_id="r2") + ) + ) + await fwd.started["r1"].wait() + await fwd.started["r2"].wait() + + await runner._control.handle(_control("interrupt", scope="room", room_id="r1")) + + await _await_cancelled(t1) + assert not t2.done() + + t2.cancel() + with contextlib.suppress(asyncio.CancelledError): + await t2 + + async def test_agent_scope_null_room_cancels_all_active_forwards(self) -> None: + fwd = _MultiRoomHangingForwarder(rooms=["r1", "r2"]) + runner, _, link = _build_runner(forwarder=fwd) + link.rest.agent_api_chats.list_agent_chats = AsyncMock( + return_value=MagicMock(data=[MagicMock(id="r1"), MagicMock(id="r2")]) + ) + t1 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m1", room_id="r1") + ) + ) + t2 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m2", room_id="r2") + ) + ) + await fwd.started["r1"].wait() + await fwd.started["r2"].wait() + + await runner._control.handle(_control("interrupt", scope="agent", room_id=None)) + + await _await_cancelled(t1) + await _await_cancelled(t2) + + async def test_bad_scope_with_no_room_id_is_noop(self) -> None: + runner, fwd, _ = _build_runner() + + await runner._control.handle( + AgentControlPayload(mode="interrupt", scope="room", agent_id="agent-1") + ) + + assert isinstance(fwd, FakeForwarder) + assert fwd.forwarded == [] + + +# --------------------------------------------------------------------------- +# Dedup +# --------------------------------------------------------------------------- + + +class TestControlDedup: + async def test_duplicate_correlation_id_does_not_clobber_a_later_task(self) -> None: + """The real risk dedup guards against: a stale duplicate delivery must + not reach out and cancel whatever unrelated new task now occupies the + same room slot.""" + fwd = _HangingForwarder() + runner, _, _ = _build_runner(forwarder=fwd) + sig = _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-1") + + task1 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m1", room_id="r1") + ) + ) + await fwd.started.wait() + await runner._control.handle(sig) + await _await_cancelled(task1) + + # A new, unrelated forward now occupies the room slot. + fwd.started = asyncio.Event() + fwd.hang = True + task2 = asyncio.create_task( + runner._safe_handle_event( + _make_message_event(message_id="m2", room_id="r1") + ) + ) + await fwd.started.wait() + + # Duplicate delivery of the SAME signal must not cancel task2. + await runner._control.handle(sig) + await asyncio.sleep(0.01) + assert not task2.done() + + task2.cancel() + with contextlib.suppress(asyncio.CancelledError): + await task2 + + async def test_duplicate_correlation_id_dropped(self) -> None: + runner, _, _ = _build_runner() + runner._cancel_active_forward = MagicMock() + sig = _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-1") + + await runner._control.handle(sig) + await runner._control.handle(sig) + + runner._cancel_active_forward.assert_called_once() + + async def test_distinct_correlation_ids_both_applied(self) -> None: + runner, _, _ = _build_runner() + runner._cancel_active_forward = MagicMock() + + await runner._control.handle( + _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-1") + ) + await runner._control.handle( + _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-2") + ) + + assert runner._cancel_active_forward.call_count == 2 + + async def test_missing_correlation_id_not_deduped(self) -> None: + runner, _, _ = _build_runner() + runner._cancel_active_forward = MagicMock() + sig = _control("interrupt", scope="room", room_id="r1") + + await runner._control.handle(sig) + await runner._control.handle(sig) + + assert runner._cancel_active_forward.call_count == 2 + + +# --------------------------------------------------------------------------- +# Play — proactive /next nudge +# --------------------------------------------------------------------------- + + +class TestControlPlayNudge: + async def test_room_scope_play_nudges_next_message(self) -> None: + runner, fwd, link = _build_runner() + link.get_next_message = AsyncMock( + return_value=_make_platform_message("m-a", "r1") + ) + + await runner._control.handle(_control("play", scope="room", room_id="r1")) + await _drain_control_tasks(runner) + + assert isinstance(fwd, FakeForwarder) + ids = {(p.get("payload") or {}).get("id") for p in fwd.forwarded} + assert ids == {"m-a"} + + async def test_agent_scope_play_nudges_all_rooms(self) -> None: + runner, fwd, link = _build_runner() + link.rest.agent_api_chats.list_agent_chats = AsyncMock( + return_value=MagicMock(data=[MagicMock(id="r1"), MagicMock(id="r2")]) + ) + link.get_next_message.side_effect = [ + _make_platform_message("m-a", "r1"), + _make_platform_message("m-b", "r2"), + ] + + await runner._control.handle(_control("play", scope="agent", room_id=None)) + await _drain_control_tasks(runner) + + assert isinstance(fwd, FakeForwarder) + ids = {(p.get("payload") or {}).get("id") for p in fwd.forwarded} + assert ids == {"m-a", "m-b"} + + async def test_play_with_no_backlog_forwards_nothing(self) -> None: + runner, fwd, link = _build_runner() + link.get_next_message = AsyncMock(return_value=None) + + await runner._control.handle(_control("play", scope="room", room_id="r1")) + await _drain_control_tasks(runner) + + assert isinstance(fwd, FakeForwarder) + assert fwd.forwarded == [] + + async def test_play_nudge_does_not_block_receive_path(self) -> None: + """play must return on the receive task without awaiting the /next + + forward, so a following stop/interrupt isn't queued behind it.""" + started = asyncio.Event() + + async def slow_next(room_id: str) -> None: + started.set() + await asyncio.sleep(120) # would block the receive task if awaited + + runner, _, link = _build_runner() + link.get_next_message = AsyncMock(side_effect=slow_next) + + # handle() returns promptly even though the nudge is still running. + await asyncio.wait_for( + runner._control.handle(_control("play", scope="room", room_id="r1")), + timeout=1.0, + ) + await started.wait() # the nudge did start, in the background + assert runner._control_tasks # tracked for cancellation on close() + + for task in tuple(runner._control_tasks): + task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await asyncio.gather(*runner._control_tasks) + + +# --------------------------------------------------------------------------- +# Graceful degradation +# --------------------------------------------------------------------------- + + +class TestControlGracefulDegradation: + async def test_unknown_mode_combinations_never_raise(self) -> None: + """Defensive: any well-formed payload must not raise, even degenerate + combinations (e.g. play for a room with nothing pending).""" + runner, _, link = _build_runner() + link.get_next_message = AsyncMock(return_value=None) + + for mode in ("interrupt", "stop", "play"): + await runner._control.handle( + _control(mode, scope="room", room_id="ghost-room") + ) + await _drain_control_tasks(runner) diff --git a/tests/conftest.py b/tests/conftest.py index 552052d7d..9c6acbb99 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,6 +15,7 @@ from __future__ import annotations +import asyncio from datetime import datetime, timezone from pathlib import Path from unittest.mock import AsyncMock @@ -159,6 +160,61 @@ def isolated_guard(agent_id): guard.release() +# ============================================================================= +# Controllable on_execute handler (interrupt/stop tests) +# ============================================================================= + + +class BlockingHandler: + """Deterministic ``on_execute`` stand-in for interrupt/stop tests. + + Replaces the started/cancelled-event + hang-until-cancelled pattern that + interrupt/stop tests otherwise hand-roll. On each cycle it records the + message id, sets ``started``, optionally hangs until cancelled so a control + signal can land mid-cycle, and sets ``cancelled`` if the cycle is aborted. + + ``block`` controls the hang: + - ``True`` — every cycle hangs until cancelled. + - ``False`` — no cycle hangs; each completes immediately. + - ``int N`` — only the first ``N`` cycles hang (later ones complete), for + the stop-then-replay case where the first attempt is aborted and the + redelivered message must run to completion. + + ``invoked`` lists the message ids every cycle *entered*; ``completed`` + lists only those that ran to completion (never populated for an aborted + hanging cycle). + """ + + def __init__(self, *, block: bool | int = True, block_seconds: float = 60) -> None: + self.block = block + self.block_seconds = block_seconds + self.started = asyncio.Event() + self.cancelled = asyncio.Event() + self.invoked: list[str] = [] + self.completed: list[str] = [] + self.invocations = 0 + + def _should_block(self) -> bool: + if isinstance(self.block, bool): + return self.block + return self.invocations <= self.block + + async def __call__(self, ctx, event) -> None: + self.invocations += 1 + msg_id = getattr(getattr(event, "payload", None), "id", None) + if msg_id is not None: + self.invoked.append(msg_id) + self.started.set() + try: + if self._should_block(): + await asyncio.sleep(self.block_seconds) + if msg_id is not None: + self.completed.append(msg_id) + except asyncio.CancelledError: + self.cancelled.set() + raise + + # ============================================================================= # Event Factory Helpers (must return SDK-native types for pattern matching) # ============================================================================= diff --git a/tests/conftest_integration.py b/tests/conftest_integration.py index 6825432e0..30f1de049 100644 --- a/tests/conftest_integration.py +++ b/tests/conftest_integration.py @@ -20,8 +20,11 @@ from __future__ import annotations +import asyncio import logging import os +import time +from collections.abc import Callable from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -30,6 +33,9 @@ from dotenv import load_dotenv from band_rest import AsyncRestClient, ChatRoomRequest from band_rest.core.api_error import ApiError +from band_rest.human_api_chats.types.create_my_chat_room_request_chat import ( + CreateMyChatRoomRequestChat, +) from band_rest.types import ParticipantRequest from thenvoi_testing.markers import skip_without_env, skip_without_envs from thenvoi_testing.settings import BaseTestSettings @@ -381,6 +387,23 @@ async def is_room_alive(api_client: AsyncRestClient, chat_id: str) -> bool: return False +async def wait_until( + pred: Callable[[], bool], *, timeout: float, interval: float = 1.0 +) -> bool: + """Poll ``pred`` until it returns True or ``timeout`` seconds elapse. + + Generic polling helper for integration tests that wait on eventual state + (e.g. a message replayed via /next, a cycle cancelled). Returns the final + value of ``pred`` so callers can assert on it directly. + """ + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return True + await asyncio.sleep(interval) + return pred() + + @pytest_asyncio.fixture(scope="session") async def shared_room( session_api_client: AsyncRestClient | None, @@ -480,6 +503,65 @@ async def shared_multi_agent_room( return chat_id +# Marker title stamped on rooms this fixture creates. list_my_chats returns +# every room the user is *in* (owned or merely added to), and the model exposes +# no ownership field — but room-scope stop/play is owner-only. Reusing only +# rooms bearing this marker guarantees we picked one the user created (owns), +# so those signals don't 403. +_USER_OWNED_ROOM_TITLE = "sdk-integration-user-owned" + + +@pytest_asyncio.fixture +async def shared_user_owned_room( + user_api_client: AsyncRestClient | None, + session_api_client: AsyncRestClient | None, + shared_agent1_info: AgentInfo | None, +) -> str | None: + """A **user-owned** chat with Agent 1 as a participant (reused across runs). + + User ownership is what lets the user issue room-scope signals (stop/play) + and other owner-only operations, so control- and permission-oriented + integration tests share this instead of each re-implementing the + reuse-or-create + ``is_room_alive`` dance. The reuse logic keeps it within + the platform's per-agent room limit even though the fixture is + function-scoped. Returns None when credentials/agent are unavailable so + callers can skip cleanly. + """ + if ( + user_api_client is None + or session_api_client is None + or shared_agent1_info is None + ): + return None + + agent_id = shared_agent1_info.id + + # Reuse a live room WE created (marker title => user-owned) that still has + # the agent. Reusing a merely-participating room would 403 on stop/play. + chats = await user_api_client.human_api_chats.list_my_chats() + for room in reversed(chats.data or []): + if room.title != _USER_OWNED_ROOM_TITLE: + continue + if not await is_room_alive(session_api_client, room.id): + continue + parts = await user_api_client.human_api_participants.list_my_chat_participants( + room.id + ) + if any(p.id == agent_id for p in (parts.data or [])): + logger.info("Reusing user-owned room: %s", room.id) + return room.id + + created = await user_api_client.human_api_chats.create_my_chat_room( + chat=CreateMyChatRoomRequestChat(title=_USER_OWNED_ROOM_TITLE) + ) + chat_id = created.data.id + await user_api_client.human_api_participants.add_my_chat_participant( + chat_id, participant=ParticipantRequest(participant_id=agent_id, role="member") + ) + logger.info("Created user-owned room: %s", chat_id) + return chat_id + + # ============================================================================= # Function-Scoped Test Resource Fixtures # ============================================================================= diff --git a/tests/integration/conftest.py b/tests/integration/conftest.py index 033e1c7de..24c1f438e 100644 --- a/tests/integration/conftest.py +++ b/tests/integration/conftest.py @@ -36,6 +36,7 @@ shared_agent2_info, shared_multi_agent_room, shared_room, + shared_user_owned_room, shared_user_peer, # Test fixtures test_chat, @@ -48,6 +49,7 @@ # Helpers fetch_all_context, is_room_alive, + wait_until, ) # NOTE: pytestmark in conftest.py is NOT applied to collected tests. @@ -94,6 +96,7 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: "shared_agent2_info", "shared_multi_agent_room", "shared_room", + "shared_user_owned_room", "shared_user_peer", # Test fixtures "test_chat", @@ -106,4 +109,5 @@ def pytest_collection_modifyitems(items: list[pytest.Item]) -> None: # Helpers "fetch_all_context", "is_room_alive", + "wait_until", ] diff --git a/tests/integration/test_control_signals.py b/tests/integration/test_control_signals.py new file mode 100644 index 000000000..d70ca1b09 --- /dev/null +++ b/tests/integration/test_control_signals.py @@ -0,0 +1,212 @@ +"""Integration tests for control signals (interrupt / stop / play). + +These exercise the one thing the unit tests mock out: the real platform->SDK +round trip. A real ``agent.control`` push emitted by the platform must reach +``BandLink._on_control`` -> ``AgentRuntime.handle_control`` -> the per-room +``ExecutionContext`` and take effect. + +Design notes: +- A real ``AgentRuntime`` is driven with a *controllable* ``on_execute`` handler + (no LLM) so cycle timing is deterministic — this stays an integration test, + not an e2e/LLM test. +- The room is **user-owned** so the user may issue room-scope stop/play + (``POST /me/chats/{id}/agents/{stop,play}``). Interrupt is agent-scope + (``POST /me/agents/{id}/executions/{exec}/interrupt``) and requires the user + to own the agent; the test skips cleanly if that authorization is absent. +- The control REST endpoints are not in the generated human client yet + (regeneration tracked separately), so they are called via raw ``httpx``, + mirroring the documented control-endpoint contract. Swap to the typed + client once regenerated. + +Gated on ``BAND_API_KEY`` + ``BAND_API_KEY_USER``; not run in CI. + +Run with: + uv run pytest tests/integration/test_control_signals.py -v -s +""" + +from __future__ import annotations + +import asyncio +import logging + +import httpx +import pytest +import pytest_asyncio +from band_rest import ChatMessageRequest +from band_rest.types import ChatMessageRequestMentionsItem + +from band.platform.link import BandLink +from band.runtime.runtime import AgentRuntime +from tests.conftest import BlockingHandler +from tests.integration.conftest import ( + get_api_key, + get_base_url, + get_user_api_key, + get_ws_url, + requires_api, + requires_user_api, + wait_until, +) + +logger = logging.getLogger(__name__) + +_API = "/api/v1" + + +# --------------------------------------------------------------------------- # +# Raw HTTP helpers (user key) for the control endpoints not yet in the client +# --------------------------------------------------------------------------- # + + +def _user_headers() -> dict[str, str]: + return {"X-API-Key": get_user_api_key() or ""} + + +async def _post(path: str) -> httpx.Response: + async with httpx.AsyncClient(timeout=30.0) as client: + return await client.post(f"{get_base_url()}{path}", headers=_user_headers()) + + +async def _get(path: str) -> httpx.Response: + async with httpx.AsyncClient(timeout=30.0) as client: + return await client.get(f"{get_base_url()}{path}", headers=_user_headers()) + + +async def _send_user_mention( + user_api_client, chat_id: str, agent_id: str, text: str +) -> str: + """Post a message from the user that @mentions the agent; return message id.""" + resp = await user_api_client.human_api_messages.send_my_chat_message( + chat_id, + message=ChatMessageRequest( + content=text, + mentions=[ChatMessageRequestMentionsItem(id=agent_id)], + ), + ) + return resp.data.id + + +# --------------------------------------------------------------------------- # +# Fixtures +# --------------------------------------------------------------------------- # + + +@pytest_asyncio.fixture +async def control_runtime(shared_user_owned_room, shared_agent1_info): + """A live AgentRuntime for the agent with a controllable handler. + + Wires ``link.on_control`` exactly as PlatformRuntime does. Always resumes the + room and stops the runtime on teardown so a failed test can't leave the agent + parked in a stopped state. + """ + if shared_user_owned_room is None or shared_agent1_info is None: + pytest.skip("user-owned room/agent unavailable") + + agent_id = shared_agent1_info.id + link = BandLink( + agent_id=agent_id, + api_key=get_api_key() or "", + ws_url=get_ws_url(), + rest_url=get_base_url(), + ) + # Non-blocking by default: the stop/play test needs cycles to complete so + # the replayed message lands in ``completed``. The interrupt test flips + # ``handler.block = True`` to make a cycle hang mid-flight. + handler = BlockingHandler(block=False) + runtime = AgentRuntime(link=link, agent_id=agent_id, on_execute=handler) + link.on_control = runtime.handle_control # mirror PlatformRuntime wiring + + await runtime.start() + await asyncio.sleep(2.0) # let presence subscribe to the room + try: + yield runtime, handler, agent_id, shared_user_owned_room + finally: + try: + await _post(f"{_API}/me/chats/{shared_user_owned_room}/agents/play") + except Exception: # noqa: BLE001 - best-effort un-park + logger.warning("teardown play failed", exc_info=True) + await runtime.stop() + + +# --------------------------------------------------------------------------- # +# Tests +# --------------------------------------------------------------------------- # + + +@requires_api +@requires_user_api +@pytest.mark.asyncio(loop_scope="session") +class TestControlSignalsIntegration: + async def test_stop_suppresses_then_play_replays( + self, control_runtime, user_api_client + ): + """Stop silences the agent in the room; play replays the missed mention.""" + runtime, handler, agent_id, chat_id = control_runtime + + # STOP (room-scope; the user owns this room) + stop = await _post(f"{_API}/me/chats/{chat_id}/agents/stop") + if stop.status_code == 404: + pytest.skip("control endpoints not deployed on this platform") + assert stop.status_code == 200, f"stop failed: {stop.status_code} {stop.text}" + + # A mention sent while stopped must NOT be processed by the handler. + stopped_msg_id = await _send_user_mention( + user_api_client, chat_id, agent_id, "ping while stopped — please stay quiet" + ) + await asyncio.sleep(10.0) + assert stopped_msg_id not in handler.completed, ( + "agent processed a message while stopped" + ) + + # PLAY -> the platform replays the queued mention; the SDK catches up + # via /next and the handler now processes it. + play = await _post(f"{_API}/me/chats/{chat_id}/agents/play") + assert play.status_code == 200, f"play failed: {play.status_code} {play.text}" + + replayed = await wait_until( + lambda: stopped_msg_id in handler.completed, timeout=90.0 + ) + assert replayed, "play did not replay the message missed while stopped" + + async def test_interrupt_aborts_in_flight_cycle( + self, control_runtime, user_api_client + ): + """Interrupt cancels a cycle already in flight; nothing is delivered.""" + runtime, handler, agent_id, chat_id = control_runtime + + # Make sure we start un-stopped, then arm the handler to hang mid-cycle. + await _post(f"{_API}/me/chats/{chat_id}/agents/play") + handler.block = True + handler.started.clear() + + msg_id = await _send_user_mention( + user_api_client, chat_id, agent_id, "begin a long task and keep working" + ) + assert await wait_until(handler.started.is_set, timeout=90.0), ( + "cycle never started — cannot test interrupt" + ) + + # Interrupt is agent-scope and needs an execution_id; it requires the + # user to own the agent. Skip cleanly if not authorized. + ex = await _get(f"{_API}/me/agents/{agent_id}/executions") + if ex.status_code in (401, 403, 404): + pytest.skip("user not authorized for agent-scope interrupt") + executions = ex.json().get("data") or [] + if not executions: + pytest.skip("no executions available to target") + + # The agent-scope signal fans out to all of the agent's rooms, so any + # execution id cancels the in-flight cycle in our room. + interrupt = await _post( + f"{_API}/me/agents/{agent_id}/executions/{executions[0]['id']}/interrupt" + ) + assert interrupt.status_code == 200, ( + f"interrupt failed: {interrupt.status_code} {interrupt.text}" + ) + + assert await wait_until(handler.cancelled.is_set, timeout=45.0), ( + "in-flight cycle was not interrupted" + ) + assert msg_id not in handler.completed, ( + "interrupted cycle still delivered output" + ) diff --git a/tests/platform/test_link.py b/tests/platform/test_link.py index 02d93b187..e5429774f 100644 --- a/tests/platform/test_link.py +++ b/tests/platform/test_link.py @@ -116,6 +116,7 @@ async def test_connect_creates_websocket(self, mock_ws_class, mock_ws_client): mock_ws_client.join_agent_control_channel.assert_called_once_with( link.agent_id, on_supersede=link._on_supersede, + on_control=link._on_control, ) assert link.is_connected is True diff --git a/tests/platform/test_link_control.py b/tests/platform/test_link_control.py new file mode 100644 index 000000000..df1ab9f78 --- /dev/null +++ b/tests/platform/test_link_control.py @@ -0,0 +1,84 @@ +"""Tests for BandLink agent.control signal routing. + +Control signals must be delivered preemptively, directly from the WebSocket +receive task — NOT enqueued on the serialized _event_queue (which would make a +control signal wait behind the very message cycle it is meant to interrupt). +""" + +from __future__ import annotations + +from unittest.mock import AsyncMock + +import pytest + +from band.client.streaming import AgentControlPayload +from band.platform.link import BandLink + + +@pytest.fixture +def link() -> BandLink: + return BandLink( + agent_id="agent-123", api_key="k", ws_url="wss://x/ws", rest_url="https://x" + ) + + +class TestOnControlRouting: + async def test_on_control_invokes_registered_hook(self, link: BandLink): + """_on_control should call the registered on_control hook with the payload.""" + received: list[AgentControlPayload] = [] + + async def hook(payload: AgentControlPayload) -> None: + received.append(payload) + + link.on_control = hook + payload = AgentControlPayload( + mode="interrupt", scope="room", agent_id="agent-123", room_id="room-1" + ) + + await link._on_control(payload) + + assert received == [payload] + + async def test_on_control_does_not_enqueue(self, link: BandLink): + """Control must bypass the serialized event queue entirely.""" + link.on_control = AsyncMock() + payload = AgentControlPayload(mode="stop", scope="agent", agent_id="agent-123") + + await link._on_control(payload) + + assert link._event_queue.empty() + + async def test_on_control_without_hook_is_safe_noop(self, link: BandLink): + """A control push with no hook registered must not raise or enqueue.""" + payload = AgentControlPayload(mode="play", scope="agent", agent_id="agent-123") + + await link._on_control(payload) # should not raise + + assert link._event_queue.empty() + + async def test_connect_wires_on_control_to_channel(self): + """connect() must pass _on_control into join_agent_control_channel.""" + link = BandLink( + agent_id="agent-123", api_key="k", ws_url="wss://x/ws", rest_url="https://x" + ) + + fake_ws = AsyncMock() + fake_ws.__aenter__ = AsyncMock(return_value=fake_ws) + fake_ws.join_agent_control_channel = AsyncMock() + + async def fake_factory(*args, **kwargs): + return fake_ws + + # Patch WebSocketClient construction to return our fake. + import band.platform.link as link_mod + + orig = link_mod.WebSocketClient + link_mod.WebSocketClient = lambda *a, **k: fake_ws # type: ignore[assignment] + try: + await link.connect() + finally: + link_mod.WebSocketClient = orig + + fake_ws.join_agent_control_channel.assert_awaited_once() + kwargs = fake_ws.join_agent_control_channel.await_args.kwargs + assert kwargs.get("on_control") == link._on_control diff --git a/tests/runtime/test_execution_interrupt.py b/tests/runtime/test_execution_interrupt.py new file mode 100644 index 000000000..8995b8f1d --- /dev/null +++ b/tests/runtime/test_execution_interrupt.py @@ -0,0 +1,450 @@ +"""Tests for ExecutionContext per-cycle interrupt/stop. + +The reasoning cycle runs as a cancellable child task so a control signal can +abort just one turn without killing the room's process loop. These tests drive +``_process_event`` directly with a controllable ``on_execute`` so we can cancel +mid-cycle deterministically. +""" + +from __future__ import annotations + +import asyncio +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from band.runtime.execution import ExecutionContext, _BacklogProcessResult +from band.runtime.types import PlatformMessage, SessionConfig +from tests.conftest import BlockingHandler, make_message_event + + +@pytest.fixture +def mock_link(): + link = MagicMock() + link.agent_id = "agent-123" + link.rest = MagicMock() + link.rest.agent_api_participants = MagicMock() + link.rest.agent_api_participants.list_agent_chat_participants = AsyncMock( + return_value=MagicMock(data=[]) + ) + link.rest.agent_api_context = MagicMock() + link.rest.agent_api_context.get_agent_chat_context = AsyncMock( + return_value=MagicMock(data=[]) + ) + link.mark_processing = AsyncMock(return_value=True) + link.mark_processed = AsyncMock(return_value=True) + link.mark_failed = AsyncMock(return_value=True) + link.get_next_message = AsyncMock(return_value=None) + link.get_stale_processing_messages = AsyncMock(return_value=[]) + return link + + +def _backlog_message(msg_id: str = "msg-bk") -> PlatformMessage: + return PlatformMessage( + id=msg_id, + room_id="room-123", + content="hi", + sender_id="user-1", + sender_type="User", + sender_name="User One", + message_type="text", + metadata={}, + created_at=None, + ) + + +class TestInterruptInFlightCycle: + async def test_interrupt_cancels_cycle_marks_processed_loop_alive(self, mock_link): + """Interrupt aborts the cycle, sends nothing, consumes the message, and + the loop stays alive to process a fresh message afterward.""" + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + + event = make_message_event(msg_id="msg-1") + proc = asyncio.create_task(ctx._process_event(event)) + await handler.started.wait() + + # Interrupt from the "receive task" side. + assert ctx.interrupt() is True + + result = await proc + assert result is True # loop continues, not a failure + assert handler.cancelled.is_set() # cycle was aborted + assert "msg-1" not in handler.completed # nothing was sent + # Consumed: durable mark + local dedupe. + mock_link.mark_processed.assert_awaited_once_with("room-123", "msg-1") + assert "msg-1" in ctx.claims.completed_ids(ctx.room_id) + assert ctx._interrupt_kind is None # flag cleared + assert ctx._active_cycle_task is None + + # Loop still alive: a fresh message processes normally. + handler2 = BlockingHandler(block=False) + ctx._on_execute = handler2 + result2 = await ctx._process_event(make_message_event(msg_id="msg-2")) + assert result2 is True + assert handler2.completed == ["msg-2"] + + async def test_interrupt_during_tool_call_drops_result(self, mock_link): + """A tool call already executing is abandoned (await dropped); its + result is never sent.""" + in_tool = asyncio.Event() + tool_results: list[str] = [] + + async def fake_tool(): + await asyncio.sleep(60) + return "tool-output" + + async def on_execute(ctx, event): + in_tool.set() + result = await fake_tool() + tool_results.append(result) # must never run + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="m"))) + await in_tool.wait() + + ctx.interrupt() + await proc + + assert tool_results == [] # tool result abandoned, not delivered + + async def test_interrupt_between_cycles_is_noop(self, mock_link): + """Interrupt with no active cycle is a clean no-op and must not set the + flag (which would mis-flag the next cycle).""" + ctx = ExecutionContext("room-123", mock_link, AsyncMock(), agent_id="agent-123") + assert ctx.interrupt() is False + assert ctx._interrupt_kind is None + + # Next cycle runs normally. + result = await ctx._process_event(make_message_event(msg_id="m1")) + assert result is True + mock_link.mark_processed.assert_awaited_once_with("room-123", "m1") + + +class TestStopInFlightCycle: + async def test_stop_leaves_message_actionable(self, mock_link): + """Stop aborts the cycle but leaves the message in 'processing' (no + mark_processed, not remembered) so the platform replays it on play.""" + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="s1"))) + await handler.started.wait() + + ctx.interrupt(kind="stop") + result = await proc + + assert result is True + mock_link.mark_processed.assert_not_awaited() + assert "s1" not in ctx.claims.completed_ids(ctx.room_id) + # Local in-flight claim released so it can be reprocessed on play. + assert "s1" not in ctx.claims.inflight_ids(ctx.room_id) + + +class TestShutdownVsInterrupt: + async def test_shutdown_cancels_cycle_without_marking(self, mock_link): + """stop() (shutdown) cancels an in-flight cycle, does NOT mark it + processed, and leaves no orphaned child task.""" + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + await ctx.start() + + # Feed a message through the running loop. + await ctx.on_event(make_message_event(msg_id="sd1")) + await handler.started.wait() + + cycle_task = ctx._active_cycle_task + assert cycle_task is not None + + graceful = await ctx.stop() + assert graceful is True + mock_link.mark_processed.assert_not_awaited() # shutdown != consume + assert cycle_task.cancelled() or cycle_task.done() + assert ctx._active_cycle_task is None + + +class TestStopRoomResumeRoom: + async def test_stop_room_sets_flag_and_interrupts(self, mock_link): + """stop_room aborts the in-flight cycle and sets the local _stopped flag.""" + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="x"))) + await handler.started.wait() + + ctx.stop_room() + result = await proc + + assert ctx._stopped is True + assert result is True + mock_link.mark_processed.assert_not_awaited() + + async def test_stopped_room_skips_new_message(self, mock_link): + """A WS trigger arriving while stopped is left actionable (never claimed + or marked), not processed.""" + handler = BlockingHandler(block=False) + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + ctx._stopped = True + + result = await ctx._process_event(make_message_event(msg_id="while-stopped")) + + assert result is True + assert handler.invoked == [] # adapter never invoked + mock_link.mark_processing.assert_not_awaited() + mock_link.mark_processed.assert_not_awaited() + + async def test_resume_room_clears_flag_and_requests_resync(self, mock_link): + """play clears _stopped and enqueues a resync sentinel (the /next + rehydration catch-up).""" + ctx = ExecutionContext("room-123", mock_link, AsyncMock(), agent_id="agent-123") + ctx._stopped = True + + await ctx.resume_room() + + assert ctx._stopped is False + # A resync sentinel was enqueued (request_resync) for the loop to catch up. + assert ctx.queue.qsize() == 1 + + async def test_stale_recovery_skipped_while_stopped(self, mock_link): + """Reconnect-while-stopped must NOT resurrect the interrupted message via + the stale-processing recovery sweep (stop-survives-reconnect, locally + guaranteed — not reliant on the platform mark gate).""" + mock_link.get_stale_processing_messages = AsyncMock( + return_value=[_backlog_message("stuck-in-processing")] + ) + handler = BlockingHandler(block=False) + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + ctx._stopped = True + + ok = await ctx._recover_stale_processing_messages() + + assert ok is True + assert handler.invoked == [] # adapter not invoked while stopped + mock_link.get_stale_processing_messages.assert_not_awaited() + + async def test_resume_replays_and_responds(self, mock_link): + """After play, the loop catches up via /next and the adapter runs for + the replayed backlog message.""" + # Backlog has one message waiting (the one left actionable while stopped). + replayed = _backlog_message("replayed-1") + calls = {"n": 0} + + async def get_next(room_id): + calls["n"] += 1 + return replayed if calls["n"] == 1 else None + + mock_link.get_next_message = AsyncMock(side_effect=get_next) + handler = BlockingHandler(block=False) + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + # Simulate: was stopped, now resuming. + ctx._stopped = False + ok = await ctx._resync_pending_messages() + + assert ok is True + assert handler.completed == ["replayed-1"] + + async def test_stop_does_not_poison_retry_budget(self, mock_link): + """A cycle aborted by stop must not count against the message's retry + budget — at the real default of one retry, the replayed backlog + message (delivered again via /next after play) must still reach the + handler instead of immediately landing in permanently_failed.""" + # First attempt hangs (gets stopped); the redelivered replay completes. + handler = BlockingHandler(block=1) + + # Real default (SessionConfig().max_message_retries == 1) — no + # generous override, so a poisoned attempt count would trip this. + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="p1"))) + await handler.started.wait() + ctx.stop_room() + assert await proc is True + mock_link.mark_processed.assert_not_awaited() + + # Resume, then the platform replays the still-"processing" message via + # /next — same message id, now delivered through the backlog path. + handler.started.clear() + result = await ctx._process_backlog_message(_backlog_message("p1")) + + assert result == _BacklogProcessResult.ADVANCED + assert handler.completed == ["p1"] # handler actually ran this time + assert not ctx._retry_tracker.is_permanently_failed("p1") + + +class TestBacklogInterrupt: + async def test_interrupt_during_backlog_consumes_and_advances(self, mock_link): + """Interrupt during a /next backlog cycle consumes the message and the + sync advances (does not retry the interrupted turn).""" + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + msg = _backlog_message("bk1") + proc = asyncio.create_task(ctx._process_backlog_message(msg)) + await handler.started.wait() + + ctx.interrupt() + result = await proc + + assert result == _BacklogProcessResult.ADVANCED + mock_link.mark_processed.assert_awaited_once_with("room-123", "bk1") + assert "bk1" in ctx.claims.completed_ids(ctx.room_id) + + async def test_stop_during_backlog_leaves_actionable(self, mock_link): + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + proc = asyncio.create_task( + ctx._process_backlog_message(_backlog_message("bk2")) + ) + await handler.started.wait() + + ctx.interrupt(kind="stop") + result = await proc + + assert result == _BacklogProcessResult.ADVANCED + mock_link.mark_processed.assert_not_awaited() + assert "bk2" not in ctx.claims.completed_ids(ctx.room_id) + + async def test_stop_mid_resync_loop_does_not_reprocess(self, mock_link): + """A stop landing mid-cycle during ``_resync_pending_messages`` must not + be undone by the very next /next call in the same loop. + + The platform's /next excludes only 'processed' messages, so a + 'processing' message left behind by stop is returned again on the next + poll. If the enclosing loop doesn't notice ``_stopped``, it will + re-claim and fully run the very cycle stop just aborted -- silently + breaking the "stop -> goes quiet until play" contract. + """ + processed_ids: set[str] = set() + + async def fake_mark_processed(room_id, msg_id): + processed_ids.add(msg_id) + return True + + mock_link.mark_processed = AsyncMock(side_effect=fake_mark_processed) + + async def fake_get_next(room_id): + # Mirrors the real /next contract: keep returning the message until + # it's actually marked processed. + if "loop1" in processed_ids: + return None + return _backlog_message("loop1") + + mock_link.get_next_message = AsyncMock(side_effect=fake_get_next) + + handler = BlockingHandler() # hangs until cancelled by stop_room() below + ctx = ExecutionContext( + "room-123", + mock_link, + handler, + agent_id="agent-123", + # A high retry cap so the retry tracker can't itself block a second + # attempt -- the test must fail (or pass) on the _stopped guard + # alone, not be masked by the unrelated max-retries limit. + config=SessionConfig(max_message_retries=10), + ) + + resync_task = asyncio.create_task(ctx._resync_pending_messages()) + await handler.started.wait() + + ctx.stop_room() + result = await asyncio.wait_for(resync_task, timeout=5) + + assert result is True + # The adapter must not run a second time on the very next /next poll. + assert handler.invoked == ["loop1"] + mock_link.mark_processed.assert_not_awaited() + assert "loop1" not in ctx.claims.completed_ids(ctx.room_id) + + +class TestControlSignalInClaimWindow: + """Signals landing after a message is claimed but before its cancellable + cycle task exists (the mark_processing/hydration window) must not be lost. + """ + + @staticmethod + def _gated_mark_processing( + claiming: asyncio.Event, release: asyncio.Event + ) -> AsyncMock: + """mark_processing that parks inside the claim so a test can fire a + signal while the cycle task does not yet exist.""" + + async def _gate(room_id: str, msg_id: str) -> bool: + claiming.set() + await release.wait() + return True + + return AsyncMock(side_effect=_gate) + + async def test_interrupt_in_window_aborts_before_handler(self, mock_link): + claiming, release = asyncio.Event(), asyncio.Event() + mock_link.mark_processing = self._gated_mark_processing(claiming, release) + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="w1"))) + await claiming.wait() + + # No cycle task to cancel yet, but the signal must still take effect. + assert ctx._active_cycle_task is None + assert ctx.interrupt() is True + + release.set() + result = await proc + + assert result is True + assert handler.invoked == [] # handler never ran + mock_link.mark_processed.assert_awaited_once_with("room-123", "w1") # consumed + assert "w1" in ctx.claims.completed_ids(ctx.room_id) + assert ctx._pending_interrupt is None + assert ctx._cycle_armed is False + + async def test_stop_in_window_leaves_message_actionable(self, mock_link): + claiming, release = asyncio.Event(), asyncio.Event() + mock_link.mark_processing = self._gated_mark_processing(claiming, release) + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="w2"))) + await claiming.wait() + + ctx.stop_room() # sets _stopped and records the pending stop + + release.set() + result = await proc + + assert result is True + assert handler.invoked == [] + assert ctx._stopped is True + mock_link.mark_processed.assert_not_awaited() # left for replay on play + assert "w2" not in ctx.claims.completed_ids(ctx.room_id) + + async def test_interrupt_in_backlog_window_aborts_and_advances(self, mock_link): + claiming, release = asyncio.Event(), asyncio.Event() + mock_link.mark_processing = self._gated_mark_processing(claiming, release) + handler = BlockingHandler() + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + + proc = asyncio.create_task( + ctx._process_backlog_message(_backlog_message("bw1")) + ) + await claiming.wait() + + assert ctx.interrupt() is True + + release.set() + result = await proc + + assert result == _BacklogProcessResult.ADVANCED + assert handler.invoked == [] + mock_link.mark_processed.assert_awaited_once_with("room-123", "bw1") + + async def test_interrupt_between_cycles_arms_nothing(self, mock_link): + """A truly idle interrupt (no claim in flight) stays a no-op and does not + arm a pending signal that would mis-flag the next message.""" + handler = BlockingHandler(block=False) + ctx = ExecutionContext("room-123", mock_link, handler, agent_id="agent-123") + + assert ctx.interrupt() is False + assert ctx._pending_interrupt is None + + # The next message runs to completion, unaffected. + result = await ctx._process_event(make_message_event(msg_id="n1")) + assert result is True + assert handler.completed == ["n1"] diff --git a/tests/runtime/test_runtime_control.py b/tests/runtime/test_runtime_control.py new file mode 100644 index 000000000..b125325bc --- /dev/null +++ b/tests/runtime/test_runtime_control.py @@ -0,0 +1,210 @@ +"""Tests for AgentRuntime.handle_control routing/dedup.""" + +from __future__ import annotations + +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from band.client.streaming import AgentControlPayload +from band.runtime.runtime import AgentRuntime + + +def _control(mode: str, scope: str = "agent", **kw) -> AgentControlPayload: + return AgentControlPayload(mode=mode, scope=scope, agent_id="agent-123", **kw) + + +def _fake_execution(room_id: str) -> MagicMock: + """A control-capable execution: interrupt/stop_room sync, resume_room async.""" + ex = MagicMock() + ex.room_id = room_id + ex.interrupt = MagicMock(return_value=True) + ex.stop_room = MagicMock() + ex.resume_room = AsyncMock() + return ex + + +@pytest.fixture +def runtime() -> AgentRuntime: + link = MagicMock() + return AgentRuntime(link=link, agent_id="agent-123", on_execute=AsyncMock()) + + +class TestRouting: + async def test_agent_scope_null_room_fans_out_to_all(self, runtime): + a, b = _fake_execution("r1"), _fake_execution("r2") + runtime.executions = {"r1": a, "r2": b} + + await runtime.handle_control(_control("interrupt", scope="agent", room_id=None)) + + a.interrupt.assert_called_once() + b.interrupt.assert_called_once() + + async def test_room_scope_targets_single_room(self, runtime): + a, b = _fake_execution("r1"), _fake_execution("r2") + runtime.executions = {"r1": a, "r2": b} + + await runtime.handle_control(_control("stop", scope="room", room_id="r2")) + + a.stop_room.assert_not_called() + b.stop_room.assert_called_once() + + async def test_agent_scope_with_room_id_targets_that_room(self, runtime): + a, b = _fake_execution("r1"), _fake_execution("r2") + runtime.executions = {"r1": a, "r2": b} + + await runtime.handle_control(_control("interrupt", scope="agent", room_id="r1")) + + a.interrupt.assert_called_once() + b.interrupt.assert_not_called() + + async def test_unknown_room_is_noop(self, runtime): + a = _fake_execution("r1") + runtime.executions = {"r1": a} + + await runtime.handle_control( + _control("interrupt", scope="room", room_id="ghost") + ) + + a.interrupt.assert_not_called() + + async def test_play_routes_to_resume_room(self, runtime): + a = _fake_execution("r1") + runtime.executions = {"r1": a} + + await runtime.handle_control(_control("play", scope="room", room_id="r1")) + + a.resume_room.assert_awaited_once() + + +class TestDedup: + async def test_duplicate_correlation_id_dropped(self, runtime): + a = _fake_execution("r1") + runtime.executions = {"r1": a} + sig = _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-1") + + await runtime.handle_control(sig) + await runtime.handle_control(sig) # duplicate + + a.interrupt.assert_called_once() + + async def test_distinct_correlation_ids_both_applied(self, runtime): + a = _fake_execution("r1") + runtime.executions = {"r1": a} + + await runtime.handle_control( + _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-1") + ) + await runtime.handle_control( + _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-2") + ) + + assert a.interrupt.call_count == 2 + + async def test_play_after_stop_not_dropped(self, runtime): + """Distinct signals have distinct correlation_ids; play after stop runs.""" + a = _fake_execution("r1") + runtime.executions = {"r1": a} + + await runtime.handle_control( + _control("stop", scope="room", room_id="r1", correlation_id="ctl-stop") + ) + await runtime.handle_control( + _control("play", scope="room", room_id="r1", correlation_id="ctl-play") + ) + + a.stop_room.assert_called_once() + a.resume_room.assert_awaited_once() + + async def test_missing_correlation_id_not_deduped(self, runtime): + a = _fake_execution("r1") + runtime.executions = {"r1": a} + sig = _control("interrupt", scope="room", room_id="r1") # no correlation_id + + await runtime.handle_control(sig) + await runtime.handle_control(sig) + + assert a.interrupt.call_count == 2 # both applied (cannot dedup) + + +class TestStopSurvivesReconnect: + async def test_reconnect_while_stopped_does_not_invoke_adapter(self): + """After stop, a reconnect (ReconnectedEvent -> /next sync incl. stale + recovery) must NOT re-fire the adapter. Needs no SDK persistence — the + local _stopped guard + platform /next->204 keep the room quiet. + + Asserts the adapter is NOT invoked (covers both the no-persistence claim + and the recovery-sweep guard), per architect's Step-4 should-fix. + """ + from band.platform.event import ReconnectedEvent + from band.runtime.execution import ExecutionContext + + link = MagicMock() + link.agent_id = "agent-123" + link.rest = MagicMock() + link.rest.agent_api_participants = MagicMock() + link.rest.agent_api_participants.list_agent_chat_participants = AsyncMock( + return_value=MagicMock(data=[]) + ) + link.rest.agent_api_context = MagicMock() + link.rest.agent_api_context.get_agent_chat_context = AsyncMock( + return_value=MagicMock(data=[]) + ) + link.mark_processing = AsyncMock(return_value=True) + link.mark_processed = AsyncMock(return_value=True) + link.mark_failed = AsyncMock(return_value=True) + # Platform /next gate returns 204 (None) for a stopped agent — that path + # is platform-authoritative. The LOCAL risk is the recovery sweep, which + # fetches 'processing' messages DIRECTLY (bypassing /next): the stop path + # leaves the interrupted message there. The _stopped guard must skip it. + from band.runtime.types import PlatformMessage + + stuck = PlatformMessage( + id="stuck", + room_id="room-123", + content="hi", + sender_id="u1", + sender_type="User", + sender_name="U1", + message_type="text", + metadata={}, + created_at=None, + ) + link.get_stale_processing_messages = AsyncMock(return_value=[stuck]) + link.get_next_message = AsyncMock(return_value=None) # /next gate: 204 + + executed: list[str] = [] + + async def on_execute(ctx, event): + executed.append(event.payload.id) + + ctx = ExecutionContext("room-123", link, on_execute, agent_id="agent-123") + ctx._stopped = True + ctx._reconnect_sync_requested = True + + # Drive the reconnect sync path directly. + await ctx._process_event(ReconnectedEvent()) + + assert executed == [] # adapter never invoked while stopped + link.mark_processing.assert_not_awaited() + # Recovery sweep (the gate-bypassing path) was skipped locally. + link.get_stale_processing_messages.assert_not_awaited() + # Efficiency: the reconnect sync must short-circuit on _stopped locally, + # same as the idle-timeout and resync-sentinel paths, instead of making + # a /next call that's guaranteed to come back empty. + link.get_next_message.assert_not_awaited() + + +class TestGracefulDegradation: + async def test_custom_execution_without_methods_is_skipped(self, runtime): + """A custom Execution lacking the control methods degrades to no-op.""" + + class BareExecution: + room_id = "r1" + + runtime.executions = {"r1": BareExecution()} + + # Must not raise for any mode. + await runtime.handle_control(_control("interrupt", scope="room", room_id="r1")) + await runtime.handle_control(_control("stop", scope="room", room_id="r1")) + await runtime.handle_control(_control("play", scope="room", room_id="r1")) diff --git a/tests/test_platform_runtime.py b/tests/test_platform_runtime.py index 08f9f5acb..7be033c51 100644 --- a/tests/test_platform_runtime.py +++ b/tests/test_platform_runtime.py @@ -293,6 +293,22 @@ async def test_starts_agent_runtime(self, mock_link, mock_runtime): mock_runtime.start.assert_awaited_once() + @pytest.mark.asyncio + async def test_wires_control_hook_to_runtime(self, mock_link, mock_runtime): + """start() should route control signals to AgentRuntime.handle_control.""" + mock_runtime.handle_control = AsyncMock() + with patch("band.runtime.platform_runtime.BandLink") as mock_link_class: + mock_link_class.return_value = mock_link + with patch( + "band.runtime.platform_runtime.AgentRuntime" + ) as mock_runtime_class: + mock_runtime_class.return_value = mock_runtime + + runtime = PlatformRuntime(agent_id="agent-123", api_key="test-key") + await runtime.start(on_execute=AsyncMock()) + + assert mock_link.on_control == mock_runtime.handle_control + @pytest.mark.asyncio async def test_raises_on_missing_metadata(self, mock_link, mock_runtime): """Should raise if agent metadata response is empty.""" diff --git a/tests/websocket/test_control_payload.py b/tests/websocket/test_control_payload.py new file mode 100644 index 000000000..d0099039c --- /dev/null +++ b/tests/websocket/test_control_payload.py @@ -0,0 +1,83 @@ +"""Unit tests for the agent.control WebSocket payload model.""" + +import pytest +from pydantic import ValidationError + +from band.client.streaming import AgentControlPayload +from band.client.streaming.client import _PAYLOAD_MODELS + + +class TestAgentControlPayload: + """Tests for AgentControlPayload model.""" + + @pytest.mark.parametrize("mode", ["interrupt", "stop", "play"]) + def test_valid_modes(self, mode): + """Should accept each of the three control modes.""" + payload = AgentControlPayload( + type="agent.control", + mode=mode, + scope="agent", + agent_id="agent-123", + execution_id="exec-1", + room_id="room-1", + reason="user_requested", + correlation_id="ctl-abc", + ) + assert payload.mode == mode + assert payload.scope == "agent" + assert payload.agent_id == "agent-123" + assert payload.execution_id == "exec-1" + assert payload.room_id == "room-1" + assert payload.correlation_id == "ctl-abc" + + def test_room_scope(self): + """Should accept room scope.""" + payload = AgentControlPayload( + mode="interrupt", + scope="room", + agent_id="agent-123", + room_id="room-9", + ) + assert payload.scope == "room" + assert payload.room_id == "room-9" + + def test_null_room_and_execution(self): + """Agent-scoped fan-out: room_id and execution_id may be omitted/null.""" + payload = AgentControlPayload( + mode="stop", + scope="agent", + agent_id="agent-123", + ) + assert payload.room_id is None + assert payload.execution_id is None + assert payload.reason is None + assert payload.correlation_id is None + + def test_bad_mode_rejected(self): + """Should reject unknown modes (e.g. legacy 'pause').""" + with pytest.raises(ValidationError): + AgentControlPayload(mode="pause", scope="agent", agent_id="a") + + def test_bad_scope_rejected(self): + """Should reject unknown scopes.""" + with pytest.raises(ValidationError): + AgentControlPayload(mode="stop", scope="cluster", agent_id="a") + + def test_missing_agent_id_rejected(self): + """Should require agent_id.""" + with pytest.raises(ValidationError): + AgentControlPayload(mode="stop", scope="agent") + + def test_extra_fields_allowed(self): + """Should allow extra fields for forward compatibility.""" + payload = AgentControlPayload( + mode="play", + scope="agent", + agent_id="agent-123", + future_field="allowed", + ) + assert payload.agent_id == "agent-123" + + def test_registered_in_payload_models(self): + """The dispatcher must know how to parse agent.control events.""" + assert _PAYLOAD_MODELS.get("agent.control") is AgentControlPayload