| 12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619162016211622162316241625162616271628162916301631163216331634163516361637163816391640164116421643164416451646164716481649165016511652165316541655165616571658165916601661166216631664166516661667166816691670167116721673167416751676167716781679168016811682168316841685168616871688168916901691169216931694169516961697169816991700170117021703170417051706170717081709171017111712171317141715171617171718171917201721172217231724172517261727172817291730173117321733173417351736173717381739174017411742174317441745174617471748174917501751175217531754175517561757175817591760176117621763176417651766176717681769177017711772177317741775177617771778177917801781178217831784178517861787178817891790179117921793179417951796179717981799180018011802180318041805180618071808180918101811181218131814181518161817181818191820182118221823182418251826182718281829183018311832183318341835183618371838183918401841184218431844184518461847184818491850185118521853185418551856185718581859186018611862186318641865186618671868186918701871187218731874187518761877187818791880188118821883188418851886188718881889189018911892189318941895189618971898189919001901190219031904190519061907190819091910191119121913191419151916191719181919192019211922192319241925192619271928192919301931193219331934193519361937193819391940194119421943194419451946194719481949195019511952195319541955195619571958195919601961196219631964196519661967196819691970197119721973197419751976197719781979198019811982198319841985198619871988198919901991199219931994199519961997199819992000200120022003200420052006200720082009201020112012201320142015201620172018201920202021202220232024202520262027202820292030203120322033203420352036203720382039204020412042204320442045204620472048204920502051205220532054205520562057205820592060206120622063206420652066206720682069207020712072207320742075207620772078207920802081208220832084208520862087208820892090209120922093209420952096209720982099210021012102210321042105210621072108210921102111211221132114211521162117 |
- package service
- import (
- "bytes"
- "context"
- "encoding/json"
- "fmt"
- "io"
- "log/slog"
- "net/http"
- "strings"
- "time"
- "github.com/Wei-Shaw/sub2api/internal/config"
- "github.com/Wei-Shaw/sub2api/internal/pkg/apicompat"
- "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
- "github.com/gin-gonic/gin"
- "github.com/tidwall/gjson"
- "github.com/tidwall/sjson"
- )
- const (
- grokComposerImageBridgeVisionModel = "grok-build-0.1"
- grokComposerImageBridgeMaxOutputTokens = 512
- grokCLIVersion = xai.CLIClientVersion
- grokDefaultResponsesModel = "grok-4.5"
- grokRateLimitFallbackCooldown = 2 * time.Minute
- grokRateLimitRepeatCooldown = 10 * time.Minute
- grokRateLimitSustainedCooldown = 30 * time.Minute
- grokRateLimitMaxAdaptiveCooldown = time.Hour
- grokRateLimitBackoffQuietPeriod = time.Hour
- )
- func (s *OpenAIGatewayService) forwardGrokResponses(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- body []byte,
- originalModel string,
- reqStream bool,
- startTime time.Time,
- ) (*OpenAIForwardResult, error) {
- if account.Type != AccountTypeOAuth && account.Type != AccountTypeAPIKey {
- return nil, fmt.Errorf("grok account type %s is not supported by Responses forwarding", account.Type)
- }
- upstreamModel := account.GetMappedModel(originalModel)
- if strings.TrimSpace(upstreamModel) == "" {
- upstreamModel = grokDefaultResponsesModel
- }
- // Account mappings are optional. Canonicalize client aliases even when the
- // account has no model_mapping, matching the Chat Completions path and xAI's
- // actual Responses model IDs.
- upstreamModel = xai.ResolveGrokTextResponsesModelID(upstreamModel, grokDefaultResponsesModel)
- if isGrokImageGenerationModel(upstreamModel) {
- return nil, fmt.Errorf("model %s is an image model and is not available on the Responses endpoint; use /v1/images/generations instead", upstreamModel)
- }
- patchedBody, clientToolMapping, err := patchGrokResponsesBodyWithClientTools(body, upstreamModel)
- if err != nil {
- setOpsUpstreamError(c, http.StatusBadRequest, err.Error(), "")
- c.JSON(http.StatusBadRequest, gin.H{"error": gin.H{
- "type": "invalid_request_error", "message": err.Error(), "param": "tools",
- }})
- return nil, err
- }
- setGrokResponsesClientToolMapping(c, clientToolMapping)
- // OpenAI /responses/compact is not a native xAI endpoint. Convert it into a
- // normal Grok Responses turn that asks for a structured summary, then map the
- // reply back to an OpenAI compaction item on the way out.
- if isOpenAIResponsesCompactPath(c) {
- patchedBody, err = buildGrokCompactRequestBody(patchedBody)
- if err != nil {
- return nil, err
- }
- }
- // Derive the identity from the request xAI will actually see. This makes
- // Codex Responses Lite additional_tools part of the stable tool prefix. If
- // the client supplied a Claude Code session only through metadata.user_id,
- // keep using that identity even though metadata is stripped before xAI.
- cacheIdentityBody := patchedBody
- if extractClaudeCodeSessionIDFromPayload(body) != "" {
- cacheIdentityBody = body
- }
- cacheIdentity := resolveGrokCacheIdentity(c, cacheIdentityBody, "", upstreamModel)
- mixedCacheIntentBody := append([]byte(nil), patchedBody...)
- patchedBody, err = applyGrokResponsesCacheIdentity(patchedBody, body, cacheIdentity, account.IsGrokOAuth())
- if err != nil {
- return nil, fmt.Errorf("apply grok prompt cache identity: %w", err)
- }
- // Free OAuth + client function tools: reuse Messages mixed-tools cache route
- // (append web_search/x_search so xAI does not force non-cacheable build-free).
- patchedBody, err = applyGrokFreeRequestToolCacheRoute(c, patchedBody, mixedCacheIntentBody, account, cacheIdentity)
- if err != nil {
- return nil, fmt.Errorf("apply grok Free function-tool cache route: %w", err)
- }
- token, _, err := s.getRequestCredential(ctx, c, account)
- if err != nil {
- return nil, err
- }
- upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
- defer releaseUpstreamCtx()
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- upstreamStart := time.Now()
- var resp *http.Response
- for attempt := 0; ; attempt++ {
- upstreamReq, buildErr := buildGrokResponsesRequest(upstreamCtx, c, account, patchedBody, token, cacheIdentity, s.cfg, s.settingService)
- if buildErr != nil {
- return nil, buildErr
- }
- resp, err = s.doOpenAIUpstream(upstreamReq, proxyURL, account)
- SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
- if err != nil {
- return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
- }
- // xAI can reject encrypted reasoning or a compaction blob copied from a
- // different decoder/cache context. Retry once on the same account after
- // preserving visible summaries and removing only opaque replay state.
- if attempt > 0 || (resp.StatusCode != http.StatusBadRequest && resp.StatusCode != http.StatusUnprocessableEntity) {
- break
- }
- respBody := s.readUpstreamErrorBody(resp)
- if resp.Body != nil {
- _ = resp.Body.Close()
- }
- invalidEncryptedContent := isGrokInvalidEncryptedContentResponse(resp.StatusCode, respBody)
- if !invalidEncryptedContent && !isGrokCompactionReplayDecodeError(resp.StatusCode, respBody) {
- resp.Body = io.NopCloser(bytes.NewReader(respBody))
- break
- }
- var retryBody []byte
- var changed bool
- var trimErr error
- if invalidEncryptedContent {
- retryBody, changed, trimErr = trimGrokInvalidEncryptedContentRetryBody(patchedBody)
- } else {
- retryBody, changed, trimErr = sanitizeGrokCompactionReplayBody(patchedBody)
- }
- if trimErr != nil {
- return nil, fmt.Errorf("prepare Grok replay decode retry: %w", trimErr)
- }
- if !changed {
- resp.Body = io.NopCloser(bytes.NewReader(respBody))
- break
- }
- patchedBody = retryBody
- slog.Info("grok_replay_decode_retry", "account_id", account.ID, "cache_identity_present", cacheIdentity != "")
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode >= 400 {
- respBody := s.readUpstreamErrorBody(resp)
- resp.Body = io.NopCloser(bytes.NewReader(respBody))
- upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody))
- if upstreamMsg == "" {
- upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode)
- }
- kind := "http_error"
- if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
- kind = "failover"
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
- Kind: kind,
- Message: upstreamMsg,
- })
- errCtx := withGrokTeamRateLimitModel(ctx, upstreamModel)
- s.handleGrokAccountUpstreamError(errCtx, account, resp.StatusCode, resp.Header, respBody)
- // Quota/rate-limit responses stamp the team+model overlay. Capacity is
- // request pressure and must not hide sibling accounts.
- if shouldMarkGrokTeamModelRateLimit(resp.StatusCode, respBody) {
- markGrokTeamModelRateLimit(account, upstreamModel, resolveGrokTeamRateLimitUntil(time.Now().Add(grokTeamRateLimitDefaultTTL), time.Now()))
- }
- if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
- retryable, retryDelay, retryDeadline, retryMax := grokSameAccountRetryMetadata(account, resp.StatusCode, respBody)
- return nil, &UpstreamFailoverError{
- StatusCode: resp.StatusCode,
- ResponseBody: respBody,
- ResponseHeaders: resp.Header.Clone(),
- RetryableOnSameAccount: retryable,
- RequestScopedTransient: retryable && resp.StatusCode == http.StatusTooManyRequests,
- SameAccountRetryDelay: retryDelay,
- SameAccountRetryDeadline: retryDeadline,
- SameAccountRetryMax: retryMax,
- }
- }
- return s.handleErrorResponse(ctx, resp, c, account, patchedBody, upstreamModel)
- }
- // Attach model so rate-limit snapshots can fan out a team+model cool.
- stateCtx := withGrokTeamRateLimitModel(ctx, upstreamModel)
- s.updateGrokUsageFromResponse(stateCtx, account, resp.Header, resp.StatusCode)
- var usage *OpenAIUsage
- var firstTokenMs *int
- responseID := ""
- searchCount := 0
- imageCount := 0
- var imageOutputSizes []string
- if reqStream {
- maxLineSize := defaultMaxLineSize
- if s.cfg != nil && s.cfg.Gateway.MaxLineSize > 0 {
- maxLineSize = s.cfg.Gateway.MaxLineSize
- }
- resp.Body = newGrokResponsesBillingPingFilterBody(resp.Body, account, maxLineSize)
- if hasGrokResponsesClientToolMapping(clientToolMapping) {
- resp.Body = newGrokResponsesClientToolStreamBody(resp.Body, clientToolMapping, maxLineSize)
- }
- streamResult, err := s.handleStreamingResponse(ctx, resp, c, account, startTime, originalModel, upstreamModel)
- if err != nil {
- return nil, err
- }
- usage = streamResult.usage
- firstTokenMs = streamResult.firstTokenMs
- responseID = strings.TrimSpace(streamResult.responseID)
- searchCount = streamResult.searchCount
- imageCount = streamResult.imageCount
- imageOutputSizes = streamResult.imageOutputSizes
- } else {
- nonStreamResult, err := s.handleNonStreamingResponse(ctx, resp, c, account, originalModel, upstreamModel)
- if err != nil {
- return nil, err
- }
- usage = nonStreamResult.usage
- responseID = strings.TrimSpace(nonStreamResult.responseID)
- searchCount = nonStreamResult.searchCount
- imageCount = nonStreamResult.imageCount
- imageOutputSizes = nonStreamResult.imageOutputSizes
- }
- if usage == nil {
- usage = &OpenAIUsage{}
- }
- reasoningEffort := extractOpenAIReasoningEffortFromBody(patchedBody, originalModel)
- result := &OpenAIForwardResult{
- RequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
- UpstreamHeaders: resp.Header,
- ResponseID: responseID,
- Usage: *usage,
- Model: originalModel,
- UpstreamModel: upstreamModel,
- ReasoningEffort: reasoningEffort,
- Stream: reqStream,
- OpenAIWSMode: false,
- ResponseHeaders: resp.Header.Clone(),
- Duration: time.Since(startTime),
- FirstTokenMs: firstTokenMs,
- }
- // Propagate search/image counters from the shared Responses handler — without
- // this, stream/JSON counting runs but search_price_per_1k / image bills never apply.
- if searchCount > 0 {
- result.SearchCount = searchCount
- }
- if imageCount > 0 {
- result.ImageCount = imageCount
- result.ImageOutputSizes = imageOutputSizes
- }
- return result, nil
- }
- func isGrokInvalidEncryptedContentResponse(statusCode int, body []byte) bool {
- if statusCode != http.StatusBadRequest && statusCode != http.StatusUnprocessableEntity {
- return false
- }
- // xAI has used both flat and nested error envelopes:
- // {"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}
- // {"error":{"message":"Could not decrypt the provided encrypted_content."}}
- code := strings.TrimSpace(gjson.GetBytes(body, "code").String())
- errNode := gjson.GetBytes(body, "error")
- if code == "" && errNode.IsObject() {
- code = strings.TrimSpace(errNode.Get("code").String())
- }
- if strings.EqualFold(code, "invalid_encrypted_content") || strings.EqualFold(code, "invalid_compaction") || strings.EqualFold(code, "compaction_decode_error") {
- return true
- }
- // Keep the official xAI flat-code gate so unrelated 400s are not retried.
- if !strings.EqualFold(code, "invalid-argument") && code != "" {
- return false
- }
- for _, candidate := range grokStructuredErrorMessageCandidates(body) {
- normalizedMessage := strings.ToLower(candidate)
- // Nested OpenAI-style envelopes may omit top-level code; require decrypt text.
- if code == "" && !strings.Contains(normalizedMessage, "decrypt") && !strings.Contains(normalizedMessage, "decode the compaction blob") {
- continue
- }
- if strings.Contains(normalizedMessage, "encrypted_content") &&
- (strings.Contains(normalizedMessage, "decrypt") || strings.Contains(normalizedMessage, "unmodified")) {
- return true
- }
- if strings.Contains(normalizedMessage, "decode the compaction blob") {
- return true
- }
- }
- return false
- }
- func isGrokCompactionReplayDecodeError(statusCode int, body []byte) bool {
- if (statusCode != http.StatusBadRequest && statusCode != http.StatusUnprocessableEntity) || len(body) == 0 {
- return false
- }
- for _, candidate := range grokStructuredErrorMessageCandidates(body) {
- message := strings.ToLower(candidate)
- decodeSignal := strings.Contains(message, "decode") ||
- strings.Contains(message, "deserialize") ||
- strings.Contains(message, "decoder")
- replaySignal := strings.Contains(message, "compaction") ||
- strings.Contains(message, "summary") ||
- strings.Contains(message, "encrypted_content") ||
- strings.Contains(message, "response history")
- if decodeSignal && replaySignal {
- return true
- }
- }
- return false
- }
- func sanitizeGrokCompactionReplayBody(body []byte) ([]byte, bool, error) {
- converted, err := convertOpenAICompactInputsForGrok(body)
- if err != nil {
- return nil, false, fmt.Errorf("convert Grok compaction replay: %w", err)
- }
- var requestBody map[string]any
- decoder := json.NewDecoder(bytes.NewReader(converted))
- decoder.UseNumber()
- if err := decoder.Decode(&requestBody); err != nil {
- return nil, false, err
- }
- changed := !bytes.Equal(converted, body)
- if trimOpenAIEncryptedReasoningItems(requestBody) {
- changed = true
- }
- if dropEmptyGrokReplayReasoning(requestBody) {
- changed = true
- }
- if previousID, _ := requestBody["previous_response_id"].(string); strings.TrimSpace(previousID) != "" && !HasFunctionCallOutput(requestBody) {
- delete(requestBody, "previous_response_id")
- if _, exists := requestBody["store"]; !exists {
- requestBody["store"] = false
- }
- changed = true
- }
- if !changed {
- return body, false, nil
- }
- retryBody, err := marshalOpenAIUpstreamJSON(requestBody)
- if err != nil {
- return nil, false, err
- }
- return retryBody, true, nil
- }
- func dropEmptyGrokReplayReasoning(requestBody map[string]any) bool {
- items, ok := requestBody["input"].([]any)
- if !ok {
- return false
- }
- filtered := items[:0]
- changed := false
- for _, rawItem := range items {
- item, ok := rawItem.(map[string]any)
- if !ok || strings.TrimSpace(grokStringValue(item["type"])) != "reasoning" {
- filtered = append(filtered, rawItem)
- continue
- }
- summary, _ := item["summary"].([]any)
- content, hasContent := item["content"]
- _, hasEncrypted := item["encrypted_content"]
- if hasEncrypted || len(summary) > 0 || (hasContent && content != nil) {
- filtered = append(filtered, rawItem)
- continue
- }
- changed = true
- }
- if changed {
- requestBody["input"] = filtered
- }
- return changed
- }
- // requestHasGrokEncryptedReasoning reports whether the outbound Responses body
- // still carries reasoning.encrypted_content that can be stripped for retry.
- func requestHasGrokEncryptedReasoning(body []byte) bool {
- input := gjson.GetBytes(body, "input")
- if !input.Exists() {
- return false
- }
- items := input.Array()
- if input.IsObject() {
- items = []gjson.Result{input}
- }
- for _, item := range items {
- if strings.TrimSpace(item.Get("type").String()) != "reasoning" {
- continue
- }
- enc := item.Get("encrypted_content")
- if enc.Exists() && enc.Type != gjson.Null && strings.TrimSpace(enc.String()) != "" {
- return true
- }
- }
- return false
- }
- type grokEncryptedContentStripRetriedKey struct{}
- func markGrokEncryptedContentStripRetried(ctx context.Context) context.Context {
- return context.WithValue(ctx, grokEncryptedContentStripRetriedKey{}, true)
- }
- func grokEncryptedContentStripRetried(ctx context.Context) bool {
- v, _ := ctx.Value(grokEncryptedContentStripRetriedKey{}).(bool)
- return v
- }
- // stripAnthropicThinkingSignatures removes thinking.signature from Claude
- // history so a different Grok OAuth account can accept multi-turn tool
- // continuations after decrypt failures. Returns ok=false when nothing changed.
- func stripAnthropicThinkingSignatures(body []byte) ([]byte, bool) {
- if len(body) == 0 || !bytes.Contains(body, []byte(`"signature"`)) {
- return body, false
- }
- var req map[string]any
- if err := decodeOpenAIJSONUseNumber(body, &req); err != nil {
- return body, false
- }
- messages, ok := req["messages"].([]any)
- if !ok || len(messages) == 0 {
- return body, false
- }
- changed := false
- for _, rawMsg := range messages {
- msg, ok := rawMsg.(map[string]any)
- if !ok {
- continue
- }
- content, ok := msg["content"].([]any)
- if !ok {
- continue
- }
- for _, rawBlock := range content {
- block, ok := rawBlock.(map[string]any)
- if !ok {
- continue
- }
- if typ, _ := block["type"].(string); typ != "thinking" {
- continue
- }
- if _, has := block["signature"]; has {
- delete(block, "signature")
- changed = true
- }
- }
- }
- if !changed {
- return body, false
- }
- out, err := marshalOpenAIUpstreamJSON(req)
- if err != nil {
- return body, false
- }
- return out, true
- }
- func trimGrokInvalidEncryptedContentRetryBody(body []byte) ([]byte, bool, error) {
- input := gjson.GetBytes(body, "input")
- items := input.Array()
- if input.IsObject() {
- items = []gjson.Result{input}
- }
- hasEncryptedReasoning := false
- for _, item := range items {
- if (strings.TrimSpace(item.Get("type").String()) == "reasoning" && item.Get("encrypted_content").Exists()) ||
- (isOpenAICompactionType(strings.TrimSpace(item.Get("type").String())) && item.Get("encrypted_content").Exists()) {
- hasEncryptedReasoning = true
- break
- }
- }
- if !hasEncryptedReasoning {
- return body, false, nil
- }
- var requestBody map[string]any
- decoder := json.NewDecoder(bytes.NewReader(body))
- decoder.UseNumber()
- if err := decoder.Decode(&requestBody); err != nil {
- return nil, false, err
- }
- if !trimOpenAIEncryptedReasoningItems(requestBody) {
- return body, false, nil
- }
- retryBody, err := marshalOpenAIUpstreamJSON(requestBody)
- if err != nil {
- return nil, false, err
- }
- return retryBody, true, nil
- }
- func patchGrokResponsesBody(body []byte, upstreamModel string) ([]byte, error) {
- return patchGrokResponsesBodyBase(body, upstreamModel)
- }
- func patchGrokResponsesBodyWithClientTools(body []byte, upstreamModel string) ([]byte, apicompat.ResponsesClientToolMapping, error) {
- if !json.Valid(body) {
- return nil, apicompat.ResponsesClientToolMapping{}, fmt.Errorf("invalid json request body")
- }
- promoted, err := sanitizeGrokResponsesInput(body)
- if err != nil {
- return nil, apicompat.ResponsesClientToolMapping{}, err
- }
- adapted, mapping, err := adaptGrokResponsesClientTools(promoted)
- if err != nil {
- return nil, apicompat.ResponsesClientToolMapping{}, err
- }
- patched, err := patchGrokResponsesBodyBase(adapted, upstreamModel)
- if err != nil {
- return nil, apicompat.ResponsesClientToolMapping{}, err
- }
- return patched, mapping, nil
- }
- func patchGrokResponsesBodyBase(body []byte, upstreamModel string) ([]byte, error) {
- if !json.Valid(body) {
- return nil, fmt.Errorf("invalid json request body")
- }
- // sjson may reuse the input backing array; keep the caller's request bytes
- // unchanged because the same body can be inspected for billing/retry paths.
- out, err := sjson.SetBytes(append([]byte(nil), body...), "model", upstreamModel)
- if err != nil {
- return nil, err
- }
- out, err = normalizeGrokResponsesReasoningEffort(out, upstreamModel)
- if err != nil {
- return nil, err
- }
- out, err = sanitizeGrokResponsesModelCapabilities(out, upstreamModel)
- if err != nil {
- return nil, err
- }
- for _, unsupportedField := range []string{"prompt_cache_retention", "safety_identifier", "metadata"} {
- if gjson.GetBytes(out, unsupportedField).Exists() {
- out, err = sjson.DeleteBytes(out, unsupportedField)
- if err != nil {
- return nil, err
- }
- }
- }
- if strings.EqualFold(upstreamModel, "grok-4.5") {
- for _, unsupportedField := range []string{"presence_penalty", "presencePenalty", "frequency_penalty", "frequencyPenalty", "stop"} {
- if gjson.GetBytes(out, unsupportedField).Exists() {
- out, err = sjson.DeleteBytes(out, unsupportedField)
- if err != nil {
- return nil, err
- }
- }
- }
- }
- if grokModelRejectsLogprobs(upstreamModel) {
- for _, unsupportedField := range []string{"logprobs", "top_logprobs"} {
- if gjson.GetBytes(out, unsupportedField).Exists() {
- out, err = sjson.DeleteBytes(out, unsupportedField)
- if err != nil {
- return nil, err
- }
- }
- }
- }
- out, err = sanitizeGrokResponsesUnsupportedFields(out)
- if err != nil {
- return nil, err
- }
- out, err = convertOpenAICompactInputsForGrok(out)
- if err != nil {
- return nil, err
- }
- out, err = sanitizeGrokResponsesInput(out)
- if err != nil {
- return nil, err
- }
- out, err = sanitizeGrokResponsesModelInput(out)
- if err != nil {
- return nil, err
- }
- out, err = stripRedundantGrokViewImageTool(out)
- if err != nil {
- return nil, err
- }
- out, err = sanitizeGrokReasoningNullContent(out)
- if err != nil {
- return nil, err
- }
- out, err = sanitizeGrokResponsesTools(out)
- if err != nil {
- return nil, err
- }
- return out, nil
- }
- // xAI's Grok 4.20 family and newer models do not support OpenAI's logprobs
- // fields. Remove them before egress instead of forwarding a request the
- // upstream rejects. Older Grok models retain the fields for compatibility.
- func grokModelRejectsLogprobs(model string) bool {
- model = strings.ToLower(strings.TrimSpace(model))
- if slash := strings.LastIndex(model, "/"); slash >= 0 {
- model = strings.TrimSpace(model[slash+1:])
- }
- return strings.HasPrefix(model, "grok-4.20")
- }
- func sanitizeGrokResponsesModelCapabilities(body []byte, upstreamModel string) ([]byte, error) {
- if !grokModelRejectsReasoningEffort(upstreamModel) {
- return body, nil
- }
- out := body
- for _, field := range []string{"reasoning", "reasoning_effort", "reasoningEffort"} {
- if !gjson.GetBytes(out, field).Exists() {
- continue
- }
- var err error
- out, err = sjson.DeleteBytes(out, field)
- if err != nil {
- return nil, fmt.Errorf("remove unsupported Grok Composer %s: %w", field, err)
- }
- }
- return out, nil
- }
- func grokModelRejectsReasoningEffort(model string) bool {
- model = strings.TrimSpace(strings.ToLower(model))
- if slash := strings.LastIndex(model, "/"); slash >= 0 {
- model = strings.TrimSpace(model[slash+1:])
- }
- switch model {
- case "grok-composer", "grok-composer-2.5-fast", "composer-2.5":
- return true
- default:
- return false
- }
- }
- func normalizeGrokResponsesReasoningEffort(body []byte, upstreamModel string) ([]byte, error) {
- supportsEffort := grokSupportsReasoningEffort(upstreamModel)
- out := body
- var err error
- for _, field := range []string{"reasoning.effort", "reasoning_effort"} {
- value := gjson.GetBytes(out, field)
- if !value.Exists() {
- continue
- }
- normalized, keep := normalizeGrokReasoningEffortValue(value.String(), upstreamModel)
- if !supportsEffort || !keep {
- out, err = sjson.DeleteBytes(out, field)
- } else {
- out, err = sjson.SetBytes(out, field, normalized)
- }
- if err != nil {
- return nil, fmt.Errorf("normalize Grok reasoning field %s: %w", field, err)
- }
- }
- if camel := gjson.GetBytes(out, "reasoningEffort"); camel.Exists() {
- normalized, keep := normalizeGrokReasoningEffortValue(camel.String(), upstreamModel)
- out, err = sjson.DeleteBytes(out, "reasoningEffort")
- if err != nil {
- return nil, fmt.Errorf("remove Grok reasoningEffort: %w", err)
- }
- if supportsEffort && keep && !gjson.GetBytes(out, "reasoning_effort").Exists() {
- out, err = sjson.SetBytes(out, "reasoning_effort", normalized)
- if err != nil {
- return nil, fmt.Errorf("set Grok reasoning_effort: %w", err)
- }
- }
- }
- if reasoning := gjson.GetBytes(out, "reasoning"); reasoning.Exists() && reasoning.IsObject() && len(reasoning.Map()) == 0 {
- out, err = sjson.DeleteBytes(out, "reasoning")
- if err != nil {
- return nil, fmt.Errorf("remove empty Grok reasoning: %w", err)
- }
- }
- return out, nil
- }
- func normalizeGrokChatReasoningEffort(body []byte, upstreamModel string) ([]byte, error) {
- raw := strings.TrimSpace(gjson.GetBytes(body, "reasoning_effort").String())
- if raw == "" {
- raw = strings.TrimSpace(gjson.GetBytes(body, "reasoningEffort").String())
- }
- normalized, keep := normalizeGrokReasoningEffortValue(raw, upstreamModel)
- keep = keep && grokSupportsReasoningEffort(upstreamModel)
- out := body
- var err error
- if gjson.GetBytes(out, "reasoningEffort").Exists() {
- out, err = sjson.DeleteBytes(out, "reasoningEffort")
- if err != nil {
- return nil, err
- }
- }
- if !keep {
- if gjson.GetBytes(out, "reasoning_effort").Exists() {
- out, err = sjson.DeleteBytes(out, "reasoning_effort")
- }
- return out, err
- }
- out, err = sjson.SetBytes(out, "reasoning_effort", normalized)
- return out, err
- }
- func normalizeGrokReasoningEffortValue(raw, model string) (string, bool) {
- value := strings.NewReplacer("-", "", "_", "", " ", "").Replace(strings.ToLower(strings.TrimSpace(raw)))
- switch value {
- case "none", "low", "medium", "high":
- return value, true
- case "minimal":
- return "low", true
- case "xhigh", "extrahigh":
- if GrokSupportsXHighReasoningEffort(model) {
- return "xhigh", true
- }
- return "high", true
- case "max", "ultra":
- return "high", true
- default:
- return "", false
- }
- }
- // GrokSupportsXHighReasoningEffort reports whether the model advertises and
- // forwards the xhigh reasoning effort (Grok 4.6 and its undated alias).
- func GrokSupportsXHighReasoningEffort(model string) bool {
- model = strings.ToLower(xai.StripGrokProviderPrefix(strings.TrimSpace(model)))
- return model == "grok-4.6" || model == "grok-4.6-latest"
- }
- func grokSupportsReasoningEffort(model string) bool {
- model = strings.ToLower(xai.StripGrokProviderPrefix(strings.TrimSpace(model)))
- switch model {
- case "grok-4.5", "grok-4.5-latest", "grok-4.6", "grok-4.6-latest",
- "grok-4.3", "grok-4.3-latest",
- "grok-3-mini", "grok-3-mini-fast", "grok-4.20-0309-reasoning",
- "grok-4.20-reasoning", "grok-4.20-multi-agent-0309":
- return true
- default:
- return false
- }
- }
- // grokUnsupportedRecursiveFields 定义 Grok 平台(Responses 和 Chat Completions)不支持的字段
- var grokUnsupportedRecursiveFields = map[string]struct{}{
- "external_web_access": {},
- }
- // sanitizeGrokUnsupportedFields 递归移除 Grok 平台不支持的字段
- // 适用于 Responses API 和 Chat Completions API
- func sanitizeGrokUnsupportedFields(body []byte) ([]byte, error) {
- if !bytes.Contains(body, []byte(`"external_web_access"`)) {
- return body, nil
- }
- var payload any
- if err := decodeOpenAIJSONUseNumber(body, &payload); err != nil {
- return nil, err
- }
- if !deleteJSONFields(payload, grokUnsupportedRecursiveFields) {
- return body, nil
- }
- return marshalOpenAIUpstreamJSON(payload)
- }
- // sanitizeGrokResponsesUnsupportedFields 保留旧函数名作为别名,向后兼容
- var sanitizeGrokResponsesUnsupportedFields = sanitizeGrokUnsupportedFields
- func deleteJSONFields(value any, fields map[string]struct{}) bool {
- switch typed := value.(type) {
- case map[string]any:
- changed := false
- for field := range fields {
- if _, ok := typed[field]; ok {
- delete(typed, field)
- changed = true
- }
- }
- for _, child := range typed {
- if deleteJSONFields(child, fields) {
- changed = true
- }
- }
- return changed
- case []any:
- changed := false
- for _, child := range typed {
- if deleteJSONFields(child, fields) {
- changed = true
- }
- }
- return changed
- default:
- return false
- }
- }
- // additional_tools is a Codex/Responses Lite private input carrier. xAI's
- // Responses schema rejects the carrier itself, but accepts supported tools at
- // the top level. Preserve top-level order, append newly discovered tools in
- // carrier order, then let sanitizeGrokResponsesTools filter unsupported types.
- func sanitizeGrokResponsesInput(body []byte) ([]byte, error) {
- if !bytes.Contains(body, []byte(`"additional_tools"`)) {
- return body, nil
- }
- input := gjson.GetBytes(body, "input")
- if !input.Exists() || !input.IsArray() {
- return body, nil
- }
- rawItems := input.Array()
- filtered := make([]json.RawMessage, 0, len(rawItems))
- topLevelTools := gjson.GetBytes(body, "tools")
- mergedTools := make([]json.RawMessage, 0)
- seenTools := make(map[string]struct{})
- appendTool := func(tool gjson.Result) bool {
- key := grokResponsesToolDedupKey(tool)
- if _, exists := seenTools[key]; exists {
- return false
- }
- seenTools[key] = struct{}{}
- mergedTools = append(mergedTools, json.RawMessage(tool.Raw))
- return true
- }
- if topLevelTools.IsArray() {
- for _, tool := range topLevelTools.Array() {
- seenTools[grokResponsesToolDedupKey(tool)] = struct{}{}
- mergedTools = append(mergedTools, json.RawMessage(tool.Raw))
- }
- }
- promoted := false
- for _, item := range rawItems {
- if strings.TrimSpace(item.Get("type").String()) == "additional_tools" {
- tools := item.Get("tools")
- if tools.IsArray() {
- for _, tool := range tools.Array() {
- if appendTool(tool) {
- promoted = true
- }
- }
- }
- continue
- }
- filtered = append(filtered, json.RawMessage(item.Raw))
- }
- if len(filtered) == len(rawItems) {
- return body, nil
- }
- encoded, err := json.Marshal(filtered)
- if err != nil {
- return nil, err
- }
- body, err = sjson.SetRawBytes(body, "input", encoded)
- if err != nil || !promoted {
- return body, err
- }
- encodedTools, err := json.Marshal(mergedTools)
- if err != nil {
- return nil, err
- }
- return sjson.SetRawBytes(body, "tools", encodedTools)
- }
- // An inline input_image is already visible to Grok. Keeping Codex's local
- // view_image tool in the same turn can make Grok announce a tool call without
- // actually calling it, so remove only that redundant automatic choice.
- func stripRedundantGrokViewImageTool(body []byte) ([]byte, error) {
- input := gjson.GetBytes(body, "input")
- if !input.IsArray() {
- return body, nil
- }
- items := input.Array()
- if len(items) == 0 {
- return body, nil
- }
- current := items[len(items)-1]
- if strings.TrimSpace(current.Get("role").String()) != "user" ||
- !openAIJSONValueMayContainImageInput(current) {
- return body, nil
- }
- toolChoice := gjson.GetBytes(body, "tool_choice")
- if toolChoice.IsObject() && strings.TrimSpace(toolChoice.Get("type").String()) == "function" {
- choiceName := strings.TrimSpace(toolChoice.Get("name").String())
- if choiceName == "" {
- choiceName = strings.TrimSpace(toolChoice.Get("function.name").String())
- }
- if choiceName == "view_image" {
- return body, nil
- }
- }
- tools := gjson.GetBytes(body, "tools")
- if !tools.IsArray() {
- return body, nil
- }
- filtered := make([]json.RawMessage, 0, len(tools.Array()))
- changed := false
- for _, tool := range tools.Array() {
- if strings.TrimSpace(tool.Get("type").String()) == "function" &&
- strings.TrimSpace(tool.Get("name").String()) == "view_image" {
- changed = true
- continue
- }
- filtered = append(filtered, json.RawMessage(tool.Raw))
- }
- if !changed {
- return body, nil
- }
- if len(filtered) == 0 && strings.TrimSpace(toolChoice.String()) == "required" {
- return body, nil
- }
- if len(filtered) == 0 {
- out, err := sjson.DeleteBytes(body, "tools")
- if err != nil {
- return nil, err
- }
- return sjson.DeleteBytes(out, "parallel_tool_calls")
- }
- encoded, err := json.Marshal(filtered)
- if err != nil {
- return nil, err
- }
- return sjson.SetRawBytes(body, "tools", encoded)
- }
- func grokResponsesToolDedupKey(tool gjson.Result) string {
- toolType := strings.TrimSpace(tool.Get("type").String())
- if toolType != "" {
- if name := strings.TrimSpace(tool.Get("name").String()); name != "" {
- return "type:" + toolType + "\x00name:" + name
- }
- if toolType == "mcp" {
- if label := strings.TrimSpace(tool.Get("server_label").String()); label != "" {
- return "type:mcp\x00server_label:" + label
- }
- }
- }
- return "json:" + normalizeCompatSeedJSON(json.RawMessage(tool.Raw))
- }
- // sanitizeGrokReasoningNullContent drops explicit JSON nulls from Responses
- // input items. xAI's untagged ModelInput decoder 422s on those fields.
- // Compaction items stay unmodified per the compact contract.
- func sanitizeGrokReasoningNullContent(body []byte) ([]byte, error) {
- input := gjson.GetBytes(body, "input")
- if !input.Exists() || (!input.IsArray() && !input.IsObject()) {
- return body, nil
- }
- var decoded map[string]any
- decoder := json.NewDecoder(bytes.NewReader(body))
- decoder.UseNumber()
- if err := decoder.Decode(&decoded); err != nil {
- return body, nil
- }
- rawInput, ok := decoded["input"]
- if !ok {
- return body, nil
- }
- cleaned, changed := stripExplicitNullsFromGrokInput(rawInput)
- if !changed {
- return body, nil
- }
- decoded["input"] = cleaned
- out, err := marshalOpenAIUpstreamJSON(decoded)
- if err != nil {
- return nil, err
- }
- return out, nil
- }
- func stripExplicitNullsFromGrokInput(value any) (any, bool) {
- switch node := value.(type) {
- case []any:
- changed := false
- for i, item := range node {
- itemMap, ok := item.(map[string]any)
- if !ok {
- next, childChanged := stripExplicitNullsFromGrokInput(item)
- if childChanged {
- node[i] = next
- changed = true
- }
- continue
- }
- if isOpenAICompactionType(stringValue(itemMap["type"])) {
- continue
- }
- next, childChanged := stripExplicitNullsFromJSONObject(itemMap)
- if childChanged {
- node[i] = next
- changed = true
- }
- }
- return node, changed
- case map[string]any:
- if isOpenAICompactionType(stringValue(node["type"])) {
- return node, false
- }
- return stripExplicitNullsFromJSONObject(node)
- default:
- return value, false
- }
- }
- func stripExplicitNullsFromJSONObject(node map[string]any) (map[string]any, bool) {
- if node == nil {
- return node, false
- }
- changed := false
- for key, child := range node {
- if child == nil {
- delete(node, key)
- changed = true
- continue
- }
- switch typed := child.(type) {
- case map[string]any:
- next, childChanged := stripExplicitNullsFromJSONObject(typed)
- if childChanged {
- node[key] = next
- changed = true
- }
- case []any:
- next, childChanged := stripExplicitNullsFromGrokInput(typed)
- if childChanged {
- node[key] = next
- changed = true
- }
- }
- }
- return node, changed
- }
- var grokResponsesSupportedToolTypes = map[string]struct{}{
- "code_execution": {},
- "code_interpreter": {},
- "collections_search": {},
- "file_search": {},
- "function": {},
- "mcp": {},
- "shell": {},
- "web_search": {},
- "x_search": {},
- }
- const grokSafeFunctionParameters = `{"type":"object","properties":{},"additionalProperties":true}`
- func sanitizeGrokResponsesTools(body []byte) ([]byte, error) {
- tools := gjson.GetBytes(body, "tools")
- if !tools.Exists() {
- return deleteGrokOrphanToolControls(body)
- }
- if !tools.IsArray() {
- // xAI rejects tool_choice when tools is null/object. Drop the malformed
- // collection and any orphan tool controls instead of forwarding a pair
- // the Grok Responses endpoint cannot interpret.
- body, err := sjson.DeleteBytes(body, "tools")
- if err != nil {
- return nil, err
- }
- return deleteGrokOrphanToolControls(body)
- }
- rawTools := tools.Array()
- filteredTools := make([]json.RawMessage, 0, len(rawTools))
- toolsChanged := false
- for _, tool := range rawTools {
- toolType := strings.TrimSpace(tool.Get("type").String())
- if _, ok := grokResponsesSupportedToolTypes[toolType]; ok {
- raw := json.RawMessage(tool.Raw)
- if toolType == "function" && (!tool.Get("parameters").Exists() || tool.Get("parameters").Type == gjson.Null) {
- var payload map[string]any
- if err := decodeOpenAIJSONUseNumber(raw, &payload); err != nil {
- return nil, err
- }
- payload["parameters"] = map[string]any{"type": "object", "properties": map[string]any{}}
- encoded, err := marshalOpenAIUpstreamJSON(payload)
- if err != nil {
- return nil, err
- }
- raw = encoded
- toolsChanged = true
- } else if toolType == "function" && grokFunctionParametersHaveInvalidUnionRoot(tool.Get("parameters")) {
- var err error
- raw, err = sjson.SetRawBytes(raw, "parameters", []byte(grokSafeFunctionParameters))
- if err != nil {
- return nil, err
- }
- if strict := tool.Get("strict"); strict.Exists() && strict.Bool() {
- raw, err = sjson.SetBytes(raw, "strict", false)
- if err != nil {
- return nil, err
- }
- }
- toolsChanged = true
- }
- filteredTools = append(filteredTools, raw)
- }
- }
- if !grokRawToolsContainType(filteredTools, "tool_search") {
- for index, raw := range filteredTools {
- if !gjson.GetBytes(raw, "defer_loading").Exists() {
- continue
- }
- cleaned, deleteErr := sjson.DeleteBytes(raw, "defer_loading")
- if deleteErr != nil {
- return nil, deleteErr
- }
- filteredTools[index] = cleaned
- toolsChanged = true
- }
- }
- var err error
- if len(filteredTools) != len(rawTools) || toolsChanged {
- if len(filteredTools) == 0 {
- body, err = sjson.DeleteBytes(body, "tools")
- } else {
- var encoded []byte
- encoded, err = json.Marshal(filteredTools)
- if err != nil {
- return nil, err
- }
- body, err = sjson.SetRawBytes(body, "tools", encoded)
- }
- if err != nil {
- return nil, err
- }
- }
- if len(filteredTools) == 0 {
- return deleteGrokOrphanToolControls(body)
- }
- toolChoice := gjson.GetBytes(body, "tool_choice")
- if !toolChoice.Exists() {
- return body, nil
- }
- if shouldDropGrokToolChoice(toolChoice, filteredTools) {
- body, err = sjson.DeleteBytes(body, "tool_choice")
- if err != nil {
- return nil, err
- }
- }
- return body, nil
- }
- func grokFunctionParametersHaveInvalidUnionRoot(parameters gjson.Result) bool {
- if !parameters.Exists() || !parameters.IsObject() {
- return false
- }
- for _, keyword := range []string{"anyOf", "oneOf"} {
- branches := parameters.Get(keyword)
- if !branches.IsArray() {
- continue
- }
- values := branches.Array()
- if len(values) == 0 {
- continue
- }
- for _, branch := range values {
- if !strings.EqualFold(strings.TrimSpace(branch.Get("type").String()), "object") {
- return true
- }
- }
- }
- return false
- }
- func grokRawToolsContainType(tools []json.RawMessage, want string) bool {
- for _, tool := range tools {
- if strings.TrimSpace(gjson.GetBytes(tool, "type").String()) == want {
- return true
- }
- }
- return false
- }
- func deleteGrokOrphanToolControls(body []byte) ([]byte, error) {
- var err error
- for _, field := range []string{"tool_choice", "parallel_tool_calls"} {
- if !gjson.GetBytes(body, field).Exists() {
- continue
- }
- body, err = sjson.DeleteBytes(body, field)
- if err != nil {
- return nil, err
- }
- }
- return body, nil
- }
- func shouldDropGrokToolChoice(toolChoice gjson.Result, tools []json.RawMessage) bool {
- if len(tools) == 0 {
- return true
- }
- if !toolChoice.IsObject() {
- return false
- }
- choiceType := strings.TrimSpace(toolChoice.Get("type").String())
- if choiceType == "" {
- return false
- }
- if _, ok := grokResponsesSupportedToolTypes[choiceType]; !ok {
- return true
- }
- if choiceType == "function" {
- choiceName := strings.TrimSpace(toolChoice.Get("name").String())
- if choiceName == "" {
- choiceName = strings.TrimSpace(toolChoice.Get("function.name").String())
- }
- if choiceName == "" {
- return false
- }
- for _, tool := range tools {
- var item struct {
- Type string `json:"type"`
- Name string `json:"name"`
- Function struct {
- Name string `json:"name"`
- } `json:"function"`
- }
- if err := json.Unmarshal(tool, &item); err != nil {
- continue
- }
- name := strings.TrimSpace(item.Name)
- if name == "" {
- name = strings.TrimSpace(item.Function.Name)
- }
- if strings.TrimSpace(item.Type) == "function" && name == choiceName {
- return false
- }
- }
- return true
- }
- return false
- }
- func (s *OpenAIGatewayService) bridgeGrokComposerImageInputs(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- body []byte,
- token string,
- ) ([]byte, OpenAIUsage, bool, error) {
- if !shouldBridgeGrokComposerImageInputs(body) {
- return body, OpenAIUsage{}, false, nil
- }
- var reqBody map[string]any
- if err := decodeOpenAIJSONUseNumber(body, &reqBody); err != nil {
- return body, OpenAIUsage{}, false, fmt.Errorf("parse grok composer image bridge request: %w", err)
- }
- imageURLs := collectGrokComposerImageURLs(reqBody)
- if len(imageURLs) == 0 {
- return body, OpenAIUsage{}, false, nil
- }
- descriptions := make([]string, 0, len(imageURLs))
- var bridgeUsage OpenAIUsage
- for index, imageURL := range imageURLs {
- description, usage, err := s.describeGrokComposerImage(ctx, c, account, token, imageURL, index+1)
- if err != nil {
- return body, bridgeUsage, false, err
- }
- descriptions = append(descriptions, description)
- addOpenAIUsage(&bridgeUsage, usage)
- }
- if !rewriteGrokComposerImagesAsText(reqBody, descriptions) {
- return body, bridgeUsage, false, nil
- }
- bridgedBody, err := marshalOpenAIUpstreamJSON(reqBody)
- if err != nil {
- return body, bridgeUsage, false, fmt.Errorf("serialize grok composer image bridge request: %w", err)
- }
- return bridgedBody, bridgeUsage, true, nil
- }
- func shouldBridgeGrokComposerImageInputs(body []byte) bool {
- if len(body) == 0 || !isGrokComposerModel(gjson.GetBytes(body, "model").String()) {
- return false
- }
- messages := gjson.GetBytes(body, "messages")
- if !messages.Exists() {
- return false
- }
- return openAIJSONValueMayContainImageInput(messages)
- }
- func isGrokComposerModel(model string) bool {
- model = strings.TrimSpace(strings.ToLower(model))
- if model == "" {
- return false
- }
- if strings.Contains(model, "/") {
- parts := strings.Split(model, "/")
- model = strings.TrimSpace(parts[len(parts)-1])
- }
- return strings.Contains(model, "composer")
- }
- func collectGrokComposerImageURLs(reqBody map[string]any) []string {
- messages, ok := reqBody["messages"].([]any)
- if !ok {
- return nil
- }
- var imageURLs []string
- for _, msg := range messages {
- msgMap, ok := msg.(map[string]any)
- if !ok {
- continue
- }
- parts, ok := msgMap["content"].([]any)
- if !ok {
- continue
- }
- for _, part := range parts {
- if imageURL := grokComposerImageURLFromPart(part); imageURL != "" {
- imageURLs = append(imageURLs, imageURL)
- }
- }
- }
- return imageURLs
- }
- func grokComposerImageURLFromPart(part any) string {
- partMap, ok := part.(map[string]any)
- if !ok {
- return ""
- }
- if strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"]))) != "image_url" {
- return ""
- }
- switch imageURL := partMap["image_url"].(type) {
- case string:
- return normalizeGrokComposerImageURL(imageURL)
- case map[string]any:
- raw, _ := imageURL["url"].(string)
- return normalizeGrokComposerImageURL(raw)
- default:
- return ""
- }
- }
- func normalizeGrokComposerImageURL(raw string) string {
- trimmed := strings.TrimSpace(raw)
- if trimmed == "" || isEmptyBase64DataURI(trimmed) {
- return ""
- }
- return trimmed
- }
- func (s *OpenAIGatewayService) describeGrokComposerImage(
- ctx context.Context,
- c *gin.Context,
- account *Account,
- token string,
- imageURL string,
- index int,
- ) (string, OpenAIUsage, error) {
- body, err := buildGrokComposerImageDescriptionBody(imageURL, index)
- if err != nil {
- return "", OpenAIUsage{}, err
- }
- upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
- // Image-description probes are auxiliary requests, not conversation turns.
- // Do not bind them to the caller's Grok prompt-cache identity.
- upstreamReq, err := buildGrokResponsesRequest(upstreamCtx, c, account, body, token, "", s.cfg, s.settingService)
- releaseUpstreamCtx()
- if err != nil {
- return "", OpenAIUsage{}, fmt.Errorf("build grok composer image bridge request: %w", err)
- }
- proxyURL := ""
- if account.ProxyID != nil && account.Proxy != nil {
- proxyURL = account.Proxy.URL()
- }
- resp, err := s.doOpenAIUpstream(upstreamReq, proxyURL, account)
- if err != nil {
- return "", OpenAIUsage{}, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
- }
- defer func() { _ = resp.Body.Close() }()
- if resp.StatusCode >= 400 {
- respBody := s.readUpstreamErrorBody(resp)
- upstreamMsg := sanitizeUpstreamErrorMessage(extractUpstreamErrorMessage(respBody))
- if upstreamMsg == "" {
- upstreamMsg = fmt.Sprintf("xAI image bridge upstream returned status %d", resp.StatusCode)
- }
- kind := "http_error"
- if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
- kind = "failover"
- }
- appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
- ProxyID: opsUpstreamProxyID(account),
- ProxyName: opsUpstreamProxyName(account),
- Platform: account.Platform,
- AccountID: account.ID,
- AccountName: account.Name,
- UpstreamStatusCode: resp.StatusCode,
- UpstreamRequestID: firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id")),
- Kind: kind,
- Message: upstreamMsg,
- })
- s.handleGrokAccountUpstreamError(withGrokTeamRateLimitModel(ctx, grokComposerImageBridgeVisionModel), account, resp.StatusCode, resp.Header, respBody)
- if s.shouldFailoverGrokUpstreamError(resp.StatusCode, respBody) {
- retryable, retryDelay, retryDeadline, retryMax := grokSameAccountRetryMetadata(account, resp.StatusCode, respBody)
- return "", OpenAIUsage{}, &UpstreamFailoverError{
- StatusCode: resp.StatusCode,
- ResponseBody: respBody,
- ResponseHeaders: resp.Header.Clone(),
- RetryableOnSameAccount: retryable,
- RequestScopedTransient: retryable && resp.StatusCode == http.StatusTooManyRequests,
- SameAccountRetryDelay: retryDelay,
- SameAccountRetryDeadline: retryDeadline,
- SameAccountRetryMax: retryMax,
- }
- }
- return "", OpenAIUsage{}, fmt.Errorf("grok composer image bridge upstream error: %s", upstreamMsg)
- }
- s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, grokComposerImageBridgeVisionModel), account, resp.Header, resp.StatusCode)
- respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, nil)
- if err != nil {
- return "", OpenAIUsage{}, fmt.Errorf("read grok composer image bridge response: %w", err)
- }
- var parsed apicompat.ResponsesResponse
- if err := json.Unmarshal(respBody, &parsed); err != nil {
- return "", OpenAIUsage{}, fmt.Errorf("parse grok composer image bridge response: %w", err)
- }
- description := strings.TrimSpace(grokResponsesOutputText(&parsed))
- if description == "" {
- return "", copyOpenAIUsageFromResponsesUsage(parsed.Usage), fmt.Errorf("grok composer image bridge returned empty description")
- }
- return description, copyOpenAIUsageFromResponsesUsage(parsed.Usage), nil
- }
- func buildGrokComposerImageDescriptionBody(imageURL string, index int) ([]byte, error) {
- 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)
- req := map[string]any{
- "model": grokComposerImageBridgeVisionModel,
- "stream": false,
- "store": false,
- "max_output_tokens": grokComposerImageBridgeMaxOutputTokens,
- "input": []any{
- map[string]any{
- "type": "message",
- "role": "user",
- "content": []any{
- map[string]any{"type": "input_text", "text": prompt},
- map[string]any{"type": "input_image", "image_url": imageURL},
- },
- },
- },
- }
- return marshalOpenAIUpstreamJSON(req)
- }
- func grokResponsesOutputText(resp *apicompat.ResponsesResponse) string {
- if resp == nil {
- return ""
- }
- var parts []string
- for _, output := range resp.Output {
- for _, content := range output.Content {
- if content.Type == "output_text" || content.Type == "text" || content.Type == "input_text" {
- if text := strings.TrimSpace(content.Text); text != "" {
- parts = append(parts, text)
- }
- }
- }
- }
- return strings.Join(parts, "\n\n")
- }
- func rewriteGrokComposerImagesAsText(reqBody map[string]any, descriptions []string) bool {
- messages, ok := reqBody["messages"].([]any)
- if !ok {
- return false
- }
- imageIndex := 0
- changed := false
- for _, msg := range messages {
- msgMap, ok := msg.(map[string]any)
- if !ok {
- continue
- }
- parts, ok := msgMap["content"].([]any)
- if !ok {
- continue
- }
- var textParts []string
- messageChanged := false
- for _, part := range parts {
- if imageURL := grokComposerImageURLFromPart(part); imageURL != "" {
- if imageIndex < len(descriptions) {
- textParts = append(textParts, fmt.Sprintf("Image %d description: %s", imageIndex+1, strings.TrimSpace(descriptions[imageIndex])))
- }
- imageIndex++
- messageChanged = true
- continue
- }
- if text := grokComposerTextFromPart(part); text != "" {
- textParts = append(textParts, text)
- }
- }
- if messageChanged {
- msgMap["content"] = strings.Join(textParts, "\n\n")
- changed = true
- }
- }
- return changed
- }
- func grokComposerTextFromPart(part any) string {
- partMap, ok := part.(map[string]any)
- if !ok {
- return ""
- }
- partType := strings.TrimSpace(strings.ToLower(fmt.Sprint(partMap["type"])))
- switch partType {
- case "text", "input_text":
- text, _ := partMap["text"].(string)
- return strings.TrimSpace(text)
- default:
- return ""
- }
- }
- func addOpenAIUsage(dst *OpenAIUsage, usage OpenAIUsage) {
- if dst == nil {
- return
- }
- dst.InputTokens += usage.InputTokens
- dst.ImageInputTokens += usage.ImageInputTokens
- dst.OutputTokens += usage.OutputTokens
- dst.CacheCreationInputTokens += usage.CacheCreationInputTokens
- dst.CacheReadInputTokens += usage.CacheReadInputTokens
- dst.ImageOutputTokens += usage.ImageOutputTokens
- }
- func buildGrokResponsesRequest(ctx context.Context, c *gin.Context, account *Account, body []byte, token, cacheIdentity string, cfg *config.Config, settings ...*SettingService) (*http.Request, error) {
- targetURL, err := buildGrokResponsesURL(account, cfg, settings...)
- if err != nil {
- return nil, err
- }
- req, err := http.NewRequestWithContext(ctx, http.MethodPost, targetURL, bytes.NewReader(body))
- if err != nil {
- return nil, err
- }
- req = req.WithContext(WithHTTPUpstreamProfile(req.Context(), HTTPUpstreamProfileGrok))
- req.Header.Set("Authorization", "Bearer "+token)
- req.Header.Set("Content-Type", "application/json")
- req.Header.Set("Accept", "application/json, text/event-stream")
- if account.IsGrokOAuth() {
- applyGrokCLIHeaders(req.Header)
- }
- applyGrokCacheHeaders(req.Header, cacheIdentity)
- if c != nil {
- if v := c.GetHeader("OpenAI-Beta"); strings.TrimSpace(v) != "" {
- req.Header.Set("OpenAI-Beta", v)
- }
- }
- // 账号级请求头覆写最后应用,使配置值优先于上面的内置默认头;
- // 打到官方 CLI 网关时身份头仍由共享传输层最终强制。
- account.ApplyHeaderOverrides(req.Header)
- return req, nil
- }
- // applyGrokCLIHeaders identifies subscription traffic as a supported Grok CLI
- // version. The CLI gateway rejects otherwise valid OAuth requests without it.
- // Identity pins come from package xai so service-layer headers match the final
- // transport rewrite on cli-chat-proxy.grok.com.
- func applyGrokCLIHeaders(headers http.Header) {
- if headers == nil {
- return
- }
- version := xai.ResolveCLIVersion()
- headers.Set("User-Agent", xai.CLIUserAgent(version))
- headers.Set("X-Grok-Client-Version", version)
- headers.Set("x-grok-client-version", version)
- headers.Set("x-grok-client-identifier", xai.CLIClientIdentifier)
- // Historical mode value expected by some unit tests / older CLI probes.
- headers.Set("X-Grok-Client-Mode", "interactive")
- }
- func (s *OpenAIGatewayService) updateGrokUsageSnapshot(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot) {
- s.updateGrokUsageSnapshotWithRateLimit(ctx, account, snapshot, true)
- }
- func (s *OpenAIGatewayService) updateGrokUsageSnapshotWithRateLimit(ctx context.Context, account *Account, snapshot *xai.QuotaSnapshot, installRateLimit bool) {
- if s == nil || account == nil || account.ID <= 0 || snapshot == nil {
- return
- }
- accountID := account.ID
- now := time.Now()
- resetAt, hasActiveLimit := grokRateLimitResetAtForAccount(account, snapshot, now)
- if hasActiveLimit {
- normalizeGrokExhaustedWindowResets(snapshot, resetAt, now)
- }
- recovery := isSuccessfulGrokRateLimitRecovery(account, snapshot)
- critical := snapshot.StatusCode == http.StatusTooManyRequests || hasActiveLimit || recovery
- if s.codexSnapshotThrottle != nil {
- allowed := s.codexSnapshotThrottle.Allow(accountID, now)
- if !critical && !allowed {
- return
- }
- }
- updates := map[string]any{
- grokQuotaSnapshotExtraKey: snapshot,
- }
- // Also derive the scheduling-threshold extras (grok_sched_*) the evaluator
- // reads in grokThresholdCandidates. Without this writer the admin-configured
- // Grok auto-pause threshold could never fire (the read side was dead config).
- for k, v := range buildGrokSchedulerExtraUpdates(snapshot) {
- updates[k] = v
- }
- stateCtx := ctx
- if hasActiveLimit {
- var cancel context.CancelFunc
- stateCtx, cancel = openAIAccountStateContext(ctx)
- defer cancel()
- }
- // Account pointers on the request path are per-request copies (Redis/DB decode),
- // not a shared in-process cache. Mutating Extra here matches token refresh /
- // rate-limit writers; do not reuse the same *Account across goroutines.
- if account.Extra == nil {
- account.Extra = map[string]any{}
- }
- account.Extra[grokQuotaSnapshotExtraKey] = snapshot
- if s.accountRepo != nil {
- _ = s.accountRepo.UpdateExtra(stateCtx, accountID, updates)
- }
- // Error responses are reconciled by handleGrokAccountUpstreamError. Pool-mode
- // API keys retain the snapshot for observability but leave account health to
- // the upstream pool. Other accounts install the immediate runtime and durable
- // rate-limit state when the observed window is exhausted.
- if installRateLimit && hasActiveLimit && !account.IsPoolMode() {
- s.rateLimitGrok(stateCtx, account, resetAt)
- } else if recovery {
- clearGrokRateLimitAfterRecovery(stateCtx, s.accountRepo, account)
- }
- }
- func (s *OpenAIGatewayService) updateGrokUsageFromResponse(ctx context.Context, account *Account, headers http.Header, statusCode int) {
- snapshot := parseGrokQuotaSnapshot(headers, statusCode, time.Now())
- if snapshot != nil {
- stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
- s.updateGrokUsageSnapshot(ctx, account, snapshot)
- return
- }
- // Successful responses are recovery evidence even when the upstream omits
- // optional quota headers. Do not replace an informative stored snapshot with
- // an empty one; only clear the exact observed cooldown generation.
- recoverySnapshot := &xai.QuotaSnapshot{StatusCode: statusCode}
- if isSuccessfulGrokRateLimitRecovery(account, recoverySnapshot) {
- clearGrokRateLimitAfterRecovery(ctx, s.accountRepo, account)
- }
- }
- func parseGrokQuotaSnapshot(headers http.Header, statusCode int, now time.Time) *xai.QuotaSnapshot {
- snapshot := xai.ParseQuotaHeaders(headers, statusCode)
- if snapshot == nil && statusCode == http.StatusTooManyRequests {
- return &xai.QuotaSnapshot{
- StatusCode: statusCode,
- UpdatedAt: now.UTC().Format(time.RFC3339),
- }
- }
- return snapshot
- }
- func normalizeGrokExhaustedWindowResets(snapshot *xai.QuotaSnapshot, resetAt, now time.Time) {
- if snapshot == nil || !resetAt.After(now) {
- return
- }
- for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} {
- if window == nil || window.Remaining == nil || *window.Remaining > 0 {
- continue
- }
- candidate := time.Time{}
- if window.ResetUnix != nil && *window.ResetUnix > 0 {
- candidate = time.Unix(*window.ResetUnix, 0)
- } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil {
- candidate = parsed
- }
- if !candidate.After(now) {
- candidate = resetAt
- }
- resetUnix := candidate.Unix()
- window.ResetUnix = &resetUnix
- window.ResetAt = candidate.UTC().Format(time.RFC3339)
- }
- }
- func grokRateLimitResetAt(snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) {
- if snapshot == nil {
- return time.Time{}, false
- }
- // Retry-After is xAI's explicit retry boundary. Use the observation time so
- // a persisted snapshot does not start a fresh cooldown every time it is read.
- retryAfterExpired := false
- var resetAt time.Time
- if snapshot.RetryAfterSeconds != nil && *snapshot.RetryAfterSeconds > 0 {
- observedAt := now
- if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(snapshot.UpdatedAt)); err == nil {
- observedAt = parsed
- }
- retryAfterResetAt := observedAt.Add(time.Duration(*snapshot.RetryAfterSeconds) * time.Second)
- if retryAfterResetAt.After(now) {
- resetAt = retryAfterResetAt
- } else {
- retryAfterExpired = true
- }
- }
- exhausted := false
- for _, window := range []*xai.QuotaWindow{snapshot.Requests, snapshot.Tokens} {
- if window == nil || window.Remaining == nil || *window.Remaining > 0 {
- continue
- }
- exhausted = true
- candidate := time.Time{}
- if window.ResetUnix != nil && *window.ResetUnix > 0 {
- candidate = time.Unix(*window.ResetUnix, 0)
- } else if parsed, err := time.Parse(time.RFC3339, strings.TrimSpace(window.ResetAt)); err == nil {
- candidate = parsed
- }
- if candidate.After(now) && candidate.After(resetAt) {
- resetAt = candidate
- }
- }
- if !resetAt.IsZero() {
- return resetAt, true
- }
- // An observed Retry-After is an absolute boundary once combined with the
- // snapshot timestamp. Do not turn an expired persisted snapshot into a new
- // rolling fallback cooldown, but still allow a later explicit window reset.
- if retryAfterExpired {
- return time.Time{}, false
- }
- if exhausted || snapshot.StatusCode == http.StatusTooManyRequests {
- return now.Add(grokRateLimitFallbackCooldown), true
- }
- return time.Time{}, false
- }
- func grokRateLimitResetAtForAccount(account *Account, snapshot *xai.QuotaSnapshot, now time.Time) (time.Time, bool) {
- resetAt, limited := grokRateLimitResetAt(snapshot, now)
- if !limited || !isGrokOAuthAccount(account) || snapshot == nil || snapshot.StatusCode != http.StatusTooManyRequests {
- return resetAt, limited
- }
- if account.RateLimitedAt == nil || account.RateLimitResetAt == nil {
- return resetAt, true
- }
- previousResetAt := *account.RateLimitResetAt
- if previousResetAt.After(now) || now.Sub(previousResetAt) > grokRateLimitBackoffQuietPeriod {
- return resetAt, true
- }
- previousCooldown := previousResetAt.Sub(*account.RateLimitedAt)
- if previousCooldown <= 0 {
- return resetAt, true
- }
- adaptiveCooldown := grokRateLimitRepeatCooldown
- switch {
- case previousCooldown >= grokRateLimitSustainedCooldown:
- adaptiveCooldown = grokRateLimitMaxAdaptiveCooldown
- case previousCooldown >= grokRateLimitRepeatCooldown:
- adaptiveCooldown = grokRateLimitSustainedCooldown
- }
- adaptiveResetAt := now.Add(adaptiveCooldown)
- if adaptiveResetAt.After(resetAt) {
- resetAt = adaptiveResetAt
- }
- return resetAt, true
- }
- func normalizeGrokRateLimitResetAt(account *Account, resetAt, now time.Time) time.Time {
- if !resetAt.After(now) {
- resetAt = now.Add(grokRateLimitFallbackCooldown)
- }
- if account != nil && account.RateLimitResetAt != nil && account.RateLimitResetAt.After(resetAt) {
- resetAt = *account.RateLimitResetAt
- }
- return resetAt
- }
- type grokRateLimitExtendingRepository interface {
- SetRateLimitedIfLater(ctx context.Context, id int64, resetAt time.Time) error
- }
- type grokRateLimitRecoveryRepository interface {
- ClearRateLimitIfObserved(ctx context.Context, id int64, observedLimitedAt, observedResetAt time.Time) (bool, error)
- }
- func isSuccessfulGrokRateLimitRecovery(account *Account, snapshot *xai.QuotaSnapshot) bool {
- return isGrokOAuthAccount(account) &&
- account.RateLimitedAt != nil &&
- account.RateLimitResetAt != nil &&
- snapshot != nil &&
- snapshot.StatusCode >= http.StatusOK &&
- snapshot.StatusCode < http.StatusMultipleChoices
- }
- func clearGrokRateLimitAfterRecovery(ctx context.Context, repo AccountRepository, account *Account) {
- if repo == nil || account == nil || account.RateLimitedAt == nil || account.RateLimitResetAt == nil || ctx.Err() != nil {
- return
- }
- recoveryRepo, ok := repo.(grokRateLimitRecoveryRepository)
- if !ok {
- return
- }
- _, err := recoveryRepo.ClearRateLimitIfObserved(ctx, account.ID, *account.RateLimitedAt, *account.RateLimitResetAt)
- if err != nil {
- slog.Warn("grok_rate_limit_recovery_clear_failed", "account_id", account.ID, "error", err)
- }
- }
- func persistGrokRateLimit(ctx context.Context, repo AccountRepository, account *Account, resetAt time.Time) {
- if repo == nil || account == nil || account.ID <= 0 {
- return
- }
- resetAt = normalizeGrokRateLimitResetAt(account, resetAt, time.Now())
- stateCtx, cancel := openAIAccountStateContext(ctx)
- defer cancel()
- var err error
- if extendingRepo, ok := repo.(grokRateLimitExtendingRepository); ok {
- err = extendingRepo.SetRateLimitedIfLater(stateCtx, account.ID, resetAt)
- } else {
- err = repo.SetRateLimited(stateCtx, account.ID, resetAt)
- }
- if err != nil {
- slog.Warn("persist_grok_rate_limit_failed", "account_id", account.ID, "reset_at", resetAt.UTC(), "error", err)
- }
- }
- func (s *OpenAIGatewayService) rateLimitGrok(ctx context.Context, account *Account, resetAt time.Time) {
- if s == nil || account == nil {
- return
- }
- now := time.Now()
- resetAt = normalizeGrokRateLimitResetAt(account, resetAt, now)
- runtimeUntil := resetAt
- if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(runtimeUntil) {
- runtimeUntil = *account.TempUnschedulableUntil
- }
- s.BlockAccountScheduling(account, runtimeUntil, "429")
- persistGrokRateLimit(ctx, s.accountRepo, account, resetAt)
- // Propagate a short team+model cool so sibling OAuth accounts on the same
- // xAI team skip the hot model without waiting for each to hit 429 alone.
- // Model is taken from the latest request context when available; empty is a
- // no-op inside markGrokTeamModelRateLimit.
- if model, _ := ctx.Value(grokTeamRateLimitModelContextKey{}).(string); model != "" {
- markGrokTeamModelRateLimit(account, model, resolveGrokTeamRateLimitUntil(resetAt, now))
- }
- }
- // buildGrokSchedulerExtraUpdates derives the grok_sched_* scheduling snapshot
- // (utilization percent + reset time) consumed by EvaluateAccountSchedulingThreshold.
- // Utilization is the most-constrained of the requests/tokens windows.
- func buildGrokSchedulerExtraUpdates(snapshot *xai.QuotaSnapshot) map[string]any {
- if snapshot == nil {
- return nil
- }
- util, reset, ok := grokSnapshotUtilization(snapshot)
- if !ok {
- return nil
- }
- updates := map[string]any{
- "grok_sched_utilization": util,
- "grok_sched_usage_updated_at": time.Now().UTC().Format(time.RFC3339),
- }
- if reset != nil {
- // 防御:调度阈值暂停时长由 grok_sched_reset_at 决定。若上游返回脏的
- // reset 头(例如把相对毫秒 "6000" 误当相对秒解析出 ~33h 的未来时刻),
- // 不设上限会把耗尽账号长时间锁死。xAI 配额窗口不会超过一天,因此对
- // 未来时刻做 grokMaxSchedulingResetHorizon 钳制;过去/无效值直接不写。
- now := time.Now()
- if reset.After(now) {
- capped := *reset
- if horizon := now.Add(grokMaxSchedulingResetHorizon); capped.After(horizon) {
- capped = horizon
- }
- updates["grok_sched_reset_at"] = capped.UTC().Format(time.RFC3339)
- }
- }
- return updates
- }
- // grokSnapshotUtilization returns the highest window utilization (0-100) across
- // the requests/tokens quota windows and the reset time of that window.
- func grokSnapshotUtilization(snapshot *xai.QuotaSnapshot) (float64, *time.Time, bool) {
- if snapshot == nil {
- return 0, nil, false
- }
- best := -1.0
- var bestReset *time.Time
- consider := func(window *xai.QuotaWindow) {
- if window == nil || window.Limit == nil || *window.Limit <= 0 || window.Remaining == nil {
- return
- }
- remaining := *window.Remaining
- if remaining < 0 {
- remaining = 0
- }
- util := (1 - float64(remaining)/float64(*window.Limit)) * 100
- if util < 0 {
- util = 0
- }
- if util > 100 {
- util = 100
- }
- if util > best {
- best = util
- if window.ResetUnix != nil {
- t := time.Unix(*window.ResetUnix, 0).UTC()
- bestReset = &t
- } else {
- bestReset = nil
- }
- }
- }
- consider(snapshot.Requests)
- consider(snapshot.Tokens)
- if best < 0 {
- return 0, nil, false
- }
- return best, bestReset, true
- }
- // grokMaxSchedulingResetHorizon bounds how far into the future a Grok
- // scheduling-threshold pause (grok_sched_reset_at) may be set, so a malformed
- // upstream reset header can't park an over-threshold account for days. xAI quota
- // windows do not exceed ~a day.
- const grokMaxSchedulingResetHorizon = 25 * time.Hour
- // grokTeamRateLimitModelContextKey carries the upstream model for team cools.
- type grokTeamRateLimitModelContextKey struct{}
- // withGrokTeamRateLimitModel attaches the upstream model name for rate-limit
- // side effects (team+model cool). Safe when model is empty.
- func withGrokTeamRateLimitModel(ctx context.Context, model string) context.Context {
- model = strings.TrimSpace(model)
- if model == "" || ctx == nil {
- return ctx
- }
- return context.WithValue(ctx, grokTeamRateLimitModelContextKey{}, model)
- }
- func grokRequestedModelFromCtx(ctx context.Context) string {
- if ctx == nil {
- return ""
- }
- model, _ := ctx.Value(grokTeamRateLimitModelContextKey{}).(string)
- return strings.TrimSpace(model)
- }
- func persistGrokTransientModelCooldown(account *Account, decision GrokUpstreamFailureDecision) bool {
- if account == nil {
- return false
- }
- model := strings.TrimSpace(decision.Model)
- if model == "" {
- return false
- }
- cooldown := decision.Cooldown
- if cooldown <= 0 {
- cooldown = 3 * time.Minute
- }
- markGrokModelTransientBlock(account.ID, model, time.Now().Add(cooldown))
- return true
- }
- func (s *OpenAIGatewayService) handleGrokAccountUpstreamError(ctx context.Context, account *Account, statusCode int, headers http.Header, responseBody []byte) {
- if s == nil || account == nil {
- return
- }
- if isGrokContentPolicyRejection(statusCode, responseBody) {
- return
- }
- now := time.Now()
- decision := classifyGrokUpstreamFailure(statusCode, responseBody, grokRequestedModelFromCtx(ctx))
- snapshot := parseGrokQuotaSnapshot(headers, statusCode, now)
- stampGrokQuotaSnapshotForPlan(account, snapshot, grokRequestedModelFromCtx(ctx))
- // Capacity 429 is model pressure, not account quota exhaustion. Keep the
- // snapshot for observability but do not install account-level rate limiting;
- // the failover decision below applies a bounded model-scoped block instead.
- s.updateGrokUsageSnapshotWithRateLimit(ctx, account, snapshot, decision.Class != GrokFailureModelCapacity)
- // Body-first free-usage / empty / billing / capacity must run before the
- // status switch so non-429 free-usage bodies still cool the account.
- // Pool-mode still skips durable mutation unless an explicit temp rule matches.
- if decision.ShouldCooldown && decision.Class != GrokFailureNone && decision.Class != GrokFailureRateLimit {
- if account.IsPoolMode() {
- // Allow configured temp rules (403) below; skip default body cools.
- } else {
- // A free-tier exhaustion message describes a rolling usage window. Use
- // an upstream absolute reset (or Retry-After) when available; otherwise
- // apply only a short probe cooldown. Never start a fabricated 24h window
- // at the instant this error was received.
- if decision.Class == GrokFailureFreeUsage {
- if resetAt, limited := grokRateLimitResetAtForAccount(account, parseGrokQuotaSnapshot(headers, statusCode, now), now); limited && resetAt.After(now) {
- if decision.Model != "" && isGrokModelSpecificFreeUsage(strings.ToLower(decision.Reason), decision.Model) {
- markGrokModelQuotaBlock(account.ID, decision.Model, resetAt)
- return
- }
- s.rateLimitGrok(ctx, account, resetAt)
- return
- }
- }
- if s.applyGrokUpstreamFailureDecision(ctx, account, decision) {
- return
- }
- }
- }
- if statusCode == http.StatusForbidden && s.applyGrokForbiddenPolicy(ctx, account, responseBody) {
- return
- }
- if account.IsPoolMode() {
- slog.Info("grok_pool_mode_error_state_skipped", "account_id", account.ID, "status_code", statusCode)
- return
- }
- switch statusCode {
- case http.StatusUnauthorized:
- s.tempUnscheduleGrok(ctx, account, 10*time.Minute, "grok credentials unauthorized")
- case http.StatusPaymentRequired:
- // 402 without a body-classified billing decision: keep the legacy 30m cool.
- s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok payment required")
- case http.StatusForbidden:
- // Spending-limit already handled by body classifier when phrasing matches.
- if isGrokSpendingLimitError(responseBody) {
- s.rateLimitGrok(ctx, account, grokSpendingLimitResetAt(account, time.Now()))
- return
- }
- s.tempUnscheduleGrok(ctx, account, 30*time.Minute, "grok access or entitlement denied")
- case http.StatusTooManyRequests:
- // updateGrokUsageSnapshot installs rate-limit state for non-pool accounts.
- // Free-usage 429 was already cooled above via body classification.
- default:
- if statusCode >= 500 {
- s.tempUnscheduleGrok(ctx, account, 2*time.Minute, "grok upstream temporary error")
- }
- }
- }
- // isGrokSpendingLimitError detects xAI billing exhaustion bodies (often 403, sometimes 402).
- func isGrokSpendingLimitError(responseBody []byte) bool {
- if len(responseBody) == 0 {
- return false
- }
- code := strings.ToLower(strings.TrimSpace(firstNonEmpty(
- gjson.GetBytes(responseBody, "code").String(),
- gjson.GetBytes(responseBody, "error.code").String(),
- )))
- if code == "personal-team-blocked:spending-limit" {
- return true
- }
- message := strings.ToLower(strings.TrimSpace(firstNonEmpty(
- gjson.GetBytes(responseBody, "error").String(),
- gjson.GetBytes(responseBody, "error.message").String(),
- gjson.GetBytes(responseBody, "message").String(),
- )))
- return strings.Contains(message, "spending limit") ||
- strings.Contains(message, "run out of credits")
- }
- func (s *OpenAIGatewayService) tempUnscheduleGrok(ctx context.Context, account *Account, cooldown time.Duration, reason string) {
- if s == nil || account == nil {
- return
- }
- until := time.Now().Add(cooldown)
- if account.TempUnschedulableUntil != nil && account.TempUnschedulableUntil.After(until) {
- until = *account.TempUnschedulableUntil
- }
- s.BlockAccountScheduling(account, until, reason)
- if s.accountRepo != nil {
- stateCtx, cancel := openAIAccountStateContext(ctx)
- defer cancel()
- _ = s.accountRepo.SetTempUnschedulable(stateCtx, account.ID, until, reason)
- }
- }
|