Skip to content

Commit 2b41cec

Browse files
Wei-ShawWei-Shaw
authored
Merge pull request Wei-Shaw#1076 from touwaeriol/pr/antigravity-test-connection-unify
refactor(antigravity): unify TestConnection with dispatch retry loop
2 parents 6cf7704 + a6f99cf commit 2b41cec

3 files changed

Lines changed: 111 additions & 73 deletions

File tree

‎backend/internal/service/antigravity_gateway_service.go‎

Lines changed: 55 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -930,7 +930,7 @@ func (s *AntigravityGatewayService) applyErrorPolicy(p antigravityRetryLoopParam
930930
case ErrorPolicyTempUnscheduled:
931931
slog.Info("temp_unschedulable_matched",
932932
"prefix", p.prefix, "status_code", statusCode, "account_id", p.account.ID)
933-
return true, statusCode, &AntigravityAccountSwitchError{OriginalAccountID: p.account.ID, IsStickySession: p.isStickySession}
933+
return true, statusCode, &AntigravityAccountSwitchError{OriginalAccountID: p.account.ID, RateLimitedModel: p.requestedModel, IsStickySession: p.isStickySession}
934934
}
935935
return false, statusCode, nil
936936
}
@@ -1001,8 +1001,9 @@ type TestConnectionResult struct {
10011001
MappedModel string // 实际使用的模型
10021002
}
10031003

1004-
// TestConnection 测试 Antigravity 账号连接（非流式，无重试、无计费）
1005-
// 支持 Claude 和 Gemini 两种协议，根据 modelID 前缀自动选择
1004+
// TestConnection 测试 Antigravity 账号连接。
1005+
// 复用 antigravityRetryLoop 的完整重试 / credits overages / 智能重试逻辑，
1006+
// 与真实调度行为一致。差异：不做账号切换（测试指定账号）、不记录 ops 错误。
10061007
func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account *Account, modelID string) (*TestConnectionResult, error) {
10071008

10081009
// 获取 token
@@ -1026,10 +1027,8 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
10261027
// 构建请求体
10271028
var requestBody []byte
10281029
if strings.HasPrefix(modelID, "gemini-") {
1029-
// Gemini 模型：直接使用 Gemini 格式
10301030
requestBody, err = s.buildGeminiTestRequest(projectID, mappedModel)
10311031
} else {
1032-
// Claude 模型：使用协议转换
10331032
requestBody, err = s.buildClaudeTestRequest(projectID, mappedModel)
10341033
}
10351034
if err != nil {
@@ -1042,64 +1041,63 @@ func (s *AntigravityGatewayService) TestConnection(ctx context.Context, account
10421041
proxyURL = account.Proxy.URL()
10431042
}
10441043

1045-
baseURL := resolveAntigravityForwardBaseURL()
1046-
if baseURL == "" {
1047-
return nil, errors.New("no antigravity forward base url configured")
1048-
}
1049-
availableURLs := []string{baseURL}
1050-
1051-
var lastErr error
1052-
for urlIdx, baseURL := range availableURLs {
1053-
// 构建 HTTP 请求（总是使用流式 endpoint，与官方客户端一致）
1054-
req, err := antigravity.NewAPIRequestWithURL(ctx, baseURL, "streamGenerateContent", accessToken, requestBody)
1055-
if err != nil {
1056-
lastErr = err
1057-
continue
1058-
}
1059-
1060-
// 调试日志：Test 请求信息
1061-
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Test] account=%s request_size=%d url=%s", account.Name, len(requestBody), req.URL.String())
1062-
1063-
// 发送请求
1064-
resp, err := s.httpUpstream.Do(req, proxyURL, account.ID, account.Concurrency)
1065-
if err != nil {
1066-
lastErr = fmt.Errorf("请求失败: %w", err)
1067-
if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
1068-
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Test] URL fallback: %s -> %s", baseURL, availableURLs[urlIdx+1])
1069-
continue
1070-
}
1071-
return nil, lastErr
1072-
}
1073-
1074-
// 读取响应
1075-
respBody, err := io.ReadAll(io.LimitReader(resp.Body, 2<<20))
1076-
_ = resp.Body.Close() // 立即关闭，避免循环内 defer 导致的资源泄漏
1077-
if err != nil {
1078-
return nil, fmt.Errorf("读取响应失败: %w", err)
1079-
}
1080-
1081-
// 检查是否需要 URL 降级
1082-
if shouldAntigravityFallbackToNextURL(nil, resp.StatusCode) && urlIdx < len(availableURLs)-1 {
1083-
logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Test] URL fallback (HTTP %d): %s -> %s", resp.StatusCode, baseURL, availableURLs[urlIdx+1])
1084-
continue
1044+
// 复用 antigravityRetryLoop：完整的重试 / credits overages / 智能重试
1045+
prefix := fmt.Sprintf("[antigravity-Test] account=%d(%s)", account.ID, account.Name)
1046+
p := antigravityRetryLoopParams{
1047+
ctx: ctx,
1048+
prefix: prefix,
1049+
account: account,
1050+
proxyURL: proxyURL,
1051+
accessToken: accessToken,
1052+
action: "streamGenerateContent",
1053+
body: requestBody,
1054+
c: nil, // 无 gin.Context → 跳过 ops 追踪
1055+
httpUpstream: s.httpUpstream,
1056+
settingService: s.settingService,
1057+
accountRepo: s.accountRepo,
1058+
requestedModel: modelID,
1059+
handleError: testConnectionHandleError,
1060+
}
1061+
1062+
result, err := s.antigravityRetryLoop(p)
1063+
if err != nil {
1064+
// AccountSwitchError → 测试时不切换账号，返回友好提示
1065+
var switchErr *AntigravityAccountSwitchError
1066+
if errors.As(err, &switchErr) {
1067+
return nil, fmt.Errorf("该账号模型 %s 当前限流中，请稍后重试", switchErr.RateLimitedModel)
10851068
}
1069+
return nil, err
1070+
}
10861071

1087-
if resp.StatusCode >= 400 {
1088-
return nil, fmt.Errorf("API 返回 %d: %s", resp.StatusCode, string(respBody))
1089-
}
1072+
if result == nil || result.resp == nil {
1073+
return nil, errors.New("upstream returned empty response")
1074+
}
1075+
defer func() { _ = result.resp.Body.Close() }()
10901076

1091-
// 解析流式响应，提取文本
1092-
text := extractTextFromSSEResponse(respBody)
1077+
respBody, err := io.ReadAll(io.LimitReader(result.resp.Body, 2<<20))
1078+
if err != nil {
1079+
return nil, fmt.Errorf("读取响应失败: %w", err)
1080+
}
10931081

1094-
// 标记成功的 URL，下次优先使用
1095-
antigravity.DefaultURLAvailability.MarkSuccess(baseURL)
1096-
return &TestConnectionResult{
1097-
Text: text,
1098-
MappedModel: mappedModel,
1099-
}, nil
1082+
if result.resp.StatusCode >= 400 {
1083+
return nil, fmt.Errorf("API 返回 %d: %s", result.resp.StatusCode, string(respBody))
11001084
}
11011085

1102-
return nil, lastErr
1086+
text := extractTextFromSSEResponse(respBody)
1087+
return &TestConnectionResult{Text: text, MappedModel: mappedModel}, nil
1088+
}
1089+
1090+
// testConnectionHandleError 是 TestConnection 使用的轻量 handleError 回调。
1091+
// 仅记录日志，不做 ops 错误追踪或粘性会话清除。
1092+
func testConnectionHandleError(
1093+
_ context.Context, prefix string, account *Account,
1094+
statusCode int, _ http.Header, body []byte,
1095+
requestedModel string, _ int64, _ string, _ bool,
1096+
) *handleModelRateLimitResult {
1097+
logger.LegacyPrintf("service.antigravity_gateway",
1098+
"%s test_handle_error status=%d model=%s account=%d body=%s",
1099+
prefix, statusCode, requestedModel, account.ID, truncateForLog(body, 200))
1100+
return nil
11031101
}
11041102

11051103
// buildGeminiTestRequest 构建 Gemini 格式测试请求

‎backend/internal/service/antigravity_single_account_retry_test.go‎

Lines changed: 11 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -260,14 +260,15 @@ func TestHandleSmartRetry_429_LongDelay_SingleAccountRetry_StillSwitches(t *test
260260

261261
// TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit
262262
// 503 + retryDelay < 7s + SingleAccountRetry → 智能重试耗尽后直接返回 503，不设限流
263+
// 使用 RATE_LIMIT_EXCEEDED（走 1 次智能重试），避免 MODEL_CAPACITY_EXHAUSTED 的 60 次重试导致测试超时
263264
func TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit(t *testing.T) {
264265
// 智能重试也返回 503
265266
failRespBody := `{
266267
"error": {
267268
"code": 503,
268-
"status": "UNAVAILABLE",
269+
"status": "RESOURCE_EXHAUSTED",
269270
"details": [
270-
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
271+
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
271272
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
272273
]
273274
}
@@ -278,8 +279,9 @@ func TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit(t *testi
278279
Body: io.NopCloser(strings.NewReader(failRespBody)),
279280
}
280281
upstream := &mockSmartRetryUpstream{
281-
responses: []*http.Response{failResp},
282-
errors: []error{nil},
282+
responses: []*http.Response{failResp},
283+
errors: []error{nil},
284+
repeatLast: true,
283285
}
284286

285287
repo := &stubAntigravityAccountRepo{}
@@ -294,9 +296,9 @@ func TestHandleSmartRetry_503_ShortDelay_SingleAccountRetry_NoRateLimit(t *testi
294296
respBody := []byte(`{
295297
"error": {
296298
"code": 503,
297-
"status": "UNAVAILABLE",
299+
"status": "RESOURCE_EXHAUSTED",
298300
"details": [
299-
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "MODEL_CAPACITY_EXHAUSTED"},
301+
{"@type": "type.googleapis.com/google.rpc.ErrorInfo", "metadata": {"model": "gemini-3-flash"}, "reason": "RATE_LIMIT_EXCEEDED"},
300302
{"@type": "type.googleapis.com/google.rpc.RetryInfo", "retryDelay": "0.1s"}
301303
]
302304
}
@@ -569,8 +571,9 @@ func TestHandleSingleAccountRetryInPlace_WaitDurationClamped(t *testing.T) {
569571

570572
svc := &AntigravityGatewayService{}
571573

572-
// 等待时间过大应被 clamp 到 antigravitySingleAccountSmartRetryMaxWait
573-
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 999*time.Second, "gemini-3-pro")
574+
// waitDuration=0 会被 clamp 到 antigravitySmartRetryMinWait=1s。
575+
// 首次重试即成功（200），总耗时 ~1s。
576+
result := svc.handleSingleAccountRetryInPlace(params, resp, nil, "https://ag-1.test", 0, "gemini-3-pro")
574577
require.NotNil(t, result)
575578
require.Equal(t, smartRetryActionBreakWithResp, result.action)
576579
require.NotNil(t, result.resp)

‎backend/internal/service/antigravity_smart_retry_test.go‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -32,11 +32,13 @@ func (c *stubSmartRetryCache) DeleteSessionAccountID(_ context.Context, groupID
3232

3333
// mockSmartRetryUpstream 用于 handleSmartRetry 测试的 mock upstream
3434
type mockSmartRetryUpstream struct {
35-
responses []*http.Response
36-
errors []error
37-
callIdx int
38-
calls []string
39-
requestBodies [][]byte
35+
responses []*http.Response
36+
responseBodies [][]byte // 缓存的 response body 字节（用于 repeatLast 重建）
37+
errors []error
38+
callIdx int
39+
calls []string
40+
requestBodies [][]byte
41+
repeatLast bool // 超出范围时重复最后一个响应
4042
}
4143

4244
func (m *mockSmartRetryUpstream) Do(req *http.Request, proxyURL string, accountID int64, accountConcurrency int) (*http.Response, error) {
@@ -50,10 +52,45 @@ func (m *mockSmartRetryUpstream) Do(req *http.Request, proxyURL string, accountI
5052
m.requestBodies = append(m.requestBodies, nil)
5153
}
5254
m.callIdx++
53-
if idx < len(m.responses) {
54-
return m.responses[idx], m.errors[idx]
55+
56+
// 确定使用哪个索引
57+
respIdx := idx
58+
if respIdx >= len(m.responses) {
59+
if !m.repeatLast || len(m.responses) == 0 {
60+
return nil, nil
61+
}
62+
respIdx = len(m.responses) - 1
63+
}
64+
65+
resp := m.responses[respIdx]
66+
respErr := m.errors[respIdx]
67+
if resp == nil {
68+
return nil, respErr
5569
}
56-
return nil, nil
70+
71+
// 首次调用时缓存 body 字节
72+
if respIdx >= len(m.responseBodies) {
73+
for len(m.responseBodies) <= respIdx {
74+
m.responseBodies = append(m.responseBodies, nil)
75+
}
76+
}
77+
if m.responseBodies[respIdx] == nil && resp.Body != nil {
78+
bodyBytes, _ := io.ReadAll(resp.Body)
79+
_ = resp.Body.Close()
80+
m.responseBodies[respIdx] = bodyBytes
81+
}
82+
83+
// 用缓存的 body 字节重建新的 reader
84+
var body io.ReadCloser
85+
if m.responseBodies[respIdx] != nil {
86+
body = io.NopCloser(bytes.NewReader(m.responseBodies[respIdx]))
87+
}
88+
89+
return &http.Response{
90+
StatusCode: resp.StatusCode,
91+
Header: resp.Header.Clone(),
92+
Body: body,
93+
}, respErr
5794
}
5895

5996
func (m *mockSmartRetryUpstream) DoWithTLS(req *http.Request, proxyURL string, accountID int64, accountConcurrency int, enableTLSFingerprint bool) (*http.Response, error) {

0 commit comments

Comments
 (0)