update

backend/controller/agent_cli.go

@@ -19,6 +19,7 @@ import (
 	"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"
@@ -86,7 +87,7 @@ type agentConn struct {
 	server        model.Server
 	sshCfg        remoteexec.SSHConfig
 	kubernetesCtx *kubernetesAgentContext
-	session       *model.AgentCLISession
+	session       *model.AIRun
 	taskCtx       context.Context
 	cancel        context.CancelFunc
 	totalTokens   int
@@ -97,6 +98,9 @@ type agentConn struct {
 	convHistory   []llm.Message
 	running       bool
 	runningMu     sync.Mutex
+	pendingSpans  map[string]string
+	seq           int
+	seqMu         sync.Mutex
 }
 
 // Run WebSocket 入口 // Logan.Li
@@ -144,7 +148,7 @@ func (h *AgentCLIHandler) runHostWS(c *gin.Context) {
 	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),
+		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)})
@@ -181,7 +185,7 @@ func (h *AgentCLIHandler) runKubernetesWS(c *gin.Context) {
 	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),
+		taskCtx: context.Background(), confirms: make(map[string]chan bool), pendingSpans: make(map[string]string),
 	}
 	target := buildKubernetesTargetInfo(kCtx)
 	ac.write(wsEvent{Type: "session", SessionID: 0,
@@ -194,7 +198,9 @@ func (ac *agentConn) wsReadLoop() {
 	for {
 		_, data, err := ac.conn.ReadMessage()
 		if err != nil {
-			ac.abort()
+			if ac.isRunning() {
+				ac.abort()
+			}
 			return
 		}
 		var ctrl wsControl
@@ -256,20 +262,27 @@ func (ac *agentConn) startTask(query, safetyMode string) {
 	if r := []rune(title); len(r) > 60 {
 		title = string(r[:60]) + "..."
 	}
-	sess := &model.AgentCLISession{
-		UserID: ac.userID, Mode: ac.mode, Title: title, Status: "running",
+	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.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
+		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.ServerID = ac.server.ID
-		sess.ServerName = ac.server.Hostname
-		sess.ServerHost = ac.server.IP
+		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()})
@@ -287,15 +300,15 @@ func (ac *agentConn) resumeSession(sessionID uint) {
 		ac.write(wsEvent{Type: "error", Message: "任务进行中，请等待完成或停止"})
 		return
 	}
-	var sess model.AgentCLISession
+	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("cluster_id = ? AND namespace = ?", ac.kubernetesCtx.cluster.ID, ac.kubernetesCtx.namespace)
+			q = q.Where("subject_id = ? AND subject_scope = ?", fmt.Sprintf("%d", ac.kubernetesCtx.cluster.ID), ac.kubernetesCtx.namespace)
 		}
 	} else {
-		q = q.Where("mode = ? AND server_id = ?", ac.mode, ac.server.ID)
+		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: "会话不存在或无权访问"})
@@ -305,8 +318,8 @@ func (ac *agentConn) resumeSession(sessionID uint) {
 		ac.write(wsEvent{Type: "error", Message: "该会话仍在执行中，无法恢复"})
 		return
 	}
-	var msgs []model.AgentCLIMessage
-	ac.db.Where("session_id = ?", sessionID).Order("id ASC").Find(&msgs)
+	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)
@@ -340,7 +353,7 @@ func (ac *agentConn) continueTask(query, safetyMode string) {
 }
 
 // runTask 执行一轮 ReAct（新建或追问共用） // Logan.Li
-func (ac *agentConn) runTask(sess *model.AgentCLISession, query, safetyMode string, llmCfg model.LLMConfig, isContinue bool) {
+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
@@ -393,9 +406,9 @@ func (ac *agentConn) runTask(sess *model.AgentCLISession, query, safetyMode stri
 	ac.taskCtx = ctx
 	ac.cancel = cancel
 	ac.setRunning(true)
-	ac.db.Model(sess).Update("status", "running")
+	ac.db.Model(sess).Updates(map[string]interface{}{"status": "running", "safety_mode": safetyMode})
 
-	ac.persistMsg(sess.ID, 0, "user", "user", query, "", "", "", "", false, "", 0)
+	ac.persistMsg(sess.RunID, 0, "user", "user", query, "", "", "", "", false, "", 0)
 
 	targetInfo := ac.buildTargetInfo()
 
@@ -457,6 +470,9 @@ 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,
@@ -474,23 +490,25 @@ func (ac *agentConn) onEvent(e core.Event) {
 	if ac.session == nil {
 		return
 	}
-	sid := ac.session.ID
+	runID := ac.session.RunID
 	switch e.Type {
 	case core.EventStatus:
 		// status 事件不改 DB schema：duration 字段在 type=status 时复用为 phase。
-		ac.persistMsg(sid, e.Round, "status", "system", e.Summary, e.ToolName, "", "", e.Risk, e.Phase != core.PhaseError && e.Phase != core.PhaseBlocked, e.Phase, e.Tokens)
+		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(sid, e.Round, "thought", "assistant", e.Thought, "", "", "", "", false, "", 0)
+		ac.persistMsg(runID, 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)
+		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(sid, e.Round, "tool_result", "tool", e.Output, e.ToolName, "", "", "", e.Success, e.Duration, 0)
+		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(sid, e.Round, "done", "assistant", e.Answer, "", "", "", "", true, "", e.Tokens)
+		ac.persistMsg(runID, e.Round, "done", "assistant", e.Answer, "", "", "", "", true, "", e.Tokens)
 	case core.EventError:
-		ac.persistMsg(sid, e.Round, "error", "assistant", e.Message, "", "", "", "", false, "", 0)
+		ac.persistMsg(runID, e.Round, "error", "assistant", e.Message, "", "", "", "", false, "", 0)
 	}
 }
 
@@ -519,6 +537,7 @@ func (ac *agentConn) confirmFn(toolName string, args map[string]interface{}, ris
 
 	select {
 	case approved := <-ch:
+		ac.recordConfirmSpan(toolName, args, risk, approved)
 		return approved, nil
 	case <-ac.taskCtx.Done():
 		ac.delConfirm(confirmID)
@@ -529,6 +548,25 @@ func (ac *agentConn) confirmFn(toolName string, args map[string]interface{}, ris
 	}
 }
 
+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()
@@ -571,13 +609,13 @@ func (ac *agentConn) lastRound() int {
 		return 0
 	}
 	var maxRound int
-	ac.db.Model(&model.AgentCLIMessage{}).Where("session_id = ?", ac.session.ID).
+	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.AgentCLISession, status, errMsg string, rounds int, dur time.Duration) {
+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
@@ -588,6 +626,10 @@ func (ac *agentConn) finish(sess *model.AgentCLISession, status, errMsg string,
 		"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
 	}
@@ -602,15 +644,145 @@ func (ac *agentConn) write(ev wsEvent) {
 	_ = 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,
+// 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()
+}