feat(router): 添加用户追踪和节点通知功能

在ROUTER_MODE启用时跟踪用户消息,并在节点注册/注销时通知相关用户。新增_known_users集合记录活跃用户,重构通知逻辑以支持所有已知用户或特定服务用户的通知。
This commit is contained in:
Yuyao Huang (Sam) 2026-03-29 18:11:00 +08:00
parent ed7bbb1497
commit 2a8f745b3d
2 changed files with 35 additions and 3 deletions

View File

@ -129,6 +129,11 @@ async def _process_message(user_id: str, chat_id: str, text: str) -> None:
await send_text(chat_id, "chat_id", "Sorry, you are not authorized to use this bot.") await send_text(chat_id, "chat_id", "Sorry, you are not authorized to use this bot.")
return return
from config import ROUTER_MODE
if ROUTER_MODE:
from router.nodes import get_node_registry
get_node_registry().track_user(user_id)
reply = await handle_command(user_id, text) reply = await handle_command(user_id, text)
if reply is not None: if reply is not None:
if reply: if reply:

View File

@ -59,6 +59,7 @@ class NodeRegistry:
self._nodes: dict[str, NodeConnection] = {} self._nodes: dict[str, NodeConnection] = {}
self._user_nodes: dict[str, Set[str]] = {} self._user_nodes: dict[str, Set[str]] = {}
self._active_node: dict[str, str] = {} self._active_node: dict[str, str] = {}
self._known_users: Set[str] = set()
self._secret = router_secret self._secret = router_secret
self._lock = asyncio.Lock() self._lock = asyncio.Lock()
@ -68,6 +69,20 @@ class NodeRegistry:
return True return True
return secret == self._secret return secret == self._secret
def track_user(self, user_id: str) -> None:
"""Record a user as known (has sent at least one message)."""
self._known_users.add(user_id)
def _get_notifiable_users(self, node: NodeConnection) -> Set[str]:
"""Get users to notify about a node event.
If the node has explicit serves_users, return those.
Otherwise (serves everyone), return all known users.
"""
if node.serves_users:
return node.serves_users
return self._known_users.copy()
async def register(self, ws: Any, msg: RegisterMessage) -> NodeConnection: async def register(self, ws: Any, msg: RegisterMessage) -> NodeConnection:
"""Register a new node connection.""" """Register a new node connection."""
async with self._lock: async with self._lock:
@ -95,12 +110,24 @@ class NodeRegistry:
msg.capabilities, msg.capabilities,
) )
users_to_notify = self._get_notifiable_users(node)
if is_reconnect: if is_reconnect:
for user_id in msg.serves_users: for user_id in users_to_notify:
asyncio.create_task(self._notify_reconnect(user_id, node.display_name)) asyncio.create_task(self._notify_reconnect(user_id, node.display_name))
else:
for user_id in users_to_notify:
asyncio.create_task(self._notify_new_node(user_id, node.display_name))
return node return node
async def _notify_new_node(self, user_id: str, node_name: str) -> None:
"""Notify user about a new node coming online."""
try:
from bot.feishu import send_text
await send_text(user_id, "open_id", f"🟢 Node \"{node_name}\" is online.")
except Exception as e:
logger.error("Failed to send new node notification: %s", e)
async def _notify_reconnect(self, user_id: str, node_name: str) -> None: async def _notify_reconnect(self, user_id: str, node_name: str) -> None:
"""Notify user about node reconnect.""" """Notify user about node reconnect."""
try: try:
@ -114,7 +141,7 @@ class NodeRegistry:
async with self._lock: async with self._lock:
node = self._nodes.pop(node_id, None) node = self._nodes.pop(node_id, None)
if node: if node:
affected_users = list(node.serves_users) users_to_notify = self._get_notifiable_users(node)
for user_id in node.serves_users: for user_id in node.serves_users:
if user_id in self._user_nodes: if user_id in self._user_nodes:
@ -128,7 +155,7 @@ class NodeRegistry:
logger.info("Node unregistered: %s", node_id) logger.info("Node unregistered: %s", node_id)
for user_id in affected_users: for user_id in users_to_notify:
asyncio.create_task(self._notify_disconnect(user_id, node.display_name)) asyncio.create_task(self._notify_disconnect(user_id, node.display_name))
async def _notify_disconnect(self, user_id: str, node_name: str) -> None: async def _notify_disconnect(self, user_id: str, node_name: str) -> None: