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) || {}) @@ -100,6 +107,25 @@ export function CacheStatsDialog(props: Props) { key: t('Key Fingerprint'), value: (s.key_fp || props.target?.key_fp || '') as string, }) + const channelId = props.target?.channel_id + if ( + channelId != null && + Number.isFinite(channelId) && + channelId >= 0 + ) { + data.push({ key: t('Channel'), value: `#${channelId}` }) + } + const keyIndex = props.target?.key_index + if ( + keyIndex != null && + Number.isFinite(keyIndex) && + keyIndex >= 0 + ) { + data.push({ + key: t('Channel key'), + value: `#${keyIndex + 1} (index ${keyIndex})`, + }) + } if (Number(s.window_seconds || 0) > 0) data.push({ key: t('TTL (seconds)'), value: s.window_seconds as number }) if (total > 0) diff --git a/web/default/src/features/usage-logs/components/columns/common-logs-columns.tsx b/web/default/src/features/usage-logs/components/columns/common-logs-columns.tsx index d56ba10ac3bf..0c1619048a4b 100644 --- a/web/default/src/features/usage-logs/components/columns/common-logs-columns.tsx +++ b/web/default/src/features/usage-logs/components/columns/common-logs-columns.tsx @@ -72,6 +72,13 @@ function formatRatioCompact(ratio: number | undefined): string { : ratio.toFixed(4).replace(/\.?0+$/, '') } +function formatKeyIndexLabel(keyIndex?: number): string | null { + if (keyIndex == null || !Number.isFinite(keyIndex) || keyIndex < 0) { + return null + } + return `#${keyIndex + 1} (index ${keyIndex})` +} + function getGroupRatioText(other: LogOtherData | null): string | null { const userGroupRatio = other?.user_group_ratio if ( @@ -324,6 +331,7 @@ export function useCommonLogsColumns(isAdmin: boolean): ColumnDef[] { ? `${log.channel_name} #${log.channel}` : `#${log.channel}` const channelIdDisplay = `#${log.channel}` + const affinityKeyDisplay = formatKeyIndexLabel(affinity?.key_index) const channelName = sensitiveVisible ? log.channel_name : '••••' return ( @@ -357,6 +365,8 @@ export function useCommonLogsColumns(isAdmin: boolean): ColumnDef[] { '', key_hint: affinity.key_hint || '', key_fp: affinity.key_fp || '', + channel_id: affinity.channel_id, + key_index: affinity.key_index, }) setAffinityDialogOpen(true) }} @@ -395,6 +405,11 @@ export function useCommonLogsColumns(isAdmin: boolean): ColumnDef[] { '-' : '••••'}

+ {affinityKeyDisplay && ( +

+ {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 && ( + + {affinity.rule_name && ( + + )} + {affinity.channel_id != null && ( + + )} + {affinityKeyDisplay && ( + + )} + {(affinity.using_group || affinity.selected_group) && ( + + )} + + )} + {props.log.token_name && ( )} diff --git a/web/default/src/features/usage-logs/index.tsx b/web/default/src/features/usage-logs/index.tsx index b6d28b7c8a61..3d3d0e1a031d 100644 --- a/web/default/src/features/usage-logs/index.tsx +++ b/web/default/src/features/usage-logs/index.tsx @@ -151,6 +151,8 @@ function UsageLogsContent() { '', key_hint: affinityTarget.key_hint || '', key_fp: affinityTarget.key_fp || '', + channel_id: affinityTarget.channel_id, + key_index: affinityTarget.key_index, } : null } diff --git a/web/default/src/features/usage-logs/types.ts b/web/default/src/features/usage-logs/types.ts index 8db88c499107..6846ff4e088a 100644 --- a/web/default/src/features/usage-logs/types.ts +++ b/web/default/src/features/usage-logs/types.ts @@ -83,6 +83,8 @@ export type LogFilters = CommonLogFilters | DrawingLogFilters | TaskLogFilters */ export interface ChannelAffinityInfo { rule_name?: string + channel_id?: number + key_index?: number selected_group?: string key_source?: string key_path?: string