@agentunion/kite 1.3.2 → 1.5.0

package/extensions/services/proxy/agentcp/msg/connection_manager.py

@@ -0,0 +1,456 @@
+# -*- coding: utf-8 -*-
+# Copyright 2025 AgentUnion Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+ConnectionManager - WebSocket 连接管理器
+
+提供 MessageClient 的生命周期管理，支持：
+1. 连接状态查询
+2. 主动销毁和重建连接
+3. 断开/恢复回调
+4. 连接健康监控
+
+每个 AgentID 对应一个 ConnectionManager，管理该 Agent 的所有 MessageClient。
+"""
+
+import threading
+import time
+from typing import Dict, Optional, Callable, List
+from dataclasses import dataclass, field
+from enum import Enum
+
+from agentcp.base.log import log_info, log_error, log_warning, log_debug
+from agentcp.msg.message_client import MessageClient, MessageClientConfig, ConnectionState
+
+
+class ConnectionEvent(Enum):
+    """连接事件类型"""
+    CONNECTED = "connected"           # 连接建立
+    DISCONNECTED = "disconnected"     # 连接断开
+    RECONNECTING = "reconnecting"     # 正在重连
+    RECONNECTED = "reconnected"       # 重连成功
+    DESTROYED = "destroyed"           # 连接销毁
+    CREATED = "created"               # 连接创建
+
+
+@dataclass
+class ConnectionInfo:
+    """连接信息"""
+    server_url: str
+    message_client: Optional[MessageClient] = None
+    created_at: float = 0.0
+    last_connected_at: float = 0.0
+    last_disconnected_at: float = 0.0
+    disconnect_count: int = 0
+    reconnect_count: int = 0
+    is_destroyed: bool = False
+
+
+class ConnectionManager:
+    """WebSocket 连接管理器
+
+    管理一个 AgentID 下的所有 MessageClient 连接。
+
+    主要功能：
+    1. 连接生命周期管理（创建、销毁、重建）
+    2. 连接状态查询
+    3. 事件回调（断开、恢复）
+    4. 健康监控
+
+    使用示例：
+        ```python
+        # 创建连接管理器
+        conn_mgr = ConnectionManager(agent_id, aid_path, seed_password)
+
+        # 设置回调
+        conn_mgr.set_event_callback(on_connection_event)
+
+        # 获取或创建连接
+        mc = conn_mgr.get_or_create_connection(server_url)
+
+        # 查询连接状态
+        if conn_mgr.is_healthy(server_url):
+            mc.send_msg(...)
+
+        # 销毁并重建连接
+        conn_mgr.rebuild_connection(server_url)
+
+        # 关闭所有连接
+        conn_mgr.destroy_all()
+        ```
+    """
+
+    def __init__(
+        self,
+        agent_id: str,
+        aid_path: str,
+        seed_password: str,
+        config: Optional[MessageClientConfig] = None
+    ):
+        """初始化连接管理器
+
+        Args:
+            agent_id: Agent ID
+            aid_path: AID 证书路径
+            seed_password: 种子密码
+            config: MessageClient 配置（可选，不传则使用默认配置）
+        """
+        self.agent_id = agent_id
+        self.aid_path = aid_path
+        self.seed_password = seed_password
+        self.config = config or MessageClientConfig()
+
+        # 连接映射：server_url -> ConnectionInfo
+        self._connections: Dict[str, ConnectionInfo] = {}
+        self._lock = threading.RLock()
+
+        # 事件回调
+        self._event_callback: Optional[Callable[[str, ConnectionEvent, dict], None]] = None
+
+        # 是否已关闭
+        self._shutdown = False
+
+        log_info(f"[ConnectionManager] 初始化: agent_id={agent_id}")
+
+    def set_event_callback(self, callback: Callable[[str, ConnectionEvent, dict], None]) -> None:
+        """设置连接事件回调
+
+        回调函数签名: callback(server_url: str, event: ConnectionEvent, info: dict)
+
+        Args:
+            callback: 事件回调函数
+        """
+        self._event_callback = callback
+        log_info(f"[ConnectionManager] 已设置事件回调")
+
+    def _fire_event(self, server_url: str, event: ConnectionEvent, extra_info: dict = None) -> None:
+        """触发事件回调"""
+        if self._event_callback:
+            try:
+                info = extra_info or {}
+                info["agent_id"] = self.agent_id
+                info["server_url"] = server_url
+                info["timestamp"] = time.time()
+                self._event_callback(server_url, event, info)
+            except Exception as e:
+                log_error(f"[ConnectionManager] 事件回调异常: {e}")
+
+    def get_or_create_connection(
+        self,
+        server_url: str,
+        cache_auth_client=None,
+        message_handler=None
+    ) -> Optional[MessageClient]:
+        """获取或创建连接
+
+        如果连接已存在且未销毁，返回现有连接。
+        如果连接不存在或已销毁，创建新连接。
+
+        Args:
+            server_url: 消息服务器 URL
+            cache_auth_client: 缓存的认证客户端（可选）
+            message_handler: 消息处理器（可选）
+
+        Returns:
+            MessageClient 实例，失败返回 None
+        """
+        if self._shutdown:
+            log_warning(f"[ConnectionManager] 已关闭，无法创建连接")
+            return None
+
+        server_url = server_url.rstrip("/")
+
+        with self._lock:
+            conn_info = self._connections.get(server_url)
+
+            # 已存在且未销毁，复用
+            if conn_info and not conn_info.is_destroyed and conn_info.message_client:
+                log_debug(f"[ConnectionManager] 复用现有连接: {server_url}")
+                return conn_info.message_client
+
+            # 需要创建新连接
+            log_info(f"[ConnectionManager] 创建新连接: {server_url}")
+
+            try:
+                mc = MessageClient(
+                    agent_id=self.agent_id,
+                    server_url=server_url,
+                    aid_path=self.aid_path,
+                    seed_password=self.seed_password,
+                    cache_auth_client=cache_auth_client,
+                    config=self.config
+                )
+                mc.initialize()
+
+                if message_handler:
+                    mc.set_message_handler(message_handler)
+
+                # 设置回调
+                mc.set_disconnect_callback(self._on_disconnect)
+                mc.set_reconnect_callback(self._on_reconnect)
+
+                # 记录连接信息
+                conn_info = ConnectionInfo(
+                    server_url=server_url,
+                    message_client=mc,
+                    created_at=time.time(),
+                    is_destroyed=False
+                )
+                self._connections[server_url] = conn_info
+
+                self._fire_event(server_url, ConnectionEvent.CREATED)
+                return mc
+
+            except Exception as e:
+                log_error(f"[ConnectionManager] 创建连接失败: {server_url}, error={e}")
+                return None
+
+    def get_connection(self, server_url: str) -> Optional[MessageClient]:
+        """获取现有连接（不创建）
+
+        Args:
+            server_url: 消息服务器 URL
+
+        Returns:
+            MessageClient 实例，不存在返回 None
+        """
+        server_url = server_url.rstrip("/")
+        with self._lock:
+            conn_info = self._connections.get(server_url)
+            if conn_info and not conn_info.is_destroyed:
+                return conn_info.message_client
+            return None
+
+    def is_healthy(self, server_url: str) -> bool:
+        """检查连接是否健康
+
+        Args:
+            server_url: 消息服务器 URL
+
+        Returns:
+            True: 连接健康可用
+            False: 连接不可用
+        """
+        mc = self.get_connection(server_url)
+        if mc:
+            return mc.is_healthy()
+        return False
+
+    def get_connection_info(self, server_url: str) -> Optional[dict]:
+        """获取连接详细信息
+
+        Args:
+            server_url: 消息服务器 URL
+
+        Returns:
+            连接信息字典，不存在返回 None
+        """
+        server_url = server_url.rstrip("/")
+        with self._lock:
+            conn_info = self._connections.get(server_url)
+            if not conn_info:
+                return None
+
+            mc = conn_info.message_client
+            mc_info = mc.get_connection_info() if mc else {}
+
+            return {
+                "server_url": server_url,
+                "created_at": conn_info.created_at,
+                "last_connected_at": conn_info.last_connected_at,
+                "last_disconnected_at": conn_info.last_disconnected_at,
+                "disconnect_count": conn_info.disconnect_count,
+                "reconnect_count": conn_info.reconnect_count,
+                "is_destroyed": conn_info.is_destroyed,
+                **mc_info
+            }
+
+    def get_all_connections_info(self) -> List[dict]:
+        """获取所有连接的信息
+
+        Returns:
+            连接信息列表
+        """
+        with self._lock:
+            result = []
+            for server_url in self._connections:
+                info = self.get_connection_info(server_url)
+                if info:
+                    result.append(info)
+            return result
+
+    def get_health_summary(self) -> dict:
+        """获取健康状态摘要
+
+        Returns:
+            健康状态摘要字典
+        """
+        with self._lock:
+            total = len(self._connections)
+            healthy = sum(1 for url in self._connections if self.is_healthy(url))
+            destroyed = sum(1 for info in self._connections.values() if info.is_destroyed)
+
+            return {
+                "agent_id": self.agent_id,
+                "total_connections": total,
+                "healthy_connections": healthy,
+                "unhealthy_connections": total - healthy - destroyed,
+                "destroyed_connections": destroyed,
+                "is_all_healthy": healthy == total and total > 0
+            }
+
+    def destroy_connection(self, server_url: str, wait: bool = True) -> bool:
+        """销毁指定连接
+
+        Args:
+            server_url: 消息服务器 URL
+            wait: 是否等待连接完全关闭
+
+        Returns:
+            True: 销毁成功
+            False: 连接不存在
+        """
+        server_url = server_url.rstrip("/")
+
+        with self._lock:
+            conn_info = self._connections.get(server_url)
+            if not conn_info:
+                log_warning(f"[ConnectionManager] 连接不存在: {server_url}")
+                return False
+
+            if conn_info.is_destroyed:
+                log_debug(f"[ConnectionManager] 连接已销毁: {server_url}")
+                return True
+
+            log_info(f"[ConnectionManager] 销毁连接: {server_url}")
+
+            mc = conn_info.message_client
+            if mc:
+                try:
+                    mc.stop_websocket_client()
+                except Exception as e:
+                    log_error(f"[ConnectionManager] 停止 WebSocket 异常: {e}")
+
+            conn_info.is_destroyed = True
+            conn_info.message_client = None
+
+            self._fire_event(server_url, ConnectionEvent.DESTROYED)
+            return True
+
+    def rebuild_connection(
+        self,
+        server_url: str,
+        cache_auth_client=None,
+        message_handler=None
+    ) -> Optional[MessageClient]:
+        """销毁并重建连接
+
+        Args:
+            server_url: 消息服务器 URL
+            cache_auth_client: 缓存的认证客户端（可选）
+            message_handler: 消息处理器（可选）
+
+        Returns:
+            新的 MessageClient 实例，失败返回 None
+        """
+        log_info(f"[ConnectionManager] 重建连接: {server_url}")
+
+        # 先销毁
+        self.destroy_connection(server_url, wait=True)
+
+        # 短暂等待确保资源释放
+        time.sleep(0.3)
+
+        # 再创建
+        return self.get_or_create_connection(server_url, cache_auth_client, message_handler)
+
+    def destroy_all(self, wait: bool = True) -> None:
+        """销毁所有连接
+
+        Args:
+            wait: 是否等待所有连接完全关闭
+        """
+        log_info(f"[ConnectionManager] 销毁所有连接...")
+        self._shutdown = True
+
+        with self._lock:
+            for server_url in list(self._connections.keys()):
+                self.destroy_connection(server_url, wait=wait)
+
+        log_info(f"[ConnectionManager] 所有连接已销毁")
+
+    def force_reconnect(self, server_url: str) -> bool:
+        """强制触发重连
+
+        不销毁连接，而是触发 MessageClient 的重连逻辑。
+
+        Args:
+            server_url: 消息服务器 URL
+
+        Returns:
+            True: 触发成功
+            False: 连接不存在
+        """
+        mc = self.get_connection(server_url)
+        if not mc:
+            return False
+
+        log_info(f"[ConnectionManager] 强制重连: {server_url}")
+
+        # 停止当前连接
+        mc.stop_websocket_client()
+
+        # 短暂等待
+        time.sleep(0.2)
+
+        # 重新启动
+        return mc.start_websocket_client()
+
+    def _on_disconnect(self, agent_id: str, server_url: str, code: int, reason: str) -> None:
+        """连接断开回调"""
+        server_url = server_url.rstrip("/")
+        log_warning(f"[ConnectionManager] 连接断开: {server_url}, code={code}, reason={reason}")
+
+        with self._lock:
+            conn_info = self._connections.get(server_url)
+            if conn_info:
+                conn_info.last_disconnected_at = time.time()
+                conn_info.disconnect_count += 1
+
+        self._fire_event(server_url, ConnectionEvent.DISCONNECTED, {
+            "code": code,
+            "reason": reason
+        })
+
+    def _on_reconnect(self, agent_id: str, server_url: str) -> None:
+        """连接恢复回调"""
+        server_url = server_url.rstrip("/")
+        log_info(f"[ConnectionManager] 连接恢复: {server_url}")
+
+        with self._lock:
+            conn_info = self._connections.get(server_url)
+            if conn_info:
+                conn_info.last_connected_at = time.time()
+                conn_info.reconnect_count += 1
+
+        self._fire_event(server_url, ConnectionEvent.RECONNECTED)
+
+    def __del__(self):
+        """析构时清理"""
+        try:
+            if not self._shutdown:
+                self.destroy_all(wait=False)
+        except Exception:
+            pass