From 56ddbb842a9d68a30a0228436d39a888f0d8d87b Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Sun, 24 Mar 2024 22:20:41 +0800
Subject: [PATCH] fix: return pre-consumed quota when error happened for audio
 (close #1217)

---
 relay/controller/audio.go | 31 +++++++++++++++++++------------
 1 file changed, 19 insertions(+), 12 deletions(-)

diff --git a/relay/controller/audio.go b/relay/controller/audio.go
index e69c47b3..cd118985 100644
--- a/relay/controller/audio.go
+++ b/relay/controller/audio.go
@@ -83,6 +83,24 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
 			return openai.ErrorWrapper(err, "pre_consume_token_quota_failed", http.StatusForbidden)
 		}
 	}
+	succeed := false
+	defer func() {
+		if succeed {
+			return
+		}
+		if preConsumedQuota > 0 {
+			// we need to roll back the pre-consumed quota
+			defer func(ctx context.Context) {
+				go func() {
+					// negative means add quota back for token & user
+					err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
+					if err != nil {
+						logger.Error(ctx, fmt.Sprintf("error rollback pre-consumed quota: %s", err.Error()))
+					}
+				}()
+			}(c.Request.Context())
+		}
+	}()
 
 	// map model name
 	modelMapping := c.GetString("model_mapping")
@@ -193,20 +211,9 @@ func RelayAudioHelper(c *gin.Context, relayMode int) *relaymodel.ErrorWithStatus
 		resp.Body = io.NopCloser(bytes.NewBuffer(responseBody))
 	}
 	if resp.StatusCode != http.StatusOK {
-		if preConsumedQuota > 0 {
-			// we need to roll back the pre-consumed quota
-			defer func(ctx context.Context) {
-				go func() {
-					// negative means add quota back for token & user
-					err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
-					if err != nil {
-						logger.Error(ctx, fmt.Sprintf("error rollback pre-consumed quota: %s", err.Error()))
-					}
-				}()
-			}(c.Request.Context())
-		}
 		return util.RelayErrorHandler(resp)
 	}
+	succeed = true
 	quotaDelta := quota - preConsumedQuota
 	defer func(ctx context.Context) {
 		go util.PostConsumeQuota(ctx, tokenId, quotaDelta, quota, userId, channelId, modelRatio, groupRatio, audioModel, tokenName)