diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index f83494ef9a..165442ddd0 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -1162,7 +1162,7 @@ func (h *GatewayHandler) Models(c *gin.Context) { } if platform == service.PlatformComposite { - availableModels := h.compositeAvailableModels(c.Request.Context(), groupID) + availableModels := h.compositeAvailableModels(c.Request.Context(), groupID, "") if apiKey != nil && apiKey.Group != nil && apiKey.Group.ModelAllowlistEnabled() { source := availableModels if len(source) == 0 { @@ -1257,7 +1257,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi platform = group.Platform } if platform == service.PlatformComposite { - availableModels := h.compositeAvailableModels(ctx, groupID) + availableModels := h.compositeAvailableModels(ctx, groupID, service.CompositeRouteEndpointResponses) fallbackModels := defaultCodexModelIDsForPlatform(service.PlatformComposite) if group.ModelAllowlistEnabled() { source := availableModels @@ -1283,7 +1283,7 @@ func (h *GatewayHandler) codexModelIDsForGroup(ctx context.Context, group *servi return fallbackModels } -func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64) []string { +func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID *int64, endpoint string) []string { if h == nil || h.gatewayService == nil { return nil } @@ -1311,6 +1311,16 @@ func (h *GatewayHandler) compositeAvailableModels(ctx context.Context, groupID * models = append(models, model) } } + // A route can expose a public ID that no account model mapping contains. + // On lookup failure, retain the existing account-derived catalog only. + if routeModels, err := h.gatewayService.GetCompositeRouteModels(ctx, groupID, endpoint); err == nil { + for _, model := range routeModels { + if _, ok := seen[model]; !ok { + seen[model] = struct{}{} + models = append(models, model) + } + } + } return models } diff --git a/backend/internal/handler/gateway_models_test.go b/backend/internal/handler/gateway_models_test.go index 0b49ca5f35..793f22742f 100644 --- a/backend/internal/handler/gateway_models_test.go +++ b/backend/internal/handler/gateway_models_test.go @@ -3,6 +3,7 @@ package handler import ( "context" "encoding/json" + "errors" "net/http" "net/http/httptest" "strings" @@ -22,6 +23,16 @@ type gatewayModelsAccountRepoStub struct { byGroup map[int64][]service.Account } +type gatewayModelsRouteRepoStub struct { + service.CompositeModelRouteRepository + routes []service.CompositeModelRoute + err error +} + +func (s *gatewayModelsRouteRepoStub) ListByGroup(_ context.Context, _ int64, _ bool) ([]service.CompositeModelRoute, error) { + return s.routes, s.err +} + type gatewayModelsResponseForTest struct { Object string `json:"object"` Data []gatewayModelItemForTest `json:"data"` @@ -81,16 +92,79 @@ func (s *gatewayModelsAccountRepoStub) ListModelAvailabilityCandidates(ctx conte return s.ListSchedulableByGroupID(ctx, *groupID) } -func newGatewayModelsHandlerForTest(repo service.AccountRepository) *GatewayHandler { +func newGatewayModelsHandlerForTest(repo service.AccountRepository, routeRepo ...service.CompositeModelRouteRepository) *GatewayHandler { + var resolver *service.CompositeRouteResolver + if len(routeRepo) > 0 { + resolver = service.NewCompositeRouteResolver(routeRepo[0]) + } return &GatewayHandler{ gatewayService: service.NewGatewayService( repo, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, - nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, + nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, resolver, nil, nil, ), } } +func TestGatewayCompositeRouteModelsAppearInBothCatalogs(t *testing.T) { + gin.SetMode(gin.TestMode) + const groupID int64 = 7784 + accounts := &gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{ + groupID: {{ID: 1, Platform: service.PlatformOpenAI, Credentials: map[string]any{ + "model_mapping": map[string]any{"mapped-model": "upstream"}, + }}}, + }} + routes := &gatewayModelsRouteRepoStub{routes: []service.CompositeModelRoute{ + {PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true}, + {PublicModel: "route-only", MatchType: service.CompositeRouteMatchExact, Enabled: true}, + {PublicModel: "mapped-model", MatchType: service.CompositeRouteMatchExact, Enabled: true}, + {PublicModel: "image-only", MatchType: service.CompositeRouteMatchExact, Endpoint: service.CompositeRouteEndpointImages, Enabled: true}, + {PublicModel: "prefix-", MatchType: service.CompositeRouteMatchPrefix, Enabled: true}, + {PublicModel: "disabled", MatchType: service.CompositeRouteMatchExact, Enabled: false}, + }} + for _, tc := range []struct { + name string + repo service.CompositeModelRouteRepository + allowlist []string + want []string + }{ + {"routes", routes, nil, []string{"mapped-model", "route-only", "image-only"}}, + {"allowlist", routes, []string{"route-only", "disabled", "unknown"}, []string{"route-only"}}, + {"lookup failure", &gatewayModelsRouteRepoStub{err: errors.New("unavailable")}, nil, []string{"mapped-model"}}, + {"no route repo", nil, nil, []string{"mapped-model"}}, + } { + t.Run(tc.name, func(t *testing.T) { + h := newGatewayModelsHandlerForTest(accounts, tc.repo) + group := &service.Group{ID: groupID, Platform: service.PlatformComposite} + if tc.allowlist != nil { + group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: tc.allowlist} + } + for _, endpoint := range []string{"models", "codex"} { + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil) + c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{Group: group}) + if endpoint == "models" { + h.Models(c) + var got gatewayModelsResponseForTest + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + require.Equal(t, tc.want, modelIDsForTest(got.Data)) + } else { + h.CodexModels(c) + var got codexModelsResponseForTest + require.NoError(t, json.Unmarshal(rec.Body.Bytes(), &got)) + want := tc.want + if tc.name == "routes" { + want = []string{"mapped-model", "route-only"} + } + require.ElementsMatch(t, want, codexModelSlugsForTest(got.Models)) + } + require.Equal(t, http.StatusOK, rec.Code) + } + }) + } +} + func TestDefaultModelIDsForCompositeIncludesAntigravityDefaults(t *testing.T) { antigravityIDs := defaultModelIDsForPlatform(service.PlatformAntigravity) require.NotEmpty(t, antigravityIDs) diff --git a/backend/internal/service/composite_route_resolver.go b/backend/internal/service/composite_route_resolver.go index f2e1928ace..450175c5a0 100644 --- a/backend/internal/service/composite_route_resolver.go +++ b/backend/internal/service/composite_route_resolver.go @@ -22,6 +22,27 @@ func (r *CompositeRouteResolver) SetModelOwnershipResolver(resolver CompositeMod } } +// ListExactPublicModels returns enabled, concrete route IDs suitable for a model catalog. +func (r *CompositeRouteResolver) ListExactPublicModels(ctx context.Context, groupID int64, endpoint string) ([]string, error) { + if r == nil || r.repo == nil || groupID <= 0 { + return nil, nil + } + routes, err := r.repo.ListByGroup(ctx, groupID, false) + if err != nil { + return nil, err + } + models := make([]string, 0, len(routes)) + for _, route := range routes { + if route.Enabled && route.MatchType == CompositeRouteMatchExact && + (endpoint == "" || normalizeCompositeRouteEndpoint(route.Endpoint) == CompositeRouteEndpointAny || normalizeCompositeRouteEndpoint(route.Endpoint) == endpoint) { + if model := strings.TrimSpace(route.PublicModel); model != "" { + models = append(models, model) + } + } + } + return models, nil +} + func (r *CompositeRouteResolver) Resolve(ctx context.Context, groupID int64, model, endpoint string) (CompositeRouteDecision, error) { model = strings.TrimSpace(model) endpoint = normalizeCompositeRouteEndpoint(endpoint) diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index eb782d9a1e..db0e095a86 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -1535,6 +1535,14 @@ func explicitModelMappingClaims(account Account, model string) bool { return ok && strings.TrimSpace(mapped) != "" } +// GetCompositeRouteModels returns public IDs from enabled exact composite routes. +func (s *GatewayService) GetCompositeRouteModels(ctx context.Context, groupID *int64, endpoint string) ([]string, error) { + if s == nil || s.compositeResolver == nil || groupID == nil { + return nil, nil + } + return s.compositeResolver.ListExactPublicModels(ctx, *groupID, endpoint) +} + // 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{} {