"""Live database reads on one daemon worker; GUI never waits for key acquisition.""" from concurrent.futures import Future, TimeoutError import queue import threading import time from runtime_paths import application_data_dir from wxwork_db import DatabaseReadError, WXWorkDB, detect_wxwork_dir, load_keys class LiveReplyDatabase: def __init__(self, *, factory=None, read_timeout=0.35, retry_seconds=15.0): self._factory = factory or self._open_database self._read_timeout = max(0.0, float(read_timeout)) self._retry_seconds = max(0.0, float(retry_seconds)) self._jobs = queue.Queue(maxsize=8) self._pending = {} self._lock = threading.RLock() self._stop = threading.Event() self._thread = None self._ready = False self.last_error = "数据库等待初始化" self._since = time.time() - 120.0 @staticmethod def _open_database(*, account="", initialize=True, source_directory=None, load_metadata=True): from wxwork_local_setup import acquire_local_keys, select_source_directory # The protocol probe already located this account; reuse only that path, # never a readiness result. health_check below still verifies the account. source = source_directory if source_directory is not None else (detect_wxwork_dir(account=account) if account else detect_wxwork_dir()) if not source: raise DatabaseReadError("未找到当前账号的企业微信数据目录，请选择本机聊天数据所在目录") try: keys = load_keys() except ValueError: keys = {} cache = str(application_data_dir() / "wxwork_reply_cache") database = WXWorkDB(source, keys, cache_dir=cache, account=account, load_metadata=load_metadata) if database.health_check(account=account or None): return database database.close() if not initialize: raise DatabaseReadError("已找到当前账号消息库，尚未完成本机初始化；请点击检测本机协议，自动获取并校验本机密钥") select_source_directory(source) acquire_local_keys() database = WXWorkDB(source, load_keys(), cache_dir=cache, account=account, load_metadata=load_metadata) if not database.health_check(account=account or None): database.close() raise DatabaseReadError("本机初始化后，当前账号消息库仍不可读取。请核对数据目录、当前登录账号及本机密钥") return database def _run(self): database = None retry_at = 0.0 checkpoint = None try: while not self._stop.is_set(): try: key, operation, args, kwargs, future = self._jobs.get(timeout=0.25) except queue.Empty: continue if not future.set_running_or_notify_cancel(): continue if database is None and time.monotonic() < retry_at: future.set_exception(DatabaseReadError(self.last_error)) continue try: if database is None: database = self._factory() if checkpoint is not None and isinstance(database, WXWorkDB): database._message_cursors = {key: dict(value) for key, value in checkpoint.items()} database._bootstrap_since = self._since if self._stop.is_set(): raise DatabaseReadError("数据库读取已停止") result = getattr(database, operation)(*args, **kwargs) self._ready = True self.last_error = "" future.set_result((time.monotonic(), result)) except Exception as exc: self._ready = False self.last_error = str(exc) future.set_exception(DatabaseReadError(str(exc))) if database is not None: if isinstance(database, WXWorkDB): checkpoint = {key: dict(value) for key, value in database._message_cursors.items()} database.close() database = None retry_at = time.monotonic() + self._retry_seconds finally: if database is not None: database.close() self._ready = False def _request(self, key, operation, *args, **kwargs): with self._lock: if self._stop.is_set(): raise DatabaseReadError("数据库读取已停止") if self._thread is None: self._thread = threading.Thread(target=self._run, daemon=True, name="reply-database-reader") self._thread.start() # Do not retain completed snapshots for contacts that are never revisited. for pending_key, pending_future in list(self._pending.items()): if pending_key == "poll" or pending_key == key or not pending_future.done(): continue try: completed_at, _ = pending_future.result() expired = time.monotonic() - completed_at > 2.0 except Exception: expired = True if expired: self._pending.pop(pending_key, None) future = self._pending.get(key) # 文字快照过期时重新读；增量消息结果必须消费，不能丢弃已推进的游标。 if future is not None and future.done() and operation != "get_new_messages": try: completed_at, _ = future.result() if time.monotonic() - completed_at > 2.0: self._pending.pop(key, None) future = None except Exception: self._pending.pop(key, None) future = None if future is None: future = Future() try: self._jobs.put_nowait((key, operation, args, kwargs, future)) except queue.Full as exc: raise DatabaseReadError("数据库正在读取其他消息") from exc self._pending[key] = future consumed = False try: _, result = future.result(timeout=self._read_timeout) consumed = True except TimeoutError as exc: raise DatabaseReadError("数据库正在初始化或刷新，暂用备用读取") from exc except Exception: consumed = True raise finally: if consumed: with self._lock: if self._pending.get(key) is future: self._pending.pop(key, None) return result def get_new_messages(self, since_ts): # 首次启动的时间边界固定，恢复连接时不会被其他账号的较晚时间推走。 return self._request("poll", "get_new_messages", self._since) def get_conversation_context(self, display_name, *, account="", conv_id="", limit=100): key = ("context", str(display_name), str(account), str(conv_id)) try: return self._request(key, "get_conversation_context", display_name, account=account, conv_id=conv_id, limit=limit) except DatabaseReadError: return None def get_conversation_context_by_id(self, account, conv_id, *, limit=100): if not isinstance(account, str) or not account.strip(): raise ValueError("读取协议会话必须提供账号 ID") if not isinstance(conv_id, str) or not conv_id.strip(): raise ValueError("读取协议会话必须提供会话 ID") limit = max(1, min(500, int(limit))) key = ("context_by_id", account, conv_id, limit) # Preserve errors/timeouts; protocol readers must never trigger clipboard fallback. return self._request(key, "get_conversation_context_by_id", account, conv_id, limit=limit) def health_check(self): return self._ready and not self._stop.is_set() def close(self): self._stop.set() with self._lock: for future in self._pending.values(): future.cancel() self._pending.clear()