feat: 实现多主机架构的核心组件

新增路由器、主机客户端和共享协议模块，支持多主机部署模式： - 路由器作为中央节点管理主机连接和消息路由 - 主机客户端作为工作节点运行本地代理 - 共享协议定义通信消息格式 - 新增独立运行模式standalone.py - 更新配置系统支持路由模式
2026-03-28 14:08:47 +08:00
parent 8ecc701d5e
commit 64297e5e27
17 changed files with 1338 additions and 6 deletions
@@ -0,0 +1,15 @@
+"""Router module - public-facing component of PhoneWork."""
+
+from router.nodes import NodeRegistry, NodeConnection, get_node_registry
+from router.main import create_app
+from router.rpc import forward
+from router.routing_agent import route
+
+__all__ = [
+    "NodeRegistry",
+    "NodeConnection",
+    "get_node_registry",
+    "create_app",
+    "forward",
+    "route",
+]
@@ -0,0 +1,75 @@
+"""Router main module - FastAPI app factory.
+
+Creates the FastAPI application with:
+- Feishu WebSocket client
+- Node WebSocket endpoint
+- Health check endpoints
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+from fastapi import FastAPI, WebSocket
+from fastapi.middleware.cors import CORSMiddleware
+
+from bot.handler import start_websocket_client
+from router.nodes import NodeRegistry, get_node_registry
+from router.ws import ws_node_endpoint
+
+logger = logging.getLogger(__name__)
+
+
+def create_app(router_secret: Optional[str] = None) -> FastAPI:
+    """Create the FastAPI application.
+
+    Args:
+        router_secret: Secret for authenticating host client connections
+    """
+    app = FastAPI(title="PhoneWork Router", version="3.0.0")
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=["*"],
+        allow_credentials=True,
+        allow_methods=["*"],
+        allow_headers=["*"],
+    )
+
+    registry = get_node_registry()
+    if router_secret:
+        registry._secret = router_secret
+
+    @app.get("/health")
+    async def health():
+        nodes = registry.list_nodes()
+        online_nodes = [n for n in nodes if n["status"] == "online"]
+        return {
+            "status": "ok",
+            "nodes": nodes,
+            "online_nodes": len(online_nodes),
+            "total_nodes": len(nodes),
+            "pending_requests": 0,
+        }
+
+    @app.get("/nodes")
+    async def list_nodes():
+        return registry.list_nodes()
+
+    @app.websocket("/ws/node")
+    async def ws_node(websocket: WebSocket):
+        await ws_node_endpoint(websocket)
+
+    @app.on_event("startup")
+    async def startup():
+        import asyncio
+        loop = asyncio.get_event_loop()
+        start_websocket_client(loop)
+        logger.info("Router started")
+
+    @app.on_event("shutdown")
+    async def shutdown():
+        logger.info("Router shut down")
+
+    return app
@@ -0,0 +1,209 @@
+"""Node registry for managing connected host clients.
+
+Maintains:
+- Connected nodes with their WebSocket connections
+- User-to-node mapping (which users each node serves)
+- Active node preference per user
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import time
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Set
+
+from shared import RegisterMessage, NodeStatus
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class NodeConnection:
+    """Represents a connected host client."""
+    node_id: str
+    ws: Any
+    display_name: str = ""
+    serves_users: Set[str] = field(default_factory=set)
+    working_dir: str = ""
+    capabilities: List[str] = field(default_factory=list)
+    connected_at: float = field(default_factory=time.time)
+    last_heartbeat: float = field(default_factory=time.time)
+    sessions: int = 0
+    active_sessions: List[Dict[str, Any]] = field(default_factory=list)
+
+    @property
+    def is_online(self) -> bool:
+        """Check if node is still considered online (heartbeat within 60s)."""
+        return time.time() - self.last_heartbeat < 60
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Serialize for API responses."""
+        return {
+            "node_id": self.node_id,
+            "display_name": self.display_name,
+            "status": "online" if self.is_online else "offline",
+            "users": len(self.serves_users),
+            "sessions": self.sessions,
+            "capabilities": self.capabilities,
+            "connected_at": self.connected_at,
+            "last_heartbeat": self.last_heartbeat,
+        }
+
+
+class NodeRegistry:
+    """Registry of connected host clients."""
+
+    def __init__(self, router_secret: str = ""):
+        self._nodes: Dict[str, NodeConnection] = {}
+        self._user_nodes: Dict[str, Set[str]] = {}
+        self._active_node: Dict[str, str] = {}
+        self._secret = router_secret
+        self._lock = asyncio.Lock()
+
+    def validate_secret(self, secret: str) -> bool:
+        """Validate router secret."""
+        if not self._secret:
+            return True
+        return secret == self._secret
+
+    async def register(self, ws: Any, msg: RegisterMessage) -> NodeConnection:
+        """Register a new node connection."""
+        async with self._lock:
+            is_reconnect = msg.node_id in self._nodes
+
+            node = NodeConnection(
+                node_id=msg.node_id,
+                ws=ws,
+                display_name=msg.display_name or msg.node_id,
+                serves_users=set(msg.serves_users),
+                working_dir=msg.working_dir,
+                capabilities=msg.capabilities,
+            )
+            self._nodes[msg.node_id] = node
+
+            for user_id in msg.serves_users:
+                if user_id not in self._user_nodes:
+                    self._user_nodes[user_id] = set()
+                self._user_nodes[user_id].add(msg.node_id)
+
+            logger.info(
+                "Node registered: %s (users: %s, capabilities: %s)",
+                msg.node_id,
+                msg.serves_users,
+                msg.capabilities,
+            )
+
+            if is_reconnect:
+                for user_id in msg.serves_users:
+                    asyncio.create_task(self._notify_reconnect(user_id, node.display_name))
+
+            return node
+
+    async def _notify_reconnect(self, user_id: str, node_name: str) -> None:
+        """Notify user about node reconnect."""
+        try:
+            from bot.feishu import send_text
+            await send_text(user_id, "open_id", f"✅ Node \"{node_name}\" reconnected.")
+        except Exception as e:
+            logger.error("Failed to send reconnect notification: %s", e)
+
+    async def unregister(self, node_id: str) -> None:
+        """Unregister a node connection."""
+        async with self._lock:
+            node = self._nodes.pop(node_id, None)
+            if node:
+                affected_users = list(node.serves_users)
+
+                for user_id in node.serves_users:
+                    if user_id in self._user_nodes:
+                        self._user_nodes[user_id].discard(node_id)
+                        if not self._user_nodes[user_id]:
+                            del self._user_nodes[user_id]
+
+                for user_id in list(self._active_node.keys()):
+                    if self._active_node[user_id] == node_id:
+                        del self._active_node[user_id]
+
+                logger.info("Node unregistered: %s", node_id)
+
+                for user_id in affected_users:
+                    asyncio.create_task(self._notify_disconnect(user_id, node.display_name))
+
+    async def _notify_disconnect(self, user_id: str, node_name: str) -> None:
+        """Notify user about node disconnect."""
+        try:
+            from bot.feishu import send_text
+            await send_text(user_id, "open_id", f"⚠️ Node \"{node_name}\" disconnected.")
+        except Exception as e:
+            logger.error("Failed to send disconnect notification: %s", e)
+
+    async def update_status(self, msg: NodeStatus) -> None:
+        """Update node status from heartbeat."""
+        async with self._lock:
+            node = self._nodes.get(msg.node_id)
+            if node:
+                node.sessions = msg.sessions
+                node.active_sessions = msg.active_sessions
+                node.last_heartbeat = time.time()
+
+    async def update_heartbeat(self, node_id: str) -> None:
+        """Update node heartbeat timestamp."""
+        async with self._lock:
+            node = self._nodes.get(node_id)
+            if node:
+                node.last_heartbeat = time.time()
+
+    def get_node(self, node_id: str) -> Optional[NodeConnection]:
+        """Get a node by ID."""
+        return self._nodes.get(node_id)
+
+    def get_nodes_for_user(self, user_id: str) -> List[NodeConnection]:
+        """Get all nodes that serve a user."""
+        node_ids = self._user_nodes.get(user_id, set())
+        return [self._nodes[nid] for nid in node_ids if nid in self._nodes]
+
+    def get_active_node(self, user_id: str) -> Optional[NodeConnection]:
+        """Get the active node for a user."""
+        node_id = self._active_node.get(user_id)
+        if node_id:
+            return self._nodes.get(node_id)
+
+        nodes = self.get_nodes_for_user(user_id)
+        if nodes:
+            online = [n for n in nodes if n.is_online]
+            if online:
+                return online[0]
+
+        return None
+
+    def set_active_node(self, user_id: str, node_id: str) -> bool:
+        """Set the active node for a user."""
+        if node_id not in self._nodes:
+            return False
+        self._active_node[user_id] = node_id
+        logger.info("Active node for user %s set to %s", user_id, node_id)
+        return True
+
+    def list_nodes(self) -> List[Dict[str, Any]]:
+        """List all nodes with their status."""
+        return [node.to_dict() for node in self._nodes.values()]
+
+    def get_affected_users(self, node_id: str) -> List[str]:
+        """Get users affected by a node disconnect."""
+        node = self._nodes.get(node_id)
+        if node:
+            return list(node.serves_users)
+        return []
+
+
+node_registry: Optional[NodeRegistry] = None
+
+
+def get_node_registry() -> NodeRegistry:
+    """Get the global node registry instance."""
+    global node_registry
+    if node_registry is None:
+        node_registry = NodeRegistry()
+    return node_registry
@@ -0,0 +1,128 @@
+"""Routing LLM for deciding which node to forward messages to.
+
+This is a lightweight, one-shot LLM call that decides routing.
+No history, no multi-step loop. Single call with one tool.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from typing import List, Optional
+
+from langchain_core.messages import HumanMessage, SystemMessage
+from langchain_openai import ChatOpenAI
+
+from config import OPENAI_API_KEY, OPENAI_BASE_URL, OPENAI_MODEL
+from router.nodes import NodeConnection, get_node_registry
+
+logger = logging.getLogger(__name__)
+
+ROUTING_SYSTEM_PROMPT = """You are a routing assistant. A user has sent a message. \
+Choose which node to forward it to.
+
+Connected nodes for this user:
+{nodes_info}
+
+Rules:
+- If the message references an active session on a node, route to that node.
+- If the user names a machine explicitly ("on work-server", "@home-pc"), route there.
+- If only one node is connected, route there without asking.
+- If ambiguous with multiple idle nodes, ask the user to clarify.
+- For meta commands (/nodes, /help, /status), respond with "meta" as the node_id.
+
+Respond with a JSON object:
+{{"node_id": "<node_id>", "reason": "<brief reason>"}}
+"""
+
+
+def _format_nodes_info(nodes: List[NodeConnection], active_node_id: Optional[str] = None) -> str:
+    """Format node information for the routing prompt."""
+    lines = []
+    for node in nodes:
+        marker = " [ACTIVE]" if node.node_id == active_node_id else ""
+        sessions = ", ".join(
+            s.get("working_dir", "unknown") for s in node.active_sessions[:3]
+        ) or "none"
+        lines.append(
+            f"- {node.display_name or node.node_id}{marker}: "
+            f"sessions=[{sessions}], capabilities={node.capabilities}"
+        )
+    return "\n".join(lines)
+
+
+async def route(user_id: str, chat_id: str, text: str) -> tuple[Optional[str], str]:
+    """Determine which node to route a message to.
+
+    Args:
+        user_id: User's Feishu open_id
+        chat_id: Chat ID for context
+        text: User's message text
+
+    Returns:
+        Tuple of (node_id, reason). node_id is None if no suitable node found.
+    """
+    registry = get_node_registry()
+    nodes = registry.get_nodes_for_user(user_id)
+
+    if not nodes:
+        return None, "No nodes available for this user"
+
+    online_nodes = [n for n in nodes if n.is_online]
+    if not online_nodes:
+        return None, "All nodes for this user are offline"
+
+    if len(online_nodes) == 1:
+        return online_nodes[0].node_id, "Only one node available"
+
+    if text.strip().startswith("/"):
+        return "meta", "Meta command"
+
+    active_node = registry.get_active_node(user_id)
+    active_node_id = active_node.node_id if active_node else None
+
+    nodes_info = _format_nodes_info(online_nodes, active_node_id)
+
+    try:
+        llm = ChatOpenAI(
+            model=OPENAI_MODEL,
+            openai_api_key=OPENAI_API_KEY,
+            openai_api_base=OPENAI_BASE_URL,
+            temperature=0,
+        )
+
+        prompt = ROUTING_SYSTEM_PROMPT.format(nodes_info=nodes_info)
+        messages = [
+            SystemMessage(content=prompt),
+            HumanMessage(content=text),
+        ]
+
+        response = await llm.ainvoke(messages)
+        content = response.content.strip()
+
+        if content.startswith("```"):
+            content = content.split("\n", 1)[1]
+            content = content.rsplit("```", 1)[0]
+
+        result = json.loads(content)
+        node_id = result.get("node_id")
+        reason = result.get("reason", "")
+
+        if node_id == "meta":
+            return "meta", reason
+
+        for node in online_nodes:
+            if node.node_id == node_id or node.display_name == node_id:
+                return node.node_id, reason
+
+        if active_node:
+            return active_node.node_id, f"Defaulting to active node (LLM suggested unavailable: {node_id})"
+
+        return online_nodes[0].node_id, f"Defaulting to first available node (LLM suggested: {node_id})"
+
+    except Exception as e:
+        logger.warning("Routing LLM failed: %s, falling back to active node", e)
+
+        if active_node:
+            return active_node.node_id, "Fallback to active node"
+        return online_nodes[0].node_id, "Fallback to first available node"
@@ -0,0 +1,109 @@
+"""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)
@@ -0,0 +1,102 @@
+"""WebSocket endpoint for host client connections.
+
+Handles:
+- Connection authentication
+- Node registration
+- Message forwarding
+- Heartbeat
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import Optional
+
+from fastapi import WebSocket, WebSocketDisconnect, WebSocketException
+
+from router.nodes import get_node_registry
+from router.rpc import handle_task_complete
+from shared import (
+    RegisterMessage,
+    ForwardRequest,
+    ForwardResponse,
+    TaskComplete,
+    Heartbeat,
+    NodeStatus,
+    decode,
+    encode,
+)
+
+logger = logging.getLogger(__name__)
+
+
+async def ws_node_endpoint(websocket: WebSocket) -> None:
+    """WebSocket endpoint for host client connections."""
+    await websocket.accept()
+
+    registry = get_node_registry()
+
+    secret = websocket.headers.get("authorization", "")
+    if secret.startswith("Bearer "):
+        secret = secret[7:]
+
+    if not registry.validate_secret(secret):
+        logger.warning("Invalid router secret, rejecting connection")
+        await websocket.close(code=4001, reason="Invalid secret")
+        return
+
+    node_id: Optional[str] = None
+    heartbeat_task: Optional[asyncio.Task] = None
+
+    async def send_heartbeat():
+        """Send periodic pings to the host client."""
+        try:
+            while True:
+                await asyncio.sleep(30)
+                try:
+                    await websocket.send_text(encode(Heartbeat(type="ping")))
+                except Exception:
+                    break
+        except asyncio.CancelledError:
+            pass
+
+    try:
+        async for data in websocket.iter_text():
+            try:
+                msg = decode(data)
+            except Exception as e:
+                logger.error("Failed to decode message: %s", e)
+                continue
+
+            if isinstance(msg, RegisterMessage):
+                node_id = msg.node_id
+                await registry.register(websocket, msg)
+                heartbeat_task = asyncio.create_task(send_heartbeat())
+
+            elif isinstance(msg, ForwardResponse):
+                from router.rpc import resolve_response
+                await resolve_response(msg)
+
+            elif isinstance(msg, TaskComplete):
+                await handle_task_complete(msg)
+
+            elif isinstance(msg, Heartbeat):
+                if msg.type == "pong" and node_id:
+                    await registry.update_heartbeat(node_id)
+
+            elif isinstance(msg, NodeStatus):
+                await registry.update_status(msg)
+
+            else:
+                logger.debug("Received unhandled message type: %s", type(msg).__name__)
+
+    except WebSocketDisconnect:
+        logger.info("WebSocket disconnected")
+    except Exception as e:
+        logger.exception("WebSocket error: %s", e)
+    finally:
+        if heartbeat_task:
+            heartbeat_task.cancel()
+        if node_id:
+            await registry.unregister(node_id)