From f6cfab990145404bbdb0509a7d12530292dc4c69 Mon Sep 17 00:00:00 2001 From: Rose Ding Date: Mon, 9 Feb 2026 14:26:01 +0800 Subject: [PATCH 1/3] =?UTF-8?q?feat:=20=E6=B7=BB=E5=8A=A0=20Antigravity=20?= =?UTF-8?q?=E5=8D=95=E8=B4=A6=E5=8F=B7=20503=20=E9=80=80=E9=81=BF=E9=87=8D?= =?UTF-8?q?=E8=AF=95=E6=9C=BA=E5=88=B6?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 当分组内只有一个可用账号且上游返回 503 (MODEL_CAPACITY_EXHAUSTED) 时， 不再设置模型限流+切换账号（因为切换回来还是同一个账号），而是在 Service 层 原地等待+重试，避免双重等待问题。 主要变更： - Handler 层：检测单账号 503 场景，清除排除列表并设置 SingleAccountRetry 标记 - Service 层：新增 handleSingleAccountRetryInPlace 原地重试逻辑 - Service 层：预检查跳过单账号模式下的限流检查 - 新增 ctxkey.SingleAccountRetry 上下文标记 --- backend/internal/handler/gateway_handler.go | 47 +++++ .../internal/handler/gemini_v1beta_handler.go | 13 ++ backend/internal/pkg/ctxkey/ctxkey.go | 4 + .../service/antigravity_gateway_service.go | 196 +++++++++++++++++- 4 files changed, 253 insertions(+), 7 deletions(-) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 6900fa55..82181948 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -245,6 +245,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted) 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 { h.handleFailoverExhausted(c, lastFailoverErr, service.PlatformGemini, streamStarted) } else { @@ -412,6 +425,19 @@ func (h *GatewayHandler) Messages(c *gin.Context) { h.handleStreamingAwareError(c, http.StatusServiceUnavailable, "api_error", "No available accounts: "+err.Error(), streamStarted) 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 { h.handleFailoverExhausted(c, lastFailoverErr, platform, streamStarted) } else { @@ -838,6 +864,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) { statusCode := failoverErr.StatusCode responseBody := failoverErr.ResponseBody diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index d5149f22..2b67cb1f 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -334,6 +334,19 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { googleError(c, http.StatusServiceUnavailable, "No available Gemini accounts: "+err.Error()) 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) return } diff --git a/backend/internal/pkg/ctxkey/ctxkey.go b/backend/internal/pkg/ctxkey/ctxkey.go index 9bf563e7..0c4d82f7 100644 --- a/backend/internal/pkg/ctxkey/ctxkey.go +++ b/backend/internal/pkg/ctxkey/ctxkey.go @@ -28,4 +28,8 @@ const ( // IsMaxTokensOneHaikuRequest 标识当前请求是否为 max_tokens=1 + haiku 模型的探测请求 // 用于 ClaudeCodeOnly 验证绕过（绕过 system prompt 检查，但仍需验证 User-Agent） IsMaxTokensOneHaikuRequest Key = "ctx_is_max_tokens_one_haiku" + + // SingleAccountRetry 标识当前请求处于单账号 503 退避重试模式。 + // 在此模式下，Service 层的模型限流预检查将等待限流过期而非直接切换账号。 + SingleAccountRetry Key = "ctx_single_account_retry" ) diff --git a/backend/internal/service/antigravity_gateway_service.go b/backend/internal/service/antigravity_gateway_service.go index 014b3c86..11f975fe 100644 --- a/backend/internal/service/antigravity_gateway_service.go +++ b/backend/internal/service/antigravity_gateway_service.go @@ -20,6 +20,7 @@ import ( "time" "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity" + "github.com/Wei-Shaw/sub2api/internal/pkg/ctxkey" "github.com/gin-gonic/gin" "github.com/google/uuid" ) @@ -46,6 +47,23 @@ const ( googleRPCTypeErrorInfo = "type.googleapis.com/google.rpc.ErrorInfo" googleRPCReasonModelCapacityExhausted = "MODEL_CAPACITY_EXHAUSTED" 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 透传给客户端的错误消息白名单（小写） @@ -148,6 +166,13 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam // 情况1: retryDelay >= 阈值，限流模型并切换账号 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 if rateLimitDuration <= 0 { rateLimitDuration = antigravityDefaultRateLimitDuration @@ -236,7 +261,7 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam } } - // 所有重试都失败，限流当前模型并切换账号 + // 所有重试都失败 rateLimitDuration := waitDuration if rateLimitDuration <= 0 { rateLimitDuration = antigravityDefaultRateLimitDuration @@ -245,6 +270,22 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam if retryBody == nil { 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)", p.prefix, resp.StatusCode, antigravitySmartRetryMaxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200)) @@ -279,17 +320,152 @@ func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParam 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 的重试循环 func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) { // 预检查：如果账号已限流，直接返回切换信号 if p.requestedModel != "" { 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", - p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID) - return nil, &AntigravityAccountSwitchError{ - OriginalAccountID: p.account.ID, - RateLimitedModel: p.requestedModel, - IsStickySession: p.isStickySession, + // 单账号 503 退避重试模式：跳过限流预检查，直接发请求。 + // 首次请求设的限流是为了多账号调度器跳过该账号，在单账号模式下无意义。 + // 如果上游确实还不可用，handleSmartRetry → handleSingleAccountRetryInPlace + // 会在 Service 层原地等待+重试，不需要在预检查这里等。 + if isSingleAccountRetry(p.ctx) { + 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, + } } } } @@ -1943,6 +2119,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 设置模型级限流 // 直接使用上游返回的模型 ID（如 claude-sonnet-4-5）作为限流 key // 返回是否已成功设置（若模型名为空或 repo 为 nil 将返回 false） From 021abfca181af4f6f52f594200d948de04070119 Mon Sep 17 00:00:00 2001 From: Rose Ding Date: Mon, 9 Feb 2026 17:25:36 +0800 Subject: [PATCH 2/3] =?UTF-8?q?fix:=20=E5=8D=95=E8=B4=A6=E5=8F=B7=E5=88=86?= =?UTF-8?q?=E7=BB=84=E9=A6=96=E6=AC=A1=20503=20=E4=B8=8D=E8=AE=BE=E6=A8=A1?= =?UTF-8?q?=E5=9E=8B=E9=99=90=E6=B5=81=E6=A0=87=E8=AE=B0=EF=BC=8C=E9=81=BF?= =?UTF-8?q?=E5=85=8D=E5=90=8E=E7=BB=AD=E8=AF=B7=E6=B1=82=E9=9B=AA=E5=B4=A9?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 单账号 antigravity 分组收到 503 (MODEL_CAPACITY_EXHAUSTED) 时， 原逻辑会设置 ~29s 模型限流标记。由于只有一个账号无法切换， 后续所有新请求在预检查时命中限流 → 几毫秒内直接返回 503， 导致约 30 秒的雪崩窗口。 修复：在 Handler 入口处检查分组是否只有单个 antigravity 账号， 如果是则提前设置 SingleAccountRetry context 标记，让 Service 层 首次 503 就走原地重试逻辑（不设限流标记），避免污染后续请求。 --- backend/internal/handler/gateway_handler.go | 14 ++++++++++++++ backend/internal/handler/gemini_v1beta_handler.go | 7 +++++++ backend/internal/service/gateway_service.go | 11 +++++++++++ 3 files changed, 32 insertions(+) diff --git a/backend/internal/handler/gateway_handler.go b/backend/internal/handler/gateway_handler.go index 82181948..2b3703b4 100644 --- a/backend/internal/handler/gateway_handler.go +++ b/backend/internal/handler/gateway_handler.go @@ -238,6 +238,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) { var lastFailoverErr *service.UpstreamFailoverError 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 { selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, reqModel, failedAccountIDs, "") // Gemini 不使用会话限制 if err != nil { @@ -409,6 +416,13 @@ func (h *GatewayHandler) Messages(c *gin.Context) { } 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 { maxAccountSwitches := h.maxAccountSwitches switchCount := 0 diff --git a/backend/internal/handler/gemini_v1beta_handler.go b/backend/internal/handler/gemini_v1beta_handler.go index 2b67cb1f..f8fb0dcb 100644 --- a/backend/internal/handler/gemini_v1beta_handler.go +++ b/backend/internal/handler/gemini_v1beta_handler.go @@ -327,6 +327,13 @@ func (h *GatewayHandler) GeminiV1BetaModels(c *gin.Context) { var lastFailoverErr *service.UpstreamFailoverError 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 { selection, err := h.gatewayService.SelectAccountWithLoadAwareness(c.Request.Context(), apiKey.GroupID, sessionKey, modelName, failedAccountIDs, "") // Gemini 不使用会话限制 if err != nil { diff --git a/backend/internal/service/gateway_service.go b/backend/internal/service/gateway_service.go index 4e723232..2c04ae14 100644 --- a/backend/internal/service/gateway_service.go +++ b/backend/internal/service/gateway_service.go @@ -1683,6 +1683,17 @@ func (s *GatewayService) listSchedulableAccounts(ctx context.Context, groupID *i 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 { if account == nil { return false From e4bc35151f4feb41e8283fda1dd8e074a20289a0 Mon Sep 17 00:00:00 2001 From: Rose Ding Date: Mon, 9 Feb 2026 22:06:06 +0800 Subject: [PATCH 3/3] =?UTF-8?q?test:=20=E6=B7=BB=E5=8A=A0=E5=8D=95?= =?UTF-8?q?=E8=B4=A6=E5=8F=B7=20503=20=E9=80=80=E9=81=BF=E9=87=8D=E8=AF=95?= =?UTF-8?q?=E6=9C=BA=E5=88=B6=E7=9A=84=E5=8D=95=E5=85=83=E6=B5=8B=E8=AF=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 覆盖 Service 层和 Handler 层的所有新增逻辑： - isSingleAccountRetry context 标记检查 - handleSmartRetry 中 503 + SingleAccountRetry 分支 - handleSingleAccountRetryInPlace 原地重试逻辑 - antigravityRetryLoop 预检查跳过限流 - sleepAntigravitySingleAccountBackoff 固定延迟退避 - 端到端集成场景验证 Co-Authored-By: Claude Opus 4.6 --- ...teway_handler_single_account_retry_test.go | 51 + .../antigravity_single_account_retry_test.go | 904 ++++++++++++++++++ 2 files changed, 955 insertions(+) create mode 100644 backend/internal/handler/gateway_handler_single_account_retry_test.go create mode 100644 backend/internal/service/antigravity_single_account_retry_test.go diff --git a/backend/internal/handler/gateway_handler_single_account_retry_test.go b/backend/internal/handler/gateway_handler_single_account_retry_test.go new file mode 100644 index 00000000..96aa14c6 --- /dev/null +++ b/backend/internal/handler/gateway_handler_single_account_retry_test.go @@ -0,0 +1,51 @@ +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) +} diff --git a/backend/internal/service/antigravity_single_account_retry_test.go b/backend/internal/service/antigravity_single_account_retry_test.go new file mode 100644 index 00000000..0950b728 --- /dev/null +++ b/backend/internal/service/antigravity_single_account_retry_test.go @@ -0,0 +1,904 @@ +//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") +}