openai_gateway_grok_test.go 159 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414141514161417141814191420142114221423142414251426142714281429143014311432143314341435143614371438143914401441144214431444144514461447144814491450145114521453145414551456145714581459146014611462146314641465146614671468146914701471147214731474147514761477147814791480148114821483148414851486148714881489149014911492149314941495149614971498149915001501150215031504150515061507150815091510151115121513151415151516151715181519152015211522152315241525152615271528152915301531153215331534153515361537153815391540154115421543154415451546154715481549155015511552155315541555155615571558155915601561156215631564156515661567156815691570157115721573157415751576157715781579158015811582158315841585158615871588158915901591159215931594159515961597159815991600160116021603160416051606160716081609161016111612161316141615161616171618161916201621162216231624162516261627162816291630163116321633163416351636163716381639164016411642164316441645164616471648164916501651165216531654165516561657165816591660166116621663166416651666166716681669167016711672167316741675167616771678167916801681168216831684168516861687168816891690169116921693169416951696169716981699170017011702170317041705170617071708170917101711171217131714171517161717171817191720172117221723172417251726172717281729173017311732173317341735173617371738173917401741174217431744174517461747174817491750175117521753175417551756175717581759176017611762176317641765176617671768176917701771177217731774177517761777177817791780178117821783178417851786178717881789179017911792179317941795179617971798179918001801180218031804180518061807180818091810181118121813181418151816181718181819182018211822182318241825182618271828182918301831183218331834183518361837183818391840184118421843184418451846184718481849185018511852185318541855185618571858185918601861186218631864186518661867186818691870187118721873187418751876187718781879188018811882188318841885188618871888188918901891189218931894189518961897189818991900190119021903190419051906190719081909191019111912191319141915191619171918191919201921192219231924192519261927192819291930193119321933193419351936193719381939194019411942194319441945194619471948194919501951195219531954195519561957195819591960196119621963196419651966196719681969197019711972197319741975197619771978197919801981198219831984198519861987198819891990199119921993199419951996199719981999200020012002200320042005200620072008200920102011201220132014201520162017201820192020202120222023202420252026202720282029203020312032203320342035203620372038203920402041204220432044204520462047204820492050205120522053205420552056205720582059206020612062206320642065206620672068206920702071207220732074207520762077207820792080208120822083208420852086208720882089209020912092209320942095209620972098209921002101210221032104210521062107210821092110211121122113211421152116211721182119212021212122212321242125212621272128212921302131213221332134213521362137213821392140214121422143214421452146214721482149215021512152215321542155215621572158215921602161216221632164216521662167216821692170217121722173217421752176217721782179218021812182218321842185218621872188218921902191219221932194219521962197219821992200220122022203220422052206220722082209221022112212221322142215221622172218221922202221222222232224222522262227222822292230223122322233223422352236223722382239224022412242224322442245224622472248224922502251225222532254225522562257225822592260226122622263226422652266226722682269227022712272227322742275227622772278227922802281228222832284228522862287228822892290229122922293229422952296229722982299230023012302230323042305230623072308230923102311231223132314231523162317231823192320232123222323232423252326232723282329233023312332233323342335233623372338233923402341234223432344234523462347234823492350235123522353235423552356235723582359236023612362236323642365236623672368236923702371237223732374237523762377237823792380238123822383238423852386238723882389239023912392239323942395239623972398239924002401240224032404240524062407240824092410241124122413241424152416241724182419242024212422242324242425242624272428242924302431243224332434243524362437243824392440244124422443244424452446244724482449245024512452245324542455245624572458245924602461246224632464246524662467246824692470247124722473247424752476247724782479248024812482248324842485248624872488248924902491249224932494249524962497249824992500250125022503250425052506250725082509251025112512251325142515251625172518251925202521252225232524252525262527252825292530253125322533253425352536253725382539254025412542254325442545254625472548254925502551255225532554255525562557255825592560256125622563256425652566256725682569257025712572257325742575257625772578257925802581258225832584258525862587258825892590259125922593259425952596259725982599260026012602260326042605260626072608260926102611261226132614261526162617261826192620262126222623262426252626262726282629263026312632263326342635263626372638263926402641264226432644264526462647264826492650265126522653265426552656265726582659266026612662266326642665266626672668266926702671267226732674267526762677267826792680268126822683268426852686268726882689269026912692269326942695269626972698269927002701270227032704270527062707270827092710271127122713271427152716271727182719272027212722272327242725272627272728272927302731273227332734273527362737273827392740274127422743274427452746274727482749275027512752275327542755275627572758275927602761276227632764276527662767276827692770277127722773277427752776277727782779278027812782278327842785278627872788278927902791279227932794279527962797279827992800280128022803280428052806280728082809281028112812281328142815281628172818281928202821282228232824282528262827282828292830283128322833283428352836283728382839284028412842284328442845284628472848284928502851285228532854285528562857285828592860286128622863286428652866286728682869287028712872287328742875287628772878287928802881288228832884288528862887288828892890289128922893289428952896289728982899290029012902290329042905290629072908290929102911291229132914291529162917291829192920292129222923292429252926292729282929293029312932293329342935293629372938293929402941294229432944294529462947294829492950295129522953295429552956295729582959296029612962296329642965296629672968296929702971297229732974297529762977297829792980298129822983298429852986298729882989299029912992299329942995299629972998299930003001300230033004300530063007300830093010301130123013301430153016301730183019302030213022302330243025302630273028302930303031303230333034303530363037303830393040304130423043304430453046304730483049305030513052305330543055305630573058305930603061306230633064306530663067306830693070307130723073307430753076307730783079308030813082308330843085308630873088308930903091309230933094309530963097309830993100310131023103310431053106310731083109311031113112311331143115311631173118311931203121312231233124312531263127312831293130313131323133313431353136313731383139314031413142314331443145314631473148314931503151315231533154315531563157315831593160316131623163316431653166316731683169317031713172317331743175317631773178317931803181318231833184318531863187318831893190319131923193319431953196319731983199320032013202320332043205320632073208320932103211321232133214321532163217321832193220322132223223322432253226322732283229323032313232323332343235323632373238323932403241324232433244324532463247324832493250325132523253325432553256325732583259326032613262326332643265326632673268326932703271327232733274327532763277327832793280328132823283328432853286328732883289329032913292329332943295329632973298329933003301330233033304330533063307330833093310331133123313331433153316331733183319332033213322332333243325332633273328332933303331333233333334333533363337333833393340334133423343334433453346334733483349335033513352335333543355335633573358335933603361336233633364336533663367336833693370337133723373337433753376337733783379338033813382338333843385338633873388338933903391339233933394339533963397339833993400340134023403340434053406340734083409341034113412341334143415341634173418341934203421342234233424342534263427342834293430343134323433343434353436343734383439344034413442344334443445344634473448344934503451345234533454345534563457345834593460346134623463346434653466346734683469347034713472347334743475347634773478347934803481348234833484348534863487348834893490349134923493349434953496349734983499350035013502350335043505350635073508350935103511351235133514351535163517351835193520352135223523352435253526352735283529353035313532353335343535353635373538353935403541354235433544354535463547354835493550355135523553355435553556355735583559356035613562356335643565356635673568356935703571357235733574357535763577357835793580358135823583358435853586358735883589359035913592359335943595359635973598359936003601360236033604360536063607360836093610361136123613361436153616361736183619362036213622362336243625362636273628362936303631363236333634363536363637363836393640364136423643364436453646364736483649365036513652365336543655365636573658365936603661366236633664366536663667366836693670367136723673367436753676367736783679368036813682368336843685368636873688368936903691369236933694369536963697369836993700370137023703370437053706370737083709371037113712371337143715371637173718371937203721372237233724372537263727372837293730373137323733373437353736373737383739374037413742374337443745374637473748374937503751375237533754375537563757375837593760376137623763376437653766376737683769377037713772377337743775377637773778377937803781378237833784378537863787378837893790379137923793379437953796379737983799
  1. //go:build unit
  2. package service
  3. import (
  4. "bytes"
  5. "context"
  6. "encoding/json"
  7. "fmt"
  8. "io"
  9. "mime/multipart"
  10. "net/http"
  11. "net/http/httptest"
  12. "net/textproto"
  13. "strings"
  14. "testing"
  15. "time"
  16. "github.com/Wei-Shaw/sub2api/internal/config"
  17. "github.com/Wei-Shaw/sub2api/internal/pkg/xai"
  18. "github.com/gin-gonic/gin"
  19. "github.com/stretchr/testify/require"
  20. "github.com/tidwall/gjson"
  21. )
  22. func TestPatchGrokResponsesBodySetsMappedModelAndDropsUnsupportedFields(t *testing.T) {
  23. t.Parallel()
  24. body := []byte(`{
  25. "model": "grok",
  26. "input": "hello",
  27. "prompt_cache_retention": "24h",
  28. "safety_identifier": "user-1",
  29. "reasoning": {"effort": "high"}
  30. }`)
  31. patched, err := patchGrokResponsesBody(body, "grok-4.3")
  32. require.NoError(t, err)
  33. require.True(t, json.Valid(patched))
  34. require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
  35. require.False(t, gjson.GetBytes(patched, "prompt_cache_retention").Exists())
  36. require.False(t, gjson.GetBytes(patched, "safety_identifier").Exists())
  37. require.Equal(t, "high", gjson.GetBytes(patched, "reasoning.effort").String())
  38. }
  39. func TestPatchGrokResponsesBodyDropsRedundantViewImageForCurrentInlineImage(t *testing.T) {
  40. t.Parallel()
  41. tests := []struct {
  42. name string
  43. body string
  44. }{
  45. {
  46. name: "top-level tools",
  47. body: `{
  48. "model":"grok-4.6",
  49. "input":[{"type":"message","role":"user","content":[
  50. {"type":"input_text","text":"What text is in this image?"},
  51. {"type":"input_image","image_url":"data:image/png;base64,AA=="}
  52. ]}],
  53. "tools":[
  54. {"type":"function","name":"view_image","parameters":{"type":"object"}},
  55. {"type":"function","name":"shell_command","parameters":{"type":"object"}}
  56. ]
  57. }`,
  58. },
  59. {
  60. name: "Responses Lite additional tools",
  61. body: `{
  62. "model":"grok-4.6",
  63. "input":[
  64. {"type":"additional_tools","role":"developer","tools":[
  65. {"type":"function","name":"view_image","parameters":{"type":"object"}},
  66. {"type":"function","name":"shell_command","parameters":{"type":"object"}}
  67. ]},
  68. {"type":"message","role":"user","content":[
  69. {"type":"input_text","text":"What text is in this image?"},
  70. {"type":"input_image","image_url":"data:image/png;base64,AA=="}
  71. ]}
  72. ]
  73. }`,
  74. },
  75. }
  76. for _, tt := range tests {
  77. t.Run(tt.name, func(t *testing.T) {
  78. t.Parallel()
  79. patched, err := patchGrokResponsesBody([]byte(tt.body), "grok-4.6")
  80. require.NoError(t, err)
  81. require.False(t, gjson.GetBytes(patched, `tools.#(name=="view_image")`).Exists())
  82. require.Equal(t, "shell_command", gjson.GetBytes(patched, "tools.0.name").String())
  83. })
  84. }
  85. }
  86. func TestPatchGrokResponsesBodyKeepsNonRedundantViewImage(t *testing.T) {
  87. t.Parallel()
  88. tests := []struct {
  89. name string
  90. body string
  91. }{
  92. {
  93. name: "current turn has no inline image",
  94. body: `{"input":[{"role":"user","content":[{"type":"input_text","text":"Inspect a local image"}]}],"tools":[{"type":"function","name":"view_image"}]}`,
  95. },
  96. {
  97. name: "inline image is only historical",
  98. body: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]},{"role":"assistant","content":[{"type":"output_text","text":"Done"}]},{"role":"user","content":[{"type":"input_text","text":"Inspect another local image"}]}],"tools":[{"type":"function","name":"view_image"}]}`,
  99. },
  100. {
  101. name: "view image is explicitly selected",
  102. body: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}],"tools":[{"type":"function","name":"view_image"}],"tool_choice":{"type":"function","name":"view_image"}}`,
  103. },
  104. {
  105. name: "required with view image as the only tool",
  106. body: `{"input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}],"tools":[{"type":"function","name":"view_image"}],"tool_choice":"required"}`,
  107. },
  108. }
  109. for _, tt := range tests {
  110. t.Run(tt.name, func(t *testing.T) {
  111. t.Parallel()
  112. patched, err := patchGrokResponsesBody([]byte(tt.body), "grok-4.6")
  113. require.NoError(t, err)
  114. require.Equal(t, "view_image", gjson.GetBytes(patched, "tools.0.name").String())
  115. })
  116. }
  117. }
  118. func TestPatchGrokResponsesBodyDropsViewImageOnlyToolMetadata(t *testing.T) {
  119. t.Parallel()
  120. body := []byte(`{
  121. "input":[{"role":"user","content":[{"type":"input_image","image_url":"data:image/png;base64,AA=="}]}],
  122. "tools":[{"type":"function","name":"view_image"}],
  123. "tool_choice":"auto",
  124. "parallel_tool_calls":true
  125. }`)
  126. patched, err := patchGrokResponsesBody(body, "grok-4.6")
  127. require.NoError(t, err)
  128. require.False(t, gjson.GetBytes(patched, "tools").Exists())
  129. require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
  130. require.False(t, gjson.GetBytes(patched, "parallel_tool_calls").Exists())
  131. }
  132. func TestPatchGrokResponsesBodyPreservesGrokShellFunctionOutputImages(t *testing.T) {
  133. t.Parallel()
  134. body := []byte(`{
  135. "model":"grok-4.6",
  136. "input":[
  137. {"type":"function_call","call_id":"call_read","name":"read_file","arguments":"{}"},
  138. {"type":"function_call_output","call_id":"call_read","content":"Read image file: /tmp/example.png","images":[{"type":"image","url":"data:image/png;base64,QUE="}]}
  139. ]
  140. }`)
  141. patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.6")
  142. require.NoError(t, err)
  143. require.Equal(t, "function_call_output", gjson.GetBytes(patched, "input.1.type").String())
  144. require.Equal(t, "Read image file: /tmp/example.png", gjson.GetBytes(patched, "input.1.output").String())
  145. require.Equal(t, "message", gjson.GetBytes(patched, "input.2.type").String())
  146. require.Equal(t, "input_image", gjson.GetBytes(patched, "input.2.content.1.type").String())
  147. require.Equal(t, "data:image/png;base64,QUE=", gjson.GetBytes(patched, "input.2.content.1.image_url").String())
  148. }
  149. func TestPatchGrokResponsesBodyPreservesGrok105StructuredFunctionOutputImages(t *testing.T) {
  150. t.Parallel()
  151. body := []byte(`{
  152. "model":"grok-4.6",
  153. "input":[
  154. {"type":"function_call","call_id":"call_read","name":"read_file","arguments":"{}"},
  155. {"type":"function_call_output","call_id":"call_read","output":[
  156. {"type":"input_text","text":"Read image file: /tmp/example.png"},
  157. {"type":"input_image","detail":"auto","image_url":"data:image/png;base64,QUE="}
  158. ]}
  159. ]
  160. }`)
  161. patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.6")
  162. require.NoError(t, err)
  163. require.Equal(t, "Read image file: /tmp/example.png", gjson.GetBytes(patched, "input.1.output").String())
  164. require.Equal(t, "input_image", gjson.GetBytes(patched, "input.2.content.1.type").String())
  165. require.Equal(t, "data:image/png;base64,QUE=", gjson.GetBytes(patched, "input.2.content.1.image_url").String())
  166. }
  167. func TestPatchGrokResponsesBodySanitizesComposerReasoningParameters(t *testing.T) {
  168. t.Parallel()
  169. tests := []struct {
  170. name string
  171. upstreamModel string
  172. wantReasoning bool
  173. }{
  174. {name: "composer fast", upstreamModel: "grok-composer-2.5-fast"},
  175. {name: "composer shorthand", upstreamModel: "grok-composer"},
  176. {name: "composer legacy alias", upstreamModel: "composer-2.5"},
  177. {name: "provider-prefixed composer", upstreamModel: "xai/grok-composer-2.5-fast"},
  178. {name: "grok 4.5", upstreamModel: "grok-4.5", wantReasoning: true},
  179. {name: "grok 4.6", upstreamModel: "grok-4.6", wantReasoning: true},
  180. {name: "grok 4.6 latest", upstreamModel: "grok-4.6-latest", wantReasoning: true},
  181. }
  182. bodyTemplate := []byte(`{
  183. "model": "grok",
  184. "input": "hello",
  185. "reasoning": {"effort": "medium", "summary": "auto"},
  186. "reasoning_effort": "medium",
  187. "reasoningEffort": "medium"
  188. }`)
  189. for _, tt := range tests {
  190. t.Run(tt.name, func(t *testing.T) {
  191. patched, err := patchGrokResponsesBody(append([]byte(nil), bodyTemplate...), tt.upstreamModel)
  192. require.NoError(t, err)
  193. require.True(t, json.Valid(patched))
  194. require.Equal(t, tt.upstreamModel, gjson.GetBytes(patched, "model").String())
  195. if tt.wantReasoning {
  196. require.Equal(t, "medium", gjson.GetBytes(patched, "reasoning.effort").String())
  197. require.Equal(t, "medium", gjson.GetBytes(patched, "reasoning_effort").String())
  198. require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
  199. return
  200. }
  201. require.False(t, gjson.GetBytes(patched, "reasoning").Exists())
  202. require.False(t, gjson.GetBytes(patched, "reasoning_effort").Exists())
  203. require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
  204. })
  205. }
  206. }
  207. func TestExtractGrokResponsesReasoningEffortSupportsOpenAICompatibleField(t *testing.T) {
  208. t.Parallel()
  209. effort := extractOpenAIReasoningEffortFromBody(
  210. []byte(`{"model":"grok-4.3","reasoning_effort":"high"}`),
  211. "grok-4.3",
  212. )
  213. require.NotNil(t, effort)
  214. require.Equal(t, "high", *effort)
  215. }
  216. func TestPatchGrokResponsesBodyDropsGrok45ReasoningUnsupportedFields(t *testing.T) {
  217. t.Parallel()
  218. body := []byte(`{
  219. "model": "grok-latest",
  220. "input": "hello",
  221. "presence_penalty": 0.1,
  222. "presencePenalty": 0.2,
  223. "frequency_penalty": 0.3,
  224. "frequencyPenalty": 0.4,
  225. "stop": ["done"]
  226. }`)
  227. patched, err := patchGrokResponsesBody(body, "grok-4.5")
  228. require.NoError(t, err)
  229. require.True(t, json.Valid(patched))
  230. require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String())
  231. require.False(t, gjson.GetBytes(patched, "presence_penalty").Exists())
  232. require.False(t, gjson.GetBytes(patched, "presencePenalty").Exists())
  233. require.False(t, gjson.GetBytes(patched, "frequency_penalty").Exists())
  234. require.False(t, gjson.GetBytes(patched, "frequencyPenalty").Exists())
  235. require.False(t, gjson.GetBytes(patched, "stop").Exists())
  236. }
  237. func TestPatchGrokResponsesBodyKeepsPenaltyAndStopFieldsForNon45Models(t *testing.T) {
  238. t.Parallel()
  239. body := []byte(`{
  240. "model": "grok-4.3",
  241. "input": "hello",
  242. "presence_penalty": 0.1,
  243. "frequency_penalty": 0.2,
  244. "stop": ["done"]
  245. }`)
  246. patched, err := patchGrokResponsesBody(body, "grok-4.3")
  247. require.NoError(t, err)
  248. require.True(t, json.Valid(patched))
  249. require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
  250. require.Equal(t, 0.1, gjson.GetBytes(patched, "presence_penalty").Float())
  251. require.Equal(t, 0.2, gjson.GetBytes(patched, "frequency_penalty").Float())
  252. require.Len(t, gjson.GetBytes(patched, "stop").Array(), 1)
  253. }
  254. func TestPatchGrokResponsesBodyDropsLogprobsForGrok420Family(t *testing.T) {
  255. t.Parallel()
  256. body := []byte(`{"model":"grok-4.20-0309-reasoning","input":"hello","logprobs":true,"top_logprobs":5}`)
  257. patched, err := patchGrokResponsesBody(body, "grok-4.20-0309-reasoning")
  258. require.NoError(t, err)
  259. require.False(t, gjson.GetBytes(patched, "logprobs").Exists())
  260. require.False(t, gjson.GetBytes(patched, "top_logprobs").Exists())
  261. }
  262. func TestPatchGrokResponsesBodyNormalizesReasoningEffortAliases(t *testing.T) {
  263. t.Parallel()
  264. tests := []struct {
  265. name string
  266. body string
  267. upstreamModel string
  268. path string
  269. want string
  270. }{
  271. {name: "minimal nested", body: `{"input":"hi","reasoning":{"effort":"minimal"}}`, upstreamModel: "grok-4.5", path: "reasoning.effort", want: "low"},
  272. {name: "xhigh stays high for 4.5", body: `{"input":"hi","reasoning_effort":"xhigh"}`, upstreamModel: "grok-4.5", path: "reasoning_effort", want: "high"},
  273. {name: "xhigh nested for 4.6", body: `{"input":"hi","reasoning":{"effort":"xhigh"}}`, upstreamModel: "grok-4.6", path: "reasoning.effort", want: "xhigh"},
  274. {name: "xhigh snake for 4.6 latest", body: `{"input":"hi","reasoning_effort":"xhigh"}`, upstreamModel: "grok-4.6-latest", path: "reasoning_effort", want: "xhigh"},
  275. }
  276. for _, tt := range tests {
  277. t.Run(tt.name, func(t *testing.T) {
  278. patched, err := patchGrokResponsesBody([]byte(tt.body), tt.upstreamModel)
  279. require.NoError(t, err)
  280. require.Equal(t, tt.want, gjson.GetBytes(patched, tt.path).String(), string(patched))
  281. require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
  282. })
  283. }
  284. }
  285. func TestPatchGrokResponsesBodyAddsDefaultFunctionParameters(t *testing.T) {
  286. patched, err := patchGrokResponsesBody(
  287. []byte(`{"input":"hi","tools":[{"type":"function","name":"lookup","large_id":9007199254740993},{"type":"function","name":"wait","parameters":null}]}`),
  288. "grok-4.5",
  289. )
  290. require.NoError(t, err)
  291. for _, tool := range gjson.GetBytes(patched, "tools").Array() {
  292. require.Equal(t, "object", tool.Get("parameters.type").String(), string(patched))
  293. require.True(t, tool.Get("parameters.properties").IsObject(), string(patched))
  294. }
  295. require.Equal(t, "9007199254740993", gjson.GetBytes(patched, "tools.0.large_id").Raw)
  296. }
  297. func TestNormalizeGrokChatReasoningEffort(t *testing.T) {
  298. patched, err := normalizeGrokChatReasoningEffort([]byte(`{"reasoningEffort":"ultra"}`), "grok-4.3")
  299. require.NoError(t, err)
  300. require.Equal(t, "high", gjson.GetBytes(patched, "reasoning_effort").String())
  301. require.False(t, gjson.GetBytes(patched, "reasoningEffort").Exists())
  302. patched, err = normalizeGrokChatReasoningEffort([]byte(`{"reasoning_effort":"xhigh"}`), "grok-4.6")
  303. require.NoError(t, err)
  304. require.Equal(t, "xhigh", gjson.GetBytes(patched, "reasoning_effort").String())
  305. patched, err = normalizeGrokChatReasoningEffort([]byte(`{"reasoning_effort":"high"}`), "grok-composer-2.5-fast")
  306. require.NoError(t, err)
  307. require.False(t, gjson.GetBytes(patched, "reasoning_effort").Exists())
  308. }
  309. func TestPatchGrokResponsesBodyDropsNestedUnsupportedFields(t *testing.T) {
  310. t.Parallel()
  311. body := []byte(`{
  312. "model": "grok",
  313. "input": "hello",
  314. "external_web_access": true,
  315. "tools": [
  316. {"type": "function", "name": "kept_fn", "external_web_access": true, "parameters": {"type": "object", "properties": {"q": {"type": "string", "external_web_access": true}}}}
  317. ],
  318. "metadata": {"external_web_access": false, "large_id": 9007199254740993}
  319. }`)
  320. patched, err := patchGrokResponsesBody(body, "grok-4.3")
  321. require.NoError(t, err)
  322. require.True(t, json.Valid(patched))
  323. require.False(t, strings.Contains(string(patched), "external_web_access"))
  324. require.Equal(t, "kept_fn", gjson.GetBytes(patched, "tools.0.name").String())
  325. require.False(t, gjson.GetBytes(patched, "metadata").Exists())
  326. }
  327. func TestStripAnthropicThinkingSignaturesPreservesLargeIntegers(t *testing.T) {
  328. body := []byte(`{"messages":[{"role":"assistant","content":[{"type":"thinking","thinking":"work","signature":"opaque"}]}],"metadata":{"large_id":9007199254740993}}`)
  329. patched, changed := stripAnthropicThinkingSignatures(body)
  330. require.True(t, changed)
  331. require.False(t, gjson.GetBytes(patched, "messages.0.content.0.signature").Exists())
  332. require.Equal(t, "9007199254740993", gjson.GetBytes(patched, "metadata.large_id").Raw)
  333. }
  334. func TestPatchGrokResponsesBodyFlattensNamespaceTools(t *testing.T) {
  335. t.Parallel()
  336. body := []byte(`{
  337. "model": "grok",
  338. "input": "hello",
  339. "tools": [
  340. {"type": "namespace", "name": "functions", "tools": [{"type": "function", "name": "inner"}]},
  341. {"type": "function", "name": "kept_fn", "parameters": {"type": "object"}},
  342. {"type": "shell", "name": "kept_shell"}
  343. ],
  344. "tool_choice": {"type": "function", "namespace": "functions", "name": "inner"}
  345. }`)
  346. patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.3")
  347. require.NoError(t, err)
  348. require.True(t, json.Valid(patched))
  349. require.Equal(t, "grok-4.3", gjson.GetBytes(patched, "model").String())
  350. require.Len(t, gjson.GetBytes(patched, "tools").Array(), 3)
  351. require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists())
  352. require.True(t, gjson.GetBytes(patched, `tools.#(type=="function")`).Exists())
  353. require.True(t, gjson.GetBytes(patched, `tools.#(type=="shell")`).Exists())
  354. require.Equal(t, "functions__inner", gjson.GetBytes(patched, "tools.0.name").String())
  355. require.Equal(t, "functions__inner", gjson.GetBytes(patched, "tool_choice.name").String())
  356. require.False(t, gjson.GetBytes(patched, "tool_choice.namespace").Exists())
  357. }
  358. func TestPatchGrokResponsesBodyDropsToolChoiceWhenNoSupportedToolsRemain(t *testing.T) {
  359. t.Parallel()
  360. body := []byte(`{
  361. "model": "grok",
  362. "input": "hello",
  363. "tools": [
  364. {"type": "namespace", "namespace": "functions"},
  365. {"type": "image_generation", "model": "gpt-image-2"}
  366. ],
  367. "tool_choice": {"type": "namespace", "namespace": "functions"}
  368. }`)
  369. patched, err := patchGrokResponsesBody(body, "grok-4.3")
  370. require.NoError(t, err)
  371. require.True(t, json.Valid(patched))
  372. require.False(t, gjson.GetBytes(patched, "tools").Exists())
  373. require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
  374. }
  375. func TestSanitizeGrokResponsesToolsRemovesDeferredFlagsWithToolSearch(t *testing.T) {
  376. body := []byte(`{"tools":[{"type":"tool_search"},{"type":"function","name":"shell","defer_loading":true},{"type":"function","name":"apply_patch"}]}`)
  377. patched, err := sanitizeGrokResponsesTools(body)
  378. require.NoError(t, err)
  379. require.False(t, gjson.GetBytes(patched, `tools.#(type=="tool_search")`).Exists())
  380. require.False(t, gjson.GetBytes(patched, `tools.#(name=="shell").defer_loading`).Exists())
  381. require.True(t, gjson.GetBytes(patched, `tools.#(name=="apply_patch")`).Exists())
  382. }
  383. func TestSanitizeGrokResponsesToolsSimplifiesInvalidRootUnion(t *testing.T) {
  384. body := []byte(`{"tools":[
  385. {"type":"function","name":"mcp__codex_app__automation_update","strict":true,"parameters":{"oneOf":[{"type":"object","properties":{"id":{"type":"string"}}},{"type":"null"}]}},
  386. {"type":"function","name":"object_only","strict":true,"parameters":{"type":"object","anyOf":[{"type":"object","properties":{"a":{"type":"string"}}},{"type":"object","properties":{"b":{"type":"integer"}}}]}}
  387. ]}`)
  388. patched, err := sanitizeGrokResponsesTools(body)
  389. require.NoError(t, err)
  390. require.True(t, json.Valid(patched))
  391. mixed := gjson.GetBytes(patched, `tools.#(name=="mcp__codex_app__automation_update")`)
  392. require.Equal(t, "object", mixed.Get("parameters.type").String())
  393. require.True(t, mixed.Get("parameters.properties").IsObject())
  394. require.True(t, mixed.Get("parameters.additionalProperties").Bool())
  395. require.False(t, mixed.Get("parameters.oneOf").Exists())
  396. require.Equal(t, gjson.False, mixed.Get("strict").Type)
  397. objectOnly := gjson.GetBytes(patched, `tools.#(name=="object_only")`)
  398. require.True(t, objectOnly.Get("parameters.anyOf").Exists())
  399. require.Equal(t, gjson.True, objectOnly.Get("strict").Type)
  400. }
  401. func TestPatchGrokResponsesBodySimplifiesTypedInvalidRootUnion(t *testing.T) {
  402. body := []byte(`{
  403. "model":"grok-4.6",
  404. "metadata":{"session_id":"abc"},
  405. "tools":[{
  406. "type":"namespace",
  407. "name":"mcp__codex_app",
  408. "tools":[{
  409. "type":"function",
  410. "name":"automation_update",
  411. "strict":true,
  412. "parameters":{
  413. "type":"object",
  414. "oneOf":[{"$ref":"#/$defs/update"},{"type":"null"}],
  415. "$defs":{"update":{"type":"object","properties":{"id":{"type":"string"}}}}
  416. }
  417. }]
  418. }]
  419. }`)
  420. patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.6")
  421. require.NoError(t, err)
  422. require.True(t, json.Valid(patched))
  423. require.False(t, gjson.GetBytes(patched, "metadata").Exists())
  424. tool := gjson.GetBytes(patched, `tools.#(name=="mcp__codex_app__automation_update")`)
  425. require.Equal(t, "object", tool.Get("parameters.type").String())
  426. require.True(t, tool.Get("parameters.properties").IsObject())
  427. require.True(t, tool.Get("parameters.additionalProperties").Bool())
  428. require.False(t, tool.Get("parameters.oneOf").Exists())
  429. require.False(t, tool.Get("parameters.$defs").Exists())
  430. require.Equal(t, gjson.False, tool.Get("strict").Type)
  431. }
  432. func TestSanitizeGrokResponsesToolsKeepsToolChoiceOnlyWithSupportedTools(t *testing.T) {
  433. t.Parallel()
  434. tests := []struct {
  435. name string
  436. body string
  437. wantTools bool
  438. wantToolChoice bool
  439. }{
  440. {
  441. name: "missing tools with string tool choice",
  442. body: `{"input":"hello","tool_choice":"auto"}`,
  443. },
  444. {
  445. name: "missing tools with object tool choice",
  446. body: `{"input":"hello","tool_choice":{"type":"function","name":"lookup"}}`,
  447. },
  448. {
  449. name: "empty tools",
  450. body: `{"input":"hello","tools":[],"tool_choice":"auto"}`,
  451. wantTools: true,
  452. },
  453. {
  454. name: "all tools unsupported",
  455. body: `{"input":"hello","tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":"auto"}`,
  456. },
  457. {
  458. name: "supported tool",
  459. body: `{"input":"hello","tools":[{"type":"function","name":"lookup"}],"tool_choice":"auto"}`,
  460. wantTools: true,
  461. wantToolChoice: true,
  462. },
  463. {
  464. name: "malformed non-array tools are removed",
  465. body: `{"input":"hello","tools":{"type":"function","name":"lookup"},"tool_choice":"auto"}`,
  466. wantTools: false,
  467. wantToolChoice: false,
  468. },
  469. }
  470. for _, tt := range tests {
  471. t.Run(tt.name, func(t *testing.T) {
  472. patched, err := sanitizeGrokResponsesTools([]byte(tt.body))
  473. require.NoError(t, err)
  474. require.True(t, json.Valid(patched))
  475. require.Equal(t, tt.wantTools, gjson.GetBytes(patched, "tools").Exists())
  476. require.Equal(t, tt.wantToolChoice, gjson.GetBytes(patched, "tool_choice").Exists())
  477. if tt.wantToolChoice {
  478. require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String())
  479. }
  480. })
  481. }
  482. }
  483. func TestPatchGrokResponsesBodyPromotesCodexAdditionalTools(t *testing.T) {
  484. t.Parallel()
  485. body := []byte(`{
  486. "model": "grok",
  487. "tools": [
  488. {"type": "function", "name": "existing", "description": "top-level wins"},
  489. {"type": "web_search"}
  490. ],
  491. "tool_choice": "auto",
  492. "input": [
  493. {
  494. "type": "additional_tools",
  495. "role": "developer",
  496. "tools": [
  497. {"type": "function", "name": "existing", "description": "duplicate carrier definition"},
  498. {"type": "function", "name": "wait"},
  499. {"type": "web_search"},
  500. {"type": "shell"},
  501. {"type": "custom", "name": "apply_patch"},
  502. {"type": "namespace", "name": "collaboration"}
  503. ]
  504. },
  505. {
  506. "type": "message",
  507. "role": "developer",
  508. "content": [{"type": "input_text", "text": "system prompt"}]
  509. },
  510. {
  511. "type": "message",
  512. "role": "user",
  513. "content": [{"type": "input_text", "text": "hello"}]
  514. }
  515. ]
  516. }`)
  517. patched, _, err := patchGrokResponsesBodyWithClientTools(body, "grok-4.5")
  518. require.NoError(t, err)
  519. require.True(t, json.Valid(patched))
  520. require.Equal(t, "grok-4.5", gjson.GetBytes(patched, "model").String())
  521. require.Equal(t, 2, len(gjson.GetBytes(patched, "input").Array()))
  522. require.False(t, gjson.GetBytes(patched, `input.#(type=="additional_tools")`).Exists())
  523. tools := gjson.GetBytes(patched, "tools").Array()
  524. require.Len(t, tools, 5)
  525. require.Equal(t, "existing", tools[0].Get("name").String())
  526. require.Equal(t, "top-level wins", tools[0].Get("description").String())
  527. require.Equal(t, "web_search", tools[1].Get("type").String())
  528. require.Equal(t, "wait", tools[2].Get("name").String())
  529. require.Equal(t, "shell", tools[3].Get("type").String())
  530. require.Equal(t, "function", tools[4].Get("type").String())
  531. require.Equal(t, "apply_patch", tools[4].Get("name").String())
  532. require.Equal(t, "string", tools[4].Get("parameters.properties.input.type").String())
  533. require.False(t, gjson.GetBytes(patched, `tools.#(type=="custom")`).Exists())
  534. require.False(t, gjson.GetBytes(patched, `tools.#(type=="namespace")`).Exists())
  535. require.Equal(t, "auto", gjson.GetBytes(patched, "tool_choice").String())
  536. require.Equal(t, "developer", gjson.GetBytes(patched, "input.0.role").String())
  537. require.Equal(t, "system prompt", gjson.GetBytes(patched, "input.0.content.0.text").String())
  538. require.Equal(t, "user", gjson.GetBytes(patched, "input.1.role").String())
  539. require.Equal(t, "hello", gjson.GetBytes(patched, "input.1.content.0.text").String())
  540. }
  541. func TestForwardGrokResponsesCodexAdditionalToolsUsesMixedCacheIntent(t *testing.T) {
  542. gin.SetMode(gin.TestMode)
  543. body := []byte(`{
  544. "model":"grok",
  545. "stream":false,
  546. "prompt_cache_key":"codex-session",
  547. "input":[
  548. {"type":"additional_tools","role":"developer","tools":[
  549. {"type":"function","name":"lookup","description":"look up a key","parameters":{"type":"object"}},
  550. {"type":"function","name":"web_search","description":"search","parameters":{"type":"object"}},
  551. {"type":"custom","name":"apply_patch"},
  552. {"type":"namespace","name":"collaboration"}
  553. ]},
  554. {"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}
  555. ]
  556. }`)
  557. recorder := httptest.NewRecorder()
  558. c, _ := gin.CreateTestContext(recorder)
  559. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  560. c.Request.Header.Set(grokClientToolCacheOptInHeader, "prefer-cache")
  561. c.Set("api_key", &APIKey{ID: 4501})
  562. account := healthyGrokOAuthGatewayTestAccount(4501, "access-token")
  563. account.Credentials["subscription_tier"] = "free"
  564. repo := &grokQuotaAccountRepo{
  565. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  566. accountsByID: map[int64]*Account{account.ID: account},
  567. },
  568. }
  569. upstream := &httpUpstreamRecorder{resp: &http.Response{
  570. StatusCode: http.StatusOK,
  571. Header: http.Header{"Content-Type": []string{"application/json"}},
  572. Body: io.NopCloser(strings.NewReader(`{
  573. "id":"resp_codex_lite","object":"response","model":"grok-4.5","status":"completed",
  574. "output":[],"usage":{"input_tokens":10,"output_tokens":1}
  575. }`)),
  576. }}
  577. svc := &OpenAIGatewayService{
  578. httpUpstream: upstream,
  579. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  580. accountRepo: repo,
  581. }
  582. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
  583. require.NoError(t, err)
  584. require.NotNil(t, result)
  585. require.Equal(t, "resp_codex_lite", result.ResponseID)
  586. require.False(t, gjson.GetBytes(upstream.lastBody, `input.#(type=="additional_tools")`).Exists())
  587. tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
  588. require.Len(t, tools, 4)
  589. require.Equal(t, "function", tools[0].Get("type").String())
  590. require.Equal(t, "lookup", tools[0].Get("name").String())
  591. require.Equal(t, "web_search", tools[1].Get("type").String())
  592. require.Equal(t, "function", tools[2].Get("type").String())
  593. require.Equal(t, "apply_patch", tools[2].Get("name").String())
  594. require.Equal(t, "x_search", tools[3].Get("type").String())
  595. require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
  596. require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="custom")`).Exists())
  597. require.False(t, gjson.GetBytes(upstream.lastBody, `tools.#(type=="namespace")`).Exists())
  598. identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
  599. require.NotEmpty(t, identity)
  600. require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
  601. require.Empty(t, upstream.lastReq.Header.Get(grokClientToolCacheOptInHeader))
  602. }
  603. func TestForwardGrokResponsesClaudeDesktopClientToolsUseCacheRoute(t *testing.T) {
  604. gin.SetMode(gin.TestMode)
  605. firstBody := []byte(`{
  606. "model":"grok","stream":false,"instructions":"You are Claude Desktop.",
  607. "tools":[
  608. {"type":"function","name":"Read","parameters":{"type":"object"}},
  609. {"type":"function","name":"Edit","parameters":{"type":"object"}},
  610. {"type":"function","name":"WebSearch","parameters":{"type":"object"}},
  611. {"type":"function","name":"mcp__workspace__bash","parameters":{"type":"object"}}
  612. ],
  613. "input":[{"role":"user","content":[{"type":"input_text","text":"first turn"}]}]
  614. }`)
  615. secondBody := []byte(`{
  616. "model":"grok","stream":false,"instructions":"You are Claude Desktop.",
  617. "tools":[
  618. {"type":"function","name":"Read","parameters":{"type":"object"}},
  619. {"type":"function","name":"Edit","parameters":{"type":"object"}},
  620. {"type":"function","name":"WebSearch","parameters":{"type":"object"}},
  621. {"type":"function","name":"mcp__workspace__bash","parameters":{"type":"object"}}
  622. ],
  623. "input":[
  624. {"role":"user","content":[{"type":"input_text","text":"first turn"}]},
  625. {"role":"assistant","content":[{"type":"output_text","text":"first answer"}]},
  626. {"role":"user","content":[{"type":"input_text","text":"second turn"}]}
  627. ]
  628. }`)
  629. account := healthyGrokOAuthGatewayTestAccount(4504, "access-token")
  630. account.Credentials["subscription_tier"] = "free"
  631. repo := &grokQuotaAccountRepo{
  632. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  633. accountsByID: map[int64]*Account{account.ID: account},
  634. },
  635. }
  636. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  637. {
  638. StatusCode: http.StatusOK,
  639. Header: http.Header{"Content-Type": []string{"application/json"}},
  640. Body: io.NopCloser(strings.NewReader(`{
  641. "id":"resp_claude_desktop_1","object":"response","model":"grok-4.5","status":"completed",
  642. "output":[],"usage":{"input_tokens":30000,"output_tokens":10,"input_tokens_details":{"cached_tokens":0}}
  643. }`)),
  644. },
  645. {
  646. StatusCode: http.StatusOK,
  647. Header: http.Header{"Content-Type": []string{"application/json"}},
  648. Body: io.NopCloser(strings.NewReader(`{
  649. "id":"resp_claude_desktop_2","object":"response","model":"grok-4.5","status":"completed",
  650. "output":[],"usage":{"input_tokens":30100,"output_tokens":12,"input_tokens_details":{"cached_tokens":28672}}
  651. }`)),
  652. },
  653. }}
  654. svc := &OpenAIGatewayService{
  655. httpUpstream: upstream,
  656. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  657. accountRepo: repo,
  658. }
  659. newContext := func(body []byte) *gin.Context {
  660. recorder := httptest.NewRecorder()
  661. c, _ := gin.CreateTestContext(recorder)
  662. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  663. c.Request.Header.Set("User-Agent", "claude-cli/2.1.215 (external, claude-desktop-3p, agent-sdk/0.3.215)")
  664. c.Request.Header.Set("X-App", "cli")
  665. c.Request.Header.Set("anthropic-client-platform", "desktop_app")
  666. c.Request.Header.Set("X-Claude-Code-Session-Id", "claude-desktop-session")
  667. c.Set("api_key", &APIKey{ID: 4504})
  668. return c
  669. }
  670. first, err := svc.forwardGrokResponses(context.Background(), newContext(firstBody), account, firstBody, "grok", false, time.Now())
  671. require.NoError(t, err)
  672. second, err := svc.forwardGrokResponses(context.Background(), newContext(secondBody), account, secondBody, "grok", false, time.Now())
  673. require.NoError(t, err)
  674. require.Equal(t, 0, first.Usage.CacheReadInputTokens)
  675. require.Equal(t, 28672, second.Usage.CacheReadInputTokens)
  676. require.Len(t, upstream.bodies, 2)
  677. require.Len(t, upstream.requests, 2)
  678. for i := range upstream.bodies {
  679. tools := gjson.GetBytes(upstream.bodies[i], "tools").Array()
  680. require.Len(t, tools, 6)
  681. require.Equal(t, "Read", tools[0].Get("name").String())
  682. require.Equal(t, "Edit", tools[1].Get("name").String())
  683. require.Equal(t, "WebSearch", tools[2].Get("name").String())
  684. require.Equal(t, "mcp__workspace__bash", tools[3].Get("name").String())
  685. require.Equal(t, "web_search", tools[4].Get("type").String())
  686. require.Equal(t, "x_search", tools[5].Get("type").String())
  687. require.False(t, gjson.GetBytes(upstream.bodies[i], "tool_choice").Exists())
  688. require.Empty(t, upstream.requests[i].Header.Get("X-App"))
  689. require.Empty(t, upstream.requests[i].Header.Get("anthropic-client-platform"))
  690. require.Empty(t, upstream.requests[i].Header.Get("X-Claude-Code-Session-Id"))
  691. }
  692. firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
  693. secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
  694. require.NotEmpty(t, firstIdentity)
  695. require.Equal(t, firstIdentity, secondIdentity)
  696. require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
  697. require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
  698. }
  699. func TestGrokResponsesCacheIdentityIncludesPromotedCodexTools(t *testing.T) {
  700. c := newGrokCacheTestContext(4503)
  701. lookupBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"lookup","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`)
  702. readBody := []byte(`{"model":"grok","input":[{"type":"additional_tools","tools":[{"type":"function","name":"read_file","parameters":{"type":"object"}}]},{"type":"message","role":"user","content":"same prompt"}]}`)
  703. patchedLookup, err := patchGrokResponsesBody(lookupBody, "grok-4.5")
  704. require.NoError(t, err)
  705. patchedRead, err := patchGrokResponsesBody(readBody, "grok-4.5")
  706. require.NoError(t, err)
  707. lookupIdentity := resolveGrokCacheIdentity(c, patchedLookup, "", "grok-4.5")
  708. readIdentity := resolveGrokCacheIdentity(c, patchedRead, "", "grok-4.5")
  709. require.NotEmpty(t, lookupIdentity)
  710. require.NotEmpty(t, readIdentity)
  711. require.NotEqual(t, lookupIdentity, readIdentity)
  712. }
  713. func TestCodexUnsupportedAdditionalToolsDoNotBecomeToolFreeCacheIntent(t *testing.T) {
  714. body := []byte(`{
  715. "model":"grok","tool_choice":"auto",
  716. "input":[
  717. {"type":"additional_tools","role":"developer","tools":[
  718. {"type":"custom","name":"apply_patch"},
  719. {"type":"namespace","name":"collaboration"}
  720. ]},
  721. {"type":"message","role":"user","content":"hello"}
  722. ]
  723. }`)
  724. patched, err := patchGrokResponsesBody(body, "grok-4.5")
  725. require.NoError(t, err)
  726. require.False(t, gjson.GetBytes(patched, "tools").Exists())
  727. require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
  728. mixedCacheIntent := patched
  729. patched, err = applyGrokResponsesCacheIdentity(patched, body, "isolated-id", true)
  730. require.NoError(t, err)
  731. account := healthyGrokOAuthGatewayTestAccount(4502, "access-token")
  732. account.Credentials["subscription_tier"] = "free"
  733. patched, err = applyGrokFreeRequestToolCacheRoute(nil, patched, mixedCacheIntent, account, "isolated-id")
  734. require.NoError(t, err)
  735. require.False(t, gjson.GetBytes(patched, "tools").Exists())
  736. require.False(t, gjson.GetBytes(patched, "tool_choice").Exists())
  737. require.Equal(t, "isolated-id", gjson.GetBytes(patched, "prompt_cache_key").String())
  738. }
  739. func TestBuildGrokResponsesRequestUsesAccountBaseURLAndBearerToken(t *testing.T) {
  740. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  741. account := &Account{
  742. Platform: PlatformGrok,
  743. Type: AccountTypeOAuth,
  744. Credentials: map[string]any{
  745. "base_url": "https://xai.test/v1/",
  746. },
  747. }
  748. req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "isolated-cache-id", nil)
  749. require.NoError(t, err)
  750. require.Equal(t, http.MethodPost, req.Method)
  751. require.Equal(t, "https://xai.test/v1/responses", req.URL.String())
  752. require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
  753. require.Equal(t, "application/json", req.Header.Get("Content-Type"))
  754. require.Contains(t, req.Header.Get("Accept"), "text/event-stream")
  755. require.Equal(t, grokCLIVersion, req.Header.Get("X-Grok-Client-Version"))
  756. require.Equal(t, "isolated-cache-id", req.Header.Get(grokConversationIDHeader))
  757. data, err := io.ReadAll(req.Body)
  758. require.NoError(t, err)
  759. require.Equal(t, `{"model":"grok-4.3"}`, strings.TrimSpace(string(data)))
  760. }
  761. func TestBuildGrokCompactRequestBodyUsesResponsesCompactionTurn(t *testing.T) {
  762. body := []byte(`{"model":"grok-4.5","input":[{"type":"message","role":"user","content":[{"type":"input_text","text":"hello"}]}],"tools":[{"type":"function","name":"shell"}],"stream":true}`)
  763. patched, err := buildGrokCompactRequestBody(body)
  764. require.NoError(t, err)
  765. require.False(t, gjson.GetBytes(patched, "stream").Bool())
  766. require.False(t, gjson.GetBytes(patched, "store").Bool())
  767. require.Equal(t, "none", gjson.GetBytes(patched, "tool_choice").String())
  768. require.Equal(t, "reasoning.encrypted_content", gjson.GetBytes(patched, "include.0").String())
  769. require.Equal(t, "hello", gjson.GetBytes(patched, "input.0.content.0.text").String())
  770. prompt := gjson.GetBytes(patched, "input.1.content.0.text").String()
  771. require.Contains(t, prompt, "1. Primary Request and Intent")
  772. require.Contains(t, prompt, "9. Optional Next Step")
  773. require.Contains(t, prompt, "Respond with ONLY the <summary>...</summary> block")
  774. require.NotContains(t, prompt, "<summary_request>")
  775. }
  776. func TestConvertGrokResponseToOpenAICompact(t *testing.T) {
  777. body := []byte(`{
  778. "id":"resp_grok_1",
  779. "object":"response",
  780. "status":"completed",
  781. "model":"grok-4.5",
  782. "output":[
  783. {"id":"rs_1","type":"reasoning","summary":[],"encrypted_content":"grok-encrypted-state"},
  784. {"id":"msg_1","type":"message","role":"assistant","content":[{"type":"output_text","text":"summary text"}]}
  785. ],
  786. "usage":{"input_tokens":10,"output_tokens":4,"total_tokens":14}
  787. }`)
  788. converted, err := convertGrokResponseToOpenAICompact(body)
  789. require.NoError(t, err)
  790. require.Equal(t, "resp_grok_1", gjson.GetBytes(converted, "id").String())
  791. require.Len(t, gjson.GetBytes(converted, "output").Array(), 1)
  792. require.Equal(t, "compaction", gjson.GetBytes(converted, "output.0.type").String())
  793. require.Equal(t, "grok-encrypted-state", gjson.GetBytes(converted, "output.0.encrypted_content").String())
  794. require.Equal(t, "summary text", gjson.GetBytes(converted, "output.0.summary.0.text").String())
  795. require.Equal(t, int64(14), gjson.GetBytes(converted, "usage.total_tokens").Int())
  796. }
  797. func TestPatchGrokResponsesBodyRestoresCompactInput(t *testing.T) {
  798. body := []byte(`{
  799. "model":"grok-4.5",
  800. "input":[
  801. {"id":"cmp_1","type":"compaction","status":"completed","encrypted_content":"grok-encrypted-state","summary":[{"type":"summary_text","text":"summary text"}]},
  802. {"type":"message","role":"user","content":[{"type":"input_text","text":"continue"}]}
  803. ]
  804. }`)
  805. patched, err := patchGrokResponsesBody(body, "grok-4.5")
  806. require.NoError(t, err)
  807. require.Equal(t, "reasoning", gjson.GetBytes(patched, "input.0.type").String())
  808. require.Equal(t, "grok-encrypted-state", gjson.GetBytes(patched, "input.0.encrypted_content").String())
  809. require.Equal(t, "message", gjson.GetBytes(patched, "input.1.type").String())
  810. require.Contains(t, gjson.GetBytes(patched, "input.1.content.0.text").String(), "summary text")
  811. require.Equal(t, "continue", gjson.GetBytes(patched, "input.2.content.0.text").String())
  812. }
  813. func TestConvertGrokResponseToOpenAICompactRequiresEncryptedContent(t *testing.T) {
  814. _, err := convertGrokResponseToOpenAICompact([]byte(`{"output":[{"type":"message","content":[{"type":"output_text","text":"summary"}]}]}`))
  815. require.ErrorContains(t, err, "reasoning.encrypted_content")
  816. }
  817. func TestBuildGrokResponsesRequestAllowsPublicAPIKeyBaseURLByDefault(t *testing.T) {
  818. account := &Account{
  819. Platform: PlatformGrok,
  820. Type: AccountTypeAPIKey,
  821. Credentials: map[string]any{
  822. "base_url": "https://grok.example.test/v1/",
  823. },
  824. }
  825. req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "", nil)
  826. require.NoError(t, err)
  827. require.Equal(t, "https://grok.example.test/v1/responses", req.URL.String())
  828. require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
  829. require.Empty(t, req.Header.Get("X-Grok-Client-Version"))
  830. require.NotEqual(t, defaultGrokUpstreamUserAgent(), req.Header.Get("User-Agent"))
  831. }
  832. func TestBuildGrokResponsesRequestHonorsOAuthOfficialEndpointSwitch(t *testing.T) {
  833. t.Parallel()
  834. account := &Account{
  835. Platform: PlatformGrok,
  836. Type: AccountTypeOAuth,
  837. Credentials: map[string]any{
  838. "base_url": xai.DefaultBaseURL,
  839. },
  840. }
  841. req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "", nil)
  842. require.NoError(t, err)
  843. require.Equal(t, xai.DefaultBaseURL+"/responses", req.URL.String())
  844. }
  845. func TestBuildGrokResponsesRequestAppliesHeaderOverridesLast(t *testing.T) {
  846. t.Parallel()
  847. account := &Account{
  848. Platform: PlatformGrok,
  849. Type: AccountTypeOAuth,
  850. Credentials: map[string]any{
  851. "base_url": "https://relay.example.test/v1",
  852. "header_override_enabled": true,
  853. "header_overrides": map[string]any{
  854. "User-Agent": "relay-client/2.0",
  855. "X-Grok-Client-Version": "9.9.9",
  856. "X-Relay-Token": "relay-secret",
  857. },
  858. },
  859. }
  860. req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "access-token", "conv-1", nil)
  861. require.NoError(t, err)
  862. require.Equal(t, "https://relay.example.test/v1/responses", req.URL.String())
  863. // 覆写值优先于内置 CLI 身份头。名字不在 wire casing 映射中的覆写头
  864. // 以小写键直写(HTTP/2 线上语义),需按写入形态断言。
  865. require.Equal(t, "relay-client/2.0", req.Header.Get("User-Agent"))
  866. require.Equal(t, []string{"9.9.9"}, req.Header["x-grok-client-version"])
  867. require.Empty(t, req.Header.Get("X-Grok-Client-Version"))
  868. require.Equal(t, []string{"relay-secret"}, req.Header["x-relay-token"])
  869. // 会话路由头与认证头不受覆写影响。
  870. require.Equal(t, "conv-1", req.Header.Get(grokConversationIDHeader))
  871. require.Equal(t, "Bearer access-token", req.Header.Get("Authorization"))
  872. }
  873. func TestBuildGrokResponsesRequestIgnoresBlockedHeaderOverrides(t *testing.T) {
  874. t.Parallel()
  875. account := &Account{
  876. Platform: PlatformGrok,
  877. Type: AccountTypeAPIKey,
  878. Credentials: map[string]any{
  879. "header_override_enabled": true,
  880. "header_overrides": map[string]any{
  881. "Authorization": "Bearer stolen",
  882. "x-grok-conv-id": "pinned-conversation",
  883. },
  884. },
  885. }
  886. req, err := buildGrokResponsesRequest(context.Background(), nil, account, []byte(`{"model":"grok-4.3"}`), "api-key", "conv-2", nil)
  887. require.NoError(t, err)
  888. require.Equal(t, "Bearer api-key", req.Header.Get("Authorization"))
  889. require.Equal(t, "conv-2", req.Header.Get(grokConversationIDHeader))
  890. }
  891. func TestGrokMediaGenerationGateCoversImagesAndVideo(t *testing.T) {
  892. tests := []struct {
  893. name string
  894. endpoint GrokMediaEndpoint
  895. want bool
  896. }{
  897. {name: "image generation", endpoint: GrokMediaEndpointImagesGenerations, want: true},
  898. {name: "image edit", endpoint: GrokMediaEndpointImagesEdits, want: true},
  899. {name: "video generation", endpoint: GrokMediaEndpointVideosGenerations, want: true},
  900. {name: "video status", endpoint: GrokMediaEndpointVideoStatus, want: false},
  901. }
  902. for _, tt := range tests {
  903. t.Run(tt.name, func(t *testing.T) {
  904. require.Equal(t, tt.want, tt.endpoint.IsGenerationRequest())
  905. })
  906. }
  907. }
  908. func TestExtractGrokMediaModelSupportsJSONAndMultipart(t *testing.T) {
  909. require.Equal(t, "grok-imagine", ExtractGrokMediaModel("application/json", []byte(`{"model":"grok-imagine"}`)))
  910. var buf bytes.Buffer
  911. writer := multipart.NewWriter(&buf)
  912. require.NoError(t, writer.WriteField("prompt", "draw a cat"))
  913. require.NoError(t, writer.WriteField("model", "grok-imagine-edit"))
  914. require.NoError(t, writer.Close())
  915. require.Equal(t, "grok-imagine-edit", ExtractGrokMediaModel(writer.FormDataContentType(), buf.Bytes()))
  916. }
  917. func TestParseGrokMediaRequestBuildsMultipartModerationBody(t *testing.T) {
  918. var buf bytes.Buffer
  919. writer := multipart.NewWriter(&buf)
  920. require.NoError(t, writer.WriteField("prompt", "edit this private image"))
  921. require.NoError(t, writer.WriteField("model", "grok-imagine-edit"))
  922. partHeader := textproto.MIMEHeader{}
  923. partHeader.Set("Content-Disposition", `form-data; name="image"; filename="input.png"`)
  924. partHeader.Set("Content-Type", "image/png")
  925. part, err := writer.CreatePart(partHeader)
  926. require.NoError(t, err)
  927. _, err = part.Write([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a})
  928. require.NoError(t, err)
  929. require.NoError(t, writer.Close())
  930. info := ParseGrokMediaRequest(writer.FormDataContentType(), buf.Bytes())
  931. require.Equal(t, "grok-imagine-edit", info.Model)
  932. require.Equal(t, "edit this private image", info.Prompt)
  933. moderationBody := info.ModerationBody()
  934. require.NotEmpty(t, moderationBody)
  935. require.Equal(t, "edit this private image", gjson.GetBytes(moderationBody, "prompt").String())
  936. require.True(t, strings.HasPrefix(gjson.GetBytes(moderationBody, "images.0.image_url").String(), "data:image/"))
  937. }
  938. func TestParseGrokMediaVideoRequestResolution(t *testing.T) {
  939. info := ParseGrokMediaRequest("application/json", []byte(`{"model":"grok-imagine-video","prompt":"waves","resolution":"720p"}`))
  940. require.Equal(t, "grok-imagine-video", info.Model)
  941. require.Equal(t, "720p", info.Resolution)
  942. }
  943. func TestParseGrokMediaRequestAcceptsOfficialImageURLFields(t *testing.T) {
  944. body := []byte(`{
  945. "model":"grok-imagine-video-1.5",
  946. "image":{"url":"https://example.com/source.png"},
  947. "reference_images":[{"url":"https://example.com/reference.png"}]
  948. }`)
  949. info := ParseGrokMediaRequest("application/json", body)
  950. require.Equal(t, []string{
  951. "https://example.com/source.png",
  952. "https://example.com/reference.png",
  953. }, info.InputImageURLs)
  954. require.True(t, info.HasInputImage())
  955. }
  956. func TestNormalizeGrokMediaForwardBodyCanonicalizesImageURLAlias(t *testing.T) {
  957. body := []byte(`{
  958. "model":"grok-imagine-video-1.5",
  959. "prompt":"animate",
  960. "image":{"image_url":"https://example.com/source.png"},
  961. "duration":8
  962. }`)
  963. out, contentType, err := normalizeGrokMediaForwardBody(GrokMediaEndpointVideosGenerations, body, "application/json")
  964. require.NoError(t, err)
  965. require.Equal(t, "application/json", contentType)
  966. require.Equal(t, "grok-imagine-video-1.5", gjson.GetBytes(out, "model").String())
  967. require.Equal(t, "https://example.com/source.png", gjson.GetBytes(out, "image.url").String())
  968. require.False(t, gjson.GetBytes(out, "image.image_url").Exists())
  969. }
  970. func TestNormalizeGrokMediaForwardBodyPreservesImageToVideoModelForOfficialURL(t *testing.T) {
  971. body := []byte(`{
  972. "model":"grok-imagine-video-1.5",
  973. "prompt":"animate",
  974. "image":{"url":"https://example.com/source.png"}
  975. }`)
  976. out, _, err := normalizeGrokMediaForwardBody(GrokMediaEndpointVideosGenerations, body, "application/json")
  977. require.NoError(t, err)
  978. require.Equal(t, "grok-imagine-video-1.5", gjson.GetBytes(out, "model").String())
  979. require.Equal(t, "https://example.com/source.png", gjson.GetBytes(out, "image.url").String())
  980. }
  981. func TestCanonicalizeGrokMediaImageURLFieldsPreservesOfficialURL(t *testing.T) {
  982. body := []byte(`{
  983. "image":{"url":"https://example.com/official.png","image_url":"https://example.com/legacy.png"},
  984. "images":[
  985. {"image_url":"https://example.com/first.png"},
  986. {"url":"https://example.com/second.png"}
  987. ],
  988. "reference_images":[{"image_url":"https://example.com/reference.png"}],
  989. "mask":{"image_url":"https://example.com/mask.png"}
  990. }`)
  991. out, err := canonicalizeGrokMediaImageURLFields(body, "image", "images", "reference_images", "mask")
  992. require.NoError(t, err)
  993. require.Equal(t, "https://example.com/official.png", gjson.GetBytes(out, "image.url").String())
  994. require.False(t, gjson.GetBytes(out, "image.image_url").Exists())
  995. require.Equal(t, "https://example.com/first.png", gjson.GetBytes(out, "images.0.url").String())
  996. require.False(t, gjson.GetBytes(out, "images.0.image_url").Exists())
  997. require.Equal(t, "https://example.com/second.png", gjson.GetBytes(out, "images.1.url").String())
  998. require.Equal(t, "https://example.com/reference.png", gjson.GetBytes(out, "reference_images.0.url").String())
  999. require.False(t, gjson.GetBytes(out, "reference_images.0.image_url").Exists())
  1000. require.Equal(t, "https://example.com/mask.png", gjson.GetBytes(out, "mask.url").String())
  1001. require.False(t, gjson.GetBytes(out, "mask.image_url").Exists())
  1002. }
  1003. func TestCanonicalizeGrokMediaImageURLFieldsReplacesEmptyOfficialURL(t *testing.T) {
  1004. body := []byte(`{"image":{"url":" ","image_url":"https://example.com/legacy.png"}}`)
  1005. out, err := canonicalizeGrokMediaImageURLFields(body, "image")
  1006. require.NoError(t, err)
  1007. require.Equal(t, "https://example.com/legacy.png", gjson.GetBytes(out, "image.url").String())
  1008. require.False(t, gjson.GetBytes(out, "image.image_url").Exists())
  1009. }
  1010. func TestPrepareGrokImageEditNormalizesOfficialImageObjects(t *testing.T) {
  1011. body := []byte(`{
  1012. "model":"grok-imagine-image-quality",
  1013. "image":{"image_url":{"url":"https://example.com/first.png"}},
  1014. "images":["https://example.com/second.png"],
  1015. "mask":{"image_url":"https://example.com/mask.png"}
  1016. }`)
  1017. out, contentType, err := prepareGrokMediaForwardBody(GrokMediaEndpointImagesEdits, body, "application/json")
  1018. require.NoError(t, err)
  1019. require.Equal(t, "application/json", contentType)
  1020. for _, path := range []string{"image", "images.0", "mask"} {
  1021. require.Equal(t, "image_url", gjson.GetBytes(out, path+".type").String())
  1022. require.NotEmpty(t, gjson.GetBytes(out, path+".url").String())
  1023. require.False(t, gjson.GetBytes(out, path+".image_url").Exists())
  1024. }
  1025. }
  1026. func TestPrepareGrokImageEditRejectsMoreThanThreeSources(t *testing.T) {
  1027. body := []byte(`{"images":["https://example.com/1.png","https://example.com/2.png","https://example.com/3.png","https://example.com/4.png"]}`)
  1028. out, _, err := prepareGrokMediaForwardBody(GrokMediaEndpointImagesEdits, body, "application/json")
  1029. require.Error(t, err)
  1030. require.Nil(t, out)
  1031. require.Contains(t, err.Error(), "maximum of 3 source images")
  1032. }
  1033. func TestNormalizeGrokMediaModelForEndpoint(t *testing.T) {
  1034. tests := []struct {
  1035. name string
  1036. endpoint GrokMediaEndpoint
  1037. model string
  1038. hasInputImage bool
  1039. want string
  1040. }{
  1041. {name: "image generation alias", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine", want: "grok-imagine-image-quality"},
  1042. {name: "image edit alias", endpoint: GrokMediaEndpointImagesEdits, model: "grok-imagine", want: "grok-imagine-image-quality"},
  1043. {name: "image quality passthrough", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine-image-quality", want: "grok-imagine-image-quality"},
  1044. {name: "image fast passthrough", endpoint: GrokMediaEndpointImagesGenerations, model: "grok-imagine-image", want: "grok-imagine-image"},
  1045. {name: "video passthrough", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video", want: "grok-imagine-video"},
  1046. {name: "video 1.5 text-only remains explicit", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video-1.5", want: "grok-imagine-video-1.5"},
  1047. {name: "video 1.5 image-to-video passthrough", endpoint: GrokMediaEndpointVideosGenerations, model: "grok-imagine-video-1.5", hasInputImage: true, want: "grok-imagine-video-1.5"},
  1048. }
  1049. for _, tt := range tests {
  1050. t.Run(tt.name, func(t *testing.T) {
  1051. require.Equal(t, tt.want, NormalizeGrokMediaModelForEndpoint(tt.endpoint, tt.model, tt.hasInputImage))
  1052. })
  1053. }
  1054. }
  1055. func TestForwardGrokMediaImagesGenerationNormalizesImagineAlias(t *testing.T) {
  1056. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1057. gin.SetMode(gin.TestMode)
  1058. recorder := httptest.NewRecorder()
  1059. c, _ := gin.CreateTestContext(recorder)
  1060. body := []byte(`{"model":"grok-imagine","prompt":"draw a cat"}`)
  1061. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
  1062. c.Request.Header.Set("Content-Type", "application/json")
  1063. account := &Account{
  1064. ID: 61,
  1065. Name: "grok",
  1066. Platform: PlatformGrok,
  1067. Type: AccountTypeAPIKey,
  1068. Concurrency: 1,
  1069. Credentials: map[string]any{
  1070. "api_key": "api-key",
  1071. "base_url": "https://xai.test/v1",
  1072. },
  1073. }
  1074. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1075. StatusCode: http.StatusOK,
  1076. Header: http.Header{
  1077. "Content-Type": []string{"application/json"},
  1078. "Xai-Request-Id": []string{"xai-image-req"},
  1079. },
  1080. Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)),
  1081. }}
  1082. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1083. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
  1084. require.NoError(t, err)
  1085. require.Equal(t, "https://xai.test/v1/images/generations", upstream.lastReq.URL.String())
  1086. require.Equal(t, http.MethodPost, upstream.lastReq.Method)
  1087. require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization"))
  1088. require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
  1089. require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  1090. require.NotEqual(t, defaultGrokUpstreamUserAgent(), upstream.lastReq.Header.Get("User-Agent"))
  1091. require.JSONEq(t, `{"model":"grok-imagine-image-quality","prompt":"draw a cat"}`, string(upstream.lastBody))
  1092. require.Equal(t, http.StatusOK, recorder.Code)
  1093. require.JSONEq(t, `{"data":[{"url":"https://images.test/cat.png"}]}`, recorder.Body.String())
  1094. require.Equal(t, "xai-image-req", result.RequestID)
  1095. require.Equal(t, "grok-imagine-image-quality", result.Model)
  1096. require.Equal(t, "grok-imagine-image-quality", result.BillingModel)
  1097. require.Equal(t, 1, result.ImageCount)
  1098. require.Equal(t, ImageBillingSize2K, result.ImageSize)
  1099. }
  1100. func TestForwardGrokMediaAppliesAccountModelMappingAfterEndpointNormalization(t *testing.T) {
  1101. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1102. gin.SetMode(gin.TestMode)
  1103. tests := []struct {
  1104. name string
  1105. endpoint GrokMediaEndpoint
  1106. path string
  1107. body string
  1108. modelMapping map[string]any
  1109. wantRequestModel string
  1110. wantUpstream string
  1111. wantBody string
  1112. responseBody string
  1113. }{
  1114. {
  1115. name: "image generation maps normalized image alias",
  1116. endpoint: GrokMediaEndpointImagesGenerations,
  1117. path: "/v1/images/generations",
  1118. body: `{"model":"grok-imagine","prompt":"draw a cat"}`,
  1119. modelMapping: map[string]any{"grok-imagine-image-quality": "vendor-image-model"},
  1120. wantRequestModel: "grok-imagine-image-quality",
  1121. wantUpstream: "vendor-image-model",
  1122. wantBody: `{"model":"vendor-image-model","prompt":"draw a cat"}`,
  1123. responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`,
  1124. },
  1125. {
  1126. name: "video generation maps text-only fallback model",
  1127. endpoint: GrokMediaEndpointVideosGenerations,
  1128. path: "/v1/videos/generations",
  1129. body: `{"model":"grok-imagine-video-1.5","prompt":"waves"}`,
  1130. modelMapping: map[string]any{"grok-imagine-video": "grok-image-video"},
  1131. wantRequestModel: "grok-imagine-video-1.5",
  1132. wantUpstream: "grok-imagine-video-1.5",
  1133. wantBody: `{"model":"grok-imagine-video-1.5","prompt":"waves"}`,
  1134. responseBody: `{"request_id":"video-request-mapped"}`,
  1135. },
  1136. {
  1137. name: "image-to-video preserves then maps the requested model",
  1138. endpoint: GrokMediaEndpointVideosGenerations,
  1139. path: "/v1/videos/generations",
  1140. body: `{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"url":"https://example.com/input.png"}}`,
  1141. modelMapping: map[string]any{"grok-imagine-video-1.5": "vendor-image-video"},
  1142. wantRequestModel: "grok-imagine-video-1.5",
  1143. wantUpstream: "vendor-image-video",
  1144. wantBody: `{"model":"vendor-image-video","prompt":"animate","image":{"url":"https://example.com/input.png"}}`,
  1145. responseBody: `{"request_id":"image-video-request-mapped"}`,
  1146. },
  1147. {
  1148. name: "mapping and image sanitization compose",
  1149. endpoint: GrokMediaEndpointImagesGenerations,
  1150. path: "/v1/images/generations",
  1151. body: `{"model":"grok-imagine","prompt":"draw","size":"1024x1024"}`,
  1152. modelMapping: map[string]any{"grok-imagine-image-quality": "vendor-image-model"},
  1153. wantRequestModel: "grok-imagine-image-quality",
  1154. wantUpstream: "vendor-image-model",
  1155. wantBody: `{"model":"vendor-image-model","prompt":"draw","resolution":"1k","aspect_ratio":"1:1"}`,
  1156. responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`,
  1157. },
  1158. {
  1159. name: "whitespace mapping target safely preserves normalized model",
  1160. endpoint: GrokMediaEndpointImagesGenerations,
  1161. path: "/v1/images/generations",
  1162. body: `{"model":"grok-imagine","prompt":"draw"}`,
  1163. modelMapping: map[string]any{"grok-imagine-image-quality": " "},
  1164. wantRequestModel: "grok-imagine-image-quality",
  1165. wantUpstream: "grok-imagine-image-quality",
  1166. wantBody: `{"model":"grok-imagine-image-quality","prompt":"draw"}`,
  1167. responseBody: `{"data":[{"url":"https://images.test/mapped.png"}]}`,
  1168. },
  1169. }
  1170. for _, tt := range tests {
  1171. t.Run(tt.name, func(t *testing.T) {
  1172. recorder := httptest.NewRecorder()
  1173. c, _ := gin.CreateTestContext(recorder)
  1174. c.Request = httptest.NewRequest(http.MethodPost, tt.path, strings.NewReader(tt.body))
  1175. c.Request.Header.Set("Content-Type", "application/json")
  1176. account := &Account{
  1177. ID: 66,
  1178. Name: "grok-mapped",
  1179. Platform: PlatformGrok,
  1180. Type: AccountTypeAPIKey,
  1181. Concurrency: 1,
  1182. Credentials: map[string]any{
  1183. "api_key": "api-key",
  1184. "base_url": "https://xai.test/v1",
  1185. "model_mapping": tt.modelMapping,
  1186. },
  1187. }
  1188. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1189. StatusCode: http.StatusOK,
  1190. Header: http.Header{"Content-Type": []string{"application/json"}},
  1191. Body: io.NopCloser(strings.NewReader(tt.responseBody)),
  1192. }}
  1193. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1194. result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", []byte(tt.body), "application/json")
  1195. require.NoError(t, err)
  1196. require.JSONEq(t, tt.wantBody, string(upstream.lastBody))
  1197. require.Equal(t, tt.wantRequestModel, result.Model)
  1198. require.Equal(t, tt.wantRequestModel, result.BillingModel)
  1199. require.Equal(t, tt.wantUpstream, result.UpstreamModel)
  1200. })
  1201. }
  1202. }
  1203. func TestForwardGrokMediaImagesGenerationRejectsEmptySuccessfulResponse(t *testing.T) {
  1204. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1205. gin.SetMode(gin.TestMode)
  1206. recorder := httptest.NewRecorder()
  1207. c, _ := gin.CreateTestContext(recorder)
  1208. body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat"}`)
  1209. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
  1210. c.Request.Header.Set("Content-Type", "application/json")
  1211. account := &Account{
  1212. ID: 66,
  1213. Name: "grok",
  1214. Platform: PlatformGrok,
  1215. Type: AccountTypeAPIKey,
  1216. Concurrency: 1,
  1217. Credentials: map[string]any{
  1218. "api_key": "api-key",
  1219. "base_url": "https://xai.test/v1",
  1220. },
  1221. }
  1222. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1223. StatusCode: http.StatusOK,
  1224. Header: http.Header{"Content-Type": []string{"application/json"}},
  1225. Body: io.NopCloser(strings.NewReader(`{"data":[]}`)),
  1226. }}
  1227. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1228. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
  1229. require.Nil(t, result)
  1230. var failoverErr *UpstreamFailoverError
  1231. require.ErrorAs(t, err, &failoverErr)
  1232. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  1233. require.JSONEq(t, `{"data":[]}`, string(failoverErr.ResponseBody))
  1234. require.Empty(t, recorder.Body.String())
  1235. }
  1236. func TestForwardGrokMediaImagesGenerationStripsUnsupportedSize(t *testing.T) {
  1237. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1238. gin.SetMode(gin.TestMode)
  1239. recorder := httptest.NewRecorder()
  1240. c, _ := gin.CreateTestContext(recorder)
  1241. body := []byte(`{"model":"grok-imagine-image","prompt":"draw a cat","size":"1024x1024"}`)
  1242. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
  1243. c.Request.Header.Set("Content-Type", "application/json")
  1244. account := &Account{
  1245. ID: 65,
  1246. Name: "grok",
  1247. Platform: PlatformGrok,
  1248. Type: AccountTypeAPIKey,
  1249. Concurrency: 1,
  1250. Credentials: map[string]any{
  1251. "api_key": "api-key",
  1252. "base_url": "https://xai.test/v1",
  1253. "model_mapping": map[string]any{"grok-imagine-edit": "vendor-image-edit"},
  1254. },
  1255. }
  1256. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1257. StatusCode: http.StatusOK,
  1258. Header: http.Header{
  1259. "Content-Type": []string{"application/json"},
  1260. },
  1261. Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/cat.png"}]}`)),
  1262. }}
  1263. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1264. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
  1265. require.NoError(t, err)
  1266. require.JSONEq(t, `{"model":"grok-imagine-image","prompt":"draw a cat","resolution":"1k","aspect_ratio":"1:1"}`, string(upstream.lastBody))
  1267. require.False(t, gjson.GetBytes(upstream.lastBody, "size").Exists())
  1268. require.Equal(t, ImageBillingSize1K, result.ImageSize)
  1269. require.Equal(t, "1024x1024", result.ImageInputSize)
  1270. }
  1271. func TestForwardGrokMediaImagesEditMultipartConvertsToJSON(t *testing.T) {
  1272. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1273. gin.SetMode(gin.TestMode)
  1274. var buf bytes.Buffer
  1275. writer := multipart.NewWriter(&buf)
  1276. require.NoError(t, writer.WriteField("model", "grok-imagine-edit"))
  1277. require.NoError(t, writer.WriteField("prompt", "edit this private image"))
  1278. partHeader := textproto.MIMEHeader{}
  1279. partHeader.Set("Content-Disposition", `form-data; name="image"; filename="input.png"`)
  1280. partHeader.Set("Content-Type", "image/png")
  1281. part, err := writer.CreatePart(partHeader)
  1282. require.NoError(t, err)
  1283. _, err = part.Write([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a})
  1284. require.NoError(t, err)
  1285. require.NoError(t, writer.Close())
  1286. recorder := httptest.NewRecorder()
  1287. c, _ := gin.CreateTestContext(recorder)
  1288. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(buf.Bytes()))
  1289. c.Request.Header.Set("Content-Type", writer.FormDataContentType())
  1290. account := &Account{
  1291. ID: 62,
  1292. Name: "grok",
  1293. Platform: PlatformGrok,
  1294. Type: AccountTypeAPIKey,
  1295. Concurrency: 1,
  1296. Credentials: map[string]any{
  1297. "api_key": "api-key",
  1298. "base_url": "https://xai.test/v1",
  1299. "model_mapping": map[string]any{"grok-imagine-edit": "vendor-image-edit"},
  1300. },
  1301. }
  1302. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1303. StatusCode: http.StatusOK,
  1304. Header: http.Header{
  1305. "Content-Type": []string{"application/json"},
  1306. },
  1307. Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/edited.png"}]}`)),
  1308. }}
  1309. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1310. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesEdits, "", buf.Bytes(), writer.FormDataContentType())
  1311. require.NoError(t, err)
  1312. require.Equal(t, "https://xai.test/v1/images/edits", upstream.lastReq.URL.String())
  1313. require.Equal(t, "application/json", upstream.lastReq.Header.Get("Content-Type"))
  1314. require.True(t, json.Valid(upstream.lastBody))
  1315. require.Equal(t, "vendor-image-edit", gjson.GetBytes(upstream.lastBody, "model").String())
  1316. require.Equal(t, "edit this private image", gjson.GetBytes(upstream.lastBody, "prompt").String())
  1317. require.True(t, strings.HasPrefix(gjson.GetBytes(upstream.lastBody, "image.url").String(), "data:image/png;base64,"))
  1318. require.False(t, gjson.GetBytes(upstream.lastBody, "image.image_url").Exists())
  1319. require.Equal(t, "grok-imagine-edit", result.BillingModel)
  1320. require.Equal(t, "vendor-image-edit", result.UpstreamModel)
  1321. }
  1322. func TestForwardGrokMediaImagesEditMultipartPreservesExplicitGeometry(t *testing.T) {
  1323. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1324. gin.SetMode(gin.TestMode)
  1325. var buf bytes.Buffer
  1326. writer := multipart.NewWriter(&buf)
  1327. require.NoError(t, writer.WriteField("model", "grok-imagine-edit"))
  1328. require.NoError(t, writer.WriteField("prompt", "edit this private image"))
  1329. require.NoError(t, writer.WriteField("size", "1024x1024"))
  1330. require.NoError(t, writer.WriteField("resolution", "2k"))
  1331. require.NoError(t, writer.WriteField("aspect_ratio", "16:9"))
  1332. partHeader := textproto.MIMEHeader{}
  1333. partHeader.Set("Content-Disposition", `form-data; name="image"; filename="input.png"`)
  1334. partHeader.Set("Content-Type", "image/png")
  1335. part, err := writer.CreatePart(partHeader)
  1336. require.NoError(t, err)
  1337. _, err = part.Write([]byte{0x89, 0x50, 0x4e, 0x47, 0x0d, 0x0a, 0x1a, 0x0a})
  1338. require.NoError(t, err)
  1339. require.NoError(t, writer.Close())
  1340. recorder := httptest.NewRecorder()
  1341. c, _ := gin.CreateTestContext(recorder)
  1342. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/edits", bytes.NewReader(buf.Bytes()))
  1343. c.Request.Header.Set("Content-Type", writer.FormDataContentType())
  1344. account := &Account{
  1345. ID: 67,
  1346. Name: "grok",
  1347. Platform: PlatformGrok,
  1348. Type: AccountTypeAPIKey,
  1349. Concurrency: 1,
  1350. Credentials: map[string]any{
  1351. "api_key": "api-key",
  1352. "base_url": "https://xai.test/v1",
  1353. },
  1354. }
  1355. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1356. StatusCode: http.StatusOK,
  1357. Header: http.Header{"Content-Type": []string{"application/json"}},
  1358. Body: io.NopCloser(strings.NewReader(`{"data":[{"url":"https://images.test/edited.png"}]}`)),
  1359. }}
  1360. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1361. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesEdits, "", buf.Bytes(), writer.FormDataContentType())
  1362. require.NoError(t, err)
  1363. require.False(t, gjson.GetBytes(upstream.lastBody, "size").Exists())
  1364. require.Equal(t, "2k", gjson.GetBytes(upstream.lastBody, "resolution").String())
  1365. require.Equal(t, "16:9", gjson.GetBytes(upstream.lastBody, "aspect_ratio").String())
  1366. require.Equal(t, ImageBillingSize1K, result.ImageSize)
  1367. require.Equal(t, "1024x1024", result.ImageInputSize)
  1368. }
  1369. func TestForwardGrokMediaVideoGenerationReturnsUsageAndResponseID(t *testing.T) {
  1370. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1371. gin.SetMode(gin.TestMode)
  1372. recorder := httptest.NewRecorder()
  1373. c, _ := gin.CreateTestContext(recorder)
  1374. body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"waves","resolution":"720p","duration":10}`)
  1375. c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
  1376. c.Request.Header.Set("Content-Type", "application/json")
  1377. account := &Account{
  1378. ID: 63,
  1379. Name: "grok",
  1380. Platform: PlatformGrok,
  1381. Type: AccountTypeAPIKey,
  1382. Concurrency: 1,
  1383. Credentials: map[string]any{
  1384. "api_key": "api-key",
  1385. "base_url": "https://xai.test/v1",
  1386. },
  1387. }
  1388. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1389. StatusCode: http.StatusOK,
  1390. Header: http.Header{
  1391. "Content-Type": []string{"application/json"},
  1392. "Xai-Request-Id": []string{"xai-video-generate-req"},
  1393. },
  1394. Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-123","usage":{"prompt_tokens":3,"completion_tokens":4}}`)),
  1395. }}
  1396. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1397. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
  1398. require.NoError(t, err)
  1399. require.Equal(t, "https://xai.test/v1/videos/generations", upstream.lastReq.URL.String())
  1400. require.JSONEq(t, `{"model":"grok-imagine-video-1.5","prompt":"waves","resolution":"720p","duration":10}`, string(upstream.lastBody))
  1401. require.Equal(t, "video-request-123", result.ResponseID)
  1402. require.Equal(t, "grok-imagine-video-1.5", result.BillingModel)
  1403. require.Equal(t, 3, result.Usage.InputTokens)
  1404. require.Equal(t, 4, result.Usage.OutputTokens)
  1405. // Create accepts the job only — VideoCount stays 0 until status returns video.url.
  1406. require.Equal(t, 0, result.ImageCount)
  1407. require.Empty(t, result.ImageSize)
  1408. require.Equal(t, 0, result.VideoCount)
  1409. require.Equal(t, VideoBillingResolution720P, result.VideoResolution)
  1410. require.Equal(t, 10, result.VideoDurationSeconds)
  1411. }
  1412. func TestForwardGrokMediaVideoGenerationReturnsTaskIDAsResponseID(t *testing.T) {
  1413. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1414. gin.SetMode(gin.TestMode)
  1415. recorder := httptest.NewRecorder()
  1416. c, _ := gin.CreateTestContext(recorder)
  1417. body := []byte(`{"model":"grok-imagine-video","prompt":"waves"}`)
  1418. c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
  1419. c.Request.Header.Set("Content-Type", "application/json")
  1420. account := &Account{
  1421. ID: 63,
  1422. Name: "grok",
  1423. Platform: PlatformGrok,
  1424. Type: AccountTypeAPIKey,
  1425. Concurrency: 1,
  1426. Credentials: map[string]any{
  1427. "api_key": "api-key",
  1428. "base_url": "https://xai.test/v1",
  1429. },
  1430. }
  1431. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1432. StatusCode: http.StatusOK,
  1433. Header: http.Header{"Content-Type": []string{"application/json"}},
  1434. Body: io.NopCloser(strings.NewReader(`{"task_id":"video-task-123"}`)),
  1435. }}
  1436. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1437. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
  1438. require.NoError(t, err)
  1439. require.Equal(t, "video-task-123", result.ResponseID)
  1440. }
  1441. func TestExtractGrokMediaVideoRequestIDPreservesExistingPrecedence(t *testing.T) {
  1442. body := []byte(`{
  1443. "request_id":"request-id",
  1444. "id":"id",
  1445. "task_id":"task-id",
  1446. "data":{"request_id":"data-request-id","id":"data-id","task_id":"data-task-id"},
  1447. "video":{"request_id":"video-request-id","id":"video-id","task_id":"video-task-id"}
  1448. }`)
  1449. require.Equal(t, "request-id", extractGrokMediaVideoRequestID(body))
  1450. }
  1451. func TestForwardGrokMediaVideoGenerationPreservesImageToVideoModel(t *testing.T) {
  1452. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1453. gin.SetMode(gin.TestMode)
  1454. recorder := httptest.NewRecorder()
  1455. c, _ := gin.CreateTestContext(recorder)
  1456. body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"image_url":"data:image/png;base64,aW1n"}}`)
  1457. c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
  1458. c.Request.Header.Set("Content-Type", "application/json")
  1459. account := &Account{
  1460. ID: 63,
  1461. Name: "grok",
  1462. Platform: PlatformGrok,
  1463. Type: AccountTypeAPIKey,
  1464. Concurrency: 1,
  1465. Credentials: map[string]any{
  1466. "api_key": "api-key",
  1467. "base_url": "https://xai.test/v1",
  1468. },
  1469. }
  1470. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1471. StatusCode: http.StatusOK,
  1472. Header: http.Header{
  1473. "Content-Type": []string{"application/json"},
  1474. },
  1475. Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-456"}`)),
  1476. }}
  1477. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1478. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
  1479. require.NoError(t, err)
  1480. require.Equal(t, "https://xai.test/v1/videos/generations", upstream.lastReq.URL.String())
  1481. require.JSONEq(t, `{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"url":"data:image/png;base64,aW1n"}}`, string(upstream.lastBody))
  1482. require.Equal(t, "video-request-456", result.ResponseID)
  1483. require.Equal(t, "grok-imagine-video-1.5", result.BillingModel)
  1484. // 未指定 duration 时按上游默认 8 秒计费。
  1485. require.Equal(t, VideoBillingDefaultDurationSeconds, result.VideoDurationSeconds)
  1486. }
  1487. func TestForwardGrokMediaOAuthImageToVideoUsesOfficialAPIForLargeBody(t *testing.T) {
  1488. gin.SetMode(gin.TestMode)
  1489. recorder := httptest.NewRecorder()
  1490. c, _ := gin.CreateTestContext(recorder)
  1491. imageData := strings.Repeat("A", 2*1024*1024)
  1492. body := []byte(`{"model":"grok-imagine-video-1.5","prompt":"animate","image":{"image_url":"data:image/png;base64,` + imageData + `"}}`)
  1493. c.Request = httptest.NewRequest(http.MethodPost, "/v1/videos/generations", bytes.NewReader(body))
  1494. c.Request.Header.Set("Content-Type", "application/json")
  1495. account := &Account{
  1496. ID: 66,
  1497. Name: "grok-oauth",
  1498. Platform: PlatformGrok,
  1499. Type: AccountTypeOAuth,
  1500. Status: StatusActive,
  1501. Schedulable: true,
  1502. Concurrency: 1,
  1503. Credentials: map[string]any{
  1504. "access_token": "oauth-access-token",
  1505. "refresh_token": "oauth-refresh-token",
  1506. "expires_at": time.Now().Add(2 * time.Hour).UTC().Format(time.RFC3339),
  1507. "base_url": xai.DefaultCLIBaseURL,
  1508. },
  1509. }
  1510. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1511. StatusCode: http.StatusOK,
  1512. Header: http.Header{
  1513. "Content-Type": []string{"application/json"},
  1514. },
  1515. Body: io.NopCloser(strings.NewReader(`{"request_id":"video-request-oauth"}`)),
  1516. }}
  1517. svc := &OpenAIGatewayService{httpUpstream: upstream, grokTokenProvider: NewGrokTokenProvider(nil, nil)}
  1518. _, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideosGenerations, "", body, "application/json")
  1519. require.NoError(t, err)
  1520. require.Equal(t, xai.DefaultBaseURL+"/videos/generations", upstream.lastReq.URL.String())
  1521. require.Empty(t, upstream.lastReq.Header.Get("X-XAI-Token-Auth"))
  1522. require.Empty(t, upstream.lastReq.Header.Get("x-grok-client-version"))
  1523. require.Equal(t, "data:image/png;base64,"+imageData, gjson.GetBytes(upstream.lastBody, "image.url").String())
  1524. require.False(t, gjson.GetBytes(upstream.lastBody, "image.image_url").Exists())
  1525. }
  1526. func TestForwardGrokMediaVideoStatusUsesGETWithoutBody(t *testing.T) {
  1527. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1528. gin.SetMode(gin.TestMode)
  1529. recorder := httptest.NewRecorder()
  1530. c, _ := gin.CreateTestContext(recorder)
  1531. c.Request = httptest.NewRequest(http.MethodGet, "/v1/videos/request-123", nil)
  1532. account := &Account{
  1533. ID: 62,
  1534. Name: "grok",
  1535. Platform: PlatformGrok,
  1536. Type: AccountTypeAPIKey,
  1537. Concurrency: 1,
  1538. Credentials: map[string]any{
  1539. "api_key": "api-key",
  1540. "base_url": "https://xai.test/v1",
  1541. },
  1542. }
  1543. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1544. StatusCode: http.StatusOK,
  1545. Header: http.Header{
  1546. "Content-Type": []string{"application/json"},
  1547. "Xai-Request-Id": []string{"xai-video-req"},
  1548. },
  1549. Body: io.NopCloser(strings.NewReader(`{"id":"request-123","status":"completed"}`)),
  1550. }}
  1551. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1552. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointVideoStatus, "request-123", nil, "")
  1553. require.NoError(t, err)
  1554. require.Equal(t, "https://xai.test/v1/videos/request-123", upstream.lastReq.URL.String())
  1555. require.Equal(t, http.MethodGet, upstream.lastReq.Method)
  1556. require.Equal(t, "Bearer api-key", upstream.lastReq.Header.Get("Authorization"))
  1557. require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  1558. require.NotEqual(t, defaultGrokUpstreamUserAgent(), upstream.lastReq.Header.Get("User-Agent"))
  1559. require.Empty(t, upstream.lastReq.Header.Get("Content-Type"))
  1560. require.Empty(t, upstream.lastBody)
  1561. require.Equal(t, http.StatusOK, recorder.Code)
  1562. require.JSONEq(t, `{"id":"request-123","status":"completed"}`, recorder.Body.String())
  1563. require.Equal(t, "xai-video-req", result.RequestID)
  1564. }
  1565. func TestForwardGrokMediaVideoMutationEndpoints(t *testing.T) {
  1566. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1567. gin.SetMode(gin.TestMode)
  1568. tests := []struct {
  1569. name string
  1570. endpoint GrokMediaEndpoint
  1571. path string
  1572. }{
  1573. {name: "edit", endpoint: GrokMediaEndpointVideosEdits, path: "/videos/edits"},
  1574. {name: "extension", endpoint: GrokMediaEndpointVideosExtensions, path: "/videos/extensions"},
  1575. }
  1576. for _, tt := range tests {
  1577. t.Run(tt.name, func(t *testing.T) {
  1578. recorder := httptest.NewRecorder()
  1579. c, _ := gin.CreateTestContext(recorder)
  1580. body := []byte(`{"model":"grok-imagine-video","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`)
  1581. c.Request = httptest.NewRequest(http.MethodPost, "/v1"+tt.path, bytes.NewReader(body))
  1582. c.Request.Header.Set("Content-Type", "application/json")
  1583. account := &Account{
  1584. ID: 71, Name: "grok", Platform: PlatformGrok, Type: AccountTypeAPIKey, Concurrency: 1,
  1585. Credentials: map[string]any{
  1586. "api_key": "api-key",
  1587. "base_url": "https://xai.test/v1",
  1588. "model_mapping": map[string]any{"grok-imagine-video": "vendor-video-mutation"},
  1589. },
  1590. }
  1591. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1592. StatusCode: http.StatusOK,
  1593. Header: http.Header{"Content-Type": []string{"application/json"}},
  1594. Body: io.NopCloser(strings.NewReader(`{"request_id":"video-mutation-123"}`)),
  1595. }}
  1596. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1597. result, err := svc.ForwardGrokMedia(context.Background(), c, account, tt.endpoint, "", body, "application/json")
  1598. require.NoError(t, err)
  1599. require.Equal(t, "https://xai.test/v1"+tt.path, upstream.lastReq.URL.String())
  1600. require.Equal(t, http.MethodPost, upstream.lastReq.Method)
  1601. require.JSONEq(t, `{"model":"vendor-video-mutation","prompt":"continue","video":{"url":"https://example.com/in.mp4"},"duration":6}`, string(upstream.lastBody))
  1602. require.Equal(t, "video-mutation-123", result.ResponseID)
  1603. require.Equal(t, 0, result.VideoCount)
  1604. require.Equal(t, 6, result.VideoDurationSeconds)
  1605. require.Equal(t, "grok-imagine-video", result.BillingModel)
  1606. require.Equal(t, "vendor-video-mutation", result.UpstreamModel)
  1607. })
  1608. }
  1609. }
  1610. func TestGrokMediaVideoRequestBindingIsScopedToUserAndAPIKey(t *testing.T) {
  1611. gin.SetMode(gin.TestMode)
  1612. c, _ := gin.CreateTestContext(httptest.NewRecorder())
  1613. c.Request = httptest.NewRequest(http.MethodGet, "/v1/videos/video-request-123", nil)
  1614. c.Request.Header.Set("session_id", "shared-client-session")
  1615. groupID := int64(7)
  1616. cache := &stubGatewayCache{}
  1617. svc := &OpenAIGatewayService{cache: cache}
  1618. const userID int64 = 41
  1619. const apiKeyID int64 = 51
  1620. require.NotEmpty(t, svc.GenerateExplicitSessionHash(c, nil))
  1621. ctx := c.Request.Context()
  1622. hash := GrokMediaVideoRequestSessionHash("video-request-123", userID, apiKeyID)
  1623. require.NotEmpty(t, hash)
  1624. require.NoError(t, svc.BindGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID, 63))
  1625. accountID, err := svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID)
  1626. require.NoError(t, err)
  1627. require.Equal(t, int64(63), accountID)
  1628. accountID, err = svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID+1, apiKeyID)
  1629. require.Error(t, err)
  1630. require.Zero(t, accountID)
  1631. accountID, err = svc.ResolveGrokMediaVideoRequestAccount(ctx, &groupID, "video-request-123", userID, apiKeyID+1)
  1632. require.Error(t, err)
  1633. require.Zero(t, accountID)
  1634. }
  1635. func TestForwardGrokMedia429ReconcilesRateLimitBeforeCustomErrorBypass(t *testing.T) {
  1636. t.Setenv(xai.EnvAllowUnsafeURLOverrides, "true")
  1637. gin.SetMode(gin.TestMode)
  1638. recorder := httptest.NewRecorder()
  1639. c, _ := gin.CreateTestContext(recorder)
  1640. body := []byte(`{"model":"grok-imagine","prompt":"draw a cat"}`)
  1641. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", bytes.NewReader(body))
  1642. c.Request.Header.Set("Content-Type", "application/json")
  1643. account := &Account{
  1644. ID: 64,
  1645. Name: "grok",
  1646. Platform: PlatformGrok,
  1647. Type: AccountTypeAPIKey,
  1648. Concurrency: 1,
  1649. Credentials: map[string]any{
  1650. "api_key": "api-key",
  1651. "base_url": "https://xai.test/v1",
  1652. "custom_error_codes_enabled": true,
  1653. "custom_error_codes": []any{float64(http.StatusBadRequest)},
  1654. },
  1655. }
  1656. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1657. StatusCode: http.StatusTooManyRequests,
  1658. Header: http.Header{
  1659. "Content-Type": []string{"application/json"},
  1660. "Xai-Request-Id": []string{"xai-error-req"},
  1661. "Retry-After": []string{"45"},
  1662. },
  1663. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"do not expose this upstream detail"}}`)),
  1664. }}
  1665. repo := &grokQuotaAccountRepo{}
  1666. svc := &OpenAIGatewayService{httpUpstream: upstream, accountRepo: repo}
  1667. result, err := svc.ForwardGrokMedia(context.Background(), c, account, GrokMediaEndpointImagesGenerations, "", body, "application/json")
  1668. require.Error(t, err)
  1669. require.Nil(t, result)
  1670. require.Equal(t, http.StatusInternalServerError, recorder.Code)
  1671. require.Contains(t, recorder.Body.String(), "Upstream gateway error")
  1672. require.NotContains(t, recorder.Body.String(), "do not expose")
  1673. require.Equal(t, 1, repo.rateLimitedCalls)
  1674. require.Zero(t, repo.tempUnschedCalls)
  1675. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  1676. }
  1677. func TestGrokMedia429FailoverPreservesRetryAfter(t *testing.T) {
  1678. gin.SetMode(gin.TestMode)
  1679. recorder := httptest.NewRecorder()
  1680. c, _ := gin.CreateTestContext(recorder)
  1681. c.Request = httptest.NewRequest(http.MethodPost, "/v1/images/generations", nil)
  1682. account := &Account{
  1683. ID: 641, Name: "grok-oauth", Platform: PlatformGrok, Type: AccountTypeOAuth,
  1684. Status: StatusActive, Schedulable: true,
  1685. Credentials: map[string]any{
  1686. "custom_error_codes_enabled": true,
  1687. "custom_error_codes": []any{float64(http.StatusTooManyRequests)},
  1688. },
  1689. }
  1690. repo := &grokQuotaAccountRepo{}
  1691. svc := &OpenAIGatewayService{accountRepo: repo}
  1692. resp := &http.Response{
  1693. StatusCode: http.StatusTooManyRequests,
  1694. Header: http.Header{"Retry-After": []string{"45"}},
  1695. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
  1696. }
  1697. result, err := svc.handleGrokMediaErrorResponse(context.Background(), resp, c, account, "request-id", "grok-imagine")
  1698. require.Nil(t, result)
  1699. var failoverErr *UpstreamFailoverError
  1700. require.ErrorAs(t, err, &failoverErr)
  1701. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  1702. require.Equal(t, "45", failoverErr.ResponseHeaders.Get("Retry-After"))
  1703. }
  1704. func healthyGrokOAuthGatewayTestAccount(id int64, token string) *Account {
  1705. return &Account{
  1706. ID: id,
  1707. Name: "grok",
  1708. Platform: PlatformGrok,
  1709. Type: AccountTypeOAuth,
  1710. Status: StatusActive,
  1711. Schedulable: true,
  1712. Concurrency: 1,
  1713. Credentials: map[string]any{
  1714. "access_token": token,
  1715. "refresh_token": "refresh-token",
  1716. "expires_at": time.Now().Add(2 * grokTokenRefreshSkew).UTC().Format(time.RFC3339),
  1717. "base_url": xai.DefaultCLIBaseURL,
  1718. },
  1719. }
  1720. }
  1721. func TestForwardAsChatCompletionsForGrokStopFallsBackToXAIChatCompletions(t *testing.T) {
  1722. gin.SetMode(gin.TestMode)
  1723. recorder := httptest.NewRecorder()
  1724. c, _ := gin.CreateTestContext(recorder)
  1725. body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false,"stop":"done","prompt_cache_key":"raw-client-cache-key"}`)
  1726. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  1727. c.Set("api_key", &APIKey{ID: 5101})
  1728. account := healthyGrokOAuthGatewayTestAccount(51, "access-token")
  1729. repo := &grokQuotaAccountRepo{
  1730. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  1731. accountsByID: map[int64]*Account{51: account},
  1732. },
  1733. }
  1734. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1735. StatusCode: http.StatusOK,
  1736. Header: http.Header{
  1737. "Content-Type": []string{"application/json"},
  1738. "Xai-Request-Id": []string{"xai-req"},
  1739. "X-Ratelimit-Limit-Requests": []string{"10"},
  1740. "X-Ratelimit-Remaining-Requests": []string{"9"},
  1741. "X-Ratelimit-Limit-Tokens": []string{"1000"},
  1742. "X-Ratelimit-Remaining-Tokens": []string{"990"},
  1743. },
  1744. Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":2,"prompt_tokens_details":{"cached_tokens":1}}}`)),
  1745. }}
  1746. svc := &OpenAIGatewayService{
  1747. httpUpstream: upstream,
  1748. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  1749. accountRepo: repo,
  1750. }
  1751. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  1752. require.NoError(t, err)
  1753. require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
  1754. require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
  1755. require.NotEmpty(t, upstream.lastReq.Header.Get(grokConversationIDHeader))
  1756. require.NotEqual(t, "raw-client-cache-key", upstream.lastReq.Header.Get(grokConversationIDHeader))
  1757. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.lastBody, "model").String())
  1758. require.False(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").Exists())
  1759. require.Equal(t, "grok", result.Model)
  1760. require.Equal(t, "grok-4.6", result.UpstreamModel)
  1761. require.Equal(t, 1, result.Usage.InputTokens)
  1762. require.Equal(t, 2, result.Usage.OutputTokens)
  1763. require.Equal(t, 1, result.Usage.CacheReadInputTokens)
  1764. require.NotNil(t, repo.updates[51][grokQuotaSnapshotExtraKey])
  1765. require.Equal(t, http.StatusOK, recorder.Code)
  1766. }
  1767. func TestForwardGrokResponsesStreamingDefaultsEmptyModelTo45AndSnapshots(t *testing.T) {
  1768. gin.SetMode(gin.TestMode)
  1769. recorder := httptest.NewRecorder()
  1770. c, _ := gin.CreateTestContext(recorder)
  1771. body := []byte(`{"input":"hi","stream":true,"reasoning_effort":"high"}`)
  1772. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  1773. c.Request.Header.Set("Content-Type", "application/json")
  1774. c.Request.Header.Set("OpenAI-Beta", "responses=experimental")
  1775. c.Set("api_key", &APIKey{ID: 5201})
  1776. account := healthyGrokOAuthGatewayTestAccount(52, "access-token")
  1777. repo := &grokQuotaAccountRepo{
  1778. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  1779. accountsByID: map[int64]*Account{52: account},
  1780. },
  1781. }
  1782. upstreamBody := strings.Join([]string{
  1783. `data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`,
  1784. "",
  1785. `data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_grok","model":"grok-4.3","usage":{"input_tokens":5,"output_tokens":3,"input_tokens_details":{"cached_tokens":2}}}}`,
  1786. "",
  1787. }, "\n")
  1788. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1789. StatusCode: http.StatusOK,
  1790. Header: http.Header{
  1791. "Content-Type": []string{"text/event-stream"},
  1792. "Xai-Request-Id": []string{"xai-stream-req"},
  1793. "X-Ratelimit-Limit-Requests": []string{"10"},
  1794. "X-Ratelimit-Remaining-Requests": []string{"8"},
  1795. "X-Ratelimit-Limit-Tokens": []string{"1000"},
  1796. "X-Ratelimit-Remaining-Tokens": []string{"990"},
  1797. },
  1798. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  1799. }}
  1800. svc := &OpenAIGatewayService{
  1801. httpUpstream: upstream,
  1802. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  1803. accountRepo: repo,
  1804. }
  1805. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "", true, time.Now())
  1806. require.NoError(t, err)
  1807. require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
  1808. require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
  1809. require.Equal(t, "responses=experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
  1810. require.Equal(t, "grok-4.5", gjson.GetBytes(upstream.lastBody, "model").String())
  1811. require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
  1812. require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader))
  1813. require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
  1814. require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String())
  1815. require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
  1816. require.Equal(t, "high", gjson.GetBytes(upstream.lastBody, "reasoning_effort").String())
  1817. require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
  1818. require.True(t, result.Stream)
  1819. require.Equal(t, "resp_grok", result.ResponseID)
  1820. require.Equal(t, "xai-stream-req", result.RequestID)
  1821. require.Equal(t, 5, result.Usage.InputTokens)
  1822. require.Equal(t, 3, result.Usage.OutputTokens)
  1823. require.Equal(t, 2, result.Usage.CacheReadInputTokens)
  1824. require.NotNil(t, result.ReasoningEffort)
  1825. require.Equal(t, "high", *result.ReasoningEffort)
  1826. require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
  1827. require.Contains(t, recorder.Body.String(), "response.output_text.delta")
  1828. require.NotNil(t, repo.updates[52][grokQuotaSnapshotExtraKey])
  1829. }
  1830. func TestForwardGrokResponsesAPIKeyUsesXAIResponses(t *testing.T) {
  1831. gin.SetMode(gin.TestMode)
  1832. recorder := httptest.NewRecorder()
  1833. c, _ := gin.CreateTestContext(recorder)
  1834. body := []byte(`{"model":"grok","input":"hi","metadata":{"session_id":"abc"},"stream":true}`)
  1835. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  1836. c.Request.Header.Set("Content-Type", "application/json")
  1837. account := &Account{
  1838. ID: 53,
  1839. Name: "grok-api-key",
  1840. Platform: PlatformGrok,
  1841. Type: AccountTypeAPIKey,
  1842. Concurrency: 2,
  1843. Credentials: map[string]any{
  1844. "api_key": "xai-test-key",
  1845. "base_url": "https://api.x.ai/v1",
  1846. },
  1847. }
  1848. upstreamBody := strings.Join([]string{
  1849. `data: {"type":"response.output_text.delta","sequence_number":0,"delta":"ok"}`,
  1850. "",
  1851. `data: {"type":"response.completed","sequence_number":1,"response":{"id":"resp_grok_api_key","model":"grok-4.5","usage":{"input_tokens":2,"output_tokens":1}}}`,
  1852. "",
  1853. }, "\n")
  1854. upstream := &httpUpstreamRecorder{resp: &http.Response{
  1855. StatusCode: http.StatusOK,
  1856. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  1857. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  1858. }}
  1859. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1860. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", true, time.Now())
  1861. require.NoError(t, err)
  1862. require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
  1863. require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization"))
  1864. require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  1865. require.NotEqual(t, defaultGrokUpstreamUserAgent(), upstream.lastReq.Header.Get("User-Agent"))
  1866. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.lastBody, "model").String())
  1867. require.False(t, gjson.GetBytes(upstream.lastBody, "metadata").Exists())
  1868. require.Equal(t, "resp_grok_api_key", result.ResponseID)
  1869. require.Equal(t, 2, result.Usage.InputTokens)
  1870. require.Equal(t, 1, result.Usage.OutputTokens)
  1871. }
  1872. func TestForwardGrokResponsesUsesMetadataSessionForCacheIdentityWithoutForwardingMetadata(t *testing.T) {
  1873. gin.SetMode(gin.TestMode)
  1874. firstBody := []byte(`{"model":"grok","input":"first turn","metadata":{"user_id":"{\"session_id\":\"metadata-session\"}"},"stream":false}`)
  1875. secondBody := []byte(`{"model":"grok","input":"different second turn","metadata":{"user_id":"{\"session_id\":\"metadata-session\"}"},"stream":false}`)
  1876. account := &Account{
  1877. ID: 5401,
  1878. Name: "grok-api-key",
  1879. Platform: PlatformGrok,
  1880. Type: AccountTypeAPIKey,
  1881. Concurrency: 2,
  1882. Credentials: map[string]any{
  1883. "api_key": "xai-test-key",
  1884. "base_url": "https://api.x.ai/v1",
  1885. },
  1886. }
  1887. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  1888. {
  1889. StatusCode: http.StatusOK,
  1890. Header: http.Header{"Content-Type": []string{"application/json"}},
  1891. Body: io.NopCloser(strings.NewReader(`{"id":"resp_first","object":"response","model":"grok-4.6","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1}}`)),
  1892. },
  1893. {
  1894. StatusCode: http.StatusOK,
  1895. Header: http.Header{"Content-Type": []string{"application/json"}},
  1896. Body: io.NopCloser(strings.NewReader(`{"id":"resp_second","object":"response","model":"grok-4.6","status":"completed","output":[],"usage":{"input_tokens":3,"output_tokens":1}}`)),
  1897. },
  1898. }}
  1899. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1900. newContext := func(body []byte) *gin.Context {
  1901. recorder := httptest.NewRecorder()
  1902. c, _ := gin.CreateTestContext(recorder)
  1903. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  1904. c.Set("api_key", &APIKey{ID: 5401})
  1905. return c
  1906. }
  1907. _, err := svc.forwardGrokResponses(context.Background(), newContext(firstBody), account, firstBody, "grok", false, time.Now())
  1908. require.NoError(t, err)
  1909. _, err = svc.forwardGrokResponses(context.Background(), newContext(secondBody), account, secondBody, "grok", false, time.Now())
  1910. require.NoError(t, err)
  1911. require.Len(t, upstream.bodies, 2)
  1912. firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
  1913. secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
  1914. require.NotEmpty(t, firstIdentity)
  1915. require.Equal(t, firstIdentity, secondIdentity)
  1916. require.False(t, gjson.GetBytes(upstream.bodies[0], "metadata").Exists())
  1917. require.False(t, gjson.GetBytes(upstream.bodies[1], "metadata").Exists())
  1918. }
  1919. func TestForwardGrokResponsesRetriesInvalidEncryptedContentOnce(t *testing.T) {
  1920. gin.SetMode(gin.TestMode)
  1921. recorder := httptest.NewRecorder()
  1922. c, _ := gin.CreateTestContext(recorder)
  1923. body := []byte(`{
  1924. "model":"grok",
  1925. "previous_response_id":"resp_valid_history",
  1926. "input":[
  1927. {"type":"reasoning","summary":[{"type":"summary_text","text":"keep this summary"}],"encrypted_content":"encrypted-reasoning"},
  1928. {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]}
  1929. ],
  1930. "metadata":{"large_id":9007199254740993},
  1931. "stream":false
  1932. }`)
  1933. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  1934. c.Set("api_key", &APIKey{ID: 4535})
  1935. account := &Account{
  1936. ID: 4535,
  1937. Name: "grok-api-key",
  1938. Platform: PlatformGrok,
  1939. Type: AccountTypeAPIKey,
  1940. Status: StatusActive,
  1941. Schedulable: true,
  1942. Concurrency: 2,
  1943. Credentials: map[string]any{
  1944. "api_key": "same-token",
  1945. "base_url": "https://api.x.ai/v1",
  1946. },
  1947. }
  1948. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  1949. {
  1950. StatusCode: http.StatusBadRequest,
  1951. Header: http.Header{
  1952. "Content-Type": []string{"application/json"},
  1953. "Xai-Request-Id": []string{"recoverable-first"},
  1954. },
  1955. Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content. Ensure the value is unmodified."}`)),
  1956. },
  1957. {
  1958. StatusCode: http.StatusOK,
  1959. Header: http.Header{
  1960. "Content-Type": []string{"application/json"},
  1961. "Xai-Request-Id": []string{"recovered-second"},
  1962. },
  1963. Body: io.NopCloser(strings.NewReader(`{"id":"resp_recovered","object":"response","model":"grok-4.5","status":"completed","output":[],"usage":{"input_tokens":2,"output_tokens":1}}`)),
  1964. },
  1965. }}
  1966. svc := &OpenAIGatewayService{httpUpstream: upstream}
  1967. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
  1968. require.NoError(t, err)
  1969. require.NotNil(t, result)
  1970. require.Equal(t, "resp_recovered", result.ResponseID)
  1971. require.Equal(t, "recovered-second", result.RequestID)
  1972. require.Len(t, upstream.requests, 2)
  1973. require.Len(t, upstream.bodies, 2)
  1974. require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.type").String())
  1975. require.Equal(t, "resp_valid_history", gjson.GetBytes(upstream.bodies[0], "previous_response_id").String())
  1976. require.Equal(t, "encrypted-reasoning", gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").String())
  1977. require.Equal(t, "reasoning", gjson.GetBytes(upstream.bodies[1], "input.0.type").String())
  1978. require.Equal(t, "resp_valid_history", gjson.GetBytes(upstream.bodies[1], "previous_response_id").String())
  1979. require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists())
  1980. require.Equal(t, "keep this summary", gjson.GetBytes(upstream.bodies[1], "input.0.summary.0.text").String())
  1981. require.Equal(t, "message", gjson.GetBytes(upstream.bodies[1], "input.1.type").String())
  1982. require.False(t, gjson.GetBytes(upstream.bodies[0], "metadata").Exists())
  1983. require.False(t, gjson.GetBytes(upstream.bodies[1], "metadata").Exists())
  1984. firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
  1985. secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
  1986. require.NotEmpty(t, firstIdentity)
  1987. require.Equal(t, firstIdentity, secondIdentity)
  1988. for _, req := range upstream.requests {
  1989. require.Equal(t, "Bearer same-token", req.Header.Get("Authorization"))
  1990. require.Equal(t, firstIdentity, req.Header.Get(grokConversationIDHeader))
  1991. }
  1992. require.Equal(t, StatusActive, account.Status)
  1993. _, hasUpstreamErrors := c.Get(OpsUpstreamErrorsKey)
  1994. require.False(t, hasUpstreamErrors)
  1995. _, hasTerminalStatus := c.Get(OpsUpstreamStatusCodeKey)
  1996. require.False(t, hasTerminalStatus)
  1997. }
  1998. func TestForwardGrokResponsesInvalidEncryptedContentRecoveryDoesNotOvermatch(t *testing.T) {
  1999. gin.SetMode(gin.TestMode)
  2000. matchingError := `{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}`
  2001. tests := []struct {
  2002. name string
  2003. requestBody string
  2004. responseBody string
  2005. }{
  2006. {
  2007. name: "different top-level code",
  2008. requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`,
  2009. responseBody: `{"code":"bad-request","error":"Could not decrypt the provided encrypted_content."}`,
  2010. },
  2011. {
  2012. name: "message does not mention decryption",
  2013. requestBody: `{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"}],"stream":false}`,
  2014. responseBody: `{"code":"invalid-argument","error":"The provided encrypted_content is invalid."}`,
  2015. },
  2016. {
  2017. name: "request has no encrypted reasoning",
  2018. requestBody: `{"model":"grok","input":[{"type":"message","role":"user","content":"hi"}],"stream":false}`,
  2019. responseBody: matchingError,
  2020. },
  2021. }
  2022. for _, tt := range tests {
  2023. t.Run(tt.name, func(t *testing.T) {
  2024. recorder := httptest.NewRecorder()
  2025. c, _ := gin.CreateTestContext(recorder)
  2026. body := []byte(tt.requestBody)
  2027. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  2028. account := &Account{
  2029. ID: 4536,
  2030. Name: "grok-api-key",
  2031. Platform: PlatformGrok,
  2032. Type: AccountTypeAPIKey,
  2033. Concurrency: 1,
  2034. Credentials: map[string]any{"api_key": "token", "base_url": "https://api.x.ai/v1"},
  2035. }
  2036. upstream := &httpUpstreamRecorder{responses: []*http.Response{{
  2037. StatusCode: http.StatusBadRequest,
  2038. Header: http.Header{"Content-Type": []string{"application/json"}},
  2039. Body: io.NopCloser(strings.NewReader(tt.responseBody)),
  2040. }}}
  2041. svc := &OpenAIGatewayService{httpUpstream: upstream}
  2042. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
  2043. require.Nil(t, result)
  2044. require.Error(t, err)
  2045. require.Len(t, upstream.requests, 1)
  2046. require.Len(t, upstream.bodies, 1)
  2047. })
  2048. }
  2049. }
  2050. func TestGrokCompactionBlobRecoveryStripsCompactionItem(t *testing.T) {
  2051. body := []byte(`{"model":"grok","input":[{"type":"compaction","id":"cmp_1","encrypted_content":"blob"},{"type":"message","role":"user","content":"hi"}]}`)
  2052. require.True(t, isGrokInvalidEncryptedContentResponse(http.StatusUnprocessableEntity, []byte(`{"code":"invalid_compaction","error":"could not decode the compaction blob"}`)))
  2053. retry, changed, err := trimGrokInvalidEncryptedContentRetryBody(body)
  2054. require.NoError(t, err)
  2055. require.True(t, changed)
  2056. require.Equal(t, "message", gjson.GetBytes(retry, "input.0.type").String())
  2057. require.False(t, gjson.GetBytes(retry, "input.#(type==\"compaction\")").Exists())
  2058. }
  2059. func TestForwardGrokResponsesInvalidEncryptedContentRecoveryNestedErrorShape(t *testing.T) {
  2060. gin.SetMode(gin.TestMode)
  2061. recorder := httptest.NewRecorder()
  2062. c, _ := gin.CreateTestContext(recorder)
  2063. body := []byte(`{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"},{"type":"message","role":"user","content":"hi"}],"stream":false}`)
  2064. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  2065. account := &Account{
  2066. ID: 4538,
  2067. Name: "grok-api-key",
  2068. Platform: PlatformGrok,
  2069. Type: AccountTypeAPIKey,
  2070. Concurrency: 1,
  2071. Credentials: map[string]any{"api_key": "token", "base_url": "https://api.x.ai/v1"},
  2072. }
  2073. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  2074. {
  2075. StatusCode: http.StatusBadRequest,
  2076. Header: http.Header{"Content-Type": []string{"application/json"}},
  2077. Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":{"message":"Could not decrypt the provided encrypted_content."}}`)),
  2078. },
  2079. {
  2080. StatusCode: http.StatusOK,
  2081. Header: http.Header{"Content-Type": []string{"application/json"}},
  2082. Body: io.NopCloser(strings.NewReader(`{"id":"resp_ok","output":[],"usage":{"input_tokens":1,"output_tokens":1}}`)),
  2083. },
  2084. }}
  2085. svc := &OpenAIGatewayService{httpUpstream: upstream}
  2086. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
  2087. require.NoError(t, err)
  2088. require.NotNil(t, result)
  2089. require.Len(t, upstream.requests, 2)
  2090. require.True(t, gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").Exists())
  2091. require.False(t, gjson.GetBytes(upstream.bodies[1], "input.0.encrypted_content").Exists())
  2092. }
  2093. func TestForwardGrokResponsesInvalidEncryptedContentRetryFailureIsTerminal(t *testing.T) {
  2094. gin.SetMode(gin.TestMode)
  2095. recorder := httptest.NewRecorder()
  2096. c, _ := gin.CreateTestContext(recorder)
  2097. body := []byte(`{"model":"grok","input":[{"type":"reasoning","encrypted_content":"cipher"},{"type":"message","role":"user","content":"hi"}],"stream":false}`)
  2098. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  2099. account := &Account{
  2100. ID: 4537,
  2101. Name: "grok-api-key",
  2102. Platform: PlatformGrok,
  2103. Type: AccountTypeAPIKey,
  2104. Concurrency: 1,
  2105. Credentials: map[string]any{"api_key": "same-token", "base_url": "https://api.x.ai/v1"},
  2106. }
  2107. newInvalidEncryptedResponse := func(requestID string) *http.Response {
  2108. return &http.Response{
  2109. StatusCode: http.StatusBadRequest,
  2110. Header: http.Header{
  2111. "Content-Type": []string{"application/json"},
  2112. "Xai-Request-Id": []string{requestID},
  2113. },
  2114. Body: io.NopCloser(strings.NewReader(`{"code":"invalid-argument","error":"Could not decrypt the provided encrypted_content."}`)),
  2115. }
  2116. }
  2117. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  2118. newInvalidEncryptedResponse("recoverable-first"),
  2119. newInvalidEncryptedResponse("terminal-second"),
  2120. }}
  2121. svc := &OpenAIGatewayService{httpUpstream: upstream}
  2122. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
  2123. require.Nil(t, result)
  2124. require.Error(t, err)
  2125. require.Len(t, upstream.requests, 2)
  2126. require.Len(t, upstream.bodies, 2)
  2127. require.True(t, gjson.GetBytes(upstream.bodies[0], "input.0.encrypted_content").Exists())
  2128. require.False(t, gjson.GetBytes(upstream.bodies[1], `input.#(type=="reasoning")`).Exists())
  2129. rawEvents, ok := c.Get(OpsUpstreamErrorsKey)
  2130. require.True(t, ok)
  2131. events, ok := rawEvents.([]*OpsUpstreamErrorEvent)
  2132. require.True(t, ok)
  2133. require.NotEmpty(t, events)
  2134. for _, event := range events {
  2135. require.NotEqual(t, "recoverable-first", event.UpstreamRequestID)
  2136. }
  2137. require.Equal(t, http.StatusBadRequest, c.GetInt(OpsUpstreamStatusCodeKey))
  2138. }
  2139. func TestForwardAsChatCompletionsForGrokAPIKeyUsesConfiguredRawEndpointWithoutOAuthIdentity(t *testing.T) {
  2140. gin.SetMode(gin.TestMode)
  2141. recorder := httptest.NewRecorder()
  2142. c, _ := gin.CreateTestContext(recorder)
  2143. body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
  2144. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  2145. c.Request.Header.Set("Content-Type", "application/json")
  2146. account := &Account{
  2147. ID: 706,
  2148. Platform: PlatformGrok,
  2149. Type: AccountTypeAPIKey,
  2150. Concurrency: 1,
  2151. Credentials: map[string]any{
  2152. "api_key": "third-party-key",
  2153. "base_url": "https://grok.example.test/v1",
  2154. },
  2155. }
  2156. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2157. StatusCode: http.StatusOK,
  2158. Header: http.Header{"Content-Type": []string{"application/json"}},
  2159. Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl","object":"chat.completion","model":"grok-4.5","choices":[],"usage":{"prompt_tokens":1,"completion_tokens":1}}`)),
  2160. }}
  2161. svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
  2162. _, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  2163. require.NoError(t, err)
  2164. require.Equal(t, "https://grok.example.test/v1/chat/completions", upstream.lastReq.URL.String())
  2165. require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization"))
  2166. require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  2167. require.NotEqual(t, defaultGrokUpstreamUserAgent(), upstream.lastReq.Header.Get("User-Agent"))
  2168. }
  2169. func TestForwardAsChatCompletionsForGrokAPIKeyRejectsNonStreamingResponseWithoutUsage(t *testing.T) {
  2170. gin.SetMode(gin.TestMode)
  2171. recorder := httptest.NewRecorder()
  2172. c, _ := gin.CreateTestContext(recorder)
  2173. body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":false}`)
  2174. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  2175. c.Request.Header.Set("Content-Type", "application/json")
  2176. account := &Account{
  2177. ID: 707,
  2178. Platform: PlatformGrok,
  2179. Type: AccountTypeAPIKey,
  2180. Concurrency: 1,
  2181. Credentials: map[string]any{
  2182. "api_key": "third-party-key",
  2183. "base_url": "https://grok.example.test/v1",
  2184. },
  2185. }
  2186. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2187. StatusCode: http.StatusOK,
  2188. Header: http.Header{"Content-Type": []string{"application/json"}},
  2189. Body: io.NopCloser(strings.NewReader(
  2190. `{"id":"resp_missing_usage","object":"chat.completion","model":"grok-4.5","choices":[{"index":0,"message":{"role":"assistant","content":"ok"},"finish_reason":"stop"}]}`,
  2191. )),
  2192. }}
  2193. svc := &OpenAIGatewayService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
  2194. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  2195. require.Nil(t, result)
  2196. var failoverErr *UpstreamFailoverError
  2197. require.ErrorAs(t, err, &failoverErr)
  2198. require.Equal(t, http.StatusBadGateway, failoverErr.StatusCode)
  2199. require.Equal(t, grokMissingUsageErrorCode, gjson.GetBytes(failoverErr.ResponseBody, "error.code").String())
  2200. require.False(t, c.Writer.Written(), "an unbillable response must not be committed to the client")
  2201. require.Empty(t, recorder.Body.String())
  2202. }
  2203. func TestAccountTestServiceGrokAPIKeyUsesXAIResponses(t *testing.T) {
  2204. gin.SetMode(gin.TestMode)
  2205. account := &Account{
  2206. ID: 54,
  2207. Name: "grok-api-key",
  2208. Platform: PlatformGrok,
  2209. Type: AccountTypeAPIKey,
  2210. Concurrency: 2,
  2211. Credentials: map[string]any{
  2212. "api_key": "xai-test-key",
  2213. "base_url": "https://api.x.ai/v1",
  2214. },
  2215. }
  2216. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2217. StatusCode: http.StatusOK,
  2218. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  2219. Body: io.NopCloser(strings.NewReader(
  2220. "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
  2221. "data: {\"type\":\"response.completed\"}\n\n",
  2222. )),
  2223. }}
  2224. svc := &AccountTestService{httpUpstream: upstream}
  2225. recorder := httptest.NewRecorder()
  2226. c, _ := gin.CreateTestContext(recorder)
  2227. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/54/test", nil)
  2228. err := svc.testGrokAccountConnection(c, account, "grok", "", AccountTestModeDefault, AccountTestOptions{})
  2229. require.NoError(t, err)
  2230. require.Equal(t, "https://api.x.ai/v1/responses", upstream.lastReq.URL.String())
  2231. require.Equal(t, "Bearer xai-test-key", upstream.lastReq.Header.Get("Authorization"))
  2232. require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
  2233. }
  2234. func TestAccountTestServiceGrokAPIKeyAllowsConfiguredHTTPWhenGlobalPolicyDoes(t *testing.T) {
  2235. gin.SetMode(gin.TestMode)
  2236. account := &Account{
  2237. ID: 55,
  2238. Name: "grok-api-key-http",
  2239. Platform: PlatformGrok,
  2240. Type: AccountTypeAPIKey,
  2241. Concurrency: 1,
  2242. Credentials: map[string]any{
  2243. "api_key": "third-party-key",
  2244. "base_url": "http://grok.example.test/v1",
  2245. },
  2246. }
  2247. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2248. StatusCode: http.StatusOK,
  2249. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  2250. Body: io.NopCloser(strings.NewReader(
  2251. "data: {\"type\":\"response.output_text.delta\",\"delta\":\"ok\"}\n\n" +
  2252. "data: {\"type\":\"response.completed\"}\n\n",
  2253. )),
  2254. }}
  2255. svc := &AccountTestService{cfg: rawChatCompletionsTestConfig(), httpUpstream: upstream}
  2256. recorder := httptest.NewRecorder()
  2257. c, _ := gin.CreateTestContext(recorder)
  2258. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/55/test", nil)
  2259. err := svc.testGrokAccountConnection(c, account, "grok", "", AccountTestModeDefault, AccountTestOptions{})
  2260. require.NoError(t, err)
  2261. require.Equal(t, "http://grok.example.test/v1/responses", upstream.lastReq.URL.String())
  2262. require.Equal(t, "Bearer third-party-key", upstream.lastReq.Header.Get("Authorization"))
  2263. require.Empty(t, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  2264. require.Contains(t, recorder.Body.String(), `"type":"test_complete"`)
  2265. }
  2266. func TestAccountTestServiceGrokOAuthPaymentRequiredTemporarilyUnschedulesAccount(t *testing.T) {
  2267. gin.SetMode(gin.TestMode)
  2268. account := healthyGrokOAuthGatewayTestAccount(56, "access-token")
  2269. repo := &grokQuotaAccountRepo{}
  2270. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2271. StatusCode: http.StatusPaymentRequired,
  2272. Header: http.Header{"Content-Type": []string{"application/json"}},
  2273. Body: io.NopCloser(strings.NewReader(`{"code":"personal-team-blocked:spending-limit"}`)),
  2274. }}
  2275. svc := &AccountTestService{
  2276. accountRepo: repo,
  2277. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2278. httpUpstream: upstream,
  2279. }
  2280. recorder := httptest.NewRecorder()
  2281. c, _ := gin.CreateTestContext(recorder)
  2282. c.Request = httptest.NewRequest(http.MethodPost, "/api/v1/admin/accounts/56/test", nil)
  2283. before := time.Now()
  2284. err := svc.testGrokAccountConnection(c, account, "grok", "", AccountTestModeDefault, AccountTestOptions{})
  2285. require.Error(t, err)
  2286. require.Zero(t, repo.tempUnschedCalls)
  2287. require.Equal(t, 1, repo.rateLimitedCalls)
  2288. require.Equal(t, account.ID, repo.lastRateLimitedID)
  2289. require.WithinDuration(t, before.Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, time.Second)
  2290. require.Contains(t, recorder.Body.String(), `"type":"error"`)
  2291. require.Contains(t, recorder.Body.String(), "Grok Responses API returned 402")
  2292. }
  2293. func TestForwardAsChatCompletionsForGrokStreamingUsesRawXAIChatCompletions(t *testing.T) {
  2294. gin.SetMode(gin.TestMode)
  2295. recorder := httptest.NewRecorder()
  2296. c, _ := gin.CreateTestContext(recorder)
  2297. body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true}`)
  2298. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  2299. c.Request.Header.Set("Content-Type", "application/json")
  2300. account := healthyGrokOAuthGatewayTestAccount(53, "access-token")
  2301. repo := &grokQuotaAccountRepo{
  2302. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2303. accountsByID: map[int64]*Account{53: account},
  2304. },
  2305. }
  2306. upstreamBody := strings.Join([]string{
  2307. `data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[{"index":0,"delta":{"content":"ok"}}]}`,
  2308. "",
  2309. `data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":6,"completion_tokens":4,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":1}}}`,
  2310. "",
  2311. "data: [DONE]",
  2312. "",
  2313. }, "\n")
  2314. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2315. StatusCode: http.StatusOK,
  2316. Header: http.Header{
  2317. "Content-Type": []string{"text/event-stream"},
  2318. "X-Request-Id": []string{"chat-stream-req"},
  2319. "X-Ratelimit-Limit-Requests": []string{"10"},
  2320. "X-Ratelimit-Remaining-Requests": []string{"7"},
  2321. },
  2322. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  2323. }}
  2324. svc := &OpenAIGatewayService{
  2325. cfg: rawChatCompletionsTestConfig(),
  2326. httpUpstream: upstream,
  2327. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2328. accountRepo: repo,
  2329. }
  2330. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  2331. require.NoError(t, err)
  2332. require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
  2333. require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
  2334. require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
  2335. require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
  2336. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.lastBody, "model").String())
  2337. require.True(t, gjson.GetBytes(upstream.lastBody, "stream_options.include_usage").Bool())
  2338. require.True(t, result.Stream)
  2339. require.Equal(t, 6, result.Usage.InputTokens)
  2340. require.Equal(t, 4, result.Usage.OutputTokens)
  2341. require.Equal(t, 1, result.Usage.CacheReadInputTokens)
  2342. require.Contains(t, recorder.Body.String(), "data: [DONE]")
  2343. require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey])
  2344. }
  2345. func TestForwardGrokResponsesNonStreamingUsesCacheIdentityAndCachedUsage(t *testing.T) {
  2346. gin.SetMode(gin.TestMode)
  2347. recorder := httptest.NewRecorder()
  2348. c, _ := gin.CreateTestContext(recorder)
  2349. body := []byte(`{"model":"grok","input":"hi","stream":false,"tools":[{"type":"namespace","name":"client_tools"}],"tool_choice":{"type":"namespace","name":"client_tools"}}`)
  2350. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  2351. c.Request.Header.Set("Content-Type", "application/json")
  2352. c.Set("api_key", &APIKey{ID: 5202})
  2353. account := healthyGrokOAuthGatewayTestAccount(56, "access-token")
  2354. observedResetAt := time.Now().Add(-time.Second).UTC().Truncate(time.Second)
  2355. observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
  2356. account.RateLimitedAt = &observedLimitedAt
  2357. account.RateLimitResetAt = &observedResetAt
  2358. repo := &grokQuotaAccountRepo{
  2359. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2360. accountsByID: map[int64]*Account{56: account},
  2361. },
  2362. recoveryClearResult: true,
  2363. }
  2364. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2365. StatusCode: http.StatusOK,
  2366. Header: http.Header{
  2367. "Content-Type": []string{"application/json"},
  2368. "Xai-Request-Id": []string{"xai-non-stream-req"},
  2369. },
  2370. Body: io.NopCloser(strings.NewReader(`{"id":"resp_grok_non_stream","object":"response","model":"grok-4.3","status":"completed","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":7,"output_tokens":2,"total_tokens":9,"input_tokens_details":{"cached_tokens":4}}}`)),
  2371. }}
  2372. svc := &OpenAIGatewayService{
  2373. httpUpstream: upstream,
  2374. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2375. accountRepo: repo,
  2376. }
  2377. result, err := svc.forwardGrokResponses(context.Background(), c, account, body, "grok", false, time.Now())
  2378. require.NoError(t, err)
  2379. require.NotNil(t, result)
  2380. require.False(t, result.Stream)
  2381. require.Equal(t, "resp_grok_non_stream", result.ResponseID)
  2382. require.Equal(t, 7, result.Usage.InputTokens)
  2383. require.Equal(t, 2, result.Usage.OutputTokens)
  2384. require.Equal(t, 4, result.Usage.CacheReadInputTokens)
  2385. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.lastBody, "model").String())
  2386. identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
  2387. require.NotEmpty(t, identity)
  2388. require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
  2389. // The sanitizer drops this unsupported client tool, but its explicit intent
  2390. // must still prevent native cache-routing tools from being injected.
  2391. require.False(t, gjson.GetBytes(upstream.lastBody, "tools").Exists())
  2392. require.False(t, gjson.GetBytes(upstream.lastBody, "tool_choice").Exists())
  2393. require.Equal(t, "resp_grok_non_stream", gjson.Get(recorder.Body.String(), "id").String())
  2394. require.Equal(t, 1, repo.recoveryClearCalls)
  2395. require.Equal(t, observedLimitedAt, repo.recoveryObservedAt)
  2396. require.Equal(t, observedResetAt, repo.recoveryObservedReset)
  2397. }
  2398. func TestForwardGrokResponsesFailoverKeepsCacheIdentityAcrossAccounts(t *testing.T) {
  2399. gin.SetMode(gin.TestMode)
  2400. recorder := httptest.NewRecorder()
  2401. c, _ := gin.CreateTestContext(recorder)
  2402. body := []byte(`{"model":"grok","input":[{"role":"user","content":"stable prefix"}],"stream":false}`)
  2403. c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", bytes.NewReader(body))
  2404. c.Set("api_key", &APIKey{ID: 5203})
  2405. newAccount := func(id int64, token string) *Account {
  2406. account := healthyGrokOAuthGatewayTestAccount(id, token)
  2407. account.Name = fmt.Sprintf("grok-%d", id)
  2408. return account
  2409. }
  2410. firstAccount := newAccount(58, "access-token-a")
  2411. secondAccount := newAccount(59, "access-token-b")
  2412. repo := &grokQuotaAccountRepo{
  2413. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2414. accountsByID: map[int64]*Account{58: firstAccount, 59: secondAccount},
  2415. },
  2416. }
  2417. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  2418. {
  2419. StatusCode: http.StatusServiceUnavailable,
  2420. Header: http.Header{"Content-Type": []string{"application/json"}},
  2421. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"temporary"}}`)),
  2422. },
  2423. {
  2424. StatusCode: http.StatusOK,
  2425. Header: http.Header{"Content-Type": []string{"application/json"}},
  2426. Body: io.NopCloser(strings.NewReader(`{"id":"resp_after_failover","object":"response","model":"grok-4.3","status":"completed","output":[],"usage":{"input_tokens":5,"output_tokens":1}}`)),
  2427. },
  2428. }}
  2429. svc := &OpenAIGatewayService{
  2430. httpUpstream: upstream,
  2431. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2432. accountRepo: repo,
  2433. }
  2434. _, err := svc.forwardGrokResponses(context.Background(), c, firstAccount, body, "grok", false, time.Now())
  2435. var failoverErr *UpstreamFailoverError
  2436. require.ErrorAs(t, err, &failoverErr)
  2437. result, err := svc.forwardGrokResponses(context.Background(), c, secondAccount, body, "grok", false, time.Now())
  2438. require.NoError(t, err)
  2439. require.NotNil(t, result)
  2440. require.Len(t, upstream.requests, 2)
  2441. require.Len(t, upstream.bodies, 2)
  2442. firstIdentity := gjson.GetBytes(upstream.bodies[0], "prompt_cache_key").String()
  2443. secondIdentity := gjson.GetBytes(upstream.bodies[1], "prompt_cache_key").String()
  2444. require.NotEmpty(t, firstIdentity)
  2445. require.Equal(t, firstIdentity, secondIdentity)
  2446. require.Equal(t, firstIdentity, upstream.requests[0].Header.Get(grokConversationIDHeader))
  2447. require.Equal(t, secondIdentity, upstream.requests[1].Header.Get(grokConversationIDHeader))
  2448. require.Equal(t, "Bearer access-token-a", upstream.requests[0].Header.Get("Authorization"))
  2449. require.Equal(t, "Bearer access-token-b", upstream.requests[1].Header.Get("Authorization"))
  2450. }
  2451. func TestForwardAsChatCompletionsForGrokStreamingStopFallsBackToRawXAIChatCompletions(t *testing.T) {
  2452. gin.SetMode(gin.TestMode)
  2453. recorder := httptest.NewRecorder()
  2454. c, _ := gin.CreateTestContext(recorder)
  2455. body := []byte(`{"model":"grok","messages":[{"role":"user","content":"hi"}],"stream":true,"stop":"done"}`)
  2456. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  2457. c.Request.Header.Set("Content-Type", "application/json")
  2458. c.Request.Header.Set(grokConversationIDHeader, "native-client-conversation")
  2459. c.Set("api_key", &APIKey{ID: 5301})
  2460. account := healthyGrokOAuthGatewayTestAccount(53, "access-token")
  2461. repo := &grokQuotaAccountRepo{
  2462. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2463. accountsByID: map[int64]*Account{53: account},
  2464. },
  2465. }
  2466. upstreamBody := strings.Join([]string{
  2467. `data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[{"index":0,"delta":{"content":"ok"}}]}`,
  2468. "",
  2469. `data: {"id":"chatcmpl_grok","object":"chat.completion.chunk","model":"grok-4.3","choices":[],"usage":{"prompt_tokens":6,"completion_tokens":4,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":1}}}`,
  2470. "",
  2471. "data: [DONE]",
  2472. "",
  2473. }, "\n")
  2474. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2475. StatusCode: http.StatusOK,
  2476. Header: http.Header{
  2477. "Content-Type": []string{"text/event-stream"},
  2478. "X-Request-Id": []string{"chat-stream-req"},
  2479. "X-Ratelimit-Limit-Requests": []string{"10"},
  2480. "X-Ratelimit-Remaining-Requests": []string{"7"},
  2481. },
  2482. Body: io.NopCloser(strings.NewReader(upstreamBody)),
  2483. }}
  2484. svc := &OpenAIGatewayService{
  2485. cfg: rawChatCompletionsTestConfig(),
  2486. httpUpstream: upstream,
  2487. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2488. accountRepo: repo,
  2489. }
  2490. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  2491. require.NoError(t, err)
  2492. require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.lastReq.URL.String())
  2493. require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
  2494. require.Equal(t, "text/event-stream", upstream.lastReq.Header.Get("Accept"))
  2495. require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
  2496. require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  2497. require.NotEmpty(t, upstream.lastReq.Header.Get(grokConversationIDHeader))
  2498. require.NotEqual(t, "native-client-conversation", upstream.lastReq.Header.Get(grokConversationIDHeader))
  2499. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.lastBody, "model").String())
  2500. require.True(t, gjson.GetBytes(upstream.lastBody, "stream_options.include_usage").Bool())
  2501. require.True(t, result.Stream)
  2502. require.Equal(t, 6, result.Usage.InputTokens)
  2503. require.Equal(t, 4, result.Usage.OutputTokens)
  2504. require.Equal(t, 1, result.Usage.CacheReadInputTokens)
  2505. require.Contains(t, recorder.Body.String(), "data: [DONE]")
  2506. require.NotNil(t, repo.updates[53][grokQuotaSnapshotExtraKey])
  2507. }
  2508. func TestForwardAsChatCompletionsForGrokComposerBridgesImageInput(t *testing.T) {
  2509. gin.SetMode(gin.TestMode)
  2510. recorder := httptest.NewRecorder()
  2511. c, _ := gin.CreateTestContext(recorder)
  2512. body := []byte(`{"model":"grok-composer-2.5-fast","messages":[{"role":"system","content":"You are concise."},{"role":"user","content":[{"type":"text","text":"What is shown?"},{"type":"image_url","image_url":{"url":"data:image/png;base64,QUJD"}}]}],"metadata":{"large_id":9007199254740993},"stream":false}`)
  2513. c.Request = httptest.NewRequest(http.MethodPost, "/v1/chat/completions", bytes.NewReader(body))
  2514. c.Request.Header.Set("Content-Type", "application/json")
  2515. c.Set("api_key", &APIKey{ID: 5501})
  2516. account := healthyGrokOAuthGatewayTestAccount(55, "access-token")
  2517. repo := &grokQuotaAccountRepo{
  2518. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2519. accountsByID: map[int64]*Account{55: account},
  2520. },
  2521. }
  2522. upstream := &httpUpstreamRecorder{responses: []*http.Response{
  2523. {
  2524. StatusCode: http.StatusOK,
  2525. Header: http.Header{"Content-Type": []string{"application/json"}, "xai-request-id": []string{"vision-req"}},
  2526. Body: io.NopCloser(strings.NewReader(`{"id":"resp_vision","object":"response","model":"grok-build-0.1","output":[{"type":"message","role":"assistant","content":[{"type":"output_text","text":"A small diagram with ABC letters."}]}],"usage":{"input_tokens":11,"output_tokens":7,"total_tokens":18}}`)),
  2527. },
  2528. {
  2529. StatusCode: http.StatusOK,
  2530. Header: http.Header{
  2531. "Content-Type": []string{"application/json"},
  2532. "X-Request-Id": []string{"composer-req"},
  2533. "X-Ratelimit-Limit-Requests": []string{"10"},
  2534. "X-Ratelimit-Remaining-Requests": []string{"9"},
  2535. "X-Ratelimit-Limit-Tokens": []string{"1000"},
  2536. "X-Ratelimit-Remaining-Tokens": []string{"980"},
  2537. },
  2538. Body: io.NopCloser(strings.NewReader(`{"id":"chatcmpl_composer","object":"chat.completion","model":"grok-composer-2.5-fast","choices":[{"index":0,"message":{"role":"assistant","content":"It shows ABC."},"finish_reason":"stop"}],"usage":{"prompt_tokens":3,"completion_tokens":5,"total_tokens":8}}`)),
  2539. },
  2540. }}
  2541. svc := &OpenAIGatewayService{
  2542. cfg: rawChatCompletionsTestConfig(),
  2543. httpUpstream: upstream,
  2544. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2545. accountRepo: repo,
  2546. }
  2547. result, err := svc.ForwardAsChatCompletions(context.Background(), c, account, body, "", "")
  2548. require.NoError(t, err)
  2549. require.NotNil(t, result)
  2550. require.Len(t, upstream.requests, 2)
  2551. require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.requests[0].URL.String())
  2552. require.Empty(t, upstream.requests[0].Header.Get(grokConversationIDHeader))
  2553. require.Equal(t, "grok-build-0.1", gjson.GetBytes(upstream.bodies[0], "model").String())
  2554. require.Equal(t, "input_image", gjson.GetBytes(upstream.bodies[0], "input.0.content.1.type").String())
  2555. require.Equal(t, xai.DefaultCLIBaseURL+"/chat/completions", upstream.requests[1].URL.String())
  2556. require.NotEmpty(t, upstream.requests[1].Header.Get(grokConversationIDHeader))
  2557. require.Equal(t, "grok-composer-2.5-fast", gjson.GetBytes(upstream.bodies[1], "model").String())
  2558. require.False(t, strings.Contains(string(upstream.bodies[1]), "image_url"))
  2559. require.Equal(t, "9007199254740993", gjson.GetBytes(upstream.bodies[1], "metadata.large_id").Raw)
  2560. require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "Image 1 description")
  2561. require.Contains(t, gjson.GetBytes(upstream.bodies[1], "messages.1.content").String(), "A small diagram with ABC letters.")
  2562. require.Equal(t, 14, result.Usage.InputTokens)
  2563. require.Equal(t, 12, result.Usage.OutputTokens)
  2564. require.Equal(t, "It shows ABC.", gjson.Get(recorder.Body.String(), "choices.0.message.content").String())
  2565. require.NotNil(t, repo.updates[55][grokQuotaSnapshotExtraKey])
  2566. }
  2567. func TestForwardAsAnthropicForGrokUsesXAIResponses(t *testing.T) {
  2568. gin.SetMode(gin.TestMode)
  2569. recorder := httptest.NewRecorder()
  2570. c, _ := gin.CreateTestContext(recorder)
  2571. body := []byte(`{"model":"grok","max_tokens":32,"stream":false,"messages":[{"role":"user","content":"hi"}]}`)
  2572. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  2573. c.Set("api_key", &APIKey{ID: 5401})
  2574. c.Request.Header.Set("OpenAI-Beta", "grok-experimental")
  2575. c.Request.Header.Set("originator", "opencode")
  2576. account := healthyGrokOAuthGatewayTestAccount(54, "access-token")
  2577. repo := &grokQuotaAccountRepo{
  2578. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2579. accountsByID: map[int64]*Account{54: account},
  2580. },
  2581. }
  2582. upstream := &httpUpstreamRecorder{resp: grokMessagesSSECompletedResponse("resp_grok_messages", 3)}
  2583. svc := &OpenAIGatewayService{
  2584. httpUpstream: upstream,
  2585. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2586. accountRepo: repo,
  2587. }
  2588. result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
  2589. require.NoError(t, err)
  2590. require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
  2591. require.Equal(t, "Bearer access-token", upstream.lastReq.Header.Get("Authorization"))
  2592. require.Equal(t, xai.CLIUserAgent(xai.CLIClientVersion), upstream.lastReq.Header.Get("User-Agent"))
  2593. require.Equal(t, grokCLIVersion, upstream.lastReq.Header.Get("X-Grok-Client-Version"))
  2594. require.Equal(t, "grok-experimental", upstream.lastReq.Header.Get("OpenAI-Beta"))
  2595. require.Empty(t, upstream.lastReq.Header.Get("originator"))
  2596. require.Empty(t, upstream.lastReq.Header.Get("version"))
  2597. require.Equal(t, "grok-4.6", gjson.GetBytes(upstream.lastBody, "model").String())
  2598. require.NotEmpty(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String())
  2599. require.Equal(t, gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String(), upstream.lastReq.Header.Get(grokConversationIDHeader))
  2600. require.Equal(t, "web_search", gjson.GetBytes(upstream.lastBody, "tools.0.type").String())
  2601. require.Equal(t, "x_search", gjson.GetBytes(upstream.lastBody, "tools.1.type").String())
  2602. require.Equal(t, "none", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
  2603. require.Empty(t, upstream.lastReq.Header.Get("session_id"))
  2604. require.True(t, gjson.GetBytes(upstream.lastBody, "stream").Bool())
  2605. require.NotContains(t, string(upstream.lastBody), "chatgpt.com")
  2606. require.Equal(t, "grok", result.Model)
  2607. require.Equal(t, "grok-4.6", result.UpstreamModel)
  2608. require.Equal(t, 5, result.Usage.InputTokens)
  2609. require.Equal(t, 2, result.Usage.OutputTokens)
  2610. require.Equal(t, 3, result.Usage.CacheReadInputTokens)
  2611. require.Contains(t, recorder.Body.String(), `"type":"message"`)
  2612. require.Equal(t, int64(3), gjson.Get(recorder.Body.String(), "usage.cache_read_input_tokens").Int())
  2613. require.Contains(t, recorder.Body.String(), "ok")
  2614. }
  2615. func TestForwardAsAnthropicForGrokFunctionToolUsesCacheCapableMixedRoute(t *testing.T) {
  2616. gin.SetMode(gin.TestMode)
  2617. recorder := httptest.NewRecorder()
  2618. c, _ := gin.CreateTestContext(recorder)
  2619. body := []byte(`{
  2620. "model":"grok","max_tokens":32,"stream":false,
  2621. "messages":[{"role":"user","content":"look up alpha"}],
  2622. "tools":[{"name":"lookup","description":"look up a key","input_schema":{"type":"object","properties":{"key":{"type":"string"}},"required":["key"]}},{"name":"web_search","description":"search the web","input_schema":{"type":"object","properties":{"query":{"type":"string"}},"required":["query"]}}],
  2623. "tool_choice":{"type":"auto"}
  2624. }`)
  2625. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  2626. c.Set("api_key", &APIKey{ID: 5403})
  2627. account := healthyGrokOAuthGatewayTestAccount(58, "access-token")
  2628. account.Extra = map[string]any{grokBillingExtraKey: map[string]any{
  2629. "status_code": http.StatusOK,
  2630. "source": "billing_probe",
  2631. "monthly_updated_at": "2026-07-15T05:00:00Z",
  2632. }}
  2633. repo := &grokQuotaAccountRepo{
  2634. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2635. accountsByID: map[int64]*Account{58: account},
  2636. },
  2637. }
  2638. responseBody := strings.Join([]string{
  2639. `data: {"type":"response.completed","response":{"id":"resp_grok_function","object":"response","model":"grok-4.5","status":"completed","output":[{"type":"function_call","id":"fc_lookup","call_id":"call_lookup","name":"lookup","arguments":"{\"key\":\"alpha\"}","status":"completed"}],"usage":{"input_tokens":7000,"output_tokens":2,"total_tokens":7002,"input_tokens_details":{"cached_tokens":6144}}}}`,
  2640. "",
  2641. "data: [DONE]",
  2642. "",
  2643. }, "\n")
  2644. upstream := &httpUpstreamRecorder{resp: &http.Response{
  2645. StatusCode: http.StatusOK,
  2646. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  2647. Body: io.NopCloser(strings.NewReader(responseBody)),
  2648. }}
  2649. svc := &OpenAIGatewayService{
  2650. httpUpstream: upstream,
  2651. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2652. accountRepo: repo,
  2653. }
  2654. result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
  2655. require.NoError(t, err)
  2656. require.NotNil(t, result)
  2657. require.Equal(t, xai.DefaultCLIBaseURL+"/responses", upstream.lastReq.URL.String())
  2658. identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
  2659. require.NotEmpty(t, identity)
  2660. require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
  2661. tools := gjson.GetBytes(upstream.lastBody, "tools").Array()
  2662. require.Len(t, tools, 3)
  2663. require.Equal(t, "function", tools[0].Get("type").String())
  2664. require.Equal(t, "lookup", tools[0].Get("name").String())
  2665. require.Equal(t, "object", tools[0].Get("parameters.type").String())
  2666. require.Equal(t, "web_search", tools[1].Get("type").String())
  2667. require.Equal(t, "x_search", tools[2].Get("type").String())
  2668. require.Equal(t, "auto", gjson.GetBytes(upstream.lastBody, "tool_choice").String())
  2669. require.Equal(t, 7000, result.Usage.InputTokens)
  2670. require.Equal(t, 6144, result.Usage.CacheReadInputTokens)
  2671. clientBody := recorder.Body.String()
  2672. require.Equal(t, "tool_use", gjson.Get(clientBody, "content.0.type").String())
  2673. require.Equal(t, "call_lookup", gjson.Get(clientBody, "content.0.id").String())
  2674. require.Equal(t, "lookup", gjson.Get(clientBody, "content.0.name").String())
  2675. require.Equal(t, "alpha", gjson.Get(clientBody, "content.0.input.key").String())
  2676. require.Equal(t, "tool_use", gjson.Get(clientBody, "stop_reason").String())
  2677. require.Equal(t, int64(856), gjson.Get(clientBody, "usage.input_tokens").Int())
  2678. require.Equal(t, int64(6144), gjson.Get(clientBody, "usage.cache_read_input_tokens").Int())
  2679. }
  2680. func TestForwardAsAnthropicForGrokStreamingPreservesCacheUsage(t *testing.T) {
  2681. gin.SetMode(gin.TestMode)
  2682. recorder := httptest.NewRecorder()
  2683. c, _ := gin.CreateTestContext(recorder)
  2684. body := []byte(`{"model":"grok","max_tokens":32,"stream":true,"messages":[{"role":"user","content":"hi"}]}`)
  2685. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", bytes.NewReader(body))
  2686. c.Set("api_key", &APIKey{ID: 5402})
  2687. account := healthyGrokOAuthGatewayTestAccount(57, "access-token")
  2688. repo := &grokQuotaAccountRepo{
  2689. mockAccountRepoForPlatform: &mockAccountRepoForPlatform{
  2690. accountsByID: map[int64]*Account{57: account},
  2691. },
  2692. }
  2693. upstream := &httpUpstreamRecorder{resp: grokMessagesSSECompletedResponse("resp_grok_messages_stream", 2)}
  2694. svc := &OpenAIGatewayService{
  2695. httpUpstream: upstream,
  2696. grokTokenProvider: NewGrokTokenProvider(repo, nil),
  2697. accountRepo: repo,
  2698. }
  2699. result, err := svc.ForwardAsAnthropic(context.Background(), c, account, body, "", "")
  2700. require.NoError(t, err)
  2701. require.NotNil(t, result)
  2702. require.Equal(t, 2, result.Usage.CacheReadInputTokens)
  2703. identity := gjson.GetBytes(upstream.lastBody, "prompt_cache_key").String()
  2704. require.NotEmpty(t, identity)
  2705. require.Equal(t, identity, upstream.lastReq.Header.Get(grokConversationIDHeader))
  2706. require.Contains(t, recorder.Header().Get("Content-Type"), "text/event-stream")
  2707. require.Contains(t, recorder.Body.String(), `"cache_read_input_tokens":2`)
  2708. }
  2709. func grokMessagesSSECompletedResponse(responseID string, cachedTokens int) *http.Response {
  2710. body := strings.Join([]string{
  2711. fmt.Sprintf(`data: {"type":"response.completed","response":{"id":%q,"object":"response","model":"grok-4.3","status":"completed","output":[{"type":"message","id":"msg_1","role":"assistant","status":"completed","content":[{"type":"output_text","text":"ok"}]}],"usage":{"input_tokens":5,"output_tokens":2,"total_tokens":7,"input_tokens_details":{"cached_tokens":%d}}}}`, responseID, cachedTokens),
  2712. "",
  2713. "data: [DONE]",
  2714. "",
  2715. }, "\n")
  2716. return &http.Response{
  2717. StatusCode: http.StatusOK,
  2718. Header: http.Header{"Content-Type": []string{"text/event-stream"}},
  2719. Body: io.NopCloser(strings.NewReader(body)),
  2720. }
  2721. }
  2722. func TestHandleGrokAccountUpstreamErrorSpendingLimitUsesRecoverableProbeCool(t *testing.T) {
  2723. repo := &grokQuotaAccountRepo{}
  2724. svc := &OpenAIGatewayService{accountRepo: repo}
  2725. account := &Account{ID: 2570, Platform: PlatformGrok, Type: AccountTypeOAuth}
  2726. before := time.Now()
  2727. body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"You have run out of credits"}`)
  2728. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body)
  2729. require.Equal(t, 1, repo.rateLimitedCalls)
  2730. require.WithinDuration(t, before.Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, 2*time.Second)
  2731. require.Zero(t, repo.tempUnschedCalls)
  2732. }
  2733. func TestHandleGrokAccountUpstreamErrorTempUnschedulesNonRateLimitStates(t *testing.T) {
  2734. tests := []struct {
  2735. name string
  2736. status int
  2737. headers http.Header
  2738. wantReason string
  2739. wantMinCooldown time.Duration
  2740. wantMaxCooldown time.Duration
  2741. }{
  2742. {
  2743. name: "unauthorized reauth",
  2744. status: http.StatusUnauthorized,
  2745. wantReason: "grok credentials unauthorized",
  2746. wantMinCooldown: 10*time.Minute - time.Second,
  2747. wantMaxCooldown: 10*time.Minute + time.Second,
  2748. },
  2749. {
  2750. name: "forbidden entitlement",
  2751. status: http.StatusForbidden,
  2752. wantReason: "grok access or entitlement denied",
  2753. wantMinCooldown: 30*time.Minute - time.Second,
  2754. wantMaxCooldown: 30*time.Minute + time.Second,
  2755. },
  2756. {
  2757. name: "payment required",
  2758. status: http.StatusPaymentRequired,
  2759. wantReason: "grok payment required",
  2760. wantMinCooldown: 30*time.Minute - time.Second,
  2761. wantMaxCooldown: 30*time.Minute + time.Second,
  2762. },
  2763. {
  2764. name: "upstream temporary error",
  2765. status: http.StatusInternalServerError,
  2766. wantReason: "grok upstream temporary error",
  2767. wantMinCooldown: 2*time.Minute - time.Second,
  2768. wantMaxCooldown: 2*time.Minute + time.Second,
  2769. },
  2770. }
  2771. for _, tt := range tests {
  2772. t.Run(tt.name, func(t *testing.T) {
  2773. account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth}
  2774. repo := &grokQuotaAccountRepo{}
  2775. svc := &OpenAIGatewayService{accountRepo: repo}
  2776. before := time.Now()
  2777. svc.handleGrokAccountUpstreamError(context.Background(), account, tt.status, tt.headers, nil)
  2778. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2779. require.Equal(t, 1, repo.tempUnschedCalls)
  2780. require.Zero(t, repo.rateLimitedCalls)
  2781. require.Equal(t, account.ID, repo.lastTempUnschedID)
  2782. require.Equal(t, tt.wantReason, repo.lastTempUnschedReason)
  2783. require.True(t, repo.lastTempUnschedUntil.After(before.Add(tt.wantMinCooldown)))
  2784. require.True(t, repo.lastTempUnschedUntil.Before(before.Add(tt.wantMaxCooldown)))
  2785. })
  2786. }
  2787. }
  2788. func TestHandleGrokAccountUpstreamErrorSpendingLimit403RateLimits(t *testing.T) {
  2789. account := &Account{ID: 614, Platform: PlatformGrok, Type: AccountTypeOAuth}
  2790. repo := &grokQuotaAccountRepo{}
  2791. svc := &OpenAIGatewayService{accountRepo: repo}
  2792. before := time.Now()
  2793. body := []byte(`{"code":"personal-team-blocked:spending-limit","error":"You have run out of credits"}`)
  2794. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusForbidden, nil, body)
  2795. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2796. require.Equal(t, 1, repo.rateLimitedCalls)
  2797. require.Equal(t, account.ID, repo.lastRateLimitedID)
  2798. require.WithinDuration(t, before.Add(grokSpendingLimitProbeCooldown), repo.lastRateLimitResetAt, 2*time.Second)
  2799. require.Zero(t, repo.tempUnschedCalls)
  2800. require.True(t, isGrokSpendingLimitError(body))
  2801. }
  2802. func TestHandleGrokAccountUpstreamError5xxRespectsPoolMode(t *testing.T) {
  2803. t.Run("pool mode keeps scheduling state", func(t *testing.T) {
  2804. account := &Account{
  2805. ID: 611,
  2806. Platform: PlatformGrok,
  2807. Type: AccountTypeAPIKey,
  2808. Credentials: map[string]any{
  2809. "pool_mode": true,
  2810. },
  2811. }
  2812. repo := &grokQuotaAccountRepo{}
  2813. svc := &OpenAIGatewayService{accountRepo: repo}
  2814. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil)
  2815. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2816. require.Zero(t, repo.tempUnschedCalls)
  2817. require.Nil(t, account.TempUnschedulableUntil)
  2818. require.Empty(t, account.TempUnschedulableReason)
  2819. })
  2820. t.Run("non-pool mode keeps two minute cooldown", func(t *testing.T) {
  2821. account := &Account{ID: 612, Platform: PlatformGrok, Type: AccountTypeAPIKey}
  2822. repo := &grokQuotaAccountRepo{}
  2823. svc := &OpenAIGatewayService{accountRepo: repo}
  2824. before := time.Now()
  2825. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusBadGateway, nil, nil)
  2826. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2827. require.Equal(t, 1, repo.tempUnschedCalls)
  2828. require.Equal(t, account.ID, repo.lastTempUnschedID)
  2829. require.Equal(t, "grok upstream temporary error", repo.lastTempUnschedReason)
  2830. require.WithinDuration(t, before.Add(2*time.Minute), repo.lastTempUnschedUntil, time.Second)
  2831. })
  2832. }
  2833. func TestHandleGrokAccountUpstreamError429SetsRateLimitedFromRetryAfter(t *testing.T) {
  2834. account := &Account{ID: 61, Platform: PlatformGrok, Type: AccountTypeOAuth}
  2835. repo := &grokQuotaAccountRepo{}
  2836. svc := &OpenAIGatewayService{accountRepo: repo}
  2837. before := time.Now()
  2838. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil)
  2839. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2840. require.Equal(t, 1, repo.rateLimitedCalls)
  2841. require.Equal(t, account.ID, repo.lastRateLimitedID)
  2842. require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
  2843. require.Zero(t, repo.tempUnschedCalls)
  2844. }
  2845. func TestHandleGrokAccountUpstreamError429PoolModeKeepsSchedulingState(t *testing.T) {
  2846. account := &Account{
  2847. ID: 613,
  2848. Platform: PlatformGrok,
  2849. Type: AccountTypeAPIKey,
  2850. Credentials: map[string]any{
  2851. "pool_mode": true,
  2852. },
  2853. }
  2854. repo := &grokQuotaAccountRepo{}
  2855. svc := &OpenAIGatewayService{accountRepo: repo}
  2856. svc.handleGrokAccountUpstreamError(
  2857. context.Background(), account, http.StatusTooManyRequests,
  2858. http.Header{"Retry-After": []string{"45"}}, nil,
  2859. )
  2860. require.Equal(t, 1, repo.updateCalls, "pool mode should retain the quota snapshot for observability")
  2861. require.Zero(t, repo.rateLimitedCalls)
  2862. require.Zero(t, repo.tempUnschedCalls)
  2863. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2864. require.Nil(t, account.RateLimitResetAt)
  2865. }
  2866. func TestHandleGrokAccountUpstreamError402RecoversAfterCooldownExpiry(t *testing.T) {
  2867. account := &Account{
  2868. ID: 610, Platform: PlatformGrok, Type: AccountTypeOAuth,
  2869. Status: StatusActive, Schedulable: true,
  2870. }
  2871. repo := &grokQuotaAccountRepo{}
  2872. svc := &OpenAIGatewayService{accountRepo: repo}
  2873. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusPaymentRequired, nil, nil)
  2874. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2875. require.Equal(t, 1, repo.tempUnschedCalls)
  2876. expired := time.Now().Add(-time.Second)
  2877. account.TempUnschedulableUntil = &expired
  2878. svc.openaiAccountRuntimeBlockUntil.Store(account.ID, expired)
  2879. require.False(t, svc.isOpenAIAccountRuntimeBlocked(account))
  2880. require.True(t, account.IsSchedulable())
  2881. }
  2882. func TestHandleGrokAccountUpstreamError429UsesLatestExhaustedWindowReset(t *testing.T) {
  2883. now := time.Now()
  2884. requestReset := now.Add(10 * time.Minute).Truncate(time.Second)
  2885. tokenReset := now.Add(20 * time.Minute).Truncate(time.Second)
  2886. headers := http.Header{
  2887. "X-Ratelimit-Limit-Requests": []string{"10"},
  2888. "X-Ratelimit-Remaining-Requests": []string{"0"},
  2889. "X-Ratelimit-Reset-Requests": []string{fmt.Sprintf("%d", requestReset.Unix())},
  2890. "X-Ratelimit-Limit-Tokens": []string{"1000"},
  2891. "X-Ratelimit-Remaining-Tokens": []string{"0"},
  2892. "X-Ratelimit-Reset-Tokens": []string{fmt.Sprintf("%d", tokenReset.Unix())},
  2893. }
  2894. account := &Account{ID: 62, Platform: PlatformGrok, Type: AccountTypeOAuth}
  2895. repo := &grokQuotaAccountRepo{}
  2896. svc := &OpenAIGatewayService{accountRepo: repo}
  2897. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, headers, nil)
  2898. require.Equal(t, 1, repo.rateLimitedCalls)
  2899. require.WithinDuration(t, tokenReset, repo.lastRateLimitResetAt, time.Second)
  2900. require.Zero(t, repo.tempUnschedCalls)
  2901. }
  2902. func TestHandleGrokAccountUpstreamError429UsesFallbackReset(t *testing.T) {
  2903. account := &Account{ID: 63, Platform: PlatformGrok, Type: AccountTypeOAuth}
  2904. repo := &grokQuotaAccountRepo{}
  2905. svc := &OpenAIGatewayService{accountRepo: repo}
  2906. before := time.Now()
  2907. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, nil, nil)
  2908. require.Equal(t, 1, repo.rateLimitedCalls)
  2909. require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second)
  2910. require.Zero(t, repo.tempUnschedCalls)
  2911. }
  2912. func TestGrokRateLimitResetAtForAccountEscalatesRepeated429s(t *testing.T) {
  2913. now := time.Now().UTC().Truncate(time.Second)
  2914. retryAfter := 45
  2915. snapshot := &xai.QuotaSnapshot{
  2916. StatusCode: http.StatusTooManyRequests,
  2917. RetryAfterSeconds: &retryAfter,
  2918. UpdatedAt: now.Format(time.RFC3339),
  2919. }
  2920. tests := []struct {
  2921. name string
  2922. previousCooldown time.Duration
  2923. wantCooldown time.Duration
  2924. }{
  2925. {name: "repeat after short boundary", previousCooldown: 45 * time.Second, wantCooldown: grokRateLimitRepeatCooldown},
  2926. {name: "sustained repeat", previousCooldown: grokRateLimitRepeatCooldown, wantCooldown: grokRateLimitSustainedCooldown},
  2927. {name: "capped repeat", previousCooldown: grokRateLimitSustainedCooldown, wantCooldown: grokRateLimitMaxAdaptiveCooldown},
  2928. }
  2929. for _, tt := range tests {
  2930. t.Run(tt.name, func(t *testing.T) {
  2931. previousReset := now.Add(-time.Second)
  2932. previousLimited := previousReset.Add(-tt.previousCooldown)
  2933. account := &Account{
  2934. ID: 630,
  2935. Platform: PlatformGrok,
  2936. Type: AccountTypeOAuth,
  2937. RateLimitedAt: &previousLimited,
  2938. RateLimitResetAt: &previousReset,
  2939. }
  2940. resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
  2941. require.True(t, limited)
  2942. require.WithinDuration(t, now.Add(tt.wantCooldown), resetAt, time.Second)
  2943. })
  2944. }
  2945. }
  2946. func TestGrokRateLimitResetAtForAccountPreservesAuthoritativeAndQuietRecovery(t *testing.T) {
  2947. now := time.Now().UTC().Truncate(time.Second)
  2948. retryAfter := 45
  2949. previousReset := now.Add(-grokRateLimitBackoffQuietPeriod - time.Second)
  2950. previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown)
  2951. account := &Account{
  2952. ID: 631,
  2953. Platform: PlatformGrok,
  2954. Type: AccountTypeOAuth,
  2955. RateLimitedAt: &previousLimited,
  2956. RateLimitResetAt: &previousReset,
  2957. }
  2958. snapshot := &xai.QuotaSnapshot{
  2959. StatusCode: http.StatusTooManyRequests,
  2960. RetryAfterSeconds: &retryAfter,
  2961. UpdatedAt: now.Format(time.RFC3339),
  2962. }
  2963. resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
  2964. require.True(t, limited)
  2965. require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second)
  2966. authoritativeReset := now.Add(2 * time.Hour)
  2967. remaining := int64(0)
  2968. snapshot.Requests = &xai.QuotaWindow{Remaining: &remaining, ResetUnix: grokInt64PtrForTest(authoritativeReset.Unix())}
  2969. recentReset := now.Add(-time.Second)
  2970. recentLimited := recentReset.Add(-grokRateLimitSustainedCooldown)
  2971. account.RateLimitResetAt = &recentReset
  2972. account.RateLimitedAt = &recentLimited
  2973. resetAt, limited = grokRateLimitResetAtForAccount(account, snapshot, now)
  2974. require.True(t, limited)
  2975. require.WithinDuration(t, authoritativeReset, resetAt, time.Second)
  2976. }
  2977. func TestGrokRateLimitResetAtForAccountLeavesAPIKey429PolicyUnchanged(t *testing.T) {
  2978. now := time.Now().UTC().Truncate(time.Second)
  2979. retryAfter := 45
  2980. previousReset := now.Add(-time.Second)
  2981. previousLimited := previousReset.Add(-grokRateLimitSustainedCooldown)
  2982. account := &Account{
  2983. ID: 632,
  2984. Platform: PlatformGrok,
  2985. Type: AccountTypeAPIKey,
  2986. RateLimitedAt: &previousLimited,
  2987. RateLimitResetAt: &previousReset,
  2988. }
  2989. snapshot := &xai.QuotaSnapshot{
  2990. StatusCode: http.StatusTooManyRequests,
  2991. RetryAfterSeconds: &retryAfter,
  2992. UpdatedAt: now.Format(time.RFC3339),
  2993. }
  2994. resetAt, limited := grokRateLimitResetAtForAccount(account, snapshot, now)
  2995. require.True(t, limited)
  2996. require.WithinDuration(t, now.Add(45*time.Second), resetAt, time.Second)
  2997. }
  2998. func TestGrokRateLimitResetAtUsesFutureWindowAfterRetryAfterExpires(t *testing.T) {
  2999. now := time.Now().UTC().Truncate(time.Second)
  3000. observedAt := now.Add(-2 * time.Minute)
  3001. windowReset := now.Add(15 * time.Minute)
  3002. retryAfter := 30
  3003. snapshot := &xai.QuotaSnapshot{
  3004. StatusCode: http.StatusTooManyRequests,
  3005. UpdatedAt: observedAt.Format(time.RFC3339),
  3006. RetryAfterSeconds: &retryAfter,
  3007. Requests: &xai.QuotaWindow{
  3008. Limit: grokInt64PtrForTest(10),
  3009. Remaining: grokInt64PtrForTest(0),
  3010. ResetUnix: grokInt64PtrForTest(windowReset.Unix()),
  3011. },
  3012. }
  3013. resetAt, limited := grokRateLimitResetAt(snapshot, now)
  3014. require.True(t, limited)
  3015. require.WithinDuration(t, windowReset, resetAt, time.Second)
  3016. }
  3017. func TestHandleGrokAccountUpstreamError429DoesNotShortenExistingPause(t *testing.T) {
  3018. existingUntil := time.Now().Add(15 * time.Minute)
  3019. account := &Account{
  3020. ID: 64,
  3021. Platform: PlatformGrok,
  3022. Type: AccountTypeOAuth,
  3023. TempUnschedulableUntil: &existingUntil,
  3024. TempUnschedulableReason: "existing pause",
  3025. }
  3026. repo := &grokQuotaAccountRepo{}
  3027. svc := &OpenAIGatewayService{accountRepo: repo}
  3028. svc.handleGrokAccountUpstreamError(context.Background(), account, http.StatusTooManyRequests, http.Header{"Retry-After": []string{"45"}}, nil)
  3029. require.Equal(t, 1, repo.rateLimitedCalls)
  3030. require.WithinDuration(t, time.Now().Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
  3031. require.Zero(t, repo.tempUnschedCalls)
  3032. value, ok := svc.openaiAccountRuntimeBlockUntil.Load(account.ID)
  3033. require.True(t, ok)
  3034. runtimeUntil, ok := value.(time.Time)
  3035. require.True(t, ok)
  3036. require.WithinDuration(t, existingUntil, runtimeUntil, time.Second)
  3037. }
  3038. func TestUpdateGrokUsageSnapshotExhaustedSuccessBypassesThrottleAndSetsRateLimited(t *testing.T) {
  3039. account := &Account{ID: 65, Platform: PlatformGrok, Type: AccountTypeOAuth}
  3040. repo := &grokQuotaAccountRepo{}
  3041. svc := &OpenAIGatewayService{
  3042. accountRepo: repo,
  3043. codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
  3044. }
  3045. now := time.Now()
  3046. // Consume the normal snapshot write allowance first.
  3047. svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
  3048. StatusCode: http.StatusOK,
  3049. Requests: &xai.QuotaWindow{
  3050. Limit: grokInt64PtrForTest(10),
  3051. Remaining: grokInt64PtrForTest(9),
  3052. },
  3053. UpdatedAt: now.UTC().Format(time.RFC3339),
  3054. })
  3055. resetAt := now.Add(30 * time.Minute).Truncate(time.Second)
  3056. svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
  3057. StatusCode: http.StatusOK,
  3058. Requests: &xai.QuotaWindow{
  3059. Limit: grokInt64PtrForTest(10),
  3060. Remaining: grokInt64PtrForTest(0),
  3061. ResetUnix: grokInt64PtrForTest(resetAt.Unix()),
  3062. ResetAt: resetAt.UTC().Format(time.RFC3339),
  3063. },
  3064. UpdatedAt: now.UTC().Format(time.RFC3339),
  3065. })
  3066. require.Equal(t, 2, repo.updateCalls)
  3067. require.Equal(t, 1, repo.rateLimitedCalls)
  3068. require.Equal(t, account.ID, repo.lastRateLimitedID)
  3069. require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second)
  3070. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  3071. }
  3072. func TestUpdateGrokUsageSnapshotAvailableSuccessDoesNotSetRateLimited(t *testing.T) {
  3073. repo := &grokQuotaAccountRepo{}
  3074. svc := &OpenAIGatewayService{accountRepo: repo}
  3075. account := &Account{ID: 66, Platform: PlatformGrok, Type: AccountTypeOAuth}
  3076. svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
  3077. StatusCode: http.StatusOK,
  3078. Requests: &xai.QuotaWindow{
  3079. Limit: grokInt64PtrForTest(10),
  3080. Remaining: grokInt64PtrForTest(1),
  3081. },
  3082. UpdatedAt: time.Now().UTC().Format(time.RFC3339),
  3083. })
  3084. require.Equal(t, 1, repo.updateCalls)
  3085. require.Zero(t, repo.rateLimitedCalls)
  3086. }
  3087. func TestUpdateGrokUsageFromResponseHeaderlessSuccessClearsObservedCooldown(t *testing.T) {
  3088. now := time.Now().UTC().Truncate(time.Second)
  3089. limitedAt := now.Add(-grokRateLimitRepeatCooldown)
  3090. observedResetAt := now.Add(-time.Second)
  3091. account := &Account{
  3092. ID: 660,
  3093. Platform: PlatformGrok,
  3094. Type: AccountTypeOAuth,
  3095. RateLimitedAt: &limitedAt,
  3096. RateLimitResetAt: &observedResetAt,
  3097. }
  3098. repo := &grokQuotaAccountRepo{recoveryClearResult: true}
  3099. svc := &OpenAIGatewayService{
  3100. accountRepo: repo,
  3101. codexSnapshotThrottle: newAccountWriteThrottle(time.Hour),
  3102. }
  3103. svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK)
  3104. require.Zero(t, repo.updateCalls, "headerless success must not overwrite an informative quota snapshot")
  3105. require.Equal(t, 1, repo.recoveryClearCalls)
  3106. require.Equal(t, limitedAt, repo.recoveryObservedAt)
  3107. require.Equal(t, observedResetAt, repo.recoveryObservedReset)
  3108. require.Same(t, &observedResetAt, account.RateLimitResetAt, "shared account snapshots must not be mutated in place")
  3109. }
  3110. func TestUpdateGrokUsageFromResponseRecoveryRespectsCancellationAndAPIKeyBoundary(t *testing.T) {
  3111. now := time.Now().UTC().Truncate(time.Second)
  3112. observedResetAt := now.Add(-time.Second)
  3113. observedLimitedAt := observedResetAt.Add(-grokRateLimitRepeatCooldown)
  3114. t.Run("parent cancellation does not mutate account state", func(t *testing.T) {
  3115. ctx, cancel := context.WithCancel(context.Background())
  3116. cancel()
  3117. account := &Account{
  3118. ID: 661,
  3119. Platform: PlatformGrok,
  3120. Type: AccountTypeOAuth,
  3121. RateLimitedAt: &observedLimitedAt,
  3122. RateLimitResetAt: &observedResetAt,
  3123. }
  3124. repo := &grokQuotaAccountRepo{recoveryClearResult: true}
  3125. svc := &OpenAIGatewayService{accountRepo: repo}
  3126. svc.updateGrokUsageFromResponse(ctx, account, nil, http.StatusOK)
  3127. require.Zero(t, repo.recoveryClearCalls)
  3128. })
  3129. t.Run("API key success does not alter OAuth cooldown state", func(t *testing.T) {
  3130. account := &Account{
  3131. ID: 662,
  3132. Platform: PlatformGrok,
  3133. Type: AccountTypeAPIKey,
  3134. RateLimitedAt: &observedLimitedAt,
  3135. RateLimitResetAt: &observedResetAt,
  3136. }
  3137. repo := &grokQuotaAccountRepo{recoveryClearResult: true}
  3138. svc := &OpenAIGatewayService{accountRepo: repo}
  3139. svc.updateGrokUsageFromResponse(context.Background(), account, nil, http.StatusOK)
  3140. require.Zero(t, repo.recoveryClearCalls)
  3141. })
  3142. }
  3143. func TestUpdateGrokUsageSnapshotExhaustedSuccessWithoutResetUsesFallback(t *testing.T) {
  3144. repo := &grokQuotaAccountRepo{}
  3145. svc := &OpenAIGatewayService{accountRepo: repo}
  3146. account := &Account{ID: 67, Platform: PlatformGrok, Type: AccountTypeOAuth}
  3147. before := time.Now()
  3148. svc.updateGrokUsageSnapshot(context.Background(), account, &xai.QuotaSnapshot{
  3149. StatusCode: http.StatusOK,
  3150. Tokens: &xai.QuotaWindow{
  3151. Limit: grokInt64PtrForTest(2_000_000),
  3152. Remaining: grokInt64PtrForTest(0),
  3153. },
  3154. UpdatedAt: before.UTC().Format(time.RFC3339),
  3155. })
  3156. require.Equal(t, 1, repo.rateLimitedCalls)
  3157. require.WithinDuration(t, before.Add(grokRateLimitFallbackCooldown), repo.lastRateLimitResetAt, time.Second)
  3158. stored, ok := repo.updates[account.ID][grokQuotaSnapshotExtraKey].(*xai.QuotaSnapshot)
  3159. require.True(t, ok)
  3160. require.NotNil(t, stored.Tokens.ResetUnix)
  3161. paused, _ := shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, before.Add(time.Second))
  3162. require.True(t, paused)
  3163. paused, _ = shouldAutoPauseGrokQuotaWindow("tokens", stored.Tokens, repo.lastRateLimitResetAt.Add(time.Second))
  3164. require.False(t, paused)
  3165. }
  3166. func TestOpenAIWSHTTPBridgeGrok429PersistsRateLimit(t *testing.T) {
  3167. repo := &grokQuotaAccountRepo{}
  3168. upstream := &httpUpstreamRecorder{resp: &http.Response{
  3169. StatusCode: http.StatusTooManyRequests,
  3170. Header: http.Header{"Retry-After": []string{"45"}},
  3171. Body: io.NopCloser(strings.NewReader(`{"error":{"message":"rate limited"}}`)),
  3172. }}
  3173. svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream}
  3174. account := &Account{ID: 68, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1}
  3175. before := time.Now()
  3176. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  3177. context.Background(), nil, account, "token",
  3178. []byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`),
  3179. 64, "grok-4.3", "", "", "", "cache-id", 1,
  3180. func([]byte) error { return nil },
  3181. )
  3182. require.Error(t, err)
  3183. require.Nil(t, result)
  3184. require.Equal(t, 1, repo.rateLimitedCalls)
  3185. require.WithinDuration(t, before.Add(45*time.Second), repo.lastRateLimitResetAt, time.Second)
  3186. require.Zero(t, repo.tempUnschedCalls)
  3187. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  3188. }
  3189. func TestOpenAIWSHTTPBridgeSSEErrorSideEffectsRunOncePerPlatform(t *testing.T) {
  3190. gin.SetMode(gin.TestMode)
  3191. for _, platform := range []string{PlatformOpenAI, PlatformGrok} {
  3192. t.Run(platform, func(t *testing.T) {
  3193. repo := &grokQuotaAccountRepo{}
  3194. cfg := &config.Config{}
  3195. upstream := &httpUpstreamRecorder{resp: &http.Response{
  3196. StatusCode: http.StatusOK,
  3197. Header: make(http.Header),
  3198. Body: io.NopCloser(strings.NewReader(
  3199. "data: {\"type\":\"error\",\"error\":{\"type\":\"rate_limit_error\",\"code\":\"rate_limit_exceeded\",\"message\":\"limited\"}}\n\n",
  3200. )),
  3201. }}
  3202. svc := &OpenAIGatewayService{
  3203. cfg: cfg,
  3204. accountRepo: repo,
  3205. httpUpstream: upstream,
  3206. }
  3207. if platform == PlatformOpenAI {
  3208. svc.rateLimitService = NewRateLimitService(repo, nil, cfg, nil, nil)
  3209. }
  3210. account := &Account{ID: 70, Platform: platform, Type: AccountTypeOAuth, Concurrency: 1}
  3211. recorder := httptest.NewRecorder()
  3212. c, _ := gin.CreateTestContext(recorder)
  3213. c.Request = httptest.NewRequest(http.MethodGet, "/v1/responses", nil)
  3214. payload := []byte(`{"type":"response.create","model":"gpt-5","input":"hi"}`)
  3215. writes := 0
  3216. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  3217. context.Background(), c, account, "sk-test", payload, len(payload),
  3218. "gpt-5", "", "", "", "", 1,
  3219. func([]byte) error {
  3220. writes++
  3221. return nil
  3222. },
  3223. )
  3224. require.Nil(t, result)
  3225. var failoverErr *UpstreamFailoverError
  3226. require.ErrorAs(t, err, &failoverErr)
  3227. require.Equal(t, http.StatusTooManyRequests, failoverErr.StatusCode)
  3228. require.Zero(t, writes)
  3229. require.Equal(t, 1, repo.rateLimitedCalls)
  3230. })
  3231. }
  3232. }
  3233. func TestOpenAIWSHTTPBridgeGrokExhaustedSuccessPersistsRateLimit(t *testing.T) {
  3234. repo := &grokQuotaAccountRepo{}
  3235. resetAt := time.Now().Add(20 * time.Minute).UTC().Truncate(time.Second)
  3236. resp := grokMessagesSSECompletedResponse("resp_ws_limited", 0)
  3237. resp.Header.Set("X-Ratelimit-Limit-Requests", "10")
  3238. resp.Header.Set("X-Ratelimit-Remaining-Requests", "0")
  3239. resp.Header.Set("X-Ratelimit-Reset-Requests", fmt.Sprintf("%d", resetAt.Unix()))
  3240. upstream := &httpUpstreamRecorder{resp: resp}
  3241. svc := &OpenAIGatewayService{accountRepo: repo, httpUpstream: upstream}
  3242. account := &Account{ID: 69, Platform: PlatformGrok, Type: AccountTypeOAuth, Concurrency: 1}
  3243. result, err := svc.proxyOpenAIWSHTTPBridgeTurn(
  3244. context.Background(), nil, account, "token",
  3245. []byte(`{"type":"response.create","model":"grok-4.3","input":"hi"}`),
  3246. 64, "grok-4.3", "", "", "", "cache-id", 1,
  3247. func([]byte) error { return nil },
  3248. )
  3249. require.NoError(t, err)
  3250. require.NotNil(t, result)
  3251. require.Equal(t, 1, repo.rateLimitedCalls)
  3252. require.WithinDuration(t, resetAt, repo.lastRateLimitResetAt, time.Second)
  3253. require.True(t, svc.isOpenAIAccountRuntimeBlocked(account))
  3254. }
  3255. func TestFailoverOpenAIUpstreamHTTPErrorUsesOnlyGrokRateLimitPolicy(t *testing.T) {
  3256. gin.SetMode(gin.TestMode)
  3257. repo := &grokQuotaAccountRepo{}
  3258. svc := &OpenAIGatewayService{accountRepo: repo}
  3259. account := &Account{ID: 70, Platform: PlatformGrok, Type: AccountTypeOAuth}
  3260. resp := &http.Response{
  3261. StatusCode: http.StatusTooManyRequests,
  3262. Header: http.Header{"Retry-After": []string{"45"}},
  3263. }
  3264. recorder := httptest.NewRecorder()
  3265. c, _ := gin.CreateTestContext(recorder)
  3266. c.Request = httptest.NewRequest(http.MethodPost, "/v1/messages", nil)
  3267. failoverErr := svc.failoverOpenAIUpstreamHTTPError(
  3268. context.Background(), c, account, resp,
  3269. []byte(`{"error":{"message":"rate limited"}}`), "rate limited", "grok-4.3",
  3270. )
  3271. require.NotNil(t, failoverErr)
  3272. require.Equal(t, 1, repo.rateLimitedCalls)
  3273. require.Zero(t, repo.tempUnschedCalls)
  3274. }
  3275. func TestPatchGrokResponsesBody_StripsReasoningContentNull(t *testing.T) {
  3276. t.Parallel()
  3277. body := []byte(`{
  3278. "model": "grok-latest",
  3279. "input": [
  3280. {"type":"message","role":"user","content":[{"type":"input_text","text":"hi"}]},
  3281. {"type":"reasoning","summary":[{"type":"summary_text","text":"thinking..."}],"content":null,"encrypted_content":null},
  3282. {"type":"message","role":"assistant","content":[{"type":"output_text","text":"Hello!"}]}
  3283. ]
  3284. }`)
  3285. patched, err := patchGrokResponsesBody(body, "grok-4.5")
  3286. require.NoError(t, err)
  3287. require.True(t, json.Valid(patched))
  3288. input := gjson.GetBytes(patched, "input")
  3289. require.True(t, input.IsArray())
  3290. items := input.Array()
  3291. require.Len(t, items, 3)
  3292. reasoning := items[1]
  3293. require.Equal(t, "reasoning", reasoning.Get("type").String())
  3294. require.True(t, reasoning.Get("summary").Exists(), "summary should be preserved")
  3295. require.False(t, reasoning.Get("content").Exists(), "content: null should be stripped")
  3296. }
  3297. func TestPatchGrokResponsesBody_StripsNestedInputNulls(t *testing.T) {
  3298. t.Parallel()
  3299. body := []byte(`{
  3300. "model": "grok-4.6",
  3301. "input": [
  3302. {"type":"message","role":"user","content":[{"type":"input_text","text":"hi","logprobs":null}]}
  3303. ]
  3304. }`)
  3305. patched, err := patchGrokResponsesBody(body, "grok-4.6")
  3306. require.NoError(t, err)
  3307. require.False(t, gjson.GetBytes(patched, "input.0.content.0.logprobs").Exists())
  3308. require.Equal(t, "hi", gjson.GetBytes(patched, "input.0.content.0.text").String())
  3309. }
  3310. func TestStripExplicitNullsFromGrokInputSkipsCompaction(t *testing.T) {
  3311. t.Parallel()
  3312. input := []any{
  3313. map[string]any{"type": "message", "role": "user", "content": []any{map[string]any{"type": "input_text", "text": "hi", "logprobs": nil}}},
  3314. map[string]any{"type": "compaction", "id": "cmp_1", "encrypted_content": "blob", "status": nil},
  3315. }
  3316. out, changed := stripExplicitNullsFromGrokInput(input)
  3317. require.True(t, changed)
  3318. items, ok := out.([]any)
  3319. require.True(t, ok)
  3320. require.Len(t, items, 2)
  3321. msg := items[0].(map[string]any)
  3322. content := msg["content"].([]any)[0].(map[string]any)
  3323. _, hasLogprobs := content["logprobs"]
  3324. require.False(t, hasLogprobs)
  3325. compaction := items[1].(map[string]any)
  3326. require.Equal(t, "blob", compaction["encrypted_content"])
  3327. _, hasStatus := compaction["status"]
  3328. require.True(t, hasStatus, "compaction items must stay unmodified")
  3329. }
  3330. func TestPatchGrokResponsesBody_KeepsReasoningContentNonNull(t *testing.T) {
  3331. t.Parallel()
  3332. body := []byte(`{
  3333. "model": "grok-latest",
  3334. "input": [
  3335. {"type":"reasoning","summary":[{"type":"summary_text","text":"ok"}],"content":"real content"}
  3336. ]
  3337. }`)
  3338. patched, err := patchGrokResponsesBody(body, "grok-4.5")
  3339. require.NoError(t, err)
  3340. reasoning := gjson.GetBytes(patched, "input.0")
  3341. require.Equal(t, "real content", reasoning.Get("content").String(), "non-null content must not be stripped")
  3342. }
  3343. func TestPatchGrokResponsesBody_MultipleReasoningContentNull(t *testing.T) {
  3344. t.Parallel()
  3345. body := []byte(`{
  3346. "model": "grok-latest",
  3347. "input": [
  3348. {"type":"reasoning","summary":[{"type":"summary_text","text":"r1"}],"content":null},
  3349. {"type":"message","role":"user","content":"hi"},
  3350. {"type":"reasoning","summary":[{"type":"summary_text","text":"r2"}],"content":null}
  3351. ]
  3352. }`)
  3353. patched, err := patchGrokResponsesBody(body, "grok-4.5")
  3354. require.NoError(t, err)
  3355. items := gjson.GetBytes(patched, "input").Array()
  3356. require.Len(t, items, 3)
  3357. require.False(t, items[0].Get("content").Exists())
  3358. require.False(t, items[2].Get("content").Exists())
  3359. }
  3360. func TestIsGrokImageGenerationModel(t *testing.T) {
  3361. t.Parallel()
  3362. tests := []struct {
  3363. model string
  3364. want bool
  3365. }{
  3366. {"grok-imagine", true},
  3367. {"grok-imagine-image-quality", true},
  3368. {"grok-imagine-edit", true},
  3369. {"grok-imagine-image-hd", true},
  3370. {"grok-4.5", false},
  3371. {"grok-composer", false},
  3372. }
  3373. for _, tt := range tests {
  3374. t.Run(tt.model, func(t *testing.T) {
  3375. require.Equal(t, tt.want, isGrokImageGenerationModel(tt.model))
  3376. })
  3377. }
  3378. }
  3379. func TestBuildGrokSchedulerExtraUpdates_FeedsThresholdEvaluator(t *testing.T) {
  3380. int64p := func(v int64) *int64 { return &v }
  3381. resetUnix := time.Now().Add(90 * time.Minute).Unix()
  3382. snapshot := &xai.QuotaSnapshot{
  3383. Requests: &xai.QuotaWindow{Limit: int64p(100), Remaining: int64p(30)}, // 70% used
  3384. Tokens: &xai.QuotaWindow{Limit: int64p(1000), Remaining: int64p(50), ResetUnix: &resetUnix}, // 95% used (most constrained)
  3385. }
  3386. updates := buildGrokSchedulerExtraUpdates(snapshot)
  3387. require.NotNil(t, updates)
  3388. require.InDelta(t, 95.0, updates["grok_sched_utilization"], 0.001, "picks the most-constrained window")
  3389. require.Contains(t, updates, "grok_sched_reset_at")
  3390. // The written extras must actually drive EvaluateAccountSchedulingThreshold
  3391. // (proves the previously-dead read side is now fed).
  3392. account := &Account{Platform: PlatformGrok, Extra: updates}
  3393. decision := EvaluateAccountSchedulingThreshold(account, map[string]int{PlatformGrok: 90}, time.Now())
  3394. require.True(t, decision.ShouldPause)
  3395. require.InDelta(t, 95.0, decision.UsedPercent, 0.001)
  3396. require.NotNil(t, decision.Until)
  3397. }
  3398. func TestBuildGrokSchedulerExtraUpdates_NilWhenNoQuotaWindows(t *testing.T) {
  3399. require.Nil(t, buildGrokSchedulerExtraUpdates(&xai.QuotaSnapshot{}))
  3400. require.Nil(t, buildGrokSchedulerExtraUpdates(nil))
  3401. }