Skip to content

feat: query usage of token #1161

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
Calcium-Ion merged 4 commits into from
Aug 23, 2025
Merged

feat: query usage of token #1161

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
530af5e
feat: `/api/token/usage`
lollipopkit Apr 29, 2025
6c4242a
Merge branch 'QuantumNous:main' into main
lollipopkit Jun 5, 2025
df1ec48
Merge branch 'alpha' into main
Calcium-Ion Aug 23, 2025
93ce48a
feat: restructure token usage routes and enhance token retrieval logic
Calcium-Ion Aug 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
52 changes: 52 additions & 0 deletions controller/token.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,7 @@ import (
"one-api/common"
"one-api/model"
"strconv"
"strings"

"github.com/gin-gonic/gin"
)
Expand Down Expand Up @@ -82,6 +83,57 @@ func GetTokenStatus(c *gin.Context) {
})
}

func GetTokenUsage(c *gin.Context) {
authHeader := c.GetHeader("Authorization")
if authHeader == "" {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": "No Authorization header",
})
return
}

parts := strings.Split(authHeader, " ")
if len(parts) != 2 || strings.ToLower(parts[0]) != "bearer" {
c.JSON(http.StatusUnauthorized, gin.H{
"success": false,
"message": "Invalid Bearer token",
})
return
}
tokenKey := parts[1]

token, err := model.GetTokenByKey(strings.TrimPrefix(tokenKey, "sk-"), false)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": err.Error(),
})
return
}

expiredAt := token.ExpiredTime
if expiredAt == -1 {
expiredAt = 0
}

c.JSON(http.StatusOK, gin.H{
"code": true,
"message": "ok",
"data": gin.H{
"object": "token_usage",
"name": token.Name,
"total_granted": token.RemainQuota + token.UsedQuota,
"total_used": token.UsedQuota,
"total_available": token.RemainQuota,
"unlimited_quota": token.UnlimitedQuota,
"model_limits": token.GetModelLimitsMap(),
"model_limits_enabled": token.ModelLimitsEnabled,
"expires_at": expiredAt,
},
})
}
Comment on lines +86 to +135

@coderabbitai coderabbitai Bot Aug 23, 2025

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Verification agent

🧩 Analysis chain

Add endpoint tests for contract and auth behavior

Cover: 401 no header (if header-parsing kept), 401 invalid scheme, happy path shape (envelope + ms for expires_at), and ownership enforcement (403) if you continue accepting arbitrary token keys.

Proposed verification script to ensure routing and consistency:

🏁 Script executed:

#!/bin/bash
# Confirm the route is protected by TokenAuth
rg -nC3 'GetTokenUsage' -- router | cat
rg -n 'TokenAuth\(' router | cat

# Confirm GetTokenStatus returns expires_at in ms for consistency reference
rg -nC2 'GetTokenStatus|expires_at.*\* *1000' controller | cat

Length of output: 1989

Add tests for GetTokenUsage and normalize expires_at units

Please address the following before merging:

  • Ensure expires_at in GetTokenUsage is returned in milliseconds, to match GetTokenStatus (which does expiredAt * 1000 in controller/token.go:82)
  • Add endpoint tests covering:
    • 401 when the Authorization header is missing
    • 401 when the scheme is not Bearer
    • 403 when a valid token key is used by a non-owner (ownership enforcement)
    • 200 “happy path” response shape, including the JSON envelope and expires_at in ms

Files/locations to update:

  • controller/token.go (normalize expires_at unit)
  • tests for token endpoints (e.g. in controller/token_test.go or equivalent)


func AddToken(c *gin.Context) {
token := model.Token{}
err := c.ShouldBindJSON(&token)
Expand Down
12 changes: 11 additions & 1 deletion router/api-router.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,6 +145,17 @@ func SetApiRouter(router *gin.Engine) {
tokenRoute.DELETE("/:id", controller.DeleteToken)
tokenRoute.POST("/batch", controller.DeleteTokenBatch)
}

usageRoute := apiRouter.Group("/usage")
usageRoute.Use(middleware.CriticalRateLimit())
{
tokenUsageRoute := usageRoute.Group("/token")
tokenUsageRoute.Use(middleware.TokenAuth())
{
tokenUsageRoute.GET("/", controller.GetTokenUsage)
}
}

redemptionRoute := apiRouter.Group("/redemption")
redemptionRoute.Use(middleware.AdminAuth())
{
Expand Down Expand Up @@ -172,7 +183,6 @@ func SetApiRouter(router *gin.Engine) {
logRoute.Use(middleware.CORS())
{
logRoute.GET("/token", controller.GetLogByKey)

}
groupRoute := apiRouter.Group("/group")
groupRoute.Use(middleware.AdminAuth())
Expand Down