aiops/backend/controller/agent_cli.go

// agent_cli.go - Agent CLI WebSocket 控制器 // Logan.Li
// 驱动 pkg/agentcli 引擎：关联主机（model.Server 凭据）→ 构建引擎 → ReAct 循环，
// 通过 WebSocket 实时推送思考/工具调用/输出/结果，高危命令经前端确认，过程持久化到 DB。
package controller

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"sync"
	"time"

	"github.com/gin-gonic/gin"
	"github.com/gorilla/websocket"
	"gorm.io/gorm"

	"ezops/middleware"
	"ezops/model"
	"ezops/pkg/agentcli"
	"ezops/pkg/agentcli/core"
	"ezops/pkg/agentcli/llm"
	"ezops/pkg/agentcli/safety"
	"ezops/pkg/agentcli/tools"
	"ezops/pkg/remoteexec"
)

// AgentCLIHandler Agent CLI 控制器 // Logan.Li
type AgentCLIHandler struct {
	db *gorm.DB
}

// NewAgentCLIHandler 创建控制器 // Logan.Li
func NewAgentCLIHandler(db *gorm.DB) *AgentCLIHandler { return &AgentCLIHandler{db: db} }

// wsEvent WebSocket 事件信封（后端→前端） // Logan.Li
type wsEvent struct {
	Type      string                 `json:"type"`
	Round     int                    `json:"round,omitempty"`
	SessionID uint                   `json:"session_id,omitempty"`
	Thought   string                 `json:"thought,omitempty"`
	ToolName  string                 `json:"tool_name,omitempty"`
	ToolArgs  map[string]interface{} `json:"tool_args,omitempty"`
	Risk      string                 `json:"risk,omitempty"`
	Output    string                 `json:"output,omitempty"`
	Success   bool                   `json:"success,omitempty"`
	Duration  string                 `json:"duration,omitempty"`
	Tokens    int                    `json:"tokens,omitempty"`
	Answer    string                 `json:"answer,omitempty"`
	Message   string                 `json:"message,omitempty"`
	ConfirmID string                 `json:"confirm_id,omitempty"`
	Command   string                 `json:"command,omitempty"`
}

// wsControl 前端→后端控制消息 // Logan.Li
type wsControl struct {
	Type      string `json:"type"`                 // start/confirm/deny/abort
	Query     string `json:"query,omitempty"`      // start 时任务描述
	ConfirmID string `json:"confirm_id,omitempty"` // confirm/deny 时
}

// agentConn 单个 WS 连接的运行态 // Logan.Li
type agentConn struct {
	conn     *websocket.Conn
	db       *gorm.DB
	userID   uint
	server   model.Server
	session  *model.AgentCLISession
	cancel   context.CancelFunc
	writeMu  sync.Mutex
	confirms map[string]chan bool
	cfMu     sync.Mutex
}

// Run WebSocket 入口：GET /ws/agent-cli?id=<serverID>&token=<jwt> // Logan.Li
func (h *AgentCLIHandler) Run(c *gin.Context) {
	serverID := c.Query("id")
	if serverID == "" {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "缺少主机 id"})
		return
	}
	var server model.Server
	if err := h.db.First(&server, serverID).Error; err != nil {
		c.JSON(http.StatusNotFound, model.Response{Code: 404, Message: "主机不存在"})
		return
	}
	sshCfg, err := buildSSHConfig(h.db, server)
	if err != nil {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: err.Error()})
		return
	}
	// LLM 配置
	var llmCfg model.LLMConfig
	if err := h.db.First(&llmCfg).Error; err != nil || !llmCfg.Enabled {
		c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "LLM 服务未配置或未启用，请先在「大模型设置」中配置"})
		return
	}

	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	ac := &agentConn{
		conn: conn, db: h.db,
		userID: middleware.GetUserID(c), server: server,
		confirms: make(map[string]chan bool),
	}
	// 推送主机上下文
	ac.write(wsEvent{Type: "session", SessionID: 0,
		Message: fmt.Sprintf("已连接主机 %s（%s），LLM: %s。请描述要执行的任务。", server.Hostname, server.IP, llmCfg.Model)})

	// 读循环：处理前端控制消息
	for {
		_, data, err := conn.ReadMessage()
		if err != nil {
			ac.abort()
			return
		}
		var ctrl wsControl
		if json.Unmarshal(data, &ctrl) != nil {
			continue
		}
		switch ctrl.Type {
		case "start":
			if ctrl.Query == "" {
				ac.write(wsEvent{Type: "error", Message: "任务描述不能为空"})
				continue
			}
			ac.startTask(sshCfg, llmCfg, ctrl.Query)
		case "confirm":
			ac.resolveConfirm(ctrl.ConfirmID, true)
		case "deny":
			ac.resolveConfirm(ctrl.ConfirmID, false)
		case "abort":
			ac.abort()
		}
	}
}

// startTask 启动一次任务（新会话 + 新引擎） // Logan.Li
func (ac *agentConn) startTask(sshCfg remoteexec.SSHConfig, llmCfg model.LLMConfig, query string) {
	ac.abort() // 取消上一次任务

	// 创建 DB 会话
	title := query
	if r := []rune(title); len(r) > 60 {
		title = string(r[:60]) + "..."
	}
	sess := &model.AgentCLISession{
		UserID: ac.userID, ServerID: ac.server.ID,
		ServerName: ac.server.Hostname, ServerHost: ac.server.IP,
		Title: title, Status: "running",
		Provider: llmCfg.Provider, Model: llmCfg.Model,
	}
	if err := ac.db.Create(sess).Error; err != nil {
		ac.write(wsEvent{Type: "error", Message: "创建会话失败: " + err.Error()})
		return
	}
	ac.session = sess

	// 存用户消息
	ac.persistMsg(sess.ID, 0, "user", "user", query, "", "", "", "", false, "", 0)

	// 构建引擎
	engine, err := agentcli.NewEngine(agentcli.EngineConfig{
		LLM: llm.Config{
			Protocol: resolveProtocol(llmCfg.Provider),
			Endpoint: resolveEndpoint(llmCfg.Provider, llmCfg.Endpoint),
			APIKey:   llmCfg.APIKey, Model: llmCfg.Model,
			Temperature: llmCfg.Temperature, MaxTokens: llmCfg.MaxTokens,
		},
		SSH:        sshCfg,
		SafetyMode: safety.ModeBalanced,
	})
	if err != nil {
		ac.finish(sess, "error", "引擎初始化失败: "+err.Error(), 0, time.Duration(0))
		return
	}
	engine.SetCallback(ac.onEvent)
	engine.SetConfirmFn(ac.confirmFn)

	ctx, cancel := context.WithCancel(context.Background())
	ac.cancel = cancel
	serverInfo := fmt.Sprintf("主机名: %s | IP: %s:%d | 用户: %s | OS: %s", ac.server.Hostname, ac.server.IP, ac.server.Port, ac.server.Username, ac.server.OS)

	ac.write(wsEvent{Type: "session", SessionID: sess.ID, Message: "任务已开始"})

	go func() {
		defer engine.Close()
		start := time.Now()
		_, err := engine.Run(ctx, query, serverInfo)
		dur := time.Since(start)
		if ctx.Err() != nil {
			ac.finish(sess, "aborted", "任务已中断", ac.lastRound(), dur)
		} else if err != nil {
			ac.finish(sess, "error", err.Error(), ac.lastRound(), dur)
		} else {
			ac.finish(sess, "done", "", ac.lastRound(), dur)
		}
	}()
}

// onEvent 引擎事件回调：推 WS + 持久化 // Logan.Li
func (ac *agentConn) onEvent(e core.Event) {
	ev := wsEvent{
		Type: string(e.Type), Round: e.Round, Thought: e.Thought,
		ToolName: e.ToolName, ToolArgs: e.ToolArgs, Risk: e.Risk,
		Output: e.Output, Success: e.Success, Duration: e.Duration,
		Tokens: e.Tokens, Answer: e.Answer, Message: e.Message,
	}
	if ac.session != nil {
		ev.SessionID = ac.session.ID
	}
	ac.write(ev)

	// 持久化（progress 太频繁不入库）
	if ac.session == nil {
		return
	}
	sid := ac.session.ID
	switch e.Type {
	case core.EventThought:
		ac.persistMsg(sid, e.Round, "thought", "assistant", e.Thought, "", "", "", "", false, "", 0)
	case core.EventToolCall:
		argsJSON, _ := json.Marshal(e.ToolArgs)
		ac.persistMsg(sid, e.Round, "tool_call", "tool", "", e.ToolName, string(argsJSON), "", e.Risk, false, "", 0)
	case core.EventToolResult:
		ac.persistMsg(sid, e.Round, "tool_result", "tool", e.Output, e.ToolName, "", "", "", e.Success, e.Duration, 0)
	case core.EventDone:
		ac.persistMsg(sid, e.Round, "done", "assistant", e.Answer, "", "", "", "", true, "", e.Tokens)
	case core.EventError:
		ac.persistMsg(sid, e.Round, "error", "assistant", e.Message, "", "", "", "", false, "", 0)
	}
}

// confirmFn 高危命令确认回调：推 confirm 事件并阻塞等前端回传 // Logan.Li
func (ac *agentConn) confirmFn(toolName string, args map[string]interface{}, risk tools.RiskLevel) (bool, error) {
	confirmID := fmt.Sprintf("cf-%d", time.Now().UnixNano())
	cmd, _ := args["command"].(string)
	ac.write(wsEvent{Type: "confirm", ConfirmID: confirmID, ToolName: toolName, Command: cmd, Risk: risk.Display(), ToolArgs: args})

	ch := make(chan bool, 1)
	ac.cfMu.Lock()
	ac.confirms[confirmID] = ch
	ac.cfMu.Unlock()

	select {
	case approved := <-ch:
		return approved, nil
	case <-time.After(5 * time.Minute):
		ac.delConfirm(confirmID)
		return false, fmt.Errorf("确认超时")
	}
}

// resolveConfirm 前端回传确认结果 // Logan.Li
func (ac *agentConn) resolveConfirm(confirmID string, approved bool) {
	ac.cfMu.Lock()
	ch, ok := ac.confirms[confirmID]
	if ok {
		delete(ac.confirms, confirmID)
	}
	ac.cfMu.Unlock()
	if ok {
		ch <- approved
	}
}

func (ac *agentConn) delConfirm(id string) {
	ac.cfMu.Lock()
	delete(ac.confirms, id)
	ac.cfMu.Unlock()
}

// abort 中断当前任务 // Logan.Li
func (ac *agentConn) abort() {
	if ac.cancel != nil {
		ac.cancel()
		ac.cancel = nil
	}
}

// lastRound 估算已完成轮次（从 DB 取该会话最大 round） // Logan.Li
func (ac *agentConn) lastRound() int {
	if ac.session == nil {
		return 0
	}
	var maxRound int
	ac.db.Model(&model.AgentCLIMessage{}).Where("session_id = ?", ac.session.ID).
		Select("COALESCE(MAX(round),0)").Scan(&maxRound)
	return maxRound
}

// finish 结束会话并更新 DB // Logan.Li
func (ac *agentConn) finish(sess *model.AgentCLISession, status, errMsg string, rounds int, dur time.Duration) {
	if status == "done" || status == "aborted" || status == "error" {
		ac.write(wsEvent{Type: status, SessionID: sess.ID, Message: errMsg, Round: rounds})
	}
	updates := map[string]interface{}{
		"status": status, "rounds": rounds, "duration": dur.Milliseconds(),
	}
	if errMsg != "" {
		updates["error"] = errMsg
	}
	ac.db.Model(sess).Updates(updates)
}

// write 线程安全写 WS 消息 // Logan.Li
func (ac *agentConn) write(ev wsEvent) {
	ac.writeMu.Lock()
	defer ac.writeMu.Unlock()
	b, _ := json.Marshal(ev)
	_ = ac.conn.WriteMessage(websocket.TextMessage, b)
}

// persistMsg 持久化一条消息 // Logan.Li
func (ac *agentConn) persistMsg(sid uint, round int, typ, role, content, toolName, args, output, risk string, success bool, dur string, tokens int) {
	msg := &model.AgentCLIMessage{
		SessionID: sid, Round: round, Type: typ, Role: role, Content: content,
		ToolName: toolName, ToolArgs: args, ToolOutput: output, Risk: risk,
		Success: success, Duration: dur, Tokens: tokens,
	}
	if err := ac.db.Create(msg).Error; err != nil {
		// 持久化失败不阻断主流程
		_ = err
	}
}