Build lightweight AI agent admin
This commit is contained in:
Codex
@@ -0,0 +1,260 @@
|
||||
#!/usr/bin/env python
|
||||
# -*- coding: utf-8 -*-
|
||||
"""
|
||||
钉钉 Stream 模式事件监听客户端
|
||||
|
||||
取代 HTTP 回调模式，通过 WebSocket 长连接接收钉钉通讯录变更事件。
|
||||
无需公网回调地址，无需配置 callback_token/callback_aes_key/corp_id。
|
||||
|
||||
使用方式：
|
||||
1. 在钉钉开放平台应用配置中，开启 Stream 模式推送
|
||||
2. 后端启动后自动建立长连接，接收事件推送
|
||||
|
||||
Stream 模式事件数据格式与 HTTP 推送不同：
|
||||
HTTP: { "EventType": "...", "UserId": [...], "DeptId": [...] }
|
||||
Stream: event.headers.event_type 为事件类型，
|
||||
event.data 为事件体 { "timeStamp": "...", "deptId": [...] / "userId": [...] }
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
from collections import OrderedDict
|
||||
from datetime import datetime
|
||||
from typing import Optional
|
||||
|
||||
import dingtalk_stream
|
||||
from dingtalk_stream import AckMessage
|
||||
from dingtalk_stream import EventMessage
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_DEDUP_MAX = 256
|
||||
|
||||
|
||||
class _EventDedup:
|
||||
"""基于 OrderedDict 的有界去重缓存，保留最近 N 个 event_id"""
|
||||
|
||||
def __init__(self, maxlen: int = _DEDUP_MAX):
|
||||
self._seen: OrderedDict[str, None] = OrderedDict()
|
||||
self._maxlen = maxlen
|
||||
|
||||
def is_duplicate(self, event_id: str) -> bool:
|
||||
if event_id in self._seen:
|
||||
self._seen.move_to_end(event_id)
|
||||
return True
|
||||
self._seen[event_id] = None
|
||||
if len(self._seen) > self._maxlen:
|
||||
self._seen.popitem(last=False)
|
||||
return False
|
||||
|
||||
|
||||
class DingtalkStreamEventHandler(dingtalk_stream.EventHandler):
|
||||
"""钉钉 Stream 事件处理器 - 处理通讯录变更事件"""
|
||||
|
||||
_dedup = _EventDedup()
|
||||
|
||||
async def process(self, event: EventMessage) -> tuple:
|
||||
try:
|
||||
event_type = event.headers.event_type if event.headers else ''
|
||||
event_data = event.data if isinstance(event.data, dict) else {}
|
||||
|
||||
event_id = getattr(event.headers, 'event_id', None) or ''
|
||||
if event_id and self._dedup.is_duplicate(event_id):
|
||||
logger.info("跳过重复事件: event_id=%s, type=%s", event_id, event_type)
|
||||
return AckMessage.STATUS_OK, 'OK'
|
||||
|
||||
DingtalkStreamManager._record_event(event_type)
|
||||
|
||||
if not event_type:
|
||||
logger.warning(f"收到未知事件: {event_data}")
|
||||
return AckMessage.STATUS_OK, 'OK'
|
||||
|
||||
event_body = event_data
|
||||
|
||||
logger.info(f"收到钉钉 Stream 事件: {event_type}")
|
||||
|
||||
from core.dingtalk_sync.callback_handler import DingtalkCallbackHandler
|
||||
from core.dingtalk_sync.callback_handler import DEPT_EVENTS, USER_EVENTS
|
||||
|
||||
from app.config_manager import config_manager
|
||||
config = await config_manager.get_group("sync_dingtalk")
|
||||
enable_dept = config.get("enable_dept_event") == "true"
|
||||
enable_user = config.get("enable_user_event") == "true"
|
||||
|
||||
if event_type in DEPT_EVENTS and enable_dept:
|
||||
normalized = self._normalize_dept_event(event_body)
|
||||
await DingtalkCallbackHandler.handle_event(event_type, normalized)
|
||||
|
||||
elif event_type in USER_EVENTS and enable_user:
|
||||
normalized = self._normalize_user_event(event_body)
|
||||
await DingtalkCallbackHandler.handle_event(event_type, normalized)
|
||||
|
||||
else:
|
||||
logger.info(f"忽略未启用的事件: {event_type}")
|
||||
|
||||
except Exception as e:
|
||||
logger.error(f"处理 Stream 事件失败: {e}", exc_info=True)
|
||||
|
||||
return AckMessage.STATUS_OK, 'OK'
|
||||
|
||||
@staticmethod
|
||||
def _normalize_dept_event(event_body: dict) -> dict:
|
||||
"""将 Stream 模式的部门事件字段转为 HTTP 回调格式"""
|
||||
dept_ids = event_body.get('deptId') or event_body.get('DeptId') or []
|
||||
return {"DeptId": dept_ids}
|
||||
|
||||
@staticmethod
|
||||
def _normalize_user_event(event_body: dict) -> dict:
|
||||
"""将 Stream 模式的用户事件字段转为 HTTP 回调格式"""
|
||||
user_ids = event_body.get('userId') or event_body.get('UserId') or []
|
||||
return {"UserId": user_ids}
|
||||
|
||||
|
||||
class DingtalkStreamManager:
|
||||
"""钉钉 Stream 模式连接管理器"""
|
||||
|
||||
_client: Optional[dingtalk_stream.DingTalkStreamClient] = None
|
||||
_task: Optional[asyncio.Task] = None
|
||||
_connected: bool = False
|
||||
_started_at: Optional[str] = None
|
||||
|
||||
# 事件统计
|
||||
_event_stats = {
|
||||
"total_events": 0,
|
||||
"last_event_type": None,
|
||||
"last_event_time": None,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def _record_event(cls, event_type: str):
|
||||
cls._event_stats["total_events"] += 1
|
||||
cls._event_stats["last_event_type"] = event_type
|
||||
cls._event_stats["last_event_time"] = datetime.now().isoformat()
|
||||
|
||||
@classmethod
|
||||
async def start(cls):
|
||||
"""启动 Stream 连接，使用同步配置中的 app_key/app_secret"""
|
||||
from app.config_manager import config_manager
|
||||
|
||||
config = await config_manager.get_group("sync_dingtalk")
|
||||
app_key = config.get("app_key")
|
||||
app_secret = config.get("app_secret")
|
||||
|
||||
if not app_key or not app_secret:
|
||||
logger.warning("钉钉同步凭证未配置（app_key/app_secret），Stream 模式未启动")
|
||||
return
|
||||
|
||||
credential = dingtalk_stream.Credential(app_key, app_secret)
|
||||
client = dingtalk_stream.DingTalkStreamClient(credential)
|
||||
|
||||
handler = DingtalkStreamEventHandler()
|
||||
client.register_all_event_handler(handler)
|
||||
|
||||
cls._client = client
|
||||
cls._connected = True
|
||||
cls._started_at = datetime.now().isoformat()
|
||||
cls._task = asyncio.create_task(cls._run(client))
|
||||
logger.info("钉钉 Stream 模式已启动（app_key=%s）", app_key[:6] + "***")
|
||||
|
||||
@classmethod
|
||||
async def _run(cls, client: dingtalk_stream.DingTalkStreamClient):
|
||||
"""
|
||||
运行 Stream 客户端（带超时和重连）。
|
||||
|
||||
SDK 的 open_connection() 使用同步 requests（短暂阻塞可接受），
|
||||
但 websockets.connect() 可能因为网络/防火墙问题长时间卡住，
|
||||
因此对整个 start() 加超时保护，超时后自动重试。
|
||||
"""
|
||||
import json as _json
|
||||
import websockets as _ws
|
||||
from urllib.parse import quote_plus
|
||||
|
||||
client.pre_start()
|
||||
while cls._connected:
|
||||
try:
|
||||
connection = await asyncio.get_event_loop().run_in_executor(
|
||||
None, client.open_connection
|
||||
)
|
||||
if not connection:
|
||||
logger.error("钉钉 Stream: open_connection 返回空")
|
||||
await asyncio.sleep(10)
|
||||
continue
|
||||
|
||||
logger.info("钉钉 Stream endpoint: %s", connection.get('endpoint', ''))
|
||||
uri = f'{connection["endpoint"]}?ticket={quote_plus(connection["ticket"])}'
|
||||
|
||||
async with asyncio.timeout(30):
|
||||
websocket = await _ws.connect(uri)
|
||||
|
||||
logger.info("钉钉 Stream WebSocket 连接成功")
|
||||
client.websocket = websocket
|
||||
|
||||
asyncio.create_task(client.keepalive(websocket))
|
||||
async for raw_message in websocket:
|
||||
json_message = _json.loads(raw_message)
|
||||
asyncio.create_task(client.background_task(json_message))
|
||||
|
||||
except asyncio.CancelledError:
|
||||
logger.info("钉钉 Stream 任务已取消")
|
||||
break
|
||||
except TimeoutError:
|
||||
logger.warning("钉钉 Stream WebSocket 连接超时（30s），将重试")
|
||||
await asyncio.sleep(5)
|
||||
continue
|
||||
except (_ws.exceptions.ConnectionClosedError, ConnectionError, OSError) as e:
|
||||
logger.warning(f"钉钉 Stream 连接断开: {e}，10s 后重连")
|
||||
await asyncio.sleep(10)
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.error(f"钉钉 Stream 异常: {e}", exc_info=True)
|
||||
await asyncio.sleep(5)
|
||||
continue
|
||||
|
||||
cls._connected = False
|
||||
|
||||
@classmethod
|
||||
async def stop(cls):
|
||||
"""停止 Stream 连接"""
|
||||
cls._connected = False
|
||||
if cls._task:
|
||||
cls._task.cancel()
|
||||
cls._task = None
|
||||
if cls._client:
|
||||
cls._client = None
|
||||
cls._started_at = None
|
||||
logger.info("钉钉 Stream 模式已停止")
|
||||
|
||||
@classmethod
|
||||
def is_running(cls) -> bool:
|
||||
"""
|
||||
检查 Stream 是否正在运行。
|
||||
dingtalk_stream SDK 的 start() 可能在连接建立后返回（task done），
|
||||
但连接实际仍然活跃，因此用 _connected 标志而非仅依赖 task 状态。
|
||||
"""
|
||||
if cls._connected and cls._task is not None:
|
||||
if cls._task.done():
|
||||
ex = cls._task.exception() if not cls._task.cancelled() else None
|
||||
if ex or cls._task.cancelled():
|
||||
cls._connected = False
|
||||
return False
|
||||
return True
|
||||
return True
|
||||
return False
|
||||
|
||||
@classmethod
|
||||
def status(cls) -> dict:
|
||||
"""获取 Stream 模式状态"""
|
||||
return {
|
||||
"stream_mode": True,
|
||||
"running": cls.is_running(),
|
||||
"started_at": cls._started_at,
|
||||
}
|
||||
|
||||
@classmethod
|
||||
def get_event_stats(cls) -> dict:
|
||||
"""获取 Stream 事件统计"""
|
||||
return {
|
||||
**cls.status(),
|
||||
"total_events": cls._event_stats.get("total_events", 0),
|
||||
"last_event_type": cls._event_stats.get("last_event_type"),
|
||||
"last_event_time": cls._event_stats.get("last_event_time"),
|
||||
}
|
||||
Reference in New Issue
Block a user