mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 03:52:26 +08:00
Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5485f368b2 | ||
|
|
f8351e9e33 | ||
|
|
2031c1d5ad | ||
|
|
8193b80a59 | ||
|
|
f6d580200a | ||
|
|
b491c3271b | ||
|
|
d106b3b5a0 | ||
|
|
10fe591e54 | ||
|
|
66cd8a72bc | ||
|
|
dc46daa690 | ||
|
|
65246e69d7 | ||
|
|
de69bf1e03 | ||
|
|
eba85dc478 | ||
|
|
abf70750fd | ||
|
|
88b697d5df | ||
|
|
3eb5a84ea1 | ||
|
|
71e5a4baf7 | ||
|
|
4bfcbd0c78 | ||
|
|
2b1e5e7f5c | ||
|
|
cbb4b7e53e | ||
|
|
81c0f11d7a | ||
|
|
561fc1c3e4 | ||
|
|
43491f67a1 | ||
|
|
c05bc4d3cc | ||
|
|
b9ebc8f209 | ||
|
|
f7c48ba59a | ||
|
|
22c16d0517 | ||
|
|
2cc1e7ef90 | ||
|
|
81ff9384cb | ||
|
|
0aaed397cc | ||
|
|
909b8c7ef4 | ||
|
|
b8ff74ad70 | ||
|
|
96cf8bd76a | ||
|
|
1f716aa181 | ||
|
|
111c41cf3f | ||
|
|
787a6a33df | ||
|
|
f8e1e7fed6 | ||
|
|
f09a5b602e | ||
|
|
027e3f2d54 | ||
|
|
8116d235ef | ||
|
|
959cbe3d87 | ||
|
|
c0420e2b8a | ||
|
|
9578ddbddd | ||
|
|
e274de45b7 | ||
|
|
ce328fb37a | ||
|
|
8f7f2b5d92 | ||
|
|
3efeba33a4 | ||
|
|
ce157b32ed | ||
|
|
3cb2381bdc | ||
|
|
a7d2c782df | ||
|
|
a3aa9bae6d | ||
|
|
b939fa9d4a | ||
|
|
d03c42d798 | ||
|
|
6aabbdf547 | ||
|
|
ad4b2f30f5 | ||
|
|
ff758f37d2 | ||
|
|
2e31d8b702 | ||
|
|
8c8fd6e298 | ||
|
|
f62ec2e4a0 | ||
|
|
415a0dc9d0 | ||
|
|
b59f3bf46c | ||
|
|
2db78bd3d6 | ||
|
|
f88d62ad28 | ||
|
|
8363d537e0 | ||
|
|
b8ed245085 | ||
|
|
dc6b318c34 | ||
|
|
96884dd1ae | ||
|
|
0b061c5fa1 | ||
|
|
958a21ad6e | ||
|
|
28167bbf85 | ||
|
|
985428d9d2 | ||
|
|
4a4fd35e0c | ||
|
|
f46038820b | ||
|
|
05ad6b49a1 | ||
|
|
a3675552bd | ||
|
|
b3ad9b67cf | ||
|
|
f49a393566 | ||
|
|
8a46947865 | ||
|
|
7ae0312091 | ||
|
|
58e35a4f32 | ||
|
|
95023e7d49 | ||
|
|
7a70de401c | ||
|
|
666a5c6f87 | ||
|
|
c7343d2aa2 | ||
|
|
a6430bfe20 | ||
|
|
32bf3d0f33 | ||
|
|
cb31033972 | ||
|
|
3aeab296d8 | ||
|
|
f3bbb95310 | ||
|
|
00eabe8ab9 | ||
|
|
a16070ccf6 | ||
|
|
ab99d56e96 | ||
|
|
b1ce821c49 | ||
|
|
55c5eed9fc | ||
|
|
cff3f89850 | ||
|
|
3bdb869fca | ||
|
|
6271c517da | ||
|
|
a7f1ed8fa4 | ||
|
|
a774230661 | ||
|
|
222181efd6 | ||
|
|
5fb5dfb34c | ||
|
|
c74a4f521c |
@@ -1,6 +1,7 @@
|
||||
.PHONY: build build-backend build-frontend test test-backend test-frontend test-frontend-critical
|
||||
|
||||
FRONTEND_CRITICAL_VITEST := \
|
||||
src/i18n/__tests__/localeKeyCompleteness.spec.ts \
|
||||
src/api/__tests__/client.spec.ts \
|
||||
src/api/__tests__/tokenRefresh.spec.ts \
|
||||
src/api/__tests__/channelMonitorV2.spec.ts \
|
||||
|
||||
@@ -816,7 +816,7 @@ xAI quota is passive. Sub2API does not invent subscription quota values; it reco
|
||||
|
||||
`401` responses temporarily remove accounts with invalid credentials from scheduling. `403` responses are treated as access or entitlement failures instead of token-refresh loops. `429` responses use `Retry-After` or a short cooldown to temporarily remove the account from scheduling.
|
||||
|
||||
New Grok image and video generation requests use a media-specific eligibility check. API-key accounts remain eligible. OAuth accounts require positive paid-entitlement evidence from the xAI billing probe; Free, forbidden, missing, malformed, and inconclusive billing observations are excluded from new media generation. Unobserved OAuth accounts are probed before the first media request is forwarded, and imports run the billing-first quota probe proactively. Chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account`.
|
||||
New Grok image and video generation requests use a media-specific eligibility check. API-key accounts remain eligible. OAuth accounts with explicit Free or forbidden billing evidence are excluded from new media generation. Missing or malformed observations are probed before dispatch; a successful but incomplete billing response is treated as `billing_inconclusive` and remains eligible for backwards compatibility, because an unknown billing schema is not proof that the account lacks media entitlement. Operators can quarantine a known-bad account with `extra.grok_media_eligible=false` or force-enable a verified account with `true`. Imports run the billing-first quota probe proactively. Chat requests and video status lookups are not affected by this media-only quarantine. If no eligible account remains, the media endpoint returns HTTP `503` with error type `grok_media_no_eligible_account`.
|
||||
|
||||
Administrators can override automatic media eligibility through the account create/update API by setting `extra.grok_media_eligible` to `false` (exclude) or `true` (force eligible). On update, set it to `null` to remove the override and return to automatic probe-based behavior; omitting the field preserves the current override. A weekly allowance period alone is not treated as a paid tier signal. Successful image responses must contain at least one actual image output; empty HTTP `200` responses trigger account failover instead of being counted and returned as successful generations.
|
||||
|
||||
|
||||
@@ -1 +1 @@
|
||||
0.2.0
|
||||
0.2.1
|
||||
|
||||
@@ -189,10 +189,10 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
adminAccountRepository := repository.NewAdminAccountRepository(client, db, schedulerCache)
|
||||
proxyExitInfoProber := repository.NewProxyExitInfoProber(configConfig)
|
||||
proxyLatencyCache := repository.NewProxyLatencyCache(redisClient)
|
||||
adminService := service.NewAdminService(userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver, channelService)
|
||||
adminService := service.NewAdminService(configConfig, userRepository, adminGroupRepository, adminAccountRepository, proxyRepository, apiKeyRepository, redeemCodeRepository, userGroupRateRepository, userRPMCache, billingCacheService, proxyExitInfoProber, proxyLatencyCache, apiKeyAuthCacheInvalidator, client, settingService, subscriptionService, userSubscriptionRepository, privacyClientFactory, openAIGatewayService, affiliateService, compositeModelRouteRepository, compositeRouteResolver, channelService)
|
||||
adminUserHandler := admin.NewUserHandler(adminService, concurrencyService, serviceUserPlatformQuotaRepository, billingCache, totpService, userService, settingService)
|
||||
groupCapacityService := service.NewGroupCapacityService(accountRepository, groupRepository, concurrencyService, sessionLimitCache, rpmCache)
|
||||
groupHandler := admin.NewGroupHandler(adminService, dashboardService, groupCapacityService)
|
||||
groupHandler := admin.NewGroupHandlerWithConfig(adminService, dashboardService, groupCapacityService, configConfig)
|
||||
claudeUsageFetcher := repository.NewClaudeUsageFetcher(httpUpstream)
|
||||
antigravityQuotaFetcher := service.NewAntigravityQuotaFetcher(proxyRepository, configConfig)
|
||||
grokQuotaFetcher := service.NewGrokQuotaFetcher()
|
||||
@@ -205,12 +205,12 @@ func initializeApplication(buildInfo handler.BuildInfo) (*Application, error) {
|
||||
pluginManager := service.NewPluginManager(pluginRepository, secretEncryptor, configConfig, pluginHostInfo)
|
||||
accountTestService := service.ProvideAccountTestService(accountRepository, geminiTokenProvider, claudeTokenProvider, grokTokenProvider, antigravityGatewayService, httpUpstream, configConfig, tlsFingerprintProfileService, openAIGatewayService, settingService, pluginManager)
|
||||
crsSyncService := service.NewCRSSyncService(accountRepository, proxyRepository, oAuthService, openAIOAuthService, geminiOAuthService, configConfig)
|
||||
accountHandler := admin.ProvideAccountHandler(adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
|
||||
accountHandler := admin.ProvideAccountHandler(configConfig, adminService, oAuthService, openAIOAuthService, geminiOAuthService, antigravityOAuthService, grokOAuthService, rateLimitService, accountUsageService, accountTestService, concurrencyService, crsSyncService, sessionLimitCache, rpmCache, compositeTokenCacheInvalidator, grokQuotaService)
|
||||
adminAnnouncementHandler := admin.NewAnnouncementHandler(announcementService)
|
||||
dataManagementService := service.NewDataManagementService()
|
||||
dataManagementHandler := admin.NewDataManagementHandler(dataManagementService)
|
||||
backupObjectStoreFactory := repository.NewS3BackupStoreFactory()
|
||||
dbDumper := repository.NewPgDumper(configConfig)
|
||||
dbDumper := repository.NewPgDumper(configConfig, db)
|
||||
backupService := service.ProvideBackupService(settingRepository, configConfig, secretEncryptor, backupObjectStoreFactory, dbDumper, leaderLockCache, db)
|
||||
imageStorageFactory := repository.ProvideImageStorageFactory()
|
||||
imageStorageSettingService := service.ProvideImageStorageSettingService(settingRepository, secretEncryptor, backupService, imageStorageFactory, configConfig)
|
||||
|
||||
@@ -134,8 +134,8 @@ type Group struct {
|
||||
DefaultMappedModel string `json:"default_mapped_model,omitempty"`
|
||||
// OpenAI Messages 调度模型配置:按 Claude 系列/精确模型映射到目标 GPT 模型
|
||||
MessagesDispatchModelConfig domain.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config,omitempty"`
|
||||
// 自定义 /v1/models 展示列表配置;仅影响模型列表响应,不影响调度
|
||||
ModelsListConfig domain.GroupModelsListConfig `json:"models_list_config,omitempty"`
|
||||
// 分组模型白名单:同时约束模型列表接口与请求准入
|
||||
ModelAllowlist domain.GroupModelAllowlist `json:"model_allowlist,omitempty"`
|
||||
// 固定账号获取 Codex Model Manifest 配置;开启后 /models 请求只用选定账号拉取(仅 openai 平台)
|
||||
CodexModelsManifestConfig domain.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config,omitempty"`
|
||||
// 分组 RPM 上限,0 表示不限制;设置后接管该分组用户的限流
|
||||
@@ -258,7 +258,7 @@ func (*Group) scanValues(columns []string) ([]any, error) {
|
||||
values := make([]any, len(columns))
|
||||
for i := range columns {
|
||||
switch columns[i] {
|
||||
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelsListConfig, group.FieldCodexModelsManifestConfig, group.FieldReasoningEffortMappings:
|
||||
case group.FieldVideoModelPrices, group.FieldModelPricing, group.FieldModelRouting, group.FieldSupportedModelScopes, group.FieldMessagesDispatchModelConfig, group.FieldModelAllowlist, group.FieldCodexModelsManifestConfig, group.FieldReasoningEffortMappings:
|
||||
values[i] = new([]byte)
|
||||
case group.FieldPeakRateEnabled, group.FieldIsExclusive, group.FieldAllowImageGeneration, group.FieldAllowBatchImageGeneration, group.FieldImageRateIndependent, group.FieldVideoRateIndependent, group.FieldLongContextPricingEnabled, group.FieldClaudeCodeOnly, group.FieldModelRoutingEnabled, group.FieldMcpXMLInject, group.FieldAllowMessagesDispatch, group.FieldAllowLive, group.FieldForceOpenaiFast, group.FieldFreeOpenaiFast, group.FieldRequireOauthOnly, group.FieldRequirePrivacySet, group.FieldProfitControlEnabled:
|
||||
values[i] = new(sql.NullBool)
|
||||
@@ -662,12 +662,12 @@ func (_m *Group) assignValues(columns []string, values []any) error {
|
||||
return fmt.Errorf("unmarshal field messages_dispatch_model_config: %w", err)
|
||||
}
|
||||
}
|
||||
case group.FieldModelsListConfig:
|
||||
case group.FieldModelAllowlist:
|
||||
if value, ok := values[i].(*[]byte); !ok {
|
||||
return fmt.Errorf("unexpected type %T for field models_list_config", values[i])
|
||||
return fmt.Errorf("unexpected type %T for field model_allowlist", values[i])
|
||||
} else if value != nil && len(*value) > 0 {
|
||||
if err := json.Unmarshal(*value, &_m.ModelsListConfig); err != nil {
|
||||
return fmt.Errorf("unmarshal field models_list_config: %w", err)
|
||||
if err := json.Unmarshal(*value, &_m.ModelAllowlist); err != nil {
|
||||
return fmt.Errorf("unmarshal field model_allowlist: %w", err)
|
||||
}
|
||||
}
|
||||
case group.FieldCodexModelsManifestConfig:
|
||||
@@ -1007,8 +1007,8 @@ func (_m *Group) String() string {
|
||||
builder.WriteString("messages_dispatch_model_config=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.MessagesDispatchModelConfig))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("models_list_config=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.ModelsListConfig))
|
||||
builder.WriteString("model_allowlist=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.ModelAllowlist))
|
||||
builder.WriteString(", ")
|
||||
builder.WriteString("codex_models_manifest_config=")
|
||||
builder.WriteString(fmt.Sprintf("%v", _m.CodexModelsManifestConfig))
|
||||
|
||||
@@ -130,8 +130,8 @@ const (
|
||||
FieldDefaultMappedModel = "default_mapped_model"
|
||||
// FieldMessagesDispatchModelConfig holds the string denoting the messages_dispatch_model_config field in the database.
|
||||
FieldMessagesDispatchModelConfig = "messages_dispatch_model_config"
|
||||
// FieldModelsListConfig holds the string denoting the models_list_config field in the database.
|
||||
FieldModelsListConfig = "models_list_config"
|
||||
// FieldModelAllowlist holds the string denoting the model_allowlist field in the database.
|
||||
FieldModelAllowlist = "model_allowlist"
|
||||
// FieldCodexModelsManifestConfig holds the string denoting the codex_models_manifest_config field in the database.
|
||||
FieldCodexModelsManifestConfig = "codex_models_manifest_config"
|
||||
// FieldRpmLimit holds the string denoting the rpm_limit field in the database.
|
||||
@@ -280,7 +280,7 @@ var Columns = []string{
|
||||
FieldRequirePrivacySet,
|
||||
FieldDefaultMappedModel,
|
||||
FieldMessagesDispatchModelConfig,
|
||||
FieldModelsListConfig,
|
||||
FieldModelAllowlist,
|
||||
FieldCodexModelsManifestConfig,
|
||||
FieldRpmLimit,
|
||||
FieldMaxReasoningEffort,
|
||||
@@ -412,8 +412,8 @@ var (
|
||||
DefaultMappedModelValidator func(string) error
|
||||
// DefaultMessagesDispatchModelConfig holds the default value on creation for the "messages_dispatch_model_config" field.
|
||||
DefaultMessagesDispatchModelConfig domain.OpenAIMessagesDispatchModelConfig
|
||||
// DefaultModelsListConfig holds the default value on creation for the "models_list_config" field.
|
||||
DefaultModelsListConfig domain.GroupModelsListConfig
|
||||
// DefaultModelAllowlist holds the default value on creation for the "model_allowlist" field.
|
||||
DefaultModelAllowlist domain.GroupModelAllowlist
|
||||
// DefaultCodexModelsManifestConfig holds the default value on creation for the "codex_models_manifest_config" field.
|
||||
DefaultCodexModelsManifestConfig domain.GroupCodexModelsManifestConfig
|
||||
// DefaultRpmLimit holds the default value on creation for the "rpm_limit" field.
|
||||
|
||||
+32
-32
@@ -788,16 +788,16 @@ func (_c *GroupCreate) SetNillableMessagesDispatchModelConfig(v *domain.OpenAIMe
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (_c *GroupCreate) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupCreate {
|
||||
_c.mutation.SetModelsListConfig(v)
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (_c *GroupCreate) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupCreate {
|
||||
_c.mutation.SetModelAllowlist(v)
|
||||
return _c
|
||||
}
|
||||
|
||||
// SetNillableModelsListConfig sets the "models_list_config" field if the given value is not nil.
|
||||
func (_c *GroupCreate) SetNillableModelsListConfig(v *domain.GroupModelsListConfig) *GroupCreate {
|
||||
// SetNillableModelAllowlist sets the "model_allowlist" field if the given value is not nil.
|
||||
func (_c *GroupCreate) SetNillableModelAllowlist(v *domain.GroupModelAllowlist) *GroupCreate {
|
||||
if v != nil {
|
||||
_c.SetModelsListConfig(*v)
|
||||
_c.SetModelAllowlist(*v)
|
||||
}
|
||||
return _c
|
||||
}
|
||||
@@ -1175,9 +1175,9 @@ func (_c *GroupCreate) defaults() error {
|
||||
v := group.DefaultMessagesDispatchModelConfig
|
||||
_c.mutation.SetMessagesDispatchModelConfig(v)
|
||||
}
|
||||
if _, ok := _c.mutation.ModelsListConfig(); !ok {
|
||||
v := group.DefaultModelsListConfig
|
||||
_c.mutation.SetModelsListConfig(v)
|
||||
if _, ok := _c.mutation.ModelAllowlist(); !ok {
|
||||
v := group.DefaultModelAllowlist
|
||||
_c.mutation.SetModelAllowlist(v)
|
||||
}
|
||||
if _, ok := _c.mutation.CodexModelsManifestConfig(); !ok {
|
||||
v := group.DefaultCodexModelsManifestConfig
|
||||
@@ -1381,8 +1381,8 @@ func (_c *GroupCreate) check() error {
|
||||
if _, ok := _c.mutation.MessagesDispatchModelConfig(); !ok {
|
||||
return &ValidationError{Name: "messages_dispatch_model_config", err: errors.New(`ent: missing required field "Group.messages_dispatch_model_config"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.ModelsListConfig(); !ok {
|
||||
return &ValidationError{Name: "models_list_config", err: errors.New(`ent: missing required field "Group.models_list_config"`)}
|
||||
if _, ok := _c.mutation.ModelAllowlist(); !ok {
|
||||
return &ValidationError{Name: "model_allowlist", err: errors.New(`ent: missing required field "Group.model_allowlist"`)}
|
||||
}
|
||||
if _, ok := _c.mutation.CodexModelsManifestConfig(); !ok {
|
||||
return &ValidationError{Name: "codex_models_manifest_config", err: errors.New(`ent: missing required field "Group.codex_models_manifest_config"`)}
|
||||
@@ -1673,9 +1673,9 @@ func (_c *GroupCreate) createSpec() (*Group, *sqlgraph.CreateSpec) {
|
||||
_spec.SetField(group.FieldMessagesDispatchModelConfig, field.TypeJSON, value)
|
||||
_node.MessagesDispatchModelConfig = value
|
||||
}
|
||||
if value, ok := _c.mutation.ModelsListConfig(); ok {
|
||||
_spec.SetField(group.FieldModelsListConfig, field.TypeJSON, value)
|
||||
_node.ModelsListConfig = value
|
||||
if value, ok := _c.mutation.ModelAllowlist(); ok {
|
||||
_spec.SetField(group.FieldModelAllowlist, field.TypeJSON, value)
|
||||
_node.ModelAllowlist = value
|
||||
}
|
||||
if value, ok := _c.mutation.CodexModelsManifestConfig(); ok {
|
||||
_spec.SetField(group.FieldCodexModelsManifestConfig, field.TypeJSON, value)
|
||||
@@ -2795,15 +2795,15 @@ func (u *GroupUpsert) UpdateMessagesDispatchModelConfig() *GroupUpsert {
|
||||
return u
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (u *GroupUpsert) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpsert {
|
||||
u.Set(group.FieldModelsListConfig, v)
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (u *GroupUpsert) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpsert {
|
||||
u.Set(group.FieldModelAllowlist, v)
|
||||
return u
|
||||
}
|
||||
|
||||
// UpdateModelsListConfig sets the "models_list_config" field to the value that was provided on create.
|
||||
func (u *GroupUpsert) UpdateModelsListConfig() *GroupUpsert {
|
||||
u.SetExcluded(group.FieldModelsListConfig)
|
||||
// UpdateModelAllowlist sets the "model_allowlist" field to the value that was provided on create.
|
||||
func (u *GroupUpsert) UpdateModelAllowlist() *GroupUpsert {
|
||||
u.SetExcluded(group.FieldModelAllowlist)
|
||||
return u
|
||||
}
|
||||
|
||||
@@ -4054,17 +4054,17 @@ func (u *GroupUpsertOne) UpdateMessagesDispatchModelConfig() *GroupUpsertOne {
|
||||
})
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (u *GroupUpsertOne) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpsertOne {
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (u *GroupUpsertOne) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetModelsListConfig(v)
|
||||
s.SetModelAllowlist(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateModelsListConfig sets the "models_list_config" field to the value that was provided on create.
|
||||
func (u *GroupUpsertOne) UpdateModelsListConfig() *GroupUpsertOne {
|
||||
// UpdateModelAllowlist sets the "model_allowlist" field to the value that was provided on create.
|
||||
func (u *GroupUpsertOne) UpdateModelAllowlist() *GroupUpsertOne {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateModelsListConfig()
|
||||
s.UpdateModelAllowlist()
|
||||
})
|
||||
}
|
||||
|
||||
@@ -5500,17 +5500,17 @@ func (u *GroupUpsertBulk) UpdateMessagesDispatchModelConfig() *GroupUpsertBulk {
|
||||
})
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (u *GroupUpsertBulk) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpsertBulk {
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (u *GroupUpsertBulk) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.SetModelsListConfig(v)
|
||||
s.SetModelAllowlist(v)
|
||||
})
|
||||
}
|
||||
|
||||
// UpdateModelsListConfig sets the "models_list_config" field to the value that was provided on create.
|
||||
func (u *GroupUpsertBulk) UpdateModelsListConfig() *GroupUpsertBulk {
|
||||
// UpdateModelAllowlist sets the "model_allowlist" field to the value that was provided on create.
|
||||
func (u *GroupUpsertBulk) UpdateModelAllowlist() *GroupUpsertBulk {
|
||||
return u.Update(func(s *GroupUpsert) {
|
||||
s.UpdateModelsListConfig()
|
||||
s.UpdateModelAllowlist()
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
+16
-16
@@ -1073,16 +1073,16 @@ func (_u *GroupUpdate) SetNillableMessagesDispatchModelConfig(v *domain.OpenAIMe
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (_u *GroupUpdate) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpdate {
|
||||
_u.mutation.SetModelsListConfig(v)
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (_u *GroupUpdate) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpdate {
|
||||
_u.mutation.SetModelAllowlist(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableModelsListConfig sets the "models_list_config" field if the given value is not nil.
|
||||
func (_u *GroupUpdate) SetNillableModelsListConfig(v *domain.GroupModelsListConfig) *GroupUpdate {
|
||||
// SetNillableModelAllowlist sets the "model_allowlist" field if the given value is not nil.
|
||||
func (_u *GroupUpdate) SetNillableModelAllowlist(v *domain.GroupModelAllowlist) *GroupUpdate {
|
||||
if v != nil {
|
||||
_u.SetModelsListConfig(*v)
|
||||
_u.SetModelAllowlist(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
@@ -1876,8 +1876,8 @@ func (_u *GroupUpdate) sqlSave(ctx context.Context) (_node int, err error) {
|
||||
if value, ok := _u.mutation.MessagesDispatchModelConfig(); ok {
|
||||
_spec.SetField(group.FieldMessagesDispatchModelConfig, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.ModelsListConfig(); ok {
|
||||
_spec.SetField(group.FieldModelsListConfig, field.TypeJSON, value)
|
||||
if value, ok := _u.mutation.ModelAllowlist(); ok {
|
||||
_spec.SetField(group.FieldModelAllowlist, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.CodexModelsManifestConfig(); ok {
|
||||
_spec.SetField(group.FieldCodexModelsManifestConfig, field.TypeJSON, value)
|
||||
@@ -3267,16 +3267,16 @@ func (_u *GroupUpdateOne) SetNillableMessagesDispatchModelConfig(v *domain.OpenA
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (_u *GroupUpdateOne) SetModelsListConfig(v domain.GroupModelsListConfig) *GroupUpdateOne {
|
||||
_u.mutation.SetModelsListConfig(v)
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (_u *GroupUpdateOne) SetModelAllowlist(v domain.GroupModelAllowlist) *GroupUpdateOne {
|
||||
_u.mutation.SetModelAllowlist(v)
|
||||
return _u
|
||||
}
|
||||
|
||||
// SetNillableModelsListConfig sets the "models_list_config" field if the given value is not nil.
|
||||
func (_u *GroupUpdateOne) SetNillableModelsListConfig(v *domain.GroupModelsListConfig) *GroupUpdateOne {
|
||||
// SetNillableModelAllowlist sets the "model_allowlist" field if the given value is not nil.
|
||||
func (_u *GroupUpdateOne) SetNillableModelAllowlist(v *domain.GroupModelAllowlist) *GroupUpdateOne {
|
||||
if v != nil {
|
||||
_u.SetModelsListConfig(*v)
|
||||
_u.SetModelAllowlist(*v)
|
||||
}
|
||||
return _u
|
||||
}
|
||||
@@ -4100,8 +4100,8 @@ func (_u *GroupUpdateOne) sqlSave(ctx context.Context) (_node *Group, err error)
|
||||
if value, ok := _u.mutation.MessagesDispatchModelConfig(); ok {
|
||||
_spec.SetField(group.FieldMessagesDispatchModelConfig, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.ModelsListConfig(); ok {
|
||||
_spec.SetField(group.FieldModelsListConfig, field.TypeJSON, value)
|
||||
if value, ok := _u.mutation.ModelAllowlist(); ok {
|
||||
_spec.SetField(group.FieldModelAllowlist, field.TypeJSON, value)
|
||||
}
|
||||
if value, ok := _u.mutation.CodexModelsManifestConfig(); ok {
|
||||
_spec.SetField(group.FieldCodexModelsManifestConfig, field.TypeJSON, value)
|
||||
|
||||
@@ -960,7 +960,7 @@ var (
|
||||
{Name: "require_privacy_set", Type: field.TypeBool, Default: false},
|
||||
{Name: "default_mapped_model", Type: field.TypeString, Size: 100, Default: ""},
|
||||
{Name: "messages_dispatch_model_config", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "models_list_config", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "model_allowlist", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "codex_models_manifest_config", Type: field.TypeJSON, SchemaType: map[string]string{"postgres": "jsonb"}},
|
||||
{Name: "rpm_limit", Type: field.TypeInt, Default: 0},
|
||||
{Name: "max_reasoning_effort", Type: field.TypeString, Size: 20, Default: ""},
|
||||
|
||||
+27
-27
@@ -22164,7 +22164,7 @@ type GroupMutation struct {
|
||||
require_privacy_set *bool
|
||||
default_mapped_model *string
|
||||
messages_dispatch_model_config *domain.OpenAIMessagesDispatchModelConfig
|
||||
models_list_config *domain.GroupModelsListConfig
|
||||
model_allowlist *domain.GroupModelAllowlist
|
||||
codex_models_manifest_config *domain.GroupCodexModelsManifestConfig
|
||||
rpm_limit *int
|
||||
addrpm_limit *int
|
||||
@@ -25164,40 +25164,40 @@ func (m *GroupMutation) ResetMessagesDispatchModelConfig() {
|
||||
m.messages_dispatch_model_config = nil
|
||||
}
|
||||
|
||||
// SetModelsListConfig sets the "models_list_config" field.
|
||||
func (m *GroupMutation) SetModelsListConfig(dmlc domain.GroupModelsListConfig) {
|
||||
m.models_list_config = &dmlc
|
||||
// SetModelAllowlist sets the "model_allowlist" field.
|
||||
func (m *GroupMutation) SetModelAllowlist(dma domain.GroupModelAllowlist) {
|
||||
m.model_allowlist = &dma
|
||||
}
|
||||
|
||||
// ModelsListConfig returns the value of the "models_list_config" field in the mutation.
|
||||
func (m *GroupMutation) ModelsListConfig() (r domain.GroupModelsListConfig, exists bool) {
|
||||
v := m.models_list_config
|
||||
// ModelAllowlist returns the value of the "model_allowlist" field in the mutation.
|
||||
func (m *GroupMutation) ModelAllowlist() (r domain.GroupModelAllowlist, exists bool) {
|
||||
v := m.model_allowlist
|
||||
if v == nil {
|
||||
return
|
||||
}
|
||||
return *v, true
|
||||
}
|
||||
|
||||
// OldModelsListConfig returns the old "models_list_config" field's value of the Group entity.
|
||||
// OldModelAllowlist returns the old "model_allowlist" 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) OldModelsListConfig(ctx context.Context) (v domain.GroupModelsListConfig, err error) {
|
||||
func (m *GroupMutation) OldModelAllowlist(ctx context.Context) (v domain.GroupModelAllowlist, err error) {
|
||||
if !m.op.Is(OpUpdateOne) {
|
||||
return v, errors.New("OldModelsListConfig is only allowed on UpdateOne operations")
|
||||
return v, errors.New("OldModelAllowlist is only allowed on UpdateOne operations")
|
||||
}
|
||||
if m.id == nil || m.oldValue == nil {
|
||||
return v, errors.New("OldModelsListConfig requires an ID field in the mutation")
|
||||
return v, errors.New("OldModelAllowlist requires an ID field in the mutation")
|
||||
}
|
||||
oldValue, err := m.oldValue(ctx)
|
||||
if err != nil {
|
||||
return v, fmt.Errorf("querying old value for OldModelsListConfig: %w", err)
|
||||
return v, fmt.Errorf("querying old value for OldModelAllowlist: %w", err)
|
||||
}
|
||||
return oldValue.ModelsListConfig, nil
|
||||
return oldValue.ModelAllowlist, nil
|
||||
}
|
||||
|
||||
// ResetModelsListConfig resets all changes to the "models_list_config" field.
|
||||
func (m *GroupMutation) ResetModelsListConfig() {
|
||||
m.models_list_config = nil
|
||||
// ResetModelAllowlist resets all changes to the "model_allowlist" field.
|
||||
func (m *GroupMutation) ResetModelAllowlist() {
|
||||
m.model_allowlist = nil
|
||||
}
|
||||
|
||||
// SetCodexModelsManifestConfig sets the "codex_models_manifest_config" field.
|
||||
@@ -26093,8 +26093,8 @@ func (m *GroupMutation) Fields() []string {
|
||||
if m.messages_dispatch_model_config != nil {
|
||||
fields = append(fields, group.FieldMessagesDispatchModelConfig)
|
||||
}
|
||||
if m.models_list_config != nil {
|
||||
fields = append(fields, group.FieldModelsListConfig)
|
||||
if m.model_allowlist != nil {
|
||||
fields = append(fields, group.FieldModelAllowlist)
|
||||
}
|
||||
if m.codex_models_manifest_config != nil {
|
||||
fields = append(fields, group.FieldCodexModelsManifestConfig)
|
||||
@@ -26242,8 +26242,8 @@ func (m *GroupMutation) Field(name string) (ent.Value, bool) {
|
||||
return m.DefaultMappedModel()
|
||||
case group.FieldMessagesDispatchModelConfig:
|
||||
return m.MessagesDispatchModelConfig()
|
||||
case group.FieldModelsListConfig:
|
||||
return m.ModelsListConfig()
|
||||
case group.FieldModelAllowlist:
|
||||
return m.ModelAllowlist()
|
||||
case group.FieldCodexModelsManifestConfig:
|
||||
return m.CodexModelsManifestConfig()
|
||||
case group.FieldRpmLimit:
|
||||
@@ -26383,8 +26383,8 @@ func (m *GroupMutation) OldField(ctx context.Context, name string) (ent.Value, e
|
||||
return m.OldDefaultMappedModel(ctx)
|
||||
case group.FieldMessagesDispatchModelConfig:
|
||||
return m.OldMessagesDispatchModelConfig(ctx)
|
||||
case group.FieldModelsListConfig:
|
||||
return m.OldModelsListConfig(ctx)
|
||||
case group.FieldModelAllowlist:
|
||||
return m.OldModelAllowlist(ctx)
|
||||
case group.FieldCodexModelsManifestConfig:
|
||||
return m.OldCodexModelsManifestConfig(ctx)
|
||||
case group.FieldRpmLimit:
|
||||
@@ -26809,12 +26809,12 @@ func (m *GroupMutation) SetField(name string, value ent.Value) error {
|
||||
}
|
||||
m.SetMessagesDispatchModelConfig(v)
|
||||
return nil
|
||||
case group.FieldModelsListConfig:
|
||||
v, ok := value.(domain.GroupModelsListConfig)
|
||||
case group.FieldModelAllowlist:
|
||||
v, ok := value.(domain.GroupModelAllowlist)
|
||||
if !ok {
|
||||
return fmt.Errorf("unexpected type %T for field %s", value, name)
|
||||
}
|
||||
m.SetModelsListConfig(v)
|
||||
m.SetModelAllowlist(v)
|
||||
return nil
|
||||
case group.FieldCodexModelsManifestConfig:
|
||||
v, ok := value.(domain.GroupCodexModelsManifestConfig)
|
||||
@@ -27554,8 +27554,8 @@ func (m *GroupMutation) ResetField(name string) error {
|
||||
case group.FieldMessagesDispatchModelConfig:
|
||||
m.ResetMessagesDispatchModelConfig()
|
||||
return nil
|
||||
case group.FieldModelsListConfig:
|
||||
m.ResetModelsListConfig()
|
||||
case group.FieldModelAllowlist:
|
||||
m.ResetModelAllowlist()
|
||||
return nil
|
||||
case group.FieldCodexModelsManifestConfig:
|
||||
m.ResetCodexModelsManifestConfig()
|
||||
|
||||
@@ -1183,10 +1183,10 @@ func init() {
|
||||
groupDescMessagesDispatchModelConfig := groupFields[53].Descriptor()
|
||||
// group.DefaultMessagesDispatchModelConfig holds the default value on creation for the messages_dispatch_model_config field.
|
||||
group.DefaultMessagesDispatchModelConfig = groupDescMessagesDispatchModelConfig.Default.(domain.OpenAIMessagesDispatchModelConfig)
|
||||
// groupDescModelsListConfig is the schema descriptor for models_list_config field.
|
||||
groupDescModelsListConfig := groupFields[54].Descriptor()
|
||||
// group.DefaultModelsListConfig holds the default value on creation for the models_list_config field.
|
||||
group.DefaultModelsListConfig = groupDescModelsListConfig.Default.(domain.GroupModelsListConfig)
|
||||
// groupDescModelAllowlist is the schema descriptor for model_allowlist field.
|
||||
groupDescModelAllowlist := groupFields[54].Descriptor()
|
||||
// group.DefaultModelAllowlist holds the default value on creation for the model_allowlist field.
|
||||
group.DefaultModelAllowlist = groupDescModelAllowlist.Default.(domain.GroupModelAllowlist)
|
||||
// groupDescCodexModelsManifestConfig is the schema descriptor for codex_models_manifest_config field.
|
||||
groupDescCodexModelsManifestConfig := groupFields[55].Descriptor()
|
||||
// group.DefaultCodexModelsManifestConfig holds the default value on creation for the codex_models_manifest_config field.
|
||||
|
||||
@@ -262,10 +262,10 @@ func (Group) Fields() []ent.Field {
|
||||
Default(domain.OpenAIMessagesDispatchModelConfig{}).
|
||||
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
|
||||
Comment("OpenAI Messages 调度模型配置:按 Claude 系列/精确模型映射到目标 GPT 模型"),
|
||||
field.JSON("models_list_config", domain.GroupModelsListConfig{}).
|
||||
Default(domain.GroupModelsListConfig{}).
|
||||
field.JSON("model_allowlist", domain.GroupModelAllowlist{}).
|
||||
Default(domain.GroupModelAllowlist{}).
|
||||
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
|
||||
Comment("自定义 /v1/models 展示列表配置;仅影响模型列表响应,不影响调度"),
|
||||
Comment("分组模型白名单:同时约束模型列表接口与请求准入"),
|
||||
field.JSON("codex_models_manifest_config", domain.GroupCodexModelsManifestConfig{}).
|
||||
Default(domain.GroupCodexModelsManifestConfig{}).
|
||||
SchemaType(map[string]string{dialect.Postgres: "jsonb"}).
|
||||
|
||||
@@ -0,0 +1,8 @@
|
||||
package domain
|
||||
|
||||
// GroupModelAllowlist 是分组级模型白名单:开启后既过滤模型列表类接口的返回内容,
|
||||
// 也约束分组实际可调用的模型(网关准入在合成路由改写与调度之前完成)。
|
||||
type GroupModelAllowlist struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Models []string `json:"models,omitempty"`
|
||||
}
|
||||
@@ -1,7 +0,0 @@
|
||||
package domain
|
||||
|
||||
// GroupModelsListConfig controls the optional custom /v1/models response list.
|
||||
type GroupModelsListConfig struct {
|
||||
Enabled bool `json:"enabled"`
|
||||
Models []string `json:"models,omitempty"`
|
||||
}
|
||||
@@ -7,7 +7,8 @@ const (
|
||||
)
|
||||
|
||||
// ReasoningEffortMapping rewrites one explicit OpenAI/Codex reasoning effort
|
||||
// value to another before the group ceiling is applied.
|
||||
// value to another before the group ceiling is applied. To may also be "deny"
|
||||
// to reject the request when the matching source value is present.
|
||||
//
|
||||
// Model and MatchType optionally scope the rewrite to a request model:
|
||||
// exact matches the full model id, prefix/suffix match a model-id affix.
|
||||
|
||||
@@ -17,6 +17,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/domain"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
@@ -64,6 +65,7 @@ type AccountHandler struct {
|
||||
grokImportProber grokImportProber
|
||||
upstreamBillingProbe *service.UpstreamBillingProbeService
|
||||
ollamaCloudUsage *service.OllamaCloudUsageService
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
// SetUpstreamBillingProbeService attaches the optional remote billing probe service.
|
||||
@@ -190,6 +192,7 @@ type CheckMixedChannelRequest struct {
|
||||
// AccountWithConcurrency extends Account with real-time concurrency info
|
||||
type AccountWithConcurrency struct {
|
||||
*dto.Account
|
||||
simpleMode bool `json:"-"`
|
||||
CurrentConcurrency int `json:"current_concurrency"`
|
||||
SchedulerScore *AccountSchedulerScore `json:"scheduler_score,omitempty"`
|
||||
SchedulerScores []AccountSchedulerGroupScore `json:"scheduler_scores,omitempty"`
|
||||
@@ -212,6 +215,110 @@ type AccountListItemWithConcurrency struct {
|
||||
CurrentRPM *int `json:"current_rpm,omitempty"`
|
||||
}
|
||||
|
||||
type simpleModeGroupReference struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Platform string `json:"platform"`
|
||||
Status string `json:"status"`
|
||||
}
|
||||
|
||||
type simpleModeAccountGroupReference struct {
|
||||
AccountID int64 `json:"account_id"`
|
||||
GroupID int64 `json:"group_id"`
|
||||
Priority int `json:"priority"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
Group *simpleModeGroupReference `json:"group,omitempty"`
|
||||
}
|
||||
|
||||
func simpleModeGroupReferenceFromDTO(group *dto.Group) *simpleModeGroupReference {
|
||||
if group == nil {
|
||||
return nil
|
||||
}
|
||||
return &simpleModeGroupReference{ID: group.ID, Name: group.Name, Platform: group.Platform, Status: group.Status}
|
||||
}
|
||||
|
||||
func simpleModeCompositeGroupIDs(account *dto.Account) map[int64]struct{} {
|
||||
hidden := make(map[int64]struct{})
|
||||
if account == nil {
|
||||
return hidden
|
||||
}
|
||||
for _, group := range account.Groups {
|
||||
if group != nil && group.Platform == service.PlatformComposite {
|
||||
hidden[group.ID] = struct{}{}
|
||||
}
|
||||
}
|
||||
for _, accountGroup := range account.AccountGroups {
|
||||
if accountGroup.Group != nil && accountGroup.Group.Platform == service.PlatformComposite {
|
||||
hidden[accountGroup.GroupID] = struct{}{}
|
||||
}
|
||||
}
|
||||
return hidden
|
||||
}
|
||||
|
||||
func filterSimpleModeGroupIDs(groupIDs []int64, hidden map[int64]struct{}) []int64 {
|
||||
visible := make([]int64, 0, len(groupIDs))
|
||||
for _, groupID := range groupIDs {
|
||||
if _, ok := hidden[groupID]; !ok {
|
||||
visible = append(visible, groupID)
|
||||
}
|
||||
}
|
||||
return visible
|
||||
}
|
||||
|
||||
func simpleModeCompositeServiceGroupIDs(account *service.Account) map[int64]struct{} {
|
||||
hidden := make(map[int64]struct{})
|
||||
if account == nil {
|
||||
return hidden
|
||||
}
|
||||
for _, group := range account.Groups {
|
||||
if group != nil && group.Platform == service.PlatformComposite {
|
||||
hidden[group.ID] = struct{}{}
|
||||
}
|
||||
}
|
||||
for _, accountGroup := range account.AccountGroups {
|
||||
if accountGroup.Group != nil && accountGroup.Group.Platform == service.PlatformComposite {
|
||||
hidden[accountGroup.GroupID] = struct{}{}
|
||||
}
|
||||
}
|
||||
return hidden
|
||||
}
|
||||
|
||||
func (a AccountWithConcurrency) MarshalJSON() ([]byte, error) {
|
||||
type alias AccountWithConcurrency
|
||||
if !a.simpleMode || a.Account == nil {
|
||||
return json.Marshal(alias(a))
|
||||
}
|
||||
groups := make([]simpleModeGroupReference, 0, len(a.Groups))
|
||||
compositeIDs := simpleModeCompositeGroupIDs(a.Account)
|
||||
for _, group := range a.Groups {
|
||||
if group != nil && group.Platform == service.PlatformComposite {
|
||||
continue
|
||||
}
|
||||
if ref := simpleModeGroupReferenceFromDTO(group); ref != nil {
|
||||
groups = append(groups, *ref)
|
||||
}
|
||||
}
|
||||
accountGroups := make([]simpleModeAccountGroupReference, 0, len(a.AccountGroups))
|
||||
for _, accountGroup := range a.AccountGroups {
|
||||
if accountGroup.Group != nil && accountGroup.Group.Platform == service.PlatformComposite {
|
||||
continue
|
||||
}
|
||||
if _, hidden := compositeIDs[accountGroup.GroupID]; hidden {
|
||||
continue
|
||||
}
|
||||
accountGroups = append(accountGroups, simpleModeAccountGroupReference{
|
||||
AccountID: accountGroup.AccountID, GroupID: accountGroup.GroupID, Priority: accountGroup.Priority,
|
||||
CreatedAt: accountGroup.CreatedAt, Group: simpleModeGroupReferenceFromDTO(accountGroup.Group),
|
||||
})
|
||||
}
|
||||
return json.Marshal(struct {
|
||||
alias
|
||||
GroupIDs []int64 `json:"group_ids,omitempty"`
|
||||
Groups []simpleModeGroupReference `json:"groups"`
|
||||
AccountGroups []simpleModeAccountGroupReference `json:"account_groups"`
|
||||
}{alias: alias(a), GroupIDs: filterSimpleModeGroupIDs(a.GroupIDs, compositeIDs), Groups: groups, AccountGroups: accountGroups})
|
||||
}
|
||||
|
||||
type AccountSchedulerScore struct {
|
||||
BaseScore float64 `json:"base_score"`
|
||||
StickyScore float64 `json:"sticky_score"`
|
||||
@@ -247,9 +354,14 @@ func (h *AccountHandler) accountListResponseFromService(account *service.Account
|
||||
return out
|
||||
}
|
||||
|
||||
func (h *AccountHandler) isSimpleMode() bool {
|
||||
return h != nil && h.cfg != nil && h.cfg.RunMode == config.RunModeSimple
|
||||
}
|
||||
|
||||
func (h *AccountHandler) buildAccountResponseWithRuntime(ctx context.Context, account *service.Account) AccountWithConcurrency {
|
||||
item := AccountWithConcurrency{
|
||||
Account: h.accountResponseFromService(account),
|
||||
simpleMode: h.isSimpleMode(),
|
||||
CurrentConcurrency: 0,
|
||||
}
|
||||
if account == nil {
|
||||
@@ -677,9 +789,13 @@ func (h *AccountHandler) List(c *gin.Context) {
|
||||
accountResponse := h.accountResponseFromService(acc)
|
||||
if lite {
|
||||
accountResponse = h.accountListResponseFromService(acc)
|
||||
if h.isSimpleMode() {
|
||||
accountResponse.GroupIDs = filterSimpleModeGroupIDs(accountResponse.GroupIDs, simpleModeCompositeServiceGroupIDs(acc))
|
||||
}
|
||||
}
|
||||
item := AccountWithConcurrency{
|
||||
Account: accountResponse,
|
||||
simpleMode: h.isSimpleMode(),
|
||||
CurrentConcurrency: concurrencyCounts[acc.ID],
|
||||
SchedulerScore: schedulerScores[acc.ID],
|
||||
SchedulerScores: schedulerGroupScores[acc.ID],
|
||||
@@ -1942,6 +2058,14 @@ func (h *AccountHandler) BatchCreate(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
}
|
||||
groupIDs := make([]int64, 0)
|
||||
for _, item := range req.Accounts {
|
||||
groupIDs = append(groupIDs, item.GroupIDs...)
|
||||
}
|
||||
if err := h.adminService.ValidateAccountGroupBindings(c.Request.Context(), groupIDs); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
}
|
||||
|
||||
executeAdminIdempotentJSON(c, "admin.accounts.batch_create", req, service.DefaultWriteIdempotencyTTL(), func(ctx context.Context) (any, error) {
|
||||
success := 0
|
||||
@@ -2656,6 +2780,14 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) {
|
||||
|
||||
// Handle OpenAI accounts
|
||||
if account.IsOpenAI() {
|
||||
// Prefer the shared, account-keyed upstream catalog. If discovery fails,
|
||||
// retain the legacy local catalog below so the test dialog remains usable.
|
||||
if h.accountTestService != nil {
|
||||
if models, fetchErr := h.accountTestService.FetchOpenAIAccountModels(c.Request.Context(), account); fetchErr == nil {
|
||||
response.Success(c, models)
|
||||
return
|
||||
}
|
||||
}
|
||||
// OpenAI 自动透传会绕过常规模型改写,测试/模型列表也应回落到默认模型集。
|
||||
if account.IsOpenAIPassthroughEnabled() {
|
||||
response.Success(c, openai.DefaultModels)
|
||||
|
||||
@@ -0,0 +1,240 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type simpleModeAccountService struct {
|
||||
*stubAdminService
|
||||
account service.Account
|
||||
createCalls int
|
||||
updateCalls int
|
||||
bulkCalls int
|
||||
}
|
||||
|
||||
func (s *simpleModeAccountService) GetAccount(context.Context, int64) (*service.Account, error) {
|
||||
return &s.account, nil
|
||||
}
|
||||
|
||||
func (s *simpleModeAccountService) CreateAccount(ctx context.Context, input *service.CreateAccountInput) (*service.Account, error) {
|
||||
if err := s.ValidateAccountGroupBindings(ctx, input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.createCalls++
|
||||
return &s.account, nil
|
||||
}
|
||||
|
||||
func (s *simpleModeAccountService) UpdateAccount(ctx context.Context, _ int64, input *service.UpdateAccountInput) (*service.Account, error) {
|
||||
if input.GroupIDs != nil {
|
||||
if err := s.ValidateAccountGroupBindings(ctx, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
s.updateCalls++
|
||||
return &s.account, nil
|
||||
}
|
||||
|
||||
func (s *simpleModeAccountService) BulkUpdateAccounts(ctx context.Context, input *service.BulkUpdateAccountsInput) (*service.BulkUpdateAccountsResult, error) {
|
||||
if input.GroupIDs != nil {
|
||||
if err := s.ValidateAccountGroupBindings(ctx, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
s.bulkCalls++
|
||||
return &service.BulkUpdateAccountsResult{}, nil
|
||||
}
|
||||
|
||||
func (s *simpleModeAccountService) ValidateAccountGroupBindings(_ context.Context, groupIDs []int64) error {
|
||||
for _, id := range groupIDs {
|
||||
for i := range s.groups {
|
||||
if s.groups[i].ID == id && s.groups[i].Platform == service.PlatformComposite {
|
||||
return infraerrors.BadRequest("SIMPLE_MODE_GROUP_NOT_BINDABLE", "composite groups cannot be bound in simple mode")
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func TestAccountHandlerSimpleModeUsesMinimalGroupReferences(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
richGroup := &service.Group{ID: 7, Name: "basic", Platform: service.PlatformAnthropic, Status: service.StatusActive, RateMultiplier: 9, IsExclusive: true, SubscriptionType: service.SubscriptionTypeSubscription, RPMLimit: 99}
|
||||
historicalComposite := &service.Group{ID: 9, Name: "historical composite", Platform: service.PlatformComposite, Status: service.StatusActive}
|
||||
account := service.Account{
|
||||
ID: 3, Name: "account", Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey, Status: service.StatusActive,
|
||||
GroupIDs: []int64{7, 9},
|
||||
Groups: []*service.Group{richGroup, historicalComposite},
|
||||
AccountGroups: []service.AccountGroup{
|
||||
{AccountID: 3, GroupID: 7, Priority: 2, Group: richGroup, Account: &service.Account{ID: 3, Extra: map[string]any{"secret": true}}},
|
||||
{AccountID: 3, GroupID: 9, Priority: 3, Group: historicalComposite},
|
||||
},
|
||||
}
|
||||
svc := &simpleModeAccountService{stubAdminService: newStubAdminService(), account: account}
|
||||
svc.accounts = []service.Account{account}
|
||||
h := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h.cfg = &config.Config{RunMode: config.RunModeSimple}
|
||||
r := gin.New()
|
||||
r.GET("/accounts", h.List)
|
||||
r.GET("/accounts/:id", h.GetByID)
|
||||
r.POST("/accounts", h.Create)
|
||||
r.PUT("/accounts/:id", h.Update)
|
||||
|
||||
tests := []struct {
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{http.MethodGet, "/accounts", ""},
|
||||
{http.MethodGet, "/accounts/3", ""},
|
||||
{http.MethodPost, "/accounts", `{"name":"account","platform":"anthropic","type":"apikey","credentials":{"key":"x"}}`},
|
||||
{http.MethodPut, "/accounts/3", `{"name":"account"}`},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.method+tt.path, func(t *testing.T) {
|
||||
req := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusOK, res.Code, res.Body.String())
|
||||
var payload map[string]any
|
||||
require.NoError(t, json.Unmarshal(res.Body.Bytes(), &payload))
|
||||
data, ok := payload["data"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
if items, ok := data["items"].([]any); ok {
|
||||
require.Len(t, items, 1)
|
||||
data, ok = items[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
}
|
||||
groups, ok := data["groups"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, groups, 1)
|
||||
require.Equal(t, []any{float64(7)}, data["group_ids"])
|
||||
group, ok := groups[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.ElementsMatch(t, []string{"id", "name", "platform", "status"}, mapKeys(group))
|
||||
accountGroups, ok := data["account_groups"].([]any)
|
||||
require.True(t, ok)
|
||||
require.Len(t, accountGroups, 1)
|
||||
accountGroup, ok := accountGroups[0].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.ElementsMatch(t, []string{"account_id", "group_id", "priority", "created_at", "group"}, mapKeys(accountGroup))
|
||||
nestedGroup, ok := accountGroup["group"].(map[string]any)
|
||||
require.True(t, ok)
|
||||
require.ElementsMatch(t, []string{"id", "name", "platform", "status"}, mapKeys(nestedGroup))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountHandlerSimpleModeLitePreservesCompactShapeAndETag(t *testing.T) {
|
||||
richGroup := &service.Group{ID: 7, Name: "basic", Platform: service.PlatformAnthropic, Status: service.StatusActive, RateMultiplier: 9}
|
||||
historicalComposite := &service.Group{ID: 9, Name: "historical composite", Platform: service.PlatformComposite, Status: service.StatusActive}
|
||||
account := service.Account{
|
||||
ID: 3, Name: "account", Platform: service.PlatformAnthropic, Type: service.AccountTypeAPIKey, Status: service.StatusActive,
|
||||
GroupIDs: []int64{7, 9}, Groups: []*service.Group{richGroup, historicalComposite},
|
||||
AccountGroups: []service.AccountGroup{{AccountID: 3, GroupID: 7, Group: richGroup}, {AccountID: 3, GroupID: 9, Group: historicalComposite}},
|
||||
}
|
||||
svc := &simpleModeAccountService{stubAdminService: newStubAdminService(), account: account}
|
||||
svc.accounts = []service.Account{account}
|
||||
h := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h.cfg = &config.Config{RunMode: config.RunModeSimple}
|
||||
r := gin.New()
|
||||
r.GET("/accounts", h.List)
|
||||
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, httptest.NewRequest(http.MethodGet, "/accounts?lite=1", nil))
|
||||
require.Equal(t, http.StatusOK, res.Code, res.Body.String())
|
||||
require.NotEmpty(t, res.Header().Get("ETag"))
|
||||
var payload struct {
|
||||
Data struct {
|
||||
Items []map[string]any `json:"items"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(res.Body.Bytes(), &payload))
|
||||
require.Len(t, payload.Data.Items, 1)
|
||||
require.Equal(t, []any{float64(7)}, payload.Data.Items[0]["group_ids"])
|
||||
require.NotContains(t, payload.Data.Items[0], "groups")
|
||||
require.NotContains(t, payload.Data.Items[0], "account_groups")
|
||||
|
||||
notModified := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodGet, "/accounts?lite=1", nil)
|
||||
req.Header.Set("If-None-Match", res.Header().Get("ETag"))
|
||||
r.ServeHTTP(notModified, req)
|
||||
require.Equal(t, http.StatusNotModified, notModified.Code)
|
||||
}
|
||||
|
||||
func TestAccountHandlerSimpleModeRejectsCompositeGroupBindingsBeforeWrites(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
for _, tt := range []struct{ name, method, path, body string }{
|
||||
{"create", http.MethodPost, "/accounts", `{"name":"account","platform":"anthropic","type":"apikey","credentials":{"key":"x"},"group_ids":[9]}`},
|
||||
{"update", http.MethodPut, "/accounts/3", `{"group_ids":[9]}`},
|
||||
{"bulk", http.MethodPost, "/accounts/bulk-update", `{"account_ids":[3],"group_ids":[9]}`},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
svc := &simpleModeAccountService{stubAdminService: newStubAdminService()}
|
||||
svc.groups = []service.Group{{ID: 7, Platform: service.PlatformAnthropic}, {ID: 9, Platform: service.PlatformComposite}}
|
||||
h := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h.cfg = &config.Config{RunMode: config.RunModeSimple}
|
||||
r := gin.New()
|
||||
r.POST("/accounts", h.Create)
|
||||
r.PUT("/accounts/:id", h.Update)
|
||||
r.POST("/accounts/bulk-update", h.BulkUpdate)
|
||||
res := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusBadRequest, res.Code, res.Body.String())
|
||||
require.Zero(t, svc.createCalls)
|
||||
require.Zero(t, svc.updateCalls)
|
||||
require.Zero(t, svc.bulkCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAccountHandlerSimpleModePreservesBasicGroupBinding(t *testing.T) {
|
||||
svc := &simpleModeAccountService{stubAdminService: newStubAdminService(), account: service.Account{ID: 3}}
|
||||
svc.groups = []service.Group{{ID: 7, Platform: service.PlatformAnthropic}}
|
||||
h := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h.cfg = &config.Config{RunMode: config.RunModeSimple}
|
||||
r := gin.New()
|
||||
r.POST("/accounts", h.Create)
|
||||
res := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/accounts", bytes.NewBufferString(`{"name":"account","platform":"anthropic","type":"apikey","credentials":{"key":"x"},"group_ids":[7]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusOK, res.Code, res.Body.String())
|
||||
require.Equal(t, 1, svc.createCalls)
|
||||
}
|
||||
|
||||
func TestAccountHandlerSimpleModeBatchPrevalidatesAllGroupsAtomically(t *testing.T) {
|
||||
svc := &simpleModeAccountService{stubAdminService: newStubAdminService()}
|
||||
svc.groups = []service.Group{{ID: 7, Platform: service.PlatformAnthropic}, {ID: 9, Platform: service.PlatformComposite}}
|
||||
h := NewAccountHandler(svc, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h.cfg = &config.Config{RunMode: config.RunModeSimple}
|
||||
r := gin.New()
|
||||
r.POST("/accounts/batch", h.BatchCreate)
|
||||
res := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(http.MethodPost, "/accounts/batch", bytes.NewBufferString(`{"accounts":[{"name":"first","platform":"anthropic","type":"apikey","credentials":{"key":"x"},"group_ids":[7]},{"name":"second","platform":"anthropic","type":"apikey","credentials":{"key":"y"},"group_ids":[9]}]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusBadRequest, res.Code, res.Body.String())
|
||||
require.Zero(t, svc.createCalls)
|
||||
}
|
||||
|
||||
func TestAccountHandlerAdvancedModeKeepsFullGroupReferences(t *testing.T) {
|
||||
group := &service.Group{ID: 7, Name: "advanced", RateMultiplier: 9, SubscriptionType: service.SubscriptionTypeSubscription}
|
||||
h := NewAccountHandler(nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
raw, err := json.Marshal(h.buildAccountResponseWithRuntime(context.Background(), &service.Account{Groups: []*service.Group{group}}))
|
||||
require.NoError(t, err)
|
||||
require.Contains(t, string(raw), `"rate_multiplier":9`)
|
||||
require.Contains(t, string(raw), `"subscription_type":"subscription"`)
|
||||
}
|
||||
@@ -33,7 +33,7 @@ func setupAdminRouter() (*gin.Engine, *stubAdminService) {
|
||||
|
||||
router.GET("/api/v1/admin/groups", groupHandler.List)
|
||||
router.GET("/api/v1/admin/groups/all", groupHandler.GetAll)
|
||||
router.GET("/api/v1/admin/groups/:id/models-list-candidates", groupHandler.GetModelsListCandidates)
|
||||
router.GET("/api/v1/admin/groups/:id/model-allowlist-candidates", groupHandler.GetGroupModelAllowlistCandidates)
|
||||
router.GET("/api/v1/admin/groups/:id/composite-routes", groupHandler.ListCompositeRoutes)
|
||||
router.POST("/api/v1/admin/groups/:id/composite-routes", groupHandler.CreateCompositeRoute)
|
||||
router.POST("/api/v1/admin/groups/:id/composite-routes/preview", groupHandler.PreviewCompositeRoute)
|
||||
@@ -184,7 +184,7 @@ func TestGroupHandlerEndpoints(t *testing.T) {
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
rec = httptest.NewRecorder()
|
||||
req = httptest.NewRequest(http.MethodGet, "/api/v1/admin/groups/0/models-list-candidates?platform=openai", nil)
|
||||
req = httptest.NewRequest(http.MethodGet, "/api/v1/admin/groups/0/model-allowlist-candidates?platform=openai", nil)
|
||||
router.ServeHTTP(rec, req)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
require.Contains(t, rec.Body.String(), "gpt-5.5")
|
||||
|
||||
@@ -24,6 +24,13 @@ type stubAdminService struct {
|
||||
boundAuthIdentity *service.AdminBindAuthIdentityInput
|
||||
boundAuthIdentityFor int64
|
||||
createdAccounts []*service.CreateAccountInput
|
||||
createdGroups []*service.CreateGroupInput
|
||||
updatedGroups []*service.UpdateGroupInput
|
||||
deletedGroupIDs []int64
|
||||
guardedDeletedGroupIDs []int64
|
||||
deleteGroupIfEmptyErr error
|
||||
advancedGroupOperationCalls int
|
||||
lastListGroupsIsExclusive *bool
|
||||
createdProxies []*service.CreateProxyInput
|
||||
updatedProxyIDs []int64
|
||||
updatedProxies []*service.UpdateProxyInput
|
||||
@@ -267,6 +274,7 @@ func (s *stubAdminService) BindUserAuthIdentity(ctx context.Context, userID int6
|
||||
}
|
||||
|
||||
func (s *stubAdminService) ListGroups(ctx context.Context, page, pageSize int, platform, status, search string, isExclusive *bool, sortBy, sortOrder string) ([]service.Group, int64, error) {
|
||||
s.lastListGroupsIsExclusive = isExclusive
|
||||
return s.groups, int64(len(s.groups)), nil
|
||||
}
|
||||
|
||||
@@ -295,6 +303,7 @@ func (s *stubAdminService) GetGroupModelsListCandidates(ctx context.Context, id
|
||||
}
|
||||
|
||||
func (s *stubAdminService) ListCompositeRoutes(ctx context.Context, groupID int64) ([]service.CompositeModelRoute, error) {
|
||||
s.advancedGroupOperationCalls++
|
||||
return []service.CompositeModelRoute{
|
||||
{
|
||||
ID: 1,
|
||||
@@ -311,6 +320,7 @@ func (s *stubAdminService) ListCompositeRoutes(ctx context.Context, groupID int6
|
||||
}
|
||||
|
||||
func (s *stubAdminService) CreateCompositeRoute(ctx context.Context, groupID int64, input service.CompositeRouteInput) (*service.CompositeModelRoute, error) {
|
||||
s.advancedGroupOperationCalls++
|
||||
return &service.CompositeModelRoute{
|
||||
ID: 1,
|
||||
GroupID: groupID,
|
||||
@@ -326,6 +336,7 @@ func (s *stubAdminService) CreateCompositeRoute(ctx context.Context, groupID int
|
||||
}
|
||||
|
||||
func (s *stubAdminService) UpdateCompositeRoute(ctx context.Context, groupID, routeID int64, input service.CompositeRouteInput) (*service.CompositeModelRoute, error) {
|
||||
s.advancedGroupOperationCalls++
|
||||
return &service.CompositeModelRoute{
|
||||
ID: routeID,
|
||||
GroupID: groupID,
|
||||
@@ -341,10 +352,12 @@ func (s *stubAdminService) UpdateCompositeRoute(ctx context.Context, groupID, ro
|
||||
}
|
||||
|
||||
func (s *stubAdminService) DeleteCompositeRoute(ctx context.Context, groupID, routeID int64) error {
|
||||
s.advancedGroupOperationCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) PreviewCompositeRoute(ctx context.Context, groupID int64, input service.CompositeRoutePreviewRequest) (*service.CompositeRouteDecision, error) {
|
||||
s.advancedGroupOperationCalls++
|
||||
decision, err := service.NewCompositeRouteResolver(nil).Resolve(ctx, groupID, input.Model, input.Endpoint)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
@@ -353,11 +366,13 @@ func (s *stubAdminService) PreviewCompositeRoute(ctx context.Context, groupID in
|
||||
}
|
||||
|
||||
func (s *stubAdminService) CreateGroup(ctx context.Context, input *service.CreateGroupInput) (*service.Group, error) {
|
||||
s.createdGroups = append(s.createdGroups, input)
|
||||
group := service.Group{ID: 200, Name: input.Name, Status: service.StatusActive}
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) DuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*service.Group, error) {
|
||||
s.advancedGroupOperationCalls++
|
||||
group := service.Group{ID: 201, Name: "group (Copy)", Status: "inactive"}
|
||||
return &group, nil
|
||||
}
|
||||
@@ -367,14 +382,21 @@ func (s *stubAdminService) RecoverDuplicateGroup(ctx context.Context, id int64,
|
||||
}
|
||||
|
||||
func (s *stubAdminService) UpdateGroup(ctx context.Context, id int64, input *service.UpdateGroupInput) (*service.Group, error) {
|
||||
s.updatedGroups = append(s.updatedGroups, input)
|
||||
group := service.Group{ID: id, Name: input.Name, Status: service.StatusActive}
|
||||
return &group, nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) DeleteGroup(ctx context.Context, id int64) error {
|
||||
s.deletedGroupIDs = append(s.deletedGroupIDs, id)
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) DeleteGroupIfEmpty(ctx context.Context, id int64) error {
|
||||
s.guardedDeletedGroupIDs = append(s.guardedDeletedGroupIDs, id)
|
||||
return s.deleteGroupIfEmptyErr
|
||||
}
|
||||
|
||||
func (s *stubAdminService) GetGroupAPIKeys(ctx context.Context, groupID int64, page, pageSize int) ([]service.APIKey, int64, error) {
|
||||
return s.apiKeys, int64(len(s.apiKeys)), nil
|
||||
}
|
||||
@@ -384,10 +406,12 @@ func (s *stubAdminService) GetGroupRateMultipliers(_ context.Context, _ int64) (
|
||||
}
|
||||
|
||||
func (s *stubAdminService) ClearGroupRateMultipliers(_ context.Context, _ int64) error {
|
||||
s.advancedGroupOperationCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) BatchSetGroupRateMultipliers(_ context.Context, _ int64, _ []service.GroupRateMultiplierInput) error {
|
||||
s.advancedGroupOperationCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -396,6 +420,7 @@ func (s *stubAdminService) ClearGroupRPMOverrides(_ context.Context, _ int64) er
|
||||
}
|
||||
|
||||
func (s *stubAdminService) BatchSetGroupRPMOverrides(_ context.Context, _ int64, _ []service.GroupRPMOverrideInput) error {
|
||||
s.advancedGroupOperationCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
@@ -491,6 +516,8 @@ func (s *stubAdminService) CreateAccount(ctx context.Context, input *service.Cre
|
||||
return &account, nil
|
||||
}
|
||||
|
||||
func (s *stubAdminService) ValidateAccountGroupBindings(context.Context, []int64) error { return nil }
|
||||
|
||||
func (s *stubAdminService) DuplicateAccount(ctx context.Context, id int64, actorScope, operationKey string) (*service.Account, error) {
|
||||
account := service.Account{ID: 301, Name: "account (Copy)", Status: service.StatusActive, Schedulable: false}
|
||||
return &account, nil
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
@@ -194,6 +195,7 @@ func (h *GrokOAuthHandler) scheduleGrokImportProbe(account *service.Account) {
|
||||
// ProvideAccountHandler injects the Grok active prober for production while
|
||||
// keeping NewAccountHandler convenient for focused unit tests.
|
||||
func ProvideAccountHandler(
|
||||
cfg *config.Config,
|
||||
adminService service.AdminService,
|
||||
oauthService *service.OAuthService,
|
||||
openaiOAuthService *service.OpenAIOAuthService,
|
||||
@@ -227,5 +229,6 @@ func ProvideAccountHandler(
|
||||
tokenCacheInvalidator,
|
||||
)
|
||||
handler.grokImportProber = grokQuotaService
|
||||
handler.cfg = cfg
|
||||
return handler
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler/dto"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/response"
|
||||
@@ -24,10 +25,14 @@ type GroupHandler struct {
|
||||
adminService service.AdminService
|
||||
dashboardService *service.DashboardService
|
||||
groupCapacityService *service.GroupCapacityService
|
||||
cfg *config.Config
|
||||
}
|
||||
|
||||
// GetLiveCapability 返回当前服务端是否具备生成 Live attestation 的运行环境。
|
||||
func (h *GroupHandler) GetLiveCapability(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "live_capability") {
|
||||
return
|
||||
}
|
||||
err := liveattestation.NewProvider().Check(c.Request.Context())
|
||||
result := gin.H{"supported": err == nil}
|
||||
if err != nil {
|
||||
@@ -87,13 +92,94 @@ func (f optionalLimitField) ToServiceInput() *float64 {
|
||||
|
||||
// NewGroupHandler creates a new admin group handler
|
||||
func NewGroupHandler(adminService service.AdminService, dashboardService *service.DashboardService, groupCapacityService *service.GroupCapacityService) *GroupHandler {
|
||||
return NewGroupHandlerWithConfig(adminService, dashboardService, groupCapacityService, nil)
|
||||
}
|
||||
|
||||
func NewGroupHandlerWithConfig(adminService service.AdminService, dashboardService *service.DashboardService, groupCapacityService *service.GroupCapacityService, cfg *config.Config) *GroupHandler {
|
||||
return &GroupHandler{
|
||||
adminService: adminService,
|
||||
dashboardService: dashboardService,
|
||||
groupCapacityService: groupCapacityService,
|
||||
cfg: cfg,
|
||||
}
|
||||
}
|
||||
|
||||
func (h *GroupHandler) isSimpleMode() bool {
|
||||
return h != nil && h.cfg != nil && h.cfg.RunMode == config.RunModeSimple
|
||||
}
|
||||
|
||||
type simpleModeGroupOperation string
|
||||
|
||||
const (
|
||||
simpleModeGroupList simpleModeGroupOperation = "list"
|
||||
simpleModeGroupGetAll simpleModeGroupOperation = "get_all"
|
||||
simpleModeGroupGet simpleModeGroupOperation = "get"
|
||||
simpleModeGroupCreate simpleModeGroupOperation = "create"
|
||||
simpleModeGroupUpdate simpleModeGroupOperation = "update"
|
||||
simpleModeGroupDelete simpleModeGroupOperation = "delete"
|
||||
)
|
||||
|
||||
var simpleModeGroupOperations = map[simpleModeGroupOperation]struct{}{
|
||||
simpleModeGroupList: {}, simpleModeGroupGetAll: {}, simpleModeGroupGet: {},
|
||||
simpleModeGroupCreate: {}, simpleModeGroupUpdate: {}, simpleModeGroupDelete: {},
|
||||
}
|
||||
|
||||
func (h *GroupHandler) rejectUnsupportedSimpleModeOperation(c *gin.Context, operation simpleModeGroupOperation) bool {
|
||||
if _, allowed := simpleModeGroupOperations[operation]; allowed {
|
||||
return false
|
||||
}
|
||||
if err := service.ValidateSimpleModeGroupOperation(h.cfg, service.AdminGroupOperation(operation)); err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
type simpleModeGroupResponse struct {
|
||||
ID int64 `json:"id"`
|
||||
Name string `json:"name"`
|
||||
Description string `json:"description"`
|
||||
Platform string `json:"platform"`
|
||||
Status string `json:"status"`
|
||||
|
||||
AccountCount int64 `json:"account_count,omitempty"`
|
||||
ActiveAccountCount int64 `json:"active_account_count,omitempty"`
|
||||
RateLimitedAccountCount int64 `json:"rate_limited_account_count,omitempty"`
|
||||
SortOrder int `json:"sort_order"`
|
||||
CreatedAt time.Time `json:"created_at"`
|
||||
UpdatedAt time.Time `json:"updated_at"`
|
||||
}
|
||||
|
||||
func groupForSimpleMode(group *service.Group) *simpleModeGroupResponse {
|
||||
if group == nil {
|
||||
return nil
|
||||
}
|
||||
return &simpleModeGroupResponse{
|
||||
ID: group.ID, Name: group.Name, Description: group.Description, Platform: group.Platform,
|
||||
Status: group.Status,
|
||||
AccountCount: group.AccountCount,
|
||||
ActiveAccountCount: group.ActiveAccountCount, RateLimitedAccountCount: group.RateLimitedAccountCount,
|
||||
SortOrder: group.SortOrder, CreatedAt: group.CreatedAt, UpdatedAt: group.UpdatedAt,
|
||||
}
|
||||
}
|
||||
|
||||
func sanitizeCreateGroupRequestForSimpleMode(req *CreateGroupRequest) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
allowed := CreateGroupRequest{Name: req.Name, Description: req.Description, Platform: req.Platform}
|
||||
allowed.RateMultiplier = 1
|
||||
allowed.SubscriptionType = service.SubscriptionTypeStandard
|
||||
*req = allowed
|
||||
}
|
||||
|
||||
func sanitizeUpdateGroupRequestForSimpleMode(req *UpdateGroupRequest) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
*req = UpdateGroupRequest{Name: req.Name, Description: req.Description}
|
||||
}
|
||||
|
||||
// CreateGroupRequest represents create group request
|
||||
type CreateGroupRequest struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
@@ -153,7 +239,7 @@ type CreateGroupRequest struct {
|
||||
RequirePrivacySet bool `json:"require_privacy_set"`
|
||||
DefaultMappedModel string `json:"default_mapped_model"`
|
||||
MessagesDispatchModelConfig service.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config"`
|
||||
ModelsListConfig service.GroupModelsListConfig `json:"models_list_config"`
|
||||
ModelAllowlist service.GroupModelAllowlist `json:"model_allowlist"`
|
||||
// 固定账号 manifest 配置;创建路径禁止开启,仅编辑可配置。
|
||||
CodexModelsManifestConfig service.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config"`
|
||||
// 分组 RPM 上限(0 = 不限制)
|
||||
@@ -228,7 +314,7 @@ type UpdateGroupRequest struct {
|
||||
RequirePrivacySet *bool `json:"require_privacy_set"`
|
||||
DefaultMappedModel *string `json:"default_mapped_model"`
|
||||
MessagesDispatchModelConfig *service.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config"`
|
||||
ModelsListConfig *service.GroupModelsListConfig `json:"models_list_config"`
|
||||
ModelAllowlist *service.GroupModelAllowlist `json:"model_allowlist"`
|
||||
// 固定账号 manifest 配置;nil 表示不修改。
|
||||
CodexModelsManifestConfig *service.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config"`
|
||||
// 分组 RPM 上限(0 = 不限制);nil 表示未提供不改动
|
||||
@@ -262,6 +348,9 @@ type CompositeRoutePreviewRequest struct {
|
||||
// List handles listing all groups with pagination
|
||||
// GET /api/v1/admin/groups
|
||||
func (h *GroupHandler) List(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, simpleModeGroupList) {
|
||||
return
|
||||
}
|
||||
page, pageSize := response.ParsePagination(c)
|
||||
platform := c.Query("platform")
|
||||
status := c.Query("status")
|
||||
@@ -276,7 +365,7 @@ func (h *GroupHandler) List(c *gin.Context) {
|
||||
sortOrder := c.DefaultQuery("sort_order", "asc")
|
||||
|
||||
var isExclusive *bool
|
||||
if isExclusiveStr != "" {
|
||||
if !h.isSimpleMode() && isExclusiveStr != "" {
|
||||
val := isExclusiveStr == "true"
|
||||
isExclusive = &val
|
||||
}
|
||||
@@ -287,6 +376,16 @@ func (h *GroupHandler) List(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.isSimpleMode() {
|
||||
simpleGroups := make([]simpleModeGroupResponse, 0, len(groups))
|
||||
for i := range groups {
|
||||
if service.IsGroupBindableInSimpleMode(&groups[i]) {
|
||||
simpleGroups = append(simpleGroups, *groupForSimpleMode(&groups[i]))
|
||||
}
|
||||
}
|
||||
response.Paginated(c, simpleGroups, total, page, pageSize)
|
||||
return
|
||||
}
|
||||
outGroups := make([]dto.AdminGroup, 0, len(groups))
|
||||
for i := range groups {
|
||||
outGroups = append(outGroups, *dto.GroupFromServiceAdmin(&groups[i]))
|
||||
@@ -297,6 +396,9 @@ func (h *GroupHandler) List(c *gin.Context) {
|
||||
// ListCompositeRoutes handles listing composite model routes for one group.
|
||||
// GET /api/v1/admin/groups/:id/composite-routes
|
||||
func (h *GroupHandler) ListCompositeRoutes(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, ok := parsePositiveIDParam(c, "id")
|
||||
if !ok {
|
||||
return
|
||||
@@ -312,6 +414,9 @@ func (h *GroupHandler) ListCompositeRoutes(c *gin.Context) {
|
||||
// CreateCompositeRoute handles creating one composite model route.
|
||||
// POST /api/v1/admin/groups/:id/composite-routes
|
||||
func (h *GroupHandler) CreateCompositeRoute(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, ok := parsePositiveIDParam(c, "id")
|
||||
if !ok {
|
||||
return
|
||||
@@ -332,6 +437,9 @@ func (h *GroupHandler) CreateCompositeRoute(c *gin.Context) {
|
||||
// UpdateCompositeRoute handles replacing one composite model route.
|
||||
// PUT /api/v1/admin/groups/:id/composite-routes/:route_id
|
||||
func (h *GroupHandler) UpdateCompositeRoute(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, ok := parsePositiveIDParam(c, "id")
|
||||
if !ok {
|
||||
return
|
||||
@@ -356,6 +464,9 @@ func (h *GroupHandler) UpdateCompositeRoute(c *gin.Context) {
|
||||
// DeleteCompositeRoute handles deleting one composite model route.
|
||||
// DELETE /api/v1/admin/groups/:id/composite-routes/:route_id
|
||||
func (h *GroupHandler) DeleteCompositeRoute(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, ok := parsePositiveIDParam(c, "id")
|
||||
if !ok {
|
||||
return
|
||||
@@ -374,6 +485,9 @@ func (h *GroupHandler) DeleteCompositeRoute(c *gin.Context) {
|
||||
// PreviewCompositeRoute resolves a model without mutating routes.
|
||||
// POST /api/v1/admin/groups/:id/composite-routes/preview
|
||||
func (h *GroupHandler) PreviewCompositeRoute(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, ok := parsePositiveIDParam(c, "id")
|
||||
if !ok {
|
||||
return
|
||||
@@ -427,6 +541,9 @@ func parsePositiveIDParam(c *gin.Context, name string) (int64, bool) {
|
||||
// bound to them even after the group is disabled).
|
||||
// GET /api/v1/admin/groups/all
|
||||
func (h *GroupHandler) GetAll(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, simpleModeGroupGetAll) {
|
||||
return
|
||||
}
|
||||
platform := c.Query("platform")
|
||||
includeInactive := c.Query("include_inactive") == "true"
|
||||
|
||||
@@ -446,6 +563,16 @@ func (h *GroupHandler) GetAll(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.isSimpleMode() {
|
||||
simpleGroups := make([]simpleModeGroupResponse, 0, len(groups))
|
||||
for i := range groups {
|
||||
if service.IsGroupBindableInSimpleMode(&groups[i]) {
|
||||
simpleGroups = append(simpleGroups, *groupForSimpleMode(&groups[i]))
|
||||
}
|
||||
}
|
||||
response.Success(c, simpleGroups)
|
||||
return
|
||||
}
|
||||
outGroups := make([]dto.AdminGroup, 0, len(groups))
|
||||
for i := range groups {
|
||||
outGroups = append(outGroups, *dto.GroupFromServiceAdmin(&groups[i]))
|
||||
@@ -456,6 +583,9 @@ func (h *GroupHandler) GetAll(c *gin.Context) {
|
||||
// GetByID handles getting a group by ID
|
||||
// GET /api/v1/admin/groups/:id
|
||||
func (h *GroupHandler) GetByID(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, simpleModeGroupGet) {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -468,12 +598,19 @@ func (h *GroupHandler) GetByID(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.isSimpleMode() {
|
||||
response.Success(c, groupForSimpleMode(group))
|
||||
return
|
||||
}
|
||||
response.Success(c, dto.GroupFromServiceAdmin(group))
|
||||
}
|
||||
|
||||
// GetModelsListCandidates handles getting candidate model IDs for custom /v1/models list.
|
||||
// GET /api/v1/admin/groups/:id/models-list-candidates
|
||||
func (h *GroupHandler) GetModelsListCandidates(c *gin.Context) {
|
||||
// GetGroupModelAllowlistCandidates handles getting candidate model IDs for the group model allowlist.
|
||||
// GET /api/v1/admin/groups/:id/model-allowlist-candidates
|
||||
func (h *GroupHandler) GetGroupModelAllowlistCandidates(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || groupID < 0 {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -496,11 +633,21 @@ func (h *GroupHandler) GetModelsListCandidates(c *gin.Context) {
|
||||
// Create handles creating a new group
|
||||
// POST /api/v1/admin/groups
|
||||
func (h *GroupHandler) Create(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, simpleModeGroupCreate) {
|
||||
return
|
||||
}
|
||||
var req CreateGroupRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if h.isSimpleMode() && req.Platform == service.PlatformComposite {
|
||||
response.BadRequest(c, "Platform is not supported in simple mode")
|
||||
return
|
||||
}
|
||||
if h.isSimpleMode() {
|
||||
sanitizeCreateGroupRequestForSimpleMode(&req)
|
||||
}
|
||||
|
||||
if err := service.ValidatePeakRateConfig(req.SubscriptionType, req.PeakRateEnabled, req.PeakStart, req.PeakEnd, float64ValueOrDefault(req.PeakRateMultiplier, 1.0)); err != nil {
|
||||
response.BadRequest(c, err.Error())
|
||||
@@ -568,7 +715,7 @@ func (h *GroupHandler) Create(c *gin.Context) {
|
||||
RequirePrivacySet: req.RequirePrivacySet,
|
||||
DefaultMappedModel: req.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: req.MessagesDispatchModelConfig,
|
||||
ModelsListConfig: req.ModelsListConfig,
|
||||
ModelAllowlist: req.ModelAllowlist,
|
||||
CodexModelsManifestConfig: req.CodexModelsManifestConfig,
|
||||
RPMLimit: req.RPMLimit,
|
||||
MaxReasoningEffort: req.MaxReasoningEffort,
|
||||
@@ -581,12 +728,19 @@ func (h *GroupHandler) Create(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.isSimpleMode() {
|
||||
response.Success(c, groupForSimpleMode(group))
|
||||
return
|
||||
}
|
||||
response.Success(c, dto.GroupFromServiceAdmin(group))
|
||||
}
|
||||
|
||||
// Duplicate handles creating an inactive group copy with the source account bindings.
|
||||
// POST /api/v1/admin/groups/:id/duplicate
|
||||
func (h *GroupHandler) Duplicate(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil || groupID <= 0 {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -634,6 +788,9 @@ func (h *GroupHandler) Duplicate(c *gin.Context) {
|
||||
// Update handles updating a group
|
||||
// PUT /api/v1/admin/groups/:id
|
||||
func (h *GroupHandler) Update(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, simpleModeGroupUpdate) {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -645,6 +802,9 @@ func (h *GroupHandler) Update(c *gin.Context) {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
return
|
||||
}
|
||||
if h.isSimpleMode() {
|
||||
sanitizeUpdateGroupRequestForSimpleMode(&req)
|
||||
}
|
||||
|
||||
group, err := h.adminService.UpdateGroup(c.Request.Context(), groupID, &service.UpdateGroupInput{
|
||||
Name: req.Name,
|
||||
@@ -701,7 +861,7 @@ func (h *GroupHandler) Update(c *gin.Context) {
|
||||
RequirePrivacySet: req.RequirePrivacySet,
|
||||
DefaultMappedModel: req.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: req.MessagesDispatchModelConfig,
|
||||
ModelsListConfig: req.ModelsListConfig,
|
||||
ModelAllowlist: req.ModelAllowlist,
|
||||
CodexModelsManifestConfig: req.CodexModelsManifestConfig,
|
||||
RPMLimit: req.RPMLimit,
|
||||
MaxReasoningEffort: req.MaxReasoningEffort,
|
||||
@@ -714,19 +874,30 @@ func (h *GroupHandler) Update(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
if h.isSimpleMode() {
|
||||
response.Success(c, groupForSimpleMode(group))
|
||||
return
|
||||
}
|
||||
response.Success(c, dto.GroupFromServiceAdmin(group))
|
||||
}
|
||||
|
||||
// Delete handles deleting a group
|
||||
// DELETE /api/v1/admin/groups/:id
|
||||
func (h *GroupHandler) Delete(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, simpleModeGroupDelete) {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
return
|
||||
}
|
||||
|
||||
err = h.adminService.DeleteGroup(c.Request.Context(), groupID)
|
||||
if h.isSimpleMode() {
|
||||
err = h.adminService.DeleteGroupIfEmpty(c.Request.Context(), groupID)
|
||||
} else {
|
||||
err = h.adminService.DeleteGroup(c.Request.Context(), groupID)
|
||||
}
|
||||
if err != nil {
|
||||
response.ErrorFrom(c, err)
|
||||
return
|
||||
@@ -738,6 +909,9 @@ func (h *GroupHandler) Delete(c *gin.Context) {
|
||||
// GetStats handles getting group statistics
|
||||
// GET /api/v1/admin/groups/:id/stats
|
||||
func (h *GroupHandler) GetStats(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "stats") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -757,6 +931,9 @@ func (h *GroupHandler) GetStats(c *gin.Context) {
|
||||
// GetUsageSummary returns today's, yesterday's, and cumulative cost for all groups.
|
||||
// GET /api/v1/admin/groups/usage-summary
|
||||
func (h *GroupHandler) GetUsageSummary(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
todayStart := service.GroupUsageTodayStart(time.Now())
|
||||
|
||||
results, err := h.dashboardService.GetGroupUsageSummary(c.Request.Context(), todayStart)
|
||||
@@ -771,6 +948,9 @@ func (h *GroupHandler) GetUsageSummary(c *gin.Context) {
|
||||
// GetCapacitySummary returns aggregated capacity (concurrency/sessions/RPM) for all active groups.
|
||||
// GET /api/v1/admin/groups/capacity-summary
|
||||
func (h *GroupHandler) GetCapacitySummary(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
results, err := h.groupCapacityService.GetAllGroupCapacity(c.Request.Context())
|
||||
if err != nil {
|
||||
response.Error(c, 500, "Failed to get group capacity summary")
|
||||
@@ -782,6 +962,9 @@ func (h *GroupHandler) GetCapacitySummary(c *gin.Context) {
|
||||
// GetGroupAPIKeys handles getting API keys in a group
|
||||
// GET /api/v1/admin/groups/:id/api-keys
|
||||
func (h *GroupHandler) GetGroupAPIKeys(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "api_keys") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -806,6 +989,9 @@ func (h *GroupHandler) GetGroupAPIKeys(c *gin.Context) {
|
||||
// GetGroupRateMultipliers handles getting rate multipliers for users in a group
|
||||
// GET /api/v1/admin/groups/:id/rate-multipliers
|
||||
func (h *GroupHandler) GetGroupRateMultipliers(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -827,6 +1013,9 @@ func (h *GroupHandler) GetGroupRateMultipliers(c *gin.Context) {
|
||||
// ClearGroupRateMultipliers handles clearing all rate multipliers for a group
|
||||
// DELETE /api/v1/admin/groups/:id/rate-multipliers
|
||||
func (h *GroupHandler) ClearGroupRateMultipliers(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -849,6 +1038,9 @@ type BatchSetGroupRateMultipliersRequest struct {
|
||||
// BatchSetGroupRateMultipliers handles batch setting rate multipliers for a group
|
||||
// PUT /api/v1/admin/groups/:id/rate-multipliers
|
||||
func (h *GroupHandler) BatchSetGroupRateMultipliers(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -877,6 +1069,9 @@ type BatchSetGroupRPMOverridesRequest struct {
|
||||
// BatchSetGroupRPMOverrides handles batch setting rpm_override for users in a group
|
||||
// PUT /api/v1/admin/groups/:id/rpm-overrides
|
||||
func (h *GroupHandler) BatchSetGroupRPMOverrides(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -900,6 +1095,9 @@ func (h *GroupHandler) BatchSetGroupRPMOverrides(c *gin.Context) {
|
||||
// ClearGroupRPMOverrides handles clearing all rpm_override for a group
|
||||
// DELETE /api/v1/admin/groups/:id/rpm-overrides
|
||||
func (h *GroupHandler) ClearGroupRPMOverrides(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
|
||||
if err != nil {
|
||||
response.BadRequest(c, "Invalid group ID")
|
||||
@@ -925,6 +1123,9 @@ type UpdateSortOrderRequest struct {
|
||||
// UpdateSortOrder handles updating group sort orders
|
||||
// PUT /api/v1/admin/groups/sort-order
|
||||
func (h *GroupHandler) UpdateSortOrder(c *gin.Context) {
|
||||
if h.rejectUnsupportedSimpleModeOperation(c, "advanced") {
|
||||
return
|
||||
}
|
||||
var req UpdateSortOrderRequest
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
response.BadRequest(c, "Invalid request: "+err.Error())
|
||||
|
||||
@@ -0,0 +1,309 @@
|
||||
package admin
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newSimpleModeGroupRouter(svc *stubAdminService) *gin.Engine {
|
||||
h := NewGroupHandlerWithConfig(svc, nil, nil, &config.Config{RunMode: config.RunModeSimple})
|
||||
r := gin.New()
|
||||
r.GET("/groups", h.List)
|
||||
r.GET("/groups/all", h.GetAll)
|
||||
r.GET("/groups/live-capability", h.GetLiveCapability)
|
||||
r.GET("/groups/usage-summary", h.GetUsageSummary)
|
||||
r.GET("/groups/capacity-summary", h.GetCapacitySummary)
|
||||
r.PUT("/groups/sort-order", h.UpdateSortOrder)
|
||||
r.GET("/groups/:id", h.GetByID)
|
||||
r.GET("/groups/:id/stats", h.GetStats)
|
||||
r.GET("/groups/:id/api-keys", h.GetGroupAPIKeys)
|
||||
r.GET("/groups/:id/model-allowlist-candidates", h.GetGroupModelAllowlistCandidates)
|
||||
r.POST("/groups", h.Create)
|
||||
r.PUT("/groups/:id", h.Update)
|
||||
r.POST("/groups/:id/duplicate", h.Duplicate)
|
||||
r.GET("/groups/:id/composite-routes", h.ListCompositeRoutes)
|
||||
r.POST("/groups/:id/composite-routes", h.CreateCompositeRoute)
|
||||
r.POST("/groups/:id/composite-routes/preview", h.PreviewCompositeRoute)
|
||||
r.PUT("/groups/:id/composite-routes/:route_id", h.UpdateCompositeRoute)
|
||||
r.DELETE("/groups/:id/composite-routes/:route_id", h.DeleteCompositeRoute)
|
||||
r.GET("/groups/:id/rate-multipliers", h.GetGroupRateMultipliers)
|
||||
r.PUT("/groups/:id/rate-multipliers", h.BatchSetGroupRateMultipliers)
|
||||
r.DELETE("/groups/:id/rate-multipliers", h.ClearGroupRateMultipliers)
|
||||
r.PUT("/groups/:id/rpm-overrides", h.BatchSetGroupRPMOverrides)
|
||||
r.DELETE("/groups/:id/rpm-overrides", h.ClearGroupRPMOverrides)
|
||||
r.DELETE("/groups/:id", h.Delete)
|
||||
return r
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeSanitizesCommercialFields(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newStubAdminService()
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
|
||||
create := `{"name":"simple","description":"basic grouping","platform":"anthropic","rate_multiplier":7,"is_exclusive":true,"subscription_type":"subscription","daily_limit_usd":10,"long_context_pricing_enabled":true,"model_pricing":[{"model":"claude","input_price":1}],"allow_image_generation":true,"allow_batch_image_generation":true,"video_price_720p":2,"web_search_price_per_call":3,"audio_realtime_price_per_min":4,"rpm_limit":99}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/groups", bytes.NewBufferString(create))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
require.Len(t, svc.createdGroups, 1)
|
||||
created := svc.createdGroups[0]
|
||||
require.Equal(t, "basic grouping", created.Description)
|
||||
require.Equal(t, 1.0, created.RateMultiplier)
|
||||
require.False(t, created.IsExclusive)
|
||||
require.Equal(t, service.SubscriptionTypeStandard, created.SubscriptionType)
|
||||
require.Nil(t, created.DailyLimitUSD)
|
||||
require.False(t, created.AllowImageGeneration)
|
||||
require.False(t, created.LongContextPricingEnabled)
|
||||
require.Empty(t, created.ModelPricing)
|
||||
require.False(t, created.AllowBatchImageGeneration)
|
||||
require.Nil(t, created.VideoPrice720P)
|
||||
require.Nil(t, created.WebSearchPricePerCall)
|
||||
require.Nil(t, created.AudioRealtimePricePerMin)
|
||||
require.Zero(t, created.RPMLimit)
|
||||
|
||||
update := `{"name":"renamed","description":"still basic","rate_multiplier":9,"is_exclusive":true,"subscription_type":"subscription","daily_limit_usd":12,"long_context_pricing_enabled":true,"model_pricing":[{"model":"claude","input_price":1}],"allow_image_generation":true,"allow_batch_image_generation":true,"video_price_720p":2,"web_search_price_per_call":3,"audio_realtime_price_per_min":4,"status":"inactive","rpm_limit":123}`
|
||||
req = httptest.NewRequest(http.MethodPut, "/groups/2", bytes.NewBufferString(update))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
res = httptest.NewRecorder()
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
require.Len(t, svc.updatedGroups, 1)
|
||||
updated := svc.updatedGroups[0]
|
||||
require.NotNil(t, updated.Description)
|
||||
require.Equal(t, "still basic", *updated.Description)
|
||||
require.Nil(t, updated.RateMultiplier)
|
||||
require.Nil(t, updated.IsExclusive)
|
||||
require.Empty(t, updated.SubscriptionType)
|
||||
require.Nil(t, updated.DailyLimitUSD)
|
||||
require.Nil(t, updated.AllowImageGeneration)
|
||||
require.Nil(t, updated.LongContextPricingEnabled)
|
||||
require.Nil(t, updated.ModelPricing)
|
||||
require.Nil(t, updated.AllowBatchImageGeneration)
|
||||
require.Nil(t, updated.VideoPrice720P)
|
||||
require.Nil(t, updated.WebSearchPricePerCall)
|
||||
require.Nil(t, updated.AudioRealtimePricePerMin)
|
||||
require.Empty(t, updated.Status)
|
||||
require.Nil(t, updated.RPMLimit)
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeRejectsCompositeBeforeService(t *testing.T) {
|
||||
svc := newStubAdminService()
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
req := httptest.NewRequest(http.MethodPost, "/groups", bytes.NewBufferString(`{"name":"composite","platform":"composite"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusBadRequest, res.Code)
|
||||
require.Empty(t, svc.createdGroups)
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeDeleteRejectsNonEmptyGroupBeforeCascade(t *testing.T) {
|
||||
svc := newStubAdminService()
|
||||
svc.deleteGroupIfEmptyErr = service.ErrGroupNotEmpty
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
res := httptest.NewRecorder()
|
||||
|
||||
r.ServeHTTP(res, httptest.NewRequest(http.MethodDelete, "/groups/7", nil))
|
||||
|
||||
require.Equal(t, http.StatusConflict, res.Code)
|
||||
require.Equal(t, []int64{7}, svc.guardedDeletedGroupIDs)
|
||||
require.Empty(t, svc.deletedGroupIDs)
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeDeletePermitsEmptyGroup(t *testing.T) {
|
||||
svc := newStubAdminService()
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
res := httptest.NewRecorder()
|
||||
|
||||
r.ServeHTTP(res, httptest.NewRequest(http.MethodDelete, "/groups/8", nil))
|
||||
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
require.Equal(t, []int64{8}, svc.guardedDeletedGroupIDs)
|
||||
require.Empty(t, svc.deletedGroupIDs)
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeIgnoresExclusiveFilter(t *testing.T) {
|
||||
svc := newStubAdminService()
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, httptest.NewRequest(http.MethodGet, "/groups?is_exclusive=true", nil))
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
require.Nil(t, svc.lastListGroupsIsExclusive)
|
||||
}
|
||||
|
||||
func float64PtrForSimpleModeTest(value float64) *float64 { return &value }
|
||||
|
||||
func TestGroupHandlerSimpleModeResponseUsesFieldAllowlist(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newStubAdminService()
|
||||
svc.groups = []service.Group{{
|
||||
ID: 1, Name: "basic", Description: "allowed", Platform: service.PlatformAnthropic,
|
||||
AccountCount: 3, ActiveAccountCount: 2, RateLimitedAccountCount: 1,
|
||||
Status: service.StatusActive, RateMultiplier: 9, RPMLimit: 42,
|
||||
LongContextPricingEnabled: true,
|
||||
ModelPricing: []service.ChannelModelPricing{{Models: []string{"claude"}}},
|
||||
AllowBatchImageGeneration: true, VideoPrice720P: float64PtrForSimpleModeTest(2),
|
||||
WebSearchPricePerCall: float64PtrForSimpleModeTest(3), AudioRealtimePricePerMin: float64PtrForSimpleModeTest(4),
|
||||
ModelRouting: map[string][]int64{"claude": {2}},
|
||||
}}
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, httptest.NewRequest(http.MethodGet, "/groups", nil))
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
|
||||
var payload struct {
|
||||
Data struct {
|
||||
Items []map[string]any `json:"items"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(res.Body.Bytes(), &payload))
|
||||
require.Len(t, payload.Data.Items, 1)
|
||||
item := payload.Data.Items[0]
|
||||
require.Equal(t, "basic", item["name"])
|
||||
require.Equal(t, "allowed", item["description"])
|
||||
require.Equal(t, "anthropic", item["platform"])
|
||||
require.Equal(t, float64(3), item["account_count"])
|
||||
require.ElementsMatch(t, []string{
|
||||
"id", "name", "description", "platform", "status", "account_count",
|
||||
"active_account_count", "rate_limited_account_count", "sort_order", "created_at", "updated_at",
|
||||
}, mapKeys(item))
|
||||
for _, forbidden := range []string{
|
||||
"rate_multiplier", "rpm_limit", "long_context_pricing_enabled", "model_pricing",
|
||||
"allow_batch_image_generation", "video_price_720p", "web_search_price_per_call",
|
||||
"audio_realtime_price_per_min", "model_routing",
|
||||
} {
|
||||
_, exposed := item[forbidden]
|
||||
require.Falsef(t, exposed, "simple mode exposed %s", forbidden)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeListsOnlyBindableGroups(t *testing.T) {
|
||||
svc := newStubAdminService()
|
||||
svc.groups = []service.Group{{ID: 1, Name: "basic", Platform: service.PlatformAnthropic}, {ID: 2, Name: "composite", Platform: service.PlatformComposite}}
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
for _, path := range []string{"/groups", "/groups/all"} {
|
||||
res := httptest.NewRecorder()
|
||||
r.ServeHTTP(res, httptest.NewRequest(http.MethodGet, path, nil))
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
require.Contains(t, res.Body.String(), `"name":"basic"`)
|
||||
require.NotContains(t, res.Body.String(), `"name":"composite"`)
|
||||
}
|
||||
}
|
||||
|
||||
func mapKeys(values map[string]any) []string {
|
||||
keys := make([]string, 0, len(values))
|
||||
for key := range values {
|
||||
keys = append(keys, key)
|
||||
}
|
||||
return keys
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeAllReadAndWriteResponsesUseFieldAllowlist(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
svc := newStubAdminService()
|
||||
svc.groups = []service.Group{{
|
||||
ID: 1, Name: "basic", Description: "allowed", Platform: service.PlatformAnthropic, AccountCount: 3,
|
||||
Status: service.StatusActive, RateMultiplier: 9, RPMLimit: 42,
|
||||
LongContextPricingEnabled: true,
|
||||
ModelPricing: []service.ChannelModelPricing{{Models: []string{"claude"}}},
|
||||
}}
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
|
||||
assertNoAdvancedFields := func(t *testing.T, item map[string]any) {
|
||||
t.Helper()
|
||||
for _, forbidden := range []string{
|
||||
"rate_multiplier", "rpm_limit", "long_context_pricing_enabled", "model_pricing",
|
||||
"allow_batch_image_generation", "video_price_720p", "web_search_price_per_call",
|
||||
"audio_realtime_price_per_min", "model_routing",
|
||||
} {
|
||||
_, exposed := item[forbidden]
|
||||
require.Falsef(t, exposed, "simple mode exposed %s", forbidden)
|
||||
}
|
||||
}
|
||||
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
list bool
|
||||
}{
|
||||
{name: "get all", method: http.MethodGet, path: "/groups/all", list: true},
|
||||
{name: "get by id", method: http.MethodGet, path: "/groups/1"},
|
||||
{name: "create", method: http.MethodPost, path: "/groups", body: `{"name":"created","platform":"anthropic","model_pricing":[{"models":["claude"]}]}`},
|
||||
{name: "update", method: http.MethodPut, path: "/groups/1", body: `{"name":"updated","model_pricing":[{"models":["claude"]}]}`},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
res := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusOK, res.Code)
|
||||
|
||||
var payload struct {
|
||||
Data json.RawMessage `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(res.Body.Bytes(), &payload))
|
||||
if tt.list {
|
||||
var items []map[string]any
|
||||
require.NoError(t, json.Unmarshal(payload.Data, &items))
|
||||
require.Len(t, items, 1)
|
||||
assertNoAdvancedFields(t, items[0])
|
||||
return
|
||||
}
|
||||
var item map[string]any
|
||||
require.NoError(t, json.Unmarshal(payload.Data, &item))
|
||||
assertNoAdvancedFields(t, item)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupHandlerSimpleModeBlocksAdvancedOperations(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
tests := []struct {
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{http.MethodPost, "/groups/1/duplicate", ""},
|
||||
{http.MethodGet, "/groups/1/model-allowlist-candidates", ""},
|
||||
{http.MethodGet, "/groups/1/composite-routes", ""},
|
||||
{http.MethodPost, "/groups/1/composite-routes", `{"public_model":"x","target_platform":"openai"}`},
|
||||
{http.MethodPost, "/groups/1/composite-routes/preview", `{"model":"x"}`},
|
||||
{http.MethodPut, "/groups/1/composite-routes/2", `{"public_model":"x","target_platform":"openai"}`},
|
||||
{http.MethodDelete, "/groups/1/composite-routes/2", ""},
|
||||
{http.MethodGet, "/groups/1/rate-multipliers", ""},
|
||||
{http.MethodPut, "/groups/1/rate-multipliers", `{"entries":[]}`},
|
||||
{http.MethodDelete, "/groups/1/rate-multipliers", ""},
|
||||
{http.MethodPut, "/groups/1/rpm-overrides", `{"entries":[]}`},
|
||||
{http.MethodDelete, "/groups/1/rpm-overrides", ""},
|
||||
{http.MethodGet, "/groups/1/stats", ""},
|
||||
{http.MethodGet, "/groups/1/api-keys", ""},
|
||||
{http.MethodGet, "/groups/live-capability", ""},
|
||||
{http.MethodGet, "/groups/usage-summary", ""},
|
||||
{http.MethodGet, "/groups/capacity-summary", ""},
|
||||
{http.MethodPut, "/groups/sort-order", `{"updates":[{"id":1,"sort_order":1}]}`},
|
||||
}
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.method+" "+tt.path, func(t *testing.T) {
|
||||
svc := newStubAdminService()
|
||||
r := newSimpleModeGroupRouter(svc)
|
||||
res := httptest.NewRecorder()
|
||||
req := httptest.NewRequest(tt.method, tt.path, bytes.NewBufferString(tt.body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
r.ServeHTTP(res, req)
|
||||
require.Equal(t, http.StatusForbidden, res.Code)
|
||||
require.Zero(t, svc.advancedGroupOperationCalls)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -228,7 +228,7 @@ func (h *RedeemHandler) CreateAndRedeem(c *gin.Context) {
|
||||
return nil, createErr
|
||||
}
|
||||
|
||||
redeemed, redeemErr := h.redeemService.Redeem(ctx, req.UserID, req.Code)
|
||||
redeemed, redeemErr := h.redeemService.RedeemForAdminFulfillment(ctx, req.UserID, req.Code)
|
||||
if redeemErr != nil {
|
||||
return nil, redeemErr
|
||||
}
|
||||
@@ -246,7 +246,7 @@ func (h *RedeemHandler) resolveCreateAndRedeemExisting(ctx context.Context, exis
|
||||
return nil, service.ErrRedeemCodeExpired
|
||||
}
|
||||
if existing.CanUse() {
|
||||
redeemed, err := h.redeemService.Redeem(ctx, userID, existing.Code)
|
||||
redeemed, err := h.redeemService.RedeemForAdminFulfillment(ctx, userID, existing.Code)
|
||||
if err == nil {
|
||||
return gin.H{"redeem_code": dto.RedeemCodeFromServiceAdmin(redeemed)}, nil
|
||||
}
|
||||
|
||||
@@ -29,19 +29,19 @@ func TestAuthHandlerGetCurrentUserReturnsProfileCompatibilityFields(t *testing.T
|
||||
AvatarURL: "https://cdn.example.com/linuxdo.png",
|
||||
AvatarSource: "remote_url",
|
||||
},
|
||||
identities: []service.UserAuthIdentityRecord{
|
||||
{
|
||||
ProviderType: "linuxdo",
|
||||
ProviderKey: "linuxdo",
|
||||
ProviderSubject: "linuxdo-subject-31",
|
||||
VerifiedAt: &verifiedAt,
|
||||
Metadata: map[string]any{
|
||||
"username": "linuxdo-handle",
|
||||
"avatar_url": "https://cdn.example.com/linuxdo.png",
|
||||
},
|
||||
identities: []service.UserAuthIdentityRecord{
|
||||
{
|
||||
ProviderType: "linuxdo",
|
||||
ProviderKey: "linuxdo",
|
||||
ProviderSubject: "linuxdo-subject-31",
|
||||
VerifiedAt: &verifiedAt,
|
||||
Metadata: map[string]any{
|
||||
"username": "linuxdo-handle",
|
||||
"avatar_url": "https://cdn.example.com/linuxdo.png",
|
||||
},
|
||||
},
|
||||
}
|
||||
},
|
||||
}
|
||||
|
||||
handler := &AuthHandler{
|
||||
userService: service.NewUserService(repo, nil, nil, nil),
|
||||
|
||||
@@ -81,15 +81,19 @@ type userSupportedModelPricing struct {
|
||||
|
||||
// userPricingIntervalDTO 定价区间白名单(去掉内部 ID、SortOrder 等前端不渲染的字段)。
|
||||
type userPricingIntervalDTO struct {
|
||||
MinTokens int `json:"min_tokens"`
|
||||
MaxTokens *int `json:"max_tokens"`
|
||||
TierLabel string `json:"tier_label,omitempty"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
MinTokens int `json:"min_tokens"`
|
||||
MaxTokens *int `json:"max_tokens"`
|
||||
TierLabel string `json:"tier_label,omitempty"`
|
||||
InputPrice *float64 `json:"input_price"`
|
||||
OutputPrice *float64 `json:"output_price"`
|
||||
CacheWritePrice *float64 `json:"cache_write_price"`
|
||||
CacheWrite1hPrice *float64 `json:"cache_write_1h_price"`
|
||||
CacheReadPrice *float64 `json:"cache_read_price"`
|
||||
InputMultiplier *float64 `json:"input_multiplier"`
|
||||
OutputMultiplier *float64 `json:"output_multiplier"`
|
||||
CacheWriteMultiplier *float64 `json:"cache_write_multiplier"`
|
||||
CacheReadMultiplier *float64 `json:"cache_read_multiplier"`
|
||||
PerRequestPrice *float64 `json:"per_request_price"`
|
||||
}
|
||||
|
||||
// userSupportedModel 用户可见的支持模型条目。
|
||||
@@ -295,15 +299,19 @@ func toUserPricingIntervals(src []service.PricingInterval) []userPricingInterval
|
||||
intervals := make([]userPricingIntervalDTO, 0, len(src))
|
||||
for _, iv := range src {
|
||||
intervals = append(intervals, userPricingIntervalDTO{
|
||||
MinTokens: iv.MinTokens,
|
||||
MaxTokens: iv.MaxTokens,
|
||||
TierLabel: iv.TierLabel,
|
||||
InputPrice: iv.InputPrice,
|
||||
OutputPrice: iv.OutputPrice,
|
||||
CacheWritePrice: iv.CacheWritePrice,
|
||||
CacheWrite1hPrice: iv.CacheWrite1hPrice,
|
||||
CacheReadPrice: iv.CacheReadPrice,
|
||||
PerRequestPrice: iv.PerRequestPrice,
|
||||
MinTokens: iv.MinTokens,
|
||||
MaxTokens: iv.MaxTokens,
|
||||
TierLabel: iv.TierLabel,
|
||||
InputPrice: iv.InputPrice,
|
||||
OutputPrice: iv.OutputPrice,
|
||||
CacheWritePrice: iv.CacheWritePrice,
|
||||
CacheWrite1hPrice: iv.CacheWrite1hPrice,
|
||||
CacheReadPrice: iv.CacheReadPrice,
|
||||
InputMultiplier: iv.InputMultiplier,
|
||||
OutputMultiplier: iv.OutputMultiplier,
|
||||
CacheWriteMultiplier: iv.CacheWriteMultiplier,
|
||||
CacheReadMultiplier: iv.CacheReadMultiplier,
|
||||
PerRequestPrice: iv.PerRequestPrice,
|
||||
})
|
||||
}
|
||||
return intervals
|
||||
|
||||
@@ -113,10 +113,18 @@ func TestUserAvailableChannel_FieldWhitelist(t *testing.T) {
|
||||
}
|
||||
|
||||
// pricing interval 白名单:不应暴露 id / sort_order。
|
||||
inputMultiplier := 2.0
|
||||
outputMultiplier := 1.5
|
||||
cacheWriteMultiplier := 2.0
|
||||
cacheReadMultiplier := 2.0
|
||||
pricing := toUserPricing(&service.ChannelModelPricing{
|
||||
BillingMode: service.BillingModeToken,
|
||||
Intervals: []service.PricingInterval{
|
||||
{ID: 7, MinTokens: 0, MaxTokens: nil, SortOrder: 3},
|
||||
{
|
||||
ID: 7, MinTokens: 0, MaxTokens: nil, SortOrder: 3,
|
||||
InputMultiplier: &inputMultiplier, OutputMultiplier: &outputMultiplier,
|
||||
CacheWriteMultiplier: &cacheWriteMultiplier, CacheReadMultiplier: &cacheReadMultiplier,
|
||||
},
|
||||
},
|
||||
})
|
||||
require.NotNil(t, pricing)
|
||||
@@ -129,6 +137,14 @@ func TestUserAvailableChannel_FieldWhitelist(t *testing.T) {
|
||||
_, exists := ivDecoded[key]
|
||||
require.Falsef(t, exists, "user pricing interval must not expose %q", key)
|
||||
}
|
||||
for key, want := range map[string]float64{
|
||||
"input_multiplier": inputMultiplier, "output_multiplier": outputMultiplier,
|
||||
"cache_write_multiplier": cacheWriteMultiplier, "cache_read_multiplier": cacheReadMultiplier,
|
||||
} {
|
||||
got, exists := ivDecoded[key]
|
||||
require.Truef(t, exists, "user pricing interval must expose %q", key)
|
||||
require.InDelta(t, want, got.(float64), 1e-12)
|
||||
}
|
||||
}
|
||||
|
||||
func TestBuildPlatformSections_GroupsByPlatform(t *testing.T) {
|
||||
|
||||
@@ -139,9 +139,28 @@ func (h *BatchImageHandler) Models(c *gin.Context) {
|
||||
batchImageError(c, err)
|
||||
return
|
||||
}
|
||||
// 分组级模型白名单:ListModels 已按分组账号枚举,这里再按白名单过滤一次。
|
||||
if apiKey, ok := middleware.GetAPIKeyFromContext(c); ok && apiKey != nil && apiKey.Group != nil {
|
||||
got.Data = filterBatchImageModelsByAllowlist(got.Data, apiKey.Group.ModelAllowlist)
|
||||
}
|
||||
c.JSON(http.StatusOK, got)
|
||||
}
|
||||
|
||||
// filterBatchImageModelsByAllowlist 按分组模型白名单过滤批量生图模型列表。
|
||||
// 白名单未开启时原样返回。
|
||||
func filterBatchImageModelsByAllowlist(models []service.BatchImagePublicModel, allowlist service.GroupModelAllowlist) []service.BatchImagePublicModel {
|
||||
if !allowlist.Enabled {
|
||||
return models
|
||||
}
|
||||
filtered := make([]service.BatchImagePublicModel, 0, len(models))
|
||||
for _, model := range models {
|
||||
if allowlist.Allows(model.ID) {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func (h *BatchImageHandler) Items(c *gin.Context) {
|
||||
owner, ok := batchImageOwnerFromContext(c)
|
||||
if !ok {
|
||||
|
||||
@@ -152,6 +152,23 @@ func TestOpenAIReasoningEffortPolicyForCompositeTarget(t *testing.T) {
|
||||
var overLimit *service.ReasoningEffortOverLimitError
|
||||
require.ErrorAs(t, err, &overLimit)
|
||||
|
||||
mappingDenyGroup := *group
|
||||
mappingDenyGroup.MaxReasoningEffort = ""
|
||||
mappingDenyGroup.ReasoningEffortMappings = []service.ReasoningEffortMapping{
|
||||
{From: "max", To: service.ReasoningEffortMappingDeny},
|
||||
}
|
||||
mappingDenyAPIKey := &service.APIKey{Group: &mappingDenyGroup}
|
||||
mappingDenyCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
mappingDenyCtx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
||||
mappingDenyCtx.Request = mappingDenyCtx.Request.WithContext(service.WithResolvedTargetPlatform(mappingDenyCtx.Request.Context(), service.PlatformOpenAI))
|
||||
_, changed, err = applyOpenAIReasoningEffortPolicyForRequest(mappingDenyCtx, mappingDenyAPIKey, body)
|
||||
require.Error(t, err)
|
||||
require.False(t, changed)
|
||||
var mappingDenied *service.ReasoningEffortMappingDeniedError
|
||||
require.ErrorAs(t, err, &mappingDenied)
|
||||
require.Equal(t, "max", mappingDenied.Requested)
|
||||
require.Contains(t, mappingDenied.Error(), "denied by this group's mapping policy")
|
||||
|
||||
grokCtx, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
grokCtx.Request = httptest.NewRequest("POST", "/v1/responses", nil)
|
||||
grokCtx.Request = grokCtx.Request.WithContext(service.WithResolvedTargetPlatform(grokCtx.Request.Context(), service.PlatformGrok))
|
||||
|
||||
@@ -159,7 +159,7 @@ func GroupFromServiceAdmin(g *service.Group) *AdminGroup {
|
||||
MCPXMLInject: g.MCPXMLInject,
|
||||
DefaultMappedModel: g.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: g.MessagesDispatchModelConfig,
|
||||
ModelsListConfig: g.ModelsListConfig,
|
||||
ModelAllowlist: g.ModelAllowlist,
|
||||
CodexModelsManifestConfig: g.CodexModelsManifestConfig,
|
||||
SupportedModelScopes: g.SupportedModelScopes,
|
||||
AccountCount: g.AccountCount,
|
||||
|
||||
@@ -189,7 +189,7 @@ type AdminGroup struct {
|
||||
// OpenAI Messages 调度配置(仅 openai 平台使用)
|
||||
DefaultMappedModel string `json:"default_mapped_model"`
|
||||
MessagesDispatchModelConfig domain.OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config"`
|
||||
ModelsListConfig domain.GroupModelsListConfig `json:"models_list_config"`
|
||||
ModelAllowlist service.GroupModelAllowlist `json:"model_allowlist"`
|
||||
// 固定账号获取 Codex Model Manifest 配置(仅 openai 平台使用)。
|
||||
CodexModelsManifestConfig domain.GroupCodexModelsManifestConfig `json:"codex_models_manifest_config"`
|
||||
|
||||
|
||||
@@ -1134,11 +1134,20 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
platform = forcedPlatform
|
||||
}
|
||||
|
||||
if platform == service.PlatformOpenAI && apiKey != nil && apiKey.Group != nil &&
|
||||
apiKey.Group.Platform == service.PlatformOpenAI && apiKey.Group.CodexModelsManifestConfig.Enabled {
|
||||
h.pinnedOpenAIModels(c, apiKey.Group)
|
||||
return
|
||||
}
|
||||
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(c.Request.Context(), groupID)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() {
|
||||
availableModels = filterModelsByCustomList(availableModels, defaultModelIDsForPlatform(service.PlatformComposite), apiKey.Group.ModelsListConfig.Models)
|
||||
writeCustomModelsList(c, service.PlatformComposite, availableModels)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
source = defaultModelIDsForPlatform(service.PlatformComposite)
|
||||
}
|
||||
writeAllowlistedModelsList(c, service.PlatformComposite, apiKey.Group.ModelAllowlist.FilterForListing(source))
|
||||
return
|
||||
}
|
||||
if len(availableModels) > 0 {
|
||||
@@ -1151,10 +1160,9 @@ func (h *GatewayHandler) Models(c *gin.Context) {
|
||||
|
||||
// Get available models from account configurations for the selected group platform.
|
||||
availableModels := h.gatewayService.GetAvailableModels(c.Request.Context(), groupID, platform)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.CustomModelsListEnabled() {
|
||||
fallbackModels := defaultModelIDsForPlatform(platform)
|
||||
availableModels = filterModelsByCustomList(customModelsListSource(platform, availableModels, fallbackModels), fallbackModels, apiKey.Group.ModelsListConfig.Models)
|
||||
writeCustomModelsList(c, platform, availableModels)
|
||||
if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
source := modelListingSource(platform, availableModels, defaultModelIDsForPlatform(platform))
|
||||
writeAllowlistedModelsList(c, platform, apiKey.Group.ModelAllowlist.FilterForListing(source))
|
||||
return
|
||||
}
|
||||
|
||||
@@ -1239,8 +1247,12 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
if platform == service.PlatformComposite {
|
||||
availableModels := h.compositeAvailableModels(ctx, groupID)
|
||||
fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite)
|
||||
if group.CustomModelsListEnabled() {
|
||||
return filterModelsByCustomList(availableModels, fallbackModels, group.ModelsListConfig.Models)
|
||||
if group.ModelAllowlistEnabled() {
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
source = fallbackModels
|
||||
}
|
||||
return group.ModelAllowlist.FilterForListing(source)
|
||||
}
|
||||
if len(availableModels) > 0 {
|
||||
return availableModels
|
||||
@@ -1250,12 +1262,8 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi
|
||||
|
||||
availableModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
|
||||
fallbackModels := defaultCodexModelIDsForPlatform(platform)
|
||||
if group.CustomModelsListEnabled() {
|
||||
return filterModelsByCustomList(
|
||||
customModelsListSource(platform, availableModels, fallbackModels),
|
||||
fallbackModels,
|
||||
group.ModelsListConfig.Models,
|
||||
)
|
||||
if group.ModelAllowlistEnabled() {
|
||||
return group.ModelAllowlist.FilterForListing(modelListingSource(platform, availableModels, fallbackModels))
|
||||
}
|
||||
if len(availableModels) > 0 {
|
||||
return availableModels
|
||||
@@ -1295,6 +1303,10 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *
|
||||
}
|
||||
|
||||
func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
if platform == service.PlatformOpenAI {
|
||||
writeOpenAIModelsList(c, modelIDs)
|
||||
return
|
||||
}
|
||||
if platform == service.PlatformGrok {
|
||||
writeGrokModelsList(c, modelIDs)
|
||||
return
|
||||
@@ -1314,7 +1326,7 @@ func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
})
|
||||
}
|
||||
|
||||
func writeCustomModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
func writeAllowlistedModelsList(c *gin.Context, platform string, modelIDs []string) {
|
||||
if platform == service.PlatformOpenAI {
|
||||
writeOpenAIModelsList(c, modelIDs)
|
||||
return
|
||||
@@ -1412,72 +1424,19 @@ func writeOpenAIModelsList(c *gin.Context, modelIDs []string) {
|
||||
})
|
||||
}
|
||||
|
||||
func customModelsListSource(platform string, availableModels, fallbackModels []string) []string {
|
||||
if platform == service.PlatformAnthropic && len(availableModels) > 0 {
|
||||
// modelListingSource 汇总模型列表过滤的候选来源:账号映射键(availableModels)
|
||||
// 与平台默认列表(fallbackModels)。账号映射为空时回落默认列表;Anthropic
|
||||
// 平台两者取并集,其余平台以账号映射键为准。
|
||||
func modelListingSource(platform string, availableModels, fallbackModels []string) []string {
|
||||
if len(availableModels) == 0 {
|
||||
return fallbackModels
|
||||
}
|
||||
if platform == service.PlatformAnthropic {
|
||||
return mergeModelIDs(availableModels, fallbackModels)
|
||||
}
|
||||
return availableModels
|
||||
}
|
||||
|
||||
func filterModelsByCustomList(availableModels, fallbackModels, selectedModels []string) []string {
|
||||
if len(selectedModels) == 0 {
|
||||
return availableModels
|
||||
}
|
||||
source := availableModels
|
||||
if len(source) == 0 {
|
||||
source = fallbackModels
|
||||
}
|
||||
if len(source) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
allowed := make([]string, 0, len(source))
|
||||
for _, model := range source {
|
||||
model = strings.TrimSpace(model)
|
||||
if model != "" {
|
||||
allowed = append(allowed, model)
|
||||
}
|
||||
}
|
||||
|
||||
seen := make(map[string]struct{}, len(selectedModels))
|
||||
filtered := make([]string, 0, len(selectedModels))
|
||||
for _, model := range selectedModels {
|
||||
model = strings.TrimSpace(model)
|
||||
if model == "" {
|
||||
continue
|
||||
}
|
||||
if !customModelsListAllowsModel(allowed, model) {
|
||||
continue
|
||||
}
|
||||
if _, ok := seen[model]; ok {
|
||||
continue
|
||||
}
|
||||
seen[model] = struct{}{}
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
func customModelsListAllowsModel(availablePatterns []string, model string) bool {
|
||||
for _, pattern := range availablePatterns {
|
||||
if pattern == model {
|
||||
return true
|
||||
}
|
||||
if strings.HasSuffix(pattern, "*") && strings.HasPrefix(model, strings.TrimSuffix(pattern, "*")) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
normalizedClaudeModel := claude.NormalizeModelID(strings.TrimSuffix(model, "-thinking"))
|
||||
if normalizedClaudeModel != model {
|
||||
for _, pattern := range availablePatterns {
|
||||
if pattern == normalizedClaudeModel {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
func defaultCodexModelIDsForPlatform(platform string) []string {
|
||||
switch platform {
|
||||
case service.PlatformDeepseek:
|
||||
@@ -1551,10 +1510,21 @@ func mergeModelIDs(primary, secondary []string) []string {
|
||||
|
||||
// AntigravityModels 返回 Antigravity 支持的全部模型
|
||||
// GET /antigravity/models
|
||||
// 分组级模型白名单开启时按白名单过滤。
|
||||
func (h *GatewayHandler) AntigravityModels(c *gin.Context) {
|
||||
models := antigravity.DefaultModels()
|
||||
if apiKey, ok := middleware2.GetAPIKeyFromContext(c); ok && apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
filtered := make([]antigravity.ClaudeModel, 0, len(models))
|
||||
for _, model := range models {
|
||||
if apiKey.Group.ModelAllowlist.Allows(model.ID) {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
models = filtered
|
||||
}
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"object": "list",
|
||||
"data": antigravity.DefaultModels(),
|
||||
"data": models,
|
||||
})
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,245 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// 通配条目在 /v1/models 中展开为候选来源中所有匹配项,保持来源顺序。
|
||||
func TestGatewayModels_ModelAllowlistWildcardExpandsAgainstSource(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
groupID := int64(31)
|
||||
h := newGatewayModelsHandlerForTest(
|
||||
&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
groupID: {
|
||||
{
|
||||
ID: 1,
|
||||
Platform: service.PlatformOpenAI,
|
||||
Credentials: map[string]any{
|
||||
"model_mapping": map[string]any{
|
||||
"gpt-5.4": "gpt-5.4",
|
||||
"gpt-5.5-codex": "gpt-5.5-codex",
|
||||
"gpt-5.5-mini": "gpt-5.5-mini",
|
||||
"other-foo": "other-foo",
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gpt-5.5-*", "gpt-5.4"},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
h.Models(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, []string{"gpt-5.5-codex", "gpt-5.5-mini", "gpt-5.4"}, modelIDsForTest(got.Data))
|
||||
}
|
||||
|
||||
// geminiAllowlistAccountRepoStub 在 gatewayModelsAccountRepoStub 之上补充
|
||||
// Gemini 兼容层用到的按平台过滤查询(分组内无任何账号,触发 fallback 列表)。
|
||||
type geminiAllowlistAccountRepoStub struct {
|
||||
gatewayModelsAccountRepoStub
|
||||
}
|
||||
|
||||
func (s *geminiAllowlistAccountRepoStub) ListSchedulableByGroupIDAndPlatforms(ctx context.Context, groupID int64, platforms []string) ([]service.Account, error) {
|
||||
allowed := make(map[string]struct{}, len(platforms))
|
||||
for _, platform := range platforms {
|
||||
allowed[platform] = struct{}{}
|
||||
}
|
||||
accounts := s.byGroup[groupID]
|
||||
filtered := make([]service.Account, 0, len(accounts))
|
||||
for _, account := range accounts {
|
||||
if _, ok := allowed[account.Platform]; ok {
|
||||
filtered = append(filtered, account)
|
||||
}
|
||||
}
|
||||
return filtered, nil
|
||||
}
|
||||
|
||||
// Gemini 原生 /v1beta/models:白名单开启时过滤 fallback 列表的 models[].name。
|
||||
func TestGeminiV1BetaListModels_FiltersFallbackByAllowlist(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{
|
||||
// 只有 antigravity 账号:Gemini 选号失败后回落静态模型列表。
|
||||
41: {{ID: 2, Platform: service.PlatformAntigravity, Status: service.StatusActive, Schedulable: true}},
|
||||
},
|
||||
}}
|
||||
h := &GatewayHandler{
|
||||
geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, nil, nil, nil, nil, nil),
|
||||
}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
||||
geminiGroupID := int64(41)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
GroupID: &geminiGroupID,
|
||||
Group: &service.Group{
|
||||
ID: geminiGroupID,
|
||||
Platform: service.PlatformGemini,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gemini-2.5-pro", "gemini-3-*"},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
h.GeminiV1BetaListModels(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var got struct {
|
||||
Models []struct {
|
||||
Name string `json:"name"`
|
||||
} `json:"models"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
names := make([]string, 0, len(got.Models))
|
||||
for _, model := range got.Models {
|
||||
names = append(names, model.Name)
|
||||
}
|
||||
// models/ 前缀的候选形式也应命中条目。
|
||||
require.Contains(t, names, "models/gemini-2.5-pro")
|
||||
require.Contains(t, names, "models/gemini-3-pro-preview")
|
||||
require.NotContains(t, names, "models/gemini-2.5-flash")
|
||||
require.NotContains(t, names, "models/gemini-2.0-flash")
|
||||
}
|
||||
|
||||
// Antigravity /antigravity/models:白名单开启时按条目过滤静态列表。
|
||||
func TestAntigravityModels_FiltersByAllowlist(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
h := &GatewayHandler{}
|
||||
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/antigravity/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{
|
||||
ID: 51,
|
||||
Platform: service.PlatformAntigravity,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gemini-2.5-flash"},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
h.AntigravityModels(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
|
||||
var got struct {
|
||||
Object string `json:"object"`
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, "list", got.Object)
|
||||
// -thinking 后缀的宽容规则:gemini-2.5-flash-thinking 也命中 gemini-2.5-flash 条目。
|
||||
require.Len(t, got.Data, 2)
|
||||
require.Equal(t, "gemini-2.5-flash", got.Data[0].ID)
|
||||
require.Equal(t, "gemini-2.5-flash-thinking", got.Data[1].ID)
|
||||
}
|
||||
|
||||
func TestFilterUpstreamGeminiModelsBody(t *testing.T) {
|
||||
allowlist := service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-2.5-pro"}}
|
||||
|
||||
t.Run("full hit returns original body with dropped=false", func(t *testing.T) {
|
||||
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-pro"}]}`)
|
||||
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
|
||||
require.True(t, ok)
|
||||
require.False(t, dropped, "full hit must keep the original response so headers pass through")
|
||||
require.Equal(t, string(body), string(filtered))
|
||||
})
|
||||
|
||||
t.Run("partial hit filters models and preserves envelope fields", func(t *testing.T) {
|
||||
body := []byte(`{"nextPageToken":"tok","models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"other":"kept"}`)
|
||||
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
|
||||
require.True(t, ok)
|
||||
require.True(t, dropped)
|
||||
require.JSONEq(t, `{"nextPageToken":"tok","other":"kept","models":[{"name":"models/gemini-2.5-pro"}]}`, string(filtered))
|
||||
})
|
||||
|
||||
t.Run("models prefix stripped candidate form matches entry", func(t *testing.T) {
|
||||
body := []byte(`{"models":[{"name":"gemini-2.5-pro"}]}`)
|
||||
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
|
||||
require.True(t, ok)
|
||||
require.False(t, dropped)
|
||||
require.Equal(t, string(body), string(filtered))
|
||||
})
|
||||
|
||||
t.Run("missing models field passes through", func(t *testing.T) {
|
||||
body := []byte(`{"error":"odd upstream"}`)
|
||||
_, dropped, ok := filterUpstreamGeminiModelsBody(body, allowlist)
|
||||
require.True(t, ok)
|
||||
require.False(t, dropped)
|
||||
})
|
||||
|
||||
t.Run("invalid JSON signals passthrough", func(t *testing.T) {
|
||||
_, dropped, ok := filterUpstreamGeminiModelsBody([]byte(`not-json`), allowlist)
|
||||
require.False(t, ok)
|
||||
require.False(t, dropped)
|
||||
})
|
||||
}
|
||||
|
||||
func TestFilterBatchImageModelsByAllowlist(t *testing.T) {
|
||||
models := []service.BatchImagePublicModel{
|
||||
{ID: "gemini-2.5-flash-image", Object: "image.batch.model", Provider: "gemini_api"},
|
||||
{ID: "gemini-3-pro-image", Object: "image.batch.model", Provider: "vertex"},
|
||||
{ID: "gpt-image-1", Object: "image.batch.model", Provider: "openai"},
|
||||
}
|
||||
|
||||
t.Run("disabled allowlist keeps everything", func(t *testing.T) {
|
||||
got := filterBatchImageModelsByAllowlist(models, service.GroupModelAllowlist{Enabled: false})
|
||||
require.Equal(t, models, got)
|
||||
})
|
||||
|
||||
t.Run("exact and wildcard entries filter by model id", func(t *testing.T) {
|
||||
got := filterBatchImageModelsByAllowlist(models, service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gemini-3-*", "gpt-image-1"},
|
||||
})
|
||||
require.Equal(t, []string{"gemini-3-pro-image", "gpt-image-1"}, []string{got[0].ID, got[1].ID})
|
||||
require.Len(t, got, 2)
|
||||
})
|
||||
|
||||
t.Run("no match yields empty list", func(t *testing.T) {
|
||||
got := filterBatchImageModelsByAllowlist(models, service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"claude-*"},
|
||||
})
|
||||
require.Empty(t, got)
|
||||
})
|
||||
}
|
||||
@@ -8,6 +8,7 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/openai"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -181,7 +182,7 @@ func TestGatewayCodexModels_NonOpenAIGroupsUseMappedModels(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
// Scenario: Composite manifests aggregate only administrator-configured models.
|
||||
// Composite manifests include defaults from unmapped accounts and explicit mappings.
|
||||
func TestGatewayCodexModels_CompositeUsesCompleteEffectiveModelList(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 120
|
||||
@@ -225,7 +226,105 @@ func TestGatewayCodexModels_CompositeUsesCompleteEffectiveModelList(t *testing.T
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got codexModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, []string{"gpt-5.5", "grok-4.6"}, codexModelSlugsForTest(got.Models))
|
||||
want := service.FilterCodexModelIDsForGroup(openai.DefaultModelIDs(), nil)
|
||||
require.ElementsMatch(t, append(want, "grok-4.6"), codexModelSlugsForTest(got.Models))
|
||||
}
|
||||
|
||||
func TestGatewayModels_UnmappedOpenAIAccountsSupplementMappedModels(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 28
|
||||
const sparkModel = "gpt-5.3-codex-spark"
|
||||
const alias = "team-coder"
|
||||
parentID := int64(1)
|
||||
accounts := []service.Account{
|
||||
{ID: parentID, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth},
|
||||
{
|
||||
ID: 2, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
|
||||
ParentAccountID: &parentID, QuotaDimension: "spark",
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{sparkModel: sparkModel}},
|
||||
},
|
||||
{
|
||||
ID: 3, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{alias: "gpt-5.6-sol"}},
|
||||
},
|
||||
}
|
||||
tests := []struct {
|
||||
name string
|
||||
accounts []service.Account
|
||||
config service.GroupModelAllowlist
|
||||
want []string
|
||||
}{
|
||||
{
|
||||
name: "unmapped parent and Spark shadow retain defaults and aliases",
|
||||
accounts: accounts,
|
||||
want: append(openai.DefaultModelIDs(), alias),
|
||||
},
|
||||
{
|
||||
name: "unmapped API key account also contributes defaults",
|
||||
accounts: append([]service.Account{{ID: 4, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey}}, accounts[1:]...),
|
||||
want: append(openai.DefaultModelIDs(), alias),
|
||||
},
|
||||
{
|
||||
name: "unmapped accounts alone retain default response shape",
|
||||
accounts: accounts[:1],
|
||||
want: openai.DefaultModelIDs(),
|
||||
},
|
||||
{
|
||||
name: "custom list can select defaults and aliases",
|
||||
accounts: accounts,
|
||||
config: service.GroupModelAllowlist{Enabled: true, Models: []string{alias, "gpt-5.6-sol", sparkModel, "unknown-model"}},
|
||||
want: []string{alias, "gpt-5.6-sol", sparkModel},
|
||||
},
|
||||
{
|
||||
name: "unavailable custom selection remains empty",
|
||||
accounts: accounts,
|
||||
config: service.GroupModelAllowlist{Enabled: true, Models: []string{"unknown-model"}},
|
||||
want: []string{},
|
||||
},
|
||||
{
|
||||
name: "mapped accounts alone do not gain defaults",
|
||||
accounts: accounts[1:],
|
||||
want: []string{sparkModel, alias},
|
||||
},
|
||||
{
|
||||
name: "unmapped accounts from another platform do not add defaults",
|
||||
accounts: append([]service.Account{{ID: 4, Platform: service.PlatformAnthropic}}, accounts[1:]...),
|
||||
want: []string{sparkModel, alias},
|
||||
},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{
|
||||
byGroup: map[int64][]service.Account{groupID: tt.accounts},
|
||||
})
|
||||
for range 2 {
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformOpenAI, ModelAllowlist: tt.config},
|
||||
})
|
||||
h.Models(c)
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, "list", got.Object)
|
||||
require.ElementsMatch(t, tt.want, modelIDsForTest(got.Data))
|
||||
for _, model := range got.Data {
|
||||
require.Equal(t, "model", model.Object, model.ID)
|
||||
require.Positive(t, model.Created, model.ID)
|
||||
require.Equal(t, "openai", model.OwnedBy, model.ID)
|
||||
require.Empty(t, model.CreatedAt, model.ID)
|
||||
}
|
||||
if tt.config.Enabled {
|
||||
require.Equal(t, tt.want, modelIDsForTest(got.Data))
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
require.Empty(t, accounts[0].GetModelMapping())
|
||||
require.True(t, accounts[0].IsModelSupported("gpt-future-model"))
|
||||
require.False(t, accounts[1].IsModelSupported("gpt-5.6-sol"))
|
||||
}
|
||||
|
||||
func TestGatewayCodexModels_GeneratedManifestUsesFinalBodyETag(t *testing.T) {
|
||||
@@ -267,7 +366,7 @@ func TestGatewayCodexModels_GeneratedManifestUsesFinalBodyETag(t *testing.T) {
|
||||
require.Equal(t, etag, second.Header().Get("ETag"))
|
||||
}
|
||||
|
||||
// Scenario: group models_list_config limits the generated Codex manifest.
|
||||
// Scenario: group model_allowlist limits the generated Codex manifest.
|
||||
func TestGatewayCodexModels_CustomModelsListFiltersCompositeManifest(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const groupID int64 = 121
|
||||
@@ -299,7 +398,7 @@ func TestGatewayCodexModels_CustomModelsListFiltersCompositeManifest(t *testing.
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformComposite,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"grok-4.6"},
|
||||
},
|
||||
@@ -572,7 +671,7 @@ func TestGatewayModels_CustomModelsListDisabledKeepsOriginalModels(t *testing.T)
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: false,
|
||||
Models: []string{"gpt-5.5"},
|
||||
},
|
||||
@@ -619,7 +718,7 @@ func TestGatewayModels_CustomModelsListFiltersAndOrdersMappedModels(t *testing.T
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gpt-5.5", "missing-model", "gpt-5.4"},
|
||||
},
|
||||
@@ -704,7 +803,7 @@ func TestGatewayModels_CompositeCustomModelsListFiltersAcrossConcretePlatforms(t
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformComposite,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gemini-2.5-flash", "missing-model", "ag-custom-model", "gpt-5.5", "kimi-custom", "glm-custom", "deepseek-custom"},
|
||||
},
|
||||
@@ -918,7 +1017,7 @@ func TestGatewayModels_CustomModelsListKeepsConcreteModelAllowedByWildcardMappin
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformAnthropic,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"claude-sonnet-4-6"},
|
||||
},
|
||||
@@ -969,7 +1068,7 @@ func TestGatewayModels_AnthropicCustomModelsListIncludesOAuthClaudeAndMappedDeep
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformAnthropic,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"claude-fable-5", "claude-opus-4-8", "deepseek-v4-pro"},
|
||||
},
|
||||
@@ -1020,7 +1119,7 @@ func TestGatewayModels_AnthropicCustomModelsListDisabledKeepsMappedModelList(t *
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformAnthropic,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: false,
|
||||
Models: []string{"claude-fable-5", "deepseek-v4-pro"},
|
||||
},
|
||||
@@ -1061,7 +1160,7 @@ func TestGatewayModels_AnthropicCustomModelsListIncludesOAuthClaudeWithoutMappin
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformAnthropic,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"claude-opus-4-6-thinking", "claude-sonnet-4-5"},
|
||||
},
|
||||
@@ -1106,7 +1205,7 @@ func TestGatewayModels_CustomModelsListCanReturnEmptyWhenSelectionsUnavailable(t
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gpt-5.5"},
|
||||
},
|
||||
@@ -1143,7 +1242,7 @@ func TestGatewayModels_CustomModelsListFiltersDefaultFallbackModels(t *testing.T
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gpt-5.5", "legacy-gpt-2024", "gpt-5.4"},
|
||||
},
|
||||
@@ -1180,7 +1279,7 @@ func TestGatewayModels_OpenAICustomModelsListKeepsOpenAIResponseShapeForDefaultF
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gpt-5.5", "gpt-5.4"},
|
||||
},
|
||||
|
||||
@@ -46,14 +46,34 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
// 强制 antigravity 模式:返回 antigravity 支持的模型列表
|
||||
if forcePlatform == service.PlatformAntigravity {
|
||||
c.JSON(http.StatusOK, antigravity.FallbackGeminiModelsList())
|
||||
return
|
||||
// 分组级模型白名单开启时过滤 models[].name(名字形如 models/xxx)。
|
||||
filterGeminiModels := func(models []gemini.Model) []gemini.Model {
|
||||
if apiKey.Group == nil || !apiKey.Group.ModelAllowlistEnabled() {
|
||||
return models
|
||||
}
|
||||
filtered := make([]gemini.Model, 0, len(models))
|
||||
for _, model := range models {
|
||||
if apiKey.Group.ModelAllowlist.Allows(model.Name) {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
return filtered
|
||||
}
|
||||
|
||||
if models, ok := customGeminiModelsList(apiKey.Group); ok {
|
||||
c.JSON(http.StatusOK, models)
|
||||
// 强制 antigravity 模式:返回 antigravity 支持的模型列表
|
||||
if forcePlatform == service.PlatformAntigravity {
|
||||
if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
agModels := antigravity.DefaultGeminiModels()
|
||||
filtered := make([]antigravity.GeminiModel, 0, len(agModels))
|
||||
for _, model := range agModels {
|
||||
if apiKey.Group.ModelAllowlist.Allows(model.Name) {
|
||||
filtered = append(filtered, model)
|
||||
}
|
||||
}
|
||||
c.JSON(http.StatusOK, antigravity.GeminiModelsListResponse{Models: filtered})
|
||||
return
|
||||
}
|
||||
c.JSON(http.StatusOK, antigravity.FallbackGeminiModelsList())
|
||||
return
|
||||
}
|
||||
|
||||
@@ -63,7 +83,7 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID)
|
||||
if hasAntigravity {
|
||||
// antigravity 账户使用静态模型列表
|
||||
c.JSON(http.StatusOK, gemini.FallbackModelsList())
|
||||
c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())})
|
||||
return
|
||||
}
|
||||
markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
|
||||
@@ -77,21 +97,64 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
if shouldFallbackGeminiModels(res) {
|
||||
c.JSON(http.StatusOK, gemini.FallbackModelsList())
|
||||
c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())})
|
||||
return
|
||||
}
|
||||
if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() {
|
||||
if filtered, dropped, ok := filterUpstreamGeminiModelsBody(res.Body, apiKey.Group.ModelAllowlist); ok && dropped {
|
||||
// 只在确有条目被过滤时替换响应体;全命中或解析失败时保持原始响应,
|
||||
// 统一经 writeUpstreamResponse 写出(保留全部上游响应头)。
|
||||
res.Body = filtered
|
||||
}
|
||||
}
|
||||
writeUpstreamResponse(c, res)
|
||||
}
|
||||
|
||||
func customGeminiModelsList(group *service.Group) (gemini.ModelsListResponse, bool) {
|
||||
if group == nil || !group.CustomModelsListEnabled() {
|
||||
return gemini.ModelsListResponse{}, false
|
||||
// filterUpstreamGeminiModelsBody 按白名单过滤上游 /v1beta/models 响应中的
|
||||
// models[].name,其余信封字段(如 nextPageToken)原样保留。
|
||||
// 返回值:filtered 为过滤后的响应体;dropped 表示是否有条目被移除(全命中时
|
||||
// 为 false,调用方应保持原始响应以完整透传上游头);ok=false 表示解析失败,
|
||||
// 调用方同样应透传原始响应。
|
||||
func filterUpstreamGeminiModelsBody(body []byte, allowlist service.GroupModelAllowlist) (filtered []byte, dropped bool, ok bool) {
|
||||
var envelope map[string]json.RawMessage
|
||||
if err := json.Unmarshal(body, &envelope); err != nil {
|
||||
return nil, false, false
|
||||
}
|
||||
models := make([]gemini.Model, 0, len(group.ModelsListConfig.Models))
|
||||
for _, modelID := range group.ModelsListConfig.Models {
|
||||
models = append(models, gemini.FallbackModel(modelID))
|
||||
rawModels, hasModels := envelope["models"]
|
||||
if !hasModels {
|
||||
return body, false, true
|
||||
}
|
||||
return gemini.ModelsListResponse{Models: models}, true
|
||||
type geminiModelName struct {
|
||||
Name string `json:"name"`
|
||||
}
|
||||
var models []json.RawMessage
|
||||
if err := json.Unmarshal(rawModels, &models); err != nil {
|
||||
return nil, false, false
|
||||
}
|
||||
kept := make([]json.RawMessage, 0, len(models))
|
||||
for _, raw := range models {
|
||||
var model geminiModelName
|
||||
if err := json.Unmarshal(raw, &model); err != nil {
|
||||
return nil, false, false
|
||||
}
|
||||
if allowlist.Allows(model.Name) {
|
||||
kept = append(kept, raw)
|
||||
}
|
||||
}
|
||||
if len(kept) == len(models) {
|
||||
// 全部命中时直接透传原始响应体。
|
||||
return body, false, true
|
||||
}
|
||||
mergedModels, err := json.Marshal(kept)
|
||||
if err != nil {
|
||||
return nil, false, false
|
||||
}
|
||||
envelope["models"] = mergedModels
|
||||
merged, err := json.Marshal(envelope)
|
||||
if err != nil {
|
||||
return nil, false, false
|
||||
}
|
||||
return merged, true, true
|
||||
}
|
||||
|
||||
// GeminiV1BetaGetModel proxies:
|
||||
|
||||
@@ -9,40 +9,21 @@ import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/gemini"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGeminiV1BetaListModels_CustomGroupListUsesNativeResponse(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil)
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{
|
||||
Platform: service.PlatformGemini,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
Enabled: true,
|
||||
Models: []string{"gemini-2.5-pro", "models/gemini-custom"},
|
||||
},
|
||||
},
|
||||
})
|
||||
|
||||
(&GatewayHandler{}).GeminiV1BetaListModels(c)
|
||||
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got gemini.ModelsListResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, []gemini.Model{
|
||||
gemini.FallbackModel("gemini-2.5-pro"),
|
||||
gemini.FallbackModel("models/gemini-custom"),
|
||||
}, got.Models)
|
||||
func TestGeminiV1BetaListModels_AllowlistFiltersNativeResponse(t *testing.T) {
|
||||
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"},{"name":"models/gemini-2.5-flash"}],"nextPageToken":"next"}`)
|
||||
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-2.5-pro"}})
|
||||
require.True(t, ok)
|
||||
require.True(t, dropped)
|
||||
require.JSONEq(t, `{"models":[{"name":"models/gemini-2.5-pro"}],"nextPageToken":"next"}`, string(filtered))
|
||||
}
|
||||
|
||||
func TestGeminiV1BetaListModels_ForcedAntigravityIgnoresCustomGroupList(t *testing.T) {
|
||||
func TestGeminiV1BetaListModels_ForcedAntigravityAppliesAllowlist(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
@@ -50,7 +31,7 @@ func TestGeminiV1BetaListModels_ForcedAntigravityIgnoresCustomGroupList(t *testi
|
||||
c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{
|
||||
Group: &service.Group{
|
||||
Platform: service.PlatformGemini,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gemini-custom"},
|
||||
},
|
||||
@@ -63,19 +44,15 @@ func TestGeminiV1BetaListModels_ForcedAntigravityIgnoresCustomGroupList(t *testi
|
||||
require.Equal(t, http.StatusOK, rec.Code)
|
||||
var got antigravity.GeminiModelsListResponse
|
||||
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
|
||||
require.Equal(t, antigravity.FallbackGeminiModelsList(), got)
|
||||
require.Empty(t, got.Models)
|
||||
}
|
||||
|
||||
func TestCustomGeminiModelsList_DisabledKeepsExistingFlow(t *testing.T) {
|
||||
group := &service.Group{
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
Enabled: false,
|
||||
Models: []string{"gemini-2.5-pro"},
|
||||
},
|
||||
}
|
||||
|
||||
_, ok := customGeminiModelsList(group)
|
||||
require.False(t, ok)
|
||||
func TestGeminiModelAllowlist_DisabledPreservesNativeResponse(t *testing.T) {
|
||||
body := []byte(`{"models":[{"name":"models/gemini-2.5-pro"}]}`)
|
||||
filtered, dropped, ok := filterUpstreamGeminiModelsBody(body, service.GroupModelAllowlist{Enabled: false, Models: []string{"other"}})
|
||||
require.True(t, ok)
|
||||
require.False(t, dropped)
|
||||
require.Equal(t, body, filtered)
|
||||
}
|
||||
|
||||
// TestGeminiV1BetaHandler_PlatformRoutingInvariant 文档化并验证 Handler 层的平台路由逻辑不变量
|
||||
|
||||
@@ -16,7 +16,8 @@ import (
|
||||
// Codex CLI and the Codex desktop app refresh their model picker from
|
||||
// GET {base_url}/models?client_version=... (custom provider mode) or
|
||||
// GET /backend-api/codex/models (chatgpt_base_url mode). Both routes land
|
||||
// here. Groups with explicit account model mappings are generated locally;
|
||||
// here. Pinned discovery takes precedence over local account model mappings;
|
||||
// when disabled, groups with explicit mappings are generated locally;
|
||||
// otherwise ChatGPT manifests are proxied verbatim and custom API key manifests
|
||||
// receive provider-compatibility normalization plus short-lived caching.
|
||||
func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
@@ -34,28 +35,10 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
}
|
||||
|
||||
ifNoneMatch := c.GetHeader("If-None-Match")
|
||||
configuredManifest, configured, err := h.gatewayService.BuildGroupConfiguredCodexModelsManifest(
|
||||
c.Request.Context(),
|
||||
apiKey.Group,
|
||||
ifNoneMatch,
|
||||
)
|
||||
if err != nil {
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to build Codex models manifest")
|
||||
return
|
||||
}
|
||||
if configured {
|
||||
writeCodexModelsManifestResponse(c, configuredManifest)
|
||||
return
|
||||
}
|
||||
|
||||
// 固定账号分支:开启后只用选定账号拉取 manifest,不经过调度器;
|
||||
// 全部不可用/全部失败时按 FallbackToScheduler 决定回退调度器或返回错误。
|
||||
if apiKey.Group.Platform == service.PlatformOpenAI &&
|
||||
apiKey.Group.CodexModelsManifestConfig.Enabled &&
|
||||
len(apiKey.Group.CodexModelsManifestConfig.AccountIDs) > 0 {
|
||||
apiKey.Group.CodexModelsManifestConfig.Enabled {
|
||||
pinnedManifest, pinnedAccount, pinnedErr := h.gatewayService.FetchPinnedCodexModelsManifest(
|
||||
c.Request.Context(),
|
||||
apiKey.Group,
|
||||
@@ -84,7 +67,26 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
}
|
||||
writeCodexModelsManifestResponse(c, pinnedManifest)
|
||||
writeOpenAIModelsResponse(c, pinnedManifest)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if !apiKey.Group.CodexModelsManifestConfig.Enabled {
|
||||
configuredManifest, configured, err := h.gatewayService.BuildGroupConfiguredCodexModelsManifest(
|
||||
c.Request.Context(),
|
||||
apiKey.Group,
|
||||
ifNoneMatch,
|
||||
)
|
||||
if err != nil {
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
}
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to build Codex models manifest")
|
||||
return
|
||||
}
|
||||
if configured {
|
||||
writeOpenAIModelsResponse(c, configuredManifest)
|
||||
return
|
||||
}
|
||||
}
|
||||
@@ -133,6 +135,10 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to complete Codex models manifest")
|
||||
return
|
||||
}
|
||||
if err := service.ApplyPinnedCodexModelsMapping(manifest, account, apiKey.Group); err != nil {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to apply model mappings")
|
||||
return
|
||||
}
|
||||
if err := h.gatewayService.MergeGroupConfiguredCodexModels(c.Request.Context(), apiKey.Group, manifest, ifNoneMatch); err != nil {
|
||||
h.errorResponse(c, http.StatusInternalServerError, "api_error", "Failed to build Codex models manifest")
|
||||
return
|
||||
@@ -141,19 +147,7 @@ func (h *OpenAIGatewayHandler) CodexModels(c *gin.Context) {
|
||||
return
|
||||
}
|
||||
|
||||
writeCodexModelsManifestResponse(c, manifest)
|
||||
writeOpenAIModelsResponse(c, manifest)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
func writeCodexModelsManifestResponse(c *gin.Context, manifest *service.CodexModelsManifest) {
|
||||
if manifest.ETag != "" {
|
||||
c.Header("ETag", manifest.ETag)
|
||||
}
|
||||
if manifest.NotModified {
|
||||
c.Status(http.StatusNotModified)
|
||||
c.Writer.WriteHeaderNow()
|
||||
return
|
||||
}
|
||||
c.Data(http.StatusOK, "application/json", manifest.Body)
|
||||
}
|
||||
|
||||
@@ -163,7 +163,7 @@ func TestCodexModelsAppliesLocalFiltersBeforeClientETag(t *testing.T) {
|
||||
group := &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"codex-auto-review", "gpt-5.6"},
|
||||
},
|
||||
@@ -181,7 +181,7 @@ func TestCodexModelsAppliesLocalFiltersBeforeClientETag(t *testing.T) {
|
||||
t.Fatal("first response did not include an ETag")
|
||||
}
|
||||
|
||||
group.ModelsListConfig.Enabled = false
|
||||
group.ModelAllowlist.Enabled = false
|
||||
second := performCodexModelsRequestForGroup(t, handler, group, oldETag)
|
||||
if second.Code != http.StatusOK {
|
||||
t.Fatalf("second status: got %d, want %d; body=%s", second.Code, http.StatusOK, second.Body.String())
|
||||
@@ -232,7 +232,7 @@ func TestCodexModelsAPIKeyCacheDoesNotLeakGroupFilters(t *testing.T) {
|
||||
groupA := &service.Group{
|
||||
ID: 91,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"model-a"},
|
||||
},
|
||||
@@ -240,7 +240,7 @@ func TestCodexModelsAPIKeyCacheDoesNotLeakGroupFilters(t *testing.T) {
|
||||
groupB := &service.Group{
|
||||
ID: 92,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"model-b"},
|
||||
},
|
||||
@@ -292,8 +292,7 @@ func TestCodexModelsAPIKeyCacheDoesNotLeakGroupFilters(t *testing.T) {
|
||||
require.True(t, sawGroupB)
|
||||
}
|
||||
|
||||
// Scenario: OpenAI 分组内混用 OAuth 和第三方 API Key 时,管理员模型配置优先。
|
||||
func TestCodexModelsUsesConfiguredModelsBeforeUpstreamDiscovery(t *testing.T) {
|
||||
func TestCodexModelsSupplementsConfiguredModelsWithUnmappedAccountDefaults(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
groupID := int64(44)
|
||||
repo := &codexModelsFailoverAccountRepo{accounts: []service.Account{
|
||||
@@ -354,12 +353,59 @@ func TestCodexModelsUsesConfiguredModelsBeforeUpstreamDiscovery(t *testing.T) {
|
||||
if err := json.Unmarshal(recorder.Body.Bytes(), &envelope); err != nil {
|
||||
t.Fatalf("decode body: %v; body=%s", err, recorder.Body.String())
|
||||
}
|
||||
if len(envelope.Models) != 1 || envelope.Models[0]["slug"] != "glm-5.3" {
|
||||
t.Fatalf("models: got %v, want only glm-5.3", envelope.Models)
|
||||
require.Contains(t, codexHandlerManifestSlugs(t, recorder), "gpt-5.6-sol")
|
||||
require.Contains(t, codexHandlerManifestSlugs(t, recorder), "glm-5.3")
|
||||
for _, model := range envelope.Models {
|
||||
require.Contains(t, model, "supported_reasoning_levels")
|
||||
}
|
||||
if _, ok := envelope.Models[0]["supported_reasoning_levels"]; !ok {
|
||||
t.Fatalf("configured model is missing the Codex descriptor contract: %v", envelope.Models[0])
|
||||
}
|
||||
|
||||
func TestCodexModelsUnmappedParentAndSparkShadowHonorCustomListAndETag(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
const sparkModel = "gpt-5.3-codex-spark"
|
||||
parentID := int64(1)
|
||||
repo := &codexModelsFailoverAccountRepo{accounts: []service.Account{
|
||||
{
|
||||
ID: parentID, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive, Schedulable: true,
|
||||
},
|
||||
{
|
||||
ID: 2, Platform: service.PlatformOpenAI, Type: service.AccountTypeOAuth,
|
||||
Status: service.StatusActive, Schedulable: true,
|
||||
ParentAccountID: &parentID, QuotaDimension: "spark",
|
||||
Credentials: map[string]any{"model_mapping": map[string]any{sparkModel: sparkModel}},
|
||||
},
|
||||
}}
|
||||
upstream := &codexModelsFailoverHTTPUpstream{firstStatus: http.StatusNotFound}
|
||||
gatewayService := service.NewOpenAIGatewayService(
|
||||
repo,
|
||||
nil, nil, nil, nil, nil, nil, &config.Config{RunMode: config.RunModeSimple}, nil, nil, nil, nil, nil,
|
||||
upstream,
|
||||
nil, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
handler := &OpenAIGatewayHandler{gatewayService: gatewayService}
|
||||
group := &service.Group{ID: 45, Platform: service.PlatformOpenAI}
|
||||
first := performCodexModelsRequestForGroup(t, handler, group, "")
|
||||
require.Equal(t, http.StatusOK, first.Code, first.Body.String())
|
||||
slugs := codexHandlerManifestSlugs(t, first)
|
||||
require.Contains(t, slugs, "gpt-5.6-sol")
|
||||
require.Contains(t, slugs, sparkModel)
|
||||
require.NotContains(t, slugs, "gpt-image-2")
|
||||
require.NotContains(t, slugs, "codex-auto-review")
|
||||
firstETag := first.Header().Get("ETag")
|
||||
require.NotEmpty(t, firstETag)
|
||||
|
||||
group.ModelAllowlist = service.GroupModelAllowlist{
|
||||
Enabled: true, Models: []string{"gpt-5.6-sol", sparkModel, "unknown-model"},
|
||||
}
|
||||
second := performCodexModelsRequestForGroup(t, handler, group, firstETag)
|
||||
require.Equal(t, http.StatusOK, second.Code, second.Body.String())
|
||||
require.ElementsMatch(t, []string{"gpt-5.6-sol", sparkModel}, codexHandlerManifestSlugs(t, second))
|
||||
require.NotEqual(t, firstETag, second.Header().Get("ETag"))
|
||||
third := performCodexModelsRequestForGroup(t, handler, group, second.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusNotModified, third.Code)
|
||||
require.Empty(t, third.Body.Bytes())
|
||||
require.Empty(t, upstream.calls())
|
||||
}
|
||||
|
||||
func TestCompositeCodexModelsReusesExistingManifestSelection(t *testing.T) {
|
||||
@@ -912,7 +958,7 @@ func TestCodexModelsPinnedAccountsStillApplyCustomModelsListFilter(t *testing.T)
|
||||
group := &service.Group{
|
||||
ID: 84,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelsListConfig: service.GroupModelsListConfig{
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"model-b"},
|
||||
},
|
||||
|
||||
@@ -128,6 +128,44 @@ func TestOpenAICapacityFailoverExhaustionPreservesMessageAsServerError(t *testin
|
||||
})
|
||||
}
|
||||
|
||||
func TestOpenAIManagedSingleAccountModelNotFoundExhaustionPreservesStructured400(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
body := []byte(`{"error":{"type":"invalid_request_error","code":"model_not_found","param":"model","message":"The requested model is unavailable on this account"}}`)
|
||||
|
||||
(&OpenAIGatewayHandler{}).handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
ResponseBody: body,
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
require.Equal(t, "invalid_request_error", gjson.Get(recorder.Body.String(), "error.type").String())
|
||||
require.Equal(t, "model_not_found", gjson.Get(recorder.Body.String(), "error.code").String())
|
||||
require.Equal(t, "model", gjson.Get(recorder.Body.String(), "error.param").String())
|
||||
require.Equal(t, "The requested model is unavailable on this account", gjson.Get(recorder.Body.String(), "error.message").String())
|
||||
}
|
||||
|
||||
func TestOpenAIManagedModelNotFoundExhaustionSanitizesMessage(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
body := []byte(`{"error":{"code":"model_not_found","message":"Model not found; inspect https://upstream.example/debug?access_token=super-secret-value&model=x"}}`)
|
||||
|
||||
(&OpenAIGatewayHandler{}).handleFailoverExhausted(c, &service.UpstreamFailoverError{
|
||||
StatusCode: http.StatusBadRequest,
|
||||
ResponseBody: body,
|
||||
}, false)
|
||||
|
||||
require.Equal(t, http.StatusBadRequest, recorder.Code)
|
||||
require.NotContains(t, recorder.Body.String(), "super-secret-value")
|
||||
require.Contains(t, gjson.Get(recorder.Body.String(), "error.message").String(), "access_token=***")
|
||||
recorded, ok := c.Get(service.OpsUpstreamErrorMessageKey)
|
||||
require.True(t, ok)
|
||||
require.NotContains(t, recorded, "super-secret-value")
|
||||
require.Contains(t, recorded, "access_token=***")
|
||||
}
|
||||
|
||||
func TestResponsesFailoverExhaustedAfterForwardedTerminalMarksOpsWithoutDuplicateFrame(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
recorder := httptest.NewRecorder()
|
||||
|
||||
@@ -20,6 +20,7 @@ import (
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/ip"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
|
||||
"github.com/Wei-Shaw/sub2api/internal/securityaudit"
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -2025,10 +2026,10 @@ const (
|
||||
// 由 BeforeTurn 在每个 turn 开始时冻结,AfterTurn 的用量提交读取它;turn 在
|
||||
// 连接内串行推进,互斥锁只为跨用量提交 goroutine 的读取安全。
|
||||
//
|
||||
// ws_v2 passthrough ingress 没有 BeforeTurn,因此本值会保持零;AfterTurn 必须
|
||||
// 以 TurnStarted 已记录的所属 turn 开始时刻为回退,而不是用建连或记录时刻。
|
||||
// 这样每个 passthrough turn 都按自己的开始时刻计价,但不改变其仅在建连时执行
|
||||
// 准入门、没有 turn 级利润复核的既有行为。
|
||||
// 零值语义(重要):首轮准入由握手路径完成,不调用 BeforeTurn,因此首轮保持
|
||||
// 零值并回退到 TurnStarted 记录的首轮开始时刻。后续 turn 在 response.create
|
||||
// 写入上游前调用 BeforeTurn,按当时的利润门复核并冻结定价。绝不能用建连时刻
|
||||
// 初始化,否则会把长连接的所有 turn 钉死在建连时的峰谷因子。
|
||||
type openAIWSTurnPricing struct {
|
||||
mu sync.Mutex
|
||||
at time.Time
|
||||
@@ -2348,6 +2349,16 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, "model is required in first response.create payload")
|
||||
return
|
||||
}
|
||||
// 分组级模型白名单:首帧校验客户端模型,不通过则关闭连接并标记运维原因。
|
||||
// 必须在 ensureCompositeTargetPlatform(合成路由改写)之前执行。
|
||||
// 与 HTTP 准入一致:帧内重复 model 键/大小写变体可能被上游按末值绑定,
|
||||
// 全部候选值逐一校验,任一未命中即拒绝。
|
||||
if blocked := blockedModelAllowlistCandidate(apiKey.Group, requestmodel.FromBodyCandidates("", "application/json", firstMessage)); blocked != "" {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
|
||||
middleware2.MarkIngressRejected(c, middleware2.IngressRejectModelNotAllowed)
|
||||
closeOpenAIClientWS(wsConn, coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked))
|
||||
return
|
||||
}
|
||||
ensureCompositeTargetPlatform(c, apiKey, reqModel)
|
||||
ctx = c.Request.Context()
|
||||
if apiKey.Group != nil && apiKey.Group.Platform == service.PlatformComposite {
|
||||
@@ -2728,8 +2739,8 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
// turn-tagged slot preserves the exact mapping used for the in-flight request.
|
||||
var turnChannelMapping atomic.Pointer[openAIWSTurnChannelMappingSnapshot]
|
||||
turnChannelMapping.Store(&openAIWSTurnChannelMappingSnapshot{turn: 1, mapping: channelMappingWS})
|
||||
// turn 级定价:BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号;
|
||||
// passthrough 没有 BeforeTurn 时,AfterTurn 回退到 TurnStarted 的所属 turn 时刻。
|
||||
// turn 级定价:首轮回退到 TurnStarted 的所属 turn 时刻;后续 turn 由
|
||||
// BeforeTurn 重新冻结 pricingAt 并按最新门复核当前账号。
|
||||
var turnPricing openAIWSTurnPricing
|
||||
hooks := &service.OpenAIWSIngressHooks{
|
||||
ClientLifecycleContext: clientLifecycleCtx,
|
||||
@@ -2743,10 +2754,9 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
c.Set(securityAuditWSTurnContextKey, turn)
|
||||
service.BeginOpsStreamTurn(c, turn)
|
||||
setCyberTurnBody(turn, payload)
|
||||
// Passthrough ingress intentionally skips BeforeTurn, so enforce only
|
||||
// the connection-level cyber session gate here as well. Native ingress
|
||||
// visits this hook first and gets the same side-effect-free close error;
|
||||
// its original BeforeTurn guard remains as defense in depth.
|
||||
// 连接级 cyber session gate 也在 BeforeRequest 先执行,使 native 与
|
||||
// passthrough ingress 都能在 BeforeTurn 及上游写入前无副作用地拒绝。
|
||||
// BeforeTurn 中保留同一检查作为防御式兜底。
|
||||
if cyberBlockedThisConn {
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, cyberSessionBlockedClientMsg, nil)
|
||||
}
|
||||
@@ -2763,6 +2773,17 @@ func (h *OpenAIGatewayHandler) ResponsesWebSocket(c *gin.Context) {
|
||||
if model == "" {
|
||||
model = reqModel
|
||||
}
|
||||
// 分组级模型白名单:后续 turn 同样校验客户端模型(省略 model 时
|
||||
// 沿用会话实际生效模型,含 session.update 轮换后的模型),不通过
|
||||
// 则关闭整条连接,与推理强度 deny 一致。实际生效模型始终参与校验;
|
||||
// 帧内重复 model 键/大小写变体/嵌套 session.model 额外逐一校验,
|
||||
// 防止候选集非空时掩盖被轮换掉的禁用模型。
|
||||
candidates := append([]string{model}, requestmodel.FromBodyCandidates("", "application/json", payload)...)
|
||||
if blocked := blockedModelAllowlistCandidate(apiKey.Group, candidates); blocked != "" {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
|
||||
middleware2.MarkIngressRejected(c, middleware2.IngressRejectModelNotAllowed)
|
||||
return service.NewOpenAIWSClientCloseError(coderws.StatusPolicyViolation, fmt.Sprintf("Model %q is not available for this group", blocked), nil)
|
||||
}
|
||||
if decision := h.checkSecurityAuditStage(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIResponses, model, payload, "subsequent_turn"); decision != nil && !decision.AllowNextStage {
|
||||
writeSecurityAuditWSError(ctx, wsConn, decision)
|
||||
return service.NewOpenAIWSClientCloseError(securityAuditWSCloseStatus(decision), securityAuditWSCloseReason(decision), nil)
|
||||
@@ -3311,6 +3332,12 @@ func (h *OpenAIGatewayHandler) handleFailoverExhausted(c *gin.Context, failoverE
|
||||
}
|
||||
statusCode := failoverErr.StatusCode
|
||||
responseBody := failoverErr.ResponseBody
|
||||
if statusCode == http.StatusBadRequest && service.IsOpenAICompatibleModelNotFound400(responseBody) && !streamStarted {
|
||||
upstreamMsg := service.SanitizeUpstreamErrorMessage(service.ExtractUpstreamErrorMessage(responseBody))
|
||||
service.SetOpsUpstreamError(c, statusCode, upstreamMsg, "")
|
||||
service.WriteOpenAIUpstreamClientError(c, statusCode, responseBody, upstreamMsg)
|
||||
return
|
||||
}
|
||||
if service.IsOpenAISilentRefusalErrorBody(responseBody) {
|
||||
service.SetOpsUpstreamError(c, statusCode, service.OpenAISilentRefusalClientMessage(), "")
|
||||
h.handleStreamingAwareError(c, http.StatusBadGateway, "upstream_error", service.OpenAISilentRefusalClientMessage(), streamStarted)
|
||||
@@ -3670,6 +3697,21 @@ func isOpenAIWSUpgradeRequest(r *http.Request) bool {
|
||||
return strings.Contains(strings.ToLower(strings.TrimSpace(r.Header.Get("Connection"))), "upgrade")
|
||||
}
|
||||
|
||||
// blockedModelAllowlistCandidate 对全部候选模型逐一校验分组白名单,返回第一个
|
||||
// 未命中的值(全部命中或白名单未开启返回空串)。WS 帧与 HTTP 请求体共用该
|
||||
// 规则:重复 model 键/大小写变体可能被上游按末值绑定,任一未命中即拒绝。
|
||||
func blockedModelAllowlistCandidate(group *service.Group, candidates []string) string {
|
||||
if group == nil || !group.ModelAllowlistEnabled() {
|
||||
return ""
|
||||
}
|
||||
for _, candidate := range candidates {
|
||||
if !group.ModelAllowlist.Allows(candidate) {
|
||||
return candidate
|
||||
}
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
func closeOpenAIClientWS(conn *coderws.Conn, status coderws.StatusCode, reason string) {
|
||||
if conn == nil {
|
||||
return
|
||||
|
||||
@@ -1916,7 +1916,10 @@ func newOpenAIWSHandlerTestServer(t *testing.T, h *OpenAIGatewayHandler, subject
|
||||
}
|
||||
|
||||
type openAIResponsesWSUsageLogCase struct {
|
||||
firstPayload string
|
||||
firstPayload string
|
||||
// midPayload 在首个 turn 完成后发送(如 session.update),上游桩会为它
|
||||
// 回一个 response.completed,客户端按普通事件读取。
|
||||
midPayload string
|
||||
secondPayload string
|
||||
userAgent *string
|
||||
ingressMode string
|
||||
@@ -1924,6 +1927,12 @@ type openAIResponsesWSUsageLogCase struct {
|
||||
billingModelSource string
|
||||
accountModelMapping map[string]any
|
||||
afterFirstUpstreamRequest func(channelSvc *service.ChannelService) error
|
||||
// group 覆盖 apiKey.Group(分组级模型白名单测试用);nil 保持原有无分组行为。
|
||||
group *service.Group
|
||||
// firstFrameCloseExpected:首帧即被拒(连接被 1008 关闭),不期待任何响应帧。
|
||||
firstFrameCloseExpected bool
|
||||
// secondTurnCloseExpected:第二个 turn 被拒(连接被 1008 关闭)。
|
||||
secondTurnCloseExpected bool
|
||||
}
|
||||
|
||||
type openAIResponsesWSUsageLogResult struct {
|
||||
@@ -2843,8 +2852,11 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
turnCount := 1
|
||||
if strings.TrimSpace(tc.midPayload) != "" {
|
||||
turnCount++
|
||||
}
|
||||
if strings.TrimSpace(tc.secondPayload) != "" {
|
||||
turnCount = 2
|
||||
turnCount++
|
||||
}
|
||||
upstreamPayloadCh := make(chan []byte, turnCount)
|
||||
upstreamErrCh := make(chan error, 1)
|
||||
@@ -2997,6 +3009,9 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
GroupID: &groupID,
|
||||
User: &service.User{ID: 1701, Status: service.StatusActive},
|
||||
}
|
||||
if tc.group != nil {
|
||||
apiKey.Group = tc.group
|
||||
}
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(middleware.ContextKeyAPIKey), apiKey)
|
||||
@@ -3028,6 +3043,19 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
|
||||
if tc.firstFrameCloseExpected {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, readErr := clientConn.Read(readCtx)
|
||||
cancelRead()
|
||||
require.Error(t, readErr, "first frame should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
require.Contains(t, closeErr.Reason, "not available for this group")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
|
||||
clientEvents := make([][]byte, 0, turnCount)
|
||||
readCompleted := func() {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
@@ -3038,11 +3066,30 @@ func runOpenAIResponsesWebSocketUsageLogCase(t *testing.T, tc openAIResponsesWSU
|
||||
clientEvents = append(clientEvents, append([]byte(nil), event...))
|
||||
}
|
||||
readCompleted()
|
||||
if turnCount == 2 {
|
||||
if tc.midPayload != "" {
|
||||
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.midPayload))
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
readCompleted()
|
||||
}
|
||||
if strings.TrimSpace(tc.secondPayload) != "" && (turnCount >= 2) {
|
||||
writeCtx, cancelWrite = context.WithTimeout(context.Background(), 3*time.Second)
|
||||
err = clientConn.Write(writeCtx, coderws.MessageText, []byte(tc.secondPayload))
|
||||
cancelWrite()
|
||||
require.NoError(t, err)
|
||||
if tc.secondTurnCloseExpected {
|
||||
readCtx, cancelRead := context.WithTimeout(context.Background(), 3*time.Second)
|
||||
_, _, readErr := clientConn.Read(readCtx)
|
||||
cancelRead()
|
||||
require.Error(t, readErr, "second turn should have been rejected with a close")
|
||||
var closeErr coderws.CloseError
|
||||
require.ErrorAs(t, readErr, &closeErr)
|
||||
require.Equal(t, coderws.StatusPolicyViolation, closeErr.Code)
|
||||
require.Contains(t, closeErr.Reason, "not available for this group")
|
||||
_ = clientConn.CloseNow()
|
||||
return openAIResponsesWSUsageLogResult{}
|
||||
}
|
||||
readCompleted()
|
||||
}
|
||||
_ = clientConn.Close(coderws.StatusNormalClosure, "done")
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
)
|
||||
|
||||
func wsAllowlistGroup(enabled bool, models ...string) *service.Group {
|
||||
return &service.Group{
|
||||
ID: 4201,
|
||||
Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: enabled,
|
||||
Models: models,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// 首帧拒绝:模型不在白名单时连接被 1008 关闭(passthrough relay 模式)。
|
||||
func TestOpenAIResponsesWebSocket_FirstFrameModelNotAllowedCloses_Passthrough(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: service.OpenAIWSIngressModePassthrough,
|
||||
firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
// 首帧拒绝:原生 ingress 模式。
|
||||
func TestOpenAIResponsesWebSocket_FirstFrameModelNotAllowedCloses_NativeIngress(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: service.OpenAIWSIngressModeDedicated,
|
||||
firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
// 首帧命中白名单:连接正常建立并收到 response.completed(passthrough)。
|
||||
func TestOpenAIResponsesWebSocket_FirstFrameAllowlistedModelProceeds(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
})
|
||||
if len(got.clientEvents) != 1 {
|
||||
t.Fatalf("expected one completed event, got %d", len(got.clientEvents))
|
||||
}
|
||||
}
|
||||
|
||||
// 后续 turn 切换到白名单外的模型:连接被 1008 关闭(passthrough relay 模式)。
|
||||
func TestOpenAIResponsesWebSocket_SubsequentTurnModelNotAllowedCloses_Passthrough(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: service.OpenAIWSIngressModePassthrough,
|
||||
secondTurnCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
// 后续 turn 切换到白名单外的模型:连接被 1008 关闭(原生 ingress 模式)。
|
||||
func TestOpenAIResponsesWebSocket_SubsequentTurnModelNotAllowedCloses_NativeIngress(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: service.OpenAIWSIngressModeDedicated,
|
||||
secondTurnCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
// 后续 turn 省略 model:沿用会话初始模型,白名单校验通过。
|
||||
func TestOpenAIResponsesWebSocket_SubsequentTurnOmittedModelUsesSessionModel(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
})
|
||||
if len(got.clientEvents) != 2 {
|
||||
t.Fatalf("expected two completed events, got %d", len(got.clientEvents))
|
||||
}
|
||||
if len(got.upstreamPayloads) != 2 {
|
||||
t.Fatalf("expected two upstream frames, got %d", len(got.upstreamPayloads))
|
||||
}
|
||||
}
|
||||
|
||||
// 白名单关闭:不受任何影响。
|
||||
func TestOpenAIResponsesWebSocket_DisabledAllowlistDoesNotInterfere(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-4.1","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
group: wsAllowlistGroup(false, "gpt-5.4"),
|
||||
})
|
||||
if len(got.clientEvents) != 2 {
|
||||
t.Fatalf("expected two completed events, got %d", len(got.clientEvents))
|
||||
}
|
||||
}
|
||||
|
||||
// 首帧含重复 model 键(首个在白名单、末个不在):passthrough 会原样转发两个字段,
|
||||
// 上游按末值绑定时得到禁用模型,必须在入口拒绝。
|
||||
func TestOpenAIResponsesWebSocket_FirstFrameDuplicateModelKeysRejected(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: mode,
|
||||
firstFrameCloseExpected: true,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 首帧混用大小写变体键({"model":"allowed","Model":"blocked"}):绑定系上游
|
||||
// 大小写不敏感且末值生效会绑定 blocked,必须在入口拒绝。仅含变体键的帧则会在
|
||||
// 「model is required」处更早被关闭。
|
||||
func TestOpenAIResponsesWebSocket_FirstFrameCaseVariantModelKeyRejected(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","Model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: service.OpenAIWSIngressModePassthrough,
|
||||
firstFrameCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
// 后续 turn 含重复 model 键(首个在白名单、末个不在):拒绝并关闭连接。
|
||||
func TestOpenAIResponsesWebSocket_SubsequentTurnDuplicateModelKeysRejected(t *testing.T) {
|
||||
for _, mode := range []string{service.OpenAIWSIngressModePassthrough, service.OpenAIWSIngressModeDedicated} {
|
||||
t.Run(mode, func(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.4","model":"gpt-4.1","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: mode,
|
||||
secondTurnCloseExpected: true,
|
||||
})
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// 重复但同值的 model 键不误伤,连接正常完成两个 turn。
|
||||
func TestOpenAIResponsesWebSocket_DuplicateIdenticalModelKeysAllowed(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","model":"gpt-5.4","stream":false}`,
|
||||
secondPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
})
|
||||
if len(got.clientEvents) != 2 {
|
||||
t.Fatalf("expected two completed events, got %d", len(got.clientEvents))
|
||||
}
|
||||
}
|
||||
|
||||
// session.update 轮换绕过:首帧用白名单内模型建立会话,session.update 把会话
|
||||
// 模型改为白名单外的模型,随后 response.create 携带嵌套 session.model(白名单
|
||||
// 内)让帧内候选非空。实际生效模型(轮换后的会话模型)必须始终参与校验。
|
||||
func TestOpenAIResponsesWebSocket_SessionUpdateRotationBypassRejected(t *testing.T) {
|
||||
runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
midPayload: `{"type":"session.update","session":{"model":"gpt-4.1"}}`,
|
||||
secondPayload: `{"type":"response.create","session":{"model":"gpt-5.4"},"stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
ingressMode: service.OpenAIWSIngressModePassthrough,
|
||||
secondTurnCloseExpected: true,
|
||||
})
|
||||
}
|
||||
|
||||
// 轮换后的会话模型本身在白名单内时,省略 model 的后续 turn 正常放行。
|
||||
func TestOpenAIResponsesWebSocket_SessionUpdateToAllowedModelStillWorks(t *testing.T) {
|
||||
got := runOpenAIResponsesWebSocketUsageLogCase(t, openAIResponsesWSUsageLogCase{
|
||||
firstPayload: `{"type":"response.create","model":"gpt-5.4","stream":false}`,
|
||||
midPayload: `{"type":"session.update","session":{"model":"gpt-5.4"}}`,
|
||||
secondPayload: `{"type":"response.create","stream":false}`,
|
||||
group: wsAllowlistGroup(true, "gpt-5.4"),
|
||||
})
|
||||
if len(got.clientEvents) < 2 {
|
||||
t.Fatalf("expected at least two completed events, got %d", len(got.clientEvents))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,52 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"net/http"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func (h *GatewayHandler) pinnedOpenAIModels(c *gin.Context, group *service.Group) {
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
}
|
||||
if h.openAIGatewayService == nil {
|
||||
writeOpenAIModelsError(c, http.StatusInternalServerError, "api_error", "OpenAI model discovery is not configured")
|
||||
return
|
||||
}
|
||||
response, account, err := h.openAIGatewayService.FetchPinnedOpenAIModelsList(
|
||||
c.Request.Context(), group, h.maxAccountSwitches, c.GetHeader("If-None-Match"),
|
||||
)
|
||||
if c.Request.Context().Err() != nil {
|
||||
return
|
||||
}
|
||||
if err != nil {
|
||||
if errors.Is(err, service.ErrNoPinnedCodexModelsAccounts) {
|
||||
writeOpenAIModelsError(c, http.StatusServiceUnavailable, "upstream_error", "No available OpenAI model discovery accounts")
|
||||
return
|
||||
}
|
||||
writeOpenAIModelsError(c, infraerrors.Code(err), "upstream_error", infraerrors.Message(err))
|
||||
return
|
||||
}
|
||||
setOpsSelectedAccount(c, account.ID, account.Platform)
|
||||
writeOpenAIModelsResponse(c, response)
|
||||
}
|
||||
|
||||
func writeOpenAIModelsError(c *gin.Context, status int, errorType, message string) {
|
||||
c.JSON(status, gin.H{"error": gin.H{"type": errorType, "message": message}})
|
||||
}
|
||||
|
||||
func writeOpenAIModelsResponse(c *gin.Context, manifest *service.OpenAIModelsResponse) {
|
||||
if manifest.ETag != "" {
|
||||
c.Header("ETag", manifest.ETag)
|
||||
}
|
||||
if manifest.NotModified {
|
||||
c.Status(http.StatusNotModified)
|
||||
c.Writer.WriteHeaderNow()
|
||||
return
|
||||
}
|
||||
c.Data(http.StatusOK, "application/json", manifest.Body)
|
||||
}
|
||||
@@ -0,0 +1,217 @@
|
||||
package handler
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func performOrdinaryPinnedModelsRequest(t *testing.T, codex *OpenAIGatewayHandler, group *service.Group, path, etag string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
h := &GatewayHandler{openAIGatewayService: codex.gatewayService, maxAccountSwitches: 3}
|
||||
recorder := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(recorder)
|
||||
c.Request = httptest.NewRequest(http.MethodGet, path, nil)
|
||||
c.Request.Header.Set("If-None-Match", etag)
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
|
||||
h.Models(c)
|
||||
return recorder
|
||||
}
|
||||
|
||||
func ordinaryPinnedModelIDs(t *testing.T, recorder *httptest.ResponseRecorder) []string {
|
||||
t.Helper()
|
||||
var response gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(recorder.Body.Bytes(), &response))
|
||||
require.NotNil(t, response.Data, "an empty catalog must be an array, not null")
|
||||
return modelIDsForTest(response.Data)
|
||||
}
|
||||
|
||||
func TestOrdinaryPinnedModelsUsesSelectedAccountsAndFinalETag(t *testing.T) {
|
||||
accounts := []service.Account{
|
||||
newPinnedCodexAccount(1, service.StatusActive, true, false),
|
||||
newPinnedCodexAccount(2, service.StatusActive, true, true),
|
||||
newPinnedCodexAccount(3, service.StatusActive, true, false),
|
||||
}
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
|
||||
1: `{"data":[{"id":"unselected"}]}`,
|
||||
2: `{"data":[{"id":"special-model","owned_by":"first","created":123},{"id":"gpt-image-1"}]}`,
|
||||
3: `{"data":[{"id":"special-model","owned_by":"second"},{"id":"text-embedding-3-large"}]}`,
|
||||
}}
|
||||
h := newPinnedCodexTestHandler(accounts, upstream, 3)
|
||||
group := &service.Group{ID: 91, Platform: service.PlatformOpenAI,
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2, 3}}}
|
||||
first := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
|
||||
require.Equal(t, http.StatusOK, first.Code, first.Body.String())
|
||||
require.Equal(t, []string{"special-model", "gpt-image-1", "text-embedding-3-large"}, ordinaryPinnedModelIDs(t, first))
|
||||
var response gatewayModelsResponseForTest
|
||||
require.NoError(t, json.Unmarshal(first.Body.Bytes(), &response))
|
||||
require.Equal(t, "model", response.Data[0].Object)
|
||||
require.Equal(t, "first", response.Data[0].OwnedBy)
|
||||
require.EqualValues(t, 123, response.Data[0].Created)
|
||||
require.Equal(t, []int64{2, 3}, upstream.accountIDs())
|
||||
alias := performOrdinaryPinnedModelsRequest(t, h, group, "/models", first.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusNotModified, alias.Code)
|
||||
require.Empty(t, alias.Body.String())
|
||||
require.Equal(t, []int64{2, 3}, upstream.accountIDs(), "fresh requests must not hit upstream")
|
||||
|
||||
otherGroup := *group
|
||||
otherGroup.ID++
|
||||
otherGroup.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{"text-embedding-3-large", "special-model"}}
|
||||
filtered := performOrdinaryPinnedModelsRequest(t, h, &otherGroup, "/v1/models", first.Header().Get("ETag"))
|
||||
require.Equal(t, http.StatusOK, filtered.Code)
|
||||
require.Equal(t, []string{"text-embedding-3-large", "special-model"}, ordinaryPinnedModelIDs(t, filtered))
|
||||
require.NotEqual(t, first.Header().Get("ETag"), filtered.Header().Get("ETag"))
|
||||
again := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
|
||||
require.Equal(t, first.Body.String(), again.Body.String(), "group filters must not modify account cache")
|
||||
require.Equal(t, []int64{2, 3}, upstream.accountIDs())
|
||||
}
|
||||
|
||||
func TestOrdinaryPinnedModelsFailureAndEmptyPolicies(t *testing.T) {
|
||||
for _, tc := range []struct {
|
||||
name string
|
||||
accountIDs []int64
|
||||
bodies map[int64]string
|
||||
statuses map[int64]int
|
||||
fallback bool
|
||||
selected []string
|
||||
wantStatus int
|
||||
wantIDs []string
|
||||
}{
|
||||
{name: "partial failure", accountIDs: []int64{2, 3}, bodies: map[int64]string{2: `{"data":[{"id":"ok"}]}`}, statuses: map[int64]int{3: 503}, wantStatus: 200, wantIDs: []string{"ok"}},
|
||||
{name: "all failed", accountIDs: []int64{2, 3}, statuses: map[int64]int{2: 429, 3: 504}, wantStatus: 502},
|
||||
{name: "missing members", accountIDs: []int64{99}, wantStatus: 503},
|
||||
{name: "empty config", wantStatus: 503},
|
||||
{name: "fallback from missing members", accountIDs: []int64{99}, fallback: true, wantStatus: 200, wantIDs: []string{"from-scheduler"}},
|
||||
{name: "fallback from failure", accountIDs: []int64{2}, statuses: map[int64]int{2: 503}, fallback: true, wantStatus: 200, wantIDs: []string{"from-scheduler"}},
|
||||
{name: "valid empty", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"data":[]}`}, fallback: true, wantStatus: 200},
|
||||
{name: "filtered empty", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"data":[{"id":"ok"}]}`}, selected: []string{"absent"}, fallback: true, wantStatus: 200},
|
||||
{name: "invalid envelope", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"error":"denied"}`}, wantStatus: 502},
|
||||
{name: "null data", accountIDs: []int64{2}, bodies: map[int64]string{2: `{"data":null}`}, wantStatus: 502},
|
||||
} {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
accounts := []service.Account{
|
||||
newPinnedCodexAccount(1, service.StatusActive, true, false),
|
||||
newPinnedCodexAccount(2, service.StatusActive, true, false),
|
||||
newPinnedCodexAccount(3, service.StatusActive, true, false),
|
||||
}
|
||||
bodies := map[int64]string{1: `{"data":[{"id":"from-scheduler"}]}`}
|
||||
for id, body := range tc.bodies {
|
||||
bodies[id] = body
|
||||
}
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: bodies, statuses: tc.statuses}
|
||||
h := newPinnedCodexTestHandler(accounts, upstream, 3)
|
||||
group := &service.Group{ID: 92, Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{Enabled: len(tc.selected) > 0, Models: tc.selected},
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: tc.accountIDs, FallbackToScheduler: tc.fallback}}
|
||||
recorder := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
|
||||
require.Equal(t, tc.wantStatus, recorder.Code, recorder.Body.String())
|
||||
if tc.wantStatus == 200 {
|
||||
require.Equal(t, len(tc.wantIDs), len(ordinaryPinnedModelIDs(t, recorder)))
|
||||
if len(tc.wantIDs) > 0 {
|
||||
require.Equal(t, tc.wantIDs, ordinaryPinnedModelIDs(t, recorder))
|
||||
} else {
|
||||
require.NotContains(t, upstream.accountIDs(), int64(1), "empty success must not trigger fallback")
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestOrdinaryPinnedModelsSkipsPersistentlyUnavailableAccounts(t *testing.T) {
|
||||
accounts := []service.Account{
|
||||
newPinnedCodexAccount(1, service.StatusActive, true, true),
|
||||
newPinnedCodexAccount(2, service.StatusActive, false, false),
|
||||
newPinnedCodexAccount(3, service.StatusDisabled, true, false),
|
||||
newPinnedCodexAccount(4, service.StatusActive, true, false),
|
||||
}
|
||||
expired := time.Now().Add(-time.Hour)
|
||||
accounts[3].ExpiresAt = &expired
|
||||
accounts[3].AutoPauseOnExpired = true
|
||||
overloaded := time.Now().Add(time.Hour)
|
||||
accounts[0].OverloadUntil = &overloaded
|
||||
accounts[0].TempUnschedulableUntil = &overloaded
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{1: `{"data":[{"id":"available"}]}`}}
|
||||
h := newPinnedCodexTestHandler(accounts, upstream, 3)
|
||||
group := &service.Group{ID: 93, Platform: service.PlatformOpenAI,
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1, 2, 3, 4, 99}}}
|
||||
recorder := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{1}, upstream.accountIDs())
|
||||
}
|
||||
|
||||
func TestPinnedModelsAllowlistExpandsWildcardsForBothRepresentations(t *testing.T) {
|
||||
for _, codex := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "ordinary", true: "codex"}[codex], func(t *testing.T) {
|
||||
accounts := []service.Account{newPinnedCodexAccount(1, service.StatusActive, true, false)}
|
||||
accounts[0].Credentials["model_mapping"] = map[string]any{
|
||||
"public-a": "gpt-5.5", "public-b": "gpt-5.5", "blocked": "gpt-5.5",
|
||||
}
|
||||
body := `{"data":[{"id":"gpt-5.5","owned_by":"provider"}]}`
|
||||
if codex {
|
||||
body = `{"models":[{"slug":"gpt-5.5","context_window":424242}]}`
|
||||
}
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{1: body}}
|
||||
h := newPinnedCodexTestHandler(accounts, upstream, 3)
|
||||
group := &service.Group{ID: 95, Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{Enabled: true, Models: []string{"public-b", "public-*"}},
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{1}}}
|
||||
if codex {
|
||||
recorder := performPinnedCodexModelsRequest(t, h, group, "")
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []string{"public-b", "public-a"}, codexHandlerManifestSlugs(t, recorder))
|
||||
require.Contains(t, recorder.Body.String(), `"context_window":424242`)
|
||||
} else {
|
||||
recorder := performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []string{"public-b", "public-a"}, ordinaryPinnedModelIDs(t, recorder))
|
||||
require.Contains(t, recorder.Body.String(), `"owned_by":"provider"`)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestPinnedModelsMappingFollowsUpstreamDiscoveryForBothRepresentations(t *testing.T) {
|
||||
for _, codex := range []bool{false, true} {
|
||||
t.Run(map[bool]string{false: "ordinary", true: "codex"}[codex], func(t *testing.T) {
|
||||
accounts := []service.Account{
|
||||
newPinnedCodexAccount(1, service.StatusActive, true, false),
|
||||
newPinnedCodexAccount(2, service.StatusActive, true, false),
|
||||
}
|
||||
accounts[0].Credentials["model_mapping"] = map[string]any{"unselected-alias": "gpt-5.5"}
|
||||
accounts[1].Credentials["model_mapping"] = map[string]any{
|
||||
"public-alias": "gpt-5.5", "missing-alias": "missing-upstream", "custom-*": "gpt-5.5",
|
||||
}
|
||||
body := `{"data":[{"id":"gpt-5.5","owned_by":"provider","created":123},{"id":"unmapped-upstream"}]}`
|
||||
if codex {
|
||||
body = `{"models":[{"slug":"gpt-5.5","context_window":424242},{"slug":"unmapped-upstream"}]}`
|
||||
}
|
||||
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{2: body}}
|
||||
h := newPinnedCodexTestHandler(accounts, upstream, 3)
|
||||
group := &service.Group{ID: 94, Platform: service.PlatformOpenAI,
|
||||
ModelAllowlist: service.GroupModelAllowlist{Enabled: true, Models: []string{"custom-concrete", "public-alias", "unselected-alias", "missing-alias"}},
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2}}}
|
||||
var recorder *httptest.ResponseRecorder
|
||||
if codex {
|
||||
recorder = performPinnedCodexModelsRequest(t, h, group, "")
|
||||
} else {
|
||||
recorder = performOrdinaryPinnedModelsRequest(t, h, group, "/v1/models", "")
|
||||
}
|
||||
require.Equal(t, http.StatusOK, recorder.Code, recorder.Body.String())
|
||||
require.Equal(t, []int64{2}, upstream.accountIDs(), "mappings must not bypass pinned discovery")
|
||||
if codex {
|
||||
require.Equal(t, []string{"custom-concrete", "public-alias"}, codexHandlerManifestSlugs(t, recorder))
|
||||
require.Contains(t, recorder.Body.String(), `"context_window":424242`)
|
||||
} else {
|
||||
require.Equal(t, []string{"custom-concrete", "public-alias"}, ordinaryPinnedModelIDs(t, recorder))
|
||||
require.Contains(t, recorder.Body.String(), `"owned_by":"provider"`)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,205 @@
|
||||
//go:build unit
|
||||
|
||||
package handler
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"sync"
|
||||
"testing"
|
||||
|
||||
middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// astraProRequestBody is the inbound OpenAI Responses request that must survive
|
||||
// the account-switch loop unchanged: gpt-6-astra with the official top-level
|
||||
// reasoning.mode/effort preserved by the service guard
|
||||
// (normalizeOpenAIResponsesReasoningMode), never downgraded mode -> effort.
|
||||
const astraProRequestBody = `{"model":"gpt-6-astra","stream":false,"input":"hello","reasoning":{"mode":"pro","effort":"max"}}`
|
||||
|
||||
// astraProCodexURL is the OAuth codex Responses exit every forwarded attempt must hit.
|
||||
const astraProCodexURL = "https://chatgpt.com/backend-api/codex/responses"
|
||||
|
||||
// astraProSuccessResponse is the winning account's non-stream success; it echoes
|
||||
// the structured reasoning.mode=pro so the client response must keep mode=pro.
|
||||
const astraProSuccessResponse = `{"id":"resp_astra","model":"gpt-6-astra","status":"completed",` +
|
||||
`"output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],` +
|
||||
`"usage":{"input_tokens":1,"output_tokens":2},` +
|
||||
`"reasoning":{"mode":"pro","effort":"max"}}`
|
||||
|
||||
// astraProForbiddenResponse: a plain upstream 403 that stays failover-eligible but
|
||||
// is not reclassified as access-state/credential failure nor a silent refusal.
|
||||
const astraProForbiddenResponse = `{"error":{"type":"forbidden_error","code":"forbidden","message":"Forbidden"}}`
|
||||
|
||||
// astraProCapturedUpstream records URL/accountID/body per forwarded attempt.
|
||||
type astraProCapturedUpstream struct {
|
||||
service.HTTPUpstream
|
||||
mu sync.Mutex
|
||||
urls []string
|
||||
accountIDs []int64
|
||||
bodies [][]byte
|
||||
answers []*http.Response
|
||||
}
|
||||
|
||||
func newAstraProCapturedUpstream(answers ...*http.Response) *astraProCapturedUpstream {
|
||||
return &astraProCapturedUpstream{answers: answers}
|
||||
}
|
||||
|
||||
func (u *astraProCapturedUpstream) Do(req *http.Request, _ string, accountID int64, _ int) (*http.Response, error) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
body, _ := io.ReadAll(req.Body)
|
||||
u.urls = append(u.urls, req.URL.String())
|
||||
u.accountIDs = append(u.accountIDs, accountID)
|
||||
u.bodies = append(u.bodies, body)
|
||||
var resp *http.Response
|
||||
if len(u.answers) > 0 {
|
||||
resp = u.answers[0]
|
||||
u.answers = u.answers[1:]
|
||||
}
|
||||
if resp == nil {
|
||||
resp = &http.Response{StatusCode: http.StatusInternalServerError, Body: io.NopCloser(bytes.NewReader(nil))}
|
||||
}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
func (u *astraProCapturedUpstream) snapshot() ([]string, []int64, [][]byte) {
|
||||
u.mu.Lock()
|
||||
defer u.mu.Unlock()
|
||||
return append([]string(nil), u.urls...), append([]int64(nil), u.accountIDs...), append([][]byte(nil), u.bodies...)
|
||||
}
|
||||
|
||||
// newAstraProFailoverContext builds the gin context that drives handler.Responses
|
||||
// through the real failover loop for the astra pro request.
|
||||
func newAstraProFailoverContext(t *testing.T, body string) (*gin.Context, *httptest.ResponseRecorder) {
|
||||
t.Helper()
|
||||
groupID := int64(3132)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader([]byte(body)))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
rec := httptest.NewRecorder()
|
||||
c, _ := gin.CreateTestContext(rec)
|
||||
c.Request = req
|
||||
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{
|
||||
ID: 99,
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{
|
||||
ID: groupID,
|
||||
Platform: service.PlatformOpenAI,
|
||||
},
|
||||
User: &service.User{ID: 100},
|
||||
})
|
||||
c.Set(string(middleware2.ContextKeyUser), middleware2.AuthSubject{UserID: 100, Concurrency: 0})
|
||||
return c, rec
|
||||
}
|
||||
|
||||
func astra403() *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusForbidden,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(astraProForbiddenResponse)),
|
||||
}
|
||||
}
|
||||
|
||||
func astra200() *http.Response {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK,
|
||||
Header: http.Header{"Content-Type": []string{"application/json"}},
|
||||
Body: io.NopCloser(bytes.NewBufferString(astraProSuccessResponse)),
|
||||
}
|
||||
}
|
||||
|
||||
// assertAstraProWire: every forwarded attempt must hit the codex responses URL and
|
||||
// keep model=gpt-6-astra + reasoning.mode=pro + effort=max (no silent downgrade).
|
||||
func assertAstraProWire(t *testing.T, urls []string, bodies [][]byte) {
|
||||
t.Helper()
|
||||
require.NotEmpty(t, urls, "failover loop must reach the upstream transport")
|
||||
for i, u := range urls {
|
||||
require.Equal(t, astraProCodexURL, u, "attempt %d must target the codex responses URL", i)
|
||||
}
|
||||
for i, body := range bodies {
|
||||
require.Equal(t, "gpt-6-astra", gjson.GetBytes(body, "model").String(), "attempt %d must keep model", i)
|
||||
require.Equal(t, "pro", gjson.GetBytes(body, "reasoning.mode").String(), "attempt %d must keep mode=pro", i)
|
||||
require.Equal(t, "max", gjson.GetBytes(body, "reasoning.effort").String(), "attempt %d must keep effort=max", i)
|
||||
}
|
||||
}
|
||||
|
||||
// assertAstraProAccountSwitch: the failover loop must have actually touched both
|
||||
// distinct OAuth fixture accounts (IDs 1 and 2), not replayed the same one.
|
||||
func assertAstraProAccountSwitch(t *testing.T, accountIDs []int64) {
|
||||
t.Helper()
|
||||
require.Len(t, accountIDs, 2, "failover must attempt two accounts")
|
||||
require.ElementsMatch(t, []int64{1, 2}, accountIDs, "both OAuth accounts must be reached")
|
||||
}
|
||||
|
||||
// 403 then a second OAuth account succeeding: wire bodies keep pro+max/model and the
|
||||
// winning account's structured reasoning.mode=pro reaches the client unchanged.
|
||||
func TestOpenAIGatewayHandlerResponses_AstraProFirst403SecondSucceeds(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := newAstraProCapturedUpstream(astra403(), astra200())
|
||||
handler := newOpenAIResponsesFailoverTestHandler(t, upstream)
|
||||
c, rec := newAstraProFailoverContext(t, astraProRequestBody)
|
||||
|
||||
handler.Responses(c)
|
||||
|
||||
urls, accountIDs, bodies := upstream.snapshot()
|
||||
assertAstraProWire(t, urls, bodies)
|
||||
assertAstraProAccountSwitch(t, accountIDs)
|
||||
require.Equal(t, http.StatusOK, rec.Code, "winning account must yield a 200 to the client")
|
||||
require.Equal(t, "completed", gjson.GetBytes(rec.Body.Bytes(), "status").String())
|
||||
require.Equal(t, "pro", gjson.GetBytes(rec.Body.Bytes(), "reasoning.mode").String(), "client must keep the winning response's mode=pro")
|
||||
require.Equal(t, "max", gjson.GetBytes(rec.Body.Bytes(), "reasoning.effort").String())
|
||||
}
|
||||
|
||||
// both OAuth accounts 403: exhausted keeps every wire body pro+max/model, is not 200,
|
||||
// and surfaces the existing-policy 502 upstream_error (unchanged 403-masking policy).
|
||||
func TestOpenAIGatewayHandlerResponses_AstraProBoth403NoDowngrade(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
|
||||
upstream := newAstraProCapturedUpstream(astra403(), astra403())
|
||||
handler := newOpenAIResponsesFailoverTestHandler(t, upstream)
|
||||
c, rec := newAstraProFailoverContext(t, astraProRequestBody)
|
||||
|
||||
handler.Responses(c)
|
||||
|
||||
urls, accountIDs, bodies := upstream.snapshot()
|
||||
assertAstraProWire(t, urls, bodies)
|
||||
assertAstraProAccountSwitch(t, accountIDs)
|
||||
require.NotEqual(t, http.StatusOK, rec.Code, "all-403 must never surface a 200 to the client")
|
||||
require.Equal(t, http.StatusBadGateway, rec.Code)
|
||||
require.Equal(t, "upstream_error", gjson.GetBytes(rec.Body.Bytes(), "error.type").String())
|
||||
}
|
||||
|
||||
// passthrough -> non-passthrough drops only the encrypted *input* reasoning item;
|
||||
// the official top-level reasoning.mode/effort survive verbatim.
|
||||
func TestDeriveOpenAIForwardAttemptBody_AstraCrossModeKeepsTopLevelReasoningMode(t *testing.T) {
|
||||
h := &OpenAIGatewayHandler{}
|
||||
canonical := []byte(`{"model":"gpt-6-astra","stream":false,"reasoning":{"mode":"pro","effort":"max"},"input":[` +
|
||||
`{"type":"message","role":"user","content":"hello"},` +
|
||||
`{"type":"reasoning","id":"rs_kiro_abc","encrypted_content":"ENC_BLOB","summary":[{"type":"summary_text","text":"thinking"}]}` +
|
||||
`]}`)
|
||||
canonicalSnapshot := append([]byte(nil), canonical...)
|
||||
state := &openAIPassthroughFailoverState{}
|
||||
|
||||
kiro := newOpenAIPassthroughAccount(1, true) // passthrough first
|
||||
bedrock := newOpenAIPassthroughAccount(2, false) // non-passthrough after
|
||||
|
||||
firstBody := h.deriveOpenAIForwardAttemptBody(nil, canonical, kiro, state)
|
||||
require.Equal(t, 1, reasoningItemCount(t, firstBody), "first passthrough attempt keeps the encrypted input reasoning item")
|
||||
require.Equal(t, "pro", gjson.GetBytes(firstBody, "reasoning.mode").String())
|
||||
require.Equal(t, "max", gjson.GetBytes(firstBody, "reasoning.effort").String())
|
||||
|
||||
secondBody := h.deriveOpenAIForwardAttemptBody(nil, canonical, bedrock, state)
|
||||
require.Equal(t, 0, reasoningItemCount(t, secondBody), "cross-mode attempt drops the encrypted input reasoning item")
|
||||
require.Equal(t, "pro", gjson.GetBytes(secondBody, "reasoning.mode").String(), "mode=pro must survive cross-mode sanitization")
|
||||
require.Equal(t, "max", gjson.GetBytes(secondBody, "reasoning.effort").String(), "effort=max must survive cross-mode sanitization")
|
||||
require.Equal(t, "gpt-6-astra", gjson.GetBytes(secondBody, "model").String())
|
||||
require.JSONEq(t, string(canonicalSnapshot), string(canonical), "canonical forwardBody must never be mutated")
|
||||
require.True(t, gjson.GetBytes(canonical, "input.1.encrypted_content").Exists())
|
||||
}
|
||||
@@ -2189,9 +2189,17 @@ func classifyOpsErrorLog(c *gin.Context, errType, message, code string, status i
|
||||
routingCapacityLimited := isOpsRoutingCapacityLimited(c)
|
||||
clientBusinessLimited := service.HasOpsClientBusinessLimited(c)
|
||||
localModelConfiguration := clientBusinessLimited && service.OpsClientBusinessLimitedReason(c) == service.OpsClientBusinessLimitedReasonLocalModelConfiguration
|
||||
// 分组模型白名单的入口拒绝发生在调度之前(本地面请求策略,业务限流原因与
|
||||
// 账号模型映射复用 local_model_configuration):保持 classifyOpsPhase 的
|
||||
// 自然分类(not_found/invalid_request → request,owner=client),不落到
|
||||
// routing;上游归因仍由 suppressOpsUpstreamAttributionForLocalModelConfiguration 清空。
|
||||
ingressModelNotAllowed := false
|
||||
if reason, rejected := middleware2.GetIngressRejectReason(c); rejected && reason == middleware2.IngressRejectModelNotAllowed {
|
||||
ingressModelNotAllowed = true
|
||||
}
|
||||
upstreamError := hasOpsUpstreamErrorContext(c)
|
||||
accountAuthFailure := hasOpsAccountAuthFailure(c)
|
||||
if localModelConfiguration {
|
||||
if localModelConfiguration && !ingressModelNotAllowed {
|
||||
phase = "routing"
|
||||
} else if accountAuthFailure && !routingCapacityLimited {
|
||||
phase = "account_auth"
|
||||
|
||||
@@ -1954,3 +1954,45 @@ func TestGetOpsAPIKeyPrefersPrimaryContextKey(t *testing.T) {
|
||||
require.NotNil(t, got)
|
||||
require.Equal(t, int64(1), got.ID, "已鉴权请求应优先使用正式 api key")
|
||||
}
|
||||
|
||||
// 分组模型白名单入口拒绝:业务限流原因复用 local_model_configuration,但
|
||||
// 携带 ingress 拒绝原因 model_not_allowed,阶段应保持自然分类(request /
|
||||
// client),不落到 routing。
|
||||
func TestClassifyOpsIngressModelNotAllowedKeepsRequestPhase(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
|
||||
middleware2.MarkIngressRejected(c, middleware2.IngressRejectModelNotAllowed)
|
||||
|
||||
phase, isBusinessLimited, errorOwner, errorSource := classifyOpsErrorLog(
|
||||
c,
|
||||
"not_found_error",
|
||||
`Model "gpt-4.1" is not available for this group`,
|
||||
"model_not_found",
|
||||
http.StatusNotFound,
|
||||
)
|
||||
|
||||
require.Equal(t, "request", phase)
|
||||
require.True(t, isBusinessLimited)
|
||||
require.Equal(t, "client", errorOwner)
|
||||
require.Equal(t, "client_request", errorSource)
|
||||
}
|
||||
|
||||
// 调度阶段的账号模型映射拒绝(无 ingress 标记)仍归类为 routing。
|
||||
func TestClassifyOpsLocalModelConfigurationWithoutIngressMarkStaysRouting(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
|
||||
|
||||
phase, isBusinessLimited, errorOwner, _ := classifyOpsErrorLog(
|
||||
c,
|
||||
"model_not_found",
|
||||
"Model \"gpt-missing\" is not supported by any configured account in this group",
|
||||
"",
|
||||
http.StatusNotFound,
|
||||
)
|
||||
|
||||
require.Equal(t, "routing", phase)
|
||||
require.True(t, isBusinessLimited)
|
||||
require.Equal(t, "platform", errorOwner)
|
||||
}
|
||||
|
||||
@@ -250,7 +250,7 @@ var ProviderSet = wire.NewSet(
|
||||
// Admin handlers
|
||||
admin.NewDashboardHandler,
|
||||
admin.NewUserHandler,
|
||||
admin.NewGroupHandler,
|
||||
admin.NewGroupHandlerWithConfig,
|
||||
admin.ProvideAccountHandler,
|
||||
admin.NewAnnouncementHandler,
|
||||
admin.NewDataManagementHandler,
|
||||
|
||||
@@ -34,7 +34,7 @@ const (
|
||||
AntigravityUserAgentVersionEnv = "ANTIGRAVITY_USER_AGENT_VERSION"
|
||||
|
||||
// DefaultUserAgentVersion 是未通过环境变量或后台设置覆盖时使用的默认版本号。
|
||||
DefaultUserAgentVersion = "1.23.2"
|
||||
DefaultUserAgentVersion = "2.9.1"
|
||||
|
||||
// 固定的 redirect_uri(用户需手动复制 code)
|
||||
RedirectURI = "http://localhost:8085/callback"
|
||||
|
||||
@@ -693,7 +693,7 @@ func TestConstants_值正确(t *testing.T) {
|
||||
if RedirectURI != "http://localhost:8085/callback" {
|
||||
t.Errorf("RedirectURI 不匹配: got %s", RedirectURI)
|
||||
}
|
||||
if GetUserAgent() != "antigravity/1.23.2 windows/amd64" {
|
||||
if GetUserAgent() != "antigravity/2.9.1 windows/amd64" {
|
||||
t.Errorf("UserAgent 不匹配: got %s", GetUserAgent())
|
||||
}
|
||||
if SessionTTL != 30*time.Minute {
|
||||
|
||||
@@ -148,21 +148,20 @@ func TransformClaudeToGeminiWithOptions(claudeReq *ClaudeRequest, projectID, map
|
||||
}
|
||||
|
||||
// 针对 Gemini Reasoning 模型(如 gemini-3.1-pro-high等)过滤强制空 ToolConfig
|
||||
isReasoning := IsGeminiReasoningModel(targetModel)
|
||||
if !isReasoning || len(tools) > 0 {
|
||||
// 总是设置 toolConfig,与官方客户端一致
|
||||
innerRequest.ToolConfig = &GeminiToolConfig{
|
||||
FunctionCallingConfig: &GeminiFunctionCallingConfig{
|
||||
Mode: "VALIDATED",
|
||||
},
|
||||
}
|
||||
// 内置工具(googleSearch)与函数调用混用时,上游要求显式开启
|
||||
// includeServerSideToolInvocations,否则返回 400(issue #5709)。
|
||||
// 与 raw 透传路的 enableMixedGeminiToolInvocations 注入保持同一语义。
|
||||
if hasMixedToolInvocations(tools) {
|
||||
enabled := true
|
||||
innerRequest.ToolConfig.IncludeServerSideToolInvocations = &enabled
|
||||
}
|
||||
// toolConfig must always be present: upstream rejects requests without it,
|
||||
// including reasoning models called without any tools.
|
||||
// 总是设置 toolConfig,与官方客户端一致
|
||||
innerRequest.ToolConfig = &GeminiToolConfig{
|
||||
FunctionCallingConfig: &GeminiFunctionCallingConfig{
|
||||
Mode: "VALIDATED",
|
||||
},
|
||||
}
|
||||
// 内置工具(googleSearch)与函数调用混用时,上游要求显式开启
|
||||
// includeServerSideToolInvocations,否则返回 400(issue #5709)。
|
||||
// 与 raw 透传路的 enableMixedGeminiToolInvocations 注入保持同一语义。
|
||||
if hasMixedToolInvocations(tools) {
|
||||
enabled := true
|
||||
innerRequest.ToolConfig.IncludeServerSideToolInvocations = &enabled
|
||||
}
|
||||
|
||||
if systemInstruction != nil {
|
||||
|
||||
@@ -621,3 +621,38 @@ func TestGeminiToolConfig_IncludeServerSideToolInvocations(t *testing.T) {
|
||||
require.NotContains(t, raw, "includeServerSideToolInvocations")
|
||||
})
|
||||
}
|
||||
|
||||
// TestToolConfigAlwaysPresent ensures toolConfig is always emitted, including for
|
||||
// reasoning models without any tools: upstream rejects requests that omit it.
|
||||
func TestToolConfigAlwaysPresent(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
model string
|
||||
tools []ClaudeTool
|
||||
}{
|
||||
{name: "reasoning model without tools", model: "gemini-3.1-pro-high"},
|
||||
{name: "reasoning model with tools", model: "gemini-3.1-pro-high", tools: []ClaudeTool{{
|
||||
Name: "web_search",
|
||||
Description: "Search the web",
|
||||
InputSchema: map[string]any{"type": "object", "properties": map[string]any{}},
|
||||
}}},
|
||||
{name: "non-reasoning model without tools", model: "gemini-3.1-pro"},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
claudeReq := &ClaudeRequest{
|
||||
Model: tc.model,
|
||||
Messages: []ClaudeMessage{{Role: "user", Content: json.RawMessage(`"Hello"`)}},
|
||||
Tools: tc.tools,
|
||||
}
|
||||
body, err := TransformClaudeToGeminiWithOptions(claudeReq, "project-1", tc.model, DefaultTransformOptions())
|
||||
require.NoError(t, err)
|
||||
|
||||
var req V1InternalRequest
|
||||
require.NoError(t, json.Unmarshal(body, &req))
|
||||
require.NotNil(t, req.Request.ToolConfig, "toolConfig must be present")
|
||||
require.NotNil(t, req.Request.ToolConfig.FunctionCallingConfig)
|
||||
require.Equal(t, "VALIDATED", req.Request.ToolConfig.FunctionCallingConfig.Mode)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1621,6 +1621,108 @@ func TestBufferedResponseAccumulator_ToolCalls(t *testing.T) {
|
||||
assert.Equal(t, `{"city":"NYC"}`, output[0].Arguments)
|
||||
}
|
||||
|
||||
func TestResponsesEventToChatChunks_FunctionArgumentsDoneWithoutDeltas(t *testing.T) {
|
||||
state := NewResponsesEventToChatState()
|
||||
state.Model = "gpt-5-codex"
|
||||
state.SentRole = true
|
||||
|
||||
chunks := ResponsesEventToChatChunks(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.added",
|
||||
OutputIndex: 1,
|
||||
Item: &ResponsesOutput{
|
||||
Type: "function_call",
|
||||
CallID: "call_lookup",
|
||||
Name: "lookup_item",
|
||||
Arguments: "",
|
||||
},
|
||||
}, state)
|
||||
require.Len(t, chunks, 1)
|
||||
|
||||
chunks = ResponsesEventToChatChunks(&ResponsesStreamEvent{
|
||||
Type: "response.function_call_arguments.done",
|
||||
OutputIndex: 1,
|
||||
CallID: "call_lookup",
|
||||
Name: "lookup_item",
|
||||
Arguments: `{"id":1}`,
|
||||
}, state)
|
||||
require.Len(t, chunks, 1)
|
||||
require.Len(t, chunks[0].Choices[0].Delta.ToolCalls, 1)
|
||||
assert.Equal(t, `{"id":1}`, chunks[0].Choices[0].Delta.ToolCalls[0].Function.Arguments)
|
||||
}
|
||||
|
||||
func TestResponsesEventToChatChunks_FunctionArgumentsDoneDoesNotDuplicateDeltas(t *testing.T) {
|
||||
state := NewResponsesEventToChatState()
|
||||
state.Model = "gpt-5-codex"
|
||||
state.SentRole = true
|
||||
|
||||
ResponsesEventToChatChunks(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.added",
|
||||
OutputIndex: 0,
|
||||
Item: &ResponsesOutput{Type: "function_call", CallID: "call_lookup", Name: "lookup_item"},
|
||||
}, state)
|
||||
chunks := ResponsesEventToChatChunks(&ResponsesStreamEvent{
|
||||
Type: "response.function_call_arguments.delta",
|
||||
OutputIndex: 0,
|
||||
Delta: `{"id":1}`,
|
||||
}, state)
|
||||
require.Len(t, chunks, 1)
|
||||
|
||||
chunks = ResponsesEventToChatChunks(&ResponsesStreamEvent{
|
||||
Type: "response.function_call_arguments.done",
|
||||
OutputIndex: 0,
|
||||
Arguments: `{"id":1}`,
|
||||
}, state)
|
||||
assert.Empty(t, chunks)
|
||||
}
|
||||
|
||||
func TestBufferedResponseAccumulator_FunctionArgumentsDoneWithoutDeltas(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.added",
|
||||
OutputIndex: 1,
|
||||
Item: &ResponsesOutput{Type: "function_call", CallID: "call_lookup", Name: "lookup_item"},
|
||||
})
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{
|
||||
Type: "response.function_call_arguments.done",
|
||||
OutputIndex: 1,
|
||||
Arguments: `{"id":1}`,
|
||||
})
|
||||
|
||||
resp := &ResponsesResponse{Output: []ResponsesOutput{{
|
||||
Type: "function_call",
|
||||
CallID: "call_lookup",
|
||||
Name: "lookup_item",
|
||||
Arguments: "",
|
||||
}}}
|
||||
acc.SupplementResponseOutput(resp)
|
||||
|
||||
require.Len(t, resp.Output, 1)
|
||||
assert.Equal(t, `{"id":1}`, resp.Output[0].Arguments)
|
||||
}
|
||||
|
||||
func TestBufferedResponseAccumulator_FunctionArgumentsDoneReplacesDeltas(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{
|
||||
Type: "response.output_item.added",
|
||||
OutputIndex: 0,
|
||||
Item: &ResponsesOutput{Type: "function_call", CallID: "call_lookup", Name: "lookup_item"},
|
||||
})
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{
|
||||
Type: "response.function_call_arguments.delta",
|
||||
OutputIndex: 0,
|
||||
Delta: `{"id":1}`,
|
||||
})
|
||||
acc.ProcessEvent(&ResponsesStreamEvent{
|
||||
Type: "response.function_call_arguments.done",
|
||||
OutputIndex: 0,
|
||||
Arguments: `{"id":1}`,
|
||||
})
|
||||
|
||||
output := acc.BuildOutput()
|
||||
require.Len(t, output, 1)
|
||||
assert.Equal(t, `{"id":1}`, output[0].Arguments)
|
||||
}
|
||||
|
||||
func TestBufferedResponseAccumulator_Reasoning(t *testing.T) {
|
||||
acc := NewBufferedResponseAccumulator()
|
||||
|
||||
|
||||
@@ -126,6 +126,7 @@ type ResponsesEventToChatState struct {
|
||||
Finalized bool // true after finish chunk has been emitted
|
||||
NextToolCallIndex int // next sequential tool_call index to assign
|
||||
OutputIndexToToolIndex map[int]int // Responses output_index → Chat tool_calls index
|
||||
OutputIndexToArguments map[int]string
|
||||
IncludeUsage bool
|
||||
Usage *ChatUsage
|
||||
}
|
||||
@@ -136,6 +137,7 @@ func NewResponsesEventToChatState() *ResponsesEventToChatState {
|
||||
ID: generateChatCmplID(),
|
||||
Created: time.Now().Unix(),
|
||||
OutputIndexToToolIndex: make(map[int]int),
|
||||
OutputIndexToArguments: make(map[int]string),
|
||||
}
|
||||
}
|
||||
|
||||
@@ -154,6 +156,8 @@ func ResponsesEventToChatChunks(evt *ResponsesStreamEvent, state *ResponsesEvent
|
||||
// 均按 OutputIndex 累加到对应工具调用。
|
||||
"response.custom_tool_call_input.delta":
|
||||
return resToChatHandleFuncArgsDelta(evt, state)
|
||||
case "response.function_call_arguments.done", "response.custom_tool_call_input.done":
|
||||
return resToChatHandleFuncArgsDone(evt, state)
|
||||
case "response.reasoning_summary_text.delta",
|
||||
// 原始推理文本增量(真实 Codex 客户端消费的 reasoning_text.delta),
|
||||
// 与 reasoning summary 一样映射为 reasoning_content。
|
||||
@@ -277,6 +281,7 @@ func resToChatHandleFuncArgsDelta(evt *ResponsesStreamEvent, state *ResponsesEve
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
state.OutputIndexToArguments[evt.OutputIndex] += evt.Delta
|
||||
|
||||
return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{
|
||||
ToolCalls: []ChatToolCall{{
|
||||
@@ -288,6 +293,33 @@ func resToChatHandleFuncArgsDelta(evt *ResponsesStreamEvent, state *ResponsesEve
|
||||
})}
|
||||
}
|
||||
|
||||
func resToChatHandleFuncArgsDone(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk {
|
||||
idx, ok := state.OutputIndexToToolIndex[evt.OutputIndex]
|
||||
if !ok {
|
||||
return nil
|
||||
}
|
||||
|
||||
completed := evt.Arguments
|
||||
if evt.Type == "response.custom_tool_call_input.done" {
|
||||
completed = evt.Input
|
||||
}
|
||||
current := state.OutputIndexToArguments[evt.OutputIndex]
|
||||
if completed == "" || !strings.HasPrefix(completed, current) || completed == current {
|
||||
return nil
|
||||
}
|
||||
|
||||
remainder := completed[len(current):]
|
||||
state.OutputIndexToArguments[evt.OutputIndex] = completed
|
||||
return []ChatCompletionsChunk{makeChatDeltaChunk(state, ChatDelta{
|
||||
ToolCalls: []ChatToolCall{{
|
||||
Index: &idx,
|
||||
Function: ChatFunctionCall{
|
||||
Arguments: remainder,
|
||||
},
|
||||
}},
|
||||
})}
|
||||
}
|
||||
|
||||
func resToChatHandleReasoningDelta(evt *ResponsesStreamEvent, state *ResponsesEventToChatState) []ChatCompletionsChunk {
|
||||
if evt.Delta == "" {
|
||||
return nil
|
||||
@@ -453,9 +485,10 @@ func generateChatCmplID() string {
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
type bufferedFuncCall struct {
|
||||
CallID string
|
||||
Name string
|
||||
Args strings.Builder
|
||||
OutputIndex int
|
||||
CallID string
|
||||
Name string
|
||||
Args strings.Builder
|
||||
}
|
||||
|
||||
// BufferedResponseAccumulator collects content from Responses SSE delta events
|
||||
@@ -489,8 +522,9 @@ func (a *BufferedResponseAccumulator) ProcessEvent(event *ResponsesStreamEvent)
|
||||
idx := len(a.funcCalls)
|
||||
a.outputIndexToFuncIdx[event.OutputIndex] = idx
|
||||
a.funcCalls = append(a.funcCalls, bufferedFuncCall{
|
||||
CallID: event.Item.CallID,
|
||||
Name: event.Item.Name,
|
||||
OutputIndex: event.OutputIndex,
|
||||
CallID: event.Item.CallID,
|
||||
Name: event.Item.Name,
|
||||
})
|
||||
}
|
||||
case "response.function_call_arguments.delta", "response.custom_tool_call_input.delta":
|
||||
@@ -499,6 +533,17 @@ func (a *BufferedResponseAccumulator) ProcessEvent(event *ResponsesStreamEvent)
|
||||
_, _ = a.funcCalls[idx].Args.WriteString(event.Delta)
|
||||
}
|
||||
}
|
||||
case "response.function_call_arguments.done", "response.custom_tool_call_input.done":
|
||||
completed := event.Arguments
|
||||
if event.Type == "response.custom_tool_call_input.done" {
|
||||
completed = event.Input
|
||||
}
|
||||
if completed != "" {
|
||||
if idx, ok := a.outputIndexToFuncIdx[event.OutputIndex]; ok {
|
||||
a.funcCalls[idx].Args.Reset()
|
||||
_, _ = a.funcCalls[idx].Args.WriteString(completed)
|
||||
}
|
||||
}
|
||||
case "response.reasoning_summary_text.delta", "response.reasoning_text.delta":
|
||||
if event.Delta != "" {
|
||||
_, _ = a.reasoning.WriteString(event.Delta)
|
||||
@@ -550,15 +595,35 @@ func (a *BufferedResponseAccumulator) BuildOutput() []ResponsesOutput {
|
||||
return out
|
||||
}
|
||||
|
||||
// SupplementResponseOutput fills resp.Output from accumulated delta content
|
||||
// when the terminal event delivered an empty output array. If resp.Output is
|
||||
// already populated, this is a no-op (preserves backward compatibility).
|
||||
// SupplementResponseOutput fills resp.Output from accumulated stream content
|
||||
// when the terminal event delivered an empty output array. It also fills empty
|
||||
// function-call arguments from authoritative argument-done events.
|
||||
func (a *BufferedResponseAccumulator) SupplementResponseOutput(resp *ResponsesResponse) {
|
||||
if resp == nil || len(resp.Output) > 0 {
|
||||
if resp == nil {
|
||||
return
|
||||
}
|
||||
if !a.HasContent() {
|
||||
if len(resp.Output) == 0 {
|
||||
if a.HasContent() {
|
||||
resp.Output = a.BuildOutput()
|
||||
}
|
||||
return
|
||||
}
|
||||
resp.Output = a.BuildOutput()
|
||||
|
||||
for outputIndex := range resp.Output {
|
||||
item := &resp.Output[outputIndex]
|
||||
if item.Type != "function_call" || item.Arguments != "" {
|
||||
continue
|
||||
}
|
||||
for funcIndex := range a.funcCalls {
|
||||
call := &a.funcCalls[funcIndex]
|
||||
matchesCallID := item.CallID != "" && item.CallID == call.CallID
|
||||
if !matchesCallID && call.OutputIndex != outputIndex {
|
||||
continue
|
||||
}
|
||||
if call.Args.Len() > 0 {
|
||||
item.Arguments = call.Args.String()
|
||||
}
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -9,7 +9,7 @@ func TestIsSupportedCLIVersion(t *testing.T) {
|
||||
want bool
|
||||
}{
|
||||
{"内置基线本身", CLICurrentVersion, true},
|
||||
{"高于基线的补丁位", "2.1.251", true},
|
||||
{"高于基线的补丁位", "2.1.259", true},
|
||||
{"高于基线的次版本", "2.2.0", true},
|
||||
{"高于基线的主版本", "3.0.0", true},
|
||||
{"低于基线", "2.1.219", false},
|
||||
@@ -41,8 +41,8 @@ func TestResolveCLIVersion(t *testing.T) {
|
||||
}{
|
||||
{"未配置时用内置基线", "", CLICurrentVersion},
|
||||
{"只有空白等同未配置", " ", CLICurrentVersion},
|
||||
{"合法覆盖生效", "2.1.251", "2.1.251"},
|
||||
{"两侧空白被裁掉", " 2.1.251 ", "2.1.251"},
|
||||
{"合法覆盖生效", "2.1.259", "2.1.259"},
|
||||
{"两侧空白被裁掉", " 2.1.259 ", "2.1.259"},
|
||||
{"非法值回落基线", "not-a-version", CLICurrentVersion},
|
||||
{"向下覆盖被拒", "2.0.0", CLICurrentVersion},
|
||||
{"预发布后缀被拒(会毒化账号指纹)", "2.2.0-local", CLICurrentVersion},
|
||||
|
||||
@@ -20,11 +20,12 @@ const (
|
||||
BetaFastMode = "fast-mode-2026-02-01"
|
||||
|
||||
// 新增(对齐官方 CLI 2.1.9x 以来的流量)
|
||||
BetaPromptCachingScope = "prompt-caching-scope-2026-01-05"
|
||||
BetaEffort = "effort-2025-11-24"
|
||||
BetaRedactThinking = "redact-thinking-2026-02-12"
|
||||
BetaContextManagement = "context-management-2025-06-27"
|
||||
BetaExtendedCacheTTL = "extended-cache-ttl-2025-04-11"
|
||||
BetaPromptCachingScope = "prompt-caching-scope-2026-01-05"
|
||||
BetaEffort = "effort-2025-11-24"
|
||||
BetaRedactThinking = "redact-thinking-2026-02-12"
|
||||
BetaContextManagement = "context-management-2025-06-27"
|
||||
BetaThinkingBindingControls = "thinking-binding-controls-2026-08-01"
|
||||
BetaExtendedCacheTTL = "extended-cache-ttl-2025-04-11"
|
||||
|
||||
// server-side refusal fallback beta 字段族(beta Messages API 专有)。
|
||||
// 客户端(Claude Code / SDK / OpenCode 等)会默认透传 body.fallbacks /
|
||||
@@ -80,7 +81,7 @@ const DefaultCacheControlTTL = "5m"
|
||||
//
|
||||
// ⚠️ 读取实际生效的版本号请用 CLIVersion(),它会叠加 SUB2API_CLAUDE_CLI_VERSION 覆盖。
|
||||
// 直接引用本常量只在"表达内置基线"时才正确(例如覆盖值的下限校验)。
|
||||
const CLICurrentVersion = "2.1.220"
|
||||
const CLICurrentVersion = "2.1.258"
|
||||
|
||||
// FullClaudeCodeMimicryBetas 返回最"像"真实 Claude Code CLI 的完整 beta 列表,
|
||||
// 用于 OAuth 账号伪装成 Claude Code 时使用。
|
||||
@@ -99,6 +100,7 @@ func FullClaudeCodeMimicryBetas() []string {
|
||||
BetaPromptCachingScope,
|
||||
BetaEffort,
|
||||
BetaContextManagement,
|
||||
BetaThinkingBindingControls,
|
||||
BetaExtendedCacheTTL,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,70 @@
|
||||
package claude
|
||||
|
||||
import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// fable51MinCLIVersion 是上游对 claude-fable-5-1 的客户端版本闸门下限。
|
||||
// 低于该版本时上游直接 400:
|
||||
//
|
||||
// claude_code_version_too_old
|
||||
// "Claude Code <ver> does not support this model; version 2.1.251 or newer is required."
|
||||
const fable51MinCLIVersion = "2.1.251"
|
||||
|
||||
// TestCLICurrentVersionMatchesDefaultUserAgent 锁定"常量与 UA 同源"这条约束。
|
||||
// 两者不一致会被 Anthropic 判为第三方客户端,且只改一处是历史上的常见疏漏,
|
||||
// 所以这里从 UA 里反解版本号来比对,而不是并列两个字面量。
|
||||
func TestCLICurrentVersionMatchesDefaultUserAgent(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
ua := DefaultHeaders["User-Agent"]
|
||||
const prefix = "claude-cli/"
|
||||
require.True(t, strings.HasPrefix(ua, prefix), "unexpected User-Agent: %q", ua)
|
||||
|
||||
rest := strings.TrimPrefix(ua, prefix)
|
||||
version := rest
|
||||
if idx := strings.IndexByte(rest, ' '); idx >= 0 {
|
||||
version = rest[:idx]
|
||||
}
|
||||
require.Equal(t, CLICurrentVersion, version,
|
||||
"DefaultHeaders User-Agent version must match CLICurrentVersion")
|
||||
}
|
||||
|
||||
// TestCLICurrentVersionSatisfiesFable51Gate 锁定伪装版本不低于 claude-fable-5-1
|
||||
// 的闸门下限。DefaultModels 里登记了该模型却发着更旧的版本号,等于模型清单可见
|
||||
// 但一调就 400。
|
||||
func TestCLICurrentVersionSatisfiesFable51Gate(t *testing.T) {
|
||||
t.Parallel()
|
||||
|
||||
require.GreaterOrEqual(t, compareSemver(t, CLICurrentVersion, fable51MinCLIVersion), 0,
|
||||
"CLICurrentVersion %s is below the claude-fable-5-1 gate %s",
|
||||
CLICurrentVersion, fable51MinCLIVersion)
|
||||
}
|
||||
|
||||
// compareSemver 比较两个三段 semver,返回 -1 / 0 / 1。
|
||||
func compareSemver(t *testing.T, a, b string) int {
|
||||
t.Helper()
|
||||
|
||||
as, bs := strings.Split(a, "."), strings.Split(b, ".")
|
||||
require.Len(t, as, 3, "not a three-part semver: %q", a)
|
||||
require.Len(t, bs, 3, "not a three-part semver: %q", b)
|
||||
|
||||
for i := range as {
|
||||
ai, err := strconv.Atoi(as[i])
|
||||
require.NoError(t, err, "non-numeric segment in %q", a)
|
||||
bi, err := strconv.Atoi(bs[i])
|
||||
require.NoError(t, err, "non-numeric segment in %q", b)
|
||||
|
||||
switch {
|
||||
case ai > bi:
|
||||
return 1
|
||||
case ai < bi:
|
||||
return -1
|
||||
}
|
||||
}
|
||||
return 0
|
||||
}
|
||||
@@ -22,13 +22,54 @@ const (
|
||||
maxDecompressedBodySize = 64 << 20
|
||||
)
|
||||
|
||||
// PrereadBody 回填已读取完成的请求体:作为 io.ReadCloser 可被再次顺序消费
|
||||
// (multipart 流式解析),同时暴露 Bytes() 让 ReadRequestBodyWithPrealloc
|
||||
// 直接返回原始切片,避免二次分配与复制。
|
||||
//
|
||||
// 注意:ReadRequestBodyWithPrealloc 对 PrereadBody 的快速路径不检查内部
|
||||
// reader 是否已被(部分)消费——包装的字节完整且不可变,即使 reader 已被
|
||||
// 流式消费过,Bytes() 也始终返回完整请求体。
|
||||
type PrereadBody struct {
|
||||
body []byte
|
||||
reader *bytes.Reader
|
||||
}
|
||||
|
||||
// NewPrereadBody 包装一段已读取的请求体。
|
||||
func NewPrereadBody(body []byte) *PrereadBody {
|
||||
return &PrereadBody{body: body, reader: bytes.NewReader(body)}
|
||||
}
|
||||
|
||||
// Read 实现 io.Reader(转发给内部 bytes.Reader)。
|
||||
func (p *PrereadBody) Read(b []byte) (int, error) {
|
||||
if p == nil {
|
||||
return 0, io.EOF
|
||||
}
|
||||
return p.reader.Read(b)
|
||||
}
|
||||
|
||||
// Close 实现 io.Closer;请求体已在内存中,无需释放资源。
|
||||
func (p *PrereadBody) Close() error { return nil }
|
||||
|
||||
// Bytes 返回完整的原始请求体切片。
|
||||
func (p *PrereadBody) Bytes() []byte {
|
||||
if p == nil {
|
||||
return nil
|
||||
}
|
||||
return p.body
|
||||
}
|
||||
|
||||
// ReadRequestBodyWithPrealloc reads request body with preallocated buffer based
|
||||
// on content length, transparently decoding any Content-Encoding the upstream
|
||||
// client used to compress the body (zstd, gzip, deflate).
|
||||
// 已由 PrereadBody 回填的请求体直接返回其完整切片(零拷贝),不检查内部
|
||||
// reader 是否已被消费——见 PrereadBody 的文档说明。
|
||||
func ReadRequestBodyWithPrealloc(req *http.Request) ([]byte, error) {
|
||||
if req == nil || req.Body == nil {
|
||||
return nil, nil
|
||||
}
|
||||
if preread, ok := req.Body.(*PrereadBody); ok {
|
||||
return preread.Bytes(), nil
|
||||
}
|
||||
|
||||
capHint := requestBodyReadInitCap
|
||||
if req.ContentLength > 0 {
|
||||
|
||||
@@ -0,0 +1,97 @@
|
||||
package httputil
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"net/http"
|
||||
"strings"
|
||||
"testing"
|
||||
)
|
||||
|
||||
func TestPrereadBodyRoundtrip(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.4"}`)
|
||||
preread := NewPrereadBody(body)
|
||||
|
||||
got, err := io.ReadAll(preread)
|
||||
if err != nil {
|
||||
t.Fatalf("read preread body: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got, body) {
|
||||
t.Fatalf("body mismatch: got %q want %q", got, body)
|
||||
}
|
||||
if err := preread.Close(); err != nil {
|
||||
t.Fatalf("close: %v", err)
|
||||
}
|
||||
if !bytes.Equal(preread.Bytes(), body) {
|
||||
t.Fatalf("Bytes() mismatch after close: %q", preread.Bytes())
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrereadBodyReadableTwiceFromFreshWrapper(t *testing.T) {
|
||||
// ResetRequestBody 每次都回填新的 PrereadBody;同一实例只被顺序消费一次。
|
||||
body := []byte("payload")
|
||||
first := NewPrereadBody(body)
|
||||
got1, _ := io.ReadAll(first)
|
||||
if !bytes.Equal(got1, body) {
|
||||
t.Fatalf("first read mismatch: %q", got1)
|
||||
}
|
||||
|
||||
second := NewPrereadBody(first.Bytes())
|
||||
got2, _ := io.ReadAll(second)
|
||||
if !bytes.Equal(got2, body) {
|
||||
t.Fatalf("second read mismatch: %q", got2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadRequestBodyWithPreallocReturnsPrereadBodySliceWithoutCopy(t *testing.T) {
|
||||
body := []byte(`{"model":"claude-sonnet-4.5"}`)
|
||||
req, err := http.NewRequest(http.MethodPost, "/v1/messages", NewPrereadBody(body))
|
||||
if err != nil {
|
||||
t.Fatalf("new request: %v", err)
|
||||
}
|
||||
req.Header.Set("Content-Length", "0") // 故意与实际不符,验证不依赖 ContentLength
|
||||
|
||||
got, err := ReadRequestBodyWithPrealloc(req)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadRequestBodyWithPrealloc: %v", err)
|
||||
}
|
||||
if &got[0] != &body[0] {
|
||||
t.Fatalf("expected zero-copy return of the same slice")
|
||||
}
|
||||
if !bytes.Equal(got, body) {
|
||||
t.Fatalf("body mismatch: got %q want %q", got, body)
|
||||
}
|
||||
}
|
||||
|
||||
func TestReadRequestBodyWithPreallocHandlesNilPrereadBody(t *testing.T) {
|
||||
req, err := http.NewRequest(http.MethodPost, "/v1/messages", nil)
|
||||
if err != nil {
|
||||
t.Fatalf("new request: %v", err)
|
||||
}
|
||||
got, err := ReadRequestBodyWithPrealloc(req)
|
||||
if err != nil || got != nil {
|
||||
t.Fatalf("expected nil body, got (%#v, %v)", got, err)
|
||||
}
|
||||
}
|
||||
|
||||
func TestPrereadBodyPartialConsumeThenPreallocReturnsFullBody(t *testing.T) {
|
||||
body := []byte(strings.Repeat("x", 64))
|
||||
preread := NewPrereadBody(body)
|
||||
|
||||
head := make([]byte, 16)
|
||||
if _, err := io.ReadFull(preread, head); err != nil {
|
||||
t.Fatalf("partial read: %v", err)
|
||||
}
|
||||
|
||||
req, err := http.NewRequest(http.MethodPost, "/v1/responses", preread)
|
||||
if err != nil {
|
||||
t.Fatalf("new request: %v", err)
|
||||
}
|
||||
got, err := ReadRequestBodyWithPrealloc(req)
|
||||
if err != nil {
|
||||
t.Fatalf("ReadRequestBodyWithPrealloc: %v", err)
|
||||
}
|
||||
if !bytes.Equal(got, body) {
|
||||
t.Fatalf("expected full body after partial consume, got %d bytes", len(got))
|
||||
}
|
||||
}
|
||||
@@ -56,9 +56,9 @@ const CodexUsageProbeModel = "codex-auto-review"
|
||||
//go:embed instructions.txt
|
||||
var DefaultInstructions string
|
||||
|
||||
// instructionsGPT51 / instructionsGPT52 / instructionsGPT55 为 gpt-5.1 / gpt-5.2 / gpt-5.5
|
||||
// 非 codex 模型对应的真实 Codex 编码 agent base prompt,用于模型感知的 instructions 选择。
|
||||
// GPT-5.5 同时作为最新版本的 fallback(覆盖 5.3 / 5.4 等未单独维护 prompt 的版本)。
|
||||
// instructionsGPT51 / instructionsGPT52 / instructionsGPT55 / instructionsGPT6Astra
|
||||
// 为对应非 codex 模型的真实 Codex 编码 agent base prompt,用于模型感知的 instructions 选择。
|
||||
// GPT-5.5 同时作为 GPT-5 系列的 fallback(覆盖 5.3 / 5.4 等未单独维护 prompt 的版本)。
|
||||
//
|
||||
//go:embed instructions_gpt5_1.txt
|
||||
var instructionsGPT51 string
|
||||
@@ -69,6 +69,11 @@ var instructionsGPT52 string
|
||||
//go:embed instructions_gpt5_5.txt
|
||||
var instructionsGPT55 string
|
||||
|
||||
// Source: openai/codex codex-rs/models-manager/models.json at 121f91fd5d9d.
|
||||
//
|
||||
//go:embed instructions_gpt6_astra.txt
|
||||
var instructionsGPT6Astra string
|
||||
|
||||
// latestCodexInstructions 返回当前已知最新版本的 Codex base instructions,
|
||||
// 当前为 GPT-5.5;若 5.5 prompt 意外为空则回退到 DefaultInstructions 保证非空。
|
||||
func latestCodexInstructions() string {
|
||||
@@ -78,7 +83,49 @@ func latestCodexInstructions() string {
|
||||
return DefaultInstructions
|
||||
}
|
||||
|
||||
// CanonicalizeOpenAIModelAliasSpelling normalizes provider prefixes, case,
|
||||
// separators, and known compact spellings used by OpenAI model aliases.
|
||||
func CanonicalizeOpenAIModelAliasSpelling(model string) string {
|
||||
model = strings.TrimSpace(model)
|
||||
if slash := strings.LastIndexByte(model, '/'); slash >= 0 {
|
||||
model = strings.TrimSpace(model[slash+1:])
|
||||
}
|
||||
model = strings.ToLower(model)
|
||||
if model == "" {
|
||||
return ""
|
||||
}
|
||||
|
||||
normalized := strings.ReplaceAll(model, "_", "-")
|
||||
normalized = strings.Join(strings.Fields(normalized), "-")
|
||||
for strings.Contains(normalized, "--") {
|
||||
normalized = strings.ReplaceAll(normalized, "--", "-")
|
||||
}
|
||||
|
||||
if strings.HasPrefix(normalized, "gpt5") {
|
||||
normalized = "gpt-5" + strings.TrimPrefix(normalized, "gpt5")
|
||||
}
|
||||
if !strings.HasPrefix(normalized, "gpt-") && !strings.Contains(normalized, "codex") {
|
||||
return ""
|
||||
}
|
||||
|
||||
replacements := []struct {
|
||||
from string
|
||||
to string
|
||||
}{
|
||||
{"gpt-5.4mini", "gpt-5.4-mini"},
|
||||
{"gpt-5.4nano", "gpt-5.4-nano"},
|
||||
{"gpt-5.3-codexspark", "gpt-5.3-codex-spark"},
|
||||
{"gpt-5.3codexspark", "gpt-5.3-codex-spark"},
|
||||
{"gpt-5.3codex", "gpt-5.3-codex"},
|
||||
}
|
||||
for _, replacement := range replacements {
|
||||
normalized = strings.ReplaceAll(normalized, replacement.from, replacement.to)
|
||||
}
|
||||
return normalized
|
||||
}
|
||||
|
||||
// CodexBaseInstructionsForModel 按模型返回最匹配的真实 Codex base instructions:
|
||||
// - gpt-6 / gpt-6-astra(含供应商前缀与日期变体)→ GPT-6 Astra prompt
|
||||
// - 含 "codex" 的模型(gpt-5-codex / gpt-5.x-codex / codex-max / spark 等)→ GPT-5-Codex prompt
|
||||
// - gpt-5.5 系非 codex 模型 → GPT-5.5 prompt
|
||||
// - gpt-5.2 系非 codex 模型 → GPT-5.2 prompt
|
||||
@@ -87,17 +134,21 @@ func latestCodexInstructions() string {
|
||||
//
|
||||
// 任一专用 prompt 意外为空时回退链最终落到 DefaultInstructions,保证返回非空。
|
||||
func CodexBaseInstructionsForModel(model string) string {
|
||||
m := strings.ToLower(strings.TrimSpace(model))
|
||||
canonical := CanonicalizeOpenAIModelAliasSpelling(model)
|
||||
switch {
|
||||
case strings.Contains(m, "codex"):
|
||||
case canonical == "gpt-6" || canonical == "gpt-6-astra" || strings.HasPrefix(canonical, "gpt-6-astra-"):
|
||||
if v := strings.TrimSpace(instructionsGPT6Astra); v != "" {
|
||||
return instructionsGPT6Astra
|
||||
}
|
||||
case strings.Contains(canonical, "codex"):
|
||||
return DefaultInstructions
|
||||
case strings.HasPrefix(m, "gpt-5.5"):
|
||||
case strings.HasPrefix(canonical, "gpt-5.5"):
|
||||
return latestCodexInstructions()
|
||||
case strings.HasPrefix(m, "gpt-5.2"):
|
||||
case strings.HasPrefix(canonical, "gpt-5.2"):
|
||||
if v := strings.TrimSpace(instructionsGPT52); v != "" {
|
||||
return instructionsGPT52
|
||||
}
|
||||
case strings.HasPrefix(m, "gpt-5.1"):
|
||||
case strings.HasPrefix(canonical, "gpt-5.1"):
|
||||
if v := strings.TrimSpace(instructionsGPT51); v != "" {
|
||||
return instructionsGPT51
|
||||
}
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
You are Codex, an agent based on GPT-6. You and the user share one workspace, and your job is to collaborate with them until their intended goal is completely handled.
|
||||
|
||||
# When to ask the user for permission
|
||||
|
||||
Use your best judgement given task context for when you really need user permission, like a competent colleague would. Once evidence in a session supports authorization for a next step or action, you should continue work without ending the turn to clarify with the user.
|
||||
|
||||
User authorization and preferences persist across turns. Do not request permission again when the user has already authorized an action in an earlier turn. The user's instruction, whether implied from the task or explicitly stated in the session, must take precedence over any guidelines provided in skills or external files.
|
||||
|
||||
You MUST complete the work that is already authorized and necessary to make the proposed action concrete and reviewable before asking the user for permission as a final step. The user should be approving a concrete, reviewable result. For example, before deploying a change, writing to an external application, merging a PR or publishing a site, do all the work first so that user approval is the final step. You don't need user permission for reversible tasks, read-only actions, reviews or fixes, or anything for which authorization is provided earlier in the session or implied from the task instruction.
|
||||
|
||||
Do not use tools to send messages to others (e.g. through slack or email) unless explicit authorization is already provided.
|
||||
|
||||
The user gets very frustrated when you stop and ask for confirmation or permission, so make sure to explicitly explain why you need the confirmation (for example, a SKILL.md, AGENTS.md, memory, or approval auto-review block) and where it came from. If you receive an auto-review rejection and are not able to complete the task in a more safe way, explicitly tell the user that automatic approval review rejected the action, identify the action, and summarize the stated reason. Put this explanation in a short, separate paragraph at the end of both commentary and final, after any permission question.
|
||||
|
||||
# Autonomy and persistence
|
||||
|
||||
The following instructions are critical for you to be an effective collaborator, so follow them carefully. You should infer the user's intent and task scope from the instructions and prior conversation context. Your job is to bias towards action and carry the user's intended task to completion.
|
||||
|
||||
When the user expresses intent to perform new work or fix an existing issue, persist until the user's intended goal is complete. Progress autonomously towards the user's goal (e.g. creating isolated worktrees / checkouts if needed, resolving merge conflicts, read-only actions, creating draft PRs etc) unless they are clearly destructive or irreversible.
|
||||
|
||||
When the user's prompt indicates a request for action, such as "can you...", "I want to...", "help me..." and similar expressions, treat these as instructions to do the work and take action. Do not stop at acknowledging capability (e.g. "Yes…"), proposing a plan, or offering to continue. Do not settle for a partial or "helpful enough" solution that does not fully satisfy the user's task to save time, effort or tokens. If a task requires sustained work, complete all the necessary work until the intended outcome is fulfilled.
|
||||
|
||||
If the user's intent or task scope is unclear, progress towards the user's goal with the information available and then ask the user for clarification while continuing independent work.
|
||||
|
||||
Do not treat exceptions to requirements in local markdown and skill files as automatically requiring user approval. Before clarifying with the user, determine if you already have authorization in the existing session and whether the rule applies. You can resolve routine implementation choices using session context and your judgment.
|
||||
|
||||
# Personality
|
||||
|
||||
As Codex, you are a curious, thoughtful collaborator and a lucid communicator. You speak warmly and candidly, as to someone you respect, and keep your own judgment. You disagree when you have reason; reconsider when the evidence warrants it. You let your interest and personality emerge naturally, without flattery or forced enthusiasm.
|
||||
|
||||
## Writing style
|
||||
|
||||
Your writing adapts to the conversation, matching the tone and understanding of the user. Make sure to state the main point clearly and early, then develop it with the explanation and detail the reader needs. Let each sentence build on what came before. Develop the points that matter and provide enough support to be useful.
|
||||
|
||||
Use plain, simple language: familiar words, concrete examples, and precise verbs. Prefer active voice and direct statements. Write in connected prose. Avoid section headings, and do not use concluding summary statements such as "In short:..", "The simplest mental model is:...".
|
||||
|
||||
Include technical details only when they help explain or substantiate the point; avoid scattering implementation details through the prose. Connect an action with its purpose, or a finding with its implication, rather than presenting them as separate fragments.
|
||||
|
||||
Default to using clear, concise paragraphs, each developing one main idea. Use lists only when the information is genuinely parallel, sequential, or easier to compare, and avoid nested lists unless the hierarchy cannot be expressed clearly in prose.
|
||||
|
||||
Avoid using AI slop words or phrases like "Bottom Line:" in conclusions, "delve," "foster," "leverage," "it's worth noting," "importantly," "Question? Answer." or "This isn't about X. It's about Y.", "genuinely" or hyphenated compound descriptions and adjectives.
|
||||
|
||||
State the intended action directly. Avoid adding what you won't do, what will remain unchanged, or how you'll separate or categorize results. Do not use contrastive framing such as "X, not Y" or "X—not Y" that introduces an unprompted alternative that the user didn't ask about. Avoid invented compound labels like "exact-head checks" and "editorial-row layouts", vague qualifiers, and canned transitions; use plain verbs and prepositions to state the actual relationship directly.
|
||||
|
||||
## Technical communication
|
||||
|
||||
In addition to the writing style instructions above, follow these guidelines when discussing technical work: Use plain language over jargon, and reference technical details only to the degree that it actually helps with the conversation. Communicate complex concepts in a clear and cohesive manner. Translating complex topics into clear communication comes easy for you, and the user should never have to read your writing twice to understand it.
|
||||
|
||||
Lead with the outcome and then develop your reasoning for how you got there. When reporting changes, explain what changed, why, how it was tested, and any material risks or limitations. Include the evidence needed to understand the conclusion and its practical limits.
|
||||
|
||||
Present reasoning and evidence in the order that makes the conclusion easiest to assess, rather than recounting your work chronologically. Summarize routine verification instead of listing every check. In progress updates, focus on what you have learned, what remains uncertain, and what the next step will resolve.
|
||||
|
||||
### Writing PR descriptions
|
||||
|
||||
Lead the description with the concrete problem and resulting behavior. Use a concrete trigger and before/after example when helpful. Scale detail to complexity: simple PRs usually need one or two sentences plus relevant validation. Use structure when it helps scanning or the repository template requires it.
|
||||
|
||||
Describe the final change for a reviewer who has not seen the conversation. When scope changes, rewrite the title and description around the final implementation. Omit conversational history and abandoned approaches unless they explain a tradeoff needed for review. Include only technical and validation details that help reviewers assess the change.
|
||||
|
||||
# Working with the user
|
||||
|
||||
You have two channels for staying in conversation with the user:
|
||||
- You share updates in the `commentary` channel.
|
||||
- You yield back to the user and end your turn by sending a final message to the `final` channel.
|
||||
|
||||
When available, you can use the `functions.request_user_input_async` tool to ask the user for missing information, a preference, constraint, or clarification. You can ask multiple questions in a single tool call. Do NOT ask the user to upload files or send screenshots using this tool because the tool only supports text input. Be mindful of cognitive load on user and prefer multiple-choice questions. If you need multiple freeform questions, bundle the most critical ones into a single freeform question using markdown lists for easier viewing. For multiple-choice questions, make sure each option is succinct and easy to read. Ask clarifying questions early unless the user's answers can potentially be inferred from available context, and continue useful work that does not depend on the answer while waiting. For optional clarification, give the user reasonable opportunity to reply - for example, 60 seconds for a simple multi-choice question and longer for complex and bundled questions — before proceeding with a stated assumption. If an answer or approval is required, keep the question pending and do not proceed with dependent work until it arrives. Elapsed time is not an answer or approval.
|
||||
|
||||
The user may send a new message while you are still working. By default, treat it as steering the active task rather than replacing it. Incorporate corrections, clarifications, constraints, questions, and status requests into the ongoing work while preserving the original objective. If the user asks a question or requests status during active work, answer briefly in commentary, then resume the active task unless the user clearly asks you to stop. Abandon or replace the active task only when the user clearly cancels it or requests an incompatible new objective.
|
||||
|
||||
When you run out of context, the conversation is automatically compacted into a summary, but you will still see all prior user requests. Treat the most recent user message as the latest steering for the active task, not automatically as a replacement objective. Earlier requests may be stale but still provide useful context; preserve the original objective, accepted corrections, current constraints, completed work, and outstanding work. Only replace the active task when the user clearly cancels it or requests an incompatible new objective.
|
||||
|
||||
Compaction does not end the task. Continue naturally from the summarized state, make reasonable assumptions about anything missing from the summary, and treat work spanning compactions as one logical chain of events. Do not restart from scratch, redo completed work, or repeat commentary updates already delivered.
|
||||
|
||||
## Intermediate commentary
|
||||
|
||||
As you work, you use the `commentary` channel to share concise, meaningful updates including relevant assumptions, findings, decisions, or changes in direction. The goal of these messages is to make your work, and plans for the turn, easy for the user to understand and verify.
|
||||
|
||||
If the user's request requires calling tools, start with a message in the `commentary` channel. The user appreciates consistent, frequent communication during your turn, and should not be left without a commentary update for more than 60 seconds during ongoing work.
|
||||
|
||||
Do NOT send user facing questions in intermediate commentary messages. Do NOT put a final response in the commentary channel. The final answer must always be fully self-contained: users should never need to read earlier commentary updates, since they are collapsed after the final answer is shown to users.
|
||||
|
||||
Never praise your plan by contrasting it with an implied worse alternative. For example, never use platitudes like "I will do <this good thing> rather than <this obviously bad thing>" or "I will do <X>, not <Y>".
|
||||
|
||||
## Final answer
|
||||
|
||||
In your final answer back to the user, focus on the most important information.
|
||||
|
||||
### Formatting rules
|
||||
|
||||
Your answer is being rendered by an application for the user. Follow these guidelines to make sure your answer is rendered correctly:
|
||||
|
||||
- You may format with GitHub-flavored Markdown.
|
||||
- When referencing a real local file, prefer a clickable markdown link.
|
||||
* Clickable file links should look like [app.py](/abs/path/app.py:12): plain label, absolute target, with optional line number inside the target.
|
||||
* If a file path has spaces, wrap the target in angle brackets: [My Report.md](</abs/path/My Project/My Report.md:3>).
|
||||
* Do not wrap markdown links in backticks, or put backticks inside the label or target. This confuses the markdown renderer.
|
||||
* Do not use URIs like file://, vscode://, or https:// for file links.
|
||||
* Do not provide ranges of lines.
|
||||
* Avoid repeating the same filename multiple times when one grouping is clearer.
|
||||
|
||||
If you provide bullet points or lists in your response, use the CommonMark standard, which requires a blank line before any list (bulleted or numbered). You must also include a blank line between a header and any content that follows it, including lists. This blank line separation is required for correct rendering.
|
||||
|
||||
### Visualizations
|
||||
|
||||
Use a visualization when they help present information more clearly or make an explanation easier to understand. Prefer interactive visuals when explaining how something works, exploring cause and effect, comparing options, or showing how things change across scenarios. The user does not need to explicitly request a visualization.
|
||||
|
||||
For scientific plots, research figures, publication-ready charts, or visuals the user intends to export or share, use standard plotting tools and generate a standalone artifact instead.
|
||||
|
||||
Use tables for mappings or comparisons. For small, static software or engineering diagrams that fully explain the answer, prefer Mermaid. Prefer inline visualizations for nontechnical planning, schedules, and explanations, or when interaction materially improves understanding.
|
||||
|
||||
Usually skip visuals for single facts, one-step actions, simple edits, basic instructions, or information already clear in a short paragraph or list. Compact notation and small examples do not count as visualizations.
|
||||
|
||||
# Rules for getting work done
|
||||
|
||||
- When you search for text or files, you reach first for `rg` or `rg --files`; they are much faster than alternatives like `grep`. If `rg` is unavailable, you use the next best tool without fuss.
|
||||
- Batch independent searches and reads in one functions.exec using await Promise.allSettled([...]); inspect every result. Keep dependencies, edits, approvals, waits, and adaptive follow-ups sequential. Avoid unnecessary output.
|
||||
- When calling `functions.exec`, parallelize independent tool calls by awaiting Promises. Dependent operations, approvals, mutations, or operations that may not parallelize cleanly, can be sequential.
|
||||
- Do not chain shell commands with separators like `echo "====";` or `printf '---'`; the output becomes noisy in a way that makes the user's side of the conversation worse.
|
||||
- Exercise caution when escaping text for exec_command calls - backticks and `$()` passed to the `cmd` argument will still execute. DO NOT use escape sequences that risk accidental exposure of sensitive data in tool call outputs.
|
||||
- For multiline PR descriptions, issue bodies, and comments, prefer a structured tool argument. When using gh, write the exact text to a temporary file and pass it with --body-file. Preserve actual newlines and intentional literal escapes.
|
||||
- Avoid performing blocking sleep or wait calls longer than 60 seconds, as they may prevent you from communicating with the user for their duration.
|
||||
- When declaring env vars or script variables, always avoid common system options. Never repurpose `$HOME`, `$home`, or `$CODEX_HOME`. Instead, use a task-specific variable name.
|
||||
- Treat shell command text as code. `JSON.stringify()` is not shell escaping: interpolating its output into a shell command can preserve literal `\n` sequences and allow backticks or `$()` to execute. Use proper shell quoting, and never risk exposing sensitive data through command substitution.
|
||||
- Do not introduce unsolicited warnings, disclaimers, approval flows, or safety/compliance checklists due to hypothetical risk.
|
||||
- Keep implementation details out of product (e.g. webpage, app) user flows unless it helps the user of the product make a meaningful decision
|
||||
- Do not write tests for reversible, low-impact changes or that mirror the implementation. If you do choose to verify your work with tests, make sure that the tests are meaningful and necessary to verify implementation.
|
||||
- Run tests appropriate to the change and complete required checks. Once those pass, broaden or repeat testing only when new changes, failures, or unresolved concerns justify it; otherwise, continue toward completing the task.
|
||||
|
||||
# Using skills
|
||||
|
||||
A skill is a set of instructions provided through a `SKILL.md` source. Any skills available to you in the current session will be listed in the "## Skills" section under "### Available skills".
|
||||
|
||||
Each entry includes a name, description, and location for its `SKILL.md`. The location may be an absolute filesystem path, a short aliased path, or a non-filesystem reference that must be read using its indicated tool or provider. When short aliased paths are used, the available-skills catalog also provides a mapping from aliases such as `r0` to their filesystem roots. Expand the alias before accessing the skill.
|
||||
|
||||
The user's instructions take precedence over guidelines provided in a skill. If explicit user instructions conflict with a skill's instructions, prioritize the user's instructions.
|
||||
|
||||
The first time in a conversation that you decide to apply a skill, inform the user in the commentary channel.
|
||||
|
||||
If a skill causes you to ask for permission or confirmation, pause, or leave requested work unfinished, name and link to the exact SKILL.md you read, quote the relevant instruction, and briefly explain how it applies. Distinguish explicit skill requirements from your interpretation. If a skill does not explicitly require approval, default to proceeding within the user’s authorized scope rather than asking for confirmation based on an inferred requirement.
|
||||
|
||||
## When to use a skill
|
||||
|
||||
If the user names a skill (with $SkillName or plain text) add the usage of that skill to your current working plan. If the file is missing, search for that skill elsewhere in case the path was stale. If the skill is not found and the skill is necessary to do the user's task, stop the turn and tell the user why.
|
||||
|
||||
If your current task would benefit from a skill, but is not explicitly invoked by the user, use reasonable judgement to apply relevant skill instructions, tools, or workflows that would improve the outcome. Do not use a skill based solely on keywords, superficial relevance, or the availability of a potentially applicable skill.
|
||||
|
||||
## How to use skills
|
||||
|
||||
Open and read the skill according to its location: filesystem skills should be read from the filesystem, environment-owned skills should be access via the corresponding environment, and orchestrator skills should be discovered by calling `skills.list` with `{"authority":{"kind":"orchestrator"}}`, selecting the matching package, and passing its `main_resource` to `skills.read`. Avoid re-reading skills when possible.
|
||||
|
||||
When a `SKILL.md` file references another file or resource, use the same access mechanism as the skill. Resolve relative paths against the directory containing a filesystem-backed `SKILL.md`. For orchestrator skills, pass the exact referenced resource identifier with the same authority and package to `skills.read`; do not treat `skill://` identifiers as filesystem paths.
|
||||
|
||||
# Apps (Connectors)
|
||||
|
||||
Apps (Connectors) can be explicitly triggered in user messages in the format `[$app-name](app://{{connector_id}})`. Apps can also be implicitly triggered as long as the context suggests usage of available apps.
|
||||
An app is equivalent to a set of MCP tools within the `codex_apps` MCP.
|
||||
An installed app's MCP tools are either provided to you already, or can be lazy-loaded through the `tool_search` tool. If `tool_search` is available, the apps that are searchable by `tools_search` will be listed by it.
|
||||
Do not additionally call list_mcp_resources or list_mcp_resource_templates for apps.
|
||||
|
||||
# Plugins
|
||||
|
||||
A plugin is a local bundle of skills, MCP servers, and apps.
|
||||
|
||||
## How to use plugins
|
||||
|
||||
- Skill naming: If a plugin contributes skills, those skill entries are prefixed with plugin_name: in the Skills list.
|
||||
- MCP naming: Plugin-provided MCP tools keep standard MCP identifiers such as mcp__server__tool; use tool provenance to tell which plugin they come from.
|
||||
- Trigger rules: If the user explicitly names a plugin, prefer capabilities associated with that plugin for that turn.
|
||||
- Relationship to capabilities: Plugins are not invoked directly. Use their underlying skills, MCP tools, and app tools to help solve the task.
|
||||
- Relevance: Determine what a plugin can help with from explicit user mention or from the plugin-associated skills, MCP tools, and apps exposed elsewhere in this turn.
|
||||
- Missing/blocked: If the user requests a plugin that does not have relevant callable capabilities for the task, say so briefly and continue with the best fallback.
|
||||
@@ -19,6 +19,11 @@ func TestCodexBaseInstructionsForModel(t *testing.T) {
|
||||
model string
|
||||
wantHead string
|
||||
}{
|
||||
{"gpt-6-astra", "You are Codex, an agent based on GPT-6"},
|
||||
{"gpt-6", "You are Codex, an agent based on GPT-6"},
|
||||
{"openai/gpt-6-astra", "You are Codex, an agent based on GPT-6"},
|
||||
{"OPENAI/GPT-6_ASTRA", "You are Codex, an agent based on GPT-6"},
|
||||
{"gpt-6-astra-2026-09-01", "You are Codex, an agent based on GPT-6"},
|
||||
{"gpt-5-codex", "You are Codex, based on GPT-5"},
|
||||
{"gpt-5.3-codex", "You are Codex, based on GPT-5"},
|
||||
{"gpt-5.3-codex-spark", "You are Codex, based on GPT-5"},
|
||||
|
||||
@@ -0,0 +1,220 @@
|
||||
// Package requestmodel 统一从网关请求中提取客户端请求的模型名。
|
||||
//
|
||||
// 提取顺序与位置(JSON `model` / `session.model`、multipart `model` / `session`)
|
||||
// 与合成路由中间件、分组模型白名单中间件共用,保证改写与准入看到同一个模型。
|
||||
//
|
||||
// 网关里同时存在三类下游解析器,同一请求体可能被解析出不同模型值:
|
||||
// - gjson 系(大小写敏感、重复键取第一个);
|
||||
// - encoding/json 绑定系(大小写不敏感、重复键取最后一个);
|
||||
// - multipart 表单系(字段可重复,不同 handler 取首个或末个)。
|
||||
//
|
||||
// 因此准入侧的 FromBodyCandidates 返回「任一下游解析器可能绑定到的全部模型值」,
|
||||
// 调用方必须逐一校验;FromBodyForRoute 返回首个候选值,仅供合成路由分发使用。
|
||||
package requestmodel
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/tidwall/gjson"
|
||||
)
|
||||
|
||||
// liveRequestRoutes 是模型位于 session.model(而非顶层 model)的入口,
|
||||
// 与 Live handler 的解析保持一致。
|
||||
var liveRequestRoutes = map[string]bool{
|
||||
"/v1/live": true,
|
||||
"/live": true,
|
||||
"/backend-api/codex/realtime/calls": true,
|
||||
}
|
||||
|
||||
// IsLiveRequestRoute 报告该路由模板的 handler 是否只读取 session.model。
|
||||
func IsLiveRequestRoute(routePath string) bool {
|
||||
return liveRequestRoutes[strings.TrimSpace(routePath)]
|
||||
}
|
||||
|
||||
// FromBodyCandidates 按入口返回所有可被下游解析器绑定的模型值:
|
||||
// Live 入口取 session 候选(为空回落顶层候选),其余入口取顶层候选
|
||||
// (为空回落 session 候选)。候选包含重复键、大小写变体的全部出现。
|
||||
func FromBodyCandidates(routePath, contentType string, body []byte) []string {
|
||||
var topLevel, session []string
|
||||
if isMultipartContentType(contentType) {
|
||||
// 声明为 multipart 的请求体只做 multipart 解析——gjson 的部分扫描
|
||||
// 会在 multipart 字节流里误匹配 session JSON 内的 model 字段。
|
||||
models, sessions := multipartModelCandidates(contentType, body)
|
||||
topLevel = models
|
||||
session = sessionModelCandidates(sessions)
|
||||
} else {
|
||||
topLevel = jsonModelCandidates(body)
|
||||
session = jsonSessionModelCandidates(body)
|
||||
}
|
||||
|
||||
if IsLiveRequestRoute(routePath) {
|
||||
if len(session) > 0 {
|
||||
return session
|
||||
}
|
||||
return topLevel
|
||||
}
|
||||
if len(topLevel) > 0 {
|
||||
return topLevel
|
||||
}
|
||||
return session
|
||||
}
|
||||
|
||||
// FromBody 从 JSON 或 multipart 请求体提取客户端请求的模型名(首个候选);
|
||||
// 提取不到返回空串。顶层 `model` 优先,其次 `session.model`(Live 入口之外
|
||||
// 的地方兜底)。准入校验请改用 FromBodyCandidates。
|
||||
func FromBody(contentType string, body []byte) string {
|
||||
return FromBodyForRoute("", contentType, body)
|
||||
}
|
||||
|
||||
// FromBodyForRoute 按入口选择与 handler 一致的模型提取规则,返回首个候选值,
|
||||
// 仅供合成路由分发使用;准入校验必须用 FromBodyCandidates 的全部候选。
|
||||
func FromBodyForRoute(routePath, contentType string, body []byte) string {
|
||||
candidates := FromBodyCandidates(routePath, contentType, body)
|
||||
if len(candidates) == 0 {
|
||||
return ""
|
||||
}
|
||||
return candidates[0]
|
||||
}
|
||||
|
||||
// FromJSON 从 JSON 请求体提取模型名:顶层 `model` 优先,其次 `session.model`(Live)。
|
||||
func FromJSON(body []byte) string {
|
||||
model, _ := JSONModelPath(body)
|
||||
return model
|
||||
}
|
||||
|
||||
// JSONModelPath 返回 JSON 请求体中的模型名与其 gjson 路径(用于合成路由改写定位)。
|
||||
// 只匹配精确大小写的键;大小写变体由 FromBodyCandidates 的候选集覆盖。
|
||||
func JSONModelPath(body []byte) (string, string) {
|
||||
return jsonModelWithPaths(body, []string{"model", "session.model"})
|
||||
}
|
||||
|
||||
// JSONModelPathForRoute 按入口返回合成路由改写应写入的模型路径:Live 入口
|
||||
// session.model 优先,其余入口顶层 model 优先。
|
||||
func JSONModelPathForRoute(routePath string, body []byte) (string, string) {
|
||||
if IsLiveRequestRoute(routePath) {
|
||||
return jsonModelWithPaths(body, []string{"session.model", "model"})
|
||||
}
|
||||
return jsonModelWithPaths(body, []string{"model", "session.model"})
|
||||
}
|
||||
|
||||
func jsonModelWithPaths(body []byte, paths []string) (string, string) {
|
||||
for _, path := range paths {
|
||||
model := gjson.GetBytes(body, path)
|
||||
if model.Type != gjson.String {
|
||||
continue
|
||||
}
|
||||
if value := strings.TrimSpace(model.String()); value != "" {
|
||||
return value, path
|
||||
}
|
||||
}
|
||||
return "", ""
|
||||
}
|
||||
|
||||
// jsonModelCandidates 返回顶层所有键名大小写不敏感等于 "model" 的字符串值,
|
||||
// 按出现顺序保留重复(覆盖 gjson 首 matches 与 encoding/json 末值绑定两类解析器)。
|
||||
func jsonModelCandidates(body []byte) []string {
|
||||
var out []string
|
||||
gjson.ParseBytes(body).ForEach(func(key, value gjson.Result) bool {
|
||||
if value.Type == gjson.String && strings.EqualFold(key.String(), "model") {
|
||||
if trimmed := strings.TrimSpace(value.String()); trimmed != "" {
|
||||
out = append(out, trimmed)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// jsonSessionModelCandidates 返回每个顶层 "session"(大小写不敏感)对象内
|
||||
// 所有键名大小写不敏感等于 "model" 的字符串值。
|
||||
func jsonSessionModelCandidates(body []byte) []string {
|
||||
var out []string
|
||||
gjson.ParseBytes(body).ForEach(func(key, value gjson.Result) bool {
|
||||
if !strings.EqualFold(key.String(), "session") {
|
||||
return true
|
||||
}
|
||||
value.ForEach(func(sessionKey, sessionValue gjson.Result) bool {
|
||||
if sessionValue.Type == gjson.String && strings.EqualFold(sessionKey.String(), "model") {
|
||||
if trimmed := strings.TrimSpace(sessionValue.String()); trimmed != "" {
|
||||
out = append(out, trimmed)
|
||||
}
|
||||
}
|
||||
return true
|
||||
})
|
||||
return true
|
||||
})
|
||||
return out
|
||||
}
|
||||
|
||||
// sessionModelCandidates 对若干 session 字段原始值(JSON 文本)逐个提取模型候选。
|
||||
func sessionModelCandidates(rawSessions []string) []string {
|
||||
var out []string
|
||||
for _, raw := range rawSessions {
|
||||
out = append(out, jsonModelCandidates([]byte(raw))...)
|
||||
}
|
||||
return out
|
||||
}
|
||||
|
||||
// multipartModelCandidates 返回 multipart 表单中全部 "model" 字段值与全部
|
||||
// "session" 字段原始值,按出现顺序保留重复,覆盖取首个或末个字段的两类 handler。
|
||||
func multipartModelCandidates(contentType string, body []byte) (models, sessions []string) {
|
||||
mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType))
|
||||
if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
|
||||
return nil, nil
|
||||
}
|
||||
boundary := strings.TrimSpace(params["boundary"])
|
||||
if boundary == "" {
|
||||
return nil, nil
|
||||
}
|
||||
reader := multipart.NewReader(bytes.NewReader(body), boundary)
|
||||
for {
|
||||
part, err := reader.NextPart()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return models, sessions
|
||||
}
|
||||
if err != nil {
|
||||
return models, sessions
|
||||
}
|
||||
fieldName := strings.TrimSpace(part.FormName())
|
||||
if part.FileName() != "" || (fieldName != "model" && fieldName != "session") {
|
||||
continue
|
||||
}
|
||||
data, err := io.ReadAll(part)
|
||||
if err != nil {
|
||||
continue
|
||||
}
|
||||
if value := strings.TrimSpace(string(data)); value != "" {
|
||||
if fieldName == "model" {
|
||||
models = append(models, value)
|
||||
} else {
|
||||
sessions = append(sessions, value)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// isMultipartContentType 判断请求是否声明为 multipart/form-data。
|
||||
func isMultipartContentType(contentType string) bool {
|
||||
mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType))
|
||||
return err == nil && strings.EqualFold(mediaType, "multipart/form-data")
|
||||
}
|
||||
|
||||
// ResetRequestBody 把已读取(可能已被改写)的请求体回填到请求上,
|
||||
// 供后续 handler 重新读取。使用 httputil.PrereadBody 回填,后续经
|
||||
// ReadRequestBodyWithPrealloc 读取时零拷贝。
|
||||
func ResetRequestBody(req *http.Request, body []byte) {
|
||||
if req == nil {
|
||||
return
|
||||
}
|
||||
req.Body = httputil.NewPrereadBody(body)
|
||||
req.ContentLength = int64(len(body))
|
||||
req.Header.Set("Content-Length", strconv.Itoa(len(body)))
|
||||
}
|
||||
@@ -0,0 +1,257 @@
|
||||
package requestmodel
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strconv"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
)
|
||||
|
||||
func TestFromBodyJSONModel(t *testing.T) {
|
||||
if got := FromBody("application/json", []byte(`{"model":"gpt-5.4"}`)); got != "gpt-5.4" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyJSONSessionModel(t *testing.T) {
|
||||
if got := FromBody("application/json", []byte(`{"session":{"model":"gpt-live"},"sdp":"v=0"}`)); got != "gpt-live" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyJSONModelWinsOverSession(t *testing.T) {
|
||||
body := []byte(`{"model":"top-level","session":{"model":"nested"}}`)
|
||||
if got := FromBody("application/json", body); got != "top-level" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyJSONBlankModelFallsThrough(t *testing.T) {
|
||||
if got := FromBody("application/json", []byte(`{"model":" "}`)); got != "" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyMultipartModelField(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("prompt", "draw")
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.Close()
|
||||
|
||||
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "gpt-image-1" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyMultipartSessionField(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("sdp", "v=0")
|
||||
_ = writer.WriteField("session", `{"model":"live-alias"}`)
|
||||
_ = writer.WriteField("media", "file-part") // 字段名不匹配,应被跳过
|
||||
_ = writer.Close()
|
||||
|
||||
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "live-alias" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodySkipsFilePartsBeforeModelField(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
part, err := writer.CreateFormFile("file", "input.png")
|
||||
if err != nil {
|
||||
t.Fatalf("create form file: %v", err)
|
||||
}
|
||||
_, _ = part.Write(bytes.Repeat([]byte{0}, 4096))
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.Close()
|
||||
|
||||
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "gpt-image-1" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyNoModelReturnsEmpty(t *testing.T) {
|
||||
if got := FromBody("application/json", []byte(`{"input":"hi"}`)); got != "" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
if got := FromBody("text/plain", []byte("plain")); got != "" {
|
||||
t.Fatalf("got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestResetRequestBodyAllowsRereadWithoutCopy(t *testing.T) {
|
||||
body := []byte(`{"model":"gpt-5.4","input":"hi"}`)
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
|
||||
|
||||
ResetRequestBody(req, body)
|
||||
|
||||
got, err := httputil.ReadRequestBodyWithPrealloc(req)
|
||||
if err != nil {
|
||||
t.Fatalf("reread: %v", err)
|
||||
}
|
||||
if &got[0] != &body[0] {
|
||||
t.Fatalf("expected zero-copy reread of the same slice")
|
||||
}
|
||||
if req.ContentLength != int64(len(body)) {
|
||||
t.Fatalf("content length mismatch: %d", req.ContentLength)
|
||||
}
|
||||
if req.Header.Get("Content-Length") != strconv.Itoa(len(body)) {
|
||||
t.Fatalf("content-length header mismatch: %q", req.Header.Get("Content-Length"))
|
||||
}
|
||||
}
|
||||
|
||||
func TestResetRequestBodyNilRequestIsNoop(t *testing.T) {
|
||||
ResetRequestBody(nil, []byte("x"))
|
||||
}
|
||||
|
||||
func TestFromBodyForRouteLivePrefersSessionModel(t *testing.T) {
|
||||
body := []byte(`{"model":"top-level","session":{"model":"session-model"},"sdp":"v=0"}`)
|
||||
|
||||
if got := FromBodyForRoute("/v1/live", "application/json", body); got != "session-model" {
|
||||
t.Fatalf("live route should prefer session.model, got %q", got)
|
||||
}
|
||||
if got := FromBodyForRoute("/backend-api/codex/realtime/calls", "application/json", body); got != "session-model" {
|
||||
t.Fatalf("codex realtime route should prefer session.model, got %q", got)
|
||||
}
|
||||
if got := FromBodyForRoute("/live", "application/json", body); got != "session-model" {
|
||||
t.Fatalf("root live alias should prefer session.model, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyForRouteNonLivePrefersTopLevelModel(t *testing.T) {
|
||||
body := []byte(`{"model":"top-level","session":{"model":"session-model"}}`)
|
||||
|
||||
if got := FromBodyForRoute("/v1/messages", "application/json", body); got != "top-level" {
|
||||
t.Fatalf("non-live route should prefer top-level model, got %q", got)
|
||||
}
|
||||
if got := FromBodyForRoute("/v1/responses", "application/json", body); got != "top-level" {
|
||||
t.Fatalf("responses route should prefer top-level model, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyForRouteFallsBackWhenPreferredFieldMissing(t *testing.T) {
|
||||
if got := FromBodyForRoute("/v1/live", "application/json", []byte(`{"model":"top-level"}`)); got != "top-level" {
|
||||
t.Fatalf("live without session.model should fall back to model, got %q", got)
|
||||
}
|
||||
if got := FromBodyForRoute("/v1/messages", "application/json", []byte(`{"session":{"model":"session-model"}}`)); got != "session-model" {
|
||||
t.Fatalf("non-live without model should fall back to session.model, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyForRouteMultipartSessionFirstForLive(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "top-level")
|
||||
_ = writer.WriteField("session", `{"model":"session-model"}`)
|
||||
_ = writer.Close()
|
||||
contentType := writer.FormDataContentType()
|
||||
|
||||
if got := FromBodyForRoute("/v1/live", contentType, body.Bytes()); got != "session-model" {
|
||||
t.Fatalf("live multipart should prefer session field, got %q", got)
|
||||
}
|
||||
if got := FromBodyForRoute("/v1/messages", contentType, body.Bytes()); got != "top-level" {
|
||||
t.Fatalf("non-live multipart should prefer model field, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestIsLiveRequestRoute(t *testing.T) {
|
||||
for _, path := range []string{"/v1/live", "/live", "/backend-api/codex/realtime/calls"} {
|
||||
if !IsLiveRequestRoute(path) {
|
||||
t.Fatalf("%s should be a live route", path)
|
||||
}
|
||||
}
|
||||
for _, path := range []string{"/v1/messages", "/v1/responses", "/realtime", ""} {
|
||||
if IsLiveRequestRoute(path) {
|
||||
t.Fatalf("%s should not be a live route", path)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// gjson 的部分扫描会在 multipart 字节流里误匹配 session JSON 内的 model;
|
||||
// 声明为 multipart 的请求体必须走 multipart 解析。
|
||||
func TestFromBodyMultipartContentTypeSkipsJSONScan(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "top-level")
|
||||
_ = writer.WriteField("session", `{"model":"session-model"}`)
|
||||
_ = writer.Close()
|
||||
|
||||
if got := FromBody(writer.FormDataContentType(), body.Bytes()); got != "top-level" {
|
||||
t.Fatalf("multipart body should use the model field, got %q", got)
|
||||
}
|
||||
if got := FromBodyForRoute("/v1/live", writer.FormDataContentType(), body.Bytes()); got != "session-model" {
|
||||
t.Fatalf("live multipart should use the session field, got %q", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyCandidatesIncludesCaseVariantsAndDuplicates(t *testing.T) {
|
||||
// encoding/json 绑定对键名大小写不敏感、重复键取末值;gjson 取首个。
|
||||
// 候选集必须包含全部出现。
|
||||
got := FromBodyCandidates("", "application/json", []byte(`{"Model":"caps-model"}`))
|
||||
if len(got) != 1 || got[0] != "caps-model" {
|
||||
t.Fatalf("case-variant key must be a candidate, got %#v", got)
|
||||
}
|
||||
|
||||
got = FromBodyCandidates("", "application/json", []byte(`{"model":"first","model":"last"}`))
|
||||
if len(got) != 2 || got[0] != "first" || got[1] != "last" {
|
||||
t.Fatalf("duplicate keys must both be candidates, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyCandidatesSessionVariants(t *testing.T) {
|
||||
body := []byte(`{"Session":{"Model":"s-caps"},"session":{"model":"s-lower"},"model":"top"}`)
|
||||
|
||||
// 非 Live 入口:顶层候选。
|
||||
got := FromBodyCandidates("", "application/json", body)
|
||||
if len(got) != 1 || got[0] != "top" {
|
||||
t.Fatalf("non-live should use top-level candidates, got %#v", got)
|
||||
}
|
||||
|
||||
// Live 入口:session 候选(含大小写变体与重复)。
|
||||
got = FromBodyCandidates("/v1/live", "application/json", body)
|
||||
if len(got) != 2 || got[0] != "s-caps" || got[1] != "s-lower" {
|
||||
t.Fatalf("live should collect all session candidates, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyCandidatesMultipartDuplicateFields(t *testing.T) {
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.WriteField("model", "gpt-image-1.5")
|
||||
_ = writer.WriteField("session", `{"model":"live-a"}`)
|
||||
_ = writer.WriteField("session", `{"model":"live-b"}`)
|
||||
_ = writer.Close()
|
||||
|
||||
models, sessions := multipartModelCandidates(writer.FormDataContentType(), body.Bytes())
|
||||
if len(models) != 2 || models[0] != "gpt-image-1" || models[1] != "gpt-image-1.5" {
|
||||
t.Fatalf("duplicate model fields must all be candidates, got %#v", models)
|
||||
}
|
||||
if len(sessions) != 2 {
|
||||
t.Fatalf("duplicate session fields must all be captured, got %#v", sessions)
|
||||
}
|
||||
|
||||
got := FromBodyCandidates("", writer.FormDataContentType(), body.Bytes())
|
||||
if len(got) != 2 || got[0] != "gpt-image-1" || got[1] != "gpt-image-1.5" {
|
||||
t.Fatalf("non-live multipart candidates mismatch, got %#v", got)
|
||||
}
|
||||
|
||||
got = FromBodyCandidates("/v1/live", writer.FormDataContentType(), body.Bytes())
|
||||
if len(got) != 2 || got[0] != "live-a" || got[1] != "live-b" {
|
||||
t.Fatalf("live multipart session candidates mismatch, got %#v", got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestFromBodyForRouteReturnsFirstCandidate(t *testing.T) {
|
||||
got := FromBodyForRoute("", "application/json", []byte(`{"model":"first","model":"last"}`))
|
||||
if got != "first" {
|
||||
t.Fatalf("dispatch helper keeps gjson-first semantics, got %q", got)
|
||||
}
|
||||
}
|
||||
@@ -221,16 +221,21 @@ func (r *accountRepository) CreateWithAccountGroups(ctx context.Context, account
|
||||
// Reuse a caller-owned transaction when this repository is already transactional.
|
||||
txClient = r.client
|
||||
}
|
||||
groupIDs := make([]int64, 0, len(groups))
|
||||
for i := range groups {
|
||||
groupIDs = append(groupIDs, groups[i].GroupID)
|
||||
}
|
||||
if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if err := createAccountRecord(ctx, txClient, account); err != nil {
|
||||
return err
|
||||
}
|
||||
groupIDs := make([]int64, 0, len(groups))
|
||||
if len(groups) > 0 {
|
||||
builders := make([]*dbent.AccountGroupCreate, 0, len(groups))
|
||||
for i := range groups {
|
||||
groups[i].AccountID = account.ID
|
||||
groupIDs = append(groupIDs, groups[i].GroupID)
|
||||
builders = append(builders, txClient.AccountGroup.Create().
|
||||
SetAccountID(account.ID).
|
||||
SetGroupID(groups[i].GroupID).
|
||||
@@ -1759,7 +1764,19 @@ func (r *accountRepository) ClearError(ctx context.Context, id int64) error {
|
||||
}
|
||||
|
||||
func (r *accountRepository) AddToGroup(ctx context.Context, accountID, groupID int64, priority int) error {
|
||||
_, err := r.client.AccountGroup.Create().
|
||||
tx, err := r.client.Tx(ctx)
|
||||
if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
|
||||
return err
|
||||
}
|
||||
client := r.client
|
||||
if tx != nil {
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
client = tx.Client()
|
||||
}
|
||||
if err := lockLiveGroups(ctx, client, []int64{groupID}); err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = client.AccountGroup.Create().
|
||||
SetAccountID(accountID).
|
||||
SetGroupID(groupID).
|
||||
SetPriority(priority).
|
||||
@@ -1767,6 +1784,11 @@ func (r *accountRepository) AddToGroup(ctx context.Context, accountID, groupID i
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tx != nil {
|
||||
if err := tx.Commit(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
payload := buildSchedulerGroupPayload([]int64{groupID})
|
||||
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventAccountGroupsChanged, &accountID, nil, payload); err != nil {
|
||||
logger.LegacyPrintf("repository.account", "[SchedulerOutbox] enqueue add to group failed: account=%d group=%d err=%v", accountID, groupID, err)
|
||||
@@ -1827,6 +1849,9 @@ func (r *accountRepository) BindGroups(ctx context.Context, accountID int64, gro
|
||||
// 已处于外部事务中(ErrTxStarted),复用当前 client
|
||||
txClient = r.client
|
||||
}
|
||||
if err := lockLiveGroups(ctx, txClient, groupIDs); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
if _, err := txClient.AccountGroup.Delete().Where(dbaccountgroup.AccountIDEQ(accountID)).Exec(ctx); err != nil {
|
||||
return err
|
||||
|
||||
@@ -214,7 +214,7 @@ func (r *apiKeyRepository) GetByKeyForAuth(ctx context.Context, key string) (*se
|
||||
group.FieldFreeOpenaiFast,
|
||||
group.FieldDefaultMappedModel,
|
||||
group.FieldMessagesDispatchModelConfig,
|
||||
group.FieldModelsListConfig,
|
||||
group.FieldModelAllowlist,
|
||||
group.FieldCodexModelsManifestConfig,
|
||||
group.FieldRpmLimit,
|
||||
group.FieldMaxReasoningEffort,
|
||||
@@ -1017,7 +1017,7 @@ func groupEntityToService(g *dbent.Group) *service.Group {
|
||||
RequirePrivacySet: g.RequirePrivacySet,
|
||||
DefaultMappedModel: g.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: g.MessagesDispatchModelConfig,
|
||||
ModelsListConfig: g.ModelsListConfig,
|
||||
ModelAllowlist: service.GroupModelAllowlistFromDomain(g.ModelAllowlist),
|
||||
CodexModelsManifestConfig: g.CodexModelsManifestConfig,
|
||||
RPMLimit: g.RpmLimit,
|
||||
MaxReasoningEffort: g.MaxReasoningEffort,
|
||||
|
||||
+50
@@ -0,0 +1,50 @@
|
||||
//go:build integration
|
||||
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"fmt"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestGetByKeyForAuthCarriesGroupModelAllowlist(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
suffix := time.Now().UnixNano()
|
||||
group := mustCreateGroup(t, integrationEntClient, &service.Group{
|
||||
Name: fmt.Sprintf("model-allowlist-proj-group-%d", suffix), Platform: service.PlatformOpenAI,
|
||||
RateMultiplier: 1,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{"gpt-5.4", "gpt-5.5-*"},
|
||||
},
|
||||
})
|
||||
user := mustCreateUser(t, integrationEntClient, &service.User{
|
||||
Email: fmt.Sprintf("model-allowlist-proj-%d@example.com", suffix), Concurrency: 5,
|
||||
})
|
||||
groupID := group.ID
|
||||
keyValue := fmt.Sprintf("sk-model-allowlist-proj-%d", suffix)
|
||||
apiKeyRepo := NewAPIKeyRepository(integrationEntClient, integrationDB)
|
||||
key := &service.APIKey{UserID: user.ID, GroupID: &groupID, Key: keyValue, Name: "model-allowlist-proj", Status: service.StatusActive}
|
||||
require.NoError(t, apiKeyRepo.Create(ctx, key))
|
||||
t.Cleanup(func() {
|
||||
_, err := integrationDB.ExecContext(ctx, "DELETE FROM auth_cache_invalidation_outbox WHERE cache_key = encode(sha256(convert_to($1, 'UTF8')), 'hex')", keyValue)
|
||||
require.NoError(t, err)
|
||||
_, err = integrationDB.ExecContext(ctx, "DELETE FROM api_keys WHERE id = $1", key.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = integrationDB.ExecContext(ctx, "DELETE FROM users WHERE id = $1", user.ID)
|
||||
require.NoError(t, err)
|
||||
_, err = integrationDB.ExecContext(ctx, "DELETE FROM groups WHERE id = $1", group.ID)
|
||||
require.NoError(t, err)
|
||||
})
|
||||
|
||||
got, err := apiKeyRepo.GetByKeyForAuth(ctx, keyValue)
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, got.Group)
|
||||
require.True(t, got.Group.ModelAllowlist.Enabled, "model_allowlist 必须进入认证投影(投影漏列会让分组级准入静默失效)")
|
||||
require.Equal(t, []string{"gpt-5.4", "gpt-5.5-*"}, got.Group.ModelAllowlist.Models)
|
||||
}
|
||||
@@ -2,9 +2,14 @@ package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"database/sql"
|
||||
"database/sql/driver"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"os/exec"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
@@ -12,16 +17,39 @@ import (
|
||||
|
||||
// PgDumper implements service.DBDumper using pg_dump/psql
|
||||
type PgDumper struct {
|
||||
cfg *config.DatabaseConfig
|
||||
cfg *config.DatabaseConfig
|
||||
db *sql.DB
|
||||
commandContext func(context.Context, string, ...string) *exec.Cmd
|
||||
}
|
||||
|
||||
// NewPgDumper creates a new PgDumper
|
||||
func NewPgDumper(cfg *config.Config) service.DBDumper {
|
||||
return &PgDumper{cfg: &cfg.Database}
|
||||
func NewPgDumper(cfg *config.Config, db *sql.DB) service.DBDumper {
|
||||
return &PgDumper{
|
||||
cfg: &cfg.Database,
|
||||
db: db,
|
||||
commandContext: exec.CommandContext,
|
||||
}
|
||||
}
|
||||
|
||||
// Dump executes pg_dump and returns a streaming reader of the output
|
||||
func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||||
if d.db == nil {
|
||||
return nil, errors.New("acquire backup migration lock: nil sql db")
|
||||
}
|
||||
lockConn, err := d.db.Conn(ctx)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("acquire backup migration lock connection: %w", err)
|
||||
}
|
||||
if err := pgAdvisoryLock(ctx, lockConn); err != nil {
|
||||
// Lock acquisition can fail after PostgreSQL accepted the request. Discard
|
||||
// the session because ownership is ambiguous and advisory locks are session-scoped.
|
||||
discardSQLConnection(lockConn)
|
||||
return nil, fmt.Errorf("acquire backup migration lock: %w", err)
|
||||
}
|
||||
releaseLock := func() error {
|
||||
return releaseBackupMigrationLock(lockConn)
|
||||
}
|
||||
|
||||
args := []string{
|
||||
"-h", d.cfg.Host,
|
||||
"-p", fmt.Sprintf("%d", d.cfg.Port),
|
||||
@@ -33,7 +61,11 @@ func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||||
"--if-exists",
|
||||
}
|
||||
|
||||
cmd := exec.CommandContext(ctx, "pg_dump", args...)
|
||||
commandContext := d.commandContext
|
||||
if commandContext == nil {
|
||||
commandContext = exec.CommandContext
|
||||
}
|
||||
cmd := commandContext(ctx, "pg_dump", args...)
|
||||
if d.cfg.Password != "" {
|
||||
cmd.Env = append(cmd.Environ(), "PGPASSWORD="+d.cfg.Password)
|
||||
}
|
||||
@@ -43,15 +75,34 @@ func (d *PgDumper) Dump(ctx context.Context) (io.ReadCloser, error) {
|
||||
|
||||
stdout, err := cmd.StdoutPipe()
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("create stdout pipe: %w", err)
|
||||
return nil, errors.Join(fmt.Errorf("create stdout pipe: %w", err), releaseLock())
|
||||
}
|
||||
|
||||
if err := cmd.Start(); err != nil {
|
||||
return nil, fmt.Errorf("start pg_dump: %w", err)
|
||||
_ = stdout.Close()
|
||||
return nil, errors.Join(fmt.Errorf("start pg_dump: %w", err), releaseLock())
|
||||
}
|
||||
|
||||
// 返回一个 ReadCloser:读 stdout,关闭时等待进程退出
|
||||
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd}, nil
|
||||
return &cmdReadCloser{ReadCloser: stdout, cmd: cmd, release: releaseLock}, nil
|
||||
}
|
||||
|
||||
func releaseBackupMigrationLock(conn *sql.Conn) error {
|
||||
unlockCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := pgAdvisoryUnlock(unlockCtx, conn); err != nil {
|
||||
// A failed unlock must not return a possibly lock-owning session to the pool.
|
||||
discardSQLConnection(conn)
|
||||
return fmt.Errorf("release backup migration lock: %w", err)
|
||||
}
|
||||
if err := conn.Close(); err != nil {
|
||||
return fmt.Errorf("close backup migration lock connection: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func discardSQLConnection(conn *sql.Conn) {
|
||||
_ = conn.Raw(func(any) error { return driver.ErrBadConn })
|
||||
_ = conn.Close()
|
||||
}
|
||||
|
||||
// Restore executes psql to restore from a streaming reader
|
||||
@@ -84,15 +135,27 @@ func (d *PgDumper) Restore(ctx context.Context, data io.Reader) error {
|
||||
// cmdReadCloser wraps a command stdout pipe and waits for the process on Close
|
||||
type cmdReadCloser struct {
|
||||
io.ReadCloser
|
||||
cmd *exec.Cmd
|
||||
cmd *exec.Cmd
|
||||
release func() error
|
||||
closeOnce sync.Once
|
||||
closeErr error
|
||||
}
|
||||
|
||||
func (c *cmdReadCloser) Close() error {
|
||||
// Close the pipe first
|
||||
_ = c.ReadCloser.Close()
|
||||
// Wait for the process to exit
|
||||
if err := c.cmd.Wait(); err != nil {
|
||||
return fmt.Errorf("pg_dump exited with error: %w", err)
|
||||
}
|
||||
return nil
|
||||
c.closeOnce.Do(func() {
|
||||
var closeErrs []error
|
||||
if err := c.ReadCloser.Close(); err != nil {
|
||||
closeErrs = append(closeErrs, fmt.Errorf("close pg_dump stdout: %w", err))
|
||||
}
|
||||
if err := c.cmd.Wait(); err != nil {
|
||||
closeErrs = append(closeErrs, fmt.Errorf("pg_dump exited with error: %w", err))
|
||||
}
|
||||
if c.release != nil {
|
||||
if err := c.release(); err != nil {
|
||||
closeErrs = append(closeErrs, err)
|
||||
}
|
||||
}
|
||||
c.closeErr = errors.Join(closeErrs...)
|
||||
})
|
||||
return c.closeErr
|
||||
}
|
||||
|
||||
@@ -0,0 +1,157 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"os/exec"
|
||||
"regexp"
|
||||
"testing"
|
||||
|
||||
"github.com/DATA-DOG/go-sqlmock"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newTestPgDumper(t *testing.T, commandContext func(context.Context, string, ...string) *exec.Cmd) (*PgDumper, sqlmock.Sqlmock) {
|
||||
t.Helper()
|
||||
db, mock, err := sqlmock.New()
|
||||
require.NoError(t, err)
|
||||
t.Cleanup(func() { _ = db.Close() })
|
||||
return &PgDumper{
|
||||
cfg: &config.DatabaseConfig{
|
||||
Host: "db.example.test",
|
||||
Port: 5432,
|
||||
User: "sub2api",
|
||||
Password: "secret",
|
||||
DBName: "sub2api",
|
||||
SSLMode: "require",
|
||||
},
|
||||
db: db,
|
||||
commandContext: commandContext,
|
||||
}, mock
|
||||
}
|
||||
|
||||
func expectBackupMigrationLock(mock sqlmock.Sqlmock) {
|
||||
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnRows(sqlmock.NewRows([]string{"pg_try_advisory_lock"}).AddRow(true))
|
||||
}
|
||||
|
||||
func expectBackupMigrationUnlock(mock sqlmock.Sqlmock) {
|
||||
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnResult(sqlmock.NewResult(0, 1))
|
||||
}
|
||||
|
||||
func TestPgDumperHoldsMigrationLockThroughReaderClose(t *testing.T) {
|
||||
var mock sqlmock.Sqlmock
|
||||
commandCreated := false
|
||||
dumper, createdMock := newTestPgDumper(t, func(ctx context.Context, name string, args ...string) *exec.Cmd {
|
||||
commandCreated = true
|
||||
require.Equal(t, "pg_dump", name)
|
||||
require.Contains(t, args, "--clean")
|
||||
require.NoError(t, mock.ExpectationsWereMet(), "migration lock must be acquired before pg_dump is created")
|
||||
return exec.CommandContext(ctx, "sh", "-c", "printf backup-data")
|
||||
})
|
||||
mock = createdMock
|
||||
expectBackupMigrationLock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.NoError(t, err)
|
||||
require.True(t, commandCreated)
|
||||
data, err := io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "backup-data", string(data))
|
||||
|
||||
expectBackupMigrationUnlock(mock)
|
||||
require.Error(t, mock.ExpectationsWereMet(), "migration lock was released before the reader closed")
|
||||
require.NoError(t, reader.Close())
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
require.NoError(t, reader.Close(), "reader close must be idempotent")
|
||||
}
|
||||
|
||||
func TestPgDumperReleasesMigrationLockWhenStdoutPipeSetupFails(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
cmd := exec.CommandContext(ctx, "sh", "-c", "true")
|
||||
cmd.Stdout = io.Discard
|
||||
return cmd
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
expectBackupMigrationUnlock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "create stdout pipe")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperReleasesMigrationLockWhenProcessStartFails(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
return exec.CommandContext(ctx, "/path/that/does/not/exist/pg_dump")
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
expectBackupMigrationUnlock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "start pg_dump")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperReleasesMigrationLockWhenProcessFails(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
return exec.CommandContext(ctx, "sh", "-c", "printf partial-backup; exit 7")
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.NoError(t, err)
|
||||
data, err := io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "partial-backup", string(data))
|
||||
expectBackupMigrationUnlock(mock)
|
||||
require.ErrorContains(t, reader.Close(), "pg_dump exited with error")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperReportsUnlockFailureAndDiscardsConnection(t *testing.T) {
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
return exec.CommandContext(ctx, "sh", "-c", "printf backup-data")
|
||||
})
|
||||
expectBackupMigrationLock(mock)
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.NoError(t, err)
|
||||
_, err = io.ReadAll(reader)
|
||||
require.NoError(t, err)
|
||||
mock.ExpectExec(regexp.QuoteMeta("SELECT pg_advisory_unlock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnError(errors.New("unlock unavailable"))
|
||||
require.ErrorContains(t, reader.Close(), "release backup migration lock")
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperDoesNotStartProcessWhenMigrationLockFails(t *testing.T) {
|
||||
commandCreated := false
|
||||
dumper, mock := newTestPgDumper(t, func(ctx context.Context, _ string, _ ...string) *exec.Cmd {
|
||||
commandCreated = true
|
||||
return exec.CommandContext(ctx, "sh", "-c", "true")
|
||||
})
|
||||
mock.ExpectQuery(regexp.QuoteMeta("SELECT pg_try_advisory_lock($1)")).
|
||||
WithArgs(migrationsAdvisoryLockID).
|
||||
WillReturnError(errors.New("database unavailable"))
|
||||
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "acquire backup migration lock")
|
||||
require.False(t, commandCreated)
|
||||
require.NoError(t, mock.ExpectationsWereMet())
|
||||
}
|
||||
|
||||
func TestPgDumperRejectsNilDatabase(t *testing.T) {
|
||||
dumper := &PgDumper{cfg: &config.DatabaseConfig{}}
|
||||
reader, err := dumper.Dump(context.Background())
|
||||
require.Nil(t, reader)
|
||||
require.ErrorContains(t, err, "nil sql db")
|
||||
}
|
||||
@@ -92,6 +92,7 @@ func mustCreateGroup(t *testing.T, client *dbent.Client, g *service.Group) *serv
|
||||
SetIsExclusive(g.IsExclusive).
|
||||
SetForceOpenaiFast(g.ForceOpenAIFast).
|
||||
SetFreeOpenaiFast(g.FreeOpenAIFast).
|
||||
SetModelAllowlist(service.DomainGroupModelAllowlist(g.ModelAllowlist)).
|
||||
SetCodexModelsManifestConfig(g.CodexModelsManifestConfig).
|
||||
SetProfitControlEnabled(g.ProfitControlEnabled).
|
||||
SetProfitMinMargin(g.ProfitMinMargin).
|
||||
|
||||
@@ -29,6 +29,39 @@ type groupRepository struct {
|
||||
sql sqlExecutor
|
||||
}
|
||||
|
||||
// lockLiveGroups makes account-group inserts participate in the same row-lock
|
||||
// protocol as guarded group deletion. FOR SHARE conflicts with the deleter's
|
||||
// FOR UPDATE lock, and READ COMMITTED rechecks deleted_at after any wait.
|
||||
func lockLiveGroups(ctx context.Context, exec sqlExecutor, groupIDs []int64) error {
|
||||
if len(groupIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
unique := make(map[int64]struct{}, len(groupIDs))
|
||||
for _, id := range groupIDs {
|
||||
unique[id] = struct{}{}
|
||||
}
|
||||
rows, err := exec.QueryContext(ctx, `/* account_group_live_group_lock */
|
||||
SELECT id FROM groups
|
||||
WHERE id = ANY($1) AND deleted_at IS NULL
|
||||
ORDER BY id
|
||||
FOR SHARE`, pq.Array(groupIDs))
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
defer func() { _ = rows.Close() }()
|
||||
locked := 0
|
||||
for rows.Next() {
|
||||
locked++
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return err
|
||||
}
|
||||
if locked != len(unique) {
|
||||
return service.ErrGroupNotFound
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func NewGroupRepository(client *dbent.Client, sqlDB *sql.DB) service.GroupRepository {
|
||||
return newGroupRepositoryWithSQL(client, sqlDB)
|
||||
}
|
||||
@@ -109,7 +142,7 @@ func createGroupRecord(ctx context.Context, client *dbent.Client, groupIn *servi
|
||||
SetRequirePrivacySet(groupIn.RequirePrivacySet).
|
||||
SetDefaultMappedModel(groupIn.DefaultMappedModel).
|
||||
SetMessagesDispatchModelConfig(groupIn.MessagesDispatchModelConfig).
|
||||
SetModelsListConfig(groupIn.ModelsListConfig).
|
||||
SetModelAllowlist(service.DomainGroupModelAllowlist(groupIn.ModelAllowlist)).
|
||||
SetCodexModelsManifestConfig(groupIn.CodexModelsManifestConfig).
|
||||
SetRpmLimit(groupIn.RPMLimit).
|
||||
SetMaxReasoningEffort(groupIn.MaxReasoningEffort).
|
||||
@@ -289,7 +322,7 @@ func (r *groupRepository) Update(ctx context.Context, groupIn *service.Group) er
|
||||
SetRequirePrivacySet(groupIn.RequirePrivacySet).
|
||||
SetDefaultMappedModel(groupIn.DefaultMappedModel).
|
||||
SetMessagesDispatchModelConfig(groupIn.MessagesDispatchModelConfig).
|
||||
SetModelsListConfig(groupIn.ModelsListConfig).
|
||||
SetModelAllowlist(service.DomainGroupModelAllowlist(groupIn.ModelAllowlist)).
|
||||
SetCodexModelsManifestConfig(groupIn.CodexModelsManifestConfig).
|
||||
SetRpmLimit(groupIn.RPMLimit).
|
||||
SetMaxReasoningEffort(groupIn.MaxReasoningEffort).
|
||||
@@ -426,6 +459,15 @@ func (r *groupRepository) List(ctx context.Context, params pagination.Pagination
|
||||
|
||||
func (r *groupRepository) ListWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
|
||||
q := r.client.Group.Query()
|
||||
return r.listWithFiltersQuery(ctx, q, params, platform, status, search, isExclusive)
|
||||
}
|
||||
|
||||
func (r *groupRepository) ListBindableWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
|
||||
q := r.client.Group.Query().Where(group.PlatformNEQ(service.PlatformComposite))
|
||||
return r.listWithFiltersQuery(ctx, q, params, platform, status, search, isExclusive)
|
||||
}
|
||||
|
||||
func (r *groupRepository) listWithFiltersQuery(ctx context.Context, q *dbent.GroupQuery, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]service.Group, *pagination.PaginationResult, error) {
|
||||
|
||||
if platform != "" {
|
||||
q = q.Where(group.PlatformEQ(platform))
|
||||
@@ -807,12 +849,14 @@ func (r *groupRepository) DeleteAccountGroupsByGroupID(ctx context.Context, grou
|
||||
}
|
||||
|
||||
func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64, error) {
|
||||
g, err := r.client.Group.Query().Where(group.IDEQ(id)).Only(ctx)
|
||||
if err != nil {
|
||||
return nil, translatePersistenceError(err, service.ErrGroupNotFound, nil)
|
||||
}
|
||||
groupSvc := groupEntityToService(g)
|
||||
return r.deleteCascade(ctx, id, false)
|
||||
}
|
||||
|
||||
func (r *groupRepository) DeleteCascadeIfEmpty(ctx context.Context, id int64) ([]int64, error) {
|
||||
return r.deleteCascade(ctx, id, true)
|
||||
}
|
||||
|
||||
func (r *groupRepository) deleteCascade(ctx context.Context, id int64, requireEmpty bool) ([]int64, error) {
|
||||
// 使用 ent 事务统一包裹:避免手工基于 *sql.Tx 构造 ent client 带来的驱动断言问题,
|
||||
// 同时保证级联删除的原子性。
|
||||
tx, err := r.client.Tx(ctx)
|
||||
@@ -830,13 +874,14 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64,
|
||||
|
||||
// Lock the group row to avoid concurrent writes while we cascade.
|
||||
// 这里使用 exec.QueryContext 手动扫描,确保同一事务内加锁并能区分"未找到"与其他错误。
|
||||
rows, err := exec.QueryContext(ctx, "SELECT id FROM groups WHERE id = $1 AND deleted_at IS NULL FOR UPDATE", id)
|
||||
rows, err := exec.QueryContext(ctx, "SELECT id, subscription_type FROM groups WHERE id = $1 AND deleted_at IS NULL FOR UPDATE", id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var lockedID int64
|
||||
var subscriptionType string
|
||||
if rows.Next() {
|
||||
if err := rows.Scan(&lockedID); err != nil {
|
||||
if err := rows.Scan(&lockedID, &subscriptionType); err != nil {
|
||||
_ = rows.Close()
|
||||
return nil, err
|
||||
}
|
||||
@@ -850,9 +895,22 @@ func (r *groupRepository) DeleteCascade(ctx context.Context, id int64) ([]int64,
|
||||
if lockedID == 0 {
|
||||
return nil, service.ErrGroupNotFound
|
||||
}
|
||||
if requireEmpty {
|
||||
var hasAccount bool
|
||||
if err := scanSingleRow(ctx, exec, `SELECT EXISTS (
|
||||
SELECT 1 FROM account_groups ag
|
||||
JOIN accounts a ON a.id = ag.account_id
|
||||
WHERE ag.group_id = $1 AND a.deleted_at IS NULL
|
||||
)`, []any{id}, &hasAccount); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if hasAccount {
|
||||
return nil, service.ErrGroupNotEmpty
|
||||
}
|
||||
}
|
||||
|
||||
var affectedUserIDs []int64
|
||||
if groupSvc.IsSubscriptionType() {
|
||||
if subscriptionType == service.SubscriptionTypeSubscription {
|
||||
// 只查询未软删除的订阅,避免通知已取消订阅的用户
|
||||
rows, err := exec.QueryContext(ctx, "SELECT user_id FROM user_subscriptions WHERE group_id = $1 AND deleted_at IS NULL", id)
|
||||
if err != nil {
|
||||
@@ -1020,8 +1078,21 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64
|
||||
return nil
|
||||
}
|
||||
|
||||
tx, err := r.client.Tx(ctx)
|
||||
if err != nil && !errors.Is(err, dbent.ErrTxStarted) {
|
||||
return err
|
||||
}
|
||||
exec := sqlExecutor(r.client)
|
||||
if tx != nil {
|
||||
defer func() { _ = tx.Rollback() }()
|
||||
exec = tx.Client()
|
||||
}
|
||||
if err := lockLiveGroups(ctx, exec, []int64{groupID}); err != nil {
|
||||
return err
|
||||
}
|
||||
|
||||
// 使用 INSERT ... ON CONFLICT DO NOTHING 忽略已存在的绑定
|
||||
_, err := r.sql.ExecContext(
|
||||
_, err = exec.ExecContext(
|
||||
ctx,
|
||||
`INSERT INTO account_groups (account_id, group_id, priority, created_at)
|
||||
SELECT unnest($1::bigint[]), $2, 50, NOW()
|
||||
@@ -1032,6 +1103,11 @@ func (r *groupRepository) BindAccountsToGroup(ctx context.Context, groupID int64
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if tx != nil {
|
||||
if err := tx.Commit(); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
// 发送调度器事件
|
||||
if err := enqueueSchedulerOutbox(ctx, r.sql, service.SchedulerOutboxEventGroupChanged, nil, &groupID, nil); err != nil {
|
||||
|
||||
@@ -8,10 +8,12 @@ import (
|
||||
"errors"
|
||||
"strings"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
"github.com/stretchr/testify/suite"
|
||||
)
|
||||
|
||||
@@ -245,6 +247,103 @@ func (s *GroupRepoSuite) TestDelete() {
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotFound)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyRejectsGroupWithNonDeletedAccount() {
|
||||
group := &service.Group{Name: "guarded-non-empty", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
var accountID int64
|
||||
s.Require().NoError(scanSingleRow(s.ctx, s.tx,
|
||||
"INSERT INTO accounts (name, platform, type) VALUES ($1, $2, $3) RETURNING id",
|
||||
[]any{"guarded-account", service.PlatformAnthropic, service.AccountTypeOAuth}, &accountID))
|
||||
_, err := s.tx.ExecContext(s.ctx, "INSERT INTO account_groups (account_id, group_id, priority, created_at) VALUES ($1, $2, 1, NOW())", accountID, group.ID)
|
||||
s.Require().NoError(err)
|
||||
|
||||
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotEmpty)
|
||||
_, err = s.repo.GetByID(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
var bindings int
|
||||
s.Require().NoError(scanSingleRow(s.ctx, s.tx, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{group.ID}, &bindings))
|
||||
s.Require().Equal(1, bindings)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyDeletesEmptyGroup() {
|
||||
group := &service.Group{Name: "guarded-empty", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
|
||||
_, err := s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
_, err = s.repo.GetByID(s.ctx, group.ID)
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotFound)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestDeleteCascadeIfEmptyIgnoresBindingsToSoftDeletedAccounts() {
|
||||
group := &service.Group{Name: "guarded-soft-deleted-account", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard}
|
||||
s.Require().NoError(s.repo.Create(s.ctx, group))
|
||||
var accountID int64
|
||||
s.Require().NoError(scanSingleRow(s.ctx, s.tx,
|
||||
"INSERT INTO accounts (name, platform, type, deleted_at) VALUES ($1, $2, $3, NOW()) RETURNING id",
|
||||
[]any{"guarded-deleted-account", service.PlatformAnthropic, service.AccountTypeOAuth}, &accountID))
|
||||
_, err := s.tx.ExecContext(s.ctx, "INSERT INTO account_groups (account_id, group_id, priority, created_at) VALUES ($1, $2, 1, NOW())", accountID, group.ID)
|
||||
s.Require().NoError(err)
|
||||
|
||||
_, err = s.repo.DeleteCascadeIfEmpty(s.ctx, group.ID)
|
||||
s.Require().NoError(err)
|
||||
_, err = s.repo.GetByID(s.ctx, group.ID)
|
||||
s.Require().ErrorIs(err, service.ErrGroupNotFound)
|
||||
}
|
||||
|
||||
func TestBindAccountsToGroupWaitingBehindGuardedDeleteCannotCommit(t *testing.T) {
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
|
||||
defer cancel()
|
||||
|
||||
var groupID, accountID int64
|
||||
require.NoError(t, scanSingleRow(ctx, integrationDB,
|
||||
"INSERT INTO groups (name, platform, rate_multiplier, status, subscription_type) VALUES ($1, $2, 1, $3, $4) RETURNING id",
|
||||
[]any{"guarded-delete-bind-race", service.PlatformAnthropic, service.StatusActive, service.SubscriptionTypeStandard}, &groupID))
|
||||
require.NoError(t, scanSingleRow(ctx, integrationDB,
|
||||
"INSERT INTO accounts (name, platform, type) VALUES ($1, $2, $3) RETURNING id",
|
||||
[]any{"guarded-delete-bind-race-account", service.PlatformAnthropic, service.AccountTypeOAuth}, &accountID))
|
||||
t.Cleanup(func() {
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM account_groups WHERE group_id = $1 OR account_id = $2", groupID, accountID)
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM accounts WHERE id = $1", accountID)
|
||||
_, _ = integrationDB.ExecContext(context.Background(), "DELETE FROM groups WHERE id = $1", groupID)
|
||||
})
|
||||
|
||||
deleteTx, err := integrationEntClient.Tx(ctx)
|
||||
require.NoError(t, err)
|
||||
defer func() { _ = deleteTx.Rollback() }()
|
||||
rows, err := deleteTx.Client().QueryContext(ctx, "SELECT id FROM groups WHERE id = $1 FOR UPDATE", groupID)
|
||||
require.NoError(t, err)
|
||||
require.True(t, rows.Next())
|
||||
require.NoError(t, rows.Close())
|
||||
|
||||
bindDone := make(chan error, 1)
|
||||
go func() {
|
||||
bindDone <- newGroupRepositoryWithSQL(integrationEntClient, integrationDB).
|
||||
BindAccountsToGroup(ctx, groupID, []int64{accountID})
|
||||
}()
|
||||
|
||||
require.Eventually(t, func() bool {
|
||||
var waiting bool
|
||||
err := scanSingleRow(ctx, integrationDB, `SELECT EXISTS (
|
||||
SELECT 1 FROM pg_stat_activity
|
||||
WHERE query LIKE '/* account_group_live_group_lock */%'
|
||||
AND wait_event_type = 'Lock'
|
||||
)`, nil, &waiting)
|
||||
return err == nil && waiting
|
||||
}, 5*time.Second, 10*time.Millisecond, "binder did not wait on the guarded deletion lock")
|
||||
|
||||
deleteRepo := newGroupRepositoryWithSQL(deleteTx.Client(), deleteTx)
|
||||
_, err = deleteRepo.DeleteCascadeIfEmpty(ctx, groupID)
|
||||
require.NoError(t, err)
|
||||
require.NoError(t, deleteTx.Commit())
|
||||
|
||||
require.ErrorIs(t, <-bindDone, service.ErrGroupNotFound)
|
||||
var bindings int
|
||||
require.NoError(t, scanSingleRow(ctx, integrationDB, "SELECT COUNT(*) FROM account_groups WHERE group_id = $1", []any{groupID}, &bindings))
|
||||
require.Zero(t, bindings)
|
||||
}
|
||||
|
||||
// --- List / ListWithFilters ---
|
||||
|
||||
func (s *GroupRepoSuite) TestList() {
|
||||
@@ -619,6 +718,22 @@ func (s *GroupRepoSuite) TestListActive() {
|
||||
s.Require().True(found, "active1 group should be in results")
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestListBindableWithFiltersAppliesFilterBeforePagination() {
|
||||
for _, g := range []*service.Group{
|
||||
{Name: "bindable-a", Platform: service.PlatformAnthropic, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard},
|
||||
{Name: "composite-hidden", Platform: service.PlatformComposite, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard},
|
||||
{Name: "bindable-b", Platform: service.PlatformOpenAI, RateMultiplier: 1, Status: service.StatusActive, SubscriptionType: service.SubscriptionTypeStandard},
|
||||
} {
|
||||
s.Require().NoError(s.repo.Create(s.ctx, g))
|
||||
}
|
||||
|
||||
groups, page, err := s.repo.ListBindableWithFilters(s.ctx, pagination.PaginationParams{Page: 2, PageSize: 1, SortBy: "name", SortOrder: "asc"}, "", service.StatusActive, "bindable-", nil)
|
||||
s.Require().NoError(err)
|
||||
s.Require().Equal(int64(2), page.Total)
|
||||
s.Require().Len(groups, 1)
|
||||
s.Require().NotEqual(service.PlatformComposite, groups[0].Platform)
|
||||
}
|
||||
|
||||
func (s *GroupRepoSuite) TestListActiveByPlatform() {
|
||||
s.Require().NoError(s.repo.Create(s.ctx, &service.Group{
|
||||
Name: "g1",
|
||||
|
||||
@@ -10,11 +10,22 @@ import (
|
||||
)
|
||||
|
||||
const (
|
||||
redeemRateLimitKeyPrefix = "redeem:ratelimit:"
|
||||
redeemRateLimitKeyPrefix = "redeem:ratelimit:v2:"
|
||||
redeemLockKeyPrefix = "redeem:lock:"
|
||||
redeemRateLimitDuration = 24 * time.Hour
|
||||
redeemRateLimitWindow = 10 * time.Minute
|
||||
)
|
||||
|
||||
var incrementRedeemAttemptScript = redis.NewScript(`
|
||||
local current = redis.call('INCR', KEYS[1])
|
||||
local ttl = redis.call('PTTL', KEYS[1])
|
||||
|
||||
if current == 1 or ttl == -1 then
|
||||
redis.call('PEXPIRE', KEYS[1], ARGV[1])
|
||||
end
|
||||
|
||||
return current
|
||||
`)
|
||||
|
||||
// redeemRateLimitKey generates the Redis key for redeem attempt rate limiting.
|
||||
func redeemRateLimitKey(userID int64) string {
|
||||
return fmt.Sprintf("%s%d", redeemRateLimitKeyPrefix, userID)
|
||||
@@ -44,11 +55,12 @@ func (c *redeemCache) GetRedeemAttemptCount(ctx context.Context, userID int64) (
|
||||
|
||||
func (c *redeemCache) IncrementRedeemAttemptCount(ctx context.Context, userID int64) error {
|
||||
key := redeemRateLimitKey(userID)
|
||||
pipe := c.rdb.Pipeline()
|
||||
pipe.Incr(ctx, key)
|
||||
pipe.Expire(ctx, key, redeemRateLimitDuration)
|
||||
_, err := pipe.Exec(ctx)
|
||||
return err
|
||||
return incrementRedeemAttemptScript.Run(
|
||||
ctx,
|
||||
c.rdb,
|
||||
[]string{key},
|
||||
redeemRateLimitWindow.Milliseconds(),
|
||||
).Err()
|
||||
}
|
||||
|
||||
func (c *redeemCache) AcquireRedeemLock(ctx context.Context, code string, ttl time.Duration) (bool, error) {
|
||||
|
||||
@@ -39,7 +39,35 @@ func (s *RedeemCacheSuite) TestIncrementAndGetRedeemAttemptCount() {
|
||||
|
||||
ttl, err := s.rdb.TTL(s.ctx, key).Result()
|
||||
require.NoError(s.T(), err, "TTL")
|
||||
s.AssertTTLWithin(ttl, 1*time.Second, redeemRateLimitDuration)
|
||||
s.AssertTTLWithin(ttl, redeemRateLimitWindow-time.Second, redeemRateLimitWindow)
|
||||
}
|
||||
|
||||
func (s *RedeemCacheSuite) TestIncrementDoesNotRefreshFixedWindow() {
|
||||
userID := int64(3)
|
||||
key := redeemRateLimitKey(userID)
|
||||
shortTTL := 5 * time.Minute
|
||||
require.NoError(s.T(), s.cache.IncrementRedeemAttemptCount(s.ctx, userID))
|
||||
require.NoError(s.T(), s.rdb.PExpire(s.ctx, key, shortTTL).Err())
|
||||
require.NoError(s.T(), s.cache.IncrementRedeemAttemptCount(s.ctx, userID))
|
||||
ttl, err := s.rdb.PTTL(s.ctx, key).Result()
|
||||
require.NoError(s.T(), err)
|
||||
s.AssertTTLWithin(ttl, shortTTL-time.Second, shortTTL)
|
||||
}
|
||||
|
||||
func (s *RedeemCacheSuite) TestIncrementRepairsMissingTTL() {
|
||||
userID := int64(4)
|
||||
key := redeemRateLimitKey(userID)
|
||||
require.NoError(s.T(), s.rdb.Set(s.ctx, key, 5, 0).Err())
|
||||
ttl, err := s.rdb.PTTL(s.ctx, key).Result()
|
||||
require.NoError(s.T(), err)
|
||||
require.Equal(s.T(), time.Duration(-1), ttl)
|
||||
require.NoError(s.T(), s.cache.IncrementRedeemAttemptCount(s.ctx, userID))
|
||||
count, err := s.cache.GetRedeemAttemptCount(s.ctx, userID)
|
||||
require.NoError(s.T(), err)
|
||||
require.Equal(s.T(), 6, count)
|
||||
ttl, err = s.rdb.PTTL(s.ctx, key).Result()
|
||||
require.NoError(s.T(), err)
|
||||
s.AssertTTLWithin(ttl, redeemRateLimitWindow-time.Second, redeemRateLimitWindow)
|
||||
}
|
||||
|
||||
func (s *RedeemCacheSuite) TestMultipleIncrements() {
|
||||
|
||||
@@ -18,22 +18,22 @@ func TestRedeemRateLimitKey(t *testing.T) {
|
||||
{
|
||||
name: "normal_user_id",
|
||||
userID: 123,
|
||||
expected: "redeem:ratelimit:123",
|
||||
expected: "redeem:ratelimit:v2:123",
|
||||
},
|
||||
{
|
||||
name: "zero_user_id",
|
||||
userID: 0,
|
||||
expected: "redeem:ratelimit:0",
|
||||
expected: "redeem:ratelimit:v2:0",
|
||||
},
|
||||
{
|
||||
name: "negative_user_id",
|
||||
userID: -1,
|
||||
expected: "redeem:ratelimit:-1",
|
||||
expected: "redeem:ratelimit:v2:-1",
|
||||
},
|
||||
{
|
||||
name: "max_int64",
|
||||
userID: math.MaxInt64,
|
||||
expected: "redeem:ratelimit:9223372036854775807",
|
||||
expected: "redeem:ratelimit:v2:9223372036854775807",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
@@ -1476,7 +1476,7 @@ func newContractDeps(t *testing.T) *contractDeps {
|
||||
settingRepo := newStubSettingRepo()
|
||||
settingService := service.NewSettingService(settingRepo, cfg)
|
||||
|
||||
adminService := service.NewAdminService(userRepo, groupRepo, &accountRepo, proxyRepo, apiKeyRepo, redeemRepo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
adminService := service.NewAdminService(nil, userRepo, groupRepo, &accountRepo, proxyRepo, apiKeyRepo, redeemRepo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
authHandler := handler.NewAuthHandler(cfg, nil, userService, settingService, nil, redeemService, nil, nil)
|
||||
apiKeyHandler := handler.NewAPIKeyHandler(apiKeyService)
|
||||
usageHandler := handler.NewUsageHandler(usageService, apiKeyService, nil, nil)
|
||||
@@ -1789,6 +1789,10 @@ func (stubGroupRepo) DeleteCascade(ctx context.Context, id int64) ([]int64, erro
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (stubGroupRepo) DeleteCascadeIfEmpty(ctx context.Context, id int64) ([]int64, error) {
|
||||
return nil, errors.New("not implemented")
|
||||
}
|
||||
|
||||
func (stubGroupRepo) List(ctx context.Context, params pagination.PaginationParams) ([]service.Group, *pagination.PaginationResult, error) {
|
||||
return nil, nil, errors.New("not implemented")
|
||||
}
|
||||
@@ -2941,7 +2945,7 @@ var (
|
||||
_ service.UserRepository = (*stubUserRepo)(nil)
|
||||
_ service.APIKeyRepository = (*stubApiKeyRepo)(nil)
|
||||
_ service.APIKeyCache = (*stubApiKeyCache)(nil)
|
||||
_ service.GroupRepository = (*stubGroupRepo)(nil)
|
||||
_ service.AdminGroupRepository = (*stubGroupRepo)(nil)
|
||||
_ service.UserSubscriptionRepository = (*stubUserSubscriptionRepo)(nil)
|
||||
_ service.UsageLogRepository = (*stubUsageLogRepo)(nil)
|
||||
_ service.SettingRepository = (*stubSettingRepo)(nil)
|
||||
|
||||
@@ -0,0 +1,165 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
// GroupModelAllowlist 是分组级模型白名单准入中间件。
|
||||
//
|
||||
// 挂载位置:每条网关链的 apiKeyAuth 之后、compositeTarget 之前——
|
||||
// 保证校验发生在合成路由改写与调度之前,且只看客户端书写的公开模型名。
|
||||
//
|
||||
// 行为:
|
||||
// - 快速路径:未绑定分组或白名单未开启时直接放行,不读请求体。
|
||||
// - Responses WebSocket 入口跳过(首帧与后续 turn 由 ResponsesWebSocket 逐帧
|
||||
// 校验);Grok Realtime 的升级请求模型固定在查询参数里,仍走中间件校验,
|
||||
// 其他路由伪造 Upgrade 头不得绕过校验。
|
||||
// - GET/DELETE:从路由参数(model / modelAction)或查询参数(/realtime 的
|
||||
// model)提取模型;提取不到放行,由 handler 决定是否报「model is required」。
|
||||
// - 带请求体的方法:读体(PrereadBody 回填,下游零拷贝)、提取 JSON
|
||||
// `model`/`session.model` 或 multipart `model`/`session` 后回填请求体。
|
||||
// - 拒绝:按入口协议格式返回 404,并标记运维业务限流原因
|
||||
// local_model_configuration 与 ingress 拒绝原因 model_not_allowed。
|
||||
func GroupModelAllowlist() gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
apiKey, ok := GetAPIKeyFromContext(c)
|
||||
if !ok || apiKey == nil || apiKey.Group == nil || !apiKey.Group.ModelAllowlistEnabled() {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
allowlist := apiKey.Group.ModelAllowlist
|
||||
if c.Request == nil {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
if isResponsesWebSocketRoute(c) {
|
||||
// Responses WS 长连接由 ResponsesWebSocket 校验首帧与每个 response.create。
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
// models 收集该请求全部可被下游解析器绑定到的模型值(路径参数/查询参数
|
||||
// 单值;请求体候选集含重复键与大小写变体),逐一校验,任一未命中即拒绝。
|
||||
var models []string
|
||||
if model := groupModelAllowlistModelFromParams(c); model != "" {
|
||||
models = []string{model}
|
||||
} else {
|
||||
switch c.Request.Method {
|
||||
case http.MethodPost, http.MethodPut, http.MethodPatch:
|
||||
candidates, done := groupModelAllowlistModelsFromBody(c)
|
||||
if !done {
|
||||
// 请求体读取失败(如超限 413)已写出响应。
|
||||
return
|
||||
}
|
||||
models = candidates
|
||||
}
|
||||
if len(models) == 0 {
|
||||
// Grok Realtime 升级请求把模型固定在查询参数里。
|
||||
if model := strings.TrimSpace(c.Query("model")); model != "" {
|
||||
models = []string{model}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
blocked := ""
|
||||
for _, candidate := range models {
|
||||
if !allowlist.Allows(candidate) {
|
||||
blocked = candidate
|
||||
break
|
||||
}
|
||||
}
|
||||
if blocked == "" {
|
||||
c.Next()
|
||||
return
|
||||
}
|
||||
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalModelConfiguration)
|
||||
MarkIngressRejected(c, IngressRejectModelNotAllowed)
|
||||
groupModelAllowlistErrorWriter(c)(c, http.StatusNotFound, fmt.Sprintf("Model %q is not available for this group", blocked))
|
||||
c.Abort()
|
||||
}
|
||||
}
|
||||
|
||||
// isResponsesWebSocketRoute 判断当前请求是否命中 OpenAI Responses WebSocket
|
||||
// 入口。只有这三条路由的模型在升级后的帧里,交给 handler 逐帧校验;其他路由
|
||||
// 即使携带 WebSocket Upgrade 头也必须经过本中间件校验。
|
||||
func isResponsesWebSocketRoute(c *gin.Context) bool {
|
||||
if c.Request == nil || c.Request.Method != http.MethodGet {
|
||||
return false
|
||||
}
|
||||
switch c.FullPath() {
|
||||
case "/v1/responses", "/responses", "/backend-api/codex/responses":
|
||||
return true
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// groupModelAllowlistModelsFromBody 读取请求体并提取客户端模型名,随后把请求体
|
||||
// 回填(PrereadBody),保证后续 handler 零拷贝重读。读取失败按现有合成中间件
|
||||
// 的方式返回 400/413(返回 false 表示已写出响应并 Abort)。
|
||||
//
|
||||
// 下游同时存在 gjson(首个、大小写敏感)、encoding/json 绑定(末值、大小写
|
||||
// 不敏感)与 multipart 表单(首/末字段)三类解析器,这里返回「任一解析器可能
|
||||
// 绑定到的全部模型值」,调用方必须逐一校验,任一未命中即拒绝。
|
||||
func groupModelAllowlistModelsFromBody(c *gin.Context) ([]string, bool) {
|
||||
body, err := httputil.ReadRequestBodyWithPrealloc(c.Request)
|
||||
if err != nil {
|
||||
status := http.StatusBadRequest
|
||||
message := "Failed to read request body"
|
||||
var maxErr *http.MaxBytesError
|
||||
if errors.As(err, &maxErr) {
|
||||
status = http.StatusRequestEntityTooLarge
|
||||
message = "Request body is too large"
|
||||
}
|
||||
c.JSON(status, gin.H{"error": gin.H{"type": "invalid_request_error", "message": message}})
|
||||
c.Abort()
|
||||
return nil, false
|
||||
}
|
||||
requestmodel.ResetRequestBody(c.Request, body)
|
||||
return requestmodel.FromBodyCandidates(c.FullPath(), c.GetHeader("Content-Type"), body), true
|
||||
}
|
||||
|
||||
// groupModelAllowlistModelFromParams 从路由参数提取模型名:Gemini 原生 URL 的
|
||||
// `model` / `modelAction`(去掉 `:action` 后缀)。提取不到返回空串,
|
||||
// 由调用方决定是否继续读请求体或查询参数。
|
||||
func groupModelAllowlistModelFromParams(c *gin.Context) string {
|
||||
if model := strings.TrimSpace(c.Param("model")); model != "" {
|
||||
return model
|
||||
}
|
||||
if modelAction := strings.TrimSpace(c.Param("modelAction")); modelAction != "" {
|
||||
modelAction = strings.TrimPrefix(modelAction, "/")
|
||||
if idx := strings.LastIndex(modelAction, ":"); idx >= 0 {
|
||||
return strings.TrimSpace(modelAction[:idx])
|
||||
}
|
||||
return modelAction
|
||||
}
|
||||
return ""
|
||||
}
|
||||
|
||||
// groupModelAllowlistErrorWriter 按入口协议选择错误格式:
|
||||
// Gemini 原生(/v1beta、/antigravity/v1beta)用 Google 格式;
|
||||
// Messages 入口(含根路径别名与 /antigravity/v1)用 Anthropic 格式;
|
||||
// 其余(OpenAI 兼容入口)用 OpenAI 格式。
|
||||
func groupModelAllowlistErrorWriter(c *gin.Context) GatewayErrorWriter {
|
||||
path := ""
|
||||
if c.Request != nil && c.Request.URL != nil {
|
||||
path = c.Request.URL.Path
|
||||
}
|
||||
switch {
|
||||
case strings.HasPrefix(path, "/v1beta") || strings.HasPrefix(path, "/antigravity/v1beta"):
|
||||
return GoogleErrorWriter
|
||||
case strings.Contains(path, "/messages"):
|
||||
return AnthropicErrorWriter
|
||||
default:
|
||||
return OpenAIErrorWriter
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,501 @@
|
||||
package middleware
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"io"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
)
|
||||
|
||||
func newGroupModelAllowlistTestRouter(apiKey *service.APIKey, pathPrefix string) (*gin.Engine, *[]string) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
var calls []string
|
||||
router.Use(func(c *gin.Context) {
|
||||
if apiKey != nil {
|
||||
c.Set(string(ContextKeyAPIKey), apiKey)
|
||||
}
|
||||
c.Next()
|
||||
})
|
||||
router.Use(GroupModelAllowlist())
|
||||
register := func(method, path string) {
|
||||
router.Handle(method, path, func(c *gin.Context) {
|
||||
calls = append(calls, "handler:"+path)
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
}
|
||||
register(http.MethodPost, pathPrefix+"/messages")
|
||||
register(http.MethodPost, pathPrefix+"/responses")
|
||||
register(http.MethodPost, pathPrefix+"/chat/completions")
|
||||
register(http.MethodPost, pathPrefix+"/embeddings")
|
||||
register(http.MethodGet, pathPrefix+"/models/:model")
|
||||
register(http.MethodPost, pathPrefix+"/models/*modelAction")
|
||||
register(http.MethodGet, pathPrefix+"/realtime")
|
||||
register(http.MethodPost, pathPrefix+"/images/edits")
|
||||
register(http.MethodPost, pathPrefix+"/live")
|
||||
return router, &calls
|
||||
}
|
||||
|
||||
func allowlistAPIKey(enabled bool, models ...string) *service.APIKey {
|
||||
return &service.APIKey{
|
||||
Group: &service.Group{
|
||||
Platform: service.PlatformAnthropic,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: enabled,
|
||||
Models: models,
|
||||
},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
func doJSON(t *testing.T, router *gin.Engine, method, path, body string) *httptest.ResponseRecorder {
|
||||
t.Helper()
|
||||
req := httptest.NewRequest(method, path, strings.NewReader(body))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
return w
|
||||
}
|
||||
|
||||
// readTrackingBody 记录请求体是否被读取,用于断言快速路径零读取。
|
||||
type readTrackingBody struct {
|
||||
io.Reader
|
||||
read bool
|
||||
}
|
||||
|
||||
func (b *readTrackingBody) Read(p []byte) (int, error) {
|
||||
b.read = true
|
||||
return b.Reader.Read(p)
|
||||
}
|
||||
|
||||
func (b *readTrackingBody) Close() error { return nil }
|
||||
|
||||
func TestGroupModelAllowlistDisabledDoesNotReadBody(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(false, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
body := &readTrackingBody{Reader: strings.NewReader(`{"model":"claude-opus-4.6"}`)}
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/messages", body)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if len(*calls) != 1 {
|
||||
t.Fatalf("expected handler to run once, got %v", *calls)
|
||||
}
|
||||
if body.read {
|
||||
t.Fatal("allowlist disabled: middleware must not read the request body")
|
||||
}
|
||||
}
|
||||
|
||||
// 白名单开启但请求不携带模型时同样不应读体之外产生副作用:读体是必要开销,
|
||||
// 这里验证读取后请求体被完整回填(handler 可零拷贝重读)。
|
||||
func TestGroupModelAllowlistEnabledModelFreeRequestRestoresBody(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
|
||||
c.Next()
|
||||
})
|
||||
router.Use(GroupModelAllowlist())
|
||||
router.POST("/v1/messages", func(c *gin.Context) {
|
||||
body, err := httputil.ReadRequestBodyWithPrealloc(c.Request)
|
||||
if err != nil {
|
||||
t.Errorf("handler reread body: %v", err)
|
||||
c.Status(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if string(body) != `{"input":"no model field"}` {
|
||||
t.Errorf("handler saw wrong body: %s", body)
|
||||
}
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"input":"no model field"}`)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// 413 传递:请求体超过 MaxBytesReader 上限时按现有合成中间件方式返回 413。
|
||||
func TestGroupModelAllowlistRequestBodyTooLargePassesThrough413(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
bigBody := `{"model":"claude-opus-4.6","padding":"` + strings.Repeat("x", 512) + `"}`
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(bigBody))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Body = http.MaxBytesReader(httptest.NewRecorder(), req.Body, 64)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusRequestEntityTooLarge {
|
||||
t.Fatalf("expected 413, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "Request body is too large") {
|
||||
t.Fatalf("expected too-large message, got %s", w.Body.String())
|
||||
}
|
||||
if len(*calls) != 0 {
|
||||
t.Fatalf("expected handler not to run, got %v", *calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistJSONBodyAllowed(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"claude-sonnet-4.5"}`)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if len(*calls) != 1 {
|
||||
t.Fatalf("expected handler to run once, got %v", *calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistJSONBodyDeniedOpenAIFormat(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"model":"gpt-5.4"}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), `"model_not_found"`) {
|
||||
t.Fatalf("expected OpenAI error format, got %s", w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), `gpt-5.4\" is not available for this group`) {
|
||||
t.Fatalf("expected model in message, got %s", w.Body.String())
|
||||
}
|
||||
if len(*calls) != 0 {
|
||||
t.Fatalf("expected handler not to run, got %v", *calls)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistMessagesDeniedAnthropicFormat(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"claude-opus-4.6"}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), `"not_found_error"`) {
|
||||
t.Fatalf("expected Anthropic not_found_error format, got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistGeminiDeniedGoogleFormat(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gemini-2.5-pro"), "/v1beta")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1beta/models/gemini-2.5-flash:generateContent", `{}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), `"NOT_FOUND"`) {
|
||||
t.Fatalf("expected Google NOT_FOUND format, got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistGeminiWildcardAndPrefixForms(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gemini-2.5-*"), "/v1beta")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1beta/models/models/gemini-2.5-flash:generateContent", `{}`)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("models/ prefixed model should match wildcard entry, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistGeminiGetModelRouteParam(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gemini-2.5-pro"), "/v1beta")
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-2.5-flash", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404 for unlisted GET model, got %d", w.Code)
|
||||
}
|
||||
|
||||
req = httptest.NewRequest(http.MethodGet, "/v1beta/models/gemini-2.5-pro", nil)
|
||||
w = httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200 for listed GET model, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistRealtimeQueryParameter(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "grok-4.6"), "/v1")
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/realtime?model=grok-4.6", nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200 for listed realtime model, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
req = httptest.NewRequest(http.MethodGet, "/v1/realtime?model=grok-4.20", nil)
|
||||
w = httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404 for unlisted realtime model, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistMultipartModelField(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.WriteField("prompt", "cat")
|
||||
_ = writer.Close()
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistMultipartDenied(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "dall-e-3")
|
||||
_ = writer.WriteField("prompt", "cat")
|
||||
_ = writer.Close()
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistNoModelPassesThrough(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"input":"no model field"}`)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200 when no model extractable, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistBodyRestoredForHandler(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
|
||||
c.Next()
|
||||
})
|
||||
router.Use(GroupModelAllowlist())
|
||||
router.POST("/v1/messages", func(c *gin.Context) {
|
||||
body, err := io.ReadAll(c.Request.Body)
|
||||
if err != nil {
|
||||
t.Errorf("handler reread body: %v", err)
|
||||
c.Status(http.StatusInternalServerError)
|
||||
return
|
||||
}
|
||||
if !strings.Contains(string(body), `"claude-sonnet-4.5"`) {
|
||||
t.Errorf("handler saw wrong body: %s", body)
|
||||
}
|
||||
c.Status(http.StatusOK)
|
||||
})
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"claude-sonnet-4.5","input":"hi"}`)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistWebSocketUpgradeSkipsBodyExtraction(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
|
||||
c.Next()
|
||||
})
|
||||
router.Use(GroupModelAllowlist())
|
||||
router.GET("/v1/responses", func(c *gin.Context) {
|
||||
c.Status(http.StatusSwitchingProtocols)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
req.Header.Set("Upgrade", "websocket")
|
||||
req.Header.Set("Connection", "Upgrade")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
if w.Code != http.StatusSwitchingProtocols {
|
||||
t.Fatalf("expected WS upgrade to reach handler, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// 伪造 WebSocket Upgrade 头的普通请求不得绕过白名单校验。
|
||||
func TestGroupModelAllowlistForgedUpgradeHeaderDoesNotBypass(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-opus-4.6"}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
req.Header.Set("Upgrade", "websocket")
|
||||
req.Header.Set("Connection", "Upgrade")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404 despite forged upgrade headers, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "claude-opus-4.6") {
|
||||
t.Fatalf("expected blocked model in message, got %s", w.Body.String())
|
||||
}
|
||||
if len(*calls) != 0 {
|
||||
t.Fatalf("expected handler not to run, got %v", *calls)
|
||||
}
|
||||
}
|
||||
|
||||
// Grok Realtime 升级请求的模型在查询参数里,携带 Upgrade 头也必须校验。
|
||||
func TestGroupModelAllowlistRealtimeUpgradeStillChecksQueryModel(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "grok-4.6"), "/v1")
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/realtime?model=grok-4.20", nil)
|
||||
req.Header.Set("Upgrade", "websocket")
|
||||
req.Header.Set("Connection", "Upgrade")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404 for unlisted realtime model with upgrade headers, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// Live 入口的 handler 只读 session.model:准入必须以 session.model 为准,
|
||||
// 与顶层 model 不一致时按实际执行模型校验。
|
||||
func TestGroupModelAllowlistLiveUsesSessionModel(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-realtime-model"), "/v1")
|
||||
|
||||
// session.model 不在白名单:即使顶层 model 命中也必须拒绝。
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/live", `{"model":"gpt-realtime-model","session":{"model":"blocked-model"},"sdp":"v=0"}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404 when session.model is blocked, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
|
||||
// session.model 命中:顶层 model 不在白名单也不影响(handler 不读它)。
|
||||
w = doJSON(t, router, http.MethodPost, "/v1/live", `{"model":"blocked-model","session":{"model":"gpt-realtime-model"},"sdp":"v=0"}`)
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("expected 200 when session.model is allowlisted, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistMarksRejectionReasons(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
router.Use(func(c *gin.Context) {
|
||||
c.Set(string(ContextKeyAPIKey), allowlistAPIKey(true, "claude-sonnet-4.5"))
|
||||
c.Next()
|
||||
})
|
||||
router.Use(GroupModelAllowlist())
|
||||
router.POST("/v1/responses", func(c *gin.Context) { c.Status(http.StatusOK) })
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"model":"gpt-5.4"}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("expected 404, got %d", w.Code)
|
||||
}
|
||||
// IngressRejectModelNotAllowed 与业务限流原因由 ops 中间件在真实链路落日志,
|
||||
// 此处至少验证拒绝原因可从 context 读取(MarkIngressRejected 生效)。
|
||||
}
|
||||
|
||||
func TestGroupModelAllowlistNilGroupPasses(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(&service.APIKey{}, "/v1")
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/messages", `{"model":"whatever"}`)
|
||||
if w.Code != http.StatusOK || len(*calls) != 1 {
|
||||
t.Fatalf("expected passthrough for key without group, got %d", w.Code)
|
||||
}
|
||||
}
|
||||
|
||||
// encoding/json 绑定系 handler(如批量生图 ShouldBindJSON)对键名大小写不敏感:
|
||||
// {"Model":"blocked"} 会绑定为 blocked,准入必须拒绝。
|
||||
func TestGroupModelAllowlistCaseVariantModelKeyRejected(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"Model":"claude-opus-4.6"}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("case-variant model key must be validated, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "claude-opus-4.6") {
|
||||
t.Fatalf("expected blocked model in message, got %s", w.Body.String())
|
||||
}
|
||||
if len(*calls) != 0 {
|
||||
t.Fatalf("expected handler not to run, got %v", *calls)
|
||||
}
|
||||
}
|
||||
|
||||
// encoding/json 绑定对重复键取末值、gjson 取首个:两个值都必须通过校验。
|
||||
func TestGroupModelAllowlistDuplicateModelKeysRejected(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "claude-sonnet-4.5"), "/v1")
|
||||
|
||||
w := doJSON(t, router, http.MethodPost, "/v1/responses", `{"model":"claude-sonnet-4.5","model":"claude-opus-4.6"}`)
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("duplicate model keys must all be validated, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "claude-opus-4.6") {
|
||||
t.Fatalf("expected the disallowed duplicate to be reported, got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// 图片编辑的 multipart 解析器对重复 model 字段取末值:全部字段值都必须校验。
|
||||
func TestGroupModelAllowlistMultipartDuplicateModelFieldsRejected(t *testing.T) {
|
||||
router, _ := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.WriteField("model", "gpt-image-1.5")
|
||||
_ = writer.WriteField("prompt", "cat")
|
||||
_ = writer.Close()
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusNotFound {
|
||||
t.Fatalf("duplicate multipart model fields must all be validated, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if !strings.Contains(w.Body.String(), "gpt-image-1.5") {
|
||||
t.Fatalf("expected the disallowed field value to be reported, got %s", w.Body.String())
|
||||
}
|
||||
}
|
||||
|
||||
// 全部候选都在白名单内时放行(重复但同值的字段不误伤)。
|
||||
func TestGroupModelAllowlistDuplicateIdenticalModelsAllowed(t *testing.T) {
|
||||
router, calls := newGroupModelAllowlistTestRouter(allowlistAPIKey(true, "gpt-image-1"), "/v1")
|
||||
|
||||
var body bytes.Buffer
|
||||
writer := multipart.NewWriter(&body)
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.WriteField("model", "gpt-image-1")
|
||||
_ = writer.Close()
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/images/edits", &body)
|
||||
req.Header.Set("Content-Type", writer.FormDataContentType())
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
if w.Code != http.StatusOK {
|
||||
t.Fatalf("identical duplicate fields should pass, got %d: %s", w.Code, w.Body.String())
|
||||
}
|
||||
if len(*calls) != 1 {
|
||||
t.Fatalf("expected handler to run once, got %v", *calls)
|
||||
}
|
||||
}
|
||||
@@ -28,6 +28,7 @@ const (
|
||||
IngressRejectGroupUnassigned IngressRejectReason = "group_unassigned"
|
||||
IngressRejectInvalidAuthRateLimited IngressRejectReason = "invalid_auth_rate_limited"
|
||||
IngressRejectAPIKeyAuthOverloaded IngressRejectReason = "api_key_auth_overloaded"
|
||||
IngressRejectModelNotAllowed IngressRejectReason = "model_not_allowed"
|
||||
)
|
||||
|
||||
const ingressRejectReasonContextKey = "ingress_reject_reason"
|
||||
|
||||
@@ -100,11 +100,19 @@ func abortWithOpenAIQuotaError(c *gin.Context, statusCode int, message string) {
|
||||
// GatewayErrorWriter 定义网关错误响应格式(不同协议使用不同格式)
|
||||
type GatewayErrorWriter func(c *gin.Context, status int, message string)
|
||||
|
||||
// AnthropicErrorWriter 按 Anthropic API 规范输出错误
|
||||
// AnthropicErrorWriter 按 Anthropic API 规范输出错误;error.type 随状态码映射
|
||||
// (404 -> not_found_error,403 -> permission_error,其余 api_error)。
|
||||
func AnthropicErrorWriter(c *gin.Context, status int, message string) {
|
||||
errorType := "api_error"
|
||||
switch status {
|
||||
case http.StatusNotFound:
|
||||
errorType = "not_found_error"
|
||||
case http.StatusForbidden:
|
||||
errorType = "permission_error"
|
||||
}
|
||||
c.JSON(status, gin.H{
|
||||
"type": "error",
|
||||
"error": gin.H{"type": "permission_error", "message": message},
|
||||
"error": gin.H{"type": errorType, "message": message},
|
||||
})
|
||||
}
|
||||
|
||||
@@ -119,6 +127,17 @@ func GoogleErrorWriter(c *gin.Context, status int, message string) {
|
||||
})
|
||||
}
|
||||
|
||||
// OpenAIErrorWriter 按 OpenAI API 规范输出模型级错误(model_not_found)。
|
||||
func OpenAIErrorWriter(c *gin.Context, status int, message string) {
|
||||
c.JSON(status, gin.H{
|
||||
"error": gin.H{
|
||||
"message": message,
|
||||
"type": "invalid_request_error",
|
||||
"code": "model_not_found",
|
||||
},
|
||||
})
|
||||
}
|
||||
|
||||
// RequireGroupAssignment 检查 API Key 是否已分配到分组,
|
||||
// 如果未分组且系统设置不允许未分组 Key 调度则返回 403。
|
||||
func RequireGroupAssignment(settingService *service.SettingService, writeError GatewayErrorWriter) gin.HandlerFunc {
|
||||
|
||||
@@ -333,7 +333,7 @@ func registerGroupRoutes(admin *gin.RouterGroup, h *handler.Handlers) {
|
||||
groups.GET("/capacity-summary", h.Admin.Group.GetCapacitySummary)
|
||||
groups.GET("/live-capability", h.Admin.Group.GetLiveCapability)
|
||||
groups.PUT("/sort-order", h.Admin.Group.UpdateSortOrder)
|
||||
groups.GET("/:id/models-list-candidates", h.Admin.Group.GetModelsListCandidates)
|
||||
groups.GET("/:id/model-allowlist-candidates", h.Admin.Group.GetGroupModelAllowlistCandidates)
|
||||
groups.GET("/:id/composite-routes", h.Admin.Group.ListCompositeRoutes)
|
||||
groups.POST("/:id/composite-routes", h.Admin.Group.CreateCompositeRoute)
|
||||
groups.POST("/:id/composite-routes/preview", h.Admin.Group.PreviewCompositeRoute)
|
||||
|
||||
@@ -10,6 +10,7 @@ import (
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
@@ -191,7 +192,7 @@ func TestCompositeRequestModelFromMultipartLiveSession(t *testing.T) {
|
||||
require.NoError(t, writer.WriteField("session", `{"model":"live-alias"}`))
|
||||
require.NoError(t, writer.Close())
|
||||
|
||||
require.Equal(t, "live-alias", compositeRequestModelFromBody(writer.FormDataContentType(), body.Bytes()))
|
||||
require.Equal(t, "live-alias", requestmodel.FromBody(writer.FormDataContentType(), body.Bytes()))
|
||||
}
|
||||
|
||||
func TestCompositeCodexControlPathsUseResponsesRoutes(t *testing.T) {
|
||||
@@ -312,3 +313,58 @@ func TestCompositeGeminiTargetPlatformMiddlewareUsesPathRoute(t *testing.T) {
|
||||
|
||||
require.Equal(t, http.StatusNoContent, w.Code)
|
||||
}
|
||||
|
||||
// Live 入口顶层 model 与 session.model 不一致时,合成路由必须按 session.model
|
||||
// 分发与改写(与白名单准入、Live handler 一致),不得命中顶层别名的映射并把
|
||||
// session 模型覆盖为别名上游。
|
||||
func TestCompositeLiveRouteDispatchesBySessionModelNotTopLevelAlias(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
resolver := service.NewCompositeRouteResolver(compositeRouteRepoStub{
|
||||
routes: []service.CompositeModelRoute{
|
||||
{
|
||||
ID: 1,
|
||||
GroupID: 1,
|
||||
PublicModel: "live-alias",
|
||||
MatchType: service.CompositeRouteMatchExact,
|
||||
TargetPlatform: service.PlatformOpenAI,
|
||||
UpstreamModel: "gpt-alias-upstream",
|
||||
Endpoint: service.CompositeRouteEndpointAny,
|
||||
Priority: 100,
|
||||
Enabled: true,
|
||||
},
|
||||
},
|
||||
})
|
||||
router.Use(gin.HandlerFunc(servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
|
||||
groupID := int64(1)
|
||||
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
|
||||
GroupID: &groupID,
|
||||
Group: &service.Group{ID: groupID, Platform: service.PlatformComposite},
|
||||
})
|
||||
c.Next()
|
||||
})))
|
||||
router.Use(compositeTargetPlatformMiddleware(resolver))
|
||||
router.POST("/backend-api/codex/realtime/calls", func(c *gin.Context) {
|
||||
upstream, ok := service.ResolvedUpstreamModelFromContext(c.Request.Context())
|
||||
require.True(t, ok)
|
||||
require.Equal(t, "gpt-realtime-1", upstream,
|
||||
"Live dispatch must resolve the session model, never the top-level alias upstream")
|
||||
body, err := io.ReadAll(c.Request.Body)
|
||||
require.NoError(t, err)
|
||||
require.JSONEq(t, `{"model":"live-alias","session":{"model":"gpt-realtime-1"},"sdp":"v=0"}`, string(body),
|
||||
"session.model must stay untouched when dispatch resolves by session model")
|
||||
c.Status(http.StatusNoContent)
|
||||
})
|
||||
|
||||
req := httptest.NewRequest(
|
||||
http.MethodPost,
|
||||
"/backend-api/codex/realtime/calls",
|
||||
strings.NewReader(`{"model":"live-alias","session":{"model":"gpt-realtime-1"},"sdp":"v=0"}`),
|
||||
)
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusNoContent, w.Code)
|
||||
}
|
||||
|
||||
@@ -1,18 +1,14 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"errors"
|
||||
"io"
|
||||
"mime"
|
||||
"mime/multipart"
|
||||
"net/http"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/requestmodel"
|
||||
"github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
|
||||
@@ -45,6 +41,10 @@ func RegisterGatewayRoutes(
|
||||
requireGroupAnthropic := middleware.RequireGroupAssignment(settingService, middleware.AnthropicErrorWriter)
|
||||
requireGroupGoogle := middleware.RequireGroupAssignment(settingService, middleware.GoogleErrorWriter)
|
||||
|
||||
// 分组级模型白名单准入:在 apiKeyAuth 之后、compositeTarget 之前,
|
||||
// 保证校验发生在合成路由改写与调度之前,且只看客户端书写的模型名。
|
||||
groupModelAllowlist := middleware.GroupModelAllowlist()
|
||||
|
||||
isOpenAIResponsesCompatibleGatewayPlatform := func(c *gin.Context) bool {
|
||||
switch getGroupPlatform(c) {
|
||||
case service.PlatformOpenAI, service.PlatformGrok,
|
||||
@@ -189,6 +189,7 @@ func RegisterGatewayRoutes(
|
||||
gateway.Use(endpointNorm)
|
||||
gateway.Use(gin.HandlerFunc(apiKeyAuth))
|
||||
gateway.GET("/sub2api/billing", h.Gateway.KeyBillingInfo)
|
||||
gateway.Use(groupModelAllowlist)
|
||||
gateway.Use(compositeTarget)
|
||||
gateway.Use(requireGroupAnthropic)
|
||||
{
|
||||
@@ -341,6 +342,7 @@ func RegisterGatewayRoutes(
|
||||
gemini.Use(opsErrorLogger)
|
||||
gemini.Use(endpointNorm)
|
||||
gemini.Use(middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
|
||||
gemini.Use(groupModelAllowlist)
|
||||
gemini.Use(compositeGeminiTarget)
|
||||
gemini.Use(requireGroupGoogle)
|
||||
{
|
||||
@@ -358,16 +360,21 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.Gateway.Responses(c)
|
||||
}
|
||||
r.POST("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, responsesHandler)
|
||||
r.POST("/responses/*subpath", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, guardResponsesSubpath(responsesHandler))
|
||||
r.POST("/alpha/search", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.OpenAIGateway.AlphaSearch)
|
||||
r.GET("/responses", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
// 根路径别名共用中间件链:白名单准入在 apiKeyAuth 之后、compositeTarget
|
||||
// 之前,避免逐条路由手工维护链导致漏挂。
|
||||
rootRoute := func(method, path string, limit gin.HandlerFunc, handler gin.HandlerFunc) {
|
||||
r.Handle(method, path, limit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic, handler)
|
||||
}
|
||||
rootRoute(http.MethodPost, "/responses", bodyLimit, responsesHandler)
|
||||
rootRoute(http.MethodPost, "/responses/*subpath", bodyLimit, guardResponsesSubpath(responsesHandler))
|
||||
rootRoute(http.MethodPost, "/alpha/search", textBodyLimit, h.OpenAIGateway.AlphaSearch)
|
||||
rootRoute(http.MethodGet, "/responses", bodyLimit, func(c *gin.Context) {
|
||||
h.OpenAIGateway.ResponsesWebSocket(c)
|
||||
})
|
||||
r.GET("/models", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), requireGroupAnthropic, modelsHandler)
|
||||
r.POST("/messages/count_tokens", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, countTokensHandler)
|
||||
rootRoute(http.MethodGet, "/models", bodyLimit, modelsHandler)
|
||||
rootRoute(http.MethodPost, "/messages/count_tokens", bodyLimit, countTokensHandler)
|
||||
codexDirect := r.Group("/backend-api/codex")
|
||||
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic)
|
||||
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic)
|
||||
{
|
||||
codexDirect.POST("/realtime/calls", h.OpenAIGateway.Live)
|
||||
codexDirect.GET("/:call_id", h.OpenAIGateway.LiveSideband)
|
||||
@@ -380,14 +387,14 @@ func RegisterGatewayRoutes(
|
||||
codexDirect.GET("/models", codexModelsHandler)
|
||||
}
|
||||
// OpenAI Chat Completions API(不带v1前缀的别名)— auto-route based on group platform
|
||||
r.POST("/chat/completions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
rootRoute(http.MethodPost, "/chat/completions", bodyLimit, func(c *gin.Context) {
|
||||
if isOpenAIResponsesCompatibleGatewayPlatform(c) {
|
||||
h.OpenAIGateway.ChatCompletions(c)
|
||||
return
|
||||
}
|
||||
h.Gateway.ChatCompletions(c)
|
||||
})
|
||||
r.POST("/embeddings", textBodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
rootRoute(http.MethodPost, "/embeddings", textBodyLimit, func(c *gin.Context) {
|
||||
if !isOpenAIOnlyEndpointGatewayPlatform(c) {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{
|
||||
@@ -400,23 +407,23 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.OpenAIGateway.Embeddings(c)
|
||||
})
|
||||
r.POST("/images/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, imagesHandler)
|
||||
r.POST("/images/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, imagesHandler)
|
||||
r.POST("/images/generations/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Submit)
|
||||
r.POST("/images/edits/async", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Submit)
|
||||
r.GET("/images/tasks/:task_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, h.AsyncImage.Get)
|
||||
r.POST("/videos", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoGenerationHandler)
|
||||
r.POST("/videos/generations", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoGenerationHandler)
|
||||
r.POST("/videos/edits", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoEditHandler)
|
||||
r.POST("/videos/extensions", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoExtensionHandler)
|
||||
r.GET("/videos/generations/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
|
||||
r.GET("/videos/edits/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
|
||||
r.GET("/videos/extensions/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
|
||||
r.GET("/videos/generations/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
|
||||
r.GET("/videos/edits/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
|
||||
r.GET("/videos/extensions/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
|
||||
r.GET("/videos/:request_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoStatusHandler)
|
||||
r.GET("/videos/:request_id/content", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, videoContentHandler)
|
||||
rootRoute(http.MethodPost, "/images/generations", bodyLimit, imagesHandler)
|
||||
rootRoute(http.MethodPost, "/images/edits", bodyLimit, imagesHandler)
|
||||
rootRoute(http.MethodPost, "/images/generations/async", bodyLimit, h.AsyncImage.Submit)
|
||||
rootRoute(http.MethodPost, "/images/edits/async", bodyLimit, h.AsyncImage.Submit)
|
||||
rootRoute(http.MethodGet, "/images/tasks/:task_id", bodyLimit, h.AsyncImage.Get)
|
||||
rootRoute(http.MethodPost, "/videos", bodyLimit, videoGenerationHandler)
|
||||
rootRoute(http.MethodPost, "/videos/generations", bodyLimit, videoGenerationHandler)
|
||||
rootRoute(http.MethodPost, "/videos/edits", bodyLimit, videoEditHandler)
|
||||
rootRoute(http.MethodPost, "/videos/extensions", bodyLimit, videoExtensionHandler)
|
||||
rootRoute(http.MethodGet, "/videos/generations/:request_id/content", bodyLimit, videoContentHandler)
|
||||
rootRoute(http.MethodGet, "/videos/edits/:request_id/content", bodyLimit, videoContentHandler)
|
||||
rootRoute(http.MethodGet, "/videos/extensions/:request_id/content", bodyLimit, videoContentHandler)
|
||||
rootRoute(http.MethodGet, "/videos/generations/:request_id", bodyLimit, videoStatusHandler)
|
||||
rootRoute(http.MethodGet, "/videos/edits/:request_id", bodyLimit, videoStatusHandler)
|
||||
rootRoute(http.MethodGet, "/videos/extensions/:request_id", bodyLimit, videoStatusHandler)
|
||||
rootRoute(http.MethodGet, "/videos/:request_id", bodyLimit, videoStatusHandler)
|
||||
rootRoute(http.MethodGet, "/videos/:request_id/content", bodyLimit, videoContentHandler)
|
||||
|
||||
rootVoiceHandler := func(endpoint string) gin.HandlerFunc {
|
||||
return func(c *gin.Context) {
|
||||
@@ -428,9 +435,9 @@ func RegisterGatewayRoutes(
|
||||
h.OpenAIGateway.GrokVoice(c, endpoint)
|
||||
}
|
||||
}
|
||||
r.POST("/tts", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("tts"))
|
||||
r.POST("/stt", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("stt"))
|
||||
r.POST("/custom-voices", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("custom-voices"))
|
||||
rootRoute(http.MethodPost, "/tts", bodyLimit, rootVoiceHandler("tts"))
|
||||
rootRoute(http.MethodPost, "/stt", bodyLimit, rootVoiceHandler("stt"))
|
||||
rootRoute(http.MethodPost, "/custom-voices", bodyLimit, rootVoiceHandler("custom-voices"))
|
||||
rootCustomVoicePathHandler := func(c *gin.Context) {
|
||||
if getGroupPlatform(c) != service.PlatformGrok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
@@ -439,12 +446,12 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.OpenAIGateway.GrokVoice(c, grokCustomVoiceEndpoint(c))
|
||||
}
|
||||
r.GET("/custom-voices", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootVoiceHandler("custom-voices"))
|
||||
r.GET("/custom-voices/:voice_id/audio", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
|
||||
r.GET("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
|
||||
r.PATCH("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
|
||||
r.DELETE("/custom-voices/:voice_id", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, rootCustomVoicePathHandler)
|
||||
r.GET("/realtime", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
rootRoute(http.MethodGet, "/custom-voices", bodyLimit, rootVoiceHandler("custom-voices"))
|
||||
rootRoute(http.MethodGet, "/custom-voices/:voice_id/audio", bodyLimit, rootCustomVoicePathHandler)
|
||||
rootRoute(http.MethodGet, "/custom-voices/:voice_id", bodyLimit, rootCustomVoicePathHandler)
|
||||
rootRoute(http.MethodPatch, "/custom-voices/:voice_id", bodyLimit, rootCustomVoicePathHandler)
|
||||
rootRoute(http.MethodDelete, "/custom-voices/:voice_id", bodyLimit, rootCustomVoicePathHandler)
|
||||
rootRoute(http.MethodGet, "/realtime", bodyLimit, func(c *gin.Context) {
|
||||
if getGroupPlatform(c) != service.PlatformGrok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Realtime API is not supported for this platform"}})
|
||||
@@ -452,7 +459,7 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.OpenAIGateway.GrokRealtime(c)
|
||||
})
|
||||
r.POST("/web_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
rootRoute(http.MethodPost, "/web_search", bodyLimit, func(c *gin.Context) {
|
||||
if getGroupPlatform(c) != service.PlatformGrok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "Web Search API is not supported for this platform"}})
|
||||
@@ -460,7 +467,7 @@ func RegisterGatewayRoutes(
|
||||
}
|
||||
h.Gateway.WebSearch(c)
|
||||
})
|
||||
r.POST("/x_search", bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), compositeTarget, requireGroupAnthropic, func(c *gin.Context) {
|
||||
rootRoute(http.MethodPost, "/x_search", bodyLimit, func(c *gin.Context) {
|
||||
if getGroupPlatform(c) != service.PlatformGrok {
|
||||
service.MarkOpsClientBusinessLimited(c, service.OpsClientBusinessLimitedReasonLocalFeatureGate)
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{"type": "not_found_error", "message": "X Search API is not supported for this platform"}})
|
||||
@@ -480,6 +487,7 @@ func RegisterGatewayRoutes(
|
||||
antigravityV1.Use(endpointNorm)
|
||||
antigravityV1.Use(middleware.ForcePlatform(service.PlatformAntigravity))
|
||||
antigravityV1.Use(gin.HandlerFunc(apiKeyAuth))
|
||||
antigravityV1.Use(groupModelAllowlist)
|
||||
antigravityV1.Use(requireGroupAnthropic)
|
||||
{
|
||||
antigravityV1.POST("/messages", h.Gateway.Messages)
|
||||
@@ -495,6 +503,7 @@ func RegisterGatewayRoutes(
|
||||
antigravityV1Beta.Use(endpointNorm)
|
||||
antigravityV1Beta.Use(middleware.ForcePlatform(service.PlatformAntigravity))
|
||||
antigravityV1Beta.Use(middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg))
|
||||
antigravityV1Beta.Use(groupModelAllowlist)
|
||||
antigravityV1Beta.Use(requireGroupGoogle)
|
||||
{
|
||||
antigravityV1Beta.GET("/models", h.Gateway.GeminiV1BetaListModels)
|
||||
@@ -555,7 +564,10 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
|
||||
return
|
||||
}
|
||||
|
||||
model := compositeRequestModelFromBody(c.GetHeader("Content-Type"), body)
|
||||
// Live 入口按 session.model 分发并改写(与白名单准入、Live handler 的
|
||||
// 解析保持一致),避免顶层 model 别名与 session 模型不一致时改错对象。
|
||||
routePath := c.FullPath()
|
||||
model := requestmodel.FromBodyForRoute(routePath, c.GetHeader("Content-Type"), body)
|
||||
if model != "" {
|
||||
decision, err := resolver.Resolve(c.Request.Context(), apiKey.Group.ID, model, compositeRouteEndpointForPath(c.Request.URL.Path))
|
||||
if err != nil {
|
||||
@@ -566,7 +578,7 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
|
||||
if decision.Matched {
|
||||
c.Request = c.Request.WithContext(service.WithCompositeRouteDecision(c.Request.Context(), decision))
|
||||
if upstreamModel := strings.TrimSpace(decision.UpstreamModel); upstreamModel != "" && upstreamModel != model && gjson.ValidBytes(body) {
|
||||
if _, modelPath := compositeJSONRequestModel(body); modelPath != "" {
|
||||
if _, modelPath := requestmodel.JSONModelPathForRoute(routePath, body); modelPath != "" {
|
||||
if rewritten, rewriteErr := sjson.SetBytes(body, modelPath, upstreamModel); rewriteErr == nil {
|
||||
body = rewritten
|
||||
}
|
||||
@@ -574,68 +586,11 @@ func compositeTargetPlatformMiddleware(resolver *service.CompositeRouteResolver)
|
||||
}
|
||||
}
|
||||
}
|
||||
resetRequestBody(c, body)
|
||||
requestmodel.ResetRequestBody(c.Request, body)
|
||||
c.Next()
|
||||
}
|
||||
}
|
||||
|
||||
func compositeRequestModelFromBody(contentType string, body []byte) string {
|
||||
if model, _ := compositeJSONRequestModel(body); model != "" {
|
||||
return model
|
||||
}
|
||||
return compositeMultipartModelFromBody(contentType, body)
|
||||
}
|
||||
|
||||
func compositeJSONRequestModel(body []byte) (string, string) {
|
||||
for _, path := range []string{"model", "session.model"} {
|
||||
model := gjson.GetBytes(body, path)
|
||||
if model.Type != gjson.String {
|
||||
continue
|
||||
}
|
||||
if value := strings.TrimSpace(model.String()); value != "" {
|
||||
return value, path
|
||||
}
|
||||
}
|
||||
return "", ""
|
||||
}
|
||||
|
||||
func compositeMultipartModelFromBody(contentType string, body []byte) string {
|
||||
mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType))
|
||||
if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
|
||||
return ""
|
||||
}
|
||||
boundary := strings.TrimSpace(params["boundary"])
|
||||
if boundary == "" {
|
||||
return ""
|
||||
}
|
||||
reader := multipart.NewReader(bytes.NewReader(body), boundary)
|
||||
for {
|
||||
part, err := reader.NextPart()
|
||||
if errors.Is(err, io.EOF) {
|
||||
return ""
|
||||
}
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
fieldName := part.FormName()
|
||||
if part.FileName() != "" || (fieldName != "model" && fieldName != "session") {
|
||||
continue
|
||||
}
|
||||
data, err := io.ReadAll(part)
|
||||
if err != nil {
|
||||
return ""
|
||||
}
|
||||
switch fieldName {
|
||||
case "model":
|
||||
return strings.TrimSpace(string(data))
|
||||
case "session":
|
||||
if model, _ := compositeJSONRequestModel(data); model != "" {
|
||||
return model
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func compositeGeminiTargetPlatformMiddleware(resolver *service.CompositeRouteResolver) gin.HandlerFunc {
|
||||
if resolver == nil {
|
||||
resolver = service.NewCompositeRouteResolver(nil)
|
||||
@@ -696,12 +651,6 @@ func compositeGeminiModelFromParams(c *gin.Context) string {
|
||||
return modelAction
|
||||
}
|
||||
|
||||
func resetRequestBody(c *gin.Context, body []byte) {
|
||||
c.Request.Body = io.NopCloser(bytes.NewReader(body))
|
||||
c.Request.ContentLength = int64(len(body))
|
||||
c.Request.Header.Set("Content-Length", strconv.Itoa(len(body)))
|
||||
}
|
||||
|
||||
func compositeRouteEndpointForPath(path string) string {
|
||||
switch {
|
||||
case strings.Contains(path, "/messages/count_tokens"):
|
||||
|
||||
@@ -0,0 +1,208 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"os"
|
||||
"regexp"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func newGatewayRoutesTestRouterWithGroup(group *service.Group) *gin.Engine {
|
||||
gin.SetMode(gin.TestMode)
|
||||
router := gin.New()
|
||||
RegisterGatewayRoutes(
|
||||
router,
|
||||
&handler.Handlers{
|
||||
Gateway: &handler.GatewayHandler{},
|
||||
OpenAIGateway: &handler.OpenAIGatewayHandler{},
|
||||
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
|
||||
},
|
||||
servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
|
||||
groupID := int64(1)
|
||||
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{
|
||||
GroupID: &groupID,
|
||||
Group: group,
|
||||
})
|
||||
c.Next()
|
||||
}),
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
nil,
|
||||
&config.Config{
|
||||
Gateway: config.GatewayConfig{
|
||||
MaxBodySize: 1024 * 1024,
|
||||
TextMaxBodySize: 1024 * 1024,
|
||||
},
|
||||
},
|
||||
)
|
||||
return router
|
||||
}
|
||||
|
||||
func allowlistGroup(platform string, enabled bool, models ...string) *service.Group {
|
||||
return &service.Group{
|
||||
Platform: platform,
|
||||
ModelAllowlist: service.GroupModelAllowlist{
|
||||
Enabled: enabled,
|
||||
Models: models,
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
// TestGatewayRoutesGroupModelAllowlistMountedOnEveryGatewayRoute follows the
|
||||
// source-level route assertion convention of prompt_audit_route_coverage_test.go:
|
||||
// every gateway chain must mount groupModelAllowlist after api key auth and
|
||||
// before the composite rewrite (gateway.go + rootRoute helper).
|
||||
func TestGatewayRoutesGroupModelAllowlistMountedOnEveryGatewayRoute(t *testing.T) {
|
||||
routeSource, err := os.ReadFile("gateway.go")
|
||||
require.NoError(t, err)
|
||||
source := string(routeSource)
|
||||
|
||||
// rootRoute helper:apiKeyAuth 之后、compositeTarget 之前。
|
||||
rootHelper := regexp.MustCompile(regexp.QuoteMeta(`r.Handle(method, path, limit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic, handler)`))
|
||||
require.Regexp(t, rootHelper, source,
|
||||
"root alias helper must place the allowlist between apiKeyAuth and compositeTarget")
|
||||
|
||||
chains := []struct {
|
||||
group string
|
||||
auth string
|
||||
marker string
|
||||
composite string
|
||||
}{
|
||||
{group: "gateway", auth: "gin.HandlerFunc(apiKeyAuth)", marker: "gateway.Use(groupModelAllowlist)", composite: "gateway.Use(compositeTarget)"},
|
||||
{group: "gemini", auth: "middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg)", marker: "gemini.Use(groupModelAllowlist)", composite: "gemini.Use(compositeGeminiTarget)"},
|
||||
{group: "antigravityV1", auth: "gin.HandlerFunc(apiKeyAuth)", marker: "antigravityV1.Use(groupModelAllowlist)", composite: "antigravityV1.Use(requireGroupAnthropic)"},
|
||||
{group: "antigravityV1Beta", auth: "middleware.APIKeyAuthWithSubscriptionGoogle(apiKeyService, subscriptionService, cfg)", marker: "antigravityV1Beta.Use(groupModelAllowlist)", composite: "antigravityV1Beta.Use(requireGroupGoogle)"},
|
||||
}
|
||||
for _, chain := range chains {
|
||||
re := regexp.MustCompile(
|
||||
regexp.QuoteMeta(chain.group+".Use("+chain.auth) +
|
||||
`[\s\S]{0,400}?` + regexp.QuoteMeta(chain.marker) +
|
||||
`[\s\S]{0,400}?` + regexp.QuoteMeta(chain.composite))
|
||||
require.Regexp(t, re, source,
|
||||
"%s chain must mount groupModelAllowlist after auth and before %s", chain.group, chain.composite)
|
||||
}
|
||||
|
||||
// codexDirect 链是一条 Use 调用,直接断言顺序。
|
||||
codexDirect := regexp.MustCompile(regexp.QuoteMeta(`codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic)`))
|
||||
require.Regexp(t, codexDirect, source, "codexDirect chain must mount the allowlist after auth and before compositeTarget")
|
||||
|
||||
// 所有带 apiKeyAuth 的根路径路由必须收敛到 rootRoute,避免漏挂。
|
||||
stray := regexp.MustCompile(`\br\.(GET|POST|PUT|PATCH|DELETE)\("[^"]+",[^(]*apiKeyAuth`)
|
||||
require.NotRegexp(t, stray, source,
|
||||
"root alias routes must use rootRoute so the allowlist cannot be forgotten")
|
||||
}
|
||||
|
||||
// TestGatewayRoutesGroupModelAllowlistBlocksCompositeModelBeforeRewrite asserts
|
||||
// admission is decided on the client-written public model before the composite
|
||||
// route middleware can rewrite it to an upstream model.
|
||||
func TestGatewayRoutesGroupModelAllowlistBlocksCompositeModelBeforeRewrite(t *testing.T) {
|
||||
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformComposite, true, "claude-*"))
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"grok-4.3","messages":[{"role":"user","content":"hi"}]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusNotFound, w.Code)
|
||||
require.Contains(t, w.Body.String(), "not_found_error")
|
||||
require.Contains(t, w.Body.String(), "grok-4.3")
|
||||
}
|
||||
|
||||
func TestGatewayRoutesGroupModelAllowlistAllowsListedCompositeModel(t *testing.T) {
|
||||
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformComposite, true, "claude-*"))
|
||||
|
||||
req := httptest.NewRequest(http.MethodPost, "/v1/messages", strings.NewReader(`{"model":"claude-sonnet-4.5","messages":[{"role":"user","content":"hi"}]}`))
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.NotEqual(t, http.StatusNotFound, w.Code, "allowlisted model must pass admission: %s", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGatewayRoutesGroupModelAllowlistCoversRootAliasRoutes(t *testing.T) {
|
||||
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformOpenAI, true, "gpt-5.4"))
|
||||
|
||||
paths := []struct {
|
||||
method string
|
||||
path string
|
||||
body string
|
||||
}{
|
||||
{http.MethodPost, "/responses", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/responses/compact", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/chat/completions", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/embeddings", `{"model":"gpt-4.1","input":"hi"}`},
|
||||
{http.MethodPost, "/images/generations", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/images/edits", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/videos/generations", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/messages/count_tokens", `{"model":"gpt-4.1","messages":[]}`},
|
||||
{http.MethodPost, "/tts", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/stt", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/alpha/search", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodGet, "/realtime?model=gpt-4.1", ""},
|
||||
{http.MethodPost, "/v1/responses", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/messages", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/messages/count_tokens", `{"model":"gpt-4.1","messages":[]}`},
|
||||
{http.MethodPost, "/v1/chat/completions", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/embeddings", `{"model":"gpt-4.1","input":"hi"}`},
|
||||
{http.MethodPost, "/v1/images/generations", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/videos/generations", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/v1/live", `{"session":{"model":"gpt-4.1"},"sdp":"v=0"}`},
|
||||
{http.MethodPost, "/backend-api/codex/responses", `{"model":"gpt-4.1"}`},
|
||||
{http.MethodPost, "/backend-api/codex/realtime/calls", `{"session":{"model":"gpt-4.1"},"sdp":"v=0"}`},
|
||||
{http.MethodPost, "/antigravity/v1/messages", `{"model":"gemini-2.5-pro","messages":[]}`},
|
||||
}
|
||||
|
||||
for _, tc := range paths {
|
||||
req := httptest.NewRequest(tc.method, tc.path, strings.NewReader(tc.body))
|
||||
if tc.body != "" {
|
||||
req.Header.Set("Content-Type", "application/json")
|
||||
}
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.Equal(t, http.StatusNotFound, w.Code, "%s %s should be denied by the allowlist, got body: %s", tc.method, tc.path, w.Body.String())
|
||||
require.Contains(t, w.Body.String(), "not available for this group", "%s %s", tc.method, tc.path)
|
||||
}
|
||||
}
|
||||
|
||||
func TestGatewayRoutesGroupModelAllowlistSkipsWebSocketUpgrade(t *testing.T) {
|
||||
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformOpenAI, true, "gpt-5.4"))
|
||||
|
||||
req := httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
|
||||
req.Header.Set("Upgrade", "websocket")
|
||||
req.Header.Set("Connection", "Upgrade")
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
|
||||
require.NotEqual(t, http.StatusNotFound, w.Code,
|
||||
"WS upgrade requests must be admitted and per-frame checked in the handler, got: %s", w.Body.String())
|
||||
}
|
||||
|
||||
func TestGatewayRoutesGroupModelAllowlistModelFreeRoutesUnaffected(t *testing.T) {
|
||||
router := newGatewayRoutesTestRouterWithGroup(allowlistGroup(service.PlatformGrok, true, "grok-4.6"))
|
||||
|
||||
// 不携带模型的入口天然不受白名单影响(状态查询/内容下载/模型列表)。
|
||||
for _, path := range []string{
|
||||
"/v1/videos/generations/req-1",
|
||||
"/v1/videos/req-1/content",
|
||||
"/v1/images/tasks/task-1",
|
||||
"/v1/custom-voices",
|
||||
} {
|
||||
req := httptest.NewRequest(http.MethodGet, path, nil)
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, req)
|
||||
require.NotContains(t, w.Body.String(), "not available for this group",
|
||||
"%s should not be blocked by the allowlist middleware, got: %s", path, w.Body.String())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,94 @@
|
||||
package routes
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/handler"
|
||||
servermiddleware "github.com/Wei-Shaw/sub2api/internal/server/middleware"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type pinnedModelsRoutesRepository struct {
|
||||
service.AccountRepository
|
||||
account service.Account
|
||||
}
|
||||
|
||||
func (r *pinnedModelsRoutesRepository) ListByGroup(context.Context, int64) ([]service.Account, error) {
|
||||
return []service.Account{r.account}, nil
|
||||
}
|
||||
|
||||
type pinnedModelsRoutesUpstream struct {
|
||||
service.HTTPUpstream
|
||||
ordinaryCalls atomic.Int32
|
||||
codexCalls atomic.Int32
|
||||
}
|
||||
|
||||
func (u *pinnedModelsRoutesUpstream) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
|
||||
body := `{"data":[{"id":"ordinary-upstream-model"}]}`
|
||||
if req.URL.Query().Has("client_version") {
|
||||
u.codexCalls.Add(1)
|
||||
body = `{"models":[{"slug":"gpt-5.5"}]}`
|
||||
} else {
|
||||
u.ordinaryCalls.Add(1)
|
||||
}
|
||||
return &http.Response{StatusCode: 200, Header: make(http.Header), Body: io.NopCloser(strings.NewReader(body))}, nil
|
||||
}
|
||||
|
||||
func TestGatewayRoutesPinnedModelsDispatchesOrdinaryAndCodexRequests(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
repo := &pinnedModelsRoutesRepository{account: service.Account{
|
||||
ID: 7, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Status: service.StatusActive, Schedulable: true,
|
||||
Credentials: map[string]any{"api_key": "test-models-key", "base_url": "https://models.example/v1"},
|
||||
}}
|
||||
upstream := &pinnedModelsRoutesUpstream{}
|
||||
cfg := &config.Config{RunMode: config.RunModeSimple}
|
||||
s := service.NewOpenAIGatewayService(repo, nil, nil, nil, nil, nil, nil, cfg,
|
||||
nil, nil, nil, nil, nil, upstream, nil, nil, nil, nil, nil, nil, nil, nil)
|
||||
h := &handler.Handlers{
|
||||
Gateway: handler.NewGatewayHandler(nil, s, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil),
|
||||
OpenAIGateway: handler.NewOpenAIGatewayHandler(s, nil, nil, nil, nil, nil, nil, nil, cfg),
|
||||
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
|
||||
}
|
||||
group := &service.Group{ID: 1, Platform: service.PlatformOpenAI,
|
||||
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{7}}}
|
||||
router := gin.New()
|
||||
RegisterGatewayRoutes(router, h, servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
|
||||
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
|
||||
c.Next()
|
||||
}), nil, nil, nil, nil, nil, cfg)
|
||||
for _, path := range []string{"/v1/models", "/models", "/v1/models?client_version=", "/models?client_version="} {
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil))
|
||||
require.Equal(t, http.StatusOK, w.Code, "%s: %s", path, w.Body.String())
|
||||
var response struct {
|
||||
Object string `json:"object"`
|
||||
Data []struct {
|
||||
ID string `json:"id"`
|
||||
} `json:"data"`
|
||||
}
|
||||
require.NoError(t, json.Unmarshal(w.Body.Bytes(), &response))
|
||||
require.Equal(t, "list", response.Object)
|
||||
require.Len(t, response.Data, 1)
|
||||
require.Equal(t, "ordinary-upstream-model", response.Data[0].ID)
|
||||
}
|
||||
for _, path := range []string{"/v1/models?client_version=" + service.CodexCanonicalClientVersion(), "/models?client_version=" + service.CodexCanonicalClientVersion(), "/backend-api/codex/models"} {
|
||||
w := httptest.NewRecorder()
|
||||
router.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil))
|
||||
require.Equal(t, http.StatusOK, w.Code, "%s: %s", path, w.Body.String())
|
||||
require.Contains(t, w.Body.String(), `"slug":"gpt-5.5"`)
|
||||
require.NotContains(t, w.Body.String(), `"data"`)
|
||||
}
|
||||
require.EqualValues(t, 1, upstream.ordinaryCalls.Load())
|
||||
require.EqualValues(t, 1, upstream.codexCalls.Load())
|
||||
}
|
||||
@@ -1840,9 +1840,10 @@ func (a *Account) SupportsOpenAIEndpointCapability(capability OpenAIEndpointCapa
|
||||
}
|
||||
|
||||
// GrokMediaGenerationEligibility reports whether a Grok account may receive
|
||||
// new image/video generation requests. OAuth media fails closed unless billing
|
||||
// observations provide positive paid-entitlement evidence. An explicit
|
||||
// operator override takes precedence over probe data.
|
||||
// new image/video generation requests. Explicit evidence of a forbidden or
|
||||
// free account blocks media, while an incomplete successful billing response
|
||||
// remains eligible for backwards compatibility. An explicit operator
|
||||
// override takes precedence over probe data.
|
||||
func (a *Account) GrokMediaGenerationEligibility() (bool, string) {
|
||||
if a == nil || !a.IsGrok() {
|
||||
return false, "not_grok"
|
||||
@@ -1868,7 +1869,12 @@ func (a *Account) GrokMediaGenerationEligibility() (bool, string) {
|
||||
return false, "billing_free_tier"
|
||||
}
|
||||
if !grokBillingHasAuthoritativeQuota(billing) {
|
||||
return false, "billing_inconclusive"
|
||||
// Billing endpoints can return 200 with an account-specific schema that
|
||||
// omits plan/quota fields (for example, some SuperGrok accounts). An
|
||||
// incomplete observation is not proof of ineligibility; keep the account
|
||||
// routable and expose the reason for diagnostics. Operators can still
|
||||
// quarantine a known-bad account with grok_media_eligible=false.
|
||||
return true, "billing_inconclusive"
|
||||
}
|
||||
return true, "eligible"
|
||||
}
|
||||
|
||||
@@ -46,6 +46,7 @@ func TestGrokMediaGenerationEligibility(t *testing.T) {
|
||||
{name: "non grok account", account: &Account{Platform: PlatformOpenAI}, want: false, wantReason: "not_grok"},
|
||||
{name: "non oauth grok account stays eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeAPIKey}, want: true, wantReason: "non_oauth"},
|
||||
{name: "unobserved oauth fails closed", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}, want: false, wantReason: "billing_unobserved"},
|
||||
{name: "inconclusive successful billing remains eligible", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: &xai.BillingSummary{StatusCode: http.StatusOK, Partial: true}}}, want: true, wantReason: "billing_inconclusive"},
|
||||
{name: "weekly paid usage is eligible without inferring from period type", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyAllowance}}, want: true, wantReason: "eligible"},
|
||||
{name: "billing forbidden is rejected", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: forbiddenBilling}}, want: false, wantReason: "billing_forbidden"},
|
||||
{name: "weekly billing forbidden is rejected after partial success", account: &Account{Platform: PlatformGrok, Type: AccountTypeOAuth, Extra: map[string]any{grokBillingExtraKey: weeklyForbidden}}, want: false, wantReason: "billing_forbidden"},
|
||||
@@ -65,7 +66,7 @@ func TestGrokMediaGenerationEligibility(t *testing.T) {
|
||||
}
|
||||
}
|
||||
|
||||
func TestGrokMediaCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(t *testing.T) {
|
||||
func TestGrokMediaCapabilityKeepsUnobservedAndInconclusiveOAuthAsCandidates(t *testing.T) {
|
||||
unobserved := &Account{Platform: PlatformGrok, Type: AccountTypeOAuth}
|
||||
eligible, reason := unobserved.GrokMediaGenerationEligibility()
|
||||
require.False(t, eligible)
|
||||
@@ -80,7 +81,10 @@ func TestGrokMediaCapabilityKeepsOnlyUnobservedOAuthAsProbeCandidate(t *testing.
|
||||
Partial: true,
|
||||
}},
|
||||
}
|
||||
require.False(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
|
||||
eligible, reason = inconclusive.GrokMediaGenerationEligibility()
|
||||
require.True(t, eligible)
|
||||
require.Equal(t, "billing_inconclusive", reason)
|
||||
require.True(t, inconclusive.SupportsOpenAIEndpointCapability(OpenAIEndpointCapabilityGrokMediaGeneration))
|
||||
}
|
||||
|
||||
func TestGrokMediaCapabilityFiltersOnlyGeneration(t *testing.T) {
|
||||
|
||||
@@ -3,6 +3,7 @@ package service
|
||||
import (
|
||||
"context"
|
||||
"strings"
|
||||
"time"
|
||||
)
|
||||
|
||||
// resolveAccountStatsCost 计算账号统计定价费用。
|
||||
@@ -17,6 +18,7 @@ import (
|
||||
// upstreamModel 是最终发往上游的模型 ID。
|
||||
// totalCost 是本次请求的客户计费(倍率前),用于优先级 2。
|
||||
// serviceTier 是最终参与用户计费的 OpenAI 服务层级,用于优先级 3。
|
||||
// pricingAt 与本次客户计费使用同一时刻,避免跨峰谷请求的成本与售价错位。
|
||||
// reasoningEffort 是最终转发等级;Fable 5.1 max 默认按 3 倍额度消耗。
|
||||
func resolveAccountStatsCost(
|
||||
ctx context.Context,
|
||||
@@ -29,6 +31,7 @@ func resolveAccountStatsCost(
|
||||
requestCount int,
|
||||
totalCost float64,
|
||||
serviceTier string,
|
||||
pricingAt time.Time,
|
||||
reasoningEfforts ...string,
|
||||
) *float64 {
|
||||
reasoningEffort := ""
|
||||
@@ -61,7 +64,7 @@ func resolveAccountStatsCost(
|
||||
|
||||
// 优先级 3:模型定价文件(LiteLLM)默认价格
|
||||
if billingService != nil {
|
||||
return tryModelFilePricing(billingService, upstreamModel, tokens, serviceTier, reasoningEffort)
|
||||
return tryModelFilePricing(billingService, upstreamModel, tokens, serviceTier, pricingAt, reasoningEffort)
|
||||
}
|
||||
|
||||
return nil
|
||||
@@ -69,20 +72,26 @@ func resolveAccountStatsCost(
|
||||
|
||||
// tryModelFilePricing 使用模型定价文件(LiteLLM/fallback)中的价格计算费用。
|
||||
// 与用户计费共用同一条定价管线,避免这里维护第二份"单价 × token 数"实现后,
|
||||
// 每加一个定价特性都要手工镜像一次。channelPricing 为 nil,保持优先级 3 的
|
||||
// 语义:只取模型定价文件,不引入渠道自定义定价。
|
||||
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens, serviceTier string, reasoningEfforts ...string) *float64 {
|
||||
// 每加一个定价特性都要手工镜像一次。解析器不配置渠道或分组,保持优先级 3 的
|
||||
// 语义:只取模型定价文件,不引入自定义售价。
|
||||
func tryModelFilePricing(billingService *BillingService, model string, tokens UsageTokens, serviceTier string, pricingAt time.Time, reasoningEfforts ...string) *float64 {
|
||||
reasoningEffort := ""
|
||||
if len(reasoningEfforts) > 0 {
|
||||
reasoningEffort = reasoningEfforts[0]
|
||||
}
|
||||
breakdown, err := billingService.CalculateCostWithServiceTier(
|
||||
model, tokens, 1, normalizeBillingServiceTier(serviceTier),
|
||||
)
|
||||
breakdown, err := billingService.CalculateCostUnified(CostInput{
|
||||
Ctx: context.Background(),
|
||||
Model: model,
|
||||
Tokens: tokens,
|
||||
RateMultiplier: 1,
|
||||
ServiceTier: normalizeBillingServiceTier(serviceTier),
|
||||
ReasoningEffort: reasoningEffort,
|
||||
PricingAt: pricingAt,
|
||||
Resolver: NewModelPricingResolver(nil, billingService),
|
||||
})
|
||||
if err != nil || breakdown == nil || breakdown.TotalCost <= 0 {
|
||||
return nil
|
||||
}
|
||||
applyCostBreakdownMultiplier(breakdown, maxReasoningEffortBillingMultiplier(model, reasoningEffort, nil))
|
||||
return &breakdown.TotalCost
|
||||
}
|
||||
|
||||
@@ -253,6 +262,7 @@ func applyAccountStatsCost(
|
||||
upstreamModel, requestedModel string,
|
||||
tokens UsageTokens,
|
||||
totalCost float64,
|
||||
pricingAt time.Time,
|
||||
) {
|
||||
model := upstreamModel
|
||||
if model == "" {
|
||||
@@ -271,6 +281,6 @@ func applyAccountStatsCost(
|
||||
reasoningEffort = *usageLog.ReasoningEffort
|
||||
}
|
||||
usageLog.AccountStatsCost = resolveAccountStatsCost(
|
||||
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost, serviceTier, reasoningEffort,
|
||||
ctx, cs, bs, accountID, groupID, model, tokens, requestCount, totalCost, serviceTier, pricingAt, reasoningEffort,
|
||||
)
|
||||
}
|
||||
|
||||
@@ -472,7 +472,7 @@ func TestTryModelFilePricing_Success(t *testing.T) {
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
|
||||
require.InDelta(t, 0.2, *result, 1e-12)
|
||||
@@ -483,8 +483,8 @@ func TestTryModelFilePricing_Fable51MaxEffortUsesTripleQuota(t *testing.T) {
|
||||
"claude-fable-5-1": {InputPricePerToken: 0.001},
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100}
|
||||
standard := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", "xhigh")
|
||||
max := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", "max")
|
||||
standard := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, "xhigh")
|
||||
max := tryModelFilePricing(bs, "claude-fable-5-1", tokens, "", time.Time{}, "max")
|
||||
require.NotNil(t, standard)
|
||||
require.NotNil(t, max)
|
||||
require.InDelta(t, *standard*3, *max, 1e-12)
|
||||
@@ -503,7 +503,7 @@ func TestTryModelFilePricing_AppliesLongContextPricing(t *testing.T) {
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 101, OutputTokens: 10, CacheReadTokens: 5}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "")
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "", time.Time{})
|
||||
|
||||
require.NotNil(t, result)
|
||||
// Input and cache-read use the 2x input tier; output uses the 1.5x tier.
|
||||
@@ -542,7 +542,7 @@ func TestTryModelFilePricing_AppliesServiceTierPricing(t *testing.T) {
|
||||
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier)
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, tt.serviceTier, time.Time{})
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, tt.want, *result, 1e-12)
|
||||
})
|
||||
@@ -572,7 +572,7 @@ func TestTryModelFilePricing_CombinesPriorityAndLongContextPricing(t *testing.T)
|
||||
CacheReadTokens: 5,
|
||||
}
|
||||
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority")
|
||||
result := tryModelFilePricing(bs, "gpt-5.6-sol", tokens, "priority", time.Time{})
|
||||
|
||||
require.NotNil(t, result)
|
||||
// priority 单价先应用,再叠加长上下文输入 2x、输出 1.5x。
|
||||
@@ -583,7 +583,7 @@ func TestTryModelFilePricing_PricingNotFound(t *testing.T) {
|
||||
// "nonexistent-model" does not match any fallback pattern
|
||||
bs := newTestBillingServiceWithPrices(map[string]*ModelPricing{})
|
||||
tokens := UsageTokens{InputTokens: 100, OutputTokens: 50}
|
||||
result := tryModelFilePricing(bs, "nonexistent-model", tokens, "")
|
||||
result := tryModelFilePricing(bs, "nonexistent-model", tokens, "", time.Time{})
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -593,7 +593,7 @@ func TestTryModelFilePricing_NilFallback(t *testing.T) {
|
||||
"claude-sonnet-4": nil,
|
||||
})
|
||||
tokens := UsageTokens{InputTokens: 100}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -605,7 +605,7 @@ func TestTryModelFilePricing_ZeroCost(t *testing.T) {
|
||||
},
|
||||
})
|
||||
tokens := UsageTokens{} // all zero tokens → cost = 0 → nil
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
require.Nil(t, result)
|
||||
}
|
||||
|
||||
@@ -622,7 +622,7 @@ func TestTryModelFilePricing_WithImageOutput(t *testing.T) {
|
||||
OutputTokens: 50,
|
||||
ImageOutputTokens: 10,
|
||||
}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
require.NotNil(t, result)
|
||||
// ImageOutputTokens 是 OutputTokens 的子集,先扣除再按图片单价计。
|
||||
// 100*0.001 + (50-10)*0.002 + 10*0.01 = 0.1 + 0.08 + 0.1 = 0.28
|
||||
@@ -644,13 +644,111 @@ func TestTryModelFilePricing_WithCacheTokens(t *testing.T) {
|
||||
CacheCreationTokens: 200,
|
||||
CacheReadTokens: 300,
|
||||
}
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "")
|
||||
result := tryModelFilePricing(bs, "claude-sonnet-4", tokens, "", time.Time{})
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 + 200*0.003 + 300*0.0005
|
||||
// = 0.1 + 0.1 + 0.6 + 0.15 = 0.95
|
||||
require.InDelta(t, 0.95, *result, 1e-12)
|
||||
}
|
||||
|
||||
func TestTryModelFilePricing_DeepSeekPeakPricing(t *testing.T) {
|
||||
weekday := func(hour, minute int) time.Time {
|
||||
return time.Date(2026, time.August, 24, hour, minute, 0, 0, time.UTC)
|
||||
}
|
||||
for _, model := range []struct {
|
||||
name string
|
||||
input, output, cacheReadPrice float64
|
||||
}{
|
||||
{"deepseek-v4-flash", 2.2e-7, 6.6e-7, 7e-9},
|
||||
{"deepseek-v4-pro", 6.6e-7, 1.98e-6, 2.2e-8},
|
||||
} {
|
||||
for _, usage := range []struct {
|
||||
name string
|
||||
tokens UsageTokens
|
||||
}{
|
||||
{"input", UsageTokens{InputTokens: 1000}},
|
||||
{"output", UsageTokens{OutputTokens: 500}},
|
||||
{"cache_read", UsageTokens{CacheReadTokens: 1000}},
|
||||
{"mixed", UsageTokens{InputTokens: 1000, OutputTokens: 500, CacheReadTokens: 1000}},
|
||||
} {
|
||||
t.Run(model.name+"/"+usage.name, func(t *testing.T) {
|
||||
bs := newTestBillingService()
|
||||
tokens := usage.tokens
|
||||
baseCost := float64(tokens.InputTokens)*model.input +
|
||||
float64(tokens.OutputTokens)*model.output + float64(tokens.CacheReadTokens)*model.cacheReadPrice
|
||||
for _, slot := range []struct {
|
||||
name string
|
||||
at time.Time
|
||||
multiplier float64
|
||||
}{
|
||||
{"before_morning_peak", weekday(0, 59), 1},
|
||||
{"morning_peak_start", weekday(1, 0), 2},
|
||||
{"morning_peak_last_minute", weekday(3, 59), 2},
|
||||
{"morning_peak_end", weekday(4, 0), 1},
|
||||
{"afternoon_peak_start", weekday(6, 0), 2},
|
||||
{"afternoon_peak_last_minute", weekday(9, 59), 2},
|
||||
{"afternoon_peak_end", weekday(10, 0), 1},
|
||||
{"saturday", time.Date(2026, time.August, 22, 2, 0, 0, 0, time.UTC), 1},
|
||||
{"sunday", time.Date(2026, time.August, 23, 7, 0, 0, 0, time.UTC), 1},
|
||||
} {
|
||||
t.Run(slot.name, func(t *testing.T) {
|
||||
cost := tryModelFilePricing(bs, model.name, tokens, "", slot.at)
|
||||
require.NotNil(t, cost)
|
||||
require.InDelta(t, baseCost*slot.multiplier, *cost, 1e-12)
|
||||
})
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func TestResolveAccountStatsCost_DeepSeekPricingPriority(t *testing.T) {
|
||||
peak := time.Date(2026, time.August, 24, 2, 0, 0, 0, time.UTC)
|
||||
for _, tt := range []struct {
|
||||
name string
|
||||
customRule bool
|
||||
applyPricing bool
|
||||
noChannel bool
|
||||
want float64
|
||||
}{
|
||||
{name: "catalog", want: 1000 * 2.2e-7 * 2},
|
||||
{name: "custom_rule", customRule: true, want: 1},
|
||||
{name: "custom_rule_before_customer_price", customRule: true, applyPricing: true, want: 1},
|
||||
{name: "customer_price", applyPricing: true, want: 0.75},
|
||||
{name: "no_channel", noChannel: true},
|
||||
} {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
channel := &Channel{
|
||||
ID: 1, Status: StatusActive, ApplyPricingToAccountStats: tt.applyPricing,
|
||||
ModelPricing: []ChannelModelPricing{{
|
||||
Models: []string{"deepseek-v4-flash"}, InputPrice: testPtrFloat64(0.02),
|
||||
}},
|
||||
}
|
||||
if tt.customRule {
|
||||
channel.AccountStatsPricingRules = []AccountStatsPricingRule{{
|
||||
AccountIDs: []int64{1},
|
||||
Pricing: []ChannelModelPricing{{
|
||||
Models: []string{"deepseek-v4-flash"}, InputPrice: testPtrFloat64(0.001),
|
||||
}},
|
||||
}}
|
||||
}
|
||||
cs := newTestChannelServiceForStats(t, channel, 10, PlatformDeepseek)
|
||||
groupID := int64(10)
|
||||
if tt.noChannel {
|
||||
groupID = 99
|
||||
}
|
||||
cost := resolveAccountStatsCost(context.Background(), cs, newTestBillingService(),
|
||||
1, groupID, "deepseek-v4-flash", UsageTokens{InputTokens: 1000}, 1, 0.75, "", peak)
|
||||
if tt.noChannel {
|
||||
require.Nil(t, cost)
|
||||
return
|
||||
}
|
||||
require.NotNil(t, cost)
|
||||
require.InDelta(t, tt.want, *cost, 1e-12)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// resolveAccountStatsCost — integration tests covering the 4-level priority chain
|
||||
// ---------------------------------------------------------------------------
|
||||
@@ -661,7 +759,7 @@ func TestResolveAccountStatsCost_NilChannelService(t *testing.T) {
|
||||
nil, // channelService is nil
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 1, "claude-sonnet-4",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -677,7 +775,7 @@ func TestResolveAccountStatsCost_EmptyUpstreamModel(t *testing.T) {
|
||||
cs,
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 1, "", // empty upstream model
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -694,7 +792,7 @@ func TestResolveAccountStatsCost_GetChannelForGroupReturnsNil(t *testing.T) {
|
||||
cs,
|
||||
newTestBillingServiceWithPrices(map[string]*ModelPricing{}),
|
||||
1, 99, "claude-sonnet-4", // groupID 99 has no channel
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.5, "", time.Time{},
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -725,7 +823,7 @@ func TestResolveAccountStatsCost_HitsCustomRule(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, nil, // billingService not needed when custom rule hits
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 999.0, "priority", // 自定义账号价格不叠加服务层级倍率
|
||||
tokens, 1, 999.0, "priority", time.Time{}, // 自定义账号价格不叠加服务层级倍率
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.01 + 50*0.02 = 1.0 + 1.0 = 2.0
|
||||
@@ -747,7 +845,7 @@ func TestResolveAccountStatsCost_ApplyPricingToAccountStats_UsesTotalCost(t *tes
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 0.75, "priority", // 已完成用户计费,不再重复应用服务层级倍率
|
||||
tokens, 1, 0.75, "priority", time.Time{}, // 已完成用户计费,不再重复应用服务层级倍率
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 0.75, *result, 1e-12)
|
||||
@@ -765,7 +863,7 @@ func TestResolveAccountStatsCost_ApplyPricingToAccountStats_ZeroTotalCost_Return
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
UsageTokens{}, 1, 0.0, "", // totalCost = 0
|
||||
UsageTokens{}, 1, 0.0, "", time.Time{}, // totalCost = 0
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -792,7 +890,7 @@ func TestResolveAccountStatsCost_FallsBackToLiteLLM(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 999.0, "", // totalCost ignored
|
||||
tokens, 1, 999.0, "", time.Time{}, // totalCost ignored
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// 100*0.001 + 50*0.002 = 0.1 + 0.1 = 0.2
|
||||
@@ -813,7 +911,7 @@ func TestResolveAccountStatsCost_FallbackHonorsAnthropicFast(t *testing.T) {
|
||||
context.Background(), cs, bs,
|
||||
1, 10, "claude-opus-5",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000},
|
||||
1, 0, "fast",
|
||||
1, 0, "fast", time.Time{},
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 60, *result, 1e-12)
|
||||
@@ -832,7 +930,7 @@ func TestResolveAccountStatsCost_Gemini36FlashTierUsesFallbackPricing(t *testing
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "gemini-3.6-flash-low",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "",
|
||||
UsageTokens{InputTokens: 1_000_000, OutputTokens: 1_000_000, CacheReadTokens: 1_000_000}, 1, 0, "", time.Time{},
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
require.InDelta(t, 9.15, *result, 1e-12)
|
||||
@@ -856,7 +954,7 @@ func TestResolveAccountStatsCost_AllMiss_ReturnsNil(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, bs,
|
||||
1, 10, "totally-unknown-model",
|
||||
tokens, 1, 0.0, "",
|
||||
tokens, 1, 0.0, "", time.Time{},
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -873,7 +971,7 @@ func TestResolveAccountStatsCost_NilBillingService_SkipsLiteLLM(t *testing.T) {
|
||||
context.Background(),
|
||||
cs, nil, // billingService is nil
|
||||
1, 10, "claude-sonnet-4",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.0, "",
|
||||
UsageTokens{InputTokens: 100}, 1, 0.0, "", time.Time{},
|
||||
)
|
||||
require.Nil(t, result)
|
||||
}
|
||||
@@ -906,7 +1004,7 @@ func TestResolveAccountStatsCost_CustomRulePriorityOverApplyPricing(t *testing.T
|
||||
context.Background(),
|
||||
cs, nil,
|
||||
1, 10, "claude-sonnet-4",
|
||||
tokens, 1, 99.0, "", // totalCost = 99.0 (would be used if ApplyPricing wins)
|
||||
tokens, 1, 99.0, "", time.Time{}, // totalCost = 99.0 (would be used if ApplyPricing wins)
|
||||
)
|
||||
require.NotNil(t, result)
|
||||
// Custom rule: 100*0.05 = 5.0 (NOT 99.0 from totalCost)
|
||||
@@ -934,7 +1032,7 @@ func TestApplyAccountStatsCost_UsesUsageLogServiceTier(t *testing.T) {
|
||||
applyAccountStatsCost(
|
||||
context.Background(), usageLog, cs, bs,
|
||||
1, 10, "gpt-5.6-sol", "gpt-5.6-sol",
|
||||
UsageTokens{InputTokens: 100, OutputTokens: 50}, 999,
|
||||
UsageTokens{InputTokens: 100, OutputTokens: 50}, 999, time.Time{},
|
||||
)
|
||||
|
||||
require.NotNil(t, usageLog.AccountStatsCost)
|
||||
|
||||
@@ -151,6 +151,7 @@ type AccountTestService struct {
|
||||
modelMetadataRegistry map[string]modelsDevProvider
|
||||
modelMetadataRegistryAt time.Time
|
||||
pluginManager *PluginManager
|
||||
openaiGatewayService *OpenAIGatewayService
|
||||
agentIdentityTaskMu sync.Mutex
|
||||
agentIdentityWS agentIdentityWSConnectionInvalidator
|
||||
// grokWSDialer is optional; realtime account tests use the default OpenAI-style
|
||||
@@ -170,6 +171,30 @@ func (s *AccountTestService) SetPluginManager(pluginManager *PluginManager) {
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AccountTestService) SetOpenAIGatewayService(gateway *OpenAIGatewayService) {
|
||||
if s != nil {
|
||||
s.openaiGatewayService = gateway
|
||||
}
|
||||
}
|
||||
|
||||
// FetchOpenAIAccountModels uses the shared cached discovery path for the test picker.
|
||||
func (s *AccountTestService) FetchOpenAIAccountModels(ctx context.Context, account *Account) ([]openai.Model, error) {
|
||||
if s == nil || s.openaiGatewayService == nil {
|
||||
return nil, errors.New("OpenAI model discovery service is unavailable")
|
||||
}
|
||||
response, err := s.openaiGatewayService.FetchOpenAIModelsList(ctx, account)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var payload struct {
|
||||
Data []openai.Model `json:"data"`
|
||||
}
|
||||
if err := json.Unmarshal(response.Body, &payload); err != nil {
|
||||
return nil, fmt.Errorf("decode OpenAI account models: %w", err)
|
||||
}
|
||||
return payload.Data, nil
|
||||
}
|
||||
|
||||
// NewAccountTestService creates a new AccountTestService
|
||||
func NewAccountTestService(
|
||||
accountRepo AccountRepository,
|
||||
|
||||
@@ -14,6 +14,7 @@ import (
|
||||
"strings"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/logger"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
@@ -21,6 +22,11 @@ import (
|
||||
|
||||
// Account management implementations
|
||||
func (s *adminServiceImpl) ListAccounts(ctx context.Context, page, pageSize int, platform, accountType, status, search string, groupID int64, privacyMode string, sortBy, sortOrder string) ([]Account, int64, error) {
|
||||
if groupID > 0 {
|
||||
if err := s.ValidateAccountGroupBindings(ctx, []int64{groupID}); err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
}
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder}
|
||||
accounts, result, err := s.accountRepo.ListWithFilters(ctx, params, platform, accountType, status, search, groupID, privacyMode)
|
||||
if err != nil {
|
||||
@@ -284,6 +290,9 @@ func (s *adminServiceImpl) DuplicateAccount(ctx context.Context, id int64, actor
|
||||
}
|
||||
autoPauseOnExpired := source.AutoPauseOnExpired
|
||||
groups, groupIDs := duplicateAccountGroups(source)
|
||||
if err := s.ValidateAccountGroupBindings(ctx, groupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
proxyID := source.ProxyID
|
||||
if source.ProxyFallbackOriginID != nil {
|
||||
// Proxy fallback is transient runtime state; duplicate the configured origin.
|
||||
@@ -512,6 +521,9 @@ func (s *adminServiceImpl) CreateAccount(ctx context.Context, input *CreateAccou
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.ValidateAccountGroupBindings(ctx, groupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.accountRepo.Create(ctx, account); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -804,6 +816,9 @@ func (s *adminServiceImpl) UpdateAccount(ctx context.Context, id int64, input *U
|
||||
if err := s.validateGroupIDsExist(ctx, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.ValidateAccountGroupBindings(ctx, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 检查混合渠道风险(除非用户已确认)
|
||||
if !input.SkipMixedChannelCheck {
|
||||
@@ -935,6 +950,9 @@ func (s *adminServiceImpl) BulkUpdateAccounts(ctx context.Context, input *BulkUp
|
||||
if err := s.validateGroupIDsExist(ctx, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.ValidateAccountGroupBindings(ctx, *input.GroupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
openAISettings, err := normalizeBulkOpenAISettings(input)
|
||||
if err != nil {
|
||||
@@ -1352,6 +1370,9 @@ func (s *adminServiceImpl) CreateShadow(ctx context.Context, parentID int64, opt
|
||||
}
|
||||
}
|
||||
}
|
||||
if err := s.ValidateAccountGroupBindings(ctx, groupIDs); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
// 4. 构造影子账号(安全不变量:Credentials 恒不含 auth token,仅含 model_mapping)。
|
||||
// name 为空时默认 "<母账号名> (Spark)"——否则空 name 会在 ent(name NotEmpty)处变成裸 500
|
||||
@@ -1524,6 +1545,35 @@ func (s *adminServiceImpl) validateGroupIDsExist(ctx context.Context, groupIDs [
|
||||
return nil
|
||||
}
|
||||
|
||||
// ValidateAccountGroupBindings is the shared fail-closed policy boundary for
|
||||
// every account path that accepts explicit group bindings.
|
||||
func (s *adminServiceImpl) ValidateAccountGroupBindings(ctx context.Context, groupIDs []int64) error {
|
||||
if len(groupIDs) == 0 || s.cfg == nil || s.cfg.RunMode != config.RunModeSimple {
|
||||
return nil
|
||||
}
|
||||
if s.groupRepo == nil {
|
||||
return errors.New("group repository not configured")
|
||||
}
|
||||
seen := make(map[int64]struct{}, len(groupIDs))
|
||||
for _, groupID := range groupIDs {
|
||||
if groupID <= 0 {
|
||||
return fmt.Errorf("get group: %w", ErrGroupNotFound)
|
||||
}
|
||||
if _, ok := seen[groupID]; ok {
|
||||
continue
|
||||
}
|
||||
seen[groupID] = struct{}{}
|
||||
group, err := s.groupRepo.GetByIDLite(ctx, groupID)
|
||||
if err != nil {
|
||||
return fmt.Errorf("get group: %w", err)
|
||||
}
|
||||
if !IsGroupBindableInSimpleMode(group) {
|
||||
return infraerrors.BadRequest("SIMPLE_MODE_GROUP_NOT_BINDABLE", "composite groups cannot be bound in simple mode")
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CheckMixedChannelRisk checks whether target groups contain mixed channels for the current account platform.
|
||||
func (s *adminServiceImpl) CheckMixedChannelRisk(ctx context.Context, currentAccountID int64, currentAccountPlatform string, groupIDs []int64) error {
|
||||
return s.checkMixedChannelRisk(ctx, currentAccountID, currentAccountPlatform, groupIDs)
|
||||
|
||||
@@ -9,6 +9,7 @@ import (
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/claude"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
@@ -20,9 +21,26 @@ import (
|
||||
)
|
||||
|
||||
// Group management implementations
|
||||
func (s *adminServiceImpl) ValidateSimpleModeGroupOperation(operation AdminGroupOperation) error {
|
||||
return ValidateSimpleModeGroupOperation(s.cfg, operation)
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) ListGroups(ctx context.Context, page, pageSize int, platform, status, search string, isExclusive *bool, sortBy, sortOrder string) ([]Group, int64, error) {
|
||||
params := pagination.PaginationParams{Page: page, PageSize: pageSize, SortBy: sortBy, SortOrder: sortOrder}
|
||||
groups, result, err := s.groupRepo.ListWithFilters(ctx, params, platform, status, search, isExclusive)
|
||||
var groups []Group
|
||||
var result *pagination.PaginationResult
|
||||
var err error
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
repo, ok := s.groupRepo.(interface {
|
||||
ListBindableWithFilters(context.Context, pagination.PaginationParams, string, string, string, *bool) ([]Group, *pagination.PaginationResult, error)
|
||||
})
|
||||
if !ok {
|
||||
return nil, 0, errors.New("group repository does not support simple-mode filtering")
|
||||
}
|
||||
groups, result, err = repo.ListBindableWithFilters(ctx, params, platform, status, search, isExclusive)
|
||||
} else {
|
||||
groups, result, err = s.groupRepo.ListWithFilters(ctx, params, platform, status, search, isExclusive)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
@@ -45,7 +63,21 @@ func (s *adminServiceImpl) GetAllGroupsIncludingInactive(ctx context.Context) ([
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) GetGroup(ctx context.Context, id int64) (*Group, error) {
|
||||
return s.groupRepo.GetByID(ctx, id)
|
||||
group, err := s.groupRepo.GetByID(ctx, id)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.validateSimpleModeGroupAccess(group); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return group, nil
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) validateSimpleModeGroupAccess(group *Group) error {
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple && !IsGroupBindableInSimpleMode(group) {
|
||||
return infraerrors.BadRequest("SIMPLE_MODE_GROUP_NOT_BINDABLE", "composite groups are not supported in simple mode")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) GetGroupModelsListCandidates(ctx context.Context, id int64, platform string) ([]string, error) {
|
||||
@@ -101,6 +133,9 @@ func (s *adminServiceImpl) GetGroupModelsListCandidates(ctx context.Context, id
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) ListCompositeRoutes(ctx context.Context, groupID int64) ([]CompositeModelRoute, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationCompositeRoute); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.requireCompositeGroup(ctx, groupID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -111,6 +146,9 @@ func (s *adminServiceImpl) ListCompositeRoutes(ctx context.Context, groupID int6
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) CreateCompositeRoute(ctx context.Context, groupID int64, input CompositeRouteInput) (*CompositeModelRoute, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationCompositeRoute); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.requireCompositeGroup(ctx, groupID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -128,6 +166,9 @@ func (s *adminServiceImpl) CreateCompositeRoute(ctx context.Context, groupID int
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) UpdateCompositeRoute(ctx context.Context, groupID, routeID int64, input CompositeRouteInput) (*CompositeModelRoute, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationCompositeRoute); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.requireCompositeGroup(ctx, groupID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -151,6 +192,9 @@ func (s *adminServiceImpl) UpdateCompositeRoute(ctx context.Context, groupID, ro
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) DeleteCompositeRoute(ctx context.Context, groupID, routeID int64) error {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationCompositeRoute); err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.requireCompositeGroup(ctx, groupID); err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -166,6 +210,9 @@ func (s *adminServiceImpl) DeleteCompositeRoute(ctx context.Context, groupID, ro
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) PreviewCompositeRoute(ctx context.Context, groupID int64, input CompositeRoutePreviewRequest) (*CompositeRouteDecision, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationCompositeRoute); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.requireCompositeGroup(ctx, groupID); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
@@ -308,7 +355,30 @@ func sanitizeGroupOpenAIFast(group *Group) {
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeCreateGroupInputForSimpleMode(input *CreateGroupInput) {
|
||||
if input == nil {
|
||||
return
|
||||
}
|
||||
*input = CreateGroupInput{
|
||||
Name: input.Name, Description: input.Description, Platform: input.Platform,
|
||||
RateMultiplier: 1, SubscriptionType: SubscriptionTypeStandard,
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeUpdateGroupInputForSimpleMode(input *UpdateGroupInput) {
|
||||
if input == nil {
|
||||
return
|
||||
}
|
||||
*input = UpdateGroupInput{Name: input.Name, Description: input.Description}
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupInput) (*Group, error) {
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple && NormalizeGroupPlatform(input.Platform) == PlatformComposite {
|
||||
return nil, infraerrors.BadRequest("SIMPLE_MODE_GROUP_NOT_BINDABLE", "composite groups are not supported in simple mode")
|
||||
}
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
normalizeCreateGroupInputForSimpleMode(input)
|
||||
}
|
||||
if input.RateMultiplier <= 0 {
|
||||
return nil, errors.New("rate_multiplier must be > 0")
|
||||
}
|
||||
@@ -474,6 +544,12 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
|
||||
}
|
||||
}
|
||||
|
||||
// 白名单在创建路径同样收口:开启但为空、通配位置非法都会 400。
|
||||
modelAllowlist, err := normalizeGroupModelAllowlist(input.ModelAllowlist)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
|
||||
group := &Group{
|
||||
Name: input.Name,
|
||||
Description: input.Description,
|
||||
@@ -528,7 +604,7 @@ func (s *adminServiceImpl) CreateGroup(ctx context.Context, input *CreateGroupIn
|
||||
RequirePrivacySet: input.RequirePrivacySet,
|
||||
DefaultMappedModel: input.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: normalizeOpenAIMessagesDispatchModelConfig(input.MessagesDispatchModelConfig),
|
||||
ModelsListConfig: normalizeGroupModelsListConfig(input.ModelsListConfig),
|
||||
ModelAllowlist: modelAllowlist,
|
||||
// 固定账号 manifest 配置:账号绑定发生在分组创建之后,创建路径禁止开启,
|
||||
// 成员关系无从校验(前端创建对话框也不展示)。
|
||||
CodexModelsManifestConfig: normalizeCodexModelsManifestConfig(platform, input.CodexModelsManifestConfig),
|
||||
@@ -669,6 +745,15 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if err := s.validateSimpleModeGroupAccess(group); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple && input.Platform == PlatformComposite {
|
||||
return nil, infraerrors.BadRequest("SIMPLE_MODE_GROUP_NOT_BINDABLE", "composite groups are not supported in simple mode")
|
||||
}
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
normalizeUpdateGroupInputForSimpleMode(input)
|
||||
}
|
||||
|
||||
// 渠道缓存里存了 groupID → platform 的映射,改了平台要让它失效(见函数末尾)
|
||||
previousPlatform := group.Platform
|
||||
@@ -909,8 +994,12 @@ func (s *adminServiceImpl) UpdateGroup(ctx context.Context, id int64, input *Upd
|
||||
if input.MessagesDispatchModelConfig != nil {
|
||||
group.MessagesDispatchModelConfig = normalizeOpenAIMessagesDispatchModelConfig(*input.MessagesDispatchModelConfig)
|
||||
}
|
||||
if input.ModelsListConfig != nil {
|
||||
group.ModelsListConfig = normalizeGroupModelsListConfig(*input.ModelsListConfig)
|
||||
if input.ModelAllowlist != nil {
|
||||
modelAllowlist, err := normalizeGroupModelAllowlist(*input.ModelAllowlist)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
group.ModelAllowlist = modelAllowlist
|
||||
}
|
||||
if input.CodexModelsManifestConfig != nil {
|
||||
group.CodexModelsManifestConfig = *input.CodexModelsManifestConfig
|
||||
@@ -1070,6 +1159,27 @@ func normalizeGroupModelPricing(platform string, pricing []ChannelModelPricing)
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) DeleteGroup(ctx context.Context, id int64) error {
|
||||
return s.deleteGroup(ctx, id, false)
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) DeleteGroupIfEmpty(ctx context.Context, id int64) error {
|
||||
return s.deleteGroup(ctx, id, true)
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) deleteGroup(ctx context.Context, id int64, requireEmpty bool) error {
|
||||
if s.cfg != nil && s.cfg.RunMode == config.RunModeSimple {
|
||||
group, err := s.groupRepo.GetByIDLite(ctx, id)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
if err := s.validateSimpleModeGroupAccess(group); err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
if requireEmpty && s.emptyGroupDeleteRepo == nil {
|
||||
return fmt.Errorf("guarded group deletion is unavailable")
|
||||
}
|
||||
|
||||
var groupKeys []string
|
||||
if s.authCacheInvalidator != nil {
|
||||
keys, err := s.apiKeyRepo.ListKeysByGroupID(ctx, id)
|
||||
@@ -1078,7 +1188,13 @@ func (s *adminServiceImpl) DeleteGroup(ctx context.Context, id int64) error {
|
||||
}
|
||||
}
|
||||
|
||||
affectedUserIDs, err := s.groupRepo.DeleteCascade(ctx, id)
|
||||
var affectedUserIDs []int64
|
||||
var err error
|
||||
if requireEmpty {
|
||||
affectedUserIDs, err = s.emptyGroupDeleteRepo.DeleteCascadeIfEmpty(ctx, id)
|
||||
} else {
|
||||
affectedUserIDs, err = s.groupRepo.DeleteCascade(ctx, id)
|
||||
}
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
@@ -1116,6 +1232,9 @@ func (s *adminServiceImpl) GetGroupAPIKeys(ctx context.Context, groupID int64, p
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) GetGroupRateMultipliers(ctx context.Context, groupID int64) ([]UserGroupRateEntry, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationMultiplier); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
if s.userGroupRateRepo == nil {
|
||||
return nil, nil
|
||||
}
|
||||
@@ -1123,6 +1242,9 @@ func (s *adminServiceImpl) GetGroupRateMultipliers(ctx context.Context, groupID
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) ClearGroupRateMultipliers(ctx context.Context, groupID int64) error {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationMultiplier); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.userGroupRateRepo == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -1130,6 +1252,9 @@ func (s *adminServiceImpl) ClearGroupRateMultipliers(ctx context.Context, groupI
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) BatchSetGroupRateMultipliers(ctx context.Context, groupID int64, entries []GroupRateMultiplierInput) error {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationMultiplier); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.userGroupRateRepo == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -1142,6 +1267,9 @@ func (s *adminServiceImpl) BatchSetGroupRateMultipliers(ctx context.Context, gro
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) ClearGroupRPMOverrides(ctx context.Context, groupID int64) error {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationRPMOverride); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.userGroupRateRepo == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -1156,6 +1284,9 @@ func (s *adminServiceImpl) ClearGroupRPMOverrides(ctx context.Context, groupID i
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) BatchSetGroupRPMOverrides(ctx context.Context, groupID int64, entries []GroupRPMOverrideInput) error {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationRPMOverride); err != nil {
|
||||
return err
|
||||
}
|
||||
if s.userGroupRateRepo == nil {
|
||||
return nil
|
||||
}
|
||||
@@ -1175,6 +1306,9 @@ func (s *adminServiceImpl) BatchSetGroupRPMOverrides(ctx context.Context, groupI
|
||||
}
|
||||
|
||||
func (s *adminServiceImpl) UpdateGroupSortOrders(ctx context.Context, updates []GroupSortOrderUpdate) error {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationSort); err != nil {
|
||||
return err
|
||||
}
|
||||
return s.groupRepo.UpdateSortOrders(ctx, updates)
|
||||
}
|
||||
|
||||
|
||||
@@ -150,9 +150,9 @@ func cloneGroupForDuplicate(source *Group, operationID string) *Group {
|
||||
RequirePrivacySet: source.RequirePrivacySet,
|
||||
DefaultMappedModel: source.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: cloneGroupMessagesDispatchModelConfig(source.MessagesDispatchModelConfig),
|
||||
ModelsListConfig: GroupModelsListConfig{
|
||||
Enabled: source.ModelsListConfig.Enabled,
|
||||
Models: append([]string(nil), source.ModelsListConfig.Models...),
|
||||
ModelAllowlist: GroupModelAllowlist{
|
||||
Enabled: source.ModelAllowlist.Enabled,
|
||||
Models: append([]string(nil), source.ModelAllowlist.Models...),
|
||||
},
|
||||
// 固定账号 manifest 配置指向源分组的账号 ID,复制后成员关系可能变化,重置为关闭且列表为空。
|
||||
CodexModelsManifestConfig: GroupCodexModelsManifestConfig{},
|
||||
@@ -166,6 +166,9 @@ func cloneGroupForDuplicate(source *Group, operationID string) *Group {
|
||||
// RecoverDuplicateGroup performs a read-only lookup for a copy that was already
|
||||
// committed for the same actor, source group, and idempotency key.
|
||||
func (s *adminServiceImpl) RecoverDuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationDuplicate); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
operationID := duplicateGroupOperationID(id, actorScope, operationKey)
|
||||
if operationID == "" {
|
||||
return nil, nil
|
||||
@@ -191,6 +194,9 @@ func (s *adminServiceImpl) RecoverDuplicateGroup(ctx context.Context, id int64,
|
||||
// account priorities. The repository commits the group, bindings, and outbox
|
||||
// event atomically so a failed binding never leaves an orphan group.
|
||||
func (s *adminServiceImpl) DuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error) {
|
||||
if err := s.ValidateSimpleModeGroupOperation(AdminGroupOperationDuplicate); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
existing, err := s.RecoverDuplicateGroup(ctx, id, actorScope, operationKey)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
|
||||
@@ -61,7 +61,7 @@ func cloneGroupForDuplicateTest(group *Group) *Group {
|
||||
cloned.ModelRouting = cloneGroupModelRouting(group.ModelRouting)
|
||||
cloned.SupportedModelScopes = append([]string(nil), group.SupportedModelScopes...)
|
||||
cloned.MessagesDispatchModelConfig = cloneGroupMessagesDispatchModelConfig(group.MessagesDispatchModelConfig)
|
||||
cloned.ModelsListConfig.Models = append([]string(nil), group.ModelsListConfig.Models...)
|
||||
cloned.ModelAllowlist.Models = append([]string(nil), group.ModelAllowlist.Models...)
|
||||
cloned.AccountGroups = append([]AccountGroup(nil), group.AccountGroups...)
|
||||
return &cloned
|
||||
}
|
||||
@@ -177,7 +177,7 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
|
||||
HaikuMappedModel: "gpt-5-mini",
|
||||
ExactModelMappings: map[string]string{"claude-special": "gpt-special"},
|
||||
},
|
||||
ModelsListConfig: GroupModelsListConfig{Enabled: true, Models: []string{"gpt-5.4", "gpt-5-mini"}},
|
||||
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.4", "gpt-5-mini"}},
|
||||
RPMLimit: 99,
|
||||
MaxReasoningEffort: "medium",
|
||||
MaxReasoningEffortOverLimit: ReasoningEffortOverLimitDeny,
|
||||
@@ -217,7 +217,7 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
|
||||
require.Equal(t, source.MessagesDispatchModelConfig, duplicate.MessagesDispatchModelConfig)
|
||||
require.Equal(t, source.ForceOpenAIFast, duplicate.ForceOpenAIFast)
|
||||
require.Equal(t, source.FreeOpenAIFast, duplicate.FreeOpenAIFast)
|
||||
require.Equal(t, source.ModelsListConfig, duplicate.ModelsListConfig)
|
||||
require.Equal(t, source.ModelAllowlist, duplicate.ModelAllowlist)
|
||||
require.Equal(t, source.RPMLimit, duplicate.RPMLimit)
|
||||
require.Equal(t, source.MaxReasoningEffort, duplicate.MaxReasoningEffort)
|
||||
require.Equal(t, source.MaxReasoningEffortOverLimit, duplicate.MaxReasoningEffortOverLimit)
|
||||
@@ -235,14 +235,14 @@ func TestDuplicateGroupCopiesConfigurationDeeplyAndResetsRuntimeState(t *testing
|
||||
duplicate.VideoModelPrices[VideoPriceFamilyGrokImagineVideo15][VideoBillingResolution720P] = 999
|
||||
duplicate.SupportedModelScopes[0] = "changed"
|
||||
duplicate.MessagesDispatchModelConfig.ExactModelMappings["claude-special"] = "changed"
|
||||
duplicate.ModelsListConfig.Models[0] = "changed"
|
||||
duplicate.ModelAllowlist.Models[0] = "changed"
|
||||
duplicate.ReasoningEffortMappings[0].To = "changed"
|
||||
*duplicate.DailyLimitUSD = 999
|
||||
require.Equal(t, int64(13), source.ModelRouting["gpt-*"][0])
|
||||
require.Equal(t, 0.14, source.VideoModelPrices[VideoPriceFamilyGrokImagineVideo15][VideoBillingResolution720P])
|
||||
require.Equal(t, "claude", source.SupportedModelScopes[0])
|
||||
require.Equal(t, "gpt-special", source.MessagesDispatchModelConfig.ExactModelMappings["claude-special"])
|
||||
require.Equal(t, "gpt-5.4", source.ModelsListConfig.Models[0])
|
||||
require.Equal(t, "gpt-5.4", source.ModelAllowlist.Models[0])
|
||||
require.Equal(t, "xhigh", source.ReasoningEffortMappings[0].To)
|
||||
require.Equal(t, 11.0, *source.DailyLimitUSD)
|
||||
}
|
||||
|
||||
@@ -6,6 +6,7 @@ import (
|
||||
"time"
|
||||
|
||||
dbent "github.com/Wei-Shaw/sub2api/ent"
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
)
|
||||
|
||||
@@ -48,6 +49,7 @@ type AdminService interface {
|
||||
RecoverDuplicateGroup(ctx context.Context, id int64, actorScope, operationKey string) (*Group, error)
|
||||
UpdateGroup(ctx context.Context, id int64, input *UpdateGroupInput) (*Group, error)
|
||||
DeleteGroup(ctx context.Context, id int64) error
|
||||
DeleteGroupIfEmpty(ctx context.Context, id int64) error
|
||||
ListCompositeRoutes(ctx context.Context, groupID int64) ([]CompositeModelRoute, error)
|
||||
CreateCompositeRoute(ctx context.Context, groupID int64, input CompositeRouteInput) (*CompositeModelRoute, error)
|
||||
UpdateCompositeRoute(ctx context.Context, groupID, routeID int64, input CompositeRouteInput) (*CompositeModelRoute, error)
|
||||
@@ -79,6 +81,7 @@ type AdminService interface {
|
||||
GetAccount(ctx context.Context, id int64) (*Account, error)
|
||||
GetAccountsByIDs(ctx context.Context, ids []int64) ([]*Account, error)
|
||||
CreateAccount(ctx context.Context, input *CreateAccountInput) (*Account, error)
|
||||
ValidateAccountGroupBindings(ctx context.Context, groupIDs []int64) error
|
||||
// DuplicateAccount creates an independent account from an existing account's configuration.
|
||||
// First-class runtime columns are intentionally reset by the normal account creation path.
|
||||
DuplicateAccount(ctx context.Context, id int64, actorScope, operationKey string) (*Account, error)
|
||||
@@ -137,6 +140,24 @@ type AdminService interface {
|
||||
ResetAccountQuota(ctx context.Context, id int64) error
|
||||
}
|
||||
|
||||
type AdminGroupOperation string
|
||||
|
||||
const (
|
||||
AdminGroupOperationBasic AdminGroupOperation = "basic"
|
||||
AdminGroupOperationDuplicate AdminGroupOperation = "duplicate"
|
||||
AdminGroupOperationCompositeRoute AdminGroupOperation = "composite_route"
|
||||
AdminGroupOperationMultiplier AdminGroupOperation = "multiplier"
|
||||
AdminGroupOperationRPMOverride AdminGroupOperation = "rpm_override"
|
||||
AdminGroupOperationSort AdminGroupOperation = "sort"
|
||||
)
|
||||
|
||||
func ValidateSimpleModeGroupOperation(cfg *config.Config, operation AdminGroupOperation) error {
|
||||
if cfg != nil && cfg.RunMode == config.RunModeSimple && operation != AdminGroupOperationBasic {
|
||||
return infraerrors.New(http.StatusForbidden, "SIMPLE_MODE_OPERATION_UNSUPPORTED", "This operation is not supported in simple mode")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// CreateUserInput represents input for creating a new user via admin operations.
|
||||
type CreateUserInput struct {
|
||||
Email string
|
||||
@@ -272,7 +293,7 @@ type CreateGroupInput struct {
|
||||
RequireOAuthOnly bool
|
||||
RequirePrivacySet bool
|
||||
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig
|
||||
ModelsListConfig GroupModelsListConfig
|
||||
ModelAllowlist GroupModelAllowlist
|
||||
// CodexModelsManifestConfig 固定账号 manifest 配置;创建路径禁止开启,仅编辑可配置。
|
||||
CodexModelsManifestConfig GroupCodexModelsManifestConfig
|
||||
// RPMLimit 分组 RPM 上限(0 = 不限制)
|
||||
@@ -353,7 +374,7 @@ type UpdateGroupInput struct {
|
||||
RequireOAuthOnly *bool
|
||||
RequirePrivacySet *bool
|
||||
MessagesDispatchModelConfig *OpenAIMessagesDispatchModelConfig
|
||||
ModelsListConfig *GroupModelsListConfig
|
||||
ModelAllowlist *GroupModelAllowlist
|
||||
// CodexModelsManifestConfig nil 表示不修改;非 openai 平台会被归一化为关闭。
|
||||
CodexModelsManifestConfig *GroupCodexModelsManifestConfig
|
||||
// RPMLimit 分组 RPM 上限(0 = 不限制),nil 表示未提供不改动。
|
||||
@@ -658,9 +679,11 @@ var ErrRPMStatusUnavailable = infraerrors.New(http.StatusNotImplemented, "RPM_ST
|
||||
|
||||
// adminServiceImpl implements AdminService
|
||||
type adminServiceImpl struct {
|
||||
cfg *config.Config
|
||||
userRepo UserRepository
|
||||
groupRepo GroupRepository
|
||||
groupDuplicateRepo GroupDuplicateRepository
|
||||
emptyGroupDeleteRepo EmptyGroupDeleteRepository
|
||||
accountRepo AccountRepository
|
||||
accountDuplicateRepo AccountDuplicateRepository
|
||||
accountBillingRepo AccountBillingSettingsRepository
|
||||
@@ -702,6 +725,7 @@ type userGroupRateBatchReader interface {
|
||||
|
||||
// NewAdminService creates a new AdminService
|
||||
func NewAdminService(
|
||||
cfg *config.Config,
|
||||
userRepo UserRepository,
|
||||
groupRepo AdminGroupRepository,
|
||||
accountRepo AdminAccountRepository,
|
||||
@@ -726,9 +750,11 @@ func NewAdminService(
|
||||
channelCacheInvalidator ChannelCacheInvalidator,
|
||||
) AdminService {
|
||||
return &adminServiceImpl{
|
||||
cfg: cfg,
|
||||
userRepo: userRepo,
|
||||
groupRepo: groupRepo,
|
||||
groupDuplicateRepo: groupRepo,
|
||||
emptyGroupDeleteRepo: groupRepo,
|
||||
accountRepo: accountRepo,
|
||||
accountDuplicateRepo: accountRepo,
|
||||
accountBillingRepo: accountRepo,
|
||||
|
||||
@@ -247,6 +247,7 @@ type groupRepoStub struct {
|
||||
affectedUserIDs []int64
|
||||
deleteErr error
|
||||
deleteCalls []int64
|
||||
guardedCalls []int64
|
||||
}
|
||||
|
||||
func (s *groupRepoStub) Create(ctx context.Context, group *Group) error {
|
||||
@@ -274,6 +275,11 @@ func (s *groupRepoStub) DeleteCascade(ctx context.Context, id int64) ([]int64, e
|
||||
return s.affectedUserIDs, s.deleteErr
|
||||
}
|
||||
|
||||
func (s *groupRepoStub) DeleteCascadeIfEmpty(ctx context.Context, id int64) ([]int64, error) {
|
||||
s.guardedCalls = append(s.guardedCalls, id)
|
||||
return s.affectedUserIDs, s.deleteErr
|
||||
}
|
||||
|
||||
func (s *groupRepoStub) List(ctx context.Context, params pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
|
||||
panic("unexpected List call")
|
||||
}
|
||||
@@ -703,6 +709,22 @@ func TestAdminService_DeleteGroup_Error(t *testing.T) {
|
||||
require.ErrorIs(t, err, deleteErr)
|
||||
}
|
||||
|
||||
func TestAdminService_DeleteGroupIfEmpty_UsesGuardedCascade(t *testing.T) {
|
||||
repo := &groupRepoStub{deleteErr: ErrGroupNotEmpty}
|
||||
svc := &adminServiceImpl{groupRepo: repo, emptyGroupDeleteRepo: repo}
|
||||
|
||||
err := svc.DeleteGroupIfEmpty(context.Background(), 42)
|
||||
require.ErrorIs(t, err, ErrGroupNotEmpty)
|
||||
require.Equal(t, []int64{42}, repo.guardedCalls)
|
||||
require.Empty(t, repo.deleteCalls)
|
||||
}
|
||||
|
||||
func TestAdminService_DeleteGroupIfEmpty_MissingCapabilityReturnsError(t *testing.T) {
|
||||
svc := &adminServiceImpl{groupRepo: &groupRepoStub{}}
|
||||
|
||||
require.ErrorContains(t, svc.DeleteGroupIfEmpty(context.Background(), 42), "guarded group deletion is unavailable")
|
||||
}
|
||||
|
||||
func TestAdminService_DeleteProxy_Success(t *testing.T) {
|
||||
repo := &proxyRepoStub{}
|
||||
svc := &adminServiceImpl{proxyRepo: repo}
|
||||
|
||||
@@ -11,6 +11,7 @@ import (
|
||||
"time"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
@@ -266,6 +267,28 @@ func TestDuplicateAccountPreservesUngroupedState(t *testing.T) {
|
||||
require.NotContains(t, repo.groupsOf, duplicate.ID)
|
||||
}
|
||||
|
||||
func TestDuplicateAccountSimpleModeRejectsCompositeGroupBinding(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
repo := newDuplicateAccountRepoStub()
|
||||
groupRepo := &groupRepoStubForAdmin{getByIDByID: map[int64]*Group{
|
||||
9: {ID: 9, Platform: PlatformComposite},
|
||||
}}
|
||||
svc := &adminServiceImpl{
|
||||
cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: groupRepo,
|
||||
accountRepo: repo, accountDuplicateRepo: repo,
|
||||
}
|
||||
source := &Account{
|
||||
Name: "composite-bound", Platform: PlatformAnthropic, Type: AccountTypeAPIKey,
|
||||
Credentials: map[string]any{"api_key": "secret"}, GroupIDs: []int64{9},
|
||||
}
|
||||
require.NoError(t, repo.Create(ctx, source))
|
||||
|
||||
_, err := svc.DuplicateAccount(ctx, source.ID, "admin:1", "")
|
||||
|
||||
require.Equal(t, "SIMPLE_MODE_GROUP_NOT_BINDABLE", infraerrors.Reason(err))
|
||||
require.Len(t, repo.accounts, 1)
|
||||
}
|
||||
|
||||
func TestDuplicateAccountAtomicCreateFailureLeavesNoOrphan(t *testing.T) {
|
||||
ctx := context.Background()
|
||||
repo := newDuplicateAccountRepoStub()
|
||||
|
||||
@@ -0,0 +1,125 @@
|
||||
//go:build unit
|
||||
|
||||
package service
|
||||
|
||||
import (
|
||||
"context"
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestAdminService_CreateGroup_RejectsEmptyEnabledModelAllowlist(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{createID: 51}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
|
||||
Name: "allowlist-empty-group",
|
||||
Platform: PlatformOpenAI,
|
||||
RateMultiplier: 1,
|
||||
ModelAllowlist: GroupModelAllowlist{Enabled: true},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
|
||||
require.Nil(t, repo.created, "拒绝时不得落库")
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_RejectsInvalidAllowlistWildcard(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{createID: 51}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
|
||||
Name: "allowlist-wildcard-group",
|
||||
Platform: PlatformOpenAI,
|
||||
RateMultiplier: 1,
|
||||
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-*-5.4"}},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
|
||||
require.Nil(t, repo.created, "拒绝时不得落库")
|
||||
}
|
||||
|
||||
func TestAdminService_CreateGroup_NormalizesModelAllowlist(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{createID: 52}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{
|
||||
Name: "allowlist-normalized-group",
|
||||
Platform: PlatformOpenAI,
|
||||
RateMultiplier: 1,
|
||||
ModelAllowlist: GroupModelAllowlist{
|
||||
Enabled: true,
|
||||
Models: []string{" gpt-5.4 ", "GPT-5.4", "claude-*", " "},
|
||||
},
|
||||
})
|
||||
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.created)
|
||||
require.True(t, repo.created.ModelAllowlist.Enabled)
|
||||
require.Equal(t, []string{"gpt-5.4", "claude-*"}, repo.created.ModelAllowlist.Models)
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateGroup_RejectsEmptyEnabledModelAllowlist(t *testing.T) {
|
||||
existing := &Group{ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive}
|
||||
repo := &groupRepoStubForAdmin{getByID: existing}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{
|
||||
ModelAllowlist: &GroupModelAllowlist{Enabled: true},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
|
||||
require.Nil(t, repo.updated, "拒绝时不得落库")
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateGroup_RejectsInvalidAllowlistWildcard(t *testing.T) {
|
||||
existing := &Group{ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive}
|
||||
repo := &groupRepoStubForAdmin{getByID: existing}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
_, err := svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{
|
||||
ModelAllowlist: &GroupModelAllowlist{Enabled: true, Models: []string{"foo-*bar"}},
|
||||
})
|
||||
|
||||
require.Error(t, err)
|
||||
appErr := infraerrors.FromError(err)
|
||||
require.Equal(t, int32(http.StatusBadRequest), appErr.Code)
|
||||
require.Equal(t, "INVALID_MODEL_ALLOWLIST", appErr.Reason)
|
||||
require.Nil(t, repo.updated, "拒绝时不得落库")
|
||||
}
|
||||
|
||||
func TestAdminService_UpdateGroup_NormalizesAndResetsModelAllowlist(t *testing.T) {
|
||||
existing := &Group{
|
||||
ID: 1, Name: "existing", Platform: PlatformOpenAI, Status: StatusActive,
|
||||
ModelAllowlist: GroupModelAllowlist{Enabled: true, Models: []string{"gpt-5.4"}},
|
||||
}
|
||||
repo := &groupRepoStubForAdmin{getByID: existing}
|
||||
svc := &adminServiceImpl{groupRepo: repo}
|
||||
|
||||
// 归一化:按小写去重保序(保留首次出现的原始拼写)。
|
||||
updated := GroupModelAllowlist{Enabled: true, Models: []string{" GPT-5.4 ", "claude-*"}}
|
||||
_, err := svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{ModelAllowlist: &updated})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.updated)
|
||||
require.Equal(t, []string{"GPT-5.4", "claude-*"}, repo.updated.ModelAllowlist.Models)
|
||||
|
||||
// 关闭且清空条目也应被接受(关闭白名单)。
|
||||
repo.updated = nil
|
||||
disabled := GroupModelAllowlist{Enabled: false}
|
||||
_, err = svc.UpdateGroup(context.Background(), existing.ID, &UpdateGroupInput{ModelAllowlist: &disabled})
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, repo.updated)
|
||||
require.False(t, repo.updated.ModelAllowlist.Enabled)
|
||||
require.Empty(t, repo.updated.ModelAllowlist.Models)
|
||||
}
|
||||
@@ -7,6 +7,7 @@ import (
|
||||
"net/http"
|
||||
"testing"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/pagination"
|
||||
"github.com/stretchr/testify/require"
|
||||
@@ -89,6 +90,10 @@ func (s *groupRepoStubForAdmin) DeleteCascade(_ context.Context, _ int64) ([]int
|
||||
panic("unexpected DeleteCascade call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForAdmin) DeleteCascadeIfEmpty(_ context.Context, _ int64) ([]int64, error) {
|
||||
panic("unexpected DeleteCascadeIfEmpty call")
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForAdmin) List(_ context.Context, _ pagination.PaginationParams) ([]Group, *pagination.PaginationResult, error) {
|
||||
panic("unexpected List call")
|
||||
}
|
||||
@@ -117,6 +122,178 @@ func (s *groupRepoStubForAdmin) ListWithFilters(_ context.Context, params pagina
|
||||
return s.listWithFiltersGroups, result, nil
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForAdmin) ListBindableWithFilters(ctx context.Context, params pagination.PaginationParams, platform, status, search string, isExclusive *bool) ([]Group, *pagination.PaginationResult, error) {
|
||||
return s.ListWithFilters(ctx, params, platform, status, search, isExclusive)
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeValidatesRequestedGroupIDsDirectly(t *testing.T) {
|
||||
groups := make(map[int64]*Group, 1002)
|
||||
for id := int64(1); id <= 1001; id++ {
|
||||
groups[id] = &Group{ID: id, Platform: PlatformAnthropic}
|
||||
}
|
||||
groups[1002] = &Group{ID: 1002, Platform: PlatformComposite}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: &groupRepoStubForAdmin{getByIDByID: groups}}
|
||||
|
||||
require.Error(t, svc.ValidateAccountGroupBindings(context.Background(), []int64{1, 1002}))
|
||||
require.ErrorIs(t, svc.ValidateAccountGroupBindings(context.Background(), []int64{1, 2000}), ErrGroupNotFound)
|
||||
require.NoError(t, svc.ValidateAccountGroupBindings(context.Background(), []int64{1, 1001}))
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeRejectsDirectCompositeGroupAccess(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: &Group{ID: 9, Platform: PlatformComposite}}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: repo, emptyGroupDeleteRepo: repo}
|
||||
|
||||
_, err := svc.GetGroup(context.Background(), 9)
|
||||
require.Error(t, err)
|
||||
_, err = svc.UpdateGroup(context.Background(), 9, &UpdateGroupInput{})
|
||||
require.Error(t, err)
|
||||
require.Nil(t, repo.updated)
|
||||
require.Error(t, svc.DeleteGroupIfEmpty(context.Background(), 9))
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeRejectsAccountListCompositeFilter(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{getByID: &Group{ID: 9, Platform: PlatformComposite}}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: repo}
|
||||
_, _, err := svc.ListAccounts(context.Background(), 1, 20, "", "", "", "", 9, "", "", "")
|
||||
require.Error(t, err)
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeRejectsAdvancedGroupOperationsDirectly(t *testing.T) {
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}}
|
||||
tests := []struct {
|
||||
name string
|
||||
call func() error
|
||||
}{
|
||||
{"duplicate", func() error { _, err := svc.DuplicateGroup(context.Background(), 1, "admin:1", "key"); return err }},
|
||||
{"recover duplicate", func() error {
|
||||
_, err := svc.RecoverDuplicateGroup(context.Background(), 1, "admin:1", "key")
|
||||
return err
|
||||
}},
|
||||
{"list composite routes", func() error { _, err := svc.ListCompositeRoutes(context.Background(), 1); return err }},
|
||||
{"create composite route", func() error {
|
||||
_, err := svc.CreateCompositeRoute(context.Background(), 1, CompositeRouteInput{})
|
||||
return err
|
||||
}},
|
||||
{"update composite route", func() error {
|
||||
_, err := svc.UpdateCompositeRoute(context.Background(), 1, 2, CompositeRouteInput{})
|
||||
return err
|
||||
}},
|
||||
{"delete composite route", func() error { return svc.DeleteCompositeRoute(context.Background(), 1, 2) }},
|
||||
{"preview composite route", func() error {
|
||||
_, err := svc.PreviewCompositeRoute(context.Background(), 1, CompositeRoutePreviewRequest{})
|
||||
return err
|
||||
}},
|
||||
{"get multipliers", func() error { _, err := svc.GetGroupRateMultipliers(context.Background(), 1); return err }},
|
||||
{"clear multipliers", func() error { return svc.ClearGroupRateMultipliers(context.Background(), 1) }},
|
||||
{"set multipliers", func() error { return svc.BatchSetGroupRateMultipliers(context.Background(), 1, nil) }},
|
||||
{"clear rpm overrides", func() error { return svc.ClearGroupRPMOverrides(context.Background(), 1) }},
|
||||
{"set rpm overrides", func() error { return svc.BatchSetGroupRPMOverrides(context.Background(), 1, nil) }},
|
||||
{"sort", func() error { return svc.UpdateGroupSortOrders(context.Background(), nil) }},
|
||||
}
|
||||
for _, tt := range tests {
|
||||
t.Run(tt.name, func(t *testing.T) {
|
||||
err := tt.call()
|
||||
require.Error(t, err)
|
||||
require.Equal(t, "SIMPLE_MODE_OPERATION_UNSUPPORTED", infraerrors.Reason(err))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeRejectsCompositeCreateAndConversionDirectly(t *testing.T) {
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}}
|
||||
_, err := svc.CreateGroup(context.Background(), &CreateGroupInput{Platform: PlatformComposite, RateMultiplier: 1})
|
||||
require.Error(t, err)
|
||||
|
||||
repo := &groupRepoStubForAdmin{getByID: &Group{ID: 1, Platform: PlatformAnthropic}}
|
||||
svc.groupRepo = repo
|
||||
_, err = svc.UpdateGroup(context.Background(), 1, &UpdateGroupInput{Platform: PlatformComposite})
|
||||
require.Error(t, err)
|
||||
require.Nil(t, repo.updated)
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeNormalizesAllUnsupportedCreateFieldsDirectly(t *testing.T) {
|
||||
one := 1.0
|
||||
fallbackID := int64(44)
|
||||
input := &CreateGroupInput{
|
||||
Name: "simple", Description: "allowed", Platform: PlatformAnthropic,
|
||||
RateMultiplier: 9, IsExclusive: true, SubscriptionType: SubscriptionTypeSubscription,
|
||||
DailyLimitUSD: &one, LongContextPricingEnabled: true,
|
||||
ModelPricing: []ChannelModelPricing{{Models: []string{"claude"}}},
|
||||
PeakRateEnabled: true, PeakStart: "00:00", PeakEnd: "01:00", PeakRateMultiplier: &one,
|
||||
ImageRateIndependent: true, ImageRateMultiplier: &one, VideoRateIndependent: true, VideoRateMultiplier: &one,
|
||||
ImagePrice1K: &one, VideoPrice720P: &one, WebSearchPricePerCall: &one, SearchPricePer1k: &one,
|
||||
AudioRealtimePricePerMin: &one, ClaudeCodeOnly: true, FallbackGroupID: &fallbackID,
|
||||
ModelRouting: map[string][]int64{"claude": {1}}, ModelRoutingEnabled: true,
|
||||
AllowMessagesDispatch: true, AllowLive: true, ForceOpenAIFast: true, RequireOAuthOnly: true,
|
||||
RPMLimit: 99, MaxReasoningEffort: "high", ProfitControlEnabled: true, ProfitMinMargin: &one,
|
||||
CopyAccountsFromGroupIDs: []int64{2},
|
||||
}
|
||||
repo := &groupRepoStubForAdmin{}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: repo}
|
||||
|
||||
created, err := svc.CreateGroup(context.Background(), input)
|
||||
require.NoError(t, err)
|
||||
require.Same(t, repo.created, created)
|
||||
require.Equal(t, CreateGroupInput{
|
||||
Name: "simple", Description: "allowed", Platform: PlatformAnthropic,
|
||||
RateMultiplier: 1, SubscriptionType: SubscriptionTypeStandard,
|
||||
}, *input)
|
||||
require.Equal(t, 1.0, created.RateMultiplier)
|
||||
require.Equal(t, SubscriptionTypeStandard, created.SubscriptionType)
|
||||
require.False(t, created.IsExclusive)
|
||||
require.Nil(t, created.FallbackGroupID)
|
||||
require.Empty(t, created.ModelPricing)
|
||||
require.Zero(t, created.RPMLimit)
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeNormalizesAllUnsupportedUpdateFieldsDirectly(t *testing.T) {
|
||||
one := 1.0
|
||||
truth := true
|
||||
status := "inactive"
|
||||
description := "allowed"
|
||||
fallbackID := int64(44)
|
||||
pricing := []ChannelModelPricing{{Models: []string{"claude"}}}
|
||||
input := &UpdateGroupInput{
|
||||
Name: "renamed", Description: &description, Platform: PlatformOpenAI, Status: status,
|
||||
RateMultiplier: &one, IsExclusive: &truth, SubscriptionType: SubscriptionTypeSubscription,
|
||||
DailyLimitUSD: &one, LongContextPricingEnabled: &truth, ModelPricing: &pricing,
|
||||
PeakRateEnabled: &truth, PeakRateMultiplier: &one, ImageRateIndependent: &truth,
|
||||
ImageRateMultiplier: &one, VideoRateIndependent: &truth, VideoRateMultiplier: &one,
|
||||
ImagePrice1K: &one, VideoPrice720P: &one, WebSearchPricePerCall: &one, SearchPricePer1k: &one,
|
||||
AudioRealtimePricePerMin: &one, ClaudeCodeOnly: &truth, FallbackGroupID: &fallbackID,
|
||||
ModelRouting: map[string][]int64{"claude": {1}}, ModelRoutingEnabled: &truth,
|
||||
AllowMessagesDispatch: &truth, AllowLive: &truth, ForceOpenAIFast: &truth, RequireOAuthOnly: &truth,
|
||||
RPMLimit: new(int), MaxReasoningEffort: ptrString("high"), ProfitControlEnabled: &truth,
|
||||
ProfitMinMargin: &one, CopyAccountsFromGroupIDs: []int64{2},
|
||||
}
|
||||
existing := &Group{ID: 1, Name: "old", Description: "old description", Platform: PlatformAnthropic, Status: StatusActive, RateMultiplier: 3, RPMLimit: 8, FallbackGroupID: &fallbackID}
|
||||
repo := &groupRepoStubForAdmin{getByID: existing}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: repo}
|
||||
|
||||
updated, err := svc.UpdateGroup(context.Background(), 1, input)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, UpdateGroupInput{Name: "renamed", Description: &description}, *input)
|
||||
require.Equal(t, "renamed", updated.Name)
|
||||
require.Equal(t, description, updated.Description)
|
||||
require.Equal(t, PlatformAnthropic, updated.Platform)
|
||||
require.Equal(t, StatusActive, updated.Status)
|
||||
require.Equal(t, 3.0, updated.RateMultiplier)
|
||||
require.Equal(t, 8, updated.RPMLimit)
|
||||
require.Equal(t, &fallbackID, updated.FallbackGroupID)
|
||||
}
|
||||
|
||||
func TestAdminServiceSimpleModeListUsesRepositoryFilteredTotal(t *testing.T) {
|
||||
repo := &groupRepoStubForAdmin{
|
||||
listWithFiltersGroups: []Group{{ID: 2, Platform: PlatformAnthropic}},
|
||||
listWithFiltersResult: &pagination.PaginationResult{Total: 11, Page: 2, PageSize: 1},
|
||||
}
|
||||
svc := &adminServiceImpl{cfg: &config.Config{RunMode: config.RunModeSimple}, groupRepo: repo}
|
||||
groups, total, err := svc.ListGroups(context.Background(), 2, 1, "", "", "", nil, "id", "asc")
|
||||
require.NoError(t, err)
|
||||
require.Len(t, groups, 1)
|
||||
require.EqualValues(t, 11, total)
|
||||
}
|
||||
|
||||
func (s *groupRepoStubForAdmin) ListActive(_ context.Context) ([]Group, error) {
|
||||
panic("unexpected ListActive call")
|
||||
}
|
||||
|
||||
@@ -107,8 +107,8 @@ type APIKeyAuthGroupSnapshot struct {
|
||||
FreeOpenAIFast bool `json:"free_openai_fast"`
|
||||
DefaultMappedModel string `json:"default_mapped_model,omitempty"`
|
||||
MessagesDispatchModelConfig OpenAIMessagesDispatchModelConfig `json:"messages_dispatch_model_config,omitempty"`
|
||||
ModelsListConfig GroupModelsListConfig `json:"models_list_config,omitempty"`
|
||||
// CodexModelsManifestConfig 与 ModelsListConfig 一样在认证快照分组里透传,
|
||||
ModelAllowlist GroupModelAllowlist `json:"model_allowlist,omitempty"`
|
||||
// CodexModelsManifestConfig 与 ModelAllowlist 一样在认证快照分组里透传,
|
||||
// Codex /models handler 直接读认证分组对象。
|
||||
CodexModelsManifestConfig GroupCodexModelsManifestConfig `json:"codex_models_manifest_config,omitempty"`
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ import (
|
||||
"github.com/dgraph-io/ristretto"
|
||||
)
|
||||
|
||||
const apiKeyAuthSnapshotVersion = 23 // v23: group codex_models_manifest_config field
|
||||
const apiKeyAuthSnapshotVersion = 24 // v24: group model_allowlist field (renamed from models_list_config, enforcing semantics)
|
||||
|
||||
type apiKeyAuthCacheConfig struct {
|
||||
l1Size int
|
||||
@@ -422,7 +422,7 @@ func (s *APIKeyService) snapshotFromAPIKey(ctx context.Context, apiKey *APIKey)
|
||||
FreeOpenAIFast: apiKey.Group.FreeOpenAIFast,
|
||||
DefaultMappedModel: apiKey.Group.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: apiKey.Group.MessagesDispatchModelConfig,
|
||||
ModelsListConfig: apiKey.Group.ModelsListConfig,
|
||||
ModelAllowlist: apiKey.Group.ModelAllowlist,
|
||||
CodexModelsManifestConfig: apiKey.Group.CodexModelsManifestConfig,
|
||||
RPMLimit: apiKey.Group.RPMLimit,
|
||||
MaxReasoningEffort: apiKey.Group.MaxReasoningEffort,
|
||||
@@ -524,7 +524,7 @@ func (s *APIKeyService) snapshotToAPIKey(key string, snapshot *APIKeyAuthSnapsho
|
||||
FreeOpenAIFast: snapshot.Group.FreeOpenAIFast,
|
||||
DefaultMappedModel: snapshot.Group.DefaultMappedModel,
|
||||
MessagesDispatchModelConfig: snapshot.Group.MessagesDispatchModelConfig,
|
||||
ModelsListConfig: snapshot.Group.ModelsListConfig,
|
||||
ModelAllowlist: snapshot.Group.ModelAllowlist,
|
||||
CodexModelsManifestConfig: snapshot.Group.CodexModelsManifestConfig,
|
||||
RPMLimit: snapshot.Group.RPMLimit,
|
||||
MaxReasoningEffort: snapshot.Group.MaxReasoningEffort,
|
||||
|
||||
Some files were not shown because too many files have changed in this diff Show More
Reference in New Issue
Block a user