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