09b63341cd
- 将 Dict、List 等传统类型注解替换为 dict、list 等现代类型注解 - 更新类型注解以更精确地反映变量类型 - 修复部分类型注解与实际使用不匹配的问题 - 优化部分代码逻辑以提高类型安全性
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[str]] = {}
|
|
_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[str] = 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)
|