fix(models): list composite exact route aliases

This commit is contained in:
wucm667
2026-10-01 22:54:05 +08:00
parent d6adebd22d
commit 1888de135d
4 changed files with 118 additions and 5 deletions
+13 -3
View File
@@ -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
}
@@ -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)
@@ -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)
@@ -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{} {