diff --git a/docs/adapters/opencode.md b/docs/adapters/opencode.md index 3b397597e..c630af04e 100644 --- a/docs/adapters/opencode.md +++ b/docs/adapters/opencode.md @@ -36,7 +36,12 @@ since the ask carries its own `approval_wait_timeout_s` expiry. A manual ask also releases the turn's delivery early (it is marked processed while the human decides), so the adapter judges that turn itself when OpenCode finishes it: a turn that ends without a reply then posts one missing-reply -failure. A turn cancelled by cleanup or an interrupt posts nothing. +failure. A turn ended by cleanup, a STOP or an interrupt, even one already +released for approval, posts nothing: its watcher is cancelled first, then its +OpenCode session is aborted once through the adapter's `backend_turns`. A turn +that already ended on its own, by finishing or by timing out, still posts that +outcome when a STOP lands while it is being posted; the STOP only finishes the +abort. New ordinary work remains outstanding while the session is occupied. Approval, status and question replies remain reachable through the shared control hook; diff --git a/pyproject.toml b/pyproject.toml index 29beace4f..4f7f90a6d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -63,7 +63,10 @@ opencode = [ "mcp>=1.28.1,<2", ] letta = [ - "letta-client>=1.0.0", + # Floor is `conversations.cancel` (shared-mode STOP), first shipped in + # 1.7.1; the `conversations` resource itself arrived in 1.6.5 -- checked + # against the PyPI wheels 2026-10-19. + "letta-client>=1.7.1", # Floor pinned to crewai's own transitive pin (mcp~=1.28.1) -- see `desktop` above. "mcp>=1.28.1,<2", ] @@ -245,7 +248,7 @@ dev = [ # Include Agno for testing "agno>=2.6.0", # Include letta-client for testing - "letta-client>=1.0.0", + "letta-client>=1.7.1", # Include bridge deps for testing "aiohttp>=3.9,<4", "python-dotenv>=1.2.2", diff --git a/src/band/adapters/copilot_sdk.py b/src/band/adapters/copilot_sdk.py index 5b48037d5..5a3c8fde1 100644 --- a/src/band/adapters/copilot_sdk.py +++ b/src/band/adapters/copilot_sdk.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio import json import logging import re @@ -33,6 +34,7 @@ from band.core.protocols import GENERIC_PROVIDER_FAILURE_MESSAGE, send_event_safe from band.core.simple_adapter import SimpleAdapter from band.core.tool_filter import filter_tool_schemas +from band.core.turn_lifecycle import BackendTurns from band.core.types import Capability, Emit, MessageType, ToolEventKey, TurnUsage from band.integrations.copilot_sdk import CopilotSessionManager from band.integrations.copilot_sdk.prompts import TURN_COMPLETION_GUIDANCE @@ -72,6 +74,7 @@ AssistantTurnRetryData, AssistantUsageData, ModelCallFailureData, + SessionIdleData, SessionWarningData, ToolExecutionCompleteData, ) @@ -170,6 +173,16 @@ class TurnState: incidents: list[str] = field(default_factory=list) +@dataclass +class SessionTurn: + """The session a stop aborts, and whether its prompt is still running.""" + + session: CopilotSession + # Set when the prompt is sent, cleared once it returns; a prompt that is + # cancelled or fails leaves it set. + prompting: bool = False + + def _turn_incident(data: object) -> str | None: """A one-line account of a session event that can explain a silent turn.""" match data: @@ -319,6 +332,9 @@ def __init__( self._room_tools: dict[str, AgentToolsProtocol] = {} self._session_ids: dict[str, RoomSessionIds] = {} self._turn_state: dict[str, TurnState] = {} + self._backend_turns: BackendTurns[SessionTurn] = BackendTurns( + self._abort_turn, provider=_PROVIDER + ) self._system_prompt: str = "" # --- Lifecycle ------------------------------------------------------- @@ -395,6 +411,7 @@ async def on_cleanup(self, room_id: str) -> None: rooms); ``cleanup_all`` stops it at shutdown (unless it was borrowed via ``client=`` — then its owner stops it). """ + await self.backend_turns.settled(room_id) if self._session_manager: await self._session_manager.cleanup_session(room_id) self._room_tools.pop(room_id, None) @@ -403,6 +420,7 @@ async def on_cleanup(self, room_id: str) -> None: async def cleanup_all(self) -> None: """Stop all sessions and, if adapter-owned, the Copilot client.""" + await self.backend_turns.drain() if self._session_manager: await self._session_manager.stop() self._room_tools.clear() @@ -411,6 +429,10 @@ async def cleanup_all(self) -> None: # --- Message handling -------------------------------------------------- + @property + def backend_turns(self) -> BackendTurns[SessionTurn]: + return self._backend_turns + async def on_message( self, msg: PlatformMessage, @@ -433,52 +455,61 @@ async def on_message( # Same-session calls must not interleave; other rooms run concurrently. async with self._session_manager.turn_lock(room_id): - # Bound only once this message holds the room, so a queued - # message never redirects the running turn's tool calls. - self._room_tools[room_id] = tools - try: - session, inject_text = await self._obtain_session( - room_id, history, tools, is_session_bootstrap=is_session_bootstrap - ) - except Exception: - logger.exception("Room %s: Copilot session setup failed", room_id) - await tools.send_failure( - AgentFailure(_PROVIDER, GENERIC_PROVIDER_FAILURE_MESSAGE) + # Abort is what stops the runtime's work for a cancelled turn; the + # session stays valid for the next message. + async with self.backend_turns.running(room_id): + # Bound once the previous turn's abort has settled and this + # message holds the room, so neither a queued message nor a + # late call from the stopped turn reaches this turn's tools. + self._room_tools[room_id] = tools + try: + session, inject_text = await self._obtain_session( + room_id, + history, + tools, + is_session_bootstrap=is_session_bootstrap, + ) + except Exception: + logger.exception("Room %s: Copilot session setup failed", room_id) + await tools.send_failure( + AgentFailure(_PROVIDER, GENERIC_PROVIDER_FAILURE_MESSAGE) + ) + raise + live = SessionTurn(session) + self.backend_turns.bind(room_id, live) + prompt = self._compose_prompt( + msg, + participants_msg, + contacts_msg, + inject_text=inject_text, ) - raise - prompt = self._compose_prompt( - msg, - participants_msg, - contacts_msg, - inject_text=inject_text, - ) - # The session's ask_user handler reads this slot; safe because the - # turn lock serializes turns per room. - turn = TurnState( - sender_mention={ - "id": msg.sender_id, - "name": msg.sender_name or msg.sender_type, - } - ) - self._turn_state[room_id] = turn - try: - final_text = await self._run_turn(session, prompt, turn) - except Exception: - logger.exception("Room %s: Copilot turn failed", room_id) - # Abort any work the runtime is still doing for this turn and - # drop the session; the next message resumes it fresh by id. - await self._session_manager.evict_session(room_id) - await tools.send_failure( - AgentFailure(_PROVIDER, GENERIC_PROVIDER_FAILURE_MESSAGE) + # The session's ask_user handler reads this slot; safe because the + # turn lock serializes turns per room. + turn = TurnState( + sender_mention={ + "id": msg.sender_id, + "name": msg.sender_name or msg.sender_type, + } ) - raise - finally: - self._turn_state.pop(room_id, None) - # No-op unless Emit.USAGE is on / usage is non-empty; best-effort, - # never raises. In the finally so a turn that errors or times out - # AFTER a model call still reports the tokens it spent. - await self.emit_usage(tools, turn.usage) + self._turn_state[room_id] = turn + try: + final_text = await self._run_turn(live, prompt, turn) + except Exception: + logger.exception("Room %s: Copilot turn failed", room_id) + # Abort any work the runtime is still doing for this turn and + # drop the session; the next message resumes it fresh by id. + await self._session_manager.evict_session(room_id) + await tools.send_failure( + AgentFailure(_PROVIDER, GENERIC_PROVIDER_FAILURE_MESSAGE) + ) + raise + finally: + self._turn_state.pop(room_id, None) + # No-op unless Emit.USAGE is on / usage is non-empty; best-effort, + # never raises. In the finally so a turn that errors or times out + # AFTER a model call still reports the tokens it spent. + await self.emit_usage(tools, turn.usage) await self._emit_thoughts(turn, tools) @@ -889,8 +920,31 @@ def _compose_prompt( parts.append(msg.format_for_llm()) return "\n\n".join(parts) + @staticmethod + async def _abort_turn(live: SessionTurn) -> None: + """Abort the prompt and, while one is running, wait for the idle that ends it. + + The CLI emits that idle just after acknowledging the abort, and + ``send_and_wait`` ends on any idle: a next prompt sent before it + arrives returns on it with no reply. Subscribing first also catches + an idle that arrives ahead of the acknowledgement. + """ + idled = asyncio.Event() + + def on_event(event: SessionEvent) -> None: + if isinstance(event.data, SessionIdleData): + idled.set() + + unsubscribe = live.session.on(on_event) + try: + await live.session.abort() + if live.prompting: + await idled.wait() + finally: + unsubscribe() + async def _run_turn( - self, session: CopilotSession, prompt: str, turn: TurnState + self, live: SessionTurn, prompt: str, turn: TurnState ) -> str | None: """Send one prompt, wait for idle, and return the final assistant text. @@ -908,13 +962,15 @@ def collect(event: SessionEvent) -> None: elif incident := _turn_incident(data): turn.incidents.append(incident) - unsubscribe = session.on(collect) + unsubscribe = live.session.on(collect) + live.prompting = True try: - final_event = await session.send_and_wait( + final_event = await live.session.send_and_wait( prompt, timeout=self.config.turn_timeout_s ) finally: unsubscribe() + live.prompting = False content = getattr(final_event.data, "content", None) if final_event else None # An empty final message means no reply — same as no final event. diff --git a/src/band/adapters/letta.py b/src/band/adapters/letta.py index 4fec51716..28dcb83e6 100644 --- a/src/band/adapters/letta.py +++ b/src/band/adapters/letta.py @@ -26,6 +26,7 @@ TurnResultAlreadyReported, ) from band.core.simple_adapter import SimpleAdapter +from band.core.turn_lifecycle import BackendTurns from band.core.types import ( AdapterFeatures, Capability, @@ -64,6 +65,16 @@ _PROVIDER = "letta" +@dataclass(frozen=True) +class LettaRun: + """What cancelling one room's runs needs: the room's conversation in + shared mode, its own agent otherwise.""" + + client: Any + agent_id: str + conversation_id: str | None + + @dataclass class RoomContext: """Per-room state for a Letta agent.""" @@ -163,6 +174,9 @@ def __init__( # Per-room state self._rooms: dict[str, RoomContext] = {} + self._backend_turns = BackendTurns[LettaRun]( + self._cancel_runs, provider=_PROVIDER + ) # Shared mode: single agent ID used across all rooms self._shared_agent_id: str | None = None @@ -277,6 +291,10 @@ async def _build_client(self, agent_name: str) -> Any: # Message handling # ------------------------------------------------------------------ + @property + def backend_turns(self) -> BackendTurns[LettaRun]: + return self._backend_turns + async def on_message( self, msg: PlatformMessage, @@ -342,10 +360,6 @@ async def _handle_message( await tools.send_failure(AgentFailure(_PROVIDER, message)) raise RuntimeError(message) - # Point the MCP resolver at this room's current tools for the - # server-side tool calls this turn will make. - room_ctx.tools = tools - content = self._compose_turn_content( msg, room_ctx, @@ -483,15 +497,25 @@ async def _send_message( # Emitted on every exit via the finally. turn_usage = TurnUsage() + run = LettaRun(self._client, agent_id, self._shared_conversation(room_ctx)) try: try: # turn_timeout_s bounds only the round-trip to Letta -- response # processing (including deliver_reply, below) runs unbounded so a # slow Band-side delivery is never mislabeled as a Letta timeout. - response_messages, turn_usage = await asyncio.wait_for( - self._call_provider(agent_id, messages, room_ctx), - timeout=self.config.turn_timeout_s, - ) + async with self.backend_turns.running(room_id, run): + # Point the MCP resolver at this turn's tools only once the + # previous run's cancel has settled, so a late call from + # that run cannot reach them. + room_ctx.tools = tools + try: + response_messages, turn_usage = await asyncio.wait_for( + self._call_provider(run, messages), + timeout=self.config.turn_timeout_s, + ) + except TimeoutError: + await self.backend_turns.stop(room_id) + raise room_ctx.pending_seed = [] room_ctx.last_interaction = datetime.now(UTC) except TimeoutError: @@ -527,28 +551,36 @@ async def _send_message( # No-op unless Emit.USAGE is on; best-effort, never raises. await self.emit_usage(tools, turn_usage) + def _shared_conversation(self, room_ctx: RoomContext) -> str | None: + """The room's conversation when turns go through the Conversations API.""" + return room_ctx.conversation_id if self.config.mode == "shared" else None + async def _call_provider( - self, - agent_id: str, - messages: list[dict[str, str]], - room_ctx: RoomContext, + self, run: LettaRun, messages: list[dict[str, str]] ) -> tuple[list[Any], TurnUsage]: """Round-trip to the Letta API -- no response processing or delivery.""" - # Use Conversations API in shared mode, direct agent API in per_room mode - if self.config.mode == "shared" and room_ctx.conversation_id: - conversation_stream = await self._client.conversations.messages.create( - conversation_id=room_ctx.conversation_id, + if run.conversation_id: + conversation_stream = await run.client.conversations.messages.create( + conversation_id=run.conversation_id, messages=messages, ) response_messages = [resp_msg async for resp_msg in conversation_stream] return response_messages, TurnUsage() - response = await self._client.agents.messages.create( - agent_id=agent_id, + response = await run.client.agents.messages.create( + agent_id=run.agent_id, messages=messages, ) return list(response.messages), self._usage_from_response(response) + @staticmethod + async def _cancel_runs(run: LettaRun) -> None: + """Cancel the room's in-flight runs; the server needs Redis for this.""" + if run.conversation_id: + await run.client.conversations.cancel(run.conversation_id) + else: + await run.client.agents.messages.cancel(run.agent_id) + async def _process_response_messages( self, response_messages: list[Any], @@ -1110,6 +1142,7 @@ async def on_cleanup(self, room_id: str) -> None: """ async with self._rpc_lock: room_ctx = self._rooms.pop(room_id, None) + await self.backend_turns.settled(room_id) if room_ctx and self._client: if self.config.delete_agents_on_cleanup and self.config.mode == "per_room": @@ -1133,6 +1166,7 @@ async def cleanup_all(self) -> None: marking needed here. (A runtime re-registration triggered by a stale 404 is what marks rooms stale; see ``_reregister_and_mark_stale``.) """ + await self.backend_turns.drain() await self._mcp.release(self._client) async def _delete_agent(self, agent_id: str, room_id: str) -> None: diff --git a/src/band/adapters/opencode/adapter.py b/src/band/adapters/opencode/adapter.py index 8de9c2381..7d510a6dc 100644 --- a/src/band/adapters/opencode/adapter.py +++ b/src/band/adapters/opencode/adapter.py @@ -28,6 +28,7 @@ ) from band.core.simple_adapter import SimpleAdapter from band.core.turn import judge_detached_turn +from band.core.turn_lifecycle import BackendTurns from band.core.types import ( AdapterFeatures, Capability, @@ -104,6 +105,11 @@ class TurnState: # since message.updated streams repeatedly). Summed across messages at turn # end — a tool loop produces several assistant messages. usage_by_message: dict[str, TurnUsage] = field(default_factory=dict) + # The task that waits out this turn and posts its result. + watcher: asyncio.Task[None] | None = None + # Set while that watcher reports a timeout, so the stop it requested + # aborts the session without cancelling the task that still posts the report. + reports_timeout: bool = False # Resolves to the TurnState an ApprovalPorts bundle acts on -- either pinned @@ -120,7 +126,6 @@ class RoomState: session_id: str | None = None tools: AgentToolsProtocol | None = None turn: TurnState | None = None - turn_task: asyncio.Task[None] | None = None # Bound in _get_or_create_room_state at construction, and rebound by # _begin_turn on every subsequent turn. approvals: RoomApprovals = field(init=False) @@ -149,7 +154,6 @@ def begin_turn( turn_release_future=loop.create_future(), pending_mentions=[{"id": sender_id}] if sender_id else [], ) - self.turn_task = None return self.turn def record_message( @@ -293,6 +297,9 @@ def __init__( self._event_task: asyncio.Task[None] | None = None self._mcp = SharedBandMCPBackend(self._mcp_settings) self._rooms: dict[str, RoomState] = {} + self._backend_turns = BackendTurns[TurnState]( + self._stop_turn, provider=_PROVIDER + ) self._room_by_session: dict[str, str] = {} self._state_lock = asyncio.Lock() self._mcp_lifecycle_lock = asyncio.Lock() @@ -424,6 +431,10 @@ async def try_handle_control( tools.turn.settle() return True + @property + def backend_turns(self) -> BackendTurns[TurnState]: + return self._backend_turns + async def on_message( self, msg: PlatformMessage, @@ -436,6 +447,9 @@ async def on_message( room_id: str, ) -> None: room_state = await self._get_or_create_room_state(room_id) + # The room's previous turn may still be aborting its session; nothing + # of this turn is bound, and no prompt is sent, until that settles. + await self.backend_turns.settled(room_id) if self._client is None: agent_id = getattr(tools, "agent_id", None) @@ -458,7 +472,6 @@ async def on_message( raise RuntimeError("OpenCode client is not initialized") turn: TurnState | None = None - turn_future: asyncio.Future[None] | None = None try: session_id, created = await self._ensure_session(room_state, history) if Emit.TASK_EVENTS in self.features.emit and ( @@ -476,18 +489,17 @@ async def on_message( tools=tools, sender_id=msg.sender_id, ) - # Snapshot THIS turn's state before the prompt await: prompt_async - # can span the whole turn (session.idle may arrive mid-POST), and a - # message racing in during that window would _begin_turn again; - # reading room_state afterwards would wire this turn's watch task - # to the wrong turn's future and usage dict. - release_future = turn.turn_release_future - turn_future = turn.turn_future + # The watcher owns the turn's backend work, so it is registered + # before the prompt can start any. + turn_task = asyncio.create_task( + self._watch_turn_completion(room_state, room_id, turn) + ) + turn.watcher = turn_task try: # Turn-phase diagnostics (classify a stuck turn from CI logs): # if 'returned' never follows 'start', prompt_async is blocking - # (submission/scheduling), which no watcher bounds -- see the - # note below on _watch_turn_completion owning the timeout. + # (submission/scheduling) and the watcher, already running, ends + # the turn at its timeout -- see the note below. logger.info( "OpenCode turn: prompt_async start room=%s session=%s", room_id, @@ -525,22 +537,8 @@ async def on_message( self._clear_turn_state(room_state, expected_turn=turn) raise - turn_task = asyncio.create_task( - self._watch_turn_completion( - room_state, - room_id, - turn, - ) - ) - # Register the watcher only while this turn is still current; a - # superseded turn's task must not clobber (or be cancelled through) - # the next turn's ambient pointer. - if room_state.turn is turn: - room_state.turn_task = turn_task - - if release_future is not None: - await release_future - if turn_future is not None and turn_future.done(): + await turn.turn_release_future + if turn.turn_future.done(): await turn_task # NOTE: the turn timeout is owned solely by _watch_turn_completion (via # asyncio.wait_for), which aborts the session and emits the error event. @@ -548,20 +546,10 @@ async def on_message( # timeout handler of its own. except asyncio.CancelledError: # The runtime interrupts a turn by cancelling this coroutine, but the - # watcher runs detached. Left alone it outlives the interrupt: it - # posts the very reply the user stopped, and holds the room's busy - # guard until turn_timeout_s. Drop the turn state even if a second - # cancellation (e.g. a double-tap stop) cuts off the abandon/abort - # awaits below -- otherwise the room stays wedged on a turn whose - # state never clears. - if turn_future is not None: - assert turn is not None - try: - already_aborted = await turn.approvals.abandon() - finally: - self._clear_turn_state(room_state, expected_turn=turn) - if not already_aborted: - await self._abort_turn(turn, "interrupted") + # watcher runs detached; left alone it would post the very reply the + # user stopped. Cancelling it stops the session through backend_turns. + if turn is not None and turn.watcher is not None: + turn.watcher.cancel() raise except httpx.HTTPStatusError as exc: logger.exception("OpenCode request failed for room %s", room_id) @@ -593,6 +581,7 @@ async def on_cleanup(self, room_id: str) -> None: try: if room_state: try: + await self.backend_turns.stop(room_id) await room_state.approvals.abandon() finally: self._clear_turn_state(room_state) @@ -771,6 +760,7 @@ async def _shutdown_client(self) -> None: self._mcp_server_name, ) + await self.backend_turns.drain() # What follows acts only on objects already detached from ``self``, so # no successor can be affected — and the lock must not be held across a # stop that waits out OpenCode's open SSE read. @@ -1018,42 +1008,33 @@ async def _watch_turn_completion( room_id: str, turn: TurnState, ) -> None: - # 'watcher started' after 'prompt_async returned' but no later - # 'session.idle' points at a lost/late SSE terminal event, not a slow - # model -- distinct from the 'timed out' branch below firing at 300s. + # 'prompt_async returned' but no later 'session.idle' points at a + # lost/late SSE terminal event, not a slow model -- distinct from the + # 'timed out' branch below firing at 300s. 'watcher started' only says + # the watcher is running: it starts before the prompt is sent. logger.info( "OpenCode turn: watcher started room=%s session=%s timeout=%ss", room_id, turn.session_id, self.config.turn_timeout_s, ) + # The timeout abort stays inside running(), so a STOP during that + # HTTP call still finds the turn and the abort keeps running in the + # owned stop task after this watcher is cancelled. try: - await self._await_turn(turn) - except TimeoutError: - logger.warning( - "OpenCode turn timed out for room %s (session=%s)", - room_id, - turn.session_id, - ) - await self._abort_turn(turn, "timed-out") - try: - if turn.tools: - await turn.tools.send_failure( - AgentFailure( - _PROVIDER, - "OpenCode timed out before completing the turn.", - FAILURE_CODE_TIMEOUT, - ) + async with self.backend_turns.running(room_id, turn): + try: + await self._await_turn(turn) + except TimeoutError: + logger.warning( + "OpenCode turn timed out for room %s (session=%s)", + room_id, + turn.session_id, ) - # Tokens spent before the timeout were still spent — emit them, - # same as the success path (best-effort; no-op if none captured). - await self._emit_turn_usage(room_id, turn) - except Exception: - logger.exception( - "Failed to report the OpenCode timeout for room %s", room_id - ) - await self._report_delivery_failure(turn) - else: + turn.reports_timeout = True + await self.backend_turns.stop(room_id) + await self._report_timeout(room_id, turn) + return try: await self._deliver_fallback_text(room_state.room_id, turn) await judge_detached_turn(turn.tools, room_id=room_state.room_id) @@ -1067,7 +1048,7 @@ async def _watch_turn_completion( # Release the on_message waiter even if delivering the reply or # emitting usage raised (e.g. a sender-less turn has no one to # @mention, which the platform rejects) — otherwise on_message - # waits on the captured release_future forever. + # waits on turn_release_future forever. self._release_turn_wait_for(turn) self._clear_turn_state( room_state, @@ -1075,6 +1056,38 @@ async def _watch_turn_completion( expected_task=asyncio.current_task(), ) + async def _stop_turn(self, turn: TurnState) -> None: + """Stop a cancelled turn's backend work, posting nothing for it. + + The watcher goes first, so the abort's ``session.idle`` cannot make it + post fallback text. A watcher that is itself reporting a timeout has + already left that path; cancelling it would drop the report. + """ + if turn.watcher is not None and not turn.reports_timeout: + turn.watcher.cancel() + reason = "timed-out" if turn.reports_timeout else "stopped" + if not await turn.approvals.abandon(): + await self._abort_turn(turn, reason) + + async def _report_timeout(self, room_id: str, turn: TurnState) -> None: + try: + if turn.tools: + await turn.tools.send_failure( + AgentFailure( + _PROVIDER, + "OpenCode timed out before completing the turn.", + FAILURE_CODE_TIMEOUT, + ) + ) + # Tokens spent before the timeout were still spent — emit them, + # same as the success path (best-effort; no-op if none captured). + await self._emit_turn_usage(room_id, turn) + except Exception: + logger.exception( + "Failed to report the OpenCode timeout for room %s", room_id + ) + await self._report_delivery_failure(turn) + async def _abort_turn(self, turn: TurnState, reason: str) -> None: """Best-effort: tell OpenCode to stop working on this room's session.""" try: @@ -1172,14 +1185,12 @@ def _clear_turn_state( if expected_turn is not None and room_state.turn is not expected_turn: return - turn_task = room_state.turn_task - if turn_task is not None and turn_task is not expected_task: - turn_task.cancel() - if room_state.turn is not None: + watcher = room_state.turn.watcher + if watcher is not None and watcher is not expected_task: + watcher.cancel() room_state.turn.approvals.cancel() room_state.turn = None - room_state.turn_task = None @staticmethod def _resolve_future(future: asyncio.Future[None] | None) -> None: diff --git a/src/band/integrations/a2a/adapter.py b/src/band/integrations/a2a/adapter.py index fa6308bdf..3ff55dc9f 100644 --- a/src/band/integrations/a2a/adapter.py +++ b/src/band/integrations/a2a/adapter.py @@ -3,17 +3,16 @@ from __future__ import annotations import logging -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, AsyncIterator from contextlib import aclosing +from dataclasses import dataclass from typing import ClassVar, cast import httpx from a2a.client import Client, ClientConfig, ClientFactory from a2a.helpers import get_message_text, new_text_message from a2a.types import ( - Message as A2AMessage, -) -from a2a.types import ( + CancelTaskRequest, Role, SendMessageRequest, StreamResponse, @@ -21,6 +20,9 @@ Task, TaskState, ) +from a2a.types import ( + Message as A2AMessage, +) from band_sdk_core import AgentFailure from typing_extensions import Unpack @@ -37,6 +39,7 @@ send_event_safe, ) from band.core.simple_adapter import SimpleAdapter +from band.core.turn_lifecycle import BackendTurns from band.core.types import Capability, Emit, FeatureKwargs, PlatformMessage from band.integrations.a2a.protocol import ( RETRYABLE_TASK_FAILURE_STATES, @@ -53,6 +56,52 @@ _PROVIDER = "a2a" + +@dataclass +class A2ATaskRun: + """The remote task a room's turn is running, and where to cancel it.""" + + client: Client + room_id: str + task_id: str + # Set once the remote reports a terminal state: nothing is left to cancel, + # even while this turn is still delivering that state. + finished: bool = False + + +async def _prepended( + read: list[StreamResponse], rest: AsyncIterator[StreamResponse] +) -> AsyncIterator[StreamResponse]: + """Replay events already read from a stream ahead of the rest of it.""" + for event in read: + yield event + async for event in rest: + yield event + + +async def _reveal_task( + client: Client, + room_id: str, + events: AsyncIterator[StreamResponse], + read: list[StreamResponse], +) -> A2ATaskRun | None: + """Read the stream's first event, which reveals a new turn's task. + + The event is appended to ``read`` for the turn to deliver. A message + reply reveals no task. + """ + first = await anext(events, None) + if first is None: + return None + read.append(first) + task = apply_task_stream_event(None, first) + if task is None: + return None + return A2ATaskRun( + client, room_id, task.id, finished=task.status.state in TERMINAL_TASK_STATES + ) + + # httpx's read timeout resets on every chunk received, so this bounds the gap # between SSE events, not the turn as a whole. Generous enough for the # multi-second silences of a live LLM call or tool loop; still finite, so a @@ -119,6 +168,9 @@ def __init__( self._http_client: httpx.AsyncClient | None = None self._contexts: dict[str, str] = {} # room_id → A2A context_id self._tasks: dict[str, str] = {} # room_id → last task_id + self._backend_turns = BackendTurns[A2ATaskRun]( + self._cancel_task, provider=_PROVIDER + ) self._task_cache: dict[tuple[str, str], Task] = {} # Track sender per task for mentions: (room_id, task_id) → sender info self._task_senders: dict[tuple[str, str], dict[str, str]] = {} @@ -151,6 +203,10 @@ async def on_started(self, agent_name: str, agent_description: str) -> None: bool(headers), ) + @property + def backend_turns(self) -> BackendTurns[A2ATaskRun]: + return self._backend_turns + async def on_message( self, msg: PlatformMessage, @@ -173,21 +229,39 @@ async def on_message( history, ) - # Rehydrate session state on bootstrap + # A pending tasks/cancel deletes the local continuation, then waits + # on the remote. Rehydration and the outgoing message must run after + # that wait, or a still-nonterminal resubscribe restores the task + # the cancel is about to finish. + await self.backend_turns.settled(room_id) if is_session_bootstrap and history: await self._rehydrate_from_history(room_id, history) - # Convert Band message to A2A format a2a_message = self._to_a2a_message(msg, room_id) + # A continued task is known before the remote answers; a new one only + # once its first event arrives. + continued = ( + A2ATaskRun(self._client, room_id, a2a_message.task_id) + if a2a_message.task_id + else None + ) try: - # Send to remote A2A agent and process events - async for event in self._client.send_message( - SendMessageRequest(message=a2a_message) - ): - await self._handle_event( - event, tools, room_id, msg.sender_id, msg.sender_name + # Send to remote A2A agent and process events. A task paused on + # INPUT_REQUIRED ends the stream normally and stays for the next turn. + async with self.backend_turns.running(room_id, continued): + events = self._client.send_message( + SendMessageRequest(message=a2a_message) ) + run, read = continued, [] + if run is None: + run = await self.backend_turns.started( + room_id, _reveal_task(self._client, room_id, events, read) + ) + async for event in _prepended(read, events): + await self._handle_event( + event, tools, room_id, msg.sender_id, msg.sender_name, run + ) except DeliveryFailedError as e: reraise_delivery_cause(e) @@ -207,6 +281,7 @@ async def _handle_event( room_id: str, sender_id: str, sender_name: str | None, + run: A2ATaskRun | None = None, ) -> None: """Handle A2A event and forward to Band platform.""" if event.HasField("message"): @@ -221,6 +296,8 @@ async def _handle_event( self._task_senders.setdefault(key, {"id": sender_id, "name": sender_name or ""}) state = task.status.state + if run is not None and state in TERMINAL_TASK_STATES: + run.finished = True try: await self._deliver_task_update(task, tools, self._task_senders[key]) if state == TaskState.TASK_STATE_INPUT_REQUIRED: @@ -269,6 +346,13 @@ def _reduce_task_event(self, room_id: str, event: StreamResponse) -> Task | None self._task_cache[(room_id, task.id)] = task return task + async def _cancel_task(self, run: A2ATaskRun) -> None: + """Release the turn's task and cancel it remotely unless it finished; + a cancelled task is never continued.""" + self._finalize_task(run.room_id, run.task_id) + if not run.finished: + await run.client.cancel_task(CancelTaskRequest(id=run.task_id)) + def _remember_task(self, room_id: str, task: Task) -> None: """Record task identity needed for subsequent turns and resumption.""" self._tasks[room_id] = task.id @@ -311,10 +395,12 @@ async def _deliver_task_update( raise TurnResultAlreadyReported(error_text) def _finalize_task(self, room_id: str, task_id: str) -> None: - """Release a terminal task after its Band output and state are persisted.""" + """Release a finished or cancelled task; the room's task pointer moves + on only if it still names this task.""" self._task_senders.pop((room_id, task_id), None) self._task_cache.pop((room_id, task_id), None) - self._tasks.pop(room_id, None) + if self._tasks.get(room_id) == task_id: + del self._tasks[room_id] def _to_a2a_message(self, msg: PlatformMessage, room_id: str) -> A2AMessage: """Convert Band message to A2A format.""" @@ -351,6 +437,7 @@ def _extract_response(self, task: Task) -> str: async def on_cleanup(self, room_id: str) -> None: """Clean up A2A context for room.""" + await self.backend_turns.settled(room_id) self._contexts.pop(room_id, None) self._tasks.pop(room_id, None) # A resubscribed task may live in the cache without a sender entry, @@ -367,6 +454,7 @@ async def on_cleanup(self, room_id: str) -> None: async def cleanup_all(self) -> None: """Close the owned A2A client and its HTTP transport.""" + await self.backend_turns.drain() client, self._client = self._client, None http_client, self._http_client = self._http_client, None try: diff --git a/src/band/integrations/a2a/types.py b/src/band/integrations/a2a/types.py index 343c72956..65761d1f6 100644 --- a/src/band/integrations/a2a/types.py +++ b/src/band/integrations/a2a/types.py @@ -51,7 +51,9 @@ class A2AAdapterConfig(BaseAdapterConfig): Attributes: remote_url: Base URL of the remote A2A agent. auth: Credentials sent with every request to the remote agent. - streaming: Whether to use streaming mode (SSE) for responses. + streaming: Whether to use streaming mode (SSE) for responses. Without + it (or when the remote's card does not advertise it) a new task's + first event is its final answer, so a STOP cannot cancel that task. """ remote_url: str diff --git a/tests/adapters/copilot_sdk/fakes.py b/tests/adapters/copilot_sdk/fakes.py index 96f84ec20..580edbaf0 100644 --- a/tests/adapters/copilot_sdk/fakes.py +++ b/tests/adapters/copilot_sdk/fakes.py @@ -8,6 +8,7 @@ from __future__ import annotations +import asyncio from datetime import UTC, datetime from types import SimpleNamespace from typing import Any @@ -31,7 +32,7 @@ ) if _COPILOT_SDK_AVAILABLE: - from copilot.generated.session_events import SessionErrorData + from copilot.generated.session_events import SessionErrorData, SessionIdleData def make_platform_message( @@ -101,6 +102,14 @@ def __init__( self.handlers: list[Any] = [] self.aborted = False self.disconnected = False + # Every prompt, abort and disconnect, in the order the session saw them. + self.journal: list[str] = [] + # While cleared, an abort stays in flight. + self.abort_released = asyncio.Event() + self.abort_released.set() + # While set, an abort is acknowledged but the idle that ends the aborted + # turn is withheld until ``deliver_aborted_idle``. + self.holds_aborted_idle = False def on(self, handler: Any) -> Any: self.handlers.append(handler) @@ -113,6 +122,7 @@ async def send_and_wait( self, prompt: str, *, timeout: float = 60.0, **_: Any ) -> Any: self.prompts.append(prompt) + self.journal.append("prompt") if self.send_error: raise self.send_error for data in self.turn_events: @@ -135,10 +145,22 @@ async def send_and_wait( return SimpleNamespace(data=SimpleNamespace(content=self.reply_content)) async def abort(self) -> None: + await self.abort_released.wait() self.aborted = True + self.journal.append("abort") + if not self.holds_aborted_idle: + self.deliver_aborted_idle() + + def deliver_aborted_idle(self) -> None: + """The CLI ends an aborted turn with an idle, as the real one does.""" + self.journal.append("idle") + event = SimpleNamespace(data=SessionIdleData(aborted=True)) + for handler in list(self.handlers): + handler(event) async def disconnect(self) -> None: self.disconnected = True + self.journal.append("disconnect") class FakeCopilotClient: diff --git a/tests/adapters/lettakit.py b/tests/adapters/lettakit.py index b9b7f8103..5ee801992 100644 --- a/tests/adapters/lettakit.py +++ b/tests/adapters/lettakit.py @@ -8,9 +8,11 @@ from __future__ import annotations +import asyncio import json from collections.abc import Awaitable, Callable, Sequence from datetime import UTC, datetime +from types import SimpleNamespace from typing import Any, Literal from unittest.mock import AsyncMock, MagicMock, create_autospec from uuid import uuid4 @@ -169,6 +171,18 @@ def ready_letta_adapter( return adapter, client +def band_mcp_registrations(adapter: LettaAdapter) -> dict[str, Any]: + """The adapter's real Band MCP tools by name: the path Letta's MCP client + reaches.""" + return { + registration.name: registration + for registration in build_resolved_band_mcp_tool_registrations( + get_tools=adapter._get_room_tools, + capabilities=adapter.features.capabilities, + ) + } + + def scripted_letta_turn( adapter: LettaAdapter, *, @@ -182,13 +196,7 @@ def scripted_letta_turn( path Letta's MCP client reaches), then the turn's messages come back as Letta's response. """ - registrations = { - registration.name: registration - for registration in build_resolved_band_mcp_tool_registrations( - get_tools=adapter._get_room_tools, - capabilities=adapter.features.capabilities, - ) - } + registrations = band_mcp_registrations(adapter) async def create(**_kwargs: Any) -> MagicMock: messages: list[Any] = [] @@ -332,3 +340,78 @@ async def __aiter__(self) -> Any: yield stream_message return _AsyncStream(list(messages)) + + +class RecordingLettaClient: + """A Letta server that records each turn sent and each run cancelled. + + ``hold()`` keeps every create in flight until ``release`` is set, as a + Letta agent mid-turn does. ``hold_cancel()`` does the same for each run + cancel, which is journaled only once it finishes. ``reject_cancels`` + answers each cancel with the server error Letta returns without Redis. + """ + + def __init__(self) -> None: + self.journal: list[str] = [] + self.release = asyncio.Event() + self.release.set() + self.reject_cancels = False + self.cancel_started = asyncio.Event() + self._cancel_release: asyncio.Event | None = None + self.agents = SimpleNamespace( + messages=SimpleNamespace( + create=self._create_agent_turn, cancel=self._cancel_agent_runs + ), + delete=self._delete_agent, + ) + self.conversations = SimpleNamespace( + messages=SimpleNamespace(create=self._create_conversation_turn), + cancel=self._cancel_conversation_runs, + ) + + def hold(self) -> None: + self.release.clear() + + def hold_cancel(self) -> asyncio.Event: + """Each cancel waits for the returned event after it has started.""" + self._cancel_release = asyncio.Event() + return self._cancel_release + + @property + def sent(self) -> list[str]: + return [entry for entry in self.journal if entry.startswith("send ")] + + @property + def cancels(self) -> list[str]: + return [entry for entry in self.journal if entry.startswith("cancel ")] + + async def _create_agent_turn( + self, *, agent_id: str, messages: list[dict[str, str]] + ) -> MagicMock: + self.journal.append(f"send agent {agent_id}") + await self.release.wait() + return make_letta_response(make_assistant_message()) + + async def _create_conversation_turn( + self, *, conversation_id: str, messages: list[dict[str, str]] + ) -> Any: + self.journal.append(f"send conversation {conversation_id}") + await self.release.wait() + return make_mock_async_stream(make_assistant_message()) + + async def _cancel_agent_runs(self, agent_id: str) -> None: + await self._cancel(f"cancel agent {agent_id}") + + async def _cancel_conversation_runs(self, conversation_id: str) -> None: + await self._cancel(f"cancel conversation {conversation_id}") + + async def _delete_agent(self, agent_id: str) -> None: + self.journal.append(f"delete agent {agent_id}") + + async def _cancel(self, entry: str) -> None: + self.cancel_started.set() + if self._cancel_release is not None: + await self._cancel_release.wait() + self.journal.append(entry) + if self.reject_cancels: + raise RuntimeError("Run cancellation requires Redis on the Letta server") diff --git a/tests/adapters/opencode/test_lifecycle.py b/tests/adapters/opencode/test_lifecycle.py index f1ba0f5dd..568326a2f 100644 --- a/tests/adapters/opencode/test_lifecycle.py +++ b/tests/adapters/opencode/test_lifecycle.py @@ -494,10 +494,17 @@ async def test_interrupting_a_turn_stops_the_reply_and_frees_the_room(tools) -> ) await wait_for(lambda: bool(fake_client.prompt_calls)) + watcher = (await adapter._get_or_create_room_state("room-1")).turn.watcher interrupted.cancel() with suppress(asyncio.CancelledError): await interrupted + # Cancelling on_message only schedules the watcher's cancellation; the abort + # is spawned once the watcher has unwound. + with suppress(asyncio.CancelledError): + await watcher + # The abort is sent from its own task, which the room's next turn joins. + await adapter.backend_turns.drain() assert fake_client.aborted_sessions == ["sess-1"], "OpenCode kept working" assert tools.messages_sent == [], "posted the reply the user interrupted" @@ -518,68 +525,6 @@ async def test_interrupting_a_turn_stops_the_reply_and_frees_the_room(tools) -> await adapter.on_cleanup("room-1") -async def test_double_stop_while_abandon_awaits_abort_still_clears_turn_state( - tools, -) -> None: - """Regression: a second cancellation (a double-tap stop) landing while - the interrupt handler's ``turn.approvals.abandon()`` await is itself - parked inside ``abort_session`` -- aborting a permission still waiting - on a human -- must not skip ``_clear_turn_state``. Otherwise the room - stays wedged forever on a turn whose state never clears. - """ - - class BlockingAbortClient(FakeOpencodeClient): - def __init__(self) -> None: - super().__init__(prompt_event_sequences=[[]]) - self.abort_started = asyncio.Event() - self.release_abort = asyncio.Event() - - async def abort_session(self, session_id: str) -> None: - self.abort_started.set() - await self.release_abort.wait() - await super().abort_session(session_id) - - fake_client = BlockingAbortClient() - adapter = OpencodeAdapter(client_factory=lambda _config: fake_client) - - await adapter.on_started("OpenCode Agent", "A coding agent") - interrupted = asyncio.create_task( - deliver_message( - adapter, - make_platform_message(content="long task"), - tools_protocol(tools), - OpencodeSessionState(), - participants_msg=None, - contacts_msg=None, - is_session_bootstrap=True, - room_id="room-1", - ) - ) - await wait_for(lambda: bool(fake_client.prompt_calls)) - - room_state = await adapter._get_or_create_room_state("room-1") - assert room_state.turn is not None - await room_state.turn.approvals.on_permission_asked( - OpencodePermissionRequest(id="perm-1", permission="bash") - ) - - # First stop: on_message's CancelledError handler finds the permission - # still parked on a human and starts abandon(), which awaits - # abort_session for it. - interrupted.cancel() - await fake_client.abort_started.wait() - - # Second stop (double-tap) cuts off that abort_session await mid-flight. - interrupted.cancel() - with suppress(asyncio.CancelledError): - await interrupted - - assert room_state.turn is None, "turn state must clear despite the double cancel" - assert room_state.turn_task is None - - await adapter.on_cleanup("room-1") - - async def test_teardown_does_not_disconnect_a_successor_registration(tools) -> None: """A superseded teardown must not strip the next client's MCP registration. diff --git a/tests/framework_conformance/stopprobes/__init__.py b/tests/framework_conformance/stopprobes/__init__.py index ebefa6461..1f9b46b8c 100644 --- a/tests/framework_conformance/stopprobes/__init__.py +++ b/tests/framework_conformance/stopprobes/__init__.py @@ -19,11 +19,16 @@ from band.client.streaming import ControlMode from band.core.simple_adapter import SimpleAdapter +from band.core.turn_lifecycle import BackendTurns from tests.adapters.waiting import until #: A control probe fails rather than hanging if its task never becomes ready. STARTUP_TIMEOUT_SECONDS = 5.0 +#: The turn timeout of a probe that times a turn out. Advanced by looptime; +#: provider cancellation retains its own shorter bound. +TURN_TIMEOUT_SECONDS = 60.0 + #: The controls that end a running turn. ENDING_CONTROLS = [ControlMode.STOP, ControlMode.INTERRUPT] @@ -44,16 +49,19 @@ class StopCase: expected: list[str] +STOP_WAITS = frozenset({BackendTurns.settled.__code__, BackendTurns.drain.__code__}) + + def waiting_on_stop(task: asyncio.Task[None]) -> bool: - """The task is blocked in ``BackendTurns.settled``. + """The task is blocked waiting for a backend stop, in + ``BackendTurns.settled`` or ``BackendTurns.drain``. - ``Task.get_stack`` stops at the outermost coroutine, so ``settled`` is + ``Task.get_stack`` stops at the outermost coroutine, so the wait is reached through the chain of awaited coroutines. """ coro = task.get_coro() while coro is not None: - code = getattr(coro, "cr_code", None) - if code is not None and code.co_name == "settled": + if getattr(coro, "cr_code", None) in STOP_WAITS: return True coro = getattr(coro, "cr_await", None) return False diff --git a/tests/framework_conformance/stopprobes/a2a.py b/tests/framework_conformance/stopprobes/a2a.py new file mode 100644 index 000000000..82426621f --- /dev/null +++ b/tests/framework_conformance/stopprobes/a2a.py @@ -0,0 +1,284 @@ +"""Backend-stop probes for the a2a adapter; see ``stopprobes``.""" + +from __future__ import annotations + +import asyncio +from contextlib import suppress +from dataclasses import replace + +from a2a.types import StreamResponse, TaskState + +from band.client.streaming import ControlMode +from band.core.types import HistoryProvider, MessageType +from band.integrations.a2a.adapter import A2AAdapter +from band.integrations.a2a.types import A2AAdapterConfig +from band.testing.fake_tools import FakeAgentTools +from tests.framework_conformance.stopprobes import ( + ENDING_CONTROLS, + StopCase, + cancel_turn, + end_turn, + stop_room, + wait_until, + waiting_on_stop, +) +from tests.framework_conformance.turnprobes import ROOM_ID, turn_input +from tests.integrations.a2a.kit import WORKING, RecordingA2AClient, task_update + +TASK_ID = "task-1" +INPUT_REQUIRED = TaskState.Name(TaskState.TASK_STATE_INPUT_REQUIRED) +COMPLETED = TaskState.Name(TaskState.TASK_STATE_COMPLETED) +ANSWER = "the answer" +CANCEL_TASK = f"tasks/cancel {TASK_ID}" +NO_TASK = "" +RESTORED_CANCELLED_TASK = "bootstrap restored the cancelled task" +CLIENT_CLOSED_MID_CANCEL = "client closed while tasks/cancel was in flight" +CLEANUP_RETURNED_MID_CANCEL = "cleanup returned while tasks/cancel was in flight" + + +def continues(task_id: str) -> str: + return f"next message continues {task_id or 'no task'}" + + +def started() -> tuple[A2AAdapter, RecordingA2AClient]: + adapter = A2AAdapter(A2AAdapterConfig(remote_url="http://a2a.test")) + client = RecordingA2AClient() + adapter._client = client + return adapter, client + + +def task_in(state_name: str, *, text: str | None = None) -> StreamResponse: + """A stream event reporting the probe's task in the named ``TaskState``.""" + return task_update(TASK_ID, state_name, text=text) + + +async def next_message(adapter: A2AAdapter, client: RecordingA2AClient) -> str: + client.script(task_in(COMPLETED)) + await adapter.on_event(turn_input(FakeAgentTools())) + return continues(client.continued[-1]) + + +async def end_running_task(mode: ControlMode) -> list[str]: + adapter, client = started() + client.script(task_in(WORKING), hold=True) + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: ROOM_ID in adapter._tasks) + await end_turn(adapter, turn, ROOM_ID, mode) + return [*client.cancelled, await next_message(adapter, client)] + + +async def stop_before_the_remote_answers( + mode: ControlMode, answer: StreamResponse +) -> list[str]: + """A STOP that lands before a new task's first event still reaches the task + that event reveals, unless that task has already finished.""" + adapter, client = started() + answered = asyncio.Event() + client.script(answer, hold=True, gate=answered) + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: len(client.continued) == 1) + await cancel_turn(turn) + await adapter.on_interrupt(ROOM_ID, mode) + answered.set() + await adapter.backend_turns.drain() + return [*client.cancelled, await next_message(adapter, client)] + + +async def stop_while_a_finished_task_is_delivered( + mode: ControlMode, *, resumed: bool = False +) -> list[str]: + """A STOP that lands while a finished task's reply is being posted has no + task left to cancel, whether the task is new or resumed from INPUT_REQUIRED.""" + adapter, client = started() + if resumed: + client.script(task_in(INPUT_REQUIRED)) + await adapter.on_event(turn_input(FakeAgentTools())) + client.script(task_in(COMPLETED, text=ANSWER)) + tools = FakeAgentTools() + delivering = tools.hold_message(ANSWER) + turn = asyncio.create_task(adapter.on_event(turn_input(tools))) + async with delivering: + await end_turn(adapter, turn, ROOM_ID, mode) + return [*client.cancelled, await next_message(adapter, client)] + + +async def stop_continued_task_before_it_answers(mode: ControlMode) -> list[str]: + """A continued task is cancelled even before the remote's first event.""" + adapter, client = started() + client.script(task_in(INPUT_REQUIRED)) + await adapter.on_event(turn_input(FakeAgentTools())) + client.script(hold=True) + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: len(client.continued) == 2) + await end_turn(adapter, turn, ROOM_ID, mode) + return [*client.cancelled, await next_message(adapter, client)] + + +def input_required_history() -> list[dict[str, object]]: + return [ + { + "message_type": MessageType.TASK, + "content": f"A2A task {INPUT_REQUIRED}", + "metadata": { + "a2a_context_id": "ctx-1", + "a2a_task_id": TASK_ID, + "a2a_task_state": INPUT_REQUIRED, + }, + } + ] + + +async def bootstrap_during_cancel() -> list[str]: + """A bootstrap that starts while tasks/cancel is held does not resume that task.""" + adapter, client = started() + release = client.hold_cancel() + client.script(task_in(INPUT_REQUIRED)) + await adapter.on_event(turn_input(FakeAgentTools())) + client.script(hold=True) + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: len(client.continued) == 2) + turn.cancel() + with suppress(asyncio.CancelledError): + await turn + await adapter.on_interrupt(ROOM_ID, ControlMode.STOP) + await wait_until(lambda: client.cancel_started.is_set()) + client.script(task_in(COMPLETED)) + bootstrap = asyncio.create_task( + adapter.on_event( + replace( + turn_input(FakeAgentTools()), + history=HistoryProvider(raw=input_required_history()), + ) + ) + ) + await wait_until(lambda: waiting_on_stop(bootstrap)) + restored = adapter._tasks.get(ROOM_ID) == TASK_ID + release.set() + await bootstrap + return [ + *client.cancelled, + *([RESTORED_CANCELLED_TASK] if restored else []), + continues(client.continued[-1]), + ] + + +async def stop_holding_cancel() -> tuple[A2AAdapter, RecordingA2AClient, asyncio.Event]: + """Stop a running task and leave its tasks/cancel in flight until the event is set.""" + adapter, client = started() + release = client.hold_cancel() + client.script(task_in(WORKING), hold=True) + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: ROOM_ID in adapter._tasks) + await cancel_turn(turn) + await adapter.on_interrupt(ROOM_ID, ControlMode.STOP) + await wait_until(client.cancel_started.is_set) + return adapter, client, release + + +async def cleanup_waits_for_the_cancel() -> list[str]: + """Leaving the room does not return while tasks/cancel is in flight.""" + adapter, client, release = await stop_holding_cancel() + cleanup = asyncio.create_task(adapter.on_cleanup(ROOM_ID)) + await wait_until(lambda: waiting_on_stop(cleanup) or cleanup.done()) + returned_mid_cancel = cleanup.done() + release.set() + await cleanup + return [ + *client.cancelled, + *([CLEANUP_RETURNED_MID_CANCEL] if returned_mid_cancel else []), + ] + + +async def shutdown_waits_for_the_cancel() -> list[str]: + """Shutdown does not close the client while tasks/cancel is in flight.""" + adapter, client, release = await stop_holding_cancel() + shutdown = asyncio.create_task(adapter.cleanup_all()) + await wait_until(lambda: waiting_on_stop(shutdown) or shutdown.done()) + closed_mid_cancel = client.closed + release.set() + await shutdown + return [ + *client.cancelled, + *([CLIENT_CLOSED_MID_CANCEL] if closed_mid_cancel else []), + ] + + +async def input_required_is_not_cancelled() -> list[str]: + """A task paused for input ends its turn normally and stays continuable.""" + adapter, client = started() + client.script(task_in(INPUT_REQUIRED)) + await adapter.on_event(turn_input(FakeAgentTools())) + await stop_room(adapter, ROOM_ID) + return [*client.cancelled, await next_message(adapter, client)] + + +async def idle() -> list[str]: + adapter, client = started() + await stop_room(adapter, ROOM_ID) + return client.cancelled + + +CASES: dict[str, StopCase] = { + **{ + f"{mode}-mid-turn": StopCase( + run=lambda mode=mode: end_running_task(mode), + expected=[CANCEL_TASK, continues(NO_TASK)], + ) + for mode in ENDING_CONTROLS + }, + **{ + f"{mode}-before-a-new-task-answers": StopCase( + run=lambda mode=mode: stop_before_the_remote_answers( + mode, task_in(WORKING) + ), + expected=[CANCEL_TASK, continues(NO_TASK)], + ) + for mode in ENDING_CONTROLS + }, + **{ + f"{mode}-before-a-finished-task-answers": StopCase( + run=lambda mode=mode: stop_before_the_remote_answers( + mode, task_in(COMPLETED) + ), + expected=[continues(NO_TASK)], + ) + for mode in ENDING_CONTROLS + }, + **{ + f"{mode}-while-a-finished-task-is-delivered": StopCase( + run=lambda mode=mode: stop_while_a_finished_task_is_delivered(mode), + expected=[continues(NO_TASK)], + ) + for mode in ENDING_CONTROLS + }, + **{ + f"{mode}-while-a-continued-task-finishes": StopCase( + run=lambda mode=mode: stop_while_a_finished_task_is_delivered( + mode, resumed=True + ), + expected=[continues(NO_TASK)], + ) + for mode in ENDING_CONTROLS + }, + **{ + f"{mode}-continued-task-before-it-answers": StopCase( + run=lambda mode=mode: stop_continued_task_before_it_answers(mode), + expected=[CANCEL_TASK, continues(NO_TASK)], + ) + for mode in ENDING_CONTROLS + }, + "input-required-is-not-cancelled": StopCase( + run=input_required_is_not_cancelled, expected=[continues(TASK_ID)] + ), + "bootstrap-during-cancel": StopCase( + run=bootstrap_during_cancel, + expected=[CANCEL_TASK, continues(NO_TASK)], + ), + "cleanup-waits-for-the-cancel": StopCase( + run=cleanup_waits_for_the_cancel, expected=[CANCEL_TASK] + ), + "shutdown-waits-for-the-cancel": StopCase( + run=shutdown_waits_for_the_cancel, expected=[CANCEL_TASK] + ), + "idle-room": StopCase(run=idle, expected=[]), +} diff --git a/tests/framework_conformance/stopprobes/copilotsdk.py b/tests/framework_conformance/stopprobes/copilotsdk.py new file mode 100644 index 000000000..da9c89fa4 --- /dev/null +++ b/tests/framework_conformance/stopprobes/copilotsdk.py @@ -0,0 +1,244 @@ +"""Backend-stop probes for the copilot_sdk adapter; see ``stopprobes``. + +The adapter stops a room's backend turn with ``session.abort()``, which keeps +the session usable, so each probe reports what the room's sessions saw. +""" + +from __future__ import annotations + +import asyncio +from collections.abc import Callable + +from band.adapters.copilot_sdk import CopilotSDKAdapter +from band.client.streaming import ControlMode +from band.runtime.tools import BandTool +from band.testing.fake_tools import FakeAgentTools +from tests.adapters.copilot_sdk.fakes import ( + FakeCopilotClient, + FakeCopilotSession, + ToolSchemaFakeTools, + make_started_adapter, +) +from tests.adapters.waiting import spin +from tests.framework_conformance.stopprobes import ( + ENDING_CONTROLS, + StopCase, + cancel_turn, + end_turn, + stop_room, + wait_until, + waiting_on_stop, +) +from tests.framework_conformance.turnprobes import ROOM_ID, turn_input + + +class FirstTurnHangs: + """A scripted mid-turn step: the first turn never finishes, later ones do.""" + + def __init__(self) -> None: + self.entered = asyncio.Event() + + async def __call__(self, session: FakeCopilotSession) -> None: + if self.entered.is_set(): + return + self.entered.set() + await asyncio.Future[None]() + + +def journal(client: FakeCopilotClient) -> list[str]: + """Every session's prompts, aborts and disconnects, by session number.""" + return [ + f"session {number} {entry}" + for number, session in enumerate(client.sessions, start=1) + for entry in session.journal + ] + + +async def started() -> tuple[CopilotSDKAdapter, FakeCopilotClient, FirstTurnHangs]: + hang = FirstTurnHangs() + client = FakeCopilotClient(turn_events=[hang]) + return await make_started_adapter(client), client, hang + + +async def stop_running_turn( + adapter: CopilotSDKAdapter, hang: FirstTurnHangs, mode: ControlMode +) -> None: + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(hang.entered.is_set) + await end_turn(adapter, turn, ROOM_ID, mode) + + +def hold_abort(session: FakeCopilotSession) -> None: + session.abort_released.clear() + + +def hold_aborted_idle(session: FakeCopilotSession) -> None: + session.holds_aborted_idle = True + + +async def stop_holding( + adapter: CopilotSDKAdapter, + client: FakeCopilotClient, + hang: FirstTurnHangs, + tools: FakeAgentTools, + hold: Callable[[FakeCopilotSession], None], +) -> FakeCopilotSession: + """Stop the running turn and leave part of its abort in flight, as ``hold`` says.""" + turn = asyncio.create_task(adapter.on_event(turn_input(tools))) + await wait_until(hang.entered.is_set) + session = client.sessions[0] + hold(session) + await cancel_turn(turn) + await adapter.on_interrupt(ROOM_ID, ControlMode.STOP) + return session + + +async def end_running_turn(mode: ControlMode) -> list[str]: + adapter, client, hang = await started() + await stop_running_turn(adapter, hang, mode) + return journal(client) + + +async def next_message_waits_for_the_abort() -> list[str]: + """The next message waits for the abort, then runs on the same session.""" + adapter, client, hang = await started() + session = await stop_holding(adapter, client, hang, FakeAgentTools(), hold_abort) + next_turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + # The turn blocks inside the context manager that owns the wait, where + # ``waiting_on_stop`` cannot see it; a turn that ignored the abort would + # prompt the session within these loop turns. + await spin() + session.abort_released.set() + await next_turn + return journal(client) + + +async def next_message_waits_for_the_aborted_idle() -> list[str]: + """The next message waits for the idle that ends the aborted turn: arriving + after the abort is acknowledged, it would end that message's prompt with no + reply.""" + adapter, client, hang = await started() + session = await stop_holding( + adapter, client, hang, FakeAgentTools(), hold_aborted_idle + ) + next_turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + # As in ``next_message_waits_for_the_abort``, the wait is out of + # ``waiting_on_stop``'s sight: a turn that ignored the idle would prompt the + # session within these loop turns. + await spin() + session.deliver_aborted_idle() + await next_turn + return journal(client) + + +async def late_tool_call_during_held_abort() -> list[str]: + """A tool call from the stopped turn that lands while its abort is pending + is answered on that turn's tools, not on the replacement turn's.""" + from copilot import ( # noqa: PLC0415 -- copilot extra, absent from some lane venvs + ToolInvocation, + ) + + adapter, client, hang = await started() + stopped, replacement = ToolSchemaFakeTools(), ToolSchemaFakeTools() + session = await stop_holding(adapter, client, hang, stopped, hold_abort) + next_turn = asyncio.create_task(adapter.on_event(turn_input(replacement))) + await spin() + await session.find_tool(BandTool.SEND_MESSAGE).handler( + ToolInvocation( + tool_call_id="late-call", + tool_name=BandTool.SEND_MESSAGE, + arguments={"content": "from the stopped turn", "mentions": []}, + ) + ) + session.abort_released.set() + await next_turn + return [ + *( + f"stopped turn's tools got {call['tool_name']}" + for call in stopped.tool_calls + ), + *( + f"replacement turn's tools got {call['tool_name']}" + for call in replacement.tool_calls + ), + ] + + +async def cleanup_waits_for_the_abort() -> list[str]: + """Leaving the room waits for the abort before it disconnects the session.""" + adapter, client, hang = await started() + session = await stop_holding(adapter, client, hang, FakeAgentTools(), hold_abort) + cleanup = asyncio.create_task(adapter.on_cleanup(ROOM_ID)) + await wait_until(lambda: waiting_on_stop(cleanup) or cleanup.done()) + session.abort_released.set() + await cleanup + return journal(client) + + +async def shutdown_right_after_stop() -> list[str]: + """Shutdown waits for the room's abort before it disconnects the session.""" + adapter, client, hang = await started() + session = await stop_holding(adapter, client, hang, FakeAgentTools(), hold_abort) + shutdown = asyncio.create_task(adapter.cleanup_all()) + await wait_until(lambda: waiting_on_stop(shutdown) or shutdown.done()) + session.abort_released.set() + await shutdown + return journal(client) + + +async def idle() -> list[str]: + adapter, client, _ = await started() + await stop_room(adapter, ROOM_ID) + return journal(client) + + +CASES: dict[str, StopCase] = { + **{ + f"{mode}-mid-turn": StopCase( + run=lambda mode=mode: end_running_turn(mode), + expected=["session 1 prompt", "session 1 abort", "session 1 idle"], + ) + for mode in ENDING_CONTROLS + }, + "next-message-waits-for-the-abort": StopCase( + run=next_message_waits_for_the_abort, + expected=[ + "session 1 prompt", + "session 1 abort", + "session 1 idle", + "session 1 prompt", + ], + ), + "next-message-waits-for-the-aborted-idle": StopCase( + run=next_message_waits_for_the_aborted_idle, + expected=[ + "session 1 prompt", + "session 1 abort", + "session 1 idle", + "session 1 prompt", + ], + ), + "late-tool-call-during-held-abort": StopCase( + run=late_tool_call_during_held_abort, + expected=[f"stopped turn's tools got {BandTool.SEND_MESSAGE}"], + ), + "cleanup-waits-for-the-abort": StopCase( + run=cleanup_waits_for_the_abort, + expected=[ + "session 1 prompt", + "session 1 abort", + "session 1 idle", + "session 1 disconnect", + ], + ), + "shutdown-right-after-stop": StopCase( + run=shutdown_right_after_stop, + expected=[ + "session 1 prompt", + "session 1 abort", + "session 1 idle", + "session 1 disconnect", + ], + ), + "idle-room": StopCase(run=idle, expected=[]), +} diff --git a/tests/framework_conformance/stopprobes/letta.py b/tests/framework_conformance/stopprobes/letta.py new file mode 100644 index 000000000..bc559df7a --- /dev/null +++ b/tests/framework_conformance/stopprobes/letta.py @@ -0,0 +1,255 @@ +"""Backend-stop probes for the letta adapter; see ``stopprobes``.""" + +from __future__ import annotations + +import asyncio +import logging +from typing import Any + +import band_sdk_core +import looptime +import pytest + +from band.adapters.letta import LettaAdapterConfig +from band.client.streaming import ControlMode +from band.core.protocols import TurnResultAlreadyReported +from band.core.retention import cleanup +from band.runtime.tools import CHAT_ID_FIELD_NAME, BandTool +from band.testing.fake_tools import FakeAgentTools +from tests.adapters.waiting import spin +from tests.framework_conformance.stopprobes import ( + ENDING_CONTROLS, + TURN_TIMEOUT_SECONDS, + StopCase, + cancel_turn, + end_turn, + stop_room, + wait_until, + waiting_on_stop, +) +from tests.framework_conformance.turnprobes import ROOM_ID, turn_input + +AGENT_ID = "agent-1" +CONVERSATION_ID = "conversation-1" +CANCEL_AGENT_RUNS = f"cancel agent {AGENT_ID}" +CANCEL_CONVERSATION_RUNS = f"cancel conversation {CONVERSATION_ID}" +DELETE_AGENT = f"delete agent {AGENT_ID}" +FAILURE_POSTED = "failure posted" +NEXT_TURN_ANSWERED = "next turn answered" +CANCEL_FAILED_WARNING = "cancel failure logged at WARNING" +SHUTDOWN_RETURNED = "shutdown returned" + + +class JournaledTools(FakeAgentTools): + """Room tools that log each failure post into the Letta client's journal.""" + + def __init__(self, journal: list[str]) -> None: + super().__init__() + self.journal = journal + + async def send_failure(self, failure: band_sdk_core.AgentFailure) -> Any: + self.journal.append(FAILURE_POSTED) + return await super().send_failure(failure) + + +class WarningRecords(logging.Handler): + """Collects the WARNING records of the logger it is attached to.""" + + def __init__(self) -> None: + super().__init__(level=logging.WARNING) + self.records: list[logging.LogRecord] = [] + + def emit(self, record: logging.LogRecord) -> None: + self.records.append(record) + + +def held_adapter( + *, + shared: bool = False, + timeout_s: float | None = None, + delete_agents: bool = False, +) -> Any: + from tests.adapters.lettakit import ( # noqa: PLC0415 -- letta extra, absent from some lane venvs + RecordingLettaClient, + ready_letta_adapter, + ) + + settings: dict[str, Any] = {"mode": "shared"} if shared else {} + if timeout_s is not None: + settings["turn_timeout_s"] = timeout_s + if delete_agents: + settings["delete_agents_on_cleanup"] = True + adapter, _ = ready_letta_adapter(LettaAdapterConfig(**settings), room_id=ROOM_ID) + client = RecordingLettaClient() + client.hold() + adapter._client = client + room = adapter._rooms[ROOM_ID] + room.agent_id = AGENT_ID + if shared: + room.conversation_id = CONVERSATION_ID + return adapter, client + + +async def end_running_turn(mode: ControlMode, *, shared: bool = False) -> list[str]: + adapter, client = held_adapter(shared=shared) + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: bool(client.sent)) + await end_turn(adapter, turn, ROOM_ID, mode) + return client.cancels + + +async def server_rejects_cancel() -> list[str]: + """Without Redis the cancel fails, is logged, and the room takes its next turn.""" + adapter, client = held_adapter() + client.reject_cancels = True + turn = asyncio.create_task(adapter.on_event(turn_input(FakeAgentTools()))) + await wait_until(lambda: bool(client.sent)) + warnings = WarningRecords() + cleanup.logger.addHandler(warnings) + try: + await end_turn(adapter, turn, ROOM_ID, ControlMode.STOP) + finally: + cleanup.logger.removeHandler(warnings) + stop_failures = [ + record + for record in warnings.records + if f"backend stop for room {ROOM_ID}" in record.getMessage() + ] + client.release.set() + next_tools = FakeAgentTools() + await adapter.on_event(turn_input(next_tools)) + return [ + *client.cancels, + *([CANCEL_FAILED_WARNING] if stop_failures else []), + *([NEXT_TURN_ANSWERED] if next_tools.messages_sent else []), + ] + + +async def stop_holding_cancel( + adapter: Any, client: Any, tools: FakeAgentTools +) -> asyncio.Event: + """Stop the running turn and leave its run cancel in flight; the returned + event lets that cancel finish.""" + release_cancel = client.hold_cancel() + turn = asyncio.create_task(adapter.on_event(turn_input(tools))) + await wait_until(lambda: bool(client.sent)) + await cancel_turn(turn) + await adapter.on_interrupt(ROOM_ID, ControlMode.STOP) + await wait_until(client.cancel_started.is_set) + return release_cancel + + +def effects(client: Any) -> list[str]: + """What the server was asked to do besides take the turn.""" + return [entry for entry in client.journal if entry not in client.sent] + + +async def late_mcp_call_during_held_cancel() -> list[str]: + """A tool call from the stopped run that lands while its cancel is pending + posts on that run's tools, not on the replacement turn's.""" + from tests.adapters.lettakit import ( # noqa: PLC0415 -- letta extra, absent from some lane venvs + band_mcp_registrations, + ) + + adapter, client = held_adapter() + stopped, replacement = FakeAgentTools(), FakeAgentTools() + release_cancel = await stop_holding_cancel(adapter, client, stopped) + next_turn = asyncio.create_task(adapter.on_event(turn_input(replacement))) + # The turn blocks inside the context manager that owns the wait, where + # ``waiting_on_stop`` cannot see it; a turn that ignored the cancel would + # publish its tools within these loop turns. + await spin() + await band_mcp_registrations(adapter)[BandTool.SEND_MESSAGE].execute( + { + "content": "from the stopped run", + "mentions": ["@alice"], + CHAT_ID_FIELD_NAME: ROOM_ID, + } + ) + landed = [ + *(f"stopped turn's tools got {m['content']}" for m in stopped.messages_sent), + *( + f"replacement turn's tools got {m['content']}" + for m in replacement.messages_sent + ), + ] + release_cancel.set() + client.release.set() + await next_turn + return landed + + +async def cleanup_waits_for_the_cancel() -> list[str]: + """Leaving the room waits for the run cancel before it deletes the agent.""" + adapter, client = held_adapter(delete_agents=True) + release_cancel = await stop_holding_cancel(adapter, client, FakeAgentTools()) + cleanup_room = asyncio.create_task(adapter.on_cleanup(ROOM_ID)) + await wait_until(lambda: waiting_on_stop(cleanup_room) or cleanup_room.done()) + release_cancel.set() + await cleanup_room + return effects(client) + + +async def shutdown_waits_for_the_cancel() -> list[str]: + """Shutdown does not return while a run cancel is still in flight.""" + adapter, client = held_adapter() + release_cancel = await stop_holding_cancel(adapter, client, FakeAgentTools()) + + async def shut_down() -> None: + await adapter.cleanup_all() + client.journal.append(SHUTDOWN_RETURNED) + + shutdown = asyncio.create_task(shut_down()) + await wait_until(lambda: waiting_on_stop(shutdown) or shutdown.done()) + release_cancel.set() + await shutdown + return effects(client) + + +@looptime.enabled(strict=True) +async def timeout() -> list[str]: + """The cancel is sent once, before the timeout is reported.""" + adapter, client = held_adapter(timeout_s=TURN_TIMEOUT_SECONDS) + tools = JournaledTools(client.journal) + with pytest.raises(TurnResultAlreadyReported): + await adapter.on_event(turn_input(tools)) + await adapter.backend_turns.drain() + return effects(client) + + +async def idle() -> list[str]: + adapter, client = held_adapter() + await stop_room(adapter, ROOM_ID) + return client.cancels + + +CASES: dict[str, StopCase] = { + **{ + f"{mode}-mid-turn-per-room": StopCase( + run=lambda mode=mode: end_running_turn(mode), + expected=[CANCEL_AGENT_RUNS], + ) + for mode in ENDING_CONTROLS + }, + "stop-mid-turn-shared": StopCase( + run=lambda: end_running_turn(ControlMode.STOP, shared=True), + expected=[CANCEL_CONVERSATION_RUNS], + ), + "server-rejects-cancel": StopCase( + run=server_rejects_cancel, + expected=[CANCEL_AGENT_RUNS, CANCEL_FAILED_WARNING, NEXT_TURN_ANSWERED], + ), + "late-mcp-call-during-held-cancel": StopCase( + run=late_mcp_call_during_held_cancel, + expected=["stopped turn's tools got from the stopped run"], + ), + "cleanup-waits-for-the-cancel": StopCase( + run=cleanup_waits_for_the_cancel, expected=[CANCEL_AGENT_RUNS, DELETE_AGENT] + ), + "shutdown-waits-for-the-cancel": StopCase( + run=shutdown_waits_for_the_cancel, + expected=[CANCEL_AGENT_RUNS, SHUTDOWN_RETURNED], + ), + "timeout": StopCase(run=timeout, expected=[CANCEL_AGENT_RUNS, FAILURE_POSTED]), + "idle-room": StopCase(run=idle, expected=[]), +} diff --git a/tests/framework_conformance/stopprobes/opencode.py b/tests/framework_conformance/stopprobes/opencode.py new file mode 100644 index 000000000..190aee030 --- /dev/null +++ b/tests/framework_conformance/stopprobes/opencode.py @@ -0,0 +1,372 @@ +"""Backend-stop probes for the opencode adapter; see ``stopprobes``.""" + +from __future__ import annotations + +import asyncio +from typing import Any + +import looptime + +from band.adapters.opencode import OpencodeAdapter, OpencodeAdapterConfig +from band.client.streaming import ControlMode +from band.core.protocols import FAILURE_CODE_TIMEOUT +from band.integrations.opencode.types import OpencodeSessionState +from band.runtime.tools import BandTool +from band.testing import FakeAgentTools, reported_failures +from tests.adapters.dispatch import deliver_message +from tests.adapters.opencode.helpers import ( + BandMCPCalls, + FakeOpencodeClient, + event_message_updated, + event_permission, + event_session_idle, + event_text_part, + make_platform_message, + tools_protocol, +) +from tests.adapters.waiting import spin +from tests.framework_conformance.stopprobes import ( + ENDING_CONTROLS, + STARTUP_TIMEOUT_SECONDS, + TURN_TIMEOUT_SECONDS, + StopCase, + cancel_turn, + end_turn, + stop_room, + wait_until, + waiting_on_stop, +) +from tests.framework_conformance.turnprobes import ROOM_ID +from tests.mcpbackends import backends_created_by + +# The fake server numbers its first session. +SESSION_ID = "sess-1" +PROMPT_FIRST_SESSION = f"prompt {SESSION_ID}" +ABORT_FIRST_SESSION = f"abort {SESSION_ID}" +CLIENT_CLOSED = "client closed" +TIMEOUT_REPORTED = "timeout reported" +# What the server streams once an aborted turn winds down; a watcher still +# running would post it as the turn's fallback reply. +LATE_TEXT = "text from the stopped turn" +LATE_CALL = "call from the stopped turn" + + +class JournalingOpencodeClient(FakeOpencodeClient): + """Records session prompts, aborts and its own close, in the order they happen.""" + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) # type: ignore[arg-type] + self.journal: list[str] = [] + + async def prompt_async(self, session_id: str, **kwargs: Any) -> None: + self.journal.append(f"prompt {session_id}") + await super().prompt_async(session_id, **kwargs) + + async def abort_session(self, session_id: str) -> None: + await super().abort_session(session_id) + self.journal.append(f"abort {session_id}") + + async def close(self) -> None: + self.journal.append(CLIENT_CLOSED) + await super().close() + + +class HoldingAbortClient(JournalingOpencodeClient): + """An abort that waits until released, and records only one that finished.""" + + def __init__(self, **kwargs: object) -> None: + super().__init__(**kwargs) # type: ignore[arg-type] + self.entered = asyncio.Event() + self.release = asyncio.Event() + self.completed: list[str] = [] + + async def abort_session(self, session_id: str) -> None: + self.entered.set() + await self.release.wait() + self.completed.append(session_id) + await super().abort_session(session_id) + + +async def started( + client: JournalingOpencodeClient, config: OpencodeAdapterConfig | None = None +) -> OpencodeAdapter: + adapter = OpencodeAdapter( + config=config or OpencodeAdapterConfig(), client_factory=lambda _: client + ) + await adapter.on_started("OpenCode Agent", "A coding agent") + return adapter + + +def deliver(adapter: OpencodeAdapter, tools: FakeAgentTools) -> asyncio.Task[None]: + return asyncio.create_task( + deliver_message( + adapter, + make_platform_message(room_id=ROOM_ID, content="long task"), + tools_protocol(tools), + OpencodeSessionState(), + participants_msg=None, + contacts_msg=None, + is_session_bootstrap=True, + room_id=ROOM_ID, + ) + ) + + +async def stopped_turn_winds_down(client: JournalingOpencodeClient) -> None: + """The server streams the stopped turn's last text and goes idle.""" + for event in ( + event_message_updated(SESSION_ID, "msg-late"), + event_text_part(SESSION_ID, "msg-late", LATE_TEXT), + event_session_idle(SESSION_ID), + ): + await client.push_event(event) + consumed = asyncio.Event() + + async def all_consumed() -> None: + consumed.set() + + # The event loop reaches this step only after handling every event above. + await client.push_event(all_consumed) + async with asyncio.timeout(STARTUP_TIMEOUT_SECONDS): + await consumed.wait() + # A watcher that was not cancelled resumes on the idle event and needs + # more loop turns to relay the late text; the turn exposes no state to + # wait on once it is stopped. + await spin() + + +def posted(tools: FakeAgentTools, before: int) -> list[str]: + return [f"posted {message['content']}" for message in tools.messages_sent[before:]] + + +async def end_running_turn(mode: ControlMode) -> list[str]: + client = JournalingOpencodeClient(prompt_event_sequences=[[]]) + adapter = await started(client) + tools = FakeAgentTools() + turn = deliver(adapter, tools) + await wait_until(lambda: bool(client.prompt_calls)) + await end_turn(adapter, turn, ROOM_ID, mode) + await stopped_turn_winds_down(client) + recorded = client.journal + posted(tools, 0) + await adapter.on_cleanup(ROOM_ID) + return recorded + + +async def stop_detached_after_approval() -> list[str]: + """A manual ask releases on_message; a later STOP still stops the turn.""" + client = JournalingOpencodeClient( + prompt_event_sequences=[[event_permission(SESSION_ID, "perm-1")]] + ) + adapter = await started(client, OpencodeAdapterConfig(approval_mode="manual")) + tools = FakeAgentTools() + await deliver(adapter, tools) + asked = len(tools.messages_sent) + await stop_room(adapter, ROOM_ID) + await stopped_turn_winds_down(client) + recorded = client.journal + posted(tools, asked) + await adapter.on_cleanup(ROOM_ID) + return recorded + + +async def stop_holding_abort( + adapter: OpencodeAdapter, client: HoldingAbortClient, tools: FakeAgentTools +) -> None: + """Stop the running turn and leave its session's abort in flight.""" + turn = deliver(adapter, tools) + await wait_until(lambda: bool(client.prompt_calls)) + await cancel_turn(turn) + await adapter.on_interrupt(ROOM_ID, ControlMode.STOP) + await wait_until(client.entered.is_set) + + +async def wind_down(adapter: OpencodeAdapter, turn: asyncio.Task[None]) -> None: + await cancel_turn(turn) + await adapter.backend_turns.drain() + await adapter.on_cleanup(ROOM_ID) + + +async def next_message_waits_for_the_abort() -> list[str]: + """The next message prompts the session only after its abort has finished.""" + client = HoldingAbortClient(prompt_event_sequences=[[], []]) + adapter = await started(client) + await stop_holding_abort(adapter, client, FakeAgentTools()) + next_turn = deliver(adapter, FakeAgentTools()) + await wait_until(lambda: waiting_on_stop(next_turn) or len(client.prompt_calls) > 1) + client.release.set() + await wait_until(lambda: len(client.prompt_calls) == 2) + recorded = list(client.journal) + await wind_down(adapter, next_turn) + return recorded + + +async def late_mcp_call_during_held_abort() -> list[str]: + """A tool call from the stopped turn that lands while its abort is pending + posts on that turn's tools, not on the replacement turn's.""" + with backends_created_by() as starts: + client = HoldingAbortClient(prompt_event_sequences=[[], []]) + adapter = await started(client) + stopped, replacement = FakeAgentTools(), FakeAgentTools() + await stop_holding_abort(adapter, client, stopped) + next_turn = deliver(adapter, replacement) + await wait_until( + lambda: waiting_on_stop(next_turn) or len(client.prompt_calls) > 1 + ) + await BandMCPCalls(starts).band_tool_call( + BandTool.SEND_MESSAGE, + {"content": LATE_CALL, "mentions": ["@alice"]}, + room_id=ROOM_ID, + )() + landed = [ + *( + f"stopped turn's tools got {m['content']}" + for m in stopped.messages_sent + ), + *( + f"replacement turn's tools got {m['content']}" + for m in replacement.messages_sent + ), + ] + client.release.set() + await wait_until(lambda: len(client.prompt_calls) == 2) + await wind_down(adapter, next_turn) + return landed + + +async def cleanup_waits_for_the_abort() -> list[str]: + """Teardown right after STOP waits for the session abort before it closes + the client.""" + client = HoldingAbortClient(prompt_event_sequences=[[]]) + adapter = await started(client) + await stop_holding_abort(adapter, client, FakeAgentTools()) + cleanup = asyncio.create_task(adapter.on_cleanup(ROOM_ID)) + await wait_until(lambda: waiting_on_stop(cleanup) or cleanup.done()) + client.release.set() + await cleanup + return client.journal + + +def reported_timeout(tools: FakeAgentTools) -> list[str]: + return [ + TIMEOUT_REPORTED + for failure in reported_failures(tools) + if failure["code"] == FAILURE_CODE_TIMEOUT + ] + + +@looptime.enabled(strict=True) +async def timeout_reports() -> list[str]: + """A timed-out turn aborts the session and still posts the timeout.""" + client = JournalingOpencodeClient(prompt_event_sequences=[[]]) + adapter = await started( + client, OpencodeAdapterConfig(turn_timeout_s=TURN_TIMEOUT_SECONDS) + ) + tools = FakeAgentTools() + await deliver(adapter, tools) + recorded = client.journal + reported_timeout(tools) + await adapter.on_cleanup(ROOM_ID) + return recorded + + +@looptime.enabled(strict=True) +async def stop_during_timeout_abort(mode: ControlMode) -> list[str]: + """A STOP while the timeout abort is in flight still finishes that abort.""" + client = HoldingAbortClient(prompt_event_sequences=[[]]) + adapter = await started( + client, OpencodeAdapterConfig(turn_timeout_s=TURN_TIMEOUT_SECONDS) + ) + turn = deliver(adapter, FakeAgentTools()) + await client.entered.wait() + await cancel_turn(turn) + await adapter.on_interrupt(ROOM_ID, mode) + client.release.set() + await adapter.backend_turns.drain() + recorded = list(client.completed) + await adapter.on_cleanup(ROOM_ID) + return recorded + + +async def approve(adapter: OpencodeAdapter) -> None: + """The human approves the pending ask, as a room message the control hook takes.""" + message = make_platform_message(room_id=ROOM_ID, content="approve req-1") + assert await adapter.try_handle_control(message, tools_protocol(FakeAgentTools())) + + +@looptime.enabled(strict=True) +async def stop_released_turn_during_timeout_abort(mode: ControlMode) -> list[str]: + """A released turn that timed out has already ended on its own: a STOP during + its timeout abort only finishes the abort, and the timeout is still reported.""" + client = HoldingAbortClient( + prompt_event_sequences=[[event_permission(SESSION_ID, "req-1")]] + ) + adapter = await started( + client, + OpencodeAdapterConfig( + approval_mode="manual", turn_timeout_s=TURN_TIMEOUT_SECONDS + ), + ) + tools = FakeAgentTools() + await deliver(adapter, tools) + await approve(adapter) + await client.entered.wait() + await adapter.on_interrupt(ROOM_ID, mode) + client.release.set() + await adapter.backend_turns.drain() + await wait_until(lambda: bool(reported_timeout(tools))) + recorded = reported_timeout(tools) + await adapter.on_cleanup(ROOM_ID) + return recorded + + +async def idle() -> list[str]: + client = JournalingOpencodeClient() + adapter = await started(client) + await stop_room(adapter, ROOM_ID) + recorded = list(client.journal) + await adapter.on_cleanup(ROOM_ID) + return recorded + + +CASES: dict[str, StopCase] = { + **{ + f"{mode}-mid-turn": StopCase( + run=lambda mode=mode: end_running_turn(mode), + expected=[PROMPT_FIRST_SESSION, ABORT_FIRST_SESSION], + ) + for mode in ENDING_CONTROLS + }, + "stop-detached-after-approval": StopCase( + run=stop_detached_after_approval, + expected=[PROMPT_FIRST_SESSION, ABORT_FIRST_SESSION], + ), + "next-message-waits-for-the-abort": StopCase( + run=next_message_waits_for_the_abort, + expected=[PROMPT_FIRST_SESSION, ABORT_FIRST_SESSION, PROMPT_FIRST_SESSION], + ), + "late-mcp-call-during-held-abort": StopCase( + run=late_mcp_call_during_held_abort, + expected=[f"stopped turn's tools got {LATE_CALL}"], + ), + "cleanup-waits-for-the-abort": StopCase( + run=cleanup_waits_for_the_abort, + expected=[PROMPT_FIRST_SESSION, ABORT_FIRST_SESSION, CLIENT_CLOSED], + ), + "timeout": StopCase( + run=timeout_reports, + expected=[PROMPT_FIRST_SESSION, ABORT_FIRST_SESSION, TIMEOUT_REPORTED], + ), + **{ + f"{mode}-during-timeout-abort": StopCase( + run=lambda mode=mode: stop_during_timeout_abort(mode), + expected=[SESSION_ID], + ) + for mode in ENDING_CONTROLS + }, + **{ + f"{mode}-released-turn-during-timeout-abort": StopCase( + run=lambda mode=mode: stop_released_turn_during_timeout_abort(mode), + expected=[TIMEOUT_REPORTED], + ) + for mode in ENDING_CONTROLS + }, + "idle-room": StopCase(run=idle, expected=[]), +} diff --git a/tests/framework_conformance/test_backend_stop.py b/tests/framework_conformance/test_backend_stop.py index 9d898ed31..7c4efe781 100644 --- a/tests/framework_conformance/test_backend_stop.py +++ b/tests/framework_conformance/test_backend_stop.py @@ -7,17 +7,25 @@ from band.core.simple_adapter import SimpleAdapter from tests.framework_conformance.outcomeprobes import band_adapters from tests.framework_conformance.stopprobes import StopCase +from tests.framework_conformance.stopprobes import a2a as a2a_probe from tests.framework_conformance.stopprobes import acp as acp_probe from tests.framework_conformance.stopprobes import claudesdk as claudesdk_probe from tests.framework_conformance.stopprobes import codex as codex_probe +from tests.framework_conformance.stopprobes import copilotsdk as copilotsdk_probe from tests.framework_conformance.stopprobes import cursoracp as cursoracp_probe +from tests.framework_conformance.stopprobes import letta as letta_probe +from tests.framework_conformance.stopprobes import opencode as opencode_probe #: The stop cases of each adapter that drives a backend, by class name. STOP_PROBES: dict[str, dict[str, StopCase]] = { + "A2AAdapter": a2a_probe.CASES, "ACPClientAdapter": acp_probe.CASES, "ClaudeSDKAdapter": claudesdk_probe.CASES, "CodexAdapter": codex_probe.CASES, + "CopilotSDKAdapter": copilotsdk_probe.CASES, "CursorACPAdapter": cursoracp_probe.CASES, + "LettaAdapter": letta_probe.CASES, + "OpencodeAdapter": opencode_probe.CASES, } #: Subclasses that keep a probed parent's ``backend_turns`` and ``on_message``, @@ -32,7 +40,9 @@ REQUIRED_MODULES: dict[str, str] = { "ACPClientAdapter": "acp", "ClaudeSDKAdapter": "claude_agent_sdk", + "CopilotSDKAdapter": "copilot", "CursorACPAdapter": "acp", + "LettaAdapter": "letta_client", } CASES = [ diff --git a/tests/integrations/a2a/kit.py b/tests/integrations/a2a/kit.py new file mode 100644 index 000000000..2a4de0f61 --- /dev/null +++ b/tests/integrations/a2a/kit.py @@ -0,0 +1,117 @@ +"""A remote A2A agent for adapter tests: scripted streams, recorded cancels.""" + +from __future__ import annotations + +import asyncio +from collections import deque +from collections.abc import AsyncIterator +from dataclasses import dataclass +from typing import Any + +from a2a.helpers import new_text_message +from a2a.types import ( + CancelTaskRequest, + Role, + SendMessageRequest, + StreamResponse, + Task, + TaskState, + TaskStatus, +) + +CANCELLED = "tasks/cancel" +WORKING = TaskState.Name(TaskState.TASK_STATE_WORKING) +CANCELED = TaskState.Name(TaskState.TASK_STATE_CANCELED) + + +def task_update( + task_id: str, + state_name: str, + *, + context_id: str = "ctx-1", + text: str | None = None, +) -> StreamResponse: + """A stream event reporting ``task_id`` in the named ``TaskState``, with + ``text`` as the agent's status message.""" + status = TaskStatus( + state=TaskState.Value(state_name), + message=new_text_message(text, role=Role.ROLE_AGENT) if text else None, + ) + return StreamResponse(task=Task(id=task_id, context_id=context_id, status=status)) + + +@dataclass(frozen=True) +class ScriptedStream: + """One reply stream: these events, then either the end or a held turn. + + A ``gate`` keeps the remote silent after it accepted the request, until + the gate is set. + """ + + events: tuple[StreamResponse, ...] + hold: bool = False + gate: asyncio.Event | None = None + + +class RecordingA2AClient: + """An A2A client whose remote agent follows scripted streams. + + It records the task each request continues and each task it is asked to + cancel. + """ + + def __init__(self) -> None: + self.continued: list[str] = [] + self.cancelled: list[str] = [] + self.closed = False + self.cancel_started = asyncio.Event() + self._cancel_release: asyncio.Event | None = None + self._cancelled_ids: set[str] = set() + self._streams: deque[ScriptedStream] = deque() + + def hold_cancel(self) -> asyncio.Event: + """``cancel_task`` waits for the returned event after it has started.""" + self._cancel_release = asyncio.Event() + return self._cancel_release + + def script( + self, + *events: StreamResponse, + hold: bool = False, + gate: asyncio.Event | None = None, + ) -> None: + self._streams.append(ScriptedStream(events, hold, gate)) + + async def send_message( + self, request: SendMessageRequest, *, context: Any = None + ) -> AsyncIterator[StreamResponse]: + self.continued.append(request.message.task_id) + stream = self._streams.popleft() + if stream.gate is not None: + await stream.gate.wait() + for event in stream.events: + yield event + if stream.hold: + await asyncio.Future[None]() + + async def cancel_task( + self, request: CancelTaskRequest, *, context: Any = None + ) -> Task: + self.cancel_started.set() + if self._cancel_release is not None: + await self._cancel_release.wait() + self._cancelled_ids.add(request.id) + self.cancelled.append(f"{CANCELLED} {request.id}") + return Task( + id=request.id, status=TaskStatus(state=TaskState.TASK_STATE_CANCELED) + ) + + async def subscribe( + self, request: Any, *, context: Any = None + ) -> AsyncIterator[StreamResponse]: + """The task stays non-terminal until its cancel has finished.""" + state = CANCELED if request.id in self._cancelled_ids else WORKING + yield task_update(request.id, state) + + async def close(self) -> None: + self.closed = True diff --git a/uv.lock b/uv.lock index e813e3c6e..9a9e11257 100644 --- a/uv.lock +++ b/uv.lock @@ -1052,8 +1052,8 @@ requires-dist = [ { name = "langgraph", marker = "extra == 'dev'", specifier = ">=1.0.0" }, { name = "langgraph", marker = "extra == 'dev-strands-nooa'", specifier = ">=1.0.0" }, { name = "langgraph", marker = "extra == 'langgraph'", specifier = ">=1.0.0" }, - { name = "letta-client", marker = "extra == 'dev'", specifier = ">=1.0.0" }, - { name = "letta-client", marker = "extra == 'letta'", specifier = ">=1.0.0" }, + { name = "letta-client", marker = "extra == 'dev'", specifier = ">=1.7.1" }, + { name = "letta-client", marker = "extra == 'letta'", specifier = ">=1.7.1" }, { name = "looptime", marker = "extra == 'dev'", specifier = ">=0.7" }, { name = "looptime", marker = "extra == 'dev-crewai'", specifier = ">=0.7" }, { name = "looptime", marker = "extra == 'dev-parlant'", specifier = ">=0.7" },