xinmi/internal/workbench/operations.go

package workbench

import (
	"context"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/Tencent/WeKnora/internal/accesspolicy"
	"github.com/Tencent/WeKnora/internal/application/repository"
	apperrors "github.com/Tencent/WeKnora/internal/errors"
	"github.com/Tencent/WeKnora/internal/types"
	"github.com/gin-gonic/gin"
	"github.com/google/uuid"
	"gorm.io/gorm"
	"gorm.io/gorm/clause"
)

func actionName(action string) string {
	if action == "review-settings" {
		return "保存质检规则与风格"
	}
	if name := modelActionNames[action]; name != "" {
		return name
	}
	if name := map[string]string{"generate": "生成内容", "generate-day": "生成每日连线框架", "decompose": "拆解内容段", "import-image": "识别导图图片", "revise": "按建议修改", "review": "跨稿质检", "compose": "组合成稿", "accept": "采用内容段", "skip": "跳过内容段", "confirm": "确认训练营方案", "confirm-day": "确认每日框架", "create": "创建", "update": "修改", "delete": "删除", "cancel": "取消任务", "pause": "暂停任务", "resume": "继续任务"}[action]; name != "" {
		return name
	}
	return action
}

var modelActionNames = map[string]string{
	"knowledge_classification_recommend": "AI 推荐知识分类",
	"digital_human_distillation":         "蒸馏数字人",
	"digital_human_trial":                "数字人版本试写",
	"document_summary":                   "生成资料摘要",
	"document_auto_tag":                  "自动分类知识资料",
	"question_generation":                "生成资料问题",
	"knowledge_qa":                       "知识问答",
	"agent_round":                        "Agent 推理与回答",
	"agent_context_compaction":           "整理对话上下文",
	"follow_up.suggestions":              "生成后续问题建议",
	"web_fetch_summary":                  "生成网页摘要",
	"query_rewrite":                      "优化检索问题",
	"data_analysis_plan":                 "生成数据分析计划",
	"entity_extraction":                  "提取知识实体",
	"wiki_page_modify":                   "修改 Wiki 页面",
	"wiki_chunk_citation":                "整理 Wiki 来源引用",
	"wiki_candidate_slug":                "整理 Wiki 页面名称",
	"wiki_summary":                       "生成 Wiki 摘要",
	"wiki_knowledge_extract":             "提取 Wiki 知识",
	"wiki_taxonomy_plan":                 "规划 Wiki 分类",
	"wiki_deduplication":                 "合并 Wiki 重复内容",
	"wiki_index_intro":                   "生成 Wiki 目录介绍",
	"wiki_generation":                    "生成 Wiki 内容",
}

func operationsScope(ctx context.Context, db *gorm.DB, adminAll bool) (*gorm.DB, error) {
	t, u, admin, e := identity(ctx)
	if e != nil {
		return nil, e
	}
	q := db.Where("tenant_id = ?", t)
	if !admin || !adminAll {
		q = q.Where("owner_id = ?", u)
	}
	return q, nil
}

func (h *Handler) Balance(c *gin.Context) {
	t, u, _, e := identity(c.Request.Context())
	if e != nil {
		respondError(c, e)
		return
	}
	row := Balance{TenantID: t, UserID: u}
	e = h.service.db.WithContext(c.Request.Context()).Where("tenant_id=? AND user_id=?", t, u).Take(&row).Error
	if e != nil && !errors.Is(e, gorm.ErrRecordNotFound) {
		respondError(c, e)
		return
	}
	result(c, gin.H{"balance": row.Balance, "available": max(int64(0), row.Balance-row.Reserved), "reserved": row.Reserved})
}

func (h *Handler) Accounts(c *gin.Context) {
	t, _, admin, e := identity(c.Request.Context())
	if e != nil || !admin {
		respondError(c, ErrForbidden)
		return
	}
	rows := []struct {
		UserID   string `json:"user_id"`
		UserName string `json:"user_name"`
		Balance  int64  `json:"balance"`
		Reserved int64  `json:"reserved"`
	}{}
	e = repository.TenantUsersQuery(h.service.db.WithContext(c.Request.Context()), t).
		Select("m.user_id, u.username AS user_name, COALESCE(b.balance,0) AS balance, COALESCE(b.reserved,0) AS reserved").
		Joins("LEFT JOIN workbench_balances b ON b.tenant_id=m.tenant_id AND b.user_id=m.user_id").
		Order("u.username,u.id").Scan(&rows).Error
	if e != nil {
		respondError(c, e)
		return
	}
	result(c, rows)
}

func (s *Service) grant(ctx context.Context, user string, amount int64, reason, key string) (*Transaction, error) {
	t, actor, admin, e := identity(ctx)
	if e != nil || !admin {
		return nil, ErrForbidden
	}
	if user == "" || amount <= 0 || amount > 1_000_000_000_000 || len(key) < 8 || len(key) > 120 || len([]rune(reason)) > 1000 {
		return nil, ErrInvalid
	}
	var result Transaction
	e = s.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var member types.TenantMember
		if e := repository.TenantUsersQuery(tx, t).Select("m.*").Where("m.user_id=?", user).Take(&member).Error; e != nil {
			return e
		}
		if e := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&Balance{TenantID: t, UserID: user}).Error; e != nil {
			return e
		}
		var account Balance
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id=? AND user_id=?", t, user).Take(&account).Error; e != nil {
			return e
		}
		key = "grant:" + key
		e := tx.Where("tenant_id=? AND idempotency_key=?", t, key).Take(&result).Error
		if e == nil {
			if result.UserID != user || result.Delta != amount || result.Reason != reason {
				return ErrConflict
			}
			return nil
		}
		if !errors.Is(e, gorm.ErrRecordNotFound) {
			return e
		}
		if account.Balance > 9_000_000_000_000_000-amount {
			return ErrInvalid
		}
		result = Transaction{ID: uuid.NewString(), TenantID: t, UserID: user, Delta: amount, Kind: "credit", Reason: reason, ActorID: actor, IdempotencyKey: key, Metadata: JSON{}, CreatedAt: time.Now()}
		if e = tx.Create(&result).Error; e != nil {
			return e
		}
		return tx.Model(&account).Updates(map[string]any{"balance": account.Balance + amount, "updated_at": time.Now()}).Error
	})
	return &result, e
}

func (h *Handler) GrantBalance(c *gin.Context) {
	var body struct {
		UserID string `json:"user_id"`
		Amount int64  `json:"amount"`
		Reason string `json:"reason"`
		Key    string `json:"idempotency_key"`
	}
	if e := c.ShouldBindJSON(&body); e != nil {
		respondError(c, ErrInvalid)
		return
	}
	row, e := h.service.grant(c.Request.Context(), body.UserID, body.Amount, strings.TrimSpace(body.Reason), body.Key)
	if e != nil {
		respondError(c, e)
		return
	}
	result(c, row)
}

func operationsPage(c *gin.Context) (int, int) {
	p, _ := strconv.Atoi(c.Query("page"))
	n, _ := strconv.Atoi(c.Query("page_size"))
	if p < 1 {
		p = 1
	}
	if n < 1 {
		n = 50
	}
	return min(p, 100000), min(n, 200)
}

func (h *Handler) OperationRecords(c *gin.Context) {
	ctx := c.Request.Context()
	t, u, admin, e := identity(ctx)
	if e != nil {
		respondError(c, e)
		return
	}
	kind := c.Param("recordKind")
	page, size := operationsPage(c)
	db := h.service.db.WithContext(ctx)
	var rows any
	var q *gorm.DB
	switch kind {
	case "transactions":
		if !admin {
			respondError(c, ErrForbidden)
			return
		}
		rows = &[]Transaction{}
		q = db.Model(&Transaction{}).Where("tenant_id=?", t)
		if user := c.Query("user_id"); user != "" {
			q = q.Where("user_id=?", user)
		}
	case "usage":
		if !admin {
			respondError(c, ErrForbidden)
			return
		}
		rows = &[]Usage{}
		q = db.Model(&Usage{}).Where("tenant_id=?", t)
		if user := c.Query("user_id"); user != "" {
			q = q.Where("owner_id=?", user)
		}
	case "audit", "system":
		if !admin {
			respondError(c, ErrForbidden)
			return
		}
		rows = &[]OperationLog{}
		q = db.Model(&OperationLog{}).Where("tenant_id=? AND kind=?", t, kind)
		if kind == "system" {
			q = q.Where("level NOT IN ?", []string{"warn", "warning"})
		}
		if search := strings.TrimSpace(c.Query("q")); search != "" {
			q = q.Where("message LIKE ? OR action LIKE ? OR owner_name LIKE ? OR trace_id LIKE ?", "%"+search+"%", "%"+search+"%", "%"+search+"%", "%"+search+"%")
		}
	case "consumption":
		if !accesspolicy.Allowed(ctx, "knowledge.read") {
			respondError(c, ErrForbidden)
			return
		}
		rows = &[]Consumption{}
		q = db.Model(&Consumption{}).Where("tenant_id=?", t)
		if !admin || c.Query("all_users") != "true" {
			q = q.Where("owner_id=?", u)
		}
		if kb := c.Query("kb_id"); kb != "" {
			q = q.Where("kb_id=?", kb)
		}
		if state := c.Query("status"); state != "" {
			q = q.Where("status=?", state)
		}
	case "memories":
		if !accesspolicy.Allowed(ctx, "agents.read") {
			respondError(c, ErrForbidden)
			return
		}
		rows = &[]PendingMemory{}
		q = db.Model(&PendingMemory{}).Where("tenant_id=? AND owner_id=?", t, u)
		if state := c.Query("status"); state != "" {
			q = q.Where("status=?", state)
		}
	default:
		respondError(c, ErrInvalid)
		return
	}
	var total int64
	if e = q.Count(&total).Error; e == nil {
		e = q.Order("created_at DESC,id DESC").Offset((page - 1) * size).Limit(size).Find(rows).Error
	}
	if e != nil {
		respondError(c, e)
		return
	}
	if memories, ok := rows.(*[]PendingMemory); ok {
		h.service.decorateLegacyMemories(ctx, *memories)
	}
	c.JSON(http.StatusOK, gin.H{"success": true, "data": rows, "total": total})
}

func (h *Handler) UpdateConsumption(c *gin.Context) {
	ctx := c.Request.Context()
	t, u, _, e := identity(ctx)
	if e != nil {
		respondError(c, e)
		return
	}
	if !accesspolicy.Allowed(ctx, "knowledge.consumption") {
		respondError(c, ErrForbidden)
		return
	}
	var body struct {
		IDs    []string `json:"ids"`
		Action string   `json:"action"`
	}
	if c.ShouldBindJSON(&body) != nil || len(body.IDs) < 1 || len(body.IDs) > 1000 || (body.Action != "open" && body.Action != "ban") {
		respondError(c, ErrInvalid)
		return
	}
	e = h.service.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		var rows []Consumption
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("tenant_id=? AND owner_id=? AND id IN ?", t, u, body.IDs).Find(&rows).Error; e != nil {
			return e
		}
		if len(rows) != len(body.IDs) {
			return ErrForbidden
		}
		for _, r := range rows {
			if r.Status == "reserved" && r.LeaseUntil != nil && r.LeaseUntil.After(time.Now()) {
				return fmt.Errorf("%w: 资料正在生成中使用，请等待任务结束", ErrConflict)
			}
		}
		return tx.Model(&Consumption{}).Where("tenant_id=? AND owner_id=? AND id IN ?", t, u, body.IDs).Updates(map[string]any{"status": "open", "banned": body.Action == "ban", "job_id": "", "lease_until": nil, "updated_at": time.Now()}).Error
	})
	if e != nil {
		respondError(c, e)
		return
	}
	result(c, gin.H{"updated": len(body.IDs)})
}

func (h *Handler) ReviewMemory(c *gin.Context) {
	ctx := c.Request.Context()
	t, u, _, e := identity(ctx)
	if e != nil {
		respondError(c, e)
		return
	}
	if !accesspolicy.Allowed(ctx, "agents.memory") {
		respondError(c, ErrForbidden)
		return
	}
	var body memoryReviewRequest
	if c.ShouldBindJSON(&body) != nil || (body.Status != "accepted" && body.Status != "rejected") {
		respondError(c, ErrInvalid)
		return
	}
	var row PendingMemory
	e = h.service.db.WithContext(ctx).Transaction(func(tx *gorm.DB) error {
		if e := tx.Where("id=? AND tenant_id=? AND owner_id=?", c.Param("id"), t, u).Take(&row).Error; e != nil {
			return e
		}
		if body.ReconfirmLegacy {
			return h.service.reconfirmLegacyMemory(ctx, tx, &row, body)
		}
		if body.Status == "accepted" {
			if e := lockContentFeedbackSource(tx, &row); e != nil {
				return e
			}
		}
		if e := tx.Clauses(clause.Locking{Strength: "UPDATE"}).Where("id=? AND tenant_id=? AND owner_id=?", c.Param("id"), t, u).Take(&row).Error; e != nil {
			return e
		}
		if row.Status != "pending" {
			return ErrConflict
		}
		if body.Content != "" {
			row.Content = strings.TrimSpace(body.Content)
		}
		if row.Content == "" || len([]rune(row.Content)) > 12000 {
			return ErrInvalid
		}
		if body.Status == "accepted" {
			if len([]rune(row.Content)) > types.MemoryContentMaxRunes {
				return fmt.Errorf("%w: 长期记忆最多 %d 字，请精简后再确认；当前候选仍保留", ErrInvalid, types.MemoryContentMaxRunes)
			}
			if h.service.rememberInTransaction == nil {
				return fmt.Errorf("%w: 记忆服务未初始化", ErrInvalid)
			}
			item := types.MemoryItem{Kind: types.MemoryKindFact, Content: row.Content, Origin: types.MemoryOriginExplicit, Importance: 3}
			if strings.HasPrefix(row.SourceRef, contentMemorySourcePrefix) {
				item.Kind = types.MemoryKindPreference
				item.SourceMessageID = row.ID
			}
			stored, err := h.service.rememberInTransaction(ctx, tx, item)
			if err != nil {
				return err
			}
			if stored == nil || stored.Status != types.MemoryStatusActive {
				return fmt.Errorf("%w: 记忆尚未写入，请重试", ErrInvalid)
			}
			// Approval history must show the sanitized content actually retained.
			row.Content = stored.Content
		}
		row.Status, row.ReviewedBy, row.UpdatedAt = body.Status, u, time.Now()
		return tx.Save(&row).Error
	})
	if e != nil {
		respondError(c, e)
		return
	}
	result(c, row)
}

func (s *Service) recordOperation(ctx context.Context, row OperationLog) error {
	t, ok := types.TenantIDFromContext(ctx)
	if !ok || t == 0 {
		return ErrForbidden
	}
	u, _ := types.UserIDFromContext(ctx)
	row.ID, row.TenantID, row.OwnerID, row.CreatedAt = uuid.NewString(), t, u, time.Now()
	if u != "" && !types.IsSyntheticUserID(u) {
		row.OwnerName = u
		var actor struct{ Username string }
		if e := s.db.WithContext(ctx).Table("users").Select("username").Where("id=?", u).Take(&actor).Error; e == nil && actor.Username != "" {
			row.OwnerName = actor.Username
		}
	} else {
		row.OwnerName = "系统任务"
	}
	if row.Metadata == nil {
		row.Metadata = JSON{}
	}
	return s.db.WithContext(ctx).Create(&row).Error
}

// The outer ErrorHandler may not have rendered a pending error yet.
func operationHTTPStatus(c *gin.Context) int {
	if c.Writer.Written() || len(c.Errors) == 0 {
		return c.Writer.Status()
	}
	var appError *apperrors.AppError
	if errors.As(c.Errors.Last().Err, &appError) && appError.HTTPCode >= 400 && appError.HTTPCode <= 599 {
		return appError.HTTPCode
	}
	return http.StatusInternalServerError
}

// Request logs store route names and safe diagnostics, never request bodies or keys.
func (h *Handler) OperationAudit() gin.HandlerFunc {
	return func(c *gin.Context) {
		started := time.Now()
		c.Next()
		ctx := c.Request.Context()
		_, _, _, e := identity(ctx)
		if e != nil {
			return
		}
		path := c.FullPath()
		if path == "" {
			return
		}
		status := operationHTTPStatus(c)
		write := c.Request.Method != "GET" && c.Request.Method != "HEAD" && c.Request.Method != "OPTIONS"
		if !write && status < 400 {
			return
		}
		if strings.Contains(path, "/notifications/") || strings.Contains(path, "/auth/") {
			return
		}
		action := map[string]string{"POST": "创建或执行", "PUT": "修改", "PATCH": "修改", "DELETE": "删除", "GET": "访问"}[c.Request.Method]
		module := "系统"
		for _, entry := range [][2]string{{"training-maps", "训练营导图"}, {"content-tasks", "内容任务"}, {"digital-humans", "数字人"}, {"knowledge", "知识库"}, {"agents", "Agent"}, {"skills", "技能"}, {"memory", "记忆"}, {"access", "访问控制"}, {"models", "模型"}, {"flows", "画布"}, {"operations", "业务管理"}} {
			if strings.Contains(c.Request.URL.Path, entry[0]) {
				module = entry[1]
				break
			}
		}
		if a := c.Param("action"); a != "" {
			action = actionName(a)
		}
		message := module + "：" + action
		level := "info"
		if status >= 400 {
			level = "warn"
			if status >= 500 {
				level = "error"
			}
			message += fmt.Sprintf("失败（HTTP %d）", status)
		}
		trace := c.Writer.Header().Get("X-Request-ID")
		if trace == "" {
			trace = uuid.NewString()
		}
		finish, cancel := context.WithTimeout(context.WithoutCancel(ctx), 3*time.Second)
		defer cancel()
		row := OperationLog{Kind: "audit", Level: level, Action: action, Message: message, ResourceID: c.Param("id"), TraceID: trace, Status: status, Metadata: JSON{"method": c.Request.Method, "route": path, "duration_ms": time.Since(started).Milliseconds()}}
		if status >= 400 {
			row.Metadata["error_class"] = fmt.Sprintf("http_%d", status)
			if last := c.Errors.Last(); last != nil {
				for key, value := range jobErrorMetadata(last.Err) {
					row.Metadata[key] = value
				}
				row.Message += "；" + row.Metadata["error_message"].(string)
			}
		}
		if e = h.service.recordOperation(finish, row); e != nil {
			logJobFailure(finish, nil, "audit", e)
		}
		if status >= 400 {
			row.Kind = "system"
			if status == 403 {
				row.Message = "权限不足：" + module + " / " + action
				if label := c.GetString("required_permission_label"); label != "" {
					row.Message = "权限不足：" + label
				}
			}
			if e = h.service.recordOperation(finish, row); e != nil {
				logJobFailure(finish, nil, "system_log", e)
			}
		}
	}
}