diff --git a/model/subscription.go b/model/subscription.go index 52a3a381f54a..81d9b0608cc9 100644 --- a/model/subscription.go +++ b/model/subscription.go @@ -273,7 +273,8 @@ type UserSubscription struct { // Downgrade target group on expiry (snapshot from plan; empty = revert to PrevUserGroup) DowngradeGroup string `json:"downgrade_group" gorm:"type:varchar(64);default:''"` - // Whether wallet fallback is allowed after this subscription's quota is exhausted (snapshot from plan) + // Whether wallet fallback is allowed after this subscription's quota is exhausted (snapshot from plan). + // Keep field stable across usage/reset updates; write only touched columns. AllowWalletOverflow bool `json:"allow_wallet_overflow"` CreatedAt int64 `json:"created_at" gorm:"bigint"` @@ -1128,16 +1129,29 @@ func maybeResetUserSubscriptionWithPlanTx(tx *gorm.DB, sub *UserSubscription, pl } if !advanced { if sub.NextResetTime == 0 && next > 0 { - sub.NextResetTime = next sub.LastResetTime = base.Unix() - return tx.Save(sub).Error + sub.NextResetTime = next + return tx.Model(&UserSubscription{}). + Where("id = ?", sub.Id). + Updates(map[string]interface{}{ + "last_reset_time": sub.LastResetTime, + "next_reset_time": sub.NextResetTime, + "updated_at": common.GetTimestamp(), + }).Error } return nil } sub.AmountUsed = 0 sub.LastResetTime = base.Unix() sub.NextResetTime = next - return tx.Save(sub).Error + return tx.Model(&UserSubscription{}). + Where("id = ?", sub.Id). + Updates(map[string]interface{}{ + "amount_used": sub.AmountUsed, + "last_reset_time": sub.LastResetTime, + "next_reset_time": sub.NextResetTime, + "updated_at": common.GetTimestamp(), + }).Error } // PreConsumeUserSubscription pre-consumes from any active subscription total quota. @@ -1226,7 +1240,12 @@ func PreConsumeUserSubscription(requestId string, userId int, modelName string, return err } sub.AmountUsed += amount - if err := tx.Save(&sub).Error; err != nil { + if err := tx.Model(&UserSubscription{}). + Where("id = ?", sub.Id). + Updates(map[string]interface{}{ + "amount_used": sub.AmountUsed, + "updated_at": common.GetTimestamp(), + }).Error; err != nil { return err } returnValue.UserSubscriptionId = sub.Id @@ -1374,7 +1393,11 @@ func PostConsumeUserSubscriptionDelta(userSubscriptionId int, delta int64) error if sub.AmountTotal > 0 && newUsed > sub.AmountTotal { return fmt.Errorf("subscription used exceeds total, used=%d total=%d", newUsed, sub.AmountTotal) } - sub.AmountUsed = newUsed - return tx.Save(&sub).Error + return tx.Model(&UserSubscription{}). + Where("id = ?", sub.Id). + Updates(map[string]interface{}{ + "amount_used": newUsed, + "updated_at": common.GetTimestamp(), + }).Error }) } diff --git a/model/subscription_test.go b/model/subscription_test.go new file mode 100644 index 000000000000..18a766abce59 --- /dev/null +++ b/model/subscription_test.go @@ -0,0 +1,117 @@ +package model + +import ( + "database/sql" + "fmt" + "path/filepath" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupSubscriptionTestDB(t *testing.T) func() { + t.Helper() + + oldDB := DB + oldLogDB := LOG_DB + oldUsingSQLite := common.UsingSQLite + oldUsingPostgreSQL := common.UsingPostgreSQL + oldUsingMySQL := common.UsingMySQL + oldRedisEnabled := common.RedisEnabled + + testDB, err := gorm.Open(sqlite.Open(filepath.Join(t.TempDir(), "subscription-test.db")), &gorm.Config{}) + require.NoError(t, err) + + DB = testDB + LOG_DB = testDB + common.UsingSQLite = true + common.UsingPostgreSQL = false + common.UsingMySQL = false + common.RedisEnabled = false + initCol() + + return func() { + DB = oldDB + LOG_DB = oldLogDB + common.UsingSQLite = oldUsingSQLite + common.UsingPostgreSQL = oldUsingPostgreSQL + common.UsingMySQL = oldUsingMySQL + common.RedisEnabled = oldRedisEnabled + initCol() + } +} + +func TestSubscriptionUsageUpdatesDoNotOverwriteWalletOverflowSnapshot(t *testing.T) { + cleanup := setupSubscriptionTestDB(t) + defer cleanup() + + require.NoError(t, DB.AutoMigrate(&SubscriptionPlan{}, &UserSubscription{}, &SubscriptionPreConsumeRecord{})) + + plan := &SubscriptionPlan{ + Title: "Reset Snapshot", + PriceAmount: 9.9, + Currency: "USD", + DurationUnit: SubscriptionDurationMonth, + DurationValue: 1, + Enabled: true, + TotalAmount: 100, + QuotaResetPeriod: SubscriptionResetDaily, + } + require.NoError(t, DB.Create(plan).Error) + + now := GetDBTimestamp() + sub := &UserSubscription{ + UserId: 1, + PlanId: plan.Id, + AmountTotal: 100, + AmountUsed: 60, + StartTime: now - 86400, + EndTime: now + 86400, + Status: "active", + LastResetTime: now - 86400, + NextResetTime: now - 1, + CreatedAt: common.GetTimestamp(), + UpdatedAt: common.GetTimestamp(), + } + require.NoError(t, DB.Create(sub).Error) + require.NoError(t, DB.Model(&UserSubscription{}). + Where("id = ?", sub.Id). + Update("allow_wallet_overflow", gorm.Expr("NULL")).Error) + + require.NoError(t, DB.Transaction(func(tx *gorm.DB) error { + var locked UserSubscription + if err := tx.Set("gorm:query_option", "FOR UPDATE"). + First(&locked, sub.Id).Error; err != nil { + return err + } + return maybeResetUserSubscriptionWithPlanTx(tx, &locked, plan, GetDBTimestamp()) + })) + + var updated UserSubscription + require.NoError(t, DB.First(&updated, sub.Id).Error) + require.EqualValues(t, 0, updated.AmountUsed) + requireWalletOverflowSnapshotNull(t, sub.Id) + + _, err := PreConsumeUserSubscription(fmt.Sprintf("req-reset-%s", t.Name()), 1, "gpt-4o", 0, 10) + require.NoError(t, err) + require.NoError(t, DB.First(&updated, sub.Id).Error) + require.EqualValues(t, 10, updated.AmountUsed) + requireWalletOverflowSnapshotNull(t, sub.Id) + + require.NoError(t, PostConsumeUserSubscriptionDelta(sub.Id, -5)) + require.NoError(t, DB.First(&updated, sub.Id).Error) + require.EqualValues(t, 5, updated.AmountUsed) + requireWalletOverflowSnapshotNull(t, sub.Id) +} + +func requireWalletOverflowSnapshotNull(t *testing.T, userSubscriptionId int) { + t.Helper() + + var overflowSnapshot sql.NullBool + require.NoError(t, DB.Raw("SELECT allow_wallet_overflow FROM user_subscriptions WHERE id = ?", userSubscriptionId). + Scan(&overflowSnapshot).Error) + require.False(t, overflowSnapshot.Valid) +}