feat(gateway): retrieve models from the visible catalogue

This commit is contained in:
r266-tech
2026-09-10 15:57:17 +08:00
parent 98d86915be
commit 1c0932e173
5 changed files with 239 additions and 27 deletions
+8 -26
View File
@@ -1116,8 +1116,8 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
}
}
// Models handles listing available models
// GET /v1/models
// Models lists visible models, or retrieves the exact list entry for a model path parameter.
// GET /v1/models and /v1/models/:model (also exposed through root aliases)
// Returns models based on account configurations (model_mapping whitelist)
// Falls back to default models if no whitelist is configured
func (h *GatewayHandler) Models(c *gin.Context) {
@@ -1173,18 +1173,12 @@ func (h *GatewayHandler) Models(c *gin.Context) {
// Fallback to default models
if platform == service.PlatformOpenAI {
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": openai.DefaultModels,
})
writeModelsListResponse(c, openai.DefaultModels)
return
}
if platform == service.PlatformGemini {
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": geminicli.DefaultModels,
})
writeModelsListResponse(c, geminicli.DefaultModels)
return
}
if platform == service.PlatformGrok {
@@ -1192,10 +1186,7 @@ func (h *GatewayHandler) Models(c *gin.Context) {
return
}
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": claude.DefaultModels,
})
writeModelsListResponse(c, claude.DefaultModels)
}
// CodexModels returns the effective group model list using the manifest shape
@@ -1320,10 +1311,7 @@ func writeModelsList(c *gin.Context, platform string, modelIDs []string) {
CreatedAt: "2024-01-01T00:00:00Z",
})
}
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": models,
})
writeModelsListResponse(c, models)
}
func writeAllowlistedModelsList(c *gin.Context, platform string, modelIDs []string) {
@@ -1382,10 +1370,7 @@ func writeGrokModelsList(c *gin.Context, modelIDs []string) {
models = append(models, item)
}
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": models,
})
writeModelsListResponse(c, models)
}
func grokModelSupportsConfigurableReasoning(modelID string) bool {
@@ -1418,10 +1403,7 @@ func writeOpenAIModelsList(c *gin.Context, modelIDs []string) {
DisplayName: modelID,
})
}
c.JSON(http.StatusOK, gin.H{
"object": "list",
"data": models,
})
writeModelsListResponse(c, models)
}
// modelListingSource 汇总模型列表过滤的候选来源:账号映射键(availableModels)
@@ -0,0 +1,114 @@
package handler
import (
"encoding/json"
"net/http"
"net/http/httptest"
"testing"
middleware2 "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"
)
func requestModelForTest(h *GatewayHandler, group *service.Group, modelID, etag string) *httptest.ResponseRecorder {
rec := httptest.NewRecorder()
c, _ := gin.CreateTestContext(rec)
c.Request = httptest.NewRequest(http.MethodGet, "/v1/models", nil)
c.Request.Header.Set("If-None-Match", etag)
if modelID != "" {
c.Params = gin.Params{{Key: "model", Value: modelID}}
}
c.Set(string(middleware2.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
h.Models(c)
return rec
}
func TestRetrieveModelMatchesVisibleCatalogue(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, platform := range []string{service.PlatformOpenAI, service.PlatformAnthropic, service.PlatformGemini, service.PlatformGrok, service.PlatformComposite} {
for _, mapped := range []bool{false, true} {
name := platform + "/fallback"
if mapped {
name = platform + "/mapped"
}
t.Run(name, func(t *testing.T) {
accountPlatform := platform
if platform == service.PlatformComposite {
accountPlatform = service.PlatformOpenAI
}
credentials := map[string]any{}
if mapped {
credentials["model_mapping"] = map[string]any{"custom-model": "upstream-only-model"}
}
group := &service.Group{ID: 71, Platform: platform}
h := newGatewayModelsHandlerForTest(&gatewayModelsAccountRepoStub{byGroup: map[int64][]service.Account{
group.ID: {{ID: 1, Platform: accountPlatform, Type: service.AccountTypeAPIKey, Status: service.StatusActive, Schedulable: true, Credentials: credentials}},
}})
list := requestModelForTest(h, group, "", "")
require.Equal(t, http.StatusOK, list.Code, list.Body.String())
var catalog struct {
Data []json.RawMessage `json:"data"`
}
require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
require.NotEmpty(t, catalog.Data)
var model struct {
ID string `json:"id"`
}
require.NoError(t, json.Unmarshal(catalog.Data[0], &model))
retrieved := requestModelForTest(h, group, model.ID, "")
require.Equal(t, http.StatusOK, retrieved.Code, retrieved.Body.String())
require.JSONEq(t, string(catalog.Data[0]), retrieved.Body.String())
require.Equal(t, http.StatusNotFound, requestModelForTest(h, group, "unknown-model", "").Code)
group.ModelAllowlist = service.GroupModelAllowlist{Enabled: true, Models: []string{"absent-from-source"}}
hidden := requestModelForTest(h, group, model.ID, "")
require.Equal(t, http.StatusNotFound, hidden.Code, hidden.Body.String())
require.Contains(t, hidden.Body.String(), `"code":"model_not_found"`)
})
}
}
}
func TestRetrievePinnedModelPreservesMetadataFilteringAndErrors(t *testing.T) {
for _, tc := range []struct {
name string
status int
selected []string
wantStatus int
}{
{name: "metadata", wantStatus: http.StatusOK},
{name: "hidden", selected: []string{"other-model"}, wantStatus: http.StatusNotFound},
{name: "upstream failure", status: http.StatusServiceUnavailable, wantStatus: http.StatusBadGateway},
} {
t.Run(tc.name, func(t *testing.T) {
upstream := &codexModelsPinnedHTTPUpstream{bodies: map[int64]string{
2: `{"data":[{"id":"special-model","ID":"not-a-model-id","owned_by":"source-owner","created":123,"extra":{"context":999}}]}`,
}, statuses: map[int64]int{}}
if tc.status != 0 {
upstream.statuses[2] = tc.status
}
codex := newPinnedCodexTestHandler([]service.Account{newPinnedCodexAccount(2, service.StatusActive, true, false)}, upstream, 3)
h := &GatewayHandler{openAIGatewayService: codex.gatewayService, maxAccountSwitches: 3}
group := &service.Group{ID: 72, Platform: service.PlatformOpenAI,
ModelAllowlist: service.GroupModelAllowlist{Enabled: len(tc.selected) > 0, Models: tc.selected},
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{2}}}
got := requestModelForTest(h, group, "special-model", "collection-etag")
require.Equal(t, tc.wantStatus, got.Code, got.Body.String())
if tc.wantStatus == http.StatusOK {
list := requestModelForTest(h, group, "", "")
var catalog struct {
Data []json.RawMessage `json:"data"`
}
require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
require.JSONEq(t, string(catalog.Data[0]), got.Body.String())
require.Contains(t, got.Body.String(), `"owned_by":"source-owner"`)
require.Contains(t, got.Body.String(), `"created":123`)
again := requestModelForTest(h, group, "special-model", list.Header().Get("ETag"))
require.Equal(t, http.StatusOK, again.Code)
require.Empty(t, again.Header().Get("ETag"))
require.Equal(t, http.StatusNotFound, requestModelForTest(h, group, "not-a-model-id", "").Code)
}
})
}
}
@@ -1,7 +1,9 @@
package handler
import (
"encoding/json"
"errors"
"fmt"
"net/http"
infraerrors "github.com/Wei-Shaw/sub2api/internal/pkg/errors"
@@ -17,8 +19,12 @@ func (h *GatewayHandler) pinnedOpenAIModels(c *gin.Context, group *service.Group
writeOpenAIModelsError(c, http.StatusInternalServerError, "api_error", "OpenAI model discovery is not configured")
return
}
etag := c.GetHeader("If-None-Match")
if c.Param("model") != "" {
etag = "" // A collection ETag cannot validate a single-model representation.
}
response, account, err := h.openAIGatewayService.FetchPinnedOpenAIModelsList(
c.Request.Context(), group, h.maxAccountSwitches, c.GetHeader("If-None-Match"),
c.Request.Context(), group, h.maxAccountSwitches, etag,
)
if c.Request.Context().Err() != nil {
return
@@ -40,6 +46,10 @@ func writeOpenAIModelsError(c *gin.Context, status int, errorType, message strin
}
func writeOpenAIModelsResponse(c *gin.Context, manifest *service.OpenAIModelsResponse) {
if c.Param("model") != "" {
writeRetrievedModel(c, manifest.Body)
return
}
if manifest.ETag != "" {
c.Header("ETag", manifest.ETag)
}
@@ -50,3 +60,50 @@ func writeOpenAIModelsResponse(c *gin.Context, manifest *service.OpenAIModelsRes
}
c.Data(http.StatusOK, "application/json", manifest.Body)
}
// Both discovery endpoints consume the same final catalogue, after group/platform
// selection and allowlist filtering. Preserve every field on the selected entry.
func writeModelsListResponse(c *gin.Context, models any) {
response := gin.H{"object": "list", "data": models}
if c.Param("model") == "" {
c.JSON(http.StatusOK, response)
return
}
body, err := json.Marshal(response)
if err != nil {
writeOpenAIModelsError(c, http.StatusInternalServerError, "api_error", "Failed to encode model catalogue")
return
}
writeRetrievedModel(c, body)
}
func writeRetrievedModel(c *gin.Context, body []byte) {
var catalog struct {
Data []json.RawMessage `json:"data"`
}
if err := json.Unmarshal(body, &catalog); err != nil {
writeOpenAIModelsError(c, http.StatusBadGateway, "upstream_error", "Invalid model catalogue")
return
}
modelID := c.Param("model")
for _, raw := range catalog.Data {
var model map[string]json.RawMessage
if err := json.Unmarshal(raw, &model); err != nil {
writeOpenAIModelsError(c, http.StatusBadGateway, "upstream_error", "Invalid model catalogue entry")
return
}
var id string
if err := json.Unmarshal(model["id"], &id); err != nil {
writeOpenAIModelsError(c, http.StatusBadGateway, "upstream_error", "Invalid model catalogue ID")
return
}
if id == modelID {
c.Data(http.StatusOK, "application/json", raw)
return
}
}
c.JSON(http.StatusNotFound, gin.H{"error": gin.H{
"type": "invalid_request_error", "code": "model_not_found", "param": "model",
"message": fmt.Sprintf("Model %q does not exist or is not available for this group", modelID),
}})
}
@@ -208,6 +208,8 @@ func RegisterGatewayRoutes(
// /models endpoint with a client_version query and expect the ChatGPT
// Codex manifest format; other clients keep the OpenAI-style list.
gateway.GET("/models", modelsHandler)
// Single-model discovery never selects the Codex client_version manifest.
gateway.GET("/models/:model", h.Gateway.Models)
gateway.GET("/usage", h.Gateway.Usage)
gateway.POST("/live", h.OpenAIGateway.Live)
gateway.GET("/live/:call_id", h.OpenAIGateway.LiveSideband)
@@ -372,6 +374,7 @@ func RegisterGatewayRoutes(
h.OpenAIGateway.ResponsesWebSocket(c)
})
rootRoute(http.MethodGet, "/models", bodyLimit, modelsHandler)
rootRoute(http.MethodGet, "/models/:model", bodyLimit, h.Gateway.Models)
rootRoute(http.MethodPost, "/messages/count_tokens", bodyLimit, countTokensHandler)
codexDirect := r.Group("/backend-api/codex")
codexDirect.Use(bodyLimit, clientRequestID, opsErrorLogger, endpointNorm, gin.HandlerFunc(apiKeyAuth), groupModelAllowlist, compositeTarget, requireGroupAnthropic)
@@ -92,3 +92,59 @@ func TestGatewayRoutesPinnedModelsDispatchesOrdinaryAndCodexRequests(t *testing.
require.EqualValues(t, 1, upstream.ordinaryCalls.Load())
require.EqualValues(t, 1, upstream.codexCalls.Load())
}
func TestGatewayRoutesRetrievePinnedModel(t *testing.T) {
gin.SetMode(gin.TestMode)
repo := &pinnedModelsRoutesRepository{account: service.Account{
ID: 7, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
Status: service.StatusActive, Schedulable: true,
Credentials: map[string]any{"api_key": "test-models-key", "base_url": "https://models.example/v1"},
}}
upstream := &pinnedModelsRoutesUpstream{}
cfg := &config.Config{RunMode: config.RunModeSimple}
s := service.NewOpenAIGatewayService(repo, nil, nil, nil, nil, nil, nil, cfg,
nil, nil, nil, nil, nil, upstream, nil, nil, nil, nil, nil, nil, nil, nil)
h := &handler.Handlers{
Gateway: handler.NewGatewayHandler(nil, s, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, nil, cfg, nil),
OpenAIGateway: handler.NewOpenAIGatewayHandler(s, nil, nil, nil, nil, nil, nil, nil, cfg),
AsyncImage: handler.NewAsyncImageHandler(nil, nil),
}
group := &service.Group{ID: 1, Platform: service.PlatformOpenAI,
CodexModelsManifestConfig: service.GroupCodexModelsManifestConfig{Enabled: true, AccountIDs: []int64{7}}}
router := gin.New()
RegisterGatewayRoutes(router, h, servermiddleware.APIKeyAuthMiddleware(func(c *gin.Context) {
if c.GetHeader("Authorization") != "Bearer client-key" {
c.AbortWithStatus(http.StatusUnauthorized)
return
}
c.Set(string(servermiddleware.ContextKeyAPIKey), &service.APIKey{GroupID: &group.ID, Group: group})
c.Next()
}), nil, nil, nil, nil, nil, cfg)
request := func(path, key, etag string) *httptest.ResponseRecorder {
w := httptest.NewRecorder()
req := httptest.NewRequest(http.MethodGet, path, nil)
req.Header.Set("Authorization", key)
req.Header.Set("If-None-Match", etag)
router.ServeHTTP(w, req)
return w
}
list := request("/v1/models", "Bearer client-key", "")
require.Equal(t, http.StatusOK, list.Code)
var catalog struct {
Data []json.RawMessage `json:"data"`
}
require.NoError(t, json.Unmarshal(list.Body.Bytes(), &catalog))
require.Len(t, catalog.Data, 1)
for _, base := range []string{"/v1/models/", "/models/"} {
path := base + "ordinary-upstream-model?client_version=ignored"
require.Equal(t, http.StatusUnauthorized, request(path, "", "").Code)
got := request(path, "Bearer client-key", list.Header().Get("ETag"))
require.Equal(t, http.StatusOK, got.Code, got.Body.String())
require.JSONEq(t, string(catalog.Data[0]), got.Body.String())
require.Empty(t, got.Header().Get("ETag"), "a collection validator cannot validate one model")
missing := request(base+"unknown-model", "Bearer client-key", "")
require.Equal(t, http.StatusNotFound, missing.Code)
require.Contains(t, missing.Body.String(), `"error"`)
}
require.Zero(t, upstream.codexCalls.Load(), "retrieve never dispatches a Codex manifest")
}