diff --git a/controller/relay.go b/controller/relay.go
index 6118ca8..a44b113 100644
--- a/controller/relay.go
+++ b/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)
diff --git a/controller/token.go b/controller/token.go
index c80ed4b..d7db8f2 100644
--- a/controller/token.go
+++ b/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{
diff --git a/middleware/auth.go b/middleware/auth.go
index 22bf9fe..dade849 100644
--- a/middleware/auth.go
+++ b/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])
}
diff --git a/model/token.go b/model/token.go
index 857a515..c318206 100644
--- a/model/token.go
+++ b/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
+}
diff --git a/web/src/components/TokensTable.js b/web/src/components/TokensTable.js
index 8d1cbcc..4ac2768 100644
--- a/web/src/components/TokensTable.js
+++ b/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 = () => {
{token.id}
{token.name ? token.name : '无'}
{renderStatus(token.status)}
- {token.remain_times === -1 ? "无限制" : token.remain_times}
+ {token.unlimited_times ? "无限制" : token.remain_times}
{renderTimestamp(token.created_time)}
{renderTimestamp(token.accessed_time)}
{token.expired_time === -1 ? "永不过期" : renderTimestamp(token.expired_time)}
diff --git a/web/src/pages/Token/EditToken.js b/web/src/pages/Token/EditToken.js
index 4e6d171..dcd3180 100644
--- a/web/src/pages/Token/EditToken.js
+++ b/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 = () => {
+