Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
35 changes: 33 additions & 2 deletions model/channel.go
Original file line number Diff line number Diff line change
Expand Up @@ -643,13 +643,25 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason
if len(keys) == 0 {
channel.Status = status
} else {
var keyIndex int
keyIndex := -1
for i, key := range keys {
if key == usingKey {
keyIndex = i
break
}
}
if keyIndex < 0 {
if usingKey != "" {
common.SysLog(fmt.Sprintf("failed to update multi-key status: channel_id=%d, using key not found", channel.Id))
return
}
channel.Status = status
info := channel.GetOtherInfo()
info["status_reason"] = reason
info["status_time"] = common.GetTimestamp()
channel.SetOtherInfo(info)
return
Comment on lines +653 to +663
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Don't re-enable the aggregate channel without reconciling per-key state.

When usingKey == "", this branch flips channel.Status and returns without touching MultiKeyStatusList. If a fully auto-disabled multi-key channel is re-enabled through this path, GetNextEnabledKey() will still see every key as disabled while the route cache now treats the channel as enabled. Please clear or recompute the per-key disabled maps before returning from the whole-channel enable path.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model/channel.go` around lines 653 - 663, When handling the whole-channel
enable path (when usingKey == ""), do not just flip channel.Status and return;
you must clear or recompute the per-key disabled state so GetNextEnabledKey()
won't still see all keys as disabled. In the branch that sets channel.Status,
update channel.MultiKeyStatusList (or the structure holding per-key disabled
maps) to reflect the new enabled state (e.g., reset/clear the list or set each
key's status to enabled) and persist any changes (same place
GetOtherInfo()/SetOtherInfo() are used for channel metadata) before returning so
route cache and per-key lookup remain consistent.

}
if channel.ChannelInfo.MultiKeyStatusList == nil {
channel.ChannelInfo.MultiKeyStatusList = make(map[int]int)
}
Expand All @@ -666,16 +678,31 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason
channel.ChannelInfo.MultiKeyDisabledReason[keyIndex] = reason
channel.ChannelInfo.MultiKeyDisabledTime[keyIndex] = common.GetTimestamp()
}
if len(channel.ChannelInfo.MultiKeyStatusList) >= channel.ChannelInfo.MultiKeySize {
if !hasEnabledMultiKey(keys, channel.ChannelInfo.MultiKeyStatusList) {
channel.Status = common.ChannelStatusAutoDisabled
info := channel.GetOtherInfo()
info["status_reason"] = "All keys are disabled"
info["status_time"] = common.GetTimestamp()
channel.SetOtherInfo(info)
} else if status == common.ChannelStatusEnabled {
channel.Status = common.ChannelStatusEnabled
}
}
}

func hasEnabledMultiKey(keys []string, statusList map[int]int) bool {
for i := range keys {
if statusList == nil {
return true
}
status, ok := statusList[i]
if !ok || status == common.ChannelStatusEnabled {
return true
}
}
return false
}

func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool {
if common.MemoryCacheEnabled {
channelStatusLock.Lock()
Expand All @@ -687,11 +714,15 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri
}
if channelCache.ChannelInfo.IsMultiKey {
// Use per-channel lock to prevent concurrent map read/write with GetNextEnabledKey
beforeStatus := channelCache.Status
pollingLock := GetChannelPollingLock(channelId)
pollingLock.Lock()
// 如果是多Key模式,更新缓存中的状态
handlerMultiKeyUpdate(channelCache, usingKey, status, reason)
pollingLock.Unlock()
if beforeStatus != channelCache.Status {
CacheUpdateChannelStatus(channelId, channelCache.Status)
}
Comment on lines +717 to +725
Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Persist per-key recoveries even when the aggregate status stays enabled.

This cache-path change updates channelCache.ChannelInfo before comparing beforeStatus, but the DB path still short-circuits at Line 750 when channel.Status == status. Re-enabling one key on an already-enabled multi-key channel will clear the in-memory MultiKeyStatusList entry and then return before the DB copy is updated, so the key comes back as disabled after a reload. The no-op check for multi-key channels needs to run after handlerMultiKeyUpdate() and consider ChannelInfo changes, not just channel.Status.

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@model/channel.go` around lines 717 - 725, The DB-path early return that
checks channel.Status == status must be changed so the no-op check runs after
handlerMultiKeyUpdate(); instead of returning based solely on channel.Status,
reload or compare the persisted ChannelInfo/MultiKeyStatusList against the
updated channelCache.ChannelInfo (or compare MultiKeyStatusList entries) and
only short-circuit if the per-key status lists are identical; this ensures
handlerMultiKeyUpdate(channelCache, usingKey, status, reason) updates are
persisted via CacheUpdateChannelStatus/DB write even when the aggregate
channel.Status remains enabled.

//CacheUpdateChannel(channelCache)
//return true
} else {
Expand Down
71 changes: 71 additions & 0 deletions model/channel_cache_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
package model

import (
"testing"

"github.com/QuantumNous/new-api/common"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)

func withChannelCacheTestState(t *testing.T) {
t.Helper()
originalMemoryCacheEnabled := common.MemoryCacheEnabled
originalGroup2Model2Channels := group2model2channels
originalChannelsIDM := channelsIDM

common.MemoryCacheEnabled = true
group2model2channels = map[string]map[string][]int{}
channelsIDM = map[int]*Channel{}

t.Cleanup(func() {
common.MemoryCacheEnabled = originalMemoryCacheEnabled
group2model2channels = originalGroup2Model2Channels
channelsIDM = originalChannelsIDM
})
}

func TestUpdateChannelStatusEvictsMultiKeyChannelFromRouteCache(t *testing.T) {
withChannelCacheTestState(t)

channelsIDM = map[int]*Channel{
1: {
Id: 1,
Status: common.ChannelStatusEnabled,
Key: "k1",
Group: "default",
Models: "gpt-test",
ChannelInfo: ChannelInfo{
IsMultiKey: true,
MultiKeySize: 1,
MultiKeyStatusList: map[int]int{},
},
},
2: {
Id: 2,
Status: common.ChannelStatusEnabled,
Group: "default",
Models: "gpt-test",
},
}
group2model2channels = map[string]map[string][]int{
"default": {"gpt-test": {1, 2}},
}

cache := channelsIDM[1]
pollingLock := GetChannelPollingLock(cache.Id)
pollingLock.Lock()
beforeStatus := cache.Status
handlerMultiKeyUpdate(cache, "k1", common.ChannelStatusAutoDisabled, "test reason")
pollingLock.Unlock()
require.NotEqual(t, beforeStatus, cache.Status, "channel should auto-disable when all keys are disabled")
CacheUpdateChannelStatus(cache.Id, cache.Status)

assert.NotContains(t, group2model2channels["default"]["gpt-test"], 1,
"auto-disabled multi-key channel should be removed from route cache")

channel, err := GetRandomSatisfiedChannel("default", "gpt-test", 0)
require.NoError(t, err)
require.NotNil(t, channel)
assert.Equal(t, 2, channel.Id)
}