Skip to content

feat: support multi-key channel key testing #5646

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
SwillerDawn wants to merge 2 commits into from
Closed

feat: support multi-key channel key testing #5646

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
f0b68fd
feat: support multi-key channel key testing
SwillerDawn Jun 21, 2026
0871a9c
fix: mask multi-key status previews
SwillerDawn Jun 21, 2026
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
23 changes: 20 additions & 3 deletions controller/channel-test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,7 @@ func resolveChannelTestUserID(c *gin.Context) (int, error) {
return rootUser.Id, nil
}

func testChannel(channel *model.Channel, testUserID int, testModel string, endpointType string, isStream bool) testResult {
func testChannel(channel *model.Channel, testUserID int, testModel string, endpointType string, isStream bool, keyIndex *int) testResult {
tik := time.Now()
var unsupportedTestChannelTypes = []int{
constant.ChannelTypeMidjourney,
Expand Down Expand Up @@ -176,6 +176,9 @@ func testChannel(channel *model.Channel, testUserID int, testModel string, endpo
c.Set("base_url", channel.GetBaseURL())
group, _ := model.GetUserGroup(testUserID, false)
c.Set("group", group)
if keyIndex != nil {
middleware.SetChannelTestMultiKeyIndex(c, *keyIndex)
}

newAPIError := middleware.SetupContextForSelectedChannel(c, channel, testModel)
if newAPIError != nil {
Expand Down Expand Up @@ -851,13 +854,27 @@ func TestChannel(c *gin.Context) {
testModel := c.Query("model")
endpointType := c.Query("endpoint_type")
isStream, _ := strconv.ParseBool(c.Query("stream"))
var keyIndex *int
keyIndexText := strings.TrimSpace(c.Query("key_index"))
if keyIndexText != "" {
parsedKeyIndex, err := strconv.Atoi(keyIndexText)
if err != nil {
c.JSON(http.StatusOK, gin.H{
"success": false,
"message": "key_index must be an integer",
"time": 0.0,
})
return
}
keyIndex = &parsedKeyIndex
}
testUserID, err := resolveChannelTestUserID(c)
if err != nil {
common.ApiError(c, err)
return
}
tik := time.Now()
result := testChannel(channel, testUserID, testModel, endpointType, isStream)
result := testChannel(channel, testUserID, testModel, endpointType, isStream, keyIndex)
if result.localErr != nil {
resp := gin.H{
"success": false,
Expand Down Expand Up @@ -928,7 +945,7 @@ func testAllChannels(notify bool) error {
}
isChannelEnabled := channel.Status == common.ChannelStatusEnabled
tik := time.Now()
result := testChannel(channel, testUserID, "", "", shouldUseStreamForAutomaticChannelTest(channel))
result := testChannel(channel, testUserID, "", "", shouldUseStreamForAutomaticChannelTest(channel), nil)
tok := time.Now()
milliseconds := tok.Sub(tik).Milliseconds()

Expand Down
10 changes: 2 additions & 8 deletions controller/channel.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1336,7 +1336,7 @@ type KeyStatus struct {
Status int `json:"status"` // 1: enabled, 2: disabled
DisabledTime int64 `json:"disabled_time,omitempty"`
Reason string `json:"reason,omitempty"`
KeyPreview string `json:"key_preview"` // first 10 chars of key for identification
KeyPreview string `json:"key_preview"` // masked key for identification
}

// ManageMultiKeys handles multi-key management operations
Expand Down Expand Up @@ -1428,18 +1428,12 @@ func ManageMultiKeys(c *gin.Context) {
}
}

// Create key preview (first 10 chars)
keyPreview := key
if len(key) > 10 {
keyPreview = key[:10] + "..."
}

allKeyStatusList = append(allKeyStatusList, KeyStatus{
Index: i,
Status: status,
DisabledTime: disabledTime,
Reason: reason,
KeyPreview: keyPreview,
KeyPreview: model.MaskTokenKey(key),
})
}

Expand Down
39 changes: 36 additions & 3 deletions middleware/distributor.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -440,6 +440,21 @@ func getTaskOriginModelName(c *gin.Context) string {
return ""
}

const ginKeyChannelTestMultiKeyIndex = "channel_test_multi_key_index"

func SetChannelTestMultiKeyIndex(c *gin.Context, index int) {
c.Set(ginKeyChannelTestMultiKeyIndex, index)
}

func getChannelTestMultiKeyIndex(c *gin.Context) (int, bool) {
value, ok := c.Get(ginKeyChannelTestMultiKeyIndex)
if !ok {
return 0, false
}
index, ok := value.(int)
return index, ok
}

func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, modelName string) *types.NewAPIError {
c.Set("original_model", modelName) // for retry
if channel == nil {
Expand All @@ -465,9 +480,27 @@ 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 testKeyIndex, ok := getChannelTestMultiKeyIndex(c); ok {
if !channel.ChannelInfo.IsMultiKey {
return types.NewError(errors.New("key_index is only supported for multi-key channels"), types.ErrorCodeChannelNoAvailableKey, types.ErrOptionWithSkipRetry())
}
selectedKey, selectedIndex, reason, valid := channel.GetEnabledKeyByIndex(testKeyIndex)
if !valid {
if reason == "" {
reason = "specified_key_unavailable"
}
return types.NewError(fmt.Errorf("specified key unavailable: %s", reason), types.ErrorCodeChannelNoAvailableKey, types.ErrOptionWithSkipRetry())
}
key = selectedKey
index = selectedIndex
} else {
key, index, newAPIError = channel.GetNextEnabledKey()
if newAPIError != nil {
return newAPIError
}
}
if channel.ChannelInfo.IsMultiKey {
common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true)
Expand Down
20 changes: 20 additions & 0 deletions model/channel.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -282,6 +282,26 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) {
}
}

func (channel *Channel) GetEnabledKeyByIndex(index int) (string, int, string, bool) {
if index < 0 {
return "", index, "invalid_key_index", false
}

keys := channel.GetKeys()
if index >= len(keys) {
return "", index, "key_index_out_of_range", false
}

statusList := channel.ChannelInfo.MultiKeyStatusList
if statusList != nil {
if status, ok := statusList[index]; ok && status != common.ChannelStatusEnabled {
return "", index, fmt.Sprintf("key_status_%d", status), false
}
}
Comment on lines +295 to +300

@coderabbitai coderabbitai Bot Jun 21, 2026

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

Synchronize forced-index key-status reads with the channel lock.

Line 295 reads MultiKeyStatusList without GetChannelPollingLock(channel.Id). Other status reads/writes (GetNextEnabledKey, ManageMultiKeys) use this lock, so this path can race and panic on concurrent map access.

Proposed fix 
 func (channel *Channel) GetEnabledKeyByIndex(index int) (string, int, string, bool) {
 	if index < 0 {
 		return "", index, "invalid_key_index", false
 	}
+
+	lock := GetChannelPollingLock(channel.Id)
+	lock.Lock()
+	defer lock.Unlock()
 
 	keys := channel.GetKeys()
 	if index >= len(keys) {
 		return "", index, "key_index_out_of_range", false
 	}
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
statusList := channel.ChannelInfo.MultiKeyStatusList
if statusList != nil {
if status, ok := statusList[index]; ok && status != common.ChannelStatusEnabled {
return "", index, fmt.Sprintf("key_status_%d", status), false
}
}
func (channel *Channel) GetEnabledKeyByIndex(index int) (string, int, string, bool) {
if index < 0 {
return "", index, "invalid_key_index", false
}
lock := GetChannelPollingLock(channel.Id)
lock.Lock()
defer lock.Unlock()
keys := channel.GetKeys()
if index >= len(keys) {
return "", index, "key_index_out_of_range", false
}
statusList := channel.ChannelInfo.MultiKeyStatusList
if statusList != nil {
if status, ok := statusList[index]; ok && status != common.ChannelStatusEnabled {
return "", index, fmt.Sprintf("key_status_%d", status), false
}
}
}
🤖 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 295 - 300, The code at line 295 reads
channel.ChannelInfo.MultiKeyStatusList without acquiring the channel polling
lock, while other operations like GetNextEnabledKey and ManageMultiKeys use
GetChannelPollingLock(channel.Id) to protect concurrent access to this map. Wrap
the statusList access and the subsequent if block with
GetChannelPollingLock(channel.Id) to synchronize this read with other status
reads and writes, preventing potential race conditions and panic on concurrent
map access.


return keys[index], index, "", true
}

func (channel *Channel) SaveChannelInfo() error {
return DB.Model(channel).Update("channel_info", channel.ChannelInfo).Error
}
Expand Down
7 changes: 6 additions & 1 deletion web/default/src/features/channels/api.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,7 +181,12 @@ export async function batchSetChannelTag(
*/
export async function testChannel(
id: number,
params?: { model?: string; endpoint_type?: string; stream?: boolean }
params?: {
model?: string
endpoint_type?: string
stream?: boolean
key_index?: number
}
): Promise<ChannelTestResponse> {
const res = await api.get(
`/api/channel/test/${id}`,
Expand Down
Loading