Skip to content

Commit 826090e

Browse files
Wei-ShawWei-Shaw
authored
Merge pull request Wei-Shaw#946 from StarryKira/antigravity-gemini-thought-signature-fix
fix Antigravity gemini thought signature fix
2 parents 7399de6 + 25cb5e7 commit 826090e

3 files changed

Lines changed: 393 additions & 0 deletions

File tree

‎backend/internal/service/antigravity_gateway_service.go‎

Lines changed: 106 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2164,6 +2164,112 @@ func (s *AntigravityGatewayService) ForwardGemini(ctx context.Context, c *gin.Co
21642164
}
21652165
}
21662166

2167+
// Gemini 原生请求中的 thoughtSignature 可能来自旧上下文/旧账号，触发上游严格校验后返回
2168+
// "Corrupted thought signature."。检测到此类 400 时，将 thoughtSignature 清理为 dummy 值后重试一次。
2169+
signatureCheckBody := respBody
2170+
if unwrapped, unwrapErr := s.unwrapV1InternalResponse(respBody); unwrapErr == nil && len(unwrapped) > 0 {
2171+
signatureCheckBody = unwrapped
2172+
}
2173+
if resp.StatusCode == http.StatusBadRequest &&
2174+
s.settingService != nil &&
2175+
s.settingService.IsSignatureRectifierEnabled(ctx) &&
2176+
isSignatureRelatedError(signatureCheckBody) &&
2177+
bytes.Contains(injectedBody, []byte(`"thoughtSignature"`)) {
2178+
upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(signatureCheckBody)))
2179+
upstreamDetail := s.getUpstreamErrorDetail(signatureCheckBody)
2180+
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
2181+
Platform: account.Platform,
2182+
AccountID: account.ID,
2183+
AccountName: account.Name,
2184+
UpstreamStatusCode: resp.StatusCode,
2185+
UpstreamRequestID: resp.Header.Get("x-request-id"),
2186+
Kind: "signature_error",
2187+
Message: upstreamMsg,
2188+
Detail: upstreamDetail,
2189+
})
2190+
2191+
logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: detected signature-related 400, retrying with cleaned thought signatures", account.ID)
2192+
2193+
cleanedInjectedBody := CleanGeminiNativeThoughtSignatures(injectedBody)
2194+
retryWrappedBody, wrapErr := s.wrapV1InternalRequest(projectID, mappedModel, cleanedInjectedBody)
2195+
if wrapErr == nil {
2196+
retryResult, retryErr := s.antigravityRetryLoop(antigravityRetryLoopParams{
2197+
ctx: ctx,
2198+
prefix: prefix,
2199+
account: account,
2200+
proxyURL: proxyURL,
2201+
accessToken: accessToken,
2202+
action: upstreamAction,
2203+
body: retryWrappedBody,
2204+
c: c,
2205+
httpUpstream: s.httpUpstream,
2206+
settingService: s.settingService,
2207+
accountRepo: s.accountRepo,
2208+
handleError: s.handleUpstreamError,
2209+
requestedModel: originalModel,
2210+
isStickySession: isStickySession,
2211+
groupID: 0,
2212+
sessionHash: "",
2213+
})
2214+
if retryErr == nil {
2215+
retryResp := retryResult.resp
2216+
if retryResp.StatusCode < 400 {
2217+
resp = retryResp
2218+
} else {
2219+
retryRespBody, _ := io.ReadAll(io.LimitReader(retryResp.Body, 2<<20))
2220+
_ = retryResp.Body.Close()
2221+
retryOpsBody := retryRespBody
2222+
if retryUnwrapped, unwrapErr := s.unwrapV1InternalResponse(retryRespBody); unwrapErr == nil && len(retryUnwrapped) > 0 {
2223+
retryOpsBody = retryUnwrapped
2224+
}
2225+
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
2226+
Platform: account.Platform,
2227+
AccountID: account.ID,
2228+
AccountName: account.Name,
2229+
UpstreamStatusCode: retryResp.StatusCode,
2230+
UpstreamRequestID: retryResp.Header.Get("x-request-id"),
2231+
Kind: "signature_retry",
2232+
Message: sanitizeUpstreamErrorMessage(strings.TrimSpace(extractAntigravityErrorMessage(retryOpsBody))),
2233+
Detail: s.getUpstreamErrorDetail(retryOpsBody),
2234+
})
2235+
respBody = retryRespBody
2236+
resp = &http.Response{
2237+
StatusCode: retryResp.StatusCode,
2238+
Header: retryResp.Header.Clone(),
2239+
Body: io.NopCloser(bytes.NewReader(retryRespBody)),
2240+
}
2241+
contentType = resp.Header.Get("Content-Type")
2242+
}
2243+
} else {
2244+
if switchErr, ok := IsAntigravityAccountSwitchError(retryErr); ok {
2245+
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
2246+
Platform: account.Platform,
2247+
AccountID: account.ID,
2248+
AccountName: account.Name,
2249+
UpstreamStatusCode: http.StatusServiceUnavailable,
2250+
Kind: "failover",
2251+
Message: sanitizeUpstreamErrorMessage(retryErr.Error()),
2252+
})
2253+
return nil, &UpstreamFailoverError{
2254+
StatusCode: http.StatusServiceUnavailable,
2255+
ForceCacheBilling: switchErr.IsStickySession,
2256+
}
2257+
}
2258+
appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
2259+
Platform: account.Platform,
2260+
AccountID: account.ID,
2261+
AccountName: account.Name,
2262+
UpstreamStatusCode: 0,
2263+
Kind: "signature_retry_request_error",
2264+
Message: sanitizeUpstreamErrorMessage(retryErr.Error()),
2265+
})
2266+
logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry request failed: %v", account.ID, retryErr)
2267+
}
2268+
} else {
2269+
logger.LegacyPrintf("service.antigravity_gateway", "Antigravity Gemini account %d: signature retry wrap failed: %v", account.ID, wrapErr)
2270+
}
2271+
}
2272+
21672273
// fallback 成功：继续按正常响应处理
21682274
if resp.StatusCode < 400 {
21692275
goto handleSuccess

‎backend/internal/service/antigravity_gateway_service_test.go‎

Lines changed: 212 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,47 @@ func (s *httpUpstreamStub) DoWithTLS(_ *http.Request, _ string, _ int64, _ int,
134134
return s.resp, s.err
135135
}
136136

137+
type queuedHTTPUpstreamStub struct {
138+
responses []*http.Response
139+
errors []error
140+
requestBodies [][]byte
141+
callCount int
142+
onCall func(*http.Request, *queuedHTTPUpstreamStub)
143+
}
144+
145+
func (s *queuedHTTPUpstreamStub) Do(req *http.Request, _ string, _ int64, _ int) (*http.Response, error) {
146+
if req != nil && req.Body != nil {
147+
body, _ := io.ReadAll(req.Body)
148+
s.requestBodies = append(s.requestBodies, body)
149+
req.Body = io.NopCloser(bytes.NewReader(body))
150+
} else {
151+
s.requestBodies = append(s.requestBodies, nil)
152+
}
153+
154+
idx := s.callCount
155+
s.callCount++
156+
if s.onCall != nil {
157+
s.onCall(req, s)
158+
}
159+
160+
var resp *http.Response
161+
if idx < len(s.responses) {
162+
resp = s.responses[idx]
163+
}
164+
var err error
165+
if idx < len(s.errors) {
166+
err = s.errors[idx]
167+
}
168+
if resp == nil && err == nil {
169+
return nil, errors.New("unexpected upstream call")
170+
}
171+
return resp, err
172+
}
173+
174+
func (s *queuedHTTPUpstreamStub) DoWithTLS(req *http.Request, proxyURL string, accountID int64, concurrency int, _ bool) (*http.Response, error) {
175+
return s.Do(req, proxyURL, accountID, concurrency)
176+
}
177+
137178
type antigravitySettingRepoStub struct{}
138179

139180
func (s *antigravitySettingRepoStub) Get(ctx context.Context, key string) (*Setting, error) {
@@ -556,6 +597,177 @@ func TestAntigravityGatewayService_ForwardGemini_BillsWithMappedModel(t *testing
556597
require.Equal(t, mappedModel, result.Model)
557598
}
558599

600+
func TestAntigravityGatewayService_ForwardGemini_RetriesCorruptedThoughtSignature(t *testing.T) {
601+
gin.SetMode(gin.TestMode)
602+
writer := httptest.NewRecorder()
603+
c, _ := gin.CreateTestContext(writer)
604+
605+
body, err := json.Marshal(map[string]any{
606+
"contents": []map[string]any{
607+
{"role": "user", "parts": []map[string]any{{"text": "hello"}}},
608+
{"role": "model", "parts": []map[string]any{{"text": "thinking", "thought": true, "thoughtSignature": "sig_bad_1"}}},
609+
{"role": "model", "parts": []map[string]any{{"functionCall": map[string]any{"name": "toolA", "args": map[string]any{"x": 1}}, "thoughtSignature": "sig_bad_2"}}},
610+
},
611+
})
612+
require.NoError(t, err)
613+
614+
req := httptest.NewRequest(http.MethodPost, "/antigravity/v1beta/models/gemini-3.1-pro-preview:streamGenerateContent", bytes.NewReader(body))
615+
c.Request = req
616+
617+
firstRespBody := []byte(`{"response":{"error":{"code":400,"message":"Corrupted thought signature.","status":"INVALID_ARGUMENT"}}}`)
618+
secondRespBody := []byte("data: {\"response\":{\"candidates\":[{\"content\":{\"parts\":[{\"text\":\"ok\"}]},\"finishReason\":\"STOP\"}],\"usageMetadata\":{\"promptTokenCount\":8,\"candidatesTokenCount\":3}}}\n\n")
619+
620+
upstream := &queuedHTTPUpstreamStub{
621+
responses: []*http.Response{
622+
{
623+
StatusCode: http.StatusBadRequest,
624+
Header: http.Header{
625+
"Content-Type": []string{"application/json"},
626+
"X-Request-Id": []string{"req-sig-1"},
627+
},
628+
Body: io.NopCloser(bytes.NewReader(firstRespBody)),
629+
},
630+
{
631+
StatusCode: http.StatusOK,
632+
Header: http.Header{
633+
"Content-Type": []string{"text/event-stream"},
634+
"X-Request-Id": []string{"req-sig-2"},
635+
},
636+
Body: io.NopCloser(bytes.NewReader(secondRespBody)),
637+
},
638+
},
639+
}
640+
641+
svc := &AntigravityGatewayService{
642+
settingService: NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}),
643+
tokenProvider: &AntigravityTokenProvider{},
644+
httpUpstream: upstream,
645+
}
646+
647+
const originalModel = "gemini-3.1-pro-preview"
648+
const mappedModel = "gemini-3.1-pro-high"
649+
account := &Account{
650+
ID: 7,
651+
Name: "acc-gemini-signature",
652+
Platform: PlatformAntigravity,
653+
Type: AccountTypeOAuth,
654+
Status: StatusActive,
655+
Concurrency: 1,
656+
Credentials: map[string]any{
657+
"access_token": "token",
658+
"model_mapping": map[string]any{
659+
originalModel: mappedModel,
660+
},
661+
},
662+
}
663+
664+
result, err := svc.ForwardGemini(context.Background(), c, account, originalModel, "streamGenerateContent", true, body, false)
665+
require.NoError(t, err)
666+
require.NotNil(t, result)
667+
require.Equal(t, mappedModel, result.Model)
668+
require.Len(t, upstream.requestBodies, 2, "signature error should trigger exactly one retry")
669+
670+
firstReq := string(upstream.requestBodies[0])
671+
secondReq := string(upstream.requestBodies[1])
672+
require.Contains(t, firstReq, `"thoughtSignature":"sig_bad_1"`)
673+
require.Contains(t, firstReq, `"thoughtSignature":"sig_bad_2"`)
674+
require.Contains(t, secondReq, `"thoughtSignature":"skip_thought_signature_validator"`)
675+
require.NotContains(t, secondReq, `"thoughtSignature":"sig_bad_1"`)
676+
require.NotContains(t, secondReq, `"thoughtSignature":"sig_bad_2"`)
677+
678+
raw, ok := c.Get(OpsUpstreamErrorsKey)
679+
require.True(t, ok)
680+
events, ok := raw.([]*OpsUpstreamErrorEvent)
681+
require.True(t, ok)
682+
require.NotEmpty(t, events)
683+
require.Equal(t, "signature_error", events[0].Kind)
684+
}
685+
686+
func TestAntigravityGatewayService_ForwardGemini_SignatureRetryPropagatesFailover(t *testing.T) {
687+
gin.SetMode(gin.TestMode)
688+
writer := httptest.NewRecorder()
689+
c, _ := gin.CreateTestContext(writer)
690+
691+
body, err := json.Marshal(map[string]any{
692+
"contents": []map[string]any{
693+
{"role": "user", "parts": []map[string]any{{"text": "hello"}}},
694+
{"role": "model", "parts": []map[string]any{{"text": "thinking", "thought": true, "thoughtSignature": "sig_bad_1"}}},
695+
},
696+
})
697+
require.NoError(t, err)
698+
699+
req := httptest.NewRequest(http.MethodPost, "/antigravity/v1beta/models/gemini-3.1-pro-preview:streamGenerateContent", bytes.NewReader(body))
700+
c.Request = req
701+
702+
firstRespBody := []byte(`{"response":{"error":{"code":400,"message":"Corrupted thought signature.","status":"INVALID_ARGUMENT"}}}`)
703+
704+
const originalModel = "gemini-3.1-pro-preview"
705+
const mappedModel = "gemini-3.1-pro-high"
706+
account := &Account{
707+
ID: 8,
708+
Name: "acc-gemini-signature-failover",
709+
Platform: PlatformAntigravity,
710+
Type: AccountTypeOAuth,
711+
Status: StatusActive,
712+
Concurrency: 1,
713+
Credentials: map[string]any{
714+
"access_token": "token",
715+
"model_mapping": map[string]any{
716+
originalModel: mappedModel,
717+
},
718+
},
719+
}
720+
721+
upstream := &queuedHTTPUpstreamStub{
722+
responses: []*http.Response{
723+
{
724+
StatusCode: http.StatusBadRequest,
725+
Header: http.Header{
726+
"Content-Type": []string{"application/json"},
727+
"X-Request-Id": []string{"req-sig-failover-1"},
728+
},
729+
Body: io.NopCloser(bytes.NewReader(firstRespBody)),
730+
},
731+
},
732+
onCall: func(_ *http.Request, stub *queuedHTTPUpstreamStub) {
733+
if stub.callCount != 1 {
734+
return
735+
}
736+
futureResetAt := time.Now().Add(30 * time.Second).Format(time.RFC3339)
737+
account.Extra = map[string]any{
738+
modelRateLimitsKey: map[string]any{
739+
mappedModel: map[string]any{
740+
"rate_limit_reset_at": futureResetAt,
741+
},
742+
},
743+
}
744+
},
745+
}
746+
747+
svc := &AntigravityGatewayService{
748+
settingService: NewSettingService(&antigravitySettingRepoStub{}, &config.Config{Gateway: config.GatewayConfig{MaxLineSize: defaultMaxLineSize}}),
749+
tokenProvider: &AntigravityTokenProvider{},
750+
httpUpstream: upstream,
751+
}
752+
753+
result, err := svc.ForwardGemini(context.Background(), c, account, originalModel, "streamGenerateContent", true, body, true)
754+
require.Nil(t, result)
755+
756+
var failoverErr *UpstreamFailoverError
757+
require.ErrorAs(t, err, &failoverErr, "signature retry should propagate failover instead of falling back to the original 400")
758+
require.Equal(t, http.StatusServiceUnavailable, failoverErr.StatusCode)
759+
require.True(t, failoverErr.ForceCacheBilling)
760+
require.Len(t, upstream.requestBodies, 1, "retry should stop at preflight failover and not issue a second upstream request")
761+
762+
raw, ok := c.Get(OpsUpstreamErrorsKey)
763+
require.True(t, ok)
764+
events, ok := raw.([]*OpsUpstreamErrorEvent)
765+
require.True(t, ok)
766+
require.Len(t, events, 2)
767+
require.Equal(t, "signature_error", events[0].Kind)
768+
require.Equal(t, "failover", events[1].Kind)
769+
}
770+
559771
// TestStreamUpstreamResponse_UsageAndFirstToken
560772
// 验证：usage 字段可被累积/覆盖更新，并且能记录首 token 时间
561773
func TestStreamUpstreamResponse_UsageAndFirstToken(t *testing.T) {

0 commit comments

Comments
 (0)