789 lines
25 KiB
Go
789 lines
25 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/guardrail"
|
||
"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"`
|
||
Phase string `json:"phase,omitempty"`
|
||
Summary string `json:"summary,omitempty"`
|
||
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"`
|
||
OutputSize int `json:"output_size,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.AIRun
|
||
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
|
||
pendingSpans map[string]string
|
||
seq int
|
||
seqMu 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), pendingSpans: make(map[string]string),
|
||
}
|
||
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), pendingSpans: make(map[string]string),
|
||
}
|
||
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 {
|
||
if ac.isRunning() {
|
||
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.AIRun{
|
||
RunID: guardrail.NewTraceID(), UserID: ac.userID, Feature: "agent_cli",
|
||
Mode: ac.mode, Title: title, Status: "running",
|
||
Provider: llmCfg.Provider, Model: llmCfg.Model,
|
||
SafetyMode: safetyMode, StartedAt: time.Now(),
|
||
}
|
||
if ac.mode == agentcli.RunModeKubernetes && ac.kubernetesCtx != nil {
|
||
sess.SubjectType = "k8s_cluster"
|
||
if ac.kubernetesCtx.podName != "" {
|
||
sess.SubjectType = "pod"
|
||
} else if ac.kubernetesCtx.namespace != "" {
|
||
sess.SubjectType = "namespace"
|
||
}
|
||
sess.SubjectID = fmt.Sprintf("%d", ac.kubernetesCtx.cluster.ID)
|
||
sess.SubjectName = ac.kubernetesCtx.cluster.Name
|
||
sess.SubjectScope = ac.kubernetesCtx.namespace
|
||
} else {
|
||
sess.SubjectType = "host"
|
||
sess.SubjectID = fmt.Sprintf("%d", ac.server.ID)
|
||
sess.SubjectName = ac.server.Hostname
|
||
sess.SubjectScope = 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.AIRun
|
||
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("subject_id = ? AND subject_scope = ?", fmt.Sprintf("%d", ac.kubernetesCtx.cluster.ID), ac.kubernetesCtx.namespace)
|
||
}
|
||
} else {
|
||
q = q.Where("mode = ? AND subject_type = ? AND subject_id = ?", ac.mode, "host", fmt.Sprintf("%d", 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.AIRunMessage
|
||
ac.db.Where("run_id = ?", sess.RunID).Order("seq ASC, 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.AIRun, 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).Updates(map[string]interface{}{"status": "running", "safety_mode": safetyMode})
|
||
|
||
ac.persistMsg(sess.RunID, 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)
|
||
}
|
||
if e.Type == core.EventDone || e.Type == core.EventError {
|
||
ac.setRunning(false)
|
||
}
|
||
|
||
ev := wsEvent{
|
||
Type: string(e.Type), Phase: e.Phase, Summary: e.Summary,
|
||
Round: e.Round, Thought: e.Thought,
|
||
ToolName: e.ToolName, ToolArgs: e.ToolArgs, Risk: e.Risk,
|
||
Output: e.Output, Success: e.Success, Duration: e.Duration, OutputSize: e.OutputSize,
|
||
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
|
||
}
|
||
runID := ac.session.RunID
|
||
switch e.Type {
|
||
case core.EventStatus:
|
||
// status 事件不改 DB schema：duration 字段在 type=status 时复用为 phase。
|
||
ac.persistMsg(runID, e.Round, "status", "system", e.Summary, e.ToolName, "", "", e.Risk, e.Phase != core.PhaseError && e.Phase != core.PhaseBlocked, e.Phase, e.Tokens)
|
||
case core.EventThought:
|
||
ac.persistMsg(runID, e.Round, "thought", "assistant", e.Thought, "", "", "", "", false, "", 0)
|
||
case core.EventToolCall:
|
||
argsJSON, _ := json.Marshal(e.ToolArgs)
|
||
ac.persistMsg(runID, e.Round, "tool_call", "tool", string(argsJSON), e.ToolName, string(argsJSON), "", e.Risk, false, "", 0)
|
||
ac.startToolSpan(e, string(argsJSON))
|
||
case core.EventToolResult:
|
||
ac.persistMsg(runID, e.Round, "tool_result", "tool", e.Output, e.ToolName, "", "", "", e.Success, e.Duration, 0)
|
||
ac.finishToolSpan(e)
|
||
case core.EventDone:
|
||
ac.setTotalTokens(e.Tokens)
|
||
ac.persistMsg(runID, e.Round, "done", "assistant", e.Answer, "", "", "", "", true, "", e.Tokens)
|
||
case core.EventError:
|
||
ac.persistMsg(runID, 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", Phase: core.PhaseApproval, Summary: "等待用户确认", 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:
|
||
ac.recordConfirmSpan(toolName, args, risk, approved)
|
||
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("确认超时")
|
||
}
|
||
}
|
||
|
||
func (ac *agentConn) recordConfirmSpan(toolName string, args map[string]interface{}, risk tools.RiskLevel, approved bool) {
|
||
if ac.session == nil {
|
||
return
|
||
}
|
||
decision := "rejected"
|
||
status := "skipped"
|
||
if approved {
|
||
decision = "approved"
|
||
status = "success"
|
||
}
|
||
argsJSON, _ := json.Marshal(args)
|
||
span := &model.AIRunSpan{
|
||
RunID: ac.session.RunID, SpanID: guardrail.NewSpanID(), Kind: "confirm", Name: toolName,
|
||
ToolSource: ac.toolSource(toolName), Input: guardrail.SanitizeJSON(string(argsJSON)),
|
||
Risk: risk.String(), Decision: decision, Status: status, StartedAt: time.Now(), CreatedAt: time.Now(),
|
||
}
|
||
_ = ac.db.Create(span).Error
|
||
}
|
||
|
||
// 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.AIRunMessage{}).Where("run_id = ?", ac.session.RunID).
|
||
Select("COALESCE(MAX(round),0)").Scan(&maxRound)
|
||
return maxRound
|
||
}
|
||
|
||
// finish 结束本轮执行并更新会话 DB // Logan.Li
|
||
func (ac *agentConn) finish(sess *model.AIRun, 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 status == "done" || status == "error" || status == "aborted" {
|
||
now := time.Now()
|
||
updates["finished_at"] = &now
|
||
}
|
||
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 持久化一条 AI Run 回放消息。 // Logan.Li
|
||
func (ac *agentConn) persistMsg(runID string, round int, typ, role, content, toolName, args, output, risk string, success bool, dur string, tokens int) {
|
||
content = guardrail.SanitizeText(content)
|
||
if typ == "tool_call" {
|
||
content = guardrail.SanitizeJSON(content)
|
||
}
|
||
msg := &model.AIRunMessage{
|
||
RunID: runID, Round: round, Seq: ac.nextSeq(),
|
||
Type: typ, Role: role, Content: content,
|
||
ToolName: toolName, Tokens: tokens,
|
||
}
|
||
if err := ac.db.Create(msg).Error; err != nil {
|
||
// 持久化失败不阻断主流程
|
||
_ = err
|
||
}
|
||
}
|
||
|
||
func (ac *agentConn) nextSeq() int {
|
||
ac.seqMu.Lock()
|
||
defer ac.seqMu.Unlock()
|
||
ac.seq++
|
||
return ac.seq
|
||
}
|
||
|
||
func (ac *agentConn) startToolSpan(e core.Event, argsJSON string) {
|
||
if ac.session == nil {
|
||
return
|
||
}
|
||
spanID := guardrail.NewSpanID()
|
||
key := ac.pendingSpanKey(e.Round, e.ToolName)
|
||
ac.runningMu.Lock()
|
||
ac.pendingSpans[key] = spanID
|
||
ac.runningMu.Unlock()
|
||
span := &model.AIRunSpan{
|
||
RunID: ac.session.RunID, SpanID: spanID, Kind: "tool", Name: e.ToolName,
|
||
ToolSource: ac.toolSource(e.ToolName), Input: guardrail.SanitizeJSON(argsJSON),
|
||
Risk: e.Risk, Decision: "auto", Status: "running",
|
||
StartedAt: time.Now(),
|
||
}
|
||
_ = ac.db.Create(span).Error
|
||
}
|
||
|
||
func (ac *agentConn) finishToolSpan(e core.Event) {
|
||
if ac.session == nil {
|
||
return
|
||
}
|
||
key := ac.pendingSpanKey(e.Round, e.ToolName)
|
||
ac.runningMu.Lock()
|
||
spanID := ac.pendingSpans[key]
|
||
delete(ac.pendingSpans, key)
|
||
ac.runningMu.Unlock()
|
||
if spanID == "" {
|
||
return
|
||
}
|
||
status := "success"
|
||
if !e.Success {
|
||
status = "error"
|
||
}
|
||
now := time.Now()
|
||
ac.db.Model(&model.AIRunSpan{}).Where("span_id = ?", spanID).Updates(map[string]interface{}{
|
||
"output": guardrail.SanitizeText(e.Output), "output_size": e.OutputSize,
|
||
"duration": parseDurationMillis(e.Duration), "status": status, "ended_at": &now,
|
||
})
|
||
ac.recordOpsCommand(e, spanID, status)
|
||
}
|
||
|
||
func (ac *agentConn) pendingSpanKey(round int, toolName string) string {
|
||
return fmt.Sprintf("%d:%s", round, toolName)
|
||
}
|
||
|
||
func (ac *agentConn) toolSource(toolName string) string {
|
||
switch toolName {
|
||
case "ssh_execute":
|
||
return "shell"
|
||
case "kubernetes_diagnose", "kubernetes_get", "pod_logs", "pod_execute", "kubernetes_action":
|
||
return "kubernetes"
|
||
case "aiops_asset_search", "aiops_alert_list", "aiops_metric_summary":
|
||
return "aiops"
|
||
default:
|
||
return "builtin"
|
||
}
|
||
}
|
||
|
||
func (ac *agentConn) recordOpsCommand(e core.Event, spanID, status string) {
|
||
if ac.session == nil {
|
||
return
|
||
}
|
||
command := e.ToolName
|
||
if e.ToolArgs != nil {
|
||
if cmd, ok := e.ToolArgs["command"].(string); ok && cmd != "" {
|
||
command = cmd
|
||
} else if action, ok := e.ToolArgs["action"]; ok {
|
||
command = fmt.Sprintf("%v %v", action, e.ToolArgs["name"])
|
||
}
|
||
}
|
||
targetType, targetID := ac.currentTarget()
|
||
cmd := &model.OpsCommand{
|
||
CommandID: guardrail.NewSpanID(), RunID: ac.session.RunID, SpanID: spanID,
|
||
UserID: ac.userID, ActorType: "ai", TargetType: targetType, TargetID: targetID,
|
||
CommandType: ac.commandType(e.ToolName), Command: guardrail.SanitizeText(command),
|
||
Risk: e.Risk, Decision: "auto", Status: status, OutputSize: e.OutputSize,
|
||
Duration: parseDurationMillis(e.Duration), ExecutedAt: time.Now(),
|
||
}
|
||
if !e.Success {
|
||
cmd.Error = guardrail.SanitizeText(e.Output)
|
||
}
|
||
_ = ac.db.Create(cmd).Error
|
||
}
|
||
|
||
func (ac *agentConn) currentTarget() (string, string) {
|
||
if ac.mode == agentcli.RunModeKubernetes && ac.kubernetesCtx != nil {
|
||
if ac.kubernetesCtx.podName != "" {
|
||
return "pod", fmt.Sprintf("%d/%s/%s", ac.kubernetesCtx.cluster.ID, ac.kubernetesCtx.namespace, ac.kubernetesCtx.podName)
|
||
}
|
||
return "k8s_cluster", fmt.Sprintf("%d", ac.kubernetesCtx.cluster.ID)
|
||
}
|
||
return "host", fmt.Sprintf("%d", ac.server.ID)
|
||
}
|
||
|
||
func (ac *agentConn) commandType(toolName string) string {
|
||
switch toolName {
|
||
case "ssh_execute", "pod_execute":
|
||
return "shell"
|
||
case "kubernetes_diagnose", "kubernetes_action", "kubernetes_get", "pod_logs":
|
||
return "kubernetes"
|
||
case "aiops_asset_search", "aiops_alert_list", "aiops_metric_summary":
|
||
return "api"
|
||
default:
|
||
return "tool"
|
||
}
|
||
}
|
||
|
||
func parseDurationMillis(raw string) int64 {
|
||
if raw == "" {
|
||
return 0
|
||
}
|
||
d, err := time.ParseDuration(raw)
|
||
if err != nil {
|
||
return 0
|
||
}
|
||
return d.Milliseconds()
|
||
}
|