diff --git a/src/band/integrations/claude_sdk/dedup_tools.py b/src/band/integrations/claude_sdk/dedup_tools.py index 59d2c89ea..ca069dd78 100644 --- a/src/band/integrations/claude_sdk/dedup_tools.py +++ b/src/band/integrations/claude_sdk/dedup_tools.py @@ -13,9 +13,9 @@ chat). The platform happily accepts every POST, so the duplicate becomes a visible -chat message and a charged LLM call. ``ExecutionContext._processed_ids`` only -dedupes *inbound* user messages from the platform; it has no view of outbound -tool calls produced by Claude. +chat message and a charged LLM call. The runtime's ``MessageClaimRegistry`` +only dedupes *inbound* user messages from the platform; it has no view of +outbound tool calls produced by Claude. This wrapper sits in front of the per-room ``AgentToolsProtocol`` that the ``ClaudeSDKAdapter`` registers in ``_room_tools[room_id]``. It dedupes diff --git a/src/band/runtime/claims.py b/src/band/runtime/claims.py new file mode 100644 index 000000000..49e8e6257 --- /dev/null +++ b/src/band/runtime/claims.py @@ -0,0 +1,128 @@ +"""Shared message-claim ledger for the inbound delivery lifecycle.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterator +from contextlib import contextmanager +from typing import TypeAlias + +DEFAULT_COMPLETED_CACHE_SIZE = 500 +ClaimKey: TypeAlias = tuple[str, str] + + +class MessageClaimRegistry: + """Ownership ledger ensuring one execution per inbound message ID. + + Lifecycle states: + + - **in flight** — claimed by exactly one execution via ``try_claim``; + released on failure or cancellation, never evicted. + - **ack pending** — the handler completed but the durable processed ack + failed; redelivery retries only the ack, never the handler. Never + evicted (losing it would replay side effects); drains through the ack + retry budget instead. If a room never returns, this state intentionally + remains for the registry's lifetime rather than risk replaying effects. + - **completed** — durably processed; kept in a bounded LRU for dedup. + + ``AgentRuntime`` owns one registry and passes it to every context it + creates, so recreated contexts retain ownership state. Claims are keyed by + room and message ID, preventing work in one room from affecting another. + A context constructed standalone gets a private registry, preserving the + previous per-context behavior. + + ``try_claim`` checks and inserts without an ``await``, so event-loop + scheduling makes claims atomic within one runtime. + + Scope, honestly stated: this coordinates executions within one runtime + (one event loop). It cannot coordinate separate processes, containers, + or hosts — the platform re-serves in-flight messages to fresh actors and + ``mark_processing`` is not an exclusive claim, so deployments that shard + one agent id across such boundaries need a platform-level claim (see + ``band.runtime.single_instance`` for the same boundary statement). + """ + + def __init__(self, max_completed: int = DEFAULT_COMPLETED_CACHE_SIZE) -> None: + self.max_completed = max_completed + self._inflight: set[ClaimKey] = set() + self._ack_pending: OrderedDict[ClaimKey, bool] = OrderedDict() + self._ack_retries: dict[ClaimKey, int] = {} + self._completed_by_room: dict[str, OrderedDict[str, bool]] = {} + + def _try_claim(self, room_id: str, message_id: str) -> bool: + """Claim a message for one execution; False if another owner holds it.""" + key = (room_id, message_id) + if key in self._inflight: + return False + self._inflight.add(key) + return True + + def _release(self, room_id: str, message_id: str) -> None: + """Release an in-flight claim (failure, cancellation, or completion).""" + self._inflight.discard((room_id, message_id)) + + @contextmanager + def claim(self, room_id: str, message_id: str) -> Iterator[bool]: + """Yield whether a claim was acquired and always release acquired claims. + + Exceptions propagate to the delivery path's existing error handling; + this context only guarantees cleanup. + """ + acquired = self._try_claim(room_id, message_id) + try: + yield acquired + finally: + if acquired: + self._release(room_id, message_id) + + def inflight_ids(self, room_id: str) -> set[str]: + """Currently claimed message IDs for a room (copy).""" + return {message_id for room, message_id in self._inflight if room == room_id} + + def is_completed(self, room_id: str, message_id: str) -> bool: + """Whether the message already completed; a hit refreshes LRU recency.""" + completed = self._completed_by_room.get(room_id) + if completed is None or message_id not in completed: + return False + completed.move_to_end(message_id) + return True + + def completed_ids(self, room_id: str) -> list[str]: + """Completed message IDs for a room, oldest first.""" + return list(self._completed_by_room.get(room_id, ())) + + def discard_completed(self, room_id: str) -> None: + """Discard a removed room's durably acknowledged local cache.""" + self._completed_by_room.pop(room_id, None) + + def remember_completed(self, room_id: str, message_id: str) -> None: + """Record durable completion and clear any pending-ack state.""" + completed = self._completed_by_room.setdefault(room_id, OrderedDict()) + completed[message_id] = True + completed.move_to_end(message_id) + key = (room_id, message_id) + self._ack_pending.pop(key, None) + self._ack_retries.pop(key, None) + if len(completed) > self.max_completed: + completed.popitem(last=False) + + def is_ack_pending(self, room_id: str, message_id: str) -> bool: + """Whether the message completed locally but lacks a durable ack.""" + return (room_id, message_id) in self._ack_pending + + def remember_ack_pending(self, room_id: str, message_id: str) -> None: + """Record local completion awaiting a durable processed ack.""" + key = (room_id, message_id) + self._ack_pending[key] = True + self._ack_retries.setdefault(key, 0) + + def pending_ack_ids(self, room_id: str) -> list[str]: + """Message IDs awaiting a durable processed ack, oldest first.""" + return [message_id for room, message_id in self._ack_pending if room == room_id] + + def record_ack_retry(self, room_id: str, message_id: str) -> int: + """Count a failed ack retry; returns the total for budget checks.""" + key = (room_id, message_id) + retries = self._ack_retries.get(key, 0) + 1 + self._ack_retries[key] = retries + return retries diff --git a/src/band/runtime/execution.py b/src/band/runtime/execution.py index c825d4655..7e97c5d70 100644 --- a/src/band/runtime/execution.py +++ b/src/band/runtime/execution.py @@ -16,15 +16,13 @@ import asyncio import logging -from collections import OrderedDict from datetime import datetime, timezone -from enum import Enum +from enum import Enum, StrEnum from typing import ( TYPE_CHECKING, Any, Awaitable, Callable, - Literal, Protocol, runtime_checkable, ) @@ -48,9 +46,10 @@ SYNTHETIC_SENDER_TYPE, SYNTHETIC_CONTACT_EVENTS_SENDER_ID, ) -from .retry_tracker import MessageRetryTracker +from band.runtime.claims import MessageClaimRegistry from band.runtime.context_serialization import context_item_to_dict -from .working_state import WorkingStateReporter +from band.runtime.retry_tracker import MessageRetryTracker +from band.runtime.working_state import WorkingStateReporter if TYPE_CHECKING: from band.platform.link import BandLink @@ -73,6 +72,14 @@ class _BacklogProcessResult(Enum): RETRY_LATER = "retry_later" +class ExecutionState(StrEnum): + """Lifecycle state for one room execution.""" + + STARTING = "starting" + IDLE = "idle" + PROCESSING = "processing" + + def _error_label(e: Exception) -> str: """Return a non-empty label for an exception, falling back to the class name.""" return str(e).strip() or type(e).__name__ @@ -189,6 +196,7 @@ def __init__( on_participant_removed: ParticipantRemovedCallback | None = None, *, hub_room_id: str | None = None, + claim_registry: MessageClaimRegistry | None = None, ): """ Initialize execution context for a specific room. @@ -204,6 +212,10 @@ def __init__( hub_room_id: Optional hub-room ID. Forwarded to AgentTools so the schema methods can auto-enable contact tools when this context belongs to the hub room. + claim_registry: Optional shared message-claim registry. AgentRuntime + passes one registry to its default contexts so a room/message + pair executes at most once per runtime. Defaults to a private + instance for standalone contexts. """ self.room_id = room_id self.link = link @@ -228,7 +240,7 @@ def __init__( # Per-room state self.queue: asyncio.Queue[PlatformEvent] = asyncio.Queue() - self.state: Literal["starting", "idle", "processing"] = "starting" + self.state = ExecutionState.STARTING self._is_running = False self._process_loop_task: asyncio.Task[None] | None = None self._context_cache: ConversationContext | None = None @@ -242,19 +254,10 @@ def __init__( # LLM context tracking self._llm_initialized = False - # Dedupe cache (LRU for detecting duplicates during sync) - self._processed_ids: OrderedDict[str, bool] = OrderedDict() - self._max_processed_ids: int = 500 - - # Local in-flight guard shared by /next and WebSocket processing. - # This prevents the same message ID from executing twice before the - # durable processed status becomes visible to both paths. - self._inflight_message_ids: set[str] = set() - - # Messages whose local handler completed but whose durable processed ack - # failed. Redelivery should retry only the ack, not the side effects. - self._processed_ack_pending_ids: OrderedDict[str, bool] = OrderedDict() - self._processed_ack_retry_counts: dict[str, int] = {} + # Message ownership ledger (in-flight claims, completed LRU, pending + # acks) shared by /next and WebSocket processing. Runtime-provided so + # all contexts of one agent coordinate; private otherwise. + self.claims = claim_registry or MessageClaimRegistry() # Crash recovery: sync point marker and retry tracking self._first_ws_msg_id: str | None = None # First WS message = sync point @@ -280,9 +283,9 @@ def thread_id(self) -> str: @property def is_processing(self) -> bool: """Check if context is currently processing an event.""" - return self.state == "processing" + return self.state is ExecutionState.PROCESSING - def _set_state(self, new_state: Literal["starting", "idle", "processing"]) -> None: + def _set_state(self, new_state: ExecutionState) -> None: """ Set the execution state and update the idle event accordingly. @@ -290,7 +293,7 @@ def _set_state(self, new_state: Literal["starting", "idle", "processing"]) -> No for graceful shutdown coordination. """ self.state = new_state - if new_state == "processing": + if new_state is ExecutionState.PROCESSING: self._idle_event.clear() else: self._idle_event.set() @@ -374,37 +377,18 @@ def _message_processed_for_agent(self, message_id: str, metadata: Any) -> bool: == DeliveryStatus.PROCESSED ) - def _remember_processed_message(self, message_id: str) -> None: - """Track a processed message ID in the local LRU dedupe cache.""" - self._processed_ids[message_id] = True - self._processed_ids.move_to_end(message_id) - self._processed_ack_pending_ids.pop(message_id, None) - self._processed_ack_retry_counts.pop(message_id, None) - if len(self._processed_ids) > self._max_processed_ids: - self._processed_ids.popitem(last=False) - - def _remember_processed_ack_pending(self, message_id: str) -> None: - """Track local completion while waiting for durable processed ack.""" - self._processed_ack_pending_ids[message_id] = True - self._processed_ack_pending_ids.move_to_end(message_id) - self._processed_ack_retry_counts.setdefault(message_id, 0) - if len(self._processed_ack_pending_ids) > self._max_processed_ids: - self._processed_ack_pending_ids.popitem(last=False) - async def _retry_processed_ack(self, message_id: str) -> bool: """Retry durable processed ack for a locally completed message.""" - if message_id not in self._processed_ack_pending_ids: + if not self.claims.is_ack_pending(self.room_id, message_id): return False - self._processed_ack_pending_ids.move_to_end(message_id) durable_processed = await self.link.mark_processed(self.room_id, message_id) if durable_processed: self._retry_tracker.mark_success(message_id) - self._remember_processed_message(message_id) + self.claims.remember_completed(self.room_id, message_id) return True - retries = self._processed_ack_retry_counts.get(message_id, 0) + 1 - self._processed_ack_retry_counts[message_id] = retries + retries = self.claims.record_ack_retry(self.room_id, message_id) if retries >= self._retry_tracker.max_retries: logger.warning( "ExecutionContext %s: processed ack retry budget exhausted for message %s; keeping local completion marker", @@ -412,22 +396,11 @@ async def _retry_processed_ack(self, message_id: str) -> bool: message_id, ) self._retry_tracker.mark_success(message_id) - self._remember_processed_message(message_id) + self.claims.remember_completed(self.room_id, message_id) return True return False - def _try_claim_local_message(self, message_id: str) -> bool: - """Claim a message ID for local processing before hydration/status writes.""" - if message_id in self._inflight_message_ids: - return False - self._inflight_message_ids.add(message_id) - return True - - def _release_local_message(self, message_id: str) -> None: - """Release a local in-flight message claim.""" - self._inflight_message_ids.discard(message_id) - # --- Execution protocol implementation --- async def start(self) -> None: @@ -897,10 +870,10 @@ async def _process_loop(self) -> None: # If a pending message cannot be claimed yet, stay in startup sync # instead of processing newer WebSocket events out of order. while not await self._synchronize_with_next(): - self._set_state("idle") + self._set_state(ExecutionState.IDLE) await asyncio.sleep(self.config.idle_resync_seconds) - self._set_state("idle") + self._set_state(ExecutionState.IDLE) logger.info( "ExecutionContext %s: Synchronized, switching to WebSocket", self.room_id, @@ -952,7 +925,7 @@ async def _process_loop(self) -> None: async def _retry_pending_processed_acks(self) -> bool: """Retry durable processed acks for locally completed messages.""" - for msg_id in list(self._processed_ack_pending_ids): + for msg_id in self.claims.pending_ack_ids(self.room_id): if not await self._retry_processed_ack(msg_id): return False return True @@ -965,7 +938,7 @@ async def _wait_until_resync_complete(self) -> None: and await self._resync_pending_messages() ): return - self._set_state("idle") + self._set_state(ExecutionState.IDLE) await asyncio.sleep(self.config.idle_resync_seconds) async def _synchronize_with_next(self) -> bool: @@ -1200,22 +1173,28 @@ async def _process_backlog_message( return _BacklogProcessResult.ADVANCED # Skip if already processed (dedupe) - if msg_id in self._processed_ids: - self._processed_ids.move_to_end(msg_id) + if self.claims.is_completed(self.room_id, msg_id): logger.debug("Skipping duplicate backlog message: %s", msg_id) return _BacklogProcessResult.ADVANCED - if msg_id in self._processed_ack_pending_ids: + if self.claims.is_ack_pending(self.room_id, msg_id): logger.debug("Retrying processed ack for backlog message: %s", msg_id) if await self._retry_processed_ack(msg_id): return _BacklogProcessResult.ADVANCED return _BacklogProcessResult.RETRY_LATER - if not self._try_claim_local_message(msg_id): - logger.debug("Skipping already in-flight backlog message: %s", msg_id) - return _BacklogProcessResult.RETRY_LATER + with self.claims.claim(self.room_id, msg_id) as acquired: + if not acquired: + logger.debug("Deferring in-flight backlog message: %s", msg_id) + return _BacklogProcessResult.RETRY_LATER + return await self._process_claimed_backlog_message(msg) - self._set_state("processing") + async def _process_claimed_backlog_message( + self, msg: PlatformMessage + ) -> _BacklogProcessResult: + """Process a backlog message while its in-flight claim is held.""" + msg_id = msg.id + self._set_state(ExecutionState.PROCESSING) logger.info("Processing backlog message %s in room %s", msg_id, self.room_id) try: @@ -1228,7 +1207,7 @@ async def _process_backlog_message( msg_id, self.room_id, ) - self._remember_processed_message(msg_id) + self.claims.remember_completed(self.room_id, msg_id) return _BacklogProcessResult.ADVANCED # Track attempts - check if exceeded BEFORE processing @@ -1308,9 +1287,9 @@ async def _process_backlog_message( durable_processed = await self.link.mark_processed(self.room_id, msg_id) if durable_processed: self._retry_tracker.mark_success(msg_id) - self._remember_processed_message(msg_id) + self.claims.remember_completed(self.room_id, msg_id) else: - self._remember_processed_ack_pending(msg_id) + self.claims.remember_ack_pending(self.room_id, msg_id) logger.warning( "ExecutionContext %s: Local execution completed but durable processed mark failed for backlog message %s", self.room_id, @@ -1335,8 +1314,7 @@ async def _process_backlog_message( return _BacklogProcessResult.ADVANCED finally: - self._release_local_message(msg_id) - self._set_state("idle") + self._set_state(ExecutionState.IDLE) def _drain_duplicate_from_queue(self, msg_id: str) -> None: """ @@ -1378,7 +1356,7 @@ async def _execute_message_cycle(self, event: PlatformEvent) -> None: ``start()`` emits working:true (+ keep-alive); ``stop()`` in the finally emits the authoritative working:false on success, exception, and cancel — - mirroring the ``_set_state('idle')`` placement. The reporter is a no-op + mirroring the final idle-state transition. The reporter is a no-op when disabled or for the hub room, so call sites need no extra gating. """ await self._working_reporter.start() @@ -1399,24 +1377,22 @@ async def _process_event(self, event: PlatformEvent) -> bool: 5. Mark as processed (success) or failed (exception) """ if isinstance(event, ReconnectedEvent): - self._set_state("processing") + self._set_state(ExecutionState.PROCESSING) logger.debug("Processing %s in room %s", event.type, self.room_id) try: if self._reconnect_sync_requested: self._reconnect_sync_requested = False while not await self._synchronize_with_next(): - self._set_state("idle") + self._set_state(ExecutionState.IDLE) await asyncio.sleep(self.config.idle_resync_seconds) logger.debug("Event %s processed successfully", event.type) finally: - self._set_state("idle") + self._set_state(ExecutionState.IDLE) return True payload = event.payload if isinstance(event, MessageEvent) else None msg_id = payload.id if payload else None - claimed_msg_id: str | None = None - # For messages: check if we should skip if isinstance(event, MessageEvent) and msg_id and payload: # Skip messages from self (agent's own messages) to avoid infinite loops @@ -1446,33 +1422,44 @@ async def _process_event(self, event: PlatformEvent) -> bool: return True # Skip duplicates - if msg_id in self._processed_ids: - self._processed_ids.move_to_end(msg_id) + if self.claims.is_completed(self.room_id, msg_id): logger.debug("Skipping duplicate message %s", msg_id) return True - if msg_id in self._processed_ack_pending_ids: + if self.claims.is_ack_pending(self.room_id, msg_id): logger.debug("Retrying processed ack for message %s", msg_id) if await self._retry_processed_ack(msg_id): return True return False - if not self._try_claim_local_message(msg_id): - logger.debug("Skipping already in-flight message %s", msg_id) - return True - claimed_msg_id = msg_id + with self.claims.claim(self.room_id, msg_id) as acquired: + if not acquired: + # The resync safety net re-checks deferred work, so an + # owner failure never silently loses the message. + logger.debug( + "Message %s owned by another execution; deferring", + msg_id, + ) + return False + return await self._process_event_body(event, msg_id, payload) - if self._message_processed_for_agent(msg_id, payload.metadata): - logger.info( - "Skipping processed replay message %s in room %s", - msg_id, - self.room_id, - ) - self._remember_processed_message(msg_id) - self._release_local_message(msg_id) - return True + return await self._process_event_body(event, msg_id, payload) + + async def _process_event_body( + self, event: PlatformEvent, msg_id: str | None, payload: Any + ) -> bool: + """Process an event after any required in-flight claim is acquired.""" + if isinstance(event, MessageEvent) and msg_id and payload: + if self._message_processed_for_agent(msg_id, payload.metadata): + logger.info( + "Skipping processed replay message %s in room %s", + msg_id, + self.room_id, + ) + self.claims.remember_completed(self.room_id, msg_id) + return True - self._set_state("processing") + self._set_state(ExecutionState.PROCESSING) logger.debug("Processing %s in room %s", event.type, self.room_id) try: @@ -1488,7 +1475,7 @@ async def _process_event(self, event: PlatformEvent) -> bool: msg_id, self.room_id, ) - self._remember_processed_message(msg_id) + self.claims.remember_completed(self.room_id, msg_id) return True # Track attempts @@ -1534,9 +1521,9 @@ async def _process_event(self, event: PlatformEvent) -> bool: durable_processed = await self.link.mark_processed(self.room_id, msg_id) if durable_processed: self._retry_tracker.mark_success(msg_id) - self._remember_processed_message(msg_id) + self.claims.remember_completed(self.room_id, msg_id) else: - self._remember_processed_ack_pending(msg_id) + self.claims.remember_ack_pending(self.room_id, msg_id) logger.warning( "ExecutionContext %s: Local execution completed but durable processed mark failed for message %s", self.room_id, @@ -1562,6 +1549,4 @@ async def _process_event(self, event: PlatformEvent) -> bool: return True finally: - if claimed_msg_id: - self._release_local_message(claimed_msg_id) - self._set_state("idle") + self._set_state(ExecutionState.IDLE) diff --git a/src/band/runtime/runtime.py b/src/band/runtime/runtime.py index 04c5a6383..ab5262862 100644 --- a/src/band/runtime/runtime.py +++ b/src/band/runtime/runtime.py @@ -12,6 +12,7 @@ from band.platform.event import PlatformEvent +from .claims import MessageClaimRegistry from .execution import Execution, ExecutionContext, ExecutionHandler from .presence import RoomPresence from .types import ( @@ -131,6 +132,10 @@ def __init__( # Per-room executions self.executions: dict[str, Execution] = {} + # Shared by default contexts so a room/message pair executes at most + # once per runtime, including across context recreation. + self._claim_registry = MessageClaimRegistry() + # Set up presence callbacks self.presence.on_room_joined = self._on_room_joined self.presence.on_room_left = self._on_room_left @@ -273,6 +278,7 @@ async def _create_execution(self, room_id: str) -> Execution: on_participant_added=self._on_participant_added, on_participant_removed=self._on_participant_removed, hub_room_id=self._hub_room_id, + claim_registry=self._claim_registry, ) self.executions[room_id] = execution @@ -300,6 +306,11 @@ async def _destroy_execution( execution = self.executions.pop(room_id) graceful = await execution.stop(timeout=timeout) + # Durable completion state is safe to release with the room. Pending + # acknowledgements remain in the shared registry so a later rejoin + # retries only the ack instead of replaying handler side effects. + self._claim_registry.discard_completed(room_id) + # Call cleanup callback (for adapter to clean up checkpointer, etc.) if self._on_session_cleanup: try: diff --git a/tests/runtime/test_claims.py b/tests/runtime/test_claims.py new file mode 100644 index 000000000..84f6f263f --- /dev/null +++ b/tests/runtime/test_claims.py @@ -0,0 +1,33 @@ +"""Tests for the shared inbound message claim registry.""" + +from __future__ import annotations + +from band.runtime.claims import MessageClaimRegistry + + +def test_claim_reports_contention_without_releasing_owner() -> None: + registry = MessageClaimRegistry() + + with registry.claim("room-1", "message-1") as owner: + assert owner + with registry.claim("room-1", "message-1") as contender: + assert not contender + assert registry.inflight_ids("room-1") == {"message-1"} + + +def test_same_message_id_is_independent_between_rooms() -> None: + registry = MessageClaimRegistry() + + with registry.claim("room-1", "message-1") as room_one: + with registry.claim("room-2", "message-1") as room_two: + assert room_one and room_two + + +def test_completed_cache_pressure_is_isolated_per_room() -> None: + registry = MessageClaimRegistry(max_completed=1) + registry.remember_completed("room-1", "old") + registry.remember_completed("room-2", "kept") + registry.remember_completed("room-1", "new") + + assert registry.completed_ids("room-1") == ["new"] + assert registry.completed_ids("room-2") == ["kept"] diff --git a/tests/runtime/test_execution.py b/tests/runtime/test_execution.py index cb59f9f48..d578aaffe 100644 --- a/tests/runtime/test_execution.py +++ b/tests/runtime/test_execution.py @@ -8,9 +8,11 @@ import pytest +from band.runtime.claims import MessageClaimRegistry from band.runtime.execution import ( Execution, ExecutionContext, + ExecutionState, _BacklogProcessResult, _error_label, ) @@ -87,7 +89,7 @@ def test_init_starts_idle(self, mock_link, mock_handler): """Should start in starting state, not running.""" ctx = ExecutionContext("room-123", mock_link, mock_handler) - assert ctx.state == "starting" + assert ctx.state is ExecutionState.STARTING assert ctx.is_running is False assert ctx.is_processing is False @@ -553,7 +555,7 @@ async def test_sync_point_clears_marker_and_keeps_dedupe_cache( # Marker should be cleared assert ctx._first_ws_msg_id is None # Dedupe cache should keep processed sync id to avoid WS reprocessing - assert "msg-sync-001" in ctx._processed_ids + assert "msg-sync-001" in ctx.claims.completed_ids(ctx.room_id) await ctx.stop() @@ -661,7 +663,7 @@ async def test_ws_replay_with_processed_metadata_is_not_reopened( mock_handler.assert_not_called() mock_link_with_next.mark_processing.assert_not_called() - assert "msg-processed-replay" in ctx._processed_ids + assert "msg-processed-replay" in ctx.claims.completed_ids(ctx.room_id) await ctx.stop() @@ -700,7 +702,7 @@ async def test_ws_replay_uses_hydrated_context_delivery_status( mock_handler.assert_not_called() mock_link_with_next.mark_processing.assert_not_called() - assert "msg-stale-replay" in ctx._processed_ids + assert "msg-stale-replay" in ctx.claims.completed_ids(ctx.room_id) await ctx.stop() @@ -811,7 +813,114 @@ async def delayed_mark_processing(room_id: str, message_id: str) -> bool: await backlog_task assert mock_handler.await_count == 1 - assert ctx._inflight_message_ids == set() + assert ctx.claims.inflight_ids(ctx.room_id) == set() + + async def test_first_message_to_fresh_room_executes_once(self, mock_link_with_next): + """A message posted before the room's context was live executes once. + + The gateway sequence: create room, add peer, post immediately. The + peer's startup sync receives the message from /next, and the + WebSocket copy arrives while that execution is still in flight. The + second delivery must be deduplicated, not re-executed. + """ + from band.runtime.types import PlatformMessage + + handler_started = asyncio.Event() + release_handler = asyncio.Event() + handled_message_ids: list[str] = [] + + async def blocking_handler(_context: ExecutionContext, event) -> None: + handled_message_ids.append(event.payload.id) + handler_started.set() + await release_handler.wait() + + first_message = PlatformMessage( + id="msg-first", + room_id="room-123", + content="posted before the peer subscribed", + sender_id="user-1", + sender_type="User", + sender_name="User One", + message_type="text", + metadata={}, + created_at=datetime.now(timezone.utc), + ) + mock_link_with_next.get_next_message = AsyncMock( + side_effect=[first_message, None] + ) + + ctx = ExecutionContext( + "room-123", + mock_link_with_next, + blocking_handler, + agent_id="agent-123", + config=SessionConfig(enable_context_hydration=False), + ) + + await ctx.start() + await asyncio.wait_for(handler_started.wait(), timeout=1.0) + + # WebSocket copy of the same message arrives mid-execution. + await ctx.on_event(make_message_event(room_id="room-123", msg_id="msg-first")) + + release_handler.set() + await asyncio.sleep(0.2) # Let sync finish and Phase 2 drain the WS copy + + assert handled_message_ids == ["msg-first"] + assert mock_link_with_next.mark_processing.await_count == 1 + assert mock_link_with_next.mark_processed.await_count == 1 + + await ctx.stop() + + async def test_fresh_contexts_do_not_execute_the_same_message_twice( + self, mock_link_with_next + ): + """Contexts sharing a runtime's registry execute a shared message once. + + The losing delivery is deferred (not treated as durably handled), so + the resync safety net re-checks it if the owner later fails. + """ + handler_started = asyncio.Event() + release_handler = asyncio.Event() + handler_calls = 0 + + async def blocking_handler(_context: ExecutionContext, _event: object) -> None: + nonlocal handler_calls + handler_calls += 1 + handler_started.set() + await release_handler.wait() + + registry = MessageClaimRegistry() + + def fresh_context() -> ExecutionContext: + return ExecutionContext( + "room-123", + mock_link_with_next, + blocking_handler, + agent_id="agent-123", + config=SessionConfig( + enable_context_hydration=False, + enable_working_state=False, + ), + claim_registry=registry, + ) + + first_context = fresh_context() + second_context = fresh_context() + event = make_message_event(room_id="room-123", msg_id="msg-fresh-peer") + + first_task = asyncio.create_task(first_context._process_event(event)) + await asyncio.wait_for(handler_started.wait(), timeout=1.0) + second_task = asyncio.create_task(second_context._process_event(event)) + + await asyncio.sleep(0) + release_handler.set() + owner_handled, duplicate_handled = await asyncio.gather(first_task, second_task) + + assert handler_calls == 1 + assert mock_link_with_next.mark_processing.await_count == 1 + assert owner_handled is True + assert duplicate_handled is False async def test_mark_processing_failure_does_not_execute_message( self, mock_link_with_next, mock_handler @@ -833,7 +942,32 @@ async def test_mark_processing_failure_does_not_execute_message( ) mock_handler.assert_not_called() mock_link_with_next.mark_processed.assert_not_called() - assert ctx._inflight_message_ids == set() + assert ctx.claims.inflight_ids(ctx.room_id) == set() + + async def test_handler_failure_marks_failed_and_releases_claim( + self, mock_link_with_next + ): + """A real handler failure must not strand ownership of the message.""" + + async def failing_handler(ctx, event): + raise RuntimeError("handler failed") + + mock_link_with_next.mark_processing = AsyncMock(return_value=True) + mock_link_with_next.mark_failed = AsyncMock(return_value=True) + ctx = ExecutionContext( + "room-123", + mock_link_with_next, + failing_handler, + config=SessionConfig(enable_context_hydration=False), + ) + event = make_message_event(room_id="room-123", msg_id="msg-handler-fails") + + assert await ctx._process_event(event) is True + + mock_link_with_next.mark_failed.assert_awaited_once_with( + "room-123", "msg-handler-fails", "handler failed" + ) + assert ctx.claims.inflight_ids(ctx.room_id) == set() async def test_backlog_processed_ack_failure_is_not_remembered( self, mock_link_with_next, mock_handler @@ -868,8 +1002,9 @@ async def test_backlog_processed_ack_failure_is_not_remembered( mock_link_with_next.mark_processed.assert_awaited_once_with( "room-123", "msg-ack-fails" ) - assert "msg-ack-fails" not in ctx._processed_ids - assert "msg-ack-fails" in ctx._processed_ack_pending_ids + assert "msg-ack-fails" not in ctx.claims.completed_ids(ctx.room_id) + assert ctx.claims.is_ack_pending(ctx.room_id, "msg-ack-fails") + assert ctx.claims.inflight_ids(ctx.room_id) == set() async def test_websocket_processed_ack_failure_is_not_remembered( self, mock_link_with_next, mock_handler @@ -891,8 +1026,8 @@ async def test_websocket_processed_ack_failure_is_not_remembered( mock_link_with_next.mark_processed.assert_awaited_once_with( "room-123", "msg-ws-ack-fails" ) - assert "msg-ws-ack-fails" not in ctx._processed_ids - assert "msg-ws-ack-fails" in ctx._processed_ack_pending_ids + assert "msg-ws-ack-fails" not in ctx.claims.completed_ids(ctx.room_id) + assert ctx.claims.is_ack_pending(ctx.room_id, "msg-ws-ack-fails") async def test_backlog_processed_ack_failure_retries_ack_without_handler_replay( self, mock_link_with_next, mock_handler @@ -927,8 +1062,8 @@ async def test_backlog_processed_ack_failure_retries_ack_without_handler_replay( mock_handler.assert_awaited_once() assert mock_link_with_next.mark_processed.await_count == 2 - assert "msg-backlog-ack-retry" in ctx._processed_ids - assert "msg-backlog-ack-retry" not in ctx._processed_ack_pending_ids + assert "msg-backlog-ack-retry" in ctx.claims.completed_ids(ctx.room_id) + assert not ctx.claims.is_ack_pending(ctx.room_id, "msg-backlog-ack-retry") async def test_processed_ack_retry_budget_exhaustion_keeps_local_completion( self, mock_link_with_next, mock_handler @@ -966,8 +1101,8 @@ async def test_processed_ack_retry_budget_exhaustion_keeps_local_completion( mock_handler.assert_awaited_once() assert mock_link_with_next.mark_processed.await_count == 2 - assert "msg-ack-budget" in ctx._processed_ids - assert "msg-ack-budget" not in ctx._processed_ack_pending_ids + assert "msg-ack-budget" in ctx.claims.completed_ids(ctx.room_id) + assert not ctx.claims.is_ack_pending(ctx.room_id, "msg-ack-budget") async def test_websocket_processed_ack_failure_retries_ack_without_handler_replay( self, mock_link_with_next, mock_handler @@ -988,8 +1123,54 @@ async def test_websocket_processed_ack_failure_retries_ack_without_handler_repla mock_handler.assert_awaited_once() assert mock_link_with_next.mark_processed.await_count == 2 - assert "msg-ws-ack-retry" in ctx._processed_ids - assert "msg-ws-ack-retry" not in ctx._processed_ack_pending_ids + assert "msg-ws-ack-retry" in ctx.claims.completed_ids(ctx.room_id) + assert not ctx.claims.is_ack_pending(ctx.room_id, "msg-ws-ack-retry") + + async def test_ack_pending_message_survives_lru_pressure_without_replay( + self, mock_link_with_next, mock_handler + ): + """Cache pressure must never evict ACK_PENDING into a side-effect replay. + + A message whose handler completed but whose durable processed ack + failed may only retry the ack on redelivery — even after enough later + completions to overflow the dedupe cache. + """ + mock_link_with_next.mark_processing = AsyncMock(return_value=True) + mock_link_with_next.mark_processed = AsyncMock(return_value=False) + ctx = ExecutionContext( + "room-123", + mock_link_with_next, + mock_handler, + agent_id="agent-123", + config=SessionConfig( + enable_context_hydration=False, + enable_working_state=False, + # Retry headroom so eviction shows up as a handler replay, + # not as the retry budget silently dropping the redelivery. + max_message_retries=5, + ), + ) + + async def deliver(message_id: str) -> None: + await ctx._process_event( + make_message_event(room_id="room-123", msg_id=message_id) + ) + + # Handler completes but the durable ack fails → ACK_PENDING. + await deliver("msg-unacked") + # Enough later completions to overflow the dedupe cache capacity. + for i in range(ctx.claims.max_completed): + await deliver(f"msg-filler-{i}") + + executions_before_redelivery = mock_handler.await_count + mock_link_with_next.mark_processed = AsyncMock(return_value=True) + + await deliver("msg-unacked") + + assert mock_handler.await_count == executions_before_redelivery + mock_link_with_next.mark_processed.assert_awaited_once_with( + "room-123", "msg-unacked" + ) async def test_websocket_processed_ack_failure_retries_ack_before_newer_queue( self, mock_link_with_next, mock_handler @@ -1015,9 +1196,9 @@ async def test_websocket_processed_ack_failure_retries_ack_before_newer_queue( "msg-new-ws", ] assert mock_link_with_next.mark_processed.await_count == 3 - assert "msg-old-ws" in ctx._processed_ids - assert "msg-new-ws" in ctx._processed_ids - assert ctx._processed_ack_pending_ids == {} + assert "msg-old-ws" in ctx.claims.completed_ids(ctx.room_id) + assert "msg-new-ws" in ctx.claims.completed_ids(ctx.room_id) + assert ctx.claims.pending_ack_ids(ctx.room_id) == [] await ctx.stop() @@ -1055,7 +1236,7 @@ async def test_sync_point_claim_failure_does_not_clear_marker( assert ctx._first_ws_msg_id == "msg-sync-claim-fails" mock_handler.assert_not_called() - assert "msg-sync-claim-fails" not in ctx._processed_ids + assert "msg-sync-claim-fails" not in ctx.claims.completed_ids(ctx.room_id) await ctx.stop() @@ -1094,7 +1275,7 @@ async def test_startup_backlog_claim_failure_does_not_spin( "room-123", "msg-startup-claim-fails" ) mock_handler.assert_not_called() - assert "msg-startup-claim-fails" not in ctx._processed_ids + assert "msg-startup-claim-fails" not in ctx.claims.completed_ids(ctx.room_id) async def test_startup_backlog_claim_failure_does_not_process_newer_ws_event( self, mock_link_with_next, mock_handler @@ -1174,7 +1355,7 @@ async def test_resync_claim_failure_does_not_spin( "room-123", "msg-resync-claim-fails" ) mock_handler.assert_not_called() - assert "msg-resync-claim-fails" not in ctx._processed_ids + assert "msg-resync-claim-fails" not in ctx.claims.completed_ids(ctx.room_id) async def test_resync_claim_failure_does_not_process_newer_ws_event( self, mock_link_with_next, mock_handler @@ -1217,7 +1398,7 @@ async def test_resync_claim_failure_does_not_process_newer_ws_event( "room-123", "msg-resync-older-claim-fails" ) mock_handler.assert_not_called() - assert "msg-resync-newer-ws" not in ctx._processed_ids + assert "msg-resync-newer-ws" not in ctx.claims.completed_ids(ctx.room_id) await ctx.stop() @@ -1362,7 +1543,7 @@ class TestCancellationDuringProcessing: """Tests for cancellation during message processing.""" async def test_stop_cancels_slow_processing(self, mock_link): - """stop() should cancel message processing.""" + """stop() should cancel processing and release the in-flight claim.""" async def slow_handler(ctx, event): await asyncio.sleep(10) # Would take 10 seconds @@ -1393,6 +1574,9 @@ async def slow_handler(ctx, event): # Should complete quickly (not wait 10 seconds for handler) assert elapsed < 1.0, f"stop() took {elapsed}s - should cancel processing" + # Cancellation must release the local in-flight claim + assert ctx.claims.inflight_ids(ctx.room_id) == set() + class TestContextHydrationConfig: """Test context hydration behavior with config.""" @@ -1837,7 +2021,7 @@ async def test_wait_for_idle_returns_true_when_already_idle( ): """_wait_for_idle should return True immediately when idle.""" ctx = ExecutionContext("room-123", mock_link, mock_handler) - ctx.state = "idle" + ctx.state = ExecutionState.IDLE result = await ctx._wait_for_idle(timeout=1.0) @@ -1848,7 +2032,7 @@ async def test_wait_for_idle_returns_false_on_timeout( ): """_wait_for_idle should return False when timeout exceeded.""" ctx = ExecutionContext("room-123", mock_link, mock_handler) - ctx._set_state("processing") # Use _set_state to properly clear idle event + ctx._set_state(ExecutionState.PROCESSING) start = asyncio.get_running_loop().time() result = await ctx._wait_for_idle(timeout=0.1) diff --git a/tests/runtime/test_resync.py b/tests/runtime/test_resync.py index 549eb6c18..51d9f26be 100644 --- a/tests/runtime/test_resync.py +++ b/tests/runtime/test_resync.py @@ -245,13 +245,13 @@ async def test_processes_single_missed_message(self, mock_link, mock_handler): await ctx.stop() async def test_skips_duplicate_message(self, mock_link, mock_handler): - """A message whose ID is already in _processed_ids should be skipped.""" + """A message already recorded as completed should be skipped.""" msg = make_platform_message(msg_id="dup-1", room_id="room-1") mock_link.get_next_message.side_effect = [msg, None] ctx = ExecutionContext("room-1", mock_link, mock_handler) # Mark the message as already processed - ctx._processed_ids["dup-1"] = True + ctx.claims.remember_completed(ctx.room_id, "dup-1") await ctx._resync_pending_messages() diff --git a/tests/runtime/test_runtime.py b/tests/runtime/test_runtime.py index 40a9f9705..7c0eaecfe 100644 --- a/tests/runtime/test_runtime.py +++ b/tests/runtime/test_runtime.py @@ -185,6 +185,60 @@ async def test_execution_idempotent(self, mock_link, mock_handler): await runtime.stop() + async def test_all_contexts_share_one_claim_registry(self, mock_link, mock_handler): + """Message ownership must be shared across every context the runtime + creates — including a context recreated after room removal — while + retaining room-local ownership state.""" + runtime = AgentRuntime(mock_link, "agent-123", mock_handler) + + room_a = await runtime._create_execution("room-a") + room_b = await runtime._create_execution("room-b") + await runtime._on_room_left("room-a") + room_a_again = await runtime._create_execution("room-a") + + assert room_a.claims is room_b.claims + assert room_a_again.claims is room_a.claims + + await runtime.stop() + + async def test_pending_ack_is_retried_only_by_its_room( + self, mock_link, mock_handler + ): + """A shared registry must not acknowledge a different room's message.""" + mock_link.mark_processed = AsyncMock(return_value=True) + runtime = AgentRuntime(mock_link, "agent-123", mock_handler) + room_a = await runtime._create_execution("room-a") + room_b = await runtime._create_execution("room-b") + room_a.claims.remember_ack_pending("room-a", "msg-1") + + assert await room_b._retry_pending_processed_acks() + mock_link.mark_processed.assert_not_awaited() + + assert await room_a._retry_pending_processed_acks() + mock_link.mark_processed.assert_awaited_once_with("room-a", "msg-1") + + await runtime.stop() + + async def test_room_recreation_reclaims_completed_and_retries_pending_ack( + self, mock_link, mock_handler + ): + """Room teardown may drop durable cache, but not unacknowledged success.""" + runtime = AgentRuntime(mock_link, "agent-123", mock_handler) + original = await runtime._create_execution("room-a") + original.claims.remember_completed("room-a", "msg-completed") + original.claims.remember_ack_pending("room-a", "msg-ack-pending") + + await runtime._on_room_left("room-a") + recreated = await runtime._create_execution("room-a") + + assert not recreated.claims.is_completed("room-a", "msg-completed") + + mock_link.mark_processed = AsyncMock(return_value=True) + assert await recreated._retry_pending_processed_acks() + mock_link.mark_processed.assert_awaited_once_with("room-a", "msg-ack-pending") + + await runtime.stop() + async def test_active_sessions_returns_copy(self, mock_link, mock_handler): """active_sessions should return a copy.""" runtime = AgentRuntime(mock_link, "agent-123", mock_handler) diff --git a/tests/test_session_sync.py b/tests/test_session_sync.py index 7c9a435b5..4332d4b6d 100644 --- a/tests/test_session_sync.py +++ b/tests/test_session_sync.py @@ -135,7 +135,7 @@ async def test_processed_event_added_to_cache(self, ctx): await ctx._process_event(event) - assert "msg-001" in ctx._processed_ids + assert "msg-001" in ctx.claims.completed_ids(ctx.room_id) @pytest.mark.asyncio async def test_duplicate_event_skipped(self, ctx): @@ -153,26 +153,26 @@ async def test_duplicate_event_skipped(self, ctx): @pytest.mark.asyncio async def test_lru_cache_evicts_oldest(self, ctx): """LRU cache should evict oldest entries when full.""" - limit = ctx._max_processed_ids + limit = ctx.claims.max_completed for i in range(limit): event = make_message_event(msg_id=f"msg-{i:03d}") await ctx._process_event(event) - assert len(ctx._processed_ids) == limit - assert "msg-000" in ctx._processed_ids + assert len(ctx.claims.completed_ids(ctx.room_id)) == limit + assert "msg-000" in ctx.claims.completed_ids(ctx.room_id) event_overflow = make_message_event(msg_id=f"msg-{limit:03d}") await ctx._process_event(event_overflow) - assert len(ctx._processed_ids) == limit - assert "msg-000" not in ctx._processed_ids - assert f"msg-{limit:03d}" in ctx._processed_ids + assert len(ctx.claims.completed_ids(ctx.room_id)) == limit + assert "msg-000" not in ctx.claims.completed_ids(ctx.room_id) + assert f"msg-{limit:03d}" in ctx.claims.completed_ids(ctx.room_id) @pytest.mark.asyncio async def test_duplicate_refreshes_lru_position(self, ctx): """Accessing duplicate should refresh its LRU position.""" - limit = ctx._max_processed_ids + limit = ctx.claims.max_completed for i in range(limit): event = make_message_event(msg_id=f"msg-{i:03d}") @@ -184,8 +184,8 @@ async def test_duplicate_refreshes_lru_position(self, ctx): event_overflow = make_message_event(msg_id=f"msg-{limit:03d}") await ctx._process_event(event_overflow) - assert "msg-000" in ctx._processed_ids - assert "msg-001" not in ctx._processed_ids + assert "msg-000" in ctx.claims.completed_ids(ctx.room_id) + assert "msg-001" not in ctx.claims.completed_ids(ctx.room_id) class TestSynchronizeWithNext: @@ -251,7 +251,7 @@ async def test_sync_point_reached_clears_marker_and_keeps_cache( # Queue should be empty (duplicate removed) assert ctx.queue.empty() # Dedupe cache keeps processed IDs to avoid immediate WS reprocessing - assert "sync-001" in ctx._processed_ids + assert "sync-001" in ctx.claims.completed_ids(ctx.room_id) @pytest.mark.asyncio async def test_sync_point_with_non_message_head_processes_once(