fix(gemini): 每个上游错误响应只检查一次错误策略

- 三条 Gemini 转发路径在重试循环内和循环外各对最终上游错误响应检查一次错误策略。
  配了账号级 401 临时不可调度规则的 API Key / Vertex 账号首次收到 401 时,第一次
  检查写入 401 暂停记录,第二次检查读回该记录判定为重复 401,同一请求内即升级为
  error 永久停用
- 循环内的检查结果带到循环外复用;循环内未检查的最终响应(Claude 兼容层的签名 400
  分支、原生路径的 cachedContent 透传分支)仍在循环外检查
This commit is contained in:
feeeei
2026-10-14 12:00:44 +08:00
parent 20c4bf5ae5
commit 14cd587175
3 changed files with 243 additions and 21 deletions
@@ -112,6 +112,8 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
)
var resp *http.Response
// checkedPolicy 是循环内对当前 resp 的错误策略检查结果,未检查时为 nil。
var checkedPolicy *ErrorPolicyResult
for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
upstreamReq, idHeader, err := buildReq(ctx)
if err != nil {
@@ -130,11 +132,11 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
return nil, s.handleUpstreamTransportError(ctx, c, account, err)
}
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
resp = rebuilt
checked, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel)
resp = rebuilt
checkedPolicy = checked
if checked != nil && *checked != ErrorPolicyNone {
break
} else {
resp = rebuilt
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
@@ -199,7 +201,7 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions(
respBody := s.readUpstreamErrorBody(resp)
policy := ErrorPolicyNone
if s.rateLimitService != nil {
policy = s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
policy = s.finalErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel, checkedPolicy)
}
// 与 messages 兼容层一致:只有 None / Matched 才走账号状态处理。
// Skipped(池模式、或自定义错误码未命中)与 TempUnscheduled 已由策略层裁决完毕。
@@ -0,0 +1,201 @@
//go:build unit
package service
import (
"bytes"
"context"
"errors"
"io"
"net/http"
"net/http/httptest"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/config"
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
"github.com/gin-gonic/gin"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// Gemini 三条转发路径对每个上游错误响应只检查一次错误策略:
// - 重试循环内已检查的最终响应,循环外的错误处理复用该结果;
// - 账号级 401 规则首次命中只临时不可调度,账号上已有 401 暂停记录时才升级为 error;
// - 循环内未检查的最终响应(如 Claude 兼容层的非签名 400)在循环外检查。
// ---------------------------------------------------------------------------
type geminiSingleCheckRepoStub struct {
errorPolicyRepoStub
tempReason string
}
func (r *geminiSingleCheckRepoStub) SetTempUnschedulable(_ context.Context, _ int64, _ time.Time, reason string) error {
r.tempCalls++
r.tempReason = reason
return nil
}
func (r *geminiSingleCheckRepoStub) GetByID(_ context.Context, id int64) (*Account, error) {
return &Account{ID: id, TempUnschedulableReason: r.tempReason}, nil
}
type geminiSequenceUpstreamResponse struct {
status int
body string
}
type geminiSequenceUpstreamStub struct {
responses []geminiSequenceUpstreamResponse
calls int
}
func (s *geminiSequenceUpstreamStub) Do(*http.Request, string, int64, int) (*http.Response, error) {
r := s.responses[min(s.calls, len(s.responses)-1)]
s.calls++
return &http.Response{
StatusCode: r.status,
Header: http.Header{"Content-Type": []string{"application/json"}},
Body: io.NopCloser(strings.NewReader(r.body)),
}, nil
}
func (s *geminiSequenceUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ *tlsfingerprint.Profile) (*http.Response, error) {
return s.Do(req, proxyURL, accountID, concurrency)
}
func newGeminiSingleCheckService(responses ...geminiSequenceUpstreamResponse) (*GeminiMessagesCompatService, *geminiSequenceUpstreamStub, *geminiSingleCheckRepoStub) {
repo := &geminiSingleCheckRepoStub{}
rateLimit := NewRateLimitService(repo, nil, &config.Config{}, nil, nil)
upstream := &geminiSequenceUpstreamStub{responses: responses}
svc := &GeminiMessagesCompatService{
accountRepo: repo,
httpUpstream: upstream,
rateLimitService: rateLimit,
cfg: &config.Config{},
}
return svc, upstream, repo
}
func geminiAPIKeyAccountWithTempRule(rule map[string]any) *Account {
return &Account{
ID: 830,
Platform: PlatformGemini,
Type: AccountTypeAPIKey,
Credentials: map[string]any{
"api_key": "test-key",
"temp_unschedulable_enabled": true,
"temp_unschedulable_rules": []any{rule},
},
}
}
const geminiSingleCheck401Body = `{"error":{"code":401,"message":"API key expired. Please renew the API key.","status":"UNAUTHENTICATED"}}`
func geminiAccountScoped401Rule() map[string]any {
return map[string]any{
"error_code": float64(401),
"keywords": []any{"expired"},
"duration_minutes": float64(30),
"scope": TempUnschedScopeAccount,
}
}
// geminiSingleCheckForwarders 以三条转发路径各发一次请求,返回转发错误与 gin.Context。
var geminiSingleCheckForwarders = []struct {
name string
forward func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context)
}{
{
name: "native",
forward: func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context) {
c, _ := newGeminiNativeTestContext(t)
_, err := svc.ForwardNative(context.Background(), c, account, "gemini-2.5-flash", "generateContent", false,
[]byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`))
return err, c
},
},
{
name: "claude_compat",
forward: func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
body := []byte(`{"model":"gemini-2.5-flash","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
_, err := svc.Forward(context.Background(), c, account, body)
return err, c
},
},
{
name: "chat_completions_compat",
forward: func(t *testing.T, svc *GeminiMessagesCompatService, account *Account) (error, *gin.Context) {
c, _ := gin.CreateTestContext(httptest.NewRecorder())
body := []byte(`{"model":"gemini-2.5-flash","messages":[{"role":"user","content":"hi"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
_, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body)
return err, c
},
},
}
func TestGeminiForward_First401RuleHitOnlyTempUnschedules(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, fw := range geminiSingleCheckForwarders {
t.Run(fw.name, func(t *testing.T) {
svc, upstream, repo := newGeminiSingleCheckService(geminiSequenceUpstreamResponse{status: http.StatusUnauthorized, body: geminiSingleCheck401Body})
err, _ := fw.forward(t, svc, geminiAPIKeyAccountWithTempRule(geminiAccountScoped401Rule()))
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr), "规则命中应换号,got %T: %v", err, err)
require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode)
require.Equal(t, 1, upstream.calls)
require.Equal(t, 1, repo.tempCalls, "首次 401 命中规则应临时不可调度")
require.Zero(t, repo.setErrCalls, "首次 401 命中规则不应升级为 error")
})
}
}
func TestGeminiForward_Repeated401RuleHitEscalatesToError(t *testing.T) {
gin.SetMode(gin.TestMode)
for _, fw := range geminiSingleCheckForwarders {
t.Run(fw.name, func(t *testing.T) {
svc, _, repo := newGeminiSingleCheckService(geminiSequenceUpstreamResponse{status: http.StatusUnauthorized, body: geminiSingleCheck401Body})
repo.tempReason = `{"until_unix":1,"triggered_at_unix":1,"status_code":401,"matched_keyword":"expired","rule_index":0,"scope":"account"}`
_, _ = fw.forward(t, svc, geminiAPIKeyAccountWithTempRule(geminiAccountScoped401Rule()))
require.Zero(t, repo.tempCalls, "已有 401 暂停记录时直接升级,不临时不可调度")
require.Equal(t, 1, repo.setErrCalls, "已有 401 暂停记录时应升级为 error")
require.Contains(t, repo.lastErrorMsg, "Authentication failed (401)")
})
}
}
func TestGeminiForward_UncheckedFinalResponseIsCheckedAfterLoop(t *testing.T) {
gin.SetMode(gin.TestMode)
if testing.Short() {
t.Skip("包含一次 1s 的重试退避")
}
// 第一次 503 在循环内检查后重试;第二次非签名 400 走 Claude 兼容层的签名分支,循环内不检查。
svc, upstream, repo := newGeminiSingleCheckService(
geminiSequenceUpstreamResponse{status: http.StatusServiceUnavailable, body: `{"error":{"code":503,"message":"backend busy","status":"UNAVAILABLE"}}`},
geminiSequenceUpstreamResponse{status: http.StatusBadRequest, body: `{"error":{"code":400,"message":"quota project misconfigured","status":"FAILED_PRECONDITION"}}`},
)
account := geminiAPIKeyAccountWithTempRule(map[string]any{
"error_code": float64(400),
"keywords": []any{"misconfigured"},
"duration_minutes": float64(30),
"scope": TempUnschedScopeAccount,
})
c, _ := gin.CreateTestContext(httptest.NewRecorder())
body := []byte(`{"model":"gemini-2.5-flash","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`)
c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
_, err := svc.Forward(context.Background(), c, account, body)
var failoverErr *UpstreamFailoverError
require.True(t, errors.As(err, &failoverErr), "400 规则命中应换号,got %T: %v", err, err)
require.Equal(t, 2, upstream.calls)
require.Equal(t, 1, repo.tempCalls, "循环内未检查的最终响应应在循环外检查")
}
@@ -709,8 +709,11 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
}
var resp *http.Response
// checkedPolicy 是循环内对当前 resp 的错误策略检查结果,未检查时为 nil。
var checkedPolicy *ErrorPolicyResult
signatureRetryStage := 0
for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
checkedPolicy = nil
upstreamReq, idHeader, err := buildReq(ctx)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
@@ -803,11 +806,11 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
}
// 错误策略优先:匹配则跳过重试直接处理。
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
resp = rebuilt
checked, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel)
resp = rebuilt
checkedPolicy = checked
if checked != nil && *checked != ErrorPolicyNone {
break
} else {
resp = rebuilt
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
@@ -875,7 +878,7 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex
respBody := s.readUpstreamErrorBody(resp)
// 统一错误策略:自定义错误码 + 临时不可调度
if s.rateLimitService != nil {
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
policy := s.finalErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel, checkedPolicy)
switch policy {
case ErrorPolicySkipped:
upstreamReqID := resp.Header.Get(requestIDHeader)
@@ -1242,7 +1245,10 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}
var resp *http.Response
// checkedPolicy 是循环内对当前 resp 的错误策略检查结果,未检查时为 nil。
var checkedPolicy *ErrorPolicyResult
for attempt := 1; attempt <= geminiMaxRetries; attempt++ {
checkedPolicy = nil
upstreamReq, idHeader, err := buildReq(ctx)
if err != nil {
if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
@@ -1290,11 +1296,11 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}
// 错误策略优先:匹配则跳过重试直接处理。
if matched, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel); matched {
resp = rebuilt
checked, rebuilt := s.checkErrorPolicyInLoop(ctx, account, resp, mappedModel)
resp = rebuilt
checkedPolicy = checked
if checked != nil && *checked != ErrorPolicyNone {
break
} else {
resp = rebuilt
}
if resp.StatusCode >= 400 && s.shouldRetryGeminiUpstreamError(account, resp.StatusCode) {
@@ -1410,7 +1416,7 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
// 统一错误策略:自定义错误码 + 临时不可调度
if s.rateLimitService != nil {
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel)
policy := s.finalErrorPolicy(ctx, account, resp.StatusCode, respBody, mappedModel, checkedPolicy)
switch policy {
case ErrorPolicySkipped:
if failoverErr := s.skippedErrorPolicyFailoverError(c, account, resp.StatusCode, respBody, requestID); failoverErr != nil {
@@ -1579,14 +1585,14 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin.
}, nil
}
// checkErrorPolicyInLoop 在重试循环内预检查错误策略。
// 返回 true 表示策略已匹配(调用者应 break),resp 已重建可直接使用。
// 返回 false 表示 ErrorPolicyNone,resp 已重建,调用者继续走重试逻辑。
// checkErrorPolicyInLoop 在重试循环内对上游错误响应检查错误策略,resp 已重建可直接使用。
// checked 为 nil 表示未检查(非错误响应或未配置限流服务);否则指向检查结果,
// 非 ErrorPolicyNone 时调用者应 break 并由循环外的错误处理复用该结果。
func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
ctx context.Context, account *Account, resp *http.Response, mappedModel string,
) (matched bool, rebuilt *http.Response) {
) (checked *ErrorPolicyResult, rebuilt *http.Response) {
if resp.StatusCode < 400 || s.rateLimitService == nil {
return false, resp
return nil, resp
}
body := s.readUpstreamErrorBody(resp)
_ = resp.Body.Close()
@@ -1596,7 +1602,20 @@ func (s *GeminiMessagesCompatService) checkErrorPolicyInLoop(
Body: io.NopCloser(bytes.NewReader(body)),
}
policy := s.rateLimitService.CheckErrorPolicy(ctx, account, resp.StatusCode, body, mappedModel)
return policy != ErrorPolicyNone, rebuilt
return &policy, rebuilt
}
// finalErrorPolicy 返回循环外处理最终上游错误响应时使用的错误策略。每个上游响应只检查一次:
// 错误策略检查会触发临时不可调度规则,并据账号上已记录的 401 暂停把 401 升级为 error。
// checked 非 nil 时它就是重试循环内对同一响应的检查结果,直接复用;否则在此检查。
// 调用方保证 s.rateLimitService 非 nil。
func (s *GeminiMessagesCompatService) finalErrorPolicy(
ctx context.Context, account *Account, statusCode int, body []byte, mappedModel string, checked *ErrorPolicyResult,
) ErrorPolicyResult {
if checked != nil {
return *checked
}
return s.rateLimitService.CheckErrorPolicy(ctx, account, statusCode, body, mappedModel)
}
// cachedContentPassthroughInLoop 判断上游响应是否跳过同号重试与错误策略直接处理:引用缓存不存在,