From 2024f633a3303b04a571fa196c8b80195c8412ea Mon Sep 17 00:00:00 2001 From: Amit Gazal Date: Wed, 15 Jul 2026 15:02:44 +0300 Subject: [PATCH 1/5] wip --- INT-829-plan.md | 118 ++++++++ MANUAL_TEST_INT829.md | 148 ++++++++++ src/band/client/streaming/__init__.py | 2 + src/band/client/streaming/client.py | 38 ++- src/band/platform/link.py | 25 ++ src/band/runtime/execution.py | 254 ++++++++++++++++- src/band/runtime/platform_runtime.py | 5 + src/band/runtime/runtime.py | 108 +++++++ tests/integration/test_control_signals.py | 287 +++++++++++++++++++ tests/platform/test_link.py | 1 + tests/platform/test_link_control.py | 84 ++++++ tests/runtime/test_execution_interrupt.py | 326 ++++++++++++++++++++++ tests/runtime/test_runtime_control.py | 206 ++++++++++++++ tests/test_platform_runtime.py | 16 ++ tests/websocket/test_control_payload.py | 83 ++++++ uv.lock | 2 +- 16 files changed, 1696 insertions(+), 7 deletions(-) create mode 100644 INT-829-plan.md create mode 100644 MANUAL_TEST_INT829.md create mode 100644 tests/integration/test_control_signals.py create mode 100644 tests/platform/test_link_control.py create mode 100644 tests/runtime/test_execution_interrupt.py create mode 100644 tests/runtime/test_runtime_control.py create mode 100644 tests/websocket/test_control_payload.py diff --git a/INT-829-plan.md b/INT-829-plan.md new file mode 100644 index 000000000..918679373 --- /dev/null +++ b/INT-829-plan.md @@ -0,0 +1,118 @@ + +# INT-829 — Remote agents honor interrupt / stop / play control signals (Python SDK) + +> **Status: FINAL — architect-reviewed, full consensus reached over Band.** +> SDK side only. Signal emission, dispatch gating, and stopped-flag persistence +> are platform-side (Eric's Ticket 01 / PLT-944). TDD: write tests first. + +## 1. Platform contract we build against (PLT-944 — fixed, not ours to change) + +- **Channel:** `agent_control:{agent_id}` (server→SDK; the SDK already joins it today for `supersede`). +- **Event:** `"agent.control"`, best-effort push. +- **Payload:** + ```json + { + "type": "agent.control", + "mode": "interrupt|stop|play", + "scope": "agent|room", + "agent_id": "", + "execution_id": "", + "room_id": "", + "reason": "user_requested", + "correlation_id": "ctl-" + } + ``` + `room_id` null = agent-scoped (fan-out across that agent's rooms). Server does **not** dedup → SDK dedups on `correlation_id`. +- **Semantics:** + - **interrupt** (transient): abort in-flight turn, drop response, back to listening. + - **stop** (durable): interrupt + stay silent. **Trigger suppression is platform-authoritative** — GET `/next`→204, POST mark→204, POST reply→403; stopped flag persists across reconnect. SDK need not persist stop state. WS push lands ~37–111ms *after* the platform set the durable flag. + - **play**: explicit resume; platform clears the gate and **replays queued triggers oldest-first** → SDK catches up rehydration-style via `/next`. +- **Replay dependency (regression-guard this):** stop→play replay works because `/next` (`Chat.get_next_actionable_message`) **excludes only `processed`** — a message left in `processing` is re-returned. If the platform ever changes `/next` to also exclude `processing`, stop→play silently drops the message. Note + assert this. + +## 2. Current SDK architecture (verified) + +- `agent_control:{agent_id}` joined in `BandLink.connect()` → `WebSocketClient.join_agent_control_channel(...)`, today only wiring `supersede` (`streaming/client.py:402`, `platform/link.py:146`). +- WS events parsed vs `_PAYLOAD_MODELS` then dispatched per-event in `WebSocketClient._handle_events` (`client.py:341`). The channel `message_handler` runs in the **PHX receive task**, separate from per-room execution loops — the natural preemption seam. +- Per-room reasoning runs in `ExecutionContext._process_loop` (`execution.py:855`); each message handled by `await self._on_execute(self, event)` **inline** in `_process_event` (`execution.py:1467`, `mark_processed` at `:1471`, `except Exception` at `:1487`, loop `except CancelledError` at `:914`) and `_process_backlog_message` (`execution.py:1269`, `mark_processed` at `:1272`, `except Exception` at `:1288`). +- Existing cancel primitive `ExecutionContext.stop(timeout)` cancels the **whole** loop task (`execution.py:463`) — too coarse for per-cycle interrupt. +- `request_resync()` (`execution.py:526`) already drives a `/next` catch-up — reuse for **play**. +- Rooms→executions owned by `AgentRuntime.executions`; wired by `PlatformRuntime`. `request_resync` degrades via `hasattr` for custom executions (`runtime.py:235`). + +## 3. Design (consensus) + +### 3a. Ingestion (WS layer) +- Add `AgentControlPayload(BaseModel, extra="allow")`: `mode`/`scope` as `Literal`s, `agent_id` required, `execution_id`/`room_id`/`reason`/`correlation_id` optional. +- Register `"agent.control": AgentControlPayload` in `_PAYLOAD_MODELS`. +- Add `on_control` callback param to `join_agent_control_channel`; handler dict → `{"supersede": on_supersede, "agent.control": on_control}`. + +### 3b. Routing — preemptive, out-of-band (Q-route: APPROVED) +- **Do NOT** push control onto `BandLink._event_queue` (serialized behind message processing — would defeat preemption). +- Add `BandLink.on_control` hook + `_on_control` handler, called **directly** from the receive task. +- `AgentRuntime.handle_control(payload)`: + - Dedup on `correlation_id` (bounded LRU). Distinct signals have distinct correlation_ids, so dedup never drops a `play` that follows a `stop`. + - Resolve targets: `scope=="agent"` & `room_id is None` → all `self.executions`; else the single room (unknown room → log + no-op). + - Dispatch by mode → `execution.interrupt()` / `execution.stop_room()` / `execution.resume_room()`, **guarded by `hasattr`** (custom executions degrade: log + skip). +- `PlatformRuntime` sets `link.on_control = runtime.handle_control`. + +### 3c. Per-cycle interrupt (Q1 APPROVED; Q2 RESOLVED; blockers a–d adopted) +- Run the cycle as a child task: `self._active_cycle_task = asyncio.create_task(self._on_execute(self, event))`, then `await` it. No platform `execution_id` matching — interrupt acts on whatever cycle is in flight for the room. If `payload.execution_id` is present, **log** it alongside `correlation_id` for platform-side debuggability. +- `interrupt(kind)` (receive-task side, **state surface = cancel + flags only**): set `self._interrupt_kind`, then `if task and not task.done(): task.cancel()`. Between-cycles (`None`/done) → no-op (**blocker c**). Do **not** touch `_processed_ids`/`self.queue` from the receive side (**blocker d**). +- asyncio cancellation drops in-flight tool `await`s (abandoned, not rolled back) — satisfies the tool-call AC for free. +- **CancelledError routing (blocker a — THE fix that makes it hold):** `CancelledError` is a `BaseException` subclass, so it bypasses the existing `except Exception`, propagates past `_process_event`/`_process_backlog_message`, and hits the loop's `except CancelledError` at `:914` → **kills the room loop permanently.** Add an explicit `except asyncio.CancelledError` **tightly around `await self._active_cycle_task`** in BOTH `_process_event` and `_process_backlog_message`: + - if `self._interrupt_kind is not None` (interrupt/stop): swallow; handle per mode below; `return True`. **Do not re-raise.** + - else (real shutdown cancel of the loop task propagating through): re-raise so `:914` still exits for shutdown. +- **No `asyncio.shield` needed** for the interrupt mark: only the *child* task is cancelled, so the loop task's cancel-state stays clean and a subsequent inline `await mark_processed` runs normally. (shield would only matter if the loop task itself were the cancel target = shutdown, where we don't mark anyway.) +- **Shutdown orphan (blocker b):** `ExecutionContext.stop()` cancels the loop task but **not** the awaited child → orphaned task. `stop()` must explicitly cancel (and await) `_active_cycle_task` too. + +### 3d. Message status on cancel (Q2 RESOLVED) +- **interrupt** (transient): in the loop's `except CancelledError` branch → `await mark_processed` + `_remember_processed` + release claim. **Consumes** the message so the Phase-2 idle `/next` (`idle_resync_seconds`) doesn't re-return it (excludes-only-processed) and re-fire the killed cycle. Send nothing. +- **stop**: leave the message in `processing` (do **not** mark, do **not** add to `_processed_ids`), release the local in-flight claim, set `_stopped`. Platform is already stopped (push lags the durable flag), and `/next` replays the `processing` message on **play**. + +### 3e. stop / play (Q3 APPROVED — minimal, efficiency-only) +- `stop_room()`: interrupt the in-flight cycle (kind=stop) + set local `_stopped=True`. `_stopped` is a **pure efficiency cache**, not an authoritative suppression decision: (1) pause Phase-2 idle `/next` polling so a stopped room isn't hammering `/next`→204; (2) short-circuit a WS trigger for that room to avoid mark→204/reply→403 churn. Nothing else. **Not persisted** across reconnect (platform gate keeps the room quiet via `/next`→204 for free). +- `resume_room()` (**play**): set `_stopped=False`, then `request_resync()` (existing) → `/next` replays the backlog (the stop-interrupted message + anything received while stopped) = rehydration-style catch-up. +- Stale `_stopped=True` only costs a quiet room until the next `play`, which always `request_resync`s — bounded divergence. + +### 3f. Activity-state clearing (Q4 — seam only) +- `interrupt()` calls an optional `self._on_activity_clear` hook (default `None`/no-op). Real clearing lands with the activity-signal ticket; we only expose the seam here. + +### 3g. Config (Q5 — always-on, no config) +- No `AgentControlConfig` (YAGNI). Control handling is core behavior; the channel is already joined. Add an observability hook later only if needed. + +### 3h. Protocol & custom executions (Q6) +- Add `interrupt()/stop_room()/resume_room()` to the `Execution` Protocol (typed conformance, as `request_resync` was added at 0.2.0). `AgentRuntime.handle_control` `hasattr`-guards + logs-and-skips. `LettaExecution` gets no-op stubs. + +## 4. Files to touch +- `src/band/client/streaming/client.py` — `AgentControlPayload` + register + `on_control` param. +- `src/band/platform/event.py` — `AgentControlEvent`/typing if needed. +- `src/band/platform/link.py` — `on_control` hook + `_on_control`. +- `src/band/runtime/execution.py` — child-task cycle; `interrupt/stop_room/resume_room`; `except CancelledError` routing in `_process_event` + `_process_backlog_message`; shutdown orphan cancel in `stop()`; `_stopped` gating in `_process_loop`; `_on_activity_clear` seam; Protocol additions. +- `src/band/runtime/runtime.py` — `handle_control` routing (dedup, target resolution, hasattr-degrade) + wiring. +- `src/band/runtime/platform_runtime.py` — set `link.on_control`. +- `src/band/integrations/.../letta` (or wherever `LettaExecution` lives) — no-op stubs. + +## 5. TDD test plan (write tests first) +1. **Payload** (`tests/client/streaming/`): parse all 3 modes; null `room_id`/`execution_id`; extra fields allowed; bad `mode` rejected. +2. **Link** (`tests/platform/`): `_on_control` calls the registered hook with the parsed payload; control is **not** enqueued on `_event_queue`. +3. **Interrupt** (`tests/runtime/test_execution_interrupt.py`): + - in-flight cycle cancelled; fake tools assert **nothing sent**; `mark_processed` + `_remember_processed` called. + - interrupt during an awaiting tool → tool result **not** sent. + - **REGRESSION: loop still alive** — after an interrupt, a fresh message in the same room is processed normally (guards the `CancelledError`-routing bug at `:914`). + - interrupt between cycles (no active task) → safe no-op. +4. **Stop/Play**: + - stop interrupts + sets `_stopped`; Phase-2 idle polling paused; message left in `processing` (not marked, not in `_processed_ids`). + - reconnect while stopped: `/next`→204 ⇒ no processing/sends (no SDK persistence needed). + - play clears `_stopped` + `request_resync` ⇒ backlog (incl. the stop-interrupted message) replays ⇒ cycle runs + responds. + - **stop→play replay platform dependency**: assert/document that replay relies on `/next` excluding only `processed`. +5. **Shutdown** : `ExecutionContext.stop()` with an in-flight cycle cancels + awaits the child (no orphaned task warning); returns graceful/false correctly. +6. **Routing** (`tests/runtime/test_runtime_control.py`): agent-scope null room → all rooms; room-scope → one room; unknown room → no-op; `correlation_id` dedup; play-before-stop ordering harmless. +7. **Conformance**: custom `Execution` without the new methods degrades gracefully (log + skip). +8. Manual/E2E note: real control via platform REST (human key) — likely manual verification. + +## 6. Resolved questions (consensus) +- **Q1**: interrupt current in-flight cycle; no `execution_id` matching; log it + `correlation_id`. +- **Q2**: interrupt **consumes** (`mark_processed`, inline awaited in loop except, no shield); stop **leaves `processing`** + `_stopped` for replay-on-play. Plus the explicit `except CancelledError` routing fix. +- **Q3**: minimal `_stopped` efficiency-only flag; not persisted. +- **Q4**: `_on_activity_clear` seam only. +- **Q5**: always-on; no config class. +- **Q6**: Protocol methods + `hasattr`-degrade + Letta no-op stubs. diff --git a/MANUAL_TEST_INT829.md b/MANUAL_TEST_INT829.md new file mode 100644 index 000000000..1c4d30bb8 --- /dev/null +++ b/MANUAL_TEST_INT829.md @@ -0,0 +1,148 @@ +# INT-829 manual test plan — interrupt / stop / play + +Verified against `https://app.band.ai` on 2026-06-15: PLT-944 (platform side) is +deployed, the three control routes are live, and auth uses the `X-API-Key` +header (NOT `Authorization: Bearer`). + +## Fixtures (verified) + +| Thing | Value | +|---|---| +| Base URL | `https://app.band.ai` | +| Auth header | `X-API-Key: $BAND_API_KEY_USER` | +| Chat (you are **owner**; Tom+Jerry are members) | `430fb827-9d48-476d-8b9f-ec450633c47c` | +| Tom agent_id | `3d5bd75e-6503-40e6-ac49-5acae495e880` | +| Jerry agent_id | `d9378533-dc52-4826-8417-aa9f145f0e1a` | + +Control routes (confirmed live): +- Room-scope (room owner; fans out to participating agents): `POST /api/v1/me/chats/{chat_id}/agents/{stop,play}` +- Agent-scope (agent owner; targets one execution): `POST /api/v1/me/agents/{agent_id}/executions/{execution_id}/{stop,play,interrupt}` +- Executions list: `GET /api/v1/me/agents/{agent_id}/executions` → items have `id`, `status`, `stopped_at`. + +Tom & Jerry are already running against the worktree build. Tail their logs: +```bash +tail -f /tmp/tom.log # and in another pane: +tail -f /tmp/jerry.log +``` + +## 0. Setup +```bash +export BASE="https://app.band.ai" +export BAND_API_KEY_USER="" +export CHAT="430fb827-9d48-476d-8b9f-ec450633c47c" +export TOM="3d5bd75e-6503-40e6-ac49-5acae495e880" +export JERRY="d9378533-dc52-4826-8417-aa9f145f0e1a" + +# sanity: should print your user record +curl -s "$BASE/api/v1/me" -H "X-API-Key: $BAND_API_KEY_USER" | python3 -m json.tool | head +``` + +## Helper: list executions (find the in-flight one) +```bash +curl -s "$BASE/api/v1/me/agents/$TOM/executions" -H "X-API-Key: $BAND_API_KEY_USER" \ + | python3 -c "import sys,json;[print(e['status'].ljust(12), e['id'], 'stopped_at='+str(e['stopped_at'])) for e in json.load(sys.stdin)['data']]" +``` +Idle executions show `status: waiting`. A cycle in flight shows a non-waiting, +freshly-`updated_at` status — that's the `execution_id` to interrupt. + +## Trigger a reasoning cycle (needed for interrupt) +Send a mention from you (the user) so the agent starts a turn. Use a prompt that +takes a few seconds so you have a window to interrupt: +```bash +# NOTE (verified live): body is nested under "message", and mentions require the +# agent UUID "id" (handle is rejected). +curl -s -XPOST "$BASE/api/v1/me/chats/$CHAT/messages" \ + -H "X-API-Key: $BAND_API_KEY_USER" -H "Content-Type: application/json" \ + -d "{\"message\":{\"content\":\"@Tom write a long, detailed multi-paragraph plan to catch Jerry, step by step\",\"mentions\":[{\"id\":\"$TOM\"}]}}" +``` +(Or just type the mention in the Band UI. Letting Tom & Jerry mention each other +creates the runaway loop that stop/play is designed to kill.) + +--- + +## TEST 1 — INTERRUPT (transient: kill one in-flight turn) +Interrupt is **agent-scope only** (no room-scope variant; no UI affordance yet). + +1. Trigger a cycle (above). +2. Immediately grab the in-flight `execution_id` with the helper. +3. Fire interrupt: +```bash +export EXEC="" +curl -s -w "\n[HTTP %{http_code}]\n" -XPOST \ + "$BASE/api/v1/me/agents/$TOM/executions/$EXEC/interrupt" \ + -H "X-API-Key: $BAND_API_KEY_USER" +``` +**PASS criteria** +- `/tmp/tom.log`: `interrupt requested, cancelling in-flight cycle` then + `cycle interrupted (message …) — nothing sent`. +- Tom posts **no** reply for that turn (partial work dropped). +- Tom answers the **next** mention normally (back to listening) — send another + mention to confirm. + +## TEST 2 — STOP (durable: go quiet until play) +Room-scope (stops every agent in the room) — simplest since you own the chat: +```bash +curl -s -w "\n[HTTP %{http_code}]\n" -XPOST \ + "$BASE/api/v1/me/chats/$CHAT/agents/stop" -H "X-API-Key: $BAND_API_KEY_USER" +``` +(Agent-scope variant — stop just Tom's execution: +`POST $BASE/api/v1/me/agents/$TOM/executions/$EXEC/stop`.) + +**PASS criteria** +- Executions helper now shows Tom with `stopped_at` non-null. +- `/tmp/tom.log` (if a cycle was running): `stop requested, cancelling in-flight cycle`. +- Send a new mention to Tom → **no response**, and the log shows + `stopped, skipping message … (left for replay)` and periodically + `stopped, skipping idle /next poll`. +- Jerry unaffected unless you stopped room-scope (which stops both). + +## TEST 3 — PLAY (resume + catch up on what was missed) +```bash +curl -s -w "\n[HTTP %{http_code}]\n" -XPOST \ + "$BASE/api/v1/me/chats/$CHAT/agents/play" -H "X-API-Key: $BAND_API_KEY_USER" +``` +**PASS criteria** +- Executions helper shows Tom `stopped_at` back to `null`. +- `/tmp/tom.log`: `Applying control mode=play …` → `Resync sentinel enqueued` → + `Catching up missed message …`. +- The mention(s) you sent **while stopped** now get answered (rehydration catch-up, + not dropped). + +## TEST 4 — STOP survives reconnect +1. Stop Tom (Test 2). Confirm `stopped_at` set and Tom silent. +2. Restart Tom's process (simulates a reconnect/redeploy): + ```bash + pkill -f 03_tom_agent.py + ( cd /Users/amitgazal/Workspace/thenvoi-sdk-python/.claude/worktrees/gleaming-inventing-twilight \ + && uv run python examples/anthropic/03_tom_agent.py > /tmp/tom.log 2>&1 & ) + ``` +3. Send Tom a mention → **still no response** (platform `/next`→204 keeps it quiet; + the SDK keeps nothing locally). The recovery sweep is also skipped — log will + NOT show the stuck message being re-processed. +4. `play` (Test 3) → Tom resumes and answers the backlog. + +## TEST 5 — dedup / fan-out (optional) +- Re-POST the same control twice quickly. The platform sends one push per call; + if the same `correlation_id` is delivered twice, `/tmp/tom.log` shows + `Duplicate control signal … ignored` on the repeat. +- Room-scope stop with both agents present → both Tom and Jerry go quiet + (`Applying control mode=stop scope=… to N room(s)` in each log). + +## Log line cheat-sheet (what the new SDK code emits) +| Event | Logger / line | +|---|---| +| any control arrives | `band.runtime.runtime`: `Applying control mode=… scope=… to N room(s) (correlation_id=… execution_id=…)` | +| interrupt | `band.runtime.execution`: `interrupt requested, cancelling in-flight cycle` → `cycle interrupted (message …) — nothing sent` | +| stop | `stop requested, cancelling in-flight cycle` → `cycle stopped (…)` ; then `stopped, skipping message … (left for replay)` / `stopped, skipping idle /next poll` | +| play | `Applying control mode=play …` → `Resync sentinel enqueued` → `Catching up missed message …` | +| duplicate | `Duplicate control signal … ignored` | +| unknown room | `Control signal (mode=…) for unknown room …; no-op` | + +## Cleanup +```bash +pkill -f 03_tom_agent.py ; pkill -f 04_jerry_agent.py +# if you stopped them on the platform, run play first so they aren't left parked. +``` + +> Note: Tom & Jerry run against **production** and spend Anthropic tokens while +> chatting. Stop the processes when done. diff --git a/src/band/client/streaming/__init__.py b/src/band/client/streaming/__init__.py index 5fde7380f..7af39088e 100644 --- a/src/band/client/streaming/__init__.py +++ b/src/band/client/streaming/__init__.py @@ -22,6 +22,7 @@ ContactAddedPayload, ContactRemovedPayload, SupersedePayload, + AgentControlPayload, ) from band.client.streaming.errors import WebSocketUpgradeError @@ -42,4 +43,5 @@ "ContactAddedPayload", "ContactRemovedPayload", "SupersedePayload", + "AgentControlPayload", ] diff --git a/src/band/client/streaming/client.py b/src/band/client/streaming/client.py index 06352f5c2..2d5c2cc81 100644 --- a/src/band/client/streaming/client.py +++ b/src/band/client/streaming/client.py @@ -5,6 +5,7 @@ from dataclasses import dataclass import logging import random +from typing import Literal from phoenix_channels_python_client.client import ( PHXChannelsClient, @@ -205,6 +206,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.""" @@ -240,6 +262,7 @@ def to_disconnect_reason(self) -> WebSocketDisconnectReason: "contact_added": ContactAddedPayload, "contact_removed": ContactRemovedPayload, "supersede": SupersedePayload, + "agent.control": AgentControlPayload, } @@ -403,13 +426,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 b362e9ca5..7654aa7d8 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 + @property def is_connected(self) -> bool: return self._is_connected @@ -146,6 +154,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) @@ -327,6 +336,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 337acd754..9ac5723e0 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 collections import OrderedDict from datetime import datetime, timezone @@ -149,6 +150,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]] @@ -258,6 +283,26 @@ 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 + + # 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.""" @@ -458,6 +503,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() @@ -466,6 +519,12 @@ async def stop(self, timeout: float | None = None) -> bool: except asyncio.CancelledError: 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 return graceful async def _wait_for_idle(self, timeout: float) -> bool: @@ -534,6 +593,71 @@ 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 a cycle was actually in flight and got cancelled. Between + cycles this is a clean no-op and does NOT set ``_interrupt_kind`` + (which would otherwise mis-flag the next cycle). + """ + task = self._active_cycle_task + if task is None or task.done(): + return False + self._interrupt_kind = kind + task.cancel() + logger.info( + "ExecutionContext %s: %s requested, cancelling in-flight cycle", + self.room_id, + kind, + ) + return True + + 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: @@ -892,6 +1016,13 @@ async def _process_loop(self) -> None: timeout=self.config.idle_resync_seconds, ) 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, @@ -901,6 +1032,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, @@ -1028,7 +1168,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 @@ -1265,8 +1418,11 @@ async def _process_backlog_message( ), ) - # Call execution handler - 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 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) @@ -1328,6 +1484,80 @@ 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``).""" + await self._on_execute(self, event) + + 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. + """ + 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 + # Interrupt/stop: drop all accumulated work, send nothing. + 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._retry_tracker.mark_success(msg_id) + self._remember_processed_message(msg_id) + else: + logger.warning( + "ExecutionContext %s: durable mark_processed failed for " + "interrupted message %s; idle resync may re-deliver it", + 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 + finally: + self._active_cycle_task = None + async def _process_event(self, event: PlatformEvent) -> bool: """ Process single event through execution callback. @@ -1353,6 +1583,19 @@ async def _process_event(self, event: PlatformEvent) -> bool: self._set_state("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 @@ -1463,8 +1706,11 @@ async def _process_event(self, event: PlatformEvent) -> bool: self.remove_participant(event.payload.id) await self._notify_participant_removed(event) - # Call execution handler - 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. + 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: diff --git a/src/band/runtime/platform_runtime.py b/src/band/runtime/platform_runtime.py index 9c9e3eb46..d0fe6f829 100644 --- a/src/band/runtime/platform_runtime.py +++ b/src/band/runtime/platform_runtime.py @@ -166,6 +166,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/runtime.py b/src/band/runtime/runtime.py index 04c5a6383..603a0b6a9 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 @@ -21,6 +22,7 @@ ) if TYPE_CHECKING: + from band.client.streaming import AgentControlPayload from band.platform.link import BandLink logger = logging.getLogger(__name__) @@ -131,6 +133,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 + # Set up presence callbacks self.presence.on_room_joined = self._on_room_joined self.presence.on_room_left = self._on_room_left @@ -243,6 +251,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/integration/test_control_signals.py b/tests/integration/test_control_signals.py new file mode 100644 index 000000000..5f5210713 --- /dev/null +++ b/tests/integration/test_control_signals.py @@ -0,0 +1,287 @@ +"""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 time +from collections.abc import Callable + +import httpx +import pytest +import pytest_asyncio +from thenvoi_rest.human_api_chats.types.create_my_chat_room_request_chat import ( + CreateMyChatRoomRequestChat, +) +from thenvoi_rest.types import ParticipantRequest + +from band.platform.link import BandLink +from band.runtime.runtime import AgentRuntime +from tests.integration.conftest import ( + get_api_key, + get_base_url, + get_user_api_key, + get_ws_url, + is_room_alive, + requires_api, + requires_user_api, +) + +logger = logging.getLogger(__name__) + +_API = "/api/v1" + + +# --------------------------------------------------------------------------- # +# Raw HTTP helpers (user key) for endpoints not yet in the generated 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(chat_id: str, agent_id: str, text: str) -> str: + """Post a message from the user that @mentions the agent; return message id.""" + body = {"message": {"content": text, "mentions": [{"id": agent_id}]}} + async with httpx.AsyncClient(timeout=30.0) as client: + resp = await client.post( + f"{get_base_url()}{_API}/me/chats/{chat_id}/messages", + headers=_user_headers(), + json=body, + ) + assert resp.status_code in (200, 201), ( + f"send failed: {resp.status_code} {resp.text}" + ) + return resp.json()["data"]["id"] + + +async def _wait_until( + pred: Callable[[], bool], *, timeout: float, interval: float = 1.0 +) -> bool: + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if pred(): + return True + await asyncio.sleep(interval) + return pred() + + +class _ControllableHandler: + """Stand-in for an adapter's on_execute with deterministic timing. + + Records which message ids it fully processed, signals when a cycle starts, + and (when ``block`` is set) hangs until cancelled so an interrupt can land + mid-cycle. + """ + + def __init__(self) -> None: + self.block = False + self.started = asyncio.Event() + self.cancelled = asyncio.Event() + self.completed: list[str] = [] + + async def __call__(self, ctx, event) -> None: + payload = getattr(event, "payload", None) + msg_id = getattr(payload, "id", None) + self.started.set() + try: + if self.block: + await asyncio.sleep(120) # hang until interrupted + if msg_id: + self.completed.append(msg_id) + except asyncio.CancelledError: + self.cancelled.set() + raise + + +# --------------------------------------------------------------------------- # +# Fixtures +# --------------------------------------------------------------------------- # + + +@pytest_asyncio.fixture +async def control_room(user_api_client, session_api_client, shared_agent1_info): + """A **user-owned** chat with the agent as a participant (reused across runs). + + User ownership is required so the user may issue room-scope stop/play. The + reuse-or-create logic keeps this within the platform's per-agent room limit + even though the fixture is function-scoped. + """ + 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 user-owned room that already has the agent (10-room limit). + chats = await user_api_client.human_api_chats.list_my_chats() + for room in reversed(chats.data or []): + 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 control room: %s", room.id) + return room.id + + created = await user_api_client.human_api_chats.create_my_chat_room( + chat=CreateMyChatRoomRequestChat() + ) + 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 control room: %s", chat_id) + return chat_id + + +@pytest_asyncio.fixture +async def control_runtime(control_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 control_room is None or shared_agent1_info is None: + pytest.skip("control_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(), + ) + handler = _ControllableHandler() + 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, control_room + finally: + try: + await _post(f"{_API}/me/chats/{control_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 +class TestControlSignalsIntegration: + async def test_stop_suppresses_then_play_replays(self, control_runtime): + """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( + 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): + """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( + 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 be1166b7b..eaa32b8f6 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..1332eaeea --- /dev/null +++ b/tests/runtime/test_execution_interrupt.py @@ -0,0 +1,326 @@ +"""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 +from tests.conftest import 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.""" + started = asyncio.Event() + sent: list[str] = [] + + async def on_execute(ctx, event): + started.set() + try: + await asyncio.sleep(60) # simulate a long reasoning cycle + sent.append(event.payload.id) # would "send" — must not happen + except asyncio.CancelledError: + sent.append("CANCELLED") + raise + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + + event = make_message_event(msg_id="msg-1") + proc = asyncio.create_task(ctx._process_event(event)) + await 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 "msg-1" not in sent # 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._processed_ids + assert ctx._interrupt_kind is None # flag cleared + assert ctx._active_cycle_task is None + + # Loop still alive: a fresh message processes normally. + completed: list[str] = [] + + async def on_execute2(ctx, event): + completed.append(event.payload.id) + + ctx._on_execute = on_execute2 + result2 = await ctx._process_event(make_message_event(msg_id="msg-2")) + assert result2 is True + assert 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.""" + started = asyncio.Event() + + async def on_execute(ctx, event): + started.set() + await asyncio.sleep(60) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="s1"))) + await 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._processed_ids + # Local in-flight claim released so it can be reprocessed on play. + assert "s1" not in ctx._inflight_message_ids + + +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.""" + started = asyncio.Event() + + async def on_execute(ctx, event): + started.set() + await asyncio.sleep(60) + + ctx = ExecutionContext("room-123", mock_link, on_execute, 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 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.""" + started = asyncio.Event() + + async def on_execute(ctx, event): + started.set() + await asyncio.sleep(60) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="x"))) + await 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.""" + executed: list[str] = [] + + async def on_execute(ctx, event): + executed.append(event.payload.id) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + ctx._stopped = True + + result = await ctx._process_event(make_message_event(msg_id="while-stopped")) + + assert result is True + assert executed == [] # 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")] + ) + executed: list[str] = [] + + async def on_execute(ctx, event): + executed.append(event.payload.id) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + ctx._stopped = True + + ok = await ctx._recover_stale_processing_messages() + + assert ok is True + assert executed == [] # 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) + executed: list[str] = [] + + async def on_execute(ctx, event): + executed.append(event.payload.id) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + # Simulate: was stopped, now resuming. + ctx._stopped = False + ok = await ctx._resync_pending_messages() + + assert ok is True + assert executed == ["replayed-1"] + + +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).""" + started = asyncio.Event() + + async def on_execute(ctx, event): + started.set() + await asyncio.sleep(60) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + msg = _backlog_message("bk1") + proc = asyncio.create_task(ctx._process_backlog_message(msg)) + await 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._processed_ids + + async def test_stop_during_backlog_leaves_actionable(self, mock_link): + started = asyncio.Event() + + async def on_execute(ctx, event): + started.set() + await asyncio.sleep(60) + + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + proc = asyncio.create_task( + ctx._process_backlog_message(_backlog_message("bk2")) + ) + await 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._processed_ids diff --git a/tests/runtime/test_runtime_control.py b/tests/runtime/test_runtime_control.py new file mode 100644 index 000000000..1db01d478 --- /dev/null +++ b/tests/runtime/test_runtime_control.py @@ -0,0 +1,206 @@ +"""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() + + +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 b9b20f6a5..2fe3dd91a 100644 --- a/tests/test_platform_runtime.py +++ b/tests/test_platform_runtime.py @@ -292,6 +292,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 diff --git a/uv.lock b/uv.lock index 07546e263..bf8a55965 100644 --- a/uv.lock +++ b/uv.lock @@ -455,7 +455,7 @@ wheels = [ [[package]] name = "band-sdk" -version = "0.2.9" +version = "0.2.11" source = { editable = "." } dependencies = [ { name = "cryptography" }, From b784fa661d4c8e3f76726fd40b8cb32a4fe155f1 Mon Sep 17 00:00:00 2001 From: Amit Gazal Date: Thu, 16 Jul 2026 09:45:04 +0300 Subject: [PATCH 2/5] feat: bridge honors interrupt/stop control signals for one-shot agents MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wires agent_control into the bridge's AgentRunner: interrupt/stop cancel whatever forward task is currently in flight for a room, and play proactively nudges the room via /next instead of waiting for the next natural event. The bridge holds no Band lifecycle logic, so interrupt and stop are handled identically here — the consume-vs-replay distinction and durable per-room suppression already happen correctly downstream via OneShotInvoker's /next claim check and the platform's stop gate. --- band-bridge/bridge_core/bridge.py | 194 +++++++++++--- tests/bridge/test_bridge_control.py | 386 ++++++++++++++++++++++++++++ 2 files changed, 544 insertions(+), 36 deletions(-) create mode 100644 tests/bridge/test_bridge_control.py diff --git a/band-bridge/bridge_core/bridge.py b/band-bridge/bridge_core/bridge.py index 29f1968fd..3cffce45b 100644 --- a/band-bridge/bridge_core/bridge.py +++ b/band-bridge/bridge_core/bridge.py @@ -26,7 +26,7 @@ from typing import Any from band.client.rest import DEFAULT_REQUEST_OPTIONS -from band.client.streaming import MessageCreatedPayload +from band.client.streaming import AgentControlPayload, MessageCreatedPayload from band.runtime.types import PlatformMessage from band.platform.event import ( ContactAddedEvent, @@ -49,6 +49,7 @@ logger = logging.getLogger(__name__) _DEDUP_MAX_SIZE = 10_000 +_CONTROL_DEDUP_MAX_SIZE = 256 _FORWARDABLE_EVENTS: tuple[type, ...] = ( MessageEvent, @@ -105,6 +106,16 @@ def __init__( self._shutdown_event = shutdown_event self._connected_event = asyncio.Event() self._processed_message_ids: OrderedDict[str, None] = OrderedDict() + # In-flight forward task per room, so a control signal (interrupt/stop) + # can cancel whatever forward is currently running for a room. Set by + # _safe_handle_event at the top of processing, cleared in its finally. + self._active_room_tasks: dict[str, asyncio.Task[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). + self._seen_control_ids: OrderedDict[str, bool] = OrderedDict() # 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). @@ -236,6 +247,11 @@ async def _connect_and_consume(self) -> None: await self._link.connect() self._connected_event.set() + # Route control signals (interrupt/stop/play) to this runner. Set as + # soon as connected so the hook is live for the agent_control channel + # BandLink.connect() already joined. + self._link.on_control = self._handle_control + await self._link.subscribe_agent_rooms(self._config.agent_id) existing_rooms = await self._fetch_existing_rooms() @@ -331,31 +347,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 (:meth:`_handle_control`) — 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 +403,33 @@ 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: - 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, - ) + # Track this task as the in-flight forward for its room so a control + # signal (interrupt/stop) can cancel it. Room-less events (e.g. + # contact events) aren't tracked — there's nothing room-scoped to + # cancel for them. Guarded pop in `finally`: a newer event for the + # same room may already have replaced this entry. + room_id = getattr(event, "room_id", None) + task = asyncio.current_task() + if room_id and task is not None: + self._active_room_tasks[room_id] = 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: + try: + 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, + ) + finally: + if room_id and self._active_room_tasks.get(room_id) is task: + self._active_room_tasks.pop(room_id, None) async def _handle_event(self, event: object) -> None: """Manage room subscriptions; forward forwardable events. @@ -533,6 +567,94 @@ 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) --- + + async def _handle_control(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). + + 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. 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. + """ + cid = payload.correlation_id + if cid is not None: + if cid in self._seen_control_ids: + logger.debug( + "Agent %s: duplicate control signal %s ignored", + self._config.agent_id, + cid, + ) + return + 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) + else: + logger.debug( + "Agent %s: control signal mode=%s has no correlation_id; not deduped", + self._config.agent_id, + payload.mode, + ) + + if payload.room_id is not None: + room_ids = [payload.room_id] + elif payload.scope == "agent": + room_ids = await self._fetch_existing_rooms() + else: + logger.warning( + "Agent %s: control signal scope=%s with no room_id; no-op", + self._config.agent_id, + payload.scope, + ) + return + + logger.info( + "Agent %s: applying control mode=%s scope=%s to %d room(s) " + "(correlation_id=%s)", + self._config.agent_id, + payload.mode, + payload.scope, + len(room_ids), + cid, + ) + + if payload.mode in ("interrupt", "stop"): + for room_id in room_ids: + self._cancel_active_forward(room_id) + elif payload.mode == "play": + await asyncio.gather(*(self._nudge_room(rid) for rid in room_ids)) + + def _cancel_active_forward(self, room_id: str) -> None: + """Cancel the in-flight forward task for a room, if any. + + 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. + """ + task = self._active_room_tasks.get(room_id) + if task is None or task.done(): + return + task.cancel() + logger.info( + "Agent %s: cancelled in-flight forward for room %s", + self._config.agent_id, + room_id, + ) + class BandBridge: """Bridge orchestrator: N AgentRunners + signal handling + health server.""" diff --git a/tests/bridge/test_bridge_control.py b/tests/bridge/test_bridge_control.py new file mode 100644 index 000000000..22e07c6aa --- /dev/null +++ b/tests/bridge/test_bridge_control.py @@ -0,0 +1,386 @@ +"""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() + + +# --------------------------------------------------------------------------- +# 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._handle_control + + +# --------------------------------------------------------------------------- +# 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._handle_control(_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._handle_control(_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._handle_control(_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._handle_control(_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._handle_control(_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._handle_control(_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._handle_control(_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._handle_control( + 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._handle_control(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._handle_control(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._handle_control(sig) + await runner._handle_control(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._handle_control( + _control("interrupt", scope="room", room_id="r1", correlation_id="ctl-1") + ) + await runner._handle_control( + _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._handle_control(sig) + await runner._handle_control(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._handle_control(_control("play", scope="room", room_id="r1")) + + 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._handle_control(_control("play", scope="agent", room_id=None)) + + 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._handle_control(_control("play", scope="room", room_id="r1")) + + assert isinstance(fwd, FakeForwarder) + assert fwd.forwarded == [] + + +# --------------------------------------------------------------------------- +# 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._handle_control( + _control(mode, scope="room", room_id="ghost-room") + ) From 464482031d5c14fc23e85e7a09db1060f426981a Mon Sep 17 00:00:00 2001 From: Amit Gazal Date: Thu, 16 Jul 2026 10:01:19 +0300 Subject: [PATCH 3/5] fix: skip reconnect /next sync locally when room is stopped The ReconnectedEvent path in ExecutionContext._process_event always ran the full /next sync, unlike the idle-timeout and resync-sentinel paths which already short-circuit on _stopped. The platform gate made this safe (a stopped room's /next is guaranteed empty), but it wasted a call on every reconnect while stopped. Skip locally for consistency with the other _stopped short-circuits. --- src/band/runtime/execution.py | 16 +++++++++++++--- tests/runtime/test_runtime_control.py | 4 ++++ 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index 9ac5723e0..a849a302d 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -1575,9 +1575,19 @@ 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("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("idle") + await asyncio.sleep(self.config.idle_resync_seconds) logger.debug("Event %s processed successfully", event.type) finally: self._set_state("idle") diff --git a/tests/runtime/test_runtime_control.py b/tests/runtime/test_runtime_control.py index 1db01d478..b125325bc 100644 --- a/tests/runtime/test_runtime_control.py +++ b/tests/runtime/test_runtime_control.py @@ -189,6 +189,10 @@ async def on_execute(ctx, event): 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: From 3c13419491fc44c32a35131eb62fe39c56bd6eba Mon Sep 17 00:00:00 2001 From: Amit Gazal Date: Thu, 16 Jul 2026 16:11:32 +0300 Subject: [PATCH 4/5] fix: prevent stop from being undone by the next /next poll in backlog sync A stop control signal landing mid-cycle during _synchronize_with_next or _resync_pending_messages left the message in 'processing' for replay, but the enclosing loop kept polling /next, which returns 'processing' messages same as before. The very next iteration re-claimed and fully re-ran the cycle stop just aborted. Both loops now check self._stopped after processing a backlog message and stop polling until play resumes. Co-Authored-By: Claude Sonnet 5 --- src/band/runtime/execution.py | 25 ++++++++++ tests/runtime/test_execution_interrupt.py | 60 ++++++++++++++++++++++- 2 files changed, 84 insertions(+), 1 deletion(-) diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index a849a302d..347264260 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -1144,6 +1144,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", @@ -1266,6 +1280,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 diff --git a/tests/runtime/test_execution_interrupt.py b/tests/runtime/test_execution_interrupt.py index 1332eaeea..93660b4c5 100644 --- a/tests/runtime/test_execution_interrupt.py +++ b/tests/runtime/test_execution_interrupt.py @@ -14,7 +14,7 @@ import pytest from band.runtime.execution import ExecutionContext, _BacklogProcessResult -from band.runtime.types import PlatformMessage +from band.runtime.types import PlatformMessage, SessionConfig from tests.conftest import make_message_event @@ -324,3 +324,61 @@ async def on_execute(ctx, event): assert result == _BacklogProcessResult.ADVANCED mock_link.mark_processed.assert_not_awaited() assert "bk2" not in ctx._processed_ids + + 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) + + started = asyncio.Event() + calls: list[str] = [] + + async def on_execute(ctx, event): + calls.append(event.payload.id) + started.set() + await asyncio.sleep(60) # cancelled by stop_room() below + + ctx = ExecutionContext( + "room-123", + mock_link, + on_execute, + 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 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 calls == ["loop1"] + mock_link.mark_processed.assert_not_awaited() + assert "loop1" not in ctx._processed_ids From c718d9864fef80043a197ee503515b36cffea291 Mon Sep 17 00:00:00 2001 From: Amit Gazal Date: Thu, 16 Jul 2026 17:34:49 +0300 Subject: [PATCH 5/5] fix: don't charge stopped/interrupted cycles against the message retry budget A cycle cancelled by our own stop/interrupt control signal never actually ran the handler, but record_attempt() was still charging it as a real attempt. At the default max_message_retries=1, a message stopped once and replayed via /next on resume would immediately land in permanently_failed before ever reaching the adapter. Co-Authored-By: Claude Sonnet 5 --- src/band/runtime/execution.py | 10 +++++-- src/band/runtime/retry_tracker.py | 8 ++++++ tests/runtime/test_execution_interrupt.py | 35 +++++++++++++++++++++++ 3 files changed, 51 insertions(+), 2 deletions(-) diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index 347264260..2e3cdc151 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -1549,14 +1549,20 @@ async def _run_cycle(self, event: PlatformEvent, msg_id: str | None) -> bool: # Shutdown cancel of the loop task propagating through the child # await — let it propagate so the loop exits. raise - # Interrupt/stop: drop all accumulated work, send nothing. + # Interrupt/stop: drop all accumulated work, send nothing. 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._retry_tracker.mark_success(msg_id) self._remember_processed_message(msg_id) else: logger.warning( diff --git a/src/band/runtime/retry_tracker.py b/src/band/runtime/retry_tracker.py index d949d8031..6caaa6350 100644 --- a/src/band/runtime/retry_tracker.py +++ b/src/band/runtime/retry_tracker.py @@ -55,6 +55,14 @@ 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 an 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. + """ + 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/tests/runtime/test_execution_interrupt.py b/tests/runtime/test_execution_interrupt.py index 93660b4c5..7c45419f7 100644 --- a/tests/runtime/test_execution_interrupt.py +++ b/tests/runtime/test_execution_interrupt.py @@ -282,6 +282,41 @@ async def on_execute(ctx, event): assert ok is True assert executed == ["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.""" + started = asyncio.Event() + executed: list[str] = [] + calls = {"n": 0} + + async def on_execute(ctx, event): + calls["n"] += 1 + started.set() + if calls["n"] == 1: + await asyncio.sleep(60) # first attempt: hangs, gets stopped + executed.append(event.payload.id) + + # Real default (SessionConfig().max_message_retries == 1) — no + # generous override, so a poisoned attempt count would trip this. + ctx = ExecutionContext("room-123", mock_link, on_execute, agent_id="agent-123") + + proc = asyncio.create_task(ctx._process_event(make_message_event(msg_id="p1"))) + await 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. + started.clear() + result = await ctx._process_backlog_message(_backlog_message("p1")) + + assert result == _BacklogProcessResult.ADVANCED + assert executed == ["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):