tt/ttstd-infinite-canvas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: settle excess token usage with negative balance (#304)

Browse Source
This commit is contained in:
MasterChiefCN
2026-08-25 10:28:12 +08:00
committed by GitHub
parent a23e399ccf
commit c2a05db80f
3 changed files with 184 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

52
backend/internal/repository/finance.go
View File

@@ -21,7 +21,6 @@ var (
ErrTaskNotRetryable = errors.New("task is not retryable")
ErrBillingStateConflict = errors.New("billing state conflict")
ErrBillingUsageUnavailable = errors.New("billing usage unavailable")
ErrBillingUnderreserved = errors.New("billing amount exceeds reservation")
ErrChannelModelInUse = errors.New("channel model is in use")
)
@@ -604,7 +603,10 @@ func (r *Repository) MarkBillingUncertain(id string, errorText string) error {
}
func (r *Repository) SettleBillingOrder(id string, providerRequestID string) error {
return r.db.Transaction(func(tx *gorm.DB) error {
var observedUsage *BillingUsage
var observedActual int64
observedActualAvailable := false
err := r.db.Transaction(func(tx *gorm.DB) error {
var order model.BillingOrder
if err := tx.First(&order, "id = ?", id).Error; err != nil {
return err
@@ -620,6 +622,7 @@ func (r *Repository) SettleBillingOrder(id string, providerRequestID string) err
if err != nil {
return err
}
observedUsage = usage
reserved := order.ReservedAmountMicrocredits
if reserved <= 0 {
reserved = order.AmountMicrocredits
@@ -628,14 +631,14 @@ func (r *Repository) SettleBillingOrder(id string, providerRequestID string) err
if err != nil {
return err
}
if actual > reserved {
return ErrBillingUnderreserved
}
refund := reserved - actual
observedActual = actual
observedActualAvailable = true
refund := max(reserved-actual, int64(0))
supplement := max(actual-reserved, int64(0))
updated := tx.Model(&model.CreditAccount{}).
Where("user_id = ? AND reserved_microcredits >= ?", order.UserID, reserved).
Updates(map[string]any{
"available_microcredits": gorm.Expr("available_microcredits + ?", refund),
"available_microcredits": gorm.Expr("available_microcredits + ?", refund-supplement),
"reserved_microcredits": gorm.Expr("reserved_microcredits - ?", reserved),
"version": gorm.Expr("version + 1"), "updated_at": time.Now(),
})
@@ -660,10 +663,14 @@ func (r *Repository) SettleBillingOrder(id string, providerRequestID string) err
if err := tx.Model(&order).Updates(updates).Error; err != nil {
return err
}
consumeNote := ""
if supplement > 0 {
consumeNote = "Token 实际用量超过预授权，已补扣差额"
}
if err := tx.Create(&model.CreditLedgerEntry{ID: newRepositoryID(), UserID: order.UserID, Type: model.CreditLedgerConsume,
AmountMicrocredits: -actual, ReservedDeltaMicrocredits: -reserved,
AmountMicrocredits: -actual, AvailableDeltaMicrocredits: -supplement, ReservedDeltaMicrocredits: -reserved,
AvailableAfterMicrocredits: account.AvailableMicrocredits, ReservedAfterMicrocredits: account.ReservedMicrocredits,
BillingOrderID: order.ID, Model: order.Model, ChannelID: order.ChannelID, Scene: order.Scene}).Error; err != nil {
BillingOrderID: order.ID, Model: order.Model, ChannelID: order.ChannelID, Scene: order.Scene, Note: consumeNote}).Error; err != nil {
return err
}
if refund > 0 {
@@ -715,6 +722,26 @@ func (r *Repository) SettleBillingOrder(id string, providerRequestID string) err
Scene: order.Scene,
}).Error
})
if err != nil && observedUsage != nil {
// usage 是上游已经确认的事实；即使结算因账户状态异常回滚，也要保留给用户和管理员核对。
updates := map[string]any{
"input_tokens": observedUsage.InputTokens, "output_tokens": observedUsage.OutputTokens,
"cached_tokens": observedUsage.CachedTokens, "usage_available": true, "updated_at": time.Now(),
}
if observedActualAvailable {
updates["actual_amount_microcredits"] = observedActual
}
if providerRequestID != "" {
updates["provider_request_id"] = providerRequestID
}
usageErr := r.db.Model(&model.BillingOrder{}).
Where("id = ? AND status NOT IN ?", id, []model.BillingStatus{model.BillingStatusSettled, model.BillingStatusRefunded}).
Updates(updates).Error
if usageErr != nil {
return errors.Join(err, usageErr)
}
}
return err
}
func zeroPricedTokenOrder(order model.BillingOrder) bool {
@@ -822,8 +849,11 @@ func (r *Repository) AdjustCredits(userID string, actorUserID string, amount int
if err := tx.Clauses(clause.OnConflict{DoNothing: true}).Create(&account).Error; err != nil {
return err
}
updated := tx.Model(&model.CreditAccount{}).
Where("user_id = ? AND available_microcredits + ? >= 0", userID, amount).
accountQuery := tx.Model(&model.CreditAccount{}).Where("user_id = ?", userID)
if amount < 0 {
accountQuery = accountQuery.Where("available_microcredits + ? >= 0", amount)
}
updated := accountQuery.
Updates(map[string]any{
"available_microcredits": gorm.Expr("available_microcredits + ?", amount),
"version": gorm.Expr("version + 1"),

141
backend/internal/repository/finance_token_test.go
View File

@@ -99,3 +99,144 @@ func TestSettleArkVideoTokenOrderFromPollUsage(t *testing.T) {
t.Fatalf("settled account = %#v", account)
}
}
func TestSettleArkVideoTokenOrderSupplementsUnderreservation(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:finance-token-supplement?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
if err := db.AutoMigrate(&model.CreditAccount{}, &model.BillingOrder{}, &model.ApiCallLog{}, &model.CreditLedgerEntry{}); err != nil {
t.Fatal(err)
}
const (
reserved = int64(3_049_738)
actual = int64(3_115_222)
supplement = actual - reserved
)
if err := db.Create(&model.CreditAccount{UserID: "user-1", AvailableMicrocredits: 1_000_000, ReservedMicrocredits: reserved}).Error; err != nil {
t.Fatal(err)
}
if err := db.Create(&model.BillingOrder{
ID: "order-1", UserID: "user-1", IdempotencyKey: "task:task-1", Capability: "video", BillingMode: "token",
AmountMicrocredits: reserved, ReservedAmountMicrocredits: reserved, OutputTokenPriceMicrocredits: 18_200_000,
MultiplierBasisPoints: 10_000, Status: model.BillingStatusRunning,
}).Error; err != nil {
t.Fatal(err)
}
if err := db.Create(&model.ApiCallLog{
ID: "poll-log-1", BillingOrderID: "order-1", RequestKind: "poll", Billable: false,
Status: model.ApiCallStatusSucceeded, UsageAvailable: true, OutputTokens: 171_166,
}).Error; err != nil {
t.Fatal(err)
}
repo := &Repository{db: db}
if err := repo.SettleBillingOrder("order-1", "ark-task-1"); err != nil {
t.Fatalf("SettleBillingOrder() error = %v", err)
}
var order model.BillingOrder
if err := db.First(&order, "id = ?", "order-1").Error; err != nil {
t.Fatal(err)
}
if order.Status != model.BillingStatusSettled || order.ActualAmountMicrocredits != actual || order.OutputTokens != 171_166 || !order.UsageAvailable {
t.Fatalf("settled order = %#v", order)
}
var account model.CreditAccount
if err := db.First(&account, "user_id = ?", "user-1").Error; err != nil {
t.Fatal(err)
}
if account.AvailableMicrocredits != 1_000_000-supplement || account.ReservedMicrocredits != 0 {
t.Fatalf("settled account = %#v", account)
}
var entry model.CreditLedgerEntry
if err := db.First(&entry, "billing_order_id = ? AND type = ?", "order-1", model.CreditLedgerConsume).Error; err != nil {
t.Fatal(err)
}
if entry.AmountMicrocredits != -actual || entry.AvailableDeltaMicrocredits != -supplement || entry.ReservedDeltaMicrocredits != -reserved {
t.Fatalf("consume entry = %#v", entry)
}
}
func TestSettleArkVideoTokenOrderAllowsNegativeBalance(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:finance-token-negative-balance?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
if err := db.AutoMigrate(&model.CreditAccount{}, &model.BillingOrder{}, &model.ApiCallLog{}, &model.CreditLedgerEntry{}); err != nil {
t.Fatal(err)
}
const reserved = int64(3_049_738)
if err := db.Create(&model.CreditAccount{UserID: "user-1", AvailableMicrocredits: 10_000, ReservedMicrocredits: reserved}).Error; err != nil {
t.Fatal(err)
}
if err := db.Create(&model.BillingOrder{
ID: "order-1", UserID: "user-1", IdempotencyKey: "task:task-1", Capability: "video", BillingMode: "token",
AmountMicrocredits: reserved, ReservedAmountMicrocredits: reserved, OutputTokenPriceMicrocredits: 18_200_000,
MultiplierBasisPoints: 10_000, Status: model.BillingStatusRunning,
}).Error; err != nil {
t.Fatal(err)
}
if err := db.Create(&model.ApiCallLog{
ID: "poll-log-1", BillingOrderID: "order-1", RequestKind: "poll", Billable: false,
Status: model.ApiCallStatusSucceeded, UsageAvailable: true, OutputTokens: 171_166,
}).Error; err != nil {
t.Fatal(err)
}
repo := &Repository{db: db}
err = repo.SettleBillingOrder("order-1", "ark-task-1")
if err != nil {
t.Fatalf("SettleBillingOrder() error = %v", err)
}
var order model.BillingOrder
if err := db.First(&order, "id = ?", "order-1").Error; err != nil {
t.Fatal(err)
}
if order.Status != model.BillingStatusSettled || order.ActualAmountMicrocredits != 3_115_222 || order.OutputTokens != 171_166 || !order.UsageAvailable {
t.Fatalf("settled order = %#v", order)
}
var account model.CreditAccount
if err := db.First(&account, "user_id = ?", "user-1").Error; err != nil {
t.Fatal(err)
}
if account.AvailableMicrocredits != -55_484 || account.ReservedMicrocredits != 0 {
t.Fatalf("settled account = %#v", account)
}
var entry model.CreditLedgerEntry
if err := db.First(&entry, "billing_order_id = ? AND type = ?", "order-1", model.CreditLedgerConsume).Error; err != nil {
t.Fatal(err)
}
if entry.AvailableAfterMicrocredits != -55_484 || entry.AvailableDeltaMicrocredits != -65_484 {
t.Fatalf("consume entry = %#v", entry)
}
}
func TestNegativeBalanceBlocksNewReservationsAndAcceptsRepayment(t *testing.T) {
db, err := gorm.Open(sqlite.Open("file:finance-negative-balance-repayment?mode=memory&cache=shared"), &gorm.Config{})
if err != nil {
t.Fatal(err)
}
if err := db.AutoMigrate(&model.CreditAccount{}, &model.BillingOrder{}, &model.CreditLedgerEntry{}); err != nil {
t.Fatal(err)
}
if err := db.Create(&model.CreditAccount{UserID: "user-1", AvailableMicrocredits: -1_000_000}).Error; err != nil {
t.Fatal(err)
}
repo := &Repository{db: db}
order := model.BillingOrder{
ID: "order-1", UserID: "user-1", IdempotencyKey: "task:task-1",
AmountMicrocredits: 1, ReservedAmountMicrocredits: 1, Status: model.BillingStatusReserved,
}
if err := repo.ReserveBillingOrder(&order); !errors.Is(err, ErrInsufficientCredits) {
t.Fatalf("ReserveBillingOrder() error = %v", err)
}
account, err := repo.AdjustCredits("user-1", "admin-1", 100_000_000, "充值入账")
if err != nil {
t.Fatalf("AdjustCredits() error = %v", err)
}
if account.AvailableMicrocredits != 99_000_000 {
t.Fatalf("available balance = %d", account.AvailableMicrocredits)
}
}