from __future__ import annotations import json import logging import os import tempfile import threading import time from collections import OrderedDict from dataclasses import dataclass, field logger = logging.getLogger(__name__) DEFAULT_PERSIST_DIR = os.path.join(tempfile.gettempdir(), "yuxi_qqbot_known_users") def _resolve_known_users_dir(persist_dir: str | None = None) -> str: if persist_dir: return persist_dir for env_var in ("OPENCLAW_STATE_DIR", "CLAWDBOT_STATE_DIR"): env_val = os.environ.get(env_var) if env_val: qqbot_dir = os.path.join(env_val, "qqbot", "known_users") os.makedirs(qqbot_dir, exist_ok=True) return qqbot_dir return DEFAULT_PERSIST_DIR @dataclass class UserRecord: user_id: str username: str = "" first_seen: float = field(default_factory=time.time) last_seen: float = field(default_factory=time.time) message_count: int = 0 chat_types: set[str] = field(default_factory=set) metadata: dict = field(default_factory=dict) def touch(self, username: str = "", chat_type: str = "") -> None: self.last_seen = time.time() self.message_count += 1 if username: self.username = username if chat_type: self.chat_types.add(chat_type) def to_dict(self) -> dict: return { "user_id": self.user_id, "username": self.username, "first_seen": self.first_seen, "last_seen": self.last_seen, "message_count": self.message_count, "chat_types": list(self.chat_types), "metadata": self.metadata, } @classmethod def from_dict(cls, data: dict) -> UserRecord: return cls( user_id=data["user_id"], username=data.get("username", ""), first_seen=data.get("first_seen", time.time()), last_seen=data.get("last_seen", time.time()), message_count=data.get("message_count", 0), chat_types=set(data.get("chat_types", [])), metadata=data.get("metadata", {}), ) class KnownUserTracker: def __init__( self, app_id: str, max_users: int = 10000, persist_dir: str | None = None, persist_interval_s: int = 300, ): self._app_id = app_id self._max_users = max_users self._persist_dir = _resolve_known_users_dir(persist_dir) self._persist_path = os.path.join(self._persist_dir, f"{app_id}_users.json") self._persist_interval = persist_interval_s self._lock = threading.RLock() self._users: OrderedDict[str, UserRecord] = OrderedDict() self._last_persist: float = 0 self._dirty = False self._load_from_disk() def record(self, user_id: str, username: str = "", chat_type: str = "") -> UserRecord: with self._lock: if user_id in self._users: self._users.move_to_end(user_id) record = self._users[user_id] record.touch(username, chat_type) else: record = UserRecord(user_id=user_id, username=username) record.touch(username, chat_type) self._users[user_id] = record self._users.move_to_end(user_id) while len(self._users) > self._max_users: self._users.popitem(last=False) self._dirty = True self._maybe_persist() return record def is_known(self, user_id: str) -> bool: with self._lock: return user_id in self._users def get(self, user_id: str) -> UserRecord | None: with self._lock: return self._users.get(user_id) def remove(self, user_id: str) -> bool: with self._lock: if user_id in self._users: del self._users[user_id] self._dirty = True return True return False @property def count(self) -> int: with self._lock: return len(self._users) def persist(self) -> bool: with self._lock: if not self._dirty: return True return self._do_persist() def _maybe_persist(self) -> None: now = time.time() if now - self._last_persist < self._persist_interval: return if not self._dirty: return self._do_persist() def _do_persist(self) -> bool: try: os.makedirs(self._persist_dir, exist_ok=True) data = { "app_id": self._app_id, "updated_at": time.time(), "users": {uid: u.to_dict() for uid, u in self._users.items()}, } tmp_path = self._persist_path + ".tmp" with open(tmp_path, "w", encoding="utf-8") as f: json.dump(data, f, ensure_ascii=False) os.replace(tmp_path, self._persist_path) self._last_persist = time.time() self._dirty = False logger.debug("KnownUserTracker: persisted %d users", len(self._users)) return True except OSError: logger.exception("KnownUserTracker: persist failed") return False def _load_from_disk(self) -> None: try: if not os.path.exists(self._persist_path): return with open(self._persist_path, encoding="utf-8") as f: data = json.load(f) users_data = data.get("users", {}) for uid, udata in users_data.items(): self._users[uid] = UserRecord.from_dict(udata) logger.info("KnownUserTracker: loaded %d users from disk", len(self._users)) except (OSError, json.JSONDecodeError): logger.exception("KnownUserTracker: failed to load from disk") def clear(self) -> None: with self._lock: self._users.clear() self._dirty = True self.persist() def get_recent_users(self, limit: int = 50) -> list[UserRecord]: with self._lock: return list(reversed(self._users.values()))[:limit] def list_known_users(self) -> list[dict]: with self._lock: return [u.to_dict() for u in self._users.values()] def get_known_users_stats(self) -> dict: now = time.time() active_24h = 0 active_7d = 0 dm_count = 0 group_count = 0 with self._lock: for u in self._users.values(): if now - u.last_seen <= 86400: active_24h += 1 if now - u.last_seen <= 604800: active_7d += 1 if "direct" in u.chat_types: dm_count += 1 if "group" in u.chat_types: group_count += 1 return { "total": len(self._users), "active_24h": active_24h, "active_7d": active_7d, "dm_users": dm_count, "group_users": group_count, } def get_user_groups(self, user_id: str) -> list[str]: record = self.get(user_id) if record is None: return [] return sorted(record.chat_types) def get_group_members(self) -> dict[str, list[str]]: result: dict[str, list[str]] = {} with self._lock: for u in self._users.values(): for ct in u.chat_types: if ct not in result: result[ct] = [] result[ct].append(u.user_id) return result def flush_known_users(self) -> int: count = 0 with self._lock: count = len(self._users) self._users.clear() self._dirty = True self.persist() return count