"""微信机器人的后台轮询：登录之后由它把消息捞进来。 为什么是长轮询而不是 webhook：iLink Bot 协议就是"登录后由我们主动去拉"， 服务端只在下一条消息到达时返回（最长等 75 秒）。所以每个通道一个常驻任务。 几条必须守住的规矩： 1. **游标要落库**：`get_updates_buf` 是下次从哪里接着读的凭据，丢了就会重复收、 把同一条任务记两遍。 2. **一条消息只处理一次**：`bot_message.external_id` 有唯一索引，重复投递时先撞唯一键。 3. **token 失效要停**：服务端 errcode `-14` 表示登录态作废，这时停轮询、把通道标记成 需要重新扫码，别拿死 token 一直撞墙。 4. **回复失败不影响已经落库的事项**：任务记下了就是记下了，回不回得去另说。 """ from __future__ import annotations import asyncio import contextlib import uuid from typing import Any from sqlalchemy import select from sqlalchemy import text as sa_text from ..config import get_settings from ..db import SessionFactory, bind_tenant from ..logging_setup import get_logger from ..models import Workspace from ..models_integrations import BotChannel, BotMessage from . import wechat_ilink from .wechat_intake import handle_text logger = get_logger("wechat_bot") # 长轮询一次最多等这么久（服务端 75 秒超时，客户端留点余量）。 POLL_TIMEOUT_SECONDS = 80.0 RETRY_DELAY_SECONDS = 5.0 def site_base() -> str: """回复里的链接用哪个域名。生产必须配 PUBLIC_BASE_URL，否则只能退回 CORS 里的第一个。""" settings = get_settings() if settings.public_base_url: return settings.public_base_url.rstrip("/") if settings.cors_origins: return str(settings.cors_origins[0]).rstrip("/") return "" class BotManager: """一个进程里的轮询任务表：一个通道一个任务。""" def __init__(self) -> None: self._tasks: dict[uuid.UUID, asyncio.Task[None]] = {} # ── 生命周期 ────────────────────────────────────────── async def start(self, channel_id: uuid.UUID) -> None: task = self._tasks.get(channel_id) if task is not None and not task.done(): return self._tasks[channel_id] = asyncio.create_task( self._run(channel_id), name=f"wechat-bot-{channel_id}" ) async def stop(self, channel_id: uuid.UUID) -> None: task = self._tasks.pop(channel_id, None) if task is None: return task.cancel() with contextlib.suppress(asyncio.CancelledError): await task async def start_all(self) -> None: """进程启动时把已登录的通道都拉起来，不然重启就静默不收消息了。""" async with SessionFactory() as session: rows = await session.execute( select(BotChannel.id).where( BotChannel.status == "active", BotChannel.bot_token_enc.is_not(None), ) ) ids = [row[0] for row in rows] for channel_id in ids: await self.start(channel_id) if ids: logger.info("微信机器人已恢复", extra={"extra_fields": {"channels": len(ids)}}) async def shutdown(self) -> None: for channel_id in list(self._tasks): await self.stop(channel_id) def running(self, channel_id: uuid.UUID) -> bool: task = self._tasks.get(channel_id) return task is not None and not task.done() # ── 轮询主体 ────────────────────────────────────────── async def _load(self, channel_id: uuid.UUID) -> tuple[BotChannel, str] | None: """读通道并解出 token。返回 None 表示这个通道不该再轮询了。""" async with SessionFactory() as session: channel = await self._get_channel(session, channel_id) if channel is None or channel.status != "active" or not channel.bot_token_enc: return None try: token = wechat_ilink.decrypt_token(channel.bot_token_enc) except wechat_ilink.WechatError: await self._mark_relogin(channel_id, "登录态解不开（JWT_SECRET 可能换过）") return None session.expunge(channel) return channel, token async def _run(self, channel_id: uuid.UUID) -> None: loaded = await self._load(channel_id) if loaded is None: self._tasks.pop(channel_id, None) return channel, token = loaded cursor = channel.sync_buf or "" failures = 0 while True: try: update = await wechat_ilink.get_updates( token, cursor, base_url=channel.base_url, timeout=POLL_TIMEOUT_SECONDS ) except asyncio.CancelledError: raise except wechat_ilink.WechatError as exc: if exc.stale_token: await self._mark_relogin(channel_id, str(exc)) self._tasks.pop(channel_id, None) return failures += 1 await self._note_error(channel_id, str(exc)) await asyncio.sleep(RETRY_DELAY_SECONDS if failures < 10 else 60.0) continue except Exception as exc: # 轮询任务不能死，死了就永远收不到消息 failures += 1 logger.warning("微信长轮询异常：%s", exc) await asyncio.sleep(RETRY_DELAY_SECONDS) continue failures = 0 for message in update.messages: try: await self._handle(channel_id, message, token, channel.base_url) except Exception as exc: # 单条失败不拖垮整轮 logger.warning("微信消息处理失败：%s", exc) if update.cursor and update.cursor != cursor: cursor = update.cursor await self._save_cursor(channel_id, cursor) # ── 单条消息 ────────────────────────────────────────── async def _handle( self, channel_id: uuid.UUID, message: wechat_ilink.Inbound, token: str, base_url: str | None, ) -> None: async with SessionFactory() as session: channel = await self._get_channel(session, channel_id) if channel is None or channel.status != "active": return # 读通道时只设了 app.bot_channel；workspace 表同样有 RLS， # 得先把通道所属工作区告诉数据库，才能读到那一行。 await session.execute( sa_text("SELECT set_config('app.ws_id', :ws, true)"), {"ws": str(channel.workspace_id)}, ) workspace = await session.get(Workspace, channel.workspace_id) if workspace is None: return actor = channel.created_by or workspace.owner_id await bind_tenant(session, channel.workspace_id, actor) existing = await session.execute( select(BotMessage).where( BotMessage.channel_id == channel_id, BotMessage.external_id == message.message_id, ) ) if existing.scalars().first() is not None: logger.info("微信消息重复投递，跳过：%s", message.message_id) return result: dict[str, Any] if message.text: result = await handle_text(session, channel, message.text, base_url=site_base()) elif message.kinds: result = { "action": "ignored", "item_id": None, "reply": "图片、文件、语音暂时还没搬进来，发文字最快。", } else: result = {"action": "ignored", "item_id": None, "reply": "没看懂这条消息。"} session.add( BotMessage( workspace_id=channel.workspace_id, channel_id=channel_id, external_id=message.message_id, sender=message.sender[:128], text=message.text[:4000], action=str(result["action"]), item_id=result.get("item_id"), reply=str(result.get("reply") or "")[:2000], raw={"kinds": message.kinds, "message_id": message.message_id}, ) ) channel.last_seen_at = _now() channel.context_token = message.context_token channel.last_error = None await session.commit() reply = str(result.get("reply") or "") if reply: try: await wechat_ilink.send_text( token, message.sender, message.context_token, reply, base_url=base_url, ) except wechat_ilink.WechatError as exc: await self._note_error(channel_id, f"回复失败：{exc}") async def _get_channel(self, session, channel_id: uuid.UUID) -> BotChannel | None: """按主键取通道。 轮询任务一开始没有工作区上下文，RLS 会把行藏起来；迁移 0014 给了一条 `app.bot_channel` 的按主键口子，所以先设 GUC 再读。 """ await session.execute( sa_text("SELECT set_config('app.bot_channel', :cid, true)"), {"cid": str(channel_id)} ) return await session.get(BotChannel, channel_id) # ── 落库小工具 ──────────────────────────────────────── async def _save_cursor(self, channel_id: uuid.UUID, cursor: str) -> None: async with SessionFactory() as session: channel = await self._get_channel(session, channel_id) if channel is None: return channel.sync_buf = cursor[:4000] await session.commit() async def _note_error(self, channel_id: uuid.UUID, message: str) -> None: async with SessionFactory() as session: channel = await self._get_channel(session, channel_id) if channel is None: return channel.last_error = message[:500] await session.commit() async def _mark_relogin(self, channel_id: uuid.UUID, message: str) -> None: """登录态作废：停用这条通道并记下原因，界面上提示重新扫码。""" async with SessionFactory() as session: channel = await self._get_channel(session, channel_id) if channel is None: return channel.bot_token_enc = None channel.sync_buf = None channel.status = "expired" channel.context_token = None channel.last_error = message[:500] await session.commit() def _now(): from datetime import datetime, timezone return datetime.now(timezone.utc) manager = BotManager()