feat(channel): optional adaptive balance with circuit breaker (default off)

.env.example

@@ -97,3 +97,13 @@ LINUX_DO_USER_ENDPOINT=https://connect.linux.do/api/user
 # 用于验证支付成功/取消回调URL的域名安全性
 # 示例: example.com,myapp.io 将允许 example.com, sub.example.com, myapp.io 等
 # TRUSTED_REDIRECT_DOMAINS=example.com,myapp.io
+
+# --- Adaptive channel balance (default OFF) ---
+# When enabled without shadow, selection uses score+circuit. Prefer shadow first.
+# ADAPTIVE_BALANCE_ENABLED=false
+# ADAPTIVE_BALANCE_SHADOW_MODE=false
+# CHANNEL_CIRCUIT_BREAKER_ENABLED=false
+# EWMA_ALPHA=0.1
+# MAX_CHANNEL_CONCURRENCY=10
+# CHANNEL_COOLDOWN_SECONDS=30
+# MAX_RETRY_CHANNELS=0

common/env.go

@@ -36,3 +36,16 @@ func GetEnvOrDefaultBool(env string, defaultValue bool) bool {
 	}
 	return b
 }
+
+func GetEnvOrDefaultFloat(env string, defaultValue float64) float64 {
+	if env == "" || os.Getenv(env) == "" {
+		return defaultValue
+	}
+	num, err := strconv.ParseFloat(os.Getenv(env), 64)
+	if err != nil {
+		SysError(fmt.Sprintf("failed to parse %s: %s, using default value: %g", env, err.Error(), defaultValue))
+		return defaultValue
+	}
+	return num
+}
+

common/init.go

@@ -1,6 +1,7 @@
 package common
 
 import (
+	"math"
 	"flag"
 	"fmt"
 	"log"
@@ -110,6 +111,25 @@ func InitEnv() {
 	RelayIdleConnTimeout = GetEnvOrDefault("RELAY_IDLE_CONN_TIMEOUT", 90)
 	RelayMaxIdleConns = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS", 500)
 	RelayMaxIdleConnsPerHost = GetEnvOrDefault("RELAY_MAX_IDLE_CONNS_PER_HOST", 100)
+	constant.AdaptiveBalanceEnabled = GetEnvOrDefaultBool("ADAPTIVE_BALANCE_ENABLED", false)
+	constant.AdaptiveBalanceShadowMode = GetEnvOrDefaultBool("ADAPTIVE_BALANCE_SHADOW_MODE", false)
+	constant.ChannelCircuitBreakerEnabled = GetEnvOrDefaultBool("CHANNEL_CIRCUIT_BREAKER_ENABLED", false)
+	constant.MaxRetryChannels = GetEnvOrDefault("MAX_RETRY_CHANNELS", 0)
+	constant.ChannelCooldownSeconds = GetEnvOrDefault("CHANNEL_COOLDOWN_SECONDS", 30)
+	constant.EwmaAlpha = GetEnvOrDefaultFloat("EWMA_ALPHA", 0.1)
+	constant.MaxChannelConcurrency = GetEnvOrDefault("MAX_CHANNEL_CONCURRENCY", 10)
+	if constant.ChannelCooldownSeconds <= 0 {
+		constant.ChannelCooldownSeconds = 30
+	}
+	if math.IsNaN(constant.EwmaAlpha) || constant.EwmaAlpha <= 0 || constant.EwmaAlpha > 1 {
+		constant.EwmaAlpha = 0.1
+	}
+	if constant.MaxChannelConcurrency < 0 {
+		constant.MaxChannelConcurrency = 10
+	}
+	if constant.MaxRetryChannels < 0 {
+		constant.MaxRetryChannels = 0
+	}
 
 	// Initialize string variables with GetEnvOrDefaultString
 	GeminiSafetySetting = GetEnvOrDefaultString("GEMINI_SAFETY_SETTING", "BLOCK_NONE")

constant/env.go

@@ -25,3 +25,12 @@ var TaskPricePatches []string
 // TrustedRedirectDomains is a list of trusted domains for redirect URL validation.
 // Domains support subdomain matching (e.g., "example.com" matches "sub.example.com").
 var TrustedRedirectDomains []string
+
+// Adaptive channel balance settings
+var AdaptiveBalanceEnabled bool
+var AdaptiveBalanceShadowMode bool
+var ChannelCircuitBreakerEnabled bool
+var MaxRetryChannels int
+var ChannelCooldownSeconds int
+var EwmaAlpha float64
+var MaxChannelConcurrency int

controller/relay.go

@@ -6,6 +6,7 @@ import (
 	"io"
 	"log"
 	"net/http"
+	"net/url"
 	"strings"
 	"time"
 
@@ -200,6 +201,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		addUsedChannel(c, channel.Id)
 		bodyStorage, bodyErr := common.GetBodyStorage(c)
 		if bodyErr != nil {
+			service.ReleaseAdaptiveCircuitPermit(c, channel.Id)
 			// Ensure consistent 413 for oversized bodies even when error occurs later (e.g., retry path)
 			if common.IsRequestBodyTooLargeError(bodyErr) || errors.Is(bodyErr, common.ErrRequestBodyTooLarge) {
 				newAPIError = types.NewErrorWithStatusCode(bodyErr, types.ErrorCodeReadRequestBodyFailed, http.StatusRequestEntityTooLarge, types.ErrOptionWithSkipRetry())
@@ -210,15 +212,52 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 		}
 		c.Request.Body = io.NopCloser(bodyStorage)
 
-		switch relayFormat {
-		case types.RelayFormatOpenAIRealtime:
-			newAPIError = relay.WssHelper(c, relayInfo)
-		case types.RelayFormatClaude:
-			newAPIError = relay.ClaudeHelper(c, relayInfo)
-		case types.RelayFormatGemini:
-			newAPIError = geminiRelayHandler(c, relayInfo)
-		default:
-			newAPIError = relayHandler(c, relayInfo)
+		attemptStart := time.Now()
+		service.IncChannelConcurrency(channel.Id)
+		// Always dec even if helper panics (CustomRecovery still runs after).
+		func() {
+			defer service.DecChannelConcurrency(channel.Id)
+			defer func() {
+				if r := recover(); r != nil {
+					service.ReleaseAdaptiveCircuitPermit(c, channel.Id)
+					panic(r)
+				}
+			}()
+			switch relayFormat {
+			case types.RelayFormatOpenAIRealtime:
+				newAPIError = relay.WssHelper(c, relayInfo)
+			case types.RelayFormatClaude:
+				newAPIError = relay.ClaudeHelper(c, relayInfo)
+			case types.RelayFormatGemini:
+				newAPIError = geminiRelayHandler(c, relayInfo)
+			default:
+				newAPIError = relayHandler(c, relayInfo)
+			}
+		}()
+		{
+			statusCode := http.StatusOK
+			var recErr error
+			if newAPIError != nil {
+				statusCode = newAPIError.StatusCode
+				if statusCode == 0 {
+					statusCode = http.StatusInternalServerError
+				}
+				recErr = newAPIError
+			}
+			// Prefer UsingGroup (resolved auto group) so score buckets match selection.
+			metricGroup := relayInfo.UsingGroup
+			if metricGroup == "" {
+				metricGroup = relayInfo.TokenGroup
+			}
+			service.RecordAdaptiveResult(
+				c,
+				channel.Id,
+				metricGroup,
+				relayInfo.OriginModelName,
+				statusCode,
+				time.Since(attemptStart),
+				recErr,
+			)
 		}
 
 		if newAPIError == nil {
@@ -250,12 +289,30 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
 
 var upgrader = websocket.Upgrader{
 	Subprotocols: []string{"realtime"}, // WS 握手支持的协议，如果有使用 Sec-WebSocket-Protocol，则必须在此声明对应的 Protocol TODO add other protocol
-	CheckOrigin: func(r *http.Request) bool {
-		return true // 允许跨域
-	},
+	CheckOrigin:  isRealtimeWebSocketOriginAllowed,
+}
+
+func isRealtimeWebSocketOriginAllowed(r *http.Request) bool {
+	if r == nil {
+		return false
+	}
+	originValue := strings.TrimSpace(r.Header.Get("Origin"))
+	if originValue == "" {
+		return true
+	}
+
+	origin, err := url.Parse(originValue)
+	if err != nil || origin.Host == "" || (origin.Scheme != "http" && origin.Scheme != "https") {
+		return false
+	}
+	if strings.EqualFold(origin.Host, r.Host) {
+		return true
+	}
+	return common.ValidateRedirectURL(originValue) == nil
 }
 
 func addUsedChannel(c *gin.Context, channelId int) {
+	service.MarkChannelUsed(c, channelId)
 	useChannel := c.GetStringSlice("use_channel")
 	useChannel = append(useChannel, fmt.Sprintf("%d", channelId))
 	c.Set("use_channel", useChannel)
@@ -317,6 +374,7 @@ func getChannel(c *gin.Context, info *relaycommon.RelayInfo, retryParam *service
 
 	newAPIError := middleware.SetupContextForSelectedChannel(c, channel, info.OriginModelName)
 	if newAPIError != nil {
+		service.ReleaseAdaptiveCircuitPermit(c, channel.Id)
 		return nil, newAPIError
 	}
 	return channel, nil
@@ -329,18 +387,24 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b
 	if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
 		return false
 	}
-	if types.IsChannelError(openaiErr) {
-		return true
+	if retryTimes <= 0 {
+		return false
 	}
-	if types.IsSkipRetryError(openaiErr) {
+	if _, ok := c.Get("specific_channel_id"); ok {
 		return false
 	}
-	if retryTimes <= 0 {
+	if openaiErr.GetErrorCode() == types.ErrorCodeGetChannelFailed {
 		return false
 	}
-	if _, ok := c.Get("specific_channel_id"); ok {
+	if types.IsSkipRetryError(openaiErr) {
 		return false
 	}
+	if isUpstreamChannelQuotaError(openaiErr) {
+		return true
+	}
+	if types.IsChannelError(openaiErr) {
+		return true
+	}
 	code := openaiErr.StatusCode
 	if code >= 200 && code < 300 {
 		return false
@@ -354,6 +418,44 @@ func shouldRetry(c *gin.Context, openaiErr *types.NewAPIError, retryTimes int) b
 	return operation_setting.ShouldRetryByStatusCode(code)
 }
 
+func isUpstreamChannelQuotaError(err *types.NewAPIError) bool {
+	if err == nil {
+		return false
+	}
+	code := strings.ToLower(strings.TrimSpace(string(err.GetErrorCode())))
+	if code == string(types.ErrorCodeInsufficientUserQuota) || code == string(types.ErrorCodePreConsumeTokenQuotaFailed) {
+		return false
+	}
+	if err.StatusCode == http.StatusPaymentRequired {
+		return true
+	}
+	for _, marker := range []string{
+		"insufficient_quota",
+		"quota_exceeded",
+		"billing_hard_limit_reached",
+		"insufficient_balance",
+		"insufficient_credits",
+	} {
+		if strings.Contains(code, marker) {
+			return true
+		}
+	}
+	message := strings.ToLower(err.Error())
+	for _, marker := range []string{
+		"insufficient quota",
+		"quota exceeded",
+		"insufficient balance",
+		"insufficient credit",
+		"额度不足",
+		"余额不足",
+	} {
+		if strings.Contains(message, marker) {
+			return true
+		}
+	}
+	return false
+}
+
 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获取渠道信息，异步处理时可能会出现渠道信息不一致的情况

controller/relay_origin_test.go

@@ -0,0 +1,43 @@
+package controller
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/QuantumNous/new-api/constant"
+	"github.com/stretchr/testify/require"
+)
+
+func TestRealtimeWebSocketOriginAllowed(t *testing.T) {
+	originalDomains := append([]string(nil), constant.TrustedRedirectDomains...)
+	constant.TrustedRedirectDomains = []string{"example.com"}
+	t.Cleanup(func() {
+		constant.TrustedRedirectDomains = originalDomains
+	})
+
+	tests := []struct {
+		name   string
+		origin string
+		host   string
+		want   bool
+	}{
+		{name: "missing origin", host: "api.internal", want: true},
+		{name: "same origin", origin: "https://api.internal", host: "api.internal", want: true},
+		{name: "trusted exact domain", origin: "https://example.com", host: "api.internal", want: true},
+		{name: "trusted subdomain", origin: "https://console.example.com", host: "api.internal", want: true},
+		{name: "untrusted domain", origin: "https://evil.example.net", host: "api.internal", want: false},
+		{name: "suffix spoof", origin: "https://fakeexample.com", host: "api.internal", want: false},
+		{name: "invalid scheme", origin: "file://example.com", host: "api.internal", want: false},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			request := httptest.NewRequest(http.MethodGet, "https://"+tt.host+"/v1/realtime", nil)
+			if tt.origin != "" {
+				request.Header.Set("Origin", tt.origin)
+			}
+			require.Equal(t, tt.want, isRealtimeWebSocketOriginAllowed(request))
+		})
+	}
+}

controller/relay_retry_test.go

@@ -0,0 +1,37 @@
+package controller
+
+import (
+	"errors"
+	"net/http"
+	"testing"
+
+	"github.com/QuantumNous/new-api/types"
+	"github.com/gin-gonic/gin"
+	"github.com/stretchr/testify/require"
+)
+
+func TestShouldRetryStopsAfterChannelSelectionFailure(t *testing.T) {
+	ctx, _ := gin.CreateTestContext(nil)
+	err := types.NewError(errors.New("no eligible channel"), types.ErrorCodeGetChannelFailed, types.ErrOptionWithSkipRetry())
+	require.False(t, shouldRetry(ctx, err, 2))
+}
+
+func TestShouldRetrySwitchesChannelOnUpstreamQuotaExhaustion(t *testing.T) {
+	ctx, _ := gin.CreateTestContext(nil)
+	err := types.WithOpenAIError(types.OpenAIError{
+		Message: "upstream account has insufficient balance",
+		Code:    "insufficient_quota",
+	}, http.StatusTooManyRequests)
+	require.True(t, shouldRetry(ctx, err, 2))
+}
+
+func TestShouldRetryDoesNotSwitchForLocalUserQuota(t *testing.T) {
+	ctx, _ := gin.CreateTestContext(nil)
+	err := types.NewErrorWithStatusCode(
+		errors.New("user quota insufficient"),
+		types.ErrorCodeInsufficientUserQuota,
+		http.StatusForbidden,
+		types.ErrOptionWithSkipRetry(),
+	)
+	require.False(t, shouldRetry(ctx, err, 2))
+}

model/ability.go

@@ -106,6 +106,10 @@ func getChannelQuery(group string, model string, retry int) (*gorm.DB, error) {
 }
 
 func GetChannel(group string, model string, retry int, requestPath string) (*Channel, error) {
+	return GetChannelExcluding(group, model, retry, requestPath, nil)
+}
+
+func GetChannelExcluding(group string, model string, retry int, requestPath string, excluded map[int]struct{}) (*Channel, error) {
 	var abilities []Ability
 
 	var err error = nil
@@ -122,6 +126,12 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha
 		return nil, err
 	}
 	abilities = filterAbilitiesByRequestPathAndModel(abilities, requestPath, model)
+	if len(excluded) > 0 {
+		abilities = lo.Filter(abilities, func(ability Ability, _ int) bool {
+			_, skip := excluded[ability.ChannelId]
+			return !skip
+		})
+	}
 	channel := Channel{}
 	if len(abilities) > 0 {
 		// Randomly choose one