feat(simple-mode): allow basic account grouping

This commit is contained in:
wucm667
2026-09-06 08:10:41 +08:00
parent ab99d56e96
commit a3675552bd
11 changed files with 165 additions and 46 deletions
+1 -1
View File
@@ -192,7 +192,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver, channelService)
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService, settingService)
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
groupHandler := admin.NewGroupHandlerWithConfig(adminService, dashboardService, groupCapacityService, configConfig)
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository, configConfig)
grokQuotaFetcher := service.NewGrokQuotaFetcher()
@@ -24,6 +24,8 @@ type stubAdminService struct {
boundAuthIdentity *service.AdminBindAuthIdentityInput
boundAuthIdentityFor int64
createdAccounts []*service.CreateAccountInput
createdGroups []*service.CreateGroupInput
updatedGroups []*service.UpdateGroupInput
createdProxies []*service.CreateProxyInput
updatedProxyIDs []int64
updatedProxies []*service.UpdateProxyInput
@@ -353,6 +355,7 @@ func (s *stubAdminService) PreviewCompositeRoute(ctx context.Context, groupID in
}
func (s *stubAdminService) CreateGroup(ctx context.Context, input *service.CreateGroupInput) (*service.Group, error) {
s.createdGroups = append(s.createdGroups, input)
group := service.Group{ID: 200, Name: input.Name, Status: service.StatusActive}
return &group, nil
}
@@ -367,6 +370,7 @@ func (s *stubAdminService) RecoverDuplicateGroup(ctx context.Context, id int64,
}
func (s *stubAdminService) UpdateGroup(ctx context.Context, id int64, input *service.UpdateGroupInput) (*service.Group, error) {
s.updatedGroups = append(s.updatedGroups, input)
group := service.Group{ID: id, Name: input.Name, Status: service.StatusActive}
return &group, nil
}
@@ -10,6 +10,7 @@ import (
"strings"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
@@ -24,6 +25,7 @@ type GroupHandler struct {
adminService service.AdminService
dashboardService *service.DashboardService
groupCapacityService *service.GroupCapacityService
cfg *config.Config
}
// GetLiveCapability 返回当前服务端是否具备生成 Live attestation 的运行环境。
@@ -87,13 +89,70 @@ func (f optionalLimitField) ToServiceInput() *float64 {
// NewGroupHandler creates a new admin group handler
func NewGroupHandler(adminService service.AdminService, dashboardService *service.DashboardService, groupCapacityService *service.GroupCapacityService) *GroupHandler {
return NewGroupHandlerWithConfig(adminService, dashboardService, groupCapacityService, nil)
}
func NewGroupHandlerWithConfig(adminService service.AdminService, dashboardService *service.DashboardService, groupCapacityService *service.GroupCapacityService, cfg *config.Config) *GroupHandler {
return &GroupHandler{
adminService: adminService,
dashboardService: dashboardService,
groupCapacityService: groupCapacityService,
cfg: cfg,
}
}
func (h *GroupHandler) isSimpleMode() bool {
return h != nil && h.cfg != nil && h.cfg.RunMode == config.RunModeSimple
}
func sanitizeCreateGroupRequestForSimpleMode(req *CreateGroupRequest) {
if req == nil {
return
}
req.RateMultiplier = 1
req.IsExclusive = false
req.SubscriptionType = service.SubscriptionTypeStandard
req.DailyLimitUSD, req.WeeklyLimitUSD, req.MonthlyLimitUSD = optionalLimitField{}, optionalLimitField{}, optionalLimitField{}
req.AllowImageGeneration, req.AllowBatchImageGeneration = false, false
req.ImageRateIndependent, req.ImageRateMultiplier = false, nil
req.ImagePrice1K, req.ImagePrice2K, req.ImagePrice4K = nil, nil, nil
req.VideoRateIndependent, req.VideoRateMultiplier = false, nil
req.PeakRateEnabled, req.PeakStart, req.PeakEnd, req.PeakRateMultiplier = false, "", "", nil
req.ProfitControlEnabled, req.ProfitMinMargin, req.ProfitSafetyBuffer = false, nil, nil
req.ClaudeCodeOnly, req.FallbackGroupID, req.FallbackGroupIDOnInvalidRequest = false, nil, nil
req.ModelRouting, req.ModelRoutingEnabled, req.MCPXMLInject = nil, false, nil
req.SupportedModelScopes, req.AllowMessagesDispatch, req.AllowLive = nil, false, false
req.RequireOAuthOnly, req.RequirePrivacySet, req.DefaultMappedModel = false, false, ""
req.MessagesDispatchModelConfig, req.ModelsListConfig = service.OpenAIMessagesDispatchModelConfig{}, service.GroupModelsListConfig{}
req.RPMLimit, req.MaxReasoningEffort, req.ReasoningEffortMappings = 0, "", nil
req.CopyAccountsFromGroupIDs = nil
}
func sanitizeUpdateGroupRequestForSimpleMode(req *UpdateGroupRequest) {
if req == nil {
return
}
rate, exclusive, disabled, image, messages, live, oauth, privacy := 1.0, false, false, false, false, false, false, false
rpm := 0
req.RateMultiplier, req.IsExclusive = &rate, &exclusive
req.Status, req.SubscriptionType = "", service.SubscriptionTypeStandard
req.DailyLimitUSD, req.WeeklyLimitUSD, req.MonthlyLimitUSD = optionalLimitField{}, optionalLimitField{}, optionalLimitField{}
req.AllowImageGeneration, req.AllowBatchImageGeneration = &image, &image
req.ImageRateIndependent, req.ImageRateMultiplier = &disabled, nil
req.ImagePrice1K, req.ImagePrice2K, req.ImagePrice4K = nil, nil, nil
req.VideoRateIndependent, req.VideoRateMultiplier = &disabled, nil
req.PeakRateEnabled, req.PeakStart, req.PeakEnd, req.PeakRateMultiplier = &disabled, nil, nil, nil
req.ProfitControlEnabled, req.ProfitMinMargin, req.ProfitSafetyBuffer = &disabled, nil, nil
req.ClaudeCodeOnly, req.FallbackGroupID, req.FallbackGroupIDOnInvalidRequest = &disabled, nil, nil
req.ModelRouting, req.ModelRoutingEnabled, req.MCPXMLInject = nil, &disabled, nil
req.SupportedModelScopes = nil
req.AllowMessagesDispatch, req.AllowLive = &messages, &live
req.RequireOAuthOnly, req.RequirePrivacySet = &oauth, &privacy
req.DefaultMappedModel, req.MessagesDispatchModelConfig, req.ModelsListConfig = nil, nil, nil
req.RPMLimit, req.MaxReasoningEffort, req.ReasoningEffortMappings = &rpm, nil, nil
req.CopyAccountsFromGroupIDs = nil
}
// CreateGroupRequest represents create group request
type CreateGroupRequest struct {
Name string `json:"name" binding:"required"`
@@ -501,6 +560,9 @@ func (h *GroupHandler) Create(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if h.isSimpleMode() {
sanitizeCreateGroupRequestForSimpleMode(&req)
}
if err := service.ValidatePeakRateConfig(req.SubscriptionType, req.PeakRateEnabled, req.PeakStart, req.PeakEnd, float64ValueOrDefault(req.PeakRateMultiplier, 1.0)); err != nil {
response.BadRequest(c, err.Error())
@@ -645,6 +707,9 @@ func (h *GroupHandler) Update(c *gin.Context) {
response.BadRequest(c, "Invalid request: "+err.Error())
return
}
if h.isSimpleMode() {
sanitizeUpdateGroupRequestForSimpleMode(&req)
}
group, err := h.adminService.UpdateGroup(c.Request.Context(), groupID, &service.UpdateGroupInput{
Name: req.Name,
@@ -0,0 +1,57 @@
package admin
import (
"bytes"
"net/http"
"net/http/httptest"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/service"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
func TestGroupHandlerSimpleModeSanitizesCommercialFields(t *testing.T) {
gin.SetMode(gin.TestMode)
svc := newStubAdminService()
h := NewGroupHandlerWithConfig(svc, nil, nil, &config.Config{RunMode: config.RunModeSimple})
r := gin.New()
r.POST("/groups", h.Create)
r.PUT("/groups/:id", h.Update)
create := `{"name":"simple","platform":"anthropic","rate_multiplier":7,"is_exclusive":true,"subscription_type":"subscription","daily_limit_usd":10,"allow_image_generation":true,"rpm_limit":99}`
req := httptest.NewRequest(http.MethodPost, "/groups", bytes.NewBufferString(create))
req.Header.Set("Content-Type", "application/json")
res := httptest.NewRecorder()
r.ServeHTTP(res, req)
require.Equal(t, http.StatusOK, res.Code)
require.Len(t, svc.createdGroups, 1)
created := svc.createdGroups[0]
require.Equal(t, 1.0, created.RateMultiplier)
require.False(t, created.IsExclusive)
require.Equal(t, service.SubscriptionTypeStandard, created.SubscriptionType)
require.Nil(t, created.DailyLimitUSD)
require.False(t, created.AllowImageGeneration)
require.Zero(t, created.RPMLimit)
update := `{"name":"renamed","rate_multiplier":9,"is_exclusive":true,"subscription_type":"subscription","daily_limit_usd":12,"allow_image_generation":true,"status":"inactive","rpm_limit":123}`
req = httptest.NewRequest(http.MethodPut, "/groups/2", bytes.NewBufferString(update))
req.Header.Set("Content-Type", "application/json")
res = httptest.NewRecorder()
r.ServeHTTP(res, req)
require.Equal(t, http.StatusOK, res.Code)
require.Len(t, svc.updatedGroups, 1)
updated := svc.updatedGroups[0]
require.NotNil(t, updated.RateMultiplier)
require.Equal(t, 1.0, *updated.RateMultiplier)
require.NotNil(t, updated.IsExclusive)
require.False(t, *updated.IsExclusive)
require.Equal(t, service.SubscriptionTypeStandard, updated.SubscriptionType)
require.Nil(t, updated.DailyLimitUSD)
require.NotNil(t, updated.AllowImageGeneration)
require.False(t, *updated.AllowImageGeneration)
require.Empty(t, updated.Status)
require.NotNil(t, updated.RPMLimit)
require.Zero(t, *updated.RPMLimit)
}