diff --git a/constant/context_key.go b/constant/context_key.go index b856bc3d..ccb8010f 100644 --- a/constant/context_key.go +++ b/constant/context_key.go @@ -19,6 +19,7 @@ const ( ContextKeyTokenModelLimitEnabled ContextKey = "token_model_limit_enabled" ContextKeyTokenModelLimit ContextKey = "token_model_limit" ContextKeyTokenCrossGroupRetry ContextKey = "token_cross_group_retry" + ContextKeyTokenAutoGroups ContextKey = "token_auto_groups" /* channel related keys */ ContextKeyChannelId ContextKey = "channel_id" diff --git a/controller/model.go b/controller/model.go index b32eebd7..1d759301 100644 --- a/controller/model.go +++ b/controller/model.go @@ -20,6 +20,7 @@ import ( "github.com/QuantumNous/new-api/relaykit/types" "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/samber/lo" ) @@ -190,7 +191,7 @@ func getModelListGroups(c *gin.Context) (modelListGroups, error) { return modelListGroups{ userGroup: userGroup, tokenGroup: tokenGroup, - ownerGroups: service.GetUserAutoGroup(userGroup), + ownerGroups: service.GetRequestAutoGroups(c, userGroup), }, nil } @@ -228,32 +229,28 @@ func ListModels(c *gin.Context, modelType int) { } ownerGroups := groups.ownerGroups modelLimitEnable := common.GetContextKeyBool(c, constant.ContextKeyTokenModelLimitEnabled) + var tokenModelLimit map[string]bool if modelLimitEnable { s, ok := common.GetContextKey(c, constant.ContextKeyTokenModelLimit) - var tokenModelLimit map[string]bool if ok { - tokenModelLimit = s.(map[string]bool) - } else { + tokenModelLimit, _ = s.(map[string]bool) + } + if tokenModelLimit == nil { tokenModelLimit = map[string]bool{} } - for allowModel, _ := range tokenModelLimit { - if !acceptUnsetRatioModel { - if !helper.HasModelBillingConfig(allowModel) { - continue - } + } + models := service.GetGroupsEnabledModels(ownerGroups) + for _, modelName := range models { + if modelLimitEnable { + matchingName := ratio_setting.FormatMatchingModelName(modelName) + if !tokenModelLimit[modelName] && !tokenModelLimit[matchingName] { + continue } - userModelNames = append(userModelNames, allowModel) } - } else { - models := service.GetGroupsEnabledModels(ownerGroups) - for _, modelName := range models { - if !acceptUnsetRatioModel { - if !helper.HasModelBillingConfig(modelName) { - continue - } - } - userModelNames = append(userModelNames, modelName) + if !acceptUnsetRatioModel && !helper.HasModelBillingConfig(modelName) { + continue } + userModelNames = append(userModelNames, modelName) } ownerByModel := map[string]string{} @@ -276,11 +273,17 @@ func ListModels(c *gin.Context, modelType int) { Type: "model", } } + firstID := "" + lastID := "" + if len(useranthropicModels) > 0 { + firstID = useranthropicModels[0].ID + lastID = useranthropicModels[len(useranthropicModels)-1].ID + } c.JSON(200, gin.H{ "data": useranthropicModels, - "first_id": useranthropicModels[0].ID, + "first_id": firstID, "has_more": false, - "last_id": useranthropicModels[len(useranthropicModels)-1].ID, + "last_id": lastID, }) case constant.ChannelTypeGemini: userGeminiModels := make([]dto.GeminiModel, len(userOpenAiModels)) diff --git a/controller/model_list_test.go b/controller/model_list_test.go index b1fa9b95..812207b8 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -402,7 +402,13 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { "zz-token-tiered-visible-model": `tier("base", p * 1 + c * 2)`, "zz-token-tiered-empty-expr-model": "", }) - setupModelListControllerTestDB(t) + db := setupModelListControllerTestDB(t) + require.NoError(t, db.Create(&[]model.Ability{ + {Group: "default", Model: "zz-token-tiered-visible-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-tiered-empty-expr-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-tiered-missing-expr-model", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-token-unpriced-model", ChannelId: 1, Enabled: true}, + }).Error) recorder := httptest.NewRecorder() ctx, _ := gin.CreateTestContext(recorder) @@ -425,6 +431,68 @@ func TestListModelsTokenLimitIncludesTieredBillingModel(t *testing.T) { require.NotContains(t, ids, "zz-token-unpriced-model") } +func TestListModelsTokenLimitUsesResolvedCustomAutoGroups(t *testing.T) { + withSelfUseModeEnabled(t) + originalMax := setting.GetMaxTokenAutoGroups() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("5")) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) + + db := setupModelListControllerTestDB(t) + require.NoError(t, db.Create(&[]model.Ability{ + {Group: "vip", Model: "zz-vip-allowed", ChannelId: 1, Enabled: true}, + {Group: "vip", Model: "zz-vip-denied", ChannelId: 1, Enabled: true}, + {Group: "default", Model: "zz-default-outside-snapshot", ChannelId: 1, Enabled: true}, + }).Error) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + common.SetContextKey(ctx, constant.ContextKeyTokenModelLimitEnabled, true) + common.SetContextKey(ctx, constant.ContextKeyTokenModelLimit, map[string]bool{ + "zz-vip-allowed": true, + "zz-default-outside-snapshot": true, + "zz-not-enabled": true, + }) + + ListModels(ctx, constant.ChannelTypeOpenAI) + ids := decodeListModelsResponse(t, recorder) + require.Equal(t, map[string]struct{}{"zz-vip-allowed": {}}, ids) + + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + emptyRecorder := httptest.NewRecorder() + emptyCtx, _ := gin.CreateTestContext(emptyRecorder) + emptyCtx.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + common.SetContextKey(emptyCtx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(emptyCtx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(emptyCtx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimitEnabled, true) + common.SetContextKey(emptyCtx, constant.ContextKeyTokenModelLimit, map[string]bool{"zz-vip-allowed": true}) + + require.NotPanics(t, func() { + ListModels(emptyCtx, constant.ChannelTypeAnthropic) + }) + var anthropicResponse struct { + Data []dto.AnthropicModel `json:"data"` + FirstID string `json:"first_id"` + LastID string `json:"last_id"` + } + require.NoError(t, common.Unmarshal(emptyRecorder.Body.Bytes(), &anthropicResponse)) + require.Empty(t, anthropicResponse.Data) + require.Empty(t, anthropicResponse.FirstID) + require.Empty(t, anthropicResponse.LastID) +} + func TestCheckUpdatePasswordRequiresCurrentPassword(t *testing.T) { db := setupModelListControllerTestDB(t) hashedPassword, err := common.Password2Hash("CurrentPassword123") diff --git a/controller/model_owned_by_test.go b/controller/model_owned_by_test.go index bc2ef32f..da9bb1a0 100644 --- a/controller/model_owned_by_test.go +++ b/controller/model_owned_by_test.go @@ -1,11 +1,14 @@ package controller import ( + "fmt" "net/http/httptest" "testing" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" "github.com/gin-gonic/gin" "github.com/stretchr/testify/require" ) @@ -83,3 +86,33 @@ func TestGetModelListGroupsUsesExplicitTokenGroup(t *testing.T) { require.Equal(t, "vip", groups.tokenGroup) require.Equal(t, []string{"vip"}, groups.ownerGroups) } + +func TestGetModelListGroupsUsesFilteredTokenAutoGroupsSnapshot(t *testing.T) { + originalMax := setting.GetMaxTokenAutoGroups() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("1")) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenGroup, "auto") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"}) + + groups, err := getModelListGroups(ctx) + require.NoError(t, err) + require.Equal(t, []string{"vip"}, groups.ownerGroups) + + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + groups, err = getModelListGroups(ctx) + require.NoError(t, err) + require.Empty(t, groups.ownerGroups) +} diff --git a/controller/token.go b/controller/token.go index 836e9b29..c26d82e3 100644 --- a/controller/token.go +++ b/controller/token.go @@ -7,30 +7,115 @@ import ( "strings" "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/operation_setting" "github.com/gin-gonic/gin" ) -func buildMaskedTokenResponse(token *model.Token) *model.Token { +type tokenAutoGroupsInput struct { + Set bool + Groups []string +} + +func (input *tokenAutoGroupsInput) UnmarshalJSON(data []byte) error { + input.Set = true + if strings.TrimSpace(string(data)) == "null" { + input.Groups = nil + return nil + } + return common.Unmarshal(data, &input.Groups) +} + +type tokenRequest struct { + model.Token + AutoGroups tokenAutoGroupsInput `json:"auto_groups"` +} + +type tokenResponse struct { + *model.Token + AutoGroups []string `json:"auto_groups"` +} + +func buildMaskedTokenResponse(token *model.Token) *tokenResponse { if token == nil { return nil } maskedToken := *token maskedToken.Key = token.GetMaskedKey() - return &maskedToken + autoGroups, err := token.GetAutoGroups() + if err != nil { + common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) + autoGroups = nil + } + if len(autoGroups) == 0 { + autoGroups = nil + } + return &tokenResponse{Token: &maskedToken, AutoGroups: autoGroups} } -func buildMaskedTokenResponses(tokens []*model.Token) []*model.Token { - maskedTokens := make([]*model.Token, 0, len(tokens)) +func buildMaskedTokenResponses(tokens []*model.Token) []*tokenResponse { + maskedTokens := make([]*tokenResponse, 0, len(tokens)) for _, token := range tokens { maskedTokens = append(maskedTokens, buildMaskedTokenResponse(token)) } return maskedTokens } +func getTokenRequestUserGroup(c *gin.Context) (string, error) { + if userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup); userGroup != "" { + return userGroup, nil + } + if userGroup := c.GetString("group"); userGroup != "" { + return userGroup, nil + } + return model.GetUserGroup(c.GetInt("id"), false) +} + +func setTokenAutoGroups(c *gin.Context, token *model.Token, groups []string) bool { + if len(groups) == 0 { + if err := token.SetAutoGroups(nil); err != nil { + common.ApiError(c, err) + return false + } + return true + } + + maxCount := setting.GetMaxTokenAutoGroups() + if len(groups) > maxCount { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsTooMany, map[string]any{"Max": maxCount}) + return false + } + + userGroup, err := getTokenRequestUserGroup(c) + if err != nil { + common.ApiError(c, err) + return false + } + seen := make(map[string]struct{}, len(groups)) + for _, group := range groups { + if _, ok := seen[group]; ok { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsDuplicate, map[string]any{"Group": group}) + return false + } + seen[group] = struct{}{} + if !service.IsUserSelectableGroup(userGroup, group) { + common.ApiErrorI18n(c, i18n.MsgTokenAutoGroupsInvalid, map[string]any{"Group": group}) + return false + } + } + + if err := token.SetAutoGroups(groups); err != nil { + common.ApiError(c, err) + return false + } + return true +} + func GetAllTokens(c *gin.Context) { userId := c.GetInt("id") pageInfo := common.GetPageQuery(c) @@ -77,6 +162,18 @@ func GetToken(c *gin.Context) { common.ApiSuccess(c, buildMaskedTokenResponse(token)) } +func GetTokenAutoGroups(c *gin.Context) { + userGroup, err := getTokenRequestUserGroup(c) + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "groups": service.GetUserAutoGroup(userGroup), + "max_count": setting.GetMaxTokenAutoGroups(), + }) +} + func GetTokenKey(c *gin.Context) { id, err := strconv.Atoi(c.Param("id")) userId := c.GetInt("id") @@ -165,12 +262,13 @@ func GetTokenUsage(c *gin.Context) { } func AddToken(c *gin.Context) { - token := model.Token{} - err := c.ShouldBindJSON(&token) + request := tokenRequest{} + err := c.ShouldBindJSON(&request) if err != nil { common.ApiError(c, err) return } + token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -201,6 +299,14 @@ func AddToken(c *gin.Context) { }) return } + if token.Group == "auto" { + if !setTokenAutoGroups(c, &token, request.AutoGroups.Groups) { + return + } + } else { + token.CrossGroupRetry = false + _ = token.SetAutoGroups(nil) + } key, err := common.GenerateKey() if err != nil { common.ApiErrorI18n(c, i18n.MsgTokenGenerateFailed) @@ -221,6 +327,7 @@ func AddToken(c *gin.Context) { AllowIps: token.AllowIps, Group: token.Group, CrossGroupRetry: token.CrossGroupRetry, + AutoGroups: token.AutoGroups, } err = cleanToken.Insert() if err != nil { @@ -250,12 +357,13 @@ 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) + request := tokenRequest{} + err := c.ShouldBindJSON(&request) if err != nil { common.ApiError(c, err) return } + token := request.Token if len(token.Name) > 50 { common.ApiErrorI18n(c, i18n.MsgTokenNameTooLong) return @@ -299,6 +407,14 @@ func UpdateToken(c *gin.Context) { cleanToken.AllowIps = token.AllowIps cleanToken.Group = token.Group cleanToken.CrossGroupRetry = token.CrossGroupRetry + if token.Group != "auto" { + cleanToken.CrossGroupRetry = false + _ = cleanToken.SetAutoGroups(nil) + } else if request.AutoGroups.Set { + if !setTokenAutoGroups(c, cleanToken, request.AutoGroups.Groups) { + return + } + } } err = cleanToken.Update() if err != nil { diff --git a/controller/token_auto_groups_test.go b/controller/token_auto_groups_test.go new file mode 100644 index 00000000..3da29695 --- /dev/null +++ b/controller/token_auto_groups_test.go @@ -0,0 +1,234 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func configureTokenAutoGroupsTest(t *testing.T, maxCount string, autoGroups string) { + t.Helper() + originalMax := setting.GetMaxTokenAutoGroups() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups(maxCount)) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(autoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(stringInt(originalMax))) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) +} + +func stringInt(value int) string { + return fmt.Sprintf("%d", value) +} + +func setupTokenAutoGroupsControllerTest(t *testing.T) *model.User { + t.Helper() + db := setupTokenControllerTestDB(t) + require.NoError(t, db.AutoMigrate(&model.User{})) + user := &model.User{ + Id: 101, + Username: "token-auto-user", + Password: "password", + Group: "default", + Status: common.UserStatusEnabled, + } + require.NoError(t, db.Create(user).Error) + return user +} + +func baseAutoTokenRequest(name string) map[string]any { + return map[string]any{ + "name": name, + "expired_time": -1, + "remain_quota": 0, + "unlimited_quota": true, + "group": "auto", + "cross_group_retry": true, + } +} + +func newTokenAutoGroupsAuthenticatedContext(t *testing.T, method string, target string, body any, userID int) (*gin.Context, *httptest.ResponseRecorder) { + t.Helper() + ctx, recorder := newAuthenticatedContext(t, method, target, body, userID) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + return ctx, recorder +} + +func TestAddTokenEmptyAutoGroupsInheritGlobalAuto(t *testing.T) { + tests := []struct { + name string + includeField bool + value any + }{ + {name: "omitted"}, + {name: "null", includeField: true, value: nil}, + {name: "empty array", includeField: true, value: []string{}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("create-" + test.name) + if test.includeField { + request["auto_groups"] = test.value + } + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + var token model.Token + require.NoError(t, model.DB.Where("name = ?", request["name"]).First(&token).Error) + assert.Empty(t, token.AutoGroups) + assert.True(t, token.CrossGroupRetry) + payload, err := common.Marshal(buildMaskedTokenResponse(&token)) + require.NoError(t, err) + var responseData map[string]any + require.NoError(t, common.Unmarshal(payload, &responseData)) + assert.Nil(t, responseData["auto_groups"]) + }) + } +} + +func TestAddTokenPersistsOrderedAutoGroupsSnapshot(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("ordered-snapshot") + request["auto_groups"] = []string{"vip", "default"} + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + require.True(t, decodeAPIResponse(t, recorder).Success) + + var token model.Token + require.NoError(t, model.DB.Where("name = ?", "ordered-snapshot").First(&token).Error) + assert.JSONEq(t, `["vip","default"]`, token.AutoGroups) + + getCtx, getRecorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/"+stringInt(token.Id), nil, user.Id) + getCtx.Params = append(getCtx.Params, gin.Param{Key: "id", Value: stringInt(token.Id)}) + GetToken(getCtx) + getResponse := decodeAPIResponse(t, getRecorder) + require.True(t, getResponse.Success) + var data struct { + AutoGroups []string `json:"auto_groups"` + } + require.NoError(t, common.Unmarshal(getResponse.Data, &data)) + assert.Equal(t, []string{"vip", "default"}, data.AutoGroups) +} + +func TestUpdateTokenAutoGroupsTriStateAndNonAutoCleanup(t *testing.T) { + tests := []struct { + name string + includeField bool + value any + group string + expectedAutoGroups string + expectedRetry bool + }{ + {name: "omitted preserves", group: "auto", expectedAutoGroups: `["vip","default"]`, expectedRetry: true}, + {name: "null inherits", includeField: true, value: nil, group: "auto", expectedRetry: true}, + {name: "empty inherits", includeField: true, value: []string{}, group: "auto", expectedRetry: true}, + {name: "non auto clears and disables retry", includeField: true, value: []string{"vip"}, group: "default"}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, "5", `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + token := seedToken(t, model.DB, user.Id, "update-auto", "update-auto-key") + token.Group = "auto" + token.CrossGroupRetry = true + require.NoError(t, token.SetAutoGroups([]string{"vip", "default"})) + require.NoError(t, model.DB.Save(token).Error) + + request := baseAutoTokenRequest("updated-auto") + request["id"] = token.Id + request["status"] = common.TokenStatusEnabled + request["group"] = test.group + if test.includeField { + request["auto_groups"] = test.value + } + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPut, "/api/token/", request, user.Id) + UpdateToken(ctx) + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + + var updated model.Token + require.NoError(t, model.DB.First(&updated, token.Id).Error) + if test.expectedAutoGroups == "" { + assert.Empty(t, updated.AutoGroups) + } else { + assert.JSONEq(t, test.expectedAutoGroups, updated.AutoGroups) + } + assert.Equal(t, test.expectedRetry, updated.CrossGroupRetry) + }) + } +} + +func TestAddTokenRejectsInvalidAutoGroups(t *testing.T) { + tests := []struct { + name string + maxCount string + groups []string + }{ + {name: "over limit", maxCount: "1", groups: []string{"default", "vip"}}, + {name: "duplicate", maxCount: "5", groups: []string{"default", "default"}}, + {name: "auto pseudo group", maxCount: "5", groups: []string{"auto"}}, + {name: "unavailable", maxCount: "5", groups: []string{"missing"}}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + configureTokenAutoGroupsTest(t, test.maxCount, `["default","vip"]`) + user := setupTokenAutoGroupsControllerTest(t) + request := baseAutoTokenRequest("invalid-" + test.name) + request["auto_groups"] = test.groups + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodPost, "/api/token/", request, user.Id) + AddToken(ctx) + + response := decodeAPIResponse(t, recorder) + assert.False(t, response.Success) + var count int64 + require.NoError(t, model.DB.Model(&model.Token{}).Count(&count).Error) + assert.Zero(t, count) + }) + } +} + +func TestGetTokenAutoGroupsReturnsFullFilteredGlobalOrderAndLimit(t *testing.T) { + configureTokenAutoGroupsTest(t, "1", `["vip","missing","default"]`) + user := setupTokenAutoGroupsControllerTest(t) + + ctx, recorder := newTokenAutoGroupsAuthenticatedContext(t, http.MethodGet, "/api/token/auto-groups", nil, user.Id) + GetTokenAutoGroups(ctx) + + response := decodeAPIResponse(t, recorder) + require.True(t, response.Success, response.Message) + var data struct { + Groups []string `json:"groups"` + MaxCount int `json:"max_count"` + } + require.NoError(t, common.Unmarshal(response.Data, &data)) + assert.Equal(t, []string{"vip", "default"}, data.Groups) + assert.Equal(t, 1, data.MaxCount) +} diff --git a/controller/token_test.go b/controller/token_test.go index 12b1cbdd..9cca168a 100644 --- a/controller/token_test.go +++ b/controller/token_test.go @@ -273,6 +273,34 @@ func getTokenKeyColumnType(t *testing.T, db *gorm.DB, dialect string) string { } } +func getTokenAutoGroupsColumnType(t *testing.T, db *gorm.DB, dialect string) string { + t.Helper() + + switch dialect { + case "sqlite": + return getSQLiteColumnType(t, db, "tokens", "auto_groups") + case "mysql": + var columnType string + if err := db.Raw(`SELECT DATA_TYPE FROM information_schema.columns + WHERE table_schema = DATABASE() AND table_name = ? AND column_name = ?`, + "tokens", "auto_groups").Scan(&columnType).Error; err != nil { + t.Fatalf("failed to inspect mysql token auto_groups column: %v", err) + } + return strings.ToLower(columnType) + case "postgres": + var dataType string + if err := db.Raw(`SELECT data_type FROM information_schema.columns + WHERE table_schema = current_schema() AND table_name = ? AND column_name = ?`, + "tokens", "auto_groups").Scan(&dataType).Error; err != nil { + t.Fatalf("failed to inspect postgres token auto_groups column: %v", err) + } + return strings.ToLower(dataType) + default: + t.Fatalf("unsupported dialect %q", dialect) + return "" + } +} + func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect string, managedTokensTable *bool) { t.Helper() @@ -314,6 +342,12 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if got := getTokenKeyColumnType(t, db, dialect); got != "varchar(128)" { t.Fatalf("expected migrated key column type varchar(128), got %q", got) } + if !db.Migrator().HasColumn(&model.Token{}, "auto_groups") { + t.Fatal("expected migration to add auto_groups column") + } + if got := getTokenAutoGroupsColumnType(t, db, dialect); got != "text" { + t.Fatalf("expected migrated auto_groups column type text, got %q", got) + } var migratedToken model.Token if err := db.First(&migratedToken, "name = ?", "legacy-token").Error; err != nil { @@ -325,6 +359,9 @@ func runTokenMigrationCompatibilityTest(t *testing.T, db *gorm.DB, dialect strin if migratedToken.Name != "legacy-token" { t.Fatalf("expected migrated token name to be preserved, got %q", migratedToken.Name) } + if migratedToken.AutoGroups != "" { + t.Fatalf("expected legacy token to inherit global Auto groups, got %q", migratedToken.AutoGroups) + } inserted := model.Token{ UserId: 8, @@ -362,6 +399,9 @@ func TestTokenAutoMigrateUsesVarchar128KeyColumn(t *testing.T) { if got := getTokenKeyColumnType(t, db, "sqlite"); got != "varchar(128)" { t.Fatalf("expected key column type varchar(128), got %q", got) } + if got := getSQLiteColumnType(t, db, "tokens", "auto_groups"); got != "text" { + t.Fatalf("expected auto_groups column type text, got %q", got) + } } func TestTokenMigrationFromChar48ToVarchar128(t *testing.T) { diff --git a/i18n/keys.go b/i18n/keys.go index 8e9a4b56..64a835e1 100644 --- a/i18n/keys.go +++ b/i18n/keys.go @@ -55,6 +55,9 @@ const ( MsgTokenExhausted = "token.exhausted" MsgTokenStatusUnavailable = "token.status_unavailable" MsgTokenDbError = "token.db_error" + MsgTokenAutoGroupsTooMany = "token.auto_groups_too_many" + MsgTokenAutoGroupsDuplicate = "token.auto_groups_duplicate" + MsgTokenAutoGroupsInvalid = "token.auto_groups_invalid" ) // Redemption related messages diff --git a/i18n/locales/en.yaml b/i18n/locales/en.yaml index 3f1fd03c..c533daec 100644 --- a/i18n/locales/en.yaml +++ b/i18n/locales/en.yaml @@ -47,6 +47,9 @@ token.expired: "This token has expired" token.exhausted: "This token quota is exhausted TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "This token status is unavailable" token.db_error: "Invalid token, database query error, please contact administrator" +token.auto_groups_too_many: "A token can select at most {{.Max}} Auto groups" +token.auto_groups_duplicate: "Auto group {{.Group}} is duplicated" +token.auto_groups_invalid: "Auto group {{.Group}} is unavailable or unauthorized" # Redemption messages redemption.name_length: "Redemption code name length must be between 1-20" diff --git a/i18n/locales/zh-CN.yaml b/i18n/locales/zh-CN.yaml index fe982e59..a2f5275b 100644 --- a/i18n/locales/zh-CN.yaml +++ b/i18n/locales/zh-CN.yaml @@ -48,6 +48,9 @@ token.expired: "该令牌已过期" token.exhausted: "该令牌额度已用尽 TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "该令牌状态不可用" token.db_error: "无效的令牌，数据库查询出错，请联系管理员" +token.auto_groups_too_many: "每个令牌最多可选择 {{.Max}} 个 Auto 分组" +token.auto_groups_duplicate: "Auto 分组 {{.Group}} 重复" +token.auto_groups_invalid: "Auto 分组 {{.Group}} 不可用或无权访问" # Redemption messages redemption.name_length: "兑换码名称长度必须在1-20之间" diff --git a/i18n/locales/zh-TW.yaml b/i18n/locales/zh-TW.yaml index 27759d07..84ebd57e 100644 --- a/i18n/locales/zh-TW.yaml +++ b/i18n/locales/zh-TW.yaml @@ -48,6 +48,9 @@ token.expired: "該令牌已過期" token.exhausted: "該令牌額度已用盡 TokenStatusExhausted[sk-{{.Prefix}}***{{.Suffix}}]" token.status_unavailable: "該令牌狀態不可用" token.db_error: "無效的令牌，資料庫查詢出錯，請聯繫管理員" +token.auto_groups_too_many: "每個令牌最多可選擇 {{.Max}} 個 Auto 分組" +token.auto_groups_duplicate: "Auto 分組 {{.Group}} 重複" +token.auto_groups_invalid: "Auto 分組 {{.Group}} 不可用或無權存取" # Redemption messages redemption.name_length: "兌換碼名稱長度必須在1-20之間" diff --git a/middleware/auth.go b/middleware/auth.go index 2ad09a7a..4e1436f3 100644 --- a/middleware/auth.go +++ b/middleware/auth.go @@ -502,6 +502,16 @@ func SetupContextForToken(c *gin.Context, token *model.Token, parts ...string) e } common.SetContextKey(c, constant.ContextKeyTokenGroup, token.Group) common.SetContextKey(c, constant.ContextKeyTokenCrossGroupRetry, token.CrossGroupRetry) + if token.AutoGroups != "" { + autoGroups, err := token.GetAutoGroups() + if err != nil { + common.SysError(fmt.Sprintf("failed to parse auto groups for token %d: %v", token.Id, err)) + autoGroups = []string{} + common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups) + } else if len(autoGroups) > 0 { + common.SetContextKey(c, constant.ContextKeyTokenAutoGroups, autoGroups) + } + } if len(parts) > 1 { if model.IsAdmin(token.UserId) { c.Set("specific_channel_id", parts[1]) diff --git a/middleware/distributor.go b/middleware/distributor.go index bde639dd..7decf0e2 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -109,7 +109,7 @@ func Distribute() func(c *gin.Context) { channelSupportsRequestPath(preferred, c.Request.URL.Path, modelRequest.Model) { if usingGroup == "auto" { userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup) - autoGroups := service.GetUserAutoGroup(userGroup) + autoGroups := service.GetRequestAutoGroups(c, userGroup) for _, g := range autoGroups { if model.IsChannelEnabledForGroupModel(g, modelRequest.Model, preferred.Id) { selectGroup = g diff --git a/middleware/token_auto_groups_context_test.go b/middleware/token_auto_groups_context_test.go new file mode 100644 index 00000000..ac507d5b --- /dev/null +++ b/middleware/token_auto_groups_context_test.go @@ -0,0 +1,48 @@ +package middleware + +import ( + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func newTokenAutoGroupsContext() *gin.Context { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + return ctx +} + +func TestSetupContextForTokenPreservesCustomAutoGroupsOrder(t *testing.T) { + ctx := newTokenAutoGroupsContext() + token := &model.Token{Id: 1, UserId: 2, AutoGroups: `["vip","default"]`} + + require.NoError(t, SetupContextForToken(ctx, token)) + value, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups) + require.True(t, ok) + assert.Equal(t, []string{"vip", "default"}, value) +} + +func TestSetupContextForTokenTreatsStoredEmptyArrayAsInheritance(t *testing.T) { + ctx := newTokenAutoGroupsContext() + token := &model.Token{Id: 1, UserId: 2, AutoGroups: `[]`} + + require.NoError(t, SetupContextForToken(ctx, token)) + _, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups) + assert.False(t, ok) +} + +func TestSetupContextForTokenMalformedAutoGroupsFailsClosed(t *testing.T) { + ctx := newTokenAutoGroupsContext() + token := &model.Token{Id: 1, UserId: 2, AutoGroups: `not-json`} + + require.NoError(t, SetupContextForToken(ctx, token)) + value, ok := common.GetContextKey(ctx, constant.ContextKeyTokenAutoGroups) + require.True(t, ok) + assert.Equal(t, []string{}, value) +} diff --git a/model/option.go b/model/option.go index 7ab64e0d..e7fda523 100644 --- a/model/option.go +++ b/model/option.go @@ -120,6 +120,7 @@ func InitOptionMap() { common.OptionMap["Chats"] = setting.Chats2JsonString() common.OptionMap["AutoGroups"] = setting.AutoGroups2JsonString() common.OptionMap["DefaultUseAutoGroup"] = strconv.FormatBool(setting.DefaultUseAutoGroup) + common.OptionMap["MaxTokenAutoGroups"] = strconv.Itoa(setting.GetMaxTokenAutoGroups()) common.OptionMap["PayMethods"] = operation_setting.PayMethods2JsonString() common.OptionMap["GitHubClientId"] = "" common.OptionMap["GitHubClientSecret"] = "" @@ -208,6 +209,9 @@ func validateOptionValue(key string, value string) error { if key == operation_setting.ToolPriceOptionKey { return operation_setting.ValidateToolPricesJSON(value) } + if key == "MaxTokenAutoGroups" { + return setting.ValidateMaxTokenAutoGroups(value) + } return nil } @@ -413,6 +417,8 @@ func updateOptionMap(key string, value string) (err error) { err = setting.UpdateChatsByJsonString(value) case "AutoGroups": err = setting.UpdateAutoGroupsByJsonString(value) + case "MaxTokenAutoGroups": + err = setting.UpdateMaxTokenAutoGroups(value) case "CustomCallbackAddress": operation_setting.CustomCallbackAddress = value case "EpayId": diff --git a/model/option_auto_group_test.go b/model/option_auto_group_test.go new file mode 100644 index 00000000..c1a8f168 --- /dev/null +++ b/model/option_auto_group_test.go @@ -0,0 +1,17 @@ +package model + +import ( + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestValidateOptionValueRejectsInvalidMaxTokenAutoGroups(t *testing.T) { + for _, value := range []string{"", "0", "-1", "1.5", "invalid"} { + t.Run(value, func(t *testing.T) { + assert.Error(t, validateOptionValue("MaxTokenAutoGroups", value)) + }) + } + require.NoError(t, validateOptionValue("MaxTokenAutoGroups", "999999")) +} diff --git a/model/token.go b/model/token.go index 5d62258e..5aa8b3d5 100644 --- a/model/token.go +++ b/model/token.go @@ -28,9 +28,34 @@ type Token struct { UsedQuota int `json:"used_quota" gorm:"default:0"` // used quota Group string `json:"group" gorm:"default:''"` CrossGroupRetry bool `json:"cross_group_retry"` // 跨分组重试，仅auto分组有效 + AutoGroups string `json:"-" gorm:"type:text"` DeletedAt gorm.DeletedAt `gorm:"index"` } +func (token *Token) GetAutoGroups() ([]string, error) { + if token.AutoGroups == "" { + return nil, nil + } + var groups []string + if err := common.UnmarshalJsonStr(token.AutoGroups, &groups); err != nil { + return nil, err + } + return groups, nil +} + +func (token *Token) SetAutoGroups(groups []string) error { + if len(groups) == 0 { + token.AutoGroups = "" + return nil + } + data, err := common.Marshal(groups) + if err != nil { + return err + } + token.AutoGroups = string(data) + return nil +} + func (token *Token) Clean() { token.Key = "" } @@ -291,18 +316,16 @@ func (token *Token) Insert() error { // Update Make sure your token's fields is completed, because this will update non-zero values func (token *Token) Update() (err error) { - defer func() { - if shouldUpdateRedis(true, err) { - gopool.Go(func() { - err := cacheSetToken(*token) - if err != nil { - common.SysLog("failed to update token cache: " + err.Error()) - } - }) - } - }() err = DB.Model(token).Select("name", "status", "expired_time", "remain_quota", "unlimited_quota", - "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry").Updates(token).Error + "model_limits_enabled", "model_limits", "allow_ips", "group", "cross_group_retry", "auto_groups").Updates(token).Error + if shouldUpdateRedis(true, err) { + if cacheErr := cacheSetToken(*token); cacheErr != nil { + common.SysLog("failed to update token cache: " + cacheErr.Error()) + if deleteErr := cacheDeleteToken(token.Key); deleteErr != nil { + common.SysLog("failed to invalidate token cache after update: " + deleteErr.Error()) + } + } + } return err } diff --git a/model/token_auto_groups_cache_test.go b/model/token_auto_groups_cache_test.go new file mode 100644 index 00000000..2018502d --- /dev/null +++ b/model/token_auto_groups_cache_test.go @@ -0,0 +1,57 @@ +package model + +import ( + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestTokenAutoGroupsRoundTripThroughRedisHashCache(t *testing.T) { + useUserCacheMiniRedis(t) + token := Token{ + Id: 42, + UserId: 7, + Key: "token-auto-groups-cache-key", + Name: "auto-cache", + Group: "auto", + AutoGroups: `["vip","default"]`, + } + + require.NoError(t, cacheSetToken(token)) + cached, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.Equal(t, token.AutoGroups, cached.AutoGroups) + groups, err := cached.GetAutoGroups() + require.NoError(t, err) + assert.Equal(t, []string{"vip", "default"}, groups) +} + +func TestTokenUpdateSynchronouslyNarrowsPreheatedAutoGroupsCache(t *testing.T) { + truncateTables(t) + useUserCacheMiniRedis(t) + token := Token{ + UserId: 7, + Key: "token-auto-groups-update-cache-key", + Name: "auto-cache-update", + Status: common.TokenStatusEnabled, + ExpiredTime: -1, + UnlimitedQuota: true, + Group: "auto", + CrossGroupRetry: true, + AutoGroups: `["default","vip"]`, + } + require.NoError(t, token.Insert()) + require.NoError(t, cacheSetToken(token)) + + preheated, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.JSONEq(t, `["default","vip"]`, preheated.AutoGroups) + + require.NoError(t, token.SetAutoGroups([]string{"vip"})) + require.NoError(t, token.Update()) + immediate, err := cacheGetTokenByKey(token.Key) + require.NoError(t, err) + assert.JSONEq(t, `["vip"]`, immediate.AutoGroups) +} diff --git a/router/api-router.go b/router/api-router.go index 80fd6517..907cf1ed 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -238,6 +238,7 @@ func SetApiRouter(router *gin.Engine) { { tokenRoute.GET("/", controller.GetAllTokens) tokenRoute.GET("/search", middleware.SearchRateLimit(), controller.SearchTokens) + tokenRoute.GET("/auto-groups", controller.GetTokenAutoGroups) tokenRoute.GET("/:id", controller.GetToken) tokenRoute.POST("/:id/key", middleware.CriticalRateLimit(), middleware.DisableCache(), controller.GetTokenKey) tokenRoute.POST("/", controller.AddToken) diff --git a/service/channel_select.go b/service/channel_select.go index 24c4e252..0ab88dc8 100644 --- a/service/channel_select.go +++ b/service/channel_select.go @@ -7,7 +7,6 @@ import ( "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/logger" "github.com/QuantumNous/new-api/model" - "github.com/QuantumNous/new-api/setting" "github.com/gin-gonic/gin" ) @@ -88,10 +87,10 @@ func CacheGetRandomSatisfiedChannel(param *RetryParam) (*model.Channel, string, userGroup := common.GetContextKeyString(param.Ctx, constant.ContextKeyUserGroup) if param.TokenGroup == "auto" { - if len(setting.GetAutoGroups()) == 0 { + autoGroups := GetRequestAutoGroups(param.Ctx, userGroup) + if len(autoGroups) == 0 { return nil, selectGroup, errors.New("auto groups is not enabled") } - autoGroups := GetUserAutoGroup(userGroup) // startGroupIndex: the group index to start searching from // startGroupIndex: 开始搜索的分组索引 diff --git a/service/channel_select_auto_groups_test.go b/service/channel_select_auto_groups_test.go new file mode 100644 index 00000000..e8454b38 --- /dev/null +++ b/service/channel_select_auto_groups_test.go @@ -0,0 +1,129 @@ +package service + +import ( + "fmt" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupChannelSelectAutoGroupsTest(t *testing.T) *gorm.DB { + t.Helper() + + originalDB := model.DB + originalMemoryCacheEnabled := common.MemoryCacheEnabled + originalRetryTimes := common.RetryTimes + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalGroupRatios := ratio_setting.GroupRatio2JSONString() + originalMaxTokenAutoGroups := setting.GetMaxTokenAutoGroups() + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&model.Channel{}, &model.Ability{})) + model.DB = db + common.MemoryCacheEnabled = true + common.RetryTimes = 0 + + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":2}`)) + require.NoError(t, setting.UpdateMaxTokenAutoGroups("2")) + + t.Cleanup(func() { + model.DB = originalDB + common.MemoryCacheEnabled = originalMemoryCacheEnabled + common.RetryTimes = originalRetryTimes + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalGroupRatios)) + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMaxTokenAutoGroups))) + + if originalMemoryCacheEnabled && originalDB != nil && + originalDB.Migrator().HasTable(&model.Channel{}) && originalDB.Migrator().HasTable(&model.Ability{}) { + model.InitChannelCache() + } + sqlDB, err := db.DB() + if err == nil { + require.NoError(t, sqlDB.Close()) + } + }) + + return db +} + +func createChannelSelectAutoGroupsChannel(t *testing.T, db *gorm.DB, id int, group, modelName string) { + t.Helper() + priority := int64(0) + weight := uint(100) + require.NoError(t, db.Create(&model.Channel{ + Id: id, + Type: constant.ChannelTypeOpenAI, + Key: fmt.Sprintf("key-%d", id), + Status: common.ChannelStatusEnabled, + Name: fmt.Sprintf("channel-%d", id), + Weight: &weight, + Models: modelName, + Group: group, + Priority: &priority, + }).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: group, + Model: modelName, + ChannelId: id, + Enabled: true, + Priority: &priority, + Weight: weight, + }).Error) +} + +func TestCacheGetRandomSatisfiedChannelUsesTokenAutoGroupsWhenGlobalAutoIsEmpty(t *testing.T) { + db := setupChannelSelectAutoGroupsTest(t) + const modelName = "auto-groups-runtime-model" + createChannelSelectAutoGroupsChannel(t, db, 2101, "vip", modelName) + createChannelSelectAutoGroupsChannel(t, db, 2102, "default", modelName) + model.InitChannelCache() + + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(ctx, constant.ContextKeyUserGroup, "default") + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip", "default"}) + common.SetContextKey(ctx, constant.ContextKeyTokenCrossGroupRetry, true) + + retry := 0 + param := &RetryParam{ + Ctx: ctx, + TokenGroup: "auto", + ModelName: modelName, + RequestPath: "/v1/chat/completions", + Retry: &retry, + } + + first, selectedGroup, err := CacheGetRandomSatisfiedChannel(param) + require.NoError(t, err) + require.NotNil(t, first) + assert.Equal(t, 2101, first.Id) + assert.Equal(t, "vip", selectedGroup) + assert.Equal(t, "vip", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup)) + assert.Empty(t, setting.GetAutoGroups(), "the selection must not depend on the global Auto list") + + param.IncreaseRetry() + second, selectedGroup, err := CacheGetRandomSatisfiedChannel(param) + require.NoError(t, err) + require.NotNil(t, second) + assert.Equal(t, 2102, second.Id) + assert.Equal(t, "default", selectedGroup) + assert.Equal(t, "default", common.GetContextKeyString(ctx, constant.ContextKeyAutoGroup)) +} diff --git a/service/group.go b/service/group.go index 8cb359bc..e792083e 100644 --- a/service/group.go +++ b/service/group.go @@ -3,9 +3,12 @@ package service import ( "strings" + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/model" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" ) func GetUserUsableGroups(userGroup string) map[string]string { @@ -42,18 +45,67 @@ func GroupInUserUsableGroups(userGroup, groupName string) bool { return ok } +func IsUserSelectableGroup(userGroup, groupName string) bool { + if groupName == "" || groupName == "auto" { + return false + } + return GroupInUserUsableGroups(userGroup, groupName) && ratio_setting.ContainsGroupRatio(groupName) +} + // GetUserAutoGroup 根据用户分组获取自动分组设置 func GetUserAutoGroup(userGroup string) []string { - groups := GetUserUsableGroups(userGroup) autoGroups := make([]string, 0) + seen := make(map[string]struct{}) for _, group := range setting.GetAutoGroups() { - if _, ok := groups[group]; ok { - autoGroups = append(autoGroups, group) + if !IsUserSelectableGroup(userGroup, group) { + continue } + if _, ok := seen[group]; ok { + continue + } + seen[group] = struct{}{} + autoGroups = append(autoGroups, group) } return autoGroups } +// FilterUserTokenAutoGroups applies current permissions before the current +// per-token limit. It intentionally does not fall back to the global Auto list. +func FilterUserTokenAutoGroups(userGroup string, groups []string) []string { + maxCount := setting.GetMaxTokenAutoGroups() + filtered := make([]string, 0, min(len(groups), maxCount)) + seen := make(map[string]struct{}) + for _, group := range groups { + if !IsUserSelectableGroup(userGroup, group) { + continue + } + if _, ok := seen[group]; ok { + continue + } + seen[group] = struct{}{} + filtered = append(filtered, group) + if len(filtered) == maxCount { + break + } + } + return filtered +} + +// GetRequestAutoGroups resolves the ordered Auto groups for the current token. +// The absence of the context value means that the token inherits the complete +// global Auto list; a present (even empty) value is an explicit token snapshot. +func GetRequestAutoGroups(c *gin.Context, userGroup string) []string { + value, ok := common.GetContextKey(c, constant.ContextKeyTokenAutoGroups) + if !ok { + return GetUserAutoGroup(userGroup) + } + groups, ok := value.([]string) + if !ok { + return []string{} + } + return FilterUserTokenAutoGroups(userGroup, groups) +} + // GetGroupsEnabledModels 按 groups 顺序获取各分组启用的模型并去重 func GetGroupsEnabledModels(groups []string) []string { seen := make(map[string]struct{}) diff --git a/service/group_auto_groups_test.go b/service/group_auto_groups_test.go new file mode 100644 index 00000000..1f138ad4 --- /dev/null +++ b/service/group_auto_groups_test.go @@ -0,0 +1,72 @@ +package service + +import ( + "fmt" + "net/http/httptest" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/setting" + "github.com/QuantumNous/new-api/setting/ratio_setting" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func configureRequestAutoGroupsTest(t *testing.T) { + t.Helper() + originalMax := setting.GetMaxTokenAutoGroups() + originalAutoGroups := setting.AutoGroups2JsonString() + originalUsableGroups := setting.UserUsableGroups2JSONString() + originalRatios := ratio_setting.GroupRatio2JSONString() + require.NoError(t, setting.UpdateMaxTokenAutoGroups("2")) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`["vip","default","svip"]`)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default","vip":"VIP","svip":"SVIP"}`)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(`{"default":1,"vip":1,"svip":1}`)) + t.Cleanup(func() { + require.NoError(t, setting.UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", originalMax))) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(originalAutoGroups)) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(originalUsableGroups)) + require.NoError(t, ratio_setting.UpdateGroupRatioByJSONString(originalRatios)) + }) +} + +func newRequestAutoGroupsContext() *gin.Context { + gin.SetMode(gin.TestMode) + ctx, _ := gin.CreateTestContext(httptest.NewRecorder()) + return ctx +} + +func TestGetRequestAutoGroupsInheritedListIsNotLimited(t *testing.T) { + configureRequestAutoGroupsTest(t) + ctx := newRequestAutoGroupsContext() + + groups := GetRequestAutoGroups(ctx, "default") + + assert.Equal(t, []string{"vip", "default", "svip"}, groups) +} + +func TestGetRequestAutoGroupsFiltersBeforeApplyingCurrentLimit(t *testing.T) { + configureRequestAutoGroupsTest(t) + ctx := newRequestAutoGroupsContext() + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"revoked", "vip", "default", "svip"}) + require.NoError(t, setting.UpdateAutoGroupsByJsonString(`[]`)) + + groups := GetRequestAutoGroups(ctx, "default") + + assert.Equal(t, []string{"vip", "default"}, groups) + require.NoError(t, setting.UpdateMaxTokenAutoGroups("1")) + assert.Equal(t, []string{"vip"}, GetRequestAutoGroups(ctx, "default")) +} + +func TestGetRequestAutoGroupsDoesNotFallBackAfterPermissionChange(t *testing.T) { + configureRequestAutoGroupsTest(t) + ctx := newRequestAutoGroupsContext() + common.SetContextKey(ctx, constant.ContextKeyTokenAutoGroups, []string{"vip"}) + require.NoError(t, setting.UpdateUserUsableGroupsByJSONString(`{"default":"Default"}`)) + + groups := GetRequestAutoGroups(ctx, "default") + + assert.Empty(t, groups) +} diff --git a/setting/auto_group.go b/setting/auto_group.go index 9261286b..3b509b60 100644 --- a/setting/auto_group.go +++ b/setting/auto_group.go @@ -1,15 +1,27 @@ package setting import ( + "fmt" + "strconv" + "sync/atomic" + "github.com/QuantumNous/new-api/common" ) +const DefaultMaxTokenAutoGroups = 5 + var autoGroups = []string{ "default", } var DefaultUseAutoGroup = false +var maxTokenAutoGroups atomic.Int64 + +func init() { + maxTokenAutoGroups.Store(DefaultMaxTokenAutoGroups) +} + func ContainsAutoGroup(group string) bool { for _, autoGroup := range autoGroups { if autoGroup == group { @@ -35,3 +47,24 @@ func AutoGroups2JsonString() string { func GetAutoGroups() []string { return autoGroups } + +func GetMaxTokenAutoGroups() int { + return int(maxTokenAutoGroups.Load()) +} + +func ValidateMaxTokenAutoGroups(value string) error { + maxCount, err := strconv.Atoi(value) + if err != nil || maxCount <= 0 { + return fmt.Errorf("MaxTokenAutoGroups must be a positive integer") + } + return nil +} + +func UpdateMaxTokenAutoGroups(value string) error { + if err := ValidateMaxTokenAutoGroups(value); err != nil { + return err + } + maxCount, _ := strconv.Atoi(value) + maxTokenAutoGroups.Store(int64(maxCount)) + return nil +} diff --git a/setting/auto_group_test.go b/setting/auto_group_test.go new file mode 100644 index 00000000..414c169e --- /dev/null +++ b/setting/auto_group_test.go @@ -0,0 +1,29 @@ +package setting + +import ( + "fmt" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestUpdateMaxTokenAutoGroupsAcceptsAnyPositiveInteger(t *testing.T) { + original := GetMaxTokenAutoGroups() + t.Cleanup(func() { + require.NoError(t, UpdateMaxTokenAutoGroups(fmt.Sprintf("%d", original))) + }) + + require.NoError(t, UpdateMaxTokenAutoGroups("123456")) + assert.Equal(t, 123456, GetMaxTokenAutoGroups()) +} + +func TestUpdateMaxTokenAutoGroupsRejectsInvalidValuesWithoutChangingState(t *testing.T) { + original := GetMaxTokenAutoGroups() + for _, value := range []string{"", "0", "-1", "1.5", "not-a-number"} { + t.Run(value, func(t *testing.T) { + assert.Error(t, UpdateMaxTokenAutoGroups(value)) + assert.Equal(t, original, GetMaxTokenAutoGroups()) + }) + } +} diff --git a/web/src/features/keys/api.ts b/web/src/features/keys/api.ts index df3cc5ff..0f90490c 100644 --- a/web/src/features/keys/api.ts +++ b/web/src/features/keys/api.ts @@ -25,6 +25,7 @@ import type { GetApiKeysResponse, SearchApiKeysParams, ApiKeyFormData, + TokenAutoGroupsConfig, } from './types' // ============================================================================ @@ -60,6 +61,14 @@ export async function getApiKey(id: number): Promise> { return res.data } +// Get the current user's global Auto order and the per-token selection limit. +export async function getTokenAutoGroups(): Promise< + ApiResponse +> { + const res = await api.get('/api/token/auto-groups') + return res.data +} + // Create a new API key export async function createApiKey( data: ApiKeyFormData diff --git a/web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx b/web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx new file mode 100644 index 00000000..5cb64ae5 --- /dev/null +++ b/web/src/features/keys/components/__tests__/api-key-group-cell.test.tsx @@ -0,0 +1,236 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { TooltipProvider } = await import('@/components/ui/tooltip') +const { ApiKeyGroupCell } = await import('../api-key-group-cell') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + Auto: 'Auto', + 'Cross-group': 'Cross-group', + Ratio: 'Ratio', + 'Automatically selects the best available group with circuit breaker mechanism': + 'Automatically selects the best available group with circuit breaker mechanism', + }, + }, + }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +function CellHarness(props: { + group: string + ratio?: number | string + crossGroupRetry?: boolean + shouldReduceMotion?: boolean +}) { + return ( + + + + + + ) +} + +describe('API key group table cell', () => { + after(() => { + domWindow.close() + }) + + test('renders two unclipped rings and a localized Auto ratio when API data uses a nonlocalized string', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render( + + ) + ) + + const badgeCell = container.querySelector( + '[data-api-key-group-cell="auto"]' + ) + assert.ok(badgeCell) + assert.equal(badgeCell.classList.contains('overflow-visible'), true) + assert.equal(badgeCell.classList.contains('overflow-hidden'), false) + + const frames = container.querySelectorAll('[data-auto-group-frame]') + const movingRings = container.querySelectorAll( + '[data-auto-group-flow-border]' + ) + assert.equal(frames.length, 2) + assert.equal(movingRings.length, 2) + for (const frame of frames) { + assert.equal(frame.classList.contains('relative'), true) + assert.equal(frame.classList.contains('overflow-visible'), true) + assert.equal(frame.classList.contains('rounded-4xl'), true) + assert.equal(frame.classList.contains('p-px'), true) + } + + const ratio = container.querySelector( + '[data-auto-group-effect="ratio"]' + ) + assert.ok(ratio) + assert.equal(ratio.textContent, 'Auto Ratio') + assert.equal(ratio.textContent?.includes('x'), false) + assert.equal(container.textContent?.includes('自动'), false) + assert.equal(container.textContent?.includes('Cross-group'), true) + + const crossGroupBadge = [ + ...container.querySelectorAll('[data-slot="status-badge"]'), + ].find((badge) => badge.textContent === 'Cross-group') + assert.ok(crossGroupBadge) + assert.equal(crossGroupBadge.closest('[data-auto-group-frame]'), null) + + await act(async () => root.unmount()) + container.remove() + }) + + test('keeps static Auto frames but omits both moving layers for reduced motion', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render() + ) + + assert.equal( + container.querySelectorAll('[data-auto-group-frame]').length, + 2 + ) + assert.equal( + container.querySelectorAll('[data-auto-group-flow-border]').length, + 0 + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('shows only the Auto badge when ratio data is unavailable', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render() + ) + + assert.equal( + container.querySelectorAll('[data-auto-group-frame]').length, + 1 + ) + assert.equal( + container.querySelectorAll('[data-auto-group-flow-border]').length, + 1 + ) + assert.equal( + container.querySelector('[data-auto-group-effect="ratio"]'), + null + ) + assert.equal(container.textContent?.includes('Auto'), true) + assert.equal(container.textContent?.includes('Ratio'), false) + + await act(async () => root.unmount()) + container.remove() + }) + + test('narrows normal group ratios to numbers and never applies Auto rings', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render( + + ) + ) + + assert.equal(container.textContent?.includes('vip'), true) + assert.equal(container.textContent?.includes('自动'), false) + assert.equal(container.querySelector('[data-auto-group-frame]'), null) + assert.equal(container.querySelector('[data-auto-group-flow-border]'), null) + + await act(async () => + root.render( + + ) + ) + + assert.equal(container.textContent?.includes('3x'), true) + assert.equal(container.querySelector('[data-auto-group-frame]'), null) + + await act(async () => root.unmount()) + container.remove() + }) +}) diff --git a/web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx b/web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx new file mode 100644 index 00000000..5c3b6525 --- /dev/null +++ b/web/src/features/keys/components/__tests__/api-key-group-combobox.test.tsx @@ -0,0 +1,294 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'HTMLInputElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +let shouldReduceMotion = false +const reducedMotionMediaQuery = domWindow.matchMedia('(prefers-reduced-motion)') +Object.defineProperty(reducedMotionMediaQuery, 'matches', { + configurable: true, + get: () => shouldReduceMotion, +}) +Object.defineProperty(domWindow, 'matchMedia', { + configurable: true, + value: () => reducedMotionMediaQuery, +}) + +function setReducedMotion(value: boolean) { + shouldReduceMotion = value + reducedMotionMediaQuery.dispatchEvent(new domWindow.Event('change')) +} + +const { act, useState } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { ApiKeyGroupCombobox } = await import('../api-key-group-combobox') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + Auto: 'Auto', + Ratio: 'Ratio', + 'Search...': 'Search...', + 'No group found.': 'No group found.', + 'Select a group': 'Select a group', + }, + }, + }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +const options = [ + { + value: 'auto', + label: 'auto', + desc: 'Global automatic routing', + ratio: '自动', + }, + { value: 'default', label: 'default', desc: 'User group', ratio: 1 }, + { value: 'vip', label: 'vip', desc: 'Priority group', ratio: 3 }, +] + +function Harness(props: { initialValue: string }) { + const [value, setValue] = useState(props.initialValue) + + return ( + + + {value} + + ) +} + +function getTrigger(container: ParentNode): HTMLButtonElement { + const trigger = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(trigger) + return trigger +} + +function getCommandItem(label: string): HTMLElement { + const item = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((candidate) => candidate.textContent?.includes(label)) + assert.ok(item) + return item +} + +describe('API key group combobox Auto effect', () => { + after(() => { + domWindow.close() + }) + + test('rings the selected Auto trigger and its localized ratio without rendering the API ratio text', async () => { + setReducedMotion(false) + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const trigger = getTrigger(container) + assert.equal(trigger.getAttribute('aria-expanded'), 'false') + assert.equal(trigger.dataset.autoGroupEffect, 'trigger') + assert.equal(trigger.classList.contains('bg-linear-to-r'), false) + assert.equal(trigger.classList.contains('overflow-hidden'), false) + assert.equal(trigger.classList.contains('overflow-visible'), true) + + const triggerFlowBorder = trigger.querySelector( + '[data-auto-group-flow-border]' + ) + assert.ok(triggerFlowBorder) + assert.equal(triggerFlowBorder.getAttribute('aria-hidden'), 'true') + assert.equal( + triggerFlowBorder.classList.contains('pointer-events-none'), + true + ) + assert.equal( + triggerFlowBorder.classList.contains('auto-group-flow-border'), + true + ) + + const triggerRatio = trigger.querySelector( + '[data-auto-group-effect="ratio"]' + ) + assert.ok(triggerRatio) + assert.equal(triggerRatio.textContent, 'Auto Ratio') + assert.equal(triggerRatio.textContent?.includes('Auto'), true) + assert.equal(triggerRatio.textContent?.includes('x'), false) + assert.equal(trigger.textContent?.includes('自动'), false) + assert.equal(triggerRatio.classList.contains('relative'), true) + assert.equal(triggerRatio.classList.contains('overflow-visible'), true) + assert.equal(triggerRatio.classList.contains('rounded-4xl'), true) + assert.ok(triggerRatio.querySelector('[data-auto-group-flow-border]')) + + await act(async () => trigger.click()) + assert.equal(trigger.getAttribute('aria-expanded'), 'true') + + const autoOption = getCommandItem('Global automatic routing') + assert.equal(autoOption.dataset.autoGroupEffect, 'option') + assert.equal(autoOption.getAttribute('aria-selected'), 'true') + assert.equal(autoOption.classList.contains('bg-linear-to-r'), false) + assert.equal(autoOption.classList.contains('overflow-visible'), true) + assert.ok(autoOption.querySelector('[data-auto-group-flow-border]')) + const optionRatio = autoOption.querySelector( + '[data-auto-group-effect="ratio"]' + ) + assert.ok(optionRatio) + assert.equal(optionRatio.textContent, 'Auto Ratio') + assert.ok(optionRatio.querySelector('[data-auto-group-flow-border]')) + + const defaultOption = getCommandItem('User group') + assert.equal(defaultOption.hasAttribute('data-auto-group-effect'), false) + assert.equal( + defaultOption.querySelector('[data-auto-group-flow-border]'), + null + ) + assert.equal(defaultOption.textContent?.includes('1x Ratio'), true) + assert.equal( + defaultOption.querySelector('[data-auto-group-effect="ratio"]'), + null + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('keeps search and selection behavior while leaving normal groups unstyled', async () => { + setReducedMotion(false) + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const trigger = getTrigger(container) + await act(async () => trigger.click()) + + const searchInput = document.querySelector( + 'input[placeholder="Search..."]' + ) + assert.ok(searchInput) + await act(async () => { + const valueSetter = Object.getOwnPropertyDescriptor( + domWindow.HTMLInputElement.prototype, + 'value' + )?.set + assert.ok(valueSetter) + valueSetter.call(searchInput, 'vip') + searchInput.dispatchEvent( + new domWindow.Event('input', { bubbles: true }) as unknown as Event + ) + }) + + const visibleOptions = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ] + assert.equal( + visibleOptions.some((option) => + option.textContent?.includes('Global automatic routing') + ), + false + ) + const vipOption = getCommandItem('Priority group') + await act(async () => vipOption.click()) + + assert.equal( + container.querySelector('[data-testid="selected-group"]')?.textContent, + 'vip' + ) + assert.equal(trigger.getAttribute('aria-expanded'), 'false') + assert.equal(trigger.hasAttribute('data-auto-group-effect'), false) + assert.equal(trigger.querySelector('[data-auto-group-flow-border]'), null) + + await act(async () => root.unmount()) + container.remove() + }) + + test('preserves the static Auto treatment but omits moving layers for reduced motion', async () => { + setReducedMotion(true) + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const trigger = getTrigger(container) + assert.equal(trigger.dataset.autoGroupEffect, 'trigger') + assert.equal(trigger.querySelector('[data-auto-group-flow-border]'), null) + assert.ok(trigger.querySelector('[data-auto-group-effect="ratio"]')) + + await act(async () => trigger.click()) + const autoOption = getCommandItem('Global automatic routing') + assert.equal(autoOption.dataset.autoGroupEffect, 'option') + assert.equal( + autoOption.querySelector('[data-auto-group-flow-border]'), + null + ) + assert.ok(autoOption.querySelector('[data-auto-group-effect="ratio"]')) + + await act(async () => root.unmount()) + container.remove() + setReducedMotion(false) + }) +}) diff --git a/web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx b/web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx new file mode 100644 index 00000000..238d0b21 --- /dev/null +++ b/web/src/features/keys/components/__tests__/api-keys-mutate-drawer.test.tsx @@ -0,0 +1,371 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, afterEach, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'HTMLInputElement', + 'HTMLFormElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'MouseEvent', + 'FocusEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { QueryClient, QueryClientProvider } = + await import('@tanstack/react-query') +const { api } = await import('@/lib/api') +const { ApiKeysProvider } = await import('../api-keys-provider') +const { ApiKeysMutateDrawer } = await import('../api-keys-mutate-drawer') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { en: { translation: {} } }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +type ApiMethod = (url: string, data?: unknown) => Promise<{ data: unknown }> +type MockableApi = { + get: ApiMethod + post: ApiMethod +} +type RenderedDrawer = { + host: HTMLDivElement + queryClient: InstanceType + root: ReturnType +} + +const apiClient = api as unknown as MockableApi +const originalGet = apiClient.get +const originalPost = apiClient.post +let renderedDrawer: RenderedDrawer | null = null + +function installApiFixtures(createdPayloads: Array>) { + apiClient.get = async (url) => { + switch (url) { + case '/api/status': + return { data: { data: { default_use_auto_group: true } } } + case '/api/user/models': + return { data: { success: true, data: [] } } + case '/api/user/self/groups': + return { + data: { + success: true, + data: { + auto: { desc: 'Automatic routing', ratio: 'auto' }, + default: { desc: 'Standard access', ratio: 1 }, + vip: { desc: 'Priority access', ratio: 2 }, + }, + }, + } + case '/api/token/auto-groups': + return { + data: { + success: true, + data: { groups: ['vip', 'default'], max_count: 3 }, + }, + } + default: + throw new Error(`Unexpected GET ${url}`) + } + } + apiClient.post = async (url, data) => { + assert.equal(url, '/api/token/') + assert.ok(data && typeof data === 'object') + createdPayloads.push(data as Record) + return { data: { success: true, data: {} } } + } +} + +async function waitForCondition( + condition: () => boolean, + failureMessage: string +): Promise { + if (condition()) return + + await new Promise((resolve, reject) => { + const observer = new MutationObserver(() => { + if (!condition()) return + clearTimeout(timeoutId) + observer.disconnect() + resolve() + }) + const timeoutId = setTimeout(() => { + observer.disconnect() + reject(new Error(`${failureMessage}: ${document.body.textContent}`)) + }, 1500) + + observer.observe(document, { + attributes: true, + childList: true, + characterData: true, + subtree: true, + }) + }) +} + +async function renderCreateDrawer(): Promise { + const host = document.createElement('div') + document.body.append(host) + const root = createRoot(host) + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }) + const freshAt = Date.now() + 60_000 + queryClient.setQueryData( + ['status'], + { default_use_auto_group: true }, + { updatedAt: freshAt } + ) + queryClient.setQueryData( + ['user-models'], + { success: true, data: [] }, + { updatedAt: freshAt } + ) + queryClient.setQueryData( + ['user-groups'], + { + success: true, + data: { + auto: { desc: 'Automatic routing', ratio: 'auto' }, + default: { desc: 'Standard access', ratio: 1 }, + vip: { desc: 'Priority access', ratio: 2 }, + }, + }, + { updatedAt: freshAt } + ) + queryClient.setQueryData( + ['token-auto-groups'], + { + success: true, + data: { groups: ['vip', 'default'], max_count: 3 }, + }, + { updatedAt: freshAt } + ) + renderedDrawer = { host, queryClient, root } + + await act(async () => + root.render( + + + + undefined} /> + + + + ) + ) + await act(async () => + waitForCondition(() => { + const saveButton = findButton('Save changes', false) + return saveButton !== null && !saveButton.disabled + }, 'API key drawer did not finish initializing') + ) +} + +function findButton(text: string, required: true): HTMLButtonElement +function findButton(text: string, required: false): HTMLButtonElement | null +function findButton(text: string, required = true): HTMLButtonElement | null { + const button = [ + ...document.querySelectorAll('button'), + ].find((candidate) => candidate.textContent?.includes(text)) + if (required) assert.ok(button, `Expected button containing "${text}"`) + return button ?? null +} + +function getControlByLabel(labelText: string): T { + const label = [...document.querySelectorAll('label')].find( + (candidate) => candidate.textContent?.trim() === labelText + ) + assert.ok(label, `Expected label "${labelText}"`) + assert.ok(label.htmlFor) + const control = + label.control ?? + label + .closest('[data-slot="form-item"]') + ?.querySelector( + '[data-slot="form-control"], input, textarea, button[role="combobox"], [role="group"]' + ) + assert.ok(control) + return control as T +} + +async function changeInput(input: HTMLInputElement, value: string) { + await act(async () => { + const valueSetter = Object.getOwnPropertyDescriptor( + domWindow.HTMLInputElement.prototype, + 'value' + )?.set + assert.ok(valueSetter) + valueSetter.call(input, value) + input.dispatchEvent( + new domWindow.Event('input', { bubbles: true }) as unknown as Event + ) + }) +} + +async function selectComboboxOption( + trigger: HTMLButtonElement, + optionDescription: string +) { + await act(async () => trigger.click()) + const option = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((candidate) => candidate.textContent?.includes(optionDescription)) + assert.ok(option, `Expected option containing "${optionDescription}"`) + await act(async () => option.click()) +} + +afterEach(async () => { + apiClient.get = originalGet + apiClient.post = originalPost + domWindow.localStorage.clear() + if (renderedDrawer) { + await act(async () => renderedDrawer?.root.unmount()) + renderedDrawer.queryClient.clear() + renderedDrawer.host.remove() + renderedDrawer = null + } + document.body.replaceChildren() +}) + +after(() => { + domWindow.close() +}) + +describe('API keys mutate drawer Auto group integration', () => { + test('inherits the root Auto order and sends an empty override for every batch-created key', async () => { + const createdPayloads: Array> = [] + installApiFixtures(createdPayloads) + await renderCreateDrawer() + + const groupTrigger = getControlByLabel('Group') + assert.equal(groupTrigger.textContent?.includes('auto'), true) + assert.equal( + document.body.textContent?.includes( + 'Using the complete global Auto order (2 groups)' + ), + true + ) + assert.deepEqual( + [ + ...document.querySelectorAll('[data-slot="global-auto-order-name"]'), + ].map((item) => item.textContent), + ['vip', 'default'] + ) + assert.equal(findButton('Restore global Auto', true).disabled, true) + + await changeInput(getControlByLabel('Name'), 'batch') + await changeInput(getControlByLabel('Quantity'), '2') + await act(async () => findButton('Save changes', true).click()) + await act(async () => + waitForCondition( + () => createdPayloads.length === 2, + 'batch API keys were not created' + ) + ) + + assert.equal(createdPayloads.length, 2) + assert.equal(createdPayloads[0]?.name, 'batch') + for (const payload of createdPayloads) { + assert.equal(payload.group, 'auto') + assert.deepEqual(payload.auto_groups, []) + assert.equal(payload.cross_group_retry, true) + } + }) + + test('preserves an unsaved custom order and mode after Auto to ordinary to Auto changes', async () => { + const createdPayloads: Array> = [] + installApiFixtures(createdPayloads) + await renderCreateDrawer() + + const autoOrderControl = getControlByLabel('Auto group order') + const addGroupTrigger = autoOrderControl.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addGroupTrigger) + await selectComboboxOption(addGroupTrigger, 'Priority access') + + assert.ok(document.querySelector('button[aria-label="Remove vip"]')) + assert.equal( + document.body.textContent?.includes('1 / 3 groups selected'), + true + ) + assert.equal(findButton('Restore global Auto', true).disabled, false) + + const groupTrigger = getControlByLabel('Group') + await selectComboboxOption(groupTrigger, 'Standard access') + assert.equal( + document.querySelector('button[aria-label="Remove vip"]'), + null + ) + await selectComboboxOption(groupTrigger, 'Automatic routing') + + assert.ok(document.querySelector('button[aria-label="Remove vip"]')) + assert.equal( + document.body.textContent?.includes('1 / 3 groups selected'), + true + ) + assert.equal(findButton('Restore global Auto', true).disabled, false) + + await changeInput(getControlByLabel('Name'), 'custom') + await act(async () => findButton('Save changes', true).click()) + await act(async () => + waitForCondition( + () => createdPayloads.length === 1, + 'custom-order API key was not created' + ) + ) + assert.deepEqual(createdPayloads[0]?.auto_groups, ['vip']) + }) +}) diff --git a/web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx b/web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx new file mode 100644 index 00000000..f37f5075 --- /dev/null +++ b/web/src/features/keys/components/__tests__/auto-group-order-editor.test.tsx @@ -0,0 +1,540 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import assert from 'node:assert/strict' +import { after, describe, test } from 'node:test' + +import { Window } from 'happy-dom' + +const domWindow = new Window() +const domGlobals = [ + 'window', + 'document', + 'navigator', + 'HTMLElement', + 'HTMLButtonElement', + 'HTMLInputElement', + 'SVGElement', + 'Node', + 'Element', + 'Event', + 'KeyboardEvent', + 'PointerEvent', + 'CustomEvent', + 'MutationObserver', + 'ResizeObserver', + 'requestAnimationFrame', + 'cancelAnimationFrame', + 'getComputedStyle', +] as const + +for (const key of domGlobals) { + Object.defineProperty(globalThis, key, { + configurable: true, + value: domWindow[key], + }) +} + +const { act, useState } = await import('react') +const { createRoot } = await import('react-dom/client') +const { createInstance } = await import('i18next') +const { I18nextProvider, initReactI18next } = await import('react-i18next') +const { AutoGroupOrderEditor } = await import('../auto-group-order-editor') + +const i18n = createInstance() +await i18n.use(initReactI18next).init({ + lng: 'en', + resources: { + en: { + translation: { + '{{count}} / {{max}} groups selected': + '{{count}} / {{max}} groups selected', + 'Add Auto group': 'Add Auto group', + 'Auto group order': 'Auto group order', + 'Drag {{group}} to reorder': 'Drag {{group}} to reorder', + 'Inherit global Auto order': 'Inherit global Auto order', + 'Maximum {{max}} groups selected': 'Maximum {{max}} groups selected', + 'Move {{group}} down': 'Move {{group}} down', + 'Move {{group}} up': 'Move {{group}} up', + 'No available groups in the global Auto order.': + 'No available groups in the global Auto order.', + 'No valid custom Auto groups remain. Add a group or restore global Auto.': + 'No valid custom Auto groups remain. Add a group or restore global Auto.', + 'No custom groups. Saving will inherit the complete global Auto order.': + 'No custom groups. Saving will inherit the complete global Auto order.', + 'Remove {{group}}': 'Remove {{group}}', + 'Restore global Auto': 'Restore global Auto', + Ratio: 'Ratio', + 'Search...': 'Search...', + 'No group found.': 'No group found.', + 'Select a group': 'Select a group', + 'Using the complete global Auto order ({{count}} groups)': + 'Using the complete global Auto order ({{count}} groups)', + }, + }, + }, +}) + +const reactTestGlobals = globalThis as typeof globalThis & { + IS_REACT_ACT_ENVIRONMENT?: boolean +} +reactTestGlobals.IS_REACT_ACT_ENVIRONMENT = true + +const globalOptions = [ + { value: 'vip', label: 'VIP', desc: 'Priority access', ratio: 3 }, + { value: 'default', label: 'Default', desc: 'Standard access', ratio: 1 }, + { value: 'team', label: 'Team', desc: 'Shared access', ratio: 2 }, +] + +function Harness(props: { initialGroups?: string[] }) { + const [groups, setGroups] = useState( + props.initialGroups ?? ['default', 'vip'] + ) + const [mode, setMode] = useState<'inherit' | 'custom'>('custom') + return ( + + { + setGroups(value.groups) + setMode(value.mode) + }} + /> + {groups.join(',')} + {mode} + + ) +} + +function InheritanceHarness(props: { globalOptions?: typeof globalOptions }) { + const [groups, setGroups] = useState([]) + const [mode, setMode] = useState<'inherit' | 'custom'>('inherit') + + return ( + + { + setGroups(value.groups) + setMode(value.mode) + }} + /> + {groups.join(',')} + {mode} + + ) +} + +function CustomEmptyHarness() { + const [groups, setGroups] = useState([]) + const [mode, setMode] = useState<'inherit' | 'custom'>('custom') + + return ( + + { + setGroups(value.groups) + setMode(value.mode) + }} + /> + {groups.join(',')} + {mode} + + ) +} + +function findButton(container: ParentNode, label: string): HTMLButtonElement { + const button = container.querySelector( + `button[aria-label="${label}"]` + ) + assert.ok(button) + return button +} + +describe('Auto group order editor', () => { + after(() => { + domWindow.close() + }) + + test('enforces the limit and exposes accessible reorder controls', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const addButton = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addButton) + assert.equal(addButton.disabled, true) + assert.equal(container.textContent?.includes('2 / 2 groups selected'), true) + assert.ok( + container.querySelector('[role="group"][aria-label="Auto group order"]') + ) + assert.equal( + findButton(container, 'Drag default to reorder').type, + 'button' + ) + + await act(async () => findButton(container, 'Move default down').click()) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'vip,default' + ) + + await act(async () => { + findButton(container, 'Drag vip to reorder').dispatchEvent( + new domWindow.KeyboardEvent('keydown', { + key: 'ArrowDown', + bubbles: true, + }) as unknown as KeyboardEvent + ) + }) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'default,vip' + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('adds and removes groups, then restores inheritance as an empty value', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + await act(async () => findButton(container, 'Remove vip').click()) + + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'default' + ) + const addButton = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addButton) + assert.equal(addButton.disabled, false) + + await act(async () => addButton.click()) + const teamOption = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((option) => option.textContent?.includes('team')) + assert.ok(teamOption) + await act(async () => teamOption.click()) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'default,team' + ) + assert.equal(addButton.disabled, true) + + const restoreButton = [...container.querySelectorAll('button')].find( + (button) => button.textContent?.includes('Restore global Auto') + ) + assert.ok(restoreButton) + await act(async () => restoreButton.click()) + + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + '' + ) + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'inherit' + ) + assert.equal( + container.textContent?.includes( + 'Using the complete global Auto order (3 groups)' + ), + true + ) + + const inheritedItems = container.querySelectorAll( + '[data-slot="global-auto-order"] > li' + ) + assert.deepEqual( + [...inheritedItems].map( + (item) => + item.querySelector('[data-slot="global-auto-order-name"]') + ?.textContent + ), + ['VIP', 'Default', 'Team'] + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('shows the complete inherited order with metadata beyond the custom limit', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + assert.equal( + container.textContent?.includes( + 'Using the complete global Auto order (3 groups)' + ), + true + ) + assert.equal( + container.textContent?.includes('0 / 2 groups selected'), + false + ) + + const order = container.querySelector( + '[data-slot="global-auto-order"]' + ) + assert.ok(order) + assert.equal(order.classList.contains('overflow-y-auto'), true) + assert.equal(order.classList.contains('flex-wrap'), true) + + const items = [...order.querySelectorAll('li')] + assert.equal(items.length, 3) + assert.equal( + order.querySelectorAll('[data-slot="global-auto-order-connector"]') + .length, + 2 + ) + assert.deepEqual( + items.map((item) => ({ + index: item.querySelector('[data-slot="global-auto-order-index"]') + ?.textContent, + name: item.querySelector('[data-slot="global-auto-order-name"]') + ?.textContent, + title: item + .querySelector('[data-slot="global-auto-order-chip"]') + ?.getAttribute('title'), + description: item.querySelector( + '[data-slot="global-auto-order-description"]' + )?.textContent, + ratio: item.querySelector('[data-slot="badge"]')?.textContent, + })), + [ + { + index: '1', + name: 'VIP', + title: 'Priority access', + description: 'Priority access', + ratio: '3x Ratio', + }, + { + index: '2', + name: 'Default', + title: 'Standard access', + description: 'Standard access', + ratio: '1x Ratio', + }, + { + index: '3', + name: 'Team', + title: 'Shared access', + description: 'Shared access', + ratio: '2x Ratio', + }, + ] + ) + + for (const item of items) { + const chip = item.querySelector('[data-slot="global-auto-order-chip"]') + assert.ok(chip) + const description = item.querySelector( + '[data-slot="global-auto-order-description"]' + ) + assert.ok(description) + assert.equal(description.classList.contains('sr-only'), true) + } + + assert.equal( + items[0]?.querySelector('[data-slot="global-auto-order-connector"]'), + null + ) + for (const item of items.slice(1)) { + const connector = item.querySelector( + '[data-slot="global-auto-order-connector"]' + ) + assert.ok(connector) + assert.equal(connector.getAttribute('aria-hidden'), 'true') + } + + assert.equal(container.querySelector('[aria-label^="Drag "]'), null) + assert.equal(container.querySelector('[aria-label^="Move "]'), null) + assert.equal(container.querySelector('[aria-label^="Remove "]'), null) + + const restoreButton = [...container.querySelectorAll('button')].find( + (button) => button.textContent?.includes('Restore global Auto') + ) + assert.ok(restoreButton) + assert.equal(restoreButton.disabled, true) + + await act(async () => root.unmount()) + container.remove() + }) + + test('shows an explicit empty state when the global Auto order has no groups', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => + root.render() + ) + + assert.equal( + container.textContent?.includes( + 'Using the complete global Auto order (0 groups)' + ), + true + ) + assert.equal( + container.textContent?.includes( + 'No available groups in the global Auto order.' + ), + true + ) + assert.equal( + container.querySelector('[data-slot="global-auto-order"]'), + null + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('keeps an empty custom order distinct from global inheritance', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'custom' + ) + assert.equal( + container.textContent?.includes( + 'No valid custom Auto groups remain. Add a group or restore global Auto.' + ), + true + ) + assert.equal( + container.querySelector('[data-slot="global-auto-order"]'), + null + ) + + const restoreButton = [...container.querySelectorAll('button')].find( + (button) => button.textContent?.includes('Restore global Auto') + ) + assert.ok(restoreButton) + assert.equal(restoreButton.disabled, false) + await act(async () => restoreButton.click()) + + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'inherit' + ) + assert.ok(container.querySelector('[data-slot="global-auto-order"]')) + + await act(async () => root.unmount()) + container.remove() + }) + + test('adding a group from inheritance explicitly creates a custom order', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + + const addButton = container.querySelector( + 'button[role="combobox"]' + ) + assert.ok(addButton) + await act(async () => addButton.click()) + const vipOption = [ + ...document.querySelectorAll('[data-slot="command-item"]'), + ].find((option) => option.textContent?.includes('VIP')) + assert.ok(vipOption) + await act(async () => vipOption.click()) + + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'custom' + ) + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + 'vip' + ) + assert.equal( + container.querySelector('[data-slot="global-auto-order"]'), + null + ) + + await act(async () => root.unmount()) + container.remove() + }) + + test('removing the last custom group does not silently enable inheritance', async () => { + const container = document.createElement('div') + document.body.append(container) + const root = createRoot(container) + + await act(async () => root.render()) + await act(async () => findButton(container, 'Remove default').click()) + + assert.equal( + container.querySelector('[data-testid="order"]')?.textContent, + '' + ) + assert.equal( + container.querySelector('[data-testid="mode"]')?.textContent, + 'custom' + ) + assert.equal( + container.textContent?.includes( + 'No valid custom Auto groups remain. Add a group or restore global Auto.' + ), + true + ) + + await act(async () => root.unmount()) + container.remove() + }) +}) diff --git a/web/src/features/keys/components/api-key-group-cell.tsx b/web/src/features/keys/components/api-key-group-cell.tsx new file mode 100644 index 00000000..21a1bc7d --- /dev/null +++ b/web/src/features/keys/components/api-key-group-cell.tsx @@ -0,0 +1,90 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { useTranslation } from 'react-i18next' + +import { BadgeCell, TruncatedCell } from '@/components/data-table' +import { GroupBadge } from '@/components/group-badge' +import { StatusBadge } from '@/components/status-badge' +import { + Tooltip, + TooltipContent, + TooltipTrigger, +} from '@/components/ui/tooltip' + +import { + // AutoGroupBadge, + GroupRatioBadge, + type GroupRatio, +} from './auto-group-visuals' + +type ApiKeyGroupCellProps = { + crossGroupRetry: boolean + group: string + ratio?: GroupRatio + shouldReduceMotion: boolean +} + +export function ApiKeyGroupCell(props: ApiKeyGroupCellProps) { + const { t } = useTranslation() + + if (props.group !== 'auto') { + const ratio = typeof props.ratio === 'number' ? props.ratio : undefined + return ( + + + + ) + } + + return ( + + + } + > + + {/**/} + + + + + {t( + 'Automatically selects the best available group with circuit breaker mechanism' + )} + + + + ) +} diff --git a/web/src/features/keys/components/api-key-group-combobox.tsx b/web/src/features/keys/components/api-key-group-combobox.tsx index 2593eff0..a5dada76 100644 --- a/web/src/features/keys/components/api-key-group-combobox.tsx +++ b/web/src/features/keys/components/api-key-group-combobox.tsx @@ -20,7 +20,6 @@ import { Check, ChevronsUpDown } from 'lucide-react' import { useMemo, useState } from 'react' import { useTranslation } from 'react-i18next' -import { Badge } from '@/components/ui/badge' import { Button } from '@/components/ui/button' import { Command, @@ -35,8 +34,15 @@ import { PopoverContent, PopoverTrigger, } from '@/components/ui/popover' +import { useMediaQuery } from '@/hooks' import { cn } from '@/lib/utils' +import { + AUTO_GROUP_FRAME_CLASS_NAME, + AutoGroupFlowBorder, + GroupRatioBadge, +} from './auto-group-visuals' + export type ApiKeyGroupOption = { value: string label: string @@ -52,50 +58,6 @@ type ApiKeyGroupComboboxProps = { disabled?: boolean } -function formatGroupRatio( - ratio: ApiKeyGroupOption['ratio'], - ratioLabel: string -) { - if (ratio === undefined || ratio === null || ratio === '') return null - return `${ratio}x ${ratioLabel}` -} - -function getRatioBadgeClassName(ratio: ApiKeyGroupOption['ratio']) { - if (typeof ratio !== 'number') { - return 'border-emerald-200 bg-emerald-50 text-emerald-700 dark:border-emerald-900/60 dark:bg-emerald-950/40 dark:text-emerald-300' - } - - if (ratio > 5) { - return 'border-rose-200 bg-rose-50 text-rose-700 dark:border-rose-900/60 dark:bg-rose-950/40 dark:text-rose-300' - } - if (ratio > 3) { - return 'border-orange-200 bg-orange-50 text-orange-700 dark:border-orange-900/60 dark:bg-orange-950/40 dark:text-orange-300' - } - if (ratio > 1) { - return 'border-blue-200 bg-blue-50 text-blue-700 dark:border-blue-900/60 dark:bg-blue-950/40 dark:text-blue-300' - } - return 'border-emerald-200 bg-emerald-50 text-emerald-700 dark:border-emerald-900/60 dark:bg-emerald-950/40 dark:text-emerald-300' -} - -function GroupRatioBadge({ ratio }: { ratio: ApiKeyGroupOption['ratio'] }) { - const { t } = useTranslation() - const label = formatGroupRatio(ratio, t('Ratio')) - - if (!label) return null - - return ( - - {label} - - ) -} - export function ApiKeyGroupCombobox({ options, value, @@ -106,7 +68,9 @@ export function ApiKeyGroupCombobox({ const { t } = useTranslation() const [open, setOpen] = useState(false) const [searchValue, setSearchValue] = useState('') + const shouldReduceMotion = useMediaQuery('(prefers-reduced-motion: reduce)') const selectedOption = options.find((option) => option.value === value) + const isAutoSelected = selectedOption?.value === 'auto' const filteredOptions = useMemo(() => { const search = searchValue.trim().toLowerCase() @@ -138,11 +102,22 @@ export function ApiKeyGroupCombobox({ variant='outline' role='combobox' aria-expanded={open} + data-auto-group-effect={isAutoSelected ? 'trigger' : undefined} disabled={disabled} - className='border-input bg-muted/40 hover:bg-muted/55 hover:text-foreground active:bg-background data-popup-open:border-ring data-popup-open:bg-background data-popup-open:ring-ring/20 h-auto min-h-14 w-full justify-between gap-2 rounded-lg px-3 py-2 text-start shadow-none transition-[background-color,border-color,box-shadow] duration-150 data-popup-open:ring-[3px] sm:min-h-20 sm:gap-3 sm:px-4 sm:py-3' + className={cn( + 'border-input bg-muted/40 hover:bg-muted/55 hover:text-foreground active:bg-background data-popup-open:border-ring data-popup-open:bg-background data-popup-open:ring-ring/20 relative h-auto min-h-14 w-full justify-between gap-2 rounded-lg px-3 py-2 text-start shadow-none transition-[background-color,border-color,box-shadow] duration-150 data-popup-open:ring-[3px] sm:min-h-20 sm:gap-3 sm:px-4 sm:py-3', + isAutoSelected && + cn( + AUTO_GROUP_FRAME_CLASS_NAME, + 'hover:border-primary/55 data-popup-open:border-primary/55 data-popup-open:ring-primary/20' + ) + )} /> } > + {isAutoSelected && ( + + )} @@ -155,10 +130,17 @@ export function ApiKeyGroupCombobox({ )} - + - + + + + ) + })} diff --git a/web/src/features/keys/components/api-keys-columns.tsx b/web/src/features/keys/components/api-keys-columns.tsx index 645d0513..2880783d 100644 --- a/web/src/features/keys/components/api-keys-columns.tsx +++ b/web/src/features/keys/components/api-keys-columns.tsx @@ -20,8 +20,6 @@ import { useQuery } from '@tanstack/react-query' import type { ColumnDef } from '@tanstack/react-table' import { useTranslation } from 'react-i18next' -import { BadgeCell, TruncatedCell } from '@/components/data-table' -import { GroupBadge } from '@/components/group-badge' import { StatusBadge } from '@/components/status-badge' import { Checkbox } from '@/components/ui/checkbox' import { Progress } from '@/components/ui/progress' @@ -30,6 +28,7 @@ import { TooltipContent, TooltipTrigger, } from '@/components/ui/tooltip' +import { useMediaQuery } from '@/hooks' import { toIntlLocale } from '@/i18n/languages' import { getUserGroups } from '@/lib/api' import dayjs from '@/lib/dayjs' @@ -38,6 +37,7 @@ import { cn } from '@/lib/utils' import { API_KEY_STATUSES } from '../constants' import type { ApiKey } from '../types' +import { ApiKeyGroupCell } from './api-key-group-cell' import { ApiKeyTimestampCell } from './api-key-timestamp-cell' import { ApiKeyCell, @@ -53,16 +53,16 @@ function getQuotaProgressColor(percentage: number): string { return '[&_[data-slot=progress-indicator]]:bg-emerald-500' } -function useGroupRatios(): Record { +function useGroupRatios(): Record { const { data } = useQuery({ queryKey: ['user-groups'], queryFn: getUserGroups, staleTime: 0, select: (res) => { if (!res.success || !res.data) return {} - const ratios: Record = {} + const ratios: Record = {} for (const [group, info] of Object.entries(res.data)) { - if (typeof info.ratio === 'number') { + if (typeof info.ratio === 'number' || typeof info.ratio === 'string') { ratios[group] = info.ratio } } @@ -76,6 +76,7 @@ function useGroupRatios(): Record { export function useApiKeysColumns(now: number): ColumnDef[] { const { t, i18n } = useTranslation() const groupRatios = useGroupRatios() + const shouldReduceMotion = useMediaQuery('(prefers-reduced-motion: reduce)') const locale = toIntlLocale(i18n.resolvedLanguage || i18n.language) const justNowLabel = t('Just now') const staleAccessThreshold = dayjs(now).subtract(3, 'month').valueOf() @@ -195,44 +196,16 @@ export function useApiKeysColumns(now: number): ColumnDef[] { cell: ({ row }) => { const apiKey = row.original const group = row.getValue('group') as string - const ratio = group && group !== 'auto' ? groupRatios[group] : undefined - - if (group === 'auto') { - return ( - - } - > - - {apiKey.cross_group_retry && ( - - )} - - - - {t( - 'Automatically selects the best available group with circuit breaker mechanism' - )} - - - - ) - } return ( - - - + ) }, - size: 160, + size: 220, meta: { mobileHidden: true }, }, { diff --git a/web/src/features/keys/components/api-keys-mutate-drawer.tsx b/web/src/features/keys/components/api-keys-mutate-drawer.tsx index 9fa64e6f..ed393c53 100644 --- a/web/src/features/keys/components/api-keys-mutate-drawer.tsx +++ b/web/src/features/keys/components/api-keys-mutate-drawer.tsx @@ -19,7 +19,7 @@ For commercial licensing, please contact support@quantumnous.com import { zodResolver } from '@hookform/resolvers/zod' import { useQuery } from '@tanstack/react-query' import { ChevronDown, KeyRound, Settings2, WalletCards } from 'lucide-react' -import { useEffect, useState } from 'react' +import { useEffect, useMemo, useState } from 'react' import { useForm, type SubmitErrorHandler } from 'react-hook-form' import { useTranslation } from 'react-i18next' import { toast } from 'sonner' @@ -67,7 +67,12 @@ import { getUserModels, getUserGroups } from '@/lib/api' import { getCurrencyDisplay, getCurrencyLabel } from '@/lib/currency' import { cn } from '@/lib/utils' -import { createApiKey, updateApiKey, getApiKey } from '../api' +import { + createApiKey, + updateApiKey, + getApiKey, + getTokenAutoGroups, +} from '../api' import { ERROR_MESSAGES, SUCCESS_MESSAGES } from '../constants' import { getApiKeyFormSchema, @@ -82,6 +87,7 @@ import { type ApiKeyGroupOption, } from './api-key-group-combobox' import { useApiKeys } from './api-keys-provider' +import { AutoGroupOrderEditor } from './auto-group-order-editor' type ApiKeyMutateDrawerProps = { open: boolean @@ -96,10 +102,14 @@ export function ApiKeysMutateDrawer({ }: ApiKeyMutateDrawerProps) { const { t } = useTranslation() const isUpdate = !!currentRow + const currentRowId = currentRow?.id const { triggerRefresh } = useApiKeys() - const { status } = useStatus() + const { status, loading: statusLoading } = useStatus() const [isSubmitting, setIsSubmitting] = useState(false) const [advancedOpen, setAdvancedOpen] = useState(false) + const [initializedTarget, setInitializedTarget] = useState( + null + ) const defaultUseAutoGroup = status?.default_use_auto_group === true // Fetch models @@ -111,25 +121,77 @@ export function ApiKeysMutateDrawer({ }) // Fetch groups - const { data: groupsData } = useQuery({ + const { + data: groupsData, + isFetched: groupsFetched, + isFetching: groupsFetching, + } = useQuery({ queryKey: ['user-groups'], queryFn: getUserGroups, enabled: open, staleTime: 0, }) + const { + data: apiKeyData, + isFetched: apiKeyFetched, + isFetching: apiKeyFetching, + } = useQuery({ + queryKey: ['api-key', currentRowId], + queryFn: () => getApiKey(currentRowId ?? 0), + enabled: open && isUpdate && currentRowId !== undefined, + staleTime: 0, + }) + + const { + data: autoGroupsData, + isFetched: autoGroupsFetched, + isFetching: autoGroupsFetching, + } = useQuery({ + queryKey: ['token-auto-groups'], + queryFn: getTokenAutoGroups, + enabled: open, + staleTime: 0, + }) + const models = modelsData?.data || [] - const groupsRaw = groupsData?.data || {} - const groups: ApiKeyGroupOption[] = Object.entries(groupsRaw).map( - ([key, info]) => ({ - value: key, - label: key, - desc: info.desc || key, - ratio: info.ratio, - }) + const groups = useMemo( + () => + Object.entries(groupsData?.data || {}).map(([key, info]) => ({ + value: key, + label: key, + desc: info.desc || key, + ratio: info.ratio, + })), + [groupsData] ) const backendHasAuto = groups.some((g) => g.value === 'auto') - const schema = getApiKeyFormSchema(t) + const availableAutoGroupNames = useMemo( + () => groups.filter((group) => group.value !== 'auto').map((g) => g.value), + [groups] + ) + const globalAutoGroups = useMemo(() => { + const available = new Set(availableAutoGroupNames) + return (autoGroupsData?.data?.groups || []).filter((group) => + available.has(group) + ) + }, [autoGroupsData, availableAutoGroupNames]) + const globalAutoGroupOptions = useMemo(() => { + const groupsByValue = new Map(groups.map((group) => [group.value, group])) + return globalAutoGroups.flatMap((group) => { + const option = groupsByValue.get(group) + return option ? [option] : [] + }) + }, [globalAutoGroups, groups]) + const maxAutoGroups = + Number.isInteger(autoGroupsData?.data?.max_count) && + Number(autoGroupsData?.data?.max_count) > 0 + ? Number(autoGroupsData?.data?.max_count) + : 5 + const schema = useMemo( + () => getApiKeyFormSchema(t, maxAutoGroups), + [t, maxAutoGroups] + ) const form = useForm({ resolver: zodResolver(schema), @@ -138,23 +200,69 @@ export function ApiKeysMutateDrawer({ // Load existing data when updating useEffect(() => { - if (open && isUpdate && currentRow) { - void getApiKey(currentRow.id).then((result) => { - if (result.success && result.data) { - form.reset(transformApiKeyToFormDefaults(result.data)) - } - }) - } else if (open && !isUpdate) { + if (!open) { + setInitializedTarget(null) + return + } + if ( + !groupsFetched || + groupsFetching || + !autoGroupsFetched || + autoGroupsFetching + ) { + return + } + if (isUpdate && (!apiKeyFetched || apiKeyFetching)) return + if (!isUpdate && statusLoading) return + + const target = isUpdate && currentRow ? `update:${currentRow.id}` : 'create' + if (initializedTarget === target) return + if (isUpdate && currentRow) { + if (apiKeyData?.success && apiKeyData.data) { + form.reset( + transformApiKeyToFormDefaults( + apiKeyData.data, + availableAutoGroupNames, + maxAutoGroups + ) + ) + setInitializedTarget(target) + } + } else { form.reset( getApiKeyFormDefaultValues(defaultUseAutoGroup && backendHasAuto) ) + setInitializedTarget(target) } - }, [open, isUpdate, currentRow, form, defaultUseAutoGroup, backendHasAuto]) + }, [ + open, + isUpdate, + currentRow, + form, + defaultUseAutoGroup, + statusLoading, + backendHasAuto, + groupsFetched, + groupsFetching, + autoGroupsFetched, + autoGroupsFetching, + apiKeyData, + apiKeyFetched, + apiKeyFetching, + availableAutoGroupNames, + maxAutoGroups, + initializedTarget, + ]) + + const formTarget = + isUpdate && currentRow ? `update:${currentRow.id}` : 'create' + const isFormInitialized = initializedTarget === formTarget + const selectedGroup = form.watch('group') // Correct group after groups load: if the form value is not in available groups, fall back useEffect(() => { if (groups.length === 0) return - const currentGroup = form.getValues('group') + const currentGroup = selectedGroup if (currentGroup && !groups.some((g) => g.value === currentGroup)) { const fallback = groups.find((g) => g.value === 'default')?.value ?? @@ -162,10 +270,12 @@ export function ApiKeysMutateDrawer({ '' form.setValue('group', fallback) if (currentGroup === 'auto') { + form.setValue('auto_groups', []) + form.setValue('auto_groups_mode', 'inherit') form.setValue('cross_group_retry', false) } } - }, [groups, form]) + }, [groups, form, selectedGroup]) const onSubmit = async (data: ApiKeyFormValues) => { setIsSubmitting(true) @@ -247,7 +357,7 @@ export function ApiKeysMutateDrawer({ const quotaPlaceholder = tokensOnly ? t('Enter quota in tokens') : t('Enter quota in {{currency}}', { currency: currencyLabel }) - const selectedGroup = form.watch('group') + const autoGroupsMode = form.watch('auto_groups_mode') const unlimitedQuota = form.watch('unlimited_quota') return ( @@ -277,6 +387,8 @@ export function ApiKeysMutateDrawer({
@@ -310,7 +422,18 @@ export function ApiKeysMutateDrawer({ { + field.onChange(group) + if (group === 'auto') { + form.setValue('cross_group_retry', true, { + shouldDirty: true, + }) + return + } + form.setValue('cross_group_retry', false, { + shouldDirty: true, + }) + }} placeholder={t('Select a group')} /> @@ -319,6 +442,47 @@ export function ApiKeysMutateDrawer({ )} /> + {selectedGroup === 'auto' && ( + ( + + {t('Auto group order')} + + {t( + 'Choose and order the groups this API key will try.' + )} + + + { + form.setValue('auto_groups_mode', value.mode, { + shouldDirty: true, + shouldValidate: false, + }) + form.setValue( + 'auto_groups', + value.groups.slice(0, maxAutoGroups), + { + shouldDirty: true, + shouldValidate: true, + } + ) + }} + /> + + + + )} + /> + )} + {selectedGroup === 'auto' && ( {isSubmitting ? t('Saving...') : t('Save changes')} diff --git a/web/src/features/keys/components/auto-group-order-editor.tsx b/web/src/features/keys/components/auto-group-order-editor.tsx new file mode 100644 index 00000000..37cca7f8 --- /dev/null +++ b/web/src/features/keys/components/auto-group-order-editor.tsx @@ -0,0 +1,338 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import { + ArrowDown01Icon, + ArrowRight01Icon, + ArrowUp01Icon, + Cancel01Icon, + Drag01Icon, +} from '@hugeicons/core-free-icons' +import { HugeiconsIcon } from '@hugeicons/react' +import { Reorder, useDragControls } from 'motion/react' +import { + useMemo, + type ComponentProps, + type KeyboardEvent, + type PointerEvent, +} from 'react' +import { useTranslation } from 'react-i18next' + +import { Button } from '@/components/ui/button' +import { + Empty, + EmptyDescription, + EmptyHeader, + EmptyTitle, +} from '@/components/ui/empty' +import { cn } from '@/lib/utils' + +import { + ApiKeyGroupCombobox, + type ApiKeyGroupOption, +} from './api-key-group-combobox' +import { GroupRatioBadge } from './auto-group-visuals' + +type AutoGroupOrderEditorProps = Omit, 'onChange'> & { + value: string[] + mode: 'inherit' | 'custom' + options: ApiKeyGroupOption[] + globalOptions: ApiKeyGroupOption[] + maxCount: number + onChange: (value: { groups: string[]; mode: 'inherit' | 'custom' }) => void + 'data-slot'?: string + 'data-form-root'?: string +} + +type AutoGroupOrderItemProps = { + group: string + index: number + count: number + onMove: (index: number, direction: 'up' | 'down') => void + onRemove: (group: string) => void +} + +function AutoGroupOrderItem(props: AutoGroupOrderItemProps) { + const { t } = useTranslation() + const dragControls = useDragControls() + + const handleDragStart = (event: PointerEvent) => { + dragControls.start(event) + } + + const handleDragKeyDown = (event: KeyboardEvent) => { + if (event.key === 'ArrowUp') { + event.preventDefault() + props.onMove(props.index, 'up') + } + if (event.key === 'ArrowDown') { + event.preventDefault() + props.onMove(props.index, 'down') + } + } + + return ( + + + + {props.group} + +
+ + + +
+ + ) +} + +export function AutoGroupOrderEditor(props: AutoGroupOrderEditorProps) { + const { t } = useTranslation() + const maxCount = + Number.isInteger(props.maxCount) && props.maxCount > 0 ? props.maxCount : 5 + const isInheriting = props.mode === 'inherit' + const atLimit = props.value.length >= maxCount + const candidates = useMemo( + () => + props.options.filter( + (option) => + option.value !== 'auto' && !props.value.includes(option.value) + ), + [props.options, props.value] + ) + + const handleAdd = (group: string) => { + if (atLimit || props.value.includes(group)) return + props.onChange({ + groups: [...props.value, group], + mode: 'custom', + }) + } + + const handleRemove = (group: string) => { + props.onChange({ + groups: props.value.filter((item) => item !== group), + mode: 'custom', + }) + } + + const handleMove = (index: number, direction: 'up' | 'down') => { + const targetIndex = direction === 'up' ? index - 1 : index + 1 + if (targetIndex < 0 || targetIndex >= props.value.length) return + const next = [...props.value] + ;[next[index], next[targetIndex]] = [next[targetIndex], next[index]] + props.onChange({ groups: next, mode: 'custom' }) + } + + return ( +
+
+

+ {isInheriting + ? t('Using the complete global Auto order ({{count}} groups)', { + count: props.globalOptions.length, + }) + : t('{{count}} / {{max}} groups selected', { + count: props.value.length, + max: maxCount, + })} +

+ +
+ + + + {isInheriting && props.globalOptions.length === 0 && ( + + + {t('Inherit global Auto order')} + + {t('No available groups in the global Auto order.')} + + + + )} + + {isInheriting && props.globalOptions.length > 0 && ( +
    + {props.globalOptions.map((option, index) => ( +
  1. + {index > 0 && ( +
    2. + ))} +
+ )} + + {!isInheriting && props.value.length === 0 && ( + + + {t('Auto group order')} + + {t( + 'No valid custom Auto groups remain. Add a group or restore global Auto.' + )} + + + + )} + + {!isInheriting && props.value.length > 0 && ( + props.onChange({ groups, mode: 'custom' })} + className='flex flex-col gap-2' + > + {props.value.map((group, index) => ( + + ))} + + )} +
+ ) +} diff --git a/web/src/features/keys/components/auto-group-visuals.tsx b/web/src/features/keys/components/auto-group-visuals.tsx new file mode 100644 index 00000000..a60ef2f4 --- /dev/null +++ b/web/src/features/keys/components/auto-group-visuals.tsx @@ -0,0 +1,140 @@ +/* +Copyright (C) 2023-2026 QuantumNous + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU Affero General Public License as +published by the Free Software Foundation, either version 3 of the +License, or (at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU Affero General Public License for more details. + +You should have received a copy of the GNU Affero General Public License +along with this program. If not, see . + +For commercial licensing, please contact support@quantumnous.com +*/ +import type { ReactNode } from 'react' +import { useTranslation } from 'react-i18next' + +import { GroupBadge } from '@/components/group-badge' +import { Badge } from '@/components/ui/badge' +import { cn } from '@/lib/utils' + +export type GroupRatio = number | string | null | undefined + +export const AUTO_GROUP_FRAME_CLASS_NAME = + 'border-primary/40 relative overflow-visible border shadow-sm shadow-primary/10' + +type AutoGroupFlowBorderProps = { + shouldReduceMotion: boolean +} + +export function AutoGroupFlowBorder(props: AutoGroupFlowBorderProps) { + if (props.shouldReduceMotion) return null + + return ( +