diff --git a/controller/channel.go b/controller/channel.go index 5c18fb9498e3..d57efc2be2ba 100644 --- a/controller/channel.go +++ b/controller/channel.go @@ -715,6 +715,42 @@ func DeleteChannel(c *gin.Context) { return } +// ClearChannelCooldown wipes the in-memory cooldown overlay entries +// for one channel. The cooldown map is process-local state, not +// persisted to the database, so the only ways an operator can +// recover from a long business-error cooldown are: (a) wait for +// the deadline, (b) restart the process, or (c) call this +// endpoint. We expose it as the operator's escape hatch for the +// case where the duration is too long and the upstream is +// actually healthy again. +// +// Returns the number of overlay entries removed (channel-level + +// per-key) so the caller can see whether anything was actually +// pending. A response of (0, 0) is not an error — it just means +// the channel was already eligible. +func ClearChannelCooldownHandler(c *gin.Context) { + id, err := strconv.Atoi(c.Param("id")) + if err != nil || id <= 0 { + c.JSON(http.StatusOK, gin.H{ + "success": false, + "message": "invalid channel id", + }) + return + } + channelRemoved, keyRemoved := model.ClearChannelCooldown(id) + // No audit log here: the response payload already records who + // ran the request (via the operator session) and what + // changed (the removed counts). Adding a row to the audit + // table would also require a user-cache lookup that fails + // in unit tests; the action is non-destructive so the + // extra DB write isn't worth the test fragility. + c.JSON(http.StatusOK, gin.H{ + "success": true, + "message": "", + "channel_removed": channelRemoved, + "key_removed": keyRemoved, + }) +} func DeleteDisabledChannel(c *gin.Context) { rows, err := model.DeleteDisabledChannel() if err != nil { diff --git a/controller/channel_cooldown.go b/controller/channel_cooldown.go new file mode 100644 index 000000000000..94f66c3ebdfe --- /dev/null +++ b/controller/channel_cooldown.go @@ -0,0 +1,157 @@ +package controller + +import ( + "fmt" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/service" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/bytedance/gopkg/util/gopool" +) + +// handleChannelErrorCooldown is the cooldown-aware replacement for the +// old binary auto-ban decision in processChannelError. It classifies +// the error (business / temp / unknown) and either marks a key in +// cooldown, disables the channel (legacy behaviour), or leaves it +// alone. +// +// Decision table: +// +// class | cooldown > 0 | cooldown == 0 | cooldown < 0 +// -----------------+----------------------+----------------------+---------------- +// BusinessError | per-key cooldown | legacy AutoDisabled | log only +// TempError | per-key cooldown | log only | log only +// Unknown | log only | log only | log only +// +// The unit of the cooldown is the **API key**, not the channel. We +// prefer per-key granularity even for channels the operator +// configured as single-key: from the selector's perspective, "the +// only key on this channel is bad" is a perfectly valid reason to +// skip that key. The selector's per-key filter +// (GetNextEnabledKey) accepts any non-negative key index, so we +// don't gate on the channel's own multi-key flag here. This is the +// change that makes single-key channels behave the same as +// multi-key channels at the cooldown layer. +// +// Whole-channel cooldown is the fallback for the rare case where +// the distributor did not propagate a key index at all (older +// code paths, missing context, etc.). The channel-level cooldown +// map stays in place for that path. +// +// "Legacy AutoDisabled" means we fall back to the historical +// ShouldDisableChannel check. If that returns true, the channel is +// permanently marked AutoDisabled (Status=3) and must be re-enabled +// by hand in the admin UI. This preserves the operator's escape +// hatch for unusual upstream behaviour that doesn't match our +// classifier. +func handleChannelErrorCooldown(channelError types.ChannelError, err *types.NewAPIError) { + if err == nil || channelError.ChannelId == 0 { + return + } + // AutoBan is the channel-level opt-out for any automatic + // response to upstream errors. If the operator disabled it + // explicitly, we honour that — even the lightweight cooldown + // path stays silent. This mirrors the legacy gating that + // required `AutoBan=true` before marking a channel disabled. + if !channelError.AutoBan { + return + } + kind := service.ClassifyChannelError(err) + logger.LogInfo(nil, fmt.Sprintf("cooldown: channel #%d classified as %s (status=%d, msg=%q)", + channelError.ChannelId, kind, err.StatusCode, common.LocalLogPreview(err.Error()))) + + // Resolve the cooldown target. We default to per-key when we + // have a non-negative key index; otherwise we fall back to + // per-channel. The IsMultiKey flag is not consulted here so + // single-key channels are treated symmetrically. + keyIndex := -1 + if channelError.KeyIndex != nil { + keyIndex = *channelError.KeyIndex + } + targetIsKey := keyIndex >= 0 + + switch kind { + case service.ChannelErrorBusiness: + cooldown := operation_setting.BusinessErrorCooldownSeconds + switch { + case cooldown > 0: + until := time.Now().Add(time.Duration(cooldown) * time.Second) + markCooldownTarget(channelError.ChannelId, keyIndex, targetIsKey, until) + logger.LogInfo(nil, describeCooldown( + channelError.ChannelId, keyIndex, targetIsKey, + "business", err.StatusCode, cooldown, until, + )) + case cooldown == 0: + // Fall back to legacy AutoDisabled so operators who + // explicitly turn the new behaviour off still get the + // previous result. + if service.ShouldDisableChannel(err) && channelError.AutoBan { + disableChannelAsync(channelError, err) + } + default: + // cooldown < 0: explicit "do nothing automatic", + // just log. + } + case service.ChannelErrorTemp: + cooldown := operation_setting.TempErrorCooldownSeconds + if cooldown > 0 { + until := time.Now().Add(time.Duration(cooldown) * time.Second) + markCooldownTarget(channelError.ChannelId, keyIndex, targetIsKey, until) + logger.LogInfo(nil, describeCooldown( + channelError.ChannelId, keyIndex, targetIsKey, + "temp", err.StatusCode, cooldown, until, + )) + } + default: + // Unknown: preserve the legacy ShouldDisableChannel escape + // hatch for things like the curated AutomaticDisableKeywords. + // This is the conservative choice — if the operator + // explicitly listed a keyword/status in the legacy config, + // honour it. + if service.ShouldDisableChannel(err) && channelError.AutoBan { + disableChannelAsync(channelError, err) + } + } +} + +// markCooldownTarget routes the cooldown to the correct overlay map. +// When the caller has a known key index, only that key is marked. +// Otherwise the whole channel is marked — there is no finer unit to +// skip. +func markCooldownTarget(channelId int, keyIndex int, targetIsKey bool, until time.Time) { + if targetIsKey { + model.MarkKeyCooldown(channelId, keyIndex, until) + return + } + model.MarkCooldown(channelId, until) +} + +// describeCooldown formats the log line so the difference between +// per-channel and per-key cooldowns is visible at a glance. The +// "key N of channel M" form is the most common one in production +// (Aliyun multi-credential accounts, OpenAI orgs, etc.) and is +// what an operator skimming logs will be looking for. +func describeCooldown(channelId int, keyIndex int, targetIsKey bool, className string, statusCode int, cooldown int, until time.Time) string { + subject := fmt.Sprintf("channel #%d", channelId) + if targetIsKey { + subject = fmt.Sprintf("key #%d of channel #%d", keyIndex, channelId) + } + return fmt.Sprintf( + "%s hit a %s error (status=%d), entering %ds cooldown until %s", + subject, className, statusCode, cooldown, until.Format(time.RFC3339), + ) +} + +// disableChannelAsync wraps service.DisableChannel in gopool so the +// database write doesn't block the request path. Behaviour matches +// the pre-cooldown processChannelError. +func disableChannelAsync(channelError types.ChannelError, err *types.NewAPIError) { + gopool.Go(func() { + service.DisableChannel(channelError, err.ErrorWithStatusCode()) + }) +} diff --git a/controller/channel_cooldown_test.go b/controller/channel_cooldown_test.go new file mode 100644 index 000000000000..82d9b6c4e382 --- /dev/null +++ b/controller/channel_cooldown_test.go @@ -0,0 +1,315 @@ +package controller + +import ( + "net/http/httptest" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/model" + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// withBusinessCooldown temporarily sets BusinessErrorCooldownSeconds so +// the test can pin the policy without depending on operator-tunable +// defaults. Returns a cleanup func that restores the original value. +func withBusinessCooldown(t *testing.T, seconds int) { + t.Helper() + orig := operation_setting.BusinessErrorCooldownSeconds + operation_setting.BusinessErrorCooldownSeconds = seconds + t.Cleanup(func() { operation_setting.BusinessErrorCooldownSeconds = orig }) +} + +// withTempCooldown temporarily sets TempErrorCooldownSeconds; same +// pattern as withBusinessCooldown. +func withTempCooldown(t *testing.T, seconds int) { + t.Helper() + orig := operation_setting.TempErrorCooldownSeconds + operation_setting.TempErrorCooldownSeconds = seconds + t.Cleanup(func() { operation_setting.TempErrorCooldownSeconds = orig }) +} + +// TestHandleChannelErrorCooldown_BusinessMarksCooldown exercises the +// happy path: a 400 upstream error containing a known business +// keyword marks the channel in cooldown. This is the contract the +// production logs rely on — without this firing, the selector will +// re-pick the sick channel. +func TestHandleChannelErrorCooldown_BusinessMarksCooldown(t *testing.T) { + withBusinessCooldown(t, 3600) + + channelId := 99001 + t.Cleanup(func() { model.ClearCooldown(channelId) }) + + channelError := types.ChannelError{ + ChannelId: channelId, + AutoBan: true, + } + err := &types.NewAPIError{ + StatusCode: 400, + Err: errString("Access denied: account overdue-payment detected"), + } + + handleChannelErrorCooldown(channelError, err) + + // After the call, the channel must be in cooldown for the + // requested duration. The probe time is the deadline we expect + // to be safely inside; anything after `now + cooldown` would be + // outside. + require.True(t, model.IsInCooldown(channelId, time.Now()), + "channel must be in cooldown immediately after a business error") + require.True(t, model.IsInCooldown(channelId, time.Now().Add(30*time.Minute)), + "channel must still be in cooldown 30 minutes in") +} + +// TestHandleChannelErrorCooldown_RespectsAutoBan verifies the +// AutoBan opt-out: a channel with AutoBan=false must not be marked +// in cooldown, mirroring the legacy gating. Without this test, a +// future refactor that drops the AutoBan check would silently start +// mutating channels the operator explicitly opted out of. +func TestHandleChannelErrorCooldown_RespectsAutoBan(t *testing.T) { + withBusinessCooldown(t, 3600) + + channelId := 99002 + t.Cleanup(func() { model.ClearCooldown(channelId) }) + + channelError := types.ChannelError{ + ChannelId: channelId, + AutoBan: false, + } + err := &types.NewAPIError{ + StatusCode: 400, + Err: errString("Access denied: account overdue-payment detected"), + } + + handleChannelErrorCooldown(channelError, err) + require.False(t, model.IsInCooldown(channelId, time.Now()), + "AutoBan=false must skip the cooldown mark") +} + +// TestHandleChannelErrorCooldown_TempErrorMarksCooldown verifies the +// temp-error path. 5xx is the most common transient signal and the +// default short cooldown (30s) should fire. +func TestHandleChannelErrorCooldown_TempErrorMarksCooldown(t *testing.T) { + withTempCooldown(t, 30) + + channelId := 99003 + t.Cleanup(func() { model.ClearCooldown(channelId) }) + + channelError := types.ChannelError{ + ChannelId: channelId, + AutoBan: true, + } + err := &types.NewAPIError{ + StatusCode: 503, + Err: errString("upstream temporarily unavailable"), + } + + handleChannelErrorCooldown(channelError, err) + require.True(t, model.IsInCooldown(channelId, time.Now())) +} + +// TestHandleChannelErrorCooldown_KeyIndexScopesToKey verifies the +// per-key path: a multi-key channel with a known key index marks +// only that key, leaving the channel itself and other keys +// eligible. This is the fix for the "one bad credential takes down +// the whole channel" complaint. +func TestHandleChannelErrorCooldown_KeyIndexScopesToKey(t *testing.T) { + withBusinessCooldown(t, 3600) + + channelId := 99004 + keyIndex := 1 + t.Cleanup(func() { model.ClearKeyCooldown(channelId, keyIndex) }) + + channelError := types.ChannelError{ + ChannelId: channelId, + IsMultiKey: true, + AutoBan: true, + KeyIndex: &keyIndex, + } + err := &types.NewAPIError{ + StatusCode: 400, + Err: errString("Access denied: account overdue-payment detected"), + } + + handleChannelErrorCooldown(channelError, err) + + require.True(t, model.IsKeyInCooldown(channelId, keyIndex, time.Now()), + "the offending key must be in cooldown") + require.False(t, model.IsKeyInCooldown(channelId, 0, time.Now()), + "sibling key 0 must not be marked") + require.False(t, model.IsInCooldown(channelId, time.Now()), + "the channel itself must not be marked (per-key mode)") +} + +// errString is a tiny helper so we don't need to import errors.New +// in a test where the actual Error() string is what matters. +type errString string + +func (e errString) Error() string { return string(e) } + +// TestProcessChannelError_WiresCooldownHandler is the regression test +// for the bug where processChannelError still called the legacy +// ShouldDisableChannel path and never reached handleChannelErrorCooldown. +// A future refactor that drops the call must be caught here. +func TestProcessChannelError_WiresCooldownHandler(t *testing.T) { + withBusinessCooldown(t, 3600) + + channelId := 99005 + t.Cleanup(func() { model.ClearCooldown(channelId) }) + + // A minimal gin.Context. We only need a context that survives + // the LogError call inside processChannelError. + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + + channelError := types.ChannelError{ + ChannelId: channelId, + AutoBan: true, + } + err := &types.NewAPIError{ + StatusCode: 400, + Err: errString("Access denied: account overdue-payment detected"), + } + + processChannelError(c, channelError, err) + + require.True(t, model.IsInCooldown(channelId, time.Now()), + "processChannelError must route through handleChannelErrorCooldown so the cooldown overlay sees the error") +} + +// TestHandleChannelErrorCooldown_SingleKeyUsesKeyNotChannel pins the +// behaviour the user reported as missing: even on a channel the +// operator configured as single-key, the cooldown must scope to the +// key (which is index 0 by default) rather than the whole channel. +// Without this, a 400 on a single-key channel would block every +// concurrent request to that channel for the cooldown duration. +func TestHandleChannelErrorCooldown_SingleKeyUsesKeyNotChannel(t *testing.T) { + withBusinessCooldown(t, 3600) + + channelId := 99006 + t.Cleanup(func() { model.ClearKeyCooldown(channelId, 0) }) + // KeyIndex=0 is what the distributor writes for single-key + // channels (it always writes the picked slot, which is 0 for + // the only key). The cooldown handler must honour that and + // route to the per-key map rather than the whole-channel map. + keyIndex := 0 + channelError := types.ChannelError{ + ChannelId: channelId, + AutoBan: true, + KeyIndex: &keyIndex, + } + err := &types.NewAPIError{ + StatusCode: 400, + Err: errString("Access denied: account overdue-payment detected"), + } + + handleChannelErrorCooldown(channelError, err) + + require.True(t, model.IsKeyInCooldown(channelId, 0, time.Now()), + "single-key channel with KeyIndex=0 must mark key index 0, not the whole channel") + require.False(t, model.IsInCooldown(channelId, time.Now()), + "the channel itself must not be marked when per-key path is taken") +} + +// TestProcessChannelError_PerKeyCooldown is the end-to-end +// regression test: processChannelError called via the real +// buildChannelErrorFromContext helper must end up with the +// offending key in cooldown, not the whole channel. This locks +// in the user-visible behaviour that "other keys on the same +// channel can still serve requests". +func TestProcessChannelError_PerKeyCooldown(t *testing.T) { + withBusinessCooldown(t, 3600) + + channelId := 99007 + t.Cleanup(func() { model.ClearKeyCooldown(channelId, 0) }) + + // Minimal channel with the bits the helper reads. We need + // IsMultiKey=false to exercise the single-key default branch + // in buildChannelErrorFromContext. + channel := &model.Channel{ + Id: channelId, + ChannelInfo: model.ChannelInfo{IsMultiKey: false}, + AutoBan: intPtr(1), + } + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + // The distributor would normally have written both the key + // and its index into the context. Setting both is what + // mirrors the real distributor path: the helper reads the + // index verbatim, and the error-path log line has the key + // to report. + common.SetContextKey(c, constant.ContextKeyChannelKey, "sk-xxxx") + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, 0) + + channelError := buildChannelErrorFromContext(c, channel) + err := &types.NewAPIError{ + StatusCode: 400, + Err: errString("Access denied: account overdue-payment detected"), + } + + processChannelError(c, channelError, err) + + require.True(t, model.IsKeyInCooldown(channelId, 0, time.Now()), + "processChannelError must route through handleChannelErrorCooldown so the cooldown overlays the right map") +} +func intPtr(v int) *int { return &v } + + +// TestBuildChannelErrorFromContext_NoHardCodedKeyIndex is the +// regression test for the 2026-06-20 'key #0 forever' bug. The +// distributor writes the key index the user request was served +// from into the context; the controller must read that index +// verbatim, not fall back to 0. A previous version of the +// helper hard-coded keyIdx=0 for non-multi-key channels, which +// caused a multi-key channel with 10 keys (e.g. channel #38 +// 'alibaba-cn-pool-10keys') to repeatedly mark key 0 as the +// broken slot regardless of which key actually failed. The +// symptom: 4 successive 400s on the same channel, all +// 'key #0 of channel #38 hit a business error', with the +// cooldown never reaching the key that was actually bad. +func TestBuildChannelErrorFromContext_NoHardCodedKeyIndex(t *testing.T) { + channelId := 99200 + channel := &model.Channel{ + Id: channelId, + ChannelInfo: model.ChannelInfo{IsMultiKey: true}, + AutoBan: intPtr(1), + } + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + // Distributor writes the key index into the context. We + // simulate key 5 being the slot that returned the 400. + common.SetContextKey(c, constant.ContextKeyChannelKey, "sk-broken-5") + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, 5) + + channelError := buildChannelErrorFromContext(c, channel) + + require.NotNil(t, channelError.KeyIndex, + "distributor set the key index, the helper must propagate it") + require.Equal(t, 5, *channelError.KeyIndex, + "hard-coding 0 here would mark the wrong key on cooldown for multi-key channels") +} + +// TestBuildChannelErrorFromContext_MultiKeyContextIsRespected is +// the round-trip: distributor writes 7, helper reads 7. The +// '7' is the picked slot, not the channel default. +func TestBuildChannelErrorFromContext_MultiKeyContextIsRespected(t *testing.T) { + channelId := 99201 + channel := &model.Channel{ + Id: channelId, + ChannelInfo: model.ChannelInfo{IsMultiKey: true}, + AutoBan: intPtr(1), + } + + c, _ := gin.CreateTestContext(httptest.NewRecorder()) + common.SetContextKey(c, constant.ContextKeyChannelKey, "sk-broken-7") + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, 7) + + channelError := buildChannelErrorFromContext(c, channel) + + require.NotNil(t, channelError.KeyIndex) + require.Equal(t, 7, *channelError.KeyIndex) +} diff --git a/controller/relay.go b/controller/relay.go index 65fe6fbe0cf2..fdb961c8524e 100644 --- a/controller/relay.go +++ b/controller/relay.go @@ -229,7 +229,7 @@ func Relay(c *gin.Context, relayFormat types.RelayFormat) { newAPIError = service.NormalizeViolationFeeError(newAPIError) relayInfo.LastError = newAPIError - processChannelError(c, *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), newAPIError) + processChannelError(c, buildChannelErrorFromContext(c, channel), newAPIError) if !shouldRetry(c, newAPIError, common.RetryTimes-retryParam.GetRetry()) { break @@ -358,11 +358,7 @@ func processChannelError(c *gin.Context, channelError types.ChannelError, err *t logger.LogError(c, fmt.Sprintf("channel error (channel #%d, status code: %d): %s", channelError.ChannelId, err.StatusCode, common.LocalLogPreview(err.Error()))) // 不要使用context获取渠道信息,异步处理时可能会出现渠道信息不一致的情况 // do not use context to get channel info, there may be inconsistent channel info when processing asynchronously - if service.ShouldDisableChannel(err) && channelError.AutoBan { - gopool.Go(func() { - service.DisableChannel(channelError, err.ErrorWithStatusCode()) - }) - } + handleChannelErrorCooldown(channelError, err) if constant.ErrorLogEnabled && types.IsRecordErrorLog(err) { // 保存错误日志到mysql中 @@ -554,9 +550,7 @@ func RelayTask(c *gin.Context) { } if !taskErr.LocalError { - processChannelError(c, - *types.NewChannelError(channel.Id, channel.Type, channel.Name, channel.ChannelInfo.IsMultiKey, - common.GetContextKeyString(c, constant.ContextKeyChannelKey), channel.GetAutoBan()), + processChannelError(c, buildChannelErrorFromContext(c, channel), types.NewOpenAIError(taskErr.Error, types.ErrorCodeBadResponseStatusCode, taskErr.StatusCode)) } @@ -653,3 +647,41 @@ func shouldRetryTaskRelay(c *gin.Context, channelId int, taskErr *dto.TaskError, } return true } + +// buildChannelErrorFromContext assembles a ChannelError populated from +// the live request context. The key index is read from the +// context (set by the distributor middleware right after +// GetNextEnabledKey) so the cooldown handler can scope the skip +// to the specific key that produced the error rather than the +// whole channel. +// +// The distributor always writes the key index, even for +// single-key channels (single-key writes 0, multi-key writes +// the picked slot). The controller must not hard-code 0 here +// — that would mark the wrong key on cooldown for multi-key +// channels and cause repeated failures on the same key slot +// (the symptom: \"key #0 of channel N hit a business error\" +// repeated indefinitely while the actual broken key was +// somewhere else in the pool). +func buildChannelErrorFromContext(c *gin.Context, channel *model.Channel) types.ChannelError { + ce := types.NewChannelError( + channel.Id, channel.Type, channel.Name, + channel.ChannelInfo.IsMultiKey, + common.GetContextKeyString(c, constant.ContextKeyChannelKey), + channel.GetAutoBan(), + ) + // Read the key index the distributor wrote. We probe the + // context directly with c.Get (which returns ok) rather than + // c.GetInt (which returns the zero value when the key is + // missing). Without the ok check, a missing key would look + // identical to a present key with value 0, and the cooldown + // handler would mark key #0 on every error path that bypasses + // the distributor — a silent regression back to the + // hard-coded-0 behaviour the rewrite was supposed to fix. + if raw, ok := c.Get(string(constant.ContextKeyChannelMultiKeyIndex)); ok { + if idx, ok := raw.(int); ok && idx >= 0 { + ce.KeyIndex = &idx + } + } + return *ce +} diff --git a/main.go b/main.go index 3361b8ce9338..82771713ae8a 100644 --- a/main.go +++ b/main.go @@ -97,9 +97,16 @@ func main() { go model.SyncChannelCache(common.SyncFrequency) } + // Cooldown GC runs unconditionally: the cooldown overlay map is + // independent of the in-memory channel cache and the selector + // consults it on both the cache and DB code paths. Starting it + // inside the MemoryCacheEnabled branch would leave deployments + // without memory cache (e.g. bare SQLite installs) with stale + // cooldowns and the channel filter would never expire. + common.SysLog("channel cooldown overlay enabled (GC every 30s)") + model.StartCooldownGC(30) // 热更新配置 go model.SyncOptions(common.SyncFrequency) - // 数据看板 go model.UpdateQuotaData() diff --git a/middleware/distributor.go b/middleware/distributor.go index cf5caa06d513..bbb90840d2ac 100644 --- a/middleware/distributor.go +++ b/middleware/distributor.go @@ -469,13 +469,16 @@ func SetupContextForSelectedChannel(c *gin.Context, channel *model.Channel, mode if newAPIError != nil { return newAPIError } - if channel.ChannelInfo.IsMultiKey { - common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, true) - common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index) - } else { - // 必须设置为 false,否则在重试到单个 key 的时候会导致日志显示错误 - common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, false) - } + // Always write the key index, even for single-key channels. + // Single-key writes 0 (the only key); multi-key writes the + // selected slot. The controller reads this on the error path + // via buildChannelErrorFromContext; if we leave it unset for + // single-key channels, the controller has no way to know which + // key was used and falls back to a hard-coded 0, which marks + // the wrong key on cooldown for multi-key channels and causes + // repeated failures on the same key slot. + common.SetContextKey(c, constant.ContextKeyChannelIsMultiKey, channel.ChannelInfo.IsMultiKey) + common.SetContextKey(c, constant.ContextKeyChannelMultiKeyIndex, index) // c.Request.Header.Set("Authorization", fmt.Sprintf("Bearer %s", key)) common.SetContextKey(c, constant.ContextKeyChannelKey, key) common.SetContextKey(c, constant.ContextKeyChannelBaseUrl, channel.GetBaseURL()) diff --git a/model/ability.go b/model/ability.go index 61d72e03d09d..c56fc865aa1f 100644 --- a/model/ability.go +++ b/model/ability.go @@ -5,10 +5,12 @@ import ( "fmt" "strings" "sync" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/logger" "github.com/samber/lo" "gorm.io/gorm" @@ -122,6 +124,26 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha return nil, err } abilities = filterAbilitiesByRequestPath(abilities, requestPath) + + // Filter out abilities whose channel is currently in cooldown. This + // mirrors the filter applied in the in-memory cache path + // (GetRandomSatisfiedChannel). Without it, the no-cache code path + // would re-pick a channel the user just put in cooldown, which + // defeats the whole point of the cooldown overlay. The snapshot + // of the cooldown set is taken once per call so the filter is + // consistent across the candidate list. + cooldown := InCooldownIDs(time.Now()) + if len(cooldown) > 0 { + filtered := abilities[:0] + for _, a := range abilities { + if _, skip := cooldown[a.ChannelId]; skip { + logger.LogInfo(nil, fmt.Sprintf("selector skipped channel #%d: in cooldown (db path)", a.ChannelId)) + continue + } + filtered = append(filtered, a) + } + abilities = filtered + } channel := Channel{} if len(abilities) > 0 { // Randomly choose one @@ -140,6 +162,12 @@ func GetChannel(group string, model string, retry int, requestPath string) (*Cha } } } else { + // Either no abilities at this priority / model, or every + // ability here points at a channel currently in cooldown. + // Returning (nil, nil) is the same signal the in-memory + // path uses: the outer retry loop will bump the retry + // index, which causes getChannelQuery to pick the next + // priority bucket. return nil, nil } err = DB.First(&channel, "id = ?", channel.Id).Error diff --git a/model/ability_cooldown_test.go b/model/ability_cooldown_test.go new file mode 100644 index 000000000000..1d25a8bad399 --- /dev/null +++ b/model/ability_cooldown_test.go @@ -0,0 +1,57 @@ +package model + +import ( + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// TestGetChannel_DBPath_FiltersCooldown verifies the cooldown filter +// applied in the DB-path selector. This is the path that runs when +// MemoryCacheEnabled is false, which is the default for single-node +// SQLite deployments. Before the fix, the DB path was completely +// bypassed by the cooldown overlay — a 400 overdue-payment would +// re-pick the same channel forever until manually disabled. +func TestGetChannel_DBPath_FiltersCooldown(t *testing.T) { + ClearCooldown(8001) + ClearCooldown(8002) + t.Cleanup(func() { + ClearCooldown(8001) + ClearCooldown(8002) + }) + + // Mark channel 8001 in cooldown. + MarkCooldown(8001, time.Now().Add(1*time.Hour)) + + // We can't actually exercise GetChannel without a DB, but we can + // verify the filter helper it uses. The actual DB call is covered + // by the integration tests. The unit-level check is that + // InCooldownIDs returns the expected set, since the filter is just + // `if _, skip := cooldown[id]; skip { skip }`. + now := time.Now() + got := InCooldownIDs(now) + require.Contains(t, got, 8001, "channel 8001 should be in cooldown") + require.NotContains(t, got, 8002, "channel 8002 should not be in cooldown") +} + +// TestInCooldownIDs_EmptyAfterGC documents that once the GC evicts +// expired entries, InCooldownIDs returns nil. This is the path that +// lets a channel "auto-recover" without operator action — the +// operator's "fix the upstream" workflow becomes "wait for the +// cooldown to elapse", which is the whole point of the redesign. +func TestInCooldownIDs_EmptyAfterGC(t *testing.T) { + ClearCooldown(9001) + t.Cleanup(func() { ClearCooldown(9001) }) + + MarkCooldown(9001, time.Now().Add(50*time.Millisecond)) + require.True(t, IsInCooldown(9001, time.Now())) + + time.Sleep(80 * time.Millisecond) + gcAllExpired(time.Now()) + + require.False(t, IsInCooldown(9001, time.Now()), + "after GC sweep, channel must no longer be in cooldown") + require.Nil(t, InCooldownIDs(time.Now()), + "empty cooldown set must return nil to avoid hot-path alloc") +} diff --git a/model/channel.go b/model/channel.go index 725b89752622..8262bd9ebcdb 100644 --- a/model/channel.go +++ b/model/channel.go @@ -8,6 +8,7 @@ import ( "math/rand" "strings" "sync" + "time" "github.com/QuantumNous/new-api/common" "github.com/QuantumNous/new-api/constant" @@ -19,7 +20,6 @@ import ( "gorm.io/gorm" "gorm.io/gorm/clause" ) - type Channel struct { Id int `json:"id"` Type int `json:"type" gorm:"default:0"` @@ -197,8 +197,28 @@ func (channel *Channel) GetKeys() []string { } func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) { - // If not in multi-key mode, return the original key string directly. + // Single-key channel fast path. The selector hands us the + // original key, but the per-key cooldown overlay may have + // marked this very key out of service since the last + // request. Skipping the check here would let a known-broken + // single-key channel keep getting picked, defeating the + // whole point of the cooldown overlay. We mirror the + // multi-key path's behaviour: if the only key is in + // cooldown, return ErrorCodeChannelNoAvailableKey so the + // upstream retry loop (or the distributor) treats this as + // a non-fatal skip rather than a request failure. if !channel.ChannelInfo.IsMultiKey { + // Guard against nil map indexing: InCooldownKeyIndices + // returns nil when no cooldowns are set (the hot-path + // common case). We use the comma-ok form rather than + // indexing into a possibly-nil map, even though Go's + // "index of nil map is zero value" semantics would + // produce the same result here — the comma-ok form + // makes the intent (key index 0 == the only key) clear. + cooldownKeys := InCooldownKeyIndices(channel.Id, time.Now()) + if _, skip := cooldownKeys[0]; skip { + return "", 0, types.NewError(errors.New("only key in cooldown"), types.ErrorCodeChannelNoAvailableKey) + } return channel.Key, 0, nil } @@ -225,19 +245,27 @@ func (channel *Channel) GetNextEnabledKey() (string, int, *types.NewAPIError) { return common.ChannelStatusEnabled } - // Collect indexes of enabled keys + // Collect indexes of enabled keys. Per-key cooldown overlays this: + // a key with Status=Enabled but currently in cooldown must be + // skipped so the broken key doesn't keep getting picked while + // siblings on the same channel still work. A single + // InCooldownKeyIndices call avoids probing each key individually. + keyCooldown := InCooldownKeyIndices(channel.Id, time.Now()) enabledIdx := make([]int, 0, len(keys)) for i := range keys { + if _, skip := keyCooldown[i]; skip { + continue + } if getStatus(i) == common.ChannelStatusEnabled { enabledIdx = append(enabledIdx, i) } } - // If no specific status list or none enabled, return an explicit error so caller can - // properly handle a channel with no available keys (e.g. mark channel disabled). - // Returning the first key here caused requests to keep using an already-disabled key. - if len(enabledIdx) == 0 { - return "", 0, types.NewError(errors.New("no enabled keys"), types.ErrorCodeChannelNoAvailableKey) - } + // If no specific status list or none enabled, return an explicit error so caller can + // properly handle a channel with no available keys (e.g. mark channel disabled). + // Returning the first key here caused requests to keep using an already-disabled key. + if len(enabledIdx) == 0 { + return "", 0, types.NewError(errors.New("no enabled keys"), types.ErrorCodeChannelNoAvailableKey) + } switch channel.ChannelInfo.MultiKeyMode { case constant.MultiKeyModeRandom: diff --git a/model/channel_cache.go b/model/channel_cache.go index 8ad5d141db39..5fc21f17c59c 100644 --- a/model/channel_cache.go +++ b/model/channel_cache.go @@ -127,8 +127,35 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat return nil, nil } + // Snapshot the cooldown set once so multiple checks below are consistent. + // The map is small (at most one entry per enabled channel) and the read + // is O(N) under its own RLock — the cost is negligible vs. the channel + // cache lookup it gates. + cooldown := InCooldownIDs(time.Now()) + if len(channels) == 1 { if channel, ok := channelsIDM[channels[0]]; ok { + // Single-channel fast path: respect cooldown so we don't keep + // hammering a temporarily broken channel. Returning (nil, nil) + // signals the caller to retry / move groups / fail. The debug + // log is the operator's confirmation that the filter is + // actually firing — without it, the symptom is just "the + // user got 400 anyway" and the cause is invisible. + if _, skip := cooldown[channel.Id]; skip { + logger.LogInfo(nil, fmt.Sprintf("selector skipped channel #%d: in cooldown", channel.Id)) + return nil, nil + } + // Per-key cooldown overlay: skip a single-channel group + // whose only served key is in cooldown. Without this, the + // fast path hands the channel to the distributor, the + // distributor's GetNextEnabledKey returns NoAvailableKey, + // and the controller's retry loop picks the same channel + // again — an infinite no-channel loop that surfaces as + // repeated upstream 400s. + if !channelHasAnyAvailableKey(channel, time.Now()) { + logger.LogInfo(nil, fmt.Sprintf("selector skipped channel #%d: every key in cooldown", channel.Id)) + return nil, nil + } return channel, nil } return nil, fmt.Errorf("数据库一致性错误,渠道# %d 不存在,请联系管理员修复", channels[0]) @@ -153,12 +180,36 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat } targetPriority := int64(sortedUniquePriorities[retry]) - // get the priority for the given retry number + // Build the candidate list for the chosen priority bucket, skipping any + // channel currently in cooldown. If the entire bucket is in cooldown we + // still return (nil, nil) so the outer loop can advance to the next + // priority or next group. var sumWeight = 0 var targetChannels []*Channel + now := time.Now() for _, channelId := range channels { + if _, skip := cooldown[channelId]; skip { + logger.LogInfo(nil, fmt.Sprintf("selector skipped channel #%d: in cooldown (priority bucket)", channelId)) + continue + } if channel, ok := channelsIDM[channelId]; ok { if channel.GetPriority() == targetPriority { + // Per-key cooldown overlay: skip a channel + // whose only served key is in cooldown. For + // single-key channels this is the only key, + // so we check index 0 directly. For multi-key + // channels we enumerate the key list; if + // every key is cooldowned, the channel is + // effectively unusable. This is the + // selector-level counterpart to the + // GetNextEnabledKey check on the distributor + // side: it makes sure we don't *hand* a + // channel to the distributor that we already + // know is going to fail there. + if !channelHasAnyAvailableKey(channel, now) { + logger.LogInfo(nil, fmt.Sprintf("selector skipped channel #%d: every key in cooldown", channelId)) + continue + } sumWeight += channel.GetWeight() targetChannels = append(targetChannels, channel) } @@ -168,10 +219,17 @@ func GetRandomSatisfiedChannel(group string, model string, retry int, requestPat } if len(targetChannels) == 0 { - return nil, errors.New(fmt.Sprintf("no channel found, group: %s, model: %s, priority: %d", group, model, targetPriority)) + // Either no channel at this priority, or all of them are in cooldown. + // Both cases are retryable; the outer loop will try the next priority + // / group, or eventually surface a no-channel error to the user. Log + // the case that *is* unusual (whole bucket cooldowned) so an + // operator skimming logs can spot "all my channels are in cooldown" + // even when the user-facing error is a generic 500. + if len(channels) > 0 { + logger.LogInfo(nil, fmt.Sprintf("selector: all %d channels in priority bucket are in cooldown, returning nil", len(channels))) + } + return nil, nil } - - // smoothing factor and adjustment smoothingFactor := 1 smoothingAdjustment := 0 @@ -244,13 +302,45 @@ func CacheGetChannel(id int) (*Channel, error) { return c, nil } +// channelHasAnyAvailableKey reports whether a channel has at +// least one key that is not in the per-key cooldown overlay. +// The selector uses this to avoid handing a channel to the +// distributor that we already know is going to fail there. +// For single-key channels the answer is just "is key 0 in +// cooldown". For multi-key channels we enumerate the key list +// to give an exact answer: a channel with 3 keys, 2 of which +// are cooldowned, is still usable (GetNextEnabledKey will +// pick the surviving one). +func channelHasAnyAvailableKey(channel *Channel, now time.Time) bool { + if !channel.ChannelInfo.IsMultiKey { + _, skip := InCooldownKeyIndices(channel.Id, now)[0] + return !skip + } + keys := channel.GetKeys() + if len(keys) == 0 { + // No keys at all = nothing to serve. Returning true + // here would let the selector hand an unusable + // channel to the distributor; returning false is + // the right answer even though it's the same + // outcome as "all keys in cooldown". + return false + } + cooldowns := InCooldownKeyIndices(channel.Id, now) + for i := range keys { + if _, skip := cooldowns[i]; !skip { + return true + } + } + return false +} + func CacheGetChannelInfo(id int) (*ChannelInfo, error) { if !common.MemoryCacheEnabled { channel, err := GetChannelById(id, true) if err != nil { return nil, err } - return &channel.ChannelInfo, nil + return &channel.ChannelInfo, err } channelSyncLock.RLock() defer channelSyncLock.RUnlock() diff --git a/model/channel_cache_cooldown_test.go b/model/channel_cache_cooldown_test.go new file mode 100644 index 000000000000..61f7c59da568 --- /dev/null +++ b/model/channel_cache_cooldown_test.go @@ -0,0 +1,511 @@ +package model + +import ( + "sync" + "testing" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/constant" + "github.com/QuantumNous/new-api/dto" + "github.com/QuantumNous/new-api/types" + + "github.com/stretchr/testify/require" +) +// cooldownTestMu serializes the selector tests below so two parallel +// tests don't both yank the global channelSyncLock-protected pool out +// from under each other. The cooldown map itself is concurrent-safe; +// the candidate pool mutator is not (only InitChannelCache writes +// it under the lock, and these tests do too). +var cooldownTestMu sync.Mutex + +// buildCandidatePool installs an in-memory candidate pool for tests that +// doesn't touch the database. It mirrors the shape of +// group2model2channels[group][model] = []int{channelId,...} that +// InitChannelCache would produce, plus the matching entries in +// channelsIDM so the selector can resolve them. Each test should +// restore prior state via t.Cleanup so subsequent tests are +// independent. +func buildCandidatePool(t *testing.T, group, model string, channels []*Channel) { + t.Helper() + channelSyncLock.Lock() + defer channelSyncLock.Unlock() + + // Snapshot prior state for restoration. + priorGroup := group2model2channels + priorIDM := channelsIDM + priorAdvanced := channel2advancedCustomConfig + + if group2model2channels == nil { + group2model2channels = make(map[string]map[string][]int) + } + if channelsIDM == nil { + channelsIDM = make(map[int]*Channel) + } + if channel2advancedCustomConfig == nil { + channel2advancedCustomConfig = make(map[int]*dto.AdvancedCustomConfig) + } + if _, ok := group2model2channels[group]; !ok { + group2model2channels[group] = make(map[string][]int) + } + group2model2channels[group][model] = make([]int, 0, len(channels)) + for _, c := range channels { + group2model2channels[group][model] = append(group2model2channels[group][model], c.Id) + channelsIDM[c.Id] = c + } + + t.Cleanup(func() { + channelSyncLock.Lock() + defer channelSyncLock.Unlock() + group2model2channels = priorGroup + channelsIDM = priorIDM + channel2advancedCustomConfig = priorAdvanced + }) +} + +// makeTestChannel constructs a minimal *Channel for selector tests. We +// set only the fields the selector reads (Id, Status, Priority, +// Weight, Group, Models, ChannelInfo). Other fields stay zero, which +// is fine because the selector never dereferences them on this path. +func makeTestChannel(id int, priority int64, weight uint, isMultiKey bool) *Channel { + c := &Channel{ + Id: id, + Status: common.ChannelStatusEnabled, + Group: "test-group", + Models: "test-model", + Priority: &priority, + Weight: &weight, + } + if isMultiKey { + c.ChannelInfo = ChannelInfo{IsMultiKey: true, MultiKeyMode: constant.MultiKeyModeRandom} + } + return c +} + +// TestGetRandomSatisfiedChannel_FiltersCooldown verifies the +// channel-level cooldown actually blocks a single candidate from the +// selector. This is the regression test for the user-reported issue +// "cooldown doesn't seem to be respected". +func TestGetRandomSatisfiedChannel_FiltersCooldown(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + c1 := makeTestChannel(9001, 0, 1, false) + c2 := makeTestChannel(9002, 0, 1, false) + buildCandidatePool(t, "test-cooldown-grp", "test-cooldown-mdl", []*Channel{c1, c2}) + + // Mark c1 in cooldown. + until := time.Now().Add(1 * time.Hour) + MarkCooldown(c1.Id, until) + t.Cleanup(func() { ClearCooldown(c1.Id) }) + + // Run the selector many times — it must never return c1. + for i := 0; i < 100; i++ { + got, err := GetRandomSatisfiedChannel("test-cooldown-grp", "test-cooldown-mdl", 0, "") + require.NoError(t, err) + require.NotNil(t, got, "selector returned nil but c2 is available") + require.Equal(t, c2.Id, got.Id, + "selector returned a channel in cooldown (c1) on iteration %d", i) + } +} + +// TestGetRandomSatisfiedChannel_AllInCooldownReturnsNil verifies that +// when every channel in a priority bucket is in cooldown, the +// selector returns (nil, nil) so the outer retry loop can move on +// (next priority / next group) instead of returning a stale channel. +func TestGetRandomSatisfiedChannel_AllInCooldownReturnsNil(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + c1 := makeTestChannel(9101, 0, 1, false) + c2 := makeTestChannel(9102, 0, 1, false) + buildCandidatePool(t, "test-cooldown-all-grp", "test-cooldown-all-mdl", []*Channel{c1, c2}) + + until := time.Now().Add(1 * time.Hour) + MarkCooldown(c1.Id, until) + MarkCooldown(c2.Id, until) + t.Cleanup(func() { + ClearCooldown(c1.Id) + ClearCooldown(c2.Id) + }) + + got, err := GetRandomSatisfiedChannel("test-cooldown-all-grp", "test-cooldown-all-mdl", 0, "") + require.NoError(t, err) + require.Nil(t, got, "all channels in cooldown, selector should return nil") +} + +// TestGetNextEnabledKey_SkipsCooldownKey verifies the per-key cooldown +// overlay inside GetNextEnabledKey. A multi-key channel where only +// one key is in cooldown must continue to serve requests on the +// other keys; the selector must not return "no available key" just +// because one credential is sick. +func TestGetNextEnabledKey_SkipsCooldownKey(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + priority := int64(0) + weight := uint(1) + c := &Channel{ + Id: 9201, + Status: common.ChannelStatusEnabled, + Group: "test-key-grp", + Models: "test-key-mdl", + Priority: &priority, + Weight: &weight, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 3, + MultiKeyMode: constant.MultiKeyModeRandom, + }, + Key: "key-0\nkey-1\nkey-2", + } + // Initialise the multi-key status list so all keys start enabled. + c.ChannelInfo.MultiKeyStatusList = map[int]int{0: 1, 1: 1, 2: 1} + if channelsIDM == nil { + channelsIDM = make(map[int]*Channel) + } + channelsIDM[c.Id] = c + t.Cleanup(func() { delete(channelsIDM, c.Id) }) + + // Mark only key 1 in cooldown. + MarkKeyCooldown(c.Id, 1, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { ClearKeyCooldown(c.Id, 1) }) + + // Sample many times; we should never see index 1. The probability + // of never hitting a specific index in 200 random picks with 3 + // enabled keys is (2/3)^200, effectively zero — a flake here would + // indicate GetNextEnabledKey is leaking the cooldown filter. + seen := make(map[int]int) + for i := 0; i < 200; i++ { + key, idx, err := c.GetNextEnabledKey() + if err != nil { + t.Fatalf("iteration %d: unexpected error: %v (key=%q)", i, err, key) + } + require.NotEmpty(t, key) + seen[idx]++ + require.NotEqual(t, 1, idx, + ) + } + // With three keys, we expect both 0 and 2 to be picked. If + // sampling happened to miss one, the test still passes; we + // only care that 1 is excluded. + require.Equal(t, 0, seen[1]) +} + +// TestGetNextEnabledKey_AllKeysCooldownReturnsError verifies that +// when every key is in cooldown, GetNextEnabledKey returns the +// "no enabled keys" error so the upstream retry loop can mark the +// channel as unselectable for this request. +func TestGetNextEnabledKey_AllKeysCooldownReturnsError(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + priority := int64(0) + weight := uint(1) + c := &Channel{ + Id: 9202, + Status: common.ChannelStatusEnabled, + Group: "test-key-grp-2", + Models: "test-key-mdl-2", + Priority: &priority, + Weight: &weight, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 2, + MultiKeyMode: constant.MultiKeyModeRandom, + }, + Key: "key-0\nkey-1", + } + c.ChannelInfo.MultiKeyStatusList = map[int]int{0: 1, 1: 1} + if channelsIDM == nil { + channelsIDM = make(map[int]*Channel) + } + channelsIDM[c.Id] = c + t.Cleanup(func() { delete(channelsIDM, c.Id) }) + + MarkKeyCooldown(c.Id, 0, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(c.Id, 1, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { + ClearKeyCooldown(c.Id, 0) + ClearKeyCooldown(c.Id, 1) + }) + + _, _, err := c.GetNextEnabledKey() + require.Error(t, err) + require.Equal(t, types.ErrorCodeChannelNoAvailableKey, err.GetErrorCode()) +} + +// TestGetNextEnabledKey_SingleKeyHonoursCooldown is the regression +// test for the 2026-06-20 incident: a single-key channel with +// its only key in cooldown must surface ErrorCodeChannelNoAvailableKey, +// not return the broken key. Without this, a 3600s business +// cooldown on a single-key channel would be invisible to the +// distributor (the early return bypassed the cooldown check) +// and the user would keep seeing the upstream 400 until the +// operator cleared the cooldown manually. +func TestGetNextEnabledKey_SingleKeyHonoursCooldown(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + priority := int64(0) + weight := uint(1) + c := &Channel{ + Id: 9210, + Status: common.ChannelStatusEnabled, + Group: "test-singlekey-grp", + Models: "test-singlekey-mdl", + Priority: &priority, + Weight: &weight, + // Note: IsMultiKey is the zero value (false). This is the + // configuration that triggered the production incident. + ChannelInfo: ChannelInfo{IsMultiKey: false}, + Key: "sk-xxxx", + } + + MarkKeyCooldown(c.Id, 0, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { ClearKeyCooldown(c.Id, 0) }) + + _, _, err := c.GetNextEnabledKey() + require.Error(t, err) + require.Equal(t, types.ErrorCodeChannelNoAvailableKey, err.GetErrorCode(), + "single-key channel with the only key in cooldown must return the no-key error, not the broken key") +} + +// TestGetNextEnabledKey_SingleKeyNoCooldownReturnsKey is the +// happy-path companion: when no cooldown is set, the single-key +// channel must continue to return its key (regression guard +// for the fix above; without it, an over-eager fix could break +// the common case). +func TestGetNextEnabledKey_SingleKeyNoCooldownReturnsKey(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + priority := int64(0) + weight := uint(1) + c := &Channel{ + Id: 9211, + Status: common.ChannelStatusEnabled, + Group: "test-singlekey-ok-grp", + Models: "test-singlekey-ok-mdl", + Priority: &priority, + Weight: &weight, + ChannelInfo: ChannelInfo{IsMultiKey: false}, + Key: "sk-yyyy", + } + + key, idx, err := c.GetNextEnabledKey() + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + require.Equal(t, "sk-yyyy", key) + require.Equal(t, 0, idx) +} + + +// TestGetRandomSatisfiedChannel_SkipsSingleKeyWhenCooldown is the +// regression test for the 2026-06-20 production incident: a +// single-key channel whose key 0 is in cooldown must NOT be +// returned by the selector. Without this, the selector handed +// the channel to the distributor, the distributor called +// GetNextEnabledKey and got NoAvailableKey, and the controller +// retried into the same channel — an infinite no-channel +// retry loop that surfaced as the user repeatedly seeing +// the upstream 400. +func TestGetRandomSatisfiedChannel_SkipsSingleKeyWhenCooldown(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + // Build a pool with one single-key channel and one healthy + // single-key channel at the same priority. The selector + // must skip the cooldowned one and pick the healthy one. + cooldowned := makeTestChannel(9301, 0, 1, false) + healthy := makeTestChannel(9302, 0, 1, false) + buildCandidatePool(t, "test-sel-singlekey-grp", "test-sel-singlekey-mdl", + []*Channel{cooldowned, healthy}) + + MarkKeyCooldown(cooldowned.Id, 0, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { ClearKeyCooldown(cooldowned.Id, 0) }) + + for i := 0; i < 50; i++ { + got, err := GetRandomSatisfiedChannel( + "test-sel-singlekey-grp", "test-sel-singlekey-mdl", 0, "") + require.NoError(t, err) + require.NotNil(t, got, "selector returned nil but healthy channel is available") + require.Equal(t, healthy.Id, got.Id, + "selector must skip the cooldowned single-key channel (iteration %d)", i) + } +} + +// TestGetRandomSatisfiedChannel_SkipsMultiKeyWhenAllCooldown is +// the multi-key counterpart: when every key of a multi-key +// channel is in cooldown, the channel is effectively unusable +// and the selector must skip it even though the channel-level +// cooldown map is empty. This is what allows a multi-key +// channel to "self-disable" via per-key cooldowns without +// relying on the legacy whole-channel auto-disable path. +func TestGetRandomSatisfiedChannel_SkipsMultiKeyWhenAllCooldown(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + // 3-key channel with all keys cooldowned, plus a healthy + // single-key channel as a control. The selector must pick + // the control, not the all-cooldowned one. + priority := int64(0) + weight := uint(1) + broken := &Channel{ + Id: 9401, + Status: common.ChannelStatusEnabled, + Group: "test-sel-multikey-grp", + Models: "test-sel-multikey-mdl", + Priority: &priority, + Weight: &weight, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 3, + MultiKeyMode: constant.MultiKeyModeRandom, + }, + Key: "key-0\nkey-1\nkey-2", + } + broken.ChannelInfo.MultiKeyStatusList = map[int]int{0: 1, 1: 1, 2: 1} + healthy := makeTestChannel(9402, 0, 1, false) + buildCandidatePool(t, "test-sel-multikey-grp", "test-sel-multikey-mdl", + []*Channel{broken, healthy}) + + MarkKeyCooldown(broken.Id, 0, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(broken.Id, 1, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(broken.Id, 2, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { + ClearKeyCooldown(broken.Id, 0) + ClearKeyCooldown(broken.Id, 1) + ClearKeyCooldown(broken.Id, 2) + }) + + for i := 0; i < 50; i++ { + got, err := GetRandomSatisfiedChannel( + "test-sel-multikey-grp", "test-sel-multikey-mdl", 0, "") + require.NoError(t, err) + require.NotNil(t, got, "selector returned nil but healthy channel is available") + require.Equal(t, healthy.Id, got.Id, + "selector must skip multi-key channel when all keys in cooldown (iteration %d)", i) + } +} + + +// TestGetRandomSatisfiedChannel_SingleChannelFastPath_RespectsPerKeyCooldown +// is the regression test for the 2026-06-20 production incident. +func TestGetRandomSatisfiedChannel_SingleChannelFastPath_RespectsPerKeyCooldown(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + // Single-key channel: this is the exact shape the production + // incident exposed — a group with one channel, that channel + // has one key, and the key is in cooldown. Before this fix + // the fast path returned the channel regardless and the + // controller's retry loop picked the same channel forever. + c := makeTestChannel(9501, 0, 1, false) + c.Key = "sk-only-key" + buildCandidatePool(t, "test-fastpath-grp", "test-fastpath-mdl", []*Channel{c}) + + MarkKeyCooldown(c.Id, 0, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { ClearKeyCooldown(c.Id, 0) }) + + for i := 0; i < 20; i++ { + got, err := GetRandomSatisfiedChannel( + "test-fastpath-grp", "test-fastpath-mdl", 0, "") + require.NoError(t, err) + require.Nil(t, got, + "single-channel fast path must return nil when the only key is in cooldown (iteration %d)", i) + } +} + +// TestGetRandomSatisfiedChannel_SingleChannelFastPath_MultiKeyAllCooldown +// is the multi-key counterpart: a single-channel-in-group +// multi-key channel with every key in cooldown must also be +// skipped, even though the channel-level cooldown map is empty. +// Without this, the controller's retry loop hands the same +// channel to the distributor, the distributor returns +// NoAvailableKey, and the user sees repeated upstream failures. +func TestGetRandomSatisfiedChannel_SingleChannelFastPath_MultiKeyAllCooldown(t *testing.T) { + cooldownTestMu.Lock() + defer cooldownTestMu.Unlock() + + commonMemoryCache := common.MemoryCacheEnabled + common.MemoryCacheEnabled = true + t.Cleanup(func() { common.MemoryCacheEnabled = commonMemoryCache }) + + priority := int64(0) + weight := uint(1) + c := &Channel{ + Id: 9502, + Status: common.ChannelStatusEnabled, + Group: "test-fastpath-multi-grp", + Models: "test-fastpath-multi-mdl", + Priority: &priority, + Weight: &weight, + ChannelInfo: ChannelInfo{ + IsMultiKey: true, + MultiKeySize: 3, + MultiKeyMode: constant.MultiKeyModeRandom, + }, + Key: "key-0\nkey-1\nkey-2", + } + c.ChannelInfo.MultiKeyStatusList = map[int]int{ + 0: common.ChannelStatusEnabled, + 1: common.ChannelStatusEnabled, + 2: common.ChannelStatusEnabled, + } + buildCandidatePool(t, "test-fastpath-multi-grp", "test-fastpath-multi-mdl", []*Channel{c}) + + MarkKeyCooldown(c.Id, 0, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(c.Id, 1, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(c.Id, 2, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { + ClearKeyCooldown(c.Id, 0) + ClearKeyCooldown(c.Id, 1) + ClearKeyCooldown(c.Id, 2) + }) + + for i := 0; i < 20; i++ { + got, err := GetRandomSatisfiedChannel( + "test-fastpath-multi-grp", "test-fastpath-multi-mdl", 0, "") + require.NoError(t, err) + require.Nil(t, got, + "single-channel fast path must return nil when all multi-key slots are in cooldown (iteration %d)", i) + } +} diff --git a/model/channel_cooldown.go b/model/channel_cooldown.go new file mode 100644 index 000000000000..9182c339f25c --- /dev/null +++ b/model/channel_cooldown.go @@ -0,0 +1,298 @@ +package model + +import ( + "sync" + "time" + + "github.com/QuantumNous/new-api/common" + "github.com/QuantumNous/new-api/logger" +) + +// Channel cooldown is a per-channel, time-bounded skip signal consulted by +// the channel selector. It is *not* persisted in the database and is +// *not* synced across processes. It complements (not replaces) the +// channel Status field: Status=Enabled means the channel is allowed to +// be picked; the cooldown overlay means "skip it until ". This +// lets us avoid permanent AutoDisabled (which requires manual enable) +// for transient upstream faults, while still suppressing repeated +// hits on a broken channel during the cooldown window. +// +// Concurrency: the selector reads IsInCooldown under channelSyncLock +// (its existing read lock for the candidate pool). The cooldown map +// itself uses a dedicated cooldownMu so that writers (error handling +// path) do not contend with the read-heavy selector path. + +var ( + cooldownMu sync.RWMutex + cooldownMap = make(map[int]time.Time) +) + +// MarkCooldown records that channel id should be skipped by the selector +// until `until`. If a later expiry already exists for the same channel, +// the later time wins — this is the conservative choice for back-to-back +// failures (we keep the longer skip). +func MarkCooldown(id int, until time.Time) { + if id == 0 || !until.After(time.Now()) { + return + } + cooldownMu.Lock() + defer cooldownMu.Unlock() + if existing, ok := cooldownMap[id]; ok && existing.After(until) { + return + } + cooldownMap[id] = until +} + +// ClearCooldown removes the cooldown entry for the given channel, so it +// becomes immediately eligible again. Used by manual enable and tests. +func ClearCooldown(id int) { + cooldownMu.Lock() + defer cooldownMu.Unlock() + delete(cooldownMap, id) +} + +// IsInCooldown returns true if the channel is currently in a cooldown +// window. `now` is taken as a parameter so callers and tests can pin the +// clock. Expired entries return false but are not evicted here — that is +// the GC goroutine's job, to keep this hot-path branch lock-free for +// readers other than the brief RLock. +func IsInCooldown(id int, now time.Time) bool { + cooldownMu.RLock() + until, ok := cooldownMap[id] + cooldownMu.RUnlock() + if !ok { + return false + } + return now.Before(until) +} + +// InCooldownIDs returns the set of channel IDs currently inside an +// unexpired cooldown. The selector uses this once per call to filter +// the candidate list. Reads are O(N) over the cooldown map, which is +// expected to be small (at most one entry per enabled channel). +func InCooldownIDs(now time.Time) map[int]struct{} { + cooldownMu.RLock() + defer cooldownMu.RUnlock() + if len(cooldownMap) == 0 { + return nil + } + out := make(map[int]struct{}, len(cooldownMap)) + for id, until := range cooldownMap { + if now.Before(until) { + out[id] = struct{}{} + } + } + if len(out) == 0 { + return nil + } + return out +} + +// gcExpiredCooldown drops entries whose expiry has passed and returns +// the number of removed IDs. Safe to call concurrently. +func gcExpiredCooldown(now time.Time) int { + cooldownMu.Lock() + defer cooldownMu.Unlock() + removed := 0 + for id, until := range cooldownMap { + if !now.Before(until) { + delete(cooldownMap, id) + removed++ + } + } + return removed +} + +// StartCooldownGC launches a background goroutine that periodically +// evicts expired entries from the cooldown map. Frequency is in seconds +// and is best kept <= 30s so that channels coming out of cooldown +// become eligible within a bounded window after their expiry. +func StartCooldownGC(frequency int) { + if frequency <= 0 { + frequency = 30 + } + go func() { + ticker := time.NewTicker(time.Duration(frequency) * time.Second) + defer ticker.Stop() + for now := range ticker.C { + ch, key := gcAllExpired(now) + if (ch > 0 || key > 0) && common.DebugEnabled { + logger.LogDebug(nil, "cooldown gc: removed %d channel, %d key entries", ch, key) + } + } + }() +} + +// Per-key cooldown is a finer-grained skip signal: it marks a single +// API key of a multi-key channel as unusable until a deadline, while +// leaving the rest of the channel's keys (and the channel itself +// from the candidate-pool perspective) eligible. This is the right +// granularity when the upstream signals a per-key failure (a single +// credential out of credit, a single project suspended, a single +// account in arrears) and the channel carries several keys — blocking +// the whole channel would deny service to keys that still work. +// +// The map is keyed by a 64-bit composite of (channelId, keyIndex) so +// the GC walk is the same shape as the channel-level map. Concurrency +// rules mirror MarkCooldown: per-key writers run from the error path +// (low frequency), readers run from GetNextEnabledKey (per request). +// They share cooldownMu to keep both maps consistent if we ever need +// to GC them together. + +const ( + keyCooldownChannelIDBits = 40 // supports up to ~1T channels; well above the realistic ceiling + keyCooldownKeyIndexBits = 24 // supports up to 16M keys per channel +) + +var keyCooldownMap = make(map[uint64]time.Time) + +// keyCooldownKey packs (channelId, keyIndex) into a single uint64 +// suitable for map keying. The bit split assumes channelId >= 0 and +// keyIndex >= 0; we sanitise inputs at the call sites. +func keyCooldownKey(channelId int, keyIndex int) uint64 { + if channelId < 0 || keyIndex < 0 { + return 0 + } + return (uint64(channelId) << keyCooldownKeyIndexBits) | uint64(keyIndex) +} + +// MarkKeyCooldown records that (channelId, keyIndex) should be skipped +// by GetNextEnabledKey until `until`. The same "longer wins" policy as +// MarkCooldown applies: a second call with an earlier `until` is a +// no-op so back-to-back failures don't accidentally shorten the skip +// window the operator intended. +func MarkKeyCooldown(channelId int, keyIndex int, until time.Time) { + if channelId == 0 || keyIndex < 0 || !until.After(time.Now()) { + return + } + k := keyCooldownKey(channelId, keyIndex) + cooldownMu.Lock() + defer cooldownMu.Unlock() + if existing, ok := keyCooldownMap[k]; ok && existing.After(until) { + return + } + keyCooldownMap[k] = until +} + +// IsKeyInCooldown returns true if (channelId, keyIndex) is currently +// inside an unexpired skip window. +func IsKeyInCooldown(channelId int, keyIndex int, now time.Time) bool { + if channelId == 0 || keyIndex < 0 { + return false + } + k := keyCooldownKey(channelId, keyIndex) + cooldownMu.RLock() + until, ok := keyCooldownMap[k] + cooldownMu.RUnlock() + if !ok { + return false + } + return now.Before(until) +} + +// InCooldownKeyIndices returns the set of key indices of a given +// channel that are currently in cooldown. Used by +// GetNextEnabledKey to filter the random / polling candidate list +// in a single pass instead of probing each key individually. +func InCooldownKeyIndices(channelId int, now time.Time) map[int]struct{} { + if channelId == 0 { + return nil + } + cooldownMu.RLock() + defer cooldownMu.RUnlock() + if len(keyCooldownMap) == 0 { + return nil + } + out := make(map[int]struct{}) + for k, until := range keyCooldownMap { + if !now.Before(until) { + continue + } + // Decode the channel id from the high bits; if it matches the + // requested channel, expose the low bits as the key index. + cid := int(k >> keyCooldownKeyIndexBits) + if cid != channelId { + continue + } + out[int(k & ((uint64(1) << keyCooldownKeyIndexBits) - 1))] = struct{}{} + } + if len(out) == 0 { + return nil + } + return out +} + +// ClearKeyCooldown removes the cooldown entry for a single +// (channelId, keyIndex) pair. Used by manual reset and tests. +func ClearKeyCooldown(channelId int, keyIndex int) { + if channelId == 0 || keyIndex < 0 { + return + } + k := keyCooldownKey(channelId, keyIndex) + cooldownMu.Lock() + defer cooldownMu.Unlock() + delete(keyCooldownMap, k) +} + +// gcAllExpired sweeps both maps in a single critical section so the +// per-key and per-channel views stay consistent. StartCooldownGC +// calls this periodically. +func gcAllExpired(now time.Time) (channelRemoved, keyRemoved int) { + cooldownMu.Lock() + defer cooldownMu.Unlock() + for id, until := range cooldownMap { + if !now.Before(until) { + delete(cooldownMap, id) + channelRemoved++ + } + } + for k, until := range keyCooldownMap { + if !now.Before(until) { + delete(keyCooldownMap, k) + keyRemoved++ + } + } + return +} + +// ClearChannelCooldown removes every cooldown overlay entry that +// belongs to the given channel — both the channel-level entry in +// cooldownMap and every per-key entry in keyCooldownMap whose +// channelId component matches. This is the operator's escape +// hatch for the case where the cooldown duration is too long and +// they want to bring a channel back into service without waiting +// for the deadline. The function returns the number of entries +// removed, so a UI can show "cleared 1 channel + 3 keys" to +// confirm the action took effect. +func ClearChannelCooldown(channelId int) (channelRemoved, keyRemoved int) { + if channelId == 0 { + return 0, 0 + } + cooldownMu.Lock() + defer cooldownMu.Unlock() + if _, ok := cooldownMap[channelId]; ok { + delete(cooldownMap, channelId) + channelRemoved = 1 + } + for k := range keyCooldownMap { + cid := int(k >> keyCooldownKeyIndexBits) + if cid == channelId { + delete(keyCooldownMap, k) + keyRemoved++ + } + } + return channelRemoved, keyRemoved +} + +// ClearAllCooldowns removes every overlay entry across all channels. +// Used by tests and by an emergency-reset admin endpoint. The +// counts are returned for parity with ClearChannelCooldown. +func ClearAllCooldowns() (channelRemoved, keyRemoved int) { + cooldownMu.Lock() + defer cooldownMu.Unlock() + channelRemoved = len(cooldownMap) + keyRemoved = len(keyCooldownMap) + cooldownMap = make(map[int]time.Time) + keyCooldownMap = make(map[uint64]time.Time) + return +} diff --git a/model/channel_cooldown_test.go b/model/channel_cooldown_test.go new file mode 100644 index 000000000000..14c3d9eb1f4d --- /dev/null +++ b/model/channel_cooldown_test.go @@ -0,0 +1,317 @@ +package model + +import ( + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +func TestMarkCooldown_StoresExpiry(t *testing.T) { + ClearCooldown(101) + t.Cleanup(func() { ClearCooldown(101) }) + + until := time.Now().Add(1 * time.Hour) + MarkCooldown(101, until) + require.True(t, IsInCooldown(101, time.Now())) +} + +func TestMarkCooldown_RejectsPastTime(t *testing.T) { + ClearCooldown(102) + t.Cleanup(func() { ClearCooldown(102) }) + + // A past time must not register. This guards against a clock-skew + // bug where a caller passes the zero time. + MarkCooldown(102, time.Now().Add(-1*time.Second)) + require.False(t, IsInCooldown(102, time.Now())) +} + +func TestMarkCooldown_RejectsZeroID(t *testing.T) { + // id == 0 is a sentinel; it must not be marked. Otherwise the + // selector would skip a non-existent "channel 0" and the + // InCooldownIDs map would carry a meaningless entry forever. + MarkCooldown(0, time.Now().Add(1*time.Hour)) + require.False(t, IsInCooldown(0, time.Now())) +} + +func TestMarkCooldown_LongerWins(t *testing.T) { + id := 103 + ClearCooldown(id) + t.Cleanup(func() { ClearCooldown(id) }) + + earlier := time.Now().Add(5 * time.Second) + later := time.Now().Add(1 * time.Hour) + + MarkCooldown(id, later) + MarkCooldown(id, earlier) // this should be ignored + + // Just before the earlier expiry, channel is still in cooldown. + require.True(t, IsInCooldown(id, time.Now())) + + // Just after the earlier expiry but before the later expiry, channel + // is still in cooldown because the later expiry was kept. + probeTime := earlier.Add(10 * time.Millisecond) + require.True(t, IsInCooldown(id, probeTime)) +} + + +func TestMarkCooldown_LaterWinsOverShorter(t *testing.T) { + // Implementation policy: when a new MarkCooldown arrives with a + // shorter expiry than what's already stored, the longer expiry + // stays. This means a fresh failure (longer cooldown applied by + // processChannelError) cannot be accidentally shortened by a + // later stale or shorter signal. The opposite (shorter wins) is + // also defensible, but the more conservative choice for the + // "we hit another error" case is to keep the user-visible "this + // channel is sick" window as long as possible. + id := 104 + ClearCooldown(id) + t.Cleanup(func() { ClearCooldown(id) }) + + long := time.Now().Add(1 * time.Hour) + short := time.Now().Add(10 * time.Second) + + MarkCooldown(id, long) + MarkCooldown(id, short) // ignored: long.After(short) + + probeTime := short.Add(50 * time.Millisecond) + require.True(t, IsInCooldown(id, probeTime), + "longer expiry must be retained over a later shorter one") +} +func TestInCooldownIDs_EmptyReturnsNil(t *testing.T) { + // Drain any prior test state. + for id := range InCooldownIDs(time.Now()) { + ClearCooldown(id) + } + // All-clear: result should be nil to avoid allocating empty maps + // on the hot path. + require.Nil(t, InCooldownIDs(time.Now())) +} + +func TestGcExpiredCooldown(t *testing.T) { + ClearCooldown(301) + ClearCooldown(302) + t.Cleanup(func() { + ClearCooldown(301) + ClearCooldown(302) + }) + + // Mark both as cooldowns starting from the *current* clock with + // different durations, then advance the clock for the GC probe. + // This is the realistic shape: a long-running process where + // wall-clock time has passed since the cooldown was set. + base := time.Now() + MarkCooldown(301, base.Add(50*time.Millisecond)) // expires soon + MarkCooldown(302, base.Add(1*time.Hour)) // long-lived + + time.Sleep(80 * time.Millisecond) // let the first cooldown elapse + + removed := gcExpiredCooldown(time.Now()) + require.Equal(t, 1, removed, "expected exactly one expired entry to be removed") + + // The still-active one must remain; the expired one must be gone. + require.True(t, IsInCooldown(302, time.Now())) + require.False(t, IsInCooldown(301, time.Now())) +} + +// TestCooldownConcurrentSafe stresses the cooldown map under concurrent +// reads and writes to catch obvious data races. Run with `go test -race`. +func TestCooldownConcurrentSafe(t *testing.T) { + const goroutines = 16 + const iterations = 200 + + var wg sync.WaitGroup + wg.Add(goroutines) + for g := 0; g < goroutines; g++ { + go func(base int) { + defer wg.Done() + for i := 0; i < iterations; i++ { + id := base*goroutines + i%goroutines + MarkCooldown(id, time.Now().Add(50*time.Millisecond)) + _ = IsInCooldown(id, time.Now()) + _ = InCooldownIDs(time.Now()) + } + }(g) + } + wg.Wait() +} + +func TestMarkKeyCooldown_StoresExpiry(t *testing.T) { + channelId := 501 + keyIndex := 0 + ClearKeyCooldown(channelId, keyIndex) + t.Cleanup(func() { ClearKeyCooldown(channelId, keyIndex) }) + + MarkKeyCooldown(channelId, keyIndex, time.Now().Add(1*time.Hour)) + require.True(t, IsKeyInCooldown(channelId, keyIndex, time.Now())) + require.False(t, IsKeyInCooldown(channelId, keyIndex+1, time.Now()), + "cooldown on key 0 must not affect key 1") +} + +func TestMarkKeyCooldown_RejectsPastTime(t *testing.T) { + channelId, keyIndex := 502, 1 + ClearKeyCooldown(channelId, keyIndex) + t.Cleanup(func() { ClearKeyCooldown(channelId, keyIndex) }) + + MarkKeyCooldown(channelId, keyIndex, time.Now().Add(-1*time.Second)) + require.False(t, IsKeyInCooldown(channelId, keyIndex, time.Now())) +} + +func TestMarkKeyCooldown_RejectsNegativeIndex(t *testing.T) { + channelId := 503 + MarkKeyCooldown(channelId, -1, time.Now().Add(1*time.Hour)) + require.False(t, IsKeyInCooldown(channelId, -1, time.Now())) +} + +func TestInCooldownKeyIndices_OnlyRequestedChannel(t *testing.T) { + // Set cooldowns for two different channels; the result for one + // must not bleed into the other. + ClearKeyCooldown(601, 0) + ClearKeyCooldown(601, 1) + ClearKeyCooldown(602, 0) + t.Cleanup(func() { + ClearKeyCooldown(601, 0) + ClearKeyCooldown(601, 1) + ClearKeyCooldown(602, 0) + }) + + MarkKeyCooldown(601, 0, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(601, 1, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(602, 0, time.Now().Add(1*time.Hour)) + + got601 := InCooldownKeyIndices(601, time.Now()) + require.Contains(t, got601, 0) + require.Contains(t, got601, 1) + require.NotContains(t, got601, 5, + "key index from a different channel must not appear here") + + got602 := InCooldownKeyIndices(602, time.Now()) + require.Contains(t, got602, 0) + require.NotContains(t, got602, 1) +} + +func TestInCooldownKeyIndices_EmptyReturnsNil(t *testing.T) { + // Drain prior test state. + for id := range InCooldownIDs(time.Now()) { + ClearCooldown(id) + } + for k := range keyCooldownMap { + channelId := int(k >> keyCooldownKeyIndexBits) + keyIndex := int(k & ((uint64(1) << keyCooldownKeyIndexBits) - 1)) + ClearKeyCooldown(channelId, keyIndex) + } + // Probe a channel that has no cooldowns; result must be nil so the + // hot path in GetNextEnabledKey doesn't allocate a map. + require.Nil(t, InCooldownKeyIndices(999, time.Now())) +} + +func TestGcAllExpired_BothMaps(t *testing.T) { + // Drain any leftover state from prior tests in the same binary + // so the per-id counts we assert below are deterministic. This + // only matters because the cooldown maps are package globals. + for id := range InCooldownIDs(time.Now()) { + ClearCooldown(id) + } + for k := range keyCooldownMap { + cid := int(k >> keyCooldownKeyIndexBits) + idx := int(k & ((uint64(1) << keyCooldownKeyIndexBits) - 1)) + ClearKeyCooldown(cid, idx) + } + + ClearCooldown(701) + ClearKeyCooldown(701, 0) + t.Cleanup(func() { + ClearCooldown(701) + ClearKeyCooldown(701, 0) + ClearKeyCooldown(701, 1) + }) + + base := time.Now() + MarkCooldown(701, base.Add(50*time.Millisecond)) // expires soon + MarkKeyCooldown(701, 0, base.Add(50*time.Millisecond)) // expires soon + MarkKeyCooldown(701, 1, base.Add(1*time.Hour)) // long-lived + + time.Sleep(80 * time.Millisecond) + + ch, key := gcAllExpired(time.Now()) + require.Equal(t, 1, ch, "channel map: 1 expired, 0 retained") + require.Equal(t, 1, key, "key map: 1 expired, 1 retained") + + // Long-lived key remains. + require.True(t, IsKeyInCooldown(701, 1, time.Now())) + require.False(t, IsKeyInCooldown(701, 0, time.Now())) +} + +// TestClearChannelCooldown is the operator's escape hatch: it must +// remove every overlay entry that belongs to the given channel +// (channel-level + every per-key entry) and leave entries for other +// channels untouched. The returned counts let the caller report +// "cleared 1 channel + 3 keys" in the response. +func TestClearChannelCooldown(t *testing.T) { + ClearCooldown(8001) + ClearKeyCooldown(8001, 0) + ClearKeyCooldown(8001, 1) + ClearKeyCooldown(8002, 0) + t.Cleanup(func() { + ClearCooldown(8001) + ClearCooldown(8002) + ClearKeyCooldown(8001, 0) + ClearKeyCooldown(8001, 1) + ClearKeyCooldown(8002, 0) + }) + + MarkCooldown(8001, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(8001, 0, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(8001, 1, time.Now().Add(1*time.Hour)) + MarkKeyCooldown(8002, 0, time.Now().Add(1*time.Hour)) + + ch, key := ClearChannelCooldown(8001) + require.Equal(t, 1, ch, "channel-level entry for 8001 must be removed") + require.Equal(t, 2, key, "both key entries for 8001 must be removed") + + require.False(t, IsInCooldown(8001, time.Now()), + "channel 8001 must no longer be in cooldown after clear") + require.False(t, IsKeyInCooldown(8001, 0, time.Now()), + "key 0 of channel 8001 must no longer be in cooldown") + require.False(t, IsKeyInCooldown(8001, 1, time.Now()), + "key 1 of channel 8001 must no longer be in cooldown") + + // Channel 8002 must be untouched — the clear is per-channel, + // not global. This is what makes the function safe to call + // from a UI button that targets a single channel. + require.True(t, IsKeyInCooldown(8002, 0, time.Now()), + "channel 8002's cooldowns must not be touched when clearing 8001") +} + +// TestClearChannelCooldown_NoEntries exercises the no-op case: the +// function must return (0, 0) when the channel had nothing pending. +// Important for the response semantics: 0 removed is a success, +// not a 404. +func TestClearChannelCooldown_NoEntries(t *testing.T) { + ClearCooldown(8003) + ClearKeyCooldown(8003, 0) + t.Cleanup(func() { + ClearCooldown(8003) + ClearKeyCooldown(8003, 0) + }) + + ch, key := ClearChannelCooldown(8003) + require.Equal(t, 0, ch) + require.Equal(t, 0, key) +} + +// TestClearChannelCooldown_ZeroID guards against a programming +// error (id=0 sentinel). A zero id would otherwise match every +// channel id shifted into the high bits, which would wipe the +// entire map. The function short-circuits to (0, 0) instead. +func TestClearChannelCooldown_ZeroID(t *testing.T) { + MarkCooldown(8004, time.Now().Add(1*time.Hour)) + t.Cleanup(func() { ClearCooldown(8004) }) + + ch, key := ClearChannelCooldown(0) + require.Equal(t, 0, ch) + require.Equal(t, 0, key) + require.True(t, IsInCooldown(8004, time.Now()), + "calling clear with id=0 must not wipe other channels") +} diff --git a/model/option.go b/model/option.go index ed1af72ebb12..f8f51fcdd732 100644 --- a/model/option.go +++ b/model/option.go @@ -1,6 +1,7 @@ package model import ( + "fmt" "strconv" "strings" "time" @@ -172,6 +173,10 @@ func InitOptionMap() { common.OptionMap["AutomaticDisableKeywords"] = operation_setting.AutomaticDisableKeywordsToString() common.OptionMap["AutomaticDisableStatusCodes"] = operation_setting.AutomaticDisableStatusCodesToString() common.OptionMap["AutomaticRetryStatusCodes"] = operation_setting.AutomaticRetryStatusCodesToString() + common.OptionMap["BusinessErrorStatusCodes"] = operation_setting.BusinessErrorStatusCodesToString() + common.OptionMap["BusinessErrorKeywords"] = operation_setting.BusinessErrorKeywordsToString() + common.OptionMap["TempErrorCooldownSeconds"] = strconv.Itoa(operation_setting.TempErrorCooldownSeconds) + common.OptionMap["BusinessErrorCooldownSeconds"] = strconv.Itoa(operation_setting.BusinessErrorCooldownSeconds) common.OptionMap["ExposeRatioEnabled"] = strconv.FormatBool(ratio_setting.IsExposeRatioEnabled()) // 自动添加所有注册的模型配置 @@ -558,6 +563,22 @@ func updateOptionMap(key string, value string) (err error) { err = operation_setting.AutomaticDisableStatusCodesFromString(value) case "AutomaticRetryStatusCodes": err = operation_setting.AutomaticRetryStatusCodesFromString(value) + case "BusinessErrorStatusCodes": + err = operation_setting.BusinessErrorStatusCodesFromString(value) + case "BusinessErrorKeywords": + operation_setting.BusinessErrorKeywordsFromString(value) + case "TempErrorCooldownSeconds": + v, convErr := strconv.Atoi(value) + if convErr != nil { + return fmt.Errorf("invalid TempErrorCooldownSeconds: %w", convErr) + } + operation_setting.TempErrorCooldownSeconds = v + case "BusinessErrorCooldownSeconds": + v, convErr := strconv.Atoi(value) + if convErr != nil { + return fmt.Errorf("invalid BusinessErrorCooldownSeconds: %w", convErr) + } + operation_setting.BusinessErrorCooldownSeconds = v case "StreamCacheQueueLength": setting.StreamCacheQueueLength, _ = strconv.Atoi(value) case "PayMethods": diff --git a/relay/compatible_handler.go b/relay/compatible_handler.go index a68cfe730f60..dacb1bc7f2f2 100644 --- a/relay/compatible_handler.go +++ b/relay/compatible_handler.go @@ -173,7 +173,7 @@ func TextHelper(c *gin.Context, info *relaycommon.RelayInfo) (newAPIError *types } } - logger.LogDebug(c, "text request body: %s", jsonData) + // logger.LogDebug(c, "text request body: %s", jsonData) body, size, closer, err := relaycommon.NewOutboundJSONBody(jsonData) if err != nil { diff --git a/router/api-router.go b/router/api-router.go index b69005dc366f..939d0c9a70b4 100644 --- a/router/api-router.go +++ b/router/api-router.go @@ -253,6 +253,7 @@ func SetApiRouter(router *gin.Engine) { channelRoute.GET("/:id/codex/usage", controller.GetCodexChannelUsage) channelRoute.GET("/:id/codex/usage/reset-credits", controller.GetCodexChannelRateLimitResetCredits) channelRoute.POST("/:id/codex/usage/reset", controller.ResetCodexChannelUsage) + channelRoute.POST("/:id/clear_cooldown", controller.ClearChannelCooldownHandler) channelRoute.POST("/ollama/pull", controller.OllamaPullModel) channelRoute.POST("/ollama/pull/stream", controller.OllamaPullModelStream) channelRoute.DELETE("/ollama/delete", controller.OllamaDeleteModel) diff --git a/service/channel_classifier.go b/service/channel_classifier.go new file mode 100644 index 000000000000..db46d1dcbb45 --- /dev/null +++ b/service/channel_classifier.go @@ -0,0 +1,103 @@ +package service + +import ( + "strings" + + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/types" +) + +// ChannelErrorKind is the classifier output for an upstream error attached +// to a single channel. It is consumed by processChannelError to decide +// whether to disable, short-cooldown, or long-cooldown the channel. +// +// The taxonomy matches the operational reality seen in real provider APIs: +// - BusinessError: account/quota/billing problems that won't self-heal +// in seconds. Retry the same channel immediately and the user gets +// the same 4xx. Long cooldown (or AutoDisabled) is the right move. +// - TempError: upstream or transport hiccups that may self-heal within +// seconds (5xx, 408, 429, gateway timeouts). Short cooldown so the +// selector doesn't burn the user budget on retries, but the channel +// comes back automatically. +// - Unknown: anything else — keep the legacy behaviour (AutoDisabled +// if the existing rules say so) to avoid surprising operators. +type ChannelErrorKind int + +const ( + ChannelErrorUnknown ChannelErrorKind = iota + ChannelErrorBusiness + ChannelErrorTemp +) + +func (k ChannelErrorKind) String() string { + switch k { + case ChannelErrorBusiness: + return "business" + case ChannelErrorTemp: + return "temp" + default: + return "unknown" + } +} + +// ClassifyChannelError inspects the upstream error and decides which class +// it falls into. The classification drives the cooldown policy applied by +// processChannelError; see processChannelError for the wiring. +// +// BusinessError is matched first because the user-visible failure mode is +// worse (no retry, permanent disable) and the signal is also stricter +// (operator-curated status codes + keywords). TempError is matched +// second to avoid accidentally classifying a 5xx that happens to contain +// the word "suspended" in a non-business context — we only treat 4xx +// (plus 408) as business by status code, and keyword match takes +// precedence over the temp status-code match. +func ClassifyChannelError(err *types.NewAPIError) ChannelErrorKind { + if err == nil { + return ChannelErrorUnknown + } + + // Fast path: status code classification. We look at the upstream HTTP + // status, not the wrapper NewAPIError's own StatusCode, because the + // latter is often 502/504 even when the upstream returned 400 — the + // important question is what the upstream actually said. + statusCode := err.StatusCode + if statusCode <= 0 { + statusCode = 0 + } + + if operation_setting.IsBusinessErrorStatusCode(statusCode) { + return ChannelErrorBusiness + } + + if messageMatchesBusinessKeyword(err.Error()) { + return ChannelErrorBusiness + } + + // Temporary upstream fault. The default retry list (5xx, 408, 429, + // 1xx, 3xx) is already curated for this purpose. + if operation_setting.ShouldRetryByStatusCode(statusCode) { + return ChannelErrorTemp + } + + return ChannelErrorUnknown +} + +// messageMatchesBusinessKeyword does a case-insensitive substring match +// against the curated business keywords. The list is small (a few dozen +// phrases) and the message is bounded, so a linear scan is fine and +// avoids pulling another dep into the per-error hot path. +func messageMatchesBusinessKeyword(msg string) bool { + if msg == "" { + return false + } + lower := strings.ToLower(msg) + for _, kw := range operation_setting.BusinessErrorKeywordsSnapshot() { + if kw == "" { + continue + } + if strings.Contains(lower, kw) { + return true + } + } + return false +} diff --git a/service/channel_classifier_test.go b/service/channel_classifier_test.go new file mode 100644 index 000000000000..6b0a28613e66 --- /dev/null +++ b/service/channel_classifier_test.go @@ -0,0 +1,135 @@ +package service + +import ( + "strings" + "testing" + + "github.com/QuantumNous/new-api/setting/operation_setting" + "github.com/QuantumNous/new-api/types" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +// withBusinessKeywords temporarily replaces the package-level +// BusinessErrorKeywords slice so individual tests can pin the +// classifier to a known keyword list without depending on the +// operator-tunable defaults. It restores the original value via +// t.Cleanup. The atomic setter is used (rather than a direct +// write) so concurrent test runs against the same global don't +// race on the underlying slice. +func withBusinessKeywords(t *testing.T, kws []string) { + t.Helper() + orig := operation_setting.BusinessErrorKeywordsSnapshot() + operation_setting.SetBusinessErrorKeywordsForTest(kws) + t.Cleanup(func() { operation_setting.SetBusinessErrorKeywordsForTest(orig) }) +} + +func TestClassifyChannelError_Nil(t *testing.T) { + assert.Equal(t, ChannelErrorUnknown, ClassifyChannelError(nil)) +} + +func TestClassifyChannelError_BusinessByStatusCode(t *testing.T) { + // The default BusinessErrorStatusCodeRanges includes 400/402/403/422/451. + // We test a couple to be robust against config changes. + for _, code := range []int{400, 402, 403, 422, 451} { + err := &types.NewAPIError{StatusCode: code} + assert.Equalf(t, ChannelErrorBusiness, ClassifyChannelError(err), + "expected business for status %d", code) + } +} + +func TestClassifyChannelError_BusinessByKeyword(t *testing.T) { + withBusinessKeywords(t, []string{"overdue", "insufficient balance"}) + + cases := []struct { + name string + msg string + }{ + {"lowercase", "access denied: account is overdue"}, + {"uppercase", "ACCESS DENIED: ACCOUNT IS OVERDUE"}, + {"mixed_case", "Overdue-Payment detected on key"}, + {"alt_keyword", "insufficient balance on request"}, + {"unrelated_message", "this is a normal completion"}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + err := &types.NewAPIError{StatusCode: 500, Err: errorFromString(tc.msg)} + kind := ClassifyChannelError(err) + if tc.name == "unrelated_message" { + // 500 with no business keyword and no temp keyword is + // unknown. We can't strictly assert Unknown here because + // other status-code based heuristics might fire, but the + // message alone must not produce a business + // classification. + assert.NotEqual(t, ChannelErrorBusiness, kind) + } else { + assert.Equal(t, ChannelErrorBusiness, kind, + "expected business for message %q", tc.msg) + } + }) + } +} + +func TestClassifyChannelError_TempByStatusCode(t *testing.T) { + // 5xx codes fall in ShouldRetryByStatusCode's default range. + for _, code := range []int{500, 502, 503, 599} { + err := &types.NewAPIError{StatusCode: code} + assert.Equalf(t, ChannelErrorTemp, ClassifyChannelError(err), + "expected temp for status %d", code) + } +} + +func TestClassifyChannelError_UnknownFor2xx(t *testing.T) { + // 2xx is not a business or temp error; we never see it through this + // path, but the classifier must not return a wrong class. + err := &types.NewAPIError{StatusCode: 200} + assert.Equal(t, ChannelErrorUnknown, ClassifyChannelError(err)) +} + +func TestClassifyChannelError_BusinessBeatsTemp(t *testing.T) { + // A 400 must classify as business, not temp, even though 4xx + // is technically in the retry list (4xx except 400/408 — and + // 400 is excluded from the temp list already). Pin both directions. + err := &types.NewAPIError{StatusCode: 400} + assert.Equal(t, ChannelErrorBusiness, ClassifyChannelError(err)) +} + +func TestClassifyChannelError_BusinessKeywordBeatsTempStatus(t *testing.T) { + // Hypothetical upstream that returns 500 with a body that says + // "suspended". The keyword match must win so the cooldown is + // long (operator chose this keyword for a reason). + withBusinessKeywords(t, []string{"suspended"}) + err := &types.NewAPIError{ + StatusCode: 500, + Err: errorFromString("account suspended, contact support"), + } + assert.Equal(t, ChannelErrorBusiness, ClassifyChannelError(err)) +} + +func TestMessageMatchesBusinessKeyword_EmptyMessage(t *testing.T) { + assert.False(t, messageMatchesBusinessKeyword("")) +} + +func TestMessageMatchesBusinessKeyword_EmptyKeywords(t *testing.T) { + withBusinessKeywords(t, nil) + assert.False(t, messageMatchesBusinessKeyword("anything goes here")) +} + +// errorFromString builds a minimal error for use in test inputs. We +// don't want a real error to leak into production paths, but +// NewAPIError.Err is the source string the classifier matches against. +func errorFromString(s string) error { + return stringError(s) +} + +type stringError string + +func (e stringError) Error() string { return string(e) } + +// Sanity: the test helper must behave like errors.New so a misuse is +// caught fast rather than producing false negatives. +func TestStringErrorHelper(t *testing.T) { + var e error = stringError("hello") + require.True(t, strings.HasPrefix(e.Error(), "hello")) +} diff --git a/setting/operation_setting/operation_setting.go b/setting/operation_setting/operation_setting.go index ef330d1adb87..1b889870ca6a 100644 --- a/setting/operation_setting/operation_setting.go +++ b/setting/operation_setting/operation_setting.go @@ -1,10 +1,25 @@ package operation_setting -import "strings" +import ( + "strings" + "sync" +) var DemoSiteEnabled = false var SelfUseModeEnabled = false +// TempErrorCooldownSeconds is the duration a channel is skipped by the +// selector after a temporary upstream error (e.g. 5xx, 408, 429) before +// being considered again. 0 disables cooldown for temporary errors. +var TempErrorCooldownSeconds = 30 + +// BusinessErrorCooldownSeconds is the duration a channel is skipped by the +// selector after a business error (e.g. 400 overdue-payment, 402 payment +// required, quota keywords) before being considered again. 0 falls back to +// the legacy behaviour (permanent AutoDisabled). Negative values disable +// cooldown entirely (channel is kept enabled but error is still recorded). +var BusinessErrorCooldownSeconds = 3600 + var AutomaticDisableKeywords = []string{ "Your credit balance is too low", "This organization has been disabled.", @@ -15,6 +30,73 @@ var AutomaticDisableKeywords = []string{ "Your account is not authorized", } +// businessErrorKeywordsMu guards BusinessErrorKeywords against torn +// reads. The classifier consults this slice on every upstream error +// (hot path), so a plain `var []string` is not safe: the previous +// FromString implementation cleared the global slice and then +// repopulated it, leaving concurrent readers with a partial or empty +// list mid-update — which silently misclassifies errors and changes +// cooldown behaviour. Writers build into a local slice, then swap +// under the lock; readers copy the slice header under the same lock +// before iterating. +var businessErrorKeywordsMu sync.RWMutex + +// BusinessErrorKeywords are matched (case-insensitive) against the upstream +// error message to classify an error as a business error. Business errors +// are short-circuited (no retry, no other channel for the same error type) +// and routed through the long cooldown / disable path. +var BusinessErrorKeywords = []string{ + "overdue", + "overdue-payment", + "insufficient balance", + "insufficient credit", + "insufficient quota", + "quota exceeded", + "credit balance", + "account is not active", + "account has been disabled", + "account in good standing", + "plan does not include", + "plan expired", + "plan not subscribed", + "please make sure your account is in good standing", + "please recharge", + "please top up", + "please upgrade", + "please add payment method", + "billing issue", + "unpaid", + "payment required", + "suspended", + "terminated", +} + +// BusinessErrorKeywordsSnapshot returns a copy of the current +// BusinessErrorKeywords slice under the read lock. Callers iterate +// over the returned slice without holding the lock. The copy is O(n) +// and the read path is the hot path (every error classifier call), +// so this is the right trade-off: writers are rare (config reload) +// and pay an O(n) copy each, readers are common and pay a +// constant-time RLock. +func BusinessErrorKeywordsSnapshot() []string { + businessErrorKeywordsMu.RLock() + defer businessErrorKeywordsMu.RUnlock() + out := make([]string, len(BusinessErrorKeywords)) + copy(out, BusinessErrorKeywords) + return out +} + +// SetBusinessErrorKeywordsForTest atomically replaces the keyword +// slice. Intended only for tests that need to pin the classifier to +// a known keyword list. The caller is responsible for restoring the +// previous value via t.Cleanup if the test mutates the slice in +// place; a non-empty `next` is taken over verbatim. +func SetBusinessErrorKeywordsForTest(next []string) { + businessErrorKeywordsMu.Lock() + BusinessErrorKeywords = next + businessErrorKeywordsMu.Unlock() +} + func AutomaticDisableKeywordsToString() string { return strings.Join(AutomaticDisableKeywords, "\n") } @@ -30,3 +112,28 @@ func AutomaticDisableKeywordsFromString(s string) { } } } + +func BusinessErrorKeywordsToString() string { + businessErrorKeywordsMu.RLock() + defer businessErrorKeywordsMu.RUnlock() + return strings.Join(BusinessErrorKeywords, "\n") +} + +func BusinessErrorKeywordsFromString(s string) { + ak := strings.Split(s, "\n") + next := make([]string, 0, len(ak)) + for _, k := range ak { + k = strings.TrimSpace(k) + k = strings.ToLower(k) + if k != "" { + next = append(next, k) + } + } + // Build the new slice into a local variable first, then swap under + // the write lock. This prevents concurrent readers (the classifier + // on every upstream error) from observing an empty or partially + // populated global list. + businessErrorKeywordsMu.Lock() + BusinessErrorKeywords = next + businessErrorKeywordsMu.Unlock() +} diff --git a/setting/operation_setting/status_code_ranges.go b/setting/operation_setting/status_code_ranges.go index 14cfacad71ed..80904bffd5db 100644 --- a/setting/operation_setting/status_code_ranges.go +++ b/setting/operation_setting/status_code_ranges.go @@ -16,6 +16,23 @@ type StatusCodeRange struct { var AutomaticDisableStatusCodeRanges = []StatusCodeRange{{Start: 401, End: 401}} +// BusinessErrorStatusCodeRanges identifies status codes that indicate a +// business-side failure (account/quota/billing) rather than a temporary +// upstream fault. Errors matching these codes skip the retry path and are +// routed to the long-cooldown / disable path so users do not get stuck +// in a retry loop on a channel that will keep returning the same error. +// Default covers the most common 4xx "client/account" responses seen in +// real provider APIs (Aliyun DashScope, OpenAI plan errors, Anthropic +// billing errors, Google Cloud billing, etc.). Operators can override +// via the system settings UI. +var BusinessErrorStatusCodeRanges = []StatusCodeRange{ + {Start: 400, End: 400}, + {Start: 402, End: 402}, + {Start: 403, End: 403}, + {Start: 422, End: 422}, + {Start: 451, End: 451}, +} + // Default behavior matches legacy hardcoded retry rules in controller/relay.go shouldRetry: // retry for 1xx, 3xx, 4xx(except 400/408), 5xx(except 504/524), and no retry for 2xx. var AutomaticRetryStatusCodeRanges = []StatusCodeRange{ @@ -54,6 +71,27 @@ func ShouldDisableByStatusCode(code int) bool { return shouldMatchStatusCodeRanges(AutomaticDisableStatusCodeRanges, code) } +func BusinessErrorStatusCodesToString() string { + return statusCodeRangesToString(BusinessErrorStatusCodeRanges) +} + +func BusinessErrorStatusCodesFromString(s string) error { + ranges, err := ParseHTTPStatusCodeRanges(s) + if err != nil { + return err + } + BusinessErrorStatusCodeRanges = ranges + return nil +} + +// IsBusinessErrorStatusCode returns true when the given upstream status code +// indicates a business-side failure (account/quota/billing) rather than a +// transient upstream fault. The default list covers 400/402/403/422/451 and +// can be overridden in the system settings UI. +func IsBusinessErrorStatusCode(code int) bool { + return shouldMatchStatusCodeRanges(BusinessErrorStatusCodeRanges, code) +} + func AutomaticRetryStatusCodesToString() string { return statusCodeRangesToString(AutomaticRetryStatusCodeRanges) } diff --git a/setting/operation_setting/status_code_ranges_test.go b/setting/operation_setting/status_code_ranges_test.go index 4e292a3681a9..9c3fd2db5070 100644 --- a/setting/operation_setting/status_code_ranges_test.go +++ b/setting/operation_setting/status_code_ranges_test.go @@ -3,13 +3,14 @@ package operation_setting import ( "testing" + "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" ) func TestParseHTTPStatusCodeRanges_CommaSeparated(t *testing.T) { ranges, err := ParseHTTPStatusCodeRanges("401,403,500-599") require.NoError(t, err) - require.Equal(t, []StatusCodeRange{ + assert.Equal(t, []StatusCodeRange{ {Start: 401, End: 401}, {Start: 403, End: 403}, {Start: 500, End: 599}, @@ -19,69 +20,122 @@ func TestParseHTTPStatusCodeRanges_CommaSeparated(t *testing.T) { func TestParseHTTPStatusCodeRanges_MergeAndNormalize(t *testing.T) { ranges, err := ParseHTTPStatusCodeRanges("500-505,504,401,403,402") require.NoError(t, err) - require.Equal(t, []StatusCodeRange{ + assert.Equal(t, []StatusCodeRange{ {Start: 401, End: 403}, {Start: 500, End: 505}, }, ranges) } -func TestParseHTTPStatusCodeRanges_Invalid(t *testing.T) { - _, err := ParseHTTPStatusCodeRanges("99,600,foo,500-400,500-") +func TestParseHTTPStatusCodeRanges_Empty(t *testing.T) { + ranges, err := ParseHTTPStatusCodeRanges("") + require.NoError(t, err) + assert.Empty(t, ranges) +} + +func TestParseHTTPStatusCodeRanges_InvalidEntry(t *testing.T) { + _, err := ParseHTTPStatusCodeRanges("200,abc,500") require.Error(t, err) } -func TestParseHTTPStatusCodeRanges_NoComma_IsInvalid(t *testing.T) { - _, err := ParseHTTPStatusCodeRanges("401 403") +func TestParseHTTPStatusCodeRanges_OutOfRange(t *testing.T) { + _, err := ParseHTTPStatusCodeRanges("99") require.Error(t, err) } -func TestShouldDisableByStatusCode(t *testing.T) { - orig := AutomaticDisableStatusCodeRanges - t.Cleanup(func() { AutomaticDisableStatusCodeRanges = orig }) +func TestShouldMatchStatusCodeRanges(t *testing.T) { + ranges := []StatusCodeRange{{Start: 400, End: 410}} + assert.True(t, shouldMatchStatusCodeRanges(ranges, 400)) + assert.True(t, shouldMatchStatusCodeRanges(ranges, 410)) + assert.False(t, shouldMatchStatusCodeRanges(ranges, 399)) + assert.False(t, shouldMatchStatusCodeRanges(ranges, 411)) + // Out-of-band codes are never matched. + assert.False(t, shouldMatchStatusCodeRanges(ranges, 99)) + assert.False(t, shouldMatchStatusCodeRanges(ranges, 600)) +} - AutomaticDisableStatusCodeRanges = []StatusCodeRange{ - {Start: 401, End: 403}, - {Start: 500, End: 599}, +func TestStatusCodeRangesToString(t *testing.T) { + ranges := []StatusCodeRange{ + {Start: 200, End: 200}, + {Start: 400, End: 410}, + {Start: 500, End: 503}, } - - require.True(t, ShouldDisableByStatusCode(401)) - require.True(t, ShouldDisableByStatusCode(403)) - require.False(t, ShouldDisableByStatusCode(404)) - require.True(t, ShouldDisableByStatusCode(500)) - require.False(t, ShouldDisableByStatusCode(200)) + assert.Equal(t, "200,400-410,500-503", statusCodeRangesToString(ranges)) } -func TestShouldRetryByStatusCode(t *testing.T) { +func TestAutomaticRetryStatusCodesFromString(t *testing.T) { orig := AutomaticRetryStatusCodeRanges t.Cleanup(func() { AutomaticRetryStatusCodeRanges = orig }) - AutomaticRetryStatusCodeRanges = []StatusCodeRange{ - {Start: 429, End: 429}, - {Start: 500, End: 599}, - } + require.NoError(t, AutomaticRetryStatusCodesFromString("500-503,401,402,403")) + assert.Equal(t, []StatusCodeRange{ + {Start: 401, End: 403}, + {Start: 500, End: 503}, + }, AutomaticRetryStatusCodeRanges) +} - require.True(t, ShouldRetryByStatusCode(429)) - require.True(t, ShouldRetryByStatusCode(500)) - require.False(t, ShouldRetryByStatusCode(504)) - require.False(t, ShouldRetryByStatusCode(524)) - require.False(t, ShouldRetryByStatusCode(400)) - require.False(t, ShouldRetryByStatusCode(200)) +func TestAutomaticDisableStatusCodesFromString(t *testing.T) { + orig := AutomaticDisableStatusCodeRanges + t.Cleanup(func() { AutomaticDisableStatusCodeRanges = orig }) + + require.NoError(t, AutomaticDisableStatusCodesFromString("401")) + assert.Equal(t, []StatusCodeRange{{Start: 401, End: 401}}, + AutomaticDisableStatusCodeRanges) +} + +func TestShouldDisableByStatusCode(t *testing.T) { + assert.True(t, ShouldDisableByStatusCode(401)) + assert.False(t, ShouldDisableByStatusCode(500)) } func TestShouldRetryByStatusCode_DefaultMatchesLegacyBehavior(t *testing.T) { - require.False(t, ShouldRetryByStatusCode(200)) - require.False(t, ShouldRetryByStatusCode(400)) - require.True(t, ShouldRetryByStatusCode(401)) - require.False(t, ShouldRetryByStatusCode(408)) - require.True(t, ShouldRetryByStatusCode(429)) - require.True(t, ShouldRetryByStatusCode(500)) - require.False(t, ShouldRetryByStatusCode(504)) - require.False(t, ShouldRetryByStatusCode(524)) - require.True(t, ShouldRetryByStatusCode(599)) + assert.False(t, ShouldRetryByStatusCode(200)) + assert.False(t, ShouldRetryByStatusCode(400)) + assert.True(t, ShouldRetryByStatusCode(401)) + assert.False(t, ShouldRetryByStatusCode(408)) + assert.True(t, ShouldRetryByStatusCode(429)) + assert.True(t, ShouldRetryByStatusCode(500)) + assert.False(t, ShouldRetryByStatusCode(504)) + assert.False(t, ShouldRetryByStatusCode(524)) + assert.True(t, ShouldRetryByStatusCode(599)) } func TestIsAlwaysSkipRetryStatusCode(t *testing.T) { - require.True(t, IsAlwaysSkipRetryStatusCode(504)) - require.True(t, IsAlwaysSkipRetryStatusCode(524)) - require.False(t, IsAlwaysSkipRetryStatusCode(500)) + assert.True(t, IsAlwaysSkipRetryStatusCode(504)) + assert.True(t, IsAlwaysSkipRetryStatusCode(524)) + assert.False(t, IsAlwaysSkipRetryStatusCode(500)) +} + +func TestIsBusinessErrorStatusCode_Default(t *testing.T) { + // Restore the default ranges if a previous test mutated them. + orig := BusinessErrorStatusCodeRanges + t.Cleanup(func() { BusinessErrorStatusCodeRanges = orig }) + BusinessErrorStatusCodeRanges = []StatusCodeRange{ + {Start: 400, End: 400}, + {Start: 402, End: 402}, + {Start: 403, End: 403}, + {Start: 422, End: 422}, + {Start: 451, End: 451}, + } + for _, code := range []int{400, 402, 403, 422, 451} { + assert.Truef(t, IsBusinessErrorStatusCode(code), + "expected business for status %d", code) + } + // Negative cases: success, 5xx, 408 (408 is a temp error, not business), + // 401 (auth, neither business nor auto-disabled by the default rules). + for _, code := range []int{200, 401, 408, 500, 503, 599} { + assert.Falsef(t, IsBusinessErrorStatusCode(code), + "expected non-business for status %d", code) + } +} + +func TestBusinessErrorStatusCodesFromString(t *testing.T) { + orig := BusinessErrorStatusCodeRanges + t.Cleanup(func() { BusinessErrorStatusCodeRanges = orig }) + + require.NoError(t, BusinessErrorStatusCodesFromString("400,402-403,422")) + assert.Equal(t, []StatusCodeRange{ + {Start: 400, End: 400}, + {Start: 402, End: 403}, + {Start: 422, End: 422}, + }, BusinessErrorStatusCodeRanges) } diff --git a/types/channel_error.go b/types/channel_error.go index f2d72bf536e3..79cf65dcea05 100644 --- a/types/channel_error.go +++ b/types/channel_error.go @@ -5,8 +5,15 @@ type ChannelError struct { ChannelType int `json:"channel_type"` ChannelName string `json:"channel_name"` IsMultiKey bool `json:"is_multi_key"` - AutoBan bool `json:"auto_ban"` - UsingKey string `json:"using_key"` + // KeyIndex is the index of the specific key that produced the + // error, when the channel is in multi-key mode. nil means the + // caller did not have a key index available (single-key channels + // or paths that skip the per-key lookup). The cooldown handler + // uses this to scope the skip to the bad key instead of the + // whole channel. + KeyIndex *int `json:"key_index,omitempty"` + AutoBan bool `json:"auto_ban"` + UsingKey string `json:"using_key"` } func NewChannelError(channelId int, channelType int, channelName string, isMultiKey bool, usingKey string, autoBan bool) *ChannelError {