from __future__ import annotations import time from typing import Any from yuxi.channels.models import ChatType DEFAULT_SESSION_TTL_SEC = 3600 DEFAULT_BIND_TTL_SEC = 86400 class SessionRouter: def __init__(self, ttl_sec: int = DEFAULT_SESSION_TTL_SEC): self._ttl_sec = ttl_sec self._sessions: dict[str, dict[str, Any]] = {} self._bindings: dict[str, dict[str, Any]] = {} def resolve_thread_key(self, agent_id: str, follower_id: str, account_id: str = "default") -> str: return f"agent:{agent_id}:zalo_oa:{account_id}:{ChatType.DIRECT.value}:{follower_id}" def resolve_bind_key(self, agent_id: str, account_id: str = "default") -> str: return f"agent:{agent_id}:zalo_oa:{account_id}" def resolve_session_params( self, source_id: str, sender_name: str = "", avatar: str = "", dm_policy: str = "", account_id: str = "default", ) -> dict[str, Any]: params: dict[str, Any] = { "chat_type": ChatType.DIRECT, "channel_chat_id": source_id, } if sender_name: params["sender_name"] = sender_name if avatar: params["avatar"] = avatar if dm_policy: params["dm_policy"] = dm_policy if account_id: params["account_id"] = account_id return params def get_session(self, thread_key: str) -> dict[str, Any] | None: entry = self._sessions.get(thread_key) if entry is None: return None if time.time() - entry.get("created_at", 0) > self._ttl_sec: del self._sessions[thread_key] return None return entry.get("data") def set_session(self, thread_key: str, data: dict[str, Any]): self._sessions[thread_key] = { "data": data, "created_at": time.time(), } def clear_session(self, thread_key: str): self._sessions.pop(thread_key, None) def bind_agent_account( self, agent_id: str, account_id: str, bind_data: dict[str, Any] | None = None, ): bind_key = self.resolve_bind_key(agent_id, account_id) self._bindings[bind_key] = { "agent_id": agent_id, "account_id": account_id, "data": bind_data or {}, "created_at": time.time(), "ttl_sec": DEFAULT_BIND_TTL_SEC, } def unbind_agent_account(self, agent_id: str, account_id: str): bind_key = self.resolve_bind_key(agent_id, account_id) self._bindings.pop(bind_key, None) def is_bound(self, agent_id: str, account_id: str) -> bool: bind_key = self.resolve_bind_key(agent_id, account_id) entry = self._bindings.get(bind_key) if entry is None: return False ttl = entry.get("ttl_sec", DEFAULT_BIND_TTL_SEC) if time.time() - entry.get("created_at", 0) > ttl: del self._bindings[bind_key] return False return True def get_binding(self, agent_id: str, account_id: str) -> dict[str, Any] | None: bind_key = self.resolve_bind_key(agent_id, account_id) entry = self._bindings.get(bind_key) if entry is None: return None ttl = entry.get("ttl_sec", DEFAULT_BIND_TTL_SEC) if time.time() - entry.get("created_at", 0) > ttl: del self._bindings[bind_key] return None return entry def get_bound_sessions(self, agent_id: str, account_id: str) -> list[dict[str, Any]]: bind_key = self.resolve_bind_key(agent_id, account_id) prefix = f"{bind_key}:{ChatType.DIRECT.value}:" sessions = [] now = time.time() for key, entry in self._sessions.items(): if key.startswith(prefix): if now - entry.get("created_at", 0) <= self._ttl_sec: session_data = entry.get("data", {}) if session_data: sessions.append(session_data) return sessions def cleanup_expired(self) -> int: now = time.time() expired_sessions = [k for k, v in self._sessions.items() if now - v.get("created_at", 0) > self._ttl_sec] for k in expired_sessions: del self._sessions[k] expired_bindings = [ k for k, v in self._bindings.items() if now - v.get("created_at", 0) > v.get("ttl_sec", DEFAULT_BIND_TTL_SEC) ] for k in expired_bindings: del self._bindings[k] return len(expired_sessions) + len(expired_bindings) @property def session_count(self) -> int: return len(self._sessions) @property def binding_count(self) -> int: return len(self._bindings) _router = SessionRouter() def resolve_thread_key(agent_id: str, follower_id: str, account_id: str = "default") -> str: return _router.resolve_thread_key(agent_id, follower_id, account_id) def resolve_session_params( source_id: str, sender_name: str = "", avatar: str = "", dm_policy: str = "", account_id: str = "default", ) -> dict[str, Any]: return _router.resolve_session_params(source_id, sender_name, avatar, dm_policy, account_id) def bind_agent_account(agent_id: str, account_id: str, bind_data: dict[str, Any] | None = None): _router.bind_agent_account(agent_id, account_id, bind_data) def unbind_agent_account(agent_id: str, account_id: str): _router.unbind_agent_account(agent_id, account_id) def is_bound(agent_id: str, account_id: str) -> bool: return _router.is_bound(agent_id, account_id) def get_session_router() -> SessionRouter: return _router