修复余额不足邮件提醒bug

controller/relay-audio.go

@@ -44,7 +44,7 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		preConsumedQuota = 0
 	}
 	if preConsumedQuota > 0 {
-		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
+		userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
 		if err != nil {
 			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
@@ -99,7 +99,7 @@ func relayAudioHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 		go func() {
 			quota := countTokenText(audioResponse.Text, audioModel)
 			quotaDelta := quota - preConsumedQuota
-			err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
+			err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}

controller/relay-image.go

@@ -147,7 +147,7 @@ func relayImageHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode
 	var textResponse ImageResponse
 	defer func(ctx context.Context) {
 		if consumeQuota {
-			err := model.PostConsumeTokenQuota(tokenId, quota)
+			err := model.PostConsumeTokenQuota(tokenId, userId, quota, 0)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}

controller/relay-mj.go

@@ -359,7 +359,7 @@ func relayMidjourneySubmit(c *gin.Context, relayMode int) *MidjourneyResponse {
 
 	defer func(ctx context.Context) {
 		if consumeQuota {
-			err := model.PostConsumeTokenQuota(tokenId, quota)
+			err := model.PostConsumeTokenQuota(tokenId, userQuota, quota, 0)
 			if err != nil {
 				common.SysError("error consuming token remain quota: " + err.Error())
 			}

controller/relay-text.go

@@ -233,7 +233,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 		//common.LogInfo(c.Request.Context(), fmt.Sprintf("user %d has enough quota %d, trusted and no need to pre-consume", userId, userQuota))
 	}
 	if consumeQuota && preConsumedQuota > 0 {
-		err := model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
+		userQuota, err = model.PreConsumeTokenQuota(tokenId, preConsumedQuota)
 		if err != nil {
 			return errorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
@@ -400,7 +400,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 			if preConsumedQuota != 0 {
 				go func(ctx context.Context) {
 					// return pre-consumed quota
-					err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
+					err := model.PostConsumeTokenQuota(tokenId, userQuota, -preConsumedQuota, 0)
 					if err != nil {
 						common.LogError(ctx, "error return pre-consumed quota: "+err.Error())
 					}
@@ -434,7 +434,7 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 					quota = 0
 				}
 				quotaDelta := quota - preConsumedQuota
-				err := model.PostConsumeTokenQuota(tokenId, quotaDelta)
+				err := model.PostConsumeTokenQuota(tokenId, userQuota, quotaDelta, preConsumedQuota)
 				if err != nil {
 					common.LogError(ctx, "error consuming token remain quota: "+err.Error())
 				}

model/token.go

@@ -178,59 +178,60 @@ func decreaseTokenQuota(id int, quota int) (err error) {
 	return err
 }
 
-func PreConsumeTokenQuota(tokenId int, quota int) (err error) {
+func PreConsumeTokenQuota(tokenId int, quota int) (userQuota int, err error) {
 	if quota < 0 {
-		return errors.New("quota 不能为负数！")
+		return 0, errors.New("quota 不能为负数！")
 	}
 	token, err := GetTokenById(tokenId)
 	if err != nil {
-		return err
+		return 0, err
 	}
 	if !token.UnlimitedQuota && token.RemainQuota < quota {
-		return errors.New("令牌额度不足")
+		return 0, errors.New("令牌额度不足")
 	}
-	userQuota, err := GetUserQuota(token.UserId)
+	userQuota, err = GetUserQuota(token.UserId)
 	if err != nil {
-		return err
+		return 0, err
 	}
 	if userQuota < quota {
-		return errors.New("用户额度不足")
-	}
-	quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-quota < common.QuotaRemindThreshold
-	noMoreQuota := userQuota-quota <= 0
-	if quotaTooLow || noMoreQuota {
-		go func() {
-			email, err := GetUserEmail(token.UserId)
-			if err != nil {
-				common.SysError("failed to fetch user email: " + err.Error())
-			}
-			prompt := "您的额度即将用尽"
-			if noMoreQuota {
-				prompt = "您的额度已用尽"
-			}
-			if email != "" {
-				topUpLink := fmt.Sprintf("%s/topup", common.ServerAddress)
-				err = common.SendEmail(prompt, email,
-					fmt.Sprintf("%s，当前剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota, topUpLink, topUpLink))
-				if err != nil {
-					common.SysError("failed to send email" + err.Error())
-				}
-			}
-		}()
+		return userQuota, errors.New(fmt.Sprintf("用户额度不足，剩余额度为 %d", userQuota))
 	}
 	if !token.UnlimitedQuota {
 		err = DecreaseTokenQuota(tokenId, quota)
 		if err != nil {
-			return err
+			return userQuota, err
 		}
 	}
 	err = DecreaseUserQuota(token.UserId, quota)
-	return err
+	return userQuota, err
 }
 
-func PostConsumeTokenQuota(tokenId int, quota int) (err error) {
+func PostConsumeTokenQuota(tokenId int, userQuota int, quota int, preConsumedQuota int) (err error) {
 	token, err := GetTokenById(tokenId)
+
 	if quota > 0 {
+		quotaTooLow := userQuota >= common.QuotaRemindThreshold && userQuota-(quota+preConsumedQuota) < common.QuotaRemindThreshold
+		noMoreQuota := userQuota-(quota+preConsumedQuota) <= 0
+		if quotaTooLow || noMoreQuota {
+			go func() {
+				email, err := GetUserEmail(token.UserId)
+				if err != nil {
+					common.SysError("failed to fetch user email: " + err.Error())
+				}
+				prompt := "您的额度即将用尽"
+				if noMoreQuota {
+					prompt = "您的额度已用尽"
+				}
+				if email != "" {
+					topUpLink := fmt.Sprintf("%s/topup", common.ServerAddress)
+					err = common.SendEmail(prompt, email,
+						fmt.Sprintf("%s，当前剩余额度为 %d，为了不影响您的使用，请及时充值。<br/>充值链接：<a href='%s'>%s</a>", prompt, userQuota, topUpLink, topUpLink))
+					if err != nil {
+						common.SysError("failed to send email" + err.Error())
+					}
+				}
+			}()
+		}
 		err = DecreaseUserQuota(token.UserId, quota)
 	} else {
 		err = IncreaseUserQuota(token.UserId, -quota)