Logan.Li
e0d652d211
- 支持 Windows Agent 构建、安装脚本和统一上报协议 - 增强 RDP 剪贴板、虚拟盘和文件传输能力 - 新增 AIOps 上下文工具、告警 RCA trace 落库和回放接口 - 重构 docs 模块文档并补充终端、堡垒机、监控、AI Agent 设计
610 lines
19 KiB
Go
610 lines
19 KiB
Go
// agent_cli.go - Agent CLI WebSocket 控制器 // Logan.Li
|
||
// 驱动 pkg/agentcli 引擎：主机模式（SSH）或 K8s 模式（集群查询/Pod 执行）→ 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/aiopscontext"
|
||
"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
|
||
mode string
|
||
server model.Server
|
||
sshCfg remoteexec.SSHConfig
|
||
kubernetesCtx *kubernetesAgentContext
|
||
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
|
||
running bool
|
||
runningMu sync.Mutex
|
||
}
|
||
|
||
// Run WebSocket 入口 // Logan.Li
|
||
// 主机: GET /ws/agent-cli?mode=host&id=<serverID>&token=<jwt>
|
||
// Kubernetes: GET /ws/agent-cli?mode=kubernetes&cluster=<code>&namespace=<ns>&pod=<可选>&container=<可选>&token=<jwt>
|
||
func (h *AgentCLIHandler) Run(c *gin.Context) {
|
||
mode := agentcli.NormalizeRunMode(c.DefaultQuery("mode", agentcli.RunModeHost))
|
||
if mode == agentcli.RunModeKubernetes {
|
||
h.runKubernetesWS(c)
|
||
return
|
||
}
|
||
h.runHostWS(c)
|
||
}
|
||
|
||
// runHostWS 主机模式 WebSocket // Logan.Li
|
||
func (h *AgentCLIHandler) runHostWS(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
|
||
}
|
||
|
||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer conn.Close()
|
||
|
||
sshCfg, err := buildSSHConfig(h.db, server)
|
||
if err != nil {
|
||
writeWSCloseError(conn, "SSH 连接失败: "+err.Error())
|
||
return
|
||
}
|
||
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, mode: agentcli.RunModeHost,
|
||
userID: middleware.GetUserID(c), server: server, sshCfg: sshCfg,
|
||
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)})
|
||
ac.wsReadLoop()
|
||
}
|
||
|
||
// runKubernetesWS Kubernetes 模式 WebSocket // Logan.Li
|
||
func (h *AgentCLIHandler) runKubernetesWS(c *gin.Context) {
|
||
clusterCode, namespace, podName, container := parseKubernetesWSParams(c)
|
||
if clusterCode == "" || namespace == "" {
|
||
c.JSON(http.StatusBadRequest, model.Response{Code: 400, Message: "缺少 cluster 或 namespace 参数"})
|
||
return
|
||
}
|
||
|
||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||
if err != nil {
|
||
return
|
||
}
|
||
defer conn.Close()
|
||
|
||
kCtx, err := loadKubernetesAgentContext(h.db, clusterCode, namespace, podName)
|
||
if err != nil {
|
||
writeWSCloseError(conn, "Kubernetes 连接失败: "+err.Error())
|
||
return
|
||
}
|
||
kCtx.container = container
|
||
|
||
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, mode: agentcli.RunModeKubernetes,
|
||
userID: middleware.GetUserID(c), kubernetesCtx: kCtx,
|
||
taskCtx: context.Background(), confirms: make(map[string]chan bool),
|
||
}
|
||
target := buildKubernetesTargetInfo(kCtx)
|
||
ac.write(wsEvent{Type: "session", SessionID: 0,
|
||
Message: fmt.Sprintf("已连接 Kubernetes 集群，LLM: %s。\n%s\n请描述要执行的任务。", llmCfg.Model, target)})
|
||
ac.wsReadLoop()
|
||
}
|
||
|
||
// wsReadLoop 处理前端控制消息 // Logan.Li
|
||
func (ac *agentConn) wsReadLoop() {
|
||
for {
|
||
_, data, err := ac.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(ctrl.Query, ctrl.SafetyMode)
|
||
case "continue":
|
||
if ctrl.Query == "" {
|
||
ac.write(wsEvent{Type: "error", Message: "追问内容不能为空"})
|
||
continue
|
||
}
|
||
ac.continueTask(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(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, Mode: ac.mode, Title: title, Status: "running",
|
||
Provider: llmCfg.Provider, Model: llmCfg.Model,
|
||
}
|
||
if ac.mode == agentcli.RunModeKubernetes && ac.kubernetesCtx != nil {
|
||
sess.ClusterID = ac.kubernetesCtx.cluster.ID
|
||
sess.ClusterName = ac.kubernetesCtx.cluster.Name
|
||
sess.Namespace = ac.kubernetesCtx.namespace
|
||
sess.PodName = ac.kubernetesCtx.podName
|
||
sess.Container = ac.kubernetesCtx.container
|
||
} else {
|
||
sess.ServerID = ac.server.ID
|
||
sess.ServerName = ac.server.Hostname
|
||
sess.ServerHost = ac.server.IP
|
||
}
|
||
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(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
|
||
q := ac.db.Where("user_id = ?", ac.userID)
|
||
if ac.mode == agentcli.RunModeKubernetes {
|
||
q = q.Where("mode IN ?", []string{agentcli.RunModeKubernetes, "k8s"})
|
||
if ac.kubernetesCtx != nil {
|
||
q = q.Where("cluster_id = ? AND namespace = ?", ac.kubernetesCtx.cluster.ID, ac.kubernetesCtx.namespace)
|
||
}
|
||
} else {
|
||
q = q.Where("mode = ? AND server_id = ?", ac.mode, ac.server.ID)
|
||
}
|
||
if err := q.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(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(ac.session, query, safetyMode, llmCfg, true)
|
||
}
|
||
|
||
// runTask 执行一轮 ReAct（新建或追问共用） // Logan.Li
|
||
func (ac *agentConn) runTask(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
|
||
coreCfg.Mode = ac.mode
|
||
|
||
engCfg := agentcli.EngineConfig{
|
||
Mode: ac.mode,
|
||
LLM: buildLLMConfigFromModel(llmCfg),
|
||
SafetyMode: parseSafetyMode(safetyMode),
|
||
MemoryDir: fmt.Sprintf("data/agentcli/user_%d", ac.userID),
|
||
Core: coreCfg,
|
||
ExtraTools: aiopscontext.NewAgentTools(aiopscontext.NewService(ac.db)),
|
||
}
|
||
if ac.mode == agentcli.RunModeKubernetes {
|
||
if ac.kubernetesCtx == nil {
|
||
ac.write(wsEvent{Type: "error", Message: "Kubernetes 上下文丢失"})
|
||
return
|
||
}
|
||
engCfg.Kubernetes = &agentcli.KubernetesTarget{
|
||
Client: ac.kubernetesCtx.client, ClusterName: ac.kubernetesCtx.cluster.Name,
|
||
Namespace: ac.kubernetesCtx.namespace, PodName: ac.kubernetesCtx.podName, Container: ac.kubernetesCtx.container,
|
||
}
|
||
} else {
|
||
engCfg.SSH = ac.sshCfg
|
||
}
|
||
engine, err := agentcli.NewEngine(engCfg)
|
||
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)
|
||
|
||
targetInfo := ac.buildTargetInfo()
|
||
|
||
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, targetInfo, 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)
|
||
}
|
||
}
|
||
|
||
// buildTargetInfo 构建 Prompt 目标上下文 // Logan.Li
|
||
func (ac *agentConn) buildTargetInfo() string {
|
||
if ac.mode == agentcli.RunModeKubernetes && ac.kubernetesCtx != nil {
|
||
return buildKubernetesTargetInfo(ac.kubernetesCtx)
|
||
}
|
||
return fmt.Sprintf("主机名: %s | IP: %s:%d | 用户: %s | OS: %s",
|
||
ac.server.Hostname, ac.server.IP, ac.server.Port, ac.server.Username, ac.server.OS)
|
||
}
|
||
|
||
// 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)
|
||
if cmd == "" && (toolName == "kubernetes_action" || toolName == "k8s_action") {
|
||
cmd = fmt.Sprintf("%v %v", args["action"], args["name"])
|
||
}
|
||
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
|
||
}
|
||
}
|