feat: support upstream model fetch for advanced custom channels (#5971)

* feat: support upstream model fetch for advanced custom channels * fix: add advanced custom routes as separate groups * fix: select advanced custom route entry before adding --------- Co-authored-by: CaIon <i@caion.me>
2026-07-18 13:39:53 +08:00
parent 57746fc972
commit a6cf42c0f1
24 changed files with 1324 additions and 135 deletions
@@ -1,9 +1,11 @@
 package controller

 import (
+	"bytes"
+	"errors"
 	"net/http"
 	"net/http/httptest"
-	"strings"
+	"net/url"
 	"testing"

 	"github.com/QuantumNous/new-api/common"
@@ -14,6 +16,340 @@ import (
 	"github.com/stretchr/testify/require"
 )

+func newAdvancedCustomModelListChannel(baseURL string, key string, upstreamPath string, auth *dto.AdvancedCustomRouteAuth) *model.Channel {
+	config := &dto.AdvancedCustomConfig{
+		Routes: []dto.AdvancedCustomRoute{
+			{
+				IncomingPath: dto.AdvancedCustomModelListPath,
+				UpstreamPath: upstreamPath,
+				Converter:    "none",
+				Auth:         auth,
+			},
+		},
+	}
+	channel := &model.Channel{
+		Type:    constant.ChannelTypeAdvancedCustom,
+		Key:     key,
+		BaseURL: &baseURL,
+	}
+	channel.SetOtherSettings(dto.ChannelOtherSettings{AdvancedCustom: config})
+	return channel
+}
+
+func TestParseOpenAIModelIDsStrictResponseContract(t *testing.T) {
+	tests := []struct {
+		name      string
+		body      string
+		want      []string
+		wantError string
+	}{
+		{name: "malformed JSON", body: `{"data":`, wantError: "invalid OpenAI Models response"},
+		{name: "missing data", body: `{"object":"list"}`, wantError: "data is required"},
+		{name: "null data", body: `{"data":null}`, wantError: "data is required"},
+		{name: "empty data", body: `{"data":[]}`, wantError: "no valid model IDs"},
+		{name: "all IDs empty", body: `{"data":[{"id":""},{"id":"   "}]}`, wantError: "no valid model IDs"},
+		{
+			name: "filters empty IDs and normalizes valid IDs",
+			body: `{"data":[{"id":" gpt-4.1 "},{"id":""},{"id":"gpt-4.1"},{"id":"o3"}]}`,
+			want: []string{"gpt-4.1", "o3"},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			models, err := parseOpenAIModelIDs([]byte(test.body))
+			if test.wantError != "" {
+				require.ErrorContains(t, err, test.wantError)
+				require.Nil(t, models)
+				return
+			}
+			require.NoError(t, err)
+			require.Equal(t, test.want, models)
+		})
+	}
+}
+
+func TestFetchAdvancedCustomModelsAppliesHeaderOverrideAfterRouteAuth(t *testing.T) {
+	type receivedRequest struct {
+		Headers http.Header
+		Host    string
+	}
+	received := make(chan receivedRequest, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		received <- receivedRequest{Headers: r.Header.Clone(), Host: r.Host}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1"}]}`))
+	}))
+	defer server.Close()
+
+	channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/provider/models", &dto.AdvancedCustomRouteAuth{
+		Type:  dto.AdvancedCustomAuthTypeHeader,
+		Name:  "X-Route-Key",
+		Value: "route-{api_key}",
+	})
+	headerOverride := `{
+		"X-Route-Key":"global-{api_key}",
+		"X-Static":"static-value",
+		"X-Client":"{client_header:X-Client}",
+		"Host":"models.example.test",
+		"*":""
+	}`
+	channel.HeaderOverride = &headerOverride
+
+	models, err := fetchChannelUpstreamModelIDs(channel)
+	require.NoError(t, err)
+	require.Equal(t, []string{"gpt-4.1"}, models)
+
+	request := <-received
+	require.Equal(t, "global-secret-key", request.Headers.Get("X-Route-Key"))
+	require.Equal(t, "static-value", request.Headers.Get("X-Static"))
+	require.Empty(t, request.Headers.Get("X-Client"))
+	require.Equal(t, "models.example.test", request.Host)
+}
+
+func TestFetchAdvancedCustomModelsUsesEnabledSavedMultiKey(t *testing.T) {
+	authorization := make(chan string, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		authorization <- r.Header.Get("Authorization")
+		_, _ = w.Write([]byte(`{"data":[{"id":"gpt-4.1-mini"}]}`))
+	}))
+	defer server.Close()
+
+	channel := newAdvancedCustomModelListChannel(server.URL, "disabled-key\nenabled-key", "/v1/models", nil)
+	channel.ChannelInfo = model.ChannelInfo{
+		IsMultiKey: true,
+		MultiKeyStatusList: map[int]int{
+			0: common.ChannelStatusManuallyDisabled,
+			1: common.ChannelStatusEnabled,
+		},
+	}
+
+	models, err := fetchChannelUpstreamModelIDs(channel)
+	require.NoError(t, err)
+	require.Equal(t, []string{"gpt-4.1-mini"}, models)
+	require.Equal(t, "Bearer enabled-key", <-authorization)
+}
+
+func TestFetchAdvancedCustomModelsRejectsNonOKResponse(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusBadGateway)
+		_, _ = w.Write([]byte(`{"data":[{"id":"must-not-be-used"}]}`))
+	}))
+	defer server.Close()
+
+	channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/v1/models", nil)
+	models, err := fetchChannelUpstreamModelIDs(channel)
+	require.ErrorContains(t, err, "status code: 502")
+	require.Nil(t, models)
+}
+
+func TestFetchAdvancedCustomModelsRedactsQueryKeyFromTransportErrors(t *testing.T) {
+	const secret = "secret key/+"
+	server := httptest.NewServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {}))
+	baseURL := server.URL
+	server.Close()
+
+	channel := newAdvancedCustomModelListChannel(baseURL, secret, "/v1/models", &dto.AdvancedCustomRouteAuth{
+		Type:  dto.AdvancedCustomAuthTypeQuery,
+		Name:  "custom-token",
+		Value: "prefix-{api_key}",
+	})
+
+	_, err := fetchChannelUpstreamModelIDs(channel)
+	require.Error(t, err)
+	require.NotContains(t, err.Error(), secret)
+	require.NotContains(t, err.Error(), "custom-token")
+	require.NotContains(t, err.Error(), "prefix-")
+
+	direct := sanitizeFetchModelsError(&url.Error{
+		Op:  http.MethodGet,
+		URL: baseURL + "/v1/models?custom-token=prefix-" + url.QueryEscape(secret),
+		Err: errors.New("connection refused"),
+	}, secret)
+	require.EqualError(t, direct, "connection refused")
+}
+
+func TestFetchOrdinaryOpenAIModelsKeepsExistingEmptyDataBehavior(t *testing.T) {
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = w.Write([]byte(`{"object":"list"}`))
+	}))
+	defer server.Close()
+
+	baseURL := server.URL
+	channel := &model.Channel{
+		Type:    constant.ChannelTypeOpenAI,
+		Key:     "ordinary-key",
+		BaseURL: &baseURL,
+	}
+	models, err := fetchChannelUpstreamModelIDs(channel)
+	require.NoError(t, err)
+	require.Empty(t, models)
+}
+
+func TestFetchModelsAdvancedCustomCreatePreview(t *testing.T) {
+	receivedAuthorization := make(chan string, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		receivedAuthorization <- r.Header.Get("Authorization")
+		_, _ = w.Write([]byte(`{"data":[{"id":"preview-model"}]}`))
+	}))
+	defer server.Close()
+
+	config := dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
+		IncomingPath: dto.AdvancedCustomModelListPath,
+		UpstreamPath: "/preview/models",
+		Converter:    "none",
+	}}}
+	configBytes, err := common.Marshal(config)
+	require.NoError(t, err)
+	rawConfig := string(configBytes)
+	baseURL := server.URL
+	emptyProxy := ""
+	req := fetchModelsRequest{
+		BaseURL:        &baseURL,
+		Type:           constant.ChannelTypeAdvancedCustom,
+		Key:            "create-preview-key",
+		AdvancedCustom: &rawConfig,
+		Proxy:          &emptyProxy,
+	}
+	body, err := common.Marshal(req)
+	require.NoError(t, err)
+
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", bytes.NewReader(body))
+	ctx.Request.Header.Set("Content-Type", "application/json")
+	FetchModels(ctx)
+
+	var response struct {
+		Success bool     `json:"success"`
+		Message string   `json:"message"`
+		Data    []string `json:"data"`
+	}
+	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
+	require.True(t, response.Success, response.Message)
+	require.Equal(t, []string{"preview-model"}, response.Data)
+	require.Equal(t, "Bearer create-preview-key", <-receivedAuthorization)
+}
+
+func TestFetchModelsAdvancedCustomEditPreviewUsesSavedKeyAndExplicitClears(t *testing.T) {
+	db := setupModelListControllerTestDB(t)
+	receivedHeaders := make(chan http.Header, 1)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		receivedHeaders <- r.Header.Clone()
+		_, _ = w.Write([]byte(`{"data":[{"id":"edited-preview-model"}]}`))
+	}))
+	defer server.Close()
+
+	savedChannel := newAdvancedCustomModelListChannel("http://127.0.0.1:1", "disabled-saved-key\nenabled-saved-key", "/saved/models", nil)
+	savedChannel.Name = "saved advanced channel"
+	savedChannel.Models = "old-model"
+	savedChannel.ChannelInfo = model.ChannelInfo{
+		IsMultiKey: true,
+		MultiKeyStatusList: map[int]int{
+			0: common.ChannelStatusManuallyDisabled,
+			1: common.ChannelStatusEnabled,
+		},
+	}
+	savedHeaderOverride := `{"X-Saved":"must-not-be-sent"}`
+	savedChannel.HeaderOverride = &savedHeaderOverride
+	savedChannel.SetSetting(dto.ChannelSettings{Proxy: "http://127.0.0.1:1"})
+	require.NoError(t, db.Create(savedChannel).Error)
+
+	preserved, err := buildAdvancedCustomModelPreviewChannel(fetchModelsRequest{ChannelID: savedChannel.Id})
+	require.NoError(t, err)
+	require.Equal(t, "http://127.0.0.1:1", preserved.GetBaseURL())
+	require.Equal(t, savedHeaderOverride, *preserved.HeaderOverride)
+	require.Equal(t, "http://127.0.0.1:1", preserved.GetSetting().Proxy)
+
+	previewConfig := dto.AdvancedCustomConfig{Routes: []dto.AdvancedCustomRoute{{
+		IncomingPath: dto.AdvancedCustomModelListPath,
+		UpstreamPath: "/edited/models",
+		Converter:    "none",
+	}}}
+	configBytes, err := common.Marshal(previewConfig)
+	require.NoError(t, err)
+	rawConfig := string(configBytes)
+	baseURL := server.URL
+	explicitEmpty := ""
+	req := fetchModelsRequest{
+		ChannelID:      savedChannel.Id,
+		BaseURL:        &baseURL,
+		Type:           constant.ChannelTypeAdvancedCustom,
+		Key:            "request-key-must-be-ignored",
+		AdvancedCustom: &rawConfig,
+		HeaderOverride: &explicitEmpty,
+		Proxy:          &explicitEmpty,
+	}
+	cleared, err := buildAdvancedCustomModelPreviewChannel(fetchModelsRequest{
+		ChannelID:      savedChannel.Id,
+		BaseURL:        &explicitEmpty,
+		AdvancedCustom: &rawConfig,
+		HeaderOverride: &explicitEmpty,
+		Proxy:          &explicitEmpty,
+	})
+	require.NoError(t, err)
+	require.NotNil(t, cleared.BaseURL)
+	require.Empty(t, *cleared.BaseURL)
+	require.NotNil(t, cleared.HeaderOverride)
+	require.Empty(t, *cleared.HeaderOverride)
+	require.Empty(t, cleared.GetSetting().Proxy)
+
+	body, err := common.Marshal(req)
+	require.NoError(t, err)
+
+	recorder := httptest.NewRecorder()
+	ctx, _ := gin.CreateTestContext(recorder)
+	ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", bytes.NewReader(body))
+	ctx.Request.Header.Set("Content-Type", "application/json")
+	FetchModels(ctx)
+
+	var response struct {
+		Success bool     `json:"success"`
+		Message string   `json:"message"`
+		Data    []string `json:"data"`
+	}
+	require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response))
+	require.True(t, response.Success, response.Message)
+	require.Equal(t, []string{"edited-preview-model"}, response.Data)
+	require.NotContains(t, recorder.Body.String(), "enabled-saved-key")
+	require.NotContains(t, recorder.Body.String(), "request-key-must-be-ignored")
+
+	headers := <-receivedHeaders
+	require.Equal(t, "Bearer enabled-saved-key", headers.Get("Authorization"))
+	require.Empty(t, headers.Get("X-Saved"))
+}
+
+func TestFailedAdvancedCustomDetectionDoesNotStageFullRemoval(t *testing.T) {
+	db := setupModelListControllerTestDB(t)
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = w.Write([]byte(`{"data":[]}`))
+	}))
+	defer server.Close()
+
+	channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/v1/models", nil)
+	channel.Name = "empty discovery response"
+	channel.Models = "gpt-4.1,o3"
+	settings := channel.GetOtherSettings()
+	settings.UpstreamModelUpdateCheckEnabled = true
+	settings.UpstreamModelUpdateAutoSyncEnabled = true
+	channel.SetOtherSettings(settings)
+	require.NoError(t, db.Create(channel).Error)
+
+	modelsChanged, autoAdded, err := checkAndPersistChannelUpstreamModelUpdates(channel, &settings, true, true)
+	require.ErrorContains(t, err, "no valid model IDs")
+	require.False(t, modelsChanged)
+	require.Zero(t, autoAdded)
+	require.Empty(t, settings.UpstreamModelUpdateLastDetectedModels)
+	require.Empty(t, settings.UpstreamModelUpdateLastRemovedModels)
+
+	reloaded, err := model.GetChannelById(channel.Id, true)
+	require.NoError(t, err)
+	persistedSettings := reloaded.GetOtherSettings()
+	require.Empty(t, persistedSettings.UpstreamModelUpdateLastDetectedModels)
+	require.Empty(t, persistedSettings.UpstreamModelUpdateLastRemovedModels)
+	require.Equal(t, "gpt-4.1,o3", reloaded.Models)
+}
+
 func TestFetchModelsUsesSharedChannelFetchBehavior(t *testing.T) {
 	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if r.URL.Path != "/v1/models" {
@@ -39,7 +375,7 @@ func TestFetchModelsUsesSharedChannelFetchBehavior(t *testing.T) {

 	recorder := httptest.NewRecorder()
 	ctx, _ := gin.CreateTestContext(recorder)
-	ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", strings.NewReader(string(body)))
+	ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/fetch_models", bytes.NewReader(body))
 	ctx.Request.Header.Set("Content-Type", "application/json")

 	FetchModels(ctx)