Fix composite model defaults for linked platforms

This commit is contained in:
Heatherm Huang
2026-07-23 09:20:52 +08:00
parent 2e774a48b3
commit ee332cee64
3 changed files with 72 additions and 1 deletions
+8 -1
View File
@@ -1085,8 +1085,15 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *
}
seen := make(map[string]struct{})
models := make([]string, 0)
schedulablePlatforms := h.gatewayService.GetSchedulablePlatforms(ctx, groupID)
for _, platform := range []string{service.PlatformAnthropic, service.PlatformGemini, service.PlatformOpenAI, service.PlatformAntigravity, service.PlatformGrok} {
for _, model := range h.gatewayService.GetAvailableModels(ctx, groupID, platform) {
platformModels := h.gatewayService.GetAvailableModels(ctx, groupID, platform)
if len(platformModels) == 0 {
if _, ok := schedulablePlatforms[platform]; ok {
platformModels = defaultModelIDsForPlatform(platform)
}
}
for _, model := range platformModels {
model = strings.TrimSpace(model)
if model == "" {
continue
@@ -349,6 +349,42 @@ func TestGatewayModels_CompositeCustomModelsListFiltersAcrossConcretePlatforms(t
require.Equal(t, []string{"gemini-2.5-flash", "ag-custom-model", "gpt-5.5"}, modelIDsForTest(got.Data))
}
func TestGatewayModels_CompositeUnmappedAccountsFallbackToLinkedPlatformsOnly(t *testing.T) {
gin.SetMode(gin.TestMode)
groupID := int64(34)
h := newGatewayModelsHandlerForTest(
&gatewayModelsAccountRepoStub{
byGroup: map[int64][]service.Account{
groupID: {
{ID: 1, Platform: service.PlatformOpenAI},
{ID: 2, Platform: service.PlatformGrok},
},
},
},
)
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.PlatformComposite},
})
h.Models(c)
require.Equal(t, http.StatusOK, rec.Code)
var got gatewayModelsResponseForTest
require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got))
ids := modelIDsForTest(got.Data)
require.Contains(t, ids, "gpt-5.5")
require.Contains(t, ids, "grok-4.3")
require.NotContains(t, ids, "claude-sonnet-4-6")
require.NotContains(t, ids, "gemini-2.5-flash")
}
func TestGatewayModels_CustomModelsListKeepsConcreteModelAllowedByWildcardMapping(t *testing.T) {
gin.SetMode(gin.TestMode)
@@ -1244,6 +1244,34 @@ func (s *GatewayService) GetAvailableModels(ctx context.Context, groupID *int64,
return cloneStringSlice(models)
}
// GetSchedulablePlatforms returns the concrete platforms that currently have
// schedulable accounts in the target group.
func (s *GatewayService) GetSchedulablePlatforms(ctx context.Context, groupID *int64) map[string]struct{} {
platforms := make(map[string]struct{})
if s == nil || s.accountRepo == nil {
return platforms
}
var accounts []Account
var err error
if groupID != nil {
accounts, err = s.accountRepo.ListSchedulableByGroupID(ctx, *groupID)
} else {
accounts, err = s.accountRepo.ListSchedulable(ctx)
}
if err != nil {
return platforms
}
for _, acc := range accounts {
platform := strings.TrimSpace(acc.Platform)
if platform != "" {
platforms[platform] = struct{}{}
}
}
return platforms
}
func (s *GatewayService) InvalidateAvailableModelsCache(groupID *int64, platform string) {
if s == nil || s.modelsListCache == nil {
return