diff --git a/band-bridge/bridge_core/bridge.py b/band-bridge/bridge_core/bridge.py index 7f020f2c6..465f69d4a 100644 --- a/band-bridge/bridge_core/bridge.py +++ b/band-bridge/bridge_core/bridge.py @@ -120,8 +120,8 @@ def __init__( # 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 + # Background tasks spawned by control signals (play nudges) so the + # next control signal isn't queued behind 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() @@ -234,8 +234,6 @@ 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: @@ -244,6 +242,9 @@ async def close(self) -> None: self._config.agent_id, exc_info=True, ) + # After disconnect: the link applies control signals until it returns, + # so a PLAY can still spawn a nudge during it. + await self._cancel_control_tasks() try: await self._forwarder.close() except Exception: @@ -253,6 +254,13 @@ async def close(self) -> None: exc_info=True, ) + async def _cancel_control_tasks(self) -> None: + pending = tuple(self._control_tasks) + for task in pending: + task.cancel() + if pending: + await asyncio.wait(pending) + async def _connect_and_consume(self) -> None: """Connect, subscribe, and consume events until shutdown.""" # BandLink.connect() joins the control channel, whose pushes may arrive @@ -589,10 +597,10 @@ def _remember_processed_message(self, message_id: str) -> None: 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`` + Runs off BandLink's control lane 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: diff --git a/band-bridge/bridge_core/control.py b/band-bridge/bridge_core/control.py index 8fa6b9900..773374cf5 100644 --- a/band-bridge/bridge_core/control.py +++ b/band-bridge/bridge_core/control.py @@ -40,8 +40,8 @@ def __init__(self, runner: AgentRunner) -> None: # 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). + # BandLink delivers control signals one at a time (as for + # AgentRuntime's ControlDeduplicator on the long-running SDK path). self._seen_control_ids: OrderedDict[str, bool] = OrderedDict() async def handle(self, payload: AgentControlPayload) -> None: @@ -76,7 +76,8 @@ async def handle(self, payload: AgentControlPayload) -> None: 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. + # queued behind this /next + forward: BandLink applies control + # signals one at a time. self._runner._spawn_play_nudge(room_ids) case _: logger.warning( diff --git a/src/band/client/streaming/client.py b/src/band/client/streaming/client.py index 3ab5acabd..e0cb623da 100644 --- a/src/band/client/streaming/client.py +++ b/src/band/client/streaming/client.py @@ -261,6 +261,9 @@ class AgentControlPayload(WirePayload): correlation_id: str | None = None +ControlHook = Callable[[AgentControlPayload], Awaitable[None]] + + class SupersedePayload(WirePayload): """Payload for terminal agent_control supersede events.""" @@ -764,7 +767,7 @@ async def join_agent_control_channel( self, agent_id: str, on_supersede: Callable[[SupersedePayload], Awaitable[None]], - on_control: Callable[[AgentControlPayload], Awaitable[None]] | None = None, + on_control: ControlHook | None = None, ): """Subscribe to agent-control events for this agent. diff --git a/src/band/platform/link.py b/src/band/platform/link.py index 524c3cd07..7edc3605f 100644 --- a/src/band/platform/link.py +++ b/src/band/platform/link.py @@ -9,7 +9,7 @@ import asyncio import logging -from collections.abc import Awaitable, Callable +from collections.abc import Callable from typing import TYPE_CHECKING from band_sdk_core import AgentTopicKind, SessionState @@ -20,6 +20,7 @@ WebSocketDisconnectReason, WebSocketUpgradeError, ) +from band.client.streaming.client import ControlHook from band.config.settings import DEFAULT_REST_URL, DEFAULT_WS_URL from band.core.exceptions import AgentAlreadyRunningError from band.core.types import ConflictPolicy, PlatformConnection @@ -134,10 +135,13 @@ def __init__( 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 + # runtime. Never run via the serialized _event_queue, so a control + # signal can act on a cycle already in flight instead of queuing + # behind it. Signals apply one at a time, so a hook that never + # returns holds every later signal until disconnect(). + self.on_control: ControlHook | None = None + # Control signals still being applied, in arrival order. + self._control_lane: list[asyncio.Task[None]] = [] @property def is_connected(self) -> bool: @@ -232,6 +236,7 @@ async def disconnect(self) -> None: except Exception as e: # noqa: BLE001 -- best-effort event emission must not crash the turn/link logger.warning("Error unsubscribing from agent_control: %s", e) + await self._cancel_control_lane() await self._ws.__aexit__(None, None, None) self._ws = None self._is_connected = False @@ -416,18 +421,57 @@ async def _on_supersede(self, payload: SupersedePayload) -> None: 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. + Runs the ``on_control`` hook on a task ordered after every earlier + signal. The hook can take seconds (a STOP drains every room's turn), + and awaiting it here would hold a ``supersede`` queued behind it on + the same topic, where a reconnect discards queued messages unseen. + If no hook is registered, the push is a safe no-op. """ - if self.on_control is None: + hook = self.on_control + if hook is None: logger.debug( "agent.control received (mode=%s) but no on_control hook registered", payload.mode, ) return - await self.on_control(payload) + previous = self._control_lane[-1] if self._control_lane else None + task = asyncio.create_task( + self._apply_control(previous, hook, payload), + name=f"agent.control {payload.mode} correlation_id={payload.correlation_id}", + ) + self._control_lane.append(task) + task.add_done_callback(self._control_lane.remove) + + @staticmethod + async def _apply_control( + previous: asyncio.Task[None] | None, + hook: ControlHook, + payload: AgentControlPayload, + ) -> None: + if previous is not None: + await asyncio.wait({previous}) + try: + await hook(payload) + except Exception: + logger.exception( + "agent.control %s failed (correlation_id=%s room_id=%s)", + payload.mode, + payload.correlation_id, + payload.room_id, + ) + + async def _cancel_control_lane(self) -> None: + if not self._control_lane: + return + pending = list(self._control_lane) + logger.info( + "Cancelling %d unfinished control signal(s) on disconnect: %s", + len(pending), + ", ".join(task.get_name() for task in pending), + ) + for task in pending: + task.cancel() + await asyncio.wait(pending) async def _on_disconnected(self, error: Exception | None) -> None: """Handle PHX client disconnection.""" diff --git a/src/band/runtime/control.py b/src/band/runtime/control.py index ad901167d..272fe7f80 100644 --- a/src/band/runtime/control.py +++ b/src/band/runtime/control.py @@ -23,7 +23,7 @@ class ControlDeduplicator: """Drops repeated signals by correlation id; the server does not. - Bounded LRU, only touched from the WebSocket receive task. + Bounded LRU; ``BandLink`` delivers control signals one at a time. """ def __init__(self) -> None: diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index 0fdd40dcf..701180cca 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -252,9 +252,10 @@ async def request_resync(self) -> None: def interrupt(self, *, kind: ControlMode | str = ControlMode.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. + Called preemptively, outside the room's message queue, on an + ``interrupt`` or ``stop`` control signal. ``AgentRuntime`` + ``hasattr``-guards this, so custom ``Execution`` implementations that + omit it degrade to a no-op. ``kind`` keeps ``str`` in the union because ``Protocol`` signatures aren't runtime-enforced and a custom implementation may still pass a plain string. @@ -691,8 +692,8 @@ async def request_resync(self) -> None: def interrupt(self, *, kind: ControlMode | str = ControlMode.INTERRUPT) -> bool: """Abort the in-flight reasoning cycle, if any. - Called from the WebSocket receive task. It only terminates the running - turn scope; all message-status / dedupe bookkeeping happens once in + Called on a control signal, outside the cycle. It only terminates the + running turn scope; all message-status / dedupe bookkeeping happens once in ``_run_cycle``, never here, so the two coroutines do not race on shared state. diff --git a/src/band/runtime/runtime.py b/src/band/runtime/runtime.py index 35cbdad7c..9035072a4 100644 --- a/src/band/runtime/runtime.py +++ b/src/band/runtime/runtime.py @@ -307,8 +307,8 @@ async def _on_reconnected(self) -> None: 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. + Invoked via ``BandLink.on_control``, outside the serialized event + queue, so it can preempt a cycle already in flight. Routing: - ``scope == "agent"`` with no ``room_id`` → fan out to all executions. @@ -318,8 +318,8 @@ async def handle_control(self, payload: AgentControlPayload) -> None: Dedupes on ``correlation_id`` (the server does not). A signal that reached no room is not recorded, so a later copy still applies once the room is routed. Note: exceptions raised here are - swallowed-and-logged by the channel handler, so this method must not - rely on propagating errors to signal anything. + only logged by ``BandLink``, so this method must not rely on + propagating errors to signal anything. """ targets = control_targets(payload, self.executions) if not targets or not self._control_ids.first_sighting(payload): diff --git a/tests/bridge/test_close_after_control.py b/tests/bridge/test_close_after_control.py new file mode 100644 index 000000000..2a7006027 --- /dev/null +++ b/tests/bridge/test_close_after_control.py @@ -0,0 +1,43 @@ +"""``AgentRunner.close()`` cancels a play nudge spawned while its link is still +disconnecting: the link applies control signals until ``disconnect()`` returns.""" + +from __future__ import annotations + +import asyncio + +from band.client.streaming import AgentControlPayload +from tests.bridge.test_bridge import _build_runner + + +class TestCloseAfterControl: + async def test_a_play_applied_during_disconnect_never_outlives_close( + self, + ) -> None: + runner, _, link = _build_runner() + story: list[str] = [] + nudge_started = asyncio.Event() + + async def next_message(room_id: str) -> None: + story.append(f"nudge {room_id}") + nudge_started.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + story.append(f"nudge {room_id} cancelled") + raise + + async def disconnect() -> None: + """A PLAY queued in the link's control lane applies before it ends.""" + await runner._control.handle( + AgentControlPayload( + mode="play", scope="room", agent_id="agent-1", room_id="r1" + ) + ) + await nudge_started.wait() + + link.get_next_message.side_effect = next_message + link.disconnect.side_effect = disconnect + + await asyncio.wait_for(runner.close(), timeout=5.0) + + assert story == ["nudge r1", "nudge r1 cancelled"] diff --git a/tests/conftest.py b/tests/conftest.py index 90cf0cad0..f2a7061ab 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -354,6 +354,18 @@ async def __call__(self, ctx, event) -> None: # ============================================================================= +# The platform's eviction of an older same-key session (``agent_control`` +# ``supersede``), as the wire carries it. +TERMINAL_SUPERSEDE: dict[str, Any] = { + "reason": "session.already_connected", + "message": "Superseded by a newer session", + "retryable": False, + "retry_after": 15, + "target_socket_id": "agent_socket:agent-123", + "correlation_id": "evict-1", +} + + def spy_on_reconciliation_drain(link: BandLink) -> asyncio.Event: """Arm a spy on ``_drain_reconciliation`` -- the last step ``_on_reconnected`` runs -- and return an event set the instant the real diff --git a/tests/platform/conftest.py b/tests/platform/conftest.py index c88f2b20a..26f2bbf35 100644 --- a/tests/platform/conftest.py +++ b/tests/platform/conftest.py @@ -12,6 +12,7 @@ from band_sdk_core import DeadReason, SessionState from band.client.streaming import WebSocketClient +from band.platform.link import BandLink class AllTopicsJoined: @@ -34,6 +35,14 @@ def __contains__(self, topic: object) -> bool: return topic not in self._missing +@pytest.fixture +def link() -> BandLink: + """An unconnected BandLink, for driving its handlers directly.""" + return BandLink( + agent_id="agent-123", api_key="k", ws_url="wss://x/ws", rest_url="https://x" + ) + + @pytest.fixture def mock_ws_client(): """Autospecced WebSocketClient for testing BandLink. diff --git a/tests/platform/linkclose/__init__.py b/tests/platform/linkclose/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/tests/platform/linkclose/kit.py b/tests/platform/linkclose/kit.py new file mode 100644 index 000000000..ded730406 --- /dev/null +++ b/tests/platform/linkclose/kit.py @@ -0,0 +1,66 @@ +"""Shared helpers for driving a live BandLink's closes and supersedes over a +real (fake-server-backed) wire.""" + +from __future__ import annotations + +import asyncio +import logging +from collections.abc import AsyncIterator +from contextlib import asynccontextmanager + +from band_sdk_core import AgentTopicKind, DeadReason + +from band.client.streaming import AgentControlPayload, ControlMode +from band.client.streaming.client import WebSocketDisconnectReason +from band.platform.link import BandLink +from tests.conftest import TERMINAL_SUPERSEDE + +AGENT_ID = "agent-123" +CONTROL_TOPIC = AgentTopicKind.Control.topic(AGENT_ID) +HANG_GUARD_S = 5.0 + +SUPERSEDED = WebSocketDisconnectReason( + **TERMINAL_SUPERSEDE, dead_reason=DeadReason.Classified +) + + +def make_link(server_url: str) -> BandLink: + return BandLink( + agent_id=AGENT_ID, + api_key="test-key", + ws_url=server_url, + rest_url="https://test.invalid", + ) + + +@asynccontextmanager +async def connected(link: BandLink) -> AsyncIterator[BandLink]: + await link.connect() + try: + yield link + finally: + await link.disconnect() + + +@asynccontextmanager +async def running(link: BandLink) -> AsyncIterator[asyncio.Task[None]]: + """``run_forever`` as a task the host is already blocked on.""" + run = asyncio.create_task(link.run_forever(install_signal_handlers=False)) + try: + yield run + finally: + run.cancel() + await asyncio.wait({run}) + + +def control_cancellations(records: list[logging.LogRecord]) -> list[str]: + """What ``disconnect()`` reported cancelling from the control lane.""" + return [ + r.getMessage() for r in records if r.msg.startswith("Cancelling %d unfinished") + ] + + +def agent_control(mode: ControlMode, correlation_id: str) -> AgentControlPayload: + return AgentControlPayload( + mode=mode, scope="agent", agent_id=AGENT_ID, correlation_id=correlation_id + ) diff --git a/tests/platform/linkclose/test_disconnect_cancels_control.py b/tests/platform/linkclose/test_disconnect_cancels_control.py new file mode 100644 index 000000000..bffb60e90 --- /dev/null +++ b/tests/platform/linkclose/test_disconnect_cancels_control.py @@ -0,0 +1,80 @@ +"""``disconnect()`` cancels every control signal not yet finished, the running +one and those queued behind it, and leaves applied ones alone.""" + +from __future__ import annotations + +import asyncio +import logging + +import pytest + +from band.client.streaming import AgentControlPayload, ControlMode +from band.testing import fake_phoenix_server +from tests.platform.linkclose.kit import ( + HANG_GUARD_S, + agent_control, + connected, + control_cancellations, + make_link, +) + +LINK_LOGGER = "band.platform.link" + + +class TestDisconnectCancelsControl: + async def test_a_signal_queued_behind_a_running_one_never_applies( + self, caplog: pytest.LogCaptureFixture + ) -> None: + stop = agent_control(ControlMode.STOP, "stop-1") + interrupt = agent_control(ControlMode.INTERRUPT, "interrupt-1") + story: list[str] = [] + stop_entered = asyncio.Event() + + async def hook(payload: AgentControlPayload) -> None: + story.append(f"{payload.mode} applying") + stop_entered.set() + try: + await asyncio.Future() + except asyncio.CancelledError: + story.append(f"{payload.mode} cancelled") + raise + + async with fake_phoenix_server() as server: + link = make_link(server.url) + link.on_control = hook + async with connected(link): + await link._on_control(stop) + await asyncio.wait_for(stop_entered.wait(), HANG_GUARD_S) + await link._on_control(interrupt) + + with caplog.at_level(logging.INFO, logger=LINK_LOGGER): + await link.disconnect() + + assert story == ["stop applying", "stop cancelled"] + assert control_cancellations(caplog.records) == [ + ( + "Cancelling 2 unfinished control signal(s) on disconnect: " + "agent.control stop correlation_id=stop-1, " + "agent.control interrupt correlation_id=interrupt-1" + ) + ] + + async def test_an_applied_signal_leaves_nothing_to_cancel( + self, caplog: pytest.LogCaptureFixture + ) -> None: + applied = asyncio.Event() + + async def hook(payload: AgentControlPayload) -> None: + applied.set() + + async with fake_phoenix_server() as server: + link = make_link(server.url) + link.on_control = hook + async with connected(link): + await link._on_control(agent_control(ControlMode.PLAY, "play-1")) + await asyncio.wait_for(applied.wait(), HANG_GUARD_S) + + with caplog.at_level(logging.INFO, logger=LINK_LOGGER): + await link.disconnect() + + assert control_cancellations(caplog.records) == [] diff --git a/tests/platform/linkclose/test_supersede_during_control.py b/tests/platform/linkclose/test_supersede_during_control.py new file mode 100644 index 000000000..ff74a12b9 --- /dev/null +++ b/tests/platform/linkclose/test_supersede_during_control.py @@ -0,0 +1,62 @@ +"""A ``supersede`` queued behind a still-running control signal on the shared +``agent_control`` topic is handled at once, before the 1001 close can +reconnect.""" + +from __future__ import annotations + +import asyncio + +from websockets.frames import CloseCode + +from band.client.streaming.client import AgentControlPayload, ControlMode, WireEvent +from band.platform.event import WebSocketDisconnectedEvent +from band.testing import fake_phoenix_server +from tests.conftest import TERMINAL_SUPERSEDE +from tests.platform.linkclose.kit import ( + AGENT_ID, + CONTROL_TOPIC, + HANG_GUARD_S, + SUPERSEDED, + connected, + make_link, + running, +) + +STOP_CONTROL = { + "mode": ControlMode.STOP, + "scope": "agent", + "agent_id": AGENT_ID, + "execution_id": None, + "room_id": None, + "correlation_id": "control-1", +} + + +class TestSupersedeDuringControl: + async def test_a_supersede_is_handled_while_a_control_handler_runs( + self, + ) -> None: + """A STOP can take seconds; the supersede behind it must not wait.""" + async with fake_phoenix_server() as server: + link = make_link(server.url) + control_entered = asyncio.Event() + release_control = asyncio.Event() + + async def held_control(payload: AgentControlPayload) -> None: + control_entered.set() + await release_control.wait() + + link.on_control = held_control + + async with connected(link), running(link) as run: + await server.push(CONTROL_TOPIC, WireEvent.AGENT_CONTROL, STOP_CONTROL) + await asyncio.wait_for(control_entered.wait(), HANG_GUARD_S) + await server.push( + CONTROL_TOPIC, WireEvent.SUPERSEDE, TERMINAL_SUPERSEDE + ) + event = await asyncio.wait_for(anext(link), HANG_GUARD_S) + await server.close_connection(code=CloseCode.GOING_AWAY) + await asyncio.wait_for(run, HANG_GUARD_S) + + assert event == WebSocketDisconnectedEvent(payload=SUPERSEDED) + assert server.connection_count == 1 diff --git a/tests/platform/test_link_control.py b/tests/platform/test_link_control.py index 3c5c7df0e..50116c861 100644 --- a/tests/platform/test_link_control.py +++ b/tests/platform/test_link_control.py @@ -1,35 +1,29 @@ """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). +Control signals must be delivered preemptively — 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 +import asyncio from unittest.mock import AsyncMock -import pytest - import band.platform.link as link_mod 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] = [] + received: asyncio.Future[AgentControlPayload] = ( + asyncio.get_running_loop().create_future() + ) async def hook(payload: AgentControlPayload) -> None: - received.append(payload) + received.set_result(payload) link.on_control = hook payload = AgentControlPayload( @@ -38,14 +32,20 @@ async def hook(payload: AgentControlPayload) -> None: await link._on_control(payload) - assert received == [payload] + assert await asyncio.wait_for(received, timeout=5.0) == payload async def test_on_control_does_not_enqueue(self, link: BandLink): """Control must bypass the serialized event queue entirely.""" - link.on_control = AsyncMock() + applied = asyncio.Event() + + async def hook(payload: AgentControlPayload) -> None: + applied.set() + + link.on_control = hook payload = AgentControlPayload(mode="stop", scope="agent", agent_id="agent-123") await link._on_control(payload) + await asyncio.wait_for(applied.wait(), timeout=5.0) assert link._event_queue.empty() @@ -57,12 +57,8 @@ async def test_on_control_without_hook_is_safe_noop(self, link: BandLink): assert link._event_queue.empty() - async def test_connect_wires_on_control_to_channel(self): + async def test_connect_wires_on_control_to_channel(self, link: BandLink): """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() diff --git a/tests/platform/test_link_control_lane.py b/tests/platform/test_link_control_lane.py new file mode 100644 index 000000000..66bab0d67 --- /dev/null +++ b/tests/platform/test_link_control_lane.py @@ -0,0 +1,67 @@ +"""BandLink applies control signals one at a time, in arrival order, off the +``agent_control`` topic processor; one failed signal never blocks the next.""" + +from __future__ import annotations + +import asyncio +import logging + +import pytest + +from band.client.streaming import AgentControlPayload +from band.platform.link import BandLink + + +class TestControlLane: + async def test_a_signal_waits_for_the_one_before_it(self, link: BandLink): + """A PLAY behind a STOP that is still draining applies after it.""" + stop = AgentControlPayload(mode="stop", scope="agent", agent_id="agent-123") + play = AgentControlPayload(mode="play", scope="agent", agent_id="agent-123") + applied: list[AgentControlPayload] = [] + stop_entered = asyncio.Event() + release_stop = asyncio.Event() + play_applied = asyncio.Event() + + async def hook(payload: AgentControlPayload) -> None: + if payload is stop: + stop_entered.set() + await release_stop.wait() + applied.append(payload) + if payload is play: + play_applied.set() + + link.on_control = hook + await link._on_control(stop) + await asyncio.wait_for(stop_entered.wait(), timeout=5.0) + await link._on_control(play) + release_stop.set() + await asyncio.wait_for(play_applied.wait(), timeout=5.0) + + assert applied == [stop, play] + + async def test_a_failed_signal_is_logged_and_the_next_one_still_applies( + self, link: BandLink, caplog: pytest.LogCaptureFixture + ): + stop = AgentControlPayload( + mode="stop", scope="agent", agent_id="agent-123", correlation_id="stop-1" + ) + play = AgentControlPayload(mode="play", scope="agent", agent_id="agent-123") + play_applied = asyncio.Event() + + async def hook(payload: AgentControlPayload) -> None: + if payload is stop: + raise RuntimeError("adapter cancel failed") + play_applied.set() + + link.on_control = hook + with caplog.at_level(logging.ERROR, logger="band.platform.link"): + await link._on_control(stop) + await link._on_control(play) + await asyncio.wait_for(play_applied.wait(), timeout=5.0) + + assert [(r.levelno, r.getMessage()) for r in caplog.records] == [ + ( + logging.ERROR, + "agent.control stop failed (correlation_id=stop-1 room_id=None)", + ) + ] diff --git a/tests/scenarios/script.py b/tests/scenarios/script.py index a9a89934d..77a8e8027 100644 --- a/tests/scenarios/script.py +++ b/tests/scenarios/script.py @@ -23,6 +23,8 @@ def __init__(self, observations: Observations) -> None: self._held_turns: dict[str, asyncio.Event] = {} self._cleanup_release = asyncio.Event() self._cleanup_release.set() + self._control_release = asyncio.Event() + self._control_release.set() def hold_turn(self, text: str) -> None: self._held_turns[text] = asyncio.Event() @@ -36,12 +38,20 @@ def hold_cleanups(self) -> None: def release_cleanups(self) -> None: self._cleanup_release.set() + def hold_controls(self) -> None: + """Hold the adapter's control hook, as a backend slow to cancel does.""" + self._control_release.clear() + + def release_controls(self) -> None: + self._control_release.set() + def release_everything(self) -> None: for held in self._held_turns.values(): held.set() self._held_turns.clear() self.failing_cleanups.clear() self._cleanup_release.set() + self._control_release.set() async def on_execute(self, ctx: ExecutionContext, event: PlatformEvent) -> None: if not isinstance(event, MessageEvent) or event.payload is None: @@ -58,6 +68,11 @@ async def on_execute(self, ctx: ExecutionContext, event: PlatformEvent) -> None: async def on_control(self, room_id: str, mode: ControlMode) -> None: """The adapter's own control hook, as the runtime notifies it.""" self.observations.record(f"{room_id}: control {mode.value}") + try: + await self._control_release.wait() + except asyncio.CancelledError: + self.observations.record(f"{room_id}: control {mode.value} cancelled") + raise async def cleanup(self, room_id: str) -> None: self.observations.record(f"{room_id}: cleanup started") diff --git a/tests/scenarios/test_session_end_during_control.py b/tests/scenarios/test_session_end_during_control.py new file mode 100644 index 000000000..c219baa3e --- /dev/null +++ b/tests/scenarios/test_session_end_during_control.py @@ -0,0 +1,70 @@ +"""The session ends while a control signal is still being applied: a +supersede or the agent's own stop must not wait for it.""" + +from __future__ import annotations + +from websockets.frames import CloseCode + +from band.client.streaming import ControlMode +from tests.scenarios.world import SESSION_SUBJECT, BandWorld + + +class TestSessionEndDuringControl: + async def test_a_supersede_behind_a_draining_stop_ends_the_session( + self, world: BandWorld + ) -> None: + """The incident order: an agent-wide STOP whose adapter is slow to + cancel, then the platform's supersede and its 1001 close.""" + await world.join("r1") + await world.join("r2") + world.script.hold_turn("a") + world.script.hold_turn("b") + a = await world.say("r1", "a") + b = await world.say("r2", "b") + await world.until_told("r1", "turn a") + await world.until_told("r2", "turn b") + world.script.hold_controls() + + async with world.running(): + await world.control(ControlMode.STOP, correlation_id="stop-all") + await world.until_told("r1", "control stop") + await world.supersede() + await world.until_told(SESSION_SUBJECT, "disconnected") + await world.close_connection(CloseCode.GOING_AWAY) + await world.until_told(SESSION_SUBJECT, "ended") + + world.script.release_controls() + await world.until_told("r2", "control stop") + await world.until_told("r2", "turn b cancelled") + + assert world.server.connection_count == 1 + assert world.story(SESSION_SUBJECT) == ["disconnected", "ended"] + assert world.story("r1") == [ + f"processing {a}", + "turn a", + "control stop", + "turn a cancelled", + ] + assert world.story("r2") == [ + f"processing {b}", + "turn b", + "control stop", + "turn b cancelled", + ] + + async def test_stopping_the_agent_cancels_a_control_still_applying( + self, world: BandWorld + ) -> None: + await world.join("r1") + world.script.hold_controls() + await world.control(ControlMode.STOP, room_id="r1", correlation_id="stop-1") + await world.until_told("r1", "control stop") + + await world.stop() + + assert world.story("r1") == [ + "control stop", + "cleanup started", + "cleanup done", + "control stop cancelled", + ] diff --git a/tests/scenarios/world.py b/tests/scenarios/world.py index 4a95a9e89..c190a2c79 100644 --- a/tests/scenarios/world.py +++ b/tests/scenarios/world.py @@ -3,17 +3,19 @@ from __future__ import annotations +import asyncio from collections.abc import AsyncIterator, Callable from contextlib import asynccontextmanager from typing import Any -from band.client.streaming import ControlMode +from band.client.streaming import ControlMode, WireEvent from band.platform.link import BandLink from band.runtime.execution import Execution, ExecutionContext from band.runtime.runtime import AgentRuntime from band.runtime.sessionconfig import SessionConfig from band.testing import fake_phoenix_server from band.testing.phoenix_server import FakePhoenixServer +from tests.conftest import TERMINAL_SUPERSEDE from tests.runtime.helpers import rest_client_over from tests.scenarios.observations import Observations from tests.scenarios.peer import AGENT_ID, BandPeer, room_json @@ -21,6 +23,8 @@ ROOMS_TOPIC = f"agent_rooms:{AGENT_ID}" CONTROL_TOPIC = f"agent_control:{AGENT_ID}" +# Journal subject for the agent's connection, beside the per-room stories. +SESSION_SUBJECT = "session" SESSION = SessionConfig(enable_context_hydration=False, enable_working_state=False) @@ -107,6 +111,30 @@ async def control( } await self.server.push(CONTROL_TOPIC, "agent.control", payload) + async def supersede(self) -> None: + """A newer process connects with the same key; the platform evicts + this one with a terminal supersede.""" + await self.server.push(CONTROL_TOPIC, WireEvent.SUPERSEDE, TERMINAL_SUPERSEDE) + + async def close_connection(self, code: int) -> None: + """The platform closes the socket with ``code``.""" + await self.server.close_connection(code=code) + + @asynccontextmanager + async def running(self) -> AsyncIterator[None]: + """The host blocked on ``run_forever``, journaled when it returns.""" + + async def run() -> None: + await self.runtime.link.run_forever(install_signal_handlers=False) + self.observations.record(f"{SESSION_SUBJECT}: ended") + + task = asyncio.create_task(run()) + try: + yield + finally: + task.cancel() + await asyncio.wait({task}) + async def drop_connection(self) -> None: """Sever the socket; returns once the agent has reconciled its rooms.""" listings = self.peer.listings @@ -163,6 +191,11 @@ async def band_world( on_control=script.on_control, ) link.on_control = runtime.handle_control + + async def disconnected() -> None: + observations.record(f"{SESSION_SUBJECT}: disconnected") + + runtime.presence.on_disconnected = disconnected await runtime.start() try: yield BandWorld(