@agentunion/kite 1.4.0 → 1.6.0

package/extensions/services/proxy/agentcp/heartbeat/heartbeat_client.py

@@ -0,0 +1,360 @@
+# -*- 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.
+import requests
+import datetime
+import socket
+import threading
+import time
+from typing import Optional
+from agentcp.base.log import log_debug, log_error, log_exception, log_info, log_warning
+from agentcp.base.auth_client import AuthClient
+
+from agentcp.msg.message_serialize import *
+from ..context import ErrorContext, exceptions
+
+
+class HeartbeatClient:
+    # 重连相关常量
+    MAX_SEND_FAILURES = 3           # 发送失败触发重连的阈值
+    MAX_RECV_FAILURES = 3           # 接收失败触发重连的阈值
+    MAX_MISSED_HEARTBEATS = 3       # 心跳响应超时阈值（错过次数）
+    RECONNECT_BACKOFF_MAX = 30      # 重连退避上限（秒）
+    SOCKET_TIMEOUT = 1.0            # socket 超时时间（秒）
+
+    def __init__(self, agent_id: str, server_url: str, aid_path: str, seed_password: str):
+        self.agent_id = agent_id
+        self.server_url = server_url
+        self.seed_password = seed_password
+        self.port = 0  # server_port
+        self.sign_cookie = 0
+        self.udp_socket = None
+        self.local_ip = "0.0.0.0"
+        self.local_port = 0
+        self.server_ip = "127.0.0.1"
+        self.heartbeat_interval = 5000
+        self.is_running = False
+        self.is_sending_heartbeat = False
+        self.send_thread: Optional[threading.Thread] = None
+        self.receive_thread: Optional[threading.Thread] = None
+        self.msg_seq = 0
+        self.last_hb = 0
+        self.message_listener = None
+        self.auth_client = AuthClient(agent_id, server_url, aid_path, seed_password)
+        self.on_recv_invite = None
+
+        # 新增：用于自动恢复的状态
+        self._socket_lock = threading.Lock()        # 保护 socket 操作
+        self._reconnect_lock = threading.Lock()     # 防止并发重连
+        self._last_reconnect_ts = 0                 # 上次重连时间戳
+        self._last_hb_recv = 0                      # 上次收到心跳响应的时间戳
+        self._send_failures = 0                     # 连续发送失败次数
+        self._recv_failures = 0                     # 连续接收失败次数
+
+    def initialize(self):
+        self.sign_in()
+
+    def sign_in(self) -> bool:
+        data = self.auth_client.sign_in()
+        if data is None:
+            log_error("sign_in failed: data is None")
+            return False
+        self.server_ip = data.get("server_ip")
+        self.port = int(data.get("port", 0))
+        self.sign_cookie = data.get("sign_cookie")
+        log_info(f'signin {self.server_ip} {self.port} {self.sign_cookie}')
+
+        return self.server_ip is not None and self.port != 0 and self.sign_cookie is not None
+
+    def sign_out(self):
+        self.auth_client.sign_out()
+
+    def set_on_recv_invite(self, listener):
+        """设置消息监听器"""
+        self.on_recv_invite = listener
+
+    # ========== 新增：Socket 生命周期管理 ==========
+
+    def _create_socket(self):
+        """创建并绑定 UDP socket，设置超时"""
+        with self._socket_lock:
+            self._close_socket_internal()
+            self.udp_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+            self.udp_socket.settimeout(self.SOCKET_TIMEOUT)
+            self.udp_socket.bind((self.local_ip, 0))  # 使用新端口
+            self.local_ip, self.local_port = self.udp_socket.getsockname()
+            log_info(f"UDP socket created and bound to {self.local_ip}:{self.local_port}")
+
+    def _close_socket_internal(self):
+        """内部方法：关闭 socket（不加锁，由调用方保证线程安全）"""
+        if self.udp_socket is not None:
+            try:
+                self.udp_socket.close()
+            except Exception as e:
+                log_warning(f"Close socket error: {e}")
+            self.udp_socket = None
+
+    def _close_socket(self):
+        """安全关闭 socket"""
+        with self._socket_lock:
+            self._close_socket_internal()
+
+    def _reconnect(self, reason: str):
+        """重连：限流/退避后执行 sign_in() + _create_socket()"""
+        if not self.is_running:
+            log_debug(f"Reconnect skipped (client offline): {reason}")
+            return False
+        if not self._reconnect_lock.acquire(blocking=False):
+            log_debug(f"Reconnect already in progress, skip: {reason}")
+            return False
+
+        try:
+            now = time.time()
+            # 限流：距离上次重连至少间隔 5 秒
+            elapsed = now - self._last_reconnect_ts
+            if elapsed < 5:
+                backoff = min(5 - elapsed, self.RECONNECT_BACKOFF_MAX)
+                log_info(f"Reconnect backoff: waiting {backoff:.1f}s")
+                time.sleep(backoff)
+
+            log_info(f"Reconnecting due to: {reason}")
+            self._last_reconnect_ts = time.time()
+
+            # 重新登录
+            if not self.sign_in():
+                log_error("Reconnect failed: sign_in returned False")
+                return False
+
+            # 重建 socket
+            self._create_socket()
+
+            # 重置失败计数
+            self._send_failures = 0
+            self._recv_failures = 0
+            self._last_hb_recv = int(time.time() * 1000)
+
+            log_info("Reconnect successful")
+            return True
+        except Exception as e:
+            log_error(f"Reconnect exception: {e}")
+            return False
+        finally:
+            self._reconnect_lock.release()
+
+    # ========== 发送心跳（带异常恢复和超时检测） ==========
+
+    def __send_heartbeat(self):
+        backoff = 1  # 初始退避时间（秒）
+
+        while self.is_sending_heartbeat and self.is_running:
+            try:
+                current_time_ms = int(datetime.datetime.now().timestamp() * 1000)
+
+                # 检查心跳响应超时
+                if self._last_hb_recv > 0:
+                    timeout_threshold = self.MAX_MISSED_HEARTBEATS * self.heartbeat_interval
+                    if current_time_ms - self._last_hb_recv > timeout_threshold:
+                        log_warning(f"Heartbeat response timeout: {current_time_ms - self._last_hb_recv}ms > {timeout_threshold}ms")
+                        self._reconnect("heartbeat_response_timeout")
+                        backoff = 1
+                        continue
+
+                # 发送心跳
+                if current_time_ms > (self.last_hb + self.heartbeat_interval):
+                    log_debug(f'send heartbeat message to {self.server_ip}:{self.port}')
+                    self.last_hb = current_time_ms
+                    self.msg_seq = self.msg_seq + 1
+                    req = HeartbeatMessageReq()
+                    req.header.MessageMask = 0
+                    req.header.MessageSeq = self.msg_seq
+                    req.header.MessageType = 513
+                    req.header.PayloadSize = 100
+                    req.AgentId = self.agent_id
+                    req.SignCookie = self.sign_cookie
+                    buf = io.BytesIO()
+                    req.serialize(buf)
+                    data = buf.getvalue()
+
+                    with self._socket_lock:
+                        if self.udp_socket is not None:
+                            self.udp_socket.sendto(data, (self.server_ip, self.port))
+                        else:
+                            raise Exception("UDP socket is None")
+
+                    # 发送成功，重置失败计数和退避
+                    self._send_failures = 0
+                    backoff = 1
+
+                time.sleep(1)
+
+            except Exception as e:
+                self._send_failures += 1
+                log_error(f"Heartbeat send error (failures={self._send_failures}): {e}")
+                ErrorContext.publish(exceptions.SDKError(f"Heartbeat send error: {e}"))
+
+                # 达到阈值，触发重连
+                if self._send_failures >= self.MAX_SEND_FAILURES:
+                    log_warning(f"Send failures reached threshold ({self.MAX_SEND_FAILURES}), triggering reconnect")
+                    self._reconnect("send_failures_threshold")
+                    backoff = 1
+                else:
+                    # 指数退避
+                    time.sleep(backoff)
+                    backoff = min(backoff * 2, self.RECONNECT_BACKOFF_MAX)
+
+    # ========== 接收消息（可中断、可恢复） ==========
+
+    def _receive_messages(self):
+        while self.is_running:
+            try:
+                # 使用 socket 超时，确保能定期检查 is_running
+                with self._socket_lock:
+                    if self.udp_socket is None:
+                        time.sleep(0.5)
+                        continue
+                    sock = self.udp_socket
+
+                try:
+                    data, addr = sock.recvfrom(1536)
+                except socket.timeout:
+                    # 超时是正常的，继续循环检查 is_running
+                    continue
+
+                # 接收成功，重置失败计数
+                self._recv_failures = 0
+
+                udp_header, offset = UdpMessageHeader.deserialize(data, 0)
+
+                if udp_header.MessageType == 258:
+                    hb_resp, offset = HeartbeatMessageResp.deserialize(data, 0)
+                    self.heartbeat_interval = hb_resp.NextBeat
+
+                    # 更新最后收到心跳响应的时间
+                    self._last_hb_recv = int(datetime.datetime.now().timestamp() * 1000)
+
+                    # 服务器端身份验证失败(比如服务器发生了异常重启)，需要重新登录
+                    if hb_resp.NextBeat == 401:
+                        log_warning(f"Heartbeat failed: {hb_resp.NextBeat}, triggering reconnect")
+                        ErrorContext.publish(exceptions.SDKError(f"401,心跳", code=0))
+                        self._reconnect("401_auth_failed")
+                        continue
+
+                    if self.heartbeat_interval <= 5000:
+                        self.heartbeat_interval = 5000
+
+                elif udp_header.MessageType == 259:
+                    invite_req, offset = InviteMessageReq.deserialize(data, 0)
+                    if self.on_recv_invite is not None:
+                        ErrorContext.publish(exceptions.SDKError(f"收到邀请，加入session: {invite_req}", code=0))
+                        self.on_recv_invite(invite_req)
+
+                    resp = InviteMessageResp()
+                    self.msg_seq = self.msg_seq + 1
+                    resp.header.MessageMask = 0
+                    resp.header.MessageSeq = self.msg_seq
+                    resp.header.MessageType = 516
+                    resp.AgentId = self.agent_id
+                    resp.InviterAgentId = invite_req.InviterAgentId
+                    resp.SignCookie = self.sign_cookie
+                    buf = io.BytesIO()
+                    resp.serialize(buf)
+                    resp_data = buf.getvalue()
+
+                    with self._socket_lock:
+                        if self.udp_socket is not None:
+                            self.udp_socket.sendto(resp_data, (self.server_ip, self.port))
+
+            except socket.timeout:
+                # 超时是正常的，继续循环
+                continue
+            except Exception as e:
+                if not self.is_running:
+                    break
+
+                self._recv_failures += 1
+                log_error(f"Receive message exception (failures={self._recv_failures}): {e}")
+                ErrorContext.publish(exceptions.SDKError(f"Receive message exception: {e}"))
+
+                # 达到阈值，触发重连
+                if self._recv_failures >= self.MAX_RECV_FAILURES:
+                    log_warning(f"Recv failures reached threshold ({self.MAX_RECV_FAILURES}), triggering reconnect")
+                    self._reconnect("recv_failures_threshold")
+                else:
+                    time.sleep(1.5)
+
+    def online(self):
+        """开始心跳"""
+        if self.is_running:
+            return
+
+        # 使用统一的 socket 创建方法
+        self._create_socket()
+
+        # 初始化心跳响应时间
+        self._last_hb_recv = int(time.time() * 1000)
+
+        self.is_running = True
+        self.is_sending_heartbeat = True
+
+        self.send_thread = threading.Thread(target=self.__send_heartbeat, daemon=True)
+        self.receive_thread = threading.Thread(target=self._receive_messages, daemon=True)
+
+        self.send_thread.start()
+        self.receive_thread.start()
+        log_info("Successfully went online")
+
+    def offline(self):
+        """停止心跳"""
+        log_info("Going offline...")
+
+        # 1. 先设置标志位，通知线程退出
+        self.is_running = False
+        self.is_sending_heartbeat = False
+
+        # 2. 关闭 socket（会使阻塞的 recvfrom 抛出异常）
+        self._close_socket()
+
+        # 3. 等待线程退出
+        if self.send_thread is not None and self.send_thread.is_alive():
+            self.send_thread.join(timeout=3)
+            if self.send_thread.is_alive():
+                log_warning("Send thread did not exit in time")
+
+        if self.receive_thread is not None and self.receive_thread.is_alive():
+            self.receive_thread.join(timeout=3)
+            if self.receive_thread.is_alive():
+                log_warning("Receive thread did not exit in time")
+
+        self.send_thread = None
+        self.receive_thread = None
+        log_info("Successfully went offline")
+    
+    def get_online_status(self, aids):
+        try:
+            ep_url = self.server_url + "/query_online_state"
+            data = {
+                "agent_id": f"{self.agent_id}",
+                "signature": self.auth_client.signature,
+                "agents": aids
+            }
+            response = requests.post(ep_url, json=data, verify=False, proxies={}, timeout=(3, 10))
+            if response.status_code == 200:
+                log_info(f"get_online_status ok:{response.json()}")
+                return response.json()["data"]
+            else:
+                log_error(f"get_online_status failed:{response.json()}")
+                return []
+        except Exception as e:
+            log_exception(f"get_online_status exception: {e}")
+            return []