Logan
1acb0c49af
- Agent CLI: 多轮对话、SSH 长连接、流式推理、Evolver 按用户隔离、历史侧栏 - 终端审计: 回放播放器、命令书签、筛选分页导出、命令输出统计、录制落盘修复 - 修复 session_commands 表迁移、Kill 竞态导致录制丢失、无录制会话回放提示 - 审计历史页只读（移除删除），移除回放弹框录制搜索
479 lines
15 KiB
Go
479 lines
15 KiB
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/confirm/deny/abort
|
||
Query string `json:"query,omitempty"` // start 时任务描述
|
||
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 "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)
|
||
}
|
||
|
||
// 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.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()
|
||
if status == "done" || status == "aborted" || status == "error" {
|
||
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
|
||
}
|
||
} |