From f98a056f75e93c81e3ab1cc8623db9e17b3dc432 Mon Sep 17 00:00:00 2001 From: Hakunm <275509506+Hakunm@users.noreply.github.com> Date: Fri, 21 Aug 2026 02:33:00 +0800 Subject: [PATCH] fix(gemini): constrain Google One model catalog --- .../internal/handler/admin/account_handler.go | 8 +++- .../account_handler_available_models_test.go | 36 +++++++++++++++ backend/internal/pkg/geminicli/models.go | 19 ++++++++ backend/internal/pkg/geminicli/models_test.go | 21 +++++++++ backend/internal/service/account.go | 13 ++++++ .../internal/service/account_wildcard_test.go | 46 +++++++++++++++++++ 6 files changed, 142 insertions(+), 1 deletion(-) diff --git a/backend/internal/handler/admin/account_handler.go b/backend/internal/handler/admin/account_handler.go index 465c6a7816..93067da278 100644 --- a/backend/internal/handler/admin/account_handler.go +++ b/backend/internal/handler/admin/account_handler.go @@ -2614,8 +2614,14 @@ func (h *AccountHandler) GetAvailableModels(c *gin.Context) { // Handle Gemini accounts if account.IsGemini() { - // For OAuth accounts: return default Gemini models + // Consumer Google One OAuth still uses the legacy Gemini CLI / Code + // Assist channel. Do not advertise newer 3.x or image models that the + // channel cannot serve. if account.IsOAuth() { + if account.IsGeminiGoogleOne() { + response.Success(c, geminicli.GoogleOneModels) + return + } response.Success(c, geminicli.DefaultModels) return } diff --git a/backend/internal/handler/admin/account_handler_available_models_test.go b/backend/internal/handler/admin/account_handler_available_models_test.go index 880cb8c3ea..ef2f65d4e3 100644 --- a/backend/internal/handler/admin/account_handler_available_models_test.go +++ b/backend/internal/handler/admin/account_handler_available_models_test.go @@ -287,6 +287,42 @@ func TestAccountHandlerGetAvailableModels_OpenAISparkShadowReturnsMappingModels( }, ids, "影子可用模型由 model_mapping 派生（非写死）") } +func TestAccountHandlerGetAvailableModels_GeminiGoogleOneUsesConservativeCatalog(t *testing.T) { + svc := &availableModelsAdminService{ + stubAdminService: newStubAdminService(), + account: service.Account{ + ID: 45, + Name: "google-one", + Platform: service.PlatformGemini, + Type: service.AccountTypeOAuth, + Status: service.StatusActive, + Credentials: map[string]any{ + "oauth_type": "google_one", + }, + }, + } + router := setupAvailableModelsRouter(svc) + + rec := httptest.NewRecorder() + req := httptest.NewRequest(http.MethodGet, "/api/v1/admin/accounts/45/models", nil) + router.ServeHTTP(rec, req) + + require.Equal(t, http.StatusOK, rec.Code) + var resp struct { + Data []struct { + ID string `json:"id"` + } `json:"data"` + } + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &resp)) + ids := make([]string, 0, len(resp.Data)) + for _, model := range resp.Data { + ids = append(ids, model.ID) + } + require.ElementsMatch(t, []string{"gemini-2.0-flash", "gemini-2.5-flash", "gemini-2.5-pro"}, ids) + require.NotContains(t, ids, "gemini-3.5-flash") + require.NotContains(t, ids, "gemini-2.5-flash-image") +} + func TestAccountHandlerSyncUpstreamModels_ConfigErrorReturnsBadRequest(t *testing.T) { svc := &availableModelsAdminService{ stubAdminService: newStubAdminService(), diff --git a/backend/internal/pkg/geminicli/models.go b/backend/internal/pkg/geminicli/models.go index bbd9a6c464..cc9cbc80b2 100644 --- a/backend/internal/pkg/geminicli/models.go +++ b/backend/internal/pkg/geminicli/models.go @@ -22,5 +22,24 @@ var DefaultModels = []Model{ {ID: "gemini-3.1-flash-image", Type: "model", DisplayName: "Gemini 3.1 Flash Image", CreatedAt: ""}, } +// GoogleOneModels is the conservative model set exposed for legacy Google One +// OAuth accounts. Newer subscription models are served through Antigravity +// OAuth rather than the retired consumer Gemini CLI / Code Assist channel. +var GoogleOneModels = []Model{ + {ID: "gemini-2.5-flash", Type: "model", DisplayName: "Gemini 2.5 Flash", CreatedAt: ""}, + {ID: "gemini-2.5-pro", Type: "model", DisplayName: "Gemini 2.5 Pro", CreatedAt: ""}, + {ID: "gemini-2.0-flash", Type: "model", DisplayName: "Gemini 2.0 Flash", CreatedAt: ""}, +} + +// GoogleOneModelMapping returns a new whitelist map for each account so callers +// cannot mutate the package-level catalog. +func GoogleOneModelMapping() map[string]string { + mapping := make(map[string]string, len(GoogleOneModels)) + for _, model := range GoogleOneModels { + mapping[model.ID] = model.ID + } + return mapping +} + // DefaultTestModel is the default model to preselect in test flows. const DefaultTestModel = "gemini-2.0-flash" diff --git a/backend/internal/pkg/geminicli/models_test.go b/backend/internal/pkg/geminicli/models_test.go index c1884e2e10..e0cb57941f 100644 --- a/backend/internal/pkg/geminicli/models_test.go +++ b/backend/internal/pkg/geminicli/models_test.go @@ -21,3 +21,24 @@ func TestDefaultModels_ContainsImageModels(t *testing.T) { } } } + +func TestGoogleOneModels_ExcludeUnsupportedNewAndImageModels(t *testing.T) { + t.Parallel() + + mapping := GoogleOneModelMapping() + for _, id := range []string{"gemini-2.0-flash", "gemini-2.5-flash", "gemini-2.5-pro"} { + if mapping[id] != id { + t.Fatalf("expected Google One model %q to map to itself", id) + } + } + for _, id := range []string{"gemini-2.5-flash-image", "gemini-3.1-flash-image", "gemini-3.5-flash"} { + if _, ok := mapping[id]; ok { + t.Fatalf("did not expect unsupported Google One model %q", id) + } + } + + mapping["gemini-2.5-flash"] = "mutated" + if GoogleOneModelMapping()["gemini-2.5-flash"] != "gemini-2.5-flash" { + t.Fatal("GoogleOneModelMapping must return a defensive copy") + } +} diff --git a/backend/internal/service/account.go b/backend/internal/service/account.go index 20d4ec5823..4ee1a514f9 100644 --- a/backend/internal/service/account.go +++ b/backend/internal/service/account.go @@ -15,6 +15,7 @@ import ( "github.com/Wei-Shaw/sub2api/internal/config" "github.com/Wei-Shaw/sub2api/internal/domain" + "github.com/Wei-Shaw/sub2api/internal/pkg/geminicli" "github.com/Wei-Shaw/sub2api/internal/pkg/openai_compat" "github.com/Wei-Shaw/sub2api/internal/pkg/xai" ) @@ -324,6 +325,12 @@ func (a *Account) IsGeminiCodeAssist() bool { return oauthType == "code_assist" } +// IsGeminiGoogleOne reports whether this account uses the legacy consumer +// Gemini CLI / Code Assist OAuth channel. +func (a *Account) IsGeminiGoogleOne() bool { + return a.Platform == PlatformGemini && a.Type == AccountTypeOAuth && a.GeminiOAuthType() == "google_one" +} + func (a *Account) CanGetUsage() bool { return a.Type == AccountTypeOAuth } @@ -625,6 +632,9 @@ func (a *Account) resolveModelMapping(rawMapping map[string]any) map[string]stri return nil } if len(rawMapping) == 0 { + if a.IsGeminiGoogleOne() { + return geminicli.GoogleOneModelMapping() + } // Antigravity 平台使用默认映射 if a.Platform == domain.PlatformAntigravity { return domain.DefaultAntigravityModelMapping @@ -659,6 +669,9 @@ func (a *Account) resolveModelMapping(rawMapping map[string]any) map[string]stri } // Antigravity 平台使用默认映射 + if a.IsGeminiGoogleOne() { + return geminicli.GoogleOneModelMapping() + } if a.Platform == domain.PlatformAntigravity { return domain.DefaultAntigravityModelMapping } diff --git a/backend/internal/service/account_wildcard_test.go b/backend/internal/service/account_wildcard_test.go index 15c5a66fa1..3843237e10 100644 --- a/backend/internal/service/account_wildcard_test.go +++ b/backend/internal/service/account_wildcard_test.go @@ -537,6 +537,52 @@ func TestAccountGetModelMapping_AntigravityEnsuresGeminiDefaultPassthroughs(t *t } } +func TestAccountGetModelMapping_GoogleOneUsesConservativeDefaults(t *testing.T) { + account := &Account{ + Platform: PlatformGemini, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "oauth_type": "google_one", + }, + } + + mapping := account.GetModelMapping() + for _, model := range []string{"gemini-2.0-flash", "gemini-2.5-flash", "gemini-2.5-pro"} { + if mapping[model] != model { + t.Fatalf("expected Google One model %q to map to itself, got %q", model, mapping[model]) + } + } + for _, model := range []string{"gemini-2.5-flash-image", "gemini-3.1-flash-image", "gemini-3.5-flash"} { + if _, ok := mapping[model]; ok { + t.Fatalf("did not expect unsupported Google One model %q", model) + } + } + if account.IsModelSupported("gemini-3.5-flash") { + t.Fatal("Google One defaults must not treat unsupported models as eligible") + } +} + +func TestAccountGetModelMapping_GoogleOnePreservesExplicitMapping(t *testing.T) { + account := &Account{ + Platform: PlatformGemini, + Type: AccountTypeOAuth, + Credentials: map[string]any{ + "oauth_type": "google_one", + "model_mapping": map[string]any{ + "custom-model": "gemini-2.5-flash", + }, + }, + } + + mapping := account.GetModelMapping() + if mapping["custom-model"] != "gemini-2.5-flash" { + t.Fatalf("expected explicit Google One mapping to be preserved, got %v", mapping) + } + if _, ok := mapping["gemini-2.5-flash"]; ok { + t.Fatalf("did not expect defaults to overwrite an explicit mapping: %v", mapping) + } +} + func TestAccountGetModelMapping_AntigravityRespectsWildcardOverride(t *testing.T) { account := &Account{ Platform: PlatformAntigravity,