From 0032e4290537d8648a7bca6be42b56a025939fee Mon Sep 17 00:00:00 2001
From: dushixiang <dushixiang@typesafe.cn>
Date: Tue, 24 Feb 2026 19:48:28 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BC=98=E5=8C=96=E6=8E=A2=E9=92=88=E9=87=8D?=
 =?UTF-8?q?=E8=BF=9E=E9=80=BB=E8=BE=91?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/agent/service/agent.go | 23 +++++++++--------------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/pkg/agent/service/agent.go b/pkg/agent/service/agent.go
index 19593a0..1b6551c 100644
--- a/pkg/agent/service/agent.go
+++ b/pkg/agent/service/agent.go
@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/json"
-	"errors"
 	"fmt"
 	"log/slog"
 	"os"
@@ -27,12 +26,6 @@ import (
 	"github.com/sourcegraph/conc"
 )
 
-// 定义特殊错误类型
-var (
-	// ErrConnectionEstablished 表示连接已建立后断开（需要立即重连）
-	ErrConnectionEstablished = errors.New("connection was established")
-)
-
 // safeConn 线程安全的 WebSocket 连接包装器
 type safeConn struct {
 	conn *websocket.Conn
@@ -99,7 +92,7 @@ func (a *Agent) Start(ctx context.Context) error {
 
 	// 启动探针主循环
 	b := &backoff.Backoff{
-		Min:    5 * time.Second,
+		Min:    1 * time.Second,
 		Max:    1 * time.Minute,
 		Factor: 2,
 		Jitter: true,
@@ -147,8 +140,8 @@ func (a *Agent) Stop() {
 }
 
 // runOnce 运行一次探针连接
-// 返回 error 表示需要重连，返回 nil 可能是正常关闭或上下文取消
-func (a *Agent) runOnce(ctx context.Context, onConnected func()) error {
+// 返回 error 表示需要重连，返回 nil 表示上下文取消
+func (a *Agent) runOnce(ctx context.Context, onRegistered func()) error {
 	wsURL := a.cfg.GetWebSocketURL()
 	slog.Info("正在连接到服务器", "url", wsURL)
 
@@ -168,8 +161,6 @@ func (a *Agent) runOnce(ctx context.Context, onConnected func()) error {
 	}
 	defer rawConn.Close()
 
-	onConnected()
-
 	// 创建线程安全的连接包装器
 	conn := &safeConn{conn: rawConn}
 
@@ -187,6 +178,7 @@ func (a *Agent) runOnce(ctx context.Context, onConnected func()) error {
 	if err := a.registerAgent(conn); err != nil {
 		return fmt.Errorf("注册失败: %w", err)
 	}
+	onRegistered()
 
 	slog.Info("探针注册成功，开始监控...")
 
@@ -236,9 +228,8 @@ func (a *Agent) runOnce(ctx context.Context, onConnected func()) error {
 	var returnErr error
 	select {
 	case err := <-errChan:
-		// 连接已建立，无论什么原因断开都标记为已建立状态
 		slog.Info("连接断开", "error", err)
-		returnErr = ErrConnectionEstablished
+		returnErr = err
 	case <-ctx.Done():
 		// 收到取消信号
 		slog.Info("收到停止信号，准备关闭连接")
@@ -247,12 +238,16 @@ func (a *Agent) runOnce(ctx context.Context, onConnected func()) error {
 
 	// 关闭 done channel，通知所有 goroutine 退出
 	close(done)
+	a.setActiveConn(nil)
 
 	// 发送 WebSocket 关闭消息
 	closeMsg := websocket.FormatCloseMessage(websocket.CloseNormalClosure, "")
 	if err := conn.WriteMessage(websocket.CloseMessage, closeMsg); err != nil {
 		slog.Warn("发送关闭消息失败", "error", err)
 	}
+	if err := rawConn.Close(); err != nil && !websocket.IsCloseError(err, websocket.CloseNormalClosure) {
+		slog.Warn("关闭连接失败", "error", err)
+	}
 
 	// 等待所有 goroutine 优雅退出
 	wg.Wait()