Communication¶

def model_copy(
    self,
    *,
    update: Mapping[str, Any] | None = None,
    deep: bool = False,
) -> Self:
    """Override to re-freeze data when model_copy updates 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

async def send_message(  # noqa: PLR0913
    self,
    *,
    to: str,
    channel: str,
    content: str | None = None,
    parts: tuple[Part, ...] | None = None,
    message_type: MessageType,
    priority: MessagePriority = MessagePriority.NORMAL,
) -> Message:
    """Send a message to a channel.

    Auto-fills sender, timestamp, and message ID.  Provide either
    *content* (auto-wrapped as a ``TextPart``) or *parts* for rich
    multi-part messages.

    Args:
        to: Recipient agent or channel identifier.
        channel: Channel to publish through.
        content: Message body text (auto-wrapped as TextPart).
        parts: Explicit message parts (mutually exclusive with content).
        message_type: Message type classification.
        priority: Message priority level.

    Returns:
        The constructed and published message.

    Raises:
        ValueError: If neither or both content/parts are provided.
        ChannelNotFoundError: If the channel does not exist.
        MessageBusNotRunningError: If the bus is not running.
    """
    resolved = _resolve_parts(content, parts)
    msg = Message(
        timestamp=datetime.now(UTC),
        sender=self._agent_id,
        to=to,
        type=message_type,
        priority=priority,
        channel=channel,
        parts=resolved,
    )
    await self._bus.publish(msg)
    logger.info(
        COMM_MESSAGE_SENT,
        agent_id=self._agent_id,
        to=to,
        channel=channel,
        message_id=str(msg.id),
    )
    return msg

async def send_direct(
    self,
    *,
    to: str,
    content: str | None = None,
    parts: tuple[Part, ...] | None = None,
    message_type: MessageType,
    priority: MessagePriority = MessagePriority.NORMAL,
) -> Message:
    """Send a direct message to another agent.

    Auto-fills sender, timestamp, and message ID. The message's
    ``channel`` field is set to the deterministic ``@{a}:{b}``
    channel name computed from the sorted agent ID pair.

    Args:
        to: Recipient agent ID.
        content: Message body text (auto-wrapped as TextPart).
        parts: Explicit message parts (mutually exclusive with content).
        message_type: Message type classification.
        priority: Message priority level.

    Returns:
        The constructed and sent message.

    Raises:
        ValueError: If neither or both content/parts are provided.
        MessageBusNotRunningError: If the bus is not running.
    """
    resolved = _resolve_parts(content, parts)
    channel = _direct_channel_name(self._agent_id, to)
    msg = Message(
        timestamp=datetime.now(UTC),
        sender=self._agent_id,
        to=to,
        type=message_type,
        priority=priority,
        channel=channel,
        parts=resolved,
    )
    await self._bus.send_direct(msg, recipient=to)
    logger.info(
        COMM_MESSAGE_SENT,
        agent_id=self._agent_id,
        to=to,
        channel=channel,
        message_id=str(msg.id),
    )
    return msg

async def broadcast(
    self,
    *,
    content: str | None = None,
    parts: tuple[Part, ...] | None = None,
    message_type: MessageType,
    priority: MessagePriority = MessagePriority.NORMAL,
    channel: str = "#all-hands",
) -> Message:
    """Publish an announcement to a shared channel.

    Agents must be subscribed to the target channel to receive
    the message.  For true fan-out to all known agents, use a
    channel of type :attr:`ChannelType.BROADCAST`.

    Args:
        content: Message body text (auto-wrapped as TextPart).
        parts: Explicit message parts (mutually exclusive with content).
        message_type: Message type classification.
        priority: Message priority level.
        channel: Channel name (default ``"#all-hands"``).

    Returns:
        The constructed and published message.

    Raises:
        ValueError: If neither or both content/parts are provided.
        ChannelNotFoundError: If the channel does not exist.
        MessageBusNotRunningError: If the bus is not running.
    """
    resolved = _resolve_parts(content, parts)
    msg = Message(
        timestamp=datetime.now(UTC),
        sender=self._agent_id,
        to=channel,
        type=message_type,
        priority=priority,
        channel=channel,
        parts=resolved,
    )
    await self._bus.publish(msg)
    logger.info(
        COMM_MESSAGE_BROADCAST,
        agent_id=self._agent_id,
        channel=channel,
        message_id=str(msg.id),
    )
    return msg

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

    Args:
        channel_name: Channel to subscribe to.

    Returns:
        The subscription record.

    Raises:
        ChannelNotFoundError: If the channel does not exist.
        MessageBusNotRunningError: If the bus is not running.
    """
    sub = await self._bus.subscribe(channel_name, self._agent_id)
    logger.info(
        COMM_MESSENGER_SUBSCRIBED,
        agent_id=self._agent_id,
        channel=channel_name,
    )
    return sub

async def unsubscribe(self, channel_name: NotBlankStr) -> None:
    """Unsubscribe this agent from a channel.

    Args:
        channel_name: Channel to unsubscribe from.

    Raises:
        NotSubscribedError: If not currently subscribed.
    """
    await self._bus.unsubscribe(channel_name, self._agent_id)
    logger.info(
        COMM_MESSENGER_UNSUBSCRIBED,
        agent_id=self._agent_id,
        channel=channel_name,
    )

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

    Args:
        channel_name: Channel to receive from.
        timeout: Max seconds to wait, or ``None`` for indefinite.

    Returns:
        The next delivery envelope, or ``None`` when:

        - *timeout* expires without a message arriving.
        - The bus is shut down while waiting.
        - The subscription is cancelled via :meth:`unsubscribe`
            while a ``receive()`` call is in flight.
    """
    return await self._bus.receive(
        channel_name,
        self._agent_id,
        timeout=timeout,
    )

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

    Creates a dispatcher automatically if one was not provided
    at construction time.

    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.
    """
    if self._dispatcher is None:
        self._dispatcher = MessageDispatcher(
            agent_id=self._agent_id,
        )
    return self._dispatcher.register(
        handler,
        message_types=message_types,
        min_priority=min_priority,
        name=name,
    )

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

    Args:
        handler_id: The registration ID.

    Returns:
        True if the handler was found and removed.
    """
    if self._dispatcher is None:
        logger.debug(
            COMM_HANDLER_DEREGISTER_MISS,
            agent_id=self._agent_id,
            handler_id=handler_id,
        )
        return False
    return self._dispatcher.deregister(handler_id)

async def dispatch_message(self, message: Message) -> DispatchResult:
    """Dispatch an incoming message to registered handlers.

    Args:
        message: The message to dispatch.

    Returns:
        A :class:`DispatchResult` summarising the outcome.
    """
    if self._dispatcher is None:
        logger.debug(
            COMM_DISPATCH_NO_DISPATCHER,
            agent_id=self._agent_id,
            message_id=str(message.id),
        )
        return DispatchResult(
            message_id=message.id,
            handlers_succeeded=0,
            handlers_failed=0,
        )
    return await self._dispatcher.dispatch(message)

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.
    """
    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 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).
    """
    ...