openai_gateway_grok.go 69 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "log/slog"
  9. "net/http"
  10. "strings"
  11. "time"
  12. "github.com/Wei-Shaw/sub2api/internal/config"
  13. "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
  14. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  15. "github.com/gin-gonic/gin"
  16. "github.com/tidwall/gjson"
  17. "github.com/tidwall/sjson"
  18. )
  19. const (
  20. grokComposerImageBridgeVisionModel = "grok-build-0.1"
  21. grokComposerImageBridgeMaxOutputTokens = 512
  22. grokCLIVersion = xai.CLIClientVersion
  23. grokDefaultResponsesModel = "grok-4.5"
  24. grokRateLimitFallbackCooldown = 2 * time.Minute
  25. grokRateLimitRepeatCooldown = 10 * time.Minute
  26. grokRateLimitSustainedCooldown = 30 * time.Minute
  27. grokRateLimitMaxAdaptiveCooldown = time.Hour
  28. grokRateLimitBackoffQuietPeriod = time.Hour
  29. )
  30. func (s *OpenAIGatewayService) forwardGrokResponses(
  31. ctx context.Context,
  32. c *gin.Context,
  33. account *Account,
  34. body []byte,
  35. originalModel string,
  36. reqStream bool,
  37. startTime time.Time,
  38. ) (*OpenAIForwardResult, error) {
  39. if account.Type != AccountTypeOAuth && account.Type != AccountTypeAPIKey {
  40. return nil, fmt.Errorf("grok account type %s is not supported by Responses forwarding", account.Type)
  41. }
  42. upstreamModel := account.GetMappedModel(originalModel)
  43. if strings.TrimSpace(upstreamModel) == "" {
  44. upstreamModel = grokDefaultResponsesModel
  45. }
  46. // Account mappings are optional. Canonicalize client aliases even when the
  47. // account has no model_mapping, matching the Chat Completions path and xAI's
  48. // actual Responses model IDs.
  49. upstreamModel = xai.ResolveGrokTextResponsesModelID(upstreamModel, grokDefaultResponsesModel)
  50. if isGrokImageGenerationModel(upstreamModel) {
  51. return nil, fmt.Errorf("model %s is an image model and is not available on the Responses endpoint; use /v1/images/generations instead", upstreamModel)
  52. }
  53. patchedBody, clientToolMapping, err := patchGrokResponsesBodyWithClientTools(body, upstreamModel)
  54. if err != nil {
  55. setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "")
  56. c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
  57. "type": "invalid_request_error", "message": err.Error(), "param": "tools",
  58. }})
  59. return nil, err
  60. }
  61. setGrokResponsesClientToolMapping(c, clientToolMapping)
  62. // OpenAI /responses/compact is not a native xAI endpoint. Convert it into a
  63. // normal Grok Responses turn that asks for a structured summary, then map the
  64. // reply back to an OpenAI compaction item on the way out.
  65. if isOpenAIResponsesCompactPath(c) {
  66. patchedBody, err = buildGrokCompactRequestBody(patchedBody)
  67. if err != nil {
  68. return nil, err
  69. }
  70. }
  71. // Derive the identity from the request xAI will actually see. This makes
  72. // Codex Responses Lite additional_tools part of the stable tool prefix. If
  73. // the client supplied a Claude Code session only through metadata.user_id,
  74. // keep using that identity even though metadata is stripped before xAI.
  75. cacheIdentityBody := patchedBody
  76. if extractClaudeCodeSessionIDFromPayload(body) != "" {
  77. cacheIdentityBody = body
  78. }
  79. cacheIdentity := resolveGrokCacheIdentity(c, cacheIdentityBody, "", upstreamModel)
  80. mixedCacheIntentBody := append([]byte(nil), patchedBody...)
  81. patchedBody, err = applyGrokResponsesCacheIdentity(patchedBody, body, cacheIdentity, account.IsGrokOAuth())
  82. if err != nil {
  83. return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
  84. }
  85. // Free OAuth + client function tools: reuse Messages mixed-tools cache route
  86. // (append web_search/x_search so xAI does not force non-cacheable build-free).
  87. patchedBody, err = applyGrokFreeRequestToolCacheRoute(c, patchedBody, mixedCacheIntentBody, account, cacheIdentity)
  88. if err != nil {
  89. return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err)
  90. }
  91. token, _, err := s.getRequestCredential(ctx, c, account)
  92. if err != nil {
  93. return nil, err
  94. }
  95. upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
  96. defer releaseUpstreamCtx()
  97. proxyURL := ""
  98. if account.ProxyID != nil && account.Proxy != nil {
  99. proxyURL = account.Proxy.URL()
  100. }
  101. upstreamStart := time.Now()
  102. var resp *http.Response
  103. for attempt := 0; ; attempt++ {
  104. upstreamReq, buildErr := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg, s.settingService)
  105. if buildErr != nil {
  106. return nil, buildErr
  107. }
  108. resp, err = s.doOpenAIUpstream(upstreamReq, proxyURL, account)
  109. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  110. if err != nil {
  111. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
  112. }
  113. // xAI can reject encrypted reasoning or a compaction blob copied from a
  114. // different decoder/cache context. Retry once on the same account after
  115. // preserving visible summaries and removing only opaque replay state.
  116. if attempt > 0 || (resp.StatusCode != http.StatusBadRequest && resp.StatusCode != http.StatusUnprocessableEntity) {
  117. break
  118. }
  119. respBody := s.readUpstreamErrorBody(resp)
  120. if resp.Body != nil {
  121. _ = resp.Body.Close()
  122. }
  123. invalidEncryptedContent := isGrokInvalidEncryptedContentResponse(resp.StatusCode, respBody)
  124. if !invalidEncryptedContent && !isGrokCompactionReplayDecodeError(resp.StatusCode, respBody) {
  125. resp.Body = io.NopCloser(bytes.NewReader(respBody))
  126. break
  127. }
  128. var retryBody []byte
  129. var changed bool
  130. var trimErr error
  131. if invalidEncryptedContent {
  132. retryBody, changed, trimErr = trimGrokInvalidEncryptedContentRetryBody(patchedBody)
  133. } else {
  134. retryBody, changed, trimErr = sanitizeGrokCompactionReplayBody(patchedBody)
  135. }
  136. if trimErr != nil {
  137. return nil, fmt.Errorf("prepare Grok replay decode retry: %w", trimErr)
  138. }
  139. if !changed {
  140. resp.Body = io.NopCloser(bytes.NewReader(respBody))
  141. break
  142. }
  143. patchedBody = retryBody
  144. slog.Info("grok_replay_decode_retry", "account_id", account.ID, "cache_identity_present", cacheIdentity != "")
  145. }
  146. defer func() { _ = resp.Body.Close() }()
  147. if resp.StatusCode >= 400 {
  148. respBody := s.readUpstreamErrorBody(resp)
  149. resp.Body = io.NopCloser(bytes.NewReader(respBody))
  150. upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody))
  151. if upstreamMsg == "" {
  152. upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode)
  153. }
  154. kind := "http_error"
  155. if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
  156. kind = "failover"
  157. }
  158. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  159. ProxyID: opsUpstreamProxyID(account),
  160. ProxyName: opsUpstreamProxyName(account),
  161. Platform: account.Platform,
  162. AccountID: account.ID,
  163. AccountName: account.Name,
  164. UpstreamStatusCode: resp.StatusCode,
  165. UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
  166. Kind: kind,
  167. Message: upstreamMsg,
  168. })
  169. errCtx := withGrokTeamRateLimitModel(ctx, upstreamModel)
  170. s.handleGrokAccountUpstreamError(errCtx, account, resp.StatusCode, resp.Header, respBody)
  171. // Quota/rate-limit responses stamp the team+model overlay. Capacity is
  172. // request pressure and must not hide sibling accounts.
  173. if shouldMarkGrokTeamModelRateLimit(resp.StatusCode, respBody) {
  174. markGrokTeamModelRateLimit(account, upstreamModel, resolveGrokTeamRateLimitUntil(time.Now().Add(grokTeamRateLimitDefaultTTL), time.Now()))
  175. }
  176. if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
  177. retryable, retryDelay, retryDeadline, retryMax := grokSameAccountRetryMetadata(account, resp.StatusCode, respBody)
  178. return nil, &UpstreamFailoverError{
  179. StatusCode: resp.StatusCode,
  180. ResponseBody: respBody,
  181. ResponseHeaders: resp.Header.Clone(),
  182. RetryableOnSameAccount: retryable,
  183. RequestScopedTransient: retryable && resp.StatusCode == http.StatusTooManyRequests,
  184. SameAccountRetryDelay: retryDelay,
  185. SameAccountRetryDeadline: retryDeadline,
  186. SameAccountRetryMax: retryMax,
  187. }
  188. }
  189. return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel)
  190. }
  191. // Attach model so rate-limit snapshots can fan out a team+model cool.
  192. stateCtx := withGrokTeamRateLimitModel(ctx, upstreamModel)
  193. s.updateGrokUsageFromResponse(stateCtx, account, resp.Header, resp.StatusCode)
  194. var usage *OpenAIUsage
  195. var firstTokenMs *int
  196. responseID := ""
  197. searchCount := 0
  198. imageCount := 0
  199. var imageOutputSizes []string
  200. if reqStream {
  201. maxLineSize := defaultMaxLineSize
  202. if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
  203. maxLineSize = s.cfg.Gateway.MaxLineSize
  204. }
  205. resp.Body = newGrokResponsesBillingPingFilterBody(resp.Body, account, maxLineSize)
  206. if hasGrokResponsesClientToolMapping(clientToolMapping) {
  207. resp.Body = newGrokResponsesClientToolStreamBody(resp.Body, clientToolMapping, maxLineSize)
  208. }
  209. streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel)
  210. if err != nil {
  211. return nil, err
  212. }
  213. usage = streamResult.usage
  214. firstTokenMs = streamResult.firstTokenMs
  215. responseID = strings.TrimSpace(streamResult.responseID)
  216. searchCount = streamResult.searchCount
  217. imageCount = streamResult.imageCount
  218. imageOutputSizes = streamResult.imageOutputSizes
  219. } else {
  220. nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel)
  221. if err != nil {
  222. return nil, err
  223. }
  224. usage = nonStreamResult.usage
  225. responseID = strings.TrimSpace(nonStreamResult.responseID)
  226. searchCount = nonStreamResult.searchCount
  227. imageCount = nonStreamResult.imageCount
  228. imageOutputSizes = nonStreamResult.imageOutputSizes
  229. }
  230. if usage == nil {
  231. usage = &OpenAIUsage{}
  232. }
  233. reasoningEffort := extractOpenAIReasoningEffortFromBody(patchedBody, originalModel)
  234. result := &OpenAIForwardResult{
  235. RequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
  236. UpstreamHeaders: resp.Header,
  237. ResponseID: responseID,
  238. Usage: *usage,
  239. Model: originalModel,
  240. UpstreamModel: upstreamModel,
  241. ReasoningEffort: reasoningEffort,
  242. Stream: reqStream,
  243. OpenAIWSMode: false,
  244. ResponseHeaders: resp.Header.Clone(),
  245. Duration: time.Since(startTime),
  246. FirstTokenMs: firstTokenMs,
  247. }
  248. // Propagate search/image counters from the shared Responses handler — without
  249. // this, stream/JSON counting runs but search_price_per_1k / image bills never apply.
  250. if searchCount > 0 {
  251. result.SearchCount = searchCount
  252. }
  253. if imageCount > 0 {
  254. result.ImageCount = imageCount
  255. result.ImageOutputSizes = imageOutputSizes
  256. }
  257. return result, nil
  258. }
  259. func isGrokInvalidEncryptedContentResponse(statusCode int, body []byte) bool {
  260. if statusCode != http.StatusBadRequest && statusCode != http.StatusUnprocessableEntity {
  261. return false
  262. }
  263. // xAI has used both flat and nested error envelopes:
  264. // {"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}
  265. // {"error":{"message":"Could not decrypt the provided encrypted_content."}}
  266. code := strings.TrimSpace(gjson.GetBytes(body, "code").String())
  267. errNode := gjson.GetBytes(body, "error")
  268. if code == "" && errNode.IsObject() {
  269. code = strings.TrimSpace(errNode.Get("code").String())
  270. }
  271. if strings.EqualFold(code, "invalid_encrypted_content") || strings.EqualFold(code, "invalid_compaction") || strings.EqualFold(code, "compaction_decode_error") {
  272. return true
  273. }
  274. // Keep the official xAI flat-code gate so unrelated 400s are not retried.
  275. if !strings.EqualFold(code, "invalid-argument") && code != "" {
  276. return false
  277. }
  278. for _, candidate := range grokStructuredErrorMessageCandidates(body) {
  279. normalizedMessage := strings.ToLower(candidate)
  280. // Nested OpenAI-style envelopes may omit top-level code; require decrypt text.
  281. if code == "" && !strings.Contains(normalizedMessage, "decrypt") && !strings.Contains(normalizedMessage, "decode the compaction blob") {
  282. continue
  283. }
  284. if strings.Contains(normalizedMessage, "encrypted_content") &&
  285. (strings.Contains(normalizedMessage, "decrypt") || strings.Contains(normalizedMessage, "unmodified")) {
  286. return true
  287. }
  288. if strings.Contains(normalizedMessage, "decode the compaction blob") {
  289. return true
  290. }
  291. }
  292. return false
  293. }
  294. func isGrokCompactionReplayDecodeError(statusCode int, body []byte) bool {
  295. if (statusCode != http.StatusBadRequest && statusCode != http.StatusUnprocessableEntity) || len(body) == 0 {
  296. return false
  297. }
  298. for _, candidate := range grokStructuredErrorMessageCandidates(body) {
  299. message := strings.ToLower(candidate)
  300. decodeSignal := strings.Contains(message, "decode") ||
  301. strings.Contains(message, "deserialize") ||
  302. strings.Contains(message, "decoder")
  303. replaySignal := strings.Contains(message, "compaction") ||
  304. strings.Contains(message, "summary") ||
  305. strings.Contains(message, "encrypted_content") ||
  306. strings.Contains(message, "response history")
  307. if decodeSignal && replaySignal {
  308. return true
  309. }
  310. }
  311. return false
  312. }
  313. func sanitizeGrokCompactionReplayBody(body []byte) ([]byte, bool, error) {
  314. converted, err := convertOpenAICompactInputsForGrok(body)
  315. if err != nil {
  316. return nil, false, fmt.Errorf("convert Grok compaction replay: %w", err)
  317. }
  318. var requestBody map[string]any
  319. decoder := json.NewDecoder(bytes.NewReader(converted))
  320. decoder.UseNumber()
  321. if err := decoder.Decode(&requestBody); err != nil {
  322. return nil, false, err
  323. }
  324. changed := !bytes.Equal(converted, body)
  325. if trimOpenAIEncryptedReasoningItems(requestBody) {
  326. changed = true
  327. }
  328. if dropEmptyGrokReplayReasoning(requestBody) {
  329. changed = true
  330. }
  331. if previousID, _ := requestBody["previous_response_id"].(string); strings.TrimSpace(previousID) != "" && !HasFunctionCallOutput(requestBody) {
  332. delete(requestBody, "previous_response_id")
  333. if _, exists := requestBody["store"]; !exists {
  334. requestBody["store"] = false
  335. }
  336. changed = true
  337. }
  338. if !changed {
  339. return body, false, nil
  340. }
  341. retryBody, err := marshalOpenAIUpstreamJSON(requestBody)
  342. if err != nil {
  343. return nil, false, err
  344. }
  345. return retryBody, true, nil
  346. }
  347. func dropEmptyGrokReplayReasoning(requestBody map[string]any) bool {
  348. items, ok := requestBody["input"].([]any)
  349. if !ok {
  350. return false
  351. }
  352. filtered := items[:0]
  353. changed := false
  354. for _, rawItem := range items {
  355. item, ok := rawItem.(map[string]any)
  356. if !ok || strings.TrimSpace(grokStringValue(item["type"])) != "reasoning" {
  357. filtered = append(filtered, rawItem)
  358. continue
  359. }
  360. summary, _ := item["summary"].([]any)
  361. content, hasContent := item["content"]
  362. _, hasEncrypted := item["encrypted_content"]
  363. if hasEncrypted || len(summary) > 0 || (hasContent && content != nil) {
  364. filtered = append(filtered, rawItem)
  365. continue
  366. }
  367. changed = true
  368. }
  369. if changed {
  370. requestBody["input"] = filtered
  371. }
  372. return changed
  373. }
  374. // requestHasGrokEncryptedReasoning reports whether the outbound Responses body
  375. // still carries reasoning.encrypted_content that can be stripped for retry.
  376. func requestHasGrokEncryptedReasoning(body []byte) bool {
  377. input := gjson.GetBytes(body, "input")
  378. if !input.Exists() {
  379. return false
  380. }
  381. items := input.Array()
  382. if input.IsObject() {
  383. items = []gjson.Result{input}
  384. }
  385. for _, item := range items {
  386. if strings.TrimSpace(item.Get("type").String()) != "reasoning" {
  387. continue
  388. }
  389. enc := item.Get("encrypted_content")
  390. if enc.Exists() && enc.Type != gjson.Null && strings.TrimSpace(enc.String()) != "" {
  391. return true
  392. }
  393. }
  394. return false
  395. }
  396. type grokEncryptedContentStripRetriedKey struct{}
  397. func markGrokEncryptedContentStripRetried(ctx context.Context) context.Context {
  398. return context.WithValue(ctx, grokEncryptedContentStripRetriedKey{}, true)
  399. }
  400. func grokEncryptedContentStripRetried(ctx context.Context) bool {
  401. v, _ := ctx.Value(grokEncryptedContentStripRetriedKey{}).(bool)
  402. return v
  403. }
  404. // stripAnthropicThinkingSignatures removes thinking.signature from Claude
  405. // history so a different Grok OAuth account can accept multi-turn tool
  406. // continuations after decrypt failures. Returns ok=false when nothing changed.
  407. func stripAnthropicThinkingSignatures(body []byte) ([]byte, bool) {
  408. if len(body) == 0 || !bytes.Contains(body, []byte(`"signature"`)) {
  409. return body, false
  410. }
  411. var req map[string]any
  412. if err := decodeOpenAIJSONUseNumber(body, &req); err != nil {
  413. return body, false
  414. }
  415. messages, ok := req["messages"].([]any)
  416. if !ok || len(messages) == 0 {
  417. return body, false
  418. }
  419. changed := false
  420. for _, rawMsg := range messages {
  421. msg, ok := rawMsg.(map[string]any)
  422. if !ok {
  423. continue
  424. }
  425. content, ok := msg["content"].([]any)
  426. if !ok {
  427. continue
  428. }
  429. for _, rawBlock := range content {
  430. block, ok := rawBlock.(map[string]any)
  431. if !ok {
  432. continue
  433. }
  434. if typ, _ := block["type"].(string); typ != "thinking" {
  435. continue
  436. }
  437. if _, has := block["signature"]; has {
  438. delete(block, "signature")
  439. changed = true
  440. }
  441. }
  442. }
  443. if !changed {
  444. return body, false
  445. }
  446. out, err := marshalOpenAIUpstreamJSON(req)
  447. if err != nil {
  448. return body, false
  449. }
  450. return out, true
  451. }
  452. func trimGrokInvalidEncryptedContentRetryBody(body []byte) ([]byte, bool, error) {
  453. input := gjson.GetBytes(body, "input")
  454. items := input.Array()
  455. if input.IsObject() {
  456. items = []gjson.Result{input}
  457. }
  458. hasEncryptedReasoning := false
  459. for _, item := range items {
  460. if (strings.TrimSpace(item.Get("type").String()) == "reasoning" && item.Get("encrypted_content").Exists()) ||
  461. (isOpenAICompactionType(strings.TrimSpace(item.Get("type").String())) && item.Get("encrypted_content").Exists()) {
  462. hasEncryptedReasoning = true
  463. break
  464. }
  465. }
  466. if !hasEncryptedReasoning {
  467. return body, false, nil
  468. }
  469. var requestBody map[string]any
  470. decoder := json.NewDecoder(bytes.NewReader(body))
  471. decoder.UseNumber()
  472. if err := decoder.Decode(&requestBody); err != nil {
  473. return nil, false, err
  474. }
  475. if !trimOpenAIEncryptedReasoningItems(requestBody) {
  476. return body, false, nil
  477. }
  478. retryBody, err := marshalOpenAIUpstreamJSON(requestBody)
  479. if err != nil {
  480. return nil, false, err
  481. }
  482. return retryBody, true, nil
  483. }
  484. func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) {
  485. return patchGrokResponsesBodyBase(body, upstreamModel)
  486. }
  487. func patchGrokResponsesBodyWithClientTools(body []byte, upstreamModel string) ([]byte, apicompat.ResponsesClientToolMapping, error) {
  488. if !json.Valid(body) {
  489. return nil, apicompat.ResponsesClientToolMapping{}, fmt.Errorf("invalid json request body")
  490. }
  491. promoted, err := sanitizeGrokResponsesInput(body)
  492. if err != nil {
  493. return nil, apicompat.ResponsesClientToolMapping{}, err
  494. }
  495. adapted, mapping, err := adaptGrokResponsesClientTools(promoted)
  496. if err != nil {
  497. return nil, apicompat.ResponsesClientToolMapping{}, err
  498. }
  499. patched, err := patchGrokResponsesBodyBase(adapted, upstreamModel)
  500. if err != nil {
  501. return nil, apicompat.ResponsesClientToolMapping{}, err
  502. }
  503. return patched, mapping, nil
  504. }
  505. func patchGrokResponsesBodyBase(body []byte, upstreamModel string) ([]byte, error) {
  506. if !json.Valid(body) {
  507. return nil, fmt.Errorf("invalid json request body")
  508. }
  509. // sjson may reuse the input backing array; keep the caller's request bytes
  510. // unchanged because the same body can be inspected for billing/retry paths.
  511. out, err := sjson.SetBytes(append([]byte(nil), body...), "model", upstreamModel)
  512. if err != nil {
  513. return nil, err
  514. }
  515. out, err = normalizeGrokResponsesReasoningEffort(out, upstreamModel)
  516. if err != nil {
  517. return nil, err
  518. }
  519. out, err = sanitizeGrokResponsesModelCapabilities(out, upstreamModel)
  520. if err != nil {
  521. return nil, err
  522. }
  523. for _, unsupportedField := range []string{"prompt_cache_retention", "safety_identifier", "metadata"} {
  524. if gjson.GetBytes(out, unsupportedField).Exists() {
  525. out, err = sjson.DeleteBytes(out, unsupportedField)
  526. if err != nil {
  527. return nil, err
  528. }
  529. }
  530. }
  531. if strings.EqualFold(upstreamModel, "grok-4.5") {
  532. for _, unsupportedField := range []string{"presence_penalty", "presencePenalty", "frequency_penalty", "frequencyPenalty", "stop"} {
  533. if gjson.GetBytes(out, unsupportedField).Exists() {
  534. out, err = sjson.DeleteBytes(out, unsupportedField)
  535. if err != nil {
  536. return nil, err
  537. }
  538. }
  539. }
  540. }
  541. if grokModelRejectsLogprobs(upstreamModel) {
  542. for _, unsupportedField := range []string{"logprobs", "top_logprobs"} {
  543. if gjson.GetBytes(out, unsupportedField).Exists() {
  544. out, err = sjson.DeleteBytes(out, unsupportedField)
  545. if err != nil {
  546. return nil, err
  547. }
  548. }
  549. }
  550. }
  551. out, err = sanitizeGrokResponsesUnsupportedFields(out)
  552. if err != nil {
  553. return nil, err
  554. }
  555. out, err = convertOpenAICompactInputsForGrok(out)
  556. if err != nil {
  557. return nil, err
  558. }
  559. out, err = sanitizeGrokResponsesInput(out)
  560. if err != nil {
  561. return nil, err
  562. }
  563. out, err = sanitizeGrokResponsesModelInput(out)
  564. if err != nil {
  565. return nil, err
  566. }
  567. out, err = stripRedundantGrokViewImageTool(out)
  568. if err != nil {
  569. return nil, err
  570. }
  571. out, err = sanitizeGrokReasoningNullContent(out)
  572. if err != nil {
  573. return nil, err
  574. }
  575. out, err = sanitizeGrokResponsesTools(out)
  576. if err != nil {
  577. return nil, err
  578. }
  579. return out, nil
  580. }
  581. // xAI's Grok 4.20 family and newer models do not support OpenAI's logprobs
  582. // fields. Remove them before egress instead of forwarding a request the
  583. // upstream rejects. Older Grok models retain the fields for compatibility.
  584. func grokModelRejectsLogprobs(model string) bool {
  585. model = strings.ToLower(strings.TrimSpace(model))
  586. if slash := strings.LastIndex(model, "/"); slash >= 0 {
  587. model = strings.TrimSpace(model[slash+1:])
  588. }
  589. return strings.HasPrefix(model, "grok-4.20")
  590. }
  591. func sanitizeGrokResponsesModelCapabilities(body []byte, upstreamModel string) ([]byte, error) {
  592. if !grokModelRejectsReasoningEffort(upstreamModel) {
  593. return body, nil
  594. }
  595. out := body
  596. for _, field := range []string{"reasoning", "reasoning_effort", "reasoningEffort"} {
  597. if !gjson.GetBytes(out, field).Exists() {
  598. continue
  599. }
  600. var err error
  601. out, err = sjson.DeleteBytes(out, field)
  602. if err != nil {
  603. return nil, fmt.Errorf("remove unsupported Grok Composer %s: %w", field, err)
  604. }
  605. }
  606. return out, nil
  607. }
  608. func grokModelRejectsReasoningEffort(model string) bool {
  609. model = strings.TrimSpace(strings.ToLower(model))
  610. if slash := strings.LastIndex(model, "/"); slash >= 0 {
  611. model = strings.TrimSpace(model[slash+1:])
  612. }
  613. switch model {
  614. case "grok-composer", "grok-composer-2.5-fast", "composer-2.5":
  615. return true
  616. default:
  617. return false
  618. }
  619. }
  620. func normalizeGrokResponsesReasoningEffort(body []byte, upstreamModel string) ([]byte, error) {
  621. supportsEffort := grokSupportsReasoningEffort(upstreamModel)
  622. out := body
  623. var err error
  624. for _, field := range []string{"reasoning.effort", "reasoning_effort"} {
  625. value := gjson.GetBytes(out, field)
  626. if !value.Exists() {
  627. continue
  628. }
  629. normalized, keep := normalizeGrokReasoningEffortValue(value.String(), upstreamModel)
  630. if !supportsEffort || !keep {
  631. out, err = sjson.DeleteBytes(out, field)
  632. } else {
  633. out, err = sjson.SetBytes(out, field, normalized)
  634. }
  635. if err != nil {
  636. return nil, fmt.Errorf("normalize Grok reasoning field %s: %w", field, err)
  637. }
  638. }
  639. if camel := gjson.GetBytes(out, "reasoningEffort"); camel.Exists() {
  640. normalized, keep := normalizeGrokReasoningEffortValue(camel.String(), upstreamModel)
  641. out, err = sjson.DeleteBytes(out, "reasoningEffort")
  642. if err != nil {
  643. return nil, fmt.Errorf("remove Grok reasoningEffort: %w", err)
  644. }
  645. if supportsEffort && keep && !gjson.GetBytes(out, "reasoning_effort").Exists() {
  646. out, err = sjson.SetBytes(out, "reasoning_effort", normalized)
  647. if err != nil {
  648. return nil, fmt.Errorf("set Grok reasoning_effort: %w", err)
  649. }
  650. }
  651. }
  652. if reasoning := gjson.GetBytes(out, "reasoning"); reasoning.Exists() && reasoning.IsObject() && len(reasoning.Map()) == 0 {
  653. out, err = sjson.DeleteBytes(out, "reasoning")
  654. if err != nil {
  655. return nil, fmt.Errorf("remove empty Grok reasoning: %w", err)
  656. }
  657. }
  658. return out, nil
  659. }
  660. func normalizeGrokChatReasoningEffort(body []byte, upstreamModel string) ([]byte, error) {
  661. raw := strings.TrimSpace(gjson.GetBytes(body, "reasoning_effort").String())
  662. if raw == "" {
  663. raw = strings.TrimSpace(gjson.GetBytes(body, "reasoningEffort").String())
  664. }
  665. normalized, keep := normalizeGrokReasoningEffortValue(raw, upstreamModel)
  666. keep = keep && grokSupportsReasoningEffort(upstreamModel)
  667. out := body
  668. var err error
  669. if gjson.GetBytes(out, "reasoningEffort").Exists() {
  670. out, err = sjson.DeleteBytes(out, "reasoningEffort")
  671. if err != nil {
  672. return nil, err
  673. }
  674. }
  675. if !keep {
  676. if gjson.GetBytes(out, "reasoning_effort").Exists() {
  677. out, err = sjson.DeleteBytes(out, "reasoning_effort")
  678. }
  679. return out, err
  680. }
  681. out, err = sjson.SetBytes(out, "reasoning_effort", normalized)
  682. return out, err
  683. }
  684. func normalizeGrokReasoningEffortValue(raw, model string) (string, bool) {
  685. value := strings.NewReplacer("-", "", "_", "", " ", "").Replace(strings.ToLower(strings.TrimSpace(raw)))
  686. switch value {
  687. case "none", "low", "medium", "high":
  688. return value, true
  689. case "minimal":
  690. return "low", true
  691. case "xhigh", "extrahigh":
  692. if GrokSupportsXHighReasoningEffort(model) {
  693. return "xhigh", true
  694. }
  695. return "high", true
  696. case "max", "ultra":
  697. return "high", true
  698. default:
  699. return "", false
  700. }
  701. }
  702. // GrokSupportsXHighReasoningEffort reports whether the model advertises and
  703. // forwards the xhigh reasoning effort (Grok 4.6 and its undated alias).
  704. func GrokSupportsXHighReasoningEffort(model string) bool {
  705. model = strings.ToLower(xai.StripGrokProviderPrefix(strings.TrimSpace(model)))
  706. return model == "grok-4.6" || model == "grok-4.6-latest"
  707. }
  708. func grokSupportsReasoningEffort(model string) bool {
  709. model = strings.ToLower(xai.StripGrokProviderPrefix(strings.TrimSpace(model)))
  710. switch model {
  711. case "grok-4.5", "grok-4.5-latest", "grok-4.6", "grok-4.6-latest",
  712. "grok-4.3", "grok-4.3-latest",
  713. "grok-3-mini", "grok-3-mini-fast", "grok-4.20-0309-reasoning",
  714. "grok-4.20-reasoning", "grok-4.20-multi-agent-0309":
  715. return true
  716. default:
  717. return false
  718. }
  719. }
  720. // grokUnsupportedRecursiveFields 定义 Grok 平台(Responses 和 Chat Completions)不支持的字段
  721. var grokUnsupportedRecursiveFields = map[string]struct{}{
  722. "external_web_access": {},
  723. }
  724. // sanitizeGrokUnsupportedFields 递归移除 Grok 平台不支持的字段
  725. // 适用于 Responses API 和 Chat Completions API
  726. func sanitizeGrokUnsupportedFields(body []byte) ([]byte, error) {
  727. if !bytes.Contains(body, []byte(`"external_web_access"`)) {
  728. return body, nil
  729. }
  730. var payload any
  731. if err := decodeOpenAIJSONUseNumber(body, &payload); err != nil {
  732. return nil, err
  733. }
  734. if !deleteJSONFields(payload, grokUnsupportedRecursiveFields) {
  735. return body, nil
  736. }
  737. return marshalOpenAIUpstreamJSON(payload)
  738. }
  739. // sanitizeGrokResponsesUnsupportedFields 保留旧函数名作为别名,向后兼容
  740. var sanitizeGrokResponsesUnsupportedFields = sanitizeGrokUnsupportedFields
  741. func deleteJSONFields(value any, fields map[string]struct{}) bool {
  742. switch typed := value.(type) {
  743. case map[string]any:
  744. changed := false
  745. for field := range fields {
  746. if _, ok := typed[field]; ok {
  747. delete(typed, field)
  748. changed = true
  749. }
  750. }
  751. for _, child := range typed {
  752. if deleteJSONFields(child, fields) {
  753. changed = true
  754. }
  755. }
  756. return changed
  757. case []any:
  758. changed := false
  759. for _, child := range typed {
  760. if deleteJSONFields(child, fields) {
  761. changed = true
  762. }
  763. }
  764. return changed
  765. default:
  766. return false
  767. }
  768. }
  769. // additional_tools is a Codex/Responses Lite private input carrier. xAI's
  770. // Responses schema rejects the carrier itself, but accepts supported tools at
  771. // the top level. Preserve top-level order, append newly discovered tools in
  772. // carrier order, then let sanitizeGrokResponsesTools filter unsupported types.
  773. func sanitizeGrokResponsesInput(body []byte) ([]byte, error) {
  774. if !bytes.Contains(body, []byte(`"additional_tools"`)) {
  775. return body, nil
  776. }
  777. input := gjson.GetBytes(body, "input")
  778. if !input.Exists() || !input.IsArray() {
  779. return body, nil
  780. }
  781. rawItems := input.Array()
  782. filtered := make([]json.RawMessage, 0, len(rawItems))
  783. topLevelTools := gjson.GetBytes(body, "tools")
  784. mergedTools := make([]json.RawMessage, 0)
  785. seenTools := make(map[string]struct{})
  786. appendTool := func(tool gjson.Result) bool {
  787. key := grokResponsesToolDedupKey(tool)
  788. if _, exists := seenTools[key]; exists {
  789. return false
  790. }
  791. seenTools[key] = struct{}{}
  792. mergedTools = append(mergedTools, json.RawMessage(tool.Raw))
  793. return true
  794. }
  795. if topLevelTools.IsArray() {
  796. for _, tool := range topLevelTools.Array() {
  797. seenTools[grokResponsesToolDedupKey(tool)] = struct{}{}
  798. mergedTools = append(mergedTools, json.RawMessage(tool.Raw))
  799. }
  800. }
  801. promoted := false
  802. for _, item := range rawItems {
  803. if strings.TrimSpace(item.Get("type").String()) == "additional_tools" {
  804. tools := item.Get("tools")
  805. if tools.IsArray() {
  806. for _, tool := range tools.Array() {
  807. if appendTool(tool) {
  808. promoted = true
  809. }
  810. }
  811. }
  812. continue
  813. }
  814. filtered = append(filtered, json.RawMessage(item.Raw))
  815. }
  816. if len(filtered) == len(rawItems) {
  817. return body, nil
  818. }
  819. encoded, err := json.Marshal(filtered)
  820. if err != nil {
  821. return nil, err
  822. }
  823. body, err = sjson.SetRawBytes(body, "input", encoded)
  824. if err != nil || !promoted {
  825. return body, err
  826. }
  827. encodedTools, err := json.Marshal(mergedTools)
  828. if err != nil {
  829. return nil, err
  830. }
  831. return sjson.SetRawBytes(body, "tools", encodedTools)
  832. }
  833. // An inline input_image is already visible to Grok. Keeping Codex's local
  834. // view_image tool in the same turn can make Grok announce a tool call without
  835. // actually calling it, so remove only that redundant automatic choice.
  836. func stripRedundantGrokViewImageTool(body []byte) ([]byte, error) {
  837. input := gjson.GetBytes(body, "input")
  838. if !input.IsArray() {
  839. return body, nil
  840. }
  841. items := input.Array()
  842. if len(items) == 0 {
  843. return body, nil
  844. }
  845. current := items[len(items)-1]
  846. if strings.TrimSpace(current.Get("role").String()) != "user" ||
  847. !openAIJSONValueMayContainImageInput(current) {
  848. return body, nil
  849. }
  850. toolChoice := gjson.GetBytes(body, "tool_choice")
  851. if toolChoice.IsObject() && strings.TrimSpace(toolChoice.Get("type").String()) == "function" {
  852. choiceName := strings.TrimSpace(toolChoice.Get("name").String())
  853. if choiceName == "" {
  854. choiceName = strings.TrimSpace(toolChoice.Get("function.name").String())
  855. }
  856. if choiceName == "view_image" {
  857. return body, nil
  858. }
  859. }
  860. tools := gjson.GetBytes(body, "tools")
  861. if !tools.IsArray() {
  862. return body, nil
  863. }
  864. filtered := make([]json.RawMessage, 0, len(tools.Array()))
  865. changed := false
  866. for _, tool := range tools.Array() {
  867. if strings.TrimSpace(tool.Get("type").String()) == "function" &&
  868. strings.TrimSpace(tool.Get("name").String()) == "view_image" {
  869. changed = true
  870. continue
  871. }
  872. filtered = append(filtered, json.RawMessage(tool.Raw))
  873. }
  874. if !changed {
  875. return body, nil
  876. }
  877. if len(filtered) == 0 && strings.TrimSpace(toolChoice.String()) == "required" {
  878. return body, nil
  879. }
  880. if len(filtered) == 0 {
  881. out, err := sjson.DeleteBytes(body, "tools")
  882. if err != nil {
  883. return nil, err
  884. }
  885. return sjson.DeleteBytes(out, "parallel_tool_calls")
  886. }
  887. encoded, err := json.Marshal(filtered)
  888. if err != nil {
  889. return nil, err
  890. }
  891. return sjson.SetRawBytes(body, "tools", encoded)
  892. }
  893. func grokResponsesToolDedupKey(tool gjson.Result) string {
  894. toolType := strings.TrimSpace(tool.Get("type").String())
  895. if toolType != "" {
  896. if name := strings.TrimSpace(tool.Get("name").String()); name != "" {
  897. return "type:" + toolType + "\x00name:" + name
  898. }
  899. if toolType == "mcp" {
  900. if label := strings.TrimSpace(tool.Get("server_label").String()); label != "" {
  901. return "type:mcp\x00server_label:" + label
  902. }
  903. }
  904. }
  905. return "json:" + normalizeCompatSeedJSON(json.RawMessage(tool.Raw))
  906. }
  907. // sanitizeGrokReasoningNullContent drops explicit JSON nulls from Responses
  908. // input items. xAI's untagged ModelInput decoder 422s on those fields.
  909. // Compaction items stay unmodified per the compact contract.
  910. func sanitizeGrokReasoningNullContent(body []byte) ([]byte, error) {
  911. input := gjson.GetBytes(body, "input")
  912. if !input.Exists() || (!input.IsArray() && !input.IsObject()) {
  913. return body, nil
  914. }
  915. var decoded map[string]any
  916. decoder := json.NewDecoder(bytes.NewReader(body))
  917. decoder.UseNumber()
  918. if err := decoder.Decode(&decoded); err != nil {
  919. return body, nil
  920. }
  921. rawInput, ok := decoded["input"]
  922. if !ok {
  923. return body, nil
  924. }
  925. cleaned, changed := stripExplicitNullsFromGrokInput(rawInput)
  926. if !changed {
  927. return body, nil
  928. }
  929. decoded["input"] = cleaned
  930. out, err := marshalOpenAIUpstreamJSON(decoded)
  931. if err != nil {
  932. return nil, err
  933. }
  934. return out, nil
  935. }
  936. func stripExplicitNullsFromGrokInput(value any) (any, bool) {
  937. switch node := value.(type) {
  938. case []any:
  939. changed := false
  940. for i, item := range node {
  941. itemMap, ok := item.(map[string]any)
  942. if !ok {
  943. next, childChanged := stripExplicitNullsFromGrokInput(item)
  944. if childChanged {
  945. node[i] = next
  946. changed = true
  947. }
  948. continue
  949. }
  950. if isOpenAICompactionType(stringValue(itemMap["type"])) {
  951. continue
  952. }
  953. next, childChanged := stripExplicitNullsFromJSONObject(itemMap)
  954. if childChanged {
  955. node[i] = next
  956. changed = true
  957. }
  958. }
  959. return node, changed
  960. case map[string]any:
  961. if isOpenAICompactionType(stringValue(node["type"])) {
  962. return node, false
  963. }
  964. return stripExplicitNullsFromJSONObject(node)
  965. default:
  966. return value, false
  967. }
  968. }
  969. func stripExplicitNullsFromJSONObject(node map[string]any) (map[string]any, bool) {
  970. if node == nil {
  971. return node, false
  972. }
  973. changed := false
  974. for key, child := range node {
  975. if child == nil {
  976. delete(node, key)
  977. changed = true
  978. continue
  979. }
  980. switch typed := child.(type) {
  981. case map[string]any:
  982. next, childChanged := stripExplicitNullsFromJSONObject(typed)
  983. if childChanged {
  984. node[key] = next
  985. changed = true
  986. }
  987. case []any:
  988. next, childChanged := stripExplicitNullsFromGrokInput(typed)
  989. if childChanged {
  990. node[key] = next
  991. changed = true
  992. }
  993. }
  994. }
  995. return node, changed
  996. }
  997. var grokResponsesSupportedToolTypes = map[string]struct{}{
  998. "code_execution": {},
  999. "code_interpreter": {},
  1000. "collections_search": {},
  1001. "file_search": {},
  1002. "function": {},
  1003. "mcp": {},
  1004. "shell": {},
  1005. "web_search": {},
  1006. "x_search": {},
  1007. }
  1008. const grokSafeFunctionParameters = `{"type":"object","properties":{},"additionalProperties":true}`
  1009. func sanitizeGrokResponsesTools(body []byte) ([]byte, error) {
  1010. tools := gjson.GetBytes(body, "tools")
  1011. if !tools.Exists() {
  1012. return deleteGrokOrphanToolControls(body)
  1013. }
  1014. if !tools.IsArray() {
  1015. // xAI rejects tool_choice when tools is null/object. Drop the malformed
  1016. // collection and any orphan tool controls instead of forwarding a pair
  1017. // the Grok Responses endpoint cannot interpret.
  1018. body, err := sjson.DeleteBytes(body, "tools")
  1019. if err != nil {
  1020. return nil, err
  1021. }
  1022. return deleteGrokOrphanToolControls(body)
  1023. }
  1024. rawTools := tools.Array()
  1025. filteredTools := make([]json.RawMessage, 0, len(rawTools))
  1026. toolsChanged := false
  1027. for _, tool := range rawTools {
  1028. toolType := strings.TrimSpace(tool.Get("type").String())
  1029. if _, ok := grokResponsesSupportedToolTypes[toolType]; ok {
  1030. raw := json.RawMessage(tool.Raw)
  1031. if toolType == "function" && (!tool.Get("parameters").Exists() || tool.Get("parameters").Type == gjson.Null) {
  1032. var payload map[string]any
  1033. if err := decodeOpenAIJSONUseNumber(raw, &payload); err != nil {
  1034. return nil, err
  1035. }
  1036. payload["parameters"] = map[string]any{"type": "object", "properties": map[string]any{}}
  1037. encoded, err := marshalOpenAIUpstreamJSON(payload)
  1038. if err != nil {
  1039. return nil, err
  1040. }
  1041. raw = encoded
  1042. toolsChanged = true
  1043. } else if toolType == "function" && grokFunctionParametersHaveInvalidUnionRoot(tool.Get("parameters")) {
  1044. var err error
  1045. raw, err = sjson.SetRawBytes(raw, "parameters", []byte(grokSafeFunctionParameters))
  1046. if err != nil {
  1047. return nil, err
  1048. }
  1049. if strict := tool.Get("strict"); strict.Exists() && strict.Bool() {
  1050. raw, err = sjson.SetBytes(raw, "strict", false)
  1051. if err != nil {
  1052. return nil, err
  1053. }
  1054. }
  1055. toolsChanged = true
  1056. }
  1057. filteredTools = append(filteredTools, raw)
  1058. }
  1059. }
  1060. if !grokRawToolsContainType(filteredTools, "tool_search") {
  1061. for index, raw := range filteredTools {
  1062. if !gjson.GetBytes(raw, "defer_loading").Exists() {
  1063. continue
  1064. }
  1065. cleaned, deleteErr := sjson.DeleteBytes(raw, "defer_loading")
  1066. if deleteErr != nil {
  1067. return nil, deleteErr
  1068. }
  1069. filteredTools[index] = cleaned
  1070. toolsChanged = true
  1071. }
  1072. }
  1073. var err error
  1074. if len(filteredTools) != len(rawTools) || toolsChanged {
  1075. if len(filteredTools) == 0 {
  1076. body, err = sjson.DeleteBytes(body, "tools")
  1077. } else {
  1078. var encoded []byte
  1079. encoded, err = json.Marshal(filteredTools)
  1080. if err != nil {
  1081. return nil, err
  1082. }
  1083. body, err = sjson.SetRawBytes(body, "tools", encoded)
  1084. }
  1085. if err != nil {
  1086. return nil, err
  1087. }
  1088. }
  1089. if len(filteredTools) == 0 {
  1090. return deleteGrokOrphanToolControls(body)
  1091. }
  1092. toolChoice := gjson.GetBytes(body, "tool_choice")
  1093. if !toolChoice.Exists() {
  1094. return body, nil
  1095. }
  1096. if shouldDropGrokToolChoice(toolChoice, filteredTools) {
  1097. body, err = sjson.DeleteBytes(body, "tool_choice")
  1098. if err != nil {
  1099. return nil, err
  1100. }
  1101. }
  1102. return body, nil
  1103. }
  1104. func grokFunctionParametersHaveInvalidUnionRoot(parameters gjson.Result) bool {
  1105. if !parameters.Exists() || !parameters.IsObject() {
  1106. return false
  1107. }
  1108. for _, keyword := range []string{"anyOf", "oneOf"} {
  1109. branches := parameters.Get(keyword)
  1110. if !branches.IsArray() {
  1111. continue
  1112. }
  1113. values := branches.Array()
  1114. if len(values) == 0 {
  1115. continue
  1116. }
  1117. for _, branch := range values {
  1118. if !strings.EqualFold(strings.TrimSpace(branch.Get("type").String()), "object") {
  1119. return true
  1120. }
  1121. }
  1122. }
  1123. return false
  1124. }
  1125. func grokRawToolsContainType(tools []json.RawMessage, want string) bool {
  1126. for _, tool := range tools {
  1127. if strings.TrimSpace(gjson.GetBytes(tool, "type").String()) == want {
  1128. return true
  1129. }
  1130. }
  1131. return false
  1132. }
  1133. func deleteGrokOrphanToolControls(body []byte) ([]byte, error) {
  1134. var err error
  1135. for _, field := range []string{"tool_choice", "parallel_tool_calls"} {
  1136. if !gjson.GetBytes(body, field).Exists() {
  1137. continue
  1138. }
  1139. body, err = sjson.DeleteBytes(body, field)
  1140. if err != nil {
  1141. return nil, err
  1142. }
  1143. }
  1144. return body, nil
  1145. }
  1146. func shouldDropGrokToolChoice(toolChoice gjson.Result, tools []json.RawMessage) bool {
  1147. if len(tools) == 0 {
  1148. return true
  1149. }
  1150. if !toolChoice.IsObject() {
  1151. return false
  1152. }
  1153. choiceType := strings.TrimSpace(toolChoice.Get("type").String())
  1154. if choiceType == "" {
  1155. return false
  1156. }
  1157. if _, ok := grokResponsesSupportedToolTypes[choiceType]; !ok {
  1158. return true
  1159. }
  1160. if choiceType == "function" {
  1161. choiceName := strings.TrimSpace(toolChoice.Get("name").String())
  1162. if choiceName == "" {
  1163. choiceName = strings.TrimSpace(toolChoice.Get("function.name").String())
  1164. }
  1165. if choiceName == "" {
  1166. return false
  1167. }
  1168. for _, tool := range tools {
  1169. var item struct {
  1170. Type string `json:"type"`
  1171. Name string `json:"name"`
  1172. Function struct {
  1173. Name string `json:"name"`
  1174. } `json:"function"`
  1175. }
  1176. if err := json.Unmarshal(tool, &item); err != nil {
  1177. continue
  1178. }
  1179. name := strings.TrimSpace(item.Name)
  1180. if name == "" {
  1181. name = strings.TrimSpace(item.Function.Name)
  1182. }
  1183. if strings.TrimSpace(item.Type) == "function" && name == choiceName {
  1184. return false
  1185. }
  1186. }
  1187. return true
  1188. }
  1189. return false
  1190. }
  1191. func (s *OpenAIGatewayService) bridgeGrokComposerImageInputs(
  1192. ctx context.Context,
  1193. c *gin.Context,
  1194. account *Account,
  1195. body []byte,
  1196. token string,
  1197. ) ([]byte, OpenAIUsage, bool, error) {
  1198. if !shouldBridgeGrokComposerImageInputs(body) {
  1199. return body, OpenAIUsage{}, false, nil
  1200. }
  1201. var reqBody map[string]any
  1202. if err := decodeOpenAIJSONUseNumber(body, &reqBody); err != nil {
  1203. return body, OpenAIUsage{}, false, fmt.Errorf("parse grok composer image bridge request: %w", err)
  1204. }
  1205. imageURLs := collectGrokComposerImageURLs(reqBody)
  1206. if len(imageURLs) == 0 {
  1207. return body, OpenAIUsage{}, false, nil
  1208. }
  1209. descriptions := make([]string, 0, len(imageURLs))
  1210. var bridgeUsage OpenAIUsage
  1211. for index, imageURL := range imageURLs {
  1212. description, usage, err := s.describeGrokComposerImage(ctx, c, account, token, imageURL, index+1)
  1213. if err != nil {
  1214. return body, bridgeUsage, false, err
  1215. }
  1216. descriptions = append(descriptions, description)
  1217. addOpenAIUsage(&bridgeUsage, usage)
  1218. }
  1219. if !rewriteGrokComposerImagesAsText(reqBody, descriptions) {
  1220. return body, bridgeUsage, false, nil
  1221. }
  1222. bridgedBody, err := marshalOpenAIUpstreamJSON(reqBody)
  1223. if err != nil {
  1224. return body, bridgeUsage, false, fmt.Errorf("serialize grok composer image bridge request: %w", err)
  1225. }
  1226. return bridgedBody, bridgeUsage, true, nil
  1227. }
  1228. func shouldBridgeGrokComposerImageInputs(body []byte) bool {
  1229. if len(body) == 0 || !isGrokComposerModel(gjson.GetBytes(body, "model").String()) {
  1230. return false
  1231. }
  1232. messages := gjson.GetBytes(body, "messages")
  1233. if !messages.Exists() {
  1234. return false
  1235. }
  1236. return openAIJSONValueMayContainImageInput(messages)
  1237. }
  1238. func isGrokComposerModel(model string) bool {
  1239. model = strings.TrimSpace(strings.ToLower(model))
  1240. if model == "" {
  1241. return false
  1242. }
  1243. if strings.Contains(model, "/") {
  1244. parts := strings.Split(model, "/")
  1245. model = strings.TrimSpace(parts[len(parts)-1])
  1246. }
  1247. return strings.Contains(model, "composer")
  1248. }
  1249. func collectGrokComposerImageURLs(reqBody map[string]any) []string {
  1250. messages, ok := reqBody["messages"].([]any)
  1251. if !ok {
  1252. return nil
  1253. }
  1254. var imageURLs []string
  1255. for _, msg := range messages {
  1256. msgMap, ok := msg.(map[string]any)
  1257. if !ok {
  1258. continue
  1259. }
  1260. parts, ok := msgMap["content"].([]any)
  1261. if !ok {
  1262. continue
  1263. }
  1264. for _, part := range parts {
  1265. if imageURL := grokComposerImageURLFromPart(part); imageURL != "" {
  1266. imageURLs = append(imageURLs, imageURL)
  1267. }
  1268. }
  1269. }
  1270. return imageURLs
  1271. }
  1272. func grokComposerImageURLFromPart(part any) string {
  1273. partMap, ok := part.(map[string]any)
  1274. if !ok {
  1275. return ""
  1276. }
  1277. if strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"]))) != "image_url" {
  1278. return ""
  1279. }
  1280. switch imageURL := partMap["image_url"].(type) {
  1281. case string:
  1282. return normalizeGrokComposerImageURL(imageURL)
  1283. case map[string]any:
  1284. raw, _ := imageURL["url"].(string)
  1285. return normalizeGrokComposerImageURL(raw)
  1286. default:
  1287. return ""
  1288. }
  1289. }
  1290. func normalizeGrokComposerImageURL(raw string) string {
  1291. trimmed := strings.TrimSpace(raw)
  1292. if trimmed == "" || isEmptyBase64DataURI(trimmed) {
  1293. return ""
  1294. }
  1295. return trimmed
  1296. }
  1297. func (s *OpenAIGatewayService) describeGrokComposerImage(
  1298. ctx context.Context,
  1299. c *gin.Context,
  1300. account *Account,
  1301. token string,
  1302. imageURL string,
  1303. index int,
  1304. ) (string, OpenAIUsage, error) {
  1305. body, err := buildGrokComposerImageDescriptionBody(imageURL, index)
  1306. if err != nil {
  1307. return "", OpenAIUsage{}, err
  1308. }
  1309. upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
  1310. // Image-description probes are auxiliary requests, not conversation turns.
  1311. // Do not bind them to the caller's Grok prompt-cache identity.
  1312. upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "", s.cfg, s.settingService)
  1313. releaseUpstreamCtx()
  1314. if err != nil {
  1315. return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err)
  1316. }
  1317. proxyURL := ""
  1318. if account.ProxyID != nil && account.Proxy != nil {
  1319. proxyURL = account.Proxy.URL()
  1320. }
  1321. resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
  1322. if err != nil {
  1323. return "", OpenAIUsage{}, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
  1324. }
  1325. defer func() { _ = resp.Body.Close() }()
  1326. if resp.StatusCode >= 400 {
  1327. respBody := s.readUpstreamErrorBody(resp)
  1328. upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody))
  1329. if upstreamMsg == "" {
  1330. upstreamMsg = fmt.Sprintf("xAI image bridge upstream returned status %d", resp.StatusCode)
  1331. }
  1332. kind := "http_error"
  1333. if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
  1334. kind = "failover"
  1335. }
  1336. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  1337. ProxyID: opsUpstreamProxyID(account),
  1338. ProxyName: opsUpstreamProxyName(account),
  1339. Platform: account.Platform,
  1340. AccountID: account.ID,
  1341. AccountName: account.Name,
  1342. UpstreamStatusCode: resp.StatusCode,
  1343. UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
  1344. Kind: kind,
  1345. Message: upstreamMsg,
  1346. })
  1347. s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, grokComposerImageBridgeVisionModel), account, resp.StatusCode, resp.Header, respBody)
  1348. if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
  1349. retryable, retryDelay, retryDeadline, retryMax := grokSameAccountRetryMetadata(account, resp.StatusCode, respBody)
  1350. return "", OpenAIUsage{}, &UpstreamFailoverError{
  1351. StatusCode: resp.StatusCode,
  1352. ResponseBody: respBody,
  1353. ResponseHeaders: resp.Header.Clone(),
  1354. RetryableOnSameAccount: retryable,
  1355. RequestScopedTransient: retryable && resp.StatusCode == http.StatusTooManyRequests,
  1356. SameAccountRetryDelay: retryDelay,
  1357. SameAccountRetryDeadline: retryDeadline,
  1358. SameAccountRetryMax: retryMax,
  1359. }
  1360. }
  1361. return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg)
  1362. }
  1363. s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, grokComposerImageBridgeVisionModel), account, resp.Header, resp.StatusCode)
  1364. respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil)
  1365. if err != nil {
  1366. return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err)
  1367. }
  1368. var parsed apicompat.ResponsesResponse
  1369. if err := json.Unmarshal(respBody, &parsed); err != nil {
  1370. return "", OpenAIUsage{}, fmt.Errorf("parse grok composer image bridge response: %w", err)
  1371. }
  1372. description := strings.TrimSpace(grokResponsesOutputText(&parsed))
  1373. if description == "" {
  1374. return "", copyOpenAIUsageFromResponsesUsage(parsed.Usage), fmt.Errorf("grok composer image bridge returned empty description")
  1375. }
  1376. return description, copyOpenAIUsageFromResponsesUsage(parsed.Usage), nil
  1377. }
  1378. func buildGrokComposerImageDescriptionBody(imageURL string, index int) ([]byte, error) {
  1379. prompt := fmt.Sprintf("Describe image %d in concise, factual text for a downstream coding/composer model. Include visible text, UI elements, diagrams, errors, and spatial relationships. Do not mention that you are an image analysis bridge.", index)
  1380. req := map[string]any{
  1381. "model": grokComposerImageBridgeVisionModel,
  1382. "stream": false,
  1383. "store": false,
  1384. "max_output_tokens": grokComposerImageBridgeMaxOutputTokens,
  1385. "input": []any{
  1386. map[string]any{
  1387. "type": "message",
  1388. "role": "user",
  1389. "content": []any{
  1390. map[string]any{"type": "input_text", "text": prompt},
  1391. map[string]any{"type": "input_image", "image_url": imageURL},
  1392. },
  1393. },
  1394. },
  1395. }
  1396. return marshalOpenAIUpstreamJSON(req)
  1397. }
  1398. func grokResponsesOutputText(resp *apicompat.ResponsesResponse) string {
  1399. if resp == nil {
  1400. return ""
  1401. }
  1402. var parts []string
  1403. for _, output := range resp.Output {
  1404. for _, content := range output.Content {
  1405. if content.Type == "output_text" || content.Type == "text" || content.Type == "input_text" {
  1406. if text := strings.TrimSpace(content.Text); text != "" {
  1407. parts = append(parts, text)
  1408. }
  1409. }
  1410. }
  1411. }
  1412. return strings.Join(parts, "\n\n")
  1413. }
  1414. func rewriteGrokComposerImagesAsText(reqBody map[string]any, descriptions []string) bool {
  1415. messages, ok := reqBody["messages"].([]any)
  1416. if !ok {
  1417. return false
  1418. }
  1419. imageIndex := 0
  1420. changed := false
  1421. for _, msg := range messages {
  1422. msgMap, ok := msg.(map[string]any)
  1423. if !ok {
  1424. continue
  1425. }
  1426. parts, ok := msgMap["content"].([]any)
  1427. if !ok {
  1428. continue
  1429. }
  1430. var textParts []string
  1431. messageChanged := false
  1432. for _, part := range parts {
  1433. if imageURL := grokComposerImageURLFromPart(part); imageURL != "" {
  1434. if imageIndex < len(descriptions) {
  1435. textParts = append(textParts, fmt.Sprintf("Image %d description: %s", imageIndex+1, strings.TrimSpace(descriptions[imageIndex])))
  1436. }
  1437. imageIndex++
  1438. messageChanged = true
  1439. continue
  1440. }
  1441. if text := grokComposerTextFromPart(part); text != "" {
  1442. textParts = append(textParts, text)
  1443. }
  1444. }
  1445. if messageChanged {
  1446. msgMap["content"] = strings.Join(textParts, "\n\n")
  1447. changed = true
  1448. }
  1449. }
  1450. return changed
  1451. }
  1452. func grokComposerTextFromPart(part any) string {
  1453. partMap, ok := part.(map[string]any)
  1454. if !ok {
  1455. return ""
  1456. }
  1457. partType := strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"])))
  1458. switch partType {
  1459. case "text", "input_text":
  1460. text, _ := partMap["text"].(string)
  1461. return strings.TrimSpace(text)
  1462. default:
  1463. return ""
  1464. }
  1465. }
  1466. func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) {
  1467. if dst == nil {
  1468. return
  1469. }
  1470. dst.InputTokens += usage.InputTokens
  1471. dst.ImageInputTokens += usage.ImageInputTokens
  1472. dst.OutputTokens += usage.OutputTokens
  1473. dst.CacheCreationInputTokens += usage.CacheCreationInputTokens
  1474. dst.CacheReadInputTokens += usage.CacheReadInputTokens
  1475. dst.ImageOutputTokens += usage.ImageOutputTokens
  1476. }
  1477. func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string, cfg *config.Config, settings ...*SettingService) (*http.Request, error) {
  1478. targetURL, err := buildGrokResponsesURL(account, cfg, settings...)
  1479. if err != nil {
  1480. return nil, err
  1481. }
  1482. req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
  1483. if err != nil {
  1484. return nil, err
  1485. }
  1486. req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileGrok))
  1487. req.Header.Set("Authorization", "Bearer "+token)
  1488. req.Header.Set("Content-Type", "application/json")
  1489. req.Header.Set("Accept", "application/json, text/event-stream")
  1490. if account.IsGrokOAuth() {
  1491. applyGrokCLIHeaders(req.Header)
  1492. }
  1493. applyGrokCacheHeaders(req.Header, cacheIdentity)
  1494. if c != nil {
  1495. if v := c.GetHeader("OpenAI-Beta"); strings.TrimSpace(v) != "" {
  1496. req.Header.Set("OpenAI-Beta", v)
  1497. }
  1498. }
  1499. // 账号级请求头覆写最后应用,使配置值优先于上面的内置默认头;
  1500. // 打到官方 CLI 网关时身份头仍由共享传输层最终强制。
  1501. account.ApplyHeaderOverrides(req.Header)
  1502. return req, nil
  1503. }
  1504. // applyGrokCLIHeaders identifies subscription traffic as a supported Grok CLI
  1505. // version. The CLI gateway rejects otherwise valid OAuth requests without it.
  1506. // Identity pins come from package xai so service-layer headers match the final
  1507. // transport rewrite on cli-chat-proxy.grok.com.
  1508. func applyGrokCLIHeaders(headers http.Header) {
  1509. if headers == nil {
  1510. return
  1511. }
  1512. version := xai.ResolveCLIVersion()
  1513. headers.Set("User-Agent", xai.CLIUserAgent(version))
  1514. headers.Set("X-Grok-Client-Version", version)
  1515. headers.Set("x-grok-client-version", version)
  1516. headers.Set("x-grok-client-identifier", xai.CLIClientIdentifier)
  1517. // Historical mode value expected by some unit tests / older CLI probes.
  1518. headers.Set("X-Grok-Client-Mode", "interactive")
  1519. }
  1520. func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot) {
  1521. s.updateGrokUsageSnapshotWithRateLimit(ctx, account, snapshot, true)
  1522. }
  1523. func (s *OpenAIGatewayService) updateGrokUsageSnapshotWithRateLimit(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot, installRateLimit bool) {
  1524. if s == nil || account == nil || account.ID <= 0 || snapshot == nil {
  1525. return
  1526. }
  1527. accountID := account.ID
  1528. now := time.Now()
  1529. resetAt, hasActiveLimit := grokRateLimitResetAtForAccount(account, snapshot, now)
  1530. if hasActiveLimit {
  1531. normalizeGrokExhaustedWindowResets(snapshot, resetAt, now)
  1532. }
  1533. recovery := isSuccessfulGrokRateLimitRecovery(account, snapshot)
  1534. critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit || recovery
  1535. if s.codexSnapshotThrottle != nil {
  1536. allowed := s.codexSnapshotThrottle.Allow(accountID, now)
  1537. if !critical && !allowed {
  1538. return
  1539. }
  1540. }
  1541. updates := map[string]any{
  1542. grokQuotaSnapshotExtraKey: snapshot,
  1543. }
  1544. // Also derive the scheduling-threshold extras (grok_sched_*) the evaluator
  1545. // reads in grokThresholdCandidates. Without this writer the admin-configured
  1546. // Grok auto-pause threshold could never fire (the read side was dead config).
  1547. for k, v := range buildGrokSchedulerExtraUpdates(snapshot) {
  1548. updates[k] = v
  1549. }
  1550. stateCtx := ctx
  1551. if hasActiveLimit {
  1552. var cancel context.CancelFunc
  1553. stateCtx, cancel = openAIAccountStateContext(ctx)
  1554. defer cancel()
  1555. }
  1556. // Account pointers on the request path are per-request copies (Redis/DB decode),
  1557. // not a shared in-process cache. Mutating Extra here matches token refresh /
  1558. // rate-limit writers; do not reuse the same *Account across goroutines.
  1559. if account.Extra == nil {
  1560. account.Extra = map[string]any{}
  1561. }
  1562. account.Extra[grokQuotaSnapshotExtraKey] = snapshot
  1563. if s.accountRepo != nil {
  1564. _ = s.accountRepo.UpdateExtra(stateCtx, accountID, updates)
  1565. }
  1566. // Error responses are reconciled by handleGrokAccountUpstreamError. Pool-mode
  1567. // API keys retain the snapshot for observability but leave account health to
  1568. // the upstream pool. Other accounts install the immediate runtime and durable
  1569. // rate-limit state when the observed window is exhausted.
  1570. if installRateLimit && hasActiveLimit && !account.IsPoolMode() {
  1571. s.rateLimitGrok(stateCtx, account, resetAt)
  1572. } else if recovery {
  1573. clearGrokRateLimitAfterRecovery(stateCtx, s.accountRepo, account)
  1574. }
  1575. }
  1576. func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) {
  1577. snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now())
  1578. if snapshot != nil {
  1579. stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
  1580. s.updateGrokUsageSnapshot(ctx, account, snapshot)
  1581. return
  1582. }
  1583. // Successful responses are recovery evidence even when the upstream omits
  1584. // optional quota headers. Do not replace an informative stored snapshot with
  1585. // an empty one; only clear the exact observed cooldown generation.
  1586. recoverySnapshot := &xai.QuotaSnapshot{StatusCode: statusCode}
  1587. if isSuccessfulGrokRateLimitRecovery(account, recoverySnapshot) {
  1588. clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
  1589. }
  1590. }
  1591. func parseGrokQuotaSnapshot(headers http.Header, statusCode int, now time.Time) *xai.QuotaSnapshot {
  1592. snapshot := xai.ParseQuotaHeaders(headers, statusCode)
  1593. if snapshot == nil && statusCode == http.StatusTooManyRequests {
  1594. return &xai.QuotaSnapshot{
  1595. StatusCode: statusCode,
  1596. UpdatedAt: now.UTC().Format(time.RFC3339),
  1597. }
  1598. }
  1599. return snapshot
  1600. }
  1601. func normalizeGrokExhaustedWindowResets(snapshot *xai.QuotaSnapshot, resetAt, now time.Time) {
  1602. if snapshot == nil || !resetAt.After(now) {
  1603. return
  1604. }
  1605. for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} {
  1606. if window == nil || window.Remaining == nil || *window.Remaining > 0 {
  1607. continue
  1608. }
  1609. candidate := time.Time{}
  1610. if window.ResetUnix != nil && *window.ResetUnix > 0 {
  1611. candidate = time.Unix(*window.ResetUnix, 0)
  1612. } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil {
  1613. candidate = parsed
  1614. }
  1615. if !candidate.After(now) {
  1616. candidate = resetAt
  1617. }
  1618. resetUnix := candidate.Unix()
  1619. window.ResetUnix = &resetUnix
  1620. window.ResetAt = candidate.UTC().Format(time.RFC3339)
  1621. }
  1622. }
  1623. func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) {
  1624. if snapshot == nil {
  1625. return time.Time{}, false
  1626. }
  1627. // Retry-After is xAI's explicit retry boundary. Use the observation time so
  1628. // a persisted snapshot does not start a fresh cooldown every time it is read.
  1629. retryAfterExpired := false
  1630. var resetAt time.Time
  1631. if snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 {
  1632. observedAt := now
  1633. if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(snapshot.UpdatedAt)); err == nil {
  1634. observedAt = parsed
  1635. }
  1636. retryAfterResetAt := observedAt.Add(time.Duration(*snapshot.RetryAfterSeconds) * time.Second)
  1637. if retryAfterResetAt.After(now) {
  1638. resetAt = retryAfterResetAt
  1639. } else {
  1640. retryAfterExpired = true
  1641. }
  1642. }
  1643. exhausted := false
  1644. for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} {
  1645. if window == nil || window.Remaining == nil || *window.Remaining > 0 {
  1646. continue
  1647. }
  1648. exhausted = true
  1649. candidate := time.Time{}
  1650. if window.ResetUnix != nil && *window.ResetUnix > 0 {
  1651. candidate = time.Unix(*window.ResetUnix, 0)
  1652. } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil {
  1653. candidate = parsed
  1654. }
  1655. if candidate.After(now) && candidate.After(resetAt) {
  1656. resetAt = candidate
  1657. }
  1658. }
  1659. if !resetAt.IsZero() {
  1660. return resetAt, true
  1661. }
  1662. // An observed Retry-After is an absolute boundary once combined with the
  1663. // snapshot timestamp. Do not turn an expired persisted snapshot into a new
  1664. // rolling fallback cooldown, but still allow a later explicit window reset.
  1665. if retryAfterExpired {
  1666. return time.Time{}, false
  1667. }
  1668. if exhausted || snapshot.StatusCode == http.StatusTooManyRequests {
  1669. return now.Add(grokRateLimitFallbackCooldown), true
  1670. }
  1671. return time.Time{}, false
  1672. }
  1673. func grokRateLimitResetAtForAccount(account *Account, snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) {
  1674. resetAt, limited := grokRateLimitResetAt(snapshot, now)
  1675. if !limited || !isGrokOAuthAccount(account) || snapshot == nil || snapshot.StatusCode != http.StatusTooManyRequests {
  1676. return resetAt, limited
  1677. }
  1678. if account.RateLimitedAt == nil || account.RateLimitResetAt == nil {
  1679. return resetAt, true
  1680. }
  1681. previousResetAt := *account.RateLimitResetAt
  1682. if previousResetAt.After(now) || now.Sub(previousResetAt) > grokRateLimitBackoffQuietPeriod {
  1683. return resetAt, true
  1684. }
  1685. previousCooldown := previousResetAt.Sub(*account.RateLimitedAt)
  1686. if previousCooldown <= 0 {
  1687. return resetAt, true
  1688. }
  1689. adaptiveCooldown := grokRateLimitRepeatCooldown
  1690. switch {
  1691. case previousCooldown >= grokRateLimitSustainedCooldown:
  1692. adaptiveCooldown = grokRateLimitMaxAdaptiveCooldown
  1693. case previousCooldown >= grokRateLimitRepeatCooldown:
  1694. adaptiveCooldown = grokRateLimitSustainedCooldown
  1695. }
  1696. adaptiveResetAt := now.Add(adaptiveCooldown)
  1697. if adaptiveResetAt.After(resetAt) {
  1698. resetAt = adaptiveResetAt
  1699. }
  1700. return resetAt, true
  1701. }
  1702. func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time {
  1703. if !resetAt.After(now) {
  1704. resetAt = now.Add(grokRateLimitFallbackCooldown)
  1705. }
  1706. if account != nil && account.RateLimitResetAt != nil && account.RateLimitResetAt.After(resetAt) {
  1707. resetAt = *account.RateLimitResetAt
  1708. }
  1709. return resetAt
  1710. }
  1711. type grokRateLimitExtendingRepository interface {
  1712. SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error
  1713. }
  1714. type grokRateLimitRecoveryRepository interface {
  1715. ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error)
  1716. }
  1717. func isSuccessfulGrokRateLimitRecovery(account *Account, snapshot *xai.QuotaSnapshot) bool {
  1718. return isGrokOAuthAccount(account) &&
  1719. account.RateLimitedAt != nil &&
  1720. account.RateLimitResetAt != nil &&
  1721. snapshot != nil &&
  1722. snapshot.StatusCode >= http.StatusOK &&
  1723. snapshot.StatusCode < http.StatusMultipleChoices
  1724. }
  1725. func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository, account *Account) {
  1726. if repo == nil || account == nil || account.RateLimitedAt == nil || account.RateLimitResetAt == nil || ctx.Err() != nil {
  1727. return
  1728. }
  1729. recoveryRepo, ok := repo.(grokRateLimitRecoveryRepository)
  1730. if !ok {
  1731. return
  1732. }
  1733. _, err := recoveryRepo.ClearRateLimitIfObserved(ctx, account.ID, *account.RateLimitedAt, *account.RateLimitResetAt)
  1734. if err != nil {
  1735. slog.Warn("grok_rate_limit_recovery_clear_failed", "account_id", account.ID, "error", err)
  1736. }
  1737. }
  1738. func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) {
  1739. if repo == nil || account == nil || account.ID <= 0 {
  1740. return
  1741. }
  1742. resetAt = normalizeGrokRateLimitResetAt(account, resetAt, time.Now())
  1743. stateCtx, cancel := openAIAccountStateContext(ctx)
  1744. defer cancel()
  1745. var err error
  1746. if extendingRepo, ok := repo.(grokRateLimitExtendingRepository); ok {
  1747. err = extendingRepo.SetRateLimitedIfLater(stateCtx, account.ID, resetAt)
  1748. } else {
  1749. err = repo.SetRateLimited(stateCtx, account.ID, resetAt)
  1750. }
  1751. if err != nil {
  1752. slog.Warn("persist_grok_rate_limit_failed", "account_id", account.ID, "reset_at", resetAt.UTC(), "error", err)
  1753. }
  1754. }
  1755. func (s *OpenAIGatewayService) rateLimitGrok(ctx context.Context, account *Account, resetAt time.Time) {
  1756. if s == nil || account == nil {
  1757. return
  1758. }
  1759. now := time.Now()
  1760. resetAt = normalizeGrokRateLimitResetAt(account, resetAt, now)
  1761. runtimeUntil := resetAt
  1762. if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(runtimeUntil) {
  1763. runtimeUntil = *account.TempUnschedulableUntil
  1764. }
  1765. s.BlockAccountScheduling(account, runtimeUntil, "429")
  1766. persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
  1767. // Propagate a short team+model cool so sibling OAuth accounts on the same
  1768. // xAI team skip the hot model without waiting for each to hit 429 alone.
  1769. // Model is taken from the latest request context when available; empty is a
  1770. // no-op inside markGrokTeamModelRateLimit.
  1771. if model, _ := ctx.Value(grokTeamRateLimitModelContextKey{}).(string); model != "" {
  1772. markGrokTeamModelRateLimit(account, model, resolveGrokTeamRateLimitUntil(resetAt, now))
  1773. }
  1774. }
  1775. // buildGrokSchedulerExtraUpdates derives the grok_sched_* scheduling snapshot
  1776. // (utilization percent + reset time) consumed by EvaluateAccountSchedulingThreshold.
  1777. // Utilization is the most-constrained of the requests/tokens windows.
  1778. func buildGrokSchedulerExtraUpdates(snapshot *xai.QuotaSnapshot) map[string]any {
  1779. if snapshot == nil {
  1780. return nil
  1781. }
  1782. util, reset, ok := grokSnapshotUtilization(snapshot)
  1783. if !ok {
  1784. return nil
  1785. }
  1786. updates := map[string]any{
  1787. "grok_sched_utilization": util,
  1788. "grok_sched_usage_updated_at": time.Now().UTC().Format(time.RFC3339),
  1789. }
  1790. if reset != nil {
  1791. // 防御:调度阈值暂停时长由 grok_sched_reset_at 决定。若上游返回脏的
  1792. // reset 头(例如把相对毫秒 "6000" 误当相对秒解析出 ~33h 的未来时刻),
  1793. // 不设上限会把耗尽账号长时间锁死。xAI 配额窗口不会超过一天,因此对
  1794. // 未来时刻做 grokMaxSchedulingResetHorizon 钳制;过去/无效值直接不写。
  1795. now := time.Now()
  1796. if reset.After(now) {
  1797. capped := *reset
  1798. if horizon := now.Add(grokMaxSchedulingResetHorizon); capped.After(horizon) {
  1799. capped = horizon
  1800. }
  1801. updates["grok_sched_reset_at"] = capped.UTC().Format(time.RFC3339)
  1802. }
  1803. }
  1804. return updates
  1805. }
  1806. // grokSnapshotUtilization returns the highest window utilization (0-100) across
  1807. // the requests/tokens quota windows and the reset time of that window.
  1808. func grokSnapshotUtilization(snapshot *xai.QuotaSnapshot) (float64, *time.Time, bool) {
  1809. if snapshot == nil {
  1810. return 0, nil, false
  1811. }
  1812. best := -1.0
  1813. var bestReset *time.Time
  1814. consider := func(window *xai.QuotaWindow) {
  1815. if window == nil || window.Limit == nil || *window.Limit <= 0 || window.Remaining == nil {
  1816. return
  1817. }
  1818. remaining := *window.Remaining
  1819. if remaining < 0 {
  1820. remaining = 0
  1821. }
  1822. util := (1 - float64(remaining)/float64(*window.Limit)) * 100
  1823. if util < 0 {
  1824. util = 0
  1825. }
  1826. if util > 100 {
  1827. util = 100
  1828. }
  1829. if util > best {
  1830. best = util
  1831. if window.ResetUnix != nil {
  1832. t := time.Unix(*window.ResetUnix, 0).UTC()
  1833. bestReset = &t
  1834. } else {
  1835. bestReset = nil
  1836. }
  1837. }
  1838. }
  1839. consider(snapshot.Requests)
  1840. consider(snapshot.Tokens)
  1841. if best < 0 {
  1842. return 0, nil, false
  1843. }
  1844. return best, bestReset, true
  1845. }
  1846. // grokMaxSchedulingResetHorizon bounds how far into the future a Grok
  1847. // scheduling-threshold pause (grok_sched_reset_at) may be set, so a malformed
  1848. // upstream reset header can't park an over-threshold account for days. xAI quota
  1849. // windows do not exceed ~a day.
  1850. const grokMaxSchedulingResetHorizon = 25 * time.Hour
  1851. // grokTeamRateLimitModelContextKey carries the upstream model for team cools.
  1852. type grokTeamRateLimitModelContextKey struct{}
  1853. // withGrokTeamRateLimitModel attaches the upstream model name for rate-limit
  1854. // side effects (team+model cool). Safe when model is empty.
  1855. func withGrokTeamRateLimitModel(ctx context.Context, model string) context.Context {
  1856. model = strings.TrimSpace(model)
  1857. if model == "" || ctx == nil {
  1858. return ctx
  1859. }
  1860. return context.WithValue(ctx, grokTeamRateLimitModelContextKey{}, model)
  1861. }
  1862. func grokRequestedModelFromCtx(ctx context.Context) string {
  1863. if ctx == nil {
  1864. return ""
  1865. }
  1866. model, _ := ctx.Value(grokTeamRateLimitModelContextKey{}).(string)
  1867. return strings.TrimSpace(model)
  1868. }
  1869. func persistGrokTransientModelCooldown(account *Account, decision GrokUpstreamFailureDecision) bool {
  1870. if account == nil {
  1871. return false
  1872. }
  1873. model := strings.TrimSpace(decision.Model)
  1874. if model == "" {
  1875. return false
  1876. }
  1877. cooldown := decision.Cooldown
  1878. if cooldown <= 0 {
  1879. cooldown = 3 * time.Minute
  1880. }
  1881. markGrokModelTransientBlock(account.ID, model, time.Now().Add(cooldown))
  1882. return true
  1883. }
  1884. func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) {
  1885. if s == nil || account == nil {
  1886. return
  1887. }
  1888. if isGrokContentPolicyRejection(statusCode, responseBody) {
  1889. return
  1890. }
  1891. now := time.Now()
  1892. decision := classifyGrokUpstreamFailure(statusCode, responseBody, grokRequestedModelFromCtx(ctx))
  1893. snapshot := parseGrokQuotaSnapshot(headers, statusCode, now)
  1894. stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
  1895. // Capacity 429 is model pressure, not account quota exhaustion. Keep the
  1896. // snapshot for observability but do not install account-level rate limiting;
  1897. // the failover decision below applies a bounded model-scoped block instead.
  1898. s.updateGrokUsageSnapshotWithRateLimit(ctx, account, snapshot, decision.Class != GrokFailureModelCapacity)
  1899. // Body-first free-usage / empty / billing / capacity must run before the
  1900. // status switch so non-429 free-usage bodies still cool the account.
  1901. // Pool-mode still skips durable mutation unless an explicit temp rule matches.
  1902. if decision.ShouldCooldown && decision.Class != GrokFailureNone && decision.Class != GrokFailureRateLimit {
  1903. if account.IsPoolMode() {
  1904. // Allow configured temp rules (403) below; skip default body cools.
  1905. } else {
  1906. // A free-tier exhaustion message describes a rolling usage window. Use
  1907. // an upstream absolute reset (or Retry-After) when available; otherwise
  1908. // apply only a short probe cooldown. Never start a fabricated 24h window
  1909. // at the instant this error was received.
  1910. if decision.Class == GrokFailureFreeUsage {
  1911. if resetAt, limited := grokRateLimitResetAtForAccount(account, parseGrokQuotaSnapshot(headers, statusCode, now), now); limited && resetAt.After(now) {
  1912. if decision.Model != "" && isGrokModelSpecificFreeUsage(strings.ToLower(decision.Reason), decision.Model) {
  1913. markGrokModelQuotaBlock(account.ID, decision.Model, resetAt)
  1914. return
  1915. }
  1916. s.rateLimitGrok(ctx, account, resetAt)
  1917. return
  1918. }
  1919. }
  1920. if s.applyGrokUpstreamFailureDecision(ctx, account, decision) {
  1921. return
  1922. }
  1923. }
  1924. }
  1925. if statusCode == http.StatusForbidden && s.applyGrokForbiddenPolicy(ctx, account, responseBody) {
  1926. return
  1927. }
  1928. if account.IsPoolMode() {
  1929. slog.Info("grok_pool_mode_error_state_skipped", "account_id", account.ID, "status_code", statusCode)
  1930. return
  1931. }
  1932. switch statusCode {
  1933. case http.StatusUnauthorized:
  1934. s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok credentials unauthorized")
  1935. case http.StatusPaymentRequired:
  1936. // 402 without a body-classified billing decision: keep the legacy 30m cool.
  1937. s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok payment required")
  1938. case http.StatusForbidden:
  1939. // Spending-limit already handled by body classifier when phrasing matches.
  1940. if isGrokSpendingLimitError(responseBody) {
  1941. s.rateLimitGrok(ctx, account, grokSpendingLimitResetAt(account, time.Now()))
  1942. return
  1943. }
  1944. s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok access or entitlement denied")
  1945. case http.StatusTooManyRequests:
  1946. // updateGrokUsageSnapshot installs rate-limit state for non-pool accounts.
  1947. // Free-usage 429 was already cooled above via body classification.
  1948. default:
  1949. if statusCode >= 500 {
  1950. s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error")
  1951. }
  1952. }
  1953. }
  1954. // isGrokSpendingLimitError detects xAI billing exhaustion bodies (often 403, sometimes 402).
  1955. func isGrokSpendingLimitError(responseBody []byte) bool {
  1956. if len(responseBody) == 0 {
  1957. return false
  1958. }
  1959. code := strings.ToLower(strings.TrimSpace(firstNonEmpty(
  1960. gjson.GetBytes(responseBody, "code").String(),
  1961. gjson.GetBytes(responseBody, "error.code").String(),
  1962. )))
  1963. if code == "personal-team-blocked:spending-limit" {
  1964. return true
  1965. }
  1966. message := strings.ToLower(strings.TrimSpace(firstNonEmpty(
  1967. gjson.GetBytes(responseBody, "error").String(),
  1968. gjson.GetBytes(responseBody, "error.message").String(),
  1969. gjson.GetBytes(responseBody, "message").String(),
  1970. )))
  1971. return strings.Contains(message, "spending limit") ||
  1972. strings.Contains(message, "run out of credits")
  1973. }
  1974. func (s *OpenAIGatewayService) tempUnscheduleGrok(ctx context.Context, account *Account, cooldown time.Duration, reason string) {
  1975. if s == nil || account == nil {
  1976. return
  1977. }
  1978. until := time.Now().Add(cooldown)
  1979. if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(until) {
  1980. until = *account.TempUnschedulableUntil
  1981. }
  1982. s.BlockAccountScheduling(account, until, reason)
  1983. if s.accountRepo != nil {
  1984. stateCtx, cancel := openAIAccountStateContext(ctx)
  1985. defer cancel()
  1986. _ = s.accountRepo.SetTempUnschedulable(stateCtx, account.ID, until, reason)
  1987. }
  1988. }