mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-15 23:22:59 +08:00
Merge pull request #8076 from feeeei/fix/gemini-vertex-sa-and-error-policy
fix(gemini): Vertex SA 已删除时停用账号并换号,401 临时不可调度规则首次命中不再永久停用
This commit is contained in:
@@ -112,12 +112,17 @@ 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 {
|
||||
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
|
||||
@@ -127,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) {
|
||||
@@ -196,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,13 +709,19 @@ 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) {
|
||||
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())
|
||||
@@ -800,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) {
|
||||
@@ -872,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)
|
||||
@@ -1239,12 +1245,18 @@ 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) {
|
||||
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())
|
||||
@@ -1284,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) {
|
||||
@@ -1404,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 {
|
||||
@@ -1573,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()
|
||||
@@ -1590,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 判断上游响应是否跳过同号重试与错误策略直接处理:引用缓存不存在,
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
@@ -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."))
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user