feat(groups)!: enforce model allowlists across gateway endpoints

This commit is contained in:
Zero Clover
2026-09-05 16:05:17 +08:00
parent 07bf8b92fd
commit cff3f89850
73 changed files with 3691 additions and 778 deletions
@@ -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")
@@ -153,7 +153,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 +228,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 表示未提供不改动
@@ -471,9 +471,9 @@ func (h *GroupHandler) GetByID(c *gin.Context) {
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) {
groupID, err := strconv.ParseInt(c.Param("id"), 10, 64)
if err != nil || groupID < 0 {
response.BadRequest(c, "Invalid group ID")
@@ -568,7 +568,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,
@@ -701,7 +701,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,
@@ -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),
@@ -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 {
+1 -1
View File
@@ -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,
+1 -1
View File
@@ -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"`
+38 -78
View File
@@ -1089,9 +1089,12 @@ func (h *GatewayHandler) Models(c *gin.Context) {
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 {
@@ -1104,10 +1107,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
}
@@ -1192,8 +1194,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
@@ -1203,12 +1209,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
@@ -1267,7 +1269,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
@@ -1365,72 +1367,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:
@@ -1504,10 +1453,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)
})
}
+12 -12
View File
@@ -267,7 +267,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 +299,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 +572,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 +619,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 +704,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 +918,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 +969,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 +1020,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 +1061,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 +1106,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 +1143,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 +1180,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 层的平台路由逻辑不变量
@@ -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"},
},
@@ -912,7 +912,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"},
},
@@ -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"
@@ -2263,6 +2264,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 {
@@ -2680,6 +2691,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)
@@ -3581,6 +3603,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
@@ -1882,7 +1882,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
@@ -1890,6 +1893,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 {
@@ -2809,8 +2818,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)
@@ -2963,6 +2975,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)
@@ -2994,6 +3009,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)
@@ -3004,11 +3032,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))
}
}
+9 -1
View File
@@ -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)
}