Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Sign in / Register
Toggle navigation
Menu
Open sidebar
陈曦
sub2api
Commits
14e1aac9
Unverified
Commit
14e1aac9
authored
Feb 10, 2026
by
Wesley Liddick
Committed by
GitHub
Feb 10, 2026
Browse files
Merge pull request #533 from GuangYiDing/feat/antigravity-single-account-503-retry
feat: Antigravity 单账号分组 503 退避重试机制
parents
aa4b1021
e4bc3515
Changes
7
Hide whitespace changes
Inline
Side-by-side
backend/internal/handler/gateway_handler.go
View file @
14e1aac9
...
@@ -238,6 +238,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
...
@@ -238,6 +238,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
var
lastFailoverErr
*
service
.
UpstreamFailoverError
var
lastFailoverErr
*
service
.
UpstreamFailoverError
var
forceCacheBilling
bool
// 粘性会话切换时的缓存计费标记
var
forceCacheBilling
bool
// 粘性会话切换时的缓存计费标记
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
if
h
.
gatewayService
.
IsSingleAntigravityAccountGroup
(
c
.
Request
.
Context
(),
apiKey
.
GroupID
)
{
ctx
:=
context
.
WithValue
(
c
.
Request
.
Context
(),
ctxkey
.
SingleAccountRetry
,
true
)
c
.
Request
=
c
.
Request
.
WithContext
(
ctx
)
}
for
{
for
{
selection
,
err
:=
h
.
gatewayService
.
SelectAccountWithLoadAwareness
(
c
.
Request
.
Context
(),
apiKey
.
GroupID
,
sessionKey
,
reqModel
,
failedAccountIDs
,
""
)
// Gemini 不使用会话限制
selection
,
err
:=
h
.
gatewayService
.
SelectAccountWithLoadAwareness
(
c
.
Request
.
Context
(),
apiKey
.
GroupID
,
sessionKey
,
reqModel
,
failedAccountIDs
,
""
)
// Gemini 不使用会话限制
if
err
!=
nil
{
if
err
!=
nil
{
...
@@ -245,6 +252,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
...
@@ -245,6 +252,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h
.
handleStreamingAwareError
(
c
,
http
.
StatusServiceUnavailable
,
"api_error"
,
"No available accounts: "
+
err
.
Error
(),
streamStarted
)
h
.
handleStreamingAwareError
(
c
,
http
.
StatusServiceUnavailable
,
"api_error"
,
"No available accounts: "
+
err
.
Error
(),
streamStarted
)
return
return
}
}
// Antigravity 单账号退避重试:分组内没有其他可用账号时,
// 对 503 错误不直接返回,而是清除排除列表、等待退避后重试同一个账号。
// 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
if
lastFailoverErr
!=
nil
&&
lastFailoverErr
.
StatusCode
==
http
.
StatusServiceUnavailable
&&
switchCount
<=
maxAccountSwitches
{
if
sleepAntigravitySingleAccountBackoff
(
c
.
Request
.
Context
(),
switchCount
)
{
log
.
Printf
(
"Antigravity single-account 503 retry: clearing failed accounts, retry %d/%d"
,
switchCount
,
maxAccountSwitches
)
failedAccountIDs
=
make
(
map
[
int64
]
struct
{})
// 设置 context 标记,让 Service 层预检查等待限流过期而非直接切换
ctx
:=
context
.
WithValue
(
c
.
Request
.
Context
(),
ctxkey
.
SingleAccountRetry
,
true
)
c
.
Request
=
c
.
Request
.
WithContext
(
ctx
)
continue
}
}
if
lastFailoverErr
!=
nil
{
if
lastFailoverErr
!=
nil
{
h
.
handleFailoverExhausted
(
c
,
lastFailoverErr
,
service
.
PlatformGemini
,
streamStarted
)
h
.
handleFailoverExhausted
(
c
,
lastFailoverErr
,
service
.
PlatformGemini
,
streamStarted
)
}
else
{
}
else
{
...
@@ -396,6 +416,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
...
@@ -396,6 +416,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
}
}
fallbackUsed
:=
false
fallbackUsed
:=
false
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
if
h
.
gatewayService
.
IsSingleAntigravityAccountGroup
(
c
.
Request
.
Context
(),
currentAPIKey
.
GroupID
)
{
ctx
:=
context
.
WithValue
(
c
.
Request
.
Context
(),
ctxkey
.
SingleAccountRetry
,
true
)
c
.
Request
=
c
.
Request
.
WithContext
(
ctx
)
}
for
{
for
{
maxAccountSwitches
:=
h
.
maxAccountSwitches
maxAccountSwitches
:=
h
.
maxAccountSwitches
switchCount
:=
0
switchCount
:=
0
...
@@ -412,6 +439,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
...
@@ -412,6 +439,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) {
h
.
handleStreamingAwareError
(
c
,
http
.
StatusServiceUnavailable
,
"api_error"
,
"No available accounts: "
+
err
.
Error
(),
streamStarted
)
h
.
handleStreamingAwareError
(
c
,
http
.
StatusServiceUnavailable
,
"api_error"
,
"No available accounts: "
+
err
.
Error
(),
streamStarted
)
return
return
}
}
// Antigravity 单账号退避重试:分组内没有其他可用账号时,
// 对 503 错误不直接返回,而是清除排除列表、等待退避后重试同一个账号。
// 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
if
lastFailoverErr
!=
nil
&&
lastFailoverErr
.
StatusCode
==
http
.
StatusServiceUnavailable
&&
switchCount
<=
maxAccountSwitches
{
if
sleepAntigravitySingleAccountBackoff
(
c
.
Request
.
Context
(),
switchCount
)
{
log
.
Printf
(
"Antigravity single-account 503 retry: clearing failed accounts, retry %d/%d"
,
switchCount
,
maxAccountSwitches
)
failedAccountIDs
=
make
(
map
[
int64
]
struct
{})
// 设置 context 标记,让 Service 层预检查等待限流过期而非直接切换
ctx
:=
context
.
WithValue
(
c
.
Request
.
Context
(),
ctxkey
.
SingleAccountRetry
,
true
)
c
.
Request
=
c
.
Request
.
WithContext
(
ctx
)
continue
}
}
if
lastFailoverErr
!=
nil
{
if
lastFailoverErr
!=
nil
{
h
.
handleFailoverExhausted
(
c
,
lastFailoverErr
,
platform
,
streamStarted
)
h
.
handleFailoverExhausted
(
c
,
lastFailoverErr
,
platform
,
streamStarted
)
}
else
{
}
else
{
...
@@ -838,6 +878,27 @@ func sleepFailoverDelay(ctx context.Context, switchCount int) bool {
...
@@ -838,6 +878,27 @@ func sleepFailoverDelay(ctx context.Context, switchCount int) bool {
}
}
}
}
// sleepAntigravitySingleAccountBackoff Antigravity 平台单账号分组的 503 退避重试延时。
// 当分组内只有一个可用账号且上游返回 503(MODEL_CAPACITY_EXHAUSTED)时使用,
// 采用短固定延时策略。Service 层在 SingleAccountRetry 模式下已经做了充分的原地重试
// (最多 3 次、总等待 30s),所以 Handler 层的退避只需短暂等待即可。
// 返回 false 表示 context 已取消。
func
sleepAntigravitySingleAccountBackoff
(
ctx
context
.
Context
,
retryCount
int
)
bool
{
// 固定短延时:2s
// Service 层已经在原地等待了足够长的时间(retryDelay × 重试次数),
// Handler 层只需短暂间隔后重新进入 Service 层即可。
const
delay
=
2
*
time
.
Second
log
.
Printf
(
"Antigravity single-account 503 backoff: waiting %v before retry (attempt %d)"
,
delay
,
retryCount
)
select
{
case
<-
ctx
.
Done
()
:
return
false
case
<-
time
.
After
(
delay
)
:
return
true
}
}
func
(
h
*
GatewayHandler
)
handleFailoverExhausted
(
c
*
gin
.
Context
,
failoverErr
*
service
.
UpstreamFailoverError
,
platform
string
,
streamStarted
bool
)
{
func
(
h
*
GatewayHandler
)
handleFailoverExhausted
(
c
*
gin
.
Context
,
failoverErr
*
service
.
UpstreamFailoverError
,
platform
string
,
streamStarted
bool
)
{
statusCode
:=
failoverErr
.
StatusCode
statusCode
:=
failoverErr
.
StatusCode
responseBody
:=
failoverErr
.
ResponseBody
responseBody
:=
failoverErr
.
ResponseBody
...
...
backend/internal/handler/gateway_handler_single_account_retry_test.go
0 → 100644
View file @
14e1aac9
package
handler
import
(
"context"
"testing"
"time"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// sleepAntigravitySingleAccountBackoff 测试
// ---------------------------------------------------------------------------
func
TestSleepAntigravitySingleAccountBackoff_ReturnsTrue
(
t
*
testing
.
T
)
{
ctx
:=
context
.
Background
()
start
:=
time
.
Now
()
ok
:=
sleepAntigravitySingleAccountBackoff
(
ctx
,
1
)
elapsed
:=
time
.
Since
(
start
)
require
.
True
(
t
,
ok
,
"should return true when context is not canceled"
)
// 固定延迟 2s
require
.
GreaterOrEqual
(
t
,
elapsed
,
1500
*
time
.
Millisecond
,
"should wait approximately 2s"
)
require
.
Less
(
t
,
elapsed
,
5
*
time
.
Second
,
"should not wait too long"
)
}
func
TestSleepAntigravitySingleAccountBackoff_ContextCanceled
(
t
*
testing
.
T
)
{
ctx
,
cancel
:=
context
.
WithCancel
(
context
.
Background
())
cancel
()
// 立即取消
start
:=
time
.
Now
()
ok
:=
sleepAntigravitySingleAccountBackoff
(
ctx
,
1
)
elapsed
:=
time
.
Since
(
start
)
require
.
False
(
t
,
ok
,
"should return false when context is canceled"
)
require
.
Less
(
t
,
elapsed
,
500
*
time
.
Millisecond
,
"should return immediately on cancel"
)
}
func
TestSleepAntigravitySingleAccountBackoff_FixedDelay
(
t
*
testing
.
T
)
{
// 验证不同 retryCount 都使用固定 2s 延迟
ctx
:=
context
.
Background
()
start
:=
time
.
Now
()
ok
:=
sleepAntigravitySingleAccountBackoff
(
ctx
,
5
)
elapsed
:=
time
.
Since
(
start
)
require
.
True
(
t
,
ok
)
// 即使 retryCount=5,延迟仍然是固定的 2s
require
.
GreaterOrEqual
(
t
,
elapsed
,
1500
*
time
.
Millisecond
)
require
.
Less
(
t
,
elapsed
,
5
*
time
.
Second
)
}
backend/internal/handler/gemini_v1beta_handler.go
View file @
14e1aac9
...
@@ -327,6 +327,13 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
...
@@ -327,6 +327,13 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
var
lastFailoverErr
*
service
.
UpstreamFailoverError
var
lastFailoverErr
*
service
.
UpstreamFailoverError
var
forceCacheBilling
bool
// 粘性会话切换时的缓存计费标记
var
forceCacheBilling
bool
// 粘性会话切换时的缓存计费标记
// 单账号分组提前设置 SingleAccountRetry 标记,让 Service 层首次 503 就不设模型限流标记。
// 避免单账号分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时设 29s 限流,导致后续请求连续快速失败。
if
h
.
gatewayService
.
IsSingleAntigravityAccountGroup
(
c
.
Request
.
Context
(),
apiKey
.
GroupID
)
{
ctx
:=
context
.
WithValue
(
c
.
Request
.
Context
(),
ctxkey
.
SingleAccountRetry
,
true
)
c
.
Request
=
c
.
Request
.
WithContext
(
ctx
)
}
for
{
for
{
selection
,
err
:=
h
.
gatewayService
.
SelectAccountWithLoadAwareness
(
c
.
Request
.
Context
(),
apiKey
.
GroupID
,
sessionKey
,
modelName
,
failedAccountIDs
,
""
)
// Gemini 不使用会话限制
selection
,
err
:=
h
.
gatewayService
.
SelectAccountWithLoadAwareness
(
c
.
Request
.
Context
(),
apiKey
.
GroupID
,
sessionKey
,
modelName
,
failedAccountIDs
,
""
)
// Gemini 不使用会话限制
if
err
!=
nil
{
if
err
!=
nil
{
...
@@ -334,6 +341,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
...
@@ -334,6 +341,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) {
googleError
(
c
,
http
.
StatusServiceUnavailable
,
"No available Gemini accounts: "
+
err
.
Error
())
googleError
(
c
,
http
.
StatusServiceUnavailable
,
"No available Gemini accounts: "
+
err
.
Error
())
return
return
}
}
// Antigravity 单账号退避重试:分组内没有其他可用账号时,
// 对 503 错误不直接返回,而是清除排除列表、等待退避后重试同一个账号。
// 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
if
lastFailoverErr
!=
nil
&&
lastFailoverErr
.
StatusCode
==
http
.
StatusServiceUnavailable
&&
switchCount
<=
maxAccountSwitches
{
if
sleepAntigravitySingleAccountBackoff
(
c
.
Request
.
Context
(),
switchCount
)
{
log
.
Printf
(
"Antigravity single-account 503 retry: clearing failed accounts, retry %d/%d"
,
switchCount
,
maxAccountSwitches
)
failedAccountIDs
=
make
(
map
[
int64
]
struct
{})
// 设置 context 标记,让 Service 层预检查等待限流过期而非直接切换
ctx
:=
context
.
WithValue
(
c
.
Request
.
Context
(),
ctxkey
.
SingleAccountRetry
,
true
)
c
.
Request
=
c
.
Request
.
WithContext
(
ctx
)
continue
}
}
h
.
handleGeminiFailoverExhausted
(
c
,
lastFailoverErr
)
h
.
handleGeminiFailoverExhausted
(
c
,
lastFailoverErr
)
return
return
}
}
...
...
backend/internal/pkg/ctxkey/ctxkey.go
View file @
14e1aac9
...
@@ -28,4 +28,8 @@ const (
...
@@ -28,4 +28,8 @@ const (
// IsMaxTokensOneHaikuRequest 标识当前请求是否为 max_tokens=1 + haiku 模型的探测请求
// IsMaxTokensOneHaikuRequest 标识当前请求是否为 max_tokens=1 + haiku 模型的探测请求
// 用于 ClaudeCodeOnly 验证绕过(绕过 system prompt 检查,但仍需验证 User-Agent)
// 用于 ClaudeCodeOnly 验证绕过(绕过 system prompt 检查,但仍需验证 User-Agent)
IsMaxTokensOneHaikuRequest
Key
=
"ctx_is_max_tokens_one_haiku"
IsMaxTokensOneHaikuRequest
Key
=
"ctx_is_max_tokens_one_haiku"
// SingleAccountRetry 标识当前请求处于单账号 503 退避重试模式。
// 在此模式下,Service 层的模型限流预检查将等待限流过期而非直接切换账号。
SingleAccountRetry
Key
=
"ctx_single_account_retry"
)
)
backend/internal/service/antigravity_gateway_service.go
View file @
14e1aac9
...
@@ -20,6 +20,7 @@ import (
...
@@ -20,6 +20,7 @@ import (
"time"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/gin-gonic/gin"
"github.com/gin-gonic/gin"
"github.com/google/uuid"
"github.com/google/uuid"
)
)
...
@@ -46,6 +47,23 @@ const (
...
@@ -46,6 +47,23 @@ const (
googleRPCTypeErrorInfo
=
"type.googleapis.com/google.rpc.ErrorInfo"
googleRPCTypeErrorInfo
=
"type.googleapis.com/google.rpc.ErrorInfo"
googleRPCReasonModelCapacityExhausted
=
"MODEL_CAPACITY_EXHAUSTED"
googleRPCReasonModelCapacityExhausted
=
"MODEL_CAPACITY_EXHAUSTED"
googleRPCReasonRateLimitExceeded
=
"RATE_LIMIT_EXCEEDED"
googleRPCReasonRateLimitExceeded
=
"RATE_LIMIT_EXCEEDED"
// 单账号 503 退避重试:预检查中等待模型限流过期的最大时间
// 超过此值的限流将直接切换账号(避免请求等待过久)
antigravitySingleAccountMaxWait
=
30
*
time
.
Second
// 单账号 503 退避重试:Service 层原地重试的最大次数
// 在 handleSmartRetry 中,对于 shouldRateLimitModel(长延迟 ≥ 7s)的情况,
// 多账号模式下会设限流+切换账号;但单账号模式下改为原地等待+重试。
antigravitySingleAccountSmartRetryMaxAttempts
=
3
// 单账号 503 退避重试:原地重试时单次最大等待时间
// 防止上游返回过长的 retryDelay 导致请求卡住太久
antigravitySingleAccountSmartRetryMaxWait
=
15
*
time
.
Second
// 单账号 503 退避重试:原地重试的总累计等待时间上限
// 超过此上限将不再重试,直接返回 503
antigravitySingleAccountSmartRetryTotalMaxWait
=
30
*
time
.
Second
)
)
// antigravityPassthroughErrorMessages 透传给客户端的错误消息白名单(小写)
// antigravityPassthroughErrorMessages 透传给客户端的错误消息白名单(小写)
...
@@ -148,6 +166,13 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
...
@@ -148,6 +166,13 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
// 情况1: retryDelay >= 阈值,限流模型并切换账号
// 情况1: retryDelay >= 阈值,限流模型并切换账号
if
shouldRateLimitModel
{
if
shouldRateLimitModel
{
// 单账号 503 退避重试模式:不设限流、不切换账号,改为原地等待+重试
// 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
// 多账号场景下切换账号是最优选择,但单账号场景下设限流毫无意义(只会导致双重等待)。
if
resp
.
StatusCode
==
http
.
StatusServiceUnavailable
&&
isSingleAccountRetry
(
p
.
ctx
)
{
return
s
.
handleSingleAccountRetryInPlace
(
p
,
resp
,
respBody
,
baseURL
,
waitDuration
,
modelName
)
}
rateLimitDuration
:=
waitDuration
rateLimitDuration
:=
waitDuration
if
rateLimitDuration
<=
0
{
if
rateLimitDuration
<=
0
{
rateLimitDuration
=
antigravityDefaultRateLimitDuration
rateLimitDuration
=
antigravityDefaultRateLimitDuration
...
@@ -236,7 +261,7 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
...
@@ -236,7 +261,7 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
}
}
}
}
// 所有重试都失败
,限流当前模型并切换账号
// 所有重试都失败
rateLimitDuration
:=
waitDuration
rateLimitDuration
:=
waitDuration
if
rateLimitDuration
<=
0
{
if
rateLimitDuration
<=
0
{
rateLimitDuration
=
antigravityDefaultRateLimitDuration
rateLimitDuration
=
antigravityDefaultRateLimitDuration
...
@@ -245,6 +270,22 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
...
@@ -245,6 +270,22 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
if
retryBody
==
nil
{
if
retryBody
==
nil
{
retryBody
=
respBody
retryBody
=
respBody
}
}
// 单账号 503 退避重试模式:智能重试耗尽后不设限流、不切换账号,
// 直接返回 503 让 Handler 层的单账号退避循环做最终处理。
if
resp
.
StatusCode
==
http
.
StatusServiceUnavailable
&&
isSingleAccountRetry
(
p
.
ctx
)
{
log
.
Printf
(
"%s status=%d smart_retry_exhausted_single_account attempts=%d model=%s account=%d body=%s (return 503 directly)"
,
p
.
prefix
,
resp
.
StatusCode
,
antigravitySmartRetryMaxAttempts
,
modelName
,
p
.
account
.
ID
,
truncateForLog
(
retryBody
,
200
))
return
&
smartRetryResult
{
action
:
smartRetryActionBreakWithResp
,
resp
:
&
http
.
Response
{
StatusCode
:
resp
.
StatusCode
,
Header
:
resp
.
Header
.
Clone
(),
Body
:
io
.
NopCloser
(
bytes
.
NewReader
(
retryBody
)),
},
}
}
log
.
Printf
(
"%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)"
,
log
.
Printf
(
"%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)"
,
p
.
prefix
,
resp
.
StatusCode
,
antigravitySmartRetryMaxAttempts
,
modelName
,
p
.
account
.
ID
,
rateLimitDuration
,
truncateForLog
(
retryBody
,
200
))
p
.
prefix
,
resp
.
StatusCode
,
antigravitySmartRetryMaxAttempts
,
modelName
,
p
.
account
.
ID
,
rateLimitDuration
,
truncateForLog
(
retryBody
,
200
))
...
@@ -279,17 +320,152 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
...
@@ -279,17 +320,152 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam
return
&
smartRetryResult
{
action
:
smartRetryActionContinue
}
return
&
smartRetryResult
{
action
:
smartRetryActionContinue
}
}
}
// handleSingleAccountRetryInPlace 单账号 503 退避重试的原地重试逻辑。
//
// 在多账号场景下,收到 503 + 长 retryDelay(≥ 7s)时会设置模型限流 + 切换账号;
// 但在单账号场景下,设限流毫无意义(因为切换回来的还是同一个账号,还要等限流过期)。
// 此方法改为在 Service 层原地等待 + 重试,避免双重等待问题:
//
// 旧流程:Service 设限流 → Handler 退避等待 → Service 等限流过期 → 再请求(总耗时 = 退避 + 限流)
// 新流程:Service 直接等 retryDelay → 重试 → 成功/再等 → 重试...(总耗时 ≈ 实际 retryDelay × 重试次数)
//
// 约束:
// - 单次等待不超过 antigravitySingleAccountSmartRetryMaxWait
// - 总累计等待不超过 antigravitySingleAccountSmartRetryTotalMaxWait
// - 最多重试 antigravitySingleAccountSmartRetryMaxAttempts 次
func
(
s
*
AntigravityGatewayService
)
handleSingleAccountRetryInPlace
(
p
antigravityRetryLoopParams
,
resp
*
http
.
Response
,
respBody
[]
byte
,
baseURL
string
,
waitDuration
time
.
Duration
,
modelName
string
,
)
*
smartRetryResult
{
// 限制单次等待时间
if
waitDuration
>
antigravitySingleAccountSmartRetryMaxWait
{
waitDuration
=
antigravitySingleAccountSmartRetryMaxWait
}
if
waitDuration
<
antigravitySmartRetryMinWait
{
waitDuration
=
antigravitySmartRetryMinWait
}
log
.
Printf
(
"%s status=%d single_account_503_retry_in_place model=%s account=%d upstream_retry_delay=%v (retrying in-place instead of rate-limiting)"
,
p
.
prefix
,
resp
.
StatusCode
,
modelName
,
p
.
account
.
ID
,
waitDuration
)
var
lastRetryResp
*
http
.
Response
var
lastRetryBody
[]
byte
totalWaited
:=
time
.
Duration
(
0
)
for
attempt
:=
1
;
attempt
<=
antigravitySingleAccountSmartRetryMaxAttempts
;
attempt
++
{
// 检查累计等待是否超限
if
totalWaited
+
waitDuration
>
antigravitySingleAccountSmartRetryTotalMaxWait
{
remaining
:=
antigravitySingleAccountSmartRetryTotalMaxWait
-
totalWaited
if
remaining
<=
0
{
log
.
Printf
(
"%s single_account_503_retry: total_wait_exceeded total=%v max=%v, giving up"
,
p
.
prefix
,
totalWaited
,
antigravitySingleAccountSmartRetryTotalMaxWait
)
break
}
waitDuration
=
remaining
}
log
.
Printf
(
"%s status=%d single_account_503_retry attempt=%d/%d delay=%v total_waited=%v model=%s account=%d"
,
p
.
prefix
,
resp
.
StatusCode
,
attempt
,
antigravitySingleAccountSmartRetryMaxAttempts
,
waitDuration
,
totalWaited
,
modelName
,
p
.
account
.
ID
)
select
{
case
<-
p
.
ctx
.
Done
()
:
log
.
Printf
(
"%s status=context_canceled_during_single_account_retry"
,
p
.
prefix
)
return
&
smartRetryResult
{
action
:
smartRetryActionBreakWithResp
,
err
:
p
.
ctx
.
Err
()}
case
<-
time
.
After
(
waitDuration
)
:
}
totalWaited
+=
waitDuration
// 创建新请求
retryReq
,
err
:=
antigravity
.
NewAPIRequestWithURL
(
p
.
ctx
,
baseURL
,
p
.
action
,
p
.
accessToken
,
p
.
body
)
if
err
!=
nil
{
log
.
Printf
(
"%s single_account_503_retry: request_build_failed error=%v"
,
p
.
prefix
,
err
)
break
}
retryResp
,
retryErr
:=
p
.
httpUpstream
.
Do
(
retryReq
,
p
.
proxyURL
,
p
.
account
.
ID
,
p
.
account
.
Concurrency
)
if
retryErr
==
nil
&&
retryResp
!=
nil
&&
retryResp
.
StatusCode
!=
http
.
StatusTooManyRequests
&&
retryResp
.
StatusCode
!=
http
.
StatusServiceUnavailable
{
log
.
Printf
(
"%s status=%d single_account_503_retry_success attempt=%d/%d total_waited=%v"
,
p
.
prefix
,
retryResp
.
StatusCode
,
attempt
,
antigravitySingleAccountSmartRetryMaxAttempts
,
totalWaited
)
// 关闭之前的响应
if
lastRetryResp
!=
nil
{
_
=
lastRetryResp
.
Body
.
Close
()
}
return
&
smartRetryResult
{
action
:
smartRetryActionBreakWithResp
,
resp
:
retryResp
}
}
// 网络错误时继续重试
if
retryErr
!=
nil
||
retryResp
==
nil
{
log
.
Printf
(
"%s single_account_503_retry: network_error attempt=%d/%d error=%v"
,
p
.
prefix
,
attempt
,
antigravitySingleAccountSmartRetryMaxAttempts
,
retryErr
)
continue
}
// 关闭之前的响应
if
lastRetryResp
!=
nil
{
_
=
lastRetryResp
.
Body
.
Close
()
}
lastRetryResp
=
retryResp
lastRetryBody
,
_
=
io
.
ReadAll
(
io
.
LimitReader
(
retryResp
.
Body
,
2
<<
20
))
_
=
retryResp
.
Body
.
Close
()
// 解析新的重试信息,更新下次等待时间
if
attempt
<
antigravitySingleAccountSmartRetryMaxAttempts
&&
lastRetryBody
!=
nil
{
_
,
_
,
newWaitDuration
,
_
:=
shouldTriggerAntigravitySmartRetry
(
p
.
account
,
lastRetryBody
)
if
newWaitDuration
>
0
{
waitDuration
=
newWaitDuration
if
waitDuration
>
antigravitySingleAccountSmartRetryMaxWait
{
waitDuration
=
antigravitySingleAccountSmartRetryMaxWait
}
if
waitDuration
<
antigravitySmartRetryMinWait
{
waitDuration
=
antigravitySmartRetryMinWait
}
}
}
}
// 所有重试都失败,不设限流,直接返回 503
// Handler 层的单账号退避循环会做最终处理
retryBody
:=
lastRetryBody
if
retryBody
==
nil
{
retryBody
=
respBody
}
log
.
Printf
(
"%s status=%d single_account_503_retry_exhausted attempts=%d total_waited=%v model=%s account=%d body=%s (return 503 directly)"
,
p
.
prefix
,
resp
.
StatusCode
,
antigravitySingleAccountSmartRetryMaxAttempts
,
totalWaited
,
modelName
,
p
.
account
.
ID
,
truncateForLog
(
retryBody
,
200
))
return
&
smartRetryResult
{
action
:
smartRetryActionBreakWithResp
,
resp
:
&
http
.
Response
{
StatusCode
:
resp
.
StatusCode
,
Header
:
resp
.
Header
.
Clone
(),
Body
:
io
.
NopCloser
(
bytes
.
NewReader
(
retryBody
)),
},
}
}
// antigravityRetryLoop 执行带 URL fallback 的重试循环
// antigravityRetryLoop 执行带 URL fallback 的重试循环
func
(
s
*
AntigravityGatewayService
)
antigravityRetryLoop
(
p
antigravityRetryLoopParams
)
(
*
antigravityRetryLoopResult
,
error
)
{
func
(
s
*
AntigravityGatewayService
)
antigravityRetryLoop
(
p
antigravityRetryLoopParams
)
(
*
antigravityRetryLoopResult
,
error
)
{
// 预检查:如果账号已限流,直接返回切换信号
// 预检查:如果账号已限流,直接返回切换信号
if
p
.
requestedModel
!=
""
{
if
p
.
requestedModel
!=
""
{
if
remaining
:=
p
.
account
.
GetRateLimitRemainingTimeWithContext
(
p
.
ctx
,
p
.
requestedModel
);
remaining
>
0
{
if
remaining
:=
p
.
account
.
GetRateLimitRemainingTimeWithContext
(
p
.
ctx
,
p
.
requestedModel
);
remaining
>
0
{
log
.
Printf
(
"%s pre_check: rate_limit_switch remaining=%v model=%s account=%d"
,
// 单账号 503 退避重试模式:跳过限流预检查,直接发请求。
p
.
prefix
,
remaining
.
Truncate
(
time
.
Millisecond
),
p
.
requestedModel
,
p
.
account
.
ID
)
// 首次请求设的限流是为了多账号调度器跳过该账号,在单账号模式下无意义。
return
nil
,
&
AntigravityAccountSwitchError
{
// 如果上游确实还不可用,handleSmartRetry → handleSingleAccountRetryInPlace
OriginalAccountID
:
p
.
account
.
ID
,
// 会在 Service 层原地等待+重试,不需要在预检查这里等。
RateLimitedModel
:
p
.
requestedModel
,
if
isSingleAccountRetry
(
p
.
ctx
)
{
IsStickySession
:
p
.
isStickySession
,
log
.
Printf
(
"%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)"
,
p
.
prefix
,
remaining
.
Truncate
(
time
.
Millisecond
),
p
.
requestedModel
,
p
.
account
.
ID
)
}
else
{
log
.
Printf
(
"%s pre_check: rate_limit_switch remaining=%v model=%s account=%d"
,
p
.
prefix
,
remaining
.
Truncate
(
time
.
Millisecond
),
p
.
requestedModel
,
p
.
account
.
ID
)
return
nil
,
&
AntigravityAccountSwitchError
{
OriginalAccountID
:
p
.
account
.
ID
,
RateLimitedModel
:
p
.
requestedModel
,
IsStickySession
:
p
.
isStickySession
,
}
}
}
}
}
}
}
...
@@ -1944,6 +2120,12 @@ func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool {
...
@@ -1944,6 +2120,12 @@ func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool {
}
}
}
}
// isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记
func
isSingleAccountRetry
(
ctx
context
.
Context
)
bool
{
v
,
_
:=
ctx
.
Value
(
ctxkey
.
SingleAccountRetry
)
.
(
bool
)
return
v
}
// setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流
// setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流
// 直接使用上游返回的模型 ID(如 claude-sonnet-4-5)作为限流 key
// 直接使用上游返回的模型 ID(如 claude-sonnet-4-5)作为限流 key
// 返回是否已成功设置(若模型名为空或 repo 为 nil 将返回 false)
// 返回是否已成功设置(若模型名为空或 repo 为 nil 将返回 false)
...
...
backend/internal/service/antigravity_single_account_retry_test.go
0 → 100644
View file @
14e1aac9
//go:build unit
package
service
import
(
"bytes"
"context"
"io"
"net/http"
"strings"
"testing"
"time"
"github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey"
"github.com/stretchr/testify/require"
)
// ---------------------------------------------------------------------------
// 辅助函数:构造带 SingleAccountRetry 标记的 context
// ---------------------------------------------------------------------------
func
ctxWithSingleAccountRetry
()
context
.
Context
{
return
context
.
WithValue
(
context
.
Background
(),
ctxkey
.
SingleAccountRetry
,
true
)
}
// ---------------------------------------------------------------------------
// 1. isSingleAccountRetry 测试
// ---------------------------------------------------------------------------
func
TestIsSingleAccountRetry_True
(
t
*
testing
.
T
)
{
ctx
:=
context
.
WithValue
(
context
.
Background
(),
ctxkey
.
SingleAccountRetry
,
true
)
require
.
True
(
t
,
isSingleAccountRetry
(
ctx
))
}
func
TestIsSingleAccountRetry_False_NoValue
(
t
*
testing
.
T
)
{
require
.
False
(
t
,
isSingleAccountRetry
(
context
.
Background
()))
}
func
TestIsSingleAccountRetry_False_ExplicitFalse
(
t
*
testing
.
T
)
{
ctx
:=
context
.
WithValue
(
context
.
Background
(),
ctxkey
.
SingleAccountRetry
,
false
)
require
.
False
(
t
,
isSingleAccountRetry
(
ctx
))
}
func
TestIsSingleAccountRetry_False_WrongType
(
t
*
testing
.
T
)
{
ctx
:=
context
.
WithValue
(
context
.
Background
(),
ctxkey
.
SingleAccountRetry
,
"true"
)
require
.
False
(
t
,
isSingleAccountRetry
(
ctx
))
}
// ---------------------------------------------------------------------------
// 2. 常量验证
// ---------------------------------------------------------------------------
func
TestSingleAccountRetryConstants
(
t
*
testing
.
T
)
{
require
.
Equal
(
t
,
3
,
antigravitySingleAccountSmartRetryMaxAttempts
,
"单账号原地重试最多 3 次"
)
require
.
Equal
(
t
,
15
*
time
.
Second
,
antigravitySingleAccountSmartRetryMaxWait
,
"单次最大等待 15s"
)
require
.
Equal
(
t
,
30
*
time
.
Second
,
antigravitySingleAccountSmartRetryTotalMaxWait
,
"总累计等待不超过 30s"
)
require
.
Equal
(
t
,
30
*
time
.
Second
,
antigravitySingleAccountMaxWait
,
"预检查最大等待 30s"
)
}
// ---------------------------------------------------------------------------
// 3. handleSmartRetry + 503 + SingleAccountRetry → 走 handleSingleAccountRetryInPlace
// (而非设模型限流 + 切换账号)
// ---------------------------------------------------------------------------
// TestHandleSmartRetry_503_LongDelay_SingleAccountRetry_RetryInPlace
// 核心场景:503 + retryDelay >= 7s + SingleAccountRetry 标记
// → 不设模型限流、不切换账号,改为原地重试
func
TestHandleSmartRetry_503_LongDelay_SingleAccountRetry_RetryInPlace
(
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
},
}
repo
:=
&
stubAntigravityAccountRepo
{}
account
:=
&
Account
{
ID
:
1
,
Name
:
"acc-single"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
// 503 + 39s >= 7s 阈值 + MODEL_CAPACITY_EXHAUSTED
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
)),
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
// 关键:设置单账号标记
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
)
// 关键断言:返回 resp(原地重试成功),而非 switchError(切换账号)
require
.
NotNil
(
t
,
result
.
resp
,
"should return successful response from in-place retry"
)
require
.
Equal
(
t
,
http
.
StatusOK
,
result
.
resp
.
StatusCode
)
require
.
Nil
(
t
,
result
.
switchError
,
"should NOT return switchError in single account mode"
)
require
.
Nil
(
t
,
result
.
err
)
// 验证未设模型限流(单账号模式不应设限流)
require
.
Len
(
t
,
repo
.
modelRateLimitCalls
,
0
,
"should NOT set model rate limit in single account retry mode"
)
// 验证确实调用了 upstream(原地重试)
require
.
GreaterOrEqual
(
t
,
len
(
upstream
.
calls
),
1
,
"should have made at least one retry call"
)
}
// TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches
// 对照组:503 + retryDelay >= 7s + 无 SingleAccountRetry 标记
// → 照常设模型限流 + 切换账号
func
TestHandleSmartRetry_503_LongDelay_NoSingleAccountRetry_StillSwitches
(
t
*
testing
.
T
)
{
repo
:=
&
stubAntigravityAccountRepo
{}
account
:=
&
Account
{
ID
:
2
,
Name
:
"acc-multi"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
}
// 503 + 39s >= 7s 阈值
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
)),
}
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
)
// 对照:多账号模式返回 switchError
require
.
NotNil
(
t
,
result
.
switchError
,
"multi-account mode should return switchError for 503"
)
require
.
Nil
(
t
,
result
.
resp
,
"should not return resp when switchError is set"
)
// 对照:多账号模式应设模型限流
require
.
Len
(
t
,
repo
.
modelRateLimitCalls
,
1
,
"multi-account mode SHOULD set model rate limit"
)
}
// TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches
// 边界情况:429(非 503)+ SingleAccountRetry 标记
// → 单账号原地重试仅针对 503,429 依然走切换账号逻辑
func
TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches
(
t
*
testing
.
T
)
{
repo
:=
&
stubAntigravityAccountRepo
{}
account
:=
&
Account
{
ID
:
3
,
Name
:
"acc-429"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
}
// 429 + 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
,
// 429,不是 503
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
bytes
.
NewReader
(
respBody
)),
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
// 有单账号标记
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
)
// 429 即使有单账号标记,也应走切换账号
require
.
NotNil
(
t
,
result
.
switchError
,
"429 should still return switchError even with SingleAccountRetry"
)
require
.
Len
(
t
,
repo
.
modelRateLimitCalls
,
1
,
"429 should still set model rate limit even with SingleAccountRetry"
)
}
// ---------------------------------------------------------------------------
// 4. handleSmartRetry + 503 + 短延迟 + SingleAccountRetry → 智能重试耗尽后不设限流
// ---------------------------------------------------------------------------
// TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit
// 503 + retryDelay < 7s + SingleAccountRetry → 智能重试耗尽后直接返回 503,不设限流
func
TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit
(
t
*
testing
.
T
)
{
// 智能重试也返回 503
failRespBody
:=
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
failResp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
failRespBody
)),
}
upstream
:=
&
mockSmartRetryUpstream
{
responses
:
[]
*
http
.
Response
{
failResp
},
errors
:
[]
error
{
nil
},
}
repo
:=
&
stubAntigravityAccountRepo
{}
account
:=
&
Account
{
ID
:
4
,
Name
:
"acc-short-503"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
}
// 0.1s < 7s 阈值
respBody
:=
[]
byte
(
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
)
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
bytes
.
NewReader
(
respBody
)),
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
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
)
// 关键断言:单账号 503 模式下,智能重试耗尽后直接返回 503 响应,不切换
require
.
NotNil
(
t
,
result
.
resp
,
"should return 503 response directly for single account mode"
)
require
.
Equal
(
t
,
http
.
StatusServiceUnavailable
,
result
.
resp
.
StatusCode
)
require
.
Nil
(
t
,
result
.
switchError
,
"should NOT switch account in single account mode"
)
// 关键断言:不设模型限流
require
.
Len
(
t
,
repo
.
modelRateLimitCalls
,
0
,
"should NOT set model rate limit for 503 in single account mode"
)
}
// TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit
// 对照组:503 + retryDelay < 7s + 无 SingleAccountRetry → 智能重试耗尽后照常设限流
func
TestHandleSmartRetry_503_ShortDelay_NoSingleAccountRetry_SetsRateLimit
(
t
*
testing
.
T
)
{
failRespBody
:=
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
failResp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
failRespBody
)),
}
upstream
:=
&
mockSmartRetryUpstream
{
responses
:
[]
*
http
.
Response
{
failResp
},
errors
:
[]
error
{
nil
},
}
repo
:=
&
stubAntigravityAccountRepo
{}
account
:=
&
Account
{
ID
:
5
,
Name
:
"acc-multi-503"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
}
respBody
:=
[]
byte
(
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
)
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
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
)
// 对照:多账号模式应返回 switchError
require
.
NotNil
(
t
,
result
.
switchError
,
"multi-account mode should return switchError for 503"
)
// 对照:多账号模式应设模型限流
require
.
Len
(
t
,
repo
.
modelRateLimitCalls
,
1
,
"multi-account mode should set model rate limit"
)
}
// ---------------------------------------------------------------------------
// 5. handleSingleAccountRetryInPlace 直接测试
// ---------------------------------------------------------------------------
// TestHandleSingleAccountRetryInPlace_Success 原地重试成功
func
TestHandleSingleAccountRetryInPlace_Success
(
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
:
10
,
Name
:
"acc-inplace-ok"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
}
svc
:=
&
AntigravityGatewayService
{}
result
:=
svc
.
handleSingleAccountRetryInPlace
(
params
,
resp
,
nil
,
"https://ag-1.test"
,
1
*
time
.
Second
,
"gemini-3-pro"
)
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 switch account on success"
)
require
.
Nil
(
t
,
result
.
err
)
}
// TestHandleSingleAccountRetryInPlace_AllRetriesFail 所有重试都失败,返回 503(不设限流)
func
TestHandleSingleAccountRetryInPlace_AllRetriesFail
(
t
*
testing
.
T
)
{
// 构造 3 个 503 响应(对应 3 次原地重试)
var
responses
[]
*
http
.
Response
var
errors
[]
error
for
i
:=
0
;
i
<
antigravitySingleAccountSmartRetryMaxAttempts
;
i
++
{
responses
=
append
(
responses
,
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
)),
})
errors
=
append
(
errors
,
nil
)
}
upstream
:=
&
mockSmartRetryUpstream
{
responses
:
responses
,
errors
:
errors
,
}
account
:=
&
Account
{
ID
:
11
,
Name
:
"acc-inplace-fail"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
origBody
:=
[]
byte
(
`{"error":{"code":503,"status":"UNAVAILABLE"}}`
)
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{
"X-Test"
:
{
"original"
}},
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
}
svc
:=
&
AntigravityGatewayService
{}
result
:=
svc
.
handleSingleAccountRetryInPlace
(
params
,
resp
,
origBody
,
"https://ag-1.test"
,
1
*
time
.
Second
,
"gemini-3-pro"
)
require
.
NotNil
(
t
,
result
)
require
.
Equal
(
t
,
smartRetryActionBreakWithResp
,
result
.
action
)
// 关键:返回 503 resp,不返回 switchError
require
.
NotNil
(
t
,
result
.
resp
,
"should return 503 response directly"
)
require
.
Equal
(
t
,
http
.
StatusServiceUnavailable
,
result
.
resp
.
StatusCode
)
require
.
Nil
(
t
,
result
.
switchError
,
"should NOT return switchError - let Handler handle it"
)
require
.
Nil
(
t
,
result
.
err
)
// 验证确实重试了指定次数
require
.
Len
(
t
,
upstream
.
calls
,
antigravitySingleAccountSmartRetryMaxAttempts
,
"should have made exactly maxAttempts retry calls"
)
}
// TestHandleSingleAccountRetryInPlace_WaitDurationClamped 等待时间被限制在 [min, max] 范围
func
TestHandleSingleAccountRetryInPlace_WaitDurationClamped
(
t
*
testing
.
T
)
{
// 用短延迟的成功响应,只验证不 panic
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
:
12
,
Name
:
"acc-clamp"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
}
svc
:=
&
AntigravityGatewayService
{}
// 等待时间过大应被 clamp 到 antigravitySingleAccountSmartRetryMaxWait
result
:=
svc
.
handleSingleAccountRetryInPlace
(
params
,
resp
,
nil
,
"https://ag-1.test"
,
999
*
time
.
Second
,
"gemini-3-pro"
)
require
.
NotNil
(
t
,
result
)
require
.
Equal
(
t
,
smartRetryActionBreakWithResp
,
result
.
action
)
require
.
NotNil
(
t
,
result
.
resp
)
require
.
Equal
(
t
,
http
.
StatusOK
,
result
.
resp
.
StatusCode
)
}
// TestHandleSingleAccountRetryInPlace_ContextCanceled context 取消时立即返回
func
TestHandleSingleAccountRetryInPlace_ContextCanceled
(
t
*
testing
.
T
)
{
upstream
:=
&
mockSmartRetryUpstream
{
responses
:
[]
*
http
.
Response
{
nil
},
errors
:
[]
error
{
nil
},
}
account
:=
&
Account
{
ID
:
13
,
Name
:
"acc-cancel"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
}
ctx
,
cancel
:=
context
.
WithCancel
(
context
.
Background
())
ctx
=
context
.
WithValue
(
ctx
,
ctxkey
.
SingleAccountRetry
,
true
)
cancel
()
// 立即取消
params
:=
antigravityRetryLoopParams
{
ctx
:
ctx
,
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
}
svc
:=
&
AntigravityGatewayService
{}
result
:=
svc
.
handleSingleAccountRetryInPlace
(
params
,
resp
,
nil
,
"https://ag-1.test"
,
1
*
time
.
Second
,
"gemini-3-pro"
)
require
.
NotNil
(
t
,
result
)
require
.
Equal
(
t
,
smartRetryActionBreakWithResp
,
result
.
action
)
require
.
Error
(
t
,
result
.
err
,
"should return context error"
)
// 不应调用 upstream(因为在等待阶段就被取消了)
require
.
Len
(
t
,
upstream
.
calls
,
0
,
"should not call upstream when context is canceled"
)
}
// TestHandleSingleAccountRetryInPlace_NetworkError_ContinuesRetry 网络错误时继续重试
func
TestHandleSingleAccountRetryInPlace_NetworkError_ContinuesRetry
(
t
*
testing
.
T
)
{
successResp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusOK
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
`{"result":"ok"}`
)),
}
upstream
:=
&
mockSmartRetryUpstream
{
// 第1次网络错误(nil resp),第2次成功
responses
:
[]
*
http
.
Response
{
nil
,
successResp
},
errors
:
[]
error
{
nil
,
nil
},
}
account
:=
&
Account
{
ID
:
14
,
Name
:
"acc-net-retry"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
}
svc
:=
&
AntigravityGatewayService
{}
result
:=
svc
.
handleSingleAccountRetryInPlace
(
params
,
resp
,
nil
,
"https://ag-1.test"
,
1
*
time
.
Second
,
"gemini-3-pro"
)
require
.
NotNil
(
t
,
result
)
require
.
Equal
(
t
,
smartRetryActionBreakWithResp
,
result
.
action
)
require
.
NotNil
(
t
,
result
.
resp
,
"should return successful response after network error recovery"
)
require
.
Equal
(
t
,
http
.
StatusOK
,
result
.
resp
.
StatusCode
)
require
.
Len
(
t
,
upstream
.
calls
,
2
,
"first call fails (network error), second succeeds"
)
}
// ---------------------------------------------------------------------------
// 6. antigravityRetryLoop 预检查:单账号模式跳过限流
// ---------------------------------------------------------------------------
// TestAntigravityRetryLoop_PreCheck_SingleAccountRetry_SkipsRateLimit
// 预检查中,如果有 SingleAccountRetry 标记,即使账号已限流也跳过直接发请求
func
TestAntigravityRetryLoop_PreCheck_SingleAccountRetry_SkipsRateLimit
(
t
*
testing
.
T
)
{
// 创建一个已设模型限流的账号
upstream
:=
&
recordingOKUpstream
{}
account
:=
&
Account
{
ID
:
20
,
Name
:
"acc-rate-limited"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Schedulable
:
true
,
Status
:
StatusActive
,
Concurrency
:
1
,
Extra
:
map
[
string
]
any
{
modelRateLimitsKey
:
map
[
string
]
any
{
"claude-sonnet-4-5"
:
map
[
string
]
any
{
"rate_limit_reset_at"
:
time
.
Now
()
.
Add
(
30
*
time
.
Second
)
.
Format
(
time
.
RFC3339
),
},
},
},
}
svc
:=
&
AntigravityGatewayService
{}
result
,
err
:=
svc
.
antigravityRetryLoop
(
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
requestedModel
:
"claude-sonnet-4-5"
,
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
.
NoError
(
t
,
err
,
"should not return error"
)
require
.
NotNil
(
t
,
result
,
"should return result"
)
require
.
NotNil
(
t
,
result
.
resp
,
"should have response"
)
require
.
Equal
(
t
,
http
.
StatusOK
,
result
.
resp
.
StatusCode
)
// 关键:尽管限流了,有 SingleAccountRetry 标记时仍然到达了 upstream
require
.
Equal
(
t
,
1
,
upstream
.
calls
,
"should have reached upstream despite rate limit"
)
}
// TestAntigravityRetryLoop_PreCheck_NoSingleAccountRetry_SwitchesOnRateLimit
// 对照组:无 SingleAccountRetry + 已限流 → 预检查返回 switchError
func
TestAntigravityRetryLoop_PreCheck_NoSingleAccountRetry_SwitchesOnRateLimit
(
t
*
testing
.
T
)
{
upstream
:=
&
recordingOKUpstream
{}
account
:=
&
Account
{
ID
:
21
,
Name
:
"acc-rate-limited-multi"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Schedulable
:
true
,
Status
:
StatusActive
,
Concurrency
:
1
,
Extra
:
map
[
string
]
any
{
modelRateLimitsKey
:
map
[
string
]
any
{
"claude-sonnet-4-5"
:
map
[
string
]
any
{
"rate_limit_reset_at"
:
time
.
Now
()
.
Add
(
30
*
time
.
Second
)
.
Format
(
time
.
RFC3339
),
},
},
},
}
svc
:=
&
AntigravityGatewayService
{}
result
,
err
:=
svc
.
antigravityRetryLoop
(
antigravityRetryLoopParams
{
ctx
:
context
.
Background
(),
// 无单账号标记
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
requestedModel
:
"claude-sonnet-4-5"
,
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 on rate limit switch"
)
require
.
NotNil
(
t
,
err
,
"should return error"
)
var
switchErr
*
AntigravityAccountSwitchError
require
.
ErrorAs
(
t
,
err
,
&
switchErr
,
"should return AntigravityAccountSwitchError"
)
require
.
Equal
(
t
,
account
.
ID
,
switchErr
.
OriginalAccountID
)
require
.
Equal
(
t
,
"claude-sonnet-4-5"
,
switchErr
.
RateLimitedModel
)
// upstream 不应被调用(预检查就短路了)
require
.
Equal
(
t
,
0
,
upstream
.
calls
,
"upstream should NOT be called when pre-check blocks"
)
}
// ---------------------------------------------------------------------------
// 7. 端到端集成场景测试
// ---------------------------------------------------------------------------
// TestHandleSmartRetry_503_SingleAccount_RetryInPlace_ThenSuccess_E2E
// 端到端场景:503 + 单账号 + 原地重试第2次成功
func
TestHandleSmartRetry_503_SingleAccount_RetryInPlace_ThenSuccess_E2E
(
t
*
testing
.
T
)
{
// 第1次原地重试仍返回 503,第2次成功
fail503Body
:=
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
]
}
}`
resp503
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
fail503Body
)),
}
successResp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusOK
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
`{"result":"ok"}`
)),
}
upstream
:=
&
mockSmartRetryUpstream
{
responses
:
[]
*
http
.
Response
{
resp503
,
successResp
},
errors
:
[]
error
{
nil
,
nil
},
}
account
:=
&
Account
{
ID
:
30
,
Name
:
"acc-e2e"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Concurrency
:
1
,
}
resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
}
params
:=
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
prefix
:
"[test]"
,
account
:
account
,
accessToken
:
"token"
,
action
:
"generateContent"
,
body
:
[]
byte
(
`{"input":"test"}`
),
httpUpstream
:
upstream
,
}
svc
:=
&
AntigravityGatewayService
{}
result
:=
svc
.
handleSingleAccountRetryInPlace
(
params
,
resp
,
nil
,
"https://ag-1.test"
,
1
*
time
.
Second
,
"gemini-3-pro"
)
require
.
NotNil
(
t
,
result
)
require
.
Equal
(
t
,
smartRetryActionBreakWithResp
,
result
.
action
)
require
.
NotNil
(
t
,
result
.
resp
,
"should return successful response after 2nd attempt"
)
require
.
Equal
(
t
,
http
.
StatusOK
,
result
.
resp
.
StatusCode
)
require
.
Nil
(
t
,
result
.
switchError
)
require
.
Len
(
t
,
upstream
.
calls
,
2
,
"first 503, second OK"
)
}
// TestAntigravityRetryLoop_503_SingleAccount_InPlaceRetryUsed_E2E
// 通过 antigravityRetryLoop → handleSmartRetry → handleSingleAccountRetryInPlace 完整链路
func
TestAntigravityRetryLoop_503_SingleAccount_InPlaceRetryUsed_E2E
(
t
*
testing
.
T
)
{
// 初始请求返回 503 + 长延迟
initial503Body
:=
[]
byte
(
`{
"error": {
"code": 503,
"status": "UNAVAILABLE",
"details": [
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-pro"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "10s"}
],
"message": "No capacity available"
}
}`
)
initial503Resp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusServiceUnavailable
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
bytes
.
NewReader
(
initial503Body
)),
}
// 原地重试成功
successResp
:=
&
http
.
Response
{
StatusCode
:
http
.
StatusOK
,
Header
:
http
.
Header
{},
Body
:
io
.
NopCloser
(
strings
.
NewReader
(
`{"result":"ok"}`
)),
}
upstream
:=
&
mockSmartRetryUpstream
{
// 第1次调用(retryLoop 主循环)返回 503
// 第2次调用(handleSingleAccountRetryInPlace 原地重试)返回 200
responses
:
[]
*
http
.
Response
{
initial503Resp
,
successResp
},
errors
:
[]
error
{
nil
,
nil
},
}
repo
:=
&
stubAntigravityAccountRepo
{}
account
:=
&
Account
{
ID
:
31
,
Name
:
"acc-e2e-loop"
,
Type
:
AccountTypeOAuth
,
Platform
:
PlatformAntigravity
,
Schedulable
:
true
,
Status
:
StatusActive
,
Concurrency
:
1
,
}
svc
:=
&
AntigravityGatewayService
{}
result
,
err
:=
svc
.
antigravityRetryLoop
(
antigravityRetryLoopParams
{
ctx
:
ctxWithSingleAccountRetry
(),
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
},
})
require
.
NoError
(
t
,
err
,
"should not return error on successful retry"
)
require
.
NotNil
(
t
,
result
,
"should return result"
)
require
.
NotNil
(
t
,
result
.
resp
,
"should return response"
)
require
.
Equal
(
t
,
http
.
StatusOK
,
result
.
resp
.
StatusCode
)
// 验证未设模型限流
require
.
Len
(
t
,
repo
.
modelRateLimitCalls
,
0
,
"should NOT set model rate limit in single account retry mode"
)
}
backend/internal/service/gateway_service.go
View file @
14e1aac9
...
@@ -1683,6 +1683,17 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i
...
@@ -1683,6 +1683,17 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i
return
accounts
,
useMixed
,
nil
return
accounts
,
useMixed
,
nil
}
}
// IsSingleAntigravityAccountGroup 检查指定分组是否只有一个 antigravity 平台的可调度账号。
// 用于 Handler 层在首次请求时提前设置 SingleAccountRetry context,
// 避免单账号分组收到 503 时错误地设置模型限流标记导致后续请求连续快速失败。
func
(
s
*
GatewayService
)
IsSingleAntigravityAccountGroup
(
ctx
context
.
Context
,
groupID
*
int64
)
bool
{
accounts
,
_
,
err
:=
s
.
listSchedulableAccounts
(
ctx
,
groupID
,
PlatformAntigravity
,
true
)
if
err
!=
nil
{
return
false
}
return
len
(
accounts
)
==
1
}
func
(
s
*
GatewayService
)
isAccountAllowedForPlatform
(
account
*
Account
,
platform
string
,
useMixed
bool
)
bool
{
func
(
s
*
GatewayService
)
isAccountAllowedForPlatform
(
account
*
Account
,
platform
string
,
useMixed
bool
)
bool
{
if
account
==
nil
{
if
account
==
nil
{
return
false
return
false
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment