antigravity_gateway_retry.go 53 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "errors"
  7. "fmt"
  8. "io"
  9. "log"
  10. mathrand "math/rand"
  11. "net"
  12. "net/http"
  13. "os"
  14. "strconv"
  15. "strings"
  16. "time"
  17. "github.com/Wei-Shaw/sub2api/internal/pkg/antigravity"
  18. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  19. "github.com/gin-gonic/gin"
  20. )
  21. // antigravityRetryLoopParams 重试循环的参数
  22. type antigravityRetryLoopParams struct {
  23. ctx context.Context
  24. prefix string
  25. account *Account
  26. proxyURL string
  27. accessToken string
  28. action string
  29. body []byte
  30. c *gin.Context
  31. httpUpstream HTTPUpstream
  32. settingService *SettingService
  33. accountRepo AccountRepository // 用于智能重试的模型级别限流
  34. 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
  35. requestedModel string // 用于限流检查的原始请求模型
  36. isStickySession bool // 是否为粘性会话(用于账号切换时的缓存计费判断)
  37. groupID int64 // 用于模型级限流时清除粘性会话
  38. sessionHash string // 用于模型级限流时清除粘性会话
  39. }
  40. // antigravityRetryLoopResult 重试循环的结果
  41. type antigravityRetryLoopResult struct {
  42. resp *http.Response
  43. }
  44. // resolveAntigravityForwardBaseURL 解析转发用 base URL。
  45. //
  46. // 显式环境变量优先。未配置时,LoadCodeAssist 返回 paidTier 的付费账号使用
  47. // daily 端点,其他账号继续使用生产端点,避免免费账号的 OAuth token 出现 401。
  48. //
  49. // 历史上这里改用 ForwardBaseURLs()(把 daily/sandbox 排到首位)并默认取首个地址,
  50. // 导致网关把带生产 OAuth token 的请求发到 daily-cloudcode-pa.sandbox.googleapis.com,
  51. // 上游拒绝 → 账号被 401「Invalid bearer token」/502 打入临时不可调度且无法恢复
  52. // (见 #3611 / #2962)。后台「测试连接」用的是生产端点,所以「测试成功但网关 401」。
  53. func resolveAntigravityForwardBaseURL(account *Account) string {
  54. baseURLs := antigravity.BaseURLs
  55. if len(baseURLs) == 0 {
  56. return ""
  57. }
  58. mode := strings.ToLower(strings.TrimSpace(os.Getenv(antigravityForwardBaseURLEnv)))
  59. if (mode == "daily" || mode == "sandbox") && len(baseURLs) > 1 {
  60. return baseURLs[1]
  61. }
  62. if mode == "" && accountHasAntigravityPaidTier(account) && len(baseURLs) > 1 {
  63. return baseURLs[1]
  64. }
  65. return baseURLs[0]
  66. }
  67. func accountHasAntigravityPaidTier(account *Account) bool {
  68. if account == nil || account.Credentials == nil {
  69. return false
  70. }
  71. planType, ok := account.Credentials["plan_type"].(string)
  72. if !ok {
  73. return false
  74. }
  75. switch strings.ToLower(strings.TrimSpace(planType)) {
  76. case "pro", "ultra":
  77. return true
  78. default:
  79. return false
  80. }
  81. }
  82. // smartRetryAction 智能重试的处理结果
  83. type smartRetryAction int
  84. const (
  85. smartRetryActionContinue smartRetryAction = iota // 继续默认重试逻辑
  86. smartRetryActionBreakWithResp // 结束循环并返回 resp
  87. smartRetryActionContinueURL // 继续 URL fallback 循环
  88. )
  89. // smartRetryResult 智能重试的结果
  90. type smartRetryResult struct {
  91. action smartRetryAction
  92. resp *http.Response
  93. err error
  94. switchError *AntigravityAccountSwitchError // 模型限流时返回账号切换信号
  95. }
  96. // handleSmartRetry 处理 OAuth 账号的智能重试逻辑
  97. // 将 429/503 限流处理逻辑抽取为独立函数,减少 antigravityRetryLoop 的复杂度
  98. func (s *AntigravityGatewayService) handleSmartRetry(p antigravityRetryLoopParams, resp *http.Response, respBody []byte, baseURL string, urlIdx int, availableURLs []string) *smartRetryResult {
  99. // "Resource has been exhausted" 是 URL 级别限流,切换 URL(仅 429)
  100. if resp.StatusCode == http.StatusTooManyRequests && isURLLevelRateLimit(respBody) && urlIdx < len(availableURLs)-1 {
  101. logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (429): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1])
  102. return &smartRetryResult{action: smartRetryActionContinueURL}
  103. }
  104. category := antigravity429Unknown
  105. if resp.StatusCode == http.StatusTooManyRequests {
  106. category = classifyAntigravity429(respBody)
  107. }
  108. // 判断是否触发智能重试
  109. shouldSmartRetry, shouldRateLimitModel, waitDuration, modelName, isModelCapacityExhausted := shouldTriggerAntigravitySmartRetry(p.account, respBody)
  110. // AI Credits 超量请求:
  111. // 仅在上游明确返回免费配额耗尽时才允许切换到 credits。
  112. if resp.StatusCode == http.StatusTooManyRequests &&
  113. category == antigravity429QuotaExhausted &&
  114. p.account.IsOveragesEnabled() &&
  115. !p.account.isCreditsExhausted() {
  116. result := s.attemptCreditsOveragesRetry(p, baseURL, modelName, waitDuration, resp.StatusCode, respBody)
  117. if result.handled && result.resp != nil {
  118. return &smartRetryResult{
  119. action: smartRetryActionBreakWithResp,
  120. resp: result.resp,
  121. }
  122. }
  123. }
  124. // 情况1: retryDelay >= 阈值,限流模型并切换账号
  125. if shouldRateLimitModel {
  126. // 单账号 503 退避重试模式:不设限流、不切换账号,改为原地等待+重试
  127. // 谷歌上游 503 (MODEL_CAPACITY_EXHAUSTED) 通常是暂时性的,等几秒就能恢复。
  128. // 多账号场景下切换账号是最优选择,但单账号场景下设限流毫无意义(只会导致双重等待)。
  129. if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) {
  130. return s.handleSingleAccountRetryInPlace(p, resp, respBody, baseURL, waitDuration, modelName)
  131. }
  132. rateLimitDuration := waitDuration
  133. if rateLimitDuration <= 0 {
  134. rateLimitDuration = antigravityDefaultRateLimitDuration
  135. }
  136. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d oauth_long_delay model=%s account=%d upstream_retry_delay=%v body=%s (model rate limit, switch account)",
  137. p.prefix, resp.StatusCode, modelName, p.account.ID, rateLimitDuration, truncateForLog(respBody, 200))
  138. resetAt := time.Now().Add(rateLimitDuration)
  139. if !s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, false) {
  140. p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession)
  141. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited account=%d (no model mapping)", p.prefix, resp.StatusCode, p.account.ID)
  142. }
  143. s.clearStickySession(p.ctx, p.groupID, p.sessionHash)
  144. // 返回账号切换信号,让上层切换账号重试
  145. return &smartRetryResult{
  146. action: smartRetryActionBreakWithResp,
  147. switchError: &AntigravityAccountSwitchError{
  148. OriginalAccountID: p.account.ID,
  149. RateLimitedModel: modelName,
  150. IsStickySession: p.isStickySession,
  151. },
  152. }
  153. }
  154. // 情况2: retryDelay < 阈值(或 MODEL_CAPACITY_EXHAUSTED),智能重试
  155. if shouldSmartRetry {
  156. var lastRetryResp *http.Response
  157. var lastRetryBody []byte
  158. // MODEL_CAPACITY_EXHAUSTED 使用独立的重试参数(60 次,固定 1s 间隔)
  159. maxAttempts := antigravitySmartRetryMaxAttempts
  160. if isModelCapacityExhausted {
  161. maxAttempts = antigravityModelCapacityRetryMaxAttempts
  162. waitDuration = antigravityModelCapacityRetryWait
  163. // 全局去重:如果其他 goroutine 已在重试同一模型且尚在 cooldown 中,直接返回 503
  164. if modelName != "" {
  165. modelCapacityExhaustedMu.RLock()
  166. cooldownUntil, exists := modelCapacityExhaustedUntil[modelName]
  167. modelCapacityExhaustedMu.RUnlock()
  168. if exists && time.Now().Before(cooldownUntil) {
  169. log.Printf("%s status=%d model_capacity_exhausted_dedup model=%s account=%d cooldown_until=%v (skip retry)",
  170. p.prefix, resp.StatusCode, modelName, p.account.ID, cooldownUntil.Format("15:04:05"))
  171. return &smartRetryResult{
  172. action: smartRetryActionBreakWithResp,
  173. resp: &http.Response{
  174. StatusCode: resp.StatusCode,
  175. Header: resp.Header.Clone(),
  176. Body: io.NopCloser(bytes.NewReader(respBody)),
  177. },
  178. }
  179. }
  180. }
  181. }
  182. for attempt := 1; attempt <= maxAttempts; attempt++ {
  183. log.Printf("%s status=%d oauth_smart_retry attempt=%d/%d delay=%v model=%s account=%d",
  184. p.prefix, resp.StatusCode, attempt, maxAttempts, waitDuration, modelName, p.account.ID)
  185. timer := time.NewTimer(waitDuration)
  186. select {
  187. case <-p.ctx.Done():
  188. timer.Stop()
  189. log.Printf("%s status=context_canceled_during_smart_retry", p.prefix)
  190. return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()}
  191. case <-timer.C:
  192. }
  193. // 智能重试:创建新请求
  194. retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body)
  195. if err != nil {
  196. logger.LegacyPrintf("service.antigravity_gateway", "%s status=smart_retry_request_build_failed error=%v", p.prefix, err)
  197. p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession)
  198. return &smartRetryResult{
  199. action: smartRetryActionBreakWithResp,
  200. resp: &http.Response{
  201. StatusCode: resp.StatusCode,
  202. Header: resp.Header.Clone(),
  203. Body: io.NopCloser(bytes.NewReader(respBody)),
  204. },
  205. }
  206. }
  207. retryResp, retryErr := p.httpUpstream.Do(WithAccountTrafficRequest(retryReq, p.account), p.proxyURL, p.account.ID, p.account.Concurrency)
  208. if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable {
  209. log.Printf("%s status=%d smart_retry_success attempt=%d/%d", p.prefix, retryResp.StatusCode, attempt, maxAttempts)
  210. // 重试成功,清除 MODEL_CAPACITY_EXHAUSTED cooldown
  211. if isModelCapacityExhausted && modelName != "" {
  212. modelCapacityExhaustedMu.Lock()
  213. delete(modelCapacityExhaustedUntil, modelName)
  214. modelCapacityExhaustedMu.Unlock()
  215. }
  216. return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp}
  217. }
  218. // 网络错误时,继续重试
  219. if retryErr != nil || retryResp == nil {
  220. log.Printf("%s status=smart_retry_network_error attempt=%d/%d error=%v", p.prefix, attempt, maxAttempts, retryErr)
  221. continue
  222. }
  223. // 重试失败,关闭之前的响应
  224. if lastRetryResp != nil {
  225. _ = lastRetryResp.Body.Close()
  226. }
  227. lastRetryResp = retryResp
  228. if retryResp != nil {
  229. lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10))
  230. _ = retryResp.Body.Close()
  231. }
  232. // 解析新的重试信息,用于下次重试的等待时间(MODEL_CAPACITY_EXHAUSTED 使用固定循环,跳过)
  233. if !isModelCapacityExhausted && attempt < maxAttempts && lastRetryBody != nil {
  234. newShouldRetry, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody)
  235. if newShouldRetry && newWaitDuration > 0 {
  236. waitDuration = newWaitDuration
  237. }
  238. }
  239. }
  240. // 所有重试都失败
  241. rateLimitDuration := waitDuration
  242. if rateLimitDuration <= 0 {
  243. rateLimitDuration = antigravityDefaultRateLimitDuration
  244. }
  245. retryBody := lastRetryBody
  246. if retryBody == nil {
  247. retryBody = respBody
  248. }
  249. // MODEL_CAPACITY_EXHAUSTED:模型容量不足,切换账号无意义
  250. // 直接返回上游错误响应,不设置模型限流,不切换账号
  251. if isModelCapacityExhausted {
  252. // 设置 cooldown,让后续请求快速失败,避免重复重试
  253. if modelName != "" {
  254. modelCapacityExhaustedMu.Lock()
  255. modelCapacityExhaustedUntil[modelName] = time.Now().Add(antigravityModelCapacityCooldown)
  256. modelCapacityExhaustedMu.Unlock()
  257. }
  258. log.Printf("%s status=%d smart_retry_exhausted_model_capacity attempts=%d model=%s account=%d body=%s (model capacity exhausted, not switching account)",
  259. p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200))
  260. return &smartRetryResult{
  261. action: smartRetryActionBreakWithResp,
  262. resp: &http.Response{
  263. StatusCode: resp.StatusCode,
  264. Header: resp.Header.Clone(),
  265. Body: io.NopCloser(bytes.NewReader(retryBody)),
  266. },
  267. }
  268. }
  269. // 单账号 503 退避重试模式:智能重试耗尽后不设限流、不切换账号,
  270. // 直接返回 503 让 Handler 层的单账号退避循环做最终处理。
  271. if resp.StatusCode == http.StatusServiceUnavailable && isSingleAccountRetry(p.ctx) {
  272. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d smart_retry_exhausted_single_account attempts=%d model=%s account=%d body=%s (return 503 directly)",
  273. p.prefix, resp.StatusCode, antigravitySmartRetryMaxAttempts, modelName, p.account.ID, truncateForLog(retryBody, 200))
  274. return &smartRetryResult{
  275. action: smartRetryActionBreakWithResp,
  276. resp: &http.Response{
  277. StatusCode: resp.StatusCode,
  278. Header: resp.Header.Clone(),
  279. Body: io.NopCloser(bytes.NewReader(retryBody)),
  280. },
  281. }
  282. }
  283. log.Printf("%s status=%d smart_retry_exhausted attempts=%d model=%s account=%d upstream_retry_delay=%v body=%s (switch account)",
  284. p.prefix, resp.StatusCode, maxAttempts, modelName, p.account.ID, rateLimitDuration, truncateForLog(retryBody, 200))
  285. resetAt := time.Now().Add(rateLimitDuration)
  286. s.setAntigravityModelRateLimits(p.ctx, p.accountRepo, p.account, modelName, p.prefix, resp.StatusCode, resetAt, true)
  287. // 清除粘性会话绑定,避免下次请求仍命中限流账号
  288. s.clearStickySession(p.ctx, p.groupID, p.sessionHash)
  289. // 返回账号切换信号,让上层切换账号重试
  290. return &smartRetryResult{
  291. action: smartRetryActionBreakWithResp,
  292. switchError: &AntigravityAccountSwitchError{
  293. OriginalAccountID: p.account.ID,
  294. RateLimitedModel: modelName,
  295. IsStickySession: p.isStickySession,
  296. },
  297. }
  298. }
  299. // 未触发智能重试,继续默认重试逻辑
  300. return &smartRetryResult{action: smartRetryActionContinue}
  301. }
  302. // handleSingleAccountRetryInPlace 单账号 503 退避重试的原地重试逻辑。
  303. //
  304. // 在多账号场景下,收到 503 + 长 retryDelay(≥ 7s)时会设置模型限流 + 切换账号;
  305. // 但在单账号场景下,设限流毫无意义(因为切换回来的还是同一个账号,还要等限流过期)。
  306. // 此方法改为在 Service 层原地等待 + 重试,避免双重等待问题:
  307. //
  308. // 旧流程:Service 设限流 → Handler 退避等待 → Service 等限流过期 → 再请求(总耗时 = 退避 + 限流)
  309. // 新流程:Service 直接等 retryDelay → 重试 → 成功/再等 → 重试...(总耗时 ≈ 实际 retryDelay × 重试次数)
  310. //
  311. // 约束:
  312. // - 单次等待不超过 antigravitySingleAccountSmartRetryMaxWait
  313. // - 总累计等待不超过 antigravitySingleAccountSmartRetryTotalMaxWait
  314. // - 最多重试 antigravitySingleAccountSmartRetryMaxAttempts 次
  315. func (s *AntigravityGatewayService) handleSingleAccountRetryInPlace(
  316. p antigravityRetryLoopParams,
  317. resp *http.Response,
  318. respBody []byte,
  319. baseURL string,
  320. waitDuration time.Duration,
  321. modelName string,
  322. ) *smartRetryResult {
  323. // 限制单次等待时间
  324. if waitDuration > antigravitySingleAccountSmartRetryMaxWait {
  325. waitDuration = antigravitySingleAccountSmartRetryMaxWait
  326. }
  327. if waitDuration < antigravitySmartRetryMinWait {
  328. waitDuration = antigravitySmartRetryMinWait
  329. }
  330. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_in_place model=%s account=%d upstream_retry_delay=%v (retrying in-place instead of rate-limiting)",
  331. p.prefix, resp.StatusCode, modelName, p.account.ID, waitDuration)
  332. var lastRetryResp *http.Response
  333. var lastRetryBody []byte
  334. totalWaited := time.Duration(0)
  335. for attempt := 1; attempt <= antigravitySingleAccountSmartRetryMaxAttempts; attempt++ {
  336. // 检查累计等待是否超限
  337. if totalWaited+waitDuration > antigravitySingleAccountSmartRetryTotalMaxWait {
  338. remaining := antigravitySingleAccountSmartRetryTotalMaxWait - totalWaited
  339. if remaining <= 0 {
  340. logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: total_wait_exceeded total=%v max=%v, giving up",
  341. p.prefix, totalWaited, antigravitySingleAccountSmartRetryTotalMaxWait)
  342. break
  343. }
  344. waitDuration = remaining
  345. }
  346. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry attempt=%d/%d delay=%v total_waited=%v model=%s account=%d",
  347. p.prefix, resp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, waitDuration, totalWaited, modelName, p.account.ID)
  348. timer := time.NewTimer(waitDuration)
  349. select {
  350. case <-p.ctx.Done():
  351. timer.Stop()
  352. logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_single_account_retry", p.prefix)
  353. return &smartRetryResult{action: smartRetryActionBreakWithResp, err: p.ctx.Err()}
  354. case <-timer.C:
  355. }
  356. totalWaited += waitDuration
  357. // 创建新请求
  358. retryReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body)
  359. if err != nil {
  360. logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: request_build_failed error=%v", p.prefix, err)
  361. break
  362. }
  363. retryResp, retryErr := p.httpUpstream.Do(WithAccountTrafficRequest(retryReq, p.account), p.proxyURL, p.account.ID, p.account.Concurrency)
  364. if retryErr == nil && retryResp != nil && retryResp.StatusCode != http.StatusTooManyRequests && retryResp.StatusCode != http.StatusServiceUnavailable {
  365. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_success attempt=%d/%d total_waited=%v",
  366. p.prefix, retryResp.StatusCode, attempt, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited)
  367. // 关闭之前的响应
  368. if lastRetryResp != nil {
  369. _ = lastRetryResp.Body.Close()
  370. }
  371. return &smartRetryResult{action: smartRetryActionBreakWithResp, resp: retryResp}
  372. }
  373. // 网络错误时继续重试
  374. if retryErr != nil || retryResp == nil {
  375. logger.LegacyPrintf("service.antigravity_gateway", "%s single_account_503_retry: network_error attempt=%d/%d error=%v",
  376. p.prefix, attempt, antigravitySingleAccountSmartRetryMaxAttempts, retryErr)
  377. continue
  378. }
  379. // 关闭之前的响应
  380. if lastRetryResp != nil {
  381. _ = lastRetryResp.Body.Close()
  382. }
  383. lastRetryResp = retryResp
  384. lastRetryBody, _ = io.ReadAll(io.LimitReader(retryResp.Body, 8<<10))
  385. _ = retryResp.Body.Close()
  386. // 解析新的重试信息,更新下次等待时间
  387. if attempt < antigravitySingleAccountSmartRetryMaxAttempts && lastRetryBody != nil {
  388. _, _, newWaitDuration, _, _ := shouldTriggerAntigravitySmartRetry(p.account, lastRetryBody)
  389. if newWaitDuration > 0 {
  390. waitDuration = newWaitDuration
  391. if waitDuration > antigravitySingleAccountSmartRetryMaxWait {
  392. waitDuration = antigravitySingleAccountSmartRetryMaxWait
  393. }
  394. if waitDuration < antigravitySmartRetryMinWait {
  395. waitDuration = antigravitySmartRetryMinWait
  396. }
  397. }
  398. }
  399. }
  400. // 所有重试都失败,不设限流,直接返回 503
  401. // Handler 层的单账号退避循环会做最终处理
  402. retryBody := lastRetryBody
  403. if retryBody == nil {
  404. retryBody = respBody
  405. }
  406. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d single_account_503_retry_exhausted attempts=%d total_waited=%v model=%s account=%d body=%s (return 503 directly)",
  407. p.prefix, resp.StatusCode, antigravitySingleAccountSmartRetryMaxAttempts, totalWaited, modelName, p.account.ID, truncateForLog(retryBody, 200))
  408. return &smartRetryResult{
  409. action: smartRetryActionBreakWithResp,
  410. resp: &http.Response{
  411. StatusCode: resp.StatusCode,
  412. Header: resp.Header.Clone(),
  413. Body: io.NopCloser(bytes.NewReader(retryBody)),
  414. },
  415. }
  416. }
  417. // antigravityRetryLoop 执行带 URL fallback 的重试循环
  418. func (s *AntigravityGatewayService) antigravityRetryLoop(p antigravityRetryLoopParams) (*antigravityRetryLoopResult, error) {
  419. // 预检查:模型限流 + overages 启用 + 积分未耗尽 → 直接注入 AI Credits
  420. overagesInjected := false
  421. if p.requestedModel != "" && p.account.Platform == PlatformAntigravity &&
  422. p.account.IsOveragesEnabled() && !p.account.isCreditsExhausted() &&
  423. p.account.isModelRateLimitedWithContext(p.ctx, p.requestedModel) {
  424. if creditsBody := injectEnabledCreditTypes(p.body); creditsBody != nil {
  425. p.body = creditsBody
  426. overagesInjected = true
  427. logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: model_rate_limited_credits_inject model=%s account=%d (injecting enabledCreditTypes)",
  428. p.prefix, p.requestedModel, p.account.ID)
  429. }
  430. }
  431. // 预检查:如果账号已限流,直接返回切换信号
  432. if p.requestedModel != "" {
  433. if remaining := p.account.GetRateLimitRemainingTimeWithContext(p.ctx, p.requestedModel); remaining > 0 {
  434. // 已注入积分的请求不再受普通模型限流预检查阻断。
  435. if overagesInjected {
  436. logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: credits_injected_ignore_rate_limit remaining=%v model=%s account=%d",
  437. p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID)
  438. } else if isSingleAccountRetry(p.ctx) {
  439. // 单账号 503 退避重试模式:跳过限流预检查,直接发请求。
  440. // 首次请求设的限流是为了多账号调度器跳过该账号,在单账号模式下无意义。
  441. // 如果上游确实还不可用,handleSmartRetry → handleSingleAccountRetryInPlace
  442. // 会在 Service 层原地等待+重试,不需要在预检查这里等。
  443. logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: single_account_retry skipping rate_limit remaining=%v model=%s account=%d (will retry in-place if 503)",
  444. p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID)
  445. } else {
  446. logger.LegacyPrintf("service.antigravity_gateway", "%s pre_check: rate_limit_switch remaining=%v model=%s account=%d",
  447. p.prefix, remaining.Truncate(time.Millisecond), p.requestedModel, p.account.ID)
  448. return nil, &AntigravityAccountSwitchError{
  449. OriginalAccountID: p.account.ID,
  450. RateLimitedModel: p.requestedModel,
  451. IsStickySession: p.isStickySession,
  452. }
  453. }
  454. }
  455. }
  456. baseURL := resolveAntigravityForwardBaseURL(p.account)
  457. if baseURL == "" {
  458. return nil, errors.New("no antigravity forward base url configured")
  459. }
  460. availableURLs := []string{baseURL}
  461. var resp *http.Response
  462. var usedBaseURL string
  463. logBody := p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBody
  464. maxBytes := 2048
  465. if p.settingService != nil && p.settingService.cfg != nil && p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes > 0 {
  466. maxBytes = p.settingService.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
  467. }
  468. getUpstreamDetail := func(body []byte) string {
  469. if !logBody {
  470. return ""
  471. }
  472. return truncateString(string(body), maxBytes)
  473. }
  474. urlFallbackLoop:
  475. for urlIdx, baseURL := range availableURLs {
  476. usedBaseURL = baseURL
  477. allAttemptsInternal500 := true // 追踪本轮所有 attempt 是否全部命中 INTERNAL 500
  478. for attempt := 1; attempt <= antigravityMaxRetries; attempt++ {
  479. select {
  480. case <-p.ctx.Done():
  481. logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled error=%v", p.prefix, p.ctx.Err())
  482. return nil, p.ctx.Err()
  483. default:
  484. }
  485. upstreamReq, err := antigravity.NewAPIRequestWithURL(p.ctx, baseURL, p.action, p.accessToken, p.body)
  486. if err != nil {
  487. return nil, err
  488. }
  489. resp, err = p.httpUpstream.Do(WithAccountTrafficRequest(upstreamReq, p.account), p.proxyURL, p.account.ID, p.account.Concurrency)
  490. if err == nil && resp == nil {
  491. err = errors.New("upstream returned nil response")
  492. }
  493. if err != nil {
  494. if local := AccountTrafficFailover(err); local != nil {
  495. return nil, local
  496. }
  497. safeErr := sanitizeUpstreamErrorMessage(err.Error())
  498. appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{
  499. ProxyID: opsUpstreamProxyID(p.account),
  500. ProxyName: opsUpstreamProxyName(p.account),
  501. Platform: p.account.Platform,
  502. AccountID: p.account.ID,
  503. AccountName: p.account.Name,
  504. UpstreamStatusCode: 0,
  505. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  506. Kind: "request_error",
  507. Message: safeErr,
  508. })
  509. if shouldAntigravityFallbackToNextURL(err, 0) && urlIdx < len(availableURLs)-1 {
  510. logger.LegacyPrintf("service.antigravity_gateway", "%s URL fallback (connection error): %s -> %s", p.prefix, baseURL, availableURLs[urlIdx+1])
  511. continue urlFallbackLoop
  512. }
  513. if attempt < antigravityMaxRetries {
  514. logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retry=%d/%d error=%v", p.prefix, attempt, antigravityMaxRetries, err)
  515. if !sleepAntigravityBackoffWithContext(p.ctx, attempt) {
  516. logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix)
  517. return nil, p.ctx.Err()
  518. }
  519. continue
  520. }
  521. logger.LegacyPrintf("service.antigravity_gateway", "%s status=request_failed retries_exhausted error=%v", p.prefix, err)
  522. setOpsUpstreamError(p.c, 0, safeErr, "")
  523. return nil, fmt.Errorf("upstream request failed after retries: %w", err)
  524. }
  525. // 统一处理错误响应
  526. if resp.StatusCode >= 400 {
  527. respBody := s.readUpstreamErrorBody(resp)
  528. _ = resp.Body.Close()
  529. if overagesInjected && shouldMarkCreditsExhausted(resp, respBody, nil) {
  530. modelKey := resolveCreditsOveragesModelKey(p.ctx, p.account, "", p.requestedModel)
  531. s.handleCreditsRetryFailure(p.ctx, p.prefix, modelKey, p.account, &http.Response{
  532. StatusCode: resp.StatusCode,
  533. Header: resp.Header.Clone(),
  534. Body: io.NopCloser(bytes.NewReader(respBody)),
  535. }, nil)
  536. }
  537. // ★ 统一入口:自定义错误码 + 临时不可调度
  538. if handled, outStatus, policyErr := s.applyErrorPolicy(p, resp.StatusCode, resp.Header, respBody); handled {
  539. if policyErr != nil {
  540. return nil, policyErr
  541. }
  542. resp = &http.Response{
  543. StatusCode: outStatus,
  544. Header: resp.Header.Clone(),
  545. Body: io.NopCloser(bytes.NewReader(respBody)),
  546. }
  547. break urlFallbackLoop
  548. }
  549. // 429/503 限流处理:区分 URL 级别限流、智能重试和账户配额限流
  550. if resp.StatusCode == http.StatusTooManyRequests || resp.StatusCode == http.StatusServiceUnavailable {
  551. // 尝试智能重试处理(OAuth 账号专用)
  552. smartResult := s.handleSmartRetry(p, resp, respBody, baseURL, urlIdx, availableURLs)
  553. switch smartResult.action {
  554. case smartRetryActionContinueURL:
  555. continue urlFallbackLoop
  556. case smartRetryActionBreakWithResp:
  557. if smartResult.err != nil {
  558. return nil, smartResult.err
  559. }
  560. // 模型限流时返回切换账号信号
  561. if smartResult.switchError != nil {
  562. return nil, smartResult.switchError
  563. }
  564. resp = smartResult.resp
  565. break urlFallbackLoop
  566. }
  567. // smartRetryActionContinue: 继续默认重试逻辑
  568. // 账户/模型配额限流,重试 3 次(指数退避)- 默认逻辑(非 OAuth 账号或解析失败)
  569. if attempt < antigravityMaxRetries {
  570. upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody))
  571. upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
  572. appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{
  573. ProxyID: opsUpstreamProxyID(p.account),
  574. ProxyName: opsUpstreamProxyName(p.account),
  575. Platform: p.account.Platform,
  576. AccountID: p.account.ID,
  577. AccountName: p.account.Name,
  578. UpstreamStatusCode: resp.StatusCode,
  579. UpstreamRequestID: resp.Header.Get("x-request-id"),
  580. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  581. Kind: "retry",
  582. Message: upstreamMsg,
  583. Detail: getUpstreamDetail(respBody),
  584. })
  585. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 200))
  586. if !sleepAntigravityBackoffWithContext(p.ctx, attempt) {
  587. logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix)
  588. return nil, p.ctx.Err()
  589. }
  590. continue
  591. }
  592. // 重试用尽,标记账户限流
  593. p.handleError(p.ctx, p.prefix, p.account, resp.StatusCode, resp.Header, respBody, p.requestedModel, p.groupID, p.sessionHash, p.isStickySession)
  594. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d rate_limited base_url=%s body=%s", p.prefix, resp.StatusCode, baseURL, truncateForLog(respBody, 200))
  595. resp = &http.Response{
  596. StatusCode: resp.StatusCode,
  597. Header: resp.Header.Clone(),
  598. Body: io.NopCloser(bytes.NewReader(respBody)),
  599. }
  600. break urlFallbackLoop
  601. }
  602. // 其他可重试错误(500/502/504/529,不包括 429 和 503)
  603. if shouldRetryAntigravityError(resp.StatusCode) {
  604. if attempt < antigravityMaxRetries {
  605. upstreamMsg := strings.TrimSpace(extractAntigravityErrorMessage(respBody))
  606. upstreamMsg = sanitizeUpstreamErrorMessage(upstreamMsg)
  607. appendOpsUpstreamError(p.c, OpsUpstreamErrorEvent{
  608. ProxyID: opsUpstreamProxyID(p.account),
  609. ProxyName: opsUpstreamProxyName(p.account),
  610. Platform: p.account.Platform,
  611. AccountID: p.account.ID,
  612. AccountName: p.account.Name,
  613. UpstreamStatusCode: resp.StatusCode,
  614. UpstreamRequestID: resp.Header.Get("x-request-id"),
  615. UpstreamURL: safeUpstreamURL(upstreamReq.URL.String()),
  616. Kind: "retry",
  617. Message: upstreamMsg,
  618. Detail: getUpstreamDetail(respBody),
  619. })
  620. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d retry=%d/%d body=%s", p.prefix, resp.StatusCode, attempt, antigravityMaxRetries, truncateForLog(respBody, 500))
  621. if !sleepAntigravityBackoffWithContext(p.ctx, attempt) {
  622. logger.LegacyPrintf("service.antigravity_gateway", "%s status=context_canceled_during_backoff", p.prefix)
  623. return nil, p.ctx.Err()
  624. }
  625. // 追踪 INTERNAL 500:非匹配的 attempt 清除标记
  626. if !isAntigravityInternalServerError(resp.StatusCode, respBody) {
  627. allAttemptsInternal500 = false
  628. }
  629. continue
  630. }
  631. }
  632. // INTERNAL 500 渐进惩罚:3 次重试全部命中特定 500 时递增计数器并惩罚
  633. if allAttemptsInternal500 && isAntigravityInternalServerError(resp.StatusCode, respBody) {
  634. s.handleInternal500RetryExhausted(p.ctx, p.prefix, p.account)
  635. }
  636. // 其他 4xx 错误或重试用尽,直接返回
  637. resp = &http.Response{
  638. StatusCode: resp.StatusCode,
  639. Header: resp.Header.Clone(),
  640. Body: io.NopCloser(bytes.NewReader(respBody)),
  641. }
  642. break urlFallbackLoop
  643. }
  644. // 成功响应(< 400)
  645. break urlFallbackLoop
  646. }
  647. }
  648. if resp != nil && resp.StatusCode < 400 && usedBaseURL != "" {
  649. antigravity.DefaultURLAvailability.MarkSuccess(usedBaseURL)
  650. }
  651. // 成功响应时清零 INTERNAL 500 连续失败计数器(覆盖所有成功路径,含 smart retry)
  652. if resp != nil && resp.StatusCode < 400 {
  653. s.resetInternal500Counter(p.ctx, p.prefix, p.account.ID)
  654. }
  655. return &antigravityRetryLoopResult{resp: resp}, nil
  656. }
  657. // shouldRetryAntigravityError 判断是否应该重试
  658. func shouldRetryAntigravityError(statusCode int) bool {
  659. switch statusCode {
  660. case 429, 500, 502, 503, 504, 529:
  661. return true
  662. default:
  663. return false
  664. }
  665. }
  666. // isURLLevelRateLimit 判断是否为 URL 级别的限流(应切换 URL 重试)
  667. // "Resource has been exhausted" 是 URL/节点级别限流,切换 URL 可能成功
  668. // "exhausted your capacity on this model" 是账户/模型配额限流,切换 URL 无效
  669. func isURLLevelRateLimit(body []byte) bool {
  670. // 快速检查:包含 "Resource has been exhausted" 且不包含 "capacity on this model"
  671. bodyStr := string(body)
  672. return strings.Contains(bodyStr, "Resource has been exhausted") &&
  673. !strings.Contains(bodyStr, "capacity on this model")
  674. }
  675. // isAntigravityConnectionError 判断是否为连接错误(网络超时、DNS 失败、连接拒绝)
  676. func isAntigravityConnectionError(err error) bool {
  677. if err == nil {
  678. return false
  679. }
  680. // 检查超时错误
  681. var netErr net.Error
  682. if errors.As(err, &netErr) && netErr.Timeout() {
  683. return true
  684. }
  685. // 检查连接错误(DNS 失败、连接拒绝)
  686. var opErr *net.OpError
  687. return errors.As(err, &opErr)
  688. }
  689. // shouldAntigravityFallbackToNextURL 判断是否应切换到下一个 URL
  690. // 仅连接错误和 HTTP 429 触发 URL 降级
  691. func shouldAntigravityFallbackToNextURL(err error, statusCode int) bool {
  692. if isAntigravityConnectionError(err) {
  693. return true
  694. }
  695. return statusCode == http.StatusTooManyRequests
  696. }
  697. // getSessionID 从 gin.Context 获取 session_id(用于日志追踪)
  698. func getSessionID(c *gin.Context) string {
  699. if c == nil {
  700. return ""
  701. }
  702. return c.GetHeader("session_id")
  703. }
  704. // logPrefix 生成统一的日志前缀
  705. func logPrefix(sessionID, accountName string) string {
  706. if sessionID != "" {
  707. return fmt.Sprintf("[antigravity-Forward] session=%s account=%s", sessionID, accountName)
  708. }
  709. return fmt.Sprintf("[antigravity-Forward] account=%s", accountName)
  710. }
  711. func (s *AntigravityGatewayService) shouldFailoverUpstreamError(statusCode int) bool {
  712. switch statusCode {
  713. case 401, 403, 429, 529:
  714. return true
  715. default:
  716. return statusCode >= 500
  717. }
  718. }
  719. // isGoogleProjectConfigError 判断(已提取的小写)错误消息是否属于 Google 服务端配置类问题。
  720. // 只精确匹配已知的服务端侧错误,避免对客户端请求错误做无意义重试。
  721. // 适用于所有走 Google 后端的平台(Antigravity、Gemini)。
  722. func isGoogleProjectConfigError(lowerMsg string) bool {
  723. // Google 间歇性 Bug:Project ID 有效但被临时识别失败
  724. return strings.Contains(lowerMsg, "invalid project resource name")
  725. }
  726. // googleConfigErrorCooldown 服务端配置类 400 错误的临时封禁时长
  727. const googleConfigErrorCooldown = 1 * time.Minute
  728. // tempUnscheduleGoogleConfigError 对服务端配置类 400 错误触发临时封禁,
  729. // 避免短时间内反复调度到同一个有问题的账号。
  730. func tempUnscheduleGoogleConfigError(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) {
  731. until := time.Now().Add(googleConfigErrorCooldown)
  732. reason := "400: invalid project resource name (auto temp-unschedule 1m)"
  733. if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil {
  734. log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err)
  735. } else {
  736. log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason)
  737. }
  738. }
  739. // emptyResponseCooldown 空流式响应的临时封禁时长
  740. const emptyResponseCooldown = 1 * time.Minute
  741. // tempUnscheduleEmptyResponse 对空流式响应触发临时封禁,
  742. // 避免短时间内反复调度到同一个返回空响应的账号。
  743. func tempUnscheduleEmptyResponse(ctx context.Context, repo AccountRepository, accountID int64, logPrefix string) {
  744. until := time.Now().Add(emptyResponseCooldown)
  745. reason := "empty stream response (auto temp-unschedule 1m)"
  746. if err := repo.SetTempUnschedulable(ctx, accountID, until, reason); err != nil {
  747. log.Printf("%s temp_unschedule_failed account=%d error=%v", logPrefix, accountID, err)
  748. } else {
  749. log.Printf("%s temp_unscheduled account=%d until=%v reason=%q", logPrefix, accountID, until.Format("15:04:05"), reason)
  750. }
  751. }
  752. // sleepAntigravityBackoffWithContext 带 context 取消检查的退避等待
  753. // 返回 true 表示正常完成等待,false 表示 context 已取消
  754. func sleepAntigravityBackoffWithContext(ctx context.Context, attempt int) bool {
  755. delay := antigravityRetryBaseDelay * time.Duration(1<<uint(attempt-1))
  756. if delay > antigravityRetryMaxDelay {
  757. delay = antigravityRetryMaxDelay
  758. }
  759. // +/- 20% jitter
  760. r := mathrand.New(mathrand.NewSource(time.Now().UnixNano()))
  761. jitter := time.Duration(float64(delay) * 0.2 * (r.Float64()*2 - 1))
  762. sleepFor := delay + jitter
  763. if sleepFor < 0 {
  764. sleepFor = 0
  765. }
  766. timer := time.NewTimer(sleepFor)
  767. select {
  768. case <-ctx.Done():
  769. timer.Stop()
  770. return false
  771. case <-timer.C:
  772. return true
  773. }
  774. }
  775. // isSingleAccountRetry 检查 context 中是否设置了单账号退避重试标记
  776. func isSingleAccountRetry(ctx context.Context) bool {
  777. v, _ := SingleAccountRetryFromContext(ctx)
  778. return v
  779. }
  780. // setModelRateLimitByModelName 使用官方模型 ID 设置模型级限流
  781. // 直接使用上游返回的模型 ID(如 claude-sonnet-4-5)作为限流 key
  782. // 返回是否已成功设置(若模型名为空或 repo 为 nil 将返回 false)
  783. func setModelRateLimitByModelName(ctx context.Context, repo AccountRepository, accountID int64, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool {
  784. if repo == nil || modelName == "" {
  785. return false
  786. }
  787. // 直接使用官方模型 ID 作为 key,不再转换为 scope
  788. if err := repo.SetModelRateLimit(ctx, accountID, modelName, resetAt); err != nil {
  789. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_failed model=%s error=%v", prefix, statusCode, modelName, err)
  790. return false
  791. }
  792. if afterSmartRetry {
  793. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited_after_smart_retry model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second))
  794. } else {
  795. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v", prefix, statusCode, modelName, accountID, time.Until(resetAt).Truncate(time.Second))
  796. }
  797. return true
  798. }
  799. func (s *AntigravityGatewayService) setAntigravityModelRateLimits(ctx context.Context, repo AccountRepository, account *Account, modelName, prefix string, statusCode int, resetAt time.Time, afterSmartRetry bool) bool {
  800. if account == nil || repo == nil {
  801. return false
  802. }
  803. keys := antigravityModelRateLimitKeys(modelName)
  804. if len(keys) == 0 {
  805. return false
  806. }
  807. success := false
  808. for _, key := range keys {
  809. if setModelRateLimitByModelName(ctx, repo, account.ID, key, prefix, statusCode, resetAt, afterSmartRetry) {
  810. s.updateAccountModelRateLimitInCache(ctx, account, key, resetAt)
  811. success = true
  812. }
  813. }
  814. return success
  815. }
  816. func (s *AntigravityGatewayService) clearStickySession(ctx context.Context, groupID int64, sessionHash string) {
  817. if s == nil || s.cache == nil || strings.TrimSpace(sessionHash) == "" {
  818. return
  819. }
  820. if err := s.cache.DeleteSessionAccountID(ctx, groupID, sessionHash); err != nil {
  821. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] sticky_session_clear_failed group_id=%d session=%s err=%v", groupID, shortSessionHash(sessionHash), err)
  822. }
  823. }
  824. func antigravityFallbackCooldownSeconds() (time.Duration, bool) {
  825. raw := strings.TrimSpace(os.Getenv(antigravityFallbackSecondsEnv))
  826. if raw == "" {
  827. return 0, false
  828. }
  829. seconds, err := strconv.Atoi(raw)
  830. if err != nil || seconds <= 0 {
  831. return 0, false
  832. }
  833. return time.Duration(seconds) * time.Second, true
  834. }
  835. // antigravitySmartRetryInfo 智能重试所需的信息
  836. type antigravitySmartRetryInfo struct {
  837. RetryDelay time.Duration // 重试延迟时间
  838. ModelName string // 限流的模型名称(如 "claude-sonnet-4-5")
  839. IsModelCapacityExhausted bool // 是否为模型容量不足(MODEL_CAPACITY_EXHAUSTED)
  840. }
  841. // parseAntigravitySmartRetryInfo 解析 Google RPC RetryInfo 和 ErrorInfo 信息
  842. // 返回解析结果,如果解析失败或不满足条件返回 nil
  843. //
  844. // 支持两种情况:
  845. // 1. 429 RESOURCE_EXHAUSTED + RATE_LIMIT_EXCEEDED:
  846. // - error.status == "RESOURCE_EXHAUSTED"
  847. // - error.details[].reason == "RATE_LIMIT_EXCEEDED"
  848. //
  849. // 2. 503 UNAVAILABLE + MODEL_CAPACITY_EXHAUSTED:
  850. // - error.status == "UNAVAILABLE"
  851. // - error.details[].reason == "MODEL_CAPACITY_EXHAUSTED"
  852. //
  853. // 必须满足以下条件才会返回有效值:
  854. // - error.details[] 中存在 @type == "type.googleapis.com/google.rpc.RetryInfo" 的元素
  855. // - 该元素包含 retryDelay 字段,格式为 "数字s"(如 "0.201506475s")
  856. func parseAntigravitySmartRetryInfo(body []byte) *antigravitySmartRetryInfo {
  857. var parsed map[string]any
  858. if err := json.Unmarshal(body, &parsed); err != nil {
  859. return nil
  860. }
  861. errObj, ok := parsed["error"].(map[string]any)
  862. if !ok {
  863. return nil
  864. }
  865. // 检查 status 是否符合条件
  866. // 情况1: 429 RESOURCE_EXHAUSTED (需要进一步检查 reason == RATE_LIMIT_EXCEEDED)
  867. // 情况2: 503 UNAVAILABLE (需要进一步检查 reason == MODEL_CAPACITY_EXHAUSTED)
  868. status, _ := errObj["status"].(string)
  869. isResourceExhausted := status == googleRPCStatusResourceExhausted
  870. isUnavailable := status == googleRPCStatusUnavailable
  871. if !isResourceExhausted && !isUnavailable {
  872. return nil
  873. }
  874. details, ok := errObj["details"].([]any)
  875. if !ok {
  876. return nil
  877. }
  878. var retryDelay time.Duration
  879. var modelName string
  880. var hasRateLimitExceeded bool // 429 需要此 reason
  881. var hasModelCapacityExhausted bool // 503 需要此 reason
  882. for _, d := range details {
  883. dm, ok := d.(map[string]any)
  884. if !ok {
  885. continue
  886. }
  887. atType, _ := dm["@type"].(string)
  888. // 从 ErrorInfo 提取模型名称和 reason
  889. if atType == googleRPCTypeErrorInfo {
  890. if meta, ok := dm["metadata"].(map[string]any); ok {
  891. if model, ok := meta["model"].(string); ok {
  892. modelName = normalizeAntigravityModelName(model)
  893. }
  894. }
  895. // 检查 reason
  896. if reason, ok := dm["reason"].(string); ok {
  897. if reason == googleRPCReasonModelCapacityExhausted {
  898. hasModelCapacityExhausted = true
  899. }
  900. if reason == googleRPCReasonRateLimitExceeded {
  901. hasRateLimitExceeded = true
  902. }
  903. }
  904. continue
  905. }
  906. // 从 RetryInfo 提取重试延迟
  907. if atType == googleRPCTypeRetryInfo {
  908. delay, ok := dm["retryDelay"].(string)
  909. if !ok || delay == "" {
  910. continue
  911. }
  912. // 使用 time.ParseDuration 解析,支持所有 Go duration 格式
  913. // 例如: "0.5s", "10s", "4m50s", "1h30m", "200ms" 等
  914. dur, err := time.ParseDuration(delay)
  915. if err != nil {
  916. logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity] failed to parse retryDelay: %s error=%v", delay, err)
  917. continue
  918. }
  919. retryDelay = dur
  920. }
  921. }
  922. // 验证条件
  923. // 情况1: RESOURCE_EXHAUSTED 需要有 RATE_LIMIT_EXCEEDED reason
  924. // 情况2: UNAVAILABLE 需要有 MODEL_CAPACITY_EXHAUSTED reason
  925. if isResourceExhausted && !hasRateLimitExceeded {
  926. return nil
  927. }
  928. if isUnavailable && !hasModelCapacityExhausted {
  929. return nil
  930. }
  931. // 必须有模型名才返回有效结果
  932. if modelName == "" {
  933. return nil
  934. }
  935. // 如果上游未提供 retryDelay,使用默认限流时间
  936. if retryDelay <= 0 {
  937. retryDelay = antigravityDefaultRateLimitDuration
  938. }
  939. return &antigravitySmartRetryInfo{
  940. RetryDelay: retryDelay,
  941. ModelName: modelName,
  942. IsModelCapacityExhausted: hasModelCapacityExhausted,
  943. }
  944. }
  945. // shouldTriggerAntigravitySmartRetry 判断是否应该触发智能重试
  946. // 返回:
  947. // - shouldRetry: 是否应该智能重试(retryDelay < antigravityRateLimitThreshold,或 MODEL_CAPACITY_EXHAUSTED)
  948. // - shouldRateLimitModel: 是否应该限流模型并切换账号(仅 RATE_LIMIT_EXCEEDED 且 retryDelay >= 阈值)
  949. // - waitDuration: 等待时间
  950. // - modelName: 限流的模型名称
  951. // - isModelCapacityExhausted: 是否为模型容量不足(MODEL_CAPACITY_EXHAUSTED)
  952. func shouldTriggerAntigravitySmartRetry(account *Account, respBody []byte) (shouldRetry bool, shouldRateLimitModel bool, waitDuration time.Duration, modelName string, isModelCapacityExhausted bool) {
  953. if account.Platform != PlatformAntigravity {
  954. return false, false, 0, "", false
  955. }
  956. info := parseAntigravitySmartRetryInfo(respBody)
  957. if info == nil {
  958. return false, false, 0, "", false
  959. }
  960. // MODEL_CAPACITY_EXHAUSTED(模型容量不足):所有账号共享同一模型容量池
  961. // 切换账号无意义,使用固定 1s 间隔重试
  962. if info.IsModelCapacityExhausted {
  963. return true, false, antigravityModelCapacityRetryWait, info.ModelName, true
  964. }
  965. // RATE_LIMIT_EXCEEDED(账号级限流):
  966. // retryDelay >= 阈值:直接限流模型,不重试
  967. // 注意:如果上游未提供 retryDelay,parseAntigravitySmartRetryInfo 已设置为默认 30s
  968. if info.RetryDelay >= antigravityRateLimitThreshold {
  969. return false, true, info.RetryDelay, info.ModelName, false
  970. }
  971. // retryDelay < 阈值:智能重试
  972. waitDuration = info.RetryDelay
  973. if waitDuration < antigravitySmartRetryMinWait {
  974. waitDuration = antigravitySmartRetryMinWait
  975. }
  976. return true, false, waitDuration, info.ModelName, false
  977. }
  978. // handleModelRateLimitParams 模型级限流处理参数
  979. type handleModelRateLimitParams struct {
  980. ctx context.Context
  981. prefix string
  982. account *Account
  983. statusCode int
  984. body []byte
  985. cache GatewayCache
  986. groupID int64
  987. sessionHash string
  988. isStickySession bool
  989. }
  990. // handleModelRateLimitResult 模型级限流处理结果
  991. type handleModelRateLimitResult struct {
  992. Handled bool // 是否已处理
  993. ShouldRetry bool // 是否等待后重试
  994. WaitDuration time.Duration // 等待时间
  995. SwitchError *AntigravityAccountSwitchError // 账号切换错误
  996. }
  997. // handleModelRateLimit 处理模型级限流(在原有逻辑之前调用)
  998. // 仅处理 429/503,解析模型名和 retryDelay
  999. // - MODEL_CAPACITY_EXHAUSTED: 返回 Handled=true(实际重试由 handleSmartRetry 处理)
  1000. // - RATE_LIMIT_EXCEEDED + retryDelay < 阈值: 返回 ShouldRetry=true,由调用方等待后重试
  1001. // - RATE_LIMIT_EXCEEDED + retryDelay >= 阈值: 设置模型限流 + 清除粘性会话 + 返回 SwitchError
  1002. func (s *AntigravityGatewayService) handleModelRateLimit(p *handleModelRateLimitParams) *handleModelRateLimitResult {
  1003. if p.statusCode != 429 && p.statusCode != 503 {
  1004. return &handleModelRateLimitResult{Handled: false}
  1005. }
  1006. info := parseAntigravitySmartRetryInfo(p.body)
  1007. if info == nil || info.ModelName == "" {
  1008. return &handleModelRateLimitResult{Handled: false}
  1009. }
  1010. // MODEL_CAPACITY_EXHAUSTED:模型容量不足,所有账号共享同一容量池
  1011. // 切换账号无意义,不设置模型限流(实际重试由 handleSmartRetry 处理)
  1012. if info.IsModelCapacityExhausted {
  1013. log.Printf("%s status=%d model_capacity_exhausted model=%s (not switching account, retry handled by smart retry)",
  1014. p.prefix, p.statusCode, info.ModelName)
  1015. return &handleModelRateLimitResult{
  1016. Handled: true,
  1017. }
  1018. }
  1019. // RATE_LIMIT_EXCEEDED: < antigravityRateLimitThreshold: 等待后重试
  1020. if info.RetryDelay < antigravityRateLimitThreshold {
  1021. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limit_wait model=%s wait=%v",
  1022. p.prefix, p.statusCode, info.ModelName, info.RetryDelay)
  1023. return &handleModelRateLimitResult{
  1024. Handled: true,
  1025. ShouldRetry: true,
  1026. WaitDuration: info.RetryDelay,
  1027. }
  1028. }
  1029. // RATE_LIMIT_EXCEEDED: >= antigravityRateLimitThreshold: 设置限流 + 清除粘性会话 + 切换账号
  1030. s.setModelRateLimitAndClearSession(p, info)
  1031. return &handleModelRateLimitResult{
  1032. Handled: true,
  1033. SwitchError: &AntigravityAccountSwitchError{
  1034. OriginalAccountID: p.account.ID,
  1035. RateLimitedModel: info.ModelName,
  1036. IsStickySession: p.isStickySession,
  1037. },
  1038. }
  1039. }
  1040. // setModelRateLimitAndClearSession 设置模型限流并清除粘性会话
  1041. func (s *AntigravityGatewayService) setModelRateLimitAndClearSession(p *handleModelRateLimitParams, info *antigravitySmartRetryInfo) {
  1042. resetAt := time.Now().Add(info.RetryDelay)
  1043. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d model_rate_limited model=%s account=%d reset_in=%v",
  1044. p.prefix, p.statusCode, info.ModelName, p.account.ID, info.RetryDelay)
  1045. s.setAntigravityModelRateLimits(p.ctx, s.accountRepo, p.account, info.ModelName, p.prefix, p.statusCode, resetAt, false)
  1046. // 清除粘性会话绑定
  1047. if p.cache != nil && p.sessionHash != "" {
  1048. _ = p.cache.DeleteSessionAccountID(p.ctx, p.groupID, p.sessionHash)
  1049. }
  1050. }
  1051. // updateAccountModelRateLimitInCache 立即更新 Redis 中账号的模型限流状态
  1052. func (s *AntigravityGatewayService) updateAccountModelRateLimitInCache(ctx context.Context, account *Account, modelKey string, resetAt time.Time) {
  1053. if s.schedulerSnapshot == nil || account == nil || modelKey == "" {
  1054. return
  1055. }
  1056. // 更新账号对象的 Extra 字段
  1057. if account.Extra == nil {
  1058. account.Extra = make(map[string]any)
  1059. }
  1060. limits, _ := account.Extra["model_rate_limits"].(map[string]any)
  1061. if limits == nil {
  1062. limits = make(map[string]any)
  1063. account.Extra["model_rate_limits"] = limits
  1064. }
  1065. limits[modelKey] = map[string]any{
  1066. "rate_limited_at": time.Now().UTC().Format(time.RFC3339),
  1067. "rate_limit_reset_at": resetAt.UTC().Format(time.RFC3339),
  1068. }
  1069. // 更新 Redis 快照
  1070. if err := s.schedulerSnapshot.UpdateAccountInCache(ctx, account); err != nil {
  1071. logger.LegacyPrintf("service.antigravity_gateway", "[antigravity-Forward] cache_update_failed account=%d model=%s err=%v", account.ID, modelKey, err)
  1072. }
  1073. }
  1074. func (s *AntigravityGatewayService) handleUpstreamError(
  1075. ctx context.Context, prefix string, account *Account,
  1076. statusCode int, headers http.Header, body []byte,
  1077. requestedModel string,
  1078. groupID int64, sessionHash string, isStickySession bool,
  1079. ) *handleModelRateLimitResult {
  1080. // 遵守自定义错误码策略:未命中则跳过所有限流处理
  1081. if !account.ShouldHandleErrorCode(statusCode) {
  1082. return nil
  1083. }
  1084. // 模型级限流处理(优先)
  1085. result := s.handleModelRateLimit(&handleModelRateLimitParams{
  1086. ctx: ctx,
  1087. prefix: prefix,
  1088. account: account,
  1089. statusCode: statusCode,
  1090. body: body,
  1091. cache: s.cache,
  1092. groupID: groupID,
  1093. sessionHash: sessionHash,
  1094. isStickySession: isStickySession,
  1095. })
  1096. if result.Handled {
  1097. return result
  1098. }
  1099. // 503 仅处理模型限流(MODEL_CAPACITY_EXHAUSTED),非模型限流不做额外处理
  1100. // 避免将普通的 503 错误误判为账号问题
  1101. if statusCode == 503 {
  1102. return nil
  1103. }
  1104. // 429:尝试解析模型级限流,解析失败时兜底为账号级限流
  1105. if statusCode == 429 {
  1106. if logBody, maxBytes := s.getLogConfig(); logBody {
  1107. logger.LegacyPrintf("service.antigravity_gateway", "[Antigravity-Debug] 429 response body: %s", truncateString(string(body), maxBytes))
  1108. }
  1109. resetAt := ParseGeminiRateLimitResetTime(body)
  1110. defaultDur := s.getDefaultRateLimitDuration()
  1111. // 尝试解析模型 key 并设置模型级限流
  1112. //
  1113. // 注意:requestedModel 可能是"映射前"的请求模型名(例如 claude-opus-4-6),
  1114. // 调度与限流判定使用的是 Antigravity 最终模型名(包含映射与 thinking 后缀)。
  1115. // 因此这里必须写入最终模型 key,确保后续调度能正确避开已限流模型。
  1116. modelKey := resolveFinalAntigravityModelKey(ctx, account, requestedModel)
  1117. if strings.TrimSpace(modelKey) == "" {
  1118. // 极少数情况下无法映射(理论上不应发生:能转发成功说明映射已通过),
  1119. // 保持旧行为作为兜底,避免完全丢失模型级限流记录。
  1120. modelKey = resolveAntigravityModelKey(requestedModel)
  1121. }
  1122. if modelKey != "" {
  1123. ra := s.resolveResetTime(resetAt, defaultDur)
  1124. if !s.setAntigravityModelRateLimits(ctx, s.accountRepo, account, modelKey, prefix, statusCode, ra, false) {
  1125. logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limit_set_failed model=%s", prefix, modelKey)
  1126. } else {
  1127. logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 model_rate_limited model=%s account=%d reset_at=%v reset_in=%v",
  1128. prefix, modelKey, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second))
  1129. }
  1130. return nil
  1131. }
  1132. // 无法解析模型 key,兜底为账号级限流
  1133. ra := s.resolveResetTime(resetAt, defaultDur)
  1134. logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limited account=%d reset_at=%v reset_in=%v (fallback)",
  1135. prefix, account.ID, ra.Format("15:04:05"), time.Until(ra).Truncate(time.Second))
  1136. if err := s.accountRepo.SetRateLimited(ctx, account.ID, ra); err != nil {
  1137. logger.LegacyPrintf("service.antigravity_gateway", "%s status=429 rate_limit_set_failed account=%d error=%v", prefix, account.ID, err)
  1138. }
  1139. return nil
  1140. }
  1141. // 其他错误码继续使用 rateLimitService
  1142. if s.rateLimitService == nil {
  1143. return nil
  1144. }
  1145. shouldDisable := s.rateLimitService.HandleUpstreamError(ctx, account, statusCode, headers, body)
  1146. if shouldDisable {
  1147. logger.LegacyPrintf("service.antigravity_gateway", "%s status=%d marked_error", prefix, statusCode)
  1148. }
  1149. return nil
  1150. }
  1151. // getDefaultRateLimitDuration 获取默认限流时间
  1152. func (s *AntigravityGatewayService) getDefaultRateLimitDuration() time.Duration {
  1153. defaultDur := antigravityDefaultRateLimitDuration
  1154. if s.settingService != nil && s.settingService.cfg != nil && s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes > 0 {
  1155. defaultDur = time.Duration(s.settingService.cfg.Gateway.AntigravityFallbackCooldownMinutes) * time.Minute
  1156. }
  1157. if override, ok := antigravityFallbackCooldownSeconds(); ok {
  1158. defaultDur = override
  1159. }
  1160. return defaultDur
  1161. }
  1162. // resolveResetTime 根据解析的重置时间或默认时长计算重置时间点
  1163. func (s *AntigravityGatewayService) resolveResetTime(resetAt *int64, defaultDur time.Duration) time.Time {
  1164. if resetAt != nil {
  1165. return time.Unix(*resetAt, 0)
  1166. }
  1167. return time.Now().Add(defaultDur)
  1168. }