deer-flow/backend/app/channels/message_bus.py

"""MessageBus — async pub/sub hub that decouples channels from the agent dispatcher."""

from __future__ import annotations

import asyncio
import logging
import threading
import time
from collections.abc import Callable, Coroutine
from dataclasses import dataclass, field
from enum import StrEnum
from pathlib import Path
from typing import Any

logger = logging.getLogger(__name__)

DEFAULT_INBOUND_QUEUE_MAXSIZE = 1000

PENDING_CLARIFICATION_METADATA_KEY = "pending_clarification"
# Full ``human_input`` payload (artifact of the ask_clarification tool message),
# attached to the final outbound so channels can render richer affordances
# (e.g. Feishu button cards) for option-based clarifications.
HUMAN_INPUT_REQUEST_METADATA_KEY = "human_input_request"
# A channel-synthesized human-input answer (e.g. a Feishu card button click),
# carried on the inbound so the manager can attach it to the human message's
# additional_kwargs for the Gateway's human_input protocol.
HUMAN_INPUT_RESPONSE_METADATA_KEY = "human_input_response"
RESOLVED_FROM_PENDING_CLARIFICATION_METADATA_KEY = "resolved_from_pending_clarification"
# Adapter-owned bytes may use this transient key while crossing the channel
# boundary. ChannelManager consumes and removes it before persisting metadata.
INBOUND_FILE_CONTENT_KEY = "_content"


# ---------------------------------------------------------------------------
# Message types
# ---------------------------------------------------------------------------


class InboundMessageType(StrEnum):
    """Types of messages arriving from IM channels."""

    CHAT = "chat"
    COMMAND = "command"


@dataclass
class InboundMessage:
    """A message arriving from an IM channel toward the agent dispatcher.

    Attributes:
        channel_name: Name of the source channel (e.g. "feishu", "slack").
        chat_id: Platform-specific chat/conversation identifier.
        user_id: Platform-specific user identifier.
        text: The message text.
        msg_type: Whether this is a regular chat message or a command.
        thread_ts: Optional platform thread identifier (for threaded replies).
        topic_id: Conversation topic identifier used to map to a DeerFlow thread.
            Messages sharing the same ``topic_id`` within a ``chat_id`` will
            reuse the same DeerFlow thread.  When ``None``, each message
            creates a new thread (one-shot Q&A).
        connection_id: Optional DeerFlow channel connection id. When present,
            conversation mapping is scoped by the connection instead of the
            legacy global ``channel_name:chat_id[:topic_id]`` key.
        owner_user_id: DeerFlow user id that owns the channel connection.
            Platform user ids stay in ``user_id``.
        workspace_id: Optional external workspace/guild/team id.
        files: Optional list of file attachments (platform-specific dicts).
        metadata: Arbitrary extra data from the channel.
        created_at: Unix timestamp when the message was created.
    """

    channel_name: str
    chat_id: str
    user_id: str
    text: str
    msg_type: InboundMessageType = InboundMessageType.CHAT
    thread_ts: str | None = None
    topic_id: str | None = None
    connection_id: str | None = None
    owner_user_id: str | None = None
    workspace_id: str | None = None
    files: list[dict[str, Any]] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: float = field(default_factory=time.time)


@dataclass
class ResolvedAttachment:
    """A file attachment resolved to a host filesystem path, ready for upload.

    Attributes:
        virtual_path: Original virtual path (e.g. /mnt/user-data/outputs/report.pdf).
        actual_path: Resolved host filesystem path.
        filename: Basename of the file.
        mime_type: MIME type (e.g. "application/pdf").
        size: File size in bytes.
        is_image: True for image/* MIME types (platforms may handle images differently).
    """

    virtual_path: str
    actual_path: Path
    filename: str
    mime_type: str
    size: int
    is_image: bool


@dataclass
class OutboundMessage:
    """A message from the agent dispatcher back to a channel.

    Attributes:
        channel_name: Target channel name (used for routing).
        chat_id: Target chat/conversation identifier.
        thread_id: DeerFlow thread ID that produced this response.
        text: The response text.
        artifacts: List of artifact paths produced by the agent.
        is_final: Whether this is the final message in the response stream.
        thread_ts: Optional platform thread identifier for threaded replies.
        metadata: Arbitrary extra data.
        connection_id: Optional DeerFlow channel connection id used for
            connection-specific outbound credentials.
        owner_user_id: DeerFlow user id that owns the channel connection.
        created_at: Unix timestamp.
    """

    channel_name: str
    chat_id: str
    thread_id: str
    text: str
    artifacts: list[str] = field(default_factory=list)
    attachments: list[ResolvedAttachment] = field(default_factory=list)
    is_final: bool = True
    thread_ts: str | None = None
    connection_id: str | None = None
    owner_user_id: str | None = None
    metadata: dict[str, Any] = field(default_factory=dict)
    created_at: float = field(default_factory=time.time)


# ---------------------------------------------------------------------------
# MessageBus
# ---------------------------------------------------------------------------

OutboundCallback = Callable[[OutboundMessage], Coroutine[Any, Any, None]]


class InboundQueueFullError(RuntimeError):
    """Raised when bounded inbound admission has no capacity."""


class InboundQueueClosedError(RuntimeError):
    """Raised when inbound admission has closed during shutdown."""


class InboundReservationExpiredError(RuntimeError):
    """Raised when a reservation was already committed or invalidated."""


class InboundReservation:
    """One capacity slot reserved before a provider hands work to the bus.

    Some provider SDKs invoke DeerFlow on a foreign thread. Reserving before
    scheduling their final identity/ack preparation onto the Gateway loop
    bounds both the queue and those scheduled callbacks. A reservation must be
    committed exactly once or released in a ``finally`` block.
    """

    def __init__(self, bus: MessageBus, token: object) -> None:
        self._bus = bus
        self._token = token

    def commit(self, msg: InboundMessage) -> None:
        """Commit the reserved message from the MessageBus event loop."""
        self._bus._commit_inbound(self._token, msg)

    def release(self) -> None:
        """Release the slot if it has not already been committed or closed."""
        self._bus._release_inbound_reservation(self._token)


class MessageBus:
    """Async pub/sub hub connecting channels and the agent dispatcher.

    Channels publish inbound messages; the dispatcher consumes them.
    The dispatcher publishes outbound messages; channels receive them
    via registered callbacks.
    """

    def __init__(self, *, inbound_queue_maxsize: int = DEFAULT_INBOUND_QUEUE_MAXSIZE) -> None:
        if isinstance(inbound_queue_maxsize, bool) or not isinstance(inbound_queue_maxsize, int) or inbound_queue_maxsize <= 0:
            raise ValueError("inbound_queue_maxsize must be a positive integer")

        self._inbound_queue: asyncio.Queue[InboundMessage] = asyncio.Queue(maxsize=inbound_queue_maxsize)
        # Provider callbacks may reserve capacity from SDK-owned threads before
        # scheduling async identity/ack preparation on the Gateway loop.
        self._inbound_admission_lock = threading.Lock()
        self._inbound_queued = 0
        self._inbound_reservations: set[object] = set()
        self._accepting_inbound = True
        self._full_rejection_count = 0
        self._last_full_warning_at = 0.0
        self._outbound_listeners: list[OutboundCallback] = []

    # -- inbound -----------------------------------------------------------

    async def publish_inbound(self, msg: InboundMessage) -> None:
        """Admit a message immediately or raise instead of waiting for space.

        This deliberately uses reservation + ``put_nowait`` rather than an
        awaited ``Queue.put``. Under overload, producers get an explicit
        rejection and cannot accumulate an unbounded set of pending put tasks.
        The initial zero-delay sleep is a scheduling handoff, not a capacity
        wait: a producer that publishes a batch (notably GitHub webhook
        fan-out) gives fixed workers a chance to dequeue between entries so a
        batch larger than the queue does not repeatedly fail on the same
        prefix during redelivery. It occurs before admission so cancellation
        cannot report failure after this call already committed the message.
        """
        await asyncio.sleep(0)
        reservation = self.reserve_inbound(msg)
        try:
            reservation.commit(msg)
        finally:
            reservation.release()

    def reserve_inbound(self, msg: InboundMessage) -> InboundReservation:
        """Reserve one bounded intake slot, safely callable from SDK threads."""
        token = object()
        should_warn = False
        rejection_count = 0
        with self._inbound_admission_lock:
            if not self._accepting_inbound:
                raise InboundQueueClosedError("channel inbound intake is closed")
            admitted = self._inbound_queued + len(self._inbound_reservations)
            if admitted >= self._inbound_queue.maxsize:
                self._full_rejection_count += 1
                rejection_count = self._full_rejection_count
                now = time.monotonic()
                if now - self._last_full_warning_at >= 1.0:
                    self._last_full_warning_at = now
                    should_warn = True
            else:
                self._inbound_reservations.add(token)
                return InboundReservation(self, token)

        if should_warn:
            logger.warning(
                "[Bus] inbound capacity exhausted: channel=%s, chat_id=%s, capacity=%d, rejected_total=%d",
                msg.channel_name,
                msg.chat_id,
                self._inbound_queue.maxsize,
                rejection_count,
            )
        raise InboundQueueFullError(f"channel inbound queue is full (capacity={self._inbound_queue.maxsize})")

    def _commit_inbound(self, token: object, msg: InboundMessage) -> None:
        with self._inbound_admission_lock:
            if token not in self._inbound_reservations:
                raise InboundReservationExpiredError("inbound reservation is no longer active")
            self._inbound_reservations.remove(token)
            if not self._accepting_inbound:
                raise InboundQueueClosedError("channel inbound intake closed before reservation commit")
            try:
                self._inbound_queue.put_nowait(msg)
            except asyncio.QueueFull as exc:  # pragma: no cover - reservation accounting invariant
                raise RuntimeError("inbound reservation accounting exceeded queue capacity") from exc
            self._inbound_queued += 1

        logger.info(
            "[Bus] inbound enqueued: channel=%s, chat_id=%s, type=%s, queue_size=%d",
            msg.channel_name,
            msg.chat_id,
            msg.msg_type.value,
            self._inbound_queue.qsize(),
        )

    def _release_inbound_reservation(self, token: object) -> None:
        with self._inbound_admission_lock:
            if token in self._inbound_reservations:
                self._inbound_reservations.remove(token)

    async def get_inbound(self) -> InboundMessage:
        """Block until the next inbound message is available."""
        msg = await self._inbound_queue.get()
        with self._inbound_admission_lock:
            self._inbound_queued -= 1
        return msg

    def get_inbound_nowait(self) -> InboundMessage:
        """Return one queued message immediately and release admission capacity."""
        msg = self._inbound_queue.get_nowait()
        with self._inbound_admission_lock:
            self._inbound_queued -= 1
        return msg

    def inbound_task_done(self) -> None:
        """Mark one dequeued inbound message as fully handled."""
        self._inbound_queue.task_done()

    async def join_inbound(self) -> None:
        """Wait until every admitted queue item has completed or been dropped."""
        await self._inbound_queue.join()

    def close_inbound(self) -> int:
        """Reject new intake and invalidate uncommitted reservations.

        Returns the number of provider-side reservations invalidated. Queued
        messages remain until workers finish or ``discard_pending_inbound`` is
        called by shutdown.
        """
        with self._inbound_admission_lock:
            self._accepting_inbound = False
            invalidated = len(self._inbound_reservations)
            self._inbound_reservations.clear()
        return invalidated

    def open_inbound(self) -> None:
        """Re-open admission when a stopped manager is explicitly restarted."""
        with self._inbound_admission_lock:
            self._accepting_inbound = True

    def discard_pending_inbound(self) -> int:
        """Drop queued, not-yet-started messages during shutdown."""
        discarded = 0
        while True:
            try:
                self.get_inbound_nowait()
            except asyncio.QueueEmpty:
                break
            self._inbound_queue.task_done()
            discarded += 1
        return discarded

    @property
    def inbound_queue_maxsize(self) -> int:
        return self._inbound_queue.maxsize

    @property
    def inbound_queue(self) -> asyncio.Queue[InboundMessage]:
        """Expose the queue for read-only size/empty inspection."""
        return self._inbound_queue

    # -- outbound ----------------------------------------------------------

    def subscribe_outbound(self, callback: OutboundCallback) -> None:
        """Register an async callback for outbound messages."""
        self._outbound_listeners.append(callback)

    def unsubscribe_outbound(self, callback: OutboundCallback) -> None:
        """Remove a callback from subsequent outbound dispatches.

        This does not cancel or wait for callbacks in an in-flight dispatch.
        """
        self._outbound_listeners = [cb for cb in self._outbound_listeners if cb != callback]

    async def publish_outbound(self, msg: OutboundMessage) -> None:
        """Dispatch to listeners registered when dispatch begins.

        Subscribing or unsubscribing during an in-flight dispatch affects
        subsequent messages, not callbacks captured for the current message.
        """
        listeners = tuple(self._outbound_listeners)
        logger.info(
            "[Bus] outbound dispatching: channel=%s, chat_id=%s, listeners=%d, text_len=%d",
            msg.channel_name,
            msg.chat_id,
            len(listeners),
            len(msg.text),
        )
        for callback in listeners:
            try:
                await callback(msg)
            except Exception:
                logger.exception("Error in outbound callback for channel=%s", msg.channel_name)