[优化] go 引用调整

2026-10-03 07:06:36 +08:00 · 2026-06-06 10:26:20 +08:00
parent ee1110b752
commit 3cfefb4367
552 changed files with 1642 additions and 2185 deletions
@@ -0,0 +1,699 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"log/slog"
+	"strings"
+	"time"
+
+	"github.com/rain-kl/openflare/openflare-agent/internal/config"
+	"github.com/rain-kl/openflare/openflare-agent/internal/observability"
+	"github.com/rain-kl/openflare/openflare-agent/internal/protocol"
+	"github.com/rain-kl/openflare/openflare-agent/internal/state"
+	"github.com/rain-kl/openflare/openflare-agent/internal/wsclient"
+)
+
+type HeartbeatService interface {
+	Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error)
+	Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error)
+	SetToken(token string)
+}
+
+type SyncService interface {
+	SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error
+	SyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
+	ForceSyncOnce(ctx context.Context, target *protocol.ActiveConfigMeta) error
+	WAFIPGroupChecksums() (map[string]string, error)
+	ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error
+}
+
+type Updater interface {
+	CheckAndUpdate(ctx context.Context, repo string, options UpdateOptions) error
+}
+
+type RuntimeManager interface {
+	CheckHealth(ctx context.Context) error
+	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
+	Force   bool
+}
+
+type Runner struct {
+	Config              *config.Config
+	StateStore          *state.Store
+	ObservabilityBuffer *state.ObservabilityBufferStore
+	HeartbeatService    HeartbeatService
+	SyncService         SyncService
+	Updater             Updater
+	RuntimeManager      RuntimeManager
+	WebSocketService    WebSocketService
+
+	autoUpdate              bool
+	updateNow               bool
+	updateRepo              string
+	updateChan              string
+	updateTag               string
+	restartOpenrestyNow     bool
+	websocketUpgradeEnabled bool
+}
+
+func (r *Runner) Run(ctx context.Context) error {
+	nodeID, err := r.StateStore.EnsureNodeID()
+	if err != nil {
+		return err
+	}
+	slog.Info("agent runner started", "node_id", nodeID, "node", r.Config.NodeName, "ip", r.Config.NodeIP)
+	if r.hasAccessToken() {
+		if _, hbErr := r.performHeartbeatCycle(ctx, nodeID, true); hbErr != nil {
+			slog.Error("agent startup heartbeat failed", "error", hbErr)
+		}
+	} else if err = r.tryRegister(ctx, &nodeID); err != nil {
+		slog.Error("agent initial discovery register failed", "error", err)
+	}
+
+	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.hasAccessToken() {
+				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.hasAccessToken() {
+				if err = r.tryRegister(ctx, &nodeID); err != nil {
+					slog.Error("agent discovery register failed", "error", err)
+				}
+				continue
+			}
+			if changed, hbErr := r.performHeartbeatCycle(ctx, nodeID, false); hbErr != nil {
+				slog.Error("agent heartbeat failed", "error", hbErr)
+			} else {
+				if changed {
+					heartbeatTicker.Reset(r.Config.HeartbeatInterval.Duration())
+				}
+				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)
+	r.applyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
+	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.hasAccessToken()
+	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
+}
+
+type agentWSHandler struct {
+	runner       *Runner
+	conn         protocol.WebSocketConnection
+	nodeID       string
+	statusTicker *time.Ticker
+}
+
+func (h *agentWSHandler) OnConnect(ctx context.Context) error {
+	return h.runner.sendWebSocketStatus(ctx, h.nodeID, h.conn)
+}
+
+func (h *agentWSHandler) HandleMessage(ctx context.Context, msg wsclient.WSMessage) error {
+	var payloadBytes []byte
+	if msg.Payload != nil {
+		payloadBytes = []byte(msg.Payload)
+	}
+	protoMsg := protocol.WSMessage{
+		Type:    msg.Type,
+		Payload: payloadBytes,
+	}
+	changed, err := h.runner.handleWebSocketMessage(ctx, protoMsg, h.conn)
+	if err != nil {
+		return err
+	}
+	if changed {
+		h.statusTicker.Reset(h.runner.Config.HeartbeatInterval.Duration())
+	}
+	return nil
+}
+
+func (h *agentWSHandler) OnClose(err error) {
+	slog.Error("agent ws receive failed", "error", err)
+}
+
+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()
+
+	childCtx, cancel := context.WithCancel(ctx)
+	defer cancel()
+
+	// Start status ticker sender in background
+	go func() {
+		for {
+			select {
+			case <-childCtx.Done():
+				return
+			case <-statusTicker.C:
+				if err := r.sendWebSocketStatus(childCtx, nodeID, conn); err != nil {
+					slog.Error("agent ws send status failed", "error", err)
+					_ = conn.Close()
+					return
+				}
+			}
+		}
+	}()
+
+	wsConn, ok := conn.(*wsclient.Connection)
+	if !ok {
+		return errors.New("invalid websocket connection type")
+	}
+
+	return wsConn.RunReceiveLoop(childCtx, &agentWSHandler{
+		runner:       r,
+		conn:         conn,
+		nodeID:       nodeID,
+		statusTicker: statusTicker,
+	})
+}
+
+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.WSMessageTypeForceSyncConfig:
+		var target protocol.ActiveConfigMeta
+		if err := json.Unmarshal(message.Payload, &target); err != nil {
+			slog.Debug("agent ws force sync config decode failed", "error", err)
+			return false, nil
+		}
+		slog.Debug("agent ws force sync config received", "version", target.Version, "checksum", target.Checksum, "trigger_sync", true)
+		if err := r.SyncService.ForceSyncOnce(ctx, &target); err != nil {
+			r.recordSyncError(err)
+			slog.Error("agent ws triggered force sync failed", "version", target.Version, "error", err)
+		}
+		return false, nil
+	case protocol.WSMessageTypeWAFIPGroups:
+		var groups []protocol.WAFIPGroup
+		if err := json.Unmarshal(message.Payload, &groups); err != nil {
+			slog.Debug("agent ws waf ip groups decode failed", "error", err)
+			return false, nil
+		}
+		r.applyWAFIPGroups(ctx, groups)
+		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) hasAccessToken() bool {
+	return strings.TrimSpace(r.Config.AccessToken) != ""
+}
+
+func (r *Runner) applySettings(settings *protocol.AgentSettings) bool {
+	if settings == nil {
+		return false
+	}
+	changed := false
+	if settings.HeartbeatInterval > 0 {
+		newInterval := config.MillisecondDuration(time.Duration(settings.HeartbeatInterval) * time.Millisecond)
+		if newInterval != r.Config.HeartbeatInterval {
+			slog.Info("agent heartbeat interval updated", "from", r.Config.HeartbeatInterval, "to", newInterval)
+			r.Config.HeartbeatInterval = newInterval
+			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)
+	r.updateChan = strings.TrimSpace(settings.UpdateChannel)
+	r.updateTag = strings.TrimSpace(settings.UpdateTag)
+	r.restartOpenrestyNow = settings.RestartOpenrestyNow
+	return changed
+}
+
+func (r *Runner) tryRestartOpenresty(ctx context.Context) {
+	if !r.restartOpenrestyNow {
+		return
+	}
+	r.restartOpenrestyNow = false
+	if r.RuntimeManager == nil {
+		return
+	}
+	slog.Info("agent openresty restart requested by server")
+	if err := r.RuntimeManager.Restart(ctx); err != nil {
+		slog.Error("agent openresty restart failed", "error", err)
+		r.recordOpenrestyUnhealthy(err, false)
+		return
+	}
+	slog.Info("agent openresty restart succeeded")
+	r.recordOpenrestyHealthy()
+}
+
+func (r *Runner) tryAutoUpdate(ctx context.Context) {
+	force := r.updateNow
+	shouldCheck := r.autoUpdate || force
+	r.updateNow = false
+	r.updateTag = strings.TrimSpace(r.updateTag)
+	if !shouldCheck || r.Updater == nil || r.updateRepo == "" {
+		return
+	}
+	channel := "stable"
+	if force && r.updateChan != "" {
+		channel = r.updateChan
+	}
+	if err := r.Updater.CheckAndUpdate(ctx, r.updateRepo, UpdateOptions{
+		Channel: channel,
+		TagName: r.updateTag,
+		Force:   force,
+	}); err != nil {
+		slog.Error("agent update check failed", "error", err)
+	}
+	if force {
+		r.updateTag = ""
+		r.updateChan = ""
+	}
+}
+
+func (r *Runner) tryRegister(ctx context.Context, nodeID *string) error {
+	if strings.TrimSpace(r.Config.DiscoveryToken) == "" {
+		return errors.New("agent_token 为空且未配置 discovery_token")
+	}
+	slog.Info("agent discovery registration started")
+	response, err := r.HeartbeatService.Register(ctx, r.nodePayload(*nodeID))
+	if err != nil {
+		return err
+	}
+	if response == nil || strings.TrimSpace(response.AccessToken) == "" || strings.TrimSpace(response.NodeID) == "" {
+		return errors.New("discovery register response 缺少 node_id 或 agent_token")
+	}
+	snapshot, err := r.StateStore.Load()
+	if err != nil {
+		return err
+	}
+	snapshot.NodeID = response.NodeID
+	if err = r.StateStore.Save(snapshot); err != nil {
+		return err
+	}
+	r.Config.AccessToken = response.AccessToken
+	r.Config.DiscoveryToken = ""
+	if err = r.Config.Save(); err != nil {
+		return err
+	}
+	r.HeartbeatService.SetToken(response.AccessToken)
+	if r.WebSocketService != nil {
+		r.WebSocketService.SetToken(response.AccessToken)
+	}
+	*nodeID = response.NodeID
+	slog.Info("agent discovery registration succeeded", "node_id", response.NodeID)
+	r.refreshOpenrestyHealth(ctx)
+	payload, ackWindows := r.prepareHeartbeatPayload(*nodeID)
+	heartbeatResult, heartbeatErr := r.HeartbeatService.Heartbeat(ctx, payload)
+	if heartbeatErr != nil {
+		slog.Error("agent post-register heartbeat failed", "error", heartbeatErr)
+		return nil
+	}
+	r.ackObservabilityWindows(ackWindows)
+	if heartbeatResult == nil {
+		heartbeatResult = &protocol.HeartbeatResult{}
+	}
+	r.applySettings(heartbeatResult.AgentSettings)
+	r.applyWAFIPGroups(ctx, heartbeatResult.WAFIPGroups)
+	if err = r.SyncService.SyncOnStartup(ctx, heartbeatResult.ActiveConfig); err != nil {
+		r.recordSyncError(err)
+		slog.Error("agent post-register startup sync failed", "error", err)
+	} else {
+		slog.Debug("agent post-register startup sync completed")
+	}
+	r.tryRestartOpenresty(ctx)
+	r.tryAutoUpdate(ctx)
+	return nil
+}
+
+func (r *Runner) recordSyncError(err error) {
+	if err == nil || r.StateStore == nil {
+		return
+	}
+	snapshot, loadErr := r.StateStore.Load()
+	if loadErr != nil {
+		slog.Error("load state before recording sync error failed", "error", loadErr)
+		return
+	}
+	snapshot.LastError = err.Error()
+	slog.Warn("recording sync error into state", "error", snapshot.LastError)
+	if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
+		slog.Error("save state after sync error failed", "error", saveErr)
+	}
+}
+
+func (r *Runner) refreshOpenrestyHealth(ctx context.Context) {
+	if r.RuntimeManager == nil || r.StateStore == nil {
+		return
+	}
+	if err := r.RuntimeManager.CheckHealth(ctx); err != nil {
+		if strings.Contains(err.Error(), "openresty config not exists") {
+			return
+		}
+		r.recordOpenrestyUnhealthy(err, true)
+		return
+	}
+	r.recordOpenrestyHealthy()
+}
+
+func (r *Runner) recordOpenrestyHealthy() {
+	if r.StateStore == nil {
+		return
+	}
+	snapshot, err := r.StateStore.Load()
+	if err != nil {
+		slog.Error("load state before recording openresty health failed", "error", err)
+		return
+	}
+	if snapshot.OpenrestyStatus == protocol.OpenrestyStatusHealthy && strings.TrimSpace(snapshot.OpenrestyMessage) == "" {
+		return
+	}
+	snapshot.OpenrestyStatus = protocol.OpenrestyStatusHealthy
+	snapshot.OpenrestyMessage = ""
+	if err = r.StateStore.Save(snapshot); err != nil {
+		slog.Error("save state after recording openresty health failed", "error", err)
+	}
+}
+
+func (r *Runner) recordOpenrestyUnhealthy(err error, fallbackOnly bool) {
+	if err == nil || r.StateStore == nil {
+		return
+	}
+	snapshot, loadErr := r.StateStore.Load()
+	if loadErr != nil {
+		slog.Error("load state before recording openresty error failed", "error", loadErr)
+		return
+	}
+	message := strings.TrimSpace(err.Error())
+	if !fallbackOnly || strings.TrimSpace(snapshot.OpenrestyMessage) == "" {
+		snapshot.OpenrestyMessage = message
+	}
+	snapshot.OpenrestyStatus = protocol.OpenrestyStatusUnhealthy
+	if saveErr := r.StateStore.Save(snapshot); saveErr != nil {
+		slog.Error("save state after recording openresty error failed", "error", saveErr)
+	}
+}
+
+func (r *Runner) nodePayload(nodeID string) protocol.NodePayload {
+	snapshot, _ := r.StateStore.Load()
+	openrestyStatus := strings.TrimSpace(snapshot.OpenrestyStatus)
+	if openrestyStatus == "" {
+		openrestyStatus = protocol.OpenrestyStatusUnknown
+	}
+	profile := observability.BuildProfile(r.Config, r.StateStore)
+	managedOpenRestyMetrics := observability.CollectManagedOpenRestyMetrics(r.Config)
+	trafficReport, accessLogs, fallbackMetrics := observability.BuildTrafficObservability(r.Config, r.StateStore, managedOpenRestyMetrics)
+	if managedOpenRestyMetrics == nil {
+		managedOpenRestyMetrics = fallbackMetrics
+	}
+	metricSnapshot := observability.BuildSnapshot(r.Config, r.StateStore)
+	openrestyObservation := observability.BuildOpenrestyObservation(managedOpenRestyMetrics)
+	healthEvents := observability.BuildHealthEvents(snapshot)
+	payload := protocol.NodePayload{
+		NodeID:               nodeID,
+		Name:                 r.Config.NodeName,
+		IP:                   r.Config.NodeIP,
+		Version:              r.Config.Version,
+		ExtVersion:           r.Config.ExtVersion,
+		CurrentVersion:       snapshot.CurrentVersion,
+		LastError:            snapshot.LastError,
+		OpenrestyStatus:      openrestyStatus,
+		OpenrestyMessage:     snapshot.OpenrestyMessage,
+		Profile:              profile,
+		Snapshot:             metricSnapshot,
+		OpenrestyObservation: openrestyObservation,
+		TrafficReport:        trafficReport,
+		AccessLogs:           accessLogs,
+		HealthEvents:         healthEvents,
+	}
+	if r.SyncService != nil {
+		checksums, err := r.SyncService.WAFIPGroupChecksums()
+		if err != nil {
+			slog.Debug("load local waf ip group checksums failed", "error", err)
+		} else if len(checksums) > 0 {
+			payload.WAFIPGroupChecksums = checksums
+		}
+	}
+	return payload
+}
+
+func (r *Runner) applyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) {
+	if len(groups) == 0 || r.SyncService == nil {
+		return
+	}
+	if err := r.SyncService.ApplyWAFIPGroups(ctx, groups); err != nil {
+		r.recordSyncError(err)
+		slog.Error("agent apply waf ip groups failed", "error", err)
+	}
+}
+
+func (r *Runner) prepareHeartbeatPayload(nodeID string) (protocol.NodePayload, []int64) {
+	payload := r.nodePayload(nodeID)
+	if r.ObservabilityBuffer == nil || (payload.Snapshot == nil && payload.TrafficReport == nil && len(payload.AccessLogs) == 0) {
+		return payload, nil
+	}
+	now := time.Now().UTC()
+	retainAfterUnix := now.Add(-time.Duration(r.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
+	windowStartedAtUnix := state.ObservabilityWindowStartedAt(payload.Snapshot, payload.OpenrestyObservation, payload.TrafficReport)
+	if windowStartedAtUnix <= 0 {
+		return payload, nil
+	}
+
+	record := state.ObservabilityBufferRecord{
+		WindowStartedAtUnix:  windowStartedAtUnix,
+		Snapshot:             payload.Snapshot,
+		OpenrestyObservation: payload.OpenrestyObservation,
+		TrafficReport:        payload.TrafficReport,
+		AccessLogs:           payload.AccessLogs,
+		QueuedAtUnix:         now.Unix(),
+	}
+	if err := r.ObservabilityBuffer.Upsert(record, retainAfterUnix); err != nil {
+		slog.Error("upsert observability buffer failed", "error", err)
+		return payload, nil
+	}
+
+	records, err := r.ObservabilityBuffer.Replayable(windowStartedAtUnix, retainAfterUnix)
+	if err != nil {
+		slog.Error("load replayable observability buffer failed", "error", err)
+		return payload, []int64{windowStartedAtUnix}
+	}
+
+	ackWindows := make([]int64, 0, len(records)+1)
+	buffered := make([]protocol.BufferedObservabilityRecord, 0, len(records))
+	for _, item := range records {
+		if item.WindowStartedAtUnix <= 0 {
+			continue
+		}
+		buffered = append(buffered, protocol.BufferedObservabilityRecord{
+			WindowStartedAtUnix:  item.WindowStartedAtUnix,
+			Snapshot:             item.Snapshot,
+			OpenrestyObservation: item.OpenrestyObservation,
+			TrafficReport:        item.TrafficReport,
+			AccessLogs:           item.AccessLogs,
+		})
+		ackWindows = append(ackWindows, item.WindowStartedAtUnix)
+	}
+	payload.BufferedObservability = buffered
+	ackWindows = append(ackWindows, windowStartedAtUnix)
+	return payload, ackWindows
+}
+
+func (r *Runner) ackObservabilityWindows(windowStartedAtUnix []int64) {
+	if r.ObservabilityBuffer == nil || len(windowStartedAtUnix) == 0 {
+		return
+	}
+	retainAfterUnix := time.Now().UTC().Add(-time.Duration(r.Config.ObservabilityReplayMinutes) * time.Minute).Unix()
+	if err := r.ObservabilityBuffer.Ack(windowStartedAtUnix, retainAfterUnix); err != nil {
+		slog.Error("ack observability buffer failed", "error", err)
+	}
+}
@@ -0,0 +1,640 @@
+package agent
+
+import (
+	"context"
+	"encoding/json"
+	"errors"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/rain-kl/openflare/openflare-agent/internal/config"
+	"github.com/rain-kl/openflare/openflare-agent/internal/protocol"
+	"github.com/rain-kl/openflare/openflare-agent/internal/state"
+)
+
+type fakeHeartbeatService struct {
+	mu                sync.Mutex
+	registerCalls     int
+	heartbeatCalls    int
+	registerErr       error
+	registerResp      *protocol.RegisterNodeResponse
+	heartbeatErrs     []error
+	heartbeatResults  []*protocol.HeartbeatResult
+	heartbeatPayloads []protocol.NodePayload
+	onHeartbeat       func(int)
+	lastToken         string
+}
+
+func (f *fakeHeartbeatService) Register(ctx context.Context, payload protocol.NodePayload) (*protocol.RegisterNodeResponse, error) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.registerCalls++
+	return f.registerResp, f.registerErr
+}
+
+func (f *fakeHeartbeatService) Heartbeat(ctx context.Context, payload protocol.NodePayload) (*protocol.HeartbeatResult, error) {
+	f.mu.Lock()
+	f.heartbeatCalls++
+	callIndex := f.heartbeatCalls
+	f.heartbeatPayloads = append(f.heartbeatPayloads, payload)
+	var err error
+	if len(f.heartbeatErrs) >= callIndex {
+		err = f.heartbeatErrs[callIndex-1]
+	}
+	var result *protocol.HeartbeatResult
+	if len(f.heartbeatResults) >= callIndex {
+		result = f.heartbeatResults[callIndex-1]
+	}
+	onHeartbeat := f.onHeartbeat
+	f.mu.Unlock()
+	if onHeartbeat != nil {
+		onHeartbeat(callIndex)
+	}
+	return result, err
+}
+
+func (f *fakeHeartbeatService) SetToken(token string) {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.lastToken = token
+}
+
+type fakeSyncService struct {
+	mu             sync.Mutex
+	startupErr     error
+	syncOnceErr    error
+	startupCalls   int
+	syncOnceCalls  int
+	lastTarget     *protocol.ActiveConfigMeta
+	onSyncOnceCall func(int)
+	wafChecksums   map[string]string
+	wafGroups      []protocol.WAFIPGroup
+}
+
+type fakeRuntimeManager struct {
+	mu                   sync.Mutex
+	healthErr            error
+	restartErr           error
+	restartCalls         int
+	clearHealthOnRestart bool
+}
+
+func (f *fakeRuntimeManager) CheckHealth(ctx context.Context) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	return f.healthErr
+}
+
+func (f *fakeRuntimeManager) Restart(ctx context.Context) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.restartCalls++
+	if f.clearHealthOnRestart && f.restartErr == nil {
+		f.healthErr = nil
+	}
+	return f.restartErr
+}
+
+func (f *fakeSyncService) SyncOnStartup(ctx context.Context, target *protocol.ActiveConfigMeta) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.startupCalls++
+	return f.startupErr
+}
+
+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()
+	if callback != nil {
+		callback(callIndex)
+	}
+	return f.syncOnceErr
+}
+
+func (f *fakeSyncService) ForceSyncOnce(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()
+	if callback != nil {
+		callback(callIndex)
+	}
+	return f.syncOnceErr
+}
+
+func (f *fakeSyncService) WAFIPGroupChecksums() (map[string]string, error) {
+	if f.wafChecksums == nil {
+		return map[string]string{}, nil
+	}
+	return f.wafChecksums, nil
+}
+
+func (f *fakeSyncService) ApplyWAFIPGroups(ctx context.Context, groups []protocol.WAFIPGroup) error {
+	f.mu.Lock()
+	defer f.mu.Unlock()
+	f.wafGroups = append(f.wafGroups, groups...)
+	return nil
+}
+
+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()
+
+	stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
+	heartbeatService := &fakeHeartbeatService{
+		heartbeatResults: []*protocol.HeartbeatResult{{}},
+		onHeartbeat: func(callCount int) {
+			if callCount >= 2 {
+				cancel()
+			}
+		},
+	}
+	syncService := &fakeSyncService{
+		startupErr: errors.New("当前没有激活版本，保持当前 OpenResty 配置"),
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			AccessToken:       "agent-token",
+			NodeName:          "edge-01",
+			NodeIP:            "10.0.0.8",
+			Version:           config.Version,
+			ExtVersion:        "1.27.1.2",
+			HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
+		},
+		StateStore:       stateStore,
+		HeartbeatService: heartbeatService,
+		SyncService:      syncService,
+	}
+
+	err := runner.Run(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context cancellation, got %v", err)
+	}
+	if heartbeatService.registerCalls != 0 {
+		t.Fatalf("expected no discovery register call, got %d", heartbeatService.registerCalls)
+	}
+	if heartbeatService.heartbeatCalls < 2 {
+		t.Fatalf("expected heartbeat loop to continue, got %d heartbeat calls", heartbeatService.heartbeatCalls)
+	}
+	snapshot, loadErr := stateStore.Load()
+	if loadErr != nil {
+		t.Fatalf("failed to load state: %v", loadErr)
+	}
+	if snapshot.LastError != "当前没有激活版本，保持当前 OpenResty 配置" {
+		t.Fatalf("expected startup sync error to be recorded, got %q", snapshot.LastError)
+	}
+}
+
+func TestRunnerDoesNotExitOnHeartbeatOrSyncError(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
+	heartbeatService := &fakeHeartbeatService{
+		registerErr:   errors.New("register timeout"),
+		heartbeatErrs: []error{errors.New("heartbeat timeout")},
+		heartbeatResults: []*protocol.HeartbeatResult{
+			{},
+		},
+	}
+	syncService := &fakeSyncService{
+		syncOnceErr: errors.New("openresty reload failed"),
+		onSyncOnceCall: func(callCount int) {
+			if callCount >= 1 {
+				cancel()
+			}
+		},
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			AccessToken:       "agent-token",
+			NodeName:          "edge-01",
+			NodeIP:            "10.0.0.8",
+			Version:           config.Version,
+			ExtVersion:        "1.27.1.2",
+			HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
+		},
+		StateStore:       stateStore,
+		HeartbeatService: heartbeatService,
+		SyncService:      syncService,
+	}
+
+	err := runner.Run(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context cancellation, got %v", err)
+	}
+	if heartbeatService.registerCalls != 0 {
+		t.Fatalf("expected no register attempt, got %d", heartbeatService.registerCalls)
+	}
+	if syncService.syncOnceCalls == 0 {
+		t.Fatal("expected sync loop to continue after heartbeat/register errors")
+	}
+	snapshot, loadErr := stateStore.Load()
+	if loadErr != nil {
+		t.Fatalf("failed to load state: %v", loadErr)
+	}
+	if snapshot.LastError != "openresty reload failed" {
+		t.Fatalf("expected sync error to be recorded, got %q", snapshot.LastError)
+	}
+}
+
+func TestRunnerReportsOpenrestyHealthAndExecutesRestart(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
+	if err := stateStore.Save(&state.Snapshot{
+		OpenrestyStatus:  protocol.OpenrestyStatusUnhealthy,
+		OpenrestyMessage: "docker run openresty failed: bind 80 already allocated",
+	}); err != nil {
+		t.Fatalf("failed to seed state: %v", err)
+	}
+	heartbeatService := &fakeHeartbeatService{
+		heartbeatResults: []*protocol.HeartbeatResult{{
+			AgentSettings: &protocol.AgentSettings{RestartOpenrestyNow: true},
+		}},
+		onHeartbeat: func(callCount int) {
+			if callCount >= 1 {
+				cancel()
+			}
+		},
+	}
+	runtimeManager := &fakeRuntimeManager{
+		healthErr:            errors.New("docker openresty container is not running"),
+		clearHealthOnRestart: true,
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			AccessToken:       "agent-token",
+			NodeName:          "edge-01",
+			NodeIP:            "10.0.0.8",
+			Version:           config.Version,
+			ExtVersion:        "1.27.1.2",
+			HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
+		},
+		StateStore:       stateStore,
+		HeartbeatService: heartbeatService,
+		SyncService:      &fakeSyncService{},
+		RuntimeManager:   runtimeManager,
+	}
+
+	err := runner.Run(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context cancellation, got %v", err)
+	}
+	if len(heartbeatService.heartbeatPayloads) == 0 {
+		t.Fatal("expected at least one heartbeat payload")
+	}
+	payload := heartbeatService.heartbeatPayloads[0]
+	if payload.OpenrestyStatus != protocol.OpenrestyStatusUnhealthy {
+		t.Fatalf("expected unhealthy openresty status in heartbeat payload, got %q", payload.OpenrestyStatus)
+	}
+	if payload.OpenrestyMessage != "docker run openresty failed: bind 80 already allocated" {
+		t.Fatalf("unexpected openresty message: %q", payload.OpenrestyMessage)
+	}
+	if runtimeManager.restartCalls != 1 {
+		t.Fatalf("expected one openresty restart attempt, got %d", runtimeManager.restartCalls)
+	}
+	snapshot, loadErr := stateStore.Load()
+	if loadErr != nil {
+		t.Fatalf("failed to load state: %v", loadErr)
+	}
+	if snapshot.OpenrestyStatus != protocol.OpenrestyStatusHealthy || snapshot.OpenrestyMessage != "" {
+		t.Fatal("expected restart success to mark openresty healthy")
+	}
+}
+
+func TestRunnerHeartbeatPayloadIncludesObservabilityExtensions(t *testing.T) {
+	tempDir := t.TempDir()
+	stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
+	if err := stateStore.Save(&state.Snapshot{
+		NodeID:           "node-observe",
+		CurrentVersion:   "20260314-001",
+		LastError:        "sync failed",
+		OpenrestyStatus:  protocol.OpenrestyStatusUnhealthy,
+		OpenrestyMessage: "reload failed",
+	}); err != nil {
+		t.Fatalf("failed to seed state: %v", err)
+	}
+
+	runner := &Runner{
+		Config: &config.Config{
+			NodeName:          "edge-observe-1",
+			NodeIP:            "10.0.0.51",
+			Version:           config.Version,
+			ExtVersion:        "1.27.1.2",
+			DataDir:           tempDir,
+			RouteConfigPath:   filepath.Join(tempDir, "conf.d", "openflare_routes.conf"),
+			AccessLogPath:     filepath.Join(tempDir, "var", "log", "openflare", "access.log"),
+			HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
+		},
+		StateStore: stateStore,
+	}
+	if err := os.MkdirAll(filepath.Dir(runner.Config.AccessLogPath), 0o755); err != nil {
+		t.Fatalf("failed to prepare access log dir: %v", err)
+	}
+	if err := os.WriteFile(
+		runner.Config.AccessLogPath,
+		[]byte("{\"ts\":\""+time.Now().UTC().Format(time.RFC3339)+"\",\"host\":\"edge.example.com\",\"path\":\"/\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
+		0o644,
+	); err != nil {
+		t.Fatalf("failed to prepare access log: %v", err)
+	}
+
+	firstPayload := runner.nodePayload("node-observe")
+	if firstPayload.Profile == nil {
+		t.Fatal("expected first heartbeat payload to include system profile")
+	}
+	if firstPayload.Snapshot == nil {
+		t.Fatal("expected first heartbeat payload to include metric snapshot")
+	}
+	if firstPayload.TrafficReport == nil || firstPayload.TrafficReport.RequestCount != 1 {
+		t.Fatalf("expected first heartbeat payload to include traffic report, got %+v", firstPayload.TrafficReport)
+	}
+	if len(firstPayload.AccessLogs) != 1 || firstPayload.AccessLogs[0].Path != "/" {
+		t.Fatalf("expected first heartbeat payload to include access logs, got %+v", firstPayload.AccessLogs)
+	}
+	if len(firstPayload.HealthEvents) != 2 {
+		t.Fatalf("expected health events for openresty and sync error, got %+v", firstPayload.HealthEvents)
+	}
+
+	secondPayload := runner.nodePayload("node-observe")
+	if secondPayload.Profile != nil {
+		t.Fatal("expected unchanged profile to be omitted on subsequent heartbeat")
+	}
+	if secondPayload.Snapshot == nil {
+		t.Fatal("expected metric snapshot to continue reporting on subsequent heartbeat")
+	}
+	if secondPayload.TrafficReport != nil {
+		t.Fatalf("expected unchanged traffic window to be omitted on subsequent heartbeat, got %+v", secondPayload.TrafficReport)
+	}
+	if len(secondPayload.AccessLogs) != 0 {
+		t.Fatalf("expected unchanged access log delta to be omitted on subsequent heartbeat, got %+v", secondPayload.AccessLogs)
+	}
+}
+
+func TestRunnerReplaysBufferedObservabilityAfterHeartbeatRecovery(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	tempDir := t.TempDir()
+	stateStore := state.NewStore(filepath.Join(tempDir, "state.json"))
+	bufferStore := state.NewObservabilityBufferStore(filepath.Join(tempDir, "observability-buffer.json"))
+	nowUnix := time.Now().UTC().Unix()
+	bufferWindow := nowUnix - (nowUnix % 60) - 60
+	if err := bufferStore.Upsert(state.ObservabilityBufferRecord{
+		WindowStartedAtUnix: bufferWindow,
+		Snapshot:            &protocol.NodeMetricSnapshot{CapturedAtUnix: bufferWindow + 5, CPUUsagePercent: 30},
+		TrafficReport:       &protocol.NodeTrafficReport{WindowStartedAtUnix: bufferWindow, WindowEndedAtUnix: bufferWindow + 60, RequestCount: 8},
+		QueuedAtUnix:        bufferWindow + 60,
+	}, 0); err != nil {
+		t.Fatalf("failed to seed observability buffer: %v", err)
+	}
+	heartbeatService := &fakeHeartbeatService{
+		heartbeatErrs:    []error{errors.New("server offline"), nil},
+		heartbeatResults: []*protocol.HeartbeatResult{{}, {}},
+		onHeartbeat: func(callCount int) {
+			if callCount >= 2 {
+				cancel()
+			}
+		},
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			AccessToken:                "agent-token",
+			NodeName:                   "edge-buffer-01",
+			NodeIP:                     "10.0.0.52",
+			Version:                    config.Version,
+			ExtVersion:                 "1.27.1.2",
+			DataDir:                    tempDir,
+			RouteConfigPath:            filepath.Join(tempDir, "conf.d", "openflare_routes.conf"),
+			HeartbeatInterval:          config.MillisecondDuration(10 * time.Millisecond),
+			ObservabilityReplayMinutes: 15,
+		},
+		StateStore:          stateStore,
+		ObservabilityBuffer: bufferStore,
+		HeartbeatService:    heartbeatService,
+		SyncService:         &fakeSyncService{},
+	}
+	if err := os.MkdirAll(filepath.Dir(runner.Config.RouteConfigPath), 0o755); err != nil {
+		t.Fatalf("failed to prepare route config dir: %v", err)
+	}
+	if err := os.WriteFile(
+		filepath.Join(filepath.Dir(runner.Config.RouteConfigPath), "openflare_access.log"),
+		[]byte("{\"ts\":\""+time.Now().UTC().Format(time.RFC3339)+"\",\"host\":\"edge.example.com\",\"path\":\"/\",\"remote_addr\":\"10.0.0.8\",\"status\":200}\n"),
+		0o644,
+	); err != nil {
+		t.Fatalf("failed to prepare access log: %v", err)
+	}
+
+	runErr := runner.Run(ctx)
+	if runErr != context.Canceled {
+		t.Fatalf("expected run to stop by context cancellation, got %v", runErr)
+	}
+	if len(heartbeatService.heartbeatPayloads) != 2 {
+		t.Fatalf("expected two heartbeat payloads, got %d", len(heartbeatService.heartbeatPayloads))
+	}
+	secondPayload := heartbeatService.heartbeatPayloads[1]
+	if len(secondPayload.BufferedObservability) != 1 {
+		t.Fatalf("expected second heartbeat to replay one buffered observation, got %+v", secondPayload.BufferedObservability)
+	}
+	if len(secondPayload.BufferedObservability[0].AccessLogs) != 0 {
+		t.Fatalf("expected seeded buffered observation to keep empty access logs, got %+v", secondPayload.BufferedObservability[0].AccessLogs)
+	}
+
+	replayable, err := bufferStore.Replayable(0, 0)
+	if err != nil {
+		t.Fatalf("Replayable after recovery failed: %v", err)
+	}
+	if len(replayable) != 0 {
+		t.Fatalf("expected buffer to be acked after successful heartbeat, got %+v", replayable)
+	}
+}
+
+func TestRunnerDiscoveryRegisterUpdatesTokenAndNodeID(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	stateStore := state.NewStore(filepath.Join(t.TempDir(), "state.json"))
+	heartbeatService := &fakeHeartbeatService{
+		registerResp: &protocol.RegisterNodeResponse{
+			NodeID:      "node-server-assigned",
+			AccessToken: "agent-token-issued",
+			Name:        "edge-01",
+		},
+		heartbeatResults: []*protocol.HeartbeatResult{{}},
+		onHeartbeat: func(callCount int) {
+			if callCount >= 1 {
+				cancel()
+			}
+		},
+	}
+	syncService := &fakeSyncService{}
+	configPath := filepath.Join(t.TempDir(), "agent.json")
+	if err := os.WriteFile(configPath, []byte(`{"server_url":"http://127.0.0.1:3000","discovery_token":"discovery-token","node_name":"edge-01","node_ip":"10.0.0.8"}`), 0o644); err != nil {
+		t.Fatalf("failed to seed config file: %v", err)
+	}
+	cfg, err := config.Load(configPath)
+	if err != nil {
+		t.Fatalf("failed to load config: %v", err)
+	}
+	runner := &Runner{
+		Config: &config.Config{
+			ServerURL:         cfg.ServerURL,
+			DiscoveryToken:    cfg.DiscoveryToken,
+			NodeName:          cfg.NodeName,
+			NodeIP:            cfg.NodeIP,
+			Version:           config.Version,
+			ExtVersion:        "1.27.1.2",
+			HeartbeatInterval: config.MillisecondDuration(10 * time.Millisecond),
+		},
+		StateStore:       stateStore,
+		HeartbeatService: heartbeatService,
+		SyncService:      syncService,
+	}
+	runner.Config = cfg
+	runner.Config.Version = config.Version
+	runner.Config.ExtVersion = "1.27.1.2"
+	runner.Config.HeartbeatInterval = config.MillisecondDuration(10 * time.Millisecond)
+
+	err = runner.Run(ctx)
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("expected context cancellation, got %v", err)
+	}
+	if heartbeatService.registerCalls == 0 {
+		t.Fatal("expected discovery register to be attempted")
+	}
+	if heartbeatService.lastToken != "agent-token-issued" {
+		t.Fatalf("expected client token to be updated, got %q", heartbeatService.lastToken)
+	}
+	snapshot, loadErr := stateStore.Load()
+	if loadErr != nil {
+		t.Fatalf("failed to load state: %v", loadErr)
+	}
+	if snapshot.NodeID != "node-server-assigned" {
+		t.Fatalf("expected node id to be replaced, got %q", snapshot.NodeID)
+	}
+	if runner.Config.AccessToken != "agent-token-issued" || runner.Config.DiscoveryToken != "" {
+		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)
+	}
+}