diff --git a/backend/internal/handler/gateway_model_allowlist_listing_test.go b/backend/internal/handler/gateway_model_allowlist_listing_test.go index 55a663dbc0..cb1fd5d5fa 100644 --- a/backend/internal/handler/gateway_model_allowlist_listing_test.go +++ b/backend/internal/handler/gateway_model_allowlist_listing_test.go @@ -91,7 +91,7 @@ func TestGeminiV1BetaListModels_FiltersFallbackByAllowlist(t *testing.T) { repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{ byGroup: map[int64][]service.Account{ // 只有 antigravity 账号：Gemini 选号失败后回落静态模型列表。 - 41: {{ID: 2, Platform: service.PlatformAntigravity, Status: service.StatusActive, Schedulable: true}}, + 41: {{ID: 2, Platform: service.PlatformAntigravity, Extra: map[string]any{"mixed_scheduling": true}, Status: service.StatusActive, Schedulable: true}}, }, }} h := &GatewayHandler{ @@ -109,7 +109,7 @@ func TestGeminiV1BetaListModels_FiltersFallbackByAllowlist(t *testing.T) { Platform: service.PlatformGemini, ModelAllowlist: service.GroupModelAllowlist{ Enabled: true, - Models: []string{"gemini-2.5-pro", "gemini-3-*"}, + Models: []string{"gemini-3.8-flash-high", "gemini-3-*"}, }, }, }) @@ -129,8 +129,8 @@ func TestGeminiV1BetaListModels_FiltersFallbackByAllowlist(t *testing.T) { 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.Contains(t, names, "models/gemini-3.8-flash-high") + require.Contains(t, names, "models/gemini-3-flash") require.NotContains(t, names, "models/gemini-2.5-flash") require.NotContains(t, names, "models/gemini-2.0-flash") } diff --git a/backend/internal/handler/gemini_mixed_models_test.go b/backend/internal/handler/gemini_mixed_models_test.go new file mode 100644 index 0000000000..6a0e2c2eb5 --- /dev/null +++ b/backend/internal/handler/gemini_mixed_models_test.go @@ -0,0 +1,136 @@ +package handler + +import ( + "encoding/json" + "github.com/Wei-Shaw/sub2api/internal/config" + "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" + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" +) + +func TestGeminiNativeModelsUsesAccountMappings(t *testing.T) { + for _, tt := range []struct { + name string + forced, mixed, allowlist bool + status int + }{ + {"mixed", false, true, false, 200}, + {"mixed allowlist", false, true, true, 200}, + {"disabled mixed", false, false, false, 503}, + {"forced without mixed opt-in", true, false, false, 200}, + {"forced allowlist", true, false, true, 200}, + } { + t.Run(tt.name, func(t *testing.T) { + groupID := int64(45) + repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{ + groupID: {{ID: 1, Platform: service.PlatformAntigravity, + Extra: map[string]any{"mixed_scheduling": tt.mixed}, + Credentials: map[string]any{"model_mapping": map[string]any{"gemini-synced-custom": "gemini-3.8-flash-high", "claude-custom": "claude-sonnet-4-6"}}}}, + }}} + 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) + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{GroupID: &groupID, Group: &service.Group{ID: groupID, Platform: service.PlatformGemini, + ModelAllowlist: service.GroupModelAllowlist{Enabled: tt.allowlist, Models: []string{"gemini-synced-custom"}}, + }}) + if tt.forced { + c.Set(string(middleware.ContextKeyForcePlatform), service.PlatformAntigravity) + } + h.GeminiV1BetaListModels(c) + require.Equal(t, tt.status, rec.Code, rec.Body.String()) + if tt.status != 200 { + return + } + var got gemini.ModelsListResponse + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + names := []string{} + for _, model := range got.Models { + names = append(names, model.Name) + require.Contains(t, model.SupportedGenerationMethods, "generateContent") + } + require.Contains(t, names, "models/gemini-synced-custom") + require.NotContains(t, names, "models/claude-custom") + require.NotContains(t, names, "models/gemini-2.0-flash") + if tt.allowlist { + require.Equal(t, []string{"models/gemini-synced-custom"}, names) + } + }) + } +} + +func TestAppendUpstreamGeminiModelsPreservesMetadata(t *testing.T) { + body := []byte(`{"models":[{"name":"models/gemini-native","inputTokenLimit":123,"custom":{"a":true}}],"nextPageToken":"next","other":42}`) + extra := []gemini.Model{gemini.FallbackModel("gemini-native"), gemini.FallbackModel("gemini-synced-custom"), gemini.FallbackModel("gemini-synced-custom")} + merged, ok := appendUpstreamGeminiModels(body, extra) + require.True(t, ok) + require.JSONEq(t, `{"models":[{"name":"models/gemini-native","inputTokenLimit":123,"custom":{"a":true}},{"name":"models/gemini-synced-custom","supportedGenerationMethods":["generateContent","streamGenerateContent"]}],"nextPageToken":"next","other":42}`, string(merged)) + filtered, dropped, ok := filterUpstreamGeminiModelsBody(merged, service.GroupModelAllowlist{Enabled: true, Models: []string{"gemini-synced-*"}}) + require.True(t, ok) + require.True(t, dropped) + require.NotContains(t, string(filtered), "gemini-native") + require.Contains(t, string(filtered), "gemini-synced-custom") + for _, invalid := range []string{`null`, `not-json`, `{"error":"bad"}`, `{"models":{}}`} { + got, ok := appendUpstreamGeminiModels([]byte(invalid), extra) + require.False(t, ok) + require.Equal(t, invalid, string(got)) + } +} + +// Exercise the real handler's native-upstream branch, including scope fallback. +type geminiMixedModelsUpstream struct { + service.HTTPUpstream + status int + body string +} + +func (u *geminiMixedModelsUpstream) Do(_ *http.Request, _ string, _ int64, _ int) (*http.Response, error) { + return &http.Response{StatusCode: u.status, Header: http.Header{"Content-Type": {"application/json"}, "X-Request-Id": {"native-id"}}, Body: io.NopCloser(strings.NewReader(u.body))}, nil +} +func TestGeminiNativeModelsMergesNativeUpstream(t *testing.T) { + for _, tt := range []struct { + name string + status int + body string + expectedNative string + }{ + {"native", 200, `{"models":[{"name":"models/gemini-native","inputTokenLimit":123}],"nextPageToken":"next"}`, "models/gemini-native"}, + {"scope fallback", 403, `{"error":"insufficient authentication scopes"}`, "models/gemini-2.5-pro"}, + {"upstream error", 429, `{"error":"rate limited"}`, ""}, + } { + t.Run(tt.name, func(t *testing.T) { + id := int64(46) + repo := &geminiAllowlistAccountRepoStub{gatewayModelsAccountRepoStub: gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{id: { + {ID: 1, Platform: service.PlatformGemini, Type: service.AccountTypeAPIKey, Credentials: map[string]any{"api_key": "test"}}, + {ID: 2, Platform: service.PlatformAntigravity, Extra: map[string]any{"mixed_scheduling": true}, Credentials: map[string]any{"model_mapping": map[string]any{"gemini-synced-custom": "gemini-3.8-flash-high"}}}, + }}}} + upstream := &geminiMixedModelsUpstream{status: tt.status, body: tt.body} + h := &GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, nil, nil, upstream, nil, &config.Config{})} + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1beta/models", nil) + c.Set(string(middleware.ContextKeyAPIKey), &service.APIKey{GroupID: &id, Group: &service.Group{ID: id, Platform: service.PlatformGemini}}) + h.GeminiV1BetaListModels(c) + if tt.status == 429 { + require.Equal(t, 429, rec.Code) + require.JSONEq(t, tt.body, rec.Body.String()) + return + } + require.Equal(t, 200, rec.Code, rec.Body.String()) + require.Contains(t, rec.Body.String(), tt.expectedNative) + require.Contains(t, rec.Body.String(), "models/gemini-synced-custom") + if tt.status == 200 { + require.Contains(t, rec.Body.String(), `"inputTokenLimit":123`) + require.Contains(t, rec.Body.String(), `"nextPageToken":"next"`) + require.Equal(t, "native-id", rec.Header().Get("X-Request-Id")) + } + }) + } +} diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index d8ec24670b..a7ae9170c1 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -60,30 +60,24 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { return filtered } - // 强制 antigravity 模式：返回 antigravity 支持的模型列表 + agModelIDs, err := h.geminiCompatService.AntigravityGeminiModelIDs(c.Request.Context(), apiKey.GroupID, forcePlatform != service.PlatformAntigravity) + if err != nil { + googleError(c, http.StatusServiceUnavailable, "Unable to list Antigravity models") + return + } + agModels := make([]gemini.Model, 0, len(agModelIDs)) + for _, id := range agModelIDs { + agModels = append(agModels, gemini.FallbackModel(id)) + } 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()) + c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(agModels)}) return } account, err := h.geminiCompatService.SelectAccountForAIStudioEndpoints(c.Request.Context(), apiKey.GroupID) if err != nil { - // 没有 gemini 账户，检查是否有 antigravity 账户可用 - hasAntigravity, _ := h.geminiCompatService.HasAntigravityAccounts(c.Request.Context(), apiKey.GroupID) - if hasAntigravity { - // antigravity 账户使用静态模型列表 - c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())}) + if len(agModels) > 0 { + c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(agModels)}) return } markOpsRoutingCapacityLimitedIfNoAvailable(c, err) @@ -97,9 +91,15 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { return } if shouldFallbackGeminiModels(res) { - c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(gemini.DefaultModels())}) + c.JSON(http.StatusOK, gemini.ModelsListResponse{Models: filterGeminiModels(mergeGeminiModelLists(gemini.DefaultModels(), agModels))}) return } + if res.StatusCode == http.StatusOK && len(agModels) > 0 { + if merged, ok := appendUpstreamGeminiModels(res.Body, agModels); ok { + res.Body = merged + } + } + if apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() { if filtered, dropped, ok := filterUpstreamGeminiModelsBody(res.Body, apiKey.Group.ModelAllowlist); ok && dropped { // 只在确有条目被过滤时替换响应体；全命中或解析失败时保持原始响应， @@ -110,6 +110,64 @@ func (h *GatewayHandler) GeminiV1BetaListModels(c *gin.Context) { writeUpstreamResponse(c, res) } +// mergeGeminiModelLists keeps native metadata when both sources advertise a model. +func mergeGeminiModelLists(native, extra []gemini.Model) []gemini.Model { + result := append([]gemini.Model{}, native...) + seen := make(map[string]bool, len(native)) + for _, model := range native { + seen[model.Name] = true + } + for _, model := range extra { + if !seen[model.Name] { + result = append(result, model) + seen[model.Name] = true + } + } + return result +} + +// appendUpstreamGeminiModels preserves unknown model metadata and envelope fields. +func appendUpstreamGeminiModels(body []byte, extra []gemini.Model) ([]byte, bool) { + var envelope map[string]json.RawMessage + if json.Unmarshal(body, &envelope) != nil || envelope == nil { + return body, false + } + var models []json.RawMessage + raw, exists := envelope["models"] + if !exists || json.Unmarshal(raw, &models) != nil { + return body, false + } + seen := make(map[string]bool, len(models)) + for _, raw := range models { + var model struct { + Name string `json:"name"` + } + if json.Unmarshal(raw, &model) != nil { + return body, false + } + seen[model.Name] = true + } + changed := false + for _, model := range extra { + if seen[model.Name] { + continue + } + raw, err := json.Marshal(model) + if err != nil { + return body, false + } + models = append(models, raw) + seen[model.Name] = true + changed = true + } + if !changed { + return body, true + } + envelope["models"], _ = json.Marshal(models) + merged, err := json.Marshal(envelope) + return merged, err == nil +} + // filterUpstreamGeminiModelsBody 按白名单过滤上游 /v1beta/models 响应中的 // models[].name，其余信封字段（如 nextPageToken）原样保留。 // 返回值：filtered 为过滤后的响应体；dropped 表示是否有条目被移除（全命中时 diff --git a/backend/internal/handler/gemini_v1beta_handler_test.go b/backend/internal/handler/gemini_v1beta_handler_test.go index 2bf25a6b96..8604b07b88 100644 --- a/backend/internal/handler/gemini_v1beta_handler_test.go +++ b/backend/internal/handler/gemini_v1beta_handler_test.go @@ -39,7 +39,11 @@ func TestGeminiV1BetaListModels_ForcedAntigravityAppliesAllowlist(t *testing.T) }) c.Set(string(middleware.ContextKeyForcePlatform), service.PlatformAntigravity) - (&GatewayHandler{}).GeminiV1BetaListModels(c) + groupID := int64(42) + key, _ := middleware.GetAPIKeyFromContext(c) + key.GroupID = &groupID + repo := &geminiAllowlistAccountRepoStub{} + (&GatewayHandler{geminiCompatService: service.NewGeminiMessagesCompatService(repo, nil, nil, nil, nil, nil, nil, nil, nil)}).GeminiV1BetaListModels(c) require.Equal(t, http.StatusOK, rec.Code) var got antigravity.GeminiModelsListResponse diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index edd727ce5e..306d84d8f4 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -16,6 +16,7 @@ import ( mathrand "math/rand" "net/http" "regexp" + "sort" "strings" "time" @@ -483,6 +484,32 @@ func (s *GeminiMessagesCompatService) validateUpstreamBaseURL(raw string) (strin return normalized, nil } +// AntigravityGeminiModelIDs returns the effective Gemini mappings of schedulable +// Antigravity accounts. Mixed routes must additionally respect the account opt-in. +func (s *GeminiMessagesCompatService) AntigravityGeminiModelIDs(ctx context.Context, groupID *int64, requireMixed bool) ([]string, error) { + accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformAntigravity, false) + if err != nil { + return nil, err + } + seen := make(map[string]struct{}) + for _, account := range accounts { + if account.Platform != PlatformAntigravity || (requireMixed && !account.IsMixedSchedulingEnabled()) { + continue + } + for model := range account.GetModelMapping() { + if isAntigravityGeminiModel(model) { + seen[model] = struct{}{} + } + } + } + models := make([]string, 0, len(seen)) + for model := range seen { + models = append(models, model) + } + sort.Strings(models) + return models, nil +} + // HasAntigravityAccounts 检查是否有可用的 antigravity 账户 func (s *GeminiMessagesCompatService) HasAntigravityAccounts(ctx context.Context, groupID *int64) (bool, error) { accounts, err := s.listSchedulableAccountsOnce(ctx, groupID, PlatformAntigravity, false)