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`： 第一个抢到、第二个抢不到、第一个放手后第三个能接上。
247 lines
8.8 KiB
Python
247 lines
8.8 KiB
Python
"""外部接入：微信 clawbot 的扫码登录与状态。
|
|
|
|
真实协议是腾讯 iLink Bot（出处见 `app/services/wechat_ilink.py` 文件头）：
|
|
**我们**去取二维码、等人扫、拿到登录态后长轮询拉消息。所以这里的接口是三件事：
|
|
|
|
1. `POST /workspaces/{ws}/integrations/wechat/qr` —— 取一张二维码，页面显示给人扫；
|
|
2. `GET /workspaces/{ws}/integrations/wechat/qr/{id}` —— 问"扫了吗"，扫成功就存登录态并开始轮询；
|
|
3. `DELETE /workspaces/{ws}/integrations/{channel_id}` —— 断开：登录态作废、轮询停止。
|
|
|
|
二维码 id 放 Redis 而不是内存：问状态的那一下可能落到另一个进程，
|
|
存内存就会"扫完了但没人知道"。key 带工作区，两个工作区互不踩。
|
|
|
|
上一版这里做的是 webhook 接收端（`POST /integrations/clawbot/{token}`），
|
|
对 clawbot 的理解是错的，接口与令牌一并删掉了（迁移 0013 的注释记了原因）。
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import uuid
|
|
from typing import Any
|
|
|
|
import qrcode
|
|
import qrcode.image.svg
|
|
from fastapi import APIRouter, Query
|
|
from pydantic import BaseModel
|
|
from sqlalchemy import select
|
|
|
|
from ..db import bind_tenant
|
|
from ..deps import AdminDep, SessionDep
|
|
from ..errors import AppError, ErrorCode
|
|
from ..models_integrations import BotChannel, BotMessage
|
|
from ..schemas import ok
|
|
from ..services import wechat_ilink
|
|
from ..services.wechat_bot import manager
|
|
from ..stores import get_store
|
|
|
|
admin_router = APIRouter(prefix="/workspaces", tags=["外部接入"])
|
|
|
|
# iLink 的二维码大约 3 分钟过期，我们按 5 分钟缓存 id，过期就重新取。
|
|
QR_TTL_SECONDS = 300
|
|
|
|
|
|
class ChannelPatch(BaseModel):
|
|
status: str
|
|
|
|
|
|
def _qr_key(ws_id: uuid.UUID) -> str:
|
|
return f"wechat:qr:{ws_id}"
|
|
|
|
|
|
def render_qr_svg(payload: str) -> str:
|
|
"""把二维码内容画成 data URL：服务端画，前端不用再引一个二维码库。"""
|
|
image = qrcode.make(payload, image_factory=qrcode.image.svg.SvgPathImage, box_size=10)
|
|
buffer = io.BytesIO()
|
|
image.save(buffer)
|
|
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
return f"data:image/svg+xml;base64,{encoded}"
|
|
|
|
|
|
def _channel_out(channel: BotChannel) -> dict[str, Any]:
|
|
return {
|
|
"id": str(channel.id),
|
|
"kind": channel.kind,
|
|
"name": channel.name,
|
|
"status": channel.status,
|
|
"connected": bool(channel.bot_token_enc) and channel.status == "active",
|
|
# 只有抢到租约、真在拉消息的进程才算"接收中"；备用进程的任务在等键，不算。
|
|
"polling": manager.polling(channel.id),
|
|
"ilink_user_id": channel.ilink_user_id,
|
|
"last_seen_at": channel.last_seen_at,
|
|
"last_error": channel.last_error,
|
|
"created_at": channel.created_at,
|
|
}
|
|
|
|
|
|
async def _wechat_channel(session: SessionDep, ws_id: uuid.UUID) -> BotChannel | None:
|
|
rows = await session.execute(
|
|
select(BotChannel)
|
|
.where(BotChannel.workspace_id == ws_id, BotChannel.kind == "wechat_clawbot")
|
|
.order_by(BotChannel.created_at.desc())
|
|
)
|
|
return rows.scalars().first()
|
|
|
|
|
|
@admin_router.get("/{ws_id}/integrations")
|
|
async def list_integrations(ws_id: uuid.UUID, ctx: AdminDep, session: SessionDep):
|
|
rows = await session.execute(
|
|
select(BotChannel)
|
|
.where(BotChannel.workspace_id == ws_id)
|
|
.order_by(BotChannel.created_at.desc())
|
|
)
|
|
return ok({"items": [_channel_out(c) for c in rows.scalars().all()]})
|
|
|
|
|
|
@admin_router.post("/{ws_id}/integrations/wechat/qr")
|
|
async def begin_wechat_login(ws_id: uuid.UUID, ctx: AdminDep, session: SessionDep):
|
|
"""取一张二维码。已经登录过的通道会被重置成 pending（换号 = 重新扫）。"""
|
|
try:
|
|
code = await wechat_ilink.get_qrcode()
|
|
except wechat_ilink.WechatError as exc:
|
|
raise AppError(ErrorCode.BOT_TOKEN_INVALID, f"取二维码失败：{exc}") from exc
|
|
|
|
channel = await _wechat_channel(session, ws_id)
|
|
if channel is None:
|
|
channel = BotChannel(
|
|
workspace_id=ws_id,
|
|
kind="wechat_clawbot",
|
|
name="微信",
|
|
status="pending",
|
|
created_by=ctx.user.id,
|
|
)
|
|
session.add(channel)
|
|
if channel.bot_token_enc:
|
|
await manager.stop(channel.id)
|
|
channel.bot_token_enc = None
|
|
channel.sync_buf = None
|
|
channel.context_token = None
|
|
channel.status = "pending"
|
|
channel.last_error = None
|
|
await session.flush()
|
|
|
|
await get_store().set(_qr_key(ws_id), code.qrcode_id, QR_TTL_SECONDS)
|
|
await session.commit()
|
|
|
|
return ok(
|
|
{
|
|
"channel_id": str(channel.id),
|
|
"qrcode_id": code.qrcode_id,
|
|
"payload": code.payload,
|
|
"image": render_qr_svg(code.payload),
|
|
"expires_in": QR_TTL_SECONDS,
|
|
}
|
|
)
|
|
|
|
|
|
@admin_router.get("/{ws_id}/integrations/wechat/qr/{qrcode_id}")
|
|
async def poll_wechat_login(ws_id: uuid.UUID, qrcode_id: str, ctx: AdminDep, session: SessionDep):
|
|
"""问扫码结果。只有 confirmed 才落登录态并开轮询，其余状态原样回给页面。"""
|
|
cached = await get_store().get(_qr_key(ws_id))
|
|
if cached != qrcode_id:
|
|
# 页面拿着旧二维码在问：直接说"过期"，让它重新取一张。
|
|
return ok({"status": "expired"})
|
|
|
|
try:
|
|
status = await wechat_ilink.qrcode_status(qrcode_id)
|
|
except wechat_ilink.WechatError as exc:
|
|
return ok({"status": "error", "message": str(exc)})
|
|
|
|
if status.status != "confirmed" or not status.bot_token:
|
|
return ok({"status": status.status})
|
|
|
|
channel = await _wechat_channel(session, ws_id)
|
|
if channel is None:
|
|
raise AppError(ErrorCode.BOT_TOKEN_INVALID, "通道不见了，请重新取二维码")
|
|
|
|
await bind_tenant(session, ws_id, ctx.user.id)
|
|
channel.bot_token_enc = wechat_ilink.encrypt_token(status.bot_token)
|
|
channel.ilink_user_id = status.ilink_user_id
|
|
channel.base_url = status.base_url
|
|
channel.status = "active"
|
|
channel.sync_buf = None
|
|
channel.context_token = None
|
|
channel.last_error = None
|
|
channel.created_by = channel.created_by or ctx.user.id
|
|
await session.commit()
|
|
await get_store().delete(_qr_key(ws_id))
|
|
|
|
await manager.start(channel.id)
|
|
return ok({"status": "confirmed", "channel": _channel_out(channel)})
|
|
|
|
|
|
@admin_router.patch("/{ws_id}/integrations/{channel_id}")
|
|
async def update_integration(
|
|
ws_id: uuid.UUID,
|
|
channel_id: uuid.UUID,
|
|
ctx: AdminDep,
|
|
session: SessionDep,
|
|
body: ChannelPatch,
|
|
):
|
|
channel = await session.get(BotChannel, channel_id)
|
|
if channel is None or channel.workspace_id != ws_id:
|
|
raise AppError(ErrorCode.BOT_TOKEN_INVALID, "接入通道不存在")
|
|
wanted = body.status.strip()
|
|
if wanted not in ("active", "disabled"):
|
|
raise AppError(ErrorCode.VALIDATION_FAILED, "status 只能是 active 或 disabled")
|
|
if wanted == "active" and not channel.bot_token_enc:
|
|
raise AppError(ErrorCode.BOT_TOKEN_INVALID, "这个通道还没扫码登录")
|
|
channel.status = wanted
|
|
await session.commit()
|
|
if wanted == "active":
|
|
await manager.start(channel.id)
|
|
else:
|
|
await manager.stop(channel.id)
|
|
return ok({"channel": _channel_out(channel)})
|
|
|
|
|
|
@admin_router.delete("/{ws_id}/integrations/{channel_id}")
|
|
async def disconnect_integration(
|
|
ws_id: uuid.UUID, channel_id: uuid.UUID, ctx: AdminDep, session: SessionDep
|
|
):
|
|
"""断开：停轮询、清登录态。留痕（bot_message）不删，翻旧账用得上。"""
|
|
channel = await session.get(BotChannel, channel_id)
|
|
if channel is None or channel.workspace_id != ws_id:
|
|
raise AppError(ErrorCode.BOT_TOKEN_INVALID, "接入通道不存在")
|
|
await manager.stop(channel.id)
|
|
channel.bot_token_enc = None
|
|
channel.sync_buf = None
|
|
channel.context_token = None
|
|
channel.status = "pending"
|
|
channel.last_error = None
|
|
await session.commit()
|
|
return ok({"channel": _channel_out(channel)})
|
|
|
|
|
|
@admin_router.get("/{ws_id}/integrations/{channel_id}/messages")
|
|
async def list_messages(
|
|
ws_id: uuid.UUID,
|
|
channel_id: uuid.UUID,
|
|
ctx: AdminDep,
|
|
session: SessionDep,
|
|
limit: int = Query(default=5, ge=1, le=100),
|
|
):
|
|
rows = await session.execute(
|
|
select(BotMessage)
|
|
.where(BotMessage.channel_id == channel_id, BotMessage.workspace_id == ws_id)
|
|
.order_by(BotMessage.created_at.desc())
|
|
.limit(limit)
|
|
)
|
|
return ok(
|
|
{
|
|
"items": [
|
|
{
|
|
"id": str(m.id),
|
|
"action": m.action,
|
|
"sender": m.sender,
|
|
"text": m.text,
|
|
"reply": m.reply,
|
|
"item_id": str(m.item_id) if m.item_id else None,
|
|
"created_at": m.created_at,
|
|
}
|
|
for m in rows.scalars().all()
|
|
]
|
|
}
|
|
)
|