From 217584bd439e0d22a36fda16bf4602ac7c0728ff Mon Sep 17 00:00:00 2001
From: JustSong <songquanpeng@foxmail.com>
Date: Wed, 13 Sep 2023 21:50:45 +0800
Subject: [PATCH] fix: fix quota not return when error occurred (close #518)

---
 controller/relay-text.go | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/controller/relay-text.go b/controller/relay-text.go
index 8a01b810..d3e2a2c1 100644
--- a/controller/relay-text.go
+++ b/controller/relay-text.go
@@ -347,6 +347,13 @@ func relayTextHelper(c *gin.Context, relayMode int) *OpenAIErrorWithStatusCode {
 		isStream = isStream || strings.HasPrefix(resp.Header.Get("Content-Type"), "text/event-stream")
 
 		if resp.StatusCode != http.StatusOK {
+			go func() {
+				// return pre-consumed quota
+				err := model.PostConsumeTokenQuota(tokenId, -preConsumedQuota)
+				if err != nil {
+					common.SysError("error return pre-consumed quota: " + err.Error())
+				}
+			}()
 			return relayErrorHandler(resp)
 		}
 	}