import asyncio import logging import time from collections import OrderedDict from dataclasses import dataclass, field from yuxi.channel.routing.models import ( _PEER_KIND_ALIASES, EvaluatedBinding, NormalizedPeer, PeerConstraint, PeerConstraintState, PeerKind, RouteBinding, RouteResult, ) from yuxi.channel.routing.session_key import SessionKeyBuilder logger = logging.getLogger(__name__) MAX_CACHE_KEYS = 3000 @dataclass class VerificationContext: verified_by: str scope: str timestamp: float @dataclass class MessageContext: channel: str account_id: str peer_kind: PeerKind peer_id: str channel_config_id: str | None = None parent_peer_kind: PeerKind | None = None parent_peer_id: str | None = None guild_id: str | None = None team_id: str | None = None roles: list[str] = field(default_factory=list) member_role_ids: list[str] = field(default_factory=list) _verification_ctx: VerificationContext | None = field(default=None, repr=False) @dataclass class _BindingIndex: by_peer: dict[str, list[EvaluatedBinding]] = field(default_factory=dict) by_peer_wildcard: list[EvaluatedBinding] = field(default_factory=list) by_guild_with_roles: dict[str, list[EvaluatedBinding]] = field(default_factory=dict) by_guild: dict[str, list[EvaluatedBinding]] = field(default_factory=dict) by_team: dict[str, list[EvaluatedBinding]] = field(default_factory=dict) by_account: list[EvaluatedBinding] = field(default_factory=list) by_channel: list[EvaluatedBinding] = field(default_factory=list) def _build_index_key(msg_ctx: MessageContext) -> str: channel = msg_ctx.channel account_id = msg_ctx.account_id or "*" config_id = msg_ctx.channel_config_id or "*" return f"{channel}:{account_id}:{config_id}" def _build_route_cache_key(msg_ctx: MessageContext) -> str: channel = msg_ctx.channel account_id = msg_ctx.account_id or "*" config_id = msg_ctx.channel_config_id or "*" peer = f"{msg_ctx.peer_kind.value}:{msg_ctx.peer_id}" if msg_ctx.peer_kind and msg_ctx.peer_id else "-" parent = ( f"{msg_ctx.parent_peer_kind.value}:{msg_ctx.parent_peer_id}" if msg_ctx.parent_peer_kind and msg_ctx.parent_peer_id else "-" ) guild = msg_ctx.guild_id or "-" team = msg_ctx.team_id or "-" all_roles = sorted(set(msg_ctx.roles) | set(msg_ctx.member_role_ids)) roles = ",".join(all_roles) if all_roles else "-" return f"{channel}:{account_id}:{config_id}:{peer}:{parent}:{guild}:{team}:{roles}" def _peer_lookup_keys(kind: PeerKind, peer_id: str) -> list[str]: if kind == PeerKind.GROUP: return [f"group:{peer_id}", f"channel:{peer_id}"] if kind == PeerKind.CHANNEL: return [f"channel:{peer_id}", f"group:{peer_id}"] return [f"{kind.value}:{peer_id}"] class RouteMatcher: def __init__(self, session_builder: SessionKeyBuilder | None = None): self._lock = asyncio.Lock() self._bindings: list[RouteBinding] = [] self._default_agent_config_id: int = 0 self._cache: OrderedDict[str, RouteResult] = OrderedDict() self._index_cache: dict[str, _BindingIndex] = {} self._bindings_dirty: bool = True self._session_builder: SessionKeyBuilder = session_builder or SessionKeyBuilder() async def set_default_agent(self, agent_config_id: int) -> None: async with self._lock: self._default_agent_config_id = agent_config_id async def set_bindings(self, bindings: list[RouteBinding]) -> None: async with self._lock: self._bindings = list(bindings) self._bindings_dirty = True self._cache.clear() self._index_cache.clear() async def invalidate_cache(self) -> None: async with self._lock: self._cache.clear() self._index_cache.clear() self._bindings_dirty = True async def resolve(self, msg_ctx: MessageContext) -> RouteResult: if msg_ctx._verification_ctx is None: raise ValueError( "MessageContext must carry a VerificationContext before calling resolve(). " "RouteMatcher.resolve() requires upstream authentication (allowlist/pairing) " "via MessageDispatcher. Set msg_ctx._verification_ctx with verified_by and scope." ) cache_key = _build_route_cache_key(msg_ctx) async with self._lock: if cache_key in self._cache: self._cache.move_to_end(cache_key) return self._cache[cache_key] if self._bindings_dirty: self._rebuild_index() self._bindings_dirty = False index_key = _build_index_key(msg_ctx) index = self._index_cache.get(index_key) if index is None: wildcard_key = f"{msg_ctx.channel}:{msg_ctx.account_id or '*'}:*" index = self._index_cache.get(wildcard_key) if index is None: channel_wildcard_key = f"{msg_ctx.channel}:*:*" index = self._index_cache.get(channel_wildcard_key) if index is None: return self._build_default_result(msg_ctx) result = self._resolve_from_index(msg_ctx, index) if result is None: result = self._build_default_result(msg_ctx) logger.debug( "Route resolved: matched_by=%s agent_config_id=%s priority=%s session_key=%s", result.matched_by, result.agent_config_id, result.matched_priority, result.session_key, ) if len(self._cache) >= MAX_CACHE_KEYS: self._cache.popitem(last=False) self._cache[cache_key] = result return result def _evaluate_binding(self, binding: RouteBinding) -> EvaluatedBinding | None: m = binding.match if m is None: return None peer_state, normalized = self._normalize_peer_constraint(m.peer) if peer_state == PeerConstraintState.INVALID: return None return EvaluatedBinding( agent_config_id=binding.agent_config_id, dm_scope=binding.dm_scope, session_dm_scope=binding.session_dm_scope, normalized_peer=normalized, peer_state=peer_state, guild_id=m.guild_id, roles=m.roles, team_id=m.team_id, account_id=m.account_id, channel=m.channel, channel_config_id=m.channel_config_id, priority=binding.priority, ) def _normalize_peer_constraint( self, peer: PeerConstraint | None, ) -> tuple[PeerConstraintState, NormalizedPeer | None]: if peer is None: return PeerConstraintState.NONE, None if not hasattr(peer, "kind") or not hasattr(peer, "peer_id"): return PeerConstraintState.NONE, None kind = peer.kind peer_id = peer.peer_id if kind is None or not peer_id: return PeerConstraintState.INVALID, None if peer_id == "*": return PeerConstraintState.WILDCARD, NormalizedPeer( kind=kind, kind_lookup=_PEER_KIND_ALIASES.get(kind, {kind}), peer_id="*" ) return PeerConstraintState.VALID, NormalizedPeer( kind=kind, kind_lookup=_PEER_KIND_ALIASES.get(kind, {kind}), peer_id=peer_id ) def _rebuild_index(self) -> None: self._index_cache.clear() channel_account_groups: dict[str, list[EvaluatedBinding]] = {} for binding in self._bindings: evaluated = self._evaluate_binding(binding) if evaluated is None: continue m = binding.match channel = m.channel account_id = m.account_id or "*" config_id = m.channel_config_id or "*" key = f"{channel}:{account_id}:{config_id}" channel_account_groups.setdefault(key, []).append(evaluated) for key, evaled_list in channel_account_groups.items(): self._index_cache[key] = self._build_index(evaled_list) def _build_index(self, bindings: list[EvaluatedBinding]) -> _BindingIndex: idx = _BindingIndex() for b in bindings: state = b.peer_state if state == PeerConstraintState.VALID: np = b.normalized_peer for lookup_key in _peer_lookup_keys(np.kind, np.peer_id): idx.by_peer.setdefault(lookup_key, []).append(b) elif state == PeerConstraintState.WILDCARD: idx.by_peer_wildcard.append(b) elif b.guild_id is not None and b.roles: idx.by_guild_with_roles.setdefault(b.guild_id, []).append(b) elif b.guild_id is not None: idx.by_guild.setdefault(b.guild_id, []).append(b) elif b.team_id is not None: idx.by_team.setdefault(b.team_id, []).append(b) elif b.account_id is not None and b.account_id != "*": idx.by_account.append(b) else: idx.by_channel.append(b) idx.by_peer = {k: sorted(v, key=lambda b: (-b.priority, b.agent_config_id)) for k, v in idx.by_peer.items()} idx.by_peer_wildcard.sort(key=lambda b: (-b.priority, b.agent_config_id)) idx.by_guild_with_roles = {k: sorted(v, key=lambda b: (-b.priority, b.agent_config_id)) for k, v in idx.by_guild_with_roles.items()} idx.by_guild = {k: sorted(v, key=lambda b: (-b.priority, b.agent_config_id)) for k, v in idx.by_guild.items()} idx.by_team = {k: sorted(v, key=lambda b: (-b.priority, b.agent_config_id)) for k, v in idx.by_team.items()} idx.by_account.sort(key=lambda b: (-b.priority, b.agent_config_id)) idx.by_channel.sort(key=lambda b: (-b.priority, b.agent_config_id)) return idx def _resolve_from_index(self, msg: MessageContext, idx: _BindingIndex) -> RouteResult | None: for lookup_key in _peer_lookup_keys(msg.peer_kind, msg.peer_id): candidates = idx.by_peer.get(lookup_key) if candidates: logger.debug("Route matched by peer: %s -> agent=%s", lookup_key, candidates[0].agent_config_id) return self._build_result(msg, candidates[0], "binding.peer") if msg.parent_peer_kind and msg.parent_peer_id: for pk_key in _peer_lookup_keys(msg.parent_peer_kind, msg.parent_peer_id): candidates = idx.by_peer.get(pk_key) if candidates: np = candidates[0].normalized_peer if np is not None and msg.peer_kind.value in np.kind_lookup: logger.debug( "Route matched by parent peer: %s -> agent=%s", pk_key, candidates[0].agent_config_id, ) return self._build_result(msg, candidates[0], "binding.peer.parent") for wb in idx.by_peer_wildcard: np = wb.normalized_peer if np is not None and msg.peer_kind.value in np.kind_lookup: logger.debug("Route matched by peer wildcard -> agent=%s", wb.agent_config_id) return self._build_result(msg, wb, "binding.peer.wildcard") if msg.guild_id: gr = idx.by_guild_with_roles.get(msg.guild_id) if gr: msg_all_roles = set(msg.roles) | set(msg.member_role_ids) for b in gr: if any(r in msg_all_roles for r in b.roles): logger.debug( "Route matched by guild+roles: %s -> agent=%s", msg.guild_id, b.agent_config_id, ) return self._build_result(msg, b, "binding.guild+roles") if msg.guild_id: g = idx.by_guild.get(msg.guild_id) if g: logger.debug("Route matched by guild: %s -> agent=%s", msg.guild_id, g[0].agent_config_id) return self._build_result(msg, g[0], "binding.guild") if msg.team_id: t = idx.by_team.get(msg.team_id) if t: logger.debug("Route matched by team: %s -> agent=%s", msg.team_id, t[0].agent_config_id) return self._build_result(msg, t[0], "binding.team") for b in idx.by_account: if b.account_id == msg.account_id: logger.debug("Route matched by account: %s -> agent=%s", msg.account_id, b.agent_config_id) return self._build_result(msg, b, "binding.account") if idx.by_channel: logger.debug("Route matched by channel fallback -> agent=%s", idx.by_channel[0].agent_config_id) return self._build_result(msg, idx.by_channel[0], "binding.channel") return None def _build_result(self, msg: MessageContext, binding: EvaluatedBinding, matched_by: str) -> RouteResult: agent_id_str = str(binding.agent_config_id) effective_dm_scope = binding.session_dm_scope or binding.dm_scope if msg.peer_kind == PeerKind.DIRECT: session_key = self._session_builder.build_dm_session( agent_id_str, msg.channel, msg.account_id, msg.peer_id, dm_scope=effective_dm_scope, ) else: session_key = self._session_builder.build_group_session( agent_id_str, msg.channel, msg.peer_kind.value, msg.peer_id, ) main_session_key = self._session_builder.build_main_session(agent_id_str) last_route_policy = "main" if session_key == main_session_key else "session" return RouteResult( agent_config_id=binding.agent_config_id, channel=msg.channel, account_id=msg.account_id, session_key=session_key, main_session_key=main_session_key, matched_by=matched_by, matched_priority=binding.priority, dm_scope=binding.dm_scope, last_route_policy=last_route_policy, ) def _build_default_result(self, msg: MessageContext) -> RouteResult: agent_id_str = str(self._default_agent_config_id) if msg.peer_kind == PeerKind.DIRECT: session_key = self._session_builder.build_dm_session( agent_id_str, msg.channel, msg.account_id, msg.peer_id, dm_scope="per-channel-peer", ) else: session_key = self._session_builder.build_group_session( agent_id_str, msg.channel, msg.peer_kind.value, msg.peer_id, ) main_session_key = self._session_builder.build_main_session(agent_id_str) last_route_policy = "main" if session_key == main_session_key else "session" return RouteResult( agent_config_id=self._default_agent_config_id, channel=msg.channel, account_id=msg.account_id, session_key=session_key, main_session_key=main_session_key, matched_by="default", matched_priority=9, last_route_policy=last_route_policy, ) def resolve_agent_id(self, agent_id: str) -> str: return SessionKeyBuilder.normalize_agent_id(agent_id)