203 lines
8.8 KiB
Go
203 lines
8.8 KiB
Go
package middleware
|
|
|
|
import (
|
|
"fmt"
|
|
"net/http"
|
|
"net/http/httptest"
|
|
"testing"
|
|
|
|
"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"))
|
|
}
|
|
|
|
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 TestNoAvailableChannelMessageNamesClaimingTaskPlugin(t *testing.T) {
|
|
require.NoError(t, i18n.Init())
|
|
registry := jsplugin.NewRegistry()
|
|
plugin, err := registry.Register(distributorTaskPluginSource("claimer", constant.ChannelTypeKling), jsplugin.Options{})
|
|
require.NoError(t, err)
|
|
|
|
pinned, _ := gin.CreateTestContext(nil)
|
|
pinned.Request = httptest.NewRequest(http.MethodPost, "/v1/videos", nil)
|
|
pinned.Request.Header.Set("Accept-Language", "en")
|
|
pinned.Set(jsplugin.ContextKeyPinnedPlugin, jsplugin.PinnedPlugin{Generation: registry.Generation(), Plugin: plugin})
|
|
message := noAvailableChannelMessage(pinned, "default", "kling-v1")
|
|
assert.Contains(t, message, `"claimer"`)
|
|
assert.Contains(t, message, "disable or override")
|
|
assert.Contains(t, message, "kling-v1")
|
|
|
|
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")
|
|
}
|