"""微信机器人的后台轮询：登录之后由它把消息捞进来。 为什么是长轮询而不是 webhook：iLink Bot 协议就是"登录后由我们主动去拉"， 服务端只在下一条消息到达时返回（最长等 75 秒）。所以每个通道一个常驻任务。 几条必须守住的规矩： 1. **游标要落库**：`get_updates_buf` 是下次从哪里接着读的凭据，丢了就会重复收、 把同一条任务记两遍。 2. **一条消息只处理一次**：`bot_message.external_id` 有唯一索引，重复投递时先撞唯一键。 3. **token 失效要停**：服务端 errcode `-14` 表示登录态作废，这时停轮询、把通道标记成 需要重新扫码，别拿死 token 一直撞墙。 4. **回复失败不影响已经落库的事项**：任务记下了就是记下了，回不回得去另说。 5. **一个通道只能有一个进程在收**：生产是 gunicorn `-w 2`，两个进程各自跑 lifespan， 不加锁就会出现两条 getupdates 同时拉同一个会话——消息被其中一个消费、回复却从另一个发出， 游标还会互相覆盖。用 Postgres 的 advisory lock 做租约：抢到的那个进程收，另一个等着接班。 """ 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, engine 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 # 抢租约失败时隔多久再试：主进程挂了，备用进程要能在一两个呼吸之内接上。 LEASE_RETRY_SECONDS = 10.0 LEASE_NAMESPACE = "dadao.wechat_bot" 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 "" async def lease_held(channel_id: uuid.UUID) -> bool: """这个通道现在有没有进程在收消息。 不能看 `manager.polling()`：生产是 gunicorn 多进程，每个进程只知道自己那把锁， 请求落到哪个进程就回哪个答案，页面会在"接收中"和"已停止"之间闪。 advisory lock 是数据库级的，`pg_locks` 里看得见——谁抢到了，所有进程都认。 """ async with SessionFactory() as session: result = await session.execute( sa_text( "SELECT 1 FROM pg_locks" " WHERE locktype = 'advisory' AND granted AND objsubid = 2" " AND classid::bigint = (hashtext(:ns)::bigint & 4294967295)" " AND objid::bigint = (hashtext(:cid)::bigint & 4294967295)" " LIMIT 1" ), {"ns": LEASE_NAMESPACE, "cid": str(channel_id)}, ) return result.first() is not None async def active_channel_ids() -> list[uuid.UUID]: """列出"已扫码登录、该收消息"的通道。 启动时还没有工作区上下文，`bot_channel` 的 RLS 会把行全藏起来，所以先设 `app.bot_list='on'` 走迁移 0016 那条只读口子——不设的话重启后机器人会静默失联。 """ async with SessionFactory() as session, session.begin(): await session.execute(sa_text("SELECT set_config('app.bot_list', 'on', true)")) rows = await session.execute( select(BotChannel.id).where( BotChannel.status == "active", BotChannel.bot_token_enc.is_not(None), ) ) return [row[0] for row in rows] async def acquire_lease(channel_id: uuid.UUID): """抢"这个通道由我收消息"的租约；抢到返回连接（关掉就是释放），抢不到返回 None。 用会话级 advisory lock：锁挂在连接上，进程崩了连接断，锁自然释放， 不会留下"死进程占着锁没人收消息"的局面。 """ conn = await engine.connect() try: result = await conn.execute( sa_text("SELECT pg_try_advisory_lock(hashtext(:ns), hashtext(:cid))"), {"ns": LEASE_NAMESPACE, "cid": str(channel_id)}, ) if result.scalar(): return conn except Exception: await conn.close() raise await conn.close() return None class BotManager: """一个进程里的轮询任务表：一个通道一个任务。""" def __init__(self) -> None: self._tasks: dict[uuid.UUID, asyncio.Task[None]] = {} # 抢到租约、真正在拉消息的通道；备用进程的任务在等租约，不在这里。 self._polling: set[uuid.UUID] = set() # ── 生命周期 ────────────────────────────────────────── 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: """进程启动时把已登录的通道都拉起来，不然重启就静默不收消息了。""" ids = await active_channel_ids() 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() def polling(self, channel_id: uuid.UUID) -> bool: """只有真在拉消息的那个进程才算，"页面上的接收中"认这个。""" return channel_id in self._polling # ── 轮询主体 ────────────────────────────────────────── 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: lease = await self._wait_for_lease(channel_id) if lease is None: self._tasks.pop(channel_id, None) return try: await self._poll_forever(channel_id) finally: self._polling.discard(channel_id) self._tasks.pop(channel_id, None) await lease.close() async def _wait_for_lease(self, channel_id: uuid.UUID): """抢不到就等着——主进程一挂，最多 LEASE_RETRY_SECONDS 就由这里接上。""" while True: lease = await acquire_lease(channel_id) if lease is not None: self._polling.add(channel_id) return lease await asyncio.sleep(LEASE_RETRY_SECONDS) async def _poll_forever(self, channel_id: uuid.UUID) -> None: loaded = await self._load(channel_id) if loaded is 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)) 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()