Steven
07b01db7cc
CI / backend (push) Canceled after 0s
生产是 `gunicorn -w 2`，两个进程各自跑 lifespan，start_all() 会把同一个通道的 长轮询拉两遍：消息被其中一个消费、回复从另一个发出，游标还会互相覆盖。 改成 Postgres 会话级 advisory lock 做租约：抢到的进程收消息，另一个每 10 秒重试， 主进程一挂连接断、锁释放，备用进程立刻接上。页面上的"接收中"只认抢到租约的那个 （`manager.polling()`），不再把等锁的也算进去。 补测试 `test_only_one_process_holds_the_polling_lease`： 第一个抢到、第二个抢不到、第一个放手后第三个能接上。
335 lines
14 KiB
Python
335 lines
14 KiB
Python
"""微信机器人的后台轮询：登录之后由它把消息捞进来。
|
||
|
||
为什么是长轮询而不是 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 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:
|
||
"""进程启动时把已登录的通道都拉起来，不然重启就静默不收消息了。"""
|
||
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()
|
||
|
||
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()
|