mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-11-01 08:28:34 +00:00
feat(groups)!: enforce model allowlists across gateway endpoints
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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"`
|
||||
|
||||
|
||||
@@ -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)
|
||||
})
|
||||
}
|
||||
@@ -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))
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user