fix(subscription): preserve wallet overflow snapshots

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
 	})
 }

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)
+}