From 5548f8d89f2cab0f1950ce59805d64ecf965a22e Mon Sep 17 00:00:00 2001 From: Kris <2893855659@qq.com> Date: Tue, 7 Jul 2026 20:17:06 +0800 Subject: [PATCH] feat(server,db_reader): add SSE real-time message push feature This commit implements the full SSE real-time message push capability: 1. Add new `/api/messages/stream` SSE endpoint for real-time message delivery 2. Add MessageStreamer backend that polls DB with mtime awareness and broadcasts events 3. Add helper methods in DbReader: get_db_mtime and get_max_create_time 4. Optimize existing DB connections with `isolation_level=None` 5. Update bridge version to 1.5.0 and add `sse_push` to capability list 6. Add proper error handling and fallback logic for DB unavailability --- bridge/db_reader.py | 63 +++++++- bridge/message_streamer.py | 322 +++++++++++++++++++++++++++++++++++++ bridge/server.py | 120 +++++++++++++- 3 files changed, 492 insertions(+), 13 deletions(-) create mode 100644 bridge/message_streamer.py diff --git a/bridge/db_reader.py b/bridge/db_reader.py index 51836a4..c6cd116 100644 --- a/bridge/db_reader.py +++ b/bridge/db_reader.py @@ -444,6 +444,30 @@ class DbReader: """DB 是否可读。""" return self.check_db_status() == "ok" + def get_db_mtime(self, rel_path: str) -> Optional[tuple[float, float]]: + """返回指定 DB 及其 WAL 文件的 mtime，用于增量轮询感知变化。 + + Args: + rel_path: 相对路径，如 message/message_0.db + + Returns: + (db_mtime, wal_mtime)：DB 不存在时返回 None； + WAL 不存在时 wal_mtime 为 0.0；mtime 读取失败记为 0.0。 + """ + db_path = self._get_db_abs_path(rel_path) + if db_path is None: + return None + try: + db_mtime = os.path.getmtime(db_path) + except OSError: + db_mtime = 0.0 + wal_path = db_path + "-wal" + try: + wal_mtime = os.path.getmtime(wal_path) if os.path.exists(wal_path) else 0.0 + except OSError: + wal_mtime = 0.0 + return db_mtime, wal_mtime + # ------------------------------------------------------------------ # 辅助：表/列探测 # ------------------------------------------------------------------ @@ -511,7 +535,7 @@ class DbReader: result["wxid"] = wxid or "" try: db_path = self._ensure_decrypted("contact/contact.db") - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: table = self._find_contact_table(conn) @@ -573,7 +597,7 @@ class DbReader: def get_messages_since(self, cursor: int, limit: int = 50) -> dict: """读取 create_time > cursor 的增量消息。""" db_path = self._ensure_decrypted("message/message_0.db") - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: cur = conn.execute( @@ -602,6 +626,31 @@ class DbReader: "has_more": has_more, } + def get_max_create_time(self) -> Optional[int]: + """返回 message 表中最大的 create_time，用于初始化推送游标。 + + DB 不可读或 message 表不存在时返回 None。供 MessageStreamer 在启动时 + 把全局游标对齐到当前最新消息，避免向新连接推送全量历史。 + """ + try: + db_path = self._ensure_decrypted("message/message_0.db") + except BridgeError: + return None + conn = sqlite3.connect(db_path, isolation_level=None) + try: + cur = conn.execute( + "SELECT name FROM sqlite_master WHERE type='table' AND name='message'" + ) + if cur.fetchone() is None: + return None + cur = conn.execute("SELECT MAX(create_time) FROM message") + row = cur.fetchone() + return int(row[0]) if row and row[0] is not None else None + except Exception: + return None + finally: + conn.close() + # ------------------------------------------------------------------ # 联系人字段映射（公共） # ------------------------------------------------------------------ @@ -692,7 +741,7 @@ class DbReader: except BridgeError: return empty - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: table = self._find_contact_table(conn) @@ -749,7 +798,7 @@ class DbReader: except BridgeError: return None - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: table = self._find_contact_table(conn) @@ -783,7 +832,7 @@ class DbReader: except BridgeError: return empty - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: table = self._find_contact_table(conn) @@ -815,7 +864,7 @@ class DbReader: except BridgeError: return empty - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: table = self._find_table_by_name( @@ -923,7 +972,7 @@ class DbReader: except BridgeError: return None - conn = sqlite3.connect(db_path) + conn = sqlite3.connect(db_path, isolation_level=None) conn.row_factory = sqlite3.Row try: cols = self._table_columns(conn, "message") diff --git a/bridge/message_streamer.py b/bridge/message_streamer.py new file mode 100644 index 0000000..7b67e24 --- /dev/null +++ b/bridge/message_streamer.py @@ -0,0 +1,322 @@ +"""消息流推送器：内部轮询 message DB 并通过 SSE 广播给订阅者。 + +微信本地 DB 是 SQLite，没有 pub/sub 机制，bridge 只能轮询。本模块把轮询 +从客户端移到服务端：后台 watcher 周期性检查 message_0.db 的 mtime，变化时 +读取增量消息并广播给所有 SSE 订阅者。 + +优化点： +- 无订阅者时降频到 5s（仍维持全局游标，避免新连接被历史消息淹没） +- mtime 感知：DB/WAL mtime 未变化时跳过 SQL 查询 +- 批量读取（默认 200 条），消息洪峰时减少广播次数 +- 全局游标对齐当前最新 create_time，新连接只收"从现在起"的新消息 +- 背压：订阅队列满时丢弃旧事件并记日志，避免单客户端拖垮全局 + +事件类型（SSE）： +- sync：连接建立时发送，data.cursor 为当前全局游标 +- messages：一批新消息，data 含 messages/next_cursor/has_more +- status：DB 状态变化，data 含 db_accessible/db_error_code +- heartbeat：30 秒心跳，保活 NAT/代理连接 +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import time +from dataclasses import dataclass +from typing import Awaitable, Callable, Optional + +logger = logging.getLogger("woc-bridge") + +# 有订阅者时的轮询间隔（秒） +POLL_INTERVAL_ACTIVE = 1.0 +# 无订阅者时的轮询间隔（秒，仅维持全局游标） +POLL_INTERVAL_IDLE = 5.0 +# 单次读取上限 +BATCH_LIMIT = 200 +# 订阅队列容量（背压保护） +QUEUE_MAXSIZE = 100 +# 心跳间隔（秒） +HEARTBEAT_INTERVAL = 30.0 + +# DB 状态解析器返回类型：(db_accessible, db_error_code) +DbStateResolver = Callable[[], Awaitable[tuple[bool, Optional[str]]]] + + +@dataclass +class StreamEvent: + """SSE 事件。 + + Attributes: + event: 事件类型（sync / messages / status / heartbeat） + data: 事件数据（序列化为 JSON） + id: 可选事件 ID，客户端断线重连时可用 Last-Event-ID 补全 + """ + + event: str + data: dict + id: Optional[str] = None + + +def format_sse(event: StreamEvent) -> str: + """把 StreamEvent 格式化为 SSE 文本帧。 + + 格式遵循 SSE 规范：每行 `field: value`，事件间用空行分隔。 + """ + lines: list[str] = [] + if event.id: + lines.append(f"id: {event.id}") + if event.event: + lines.append(f"event: {event.event}") + lines.append(f"data: {json.dumps(event.data, ensure_ascii=False)}") + return "\n".join(lines) + "\n\n" + + +class MessageStreamer: + """消息流推送器。 + + 维护一个全局游标 global_cursor，后台 watcher 轮询 DB，把新消息以 + StreamEvent 广播给所有订阅者。每个订阅者对应一个 asyncio.Queue， + SSE 路由从队列取事件推给客户端。 + + 使用： + streamer = MessageStreamer(db_reader, lambda: _resolve_db_state_tuple()) + await streamer.start() # lifespan 启动 + q = streamer.subscribe() # SSE 路由订阅 + streamer.unsubscribe(q) # SSE 路由断开 + await streamer.stop() # lifespan 停止 + """ + + def __init__( + self, + db_reader: "DbReader", + resolve_db_state: DbStateResolver, + ) -> None: + """初始化推送器。 + + Args: + db_reader: DbReader 实例（用于读消息与查 mtime） + resolve_db_state: 异步回调，返回 (db_accessible, db_error_code)， + 供 watcher 判断 DB 是否可读并广播状态变化 + """ + self._db_reader = db_reader + self._resolve_db_state = resolve_db_state + self._subscribers: set[asyncio.Queue[StreamEvent]] = set() + self._global_cursor: int = 0 + self._watcher: asyncio.Task | None = None + # DB mtime 缓存：变化时才查 SQL + self._last_db_mtime: float = 0.0 + self._last_wal_mtime: float = 0.0 + # 上次广播的 DB 状态码，变化时发 status 事件 + self._last_db_error_code: Optional[str] = "__init__" + self._cursor_inited: bool = False + + @property + def global_cursor(self) -> int: + """当前全局游标（最近一次推送的 create_time）。""" + return self._global_cursor + + def subscriber_count(self) -> int: + """当前订阅者数量（供 /api/status 暴露给客户端）。""" + return len(self._subscribers) + + def subscribe(self) -> asyncio.Queue[StreamEvent]: + """订阅事件流。 + + 订阅后立即向队列放入一个 sync 事件，携带当前全局游标。客户端据此 + 判断是否需要用 /api/messages/since 补全断线期间消息。 + + Returns: + asyncio.Queue：事件队列，容量 QUEUE_MAXSIZE，满时丢弃旧事件 + """ + q: asyncio.Queue[StreamEvent] = asyncio.Queue(maxsize=QUEUE_MAXSIZE) + self._subscribers.add(q) + # 立即发送 sync 事件（只给这个订阅者） + try: + q.put_nowait(StreamEvent( + event="sync", + data={"cursor": self._global_cursor}, + )) + except asyncio.QueueFull: + pass + logger.info("SSE 订阅，当前订阅者 %d，cursor=%d", self.subscriber_count(), self._global_cursor) + return q + + def unsubscribe(self, q: asyncio.Queue[StreamEvent]) -> None: + """取消订阅。""" + self._subscribers.discard(q) + logger.info("SSE 取消订阅，当前订阅者 %d", self.subscriber_count()) + + async def start(self) -> None: + """启动 watcher 后台任务。""" + if self._watcher is None or self._watcher.done(): + self._watcher = asyncio.create_task(self._watch_loop()) + logger.info("MessageStreamer watcher 已启动") + + async def stop(self) -> None: + """停止 watcher 后台任务。""" + if self._watcher is not None and not self._watcher.done(): + self._watcher.cancel() + try: + await self._watcher + except asyncio.CancelledError: + pass + self._watcher = None + logger.info("MessageStreamer watcher 已停止") + + # ------------------------------------------------------------------ + # 内部：广播与轮询 + # ------------------------------------------------------------------ + def _broadcast(self, event: StreamEvent) -> None: + """广播事件给所有订阅者。 + + 队列满时丢弃最旧事件再放入，避免单客户端消费慢拖垮全局。 + """ + for q in list(self._subscribers): + try: + q.put_nowait(event) + except asyncio.QueueFull: + # 背压：丢掉最旧事件，腾出位置 + try: + q.get_nowait() + except asyncio.QueueEmpty: + pass + try: + q.put_nowait(event) + except asyncio.QueueFull: + logger.warning("SSE 订阅队列持续满，丢弃事件 event=%s", event.event) + + async def _init_cursor(self) -> bool: + """把全局游标对齐到当前 DB 最大 create_time。 + + 避免新连接收到全量历史消息。DB 不可读时保持 0，等可读后再对齐。 + + Returns: + True 表示成功对齐游标（或游标已对齐）；False 表示 DB 不可读， + 需要在后续轮询中重试 + """ + try: + max_ct = await asyncio.to_thread(self._db_reader.get_max_create_time) + if max_ct and max_ct > 0: + self._global_cursor = max_ct + self._cursor_inited = True + logger.info("MessageStreamer 初始游标对齐到 %d", self._global_cursor) + return True + # DB 可读但 message 表为空或返回 None：视为已对齐（cursor=0 合理） + self._cursor_inited = True + return True + except Exception as e: + logger.warning("MessageStreamer 初始化游标失败: %s", e) + return False + + async def _watch_loop(self) -> None: + """watcher 主循环：轮询 DB 并广播新消息。 + + - 首次运行时对齐全局游标 + - 按 POLL_INTERVAL_ACTIVE/IDLE 间隔轮询 + - DB 状态变化时发 status 事件 + - DB mtime 变化时读增量消息并广播 + - 每 HEARTBEAT_INTERVAL 秒发一次心跳 + """ + # 初始化游标（DB 可能尚不可读，失败则后续轮询中重试） + await self._init_cursor() + + last_heartbeat = time.monotonic() + while True: + try: + interval = ( + POLL_INTERVAL_ACTIVE if self._subscribers else POLL_INTERVAL_IDLE + ) + await asyncio.sleep(interval) + + # 心跳（仅有订阅者时发，无订阅者发心跳无意义） + now = time.monotonic() + if self._subscribers and now - last_heartbeat >= HEARTBEAT_INTERVAL: + self._broadcast(StreamEvent(event="heartbeat", data={})) + last_heartbeat = now + + await self._poll_once() + + except asyncio.CancelledError: + raise + except Exception as e: + logger.exception("MessageStreamer watch loop 异常: %s", e) + await asyncio.sleep(2.0) + + async def _poll_once(self) -> None: + """单次轮询：检查 DB 状态 + mtime + 读增量消息。""" + # 1. 解析 DB 状态 + try: + db_accessible, db_error_code = await self._resolve_db_state() + except Exception as e: + logger.warning("MessageStreamer: resolve_db_state 失败: %s", e) + return + + # 2. 状态变化通知 + if db_error_code != self._last_db_error_code: + self._last_db_error_code = db_error_code + self._broadcast(StreamEvent( + event="status", + data={ + "db_accessible": db_accessible, + "db_error_code": db_error_code, + }, + )) + + if not db_accessible: + return + + # 3. 游标未初始化（DB 刚恢复可读）时补对齐 + if not self._cursor_inited: + await self._init_cursor() + return + + # 4. mtime 感知：DB/WAL 未变化则跳过 SQL + mtimes = await asyncio.to_thread( + self._db_reader.get_db_mtime, "message/message_0.db" + ) + if mtimes is None: + return + db_mtime, wal_mtime = mtimes + + if db_mtime == self._last_db_mtime and wal_mtime == self._last_wal_mtime: + return + self._last_db_mtime = db_mtime + self._last_wal_mtime = wal_mtime + + # 5. 读取增量消息（has_more=True 时循环读完剩余批次，不受 mtime 缓存影响） + while True: + try: + result = await asyncio.to_thread( + self._db_reader.get_messages_since, self._global_cursor, BATCH_LIMIT + ) + except Exception as e: + logger.warning("MessageStreamer: get_messages_since 失败: %s", e) + return + + messages = result.get("messages", []) + if not messages: + return + + next_cursor = result.get("next_cursor", self._global_cursor) + has_more = result.get("has_more", False) + self._global_cursor = next_cursor + + # 6. 广播 + self._broadcast(StreamEvent( + event="messages", + id=str(next_cursor), + data={ + "messages": messages, + "next_cursor": next_cursor, + "has_more": has_more, + }, + )) + logger.info( + "MessageStreamer 推送 %d 条消息, cursor=%d, 订阅者=%d", + len(messages), next_cursor, self.subscriber_count(), + ) + + if not has_more: + return diff --git a/bridge/server.py b/bridge/server.py index c0353ba..c214419 100644 --- a/bridge/server.py +++ b/bridge/server.py @@ -28,10 +28,16 @@ - 所有 assert 替换为显式 BridgeError，兼容 python -O - _send_file 文件校验改用 isfile+access - 移除 _read_current_wxid 死代码 +- 1.5.0： + - 新增 GET /api/messages/stream SSE 实时消息推送 + - MessageStreamer 内部轮询 DB（mtime 感知）+ 广播给 SSE 订阅者 + - 替代客户端轮询 /api/messages/since；后者保留作断线补全兜底 + - BRIDGE_CAPABILITIES 追加 sse_push -路由总览（20 个）： +路由总览（21 个）： - GET /api/status 状态聚合（进程/窗口/登录态/DB/版本/能力） - GET /api/messages/since 增量消息拉取（按 create_time 游标） +- GET /api/messages/stream SSE 实时消息推送（长连接，主动推新消息） - POST /api/send/text 发送文本消息（placeholder=false） - POST /api/send/image 发送图片消息（占位实现，placeholder=true） - POST /api/send/file 发送文件消息（占位实现，placeholder=true） @@ -78,9 +84,10 @@ from typing import AsyncIterator, Callable, Optional import uvicorn from fastapi import FastAPI, Request -from fastapi.responses import JSONResponse, Response +from fastapi.responses import JSONResponse, Response, StreamingResponse from db_reader import DbReader +from message_streamer import MessageStreamer, StreamEvent, format_sse from models import ( BridgeError, ConnectivityResponse, @@ -138,17 +145,21 @@ logger.propagate = False # 1.3.0 → 1.4.0：_check_db_readable 非阻塞化；_auto_extract_with_lock 公共加锁提取； # login_qr_start 非阻塞激活；DB 状态 1s 缓存；display_name LRU 缓存； # assert → 显式 BridgeError；_send_file isfile+access；移除死代码 -BRIDGE_VERSION = "1.4.0" +# 1.4.0 → 1.5.0：新增 SSE 实时消息推送（GET /api/messages/stream）； +# MessageStreamer 内部轮询 DB + 广播，替代客户端轮询 /api/messages/since； +# 保留 /api/messages/since 作为断线补全兜底 +BRIDGE_VERSION = "1.5.0" # bridge 能力清单（供客户端通过 /api/status 协商）。 -# 当前已落地能力：text_send（文本发送真实可用）、db_decrypt（SQLCipher DB 解密）。 +# 当前已落地能力：text_send（文本发送真实可用）、db_decrypt（SQLCipher DB 解密）、 +# sse_push（SSE 实时消息推送）。 # 待落地能力（占位声明，未真正实现）： # - image_send：W-02 完成后追加 # - file_send：W-02 完成后追加 -# - long_poll：W-03 完成后追加 +# - long_poll：W-03 完成后追加（已被 sse_push 替代，保留兼容声明） # - avatar：W-04 完成后追加 # - group_admin：W-06 完成后追加 -BRIDGE_CAPABILITIES = ["text_send", "db_decrypt"] +BRIDGE_CAPABILITIES = ["text_send", "db_decrypt", "sse_push"] # 媒体发送是否已实现真实传输（W-02 完成后改 True）。 # 客户端据此决定是否调 /api/send/image 与 /api/send/file。 @@ -299,6 +310,7 @@ class AppState: self.db_reader: DbReader | None = None self.send_queue: SendQueue | None = None self.qr_capture: QrCapture | None = None + self.message_streamer: MessageStreamer | None = None self.start_time: float = 0.0 # DB 解密密钥缓存（env / api / auto_extract / file 注入） self.key_cache = KeyCache() @@ -361,6 +373,26 @@ def _require_qr_capture() -> QrCapture: return _state.qr_capture +def _require_message_streamer() -> MessageStreamer: + """返回消息流推送器，未初始化时抛 BridgeError。""" + if _state.message_streamer is None: + raise BridgeError( + code="BRIDGE_INTERNAL_ERROR", + message="message_streamer 未初始化", + ) + return _state.message_streamer + + +async def _resolve_db_state_tuple() -> tuple[bool, Optional[str]]: + """供 MessageStreamer 使用的 DB 状态解析适配器。 + + 把 _resolve_db_state() 的 DbState 降维为 (db_accessible, db_error_code)， + 供 MessageStreamer 解耦 server.py 内部数据结构。 + """ + state = await _resolve_db_state() + return (state.db_accessible, state.db_error_code) + + # --------------------------------------------------------------------------- # lifespan # --------------------------------------------------------------------------- @@ -380,9 +412,13 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: _state.start_time = time.monotonic() if _state.send_queue is not None: await _state.send_queue.start() + if _state.message_streamer is not None: + await _state.message_streamer.start() try: yield finally: + if _state.message_streamer is not None: + await _state.message_streamer.stop() if _state.send_queue is not None: await _state.send_queue.stop() @@ -1272,6 +1308,72 @@ async def get_messages_since(cursor: int = 0, limit: int = 50) -> MessagesRespon return MessagesResponse(**result) +# --------------------------------------------------------------------------- +# 路由：GET /api/messages/stream（SSE 实时消息推送） +# --------------------------------------------------------------------------- +@app.get("/api/messages/stream") +async def stream_messages(request: Request) -> StreamingResponse: + """SSE 实时消息推送长连接。 + + 建立 text/event-stream 连接后，服务端主动推送新消息事件，替代客户端 + 轮询 /api/messages/since。内部由 MessageStreamer 周期性轮询 message DB + （有订阅者 1s、无订阅者 5s），DB mtime 变化时读增量消息并广播。 + + 事件类型： + - sync：连接建立时立即发送，data.cursor 为当前全局游标。客户端据此 + 判断是否需要用 /api/messages/since 补全断线期间消息。 + - messages：一批新消息，data 含 messages / next_cursor / has_more， + 结构与 /api/messages/since 响应体一致。事件 id 为 next_cursor。 + - status：DB 状态变化时发送，data 含 db_accessible / db_error_code。 + DB 加密/退出/恢复可读时触发，客户端据此决定是否降级到轮询。 + - heartbeat：每 30 秒心跳，保活 NAT/代理连接，data 为空对象。 + + 断线补偿： + SSE 不保证 100% 投递。客户端断线重连后应： + 1. 连接本接口，拿 sync 事件中的 cursor + 2. 用 max(本地cursor, sync.cursor) 调 /api/messages/since 补全 + 3. 补全后继续消费本接口事件流 + + Returns: + StreamingResponse：media_type=text/event-stream，持续推送 SSE 帧 + + Notes: + - DB 加密时不关闭连接，改发 status 事件通知客户端；key 提取成功后 + 自动恢复推送 + - 响应头设置 X-Accel-Buffering: no 关闭 nginx 缓冲，避免事件被攒批 + - 客户端断开时自动取消订阅，释放队列 + - 不抛业务错误：即使 DB 不可读也建立连接，通过 status 事件告知 + """ + streamer = _require_message_streamer() + + queue = streamer.subscribe() + + async def event_generator(): + try: + while True: + if await request.is_disconnected(): + break + try: + event = await asyncio.wait_for(queue.get(), timeout=30.0) + except asyncio.TimeoutError: + # 队列 30s 无事件，发心跳保活 + event = StreamEvent(event="heartbeat", data={}) + yield format_sse(event) + finally: + streamer.unsubscribe(queue) + + headers = { + "Cache-Control": "no-cache", + "X-Accel-Buffering": "no", # nginx 关闭缓冲，确保实时推送 + "Connection": "keep-alive", + } + return StreamingResponse( + event_generator(), + media_type="text/event-stream", + headers=headers, + ) + + # --------------------------------------------------------------------------- # 路由：POST /api/send/text # --------------------------------------------------------------------------- @@ -2373,6 +2475,12 @@ def _init_state(cfg: BridgeConfig) -> None: max_calls_per_sec=cfg.max_calls_per_sec, ) + # 消息流推送器：内部轮询 DB + SSE 广播（lifespan 中 start/stop） + _state.message_streamer = MessageStreamer( + db_reader=_state.db_reader, + resolve_db_state=_resolve_db_state_tuple, + ) + # 读取 WOC_DB_KEY 环境变量，注入 DbReader # 用 set_default_key 而非 set_key：保留 woc-keys.json 已持久化的多 salt 映射， # env key 仅作为默认兜底，首次查询时按 salt 验证后入库