mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-30 16:23:31 +00:00
feat(simple-mode): allow basic account grouping
This commit is contained in:
@@ -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)
|
||||
}
|
||||
Reference in New Issue
Block a user