64297e5e27
新增路由器、主机客户端和共享协议模块，支持多主机部署模式： - 路由器作为中央节点管理主机连接和消息路由 - 主机客户端作为工作节点运行本地代理 - 共享协议定义通信消息格式 - 新增独立运行模式standalone.py - 更新配置系统支持路由模式
110 lines
3.0 KiB
Python
110 lines
3.0 KiB
Python
"""RPC module for forwarding requests to host clients.
|
|
|
|
Handles:
|
|
- Request correlation with asyncio.Future
|
|
- Timeout management
|
|
- Response routing
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import uuid
|
|
from typing import Any, Dict, Optional
|
|
|
|
from shared import ForwardRequest, ForwardResponse, TaskComplete, encode
|
|
from router.nodes import get_node_registry
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_pending_requests: Dict[str, asyncio.Future] = {}
|
|
_default_timeout = 600.0
|
|
|
|
|
|
async def forward(
|
|
node_id: str,
|
|
user_id: str,
|
|
chat_id: str,
|
|
text: str,
|
|
timeout: float = _default_timeout,
|
|
) -> str:
|
|
"""Forward a message to a host client and wait for response.
|
|
|
|
Args:
|
|
node_id: Target node ID
|
|
user_id: User's Feishu open_id
|
|
chat_id: Chat ID for context
|
|
text: Message text to forward
|
|
timeout: Timeout in seconds (default 600s for long CC tasks)
|
|
|
|
Returns:
|
|
Reply text from the host client
|
|
|
|
Raises:
|
|
asyncio.TimeoutError: If no response within timeout
|
|
RuntimeError: If node is not connected
|
|
"""
|
|
registry = get_node_registry()
|
|
node = registry.get_node(node_id)
|
|
|
|
if not node or not node.ws:
|
|
raise RuntimeError(f"Node not connected: {node_id}")
|
|
|
|
request_id = str(uuid.uuid4())
|
|
future: asyncio.Future = asyncio.get_event_loop().create_future()
|
|
_pending_requests[request_id] = future
|
|
|
|
request = ForwardRequest(
|
|
id=request_id,
|
|
user_id=user_id,
|
|
chat_id=chat_id,
|
|
text=text,
|
|
)
|
|
|
|
try:
|
|
await node.ws.send_text(encode(request))
|
|
logger.debug("Forwarded request %s to node %s", request_id, node_id)
|
|
|
|
result = await asyncio.wait_for(future, timeout=timeout)
|
|
return result
|
|
|
|
except asyncio.TimeoutError:
|
|
logger.warning("Request %s timed out after %ss", request_id, timeout)
|
|
raise
|
|
|
|
except Exception as e:
|
|
logger.error("Failed to forward request %s: %s", request_id, e)
|
|
raise
|
|
|
|
finally:
|
|
_pending_requests.pop(request_id, None)
|
|
|
|
|
|
async def resolve_response(response: ForwardResponse) -> None:
|
|
"""Resolve a pending request with a response."""
|
|
future = _pending_requests.get(response.id)
|
|
if future and not future.done():
|
|
if response.error:
|
|
future.set_exception(RuntimeError(response.error))
|
|
else:
|
|
future.set_result(response.reply)
|
|
logger.debug("Resolved request %s", response.id)
|
|
|
|
|
|
async def handle_task_complete(msg: TaskComplete) -> None:
|
|
"""Handle a task completion notification from a host client."""
|
|
logger.info("Task %s completed for user %s", msg.task_id, msg.user_id)
|
|
|
|
from bot.feishu import send_text
|
|
|
|
try:
|
|
await send_text(msg.chat_id, "chat_id", msg.result)
|
|
except Exception as e:
|
|
logger.error("Failed to send task completion notification: %s", e)
|
|
|
|
|
|
def get_pending_count() -> int:
|
|
"""Get the number of pending requests."""
|
|
return len(_pending_requests)
|