ForcePilot/backend/package/yuxi/channel/routing/session_key.py

import hashlib
import logging
import re
import uuid

logger = logging.getLogger(__name__)

DM_SCOPE_MAIN = "main"
DM_SCOPE_PER_PEER = "per-peer"
DM_SCOPE_PER_CHANNEL_PEER = "per-channel-peer"
DM_SCOPE_PER_ACCOUNT_CHANNEL_PEER = "per-account-channel-peer"

_VALID_DM_SCOPES = frozenset({
    DM_SCOPE_MAIN,
    DM_SCOPE_PER_PEER,
    DM_SCOPE_PER_CHANNEL_PEER,
    DM_SCOPE_PER_ACCOUNT_CHANNEL_PEER,
})

_NAMESPACE = uuid.UUID("a1b2c3d4-e5f6-7890-abcd-ef1234567890")
_SESSION_KEY_RE = re.compile(r"^agent:([^:]+):(.+)$")
_AGENT_ID_RE = re.compile(r"^[a-z0-9][a-z0-9_-]{0,63}$")


class SessionKeyBuilder:
    def __init__(self, identity_links: dict[str, list[str]] | None = None):
        self._identity_links = identity_links or {}
        self._linked_id_lower: dict[str, str] = {}
        for identity, ids in self._identity_links.items():
            for linked_id in ids:
                lid_lower = linked_id.lower()
                if lid_lower not in self._linked_id_lower:
                    self._linked_id_lower[lid_lower] = identity

    @staticmethod
    def normalize_agent_id(agent_id: str) -> str:
        aid = agent_id.strip().lower()
        if not aid:
            raise ValueError("agent_id must not be empty")
        if _AGENT_ID_RE.match(aid):
            return aid
        normalized = re.sub(r"[^a-z0-9_-]", "-", aid)
        if normalized and normalized[0].isdigit():
            normalized = "a-" + normalized
        if len(normalized) > 64:
            logger.warning("Agent ID truncated from %d to 64 chars: %r", len(normalized), normalized)
        return normalized[:64]

    @staticmethod
    def normalize_thread_id(thread_id: str) -> str:
        tid = thread_id.strip().lower()
        return re.sub(r"[^a-z0-9_-]", "-", tid)[:64]

    @staticmethod
    def _sanitize_key_part(value: str) -> str:
        return re.sub(r"[^a-z0-9_.@-]", "-", value.strip().lower())[:200]

    def build_main_session(self, agent_id: str) -> str:
        return f"agent:{self.normalize_agent_id(agent_id)}:main"

    def build_group_session(self, agent_id: str, channel: str, peer_kind: str, peer_id: str) -> str:
        aid = self.normalize_agent_id(agent_id)
        ch = self._sanitize_key_part(channel)
        pk = self._sanitize_key_part(peer_kind)
        pid = self._sanitize_key_part(peer_id)
        return f"agent:{aid}:{ch}:{pk}:{pid}"

    def build_dm_session(
        self,
        agent_id: str,
        channel: str,
        account_id: str,
        peer_id: str,
        dm_scope: str = DM_SCOPE_PER_CHANNEL_PEER,
    ) -> str:
        if dm_scope not in _VALID_DM_SCOPES:
            raise ValueError(f"Invalid dm_scope: {dm_scope!r}")
        aid = self.normalize_agent_id(agent_id)
        ch = self._sanitize_key_part(channel)
        acc = self._sanitize_key_part(account_id)
        resolved_peer_id = self._resolve_linked_peer_id(channel, peer_id, dm_scope)
        pid = self._sanitize_key_part(resolved_peer_id)

        if dm_scope == DM_SCOPE_MAIN:
            return f"agent:{aid}:main"
        if dm_scope == DM_SCOPE_PER_PEER:
            return f"agent:{aid}:direct:{pid}"
        if dm_scope == DM_SCOPE_PER_CHANNEL_PEER:
            return f"agent:{aid}:{ch}:direct:{pid}"
        return f"agent:{aid}:{ch}:{acc}:direct:{pid}"

    def build_thread_key(self, base_key: str, thread_id: str) -> str:
        tid = self.normalize_thread_id(thread_id)
        return f"{base_key}:thread:{tid}"

    def build_group_history_key(self, channel: str, account_id: str, peer_kind: str, peer_id: str) -> str:
        ch = self._sanitize_key_part(channel)
        acc = self._sanitize_key_part(account_id)
        pk = self._sanitize_key_part(peer_kind)
        pid = self._sanitize_key_part(peer_id)
        return f"{ch}:{acc}:{pk}:{pid}"

    def _resolve_linked_peer_id(self, channel: str, peer_id: str, dm_scope: str) -> str:
        if dm_scope == DM_SCOPE_MAIN or not self._identity_links:
            return peer_id

        result = self._linked_id_lower.get(peer_id.lower())
        if result is not None:
            return result
        result = self._linked_id_lower.get(f"{channel}:{peer_id}".lower())
        if result is not None:
            return result

        return peer_id

    @staticmethod
    def parse_agent_id(session_key: str) -> str | None:
        m = _SESSION_KEY_RE.match(session_key)
        return m.group(1) if m else None

    @staticmethod
    def build_command_target_key(
        agent_id: str,
        channel_type: str,
        target_session_part: str,
    ) -> str:
        """构建原生命令跨会话路由 Key

        当在某个会话中执行命令需要路由到另一个会话时使用。
        例如: 在群聊 A 中 /switch B → agent:agent_id:cmd-target:channel:target
        """
        aid = SessionKeyBuilder.normalize_agent_id(agent_id)
        target = SessionKeyBuilder.normalize_thread_id(target_session_part)
        return f"agent:{aid}:cmd-target:{channel_type}:{target}"

    @staticmethod
    def parse_command_target_session_key(session_key: str) -> str | None:
        """从 SessionKey 中解析出 CommandTarget 的目标部分"""
        if ":cmd-target:" not in session_key:
            return None
        parts = session_key.split(":")
        cmd_idx = parts.index("cmd-target")
        remaining = parts[cmd_idx + 1:]
        if not remaining:
            return None
        return ":".join(remaining)


def session_key_to_thread_id(session_key: str) -> str:
    sha = hashlib.sha256(session_key.encode("utf-8")).digest()
    return str(uuid.uuid5(_NAMESPACE, sha.hex()))


def resolve_thread_session_keys(base_key: str, thread_id: str) -> dict[str, str]:
    tid = SessionKeyBuilder.normalize_thread_id(thread_id)
    return {
        "session_key": f"{base_key}:thread:{tid}",
        "parent_session_key": base_key,
    }