ForcePilot/backend/package/yuxi/channel/session/manager.py

"""渠道会话管理器"""

import uuid
from hashlib import md5

from sqlalchemy import select, text
from sqlalchemy.exc import IntegrityError
from sqlalchemy.ext.asyncio import AsyncSession

from yuxi.channel.plugins.protocol import (
    BindingRoute,
    ChannelPlugin,
    InboundMessage,
    SessionConversationRef,
)
from yuxi.channel.routing.router import BindingRouter
from yuxi.channel.session.user_mapping import build_channel_user_uid
from yuxi.repositories.channel_session_repository import ChannelSessionRepository
from yuxi.repositories.conversation_repository import ConversationRepository
from yuxi.storage.postgres.model_channel import ChannelSession
from yuxi.storage.postgres.models_business import User


class SessionManager:
    def __init__(
        self,
        binding_router: BindingRouter,
    ) -> None:
        self.router = binding_router

    async def resolve(
        self,
        config: dict,
        plugin: ChannelPlugin,
        inbound: InboundMessage,
        db_session: AsyncSession,
    ) -> tuple[ChannelSession, BindingRoute | None]:
        """解析渠道会话，返回会话与首次解析出的路由（新建会话时）。

        若会话已存在，则路由为 None，由调用方按需复用缓存或重新解析。
        """
        # 1. 插件自定义解析
        ref = plugin.resolve_session_conversation(
            inbound.chat_type,
            inbound.peer_id or inbound.sender_id or "",
        )

        # 2. 通用回退
        if ref is None:
            ref = self._fallback_resolve(plugin, inbound)

        # 3. 解析插件级回复与线程策略
        reply_to_mode = plugin.resolve_reply_to_mode(config, inbound)
        auto_thread_id = plugin.resolve_auto_thread_id(config, inbound)
        if reply_to_mode is not None:
            ref.channel_metadata["reply_to_mode"] = reply_to_mode
        if auto_thread_id:
            if auto_thread_id not in ref.parent_conversation_candidates:
                ref.parent_conversation_candidates.append(auto_thread_id)
            ref.channel_metadata["auto_thread_id"] = auto_thread_id
        if ref.parent_conversation_candidates:
            ref.channel_metadata["parent_conversation_candidates"] = ref.parent_conversation_candidates

        session_repo = ChannelSessionRepository(db_session)

        # 4. Postgres 下对 session_key 加事务级咨询锁，串行化新建会话，避免并发冲突
        await self._acquire_session_key_lock(db_session, ref.session_key)

        # 5. 在锁保护下重新查询已有 ChannelSession
        session = await session_repo.get_by_session_key(ref.session_key)
        if session is not None:
            updates: dict = {}
            if reply_to_mode is not None:
                updates["reply_to_mode"] = reply_to_mode
            if auto_thread_id:
                updates["auto_thread_id"] = auto_thread_id
            if ref.parent_conversation_candidates:
                updates["parent_conversation_candidates"] = ref.parent_conversation_candidates
            if updates:
                session.update_channel_metadata(updates)
                await db_session.flush()
            return session, None

        # 6. 先确定 agent_id（必须在创建 Conversation 前）
        route = await self.router.resolve_runtime(config, plugin, ref, db_session)
        agent_id = route.agent_id or config.get("default_agent_id")
        if not agent_id:
            raise ValueError(f"Unable to resolve agent_id for session {ref.session_key}")

        # 7. 创建/获取账户级虚拟用户（复用同一事务）
        channel_user = await self._get_or_create_channel_user(
            channel_type=config["channel_type"],
            account_id=config["account_id"],
            db_session=db_session,
        )

        # 8. 创建 Conversation
        channel_metadata = {
            "account_id": config["account_id"],
            "sender_id": inbound.sender_id,
            "sender_name": inbound.sender_name,
        }
        conversation_repo = ConversationRepository(db_session)
        conversation = await conversation_repo.create_conversation(
            uid=channel_user.uid,
            agent_id=agent_id,
            thread_id=self._session_key_to_thread_id(ref.session_key),
            metadata={
                "channel_type": config["channel_type"],
                "channel_session_id": ref.session_key,
                "channel_metadata": channel_metadata,
            },
            channel_type=config["channel_type"],
            channel_session_id=ref.session_key,
            channel_metadata=channel_metadata,
            auto_commit=False,
        )

        # 9. 创建 ChannelSession（非 Postgres 环境仍可能冲突，兜底重取并丢弃当前路由）
        try:
            session = await session_repo.create(
                session_key=ref.session_key,
                channel_type=config["channel_type"],
                account_id=config["account_id"],
                chat_type=inbound.chat_type,
                channel_sender_id=inbound.sender_id,
                conversation_id=conversation.id,
                channel_metadata=ref.channel_metadata,
                auto_commit=False,
            )
        except IntegrityError:
            await db_session.rollback()
            session = await session_repo.get_by_session_key(ref.session_key)
            if session is None:
                raise
            await db_session.refresh(session)
            return session, None
        return session, route

    async def _acquire_session_key_lock(self, db_session: AsyncSession, session_key: str) -> None:
        """在 Postgres 下对 session_key 加事务级咨询锁，避免并发新建同一会话。

        非 Postgres 方言不执行任何操作，依赖外层 IntegrityError 兜底。
        """
        bind = db_session.bind
        if bind is None or bind.dialect.name != "postgresql":
            return
        lock_id = self._session_key_lock_id(session_key)
        await db_session.execute(text("SELECT pg_advisory_xact_lock(:lock_id)").bindparams(lock_id=lock_id))

    def _session_key_lock_id(self, session_key: str) -> int:
        """把 session_key 映射为 64 位有符号整数，供 pg_advisory_xact_lock 使用。"""
        return int(md5(session_key.encode()).hexdigest()[:15], 16)

    async def record_message(
        self,
        session: ChannelSession,
        inbound: InboundMessage,
        db_session: AsyncSession,
    ) -> None:
        """更新 ChannelSession 最后消息时间等轻量元数据。"""
        session_repo = ChannelSessionRepository(db_session)
        await session_repo.update_last_message_at(
            session,
            channel_message_id=inbound.channel_message_id,
        )

    async def _get_or_create_channel_user(
        self,
        channel_type: str,
        account_id: str,
        db_session: AsyncSession,
    ) -> User:
        uid = build_channel_user_uid(channel_type, account_id)
        result = await db_session.execute(select(User).where(User.uid == uid))
        user = result.scalar_one_or_none()
        if user is not None:
            return user

        user = User(
            uid=uid,
            username=f"{channel_type}:{account_id}",
            password_hash="",
            role="user",
            is_channel_user=True,
            channel_type=channel_type,
        )
        db_session.add(user)
        try:
            await db_session.flush()
        except IntegrityError:
            await db_session.rollback()
            result = await db_session.execute(select(User).where(User.uid == uid))
            user = result.scalar_one_or_none()
            if user is None:
                raise
        return user

    def _fallback_resolve(self, plugin: ChannelPlugin, inbound: InboundMessage) -> SessionConversationRef:
        return SessionConversationRef(
            session_key=plugin.parse_session_key(inbound),
            chat_type=inbound.chat_type,
            channel_sender_id=inbound.sender_id,
            channel_metadata={"fallback": True},
        )

    def _session_key_to_thread_id(self, session_key: str) -> str:
        return uuid.uuid5(uuid.NAMESPACE_OID, session_key).hex