ForcePilot/backend/package/yuxi/channel/message/dispatch.py

import logging
import time
from collections.abc import Awaitable, Callable
from typing import TYPE_CHECKING

from yuxi.channel.message.models import DispatchResult, PeerKind, UnifiedMessage
from yuxi.channel.protocols import (
    ApprovalAction,
    ApprovalDecision,
    ConfigProtocol,
    MessagingProtocol,
    SessionResolution,
)
from yuxi.channel.plugins.registry import ChannelPluginRegistry
from yuxi.channel.routing.matcher import MessageContext, RouteMatcher, VerificationContext
from yuxi.channel.routing.models import RouteBinding
from yuxi.channel.routing.session_key import SessionKeyBuilder, session_key_to_thread_id
from yuxi.channel.security.allowlist import AllowlistChecker
from yuxi.channel.security.pairing import PairingManager
from yuxi.repositories.channel_thread_mapping_repo import ChannelThreadMappingRepository
from yuxi.repositories.channel_user_mapping_repo import ChannelUserMappingRepository

if TYPE_CHECKING:
    from yuxi.channel.message.langfuse_trace import ChannelTrace

logger = logging.getLogger(__name__)

BeforeDispatchHook = Callable[[UnifiedMessage], Awaitable[bool | None]]
AfterDispatchHook = Callable[[UnifiedMessage, DispatchResult], Awaitable[None]]


class MessageDispatcher:
    def __init__(
        self,
        bindings: list[RouteBinding],
        allowlist: AllowlistChecker,
        pairing: PairingManager,
        *,
        matcher: RouteMatcher | None = None,
        session_key_builder: SessionKeyBuilder | None = None,
        approval_engine=None,
        user_mapping_repo: ChannelUserMappingRepository | None = None,
        thread_mapping_repo: ChannelThreadMappingRepository | None = None,
        on_before_dispatch: BeforeDispatchHook | None = None,
        on_after_dispatch: AfterDispatchHook | None = None,
    ):
        self._bindings = bindings
        self._matcher = matcher
        self._allowlist = allowlist
        self._pairing = pairing
        self._session_key_builder = session_key_builder or SessionKeyBuilder()
        self._approval_engine = approval_engine
        self._user_mapping_repo = user_mapping_repo
        self._thread_mapping_repo = thread_mapping_repo
        self._on_before_dispatch = on_before_dispatch
        self._on_after_dispatch = on_after_dispatch

    async def dispatch(self, msg: UnifiedMessage, *, trace: "ChannelTrace | None" = None) -> DispatchResult:
        command_result = await self._try_handle_command(msg)
        if command_result is not None:
            return command_result

        try:
            return await self._dispatch_inner(msg, trace=trace)
        except Exception:
            logger.exception("Dispatch failed for message: %s", msg.msg_id)
            return DispatchResult(success=False, error="dispatch_error")

    @staticmethod
    def _plugin_has_commands(plugin) -> bool:
        return callable(getattr(plugin, "get_commands", None)) and callable(getattr(plugin, "handle_command", None))

    async def _try_handle_command(self, msg: UnifiedMessage) -> DispatchResult | None:
        text = (msg.content or "").strip()
        if not text.startswith("/"):
            return None

        parts = text[1:].split()
        if not parts:
            return None

        command_name = parts[0].lower()
        args = parts[1:]

        plugin = ChannelPluginRegistry.get(msg.channel_type)
        if plugin is None or not self._plugin_has_commands(plugin):
            return None

        config = {}
        if isinstance(plugin, ConfigProtocol):
            try:
                account_cfg = await plugin.resolve_account(msg.account_id)
                config = account_cfg if account_cfg else {}
            except Exception:
                logger.debug("Failed to resolve account config for command: %s", command_name)

        try:
            response = await plugin.handle_command(
                config,
                command_name,
                args,
                msg,
                None,
            )
        except Exception:
            logger.exception("Command handling failed: %s", command_name)
            return DispatchResult(
                success=True,
                handled_by="command",
                command_response="命令执行失败，请稍后重试。",
            )

        return DispatchResult(
            success=True,
            handled_by="command",
            command_response=response,
        )

    @staticmethod
    def _plugin_has_approval(plugin) -> bool:
        return (
            callable(getattr(plugin, "check_approval_required", None))
            and callable(getattr(plugin, "create_approval_request", None))
            and callable(getattr(plugin, "check_approval_status", None))
        )

    async def _should_check_approval(self, msg: UnifiedMessage) -> bool:
        if self._approval_engine is None:
            return False
        plugin = ChannelPluginRegistry.get(msg.channel_type)
        return plugin is not None and self._plugin_has_approval(plugin)

    async def _dispatch_inner(self, msg: UnifiedMessage, *, trace: "ChannelTrace | None" = None) -> DispatchResult:
        plugin = ChannelPluginRegistry.get(msg.channel_type)

        sec_span_id = None
        if trace:
            sec_span_id = await trace.add_span("security_admission", metadata={"peer_kind": msg.sender.kind.value})

        if msg.sender.kind == PeerKind.DIRECT:
            try:
                dm_result = self._allowlist.check_dm(msg.sender.id)
                if not dm_result.allowed:
                    if dm_result.pairing_required:
                        req = await self._pairing.upsert_code(msg.channel_type, msg.sender.id, msg.account_id)
                        if req is not None:
                            if sec_span_id:
                                await trace.end_span(sec_span_id, level="ERROR", status_message="pairing_required")
                            return DispatchResult(success=False, error=f"pairing_required:{req.code}")
                        if sec_span_id:
                            await trace.end_span(sec_span_id, level="ERROR", status_message="pairing_check_error")
                        return DispatchResult(success=False, error="pairing_check_error")
                    logger.info(
                        "Sender not in DM allowlist: %s/%s",
                        msg.channel_type,
                        msg.sender.id,
                    )
                    if sec_span_id:
                        await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
                    return DispatchResult(success=False, error="not_in_allowlist")
            except Exception:
                logger.exception("DM allowlist check failed, blocking message")
                if sec_span_id:
                    await trace.end_span(sec_span_id, level="ERROR", status_message="allowlist_check_error")
                return DispatchResult(success=False, error="allowlist_check_error")
            auth_scope = "dm"
            auth_verified_by = "AllowlistChecker.dm"
        else:
            try:
                if msg.group and msg.group.id:
                    group_result = self._allowlist.check_group(msg.group.id)
                    if not group_result.allowed:
                        logger.info(
                            "Group not in allowlist: %s/%s",
                            msg.channel_type,
                            msg.group.id,
                        )
                        if sec_span_id:
                            await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
                        return DispatchResult(success=False, error="not_in_allowlist")
                from_result = self._allowlist.check_group_allow_from(msg.sender.id)
                if not from_result.allowed:
                    logger.info(
                        "Sender not in group allow_from: %s/%s",
                        msg.channel_type,
                        msg.sender.id,
                    )
                    if sec_span_id:
                        await trace.end_span(sec_span_id, level="ERROR", status_message="not_in_allowlist")
                    return DispatchResult(success=False, error="not_in_allowlist")
            except Exception:
                logger.exception("Group allowlist check failed, blocking message")
                if sec_span_id:
                    await trace.end_span(sec_span_id, level="ERROR", status_message="allowlist_check_error")
                return DispatchResult(success=False, error="allowlist_check_error")
            auth_scope = "group"
            auth_verified_by = "AllowlistChecker.group"

        if sec_span_id:
            await trace.end_span(sec_span_id, status_message="ok")

        if self._matcher is None:
            self._matcher = RouteMatcher()
            await self._matcher.set_bindings(self._bindings)

        explicit_target = None
        if isinstance(plugin, MessagingProtocol):
            session = plugin.resolve_session(msg)
            explicit_target = plugin.parse_explicit_target(msg.content)
        else:
            if msg.sender.kind == PeerKind.DIRECT:
                session = SessionResolution(kind="direct", conversation_id=msg.sender.id)
            else:
                gid = msg.group.id if msg.group else "unknown"
                session = SessionResolution(kind="group", conversation_id=gid)

        try:
            peer_kind = PeerKind(session.kind)
        except ValueError:
            peer_kind = msg.sender.kind

        ctx = MessageContext(
            channel=msg.channel_type,
            account_id=msg.account_id,
            peer_kind=peer_kind,
            peer_id=msg.sender.id,
            channel_config_id=msg.channel_config_id,
            member_role_ids=msg.member_role_ids,
        )
        ctx._verification_ctx = VerificationContext(
            verified_by=auth_verified_by,
            scope=auth_scope,
            timestamp=time.time(),
        )
        if msg.group:
            ctx.guild_id = msg.group.guild_id
            ctx.team_id = msg.group.team_id
            ctx.roles = msg.group.roles

            if msg.group.thread_id or msg.message_thread_id:
                if msg.group.route_peer_kind:
                    try:
                        ctx.parent_peer_kind = PeerKind(msg.group.route_peer_kind)
                    except ValueError:
                        pass
                ctx.parent_peer_id = msg.group.route_peer_id or msg.group.id

        route_span_id = None
        if trace:
            route_span_id = await trace.add_span("route_resolution")

        try:
            route = await self._matcher.resolve(ctx)
        except Exception:
            logger.exception("Route resolution failed")
            if route_span_id:
                await trace.end_span(route_span_id, level="ERROR", status_message="route_error")
            return DispatchResult(success=False, error="route_error")

        if route.agent_config_id == 0:
            if route_span_id:
                await trace.end_span(route_span_id, status_message="no_route")
            return DispatchResult(success=False, error="no_route")

        if route_span_id:
            await trace.end_span(route_span_id, metadata={"agent_config_id": str(route.agent_config_id)})

        if self._on_before_dispatch is not None:
            try:
                handled = await self._on_before_dispatch(msg)
                if handled:
                    return DispatchResult(
                        success=True,
                        handled_by="hook:before_dispatch",
                        thread_id=None,
                        agent_config_id=None,
                    )
            except Exception:
                logger.exception("before_dispatch hook failed, continuing")

        if await self._should_check_approval(msg):
            try:
                request = await self._approval_engine.request_approval(
                    msg.channel_type,
                    msg.account_id,
                    {},
                    ApprovalAction.APPROVE_EXEC,
                    msg.sender.id,
                    description=f"Agent 执行: {msg.content[:100]}",
                    context={"full_text": msg.content},
                )
                if request is not None:
                    result = await self._approval_engine.wait_for_decision(
                        msg.channel_type,
                        {},
                        request.id,
                    )
                    if result.decision != ApprovalDecision.APPROVED:
                        return DispatchResult(
                            success=False,
                            error=f"approval_denied:{result.reason or '审批未通过'}",
                        )
            except Exception:
                logger.exception("Approval check failed, blocking message")
                return DispatchResult(
                    success=False,
                    error="approval_check_error",
                )

        agent_id_str = str(route.agent_config_id)

        sk_span_id = None
        if trace:
            sk_span_id = await trace.add_span("session_key_build")

        internal_user_id, thread_id = await self._resolve_thread_id(
            msg,
            agent_id_str,
            route.dm_scope,
        )

        if sk_span_id:
            await trace.end_span(sk_span_id, metadata={
                "thread_id": thread_id,
                "has_internal_user": internal_user_id is not None,
            })

        reply_sent = False
        if explicit_target:
            try:
                target_plugin = ChannelPluginRegistry.get(explicit_target)
                if target_plugin is not None and isinstance(target_plugin, MessagingProtocol):
                    logger.info(
                        "Message %s forwarded to explicit target: %s",
                        msg.msg_id,
                        explicit_target,
                    )
            except Exception:
                logger.exception("Explicit target forwarding failed: %s", explicit_target)

        result = DispatchResult(
            success=True,
            thread_id=thread_id,
            agent_config_id=route.agent_config_id,
            internal_user_id=internal_user_id,
            reply_sent=reply_sent,
        )

        if self._on_after_dispatch is not None:
            try:
                await self._on_after_dispatch(msg, result)
            except Exception:
                logger.exception("after_dispatch hook failed")

        return result

    async def _resolve_thread_id(
        self,
        msg: UnifiedMessage,
        agent_id_str: str,
        dm_scope: str,
    ) -> tuple[str | None, str]:
        internal_user_id = None
        if self._user_mapping_repo is not None:
            try:
                user_mapping = await self._user_mapping_repo.resolve_user(
                    msg.channel_type,
                    msg.sender.id,
                )
                internal_user_id = user_mapping.internal_user_id
            except Exception:
                logger.exception("Failed to resolve user mapping")

        user_id = internal_user_id or msg.sender.id

        if msg.sender.kind == PeerKind.DIRECT:
            session_key = self._session_key_builder.build_dm_session(
                agent_id_str,
                msg.channel_type,
                msg.account_id,
                user_id,
                dm_scope=dm_scope,
            )
        else:
            session_key = self._session_key_builder.build_group_session(
                agent_id_str,
                msg.channel_type,
                msg.sender.kind.value,
                user_id,
            )

        thread_id = session_key_to_thread_id(session_key)

        if internal_user_id and self._thread_mapping_repo is not None:
            try:
                channel_chat_id = msg.group.id if msg.group and msg.group.id else msg.sender.id
                await self._thread_mapping_repo.resolve_thread(
                    msg.channel_type,
                    channel_chat_id,
                    internal_user_id,
                    agent_id=agent_id_str,
                    thread_id=thread_id,
                )
            except Exception:
                logger.exception("Failed to persist thread mapping")

        return internal_user_id, thread_id