feat: 二次 401 直接升级为错误状态,添加 DB 回退确保生效

账号首次 401 仅临时不可调度,给予 token 刷新窗口;若恢复后再次 401
说明凭证确实失效,直接升级为错误状态以避免反复无效调度。

- 缓存中 reason 为空时从 DB 回退读取,防止升级判断失效
- ClearError 同时清除临时不可调度状态,管理员恢复后重新给予一次机会
- 管理后台账号列表添加"临时不可调度"状态筛选
- 补充 DB 回退场景单元测试
This commit is contained in:
kyx236
2026-03-04 20:25:15 +08:00
parent fe1d46a8ea
commit 6aa8cbbf20
6 changed files with 236 additions and 2 deletions

View File

@@ -88,6 +88,49 @@ func TestCheckErrorPolicy(t *testing.T) {
body: []byte(`overloaded service`),
expected: ErrorPolicyTempUnscheduled,
},
{
name: "temp_unschedulable_401_first_hit_returns_temp_unscheduled",
account: &Account{
ID: 14,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(401),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 401,
body: []byte(`unauthorized`),
expected: ErrorPolicyTempUnscheduled,
},
{
name: "temp_unschedulable_401_second_hit_upgrades_to_none",
account: &Account{
ID: 15,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
TempUnschedulableReason: `{"status_code":401,"until_unix":1735689600}`,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(401),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 401,
body: []byte(`unauthorized`),
expected: ErrorPolicyNone,
},
{
name: "temp_unschedulable_body_miss_returns_none",
account: &Account{

View File

@@ -122,6 +122,28 @@ func TestCheckErrorPolicy_GeminiAccounts(t *testing.T) {
body: []byte(`overloaded service`),
expected: ErrorPolicyTempUnscheduled,
},
{
name: "gemini_apikey_temp_unschedulable_401_second_hit_returns_none",
account: &Account{
ID: 105,
Type: AccountTypeAPIKey,
Platform: PlatformGemini,
TempUnschedulableReason: `{"status_code":401,"until_unix":1735689600}`,
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(401),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
},
},
},
statusCode: 401,
body: []byte(`unauthorized`),
expected: ErrorPolicyNone,
},
{
name: "gemini_custom_codes_override_temp_unschedulable",
account: &Account{

View File

@@ -698,6 +698,22 @@ func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Ac
if !account.IsTempUnschedulableEnabled() {
return false
}
// 401 首次命中可临时不可调度(给 token 刷新窗口);
// 若历史上已因 401 进入过临时不可调度,则本次应升级为 error返回 false 交由默认错误逻辑处理)。
if statusCode == http.StatusUnauthorized {
reason := account.TempUnschedulableReason
// 缓存可能没有 reason从 DB 回退读取
if reason == "" {
if dbAcc, err := s.accountRepo.GetByID(ctx, account.ID); err == nil && dbAcc != nil {
reason = dbAcc.TempUnschedulableReason
}
}
if wasTempUnschedByStatusCode(reason, statusCode) {
slog.Info("401_escalated_to_error", "account_id", account.ID,
"reason", "previous temp-unschedulable was also 401")
return false
}
}
rules := account.GetTempUnschedulableRules()
if len(rules) == 0 {
return false
@@ -729,6 +745,22 @@ func (s *RateLimitService) tryTempUnschedulable(ctx context.Context, account *Ac
return false
}
func wasTempUnschedByStatusCode(reason string, statusCode int) bool {
if statusCode <= 0 {
return false
}
reason = strings.TrimSpace(reason)
if reason == "" {
return false
}
var state TempUnschedState
if err := json.Unmarshal([]byte(reason), &state); err != nil {
return false
}
return state.StatusCode == statusCode
}
func matchTempUnschedKeyword(bodyLower string, keywords []string) string {
if bodyLower == "" {
return ""

View File

@@ -0,0 +1,119 @@
//go:build unit
package service
import (
"context"
"net/http"
"testing"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/stretchr/testify/require"
)
// dbFallbackRepoStub extends errorPolicyRepoStub with a configurable DB account
// returned by GetByID, simulating cache miss + DB fallback.
type dbFallbackRepoStub struct {
errorPolicyRepoStub
dbAccount *Account // returned by GetByID when non-nil
}
func (r *dbFallbackRepoStub) GetByID(ctx context.Context, id int64) (*Account, error) {
if r.dbAccount != nil && r.dbAccount.ID == id {
return r.dbAccount, nil
}
return nil, nil // not found, no error
}
func TestCheckErrorPolicy_401_DBFallback_Escalates(t *testing.T) {
// Scenario: cache account has empty TempUnschedulableReason (cache miss),
// but DB account has a previous 401 record → should escalate to ErrorPolicyNone.
repo := &dbFallbackRepoStub{
dbAccount: &Account{
ID: 20,
TempUnschedulableReason: `{"status_code":401,"until_unix":1735689600}`,
},
}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 20,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
TempUnschedulableReason: "", // cache miss — reason is empty
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(401),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
},
},
}
result := svc.CheckErrorPolicy(context.Background(), account, http.StatusUnauthorized, []byte(`unauthorized`))
require.Equal(t, ErrorPolicyNone, result, "401 with DB fallback showing previous 401 should escalate to ErrorPolicyNone")
}
func TestCheckErrorPolicy_401_DBFallback_NoDBRecord_FirstHit(t *testing.T) {
// Scenario: cache account has empty TempUnschedulableReason,
// DB also has no previous 401 record → should NOT escalate (first hit → temp unscheduled).
repo := &dbFallbackRepoStub{
dbAccount: &Account{
ID: 21,
TempUnschedulableReason: "", // DB also empty
},
}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 21,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
TempUnschedulableReason: "",
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(401),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
},
},
}
result := svc.CheckErrorPolicy(context.Background(), account, http.StatusUnauthorized, []byte(`unauthorized`))
require.Equal(t, ErrorPolicyTempUnscheduled, result, "401 first hit with no DB record should temp-unschedule")
}
func TestCheckErrorPolicy_401_DBFallback_DBError_FirstHit(t *testing.T) {
// Scenario: cache account has empty TempUnschedulableReason,
// DB lookup returns nil (not found) → should treat as first hit → temp unscheduled.
repo := &dbFallbackRepoStub{
dbAccount: nil, // GetByID returns nil, nil
}
svc := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
account := &Account{
ID: 22,
Type: AccountTypeOAuth,
Platform: PlatformAntigravity,
TempUnschedulableReason: "",
Credentials: map[string]any{
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{
map[string]any{
"error_code": float64(401),
"keywords": []any{"unauthorized"},
"duration_minutes": float64(10),
},
},
},
}
result := svc.CheckErrorPolicy(context.Background(), account, http.StatusUnauthorized, []byte(`unauthorized`))
require.Equal(t, ErrorPolicyTempUnscheduled, result, "401 first hit with DB not found should temp-unschedule")
}