@@ -139,6 +139,7 @@ type queuedHTTPUpstreamStub struct {
139139 errors []error
140140 requestBodies [][]byte
141141 callCount int
142+ onCall func (* http.Request , * queuedHTTPUpstreamStub )
142143}
143144
144145func (s * queuedHTTPUpstreamStub ) Do (req * http.Request , _ string , _ int64 , _ int ) (* http.Response , error ) {
@@ -152,6 +153,9 @@ func (s *queuedHTTPUpstreamStub) Do(req *http.Request, _ string, _ int64, _ int)
152153
153154 idx := s .callCount
154155 s .callCount ++
156+ if s .onCall != nil {
157+ s .onCall (req , s )
158+ }
155159
156160 var resp * http.Response
157161 if idx < len (s .responses ) {
@@ -679,6 +683,91 @@ func TestAntigravityGatewayService_ForwardGemini_RetriesCorruptedThoughtSignatur
679683 require .Equal (t , "signature_error" , events [0 ].Kind )
680684}
681685
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+
682771// TestStreamUpstreamResponse_UsageAndFirstToken
683772// 验证:usage 字段可被累积/覆盖更新,并且能记录首 token 时间
684773func TestStreamUpstreamResponse_UsageAndFirstToken (t * testing.T ) {
0 commit comments