WechatOnCloud/bridge/woc_bridge/messaging/verify_bus.py

"""SSE 发送结果校验总线。

发送文本后，发送方注册一个 (to_wxid, content) 等待；MessageStreamer
广播新消息时调用 notify，匹配到等待即 resolve 为 True。超时自动清理
（默认 2.5s），避免内存泄漏。

设计要点：
- 完全异步、单线程事件循环（asyncio.Future）
- 精确匹配 talker == to_wxid 且 is_sender=True（自己发出的消息）
  且 content 包含发送内容（微信可能加前缀/后缀）
- 多发送并发时，每个 (to_wxid, content) 独立 Future，互不干扰
- 防泄漏：register 后无论是否被 notify，超时都会清理并 resolve False
- 不依赖 DB，无熔断器，无雪崩风险
"""

from __future__ import annotations

import asyncio
import logging
import time
from typing import Optional

logger = logging.getLogger("woc-bridge")

# 默认等待超时（秒）：SSE 是推送机制，比 DB 轮询快，
# 但 MessageStreamer 有 1s 轮询间隔 + DB mtime 检测延迟，
# 2.5s 足够覆盖正常场景。
_DEFAULT_TIMEOUT_SEC = 2.5


class SSEVerifyBus:
    """SSE 发送结果校验总线。

    线程安全说明：asyncio.Future 与 dict 操作都在同一事件循环中，
    无需加锁。notify 由 MessageStreamer._broadcast 调用（同一 loop），
    register/wait 由 SendTextFlow 调用（同一 loop）。
    """

    def __init__(self, default_timeout_sec: float = _DEFAULT_TIMEOUT_SEC) -> None:
        self._default_timeout = default_timeout_sec
        # key: (to_wxid, content_normalized) -> list[(future, expire_at)]
        # 用 list 支持同 to_wxid + 同 content 的并发发送（极少见，但需支持）
        self._pending: dict[tuple[str, str], list[tuple[asyncio.Future, float]]] = {}

    def _normalize(self, content: str) -> str:
        """规范化内容用于匹配：去除首尾空白。

        微信消息内容与发送内容可能有微小差异（如尾随换行），
        用 strip 后的子串匹配更鲁棒。
        """
        return (content or "").strip()

    def register(
        self,
        to_wxid: str,
        content: str,
        timeout_sec: Optional[float] = None,
    ) -> asyncio.Future:
        """注册一个发送结果等待。

        Args:
            to_wxid: 目标联系人 wxid
            content: 发送的消息内容
            timeout_sec: 超时秒数，None 用默认值

        Returns:
            asyncio.Future[bool]：True 表示收到匹配的 SSE 事件，
            False 表示超时未收到
        """
        loop = asyncio.get_running_loop()
        future: asyncio.Future = loop.create_future()
        timeout = timeout_sec if timeout_sec is not None else self._default_timeout
        expire_at = time.monotonic() + timeout

        key = (to_wxid, self._normalize(content))
        self._pending.setdefault(key, []).append((future, expire_at))

        # 超时自动清理：超时后 resolve False 并移除条目
        def _on_timeout() -> None:
            if not future.done():
                future.set_result(False)
            self._cleanup_done(key)

        loop.call_later(timeout, _on_timeout)
        return future

    async def wait(self, future: asyncio.Future) -> bool:
        """等待 register 返回的 future 完成。

        单独提供 wait 方法便于在 async 上下文中 await。
        """
        return await future

    def notify(
        self,
        talker: str,
        content: str,
        is_sender: bool,
    ) -> int:
        """消息广播时调用，匹配 pending 等待并 resolve。

        一条 SSE 消息只匹配一个等待（FIFO），避免同 (to_wxid, content) 的
        并发发送被同一条消息全部 resolve。后续相同等待需等下一条 SSE 消息。

        Args:
            talker: 消息的 talker 字段（会话对方 wxid）
            content: 消息内容
            is_sender: 是否是自己发出的消息

        Returns:
            匹配到的等待数量（0 或 1）
        """
        if not is_sender:
            return 0
        normalized = self._normalize(content)
        # 精确匹配 talker，但 content 用子串匹配（微信可能加前缀）
        # 遍历所有 pending key，找 talker 匹配且 content 包含目标内容的
        matched = 0
        keys_to_check = [
            k for k in self._pending.keys()
            if k[0] == talker
        ]
        for key in keys_to_check:
            if matched > 0:
                break  # 一条消息只匹配一个等待
            _, target_content = key
            # target_content 是发送时规范化的内容
            # 检查消息内容是否包含目标内容（微信群消息可能加发送者前缀）
            if target_content and target_content in normalized:
                waiters = self._pending.get(key, [])
                if not waiters:
                    self._pending.pop(key, None)
                    continue
                # FIFO：只 resolve 第一个未完成的等待
                future, expire_at = waiters[0]
                if not future.done():
                    future.set_result(True)
                    matched = 1
                # 移除已 resolve 的条目
                remaining = waiters[1:]
                if remaining:
                    self._pending[key] = remaining
                else:
                    self._pending.pop(key, None)
        return matched

    def _cleanup_done(self, key: tuple[str, str]) -> None:
        """清理已完成的等待条目（超时或已 resolve）。"""
        waiters = self._pending.get(key)
        if waiters is None:
            return
        remaining = [(f, t) for f, t in waiters if not f.done()]
        if remaining:
            self._pending[key] = remaining
        else:
            self._pending.pop(key, None)

    def pending_count(self) -> int:
        """当前等待中的校验数量（供监控/调试）。"""
        return sum(len(v) for v in self._pending.values())

    def clear(self) -> None:
        """清除所有等待（resolve False）。供 lifespan 停止时调用。"""
        for key, waiters in list(self._pending.items()):
            for future, _ in waiters:
                if not future.done():
                    future.set_result(False)
        self._pending.clear()