feat(relay): exclude-driven channel failover, adaptive retry budget, and rate-limit cooldown circuit breaker

VERSION

@@ -0,0 +1 @@
+v1.0.0-rc.21-erdev.2

controller/record_channel_failure_test.go

@@ -0,0 +1,85 @@
+package controller
+
+import (
+	"net/http"
+	"testing"
+
+	"github.com/QuantumNous/new-api/relaykit/types"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// channelErr builds a NewAPIError with an explicit status code and no
+// rate-limit hint (RetryAfterSeconds == 0) — i.e. a channel-level failure.
+func channelErr(status int) *types.NewAPIError {
+	return types.NewErrorWithStatusCode(&stubErr{"upstream backend failure"},
+		types.ErrorCode("upstream_error"), status)
+}
+
+// rateLimitErr builds a 429 NewAPIError — a per-key rate-limit signal.
+func rateLimitErr() *types.NewAPIError {
+	return types.NewErrorWithStatusCode(&stubErr{"rate limited"},
+		types.ErrorCode("rate_limited"), http.StatusTooManyRequests)
+}
+
+// A channel-level failure (5xx, auth_unavailable, etc.) must exclude the whole
+// channel on the first failure, regardless of how many keys it has, so the
+// retry fails over immediately instead of burning the remaining keys against
+// the same broken upstream backend.
+func TestRecordChannelFailure_ChannelLevelExcludesImmediately(t *testing.T) {
+	for _, status := range []int{http.StatusInternalServerError, http.StatusBadGateway, http.StatusServiceUnavailable} {
+		excludeChannels := map[int]bool{}
+		channelTries := map[int]int{}
+		// A 4-key channel; a single channel-level error must still exclude it at once.
+		recordChannelFailure(94, 4, channelErr(status), excludeChannels, channelTries)
+		require.True(t, excludeChannels[94], "status %d must exclude the whole channel", status)
+		assert.Equal(t, 0, channelTries[94], "status %d must not rotate keys", status)
+	}
+}
+
+// A per-key rate-limit must NOT exclude a multi-key channel until every enabled
+// key has been throttled: each 429 advances the rotation counter, and only the
+// attempt that exhausts the last key excludes the channel.
+func TestRecordChannelFailure_RateLimitRotatesThenExcludes(t *testing.T) {
+	excludeChannels := map[int]bool{}
+	channelTries := map[int]int{}
+	const enabledKeys = 3
+
+	// First two throttled keys: rotate, do not exclude.
+	for i := 1; i <= 2; i++ {
+		recordChannelFailure(94, enabledKeys, rateLimitErr(), excludeChannels, channelTries)
+		require.False(t, excludeChannels[94], "after %d of %d keys the channel remains selectable", i, enabledKeys)
+		assert.Equal(t, i, channelTries[94])
+	}
+
+	// Third throttled key exhausts the channel: now exclude.
+	recordChannelFailure(94, enabledKeys, rateLimitErr(), excludeChannels, channelTries)
+	require.True(t, excludeChannels[94])
+}
+
+// A single-key channel behaves identically for both failure classes: one
+// failure exhausts it, so it is excluded immediately either way.
+func TestRecordChannelFailure_SingleKeyExcludesOnFirstFailure(t *testing.T) {
+	// Rate-limit on a single-key channel: 1 try >= 1 key -> excluded.
+	excludeChannels := map[int]bool{}
+	channelTries := map[int]int{}
+	recordChannelFailure(8, 1, rateLimitErr(), excludeChannels, channelTries)
+	require.True(t, excludeChannels[8])
+
+	// Channel-level error on a single-key channel: excluded, counter untouched.
+	excludeChannels = map[int]bool{}
+	channelTries = map[int]int{}
+	recordChannelFailure(8, 1, channelErr(http.StatusServiceUnavailable), excludeChannels, channelTries)
+	require.True(t, excludeChannels[8])
+}
+
+func TestRecordChannelFailure_RetryAfterOn5xxIsChannelLevel(t *testing.T) {
+	err := channelErr(http.StatusServiceUnavailable)
+	err.RetryAfterSeconds = 60
+	excludeChannels := map[int]bool{}
+	channelTries := map[int]int{}
+
+	recordChannelFailure(94, 4, err, excludeChannels, channelTries)
+	require.True(t, excludeChannels[94])
+	assert.Equal(t, 0, channelTries[94])
+}

controller/relay.go

@@ -181,17 +181,39 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		}
 	}()
 
+	// excludeChannels tracks channel IDs that should no longer be selected in
+	// this request so each retry moves to a fresh channel. recordChannelFailure
+	// populates it: a channel-level failure excludes the channel immediately,
+	// while a per-key rate-limit only excludes it once every enabled key has
+	// been throttled (counted via channelTries), preserving the per-request key
+	// rotation that GetNextEnabledKey performs.
+	excludeChannels := make(map[int]bool)
+	// channelTries counts rate-limited attempts per channel, compared against
+	// the channel's enabled-key count by recordChannelFailure.
+	channelTries := make(map[int]int)
 	retryParam := &service.RetryParam{
-		Ctx:         c,
-		TokenGroup:  relayInfo.TokenGroup,
-		ModelName:   relayInfo.OriginModelName,
-		RequestPath: c.Request.URL.Path,
-		Retry:       common.GetPointer(0),
+		Ctx:             c,
+		TokenGroup:      relayInfo.TokenGroup,
+		ModelName:       relayInfo.OriginModelName,
+		RequestPath:     c.Request.URL.Path,
+		Retry:           common.GetPointer(0),
+		ExcludeChannels: excludeChannels,
 	}
 	relayInfo.RetryIndex = 0
 	relayInfo.LastError = nil
 
-	for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
+	// Adaptive retry budget: try every available channel before giving up.
+	// A fixed RetryTimes smaller than the channel pool would abort while
+	// healthy channels remain untried. We size the cap to cover all channels
+	// (count-1 retries after the first attempt) but never below the configured
+	// RetryTimes. The exclude-driven selection stops cleanly once the pool is
+	// exhausted, so an oversized cap costs nothing.
+	retryCap := common.RetryTimes
+	if availableChannels := countAvailableChannelsForRetry(c, relayInfo); availableChannels-1 > retryCap {
+		retryCap = availableChannels - 1
+	}
+
+	for ; retryParam.GetRetry() <= retryCap; retryParam.IncreaseRetry() {
 		relayInfo.RetryIndex = retryParam.GetRetry()
 		channel, channelErr := getChannel(c, relayInfo, retryParam)
 		if channelErr != nil {
@@ -200,6 +222,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 			break
 		}
 		addUsedChannel(c, channel.Id)
+		// Channel exclusion bookkeeping (channelTries / excludeChannels) is done
+		// by recordChannelFailure after a failed attempt; nothing to do here on
+		// the success path.
 		if billingErr := service.PrepareTieredBillingForSelectedGroup(c, relayInfo); billingErr != nil {
 			newAPIError = billingErr
 			break
@@ -230,15 +255,18 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 
 		if newAPIError == nil {
 			relayInfo.LastError = nil
+			clearCooldownForContext(c, channel.Id)
 			return
 		}
 
 		newAPIError = service.NormalizeViolationFeeError(newAPIError)
 		relayInfo.LastError = newAPIError
 
 		processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
+		markCooldownFromError(c, channel.Id, newAPIError)
+		recordChannelFailure(channel.Id, channel.CountEnabledKeys(), newAPIError, excludeChannels, channelTries)
 
-		if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
+		if !shouldRetry(c, relayInfo, newAPIError, retryCap-retryParam.GetRetry()) {
 			break
 		}
 	}
@@ -297,6 +325,31 @@ func fastTokenCountMetaForPricing(request dto.Request) *types.TokenCountMeta {
 	return meta
 }
 
+// countAvailableChannelsForRetry returns how many distinct channels can serve
+// this request, used to size the adaptive retry budget. For the "auto" token
+// group it sums channels across all of the user's auto groups (retry fails over
+// across them); otherwise it counts the single token group. When a fixed
+// channel is pinned (specific_channel_id / ChannelMeta), there is nothing to
+// fail over to, so it returns 1.
+func countAvailableChannelsForRetry(c *gin.Context, info *relaycommon.RelayInfo) int {
+	if info.ChannelMeta != nil {
+		return 1
+	}
+	if _, ok := c.Get("specific_channel_id"); ok {
+		return 1
+	}
+	requestPath := c.Request.URL.Path
+	if info.TokenGroup == "auto" {
+		userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
+		total := 0
+		for _, g := range service.GetUserAutoGroup(userGroup) {
+			total += model.CountAvailableChannels(g, info.OriginModelName, requestPath)
+		}
+		return total
+	}
+	return model.CountAvailableChannels(info.TokenGroup, info.OriginModelName, requestPath)
+}
+
 func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service.RetryParam) (*model.Channel, *types.NewAPIError) {
 	if info.ChannelMeta == nil {
 		autoBan := c.GetBool("auto_ban")
@@ -328,10 +381,18 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
 	return channel, nil
 }
 
-func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) bool {
+func shouldRetry(c *gin.Context, info *relaycommon.RelayInfo, openaiErr *types.NewAPIError, retryTimes int) bool {
 	if openaiErr == nil {
 		return false
 	}
+	// HasSendResponse means the relay has recorded first-response timing. Some
+	// adaptors set it immediately before rendering, so treat it conservatively
+	// as the boundary after which a retry could append a second response to the
+	// same writer. This is intentionally stronger than claiming bytes are known
+	// to have been written on every adaptor path.
+	if info != nil && info.HasSendResponse() {
+		return false
+	}
 	if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
 		return false
 	}
@@ -360,6 +421,60 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b
 	return operation_setting.ShouldRetryByStatusCode(code)
 }
 
+// contextKeyIndex returns the multi-key index recorded for the current request,
+// or 0 for a single-key channel. It mirrors what SetupContextForSelectedChannel
+// stored from GetNextEnabledKey, so cooldown reads/writes target the exact key
+// that was used.
+func contextKeyIndex(c *gin.Context) int {
+	if common.GetContextKeyBool(c, constant.ContextKeyChannelIsMultiKey) {
+		return common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex)
+	}
+	return 0
+}
+
+// markCooldownFromError puts the failing channel key into a short cross-request
+// cooldown on a rate-limit response, so other requests skip a just-throttled
+// upstream instead of re-hitting it. The upstream Retry-After hint (if any)
+// sizes the cooldown; otherwise a small default is used. This is called only
+// from the relay loops (not from processChannelError) so channel health checks
+// never pollute live rate-limit state.
+func markCooldownFromError(c *gin.Context, channelId int, err *types.NewAPIError) {
+	if err == nil {
+		return
+	}
+	if err.IsRateLimited() {
+		model.MarkChannelKeyCooldown(channelId, contextKeyIndex(c), err.RetryAfterSeconds)
+	}
+}
+
+// clearCooldownForContext removes any cooldown on the key just used, called
+// after a successful request so a recovered upstream becomes immediately
+// selectable instead of waiting out the remaining TTL.
+func clearCooldownForContext(c *gin.Context, channelId int) {
+	model.ClearChannelKeyCooldown(channelId, contextKeyIndex(c))
+}
+
+// recordChannelFailure updates retry bookkeeping after a channel attempt fails,
+// deciding whether the next retry should rotate to another key of the same
+// channel or fail over to a different channel entirely.
+//
+// A channel-level failure (broken upstream backend: 5xx, auth_unavailable, etc.)
+// excludes the whole channel at once, since every key of that channel talks to
+// the same broken backend and retrying the remaining keys only adds latency. A
+// per-key rate-limit instead advances the key-rotation counter and excludes the
+// channel only once all of its enabled keys have been throttled. IsRateLimited
+// is the single source of truth for the distinction, shared with cooldown.
+func recordChannelFailure(channelID, enabledKeys int, err *types.NewAPIError, excludeChannels map[int]bool, channelTries map[int]int) {
+	if !err.IsRateLimited() {
+		excludeChannels[channelID] = true
+		return
+	}
+	channelTries[channelID]++
+	if channelTries[channelID] >= enabledKeys {
+		excludeChannels[channelID] = true
+	}
+}
+
 func processChannelError(c *gin.Context, channelError types.ChannelError, err *types.NewAPIError) {
 	logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error())))
 	// 不要使用context获取渠道信息，异步处理时可能会出现渠道信息不一致的情况
@@ -513,12 +628,19 @@ func RelayTask(c *gin.Context) {
 		}
 	}()
 
+	// excludeChannels / channelTries mirror the synchronous relay loop so task
+	// submission retries also move to a fresh channel (or the next key of a
+	// multi-key channel) instead of re-hitting one that just failed. The pinned
+	// LockedChannel branch is never excluded.
+	excludeChannels := make(map[int]bool)
+	channelTries := make(map[int]int)
 	retryParam := &service.RetryParam{
-		Ctx:         c,
-		TokenGroup:  relayInfo.TokenGroup,
-		ModelName:   relayInfo.OriginModelName,
-		RequestPath: c.Request.URL.Path,
-		Retry:       common.GetPointer(0),
+		Ctx:             c,
+		TokenGroup:      relayInfo.TokenGroup,
+		ModelName:       relayInfo.OriginModelName,
+		RequestPath:     c.Request.URL.Path,
+		Retry:           common.GetPointer(0),
+		ExcludeChannels: excludeChannels,
 	}
 
 	for ; retryParam.GetRetry() <= common.RetryTimes; retryParam.IncreaseRetry() {
@@ -556,14 +678,20 @@ func RelayTask(c *gin.Context) {
 
 		result, taskErr = relay.RelayTaskSubmit(c, relayInfo)
 		if taskErr == nil {
+			// Success: clear any cooldown on the key just used so a recovered
+			// upstream becomes immediately selectable again.
+			clearCooldownForContext(c, channel.Id)
 			break
 		}
 
 		if !taskErr.LocalError {
+			taskAPIErr := types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode)
 			processChannelError(c,
 				*types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey,
 					common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()),
-				types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode))
+				taskAPIErr)
+			markCooldownFromError(c, channel.Id, taskAPIErr)
+			recordChannelFailure(channel.Id, channel.CountEnabledKeys(), taskAPIErr, excludeChannels, channelTries)
 		}
 
 		if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {