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/continue/resume/confirm/deny/abort
	Query      string `json:"query,omitempty"`         // start/continue 时任务描述
	SessionID  uint   `json:"session_id,omitempty"`    // resume 时恢复的历史会话 ID
	SafetyMode string `json:"safety_mode,omitempty"`   // strict/balanced/permissive
	ConfirmID  string `json:"confirm_id,omitempty"`    // confirm/deny 时
}

// writeWSCloseError 在 agentConn 建立前，直接通过 conn 推送一条 error 事件再关闭连接 // Logan.Li
// 用于 WS 升级后的预检失败（SSH 凭据/LLM 配置），让前端能拿到具体原因而非"一直转圈"
func writeWSCloseError(conn *websocket.Conn, msg string) {
	_ = conn.WriteJSON(wsEvent{Type: "error", Message: msg})
	// 留出短暂窗口让浏览器接收 error 帧，再发送 close 帧：否则 close 帧紧随其后、叠加
	// defer conn.Close() 触发的 TCP 关闭，可能令 error 帧被丢弃，前端表现就是"正在连接主机"
	// 永久转圈、看不到真实原因 // Logan.Li
	time.Sleep(120 * time.Millisecond)
	_ = conn.WriteMessage(websocket.CloseMessage,
		websocket.FormatCloseMessage(websocket.CloseNormalClosure, msg))
}

// agentConn 单个 WS 连接的运行态 // Logan.Li
type agentConn struct {
	conn        *websocket.Conn
	db          *gorm.DB
	userID      uint
	server      model.Server
	session     *model.AgentCLISession
	taskCtx     context.Context
	cancel      context.CancelFunc
	totalTokens int
	tokenMu     sync.Mutex
	writeMu     sync.Mutex
	confirms     map[string]chan bool
	cfMu         sync.Mutex
	convHistory  []llm.Message // 多轮追问对话历史（不含 system） // Logan.Li
	running      bool
	runningMu    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
	}

	// 先升级 WebSocket：SSH 凭据 / LLM 配置预检放到升级之后，失败时通过 WS error 事件
	// 把真实原因推给前端，避免握手 400 导致前端"一直转圈"看不到原因 // Logan.Li
	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
	if err != nil {
		return
	}
	defer conn.Close()

	// SSH 凭据预检
	sshCfg, err := buildSSHConfig(h.db, server)
	if err != nil {
		writeWSCloseError(conn, "SSH 连接失败: "+err.Error())
		return
	}
	// LLM 配置预检
	var llmCfg model.LLMConfig
	if err := h.db.First(&llmCfg).Error; err != nil || !llmCfg.Enabled {
		writeWSCloseError(conn, "LLM 服务未配置或未启用，请先在「大模型设置」中配置")
		return
	}

	ac := &agentConn{
		conn: conn, db: h.db,
		userID: middleware.GetUserID(c), server: server,
		taskCtx:  context.Background(),
		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, ctrl.Query, ctrl.SafetyMode)
		case "continue":
			if ctrl.Query == "" {
				ac.write(wsEvent{Type: "error", Message: "追问内容不能为空"})
				continue
			}
			ac.continueTask(sshCfg, ctrl.Query, ctrl.SafetyMode)
		case "resume":
			if ctrl.SessionID == 0 {
				ac.write(wsEvent{Type: "error", Message: "缺少 session_id"})
				continue
			}
			ac.resumeSession(ctrl.SessionID)
		case "confirm":
			ac.resolveConfirm(ctrl.ConfirmID, true)
		case "deny":
			ac.resolveConfirm(ctrl.ConfirmID, false)
		case "abort":
			ac.abort()
		}
	}
}

// parseSafetyMode 解析前端传入的安全模式 // Logan.Li
func parseSafetyMode(mode string) safety.Mode {
	switch mode {
	case "strict":
		return safety.ModeStrict
	case "permissive":
		return safety.ModePermissive
	default:
		return safety.ModeBalanced
	}
}

// startTask 启动新任务（新会话，清空对话历史） // Logan.Li
func (ac *agentConn) startTask(sshCfg remoteexec.SSHConfig, query, safetyMode string) {
	ac.abort()

	var llmCfg model.LLMConfig
	if err := ac.db.First(&llmCfg).Error; err != nil || !llmCfg.Enabled {
		ac.write(wsEvent{Type: "error", Message: "LLM 服务未配置或未启用，请先在「大模型设置」中配置"})
		return
	}

	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.convHistory = nil
	ac.setTotalTokens(0)
	ac.runTask(sshCfg, sess, query, safetyMode, llmCfg, false)
}

// resumeSession 从历史会话恢复上下文，供继续追问 // Logan.Li
func (ac *agentConn) resumeSession(sessionID uint) {
	if ac.isRunning() {
		ac.write(wsEvent{Type: "error", Message: "任务进行中，请等待完成或停止"})
		return
	}
	var sess model.AgentCLISession
	if err := ac.db.Where("user_id = ? AND server_id = ?", ac.userID, ac.server.ID).
		First(&sess, sessionID).Error; err != nil {
		ac.write(wsEvent{Type: "error", Message: "会话不存在或无权访问"})
		return
	}
	if sess.Status == "running" {
		ac.write(wsEvent{Type: "error", Message: "该会话仍在执行中，无法恢复"})
		return
	}
	var msgs []model.AgentCLIMessage
	ac.db.Where("session_id = ?", sessionID).Order("id ASC").Find(&msgs)
	ac.session = &sess
	ac.convHistory = buildConvHistoryFromMessages(msgs)
	ac.setTotalTokens(sess.TotalTokens)
	ac.write(wsEvent{
		Type: "session", SessionID: sess.ID,
		Message: fmt.Sprintf("已恢复历史会话「%s」，可继续追问或执行新指令", sess.Title),
	})
}

// continueTask 在同一会话内追问（保留上下文） // Logan.Li
func (ac *agentConn) continueTask(sshCfg remoteexec.SSHConfig, query, safetyMode string) {
	if ac.isRunning() {
		ac.write(wsEvent{Type: "error", Message: "任务进行中，请等待完成或停止"})
		return
	}
	if ac.session == nil {
		ac.write(wsEvent{Type: "error", Message: "请先开始新任务"})
		return
	}
	if ac.session.Status == "running" {
		ac.write(wsEvent{Type: "error", Message: "当前会话仍在执行中"})
		return
	}

	var llmCfg model.LLMConfig
	if err := ac.db.First(&llmCfg).Error; err != nil || !llmCfg.Enabled {
		ac.write(wsEvent{Type: "error", Message: "LLM 服务未配置或未启用，请先在「大模型设置」中配置"})
		return
	}
	ac.runTask(sshCfg, ac.session, query, safetyMode, llmCfg, true)
}

// runTask 执行一轮 ReAct（新建或追问共用） // Logan.Li
func (ac *agentConn) runTask(sshCfg remoteexec.SSHConfig, sess *model.AgentCLISession, query, safetyMode string, llmCfg model.LLMConfig, isContinue bool) {
	if ac.isRunning() {
		ac.write(wsEvent{Type: "error", Message: "任务进行中，请等待完成或停止"})
		return
	}

	prior := ac.convHistory
	priorTokens := ac.getTotalTokens()
	roundOffset := 0
	if !isContinue {
		prior = nil
		priorTokens = 0
	} else {
		roundOffset = ac.lastRound()
	}

	coreCfg := core.DefaultConfig()
	coreCfg.RoundOffset = roundOffset

	engine, err := agentcli.NewEngine(agentcli.EngineConfig{
		LLM:        buildLLMConfigFromModel(llmCfg),
		SSH:        sshCfg,
		SafetyMode: parseSafetyMode(safetyMode),
		MemoryDir:  fmt.Sprintf("data/agentcli/user_%d", ac.userID),
		Core:       coreCfg,
	})
	if err != nil {
		ac.write(wsEvent{Type: "error", Message: "引擎初始化失败: " + err.Error()})
		ac.finish(sess, "error", "引擎初始化失败: "+err.Error(), ac.lastRound(), 0)
		return
	}
	engine.SetCallback(ac.onEvent)
	engine.SetConfirmFn(ac.confirmFn)

	ctx, cancel := context.WithCancel(context.Background())
	ac.taskCtx = ctx
	ac.cancel = cancel
	ac.setRunning(true)
	ac.db.Model(sess).Update("status", "running")

	ac.persistMsg(sess.ID, 0, "user", "user", query, "", "", "", "", false, "", 0)

	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)

	msg := "任务已开始"
	if isContinue {
		msg = "继续追问"
	}
	ac.write(wsEvent{Type: "session", SessionID: sess.ID, Message: msg})

	go func() {
		defer func() {
			ac.setRunning(false)
			engine.Close()
		}()
		start := time.Now()
		result, err := engine.Run(ctx, query, serverInfo, prior, priorTokens)
		if result != nil && len(result.History) > 0 {
			ac.convHistory = result.History
			ac.setTotalTokens(result.TotalTokens)
		}
		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)
		}
	}()
}

func (ac *agentConn) isRunning() bool {
	ac.runningMu.Lock()
	defer ac.runningMu.Unlock()
	return ac.running
}

func (ac *agentConn) setRunning(v bool) {
	ac.runningMu.Lock()
	ac.running = v
	ac.runningMu.Unlock()
}

// setTotalTokens 更新当前任务 token 计数 // Logan.Li
func (ac *agentConn) setTotalTokens(n int) {
	ac.tokenMu.Lock()
	ac.totalTokens = n
	ac.tokenMu.Unlock()
}

func (ac *agentConn) getTotalTokens() int {
	ac.tokenMu.Lock()
	defer ac.tokenMu.Unlock()
	return ac.totalTokens
}

// onEvent 引擎事件回调：推 WS + 持久化 // Logan.Li
func (ac *agentConn) onEvent(e core.Event) {
	if e.Type == core.EventProgress && e.Tokens > 0 {
		ac.setTotalTokens(e.Tokens)
	}

	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 / thought_delta 太频繁不入库） // Logan.Li
	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.setTotalTokens(e.Tokens)
		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 <-ac.taskCtx.Done():
		ac.delConfirm(confirmID)
		return false, ac.taskCtx.Err()
	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
	}
	ac.cfMu.Lock()
	for id, ch := range ac.confirms {
		select {
		case ch <- false:
		default:
		}
		delete(ac.confirms, id)
	}
	ac.cfMu.Unlock()
}

// 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) {
	ac.setRunning(false)
	tokens := ac.getTotalTokens()
	// done/error 已由引擎 onEvent 推送；此处仅 aborted 需补推 WS（引擎初始化失败在 runTask 内单独推送） // Logan.Li
	if status == "aborted" {
		ac.write(wsEvent{Type: status, SessionID: sess.ID, Message: errMsg, Round: rounds, Tokens: tokens})
	}
	updates := map[string]interface{}{
		"status": status, "rounds": rounds, "duration": dur.Milliseconds(),
		"total_tokens": tokens,
	}
	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
	}
}