Skip to content

phsl / new-api

Go to a project
Commit d67f446b by CaIon
Options

feat: implement token key fetching and masking in API responses

parent ac72f90f
Showing with 530 additions and 98 deletions
controller/token.go
...@@ -14,6 +14,23 @@ import ( ...@@ -14,6 +14,23 @@ import (
"github.com/gin-gonic/gin" "github.com/gin-gonic/gin"
) )
func buildMaskedTokenResponse(token *model.Token) *model.Token {
if token == nil {
return nil
}
maskedToken := *token
maskedToken.Key = token.GetMaskedKey()
return &maskedToken
}
func buildMaskedTokenResponses(tokens []*model.Token) []*model.Token {
maskedTokens := make([]*model.Token, 0, len(tokens))
for _, token := range tokens {
maskedTokens = append(maskedTokens, buildMaskedTokenResponse(token))
}
return maskedTokens
}
func GetAllTokens(c *gin.Context) { func GetAllTokens(c *gin.Context) {
userId := c.GetInt("id") userId := c.GetInt("id")
pageInfo := common.GetPageQuery(c) pageInfo := common.GetPageQuery(c)
...@@ -24,9 +41,8 @@ func GetAllTokens(c *gin.Context) { ...@@ -24,9 +41,8 @@ func GetAllTokens(c *gin.Context) {
} }
total, _ := model.CountUserTokens(userId) total, _ := model.CountUserTokens(userId)
pageInfo.SetTotal(int(total)) pageInfo.SetTotal(int(total))
pageInfo.SetItems(tokens) pageInfo.SetItems(buildMaskedTokenResponses(tokens))
common.ApiSuccess(c, pageInfo) common.ApiSuccess(c, pageInfo)
return
} }
func SearchTokens(c *gin.Context) { func SearchTokens(c *gin.Context) {
...@@ -42,9 +58,8 @@ func SearchTokens(c *gin.Context) { ...@@ -42,9 +58,8 @@ func SearchTokens(c *gin.Context) {
return return
} }
pageInfo.SetTotal(int(total)) pageInfo.SetTotal(int(total))
pageInfo.SetItems(tokens) pageInfo.SetItems(buildMaskedTokenResponses(tokens))
common.ApiSuccess(c, pageInfo) common.ApiSuccess(c, pageInfo)
return
} }
func GetToken(c *gin.Context) { func GetToken(c *gin.Context) {
...@@ -59,12 +74,24 @@ func GetToken(c *gin.Context) { ...@@ -59,12 +74,24 @@ func GetToken(c *gin.Context) {
common.ApiError(c, err) common.ApiError(c, err)
return return
} }
c.JSON(http.StatusOK, gin.H{ common.ApiSuccess(c, buildMaskedTokenResponse(token))
"success": true, }
"message": "",
"data": token, func GetTokenKey(c *gin.Context) {
id, err := strconv.Atoi(c.Param("id"))
userId := c.GetInt("id")
if err != nil {
common.ApiError(c, err)
return
}
token, err := model.GetTokenByIds(id, userId)
if err != nil {
common.ApiError(c, err)
return
}
common.ApiSuccess(c, gin.H{
"key": token.GetFullKey(),
}) })
return
} }
func GetTokenStatus(c *gin.Context) { func GetTokenStatus(c *gin.Context) {
...@@ -204,7 +231,6 @@ func AddToken(c *gin.Context) { ...@@ -204,7 +231,6 @@ func AddToken(c *gin.Context) {
"success": true, "success": true,
"message": "", "message": "",
}) })
return
} }
func DeleteToken(c *gin.Context) { func DeleteToken(c *gin.Context) {
...@@ -219,7 +245,6 @@ func DeleteToken(c *gin.Context) { ...@@ -219,7 +245,6 @@ func DeleteToken(c *gin.Context) {
"success": true, "success": true,
"message": "", "message": "",
}) })
return
} }
func UpdateToken(c *gin.Context) { func UpdateToken(c *gin.Context) {
...@@ -283,7 +308,7 @@ func UpdateToken(c *gin.Context) { ...@@ -283,7 +308,7 @@ func UpdateToken(c *gin.Context) {
c.JSON(http.StatusOK, gin.H{ c.JSON(http.StatusOK, gin.H{
"success": true, "success": true,
"message": "", "message": "",
"data": cleanToken, "data": buildMaskedTokenResponse(cleanToken),
}) })
} }
......
controller/token_test.go 0 → 100644
package controller
import (
"bytes"
"encoding/json"
"fmt"
"net/http"
"net/http/httptest"
"strconv"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/model"
"github.com/gin-gonic/gin"
"github.com/glebarez/sqlite"
"gorm.io/gorm"
)
type tokenAPIResponse struct {
Success bool `json:"success"`
Message string `json:"message"`
Data json.RawMessage `json:"data"`
}
type tokenPageResponse struct {
Items []tokenResponseItem `json:"items"`
}
type tokenResponseItem struct {
ID int `json:"id"`
Name string `json:"name"`
Key string `json:"key"`
Status int `json:"status"`
}
type tokenKeyResponse struct {
Key string `json:"key"`
}
func setupTokenControllerTestDB(t *testing.T) *gorm.DB {
t.Helper()
gin.SetMode(gin.TestMode)
common.UsingSQLite = true
common.UsingMySQL = false
common.UsingPostgreSQL = false
common.RedisEnabled = false
dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_"))
db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{})
if err != nil {
t.Fatalf("failed to open sqlite db: %v", err)
}
model.DB = db
model.LOG_DB = db
if err := db.AutoMigrate(&model.Token{}); err != nil {
t.Fatalf("failed to migrate token table: %v", err)
}
t.Cleanup(func() {
sqlDB, err := db.DB()
if err == nil {
_ = sqlDB.Close()
}
})
return db
}
func seedToken(t *testing.T, db *gorm.DB, userID int, name string, rawKey string) *model.Token {
t.Helper()
token := &model.Token{
UserId: userID,
Name: name,
Key: rawKey,
Status: common.TokenStatusEnabled,
CreatedTime: 1,
AccessedTime: 1,
ExpiredTime: -1,
RemainQuota: 100,
UnlimitedQuota: true,
Group: "default",
}
if err := db.Create(token).Error; err != nil {
t.Fatalf("failed to create token: %v", err)
}
return token
}
func newAuthenticatedContext(t *testing.T, method string, target string, body any, userID int) (*gin.Context, *httptest.ResponseRecorder) {
t.Helper()
var requestBody *bytes.Reader
if body != nil {
payload, err := common.Marshal(body)
if err != nil {
t.Fatalf("failed to marshal request body: %v", err)
}
requestBody = bytes.NewReader(payload)
} else {
requestBody = bytes.NewReader(nil)
}
recorder := httptest.NewRecorder()
ctx, _ := gin.CreateTestContext(recorder)
ctx.Request = httptest.NewRequest(method, target, requestBody)
if body != nil {
ctx.Request.Header.Set("Content-Type", "application/json")
}
ctx.Set("id", userID)
return ctx, recorder
}
func decodeAPIResponse(t *testing.T, recorder *httptest.ResponseRecorder) tokenAPIResponse {
t.Helper()
var response tokenAPIResponse
if err := common.Unmarshal(recorder.Body.Bytes(), &response); err != nil {
t.Fatalf("failed to decode api response: %v", err)
}
return response
}
func TestGetAllTokensMasksKeyInResponse(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 1, "list-token", "abcd1234efgh5678")
seedToken(t, db, 2, "other-user-token", "zzzz1234yyyy5678")
ctx, recorder := newAuthenticatedContext(t, http.MethodGet, "/api/token/?p=1&size=10", nil, 1)
GetAllTokens(ctx)
response := decodeAPIResponse(t, recorder)
if !response.Success {
t.Fatalf("expected success response, got message: %s", response.Message)
}
var page tokenPageResponse
if err := common.Unmarshal(response.Data, &page); err != nil {
t.Fatalf("failed to decode token page response: %v", err)
}
if len(page.Items) != 1 {
t.Fatalf("expected exactly one token, got %d", len(page.Items))
}
if page.Items[0].Key != token.GetMaskedKey() {
t.Fatalf("expected masked key %q, got %q", token.GetMaskedKey(), page.Items[0].Key)
}
if strings.Contains(recorder.Body.String(), token.Key) {
t.Fatalf("list response leaked raw token key: %s", recorder.Body.String())
}
}
func TestSearchTokensMasksKeyInResponse(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 1, "searchable-token", "ijkl1234mnop5678")
ctx, recorder := newAuthenticatedContext(t, http.MethodGet, "/api/token/search?keyword=searchable-token&p=1&size=10", nil, 1)
SearchTokens(ctx)
response := decodeAPIResponse(t, recorder)
if !response.Success {
t.Fatalf("expected success response, got message: %s", response.Message)
}
var page tokenPageResponse
if err := common.Unmarshal(response.Data, &page); err != nil {
t.Fatalf("failed to decode search response: %v", err)
}
if len(page.Items) != 1 {
t.Fatalf("expected exactly one search result, got %d", len(page.Items))
}
if page.Items[0].Key != token.GetMaskedKey() {
t.Fatalf("expected masked search key %q, got %q", token.GetMaskedKey(), page.Items[0].Key)
}
if strings.Contains(recorder.Body.String(), token.Key) {
t.Fatalf("search response leaked raw token key: %s", recorder.Body.String())
}
}
func TestGetTokenMasksKeyInResponse(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 1, "detail-token", "qrst1234uvwx5678")
ctx, recorder := newAuthenticatedContext(t, http.MethodGet, "/api/token/"+strconv.Itoa(token.Id), nil, 1)
ctx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(token.Id)}}
GetToken(ctx)
response := decodeAPIResponse(t, recorder)
if !response.Success {
t.Fatalf("expected success response, got message: %s", response.Message)
}
var detail tokenResponseItem
if err := common.Unmarshal(response.Data, &detail); err != nil {
t.Fatalf("failed to decode token detail response: %v", err)
}
if detail.Key != token.GetMaskedKey() {
t.Fatalf("expected masked detail key %q, got %q", token.GetMaskedKey(), detail.Key)
}
if strings.Contains(recorder.Body.String(), token.Key) {
t.Fatalf("detail response leaked raw token key: %s", recorder.Body.String())
}
}
func TestUpdateTokenMasksKeyInResponse(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 1, "editable-token", "yzab1234cdef5678")
body := map[string]any{
"id": token.Id,
"name": "updated-token",
"expired_time": -1,
"remain_quota": 100,
"unlimited_quota": true,
"model_limits_enabled": false,
"model_limits": "",
"group": "default",
"cross_group_retry": false,
}
ctx, recorder := newAuthenticatedContext(t, http.MethodPut, "/api/token/", body, 1)
UpdateToken(ctx)
response := decodeAPIResponse(t, recorder)
if !response.Success {
t.Fatalf("expected success response, got message: %s", response.Message)
}
var detail tokenResponseItem
if err := common.Unmarshal(response.Data, &detail); err != nil {
t.Fatalf("failed to decode token update response: %v", err)
}
if detail.Key != token.GetMaskedKey() {
t.Fatalf("expected masked update key %q, got %q", token.GetMaskedKey(), detail.Key)
}
if strings.Contains(recorder.Body.String(), token.Key) {
t.Fatalf("update response leaked raw token key: %s", recorder.Body.String())
}
}
func TestGetTokenKeyRequiresOwnershipAndReturnsFullKey(t *testing.T) {
db := setupTokenControllerTestDB(t)
token := seedToken(t, db, 1, "owned-token", "owner1234token5678")
authorizedCtx, authorizedRecorder := newAuthenticatedContext(t, http.MethodPost, "/api/token/"+strconv.Itoa(token.Id)+"/key", nil, 1)
authorizedCtx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(token.Id)}}
GetTokenKey(authorizedCtx)
authorizedResponse := decodeAPIResponse(t, authorizedRecorder)
if !authorizedResponse.Success {
t.Fatalf("expected authorized key fetch to succeed, got message: %s", authorizedResponse.Message)
}
var keyData tokenKeyResponse
if err := common.Unmarshal(authorizedResponse.Data, &keyData); err != nil {
t.Fatalf("failed to decode token key response: %v", err)
}
if keyData.Key != token.GetFullKey() {
t.Fatalf("expected full key %q, got %q", token.GetFullKey(), keyData.Key)
}
unauthorizedCtx, unauthorizedRecorder := newAuthenticatedContext(t, http.MethodPost, "/api/token/"+strconv.Itoa(token.Id)+"/key", nil, 2)
unauthorizedCtx.Params = gin.Params{{Key: "id", Value: strconv.Itoa(token.Id)}}
GetTokenKey(unauthorizedCtx)
unauthorizedResponse := decodeAPIResponse(t, unauthorizedRecorder)
if unauthorizedResponse.Success {
t.Fatalf("expected unauthorized key fetch to fail")
}
if strings.Contains(unauthorizedRecorder.Body.String(), token.Key) {
t.Fatalf("unauthorized key response leaked raw token key: %s", unauthorizedRecorder.Body.String())
}
}
model/token.go
...@@ -35,6 +35,27 @@ func (token *Token) Clean() { ...@@ -35,6 +35,27 @@ func (token *Token) Clean() {
token.Key = "" token.Key = ""
} }
func MaskTokenKey(key string) string {
if key == "" {
return ""
}
if len(key) <= 4 {
return strings.Repeat("*", len(key))
}
if len(key) <= 8 {
return key[:2] + "****" + key[len(key)-2:]
}
return key[:4] + "**********" + key[len(key)-4:]
}
func (token *Token) GetFullKey() string {
return token.Key
}
func (token *Token) GetMaskedKey() string {
return MaskTokenKey(token.Key)
}
func (token *Token) GetIpLimits() []string { func (token *Token) GetIpLimits() []string {
// delete empty spaces // delete empty spaces
//split with \n //split with \n
...@@ -201,7 +222,7 @@ func ValidateUserToken(key string) (token *Token, err error) { ...@@ -201,7 +222,7 @@ func ValidateUserToken(key string) (token *Token, err error) {
} }
keyPrefix := key[:3] keyPrefix := key[:3]
keySuffix := key[len(key)-3:] keySuffix := key[len(key)-3:]
return token, errors.New(fmt.Sprintf("[sk-%s***%s] 该令牌额度已用尽 !token.UnlimitedQuota && token.RemainQuota = %d", keyPrefix, keySuffix, token.RemainQuota)) return token, fmt.Errorf("[sk-%s***%s] 该令牌额度已用尽 !token.UnlimitedQuota && token.RemainQuota = %d", keyPrefix, keySuffix, token.RemainQuota)
} }
return token, nil return token, nil
} }
......
router/api-router.go
...@@ -248,6 +248,7 @@ func SetApiRouter(router *gin.Engine) { ...@@ -248,6 +248,7 @@ func SetApiRouter(router *gin.Engine) {
tokenRoute.GET("/", controller.GetAllTokens) tokenRoute.GET("/", controller.GetAllTokens)
tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens) tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens)
tokenRoute.GET("/:id", controller.GetToken) tokenRoute.GET("/:id", controller.GetToken)
tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKey)
tokenRoute.POST("/", controller.AddToken) tokenRoute.POST("/", controller.AddToken)
tokenRoute.PUT("/", controller.UpdateToken) tokenRoute.PUT("/", controller.UpdateToken)
tokenRoute.DELETE("/:id", controller.DeleteToken) tokenRoute.DELETE("/:id", controller.DeleteToken)
......
web/src/components/table/tokens/TokensActions.jsx
...@@ -29,7 +29,6 @@ const TokensActions = ({ ...@@ -29,7 +29,6 @@ const TokensActions = ({
setShowEdit, setShowEdit,
batchCopyTokens, batchCopyTokens,
batchDeleteTokens, batchDeleteTokens,
copyText,
t, t,
}) => { }) => {
// Modal states // Modal states
...@@ -99,8 +98,7 @@ const TokensActions = ({ ...@@ -99,8 +98,7 @@ const TokensActions = ({
<CopyTokensModal <CopyTokensModal
visible={showCopyModal} visible={showCopyModal}
onCancel={() => setShowCopyModal(false)} onCancel={() => setShowCopyModal(false)}
selectedKeys={selectedKeys} batchCopyTokens={batchCopyTokens}
copyText={copyText}
t={t} t={t}
/> />
......
web/src/components/table/tokens/TokensColumnDefs.jsx
...@@ -108,17 +108,28 @@ const renderGroupColumn = (text, record, t) => { ...@@ -108,17 +108,28 @@ const renderGroupColumn = (text, record, t) => {
}; };
// Render token key column with show/hide and copy functionality // Render token key column with show/hide and copy functionality
const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => { const renderTokenKey = (
const fullKey = 'sk-' + record.key; text,
const maskedKey = record,
'sk-' + record.key.slice(0, 4) + '**********' + record.key.slice(-4); showKeys,
resolvedTokenKeys,
loadingTokenKeys,
toggleTokenVisibility,
copyTokenKey,
) => {
const revealed = !!showKeys[record.id]; const revealed = !!showKeys[record.id];
const loading = !!loadingTokenKeys[record.id];
const keyValue =
revealed && resolvedTokenKeys[record.id]
? resolvedTokenKeys[record.id]
: record.key || '';
const displayedKey = keyValue ? `sk-${keyValue}` : '';
return ( return (
<div className='w-[200px]'> <div className='w-[200px]'>
<Input <Input
readOnly readOnly
value={revealed ? fullKey : maskedKey} value={displayedKey}
size='small' size='small'
suffix={ suffix={
<div className='flex items-center'> <div className='flex items-center'>
...@@ -127,10 +138,11 @@ const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => { ...@@ -127,10 +138,11 @@ const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => {
size='small' size='small'
type='tertiary' type='tertiary'
icon={revealed ? <IconEyeClosed /> : <IconEyeOpened />} icon={revealed ? <IconEyeClosed /> : <IconEyeOpened />}
loading={loading}
aria-label='toggle token visibility' aria-label='toggle token visibility'
onClick={(e) => { onClick={async (e) => {
e.stopPropagation(); e.stopPropagation();
setShowKeys((prev) => ({ ...prev, [record.id]: !revealed })); await toggleTokenVisibility(record);
}} }}
/> />
<Button <Button
...@@ -138,10 +150,11 @@ const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => { ...@@ -138,10 +150,11 @@ const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => {
size='small' size='small'
type='tertiary' type='tertiary'
icon={<IconCopy />} icon={<IconCopy />}
loading={loading}
aria-label='copy token key' aria-label='copy token key'
onClick={async (e) => { onClick={async (e) => {
e.stopPropagation(); e.stopPropagation();
await copyText(fullKey); await copyTokenKey(record);
}} }}
/> />
</div> </div>
...@@ -427,8 +440,10 @@ const renderOperations = ( ...@@ -427,8 +440,10 @@ const renderOperations = (
export const getTokensColumns = ({ export const getTokensColumns = ({
t, t,
showKeys, showKeys,
setShowKeys, resolvedTokenKeys,
copyText, loadingTokenKeys,
toggleTokenVisibility,
copyTokenKey,
manageToken, manageToken,
onOpenLink, onOpenLink,
setEditingToken, setEditingToken,
...@@ -461,7 +476,15 @@ export const getTokensColumns = ({ ...@@ -461,7 +476,15 @@ export const getTokensColumns = ({
title: t('密钥'), title: t('密钥'),
key: 'token_key', key: 'token_key',
render: (text, record) => render: (text, record) =>
renderTokenKey(text, record, showKeys, setShowKeys, copyText), renderTokenKey(
text,
record,
showKeys,
resolvedTokenKeys,
loadingTokenKeys,
toggleTokenVisibility,
copyTokenKey,
),
}, },
{ {
title: t('可用模型'), title: t('可用模型'),
......
web/src/components/table/tokens/TokensTable.jsx
...@@ -39,8 +39,10 @@ const TokensTable = (tokensData) => { ...@@ -39,8 +39,10 @@ const TokensTable = (tokensData) => {
rowSelection, rowSelection,
handleRow, handleRow,
showKeys, showKeys,
setShowKeys, resolvedTokenKeys,
copyText, loadingTokenKeys,
toggleTokenVisibility,
copyTokenKey,
manageToken, manageToken,
onOpenLink, onOpenLink,
setEditingToken, setEditingToken,
...@@ -54,8 +56,10 @@ const TokensTable = (tokensData) => { ...@@ -54,8 +56,10 @@ const TokensTable = (tokensData) => {
return getTokensColumns({ return getTokensColumns({
t, t,
showKeys, showKeys,
setShowKeys, resolvedTokenKeys,
copyText, loadingTokenKeys,
toggleTokenVisibility,
copyTokenKey,
manageToken, manageToken,
onOpenLink, onOpenLink,
setEditingToken, setEditingToken,
...@@ -65,8 +69,10 @@ const TokensTable = (tokensData) => { ...@@ -65,8 +69,10 @@ const TokensTable = (tokensData) => {
}, [ }, [
t, t,
showKeys, showKeys,
setShowKeys, resolvedTokenKeys,
copyText, loadingTokenKeys,
toggleTokenVisibility,
copyTokenKey,
manageToken, manageToken,
onOpenLink, onOpenLink,
setEditingToken, setEditingToken,
......
web/src/components/table/tokens/index.jsx
...@@ -58,6 +58,7 @@ function TokensPage() { ...@@ -58,6 +58,7 @@ function TokensPage() {
t: (k) => k, t: (k) => k,
selectedModel: '', selectedModel: '',
prefillKey: '', prefillKey: '',
fetchTokenKey: async () => '',
}); });
const [modelOptions, setModelOptions] = useState([]); const [modelOptions, setModelOptions] = useState([]);
const [selectedModel, setSelectedModel] = useState(''); const [selectedModel, setSelectedModel] = useState('');
...@@ -74,6 +75,7 @@ function TokensPage() { ...@@ -74,6 +75,7 @@ function TokensPage() {
t: tokensData.t, t: tokensData.t,
selectedModel, selectedModel,
prefillKey, prefillKey,
fetchTokenKey: tokensData.fetchTokenKey,
}; };
}, [ }, [
tokensData.tokens, tokensData.tokens,
...@@ -81,6 +83,7 @@ function TokensPage() { ...@@ -81,6 +83,7 @@ function TokensPage() {
tokensData.t, tokensData.t,
selectedModel, selectedModel,
prefillKey, prefillKey,
tokensData.fetchTokenKey,
]); ]);
const loadModels = async () => { const loadModels = async () => {
...@@ -198,13 +201,14 @@ function TokensPage() { ...@@ -198,13 +201,14 @@ function TokensPage() {
openCCSwitchModalRef.current = openCCSwitchModal; openCCSwitchModalRef.current = openCCSwitchModal;
// Prefill to Fluent handler // Prefill to Fluent handler
const handlePrefillToFluent = () => { const handlePrefillToFluent = async () => {
const { const {
tokens, tokens,
selectedKeys, selectedKeys,
t, t,
selectedModel: chosenModel, selectedModel: chosenModel,
prefillKey: overrideKey, prefillKey: overrideKey,
fetchTokenKey,
} = latestRef.current; } = latestRef.current;
const container = document.getElementById('fluent-new-api-container'); const container = document.getElementById('fluent-new-api-container');
if (!container) { if (!container) {
...@@ -241,7 +245,11 @@ function TokensPage() { ...@@ -241,7 +245,11 @@ function TokensPage() {
Toast.warning(t('没有可用令牌用于填充')); Toast.warning(t('没有可用令牌用于填充'));
return; return;
} }
apiKeyToUse = 'sk-' + token.key; try {
apiKeyToUse = 'sk-' + (await fetchTokenKey(token));
} catch (_) {
return;
}
} }
const payload = { const payload = {
...@@ -351,7 +359,6 @@ function TokensPage() { ...@@ -351,7 +359,6 @@ function TokensPage() {
setShowEdit, setShowEdit,
batchCopyTokens, batchCopyTokens,
batchDeleteTokens, batchDeleteTokens,
copyText,
// Filters state // Filters state
formInitValues, formInitValues,
...@@ -401,7 +408,6 @@ function TokensPage() { ...@@ -401,7 +408,6 @@ function TokensPage() {
setShowEdit={setShowEdit} setShowEdit={setShowEdit}
batchCopyTokens={batchCopyTokens} batchCopyTokens={batchCopyTokens}
batchDeleteTokens={batchDeleteTokens} batchDeleteTokens={batchDeleteTokens}
copyText={copyText}
t={t} t={t}
/> />
......
web/src/components/table/tokens/modals/CCSwitchModal.jsx
...@@ -116,8 +116,7 @@ export default function CCSwitchModal({ ...@@ -116,8 +116,7 @@ export default function CCSwitchModal({
Toast.warning(t('请选择主模型')); Toast.warning(t('请选择主模型'));
return; return;
} }
const apiKey = 'sk-' + tokenKey; const url = buildCCSwitchURL(app, name, models, 'sk-' + tokenKey);
const url = buildCCSwitchURL(app, name, models, apiKey);
window.open(url, '_blank'); window.open(url, '_blank');
onClose(); onClose();
}; };
......
web/src/components/table/tokens/modals/CopyTokensModal.jsx
...@@ -20,24 +20,21 @@ For commercial licensing, please contact support@quantumnous.com ...@@ -20,24 +20,21 @@ For commercial licensing, please contact support@quantumnous.com
import React from 'react'; import React from 'react';
import { Modal, Button, Space } from '@douyinfe/semi-ui'; import { Modal, Button, Space } from '@douyinfe/semi-ui';
const CopyTokensModal = ({ visible, onCancel, selectedKeys, copyText, t }) => { const CopyTokensModal = ({
visible,
onCancel,
batchCopyTokens,
t,
}) => {
// Handle copy with name and key format // Handle copy with name and key format
const handleCopyWithName = async () => { const handleCopyWithName = async () => {
let content = ''; await batchCopyTokens('name+key');
for (let i = 0; i < selectedKeys.length; i++) {
content += selectedKeys[i].name + ' sk-' + selectedKeys[i].key + '\n';
}
await copyText(content);
onCancel(); onCancel();
}; };
// Handle copy with key only format // Handle copy with key only format
const handleCopyKeyOnly = async () => { const handleCopyKeyOnly = async () => {
let content = ''; await batchCopyTokens('key-only');
for (let i = 0; i < selectedKeys.length; i++) {
content += 'sk-' + selectedKeys[i].key + '\n';
}
await copyText(content);
onCancel(); onCancel();
}; };
......
web/src/components/table/usage-logs/modals/ColumnSelectorModal.jsx
...@@ -73,7 +73,7 @@ const ColumnSelectorModal = ({ ...@@ -73,7 +73,7 @@ const ColumnSelectorModal = ({
<RadioGroup <RadioGroup
type='button' type='button'
value={billingDisplayMode} value={billingDisplayMode}
onChange={(event) => setBillingDisplayMode(event.target.value)} onChange={(value) => setBillingDisplayMode(value)}
> >
<Radio value='price'> <Radio value='price'>
{isTokensDisplay ? t('价格模式') : t('价格模式（默认）')} {isTokensDisplay ? t('价格模式') : t('价格模式（默认）')}
......
web/src/helpers/token.js
...@@ -20,8 +20,22 @@ For commercial licensing, please contact support@quantumnous.com ...@@ -20,8 +20,22 @@ For commercial licensing, please contact support@quantumnous.com
import { API } from './api'; import { API } from './api';
/** /**
* 获取可用的token keys * 按需获取单个令牌的真实 key
* @returns {Promise<string[]>} 返回active状态的token key数组 * @param {number|string} tokenId
* @returns {Promise<string>} 返回不带 sk- 前缀的真实 token key
*/
export async function fetchTokenKey(tokenId) {
const response = await API.post(`/api/token/${tokenId}/key`);
const { success, data, message } = response.data || {};
if (!success || !data?.key) {
throw new Error(message || 'Failed to fetch token key');
}
return data.key;
}
/**
* 获取可用的 token keys
* @returns {Promise<string[]>} 返回 active 状态的不带 sk- 前缀的真实 token key 数组
*/ */
export async function fetchTokenKeys() { export async function fetchTokenKeys() {
try { try {
...@@ -31,7 +45,12 @@ export async function fetchTokenKeys() { ...@@ -31,7 +45,12 @@ export async function fetchTokenKeys() {
const tokenItems = Array.isArray(data) ? data : data.items || []; const tokenItems = Array.isArray(data) ? data : data.items || [];
const activeTokens = tokenItems.filter((token) => token.status === 1); const activeTokens = tokenItems.filter((token) => token.status === 1);
return activeTokens.map((token) => token.key); const keyResults = await Promise.allSettled(
activeTokens.map((token) => fetchTokenKey(token.id)),
);
return keyResults
.filter((result) => result.status === 'fulfilled' && result.value)
.map((result) => result.value);
} catch (error) { } catch (error) {
console.error('Error fetching token keys:', error); console.error('Error fetching token keys:', error);
return []; return [];
......
web/src/hooks/tokens/useTokensData.jsx
...@@ -17,7 +17,7 @@ along with this program. If not, see <https://www.gnu.org/licenses/>. ...@@ -17,7 +17,7 @@ along with this program. If not, see <https://www.gnu.org/licenses/>.
For commercial licensing, please contact support@quantumnous.com For commercial licensing, please contact support@quantumnous.com
*/ */
import { useState, useEffect } from 'react'; import { useState, useEffect, useRef } from 'react';
import { useTranslation } from 'react-i18next'; import { useTranslation } from 'react-i18next';
import { Modal } from '@douyinfe/semi-ui'; import { Modal } from '@douyinfe/semi-ui';
import { import {
...@@ -29,6 +29,7 @@ import { ...@@ -29,6 +29,7 @@ import {
} from '../../helpers'; } from '../../helpers';
import { ITEMS_PER_PAGE } from '../../constants'; import { ITEMS_PER_PAGE } from '../../constants';
import { useTableCompactMode } from '../common/useTableCompactMode'; import { useTableCompactMode } from '../common/useTableCompactMode';
import { fetchTokenKey as fetchTokenKeyById } from '../../helpers/token';
export const useTokensData = (openFluentNotification, openCCSwitchModal) => { export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
const { t } = useTranslation(); const { t } = useTranslation();
...@@ -54,6 +55,9 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -54,6 +55,9 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
// UI state // UI state
const [compactMode, setCompactMode] = useTableCompactMode('tokens'); const [compactMode, setCompactMode] = useTableCompactMode('tokens');
const [showKeys, setShowKeys] = useState({}); const [showKeys, setShowKeys] = useState({});
const [resolvedTokenKeys, setResolvedTokenKeys] = useState({});
const [loadingTokenKeys, setLoadingTokenKeys] = useState({});
const keyRequestsRef = useRef({});
// Form state // Form state
const [formApi, setFormApi] = useState(null); const [formApi, setFormApi] = useState(null);
...@@ -87,6 +91,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -87,6 +91,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
setTokenCount(payload.total || 0); setTokenCount(payload.total || 0);
setActivePage(payload.page || 1); setActivePage(payload.page || 1);
setPageSize(payload.page_size || pageSize); setPageSize(payload.page_size || pageSize);
setShowKeys({});
}; };
// Load tokens function // Load tokens function
...@@ -122,14 +127,86 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -122,14 +127,86 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
} }
}; };
const fetchTokenKey = async (tokenOrId, options = {}) => {
const { suppressError = false } = options;
const tokenId =
typeof tokenOrId === 'object' ? tokenOrId?.id : Number(tokenOrId);
if (!tokenId) {
const error = new Error(t('令牌不存在'));
if (!suppressError) {
showError(error.message);
}
throw error;
}
if (resolvedTokenKeys[tokenId]) {
return resolvedTokenKeys[tokenId];
}
if (keyRequestsRef.current[tokenId]) {
return keyRequestsRef.current[tokenId];
}
const request = (async () => {
setLoadingTokenKeys((prev) => ({ ...prev, [tokenId]: true }));
try {
const fullKey = await fetchTokenKeyById(tokenId);
setResolvedTokenKeys((prev) => ({ ...prev, [tokenId]: fullKey }));
return fullKey;
} catch (error) {
const normalizedError = new Error(
error?.message || t('获取令牌密钥失败'),
);
if (!suppressError) {
showError(normalizedError.message);
}
throw normalizedError;
} finally {
delete keyRequestsRef.current[tokenId];
setLoadingTokenKeys((prev) => {
const next = { ...prev };
delete next[tokenId];
return next;
});
}
})();
keyRequestsRef.current[tokenId] = request;
return request;
};
const toggleTokenVisibility = async (record) => {
const tokenId = record?.id;
if (!tokenId) {
return;
}
if (showKeys[tokenId]) {
setShowKeys((prev) => ({ ...prev, [tokenId]: false }));
return;
}
const fullKey = await fetchTokenKey(record);
if (fullKey) {
setShowKeys((prev) => ({ ...prev, [tokenId]: true }));
}
};
const copyTokenKey = async (record) => {
const fullKey = await fetchTokenKey(record);
await copyText(`sk-${fullKey}`);
};
// Open link function for chat integrations // Open link function for chat integrations
const onOpenLink = async (type, url, record) => { const onOpenLink = async (type, url, record) => {
const fullKey = await fetchTokenKey(record);
if (url && url.startsWith('ccswitch')) { if (url && url.startsWith('ccswitch')) {
openCCSwitchModal(record.key); openCCSwitchModal(fullKey);
return; return;
} }
if (url && url.startsWith('fluent')) { if (url && url.startsWith('fluent')) {
openFluentNotification(record.key); openFluentNotification(fullKey);
return; return;
} }
let status = localStorage.getItem('status'); let status = localStorage.getItem('status');
...@@ -145,7 +222,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -145,7 +222,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
let cherryConfig = { let cherryConfig = {
id: 'new-api', id: 'new-api',
baseUrl: serverAddress, baseUrl: serverAddress,
apiKey: 'sk-' + record.key, apiKey: `sk-${fullKey}`,
}; };
let encodedConfig = encodeURIComponent( let encodedConfig = encodeURIComponent(
encodeToBase64(JSON.stringify(cherryConfig)), encodeToBase64(JSON.stringify(cherryConfig)),
...@@ -155,7 +232,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -155,7 +232,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
let aionuiConfig = { let aionuiConfig = {
platform: 'new-api', platform: 'new-api',
baseUrl: serverAddress, baseUrl: serverAddress,
apiKey: 'sk-' + record.key, apiKey: `sk-${fullKey}`,
}; };
let encodedConfig = encodeURIComponent( let encodedConfig = encodeURIComponent(
encodeToBase64(JSON.stringify(aionuiConfig)), encodeToBase64(JSON.stringify(aionuiConfig)),
...@@ -164,7 +241,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -164,7 +241,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
} else { } else {
let encodedServerAddress = encodeURIComponent(serverAddress); let encodedServerAddress = encodeURIComponent(serverAddress);
url = url.replaceAll('{address}', encodedServerAddress); url = url.replaceAll('{address}', encodedServerAddress);
url = url.replaceAll('{key}', 'sk-' + record.key); url = url.replaceAll('{key}', `sk-${fullKey}`);
} }
window.open(url, '_blank'); window.open(url, '_blank');
...@@ -314,48 +391,28 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -314,48 +391,28 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
}; };
// Batch copy tokens // Batch copy tokens
const batchCopyTokens = (copyType) => { const batchCopyTokens = async (copyType) => {
if (selectedKeys.length === 0) { if (selectedKeys.length === 0) {
showError(t('请至少选择一个令牌！')); showError(t('请至少选择一个令牌！'));
return; return;
} }
try {
Modal.info({ const keys = await Promise.all(
title: t('复制令牌'), selectedKeys.map((token) => fetchTokenKey(token, { suppressError: true })),
icon: null, );
content: t('请选择你的复制方式'), let content = '';
footer: ( for (let i = 0; i < selectedKeys.length; i++) {
<div className='flex gap-2'> const fullKey = keys[i];
<button if (copyType === 'name+key') {
className='px-3 py-1 bg-gray-200 rounded' content += `${selectedKeys[i].name} sk-${fullKey}\n`;
onClick={async () => { } else {
let content = ''; content += `sk-${fullKey}\n`;
for (let i = 0; i < selectedKeys.length; i++) { }
content += }
selectedKeys[i].name + ' sk-' + selectedKeys[i].key + '\n'; await copyText(content);
} } catch (error) {
await copyText(content); showError(error?.message || t('复制令牌失败'));
Modal.destroyAll(); }
}}
>
{t('名称+密钥')}
</button>
<button
className='px-3 py-1 bg-blue-500 text-white rounded'
onClick={async () => {
let content = '';
for (let i = 0; i < selectedKeys.length; i++) {
content += 'sk-' + selectedKeys[i].key + '\n';
}
await copyText(content);
Modal.destroyAll();
}}
>
{t('仅密钥')}
</button>
</div>
),
});
}; };
// Initialize data // Initialize data
...@@ -392,6 +449,8 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -392,6 +449,8 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
setCompactMode, setCompactMode,
showKeys, showKeys,
setShowKeys, setShowKeys,
resolvedTokenKeys,
loadingTokenKeys,
// Form state // Form state
formApi, formApi,
...@@ -403,6 +462,9 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { ...@@ -403,6 +462,9 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => {
loadTokens, loadTokens,
refresh, refresh,
copyText, copyText,
fetchTokenKey,
toggleTokenVisibility,
copyTokenKey,
onOpenLink, onOpenLink,
manageToken, manageToken,
searchTokens, searchTokens,
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or sign in to comment