Merge pull request #7173 from HuntercodeT/fix/gemini-native-mixed-model-listing

fix(gemini): discover Antigravity mappings in native model lists

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")
 }

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"))
+			}
+		})
+	}
+}

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 表示是否有条目被移除（全命中时

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

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)