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>
This commit is contained in:
@@ -3,6 +3,8 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"strconv"
|
||||
@@ -13,6 +15,8 @@ import (
|
||||
"github.com/QuantumNous/new-api/constant"
|
||||
"github.com/QuantumNous/new-api/dto"
|
||||
"github.com/QuantumNous/new-api/logger"
|
||||
"github.com/QuantumNous/new-api/model"
|
||||
relaycommon "github.com/QuantumNous/new-api/relay/common"
|
||||
relayconstant "github.com/QuantumNous/new-api/relay/constant"
|
||||
"github.com/QuantumNous/new-api/setting"
|
||||
|
||||
@@ -27,6 +31,113 @@ func CovertMjpActionToModelName(mjAction string) string {
|
||||
return modelName
|
||||
}
|
||||
|
||||
// PrepareMidjourneyTaskBilling sets the durable refund marker before the task is inserted.
|
||||
func PrepareMidjourneyTaskBilling(relayInfo *relaycommon.RelayInfo, task *model.Midjourney, quota int, shouldBill bool) (bool, error) {
|
||||
if task == nil {
|
||||
return false, errors.New("Midjourney task is nil")
|
||||
}
|
||||
task.Quota = 0
|
||||
task.TokenId = 0
|
||||
task.BillingChannelId = 0
|
||||
if !shouldBill {
|
||||
return false, nil
|
||||
}
|
||||
if relayInfo == nil {
|
||||
return false, errors.New("relay info is nil")
|
||||
}
|
||||
if quota < 0 {
|
||||
return false, errors.New("quota cannot be negative")
|
||||
}
|
||||
if relayInfo.BillingSource == BillingSourceSubscription {
|
||||
return false, errors.New("legacy Midjourney billing does not support subscriptions")
|
||||
}
|
||||
|
||||
task.Quota = quota
|
||||
task.BillingChannelId = task.ChannelId
|
||||
if relayInfo.ChannelMeta != nil && relayInfo.ChannelId > 0 {
|
||||
task.BillingChannelId = relayInfo.ChannelId
|
||||
}
|
||||
return true, nil
|
||||
}
|
||||
|
||||
// SettleMidjourneyTaskBilling charges a persisted legacy task and records the applied stages.
|
||||
func SettleMidjourneyTaskBilling(relayInfo *relaycommon.RelayInfo, task *model.Midjourney, prepared bool) (bool, error) {
|
||||
if !prepared {
|
||||
return false, nil
|
||||
}
|
||||
if relayInfo == nil {
|
||||
return false, errors.New("relay info is nil")
|
||||
}
|
||||
if task == nil || task.Id == 0 {
|
||||
return false, errors.New("Midjourney task must be persisted before billing")
|
||||
}
|
||||
|
||||
result, billingErr := postConsumeQuotaWithResult(relayInfo, task.Quota, 0, true)
|
||||
if !result.FundingApplied {
|
||||
task.Quota = 0
|
||||
task.TokenId = 0
|
||||
task.BillingChannelId = 0
|
||||
if updateErr := task.UpdateBillingState(); updateErr != nil {
|
||||
return false, errors.Join(billingErr, fmt.Errorf("clear Midjourney billing state: %w", updateErr))
|
||||
}
|
||||
return false, billingErr
|
||||
}
|
||||
|
||||
task.TokenId = 0
|
||||
if result.TokenApplied {
|
||||
task.TokenId = relayInfo.TokenId
|
||||
}
|
||||
if updateErr := task.UpdateBillingState(); updateErr != nil {
|
||||
return true, errors.Join(billingErr, fmt.Errorf("update Midjourney billing state: %w", updateErr))
|
||||
}
|
||||
return true, billingErr
|
||||
}
|
||||
|
||||
// RefundMidjourneyQuota reverses every accounting element recorded for a billed legacy task.
|
||||
func RefundMidjourneyQuota(ctx context.Context, task *model.Midjourney, reason string) bool {
|
||||
quota := task.Quota
|
||||
if quota == 0 {
|
||||
return true
|
||||
}
|
||||
|
||||
if err := model.IncreaseUserQuota(task.UserId, quota, false); err != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 用户额度失败 task %s: %s", task.MjId, err.Error()))
|
||||
return false
|
||||
}
|
||||
|
||||
if task.TokenId > 0 {
|
||||
tokenKey := resolveTokenKey(ctx, task.TokenId, task.MjId)
|
||||
if tokenKey != "" {
|
||||
if err := model.IncreaseTokenQuota(task.TokenId, tokenKey, quota); err != nil {
|
||||
logger.LogWarn(ctx, fmt.Sprintf("退还 Midjourney 令牌额度失败 task %s: %s", task.MjId, err.Error()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
billingChannelId := task.GetBillingChannelId()
|
||||
model.UpdateUserUsedQuota(task.UserId, -quota)
|
||||
model.UpdateChannelUsedQuota(billingChannelId, -quota)
|
||||
model.RecordTaskBillingLog(model.RecordTaskBillingLogParams{
|
||||
UserId: task.UserId,
|
||||
LogType: model.LogTypeRefund,
|
||||
Content: "",
|
||||
ChannelId: billingChannelId,
|
||||
ModelName: CovertMjpActionToModelName(task.Action),
|
||||
Quota: quota,
|
||||
TokenId: task.TokenId,
|
||||
Other: map[string]interface{}{
|
||||
"task_id": task.MjId,
|
||||
"reason": reason,
|
||||
},
|
||||
})
|
||||
|
||||
task.Quota = 0
|
||||
if err := task.UpdateBillingState(); err != nil {
|
||||
logger.LogError(ctx, fmt.Sprintf("Midjourney 退款成功但清除 quota 失败 task %s: %s", task.MjId, err.Error()))
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
||||
func GetMjRequestModel(relayMode int, midjRequest *dto.MidjourneyRequest) (string, *dto.MidjourneyResponse, bool) {
|
||||
action := ""
|
||||
if relayMode == relayconstant.RelayModeMidjourneyAction {
|
||||
|
||||
+18
-6
@@ -406,17 +406,27 @@ func PreConsumeTokenQuota(relayInfo *relaycommon.RelayInfo, quota int) error {
|
||||
return nil
|
||||
}
|
||||
|
||||
func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int, sendEmail bool) (err error) {
|
||||
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 errors.New("subscription id is missing")
|
||||
return result, errors.New("subscription id is missing")
|
||||
}
|
||||
delta := int64(quota)
|
||||
if delta != 0 {
|
||||
if err := model.PostConsumeUserSubscriptionDelta(relayInfo.SubscriptionId, delta); err != nil {
|
||||
return err
|
||||
return result, err
|
||||
}
|
||||
relayInfo.SubscriptionPostDelta += delta
|
||||
}
|
||||
@@ -428,9 +438,10 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu
|
||||
err = model.IncreaseUserQuota(relayInfo.UserId, -quota, false)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
return result, err
|
||||
}
|
||||
}
|
||||
result.FundingApplied = true
|
||||
|
||||
if !relayInfo.IsPlayground {
|
||||
if quota > 0 {
|
||||
@@ -439,8 +450,9 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu
|
||||
err = model.IncreaseTokenQuota(relayInfo.TokenId, relayInfo.TokenKey, -quota)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
return result, err
|
||||
}
|
||||
result.TokenApplied = true
|
||||
}
|
||||
|
||||
if sendEmail {
|
||||
@@ -449,7 +461,7 @@ func PostConsumeQuota(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQu
|
||||
}
|
||||
}
|
||||
|
||||
return nil
|
||||
return result, nil
|
||||
}
|
||||
|
||||
func checkAndSendQuotaNotify(relayInfo *relaycommon.RelayInfo, quota int, preConsumedQuota int) {
|
||||
|
||||
+11
-5
@@ -161,7 +161,7 @@ func taskModelName(task *model.Task) string {
|
||||
}
|
||||
|
||||
// RefundTaskQuota 统一的任务失败退款逻辑。
|
||||
// 当异步任务失败时,将预扣的 quota 退还给用户(支持钱包和订阅),并退还令牌额度。
|
||||
// 当异步任务失败时,退还资金与令牌额度,并回减用户和渠道用量。
|
||||
// 返回资金来源是否已成功退还;失败时保留 quota,供显式重试或人工对账。
|
||||
func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool {
|
||||
quota := task.Quota
|
||||
@@ -178,7 +178,11 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
|
||||
// 2. 退还令牌额度
|
||||
taskAdjustTokenQuota(ctx, task, -quota)
|
||||
|
||||
// 3. 记录日志
|
||||
// 3. 回减预扣时累计的用户和渠道用量,请求次数保持不变
|
||||
model.UpdateUserUsedQuota(task.UserId, -quota)
|
||||
model.UpdateChannelUsedQuota(task.ChannelId, -quota)
|
||||
|
||||
// 4. 记录日志
|
||||
other := taskBillingOther(task)
|
||||
other["task_id"] = task.TaskID
|
||||
other["reason"] = reason
|
||||
@@ -194,7 +198,7 @@ func RefundTaskQuota(ctx context.Context, task *model.Task, reason string) bool
|
||||
Other: other,
|
||||
})
|
||||
|
||||
// 4. 资金退款完成后再清除持久化标记。
|
||||
// 5. 资金退款完成后再清除持久化标记。
|
||||
// 回写失败必须显式告警,避免漏掉潜在的重复退款风险。
|
||||
task.Quota = 0
|
||||
if err := task.UpdateQuota(); err != nil {
|
||||
@@ -242,13 +246,15 @@ func RecalculateTaskQuota(ctx context.Context, task *model.Task, actualQuota int
|
||||
logger.LogError(ctx, fmt.Sprintf("差额结算回写 quota 失败 task %s: %s", task.TaskID, err.Error()))
|
||||
}
|
||||
|
||||
// 提交阶段已经累计过一次请求;结算阶段只调整最终用量。
|
||||
model.UpdateUserUsedQuota(task.UserId, quotaDelta)
|
||||
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
|
||||
|
||||
var logType int
|
||||
var logQuota int
|
||||
if quotaDelta > 0 {
|
||||
logType = model.LogTypeConsume
|
||||
logQuota = quotaDelta
|
||||
model.UpdateUserUsedQuotaAndRequestCount(task.UserId, quotaDelta)
|
||||
model.UpdateChannelUsedQuota(task.ChannelId, quotaDelta)
|
||||
} else {
|
||||
logType = model.LogTypeRefund
|
||||
logQuota = -quotaDelta
|
||||
|
||||
@@ -45,6 +45,7 @@ func TestMain(m *testing.M) {
|
||||
&model.Token{},
|
||||
&model.Log{},
|
||||
&model.Channel{},
|
||||
&model.Midjourney{},
|
||||
&model.TopUp{},
|
||||
&model.UserSubscription{},
|
||||
&model.SystemTask{},
|
||||
@@ -68,6 +69,7 @@ func truncate(t *testing.T) {
|
||||
model.DB.Exec("DELETE FROM tokens")
|
||||
model.DB.Exec("DELETE FROM logs")
|
||||
model.DB.Exec("DELETE FROM channels")
|
||||
model.DB.Exec("DELETE FROM midjourneys")
|
||||
model.DB.Exec("DELETE FROM top_ups")
|
||||
model.DB.Exec("DELETE FROM user_subscriptions")
|
||||
model.DB.Exec("DELETE FROM system_task_locks")
|
||||
@@ -115,6 +117,20 @@ func seedChannel(t *testing.T, id int) {
|
||||
require.NoError(t, model.DB.Create(ch).Error)
|
||||
}
|
||||
|
||||
func seedChargedAccounting(t *testing.T, userID, channelID, tokenID, quota, requestCount int) {
|
||||
t.Helper()
|
||||
require.NoError(t, model.DB.Model(&model.User{}).Where("id = ?", userID).Updates(map[string]any{
|
||||
"used_quota": quota,
|
||||
"request_count": requestCount,
|
||||
}).Error)
|
||||
require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", channelID).
|
||||
Update("used_quota", quota).Error)
|
||||
if tokenID > 0 {
|
||||
require.NoError(t, model.DB.Model(&model.Token{}).Where("id = ?", tokenID).
|
||||
Update("used_quota", quota).Error)
|
||||
}
|
||||
}
|
||||
|
||||
func makeTask(userId, channelId, quota, tokenId int, billingSource string, subscriptionId int) *model.Task {
|
||||
return &model.Task{
|
||||
TaskID: "task_" + time.Now().Format("150405.000"),
|
||||
@@ -249,6 +265,20 @@ func getUserQuota(t *testing.T, id int) int {
|
||||
return user.Quota
|
||||
}
|
||||
|
||||
func getUserUsageAccounting(t *testing.T, id int) (int, int) {
|
||||
t.Helper()
|
||||
var user model.User
|
||||
require.NoError(t, model.DB.Select("used_quota", "request_count").Where("id = ?", id).First(&user).Error)
|
||||
return user.UsedQuota, user.RequestCount
|
||||
}
|
||||
|
||||
func getChannelUsedQuota(t *testing.T, id int) int64 {
|
||||
t.Helper()
|
||||
var channel model.Channel
|
||||
require.NoError(t, model.DB.Select("used_quota").Where("id = ?", id).First(&channel).Error)
|
||||
return channel.UsedQuota
|
||||
}
|
||||
|
||||
func getTokenRemainQuota(t *testing.T, id int) int {
|
||||
t.Helper()
|
||||
var token model.Token
|
||||
@@ -277,6 +307,13 @@ func getTaskQuota(t *testing.T, id int64) int {
|
||||
return task.Quota
|
||||
}
|
||||
|
||||
func getMidjourneyTask(t *testing.T, id int) model.Midjourney {
|
||||
t.Helper()
|
||||
var task model.Midjourney
|
||||
require.NoError(t, model.DB.First(&task, id).Error)
|
||||
return task
|
||||
}
|
||||
|
||||
func getLastLog(t *testing.T) *model.Log {
|
||||
t.Helper()
|
||||
var log model.Log
|
||||
@@ -294,6 +331,291 @@ func countLogs(t *testing.T) int64 {
|
||||
return count
|
||||
}
|
||||
|
||||
// ===========================================================================
|
||||
// Legacy Midjourney billing tests
|
||||
// ===========================================================================
|
||||
|
||||
func TestPrepareMidjourneyTaskBillingKeepsUnbilledMarkerClear(t *testing.T) {
|
||||
task := &model.Midjourney{Quota: 900, TokenId: 7, BillingChannelId: 8}
|
||||
|
||||
prepared, err := PrepareMidjourneyTaskBilling(&relaycommon.RelayInfo{}, task, 900, false)
|
||||
|
||||
require.NoError(t, err)
|
||||
assert.False(t, prepared)
|
||||
assert.Zero(t, task.Quota)
|
||||
assert.Zero(t, task.TokenId)
|
||||
assert.Zero(t, task.BillingChannelId)
|
||||
}
|
||||
|
||||
func TestSettleMidjourneyTaskBillingRequiresPersistedTask(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 49, 49, 49
|
||||
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
||||
seedUser(t, userID, initialUserQuota)
|
||||
seedToken(t, tokenID, userID, "sk-midjourney-unpersisted", initialTokenQuota)
|
||||
seedChannel(t, channelID)
|
||||
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
UserId: userID,
|
||||
TokenId: tokenID,
|
||||
TokenKey: "sk-midjourney-unpersisted",
|
||||
UserQuota: initialUserQuota,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelId: channelID,
|
||||
},
|
||||
}
|
||||
task := &model.Midjourney{UserId: userID, ChannelId: channelID}
|
||||
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
||||
require.NoError(t, err)
|
||||
require.True(t, prepared)
|
||||
|
||||
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
||||
|
||||
require.Error(t, err)
|
||||
assert.False(t, billed)
|
||||
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
||||
}
|
||||
|
||||
func TestMidjourneyRefundRestoresEveryAccountingElementOnBillingChannel(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID, tokenID, billingChannelID, executionChannelID = 50, 50, 50, 51
|
||||
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
||||
seedUser(t, userID, initialUserQuota)
|
||||
seedToken(t, tokenID, userID, "sk-midjourney", initialTokenQuota)
|
||||
seedChannel(t, billingChannelID)
|
||||
seedChannel(t, executionChannelID)
|
||||
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
UserId: userID,
|
||||
TokenId: tokenID,
|
||||
TokenKey: "sk-midjourney",
|
||||
UserQuota: initialUserQuota,
|
||||
UsingGroup: "default",
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelId: billingChannelID,
|
||||
},
|
||||
}
|
||||
task := &model.Midjourney{
|
||||
UserId: userID,
|
||||
Action: "IMAGINE",
|
||||
MjId: "mj-accounting-refund",
|
||||
ChannelId: executionChannelID,
|
||||
Progress: "0%",
|
||||
}
|
||||
|
||||
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
||||
require.NoError(t, err)
|
||||
require.True(t, prepared)
|
||||
assert.Equal(t, chargedQuota, task.Quota)
|
||||
assert.Zero(t, task.TokenId)
|
||||
assert.Equal(t, billingChannelID, task.BillingChannelId)
|
||||
require.NoError(t, task.Insert())
|
||||
|
||||
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
||||
require.NoError(t, err)
|
||||
require.True(t, billed)
|
||||
assert.Equal(t, initialUserQuota-chargedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota-chargedQuota, getTokenRemainQuota(t, tokenID))
|
||||
persisted := getMidjourneyTask(t, task.Id)
|
||||
assert.Equal(t, chargedQuota, persisted.Quota)
|
||||
assert.Equal(t, tokenID, persisted.TokenId)
|
||||
assert.Equal(t, billingChannelID, persisted.BillingChannelId)
|
||||
|
||||
seedChargedAccounting(t, userID, billingChannelID, tokenID, chargedQuota, 1)
|
||||
|
||||
assert.True(t, RefundMidjourneyQuota(ctx, task, "构图失败"))
|
||||
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
||||
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, billingChannelID))
|
||||
assert.Zero(t, getChannelUsedQuota(t, executionChannelID))
|
||||
|
||||
persisted = getMidjourneyTask(t, task.Id)
|
||||
assert.Zero(t, persisted.Quota)
|
||||
assert.Equal(t, tokenID, persisted.TokenId)
|
||||
assert.Equal(t, billingChannelID, persisted.BillingChannelId)
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, model.LogTypeRefund, log.Type)
|
||||
assert.Equal(t, chargedQuota, log.Quota)
|
||||
assert.Equal(t, tokenID, log.TokenId)
|
||||
assert.Equal(t, billingChannelID, log.ChannelId)
|
||||
|
||||
assert.True(t, RefundMidjourneyQuota(ctx, task, "duplicate poll"))
|
||||
assert.Equal(t, int64(1), countLogs(t))
|
||||
}
|
||||
|
||||
func TestSettleMidjourneyTaskBillingFundingFailureClearsMarkers(t *testing.T) {
|
||||
truncate(t)
|
||||
|
||||
const userID, tokenID, channelID = 52, 52, 52
|
||||
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
||||
seedUser(t, userID, initialUserQuota)
|
||||
seedToken(t, tokenID, userID, "sk-midjourney-funding-failure", initialTokenQuota)
|
||||
seedChannel(t, channelID)
|
||||
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
UserId: userID,
|
||||
TokenId: tokenID,
|
||||
TokenKey: "sk-midjourney-funding-failure",
|
||||
UserQuota: initialUserQuota,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelId: channelID,
|
||||
},
|
||||
}
|
||||
task := &model.Midjourney{UserId: userID, MjId: "mj-funding-failure", ChannelId: channelID}
|
||||
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
||||
require.NoError(t, err)
|
||||
require.True(t, prepared)
|
||||
require.NoError(t, task.Insert())
|
||||
|
||||
require.NoError(t, model.DB.Exec(`
|
||||
CREATE TRIGGER fail_midjourney_user_update
|
||||
BEFORE UPDATE ON users
|
||||
WHEN OLD.id = 52
|
||||
BEGIN
|
||||
SELECT RAISE(ABORT, 'forced user quota failure');
|
||||
END;
|
||||
`).Error)
|
||||
t.Cleanup(func() {
|
||||
model.DB.Exec("DROP TRIGGER IF EXISTS fail_midjourney_user_update")
|
||||
})
|
||||
|
||||
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
||||
|
||||
require.Error(t, err)
|
||||
assert.False(t, billed)
|
||||
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
||||
persisted := getMidjourneyTask(t, task.Id)
|
||||
assert.Zero(t, persisted.Quota)
|
||||
assert.Zero(t, persisted.TokenId)
|
||||
assert.Zero(t, persisted.BillingChannelId)
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Zero(t, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
assert.Zero(t, countLogs(t))
|
||||
}
|
||||
|
||||
func TestSettleMidjourneyTaskBillingTokenFailureKeepsFundingRefundable(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID, tokenID, channelID = 53, 53, 53
|
||||
const initialUserQuota, initialTokenQuota, chargedQuota = 10000, 5000, 3000
|
||||
seedUser(t, userID, initialUserQuota)
|
||||
seedToken(t, tokenID, userID, "sk-midjourney-token-failure", initialTokenQuota)
|
||||
seedChannel(t, channelID)
|
||||
|
||||
relayInfo := &relaycommon.RelayInfo{
|
||||
UserId: userID,
|
||||
TokenId: tokenID,
|
||||
TokenKey: "sk-midjourney-token-failure",
|
||||
UserQuota: initialUserQuota,
|
||||
ChannelMeta: &relaycommon.ChannelMeta{
|
||||
ChannelId: channelID,
|
||||
},
|
||||
}
|
||||
task := &model.Midjourney{UserId: userID, MjId: "mj-token-failure", ChannelId: channelID}
|
||||
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, chargedQuota, true)
|
||||
require.NoError(t, err)
|
||||
require.True(t, prepared)
|
||||
require.NoError(t, task.Insert())
|
||||
|
||||
require.NoError(t, model.DB.Exec(`
|
||||
CREATE TRIGGER fail_midjourney_token_update
|
||||
BEFORE UPDATE ON tokens
|
||||
WHEN OLD.id = 53
|
||||
BEGIN
|
||||
SELECT RAISE(ABORT, 'forced token quota failure');
|
||||
END;
|
||||
`).Error)
|
||||
t.Cleanup(func() {
|
||||
model.DB.Exec("DROP TRIGGER IF EXISTS fail_midjourney_token_update")
|
||||
})
|
||||
|
||||
billed, err := SettleMidjourneyTaskBilling(relayInfo, task, prepared)
|
||||
|
||||
require.Error(t, err)
|
||||
require.True(t, billed)
|
||||
assert.Equal(t, initialUserQuota-chargedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
||||
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
||||
persisted := getMidjourneyTask(t, task.Id)
|
||||
assert.Equal(t, chargedQuota, persisted.Quota)
|
||||
assert.Zero(t, persisted.TokenId)
|
||||
assert.Equal(t, channelID, persisted.BillingChannelId)
|
||||
|
||||
seedChargedAccounting(t, userID, channelID, 0, chargedQuota, 1)
|
||||
assert.True(t, RefundMidjourneyQuota(ctx, task, "token settlement failed"))
|
||||
assert.Equal(t, initialUserQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, initialTokenQuota, getTokenRemainQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Zero(t, log.TokenId)
|
||||
}
|
||||
|
||||
func TestPrepareMidjourneyTaskBillingRejectsSubscriptionBeforeCharge(t *testing.T) {
|
||||
task := &model.Midjourney{Quota: 900, TokenId: 7, BillingChannelId: 8}
|
||||
relayInfo := &relaycommon.RelayInfo{BillingSource: BillingSourceSubscription, SubscriptionId: 1}
|
||||
|
||||
prepared, err := PrepareMidjourneyTaskBilling(relayInfo, task, 900, true)
|
||||
|
||||
require.Error(t, err)
|
||||
assert.False(t, prepared)
|
||||
assert.Zero(t, task.Quota)
|
||||
assert.Zero(t, task.TokenId)
|
||||
assert.Zero(t, task.BillingChannelId)
|
||||
}
|
||||
|
||||
func TestRefundMidjourneyQuotaUsesLegacyChannelFallbackWithoutTokenAdjustment(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID, tokenID, channelID = 54, 54, 54
|
||||
const walletAfterCharge, tokenQuota, chargedQuota = 7000, 5000, 3000
|
||||
seedUser(t, userID, walletAfterCharge)
|
||||
seedToken(t, tokenID, userID, "sk-midjourney-legacy", tokenQuota)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, 0, chargedQuota, 1)
|
||||
task := &model.Midjourney{
|
||||
UserId: userID,
|
||||
MjId: "mj-legacy-fallback",
|
||||
Action: "IMAGINE",
|
||||
ChannelId: channelID,
|
||||
Quota: chargedQuota,
|
||||
TokenId: 0,
|
||||
Progress: "0%",
|
||||
}
|
||||
require.NoError(t, task.Insert())
|
||||
|
||||
assert.True(t, RefundMidjourneyQuota(ctx, task, "legacy failure"))
|
||||
|
||||
assert.Equal(t, walletAfterCharge+chargedQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenQuota, getTokenRemainQuota(t, tokenID))
|
||||
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
assert.Equal(t, channelID, log.ChannelId)
|
||||
assert.Zero(t, log.TokenId)
|
||||
}
|
||||
|
||||
// ===========================================================================
|
||||
// RefundTaskQuota tests
|
||||
// ===========================================================================
|
||||
@@ -309,6 +631,7 @@ func TestRefundTaskQuota_Wallet(t *testing.T) {
|
||||
seedUser(t, userID, initQuota)
|
||||
seedToken(t, tokenID, userID, "sk-test-key", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
@@ -320,7 +643,11 @@ func TestRefundTaskQuota_Wallet(t *testing.T) {
|
||||
|
||||
// Token remain_quota should increase, used_quota should decrease
|
||||
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, -preConsumed, getTokenUsedQuota(t, tokenID))
|
||||
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
|
||||
// A refund log should be created
|
||||
log := getLastLog(t)
|
||||
@@ -345,6 +672,7 @@ func TestRefundTaskQuota_Subscription(t *testing.T) {
|
||||
seedToken(t, tokenID, userID, "sk-sub-key", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedSubscription(t, subID, userID, subTotal, subUsed)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
@@ -356,6 +684,11 @@ func TestRefundTaskQuota_Subscription(t *testing.T) {
|
||||
|
||||
// Token should also be refunded
|
||||
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
||||
assert.Zero(t, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
@@ -390,6 +723,7 @@ func TestRefundTaskQuota_NoToken(t *testing.T) {
|
||||
|
||||
seedUser(t, userID, initQuota)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, 0, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, 0, BillingSourceWallet, 0) // TokenId=0
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
@@ -398,6 +732,10 @@ func TestRefundTaskQuota_NoToken(t *testing.T) {
|
||||
|
||||
// User quota refunded
|
||||
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
|
||||
// Log created
|
||||
log := getLastLog(t)
|
||||
@@ -406,19 +744,26 @@ func TestRefundTaskQuota_NoToken(t *testing.T) {
|
||||
assert.Zero(t, getTaskQuota(t, task.ID))
|
||||
}
|
||||
|
||||
func TestRefundTaskQuota_FundingFailureKeepsPendingMarker(t *testing.T) {
|
||||
func TestRefundTaskQuota_FundingFailureKeepsAccountingAndPendingMarker(t *testing.T) {
|
||||
truncate(t)
|
||||
ctx := context.Background()
|
||||
|
||||
const userID, preConsumed = 5, 1200
|
||||
const userID, channelID, preConsumed = 5, 5, 1200
|
||||
seedUser(t, userID, 5000)
|
||||
task := makeTask(userID, 0, preConsumed, 0, BillingSourceSubscription, 9999)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, 0, preConsumed, 1)
|
||||
task := makeTask(userID, channelID, preConsumed, 0, BillingSourceSubscription, 9999)
|
||||
task.Status = model.TaskStatusFailure
|
||||
require.NoError(t, model.DB.Create(task).Error)
|
||||
|
||||
assert.False(t, RefundTaskQuota(ctx, task, "subscription missing"))
|
||||
assert.Equal(t, 5000, getUserQuota(t, userID))
|
||||
assert.Equal(t, preConsumed, task.Quota)
|
||||
assert.Equal(t, preConsumed, getTaskQuota(t, task.ID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Equal(t, preConsumed, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Equal(t, int64(preConsumed), getChannelUsedQuota(t, channelID))
|
||||
assert.Equal(t, int64(0), countLogs(t))
|
||||
}
|
||||
|
||||
@@ -438,6 +783,7 @@ func TestRecalculate_PositiveDelta(t *testing.T) {
|
||||
seedUser(t, userID, initQuota)
|
||||
seedToken(t, tokenID, userID, "sk-recalc-pos", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
|
||||
@@ -448,6 +794,11 @@ func TestRecalculate_PositiveDelta(t *testing.T) {
|
||||
|
||||
// Token should also be charged the delta
|
||||
assert.Equal(t, tokenRemain-(actualQuota-preConsumed), getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, actualQuota, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Equal(t, actualQuota, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
||||
|
||||
// task.Quota should be updated to actualQuota
|
||||
assert.Equal(t, actualQuota, task.Quota)
|
||||
@@ -471,6 +822,7 @@ func TestRecalculate_NegativeDelta(t *testing.T) {
|
||||
seedUser(t, userID, initQuota)
|
||||
seedToken(t, tokenID, userID, "sk-recalc-neg", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
|
||||
@@ -481,6 +833,11 @@ func TestRecalculate_NegativeDelta(t *testing.T) {
|
||||
|
||||
// Token should be refunded the difference
|
||||
assert.Equal(t, tokenRemain+(preConsumed-actualQuota), getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, actualQuota, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Equal(t, actualQuota, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
||||
|
||||
// task.Quota updated
|
||||
assert.Equal(t, actualQuota, task.Quota)
|
||||
@@ -544,6 +901,7 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) {
|
||||
seedToken(t, tokenID, userID, "sk-sub-recalc", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedSubscription(t, subID, userID, subTotal, subUsed)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceSubscription, subID)
|
||||
|
||||
@@ -554,6 +912,11 @@ func TestRecalculate_Subscription_NegativeDelta(t *testing.T) {
|
||||
|
||||
// Token refunded
|
||||
assert.Equal(t, tokenRemain+(preConsumed-actualQuota), getTokenRemainQuota(t, tokenID))
|
||||
assert.Equal(t, actualQuota, getTokenUsedQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Equal(t, actualQuota, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
||||
|
||||
assert.Equal(t, actualQuota, task.Quota)
|
||||
|
||||
@@ -627,6 +990,7 @@ func TestCASGuardedRefund_Win(t *testing.T) {
|
||||
seedUser(t, userID, initQuota)
|
||||
seedToken(t, tokenID, userID, "sk-cas-refund-win", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
task.Status = model.TaskStatus(model.TaskStatusInProgress)
|
||||
@@ -643,6 +1007,10 @@ func TestCASGuardedRefund_Win(t *testing.T) {
|
||||
// Refund should have happened
|
||||
assert.Equal(t, initQuota+preConsumed, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+preConsumed, getTokenRemainQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Zero(t, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Zero(t, getChannelUsedQuota(t, channelID))
|
||||
|
||||
log := getLastLog(t)
|
||||
require.NotNil(t, log)
|
||||
@@ -660,6 +1028,7 @@ func TestCASGuardedRefund_Lose(t *testing.T) {
|
||||
seedUser(t, userID, initQuota)
|
||||
seedToken(t, tokenID, userID, "sk-cas-refund-lose", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
// Create task with IN_PROGRESS in DB
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
@@ -676,6 +1045,10 @@ func TestCASGuardedRefund_Lose(t *testing.T) {
|
||||
// CAS lost: user quota should NOT change (no double refund)
|
||||
assert.Equal(t, initQuota, getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain, getTokenRemainQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Equal(t, preConsumed, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Equal(t, int64(preConsumed), getChannelUsedQuota(t, channelID))
|
||||
|
||||
// No billing log should be created
|
||||
assert.Equal(t, int64(0), countLogs(t))
|
||||
@@ -693,6 +1066,7 @@ func TestCASGuardedSettle_Win(t *testing.T) {
|
||||
seedUser(t, userID, initQuota)
|
||||
seedToken(t, tokenID, userID, "sk-cas-settle-win", tokenRemain)
|
||||
seedChannel(t, channelID)
|
||||
seedChargedAccounting(t, userID, channelID, tokenID, preConsumed, 1)
|
||||
|
||||
task := makeTask(userID, channelID, preConsumed, tokenID, BillingSourceWallet, 0)
|
||||
task.Status = model.TaskStatus(model.TaskStatusInProgress)
|
||||
@@ -708,6 +1082,10 @@ func TestCASGuardedSettle_Win(t *testing.T) {
|
||||
// Settlement should refund the over-charge (5000 - 3000 = 2000 back to user)
|
||||
assert.Equal(t, initQuota+(preConsumed-actualQuota), getUserQuota(t, userID))
|
||||
assert.Equal(t, tokenRemain+(preConsumed-actualQuota), getTokenRemainQuota(t, tokenID))
|
||||
usedQuota, requestCount := getUserUsageAccounting(t, userID)
|
||||
assert.Equal(t, actualQuota, usedQuota)
|
||||
assert.Equal(t, 1, requestCount)
|
||||
assert.Equal(t, int64(actualQuota), getChannelUsedQuota(t, channelID))
|
||||
|
||||
// task.Quota should be updated to actualQuota
|
||||
assert.Equal(t, actualQuota, task.Quota)
|
||||
|
||||
Reference in New Issue
Block a user