feat: publish sanitized custom.97 source

middleware/distributor_test.go

@@ -168,6 +168,27 @@ func TestTokenModelLimitAllowsLegacyAliasAndModifierVariant(t *testing.T) {
 	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...)
@@ -267,6 +288,14 @@ func TestNoAvailableChannelMessageWithoutPlugin(t *testing.T) {
 	assert.Contains(t, generic, "gpt-4o")
 }
 
+func TestNoAvailableGroupPoolMessage(t *testing.T) {
+	c, _ := gin.CreateTestContext(httptest.NewRecorder())
+	c.Set("pool_retry_after", 9)
+
+	assert.Equal(t, "分组池中无条件内可用分组", noAvailableGroupPoolMessage(c))
+	assert.Equal(t, "9", c.Writer.Header().Get("Retry-After"))
+}
+
 func TestSharedEndpointRebindsToSelectedType61Plugin(t *testing.T) {
 	registry := jsplugin.NewRegistry()
 	for _, key := range []string{"alpha", "beta"} {
@@ -290,3 +319,54 @@ func TestSharedEndpointRebindsToSelectedType61Plugin(t *testing.T) {
 	assert.Equal(t, "beta", c.GetString("expected_task_plugin_key"))
 	assert.Equal(t, "beta", c.MustGet(jsplugin.ContextKeyPinnedEndpoint).(jsplugin.PinnedEndpoint).Plugin.Meta.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)
+}