diff --git a/controller/relay.go b/controller/relay.go
index 1d14dcc6f880..c788d764b10b 100644
--- a/controller/relay.go
+++ b/controller/relay.go
@@ -229,6 +229,9 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) {
relayInfo.LastError = newAPIError
processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError)
+ if service.ClearChannelAffinityAfterFailure(c) {
+ logger.LogInfo(c, fmt.Sprintf("渠道亲和命中失败,已清理缓存:channel #%d", channel.Id))
+ }
if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) {
break
@@ -556,6 +559,9 @@ func RelayTask(c *gin.Context) {
*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))
+ if service.ClearChannelAffinityAfterFailure(c) {
+ logger.LogInfo(c, fmt.Sprintf("渠道亲和命中失败,已清理缓存:channel #%d", channel.Id))
+ }
}
if !shouldRetryTaskRelay(c, channel.Id, taskErr, common.RetryTimes-retryParam.GetRetry()) {
diff --git a/middleware/distributor.go b/middleware/distributor.go
index 258aebb57037..10244510cf51 100644
--- a/middleware/distributor.go
+++ b/middleware/distributor.go
@@ -101,11 +101,16 @@ func Distribute() func(c *gin.Context) {
}
}
- if preferredChannelID, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
+ if affinitySelection, found := service.GetPreferredChannelByAffinity(c, modelRequest.Model, usingGroup); found {
affinityUsable := false
- preferred, err := model.CacheGetChannel(preferredChannelID)
- if err == nil && preferred != nil && preferred.Status == common.ChannelStatusEnabled {
- if usingGroup == "auto" {
+ preferred, err := model.CacheGetChannel(affinitySelection.ChannelID)
+ if err == nil && preferred != nil {
+ if preferred.Status != common.ChannelStatusEnabled {
+ if service.ShouldSkipRetryAfterChannelAffinityFailure(c) {
+ abortWithOpenAiMessage(c, http.StatusForbidden, i18n.T(c, i18n.MsgDistributorAffinityChannelDisabled))
+ return
+ }
+ } else if usingGroup == "auto" {
userGroup := common.GetContextKeyString(c, constant.ContextKeyUserGroup)
autoGroups := service.GetUserAutoGroup(userGroup)
for _, g := range autoGroups {
@@ -114,7 +119,7 @@ func Distribute() func(c *gin.Context) {
common.SetContextKey(c, constant.ContextKeyAutoGroup, g)
channel = preferred
affinityUsable = true
- service.MarkChannelAffinityUsed(c, g, preferred.Id)
+ service.MarkChannelAffinityUsed(c, g, affinitySelection)
break
}
}
@@ -122,7 +127,7 @@ func Distribute() func(c *gin.Context) {
channel = preferred
selectGroup = usingGroup
affinityUsable = true
- service.MarkChannelAffinityUsed(c, usingGroup, preferred.Id)
+ service.MarkChannelAffinityUsed(c, usingGroup, affinitySelection)
}
}
if !affinityUsable && !service.ShouldKeepChannelAffinityOnChannelDisabled() {
@@ -449,10 +454,19 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode
common.SetContextKey(c, constant.ContextKeyChannelModelMapping, channel.GetModelMapping())
common.SetContextKey(c, constant.ContextKeyChannelStatusCodeMapping, channel.GetStatusCodeMapping())
- key, index, newAPIError := channel.GetNextEnabledKey()
- if newAPIError != nil {
- return newAPIError
+ var key string
+ var index int
+ var newAPIError *types.NewAPIError
+ if preferredKeyIndex, ok := service.GetChannelAffinityKeyIndex(c, channel.Id); ok {
+ key, index, newAPIError = channel.GetEnabledKeyByIndex(preferredKeyIndex)
+ }
+ if key == "" || newAPIError != nil {
+ key, index, newAPIError = channel.GetNextEnabledKey()
+ if newAPIError != nil {
+ return newAPIError
+ }
}
+ service.UpdateChannelAffinitySelectedKeyIndex(c, channel.Id, index)
if channel.ChannelInfo.IsMultiKey {
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index)
diff --git a/model/channel.go b/model/channel.go
index 78a1477c327e..56e2e27f0c64 100644
--- a/model/channel.go
+++ b/model/channel.go
@@ -196,6 +196,38 @@ func (channel *Channel) GetKeys() []string {
return keys
}
+func (channel *Channel) GetEnabledKeyByIndex(index int) (string, int, *types.NewAPIError) {
+ if !channel.ChannelInfo.IsMultiKey {
+ if index != 0 {
+ return "", 0, types.NewError(errors.New("invalid key index"), types.ErrorCodeChannelNoAvailableKey)
+ }
+ return channel.Key, 0, nil
+ }
+
+ keys := channel.GetKeys()
+ if len(keys) == 0 {
+ return "", 0, types.NewError(errors.New("no keys available"), types.ErrorCodeChannelNoAvailableKey)
+ }
+ if index < 0 || index >= len(keys) {
+ return "", 0, types.NewError(errors.New("invalid key index"), types.ErrorCodeChannelNoAvailableKey)
+ }
+
+ lock := GetChannelPollingLock(channel.Id)
+ lock.Lock()
+ defer lock.Unlock()
+
+ status := common.ChannelStatusEnabled
+ if channel.ChannelInfo.MultiKeyStatusList != nil {
+ if s, ok := channel.ChannelInfo.MultiKeyStatusList[index]; ok {
+ status = s
+ }
+ }
+ if status != common.ChannelStatusEnabled {
+ return "", 0, types.NewError(errors.New("key is disabled"), types.ErrorCodeChannelNoAvailableKey)
+ }
+ return keys[index], index, nil
+}
+
func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
// If not in multi-key mode, return the original key string directly.
if !channel.ChannelInfo.IsMultiKey {
diff --git a/model/channel_key_affinity_test.go b/model/channel_key_affinity_test.go
new file mode 100644
index 000000000000..281df2cf1010
--- /dev/null
+++ b/model/channel_key_affinity_test.go
@@ -0,0 +1,48 @@
+package model
+
+import (
+ "testing"
+
+ "github.com/QuantumNous/new-api/common"
+ "github.com/QuantumNous/new-api/constant"
+ "github.com/stretchr/testify/require"
+)
+
+func TestGetEnabledKeyByIndex(t *testing.T) {
+ channel := &Channel{
+ Id: 9001,
+ Key: "key-a\nkey-b\nkey-c",
+ ChannelInfo: ChannelInfo{
+ IsMultiKey: true,
+ MultiKeyMode: constant.MultiKeyModePolling,
+ MultiKeyPollingIndex: 0,
+ MultiKeyStatusList: map[int]int{
+ 1: common.ChannelStatusEnabled,
+ },
+ },
+ }
+
+ key, index, apiErr := channel.GetEnabledKeyByIndex(1)
+ require.Nil(t, apiErr)
+ require.Equal(t, "key-b", key)
+ require.Equal(t, 1, index)
+ require.Equal(t, 0, channel.ChannelInfo.MultiKeyPollingIndex)
+}
+
+func TestGetEnabledKeyByIndexDisabled(t *testing.T) {
+ channel := &Channel{
+ Id: 9002,
+ Key: "key-a\nkey-b",
+ ChannelInfo: ChannelInfo{
+ IsMultiKey: true,
+ MultiKeyStatusList: map[int]int{
+ 1: common.ChannelStatusManuallyDisabled,
+ },
+ },
+ }
+
+ key, index, apiErr := channel.GetEnabledKeyByIndex(1)
+ require.NotNil(t, apiErr)
+ require.Empty(t, key)
+ require.Equal(t, 0, index)
+}
diff --git a/service/channel_affinity.go b/service/channel_affinity.go
index 96ec13e248cc..c5f871b46256 100644
--- a/service/channel_affinity.go
+++ b/service/channel_affinity.go
@@ -10,6 +10,7 @@ import (
"time"
"github.com/QuantumNous/new-api/common"
+ "github.com/QuantumNous/new-api/constant"
"github.com/QuantumNous/new-api/dto"
"github.com/QuantumNous/new-api/pkg/cachex"
"github.com/QuantumNous/new-api/setting/operation_setting"
@@ -25,14 +26,17 @@ const (
ginKeyChannelAffinityMeta = "channel_affinity_meta"
ginKeyChannelAffinityLogInfo = "channel_affinity_log_info"
ginKeyChannelAffinitySkipRetry = "channel_affinity_skip_retry_on_failure"
+ ginKeyChannelAffinityKeyIndex = "channel_affinity_key_index"
+ ginKeyChannelAffinityChannelID = "channel_affinity_channel_id"
+ ginKeyChannelAffinityUsed = "channel_affinity_used"
- channelAffinityCacheNamespace = "new-api:channel_affinity:v1"
- channelAffinityUsageCacheStatsNamespace = "new-api:channel_affinity_usage_cache_stats:v1"
+ channelAffinityCacheNamespace = "new-api:channel_affinity:v2"
+ channelAffinityUsageCacheStatsNamespace = "new-api:channel_affinity_usage_cache_stats:v2"
)
var (
channelAffinityCacheOnce sync.Once
- channelAffinityCache *cachex.HybridCache[int]
+ channelAffinityCache *cachex.HybridCache[ChannelAffinitySelection]
channelAffinityUsageCacheStatsOnce sync.Once
channelAffinityUsageCacheStatsCache *cachex.HybridCache[ChannelAffinityUsageCacheCounters]
@@ -40,6 +44,11 @@ var (
channelAffinityRegexCache sync.Map // map[string]*regexp.Regexp
)
+type ChannelAffinitySelection struct {
+ ChannelID int `json:"channel_id"`
+ KeyIndex int `json:"key_index"`
+}
+
type channelAffinityMeta struct {
CacheKey string
TTLSeconds int
@@ -78,7 +87,7 @@ type ChannelAffinityCacheStats struct {
CacheAlgo string `json:"cache_algo"`
}
-func getChannelAffinityCache() *cachex.HybridCache[int] {
+func getChannelAffinityCache() *cachex.HybridCache[ChannelAffinitySelection] {
channelAffinityCacheOnce.Do(func() {
setting := operation_setting.GetChannelAffinitySetting()
capacity := setting.MaxEntries
@@ -90,15 +99,15 @@ func getChannelAffinityCache() *cachex.HybridCache[int] {
defaultTTLSeconds = 3600
}
- channelAffinityCache = cachex.NewHybridCache[int](cachex.HybridCacheConfig[int]{
+ channelAffinityCache = cachex.NewHybridCache[ChannelAffinitySelection](cachex.HybridCacheConfig[ChannelAffinitySelection]{
Namespace: cachex.Namespace(channelAffinityCacheNamespace),
Redis: common.RDB,
RedisEnabled: func() bool {
return common.RedisEnabled && common.RDB != nil
},
- RedisCodec: cachex.IntCodec{},
- Memory: func() *hot.HotCache[string, int] {
- return hot.NewHotCache[string, int](hot.LRU, capacity).
+ RedisCodec: cachex.JSONCodec[ChannelAffinitySelection]{},
+ Memory: func() *hot.HotCache[string, ChannelAffinitySelection] {
+ return hot.NewHotCache[string, ChannelAffinitySelection](hot.LRU, capacity).
WithTTL(time.Duration(defaultTTLSeconds) * time.Second).
WithJanitor().
Build()
@@ -547,10 +556,10 @@ func ApplyChannelAffinityOverrideTemplate(c *gin.Context, paramOverride map[stri
return mergedParam, true
}
-func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup string) (int, bool) {
+func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup string) (ChannelAffinitySelection, bool) {
setting := operation_setting.GetChannelAffinitySetting()
if setting == nil || !setting.Enabled {
- return 0, false
+ return ChannelAffinitySelection{}, false
}
path := ""
if c != nil && c.Request != nil && c.Request.URL != nil {
@@ -610,17 +619,17 @@ func GetPreferredChannelByAffinity(c *gin.Context, modelName string, usingGroup
})
cache := getChannelAffinityCache()
- channelID, found, err := cache.Get(cacheKeySuffix)
+ selection, found, err := cache.Get(cacheKeySuffix)
if err != nil {
common.SysError(fmt.Sprintf("channel affinity cache get failed: key=%s, err=%v", cacheKeyFull, err))
- return 0, false
+ return ChannelAffinitySelection{}, false
}
- if found {
- return channelID, true
+ if found && selection.ChannelID > 0 {
+ return selection, true
}
- return 0, false
+ return ChannelAffinitySelection{}, false
}
- return 0, false
+ return ChannelAffinitySelection{}, false
}
func ShouldSkipRetryAfterChannelAffinityFailure(c *gin.Context) bool {
@@ -673,7 +682,8 @@ func ShouldKeepChannelAffinityOnChannelDisabled() bool {
return setting.KeepOnChannelDisabled
}
-func MarkChannelAffinityUsed(c *gin.Context, selectedGroup string, channelID int) {
+func MarkChannelAffinityUsed(c *gin.Context, selectedGroup string, selection ChannelAffinitySelection) {
+ channelID := selection.ChannelID
if c == nil || channelID <= 0 {
return
}
@@ -682,6 +692,9 @@ func MarkChannelAffinityUsed(c *gin.Context, selectedGroup string, channelID int
return
}
c.Set(ginKeyChannelAffinitySkipRetry, meta.SkipRetry)
+ c.Set(ginKeyChannelAffinityChannelID, channelID)
+ c.Set(ginKeyChannelAffinityKeyIndex, selection.KeyIndex)
+ c.Set(ginKeyChannelAffinityUsed, true)
info := map[string]interface{}{
"reason": meta.RuleName,
"rule_name": meta.RuleName,
@@ -690,6 +703,7 @@ func MarkChannelAffinityUsed(c *gin.Context, selectedGroup string, channelID int
"model": meta.ModelName,
"request_path": meta.RequestPath,
"channel_id": channelID,
+ "key_index": selection.KeyIndex,
"key_source": meta.KeySourceType,
"key_key": meta.KeySourceKey,
"key_path": meta.KeySourcePath,
@@ -699,6 +713,90 @@ func MarkChannelAffinityUsed(c *gin.Context, selectedGroup string, channelID int
c.Set(ginKeyChannelAffinityLogInfo, info)
}
+func clearChannelAffinitySelectedContext(c *gin.Context) {
+ c.Set(ginKeyChannelAffinityChannelID, 0)
+ c.Set(ginKeyChannelAffinityKeyIndex, -1)
+ c.Set(ginKeyChannelAffinityUsed, false)
+ if anyInfo, ok := c.Get(ginKeyChannelAffinityLogInfo); ok {
+ if info, ok := anyInfo.(map[string]interface{}); ok {
+ info["cleared_on_failure"] = true
+ c.Set(ginKeyChannelAffinityLogInfo, info)
+ }
+ }
+}
+
+func ClearChannelAffinityAfterFailure(c *gin.Context) bool {
+ if c == nil {
+ return false
+ }
+ if ShouldSkipRetryAfterChannelAffinityFailure(c) {
+ return false
+ }
+ cacheKey, _, ok := getChannelAffinityContext(c)
+ if !ok {
+ return false
+ }
+
+ usedAffinity, _ := c.Get(ginKeyChannelAffinityUsed)
+ used, _ := usedAffinity.(bool)
+ clearChannelAffinitySelectedContext(c)
+ if !used {
+ return false
+ }
+
+ cache := getChannelAffinityCache()
+ deleted, err := cache.DeleteMany([]string{cacheKey})
+ if err != nil {
+ common.SysError(fmt.Sprintf("channel affinity cache delete after failure failed: key=%s, err=%v", cacheKey, err))
+ return false
+ }
+ if deleted[cache.FullKey(cacheKey)] {
+ return true
+ }
+ return false
+}
+
+func GetChannelAffinityKeyIndex(c *gin.Context, channelID int) (int, bool) {
+ if c == nil || channelID <= 0 {
+ return 0, false
+ }
+ matchedChannelID, ok := c.Get(ginKeyChannelAffinityChannelID)
+ if !ok {
+ return 0, false
+ }
+ id, ok := matchedChannelID.(int)
+ if !ok || id != channelID {
+ return 0, false
+ }
+ keyIndexAny, ok := c.Get(ginKeyChannelAffinityKeyIndex)
+ if !ok {
+ return 0, false
+ }
+ keyIndex, ok := keyIndexAny.(int)
+ if !ok || keyIndex < 0 {
+ return 0, false
+ }
+ return keyIndex, true
+}
+
+func UpdateChannelAffinitySelectedKeyIndex(c *gin.Context, channelID int, keyIndex int) {
+ if c == nil || channelID <= 0 || keyIndex < 0 {
+ return
+ }
+ if _, ok := getChannelAffinityMeta(c); !ok {
+ return
+ }
+ c.Set(ginKeyChannelAffinityChannelID, channelID)
+ c.Set(ginKeyChannelAffinityKeyIndex, keyIndex)
+ if anyInfo, ok := c.Get(ginKeyChannelAffinityLogInfo); ok {
+ if info, ok := anyInfo.(map[string]interface{}); ok {
+ info["channel_id"] = channelID
+ info["key_index"] = keyIndex
+ c.Set(ginKeyChannelAffinityLogInfo, info)
+ }
+ }
+}
+
func AppendChannelAffinityAdminInfo(c *gin.Context, adminInfo map[string]interface{}) {
if c == nil || adminInfo == nil {
return
@@ -723,6 +821,10 @@ func RecordChannelAffinity(c *gin.Context, channelID int) {
channelID = successChannelID
}
}
+ keyIndex := 0
+ if c != nil {
+ keyIndex = common.GetContextKeyInt(c, constant.ContextKeyChannelMultiKeyIndex)
+ }
cacheKey, ttlSeconds, ok := getChannelAffinityContext(c)
if !ok {
return
@@ -734,7 +836,11 @@ func RecordChannelAffinity(c *gin.Context, channelID int) {
ttlSeconds = 3600
}
cache := getChannelAffinityCache()
- if err := cache.SetWithTTL(cacheKey, channelID, time.Duration(ttlSeconds)*time.Second); err != nil {
+ selection := ChannelAffinitySelection{
+ ChannelID: channelID,
+ KeyIndex: keyIndex,
+ }
+ if err := cache.SetWithTTL(cacheKey, selection, time.Duration(ttlSeconds)*time.Second); err != nil {
common.SysError(fmt.Sprintf("channel affinity cache set failed: key=%s, err=%v", cacheKey, err))
}
}
diff --git a/service/channel_affinity_template_test.go b/service/channel_affinity_template_test.go
index fb703a24e720..1ee9a4e550cf 100644
--- a/service/channel_affinity_template_test.go
+++ b/service/channel_affinity_template_test.go
@@ -8,6 +8,8 @@ import (
"testing"
"time"
+ "github.com/QuantumNous/new-api/common"
+ "github.com/QuantumNous/new-api/constant"
relaycommon "github.com/QuantumNous/new-api/relay/common"
"github.com/QuantumNous/new-api/setting/operation_setting"
"github.com/gin-gonic/gin"
@@ -208,7 +210,10 @@ func TestGetPreferredChannelByAffinity_RequestHeaderKeySource(t *testing.T) {
cacheKeySuffix := buildChannelAffinityCacheKeySuffix(rule, "gpt-5", "default", affinityValue)
cache := getChannelAffinityCache()
- require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9528, time.Minute))
+ require.NoError(t, cache.SetWithTTL(cacheKeySuffix, ChannelAffinitySelection{
+ ChannelID: 9528,
+ KeyIndex: 2,
+ }, time.Minute))
t.Cleanup(func() {
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
})
@@ -225,9 +230,10 @@ func TestGetPreferredChannelByAffinity_RequestHeaderKeySource(t *testing.T) {
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", nil)
ctx.Request.Header.Set("X-Affinity-Key", affinityValue)
- channelID, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
+ selection, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
require.True(t, found)
- require.Equal(t, 9528, channelID)
+ require.Equal(t, 9528, selection.ChannelID)
+ require.Equal(t, 2, selection.KeyIndex)
meta, ok := getChannelAffinityMeta(ctx)
require.True(t, ok)
@@ -242,7 +248,10 @@ func TestClearCurrentChannelAffinityCache(t *testing.T) {
cacheKeySuffix := fmt.Sprintf("codex cli trace:default:clear-current-%d", time.Now().UnixNano())
cacheKeyFull := channelAffinityCacheNamespace + ":" + cacheKeySuffix
cache := getChannelAffinityCache()
- require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute))
+ require.NoError(t, cache.SetWithTTL(cacheKeySuffix, ChannelAffinitySelection{
+ ChannelID: 9527,
+ KeyIndex: -1,
+ }, time.Minute))
t.Cleanup(func() {
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
})
@@ -283,7 +292,10 @@ func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
cacheKeySuffix := buildChannelAffinityCacheKeySuffix(*codexRule, "gpt-5", "default", affinityValue)
cache := getChannelAffinityCache()
- require.NoError(t, cache.SetWithTTL(cacheKeySuffix, 9527, time.Minute))
+ require.NoError(t, cache.SetWithTTL(cacheKeySuffix, ChannelAffinitySelection{
+ ChannelID: 9527,
+ KeyIndex: 1,
+ }, time.Minute))
t.Cleanup(func() {
_, _ = cache.DeleteMany([]string{cacheKeySuffix})
})
@@ -293,9 +305,10 @@ func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
ctx.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(fmt.Sprintf(`{"prompt_cache_key":"%s"}`, affinityValue)))
ctx.Request.Header.Set("Content-Type", "application/json")
- channelID, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
+ selection, found := GetPreferredChannelByAffinity(ctx, "gpt-5", "default")
require.True(t, found)
- require.Equal(t, 9527, channelID)
+ require.Equal(t, 9527, selection.ChannelID)
+ require.Equal(t, 1, selection.KeyIndex)
baseOverride := map[string]interface{}{
"temperature": 0.2,
@@ -332,3 +345,88 @@ func TestChannelAffinityHitCodexTemplatePassHeadersEffective(t *testing.T) {
_, exists = info.RuntimeHeadersOverride["x-codex-turn-metadata"]
require.False(t, exists)
}
+
+func TestRecordChannelAffinityStoresChannelAndKeyIndex(t *testing.T) {
+ cacheKeySuffix := fmt.Sprintf("record-key-index-%d", time.Now().UnixNano())
+ ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
+ CacheKey: cacheKeySuffix,
+ TTLSeconds: 600,
+ RuleName: "record-key-index",
+ })
+ common.SetContextKey(ctx, constant.ContextKeyChannelId, 123)
+ common.SetContextKey(ctx, constant.ContextKeyChannelMultiKeyIndex, 4)
+
+ RecordChannelAffinity(ctx, 111)
+
+ cache := getChannelAffinityCache()
+ t.Cleanup(func() {
+ _, _ = cache.DeleteMany([]string{cacheKeySuffix})
+ })
+ selection, found, err := cache.Get(cacheKeySuffix)
+ require.NoError(t, err)
+ require.True(t, found)
+ require.Equal(t, 123, selection.ChannelID)
+ require.Equal(t, 4, selection.KeyIndex)
+}
+
+func TestClearChannelAffinityAfterFailureDeletesUsedCache(t *testing.T) {
+ cacheKeySuffix := fmt.Sprintf("clear-after-failure-%d", time.Now().UnixNano())
+ ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
+ CacheKey: cacheKeySuffix,
+ TTLSeconds: 600,
+ RuleName: "clear-after-failure",
+ SkipRetry: false,
+ })
+ cache := getChannelAffinityCache()
+ require.NoError(t, cache.SetWithTTL(cacheKeySuffix, ChannelAffinitySelection{
+ ChannelID: 321,
+ KeyIndex: 2,
+ }, time.Minute))
+ t.Cleanup(func() {
+ _, _ = cache.DeleteMany([]string{cacheKeySuffix})
+ })
+
+ MarkChannelAffinityUsed(ctx, "default", ChannelAffinitySelection{
+ ChannelID: 321,
+ KeyIndex: 2,
+ })
+
+ require.True(t, ClearChannelAffinityAfterFailure(ctx))
+
+ _, found, err := cache.Get(cacheKeySuffix)
+ require.NoError(t, err)
+ require.False(t, found)
+ _, ok := GetChannelAffinityKeyIndex(ctx, 321)
+ require.False(t, ok)
+}
+
+func TestClearChannelAffinityAfterFailureKeepsCacheWhenSkipRetry(t *testing.T) {
+ cacheKeySuffix := fmt.Sprintf("keep-skip-retry-%d", time.Now().UnixNano())
+ ctx := buildChannelAffinityTemplateContextForTest(channelAffinityMeta{
+ CacheKey: cacheKeySuffix,
+ TTLSeconds: 600,
+ RuleName: "keep-skip-retry",
+ SkipRetry: true,
+ })
+ cache := getChannelAffinityCache()
+ require.NoError(t, cache.SetWithTTL(cacheKeySuffix, ChannelAffinitySelection{
+ ChannelID: 654,
+ KeyIndex: 3,
+ }, time.Minute))
+ t.Cleanup(func() {
+ _, _ = cache.DeleteMany([]string{cacheKeySuffix})
+ })
+
+ MarkChannelAffinityUsed(ctx, "default", ChannelAffinitySelection{
+ ChannelID: 654,
+ KeyIndex: 3,
+ })
+
+ require.False(t, ClearChannelAffinityAfterFailure(ctx))
+
+ selection, found, err := cache.Get(cacheKeySuffix)
+ require.NoError(t, err)
+ require.True(t, found)
+ require.Equal(t, 654, selection.ChannelID)
+ require.Equal(t, 3, selection.KeyIndex)
+}
diff --git a/service/channel_affinity_usage_cache_test.go b/service/channel_affinity_usage_cache_test.go
index 64d3d715b547..b4be4ede5eaa 100644
--- a/service/channel_affinity_usage_cache_test.go
+++ b/service/channel_affinity_usage_cache_test.go
@@ -25,10 +25,14 @@ func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string)
return ctx
}
+func uniqueChannelAffinityStatsKey(t *testing.T, prefix string) string {
+ return fmt.Sprintf("%s_%s_%d", prefix, t.Name(), time.Now().UnixNano())
+}
+
func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) {
- ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
+ ruleName := uniqueChannelAffinityStatsKey(t, "rule")
usingGroup := "default"
- keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
+ keyFP := uniqueChannelAffinityStatsKey(t, "fp")
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
usage := &dto.Usage{
@@ -53,9 +57,9 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T)
}
func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) {
- ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
+ ruleName := uniqueChannelAffinityStatsKey(t, "rule")
usingGroup := "default"
- keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
+ keyFP := uniqueChannelAffinityStatsKey(t, "fp")
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
openAIUsage := &dto.Usage{
@@ -83,9 +87,9 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) {
}
func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) {
- ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano())
+ ruleName := uniqueChannelAffinityStatsKey(t, "rule")
usingGroup := "default"
- keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano())
+ keyFP := uniqueChannelAffinityStatsKey(t, "fp")
ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP)
usage := &dto.Usage{
diff --git a/web/default/src/features/system-settings/general/channel-affinity/cache-stats-dialog.tsx b/web/default/src/features/system-settings/general/channel-affinity/cache-stats-dialog.tsx
index 633b197b16c8..d0991d9c0766 100644
--- a/web/default/src/features/system-settings/general/channel-affinity/cache-stats-dialog.tsx
+++ b/web/default/src/features/system-settings/general/channel-affinity/cache-stats-dialog.tsx
@@ -38,6 +38,8 @@ interface Props {
using_group: string
key_hint: string
key_fp: string
+ channel_id?: number
+ key_index?: number
} | null
}
@@ -60,7 +62,12 @@ export function CacheStatsDialog(props: Props) {
setStats(null)
- getAffinityUsageCache(props.target)
+ getAffinityUsageCache({
+ rule_name: props.target.rule_name,
+ using_group: props.target.using_group,
+ key_hint: props.target.key_hint,
+ key_fp: props.target.key_fp,
+ })
.then((res) => {
if (seq !== seqRef.current) return
if (res.success) setStats((res.data as Record
+ {t('Channel key')}: {affinityKeyDisplay} +
+ )} )} diff --git a/web/default/src/features/usage-logs/components/dialogs/details-dialog.tsx b/web/default/src/features/usage-logs/components/dialogs/details-dialog.tsx index 2728f124efd7..b31200528ecc 100644 --- a/web/default/src/features/usage-logs/components/dialogs/details-dialog.tsx +++ b/web/default/src/features/usage-logs/components/dialogs/details-dialog.tsx @@ -129,6 +129,13 @@ function formatRatio(ratio: number | undefined): string { return ratio.toFixed(4) } +function formatKeyIndexLabel(keyIndex?: number): string | null { + if (keyIndex == null || !Number.isFinite(keyIndex) || keyIndex < 0) { + return null + } + return `#${keyIndex + 1} (index ${keyIndex})` +} + function BillingBreakdown(props: { log: UsageLog other: LogOtherData @@ -477,6 +484,8 @@ export function DetailsDialog(props: DetailsDialogProps) { const useChannel = other?.admin_info?.use_channel const channelChain = useChannel && useChannel.length > 0 ? useChannel.join(' → ') : undefined + const affinity = adminInfo?.channel_affinity + const affinityKeyDisplay = formatKeyIndexLabel(affinity?.key_index) return ( )} + {props.isAdmin && affinity && ( +