ForcePilot/backend/package/yuxi/channel/extensions/wechat-kf/session.py

import time

from yuxi.channel.extensions.wechat_kf.types import KFAccount, KFSession, KFSessionStatus


class KFSessionManager:
    def __init__(self, account: KFAccount):
        self._account = account
        self._sessions: dict[str, KFSession] = {}
        self._last_cleanup: float = time.time()

    def _session_key(self, open_kfid: str, external_user_id: str) -> str:
        return f"{open_kfid}:{external_user_id}"

    def get_session(self, open_kfid: str, external_user_id: str) -> KFSession | None:
        self._cleanup_expired()
        return self._sessions.get(self._session_key(open_kfid, external_user_id))

    def create_session(self, open_kfid: str, external_user_id: str, session_id: str = "") -> KFSession:
        session = KFSession(
            session_id=session_id or f"sess_{int(time.time() * 1000)}",
            kf_account_id=open_kfid,
            external_user_id=external_user_id,
            status=KFSessionStatus.ACTIVE,
            created_at=time.time(),
            last_user_msg_at=time.time(),
        )
        self._sessions[self._session_key(open_kfid, external_user_id)] = session
        return session

    def mark_user_message(self, open_kfid: str, external_user_id: str):
        session = self.get_session(open_kfid, external_user_id)
        if session:
            session.last_user_msg_at = time.time()
            session.msg_count_in_round = 0
            if session.status == KFSessionStatus.CLOSED:
                session.status = KFSessionStatus.ACTIVE

    def mark_reply_sent(self, open_kfid: str, external_user_id: str):
        session = self.get_session(open_kfid, external_user_id)
        if session:
            session.last_reply_at = time.time()
            session.msg_count_in_round += 1

    def close_session(self, open_kfid: str, external_user_id: str):
        session = self.get_session(open_kfid, external_user_id)
        if session:
            session.status = KFSessionStatus.CLOSED

    def can_send_message(self, open_kfid: str, external_user_id: str) -> tuple[bool, str]:
        session = self.get_session(open_kfid, external_user_id)
        if not session:
            return False, "无活跃会话"
        if session.status == KFSessionStatus.CLOSED:
            return False, "会话已关闭"
        now = time.time()
        if now - session.last_user_msg_at > 48 * 3600:
            return False, "已超过48小时下发窗口"
        if session.msg_count_in_round >= 5:
            return False, "本轮消息数已达上限(5条)"
        return True, ""

    def update_session(self, session: KFSession):
        key = self._session_key(session.kf_account_id, session.external_user_id)
        self._sessions[key] = session

    def _cleanup_expired(self):
        now = time.time()
        if now - self._last_cleanup < 60:
            return
        self._last_cleanup = now
        ttl = self._account.session_ttl
        expired_keys = [
            key
            for key, sess in self._sessions.items()
            if sess.status == KFSessionStatus.CLOSED and now - sess.created_at > ttl
        ]
        for key in expired_keys:
            del self._sessions[key]