企业微信 api_server 并发优化方案（方案二）

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
企业微信桥接服务 - WeCom Bridge Service
========================================
负责：Webhook 接收、消息解析、调用 api_server、回复推送

Author: 爱优技术团队
Version: 1.0.0
"""

import asyncio
import hashlib
import hmac
import json
import logging
import os
import time
from dataclasses import dataclass
from typing import Any, Dict, Optional

import aiohttp
import redis.asyncio as redis
import yaml
from aiohttp import web

# ============================================================================
# 配置加载
# ============================================================================
CONFIG_PATH = os.environ.get("WECOM_BRIDGE_CONFIG", "/etc/wecom_bridge/config.yaml")


@dataclass
class WeComConfig:
    corp_id: str
    corp_secret: str
    agent_id: str
    token: str  # 回调验证 Token
    encoding_aes_key: str  # AES 加密 Key（可选，本方案不启用加密模式）
    api_base_url: str = "https://qyapi.weixin.qq.com"


@dataclass
class BridgeConfig:
    wecom: WeComConfig
    api_server_url: str  # 例如：http://127.0.0.1:8000
    api_key: str  # 访问 api_server 的 Bearer Token
    redis_url: str
    redis_session_prefix: str = "wecom:session:"
    redis_session_ttl: int = 3600
    max_concurrent_requests: int = 10
    request_timeout: int = 30
    reply_queue_maxsize: int = 1000


def load_config(path: str) -> BridgeConfig:
    with open(path) as f:
        raw = yaml.safe_load(f)

    wecom_cfg = WeComConfig(
        corp_id=raw["wecom"]["corp_id"],
        corp_secret=raw["wecom"]["corp_secret"],
        agent_id=raw["wecom"]["agent_id"],
        token=raw["wecom"]["token"],
        encoding_aes_key=raw["wecom"].get("encoding_aes_key", ""),
    )
    return BridgeConfig(
        wecom=wecom_cfg,
        api_server_url=raw["api_server"]["url"],
        api_key=raw["api_server"]["api_key"],
        redis_url=raw["redis"]["url"],
    )


# ============================================================================
# 日志配置
# ============================================================================
def setup_logging(debug: bool = False):
    logging.basicConfig(
        level=logging.DEBUG if debug else logging.INFO,
        format="[%(asctime)s] %(levelname)-8s %(name)s | %(message)s",
        datefmt="%Y-%m-%d %H:%M:%S",
    )


# ============================================================================
# 企业微信签名验证
# ============================================================================
class WeComSignatureVerifier:
    """企业微信回调签名验证器"""

    def __init__(self, token: str):
        self.token = token

    def verify(self, msg_signature: str, timestamp: str, nonce: str, echostr: str = None, post_data: str = None) -> bool:
        """验证签名是否合法"""
        # 将 token、timestamp、nonce、echostr（或 post_data）按字典序排序
        sort_list = sorted([self.token, timestamp, nonce, echostr or post_data or ""])
        sign_str = "".join(sort_list)

        # SHA1 哈希
        expected = hashlib.sha1(sign_str.encode()).hexdigest()

        # 使用 hmac.compare_digest 防止时序攻击
        return hmac.compare_digest(expected, msg_signature)


# ============================================================================
# Session 管理（基于 Redis）
# ============================================================================
class SessionManager:
    """多租户 Session 管理器"""

    def __init__(self, redis_url: str, prefix: str, ttl: int):
        self.redis_url = redis_url
        self.prefix = prefix
        self.ttl = ttl
        self._redis: Optional[redis.Redis] = None

    async def initialize(self):
        self._redis = redis.from_url(
            self.redis_url,
            encoding="utf-8",
            decode_responses=True,
            max_connections=20,
        )

    async def close(self):
        if self._redis:
            await self._redis.close()

    def _session_key(self, tenant_id: str, user_id: str) -> str:
        return f"{self.prefix}{tenant_id}:{user_id}"

    async def get_history(self, tenant_id: str, user_id: str) -> list:
        """获取历史消息"""
        key = self._session_key(tenant_id, user_id)
        data = await self._redis.get(key)
        if data:
            return json.loads(data)
        return []

    async def append_message(self, tenant_id: str, user_id: str, role: str, content: str):
        """追加消息到历史"""
        key = self._session_key(tenant_id, user_id)
        history = await self.get_history(tenant_id, user_id)
        history.append({"role": role, "content": content})

        # 保持最近 N 条消息（防止无限膨胀）
        MAX_HISTORY = 50
        if len(history) > MAX_HISTORY:
            history = history[-MAX_HISTORY:]

        await self._redis.setex(key, self.ttl, json.dumps(history, ensure_ascii=False))

    async def get_tenant_user_from_msg(self, msg: Dict) -> tuple:
        """从消息体提取 tenant_id 和 user_id

        企业微信回调消息中：
        - FromUserName 是用户的 user_id（UserID）
        - ToUserName 是企业的 corp_id 部分
        - AgentID 是应用的 agent_id

        这里简化处理：tenant_id = corp_id，user_id = FromUserName
        """
        corp_id = msg.get("ToUserName", "")
        user_id = msg.get("FromUserName", "")
        return corp_id, user_id


# ============================================================================
# API Server 客户端
# ============================================================================
class ApiServerClient:
    """调用 Hermes api_server 的客户端"""

    def __init__(self, base_url: str, api_key: str, timeout: int):
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self.timeout = aiohttp.ClientTimeout(total=timeout)
        self._session: Optional[aiohttp.ClientSession] = None

    async def _get_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession(timeout=self.timeout)
        return self._session

    async def chat_completions(
        self,
        tenant_id: str,
        user_id: str,
        messages: list,
        model: str = "MiniMax-M2.7-highspeed",
    ) -> Dict[str, Any]:
        """调用 /v1/chat/completions"""
        session = await self._get_session()

        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}",
            "X-Tenant-ID": tenant_id,
            "X-User-ID": user_id,
        }

        payload = {
            "model": model,
            "messages": messages,
            "stream": False,
        }

        async with session.post(
            f"{self.base_url}/v1/chat/completions",
            headers=headers,
            json=payload,
        ) as resp:
            if resp.status != 200:
                error_text = await resp.text()
                logging.error(f"api_server error: {resp.status} {error_text}")
                return {"error": f"api_server returned {resp.status}"}

            return await resp.json()

    async def close(self):
        if self._session and not self._session.closed:
            await self._session.close()


# ============================================================================
# 企业微信消息推送
# ============================================================================
class WeComReplyClient:
    """推送回复消息到企业微信"""

    def __init__(self, config: WeComConfig, redis_url: str):
        self.config = config
        self._redis_url = redis_url
        self._access_token: Optional[str] = None
        self._token_expires_at: float = 0
        self._session: Optional[aiohttp.ClientSession] = None

    async def _get_session(self) -> aiohttp.ClientSession:
        if self._session is None or self._session.closed:
            self._session = aiohttp.ClientSession()
        return self._session

    async def _ensure_token(self):
        """获取或刷新 access_token"""
        if time.time() < self._token_expires_at - 60:
            return

        session = await self._get_session()
        url = f"{self.config.api_base_url}/cgi-bin/gettoken"
        params = {"corpid": self.config.corp_id, "corpsecret": self.config.corp_secret}

        async with session.get(url, params=params) as resp:
            data = await resp.json()
            if data.get("errcode", 0) != 0:
                raise RuntimeError(f"Failed to get token: {data}")
            self._access_token = data["access_token"]
            self._token_expires_at = time.time() + data["expires_in"]

    async def send_text(self, to_user: str, content: str, agent_id: str = None):
        """发送文本消息"""
        await self._ensure_token()
        session = await self._get_session()

        url = f"{self.config.api_base_url}/cgi-bin/message/send"
        params = {"access_token": self._access_token}

        payload = {
            "touser": to_user,
            "msgtype": "text",
            "agentid": agent_id or self.config.agent_id,
            "text": {"content": content},
        }

        async with session.post(url, params=params, json=payload) as resp:
            data = await resp.json()
            if data.get("errcode", 0) != 0:
                logging.error(f"Failed to send message: {data}")
                return False
            return True


# ============================================================================
# Webhook 处理（核心路由）
# ============================================================================
class WeComWebhookHandler:
    """企业微信回调处理器"""

    def __init__(
        self,
        config: BridgeConfig,
        session_mgr: SessionManager,
        api_client: ApiServerClient,
        reply_client: WeComReplyClient,
        verifier: WeComSignatureVerifier,
    ):
        self.config = config
        self.session_mgr = session_mgr
        self.api_client = api_client
        self.reply_client = reply_client
        self.verifier = verifier
        self._queue: asyncio.Queue = asyncio.Queue(maxsize=config.reply_queue_maxsize)
        self._worker_task: Optional[asyncio.Task] = None

    async def start(self):
        """启动异步回复 worker"""
        self._worker_task = asyncio.create_task(self._reply_worker())
        logging.info("回复 worker 已启动")

    async def stop(self):
        """停止 worker"""
        if self._worker_task:
            self._worker_task.cancel()
            try:
                await self._worker_task
            except asyncio.CancelledError:
                pass

    async def _reply_worker(self):
        """异步回复队列 worker"""
        while True:
            try:
                item = await self._queue.get()
                await self.reply_client.send_text(
                    to_user=item["to_user"],
                    content=item["content"],
                    agent_id=item.get("agent_id"),
                )
                self._queue.task_done()
            except asyncio.CancelledError:
                break
            except Exception as e:
                logging.exception(f"回复 worker 异常: {e}")

    async def handle(self, request: web.Request) -> web.Response:
        """处理企业微信回调"""
        # ================================================================
        # Step 1: 解析消息体
        # ================================================================
        try:
            msg = await request.json()
        except Exception:
            return web.Response(text="", status=200)

        msg_type = msg.get("MsgType", "")
        event = msg.get("Event", "")

        logging.info(f"收到消息: type={msg_type} event={event} from={msg.get('FromUserName')}")

        # ================================================================
        # Step 2: URL 验证（首次配置回调地址时，企业微信会发送 GET 请求）
        # ================================================================
        if request.method == "GET":
            echo_str = msg.get("echostr", "")
            msg_signature = request.query.get("msg_signature", "")
            timestamp = request.query.get("timestamp", "")
            nonce = request.query.get("nonce", "")

            if self.verifier.verify(msg_signature, timestamp, nonce, echostr=echo_str):
                logging.info("URL 验证通过")
                return web.Response(text=echo_str)
            else:
                logging.warning("URL 验证失败")
                return web.Response(status=403, text="signature verification failed")

        # ================================================================
        # Step 3: 消息签名验证
        # ================================================================
        msg_signature = request.query.get("msg_signature", "")
        timestamp = request.query.get("timestamp", "")
        nonce = request.query.get("nonce", "")
        post_data = await request.text()

        if not self.verifier.verify(msg_signature, timestamp, nonce, post_data=post_data):
            logging.warning("消息签名验证失败")
            return web.Response(status=403, text="signature verification failed")

        # ================================================================
        # Step 4: 处理不同类型的消息
        # ================================================================
        if msg_type == "text":
            return await self._handle_text(msg)
        elif msg_type == "event":
            return await self._handle_event(msg)
        else:
            # 其他类型消息暂不处理，直接返回 success
            return web.Response(text="success")

    async def _handle_text(self, msg: Dict) -> web.Response:
        """处理文本消息"""
        from_user = msg.get("FromUserName", "")
        content = msg.get("Content", "").strip()
        agent_id = msg.get("AgentID", "")

        # 企业微信 corp_id 在 ToUserName 中
        corp_id = msg.get("ToUserName", "")

        # 获取历史会话
        history = await self.session_mgr.get_history(corp_id, from_user)

        # 构建消息列表（包含历史）
        messages = history + [{"role": "user", "content": content}]

        # 立即返回 "success"（企业微信要求 5 秒内响应）
        asyncio.create_task(self._process_and_reply(corp_id, from_user, messages, agent_id))

        return web.Response(text="success")

    async def _handle_event(self, msg: Dict) -> web.Response:
        """处理事件消息"""
        event = msg.get("Event", "")
        from_user = msg.get("FromUserName", "")

        # 关注/取消关注事件
        if event == "subscribe":
            reply_text = "🎉 感谢关注！有什么可以帮助您的吗？"
            await self._queue.put({
                "to_user": from_user,
                "content": reply_text,
                "agent_id": msg.get("AgentID", ""),
            })
        elif event == "unsubscribe":
            logging.info(f"用户 {from_user} 取消关注")
        elif event == "click":
            # 点击菜单事件
            event_key = msg.get("EventKey", "")
            reply_text = f"您点击了菜单：{event_key}"
            await self._queue.put({
                "to_user": from_user,
                "content": reply_text,
                "agent_id": msg.get("AgentID", ""),
            })

        return web.Response(text="success")

    async def _process_and_reply(
        self,
        corp_id: str,
        from_user: str,
        messages: list,
        agent_id: str,
    ):
        """异步处理消息并回复（不阻塞回调）"""
        try:
            # 保存用户消息到历史
            user_content = messages[-1]["content"]
            await self.session_mgr.append_message(corp_id, from_user, "user", user_content)

            # 追加历史消息（去掉最后一条用户消息，因为要在 api_server 中处理）
            history_for_api = messages[:-1]

            # 构建发送给 api_server 的消息（含历史）
            api_messages = history_for_api + [{"role": "user", "content": user_content}]

            # 调用 api_server
            result = await self.api_client.chat_completions(
                tenant_id=corp_id,
                user_id=from_user,
                messages=api_messages,
            )

            if "error" in result:
                reply_text = f"⚠️ 处理失败：{result['error']}"
            else:
                # 解析回复（OpenAI 兼容格式）
                choices = result.get("choices", [])
                if choices:
                    reply_text = choices[0].get("message", {}).get("content", "收到")
                    reply_text = reply_text.strip()
                else:
                    reply_text = "收到"

            # 保存助手回复到历史
            await self.session_mgr.append_message(corp_id, from_user, "assistant", reply_text)

            # 放入回复队列
            await self._queue.put({
                "to_user": from_user,
                "content": reply_text,
                "agent_id": agent_id,
            })

            logging.info(f"回复已加入队列: to={from_user} len={len(reply_text)}")

        except asyncio.TimeoutError:
            logging.error("api_server 调用超时")
            await self._queue.put({
                "to_user": from_user,
                "content": "⏱️ 处理超时，请稍后再试",
                "agent_id": agent_id,
            })
        except Exception as e:
            logging.exception(f"处理消息异常: {e}")
            await self._queue.put({
                "to_user": from_user,
                "content": f"❌ 处理异常：{str(e)[:50]}",
                "agent_id": agent_id,
            })


# ============================================================================
# 健康检查
# ============================================================================
async def health_check(request: web.Request) -> web.Response:
    return web.json_response({"status": "healthy", "service": "wecom_bridge"})


# ============================================================================
# 主程序
# ============================================================================
async def create_app() -> web.Application:
    config = load_config(CONFIG_PATH)
    setup_logging(debug=os.environ.get("DEBUG"))

    logging.info("=" * 60)
    logging.info("  企业微信桥接服务 (WeCom Bridge Service)")
    logging.info(f"  api_server: {config.api_server_url}")
    logging.info(f"  redis: {config.redis_url}")
    logging.info("=" * 60)

    # 初始化组件
    session_mgr = SessionManager(
        redis_url=config.redis_url,
        prefix=config.redis_session_prefix,
        ttl=config.redis_session_ttl,
    )
    await session_mgr.initialize()

    api_client = ApiServerClient(
        base_url=config.api_server_url,
        api_key=config.api_key,
        timeout=config.request_timeout,
    )

    reply_client = WeComReplyClient(config.wecom, config.redis_url)

    verifier = WeComSignatureVerifier(config.wecom.token)

    handler = WeComWebhookHandler(config, session_mgr, api_client, reply_client, verifier)
    await handler.start()

    # 构建 app
    app = web.Application()
    app["handler"] = handler
    app["session_mgr"] = session_mgr
    app["api_client"] = api_client

    app.router.add_get("/wecom/callback", handler.handle)       # URL 验证
    app.router.add_post("/wecom/callback", handler.handle)      # 消息回调
    app.router.add_get("/health", health_check)

    # 清理函数
    async def cleanup(app):
        logging.info("关闭桥接服务...")
        await handler.stop()
        await session_mgr.close()
        await api_client.close()

    app.on_cleanup.append(cleanup)

    return app


def main():
    app = create_app()
    port = int(os.environ.get("PORT", 8080))
    logging.info(f"启动服务，监听端口 {port}")
    web.run_app(app, host="0.0.0.0", port=port, access_log=None)


if __name__ == "__main__":
    main()

# 请求头注入
headers = {
    "Authorization": f"Bearer {config.api_key}",
    "X-Tenant-ID": corp_id,      # 多租户隔离
    "X-User-ID": from_user,      # 用户隔离
}

# 禁止打印的字段
SENSITIVE_KEYS = ["corp_secret", "api_key", "token", "password"]

def mask_sensitive(data: dict) -> dict:
    result = {}
    for k, v in data.items():
        if k in SENSITIVE_KEYS:
            result[k] = "***REDACTED***"
        elif isinstance(v, str) and len(v) > 100:
            result[k] = v[:50] + "..."
        else:
            result[k] = v
    return result