middleware/distributor_test.go
Rama main · réplica de lectura del repositorio original.
package middleware
import (
"bytes"
"fmt"
"net/http"
"net/http/httptest"
"strings"
"testing"
"github.com/QuantumNous/new-api/common"
"github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/i18n"
"github.com/QuantumNous/new-api/model"
"github.com/QuantumNous/new-api/pkg/jsplugin"
"github.com/QuantumNous/new-api/relaykit/dto"
"github.com/QuantumNous/new-api/setting/model_setting"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
func TestChannelMatchesExpectedTaskPluginUsesGenericChannelSetting(t *testing.T) {
channel := &model.Channel{Type: constant.ChannelTypeTaskPlugin}
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "generic-alpha"})
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "generic-alpha"))
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, "generic-beta"))
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, ""))
}
func TestChannelMatchesExpectedTaskPluginUsesPinnedLegacyIndex(t *testing.T) {
registry := jsplugin.NewRegistry()
alpha, err := registry.Register(distributorTaskPluginSource("legacy-alpha", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
pinnedGeneration := registry.Generation()
require.NoError(t, registry.Unregister("legacy-alpha"))
_, err = registry.Register(distributorTaskPluginSource("legacy-beta", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: pinnedGeneration,
Plugin: alpha,
})
channel := &model.Channel{Type: constant.ChannelTypeKling}
assert.True(t, channelMatchesExpectedTaskPlugin(c, channel, "legacy-alpha"))
assert.False(t, channelMatchesExpectedTaskPlugin(c, channel, "legacy-beta"))
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "legacy-alpha"))
}
func TestChannelMatchesExpectedTaskPluginRejectsUnindexedLegacyChannel(t *testing.T) {
registry := jsplugin.NewRegistry()
plugin, err := registry.Register(distributorTaskPluginSource("legacy-alpha", constant.ChannelTypeKling), jsplugin.Options{})
require.NoError(t, err)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{
Generation: registry.Generation(),
Plugin: plugin,
})
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "legacy-alpha"))
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: 0}, "legacy-alpha"))
assert.True(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeJimeng}, ""))
assert.False(t, channelMatchesExpectedTaskPlugin(nil, &model.Channel{Type: constant.ChannelTypeKling}, "legacy-alpha"))
c.Set("expected_task_plugin_key", "legacy-alpha")
setupErr := SetupContextForSelectedChannel(c, &model.Channel{Type: constant.ChannelTypeJimeng}, "task-model")
require.NotNil(t, setupErr)
assert.Contains(t, setupErr.Error(), "does not match")
}
func TestSharedEndpointRebindsToSelectedLegacyProvider(t *testing.T) {
registry := jsplugin.NewRegistry()
_, err := registry.Register(distributorEndpointPluginSource("gemini-shared", constant.ChannelTypeGemini), jsplugin.Options{})
require.NoError(t, err)
_, err = registry.Register(distributorEndpointPluginSource("vertex-shared", constant.ChannelTypeVertexAi), jsplugin.Options{})
require.NoError(t, err)
candidates := registry.Generation().LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, 2)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: registry.Generation(), Plugin: candidates[0].Plugin})
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
Generation: registry.Generation(),
Plugin: candidates[0].Plugin,
Protocol: candidates[0].Protocol,
Operation: candidates[0].Operation,
Model: "task-model",
Candidates: candidates,
})
c.Set("expected_task_plugin_key", candidates[0].Plugin.Meta.Key)
geminiChannel := &model.Channel{Id: 1, Type: constant.ChannelTypeGemini}
vertexChannel := &model.Channel{Id: 2, Type: constant.ChannelTypeVertexAi}
assert.True(t, channelMatchesExpectedTaskPlugin(c, geminiChannel, candidates[0].Plugin.Meta.Key))
assert.True(t, channelMatchesExpectedTaskPlugin(c, vertexChannel, candidates[0].Plugin.Meta.Key))
assert.False(t, channelMatchesExpectedTaskPlugin(c, &model.Channel{Type: constant.ChannelTypeKling}, candidates[0].Plugin.Meta.Key))
require.Nil(t, SetupContextForSelectedChannel(c, vertexChannel, "task-model"))
pinnedValue, exists := c.Get(jsplugin.ContextKeyPinnedEndpoint)
require.True(t, exists)
pinned, ok := pinnedValue.(jsplugin.PinnedEndpoint)
require.True(t, ok)
assert.Equal(t, "vertex-shared", pinned.Plugin.Meta.Key)
assert.Equal(t, "vertex-shared", c.GetString("expected_task_plugin_key"))
assert.Equal(t, "vertex-shared", c.GetString("task_plugin_key"))
assert.True(t, channelMatchesExpectedTaskPlugin(c, geminiChannel, "vertex-shared"), "a retry may select another declared provider")
}
func distributorTaskPluginSource(key string, channelType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
channelTypes: [%d],
models: ["task-model"],
fetchMode: "per_task",
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {taskId: "task"}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
`, key, key, channelType)
}
func distributorEndpointPluginSource(key string, channelType int) string {
return fmt.Sprintf(`
export const meta = {
apiVersion: 1,
key: %q,
name: %q,
version: "1.0.0",
author: {name: "Test"},
channelTypes: [%d],
models: ["task-model"],
fetchMode: "per_task",
protocols: [{name: "openai_responses", supports: ["stream", "sync", "background"]}],
};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {taskId: "task"}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {status: "SUCCESS"}; }
export const protocols = {openai_responses: {
decodeRequest: function(ctx) { return {kind: "submit", model: "task-model", requestBody: ctx.body.value}; },
renderEvents: function() { return {events: [], state: null, done: false}; },
renderFinal: function() { return {output: []}; },
}};
`, key, key, channelType)
}
func TestTokenModelLimitAllowsLegacyAliasAndModifierVariant(t *testing.T) {
aliasOnly := map[string]bool{"claude-3-7-sonnet-thinking": true}
assert.True(t, TokenModelLimitAllows(aliasOnly, "claude-3-7-sonnet-thinking"))
assert.False(t, TokenModelLimitAllows(aliasOnly, "claude-3-7-sonnet"))
baseOnly := map[string]bool{"claude-3-7-sonnet": true}
assert.True(t, TokenModelLimitAllows(baseOnly, "claude-3-7-sonnet@thinking:on"))
assert.True(t, TokenModelLimitAllows(baseOnly, "claude-3-7-sonnet-thinking"))
wildcard := map[string]bool{"gemini-2.5-flash-thinking-*": true}
assert.True(t, TokenModelLimitAllows(wildcard, "gemini-2.5-flash-thinking-8192"))
}
// /v1/tools/search and /v1/tools/fetch always route and bill by the DeepSeek
// routing names. A model named in the body is ignored: a chat model there
// would price the call at zero tokens.
func TestGetModelRequestAssignsDeepSeekToolRoutingNames(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, tc := range []struct{ path, body, want string }{
{"/v1/tools/search", `{"query":"deepseek","count":3}`, constant.DeepSeekSearchModel},
{"/v1/tools/fetch", `{"url":"https://example.com/"}`, constant.DeepSeekFetchModel},
{"/v1/tools/search", `{"query":"deepseek","model":"deepseek-flash"}`, constant.DeepSeekSearchModel},
{"/v1/tools/fetch", `{"url":"https://example.com/","model":"deepseek-v4-pro"}`, constant.DeepSeekFetchModel},
} {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
c.Request = httptest.NewRequest(http.MethodPost, tc.path, strings.NewReader(tc.body))
c.Request.Header.Set("Content-Type", "application/json")
modelRequest, shouldSelectChannel, err := getModelRequest(c)
require.NoError(t, err, tc.body)
assert.True(t, shouldSelectChannel, tc.body)
assert.Equal(t, tc.want, modelRequest.Model, tc.body)
}
}
func TestTokenModelLimitAllowsExemptAtNameByFullName(t *testing.T) {
settings := model_setting.GetGlobalSettings()
original := append([]string(nil), settings.ThinkingModelBlacklist...)
t.Cleanup(func() { settings.ThinkingModelBlacklist = original })
settings.ThinkingModelBlacklist = append(original, "re:.*@sha256:.*")
fullOnly := map[string]bool{"opaque@sha256:deadbeef": true}
assert.True(t, TokenModelLimitAllows(fullOnly, "opaque@sha256:deadbeef"))
baseOnly := map[string]bool{"opaque": true}
assert.False(t, TokenModelLimitAllows(baseOnly, "opaque@sha256:deadbeef"))
}
func TestDistributeHidesTaskPluginDetailsButLogsDiagnostics(t *testing.T) {
require.NoError(t, i18n.Init())
previousCacheEnabled := common.MemoryCacheEnabled
common.MemoryCacheEnabled = true
t.Cleanup(func() { common.MemoryCacheEnabled = previousCacheEnabled })
const group = "private-plugin-error-test-group"
for _, locale := range []struct{ language, message string }{
{"en", "No available channel for model task-model under group " + group + ": the model is claimed by a task plugin, which has no enabled channel serving it (distributor)"},
{"zh-CN", "分组 " + group + " 下模型 task-model 无可用渠道：该模型由任务插件认领，但当前没有启用的渠道可服务此模型（distributor）"},
{"zh-TW", "分組 " + group + " 下模型 task-model 無可用管道：該模型由任務插件認領，但目前沒有啟用的管道可服務此模型（distributor）"},
} {
for _, providerCount := range []int{1, 2} {
t.Run(fmt.Sprintf("%s/%d_providers", locale.language, providerCount), func(t *testing.T) {
registry := jsplugin.NewRegistry()
keys := []string{"private-provider-alpha", "private-provider-beta"}[:providerCount]
for index, key := range keys {
_, err := registry.Register(distributorEndpointPluginSource(key, constant.ChannelTypeKling+index), jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, providerCount)
var logs bytes.Buffer
common.LogWriterMu.Lock()
previousWriter := gin.DefaultErrorWriter
gin.DefaultErrorWriter = &logs
common.LogWriterMu.Unlock()
t.Cleanup(func() {
common.LogWriterMu.Lock()
gin.DefaultErrorWriter = previousWriter
common.LogWriterMu.Unlock()
})
router := gin.New()
router.POST("/v1/responses", RequestId(), func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyUsingGroup, group)
c.Set("resolved_task_model", "task-model")
c.Set("expected_task_plugin_key", keys[0])
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: generation, Plugin: candidates[0].Plugin})
if providerCount > 1 {
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{
Generation: generation, Plugin: candidates[0].Plugin,
Protocol: candidates[0].Protocol, Operation: candidates[0].Operation,
Model: "task-model", Candidates: candidates,
})
}
}, Distribute(), func(c *gin.Context) {
t.Error("unavailable requests must stop before the relay handler")
c.Status(http.StatusNoContent)
})
request := httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
request.Header.Set("Accept-Language", locale.language)
recorder := httptest.NewRecorder()
router.ServeHTTP(recorder, request)
require.Equal(t, http.StatusServiceUnavailable, recorder.Code)
requestID := recorder.Header().Get(common.RequestIdKey)
require.NotEmpty(t, requestID)
assert.JSONEq(t, fmt.Sprintf(`{"error":{"message":%q,"type":"new_api_error","code":"model_not_found"}}`,
locale.message+" (request id: "+requestID+")"), recorder.Body.String())
assert.NotContains(t, recorder.Body.String(), "disable or override")
for _, key := range keys {
assert.NotContains(t, recorder.Body.String(), key)
assert.Contains(t, logs.String(), key)
}
assert.Contains(t, logs.String(), requestID)
assert.Contains(t, logs.String(), `group="`+group+`"`)
assert.Contains(t, logs.String(), `model="task-model"`)
assert.Contains(t, logs.String(), "reason=no_eligible_channel")
})
}
}
}
func TestNoAvailableChannelMessageWithoutPlugin(t *testing.T) {
require.NoError(t, i18n.Init())
plain, _ := gin.CreateTestContext(nil)
plain.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", nil)
plain.Request.Header.Set("Accept-Language", "en")
generic := noAvailableChannelMessage(plain, "default", "gpt-4o")
assert.NotContains(t, generic, "task plugin")
assert.Contains(t, generic, "gpt-4o")
}
func TestSharedEndpointRebindsToSelectedType61Plugin(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"alpha", "beta"} {
source := strings.Replace(distributorEndpointPluginSource(key, 0), "channelTypes: [0],", "", 1)
_, err := registry.Register(source, jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, 2)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: generation, Plugin: candidates[0].Plugin})
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{Generation: generation, Plugin: candidates[0].Plugin, Protocol: candidates[0].Protocol, Operation: candidates[0].Operation, Model: "task-model", Candidates: candidates})
c.Set("expected_task_plugin_key", "alpha")
channel := &model.Channel{Id: 2, Type: constant.ChannelTypeTaskPlugin}
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "unrelated"})
assert.False(t, channelMatchesExpectedTaskPlugin(c, channel, "alpha"))
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "beta"})
require.Nil(t, SetupContextForSelectedChannel(c, channel, "task-model"))
assert.Equal(t, "beta", c.GetString("task_plugin_key"))
assert.Equal(t, "beta", c.GetString("expected_task_plugin_key"))
assert.Equal(t, "beta", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.Key)
}
// On a native path two plugins declare, the selected channel names the
// executing plugin, which takes over the route and the task identity.
func TestSharedNativeRouteRebindsToSelectedPlugin(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"route-alpha", "route-beta"} {
_, err := registry.Register(fmt.Sprintf(`
export const meta = {apiVersion: 1, key: %q, name: %q, version: "1.0.0", author: {name: "Test"}, models: ["task-model"], fetchMode: "per_task",
routes: [{method: "POST", path: "/probe/v1/jobs", type: "submit", decode: "decode", render: "render"}]};
export const native = {decode() { return {}; }, render() { return {}; }};
export function buildSubmitRequest() { return {}; }
export function parseSubmitResponse() { return {}; }
export function buildQueryRequest() { return {}; }
export function parseTaskResult() { return {}; }
`, key, key), jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupDeclaredRoutes("POST", "/probe/v1/jobs")
require.Len(t, candidates, 2)
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: generation, Plugin: candidates[0].Plugin})
c.Set(jsplugin.ContextKeyPinnedRoute, jsplugin.PinnedRoute{Generation: generation, Plugin: candidates[0].Plugin, Route: candidates[0].Route, Candidates: candidates})
c.Set("expected_task_plugin_key", "route-alpha")
channel := &model.Channel{Id: 3, Type: constant.ChannelTypeTaskPlugin}
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "unrelated"})
assert.False(t, channelMatchesExpectedTaskPlugin(c, channel, "route-alpha"))
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "route-beta"})
require.Nil(t, SetupContextForSelectedChannel(c, channel, "task-model"))
pinned := c.MustGet(jsplugin.ContextKeyPinnedRoute).(jsplugin.PinnedRoute)
assert.Same(t, candidates[1].Plugin, pinned.Plugin)
assert.Equal(t, candidates[1].Route, pinned.Route)
assert.Same(t, candidates[1].Plugin, c.MustGet(jsplugin.ContextKeyPinnedPlugin).(jsplugin.PinnedPlugin).Plugin)
for _, key := range []string{"expected_task_plugin_key", "task_plugin_key", "platform"} {
assert.Equal(t, "route-beta", c.GetString(key), key)
}
}
func TestChannelMatchesExpectedTaskPluginUsesNewAPIExtensions(t *testing.T) {
channel := &model.Channel{Type: constant.ChannelTypeNewAPI}
channel.SetSetting(dto.ChannelSettings{TaskPluginKey: "single", TaskExtendPluginKeys: []string{"ext-alpha", "ext-beta"}})
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "ext-alpha"))
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, "single"), "the single key stays valid on a New API channel")
assert.False(t, channelMatchesExpectedTaskPlugin(nil, channel, "ext-gamma"))
assert.True(t, channelMatchesExpectedTaskPlugin(nil, channel, ""), "requests without a pinned plugin still use the gateway")
assert.False(t, channelMatchesExpectedTaskPlugin(nil, &model.Channel{Type: constant.ChannelTypeNewAPI}, "ext-alpha"))
}
func TestSharedEndpointRebindsToBoundNewAPIExtension(t *testing.T) {
registry := jsplugin.NewRegistry()
for _, key := range []string{"alpha", "beta"} {
source := strings.Replace(distributorEndpointPluginSource(key, 0), "channelTypes: [0],", "", 1)
_, err := registry.Register(source, jsplugin.Options{})
require.NoError(t, err)
}
generation := registry.Generation()
candidates := generation.LookupEndpointCandidates("POST", "/v1/responses", "task-model")
require.Len(t, candidates, 2)
pin := func(expected string) *gin.Context {
pinned := candidates[0]
if candidates[1].Plugin.Meta.Key == expected {
pinned = candidates[1]
}
c, _ := gin.CreateTestContext(nil)
c.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: generation, Plugin: pinned.Plugin})
c.Set(jsplugin.ContextKeyPinnedEndpoint, jsplugin.PinnedEndpoint{Generation: generation, Plugin: pinned.Plugin, Protocol: pinned.Protocol, Operation: pinned.Operation, Model: "task-model", Candidates: candidates})
c.Set("expected_task_plugin_key", expected)
return c
}
channel := &model.Channel{Id: 3, Type: constant.ChannelTypeNewAPI}
channel.SetSetting(dto.ChannelSettings{TaskExtendPluginKeys: []string{"unrelated"}})
assert.False(t, channelMatchesExpectedTaskPlugin(pin("alpha"), channel, "alpha"))
channel.SetSetting(dto.ChannelSettings{TaskExtendPluginKeys: []string{"beta"}})
c := pin("alpha")
require.Nil(t, SetupContextForSelectedChannel(c, channel, "task-model"))
assert.Equal(t, "beta", c.GetString("task_plugin_key"), "the only bound candidate executes")
assert.Equal(t, "beta", c.GetString("expected_task_plugin_key"))
assert.Equal(t, "beta", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.Key)
channel.SetSetting(dto.ChannelSettings{TaskExtendPluginKeys: []string{"alpha", "beta"}})
c = pin("beta")
require.Nil(t, SetupContextForSelectedChannel(c, channel, "task-model"))
assert.Equal(t, "alpha", c.GetString("task_plugin_key"), "the first bound candidate executes regardless of the earlier pin")
assert.Equal(t, "alpha", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.Key)
}