diff --git a/.github/workflows/docker-build-fork.yml b/.github/workflows/docker-build-fork.yml new file mode 100644 index 000000000000..7ad2b5f4eef5 --- /dev/null +++ b/.github/workflows/docker-build-fork.yml @@ -0,0 +1,140 @@ +name: Publish fork image to GHCR + +on: + push: + branches: + - main + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: fork-image-main + cancel-in-progress: false + +env: + REGISTRY: ghcr.io + IMAGE_NAME: laplaceorange/new-api + +jobs: + prepare: + name: Resolve main revision + runs-on: ubuntu-latest + outputs: + sha: ${{ steps.version.outputs.sha }} + version: ${{ steps.version.outputs.version }} + + steps: + - name: Check out main + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + ref: main + + - name: Resolve build version + id: version + shell: bash + run: | + SHA=$(git rev-parse HEAD) + echo "sha=${SHA}" >> "${GITHUB_OUTPUT}" + echo "version=main-${SHA::12}" >> "${GITHUB_OUTPUT}" + + build: + name: Build and push (${{ matrix.arch }}) + needs: prepare + strategy: + fail-fast: false + matrix: + include: + - arch: amd64 + platform: linux/amd64 + runner: ubuntu-latest + - arch: arm64 + platform: linux/arm64 + runner: ubuntu-24.04-arm + runs-on: ${{ matrix.runner }} + timeout-minutes: 60 + permissions: + contents: read + packages: write + + steps: + - name: Check out + uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + ref: ${{ needs.prepare.outputs.sha }} + + - name: Write build version + shell: bash + run: echo "${{ needs.prepare.outputs.version }}" > VERSION + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 + + - name: Log in to GitHub Container Registry + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build and push image + uses: docker/build-push-action@53b7df96c91f9c12dcc8a07bcb9ccacbed38856a # v7.3.0 + with: + context: . + platforms: ${{ matrix.platform }} + push: true + tags: | + ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}:sha-${{ needs.prepare.outputs.sha }}-${{ matrix.arch }} + labels: | + org.opencontainers.image.source=https://github.com/${{ github.repository }} + org.opencontainers.image.revision=${{ needs.prepare.outputs.sha }} + org.opencontainers.image.version=${{ needs.prepare.outputs.version }} + cache-from: type=gha,scope=fork-${{ matrix.arch }} + cache-to: type=gha,mode=max,scope=fork-${{ matrix.arch }} + provenance: mode=max + sbom: true + + manifest: + name: Publish multi-arch manifest + needs: + - prepare + - build + runs-on: ubuntu-latest + timeout-minutes: 15 + permissions: + contents: read + packages: write + + steps: + - name: Log in to GitHub Container Registry + uses: docker/login-action@c99871dec2022cc055c062a10cc1a1310835ceb4 # v4.3.0 + with: + registry: ${{ env.REGISTRY }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Set up Docker Buildx + uses: docker/setup-buildx-action@bb05f3f5519dd87d3ba754cc423b652a5edd6d2c # v4.2.0 + + - name: Create and push manifests + shell: bash + run: | + docker buildx imagetools create \ + -t "${REGISTRY}/${IMAGE_NAME}:main" \ + -t "${REGISTRY}/${IMAGE_NAME}:latest" \ + -t "${REGISTRY}/${IMAGE_NAME}:sha-${{ needs.prepare.outputs.sha }}" \ + "${REGISTRY}/${IMAGE_NAME}:sha-${{ needs.prepare.outputs.sha }}-amd64" \ + "${REGISTRY}/${IMAGE_NAME}:sha-${{ needs.prepare.outputs.sha }}-arm64" + + - name: Record published image + shell: bash + run: | + docker buildx imagetools inspect "${REGISTRY}/${IMAGE_NAME}:sha-${{ needs.prepare.outputs.sha }}" \ + | tee manifest.txt + { + echo '### Published image' + echo '```text' + cat manifest.txt + echo '```' + } >> "${GITHUB_STEP_SUMMARY}" diff --git a/common/constants.go b/common/constants.go index d6b4fb52284c..6c0dbfdd3128 100644 --- a/common/constants.go +++ b/common/constants.go @@ -127,6 +127,8 @@ var QuotaForInvitee = 0 var ChannelDisableThreshold = 5.0 var AutomaticDisableChannelEnabled = false var AutomaticEnableChannelEnabled = false +var AutomaticDisableModelEnabled = false +var AutomaticEnableModelEnabled = false var QuotaRemindThreshold = 1000 var PreConsumedQuota = 500 diff --git a/controller/channel.go b/controller/channel.go index 3a1e58328923..b117566feed5 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -2,7 +2,6 @@ package controller import ( "context" - "encoding/json" "fmt" "net/http" "strconv" @@ -81,6 +80,10 @@ func applyChannelStatusFilter(query *gorm.DB, statusFilter int) *gorm.DB { return query } +func syncModelChannelAvailabilityAfterMutation(reason string) { + service.SyncModelChannelAvailabilityAfterMutation(reason) +} + func buildChannelListQuery(group string, statusFilter int, typeFilter int) *gorm.DB { query := model.DB.Model(&model.Channel{}) query = model.ApplyChannelGroupFilter(query, group) @@ -262,6 +265,7 @@ func FixChannelsAbilities(c *gin.Context) { common.ApiError(c, err) return } + syncModelChannelAvailabilityAfterMutation("channel.fix_abilities") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -485,12 +489,15 @@ func validateChannel(channel *model.Channel, isAdd bool) error { return fmt.Errorf("New API channel base URL cannot be empty") } - // 如果是添加操作，检查 channel 和 key 是否为空 - if isAdd { - if channel.Key == "" { + // Empty means "keep the existing key" on update, but an explicitly blank + // value is never a usable credential. + if isAdd || channel.Key != "" { + if !model.IsUsableChannelKey(channel.Key) { return fmt.Errorf("channel cannot be empty") } + } + if isAdd { // 检查模型名称长度是否超过 255 for _, m := range channel.GetModels() { if len(m) > 255 { @@ -593,7 +600,7 @@ func getVertexArrayKeys(keys string) ([]string, error) { case string: keyStr = strings.TrimSpace(v) default: - bytes, err := json.Marshal(v) + bytes, err := common.Marshal(v) if err != nil { return nil, fmt.Errorf("Vertex AI key JSON 编码失败: %w", err) } @@ -646,10 +653,10 @@ func AddChannel(c *gin.Context) { } else { cleanKeys := make([]string, 0) for _, key := range strings.Split(addChannelRequest.Channel.Key, "\n") { - if key == "" { + key = strings.TrimSpace(key) + if !model.IsUsableChannelKey(key) { continue } - key = strings.TrimSpace(key) cleanKeys = append(cleanKeys, key) } addChannelRequest.Channel.ChannelInfo.MultiKeySize = len(cleanKeys) @@ -668,7 +675,12 @@ func AddChannel(c *gin.Context) { return } } else { - keys = strings.Split(addChannelRequest.Channel.Key, "\n") + for _, key := range strings.Split(addChannelRequest.Channel.Key, "\n") { + key = strings.TrimSpace(key) + if model.IsUsableChannelKey(key) { + keys = append(keys, key) + } + } } case "single": keys = []string{addChannelRequest.Channel.Key} @@ -682,7 +694,7 @@ func AddChannel(c *gin.Context) { channels := make([]model.Channel, 0, len(keys)) for _, key := range keys { - if key == "" { + if !model.IsUsableChannelKey(key) { continue } localChannel := addChannelRequest.Channel @@ -696,16 +708,25 @@ func AddChannel(c *gin.Context) { } channels = append(channels, *localChannel) } + if len(channels) == 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "channel key cannot be empty", + }) + return + } err = model.BatchInsertChannels(channels) if err != nil { common.ApiError(c, err) return } + model.InitChannelCache() recordManageAudit(c, "channel.create", map[string]interface{}{ "name": addChannelRequest.Channel.Name, "type": addChannelRequest.Channel.Type, "count": len(channels), }) + syncModelChannelAvailabilityAfterMutation("channel.create") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -740,6 +761,7 @@ func DeleteChannel(c *gin.Context) { "id": id, "name": channelName, }) + syncModelChannelAvailabilityAfterMutation("channel.delete") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -760,6 +782,9 @@ func DeleteDisabledChannel(c *gin.Context) { recordManageAudit(c, "channel.delete_disabled", map[string]interface{}{ "count": rows, }) + if rows > 0 { + syncModelChannelAvailabilityAfterMutation("channel.delete_disabled") + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -799,6 +824,7 @@ func DisableTagChannels(c *gin.Context) { recordManageAudit(c, "channel.tag_disable", map[string]interface{}{ "tag": channelTag.Tag, }) + syncModelChannelAvailabilityAfterMutation("channel.tag_disable") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -825,6 +851,7 @@ func EnableTagChannels(c *gin.Context) { recordManageAudit(c, "channel.tag_enable", map[string]interface{}{ "tag": channelTag.Tag, }) + syncModelChannelAvailabilityAfterMutation("channel.tag_enable") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -856,7 +883,8 @@ func EditTagChannels(c *gin.Context) { } if channelTag.ParamOverride != nil { trimmed := strings.TrimSpace(*channelTag.ParamOverride) - if trimmed != "" && !json.Valid([]byte(trimmed)) { + var value any + if trimmed != "" && common.Unmarshal([]byte(trimmed), &value) != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "参数覆盖必须是合法的 JSON 格式", @@ -867,7 +895,8 @@ func EditTagChannels(c *gin.Context) { } if channelTag.HeaderOverride != nil { trimmed := strings.TrimSpace(*channelTag.HeaderOverride) - if trimmed != "" && !json.Valid([]byte(trimmed)) { + var value any + if trimmed != "" && common.Unmarshal([]byte(trimmed), &value) != nil { c.JSON(http.StatusOK, gin.H{ "success": false, "message": "请求头覆盖必须是合法的 JSON 格式", @@ -885,6 +914,9 @@ func EditTagChannels(c *gin.Context) { recordManageAudit(c, "channel.tag_edit", map[string]interface{}{ "tag": channelTag.Tag, }) + if channelTag.Models != nil { + syncModelChannelAvailabilityAfterMutation("channel.tag_edit") + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -919,6 +951,9 @@ func DeleteChannelBatch(c *gin.Context) { recordManageAudit(c, "channel.delete_batch", map[string]interface{}{ "count": deletedCount, }) + if deletedCount > 0 { + syncModelChannelAvailabilityAfterMutation("channel.delete_batch") + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -964,130 +999,167 @@ func UpdateChannel(c *gin.Context) { } clearChannelReadOnlyFields(&channel, requestData) - // 使用统一的校验函数 - if err := validateChannel(&channel.Channel, false); err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": err.Error(), - }) - 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 - } - originProxy := originChannel.GetSetting().Proxy - proxyChanged := false - if _, settingProvided := requestData["setting"]; settingProvided { - newProxy, _ := service.NormalizeProxyURL(channel.GetSetting().Proxy) - normalizedOriginProxy, originProxyErr := service.NormalizeProxyURL(originProxy) - proxyChanged = originProxyErr != nil || normalizedOriginProxy != newProxy - } + // ChannelInfo is also mutated by automatic multi-key health updates. Keep the + // in-process polling lock outside the database row lock to match the status + // update lock order, then apply only the explicit request fields to the latest + // persisted row inside one transaction. + pollingLock := model.GetChannelPollingLock(channel.Id) + pollingLock.Lock() + defer pollingLock.Unlock() - // Always copy the original ChannelInfo so that fields like IsMultiKey and MultiKeySize are retained. - channel.ChannelInfo = originChannel.ChannelInfo + canWriteSensitiveFields := authz.Can(c.GetInt("id"), c.GetInt("role"), authz.ChannelSensitiveWrite) + permissionDenied := false + proxyChanged := false + originProxy := "" + var originChannel *model.Channel + updatedChannel, err := model.UpdateChannelAtomically(channel.Id, func(current *model.Channel) error { + if channelHasSensitiveChanges(&channel, current, requestData) && !canWriteSensitiveFields { + permissionDenied = true + return fmt.Errorf("insufficient privilege") + } - if channelHasSensitiveChanges(&channel, originChannel, requestData) && - !authz.Can(c.GetInt("id"), c.GetInt("role"), authz.ChannelSensitiveWrite) { - common.ApiErrorI18n(c, i18n.MsgAuthInsufficientPrivilege) - return - } + origin := *current + originChannel = &origin - // If the request explicitly specifies a new MultiKeyMode, apply it on top of the original info. - if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" { - channel.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode) - } + if _, ok := requestData["type"]; ok { + current.Type = channel.Type + } + if _, ok := requestData["openai_organization"]; ok { + current.OpenAIOrganization = channel.OpenAIOrganization + } + if _, ok := requestData["test_model"]; ok { + current.TestModel = channel.TestModel + } + if _, ok := requestData["name"]; ok { + current.Name = channel.Name + } + if _, ok := requestData["weight"]; ok { + current.Weight = channel.Weight + } + if _, ok := requestData["base_url"]; ok { + current.BaseURL = channel.BaseURL + } + if _, ok := requestData["other"]; ok { + current.Other = channel.Other + } + if _, ok := requestData["models"]; ok { + current.Models = channel.Models + } + if _, ok := requestData["group"]; ok { + current.Group = channel.Group + } + if _, ok := requestData["model_mapping"]; ok { + current.ModelMapping = channel.ModelMapping + } + if _, ok := requestData["status_code_mapping"]; ok { + current.StatusCodeMapping = channel.StatusCodeMapping + } + if _, ok := requestData["priority"]; ok { + current.Priority = channel.Priority + } + if _, ok := requestData["auto_ban"]; ok { + current.AutoBan = channel.AutoBan + } + if _, ok := requestData["other_info"]; ok { + current.OtherInfo = channel.OtherInfo + } + if _, ok := requestData["tag"]; ok { + current.Tag = channel.Tag + } + if _, ok := requestData["setting"]; ok { + current.Setting = channel.Setting + } + if _, ok := requestData["param_override"]; ok { + current.ParamOverride = channel.ParamOverride + } + if _, ok := requestData["header_override"]; ok { + current.HeaderOverride = channel.HeaderOverride + } + if _, ok := requestData["remark"]; ok { + current.Remark = channel.Remark + } + if _, ok := requestData["settings"]; ok { + current.OtherSettings = channel.OtherSettings + } + if channel.MultiKeyMode != nil && *channel.MultiKeyMode != "" { + current.ChannelInfo.MultiKeyMode = constant.MultiKeyMode(*channel.MultiKeyMode) + } - // 处理多key模式下的密钥追加/覆盖逻辑 - if channel.KeyMode != nil && channel.ChannelInfo.IsMultiKey { - switch *channel.KeyMode { - case "append": - // 追加模式：将新密钥添加到现有密钥列表 - if originChannel.Key != "" { + if _, keyProvided := requestData["key"]; keyProvided && channel.Key != "" { + nextKey := channel.Key + if channel.KeyMode != nil && *channel.KeyMode == "append" && current.ChannelInfo.IsMultiKey && current.Key != "" { + existingKeys := current.GetKeys() var newKeys []string - var existingKeys []string - - // 解析现有密钥 - if strings.HasPrefix(strings.TrimSpace(originChannel.Key), "[") { - // JSON数组格式 - var arr []json.RawMessage - if err := json.Unmarshal([]byte(strings.TrimSpace(originChannel.Key)), &arr); err == nil { - existingKeys = make([]string, len(arr)) - for i, v := range arr { - existingKeys[i] = string(v) - } - } - } else { - // 换行分隔格式 - existingKeys = strings.Split(strings.Trim(originChannel.Key, "\n"), "\n") - } - - // 处理 Vertex AI 的特殊情况 - if channel.Type == constant.ChannelTypeVertexAi && channel.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { - // 尝试解析新密钥为JSON数组 + if current.Type == constant.ChannelTypeVertexAi && current.GetOtherSettings().VertexKeyType != dto.VertexKeyTypeAPIKey { if strings.HasPrefix(strings.TrimSpace(channel.Key), "[") { - array, err := getVertexArrayKeys(channel.Key) + newKeys, err = getVertexArrayKeys(channel.Key) if err != nil { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "追加密钥解析失败: " + err.Error(), - }) - return + return fmt.Errorf("追加密钥解析失败: %w", err) } - newKeys = array } else { - // 单个JSON密钥 newKeys = []string{channel.Key} } } else { - // 普通渠道的处理 - inputKeys := strings.Split(channel.Key, "\n") - for _, key := range inputKeys { - key = strings.TrimSpace(key) - if key != "" { + for _, key := range strings.Split(channel.Key, "\n") { + if key = strings.TrimSpace(key); key != "" { newKeys = append(newKeys, key) } } } seen := make(map[string]struct{}, len(existingKeys)+len(newKeys)) + allKeys := make([]string, 0, len(existingKeys)+len(newKeys)) for _, key := range existingKeys { normalized := strings.TrimSpace(key) if normalized == "" { continue } seen[normalized] = struct{}{} + allKeys = append(allKeys, key) } - dedupedNewKeys := make([]string, 0, len(newKeys)) for _, key := range newKeys { normalized := strings.TrimSpace(key) if normalized == "" { continue } - if _, ok := seen[normalized]; ok { + if _, exists := seen[normalized]; exists { continue } seen[normalized] = struct{}{} - dedupedNewKeys = append(dedupedNewKeys, normalized) + allKeys = append(allKeys, normalized) } - - allKeys := append(existingKeys, dedupedNewKeys...) - channel.Key = strings.Join(allKeys, "\n") + nextKey = strings.Join(allKeys, "\n") + } + if current.ChannelInfo.IsMultiKey { + current.ReplaceMultiKeyKeys(nextKey) + if current.Key == "" { + return fmt.Errorf("channel cannot be empty") + } + } else { + current.Key = nextKey } - case "replace": - // 覆盖模式：直接使用新密钥（默认行为，不需要特殊处理） } - } - err = channel.Update() + + if err := validateChannel(current, false); err != nil { + return err + } + originProxy = origin.GetSetting().Proxy + if _, settingProvided := requestData["setting"]; settingProvided { + newProxy, _ := service.NormalizeProxyURL(current.GetSetting().Proxy) + normalizedOriginProxy, originProxyErr := service.NormalizeProxyURL(originProxy) + proxyChanged = originProxyErr != nil || normalizedOriginProxy != newProxy + } + return nil + }) if err != nil { + if permissionDenied { + common.ApiErrorI18n(c, i18n.MsgAuthInsufficientPrivilege) + return + } common.ApiError(c, err) return } + channel.Channel = *updatedChannel model.InitChannelCache() if proxyChanged { service.InvalidateProxyClient(originProxy) @@ -1114,6 +1186,9 @@ func UpdateChannel(c *gin.Context) { "name": channel.Name, "changed_fields": changedFields, }) + // Key and ChannelInfo changes can alter multi-key availability even when the + // model list is unchanged, so every successful channel update is reconciled. + syncModelChannelAvailabilityAfterMutation("channel.update") channel.Key = "" clearChannelInfo(&channel.Channel) c.JSON(http.StatusOK, gin.H{ @@ -1135,7 +1210,11 @@ func UpdateChannelStatus(c *gin.Context) { common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } - changed := model.UpdateChannelStatus(id, "", req.Status, "manual operation") + changed, err := model.UpdateChannelStatusWithError(id, "", req.Status, "manual operation") + if err != nil { + common.ApiError(c, err) + return + } if changed { model.InitChannelCache() } @@ -1144,6 +1223,9 @@ func UpdateChannelStatus(c *gin.Context) { "status": req.Status, "changed": changed, }) + if changed { + syncModelChannelAvailabilityAfterMutation("channel.status_update") + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -1158,8 +1240,14 @@ func BatchUpdateChannelStatus(c *gin.Context) { return } changedCount := 0 + failedIds := make([]int, 0) for _, id := range req.Ids { - if model.UpdateChannelStatus(id, "", req.Status, "manual batch operation") { + changed, err := model.UpdateChannelStatusWithError(id, "", req.Status, "manual batch operation") + if err != nil { + failedIds = append(failedIds, id) + continue + } + if changed { changedCount++ } } @@ -1167,10 +1255,25 @@ func BatchUpdateChannelStatus(c *gin.Context) { model.InitChannelCache() } recordManageAudit(c, "channel.status_update_batch", map[string]interface{}{ - "count": changedCount, - "total": len(req.Ids), - "status": req.Status, + "count": changedCount, + "total": len(req.Ids), + "status": req.Status, + "failed_ids": failedIds, }) + if changedCount > 0 { + syncModelChannelAvailabilityAfterMutation("channel.status_update_batch") + } + if len(failedIds) > 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": fmt.Sprintf("failed to update channel status for ids: %v", failedIds), + "data": gin.H{ + "changed": changedCount, + "failed_ids": failedIds, + }, + }) + return + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -1455,6 +1558,7 @@ func CopyChannel(c *gin.Context) { "id": clone.Id, "name": clone.Name, }) + syncModelChannelAvailabilityAfterMutation("channel.copy") // success c.JSON(http.StatusOK, gin.H{"success": true, "message": "", "data": gin.H{"id": clone.Id}}) } @@ -1521,6 +1625,27 @@ func ManageMultiKeys(c *gin.Context) { return } + lock := model.GetChannelPollingLock(channel.Id) + lock.Lock() + defer lock.Unlock() + // Reload after taking the per-channel lock. Status polling may have updated + // key state between the authorization read and lock acquisition. + channel, err = model.GetChannelById(request.ChannelId, true) + if err != nil { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "渠道不存在", + }) + return + } + if !channel.ChannelInfo.IsMultiKey { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "该渠道不是多密钥模式", + }) + return + } + // get_key_status 为只读查询，不记录审计；其余为修改操作，记录审计并跳过中间件兜底。 if request.Action == "get_key_status" { markAuditLogged(c) @@ -1531,10 +1656,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() @@ -1555,6 +1676,9 @@ func ManageMultiKeys(c *gin.Context) { // Build all key status data first var allKeyStatusList []KeyStatus for i, key := range keys { + if !model.IsUsableChannelKey(key) { + continue + } status := 1 // default enabled var disabledTime int64 var reason string @@ -1660,33 +1784,35 @@ func ManageMultiKeys(c *gin.Context) { } keyIndex := *request.KeyIndex - if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "密钥索引超出范围", - }) - return - } - - if channel.ChannelInfo.MultiKeyStatusList == nil { - channel.ChannelInfo.MultiKeyStatusList = make(map[int]int) - } - if channel.ChannelInfo.MultiKeyDisabledTime == nil { - channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) - } - if channel.ChannelInfo.MultiKeyDisabledReason == nil { - channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) - } - - channel.ChannelInfo.MultiKeyStatusList[keyIndex] = 2 // disabled - - err = channel.Update() + channel, err = model.UpdateChannelAtomically(request.ChannelId, func(current *model.Channel) error { + if !current.ChannelInfo.IsMultiKey { + return fmt.Errorf("该渠道不是多密钥模式") + } + keys := current.GetKeys() + if keyIndex < 0 || keyIndex >= current.ChannelInfo.MultiKeySize || + keyIndex >= len(keys) || !model.IsUsableChannelKey(keys[keyIndex]) { + return fmt.Errorf("密钥索引超出范围") + } + if current.ChannelInfo.MultiKeyStatusList == nil { + current.ChannelInfo.MultiKeyStatusList = make(map[int]int) + } + if current.ChannelInfo.MultiKeyDisabledTime == nil { + current.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) + } + if current.ChannelInfo.MultiKeyDisabledReason == nil { + current.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) + } + current.ChannelInfo.MultiKeyStatusList[keyIndex] = common.ChannelStatusManuallyDisabled + reconcileMultiKeyChannelStatus(current, false) + return nil + }) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + syncModelChannelAvailabilityAfterMutation("channel.multikey.disable_key") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "密钥已禁用", @@ -1703,32 +1829,29 @@ func ManageMultiKeys(c *gin.Context) { } keyIndex := *request.KeyIndex - if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "密钥索引超出范围", - }) - return - } - - // 从状态列表中删除该密钥的记录，使其回到默认启用状态 - if channel.ChannelInfo.MultiKeyStatusList != nil { - delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex) - } - if channel.ChannelInfo.MultiKeyDisabledTime != nil { - delete(channel.ChannelInfo.MultiKeyDisabledTime, keyIndex) - } - if channel.ChannelInfo.MultiKeyDisabledReason != nil { - delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex) - } - - err = channel.Update() + channel, err = model.UpdateChannelAtomically(request.ChannelId, func(current *model.Channel) error { + if !current.ChannelInfo.IsMultiKey { + return fmt.Errorf("该渠道不是多密钥模式") + } + keys := current.GetKeys() + if keyIndex < 0 || keyIndex >= current.ChannelInfo.MultiKeySize || + keyIndex >= len(keys) || !model.IsUsableChannelKey(keys[keyIndex]) { + return fmt.Errorf("密钥索引超出范围") + } + // Removing the entry restores the default enabled state. + delete(current.ChannelInfo.MultiKeyStatusList, keyIndex) + delete(current.ChannelInfo.MultiKeyDisabledTime, keyIndex) + delete(current.ChannelInfo.MultiKeyDisabledReason, keyIndex) + reconcileMultiKeyChannelStatus(current, true) + return nil + }) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + syncModelChannelAvailabilityAfterMutation("channel.multikey.enable_key") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "密钥已启用", @@ -1736,23 +1859,30 @@ func ManageMultiKeys(c *gin.Context) { return case "enable_all_keys": - // 清空所有禁用状态，使所有密钥回到默认启用状态 var enabledCount int - if channel.ChannelInfo.MultiKeyStatusList != nil { - enabledCount = len(channel.ChannelInfo.MultiKeyStatusList) - } - - channel.ChannelInfo.MultiKeyStatusList = make(map[int]int) - channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) - channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) - - err = channel.Update() + channel, err = model.UpdateChannelAtomically(request.ChannelId, func(current *model.Channel) error { + if !current.ChannelInfo.IsMultiKey { + return fmt.Errorf("该渠道不是多密钥模式") + } + keys := current.GetKeys() + for keyIndex := range current.ChannelInfo.MultiKeyStatusList { + if keyIndex >= 0 && keyIndex < len(keys) && model.IsUsableChannelKey(keys[keyIndex]) { + enabledCount++ + } + } + current.ChannelInfo.MultiKeyStatusList = make(map[int]int) + current.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) + current.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) + reconcileMultiKeyChannelStatus(current, true) + return nil + }) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + syncModelChannelAvailabilityAfterMutation("channel.multikey.enable_all_keys") c.JSON(http.StatusOK, gin.H{ "success": true, "message": fmt.Sprintf("已启用 %d 个密钥", enabledCount), @@ -1760,46 +1890,46 @@ func ManageMultiKeys(c *gin.Context) { return case "disable_all_keys": - // 禁用所有启用的密钥 - if channel.ChannelInfo.MultiKeyStatusList == nil { - channel.ChannelInfo.MultiKeyStatusList = make(map[int]int) - } - if channel.ChannelInfo.MultiKeyDisabledTime == nil { - channel.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) - } - if channel.ChannelInfo.MultiKeyDisabledReason == nil { - channel.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) - } - var disabledCount int - for i := 0; i < channel.ChannelInfo.MultiKeySize; i++ { - status := 1 // default enabled - if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists { - status = s + channel, err = model.UpdateChannelAtomically(request.ChannelId, func(current *model.Channel) error { + if !current.ChannelInfo.IsMultiKey { + return fmt.Errorf("该渠道不是多密钥模式") } - - // 只禁用当前启用的密钥 - if status == 1 { - channel.ChannelInfo.MultiKeyStatusList[i] = 2 // disabled - disabledCount++ + if current.ChannelInfo.MultiKeyStatusList == nil { + current.ChannelInfo.MultiKeyStatusList = make(map[int]int) } - } - - if disabledCount == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "没有可禁用的密钥", - }) - return - } - - err = channel.Update() + if current.ChannelInfo.MultiKeyDisabledTime == nil { + current.ChannelInfo.MultiKeyDisabledTime = make(map[int]int64) + } + if current.ChannelInfo.MultiKeyDisabledReason == nil { + current.ChannelInfo.MultiKeyDisabledReason = make(map[int]string) + } + for i, key := range current.GetKeys() { + if !model.IsUsableChannelKey(key) { + continue + } + status := common.ChannelStatusEnabled + if persistedStatus, exists := current.ChannelInfo.MultiKeyStatusList[i]; exists { + status = persistedStatus + } + if status == common.ChannelStatusEnabled { + current.ChannelInfo.MultiKeyStatusList[i] = common.ChannelStatusManuallyDisabled + disabledCount++ + } + } + if disabledCount == 0 { + return fmt.Errorf("没有可禁用的密钥") + } + reconcileMultiKeyChannelStatus(current, false) + return nil + }) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + syncModelChannelAvailabilityAfterMutation("channel.multikey.disable_all_keys") c.JSON(http.StatusOK, gin.H{ "success": true, "message": fmt.Sprintf("已禁用 %d 个密钥", disabledCount), @@ -1816,70 +1946,54 @@ func ManageMultiKeys(c *gin.Context) { } keyIndex := *request.KeyIndex - if keyIndex < 0 || keyIndex >= channel.ChannelInfo.MultiKeySize { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "密钥索引超出范围", - }) - return - } - - keys := channel.GetKeys() - var remainingKeys []string - var newStatusList = make(map[int]int) - var newDisabledTime = make(map[int]int64) - var newDisabledReason = make(map[int]string) - - newIndex := 0 - for i, key := range keys { - // 跳过要删除的密钥 - if i == keyIndex { - continue + channel, err = model.UpdateChannelAtomically(request.ChannelId, func(current *model.Channel) error { + if !current.ChannelInfo.IsMultiKey { + return fmt.Errorf("该渠道不是多密钥模式") + } + keys := current.GetKeys() + if keyIndex < 0 || keyIndex >= current.ChannelInfo.MultiKeySize || + keyIndex >= len(keys) || !model.IsUsableChannelKey(keys[keyIndex]) { + return fmt.Errorf("密钥索引超出范围") } - remainingKeys = append(remainingKeys, key) - - // 保留其他密钥的状态信息，重新索引 - if channel.ChannelInfo.MultiKeyStatusList != nil { - if status, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists && status != 1 { + remainingKeys := make([]string, 0, len(keys)-1) + newStatusList := make(map[int]int) + newDisabledTime := make(map[int]int64) + newDisabledReason := make(map[int]string) + newIndex := 0 + for i, key := range keys { + if i == keyIndex || !model.IsUsableChannelKey(key) { + continue + } + remainingKeys = append(remainingKeys, key) + if status, exists := current.ChannelInfo.MultiKeyStatusList[i]; exists && status != common.ChannelStatusEnabled { newStatusList[newIndex] = status } - } - if channel.ChannelInfo.MultiKeyDisabledTime != nil { - if t, exists := channel.ChannelInfo.MultiKeyDisabledTime[i]; exists { - newDisabledTime[newIndex] = t + if disabledAt, exists := current.ChannelInfo.MultiKeyDisabledTime[i]; exists { + newDisabledTime[newIndex] = disabledAt } - } - if channel.ChannelInfo.MultiKeyDisabledReason != nil { - if r, exists := channel.ChannelInfo.MultiKeyDisabledReason[i]; exists { - newDisabledReason[newIndex] = r + if reason, exists := current.ChannelInfo.MultiKeyDisabledReason[i]; exists { + newDisabledReason[newIndex] = reason } + newIndex++ } - newIndex++ - } - - if len(remainingKeys) == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "不能删除最后一个密钥", - }) - return - } - - // Update channel with remaining keys - channel.Key = strings.Join(remainingKeys, "\n") - channel.ChannelInfo.MultiKeySize = len(remainingKeys) - channel.ChannelInfo.MultiKeyStatusList = newStatusList - channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime - channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason - - err = channel.Update() + if len(remainingKeys) == 0 { + return fmt.Errorf("不能删除最后一个密钥") + } + current.Key = strings.Join(remainingKeys, "\n") + current.ChannelInfo.MultiKeyStatusList = newStatusList + current.ChannelInfo.MultiKeyDisabledTime = newDisabledTime + current.ChannelInfo.MultiKeyDisabledReason = newDisabledReason + reconcileMultiKeyChannelStatus(current, true) + return nil + }) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + syncModelChannelAvailabilityAfterMutation("channel.multikey.delete_key") c.JSON(http.StatusOK, gin.H{ "success": true, "message": "密钥已删除", @@ -1887,67 +2001,61 @@ func ManageMultiKeys(c *gin.Context) { return case "delete_disabled_keys": - keys := channel.GetKeys() - var remainingKeys []string var deletedCount int - var newStatusList = make(map[int]int) - var newDisabledTime = make(map[int]int64) - var newDisabledReason = make(map[int]string) - - newIndex := 0 - for i, key := range keys { - status := 1 // default enabled - if channel.ChannelInfo.MultiKeyStatusList != nil { - if s, exists := channel.ChannelInfo.MultiKeyStatusList[i]; exists { - status = s - } + channel, err = model.UpdateChannelAtomically(request.ChannelId, func(current *model.Channel) error { + if !current.ChannelInfo.IsMultiKey { + return fmt.Errorf("该渠道不是多密钥模式") } - - // 只删除自动禁用（status == 3）的密钥，保留启用（status == 1）和手动禁用（status == 2）的密钥 - if status == 3 { - deletedCount++ - } else { + keys := current.GetKeys() + remainingKeys := make([]string, 0, len(keys)) + newStatusList := make(map[int]int) + newDisabledTime := make(map[int]int64) + newDisabledReason := make(map[int]string) + newIndex := 0 + for i, key := range keys { + if !model.IsUsableChannelKey(key) { + continue + } + status := common.ChannelStatusEnabled + if persistedStatus, exists := current.ChannelInfo.MultiKeyStatusList[i]; exists { + status = persistedStatus + } + if status == common.ChannelStatusAutoDisabled { + deletedCount++ + continue + } remainingKeys = append(remainingKeys, key) - // 保留非自动禁用密钥的状态信息，重新索引 - if status != 1 { + if status != common.ChannelStatusEnabled { newStatusList[newIndex] = status - if channel.ChannelInfo.MultiKeyDisabledTime != nil { - if t, exists := channel.ChannelInfo.MultiKeyDisabledTime[i]; exists { - newDisabledTime[newIndex] = t - } + if disabledAt, exists := current.ChannelInfo.MultiKeyDisabledTime[i]; exists { + newDisabledTime[newIndex] = disabledAt } - if channel.ChannelInfo.MultiKeyDisabledReason != nil { - if r, exists := channel.ChannelInfo.MultiKeyDisabledReason[i]; exists { - newDisabledReason[newIndex] = r - } + if reason, exists := current.ChannelInfo.MultiKeyDisabledReason[i]; exists { + newDisabledReason[newIndex] = reason } } newIndex++ } - } - - if deletedCount == 0 { - c.JSON(http.StatusOK, gin.H{ - "success": false, - "message": "没有需要删除的自动禁用密钥", - }) - return - } - - // Update channel with remaining keys - channel.Key = strings.Join(remainingKeys, "\n") - channel.ChannelInfo.MultiKeySize = len(remainingKeys) - channel.ChannelInfo.MultiKeyStatusList = newStatusList - channel.ChannelInfo.MultiKeyDisabledTime = newDisabledTime - channel.ChannelInfo.MultiKeyDisabledReason = newDisabledReason - - err = channel.Update() + if deletedCount == 0 { + return fmt.Errorf("没有需要删除的自动禁用密钥") + } + if len(remainingKeys) == 0 { + return fmt.Errorf("不能删除所有密钥") + } + current.Key = strings.Join(remainingKeys, "\n") + current.ChannelInfo.MultiKeyStatusList = newStatusList + current.ChannelInfo.MultiKeyDisabledTime = newDisabledTime + current.ChannelInfo.MultiKeyDisabledReason = newDisabledReason + reconcileMultiKeyChannelStatus(current, true) + return nil + }) if err != nil { common.ApiError(c, err) return } model.InitChannelCache() + syncModelChannelAvailabilityAfterMutation("channel.multikey.delete_disabled_keys") c.JSON(http.StatusOK, gin.H{ "success": true, "message": fmt.Sprintf("已删除 %d 个自动禁用的密钥", deletedCount), @@ -1968,6 +2076,32 @@ func multiKeyActionRequiresSensitiveWrite(action string) bool { return action == "delete_key" || action == "delete_disabled_keys" } +func reconcileMultiKeyChannelStatus(channel *model.Channel, enableWhenAvailable bool) { + if channel == nil || !channel.ChannelInfo.IsMultiKey { + return + } + if !channel.HasEnabledMultiKey() { + if channel.Status != common.ChannelStatusManuallyDisabled { + channel.Status = common.ChannelStatusAutoDisabled + info := channel.GetOtherInfo() + info["status_reason"] = "All keys are disabled" + info["status_time"] = common.GetTimestamp() + channel.SetOtherInfo(info) + } + return + } + if enableWhenAvailable && channel.Status == common.ChannelStatusAutoDisabled { + info := channel.GetOtherInfo() + statusReason, _ := info["status_reason"].(string) + if statusReason == "All keys are disabled" { + channel.Status = common.ChannelStatusEnabled + delete(info, "status_reason") + delete(info, "status_time") + channel.SetOtherInfo(info) + } + } +} + // OllamaPullModel 拉取 Ollama 模型 func OllamaPullModel(c *gin.Context) { var req struct { @@ -2088,7 +2222,7 @@ func OllamaPullModelStream(c *gin.Context) { // 创建进度回调函数 progressCallback := func(progress ollama.OllamaPullResponse) { - data, _ := json.Marshal(progress) + data, _ := common.Marshal(progress) fmt.Fprintf(c.Writer, "data: %s\n\n", string(data)) c.Writer.Flush() } @@ -2097,12 +2231,12 @@ func OllamaPullModelStream(c *gin.Context) { err = ollama.PullOllamaModelStream(baseURL, key, req.ModelName, progressCallback) if err != nil { - errorData, _ := json.Marshal(gin.H{ + errorData, _ := common.Marshal(gin.H{ "error": err.Error(), }) fmt.Fprintf(c.Writer, "data: %s\n\n", string(errorData)) } else { - successData, _ := json.Marshal(gin.H{ + successData, _ := common.Marshal(gin.H{ "message": fmt.Sprintf("Model %s pulled successfully", req.ModelName), }) fmt.Fprintf(c.Writer, "data: %s\n\n", string(successData)) diff --git a/controller/channel_multi_key_test.go b/controller/channel_multi_key_test.go new file mode 100644 index 000000000000..945bfd5e81cd --- /dev/null +++ b/controller/channel_multi_key_test.go @@ -0,0 +1,484 @@ +package controller + +import ( + "bytes" + "errors" + "net/http" + "net/http/httptest" + "strconv" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupMultiKeyControllerTest(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate( + &model.Channel{}, &model.Ability{}, &model.User{}, &model.Log{}, &model.Option{}, &model.Model{}, + )) + + originalDB := model.DB + originalLogDB := model.LOG_DB + originalMainDatabaseType := common.MainDatabaseType() + originalLogDatabaseType := common.LogDatabaseType() + originalMemoryCacheEnabled := common.MemoryCacheEnabled + originalRedisEnabled := common.RedisEnabled + originalAutomaticDisable := common.AutomaticDisableModelEnabled + originalAutomaticEnable := common.AutomaticEnableModelEnabled + model.DB = db + model.LOG_DB = db + common.SetDatabaseTypes(common.DatabaseTypeSQLite, common.DatabaseTypeSQLite) + common.MemoryCacheEnabled = false + common.RedisEnabled = false + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = false + t.Cleanup(func() { + model.DB = originalDB + model.LOG_DB = originalLogDB + common.SetDatabaseTypes(originalMainDatabaseType, originalLogDatabaseType) + common.MemoryCacheEnabled = originalMemoryCacheEnabled + common.RedisEnabled = originalRedisEnabled + common.AutomaticDisableModelEnabled = originalAutomaticDisable + common.AutomaticEnableModelEnabled = originalAutomaticEnable + }) + return db +} + +func TestUpdateChannelAppliesOnlyExplicitMutableFields(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: "before", + Key: "key-a\nkey-b", + Status: common.ChannelStatusManuallyDisabled, + Models: "gpt-4", + Group: "default", + UsedQuota: 42, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{1: common.ChannelStatusAutoDisabled}, + MultiKeyDisabledReason: map[int]string{1: "rejected"}, + MultiKeyPollingIndex: 1, + }, + } + require.NoError(t, db.Create(&channel).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + body, err := common.Marshal(map[string]any{ + "id": channel.Id, + "name": "after", + "used_quota": 9999, + "channel_info": map[string]any{ + "is_multi_key": true, + "multi_key_status_list": map[string]int{}, + }, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/channel/", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", 1) + ctx.Set("role", common.RoleRootUser) + + UpdateChannel(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + assert.Equal(t, "after", stored.Name) + assert.Equal(t, int64(42), stored.UsedQuota) + assert.Equal(t, common.ChannelStatusManuallyDisabled, stored.Status) + assert.Equal(t, channel.ChannelInfo, stored.ChannelInfo) + var ability model.Ability + require.NoError(t, db.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestUpdateChannelNormalizesMultiKeyReplacement(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: "replace-keys", + Key: "key-a\nkey-b", + Status: common.ChannelStatusEnabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{1: common.ChannelStatusAutoDisabled}, + MultiKeyDisabledReason: map[int]string{1: "rejected"}, + }, + } + require.NoError(t, db.Create(&channel).Error) + + body, err := common.Marshal(map[string]any{ + "id": channel.Id, + "key": " \nkey-b\n\"\"\nkey-c", + "key_mode": "replace", + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/channel/", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", 1) + ctx.Set("role", common.RoleRootUser) + + UpdateChannel(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + assert.Equal(t, "key-b\nkey-c", stored.Key) + assert.Equal(t, 2, stored.ChannelInfo.MultiKeySize) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.ChannelInfo.MultiKeyStatusList[0]) + assert.Equal(t, "rejected", stored.ChannelInfo.MultiKeyDisabledReason[0]) +} + +func TestUpdateChannelRejectsEmptyMultiKeyReplacement(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + key string + }{ + {name: "blank lines", key: " \n\t"}, + {name: "legacy empty JSON string", key: `""`}, + {name: "empty JSON array entry", key: `[""]`}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: test.name, + Key: "key-a", + Status: common.ChannelStatusEnabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 1, + }, + } + require.NoError(t, db.Create(&channel).Error) + + body, err := common.Marshal(map[string]any{ + "id": channel.Id, + "key": test.key, + "key_mode": "replace", + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/channel/", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", 1) + ctx.Set("role", common.RoleRootUser) + + UpdateChannel(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Contains(t, response.Message, "cannot be empty") + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + assert.Equal(t, "key-a", stored.Key) + }) + } +} + +func TestUpdateChannelStatusReturnsFailureForMissingChannel(t *testing.T) { + gin.SetMode(gin.TestMode) + setupMultiKeyControllerTest(t) + body, err := common.Marshal(ChannelStatusRequest{Status: common.ChannelStatusManuallyDisabled}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = []gin.Param{{Key: "id", Value: "99999"}} + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/99999/status", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + UpdateChannelStatus(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.NotEmpty(t, response.Message) +} + +func TestUpdateChannelStatusReturnsFailureWhenAbilityUpdateRollsBack(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: "status-rollback", + Key: "key", + Status: common.ChannelStatusEnabled, + Models: "gpt-4", + Group: "default", + } + require.NoError(t, db.Create(&channel).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: true, + }).Error) + + const callbackName = "test:fail_channel_status_ability_update" + require.NoError(t, db.Callback().Update().Before("gorm:update").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement != nil && tx.Statement.Table == "abilities" { + tx.AddError(errors.New("forced ability update failure")) + } + })) + t.Cleanup(func() { _ = db.Callback().Update().Remove(callbackName) }) + + body, err := common.Marshal(ChannelStatusRequest{Status: common.ChannelStatusManuallyDisabled}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = []gin.Param{{Key: "id", Value: strconv.Itoa(channel.Id)}} + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/status", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + UpdateChannelStatus(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Contains(t, response.Message, "forced ability update failure") + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusEnabled, stored.Status) + var ability model.Ability + require.NoError(t, db.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.True(t, ability.Enabled) +} + +func TestUpdateChannelStatusRejectsEnablingMultiKeyWithoutUsableCredential(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: "no-usable-key", + Key: " \n\t", + Status: common.ChannelStatusAutoDisabled, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + }, + } + require.NoError(t, db.Create(&channel).Error) + + body, err := common.Marshal(ChannelStatusRequest{Status: common.ChannelStatusEnabled}) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Params = []gin.Param{{Key: "id", Value: strconv.Itoa(channel.Id)}} + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/status", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + UpdateChannelStatus(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Contains(t, response.Message, "no usable key") +} + +func TestBatchUpdateChannelStatusReportsFailedIDs(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: "batch-status", + Key: "key", + Status: common.ChannelStatusEnabled, + Models: "gpt-4", + Group: "default", + } + require.NoError(t, db.Create(&channel).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: true, + }).Error) + missingID := channel.Id + 9999 + body, err := common.Marshal(ChannelStatusBatchRequest{ + Ids: []int{channel.Id, missingID}, + Status: common.ChannelStatusManuallyDisabled, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/status/batch", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + BatchUpdateChannelStatus(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + Data struct { + Changed int `json:"changed"` + FailedIDs []int `json:"failed_ids"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Equal(t, 1, response.Data.Changed) + assert.Equal(t, []int{missingID}, response.Data.FailedIDs) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusManuallyDisabled, stored.Status) +} + +func TestManageMultiKeysDoesNotRecoverUnrelatedAutoDisable(t *testing.T) { + gin.SetMode(gin.TestMode) + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: "unrelated-auto-disable", + Key: "key-a\nkey-b", + Status: common.ChannelStatusAutoDisabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{0: common.ChannelStatusAutoDisabled, 1: common.ChannelStatusAutoDisabled}, + }, + } + channel.SetOtherInfo(map[string]interface{}{"status_reason": "provider unavailable"}) + require.NoError(t, db.Create(&channel).Error) + require.NoError(t, db.Create(&model.Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + body, err := common.Marshal(MultiKeyManageRequest{ + ChannelId: channel.Id, + Action: "enable_key", + KeyIndex: common.GetPointer(0), + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/multi-key", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", 1) + ctx.Set("role", common.RoleRootUser) + + ManageMultiKeys(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + require.True(t, response.Success, response.Message) + + var stored model.Channel + require.NoError(t, db.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.Status) + assert.NotContains(t, stored.ChannelInfo.MultiKeyStatusList, 0) + var ability model.Ability + require.NoError(t, db.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestManageMultiKeysRejectsDeletingAllUsableKeys(t *testing.T) { + gin.SetMode(gin.TestMode) + tests := []struct { + name string + action string + keyIndex *int + statusList map[int]int + }{ + { + name: "delete last usable key", + action: "delete_key", + keyIndex: common.GetPointer(0), + }, + { + name: "delete disabled keys leaves only blank entries", + action: "delete_disabled_keys", + statusList: map[int]int{0: common.ChannelStatusAutoDisabled}, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + db := setupMultiKeyControllerTest(t) + channel := model.Channel{ + Name: test.name, + Key: "only-usable-key\n ", + Status: common.ChannelStatusEnabled, + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: test.statusList, + }, + } + require.NoError(t, db.Create(&channel).Error) + + body, err := common.Marshal(MultiKeyManageRequest{ + ChannelId: channel.Id, + Action: test.action, + KeyIndex: test.keyIndex, + }) + require.NoError(t, err) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/channel/multi-key", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + ctx.Set("id", 1) + ctx.Set("role", common.RoleRootUser) + + ManageMultiKeys(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Contains(t, response.Message, "不能删除") + + var persisted model.Channel + require.NoError(t, db.First(&persisted, channel.Id).Error) + assert.Equal(t, channel.Key, persisted.Key) + }) + } +} diff --git a/controller/channel_upstream_update.go b/controller/channel_upstream_update.go index 71ab0e53fafe..06c245e8e7cc 100644 --- a/controller/channel_upstream_update.go +++ b/controller/channel_upstream_update.go @@ -97,6 +97,59 @@ type upstreamModelUpdateChannelSummary struct { RemoveCount int } +var errChannelUpstreamFetchConfigChanged = errors.New("channel upstream model fetch configuration changed") + +func channelUpstreamFetchConfigFingerprint( + channel *model.Channel, + settings dto.ChannelOtherSettings, +) (string, error) { + if channel == nil { + return "", fmt.Errorf("channel is required") + } + baseURL := "" + if channel.BaseURL != nil { + baseURL = *channel.BaseURL + } + setting := "" + if channel.Setting != nil { + setting = *channel.Setting + } + headerOverride := "" + if channel.HeaderOverride != nil { + headerOverride = *channel.HeaderOverride + } + payload := struct { + Status int `json:"status"` + Type int `json:"type"` + Key string `json:"key"` + BaseURL string `json:"base_url"` + Setting string `json:"setting"` + HeaderOverride string `json:"header_override"` + IsMultiKey bool `json:"is_multi_key"` + MultiKeyMode constant.MultiKeyMode `json:"multi_key_mode"` + MultiKeyStatusList map[int]int `json:"multi_key_status_list"` + CheckEnabled bool `json:"check_enabled"` + AdvancedCustom *dto.AdvancedCustomConfig `json:"advanced_custom"` + }{ + Status: channel.Status, + Type: channel.Type, + Key: channel.Key, + BaseURL: baseURL, + Setting: setting, + HeaderOverride: headerOverride, + IsMultiKey: channel.ChannelInfo.IsMultiKey, + MultiKeyMode: channel.ChannelInfo.MultiKeyMode, + MultiKeyStatusList: channel.ChannelInfo.MultiKeyStatusList, + CheckEnabled: settings.UpstreamModelUpdateCheckEnabled, + AdvancedCustom: settings.AdvancedCustom, + } + encoded, err := common.Marshal(payload) + if err != nil { + return "", err + } + return string(encoded), nil +} + func normalizeModelNames(models []string) []string { return lo.Uniq(lo.FilterMap(models, func(model string, _ int) (string, bool) { trimmed := strings.TrimSpace(model) @@ -237,20 +290,6 @@ func collectPendingUpstreamModelChangesFromModels( return normalizeModelNames(pendingAdd), normalizeModelNames(pendingRemove) } -func collectPendingUpstreamModelChanges(channel *model.Channel, settings dto.ChannelOtherSettings) (pendingAddModels []string, pendingRemoveModels []string, err error) { - upstreamModels, err := fetchChannelUpstreamModelIDs(channel) - if err != nil { - return nil, nil, err - } - pendingAddModels, pendingRemoveModels = collectPendingUpstreamModelChangesFromModels( - channel.GetModels(), - upstreamModels, - settings.UpstreamModelUpdateIgnoredModels, - normalizeChannelModelMapping(channel), - ) - return pendingAddModels, pendingRemoveModels, nil -} - func getUpstreamModelUpdateMinCheckIntervalSeconds() int64 { interval := int64(common.GetEnvOrDefault( "CHANNEL_UPSTREAM_MODEL_UPDATE_MIN_CHECK_INTERVAL_SECONDS", @@ -460,15 +499,27 @@ func fetchAdvancedCustomUpstreamModelIDs(channel *model.Channel, baseURL string) return parseOpenAIModelIDs(body) } -func updateChannelUpstreamModelSettings(channel *model.Channel, settings dto.ChannelOtherSettings, updateModels bool) error { - channel.SetOtherSettings(settings) - updates := map[string]interface{}{ - "settings": channel.OtherSettings, +func updateChannelUpstreamModelSettings( + channel *model.Channel, + apply func(current *model.Channel, settings *dto.ChannelOtherSettings) error, +) error { + if channel == nil || channel.Id <= 0 || apply == nil { + return fmt.Errorf("invalid channel upstream model update") } - if updateModels { - updates["models"] = channel.Models + + updated, err := model.UpdateChannelAtomically(channel.Id, func(current *model.Channel) error { + settings := current.GetOtherSettings() + if err := apply(current, &settings); err != nil { + return err + } + current.SetOtherSettings(settings) + return nil + }) + if err != nil { + return err } - return model.DB.Model(&model.Channel{}).Where("id = ?", channel.Id).Updates(updates).Error + *channel = *updated + return nil } func checkAndPersistChannelUpstreamModelUpdates( @@ -485,52 +536,83 @@ func checkAndPersistChannelUpstreamModelUpdates( return false, 0, nil } } + fetchConfigFingerprint, err := channelUpstreamFetchConfigFingerprint(channel, *settings) + if err != nil { + return false, 0, err + } - pendingAddModels, pendingRemoveModels, fetchErr := collectPendingUpstreamModelChanges(channel, *settings) - settings.UpstreamModelUpdateLastCheckTime = now + upstreamModels, fetchErr := fetchChannelUpstreamModelIDs(channel) if fetchErr != nil { - if err = updateChannelUpstreamModelSettings(channel, *settings, false); err != nil { + if err = updateChannelUpstreamModelSettings(channel, func(current *model.Channel, currentSettings *dto.ChannelOtherSettings) error { + currentFingerprint, fingerprintErr := channelUpstreamFetchConfigFingerprint(current, *currentSettings) + if fingerprintErr != nil { + return fingerprintErr + } + if currentFingerprint != fetchConfigFingerprint { + *settings = *currentSettings + return errChannelUpstreamFetchConfigChanged + } + currentSettings.UpstreamModelUpdateLastCheckTime = now + *settings = *currentSettings + return nil + }); err != nil { + if errors.Is(err, errChannelUpstreamFetchConfigChanged) { + return false, 0, nil + } return false, 0, err } return false, 0, fetchErr } - if allowAutoApply && settings.UpstreamModelUpdateAutoSyncEnabled && len(pendingAddModels) > 0 { - originModels := normalizeModelNames(channel.GetModels()) - mergedModels := mergeModelNames(originModels, pendingAddModels) - if len(mergedModels) > len(originModels) { - channel.Models = strings.Join(mergedModels, ",") - autoAdded = len(mergedModels) - len(originModels) - modelsChanged = true + if err = updateChannelUpstreamModelSettings(channel, func(current *model.Channel, currentSettings *dto.ChannelOtherSettings) error { + currentFingerprint, fingerprintErr := channelUpstreamFetchConfigFingerprint(current, *currentSettings) + if fingerprintErr != nil { + return fingerprintErr } - settings.UpstreamModelUpdateLastDetectedModels = []string{} - } else { - settings.UpstreamModelUpdateLastDetectedModels = pendingAddModels - } - settings.UpstreamModelUpdateLastRemovedModels = pendingRemoveModels - - if err = updateChannelUpstreamModelSettings(channel, *settings, modelsChanged); err != nil { - return false, autoAdded, err - } - if modelsChanged { - if err = channel.UpdateAbilities(nil); err != nil { - return true, autoAdded, err + if currentFingerprint != fetchConfigFingerprint { + *settings = *currentSettings + return errChannelUpstreamFetchConfigChanged + } + pendingAddModels, pendingRemoveModels := collectPendingUpstreamModelChangesFromModels( + current.GetModels(), + upstreamModels, + currentSettings.UpstreamModelUpdateIgnoredModels, + normalizeChannelModelMapping(current), + ) + currentSettings.UpstreamModelUpdateLastCheckTime = now + if allowAutoApply && currentSettings.UpstreamModelUpdateAutoSyncEnabled && len(pendingAddModels) > 0 { + originModels := normalizeModelNames(current.GetModels()) + mergedModels := mergeModelNames(originModels, pendingAddModels) + if len(mergedModels) > len(originModels) { + current.Models = strings.Join(mergedModels, ",") + autoAdded = len(mergedModels) - len(originModels) + modelsChanged = true + } + currentSettings.UpstreamModelUpdateLastDetectedModels = []string{} + } else { + currentSettings.UpstreamModelUpdateLastDetectedModels = pendingAddModels + } + currentSettings.UpstreamModelUpdateLastRemovedModels = pendingRemoveModels + *settings = *currentSettings + return nil + }); err != nil { + if errors.Is(err, errChannelUpstreamFetchConfigChanged) { + return false, 0, nil } + return false, autoAdded, err } return modelsChanged, autoAdded, nil } func refreshChannelRuntimeCache() { - if common.MemoryCacheEnabled { - func() { - defer func() { - if r := recover(); r != nil { - common.SysLog(fmt.Sprintf("InitChannelCache panic: %v", r)) - } - }() - model.InitChannelCache() + func() { + defer func() { + if r := recover(); r != nil { + common.SysLog(fmt.Sprintf("InitChannelCache panic: %v", r)) + } }() - } + model.InitChannelCache() + }() } func shouldSendUpstreamModelUpdateNotification(now int64, changedChannels int, failedChannels int) bool { @@ -649,7 +731,7 @@ type upstreamModelUpdateSummary struct { // scheduled job calls (force=false, allowAutoApply=true); the manual "detect // all" trigger calls (force=true, allowAutoApply=false) so it always re-checks // and only stages changes for explicit review. -func runChannelUpstreamModelUpdateTaskOnce(ctx context.Context, force bool, allowAutoApply bool, report func(processed, total int)) upstreamModelUpdateSummary { +func runChannelUpstreamModelUpdateTaskOnce(ctx context.Context, force bool, allowAutoApply bool, report func(processed, total int)) (upstreamModelUpdateSummary, error) { checkedChannels := 0 failedChannels := 0 failedChannelIDs := make([]int, 0) @@ -661,6 +743,7 @@ func runChannelUpstreamModelUpdateTaskOnce(ctx context.Context, force bool, allo addModelSamples := make([]string, 0) removeModelSamples := make([]string, 0) refreshNeeded := false + var runErr error // Count the enabled channels up front so progress can be reported as a // percentage; a count error is non-fatal (progress just won't show a %). @@ -688,6 +771,7 @@ scanLoop: err := query.Find(&channels).Error if err != nil { common.SysLog(fmt.Sprintf("upstream model update task query failed: %v", err)) + runErr = err break } if len(channels) == 0 { @@ -766,6 +850,10 @@ scanLoop: if refreshNeeded { refreshChannelRuntimeCache() + service.SyncModelChannelAvailabilityAfterMutation("channel.upstream_auto_apply") + } + if runErr == nil && ctx != nil && ctx.Err() != nil { + runErr = ctx.Err() } summary := upstreamModelUpdateSummary{ @@ -796,7 +884,7 @@ scanLoop: changedChannels, failedChannels, )) - return summary + return summary, runErr } service.NotifyUpstreamModelUpdateWatchers( "上游模型巡检通知", @@ -813,7 +901,7 @@ scanLoop: ), ) } - return summary + return summary, runErr } func ApplyChannelUpstreamModelUpdates(c *gin.Context) { @@ -835,10 +923,7 @@ func ApplyChannelUpstreamModelUpdates(c *gin.Context) { common.ApiError(c, err) return } - beforeSettings := channel.GetOtherSettings() - ignoredModels := intersectModelNames(req.IgnoreModels, beforeSettings.UpstreamModelUpdateLastDetectedModels) - - addedModels, removedModels, remainingModels, remainingRemoveModels, modelsChanged, err := applyChannelUpstreamModelUpdates( + addedModels, removedModels, ignoredModels, remainingModels, remainingRemoveModels, modelsChanged, err := applyChannelUpstreamModelUpdates( channel, req.AddModels, req.IgnoreModels, @@ -856,6 +941,9 @@ func ApplyChannelUpstreamModelUpdates(c *gin.Context) { recordManageAudit(c, "channel.upstream_apply", map[string]interface{}{ "id": channel.Id, }) + if modelsChanged { + syncModelChannelAvailabilityAfterMutation("channel.upstream_apply") + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", @@ -900,6 +988,7 @@ func DetectChannelUpstreamModelUpdates(c *gin.Context) { } if modelsChanged { refreshChannelRuntimeCache() + syncModelChannelAvailabilityAfterMutation("channel.upstream_detect") } c.JSON(http.StatusOK, gin.H{ @@ -924,46 +1013,43 @@ func applyChannelUpstreamModelUpdates( ) ( addedModels []string, removedModels []string, + ignoredModels []string, remainingModels []string, remainingRemoveModels []string, modelsChanged bool, err error, ) { - settings := channel.GetOtherSettings() - pendingAddModels := normalizeModelNames(settings.UpstreamModelUpdateLastDetectedModels) - pendingRemoveModels := normalizeModelNames(settings.UpstreamModelUpdateLastRemovedModels) - addModels := intersectModelNames(addModelsInput, pendingAddModels) - ignoreModels := intersectModelNames(ignoreModelsInput, pendingAddModels) - removeModels := intersectModelNames(removeModelsInput, pendingRemoveModels) - removeModels = subtractModelNames(removeModels, addModels) - - originModels := normalizeModelNames(channel.GetModels()) - nextModels := applySelectedModelChanges(originModels, addModels, removeModels) - modelsChanged = !slices.Equal(originModels, nextModels) - if modelsChanged { - channel.Models = strings.Join(nextModels, ",") - } - - settings.UpstreamModelUpdateIgnoredModels = mergeModelNames(settings.UpstreamModelUpdateIgnoredModels, ignoreModels) - if len(addModels) > 0 { - settings.UpstreamModelUpdateIgnoredModels = subtractModelNames(settings.UpstreamModelUpdateIgnoredModels, addModels) - } - remainingModels = subtractModelNames(pendingAddModels, append(addModels, ignoreModels...)) - remainingRemoveModels = subtractModelNames(pendingRemoveModels, removeModels) - settings.UpstreamModelUpdateLastDetectedModels = remainingModels - settings.UpstreamModelUpdateLastRemovedModels = remainingRemoveModels - settings.UpstreamModelUpdateLastCheckTime = common.GetTimestamp() - - if err := updateChannelUpstreamModelSettings(channel, settings, modelsChanged); err != nil { - return nil, nil, nil, nil, false, err - } + err = updateChannelUpstreamModelSettings(channel, func(current *model.Channel, settings *dto.ChannelOtherSettings) error { + pendingAddModels := normalizeModelNames(settings.UpstreamModelUpdateLastDetectedModels) + pendingRemoveModels := normalizeModelNames(settings.UpstreamModelUpdateLastRemovedModels) + addedModels = intersectModelNames(addModelsInput, pendingAddModels) + ignoredModels = intersectModelNames(ignoreModelsInput, pendingAddModels) + removedModels = intersectModelNames(removeModelsInput, pendingRemoveModels) + removedModels = subtractModelNames(removedModels, addedModels) + + originModels := normalizeModelNames(current.GetModels()) + nextModels := applySelectedModelChanges(originModels, addedModels, removedModels) + modelsChanged = !slices.Equal(originModels, nextModels) + if modelsChanged { + current.Models = strings.Join(nextModels, ",") + } - if modelsChanged { - if err := channel.UpdateAbilities(nil); err != nil { - return addModels, removeModels, remainingModels, remainingRemoveModels, true, err + settings.UpstreamModelUpdateIgnoredModels = mergeModelNames(settings.UpstreamModelUpdateIgnoredModels, ignoredModels) + if len(addedModels) > 0 { + settings.UpstreamModelUpdateIgnoredModels = subtractModelNames(settings.UpstreamModelUpdateIgnoredModels, addedModels) } + remainingModels = subtractModelNames(pendingAddModels, append(addedModels, ignoredModels...)) + remainingRemoveModels = subtractModelNames(pendingRemoveModels, removedModels) + settings.UpstreamModelUpdateLastDetectedModels = remainingModels + settings.UpstreamModelUpdateLastRemovedModels = remainingRemoveModels + settings.UpstreamModelUpdateLastCheckTime = common.GetTimestamp() + return nil + }) + if err != nil { + return nil, nil, nil, nil, nil, false, err } - return addModels, removeModels, remainingModels, remainingRemoveModels, modelsChanged, nil + + return addedModels, removedModels, ignoredModels, remainingModels, remainingRemoveModels, modelsChanged, nil } func collectPendingApplyUpstreamModelChanges(settings dto.ChannelOtherSettings) (pendingAddModels []string, pendingRemoveModels []string) { @@ -1017,7 +1103,7 @@ func ApplyAllChannelUpstreamModelUpdates(c *gin.Context) { continue } - addedModels, removedModels, remainingModels, remainingRemoveModels, modelsChanged, err := applyChannelUpstreamModelUpdates( + addedModels, removedModels, _, remainingModels, remainingRemoveModels, modelsChanged, err := applyChannelUpstreamModelUpdates( channel, pendingAddModels, nil, @@ -1054,6 +1140,9 @@ func ApplyAllChannelUpstreamModelUpdates(c *gin.Context) { recordManageAudit(c, "channel.upstream_apply_all", map[string]interface{}{ "count": len(results), }) + if refreshNeeded { + syncModelChannelAvailabilityAfterMutation("channel.upstream_apply_all") + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/controller/channel_upstream_update_test.go b/controller/channel_upstream_update_test.go index 0a524d704bdf..0d96e6e4a550 100644 --- a/controller/channel_upstream_update_test.go +++ b/controller/channel_upstream_update_test.go @@ -6,7 +6,9 @@ import ( "net/http" "net/http/httptest" "net/url" + "sync" "testing" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -351,6 +353,256 @@ func TestFailedAdvancedCustomDetectionDoesNotStageFullRemoval(t *testing.T) { require.Equal(t, "gpt-4.1,o3", reloaded.Models) } +func TestUpdateChannelUpstreamModelSettingsUsesPersistedChannelStateForAbilities(t *testing.T) { + db := setupModelListControllerTestDB(t) + priority := int64(3) + weight := uint(7) + persisted := &model.Channel{ + Name: "disabled upstream sync channel", + Type: constant.ChannelTypeOpenAI, + Key: "secret-key", + Status: common.ChannelStatusManuallyDisabled, + Models: "old-model", + Group: "default", + Priority: &priority, + Weight: &weight, + } + require.NoError(t, db.Create(persisted).Error) + require.NoError(t, persisted.UpdateAbilities(nil)) + + stale := *persisted + stale.Status = common.ChannelStatusEnabled + + require.NoError(t, updateChannelUpstreamModelSettings(&stale, func(current *model.Channel, settings *dto.ChannelOtherSettings) error { + current.Models = "new-model" + settings.UpstreamModelUpdateLastCheckTime = 123 + return nil + })) + + var abilities []model.Ability + require.NoError(t, db.Where("channel_id = ?", persisted.Id).Find(&abilities).Error) + require.Len(t, abilities, 1) + assert.Equal(t, "new-model", abilities[0].Model) + assert.False(t, abilities[0].Enabled) + + reloaded, err := model.GetChannelById(persisted.Id, true) + require.NoError(t, err) + assert.Equal(t, common.ChannelStatusManuallyDisabled, reloaded.Status) +} + +func TestApplyChannelUpstreamModelUpdatesPreservesConcurrentChannelEdits(t *testing.T) { + db := setupModelListControllerTestDB(t) + priority := int64(3) + weight := uint(7) + persisted := &model.Channel{ + Name: "concurrently edited upstream channel", + Type: constant.ChannelTypeOpenAI, + Key: "secret-key", + Status: common.ChannelStatusEnabled, + Models: "old-model,remove-model", + Group: "default", + Priority: &priority, + Weight: &weight, + } + persisted.SetOtherSettings(dto.ChannelOtherSettings{ + UpstreamModelUpdateLastDetectedModels: []string{"new-model", "ignored-model"}, + UpstreamModelUpdateLastRemovedModels: []string{"remove-model"}, + UpstreamModelUpdateIgnoredModels: []string{"existing-ignore"}, + }) + require.NoError(t, db.Create(persisted).Error) + require.NoError(t, persisted.UpdateAbilities(nil)) + + stale := *persisted + latestSettings := persisted.GetOtherSettings() + latestSettings.AllowServiceTier = true + persisted.SetOtherSettings(latestSettings) + require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", persisted.Id).Updates(map[string]interface{}{ + "models": "old-model,manual-model,remove-model", + "settings": persisted.OtherSettings, + }).Error) + + added, removed, ignored, remaining, remainingRemoved, changed, err := applyChannelUpstreamModelUpdates( + &stale, + []string{"new-model"}, + []string{"ignored-model"}, + []string{"remove-model"}, + ) + require.NoError(t, err) + assert.True(t, changed) + assert.Equal(t, []string{"new-model"}, added) + assert.Equal(t, []string{"remove-model"}, removed) + assert.Equal(t, []string{"ignored-model"}, ignored) + assert.Empty(t, remaining) + assert.Empty(t, remainingRemoved) + assert.Equal(t, "old-model,manual-model,new-model", stale.Models) + + reloaded, err := model.GetChannelById(persisted.Id, true) + require.NoError(t, err) + assert.Equal(t, "old-model,manual-model,new-model", reloaded.Models) + reloadedSettings := reloaded.GetOtherSettings() + assert.True(t, reloadedSettings.AllowServiceTier) + assert.Equal(t, []string{"existing-ignore", "ignored-model"}, reloadedSettings.UpstreamModelUpdateIgnoredModels) + assert.Empty(t, reloadedSettings.UpstreamModelUpdateLastDetectedModels) + assert.Empty(t, reloadedSettings.UpstreamModelUpdateLastRemovedModels) + + var abilities []model.Ability + require.NoError(t, db.Where("channel_id = ?", persisted.Id).Order("model ASC").Find(&abilities).Error) + require.Len(t, abilities, 3) + assert.Equal(t, []string{"manual-model", "new-model", "old-model"}, []string{ + abilities[0].Model, + abilities[1].Model, + abilities[2].Model, + }) +} + +func TestCheckAndPersistChannelUpstreamModelUpdatesPreservesConcurrentChannelEdits(t *testing.T) { + db := setupModelListControllerTestDB(t) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = w.Write([]byte(`{"data":[{"id":"old-model"},{"id":"manual-model"},{"id":"new-model"}]}`)) + })) + t.Cleanup(server.Close) + + channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/v1/models", nil) + channel.Name = "concurrently edited detection channel" + channel.Status = common.ChannelStatusEnabled + channel.Models = "old-model" + channel.Group = "default" + settings := channel.GetOtherSettings() + settings.UpstreamModelUpdateCheckEnabled = true + settings.UpstreamModelUpdateAutoSyncEnabled = true + channel.SetOtherSettings(settings) + require.NoError(t, db.Create(channel).Error) + require.NoError(t, channel.UpdateAbilities(nil)) + + stale := *channel + staleSettings := stale.GetOtherSettings() + latestSettings := channel.GetOtherSettings() + latestSettings.AllowServiceTier = true + latestSettings.UpstreamModelUpdateIgnoredModels = []string{"preserved-ignore"} + channel.SetOtherSettings(latestSettings) + require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "models": "old-model,manual-model", + "settings": channel.OtherSettings, + }).Error) + + modelsChanged, autoAdded, err := checkAndPersistChannelUpstreamModelUpdates(&stale, &staleSettings, true, true) + require.NoError(t, err) + assert.True(t, modelsChanged) + assert.Equal(t, 1, autoAdded) + assert.Equal(t, "old-model,manual-model,new-model", stale.Models) + assert.True(t, staleSettings.AllowServiceTier) + assert.Equal(t, []string{"preserved-ignore"}, staleSettings.UpstreamModelUpdateIgnoredModels) + + reloaded, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, "old-model,manual-model,new-model", reloaded.Models) + assert.True(t, reloaded.GetOtherSettings().AllowServiceTier) +} + +func TestCheckAndPersistChannelUpstreamModelUpdatesSkipsStaleFetchConfig(t *testing.T) { + db := setupModelListControllerTestDB(t) + requestStarted := make(chan struct{}) + releaseResponse := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(releaseResponse) }) }) + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + close(requestStarted) + <-releaseResponse + _, _ = w.Write([]byte(`{"data":[{"id":"old-model"},{"id":"stale-model"}]}`)) + })) + t.Cleanup(server.Close) + + channel := newAdvancedCustomModelListChannel(server.URL, "secret-key", "/v1/models", nil) + channel.Name = "stale fetch configuration channel" + channel.Status = common.ChannelStatusEnabled + channel.Models = "old-model" + channel.Group = "default" + settings := channel.GetOtherSettings() + settings.UpstreamModelUpdateCheckEnabled = true + settings.UpstreamModelUpdateAutoSyncEnabled = true + channel.SetOtherSettings(settings) + require.NoError(t, db.Create(channel).Error) + require.NoError(t, channel.UpdateAbilities(nil)) + + type checkResult struct { + modelsChanged bool + autoAdded int + err error + } + resultCh := make(chan checkResult, 1) + go func() { + modelsChanged, autoAdded, err := checkAndPersistChannelUpstreamModelUpdates(channel, &settings, true, true) + resultCh <- checkResult{modelsChanged: modelsChanged, autoAdded: autoAdded, err: err} + }() + + select { + case <-requestStarted: + case <-time.After(5 * time.Second): + require.FailNow(t, "timed out waiting for upstream model request") + } + newBaseURL := "http://new-upstream.invalid" + require.NoError(t, db.Model(&model.Channel{}).Where("id = ?", channel.Id).Update("base_url", newBaseURL).Error) + releaseOnce.Do(func() { close(releaseResponse) }) + + var result checkResult + select { + case result = <-resultCh: + case <-time.After(5 * time.Second): + require.FailNow(t, "timed out waiting for stale fetch handling") + } + require.NoError(t, result.err) + assert.False(t, result.modelsChanged) + assert.Zero(t, result.autoAdded) + + reloaded, err := model.GetChannelById(channel.Id, true) + require.NoError(t, err) + assert.Equal(t, newBaseURL, reloaded.GetBaseURL()) + assert.Equal(t, "old-model", reloaded.Models) + assert.Empty(t, reloaded.GetOtherSettings().UpstreamModelUpdateLastDetectedModels) +} + +func TestRefreshChannelRuntimeCacheInvalidatesPricingWithoutMemoryCache(t *testing.T) { + db := setupModelListControllerTestDB(t) + originalMemoryCacheEnabled := common.MemoryCacheEnabled + common.MemoryCacheEnabled = false + model.InvalidatePricingCache() + t.Cleanup(func() { + common.MemoryCacheEnabled = originalMemoryCacheEnabled + model.InvalidatePricingCache() + }) + + channel := &model.Channel{ + Name: "pricing invalidation channel", + Type: constant.ChannelTypeOpenAI, + Key: "secret-key", + Status: common.ChannelStatusEnabled, + Models: "old-pricing-model", + Group: "default", + } + require.NoError(t, db.Create(channel).Error) + require.NoError(t, channel.UpdateAbilities(nil)) + + containsModel := func(pricing []model.Pricing, name string) bool { + for _, item := range pricing { + if item.ModelName == name { + return true + } + } + return false + } + require.True(t, containsModel(model.GetPricing(), "old-pricing-model")) + require.NoError(t, db.Model(&model.Ability{}). + Where("channel_id = ?", channel.Id). + Update("model", "new-pricing-model").Error) + require.True(t, containsModel(model.GetPricing(), "old-pricing-model")) + + refreshChannelRuntimeCache() + + pricing := model.GetPricing() + assert.True(t, containsModel(pricing, "new-pricing-model")) + assert.False(t, containsModel(pricing, "old-pricing-model")) +} + func TestFetchModelsUsesSharedChannelFetchBehavior(t *testing.T) { server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.URL.Path != "/v1/models" { diff --git a/controller/model_list_test.go b/controller/model_list_test.go index 812207b8fd44..9b6992776963 100644 --- a/controller/model_list_test.go +++ b/controller/model_list_test.go @@ -49,7 +49,14 @@ func setupModelListControllerTestDB(t *testing.T) *gorm.DB { model.DB = db model.LOG_DB = db - require.NoError(t, db.AutoMigrate(&model.User{}, &model.Channel{}, &model.Ability{}, &model.Model{}, &model.Vendor{})) + require.NoError(t, db.AutoMigrate( + &model.User{}, + &model.Channel{}, + &model.Ability{}, + &model.Model{}, + &model.Vendor{}, + &model.Option{}, + )) t.Cleanup(func() { sqlDB, err := db.DB() diff --git a/controller/model_meta.go b/controller/model_meta.go index c3d9954677e7..992168d29c4f 100644 --- a/controller/model_meta.go +++ b/controller/model_meta.go @@ -2,15 +2,19 @@ package controller import ( "encoding/json" + "fmt" "sort" "strconv" "strings" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" + "gorm.io/gorm" ) // GetAllModelsMeta 获取模型列表（分页） @@ -110,10 +114,30 @@ func CreateModelMeta(c *gin.Context) { common.ApiError(c, err) return } - model.RefreshPricing() + res := service.SyncModelChannelAvailabilityAfterMutation("model.create") + if !res.PricingRefreshed { + model.RefreshPricing() + } + if err := reloadModelMetaAfterMutation(&m); err != nil { + common.ApiError(c, err) + return + } common.ApiSuccess(c, &m) } +func reloadModelMetaAfterMutation(current *model.Model) error { + if current == nil || current.Id == 0 { + return nil + } + var persisted model.Model + if err := model.DB.First(&persisted, current.Id).Error; err != nil { + common.SysError(fmt.Sprintf("failed to reload model metadata after mutation: id=%d err=%v", current.Id, err)) + return fmt.Errorf("reload model metadata after mutation: %w", err) + } + *current = persisted + return nil +} + // UpdateModelMeta 更新模型 func UpdateModelMeta(c *gin.Context) { statusOnly := c.Query("status_only") == "true" @@ -130,27 +154,134 @@ func UpdateModelMeta(c *gin.Context) { if statusOnly { // 只更新状态，防止误清空其他字段 - if err := model.DB.Model(&model.Model{}).Where("id = ?", m.Id).Update("status", m.Status).Error; err != nil { - common.ApiError(c, err) + updateResult := model.DB.Model(&model.Model{}).Where("id = ?", m.Id).Updates(map[string]interface{}{ + "status": m.Status, + "auto_disabled_by_rule": false, + "updated_time": common.GetTimestamp(), + }) + if updateResult.Error != nil { + common.ApiError(c, updateResult.Error) return } - } else { - // 名称冲突检查 - if dup, err := model.IsModelNameDuplicated(m.Id, m.ModelName); err != nil { + if updateResult.RowsAffected == 0 { + var count int64 + if err := model.DB.Model(&model.Model{}).Where("id = ?", m.Id).Count(&count).Error; err != nil { + common.ApiError(c, err) + return + } + if count == 0 { + common.ApiError(c, gorm.ErrRecordNotFound) + return + } + } + // Re-evaluate immediately so auto-disable can correct a manual enable without channels. + res := service.SyncModelChannelAvailabilityAfterMutation("model.status_update") + if !res.PricingRefreshed { + model.RefreshPricing() + } + if err := reloadModelMetaAfterMutation(&m); err != nil { common.ApiError(c, err) return - } else if dup { - common.ApiErrorMsg(c, "模型名称已存在") - return } + common.ApiSuccess(c, &m) + return + } - if err := m.Update(); err != nil { - common.ApiError(c, err) + // 名称冲突检查 + if dup, err := model.IsModelNameDuplicated(m.Id, m.ModelName); err != nil { + common.ApiError(c, err) + return + } else if dup { + common.ApiErrorMsg(c, "模型名称已存在") + return + } + + if err := m.Update(); err != nil { + common.ApiError(c, err) + return + } + res := service.SyncModelChannelAvailabilityAfterMutation("model.update") + if !res.PricingRefreshed { + model.RefreshPricing() + } + if err := reloadModelMetaAfterMutation(&m); err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, &m) +} + +type batchUpdateModelStatusRequest struct { + Ids []int `json:"ids"` + Status int `json:"status"` +} + +// BatchUpdateModelStatus updates selected models and reconciles channel +// availability once for the whole batch. +func BatchUpdateModelStatus(c *gin.Context) { + req := batchUpdateModelStatusRequest{} + if err := c.ShouldBindJSON(&req); err != nil || len(req.Ids) == 0 || (req.Status != 0 && req.Status != 1) { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) + return + } + if len(req.Ids) > 100 { + common.ApiErrorI18n(c, i18n.MsgBatchTooMany, map[string]any{"Max": 100}) + return + } + + ids := make([]int, 0, len(req.Ids)) + seen := make(map[int]struct{}, len(req.Ids)) + for _, id := range req.Ids { + if id <= 0 { + common.ApiErrorI18n(c, i18n.MsgInvalidParams) return } + if _, ok := seen[id]; ok { + continue + } + seen[id] = struct{}{} + ids = append(ids, id) } - model.RefreshPricing() - common.ApiSuccess(c, &m) + + if err := model.DB.Model(&model.Model{}). + Where("id IN ?", ids). + Updates(map[string]interface{}{ + "status": req.Status, + "auto_disabled_by_rule": false, + "updated_time": common.GetTimestamp(), + }).Error; err != nil { + common.ApiError(c, err) + return + } + + res := service.SyncModelChannelAvailabilityAfterMutation("model.status_update_batch") + if !res.PricingRefreshed { + model.RefreshPricing() + } + + var persisted []model.Model + if err := model.DB.Select("id", "status").Where("id IN ?", ids).Find(&persisted).Error; err != nil { + common.ApiError(c, err) + return + } + statusByID := make(map[int]int, len(persisted)) + for i := range persisted { + statusByID[persisted[i].Id] = persisted[i].Status + } + updated := 0 + failedIds := make([]int, 0) + for _, id := range ids { + if status, ok := statusByID[id]; ok && status == req.Status { + updated++ + continue + } + failedIds = append(failedIds, id) + } + + common.ApiSuccess(c, gin.H{ + "updated": updated, + "failed_ids": failedIds, + }) } // DeleteModelMeta 删除模型 @@ -165,7 +296,10 @@ func DeleteModelMeta(c *gin.Context) { common.ApiError(c, err) return } - model.RefreshPricing() + res := service.SyncModelChannelAvailabilityAfterMutation("model.delete") + if !res.PricingRefreshed { + model.RefreshPricing() + } common.ApiSuccess(c, nil) } @@ -337,3 +471,33 @@ func enrichModels(models []*model.Model) { mm.MatchedCount = len(names) } } + +// BatchDisableModelsNoChannels 批量禁用无可用渠道的模型 +func BatchDisableModelsNoChannels(c *gin.Context) { + result, err := service.ManualDisableModelsWithoutChannels() + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "disabled": result.Disabled, + "enabled": result.Enabled, + "skipped": result.Skipped, + "reason": result.Reason, + }) +} + +// BatchEnableModelsWithChannels 批量启用：仅恢复被渠道可用性规则自动禁用、且现已有可用渠道的模型 +func BatchEnableModelsWithChannels(c *gin.Context) { + result, err := service.ManualEnableModelsWithChannels() + if err != nil { + common.ApiError(c, err) + return + } + common.ApiSuccess(c, gin.H{ + "disabled": result.Disabled, + "enabled": result.Enabled, + "skipped": result.Skipped, + "reason": result.Reason, + }) +} diff --git a/controller/model_meta_availability_test.go b/controller/model_meta_availability_test.go new file mode 100644 index 000000000000..47ad8cf72c35 --- /dev/null +++ b/controller/model_meta_availability_test.go @@ -0,0 +1,270 @@ +package controller + +import ( + "bytes" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/gin-gonic/gin" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func setupModelMetaStatusControllerTestDB(t *testing.T) *gorm.DB { + t.Helper() + + dsn := fmt.Sprintf("file:%s?mode=memory&cache=shared", strings.ReplaceAll(t.Name(), "/", "_")) + db, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate( + &model.Option{}, + &model.Model{}, + &model.Ability{}, + &model.Channel{}, + &model.Vendor{}, + )) + + originalDB := model.DB + originalDatabaseType := common.MainDatabaseType() + originalDisable := common.AutomaticDisableModelEnabled + originalEnable := common.AutomaticEnableModelEnabled + model.DB = db + common.SetMainDatabaseType(common.DatabaseTypeSQLite) + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = false + model.InvalidatePricingCache() + t.Cleanup(func() { + model.InvalidatePricingCache() + model.DB = originalDB + common.SetMainDatabaseType(originalDatabaseType) + common.AutomaticDisableModelEnabled = originalDisable + common.AutomaticEnableModelEnabled = originalEnable + sqlDB, dbErr := db.DB() + if dbErr == nil { + _ = sqlDB.Close() + } + }) + return db +} + +func TestCreateModelMetaReturnsReconciledAvailability(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate( + &model.Option{}, + &model.Model{}, + &model.Ability{}, + &model.Channel{}, + )) + require.NoError(t, db.Create(&[]model.Option{ + {Key: "AutomaticDisableModelEnabled", Value: "true"}, + {Key: "AutomaticEnableModelEnabled", Value: "false"}, + }).Error) + + originalDB := model.DB + originalDatabaseType := common.MainDatabaseType() + originalDisable := common.AutomaticDisableModelEnabled + originalEnable := common.AutomaticEnableModelEnabled + model.DB = db + common.SetMainDatabaseType(common.DatabaseTypeSQLite) + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = false + t.Cleanup(func() { + model.DB = originalDB + common.SetMainDatabaseType(originalDatabaseType) + common.AutomaticDisableModelEnabled = originalDisable + common.AutomaticEnableModelEnabled = originalEnable + }) + + gin.SetMode(gin.TestMode) + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/models", bytes.NewBufferString( + `{"model_name":"gpt-4","status":1,"sync_official":1,"auto_disabled_by_rule":true}`, + )) + ctx.Request.Header.Set("Content-Type", "application/json") + + CreateModelMeta(ctx) + + var response struct { + Success bool `json:"success"` + Data model.Model `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.True(t, response.Success) + assert.Equal(t, 0, response.Data.Status) + assert.True(t, response.Data.AutoDisabledByRule) +} + +func TestBatchModelAvailabilityEndpointsReturnReconcileErrors(t *testing.T) { + brokenDB, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + + originalDB := model.DB + originalDatabaseType := common.MainDatabaseType() + model.DB = brokenDB + common.SetMainDatabaseType(common.DatabaseTypeSQLite) + t.Cleanup(func() { + model.DB = originalDB + common.SetMainDatabaseType(originalDatabaseType) + }) + + gin.SetMode(gin.TestMode) + tests := []struct { + name string + handler gin.HandlerFunc + }{ + {name: "disable models without channels", handler: BatchDisableModelsNoChannels}, + {name: "enable models with channels", handler: BatchEnableModelsWithChannels}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + + test.handler(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Contains(t, response.Message, "reconcile model channel availability") + }) + } +} + +func TestUpdateModelMetaStatusOnlyReturnsNotFoundForMissingModel(t *testing.T) { + setupModelMetaStatusControllerTestDB(t) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/models?status_only=true", bytes.NewBufferString( + `{"id":404,"status":1}`, + )) + ctx.Request.Header.Set("Content-Type", "application/json") + + UpdateModelMeta(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.Equal(t, gorm.ErrRecordNotFound.Error(), response.Message) +} + +func TestUpdateModelMetaStatusOnlyReturnsReloadError(t *testing.T) { + db := setupModelMetaStatusControllerTestDB(t) + entry := &model.Model{ModelName: "reload-error-model", Status: 0, SyncOfficial: 1} + require.NoError(t, entry.Insert()) + + reloadErr := errors.New("forced model reload failure") + updateCompleted := false + require.NoError(t, db.Callback().Update().After("gorm:update").Register("test:mark_model_status_update", func(tx *gorm.DB) { + if tx.Statement.Table == "models" { + updateCompleted = true + } + })) + require.NoError(t, db.Callback().Query().Before("gorm:query").Register("test:fail_model_reload", func(tx *gorm.DB) { + if !updateCompleted || tx.Statement.Table != "models" { + return + } + if _, ok := tx.Statement.Dest.(*model.Model); ok { + tx.AddError(reloadErr) + } + })) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPut, "/api/models?status_only=true", bytes.NewBufferString(fmt.Sprintf( + `{"id":%d,"status":1}`, + entry.Id, + ))) + ctx.Request.Header.Set("Content-Type", "application/json") + + UpdateModelMeta(ctx) + + var response struct { + Success bool `json:"success"` + Message string `json:"message"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) + assert.ErrorContains(t, errors.New(response.Message), reloadErr.Error()) +} + +func TestBatchUpdateModelStatusReturnsFinalStatuses(t *testing.T) { + db := setupModelMetaStatusControllerTestDB(t) + require.NoError(t, db.Create(&model.Option{Key: "AutomaticDisableModelEnabled", Value: "true"}).Error) + + first := &model.Model{ModelName: "batch-status-first", Status: 0, SyncOfficial: 1} + second := &model.Model{ModelName: "batch-status-second", Status: 0, SyncOfficial: 1} + require.NoError(t, first.Insert()) + require.NoError(t, second.Insert()) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/models/batch_status", bytes.NewBufferString(fmt.Sprintf( + `{"ids":[%d,%d,404],"status":1}`, + first.Id, + second.Id, + ))) + ctx.Request.Header.Set("Content-Type", "application/json") + + BatchUpdateModelStatus(ctx) + + var response struct { + Success bool `json:"success"` + Data struct { + Updated int `json:"updated"` + FailedIDs []int `json:"failed_ids"` + } `json:"data"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.True(t, response.Success) + assert.Zero(t, response.Data.Updated) + assert.ElementsMatch(t, []int{first.Id, second.Id, 404}, response.Data.FailedIDs) + + var persisted []model.Model + require.NoError(t, db.Where("id IN ?", []int{first.Id, second.Id}).Order("id ASC").Find(&persisted).Error) + require.Len(t, persisted, 2) + assert.Equal(t, 0, persisted[0].Status) + assert.True(t, persisted[0].AutoDisabledByRule) + assert.Equal(t, 0, persisted[1].Status) + assert.True(t, persisted[1].AutoDisabledByRule) +} + +func TestBatchUpdateModelStatusRejectsMoreThanHundredIDs(t *testing.T) { + setupModelMetaStatusControllerTestDB(t) + ids := make([]int, 101) + for i := range ids { + ids[i] = i + 1 + } + body, err := common.Marshal(gin.H{"ids": ids, "status": 0}) + require.NoError(t, err) + + recorder := httptest.NewRecorder() + ctx, _ := gin.CreateTestContext(recorder) + ctx.Request = httptest.NewRequest(http.MethodPost, "/api/models/batch_status", bytes.NewReader(body)) + ctx.Request.Header.Set("Content-Type", "application/json") + + BatchUpdateModelStatus(ctx) + + var response struct { + Success bool `json:"success"` + } + require.NoError(t, common.Unmarshal(recorder.Body.Bytes(), &response)) + assert.False(t, response.Success) +} diff --git a/controller/model_sync.go b/controller/model_sync.go index f254dc88ee5e..906073b393d6 100644 --- a/controller/model_sync.go +++ b/controller/model_sync.go @@ -15,6 +15,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" "github.com/gin-gonic/gin" "gorm.io/gorm" @@ -434,8 +435,13 @@ func SyncUpstreamModels(c *gin.Context) { needUpdate = true } if containsField(ow.Fields, "status") { - local.Status = chooseStatus(up.Status, local.Status) - needUpdate = true + newStatus := chooseStatus(up.Status, local.Status) + // Only clear auto-disable marker when status actually changes. + if newStatus != local.Status { + local.Status = newStatus + local.AutoDisabledByRule = false + needUpdate = true + } } if !needUpdate { return nil @@ -450,6 +456,14 @@ func SyncUpstreamModels(c *gin.Context) { } } + if createdModels > 0 || updatedModels > 0 { + res := service.SyncModelChannelAvailabilityAfterMutation("model.sync_upstream") + // Rebuild pricing for new/updated metadata unless sync already did. + if !res.PricingRefreshed { + model.RefreshPricing() + } + } + c.JSON(http.StatusOK, gin.H{ "success": true, "data": gin.H{ @@ -486,13 +500,12 @@ func coalesce(a, b string) string { } func chooseStatus(primary, fallback int) int { - if primary == 0 && fallback != 0 { - return fallback - } + // 0 is a legitimate disabled status. Prefer non-zero primary, otherwise fallback + // (which may also be 0). Call sites that need a default enabled status pass 1. if primary != 0 { return primary } - return 1 + return fallback } // SyncUpstreamPreview 预览上游与本地的差异（仅用于弹窗选择） diff --git a/controller/option.go b/controller/option.go index 940bb3069023..af9f3000da91 100644 --- a/controller/option.go +++ b/controller/option.go @@ -9,6 +9,7 @@ import ( "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/i18n" "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" "github.com/QuantumNous/new-api/setting" "github.com/QuantumNous/new-api/setting/console_setting" "github.com/QuantumNous/new-api/setting/model_setting" @@ -372,6 +373,10 @@ func UpdateOption(c *gin.Context) { recordManageAudit(c, "option.update", map[string]interface{}{ "key": option.Key, }) + if err := service.MaybeSyncModelChannelAvailabilityAfterOptionChange(option.Key, option.Value.(string)); err != nil { + common.ApiError(c, err) + return + } c.JSON(http.StatusOK, gin.H{ "success": true, "message": "", diff --git a/controller/system_task_handlers.go b/controller/system_task_handlers.go index c31059d148da..07ecf87d4fa0 100644 --- a/controller/system_task_handlers.go +++ b/controller/system_task_handlers.go @@ -107,7 +107,11 @@ func (modelUpdateHandler) Run(ctx context.Context, task *model.SystemTask, runne finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, nil, err) return } - summary := runChannelUpstreamModelUpdateTaskOnce(ctx, payload.Manual, !payload.Manual, service.NewSystemTaskProgressReporter(task, runnerID)) + summary, err := runChannelUpstreamModelUpdateTaskOnce(ctx, payload.Manual, !payload.Manual, service.NewSystemTaskProgressReporter(task, runnerID)) + if err != nil { + finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusFailed, summary, err) + return + } finishSystemTaskHandler(task, runnerID, model.SystemTaskStatusSucceeded, summary, nil) } diff --git a/main.go b/main.go index 742d15515876..c9b1169a7708 100644 --- a/main.go +++ b/main.go @@ -100,6 +100,10 @@ func main() { go model.SyncChannelCache(common.SyncFrequency) } + if err := service.CalibrateModelChannelAvailabilityAtStartup(); err != nil { + common.FatalLog("failed to calibrate model channel availability at startup: " + err.Error()) + return + } // Warm pricing after channel cache initialization so Advanced Custom // endpoint inference can read cached route settings on first request. diff --git a/middleware/audit.go b/middleware/audit.go index 851e48358a97..8e58048b7c72 100644 --- a/middleware/audit.go +++ b/middleware/audit.go @@ -77,10 +77,13 @@ var auditRouteActions = map[string]string{ "DELETE /api/vendors/:id": "vendor.delete", // 模型元数据 - "POST /api/models/": "model.create", - "PUT /api/models/": "model.update", - "DELETE /api/models/:id": "model.delete", - "POST /api/models/sync_upstream": "model.sync_upstream", + "POST /api/models/": "model.create", + "PUT /api/models/": "model.update", + "DELETE /api/models/:id": "model.delete", + "POST /api/models/sync_upstream": "model.sync_upstream", + "POST /api/models/batch_status": "model.batch_status_update", + "POST /api/models/batch_disable_no_channels": "model.batch_disable_no_channels", + "POST /api/models/batch_enable_with_channels": "model.batch_enable_with_channels", // 部署 "POST /api/deployments/": "deployment.create", diff --git a/model/ability.go b/model/ability.go index d950a6adbfc4..c8824c6aec86 100644 --- a/model/ability.go +++ b/model/ability.go @@ -194,6 +194,7 @@ func filterAbilitiesByRequestPathAndModel(abilities []Ability, requestPath strin } func (channel *Channel) AddAbilities(tx *gorm.DB) error { + abilityEnabled := channelAbilitiesEnabled(channel) models_ := strings.Split(channel.Models, ",") groups_ := strings.Split(channel.Group, ",") abilitySet := make(map[string]struct{}) @@ -209,7 +210,7 @@ func (channel *Channel) AddAbilities(tx *gorm.DB) error { Group: group, Model: model, ChannelId: channel.Id, - Enabled: channel.Status == common.ChannelStatusEnabled, + Enabled: abilityEnabled, Priority: channel.Priority, Weight: uint(channel.GetWeight()), Tag: channel.Tag, @@ -266,6 +267,7 @@ func (channel *Channel) UpdateAbilities(tx *gorm.DB) error { } // Then add new abilities + abilityEnabled := channelAbilitiesEnabled(channel) models_ := strings.Split(channel.Models, ",") groups_ := strings.Split(channel.Group, ",") abilitySet := make(map[string]struct{}) @@ -281,7 +283,7 @@ func (channel *Channel) UpdateAbilities(tx *gorm.DB) error { Group: group, Model: model, ChannelId: channel.Id, - Enabled: channel.Status == common.ChannelStatusEnabled, + Enabled: abilityEnabled, Priority: channel.Priority, Weight: uint(channel.GetWeight()), Tag: channel.Tag, @@ -341,50 +343,29 @@ func FixAbility() (int, int, error) { } defer fixLock.Unlock() - // truncate abilities table - if common.UsingMainDatabase(common.DatabaseTypeSQLite) { - err := DB.Exec("DELETE FROM abilities").Error - if err != nil { - common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error())) - return 0, 0, err - } - } else { - err := DB.Exec("TRUNCATE TABLE abilities").Error - if err != nil { - common.SysLog(fmt.Sprintf("Truncate abilities failed: %s", err.Error())) - return 0, 0, err - } - } - var channels []*Channel - // Find all channels - err := DB.Model(&Channel{}).Find(&channels).Error - if err != nil { - return 0, 0, err - } - if len(channels) == 0 { - return 0, 0, nil - } successCount := 0 failCount := 0 - for _, chunk := range lo.Chunk(channels, 50) { - ids := lo.Map(chunk, func(c *Channel, _ int) int { return c.Id }) - // Delete all abilities of this channel - err = DB.Where("channel_id IN ?", ids).Delete(&Ability{}).Error - if err != nil { + err := DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Session(&gorm.Session{AllowGlobalUpdate: true}).Delete(&Ability{}).Error; err != nil { common.SysLog(fmt.Sprintf("Delete abilities failed: %s", err.Error())) - failCount += len(chunk) - continue + return err } - // Then add new abilities - for _, channel := range chunk { - err = channel.AddAbilities(nil) - if err != nil { + var channels []*Channel + if err := tx.Model(&Channel{}).Find(&channels).Error; err != nil { + return err + } + for _, channel := range channels { + if err := channel.AddAbilities(tx); err != nil { common.SysLog(fmt.Sprintf("Add abilities for channel %d failed: %s", channel.Id, err.Error())) failCount++ - } else { - successCount++ + return err } + successCount++ } + return nil + }) + if err != nil { + return 0, failCount, err } InitChannelCache() return successCount, failCount, nil diff --git a/model/channel.go b/model/channel.go index 0f8cdb101ec8..7a2747627a9a 100644 --- a/model/channel.go +++ b/model/channel.go @@ -1,6 +1,7 @@ package model import ( + "bytes" "database/sql/driver" "encoding/json" "errors" @@ -196,6 +197,27 @@ func (channel *Channel) GetKeys() []string { return keys } +// IsUsableChannelKey reports whether a stored key contains a non-blank +// credential. Legacy JSON-array storage may expose string entries as raw JSON, +// so decode those before checking for an empty value. +func IsUsableChannelKey(key string) bool { + trimmed := strings.TrimSpace(key) + if trimmed == "" { + return false + } + if strings.HasPrefix(trimmed, `"`) { + var decoded string + if err := common.Unmarshal([]byte(trimmed), &decoded); err == nil { + return strings.TrimSpace(decoded) != "" + } + } + return true +} + +func (channel *Channel) HasEnabledMultiKey() bool { + return hasEnabledMultiKey(channel.GetKeys(), channel.ChannelInfo.MultiKeyStatusList) +} + func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) { // If not in multi-key mode, return the original key string directly. if !channel.ChannelInfo.IsMultiKey { @@ -227,7 +249,10 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) { // Collect indexes of enabled keys enabledIdx := make([]int, 0, len(keys)) - for i := range keys { + for i, key := range keys { + if !IsUsableChannelKey(key) { + continue + } if getStatus(i) == common.ChannelStatusEnabled { enabledIdx = append(enabledIdx, i) } @@ -268,7 +293,7 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) { } for i := 0; i < len(keys); i++ { idx := (start + i) % len(keys) - if getStatus(idx) == common.ChannelStatusEnabled { + if IsUsableChannelKey(keys[idx]) && getStatus(idx) == common.ChannelStatusEnabled { // update polling index for next call (point to the next position) channel.ChannelInfo.MultiKeyPollingIndex = (idx + 1) % len(keys) return keys[idx], idx, nil @@ -316,7 +341,7 @@ func (channel *Channel) GetOtherInfo() map[string]interface{} { } func (channel *Channel) SetOtherInfo(otherInfo map[string]interface{}) { - otherInfoBytes, err := json.Marshal(otherInfo) + otherInfoBytes, err := common.Marshal(otherInfo) if err != nil { common.SysLog(fmt.Sprintf("failed to marshal other info: channel_id=%d, tag=%s, name=%s, error=%v", channel.Id, channel.GetTag(), channel.Name, err)) return @@ -350,6 +375,10 @@ func (channel *Channel) Save() error { // Keeping this allowlist here prevents a stale channel snapshot from // overwriting credentials, accounting counters, or channel configuration. func (channel *Channel) saveStatusState() error { + return channel.saveStatusStateWithDB(DB) +} + +func (channel *Channel) saveStatusStateWithDB(db *gorm.DB) error { if channel.Id == 0 { return errors.New("channel ID is 0") } @@ -360,7 +389,7 @@ func (channel *Channel) saveStatusState() error { if channel.ChannelInfo.IsMultiKey { updates["channel_info"] = channel.ChannelInfo } - return DB.Model(&Channel{}).Where("id = ?", channel.Id).Updates(updates).Error + return db.Model(&Channel{}).Where("id = ?", channel.Id).Updates(updates).Error } func GetAllChannels(startIdx int, num int, selectAll bool, idSort bool, sortOptions ...ChannelSortOptions) ([]*Channel, error) { @@ -437,6 +466,9 @@ func BatchInsertChannels(channels []Channel) error { if len(channels) == 0 { return nil } + for i := range channels { + channels[i].normalizeMultiKeyAvailability() + } tx := DB.Begin() if tx.Error != nil { return tx.Error @@ -530,62 +562,56 @@ func (channel *Channel) GetStatusCodeMapping() string { } func (channel *Channel) Insert() error { - var err error - err = DB.Create(channel).Error - if err != nil { - return err - } - err = channel.AddAbilities(nil) - return err + channel.normalizeMultiKeyAvailability() + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Create(channel).Error; err != nil { + return err + } + return channel.AddAbilities(tx) + }) } 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 + return DB.Transaction(func(tx *gorm.DB) error { + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Channel{}). + Where("id = ?", channel.Id). + UpdateColumn("status", gorm.Expr("status")).Error; err != nil { + return err } } - // 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 this is a multi-key channel, recalculate MultiKeySize based on the + // current key list while the channel row is locked. + if channel.ChannelInfo.IsMultiKey { + var existing Channel + if err := lockForUpdate(tx). + Select("id", "key", "status", "channel_info", "other_info"). + Where("id = ?", channel.Id). + First(&existing).Error; err != nil { + return err } - if len(keys) == 0 { // fallback to newline split - keys = strings.Split(strings.Trim(keyStr, "\n"), "\n") + if channel.Key == "" { + channel.Key = existing.Key + } else { + channel.Keys = nil + remapMultiKeyStateByKey(channel, &existing) } - } - 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) - } + if channel.Status == 0 { + channel.Status = existing.Status + } + if channel.OtherInfo == "" { + channel.OtherInfo = existing.OtherInfo } + channel.normalizeMultiKeyAvailability() } - } - var err error - err = DB.Model(channel).Updates(channel).Error - if err != nil { - return err - } - DB.Model(channel).First(channel, "id = ?", channel.Id) - err = channel.UpdateAbilities(nil) - return err + if err := tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(channel).Error; err != nil { + return err + } + if err := tx.First(channel, "id = ?", channel.Id).Error; err != nil { + return err + } + return channel.UpdateAbilities(tx) + }) } func (channel *Channel) UpdateResponseTime(responseTime int64) { @@ -609,13 +635,14 @@ func (channel *Channel) UpdateBalance(balance float64) { } func (channel *Channel) Delete() error { - var err error - err = DB.Delete(channel).Error - if err != nil { - return err - } - err = channel.DeleteAbilities() - return err + return DB.Transaction(func(tx *gorm.DB) error { + // Keep the lock order consistent with channel updates: channel first, + // derivative abilities second. + if err := tx.Delete(channel).Error; err != nil { + return err + } + return tx.Where("channel_id = ?", channel.Id).Delete(&Ability{}).Error + }) } var channelStatusLock sync.Mutex @@ -657,7 +684,12 @@ func CleanupChannelPollingLocks() { func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason string) { keys := channel.GetKeys() if len(keys) == 0 { - channel.Status = status + if usingKey != "" || status == common.ChannelStatusEnabled { + return + } + if channel.Status != common.ChannelStatusManuallyDisabled || status == common.ChannelStatusManuallyDisabled { + channel.Status = status + } } else { keyIndex := -1 for i, key := range keys { @@ -671,6 +703,9 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason common.SysLog(fmt.Sprintf("failed to update multi-key status: channel_id=%d, using key not found", channel.Id)) return } + if status == common.ChannelStatusEnabled && !channel.HasEnabledMultiKey() { + return + } channel.Status = status info := channel.GetOtherInfo() info["status_reason"] = reason @@ -683,6 +718,12 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason } if status == common.ChannelStatusEnabled { delete(channel.ChannelInfo.MultiKeyStatusList, keyIndex) + if channel.ChannelInfo.MultiKeyDisabledReason != nil { + delete(channel.ChannelInfo.MultiKeyDisabledReason, keyIndex) + } + if channel.ChannelInfo.MultiKeyDisabledTime != nil { + delete(channel.ChannelInfo.MultiKeyDisabledTime, keyIndex) + } } else { channel.ChannelInfo.MultiKeyStatusList[keyIndex] = status if channel.ChannelInfo.MultiKeyDisabledReason == nil { @@ -694,20 +735,32 @@ func handlerMultiKeyUpdate(channel *Channel, usingKey string, status int, reason channel.ChannelInfo.MultiKeyDisabledReason[keyIndex] = reason channel.ChannelInfo.MultiKeyDisabledTime[keyIndex] = common.GetTimestamp() } - if !hasEnabledMultiKey(keys, channel.ChannelInfo.MultiKeyStatusList) { - channel.Status = common.ChannelStatusAutoDisabled + if !channel.HasEnabledMultiKey() { + if channel.Status != common.ChannelStatusManuallyDisabled { + 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.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 + statusReason, _ := info["status_reason"].(string) + if statusReason == "All keys are disabled" { + channel.Status = common.ChannelStatusEnabled + delete(info, "status_reason") + delete(info, "status_time") + channel.SetOtherInfo(info) + } } } } func hasEnabledMultiKey(keys []string, statusList map[int]int) bool { - for i := range keys { + for i, key := range keys { + if !IsUsableChannelKey(key) { + continue + } if statusList == nil { return true } @@ -719,7 +772,111 @@ func hasEnabledMultiKey(keys []string, statusList map[int]int) bool { return false } +func channelAbilitiesEnabled(channel *Channel) bool { + if channel == nil || channel.Status != common.ChannelStatusEnabled { + return false + } + if !channel.ChannelInfo.IsMultiKey { + return true + } + return hasEnabledMultiKey(channel.GetKeys(), channel.ChannelInfo.MultiKeyStatusList) +} + +func (channel *Channel) normalizeMultiKeyAvailability() { + if channel == nil || !channel.ChannelInfo.IsMultiKey { + return + } + keys := channel.GetKeys() + channel.ChannelInfo.MultiKeySize = len(keys) + for idx := range channel.ChannelInfo.MultiKeyStatusList { + if idx < 0 || idx >= len(keys) { + delete(channel.ChannelInfo.MultiKeyStatusList, idx) + } + } + for idx := range channel.ChannelInfo.MultiKeyDisabledReason { + if idx < 0 || idx >= len(keys) { + delete(channel.ChannelInfo.MultiKeyDisabledReason, idx) + } + } + for idx := range channel.ChannelInfo.MultiKeyDisabledTime { + if idx < 0 || idx >= len(keys) { + delete(channel.ChannelInfo.MultiKeyDisabledTime, idx) + } + } + + hasEnabledKey := hasEnabledMultiKey(keys, channel.ChannelInfo.MultiKeyStatusList) + info := channel.GetOtherInfo() + statusReason, _ := info["status_reason"].(string) + if hasEnabledKey { + if channel.Status == common.ChannelStatusAutoDisabled && statusReason == "All keys are disabled" { + channel.Status = common.ChannelStatusEnabled + delete(info, "status_reason") + delete(info, "status_time") + channel.SetOtherInfo(info) + } + return + } + if channel.Status == common.ChannelStatusManuallyDisabled { + return + } + channel.Status = common.ChannelStatusAutoDisabled + if statusReason != "All keys are disabled" { + info["status_reason"] = "All keys are disabled" + info["status_time"] = common.GetTimestamp() + channel.SetOtherInfo(info) + } +} + +func remapMultiKeyStateByKey(channel *Channel, previous *Channel) { + if channel == nil || previous == nil || channel.Key == "" || channel.Key == previous.Key { + return + } + previousIndexes := make(map[string][]int) + for idx, key := range previous.GetKeys() { + key = strings.TrimSpace(key) + if key != "" { + previousIndexes[key] = append(previousIndexes[key], idx) + } + } + nextKeys := (&Channel{Key: channel.Key}).GetKeys() + statusList := make(map[int]int) + disabledReason := make(map[int]string) + disabledTime := make(map[int]int64) + for nextIdx, key := range nextKeys { + key = strings.TrimSpace(key) + indexes := previousIndexes[key] + if key == "" || len(indexes) == 0 { + continue + } + previousIdx := indexes[0] + previousIndexes[key] = indexes[1:] + if status, ok := previous.ChannelInfo.MultiKeyStatusList[previousIdx]; ok && status != common.ChannelStatusEnabled { + statusList[nextIdx] = status + } + if reason, ok := previous.ChannelInfo.MultiKeyDisabledReason[previousIdx]; ok { + disabledReason[nextIdx] = reason + } + if disabledAt, ok := previous.ChannelInfo.MultiKeyDisabledTime[previousIdx]; ok { + disabledTime[nextIdx] = disabledAt + } + } + channel.ChannelInfo.MultiKeyStatusList = statusList + channel.ChannelInfo.MultiKeyDisabledReason = disabledReason + channel.ChannelInfo.MultiKeyDisabledTime = disabledTime +} + func UpdateChannelStatus(channelId int, usingKey string, status int, reason string) bool { + changed, err := UpdateChannelStatusWithError(channelId, usingKey, status, reason) + if err != nil { + common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channelId, status, err)) + return false + } + return changed +} + +// UpdateChannelStatusWithError applies a status change and distinguishes a +// no-op from persistence or validation failures for management APIs. +func UpdateChannelStatusWithError(channelId int, usingKey string, status int, reason string) (bool, error) { if common.MemoryCacheEnabled { channelStatusLock.Lock() defer channelStatusLock.Unlock() @@ -732,85 +889,120 @@ func UpdateChannelStatus(channelId int, usingKey string, status int, reason stri pollingLock.Lock() defer pollingLock.Unlock() - if common.MemoryCacheEnabled { - channelCache, _ := CacheGetChannel(channelId) - if channelCache == nil { - return false - } - if channelCache.ChannelInfo.IsMultiKey { - beforeStatus := channelCache.Status - // 如果是多Key模式，更新缓存中的状态 - handlerMultiKeyUpdate(channelCache, usingKey, status, reason) - if beforeStatus != channelCache.Status { - CacheUpdateChannelStatus(channelId, channelCache.Status) - } - //CacheUpdateChannel(channelCache) - //return true - } else { - // 如果缓存渠道存在，且状态已是目标状态，直接返回 - if channelCache.Status == status { - return false + channel := &Channel{} + statusChanged := false + changed := false + err := DB.Transaction(func(tx *gorm.DB) error { + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Channel{}). + Where("id = ?", channelId). + UpdateColumn("status", gorm.Expr("status")).Error; err != nil { + return err } - CacheUpdateChannelStatus(channelId, status) } - } - - shouldUpdateAbilities := false - defer func() { - if shouldUpdateAbilities { - err := UpdateAbilityStatus(channelId, status == common.ChannelStatusEnabled) - if err != nil { - common.SysLog(fmt.Sprintf("failed to update ability status: channel_id=%d, error=%v", channelId, err)) - } + if err := lockForUpdate(tx).Where("id = ?", channelId).First(channel).Error; err != nil { + return err } - }() - channel, err := GetChannelById(channelId, true) - if err != nil { - return false - } else { - if channel.Status == status { - return false + if usingKey == "" && status == common.ChannelStatusEnabled && + channel.ChannelInfo.IsMultiKey && !channel.HasEnabledMultiKey() { + return fmt.Errorf("channel has no usable key") } + beforeStatus := channel.Status if channel.ChannelInfo.IsMultiKey { - beforeStatus := channel.Status + beforeInfo, err := common.Marshal(channel.ChannelInfo) + if err != nil { + return err + } handlerMultiKeyUpdate(channel, usingKey, status, reason) - if beforeStatus != channel.Status { - shouldUpdateAbilities = true + afterInfo, err := common.Marshal(channel.ChannelInfo) + if err != nil { + return err } + changed = beforeStatus != channel.Status || !bytes.Equal(beforeInfo, afterInfo) } else { + if beforeStatus == status { + return nil + } info := channel.GetOtherInfo() info["status_reason"] = reason info["status_time"] = common.GetTimestamp() channel.SetOtherInfo(info) channel.Status = status - shouldUpdateAbilities = true + changed = true } - err = channel.saveStatusState() - if err != nil { - common.SysLog(fmt.Sprintf("failed to update channel status: channel_id=%d, status=%d, error=%v", channel.Id, status, err)) - return false + if !changed { + return nil + } + statusChanged = beforeStatus != channel.Status + if err := channel.saveStatusStateWithDB(tx); err != nil { + return err + } + if statusChanged { + return tx.Model(&Ability{}). + Where("channel_id = ?", channelId). + Update("enabled", channelAbilitiesEnabled(channel)).Error } + return nil + }) + if err != nil { + return false, err } - return true + if !changed { + return false, nil + } + + CacheUpdateChannelState(channel) + return true, nil } func EnableChannelByTag(tag string) error { - err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusEnabled).Error - if err != nil { - return err - } - err = UpdateAbilityStatusByTag(tag, true) - return err + return DB.Transaction(func(tx *gorm.DB) error { + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Channel{}). + Where("tag = ?", tag). + UpdateColumn("status", gorm.Expr("status")).Error; err != nil { + return err + } + } + var channels []Channel + if err := lockForUpdate(tx).Where("tag = ?", tag).Order("id ASC").Find(&channels).Error; err != nil { + return err + } + for i := range channels { + channel := &channels[i] + if channel.ChannelInfo.IsMultiKey { + channel.normalizeMultiKeyAvailability() + if hasEnabledMultiKey(channel.GetKeys(), channel.ChannelInfo.MultiKeyStatusList) { + channel.Status = common.ChannelStatusEnabled + } + } else { + channel.Status = common.ChannelStatusEnabled + } + if err := tx.Model(&Channel{}).Where("id = ?", channel.Id).Updates(map[string]interface{}{ + "status": channel.Status, + "other_info": channel.OtherInfo, + "channel_info": channel.ChannelInfo, + }).Error; err != nil { + return err + } + if err := tx.Model(&Ability{}). + Where("channel_id = ?", channel.Id). + Update("enabled", channelAbilitiesEnabled(channel)).Error; err != nil { + return err + } + } + return nil + }) } func DisableChannelByTag(tag string) error { - err := DB.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error - if err != nil { - return err - } - err = UpdateAbilityStatusByTag(tag, false) - return err + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Model(&Channel{}).Where("tag = ?", tag).Update("status", common.ChannelStatusManuallyDisabled).Error; err != nil { + return err + } + return tx.Model(&Ability{}).Where("tag = ?", tag).Update("enabled", false).Error + }) } func EditChannelByTag(tag string, newTag *string, modelMapping *string, models *string, group *string, priority *int64, weight *uint, paramOverride *string, headerOverride *string) error { @@ -846,27 +1038,34 @@ func EditChannelByTag(tag string, newTag *string, modelMapping *string, models * updateData.HeaderOverride = headerOverride } - err := DB.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error - if err != nil { - return err - } - if shouldReCreateAbilities { - channels, err := GetChannelsByTag(updatedTag, false, false) - if err == nil { + return DB.Transaction(func(tx *gorm.DB) error { + if err := tx.Model(&Channel{}).Where("tag = ?", tag).Updates(updateData).Error; err != nil { + return err + } + if shouldReCreateAbilities { + var channels []*Channel + if err := tx.Where("tag = ?", updatedTag).Find(&channels).Error; err != nil { + return err + } for _, channel := range channels { - err = channel.UpdateAbilities(nil) - if err != nil { - common.SysLog(fmt.Sprintf("failed to update abilities: channel_id=%d, tag=%s, error=%v", channel.Id, channel.GetTag(), err)) + if err := channel.UpdateAbilities(tx); err != nil { + return fmt.Errorf("failed to update abilities: channel_id=%d, tag=%s: %w", channel.Id, channel.GetTag(), err) } } + return nil } - } else { - err := UpdateAbilityByTag(tag, newTag, priority, weight) - if err != nil { - return err + ability := Ability{} + if newTag != nil { + ability.Tag = newTag } - } - return nil + if priority != nil { + ability.Priority = priority + } + if weight != nil { + ability.Weight = *weight + } + return tx.Model(&Ability{}).Where("tag = ?", tag).Updates(ability).Error + }) } func UpdateChannelUsedQuota(id int, quota int) { @@ -885,13 +1084,59 @@ func updateChannelUsedQuota(id int, quota int) { } func DeleteChannelByStatus(status int64) (int64, error) { - result := DB.Where("status = ?", status).Delete(&Channel{}) - return result.RowsAffected, result.Error + return deleteChannelsByStatuses([]int64{status}) } func DeleteDisabledChannel() (int64, error) { - result := DB.Where("status = ? or status = ?", common.ChannelStatusAutoDisabled, common.ChannelStatusManuallyDisabled).Delete(&Channel{}) - return result.RowsAffected, result.Error + return deleteChannelsByStatuses([]int64{ + common.ChannelStatusAutoDisabled, + common.ChannelStatusManuallyDisabled, + }) +} + +func deleteChannelsByStatuses(statuses []int64) (int64, error) { + if len(statuses) == 0 { + return 0, nil + } + var deletedCount int64 + err := DB.Transaction(func(tx *gorm.DB) error { + // SQLite has no row-level FOR UPDATE. Acquire its writer lock before the + // status snapshot so a concurrent recovery cannot race the deletion. + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Channel{}). + Where("id = ?", 0). + UpdateColumn("status", gorm.Expr("status")).Error; err != nil { + return err + } + } + + var ids []int + if err := lockForUpdate(tx).Model(&Channel{}). + Where("status IN ?", statuses). + Order("id ASC"). + Pluck("id", &ids).Error; err != nil { + return err + } + if len(ids) == 0 { + return nil + } + + // Delete in the same channel -> ability order used by all other channel + // mutations. Recheck status defensively even though the rows are locked. + result := tx.Where("id IN ? AND status IN ?", ids, statuses).Delete(&Channel{}) + if result.Error != nil { + return result.Error + } + if result.RowsAffected != int64(len(ids)) { + return fmt.Errorf("channel status changed during deletion: selected=%d deleted=%d", len(ids), result.RowsAffected) + } + if err := tx.Where("channel_id IN ?", ids).Delete(&Ability{}).Error; err != nil { + return err + } + deletedCount = result.RowsAffected + return nil + }) + return deletedCount, err } func GetPaginatedTags(offset int, limit int) ([]*string, error) { diff --git a/model/channel_availability_atomic_test.go b/model/channel_availability_atomic_test.go new file mode 100644 index 000000000000..f2a7d06d5b6b --- /dev/null +++ b/model/channel_availability_atomic_test.go @@ -0,0 +1,238 @@ +package model + +import ( + "errors" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func useChannelAvailabilityTestDB(t *testing.T) *gorm.DB { + t.Helper() + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&Channel{}, &Ability{})) + + originalDB := DB + originalDatabaseType := common.MainDatabaseType() + originalMemoryCacheEnabled := common.MemoryCacheEnabled + DB = db + common.SetMainDatabaseType(common.DatabaseTypeSQLite) + common.MemoryCacheEnabled = false + t.Cleanup(func() { + DB = originalDB + common.SetMainDatabaseType(originalDatabaseType) + common.MemoryCacheEnabled = originalMemoryCacheEnabled + }) + return db +} + +func newMultiKeyChannel(id int, status int, key string, keyStatuses map[int]int) Channel { + tag := "multi-key" + return Channel{ + Id: id, + Type: 1, + Key: key, + Status: status, + Name: "multi-key-channel", + Models: "gpt-4", + Group: "default", + Tag: &tag, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeyStatusList: keyStatuses, + }, + } +} + +func TestBatchInsertChannelsNormalizesMultiKeyAvailability(t *testing.T) { + useChannelAvailabilityTestDB(t) + channels := []Channel{ + newMultiKeyChannel(1, common.ChannelStatusEnabled, "disabled-key", map[int]int{0: common.ChannelStatusManuallyDisabled}), + newMultiKeyChannel(2, common.ChannelStatusManuallyDisabled, "disabled-key", map[int]int{0: common.ChannelStatusManuallyDisabled}), + } + + require.NoError(t, BatchInsertChannels(channels)) + + var persisted []Channel + require.NoError(t, DB.Order("id ASC").Find(&persisted).Error) + require.Len(t, persisted, 2) + assert.Equal(t, common.ChannelStatusAutoDisabled, persisted[0].Status) + assert.Equal(t, common.ChannelStatusManuallyDisabled, persisted[1].Status) + var abilities []Ability + require.NoError(t, DB.Order("channel_id ASC").Find(&abilities).Error) + require.Len(t, abilities, 2) + assert.False(t, abilities[0].Enabled) + assert.False(t, abilities[1].Enabled) +} + +func TestEnableChannelByTagKeepsChannelDisabledWithoutUsableKey(t *testing.T) { + useChannelAvailabilityTestDB(t) + channel := newMultiKeyChannel(1, common.ChannelStatusManuallyDisabled, "disabled-key", map[int]int{0: common.ChannelStatusManuallyDisabled}) + require.NoError(t, channel.Insert()) + + require.NoError(t, EnableChannelByTag(channel.GetTag())) + + var persisted Channel + require.NoError(t, DB.First(&persisted, channel.Id).Error) + assert.Equal(t, common.ChannelStatusManuallyDisabled, persisted.Status) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestInitChannelCacheBuildsRoutesFromChannelSnapshot(t *testing.T) { + useChannelAvailabilityTestDB(t) + common.MemoryCacheEnabled = true + + channelSyncLock.Lock() + originalGroupRoutes := group2model2channels + originalChannels := channelsIDM + originalAdvancedCustom := channel2advancedCustomConfig + originalGeneration := channelCacheGeneration + channelSyncLock.Unlock() + t.Cleanup(func() { + channelSyncLock.Lock() + group2model2channels = originalGroupRoutes + channelsIDM = originalChannels + channel2advancedCustomConfig = originalAdvancedCustom + channelCacheGeneration = originalGeneration + channelSyncLock.Unlock() + }) + + channel := Channel{ + Name: "cache-snapshot-channel", + Type: 1, + Key: "key", + Status: common.ChannelStatusEnabled, + Models: "gpt-4", + Group: "default", + } + require.NoError(t, DB.Create(&channel).Error) + require.NotPanics(t, InitChannelCache) + + channelSyncLock.RLock() + routes := append([]int(nil), group2model2channels["default"]["gpt-4"]...) + channelSyncLock.RUnlock() + assert.Equal(t, []int{channel.Id}, routes) +} + +func TestChannelUpdateReplacesDisabledKeyStateByKeyIdentity(t *testing.T) { + useChannelAvailabilityTestDB(t) + channel := newMultiKeyChannel(1, common.ChannelStatusEnabled, "old-key", map[int]int{0: common.ChannelStatusManuallyDisabled}) + require.NoError(t, channel.Insert()) + require.NoError(t, DB.First(&channel, channel.Id).Error) + require.Equal(t, common.ChannelStatusAutoDisabled, channel.Status) + + channel.Key = "new-key" + require.NoError(t, channel.Update()) + + var persisted Channel + require.NoError(t, DB.First(&persisted, channel.Id).Error) + assert.Equal(t, common.ChannelStatusEnabled, persisted.Status) + assert.Empty(t, persisted.ChannelInfo.MultiKeyStatusList) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.True(t, ability.Enabled) +} + +func TestHasEnabledMultiKeyIgnoresBlankKeys(t *testing.T) { + assert.False(t, hasEnabledMultiKey([]string{"", " ", "\t"}, nil)) + assert.True(t, hasEnabledMultiKey([]string{"", "usable"}, nil)) +} + +func TestChannelInsertRollsBackWhenAbilityCreationFails(t *testing.T) { + db := useChannelAvailabilityTestDB(t) + forcedErr := errors.New("forced ability create failure") + callbackName := "test:fail-ability-create-on-insert" + require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "abilities" { + tx.AddError(forcedErr) + } + })) + t.Cleanup(func() { _ = db.Callback().Create().Remove(callbackName) }) + + channel := newMultiKeyChannel(1, common.ChannelStatusEnabled, "key", nil) + err := channel.Insert() + require.ErrorIs(t, err, forcedErr) + + var channelCount int64 + require.NoError(t, DB.Model(&Channel{}).Count(&channelCount).Error) + assert.Zero(t, channelCount) +} + +func TestChannelUpdateRollsBackChannelAndAbilitiesTogether(t *testing.T) { + db := useChannelAvailabilityTestDB(t) + channel := newMultiKeyChannel(1, common.ChannelStatusEnabled, "key", nil) + require.NoError(t, channel.Insert()) + + forcedErr := errors.New("forced ability recreate failure") + callbackName := "test:fail-ability-create-on-update" + require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "abilities" { + tx.AddError(forcedErr) + } + })) + t.Cleanup(func() { _ = db.Callback().Create().Remove(callbackName) }) + + channel.Models = "claude-3" + err := channel.Update() + require.ErrorIs(t, err, forcedErr) + + var persisted Channel + require.NoError(t, DB.First(&persisted, channel.Id).Error) + assert.Equal(t, "gpt-4", persisted.Models) + var abilities []Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).Find(&abilities).Error) + require.Len(t, abilities, 1) + assert.Equal(t, "gpt-4", abilities[0].Model) +} + +func TestUpdateChannelStatusRollsBackWhenAbilityUpdateFails(t *testing.T) { + db := useChannelAvailabilityTestDB(t) + channel := newMultiKeyChannel(1, common.ChannelStatusEnabled, "key", nil) + require.NoError(t, channel.Insert()) + + forcedErr := errors.New("forced ability status failure") + callbackName := "test:fail-ability-status-update" + require.NoError(t, db.Callback().Update().Before("gorm:update").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "abilities" { + tx.AddError(forcedErr) + } + })) + t.Cleanup(func() { _ = db.Callback().Update().Remove(callbackName) }) + + assert.False(t, UpdateChannelStatus(channel.Id, "", common.ChannelStatusManuallyDisabled, "test")) + var persisted Channel + require.NoError(t, DB.First(&persisted, channel.Id).Error) + assert.Equal(t, common.ChannelStatusEnabled, persisted.Status) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.True(t, ability.Enabled) +} + +func TestFixAbilityRollsBackAtomicRebuildOnFailure(t *testing.T) { + db := useChannelAvailabilityTestDB(t) + channel := newMultiKeyChannel(1, common.ChannelStatusEnabled, "key", nil) + require.NoError(t, channel.Insert()) + + forcedErr := errors.New("forced ability rebuild failure") + callbackName := "test:fail-ability-create-on-rebuild" + require.NoError(t, db.Callback().Create().Before("gorm:create").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "abilities" { + tx.AddError(forcedErr) + } + })) + t.Cleanup(func() { _ = db.Callback().Create().Remove(callbackName) }) + + _, _, err := FixAbility() + require.ErrorIs(t, err, forcedErr) + var abilities []Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).Find(&abilities).Error) + require.Len(t, abilities, 1) + assert.Equal(t, "gpt-4", abilities[0].Model) +} diff --git a/model/channel_cache.go b/model/channel_cache.go index 86c594384d50..d5b2fb9fa669 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -23,78 +23,91 @@ var channelsIDM map[int]*Channel // all channels include dis var channel2advancedCustomConfig map[int]*dto.AdvancedCustomConfig var channelSyncLock sync.RWMutex +// channelCacheGeneration prevents a full database snapshot from replacing a +// newer incremental cache publication that completed while the snapshot loaded. +var channelCacheGeneration uint64 + func InitChannelCache() { if !common.MemoryCacheEnabled { InvalidatePricingCache() return } - newChannelId2channel := make(map[int]*Channel) - newChannel2advancedCustomConfig := make(map[int]*dto.AdvancedCustomConfig) - var channels []*Channel - DB.Find(&channels) - for _, channel := range channels { - newChannelId2channel[channel.Id] = channel - if channel.Type == constant.ChannelTypeAdvancedCustom { - if config := channel.GetOtherSettings().AdvancedCustom; config != nil { - newChannel2advancedCustomConfig[channel.Id] = config - } + for { + channelSyncLock.RLock() + snapshotGeneration := channelCacheGeneration + channelSyncLock.RUnlock() + + newChannelId2channel := make(map[int]*Channel) + newChannel2advancedCustomConfig := make(map[int]*dto.AdvancedCustomConfig) + var channels []*Channel + if err := DB.Find(&channels).Error; err != nil { + common.SysError(fmt.Sprintf("failed to load channels for cache sync: %v", err)) + return } - } - var abilities []*Ability - DB.Find(&abilities) - groups := make(map[string]bool) - for _, ability := range abilities { - groups[ability.Group] = true - } - newGroup2model2channels := make(map[string]map[string][]int) - for group := range groups { - newGroup2model2channels[group] = make(map[string][]int) - } - for _, channel := range channels { - if channel.Status != common.ChannelStatusEnabled { - continue // skip disabled channels + for _, channel := range channels { + newChannelId2channel[channel.Id] = channel + if channel.Type == constant.ChannelTypeAdvancedCustom { + if config := channel.GetOtherSettings().AdvancedCustom; config != nil { + newChannel2advancedCustomConfig[channel.Id] = config + } + } } - groups := strings.Split(channel.Group, ",") - for _, group := range groups { - models := strings.Split(channel.Models, ",") - for _, model := range models { - if _, ok := newGroup2model2channels[group][model]; !ok { - newGroup2model2channels[group][model] = make([]int, 0) + newGroup2model2channels := make(map[string]map[string][]int) + for _, channel := range channels { + if !channelIsRoutable(channel) { + continue // skip disabled channels and multi-key channels without a usable key + } + groups := strings.Split(channel.Group, ",") + for _, group := range groups { + if _, ok := newGroup2model2channels[group]; !ok { + newGroup2model2channels[group] = make(map[string][]int) + } + models := strings.Split(channel.Models, ",") + for _, model := range models { + if _, ok := newGroup2model2channels[group][model]; !ok { + newGroup2model2channels[group][model] = make([]int, 0) + } + newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel.Id) } - newGroup2model2channels[group][model] = append(newGroup2model2channels[group][model], channel.Id) } } - } - // sort by priority - for group, model2channels := range newGroup2model2channels { - for model, channels := range model2channels { - sort.Slice(channels, func(i, j int) bool { - return newChannelId2channel[channels[i]].GetPriority() > newChannelId2channel[channels[j]].GetPriority() - }) - newGroup2model2channels[group][model] = channels + // sort by priority + for group, model2channels := range newGroup2model2channels { + for model, channels := range model2channels { + sort.Slice(channels, func(i, j int) bool { + return newChannelId2channel[channels[i]].GetPriority() > newChannelId2channel[channels[j]].GetPriority() + }) + newGroup2model2channels[group][model] = channels + } } - } - channelSyncLock.Lock() - group2model2channels = newGroup2model2channels - //channelsIDM = newChannelId2channel - for i, channel := range newChannelId2channel { - if channel.ChannelInfo.IsMultiKey { - channel.Keys = channel.GetKeys() - if channel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling { - if oldChannel, ok := channelsIDM[i]; ok { - // 存在旧的渠道，如果是多key且轮询，保留轮询索引信息 - if oldChannel.ChannelInfo.IsMultiKey && oldChannel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling { - channel.ChannelInfo.MultiKeyPollingIndex = oldChannel.ChannelInfo.MultiKeyPollingIndex + channelSyncLock.Lock() + if channelCacheGeneration != snapshotGeneration { + channelSyncLock.Unlock() + continue + } + group2model2channels = newGroup2model2channels + //channelsIDM = newChannelId2channel + for i, channel := range newChannelId2channel { + if channel.ChannelInfo.IsMultiKey { + channel.Keys = channel.GetKeys() + if channel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling { + if oldChannel, ok := channelsIDM[i]; ok { + // 存在旧的渠道，如果是多key且轮询，保留轮询索引信息 + if oldChannel.ChannelInfo.IsMultiKey && oldChannel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling { + channel.ChannelInfo.MultiKeyPollingIndex = oldChannel.ChannelInfo.MultiKeyPollingIndex + } } } } } + channelsIDM = newChannelId2channel + channel2advancedCustomConfig = newChannel2advancedCustomConfig + channelCacheGeneration++ + channelSyncLock.Unlock() + break } - channelsIDM = newChannelId2channel - channel2advancedCustomConfig = newChannel2advancedCustomConfig - channelSyncLock.Unlock() // Lock ordering: InvalidatePricingCache acquires updatePricingLock, and // GetPricing (holding updatePricingLock) nests channelSyncLock.RLock via // loadPricingAdvancedCustomConfigs. channelSyncLock MUST be released before @@ -273,22 +286,89 @@ func CacheUpdateChannelStatus(id int, status int) { return } channelSyncLock.Lock() - defer channelSyncLock.Unlock() if channel, ok := channelsIDM[id]; ok { channel.Status = status + refreshChannelRoutingCacheLocked(channel) } - if status != common.ChannelStatusEnabled { - // delete the channel from group2model2channels - for group, model2channels := range group2model2channels { - for model, channels := range model2channels { - for i, channelId := range channels { - if channelId == id { - // remove the channel from the slice - group2model2channels[group][model] = append(channels[:i], channels[i+1:]...) - break - } + channelCacheGeneration++ + channelSyncLock.Unlock() + InvalidatePricingCache() +} + +// CacheUpdateChannelState atomically publishes the status-owned channel fields +// and keeps the enabled routing index consistent with them. +func CacheUpdateChannelState(channel *Channel) { + if !common.MemoryCacheEnabled || channel == nil { + return + } + + channelSyncLock.Lock() + channelCacheGeneration++ + cached, ok := channelsIDM[channel.Id] + if !ok { + channelSyncLock.Unlock() + InvalidatePricingCache() + return + } + if channel.ChannelInfo.IsMultiKey && cached.ChannelInfo.IsMultiKey && + channel.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling && + cached.ChannelInfo.MultiKeyMode == constant.MultiKeyModePolling { + channel.ChannelInfo.MultiKeyPollingIndex = cached.ChannelInfo.MultiKeyPollingIndex + } + if channel.ChannelInfo.IsMultiKey { + channel.Keys = channel.GetKeys() + } + channelsIDM[channel.Id] = channel + refreshChannelRoutingCacheLocked(channel) + channelSyncLock.Unlock() + InvalidatePricingCache() +} + +func channelIsRoutable(channel *Channel) bool { + if channel == nil || channel.Status != common.ChannelStatusEnabled { + return false + } + return !channel.ChannelInfo.IsMultiKey || channel.HasEnabledMultiKey() +} + +func refreshChannelRoutingCacheLocked(channel *Channel) { + if group2model2channels == nil { + group2model2channels = make(map[string]map[string][]int) + } + for group, model2channels := range group2model2channels { + for model, channelIDs := range model2channels { + filtered := channelIDs[:0] + for _, channelID := range channelIDs { + if channelID != channel.Id { + filtered = append(filtered, channelID) } } + group2model2channels[group][model] = filtered + } + } + if !channelIsRoutable(channel) { + return + } + + for _, group := range strings.Split(channel.Group, ",") { + if group2model2channels[group] == nil { + group2model2channels[group] = make(map[string][]int) + } + for _, modelName := range strings.Split(channel.Models, ",") { + channelIDs := group2model2channels[group][modelName] + channelIDs = append(channelIDs, channel.Id) + sort.Slice(channelIDs, func(i, j int) bool { + left := channelsIDM[channelIDs[i]] + right := channelsIDM[channelIDs[j]] + if left == nil { + return false + } + if right == nil { + return true + } + return left.GetPriority() > right.GetPriority() + }) + group2model2channels[group][modelName] = channelIDs } } } @@ -302,6 +382,7 @@ func CacheUpdateChannel(channel *Channel) { channelSyncLock.Unlock() return } + channelCacheGeneration++ if channelsIDM == nil { channelsIDM = make(map[int]*Channel) diff --git a/model/channel_mutation.go b/model/channel_mutation.go new file mode 100644 index 000000000000..d82f05588e54 --- /dev/null +++ b/model/channel_mutation.go @@ -0,0 +1,74 @@ +package model + +import ( + "fmt" + "strings" + + "github.com/QuantumNous/new-api/common" + + "gorm.io/gorm" +) + +// ReplaceMultiKeyKeys preserves per-key health state for credentials that +// remain present after an append or replacement operation. +func (channel *Channel) ReplaceMultiKeyKeys(keys string) { + if channel == nil { + return + } + previous := *channel + parsed := (&Channel{Key: keys}).GetKeys() + cleanKeys := make([]string, 0, len(parsed)) + for _, key := range parsed { + if IsUsableChannelKey(key) { + cleanKeys = append(cleanKeys, strings.TrimSpace(key)) + } + } + channel.Key = strings.Join(cleanKeys, "\n") + channel.Keys = nil + remapMultiKeyStateByKey(channel, &previous) +} + +// UpdateChannelAtomically applies an update intent to the latest persisted +// channel state and rebuilds its abilities in the same transaction. +func UpdateChannelAtomically(channelID int, apply func(*Channel) error) (*Channel, error) { + if channelID <= 0 || apply == nil { + return nil, fmt.Errorf("invalid channel update") + } + + channel := &Channel{} + err := DB.Transaction(func(tx *gorm.DB) error { + // SQLite has no SELECT FOR UPDATE. Acquire its single writer lock before + // reading so a competing writer cannot commit between the read and write. + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Channel{}). + Where("id = ?", channelID). + UpdateColumn("status", gorm.Expr("status")).Error; err != nil { + return err + } + } + if err := lockForUpdate(tx).Where("id = ?", channelID).First(channel).Error; err != nil { + return err + } + + channel.Keys = nil + if err := apply(channel); err != nil { + return err + } + channel.Id = channelID + channel.Keys = nil + channel.normalizeMultiKeyAvailability() + + if err := tx.Model(&Channel{}). + Where("id = ?", channelID). + Select("*"). + Omit("id"). + Updates(channel).Error; err != nil { + return err + } + return channel.UpdateAbilities(tx) + }) + if err != nil { + return nil, err + } + return channel, nil +} diff --git a/model/channel_status_test.go b/model/channel_status_test.go index e4ad86f8133c..74306fcd82f1 100644 --- a/model/channel_status_test.go +++ b/model/channel_status_test.go @@ -1,7 +1,10 @@ package model import ( + "sync" + "sync/atomic" "testing" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -51,6 +54,369 @@ func TestUpdateChannelStatusPersistsMultiKeyState(t *testing.T) { assert.Equal(t, 1, stored.ChannelInfo.MultiKeyPollingIndex) } +func TestGetNextEnabledKeySkipsBlankEntries(t *testing.T) { + channel := Channel{ + Id: 991, + Key: " \nvalid-but-disabled", + Status: common.ChannelStatusEnabled, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{1: common.ChannelStatusManuallyDisabled}, + }, + } + + key, _, apiErr := channel.GetNextEnabledKey() + require.NotNil(t, apiErr) + assert.Empty(t, key) +} + +func TestInitChannelCacheExcludesMultiKeyChannelWithoutUsableKey(t *testing.T) { + setupChannelStatusTest(t) + common.MemoryCacheEnabled = true + + channel := Channel{ + Name: "blank-multi-key", + Key: " \n\t", + Status: common.ChannelStatusEnabled, + Models: "cache-model", + Group: "default", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + }, + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "cache-model", ChannelId: channel.Id, Enabled: true, + }).Error) + + InitChannelCache() + selected, err := GetRandomSatisfiedChannel("default", "cache-model", 0, "") + require.NoError(t, err) + assert.Nil(t, selected) +} + +func TestInitChannelCacheDoesNotOverwriteConcurrentIncrementalUpdate(t *testing.T) { + setupChannelStatusTest(t) + common.MemoryCacheEnabled = true + + channel := Channel{ + Name: "concurrent-cache-refresh", + Key: "key", + Status: common.ChannelStatusEnabled, + Models: "cache-model", + Group: "default", + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "cache-model", ChannelId: channel.Id, Enabled: true, + }).Error) + InitChannelCache() + + snapshotRead := make(chan struct{}) + releaseSnapshot := make(chan struct{}) + var releaseOnce sync.Once + t.Cleanup(func() { releaseOnce.Do(func() { close(releaseSnapshot) }) }) + var intercepted atomic.Bool + const callbackName = "test:block_stale_channel_cache_refresh" + require.NoError(t, DB.Callback().Query().After("gorm:query").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement != nil && tx.Statement.Table == "channels" && intercepted.CompareAndSwap(false, true) { + close(snapshotRead) + <-releaseSnapshot + } + })) + t.Cleanup(func() { _ = DB.Callback().Query().Remove(callbackName) }) + + refreshDone := make(chan struct{}) + go func() { + InitChannelCache() + close(refreshDone) + }() + + select { + case <-snapshotRead: + case <-time.After(5 * time.Second): + require.FailNow(t, "timed out waiting for the full cache refresh snapshot") + } + + require.True(t, UpdateChannelStatus( + channel.Id, "", common.ChannelStatusAutoDisabled, "provider rejected channel", + )) + + releaseOnce.Do(func() { close(releaseSnapshot) }) + select { + case <-refreshDone: + case <-time.After(5 * time.Second): + require.FailNow(t, "timed out waiting for the full cache refresh retry") + } + + cached, err := CacheGetChannel(channel.Id) + require.NoError(t, err) + assert.Equal(t, common.ChannelStatusAutoDisabled, cached.Status) + selected, err := GetRandomSatisfiedChannel("default", "cache-model", 0, "") + require.NoError(t, err) + assert.Nil(t, selected) +} + +func TestUpdateChannelStatusRestoresEnabledChannelToRoutingCache(t *testing.T) { + setupChannelStatusTest(t) + common.MemoryCacheEnabled = true + + channel := Channel{ + Name: "recover-cache-route", + Key: "key", + Status: common.ChannelStatusAutoDisabled, + Models: "cache-model", + Group: "default", + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "cache-model", ChannelId: channel.Id, Enabled: false, + }).Error) + InitChannelCache() + + require.True(t, UpdateChannelStatus(channel.Id, "", common.ChannelStatusEnabled, "recovered")) + selected, err := GetRandomSatisfiedChannel("default", "cache-model", 0, "") + require.NoError(t, err) + require.NotNil(t, selected) + assert.Equal(t, channel.Id, selected.Id) +} + +func TestUpdateChannelStatusDoesNotEnableMultiKeyChannelWithoutUsableKey(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "no-usable-key", + Key: "key-a\n ", + Status: common.ChannelStatusAutoDisabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{0: common.ChannelStatusAutoDisabled}, + }, + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + assert.False(t, UpdateChannelStatus(channel.Id, "", common.ChannelStatusEnabled, "manual operation")) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.Status) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestUpdateChannelStatusPreservesManualDisableDuringMultiKeyHealthChanges(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "manual-disable-multi-key", + Key: "key-a\nkey-b", + Status: common.ChannelStatusManuallyDisabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + }, + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + require.True(t, UpdateChannelStatus( + channel.Id, "key-a", common.ChannelStatusAutoDisabled, "provider rejected key", + )) + require.True(t, UpdateChannelStatus( + channel.Id, "key-b", common.ChannelStatusAutoDisabled, "provider rejected key", + )) + require.True(t, UpdateChannelStatus( + channel.Id, "key-a", common.ChannelStatusEnabled, "recovered", + )) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusManuallyDisabled, stored.Status) + assert.NotContains(t, stored.ChannelInfo.MultiKeyStatusList, 0) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.ChannelInfo.MultiKeyStatusList[1]) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestUpdateChannelStatusIgnoresHealthResultForMissingLegacyMultiKey(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "manual-disable-empty-multi-key", + Status: common.ChannelStatusManuallyDisabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + }, + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + assert.False(t, UpdateChannelStatus( + channel.Id, "deleted-key", common.ChannelStatusAutoDisabled, "late health result", + )) + assert.False(t, UpdateChannelStatus( + channel.Id, "deleted-key", common.ChannelStatusEnabled, "late recovery", + )) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusManuallyDisabled, stored.Status) +} + +func TestUpdateChannelStatusDoesNotEnableEmptyLegacyMultiKey(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "empty-legacy-multi-key", + Status: common.ChannelStatusAutoDisabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + }, + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + assert.False(t, UpdateChannelStatus( + channel.Id, "", common.ChannelStatusEnabled, "manual operation", + )) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.Status) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestReplaceMultiKeyKeysDropsUnusableEntriesAndRemapsState(t *testing.T) { + channel := Channel{ + Key: "key-a\nkey-b", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{1: common.ChannelStatusAutoDisabled}, + MultiKeyDisabledReason: map[int]string{1: "rejected"}, + MultiKeyDisabledTime: map[int]int64{1: 123}, + }, + } + + channel.ReplaceMultiKeyKeys(" \nkey-b\n\"\"\nkey-c") + channel.normalizeMultiKeyAvailability() + + assert.Equal(t, "key-b\nkey-c", channel.Key) + assert.Equal(t, 2, channel.ChannelInfo.MultiKeySize) + assert.Equal(t, common.ChannelStatusAutoDisabled, channel.ChannelInfo.MultiKeyStatusList[0]) + assert.Equal(t, "rejected", channel.ChannelInfo.MultiKeyDisabledReason[0]) + assert.Equal(t, int64(123), channel.ChannelInfo.MultiKeyDisabledTime[0]) +} + +func TestUpdateChannelStatusOnlyRecoversAvailabilityAutoDisable(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "unrelated-auto-disable", + Key: "key-a\nkey-b", + Status: common.ChannelStatusAutoDisabled, + Models: "gpt-4", + Group: "default", + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{0: common.ChannelStatusAutoDisabled, 1: common.ChannelStatusAutoDisabled}, + }, + } + channel.SetOtherInfo(map[string]interface{}{"status_reason": "provider unavailable"}) + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: false, + }).Error) + + require.True(t, UpdateChannelStatus(channel.Id, "key-a", common.ChannelStatusEnabled, "recovered")) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, common.ChannelStatusAutoDisabled, stored.Status) + assert.Equal(t, "provider unavailable", stored.GetOtherInfo()["status_reason"]) + var ability Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).First(&ability).Error) + assert.False(t, ability.Enabled) +} + +func TestUpdateChannelAtomicallyKeepsPersistedIdentity(t *testing.T) { + setupChannelStatusTest(t) + + channel := Channel{ + Name: "atomic-identity", + Key: "key", + Status: common.ChannelStatusEnabled, + Models: "gpt-4", + Group: "default", + } + require.NoError(t, DB.Create(&channel).Error) + require.NoError(t, DB.Create(&Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: true, + }).Error) + + updated, err := UpdateChannelAtomically(channel.Id, func(current *Channel) error { + current.Id = channel.Id + 1000 + current.Name = "updated" + return nil + }) + require.NoError(t, err) + assert.Equal(t, channel.Id, updated.Id) + + var stored Channel + require.NoError(t, DB.First(&stored, channel.Id).Error) + assert.Equal(t, "updated", stored.Name) + var abilities []Ability + require.NoError(t, DB.Where("channel_id = ?", channel.Id).Find(&abilities).Error) + require.Len(t, abilities, 1) + assert.Equal(t, channel.Id, abilities[0].ChannelId) +} + +func TestDeleteDisabledChannelPreservesEnabledChannelsAndAbilities(t *testing.T) { + setupChannelStatusTest(t) + + disabled := Channel{Name: "disabled", Status: common.ChannelStatusAutoDisabled} + enabled := Channel{Name: "enabled", Status: common.ChannelStatusEnabled} + require.NoError(t, DB.Create(&disabled).Error) + require.NoError(t, DB.Create(&enabled).Error) + require.NoError(t, DB.Create(&[]Ability{ + {Group: "default", Model: "disabled-model", ChannelId: disabled.Id, Enabled: false}, + {Group: "default", Model: "enabled-model", ChannelId: enabled.Id, Enabled: true}, + }).Error) + + deleted, err := DeleteDisabledChannel() + require.NoError(t, err) + assert.Equal(t, int64(1), deleted) + assert.ErrorIs(t, DB.First(&Channel{}, disabled.Id).Error, gorm.ErrRecordNotFound) + require.NoError(t, DB.First(&Channel{}, enabled.Id).Error) + assert.ErrorIs(t, DB.Where("channel_id = ?", disabled.Id).First(&Ability{}).Error, gorm.ErrRecordNotFound) + require.NoError(t, DB.Where("channel_id = ?", enabled.Id).First(&Ability{}).Error) +} + func TestSaveStatusStateFromSingleKeySnapshotPreservesUnownedColumns(t *testing.T) { setupChannelStatusTest(t) diff --git a/model/main.go b/model/main.go index 21445593e54e..8cf99996fcd3 100644 --- a/model/main.go +++ b/model/main.go @@ -19,8 +19,8 @@ import ( "gorm.io/gorm" ) -var commonGroupCol string -var commonKeyCol string +var commonGroupCol = "`group`" +var commonKeyCol = "`key`" var commonTrueVal string var commonFalseVal string diff --git a/model/model_channel_availability.go b/model/model_channel_availability.go new file mode 100644 index 000000000000..5f1b87eff30b --- /dev/null +++ b/model/model_channel_availability.go @@ -0,0 +1,205 @@ +package model + +import ( + "strings" + + "github.com/QuantumNous/new-api/common" + + "gorm.io/gorm" +) + +const ( + modelAvailabilityDisabled = 0 + modelAvailabilityEnabled = 1 + modelAvailabilityBatchSize = 500 +) + +// ModelChannelAvailabilityConfig controls one model/channel reconciliation. +// Automatic mode reads the persisted option pair while holding the same locks +// used by option updates. Manual mode uses the explicit Disable/Enable values. +type ModelChannelAvailabilityConfig struct { + Automatic bool + Disable bool + Enable bool +} + +// ModelChannelAvailabilityResult summarizes the committed model status changes. +type ModelChannelAvailabilityResult struct { + Disabled int + Enabled int + Skipped bool +} + +// ReconcileModelChannelAvailability atomically reconciles model metadata status +// against the latest enabled abilities and channels. +func ReconcileModelChannelAvailability(config ModelChannelAvailabilityConfig) (ModelChannelAvailabilityResult, error) { + result := ModelChannelAvailabilityResult{} + // FixAbility rebuilds the derivative table in one transaction. Share its + // in-process lock so reconciliation never reads a deliberately empty repair + // window. + fixLock.Lock() + defer fixLock.Unlock() + err := DB.Transaction(func(tx *gorm.DB) error { + // SQLite does not support SELECT FOR UPDATE. A harmless write acquires its + // database-wide writer lock before any state is read. + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Option{}). + Where(commonKeyCol+" = ?", ""). + UpdateColumn("value", gorm.Expr("value")).Error; err != nil { + return err + } + } + + optionKeys := []string{automaticDisableModelOptionKey, automaticEnableModelOptionKey} + var options []Option + if err := lockForUpdate(tx). + Where(commonKeyCol+" IN ?", optionKeys). + Order(commonKeyCol + " ASC"). + Find(&options).Error; err != nil { + return err + } + + disableEnabled := config.Disable + enableEnabled := config.Enable + if config.Automatic { + common.OptionMapRWMutex.RLock() + disableEnabled = common.AutomaticDisableModelEnabled + enableEnabled = common.AutomaticEnableModelEnabled + common.OptionMapRWMutex.RUnlock() + for _, option := range options { + switch option.Key { + case automaticDisableModelOptionKey: + disableEnabled = isEnabledOptionValue(option.Value) + case automaticEnableModelOptionKey: + enableEnabled = isEnabledOptionValue(option.Value) + } + } + // Auto-enable is subordinate to auto-disable. + enableEnabled = disableEnabled && enableEnabled + } + if !disableEnabled && !enableEnabled { + result.Skipped = true + return nil + } + + var models []*Model + if err := lockForUpdate(tx). + Select("id", "model_name", "status", "name_rule", "auto_disabled_by_rule"). + Order("id ASC"). + Find(&models).Error; err != nil { + return err + } + + // Keep ability and channel state in one statement snapshot. PostgreSQL's + // default READ COMMITTED isolation can otherwise observe a channel mutation + // between separate ability and channel queries and derive a state that never + // existed in the database. + var availableRows []struct { + Model string `gorm:"column:model"` + ChannelKey string `gorm:"column:channel_key"` + ChannelInfo ChannelInfo `gorm:"column:channel_info"` + } + if err := tx.Table("abilities"). + Select("abilities.model AS model, channels."+commonKeyCol+" AS channel_key, channels.channel_info AS channel_info"). + Joins("JOIN channels ON channels.id = abilities.channel_id"). + Where("abilities.enabled = ? AND channels.status = ?", true, common.ChannelStatusEnabled). + Scan(&availableRows).Error; err != nil { + return err + } + + availableModels := make(map[string]struct{}) + for i := range availableRows { + row := &availableRows[i] + if row.ChannelInfo.IsMultiKey { + channel := Channel{Key: row.ChannelKey, ChannelInfo: row.ChannelInfo} + if !channel.HasEnabledMultiKey() { + continue + } + } + name := strings.TrimSpace(row.Model) + if name == "" { + continue + } + availableModels[name] = struct{}{} + } + + disableIDs := make([]int, 0) + enableIDs := make([]int, 0) + for _, currentModel := range models { + if currentModel == nil { + continue + } + hasAvailableChannel := modelMatchesAvailableChannel(currentModel, availableModels) + if disableEnabled && currentModel.Status == modelAvailabilityEnabled && !hasAvailableChannel { + disableIDs = append(disableIDs, currentModel.Id) + continue + } + if enableEnabled && currentModel.Status == modelAvailabilityDisabled && currentModel.AutoDisabledByRule && hasAvailableChannel { + enableIDs = append(enableIDs, currentModel.Id) + } + } + + now := common.GetTimestamp() + for start := 0; start < len(disableIDs); start += modelAvailabilityBatchSize { + end := min(start+modelAvailabilityBatchSize, len(disableIDs)) + updateResult := tx.Model(&Model{}). + Where("id IN ? AND status = ?", disableIDs[start:end], modelAvailabilityEnabled). + Updates(map[string]interface{}{ + "status": modelAvailabilityDisabled, + "auto_disabled_by_rule": true, + "updated_time": now, + }) + if updateResult.Error != nil { + return updateResult.Error + } + result.Disabled += int(updateResult.RowsAffected) + } + for start := 0; start < len(enableIDs); start += modelAvailabilityBatchSize { + end := min(start+modelAvailabilityBatchSize, len(enableIDs)) + updateResult := tx.Model(&Model{}). + Where("id IN ? AND status = ? AND auto_disabled_by_rule = ?", enableIDs[start:end], modelAvailabilityDisabled, true). + Updates(map[string]interface{}{ + "status": modelAvailabilityEnabled, + "auto_disabled_by_rule": true, + "updated_time": now, + }) + if updateResult.Error != nil { + return updateResult.Error + } + result.Enabled += int(updateResult.RowsAffected) + } + return nil + }) + if err != nil { + return ModelChannelAvailabilityResult{}, err + } + return result, nil +} + +func modelMatchesAvailableChannel(currentModel *Model, availableModels map[string]struct{}) bool { + name := currentModel.ModelName + switch currentModel.NameRule { + case NameRulePrefix: + for available := range availableModels { + if strings.HasPrefix(available, name) { + return true + } + } + case NameRuleContains: + for available := range availableModels { + if strings.Contains(available, name) { + return true + } + } + case NameRuleSuffix: + for available := range availableModels { + if strings.HasSuffix(available, name) { + return true + } + } + default: + _, ok := availableModels[name] + return ok + } + return false +} diff --git a/model/model_meta.go b/model/model_meta.go index bd701e2bf7eb..944fa5f31c2c 100644 --- a/model/model_meta.go +++ b/model/model_meta.go @@ -22,18 +22,21 @@ type BoundChannel struct { } type Model struct { - Id int `json:"id"` - ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name_delete_at,priority:1"` - Description string `json:"description,omitempty" gorm:"type:text"` - Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"` - Tags string `json:"tags,omitempty" gorm:"type:varchar(255)"` - VendorID int `json:"vendor_id,omitempty" gorm:"index"` - Endpoints string `json:"endpoints,omitempty" gorm:"type:text"` - Status int `json:"status" gorm:"default:1"` - SyncOfficial int `json:"sync_official" gorm:"default:1"` - CreatedTime int64 `json:"created_time" gorm:"bigint"` - UpdatedTime int64 `json:"updated_time" gorm:"bigint"` - DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_model_name_delete_at,priority:2"` + Id int `json:"id"` + ModelName string `json:"model_name" gorm:"size:128;not null;uniqueIndex:uk_model_name_delete_at,priority:1"` + Description string `json:"description,omitempty" gorm:"type:text"` + Icon string `json:"icon,omitempty" gorm:"type:varchar(128)"` + Tags string `json:"tags,omitempty" gorm:"type:varchar(255)"` + VendorID int `json:"vendor_id,omitempty" gorm:"index"` + Endpoints string `json:"endpoints,omitempty" gorm:"type:text"` + Status int `json:"status" gorm:"default:1"` + // AutoDisabledByRule marks models managed by channel-availability automation. + // Used for auto re-enable protection and UI status badges/filters. + AutoDisabledByRule bool `json:"auto_disabled_by_rule" gorm:"column:auto_disabled_by_rule"` + SyncOfficial int `json:"sync_official" gorm:"default:1"` + CreatedTime int64 `json:"created_time" gorm:"bigint"` + UpdatedTime int64 `json:"updated_time" gorm:"bigint"` + DeletedAt gorm.DeletedAt `json:"-" gorm:"index;uniqueIndex:uk_model_name_delete_at,priority:2"` BoundChannels []BoundChannel `json:"bound_channels,omitempty" gorm:"-"` EnableGroups []string `json:"enable_groups,omitempty" gorm:"-"` @@ -48,21 +51,27 @@ func (mi *Model) Insert() error { now := common.GetTimestamp() mi.CreatedTime = now mi.UpdatedTime = now + // The automation marker is internal state and must never be accepted from + // model creation payloads. + mi.AutoDisabledByRule = false // 保存原始值（因为 Create 后可能被 GORM 的 default 标签覆盖为 1） originalStatus := mi.Status originalSyncOfficial := mi.SyncOfficial - // 先创建记录（GORM 会对零值字段应用默认值） - if err := DB.Create(mi).Error; err != nil { - return err - } - - // 使用保存的原始值进行更新，确保零值能正确保存 - return DB.Model(&Model{}).Where("id = ?", mi.Id).Updates(map[string]interface{}{ - "status": originalStatus, - "sync_official": originalSyncOfficial, - }).Error + return DB.Transaction(func(tx *gorm.DB) error { + // GORM applies default tags to zero values during Create. Keep the follow-up + // correction in the same transaction so reconciliation never observes the + // temporary defaults. + if err := tx.Create(mi).Error; err != nil { + return err + } + return tx.Model(&Model{}).Where("id = ?", mi.Id).Updates(map[string]interface{}{ + "status": originalStatus, + "sync_official": originalSyncOfficial, + "auto_disabled_by_rule": false, + }).Error + }) } func IsModelNameDuplicated(id int, name string) (bool, error) { @@ -76,10 +85,34 @@ func IsModelNameDuplicated(id int, name string) (bool, error) { func (mi *Model) Update() error { mi.UpdatedTime = common.GetTimestamp() - // 使用 Select 强制更新所有字段，包括零值 - return DB.Model(&Model{}).Where("id = ?", mi.Id). - Select("model_name", "description", "icon", "tags", "vendor_id", "endpoints", "status", "sync_official", "name_rule", "updated_time"). - Updates(mi).Error + return DB.Transaction(func(tx *gorm.DB) error { + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Model{}). + Where("id = ?", mi.Id). + UpdateColumn("updated_time", gorm.Expr("updated_time")).Error; err != nil { + return err + } + } + + var current Model + if err := lockForUpdate(tx). + Select("id", "status", "auto_disabled_by_rule"). + Where("id = ?", mi.Id). + First(¤t).Error; err != nil { + return err + } + + autoDisabledByRule := current.AutoDisabledByRule + if current.Status != mi.Status { + autoDisabledByRule = false + } + mi.AutoDisabledByRule = autoDisabledByRule + // 使用 Select 强制更新所有字段，包括零值。管理员显式修改状态时， + // status 与自动禁用标记在同一事务内更新，不暴露中间状态。 + return tx.Model(&Model{}).Where("id = ?", mi.Id). + Select("model_name", "description", "icon", "tags", "vendor_id", "endpoints", "status", "auto_disabled_by_rule", "sync_official", "name_rule", "updated_time"). + Updates(mi).Error + }) } func (mi *Model) Delete() error { @@ -205,8 +238,8 @@ func SearchModels(keyword string, vendor string, status string, syncOfficial str db = db.Joins("JOIN vendors ON vendors.id = models.vendor_id").Where("vendors.name LIKE ?", "%"+vendor+"%") } } - if statusValue, ok := parseModelStatusFilter(status); ok { - db = db.Where("models.status = ?", statusValue) + if statusFilter, ok := parseModelStatusFilterSpec(status); ok { + db = db.Where(statusFilter.query, statusFilter.args...) } if syncValue, ok := parseModelSyncFilter(syncOfficial); ok { db = db.Where("models.sync_official = ?", syncValue) @@ -221,22 +254,37 @@ func SearchModels(keyword string, vendor string, status string, syncOfficial str return models, total, nil } -// parseModelStatusFilter maps UI/API status values to the models.status column. +type modelStatusFilterSpec struct { + query string + args []interface{} +} + +// parseModelStatusFilterSpec maps UI/API status values to SQL filters on models. // Returns ok=false when no status filter should be applied. -func parseModelStatusFilter(status string) (value int, ok bool) { +func parseModelStatusFilterSpec(status string) (spec modelStatusFilterSpec, ok bool) { switch strings.ToLower(strings.TrimSpace(status)) { case "", "all": - return 0, false + return modelStatusFilterSpec{}, false case "enabled", "1": - return 1, true + return modelStatusFilterSpec{query: "models.status = ?", args: []interface{}{1}}, true case "disabled", "0": - return 0, true + return modelStatusFilterSpec{query: "models.status = ?", args: []interface{}{0}}, true + case "auto-enabled", "auto_enabled": + return modelStatusFilterSpec{ + query: "models.status = ? AND models.auto_disabled_by_rule = ?", + args: []interface{}{1, true}, + }, true + case "auto-disabled", "auto_disabled": + return modelStatusFilterSpec{ + query: "models.status = ? AND models.auto_disabled_by_rule = ?", + args: []interface{}{0, true}, + }, true default: n, err := strconv.Atoi(status) if err != nil { - return 0, false + return modelStatusFilterSpec{}, false } - return n, true + return modelStatusFilterSpec{query: "models.status = ?", args: []interface{}{n}}, true } } diff --git a/model/model_meta_availability_test.go b/model/model_meta_availability_test.go new file mode 100644 index 000000000000..4fc56c70b9d6 --- /dev/null +++ b/model/model_meta_availability_test.go @@ -0,0 +1,75 @@ +package model + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/common" + "github.com/glebarez/sqlite" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func TestReconcileModelChannelAvailabilityReadsAbilitiesAndChannelsInOneStatement(t *testing.T) { + db := useChannelAvailabilityTestDB(t) + require.NoError(t, db.AutoMigrate(&Model{}, &Option{})) + + channel := Channel{ + Id: 1, + Type: 1, + Key: "key", + Status: common.ChannelStatusEnabled, + Name: "channel", + Models: "gpt-4", + Group: "default", + } + require.NoError(t, channel.Insert()) + require.NoError(t, db.Create(&Model{ModelName: "gpt-4", Status: 1, SyncOfficial: 1}).Error) + + availabilityQueries := make([]string, 0, 1) + callbackName := "test:capture-model-availability-source-query" + require.NoError(t, db.Callback().Row().After("gorm:row").Register(callbackName, func(tx *gorm.DB) { + sql := strings.ToLower(tx.Statement.SQL.String()) + if strings.Contains(sql, "abilities") || strings.Contains(sql, "channels") { + availabilityQueries = append(availabilityQueries, sql) + } + })) + t.Cleanup(func() { _ = db.Callback().Row().Remove(callbackName) }) + + result, err := ReconcileModelChannelAvailability(ModelChannelAvailabilityConfig{Disable: true}) + require.NoError(t, err) + assert.Zero(t, result.Disabled) + require.Len(t, availabilityQueries, 1) + assert.Contains(t, availabilityQueries[0], "join channels") +} + +func TestModelInsertPreservesZeroValuesAndIgnoresAutoDisabledMarker(t *testing.T) { + db, err := gorm.Open(sqlite.Open(":memory:"), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, db.AutoMigrate(&Model{})) + + originalDB := DB + originalDatabaseType := common.MainDatabaseType() + DB = db + common.SetMainDatabaseType(common.DatabaseTypeSQLite) + t.Cleanup(func() { + DB = originalDB + common.SetMainDatabaseType(originalDatabaseType) + }) + + created := &Model{ + ModelName: "gpt-4", + Status: 0, + SyncOfficial: 0, + AutoDisabledByRule: true, + } + require.NoError(t, created.Insert()) + + var persisted Model + require.NoError(t, DB.First(&persisted, created.Id).Error) + assert.False(t, created.AutoDisabledByRule) + assert.Equal(t, 0, persisted.Status) + assert.Equal(t, 0, persisted.SyncOfficial) + assert.False(t, persisted.AutoDisabledByRule) +} diff --git a/model/option.go b/model/option.go index e7fda5231be7..aa04718d9c76 100644 --- a/model/option.go +++ b/model/option.go @@ -20,6 +20,11 @@ type Option struct { Value string `json:"value"` } +const ( + automaticDisableModelOptionKey = "AutomaticDisableModelEnabled" + automaticEnableModelOptionKey = "AutomaticEnableModelEnabled" +) + func AllOption() ([]*Option, error) { var options []*Option var err error @@ -47,6 +52,8 @@ func InitOptionMap() { common.OptionMap["RegisterEnabled"] = strconv.FormatBool(common.RegisterEnabled) common.OptionMap["AutomaticDisableChannelEnabled"] = strconv.FormatBool(common.AutomaticDisableChannelEnabled) common.OptionMap["AutomaticEnableChannelEnabled"] = strconv.FormatBool(common.AutomaticEnableChannelEnabled) + common.OptionMap["AutomaticDisableModelEnabled"] = strconv.FormatBool(common.AutomaticDisableModelEnabled) + common.OptionMap["AutomaticEnableModelEnabled"] = strconv.FormatBool(common.AutomaticEnableModelEnabled) common.OptionMap["LogConsumeEnabled"] = strconv.FormatBool(common.LogConsumeEnabled) common.OptionMap["DisplayInCurrencyEnabled"] = strconv.FormatBool(common.DisplayInCurrencyEnabled) common.OptionMap["DisplayTokenStatEnabled"] = strconv.FormatBool(common.DisplayTokenStatEnabled) @@ -188,13 +195,23 @@ func InitOptionMap() { } func loadOptionsFromDatabase() { - options, _ := AllOption() + options, err := AllOption() + if err != nil { + common.SysLog("failed to load options from database: " + err.Error()) + return + } for _, option := range options { + if option.Key == automaticDisableModelOptionKey || option.Key == automaticEnableModelOptionKey { + continue + } err := updateOptionMap(option.Key, option.Value) if err != nil { common.SysLog("failed to update option map: " + err.Error()) } } + if err := updateOptionsBulk(nil, true); err != nil { + common.SysLog("failed to normalize model availability options: " + err.Error()) + } } func SyncOptions(frequency int) { @@ -216,6 +233,9 @@ func validateOptionValue(key string, value string) error { } func UpdateOption(key string, value string) error { + if key == automaticDisableModelOptionKey || key == automaticEnableModelOptionKey { + return UpdateOptionsBulk(map[string]string{key: value}) + } if err := validateOptionValue(key, value); err != nil { return err } @@ -240,7 +260,11 @@ func UpdateOption(key string, value string) error { // is touched — safe for callers that must commit a set of related options // atomically (e.g. payment gateway binding). func UpdateOptionsBulk(values map[string]string) error { - if len(values) == 0 { + return updateOptionsBulk(values, false) +} + +func updateOptionsBulk(values map[string]string, normalizeStoredModelAvailability bool) error { + if len(values) == 0 && !normalizeStoredModelAvailability { return nil } for key, value := range values { @@ -248,8 +272,68 @@ func UpdateOptionsBulk(values map[string]string) error { return err } } + normalizedValues := make(map[string]string, len(values)+2) + for key, value := range values { + normalizedValues[key] = value + } + _, updatesDisableModel := values[automaticDisableModelOptionKey] + _, updatesEnableModel := values[automaticEnableModelOptionKey] + updatesModelAvailability := normalizeStoredModelAvailability || updatesDisableModel || updatesEnableModel + applyModelAvailability := false + normalizedDisableValue := "false" + normalizedEnableValue := "false" + err := DB.Transaction(func(tx *gorm.DB) error { - for k, v := range values { + if updatesModelAvailability { + if common.UsingMainDatabase(common.DatabaseTypeSQLite) { + if err := tx.Model(&Option{}). + Where(commonKeyCol+" = ?", ""). + UpdateColumn("value", gorm.Expr("value")).Error; err != nil { + return err + } + } + delete(normalizedValues, automaticDisableModelOptionKey) + delete(normalizedValues, automaticEnableModelOptionKey) + disableValue := "false" + enableValue := "false" + var existingOptions []Option + if err := lockForUpdate(tx).Where(commonKeyCol+" IN ?", []string{ + automaticDisableModelOptionKey, + automaticEnableModelOptionKey, + }).Order(commonKeyCol + " ASC").Find(&existingOptions).Error; err != nil { + return err + } + if len(existingOptions) == 0 && !updatesDisableModel && !updatesEnableModel { + return nil + } + for _, option := range existingOptions { + switch option.Key { + case automaticDisableModelOptionKey: + disableValue = option.Value + case automaticEnableModelOptionKey: + enableValue = option.Value + } + } + if value, ok := values[automaticDisableModelOptionKey]; ok { + disableValue = value + } + if value, ok := values[automaticEnableModelOptionKey]; ok { + enableValue = value + } + disableEnabled := isEnabledOptionValue(disableValue) + enableEnabled := disableEnabled && isEnabledOptionValue(enableValue) + normalizedDisableValue = strconv.FormatBool(disableEnabled) + normalizedEnableValue = strconv.FormatBool(enableEnabled) + applyModelAvailability = true + needsPairPersistence := updatesDisableModel || updatesEnableModel || len(existingOptions) < 2 || + disableValue != normalizedDisableValue || enableValue != normalizedEnableValue + if needsPairPersistence { + normalizedValues[automaticDisableModelOptionKey] = normalizedDisableValue + normalizedValues[automaticEnableModelOptionKey] = normalizedEnableValue + } + } + + for k, v := range normalizedValues { option := Option{Key: k} if err := tx.FirstOrCreate(&option, Option{Key: k}).Error; err != nil { return err @@ -264,7 +348,18 @@ func UpdateOptionsBulk(values map[string]string) error { if err != nil { return err } - for k, v := range values { + if applyModelAvailability { + if err := updateOptionMap(automaticDisableModelOptionKey, normalizedDisableValue); err != nil { + return err + } + if err := updateOptionMap(automaticEnableModelOptionKey, normalizedEnableValue); err != nil { + return err + } + } + for k, v := range normalizedValues { + if k == automaticDisableModelOptionKey || k == automaticEnableModelOptionKey { + continue + } if err := updateOptionMap(k, v); err != nil { return err } @@ -272,6 +367,10 @@ func UpdateOptionsBulk(values map[string]string) error { return nil } +func isEnabledOptionValue(value string) bool { + return value == "1" || strings.EqualFold(value, "true") +} + func updateOptionMap(key string, value string) (err error) { if key == retiredThemeOptionKey { common.OptionMapRWMutex.Lock() @@ -281,6 +380,9 @@ func updateOptionMap(key string, value string) (err error) { } common.OptionMapRWMutex.Lock() defer common.OptionMapRWMutex.Unlock() + if common.OptionMap == nil { + common.OptionMap = make(map[string]string) + } common.OptionMap[key] = value // 检查是否是模型配置 - 使用更规范的方式处理 @@ -304,6 +406,10 @@ func updateOptionMap(key string, value string) (err error) { } if strings.HasSuffix(key, "Enabled") || key == "DefaultCollapseSidebar" || key == "DefaultUseAutoGroup" || key == "SMTPForceAuthLogin" || key == "SMTPInsecureSkipVerify" { boolValue := value == "true" + if key == automaticDisableModelOptionKey || key == automaticEnableModelOptionKey { + boolValue = isEnabledOptionValue(value) + common.OptionMap[key] = strconv.FormatBool(boolValue) + } switch key { case "PasswordRegisterEnabled": common.PasswordRegisterEnabled = boolValue @@ -331,6 +437,20 @@ func updateOptionMap(key string, value string) (err error) { common.AutomaticDisableChannelEnabled = boolValue case "AutomaticEnableChannelEnabled": common.AutomaticEnableChannelEnabled = boolValue + case automaticDisableModelOptionKey: + common.AutomaticDisableModelEnabled = boolValue + if !boolValue { + // Enable is paired with disable; keep both off together. + common.AutomaticEnableModelEnabled = false + common.OptionMap[automaticEnableModelOptionKey] = "false" + } + case automaticEnableModelOptionKey: + // Reject enable-without-disable so automation stays consistent. + if boolValue && !common.AutomaticDisableModelEnabled { + boolValue = false + common.OptionMap[key] = "false" + } + common.AutomaticEnableModelEnabled = boolValue case "LogConsumeEnabled": common.LogConsumeEnabled = boolValue case "DisplayInCurrencyEnabled": diff --git a/relay/mjproxy_handler.go b/relay/mjproxy_handler.go index 99755f0f52b1..8f7aa61f6a50 100644 --- a/relay/mjproxy_handler.go +++ b/relay/mjproxy_handler.go @@ -575,9 +575,16 @@ func RelayMidjourneySubmit(c *gin.Context, relayInfo *relaycommon.RelayInfo) *dt channel, err := model.GetChannelById(midjourneyTask.ChannelId, true) if err != nil { common.SysLog("get_channel_null: " + err.Error()) - } - if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled { - model.UpdateChannelStatus(midjourneyTask.ChannelId, "", 2, "No available account instance") + } else if channel.GetAutoBan() && common.AutomaticDisableChannelEnabled { + changed := model.UpdateChannelStatus( + midjourneyTask.ChannelId, + "", + common.ChannelStatusAutoDisabled, + "No available account instance", + ) + if changed { + service.SyncModelChannelAvailabilityAfterMutation("channel.midjourney_auto_disable") + } } } if midjResponse.Code != 1 && midjResponse.Code != 21 && midjResponse.Code != 22 { diff --git a/router/api-router.go b/router/api-router.go index 31c595e00db2..5634ad73fbeb 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -352,7 +352,10 @@ func SetApiRouter(router *gin.Engine) { modelsRoute.GET("/:id", controller.GetModelMeta) modelsRoute.POST("/", controller.CreateModelMeta) modelsRoute.PUT("/", controller.UpdateModelMeta) + modelsRoute.POST("/batch_status", controller.BatchUpdateModelStatus) modelsRoute.DELETE("/:id", controller.DeleteModelMeta) + modelsRoute.POST("/batch_disable_no_channels", controller.BatchDisableModelsNoChannels) + modelsRoute.POST("/batch_enable_with_channels", controller.BatchEnableModelsWithChannels) } // Deployments (model deployment management) diff --git a/service/channel.go b/service/channel.go index f348e081e2ed..90cdea60cb2c 100644 --- a/service/channel.go +++ b/service/channel.go @@ -27,6 +27,7 @@ func DisableChannel(channelError types.ChannelError, reason string) { success := model.UpdateChannelStatus(channelError.ChannelId, channelError.UsingKey, common.ChannelStatusAutoDisabled, reason) if success { + SyncModelChannelAvailabilityAfterMutation("channel.auto_disable") subject := fmt.Sprintf("通道「%s」（#%d）已被禁用", channelError.ChannelName, channelError.ChannelId) content := fmt.Sprintf("通道「%s」（#%d）已被禁用，原因：%s", channelError.ChannelName, channelError.ChannelId, reason) NotifyRootUser(formatNotifyType(channelError.ChannelId, common.ChannelStatusAutoDisabled), subject, content) @@ -36,6 +37,7 @@ func DisableChannel(channelError types.ChannelError, reason string) { func EnableChannel(channelId int, usingKey string, channelName string) { success := model.UpdateChannelStatus(channelId, usingKey, common.ChannelStatusEnabled, "") if success { + SyncModelChannelAvailabilityAfterMutation("channel.auto_enable") subject := fmt.Sprintf("通道「%s」（#%d）已被启用", channelName, channelId) content := fmt.Sprintf("通道「%s」（#%d）已被启用", channelName, channelId) NotifyRootUser(formatNotifyType(channelId, common.ChannelStatusEnabled), subject, content) diff --git a/service/channel_affinity_usage_cache_test.go b/service/channel_affinity_usage_cache_test.go index 876297b21c05..2e169f96142d 100644 --- a/service/channel_affinity_usage_cache_test.go +++ b/service/channel_affinity_usage_cache_test.go @@ -4,7 +4,6 @@ import ( "fmt" "net/http/httptest" "testing" - "time" "github.com/QuantumNous/new-api/relaykit/dto" "github.com/QuantumNous/new-api/relaykit/types" @@ -26,9 +25,9 @@ func buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP string) } func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) + ruleName := fmt.Sprintf("rule_%s", t.Name()) usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) + keyFP := fmt.Sprintf("fp_%s", t.Name()) ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) usage := &dto.Usage{ @@ -53,9 +52,9 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_ClaudeMode(t *testing.T) } func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) + ruleName := fmt.Sprintf("rule_%s", t.Name()) usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) + keyFP := fmt.Sprintf("fp_%s", t.Name()) ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) openAIUsage := &dto.Usage{ @@ -83,9 +82,9 @@ func TestObserveChannelAffinityUsageCacheByRelayFormat_MixedMode(t *testing.T) { } func TestObserveChannelAffinityUsageCacheByRelayFormat_UnsupportedModeKeepsEmpty(t *testing.T) { - ruleName := fmt.Sprintf("rule_%d", time.Now().UnixNano()) + ruleName := fmt.Sprintf("rule_%s", t.Name()) usingGroup := "default" - keyFP := fmt.Sprintf("fp_%d", time.Now().UnixNano()) + keyFP := fmt.Sprintf("fp_%s", t.Name()) ctx := buildChannelAffinityStatsContextForTest(ruleName, usingGroup, keyFP) usage := &dto.Usage{ diff --git a/service/model_channel_availability.go b/service/model_channel_availability.go new file mode 100644 index 000000000000..c34f4662d251 --- /dev/null +++ b/service/model_channel_availability.go @@ -0,0 +1,166 @@ +package service + +import ( + "fmt" + "strings" + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" +) + +const modelChannelAvailabilityRetryMaxDelay = time.Minute + +var ( + modelChannelAvailabilityRetryOnce sync.Once + modelChannelAvailabilityRetrySignal = make(chan string, 1) + modelChannelAvailabilityRetryPending sync.WaitGroup +) + +const ( + modelStatusEnabled = 1 + modelStatusDisabled = 0 +) + +// ModelChannelAvailabilityResult summarizes one reconciliation pass. +type ModelChannelAvailabilityResult struct { + Disabled int + Enabled int + Skipped bool + Reason string + PricingRefreshed bool +} + +// SyncModelChannelAvailability reconciles model status against available channels. +func SyncModelChannelAvailability(reason string) (ModelChannelAvailabilityResult, error) { + return syncModelChannelAvailability(reason, false) +} + +// SyncModelChannelAvailabilityFull also logs a successful zero-change pass. +func SyncModelChannelAvailabilityFull(reason string) (ModelChannelAvailabilityResult, error) { + return syncModelChannelAvailability(reason, true) +} + +// SyncModelChannelAvailabilityAfterMutation runs reconciliation after a primary +// mutation has already committed. A reconciliation failure is logged and queued +// for a coalesced retry, but is not returned as failure for the committed API +// operation. Startup calibration provides an additional recovery boundary. +func SyncModelChannelAvailabilityAfterMutation(reason string) ModelChannelAvailabilityResult { + result, err := SyncModelChannelAvailability(reason) + if err == nil { + return result + } + common.SysError(fmt.Sprintf( + "model channel availability sync after committed mutation failed: reason=%s err=%v", + reason, err, + )) + scheduleModelChannelAvailabilityRetry(reason) + return result +} + +func scheduleModelChannelAvailabilityRetry(reason string) { + modelChannelAvailabilityRetryOnce.Do(func() { + go runModelChannelAvailabilityRetryWorker() + }) + modelChannelAvailabilityRetryPending.Add(1) + select { + case modelChannelAvailabilityRetrySignal <- reason: + default: + modelChannelAvailabilityRetryPending.Done() + // Reconciliation is global, so one pending pass covers all mutations that + // committed before that pass reads the database. + } +} + +func runModelChannelAvailabilityRetryWorker() { + for reason := range modelChannelAvailabilityRetrySignal { + delay := time.Second + for { + if _, err := SyncModelChannelAvailabilityFull("retry." + reason); err == nil { + break + } else { + common.SysError(fmt.Sprintf( + "model channel availability retry failed: reason=%s retry_in=%s err=%v", + reason, delay, err, + )) + } + time.Sleep(delay) + delay = min(delay*2, modelChannelAvailabilityRetryMaxDelay) + } + modelChannelAvailabilityRetryPending.Done() + } +} + +// CalibrateModelChannelAvailabilityAtStartup repairs stale model status left by +// a process interruption after channel state committed but before reconciliation. +func CalibrateModelChannelAvailabilityAtStartup() error { + _, err := SyncModelChannelAvailabilityFull("startup") + return err +} + +func syncModelChannelAvailability(reason string, forceFull bool) (ModelChannelAvailabilityResult, error) { + return reconcileModelChannelAvailability(reason, forceFull, model.ModelChannelAvailabilityConfig{Automatic: true}) +} + +func reconcileModelChannelAvailability( + reason string, + forceFull bool, + config model.ModelChannelAvailabilityConfig, +) (ModelChannelAvailabilityResult, error) { + result := ModelChannelAvailabilityResult{Reason: reason} + modelResult, err := model.ReconcileModelChannelAvailability(config) + if err != nil { + return result, fmt.Errorf("reconcile model channel availability: %w", err) + } + result.Disabled = modelResult.Disabled + result.Enabled = modelResult.Enabled + result.Skipped = modelResult.Skipped + + if result.Disabled > 0 || result.Enabled > 0 { + model.RefreshPricing() + result.PricingRefreshed = true + common.SysLog(fmt.Sprintf( + "model channel availability sync: reason=%s disabled=%d enabled=%d", + reason, result.Disabled, result.Enabled, + )) + } else if reason != "" && forceFull && !result.Skipped { + common.SysLog(fmt.Sprintf( + "model channel availability sync: reason=%s disabled=0 enabled=0", + reason, + )) + } + return result, nil +} + +// MaybeSyncModelChannelAvailabilityAfterOptionChange triggers full calibration +// when either model availability switch is enabled. +func MaybeSyncModelChannelAvailabilityAfterOptionChange(key string, value string) error { + if key != "AutomaticDisableModelEnabled" && key != "AutomaticEnableModelEnabled" { + return nil + } + if value != "1" && !strings.EqualFold(value, "true") { + return nil + } + SyncModelChannelAvailabilityAfterMutation(fmt.Sprintf("option.%s=true", key)) + return nil +} + +// ManualDisableModelsWithoutChannels disables enabled models with no usable channel. +func ManualDisableModelsWithoutChannels() (ModelChannelAvailabilityResult, error) { + return reconcileModelChannelAvailability( + "manual.batch.disable.no-channels", + false, + model.ModelChannelAvailabilityConfig{Disable: true}, + ) +} + +// ManualEnableModelsWithChannels only restores models previously disabled by +// channel-availability automation. +func ManualEnableModelsWithChannels() (ModelChannelAvailabilityResult, error) { + return reconcileModelChannelAvailability( + "manual.batch.enable.with-channels", + false, + model.ModelChannelAvailabilityConfig{Enable: true}, + ) +} diff --git a/service/model_channel_availability_test.go b/service/model_channel_availability_test.go new file mode 100644 index 000000000000..fbd8cdda3782 --- /dev/null +++ b/service/model_channel_availability_test.go @@ -0,0 +1,610 @@ +package service + +import ( + "errors" + "fmt" + "sync/atomic" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/model" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "gorm.io/gorm" +) + +func resetModelChannelAvailabilityFixtures(t *testing.T) { + t.Helper() + for _, table := range []string{"abilities", "channels", "models", "vendors", "options"} { + require.NoError(t, model.DB.Exec("DELETE FROM "+table).Error) + } + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = false + model.InvalidatePricingCache() +} + +func createChannelWithModels(t *testing.T, id int, status int, modelsCSV string, abilityEnabled bool) { + t.Helper() + ch := &model.Channel{ + Id: id, + Type: 1, + Key: fmt.Sprintf("key-%d", id), + Status: status, + Name: fmt.Sprintf("channel-%d", id), + Models: modelsCSV, + Group: "default", + } + require.NoError(t, model.DB.Create(ch).Error) + for _, name := range splitCSV(modelsCSV) { + require.NoError(t, model.DB.Create(&model.Ability{ + Group: "default", + Model: name, + ChannelId: id, + Enabled: abilityEnabled && status == common.ChannelStatusEnabled, + }).Error) + } +} + +func createMetaModel(t *testing.T, id int, name string, status int, nameRule int, autoDisabled bool) { + t.Helper() + m := &model.Model{ + Id: id, + ModelName: name, + NameRule: nameRule, + SyncOfficial: 1, + // Status/AutoDisabledByRule may be zero values; force them after Create. + Status: 1, + } + require.NoError(t, model.DB.Create(m).Error) + require.NoError(t, model.DB.Model(&model.Model{}).Where("id = ?", id).Updates(map[string]interface{}{ + "status": status, + "auto_disabled_by_rule": autoDisabled, + }).Error) +} + +func splitCSV(s string) []string { + parts := make([]string, 0) + for _, p := range splitByComma(s) { + if p != "" { + parts = append(parts, p) + } + } + return parts +} + +func splitByComma(s string) []string { + out := make([]string, 0) + start := 0 + for i := 0; i <= len(s); i++ { + if i == len(s) || s[i] == ',' { + part := s[start:i] + // trim spaces + for len(part) > 0 && (part[0] == ' ' || part[0] == '\t') { + part = part[1:] + } + for len(part) > 0 && (part[len(part)-1] == ' ' || part[len(part)-1] == '\t') { + part = part[:len(part)-1] + } + out = append(out, part) + start = i + 1 + } + } + return out +} + +func loadModel(t *testing.T, id int) model.Model { + t.Helper() + var m model.Model + require.NoError(t, model.DB.First(&m, id).Error) + return m +} + +func requireModelChannelAvailabilitySync(t *testing.T, reason string) ModelChannelAvailabilityResult { + t.Helper() + result, err := SyncModelChannelAvailability(reason) + require.NoError(t, err) + return result +} + +func TestSyncModelChannelAvailability_ExactMatchLastChannelFails(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", true) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + // still available + res := requireModelChannelAvailabilitySync(t, "test") + assert.Equal(t, 0, res.Disabled) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) + + // disable only channel + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusManuallyDisabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", false).Error) + + res = requireModelChannelAvailabilitySync(t, "last-channel-down") + assert.Equal(t, 1, res.Disabled) + m := loadModel(t, 1) + assert.Equal(t, modelStatusDisabled, m.Status) + assert.True(t, m.AutoDisabledByRule) +} + +func TestSyncModelChannelAvailability_OtherChannelStillAvailable(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", true) + createChannelWithModels(t, 2, common.ChannelStatusEnabled, "gpt-4", true) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusAutoDisabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", false).Error) + + res := requireModelChannelAvailabilitySync(t, "other-channel-ok") + assert.Equal(t, 0, res.Disabled) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) +} + +func TestSyncModelChannelAvailability_SoftDeletedChannelNotAvailable(t *testing.T) { + // Project hard-deletes channels; treat deleted channel rows as unavailable by removing them. + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", true) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + require.NoError(t, model.DB.Where("id = ?", 1).Delete(&model.Channel{}).Error) + require.NoError(t, model.DB.Where("channel_id = ?", 1).Delete(&model.Ability{}).Error) + + res := requireModelChannelAvailabilitySync(t, "channel-deleted") + assert.Equal(t, 1, res.Disabled) + assert.True(t, loadModel(t, 1).AutoDisabledByRule) +} + +func TestSyncModelChannelAvailability_ManualDisableProtected(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + common.AutomaticEnableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", true) + // already disabled manually (no auto marker) + createMetaModel(t, 1, "gpt-4", modelStatusDisabled, model.NameRuleExact, false) + + // no channel available + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusManuallyDisabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", false).Error) + + res := requireModelChannelAvailabilitySync(t, "manual-disabled") + assert.Equal(t, 0, res.Disabled) + assert.Equal(t, 0, res.Enabled) + m := loadModel(t, 1) + assert.Equal(t, modelStatusDisabled, m.Status) + assert.False(t, m.AutoDisabledByRule) + + // restore channel; still should not auto-enable manual disable + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusEnabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", true).Error) + res = requireModelChannelAvailabilitySync(t, "channel-recovered") + assert.Equal(t, 0, res.Enabled) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) +} + +func TestSyncModelChannelAvailability_RecoverOnlyWhenEnableSwitchOn(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + common.AutomaticEnableModelEnabled = false + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", true) + createMetaModel(t, 1, "gpt-4", modelStatusDisabled, model.NameRuleExact, true) + + res := requireModelChannelAvailabilitySync(t, "enable-off") + assert.Equal(t, 0, res.Enabled) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) + assert.True(t, loadModel(t, 1).AutoDisabledByRule) + + common.AutomaticEnableModelEnabled = true + res = requireModelChannelAvailabilitySync(t, "enable-on") + assert.Equal(t, 1, res.Enabled) + m := loadModel(t, 1) + assert.Equal(t, modelStatusEnabled, m.Status) + assert.True(t, m.AutoDisabledByRule) +} + +func TestSyncModelChannelAvailability_RulePrefixMatch(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4-turbo", true) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRulePrefix, false) + + res := requireModelChannelAvailabilitySync(t, "prefix-ok") + assert.Equal(t, 0, res.Disabled) + + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusAutoDisabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", false).Error) + res = requireModelChannelAvailabilitySync(t, "prefix-down") + assert.Equal(t, 1, res.Disabled) +} + +func TestSyncModelChannelAvailability_RuleContainsAndSuffix(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "claude-3-opus", true) + createMetaModel(t, 1, "opus", modelStatusEnabled, model.NameRuleContains, false) + createMetaModel(t, 2, "-opus", modelStatusEnabled, model.NameRuleSuffix, false) + + res := requireModelChannelAvailabilitySync(t, "rules-ok") + assert.Equal(t, 0, res.Disabled) + + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusManuallyDisabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", false).Error) + res = requireModelChannelAvailabilitySync(t, "rules-down") + assert.Equal(t, 2, res.Disabled) +} + +func TestSyncModelChannelAvailability_MainSwitchOffKeepsMarker(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = false + + createMetaModel(t, 1, "gpt-4", modelStatusDisabled, model.NameRuleExact, true) + res := requireModelChannelAvailabilitySync(t, "both-off") + assert.True(t, res.Skipped) + m := loadModel(t, 1) + assert.Equal(t, modelStatusDisabled, m.Status) + assert.True(t, m.AutoDisabledByRule) +} + +func TestSyncModelChannelAvailability_IdempotentDisable(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + res1 := requireModelChannelAvailabilitySync(t, "first") + assert.Equal(t, 1, res1.Disabled) + res2 := requireModelChannelAvailabilitySync(t, "second") + assert.Equal(t, 0, res2.Disabled) + assert.True(t, loadModel(t, 1).AutoDisabledByRule) +} + +func TestModelUpdateClearsAutoDisabledMarkerWithStatusChange(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + createMetaModel(t, 1, "gpt-4", modelStatusDisabled, model.NameRuleExact, true) + + updated := loadModel(t, 1) + updated.Status = modelStatusEnabled + require.NoError(t, updated.Update()) + + persisted := loadModel(t, 1) + assert.Equal(t, modelStatusEnabled, persisted.Status) + assert.False(t, persisted.AutoDisabledByRule) +} + +func TestMaybeSyncModelChannelAvailabilityAfterOptionChange(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + require.NoError(t, MaybeSyncModelChannelAvailabilityAfterOptionChange("AutomaticDisableModelEnabled", "false")) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) + + require.NoError(t, MaybeSyncModelChannelAvailabilityAfterOptionChange("AutomaticDisableModelEnabled", "true")) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) + assert.True(t, loadModel(t, 1).AutoDisabledByRule) +} + +func TestSyncModelChannelAvailability_AbilityDisabledChannelEnabledNotAvailable(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", false) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + res := requireModelChannelAvailabilitySync(t, "ability-disabled") + assert.Equal(t, 1, res.Disabled) +} + +func TestSyncModelChannelAvailability_ChannelLifecycleIntegration(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + common.AutomaticEnableModelEnabled = true + + // create channel + model (channel.create) + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4,gpt-4-turbo", true) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + createMetaModel(t, 2, "gpt-4-", modelStatusEnabled, model.NameRulePrefix, false) + + res := requireModelChannelAvailabilitySync(t, "channel.create") + assert.Equal(t, 0, res.Disabled) + + // edit models list removes exact model binding (channel.update models) + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("models", "gpt-4-turbo").Error) + require.NoError(t, model.DB.Where("channel_id = ? AND model = ?", 1, "gpt-4").Delete(&model.Ability{}).Error) + res = requireModelChannelAvailabilitySync(t, "channel.update") + assert.Equal(t, 1, res.Disabled) + assert.True(t, loadModel(t, 1).AutoDisabledByRule) + assert.Equal(t, modelStatusEnabled, loadModel(t, 2).Status) // prefix still matches gpt-4-turbo + + // status disable channel (channel.status_update / tag / batch) + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusManuallyDisabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", false).Error) + res = requireModelChannelAvailabilitySync(t, "channel.status_update") + assert.Equal(t, 1, res.Disabled) // prefix model now also disabled + assert.True(t, loadModel(t, 2).AutoDisabledByRule) + + // system auto enable channel recovery + require.NoError(t, model.DB.Model(&model.Channel{}).Where("id = ?", 1).Update("status", common.ChannelStatusEnabled).Error) + require.NoError(t, model.DB.Model(&model.Ability{}).Where("channel_id = ?", 1).Update("enabled", true).Error) + res = requireModelChannelAvailabilitySync(t, "channel.auto_enable") + assert.Equal(t, 1, res.Enabled) // only models with available exact names recover (prefix) + assert.Equal(t, modelStatusEnabled, loadModel(t, 2).Status) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) // gpt-4 still no exact ability + + // re-add ability for gpt-4 via recreate ability then recover + require.NoError(t, model.DB.Create(&model.Ability{Group: "default", Model: "gpt-4", ChannelId: 1, Enabled: true}).Error) + res = requireModelChannelAvailabilitySync(t, "channel.update") + assert.Equal(t, 1, res.Enabled) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) + + // delete channel + require.NoError(t, model.DB.Where("id = ?", 1).Delete(&model.Channel{}).Error) + require.NoError(t, model.DB.Where("channel_id = ?", 1).Delete(&model.Ability{}).Error) + res = requireModelChannelAvailabilitySync(t, "channel.delete") + assert.Equal(t, 2, res.Disabled) +} + +func TestSyncModelChannelAvailability_FullCalibrationOnSwitch(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + // models without channels stay enabled until switch opens + createMetaModel(t, 1, "lonely-model", modelStatusEnabled, model.NameRuleExact, false) + createMetaModel(t, 2, "already-off", modelStatusDisabled, model.NameRuleExact, false) + + common.AutomaticDisableModelEnabled = false + res := requireModelChannelAvailabilitySync(t, "before") + assert.True(t, res.Skipped) + + common.AutomaticDisableModelEnabled = true + require.NoError(t, MaybeSyncModelChannelAvailabilityAfterOptionChange("AutomaticDisableModelEnabled", "true")) + m1 := loadModel(t, 1) + assert.Equal(t, modelStatusDisabled, m1.Status) + assert.True(t, m1.AutoDisabledByRule) + // already disabled without marker remains manual + m2 := loadModel(t, 2) + assert.Equal(t, modelStatusDisabled, m2.Status) + assert.False(t, m2.AutoDisabledByRule) +} + +func TestSyncModelChannelAvailability_EnableRequiresDisableSwitch(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + // enable alone must not recover models + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = true + + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4", true) + createMetaModel(t, 1, "gpt-4", modelStatusDisabled, model.NameRuleExact, true) + + res := requireModelChannelAvailabilitySync(t, "enable-without-disable") + assert.True(t, res.Skipped) + assert.Equal(t, 0, res.Enabled) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) + + common.AutomaticDisableModelEnabled = true + res = requireModelChannelAvailabilitySync(t, "enable-with-disable") + assert.Equal(t, 1, res.Enabled) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) +} + +func TestManualEnableModelsWithChannels_OnlyAutoDisabled(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + createChannelWithModels(t, 1, common.ChannelStatusEnabled, "gpt-4,claude-3", true) + // auto-disabled with channels recovered + createMetaModel(t, 1, "gpt-4", modelStatusDisabled, model.NameRuleExact, true) + // manually disabled with channels available + createMetaModel(t, 2, "claude-3", modelStatusDisabled, model.NameRuleExact, false) + // enabled with no channel; the manual enable action must not disable it. + createMetaModel(t, 3, "o3", modelStatusEnabled, model.NameRuleExact, false) + + res, err := ManualEnableModelsWithChannels() + require.NoError(t, err) + assert.Equal(t, 0, res.Disabled) + assert.Equal(t, 1, res.Enabled) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) + assert.True(t, loadModel(t, 1).AutoDisabledByRule) + assert.Equal(t, modelStatusDisabled, loadModel(t, 2).Status) + assert.False(t, loadModel(t, 2).AutoDisabledByRule) + assert.Equal(t, modelStatusEnabled, loadModel(t, 3).Status) +} + +func TestUpdateOptionPairsEnableOffWhenDisableOff(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + if common.OptionMap == nil { + common.OptionMap = map[string]string{} + } + require.NoError(t, model.UpdateOptionsBulk(map[string]string{ + "AutomaticDisableModelEnabled": "1", + "AutomaticEnableModelEnabled": "1", + })) + assert.True(t, common.AutomaticDisableModelEnabled) + assert.True(t, common.AutomaticEnableModelEnabled) + + require.NoError(t, model.UpdateOption("AutomaticDisableModelEnabled", "false")) + assert.False(t, common.AutomaticDisableModelEnabled) + assert.False(t, common.AutomaticEnableModelEnabled) + var options []model.Option + require.NoError(t, model.DB.Where("key IN ?", []string{ + "AutomaticDisableModelEnabled", + "AutomaticEnableModelEnabled", + }).Order("key").Find(&options).Error) + require.Len(t, options, 2) + assert.Equal(t, "false", options[0].Value) + assert.Equal(t, "false", options[1].Value) +} + +func TestUpdateOptionRejectsModelAutoEnableWithoutAutoDisable(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + if common.OptionMap == nil { + common.OptionMap = map[string]string{} + } + + require.NoError(t, model.UpdateOption("AutomaticEnableModelEnabled", "1")) + + assert.False(t, common.AutomaticDisableModelEnabled) + assert.False(t, common.AutomaticEnableModelEnabled) + var options []model.Option + require.NoError(t, model.DB.Where("key IN ?", []string{ + "AutomaticDisableModelEnabled", + "AutomaticEnableModelEnabled", + }).Order("key").Find(&options).Error) + require.Len(t, options, 2) + assert.Equal(t, "false", options[0].Value) + assert.Equal(t, "false", options[1].Value) +} + +func TestInitOptionMapNormalizesStoredModelAvailabilityPair(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + require.NoError(t, model.DB.Create(&[]model.Option{ + {Key: "AutomaticDisableModelEnabled", Value: "false"}, + {Key: "AutomaticEnableModelEnabled", Value: "1"}, + }).Error) + + model.InitOptionMap() + + assert.False(t, common.AutomaticDisableModelEnabled) + assert.False(t, common.AutomaticEnableModelEnabled) + var options []model.Option + require.NoError(t, model.DB.Where("key IN ?", []string{ + "AutomaticDisableModelEnabled", + "AutomaticEnableModelEnabled", + }).Order("key").Find(&options).Error) + require.Len(t, options, 2) + assert.Equal(t, "false", options[0].Value) + assert.Equal(t, "false", options[1].Value) +} + +func TestSyncModelChannelAvailability_MultiKeyRequiresEnabledKey(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + common.AutomaticEnableModelEnabled = true + + channel := &model.Channel{ + Id: 1, + Type: 1, + Key: "key-a\nkey-b", + Status: common.ChannelStatusEnabled, + Name: "multi-key-channel", + Models: "gpt-4", + Group: "default", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyStatusList: map[int]int{0: 2, 1: 3}, + }, + } + require.NoError(t, model.DB.Create(channel).Error) + require.NoError(t, model.DB.Create(&model.Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: true, + }).Error) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + result := requireModelChannelAvailabilitySync(t, "all-keys-disabled") + assert.Equal(t, 1, result.Disabled) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) + + delete(channel.ChannelInfo.MultiKeyStatusList, 0) + require.NoError(t, channel.SaveChannelInfo()) + result = requireModelChannelAvailabilitySync(t, "key-reenabled") + assert.Equal(t, 1, result.Enabled) + assert.Equal(t, modelStatusEnabled, loadModel(t, 1).Status) +} + +func TestSyncModelChannelAvailability_MultiKeyBlankEntriesAreUnavailable(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + channel := &model.Channel{ + Id: 1, + Type: 1, + Key: " \n\t", + Status: common.ChannelStatusEnabled, + Name: "blank-multi-key-channel", + Models: "gpt-4", + Group: "default", + ChannelInfo: model.ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + }, + } + require.NoError(t, model.DB.Create(channel).Error) + require.NoError(t, model.DB.Create(&model.Ability{ + Group: "default", Model: "gpt-4", ChannelId: channel.Id, Enabled: true, + }).Error) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + result := requireModelChannelAvailabilitySync(t, "blank-keys") + assert.Equal(t, 1, result.Disabled) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) +} + +func TestStartupCalibrationUsesPersistedOptions(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = false + common.AutomaticEnableModelEnabled = false + require.NoError(t, model.DB.Create(&[]model.Option{ + {Key: "AutomaticDisableModelEnabled", Value: "true"}, + {Key: "AutomaticEnableModelEnabled", Value: "false"}, + }).Error) + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + require.NoError(t, CalibrateModelChannelAvailabilityAtStartup()) + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) +} + +func TestSyncModelChannelAvailability_ReturnsDatabaseError(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + + forcedErr := errors.New("forced database failure") + callbackName := "test:fail-model-availability-query" + require.NoError(t, model.DB.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) { + tx.AddError(forcedErr) + })) + t.Cleanup(func() { _ = model.DB.Callback().Query().Remove(callbackName) }) + + _, err := SyncModelChannelAvailability("database-error") + require.Error(t, err) + assert.ErrorIs(t, err, forcedErr) +} + +func TestSyncModelChannelAvailabilityAfterMutationRetriesWithoutFailingCommittedOperation(t *testing.T) { + resetModelChannelAvailabilityFixtures(t) + common.AutomaticDisableModelEnabled = true + createMetaModel(t, 1, "gpt-4", modelStatusEnabled, model.NameRuleExact, false) + + var modelQueries atomic.Int32 + forcedErr := errors.New("transient model query failure") + callbackName := "test:fail-first-model-availability-query" + require.NoError(t, model.DB.Callback().Query().Before("gorm:query").Register(callbackName, func(tx *gorm.DB) { + if tx.Statement.Table == "models" && modelQueries.Add(1) == 1 { + tx.AddError(forcedErr) + } + })) + t.Cleanup(func() { _ = model.DB.Callback().Query().Remove(callbackName) }) + + result := SyncModelChannelAvailabilityAfterMutation("committed-model-create") + assert.Zero(t, result.Disabled) + assert.Equal(t, "gpt-4", loadModel(t, 1).ModelName) + retryDone := make(chan struct{}) + go func() { + modelChannelAvailabilityRetryPending.Wait() + close(retryDone) + }() + select { + case <-retryDone: + case <-time.After(3 * time.Second): + t.Fatal("timed out waiting for model availability retry") + } + assert.Equal(t, modelStatusDisabled, loadModel(t, 1).Status) + assert.GreaterOrEqual(t, modelQueries.Load(), int32(2)) +} diff --git a/service/task_billing_test.go b/service/task_billing_test.go index 699cc1ed67ae..6086c4e0f917 100644 --- a/service/task_billing_test.go +++ b/service/task_billing_test.go @@ -45,6 +45,10 @@ func TestMain(m *testing.M) { &model.Token{}, &model.Log{}, &model.Channel{}, + &model.Ability{}, + &model.Model{}, + &model.Vendor{}, + &model.Option{}, &model.Midjourney{}, &model.TopUp{}, &model.UserSubscription{}, diff --git a/web/src/features/channels/api.ts b/web/src/features/channels/api.ts index 2818242d08c9..a0a358422caf 100644 --- a/web/src/features/channels/api.ts +++ b/web/src/features/channels/api.ts @@ -160,7 +160,11 @@ export async function updateChannelStatus( export async function batchUpdateChannelStatus( ids: number[], status: number -): Promise<{ success: boolean; message?: string; data?: number }> { +): Promise<{ + success: boolean + message?: string + data?: number | { changed: number; failed_ids: number[] } +}> { const res = await api.post( '/api/channel/status/batch', { ids, status }, diff --git a/web/src/features/channels/components/channels-provider.tsx b/web/src/features/channels/components/channels-provider.tsx index 72d212a81446..5e496430ac92 100644 --- a/web/src/features/channels/components/channels-provider.tsx +++ b/web/src/features/channels/components/channels-provider.tsx @@ -27,7 +27,7 @@ import React, { } from 'react' import { useChannelUpstreamUpdates } from '../hooks/use-channel-upstream-updates' -import { channelsQueryKeys } from '../lib' +import { invalidateChannelMutationQueries } from '../lib' import type { Channel } from '../types' // ============================================================================ @@ -94,7 +94,7 @@ export function ChannelsProvider({ children }: { children: React.ReactNode }) { const queryClient = useQueryClient() const refreshChannels = useCallback(async () => { - await queryClient.invalidateQueries({ queryKey: channelsQueryKeys.all }) + await invalidateChannelMutationQueries(queryClient) }, [queryClient]) const upstream = useChannelUpstreamUpdates(refreshChannels) diff --git a/web/src/features/channels/components/data-table-row-actions.tsx b/web/src/features/channels/components/data-table-row-actions.tsx index 7e15463cc97d..5693bbcdca17 100644 --- a/web/src/features/channels/components/data-table-row-actions.tsx +++ b/web/src/features/channels/components/data-table-row-actions.tsx @@ -61,10 +61,10 @@ import { useAuthStore } from '@/stores/auth-store' import { MODEL_FETCHABLE_TYPES } from '../constants' import { - channelsQueryKeys, handleDeleteChannel, handleTestChannel, handleToggleChannelStatus, + invalidateChannelMutationQueries, isChannelEnabled, isMultiKeyChannel, } from '../lib' @@ -111,7 +111,7 @@ export function DataTableRowActions({ row }: DataTableRowActionsProps) { setIsTesting(true) try { await handleTestChannel(channel.id, { channelName: channel.name }, () => { - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + void invalidateChannelMutationQueries(queryClient) }) } finally { setIsTesting(false) diff --git a/web/src/features/channels/components/dialogs/channel-test-dialog.tsx b/web/src/features/channels/components/dialogs/channel-test-dialog.tsx index a2d641d22c21..da1684388674 100644 --- a/web/src/features/channels/components/dialogs/channel-test-dialog.tsx +++ b/web/src/features/channels/components/dialogs/channel-test-dialog.tsx @@ -92,6 +92,7 @@ import { channelsQueryKeys, formatResponseTime, handleTestChannel, + invalidateChannelMutationQueries, } from '../../lib' import type { Channel, @@ -534,8 +535,7 @@ function ChannelTestDialogContent({ const refreshChannelLists = useCallback( (patch?: ChannelTestCachePatch) => { updateChannelTestCache(patch) - void queryClient - .invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + void invalidateChannelMutationQueries(queryClient) .then(() => updateChannelTestCache(patch)) .catch(() => undefined) }, diff --git a/web/src/features/channels/components/dialogs/edit-tag-dialog.tsx b/web/src/features/channels/components/dialogs/edit-tag-dialog.tsx index f8be304c14c6..40172b1b41f4 100644 --- a/web/src/features/channels/components/dialogs/edit-tag-dialog.tsx +++ b/web/src/features/channels/components/dialogs/edit-tag-dialog.tsx @@ -46,7 +46,7 @@ import { getAllModels, getGroups, } from '../../api' -import { channelsQueryKeys } from '../../lib' +import { invalidateChannelMutationQueries } from '../../lib' import type { TagOperationParams } from '../../types' import { useChannels } from '../channels-provider' @@ -196,7 +196,7 @@ export function EditTagDialog({ open, onOpenChange }: EditTagDialogProps) { if (response.success) { toast.success(t('Tag updated successfully')) - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + void invalidateChannelMutationQueries(queryClient) onOpenChange(false) } else { toast.error(response.message || t('Failed to update tag')) @@ -304,12 +304,10 @@ export function EditTagDialog({ open, onOpenChange }: EditTagDialogProps) {
- items={[ - ...availableModels.map((model) => ({ - value: model, - label: model, - })), - ]} + items={availableModels.map((model) => ({ + value: model, + label: model, + }))} onValueChange={(value) => { if (value === null) return if (!selectedModels.includes(value)) { diff --git a/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx b/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx index 6d7ce554eefa..330c951083a9 100644 --- a/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx +++ b/web/src/features/channels/components/dialogs/fetch-models-dialog.tsx @@ -43,7 +43,7 @@ import { fetchUpstreamModels, updateChannel } from '../../api' import { categorizeModels, categorizeModelsWithRedirect, - channelsQueryKeys, + invalidateChannelMutationQueries, normalizeModelName, parseModelsString, } from '../../lib' @@ -180,7 +180,7 @@ export function FetchModelsDialog({ }) if (response.success) { toast.success(t('Models updated successfully')) - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + void invalidateChannelMutationQueries(queryClient) onOpenChange(false) } else { toast.error(response.message || t('Failed to update models')) diff --git a/web/src/features/channels/components/dialogs/multi-key-manage-dialog.tsx b/web/src/features/channels/components/dialogs/multi-key-manage-dialog.tsx index 22f7521496a6..8705fdd452c9 100644 --- a/web/src/features/channels/components/dialogs/multi-key-manage-dialog.tsx +++ b/web/src/features/channels/components/dialogs/multi-key-manage-dialog.tsx @@ -54,11 +54,11 @@ import { } from '../../api' import { MULTI_KEY_FILTER_OPTIONS } from '../../constants' import { - channelsQueryKeys, formatTimestamp, getMultiKeyStatusConfig, getMultiKeyConfirmMessage, isDestructiveAction, + invalidateChannelMutationQueries, } from '../../lib' import type { KeyStatus, MultiKeyConfirmAction } from '../../types' import { useChannels } from '../channels-provider' @@ -149,7 +149,7 @@ export function MultiKeyManageDialog({ } const handleStatusFilterChange = (value: string) => { - const newFilter = value === 'all' ? null : parseInt(value) + const newFilter = value === 'all' ? null : Number.parseInt(value) setStatusFilter(newFilter) setCurrentPage(1) loadKeyStatus(1, pageSize, newFilter) @@ -193,7 +193,7 @@ export function MultiKeyManageDialog({ if (response?.success) { toast.success(response.message || t('Operation successful')) - queryClient.invalidateQueries({ queryKey: channelsQueryKeys.lists() }) + void invalidateChannelMutationQueries(queryClient) // Reload data - reset to page 1 for bulk actions const isBulkAction = type.includes('all') || type === 'delete-disabled' @@ -294,12 +294,10 @@ export function MultiKeyManageDialog({ {/* Toolbar */}
setNewTag(e.target.value)} - disabled={isSaving} - /> -

- {t('Leave empty to disband the tag')} -

-
- - {/* Models */} -
- -