feat: 实现多主机架构的核心组件
新增路由器、主机客户端和共享协议模块，支持多主机部署模式： - 路由器作为中央节点管理主机连接和消息路由 - 主机客户端作为工作节点运行本地代理 - 共享协议定义通信消息格式 - 新增独立运行模式standalone.py - 更新配置系统支持路由模式
This commit is contained in:
Yuyao Huang (Sam)
@@ -0,0 +1,6 @@
|
||||
"""Host client module - connects to router and runs local mailboy."""
|
||||
|
||||
from host_client.config import HostConfig, get_host_config
|
||||
from host_client.main import NodeClient
|
||||
|
||||
__all__ = ["HostConfig", "get_host_config", "NodeClient"]
|
||||
@@ -0,0 +1,97 @@
|
||||
"""Host client configuration loader.
|
||||
|
||||
Loads host_config.yaml which contains:
|
||||
- NODE_ID, DISPLAY_NAME
|
||||
- ROUTER_URL, ROUTER_SECRET
|
||||
- LLM config (OPENAI_*)
|
||||
- WORKING_DIR, METASO_API_KEY
|
||||
- SERVES_USERS list
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
import yaml
|
||||
|
||||
|
||||
class HostConfig:
|
||||
"""Configuration for a host client node."""
|
||||
|
||||
def __init__(self, config_path: Optional[Path] = None):
|
||||
config_path = config_path or Path(__file__).parent.parent / "host_config.yaml"
|
||||
self._load(config_path)
|
||||
|
||||
def _load(self, config_path: Path) -> None:
|
||||
if not config_path.exists():
|
||||
raise FileNotFoundError(f"Config file not found: {config_path}")
|
||||
|
||||
with open(config_path, "r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f) or {}
|
||||
|
||||
self.node_id: str = data.get("NODE_ID", "unknown-node")
|
||||
self.display_name: str = data.get("DISPLAY_NAME", self.node_id)
|
||||
self.router_url: str = data.get("ROUTER_URL", "ws://127.0.0.1:8000/ws/node")
|
||||
self.router_secret: str = data.get("ROUTER_SECRET", "")
|
||||
|
||||
self.openai_base_url: str = data.get(
|
||||
"OPENAI_BASE_URL", "https://open.bigmodel.cn/api/paas/v4/"
|
||||
)
|
||||
self.openai_api_key: str = data.get("OPENAI_API_KEY", "")
|
||||
self.openai_model: str = data.get("OPENAI_MODEL", "glm-4.7")
|
||||
|
||||
self.working_dir: str = data.get("WORKING_DIR", str(Path.home() / "projects"))
|
||||
self.metaso_api_key: Optional[str] = data.get("METASO_API_KEY")
|
||||
|
||||
serves_users = data.get("SERVES_USERS", [])
|
||||
self.serves_users: List[str] = serves_users if isinstance(serves_users, list) else []
|
||||
|
||||
self.capabilities: List[str] = data.get(
|
||||
"CAPABILITIES",
|
||||
["claude_code", "shell", "file_ops", "web", "scheduler"],
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_keyring(cls, keyring_path: Optional[Path] = None) -> "HostConfig":
|
||||
"""Create config from keyring.yaml (for standalone mode)."""
|
||||
keyring_path = keyring_path or Path(__file__).parent.parent / "keyring.yaml"
|
||||
if not keyring_path.exists():
|
||||
raise FileNotFoundError(f"keyring.yaml not found: {keyring_path}")
|
||||
|
||||
with open(keyring_path, "r", encoding="utf-8") as f:
|
||||
data = yaml.safe_load(f) or {}
|
||||
|
||||
config = cls.__new__(cls)
|
||||
config.node_id = data.get("NODE_ID", "local-node")
|
||||
config.display_name = data.get("DISPLAY_NAME", "Local Machine")
|
||||
config.router_url = data.get("ROUTER_URL", "ws://127.0.0.1:8000/ws/node")
|
||||
config.router_secret = data.get("ROUTER_SECRET", "")
|
||||
|
||||
config.openai_base_url = data.get(
|
||||
"OPENAI_BASE_URL", "https://open.bigmodel.cn/api/paas/v4/"
|
||||
)
|
||||
config.openai_api_key = data.get("OPENAI_API_KEY", "")
|
||||
config.openai_model = data.get("OPENAI_MODEL", "glm-4.7")
|
||||
|
||||
config.working_dir = data.get("WORKING_DIR", str(Path.home() / "projects"))
|
||||
config.metaso_api_key = data.get("METASO_API_KEY")
|
||||
|
||||
serves_users = data.get("ALLOWED_OPEN_IDS", [])
|
||||
config.serves_users = serves_users if isinstance(serves_users, list) else []
|
||||
|
||||
config.capabilities = ["claude_code", "shell", "file_ops", "web", "scheduler"]
|
||||
|
||||
return config
|
||||
|
||||
|
||||
host_config: Optional[HostConfig] = None
|
||||
|
||||
|
||||
def get_host_config() -> HostConfig:
|
||||
"""Get the global host config instance."""
|
||||
global host_config
|
||||
if host_config is None:
|
||||
host_config = HostConfig()
|
||||
return host_config
|
||||
@@ -0,0 +1,280 @@
|
||||
"""Host client main module.
|
||||
|
||||
Connects to the router via WebSocket, receives forwarded messages,
|
||||
runs the local mailboy LLM, and sends responses back.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import logging
|
||||
import secrets
|
||||
import time
|
||||
from typing import Optional
|
||||
|
||||
import websockets
|
||||
from websockets.client import WebSocketClientProtocol
|
||||
|
||||
from agent.manager import manager
|
||||
from agent.scheduler import scheduler
|
||||
from agent.task_runner import task_runner
|
||||
from host_client.config import HostConfig, get_host_config
|
||||
from orchestrator.agent import run as run_mailboy
|
||||
from orchestrator.tools import set_current_user, set_current_chat
|
||||
from shared import (
|
||||
RegisterMessage,
|
||||
ForwardRequest,
|
||||
ForwardResponse,
|
||||
TaskComplete,
|
||||
Heartbeat,
|
||||
NodeStatus,
|
||||
encode,
|
||||
decode,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class NodeClient:
|
||||
"""WebSocket client that connects to the router and handles messages."""
|
||||
|
||||
def __init__(self, config: HostConfig):
|
||||
self.config = config
|
||||
self.ws: Optional[WebSocketClientProtocol] = None
|
||||
self._running = False
|
||||
self._last_heartbeat = time.time()
|
||||
self._reconnect_delay = 1.0
|
||||
|
||||
async def connect(self) -> bool:
|
||||
"""Connect to the router WebSocket."""
|
||||
headers = {}
|
||||
if self.config.router_secret:
|
||||
headers["Authorization"] = f"Bearer {self.config.router_secret}"
|
||||
|
||||
try:
|
||||
self.ws = await websockets.connect(
|
||||
self.config.router_url,
|
||||
extra_headers=headers,
|
||||
ping_interval=30,
|
||||
ping_timeout=10,
|
||||
)
|
||||
logger.info("Connected to router: %s", self.config.router_url)
|
||||
self._reconnect_delay = 1.0
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("Failed to connect to router: %s", e)
|
||||
return False
|
||||
|
||||
async def register(self) -> bool:
|
||||
"""Send registration message to the router."""
|
||||
if not self.ws:
|
||||
return False
|
||||
|
||||
msg = RegisterMessage(
|
||||
node_id=self.config.node_id,
|
||||
display_name=self.config.display_name,
|
||||
serves_users=self.config.serves_users,
|
||||
working_dir=self.config.working_dir,
|
||||
capabilities=self.config.capabilities,
|
||||
)
|
||||
|
||||
try:
|
||||
await self.ws.send(encode(msg))
|
||||
logger.info("Sent registration for node: %s", self.config.node_id)
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error("Failed to send registration: %s", e)
|
||||
return False
|
||||
|
||||
async def handle_forward(self, request: ForwardRequest) -> None:
|
||||
"""Handle a forwarded message from the router."""
|
||||
logger.info("Received forward request %s from user %s", request.id, request.user_id)
|
||||
|
||||
set_current_user(request.user_id)
|
||||
set_current_chat(request.chat_id)
|
||||
|
||||
try:
|
||||
reply = await run_mailboy(request.user_id, request.text)
|
||||
|
||||
response = ForwardResponse(
|
||||
id=request.id,
|
||||
reply=reply,
|
||||
error="",
|
||||
)
|
||||
except Exception as e:
|
||||
logger.exception("Error processing forward request %s", request.id)
|
||||
response = ForwardResponse(
|
||||
id=request.id,
|
||||
reply="",
|
||||
error=str(e),
|
||||
)
|
||||
|
||||
if self.ws:
|
||||
try:
|
||||
await self.ws.send(encode(response))
|
||||
except Exception as e:
|
||||
logger.error("Failed to send response: %s", e)
|
||||
|
||||
async def send_heartbeat(self) -> None:
|
||||
"""Send a ping heartbeat to the router."""
|
||||
if self.ws:
|
||||
try:
|
||||
await self.ws.send(encode(Heartbeat(type="ping")))
|
||||
self._last_heartbeat = time.time()
|
||||
except Exception as e:
|
||||
logger.error("Failed to send heartbeat: %s", e)
|
||||
|
||||
async def send_status(self) -> None:
|
||||
"""Send node status update to the router."""
|
||||
if not self.ws:
|
||||
return
|
||||
|
||||
sessions = manager.list_sessions()
|
||||
active_sessions = [
|
||||
{"conv_id": s["conv_id"], "working_dir": s["working_dir"]}
|
||||
for s in sessions
|
||||
]
|
||||
|
||||
status = NodeStatus(
|
||||
node_id=self.config.node_id,
|
||||
sessions=len(sessions),
|
||||
active_sessions=active_sessions,
|
||||
)
|
||||
|
||||
try:
|
||||
await self.ws.send(encode(status))
|
||||
except Exception as e:
|
||||
logger.error("Failed to send status: %s", e)
|
||||
|
||||
async def handle_message(self, data: str) -> None:
|
||||
"""Handle an incoming message from the router."""
|
||||
try:
|
||||
msg = decode(data)
|
||||
except Exception as e:
|
||||
logger.error("Failed to decode message: %s", e)
|
||||
return
|
||||
|
||||
if isinstance(msg, ForwardRequest):
|
||||
await self.handle_forward(msg)
|
||||
elif isinstance(msg, Heartbeat):
|
||||
if msg.type == "ping":
|
||||
if self.ws:
|
||||
try:
|
||||
await self.ws.send(encode(Heartbeat(type="pong")))
|
||||
except Exception as e:
|
||||
logger.error("Failed to send pong: %s", e)
|
||||
elif msg.type == "pong":
|
||||
self._last_heartbeat = time.time()
|
||||
else:
|
||||
logger.debug("Received message type: %s", msg.type)
|
||||
|
||||
async def receive_loop(self) -> None:
|
||||
"""Main receive loop for incoming messages."""
|
||||
if not self.ws:
|
||||
return
|
||||
|
||||
try:
|
||||
async for data in self.ws:
|
||||
await self.handle_message(data)
|
||||
except websockets.ConnectionClosed as e:
|
||||
logger.warning("Connection closed: %s", e)
|
||||
except Exception as e:
|
||||
logger.exception("Error in receive loop: %s", e)
|
||||
|
||||
async def heartbeat_loop(self) -> None:
|
||||
"""Periodic heartbeat loop."""
|
||||
while self._running:
|
||||
await asyncio.sleep(30)
|
||||
if self.ws and self.ws.open:
|
||||
await self.send_heartbeat()
|
||||
|
||||
async def status_loop(self) -> None:
|
||||
"""Periodic status update loop."""
|
||||
while self._running:
|
||||
await asyncio.sleep(60)
|
||||
if self.ws and self.ws.open:
|
||||
await self.send_status()
|
||||
|
||||
async def run(self) -> None:
|
||||
"""Main run loop with reconnection."""
|
||||
self._running = True
|
||||
|
||||
await manager.start()
|
||||
await scheduler.start()
|
||||
|
||||
task_runner.set_notification_handler(self._send_task_complete)
|
||||
|
||||
while self._running:
|
||||
if await self.connect():
|
||||
if await self.register():
|
||||
try:
|
||||
await asyncio.gather(
|
||||
self.receive_loop(),
|
||||
self.heartbeat_loop(),
|
||||
self.status_loop(),
|
||||
)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if self._running:
|
||||
logger.info("Reconnecting in %.1f seconds...", self._reconnect_delay)
|
||||
await asyncio.sleep(self._reconnect_delay)
|
||||
self._reconnect_delay = min(self._reconnect_delay * 2, 60)
|
||||
|
||||
async def _send_task_complete(self, task) -> None:
|
||||
"""Send TaskComplete notification to router."""
|
||||
if not self.ws:
|
||||
return
|
||||
|
||||
from shared import TaskComplete, encode
|
||||
|
||||
msg = TaskComplete(
|
||||
task_id=task.task_id,
|
||||
user_id=task.user_id or "",
|
||||
chat_id=task.notify_chat_id or "",
|
||||
result=task.result or task.error or "",
|
||||
)
|
||||
|
||||
try:
|
||||
await self.ws.send(encode(msg))
|
||||
logger.info("Sent TaskComplete for task %s", task.task_id)
|
||||
except Exception as e:
|
||||
logger.error("Failed to send TaskComplete: %s", e)
|
||||
|
||||
async def stop(self) -> None:
|
||||
"""Stop the client."""
|
||||
self._running = False
|
||||
if self.ws:
|
||||
await self.ws.close()
|
||||
await manager.stop()
|
||||
await scheduler.stop()
|
||||
logger.info("Node client stopped")
|
||||
|
||||
@classmethod
|
||||
def from_keyring(cls, router_url: Optional[str] = None, secret: Optional[str] = None) -> "NodeClient":
|
||||
"""Create a client from keyring.yaml (for standalone mode)."""
|
||||
config = HostConfig.from_keyring()
|
||||
if router_url:
|
||||
config.router_url = router_url
|
||||
if secret:
|
||||
config.router_secret = secret
|
||||
return cls(config)
|
||||
|
||||
|
||||
async def main() -> None:
|
||||
"""Entry point for standalone host client."""
|
||||
logging.basicConfig(
|
||||
level=logging.INFO,
|
||||
format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
|
||||
)
|
||||
|
||||
client = NodeClient(get_host_config())
|
||||
|
||||
try:
|
||||
await client.run()
|
||||
except KeyboardInterrupt:
|
||||
await client.stop()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
asyncio.run(main())
|
||||
Reference in New Issue
Block a user