[新增] 实现 Agent WebSocket 连接升级功能，支持状态上报和配置广播

openflare_agent/internal/agent/runner.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"log/slog"
 	"strings"
@@ -33,6 +34,12 @@ type RuntimeManager interface {
 	Restart(ctx context.Context) error
 }
 
+type WebSocketService interface {
+	Connect(ctx context.Context) (protocol.WebSocketConnection, error)
+	SetToken(token string)
+	URL() string
+}
+
 type UpdateOptions struct {
 	Channel string
 	TagName string
@@ -47,13 +54,15 @@ type Runner struct {
 	SyncService         SyncService
 	Updater             Updater
 	RuntimeManager      RuntimeManager
+	WebSocketService    WebSocketService
 
-	autoUpdate          bool
-	updateNow           bool
-	updateRepo          string
-	updateChan          string
-	updateTag           string
-	restartOpenrestyNow bool
+	autoUpdate              bool
+	updateNow               bool
+	updateRepo              string
+	updateChan              string
+	updateTag               string
+	restartOpenrestyNow     bool
+	websocketUpgradeEnabled bool
 }
 
 func (r *Runner) Run(ctx context.Context) error {
@@ -63,26 +72,8 @@ func (r *Runner) Run(ctx context.Context) error {
 	}
 	slog.Info("agent runner started", "node_id", nodeID, "node", r.Config.NodeName, "ip", r.Config.NodeIP)
 	if r.hasAgentToken() {
-		r.refreshOpenrestyHealth(ctx)
-		payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
-		heartbeatResult, hbErr := r.HeartbeatService.Heartbeat(ctx, payload)
-		if hbErr != nil {
+		if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, true); hbErr != nil {
 			slog.Error("agent startup heartbeat failed", "error", hbErr)
-		} else {
-			r.ackObservabilityWindows(ackWindows)
-			if heartbeatResult == nil {
-				heartbeatResult = &protocol.HeartbeatResult{}
-			}
-			slog.Debug("agent startup heartbeat succeeded", "node_id", nodeID)
-			r.applySettings(heartbeatResult.AgentSettings)
-			if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
-				r.recordSyncError(err)
-				slog.Error("agent startup sync failed", "error", err)
-			} else {
-				slog.Debug("agent startup sync completed")
-			}
-			r.tryRestartOpenresty(ctx)
-			r.tryAutoUpdate(ctx)
 		}
 	} else if err = r.tryRegister(ctx, &nodeID); err != nil {
 		slog.Error("agent initial discovery register failed", "error", err)
@@ -90,43 +81,266 @@ func (r *Runner) Run(ctx context.Context) error {
 
 	heartbeatTicker := time.NewTicker(r.Config.HeartbeatInterval.Duration())
 	defer heartbeatTicker.Stop()
+	var wsDone <-chan error
+	wsBackoff := newWebSocketBackoff()
+	nextWSAttempt := time.Now()
+	tryStartWebSocket := func() {
+		if wsDone != nil || !r.shouldUseWebSocket() || time.Now().Before(nextWSAttempt) {
+			return
+		}
+		done, startErr := r.startWebSocket(ctx, nodeID)
+		if startErr != nil {
+			delay := wsBackoff.Next()
+			nextWSAttempt = time.Now().Add(delay)
+			slog.Debug("agent ws upgrade failed; falling back to http heartbeat",
+				"enabled", r.websocketUpgradeEnabled,
+				"url", r.websocketURL(),
+				"retry_after", delay,
+				"error", startErr,
+			)
+			return
+		}
+		wsBackoff.Reset()
+		wsDone = done
+		slog.Debug("agent switched to websocket mode", "url", r.websocketURL())
+	}
+	tryStartWebSocket()
 
 	for {
 		select {
 		case <-ctx.Done():
 			slog.Info("agent runner shutting down", "error", ctx.Err())
 			return ctx.Err()
+		case wsErr := <-wsDone:
+			wsDone = nil
+			delay := wsBackoff.Next()
+			nextWSAttempt = time.Now().Add(delay)
+			slog.Debug("agent ws disconnected; resuming http heartbeat", "retry_after", delay, "error", wsErr)
+			if r.hasAgentToken() {
+				if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
+					slog.Error("agent heartbeat after ws disconnect failed", "error", hbErr)
+				}
+			}
 		case <-heartbeatTicker.C:
+			if wsDone != nil {
+				continue
+			}
 			if !r.hasAgentToken() {
 				if err = r.tryRegister(ctx, &nodeID); err != nil {
 					slog.Error("agent discovery register failed", "error", err)
 				}
 				continue
 			}
-			r.refreshOpenrestyHealth(ctx)
-			payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
-			heartbeatResult, hbErr := r.HeartbeatService.Heartbeat(ctx, payload)
-			if hbErr != nil {
+			if changed, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
 				slog.Error("agent heartbeat failed", "error", hbErr)
 			} else {
-				r.ackObservabilityWindows(ackWindows)
-				if heartbeatResult == nil {
-					heartbeatResult = &protocol.HeartbeatResult{}
-				}
-				if changed := r.applySettings(heartbeatResult.AgentSettings); changed {
+				if changed {
 					heartbeatTicker.Reset(r.Config.HeartbeatInterval.Duration())
 				}
-				if err = r.SyncService.SyncOnce(ctx, heartbeatResult.ActiveConfig); err != nil {
-					r.recordSyncError(err)
-					slog.Error("agent sync failed", "error", err)
-				}
-				r.tryRestartOpenresty(ctx)
-				r.tryAutoUpdate(ctx)
+				tryStartWebSocket()
 			}
 		}
 	}
 }
 
+func (r *Runner) performHeartbeatCycle(ctx context.Context, nodeID string, startup bool) (bool, error) {
+	r.refreshOpenrestyHealth(ctx)
+	payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
+	heartbeatResult, err := r.HeartbeatService.Heartbeat(ctx, payload)
+	if err != nil {
+		return false, err
+	}
+	r.ackObservabilityWindows(ackWindows)
+	if heartbeatResult == nil {
+		heartbeatResult = &protocol.HeartbeatResult{}
+	}
+	mode := "periodic"
+	if startup {
+		mode = "startup"
+	}
+	slog.Debug("agent heartbeat succeeded", "mode", mode, "node_id", nodeID)
+	changed := r.applySettings(heartbeatResult.AgentSettings)
+	if startup {
+		if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
+			r.recordSyncError(err)
+			slog.Error("agent startup sync failed", "error", err)
+		} else {
+			slog.Debug("agent startup sync completed")
+		}
+	} else if err = r.SyncService.SyncOnce(ctx, heartbeatResult.ActiveConfig); err != nil {
+		r.recordSyncError(err)
+		slog.Error("agent sync failed", "error", err)
+	}
+	r.tryRestartOpenresty(ctx)
+	r.tryAutoUpdate(ctx)
+	return changed, nil
+}
+
+func (r *Runner) shouldUseWebSocket() bool {
+	enabled := r.WebSocketService != nil && r.websocketUpgradeEnabled && r.hasAgentToken()
+	slog.Debug("agent ws upgrade eligibility checked", "enabled", enabled, "server_enabled", r.websocketUpgradeEnabled, "url", r.websocketURL())
+	return enabled
+}
+
+func (r *Runner) websocketURL() string {
+	if r.WebSocketService == nil {
+		return ""
+	}
+	return r.WebSocketService.URL()
+}
+
+func (r *Runner) startWebSocket(ctx context.Context, nodeID string) (<-chan error, error) {
+	if r.WebSocketService == nil {
+		return nil, errors.New("websocket service is not configured")
+	}
+	conn, err := r.WebSocketService.Connect(ctx)
+	if err != nil {
+		return nil, err
+	}
+	done := make(chan error, 1)
+	go func() {
+		defer func() {
+			_ = conn.Close()
+		}()
+		done <- r.runWebSocket(ctx, nodeID, conn)
+	}()
+	return done, nil
+}
+
+func (r *Runner) runWebSocket(ctx context.Context, nodeID string, conn protocol.WebSocketConnection) error {
+	slog.Debug("agent ws connected", "url", conn.URL(), "node_id", nodeID)
+	statusTicker := time.NewTicker(r.Config.HeartbeatInterval.Duration())
+	defer statusTicker.Stop()
+
+	messages := make(chan protocol.WSMessage, 8)
+	readDone := make(chan error, 1)
+	go func() {
+		for {
+			message, err := conn.Receive()
+			if err != nil {
+				readDone <- err
+				return
+			}
+			select {
+			case messages <- message:
+			case <-ctx.Done():
+				readDone <- ctx.Err()
+				return
+			}
+		}
+	}()
+
+	if err := r.sendWebSocketStatus(ctx, nodeID, conn); err != nil {
+		return err
+	}
+
+	for {
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case err := <-readDone:
+			return err
+		case <-statusTicker.C:
+			if err := r.sendWebSocketStatus(ctx, nodeID, conn); err != nil {
+				return err
+			}
+		case message := <-messages:
+			changed, err := r.handleWebSocketMessage(ctx, message, conn)
+			if err != nil {
+				return err
+			}
+			if changed {
+				statusTicker.Reset(r.Config.HeartbeatInterval.Duration())
+			}
+		}
+	}
+}
+
+func (r *Runner) sendWebSocketStatus(ctx context.Context, nodeID string, conn protocol.WebSocketConnection) error {
+	r.refreshOpenrestyHealth(ctx)
+	payload, ackWindows := r.prepareHeartbeatPayload(nodeID)
+	if err := conn.SendStatus(payload); err != nil {
+		return err
+	}
+	r.ackObservabilityWindows(ackWindows)
+	return nil
+}
+
+func (r *Runner) handleWebSocketMessage(ctx context.Context, message protocol.WSMessage, conn protocol.WebSocketConnection) (bool, error) {
+	switch message.Type {
+	case protocol.WSMessageTypeSettings:
+		var settings protocol.AgentSettings
+		if err := json.Unmarshal(message.Payload, &settings); err != nil {
+			slog.Debug("agent ws settings decode failed", "error", err)
+			return false, nil
+		}
+		changed := r.applySettings(&settings)
+		r.tryRestartOpenresty(ctx)
+		r.tryAutoUpdate(ctx)
+		if !r.websocketUpgradeEnabled {
+			slog.Debug("agent ws disabled by server settings; falling back to http heartbeat")
+			return changed, errors.New("websocket upgrade disabled by server")
+		}
+		return changed, nil
+	case protocol.WSMessageTypeActiveConfig:
+		var target protocol.ActiveConfigMeta
+		if err := json.Unmarshal(message.Payload, &target); err != nil {
+			slog.Debug("agent ws active config decode failed", "error", err)
+			return false, nil
+		}
+		slog.Debug("agent ws active config received", "version", target.Version, "checksum", target.Checksum, "trigger_sync", true)
+		if err := r.SyncService.SyncOnce(ctx, &target); err != nil {
+			r.recordSyncError(err)
+			slog.Error("agent ws triggered sync failed", "version", target.Version, "error", err)
+		}
+		return false, nil
+	case protocol.WSMessageTypePing:
+		slog.Debug("agent ws ping received")
+		return false, conn.SendPong()
+	case protocol.WSMessageTypePong:
+		slog.Debug("agent ws pong received")
+		return false, nil
+	default:
+		slog.Debug("agent ws unsupported message type", "type", message.Type)
+		return false, nil
+	}
+}
+
+type webSocketBackoff struct {
+	delays []time.Duration
+	index  int
+}
+
+func newWebSocketBackoff() *webSocketBackoff {
+	return &webSocketBackoff{
+		delays: []time.Duration{
+			time.Second,
+			2 * time.Second,
+			5 * time.Second,
+			10 * time.Second,
+			30 * time.Second,
+		},
+	}
+}
+
+func (backoff *webSocketBackoff) Next() time.Duration {
+	if backoff == nil || len(backoff.delays) == 0 {
+		return 30 * time.Second
+	}
+	if backoff.index >= len(backoff.delays) {
+		return backoff.delays[len(backoff.delays)-1]
+	}
+	delay := backoff.delays[backoff.index]
+	backoff.index++
+	return delay
+}
+
+func (backoff *webSocketBackoff) Reset() {
+	if backoff != nil {
+		backoff.index = 0
+	}
+}
+
 func (r *Runner) hasAgentToken() bool {
 	return strings.TrimSpace(r.Config.AgentToken) != ""
 }
@@ -144,6 +358,10 @@ func (r *Runner) applySettings(settings *protocol.AgentSettings) bool {
 			changed = true
 		}
 	}
+	if settings.WebsocketUpgradeEnabled != r.websocketUpgradeEnabled {
+		slog.Debug("agent websocket upgrade setting updated", "from", r.websocketUpgradeEnabled, "to", settings.WebsocketUpgradeEnabled)
+	}
+	r.websocketUpgradeEnabled = settings.WebsocketUpgradeEnabled
 	r.autoUpdate = settings.AutoUpdate
 	r.updateNow = settings.UpdateNow
 	r.updateRepo = strings.TrimSpace(settings.UpdateRepo)
@@ -222,6 +440,9 @@ func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error {
 		return err
 	}
 	r.HeartbeatService.SetToken(response.AgentToken)
+	if r.WebSocketService != nil {
+		r.WebSocketService.SetToken(response.AgentToken)
+	}
 	*nodeID = response.NodeID
 	slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
 	r.refreshOpenrestyHealth(ctx)

openflare_agent/internal/agent/runner_test.go

@@ -2,6 +2,7 @@ package agent
 
 import (
 	"context"
+	"encoding/json"
 	"errors"
 	"os"
 	"path/filepath"
@@ -67,6 +68,7 @@ type fakeSyncService struct {
 	syncOnceErr    error
 	startupCalls   int
 	syncOnceCalls  int
+	lastTarget     *protocol.ActiveConfigMeta
 	onSyncOnceCall func(int)
 }
 
@@ -104,6 +106,10 @@ func (f *fakeSyncService) SyncOnStartup(ctx context.Context, target *protocol.Ac
 func (f *fakeSyncService) SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error {
 	f.mu.Lock()
 	f.syncOnceCalls++
+	if target != nil {
+		copied := *target
+		f.lastTarget = &copied
+	}
 	callIndex := f.syncOnceCalls
 	callback := f.onSyncOnceCall
 	f.mu.Unlock()
@@ -113,6 +119,31 @@ func (f *fakeSyncService) SyncOnce(ctx context.Context, target *protocol.ActiveC
 	return f.syncOnceErr
 }
 
+type fakeWebSocketConnection struct {
+	pongCalls int
+}
+
+func (f *fakeWebSocketConnection) URL() string {
+	return "ws://127.0.0.1/api/agent/ws"
+}
+
+func (f *fakeWebSocketConnection) SendStatus(payload protocol.NodePayload) error {
+	return nil
+}
+
+func (f *fakeWebSocketConnection) SendPong() error {
+	f.pongCalls++
+	return nil
+}
+
+func (f *fakeWebSocketConnection) Receive() (protocol.WSMessage, error) {
+	return protocol.WSMessage{}, errors.New("not implemented")
+}
+
+func (f *fakeWebSocketConnection) Close() error {
+	return nil
+}
+
 func TestRunnerKeepsHeartbeatWhenStartupSyncFails(t *testing.T) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
@@ -495,3 +526,83 @@ func TestRunnerDiscoveryRegisterUpdatesTokenAndNodeID(t *testing.T) {
 		t.Fatal("expected config token rotation to complete")
 	}
 }
+
+func TestRunnerHandlesWebSocketActiveConfigMessage(t *testing.T) {
+	syncService := &fakeSyncService{}
+	runner := &Runner{SyncService: syncService}
+	payload, err := json.Marshal(protocol.ActiveConfigMeta{
+		Version:  "20260529-001",
+		Checksum: "checksum-ws",
+	})
+	if err != nil {
+		t.Fatalf("marshal active config: %v", err)
+	}
+
+	changed, err := runner.handleWebSocketMessage(context.Background(), protocol.WSMessage{
+		Type:    protocol.WSMessageTypeActiveConfig,
+		Payload: payload,
+	}, &fakeWebSocketConnection{})
+	if err != nil {
+		t.Fatalf("handle websocket active config: %v", err)
+	}
+	if changed {
+		t.Fatal("active config message should not change heartbeat interval")
+	}
+	if syncService.syncOnceCalls != 1 {
+		t.Fatalf("expected one sync call, got %d", syncService.syncOnceCalls)
+	}
+	if syncService.lastTarget == nil || syncService.lastTarget.Version != "20260529-001" || syncService.lastTarget.Checksum != "checksum-ws" {
+		t.Fatalf("unexpected sync target: %+v", syncService.lastTarget)
+	}
+}
+
+func TestRunnerHandlesWebSocketSettingsDisabled(t *testing.T) {
+	runner := &Runner{
+		Config: &config.Config{
+			HeartbeatInterval: config.MillisecondDuration(10 * time.Second),
+		},
+		websocketUpgradeEnabled: true,
+	}
+	payload, err := json.Marshal(protocol.AgentSettings{
+		HeartbeatInterval:       15000,
+		WebsocketUpgradeEnabled: false,
+	})
+	if err != nil {
+		t.Fatalf("marshal settings: %v", err)
+	}
+
+	changed, err := runner.handleWebSocketMessage(context.Background(), protocol.WSMessage{
+		Type:    protocol.WSMessageTypeSettings,
+		Payload: payload,
+	}, &fakeWebSocketConnection{})
+	if err == nil {
+		t.Fatal("expected disabled websocket setting to request fallback")
+	}
+	if !changed {
+		t.Fatal("expected heartbeat interval change to be reported")
+	}
+	if runner.websocketUpgradeEnabled {
+		t.Fatal("expected websocket upgrade to be disabled")
+	}
+}
+
+func TestWebSocketBackoffSequence(t *testing.T) {
+	backoff := newWebSocketBackoff()
+	expected := []time.Duration{
+		time.Second,
+		2 * time.Second,
+		5 * time.Second,
+		10 * time.Second,
+		30 * time.Second,
+		30 * time.Second,
+	}
+	for _, want := range expected {
+		if got := backoff.Next(); got != want {
+			t.Fatalf("unexpected backoff: got %s want %s", got, want)
+		}
+	}
+	backoff.Reset()
+	if got := backoff.Next(); got != time.Second {
+		t.Fatalf("expected reset backoff to return 1s, got %s", got)
+	}
+}

openflare_agent/internal/protocol/agent_api.go

@@ -1,5 +1,7 @@
 package protocol
 
+import "encoding/json"
+
 type APIResponse[T any] struct {
 	Success bool   `json:"success"`
 	Message string `json:"message"`
@@ -20,13 +22,40 @@ type HeartbeatResult struct {
 }
 
 type AgentSettings struct {
-	HeartbeatInterval   int    `json:"heartbeat_interval"`
-	AutoUpdate          bool   `json:"auto_update"`
-	UpdateRepo          string `json:"update_repo"`
-	UpdateNow           bool   `json:"update_now"`
-	UpdateChannel       string `json:"update_channel"`
-	UpdateTag           string `json:"update_tag"`
-	RestartOpenrestyNow bool   `json:"restart_openresty_now"`
+	HeartbeatInterval       int    `json:"heartbeat_interval"`
+	WebsocketUpgradeEnabled bool   `json:"websocket_upgrade_enabled"`
+	AutoUpdate              bool   `json:"auto_update"`
+	UpdateRepo              string `json:"update_repo"`
+	UpdateNow               bool   `json:"update_now"`
+	UpdateChannel           string `json:"update_channel"`
+	UpdateTag               string `json:"update_tag"`
+	RestartOpenrestyNow     bool   `json:"restart_openresty_now"`
+}
+
+const (
+	WSMessageTypeStatus       = "status"
+	WSMessageTypeSettings     = "settings"
+	WSMessageTypeActiveConfig = "active_config"
+	WSMessageTypePing         = "ping"
+	WSMessageTypePong         = "pong"
+)
+
+type WSMessage struct {
+	Type    string          `json:"type"`
+	Payload json.RawMessage `json:"payload,omitempty"`
+}
+
+type WSOutboundMessage struct {
+	Type    string `json:"type"`
+	Payload any    `json:"payload,omitempty"`
+}
+
+type WebSocketConnection interface {
+	URL() string
+	SendStatus(payload NodePayload) error
+	SendPong() error
+	Receive() (WSMessage, error)
+	Close() error
 }
 
 const (

openflare_agent/internal/wsclient/client.go

@@ -0,0 +1,165 @@
+package wsclient
+
+import (
+	"context"
+	"errors"
+	"log/slog"
+	"net"
+	"net/http"
+	"net/url"
+	"strings"
+	"time"
+
+	"golang.org/x/net/websocket"
+
+	"openflare-agent/internal/protocol"
+)
+
+type Client struct {
+	baseURL string
+	token   string
+	timeout time.Duration
+}
+
+type Connection struct {
+	conn        *websocket.Conn
+	url         string
+	readTimeout time.Duration
+}
+
+func New(baseURL string, token string, timeout time.Duration) *Client {
+	return &Client{
+		baseURL: strings.TrimRight(baseURL, "/"),
+		token:   strings.TrimSpace(token),
+		timeout: timeout,
+	}
+}
+
+func (c *Client) SetToken(token string) {
+	c.token = strings.TrimSpace(token)
+	slog.Debug("agent ws client token updated")
+}
+
+func (c *Client) URL() string {
+	wsURL, err := buildWebsocketURL(c.baseURL)
+	if err != nil {
+		return ""
+	}
+	return wsURL
+}
+
+func (c *Client) Connect(ctx context.Context) (protocol.WebSocketConnection, error) {
+	wsURL, err := buildWebsocketURL(c.baseURL)
+	if err != nil {
+		return nil, err
+	}
+	if strings.TrimSpace(c.token) == "" {
+		return nil, errors.New("agent ws token is empty")
+	}
+	origin := strings.TrimSpace(c.baseURL)
+	if origin == "" {
+		origin = "http://localhost"
+	}
+	config, err := websocket.NewConfig(wsURL, origin)
+	if err != nil {
+		return nil, err
+	}
+	config.Header = http.Header{}
+	config.Header.Set("X-Agent-Token", c.token)
+	if c.timeout > 0 {
+		config.Dialer = &net.Dialer{Timeout: c.timeout}
+	}
+	slog.Debug("agent ws dialing server", "url", wsURL)
+	conn, err := config.DialContext(ctx)
+	if err != nil {
+		return nil, err
+	}
+	slog.Debug("agent ws dial succeeded", "url", wsURL)
+	return &Connection{conn: conn, url: wsURL, readTimeout: websocketReadTimeout(c.timeout)}, nil
+}
+
+func buildWebsocketURL(baseURL string) (string, error) {
+	parsed, err := url.Parse(strings.TrimRight(baseURL, "/"))
+	if err != nil {
+		return "", err
+	}
+	switch parsed.Scheme {
+	case "http":
+		parsed.Scheme = "ws"
+	case "https":
+		parsed.Scheme = "wss"
+	case "ws", "wss":
+	default:
+		return "", errors.New("server_url scheme must be http, https, ws, or wss")
+	}
+	parsed.Path = strings.TrimRight(parsed.Path, "/") + "/api/agent/ws"
+	parsed.RawQuery = ""
+	parsed.Fragment = ""
+	return parsed.String(), nil
+}
+
+func (conn *Connection) URL() string {
+	if conn == nil {
+		return ""
+	}
+	return conn.url
+}
+
+func (conn *Connection) SendStatus(payload protocol.NodePayload) error {
+	if conn == nil || conn.conn == nil {
+		return errors.New("agent ws connection is nil")
+	}
+	slog.Debug("agent ws sending status",
+		"node_id", payload.NodeID,
+		"current_version", payload.CurrentVersion,
+		"openresty_status", payload.OpenrestyStatus,
+	)
+	return websocket.JSON.Send(conn.conn, protocol.WSOutboundMessage{
+		Type:    protocol.WSMessageTypeStatus,
+		Payload: payload,
+	})
+}
+
+func (conn *Connection) SendPong() error {
+	if conn == nil || conn.conn == nil {
+		return errors.New("agent ws connection is nil")
+	}
+	slog.Debug("agent ws sending pong")
+	return websocket.JSON.Send(conn.conn, protocol.WSOutboundMessage{
+		Type: protocol.WSMessageTypePong,
+	})
+}
+
+func (conn *Connection) Receive() (protocol.WSMessage, error) {
+	var message protocol.WSMessage
+	if conn == nil || conn.conn == nil {
+		return message, errors.New("agent ws connection is nil")
+	}
+	if conn.readTimeout > 0 {
+		_ = conn.conn.SetReadDeadline(time.Now().Add(conn.readTimeout))
+	}
+	err := websocket.JSON.Receive(conn.conn, &message)
+	if err != nil {
+		if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+			slog.Debug("agent ws receive timeout waiting for server message", "timeout", conn.readTimeout)
+		}
+		return message, err
+	}
+	slog.Debug("agent ws received message", "type", message.Type)
+	return message, nil
+}
+
+func websocketReadTimeout(requestTimeout time.Duration) time.Duration {
+	timeout := requestTimeout * 6
+	if timeout < 75*time.Second {
+		return 75 * time.Second
+	}
+	return timeout
+}
+
+func (conn *Connection) Close() error {
+	if conn == nil || conn.conn == nil {
+		return nil
+	}
+	return conn.conn.Close()
+}