xieyi-weixin-/reply_database.py

"""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, cache_dir=None):
        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(cache_dir if cache_dir is not None else 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()