fix: only reduce remain times when request `/v1/chat/completions` (close #15)

BREAKING CHANGE: now remain_times is -1 doesn't mean unlimited times anymore!

controller/relay.go

@@ -7,16 +7,20 @@ import (
 	"io"
 	"net/http"
 	"one-api/common"
+	"one-api/model"
 	"strings"
 )
 
 func Relay(c *gin.Context) {
 	channelType := c.GetInt("channel")
+	tokenId := c.GetInt("token_id")
+	isUnlimitedTimes := c.GetBool("unlimited_times")
 	baseURL := common.ChannelBaseURLs[channelType]
 	if channelType == common.ChannelTypeCustom {
 		baseURL = c.GetString("base_url")
 	}
-	req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, c.Request.URL.String()), c.Request.Body)
+	requestURL := c.Request.URL.String()
+	req, err := http.NewRequest(c.Request.Method, fmt.Sprintf("%s%s", baseURL, requestURL), c.Request.Body)
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{
 			"error": gin.H{
@@ -46,7 +50,19 @@ func Relay(c *gin.Context) {
 		})
 		return
 	}
-	defer resp.Body.Close()
+
+	defer func() {
+		err := req.Body.Close()
+		if err != nil {
+			common.SysError("Error closing request body: " + err.Error())
+		}
+		if !isUnlimitedTimes && requestURL == "/v1/chat/completions" {
+			err := model.DecreaseTokenRemainTimesById(tokenId)
+			if err != nil {
+				common.SysError("Error decreasing token remain times: " + err.Error())
+			}
+		}
+	}()
 	isStream := resp.Header.Get("Content-Type") == "text/event-stream"
 	if isStream {
 		scanner := bufio.NewScanner(resp.Body)

controller/token.go

@@ -93,13 +93,14 @@ func AddToken(c *gin.Context) {
 		return
 	}
 	cleanToken := model.Token{
-		UserId:       c.GetInt("id"),
-		Name:         token.Name,
-		Key:          common.GetUUID(),
-		CreatedTime:  common.GetTimestamp(),
-		AccessedTime: common.GetTimestamp(),
-		ExpiredTime:  token.ExpiredTime,
-		RemainTimes:  token.RemainTimes,
+		UserId:         c.GetInt("id"),
+		Name:           token.Name,
+		Key:            common.GetUUID(),
+		CreatedTime:    common.GetTimestamp(),
+		AccessedTime:   common.GetTimestamp(),
+		ExpiredTime:    token.ExpiredTime,
+		RemainTimes:    token.RemainTimes,
+		UnlimitedTimes: token.UnlimitedTimes,
 	}
 	err = cleanToken.Insert()
 	if err != nil {
@@ -136,6 +137,7 @@ func DeleteToken(c *gin.Context) {
 
 func UpdateToken(c *gin.Context) {
 	userId := c.GetInt("id")
+	statusOnly := c.Query("status_only")
 	token := model.Token{}
 	err := c.ShouldBindJSON(&token)
 	if err != nil {
@@ -161,19 +163,23 @@ func UpdateToken(c *gin.Context) {
 			})
 			return
 		}
-		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes == 0 {
+		if cleanToken.Status == common.TokenStatusExhausted && cleanToken.RemainTimes <= 0 && !cleanToken.UnlimitedTimes {
 			c.JSON(http.StatusOK, gin.H{
 				"success": false,
-				"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数",
+				"message": "令牌可用次数已用尽，无法启用，请先修改令牌剩余次数，或者设置为无限次数",
 			})
 			return
 		}
 	}
-
-	cleanToken.Name = token.Name
-	cleanToken.Status = token.Status
-	cleanToken.ExpiredTime = token.ExpiredTime
-	cleanToken.RemainTimes = token.RemainTimes
+	if statusOnly != "" {
+		cleanToken.Status = token.Status
+	} else {
+		// If you add more fields, please also update token.Update()
+		cleanToken.Name = token.Name
+		cleanToken.ExpiredTime = token.ExpiredTime
+		cleanToken.RemainTimes = token.RemainTimes
+		cleanToken.UnlimitedTimes = token.UnlimitedTimes
+	}
 	err = cleanToken.Update()
 	if err != nil {
 		c.JSON(http.StatusOK, gin.H{

middleware/auth.go

@@ -80,6 +80,8 @@ func TokenAuth() func(c *gin.Context) {
 			return
 		}
 		c.Set("id", token.UserId)
+		c.Set("token_id", token.Id)
+		c.Set("unlimited_times", token.UnlimitedTimes)
 		if len(parts) > 1 {
 			c.Set("channelId", parts[1])
 		}

model/token.go

@@ -3,20 +3,22 @@ package model
 import (
 	"errors"
 	_ "gorm.io/driver/sqlite"
+	"gorm.io/gorm"
 	"one-api/common"
 	"strings"
 )
 
 type Token struct {
-	Id           int    `json:"id"`
-	UserId       int    `json:"user_id"`
-	Key          string `json:"key" gorm:"uniqueIndex"`
-	Status       int    `json:"status" gorm:"default:1"`
-	Name         string `json:"name" gorm:"index" `
-	CreatedTime  int64  `json:"created_time" gorm:"bigint"`
-	AccessedTime int64  `json:"accessed_time" gorm:"bigint"`
-	ExpiredTime  int64  `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
-	RemainTimes  int    `json:"remain_times" gorm:"default:-1"`        // -1 means infinite times
+	Id             int    `json:"id"`
+	UserId         int    `json:"user_id"`
+	Key            string `json:"key" gorm:"uniqueIndex"`
+	Status         int    `json:"status" gorm:"default:1"`
+	Name           string `json:"name" gorm:"index" `
+	CreatedTime    int64  `json:"created_time" gorm:"bigint"`
+	AccessedTime   int64  `json:"accessed_time" gorm:"bigint"`
+	ExpiredTime    int64  `json:"expired_time" gorm:"bigint;default:-1"` // -1 means never expired
+	RemainTimes    int    `json:"remain_times" gorm:"default:0"`
+	UnlimitedTimes bool   `json:"unlimited_times" gorm:"default:false"`
 }
 
 func GetAllUserTokens(userId int, startIdx int, num int) ([]*Token, error) {
@@ -50,14 +52,16 @@ func ValidateUserToken(key string) (token *Token, err error) {
 			}
 			return nil, errors.New("该 token 已过期")
 		}
+		if !token.UnlimitedTimes && token.RemainTimes <= 0 {
+			token.Status = common.TokenStatusExhausted
+			err := token.SelectUpdate()
+			if err != nil {
+				common.SysError("更新 token 状态失败：" + err.Error())
+			}
+			return nil, errors.New("该 token 可用次数已用尽")
+		}
 		go func() {
 			token.AccessedTime = common.GetTimestamp()
-			if token.RemainTimes > 0 {
-				token.RemainTimes--
-				if token.RemainTimes == 0 {
-					token.Status = common.TokenStatusExhausted
-				}
-			}
 			err := token.SelectUpdate()
 			if err != nil {
 				common.SysError("更新 token 失败：" + err.Error())
@@ -84,15 +88,16 @@ func (token *Token) Insert() error {
 	return err
 }
 
+// Update Make sure your token's fields is completed, because this will update non-zero values
 func (token *Token) Update() error {
 	var err error
-	err = DB.Model(token).Updates(token).Error
+	err = DB.Model(token).Select("name", "status", "expired_time", "remain_times", "unlimited_times").Updates(token).Error
 	return err
 }
 
 func (token *Token) SelectUpdate() error {
 	// This can update zero values
-	return DB.Model(token).Select("accessed_time", "remain_times", "status").Updates(token).Error
+	return DB.Model(token).Select("accessed_time", "status").Updates(token).Error
 }
 
 func (token *Token) Delete() error {
@@ -113,3 +118,8 @@ func DeleteTokenById(id int, userId int) (err error) {
 	}
 	return token.Delete()
 }
+
+func DecreaseTokenRemainTimesById(id int) (err error) {
+	err = DB.Model(&Token{}).Where("id = ?", id).Update("remain_times", gorm.Expr("remain_times - ?", 1)).Error
+	return err
+}

web/src/components/TokensTable.js

@@ -79,11 +79,11 @@ const TokensTable = () => {
         break;
       case 'enable':
         data.status = 1;
-        res = await API.put('/api/token/', data);
+        res = await API.put('/api/token/?status_only=true', data);
         break;
       case 'disable':
         data.status = 2;
-        res = await API.put('/api/token/', data);
+        res = await API.put('/api/token/?status_only=true', data);
         break;
     }
     const { success, message } = res.data;
@@ -230,7 +230,7 @@ const TokensTable = () => {
                   <Table.Cell>{token.id}</Table.Cell>
                   <Table.Cell>{token.name ? token.name : '无'}</Table.Cell>
                   <Table.Cell>{renderStatus(token.status)}</Table.Cell>
-                  <Table.Cell>{token.remain_times === -1 ? "无限制" : token.remain_times}</Table.Cell>
+                  <Table.Cell>{token.unlimited_times ? "无限制" : token.remain_times}</Table.Cell>
                   <Table.Cell>{renderTimestamp(token.created_time)}</Table.Cell>
                   <Table.Cell>{renderTimestamp(token.accessed_time)}</Table.Cell>
                   <Table.Cell>{token.expired_time === -1 ? "永不过期" : renderTimestamp(token.expired_time)}</Table.Cell>

web/src/pages/Token/EditToken.js

@@ -10,11 +10,12 @@ const EditToken = () => {
   const [loading, setLoading] = useState(isEdit);
   const originInputs = {
     name: '',
-    remain_times: -1,
-    expired_time: -1
+    remain_times: 0,
+    expired_time: -1,
+    unlimited_times: false,
   };
   const [inputs, setInputs] = useState(originInputs);
-  const { name, remain_times, expired_time } = inputs;
+  const { name, remain_times, expired_time, unlimited_times } = inputs;
 
   const handleInputChange = (e, { name, value }) => {
     setInputs((inputs) => ({ ...inputs, [name]: value }));
@@ -35,6 +36,10 @@ const EditToken = () => {
     }
   };
 
+  const setUnlimitedTimes = () => {
+    setInputs({ ...inputs, unlimited_times: !unlimited_times });
+  }
+
   const loadToken = async () => {
     let res = await API.get(`/api/token/${tokenId}`);
     const { success, message, data } = res.data;
@@ -105,13 +110,17 @@ const EditToken = () => {
             <Form.Input
               label='剩余次数'
               name='remain_times'
-              placeholder={'请输入剩余次数，-1 表示无限制'}
+              placeholder={'请输入剩余次数'}
               onChange={handleInputChange}
               value={remain_times}
               autoComplete='off'
               type='number'
+              disabled={unlimited_times}
             />
           </Form.Field>
+          <Button type={'button'} onClick={() => {
+            setUnlimitedTimes();
+          }}>{unlimited_times ? "取消无限次" : "设置为无限次"}</Button>
           <Form.Field>
             <Form.Input
               label='过期时间'