mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-14 12:02:27 +08:00
fix(account): 显式模型级 401 规则在网关故障转移入口生效
OpenAI 转发 fastpath、HandleUpstreamError 及两处池模式分支对 401 一律跳过
临时不可调度规则,显式 scope=model 的 401 规则在 API Key 账号上仍落到
handleAuthError 整号置 error。
- 新增 tempUnschedRulesApply:401 仅在首条命中规则按模型级生效时交给规则,
其余保持认证错误处理
- 凭据级 401(OAuth 账号与 Spark 影子、OpenAI token_invalidated/token_revoked、
{"detail":"Unauthorized"})不受规则作用范围影响,各入口沿用原有处理
- 模型级命中不参与重复 401 升级;fastpath 不再因模型级命中封整号
This commit is contained in:
@@ -158,7 +158,8 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
|
||||
// Isolate a custom temporary-unschedulable match to the known upstream
|
||||
// model before entering the generic account error path. This keeps the
|
||||
// account available to other models and avoids the account runtime blocker.
|
||||
if s.rateLimitService != nil && statusCode != http.StatusUnauthorized && len(canonicalModel) > 0 && strings.TrimSpace(canonicalModel[0]) != "" &&
|
||||
if s.rateLimitService != nil && len(canonicalModel) > 0 && strings.TrimSpace(canonicalModel[0]) != "" &&
|
||||
tempUnschedRulesApply(account, statusCode, responseBody, canonicalModel[0]) &&
|
||||
s.rateLimitService.HandleTempUnschedulable(stateCtx, account, statusCode, responseBody, canonicalModel[0]) {
|
||||
return true
|
||||
}
|
||||
@@ -173,7 +174,8 @@ func (s *OpenAIGatewayService) handleOpenAIAccountUpstreamError(ctx context.Cont
|
||||
return false
|
||||
}
|
||||
shouldDisable := s.rateLimitService.HandleUpstreamError(stateCtx, account, statusCode, headers, responseBody)
|
||||
modelTempMatched := statusCode != http.StatusUnauthorized && tempUnschedulableModel(stateCtx, nil) != "" &&
|
||||
tempModel := tempUnschedulableModel(stateCtx, nil)
|
||||
modelTempMatched := tempModel != "" && tempUnschedRulesApply(account, statusCode, responseBody, tempModel) &&
|
||||
len(matchTempUnschedulableRules(account, statusCode, responseBody)) > 0
|
||||
if shouldDisable && !modelTempMatched {
|
||||
s.BlockAccountScheduling(account, time.Time{}, "upstream_disable")
|
||||
|
||||
@@ -310,8 +310,9 @@ func (s *RateLimitService) CheckErrorPolicy(ctx context.Context, account *Accoun
|
||||
}
|
||||
if account.IsPoolMode() {
|
||||
// 池模式只跳过默认账号状态处理;管理员显式配置的临时不可调度规则仍应生效。
|
||||
// 401 保留现有认证错误语义,避免改变重复 401 的升级行为。
|
||||
if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||||
// 401 只在非凭据级且命中显式模型级规则时由规则处理,其余保留现有认证错误语义。
|
||||
if tempUnschedRulesApply(account, statusCode, responseBody, tempUnschedulableModel(ctx, nil)) &&
|
||||
s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||||
return ErrorPolicyTempUnscheduled
|
||||
}
|
||||
return ErrorPolicySkipped
|
||||
@@ -337,9 +338,10 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc
|
||||
customErrorCodesEnabled := account.IsCustomErrorCodesEnabled()
|
||||
|
||||
// 池模式默认不标记本地账号状态;但管理员显式配置的临时不可调度规则优先。
|
||||
// 401 保留现有认证错误语义,不在这里改变池模式的认证处理。
|
||||
// 401 只在非凭据级且命中显式模型级规则时由规则处理,其余保留池模式现有的认证处理。
|
||||
if account.IsPoolMode() && !customErrorCodesEnabled {
|
||||
if statusCode != http.StatusUnauthorized && s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||||
if tempUnschedRulesApply(account, statusCode, responseBody, tempUnschedulableModel(ctx, nil)) &&
|
||||
s.tryTempUnschedulable(ctx, account, statusCode, responseBody) {
|
||||
return true
|
||||
}
|
||||
slog.Info("pool_mode_error_skipped", "account_id", account.ID, "status_code", statusCode)
|
||||
@@ -385,9 +387,9 @@ func (s *RateLimitService) HandleUpstreamError(ctx context.Context, account *Acc
|
||||
}
|
||||
}
|
||||
|
||||
// 先尝试临时不可调度规则(401除外)
|
||||
// 先尝试临时不可调度规则(401 仅限非凭据级且命中显式模型级规则)
|
||||
// 如果匹配成功,直接返回,不执行后续禁用逻辑
|
||||
if statusCode != 401 {
|
||||
if tempUnschedRulesApply(account, statusCode, responseBody, tempUnschedulableModel(ctx, requestedModel)) {
|
||||
if s.tryTempUnschedulable(ctx, account, statusCode, responseBody, firstRequestedModel(requestedModel)) {
|
||||
return true
|
||||
}
|
||||
@@ -2564,6 +2566,45 @@ func firstRequestedModel(requestedModel []string) string {
|
||||
return strings.TrimSpace(requestedModel[0])
|
||||
}
|
||||
|
||||
// tempUnschedRulesApply 决定本次错误是否交给临时不可调度规则处理。
|
||||
// 401 默认留给认证错误处理;只有首条命中规则按模型级生效时,才由规则只冷却该模型。
|
||||
func tempUnschedRulesApply(account *Account, statusCode int, responseBody []byte, modelKey string) bool {
|
||||
if statusCode != http.StatusUnauthorized {
|
||||
return true
|
||||
}
|
||||
matches := matchTempUnschedulableRules(account, statusCode, responseBody)
|
||||
return len(matches) > 0 && tempUnschedHitUsesModelScope(account, matches[0].rule, modelKey, statusCode, responseBody)
|
||||
}
|
||||
|
||||
// tempUnschedHitUsesModelScope 决定一次命中是否按模型级冷却:凭据级 401 始终按账号处理,
|
||||
// 其余由规则作用范围决定。
|
||||
func tempUnschedHitUsesModelScope(account *Account, rule TempUnschedulableRule, modelKey string, statusCode int, responseBody []byte) bool {
|
||||
if statusCode == http.StatusUnauthorized && unauthorizedRequiresCredentialHandling(account, responseBody) {
|
||||
return false
|
||||
}
|
||||
return tempUnschedRuleUsesModelScope(rule, modelKey, statusCode)
|
||||
}
|
||||
|
||||
// unauthorizedRequiresCredentialHandling 识别凭据级 401:OAuth 账号与 Spark 影子
|
||||
// (失效缓存、刷新窗口、缺 refresh_token 停用),以及 OpenAI 明确的凭据失效
|
||||
// (token_invalidated / token_revoked、{"detail":"Unauthorized"})。
|
||||
func unauthorizedRequiresCredentialHandling(account *Account, responseBody []byte) bool {
|
||||
if account == nil {
|
||||
return false
|
||||
}
|
||||
if account.Type == AccountTypeOAuth || account.IsShadow() {
|
||||
return true
|
||||
}
|
||||
if account.Platform != PlatformOpenAI {
|
||||
return false
|
||||
}
|
||||
switch extractUpstreamErrorCode(responseBody) {
|
||||
case "token_invalidated", "token_revoked":
|
||||
return true
|
||||
}
|
||||
return gjson.GetBytes(responseBody, "detail").String() == "Unauthorized"
|
||||
}
|
||||
|
||||
// tempUnschedRuleUsesModelScope 决定一次命中按模型级冷却还是账号级暂停。
|
||||
// 显式 scope 完全由规则决定;scope 为空时模型已知且非 401 走模型级。
|
||||
// 模型未知时无法隔离,一律账号级。
|
||||
@@ -2645,10 +2686,17 @@ func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Ac
|
||||
if !account.IsTempUnschedulableEnabled() {
|
||||
return false
|
||||
}
|
||||
matches := matchTempUnschedulableRules(account, statusCode, responseBody)
|
||||
if len(matches) == 0 {
|
||||
return false
|
||||
}
|
||||
modelKey := tempUnschedulableModel(ctx, requestedModel)
|
||||
// 401 首次命中可临时不可调度(给 token 刷新窗口);
|
||||
// 若历史上已因 401 进入过临时不可调度,则本次应升级为 error(返回 false 交由默认错误逻辑处理)。
|
||||
// Antigravity 跳过:其 401 由 applyErrorPolicy 的 temp_unschedulable_rules 自行控制,无需升级逻辑。
|
||||
if statusCode == http.StatusUnauthorized && account.Platform != PlatformAntigravity {
|
||||
// 模型级命中只冷却模型、不写账号级状态,不参与升级。
|
||||
if statusCode == http.StatusUnauthorized && account.Platform != PlatformAntigravity &&
|
||||
!tempUnschedHitUsesModelScope(account, matches[0].rule, modelKey, statusCode, responseBody) {
|
||||
reason := account.TempUnschedulableReason
|
||||
// 缓存可能没有 reason,从 DB 回退读取
|
||||
if reason == "" {
|
||||
@@ -2662,8 +2710,8 @@ func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Ac
|
||||
return false
|
||||
}
|
||||
}
|
||||
for _, match := range matchTempUnschedulableRules(account, statusCode, responseBody) {
|
||||
if s.triggerTempUnschedulable(ctx, account, match.rule, match.ruleIndex, statusCode, match.matchedKeyword, responseBody, tempUnschedulableModel(ctx, requestedModel)) {
|
||||
for _, match := range matches {
|
||||
if s.triggerTempUnschedulable(ctx, account, match.rule, match.ruleIndex, statusCode, match.matchedKeyword, responseBody, modelKey) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
@@ -2714,7 +2762,7 @@ func (s *RateLimitService) triggerTempUnschedulable(ctx context.Context, account
|
||||
now := time.Now()
|
||||
until := now.Add(time.Duration(rule.DurationMinutes) * time.Minute)
|
||||
modelKey := firstRequestedModel(requestedModel)
|
||||
modelScoped := tempUnschedRuleUsesModelScope(rule, modelKey, statusCode)
|
||||
modelScoped := tempUnschedHitUsesModelScope(account, rule, modelKey, statusCode, responseBody)
|
||||
|
||||
state := &TempUnschedState{
|
||||
UntilUnix: until.Unix(),
|
||||
|
||||
@@ -9,12 +9,19 @@ import (
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type tempUnschedScopeRepoStub struct {
|
||||
modelNotFoundAccountRepoStub
|
||||
tempReasons []string
|
||||
errorCalls int
|
||||
}
|
||||
|
||||
func (r *tempUnschedScopeRepoStub) SetError(context.Context, int64, string) error {
|
||||
r.errorCalls++
|
||||
return nil
|
||||
}
|
||||
|
||||
func (r *tempUnschedScopeRepoStub) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
|
||||
@@ -50,6 +57,36 @@ func newTempUnschedScopeAccount(scope string, extra map[string]any) *Account {
|
||||
|
||||
var tempUnschedScopeBody = []byte(`{"error":{"message":"upstream unavailable"}}`)
|
||||
|
||||
var (
|
||||
tempUnschedScope401Body = []byte(`{"error":{"message":"unauthorized"}}`)
|
||||
tempUnschedScopeRevoked401Body = []byte(`{"error":{"code":"token_revoked","message":"unauthorized"}}`)
|
||||
tempUnschedScopeDetail401Body = []byte(`{"detail":"Unauthorized"}`)
|
||||
tempUnschedScopePrevious401Mark = `{"status_code":401,"until_unix":1735689600}`
|
||||
)
|
||||
|
||||
func tempUnschedScope401Rule(scope string) map[string]any {
|
||||
rule := map[string]any{
|
||||
"error_code": float64(http.StatusUnauthorized),
|
||||
"keywords": []any{"unauthorized"},
|
||||
"duration_minutes": float64(10),
|
||||
}
|
||||
if scope != "" {
|
||||
rule["scope"] = scope
|
||||
}
|
||||
return rule
|
||||
}
|
||||
|
||||
func newTempUnschedScope401Account(platform string, scopes ...string) *Account {
|
||||
account := newTempUnschedScopeAccount("", nil)
|
||||
account.Platform = platform
|
||||
rules := make([]any, 0, len(scopes))
|
||||
for _, scope := range scopes {
|
||||
rules = append(rules, tempUnschedScope401Rule(scope))
|
||||
}
|
||||
account.Credentials["temp_unschedulable_rules"] = rules
|
||||
return account
|
||||
}
|
||||
|
||||
func TestGetTempUnschedulableRules_ParsesScope(t *testing.T) {
|
||||
mk := func(keyword string, scope any) map[string]any {
|
||||
rule := map[string]any{
|
||||
@@ -154,3 +191,133 @@ func TestHandleTempUnschedulable_ExplicitModelScopeFallsBackToAccountWhenModelUn
|
||||
require.Equal(t, 1, repo.tempCalls)
|
||||
require.Empty(t, repo.modelRateLimitCalls)
|
||||
}
|
||||
|
||||
func TestOpenAIFailoverSideEffects_401FollowsTempUnschedScope(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
scopes []string
|
||||
body []byte
|
||||
poolMode bool
|
||||
oauth bool
|
||||
wantModel bool
|
||||
wantTemp bool
|
||||
wantSetError bool
|
||||
wantBlocked bool
|
||||
}{
|
||||
{name: "显式模型级 只冷却模型", scopes: []string{TempUnschedScopeModel}, wantModel: true},
|
||||
{name: "显式模型级 池模式 只冷却模型", scopes: []string{TempUnschedScopeModel}, poolMode: true, wantModel: true},
|
||||
{name: "缺省 走认证错误处理", scopes: []string{""}, wantSetError: true, wantBlocked: true},
|
||||
{name: "显式账号级 走认证错误处理", scopes: []string{TempUnschedScopeAccount}, wantSetError: true, wantBlocked: true},
|
||||
{name: "首条命中为账号级 走认证错误处理", scopes: []string{TempUnschedScopeAccount, TempUnschedScopeModel}, wantSetError: true, wantBlocked: true},
|
||||
{name: "缺省 池模式 保持跳过", scopes: []string{""}, poolMode: true},
|
||||
{name: "显式模型级 token_revoked 走认证错误处理", scopes: []string{TempUnschedScopeModel}, body: tempUnschedScopeRevoked401Body, wantSetError: true, wantBlocked: true},
|
||||
{name: "显式模型级 detail Unauthorized 走认证错误处理", scopes: []string{TempUnschedScopeModel}, body: tempUnschedScopeDetail401Body, wantSetError: true, wantBlocked: true},
|
||||
{name: "显式模型级 token_revoked 池模式 保持跳过", scopes: []string{TempUnschedScopeModel}, body: tempUnschedScopeRevoked401Body, poolMode: true},
|
||||
{name: "显式模型级 OAuth 走刷新窗口", scopes: []string{TempUnschedScopeModel}, oauth: true, wantTemp: true, wantBlocked: true},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
repo := &tempUnschedScopeRepoStub{}
|
||||
gateway := &OpenAIGatewayService{rateLimitService: &RateLimitService{accountRepo: repo, cfg: &config.Config{}}}
|
||||
account := newTempUnschedScope401Account(PlatformOpenAI, tc.scopes...)
|
||||
if tc.poolMode {
|
||||
account.Credentials["pool_mode"] = true
|
||||
}
|
||||
if tc.oauth {
|
||||
account.Type = AccountTypeOAuth
|
||||
account.Credentials["refresh_token"] = "rt"
|
||||
}
|
||||
body := tc.body
|
||||
if body == nil {
|
||||
body = tempUnschedScope401Body
|
||||
}
|
||||
|
||||
require.True(t, gateway.shouldFailoverOpenAIUpstreamResponse(account, http.StatusUnauthorized, "unauthorized", body))
|
||||
gateway.handleFailoverSideEffects(context.Background(), &http.Response{
|
||||
StatusCode: http.StatusUnauthorized,
|
||||
Header: http.Header{},
|
||||
}, account, body, "gpt-5.4")
|
||||
|
||||
require.Equal(t, tc.wantTemp, repo.tempCalls == 1, "SetTempUnschedulable=%d", repo.tempCalls)
|
||||
require.Equal(t, tc.wantSetError, repo.errorCalls == 1, "SetError=%d", repo.errorCalls)
|
||||
require.Equal(t, tc.wantBlocked, gateway.isOpenAIAccountRuntimeBlocked(account))
|
||||
if tc.wantModel {
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "gpt-5.4", repo.modelRateLimitCalls[0].scope)
|
||||
} else {
|
||||
require.Empty(t, repo.modelRateLimitCalls)
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHandleUpstreamError_ExplicitModelScope401(t *testing.T) {
|
||||
t.Run("模型已知 只冷却模型", func(t *testing.T) {
|
||||
repo := &tempUnschedScopeRepoStub{}
|
||||
svc := &RateLimitService{accountRepo: repo}
|
||||
account := newTempUnschedScope401Account(PlatformAnthropic, TempUnschedScopeModel)
|
||||
|
||||
shouldDisable := svc.HandleUpstreamError(context.Background(), account, http.StatusUnauthorized, http.Header{}, tempUnschedScope401Body, "claude-sonnet-4-5")
|
||||
|
||||
require.True(t, shouldDisable)
|
||||
require.Zero(t, repo.errorCalls)
|
||||
require.Zero(t, repo.tempCalls)
|
||||
require.Len(t, repo.modelRateLimitCalls, 1)
|
||||
require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].scope)
|
||||
})
|
||||
|
||||
t.Run("模型未知 按账号级规则走认证错误处理", func(t *testing.T) {
|
||||
repo := &tempUnschedScopeRepoStub{}
|
||||
svc := &RateLimitService{accountRepo: repo}
|
||||
account := newTempUnschedScope401Account(PlatformAnthropic, TempUnschedScopeModel)
|
||||
|
||||
shouldDisable := svc.HandleUpstreamError(context.Background(), account, http.StatusUnauthorized, http.Header{}, tempUnschedScope401Body)
|
||||
|
||||
require.True(t, shouldDisable)
|
||||
require.Equal(t, 1, repo.errorCalls)
|
||||
require.Empty(t, repo.modelRateLimitCalls)
|
||||
})
|
||||
}
|
||||
|
||||
func TestCheckErrorPolicy_401FollowsTempUnschedScope(t *testing.T) {
|
||||
cases := []struct {
|
||||
name string
|
||||
scope string
|
||||
body []byte
|
||||
oauth bool
|
||||
previous401 bool
|
||||
want ErrorPolicyResult
|
||||
wantModel bool
|
||||
wantTemp bool
|
||||
}{
|
||||
{name: "显式模型级 不参与重复 401 升级", scope: TempUnschedScopeModel, previous401: true, want: ErrorPolicyTempUnscheduled, wantModel: true},
|
||||
{name: "缺省 第二次 401 升级", previous401: true, want: ErrorPolicyNone},
|
||||
{name: "显式模型级 token_revoked 按账号暂停", scope: TempUnschedScopeModel, body: tempUnschedScopeRevoked401Body, want: ErrorPolicyTempUnscheduled, wantTemp: true},
|
||||
{name: "显式模型级 detail Unauthorized 按账号暂停", scope: TempUnschedScopeModel, body: tempUnschedScopeDetail401Body, want: ErrorPolicyTempUnscheduled, wantTemp: true},
|
||||
{name: "显式模型级 OAuth 按账号暂停", scope: TempUnschedScopeModel, oauth: true, want: ErrorPolicyTempUnscheduled, wantTemp: true},
|
||||
{name: "显式模型级 凭据级第二次 401 升级", scope: TempUnschedScopeModel, body: tempUnschedScopeDetail401Body, previous401: true, want: ErrorPolicyNone},
|
||||
}
|
||||
for _, tc := range cases {
|
||||
t.Run(tc.name, func(t *testing.T) {
|
||||
repo := &tempUnschedScopeRepoStub{}
|
||||
svc := &RateLimitService{accountRepo: repo}
|
||||
account := newTempUnschedScope401Account(PlatformOpenAI, tc.scope)
|
||||
if tc.oauth {
|
||||
account.Type = AccountTypeOAuth
|
||||
}
|
||||
if tc.previous401 {
|
||||
account.TempUnschedulableReason = tempUnschedScopePrevious401Mark
|
||||
}
|
||||
body := tc.body
|
||||
if body == nil {
|
||||
body = tempUnschedScope401Body
|
||||
}
|
||||
|
||||
result := svc.CheckErrorPolicy(context.Background(), account, http.StatusUnauthorized, body, "gpt-5.4")
|
||||
|
||||
require.Equal(t, tc.want, result)
|
||||
require.Equal(t, tc.wantTemp, repo.tempCalls == 1, "SetTempUnschedulable=%d", repo.tempCalls)
|
||||
require.Equal(t, tc.wantModel, len(repo.modelRateLimitCalls) == 1)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user