diff --git a/controller/token.go b/controller/token.go index 50da7e339..889b962a6 100644 --- a/controller/token.go +++ b/controller/token.go @@ -14,6 +14,23 @@ import ( "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) { userId := c.GetInt("id") pageInfo := common.GetPageQuery(c) @@ -24,9 +41,8 @@ func GetAllTokens(c *gin.Context) { } total, _ := model.CountUserTokens(userId) pageInfo.SetTotal(int(total)) - pageInfo.SetItems(tokens) + pageInfo.SetItems(buildMaskedTokenResponses(tokens)) common.ApiSuccess(c, pageInfo) - return } func SearchTokens(c *gin.Context) { @@ -42,9 +58,8 @@ func SearchTokens(c *gin.Context) { return } pageInfo.SetTotal(int(total)) - pageInfo.SetItems(tokens) + pageInfo.SetItems(buildMaskedTokenResponses(tokens)) common.ApiSuccess(c, pageInfo) - return } func GetToken(c *gin.Context) { @@ -59,12 +74,24 @@ func GetToken(c *gin.Context) { common.ApiError(c, err) return } - c.JSON(http.StatusOK, gin.H{ - "success": true, - "message": "", - "data": token, + common.ApiSuccess(c, buildMaskedTokenResponse(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) { @@ -204,7 +231,6 @@ func AddToken(c *gin.Context) { "success": true, "message": "", }) - return } func DeleteToken(c *gin.Context) { @@ -219,7 +245,6 @@ func DeleteToken(c *gin.Context) { "success": true, "message": "", }) - return } func UpdateToken(c *gin.Context) { @@ -283,7 +308,7 @@ func UpdateToken(c *gin.Context) { c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", - "data": cleanToken, + "data": buildMaskedTokenResponse(cleanToken), }) } diff --git a/controller/token_test.go b/controller/token_test.go new file mode 100644 index 000000000..3eea67305 --- /dev/null +++ b/controller/token_test.go @@ -0,0 +1,275 @@ +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()) + } +} diff --git a/model/token.go b/model/token.go index 1c9ad3ed7..91e5fe1da 100644 --- a/model/token.go +++ b/model/token.go @@ -35,6 +35,27 @@ func (token *Token) Clean() { 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 { // delete empty spaces //split with \n @@ -201,7 +222,7 @@ func ValidateUserToken(key string) (token *Token, err error) { } keyPrefix := 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 } diff --git a/router/api-router.go b/router/api-router.go index fafb99575..9836083df 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -248,6 +248,7 @@ func SetApiRouter(router *gin.Engine) { tokenRoute.GET("/", controller.GetAllTokens) tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens) tokenRoute.GET("/:id", controller.GetToken) + tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKey) tokenRoute.POST("/", controller.AddToken) tokenRoute.PUT("/", controller.UpdateToken) tokenRoute.DELETE("/:id", controller.DeleteToken) diff --git a/web/src/components/table/tokens/TokensActions.jsx b/web/src/components/table/tokens/TokensActions.jsx index 2fbb2fce2..0d6ca9b7a 100644 --- a/web/src/components/table/tokens/TokensActions.jsx +++ b/web/src/components/table/tokens/TokensActions.jsx @@ -29,7 +29,6 @@ const TokensActions = ({ setShowEdit, batchCopyTokens, batchDeleteTokens, - copyText, t, }) => { // Modal states @@ -99,8 +98,7 @@ const TokensActions = ({ setShowCopyModal(false)} - selectedKeys={selectedKeys} - copyText={copyText} + batchCopyTokens={batchCopyTokens} t={t} /> diff --git a/web/src/components/table/tokens/TokensColumnDefs.jsx b/web/src/components/table/tokens/TokensColumnDefs.jsx index ce8eab807..86c4545ad 100644 --- a/web/src/components/table/tokens/TokensColumnDefs.jsx +++ b/web/src/components/table/tokens/TokensColumnDefs.jsx @@ -108,17 +108,28 @@ const renderGroupColumn = (text, record, t) => { }; // Render token key column with show/hide and copy functionality -const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => { - const fullKey = 'sk-' + record.key; - const maskedKey = - 'sk-' + record.key.slice(0, 4) + '**********' + record.key.slice(-4); +const renderTokenKey = ( + text, + record, + showKeys, + resolvedTokenKeys, + loadingTokenKeys, + toggleTokenVisibility, + copyTokenKey, +) => { 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 (
@@ -127,10 +138,11 @@ const renderTokenKey = (text, record, showKeys, setShowKeys, copyText) => { size='small' type='tertiary' icon={revealed ? : } + loading={loading} aria-label='toggle token visibility' - onClick={(e) => { + onClick={async (e) => { e.stopPropagation(); - setShowKeys((prev) => ({ ...prev, [record.id]: !revealed })); + await toggleTokenVisibility(record); }} />
@@ -427,8 +440,10 @@ const renderOperations = ( export const getTokensColumns = ({ t, showKeys, - setShowKeys, - copyText, + resolvedTokenKeys, + loadingTokenKeys, + toggleTokenVisibility, + copyTokenKey, manageToken, onOpenLink, setEditingToken, @@ -461,7 +476,15 @@ export const getTokensColumns = ({ title: t('密钥'), key: 'token_key', render: (text, record) => - renderTokenKey(text, record, showKeys, setShowKeys, copyText), + renderTokenKey( + text, + record, + showKeys, + resolvedTokenKeys, + loadingTokenKeys, + toggleTokenVisibility, + copyTokenKey, + ), }, { title: t('可用模型'), diff --git a/web/src/components/table/tokens/TokensTable.jsx b/web/src/components/table/tokens/TokensTable.jsx index eab5707cd..a62b6f270 100644 --- a/web/src/components/table/tokens/TokensTable.jsx +++ b/web/src/components/table/tokens/TokensTable.jsx @@ -39,8 +39,10 @@ const TokensTable = (tokensData) => { rowSelection, handleRow, showKeys, - setShowKeys, - copyText, + resolvedTokenKeys, + loadingTokenKeys, + toggleTokenVisibility, + copyTokenKey, manageToken, onOpenLink, setEditingToken, @@ -54,8 +56,10 @@ const TokensTable = (tokensData) => { return getTokensColumns({ t, showKeys, - setShowKeys, - copyText, + resolvedTokenKeys, + loadingTokenKeys, + toggleTokenVisibility, + copyTokenKey, manageToken, onOpenLink, setEditingToken, @@ -65,8 +69,10 @@ const TokensTable = (tokensData) => { }, [ t, showKeys, - setShowKeys, - copyText, + resolvedTokenKeys, + loadingTokenKeys, + toggleTokenVisibility, + copyTokenKey, manageToken, onOpenLink, setEditingToken, diff --git a/web/src/components/table/tokens/index.jsx b/web/src/components/table/tokens/index.jsx index c78d8b82d..1711ecde7 100644 --- a/web/src/components/table/tokens/index.jsx +++ b/web/src/components/table/tokens/index.jsx @@ -58,6 +58,7 @@ function TokensPage() { t: (k) => k, selectedModel: '', prefillKey: '', + fetchTokenKey: async () => '', }); const [modelOptions, setModelOptions] = useState([]); const [selectedModel, setSelectedModel] = useState(''); @@ -74,6 +75,7 @@ function TokensPage() { t: tokensData.t, selectedModel, prefillKey, + fetchTokenKey: tokensData.fetchTokenKey, }; }, [ tokensData.tokens, @@ -81,6 +83,7 @@ function TokensPage() { tokensData.t, selectedModel, prefillKey, + tokensData.fetchTokenKey, ]); const loadModels = async () => { @@ -198,13 +201,14 @@ function TokensPage() { openCCSwitchModalRef.current = openCCSwitchModal; // Prefill to Fluent handler - const handlePrefillToFluent = () => { + const handlePrefillToFluent = async () => { const { tokens, selectedKeys, t, selectedModel: chosenModel, prefillKey: overrideKey, + fetchTokenKey, } = latestRef.current; const container = document.getElementById('fluent-new-api-container'); if (!container) { @@ -241,7 +245,11 @@ function TokensPage() { Toast.warning(t('没有可用令牌用于填充')); return; } - apiKeyToUse = 'sk-' + token.key; + try { + apiKeyToUse = 'sk-' + (await fetchTokenKey(token)); + } catch (_) { + return; + } } const payload = { @@ -351,7 +359,6 @@ function TokensPage() { setShowEdit, batchCopyTokens, batchDeleteTokens, - copyText, // Filters state formInitValues, @@ -401,7 +408,6 @@ function TokensPage() { setShowEdit={setShowEdit} batchCopyTokens={batchCopyTokens} batchDeleteTokens={batchDeleteTokens} - copyText={copyText} t={t} /> diff --git a/web/src/components/table/tokens/modals/CCSwitchModal.jsx b/web/src/components/table/tokens/modals/CCSwitchModal.jsx index b1567e8cf..6cf817330 100644 --- a/web/src/components/table/tokens/modals/CCSwitchModal.jsx +++ b/web/src/components/table/tokens/modals/CCSwitchModal.jsx @@ -116,8 +116,7 @@ export default function CCSwitchModal({ Toast.warning(t('请选择主模型')); return; } - const apiKey = 'sk-' + tokenKey; - const url = buildCCSwitchURL(app, name, models, apiKey); + const url = buildCCSwitchURL(app, name, models, 'sk-' + tokenKey); window.open(url, '_blank'); onClose(); }; diff --git a/web/src/components/table/tokens/modals/CopyTokensModal.jsx b/web/src/components/table/tokens/modals/CopyTokensModal.jsx index 09f37efea..bb8ae9657 100644 --- a/web/src/components/table/tokens/modals/CopyTokensModal.jsx +++ b/web/src/components/table/tokens/modals/CopyTokensModal.jsx @@ -20,24 +20,21 @@ For commercial licensing, please contact support@quantumnous.com import React from 'react'; 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 const handleCopyWithName = async () => { - let content = ''; - for (let i = 0; i < selectedKeys.length; i++) { - content += selectedKeys[i].name + ' sk-' + selectedKeys[i].key + '\n'; - } - await copyText(content); + await batchCopyTokens('name+key'); onCancel(); }; // Handle copy with key only format const handleCopyKeyOnly = async () => { - let content = ''; - for (let i = 0; i < selectedKeys.length; i++) { - content += 'sk-' + selectedKeys[i].key + '\n'; - } - await copyText(content); + await batchCopyTokens('key-only'); onCancel(); }; diff --git a/web/src/components/table/usage-logs/modals/ColumnSelectorModal.jsx b/web/src/components/table/usage-logs/modals/ColumnSelectorModal.jsx index 20f0fee0b..ca1483db8 100644 --- a/web/src/components/table/usage-logs/modals/ColumnSelectorModal.jsx +++ b/web/src/components/table/usage-logs/modals/ColumnSelectorModal.jsx @@ -73,7 +73,7 @@ const ColumnSelectorModal = ({ setBillingDisplayMode(event.target.value)} + onChange={(value) => setBillingDisplayMode(value)} > {isTokensDisplay ? t('价格模式') : t('价格模式（默认）')} diff --git a/web/src/helpers/token.js b/web/src/helpers/token.js index f37e4ebde..e4812be97 100644 --- a/web/src/helpers/token.js +++ b/web/src/helpers/token.js @@ -20,8 +20,22 @@ For commercial licensing, please contact support@quantumnous.com import { API } from './api'; /** - * 获取可用的token keys - * @returns {Promise} 返回active状态的token key数组 + * 按需获取单个令牌的真实 key + * @param {number|string} tokenId + * @returns {Promise} 返回不带 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} 返回 active 状态的不带 sk- 前缀的真实 token key 数组 */ export async function fetchTokenKeys() { try { @@ -31,7 +45,12 @@ export async function fetchTokenKeys() { const tokenItems = Array.isArray(data) ? data : data.items || []; 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) { console.error('Error fetching token keys:', error); return []; diff --git a/web/src/hooks/tokens/useTokensData.jsx b/web/src/hooks/tokens/useTokensData.jsx index 2a4692d69..7015b8401 100644 --- a/web/src/hooks/tokens/useTokensData.jsx +++ b/web/src/hooks/tokens/useTokensData.jsx @@ -17,7 +17,7 @@ along with this program. If not, see . 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 { Modal } from '@douyinfe/semi-ui'; import { @@ -29,6 +29,7 @@ import { } from '../../helpers'; import { ITEMS_PER_PAGE } from '../../constants'; import { useTableCompactMode } from '../common/useTableCompactMode'; +import { fetchTokenKey as fetchTokenKeyById } from '../../helpers/token'; export const useTokensData = (openFluentNotification, openCCSwitchModal) => { const { t } = useTranslation(); @@ -54,6 +55,9 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { // UI state const [compactMode, setCompactMode] = useTableCompactMode('tokens'); const [showKeys, setShowKeys] = useState({}); + const [resolvedTokenKeys, setResolvedTokenKeys] = useState({}); + const [loadingTokenKeys, setLoadingTokenKeys] = useState({}); + const keyRequestsRef = useRef({}); // Form state const [formApi, setFormApi] = useState(null); @@ -87,6 +91,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { setTokenCount(payload.total || 0); setActivePage(payload.page || 1); setPageSize(payload.page_size || pageSize); + setShowKeys({}); }; // Load tokens function @@ -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 const onOpenLink = async (type, url, record) => { + const fullKey = await fetchTokenKey(record); if (url && url.startsWith('ccswitch')) { - openCCSwitchModal(record.key); + openCCSwitchModal(fullKey); return; } if (url && url.startsWith('fluent')) { - openFluentNotification(record.key); + openFluentNotification(fullKey); return; } let status = localStorage.getItem('status'); @@ -145,7 +222,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { let cherryConfig = { id: 'new-api', baseUrl: serverAddress, - apiKey: 'sk-' + record.key, + apiKey: `sk-${fullKey}`, }; let encodedConfig = encodeURIComponent( encodeToBase64(JSON.stringify(cherryConfig)), @@ -155,7 +232,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { let aionuiConfig = { platform: 'new-api', baseUrl: serverAddress, - apiKey: 'sk-' + record.key, + apiKey: `sk-${fullKey}`, }; let encodedConfig = encodeURIComponent( encodeToBase64(JSON.stringify(aionuiConfig)), @@ -164,7 +241,7 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { } else { let encodedServerAddress = encodeURIComponent(serverAddress); url = url.replaceAll('{address}', encodedServerAddress); - url = url.replaceAll('{key}', 'sk-' + record.key); + url = url.replaceAll('{key}', `sk-${fullKey}`); } window.open(url, '_blank'); @@ -314,48 +391,28 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { }; // Batch copy tokens - const batchCopyTokens = (copyType) => { + const batchCopyTokens = async (copyType) => { if (selectedKeys.length === 0) { showError(t('请至少选择一个令牌！')); return; } - - Modal.info({ - title: t('复制令牌'), - icon: null, - content: t('请选择你的复制方式'), - footer: ( -
- - -
- ), - }); + try { + const keys = await Promise.all( + selectedKeys.map((token) => fetchTokenKey(token, { suppressError: true })), + ); + let content = ''; + for (let i = 0; i < selectedKeys.length; i++) { + const fullKey = keys[i]; + if (copyType === 'name+key') { + content += `${selectedKeys[i].name} sk-${fullKey}\n`; + } else { + content += `sk-${fullKey}\n`; + } + } + await copyText(content); + } catch (error) { + showError(error?.message || t('复制令牌失败')); + } }; // Initialize data @@ -392,6 +449,8 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { setCompactMode, showKeys, setShowKeys, + resolvedTokenKeys, + loadingTokenKeys, // Form state formApi, @@ -403,6 +462,9 @@ export const useTokensData = (openFluentNotification, openCCSwitchModal) => { loadTokens, refresh, copyText, + fetchTokenKey, + toggleTokenVisibility, + copyTokenKey, onOpenLink, manageToken, searchTokens,