| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404 |
- //go:build unit
- package service
- import (
- "bytes"
- "context"
- "github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
- "github.com/stretchr/testify/require"
- "io"
- "net/http"
- "strings"
- "testing"
- )
- // stubSmartRetryCache 用于 handleSmartRetry 测试的 GatewayCache mock
- // 仅关注 DeleteSessionAccountID 的调用记录
- type stubSmartRetryCache struct {
- GatewayCache // 嵌入接口,未实现的方法 panic(确保只调用预期方法)
- deleteCalls []deleteSessionCall
- }
- type deleteSessionCall struct {
- groupID int64
- sessionHash string
- }
- func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID int64, sessionHash string) error {
- c.deleteCalls = append(c.deleteCalls, deleteSessionCall{groupID: groupID, sessionHash: sessionHash})
- return nil
- }
- // mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream
- type mockSmartRetryUpstream struct {
- responses []*http.Response
- responseBodies [][]byte // 缓存的 response body 字节(用于 repeatLast 重建)
- errors []error
- callIdx int
- calls []string
- requestBodies [][]byte
- repeatLast bool // 超出范围时重复最后一个响应
- }
- func (m *mockSmartRetryUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
- idx := m.callIdx
- m.calls = append(m.calls, req.URL.String())
- if req != nil && req.Body != nil {
- body, _ := io.ReadAll(req.Body)
- m.requestBodies = append(m.requestBodies, body)
- req.Body = io.NopCloser(bytes.NewReader(body))
- } else {
- m.requestBodies = append(m.requestBodies, nil)
- }
- m.callIdx++
- // 确定使用哪个索引
- respIdx := idx
- if respIdx >= len(m.responses) {
- if !m.repeatLast || len(m.responses) == 0 {
- return nil, nil
- }
- respIdx = len(m.responses) - 1
- }
- resp := m.responses[respIdx]
- respErr := m.errors[respIdx]
- if resp == nil {
- return nil, respErr
- }
- // 首次调用时缓存 body 字节
- if respIdx >= len(m.responseBodies) {
- for len(m.responseBodies) <= respIdx {
- m.responseBodies = append(m.responseBodies, nil)
- }
- }
- if m.responseBodies[respIdx] == nil && resp.Body != nil {
- bodyBytes, _ := io.ReadAll(resp.Body)
- _ = resp.Body.Close()
- m.responseBodies[respIdx] = bodyBytes
- }
- // 用缓存的 body 重建 reader(支持重试场景多次读取)
- cloned := *resp
- if m.responseBodies[respIdx] != nil {
- cloned.Body = io.NopCloser(bytes.NewReader(m.responseBodies[respIdx]))
- }
- return &cloned, respErr
- }
- func (m *mockSmartRetryUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, profile *tlsfingerprint.Profile) (*http.Response, error) {
- return m.Do(req, proxyURL, accountID, accountConcurrency)
- }
- // TestHandleSmartRetry_URLLevelRateLimit 测试 URL 级别限流切换
- func TestHandleSmartRetry_URLLevelRateLimit(t *testing.T) {
- account := &Account{
- ID: 1,
- Name: "acc-1",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{"error":{"message":"Resource has been exhausted"}}`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test", "https://ag-2.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionContinueURL, result.action)
- require.Nil(t, result.resp)
- require.Nil(t, result.err)
- require.Nil(t, result.switchError)
- }
- // TestHandleSmartRetry_LongDelay_ReturnsSwitchError 测试 retryDelay >= 阈值时返回 switchError
- func TestHandleSmartRetry_LongDelay_ReturnsSwitchError(t *testing.T) {
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 1,
- Name: "acc-1",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 15s >= 7s 阈值,应该返回 switchError
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- accountRepo: repo,
- isStickySession: true,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.Nil(t, result.resp, "should not return resp when switchError is set")
- require.Nil(t, result.err)
- require.NotNil(t, result.switchError, "should return switchError for long delay")
- require.Equal(t, account.ID, result.switchError.OriginalAccountID)
- require.Equal(t, "claude-sonnet-4-5", result.switchError.RateLimitedModel)
- require.True(t, result.switchError.IsStickySession)
- // 验证模型限流已设置
- require.Len(t, repo.modelRateLimitCalls, 1)
- require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
- }
- // TestHandleSmartRetry_ShortDelay_SmartRetrySuccess 测试智能重试成功
- func TestHandleSmartRetry_ShortDelay_SmartRetrySuccess(t *testing.T) {
- successResp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{successResp},
- errors: []error{nil},
- }
- account := &Account{
- ID: 1,
- Name: "acc-1",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 0.5s < 7s 阈值,应该触发智能重试
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.resp, "should return successful response")
- require.Equal(t, http.StatusOK, result.resp.StatusCode)
- require.Nil(t, result.err)
- require.Nil(t, result.switchError, "should not return switchError on success")
- require.Len(t, upstream.calls, 1, "should have made one retry call")
- }
- // TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError 测试智能重试失败后返回 switchError
- func TestHandleSmartRetry_ShortDelay_SmartRetryFailed_ReturnsSwitchError(t *testing.T) {
- // 智能重试后仍然返回 429(需要提供 1 个响应,因为智能重试最多 1 次)
- failRespBody := `{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`
- failResp1 := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(failRespBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{failResp1},
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 2,
- Name: "acc-2",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 3s < 7s 阈值,应该触发智能重试(最多 1 次)
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: false,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.Nil(t, result.resp, "should not return resp when switchError is set")
- require.Nil(t, result.err)
- require.NotNil(t, result.switchError, "should return switchError after smart retry failed")
- require.Equal(t, account.ID, result.switchError.OriginalAccountID)
- require.Equal(t, "gemini-3-flash", result.switchError.RateLimitedModel)
- require.False(t, result.switchError.IsStickySession)
- // 验证模型限流已设置:Gemini 同时写入精确模型和家族级 scope
- require.Len(t, repo.modelRateLimitCalls, 2)
- require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
- require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
- require.Len(t, upstream.calls, 1, "should have made one retry call (max attempts)")
- }
- // TestHandleSmartRetry_503_ModelCapacityExhausted_RetrySuccess 测试 503 MODEL_CAPACITY_EXHAUSTED 重试成功
- // MODEL_CAPACITY_EXHAUSTED 使用固定 1s 间隔重试,不切换账号
- func TestHandleSmartRetry_503_ModelCapacityExhausted_RetrySuccess(t *testing.T) {
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 3,
- Name: "acc-3",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 503 + MODEL_CAPACITY_EXHAUSTED + 39s(上游 retryDelay 应被忽略,使用固定 1s)
- respBody := []byte(`{
- "error": {
- "code": 503,
- "status": "UNAVAILABLE",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro-high"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "39s"}
- ],
- "message": "No capacity available for model gemini-3-pro-high on the server"
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusServiceUnavailable,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- // mock: 第 1 次重试返回 200 成功
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{
- {StatusCode: http.StatusOK, Header: http.Header{}, Body: io.NopCloser(strings.NewReader(`{"ok":true}`))},
- },
- errors: []error{nil},
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- accountRepo: repo,
- httpUpstream: upstream,
- isStickySession: true,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.resp, "should return successful response")
- require.Equal(t, http.StatusOK, result.resp.StatusCode)
- require.Nil(t, result.err)
- require.Nil(t, result.switchError, "MODEL_CAPACITY_EXHAUSTED should not return switchError")
- // 不应设置模型限流
- require.Empty(t, repo.modelRateLimitCalls, "MODEL_CAPACITY_EXHAUSTED should not set model rate limit")
- require.Len(t, upstream.calls, 1, "should have made one retry call before success")
- }
- // TestHandleSmartRetry_503_ModelCapacityExhausted_ContextCancel 测试 MODEL_CAPACITY_EXHAUSTED 上下文取消
- func TestHandleSmartRetry_503_ModelCapacityExhausted_ContextCancel(t *testing.T) {
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 3,
- Name: "acc-3",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "code": 503,
- "status": "UNAVAILABLE",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro-high"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "39s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusServiceUnavailable,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- // 立即取消上下文,验证重试循环能正确退出
- ctx, cancel := context.WithCancel(context.Background())
- cancel()
- params := antigravityRetryLoopParams{
- ctx: ctx,
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- accountRepo: repo,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, []string{"https://ag-1.test"})
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.Error(t, result.err, "should return context error")
- require.Nil(t, result.switchError, "should not return switchError on context cancel")
- require.Empty(t, repo.modelRateLimitCalls, "should not set model rate limit on context cancel")
- }
- // TestHandleSmartRetry_NonAntigravityAccount_ContinuesDefaultLogic 测试非 Antigravity 平台账号走默认逻辑
- func TestHandleSmartRetry_NonAntigravityAccount_ContinuesDefaultLogic(t *testing.T) {
- account := &Account{
- ID: 4,
- Name: "acc-4",
- Type: AccountTypeAPIKey, // 非 Antigravity 平台账号
- Platform: PlatformAnthropic,
- }
- // 即使是模型限流响应,非 OAuth 账号也应该走默认逻辑
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionContinue, result.action, "non-Antigravity platform account should continue default logic")
- require.Nil(t, result.resp)
- require.Nil(t, result.err)
- require.Nil(t, result.switchError)
- }
- // TestHandleSmartRetry_NonModelRateLimit_ContinuesDefaultLogic 测试非模型限流响应走默认逻辑
- func TestHandleSmartRetry_NonModelRateLimit_ContinuesDefaultLogic(t *testing.T) {
- account := &Account{
- ID: 5,
- Name: "acc-5",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 429 但没有 RATE_LIMIT_EXCEEDED 或 MODEL_CAPACITY_EXHAUSTED
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "5s"}
- ],
- "message": "Quota exceeded"
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionContinue, result.action, "non-model rate limit should continue default logic")
- require.Nil(t, result.resp)
- require.Nil(t, result.err)
- require.Nil(t, result.switchError)
- }
- // TestHandleSmartRetry_ExactlyAtThreshold_ReturnsSwitchError 测试刚好等于阈值时返回 switchError
- func TestHandleSmartRetry_ExactlyAtThreshold_ReturnsSwitchError(t *testing.T) {
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 6,
- Name: "acc-6",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 刚好 7s = 7s 阈值,应该返回 switchError
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-pro"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "7s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- accountRepo: repo,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.Nil(t, result.resp)
- require.NotNil(t, result.switchError, "exactly at threshold should return switchError")
- require.Equal(t, "gemini-pro", result.switchError.RateLimitedModel)
- }
- // TestAntigravityRetryLoop_HandleSmartRetry_SwitchError_Propagates 测试 switchError 正确传播到上层
- func TestAntigravityRetryLoop_HandleSmartRetry_SwitchError_Propagates(t *testing.T) {
- // 模拟 429 + 长延迟的响应
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "30s"}
- ]
- }
- }`)
- rateLimitResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{rateLimitResp},
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 7,
- Name: "acc-7",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- Schedulable: true,
- Status: StatusActive,
- Concurrency: 1,
- }
- svc := &AntigravityGatewayService{}
- result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: true,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- })
- require.Nil(t, result, "should not return result when switchError")
- require.NotNil(t, err, "should return error")
- var switchErr *AntigravityAccountSwitchError
- require.ErrorAs(t, err, &switchErr, "error should be AntigravityAccountSwitchError")
- require.Equal(t, account.ID, switchErr.OriginalAccountID)
- require.Equal(t, "claude-opus-4-6", switchErr.RateLimitedModel)
- require.True(t, switchErr.IsStickySession)
- }
- // TestHandleSmartRetry_NetworkError_ExhaustsRetry 测试网络错误时(maxAttempts=1)直接耗尽重试并切换账号
- func TestHandleSmartRetry_NetworkError_ExhaustsRetry(t *testing.T) {
- // 唯一一次重试遇到网络错误(nil response)
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{nil}, // 返回 nil(模拟网络错误)
- errors: []error{nil}, // mock 不返回 error,靠 nil response 触发
- }
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 8,
- Name: "acc-8",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 0.1s < 7s 阈值,应该触发智能重试
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.Nil(t, result.resp, "should not return resp when switchError is set")
- require.NotNil(t, result.switchError, "should return switchError after network error exhausted retry")
- require.Equal(t, account.ID, result.switchError.OriginalAccountID)
- require.Equal(t, "claude-sonnet-4-5", result.switchError.RateLimitedModel)
- require.Len(t, upstream.calls, 1, "should have made one retry call")
- // 验证模型限流已设置
- require.Len(t, repo.modelRateLimitCalls, 1)
- require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
- }
- // TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit 测试无 retryDelay 时使用默认 1 分钟限流
- func TestHandleSmartRetry_NoRetryDelay_UsesDefaultRateLimit(t *testing.T) {
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 9,
- Name: "acc-9",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 429 + RATE_LIMIT_EXCEEDED + 无 retryDelay → 使用默认 1 分钟限流
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"}
- ],
- "message": "You have exhausted your capacity on this model."
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- accountRepo: repo,
- isStickySession: true,
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.Nil(t, result.resp, "should not return resp when switchError is set")
- require.NotNil(t, result.switchError, "should return switchError for no retryDelay")
- require.Equal(t, "claude-sonnet-4-5", result.switchError.RateLimitedModel)
- require.True(t, result.switchError.IsStickySession)
- // 验证模型限流已设置
- require.Len(t, repo.modelRateLimitCalls, 1)
- require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
- }
- // ---------------------------------------------------------------------------
- // 以下测试覆盖本次改动:
- // 1. antigravitySmartRetryMaxAttempts = 1(仅重试 1 次)
- // 2. 智能重试失败后清除粘性会话绑定(DeleteSessionAccountID)
- // ---------------------------------------------------------------------------
- // TestSmartRetryMaxAttempts_VerifyConstant 验证常量值为 1
- func TestSmartRetryMaxAttempts_VerifyConstant(t *testing.T) {
- require.Equal(t, 1, antigravitySmartRetryMaxAttempts,
- "antigravitySmartRetryMaxAttempts should be 1 to prevent repeated rate limiting")
- }
- // TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_ClearsSession
- // 核心场景:粘性会话 + 短延迟重试失败 → 必须清除粘性绑定
- func TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_ClearsSession(t *testing.T) {
- failRespBody := `{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`
- failResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(failRespBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{failResp},
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 10,
- Name: "acc-10",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: true,
- groupID: 42,
- sessionHash: "sticky-hash-abc",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{cache: cache}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- // 验证返回 switchError
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.switchError)
- require.True(t, result.switchError.IsStickySession, "switchError should carry IsStickySession=true")
- require.Equal(t, account.ID, result.switchError.OriginalAccountID)
- // 核心断言:DeleteSessionAccountID 被调用,且参数正确
- require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID exactly once")
- require.Equal(t, int64(42), cache.deleteCalls[0].groupID)
- require.Equal(t, "sticky-hash-abc", cache.deleteCalls[0].sessionHash)
- // 验证仅重试 1 次
- require.Len(t, upstream.calls, 1, "should make exactly 1 retry call (maxAttempts=1)")
- // 验证模型限流已设置
- require.Len(t, repo.modelRateLimitCalls, 1)
- require.Equal(t, "claude-sonnet-4-5", repo.modelRateLimitCalls[0].modelKey)
- }
- // TestHandleSmartRetry_ShortDelay_NonStickySession_FailedRetry_NoDeleteSession
- // 非粘性会话 + 短延迟重试失败 → 不应调用 DeleteSessionAccountID(sessionHash 为空)
- func TestHandleSmartRetry_ShortDelay_NonStickySession_FailedRetry_NoDeleteSession(t *testing.T) {
- failRespBody := `{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`
- failResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(failRespBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{failResp},
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 11,
- Name: "acc-11",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: false,
- groupID: 42,
- sessionHash: "", // 非粘性会话,sessionHash 为空
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{cache: cache}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.switchError)
- require.False(t, result.switchError.IsStickySession)
- // 核心断言:sessionHash 为空时不应调用 DeleteSessionAccountID
- require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID when sessionHash is empty")
- }
- // TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_NilCache_NoPanic
- // 边界:cache 为 nil 时不应 panic
- func TestHandleSmartRetry_ShortDelay_StickySession_FailedRetry_NilCache_NoPanic(t *testing.T) {
- failRespBody := `{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`
- failResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(failRespBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{failResp},
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- account := &Account{
- ID: 12,
- Name: "acc-12",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: true,
- groupID: 42,
- sessionHash: "sticky-hash-nil-cache",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- // cache 为 nil,不应 panic
- svc := &AntigravityGatewayService{cache: nil}
- require.NotPanics(t, func() {
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.switchError)
- require.True(t, result.switchError.IsStickySession)
- })
- }
- // TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession
- // 重试成功时不应清除粘性会话(只有失败才清除)
- func TestHandleSmartRetry_ShortDelay_StickySession_SuccessRetry_NoDeleteSession(t *testing.T) {
- successResp := &http.Response{
- StatusCode: http.StatusOK,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(`{"result":"ok"}`)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{successResp},
- errors: []error{nil},
- }
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 13,
- Name: "acc-13",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- isStickySession: true,
- groupID: 42,
- sessionHash: "sticky-hash-success",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{cache: cache}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.resp, "should return successful response")
- require.Equal(t, http.StatusOK, result.resp.StatusCode)
- require.Nil(t, result.switchError, "should not return switchError on success")
- // 核心断言:重试成功时不应清除粘性会话
- require.Len(t, cache.deleteCalls, 0, "should NOT call DeleteSessionAccountID on successful retry")
- }
- // TestHandleSmartRetry_LongDelay_StickySession_ClearsSession
- // 长延迟路径(情况1)应立即清除 sticky 绑定,避免下一次请求继续命中已限流账号。
- func TestHandleSmartRetry_LongDelay_StickySession_ClearsSession(t *testing.T) {
- repo := &stubAntigravityAccountRepo{}
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 14,
- Name: "acc-14",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- // 15s >= 7s 阈值 → 走长延迟路径
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-sonnet-4-5"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "15s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- accountRepo: repo,
- isStickySession: true,
- groupID: 42,
- sessionHash: "sticky-hash-long-delay",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{cache: cache}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.Equal(t, smartRetryActionBreakWithResp, result.action)
- require.NotNil(t, result.switchError)
- require.True(t, result.switchError.IsStickySession)
- require.Len(t, cache.deleteCalls, 1, "long delay path should clear sticky session in handleSmartRetry")
- require.Equal(t, int64(42), cache.deleteCalls[0].groupID)
- require.Equal(t, "sticky-hash-long-delay", cache.deleteCalls[0].sessionHash)
- }
- // TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession
- // 网络错误耗尽重试 + 粘性会话 → 也应清除粘性绑定
- func TestHandleSmartRetry_ShortDelay_NetworkError_StickySession_ClearsSession(t *testing.T) {
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{nil}, // 网络错误
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 15,
- Name: "acc-15",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: true,
- groupID: 99,
- sessionHash: "sticky-net-error",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{cache: cache}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.NotNil(t, result.switchError)
- require.True(t, result.switchError.IsStickySession)
- // 核心断言:网络错误耗尽重试后也应清除粘性绑定
- require.Len(t, cache.deleteCalls, 1, "should call DeleteSessionAccountID after network error exhausts retry")
- require.Equal(t, int64(99), cache.deleteCalls[0].groupID)
- require.Equal(t, "sticky-net-error", cache.deleteCalls[0].sessionHash)
- require.Len(t, repo.modelRateLimitCalls, 2)
- require.Equal(t, "gemini-3-flash", repo.modelRateLimitCalls[0].modelKey)
- require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
- }
- // TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession
- // 429 + 短延迟 + 粘性会话 + 重试失败 → 清除粘性绑定
- func TestHandleSmartRetry_ShortDelay_503_StickySession_FailedRetry_ClearsSession(t *testing.T) {
- failRespBody := `{
- "error": {
- "code": 429,
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
- ]
- }
- }`
- failResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(failRespBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{failResp},
- errors: []error{nil},
- }
- repo := &stubAntigravityAccountRepo{}
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 16,
- Name: "acc-16",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- }
- respBody := []byte(`{
- "error": {
- "code": 429,
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.5s"}
- ]
- }
- }`)
- resp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(respBody)),
- }
- params := antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: true,
- groupID: 77,
- sessionHash: "sticky-503-short",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- }
- availableURLs := []string{"https://ag-1.test"}
- svc := &AntigravityGatewayService{cache: cache}
- result := svc.handleSmartRetry(params, resp, respBody, "https://ag-1.test", 0, availableURLs)
- require.NotNil(t, result)
- require.NotNil(t, result.switchError)
- require.True(t, result.switchError.IsStickySession)
- // 验证粘性绑定被清除
- require.Len(t, cache.deleteCalls, 1)
- require.Equal(t, int64(77), cache.deleteCalls[0].groupID)
- require.Equal(t, "sticky-503-short", cache.deleteCalls[0].sessionHash)
- // 验证模型限流已设置:Gemini 同时写入精确模型和家族级 scope
- require.Len(t, repo.modelRateLimitCalls, 2)
- require.Equal(t, "gemini-3-pro", repo.modelRateLimitCalls[0].modelKey)
- require.Equal(t, antigravityGeminiModelRateLimitKey, repo.modelRateLimitCalls[1].modelKey)
- }
- // TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates
- // 集成测试:antigravityRetryLoop → handleSmartRetry → switchError 传播
- // 验证 IsStickySession 正确传递到上层,且粘性绑定被清除
- func TestAntigravityRetryLoop_SmartRetryFailed_StickySession_SwitchErrorPropagates(t *testing.T) {
- // 初始 429 响应
- initialRespBody := []byte(`{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`)
- initialResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(bytes.NewReader(initialRespBody)),
- }
- // 智能重试也返回 429
- retryRespBody := `{
- "error": {
- "status": "RESOURCE_EXHAUSTED",
- "details": [
- {"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "claude-opus-4-6"}, "reason": "RATE_LIMIT_EXCEEDED"},
- {"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
- ]
- }
- }`
- retryResp := &http.Response{
- StatusCode: http.StatusTooManyRequests,
- Header: http.Header{},
- Body: io.NopCloser(strings.NewReader(retryRespBody)),
- }
- upstream := &mockSmartRetryUpstream{
- responses: []*http.Response{initialResp, retryResp},
- errors: []error{nil, nil},
- }
- repo := &stubAntigravityAccountRepo{}
- cache := &stubSmartRetryCache{}
- account := &Account{
- ID: 17,
- Name: "acc-17",
- Type: AccountTypeOAuth,
- Platform: PlatformAntigravity,
- Schedulable: true,
- Status: StatusActive,
- Concurrency: 1,
- }
- svc := &AntigravityGatewayService{cache: cache}
- result, err := svc.antigravityRetryLoop(antigravityRetryLoopParams{
- ctx: context.Background(),
- prefix: "[test]",
- account: account,
- accessToken: "token",
- action: "generateContent",
- body: []byte(`{"input":"test"}`),
- httpUpstream: upstream,
- accountRepo: repo,
- isStickySession: true,
- groupID: 55,
- sessionHash: "sticky-loop-test",
- handleError: func(ctx context.Context, prefix string, account *Account, statusCode int, headers http.Header, body []byte, requestedModel string, groupID int64, sessionHash string, isStickySession bool) *handleModelRateLimitResult {
- return nil
- },
- })
- require.Nil(t, result, "should not return result when switchError")
- require.NotNil(t, err, "should return error")
- var switchErr *AntigravityAccountSwitchError
- require.ErrorAs(t, err, &switchErr, "error should be AntigravityAccountSwitchError")
- require.Equal(t, account.ID, switchErr.OriginalAccountID)
- require.Equal(t, "claude-opus-4-6", switchErr.RateLimitedModel)
- require.True(t, switchErr.IsStickySession, "IsStickySession must propagate through retryLoop")
- // 验证粘性绑定被清除
- require.Len(t, cache.deleteCalls, 1, "should clear sticky session in handleSmartRetry")
- require.Equal(t, int64(55), cache.deleteCalls[0].groupID)
- require.Equal(t, "sticky-loop-test", cache.deleteCalls[0].sessionHash)
- }
|