From 20c4bf5ae59e4cb94669165436bb57ee16c41cb3 Mon Sep 17 00:00:00 2001 From: feeeei Date: Wed, 14 Oct 2026 11:51:56 +0800 Subject: [PATCH 1/2] =?UTF-8?q?fix(gemini):=20Vertex=20SA=20=E5=B7=B2?= =?UTF-8?q?=E4=B8=8D=E5=AD=98=E5=9C=A8=E6=97=B6=E8=B4=A6=E5=8F=B7=E7=BD=AE?= =?UTF-8?q?=E4=B8=BA=20error=20=E5=B9=B6=E6=8D=A2=E5=8F=B7?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 换取 SA access token 时 token 端点返回 400 {"error":"invalid_grant","error_description":"Invalid grant: account not found"} 表示服务账号已被删除，凭据再也换不到 token；此前三条 Gemini 转发路径把它当作 本地构造请求失败直接回 502，不换号也不改账号状态，账号持续被调度 - 识别该错误后账号置为 error 并停止调度（Service account token rejected (400): …）， 记录凭据获取阶段（stage=account_auth）的 Ops 事件，请求交给 handler 换号 - 只匹配 400 + invalid_grant + account not found：invalid_grant 还涵盖时钟偏差、 签名无效等情况，其他换 token 失败保持原行为 - token 端点地址改为包级变量，测试中替换为本地服务 --- .../gemini_chat_completions_compat_service.go | 3 + .../service/gemini_messages_compat_service.go | 6 + .../service/gemini_vertex_token_error.go | 60 +++++ .../service/gemini_vertex_token_error_test.go | 207 ++++++++++++++++++ .../service/vertex_service_account.go | 39 +++- 5 files changed, 313 insertions(+), 2 deletions(-) create mode 100644 backend/internal/service/gemini_vertex_token_error.go create mode 100644 backend/internal/service/gemini_vertex_token_error_test.go diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index 1ad2b487e8..669baff99d 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -118,6 +118,9 @@ func (s *GeminiMessagesCompatService) forwardClaudeBodyAsChatCompletions( if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return nil, err } + if failoverErr := s.handleVertexServiceAccountNotFound(ctx, c, account, err); failoverErr != nil { + return nil, failoverErr + } return nil, s.writeChatCompletionsError(c, http.StatusBadGateway, "upstream_error", err.Error()) } requestIDHeader = idHeader diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 06b3e84306..4717b6404f 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -716,6 +716,9 @@ func (s *GeminiMessagesCompatService) Forward(ctx context.Context, c *gin.Contex if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return nil, err } + if failoverErr := s.handleVertexServiceAccountNotFound(ctx, c, account, err); failoverErr != nil { + return nil, failoverErr + } // Local build error: don't retry. if strings.Contains(err.Error(), "missing project_id") { return nil, s.writeClaudeError(c, http.StatusBadRequest, "invalid_request_error", err.Error()) @@ -1245,6 +1248,9 @@ func (s *GeminiMessagesCompatService) ForwardNative(ctx context.Context, c *gin. if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { return nil, err } + if failoverErr := s.handleVertexServiceAccountNotFound(ctx, c, account, err); failoverErr != nil { + return nil, failoverErr + } // Local build error: don't retry. if strings.Contains(err.Error(), "missing project_id") { return nil, s.writeGoogleError(c, http.StatusBadRequest, err.Error()) diff --git a/backend/internal/service/gemini_vertex_token_error.go b/backend/internal/service/gemini_vertex_token_error.go new file mode 100644 index 0000000000..7577dcca5b --- /dev/null +++ b/backend/internal/service/gemini_vertex_token_error.go @@ -0,0 +1,60 @@ +package service + +import ( + "context" + "log/slog" + "net/http" + + "github.com/gin-gonic/gin" +) + +// vertexServiceAccountNotFoundReason 标识换取 SA access token 时 SA 已不存在。 +const vertexServiceAccountNotFoundReason GatewayFailureReason = "vertex_service_account_not_found" + +// geminiVertexCredentialFailoverBody 是 SA 凭据失效时挂在 failover 错误上的 Google 格式错误体， +// 供 handler 在换号耗尽后做透传规则匹配与错误信息提取。 +var geminiVertexCredentialFailoverBody = []byte(`{"error":{"code":401,"message":"Upstream authentication failed","status":"UNAUTHENTICATED"}}`) + +// handleVertexServiceAccountNotFound 处理 Gemini 三条转发路径（原生 / Claude 兼容 / +// Chat Completions 兼容）构造上游请求时换取 SA access token 失败、且 SA 已不存在的情况： +// +// - 账号置为 error 并停止调度，错误信息带上 token 端点的 error_description； +// - 记录凭据获取阶段（stage=account_auth）的 Ops 事件，token 端点的状态码只出现在 Message 中； +// - 返回 *UpstreamFailoverError(401)，由 handler 的 failover 循环换号。 +// +// 不是该类失败时返回 nil，由调用方继续处理。本函数不写响应。 +func (s *GeminiMessagesCompatService) handleVertexServiceAccountNotFound(ctx context.Context, c *gin.Context, account *Account, err error) *UpstreamFailoverError { + tokenErr, ok := vertexServiceAccountNotFoundError(err) + if !ok { + return nil + } + + msg := "Service account token rejected (400): " + tokenErr.Description + stateCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), openAIAccountStateUpdateTimeout) + defer cancel() + if s.rateLimitService != nil { + s.rateLimitService.handleAuthError(stateCtx, account, msg) + } else if s.accountRepo != nil { + if setErr := s.accountRepo.SetError(stateCtx, account.ID, msg); setErr != nil { + slog.Warn("account_set_error_failed", "account_id", account.ID, "error", setErr) + } + } + + event := OpsUpstreamErrorEvent{ + Platform: account.Platform, + AccountID: account.ID, + AccountName: account.Name, + Kind: "failover", + Stage: string(GatewayFailureStageAccountAuth), + Scope: string(GatewayFailureScopeAccount), + Reason: string(vertexServiceAccountNotFoundReason), + Message: sanitizeUpstreamErrorMessage(err.Error()), + } + event.ProxyID, event.ProxyName = opsUpstreamProxyAttribution(account) + appendOpsUpstreamError(c, event) + + return &UpstreamFailoverError{ + StatusCode: http.StatusUnauthorized, + ResponseBody: geminiVertexCredentialFailoverBody, + } +} diff --git a/backend/internal/service/gemini_vertex_token_error_test.go b/backend/internal/service/gemini_vertex_token_error_test.go new file mode 100644 index 0000000000..57e98f1da1 --- /dev/null +++ b/backend/internal/service/gemini_vertex_token_error_test.go @@ -0,0 +1,207 @@ +//go:build unit + +package service + +import ( + "bytes" + "context" + "crypto/rand" + "crypto/rsa" + "crypto/x509" + "encoding/json" + "encoding/pem" + "errors" + "fmt" + "net/http" + "net/http/httptest" + "strings" + "testing" + + "github.com/Wei-Shaw/sub2api/internal/config" + "github.com/gin-gonic/gin" + "github.com/stretchr/testify/require" +) + +// --------------------------------------------------------------------------- +// 换取 SA access token 时 SA 已不存在（token 端点 400 invalid_grant +// "Invalid grant: account not found"）的契约，三条转发路径一致： +// - 账号置为 error 并停止调度，错误信息带 token 端点的 error_description； +// - 不发出上游请求，返回 *UpstreamFailoverError(401) 交给 handler 换号，service 不写响应； +// - Ops 事件记为凭据获取阶段（stage=account_auth）； +// - 其他换 token 失败不改账号状态、不换号，直接回写 502。 +// --------------------------------------------------------------------------- + +const vertexAccountNotFoundTokenBody = `{"error":"invalid_grant","error_description":"Invalid grant: account not found"}` + +func newVertexTokenTestAccount(t *testing.T) *Account { + t.Helper() + privateKey, err := rsa.GenerateKey(rand.Reader, 2048) + require.NoError(t, err) + pemBytes := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}) + saJSON, err := json.Marshal(map[string]string{ + "type": "service_account", + "project_id": "proj-1", + "private_key_id": "kid-1", + "private_key": string(pemBytes), + "client_email": "svc@proj-1.iam.gserviceaccount.com", + }) + require.NoError(t, err) + return &Account{ + ID: 820, + Platform: PlatformGemini, + Type: AccountTypeServiceAccount, + Credentials: map[string]any{ + "service_account_json": string(saJSON), + "project_id": "proj-1", + "location": "global", + }, + } +} + +func useVertexTokenEndpoint(t *testing.T, status int, body string) { + t.Helper() + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + _, _ = w.Write([]byte(body)) + })) + t.Cleanup(server.Close) + original := vertexServiceAccountTokenURL + vertexServiceAccountTokenURL = server.URL + t.Cleanup(func() { vertexServiceAccountTokenURL = original }) +} + +func newVertexTokenErrorService() (*GeminiMessagesCompatService, *geminiCompatHTTPUpstreamStub, *errorPolicyRepoStub) { + repo := &errorPolicyRepoStub{} + httpStub := &geminiCompatHTTPUpstreamStub{} + svc := &GeminiMessagesCompatService{ + accountRepo: repo, + httpUpstream: httpStub, + tokenProvider: NewGeminiTokenProvider(repo, nil, nil), + rateLimitService: NewRateLimitService(repo, nil, &config.Config{}, nil, nil), + cfg: &config.Config{}, + } + return svc, httpStub, repo +} + +func requireVertexAccountNotFoundHandled(t *testing.T, c *gin.Context, err error, httpStub *geminiCompatHTTPUpstreamStub, repo *errorPolicyRepoStub) { + t.Helper() + var failoverErr *UpstreamFailoverError + require.True(t, errors.As(err, &failoverErr), "SA 不存在应换号，got %T: %v", err, err) + require.Equal(t, http.StatusUnauthorized, failoverErr.StatusCode) + require.Equal(t, string(geminiVertexCredentialFailoverBody), string(failoverErr.ResponseBody)) + require.True(t, failoverErr.ShouldRetryNextAccount()) + require.False(t, failoverErr.RetryableOnSameAccount) + + require.Equal(t, 1, repo.setErrCalls, "SA 不存在应把账号置为 error") + require.Equal(t, "Service account token rejected (400): Invalid grant: account not found", repo.lastErrorMsg) + require.Zero(t, httpStub.calls, "换不到 token 时不应发出上游请求") + require.False(t, c.Writer.Written(), "换号场景不应写客户端响应") + + raw, ok := c.Get(OpsUpstreamErrorsKey) + require.True(t, ok, "应记录 ops 上游错误事件") + events := raw.([]*OpsUpstreamErrorEvent) + require.Len(t, events, 1) + require.Equal(t, "failover", events[0].Kind) + require.Equal(t, int64(820), events[0].AccountID) + require.Equal(t, string(GatewayFailureStageAccountAuth), events[0].Stage) + require.Equal(t, string(GatewayFailureScopeAccount), events[0].Scope) + require.Equal(t, string(vertexServiceAccountNotFoundReason), events[0].Reason) + require.Zero(t, events[0].UpstreamStatusCode, "换 token 失败不占用上游状态码") + require.Contains(t, events[0].Message, "service account token request returned 400: Invalid grant: account not found") +} + +func TestExchangeVertexServiceAccountToken_AccountNotFoundIsRecognized(t *testing.T) { + useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody) + key, err := parseVertexServiceAccountKey(newVertexTokenTestAccount(t)) + require.NoError(t, err) + + _, _, err = exchangeVertexServiceAccountToken(contextWithTestTimeout(t), key, "") + + require.EqualError(t, err, "service account token request returned 400: Invalid grant: account not found") + tokenErr, ok := vertexServiceAccountNotFoundError(fmt.Errorf("build request: %w", err)) + require.True(t, ok) + require.Equal(t, "invalid_grant", tokenErr.Code) + require.Equal(t, "Invalid grant: account not found", tokenErr.Description) +} + +func TestVertexServiceAccountNotFoundError_OnlyMatchesAccountNotFound(t *testing.T) { + tests := []struct { + name string + err error + }{ + {"other_invalid_grant", &vertexTokenRequestError{StatusCode: 400, Code: "invalid_grant", Description: "Invalid JWT Signature."}}, + {"not_400", &vertexTokenRequestError{StatusCode: 401, Code: "invalid_grant", Description: "Invalid grant: account not found"}}, + {"other_code", &vertexTokenRequestError{StatusCode: 400, Code: "invalid_request", Description: "Invalid grant: account not found"}}, + {"server_error", &vertexTokenRequestError{StatusCode: 500, Message: "backend error"}}, + {"plain_error", errors.New("service account token request returned 400: Invalid grant: account not found")}, + } + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + _, ok := vertexServiceAccountNotFoundError(tt.err) + require.False(t, ok) + }) + } +} + +func TestGeminiForwardNative_VertexServiceAccountNotFoundDisablesAndFailsOver(t *testing.T) { + gin.SetMode(gin.TestMode) + useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody) + svc, httpStub, repo := newVertexTokenErrorService() + c, _ := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, newVertexTokenTestAccount(t), + "gemini-3.8-flash", "streamGenerateContent", true, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + + require.Nil(t, result) + requireVertexAccountNotFoundHandled(t, c, err, httpStub, repo) +} + +func TestGeminiForward_VertexServiceAccountNotFoundDisablesAndFailsOver(t *testing.T) { + gin.SetMode(gin.TestMode) + useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody) + svc, httpStub, repo := newVertexTokenErrorService() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gemini-3.8-flash","max_tokens":16,"messages":[{"role":"user","content":"hello"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body)) + + result, err := svc.Forward(context.Background(), c, newVertexTokenTestAccount(t), body) + + require.Nil(t, result) + requireVertexAccountNotFoundHandled(t, c, err, httpStub, repo) +} + +func TestGeminiForwardAsChatCompletions_VertexServiceAccountNotFoundDisablesAndFailsOver(t *testing.T) { + gin.SetMode(gin.TestMode) + useVertexTokenEndpoint(t, http.StatusBadRequest, vertexAccountNotFoundTokenBody) + svc, httpStub, repo := newVertexTokenErrorService() + rec := httptest.NewRecorder() + c, _ := gin.CreateTestContext(rec) + body := []byte(`{"model":"gemini-3.8-flash","messages":[{"role":"user","content":"hi"}]}`) + c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body)) + + result, err := svc.ForwardAsChatCompletions(context.Background(), c, newVertexTokenTestAccount(t), body) + + require.Nil(t, result) + requireVertexAccountNotFoundHandled(t, c, err, httpStub, repo) +} + +func TestGeminiForwardNative_OtherVertexTokenErrorKeepsAccountState(t *testing.T) { + gin.SetMode(gin.TestMode) + useVertexTokenEndpoint(t, http.StatusBadRequest, `{"error":"invalid_grant","error_description":"Invalid JWT Signature."}`) + svc, httpStub, repo := newVertexTokenErrorService() + c, rec := newGeminiNativeTestContext(t) + + result, err := svc.ForwardNative(context.Background(), c, newVertexTokenTestAccount(t), + "gemini-3.8-flash", "generateContent", false, []byte(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + + require.Nil(t, result) + require.Error(t, err) + var failoverErr *UpstreamFailoverError + require.False(t, errors.As(err, &failoverErr), "未识别的换 token 失败不换号") + require.Zero(t, repo.setErrCalls, "未识别的换 token 失败不改账号状态") + require.Zero(t, httpStub.calls) + require.Equal(t, http.StatusBadGateway, rec.Code) + require.True(t, strings.Contains(rec.Body.String(), "Invalid JWT Signature.")) +} diff --git a/backend/internal/service/vertex_service_account.go b/backend/internal/service/vertex_service_account.go index 2a19a50e21..fdc7d41e48 100644 --- a/backend/internal/service/vertex_service_account.go +++ b/backend/internal/service/vertex_service_account.go @@ -39,6 +39,9 @@ var ( vertexAnthropicAlreadyDatedIDPattern = regexp.MustCompile(`^.+@[0-9]{8}$`) ) +// vertexServiceAccountTokenURL 是换取 SA access token 的端点，测试中替换为本地服务。 +var vertexServiceAccountTokenURL = vertexDefaultTokenURL + type vertexServiceAccountKey struct { Type string `json:"type"` ProjectID string `json:"project_id"` @@ -138,7 +141,7 @@ func parseVertexServiceAccountJSON(raw []byte) (*vertexServiceAccountKey, error) return nil, errors.New("service account json missing project_id") } // Always use the well-known Google token endpoint to prevent SSRF via crafted token_uri. - key.TokenURI = vertexDefaultTokenURL + key.TokenURI = vertexServiceAccountTokenURL return &key, nil } @@ -277,7 +280,12 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc if msg == "" { msg = string(bytes.TrimSpace(body)) } - return "", 0, fmt.Errorf("service account token request returned %d: %s", resp.StatusCode, msg) + return "", 0, &vertexTokenRequestError{ + StatusCode: resp.StatusCode, + Code: strings.TrimSpace(parsed.Error), + Description: strings.TrimSpace(parsed.ErrorDesc), + Message: msg, + } } if strings.TrimSpace(parsed.AccessToken) == "" { return "", 0, errors.New("service account token response missing access_token") @@ -292,6 +300,33 @@ func exchangeVertexServiceAccountToken(ctx context.Context, key *vertexServiceAc return parsed.AccessToken, ttl, nil } +// vertexTokenRequestError 是 token 端点对换取 SA access token 请求返回的非 2xx 响应。 +type vertexTokenRequestError struct { + StatusCode int + Code string // 响应的 error 字段，如 invalid_grant + Description string // 响应的 error_description 字段 + Message string +} + +func (e *vertexTokenRequestError) Error() string { + return fmt.Sprintf("service account token request returned %d: %s", e.StatusCode, e.Message) +} + +// vertexServiceAccountNotFoundError 识别 SA 已不存在：token 端点返回 +// 400 {"error":"invalid_grant","error_description":"Invalid grant: account not found"}。 +// 该凭据再也换不到 token，账号需要人工更换凭据。 +func vertexServiceAccountNotFoundError(err error) (*vertexTokenRequestError, bool) { + var tokenErr *vertexTokenRequestError + if !errors.As(err, &tokenErr) { + return nil, false + } + if tokenErr.StatusCode != http.StatusBadRequest || tokenErr.Code != "invalid_grant" || + !strings.Contains(strings.ToLower(tokenErr.Description), "account not found") { + return nil, false + } + return tokenErr, true +} + // buildVertexGeminiAccountURL 按账号构造 Vertex Gemini 请求地址：未配置 project_id 的 API Key 账号 // 走 Express mode 全局端点，其余账号走项目级端点。 func buildVertexGeminiAccountURL(account *Account, model, action string, stream bool) (string, error) { From 14cd5871752dabfc38069e9a7d87f282da0c9eda Mon Sep 17 00:00:00 2001 From: feeeei Date: Wed, 14 Oct 2026 11:52:02 +0800 Subject: [PATCH 2/2] =?UTF-8?q?fix(gemini):=20=E6=AF=8F=E4=B8=AA=E4=B8=8A?= =?UTF-8?q?=E6=B8=B8=E9=94=99=E8=AF=AF=E5=93=8D=E5=BA=94=E5=8F=AA=E6=A3=80?= =?UTF-8?q?=E6=9F=A5=E4=B8=80=E6=AC=A1=E9=94=99=E8=AF=AF=E7=AD=96=E7=95=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - 三条 Gemini 转发路径在重试循环内和循环外各对最终上游错误响应检查一次错误策略。 配了账号级 401 临时不可调度规则的 API Key / Vertex 账号首次收到 401 时，第一次 检查写入 401 暂停记录，第二次检查读回该记录判定为重复 401，同一请求内即升级为 error 永久停用 - 循环内的检查结果带到循环外复用；循环内未检查的最终响应（Claude 兼容层的签名 400 分支、原生路径的 cachedContent 透传分支）仍在循环外检查 --- .../gemini_chat_completions_compat_service.go | 12 +- .../gemini_error_policy_single_check_test.go | 201 ++++++++++++++++++ .../service/gemini_messages_compat_service.go | 51 +++-- 3 files changed, 243 insertions(+), 21 deletions(-) create mode 100644 backend/internal/service/gemini_error_policy_single_check_test.go diff --git a/backend/internal/service/gemini_chat_completions_compat_service.go b/backend/internal/service/gemini_chat_completions_compat_service.go index 669baff99d..c877406340 100644 --- a/backend/internal/service/gemini_chat_completions_compat_service.go +++ b/backend/internal/service/gemini_chat_completions_compat_service.go @@ -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 已由策略层裁决完毕。 diff --git a/backend/internal/service/gemini_error_policy_single_check_test.go b/backend/internal/service/gemini_error_policy_single_check_test.go new file mode 100644 index 0000000000..0ba4344c53 --- /dev/null +++ b/backend/internal/service/gemini_error_policy_single_check_test.go @@ -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, "循环内未检查的最终响应应在循环外检查") +} diff --git a/backend/internal/service/gemini_messages_compat_service.go b/backend/internal/service/gemini_messages_compat_service.go index 4717b6404f..67a5411604 100644 --- a/backend/internal/service/gemini_messages_compat_service.go +++ b/backend/internal/service/gemini_messages_compat_service.go @@ -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 判断上游响应是否跳过同号重试与错误策略直接处理：引用缓存不存在，