Files
new-api/service/quota.go
T
wans10andCaIon 58d4e9bd3b fix(billing): 异步任务退款时同步减少 used_quota (#6795)
* fix(billing): 异步任务退款时同步减少 used_quota
退款时仅恢复了 quota(剩余额度),但未同步减少 used_quota(已用额度),
导致"总额度"(quota + used_quota)随退款次数持续虚增,超出用户实际充值金额。

修复三处退款路径:
- RefundTaskQuota:任务失败完整退款
- RecalculateTaskQuota:差额结算退款分支
- controller/midjourney.go:Midjourney 任务失败退款

新增 model.UpdateUserUsedQuota 公开函数,仅调整 used_quota 不影响 request_count。

* fix(billing): 任务退款时同步扣减渠道 used_quota

* fix(billing): complete async task refund accounting

* style(model): group internal Midjourney fields

---------

Co-authored-by: CaIon <i@caion.me>
2026-08-13 22:06:40 +08:00

562 lines
20 KiB
Go
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
package service
import (
"errors"
"fmt"
"math"
"strings"
"time"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/logger"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/billingexpr"
perfmetrics "github.com/QuantumNous/new-api/pkg/perf_metrics"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/ratio_setting"
"github.com/QuantumNous/new-api/types"
"github.com/bytedance/gopkg/util/gopool"
"github.com/gin-gonic/gin"
"github.com/shopspring/decimal"
)
type TokenDetails struct {
TextTokens int
AudioTokens int
}
type QuotaInfo struct {
InputDetails TokenDetails
OutputDetails TokenDetails
ModelName string
UsePrice bool
ModelPrice float64
ModelRatio float64
GroupRatio float64
}
func hasCustomModelRatio(modelName string, currentRatio float64) bool {
defaultRatio, exists := ratio_setting.GetDefaultModelRatioMap()[modelName]
if !exists {
return true
}
return currentRatio != defaultRatio
}
func calculateAudioQuota(info QuotaInfo) (int, *common.QuotaClamp) {
if info.UsePrice {
modelPrice := decimal.NewFromFloat(info.ModelPrice)
quotaPerUnit := decimal.NewFromFloat(common.QuotaPerUnit)
groupRatio := decimal.NewFromFloat(info.GroupRatio)
quota := modelPrice.Mul(quotaPerUnit).Mul(groupRatio)
return common.QuotaFromDecimalChecked(quota)
}
completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(info.ModelName))
audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(info.ModelName))
audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(info.ModelName))
groupRatio := decimal.NewFromFloat(info.GroupRatio)
modelRatio := decimal.NewFromFloat(info.ModelRatio)
ratio := groupRatio.Mul(modelRatio)
inputTextTokens := decimal.NewFromInt(int64(info.InputDetails.TextTokens))
outputTextTokens := decimal.NewFromInt(int64(info.OutputDetails.TextTokens))
inputAudioTokens := decimal.NewFromInt(int64(info.InputDetails.AudioTokens))
outputAudioTokens := decimal.NewFromInt(int64(info.OutputDetails.AudioTokens))
quota := decimal.Zero
quota = quota.Add(inputTextTokens)
quota = quota.Add(outputTextTokens.Mul(completionRatio))
quota = quota.Add(inputAudioTokens.Mul(audioRatio))
quota = quota.Add(outputAudioTokens.Mul(audioRatio).Mul(audioCompletionRatio))
quota = quota.Mul(ratio)
// If ratio is not zero and quota is less than or equal to zero, set quota to 1
if !ratio.IsZero() && quota.LessThanOrEqual(decimal.Zero) {
quota = decimal.NewFromInt(1)
}
return common.QuotaFromDecimalChecked(quota)
}
func PreWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.RealtimeUsage) error {
if relayInfo.UsePrice {
return nil
}
userQuota, err := model.GetUserQuota(relayInfo.UserId, false)
if err != nil {
return err
}
token, err := model.GetTokenByKey(strings.TrimPrefix(relayInfo.TokenKey, "sk-"), false)
if err != nil {
return err
}
modelName := relayInfo.OriginModelName
textInputTokens := usage.InputTokenDetails.TextTokens
textOutTokens := usage.OutputTokenDetails.TextTokens
audioInputTokens := usage.InputTokenDetails.AudioTokens
audioOutTokens := usage.OutputTokenDetails.AudioTokens
groupRatio := ratio_setting.GetGroupRatio(relayInfo.UsingGroup)
modelRatio, _, _ := ratio_setting.GetModelRatio(modelName)
autoGroup, exists := common.GetContextKey(ctx, constant.ContextKeyAutoGroup)
if exists {
groupRatio = ratio_setting.GetGroupRatio(autoGroup.(string))
logger.LogDebug(ctx, "final group ratio: %f", groupRatio)
relayInfo.UsingGroup = autoGroup.(string)
}
actualGroupRatio := groupRatio
userGroupRatio, ok := ratio_setting.GetGroupGroupRatio(relayInfo.UserGroup, relayInfo.UsingGroup)
if ok {
actualGroupRatio = userGroupRatio
}
quotaInfo := QuotaInfo{
InputDetails: TokenDetails{
TextTokens: textInputTokens,
AudioTokens: audioInputTokens,
},
OutputDetails: TokenDetails{
TextTokens: textOutTokens,
AudioTokens: audioOutTokens,
},
ModelName: modelName,
UsePrice: relayInfo.UsePrice,
ModelRatio: modelRatio,
GroupRatio: actualGroupRatio,
}
quota, clamp := calculateAudioQuota(quotaInfo)
noteQuotaClamp(relayInfo, clamp)
if userQuota < quota {
return fmt.Errorf("user quota is not enough, user quota: %s, need quota: %s", logger.FormatQuota(userQuota), logger.FormatQuota(quota))
}
if !token.UnlimitedQuota && token.RemainQuota < quota {
return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(token.RemainQuota), logger.FormatQuota(quota))
}
err = PostConsumeQuota(relayInfo, quota, 0, false)
if err != nil {
return err
}
logger.LogInfo(ctx, "realtime streaming consume quota success, quota: "+fmt.Sprintf("%d", quota))
return nil
}
func PostWssConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, modelName string,
usage *dto.RealtimeUsage, extraContent string) {
var tieredResult *billingexpr.TieredResult
tieredOk, tieredQuota, tieredRes := TryTieredSettle(relayInfo, billingexpr.TokenParams{
P: float64(usage.InputTokens),
C: float64(usage.OutputTokens),
Len: float64(usage.InputTokens),
})
if tieredOk {
tieredResult = tieredRes
}
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
textInputTokens := usage.InputTokenDetails.TextTokens
textOutTokens := usage.OutputTokenDetails.TextTokens
audioInputTokens := usage.InputTokenDetails.AudioTokens
audioOutTokens := usage.OutputTokenDetails.AudioTokens
tokenName := ctx.GetString("token_name")
completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(modelName))
audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(modelName))
modelRatio := relayInfo.PriceData.ModelRatio
groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
modelPrice := relayInfo.PriceData.ModelPrice
usePrice := relayInfo.PriceData.UsePrice
quotaInfo := QuotaInfo{
InputDetails: TokenDetails{
TextTokens: textInputTokens,
AudioTokens: audioInputTokens,
},
OutputDetails: TokenDetails{
TextTokens: textOutTokens,
AudioTokens: audioOutTokens,
},
ModelName: modelName,
UsePrice: usePrice,
ModelRatio: modelRatio,
GroupRatio: groupRatio,
}
quota, clamp := calculateAudioQuota(quotaInfo)
noteQuotaClamp(relayInfo, clamp)
if tieredOk {
quota = tieredQuota
}
totalTokens := usage.TotalTokens
var logContent string
if !usePrice {
logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
} else {
logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
}
// record all the consume log even if quota is 0
if totalTokens == 0 {
// in this case, must be some error happened
// we cannot just return, because we may have to return the pre-consumed quota
quota = 0
logContent += "(可能是上游超时)"
logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
"tokenId %d, model %s pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, modelName, relayInfo.FinalPreConsumedQuota))
} else {
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
}
if err := SettleBilling(ctx, relayInfo, quota); err != nil {
logger.LogError(ctx, "error settling billing: "+err.Error())
}
logModel := modelName
if extraContent != "" {
logContent += ", " + extraContent
}
other := GenerateWssOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
if tieredResult != nil {
InjectTieredBillingInfo(other, relayInfo, tieredResult)
}
attachQuotaSaturation(ctx, relayInfo, other)
model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
ChannelId: relayInfo.ChannelId,
PromptTokens: usage.InputTokens,
CompletionTokens: usage.OutputTokens,
ModelName: logModel,
TokenName: tokenName,
Quota: quota,
Content: logContent,
TokenId: relayInfo.TokenId,
UseTimeSeconds: int(useTimeSeconds),
IsStream: relayInfo.IsStream,
Group: relayInfo.UsingGroup,
Other: other,
})
}
func CalcOpenRouterCacheCreateTokens(usage dto.Usage, priceData types.PriceData) int {
if priceData.CacheCreationRatio == 1 {
return 0
}
quotaPrice := priceData.ModelRatio / common.QuotaPerUnit
promptCacheCreatePrice := quotaPrice * priceData.CacheCreationRatio
promptCacheReadPrice := quotaPrice * priceData.CacheRatio
completionPrice := quotaPrice * priceData.CompletionRatio
cost, _ := usage.Cost.(float64)
totalPromptTokens := float64(usage.PromptTokens)
completionTokens := float64(usage.CompletionTokens)
promptCacheReadTokens := float64(usage.PromptTokensDetails.CachedTokens)
return int(math.Round((cost -
totalPromptTokens*quotaPrice +
promptCacheReadTokens*(quotaPrice-promptCacheReadPrice) -
completionTokens*completionPrice) /
(promptCacheCreatePrice - quotaPrice)))
}
func PostAudioConsumeQuota(ctx *gin.Context, relayInfo *relaycommon.RelayInfo, usage *dto.Usage, extraContent string) {
var tieredUsedVars map[string]bool
if snap := relayInfo.TieredBillingSnapshot; snap != nil {
tieredUsedVars = billingexpr.UsedVars(snap.ExprString)
}
var tieredResult *billingexpr.TieredResult
tieredOk, tieredQuota, tieredRes := TryTieredSettle(relayInfo, BuildTieredTokenParams(usage, false, tieredUsedVars))
if tieredOk {
tieredResult = tieredRes
}
useTimeSeconds := time.Now().Unix() - relayInfo.StartTime.Unix()
textInputTokens := usage.PromptTokensDetails.TextTokens
textOutTokens := usage.CompletionTokenDetails.TextTokens
audioInputTokens := usage.PromptTokensDetails.AudioTokens
audioOutTokens := usage.CompletionTokenDetails.AudioTokens
tokenName := ctx.GetString("token_name")
completionRatio := decimal.NewFromFloat(ratio_setting.GetCompletionRatio(relayInfo.OriginModelName))
audioRatio := decimal.NewFromFloat(ratio_setting.GetAudioRatio(relayInfo.OriginModelName))
audioCompletionRatio := decimal.NewFromFloat(ratio_setting.GetAudioCompletionRatio(relayInfo.OriginModelName))
modelRatio := relayInfo.PriceData.ModelRatio
groupRatio := relayInfo.PriceData.GroupRatioInfo.GroupRatio
modelPrice := relayInfo.PriceData.ModelPrice
usePrice := relayInfo.PriceData.UsePrice
quotaInfo := QuotaInfo{
InputDetails: TokenDetails{
TextTokens: textInputTokens,
AudioTokens: audioInputTokens,
},
OutputDetails: TokenDetails{
TextTokens: textOutTokens,
AudioTokens: audioOutTokens,
},
ModelName: relayInfo.OriginModelName,
UsePrice: usePrice,
ModelRatio: modelRatio,
GroupRatio: groupRatio,
}
quota, clamp := calculateAudioQuota(quotaInfo)
noteQuotaClamp(relayInfo, clamp)
if tieredOk {
quota = tieredQuota
}
totalTokens := usage.TotalTokens
var logContent string
if !usePrice {
logContent = fmt.Sprintf("模型倍率 %.2f,补全倍率 %.2f,音频倍率 %.2f,音频补全倍率 %.2f,分组倍率 %.2f",
modelRatio, completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), groupRatio)
} else {
logContent = fmt.Sprintf("模型价格 %.2f,分组倍率 %.2f", modelPrice, groupRatio)
}
// record all the consume log even if quota is 0
if totalTokens == 0 {
// in this case, must be some error happened
// we cannot just return, because we may have to return the pre-consumed quota
quota = 0
logContent += "(可能是上游超时)"
logger.LogError(ctx, fmt.Sprintf("total tokens is 0, cannot consume quota, userId %d, channelId %d, "+
"tokenId %d, model %s pre-consumed quota %d", relayInfo.UserId, relayInfo.ChannelId, relayInfo.TokenId, relayInfo.OriginModelName, relayInfo.FinalPreConsumedQuota))
} else {
model.UpdateUserUsedQuotaAndRequestCount(relayInfo.UserId, quota)
model.UpdateChannelUsedQuota(relayInfo.ChannelId, quota)
}
if err := SettleBilling(ctx, relayInfo, quota); err != nil {
logger.LogError(ctx, "error settling billing: "+err.Error())
}
logModel := relayInfo.OriginModelName
if extraContent != "" {
logContent += ", " + extraContent
}
other := GenerateAudioOtherInfo(ctx, relayInfo, usage, modelRatio, groupRatio,
completionRatio.InexactFloat64(), audioRatio.InexactFloat64(), audioCompletionRatio.InexactFloat64(), modelPrice, relayInfo.PriceData.GroupRatioInfo.GroupSpecialRatio)
if tieredResult != nil {
InjectTieredBillingInfo(other, relayInfo, tieredResult)
}
attachQuotaSaturation(ctx, relayInfo, other)
model.RecordConsumeLog(ctx, relayInfo.UserId, model.RecordConsumeLogParams{
ChannelId: relayInfo.ChannelId,
PromptTokens: usage.PromptTokens,
CompletionTokens: usage.CompletionTokens,
ModelName: logModel,
TokenName: tokenName,
Quota: quota,
Content: logContent,
TokenId: relayInfo.TokenId,
UseTimeSeconds: int(useTimeSeconds),
IsStream: relayInfo.IsStream,
Group: relayInfo.UsingGroup,
Other: other,
})
gopool.Go(func() {
perfmetrics.RecordRelaySample(relayInfo, true, int64(usage.CompletionTokens))
})
}
func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
if quota < 0 {
return errors.New("quota 不能为负数!")
}
if relayInfo.IsPlayground {
return nil
}
// 原子预扣:检查与扣减在同一操作中完成,并发请求不可能同时通过检查后超扣。
reserved, err := model.TryReserveTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota, relayInfo.TokenUnlimited)
if err != nil {
return err
}
if !reserved {
remainQuota := 0
if token, tokenErr := model.GetTokenByKey(relayInfo.TokenKey, false); tokenErr == nil && token != nil {
remainQuota = token.RemainQuota
}
return fmt.Errorf("token quota is not enough, token remain quota: %s, need quota: %s", logger.FormatQuota(remainQuota), logger.FormatQuota(quota))
}
return nil
}
type postConsumeQuotaResult struct {
FundingApplied bool
TokenApplied bool
}
func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) error {
_, err := postConsumeQuotaWithResult(relayInfo, quota, preConsumedQuota, sendEmail)
return err
}
func postConsumeQuotaWithResult(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (result postConsumeQuotaResult, err error) {
// 1) Consume from wallet quota OR subscription item
if relayInfo != nil && relayInfo.BillingSource == BillingSourceSubscription {
if relayInfo.SubscriptionId == 0 {
return result, errors.New("subscription id is missing")
}
delta := int64(quota)
if delta != 0 {
if err := model.PostConsumeUserSubscriptionDelta(relayInfo.SubscriptionId, delta); err != nil {
return result, err
}
relayInfo.SubscriptionPostDelta += delta
}
} else {
// Wallet
if quota > 0 {
err = model.DecreaseUserQuota(relayInfo.UserId, quota, false)
} else {
err = model.IncreaseUserQuota(relayInfo.UserId, -quota, false)
}
if err != nil {
return result, err
}
}
result.FundingApplied = true
if !relayInfo.IsPlayground {
if quota > 0 {
err = model.DecreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, quota)
} else {
err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
}
if err != nil {
return result, err
}
result.TokenApplied = true
}
if sendEmail {
if (quota + preConsumedQuota) != 0 {
checkAndSendQuotaNotify(relayInfo, quota, preConsumedQuota)
}
}
return result, nil
}
func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) {
gopool.Go(func() {
userSetting := relayInfo.UserSetting
threshold := common.QuotaRemindThreshold
if userSetting.QuotaWarningThreshold != 0 {
threshold = int(userSetting.QuotaWarningThreshold)
}
//noMoreQuota := userCache.Quota-(quota+preConsumedQuota) <= 0
quotaTooLow := false
consumeQuota := quota + preConsumedQuota
if relayInfo.UserQuota-consumeQuota < threshold {
quotaTooLow = true
}
if quotaTooLow {
prompt := "您的额度即将用尽"
topUpLink := PaymentReturnURL("/wallet")
// 根据通知方式生成不同的内容格式
var content string
var values []interface{}
notifyType := userSetting.NotifyType
if notifyType == "" {
notifyType = dto.NotifyTypeEmail
}
if notifyType == dto.NotifyTypeBark {
// Bark推送使用简短文本,不支持HTML
content = "{{value}},剩余额度:{{value}},请及时充值"
values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota)}
} else if notifyType == dto.NotifyTypeGotify {
content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota)}
} else {
// 默认内容格式,适用于Email和Webhook(支持HTML
content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
values = []interface{}{prompt, logger.FormatQuota(relayInfo.UserQuota), topUpLink, topUpLink}
}
err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values))
if err != nil {
common.SysError(fmt.Sprintf("failed to send quota notify to user %d: %s", relayInfo.UserId, err.Error()))
}
}
})
}
func checkAndSendSubscriptionQuotaNotify(relayInfo *relaycommon.RelayInfo) {
gopool.Go(func() {
if relayInfo == nil {
return
}
if relayInfo.SubscriptionId == 0 || relayInfo.SubscriptionAmountTotal <= 0 {
return
}
userSetting := relayInfo.UserSetting
threshold := common.QuotaRemindThreshold
if userSetting.QuotaWarningThreshold != 0 {
threshold = int(userSetting.QuotaWarningThreshold)
}
usedAfter := relayInfo.SubscriptionAmountUsedAfterPreConsume + relayInfo.SubscriptionPostDelta
remaining := relayInfo.SubscriptionAmountTotal - usedAfter
if remaining >= int64(threshold) {
return
}
prompt := "您的订阅额度即将用尽"
topUpLink := PaymentReturnURL("/wallet")
var content string
var values []interface{}
notifyType := userSetting.NotifyType
if notifyType == "" {
notifyType = dto.NotifyTypeEmail
}
if notifyType == dto.NotifyTypeBark {
content = "{{value}},剩余额度:{{value}},请及时充值"
values = []interface{}{prompt, logger.FormatQuota(int(remaining))}
} else if notifyType == dto.NotifyTypeGotify {
content = "{{value}},当前剩余额度为 {{value}},请及时充值。"
values = []interface{}{prompt, logger.FormatQuota(int(remaining))}
} else {
content = "{{value}},当前剩余额度为 {{value}},为了不影响您的使用,请及时充值。<br/>充值链接:<a href='{{value}}'>{{value}}</a>"
values = []interface{}{prompt, logger.FormatQuota(int(remaining)), topUpLink, topUpLink}
}
if err := NotifyUser(relayInfo.UserId, relayInfo.UserEmail, relayInfo.UserSetting, dto.NewNotify(dto.NotifyTypeQuotaExceed, prompt, content, values)); err != nil {
common.SysError(fmt.Sprintf("failed to send subscription quota notify to user %d: %s", relayInfo.UserId, err.Error()))
}
})
}