mulit-agent/controller/internal/controlplane/restart.go

package controlplane

import (
	"errors"
	"fmt"
	"strings"
	"time"
)

const restartCooldown = 60 * time.Second

// AgentRestart is the safe lifecycle state rendered by the controller.  It
// contains no remote process details or credentials.
type AgentRestart struct {
	RequestID   string `json:"requestID"`
	AgentID     string `json:"agentID"`
	Status      string `json:"status"`
	RequestedAt string `json:"requestedAt"`
	ConfirmedAt string `json:"confirmedAt"`
	Failure     string `json:"failure"`
}

// RestartAgent sends an authenticated command over the existing paired WSS
// session.  The controller deliberately never starts or terminates a remote
// machine process itself.
func (c *ControlPlane) RestartAgent(agentID string) (AgentRestart, error) {
	agentID = strings.TrimSpace(agentID)
	if agentID == "" {
		return AgentRestart{}, errors.New("Agent 标识无效")
	}
	c.expireStaleAgents()
	var status string
	if err := c.db.QueryRow("SELECT status FROM managed_agents WHERE agent_id = ?", agentID).Scan(&status); err != nil {
		return AgentRestart{}, errors.New("未找到 Agent")
	}
	client := c.clientFor(agentID)
	if status != "online" || client == nil {
		c.recordAudit("remote.restart_rejected", "controller", "agent:"+agentID, "failed", "目标 Agent 当前离线，无法通过已认证连接发送重启指令。")
		return AgentRestart{}, errors.New("目标 Agent 当前离线，无法重启；请等待其重新连接")
	}
	var existing AgentRestart
	err := c.db.QueryRow("SELECT request_id, agent_id, status, requested_at, COALESCE(confirmed_at, ''), COALESCE(failure, '') FROM agent_restart_operations WHERE agent_id = ?", agentID).Scan(&existing.RequestID, &existing.AgentID, &existing.Status, &existing.RequestedAt, &existing.ConfirmedAt, &existing.Failure)
	if err == nil {
		if existing.Status == "restarting" || existing.Status == "sent" {
			return existing, nil // idempotent repeated click; do not send another command.
		}
		if existing.Status == "reconnected" {
			if requested, parseErr := time.Parse(time.RFC3339Nano, existing.RequestedAt); parseErr == nil && time.Since(requested) < restartCooldown {
				return existing, errors.New("该 Agent 刚完成重启，请稍后再试")
			}
		}
	}
	requestID, err := randomHex(16)
	if err != nil {
		return AgentRestart{}, err
	}
	now := time.Now().UTC().Format(time.RFC3339Nano)
	op := AgentRestart{RequestID: requestID, AgentID: agentID, Status: "restarting", RequestedAt: now}
	if _, err := c.db.Exec(`INSERT INTO agent_restart_operations(request_id, agent_id, status, requested_at, confirmed_at, failure)
		VALUES (?, ?, ?, ?, '', '') ON CONFLICT(agent_id) DO UPDATE SET request_id=excluded.request_id, status=excluded.status, requested_at=excluded.requested_at, confirmed_at='', failure=''`, requestID, agentID, op.Status, now); err != nil {
		return AgentRestart{}, fmt.Errorf("无法记录重启操作: %w", err)
	}
	c.recordAudit("remote.restart_requested", "controller", "agent:"+agentID, "restarting", "已通过认证控制连接请求 Agent 重启。")
	if err := client.send(map[string]any{"type": "agent_restart", "schema_version": protocolVersion, "request_id": requestID, "sent_at": now}); err != nil {
		_, _ = c.db.Exec("UPDATE agent_restart_operations SET status='failed', failure=? WHERE agent_id=? AND request_id=?", "重启指令未发送成功。", agentID, requestID)
		c.recordAudit("remote.restart_delivery", "controller", "agent:"+agentID, "failed", "重启指令未发送成功。")
		return AgentRestart{}, errors.New("重启指令未发送成功，请刷新后重试")
	}
	c.recordAudit("remote.restart_delivery", "controller", "agent:"+agentID, "sent", "重启指令已下发，等待 Agent 确认。")
	return op, nil
}

func (c *ControlPlane) applyRestartAck(agentID string, message map[string]any) {
	requestID, _ := message["request_id"].(string)
	status, _ := message["status"].(string)
	errorText, _ := message["error"].(string)
	if requestID == "" || (status != "accepted" && status != "restarting" && status != "failed") {
		return
	}
	var expected string
	if c.db.QueryRow("SELECT request_id FROM agent_restart_operations WHERE agent_id=?", agentID).Scan(&expected) != nil || expected != requestID {
		return
	}
	if status == "failed" {
		detail := limit(sanitizeDisplayText(errorText), 1000)
		if detail == "" {
			detail = "Agent 未能启动替代进程。"
		}
		_, _ = c.db.Exec("UPDATE agent_restart_operations SET status='failed', confirmed_at=?, failure=? WHERE agent_id=?", time.Now().UTC().Format(time.RFC3339Nano), detail, agentID)
		c.recordAudit("remote.restart_confirmed", "agent:"+agentID, "agent:"+agentID, "failed", detail)
		return
	}
	_, _ = c.db.Exec("UPDATE agent_restart_operations SET status='restarting', confirmed_at=?, failure='' WHERE agent_id=?", time.Now().UTC().Format(time.RFC3339Nano), agentID)
	c.recordAudit("remote.restart_confirmed", "agent:"+agentID, "agent:"+agentID, "restarting", "Agent 已确认正在安全结束任务并重启。")
}

func (c *ControlPlane) markRestartReconnected(agentID string) {
	result, err := c.db.Exec("UPDATE agent_restart_operations SET status='reconnected', confirmed_at=?, failure='' WHERE agent_id=? AND status IN ('sent', 'restarting')", time.Now().UTC().Format(time.RFC3339Nano), agentID)
	if err == nil {
		if changed, _ := result.RowsAffected(); changed == 1 {
			c.recordAudit("remote.restart_reconnected", "agent:"+agentID, "agent:"+agentID, "reconnected", "替代 Agent 已重新通过认证并恢复连接。")
		}
	}
}