diff --git a/controller/channel.go b/controller/channel.go index 19ddca8e6a07..f98e88c86ab3 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -954,8 +954,9 @@ func DeleteChannelBatch(c *gin.Context) { type PatchChannel struct { model.Channel - MultiKeyMode *string `json:"multi_key_mode"` - KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加 + MultiKeyMode *string `json:"multi_key_mode"` + KeyMode *string `json:"key_mode"` // 多key模式下密钥覆盖或者追加 + KeyStorageMode *string `json:"key_storage_mode"` // single | multi，编辑时转换密钥存储形态 } type ChannelStatusRequest struct { @@ -988,8 +989,37 @@ func UpdateChannel(c *gin.Context) { return } clearChannelReadOnlyFields(&channel, requestData) + convertingStorage := channel.KeyStorageMode != nil && strings.TrimSpace(*channel.KeyStorageMode) != "" + // Every update copies the persisted ChannelInfo forward, so an ordinary edit + // that reads multi-key metadata before a multi→single conversion commits + // would restore IsMultiKey and its status maps over the converted single key. + // Hold the per-channel lock from the source read through persistence for all + // requests, not just conversions. + lock := model.GetChannelPollingLock(channel.Id) + lock.Lock() + defer lock.Unlock() + + var originChannel *model.Channel + if convertingStorage { + // Conversion checks must use a fresh source row while the per-channel + // lock is held. In particular, a partial request may omit type/settings. + originChannel, err = model.GetChannelById(channel.Id, true) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + } - if channel.Type == constant.ChannelTypeTaskPlugin && + channelTypeForChecks := channel.Type + if convertingStorage { + if _, provided := requestData["type"]; !provided { + channelTypeForChecks = originChannel.Type + } + } + if channelTypeForChecks == constant.ChannelTypeTaskPlugin && !authz.Can(c.GetInt("id"), c.GetInt("role"), authz.TaskPluginBind) { c.JSON(http.StatusOK, gin.H{ "success": false, @@ -1007,13 +1037,15 @@ func UpdateChannel(c *gin.Context) { return } // Preserve existing ChannelInfo to ensure multi-key channels keep correct state even if the client does not send ChannelInfo in the request. - originChannel, err := model.GetChannelById(channel.Id, true) - if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - return + if originChannel == nil { + originChannel, err = model.GetChannelById(channel.Id, true) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } } originProxy := originChannel.GetSetting().Proxy proxyChanged := false @@ -1033,12 +1065,36 @@ func UpdateChannel(c *gin.Context) { } // If the request explicitly specifies a new MultiKeyMode, apply it on top of the original info. - if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" { + if convertingStorage { + keyConfig := *originChannel + if _, provided := requestData["type"]; provided { + keyConfig.Type = channel.Type + } + // The conversion guard must judge the settings that will actually be + // persisted. UpdateWithConvertedKeyStorage writes only key/channel_info + // explicitly and its struct update skips an empty settings column, so a + // request that omits or blanks settings keeps the stored value — even + // when it also changes the channel type. Assuming a cleared credential + // mode here would let a Vertex API Key channel pass the multi-key gate + // and then keep vertex_key_type=api_key after commit. + if strings.TrimSpace(channel.OtherSettings) != "" { + keyConfig.OtherSettings = channel.OtherSettings + } + if err := applyKeyStorageMode(&channel, originChannel, &keyConfig); err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": err.Error(), + }) + return + } + } + + if !convertingStorage && channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" { channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode) } // 处理多key模式下的密钥追加/覆盖逻辑 - if channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey { + if !convertingStorage && channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey { switch *channel.KeyMode { case "append": // 追加模式：将新密钥添加到现有密钥列表 @@ -1117,7 +1173,14 @@ func UpdateChannel(c *gin.Context) { // 覆盖模式：直接使用新密钥（默认行为，不需要特殊处理） } } - err = channel.Update() + // 密钥形态转换必须把 key 与 channel_info 放进同一事务： + // GORM Updates(struct) 会跳过零值 ChannelInfo，若分两次写入， + // 失败时会留下「标记为多密钥、实际只有一把密钥」或反过来的半成品。 + if convertingStorage { + err = channel.UpdateWithConvertedKeyStorage() + } else { + err = channel.Update() + } if err != nil { common.ApiError(c, err) return @@ -1538,6 +1601,13 @@ func ManageMultiKeys(c *gin.Context) { return } + // Acquire the channel lock before reading ChannelInfo. A conversion can + // replace the key storage shape while this request is waiting; reading + // first would let a stale multi-key snapshot overwrite that conversion. + lock := model.GetChannelPollingLock(request.ChannelId) + lock.Lock() + defer lock.Unlock() + channel, err := model.GetChannelById(request.ChannelId, true) if err != nil { c.JSON(http.StatusOK, gin.H{ @@ -1570,10 +1640,6 @@ func ManageMultiKeys(c *gin.Context) { }) } - lock := model.GetChannelPollingLock(channel.Id) - lock.Lock() - defer lock.Unlock() - switch request.Action { case "get_key_status": keys := channel.GetKeys() diff --git a/controller/channel_authz.go b/controller/channel_authz.go index f85ffef92769..29ef53a02db3 100644 --- a/controller/channel_authz.go +++ b/controller/channel_authz.go @@ -33,6 +33,9 @@ func channelHasSensitiveChanges(channel *PatchChannel, origin *model.Channel, re if _, ok := requestData["key_mode"]; ok && channel.KeyMode != nil { return true } + if _, ok := requestData["key_storage_mode"]; ok && channel.KeyStorageMode != nil { + return true + } // Fail closed: any field present in the request that is neither a known // sensitive field (gated above) nor an explicitly classified non-sensitive // field must be treated as sensitive. This keeps a newly added channel field @@ -71,6 +74,7 @@ var channelSensitiveFields = map[string]struct{}{ "other": {}, "settings": {}, "key_mode": {}, + "key_storage_mode": {}, } // channelOperationalFields lists fields managed by operation endpoints instead diff --git a/controller/channel_key_storage.go b/controller/channel_key_storage.go new file mode 100644 index 000000000000..1484c24d53ef --- /dev/null +++ b/controller/channel_key_storage.go @@ -0,0 +1,106 @@ +package controller + +import ( + "fmt" + "strings" + + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/relaykit/dto" +) + +const ( + keyStorageModeSingle = "single" + keyStorageModeMulti = "multi" +) + +func channelSupportsMultiKey(channel *model.Channel) bool { + if channel.Type == constant.ChannelTypeCodex { + return false + } + if channel.Type == constant.ChannelTypeVertexAi && channel.GetOtherSettings().VertexKeyType == dto.VertexKeyTypeAPIKey { + return false + } + return true +} + +func parseUpdateKeyList(channel *model.Channel, key string) ([]string, error) { + if channel.Type == constant.ChannelTypeVertexAi && channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { + trimmed := strings.TrimSpace(key) + if strings.HasPrefix(trimmed, "[") { + return getVertexArrayKeys(key) + } + if trimmed == "" { + return nil, fmt.Errorf("密钥不能为空") + } + return []string{trimmed}, nil + } + keys := make([]string, 0) + for _, line := range strings.Split(key, "\n") { + line = strings.TrimSpace(line) + if line == "" { + continue + } + keys = append(keys, line) + } + return keys, nil +} + +func applyKeyStorageMode(channel *PatchChannel, origin *model.Channel, keyConfig *model.Channel) error { + if channel.KeyStorageMode == nil { + return nil + } + if keyConfig == nil { + keyConfig = origin + } + mode := strings.TrimSpace(*channel.KeyStorageMode) + if mode == "" { + return nil + } + if mode != keyStorageModeSingle && mode != keyStorageModeMulti { + return fmt.Errorf("不支持的密钥存储模式") + } + if strings.TrimSpace(channel.Key) == "" { + return fmt.Errorf("转换密钥模式必须提供新的密钥") + } + keys, err := parseUpdateKeyList(keyConfig, channel.Key) + if err != nil { + return err + } + switch mode { + case keyStorageModeSingle: + if !origin.ChannelInfo.IsMultiKey { + return fmt.Errorf("渠道已经是单密钥模式") + } + if len(keys) != 1 { + return fmt.Errorf("转换为单密钥时必须只提供一把新密钥") + } + channel.Key = keys[0] + channel.ChannelInfo = model.ChannelInfo{} + return nil + default: + if origin.ChannelInfo.IsMultiKey { + return fmt.Errorf("渠道已经是多密钥模式") + } + if !channelSupportsMultiKey(keyConfig) { + return fmt.Errorf("当前渠道类型不支持多密钥模式") + } + if len(keys) < 2 { + return fmt.Errorf("转换为多密钥时至少需要两把密钥") + } + rotation := constant.MultiKeyModeRandom + if channel.MultiKeyMode != nil { + requested := constant.MultiKeyMode(strings.TrimSpace(*channel.MultiKeyMode)) + if requested == constant.MultiKeyModeRandom || requested == constant.MultiKeyModePolling { + rotation = requested + } + } + channel.Key = strings.Join(keys, "\n") + channel.ChannelInfo = model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: len(keys), + MultiKeyMode: rotation, + } + return nil + } +} diff --git a/controller/channel_key_storage_test.go b/controller/channel_key_storage_test.go new file mode 100644 index 000000000000..b8f3a5388b62 --- /dev/null +++ b/controller/channel_key_storage_test.go @@ -0,0 +1,361 @@ +package controller + +import ( + "fmt" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func putUpdateChannel(t *testing.T, userID, role int, body string) *httptest.ResponseRecorder { + t.Helper() + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + context, _ := gin.CreateTestContext(recorder) + context.Set("id", userID) + context.Set("role", role) + context.Request = httptest.NewRequest(http.MethodPut, "/api/channel", strings.NewReader(body)) + context.Request.Header.Set("Content-Type", "application/json") + UpdateChannel(context) + return recorder +} + +func insertKeyStorageChannel(t *testing.T, channel model.Channel) *model.Channel { + t.Helper() + channel.Status = common.ChannelStatusEnabled + channel.Name = "key-storage" + channel.Models = "gpt-4o" + channel.Group = "default" + require.NoError(t, channel.Insert()) + // Insert() does not persist ChannelInfo through the multi-key recalculation path, + // so write the intended starting state explicitly. + require.NoError(t, channel.SaveKeyAndChannelInfo()) + return &channel +} + +func keyStorageUpdateBody(channelID int, channelType int, extra string) string { + return fmt.Sprintf( + `{"id":%d,"type":%d,"name":"key-storage","models":"gpt-4o","group":"default",%s}`, + channelID, channelType, extra, + ) +} + +func TestUpdateChannelConvertsSingleKeyChannelToMultiKey(t *testing.T) { + setupTaskPluginBindChannelTest(t) + channel := insertKeyStorageChannel(t, model.Channel{Type: 1, Key: "sk-old"}) + + recorder := putUpdateChannel(t, 1, common.RoleRootUser, keyStorageUpdateBody( + channel.Id, 1, + `"key_storage_mode":"multi","multi_key_mode":"polling","key":"sk-a\nsk-b\nsk-c"`, + )) + require.Contains(t, recorder.Body.String(), `"success":true`) + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.True(t, stored.ChannelInfo.IsMultiKey) + assert.Equal(t, 3, stored.ChannelInfo.MultiKeySize) + assert.Equal(t, constant.MultiKeyModePolling, stored.ChannelInfo.MultiKeyMode) + assert.Equal(t, []string{"sk-a", "sk-b", "sk-c"}, stored.GetKeys()) +} + +func TestUpdateChannelConvertsMultiKeyChannelToSingleKeyWithNewKey(t *testing.T) { + setupTaskPluginBindChannelTest(t) + channel := insertKeyStorageChannel(t, model.Channel{ + Type: 1, + Key: "sk-a\nsk-b", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyMode: constant.MultiKeyModePolling, + MultiKeyStatusList: map[int]int{1: common.ChannelStatusAutoDisabled}, + MultiKeyDisabledReason: map[int]string{1: "quota"}, + MultiKeyPollingIndex: 1, + }, + }) + + recorder := putUpdateChannel(t, 1, common.RoleRootUser, keyStorageUpdateBody( + channel.Id, 1, + `"key_storage_mode":"single","key":"sk-fresh"`, + )) + require.Contains(t, recorder.Body.String(), `"success":true`) + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.False(t, stored.ChannelInfo.IsMultiKey) + assert.Equal(t, "sk-fresh", stored.Key) + assert.Equal(t, 0, stored.ChannelInfo.MultiKeySize) + assert.Empty(t, stored.ChannelInfo.MultiKeyStatusList) + assert.Empty(t, stored.ChannelInfo.MultiKeyDisabledReason) + assert.Zero(t, stored.ChannelInfo.MultiKeyPollingIndex) + + // The single-key runtime path must hand out exactly the new key. + usedKey, index, keyErr := stored.GetNextEnabledKey() + require.Nil(t, keyErr) + assert.Equal(t, "sk-fresh", usedKey) + assert.Zero(t, index) +} + +func TestUpdateChannelRejectsInvalidKeyStorageConversions(t *testing.T) { + setupTaskPluginBindChannelTest(t) + singleKeyChannel := insertKeyStorageChannel(t, model.Channel{Type: 1, Key: "sk-old"}) + multiKeyChannel := insertKeyStorageChannel(t, model.Channel{ + Type: 1, + Key: "sk-a\nsk-b", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyMode: constant.MultiKeyModeRandom, + }, + }) + // Vertex AI in API Key mode is the type gate that reaches applyKeyStorageMode; + // Codex is rejected earlier by its credential format validation. + vertexAPIKeyChannel := insertKeyStorageChannel(t, model.Channel{ + Type: constant.ChannelTypeVertexAi, + Key: "vertex-key", + Other: `{"default":"us-central1"}`, + OtherSettings: `{"vertex_key_type":"api_key"}`, + }) + + cases := []struct { + name string + channel *model.Channel + channelType int + extra string + expectedMessage string + }{ + { + name: "converting to single without a new key is rejected", + channel: multiKeyChannel, + channelType: 1, + extra: `"key_storage_mode":"single"`, + expectedMessage: "转换密钥模式必须提供新的密钥", + }, + { + name: "converting to single with multiple keys is rejected", + channel: multiKeyChannel, + channelType: 1, + extra: `"key_storage_mode":"single","key":"sk-a\nsk-b"`, + expectedMessage: "转换为单密钥时必须只提供一把新密钥", + }, + { + name: "converting to multi with a single key is rejected", + channel: singleKeyChannel, + channelType: 1, + extra: `"key_storage_mode":"multi","key":"sk-only"`, + expectedMessage: "转换为多密钥时至少需要两把密钥", + }, + { + name: "converting an already multi-key channel to multi is rejected", + channel: multiKeyChannel, + channelType: 1, + extra: `"key_storage_mode":"multi","key":"sk-a\nsk-b"`, + expectedMessage: "渠道已经是多密钥模式", + }, + { + name: "converting an already single-key channel to single is rejected", + channel: singleKeyChannel, + channelType: 1, + extra: `"key_storage_mode":"single","key":"sk-new"`, + expectedMessage: "渠道已经是单密钥模式", + }, + { + name: "unsupported storage mode is rejected", + channel: singleKeyChannel, + channelType: 1, + extra: `"key_storage_mode":"rotate","key":"sk-new"`, + expectedMessage: "不支持的密钥存储模式", + }, + { + name: "vertex api key channels cannot become multi-key", + channel: vertexAPIKeyChannel, + channelType: constant.ChannelTypeVertexAi, + extra: `"key_storage_mode":"multi","key":"vertex-a\nvertex-b","other":"{\"default\":\"us-central1\"}","settings":"{\"vertex_key_type\":\"api_key\"}"`, + expectedMessage: "当前渠道类型不支持多密钥模式", + }, + } + + for _, testCase := range cases { + t.Run(testCase.name, func(t *testing.T) { + before, err := model.GetChannelById(testCase.channel.Id, true) + require.NoError(t, err) + + recorder := putUpdateChannel(t, 1, common.RoleRootUser, keyStorageUpdateBody( + testCase.channel.Id, testCase.channelType, testCase.extra, + )) + body := recorder.Body.String() + assert.Contains(t, body, `"success":false`) + assert.Contains(t, body, testCase.expectedMessage) + + after, err := model.GetChannelById(testCase.channel.Id, true) + require.NoError(t, err) + assert.Equal(t, before.Key, after.Key) + assert.Equal(t, before.ChannelInfo, after.ChannelInfo) + }) + } +} + +func TestUpdateChannelKeyStorageConversionRequiresSensitiveWrite(t *testing.T) { + setupTaskPluginBindChannelTest(t) + channel := insertKeyStorageChannel(t, model.Channel{Type: 1, Key: "sk-old"}) + + recorder := putUpdateChannel(t, 2, common.RoleAdminUser, keyStorageUpdateBody( + channel.Id, 1, + `"key_storage_mode":"multi","key":"sk-a\nsk-b"`, + )) + assert.Contains(t, recorder.Body.String(), `"success":false`) + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.False(t, stored.ChannelInfo.IsMultiKey) + assert.Equal(t, "sk-old", stored.Key) +} + +func TestUpdateChannelConvertsVertexJsonChannelToMultiKeyUsingJsonArray(t *testing.T) { + setupTaskPluginBindChannelTest(t) + channel := insertKeyStorageChannel(t, model.Channel{ + Type: constant.ChannelTypeVertexAi, + Key: `{"type":"service_account","project_id":"a"}`, + Other: `{"default":"us-central1"}`, + OtherSettings: `{"vertex_key_type":"json"}`, + }) + + // A pretty-printed JSON array must be split per element; newline splitting + // would shred one service account into several bogus keys. + newKeys := strings.Join([]string{ + "[", + ` {"type": "service_account", "project_id": "a"},`, + ` {"type": "service_account", "project_id": "b"}`, + "]", + }, "\n") + encodedKeys, err := common.Marshal(newKeys) + require.NoError(t, err) + + recorder := putUpdateChannel(t, 1, common.RoleRootUser, keyStorageUpdateBody( + channel.Id, constant.ChannelTypeVertexAi, + fmt.Sprintf( + `"key_storage_mode":"multi","key":%s,"other":"{\"default\":\"us-central1\"}","settings":"{\"vertex_key_type\":\"json\"}"`, + encodedKeys, + ), + )) + require.Contains(t, recorder.Body.String(), `"success":true`) + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.True(t, stored.ChannelInfo.IsMultiKey) + assert.Equal(t, 2, stored.ChannelInfo.MultiKeySize) + assert.Len(t, stored.GetKeys(), 2) +} + +func TestUpdateChannelTypeChangeWithoutSettingsKeepsPersistedCredentialMode(t *testing.T) { + setupTaskPluginBindChannelTest(t) + // A channel can carry a stale vertex_key_type from an earlier configuration. + // Changing type without sending settings does not clear the settings column, + // so the conversion guard must read the value that stays persisted. + channel := insertKeyStorageChannel(t, model.Channel{ + Type: 1, + Key: "sk-old", + Other: `{"default":"us-central1"}`, + OtherSettings: `{"vertex_key_type":"api_key"}`, + }) + + // The key is a JSON array so it survives the Vertex JSON parse and reaches the + // multi-key gate. Assuming a cleared credential mode would let it through and + // commit a multi-key Vertex channel that still stores vertex_key_type=api_key. + newKeys := `[{"type":"service_account","project_id":"a"},{"type":"service_account","project_id":"b"}]` + encodedKeys, err := common.Marshal(newKeys) + require.NoError(t, err) + + recorder := putUpdateChannel(t, 1, common.RoleRootUser, keyStorageUpdateBody( + channel.Id, constant.ChannelTypeVertexAi, + fmt.Sprintf( + `"key_storage_mode":"multi","key":%s,"other":"{\"default\":\"us-central1\"}"`, + encodedKeys, + ), + )) + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), "当前渠道类型不支持多密钥模式") + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.False(t, stored.ChannelInfo.IsMultiKey) + assert.Equal(t, "sk-old", stored.Key) + assert.Equal(t, `{"vertex_key_type":"api_key"}`, stored.OtherSettings) +} + +func TestUpdateChannelConcurrentOrdinaryUpdateDoesNotRestoreMultiKeyState(t *testing.T) { + setupTaskPluginBindChannelTest(t) + channel := insertKeyStorageChannel(t, model.Channel{ + Type: 1, + Key: "sk-a\nsk-b", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyMode: constant.MultiKeyModePolling, + MultiKeyStatusList: map[int]int{1: common.ChannelStatusAutoDisabled}, + }, + }) + + // An ordinary rename racing a multi→single conversion must not write back the + // pre-conversion ChannelInfo it read, which would leave the channel flagged + // multi-key while holding a single key. + var waitGroup sync.WaitGroup + waitGroup.Add(2) + go func() { + defer waitGroup.Done() + putUpdateChannel(t, 1, common.RoleRootUser, keyStorageUpdateBody( + channel.Id, 1, + `"key_storage_mode":"single","key":"sk-fresh"`, + )) + }() + go func() { + defer waitGroup.Done() + putUpdateChannel(t, 1, common.RoleRootUser, fmt.Sprintf( + `{"id":%d,"type":1,"name":"renamed","models":"gpt-4o","group":"default"}`, + channel.Id, + )) + }() + waitGroup.Wait() + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + // Whichever request commits last, the stored key and storage flag must agree. + if stored.ChannelInfo.IsMultiKey { + assert.Equal(t, "sk-a\nsk-b", stored.Key) + } else { + assert.Equal(t, "sk-fresh", stored.Key) + assert.Empty(t, stored.ChannelInfo.MultiKeyStatusList) + } +} + +func TestUpdateChannelDoesNotTrustOmittedVertexCredentialSettings(t *testing.T) { + setupTaskPluginBindChannelTest(t) + channel := insertKeyStorageChannel(t, model.Channel{ + Type: constant.ChannelTypeVertexAi, + Key: "vertex-key", + Other: `{"default":"us-central1"}`, + OtherSettings: `{"vertex_key_type":"api_key"}`, + }) + + // A partial update may omit type/settings. The conversion guard must use + // the persisted channel configuration instead of treating it as Vertex JSON. + recorder := putUpdateChannel(t, 1, common.RoleRootUser, fmt.Sprintf( + `{"id":%d,"key_storage_mode":"multi","key":"vertex-a\nvertex-b"}`, + channel.Id, + )) + assert.Contains(t, recorder.Body.String(), `"success":false`) + assert.Contains(t, recorder.Body.String(), "当前渠道类型不支持多密钥模式") + + stored, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, "vertex-key", stored.Key) + assert.False(t, stored.ChannelInfo.IsMultiKey) +} diff --git a/model/channel.go b/model/channel.go index 268f190e703c..d39294fb9bd0 100644 --- a/model/channel.go +++ b/model/channel.go @@ -291,6 +291,85 @@ func (channel *Channel) SaveChannelInfo() error { return DB.Model(channel).Update("channel_info", channel.ChannelInfo).Error } +// SaveKeyAndChannelInfo persists credentials and key-storage metadata together. +// GORM Updates(struct) skips zero-value ChannelInfo, so multi→single conversion +// must write these columns through a map instead of relying on Channel.Update. +func (channel *Channel) SaveKeyAndChannelInfo() error { + if channel.Id == 0 { + return errors.New("channel ID is 0") + } + return DB.Model(&Channel{}).Where("id = ?", channel.Id).Updates(channel.keyStorageUpdateMap()).Error +} + +func (channel *Channel) keyStorageUpdateMap() map[string]any { + return map[string]any{ + "key": channel.Key, + "channel_info": channel.ChannelInfo, + } +} + +func (channel *Channel) prepareMultiKeySize() { + if !channel.ChannelInfo.IsMultiKey { + return + } + var keyStr string + if channel.Key != "" { + keyStr = channel.Key + } else if existing, err := GetChannelById(channel.Id, true); err == nil { + keyStr = existing.Key + } + keys := []string{} + if keyStr != "" { + trimmed := strings.TrimSpace(keyStr) + if strings.HasPrefix(trimmed, "[") { + var arr []json.RawMessage + if err := common.Unmarshal([]byte(trimmed), &arr); err == nil { + keys = make([]string, len(arr)) + for i, v := range arr { + keys[i] = string(v) + } + } + } + if len(keys) == 0 { // fallback to newline split + keys = strings.Split(strings.Trim(keyStr, "\n"), "\n") + } + } + channel.ChannelInfo.MultiKeySize = len(keys) + if channel.ChannelInfo.MultiKeyStatusList != nil { + for idx := range channel.ChannelInfo.MultiKeyStatusList { + if idx >= channel.ChannelInfo.MultiKeySize { + delete(channel.ChannelInfo.MultiKeyStatusList, idx) + } + } + } +} + +// UpdateWithConvertedKeyStorage writes the converted key and ChannelInfo in the +// same transaction as the rest of the channel update. A later failure must not +// leave the channel with a new key and stale multi-key metadata, or vice versa. +func (channel *Channel) UpdateWithConvertedKeyStorage() error { + if channel.Id == 0 { + return errors.New("channel ID is 0") + } + channel.prepareMultiKeySize() + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(channel.keyStorageUpdateMap()).Error; err != nil { + return err + } + // The credential and storage metadata were already written explicitly + // above. Omit them from the ordinary struct update so a zero-value + // ChannelInfo (multi→single) cannot be skipped or reintroduced by a + // second write. + if err := tx.Model(channel).Omit("key", "channel_info").Updates(channel).Error; err != nil { + return err + } + if err := tx.Model(channel).First(channel, "id = ?", channel.Id).Error; err != nil { + return err + } + return channel.UpdateAbilities(tx) + }) +} + func (channel *Channel) GetModels() []string { if channel.Models == "" { return []string{} @@ -554,44 +633,7 @@ func (channel *Channel) Insert() error { } func (channel *Channel) Update() error { - // If this is a multi-key channel, recalculate MultiKeySize based on the current key list to avoid inconsistency after editing keys - if channel.ChannelInfo.IsMultiKey { - var keyStr string - if channel.Key != "" { - keyStr = channel.Key - } else { - // If key is not provided, read the existing key from the database - if existing, err := GetChannelById(channel.Id, true); err == nil { - keyStr = existing.Key - } - } - // Parse the key list (supports newline separation or JSON array) - keys := []string{} - if keyStr != "" { - trimmed := strings.TrimSpace(keyStr) - if strings.HasPrefix(trimmed, "[") { - var arr []json.RawMessage - if err := common.Unmarshal([]byte(trimmed), &arr); err == nil { - keys = make([]string, len(arr)) - for i, v := range arr { - keys[i] = string(v) - } - } - } - if len(keys) == 0 { // fallback to newline split - keys = strings.Split(strings.Trim(keyStr, "\n"), "\n") - } - } - channel.ChannelInfo.MultiKeySize = len(keys) - // Clean up status data that exceeds the new key count to prevent index out of range - if channel.ChannelInfo.MultiKeyStatusList != nil { - for idx := range channel.ChannelInfo.MultiKeyStatusList { - if idx >= channel.ChannelInfo.MultiKeySize { - delete(channel.ChannelInfo.MultiKeyStatusList, idx) - } - } - } - } + channel.prepareMultiKeySize() var err error err = DB.Model(channel).Updates(channel).Error if err != nil { diff --git a/web/src/features/channels/api.ts b/web/src/features/channels/api.ts index fa13a186c5a0..6095020ea6d3 100644 --- a/web/src/features/channels/api.ts +++ b/web/src/features/channels/api.ts @@ -19,6 +19,7 @@ For commercial licensing, please contact support@quantumnous.com import { getGroups as getUserGroups } from '@/features/users/api' import { api, type ApiRequestConfig } from '@/lib/api' +import type { ChannelUpdatePayload } from './lib/channel-form' import type { AddChannelRequest, BatchDeleteParams, @@ -139,7 +140,7 @@ export async function createChannel( */ export async function updateChannel( id: number, - data: Partial + data: ChannelUpdatePayload ): Promise<{ success: boolean; message?: string; data?: Channel }> { const res = await api.put( '/api/channel/', diff --git a/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx b/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx index 93f6a0a1763e..085c9f45a98d 100644 --- a/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx +++ b/web/src/features/channels/components/drawers/channel-mutate-drawer.tsx @@ -158,7 +158,9 @@ import { CHANNEL_TYPE_ADVANCED_CUSTOM, channelFormSchema, channelsQueryKeys, + countUpdateKeys, getAdvancedCustomStats, + resolveStorageModeConversion, transformChannelToFormDefaults, type ChannelFormValues, deduplicateKeys, @@ -278,6 +280,7 @@ const SENSITIVE_FORM_FIELDS = [ 'openai_organization', 'other', 'key_mode', + 'key_storage_mode', 'param_override', 'header_override', 'settings', @@ -729,6 +732,7 @@ export function ChannelMutateDrawer({ const multiKeyMode = form.watch('multi_key_mode') const multiKeyType = form.watch('multi_key_type') const keyMode = form.watch('key_mode') + const keyStorageMode = form.watch('key_storage_mode') const currentGroups = form.watch('group') const currentType = form.watch('type') const currentStatus = form.watch('status') @@ -866,6 +870,20 @@ export function ChannelMutateDrawer({ const isChannelDetailLoading = isEditing && isChannelLoading const supportsMultiKeyAddMode = currentType !== 57 && !(currentType === 41 && vertexKeyType === 'api_key') + const { currentStorageMode, requestedStorageMode, isConvertingStorage } = + resolveStorageModeConversion(isEditing, isMultiKeyChannel, { + key_storage_mode: keyStorageMode, + }) + const isConvertingToSingle = + isConvertingStorage && requestedStorageMode === 'single' + const isConvertingToMulti = + isConvertingStorage && requestedStorageMode === 'multi' + const showMultiLineKeys = + isBatchMode || + isConvertingToMulti || + (isMultiKeyChannel && !isConvertingToSingle) + const showMultiKeyStrategy = + (!isEditing && multiKeyMode === 'multi_to_single') || isConvertingToMulti const addModeOptions = useMemo( () => supportsMultiKeyAddMode @@ -975,6 +993,7 @@ export function ChannelMutateDrawer({ formErrors.multi_key_mode || formErrors.multi_key_type || formErrors.key_mode || + formErrors.key_storage_mode || formErrors.vertex_key_type || formErrors.aws_key_type || formErrors.azure_responses_version @@ -1653,6 +1672,40 @@ export function ChannelMutateDrawer({ return } + const { + requestedStorageMode: requestedStorageModeOnSubmit, + isConvertingStorage: isConvertingStorageOnSubmit, + } = resolveStorageModeConversion(isEditing, isMultiKeyChannel, data) + if (isConvertingStorageOnSubmit) { + if (!data.key?.trim()) { + form.setError('key', { + type: 'manual', + message: ERROR_MESSAGES.REQUIRED_KEY, + }) + return + } + const keyCount = countUpdateKeys(data.key, { + type: data.type, + vertexKeyType: data.vertex_key_type, + }) + if (requestedStorageModeOnSubmit === 'single' && keyCount !== 1) { + form.setError('key', { + type: 'manual', + message: t( + 'Enter exactly one new key to convert to single-key mode' + ), + }) + return + } + if (requestedStorageModeOnSubmit === 'multi' && keyCount < 2) { + form.setError('key', { + type: 'manual', + message: t('Enter at least two keys to convert to multi-key mode'), + }) + return + } + } + if (sensitiveLocked) { const dirtyFields = form.formState.dirtyFields as Partial< Record @@ -1744,6 +1797,7 @@ export function ChannelMutateDrawer({ }, [ isEditing, + isMultiKeyChannel, sensitiveLocked, form, confirmMissingModelMappings, @@ -2638,7 +2692,7 @@ export function ChannelMutateDrawer({ { const fileList = e.target.files const files = fileList @@ -2679,7 +2733,7 @@ export function ChannelMutateDrawer({ return } - const keyValue = isBatchMode + const keyValue = showMultiLineKeys ? JSON.stringify(keys) : JSON.stringify(keys[0]) @@ -2700,7 +2754,7 @@ export function ChannelMutateDrawer({ /> - {isBatchMode + {showMultiLineKeys ? t( 'Upload multiple JSON files in batch modes' ) @@ -3012,6 +3066,94 @@ export function ChannelMutateDrawer({ /> )} + {isEditing && ( + { + const storageModeOptions = [ + { + value: 'single', + label: 'Single Key', + }, + ...(supportsMultiKeyAddMode || + currentStorageMode === 'multi' + ? [ + { + value: 'multi', + label: + 'Multi-Key Mode (multiple keys, one channel)', + }, + ] + : []), + ] + let storageModeDescription = t( + 'Switching this mode converts the channel in place and requires a new key.' + ) + if (isConvertingToSingle) { + storageModeDescription = t( + 'Converting to single-key mode replaces the existing key pool. Enter a new single key.' + ) + } else if (isConvertingToMulti) { + storageModeDescription = t( + 'Converting to multi-key mode replaces the existing key. Enter at least two new keys, one per line.' + ) + } + return ( + + + {t('Key Storage Mode')} + + + + {storageModeDescription} + + + + ) + }} + /> + )} + - {keyModeDescription} - - )} + {isMultiKeyChannel && + !isConvertingToSingle && ( + + {keyModeDescription} + + )} ) - } else if (isBatchMode) { + } else if (showMultiLineKeys) { keyDescription = t( 'Enter one API key per line for batch creation' ) @@ -3090,14 +3249,14 @@ export function ChannelMutateDrawer({