Communication¶

@override
def model_copy(
    self,
    *,
    update: Mapping[str, object] | None = None,
    deep: bool = False,
) -> Self:
    """Override to re-freeze data when model_copy updates it.

    Returns:
        The copied model, with ``data`` re-frozen when the update
        replaced it.
    """
    result = super().model_copy(update=update, deep=deep)
    if update and "data" in update:
        raw = update["data"]
        if isinstance(raw, MappingProxyType):
            raw = dict(raw)
        frozen = freeze_recursive(deep_copy_mapping(raw))
        object.__setattr__(result, "data", frozen)
    return result

def register(
    self,
    handler: MessageHandler | MessageHandlerFunc,
    *,
    message_types: frozenset[MessageType] | None = None,
    min_priority: MessagePriority = MessagePriority.LOW,
    name: str = "unnamed",
) -> str:
    """Register a handler for incoming messages.

    If *handler* is a bare async function, it is automatically wrapped
    in a :class:`FunctionHandler`.

    Args:
        handler: The handler instance or async function.
        message_types: Message types to match (empty/None = all).
        min_priority: Minimum priority to accept.
        name: Human-readable label for debugging.

    Returns:
        The unique handler registration ID.

    Raises:
        TypeError: If a ``MessageHandler`` instance has a synchronous
            ``handle()`` method.
    """
    if not isinstance(handler, MessageHandler):
        handler = FunctionHandler(handler)
    elif not inspect.iscoroutinefunction(handler.handle):
        msg = (
            f"MessageHandler {type(handler).__name__!r} has a "
            f"synchronous handle() -- must be async"
        )
        logger.warning(
            COMM_HANDLER_INVALID,
            agent_id=self._agent_id,
            handler_name=name,
            handler_type=type(handler).__name__,
            error=msg,
        )
        raise TypeError(msg)

    registration = HandlerRegistration(
        handler=handler,
        message_types=message_types or frozenset(),
        min_priority=min_priority,
        name=name,
    )
    self._registrations[registration.handler_id] = registration
    logger.info(
        COMM_HANDLER_REGISTERED,
        agent_id=self._agent_id,
        handler_id=registration.handler_id,
        handler_name=name,
    )
    return registration.handler_id

def deregister(self, handler_id: str) -> bool:
    """Remove a previously registered handler.

    Args:
        handler_id: The registration ID returned by :meth:`register`.

    Returns:
        True if the handler was found and removed, False otherwise.
    """
    removed = self._registrations.pop(handler_id, None)
    if removed is not None:
        logger.info(
            COMM_HANDLER_DEREGISTERED,
            agent_id=self._agent_id,
            handler_id=handler_id,
            handler_name=removed.name,
        )
        return True
    logger.debug(
        COMM_HANDLER_DEREGISTER_MISS,
        agent_id=self._agent_id,
        handler_id=handler_id,
    )
    return False

async def dispatch(self, message: Message) -> DispatchResult:
    """Route a message to all matching handlers concurrently.

    Most handler ``Exception`` subclasses are isolated; their
    errors are captured into the per-handler ``errors`` slot
    without affecting siblings. ``MemoryError`` and
    ``RecursionError`` are deliberate carve-outs: the interpreter
    is in catastrophic state, so they re-raise to abort the
    ``TaskGroup`` and cancel all remaining handlers rather than
    being absorbed silently. ``BaseException`` subclasses (e.g.
    ``KeyboardInterrupt``, ``asyncio.CancelledError``) likewise
    propagate through the ``TaskGroup``, cancelling all remaining
    handlers.

    Args:
        message: The message to dispatch.

    Returns:
        A :class:`DispatchResult` summarising the outcome.
    """
    matched = [
        reg for reg in self._registrations.values() if self._matches(reg, message)
    ]

    logger.debug(
        COMM_DISPATCH_START,
        agent_id=self._agent_id,
        message_id=str(message.id),
        message_type=message.type,
    )

    if not matched:
        logger.debug(
            COMM_DISPATCH_NO_HANDLERS,
            agent_id=self._agent_id,
            message_id=str(message.id),
        )
        return DispatchResult(
            message_id=message.id,
            handlers_succeeded=0,
            handlers_failed=0,
        )

    logger.debug(
        COMM_DISPATCH_HANDLER_MATCHED,
        agent_id=self._agent_id,
        message_id=str(message.id),
        count=len(matched),
    )

    errors: list[str | None] = [None] * len(matched)

    async with asyncio.TaskGroup() as tg:
        for idx, reg in enumerate(matched):
            _ = tg.create_task(
                self._guarded_handle(reg, message, errors, idx),
            )

    error_msgs = tuple(e for e in errors if e is not None)
    succeeded = len(matched) - len(error_msgs)

    logger.info(
        COMM_DISPATCH_COMPLETE,
        agent_id=self._agent_id,
        message_id=str(message.id),
        matched=len(matched),
        succeeded=succeeded,
        failed=len(error_msgs),
    )

    return DispatchResult(
        message_id=message.id,
        handlers_succeeded=succeeded,
        handlers_failed=len(error_msgs),
        errors=error_msgs,
    )

async def alert(self, *, title: str, body: str) -> None:
    """Deliver an operator alert.

    Args:
        title: Short alert title.
        body: Human-readable alert body.
    """
    ...

async def start(self) -> None:
    """Start the bus and create pre-configured channels.

    Raises:
        MessageBusAlreadyRunningError: If the bus is already running.
    """
    ...

async def stop(self) -> None:
    """Stop the bus gracefully.  Idempotent."""
    ...

async def health_check(self) -> bool:
    """Live liveness probe against the underlying transport.

    Unlike :attr:`is_running`, this performs an actual
    round-trip to confirm the broker is reachable and the
    connection is healthy. Implementations return ``False``
    rather than raising when the probe fails, so the caller
    can treat the bus as degraded without a try/except.
    """
    ...

async def publish(
    self,
    message: Message,
    *,
    ttl_seconds: float | None = None,
) -> None:
    """Publish a message to its channel.

    The target channel is determined by ``message.channel``.

    Args:
        message: The message to publish.
        ttl_seconds: Optional per-message TTL in seconds.  When set,
            the message expires after this duration on the server.
            Requires the NATS stream to have ``allow_msg_ttl=True``.
            ``None`` (default) defers to stream retention policy.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        ChannelNotFoundError: If the target channel does not exist.
    """
    ...

async def send_direct(
    self,
    message: Message,
    *,
    recipient: NotBlankStr,
    ttl_seconds: float | None = None,
) -> None:
    """Send a direct message between two agents.

    Lazily creates a DIRECT channel named ``@{a}:{b}`` (where
    a, b are the sorted agent IDs) and subscribes both agents.

    Args:
        message: The message to send (``message.sender`` is the
            sender).
        recipient: The recipient agent ID.
        ttl_seconds: Optional per-message TTL in seconds.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        ValueError: If the recipient does not match
            ``message.to``, an agent ID contains the reserved
            separator, or the message exceeds the payload limit.
    """
    ...

async def subscribe(
    self,
    channel_name: NotBlankStr,
    subscriber_id: NotBlankStr,
) -> Subscription:
    """Subscribe an agent to a channel.

    Idempotent -- returns a fresh subscription record if already
    subscribed (the channel's subscriber list is not duplicated).

    Args:
        channel_name: Channel to subscribe to.
        subscriber_id: Agent ID of the subscriber.

    Returns:
        The subscription record.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        ChannelNotFoundError: If the channel does not exist.
    """
    ...

async def unsubscribe(
    self,
    channel_name: NotBlankStr,
    subscriber_id: NotBlankStr,
) -> None:
    """Remove an agent's subscription from a channel.

    Args:
        channel_name: Channel to unsubscribe from.
        subscriber_id: Agent ID to remove.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        NotSubscribedError: If the agent is not subscribed.
    """
    ...

async def receive(
    self,
    channel_name: NotBlankStr,
    subscriber_id: NotBlankStr,
    *,
    timeout: float | None = None,  # noqa: ASYNC109
) -> DeliveryEnvelope | None:
    """Receive the next message from a channel.

    Awaits until a message is available, the timeout expires, or
    the bus is stopped.  When ``timeout`` is ``None``, awaits
    indefinitely (or until shutdown).

    Args:
        channel_name: Channel to receive from.
        subscriber_id: Agent ID receiving.
        timeout: Seconds to wait before returning ``None``.

    Returns:
        A delivery envelope, or ``None`` on timeout or shutdown.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        ChannelNotFoundError: If the channel does not exist.
        NotSubscribedError: If the subscriber is not subscribed
            (for TOPIC and DIRECT channels).
    """
    ...

async def create_channel(self, channel: Channel) -> Channel:
    """Create a new channel.

    Args:
        channel: Channel definition to create.

    Returns:
        The created channel.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        ChannelAlreadyExistsError: If a channel with that name
            already exists.
    """
    ...

async def get_channel(self, channel_name: NotBlankStr) -> Channel:
    """Get a channel by name.

    Args:
        channel_name: Name of the channel.

    Returns:
        The channel.

    Raises:
        ChannelNotFoundError: If the channel does not exist.
    """
    ...

async def publish_batch(
    self,
    messages: Sequence[Message],
    *,
    ttl_seconds: float | None = None,
) -> None:
    """Publish multiple messages in a batch.

    Optimised for throughput.  In NATS backends, uses pipelined
    async publishes for reduced round-trip overhead.

    Args:
        messages: Messages to publish.  An empty sequence is a
            no-op.
        ttl_seconds: Optional per-message TTL applied to every
            message in the batch.

    Raises:
        MessageBusNotRunningError: If the bus is not running.
        ChannelNotFoundError: If any target channel does not exist.
    """
    ...

async def list_channels(self) -> tuple[Channel, ...]:
    """List all channels.

    Returns:
        All registered channels.
    """
    ...

def set_quadratic_alert_sink(self, sink: QuadraticAlertSink) -> None:
    """Late-bind an alert sink to the bus's quadratic enforcer, if any.

    Boot wiring calls this once the notification dispatcher exists (the
    bus is built earlier, in the construction phase, so its enforcer
    starts sink-less). Backends without a quadratic enforcer implement
    it as a no-op, so wiring never branches on the concrete backend
    type to reach an enforcer that may not exist.
    """
    ...

async def get_channel_history(
    self,
    channel_name: NotBlankStr,
    *,
    limit: int | None = None,
) -> tuple[Message, ...]:
    """Get message history for a channel.

    Args:
        channel_name: Channel to query.
        limit: Maximum number of most recent messages to return.
            Values ``<= 0`` return an empty tuple.

    Returns:
        Messages in chronological order.

    Raises:
        ChannelNotFoundError: If the channel does not exist.
    """
    ...

@override
def __str__(self) -> str:
    """Format error with optional context metadata.

    Returns:
        The message, suffixed with ``(key=value, ...)`` context when
        present.
    """
    if self.context:
        ctx = ", ".join(f"{k}={v!r}" for k, v in self.context.items())
        return f"{self.message} ({ctx})"
    return self.message