Steven
8436c0edbf
CI / backend (push) Canceled after 0s
上一版按"webhook + URL 里带令牌"自己编了一套，用户当场指出微信必须先扫码配对。 现在照腾讯 iLink Bot 协议自己做客户端：取二维码 → 手机扫码确认 → 存加密登录态 → 后端长轮询收消息 → 用 context_token 回话。独立实现，不依赖 DSH 的运行环境。 后端 - wechat_ilink.py：get_bot_qrcode / get_qrcode_status / getupdates / sendmessage， bot_token 用 Fernet（密钥由 JWT_SECRET 派生）加密入库，日志与接口都不回明文 - wechat_intake.py：一句话 → 任务或知识草稿，前缀识别、#标签、首行标题 - wechat_bot.py：每条通道一个轮询任务，启动时按 status='active' 恢复； message_id 去重；errcode -14 转 expired 让用户重新扫 - api/integrations.py 重写成扫码三接口 + 停用/断开/最近消息 - 迁移 0013（登录态列）、0014（轮询入口按 app.bot_channel 主键的 RLS 口子） Web - 接入页改成扫码流程：页面里画出二维码、每 2 秒问一次状态、连上显示微信用户与最近消息 验证 - pytest 149 passed（新增 test_wechat_bot 9 例） - ruff check / ruff format --check 干净；tsc 0；vite build ✓ - web/e2e/integrations.mjs 7 项 ALL PASS（拿的是微信真二维码） 教训：接入第三方协议先去已有实现里找协议，别从直觉倒推。
283 lines
11 KiB
Python
283 lines
11 KiB
Python
"""微信机器人的后台轮询：登录之后由它把消息捞进来。
|
||
|
||
为什么是长轮询而不是 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()
|