xinmi/internal/workbench/quota.go

package workbench

import (
	"context"
	"fmt"
	"math"
	"net/http"
	"time"

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

// The legacy application used Shanghai calendar days, independent of the new server timezone.
var quotaLocation = time.FixedZone("Asia/Shanghai", 8*60*60)

type UsageQuota struct {
	ID                string    `json:"id" gorm:"primaryKey"`
	TenantID          uint64    `json:"tenant_id" gorm:"uniqueIndex:wb_quota_scope"`
	UserID            string    `json:"user_id" gorm:"uniqueIndex:wb_quota_scope"`
	Period            string    `json:"period" gorm:"uniqueIndex:wb_quota_scope"`
	TokenLimit        int64     `json:"token_limit"`
	PointsLimit       int64     `json:"points_limit"`
	StorageLimitBytes int64     `json:"storage_limit_bytes"`
	ConcurrentLimit   int       `json:"concurrent_limit"`
	WarningRatio      float64   `json:"warning_ratio"`
	OveragePolicy     string    `json:"overage_policy"`
	Enabled           bool      `json:"enabled"`
	UpdatedAt         time.Time `json:"updated_at"`
}

func (UsageQuota) TableName() string { return "workbench_usage_quotas" }

type quotaUsage struct {
	Tokens         int64 `json:"tokens"`
	Points         int64 `json:"points"`
	ReservedTokens int64 `json:"reserved_tokens"`
	ReservedPoints int64 `json:"reserved_points"`
	ActiveCalls    int64 `json:"active_calls"`
}

type quotaState struct {
	UsageQuota
	UserName        string     `json:"user_name"`
	Used            quotaUsage `json:"used" gorm:"-"`
	Warning         bool       `json:"warning"`
	Exceeded        bool       `json:"exceeded"`
	Allowed         bool       `json:"allowed"`
	Ratio           float64    `json:"ratio"`
	PeriodStart     time.Time  `json:"period_start"`
	PeriodEnd       time.Time  `json:"period_end"`
	PeriodTimezone  string     `json:"period_timezone"`
	StorageEnforced bool       `json:"storage_enforced"`
}

func quotaPeriod(period string, now time.Time) (time.Time, time.Time) {
	local := now.In(quotaLocation)
	day := local.Day()
	if period == "monthly" {
		day = 1
	}
	start := time.Date(local.Year(), local.Month(), day, 0, 0, 0, 0, quotaLocation)
	if period == "monthly" {
		return start, start.AddDate(0, 1, 0)
	}
	return start, start.AddDate(0, 0, 1)
}

func loadQuotaState(db *gorm.DB, quota UsageQuota, now time.Time) (quotaState, error) {
	start, end := quotaPeriod(quota.Period, now)
	state := quotaState{UsageQuota: quota, Allowed: true, PeriodStart: start, PeriodEnd: end, PeriodTimezone: "Asia/Shanghai"}
	if e := db.Model(&Usage{}).Select("COALESCE(SUM(input_tokens+output_tokens),0) AS tokens, COALESCE(SUM(points),0) AS points").
		Where("tenant_id=? AND owner_id=? AND status<>? AND created_at>=? AND created_at<?", quota.TenantID, quota.UserID, "void", start.UTC(), end.UTC()).Scan(&state.Used).Error; e != nil {
		return state, e
	}
	var reservations quotaUsage
	if e := db.Model(&BalanceHold{}).Select("COALESCE(SUM((amount+?)/?),0) AS reserved_tokens, COALESCE(SUM(amount),0) AS reserved_points", pointsPerToken-1, pointsPerToken).
		Where("tenant_id=? AND user_id=? AND status=? AND expires_at>? AND created_at>=? AND created_at<?", quota.TenantID, quota.UserID, "reserved", now.UTC(), start.UTC(), end.UTC()).Scan(&reservations).Error; e != nil {
		return state, e
	}
	state.Used.ReservedTokens, state.Used.ReservedPoints = reservations.ReservedTokens, reservations.ReservedPoints
	// A call crossing midnight still occupies a concurrent slot in the new day.
	if e := db.Model(&BalanceHold{}).Where("tenant_id=? AND user_id=? AND status=? AND expires_at>?", quota.TenantID, quota.UserID, "reserved", now.UTC()).Count(&state.Used.ActiveCalls).Error; e != nil {
		return state, e
	}
	if quota.Enabled {
		tokens, points := state.Used.Tokens+state.Used.ReservedTokens, state.Used.Points+state.Used.ReservedPoints
		if quota.TokenLimit > 0 {
			state.Ratio = max(state.Ratio, float64(tokens)/float64(quota.TokenLimit))
			state.Exceeded = tokens >= quota.TokenLimit
		}
		if quota.PointsLimit > 0 {
			state.Ratio = max(state.Ratio, float64(points)/float64(quota.PointsLimit))
			state.Exceeded = state.Exceeded || points >= quota.PointsLimit
		}
		state.Warning = (quota.TokenLimit > 0 || quota.PointsLimit > 0) && state.Ratio >= quota.WarningRatio
		state.Allowed = !(state.Exceeded && quota.OveragePolicy != "warn") && (quota.ConcurrentLimit == 0 || state.Used.ActiveCalls < int64(quota.ConcurrentLimit))
	}
	return state, nil
}

// The caller holds the user's balance row lock, so quota reads and the new hold are atomic.
func checkQuotaReservation(tx *gorm.DB, hold *BalanceHold) error {
	var quotas []UsageQuota
	if e := tx.Where("tenant_id=? AND user_id=? AND enabled=?", hold.TenantID, hold.UserID, true).Find(&quotas).Error; e != nil {
		return e
	}
	for _, quota := range quotas {
		state, e := loadQuotaState(tx, quota, hold.CreatedAt)
		if e != nil {
			return e
		}
		if quota.ConcurrentLimit > 0 && state.Used.ActiveCalls >= int64(quota.ConcurrentLimit) {
			return fmt.Errorf("%w: 已达到管理员设置的并发调用上限，请等待当前调用结束", ErrInvalid)
		}
		period := "本月"
		if quota.Period == "daily" {
			period = "今日"
		}
		projectedTokens := state.Used.Tokens + state.Used.ReservedTokens + (hold.Amount+pointsPerToken-1)/pointsPerToken
		projectedPoints := state.Used.Points + state.Used.ReservedPoints + hold.Amount
		if quota.OveragePolicy != "warn" && quota.TokenLimit > 0 && projectedTokens > quota.TokenLimit {
			return fmt.Errorf("%w: %s Token 配额不足，请调低输出长度或联系管理员", ErrInvalid, period)
		}
		if quota.OveragePolicy != "warn" && quota.PointsLimit > 0 && projectedPoints > quota.PointsLimit {
			return fmt.Errorf("%w: %s点数配额不足，请调低输出长度或联系管理员", ErrInvalid, period)
		}
		warning := quota.TokenLimit > 0 && float64(projectedTokens)/float64(quota.TokenLimit) >= quota.WarningRatio || quota.PointsLimit > 0 && float64(projectedPoints)/float64(quota.PointsLimit) >= quota.WarningRatio
		if warning {
			// One member notification per quota period, including repeated concurrent calls.
			id := uuid.NewSHA1(uuid.NameSpaceOID, []byte(quota.ID+":"+state.PeriodStart.Format(time.RFC3339))).String()
			event := Event{ID: id, TenantID: quota.TenantID, OwnerID: quota.UserID, Action: "用量预警", Message: period + "用量及预留已达到管理员设置的预警比例，请留意剩余额度", Level: "warning", CreatedAt: hold.CreatedAt}
			if e := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&event).Error; e != nil {
				return e
			}
		}
	}
	return nil
}

func validateQuota(q UsageQuota) error {
	if q.Period != "daily" && q.Period != "monthly" || q.TokenLimit < 0 || q.TokenLimit > 1_000_000_000_000 || q.PointsLimit < 0 || q.PointsLimit > 1_000_000_000_000 || q.StorageLimitBytes < 0 || q.StorageLimitBytes > 1_000_000_000_000 || q.ConcurrentLimit < 0 || q.ConcurrentLimit > 10000 || math.IsNaN(q.WarningRatio) || math.IsInf(q.WarningRatio, 0) || q.WarningRatio < 0 || q.WarningRatio > 1 || q.OveragePolicy != "warn" && q.OveragePolicy != "soft_limit" && q.OveragePolicy != "hard_limit" {
		return ErrInvalid
	}
	return nil
}

func (s *Service) saveQuota(ctx context.Context, user string, q UsageQuota) (*UsageQuota, error) {
	t, _, admin, e := identity(ctx)
	if e != nil || !admin {
		return nil, ErrForbidden
	}
	if user == "" {
		return nil, ErrInvalid
	}
	if e = validateQuota(q); e != nil {
		return nil, e
	}
	q.ID, q.TenantID, q.UserID, q.UpdatedAt = uuid.NewString(), t, user, time.Now()
	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
		}
		// Serialize configuration changes with in-flight reservation checks.
		if e := tx.Model(&Balance{}).Where("tenant_id=? AND user_id=?", t, user).Update("updated_at", q.UpdatedAt).Error; e != nil {
			return e
		}
		if e := tx.Clauses(clause.OnConflict{Columns: []clause.Column{{Name: "tenant_id"}, {Name: "user_id"}, {Name: "period"}}, DoUpdates: clause.AssignmentColumns([]string{"token_limit", "points_limit", "storage_limit_bytes", "concurrent_limit", "warning_ratio", "overage_policy", "enabled", "updated_at"})}).Create(&q).Error; e != nil {
			return e
		}
		var saved UsageQuota
		if e := tx.Where("tenant_id=? AND user_id=? AND period=?", t, user, q.Period).Take(&saved).Error; e != nil {
			return e
		}
		q = saved
		return nil
	})
	return &q, e
}

func (h *Handler) Quotas(c *gin.Context) {
	t, _, admin, e := identity(c.Request.Context())
	if e != nil || !admin {
		respondError(c, ErrForbidden)
		return
	}
	db := h.service.db.WithContext(c.Request.Context())
	rows := []quotaState{}
	if e = repository.TenantUsersQuery(db, t).Select("q.*, u.username AS user_name").
		Joins("JOIN workbench_usage_quotas q ON q.user_id=m.user_id AND q.tenant_id=m.tenant_id").
		Order("u.username,u.id,q.period").Scan(&rows).Error; e != nil {
		respondError(c, e)
		return
	}
	now := time.Now()
	for i, row := range rows {
		rows[i], e = loadQuotaState(db, row.UsageQuota, now)
		if e != nil {
			respondError(c, e)
			return
		}
		rows[i].UserName = row.UserName
	}
	result(c, rows)
}

func (h *Handler) SaveQuota(c *gin.Context) {
	_, _, admin, e := identity(c.Request.Context())
	if e != nil || !admin {
		respondError(c, ErrForbidden)
		return
	}
	q := UsageQuota{Period: "monthly", Enabled: true, WarningRatio: .8, OveragePolicy: "warn"}
	c.Request.Body = http.MaxBytesReader(c.Writer, c.Request.Body, 8192)
	if c.ShouldBindJSON(&q) != nil {
		respondError(c, ErrInvalid)
		return
	}
	row, e := h.service.saveQuota(c.Request.Context(), c.Param("userID"), q)
	if e != nil {
		respondError(c, e)
		return
	}
	result(c, row)
}

func (h *Handler) DeleteQuota(c *gin.Context) {
	t, _, admin, e := identity(c.Request.Context())
	if e != nil || !admin {
		respondError(c, ErrForbidden)
		return
	}
	q := h.service.db.WithContext(c.Request.Context()).Where("id=? AND tenant_id=?", c.Param("id"), t).Delete(&UsageQuota{})
	if q.Error != nil {
		respondError(c, q.Error)
		return
	}
	if q.RowsAffected != 1 {
		respondError(c, gorm.ErrRecordNotFound)
		return
	}
	result(c, gin.H{"deleted": true})
}