diff --git a/backend/cmd/server/wire_gen.go b/backend/cmd/server/wire_gen.go index 8a745c8105..0a94cb6384 100644 --- a/backend/cmd/server/wire_gen.go +++ b/backend/cmd/server/wire_gen.go @@ -150,7 +150,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { rpmCache := repository.NewRPMCache(redisClient) groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache) groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService) - accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, gatewayCache) + accountHandler := admin.NewAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator) adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService) dataManagementService := service.NewDataManagementService() dataManagementHandler := admin.NewDataManagementHandler(dataManagementService) @@ -189,7 +189,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { soraGenerationRepository := repository.NewSoraGenerationRepository(db) soraQuotaService := service.NewSoraQuotaService(userRepository, groupRepository, settingService) soraGenerationService := service.NewSoraGenerationService(soraGenerationRepository, soraS3Storage, soraQuotaService) - settingHandler := admin.NewSettingHandler(settingService, emailService, turnstileService, opsService, soraS3Storage) + settingHandler := admin.NewSettingHandler(settingService, emailService, turnstileService, opsService) opsHandler := admin.NewOpsHandler(opsService) updateCache := repository.NewUpdateCache(redisClient) gitHubReleaseClient := repository.ProvideGitHubReleaseClient(configConfig) @@ -226,18 +226,13 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { soraSDKClient := service.ProvideSoraSDKClient(configConfig, httpUpstream, openAITokenProvider, accountRepository, soraAccountRepository) soraMediaStorage := service.ProvideSoraMediaStorage(configConfig) soraGatewayService := service.NewSoraGatewayService(soraSDKClient, rateLimitService, httpUpstream, configConfig) - soraClientHandler := handler.NewSoraClientHandler(soraGenerationService, soraQuotaService, soraS3Storage, soraGatewayService, gatewayService, soraMediaStorage, apiKeyService) - soraTaskRepository := repository.NewSoraTaskRepository(db) - soraTaskService := service.NewSoraTaskService(soraTaskRepository, accountRepository, soraSDKClient, httpUpstream) - soraTaskWorker := service.NewSoraTaskWorker(soraTaskService, accountRepository, soraS3Storage, soraMediaStorage, 60*time.Second) - soraTaskWorker.Start() - soraVideosHandler := handler.NewSoraVideosHandler(soraTaskService, gatewayService, soraS3Storage, soraMediaStorage, soraGatewayService) + soraClientHandler := handler.NewSoraClientHandler(soraGenerationService, soraQuotaService, soraGatewayService, gatewayService, soraMediaStorage, apiKeyService) soraGatewayHandler := handler.NewSoraGatewayHandler(gatewayService, soraGatewayService, concurrencyService, billingCacheService, usageRecordWorkerPool, configConfig) handlerSettingHandler := handler.ProvideSettingHandler(settingService, buildInfo) totpHandler := handler.NewTotpHandler(totpService) idempotencyCoordinator := service.ProvideIdempotencyCoordinator(idempotencyRepository, configConfig) idempotencyCleanupService := service.ProvideIdempotencyCleanupService(idempotencyRepository, configConfig) - handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, soraGatewayHandler, soraClientHandler, soraVideosHandler, handlerSettingHandler, totpHandler, idempotencyCoordinator, idempotencyCleanupService) + handlers := handler.ProvideHandlers(authHandler, userHandler, apiKeyHandler, usageHandler, redeemHandler, subscriptionHandler, announcementHandler, adminHandlers, gatewayHandler, openAIGatewayHandler, soraGatewayHandler, soraClientHandler, handlerSettingHandler, totpHandler, idempotencyCoordinator, idempotencyCleanupService) jwtAuthMiddleware := middleware.NewJWTAuthMiddleware(authService, userService) adminAuthMiddleware := middleware.NewAdminAuthMiddleware(authService, userService, settingService) apiKeyAuthMiddleware := middleware.NewAPIKeyAuthMiddleware(apiKeyService, subscriptionService, configConfig) @@ -253,7 +248,7 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) { accountExpiryService := service.ProvideAccountExpiryService(accountRepository) subscriptionExpiryService := service.ProvideSubscriptionExpiryService(userSubscriptionRepository) scheduledTestRunnerService := service.ProvideScheduledTestRunnerService(scheduledTestPlanRepository, scheduledTestService, accountTestService, rateLimitService, configConfig) - v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, soraMediaCleanupService, schedulerSnapshotService, tokenRefreshService, accountExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService, soraTaskWorker) + v := provideCleanup(client, redisClient, opsMetricsCollector, opsAggregationService, opsAlertEvaluatorService, opsCleanupService, opsScheduledReportService, opsSystemLogSink, soraMediaCleanupService, schedulerSnapshotService, tokenRefreshService, accountExpiryService, subscriptionExpiryService, usageCleanupService, idempotencyCleanupService, pricingService, emailQueueService, billingCacheService, usageRecordWorkerPool, subscriptionService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, openAIGatewayService, scheduledTestRunnerService, backupService) application := &Application{ Server: httpServer, Cleanup: v, @@ -307,7 +302,6 @@ func provideCleanup( openAIGateway *service.OpenAIGatewayService, scheduledTestRunner *service.ScheduledTestRunnerService, backupSvc *service.BackupService, - soraTaskWorker *service.SoraTaskWorker, ) func() { return func() { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) @@ -449,12 +443,6 @@ func provideCleanup( } return nil }}, - {"SoraTaskWorker", func() error { - if soraTaskWorker != nil { - soraTaskWorker.Stop() - } - return nil - }}, } infraSteps := []cleanupStep{ diff --git a/backend/cmd/server/wire_gen_test.go b/backend/cmd/server/wire_gen_test.go index 51396c83e4..9d2a54b98b 100644 --- a/backend/cmd/server/wire_gen_test.go +++ b/backend/cmd/server/wire_gen_test.go @@ -76,7 +76,6 @@ func TestProvideCleanup_WithMinimalDependencies_NoPanic(t *testing.T) { nil, // openAIGateway nil, // scheduledTestRunner nil, // backupSvc - nil, // soraTaskWorker ) require.NotPanics(t, func() { diff --git a/backend/ent/group.go b/backend/ent/group.go index 7ed4990538..8b5df0a8a2 100644 --- a/backend/ent/group.go +++ b/backend/ent/group.go @@ -82,8 +82,6 @@ type Group struct { AllowMessagesDispatch bool `json:"allow_messages_dispatch,omitempty"` // 默认映射模型 ID，当账号级映射找不到时使用此值 DefaultMappedModel string `json:"default_mapped_model,omitempty"` - // simulate claude usage as claude-max style (1h cache write) - SimulateClaudeMaxEnabled bool `json:"simulate_claude_max_enabled,omitempty"` // Edges holds the relations/edges for other nodes in the graph. // The values are being populated by the GroupQuery when eager-loading is set. Edges GroupEdges `json:"edges"` @@ -192,7 +190,7 @@ func (*Group) scanValues(columns []string) ([]any, error) { switch columns[i] { case group.FieldModelRouting, group.FieldSupportedModelScopes: values[i] = new([]byte) - case group.FieldIsExclusive, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldSimulateClaudeMaxEnabled: + case group.FieldIsExclusive, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch: values[i] = new(sql.NullBool) case group.FieldRateMultiplier, group.FieldDailyLimitUsd, group.FieldWeeklyLimitUsd, group.FieldMonthlyLimitUsd, group.FieldImagePrice1k, group.FieldImagePrice2k, group.FieldImagePrice4k, group.FieldSoraImagePrice360, group.FieldSoraImagePrice540, group.FieldSoraVideoPricePerRequest, group.FieldSoraVideoPricePerRequestHd: values[i] = new(sql.NullFloat64) @@ -433,12 +431,6 @@ func (_m *Group) assignValues(columns []string, values []any) error { } else if value.Valid { _m.DefaultMappedModel = value.String } - case group.FieldSimulateClaudeMaxEnabled: - if value, ok := values[i].(*sql.NullBool); !ok { - return fmt.Errorf("unexpected type %T for field simulate_claude_max_enabled", values[i]) - } else if value.Valid { - _m.SimulateClaudeMaxEnabled = value.Bool - } default: _m.selectValues.Set(columns[i], values[i]) } @@ -638,9 +630,6 @@ func (_m *Group) String() string { builder.WriteString(", ") builder.WriteString("default_mapped_model=") builder.WriteString(_m.DefaultMappedModel) - builder.WriteString(", ") - builder.WriteString("simulate_claude_max_enabled=") - builder.WriteString(fmt.Sprintf("%v", _m.SimulateClaudeMaxEnabled)) builder.WriteByte(')') return builder.String() } diff --git a/backend/ent/group/group.go b/backend/ent/group/group.go index 970c7a85c3..2612b6cff2 100644 --- a/backend/ent/group/group.go +++ b/backend/ent/group/group.go @@ -79,8 +79,6 @@ const ( FieldAllowMessagesDispatch = "allow_messages_dispatch" // FieldDefaultMappedModel holds the string denoting the default_mapped_model field in the database. FieldDefaultMappedModel = "default_mapped_model" - // FieldSimulateClaudeMaxEnabled holds the string denoting the simulate_claude_max_enabled field in the database. - FieldSimulateClaudeMaxEnabled = "simulate_claude_max_enabled" // EdgeAPIKeys holds the string denoting the api_keys edge name in mutations. EdgeAPIKeys = "api_keys" // EdgeRedeemCodes holds the string denoting the redeem_codes edge name in mutations. @@ -188,7 +186,6 @@ var Columns = []string{ FieldSortOrder, FieldAllowMessagesDispatch, FieldDefaultMappedModel, - FieldSimulateClaudeMaxEnabled, } var ( @@ -262,8 +259,6 @@ var ( DefaultDefaultMappedModel string // DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save. DefaultMappedModelValidator func(string) error - // DefaultSimulateClaudeMaxEnabled holds the default value on creation for the "simulate_claude_max_enabled" field. - DefaultSimulateClaudeMaxEnabled bool ) // OrderOption defines the ordering options for the Group queries. @@ -424,11 +419,6 @@ func ByDefaultMappedModel(opts ...sql.OrderTermOption) OrderOption { return sql.OrderByField(FieldDefaultMappedModel, opts...).ToFunc() } -// BySimulateClaudeMaxEnabled orders the results by the simulate_claude_max_enabled field. -func BySimulateClaudeMaxEnabled(opts ...sql.OrderTermOption) OrderOption { - return sql.OrderByField(FieldSimulateClaudeMaxEnabled, opts...).ToFunc() -} - // ByAPIKeysCount orders the results by api_keys count. func ByAPIKeysCount(opts ...sql.OrderTermOption) OrderOption { return func(s *sql.Selector) { diff --git a/backend/ent/group/where.go b/backend/ent/group/where.go index 62c91d5af0..5dd8759e5d 100644 --- a/backend/ent/group/where.go +++ b/backend/ent/group/where.go @@ -205,11 +205,6 @@ func DefaultMappedModel(v string) predicate.Group { return predicate.Group(sql.FieldEQ(FieldDefaultMappedModel, v)) } -// SimulateClaudeMaxEnabled applies equality check predicate on the "simulate_claude_max_enabled" field. It's identical to SimulateClaudeMaxEnabledEQ. -func SimulateClaudeMaxEnabled(v bool) predicate.Group { - return predicate.Group(sql.FieldEQ(FieldSimulateClaudeMaxEnabled, v)) -} - // CreatedAtEQ applies the EQ predicate on the "created_at" field. func CreatedAtEQ(v time.Time) predicate.Group { return predicate.Group(sql.FieldEQ(FieldCreatedAt, v)) @@ -1560,16 +1555,6 @@ func DefaultMappedModelContainsFold(v string) predicate.Group { return predicate.Group(sql.FieldContainsFold(FieldDefaultMappedModel, v)) } -// SimulateClaudeMaxEnabledEQ applies the EQ predicate on the "simulate_claude_max_enabled" field. -func SimulateClaudeMaxEnabledEQ(v bool) predicate.Group { - return predicate.Group(sql.FieldEQ(FieldSimulateClaudeMaxEnabled, v)) -} - -// SimulateClaudeMaxEnabledNEQ applies the NEQ predicate on the "simulate_claude_max_enabled" field. -func SimulateClaudeMaxEnabledNEQ(v bool) predicate.Group { - return predicate.Group(sql.FieldNEQ(FieldSimulateClaudeMaxEnabled, v)) -} - // HasAPIKeys applies the HasEdge predicate on the "api_keys" edge. func HasAPIKeys() predicate.Group { return predicate.Group(func(s *sql.Selector) { diff --git a/backend/ent/group_create.go b/backend/ent/group_create.go index 9418b02f7a..6db5b97452 100644 --- a/backend/ent/group_create.go +++ b/backend/ent/group_create.go @@ -452,20 +452,6 @@ func (_c *GroupCreate) SetNillableDefaultMappedModel(v *string) *GroupCreate { return _c } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (_c *GroupCreate) SetSimulateClaudeMaxEnabled(v bool) *GroupCreate { - _c.mutation.SetSimulateClaudeMaxEnabled(v) - return _c -} - -// SetNillableSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field if the given value is not nil. -func (_c *GroupCreate) SetNillableSimulateClaudeMaxEnabled(v *bool) *GroupCreate { - if v != nil { - _c.SetSimulateClaudeMaxEnabled(*v) - } - return _c -} - // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs. func (_c *GroupCreate) AddAPIKeyIDs(ids ...int64) *GroupCreate { _c.mutation.AddAPIKeyIDs(ids...) @@ -663,10 +649,6 @@ func (_c *GroupCreate) defaults() error { v := group.DefaultDefaultMappedModel _c.mutation.SetDefaultMappedModel(v) } - if _, ok := _c.mutation.SimulateClaudeMaxEnabled(); !ok { - v := group.DefaultSimulateClaudeMaxEnabled - _c.mutation.SetSimulateClaudeMaxEnabled(v) - } return nil } @@ -748,9 +730,6 @@ func (_c *GroupCreate) check() error { return &ValidationError{Name: "default_mapped_model", err: fmt.Errorf(`ent: validator failed for field "Group.default_mapped_model": %w`, err)} } } - if _, ok := _c.mutation.SimulateClaudeMaxEnabled(); !ok { - return &ValidationError{Name: "simulate_claude_max_enabled", err: errors.New(`ent: missing required field "Group.simulate_claude_max_enabled"`)} - } return nil } @@ -906,10 +885,6 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) { _spec.SetField(group.FieldDefaultMappedModel, field.TypeString, value) _node.DefaultMappedModel = value } - if value, ok := _c.mutation.SimulateClaudeMaxEnabled(); ok { - _spec.SetField(group.FieldSimulateClaudeMaxEnabled, field.TypeBool, value) - _node.SimulateClaudeMaxEnabled = value - } if nodes := _c.mutation.APIKeysIDs(); len(nodes) > 0 { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.O2M, @@ -1624,18 +1599,6 @@ func (u *GroupUpsert) UpdateDefaultMappedModel() *GroupUpsert { return u } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (u *GroupUpsert) SetSimulateClaudeMaxEnabled(v bool) *GroupUpsert { - u.Set(group.FieldSimulateClaudeMaxEnabled, v) - return u -} - -// UpdateSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field to the value that was provided on create. -func (u *GroupUpsert) UpdateSimulateClaudeMaxEnabled() *GroupUpsert { - u.SetExcluded(group.FieldSimulateClaudeMaxEnabled) - return u -} - // UpdateNewValues updates the mutable fields using the new values that were set on create. // Using this option is equivalent to using: // @@ -2332,20 +2295,6 @@ func (u *GroupUpsertOne) UpdateDefaultMappedModel() *GroupUpsertOne { }) } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (u *GroupUpsertOne) SetSimulateClaudeMaxEnabled(v bool) *GroupUpsertOne { - return u.Update(func(s *GroupUpsert) { - s.SetSimulateClaudeMaxEnabled(v) - }) -} - -// UpdateSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field to the value that was provided on create. -func (u *GroupUpsertOne) UpdateSimulateClaudeMaxEnabled() *GroupUpsertOne { - return u.Update(func(s *GroupUpsert) { - s.UpdateSimulateClaudeMaxEnabled() - }) -} - // Exec executes the query. func (u *GroupUpsertOne) Exec(ctx context.Context) error { if len(u.create.conflict) == 0 { @@ -3208,20 +3157,6 @@ func (u *GroupUpsertBulk) UpdateDefaultMappedModel() *GroupUpsertBulk { }) } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (u *GroupUpsertBulk) SetSimulateClaudeMaxEnabled(v bool) *GroupUpsertBulk { - return u.Update(func(s *GroupUpsert) { - s.SetSimulateClaudeMaxEnabled(v) - }) -} - -// UpdateSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field to the value that was provided on create. -func (u *GroupUpsertBulk) UpdateSimulateClaudeMaxEnabled() *GroupUpsertBulk { - return u.Update(func(s *GroupUpsert) { - s.UpdateSimulateClaudeMaxEnabled() - }) -} - // Exec executes the query. func (u *GroupUpsertBulk) Exec(ctx context.Context) error { if u.create.err != nil { diff --git a/backend/ent/group_update.go b/backend/ent/group_update.go index 75955f7df5..b3698596f6 100644 --- a/backend/ent/group_update.go +++ b/backend/ent/group_update.go @@ -653,20 +653,6 @@ func (_u *GroupUpdate) SetNillableDefaultMappedModel(v *string) *GroupUpdate { return _u } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (_u *GroupUpdate) SetSimulateClaudeMaxEnabled(v bool) *GroupUpdate { - _u.mutation.SetSimulateClaudeMaxEnabled(v) - return _u -} - -// SetNillableSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field if the given value is not nil. -func (_u *GroupUpdate) SetNillableSimulateClaudeMaxEnabled(v *bool) *GroupUpdate { - if v != nil { - _u.SetSimulateClaudeMaxEnabled(*v) - } - return _u -} - // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs. func (_u *GroupUpdate) AddAPIKeyIDs(ids ...int64) *GroupUpdate { _u.mutation.AddAPIKeyIDs(ids...) @@ -1163,9 +1149,6 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) { if value, ok := _u.mutation.DefaultMappedModel(); ok { _spec.SetField(group.FieldDefaultMappedModel, field.TypeString, value) } - if value, ok := _u.mutation.SimulateClaudeMaxEnabled(); ok { - _spec.SetField(group.FieldSimulateClaudeMaxEnabled, field.TypeBool, value) - } if _u.mutation.APIKeysCleared() { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.O2M, @@ -2098,20 +2081,6 @@ func (_u *GroupUpdateOne) SetNillableDefaultMappedModel(v *string) *GroupUpdateO return _u } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (_u *GroupUpdateOne) SetSimulateClaudeMaxEnabled(v bool) *GroupUpdateOne { - _u.mutation.SetSimulateClaudeMaxEnabled(v) - return _u -} - -// SetNillableSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field if the given value is not nil. -func (_u *GroupUpdateOne) SetNillableSimulateClaudeMaxEnabled(v *bool) *GroupUpdateOne { - if v != nil { - _u.SetSimulateClaudeMaxEnabled(*v) - } - return _u -} - // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by IDs. func (_u *GroupUpdateOne) AddAPIKeyIDs(ids ...int64) *GroupUpdateOne { _u.mutation.AddAPIKeyIDs(ids...) @@ -2638,9 +2607,6 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error) if value, ok := _u.mutation.DefaultMappedModel(); ok { _spec.SetField(group.FieldDefaultMappedModel, field.TypeString, value) } - if value, ok := _u.mutation.SimulateClaudeMaxEnabled(); ok { - _spec.SetField(group.FieldSimulateClaudeMaxEnabled, field.TypeBool, value) - } if _u.mutation.APIKeysCleared() { edge := &sqlgraph.EdgeSpec{ Rel: sqlgraph.O2M, diff --git a/backend/ent/migrate/schema.go b/backend/ent/migrate/schema.go index 59787db096..d85182dff3 100644 --- a/backend/ent/migrate/schema.go +++ b/backend/ent/migrate/schema.go @@ -410,7 +410,6 @@ var ( {Name: "sort_order", Type: field.TypeInt, Default: 0}, {Name: "allow_messages_dispatch", Type: field.TypeBool, Default: false}, {Name: "default_mapped_model", Type: field.TypeString, Size: 100, Default: ""}, - {Name: "simulate_claude_max_enabled", Type: field.TypeBool, Default: false}, } // GroupsTable holds the schema information for the "groups" table. GroupsTable = &schema.Table{ diff --git a/backend/ent/mutation.go b/backend/ent/mutation.go index 5f6abfa60c..a9237c2529 100644 --- a/backend/ent/mutation.go +++ b/backend/ent/mutation.go @@ -8254,7 +8254,6 @@ type GroupMutation struct { addsort_order *int allow_messages_dispatch *bool default_mapped_model *string - simulate_claude_max_enabled *bool clearedFields map[string]struct{} api_keys map[int64]struct{} removedapi_keys map[int64]struct{} @@ -10071,42 +10070,6 @@ func (m *GroupMutation) ResetDefaultMappedModel() { m.default_mapped_model = nil } -// SetSimulateClaudeMaxEnabled sets the "simulate_claude_max_enabled" field. -func (m *GroupMutation) SetSimulateClaudeMaxEnabled(b bool) { - m.simulate_claude_max_enabled = &b -} - -// SimulateClaudeMaxEnabled returns the value of the "simulate_claude_max_enabled" field in the mutation. -func (m *GroupMutation) SimulateClaudeMaxEnabled() (r bool, exists bool) { - v := m.simulate_claude_max_enabled - if v == nil { - return - } - return *v, true -} - -// OldSimulateClaudeMaxEnabled returns the old "simulate_claude_max_enabled" field's value of the Group entity. -// If the Group object wasn't provided to the builder, the object is fetched from the database. -// An error is returned if the mutation operation is not UpdateOne, or the database query fails. -func (m *GroupMutation) OldSimulateClaudeMaxEnabled(ctx context.Context) (v bool, err error) { - if !m.op.Is(OpUpdateOne) { - return v, errors.New("OldSimulateClaudeMaxEnabled is only allowed on UpdateOne operations") - } - if m.id == nil || m.oldValue == nil { - return v, errors.New("OldSimulateClaudeMaxEnabled requires an ID field in the mutation") - } - oldValue, err := m.oldValue(ctx) - if err != nil { - return v, fmt.Errorf("querying old value for OldSimulateClaudeMaxEnabled: %w", err) - } - return oldValue.SimulateClaudeMaxEnabled, nil -} - -// ResetSimulateClaudeMaxEnabled resets all changes to the "simulate_claude_max_enabled" field. -func (m *GroupMutation) ResetSimulateClaudeMaxEnabled() { - m.simulate_claude_max_enabled = nil -} - // AddAPIKeyIDs adds the "api_keys" edge to the APIKey entity by ids. func (m *GroupMutation) AddAPIKeyIDs(ids ...int64) { if m.api_keys == nil { @@ -10562,9 +10525,6 @@ func (m *GroupMutation) Fields() []string { if m.default_mapped_model != nil { fields = append(fields, group.FieldDefaultMappedModel) } - if m.simulate_claude_max_enabled != nil { - fields = append(fields, group.FieldSimulateClaudeMaxEnabled) - } return fields } @@ -10637,8 +10597,6 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) { return m.AllowMessagesDispatch() case group.FieldDefaultMappedModel: return m.DefaultMappedModel() - case group.FieldSimulateClaudeMaxEnabled: - return m.SimulateClaudeMaxEnabled() } return nil, false } @@ -10712,8 +10670,6 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e return m.OldAllowMessagesDispatch(ctx) case group.FieldDefaultMappedModel: return m.OldDefaultMappedModel(ctx) - case group.FieldSimulateClaudeMaxEnabled: - return m.OldSimulateClaudeMaxEnabled(ctx) } return nil, fmt.Errorf("unknown Group field %s", name) } @@ -10947,13 +10903,6 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error { } m.SetDefaultMappedModel(v) return nil - case group.FieldSimulateClaudeMaxEnabled: - v, ok := value.(bool) - if !ok { - return fmt.Errorf("unexpected type %T for field %s", value, name) - } - m.SetSimulateClaudeMaxEnabled(v) - return nil } return fmt.Errorf("unknown Group field %s", name) } @@ -11387,9 +11336,6 @@ func (m *GroupMutation) ResetField(name string) error { case group.FieldDefaultMappedModel: m.ResetDefaultMappedModel() return nil - case group.FieldSimulateClaudeMaxEnabled: - m.ResetSimulateClaudeMaxEnabled() - return nil } return fmt.Errorf("unknown Group field %s", name) } diff --git a/backend/ent/runtime/runtime.go b/backend/ent/runtime/runtime.go index 5724de33d4..cede5fc937 100644 --- a/backend/ent/runtime/runtime.go +++ b/backend/ent/runtime/runtime.go @@ -464,10 +464,6 @@ func init() { group.DefaultDefaultMappedModel = groupDescDefaultMappedModel.Default.(string) // group.DefaultMappedModelValidator is a validator for the "default_mapped_model" field. It is called by the builders before save. group.DefaultMappedModelValidator = groupDescDefaultMappedModel.Validators[0].(func(string) error) - // groupDescSimulateClaudeMaxEnabled is the schema descriptor for simulate_claude_max_enabled field. - groupDescSimulateClaudeMaxEnabled := groupFields[29].Descriptor() - // group.DefaultSimulateClaudeMaxEnabled holds the default value on creation for the simulate_claude_max_enabled field. - group.DefaultSimulateClaudeMaxEnabled = groupDescSimulateClaudeMaxEnabled.Default.(bool) idempotencyrecordMixin := schema.IdempotencyRecord{}.Mixin() idempotencyrecordMixinFields0 := idempotencyrecordMixin[0].Fields() _ = idempotencyrecordMixinFields0 diff --git a/backend/ent/schema/group.go b/backend/ent/schema/group.go index 0842a0f846..f5bac05166 100644 --- a/backend/ent/schema/group.go +++ b/backend/ent/schema/group.go @@ -145,9 +145,6 @@ func (Group) Fields() []ent.Field { MaxLen(100). Default(""). Comment("默认映射模型 ID，当账号级映射找不到时使用此值"), - field.Bool("simulate_claude_max_enabled"). - Default(false). - Comment("simulate claude usage as claude-max style (1h cache write)"), } } diff --git a/backend/internal/handler/admin/account_affinity_handler.go b/backend/internal/handler/admin/account_affinity_handler.go deleted file mode 100644 index 62f4c3fe80..0000000000 --- a/backend/internal/handler/admin/account_affinity_handler.go +++ /dev/null @@ -1,254 +0,0 @@ -package admin - -import ( - "context" - "log/slog" - "strconv" - "strings" - "time" - - "github.com/Wei-Shaw/sub2api/internal/pkg/response" - "github.com/Wei-Shaw/sub2api/internal/service" - "github.com/gin-gonic/gin" -) - -// 账号亲和调度 API 与契约说明见 docs/ACCOUNT_SCHEDULING_FLOW.md 。 - -// GetAffinityClients returns the list of affinity clients for an account with last active timestamps. -// GET /api/v1/admin/accounts/:id/affinity-clients -func (h *AccountHandler) GetAffinityClients(c *gin.Context) { - accountID, err := strconv.ParseInt(c.Param("id"), 10, 64) - if err != nil { - response.BadRequest(c, "Invalid account ID") - return - } - - account, err := h.adminService.GetAccount(c.Request.Context(), accountID) - if err != nil { - response.ErrorFrom(c, err) - return - } - - if !account.IsAffinityEnabled() { - response.Success(c, []service.AffinityClient{}) - return - } - - if h.gatewayCache == nil || len(account.GroupIDs) == 0 { - response.Success(c, []service.AffinityClient{}) - return - } - - clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores( - c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL, - ) - if err != nil { - response.Success(c, []service.AffinityClient{}) - return - } - - response.Success(c, clients) -} - -// clearAccountAffinity 清除指定账号在所有分组的亲和记录。 -func (h *AccountHandler) clearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) { - if h.gatewayCache == nil || len(groupIDs) == 0 { - return - } - if err := h.gatewayCache.ClearAccountAffinity(ctx, accountID, groupIDs); err != nil { - // 清理失败不影响主流程，记录日志即可 - slog.Warn("clear account affinity failed", - "account_id", accountID, - "error", err, - ) - } -} - -// countUniqueUsersFromAffinityMembers 从 "{userID}/{clientID}" 格式的成员列表中计算唯一用户数。 -func countUniqueUsersFromAffinityMembers(members []string) int64 { - users := make(map[string]struct{}, len(members)) - for _, m := range members { - if idx := strings.Index(m, "/"); idx > 0 { - users[m[:idx]] = struct{}{} - } - } - return int64(len(users)) -} - -// AffinityDetailClient 亲和详情中的单个客户端信息 -type AffinityDetailClient struct { - ClientID string `json:"client_id"` - LastActive time.Time `json:"last_active"` - IsPinned bool `json:"is_pinned"` -} - -// AffinityDetailUser 亲和详情中按用户分组的信息 -type AffinityDetailUser struct { - UserID int64 `json:"user_id"` - UserEmail string `json:"user_email"` - ClientCount int `json:"client_count"` - IsPinned bool `json:"is_pinned"` - Clients []AffinityDetailClient `json:"clients"` -} - -// AffinityDetailsResponse 亲和详情响应 -type AffinityDetailsResponse struct { - Users []AffinityDetailUser `json:"users"` - TotalUsers int `json:"total_users"` - TotalClients int `json:"total_clients"` - PinnedUsers []int64 `json:"pinned_users"` -} - -type affinityStateSnapshot struct { - enabled bool - groupIDs []int64 -} - -func (h *AccountHandler) captureAffinityStates(ctx context.Context, accountIDs []int64) map[int64]affinityStateSnapshot { - states := make(map[int64]affinityStateSnapshot, len(accountIDs)) - if h.gatewayCache == nil || len(accountIDs) == 0 { - return states - } - for _, accountID := range accountIDs { - account, err := h.adminService.GetAccount(ctx, accountID) - if err != nil || account == nil { - continue - } - states[accountID] = affinityStateSnapshot{ - enabled: account.IsAffinityEnabled(), - groupIDs: append([]int64(nil), account.GroupIDs...), - } - } - return states -} - -func (h *AccountHandler) clearAffinityCacheIfDisabled(ctx context.Context, accountID int64, oldState affinityStateSnapshot) { - if h.gatewayCache == nil || !oldState.enabled { - return - } - account, err := h.adminService.GetAccount(ctx, accountID) - if err != nil || account == nil || account.IsAffinityEnabled() { - return - } - groupIDs := oldState.groupIDs - if len(account.GroupIDs) > 0 { - groupIDs = mergeGroupIDs(groupIDs, account.GroupIDs) - } - h.clearAccountAffinity(ctx, accountID, groupIDs) -} - -func (h *AccountHandler) clearAffinityCacheForBulkIfDisabled(ctx context.Context, accountIDs []int64, oldStates map[int64]affinityStateSnapshot) { - for _, accountID := range accountIDs { - oldState, ok := oldStates[accountID] - if !ok { - continue - } - h.clearAffinityCacheIfDisabled(ctx, accountID, oldState) - } -} - -// GetAffinityDetails returns the affinity details grouped by user for an account. -// GET /api/v1/admin/accounts/:id/affinity-details -func (h *AccountHandler) GetAffinityDetails(c *gin.Context) { - accountID, err := strconv.ParseInt(c.Param("id"), 10, 64) - if err != nil { - response.BadRequest(c, "Invalid account ID") - return - } - - emptyResp := AffinityDetailsResponse{ - Users: []AffinityDetailUser{}, - TotalUsers: 0, - TotalClients: 0, - PinnedUsers: []int64{}, - } - - account, err := h.adminService.GetAccount(c.Request.Context(), accountID) - if err != nil { - response.ErrorFrom(c, err) - return - } - - if !account.IsAffinityEnabled() { - response.Success(c, emptyResp) - return - } - - pinnedUsers := account.GetPinnedUsers() - if pinnedUsers == nil { - pinnedUsers = []int64{} - } - - if h.gatewayCache == nil || len(account.GroupIDs) == 0 { - emptyResp.PinnedUsers = pinnedUsers - response.Success(c, emptyResp) - return - } - - clients, err := h.gatewayCache.GetAccountAffinityClientsWithScores( - c.Request.Context(), accountID, account.GroupIDs, service.ClientAffinityTTL, - ) - if err != nil { - emptyResp.PinnedUsers = pinnedUsers - response.Success(c, emptyResp) - return - } - - pinnedSet := make(map[int64]struct{}, len(pinnedUsers)) - for _, uid := range pinnedUsers { - pinnedSet[uid] = struct{}{} - } - - // 按 UserID 分组 - userMap := make(map[int64]*AffinityDetailUser) - var userOrder []int64 - for _, cl := range clients { - u, ok := userMap[cl.UserID] - if !ok { - _, pinned := pinnedSet[cl.UserID] - u = &AffinityDetailUser{ - UserID: cl.UserID, - UserEmail: "", - ClientCount: 0, - IsPinned: pinned, - Clients: []AffinityDetailClient{}, - } - userMap[cl.UserID] = u - userOrder = append(userOrder, cl.UserID) - } - u.Clients = append(u.Clients, AffinityDetailClient{ - ClientID: cl.ClientID, - LastActive: cl.LastActive, - IsPinned: false, // 客户端级别暂无 pinned 概念 - }) - } - - // 关联用户邮箱（查询失败时保持空字符串，不影响主流程） - for _, uid := range userOrder { - user, uErr := h.adminService.GetUser(c.Request.Context(), uid) - if uErr != nil || user == nil { - continue - } - if u, ok := userMap[uid]; ok { - if user.Email != "" { - u.UserEmail = user.Email - } else { - u.UserEmail = user.Username - } - } - } - - users := make([]AffinityDetailUser, 0, len(userOrder)) - for _, uid := range userOrder { - u := userMap[uid] - u.ClientCount = len(u.Clients) - users = append(users, *u) - } - - response.Success(c, AffinityDetailsResponse{ - Users: users, - TotalUsers: len(users), - TotalClients: len(clients), - PinnedUsers: pinnedUsers, - }) -} diff --git a/backend/internal/handler/admin/account_data_handler_test.go b/backend/internal/handler/admin/account_data_handler_test.go index e516b275f2..285033a17d 100644 --- a/backend/internal/handler/admin/account_data_handler_test.go +++ b/backend/internal/handler/admin/account_data_handler_test.go @@ -65,7 +65,6 @@ func setupAccountDataRouter() (*gin.Engine, *stubAdminService) { nil, nil, nil, - nil, ) router.GET("/api/v1/admin/accounts/data", h.ExportData) diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 508eb00a9e..e1cc062db4 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -58,7 +58,6 @@ type AccountHandler struct { sessionLimitCache service.SessionLimitCache rpmCache service.RPMCache tokenCacheInvalidator service.TokenCacheInvalidator - gatewayCache service.GatewayCache } // NewAccountHandler creates a new admin account handler @@ -76,7 +75,6 @@ func NewAccountHandler( sessionLimitCache service.SessionLimitCache, rpmCache service.RPMCache, tokenCacheInvalidator service.TokenCacheInvalidator, - gatewayCache service.GatewayCache, ) *AccountHandler { return &AccountHandler{ adminService: adminService, @@ -92,7 +90,6 @@ func NewAccountHandler( sessionLimitCache: sessionLimitCache, rpmCache: rpmCache, tokenCacheInvalidator: tokenCacheInvalidator, - gatewayCache: gatewayCache, } } @@ -212,34 +209,6 @@ func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, ac } } - // 亲和客户端数据（启用亲和的账号始终返回 count，即使为 0） - if account.IsAffinityEnabled() { - if h.gatewayCache != nil && len(account.GroupIDs) > 0 { - accountGroups := map[int64][]int64{account.ID: account.GroupIDs} - if clients, err := h.gatewayCache.GetAccountAffinityClientsBatch(ctx, accountGroups, service.ClientAffinityTTL); err == nil { - if cl, ok := clients[account.ID]; ok && len(cl) > 0 { - count := int64(len(cl)) - item.AffinityClientCount = &count - item.AffinityClients = cl - userCount := countUniqueUsersFromAffinityMembers(cl) - item.AffinityUserCount = &userCount - } else { - zero := int64(0) - item.AffinityClientCount = &zero - item.AffinityUserCount = &zero - } - } else { - zero := int64(0) - item.AffinityClientCount = &zero - item.AffinityUserCount = &zero - } - } else { - zero := int64(0) - item.AffinityClientCount = &zero - item.AffinityUserCount = &zero - } - } - return item } @@ -366,21 +335,6 @@ func (h *AccountHandler) List(c *gin.Context) { _ = g.Wait() } - // 获取亲和客户端数据（Redis Pipeline，低开销） - var affinityClients map[int64][]string - if h.gatewayCache != nil { - accountGroups := make(map[int64][]int64) - for i := range accounts { - acc := &accounts[i] - if acc.IsAffinityEnabled() && len(acc.GroupIDs) > 0 { - accountGroups[acc.ID] = acc.GroupIDs - } - } - if len(accountGroups) > 0 { - affinityClients, _ = h.gatewayCache.GetAccountAffinityClientsBatch(c.Request.Context(), accountGroups, service.ClientAffinityTTL) - } - } - // Build response with concurrency info result := make([]AccountWithConcurrency, len(accounts)) for i := range accounts { @@ -411,22 +365,6 @@ func (h *AccountHandler) List(c *gin.Context) { } } - // 注入亲和客户端数据到 DTO（启用亲和的账号始终返回 count，即使为 0） - if acc.IsAffinityEnabled() { - if clients, ok := affinityClients[acc.ID]; ok && len(clients) > 0 { - count := int64(len(clients)) - item.AffinityClientCount = &count - item.AffinityClients = clients - // 从成员列表中解析唯一用户数 - userCount := countUniqueUsersFromAffinityMembers(clients) - item.AffinityUserCount = &userCount - } else { - zero := int64(0) - item.AffinityClientCount = &zero - item.AffinityUserCount = &zero - } - } - result[i] = item } @@ -651,8 +589,6 @@ func (h *AccountHandler) Update(c *gin.Context) { // base_rpm 输入校验：负值归零，超过 10000 截断 sanitizeExtraBaseRPM(req.Extra) - oldStates := h.captureAffinityStates(c.Request.Context(), []int64{accountID}) - // 确定是否跳过混合渠道检查 skipCheck := req.ConfirmMixedChannelRisk != nil && *req.ConfirmMixedChannelRisk @@ -689,8 +625,6 @@ func (h *AccountHandler) Update(c *gin.Context) { return } - h.clearAffinityCacheForBulkIfDisabled(c.Request.Context(), []int64{accountID}, oldStates) - response.Success(c, h.buildAccountResponseWithRuntime(c.Request.Context(), account)) } @@ -1454,8 +1388,6 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { return } - oldStates := h.captureAffinityStates(c.Request.Context(), req.AccountIDs) - result, err := h.adminService.BulkUpdateAccounts(c.Request.Context(), &service.BulkUpdateAccountsInput{ AccountIDs: req.AccountIDs, Name: req.Name, @@ -1490,8 +1422,6 @@ func (h *AccountHandler) BulkUpdate(c *gin.Context) { return } - h.clearAffinityCacheForBulkIfDisabled(c.Request.Context(), req.AccountIDs, oldStates) - response.Success(c, result) } diff --git a/backend/internal/handler/admin/account_handler_affinity_cache_test.go b/backend/internal/handler/admin/account_handler_affinity_cache_test.go deleted file mode 100644 index 3c22bb543a..0000000000 --- a/backend/internal/handler/admin/account_handler_affinity_cache_test.go +++ /dev/null @@ -1,298 +0,0 @@ -package admin - -import ( - "bytes" - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "testing" - "time" - - "github.com/Wei-Shaw/sub2api/internal/service" - "github.com/gin-gonic/gin" - "github.com/stretchr/testify/require" -) - -type affinityCacheAdminService struct { - *stubAdminService - accounts map[int64]*service.Account - bulkResultIDs []int64 - bulkFailedCount int -} - -func (s *affinityCacheAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) { - if acc, ok := s.accounts[id]; ok { - cloned := *acc - return &cloned, nil - } - return s.stubAdminService.GetAccount(context.Background(), id) -} - -func (s *affinityCacheAdminService) UpdateAccount(_ context.Context, id int64, input *service.UpdateAccountInput) (*service.Account, error) { - acc := s.accounts[id] - if acc == nil { - return s.stubAdminService.UpdateAccount(context.Background(), id, input) - } - if input.Extra != nil { - acc.Extra = input.Extra - } - if input.GroupIDs != nil { - acc.GroupIDs = *input.GroupIDs - } - cloned := *acc - return &cloned, nil -} - -func (s *affinityCacheAdminService) BulkUpdateAccounts(_ context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) { - for _, accountID := range input.AccountIDs { - acc := s.accounts[accountID] - if acc == nil { - continue - } - if input.Extra != nil { - acc.Extra = input.Extra - } - if input.GroupIDs != nil { - acc.GroupIDs = *input.GroupIDs - } - } - successIDs := append([]int64(nil), input.AccountIDs...) - failed := 0 - if len(s.bulkResultIDs) > 0 { - successIDs = append([]int64(nil), s.bulkResultIDs...) - failed = s.bulkFailedCount - } - return &service.BulkUpdateAccountsResult{ - Success: len(successIDs), - Failed: failed, - SuccessIDs: successIDs, - Results: []service.BulkUpdateAccountResult{}, - }, nil -} - -type affinityCacheClearRecorder struct { - clearCalls []struct { - accountID int64 - groupIDs []int64 - } -} - -func (r *affinityCacheClearRecorder) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) { - return 0, nil -} -func (r *affinityCacheClearRecorder) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (r *affinityCacheClearRecorder) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error { - return nil -} -func (r *affinityCacheClearRecorder) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error { - return nil -} -func (r *affinityCacheClearRecorder) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (r *affinityCacheClearRecorder) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (r *affinityCacheClearRecorder) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (r *affinityCacheClearRecorder) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (r *affinityCacheClearRecorder) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) { - return nil, nil -} -func (r *affinityCacheClearRecorder) ClearAccountAffinity(_ context.Context, accountID int64, groupIDs []int64) error { - r.clearCalls = append(r.clearCalls, struct { - accountID int64 - groupIDs []int64 - }{ - accountID: accountID, - groupIDs: append([]int64(nil), groupIDs...), - }) - return nil -} -func (r *affinityCacheClearRecorder) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (users, clients, perUser int64, err error) { - return 0, 0, 0, nil -} - -func setupAffinityCacheRouter(adminSvc service.AdminService, gatewayCache service.GatewayCache) *gin.Engine { - gin.SetMode(gin.TestMode) - router := gin.New() - handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, gatewayCache) - router.PUT("/api/v1/admin/accounts/:id", handler.Update) - router.POST("/api/v1/admin/accounts/bulk-update", handler.BulkUpdate) - return router -} - -func TestAccountHandlerUpdate_DisablingAffinityClearsCache(t *testing.T) { - adminSvc := &affinityCacheAdminService{ - stubAdminService: newStubAdminService(), - accounts: map[int64]*service.Account{ - 3: { - ID: 3, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeOAuth, - Status: service.StatusActive, - GroupIDs: []int64{1, 2}, - Extra: map[string]any{ - "client_affinity_enabled": true, - }, - }, - }, - } - cache := &affinityCacheClearRecorder{} - router := setupAffinityCacheRouter(adminSvc, cache) - - body, _ := json.Marshal(map[string]any{ - "extra": map[string]any{}, - }) - rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodPut, "/api/v1/admin/accounts/3", bytes.NewReader(body)) - req.Header.Set("Content-Type", "application/json") - router.ServeHTTP(rec, req) - - require.Equal(t, http.StatusOK, rec.Code) - require.Len(t, cache.clearCalls, 1) - require.Equal(t, int64(3), cache.clearCalls[0].accountID) - require.Equal(t, []int64{1, 2}, cache.clearCalls[0].groupIDs) -} - -func TestAccountHandlerBulkUpdate_DisablingAffinityClearsCache(t *testing.T) { - adminSvc := &affinityCacheAdminService{ - stubAdminService: newStubAdminService(), - accounts: map[int64]*service.Account{ - 1: { - ID: 1, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeOAuth, - Status: service.StatusActive, - GroupIDs: []int64{11}, - Extra: map[string]any{ - "client_affinity_enabled": true, - }, - }, - 2: { - ID: 2, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeSetupToken, - Status: service.StatusActive, - GroupIDs: []int64{22}, - Extra: map[string]any{ - "client_affinity_enabled": true, - }, - }, - }, - } - cache := &affinityCacheClearRecorder{} - router := setupAffinityCacheRouter(adminSvc, cache) - - body, _ := json.Marshal(map[string]any{ - "account_ids": []int64{1, 2}, - "extra": map[string]any{ - "client_affinity_enabled": false, - }, - }) - rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body)) - req.Header.Set("Content-Type", "application/json") - router.ServeHTTP(rec, req) - - require.Equal(t, http.StatusOK, rec.Code) - require.Len(t, cache.clearCalls, 2) - require.Equal(t, int64(1), cache.clearCalls[0].accountID) - require.Equal(t, []int64{11}, cache.clearCalls[0].groupIDs) - require.Equal(t, int64(2), cache.clearCalls[1].accountID) - require.Equal(t, []int64{22}, cache.clearCalls[1].groupIDs) -} - -func TestAccountHandlerBulkUpdate_DisablingAffinityWithNewFlagClearsCache(t *testing.T) { - adminSvc := &affinityCacheAdminService{ - stubAdminService: newStubAdminService(), - accounts: map[int64]*service.Account{ - 1: { - ID: 1, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeOAuth, - Status: service.StatusActive, - GroupIDs: []int64{33}, - Extra: map[string]any{ - "affinity_enabled": true, - }, - }, - }, - } - cache := &affinityCacheClearRecorder{} - router := setupAffinityCacheRouter(adminSvc, cache) - - body, _ := json.Marshal(map[string]any{ - "account_ids": []int64{1}, - "extra": map[string]any{ - "affinity_enabled": false, - }, - }) - rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body)) - req.Header.Set("Content-Type", "application/json") - router.ServeHTTP(rec, req) - - require.Equal(t, http.StatusOK, rec.Code) - require.Len(t, cache.clearCalls, 1) - require.Equal(t, int64(1), cache.clearCalls[0].accountID) - require.Equal(t, []int64{33}, cache.clearCalls[0].groupIDs) -} - -func TestAccountHandlerBulkUpdate_PartialFailureStillClearsDisabledAffinityCache(t *testing.T) { - adminSvc := &affinityCacheAdminService{ - stubAdminService: newStubAdminService(), - accounts: map[int64]*service.Account{ - 1: { - ID: 1, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeOAuth, - Status: service.StatusActive, - GroupIDs: []int64{44}, - Extra: map[string]any{ - "affinity_enabled": true, - }, - }, - 2: { - ID: 2, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeSetupToken, - Status: service.StatusActive, - GroupIDs: []int64{55}, - Extra: map[string]any{ - "client_affinity_enabled": true, - }, - }, - }, - bulkResultIDs: []int64{1}, - bulkFailedCount: 1, - } - cache := &affinityCacheClearRecorder{} - router := setupAffinityCacheRouter(adminSvc, cache) - - body, _ := json.Marshal(map[string]any{ - "account_ids": []int64{1, 2}, - "group_ids": []int64{101}, - "extra": map[string]any{ - "affinity_enabled": false, - "client_affinity_enabled": false, - }, - }) - rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/bulk-update", bytes.NewReader(body)) - req.Header.Set("Content-Type", "application/json") - router.ServeHTTP(rec, req) - - require.Equal(t, http.StatusOK, rec.Code) - require.Len(t, cache.clearCalls, 2) - require.Equal(t, int64(1), cache.clearCalls[0].accountID) - require.Equal(t, []int64{44, 101}, cache.clearCalls[0].groupIDs) - require.Equal(t, int64(2), cache.clearCalls[1].accountID) - require.Equal(t, []int64{55, 101}, cache.clearCalls[1].groupIDs) -} diff --git a/backend/internal/handler/admin/account_handler_affinity_details_test.go b/backend/internal/handler/admin/account_handler_affinity_details_test.go deleted file mode 100644 index 0a4b52d579..0000000000 --- a/backend/internal/handler/admin/account_handler_affinity_details_test.go +++ /dev/null @@ -1,205 +0,0 @@ -package admin - -import ( - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "testing" - "time" - - "github.com/Wei-Shaw/sub2api/internal/service" - "github.com/gin-gonic/gin" - "github.com/stretchr/testify/require" -) - -type affinityDetailsAdminService struct { - *stubAdminService - account service.Account - usersByID map[int64]service.User -} - -func (s *affinityDetailsAdminService) GetAccount(_ context.Context, id int64) (*service.Account, error) { - if s.account.ID == id { - acc := s.account - return &acc, nil - } - return s.stubAdminService.GetAccount(context.Background(), id) -} - -func (s *affinityDetailsAdminService) GetUser(_ context.Context, id int64) (*service.User, error) { - if u, ok := s.usersByID[id]; ok { - user := u - return &user, nil - } - return s.stubAdminService.GetUser(context.Background(), id) -} - -type affinityDetailsGatewayCacheStub struct { - clients []service.AffinityClient -} - -func (s *affinityDetailsGatewayCacheStub) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) { - return 0, nil -} -func (s *affinityDetailsGatewayCacheStub) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (s *affinityDetailsGatewayCacheStub) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error { - return nil -} -func (s *affinityDetailsGatewayCacheStub) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error { - return nil -} -func (s *affinityDetailsGatewayCacheStub) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (s *affinityDetailsGatewayCacheStub) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (s *affinityDetailsGatewayCacheStub) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) { - return s.clients, nil -} -func (s *affinityDetailsGatewayCacheStub) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} -func (s *affinityDetailsGatewayCacheStub) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (users, clients, perUser int64, err error) { - return 0, 0, 0, nil -} - -func setupAffinityDetailsRouter(adminSvc service.AdminService, gatewayCache service.GatewayCache) *gin.Engine { - gin.SetMode(gin.TestMode) - router := gin.New() - handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, gatewayCache) - router.GET("/api/v1/admin/accounts/:id/affinity-details", handler.GetAffinityDetails) - return router -} - -func TestAccountHandlerGetAffinityDetails_ContractFields(t *testing.T) { - adminSvc := &affinityDetailsAdminService{ - stubAdminService: newStubAdminService(), - account: service.Account{ - ID: 88, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeOAuth, - Status: service.StatusActive, - GroupIDs: []int64{1}, - Extra: map[string]any{ - "client_affinity_enabled": true, - "pinned_users": []any{float64(42), float64(99)}, - }, - }, - usersByID: map[int64]service.User{ - 42: {ID: 42, Email: "user42@example.com"}, - 7: {ID: 7, Username: "user7"}, - }, - } - cache := &affinityDetailsGatewayCacheStub{ - clients: []service.AffinityClient{ - {UserID: 42, ClientID: "client-a", LastActive: time.Now().Add(-time.Minute)}, - {UserID: 7, ClientID: "client-b", LastActive: time.Now().Add(-2 * time.Minute)}, - {UserID: 42, ClientID: "client-c", LastActive: time.Now().Add(-3 * time.Minute)}, - }, - } - router := setupAffinityDetailsRouter(adminSvc, cache) - - rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/88/affinity-details", nil) - router.ServeHTTP(rec, req) - require.Equal(t, http.StatusOK, rec.Code) - - var resp struct { - Code int `json:"code"` - Data struct { - Users []struct { - UserID int64 `json:"user_id"` - UserEmail string `json:"user_email"` - ClientCount int `json:"client_count"` - IsPinned bool `json:"is_pinned"` - Clients []struct { - ClientID string `json:"client_id"` - } `json:"clients"` - } `json:"users"` - TotalUsers int `json:"total_users"` - TotalClients int `json:"total_clients"` - PinnedUsers []int64 `json:"pinned_users"` - } `json:"data"` - } - require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) - require.Equal(t, 0, resp.Code) - require.Equal(t, 2, resp.Data.TotalUsers) - require.Equal(t, 3, resp.Data.TotalClients) - require.Equal(t, []int64{42, 99}, resp.Data.PinnedUsers) - require.Len(t, resp.Data.Users, 2) - - usersByID := make(map[int64]struct { - UserEmail string - ClientCount int - IsPinned bool - ClientLen int - }, len(resp.Data.Users)) - for _, u := range resp.Data.Users { - usersByID[u.UserID] = struct { - UserEmail string - ClientCount int - IsPinned bool - ClientLen int - }{ - UserEmail: u.UserEmail, - ClientCount: u.ClientCount, - IsPinned: u.IsPinned, - ClientLen: len(u.Clients), - } - } - require.Equal(t, "user42@example.com", usersByID[42].UserEmail) - require.Equal(t, 2, usersByID[42].ClientCount) - require.True(t, usersByID[42].IsPinned) - require.Equal(t, 2, usersByID[42].ClientLen) - - require.Equal(t, "user7", usersByID[7].UserEmail) - require.Equal(t, 1, usersByID[7].ClientCount) - require.False(t, usersByID[7].IsPinned) - require.Equal(t, 1, usersByID[7].ClientLen) -} - -func TestAccountHandlerGetAffinityDetails_DisabledReturnsEmptyContract(t *testing.T) { - adminSvc := &affinityDetailsAdminService{ - stubAdminService: newStubAdminService(), - account: service.Account{ - ID: 89, - Platform: service.PlatformAnthropic, - Type: service.AccountTypeOAuth, - Status: service.StatusActive, - GroupIDs: []int64{1}, - }, - usersByID: map[int64]service.User{}, - } - router := setupAffinityDetailsRouter(adminSvc, &affinityDetailsGatewayCacheStub{}) - - rec := httptest.NewRecorder() - req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/89/affinity-details", nil) - router.ServeHTTP(rec, req) - require.Equal(t, http.StatusOK, rec.Code) - - var resp struct { - Code int `json:"code"` - Data struct { - Users []any `json:"users"` - TotalUsers int `json:"total_users"` - TotalClients int `json:"total_clients"` - PinnedUsers []int64 `json:"pinned_users"` - } `json:"data"` - } - require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) - require.Equal(t, 0, resp.Code) - require.Empty(t, resp.Data.Users) - require.Equal(t, 0, resp.Data.TotalUsers) - require.Equal(t, 0, resp.Data.TotalClients) - require.Empty(t, resp.Data.PinnedUsers) -} diff --git a/backend/internal/handler/admin/account_handler_available_models_test.go b/backend/internal/handler/admin/account_handler_available_models_test.go index 79f9689a92..c5f1e2d884 100644 --- a/backend/internal/handler/admin/account_handler_available_models_test.go +++ b/backend/internal/handler/admin/account_handler_available_models_test.go @@ -28,7 +28,7 @@ func (s *availableModelsAdminService) GetAccount(_ context.Context, id int64) (* func setupAvailableModelsRouter(adminSvc service.AdminService) *gin.Engine { gin.SetMode(gin.TestMode) router := gin.New() - handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) router.GET("/api/v1/admin/accounts/:id/models", handler.GetAvailableModels) return router } diff --git a/backend/internal/handler/admin/account_handler_mixed_channel_test.go b/backend/internal/handler/admin/account_handler_mixed_channel_test.go index 5e45619064..5b81db2aad 100644 --- a/backend/internal/handler/admin/account_handler_mixed_channel_test.go +++ b/backend/internal/handler/admin/account_handler_mixed_channel_test.go @@ -15,7 +15,7 @@ import ( func setupAccountMixedChannelRouter(adminSvc *stubAdminService) *gin.Engine { gin.SetMode(gin.TestMode) router := gin.New() - accountHandler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + accountHandler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) router.POST("/api/v1/admin/accounts/check-mixed-channel", accountHandler.CheckMixedChannel) router.POST("/api/v1/admin/accounts", accountHandler.Create) router.PUT("/api/v1/admin/accounts/:id", accountHandler.Update) diff --git a/backend/internal/handler/admin/account_handler_passthrough_test.go b/backend/internal/handler/admin/account_handler_passthrough_test.go index a11f2c5195..d86501c047 100644 --- a/backend/internal/handler/admin/account_handler_passthrough_test.go +++ b/backend/internal/handler/admin/account_handler_passthrough_test.go @@ -29,7 +29,6 @@ func TestAccountHandler_Create_AnthropicAPIKeyPassthroughExtraForwarded(t *testi nil, nil, nil, - nil, ) router := gin.New() diff --git a/backend/internal/handler/admin/batch_update_credentials_test.go b/backend/internal/handler/admin/batch_update_credentials_test.go index f2cd1e3a1c..0b1b669174 100644 --- a/backend/internal/handler/admin/batch_update_credentials_test.go +++ b/backend/internal/handler/admin/batch_update_credentials_test.go @@ -36,7 +36,7 @@ func (f *failingAdminService) UpdateAccount(ctx context.Context, id int64, input func setupAccountHandlerWithService(adminSvc service.AdminService) (*gin.Engine, *AccountHandler) { gin.SetMode(gin.TestMode) router := gin.New() - handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) + handler := NewAccountHandler(adminSvc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil) router.POST("/api/v1/admin/accounts/batch-update-credentials", handler.BatchUpdateCredentials) return router, handler } diff --git a/backend/internal/handler/admin/group_handler.go b/backend/internal/handler/admin/group_handler.go index c7ef313541..81a103c795 100644 --- a/backend/internal/handler/admin/group_handler.go +++ b/backend/internal/handler/admin/group_handler.go @@ -105,8 +105,7 @@ type CreateGroupRequest struct { // 模型路由配置（仅 anthropic 平台使用） ModelRouting map[string][]int64 `json:"model_routing"` ModelRoutingEnabled bool `json:"model_routing_enabled"` - MCPXMLInject *bool `json:"mcp_xml_inject"` - SimulateClaudeMaxEnabled *bool `json:"simulate_claude_max_enabled"` + MCPXMLInject *bool `json:"mcp_xml_inject"` // 支持的模型系列（仅 antigravity 平台使用） SupportedModelScopes []string `json:"supported_model_scopes"` // Sora 存储配额 @@ -144,8 +143,7 @@ type UpdateGroupRequest struct { // 模型路由配置（仅 anthropic 平台使用） ModelRouting map[string][]int64 `json:"model_routing"` ModelRoutingEnabled *bool `json:"model_routing_enabled"` - MCPXMLInject *bool `json:"mcp_xml_inject"` - SimulateClaudeMaxEnabled *bool `json:"simulate_claude_max_enabled"` + MCPXMLInject *bool `json:"mcp_xml_inject"` // 支持的模型系列（仅 antigravity 平台使用） SupportedModelScopes *[]string `json:"supported_model_scopes"` // Sora 存储配额 @@ -266,7 +264,6 @@ func (h *GroupHandler) Create(c *gin.Context) { ModelRouting: req.ModelRouting, ModelRoutingEnabled: req.ModelRoutingEnabled, MCPXMLInject: req.MCPXMLInject, - SimulateClaudeMaxEnabled: req.SimulateClaudeMaxEnabled, SupportedModelScopes: req.SupportedModelScopes, SoraStorageQuotaBytes: req.SoraStorageQuotaBytes, AllowMessagesDispatch: req.AllowMessagesDispatch, @@ -320,7 +317,6 @@ func (h *GroupHandler) Update(c *gin.Context) { ModelRouting: req.ModelRouting, ModelRoutingEnabled: req.ModelRoutingEnabled, MCPXMLInject: req.MCPXMLInject, - SimulateClaudeMaxEnabled: req.SimulateClaudeMaxEnabled, SupportedModelScopes: req.SupportedModelScopes, SoraStorageQuotaBytes: req.SoraStorageQuotaBytes, AllowMessagesDispatch: req.AllowMessagesDispatch, diff --git a/backend/internal/handler/admin/setting_handler.go b/backend/internal/handler/admin/setting_handler.go index 580977617e..2a54a187d5 100644 --- a/backend/internal/handler/admin/setting_handler.go +++ b/backend/internal/handler/admin/setting_handler.go @@ -41,17 +41,15 @@ type SettingHandler struct { emailService *service.EmailService turnstileService *service.TurnstileService opsService *service.OpsService - soraS3Storage *service.SoraS3Storage } // NewSettingHandler 创建系统设置处理器 -func NewSettingHandler(settingService *service.SettingService, emailService *service.EmailService, turnstileService *service.TurnstileService, opsService *service.OpsService, soraS3Storage *service.SoraS3Storage) *SettingHandler { +func NewSettingHandler(settingService *service.SettingService, emailService *service.EmailService, turnstileService *service.TurnstileService, opsService *service.OpsService) *SettingHandler { return &SettingHandler{ settingService: settingService, emailService: emailService, turnstileService: turnstileService, opsService: opsService, - soraS3Storage: soraS3Storage, } } @@ -1207,386 +1205,6 @@ func (h *SettingHandler) GetStreamTimeoutSettings(c *gin.Context) { }) } -func toSoraS3SettingsDTO(settings *service.SoraS3Settings) dto.SoraS3Settings { - if settings == nil { - return dto.SoraS3Settings{} - } - return dto.SoraS3Settings{ - Enabled: settings.Enabled, - Endpoint: settings.Endpoint, - Region: settings.Region, - Bucket: settings.Bucket, - AccessKeyID: settings.AccessKeyID, - SecretAccessKeyConfigured: settings.SecretAccessKeyConfigured, - Prefix: settings.Prefix, - ForcePathStyle: settings.ForcePathStyle, - CDNURL: settings.CDNURL, - DefaultStorageQuotaBytes: settings.DefaultStorageQuotaBytes, - } -} - -func toSoraS3ProfileDTO(profile service.SoraS3Profile) dto.SoraS3Profile { - return dto.SoraS3Profile{ - ProfileID: profile.ProfileID, - Name: profile.Name, - IsActive: profile.IsActive, - Enabled: profile.Enabled, - Endpoint: profile.Endpoint, - Region: profile.Region, - Bucket: profile.Bucket, - AccessKeyID: profile.AccessKeyID, - SecretAccessKeyConfigured: profile.SecretAccessKeyConfigured, - Prefix: profile.Prefix, - ForcePathStyle: profile.ForcePathStyle, - CDNURL: profile.CDNURL, - DefaultStorageQuotaBytes: profile.DefaultStorageQuotaBytes, - UpdatedAt: profile.UpdatedAt, - } -} - -func validateSoraS3RequiredWhenEnabled(enabled bool, endpoint, bucket, accessKeyID, secretAccessKey string, hasStoredSecret bool) error { - if !enabled { - return nil - } - if strings.TrimSpace(endpoint) == "" { - return fmt.Errorf("S3 Endpoint is required when enabled") - } - if strings.TrimSpace(bucket) == "" { - return fmt.Errorf("S3 Bucket is required when enabled") - } - if strings.TrimSpace(accessKeyID) == "" { - return fmt.Errorf("S3 Access Key ID is required when enabled") - } - if strings.TrimSpace(secretAccessKey) != "" || hasStoredSecret { - return nil - } - return fmt.Errorf("S3 Secret Access Key is required when enabled") -} - -func findSoraS3ProfileByID(items []service.SoraS3Profile, profileID string) *service.SoraS3Profile { - for idx := range items { - if items[idx].ProfileID == profileID { - return &items[idx] - } - } - return nil -} - -// GetSoraS3Settings 获取 Sora S3 存储配置（兼容旧单配置接口） -// GET /api/v1/admin/settings/sora-s3 -func (h *SettingHandler) GetSoraS3Settings(c *gin.Context) { - settings, err := h.settingService.GetSoraS3Settings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, toSoraS3SettingsDTO(settings)) -} - -// ListSoraS3Profiles 获取 Sora S3 多配置 -// GET /api/v1/admin/settings/sora-s3/profiles -func (h *SettingHandler) ListSoraS3Profiles(c *gin.Context) { - result, err := h.settingService.ListSoraS3Profiles(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - items := make([]dto.SoraS3Profile, 0, len(result.Items)) - for idx := range result.Items { - items = append(items, toSoraS3ProfileDTO(result.Items[idx])) - } - response.Success(c, dto.ListSoraS3ProfilesResponse{ - ActiveProfileID: result.ActiveProfileID, - Items: items, - }) -} - -// UpdateSoraS3SettingsRequest 更新/测试 Sora S3 配置请求（兼容旧接口） -type UpdateSoraS3SettingsRequest struct { - ProfileID string `json:"profile_id"` - Enabled bool `json:"enabled"` - Endpoint string `json:"endpoint"` - Region string `json:"region"` - Bucket string `json:"bucket"` - AccessKeyID string `json:"access_key_id"` - SecretAccessKey string `json:"secret_access_key"` - Prefix string `json:"prefix"` - ForcePathStyle bool `json:"force_path_style"` - CDNURL string `json:"cdn_url"` - DefaultStorageQuotaBytes int64 `json:"default_storage_quota_bytes"` -} - -type CreateSoraS3ProfileRequest struct { - ProfileID string `json:"profile_id"` - Name string `json:"name"` - SetActive bool `json:"set_active"` - Enabled bool `json:"enabled"` - Endpoint string `json:"endpoint"` - Region string `json:"region"` - Bucket string `json:"bucket"` - AccessKeyID string `json:"access_key_id"` - SecretAccessKey string `json:"secret_access_key"` - Prefix string `json:"prefix"` - ForcePathStyle bool `json:"force_path_style"` - CDNURL string `json:"cdn_url"` - DefaultStorageQuotaBytes int64 `json:"default_storage_quota_bytes"` -} - -type UpdateSoraS3ProfileRequest struct { - Name string `json:"name"` - Enabled bool `json:"enabled"` - Endpoint string `json:"endpoint"` - Region string `json:"region"` - Bucket string `json:"bucket"` - AccessKeyID string `json:"access_key_id"` - SecretAccessKey string `json:"secret_access_key"` - Prefix string `json:"prefix"` - ForcePathStyle bool `json:"force_path_style"` - CDNURL string `json:"cdn_url"` - DefaultStorageQuotaBytes int64 `json:"default_storage_quota_bytes"` -} - -// CreateSoraS3Profile 创建 Sora S3 配置 -// POST /api/v1/admin/settings/sora-s3/profiles -func (h *SettingHandler) CreateSoraS3Profile(c *gin.Context) { - var req CreateSoraS3ProfileRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - if req.DefaultStorageQuotaBytes < 0 { - req.DefaultStorageQuotaBytes = 0 - } - if strings.TrimSpace(req.Name) == "" { - response.BadRequest(c, "Name is required") - return - } - if strings.TrimSpace(req.ProfileID) == "" { - response.BadRequest(c, "Profile ID is required") - return - } - // S3 字段验证 - if err := validateSoraS3RequiredWhenEnabled(req.Enabled, req.Endpoint, req.Bucket, req.AccessKeyID, req.SecretAccessKey, false); err != nil { - response.BadRequest(c, err.Error()) - return - } - - created, err := h.settingService.CreateSoraS3Profile(c.Request.Context(), &service.SoraS3Profile{ - ProfileID: req.ProfileID, - Name: req.Name, - Enabled: req.Enabled, - Endpoint: req.Endpoint, - Region: req.Region, - Bucket: req.Bucket, - AccessKeyID: req.AccessKeyID, - SecretAccessKey: req.SecretAccessKey, - Prefix: req.Prefix, - ForcePathStyle: req.ForcePathStyle, - CDNURL: req.CDNURL, - DefaultStorageQuotaBytes: req.DefaultStorageQuotaBytes, - }, req.SetActive) - if err != nil { - response.ErrorFrom(c, err) - return - } - - response.Success(c, toSoraS3ProfileDTO(*created)) -} - -// UpdateSoraS3Profile 更新 Sora S3 配置 -// PUT /api/v1/admin/settings/sora-s3/profiles/:profile_id -func (h *SettingHandler) UpdateSoraS3Profile(c *gin.Context) { - profileID := strings.TrimSpace(c.Param("profile_id")) - if profileID == "" { - response.BadRequest(c, "Profile ID is required") - return - } - - var req UpdateSoraS3ProfileRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - if req.DefaultStorageQuotaBytes < 0 { - req.DefaultStorageQuotaBytes = 0 - } - if strings.TrimSpace(req.Name) == "" { - response.BadRequest(c, "Name is required") - return - } - - existingList, err := h.settingService.ListSoraS3Profiles(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - existing := findSoraS3ProfileByID(existingList.Items, profileID) - if existing == nil { - response.ErrorFrom(c, service.ErrSoraS3ProfileNotFound) - return - } - // S3 字段验证 - if err := validateSoraS3RequiredWhenEnabled(req.Enabled, req.Endpoint, req.Bucket, req.AccessKeyID, req.SecretAccessKey, existing.SecretAccessKeyConfigured); err != nil { - response.BadRequest(c, err.Error()) - return - } - - updated, updateErr := h.settingService.UpdateSoraS3Profile(c.Request.Context(), profileID, &service.SoraS3Profile{ - Name: req.Name, - Enabled: req.Enabled, - Endpoint: req.Endpoint, - Region: req.Region, - Bucket: req.Bucket, - AccessKeyID: req.AccessKeyID, - SecretAccessKey: req.SecretAccessKey, - Prefix: req.Prefix, - ForcePathStyle: req.ForcePathStyle, - CDNURL: req.CDNURL, - DefaultStorageQuotaBytes: req.DefaultStorageQuotaBytes, - }) - if updateErr != nil { - response.ErrorFrom(c, updateErr) - return - } - - response.Success(c, toSoraS3ProfileDTO(*updated)) -} - -// DeleteSoraS3Profile 删除 Sora S3 配置 -// DELETE /api/v1/admin/settings/sora-s3/profiles/:profile_id -func (h *SettingHandler) DeleteSoraS3Profile(c *gin.Context) { - profileID := strings.TrimSpace(c.Param("profile_id")) - if profileID == "" { - response.BadRequest(c, "Profile ID is required") - return - } - if err := h.settingService.DeleteSoraS3Profile(c.Request.Context(), profileID); err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, gin.H{"deleted": true}) -} - -// SetActiveSoraS3Profile 切换激活 Sora S3 配置 -// POST /api/v1/admin/settings/sora-s3/profiles/:profile_id/activate -func (h *SettingHandler) SetActiveSoraS3Profile(c *gin.Context) { - profileID := strings.TrimSpace(c.Param("profile_id")) - if profileID == "" { - response.BadRequest(c, "Profile ID is required") - return - } - active, err := h.settingService.SetActiveSoraS3Profile(c.Request.Context(), profileID) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, toSoraS3ProfileDTO(*active)) -} - -// UpdateSoraS3Settings 更新 Sora S3 存储配置（兼容旧单配置接口） -// PUT /api/v1/admin/settings/sora-s3 -func (h *SettingHandler) UpdateSoraS3Settings(c *gin.Context) { - var req UpdateSoraS3SettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - - existing, err := h.settingService.GetSoraS3Settings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - - if req.DefaultStorageQuotaBytes < 0 { - req.DefaultStorageQuotaBytes = 0 - } - if err := validateSoraS3RequiredWhenEnabled(req.Enabled, req.Endpoint, req.Bucket, req.AccessKeyID, req.SecretAccessKey, existing.SecretAccessKeyConfigured); err != nil { - response.BadRequest(c, err.Error()) - return - } - - settings := &service.SoraS3Settings{ - Enabled: req.Enabled, - Endpoint: req.Endpoint, - Region: req.Region, - Bucket: req.Bucket, - AccessKeyID: req.AccessKeyID, - SecretAccessKey: req.SecretAccessKey, - Prefix: req.Prefix, - ForcePathStyle: req.ForcePathStyle, - CDNURL: req.CDNURL, - DefaultStorageQuotaBytes: req.DefaultStorageQuotaBytes, - } - if err := h.settingService.SetSoraS3Settings(c.Request.Context(), settings); err != nil { - response.ErrorFrom(c, err) - return - } - - updatedSettings, err := h.settingService.GetSoraS3Settings(c.Request.Context()) - if err != nil { - response.ErrorFrom(c, err) - return - } - response.Success(c, toSoraS3SettingsDTO(updatedSettings)) -} - -// TestSoraS3Connection 测试 Sora S3 连接（HeadBucket） -// POST /api/v1/admin/settings/sora-s3/test -func (h *SettingHandler) TestSoraS3Connection(c *gin.Context) { - if h.soraS3Storage == nil { - response.Error(c, 500, "S3 存储服务未初始化") - return - } - - var req UpdateSoraS3SettingsRequest - if err := c.ShouldBindJSON(&req); err != nil { - response.BadRequest(c, "Invalid request: "+err.Error()) - return - } - if !req.Enabled { - response.BadRequest(c, "S3 未启用，无法测试连接") - return - } - - if req.SecretAccessKey == "" { - if req.ProfileID != "" { - profiles, err := h.settingService.ListSoraS3Profiles(c.Request.Context()) - if err == nil { - profile := findSoraS3ProfileByID(profiles.Items, req.ProfileID) - if profile != nil { - req.SecretAccessKey = profile.SecretAccessKey - } - } - } - if req.SecretAccessKey == "" { - existing, err := h.settingService.GetSoraS3Settings(c.Request.Context()) - if err == nil { - req.SecretAccessKey = existing.SecretAccessKey - } - } - } - - testCfg := &service.SoraS3Settings{ - Enabled: true, - Endpoint: req.Endpoint, - Region: req.Region, - Bucket: req.Bucket, - AccessKeyID: req.AccessKeyID, - SecretAccessKey: req.SecretAccessKey, - Prefix: req.Prefix, - ForcePathStyle: req.ForcePathStyle, - CDNURL: req.CDNURL, - } - if err := h.soraS3Storage.TestConnectionWithSettings(c.Request.Context(), testCfg); err != nil { - response.Error(c, 400, "S3 连接测试失败: "+err.Error()) - return - } - response.Success(c, gin.H{"message": "S3 连接成功"}) -} - // GetRectifierSettings 获取请求整流器配置 // GET /api/v1/admin/settings/rectifier func (h *SettingHandler) GetRectifierSettings(c *gin.Context) { diff --git a/backend/internal/handler/dto/mappers.go b/backend/internal/handler/dto/mappers.go index 45499e2c39..5fc4ae935f 100644 --- a/backend/internal/handler/dto/mappers.go +++ b/backend/internal/handler/dto/mappers.go @@ -140,7 +140,6 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup { ModelRoutingEnabled: g.ModelRoutingEnabled, MCPXMLInject: g.MCPXMLInject, DefaultMappedModel: g.DefaultMappedModel, - SimulateClaudeMaxEnabled: g.SimulateClaudeMaxEnabled, SupportedModelScopes: g.SupportedModelScopes, AccountCount: g.AccountCount, ActiveAccountCount: g.ActiveAccountCount, @@ -279,17 +278,6 @@ func AccountFromServiceShallow(a *service.Account) *Account { } } - // 客户端亲和调度（Anthropic 和 Antigravity 账号） - if a.IsAffinityEnabled() { - enabled := true - out.ClientAffinityEnabled = &enabled - allow := a.IsAffinityAllowSwitch() - out.AffinityAllowSwitch = &allow - if pinnedUsers := a.GetPinnedUsers(); len(pinnedUsers) > 0 { - out.PinnedUserIDs = pinnedUsers - } - } - // 提取账号配额限制（apikey / bedrock 类型有效） if a.IsAPIKeyOrBedrock() { if limit := a.GetQuotaLimit(); limit > 0 { diff --git a/backend/internal/handler/dto/types.go b/backend/internal/handler/dto/types.go index 61752c6f3d..ce7003c6d5 100644 --- a/backend/internal/handler/dto/types.go +++ b/backend/internal/handler/dto/types.go @@ -117,8 +117,6 @@ type AdminGroup struct { // MCP XML 协议注入（仅 antigravity 平台使用） MCPXMLInject bool `json:"mcp_xml_inject"` - // Claude usage 模拟开关（仅管理员可见） - SimulateClaudeMaxEnabled bool `json:"simulate_claude_max_enabled"` // OpenAI Messages 调度配置（仅 openai 平台使用） DefaultMappedModel string `json:"default_mapped_model"` @@ -200,23 +198,6 @@ type Account struct { CacheTTLOverrideEnabled *bool `json:"cache_ttl_override_enabled,omitempty"` CacheTTLOverrideTarget *string `json:"cache_ttl_override_target,omitempty"` - // 客户端亲和调度（Anthropic 和 Antigravity 账号有效） - // 启用后新会话会优先调度到客户端之前使用过的账号 - ClientAffinityEnabled *bool `json:"client_affinity_enabled,omitempty"` - - // 亲和允许切换（默认 true） - AffinityAllowSwitch *bool `json:"affinity_allow_switch,omitempty"` - - // 亲和用户数量（admin 列表端点注入） - AffinityUserCount *int64 `json:"affinity_user_count,omitempty"` - - // 指定亲和用户 ID 列表 - PinnedUserIDs []int64 `json:"pinned_user_ids,omitempty"` - - // 亲和客户端数据（仅 admin 列表端点注入，不由 mapper 填充） - AffinityClientCount *int64 `json:"affinity_client_count,omitempty"` - AffinityClients []string `json:"affinity_clients,omitempty"` - // 自定义 Base URL 中继转发（仅 Anthropic OAuth/SetupToken 账号有效） CustomBaseURLEnabled *bool `json:"custom_base_url_enabled,omitempty"` CustomBaseURL *string `json:"custom_base_url,omitempty"` diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 55cfedf1c3..92f8b164f4 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -521,14 +521,6 @@ func (h *GatewayHandler) Messages(c *gin.Context) { // 选择支持该模型的账号 selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), currentAPIKey.GroupID, sessionKey, reqModel, fs.FailedAccountIDs, parsedReq.MetadataUserID, subject.UserID) if err != nil { - if errors.Is(err, service.ErrAffinityNoSwitch) { - h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "Affinity account unavailable and switching is disabled", streamStarted) - return - } - if errors.Is(err, service.ErrAffinityLimitExceeded) { - h.handleStreamingAwareError(c, http.StatusTooManyRequests, "api_error", "Affinity client limit exceeded", streamStarted) - return - } if len(fs.FailedAccountIDs) == 0 { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted) return diff --git a/backend/internal/handler/handler.go b/backend/internal/handler/handler.go index 125159e1b0..ebf8d5f674 100644 --- a/backend/internal/handler/handler.go +++ b/backend/internal/handler/handler.go @@ -47,7 +47,6 @@ type Handlers struct { OpenAIGateway *OpenAIGatewayHandler SoraGateway *SoraGatewayHandler SoraClient *SoraClientHandler - SoraVideos *SoraVideosHandler Setting *SettingHandler Totp *TotpHandler } diff --git a/backend/internal/handler/sora_client_handler.go b/backend/internal/handler/sora_client_handler.go index 3cb85d18b6..6060183af4 100644 --- a/backend/internal/handler/sora_client_handler.go +++ b/backend/internal/handler/sora_client_handler.go @@ -31,7 +31,6 @@ const ( type SoraClientHandler struct { genService *service.SoraGenerationService quotaService *service.SoraQuotaService - objectStorage service.SoraObjectStorage soraGatewayService *service.SoraGatewayService gatewayService *service.GatewayService mediaStorage *service.SoraMediaStorage @@ -48,7 +47,6 @@ type SoraClientHandler struct { func NewSoraClientHandler( genService *service.SoraGenerationService, quotaService *service.SoraQuotaService, - objectStorage service.SoraObjectStorage, soraGatewayService *service.SoraGatewayService, gatewayService *service.GatewayService, mediaStorage *service.SoraMediaStorage, @@ -57,7 +55,6 @@ func NewSoraClientHandler( return &SoraClientHandler{ genService: genService, quotaService: quotaService, - objectStorage: objectStorage, soraGatewayService: soraGatewayService, gatewayService: gatewayService, mediaStorage: mediaStorage, @@ -167,7 +164,7 @@ func (h *SoraClientHandler) Generate(c *gin.Context) { } // processGeneration 后台异步执行 Sora 生成任务。 -// 流程：选择账号 → Forward → 提取媒体 URL → 三层降级存储（S3 → 本地 → 上游）→ 更新记录。 +// 流程：选择账号 → Forward → 提取媒体 URL → 两层降级存储（本地 → 上游）→ 更新记录。 func (h *SoraClientHandler) processGeneration(genID int64, userID int64, groupID *int64, model, prompt, mediaType, imageInput string, videoCount int) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Minute) defer cancel() @@ -291,11 +288,11 @@ func (h *SoraClientHandler) processGeneration(genID int64, userID int64, groupID return } - // 三层降级存储：对象存储 → 本地 → 上游临时 URL + // 两层降级存储：本地 → 上游临时 URL storedURL, storedURLs, storageType, s3Keys, fileSize := h.storeMediaWithDegradation(ctx, userID, mediaType, mediaURL, mediaURLs) usageAdded := false - if (service.IsObjectStorageType(storageType) || storageType == service.SoraStorageTypeLocal) && fileSize > 0 && h.quotaService != nil { + if storageType == service.SoraStorageTypeLocal && fileSize > 0 && h.quotaService != nil { if err := h.quotaService.AddUsage(ctx, userID, fileSize); err != nil { h.cleanupStoredMedia(ctx, storageType, s3Keys, storedURLs) var quotaErr *service.QuotaExceededError @@ -336,7 +333,7 @@ func (h *SoraClientHandler) processGeneration(genID int64, userID int64, groupID logger.LegacyPrintf("handler.sora_client", "[SoraClient] 生成完成 id=%d storage=%s size=%d", genID, storageType, fileSize) } -// storeMediaWithDegradation 实现三层降级存储链：S3 → 本地 → 上游。 +// storeMediaWithDegradation 实现两层降级存储链：本地 → 上游。 func (h *SoraClientHandler) storeMediaWithDegradation( ctx context.Context, userID int64, mediaType string, mediaURL string, mediaURLs []string, @@ -346,46 +343,7 @@ func (h *SoraClientHandler) storeMediaWithDegradation( urls = []string{mediaURL} } - // 第一层：尝试对象存储（S3 / Google Drive） - if h.objectStorage != nil && h.objectStorage.Enabled(ctx) { - keys := make([]string, 0, len(urls)) - var totalSize int64 - var actualStorageType string - allOK := true - for _, u := range urls { - key, size, st, err := h.objectStorage.UploadFromURL(ctx, userID, u) - if err != nil { - logger.LegacyPrintf("handler.sora_client", "[SoraClient] 对象存储上传失败 type=%s err=%v", h.objectStorage.StorageType(), err) - allOK = false - // 清理已上传的文件 - if len(keys) > 0 { - _ = h.objectStorage.DeleteObjects(ctx, keys) - } - break - } - keys = append(keys, key) - totalSize += size - actualStorageType = st - } - if allOK && len(keys) > 0 { - accessURLs := make([]string, 0, len(keys)) - for _, key := range keys { - accessURL, err := h.objectStorage.GetAccessURL(ctx, key) - if err != nil { - logger.LegacyPrintf("handler.sora_client", "[SoraClient] 生成访问 URL 失败 type=%s err=%v", h.objectStorage.StorageType(), err) - _ = h.objectStorage.DeleteObjects(ctx, keys) - allOK = false - break - } - accessURLs = append(accessURLs, accessURL) - } - if allOK && len(accessURLs) > 0 { - return accessURLs[0], accessURLs, actualStorageType, keys, totalSize - } - } - } - - // 第二层：尝试本地存储 + // 第一层：尝试本地存储 if h.mediaStorage != nil && h.mediaStorage.Enabled() { storedPaths, err := h.mediaStorage.StoreFromURLs(ctx, mediaType, urls) if err == nil && len(storedPaths) > 0 { @@ -399,7 +357,7 @@ func (h *SoraClientHandler) storeMediaWithDegradation( logger.LegacyPrintf("handler.sora_client", "[SoraClient] 本地存储失败 err=%v", err) } - // 第三层：保留上游临时 URL + // 第二层：保留上游临时 URL return urls[0], urls, service.SoraStorageTypeUpstream, nil, 0 } @@ -650,7 +608,7 @@ func (h *SoraClientHandler) CancelGeneration(c *gin.Context) { response.Success(c, gin.H{"message": "已取消"}) } -// SaveToStorage 手动保存 upstream 记录到 S3。 +// SaveToStorage 手动保存 upstream 记录到本地存储。 // POST /api/v1/sora/generations/:id/save func (h *SoraClientHandler) SaveToStorage(c *gin.Context) { userID := getUserIDFromContext(c) @@ -680,8 +638,8 @@ func (h *SoraClientHandler) SaveToStorage(c *gin.Context) { return } - if h.objectStorage == nil || !h.objectStorage.Enabled(c.Request.Context()) { - response.Error(c, http.StatusServiceUnavailable, "云存储未配置，请联系管理员") + if h.mediaStorage == nil || !h.mediaStorage.Enabled() { + response.Error(c, http.StatusServiceUnavailable, "本地存储未配置，请联系管理员") return } @@ -694,40 +652,18 @@ func (h *SoraClientHandler) SaveToStorage(c *gin.Context) { return } - uploadedKeys := make([]string, 0, len(sourceURLs)) - accessURLs := make([]string, 0, len(sourceURLs)) - var totalSize int64 - - for _, sourceURL := range sourceURLs { - objectKey, fileSize, _, uploadErr := h.objectStorage.UploadFromURL(c.Request.Context(), userID, sourceURL) - if uploadErr != nil { - if len(uploadedKeys) > 0 { - _ = h.objectStorage.DeleteObjects(c.Request.Context(), uploadedKeys) - } - var upstreamErr *service.UpstreamDownloadError - if errors.As(uploadErr, &upstreamErr) && (upstreamErr.StatusCode == http.StatusForbidden || upstreamErr.StatusCode == http.StatusNotFound) { - response.Error(c, http.StatusGone, "媒体链接已过期，无法保存") - return - } - response.Error(c, http.StatusInternalServerError, "上传到存储失败: "+uploadErr.Error()) - return - } - accessURL, err := h.objectStorage.GetAccessURL(c.Request.Context(), objectKey) - if err != nil { - uploadedKeys = append(uploadedKeys, objectKey) - _ = h.objectStorage.DeleteObjects(c.Request.Context(), uploadedKeys) - response.Error(c, http.StatusInternalServerError, "生成访问链接失败: "+err.Error()) - return - } - uploadedKeys = append(uploadedKeys, objectKey) - accessURLs = append(accessURLs, accessURL) - totalSize += fileSize + storedPaths, storeErr := h.mediaStorage.StoreFromURLs(c.Request.Context(), gen.MediaType, sourceURLs) + if storeErr != nil { + response.Error(c, http.StatusInternalServerError, "保存到本地存储失败: "+storeErr.Error()) + return } + totalSize, _ := h.mediaStorage.TotalSizeByRelativePaths(storedPaths) + usageAdded := false if totalSize > 0 && h.quotaService != nil { if err := h.quotaService.AddUsage(c.Request.Context(), userID, totalSize); err != nil { - _ = h.objectStorage.DeleteObjects(c.Request.Context(), uploadedKeys) + _ = h.mediaStorage.DeleteByRelativePaths(storedPaths) var quotaErr *service.QuotaExceededError if errors.As(err, "aErr) { response.Error(c, http.StatusTooManyRequests, "存储配额已满，请删除不需要的作品释放空间") @@ -742,13 +678,13 @@ func (h *SoraClientHandler) SaveToStorage(c *gin.Context) { if err := h.genService.UpdateStorageForCompleted( c.Request.Context(), id, - accessURLs[0], - accessURLs, - h.objectStorage.StorageType(), - uploadedKeys, + storedPaths[0], + storedPaths, + service.SoraStorageTypeLocal, + nil, totalSize, ); err != nil { - _ = h.objectStorage.DeleteObjects(c.Request.Context(), uploadedKeys) + _ = h.mediaStorage.DeleteByRelativePaths(storedPaths) if usageAdded && h.quotaService != nil { _ = h.quotaService.ReleaseUsage(c.Request.Context(), userID, totalSize) } @@ -757,39 +693,24 @@ func (h *SoraClientHandler) SaveToStorage(c *gin.Context) { } response.Success(c, gin.H{ - "message": "已保存到云存储", - "object_key": uploadedKeys[0], - "object_keys": uploadedKeys, + "message": "已保存到本地存储", }) } // GetStorageStatus 返回存储状态。 // GET /api/v1/sora/storage-status func (h *SoraClientHandler) GetStorageStatus(c *gin.Context) { - objectStorageEnabled := h.objectStorage != nil && h.objectStorage.Enabled(c.Request.Context()) - objectStorageHealthy := false - storageType := "" - if objectStorageEnabled { - objectStorageHealthy = h.objectStorage.IsHealthy(c.Request.Context()) - storageType = h.objectStorage.StorageType() - } localEnabled := h.mediaStorage != nil && h.mediaStorage.Enabled() response.Success(c, gin.H{ - "s3_enabled": objectStorageEnabled, // 保留字段名向后兼容 - "s3_healthy": objectStorageHealthy, - "storage_type": storageType, + "s3_enabled": false, + "s3_healthy": false, + "storage_type": "", "local_enabled": localEnabled, }) } func (h *SoraClientHandler) cleanupStoredMedia(ctx context.Context, storageType string, s3Keys []string, localPaths []string) { - if service.IsObjectStorageType(storageType) { - if h.objectStorage != nil && len(s3Keys) > 0 { - if err := h.objectStorage.DeleteObjects(ctx, s3Keys); err != nil { - logger.LegacyPrintf("handler.sora_client", "[SoraClient] 清理存储文件失败 type=%s keys=%v err=%v", storageType, s3Keys, err) - } - } - } else if storageType == service.SoraStorageTypeLocal { + if storageType == service.SoraStorageTypeLocal { if h.mediaStorage != nil && len(localPaths) > 0 { if err := h.mediaStorage.DeleteByRelativePaths(localPaths); err != nil { logger.LegacyPrintf("handler.sora_client", "[SoraClient] 清理本地文件失败 paths=%v err=%v", localPaths, err) diff --git a/backend/internal/handler/sora_videos_handler.go b/backend/internal/handler/sora_videos_handler.go deleted file mode 100644 index b2a9d79c81..0000000000 --- a/backend/internal/handler/sora_videos_handler.go +++ /dev/null @@ -1,391 +0,0 @@ -package handler - -import ( - "encoding/json" - "errors" - "fmt" - "net/http" - "unicode/utf8" - - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware" - "github.com/Wei-Shaw/sub2api/internal/service" - - "github.com/gin-gonic/gin" -) - -// SoraVideosHandler handles Sora video/image async task API. -type SoraVideosHandler struct { - taskService *service.SoraTaskService - gatewayService *service.GatewayService - objectStorage service.SoraObjectStorage - mediaStorage *service.SoraMediaStorage - soraGatewayService *service.SoraGatewayService -} - -func NewSoraVideosHandler( - taskService *service.SoraTaskService, - gatewayService *service.GatewayService, - objectStorage service.SoraObjectStorage, - mediaStorage *service.SoraMediaStorage, - soraGatewayService *service.SoraGatewayService, -) *SoraVideosHandler { - if taskService == nil { - return nil - } - return &SoraVideosHandler{ - taskService: taskService, - gatewayService: gatewayService, - objectStorage: objectStorage, - mediaStorage: mediaStorage, - soraGatewayService: soraGatewayService, - } -} - -func (h *SoraVideosHandler) CreateVideo(c *gin.Context) { - apiKey, account, release, ok := h.selectAccount(c, "") - if !ok { - return - } - defer release() - - body, err := readBody(c) - if err != nil { - return - } - - var req service.CreateVideoRequest - if err := json.Unmarshal(body, &req); err != nil { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") - return - } - if req.Model == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "model is required") - return - } - - task, err := h.taskService.CreateVideoTask(c.Request.Context(), apiKey.ID, account, &req, body) - if err != nil { - handleTaskCreateError(c, "CreateVideo", err) - return - } - - c.JSON(http.StatusOK, service.TaskToResponse(task)) -} - -func (h *SoraVideosHandler) GetVideo(c *gin.Context) { - apiKey, ok := h.getAPIKey(c) - if !ok { - return - } - - taskID := c.Param("id") - if taskID == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "video_id is required") - return - } - - task, err := h.taskService.GetTask(c.Request.Context(), taskID, apiKey.ID) - if err != nil { - soraErrorResponse(c, http.StatusNotFound, "not_found", "Task not found") - return - } - - c.JSON(http.StatusOK, service.TaskToResponse(task)) -} - -func (h *SoraVideosHandler) RemixVideo(c *gin.Context) { - apiKey, ok := h.getAPIKey(c) - if !ok { - return - } - - taskID := c.Param("id") - if taskID == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "video_id is required") - return - } - - body, err := readBody(c) - if err != nil { - return - } - - var req service.RemixRequest - if err := json.Unmarshal(body, &req); err != nil || req.Prompt == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "prompt is required") - return - } - - originalTask, err := h.taskService.GetTask(c.Request.Context(), taskID, apiKey.ID) - if err != nil { - soraErrorResponse(c, http.StatusNotFound, "not_found", "Original video task not found") - return - } - if originalTask.Status != service.SoraTaskCompleted { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Original video must be completed before remix") - return - } - - remixTargetID := originalTask.ShareID - if remixTargetID == "" { - remixTargetID = originalTask.UpstreamTaskID - } - - account, err := h.selectAccountByID(c, originalTask.AccountID) - if err != nil { - soraErrorResponse(c, http.StatusServiceUnavailable, "server_error", "Failed to get account") - return - } - - videoReq := &service.CreateVideoRequest{ - Model: originalTask.Model, - Prompt: req.Prompt, - RemixTargetID: remixTargetID, - } - reqBody, _ := json.Marshal(videoReq) - - task, err := h.taskService.CreateVideoTask(c.Request.Context(), apiKey.ID, account, videoReq, reqBody) - if err != nil { - handleTaskCreateError(c, "RemixVideo", err) - return - } - - c.JSON(http.StatusOK, service.TaskToResponse(task)) -} - -// GetVideoContent returns video content based on storage configuration. -func (h *SoraVideosHandler) GetVideoContent(c *gin.Context) { - apiKey, ok := h.getAPIKey(c) - if !ok { - return - } - - taskID := c.Param("id") - if taskID == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "video_id is required") - return - } - - task, err := h.taskService.GetTask(c.Request.Context(), taskID, apiKey.ID) - if err != nil { - soraErrorResponse(c, http.StatusNotFound, "not_found", "Task not found") - return - } - - switch task.Status { - case service.SoraTaskCompleted: - contentURL := h.resolveContentURL(c, task) - if contentURL == "" { - soraErrorResponse(c, http.StatusNotFound, "not_found", "Video URL not available") - return - } - c.Redirect(http.StatusFound, contentURL) - - case service.SoraTaskFailed: - c.JSON(http.StatusGone, gin.H{ - "id": task.ID, - "object": task.ObjectType, - "status": task.Status, - "error": gin.H{ - "message": task.ErrorMessage, - "type": task.ErrorType, - }, - }) - - default: - c.JSON(http.StatusAccepted, service.TaskToResponse(task)) - } -} - -func (h *SoraVideosHandler) resolveContentURL(c *gin.Context, task *service.SoraTask) string { - if task.StoredKey == "" { - return task.VideoURL - } - - switch task.StorageType { - case "s3", "gdrive": - if h.objectStorage != nil { - accessURL, err := h.objectStorage.GetAccessURL(c.Request.Context(), task.StoredKey) - if err != nil { - logger.LegacyPrintf("handler.sora_videos", - "[GetVideoContent] task=%s get access URL error: %v, fallback to upstream", task.ID, err) - return task.VideoURL - } - return accessURL - } - return task.VideoURL - - case "local": - return "/sora/media" + task.StoredKey - - default: - return task.VideoURL - } -} - -func (h *SoraVideosHandler) CreateImage(c *gin.Context) { - apiKey, account, release, ok := h.selectAccount(c, "") - if !ok { - return - } - defer release() - - body, err := readBody(c) - if err != nil { - return - } - - var req service.CreateImageRequest - if err := json.Unmarshal(body, &req); err != nil { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") - return - } - if req.Prompt == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "prompt is required") - return - } - if req.Model == "" { - req.Model = inferImageModel(req.Size) - } - - task, err := h.taskService.CreateImageGeneration(c.Request.Context(), apiKey.ID, account, &req, body) - if err != nil { - handleTaskCreateError(c, "CreateImage", err) - return - } - - c.JSON(http.StatusOK, service.TaskToResponse(task)) -} - -func (h *SoraVideosHandler) EditImage(c *gin.Context) { - apiKey, account, release, ok := h.selectAccount(c, "") - if !ok { - return - } - defer release() - - body, err := readBody(c) - if err != nil { - return - } - - var req service.EditImageRequest - if err := json.Unmarshal(body, &req); err != nil { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Invalid request body") - return - } - if req.Image == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "image is required") - return - } - if req.Prompt == "" { - soraErrorResponse(c, http.StatusBadRequest, "parameter_missing", "prompt is required") - return - } - if req.Model == "" { - req.Model = inferImageModel(req.Size) - } - - imageReq := &service.CreateImageRequest{ - Model: req.Model, - Prompt: req.Prompt, - Image: req.Image, - Size: req.Size, - ResponseFormat: req.ResponseFormat, - N: 1, - } - - task, err := h.taskService.CreateImageGeneration(c.Request.Context(), apiKey.ID, account, imageReq, body) - if err != nil { - handleTaskCreateError(c, "EditImage", err) - return - } - - c.JSON(http.StatusOK, service.TaskToResponse(task)) -} - -// ── Internal helpers ── - -func (h *SoraVideosHandler) getAPIKey(c *gin.Context) (*service.APIKey, bool) { - apiKey, ok := middleware2.GetAPIKeyFromContext(c) - if !ok { - soraErrorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key") - return nil, false - } - return apiKey, true -} - -func (h *SoraVideosHandler) selectAccount(c *gin.Context, model string) (*service.APIKey, *service.Account, func(), bool) { - apiKey, ok := h.getAPIKey(c) - if !ok { - return nil, nil, nil, false - } - - selection, err := h.gatewayService.SelectAccountWithLoadAwareness( - c.Request.Context(), apiKey.GroupID, "", model, nil, "", int64(0), - ) - if err != nil { - soraErrorResponse(c, http.StatusServiceUnavailable, "server_error", "No available accounts") - return nil, nil, nil, false - } - - releaseFunc := func() {} - if selection.ReleaseFunc != nil { - releaseFunc = selection.ReleaseFunc - } - return apiKey, selection.Account, releaseFunc, true -} - -func (h *SoraVideosHandler) selectAccountByID(c *gin.Context, accountID int64) (*service.Account, error) { - return h.taskService.GetAccountByID(c.Request.Context(), accountID) -} - -func readBody(c *gin.Context) ([]byte, error) { - body, err := c.GetRawData() - if err != nil { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body") - return nil, err - } - if len(body) == 0 { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty") - return nil, fmt.Errorf("empty body") - } - if !utf8.Valid(body) { - soraErrorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body must be valid UTF-8") - return nil, fmt.Errorf("invalid utf-8") - } - return body, nil -} - -func soraErrorResponse(c *gin.Context, status int, errType, message string) { - c.JSON(status, gin.H{ - "error": gin.H{ - "message": message, - "type": errType, - }, - }) -} - -// handleTaskCreateError writes the appropriate error response, transparently -// forwarding upstream HTTP status codes when available. -func handleTaskCreateError(c *gin.Context, logTag string, err error) { - logger.LegacyPrintf("handler.sora_videos", "[%s] error: %v", logTag, err) - var ue *service.SoraUpstreamError - if errors.As(err, &ue) { - c.Data(ue.StatusCode, "application/json", ue.Body) - return - } - soraErrorResponse(c, http.StatusInternalServerError, "server_error", "Failed to create task") -} - -func inferImageModel(size string) string { - switch size { - case "540x360": - return "gpt-image-landscape" - case "360x540": - return "gpt-image-portrait" - default: - return "gpt-image" - } -} diff --git a/backend/internal/handler/wire.go b/backend/internal/handler/wire.go index 6377a5c0f0..c917f24a0d 100644 --- a/backend/internal/handler/wire.go +++ b/backend/internal/handler/wire.go @@ -88,7 +88,6 @@ func ProvideHandlers( openaiGatewayHandler *OpenAIGatewayHandler, soraGatewayHandler *SoraGatewayHandler, soraClientHandler *SoraClientHandler, - soraVideosHandler *SoraVideosHandler, settingHandler *SettingHandler, totpHandler *TotpHandler, _ *service.IdempotencyCoordinator, @@ -107,7 +106,6 @@ func ProvideHandlers( OpenAIGateway: openaiGatewayHandler, SoraGateway: soraGatewayHandler, SoraClient: soraClientHandler, - SoraVideos: soraVideosHandler, Setting: settingHandler, Totp: totpHandler, } diff --git a/backend/internal/repository/api_key_repo.go b/backend/internal/repository/api_key_repo.go index 9968853aef..667193a633 100644 --- a/backend/internal/repository/api_key_repo.go +++ b/backend/internal/repository/api_key_repo.go @@ -165,7 +165,6 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se group.FieldModelRoutingEnabled, group.FieldModelRouting, group.FieldMcpXMLInject, - group.FieldSimulateClaudeMaxEnabled, group.FieldSupportedModelScopes, group.FieldAllowMessagesDispatch, group.FieldDefaultMappedModel, @@ -660,7 +659,6 @@ func groupEntityToService(g *dbent.Group) *service.Group { ModelRouting: g.ModelRouting, ModelRoutingEnabled: g.ModelRoutingEnabled, MCPXMLInject: g.McpXMLInject, - SimulateClaudeMaxEnabled: g.SimulateClaudeMaxEnabled, SupportedModelScopes: g.SupportedModelScopes, SortOrder: g.SortOrder, AllowMessagesDispatch: g.AllowMessagesDispatch, diff --git a/backend/internal/repository/gateway_cache.go b/backend/internal/repository/gateway_cache.go index e40baffdc0..280629d1df 100644 --- a/backend/internal/repository/gateway_cache.go +++ b/backend/internal/repository/gateway_cache.go @@ -2,10 +2,7 @@ package repository import ( "context" - _ "embed" "fmt" - "strconv" - "strings" "time" "github.com/Wei-Shaw/sub2api/internal/service" @@ -13,34 +10,7 @@ import ( ) const ( - stickySessionPrefix = "sticky_session:" - affinityKeyPrefix = "affinity:" - affinityRevKeyPrefix = "affinity_rev:" -) - -var ( - //go:embed lua/get_affinity.lua - getAffinityLua string - //go:embed lua/update_affinity.lua - updateAffinityLua string - //go:embed lua/get_affinity_count.lua - getAffinityCountLua string - //go:embed lua/get_affinity_clients.lua - getAffinityClientsLua string - //go:embed lua/get_affinity_clients_with_scores.lua - getAffinityClientsWithScoresLua string - //go:embed lua/clear_account_affinity.lua - clearAccountAffinityLua string - //go:embed lua/get_affinity_multi_count.lua - getAffinityMultiCountLua string - - getAffinityScript = redis.NewScript(getAffinityLua) - updateAffinityScript = redis.NewScript(updateAffinityLua) - getAffinityCountScript = redis.NewScript(getAffinityCountLua) - getAffinityClientsScript = redis.NewScript(getAffinityClientsLua) - getAffinityClientsWithScoresScript = redis.NewScript(getAffinityClientsWithScoresLua) - clearAccountAffinityScript = redis.NewScript(clearAccountAffinityLua) - getAffinityMultiCountScript = redis.NewScript(getAffinityMultiCountLua) + stickySessionPrefix = "sticky_session:" ) type gatewayCache struct { @@ -51,16 +21,6 @@ func NewGatewayCache(rdb *redis.Client) service.GatewayCache { return &gatewayCache{rdb: rdb} } -// ensureScriptLoaded 确保 Lua 脚本已加载到 Redis 服务器的脚本缓存中。 -// Pipeline 中的 Script.Run 只发送 EVALSHA，如果 Redis 重启过导致脚本缓存丢失， -// EVALSHA 会返回 NOSCRIPT 错误。此方法提前加载脚本以避免该问题。 -func ensureScriptLoaded(ctx context.Context, rdb *redis.Client, script *redis.Script) { - exists, err := script.Exists(ctx, rdb).Result() - if err != nil || len(exists) == 0 || !exists[0] { - _ = script.Load(ctx, rdb).Err() - } -} - // buildSessionKey 构建 session key，包含 groupID 实现分组隔离 // 格式: sticky_session:{groupID}:{sessionHash} func buildSessionKey(groupID int64, sessionHash string) string { @@ -87,277 +47,3 @@ func (c *gatewayCache) DeleteSessionAccountID(ctx context.Context, groupID int64 key := buildSessionKey(groupID, sessionHash) return c.rdb.Del(ctx, key).Err() } - -// buildAffinityKey 构建正向亲和 key（member → accounts） -// 格式: affinity:{groupID}:{userID}/{clientID} -func buildAffinityKey(groupID int64, userID int64, clientID string) string { - return fmt.Sprintf("%s%d:%s", affinityKeyPrefix, groupID, buildAffinityMember(userID, clientID)) -} - -// buildAffinityReverseKey 构建反向亲和 key（account → members） -// 格式: affinity_rev:{groupID}:{accountID} -func buildAffinityReverseKey(groupID int64, accountID int64) string { - return fmt.Sprintf("%s%d:%d", affinityRevKeyPrefix, groupID, accountID) -} - -// buildAffinityMember 构建亲和成员标识 -// 格式: {userID}/{clientID} -func buildAffinityMember(userID int64, clientID string) string { - return fmt.Sprintf("%d/%s", userID, clientID) -} - -// parseAffinityMember 解析亲和成员标识为 userID 和 clientID -func parseAffinityMember(member string) (userID int64, clientID string) { - idx := strings.IndexByte(member, '/') - if idx < 0 { - // 兼容旧格式（纯 clientID，无 userID 前缀） - return 0, member - } - userID, _ = strconv.ParseInt(member[:idx], 10, 64) - clientID = member[idx+1:] - return userID, clientID -} - -// GetAffinityAccounts 获取亲和账号列表（按最近使用降序），同时清理过期成员 -func (c *gatewayCache) GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) { - key := buildAffinityKey(groupID, userID, clientID) - now := time.Now().Unix() - expireThreshold := now - int64(ttl.Seconds()) - - result, err := getAffinityScript.Run(ctx, c.rdb, []string{key}, expireThreshold).StringSlice() - if err != nil { - if err == redis.Nil { - return nil, nil - } - return nil, err - } - - accountIDs := make([]int64, 0, len(result)) - for _, s := range result { - id, err := strconv.ParseInt(s, 10, 64) - if err != nil { - continue - } - accountIDs = append(accountIDs, id) - } - return accountIDs, nil -} - -// UpdateAffinity 添加/更新亲和关系（更新 score 为当前时间戳，刷新 key TTL） -func (c *gatewayCache) UpdateAffinity(ctx context.Context, groupID int64, userID int64, clientID string, accountID int64, ttl time.Duration) error { - fwdKey := buildAffinityKey(groupID, userID, clientID) - revKey := buildAffinityReverseKey(groupID, accountID) - now := time.Now().Unix() - ttlSeconds := int64(ttl.Seconds()) - expireThreshold := now - ttlSeconds - - member := buildAffinityMember(userID, clientID) - return updateAffinityScript.Run(ctx, c.rdb, []string{fwdKey, revKey}, - now, ttlSeconds, accountID, expireThreshold, member, - ).Err() -} - -// GetAccountAffinityCountBatch 批量获取账号的亲和成员数量（惰性清理过期成员） -func (c *gatewayCache) GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) { - if len(accountIDs) == 0 { - return map[int64]int64{}, nil - } - - now := time.Now().Unix() - expireThreshold := now - int64(ttl.Seconds()) - - ensureScriptLoaded(ctx, c.rdb, getAffinityCountScript) - - pipe := c.rdb.Pipeline() - cmds := make([]*redis.Cmd, len(accountIDs)) - for i, accID := range accountIDs { - key := buildAffinityReverseKey(groupID, accID) - cmds[i] = getAffinityCountScript.Run(ctx, pipe, []string{key}, expireThreshold) - } - _, err := pipe.Exec(ctx) - if err != nil && err != redis.Nil { - return nil, err - } - - result := make(map[int64]int64, len(accountIDs)) - for i, accID := range accountIDs { - count, _ := cmds[i].Int64() - result[accID] = count - } - return result, nil -} - -// GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和成员列表（去重）。 -// accountGroups: map[accountID][]groupID，对每个 (groupID, accountID) 组合查询反向索引。 -// 返回值成员格式为 {userID}/{clientID}。 -func (c *gatewayCache) GetAccountAffinityClientsBatch(ctx context.Context, accountGroups map[int64][]int64, ttl time.Duration) (map[int64][]string, error) { - if len(accountGroups) == 0 { - return map[int64][]string{}, nil - } - - now := time.Now().Unix() - expireThreshold := now - int64(ttl.Seconds()) - - // 构建所有 (accountID, groupID) 组合的查询 - type queryItem struct { - accountID int64 - groupID int64 - } - var queries []queryItem - for accID, groupIDs := range accountGroups { - for _, gID := range groupIDs { - queries = append(queries, queryItem{accountID: accID, groupID: gID}) - } - } - - ensureScriptLoaded(ctx, c.rdb, getAffinityClientsScript) - - pipe := c.rdb.Pipeline() - cmds := make([]*redis.Cmd, len(queries)) - for i, q := range queries { - key := buildAffinityReverseKey(q.groupID, q.accountID) - cmds[i] = getAffinityClientsScript.Run(ctx, pipe, []string{key}, expireThreshold) - } - _, err := pipe.Exec(ctx) - if err != nil && err != redis.Nil { - return nil, err - } - - // 合并结果：同一个 accountID 跨多个 group 的成员去重 - result := make(map[int64][]string, len(accountGroups)) - seen := make(map[int64]map[string]struct{}, len(accountGroups)) - for i, q := range queries { - members, _ := cmds[i].StringSlice() - if len(members) == 0 { - continue - } - if seen[q.accountID] == nil { - seen[q.accountID] = make(map[string]struct{}) - } - for _, member := range members { - if _, exists := seen[q.accountID][member]; !exists { - seen[q.accountID][member] = struct{}{} - result[q.accountID] = append(result[q.accountID], member) - } - } - } - return result, nil -} - -// GetAccountAffinityClientsWithScores 获取单个账号跨所有分组的亲和客户端列表（含最后活跃时间戳，去重取最近）。 -func (c *gatewayCache) GetAccountAffinityClientsWithScores( - ctx context.Context, - accountID int64, - groupIDs []int64, - ttl time.Duration, -) ([]service.AffinityClient, error) { - if len(groupIDs) == 0 { - return nil, nil - } - - now := time.Now().Unix() - expireThreshold := now - int64(ttl.Seconds()) - - ensureScriptLoaded(ctx, c.rdb, getAffinityClientsWithScoresScript) - - pipe := c.rdb.Pipeline() - cmds := make([]*redis.Cmd, len(groupIDs)) - for i, gID := range groupIDs { - key := buildAffinityReverseKey(gID, accountID) - cmds[i] = getAffinityClientsWithScoresScript.Run(ctx, pipe, []string{key}, expireThreshold) - } - _, err := pipe.Exec(ctx) - if err != nil && err != redis.Nil { - return nil, err - } - - // 合并跨组结果，同一 member 取最近的 lastActive - type memberInfo struct { - userID int64 - clientID string - ts int64 - } - seen := make(map[string]*memberInfo) // member string → info - for _, cmd := range cmds { - vals, _ := cmd.StringSlice() - // vals 格式: [member1, score1, member2, score2, ...] - for j := 0; j+1 < len(vals); j += 2 { - member := vals[j] - ts, _ := strconv.ParseInt(vals[j+1], 10, 64) - if existing, ok := seen[member]; !ok || ts > existing.ts { - uid, cid := parseAffinityMember(member) - seen[member] = &memberInfo{userID: uid, clientID: cid, ts: ts} - } - } - } - - result := make([]service.AffinityClient, 0, len(seen)) - for _, info := range seen { - result = append(result, service.AffinityClient{ - UserID: info.userID, - ClientID: info.clientID, - LastActive: time.Unix(info.ts, 0), - }) - } - - // 按最后活跃时间降序排序 - service.SortAffinityClients(result) - - return result, nil -} - -// ClearAccountAffinity 清除指定账号在所有分组的亲和记录（正向+反向索引）。 -// 对每个 groupID 执行 Lua 脚本：读取反向索引获取所有成员， -// 从每个成员的正向索引中移除该账号，然后删除反向索引。 -func (c *gatewayCache) ClearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) error { - if len(groupIDs) == 0 { - return nil - } - - ensureScriptLoaded(ctx, c.rdb, clearAccountAffinityScript) - - pipe := c.rdb.Pipeline() - for _, gID := range groupIDs { - revKey := buildAffinityReverseKey(gID, accountID) - clearAccountAffinityScript.Run(ctx, pipe, []string{revKey}, gID, accountID) - } - _, err := pipe.Exec(ctx) - if err != nil && err != redis.Nil { - return err - } - return nil -} - -// GetAffinityMultiCount 获取账号的多维度亲和计数。 -// 返回: uniqueUsers（独立用户数）, uniqueClients（独立客户端数）, perUserClients（目标用户的客户端数） -func (c *gatewayCache) GetAffinityMultiCount( - ctx context.Context, - groupID int64, - accountID int64, - targetUserID int64, - ttl time.Duration, -) (users, clients, perUser int64, err error) { - key := buildAffinityReverseKey(groupID, accountID) - now := time.Now().Unix() - expireThreshold := now - int64(ttl.Seconds()) - - targetUserStr := "" - if targetUserID > 0 { - targetUserStr = strconv.FormatInt(targetUserID, 10) - } - - result, err := getAffinityMultiCountScript.Run(ctx, c.rdb, []string{key}, expireThreshold, targetUserStr).Int64Slice() - if err != nil { - if err == redis.Nil { - return 0, 0, 0, nil - } - return 0, 0, 0, err - } - - if len(result) < 4 { - return 0, 0, 0, nil - } - - // result: {totalMembers, uniqueUsers, uniqueClients, perUserClients} - return result[1], result[2], result[3], nil -} diff --git a/backend/internal/repository/group_repo.go b/backend/internal/repository/group_repo.go index 8d24bc6f14..674c655b80 100644 --- a/backend/internal/repository/group_repo.go +++ b/backend/internal/repository/group_repo.go @@ -61,8 +61,7 @@ func (r *groupRepository) Create(ctx context.Context, groupIn *service.Group) er SetMcpXMLInject(groupIn.MCPXMLInject). SetSoraStorageQuotaBytes(groupIn.SoraStorageQuotaBytes). SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch). - SetDefaultMappedModel(groupIn.DefaultMappedModel). - SetSimulateClaudeMaxEnabled(groupIn.SimulateClaudeMaxEnabled) + SetDefaultMappedModel(groupIn.DefaultMappedModel) // 设置模型路由配置 if groupIn.ModelRouting != nil { @@ -131,8 +130,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er SetMcpXMLInject(groupIn.MCPXMLInject). SetSoraStorageQuotaBytes(groupIn.SoraStorageQuotaBytes). SetAllowMessagesDispatch(groupIn.AllowMessagesDispatch). - SetDefaultMappedModel(groupIn.DefaultMappedModel). - SetSimulateClaudeMaxEnabled(groupIn.SimulateClaudeMaxEnabled) + SetDefaultMappedModel(groupIn.DefaultMappedModel) // 显式处理可空字段：nil 需要 clear，非 nil 需要 set。 if groupIn.DailyLimitUSD != nil { diff --git a/backend/internal/repository/lua/clear_account_affinity.lua b/backend/internal/repository/lua/clear_account_affinity.lua deleted file mode 100644 index d2f0af0338..0000000000 --- a/backend/internal/repository/lua/clear_account_affinity.lua +++ /dev/null @@ -1,29 +0,0 @@ --- 清除单个账号在指定分组的所有亲和记录（正向+反向） --- KEYS[1] = affinity_rev:{groupID}:{accountID} (反向索引) --- ARGV[1] = groupID (用于构建正向 key) --- ARGV[2] = accountID (正向索引中要移除的成员) --- 返回: 清理的成员数量 -local rev_key = KEYS[1] -local group_id = ARGV[1] -local account_id = ARGV[2] - --- 获取反向索引中所有成员 ({userID}/{clientID}) -local members = redis.call('ZRANGE', rev_key, 0, -1) -if #members == 0 then - return 0 -end - --- 从每个成员的正向索引中移除该账号 -for _, member in ipairs(members) do - local fwd_key = 'affinity:' .. group_id .. ':' .. member - redis.call('ZREM', fwd_key, account_id) - -- 如果正向索引为空，删除 key - if redis.call('ZCARD', fwd_key) == 0 then - redis.call('DEL', fwd_key) - end -end - --- 删除反向索引 -redis.call('DEL', rev_key) - -return #members diff --git a/backend/internal/repository/lua/get_affinity.lua b/backend/internal/repository/lua/get_affinity.lua deleted file mode 100644 index 9db85be2a9..0000000000 --- a/backend/internal/repository/lua/get_affinity.lua +++ /dev/null @@ -1,5 +0,0 @@ --- 清理过期成员后返回亲和账号列表（按最近使用降序） --- KEYS[1] = affinity:{groupID}:{userID}/{clientID} --- ARGV[1] = 过期阈值时间戳 (now - ttl) -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) -return redis.call('ZREVRANGE', KEYS[1], 0, -1) diff --git a/backend/internal/repository/lua/get_affinity_clients.lua b/backend/internal/repository/lua/get_affinity_clients.lua deleted file mode 100644 index 7290b63758..0000000000 --- a/backend/internal/repository/lua/get_affinity_clients.lua +++ /dev/null @@ -1,6 +0,0 @@ --- 清理过期成员后返回反向索引的成员列表（按最近使用降序） --- 成员格式: {userID}/{clientID} --- KEYS[1] = affinity_rev:{groupID}:{accountID} --- ARGV[1] = 过期阈值时间戳 (now - ttl) -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) -return redis.call('ZREVRANGE', KEYS[1], 0, -1) diff --git a/backend/internal/repository/lua/get_affinity_clients_with_scores.lua b/backend/internal/repository/lua/get_affinity_clients_with_scores.lua deleted file mode 100644 index bff6b616ed..0000000000 --- a/backend/internal/repository/lua/get_affinity_clients_with_scores.lua +++ /dev/null @@ -1,7 +0,0 @@ --- 清理过期成员后返回反向索引的成员列表及其 score（最后活跃时间戳） --- 成员格式: {userID}/{clientID} --- KEYS[1] = affinity_rev:{groupID}:{accountID} --- ARGV[1] = 过期阈值时间戳 (now - ttl) --- 返回: {member1, score1, member2, score2, ...}（按最近使用降序） -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) -return redis.call('ZREVRANGEBYSCORE', KEYS[1], '+inf', '-inf', 'WITHSCORES') diff --git a/backend/internal/repository/lua/get_affinity_count.lua b/backend/internal/repository/lua/get_affinity_count.lua deleted file mode 100644 index 7cf9b5cb51..0000000000 --- a/backend/internal/repository/lua/get_affinity_count.lua +++ /dev/null @@ -1,5 +0,0 @@ --- 清理过期成员后返回反向索引的成员数量 --- KEYS[1] = affinity_rev:{groupID}:{accountID} --- ARGV[1] = 过期阈值时间戳 (now - ttl) -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) -return redis.call('ZCARD', KEYS[1]) diff --git a/backend/internal/repository/lua/get_affinity_multi_count.lua b/backend/internal/repository/lua/get_affinity_multi_count.lua deleted file mode 100644 index 166cf4f9a3..0000000000 --- a/backend/internal/repository/lua/get_affinity_multi_count.lua +++ /dev/null @@ -1,40 +0,0 @@ --- 从反向索引解析多维度计数（用户/客户端/每用户客户端） --- KEYS[1] = affinity_rev:{groupID}:{accountID} --- ARGV[1] = 过期阈值时间戳 (now - ttl) --- ARGV[2] = 目标 userID（传 "" 则不计算 perUserClients） --- 返回: {totalMembers, uniqueUsers, uniqueClients, perUserClients} -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[1]) - -local members = redis.call('ZRANGE', KEYS[1], 0, -1) -local total = #members -if total == 0 then - return {0, 0, 0, 0} -end - -local target_user = ARGV[2] -local users = {} -local clients = {} -local user_count = 0 -local client_count = 0 -local per_user_count = 0 - -for _, member in ipairs(members) do - local sep = string.find(member, '/', 1, true) - if sep then - local uid = string.sub(member, 1, sep - 1) - local cid = string.sub(member, sep + 1) - if not users[uid] then - users[uid] = true - user_count = user_count + 1 - end - if not clients[cid] then - clients[cid] = true - client_count = client_count + 1 - end - if target_user ~= '' and uid == target_user then - per_user_count = per_user_count + 1 - end - end -end - -return {total, user_count, client_count, per_user_count} diff --git a/backend/internal/repository/lua/update_affinity.lua b/backend/internal/repository/lua/update_affinity.lua deleted file mode 100644 index 6e9b38fcb8..0000000000 --- a/backend/internal/repository/lua/update_affinity.lua +++ /dev/null @@ -1,15 +0,0 @@ --- 原子双写正向+反向索引 --- KEYS[1] = affinity:{groupID}:{userID}/{clientID} (正向: member → accounts) --- KEYS[2] = affinity_rev:{groupID}:{accountID} (反向: account → members) --- ARGV[1] = 当前时间戳 (score) --- ARGV[2] = TTL 秒数 --- ARGV[3] = accountID (正向索引的成员) --- ARGV[4] = 过期阈值时间戳 (now - ttl) --- ARGV[5] = {userID}/{clientID} (反向索引的成员) -redis.call('ZREMRANGEBYSCORE', KEYS[1], '-inf', ARGV[4]) -redis.call('ZADD', KEYS[1], ARGV[1], ARGV[3]) -redis.call('EXPIRE', KEYS[1], ARGV[2]) -redis.call('ZREMRANGEBYSCORE', KEYS[2], '-inf', ARGV[4]) -redis.call('ZADD', KEYS[2], ARGV[1], ARGV[5]) -redis.call('EXPIRE', KEYS[2], ARGV[2]) -return 1 diff --git a/backend/internal/repository/sora_task_repo.go b/backend/internal/repository/sora_task_repo.go deleted file mode 100644 index d46d50fe3c..0000000000 --- a/backend/internal/repository/sora_task_repo.go +++ /dev/null @@ -1,136 +0,0 @@ -package repository - -import ( - "context" - "database/sql" - "encoding/json" - "unicode/utf8" - - "github.com/Wei-Shaw/sub2api/internal/service" -) - -type SoraTaskRepository struct { - db *sql.DB -} - -func NewSoraTaskRepository(sqlDB *sql.DB) service.SoraTaskRepository { - return &SoraTaskRepository{db: sqlDB} -} - -func (r *SoraTaskRepository) Create(ctx context.Context, task *service.SoraTask) error { - var charStr, reqStr *string - if task.CharacterInfo != nil { - b, _ := json.Marshal(task.CharacterInfo) - s := string(b) - charStr = &s - } - if len(task.RequestBody) > 0 && utf8.Valid(task.RequestBody) { - s := string(task.RequestBody) - reqStr = &s - } - - _, err := r.db.ExecContext(ctx, ` - INSERT INTO sora_tasks ( - id, account_id, api_key_id, upstream_task_id, object_type, - model, prompt, status, progress, video_url, stored_key, storage_type, - share_id, character_info, error_message, error_type, - request_body, seconds, size, created_at, completed_at - ) VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12,$13,$14,$15,$16,$17,$18,$19,$20,$21)`, - task.ID, task.AccountID, task.APIKeyID, task.UpstreamTaskID, task.ObjectType, - task.Model, task.Prompt, task.Status, task.Progress, task.VideoURL, - task.StoredKey, task.StorageType, - task.ShareID, charStr, task.ErrorMessage, task.ErrorType, - reqStr, task.Seconds, task.Size, task.CreatedAt, task.CompletedAt, - ) - return err -} - -func (r *SoraTaskRepository) GetByID(ctx context.Context, id string) (*service.SoraTask, error) { - row := r.db.QueryRowContext(ctx, `SELECT `+soraTaskColumns+` FROM sora_tasks WHERE id = $1`, id) - return scanTask(row) -} - -func (r *SoraTaskRepository) GetByIDAndAPIKey(ctx context.Context, id string, apiKeyID int64) (*service.SoraTask, error) { - row := r.db.QueryRowContext(ctx, - `SELECT `+soraTaskColumns+` FROM sora_tasks WHERE id = $1 AND api_key_id = $2`, - id, apiKeyID, - ) - return scanTask(row) -} - -func (r *SoraTaskRepository) Update(ctx context.Context, task *service.SoraTask) error { - var charStr *string - if task.CharacterInfo != nil { - b, _ := json.Marshal(task.CharacterInfo) - s := string(b) - charStr = &s - } - - _, err := r.db.ExecContext(ctx, ` - UPDATE sora_tasks SET - upstream_task_id = $2, status = $3, progress = $4, - video_url = $5, stored_key = $6, storage_type = $7, - share_id = $8, character_info = $9, - error_message = $10, error_type = $11, - completed_at = $12, seconds = $13, size = $14, object_type = $15 - WHERE id = $1`, - task.ID, task.UpstreamTaskID, task.Status, task.Progress, - task.VideoURL, task.StoredKey, task.StorageType, - task.ShareID, charStr, - task.ErrorMessage, task.ErrorType, - task.CompletedAt, task.Seconds, task.Size, task.ObjectType, - ) - return err -} - -func (r *SoraTaskRepository) ListPending(ctx context.Context) ([]*service.SoraTask, error) { - rows, err := r.db.QueryContext(ctx, - `SELECT `+soraTaskColumns+` FROM sora_tasks WHERE status IN ('queued', 'in_progress') ORDER BY created_at ASC LIMIT 200`, - ) - if err != nil { - return nil, err - } - defer func() { _ = rows.Close() }() - - var tasks []*service.SoraTask - for rows.Next() { - t, err := scanTaskFromRow(rows) - if err != nil { - return nil, err - } - tasks = append(tasks, t) - } - return tasks, rows.Err() -} - -const soraTaskColumns = `id, account_id, api_key_id, upstream_task_id, object_type, - model, prompt, status, progress, video_url, stored_key, storage_type, - share_id, character_info, error_message, error_type, - request_body, seconds, size, created_at, completed_at` - -func scanTask(s scannable) (*service.SoraTask, error) { - var t service.SoraTask - var charJSON, reqBody []byte - err := s.Scan( - &t.ID, &t.AccountID, &t.APIKeyID, &t.UpstreamTaskID, &t.ObjectType, - &t.Model, &t.Prompt, &t.Status, &t.Progress, &t.VideoURL, - &t.StoredKey, &t.StorageType, - &t.ShareID, &charJSON, &t.ErrorMessage, &t.ErrorType, - &reqBody, &t.Seconds, &t.Size, &t.CreatedAt, &t.CompletedAt, - ) - if err != nil { - return nil, err - } - if len(charJSON) > 0 { - var ch service.SoraCharacter - if json.Unmarshal(charJSON, &ch) == nil && ch.Username != "" { - t.CharacterInfo = &ch - } - } - t.RequestBody = reqBody - return &t, nil -} - -func scanTaskFromRow(rows *sql.Rows) (*service.SoraTask, error) { - return scanTask(rows) -} diff --git a/backend/internal/server/routes/admin.go b/backend/internal/server/routes/admin.go index db78021ff0..b2443dd859 100644 --- a/backend/internal/server/routes/admin.go +++ b/backend/internal/server/routes/admin.go @@ -272,8 +272,6 @@ func registerAccountRoutes(admin *gin.RouterGroup, h *handler.Handlers) { accounts.POST("/today-stats/batch", h.Admin.Account.GetBatchTodayStats) accounts.POST("/:id/clear-rate-limit", h.Admin.Account.ClearRateLimit) accounts.POST("/:id/reset-quota", h.Admin.Account.ResetQuota) - accounts.GET("/:id/affinity-clients", h.Admin.Account.GetAffinityClients) - accounts.GET("/:id/affinity-details", h.Admin.Account.GetAffinityDetails) accounts.GET("/:id/temp-unschedulable", h.Admin.Account.GetTempUnschedulable) accounts.DELETE("/:id/temp-unschedulable", h.Admin.Account.ClearTempUnschedulable) accounts.POST("/:id/schedulable", h.Admin.Account.SetSchedulable) @@ -424,24 +422,6 @@ func registerSettingsRoutes(admin *gin.RouterGroup, h *handler.Handlers) { // Beta 策略配置 adminSettings.GET("/beta-policy", h.Admin.Setting.GetBetaPolicySettings) adminSettings.PUT("/beta-policy", h.Admin.Setting.UpdateBetaPolicySettings) - // Sora S3 存储配置（旧路由，保留兼容） - adminSettings.GET("/sora-s3", h.Admin.Setting.GetSoraS3Settings) - adminSettings.PUT("/sora-s3", h.Admin.Setting.UpdateSoraS3Settings) - adminSettings.POST("/sora-s3/test", h.Admin.Setting.TestSoraS3Connection) - adminSettings.GET("/sora-s3/profiles", h.Admin.Setting.ListSoraS3Profiles) - adminSettings.POST("/sora-s3/profiles", h.Admin.Setting.CreateSoraS3Profile) - adminSettings.PUT("/sora-s3/profiles/:profile_id", h.Admin.Setting.UpdateSoraS3Profile) - adminSettings.DELETE("/sora-s3/profiles/:profile_id", h.Admin.Setting.DeleteSoraS3Profile) - adminSettings.POST("/sora-s3/profiles/:profile_id/activate", h.Admin.Setting.SetActiveSoraS3Profile) - // Sora 统一存储配置（新路由，指向相同 handler） - adminSettings.GET("/sora-storage", h.Admin.Setting.GetSoraS3Settings) - adminSettings.PUT("/sora-storage", h.Admin.Setting.UpdateSoraS3Settings) - adminSettings.POST("/sora-storage/test", h.Admin.Setting.TestSoraS3Connection) - adminSettings.GET("/sora-storage/profiles", h.Admin.Setting.ListSoraS3Profiles) - adminSettings.POST("/sora-storage/profiles", h.Admin.Setting.CreateSoraS3Profile) - adminSettings.PUT("/sora-storage/profiles/:profile_id", h.Admin.Setting.UpdateSoraS3Profile) - adminSettings.DELETE("/sora-storage/profiles/:profile_id", h.Admin.Setting.DeleteSoraS3Profile) - adminSettings.POST("/sora-storage/profiles/:profile_id/activate", h.Admin.Setting.SetActiveSoraS3Profile) } } diff --git a/backend/internal/server/routes/gateway.go b/backend/internal/server/routes/gateway.go index c5a3529882..072cfdee37 100644 --- a/backend/internal/server/routes/gateway.go +++ b/backend/internal/server/routes/gateway.go @@ -175,16 +175,6 @@ func RegisterGatewayRoutes( { soraV1.POST("/chat/completions", h.SoraGateway.ChatCompletions) soraV1.GET("/models", h.Gateway.Models) - - // Sora Videos/Images async task API - if h.SoraVideos != nil { - soraV1.POST("/videos", h.SoraVideos.CreateVideo) - soraV1.GET("/videos/:id", h.SoraVideos.GetVideo) - soraV1.POST("/videos/:id/remix", h.SoraVideos.RemixVideo) - soraV1.GET("/videos/:id/content", h.SoraVideos.GetVideoContent) - soraV1.POST("/images/generations", h.SoraVideos.CreateImage) - soraV1.POST("/images/edits", h.SoraVideos.EditImage) - } } // Sora 媒体代理（可选 API Key 验证） diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 58805ce586..a1449ffd89 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -1229,29 +1229,6 @@ func (a *Account) IsSessionIDMaskingEnabled() bool { return false } -// IsAffinityEnabled 检查是否启用亲和调度（统一入口） -// 仅适用于 Anthropic 账号，同时检查新字段 affinity_enabled 和旧字段 client_affinity_enabled（向后兼容） -func (a *Account) IsAffinityEnabled() bool { - if a.Platform != PlatformAnthropic { - return false - } - if a.Extra == nil { - return false - } - if v, ok := a.Extra["affinity_enabled"]; ok { - if enabled, ok := v.(bool); ok { - return enabled - } - return false - } - if v, ok := a.Extra["client_affinity_enabled"]; ok { - if enabled, ok := v.(bool); ok { - return enabled - } - } - return false -} - // IsCustomBaseURLEnabled 检查是否启用自定义 base URL 中继转发 // 仅适用于 Anthropic OAuth/SetupToken 类型账号 func (a *Account) IsCustomBaseURLEnabled() bool { @@ -1269,221 +1246,6 @@ func (a *Account) IsCustomBaseURLEnabled() bool { return false } -// IsClientAffinityEnabled 向后兼容别名，内部调用 IsAffinityEnabled -func (a *Account) IsClientAffinityEnabled() bool { - return a.IsAffinityEnabled() -} - -// AffinityZone 表示账号的客户端亲和分区 -type AffinityZone int - -const ( - AffinityZoneGreen AffinityZone = iota // 绿区：允许绑定新客户端，优先调度 - AffinityZoneYellow // 黄区：允许绑定，仅在无绿区账号时降级调度 - AffinityZoneRed // 红区：禁止调度 -) - -// GetAffinityBase 获取亲和基础限制（绿区上限），0 表示未配置 -func (a *Account) GetAffinityBase() int { - if a.Extra == nil { - return 0 - } - if v, ok := a.Extra["affinity_base"]; ok { - return parseExtraInt(v) - } - return 0 -} - -// GetAffinityBuffer 获取亲和缓冲区大小（黄区范围） -// 返回 (value, configured)： -// - (0, false): 未配置 → 无限黄区（超过 base 永远黄区，永不红区） -// - (0, true): 显式设为 0 → 无黄区，超过 base 直接红区 -// - (n, true): n > 0 → 黄区范围为 base+1 到 base+n -func (a *Account) GetAffinityBuffer() (int, bool) { - if a.Extra == nil { - return 0, false - } - v, ok := a.Extra["affinity_buffer"] - if !ok { - return 0, false - } - // 显式设为 null/nil → 视为未配置 - if v == nil { - return 0, false - } - return parseExtraInt(v), true -} - -// GetAffinityZone 根据当前绑定的客户端数量计算账号的亲和分区。 -// 未开启亲和或未配置 base 的账号永远返回绿区。 -func (a *Account) GetAffinityZone(clientCount int64) AffinityZone { - if !a.IsClientAffinityEnabled() { - return AffinityZoneGreen - } - base := a.GetAffinityBase() - if base <= 0 { - return AffinityZoneGreen - } - if clientCount <= int64(base) { - return AffinityZoneGreen - } - buffer, configured := a.GetAffinityBuffer() - if !configured { - return AffinityZoneYellow // 未配置 buffer → 无限黄区 - } - if buffer == 0 { - return AffinityZoneRed // buffer=0 → 无黄区，直接红区 - } - if clientCount <= int64(base+buffer) { - return AffinityZoneYellow - } - return AffinityZoneRed -} - -// IsAffinityAllowSwitch 检查亲和账号是否允许在无绿区时切换到其他账号 -// 默认 true（允许切换），设为 false 时即使全部红区也不切换 -func (a *Account) IsAffinityAllowSwitch() bool { - if a.Extra == nil { - return true - } - if v, ok := a.Extra["affinity_allow_switch"]; ok { - if allow, ok := v.(bool); ok { - return allow - } - } - return true -} - -// GetPinnedUsers 获取指定亲和用户列表 -func (a *Account) GetPinnedUsers() []int64 { - if a.Extra == nil { - return nil - } - v, ok := a.Extra["pinned_users"] - if !ok || v == nil { - return nil - } - arr, ok := v.([]any) - if !ok || len(arr) == 0 { - return nil - } - result := make([]int64, 0, len(arr)) - for _, item := range arr { - switch id := item.(type) { - case float64: - result = append(result, int64(id)) - case int64: - result = append(result, id) - case int: - result = append(result, int64(id)) - case json.Number: - if i, err := id.Int64(); err == nil { - result = append(result, i) - } - } - } - return result -} - -// IsPinnedUser 检查 userID 是否在指定亲和用户列表中 -func (a *Account) IsPinnedUser(userID int64) bool { - for _, id := range a.GetPinnedUsers() { - if id == userID { - return true - } - } - return false -} - -// GetAffinityUserBase 获取用户维度亲和基础限制（绿区上限），0 表示未配置 -func (a *Account) GetAffinityUserBase() int { - if a.Extra == nil { - return 0 - } - if v, ok := a.Extra["affinity_user_base"]; ok { - return parseExtraInt(v) - } - return 0 -} - -// GetAffinityUserBuffer 获取用户维度亲和缓冲区大小（黄区范围） -// 返回 (value, configured)，语义与 GetAffinityBuffer 相同 -func (a *Account) GetAffinityUserBuffer() (int, bool) { - if a.Extra == nil { - return 0, false - } - v, ok := a.Extra["affinity_user_buffer"] - if !ok { - return 0, false - } - if v == nil { - return 0, false - } - return parseExtraInt(v), true -} - -// GetPerUserClientLimit 获取每用户客户端限制，0 表示不限制 -func (a *Account) GetPerUserClientLimit() int { - if a.Extra == nil { - return 0 - } - if v, ok := a.Extra["per_user_client_limit"]; ok { - return parseExtraInt(v) - } - return 0 -} - -// getAffinityZoneForDim 根据单一维度的计数和三区参数计算亲和分区 -func getAffinityZoneForDim(count int64, base int, buffer int, bufferConfigured bool) AffinityZone { - if base <= 0 { - return AffinityZoneGreen - } - if count <= int64(base) { - return AffinityZoneGreen - } - if !bufferConfigured { - return AffinityZoneYellow - } - if buffer == 0 { - return AffinityZoneRed - } - if count <= int64(base+buffer) { - return AffinityZoneYellow - } - return AffinityZoneRed -} - -// GetMultiDimAffinityZone 根据多维度计数计算亲和分区，取所有维度中最严格的区域。 -// 未开启亲和的账号永远返回绿区。 -func (a *Account) GetMultiDimAffinityZone(userCount, clientCount, perUserCount int64) AffinityZone { - if !a.IsAffinityEnabled() { - return AffinityZoneGreen - } - worst := AffinityZoneGreen - - // 客户端维度 - clientBase := a.GetAffinityBase() - clientBuffer, clientBufCfg := a.GetAffinityBuffer() - if z := getAffinityZoneForDim(clientCount, clientBase, clientBuffer, clientBufCfg); z > worst { - worst = z - } - - // 用户维度 - userBase := a.GetAffinityUserBase() - userBuffer, userBufCfg := a.GetAffinityUserBuffer() - if z := getAffinityZoneForDim(userCount, userBase, userBuffer, userBufCfg); z > worst { - worst = z - } - - // 每用户客户端维度 - perUserLimit := a.GetPerUserClientLimit() - if perUserLimit > 0 && perUserCount > int64(perUserLimit) { - worst = AffinityZoneRed - } - - return worst -} - // GetCustomBaseURL 返回自定义中继服务的 base URL func (a *Account) GetCustomBaseURL() string { return a.GetExtraString("custom_base_url") diff --git a/backend/internal/service/admin_service.go b/backend/internal/service/admin_service.go index be3017c25c..fe5cb2de74 100644 --- a/backend/internal/service/admin_service.go +++ b/backend/internal/service/admin_service.go @@ -155,7 +155,6 @@ type CreateGroupInput struct { ModelRouting map[string][]int64 ModelRoutingEnabled bool // 是否启用模型路由 MCPXMLInject *bool - SimulateClaudeMaxEnabled *bool // 支持的模型系列（仅 antigravity 平台使用） SupportedModelScopes []string // Sora 存储配额 @@ -195,7 +194,6 @@ type UpdateGroupInput struct { ModelRouting map[string][]int64 ModelRoutingEnabled *bool // 是否启用模型路由 MCPXMLInject *bool - SimulateClaudeMaxEnabled *bool // 支持的模型系列（仅 antigravity 平台使用） SupportedModelScopes *[]string // Sora 存储配额 @@ -884,13 +882,6 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn if input.MCPXMLInject != nil { mcpXMLInject = *input.MCPXMLInject } - simulateClaudeMaxEnabled := false - if input.SimulateClaudeMaxEnabled != nil { - if platform != PlatformAnthropic && *input.SimulateClaudeMaxEnabled { - return nil, fmt.Errorf("simulate_claude_max_enabled only supported for anthropic groups") - } - simulateClaudeMaxEnabled = *input.SimulateClaudeMaxEnabled - } // 如果指定了复制账号的源分组，先获取账号 ID 列表 var accountIDsToCopy []int64 @@ -947,7 +938,6 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn FallbackGroupIDOnInvalidRequest: fallbackOnInvalidRequest, ModelRouting: input.ModelRouting, MCPXMLInject: mcpXMLInject, - SimulateClaudeMaxEnabled: simulateClaudeMaxEnabled, SupportedModelScopes: input.SupportedModelScopes, SoraStorageQuotaBytes: input.SoraStorageQuotaBytes, AllowMessagesDispatch: input.AllowMessagesDispatch, @@ -1154,15 +1144,6 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd if input.MCPXMLInject != nil { group.MCPXMLInject = *input.MCPXMLInject } - if input.SimulateClaudeMaxEnabled != nil { - if group.Platform != PlatformAnthropic && *input.SimulateClaudeMaxEnabled { - return nil, fmt.Errorf("simulate_claude_max_enabled only supported for anthropic groups") - } - group.SimulateClaudeMaxEnabled = *input.SimulateClaudeMaxEnabled - } - if group.Platform != PlatformAnthropic { - group.SimulateClaudeMaxEnabled = false - } // 支持的模型系列（仅 antigravity 平台使用） if input.SupportedModelScopes != nil { diff --git a/backend/internal/service/admin_service_group_test.go b/backend/internal/service/admin_service_group_test.go index 51d9b3d2fe..536be0b583 100644 --- a/backend/internal/service/admin_service_group_test.go +++ b/backend/internal/service/admin_service_group_test.go @@ -785,57 +785,3 @@ func TestAdminService_UpdateGroup_InvalidRequestFallbackAllowsAntigravity(t *tes require.NotNil(t, repo.updated) require.Equal(t, fallbackID, *repo.updated.FallbackGroupIDOnInvalidRequest) } - -func TestAdminService_CreateGroup_SimulateClaudeMaxRequiresAnthropic(t *testing.T) { - repo := &groupRepoStubForAdmin{} - svc := &adminServiceImpl{groupRepo: repo} - - enabled := true - _, err := svc.CreateGroup(context.Background(), &CreateGroupInput{ - Name: "openai-group", - Platform: PlatformOpenAI, - SimulateClaudeMaxEnabled: &enabled, - }) - require.Error(t, err) - require.Contains(t, err.Error(), "simulate_claude_max_enabled only supported for anthropic groups") - require.Nil(t, repo.created) -} - -func TestAdminService_UpdateGroup_SimulateClaudeMaxRequiresAnthropic(t *testing.T) { - existingGroup := &Group{ - ID: 1, - Name: "openai-group", - Platform: PlatformOpenAI, - Status: StatusActive, - } - repo := &groupRepoStubForAdmin{getByID: existingGroup} - svc := &adminServiceImpl{groupRepo: repo} - - enabled := true - _, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{ - SimulateClaudeMaxEnabled: &enabled, - }) - require.Error(t, err) - require.Contains(t, err.Error(), "simulate_claude_max_enabled only supported for anthropic groups") - require.Nil(t, repo.updated) -} - -func TestAdminService_UpdateGroup_ClearsSimulateClaudeMaxWhenPlatformChanges(t *testing.T) { - existingGroup := &Group{ - ID: 1, - Name: "anthropic-group", - Platform: PlatformAnthropic, - Status: StatusActive, - SimulateClaudeMaxEnabled: true, - } - repo := &groupRepoStubForAdmin{getByID: existingGroup} - svc := &adminServiceImpl{groupRepo: repo} - - group, err := svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{ - Platform: PlatformOpenAI, - }) - require.NoError(t, err) - require.NotNil(t, group) - require.NotNil(t, repo.updated) - require.False(t, repo.updated.SimulateClaudeMaxEnabled) -} diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 647cf53733..6dba547cfc 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -1763,9 +1763,6 @@ func (s *AntigravityGatewayService) Forward(ctx context.Context, c *gin.Context, firstTokenMs = streamRes.firstTokenMs } - // Claude Max cache billing: 同步 ForwardResult.Usage 与客户端响应体一致 - applyClaudeMaxCacheBillingPolicyToUsage(usage, parsedRequestFromGinContext(c), claudeMaxGroupFromGinContext(c), originalModel, account.ID) - return &ForwardResult{ RequestID: requestID, Usage: *usage, @@ -3863,9 +3860,6 @@ returnResponse: return nil, s.writeClaudeError(c, http.StatusBadGateway, "upstream_error", "Failed to parse upstream response") } - // Claude Max cache billing simulation (non-streaming) - claudeResp = applyClaudeMaxNonStreamingRewrite(c, claudeResp, agUsage, originalModel, accountID) - c.Data(http.StatusOK, "application/json", claudeResp) // 转换为 service.ClaudeUsage @@ -3893,7 +3887,6 @@ func (s *AntigravityGatewayService) handleClaudeStreamingResponse(c *gin.Context } processor := antigravity.NewStreamingProcessor(originalModel) - setupClaudeMaxStreamingHook(c, processor, originalModel, accountID) var firstTokenMs *int // 使用 Scanner 并限制单行大小，避免 ReadString 无上限导致 OOM diff --git a/backend/internal/service/antigravity_smart_retry_test.go b/backend/internal/service/antigravity_smart_retry_test.go index 5bf65e1e99..8b45c551db 100644 --- a/backend/internal/service/antigravity_smart_retry_test.go +++ b/backend/internal/service/antigravity_smart_retry_test.go @@ -31,28 +31,6 @@ func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID c.deleteCalls = append(c.deleteCalls, deleteSessionCall{groupID: groupID, sessionHash: sessionHash}) return nil } -func (c *stubSmartRetryCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (c *stubSmartRetryCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (c *stubSmartRetryCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (c *stubSmartRetryCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (c *stubSmartRetryCache) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (c *stubSmartRetryCache) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (c *stubSmartRetryCache) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream type mockSmartRetryUpstream struct { responses []*http.Response diff --git a/backend/internal/service/api_key_auth_cache.go b/backend/internal/service/api_key_auth_cache.go index 258b842bca..2031e05aad 100644 --- a/backend/internal/service/api_key_auth_cache.go +++ b/backend/internal/service/api_key_auth_cache.go @@ -62,7 +62,6 @@ type APIKeyAuthGroupSnapshot struct { ModelRouting map[string][]int64 `json:"model_routing,omitempty"` ModelRoutingEnabled bool `json:"model_routing_enabled"` MCPXMLInject bool `json:"mcp_xml_inject"` - SimulateClaudeMaxEnabled bool `json:"simulate_claude_max_enabled"` // 支持的模型系列（仅 antigravity 平台使用） SupportedModelScopes []string `json:"supported_model_scopes,omitempty"` diff --git a/backend/internal/service/api_key_auth_cache_impl.go b/backend/internal/service/api_key_auth_cache_impl.go index d874ccf2ff..f727ab10f3 100644 --- a/backend/internal/service/api_key_auth_cache_impl.go +++ b/backend/internal/service/api_key_auth_cache_impl.go @@ -244,7 +244,6 @@ func (s *APIKeyService) snapshotFromAPIKey(apiKey *APIKey) *APIKeyAuthSnapshot { ModelRouting: apiKey.Group.ModelRouting, ModelRoutingEnabled: apiKey.Group.ModelRoutingEnabled, MCPXMLInject: apiKey.Group.MCPXMLInject, - SimulateClaudeMaxEnabled: apiKey.Group.SimulateClaudeMaxEnabled, SupportedModelScopes: apiKey.Group.SupportedModelScopes, AllowMessagesDispatch: apiKey.Group.AllowMessagesDispatch, DefaultMappedModel: apiKey.Group.DefaultMappedModel, @@ -304,7 +303,6 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho ModelRouting: snapshot.Group.ModelRouting, ModelRoutingEnabled: snapshot.Group.ModelRoutingEnabled, MCPXMLInject: snapshot.Group.MCPXMLInject, - SimulateClaudeMaxEnabled: snapshot.Group.SimulateClaudeMaxEnabled, SupportedModelScopes: snapshot.Group.SupportedModelScopes, AllowMessagesDispatch: snapshot.Group.AllowMessagesDispatch, DefaultMappedModel: snapshot.Group.DefaultMappedModel, diff --git a/backend/internal/service/claude_max_cache_billing_policy.go b/backend/internal/service/claude_max_cache_billing_policy.go deleted file mode 100644 index 2381915ef9..0000000000 --- a/backend/internal/service/claude_max_cache_billing_policy.go +++ /dev/null @@ -1,450 +0,0 @@ -package service - -import ( - "encoding/json" - "strings" - - "github.com/Wei-Shaw/sub2api/internal/pkg/claude" - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" - "github.com/tidwall/gjson" -) - -type claudeMaxCacheBillingOutcome struct { - Simulated bool -} - -func applyClaudeMaxCacheBillingPolicyToUsage(usage *ClaudeUsage, parsed *ParsedRequest, group *Group, model string, accountID int64) claudeMaxCacheBillingOutcome { - var out claudeMaxCacheBillingOutcome - if usage == nil || !shouldApplyClaudeMaxBillingRulesForUsage(group, model, parsed) { - return out - } - - resolvedModel := strings.TrimSpace(model) - if resolvedModel == "" && parsed != nil { - resolvedModel = strings.TrimSpace(parsed.Model) - } - - if hasCacheCreationTokens(*usage) { - // Upstream already returned cache creation usage; keep original usage. - return out - } - - if !shouldSimulateClaudeMaxUsageForUsage(*usage, parsed) { - return out - } - beforeInputTokens := usage.InputTokens - out.Simulated = safelyProjectUsageToClaudeMax1H(usage, parsed) - if out.Simulated { - logger.LegacyPrintf("service.gateway", "simulate_claude_max_usage: model=%s account=%d input_tokens:%d->%d cache_creation_1h=%d", - resolvedModel, - accountID, - beforeInputTokens, - usage.InputTokens, - usage.CacheCreation1hTokens, - ) - } - return out -} - -func isClaudeFamilyModel(model string) bool { - normalized := strings.ToLower(strings.TrimSpace(claude.NormalizeModelID(model))) - if normalized == "" { - return false - } - return strings.Contains(normalized, "claude-") -} - -func shouldApplyClaudeMaxBillingRules(input *RecordUsageInput) bool { - if input == nil || input.Result == nil || input.APIKey == nil || input.APIKey.Group == nil { - return false - } - return shouldApplyClaudeMaxBillingRulesForUsage(input.APIKey.Group, input.Result.Model, input.ParsedRequest) -} - -func shouldApplyClaudeMaxBillingRulesForUsage(group *Group, model string, parsed *ParsedRequest) bool { - if group == nil { - return false - } - if !group.SimulateClaudeMaxEnabled || group.Platform != PlatformAnthropic { - return false - } - - resolvedModel := model - if resolvedModel == "" && parsed != nil { - resolvedModel = parsed.Model - } - if !isClaudeFamilyModel(resolvedModel) { - return false - } - return true -} - -func hasCacheCreationTokens(usage ClaudeUsage) bool { - return usage.CacheCreationInputTokens > 0 || usage.CacheCreation5mTokens > 0 || usage.CacheCreation1hTokens > 0 -} - -func shouldSimulateClaudeMaxUsage(input *RecordUsageInput) bool { - if input == nil || input.Result == nil { - return false - } - if !shouldApplyClaudeMaxBillingRules(input) { - return false - } - return shouldSimulateClaudeMaxUsageForUsage(input.Result.Usage, input.ParsedRequest) -} - -func shouldSimulateClaudeMaxUsageForUsage(usage ClaudeUsage, parsed *ParsedRequest) bool { - if usage.InputTokens <= 0 { - return false - } - if hasCacheCreationTokens(usage) { - return false - } - if !hasClaudeCacheSignals(parsed) { - return false - } - return true -} - -func safelyProjectUsageToClaudeMax1H(usage *ClaudeUsage, parsed *ParsedRequest) (changed bool) { - defer func() { - if r := recover(); r != nil { - logger.LegacyPrintf("service.gateway", "simulate_claude_max_usage skipped: panic=%v", r) - changed = false - } - }() - return projectUsageToClaudeMax1H(usage, parsed) -} - -func projectUsageToClaudeMax1H(usage *ClaudeUsage, parsed *ParsedRequest) bool { - if usage == nil { - return false - } - totalWindowTokens := usage.InputTokens + usage.CacheCreation5mTokens + usage.CacheCreation1hTokens - if totalWindowTokens <= 1 { - return false - } - - simulatedInputTokens := computeClaudeMaxProjectedInputTokens(totalWindowTokens, parsed) - if simulatedInputTokens <= 0 { - simulatedInputTokens = 1 - } - if simulatedInputTokens >= totalWindowTokens { - simulatedInputTokens = totalWindowTokens - 1 - } - - cacheCreation1hTokens := totalWindowTokens - simulatedInputTokens - if usage.InputTokens == simulatedInputTokens && - usage.CacheCreation5mTokens == 0 && - usage.CacheCreation1hTokens == cacheCreation1hTokens && - usage.CacheCreationInputTokens == cacheCreation1hTokens { - return false - } - - usage.InputTokens = simulatedInputTokens - usage.CacheCreation5mTokens = 0 - usage.CacheCreation1hTokens = cacheCreation1hTokens - usage.CacheCreationInputTokens = cacheCreation1hTokens - return true -} - -type claudeCacheProjection struct { - HasBreakpoint bool - BreakpointCount int - TotalEstimatedTokens int - TailEstimatedTokens int -} - -func computeClaudeMaxProjectedInputTokens(totalWindowTokens int, parsed *ParsedRequest) int { - if totalWindowTokens <= 1 { - return totalWindowTokens - } - - projection := analyzeClaudeCacheProjection(parsed) - if !projection.HasBreakpoint || projection.TotalEstimatedTokens <= 0 || projection.TailEstimatedTokens <= 0 { - return totalWindowTokens - } - - totalEstimate := int64(projection.TotalEstimatedTokens) - tailEstimate := int64(projection.TailEstimatedTokens) - if tailEstimate > totalEstimate { - tailEstimate = totalEstimate - } - - scaled := (int64(totalWindowTokens)*tailEstimate + totalEstimate/2) / totalEstimate - if scaled <= 0 { - scaled = 1 - } - if scaled >= int64(totalWindowTokens) { - scaled = int64(totalWindowTokens - 1) - } - return int(scaled) -} - -func hasClaudeCacheSignals(parsed *ParsedRequest) bool { - if parsed == nil { - return false - } - if hasTopLevelEphemeralCacheControl(parsed) { - return true - } - return countExplicitCacheBreakpoints(parsed) > 0 -} - -func hasTopLevelEphemeralCacheControl(parsed *ParsedRequest) bool { - if parsed == nil || len(parsed.Body) == 0 { - return false - } - cacheType := strings.TrimSpace(gjson.GetBytes(parsed.Body, "cache_control.type").String()) - return strings.EqualFold(cacheType, "ephemeral") -} - -func analyzeClaudeCacheProjection(parsed *ParsedRequest) claudeCacheProjection { - var projection claudeCacheProjection - if parsed == nil { - return projection - } - - total := 0 - lastBreakpointAt := -1 - - switch system := parsed.System.(type) { - case string: - total += claudeMaxMessageOverheadTokens + estimateClaudeTextTokens(system) - case []any: - for _, raw := range system { - block, ok := raw.(map[string]any) - if !ok { - total += claudeMaxUnknownContentTokens - continue - } - total += estimateClaudeBlockTokens(block) - if hasEphemeralCacheControl(block) { - lastBreakpointAt = total - projection.BreakpointCount++ - projection.HasBreakpoint = true - } - } - } - - for _, rawMsg := range parsed.Messages { - total += claudeMaxMessageOverheadTokens - msg, ok := rawMsg.(map[string]any) - if !ok { - total += claudeMaxUnknownContentTokens - continue - } - content, exists := msg["content"] - if !exists { - continue - } - msgTokens, msgLastBreak, msgBreakCount := estimateClaudeContentTokens(content) - total += msgTokens - if msgBreakCount > 0 { - lastBreakpointAt = total - msgTokens + msgLastBreak - projection.BreakpointCount += msgBreakCount - projection.HasBreakpoint = true - } - } - - if total <= 0 { - total = 1 - } - projection.TotalEstimatedTokens = total - - if projection.HasBreakpoint && lastBreakpointAt >= 0 { - tail := total - lastBreakpointAt - if tail <= 0 { - tail = 1 - } - projection.TailEstimatedTokens = tail - return projection - } - - if hasTopLevelEphemeralCacheControl(parsed) { - tail := estimateLastUserMessageTokens(parsed) - if tail <= 0 { - tail = 1 - } - projection.HasBreakpoint = true - projection.BreakpointCount = 1 - projection.TailEstimatedTokens = tail - } - return projection -} - -func countExplicitCacheBreakpoints(parsed *ParsedRequest) int { - if parsed == nil { - return 0 - } - total := 0 - if system, ok := parsed.System.([]any); ok { - for _, raw := range system { - if block, ok := raw.(map[string]any); ok && hasEphemeralCacheControl(block) { - total++ - } - } - } - for _, rawMsg := range parsed.Messages { - msg, ok := rawMsg.(map[string]any) - if !ok { - continue - } - content, ok := msg["content"].([]any) - if !ok { - continue - } - for _, raw := range content { - if block, ok := raw.(map[string]any); ok && hasEphemeralCacheControl(block) { - total++ - } - } - } - return total -} - -func hasEphemeralCacheControl(block map[string]any) bool { - if block == nil { - return false - } - raw, ok := block["cache_control"] - if !ok || raw == nil { - return false - } - switch cc := raw.(type) { - case map[string]any: - cacheType, _ := cc["type"].(string) - return strings.EqualFold(strings.TrimSpace(cacheType), "ephemeral") - case map[string]string: - return strings.EqualFold(strings.TrimSpace(cc["type"]), "ephemeral") - default: - return false - } -} - -func estimateClaudeContentTokens(content any) (tokens int, lastBreakAt int, breakpointCount int) { - switch value := content.(type) { - case string: - return estimateClaudeTextTokens(value), -1, 0 - case []any: - total := 0 - lastBreak := -1 - breaks := 0 - for _, raw := range value { - block, ok := raw.(map[string]any) - if !ok { - total += claudeMaxUnknownContentTokens - continue - } - total += estimateClaudeBlockTokens(block) - if hasEphemeralCacheControl(block) { - lastBreak = total - breaks++ - } - } - return total, lastBreak, breaks - default: - return estimateStructuredTokens(value), -1, 0 - } -} - -func estimateClaudeBlockTokens(block map[string]any) int { - if block == nil { - return claudeMaxUnknownContentTokens - } - tokens := claudeMaxBlockOverheadTokens - blockType, _ := block["type"].(string) - switch blockType { - case "text": - if text, ok := block["text"].(string); ok { - tokens += estimateClaudeTextTokens(text) - } - case "tool_result": - if content, ok := block["content"]; ok { - nested, _, _ := estimateClaudeContentTokens(content) - tokens += nested - } - case "tool_use": - if name, ok := block["name"].(string); ok { - tokens += estimateClaudeTextTokens(name) - } - if input, ok := block["input"]; ok { - tokens += estimateStructuredTokens(input) - } - default: - if text, ok := block["text"].(string); ok { - tokens += estimateClaudeTextTokens(text) - } else if content, ok := block["content"]; ok { - nested, _, _ := estimateClaudeContentTokens(content) - tokens += nested - } - } - if tokens <= claudeMaxBlockOverheadTokens { - tokens += claudeMaxUnknownContentTokens - } - return tokens -} - -func estimateLastUserMessageTokens(parsed *ParsedRequest) int { - if parsed == nil || len(parsed.Messages) == 0 { - return 0 - } - for i := len(parsed.Messages) - 1; i >= 0; i-- { - msg, ok := parsed.Messages[i].(map[string]any) - if !ok { - continue - } - role, _ := msg["role"].(string) - if !strings.EqualFold(strings.TrimSpace(role), "user") { - continue - } - tokens, _, _ := estimateClaudeContentTokens(msg["content"]) - return claudeMaxMessageOverheadTokens + tokens - } - return 0 -} - -func estimateStructuredTokens(v any) int { - if v == nil { - return 0 - } - raw, err := json.Marshal(v) - if err != nil { - return claudeMaxUnknownContentTokens - } - return estimateClaudeTextTokens(string(raw)) -} - -func estimateClaudeTextTokens(text string) int { - if tokens, ok := estimateTokensByThirdPartyTokenizer(text); ok { - return tokens - } - return estimateClaudeTextTokensHeuristic(text) -} - -func estimateClaudeTextTokensHeuristic(text string) int { - normalized := strings.Join(strings.Fields(strings.TrimSpace(text)), " ") - if normalized == "" { - return 0 - } - asciiChars := 0 - nonASCIIChars := 0 - for _, r := range normalized { - if r <= 127 { - asciiChars++ - } else { - nonASCIIChars++ - } - } - tokens := nonASCIIChars - if asciiChars > 0 { - tokens += (asciiChars + 3) / 4 - } - if words := len(strings.Fields(normalized)); words > tokens { - tokens = words - } - if tokens <= 0 { - return 1 - } - return tokens -} diff --git a/backend/internal/service/claude_max_simulation_test.go b/backend/internal/service/claude_max_simulation_test.go deleted file mode 100644 index 3d2ae2e674..0000000000 --- a/backend/internal/service/claude_max_simulation_test.go +++ /dev/null @@ -1,156 +0,0 @@ -package service - -import ( - "strings" - "testing" -) - -func TestProjectUsageToClaudeMax1H_Conservation(t *testing.T) { - usage := &ClaudeUsage{ - InputTokens: 1200, - CacheCreationInputTokens: 0, - CacheCreation5mTokens: 0, - CacheCreation1hTokens: 0, - } - parsed := &ParsedRequest{ - Model: "claude-sonnet-4-5", - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": strings.Repeat("cached context ", 200), - "cache_control": map[string]any{"type": "ephemeral"}, - }, - map[string]any{ - "type": "text", - "text": "summarize quickly", - }, - }, - }, - }, - } - - changed := projectUsageToClaudeMax1H(usage, parsed) - if !changed { - t.Fatalf("expected usage to be projected") - } - - total := usage.InputTokens + usage.CacheCreation5mTokens + usage.CacheCreation1hTokens - if total != 1200 { - t.Fatalf("total tokens changed: got=%d want=%d", total, 1200) - } - if usage.CacheCreation5mTokens != 0 { - t.Fatalf("cache_creation_5m should be 0, got=%d", usage.CacheCreation5mTokens) - } - if usage.InputTokens <= 0 || usage.InputTokens >= 1200 { - t.Fatalf("simulated input out of range, got=%d", usage.InputTokens) - } - if usage.InputTokens > 100 { - t.Fatalf("simulated input should stay near cache breakpoint tail, got=%d", usage.InputTokens) - } - if usage.CacheCreation1hTokens <= 0 { - t.Fatalf("cache_creation_1h should be > 0, got=%d", usage.CacheCreation1hTokens) - } - if usage.CacheCreationInputTokens != usage.CacheCreation1hTokens { - t.Fatalf("cache_creation_input_tokens mismatch: got=%d want=%d", usage.CacheCreationInputTokens, usage.CacheCreation1hTokens) - } -} - -func TestComputeClaudeMaxProjectedInputTokens_Deterministic(t *testing.T) { - parsed := &ParsedRequest{ - Model: "claude-opus-4-5", - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": "build context", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - map[string]any{ - "type": "text", - "text": "what is failing now", - }, - }, - }, - }, - } - - got1 := computeClaudeMaxProjectedInputTokens(4096, parsed) - got2 := computeClaudeMaxProjectedInputTokens(4096, parsed) - if got1 != got2 { - t.Fatalf("non-deterministic input tokens: %d != %d", got1, got2) - } -} - -func TestShouldSimulateClaudeMaxUsage(t *testing.T) { - group := &Group{ - Platform: PlatformAnthropic, - SimulateClaudeMaxEnabled: true, - } - input := &RecordUsageInput{ - Result: &ForwardResult{ - Model: "claude-sonnet-4-5", - Usage: ClaudeUsage{ - InputTokens: 3000, - CacheCreationInputTokens: 0, - CacheCreation5mTokens: 0, - CacheCreation1hTokens: 0, - }, - }, - ParsedRequest: &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": "cached", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - map[string]any{ - "type": "text", - "text": "tail", - }, - }, - }, - }, - }, - APIKey: &APIKey{Group: group}, - } - - if !shouldSimulateClaudeMaxUsage(input) { - t.Fatalf("expected simulate=true for claude group with cache signal") - } - - input.ParsedRequest = &ParsedRequest{ - Messages: []any{ - map[string]any{"role": "user", "content": "no cache signal"}, - }, - } - if shouldSimulateClaudeMaxUsage(input) { - t.Fatalf("expected simulate=false when request has no cache signal") - } - - input.ParsedRequest = &ParsedRequest{ - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": "cached", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - }, - }, - }, - } - input.Result.Usage.CacheCreationInputTokens = 100 - if shouldSimulateClaudeMaxUsage(input) { - t.Fatalf("expected simulate=false when cache creation already exists") - } -} diff --git a/backend/internal/service/claude_tokenizer.go b/backend/internal/service/claude_tokenizer.go deleted file mode 100644 index 61f5e9616e..0000000000 --- a/backend/internal/service/claude_tokenizer.go +++ /dev/null @@ -1,41 +0,0 @@ -package service - -import ( - "sync" - - tiktoken "github.com/pkoukk/tiktoken-go" - tiktokenloader "github.com/pkoukk/tiktoken-go-loader" -) - -var ( - claudeTokenizerOnce sync.Once - claudeTokenizer *tiktoken.Tiktoken -) - -func getClaudeTokenizer() *tiktoken.Tiktoken { - claudeTokenizerOnce.Do(func() { - // Use offline loader to avoid runtime dictionary download. - tiktoken.SetBpeLoader(tiktokenloader.NewOfflineLoader()) - // Use a high-capacity tokenizer as the default approximation for Claude payloads. - enc, err := tiktoken.GetEncoding(tiktoken.MODEL_O200K_BASE) - if err != nil { - enc, err = tiktoken.GetEncoding(tiktoken.MODEL_CL100K_BASE) - } - if err == nil { - claudeTokenizer = enc - } - }) - return claudeTokenizer -} - -func estimateTokensByThirdPartyTokenizer(text string) (int, bool) { - enc := getClaudeTokenizer() - if enc == nil { - return 0, false - } - tokens := len(enc.EncodeOrdinary(text)) - if tokens <= 0 { - return 0, false - } - return tokens, true -} diff --git a/backend/internal/service/gateway_affinity_flow.go b/backend/internal/service/gateway_affinity_flow.go deleted file mode 100644 index cf6e7e6e26..0000000000 --- a/backend/internal/service/gateway_affinity_flow.go +++ /dev/null @@ -1,239 +0,0 @@ -package service - -import "context" - -// gatewayAffinityFlow encapsulates affinity-specific scheduling steps so the -// main account selection flow can stay focused on generic scheduling. -type gatewayAffinityFlow struct { - svc *GatewayService - ctx context.Context - groupID *int64 - sessionHash string - requestedModel string - affinityClientID string - affinityUserID int64 - platform string - useMixed bool - accountByID map[int64]*Account - isExcluded func(int64) bool -} - -type affinityWaitCandidate struct { - account *Account -} - -func newGatewayAffinityFlow( - svc *GatewayService, - ctx context.Context, - groupID *int64, - sessionHash string, - requestedModel string, - affinityClientID string, - affinityUserID int64, - platform string, - useMixed bool, - accountByID map[int64]*Account, - isExcluded func(int64) bool, -) *gatewayAffinityFlow { - return &gatewayAffinityFlow{ - svc: svc, - ctx: ctx, - groupID: groupID, - sessionHash: sessionHash, - requestedModel: requestedModel, - affinityClientID: affinityClientID, - affinityUserID: affinityUserID, - platform: platform, - useMixed: useMixed, - accountByID: accountByID, - isExcluded: isExcluded, - } -} - -// shouldFilterAccountWithoutClientID excludes affinity-enabled Anthropic accounts -// when metadata.user_id does not provide a usable client_id. -func shouldFilterAccountWithoutClientID(account *Account, affinityClientID string) bool { - if account == nil || affinityClientID != "" { - return false - } - if account.Platform != PlatformAnthropic { - return false - } - return account.IsAffinityEnabled() -} - -func filterAccountsWithoutClientID(accounts []Account, affinityClientID string) []Account { - if affinityClientID != "" { - return accounts - } - filtered := make([]Account, 0, len(accounts)) - for _, acc := range accounts { - if shouldFilterAccountWithoutClientID(&acc, affinityClientID) { - continue - } - filtered = append(filtered, acc) - } - return filtered -} - -func (f *gatewayAffinityFlow) preprocessPinnedUsers(accounts []Account) { - if f.affinityUserID <= 0 || f.affinityClientID == "" || f.svc.cache == nil { - return - } - for i := range accounts { - if accounts[i].IsPinnedUser(f.affinityUserID) && accounts[i].IsAffinityEnabled() { - _ = f.svc.cache.UpdateAffinity( - f.ctx, - derefGroupID(f.groupID), - f.affinityUserID, - f.affinityClientID, - accounts[i].ID, - ClientAffinityTTL, - ) - } - } -} - -// trySelectAffinityAccount runs Layer 1.4 and returns: -// - result != nil: affinity path selected an account or wait plan -// - affinityHit == true: an effective affinity-enabled record was considered and should suppress sticky fallback -func (f *gatewayAffinityFlow) trySelectAffinityAccount() (*AccountSelectionResult, bool, error) { - if f.affinityClientID == "" || f.affinityUserID <= 0 || f.svc.cache == nil { - return nil, false, nil - } - - gid := derefGroupID(f.groupID) - affinityAccountIDs, err := f.svc.cache.GetAffinityAccounts(f.ctx, gid, f.affinityUserID, f.affinityClientID, ClientAffinityTTL) - if err != nil || len(affinityAccountIDs) == 0 { - return nil, false, nil - } - - noSwitchBlocked := false - anyAllowSwitch := false - effectiveAffinityHit := false - waitCandidates := make(map[int64]*affinityWaitCandidate) - - for _, affinityAccID := range affinityAccountIDs { - account, ok := f.accountByID[affinityAccID] - - if f.isExcluded != nil && f.isExcluded(affinityAccID) { - checkAcc := account - if !ok && f.svc.accountRepo != nil { - if acc, repoErr := f.svc.accountRepo.GetByID(f.ctx, affinityAccID); repoErr == nil && acc != nil { - checkAcc = acc - } - } - if checkAcc != nil && checkAcc.IsAffinityEnabled() { - effectiveAffinityHit = true - if !checkAcc.IsAffinityAllowSwitch() { - noSwitchBlocked = true - } else { - anyAllowSwitch = true - } - } - continue - } - - if !ok || !f.svc.isAccountSchedulableForSelection(account) { - checkAcc := account - if !ok && f.svc.accountRepo != nil { - if acc, repoErr := f.svc.accountRepo.GetByID(f.ctx, affinityAccID); repoErr == nil && acc != nil { - checkAcc = acc - } - } - if checkAcc != nil && checkAcc.IsAffinityEnabled() { - effectiveAffinityHit = true - if !checkAcc.IsAffinityAllowSwitch() { - noSwitchBlocked = true - } else { - anyAllowSwitch = true - } - } - continue - } - - if !account.IsAffinityEnabled() { - continue - } - effectiveAffinityHit = true - if account.IsAffinityAllowSwitch() { - anyAllowSwitch = true - } - - if !f.svc.isAccountAllowedForPlatform(account, f.platform, f.useMixed) || - (f.requestedModel != "" && !f.svc.isModelSupportedByAccountWithContext(f.ctx, account, f.requestedModel)) || - !f.svc.isAccountSchedulableForModelSelection(f.ctx, account, f.requestedModel) || - !f.svc.isAccountSchedulableForQuota(account) || - !f.svc.isAccountSchedulableForWindowCost(f.ctx, account, false) || - !f.svc.isAccountSchedulableForRPM(f.ctx, account, false) { - if !account.IsAffinityAllowSwitch() { - noSwitchBlocked = true - } - continue - } - - userCount, clientCount, perUserCount, multiErr := f.svc.cache.GetAffinityMultiCount( - f.ctx, gid, affinityAccID, f.affinityUserID, ClientAffinityTTL, - ) - if multiErr == nil { - zone := account.GetMultiDimAffinityZone(userCount, clientCount, perUserCount) - if zone == AffinityZoneRed { - if !account.IsAffinityAllowSwitch() { - noSwitchBlocked = true - } - continue - } - } - - result, acquireErr := f.svc.tryAcquireAccountSlot(f.ctx, affinityAccID, account.Concurrency) - if acquireErr == nil && result.Acquired { - if !f.svc.checkAndRegisterSession(f.ctx, account, f.sessionHash) { - result.ReleaseFunc() - continue - } - _ = f.svc.cache.UpdateAffinity(f.ctx, gid, f.affinityUserID, f.affinityClientID, affinityAccID, ClientAffinityTTL) - if f.sessionHash != "" { - _ = f.svc.cache.SetSessionAccountID(f.ctx, gid, f.sessionHash, affinityAccID, stickySessionTTL) - } - return &AccountSelectionResult{ - Account: account, - Acquired: true, - ReleaseFunc: result.ReleaseFunc, - }, true, nil - } - if acquireErr == nil && !result.Acquired && !account.IsAffinityAllowSwitch() { - noSwitchBlocked = true - waitCandidates[affinityAccID] = &affinityWaitCandidate{account: account} - } - } - - if noSwitchBlocked && !anyAllowSwitch && f.svc.concurrencyService != nil { - for _, waitAccID := range affinityAccountIDs { - candidate, ok := waitCandidates[waitAccID] - if !ok || candidate == nil || candidate.account == nil { - continue - } - acc := candidate.account - waitingCount, _ := f.svc.concurrencyService.GetAccountWaitingCount(f.ctx, waitAccID) - if waitingCount >= f.svc.schedulingConfig().StickySessionMaxWaiting { - continue - } - if !f.svc.checkAndRegisterSession(f.ctx, acc, f.sessionHash) { - continue - } - cfg := f.svc.schedulingConfig() - return &AccountSelectionResult{ - Account: acc, - WaitPlan: &AccountWaitPlan{ - AccountID: waitAccID, - MaxConcurrency: acc.Concurrency, - Timeout: cfg.StickySessionWaitTimeout, - MaxWaiting: cfg.StickySessionMaxWaiting, - }, - }, true, nil - } - return nil, effectiveAffinityHit, ErrAffinityNoSwitch - } - - return nil, effectiveAffinityHit, nil -} diff --git a/backend/internal/service/gateway_affinity_scheduling_test.go b/backend/internal/service/gateway_affinity_scheduling_test.go deleted file mode 100644 index f5bf974c0a..0000000000 --- a/backend/internal/service/gateway_affinity_scheduling_test.go +++ /dev/null @@ -1,1759 +0,0 @@ -//go:build unit - -package service - -import ( - "context" - "errors" - "sort" - "testing" - "time" - - "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/stretchr/testify/assert" - "github.com/stretchr/testify/require" -) - -// --------------------------------------------------------------------------- -// Mock: GatewayCache for affinity tests -// --------------------------------------------------------------------------- - -// mockAffinityCache 为亲和调度测试提供可控的 GatewayCache mock。 -// 通过 getCountBatchFunc 可以自定义 GetAccountAffinityCountBatch 的行为。 -type mockAffinityCache struct { - getCountBatchFunc func(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) - getCountBatchCalls int // 记录 GetAccountAffinityCountBatch 被调用次数 -} - -func (m *mockAffinityCache) GetSessionAccountID(_ context.Context, _ int64, _ string) (int64, error) { - return 0, errors.New("not found") -} -func (m *mockAffinityCache) SetSessionAccountID(_ context.Context, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (m *mockAffinityCache) RefreshSessionTTL(_ context.Context, _ int64, _ string, _ time.Duration) error { - return nil -} -func (m *mockAffinityCache) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error { - return nil -} -func (m *mockAffinityCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (m *mockAffinityCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (m *mockAffinityCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (m *mockAffinityCache) GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) { - m.getCountBatchCalls++ - if m.getCountBatchFunc != nil { - return m.getCountBatchFunc(ctx, groupID, accountIDs, ttl) - } - return map[int64]int64{}, nil -} -func (m *mockAffinityCache) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (m *mockAffinityCache) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (m *mockAffinityCache) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - -// --------------------------------------------------------------------------- -// Helper: 构造启用了客户端亲和的 Anthropic 账号 -// --------------------------------------------------------------------------- - -func newAffinityAccount(id int64, priority int, affinityEnabled bool) *Account { - acc := &Account{ - ID: id, - Platform: PlatformAnthropic, - Priority: priority, - Status: StatusActive, - } - if affinityEnabled { - acc.Extra = map[string]any{"client_affinity_enabled": true} - } - return acc -} - -func newAffinityAccountWithLoad(id int64, priority int, loadRate int, affinityCount int64, lastUsedAt *time.Time) accountWithLoad { - return accountWithLoad{ - account: newAffinityAccount(id, priority, true), - loadInfo: &AccountLoadInfo{AccountID: id, LoadRate: loadRate}, - affinityCount: affinityCount, - } -} - -// =========================================================================== -// 1. filterByMinAffinityCount 测试 -// =========================================================================== - -func TestAffinityFilterByMinAffinityCount(t *testing.T) { - t.Run("empty slice returns empty", func(t *testing.T) { - result := filterByMinAffinityCount(nil) - require.Empty(t, result) - }) - - t.Run("single element returned as-is", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1}, loadInfo: &AccountLoadInfo{}, affinityCount: 5}, - } - result := filterByMinAffinityCount(accounts) - require.Len(t, result, 1) - require.Equal(t, int64(1), result[0].account.ID) - }) - - t.Run("all same affinityCount returns all", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1}, loadInfo: &AccountLoadInfo{}, affinityCount: 3}, - {account: &Account{ID: 2}, loadInfo: &AccountLoadInfo{}, affinityCount: 3}, - {account: &Account{ID: 3}, loadInfo: &AccountLoadInfo{}, affinityCount: 3}, - } - result := filterByMinAffinityCount(accounts) - require.Len(t, result, 3) - }) - - t.Run("filters to min affinityCount only", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1}, loadInfo: &AccountLoadInfo{}, affinityCount: 10}, - {account: &Account{ID: 2}, loadInfo: &AccountLoadInfo{}, affinityCount: 2}, - {account: &Account{ID: 3}, loadInfo: &AccountLoadInfo{}, affinityCount: 5}, - {account: &Account{ID: 4}, loadInfo: &AccountLoadInfo{}, affinityCount: 2}, - } - result := filterByMinAffinityCount(accounts) - require.Len(t, result, 2) - require.Equal(t, int64(2), result[0].account.ID) - require.Equal(t, int64(4), result[1].account.ID) - }) - - t.Run("zero affinityCount is smallest", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1}, loadInfo: &AccountLoadInfo{}, affinityCount: 5}, - {account: &Account{ID: 2}, loadInfo: &AccountLoadInfo{}, affinityCount: 0}, - {account: &Account{ID: 3}, loadInfo: &AccountLoadInfo{}, affinityCount: 3}, - {account: &Account{ID: 4}, loadInfo: &AccountLoadInfo{}, affinityCount: 0}, - } - result := filterByMinAffinityCount(accounts) - require.Len(t, result, 2) - require.Equal(t, int64(2), result[0].account.ID) - require.Equal(t, int64(4), result[1].account.ID) - }) - - t.Run("preserves order within same affinityCount", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 5}, loadInfo: &AccountLoadInfo{}, affinityCount: 1}, - {account: &Account{ID: 3}, loadInfo: &AccountLoadInfo{}, affinityCount: 1}, - {account: &Account{ID: 7}, loadInfo: &AccountLoadInfo{}, affinityCount: 2}, - {account: &Account{ID: 1}, loadInfo: &AccountLoadInfo{}, affinityCount: 1}, - } - result := filterByMinAffinityCount(accounts) - require.Len(t, result, 3) - // 验证保持原始顺序 - require.Equal(t, int64(5), result[0].account.ID) - require.Equal(t, int64(3), result[1].account.ID) - require.Equal(t, int64(1), result[2].account.ID) - }) -} - -// =========================================================================== -// 2. populateAffinityCounts 测试 -// =========================================================================== - -func TestAffinityPopulateAffinityCounts(t *testing.T) { - ctx := context.Background() - - t.Run("nil cache does not panic", func(t *testing.T) { - svc := &GatewayService{cache: nil} - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - } - // 不应 panic - svc.populateAffinityCounts(ctx, accounts, 0) - // affinityCount 保持零值 - require.Equal(t, int64(0), accounts[0].affinityCount) - }) - - t.Run("empty accounts returns immediately", func(t *testing.T) { - cache := &mockAffinityCache{} - svc := &GatewayService{cache: cache} - svc.populateAffinityCounts(ctx, nil, 0) - require.Equal(t, 0, cache.getCountBatchCalls, "should not call Redis for empty accounts") - }) - - t.Run("no affinity-enabled accounts skips Redis call", func(t *testing.T) { - cache := &mockAffinityCache{} - svc := &GatewayService{cache: cache} - accounts := []accountWithLoad{ - // Anthropic 但未启用亲和 - {account: newAffinityAccount(1, 1, false), loadInfo: &AccountLoadInfo{}}, - // 非 Anthropic 平台 - {account: &Account{ID: 2, Platform: PlatformOpenAI}, loadInfo: &AccountLoadInfo{}}, - } - svc.populateAffinityCounts(ctx, accounts, 0) - require.Equal(t, 0, cache.getCountBatchCalls, "should skip Redis when no affinity-enabled accounts") - }) - - t.Run("correctly populates affinityCount from Redis", func(t *testing.T) { - cache := &mockAffinityCache{ - getCountBatchFunc: func(_ context.Context, _ int64, accountIDs []int64, _ time.Duration) (map[int64]int64, error) { - result := map[int64]int64{} - for _, id := range accountIDs { - switch id { - case 1: - result[1] = 5 - case 2: - result[2] = 0 - case 3: - result[3] = 12 - } - } - return result, nil - }, - } - svc := &GatewayService{cache: cache} - - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(2, 1, false), loadInfo: &AccountLoadInfo{}}, // 未启用，但仍在列表中 - {account: newAffinityAccount(3, 1, true), loadInfo: &AccountLoadInfo{}}, - } - - svc.populateAffinityCounts(ctx, accounts, 100) - - require.Equal(t, 1, cache.getCountBatchCalls, "should call Redis exactly once") - require.Equal(t, int64(5), accounts[0].affinityCount, "account 1 should have count 5") - require.Equal(t, int64(0), accounts[1].affinityCount, "account 2 should have count 0") - require.Equal(t, int64(12), accounts[2].affinityCount, "account 3 should have count 12") - }) - - t.Run("Redis error degrades gracefully with counts at 0", func(t *testing.T) { - cache := &mockAffinityCache{ - getCountBatchFunc: func(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return nil, errors.New("redis connection refused") - }, - } - svc := &GatewayService{cache: cache} - - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(2, 1, true), loadInfo: &AccountLoadInfo{}}, - } - - svc.populateAffinityCounts(ctx, accounts, 0) - - require.Equal(t, 1, cache.getCountBatchCalls) - require.Equal(t, int64(0), accounts[0].affinityCount, "should remain 0 on error") - require.Equal(t, int64(0), accounts[1].affinityCount, "should remain 0 on error") - }) - - t.Run("partial Redis result fills only known accounts", func(t *testing.T) { - cache := &mockAffinityCache{ - getCountBatchFunc: func(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - // 只返回部分账号的计数 - return map[int64]int64{1: 7}, nil - }, - } - svc := &GatewayService{cache: cache} - - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(2, 1, true), loadInfo: &AccountLoadInfo{}}, - } - - svc.populateAffinityCounts(ctx, accounts, 0) - - require.Equal(t, int64(7), accounts[0].affinityCount) - require.Equal(t, int64(0), accounts[1].affinityCount, "missing account should default to 0") - }) - - t.Run("queries all account IDs regardless of affinity status", func(t *testing.T) { - // 验证：只要有至少一个 affinity-enabled 账号，就查询 ALL 账号的计数 - var queriedIDs []int64 - cache := &mockAffinityCache{ - getCountBatchFunc: func(_ context.Context, _ int64, accountIDs []int64, _ time.Duration) (map[int64]int64, error) { - queriedIDs = accountIDs - return map[int64]int64{}, nil - }, - } - svc := &GatewayService{cache: cache} - - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(2, 1, false), loadInfo: &AccountLoadInfo{}}, - {account: &Account{ID: 3, Platform: PlatformOpenAI}, loadInfo: &AccountLoadInfo{}}, - } - - svc.populateAffinityCounts(ctx, accounts, 0) - - require.Equal(t, 1, cache.getCountBatchCalls) - require.Equal(t, []int64{1, 2, 3}, queriedIDs, "should query ALL account IDs, not just affinity-enabled ones") - }) -} - -// =========================================================================== -// 3. Layer 1 排序链测试（sort.SliceStable 中 affinityCount 的正确性） -// =========================================================================== - -func TestAffinityLayer1SortChain(t *testing.T) { - now := time.Now() - earlier := now.Add(-1 * time.Hour) - muchEarlier := now.Add(-2 * time.Hour) - - // 复现 Layer 1 的排序逻辑 - sortByLayer1 := func(accounts []accountWithLoad) { - sort.SliceStable(accounts, func(i, j int) bool { - a, b := accounts[i], accounts[j] - if a.account.Priority != b.account.Priority { - return a.account.Priority < b.account.Priority - } - if a.loadInfo.LoadRate != b.loadInfo.LoadRate { - return a.loadInfo.LoadRate < b.loadInfo.LoadRate - } - if a.affinityCount != b.affinityCount { - return a.affinityCount < b.affinityCount - } - switch { - case a.account.LastUsedAt == nil && b.account.LastUsedAt != nil: - return true - case a.account.LastUsedAt != nil && b.account.LastUsedAt == nil: - return false - case a.account.LastUsedAt == nil && b.account.LastUsedAt == nil: - return false - default: - return a.account.LastUsedAt.Before(*b.account.LastUsedAt) - } - }) - } - - t.Run("same priority same loadRate sorts by affinityCount asc", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 10}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 2}, - {account: &Account{ID: 3, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 5}, - } - sortByLayer1(accounts) - require.Equal(t, int64(2), accounts[0].account.ID, "lowest affinityCount first") - require.Equal(t, int64(3), accounts[1].account.ID) - require.Equal(t, int64(1), accounts[2].account.ID, "highest affinityCount last") - }) - - t.Run("priority takes precedence over affinityCount", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 2, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 0}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 100}, - } - sortByLayer1(accounts) - require.Equal(t, int64(2), accounts[0].account.ID, "lower priority wins despite higher affinityCount") - }) - - t.Run("loadRate takes precedence over affinityCount", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 80}, affinityCount: 0}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}, affinityCount: 100}, - } - sortByLayer1(accounts) - require.Equal(t, int64(2), accounts[0].account.ID, "lower loadRate wins despite higher affinityCount") - }) - - t.Run("affinityCount takes precedence over LRU", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 5}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 1}, - } - sortByLayer1(accounts) - require.Equal(t, int64(2), accounts[0].account.ID, "lower affinityCount wins despite older LRU") - }) - - t.Run("same affinityCount falls through to LRU", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 3}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 3}, - {account: &Account{ID: 3, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 3}, - } - sortByLayer1(accounts) - require.Equal(t, int64(3), accounts[0].account.ID, "LRU: oldest used first") - require.Equal(t, int64(2), accounts[1].account.ID) - require.Equal(t, int64(1), accounts[2].account.ID, "LRU: most recently used last") - }) - - t.Run("full chain: priority > loadRate > affinityCount > LRU", func(t *testing.T) { - accounts := []accountWithLoad{ - // 优先级 2 - 不管其他维度如何，排在后面 - {account: &Account{ID: 10, Priority: 2, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 0}, affinityCount: 0}, - // 优先级 1，负载 80% - 负载高 - {account: &Account{ID: 20, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 80}, affinityCount: 0}, - // 优先级 1，负载 20%，亲和 5 - {account: &Account{ID: 30, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 20}, affinityCount: 5}, - // 优先级 1，负载 20%，亲和 1，最近使用 - {account: &Account{ID: 40, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 20}, affinityCount: 1}, - // 优先级 1，负载 20%，亲和 1，更早使用（应排最前） - {account: &Account{ID: 50, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 20}, affinityCount: 1}, - } - sortByLayer1(accounts) - // 期望排序：50 → 40 → 30 → 20 → 10 - require.Equal(t, int64(50), accounts[0].account.ID, "best: p1, lr20, ac1, LRU earlier") - require.Equal(t, int64(40), accounts[1].account.ID, "second: p1, lr20, ac1, LRU now") - require.Equal(t, int64(30), accounts[2].account.ID, "third: p1, lr20, ac5") - require.Equal(t, int64(20), accounts[3].account.ID, "fourth: p1, lr80") - require.Equal(t, int64(10), accounts[4].account.ID, "last: p2") - }) -} - -// =========================================================================== -// 4. Layer 2 分层过滤链完整性测试 -// =========================================================================== - -func TestAffinityLayer2FilterChain(t *testing.T) { - now := time.Now() - earlier := now.Add(-1 * time.Hour) - muchEarlier := now.Add(-2 * time.Hour) - - // 模拟 Layer 2 的完整过滤链：Priority → LoadRate → AffinityCount → LRU - applyLayer2 := func(accounts []accountWithLoad) *accountWithLoad { - candidates := filterByMinPriority(accounts) - candidates = filterByMinLoadRate(candidates) - candidates = filterByMinAffinityCount(candidates) - return selectByLRU(candidates, false) - } - - t.Run("priority different - affinityCount does not matter", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 2, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 0}, affinityCount: 0}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 100}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(2), selected.account.ID, "higher priority dimension overrides affinityCount") - }) - - t.Run("same priority same loadRate different affinityCount", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 30}, affinityCount: 10}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 30}, affinityCount: 2}, - {account: &Account{ID: 3, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 30}, affinityCount: 5}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(2), selected.account.ID, "lowest affinityCount wins") - }) - - t.Run("same priority same loadRate same affinityCount falls through to LRU", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 30}, affinityCount: 5}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &earlier}, loadInfo: &AccountLoadInfo{LoadRate: 30}, affinityCount: 5}, - {account: &Account{ID: 3, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 30}, affinityCount: 5}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(3), selected.account.ID, "LRU selects oldest") - }) - - t.Run("loadRate different overrides affinityCount", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 80}, affinityCount: 0}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}, affinityCount: 100}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(2), selected.account.ID, "lower loadRate wins over lower affinityCount") - }) - - t.Run("full chain integration: p → lr → ac → lru", func(t *testing.T) { - accounts := []accountWithLoad{ - // p=2 淘汰 - {account: &Account{ID: 1, Priority: 2, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 0}, affinityCount: 0}, - // p=1, lr=50 淘汰 - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 0}, - // p=1, lr=10, ac=8 淘汰 - {account: &Account{ID: 3, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}, affinityCount: 8}, - // p=1, lr=10, ac=2, lru=now 淘汰 - {account: &Account{ID: 4, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 10}, affinityCount: 2}, - // p=1, lr=10, ac=2, lru=muchEarlier → 胜出 - {account: &Account{ID: 5, Priority: 1, LastUsedAt: &muchEarlier}, loadInfo: &AccountLoadInfo{LoadRate: 10}, affinityCount: 2}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(5), selected.account.ID, "full chain selects ID=5") - }) - - t.Run("empty input returns nil", func(t *testing.T) { - selected := applyLayer2(nil) - require.Nil(t, selected) - }) - - t.Run("single account always selected", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 42, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 50}, affinityCount: 100}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(42), selected.account.ID) - }) - - t.Run("affinityCount zero preferred among same p and lr", func(t *testing.T) { - accounts := []accountWithLoad{ - {account: &Account{ID: 1, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 0}, affinityCount: 5}, - {account: &Account{ID: 2, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 0}, affinityCount: 0}, - {account: &Account{ID: 3, Priority: 1, LastUsedAt: &now}, loadInfo: &AccountLoadInfo{LoadRate: 0}, affinityCount: 3}, - } - selected := applyLayer2(accounts) - require.NotNil(t, selected) - require.Equal(t, int64(2), selected.account.ID, "zero affinityCount preferred") - }) -} - -// =========================================================================== -// 5. populateAffinityCounts + filterByMinAffinityCount 联合测试 -// =========================================================================== - -func TestAffinityPopulateAndFilterIntegration(t *testing.T) { - ctx := context.Background() - - t.Run("populate then filter selects least-loaded accounts", func(t *testing.T) { - cache := &mockAffinityCache{ - getCountBatchFunc: func(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{ - 1: 10, - 2: 3, - 3: 3, - 4: 7, - }, nil - }, - } - svc := &GatewayService{cache: cache} - - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(2, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(3, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(4, 1, true), loadInfo: &AccountLoadInfo{}}, - } - - svc.populateAffinityCounts(ctx, accounts, 0) - result := filterByMinAffinityCount(accounts) - - require.Len(t, result, 2) - require.Equal(t, int64(2), result[0].account.ID) - require.Equal(t, int64(3), result[1].account.ID) - }) - - t.Run("Redis failure results in all accounts having 0 affinityCount", func(t *testing.T) { - cache := &mockAffinityCache{ - getCountBatchFunc: func(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return nil, errors.New("timeout") - }, - } - svc := &GatewayService{cache: cache} - - accounts := []accountWithLoad{ - {account: newAffinityAccount(1, 1, true), loadInfo: &AccountLoadInfo{}}, - {account: newAffinityAccount(2, 1, true), loadInfo: &AccountLoadInfo{}}, - } - - svc.populateAffinityCounts(ctx, accounts, 0) - result := filterByMinAffinityCount(accounts) - - // 全部为 0，全部返回 - require.Len(t, result, 2, "all accounts should pass filter when Redis fails (all have count 0)") - }) -} - -// =========================================================================== -// 6. IsClientAffinityEnabled 边界测试 -// =========================================================================== - -func TestAffinityIsClientAffinityEnabled(t *testing.T) { - t.Run("Anthropic with enabled flag", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"client_affinity_enabled": true}, - } - assert.True(t, acc.IsClientAffinityEnabled()) - }) - - t.Run("Anthropic with disabled flag", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"client_affinity_enabled": false}, - } - assert.False(t, acc.IsClientAffinityEnabled()) - }) - - t.Run("Anthropic with nil Extra", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: nil, - } - assert.False(t, acc.IsClientAffinityEnabled()) - }) - - t.Run("Anthropic without the key", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"other_key": true}, - } - assert.False(t, acc.IsClientAffinityEnabled()) - }) - - t.Run("non-Anthropic platform always false", func(t *testing.T) { - platforms := []string{PlatformOpenAI, PlatformGemini, PlatformAntigravity} - for _, p := range platforms { - acc := &Account{ - Platform: p, - Extra: map[string]any{"client_affinity_enabled": true}, - } - assert.False(t, acc.IsClientAffinityEnabled(), "platform=%s should not support affinity", p) - } - }) - - t.Run("wrong type for enabled value", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"client_affinity_enabled": "true"}, // string 而非 bool - } - assert.False(t, acc.IsClientAffinityEnabled(), "string 'true' should not enable affinity") - }) - - t.Run("new flag false overrides legacy flag true", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{ - "affinity_enabled": false, - "client_affinity_enabled": true, - }, - } - assert.False(t, acc.IsClientAffinityEnabled()) - }) - - t.Run("new flag true overrides legacy flag false", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{ - "affinity_enabled": true, - "client_affinity_enabled": false, - }, - } - assert.True(t, acc.IsClientAffinityEnabled()) - }) -} - -// =========================================================================== -// GetAffinityZone 测试 -// =========================================================================== - -func TestGetAffinityZone(t *testing.T) { - makeAccount := func(enabled bool, base int, buffer any) *Account { - extra := map[string]any{"client_affinity_enabled": enabled} - if base > 0 { - extra["affinity_base"] = base - } - if buffer != nil { - extra["affinity_buffer"] = buffer - } - return &Account{ - Platform: PlatformAnthropic, - Extra: extra, - } - } - - t.Run("affinity disabled always green", func(t *testing.T) { - acc := makeAccount(false, 5, 3) - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(100)) - }) - - t.Run("no base configured always green", func(t *testing.T) { - acc := makeAccount(true, 0, nil) - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(100)) - }) - - t.Run("within base limit is green", func(t *testing.T) { - acc := makeAccount(true, 5, 3) - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(0)) - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(3)) - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(5)) - }) - - t.Run("no buffer configured infinite yellow", func(t *testing.T) { - acc := makeAccount(true, 5, nil) // buffer not set → infinite yellow - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(6)) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(100)) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(9999)) - }) - - t.Run("buffer zero no yellow zone", func(t *testing.T) { - acc := makeAccount(true, 5, 0) // buffer=0 → no yellow, direct red - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(5)) - assert.Equal(t, AffinityZoneRed, acc.GetAffinityZone(6)) - assert.Equal(t, AffinityZoneRed, acc.GetAffinityZone(100)) - }) - - t.Run("within buffer is yellow", func(t *testing.T) { - acc := makeAccount(true, 5, 3) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(6)) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(7)) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(8)) // base(5)+buffer(3)=8 - }) - - t.Run("beyond buffer is red", func(t *testing.T) { - acc := makeAccount(true, 5, 3) - assert.Equal(t, AffinityZoneRed, acc.GetAffinityZone(9)) - assert.Equal(t, AffinityZoneRed, acc.GetAffinityZone(100)) - }) - - t.Run("boundary exactly at base", func(t *testing.T) { - acc := makeAccount(true, 10, 5) - assert.Equal(t, AffinityZoneGreen, acc.GetAffinityZone(10)) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(11)) - }) - - t.Run("boundary exactly at base plus buffer", func(t *testing.T) { - acc := makeAccount(true, 10, 5) - assert.Equal(t, AffinityZoneYellow, acc.GetAffinityZone(15)) - assert.Equal(t, AffinityZoneRed, acc.GetAffinityZone(16)) - }) -} - -// =========================================================================== -// classifyByAffinityZone 测试 -// =========================================================================== - -func TestClassifyByAffinityZone(t *testing.T) { - makeAWL := func(id int64, base int, buffer any, count int64) accountWithLoad { - extra := map[string]any{"client_affinity_enabled": true} - if base > 0 { - extra["affinity_base"] = base - } - if buffer != nil { - extra["affinity_buffer"] = buffer - } - return accountWithLoad{ - account: &Account{ID: id, Platform: PlatformAnthropic, Extra: extra}, - loadInfo: &AccountLoadInfo{AccountID: id}, - affinityCount: count, - } - } - - t.Run("empty input returns empty", func(t *testing.T) { - result := classifyByAffinityZone(nil) - require.Empty(t, result) - }) - - t.Run("no zone config returns all", func(t *testing.T) { - // 没有账号配置 affinity_base → 原样返回 - accs := []accountWithLoad{ - {account: newAffinityAccount(1, 50, true), loadInfo: &AccountLoadInfo{AccountID: 1}}, - {account: newAffinityAccount(2, 50, true), loadInfo: &AccountLoadInfo{AccountID: 2}}, - } - result := classifyByAffinityZone(accs) - require.Len(t, result, 2) - }) - - t.Run("greens preferred over yellows", func(t *testing.T) { - accs := []accountWithLoad{ - makeAWL(1, 5, 3, 3), // green (3 ≤ 5) - makeAWL(2, 5, 3, 7), // yellow (5 < 7 ≤ 8) - makeAWL(3, 5, 3, 2), // green (2 ≤ 5) - } - result := classifyByAffinityZone(accs) - require.Len(t, result, 2) - - ids := []int64{result[0].account.ID, result[1].account.ID} - sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) - assert.Equal(t, []int64{1, 3}, ids) - }) - - t.Run("reds excluded", func(t *testing.T) { - accs := []accountWithLoad{ - makeAWL(1, 5, 3, 10), // red (10 > 8) - makeAWL(2, 5, 3, 6), // yellow (5 < 6 ≤ 8) - makeAWL(3, 5, 3, 9), // red (9 > 8) - } - result := classifyByAffinityZone(accs) - require.Len(t, result, 1) - assert.Equal(t, int64(2), result[0].account.ID) - }) - - t.Run("all red returns empty", func(t *testing.T) { - accs := []accountWithLoad{ - makeAWL(1, 5, 0, 6), // buffer=0 → red (6 > 5) - makeAWL(2, 5, 0, 10), // buffer=0 → red (10 > 5) - } - result := classifyByAffinityZone(accs) - require.Empty(t, result) - }) - - t.Run("mixed with unconfigured accounts", func(t *testing.T) { - // 账号 1: 配置了 base=5,buffer=3 → green(3≤5) - // 账号 2: 未配置 base → 视为 green - // 账号 3: 配置了 base=5,buffer=3 → red(10>8) - accs := []accountWithLoad{ - makeAWL(1, 5, 3, 3), - {account: newAffinityAccount(2, 50, true), loadInfo: &AccountLoadInfo{AccountID: 2}, affinityCount: 20}, - makeAWL(3, 5, 3, 10), - } - result := classifyByAffinityZone(accs) - require.Len(t, result, 2) - - ids := []int64{result[0].account.ID, result[1].account.ID} - sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] }) - assert.Equal(t, []int64{1, 2}, ids) - }) - - t.Run("infinite yellow never reds", func(t *testing.T) { - // buffer 未配置 → 无限黄区，永不红区 - accs := []accountWithLoad{ - makeAWL(1, 5, nil, 100), // yellow (100 > 5, no buffer → infinite yellow) - makeAWL(2, 5, nil, 3), // green (3 ≤ 5) - } - result := classifyByAffinityZone(accs) - // green 优先 - require.Len(t, result, 1) - assert.Equal(t, int64(2), result[0].account.ID) - }) - - t.Run("only yellows when no greens", func(t *testing.T) { - accs := []accountWithLoad{ - makeAWL(1, 5, nil, 10), // yellow - makeAWL(2, 5, nil, 20), // yellow - } - result := classifyByAffinityZone(accs) - require.Len(t, result, 2) - }) -} - -// =========================================================================== -// GetMultiDimAffinityZone 测试 -// =========================================================================== - -func TestGetMultiDimAffinityZone(t *testing.T) { - makeAccount := func(opts map[string]any) *Account { - extra := map[string]any{"affinity_enabled": true} - for k, v := range opts { - extra[k] = v - } - return &Account{ - Platform: PlatformAnthropic, - Extra: extra, - } - } - - t.Run("user green + client green = green", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - "affinity_user_base": 3, - "affinity_user_buffer": 2, - }) - assert.Equal(t, AffinityZoneGreen, acc.GetMultiDimAffinityZone(2, 5, 0)) - }) - - t.Run("user green + client yellow = yellow", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - "affinity_user_base": 3, - "affinity_user_buffer": 2, - }) - assert.Equal(t, AffinityZoneYellow, acc.GetMultiDimAffinityZone(2, 12, 0)) - }) - - t.Run("user yellow + client green = yellow", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - "affinity_user_base": 3, - "affinity_user_buffer": 2, - }) - assert.Equal(t, AffinityZoneYellow, acc.GetMultiDimAffinityZone(4, 5, 0)) - }) - - t.Run("user red + client green = red", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - "affinity_user_base": 3, - "affinity_user_buffer": 0, // no yellow, direct red - }) - assert.Equal(t, AffinityZoneRed, acc.GetMultiDimAffinityZone(4, 5, 0)) - }) - - t.Run("user green + client red = red", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 0, // no yellow, direct red - "affinity_user_base": 3, - "affinity_user_buffer": 2, - }) - assert.Equal(t, AffinityZoneRed, acc.GetMultiDimAffinityZone(2, 11, 0)) - }) - - t.Run("user_base=0 only checks client dimension", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - // no affinity_user_base → 0 → user dimension always green - }) - // client in yellow range - assert.Equal(t, AffinityZoneYellow, acc.GetMultiDimAffinityZone(100, 12, 0)) - }) - - t.Run("perUserCount exceeds limit = red", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - "per_user_client_limit": 3, - }) - // client green, user green, but perUser > limit - assert.Equal(t, AffinityZoneRed, acc.GetMultiDimAffinityZone(1, 5, 4)) - }) - - t.Run("perUserCount within limit = green", func(t *testing.T) { - acc := makeAccount(map[string]any{ - "affinity_base": 10, - "affinity_buffer": 5, - "per_user_client_limit": 3, - }) - assert.Equal(t, AffinityZoneGreen, acc.GetMultiDimAffinityZone(1, 5, 3)) - }) - - t.Run("affinity disabled always green", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{ - "affinity_enabled": false, - "affinity_base": 1, - "affinity_buffer": 0, - }, - } - assert.Equal(t, AffinityZoneGreen, acc.GetMultiDimAffinityZone(100, 100, 100)) - }) -} - -// =========================================================================== -// IsAffinityAllowSwitch 测试 -// =========================================================================== - -func TestIsAffinityAllowSwitch(t *testing.T) { - t.Run("default true when no field in Extra", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"affinity_enabled": true}, - } - assert.True(t, acc.IsAffinityAllowSwitch()) - }) - - t.Run("explicit false", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"affinity_allow_switch": false}, - } - assert.False(t, acc.IsAffinityAllowSwitch()) - }) - - t.Run("Extra nil defaults to true", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: nil, - } - assert.True(t, acc.IsAffinityAllowSwitch()) - }) - - t.Run("explicit true", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{"affinity_allow_switch": true}, - } - assert.True(t, acc.IsAffinityAllowSwitch()) - }) -} - -// =========================================================================== -// GetPinnedUsers / IsPinnedUser 测试 -// =========================================================================== - -func TestGetPinnedUsers(t *testing.T) { - t.Run("normal list", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{ - "pinned_users": []any{float64(10), float64(20), float64(30)}, - }, - } - result := acc.GetPinnedUsers() - assert.Equal(t, []int64{10, 20, 30}, result) - }) - - t.Run("empty list", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{ - "pinned_users": []any{}, - }, - } - result := acc.GetPinnedUsers() - assert.Nil(t, result) - }) - - t.Run("Extra nil", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: nil, - } - result := acc.GetPinnedUsers() - assert.Nil(t, result) - }) - - t.Run("pinned_users not set", func(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{}, - } - result := acc.GetPinnedUsers() - assert.Nil(t, result) - }) -} - -func TestIsPinnedUser(t *testing.T) { - acc := &Account{ - Platform: PlatformAnthropic, - Extra: map[string]any{ - "pinned_users": []any{float64(10), float64(20), float64(30)}, - }, - } - - t.Run("hit", func(t *testing.T) { - assert.True(t, acc.IsPinnedUser(20)) - }) - - t.Run("miss", func(t *testing.T) { - assert.False(t, acc.IsPinnedUser(99)) - }) - - t.Run("nil Extra", func(t *testing.T) { - nilAcc := &Account{Platform: PlatformAnthropic, Extra: nil} - assert.False(t, nilAcc.IsPinnedUser(10)) - }) -} - -// =========================================================================== -// Enhanced mock for integration tests (supports affinity lookups + tracking) -// =========================================================================== - -type affinityIntegrationCache struct { - mockAffinityCache - getAffinityAccountsFunc func(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) - updateAffinityCalls []affinityUpdateCall - getMultiCountFunc func(ctx context.Context, groupID int64, accountID int64, targetUserID int64, ttl time.Duration) (int64, int64, int64, error) - sessionBindings map[string]int64 - getAffinityAccountsCalled bool // 标记 GetAffinityAccounts 是否已被调用，用于区分预处理/Layer 2 阶段 -} - -type affinityUpdateCall struct { - groupID int64 - userID int64 - clientID string - accountID int64 - beforeAffinityLookup bool // true 表示该调用发生在 GetAffinityAccounts 之前（预处理阶段） -} - -func (c *affinityIntegrationCache) GetSessionAccountID(_ context.Context, _ int64, hash string) (int64, error) { - if c.sessionBindings != nil { - if id, ok := c.sessionBindings[hash]; ok { - return id, nil - } - } - return 0, errors.New("not found") -} - -func (c *affinityIntegrationCache) GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) { - c.getAffinityAccountsCalled = true - if c.getAffinityAccountsFunc != nil { - return c.getAffinityAccountsFunc(ctx, groupID, userID, clientID, ttl) - } - return nil, nil -} - -func (c *affinityIntegrationCache) UpdateAffinity(_ context.Context, groupID int64, userID int64, clientID string, accountID int64, _ time.Duration) error { - c.updateAffinityCalls = append(c.updateAffinityCalls, affinityUpdateCall{ - groupID: groupID, userID: userID, clientID: clientID, accountID: accountID, - beforeAffinityLookup: !c.getAffinityAccountsCalled, - }) - return nil -} - -func (c *affinityIntegrationCache) GetAffinityMultiCount(ctx context.Context, groupID int64, accountID int64, targetUserID int64, ttl time.Duration) (int64, int64, int64, error) { - if c.getMultiCountFunc != nil { - return c.getMultiCountFunc(ctx, groupID, accountID, targetUserID, ttl) - } - return 0, 0, 0, nil -} - -// =========================================================================== -// TestAffinityPreprocessPinnedUsers — 验证 pinned_users 预处理逻辑 -// =========================================================================== - -func TestAffinityPreprocessPinnedUsers(t *testing.T) { - ctx := context.Background() - testUserID := int64(42) - testClientID := "user_" + "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" + "_account_test" - - t.Run("pinned user triggers UpdateAffinity", func(t *testing.T) { - cache := &affinityIntegrationCache{} - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "pinned_users": []any{float64(42)}, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{"affinity_enabled": true}, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{1: true, 2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - _, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - - // 验证 UpdateAffinity 至少被调用了一次（pinned preprocess for account 1） - found := false - for _, call := range cache.updateAffinityCalls { - if call.accountID == 1 && call.userID == testUserID { - found = true - break - } - } - assert.True(t, found, "UpdateAffinity should be called for pinned account 1 with userID %d", testUserID) - }) - - t.Run("non-pinned user does not trigger UpdateAffinity preprocess", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil // 无亲和记录，直接进入 Layer 2 - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "pinned_users": []any{float64(99)}, // different user - }, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{1: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - _, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - - // 验证预处理阶段（beforeAffinityLookup=true）没有为 userID=42 在 account 1 上调用 UpdateAffinity - for _, call := range cache.updateAffinityCalls { - if call.beforeAffinityLookup && call.accountID == 1 && call.userID == testUserID { - t.Errorf("UpdateAffinity should NOT be called in preprocess for non-pinned user %d on account %d", testUserID, call.accountID) - } - } - }) -} - -// =========================================================================== -// TestAffinityNoSwitchError — 验证 ErrAffinityNoSwitch 行为 -// =========================================================================== - -func TestAffinityNoSwitchError(t *testing.T) { - ctx := context.Background() - testUserID := int64(42) - testClientID := "user_" + "a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2c3d4e5f6a1b2" + "_account_test" - futureTime := time.Now().Add(1 * time.Hour) - - t.Run("affinity hit but unschedulable + allow_switch=false returns ErrAffinityNoSwitch", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, - RateLimitResetAt: &futureTime, // rate limited = unschedulable for selection - Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": false, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - _, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.Error(t, err) - assert.ErrorIs(t, err, ErrAffinityNoSwitch) - }) - - t.Run("affinity hit but unschedulable + allow_switch=true continues to Layer 2", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, - RateLimitResetAt: &futureTime, // rate limited = unschedulable - Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": true, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - require.NotNil(t, result) - assert.Equal(t, int64(2), result.Account.ID, "should fall through to Layer 2 and select account 2") - }) - - t.Run("no affinity records - not affected by allow_switch", func(t *testing.T) { - cache := &affinityIntegrationCache{} // returns nil from GetAffinityAccounts - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": false, - }, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{1: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - require.NotNil(t, result) - assert.Equal(t, int64(1), result.Account.ID, "should select normally via Layer 2") - }) - - t.Run("affinity disabled account ignores allow_switch=false", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": false, - "affinity_allow_switch": false, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 2, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - require.NotNil(t, result) - require.NotNil(t, result.Account) - assert.NotEqual(t, int64(0), result.Account.ID) - }) - - t.Run("allow_switch=false + model_unsupported returns ErrAffinityNoSwitch", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": false, - }, - // model_mapping 仅包含 claude-3-haiku → 请求 claude-3-5-sonnet 将被拒绝 - Credentials: map[string]any{ - "model_mapping": map[string]any{ - "claude-3-haiku-20240307": "claude-3-haiku-20240307", - }, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - _, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.Error(t, err) - assert.ErrorIs(t, err, ErrAffinityNoSwitch) - }) - - t.Run("allow_switch=false + red_zone returns ErrAffinityNoSwitch", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - getMultiCountFunc: func(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - // userCount=5, clientCount=5, perUserCount=0 - // 账号 affinity_base=1, affinity_buffer=0 → clientCount(5) > base(1) + buffer(0) → 红区 - return 5, 5, 0, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": false, - "affinity_base": 1, - "affinity_buffer": 0, // buffer=0 → 超过 base 直接红区 - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - _, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.Error(t, err) - assert.ErrorIs(t, err, ErrAffinityNoSwitch) - }) - - t.Run("one_allow_switch=true overrides another allow_switch=false (一票放行)", func(t *testing.T) { - // 账号 1: allow_switch=false + rate limited（不可调度） - // 账号 3: allow_switch=true + rate limited（不可调度） - // 一票放行：账号 3 允许切换 → 不阻断 → 降级到 Layer 2 选中账号 2 - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1, 3}, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, - RateLimitResetAt: &futureTime, // rate limited - Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": false, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - { - ID: 3, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, - RateLimitResetAt: &futureTime, // rate limited - Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": true, - }, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - require.NotNil(t, result) - assert.Equal(t, int64(2), result.Account.ID, "一票放行: should fall through to Layer 2") - }) - - t.Run("excluded affinity account with allow_switch=false returns ErrAffinityNoSwitch", func(t *testing.T) { - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_allow_switch": false, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - excluded := map[int64]struct{}{1: {}} - _, err := svc.SelectAccountWithLoadAwareness(ctx, nil, "", "claude-3-5-sonnet-20241022", excluded, testClientID, testUserID) - require.Error(t, err) - assert.ErrorIs(t, err, ErrAffinityNoSwitch) - }) - - t.Run("disabled affinity record does not suppress sticky fallback", func(t *testing.T) { - sessionHash := "sticky-disabled-affinity" - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - sessionBindings: map[string]int64{sessionHash: 2}, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": false, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, sessionHash, "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - require.NotNil(t, result) - require.NotNil(t, result.Account) - assert.Equal(t, int64(2), result.Account.ID) - }) - - t.Run("historical affinity record with disabled account still falls back to sticky account", func(t *testing.T) { - sessionHash := "sticky-disabled-status-affinity" - cache := &affinityIntegrationCache{ - getAffinityAccountsFunc: func(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return []int64{1}, nil - }, - sessionBindings: map[string]int64{sessionHash: 2}, - } - - repo := &mockAccountRepoForPlatform{ - accounts: []Account{ - { - ID: 1, Platform: PlatformAnthropic, Priority: 1, - Status: StatusDisabled, Schedulable: true, Concurrency: 5, - Extra: map[string]any{ - "affinity_enabled": true, - }, - }, - { - ID: 2, Platform: PlatformAnthropic, Priority: 1, - Status: StatusActive, Schedulable: true, Concurrency: 5, - }, - }, - accountsByID: map[int64]*Account{}, - } - for i := range repo.accounts { - repo.accountsByID[repo.accounts[i].ID] = &repo.accounts[i] - } - - concCache := &mockConcurrencyCache{acquireResults: map[int64]bool{2: true}} - concSvc := NewConcurrencyService(concCache) - cfg := &config.Config{RunMode: config.RunModeStandard} - cfg.Gateway.Scheduling.LoadBatchEnabled = true - - svc := &GatewayService{ - accountRepo: repo, - cache: cache, - cfg: cfg, - concurrencyService: concSvc, - } - - result, err := svc.SelectAccountWithLoadAwareness(ctx, nil, sessionHash, "claude-3-5-sonnet-20241022", nil, testClientID, testUserID) - require.NoError(t, err) - require.NotNil(t, result) - require.NotNil(t, result.Account) - assert.Equal(t, int64(2), result.Account.ID) - }) -} - -// =========================================================================== -// TestAffinityMultiDimClassify — 验证 classifyByAffinityZone 与多维度计数 -// =========================================================================== - -func TestAffinityMultiDimClassify(t *testing.T) { - t.Run("multi-dim zone classification with user and client dimensions", func(t *testing.T) { - makeAWL := func(id int64, extra map[string]any, count int64) accountWithLoad { - return accountWithLoad{ - account: &Account{ID: id, Platform: PlatformAnthropic, Extra: extra}, - loadInfo: &AccountLoadInfo{AccountID: id}, - affinityCount: count, - } - } - - // Account 1: client green (3 <= 5), classified as green by old classifyByAffinityZone - // Account 2: client red (10 > 8), classified as red - accs := []accountWithLoad{ - makeAWL(1, map[string]any{ - "affinity_enabled": true, - "affinity_base": 5, - "affinity_buffer": 3, - }, 3), - makeAWL(2, map[string]any{ - "affinity_enabled": true, - "affinity_base": 5, - "affinity_buffer": 3, - }, 10), - } - - result := classifyByAffinityZone(accs) - require.Len(t, result, 1) - assert.Equal(t, int64(1), result[0].account.ID, "only green account should remain") - }) - - t.Run("all green returns all", func(t *testing.T) { - makeAWL := func(id int64, count int64) accountWithLoad { - return accountWithLoad{ - account: &Account{ - ID: id, - Platform: PlatformAnthropic, - Extra: map[string]any{ - "affinity_enabled": true, - "affinity_base": 10, - "affinity_buffer": 5, - }, - }, - loadInfo: &AccountLoadInfo{AccountID: id}, - affinityCount: count, - } - } - - accs := []accountWithLoad{ - makeAWL(1, 3), - makeAWL(2, 5), - makeAWL(3, 10), // exactly at base - } - - result := classifyByAffinityZone(accs) - require.Len(t, result, 3) - }) -} diff --git a/backend/internal/service/gateway_claude_max_response_helpers.go b/backend/internal/service/gateway_claude_max_response_helpers.go deleted file mode 100644 index a5f5f3d2d2..0000000000 --- a/backend/internal/service/gateway_claude_max_response_helpers.go +++ /dev/null @@ -1,196 +0,0 @@ -package service - -import ( - "context" - "encoding/json" - - "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" - "github.com/gin-gonic/gin" - "github.com/tidwall/sjson" -) - -type claudeMaxResponseRewriteContext struct { - Parsed *ParsedRequest - Group *Group -} - -type claudeMaxResponseRewriteContextKeyType struct{} - -var claudeMaxResponseRewriteContextKey = claudeMaxResponseRewriteContextKeyType{} - -func withClaudeMaxResponseRewriteContext(ctx context.Context, c *gin.Context, parsed *ParsedRequest) context.Context { - if ctx == nil { - ctx = context.Background() - } - value := claudeMaxResponseRewriteContext{ - Parsed: parsed, - Group: claudeMaxGroupFromGinContext(c), - } - return context.WithValue(ctx, claudeMaxResponseRewriteContextKey, value) -} - -func claudeMaxResponseRewriteContextFromContext(ctx context.Context) claudeMaxResponseRewriteContext { - if ctx == nil { - return claudeMaxResponseRewriteContext{} - } - value, _ := ctx.Value(claudeMaxResponseRewriteContextKey).(claudeMaxResponseRewriteContext) - return value -} - -func claudeMaxGroupFromGinContext(c *gin.Context) *Group { - if c == nil { - return nil - } - raw, exists := c.Get("api_key") - if !exists { - return nil - } - apiKey, ok := raw.(*APIKey) - if !ok || apiKey == nil { - return nil - } - return apiKey.Group -} - -func parsedRequestFromGinContext(c *gin.Context) *ParsedRequest { - if c == nil { - return nil - } - raw, exists := c.Get("parsed_request") - if !exists { - return nil - } - parsed, _ := raw.(*ParsedRequest) - return parsed -} - -func applyClaudeMaxSimulationToUsage(ctx context.Context, usage *ClaudeUsage, model string, accountID int64) claudeMaxCacheBillingOutcome { - var out claudeMaxCacheBillingOutcome - if usage == nil { - return out - } - rewriteCtx := claudeMaxResponseRewriteContextFromContext(ctx) - return applyClaudeMaxCacheBillingPolicyToUsage(usage, rewriteCtx.Parsed, rewriteCtx.Group, model, accountID) -} - -func applyClaudeMaxSimulationToUsageJSONMap(ctx context.Context, usageObj map[string]any, model string, accountID int64) claudeMaxCacheBillingOutcome { - var out claudeMaxCacheBillingOutcome - if usageObj == nil { - return out - } - usage := claudeUsageFromJSONMap(usageObj) - out = applyClaudeMaxSimulationToUsage(ctx, &usage, model, accountID) - if out.Simulated { - rewriteClaudeUsageJSONMap(usageObj, usage) - } - return out -} - -func rewriteClaudeUsageJSONBytes(body []byte, usage ClaudeUsage) []byte { - updated := body - var err error - - updated, err = sjson.SetBytes(updated, "usage.input_tokens", usage.InputTokens) - if err != nil { - return body - } - updated, err = sjson.SetBytes(updated, "usage.cache_creation_input_tokens", usage.CacheCreationInputTokens) - if err != nil { - return body - } - updated, err = sjson.SetBytes(updated, "usage.cache_creation.ephemeral_5m_input_tokens", usage.CacheCreation5mTokens) - if err != nil { - return body - } - updated, err = sjson.SetBytes(updated, "usage.cache_creation.ephemeral_1h_input_tokens", usage.CacheCreation1hTokens) - if err != nil { - return body - } - return updated -} - -func claudeUsageFromJSONMap(usageObj map[string]any) ClaudeUsage { - var usage ClaudeUsage - if usageObj == nil { - return usage - } - - usage.InputTokens = usageIntFromAny(usageObj["input_tokens"]) - usage.OutputTokens = usageIntFromAny(usageObj["output_tokens"]) - usage.CacheCreationInputTokens = usageIntFromAny(usageObj["cache_creation_input_tokens"]) - usage.CacheReadInputTokens = usageIntFromAny(usageObj["cache_read_input_tokens"]) - - if ccObj, ok := usageObj["cache_creation"].(map[string]any); ok { - usage.CacheCreation5mTokens = usageIntFromAny(ccObj["ephemeral_5m_input_tokens"]) - usage.CacheCreation1hTokens = usageIntFromAny(ccObj["ephemeral_1h_input_tokens"]) - } - return usage -} - -func rewriteClaudeUsageJSONMap(usageObj map[string]any, usage ClaudeUsage) { - if usageObj == nil { - return - } - usageObj["input_tokens"] = usage.InputTokens - usageObj["cache_creation_input_tokens"] = usage.CacheCreationInputTokens - - ccObj, _ := usageObj["cache_creation"].(map[string]any) - if ccObj == nil { - ccObj = make(map[string]any, 2) - usageObj["cache_creation"] = ccObj - } - ccObj["ephemeral_5m_input_tokens"] = usage.CacheCreation5mTokens - ccObj["ephemeral_1h_input_tokens"] = usage.CacheCreation1hTokens -} - -func usageIntFromAny(v any) int { - switch value := v.(type) { - case int: - return value - case int64: - return int(value) - case float64: - return int(value) - case json.Number: - if n, err := value.Int64(); err == nil { - return int(n) - } - } - return 0 -} - -// setupClaudeMaxStreamingHook 为 Antigravity 流式路径设置 SSE usage 改写 hook。 -func setupClaudeMaxStreamingHook(c *gin.Context, processor *antigravity.StreamingProcessor, originalModel string, accountID int64) { - group := claudeMaxGroupFromGinContext(c) - parsed := parsedRequestFromGinContext(c) - if !shouldApplyClaudeMaxBillingRulesForUsage(group, originalModel, parsed) { - return - } - processor.SetUsageMapHook(func(usageMap map[string]any) { - svcUsage := claudeUsageFromJSONMap(usageMap) - outcome := applyClaudeMaxCacheBillingPolicyToUsage(&svcUsage, parsed, group, originalModel, accountID) - if outcome.Simulated { - rewriteClaudeUsageJSONMap(usageMap, svcUsage) - } - }) -} - -// applyClaudeMaxNonStreamingRewrite 为 Antigravity 非流式路径改写响应体中的 usage。 -func applyClaudeMaxNonStreamingRewrite(c *gin.Context, claudeResp []byte, agUsage *antigravity.ClaudeUsage, originalModel string, accountID int64) []byte { - group := claudeMaxGroupFromGinContext(c) - parsed := parsedRequestFromGinContext(c) - if !shouldApplyClaudeMaxBillingRulesForUsage(group, originalModel, parsed) { - return claudeResp - } - svcUsage := &ClaudeUsage{ - InputTokens: agUsage.InputTokens, - OutputTokens: agUsage.OutputTokens, - CacheCreationInputTokens: agUsage.CacheCreationInputTokens, - CacheReadInputTokens: agUsage.CacheReadInputTokens, - } - outcome := applyClaudeMaxCacheBillingPolicyToUsage(svcUsage, parsed, group, originalModel, accountID) - if outcome.Simulated { - return rewriteClaudeUsageJSONBytes(claudeResp, *svcUsage) - } - return claudeResp -} diff --git a/backend/internal/service/gateway_hotpath_optimization_test.go b/backend/internal/service/gateway_hotpath_optimization_test.go index af108c6fc2..86c457434e 100644 --- a/backend/internal/service/gateway_hotpath_optimization_test.go +++ b/backend/internal/service/gateway_hotpath_optimization_test.go @@ -143,28 +143,6 @@ func (s *stickyGatewayCacheHotpathStub) RefreshSessionTTL(ctx context.Context, g func (s *stickyGatewayCacheHotpathStub) DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error { return nil } -func (s *stickyGatewayCacheHotpathStub) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (s *stickyGatewayCacheHotpathStub) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (s *stickyGatewayCacheHotpathStub) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (s *stickyGatewayCacheHotpathStub) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (s *stickyGatewayCacheHotpathStub) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (s *stickyGatewayCacheHotpathStub) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (s *stickyGatewayCacheHotpathStub) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - func (s *modelsListAccountRepoStub) ListSchedulableByGroupID(ctx context.Context, groupID int64) ([]Account, error) { s.listByGroupCalls.Add(1) if s.err != nil { diff --git a/backend/internal/service/gateway_multiplatform_test.go b/backend/internal/service/gateway_multiplatform_test.go index 9f584f87cb..b8d3538d5a 100644 --- a/backend/internal/service/gateway_multiplatform_test.go +++ b/backend/internal/service/gateway_multiplatform_test.go @@ -235,28 +235,6 @@ func (m *mockGatewayCacheForPlatform) DeleteSessionAccountID(ctx context.Context return nil } -func (m *mockGatewayCacheForPlatform) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (m *mockGatewayCacheForPlatform) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (m *mockGatewayCacheForPlatform) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (m *mockGatewayCacheForPlatform) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (m *mockGatewayCacheForPlatform) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (m *mockGatewayCacheForPlatform) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (m *mockGatewayCacheForPlatform) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - type mockGroupRepoForGateway struct { groups map[int64]*Group getByIDCalls int diff --git a/backend/internal/service/gateway_record_usage_claude_max_test.go b/backend/internal/service/gateway_record_usage_claude_max_test.go deleted file mode 100644 index 3cd8693848..0000000000 --- a/backend/internal/service/gateway_record_usage_claude_max_test.go +++ /dev/null @@ -1,199 +0,0 @@ -package service - -import ( - "context" - "testing" - "time" - - "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/stretchr/testify/require" -) - -type usageLogRepoRecordUsageStub struct { - UsageLogRepository - - last *UsageLog - inserted bool - err error -} - -func (s *usageLogRepoRecordUsageStub) Create(_ context.Context, log *UsageLog) (bool, error) { - copied := *log - s.last = &copied - return s.inserted, s.err -} - -func newGatewayServiceForRecordUsageTest(repo UsageLogRepository) *GatewayService { - return &GatewayService{ - usageLogRepo: repo, - billingService: NewBillingService(&config.Config{}, nil), - cfg: &config.Config{RunMode: config.RunModeSimple}, - deferredService: &DeferredService{}, - } -} - -func TestRecordUsage_SimulateClaudeMaxEnabled_ProjectsUsageAndSkipsTTLOverride(t *testing.T) { - repo := &usageLogRepoRecordUsageStub{inserted: true} - svc := newGatewayServiceForRecordUsageTest(repo) - - groupID := int64(11) - input := &RecordUsageInput{ - Result: &ForwardResult{ - RequestID: "req-sim-1", - Model: "claude-sonnet-4", - Duration: time.Second, - Usage: ClaudeUsage{ - InputTokens: 160, - }, - }, - ParsedRequest: &ParsedRequest{ - Model: "claude-sonnet-4", - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": "long cached context for prior turns", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - map[string]any{ - "type": "text", - "text": "please summarize the logs and provide root cause analysis", - }, - }, - }, - }, - }, - APIKey: &APIKey{ - ID: 1, - GroupID: &groupID, - Group: &Group{ - ID: groupID, - Platform: PlatformAnthropic, - RateMultiplier: 1, - SimulateClaudeMaxEnabled: true, - }, - }, - User: &User{ID: 2}, - Account: &Account{ - ID: 3, - Platform: PlatformAnthropic, - Type: AccountTypeOAuth, - Extra: map[string]any{ - "cache_ttl_override_enabled": true, - "cache_ttl_override_target": "5m", - }, - }, - } - - err := svc.RecordUsage(context.Background(), input) - require.NoError(t, err) - require.NotNil(t, repo.last) - - log := repo.last - require.Equal(t, 80, log.InputTokens) - require.Equal(t, 80, log.CacheCreationTokens) - require.Equal(t, 0, log.CacheCreation5mTokens) - require.Equal(t, 80, log.CacheCreation1hTokens) - require.False(t, log.CacheTTLOverridden, "simulate outcome should skip account ttl override") -} - -func TestRecordUsage_SimulateClaudeMaxDisabled_AppliesTTLOverride(t *testing.T) { - repo := &usageLogRepoRecordUsageStub{inserted: true} - svc := newGatewayServiceForRecordUsageTest(repo) - - groupID := int64(12) - input := &RecordUsageInput{ - Result: &ForwardResult{ - RequestID: "req-sim-2", - Model: "claude-sonnet-4", - Duration: time.Second, - Usage: ClaudeUsage{ - InputTokens: 40, - CacheCreationInputTokens: 120, - CacheCreation1hTokens: 120, - }, - }, - APIKey: &APIKey{ - ID: 2, - GroupID: &groupID, - Group: &Group{ - ID: groupID, - Platform: PlatformAnthropic, - RateMultiplier: 1, - SimulateClaudeMaxEnabled: false, - }, - }, - User: &User{ID: 3}, - Account: &Account{ - ID: 4, - Platform: PlatformAnthropic, - Type: AccountTypeOAuth, - Extra: map[string]any{ - "cache_ttl_override_enabled": true, - "cache_ttl_override_target": "5m", - }, - }, - } - - err := svc.RecordUsage(context.Background(), input) - require.NoError(t, err) - require.NotNil(t, repo.last) - - log := repo.last - require.Equal(t, 120, log.CacheCreationTokens) - require.Equal(t, 120, log.CacheCreation5mTokens) - require.Equal(t, 0, log.CacheCreation1hTokens) - require.True(t, log.CacheTTLOverridden) -} - -func TestRecordUsage_SimulateClaudeMaxEnabled_ExistingCacheCreationBypassesSimulation(t *testing.T) { - repo := &usageLogRepoRecordUsageStub{inserted: true} - svc := newGatewayServiceForRecordUsageTest(repo) - - groupID := int64(13) - input := &RecordUsageInput{ - Result: &ForwardResult{ - RequestID: "req-sim-3", - Model: "claude-sonnet-4", - Duration: time.Second, - Usage: ClaudeUsage{ - InputTokens: 20, - CacheCreationInputTokens: 120, - CacheCreation5mTokens: 120, - }, - }, - APIKey: &APIKey{ - ID: 3, - GroupID: &groupID, - Group: &Group{ - ID: groupID, - Platform: PlatformAnthropic, - RateMultiplier: 1, - SimulateClaudeMaxEnabled: true, - }, - }, - User: &User{ID: 4}, - Account: &Account{ - ID: 5, - Platform: PlatformAnthropic, - Type: AccountTypeOAuth, - Extra: map[string]any{ - "cache_ttl_override_enabled": true, - "cache_ttl_override_target": "5m", - }, - }, - } - - err := svc.RecordUsage(context.Background(), input) - require.NoError(t, err) - require.NotNil(t, repo.last) - - log := repo.last - require.Equal(t, 20, log.InputTokens) - require.Equal(t, 120, log.CacheCreation5mTokens) - require.Equal(t, 0, log.CacheCreation1hTokens) - require.Equal(t, 120, log.CacheCreationTokens) - require.False(t, log.CacheTTLOverridden, "existing cache_creation with SimulateClaudeMax enabled should skip account ttl override") -} diff --git a/backend/internal/service/gateway_response_usage_sync_test.go b/backend/internal/service/gateway_response_usage_sync_test.go deleted file mode 100644 index 445ee8ad53..0000000000 --- a/backend/internal/service/gateway_response_usage_sync_test.go +++ /dev/null @@ -1,170 +0,0 @@ -package service - -import ( - "bytes" - "context" - "encoding/json" - "net/http" - "net/http/httptest" - "testing" - - "github.com/Wei-Shaw/sub2api/internal/config" - "github.com/gin-gonic/gin" - "github.com/stretchr/testify/require" - "github.com/tidwall/gjson" -) - -func TestHandleNonStreamingResponse_UsageAlignedWithClaudeMaxSimulation(t *testing.T) { - gin.SetMode(gin.TestMode) - - svc := &GatewayService{ - cfg: &config.Config{}, - rateLimitService: &RateLimitService{}, - } - - account := &Account{ - ID: 11, - Platform: PlatformAnthropic, - Type: AccountTypeOAuth, - Extra: map[string]any{ - "cache_ttl_override_enabled": true, - "cache_ttl_override_target": "5m", - }, - } - group := &Group{ - ID: 99, - Platform: PlatformAnthropic, - SimulateClaudeMaxEnabled: true, - } - parsed := &ParsedRequest{ - Model: "claude-sonnet-4", - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": "long cached context", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - map[string]any{ - "type": "text", - "text": "new user question", - }, - }, - }, - }, - } - - upstreamBody := []byte(`{"id":"msg_1","model":"claude-sonnet-4","usage":{"input_tokens":120,"output_tokens":8}}`) - resp := &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: ioNopCloserBytes(upstreamBody), - } - - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(nil)) - c.Set("api_key", &APIKey{Group: group}) - requestCtx := withClaudeMaxResponseRewriteContext(context.Background(), c, parsed) - - usage, err := svc.handleNonStreamingResponse(requestCtx, resp, c, account, "claude-sonnet-4", "claude-sonnet-4") - require.NoError(t, err) - require.NotNil(t, usage) - - var rendered struct { - Usage ClaudeUsage `json:"usage"` - } - require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &rendered)) - rendered.Usage.CacheCreation5mTokens = int(gjson.GetBytes(rec.Body.Bytes(), "usage.cache_creation.ephemeral_5m_input_tokens").Int()) - rendered.Usage.CacheCreation1hTokens = int(gjson.GetBytes(rec.Body.Bytes(), "usage.cache_creation.ephemeral_1h_input_tokens").Int()) - - require.Equal(t, rendered.Usage.InputTokens, usage.InputTokens) - require.Equal(t, rendered.Usage.OutputTokens, usage.OutputTokens) - require.Equal(t, rendered.Usage.CacheCreationInputTokens, usage.CacheCreationInputTokens) - require.Equal(t, rendered.Usage.CacheCreation5mTokens, usage.CacheCreation5mTokens) - require.Equal(t, rendered.Usage.CacheCreation1hTokens, usage.CacheCreation1hTokens) - require.Equal(t, rendered.Usage.CacheReadInputTokens, usage.CacheReadInputTokens) - - require.Greater(t, usage.CacheCreation1hTokens, 0) - require.Equal(t, 0, usage.CacheCreation5mTokens) - require.Less(t, usage.InputTokens, 120) -} - -func TestHandleNonStreamingResponse_ClaudeMaxDisabled_NoSimulationIntercept(t *testing.T) { - gin.SetMode(gin.TestMode) - - svc := &GatewayService{ - cfg: &config.Config{}, - rateLimitService: &RateLimitService{}, - } - - account := &Account{ - ID: 12, - Platform: PlatformAnthropic, - Type: AccountTypeOAuth, - Extra: map[string]any{ - "cache_ttl_override_enabled": true, - "cache_ttl_override_target": "5m", - }, - } - group := &Group{ - ID: 100, - Platform: PlatformAnthropic, - SimulateClaudeMaxEnabled: false, - } - parsed := &ParsedRequest{ - Model: "claude-sonnet-4", - Messages: []any{ - map[string]any{ - "role": "user", - "content": []any{ - map[string]any{ - "type": "text", - "text": "long cached context", - "cache_control": map[string]any{"type": "ephemeral"}, - }, - map[string]any{ - "type": "text", - "text": "new user question", - }, - }, - }, - }, - } - - upstreamBody := []byte(`{"id":"msg_2","model":"claude-sonnet-4","usage":{"input_tokens":120,"output_tokens":8}}`) - resp := &http.Response{ - StatusCode: http.StatusOK, - Header: http.Header{"Content-Type": []string{"application/json"}}, - Body: ioNopCloserBytes(upstreamBody), - } - - rec := httptest.NewRecorder() - c, _ := gin.CreateTestContext(rec) - c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(nil)) - c.Set("api_key", &APIKey{Group: group}) - requestCtx := withClaudeMaxResponseRewriteContext(context.Background(), c, parsed) - - usage, err := svc.handleNonStreamingResponse(requestCtx, resp, c, account, "claude-sonnet-4", "claude-sonnet-4") - require.NoError(t, err) - require.NotNil(t, usage) - - require.Equal(t, 120, usage.InputTokens) - require.Equal(t, 0, usage.CacheCreationInputTokens) - require.Equal(t, 0, usage.CacheCreation5mTokens) - require.Equal(t, 0, usage.CacheCreation1hTokens) -} - -func ioNopCloserBytes(b []byte) *readCloserFromBytes { - return &readCloserFromBytes{Reader: bytes.NewReader(b)} -} - -type readCloserFromBytes struct { - *bytes.Reader -} - -func (r *readCloserFromBytes) Close() error { - return nil -} diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 252c8e3cd9..aae4883912 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -42,9 +42,8 @@ import ( const ( claudeAPIURL = "https://api.anthropic.com/v1/messages?beta=true" claudeAPICountTokensURL = "https://api.anthropic.com/v1/messages/count_tokens?beta=true" - stickySessionTTL = time.Hour // 粘性会话TTL - ClientAffinityTTL = 24 * time.Hour // 客户端亲和TTL - defaultMaxLineSize = 500 * 1024 * 1024 + stickySessionTTL = time.Hour // 粘性会话TTL + defaultMaxLineSize = 500 * 1024 * 1024 // Canonical Claude Code banner. Keep it EXACT (no trailing whitespace/newlines) // to match real Claude CLI traffic as closely as possible. When we need a visual // separator between system blocks, we add "\n\n" at concatenation time. @@ -68,21 +67,14 @@ const ( MediaTypePrompt = "prompt" ) -const ( - claudeMaxMessageOverheadTokens = 3 - claudeMaxBlockOverheadTokens = 1 - claudeMaxUnknownContentTokens = 4 -) - // ForceCacheBillingContextKey 强制缓存计费上下文键 // 用于粘性会话切换时，将 input_tokens 转为 cache_read_input_tokens 计费 type forceCacheBillingKeyType struct{} // accountWithLoad 账号与负载信息的组合，用于负载感知调度 type accountWithLoad struct { - account *Account - loadInfo *AccountLoadInfo - affinityCount int64 // 亲和客户端数量（反向索引），越少越优先 + account *Account + loadInfo *AccountLoadInfo } var ForceCacheBillingContextKey = forceCacheBillingKeyType{} @@ -346,10 +338,6 @@ var ( sseDataRe = regexp.MustCompile(`^data:\s*`) claudeCliUserAgentRe = regexp.MustCompile(`^claude-cli/\d+\.\d+\.\d+`) - // clientIDFromMetadataRegex 从 metadata.user_id 中提取客户端 ID（64位 hex） - // 格式: user_{64位hex}_account_... - clientIDFromMetadataRegex = regexp.MustCompile(`^user_([a-f0-9]{64})_account_`) - // claudeCodePromptPrefixes 用于检测 Claude Code 系统提示词的前缀列表 // 支持多种变体：标准版、Agent SDK 版、Explore Agent 版、Compact 版等 // 注意：前缀之间不应存在包含关系，否则会导致冗余匹配 @@ -367,12 +355,6 @@ var ErrNoAvailableAccounts = errors.New("no available accounts") // ErrClaudeCodeOnly 表示分组仅允许 Claude Code 客户端访问 var ErrClaudeCodeOnly = errors.New("this group only allows Claude Code clients") -// ErrAffinityNoSwitch 表示亲和账号不可用且不允许切换到其他账号 -var ErrAffinityNoSwitch = errors.New("affinity account unavailable and switching is disabled") - -// ErrAffinityLimitExceeded 表示亲和客户端限制已达上限 -var ErrAffinityLimitExceeded = errors.New("affinity client limit exceeded") - // allowedHeaders 白名单headers（参考CRS项目） var allowedHeaders = map[string]bool{ "accept": true, @@ -416,39 +398,6 @@ type GatewayCache interface { // DeleteSessionAccountID 删除粘性会话绑定，用于账号不可用时主动清理 // Delete sticky session binding, used to proactively clean up when account becomes unavailable DeleteSessionAccountID(ctx context.Context, groupID int64, sessionHash string) error - - // GetAffinityAccounts 获取亲和账号列表（按最近使用降序），同时清理过期成员 - GetAffinityAccounts(ctx context.Context, groupID int64, userID int64, clientID string, ttl time.Duration) ([]int64, error) - // UpdateAffinity 添加/更新亲和关系（更新 score 为当前时间戳，刷新 key TTL） - UpdateAffinity(ctx context.Context, groupID int64, userID int64, clientID string, accountID int64, ttl time.Duration) error - // GetAccountAffinityCountBatch 批量获取账号的亲和成员数量（惰性清理过期成员） - GetAccountAffinityCountBatch(ctx context.Context, groupID int64, accountIDs []int64, ttl time.Duration) (map[int64]int64, error) - // GetAccountAffinityClientsBatch 批量获取每个账号跨所有分组的亲和成员列表（去重） - // accountGroups: map[accountID][]groupID - // 返回值成员格式为 {userID}/{clientID} - GetAccountAffinityClientsBatch(ctx context.Context, accountGroups map[int64][]int64, ttl time.Duration) (map[int64][]string, error) - // GetAccountAffinityClientsWithScores 获取单个账号跨所有分组的亲和客户端列表（含最后活跃时间） - GetAccountAffinityClientsWithScores(ctx context.Context, accountID int64, groupIDs []int64, ttl time.Duration) ([]AffinityClient, error) - // ClearAccountAffinity 清除指定账号在所有分组的亲和记录（正向+反向索引） - // 用于账号关闭亲和时立即清理旧绑定 - ClearAccountAffinity(ctx context.Context, accountID int64, groupIDs []int64) error - // GetAffinityMultiCount 获取账号的多维度亲和计数 - // 返回: uniqueUsers, uniqueClients, perUserClients - GetAffinityMultiCount(ctx context.Context, groupID int64, accountID int64, targetUserID int64, ttl time.Duration) (users, clients, perUser int64, err error) -} - -// AffinityClient 亲和客户端信息（含用户 ID 和最后活跃时间） -type AffinityClient struct { - UserID int64 `json:"user_id"` - ClientID string `json:"client_id"` - LastActive time.Time `json:"last_active"` -} - -// SortAffinityClients 按最后活跃时间降序排序 -func SortAffinityClients(clients []AffinityClient) { - sort.Slice(clients, func(i, j int) bool { - return clients[i].LastActive.After(clients[j].LastActive) - }) } // derefGroupID safely dereferences *int64 to int64, returning 0 if nil @@ -519,20 +468,6 @@ func shouldClearStickySession(account *Account, requestedModel string) bool { return false } -// extractClientIDFromMetadata 从 metadata.user_id 中提取客户端 ID（64位 hex）。 -// 格式: user_{64位hex}_account_..._session_... -// 返回空字符串表示无法提取（非 Claude Code/Console 客户端）。 -func extractClientIDFromMetadata(metadataUserID string) string { - if metadataUserID == "" { - return "" - } - matches := clientIDFromMetadataRegex.FindStringSubmatch(metadataUserID) - if matches == nil { - return "" - } - return matches[1] -} - type AccountWaitPlan struct { AccountID int64 MaxConcurrency int @@ -1318,10 +1253,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } } - // 提取客户端 ID（用于客户端亲和调度） - affinityClientID := extractClientIDFromMetadata(metadataUserID) - affinityUserID := sub2apiUserID - if s.debugModelRoutingEnabled() && requestedModel != "" { groupPlatform := "" if group != nil { @@ -1343,10 +1274,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if err != nil { return nil, err } - if shouldFilterAccountWithoutClientID(account, affinityClientID) { - localExcluded[account.ID] = struct{}{} - continue - } result, err := s.tryAcquireAccountSlot(ctx, account.ID, account.Concurrency) if err == nil && result.Acquired { @@ -1408,7 +1335,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if err != nil { return nil, err } - accounts = filterAccountsWithoutClientID(accounts, affinityClientID) if len(accounts) == 0 { return nil, ErrNoAvailableAccounts } @@ -1427,19 +1353,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro _, excluded := excludedIDs[accountID] return excluded } - affinityFlow := newGatewayAffinityFlow( - s, - ctx, - groupID, - sessionHash, - requestedModel, - affinityClientID, - affinityUserID, - platform, - useMixed, - accountByID, - isExcluded, - ) // 获取模型路由配置（仅 anthropic 平台） var routingAccountIDs []int64 @@ -1602,10 +1515,7 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } if len(routingAvailable) > 0 { - // 批量获取亲和客户端数量 - s.populateAffinityCounts(ctx, routingAvailable, derefGroupID(groupID)) - - // 排序：优先级 > 负载率 > 亲和客户端数 > 最后使用时间 + // 排序：优先级 > 负载率 > 最后使用时间 sort.SliceStable(routingAvailable, func(i, j int) bool { a, b := routingAvailable[i], routingAvailable[j] if a.account.Priority != b.account.Priority { @@ -1614,9 +1524,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if a.loadInfo.LoadRate != b.loadInfo.LoadRate { return a.loadInfo.LoadRate < b.loadInfo.LoadRate } - if a.affinityCount != b.affinityCount { - return a.affinityCount < b.affinityCount - } switch { case a.account.LastUsedAt == nil && b.account.LastUsedAt != nil: return true @@ -1642,9 +1549,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if sessionHash != "" && s.cache != nil { _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, item.account.ID, stickySessionTTL) } - if affinityClientID != "" && affinityUserID > 0 && s.cache != nil && item.account.IsAffinityEnabled() { - _ = s.cache.UpdateAffinity(ctx, derefGroupID(groupID), affinityUserID, affinityClientID, item.account.ID, ClientAffinityTTL) - } if s.debugModelRoutingEnabled() { logger.LegacyPrintf("service.gateway", "[ModelRoutingDebug] routed select: group_id=%v model=%s session=%s account=%d", derefGroupID(groupID), requestedModel, shortSessionHash(sessionHash), item.account.ID) } @@ -1682,22 +1586,8 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } } - // ============ Layer 1.3: 用户亲和预处理（pinned_users 自动注入） ============ - affinityFlow.preprocessPinnedUsers(accounts) - - // ============ Layer 1.4: 客户端亲和调度（优先于粘性会话） ============ - affinityHit := false - if affinityResult, hit, err := affinityFlow.trySelectAffinityAccount(); err != nil { - return nil, err - } else { - affinityHit = hit - if affinityResult != nil { - return affinityResult, nil - } - } - - // ============ Layer 1.5: 粘性会话（仅在无模型路由配置 且 亲和未命中时生效） ============ - if !affinityHit && len(routingAccountIDs) == 0 && sessionHash != "" && stickyAccountID > 0 && !isExcluded(stickyAccountID) { + // ============ Layer 1.5: 粘性会话（仅在无模型路由配置时生效） ============ + if len(routingAccountIDs) == 0 && sessionHash != "" && stickyAccountID > 0 && !isExcluded(stickyAccountID) { accountID := stickyAccountID if accountID > 0 && !isExcluded(accountID) { account, ok := accountByID[accountID] @@ -1803,9 +1693,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro loadMap, err := s.concurrencyService.GetAccountsLoadBatch(ctx, accountLoads) if err != nil { if result, ok := s.tryAcquireByLegacyOrder(ctx, candidates, groupID, sessionHash, preferOAuth); ok { - if affinityClientID != "" && affinityUserID > 0 && s.cache != nil && result.Account != nil && result.Account.IsAffinityEnabled() { - _ = s.cache.UpdateAffinity(ctx, derefGroupID(groupID), affinityUserID, affinityClientID, result.Account.ID, ClientAffinityTTL) - } return result, nil } } else { @@ -1823,37 +1710,13 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro } } - // 批量获取亲和客户端数量（用于均衡分配新客户端） - s.populateAffinityCounts(ctx, available, derefGroupID(groupID)) - - // 分层过滤选择：优先级 → 亲和三区 → 负载率 → 亲和客户端数 → LRU + // 分层过滤选择：优先级 → 负载率 → LRU for len(available) > 0 { // 1. 取优先级最小的集合 candidates := filterByMinPriority(available) - // 2. 按亲和三区过滤：绿区优先 → 黄区降级 → 红区移除（在同优先级内） - candidates = classifyByAffinityZone(candidates) - if len(candidates) == 0 { - // 当前优先级组全部在红区，移除后回退到下一优先级组 - minPri := available[0].account.Priority - for _, a := range available[1:] { - if a.account.Priority < minPri { - minPri = a.account.Priority - } - } - newAvailable := make([]accountWithLoad, 0, len(available)) - for _, a := range available { - if a.account.Priority != minPri { - newAvailable = append(newAvailable, a) - } - } - available = newAvailable - continue - } - // 3. 取负载率最低的集合 + // 2. 取负载率最低的集合 candidates = filterByMinLoadRate(candidates) - // 3. 取亲和客户端数最少的集合 - candidates = filterByMinAffinityCount(candidates) - // 4. LRU 选择最久未用的账号 + // 3. LRU 选择最久未用的账号 selected := selectByLRU(candidates, preferOAuth) if selected == nil { break @@ -1868,10 +1731,6 @@ func (s *GatewayService) SelectAccountWithLoadAwareness(ctx context.Context, gro if sessionHash != "" && s.cache != nil { _ = s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), sessionHash, selected.account.ID, stickySessionTTL) } - // 更新亲和关系 - if affinityClientID != "" && affinityUserID > 0 && s.cache != nil && selected.account.IsAffinityEnabled() { - _ = s.cache.UpdateAffinity(ctx, derefGroupID(groupID), affinityUserID, affinityClientID, selected.account.ID, ClientAffinityTTL) - } return &AccountSelectionResult{ Account: selected.account, Acquired: true, @@ -2629,36 +2488,6 @@ func (s *GatewayService) getSchedulableAccount(ctx context.Context, accountID in return s.accountRepo.GetByID(ctx, accountID) } -// populateAffinityCounts 批量获取账号的亲和客户端数量并填入 accountWithLoad 切片。 -// 仅当存在开启了客户端亲和的账号时才查询 Redis，否则跳过。 -func (s *GatewayService) populateAffinityCounts(ctx context.Context, accounts []accountWithLoad, groupID int64) { - if s.cache == nil || len(accounts) == 0 { - return - } - // 快速检查：是否有任何账号开启了亲和 - hasAffinity := false - for _, acc := range accounts { - if acc.account.IsAffinityEnabled() { - hasAffinity = true - break - } - } - if !hasAffinity { - return - } - accountIDs := make([]int64, len(accounts)) - for i, acc := range accounts { - accountIDs[i] = acc.account.ID - } - countMap, err := s.cache.GetAccountAffinityCountBatch(ctx, groupID, accountIDs, ClientAffinityTTL) - if err != nil { - return // 查询失败不影响调度，affinityCount 保持 0 - } - for i := range accounts { - accounts[i].affinityCount = countMap[accounts[i].account.ID] - } -} - // filterByMinPriority 过滤出优先级最小的账号集合 func filterByMinPriority(accounts []accountWithLoad) []accountWithLoad { if len(accounts) == 0 { @@ -2699,64 +2528,6 @@ func filterByMinLoadRate(accounts []accountWithLoad) []accountWithLoad { return result } -// filterByMinAffinityCount 过滤出亲和客户端数最少的账号集合 -func filterByMinAffinityCount(accounts []accountWithLoad) []accountWithLoad { - if len(accounts) == 0 { - return accounts - } - minCount := accounts[0].affinityCount - for _, acc := range accounts[1:] { - if acc.affinityCount < minCount { - minCount = acc.affinityCount - } - } - result := make([]accountWithLoad, 0, len(accounts)) - for _, acc := range accounts { - if acc.affinityCount == minCount { - result = append(result, acc) - } - } - return result -} - -// classifyByAffinityZone 按亲和分区对候选账号进行分类。 -// 返回值：仅绿区账号（有绿区时），否则返回黄区账号。红区账号被移除。 -// 如果没有任何账号开启了亲和三区配置（即 affinity_base <= 0），则原样返回所有账号。 -func classifyByAffinityZone(accounts []accountWithLoad) []accountWithLoad { - if len(accounts) == 0 { - return accounts - } - // 快速检查：是否有任何账号配置了 affinity_base - hasZoneConfig := false - for _, acc := range accounts { - if acc.account.IsAffinityEnabled() && acc.account.GetAffinityBase() > 0 { - hasZoneConfig = true - break - } - } - if !hasZoneConfig { - return accounts - } - - greens := make([]accountWithLoad, 0, len(accounts)) - yellows := make([]accountWithLoad, 0, len(accounts)) - for _, acc := range accounts { - zone := acc.account.GetAffinityZone(acc.affinityCount) - switch zone { - case AffinityZoneGreen: - greens = append(greens, acc) - case AffinityZoneYellow: - yellows = append(yellows, acc) - case AffinityZoneRed: - // 红区：移除，不参与调度 - } - } - if len(greens) > 0 { - return greens - } - return yellows -} - // selectByLRU 从集合中选择最久未用的账号 // 如果有多个账号具有相同的最小 LastUsedAt，则随机选择一个 func selectByLRU(accounts []accountWithLoad, preferOAuth bool) *accountWithLoad { @@ -4829,7 +4600,6 @@ func (s *GatewayService) Forward(ctx context.Context, c *gin.Context, account *A } // 处理正常响应 - ctx = withClaudeMaxResponseRewriteContext(ctx, c, parsed) // 触发上游接受回调（提前释放串行锁，不等流完成） if parsed.OnUpstreamAccepted != nil { @@ -7029,7 +6799,6 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http needModelReplace := originalModel != mappedModel clientDisconnected := false // 客户端断开标志，断开后继续读取上游以获取完整usage sawTerminalEvent := false - skipAccountTTLOverride := false pendingEventLines := make([]string, 0, 4) @@ -7091,25 +6860,17 @@ func (s *GatewayService) handleStreamingResponse(ctx context.Context, resp *http if msg, ok := event["message"].(map[string]any); ok { if u, ok := msg["usage"].(map[string]any); ok { eventChanged = reconcileCachedTokens(u) || eventChanged - claudeMaxOutcome := applyClaudeMaxSimulationToUsageJSONMap(ctx, u, originalModel, account.ID) - if claudeMaxOutcome.Simulated { - skipAccountTTLOverride = true - } } } } if eventType == "message_delta" { if u, ok := event["usage"].(map[string]any); ok { eventChanged = reconcileCachedTokens(u) || eventChanged - claudeMaxOutcome := applyClaudeMaxSimulationToUsageJSONMap(ctx, u, originalModel, account.ID) - if claudeMaxOutcome.Simulated { - skipAccountTTLOverride = true - } } } // Cache TTL Override: 重写 SSE 事件中的 cache_creation 分类 - if account.IsCacheTTLOverrideEnabled() && !skipAccountTTLOverride { + if account.IsCacheTTLOverrideEnabled() { overrideTarget := account.GetCacheTTLOverrideTarget() if eventType == "message_start" { if msg, ok := event["message"].(map[string]any); ok { @@ -7541,13 +7302,8 @@ func (s *GatewayService) handleNonStreamingResponse(ctx context.Context, resp *h } } - claudeMaxOutcome := applyClaudeMaxSimulationToUsage(ctx, &response.Usage, originalModel, account.ID) - if claudeMaxOutcome.Simulated { - body = rewriteClaudeUsageJSONBytes(body, response.Usage) - } - // Cache TTL Override: 重写 non-streaming 响应中的 cache_creation 分类 - if account.IsCacheTTLOverrideEnabled() && !claudeMaxOutcome.Simulated { + if account.IsCacheTTLOverrideEnabled() { overrideTarget := account.GetCacheTTLOverrideTarget() if applyCacheTTLOverride(&response.Usage, overrideTarget) { // 同步更新 body JSON 中的嵌套 cache_creation 对象 @@ -7925,11 +7681,10 @@ func writeUsageLogBestEffort(ctx context.Context, repo UsageLogRepository, usage // recordUsageOpts 内部选项，参数化 RecordUsage 与 RecordUsageWithLongContext 的差异点。 type recordUsageOpts struct { - // Claude Max 策略所需的 ParsedRequest（可选，仅 Claude 路径传入） + // ParsedRequest（可选，仅 Claude 路径传入） ParsedRequest *ParsedRequest // EnableClaudePath 启用 Claude 路径特有逻辑： - // - Claude Max 缓存计费策略 // - Sora 媒体类型分支（image/video/prompt） // - MediaType 字段写入使用日志 EnableClaudePath bool @@ -8022,7 +7777,6 @@ type recordUsageCoreInput struct { // recordUsageCore 是 RecordUsage 和 RecordUsageWithLongContext 的统一实现。 // opts 中的字段控制两者之间的差异行为： -// - ParsedRequest != nil → 启用 Claude Max 缓存计费策略 // - EnableSoraMedia → 启用 Sora MediaType 分支（image/video/prompt） // - LongContextThreshold > 0 → Token 计费回退走 CalculateCostWithLongContext func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsageCoreInput, opts *recordUsageOpts) error { @@ -8041,21 +7795,9 @@ func (s *GatewayService) recordUsageCore(ctx context.Context, input *recordUsage result.Usage.InputTokens = 0 } - // Claude Max cache billing policy（仅 Claude 路径启用） - cacheTTLOverridden := false - simulatedClaudeMax := false - if opts.EnableClaudePath { - var apiKeyGroup *Group - if apiKey != nil { - apiKeyGroup = apiKey.Group - } - claudeMaxOutcome := applyClaudeMaxCacheBillingPolicyToUsage(&result.Usage, opts.ParsedRequest, apiKeyGroup, result.Model, account.ID) - simulatedClaudeMax = claudeMaxOutcome.Simulated || - (shouldApplyClaudeMaxBillingRulesForUsage(apiKeyGroup, result.Model, opts.ParsedRequest) && hasCacheCreationTokens(result.Usage)) - } - // Cache TTL Override: 确保计费时 token 分类与账号设置一致 - if account.IsCacheTTLOverrideEnabled() && !simulatedClaudeMax { + cacheTTLOverridden := false + if account.IsCacheTTLOverrideEnabled() { applyCacheTTLOverride(&result.Usage, account.GetCacheTTLOverrideTarget()) cacheTTLOverridden = (result.Usage.CacheCreation5mTokens + result.Usage.CacheCreation1hTokens) > 0 } diff --git a/backend/internal/service/gateway_service_affinity_test.go b/backend/internal/service/gateway_service_affinity_test.go deleted file mode 100644 index d3dc6f4bb0..0000000000 --- a/backend/internal/service/gateway_service_affinity_test.go +++ /dev/null @@ -1,122 +0,0 @@ -//go:build unit - -// Package service 提供 API 网关核心服务。 -// 本文件包含 SortAffinityClients 函数的单元测试， -// 验证 AffinityClient 切片排序逻辑在各种输入条件下的正确行为。 -// -// This file contains unit tests for the SortAffinityClients function, -// verifying correct sorting behavior for AffinityClient slices under -// various input conditions including empty, single, sorted, reverse, -// and duplicate-timestamp scenarios. -package service - -import ( - "testing" - "time" - - "github.com/stretchr/testify/require" -) - -func TestSortAffinityClients_Empty(t *testing.T) { - var clients []AffinityClient - SortAffinityClients(clients) - require.Empty(t, clients) - - clients = []AffinityClient{} - SortAffinityClients(clients) - require.Empty(t, clients) -} - -func TestSortAffinityClients_SingleElement(t *testing.T) { - now := time.Now() - clients := []AffinityClient{ - {ClientID: "client-1", LastActive: now}, - } - SortAffinityClients(clients) - require.Len(t, clients, 1) - require.Equal(t, "client-1", clients[0].ClientID) - require.Equal(t, now, clients[0].LastActive) -} - -func TestSortAffinityClients_AlreadySorted(t *testing.T) { - now := time.Now() - clients := []AffinityClient{ - {ClientID: "newest", LastActive: now}, - {ClientID: "middle", LastActive: now.Add(-1 * time.Hour)}, - {ClientID: "oldest", LastActive: now.Add(-2 * time.Hour)}, - } - SortAffinityClients(clients) - - require.Equal(t, "newest", clients[0].ClientID) - require.Equal(t, "middle", clients[1].ClientID) - require.Equal(t, "oldest", clients[2].ClientID) -} - -func TestSortAffinityClients_ReverseOrder(t *testing.T) { - now := time.Now() - clients := []AffinityClient{ - {ClientID: "oldest", LastActive: now.Add(-2 * time.Hour)}, - {ClientID: "middle", LastActive: now.Add(-1 * time.Hour)}, - {ClientID: "newest", LastActive: now}, - } - SortAffinityClients(clients) - - require.Equal(t, "newest", clients[0].ClientID) - require.Equal(t, "middle", clients[1].ClientID) - require.Equal(t, "oldest", clients[2].ClientID) -} - -func TestSortAffinityClients_SameTimestamps(t *testing.T) { - now := time.Now() - clients := []AffinityClient{ - {ClientID: "c1", LastActive: now}, - {ClientID: "c2", LastActive: now}, - {ClientID: "c3", LastActive: now}, - } - SortAffinityClients(clients) - - // 所有时间戳相同时，排序结果应保持稳定（sort.Slice 不保证稳定性， - // 但只要结果是某种确定的顺序即可）。 - // 验证所有元素仍然存在且时间相同。 - require.Len(t, clients, 3) - ids := map[string]bool{} - for _, c := range clients { - ids[c.ClientID] = true - require.Equal(t, now, c.LastActive) - } - require.True(t, ids["c1"]) - require.True(t, ids["c2"]) - require.True(t, ids["c3"]) -} - -func TestSortAffinityClients_MixedOrder(t *testing.T) { - now := time.Now() - clients := []AffinityClient{ - {ClientID: "c3", LastActive: now.Add(-30 * time.Minute)}, - {ClientID: "c1", LastActive: now}, - {ClientID: "c5", LastActive: now.Add(-2 * time.Hour)}, - {ClientID: "c2", LastActive: now.Add(-10 * time.Minute)}, - {ClientID: "c4", LastActive: now.Add(-1 * time.Hour)}, - } - SortAffinityClients(clients) - - // 按 LastActive 降序排列 - require.Equal(t, "c1", clients[0].ClientID) // now - require.Equal(t, "c2", clients[1].ClientID) // -10m - require.Equal(t, "c3", clients[2].ClientID) // -30m - require.Equal(t, "c4", clients[3].ClientID) // -1h - require.Equal(t, "c5", clients[4].ClientID) // -2h -} - -func TestSortAffinityClients_SubSecondDifferences(t *testing.T) { - base := time.Now() - clients := []AffinityClient{ - {ClientID: "early", LastActive: base}, - {ClientID: "late", LastActive: base.Add(500 * time.Millisecond)}, - } - SortAffinityClients(clients) - - // 500ms 差异也应正确排序（更晚的在前） - require.Equal(t, "late", clients[0].ClientID) - require.Equal(t, "early", clients[1].ClientID) -} diff --git a/backend/internal/service/gemini_multiplatform_test.go b/backend/internal/service/gemini_multiplatform_test.go index 81e48988c0..5e09b95af2 100644 --- a/backend/internal/service/gemini_multiplatform_test.go +++ b/backend/internal/service/gemini_multiplatform_test.go @@ -288,28 +288,6 @@ func (m *mockGatewayCacheForGemini) DeleteSessionAccountID(ctx context.Context, return nil } -func (m *mockGatewayCacheForGemini) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (m *mockGatewayCacheForGemini) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (m *mockGatewayCacheForGemini) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (m *mockGatewayCacheForGemini) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (m *mockGatewayCacheForGemini) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (m *mockGatewayCacheForGemini) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (m *mockGatewayCacheForGemini) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - // TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform 测试 Gemini 单平台选择 func TestGeminiMessagesCompatService_SelectAccountForModelWithExclusions_GeminiPlatform(t *testing.T) { ctx := context.Background() diff --git a/backend/internal/service/group.go b/backend/internal/service/group.go index b4d7845afd..e17032e000 100644 --- a/backend/internal/service/group.go +++ b/backend/internal/service/group.go @@ -50,9 +50,6 @@ type Group struct { // MCP XML 协议注入开关（仅 antigravity 平台使用） MCPXMLInject bool - // Claude usage 模拟开关：将无写缓存 usage 模拟为 claude-max 风格 - SimulateClaudeMaxEnabled bool - // 支持的模型系列（仅 antigravity 平台使用） // 可选值: claude, gemini_text, gemini_image SupportedModelScopes []string diff --git a/backend/internal/service/openai_gateway_service_test.go b/backend/internal/service/openai_gateway_service_test.go index 2828802fe6..9e2f33f22a 100644 --- a/backend/internal/service/openai_gateway_service_test.go +++ b/backend/internal/service/openai_gateway_service_test.go @@ -282,28 +282,6 @@ func (c *stubGatewayCache) DeleteSessionAccountID(ctx context.Context, groupID i return nil } -func (c *stubGatewayCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (c *stubGatewayCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (c *stubGatewayCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (c *stubGatewayCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (c *stubGatewayCache) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (c *stubGatewayCache) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (c *stubGatewayCache) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - func TestOpenAISelectAccountWithLoadAwareness_FiltersUnschedulable(t *testing.T) { now := time.Now() resetAt := now.Add(10 * time.Minute) diff --git a/backend/internal/service/openai_ws_state_store_test.go b/backend/internal/service/openai_ws_state_store_test.go index 9f35cd7ec6..235d42331d 100644 --- a/backend/internal/service/openai_ws_state_store_test.go +++ b/backend/internal/service/openai_ws_state_store_test.go @@ -193,28 +193,6 @@ func (c *openAIWSStateStoreTimeoutProbeCache) DeleteSessionAccountID(ctx context return nil } -func (c *openAIWSStateStoreTimeoutProbeCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (c *openAIWSStateStoreTimeoutProbeCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (c *openAIWSStateStoreTimeoutProbeCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (c *openAIWSStateStoreTimeoutProbeCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (c *openAIWSStateStoreTimeoutProbeCache) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (c *openAIWSStateStoreTimeoutProbeCache) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]AffinityClient, error) { - return nil, nil -} -func (c *openAIWSStateStoreTimeoutProbeCache) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - func TestOpenAIWSStateStore_RedisOpsUseShortTimeout(t *testing.T) { probe := &openAIWSStateStoreTimeoutProbeCache{} store := NewOpenAIWSStateStore(probe) diff --git a/backend/internal/service/sora_generation.go b/backend/internal/service/sora_generation.go index 7323bb38ac..a704454b82 100644 --- a/backend/internal/service/sora_generation.go +++ b/backend/internal/service/sora_generation.go @@ -37,7 +37,6 @@ const ( // Sora 存储类型常量 const ( SoraStorageTypeS3 = "s3" - SoraStorageTypeGDrive = "gdrive" SoraStorageTypeLocal = "local" SoraStorageTypeUpstream = "upstream" SoraStorageTypeNone = "none" @@ -61,5 +60,4 @@ type SoraGenerationRepository interface { Delete(ctx context.Context, id int64) error List(ctx context.Context, params SoraGenerationListParams) ([]*SoraGeneration, int64, error) CountByUserAndStatus(ctx context.Context, userID int64, statuses []string) (int64, error) - CountByStorageType(ctx context.Context, storageType string, statuses []string) (int64, error) } diff --git a/backend/internal/service/sora_generation_service.go b/backend/internal/service/sora_generation_service.go index 0eba7f91c9..22d5b51947 100644 --- a/backend/internal/service/sora_generation_service.go +++ b/backend/internal/service/sora_generation_service.go @@ -35,21 +35,21 @@ type soraGenerationRepoConditionalUpdater interface { // SoraGenerationService 管理 Sora 客户端的生成记录 CRUD。 type SoraGenerationService struct { - genRepo SoraGenerationRepository - objectStorage SoraObjectStorage - quotaService *SoraQuotaService + genRepo SoraGenerationRepository + s3Storage *SoraS3Storage + quotaService *SoraQuotaService } // NewSoraGenerationService 创建生成记录服务。 func NewSoraGenerationService( genRepo SoraGenerationRepository, - objectStorage SoraObjectStorage, + s3Storage *SoraS3Storage, quotaService *SoraQuotaService, ) *SoraGenerationService { return &SoraGenerationService{ - genRepo: genRepo, - objectStorage: objectStorage, - quotaService: quotaService, + genRepo: genRepo, + s3Storage: s3Storage, + quotaService: quotaService, } } @@ -268,15 +268,15 @@ func (s *SoraGenerationService) Delete(ctx context.Context, id, userID int64) er return fmt.Errorf("无权删除此生成记录") } - // 清理存储文件（S3 / Google Drive） - if IsObjectStorageType(gen.StorageType) && len(gen.S3ObjectKeys) > 0 && s.objectStorage != nil { - if err := s.objectStorage.DeleteObjects(ctx, gen.S3ObjectKeys); err != nil { - logger.LegacyPrintf("service.sora_gen", "[SoraGen] 存储清理失败 id=%d type=%s err=%v", id, gen.StorageType, err) + // 清理 S3 文件 + if gen.StorageType == SoraStorageTypeS3 && len(gen.S3ObjectKeys) > 0 && s.s3Storage != nil { + if err := s.s3Storage.DeleteObjects(ctx, gen.S3ObjectKeys); err != nil { + logger.LegacyPrintf("service.sora_gen", "[SoraGen] S3 清理失败 id=%d err=%v", id, err) } } - // 释放配额（对象存储/本地均释放） - if gen.FileSizeBytes > 0 && (IsObjectStorageType(gen.StorageType) || gen.StorageType == SoraStorageTypeLocal) && s.quotaService != nil { + // 释放配额（S3/本地均释放） + if gen.FileSizeBytes > 0 && (gen.StorageType == SoraStorageTypeS3 || gen.StorageType == SoraStorageTypeLocal) && s.quotaService != nil { if err := s.quotaService.ReleaseUsage(ctx, userID, gen.FileSizeBytes); err != nil { logger.LegacyPrintf("service.sora_gen", "[SoraGen] 配额释放失败 id=%d err=%v", id, err) } @@ -290,9 +290,9 @@ func (s *SoraGenerationService) CountActiveByUser(ctx context.Context, userID in return s.genRepo.CountByUserAndStatus(ctx, userID, []string{SoraGenStatusPending, SoraGenStatusGenerating}) } -// ResolveMediaURLs 为对象存储记录动态生成访问 URL。 +// ResolveMediaURLs 为 S3 记录动态生成预签名 URL。 func (s *SoraGenerationService) ResolveMediaURLs(ctx context.Context, gen *SoraGeneration) error { - if gen == nil || !IsObjectStorageType(gen.StorageType) || s.objectStorage == nil { + if gen == nil || gen.StorageType != SoraStorageTypeS3 || s.s3Storage == nil { return nil } if len(gen.S3ObjectKeys) == 0 { @@ -308,7 +308,7 @@ func (s *SoraGenerationService) ResolveMediaURLs(ctx context.Context, gen *SoraG wg.Add(1) go func(i int, objectKey string) { defer wg.Done() - url, err := s.objectStorage.GetAccessURL(ctx, objectKey) + url, err := s.s3Storage.GetAccessURL(ctx, objectKey) if err != nil { errMu.Lock() if firstErr == nil { @@ -330,22 +330,3 @@ func (s *SoraGenerationService) ResolveMediaURLs(ctx context.Context, gen *SoraG return nil } - -// StorageVideoStats 各存储类型的视频统计。 -type StorageVideoStats struct { - Completed int64 `json:"completed"` - InProgress int64 `json:"in_progress"` -} - -// CountByStorageType 按存储类型统计视频数量（completed 和 in_progress）。 -func (s *SoraGenerationService) CountByStorageType(ctx context.Context, storageType string) (completed, inProgress int64, err error) { - completed, err = s.genRepo.CountByStorageType(ctx, storageType, []string{SoraGenStatusCompleted}) - if err != nil { - return 0, 0, fmt.Errorf("count completed: %w", err) - } - inProgress, err = s.genRepo.CountByStorageType(ctx, storageType, []string{SoraGenStatusPending, SoraGenStatusGenerating}) - if err != nil { - return 0, 0, fmt.Errorf("count in_progress: %w", err) - } - return completed, inProgress, nil -} diff --git a/backend/internal/service/sora_object_storage.go b/backend/internal/service/sora_object_storage.go deleted file mode 100644 index b40a7bcd7e..0000000000 --- a/backend/internal/service/sora_object_storage.go +++ /dev/null @@ -1,36 +0,0 @@ -package service - -import "context" - -// SoraObjectStorage 是 Sora 媒体文件的通用对象存储接口。 -type SoraObjectStorage interface { - // Enabled 返回存储是否已启用且配置有效。 - Enabled(ctx context.Context) bool - - // IsHealthy 返回存储健康状态（带短缓存）。 - IsHealthy(ctx context.Context) bool - - // TestConnection 测试存储连接。 - TestConnection(ctx context.Context) error - - // UploadFromURL 从上游 URL 下载并上传到存储。 - // 返回 object key、文件大小、实际使用的存储类型。 - UploadFromURL(ctx context.Context, userID int64, sourceURL string) (objectKey string, sizeBytes int64, storageType string, err error) - - // DeleteObjects 删除一组存储对象。 - DeleteObjects(ctx context.Context, objectKeys []string) error - - // GetAccessURL 获取存储文件的访问 URL。 - GetAccessURL(ctx context.Context, objectKey string) (string, error) - - // RefreshClient 清除缓存客户端，配置变更时调用。 - RefreshClient() - - // StorageType 返回存储类型标识。 - StorageType() string -} - -// IsObjectStorageType 判断是否为对象存储类型。 -func IsObjectStorageType(t string) bool { - return t == SoraStorageTypeS3 -} diff --git a/backend/internal/service/sora_task.go b/backend/internal/service/sora_task.go deleted file mode 100644 index 6b909b9b2b..0000000000 --- a/backend/internal/service/sora_task.go +++ /dev/null @@ -1,63 +0,0 @@ -package service - -import ( - "context" - "time" -) - -// ── 任务状态常量 ── - -const ( - SoraTaskQueued = "queued" - SoraTaskInProgress = "in_progress" - SoraTaskCompleted = "completed" - SoraTaskFailed = "failed" -) - -// ── 对象类型常量 ── - -const ( - SoraObjectVideo = "video" - SoraObjectCharacter = "character" - SoraObjectImage = "image" -) - -// SoraTask 表示一个 Sora 异步任务记录。 -type SoraTask struct { - ID string - AccountID int64 - APIKeyID *int64 - UpstreamTaskID string - ObjectType string - Model string - Prompt string - Status string - Progress int - VideoURL string // 原始上游 URL - StoredKey string // 存储后的 key（本地路径或 S3 key） - StorageType string // local / s3 / gdrive / 空 - ShareID string - CharacterInfo *SoraCharacter - ErrorMessage string - ErrorType string - RequestBody []byte - Seconds string - Size string - CreatedAt time.Time - CompletedAt *time.Time -} - -// SoraCharacter 角色信息。 -type SoraCharacter struct { - Username string `json:"username"` - DisplayName string `json:"display_name"` -} - -// SoraTaskRepository 持久层接口。 -type SoraTaskRepository interface { - Create(ctx context.Context, task *SoraTask) error - GetByID(ctx context.Context, id string) (*SoraTask, error) - GetByIDAndAPIKey(ctx context.Context, id string, apiKeyID int64) (*SoraTask, error) - Update(ctx context.Context, task *SoraTask) error - ListPending(ctx context.Context) ([]*SoraTask, error) -} diff --git a/backend/internal/service/sora_task_service.go b/backend/internal/service/sora_task_service.go deleted file mode 100644 index 93197c44c8..0000000000 --- a/backend/internal/service/sora_task_service.go +++ /dev/null @@ -1,546 +0,0 @@ -package service - -import ( - "bytes" - "context" - "crypto/rand" - "encoding/hex" - "encoding/json" - "fmt" - "io" - "net/http" - "strings" - "time" - - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" -) - -// SoraTaskService manages Sora async task creation and status queries. -type SoraTaskService struct { - repo SoraTaskRepository - accountRepo AccountRepository - soraClient SoraClient - httpUpstream HTTPUpstream -} - -func NewSoraTaskService( - repo SoraTaskRepository, - accountRepo AccountRepository, - soraClient SoraClient, - httpUpstream HTTPUpstream, -) *SoraTaskService { - return &SoraTaskService{ - repo: repo, - accountRepo: accountRepo, - soraClient: soraClient, - httpUpstream: httpUpstream, - } -} - -// ── Request structs ── - -type CreateVideoRequest struct { - Model string `json:"model"` - Prompt string `json:"prompt"` - StyleID string `json:"style_id,omitempty"` - Orientation string `json:"orientation,omitempty"` - Image string `json:"image,omitempty"` - Video string `json:"video,omitempty"` - RemixTargetID string `json:"remix_target_id,omitempty"` - Size string `json:"size,omitempty"` -} - -type CreateImageRequest struct { - Model string `json:"model,omitempty"` - Prompt string `json:"prompt"` - Image string `json:"image,omitempty"` - Size string `json:"size,omitempty"` - ResponseFormat string `json:"response_format,omitempty"` - N int `json:"n,omitempty"` -} - -type EditImageRequest struct { - Image string `json:"image"` - Prompt string `json:"prompt"` - Model string `json:"model,omitempty"` - Size string `json:"size,omitempty"` - ResponseFormat string `json:"response_format,omitempty"` -} - -type RemixRequest struct { - Prompt string `json:"prompt"` -} - -// ── Response structs ── - -type SoraTaskResponse struct { - ID string `json:"id"` - Object string `json:"object"` - CreatedAt int64 `json:"created_at"` - Status string `json:"status"` - Model string `json:"model,omitempty"` - Prompt string `json:"prompt,omitempty"` - Progress int `json:"progress"` - VideoURL string `json:"video_url,omitempty"` - ShareID string `json:"share_id,omitempty"` - Seconds string `json:"seconds,omitempty"` - Size string `json:"size,omitempty"` - Character *SoraCharacter `json:"character,omitempty"` - URL string `json:"url,omitempty"` - CompletedAt *int64 `json:"completed_at,omitempty"` - Error *SoraTaskError `json:"error,omitempty"` -} - -type SoraTaskError struct { - Message string `json:"message"` - Type string `json:"type"` -} - -func TaskToResponse(t *SoraTask) *SoraTaskResponse { - resp := &SoraTaskResponse{ - ID: t.ID, - Object: t.ObjectType, - CreatedAt: t.CreatedAt.Unix(), - Status: t.Status, - Model: t.Model, - Prompt: t.Prompt, - Progress: t.Progress, - Seconds: t.Seconds, - Size: t.Size, - } - if t.VideoURL != "" { - resp.VideoURL = t.VideoURL - } - if t.ShareID != "" { - resp.ShareID = t.ShareID - } - if t.CharacterInfo != nil { - resp.Character = t.CharacterInfo - } - if t.CompletedAt != nil { - ts := t.CompletedAt.Unix() - resp.CompletedAt = &ts - } - if t.Status == SoraTaskFailed && t.ErrorMessage != "" { - resp.Error = &SoraTaskError{ - Message: t.ErrorMessage, - Type: t.ErrorType, - } - } - return resp -} - -// ── Task creation ── - -func (s *SoraTaskService) CreateVideoTask( - ctx context.Context, - apiKeyID int64, - account *Account, - req *CreateVideoRequest, - body []byte, -) (*SoraTask, error) { - modelCfg, ok := GetSoraModelConfig(req.Model) - - objectType := SoraObjectVideo - if req.Video != "" && req.Prompt == "" { - objectType = SoraObjectCharacter - } - - taskID := generateTaskID(objectType) - - seconds := "" - size := "" - if ok && modelCfg.Type == "video" { - seconds = fmt.Sprintf("%d", modelCfg.Frames/30) - if modelCfg.Orientation == "landscape" { - size = "1920x1080" - } else { - size = "1080x1920" - } - } - - task := &SoraTask{ - ID: taskID, - AccountID: account.ID, - APIKeyID: &apiKeyID, - ObjectType: objectType, - Model: req.Model, - Prompt: req.Prompt, - Status: SoraTaskQueued, - Progress: 0, - Seconds: seconds, - Size: size, - CreatedAt: time.Now(), - } - - if account.Type == AccountTypeAPIKey { - task.RequestBody = body - upstreamResp, err := s.forwardCreateToUpstream(ctx, account, "/v1/videos", body) - if err != nil { - return nil, fmt.Errorf("forward to upstream: %w", err) - } - s.applyUpstreamResponse(task, upstreamResp) - } else { - if s.soraClient == nil || !s.soraClient.Enabled() { - return nil, fmt.Errorf("sora SDK client not configured") - } - upstreamID, err := s.createViaSdk(ctx, account, req, modelCfg, objectType) - if err != nil { - return nil, fmt.Errorf("sdk create task: %w", err) - } - task.UpstreamTaskID = upstreamID - } - - if err := s.repo.Create(ctx, task); err != nil { - return nil, fmt.Errorf("save task: %w", err) - } - return task, nil -} - -func (s *SoraTaskService) CreateImageGeneration( - ctx context.Context, - apiKeyID int64, - account *Account, - req *CreateImageRequest, - body []byte, -) (*SoraTask, error) { - taskID := generateTaskID(SoraObjectImage) - - task := &SoraTask{ - ID: taskID, - AccountID: account.ID, - APIKeyID: &apiKeyID, - ObjectType: SoraObjectImage, - Model: req.Model, - Prompt: req.Prompt, - Status: SoraTaskQueued, - Progress: 0, - Size: req.Size, - CreatedAt: time.Now(), - } - - if account.Type == AccountTypeAPIKey { - task.RequestBody = body - upstreamResp, err := s.forwardCreateToUpstream(ctx, account, "/v1/images/generations", body) - if err != nil { - return nil, fmt.Errorf("forward to upstream: %w", err) - } - s.applyUpstreamResponse(task, upstreamResp) - } else { - if s.soraClient == nil || !s.soraClient.Enabled() { - return nil, fmt.Errorf("sora SDK client not configured") - } - modelCfg, _ := GetSoraModelConfig(req.Model) - width, height := modelCfg.Width, modelCfg.Height - if width == 0 { - width = 360 - } - if height == 0 { - height = 360 - } - upstreamID, err := s.soraClient.CreateImageTask(ctx, account, SoraImageRequest{ - Prompt: req.Prompt, - Width: width, - Height: height, - }) - if err != nil { - return nil, fmt.Errorf("sdk create image task: %w", err) - } - task.UpstreamTaskID = upstreamID - } - - if err := s.repo.Create(ctx, task); err != nil { - return nil, fmt.Errorf("save task: %w", err) - } - return task, nil -} - -func (s *SoraTaskService) GetTask(ctx context.Context, taskID string, apiKeyID int64) (*SoraTask, error) { - return s.repo.GetByIDAndAPIKey(ctx, taskID, apiKeyID) -} - -func (s *SoraTaskService) GetTaskByID(ctx context.Context, taskID string) (*SoraTask, error) { - return s.repo.GetByID(ctx, taskID) -} - -// ListPendingTasks returns tasks in queued or in_progress status. -func (s *SoraTaskService) ListPendingTasks(ctx context.Context) ([]*SoraTask, error) { - return s.repo.ListPending(ctx) -} - -// UpdateTask persists task state changes. -func (s *SoraTaskService) UpdateTask(ctx context.Context, task *SoraTask) error { - return s.repo.Update(ctx, task) -} - -func (s *SoraTaskService) GetAccountByID(ctx context.Context, accountID int64) (*Account, error) { - return s.accountRepo.GetByID(ctx, accountID) -} - -// ── Internal methods ── - -func (s *SoraTaskService) createViaSdk( - ctx context.Context, - account *Account, - req *CreateVideoRequest, - modelCfg SoraModelConfig, - objectType string, -) (string, error) { - if objectType == SoraObjectCharacter { - return "", fmt.Errorf("character creation via /v1/videos is not yet supported for OAuth accounts") - } - - orientation := modelCfg.Orientation - if req.Orientation != "" { - orientation = req.Orientation - } - - videoReq := SoraVideoRequest{ - Prompt: req.Prompt, - Orientation: orientation, - Frames: modelCfg.Frames, - Model: modelCfg.Model, - Size: modelCfg.Size, - RemixTargetID: req.RemixTargetID, - } - return s.soraClient.CreateVideoTask(ctx, account, videoReq) -} - -func (s *SoraTaskService) forwardCreateToUpstream( - ctx context.Context, - account *Account, - path string, - body []byte, -) (map[string]any, error) { - apiKey := account.GetCredential("api_key") - baseURL := account.GetBaseURL() - if apiKey == "" || baseURL == "" { - return nil, fmt.Errorf("account %d missing api_key or base_url", account.ID) - } - - upstreamURL := strings.TrimRight(baseURL, "/") + path - req, err := http.NewRequestWithContext(ctx, http.MethodPost, upstreamURL, bytes.NewReader(body)) - if err != nil { - return nil, err - } - req.Header.Set("Content-Type", "application/json") - req.Header.Set("Authorization", "Bearer "+apiKey) - - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) - if err != nil { - return nil, fmt.Errorf("upstream request: %w", err) - } - defer func() { _ = resp.Body.Close() }() - - respBody, err := io.ReadAll(io.LimitReader(resp.Body, 256*1024)) - if err != nil { - return nil, fmt.Errorf("read upstream response: %w", err) - } - if resp.StatusCode >= 400 { - return nil, &SoraUpstreamError{StatusCode: resp.StatusCode, Body: respBody} - } - - var result map[string]any - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, fmt.Errorf("parse upstream response: %w", err) - } - return result, nil -} - -func (s *SoraTaskService) forwardGetToUpstream( - ctx context.Context, - account *Account, - path string, -) (map[string]any, error) { - apiKey := account.GetCredential("api_key") - baseURL := account.GetBaseURL() - if apiKey == "" || baseURL == "" { - return nil, fmt.Errorf("account %d missing api_key or base_url", account.ID) - } - - upstreamURL := strings.TrimRight(baseURL, "/") + path - req, err := http.NewRequestWithContext(ctx, http.MethodGet, upstreamURL, nil) - if err != nil { - return nil, err - } - req.Header.Set("Authorization", "Bearer "+apiKey) - - proxyURL := "" - if account.ProxyID != nil && account.Proxy != nil { - proxyURL = account.Proxy.URL() - } - - resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency) - if err != nil { - return nil, fmt.Errorf("upstream request: %w", err) - } - defer func() { _ = resp.Body.Close() }() - - respBody, err := io.ReadAll(io.LimitReader(resp.Body, 256*1024)) - if err != nil { - return nil, fmt.Errorf("read upstream response: %w", err) - } - if resp.StatusCode >= 400 { - return nil, fmt.Errorf("upstream error %d: %s", resp.StatusCode, string(respBody)) - } - - var result map[string]any - if err := json.Unmarshal(respBody, &result); err != nil { - return nil, fmt.Errorf("parse upstream response: %w", err) - } - return result, nil -} - -func (s *SoraTaskService) applyUpstreamResponse(task *SoraTask, resp map[string]any) { - if id, ok := resp["id"].(string); ok && id != "" { - task.UpstreamTaskID = id - } - if status, ok := resp["status"].(string); ok && status != "" { - task.Status = status - } - if progress, ok := resp["progress"].(float64); ok { - task.Progress = int(progress) - } - if videoURL, ok := resp["video_url"].(string); ok { - task.VideoURL = videoURL - } - if shareID, ok := resp["share_id"].(string); ok { - task.ShareID = shareID - } - if seconds, ok := resp["seconds"].(string); ok { - task.Seconds = seconds - } - if size, ok := resp["size"].(string); ok { - task.Size = size - } - if obj, ok := resp["object"].(string); ok && obj != "" { - task.ObjectType = obj - } -} - -// PollTask polls a single task's latest status (called by worker). -func (s *SoraTaskService) PollTask(ctx context.Context, task *SoraTask, account *Account) error { - if account.Type == AccountTypeAPIKey { - return s.pollUpstreamTask(ctx, task, account) - } - return s.pollSdkTask(ctx, task, account) -} - -func (s *SoraTaskService) pollUpstreamTask(ctx context.Context, task *SoraTask, account *Account) error { - upstreamID := task.UpstreamTaskID - if upstreamID == "" { - upstreamID = task.ID - } - - path := fmt.Sprintf("/v1/videos/%s", upstreamID) - resp, err := s.forwardGetToUpstream(ctx, account, path) - if err != nil { - logger.LegacyPrintf("service.sora_task", "[PollUpstream] task=%s error=%v", task.ID, err) - return err - } - - s.applyUpstreamResponse(task, resp) - now := time.Now() - if task.Status == SoraTaskCompleted || task.Status == SoraTaskFailed { - task.CompletedAt = &now - } - - if errObj, ok := resp["error"].(map[string]any); ok { - if msg, ok := errObj["message"].(string); ok { - task.ErrorMessage = msg - } - if typ, ok := errObj["type"].(string); ok { - task.ErrorType = typ - } - } - - return s.repo.Update(ctx, task) -} - -func (s *SoraTaskService) pollSdkTask(ctx context.Context, task *SoraTask, account *Account) error { - if s.soraClient == nil { - return fmt.Errorf("sora SDK client not configured") - } - upstreamID := task.UpstreamTaskID - if upstreamID == "" { - return fmt.Errorf("task %s has no upstream_task_id", task.ID) - } - - now := time.Now() - - switch task.ObjectType { - case SoraObjectImage: - status, err := s.soraClient.GetImageTask(ctx, account, upstreamID) - if err != nil { - return err - } - task.Progress = int(status.ProgressPct) - switch status.Status { - case "complete", "completed": - task.Status = SoraTaskCompleted - task.Progress = 100 - task.CompletedAt = &now - if len(status.URLs) > 0 { - task.VideoURL = status.URLs[0] - } - case "failed": - task.Status = SoraTaskFailed - task.CompletedAt = &now - task.ErrorMessage = status.ErrorMsg - task.ErrorType = "server_error" - default: - task.Status = SoraTaskInProgress - } - - case SoraObjectVideo, SoraObjectCharacter: - status, err := s.soraClient.GetVideoTask(ctx, account, upstreamID) - if err != nil { - return err - } - task.Progress = status.ProgressPct - switch status.Status { - case "complete", "completed": - task.Status = SoraTaskCompleted - task.Progress = 100 - task.CompletedAt = &now - if len(status.URLs) > 0 { - task.VideoURL = status.URLs[0] - } - if status.GenerationID != "" { - task.ShareID = status.GenerationID - } - case "failed": - task.Status = SoraTaskFailed - task.CompletedAt = &now - task.ErrorMessage = status.ErrorMsg - task.ErrorType = "server_error" - default: - task.Status = SoraTaskInProgress - } - - default: - return fmt.Errorf("unknown object type: %s", task.ObjectType) - } - - return s.repo.Update(ctx, task) -} - -func generateTaskID(objectType string) string { - b := make([]byte, 8) - if _, err := rand.Read(b); err != nil { - b = []byte(fmt.Sprintf("%x", time.Now().UnixNano())) - } - switch objectType { - case SoraObjectCharacter: - return "char_" + hex.EncodeToString(b) - case SoraObjectImage: - return "img_" + hex.EncodeToString(b) - default: - return "video_" + hex.EncodeToString(b) - } -} diff --git a/backend/internal/service/sora_task_worker.go b/backend/internal/service/sora_task_worker.go deleted file mode 100644 index 059b671031..0000000000 --- a/backend/internal/service/sora_task_worker.go +++ /dev/null @@ -1,183 +0,0 @@ -package service - -import ( - "context" - "sync" - "time" - - "github.com/Wei-Shaw/sub2api/internal/pkg/logger" -) - -// SoraTaskWorker polls unfinished Sora tasks in the background. -// When a task completes, it downloads media to configured storage if available. -type SoraTaskWorker struct { - taskService *SoraTaskService - accountRepo AccountRepository - objectStorage SoraObjectStorage - mediaStorage *SoraMediaStorage - interval time.Duration - pollTimeout time.Duration - stopCh chan struct{} - stopOnce sync.Once - wg sync.WaitGroup -} - -func NewSoraTaskWorker( - taskService *SoraTaskService, - accountRepo AccountRepository, - objectStorage SoraObjectStorage, - mediaStorage *SoraMediaStorage, - interval time.Duration, -) *SoraTaskWorker { - if interval <= 0 { - interval = 60 * time.Second - } - return &SoraTaskWorker{ - taskService: taskService, - accountRepo: accountRepo, - objectStorage: objectStorage, - mediaStorage: mediaStorage, - interval: interval, - pollTimeout: 30 * time.Second, - stopCh: make(chan struct{}), - } -} - -func (w *SoraTaskWorker) Start() { - if w == nil || w.taskService == nil { - return - } - w.wg.Add(1) - go func() { - defer w.wg.Done() - ticker := time.NewTicker(w.interval) - defer ticker.Stop() - - logger.LegacyPrintf("service.sora_task_worker", "[Start] polling interval=%s", w.interval) - w.pollAll() - - for { - select { - case <-ticker.C: - w.pollAll() - case <-w.stopCh: - return - } - } - }() -} - -func (w *SoraTaskWorker) Stop() { - if w == nil { - return - } - w.stopOnce.Do(func() { - close(w.stopCh) - }) - w.wg.Wait() -} - -func (w *SoraTaskWorker) pollAll() { - ctx, cancel := context.WithTimeout(context.Background(), 5*time.Minute) - defer cancel() - - tasks, err := w.taskService.ListPendingTasks(ctx) - if err != nil { - logger.LegacyPrintf("service.sora_task_worker", "[PollAll] list pending tasks error: %v", err) - return - } - if len(tasks) == 0 { - return - } - - logger.LegacyPrintf("service.sora_task_worker", "[PollAll] found %d pending tasks", len(tasks)) - - for _, task := range tasks { - select { - case <-w.stopCh: - return - default: - } - w.pollOne(ctx, task) - } -} - -func (w *SoraTaskWorker) pollOne(ctx context.Context, task *SoraTask) { - pollCtx, cancel := context.WithTimeout(ctx, w.pollTimeout) - defer cancel() - - account, err := w.accountRepo.GetByID(pollCtx, task.AccountID) - if err != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[PollOne] task=%s get account=%d error: %v", task.ID, task.AccountID, err) - w.markTaskFailed(pollCtx, task, "account not found") - return - } - - if err := w.taskService.PollTask(pollCtx, task, account); err != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[PollOne] task=%s poll error: %v", task.ID, err) - return - } - - logger.LegacyPrintf("service.sora_task_worker", - "[PollOne] task=%s status=%s progress=%d", task.ID, task.Status, task.Progress) - - if task.Status == SoraTaskCompleted && task.VideoURL != "" && task.StoredKey == "" { - w.tryStoreMedia(pollCtx, task) - } -} - -// tryStoreMedia downloads media to configured storage. -// Priority: S3 > local disk > keep upstream URL. -func (w *SoraTaskWorker) tryStoreMedia(ctx context.Context, task *SoraTask) { - mediaType := "video" - if task.ObjectType == SoraObjectImage { - mediaType = "image" - } - - if w.objectStorage != nil && w.objectStorage.Enabled(ctx) { - key, _, storageType, err := w.objectStorage.UploadFromURL(ctx, 0, task.VideoURL) - if err != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[StoreMedia] task=%s object storage upload error: %v", task.ID, err) - } else { - task.StoredKey = key - task.StorageType = storageType - if updateErr := w.taskService.UpdateTask(ctx, task); updateErr != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[StoreMedia] task=%s update stored_key error: %v", task.ID, updateErr) - } - return - } - } - - if w.mediaStorage != nil && w.mediaStorage.Enabled() { - stored, err := w.mediaStorage.StoreFromURLs(ctx, mediaType, []string{task.VideoURL}) - if err != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[StoreMedia] task=%s local storage error: %v", task.ID, err) - return - } - if len(stored) > 0 && stored[0] != task.VideoURL { - task.StoredKey = stored[0] - task.StorageType = "local" - if updateErr := w.taskService.UpdateTask(ctx, task); updateErr != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[StoreMedia] task=%s update stored_key error: %v", task.ID, updateErr) - } - } - } -} - -func (w *SoraTaskWorker) markTaskFailed(ctx context.Context, task *SoraTask, message string) { - now := time.Now() - task.Status = SoraTaskFailed - task.ErrorMessage = message - task.ErrorType = "server_error" - task.CompletedAt = &now - if updateErr := w.taskService.UpdateTask(ctx, task); updateErr != nil { - logger.LegacyPrintf("service.sora_task_worker", - "[PollOne] task=%s update failed: %v", task.ID, updateErr) - } -} diff --git a/backend/internal/testutil/stubs.go b/backend/internal/testutil/stubs.go index 1ae72127bc..6513a14ad1 100644 --- a/backend/internal/testutil/stubs.go +++ b/backend/internal/testutil/stubs.go @@ -100,28 +100,6 @@ func (c StubGatewayCache) RefreshSessionTTL(_ context.Context, _ int64, _ string func (c StubGatewayCache) DeleteSessionAccountID(_ context.Context, _ int64, _ string) error { return nil } -func (c StubGatewayCache) GetAffinityAccounts(_ context.Context, _ int64, _ int64, _ string, _ time.Duration) ([]int64, error) { - return nil, nil -} -func (c StubGatewayCache) UpdateAffinity(_ context.Context, _ int64, _ int64, _ string, _ int64, _ time.Duration) error { - return nil -} -func (c StubGatewayCache) GetAffinityMultiCount(_ context.Context, _ int64, _ int64, _ int64, _ time.Duration) (int64, int64, int64, error) { - return 0, 0, 0, nil -} -func (c StubGatewayCache) GetAccountAffinityCountBatch(_ context.Context, _ int64, _ []int64, _ time.Duration) (map[int64]int64, error) { - return map[int64]int64{}, nil -} -func (c StubGatewayCache) GetAccountAffinityClientsBatch(_ context.Context, _ map[int64][]int64, _ time.Duration) (map[int64][]string, error) { - return map[int64][]string{}, nil -} -func (c StubGatewayCache) GetAccountAffinityClientsWithScores(_ context.Context, _ int64, _ []int64, _ time.Duration) ([]service.AffinityClient, error) { - return nil, nil -} -func (c StubGatewayCache) ClearAccountAffinity(_ context.Context, _ int64, _ []int64) error { - return nil -} - // ============================================================ // StubSessionLimitCache — service.SessionLimitCache 的空实现 // ============================================================ diff --git a/docs/ACCOUNT_SCHEDULING_FLOW.md b/docs/ACCOUNT_SCHEDULING_FLOW.md deleted file mode 100644 index 3f7cc807ae..0000000000 --- a/docs/ACCOUNT_SCHEDULING_FLOW.md +++ /dev/null @@ -1,61 +0,0 @@ -# Account Scheduling Flow（SelectAccountWithLoadAwareness） - -本文档对应后端主调度入口 `SelectAccountWithLoadAwareness`，用于说明 Anthropic 网关的账号选择流程、亲和策略与关键过滤规则。 - -## 代码入口与引用 - -- 调度主入口：`backend/internal/service/gateway_service.go` 中的 `SelectAccountWithLoadAwareness` -- 亲和详情 API：`backend/internal/handler/admin/account_affinity_handler.go` 中的 `GetAffinityDetails` -- Admin 路由：`backend/internal/server/routes/admin.go`（`/api/v1/admin/accounts/:id/affinity-details`） - -## 流程概览（修正版） - -1. 解析输入参数 -`sessionHash`、`stickyAccountID`、`affinityClientID(metadata.user_id 提取)`、`affinityUserID(sub2api user id)`。 - -2. Claude Code 限制与分组降级 -先执行 `checkClaudeCodeRestriction`，必要时替换 `groupID`。 - -3. Legacy / Load-aware 分支 -若 `concurrencyService == nil || !LoadBatchEnabled` 走传统路径；否则走负载感知路径。 - -4. Layer 1（模型路由优先） -- 路由候选过滤（排除、平台、模型、配额、窗口费用、RPM） -- 路由范围内 sticky 优先 -- 路由内按（优先级 > 负载 > 亲和数 > LRU）尝试获取槽位 - -5. Layer 1.3（pinned_users 预处理） -仅做 `UpdateAffinity` 预热，不直接做调度决策。 - -6. Layer 1.4（客户端亲和调度） -按亲和记录尝试命中；`allow_switch=false` 且无一票放行时尝试等待计划，否则返回 `ErrAffinityNoSwitch`。 - -7. Layer 1.5（粘性会话） -仅在 `!affinityHit && routingAccountIDs==0` 时生效；先过 `shouldClearStickySession`。 - -8. Layer 2（负载感知选择） -分层过滤：优先级 -> 亲和区(单维客户端) -> 最低负载 -> 最少亲和客户端 -> LRU。 - -9. Layer 3（兜底排队） -按 `FallbackSelectionMode`（`last_used` 或 `random`）排序后返回等待计划。 -注意：此层不是“按最低负载”。 - -## 关键业务规则（2026-03） - -### 无客户端 ID 时过滤亲和 Anthropic 账号 - -当 `metadata.user_id` 无法提取 `client_id`（即 `affinityClientID == ""`）时： - -- 会直接过滤 **Anthropic 平台且开启客户端亲和** 的账号（覆盖 OAuth / SetupToken / API Key / Bedrock 等类型） -- 该规则只看平台与亲和开关，不再限定账号类型 - -设计目标：避免无客户端标识的请求误用客户端亲和账号。 - -## 亲和详情 API 契约 - -`GET /api/v1/admin/accounts/:id/affinity-details` 返回结构（后端已对齐前端）： - -- `users[]`：包含 `user_id`、`user_email`、`client_count`、`is_pinned`、`clients[]` -- `total_users` -- `total_clients` -- `pinned_users` diff --git a/frontend/src/api/admin/accounts.ts b/frontend/src/api/admin/accounts.ts index b73ac554d6..14ad56bee3 100644 --- a/frontend/src/api/admin/accounts.ts +++ b/frontend/src/api/admin/accounts.ts @@ -18,7 +18,6 @@ import type { AdminDataImportResult, CheckMixedChannelRequest, CheckMixedChannelResponse, - AffinityDetailsResponse } from '@/types' /** @@ -569,28 +568,6 @@ export async function refreshOpenAIToken( return data } -/** - * Validate Sora session token and exchange to access token - * @param sessionToken - Sora session token - * @param proxyId - Optional proxy ID - * @param endpoint - API endpoint path - * @returns Token information including access_token - */ -export async function validateSoraSessionToken( - sessionToken: string, - proxyId?: number | null, - endpoint: string = '/admin/sora/st2at' -): Promise> { - const payload: { session_token: string; proxy_id?: number } = { - session_token: sessionToken - } - if (proxyId) { - payload.proxy_id = proxyId - } - const { data } = await apiClient.post>(endpoint, payload) - return data -} - /** * Batch operation result type */ @@ -638,30 +615,6 @@ export async function setPrivacy(id: number): Promise { return data } -/** - * Get affinity clients for an account with last active timestamps - * @param id - Account ID - * @returns List of affinity clients - */ -export async function getAffinityClients(id: number): Promise<{ client_id: string; last_active: string }[]> { - const { data } = await apiClient.get<{ client_id: string; last_active: string }[]>( - `/admin/accounts/${id}/affinity-clients` - ) - return data -} - -/** - * Get affinity details for an account with user-level grouping - * @param id - Account ID - * @returns Affinity details with user groups - */ -export async function getAffinityDetails(id: number): Promise { - const { data } = await apiClient.get( - `/admin/accounts/${id}/affinity-details` - ) - return data -} - export const accountsAPI = { list, listWithEtag, @@ -688,7 +641,6 @@ export const accountsAPI = { generateAuthUrl, exchangeCode, refreshOpenAIToken, - validateSoraSessionToken, batchCreate, batchUpdateCredentials, bulkUpdate, @@ -700,8 +652,6 @@ export const accountsAPI = { batchClearError, batchRefresh, setPrivacy, - getAffinityClients, - getAffinityDetails } export default accountsAPI diff --git a/frontend/src/api/admin/settings.ts b/frontend/src/api/admin/settings.ts index a7f05e33f9..45785c5805 100644 --- a/frontend/src/api/admin/settings.ts +++ b/frontend/src/api/admin/settings.ts @@ -40,7 +40,6 @@ export interface SystemSettings { hide_ccs_import_button: boolean purchase_subscription_enabled: boolean purchase_subscription_url: string - sora_client_enabled: boolean backend_mode_enabled: boolean custom_menu_items: CustomMenuItem[] custom_endpoints: CustomEndpoint[] @@ -114,7 +113,6 @@ export interface UpdateSettingsRequest { hide_ccs_import_button?: boolean purchase_subscription_enabled?: boolean purchase_subscription_url?: string - sora_client_enabled?: boolean backend_mode_enabled?: boolean custom_menu_items?: CustomMenuItem[] custom_endpoints?: CustomEndpoint[] diff --git a/frontend/src/components/account/AccountCapacityCell.vue b/frontend/src/components/account/AccountCapacityCell.vue index aa926d5eab..3df5b24b4f 100644 --- a/frontend/src/components/account/AccountCapacityCell.vue +++ b/frontend/src/components/account/AccountCapacityCell.vue @@ -28,9 +28,6 @@ - - - @@ -43,7 +40,6 @@ import { computed } from 'vue' import { useI18n } from 'vue-i18n' import type { Account } from '@/types' import CapacityBadge from '@/components/account/CapacityBadge.vue' -import AffinityBadge from '@/components/account/AffinityBadge.vue' import QuotaBadge from '@/components/account/QuotaBadge.vue' const props = defineProps<{ @@ -173,46 +169,6 @@ const rpmTooltip = computed(() => { } }) -// ====== 客户端亲和 ====== -const showAffinity = computed(() => - props.account.platform === 'anthropic' && - props.account.client_affinity_enabled === true && - (props.account.affinity_client_count != null || props.account.affinity_user_count != null) -) - -const affinityClientCount = computed(() => props.account.affinity_client_count ?? 0) -const affinityUserCount = computed(() => props.account.affinity_user_count ?? 0) - -const affinityBase = computed(() => { - const extra = props.account.extra as Record | undefined - if (!extra) return 0 - const v = extra.affinity_base - return (typeof v === 'number' && v > 0) ? v : 0 -}) - -const affinityBuffer = computed((): number | null => { - const extra = props.account.extra as Record | undefined - if (!extra) return null - const v = extra.affinity_buffer - if (typeof v === 'number') return v - return null -}) - -const affinityUserBase = computed(() => { - const extra = props.account.extra as Record | undefined - if (!extra) return 0 - const v = extra.affinity_user_base - return (typeof v === 'number' && v > 0) ? v : 0 -}) - -const affinityUserBuffer = computed((): number | null => { - const extra = props.account.extra as Record | undefined - if (!extra) return null - const v = extra.affinity_user_buffer - if (typeof v === 'number') return v - return null -}) - // 格式化费用显示 const formatCost = (value: number | null | undefined) => { if (value === null || value === undefined) return '0' diff --git a/frontend/src/components/account/AffinityBadge.vue b/frontend/src/components/account/AffinityBadge.vue deleted file mode 100644 index 99caf3de76..0000000000 --- a/frontend/src/components/account/AffinityBadge.vue +++ /dev/null @@ -1,254 +0,0 @@ - - - diff --git a/frontend/src/components/account/AffinityConfigCard.vue b/frontend/src/components/account/AffinityConfigCard.vue deleted file mode 100644 index fb458242e3..0000000000 --- a/frontend/src/components/account/AffinityConfigCard.vue +++ /dev/null @@ -1,528 +0,0 @@ - - - diff --git a/frontend/src/components/account/CreateAccountModal.vue b/frontend/src/components/account/CreateAccountModal.vue index 82f1bbbb06..acddb9b285 100644 --- a/frontend/src/components/account/CreateAccountModal.vue +++ b/frontend/src/components/account/CreateAccountModal.vue @@ -1587,24 +1587,6 @@ @update:weeklyResetHour="editWeeklyResetHour = $event" @update:resetTimezone="editResetTimezone = $event" /> - @@ -1968,25 +1950,6 @@

- -
@@ -2986,7 +2949,6 @@ import ProxySelector from '@/components/common/ProxySelector.vue' import GroupSelector from '@/components/common/GroupSelector.vue' import ModelWhitelistSelector from '@/components/account/ModelWhitelistSelector.vue' import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' -import AffinityConfigCard from '@/components/account/AffinityConfigCard.vue' import { applyInterceptWarmup } from '@/components/account/credentialsBuilder' import { formatDateTimeLocalInput, parseDateTimeLocalInput } from '@/utils/format' import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' @@ -3181,16 +3143,6 @@ const antigravityMixedChannelConfirmed = ref(false) const showAdvancedOAuth = ref(false) const showGeminiHelpDialog = ref(false) -// Client affinity (all Anthropic accounts) -const clientAffinityEnabled = ref(false) -const affinityBase = ref(null) -const affinityBuffer = ref(null) -const affinityAllowSwitch = ref(true) -const affinityUserBase = ref(null) -const affinityUserBuffer = ref(null) -const perUserClientLimit = ref(null) -const pinnedUsers = ref([]) - // Quota control state (Anthropic OAuth/SetupToken only) const windowCostEnabled = ref(false) const windowCostLimit = ref(null) @@ -3848,14 +3800,6 @@ const resetForm = () => { editWeeklyResetDay.value = null editWeeklyResetHour.value = null editResetTimezone.value = null - clientAffinityEnabled.value = false - affinityBase.value = null - affinityBuffer.value = null - affinityAllowSwitch.value = true - affinityUserBase.value = null - affinityUserBuffer.value = null - perUserClientLimit.value = null - pinnedUsers.value = [] modelMappings.value = [] modelRestrictionMode.value = 'whitelist' allowedModels.value = [...claudeModels] // Default fill related models @@ -3965,61 +3909,10 @@ const buildAnthropicExtra = (base?: Record): Record 0 ? extra : undefined } -/** 将客户端亲和设置写入 extra（Anthropic 全类型通用） */ -const applyClientAffinity = (extra: Record) => { - if (clientAffinityEnabled.value) { - extra.affinity_enabled = true - extra.client_affinity_enabled = true - if (affinityBase.value != null && affinityBase.value > 0) { - extra.affinity_base = affinityBase.value - } else { - delete extra.affinity_base - } - if (affinityBase.value != null && affinityBase.value > 0 && affinityBuffer.value != null) { - extra.affinity_buffer = affinityBuffer.value - } else { - delete extra.affinity_buffer - } - // v2 fields - extra.affinity_allow_switch = affinityAllowSwitch.value - if (affinityUserBase.value != null && affinityUserBase.value > 0) { - extra.affinity_user_base = affinityUserBase.value - } else { - delete extra.affinity_user_base - } - if (affinityUserBase.value != null && affinityUserBase.value > 0 && affinityUserBuffer.value != null) { - extra.affinity_user_buffer = affinityUserBuffer.value - } else { - delete extra.affinity_user_buffer - } - if (perUserClientLimit.value != null && perUserClientLimit.value > 0) { - extra.per_user_client_limit = perUserClientLimit.value - } else { - delete extra.per_user_client_limit - } - if (pinnedUsers.value.length > 0) { - extra.pinned_users = pinnedUsers.value - } else { - delete extra.pinned_users - } - } else { - delete extra.affinity_enabled - delete extra.client_affinity_enabled - delete extra.affinity_base - delete extra.affinity_buffer - delete extra.affinity_allow_switch - delete extra.affinity_user_base - delete extra.affinity_user_buffer - delete extra.per_user_client_limit - delete extra.pinned_users - } -} - const buildSoraExtra = ( base?: Record, linkedOpenAIAccountId?: string | number @@ -5038,9 +4931,6 @@ const handleAnthropicExchange = async (authCode: string) => { extra.cache_ttl_override_target = cacheTTLOverrideTarget.value } - // Add client affinity setting - applyClientAffinity(extra) - // Add custom base URL settings if (customBaseUrlEnabled.value && customBaseUrl.value.trim()) { extra.custom_base_url_enabled = true diff --git a/frontend/src/components/account/EditAccountModal.vue b/frontend/src/components/account/EditAccountModal.vue index 013a89792c..f1c96ff97a 100644 --- a/frontend/src/components/account/EditAccountModal.vue +++ b/frontend/src/components/account/EditAccountModal.vue @@ -1027,7 +1027,6 @@
-
@@ -1181,24 +1180,6 @@ @update:weeklyResetHour="editWeeklyResetHour = $event" @update:resetTimezone="editResetTimezone = $event" /> -
- -
@@ -1833,7 +1795,6 @@ import ProxySelector from '@/components/common/ProxySelector.vue' import GroupSelector from '@/components/common/GroupSelector.vue' import ModelWhitelistSelector from '@/components/account/ModelWhitelistSelector.vue' import QuotaLimitCard from '@/components/account/QuotaLimitCard.vue' -import AffinityConfigCard from '@/components/account/AffinityConfigCard.vue' import { applyInterceptWarmup } from '@/components/account/credentialsBuilder' import { formatDateTimeLocalInput, parseDateTimeLocalInput } from '@/utils/format' import { createStableObjectKeyResolver } from '@/utils/stableObjectKey' @@ -1963,14 +1924,6 @@ const tlsFingerprintProfiles = ref<{ id: number; name: string }[]>([]) const sessionIdMaskingEnabled = ref(false) const cacheTTLOverrideEnabled = ref(false) const cacheTTLOverrideTarget = ref('5m') -const clientAffinityEnabled = ref(false) -const affinityBase = ref(null) -const affinityBuffer = ref(null) -const affinityAllowSwitch = ref(true) -const affinityUserBase = ref(null) -const affinityUserBuffer = ref(null) -const perUserClientLimit = ref(null) -const pinnedUsers = ref([]) const customBaseUrlEnabled = ref(false) const customBaseUrl = ref('') @@ -2601,14 +2554,6 @@ function loadQuotaControlSettings(account: Account) { sessionIdMaskingEnabled.value = false cacheTTLOverrideEnabled.value = false cacheTTLOverrideTarget.value = '5m' - clientAffinityEnabled.value = false - affinityBase.value = null - affinityBuffer.value = null - affinityAllowSwitch.value = true - affinityUserBase.value = null - affinityUserBuffer.value = null - perUserClientLimit.value = null - pinnedUsers.value = [] customBaseUrlEnabled.value = false customBaseUrl.value = '' @@ -2617,31 +2562,6 @@ function loadQuotaControlSettings(account: Account) { return } - // Client affinity (all Anthropic accounts) - if (account.client_affinity_enabled === true) { - clientAffinityEnabled.value = true - } - // Affinity base/buffer from extra - const extra = account.extra as Record | undefined - if (extra) { - const base = extra.affinity_base - affinityBase.value = (typeof base === 'number' && base > 0) ? base : null - // buffer: null = infinite yellow, 0 = no yellow, >0 = yellow range - const buf = extra.affinity_buffer - affinityBuffer.value = (typeof buf === 'number') ? buf : null - - // New v2 fields - affinityAllowSwitch.value = extra.affinity_allow_switch !== false - const ub = extra.affinity_user_base - affinityUserBase.value = (typeof ub === 'number' && ub > 0) ? ub : null - const ubuf = extra.affinity_user_buffer - affinityUserBuffer.value = (typeof ubuf === 'number') ? ubuf : null - const pul = extra.per_user_client_limit - perUserClientLimit.value = (typeof pul === 'number' && pul > 0) ? pul : null - const pu = extra.pinned_users - pinnedUsers.value = Array.isArray(pu) ? pu : [] - } - // Window cost / session limit only apply to Anthropic OAuth/SetupToken accounts if (account.type !== 'oauth' && account.type !== 'setup-token') { return @@ -3067,60 +2987,6 @@ const handleSubmit = async () => { updatePayload.extra = newExtra } - // For all Anthropic accounts, handle client_affinity in extra - if (props.account.platform === 'anthropic') { - const currentExtra = (props.account.extra as Record) || {} - const newExtra: Record = { ...currentExtra } - if (clientAffinityEnabled.value) { - newExtra.affinity_enabled = true - newExtra.client_affinity_enabled = true - if (affinityBase.value != null && affinityBase.value > 0) { - newExtra.affinity_base = affinityBase.value - } else { - delete newExtra.affinity_base - } - // buffer: null = infinite yellow, 0 = no yellow, >0 = yellow range - if (affinityBase.value != null && affinityBase.value > 0 && affinityBuffer.value != null) { - newExtra.affinity_buffer = affinityBuffer.value - } else { - delete newExtra.affinity_buffer - } - // v2 fields - newExtra.affinity_allow_switch = affinityAllowSwitch.value - if (affinityUserBase.value != null && affinityUserBase.value > 0) { - newExtra.affinity_user_base = affinityUserBase.value - } else { - delete newExtra.affinity_user_base - } - if (affinityUserBase.value != null && affinityUserBase.value > 0 && affinityUserBuffer.value != null) { - newExtra.affinity_user_buffer = affinityUserBuffer.value - } else { - delete newExtra.affinity_user_buffer - } - if (perUserClientLimit.value != null && perUserClientLimit.value > 0) { - newExtra.per_user_client_limit = perUserClientLimit.value - } else { - delete newExtra.per_user_client_limit - } - if (pinnedUsers.value.length > 0) { - newExtra.pinned_users = pinnedUsers.value - } else { - delete newExtra.pinned_users - } - } else { - newExtra.affinity_enabled = false - newExtra.client_affinity_enabled = false - delete newExtra.affinity_base - delete newExtra.affinity_buffer - delete newExtra.affinity_allow_switch - delete newExtra.affinity_user_base - delete newExtra.affinity_user_buffer - delete newExtra.per_user_client_limit - delete newExtra.pinned_users - } - updatePayload.extra = newExtra - } - // For Anthropic OAuth/SetupToken accounts, handle quota control settings in extra if (props.account.platform === 'anthropic' && (props.account.type === 'oauth' || props.account.type === 'setup-token')) { const currentExtra = (updatePayload.extra as Record) || (props.account.extra as Record) || {} diff --git a/frontend/src/components/account/__tests__/AffinityBadge.spec.ts b/frontend/src/components/account/__tests__/AffinityBadge.spec.ts deleted file mode 100644 index e7561e7f76..0000000000 --- a/frontend/src/components/account/__tests__/AffinityBadge.spec.ts +++ /dev/null @@ -1,163 +0,0 @@ -import { describe, expect, it, vi } from 'vitest' -import { mount } from '@vue/test-utils' -import AffinityBadge from '../AffinityBadge.vue' - -vi.mock('@/api/admin/accounts', () => ({ - getAffinityClients: vi.fn(), - getAffinityDetails: vi.fn() -})) - -vi.mock('vue-i18n', async () => { - const actual = await vi.importActual('vue-i18n') - return { - ...actual, - useI18n: () => ({ - t: (key: string, fallback?: string) => fallback ?? key - }) - } -}) - -function mountBadge(opts: { - clientCount: number - base?: number - buffer?: number | null - userCount?: number - userBase?: number - userBuffer?: number | null -}) { - return mount(AffinityBadge, { - props: { - accountId: 42, - clientCount: opts.clientCount, - base: opts.base ?? 5, - buffer: opts.buffer ?? 10, - userCount: opts.userCount ?? 0, - userBase: opts.userBase ?? 0, - userBuffer: opts.userBuffer ?? null - } - }) -} - -describe('AffinityBadge', () => { - // ====== Original tests (client dimension only) ====== - it('renders the correct client count number', () => { - const wrapper = mountBadge({ clientCount: 5 }) - expect(wrapper.text()).toContain('5') - }) - - it('renders configured limit text', () => { - const wrapper = mountBadge({ clientCount: 5, base: 5, buffer: 10 }) - expect(wrapper.text()).toContain('15') - }) - - it('renders infinity limit text when base is not configured', () => { - const wrapper = mountBadge({ clientCount: 5, base: 0, buffer: null }) - expect(wrapper.text()).toContain('\u221E') - }) - - it('applies red badge class when count exceeds base plus buffer', () => { - const wrapper = mountBadge({ clientCount: 16, base: 5, buffer: 10 }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-red-100') - expect(badge.classes()).toContain('text-red-700') - }) - - it('applies yellow badge class when count is in buffer range', () => { - const wrapper = mountBadge({ clientCount: 6, base: 5, buffer: 10 }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-yellow-100') - expect(badge.classes()).toContain('text-yellow-700') - }) - - it('applies yellow badge class when buffer is infinite', () => { - const wrapper = mountBadge({ clientCount: 6, base: 5, buffer: null }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-yellow-100') - expect(badge.classes()).toContain('text-yellow-700') - }) - - it('applies green badge class when count is within base', () => { - const wrapper = mountBadge({ clientCount: 5, base: 5, buffer: 10 }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-emerald-100') - expect(badge.classes()).toContain('text-emerald-700') - }) - - it('applies gray badge class when count is 0', () => { - const wrapper = mountBadge({ clientCount: 0, base: 5, buffer: 10 }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-gray-100') - expect(badge.classes()).toContain('text-gray-600') - }) - - it('does NOT show popover initially', () => { - const wrapper = mountBadge({ clientCount: 3 }) - expect(wrapper.html()).not.toContain('divide-y') - }) - - it('has mouseenter and mouseleave handlers on the badge', () => { - const wrapper = mountBadge({ clientCount: 3 }) - const badge = wrapper.find('span') - expect(badge.exists()).toBe(true) - badge.trigger('mouseenter') - badge.trigger('mouseleave') - }) - - // ====== Dual dimension tests ====== - it('dual dimension: user green + client yellow = yellow', () => { - const wrapper = mountBadge({ - clientCount: 6, base: 5, buffer: 10, - userCount: 3, userBase: 5, userBuffer: 5 - }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-yellow-100') - expect(badge.classes()).toContain('text-yellow-700') - }) - - it('dual dimension: user green + client red = red', () => { - const wrapper = mountBadge({ - clientCount: 16, base: 5, buffer: 10, - userCount: 3, userBase: 5, userBuffer: 5 - }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-red-100') - expect(badge.classes()).toContain('text-red-700') - }) - - it('dual dimension: user red + client green = red', () => { - const wrapper = mountBadge({ - clientCount: 3, base: 5, buffer: 10, - userCount: 11, userBase: 5, userBuffer: 5 - }) - const badge = wrapper.find('span') - expect(badge.classes()).toContain('bg-red-100') - expect(badge.classes()).toContain('text-red-700') - }) - - it('user only dimension shows user count text', () => { - const wrapper = mountBadge({ - clientCount: 0, base: 0, buffer: null, - userCount: 3, userBase: 5, userBuffer: 5 - }) - expect(wrapper.text()).toContain('3') - expect(wrapper.text()).toContain('10') - }) - - it('client only dimension shows client count text', () => { - const wrapper = mountBadge({ - clientCount: 4, base: 5, buffer: 10, - userCount: 0, userBase: 0, userBuffer: null - }) - expect(wrapper.text()).toContain('4') - expect(wrapper.text()).toContain('15') - }) - - it('dual dimension limit text shows both U and C prefixes', () => { - const wrapper = mountBadge({ - clientCount: 3, base: 5, buffer: 10, - userCount: 2, userBase: 4, userBuffer: 3 - }) - expect(wrapper.text()).toContain('U2/7') - expect(wrapper.text()).toContain('C3/15') - }) -}) diff --git a/frontend/src/components/layout/AppSidebar.vue b/frontend/src/components/layout/AppSidebar.vue index 80a3eed8f3..09b83f6f25 100644 --- a/frontend/src/components/layout/AppSidebar.vue +++ b/frontend/src/components/layout/AppSidebar.vue @@ -467,21 +467,6 @@ const ChevronDoubleLeftIcon = { ) } -const SoraIcon = { - render: () => - h( - 'svg', - { fill: 'none', viewBox: '0 0 24 24', stroke: 'currentColor', 'stroke-width': '1.5' }, - [ - h('path', { - 'stroke-linecap': 'round', - 'stroke-linejoin': 'round', - d: 'M9.813 15.904L9 18.75l-.813-2.846a4.5 4.5 0 00-3.09-3.09L2.25 12l2.846-.813a4.5 4.5 0 003.09-3.09L9 5.25l.813 2.846a4.5 4.5 0 003.09 3.09L15.75 12l-2.846.813a4.5 4.5 0 00-3.09 3.09z' - }) - ] - ) -} - const ChevronDoubleRightIcon = { render: () => h( @@ -504,9 +489,6 @@ const userNavItems = computed((): NavItem[] => { { path: '/keys', label: t('nav.apiKeys'), icon: KeyIcon }, { path: '/usage', label: t('nav.usage'), icon: ChartIcon, hideInSimpleMode: true }, { path: '/subscriptions', label: t('nav.mySubscriptions'), icon: CreditCardIcon, hideInSimpleMode: true }, - ...(appStore.cachedPublicSettings?.sora_client_enabled - ? [{ path: '/sora', label: t('nav.sora'), icon: SoraIcon }] - : []), ...(appStore.cachedPublicSettings?.purchase_subscription_enabled ? [ { @@ -535,9 +517,6 @@ const personalNavItems = computed((): NavItem[] => { { path: '/keys', label: t('nav.apiKeys'), icon: KeyIcon }, { path: '/usage', label: t('nav.usage'), icon: ChartIcon, hideInSimpleMode: true }, { path: '/subscriptions', label: t('nav.mySubscriptions'), icon: CreditCardIcon, hideInSimpleMode: true }, - ...(appStore.cachedPublicSettings?.sora_client_enabled - ? [{ path: '/sora', label: t('nav.sora'), icon: SoraIcon }] - : []), ...(appStore.cachedPublicSettings?.purchase_subscription_enabled ? [ { diff --git a/frontend/src/components/sora/SoraGeneratePage.vue b/frontend/src/components/sora/SoraGeneratePage.vue deleted file mode 100644 index b6ba251ba0..0000000000 --- a/frontend/src/components/sora/SoraGeneratePage.vue +++ /dev/null @@ -1,441 +0,0 @@ - - - - - diff --git a/frontend/src/components/sora/SoraMediaPreview.vue b/frontend/src/components/sora/SoraMediaPreview.vue deleted file mode 100644 index 4bcfc5dfea..0000000000 --- a/frontend/src/components/sora/SoraMediaPreview.vue +++ /dev/null @@ -1,283 +0,0 @@ - - - - - diff --git a/frontend/src/components/sora/SoraProgressCard.vue b/frontend/src/components/sora/SoraProgressCard.vue deleted file mode 100644 index 8a97946a9b..0000000000 --- a/frontend/src/components/sora/SoraProgressCard.vue +++ /dev/null @@ -1,617 +0,0 @@ - - - - - diff --git a/frontend/src/components/sora/SoraPromptBar.vue b/frontend/src/components/sora/SoraPromptBar.vue deleted file mode 100644 index 25b42b6697..0000000000 --- a/frontend/src/components/sora/SoraPromptBar.vue +++ /dev/null @@ -1,748 +0,0 @@ -