grok_media.go 49 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499
  1. package service
  2. import (
  3. "bytes"
  4. "context"
  5. "encoding/json"
  6. "fmt"
  7. "io"
  8. "mime"
  9. "mime/multipart"
  10. "net/http"
  11. "net/url"
  12. "strconv"
  13. "strings"
  14. "time"
  15. "github.com/Wei-Shaw/sub2api/internal/config"
  16. "github.com/Wei-Shaw/sub2api/internal/util/responseheaders"
  17. "github.com/gin-gonic/gin"
  18. "github.com/tidwall/gjson"
  19. "github.com/tidwall/sjson"
  20. )
  21. type GrokMediaEndpoint string
  22. const (
  23. GrokMediaEndpointImagesGenerations GrokMediaEndpoint = "images_generations"
  24. GrokMediaEndpointImagesEdits GrokMediaEndpoint = "images_edits"
  25. GrokMediaEndpointVideosGenerations GrokMediaEndpoint = "videos_generations"
  26. GrokMediaEndpointVideosEdits GrokMediaEndpoint = "videos_edits"
  27. GrokMediaEndpointVideosExtensions GrokMediaEndpoint = "videos_extensions"
  28. GrokMediaEndpointVideoStatus GrokMediaEndpoint = "video_status"
  29. GrokMediaEndpointVideoContent GrokMediaEndpoint = "video_content"
  30. // Official xAI Imagine image-edit limit.
  31. grokMediaMaxEditSourceImages = 3
  32. )
  33. func (e GrokMediaEndpoint) RequiresRequestBody() bool {
  34. return !e.IsVideoLookupRequest()
  35. }
  36. func (e GrokMediaEndpoint) IsVideoLookupRequest() bool {
  37. return e == GrokMediaEndpointVideoStatus || e == GrokMediaEndpointVideoContent
  38. }
  39. func (e GrokMediaEndpoint) IsGenerationRequest() bool {
  40. switch e {
  41. case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits, GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions:
  42. return true
  43. default:
  44. return false
  45. }
  46. }
  47. type GrokMediaRequestInfo struct {
  48. Model string
  49. Prompt string
  50. N int
  51. Size string
  52. SizeTier string
  53. AspectRatio string
  54. ImageResolution string
  55. Resolution string
  56. DurationSeconds int
  57. InputImageURLs []string
  58. MaskImageURL string
  59. Uploads []OpenAIImagesUpload
  60. MaskUpload *OpenAIImagesUpload
  61. }
  62. func (r GrokMediaRequestInfo) ModerationBody() []byte {
  63. payload := map[string]any{}
  64. if prompt := strings.TrimSpace(r.Prompt); prompt != "" {
  65. payload["prompt"] = prompt
  66. }
  67. images := make([]map[string]string, 0, len(r.InputImageURLs)+len(r.Uploads)+1)
  68. for _, imageURL := range r.InputImageURLs {
  69. if imageURL = strings.TrimSpace(imageURL); imageURL != "" {
  70. images = append(images, map[string]string{"image_url": imageURL})
  71. }
  72. }
  73. for _, upload := range r.Uploads {
  74. if dataURL := upload.ModerationDataURL(); dataURL != "" {
  75. images = append(images, map[string]string{"image_url": dataURL})
  76. }
  77. }
  78. if maskURL := strings.TrimSpace(r.MaskImageURL); maskURL != "" {
  79. images = append(images, map[string]string{"image_url": maskURL})
  80. }
  81. if r.MaskUpload != nil {
  82. if dataURL := r.MaskUpload.ModerationDataURL(); dataURL != "" {
  83. images = append(images, map[string]string{"image_url": dataURL})
  84. }
  85. }
  86. if len(images) > 0 {
  87. payload["images"] = images
  88. }
  89. if len(payload) == 0 {
  90. return nil
  91. }
  92. body, err := json.Marshal(payload)
  93. if err != nil {
  94. return nil
  95. }
  96. return body
  97. }
  98. func (e GrokMediaEndpoint) httpMethod() string {
  99. if e.IsVideoLookupRequest() {
  100. return http.MethodGet
  101. }
  102. return http.MethodPost
  103. }
  104. func ExtractGrokMediaModel(contentType string, body []byte) string {
  105. return ParseGrokMediaRequest(contentType, body).Model
  106. }
  107. func ParseGrokMediaRequest(contentType string, body []byte) GrokMediaRequestInfo {
  108. info := GrokMediaRequestInfo{N: 1}
  109. if gjson.ValidBytes(body) {
  110. parseGrokMediaJSONRequest(body, &info)
  111. } else {
  112. parseGrokMediaMultipartRequest(contentType, body, &info)
  113. }
  114. info.Model = strings.TrimSpace(info.Model)
  115. info.Prompt = strings.TrimSpace(info.Prompt)
  116. info.Size = strings.TrimSpace(info.Size)
  117. info.SizeTier = NormalizeImageBillingTierOrDefault(info.Size)
  118. info.AspectRatio = strings.TrimSpace(info.AspectRatio)
  119. info.ImageResolution = grokImagineImageResolution(info.ImageResolution)
  120. info.Resolution = NormalizeVideoBillingResolutionOrDefault(info.Resolution)
  121. info.DurationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(info.DurationSeconds)
  122. if info.N <= 0 {
  123. info.N = 1
  124. }
  125. return info
  126. }
  127. func parseGrokMediaJSONRequest(body []byte, info *GrokMediaRequestInfo) {
  128. if info == nil {
  129. return
  130. }
  131. info.Model = strings.TrimSpace(gjson.GetBytes(body, "model").String())
  132. info.Prompt = strings.TrimSpace(gjson.GetBytes(body, "prompt").String())
  133. info.Size = strings.TrimSpace(gjson.GetBytes(body, "size").String())
  134. info.AspectRatio = strings.TrimSpace(gjson.GetBytes(body, "aspect_ratio").String())
  135. assignGrokMediaResolution(strings.TrimSpace(gjson.GetBytes(body, "resolution").String()), info)
  136. if duration := gjson.GetBytes(body, "duration"); duration.Exists() && duration.Type == gjson.Number {
  137. info.DurationSeconds = int(duration.Int())
  138. }
  139. if n := gjson.GetBytes(body, "n"); n.Exists() && n.Type == gjson.Number {
  140. info.N = int(n.Int())
  141. }
  142. appendJSONImageURLs := func(value gjson.Result) {
  143. if !value.Exists() {
  144. return
  145. }
  146. switch {
  147. case value.IsArray():
  148. for _, item := range value.Array() {
  149. if imageURL := extractGrokMediaImageURL(item); imageURL != "" {
  150. info.InputImageURLs = append(info.InputImageURLs, imageURL)
  151. }
  152. }
  153. default:
  154. if imageURL := extractGrokMediaImageURL(value); imageURL != "" {
  155. info.InputImageURLs = append(info.InputImageURLs, imageURL)
  156. }
  157. }
  158. }
  159. appendJSONImageURLs(gjson.GetBytes(body, "image"))
  160. appendJSONImageURLs(gjson.GetBytes(body, "images"))
  161. appendJSONImageURLs(gjson.GetBytes(body, "reference_images"))
  162. info.MaskImageURL = extractGrokMediaImageURL(gjson.GetBytes(body, "mask"))
  163. }
  164. func extractGrokMediaImageURL(value gjson.Result) string {
  165. if !value.Exists() {
  166. return ""
  167. }
  168. if value.Type == gjson.String {
  169. return strings.TrimSpace(value.String())
  170. }
  171. if imageURL := strings.TrimSpace(value.Get("url").String()); imageURL != "" {
  172. return imageURL
  173. }
  174. if nested := value.Get("image_url"); nested.Exists() {
  175. if nested.Type == gjson.String {
  176. return strings.TrimSpace(nested.String())
  177. }
  178. if imageURL := strings.TrimSpace(nested.Get("url").String()); imageURL != "" {
  179. return imageURL
  180. }
  181. }
  182. return strings.TrimSpace(value.Get("image_url").String())
  183. }
  184. func grokMediaImageObject(imageURL string) map[string]string {
  185. return map[string]string{"url": imageURL, "type": "image_url"}
  186. }
  187. func parseGrokMediaMultipartRequest(contentType string, body []byte, info *GrokMediaRequestInfo) {
  188. if info == nil {
  189. return
  190. }
  191. mediaType, params, err := mime.ParseMediaType(strings.TrimSpace(contentType))
  192. if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
  193. return
  194. }
  195. boundary := strings.TrimSpace(params["boundary"])
  196. if boundary == "" {
  197. return
  198. }
  199. reader := multipart.NewReader(bytes.NewReader(body), boundary)
  200. for {
  201. part, err := reader.NextPart()
  202. if err == io.EOF {
  203. return
  204. }
  205. if err != nil {
  206. return
  207. }
  208. name := strings.TrimSpace(part.FormName())
  209. if name == "" {
  210. _ = part.Close()
  211. continue
  212. }
  213. data, err := io.ReadAll(io.LimitReader(part, openAIImageMaxUploadPartSize))
  214. _ = part.Close()
  215. if err != nil {
  216. return
  217. }
  218. fileName := strings.TrimSpace(part.FileName())
  219. partContentType := strings.TrimSpace(part.Header.Get("Content-Type"))
  220. if fileName != "" {
  221. upload := OpenAIImagesUpload{
  222. FieldName: name,
  223. FileName: fileName,
  224. ContentType: partContentType,
  225. Data: data,
  226. }
  227. if name == "mask" {
  228. info.MaskUpload = &upload
  229. continue
  230. }
  231. if name == "image" || strings.HasPrefix(name, "image[") {
  232. info.Uploads = append(info.Uploads, upload)
  233. }
  234. continue
  235. }
  236. value := strings.TrimSpace(string(data))
  237. switch name {
  238. case "model":
  239. info.Model = value
  240. case "prompt":
  241. info.Prompt = value
  242. case "size":
  243. info.Size = value
  244. case "aspect_ratio":
  245. info.AspectRatio = value
  246. case "resolution":
  247. assignGrokMediaResolution(value, info)
  248. case "duration":
  249. if duration, err := strconv.Atoi(value); err == nil {
  250. info.DurationSeconds = duration
  251. }
  252. case "n":
  253. if n, err := strconv.Atoi(value); err == nil {
  254. info.N = n
  255. }
  256. case "image", "image_url":
  257. if value != "" {
  258. info.InputImageURLs = append(info.InputImageURLs, value)
  259. }
  260. case "mask", "mask_image_url":
  261. info.MaskImageURL = value
  262. }
  263. }
  264. }
  265. func GrokMediaVideoRequestSessionHash(requestID string, userID, apiKeyID int64) string {
  266. requestID = strings.TrimSpace(requestID)
  267. if requestID == "" || userID <= 0 || apiKeyID <= 0 {
  268. return ""
  269. }
  270. ownerSeed := fmt.Sprintf("%d:%d:%s", userID, apiKeyID, requestID)
  271. return "grok-video:" + DeriveSessionHashFromSeed(ownerSeed)
  272. }
  273. func (s *OpenAIGatewayService) BindGrokMediaVideoRequestAccount(
  274. ctx context.Context,
  275. groupID *int64,
  276. requestID string,
  277. userID, apiKeyID, accountID int64,
  278. ) error {
  279. if s == nil || s.cache == nil {
  280. return fmt.Errorf("grok video request binding cache is unavailable")
  281. }
  282. sessionHash := GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID)
  283. cacheKey := s.openAISessionCacheKey(sessionHash)
  284. if cacheKey == "" || accountID <= 0 {
  285. return fmt.Errorf("grok video request binding is invalid")
  286. }
  287. // Video jobs may complete well after WS sticky TTL (default 1h). Bind at least
  288. // as long as the pending-billing snapshot so late status/content polls resolve.
  289. ttl := grokVideoPendingBillingTTL(s.cfg)
  290. if s.cfg != nil && s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds > 0 {
  291. if sticky := time.Duration(s.cfg.Gateway.OpenAIWS.StickySessionTTLSeconds) * time.Second; sticky > ttl {
  292. ttl = sticky
  293. }
  294. }
  295. return s.cache.SetSessionAccountID(ctx, derefGroupID(groupID), cacheKey, accountID, ttl)
  296. }
  297. func (s *OpenAIGatewayService) ResolveGrokMediaVideoRequestAccount(
  298. ctx context.Context,
  299. groupID *int64,
  300. requestID string,
  301. userID, apiKeyID int64,
  302. ) (int64, error) {
  303. if s == nil || s.cache == nil {
  304. return 0, fmt.Errorf("grok video request binding cache is unavailable")
  305. }
  306. cacheKey := s.openAISessionCacheKey(GrokMediaVideoRequestSessionHash(requestID, userID, apiKeyID))
  307. if cacheKey == "" {
  308. return 0, fmt.Errorf("grok video request binding is invalid")
  309. }
  310. return s.cache.GetSessionAccountID(ctx, derefGroupID(groupID), cacheKey)
  311. }
  312. // GrokVideoPendingBilling is the create-time snapshot used when status polling
  313. // first observes a completed video URL. Status may omit model/duration; we fall
  314. // back to this snapshot, then defaults.
  315. type GrokVideoPendingBilling struct {
  316. Model string `json:"model"`
  317. BillingModel string `json:"billing_model,omitempty"`
  318. UpstreamModel string `json:"upstream_model,omitempty"`
  319. VideoResolution string `json:"video_resolution,omitempty"`
  320. VideoDurationSeconds int `json:"video_duration_seconds,omitempty"`
  321. OriginalModel string `json:"original_model,omitempty"`
  322. // CreatedAt is when the gateway accepted the async create (RFC3339Nano UTC).
  323. // duration_ms for deferred billing is measured from this instant until the
  324. // first official done+video.url observation (status poll or content download),
  325. // not the latency of that single discovery request alone.
  326. CreatedAt string `json:"created_at,omitempty"`
  327. }
  328. // GrokVideoPendingCreatedAtNow formats a create-accept timestamp for pending billing.
  329. func GrokVideoPendingCreatedAtNow() string {
  330. return time.Now().UTC().Format(time.RFC3339Nano)
  331. }
  332. // GrokVideoE2EDuration returns wall time from create accept to discovery of completion.
  333. // Returns 0 when CreatedAt is missing or unparseable (caller keeps poll-only Duration).
  334. func GrokVideoE2EDuration(createdAt string, discoveredAt time.Time) time.Duration {
  335. createdAt = strings.TrimSpace(createdAt)
  336. if createdAt == "" {
  337. return 0
  338. }
  339. if discoveredAt.IsZero() {
  340. discoveredAt = time.Now()
  341. }
  342. var created time.Time
  343. var err error
  344. if created, err = time.Parse(time.RFC3339Nano, createdAt); err != nil {
  345. if created, err = time.Parse(time.RFC3339, createdAt); err != nil {
  346. return 0
  347. }
  348. }
  349. if created.IsZero() {
  350. return 0
  351. }
  352. d := discoveredAt.Sub(created)
  353. if d < 0 {
  354. return 0
  355. }
  356. return d
  357. }
  358. func grokVideoPendingBillingKey(requestID string, userID, apiKeyID int64) string {
  359. requestID = strings.TrimSpace(requestID)
  360. if requestID == "" || userID <= 0 || apiKeyID <= 0 {
  361. return ""
  362. }
  363. return fmt.Sprintf("%d:%d:%s", userID, apiKeyID, requestID)
  364. }
  365. func grokVideoPendingBillingTTL(cfg *config.Config) time.Duration {
  366. // Video generation can take several minutes; keep create-time pricing for a day.
  367. _ = cfg
  368. return 24 * time.Hour
  369. }
  370. func grokVideoBilledClaimTTL(cfg *config.Config) time.Duration {
  371. _ = cfg
  372. return 48 * time.Hour
  373. }
  374. // StoreGrokVideoPendingBilling persists create-time billing params for deferred status billing.
  375. func (s *OpenAIGatewayService) StoreGrokVideoPendingBilling(
  376. ctx context.Context,
  377. requestID string,
  378. userID, apiKeyID int64,
  379. pending GrokVideoPendingBilling,
  380. ) error {
  381. if s == nil || s.cache == nil {
  382. return fmt.Errorf("grok video pending billing cache is unavailable")
  383. }
  384. key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
  385. if key == "" {
  386. return fmt.Errorf("grok video pending billing key is invalid")
  387. }
  388. pending.Model = strings.TrimSpace(pending.Model)
  389. pending.BillingModel = strings.TrimSpace(pending.BillingModel)
  390. pending.UpstreamModel = strings.TrimSpace(pending.UpstreamModel)
  391. pending.OriginalModel = strings.TrimSpace(pending.OriginalModel)
  392. if pending.VideoResolution != "" {
  393. pending.VideoResolution = NormalizeVideoBillingResolutionOrDefault(pending.VideoResolution)
  394. }
  395. if pending.VideoDurationSeconds > 0 {
  396. pending.VideoDurationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(pending.VideoDurationSeconds)
  397. }
  398. // Always stamp create-accept time when missing so deferred duration_ms is E2E.
  399. if strings.TrimSpace(pending.CreatedAt) == "" {
  400. pending.CreatedAt = GrokVideoPendingCreatedAtNow()
  401. } else {
  402. pending.CreatedAt = strings.TrimSpace(pending.CreatedAt)
  403. }
  404. payload, err := json.Marshal(pending)
  405. if err != nil {
  406. return err
  407. }
  408. return s.cache.SetGrokVideoPendingBilling(ctx, key, payload, grokVideoPendingBillingTTL(s.cfg))
  409. }
  410. // LoadGrokVideoPendingBilling returns the create-time snapshot (may be nil on miss).
  411. func (s *OpenAIGatewayService) LoadGrokVideoPendingBilling(
  412. ctx context.Context,
  413. requestID string,
  414. userID, apiKeyID int64,
  415. ) (*GrokVideoPendingBilling, error) {
  416. if s == nil || s.cache == nil {
  417. return nil, fmt.Errorf("grok video pending billing cache is unavailable")
  418. }
  419. key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
  420. if key == "" {
  421. return nil, fmt.Errorf("grok video pending billing key is invalid")
  422. }
  423. payload, err := s.cache.GetGrokVideoPendingBilling(ctx, key)
  424. if err != nil || len(payload) == 0 {
  425. return nil, err
  426. }
  427. var pending GrokVideoPendingBilling
  428. if err := json.Unmarshal(payload, &pending); err != nil {
  429. return nil, err
  430. }
  431. return &pending, nil
  432. }
  433. // ClaimGrokVideoBilling returns true once for a completed video request so status
  434. // polls do not double-bill. Fail-closed: claim errors are treated as already billed.
  435. func (s *OpenAIGatewayService) ClaimGrokVideoBilling(
  436. ctx context.Context,
  437. requestID string,
  438. userID, apiKeyID int64,
  439. ) (bool, error) {
  440. if s == nil || s.cache == nil {
  441. return false, fmt.Errorf("grok video billing claim cache is unavailable")
  442. }
  443. key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
  444. if key == "" {
  445. return false, fmt.Errorf("grok video billing claim key is invalid")
  446. }
  447. return s.cache.ClaimGrokVideoBilled(ctx, key, grokVideoBilledClaimTTL(s.cfg))
  448. }
  449. // ReleaseGrokVideoBilling clears a claim after a failed durable RecordUsage so a
  450. // later status/content poll can retry billing.
  451. func (s *OpenAIGatewayService) ReleaseGrokVideoBilling(
  452. ctx context.Context,
  453. requestID string,
  454. userID, apiKeyID int64,
  455. ) error {
  456. if s == nil || s.cache == nil {
  457. return fmt.Errorf("grok video billing claim cache is unavailable")
  458. }
  459. key := grokVideoPendingBillingKey(requestID, userID, apiKeyID)
  460. if key == "" {
  461. return fmt.Errorf("grok video billing claim key is invalid")
  462. }
  463. return s.cache.ReleaseGrokVideoBilled(ctx, key)
  464. }
  465. // StableGrokVideoBillingRequestID is the durable usage_logs / dedup key for one
  466. // async video task (not the per-poll gateway request id).
  467. func StableGrokVideoBillingRequestID(taskRequestID string) string {
  468. taskRequestID = strings.TrimSpace(taskRequestID)
  469. if taskRequestID == "" {
  470. return ""
  471. }
  472. if strings.HasPrefix(taskRequestID, "grok-video:") {
  473. return taskRequestID
  474. }
  475. return "grok-video:" + taskRequestID
  476. }
  477. // Official xAI async video status success shape (docs.x.ai Video Generation):
  478. //
  479. // {"status":"done","model":"grok-imagine-video-1.5","video":{"url":"...","duration":8,"respect_moderation":true}}
  480. //
  481. // Request may include resolution ("480p"|"720p"|"1080p"); completed status does not
  482. // document a resolution field — bill resolution from the create-time request snapshot.
  483. // IsGrokVideoStatusBillable matches official success: status == "done" AND non-empty video.url.
  484. // pending / expired / failed, or done without a video URL, are not billable.
  485. func IsGrokVideoStatusBillable(statusBody []byte) bool {
  486. if len(statusBody) == 0 || !gjson.ValidBytes(statusBody) {
  487. return false
  488. }
  489. if !isOfficialGrokVideoStatusDone(statusBody) {
  490. return false
  491. }
  492. return strings.TrimSpace(gjson.GetBytes(statusBody, "video.url").String()) != ""
  493. }
  494. func isOfficialGrokVideoStatusDone(statusBody []byte) bool {
  495. // Official enum: pending | done | expired | failed.
  496. return strings.EqualFold(strings.TrimSpace(gjson.GetBytes(statusBody, "status").String()), "done")
  497. }
  498. // ExtractGrokVideoBillingFromStatusBody builds usage units from an official done status.
  499. // Field priority (official docs):
  500. // - duration: video.duration (seconds)
  501. // - model: top-level model
  502. // - resolution: not in status response → create-time pending snapshot → default 480p
  503. func ExtractGrokVideoBillingFromStatusBody(statusBody []byte, pending *GrokVideoPendingBilling, requestID string) *OpenAIForwardResult {
  504. if !IsGrokVideoStatusBillable(statusBody) {
  505. return nil
  506. }
  507. model := ""
  508. billingModel := ""
  509. upstreamModel := ""
  510. resolution := ""
  511. durationSeconds := 0
  512. if gjson.ValidBytes(statusBody) {
  513. // Official: top-level model.
  514. model = strings.TrimSpace(gjson.GetBytes(statusBody, "model").String())
  515. // Official: video.duration (number of seconds).
  516. if v := gjson.GetBytes(statusBody, "video.duration"); v.Exists() && v.Type == gjson.Number {
  517. durationSeconds = int(v.Int())
  518. if durationSeconds == 0 && v.Float() > 0 {
  519. // Sub-second values are unexpected for this API; still accept truncated int path above.
  520. durationSeconds = int(v.Float())
  521. }
  522. }
  523. }
  524. if pending != nil {
  525. if model == "" {
  526. model = firstNonEmpty(pending.BillingModel, pending.Model, pending.OriginalModel)
  527. }
  528. if billingModel == "" {
  529. billingModel = firstNonEmpty(pending.BillingModel, pending.Model)
  530. }
  531. if upstreamModel == "" {
  532. upstreamModel = pending.UpstreamModel
  533. }
  534. // Official status has no resolution — always take create request when available.
  535. resolution = pending.VideoResolution
  536. if durationSeconds <= 0 {
  537. durationSeconds = pending.VideoDurationSeconds
  538. }
  539. }
  540. if model == "" {
  541. // Official default video model family when status omits model.
  542. model = "grok-imagine-video"
  543. }
  544. if billingModel == "" {
  545. billingModel = model
  546. }
  547. // Resolution is request-only per docs; empty → handler applies official default 480p.
  548. if resolution != "" {
  549. resolution = NormalizeVideoBillingResolutionOrDefault(resolution)
  550. }
  551. if durationSeconds > 0 {
  552. durationSeconds = NormalizeVideoBillingDurationSecondsOrDefault(durationSeconds)
  553. }
  554. responseID := extractGrokMediaVideoRequestID(statusBody)
  555. if responseID == "" {
  556. responseID = strings.TrimSpace(requestID)
  557. }
  558. return &OpenAIForwardResult{
  559. ResponseID: responseID,
  560. Model: model,
  561. BillingModel: billingModel,
  562. UpstreamModel: upstreamModel,
  563. VideoCount: 1,
  564. VideoResolution: resolution,
  565. VideoDurationSeconds: durationSeconds,
  566. }
  567. }
  568. func (s *OpenAIGatewayService) ForwardGrokMedia(
  569. ctx context.Context,
  570. c *gin.Context,
  571. account *Account,
  572. endpoint GrokMediaEndpoint,
  573. requestID string,
  574. body []byte,
  575. contentType string,
  576. ) (*OpenAIForwardResult, error) {
  577. startTime := time.Now()
  578. if account == nil {
  579. return nil, fmt.Errorf("grok account is required")
  580. }
  581. if account.Platform != PlatformGrok {
  582. return nil, fmt.Errorf("account platform %s is not supported for grok media", account.Platform)
  583. }
  584. token, _, err := s.getRequestCredential(ctx, c, account)
  585. if err != nil {
  586. return nil, err
  587. }
  588. if endpoint == GrokMediaEndpointVideoContent {
  589. return s.forwardGrokMediaVideoContent(ctx, c, account, token, requestID, startTime)
  590. }
  591. targetURL, err := buildGrokMediaURL(account, s.cfg, endpoint, requestID)
  592. if err != nil {
  593. return nil, err
  594. }
  595. body, contentType, err = prepareGrokMediaForwardBody(endpoint, body, contentType)
  596. if err != nil {
  597. return nil, err
  598. }
  599. body, contentType, err = normalizeGrokMediaForwardBody(endpoint, body, contentType)
  600. if err != nil {
  601. return nil, err
  602. }
  603. requestInfo := ParseGrokMediaRequest(contentType, body)
  604. upstreamModel := requestInfo.Model
  605. if endpoint.RequiresRequestBody() && gjson.ValidBytes(body) {
  606. if mappedModel := strings.TrimSpace(account.GetMappedModel(requestInfo.Model)); mappedModel != "" {
  607. upstreamModel = mappedModel
  608. }
  609. if upstreamModel != requestInfo.Model {
  610. body, err = sjson.SetBytes(body, "model", upstreamModel)
  611. if err != nil {
  612. return nil, fmt.Errorf("rewrite grok media account mapped model: %w", err)
  613. }
  614. }
  615. }
  616. body, contentType, err = sanitizeGrokMediaForwardBody(endpoint, body, contentType)
  617. if err != nil {
  618. return nil, err
  619. }
  620. var bodyReader io.Reader
  621. if endpoint.RequiresRequestBody() {
  622. bodyReader = bytes.NewReader(body)
  623. }
  624. upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
  625. defer releaseUpstreamCtx()
  626. upstreamReq, err := http.NewRequestWithContext(upstreamCtx, endpoint.httpMethod(), targetURL, bodyReader)
  627. if err != nil {
  628. return nil, err
  629. }
  630. upstreamReq.Header.Set("Authorization", "Bearer "+token)
  631. upstreamReq.Header.Set("Accept", "application/json")
  632. if account.IsGrokOAuth() && isGrokCLIProxyTarget(targetURL) {
  633. applyGrokCLIHeaders(upstreamReq.Header)
  634. }
  635. if endpoint.RequiresRequestBody() {
  636. contentType = strings.TrimSpace(contentType)
  637. if contentType == "" {
  638. contentType = "application/json"
  639. }
  640. upstreamReq.Header.Set("Content-Type", contentType)
  641. }
  642. // 账号级请求头覆写最后应用,配置值优先于内置默认头。
  643. account.ApplyHeaderOverrides(upstreamReq.Header)
  644. proxyURL := ""
  645. if account.ProxyID != nil && account.Proxy != nil {
  646. proxyURL = account.Proxy.URL()
  647. }
  648. upstreamStart := time.Now()
  649. resp, err := s.httpUpstream.Do(WithAccountTrafficRequest(upstreamReq, account), proxyURL, account.ID, account.Concurrency)
  650. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  651. if err != nil {
  652. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
  653. }
  654. defer func() { _ = resp.Body.Close() }()
  655. requestIDHeader := firstNonEmpty(resp.Header.Get("x-request-id"), resp.Header.Get("xai-request-id"))
  656. requestModel := requestInfo.Model
  657. if resp.StatusCode >= 400 {
  658. return s.handleGrokMediaErrorResponse(ctx, resp, c, account, requestIDHeader, requestModel)
  659. }
  660. s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, requestModel), account, resp.Header, resp.StatusCode)
  661. respBody, err := ReadUpstreamResponseBody(resp.Body, s.cfg, c, openAITooLargeError)
  662. if err != nil {
  663. return nil, err
  664. }
  665. if endpoint == GrokMediaEndpointImagesGenerations || endpoint == GrokMediaEndpointImagesEdits {
  666. if countOpenAIResponseImageOutputsFromJSONBytes(respBody) <= 0 {
  667. setOpsUpstreamError(c, http.StatusBadGateway, "xAI upstream returned no image output", truncateString(string(respBody), 512))
  668. return nil, &UpstreamFailoverError{
  669. StatusCode: http.StatusBadGateway,
  670. ResponseBody: respBody,
  671. ResponseHeaders: resp.Header.Clone(),
  672. }
  673. }
  674. }
  675. if endpoint == GrokMediaEndpointVideoStatus {
  676. respBody = rewriteGrokMediaVideoContentURLs(
  677. respBody,
  678. requestID,
  679. grokMediaContentProxyURL(c, requestID),
  680. )
  681. }
  682. writeGrokMediaResponse(c, resp, respBody, s.responseHeaderFilter)
  683. usage := grokMediaUsageFromResponse(endpoint, requestInfo, respBody)
  684. resultModel := requestModel
  685. resultBillingModel := requestModel
  686. if endpoint == GrokMediaEndpointVideoStatus {
  687. // Status has no request body model; use upstream status fields when billable.
  688. if m := strings.TrimSpace(usage.Model); m != "" {
  689. resultModel = m
  690. }
  691. if m := strings.TrimSpace(usage.BillingModel); m != "" {
  692. resultBillingModel = m
  693. }
  694. }
  695. return &OpenAIForwardResult{
  696. RequestID: requestIDHeader,
  697. UpstreamHeaders: resp.Header,
  698. ResponseID: usage.ResponseID,
  699. Usage: usage.Usage,
  700. Model: resultModel,
  701. BillingModel: resultBillingModel,
  702. UpstreamModel: upstreamModel,
  703. ResponseHeaders: resp.Header.Clone(),
  704. Duration: time.Since(startTime),
  705. ImageCount: usage.ImageCount,
  706. ImageSize: usage.ImageSize,
  707. ImageInputSize: usage.ImageInputSize,
  708. ImageOutputSizes: usage.ImageOutputSizes,
  709. VideoCount: usage.VideoCount,
  710. VideoResolution: usage.VideoResolution,
  711. VideoDurationSeconds: usage.VideoDurationSeconds,
  712. }, nil
  713. }
  714. func (s *OpenAIGatewayService) forwardGrokMediaVideoContent(
  715. ctx context.Context,
  716. c *gin.Context,
  717. account *Account,
  718. token, requestID string,
  719. startTime time.Time,
  720. ) (*OpenAIForwardResult, error) {
  721. statusURL, err := buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideoStatus, requestID)
  722. if err != nil {
  723. return nil, err
  724. }
  725. upstreamCtx, releaseUpstreamCtx := detachUpstreamContext(ctx)
  726. defer releaseUpstreamCtx()
  727. statusReq, err := http.NewRequestWithContext(
  728. WithHTTPUpstreamRedirectsDisabled(upstreamCtx),
  729. http.MethodGet,
  730. statusURL,
  731. nil,
  732. )
  733. if err != nil {
  734. return nil, err
  735. }
  736. statusReq.Header.Set("Authorization", "Bearer "+token)
  737. statusReq.Header.Set("Accept", "application/json")
  738. if account.IsGrokOAuth() && isGrokCLIProxyTarget(statusURL) {
  739. applyGrokCLIHeaders(statusReq.Header)
  740. }
  741. account.ApplyHeaderOverrides(statusReq.Header)
  742. proxyURL := ""
  743. if account.ProxyID != nil && account.Proxy != nil {
  744. proxyURL = account.Proxy.URL()
  745. }
  746. upstreamStart := time.Now()
  747. statusResp, err := s.httpUpstream.Do(WithAccountTrafficRequest(statusReq, account), proxyURL, account.ID, account.Concurrency)
  748. if err != nil {
  749. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  750. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
  751. }
  752. statusRequestID := firstNonEmpty(statusResp.Header.Get("x-request-id"), statusResp.Header.Get("xai-request-id"))
  753. if statusResp.StatusCode >= 300 {
  754. defer func() { _ = statusResp.Body.Close() }()
  755. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  756. if statusResp.StatusCode < 400 {
  757. return nil, fmt.Errorf("grok media status redirect is not allowed")
  758. }
  759. return s.handleGrokMediaErrorResponse(ctx, statusResp, c, account, statusRequestID, "")
  760. }
  761. statusBody, err := ReadUpstreamResponseBody(statusResp.Body, s.cfg, c, openAITooLargeError)
  762. _ = statusResp.Body.Close()
  763. if err != nil {
  764. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  765. return nil, err
  766. }
  767. contentURL, err := grokMediaSignedVideoContentURL(statusBody, requestID)
  768. if err != nil {
  769. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  770. return nil, err
  771. }
  772. signedContent := contentURL != ""
  773. if !signedContent {
  774. contentURL, err = buildGrokMediaURL(account, s.cfg, GrokMediaEndpointVideoContent, requestID)
  775. if err != nil {
  776. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  777. return nil, err
  778. }
  779. }
  780. contentReq, err := http.NewRequestWithContext(
  781. WithHTTPUpstreamRedirectsDisabled(upstreamCtx),
  782. http.MethodGet,
  783. contentURL,
  784. nil,
  785. )
  786. if err != nil {
  787. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  788. return nil, err
  789. }
  790. contentReq.Header.Set("Accept", "*/*")
  791. if c != nil {
  792. if rangeHeader := strings.TrimSpace(c.GetHeader("Range")); rangeHeader != "" {
  793. contentReq.Header.Set("Range", rangeHeader)
  794. }
  795. }
  796. if !signedContent {
  797. contentReq.Header.Set("Authorization", "Bearer "+token)
  798. if account.IsGrokOAuth() && isGrokCLIProxyTarget(contentURL) {
  799. applyGrokCLIHeaders(contentReq.Header)
  800. }
  801. account.ApplyHeaderOverrides(contentReq.Header)
  802. }
  803. contentResp, err := s.httpUpstream.Do(WithAccountTrafficRequest(contentReq, account), proxyURL, account.ID, account.Concurrency)
  804. SetOpsLatencyMs(c, OpsUpstreamLatencyMsKey, time.Since(upstreamStart).Milliseconds())
  805. if err != nil {
  806. return nil, s.handleOpenAIUpstreamTransportError(ctx, c, account, err, false)
  807. }
  808. defer func() { _ = contentResp.Body.Close() }()
  809. contentRequestID := firstNonEmpty(contentResp.Header.Get("x-request-id"), contentResp.Header.Get("xai-request-id"), statusRequestID)
  810. if contentResp.StatusCode >= 300 && contentResp.StatusCode < 400 {
  811. return nil, fmt.Errorf("grok media signed content redirect is not allowed")
  812. }
  813. if contentResp.StatusCode >= 400 && contentResp.StatusCode != http.StatusRequestedRangeNotSatisfiable {
  814. return s.handleGrokMediaErrorResponse(ctx, contentResp, c, account, contentRequestID, "")
  815. }
  816. s.updateGrokUsageFromResponse(withGrokTeamRateLimitModel(ctx, ""), account, contentResp.Header, contentResp.StatusCode)
  817. if err := writeGrokMediaContentResponse(c, contentResp); err != nil {
  818. return nil, err
  819. }
  820. // Content download is an alternate completion observation: when status body is
  821. // official done+video.url, attach billable units so the handler can claim once
  822. // (same path as status polling). Pending snapshot is merged in the handler.
  823. result := &OpenAIForwardResult{
  824. RequestID: contentRequestID,
  825. UpstreamHeaders: contentResp.Header,
  826. ResponseHeaders: contentResp.Header.Clone(),
  827. Duration: time.Since(startTime),
  828. }
  829. if billed := ExtractGrokVideoBillingFromStatusBody(statusBody, nil, requestID); billed != nil {
  830. result.ResponseID = firstNonEmpty(billed.ResponseID, strings.TrimSpace(requestID))
  831. result.Model = billed.Model
  832. result.BillingModel = billed.BillingModel
  833. result.UpstreamModel = billed.UpstreamModel
  834. result.VideoCount = billed.VideoCount
  835. result.VideoResolution = billed.VideoResolution
  836. result.VideoDurationSeconds = billed.VideoDurationSeconds
  837. }
  838. return result, nil
  839. }
  840. func grokMediaSignedVideoContentURL(body []byte, requestID string) (string, error) {
  841. rawURL := strings.TrimSpace(gjson.GetBytes(body, "video.url").String())
  842. if rawURL == "" {
  843. return "", nil
  844. }
  845. // An upstream Sub2API rewrites protected content URLs to its own proxy
  846. // endpoint. Treat that as an authenticated relay path, not as a signed URL;
  847. // the caller will rebuild it against the configured account base URL and
  848. // attach the upstream API key.
  849. if isGrokMediaVideoContentURL(rawURL, requestID) {
  850. return "", nil
  851. }
  852. parsed, err := url.Parse(rawURL)
  853. if err != nil || !strings.EqualFold(parsed.Scheme, "https") ||
  854. !strings.EqualFold(parsed.Hostname(), "vidgen.x.ai") ||
  855. (parsed.Port() != "" && parsed.Port() != "443") || parsed.User != nil {
  856. return "", fmt.Errorf("grok media status returned an unsupported video content URL")
  857. }
  858. return parsed.String(), nil
  859. }
  860. func isGrokCLIProxyTarget(rawURL string) bool {
  861. parsed, err := url.Parse(strings.TrimSpace(rawURL))
  862. return err == nil && strings.EqualFold(parsed.Hostname(), "cli-chat-proxy.grok.com")
  863. }
  864. func prepareGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) {
  865. if endpoint != GrokMediaEndpointImagesEdits {
  866. return body, contentType, nil
  867. }
  868. if gjson.ValidBytes(body) {
  869. out, err := normalizeGrokMediaJSONImageRefs(body)
  870. return out, contentType, err
  871. }
  872. mediaType, _, err := mime.ParseMediaType(strings.TrimSpace(contentType))
  873. if err != nil || !strings.EqualFold(mediaType, "multipart/form-data") {
  874. return body, contentType, nil
  875. }
  876. info := ParseGrokMediaRequest(contentType, body)
  877. payload := make(map[string]any)
  878. if info.Model != "" {
  879. payload["model"] = info.Model
  880. }
  881. if info.Prompt != "" {
  882. payload["prompt"] = info.Prompt
  883. }
  884. if info.N > 1 {
  885. payload["n"] = info.N
  886. }
  887. if info.Size != "" {
  888. payload["size"] = info.Size
  889. }
  890. if info.ImageResolution != "" {
  891. payload["resolution"] = info.ImageResolution
  892. }
  893. if info.AspectRatio != "" {
  894. payload["aspect_ratio"] = info.AspectRatio
  895. }
  896. images := make([]map[string]string, 0, len(info.InputImageURLs)+len(info.Uploads))
  897. for _, imageURL := range info.InputImageURLs {
  898. if imageURL = strings.TrimSpace(imageURL); imageURL != "" {
  899. images = append(images, grokMediaImageObject(imageURL))
  900. }
  901. }
  902. for _, upload := range info.Uploads {
  903. dataURL, err := openAIImageUploadToDataURL(upload)
  904. if err != nil {
  905. return nil, "", err
  906. }
  907. images = append(images, grokMediaImageObject(dataURL))
  908. }
  909. if len(images) > grokMediaMaxEditSourceImages {
  910. return nil, "", fmt.Errorf("a maximum of %d source images is supported for image edits", grokMediaMaxEditSourceImages)
  911. }
  912. if len(images) > 0 {
  913. payload["image"] = images[0]
  914. if len(images) > 1 {
  915. payload["images"] = images
  916. }
  917. }
  918. maskImageURL := strings.TrimSpace(info.MaskImageURL)
  919. if info.MaskUpload != nil {
  920. dataURL, err := openAIImageUploadToDataURL(*info.MaskUpload)
  921. if err != nil {
  922. return nil, "", err
  923. }
  924. maskImageURL = dataURL
  925. }
  926. if maskImageURL != "" {
  927. payload["mask"] = grokMediaImageObject(maskImageURL)
  928. }
  929. out, err := marshalOpenAIUpstreamJSON(payload)
  930. if err != nil {
  931. return nil, "", err
  932. }
  933. return out, "application/json", nil
  934. }
  935. func normalizeGrokMediaJSONImageRefs(body []byte) ([]byte, error) {
  936. info := ParseGrokMediaRequest("application/json", body)
  937. if len(info.InputImageURLs) > grokMediaMaxEditSourceImages {
  938. return nil, fmt.Errorf("a maximum of %d source images is supported for image edits", grokMediaMaxEditSourceImages)
  939. }
  940. out := body
  941. var err error
  942. for _, field := range []string{"image", "images", "mask"} {
  943. out, err = rewriteGrokMediaJSONImageField(out, field)
  944. if err != nil {
  945. return nil, err
  946. }
  947. }
  948. return out, nil
  949. }
  950. func rewriteGrokMediaJSONImageField(body []byte, path string) ([]byte, error) {
  951. value := gjson.GetBytes(body, path)
  952. if !value.Exists() {
  953. return body, nil
  954. }
  955. if value.IsArray() {
  956. rewritten := make([]map[string]string, 0, len(value.Array()))
  957. for _, item := range value.Array() {
  958. imageURL := extractGrokMediaImageURL(item)
  959. if imageURL == "" {
  960. return body, nil
  961. }
  962. rewritten = append(rewritten, grokMediaImageObject(imageURL))
  963. }
  964. out, err := sjson.SetBytes(body, path, rewritten)
  965. if err != nil {
  966. return nil, fmt.Errorf("rewrite grok media %s: %w", path, err)
  967. }
  968. return out, nil
  969. }
  970. imageURL := extractGrokMediaImageURL(value)
  971. if imageURL == "" {
  972. return body, nil
  973. }
  974. out, err := sjson.SetBytes(body, path, grokMediaImageObject(imageURL))
  975. if err != nil {
  976. return nil, fmt.Errorf("rewrite grok media %s: %w", path, err)
  977. }
  978. return out, nil
  979. }
  980. func normalizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) {
  981. if !endpoint.RequiresRequestBody() || !gjson.ValidBytes(body) {
  982. return body, contentType, nil
  983. }
  984. var imageFields []string
  985. switch endpoint {
  986. case GrokMediaEndpointImagesEdits:
  987. imageFields = []string{"image", "images", "mask"}
  988. case GrokMediaEndpointVideosGenerations:
  989. imageFields = []string{"image", "images", "reference_images"}
  990. }
  991. var err error
  992. body, err = canonicalizeGrokMediaImageURLFields(body, imageFields...)
  993. if err != nil {
  994. return nil, "", err
  995. }
  996. info := ParseGrokMediaRequest(contentType, body)
  997. upstreamModel := NormalizeGrokMediaModelForEndpoint(endpoint, info.Model, info.HasInputImage())
  998. if upstreamModel == "" || upstreamModel == info.Model {
  999. return body, contentType, nil
  1000. }
  1001. out, err := sjson.SetBytes(body, "model", upstreamModel)
  1002. if err != nil {
  1003. return nil, "", fmt.Errorf("rewrite grok media model: %w", err)
  1004. }
  1005. return out, contentType, nil
  1006. }
  1007. func canonicalizeGrokMediaImageURLFields(body []byte, fields ...string) ([]byte, error) {
  1008. out := body
  1009. for _, field := range fields {
  1010. value := gjson.GetBytes(out, field)
  1011. if !value.Exists() {
  1012. continue
  1013. }
  1014. if value.IsArray() {
  1015. for index := range value.Array() {
  1016. var err error
  1017. out, err = canonicalizeGrokMediaImageURLObject(out, fmt.Sprintf("%s.%d", field, index))
  1018. if err != nil {
  1019. return nil, err
  1020. }
  1021. }
  1022. continue
  1023. }
  1024. var err error
  1025. out, err = canonicalizeGrokMediaImageURLObject(out, field)
  1026. if err != nil {
  1027. return nil, err
  1028. }
  1029. }
  1030. return out, nil
  1031. }
  1032. func canonicalizeGrokMediaImageURLObject(body []byte, path string) ([]byte, error) {
  1033. legacyPath := path + ".image_url"
  1034. legacy := gjson.GetBytes(body, legacyPath)
  1035. if !legacy.Exists() {
  1036. return body, nil
  1037. }
  1038. out := body
  1039. if strings.TrimSpace(gjson.GetBytes(out, path+".url").String()) == "" {
  1040. var err error
  1041. out, err = sjson.SetBytes(out, path+".url", legacy.Value())
  1042. if err != nil {
  1043. return nil, fmt.Errorf("normalize grok media image url: %w", err)
  1044. }
  1045. }
  1046. out, err := sjson.DeleteBytes(out, legacyPath)
  1047. if err != nil {
  1048. return nil, fmt.Errorf("remove legacy grok media image url: %w", err)
  1049. }
  1050. return out, nil
  1051. }
  1052. func sanitizeGrokMediaForwardBody(endpoint GrokMediaEndpoint, body []byte, contentType string) ([]byte, string, error) {
  1053. if !endpoint.RequiresRequestBody() || !gjson.ValidBytes(body) {
  1054. return body, contentType, nil
  1055. }
  1056. switch endpoint {
  1057. case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits:
  1058. out, err := applyGrokImagineImageGeometry(body)
  1059. if err != nil {
  1060. return nil, "", fmt.Errorf("sanitize grok media size: %w", err)
  1061. }
  1062. return out, contentType, nil
  1063. default:
  1064. return body, contentType, nil
  1065. }
  1066. }
  1067. func (r GrokMediaRequestInfo) HasInputImage() bool {
  1068. return len(r.InputImageURLs) > 0 || len(r.Uploads) > 0
  1069. }
  1070. // NormalizeGrokMediaModelForEndpoint resolves the built-in upstream model alias
  1071. // for a media endpoint before account-level model mapping and scheduling.
  1072. func NormalizeGrokMediaModelForEndpoint(endpoint GrokMediaEndpoint, model string, hasInputImage bool) string {
  1073. model = strings.TrimSpace(model)
  1074. switch endpoint {
  1075. case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits:
  1076. if model == "grok-imagine" {
  1077. return "grok-imagine-image-quality"
  1078. }
  1079. case GrokMediaEndpointVideosGenerations:
  1080. // xAI's 1.5 model is image-to-video only. Keep the requested model
  1081. // unchanged when the image is missing so the upstream returns its
  1082. // documented invalid-argument response instead of silently switching
  1083. // models and pricing.
  1084. _ = hasInputImage
  1085. }
  1086. return model
  1087. }
  1088. type grokMediaUsageMetadata struct {
  1089. ResponseID string
  1090. Usage OpenAIUsage
  1091. Model string
  1092. BillingModel string
  1093. ImageCount int
  1094. ImageSize string
  1095. ImageInputSize string
  1096. ImageOutputSizes []string
  1097. VideoCount int
  1098. VideoResolution string
  1099. VideoDurationSeconds int
  1100. }
  1101. func grokMediaUsageFromResponse(endpoint GrokMediaEndpoint, requestInfo GrokMediaRequestInfo, responseBody []byte) grokMediaUsageMetadata {
  1102. usage, _ := extractOpenAIUsageFromJSONBytes(responseBody)
  1103. meta := grokMediaUsageMetadata{Usage: usage}
  1104. switch endpoint {
  1105. case GrokMediaEndpointImagesGenerations, GrokMediaEndpointImagesEdits:
  1106. meta.ImageCount = countOpenAIResponseImageOutputsFromJSONBytes(responseBody)
  1107. meta.ImageSize = requestInfo.SizeTier
  1108. meta.ImageInputSize = requestInfo.Size
  1109. meta.ImageOutputSizes = collectOpenAIResponseImageOutputSizesFromJSONBytes(responseBody)
  1110. case GrokMediaEndpointVideosGenerations, GrokMediaEndpointVideosEdits, GrokMediaEndpointVideosExtensions:
  1111. // Async video: capture request_id + create-time pricing params only.
  1112. // Billable VideoCount is set later when status polling observes video.url.
  1113. meta.ResponseID = extractGrokMediaVideoRequestID(responseBody)
  1114. meta.VideoResolution = requestInfo.Resolution
  1115. meta.VideoDurationSeconds = requestInfo.DurationSeconds
  1116. case GrokMediaEndpointVideoStatus:
  1117. // Prefer status-body URL success + upstream duration/resolution when present.
  1118. if IsGrokVideoStatusBillable(responseBody) {
  1119. // provisional units; handler merges with pending snapshot before RecordUsage.
  1120. if billed := ExtractGrokVideoBillingFromStatusBody(responseBody, nil, ""); billed != nil {
  1121. meta.ResponseID = billed.ResponseID
  1122. meta.Model = billed.Model
  1123. meta.BillingModel = billed.BillingModel
  1124. meta.VideoCount = billed.VideoCount
  1125. meta.VideoResolution = billed.VideoResolution
  1126. meta.VideoDurationSeconds = billed.VideoDurationSeconds
  1127. }
  1128. }
  1129. }
  1130. return meta
  1131. }
  1132. func extractGrokMediaVideoRequestID(body []byte) string {
  1133. if len(body) == 0 || !gjson.ValidBytes(body) {
  1134. return ""
  1135. }
  1136. for _, path := range []string{"request_id", "id", "data.request_id", "data.id", "video.request_id", "video.id", "task_id", "data.task_id", "video.task_id"} {
  1137. if id := strings.TrimSpace(gjson.GetBytes(body, path).String()); id != "" {
  1138. return id
  1139. }
  1140. }
  1141. return ""
  1142. }
  1143. func (s *OpenAIGatewayService) handleGrokMediaErrorResponse(
  1144. ctx context.Context,
  1145. resp *http.Response,
  1146. c *gin.Context,
  1147. account *Account,
  1148. requestIDHeader string,
  1149. requestedModel string,
  1150. ) (*OpenAIForwardResult, error) {
  1151. body := s.readUpstreamErrorBody(resp)
  1152. // Reconcile readiness before configurable passthrough branches can return;
  1153. // otherwise a Grok 429 can remain schedulable.
  1154. s.handleGrokAccountUpstreamError(ctx, account, resp.StatusCode, resp.Header, body)
  1155. upstreamMsg := sanitizeUpstreamErrorMessage(strings.TrimSpace(extractUpstreamErrorMessage(body)))
  1156. if upstreamMsg == "" {
  1157. upstreamMsg = fmt.Sprintf("xAI upstream returned status %d", resp.StatusCode)
  1158. }
  1159. upstreamDetail := ""
  1160. if s.cfg != nil && s.cfg.Gateway.LogUpstreamErrorBody {
  1161. maxBytes := s.cfg.Gateway.LogUpstreamErrorBodyMaxBytes
  1162. if maxBytes <= 0 {
  1163. maxBytes = 2048
  1164. }
  1165. upstreamDetail = truncateString(string(body), maxBytes)
  1166. }
  1167. setOpsUpstreamError(c, resp.StatusCode, upstreamMsg, upstreamDetail)
  1168. if isGrokContentPolicyRejection(resp.StatusCode, body) {
  1169. clientMsg := grokContentPolicyClientMessage(body)
  1170. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  1171. ProxyID: opsUpstreamProxyID(account),
  1172. ProxyName: opsUpstreamProxyName(account),
  1173. Platform: account.Platform,
  1174. AccountID: account.ID,
  1175. AccountName: account.Name,
  1176. UpstreamStatusCode: resp.StatusCode,
  1177. UpstreamRequestID: requestIDHeader,
  1178. Kind: "http_error",
  1179. Message: clientMsg,
  1180. Detail: upstreamDetail,
  1181. })
  1182. MarkResponseCommitted(c)
  1183. writeGrokMediaErrorResponse(c, http.StatusForbidden, "invalid_request_error", clientMsg)
  1184. return nil, fmt.Errorf("grok content policy rejection: %s", clientMsg)
  1185. }
  1186. if status, errType, errMsg, matched := applyErrorPassthroughRule(
  1187. c,
  1188. account.Platform,
  1189. resp.StatusCode,
  1190. body,
  1191. http.StatusBadGateway,
  1192. "upstream_error",
  1193. "Upstream request failed",
  1194. ); matched {
  1195. MarkResponseCommitted(c)
  1196. writeGrokMediaErrorResponse(c, status, errType, errMsg)
  1197. return nil, fmt.Errorf("upstream error: %d (passthrough rule matched) message=%s", resp.StatusCode, upstreamMsg)
  1198. }
  1199. if !account.ShouldHandleErrorCode(resp.StatusCode) {
  1200. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  1201. ProxyID: opsUpstreamProxyID(account),
  1202. ProxyName: opsUpstreamProxyName(account),
  1203. Platform: account.Platform,
  1204. AccountID: account.ID,
  1205. AccountName: account.Name,
  1206. UpstreamStatusCode: resp.StatusCode,
  1207. UpstreamRequestID: requestIDHeader,
  1208. Kind: "http_error",
  1209. Message: upstreamMsg,
  1210. Detail: upstreamDetail,
  1211. })
  1212. MarkResponseCommitted(c)
  1213. writeGrokMediaErrorResponse(c, http.StatusInternalServerError, "upstream_error", "Upstream gateway error")
  1214. return nil, fmt.Errorf("upstream error: %d (not in custom error codes) message=%s", resp.StatusCode, upstreamMsg)
  1215. }
  1216. kind := "http_error"
  1217. if s.shouldFailoverGrokUpstreamError(resp.StatusCode, body) {
  1218. kind = "failover"
  1219. }
  1220. appendOpsUpstreamError(c, OpsUpstreamErrorEvent{
  1221. ProxyID: opsUpstreamProxyID(account),
  1222. ProxyName: opsUpstreamProxyName(account),
  1223. Platform: account.Platform,
  1224. AccountID: account.ID,
  1225. AccountName: account.Name,
  1226. UpstreamStatusCode: resp.StatusCode,
  1227. UpstreamRequestID: requestIDHeader,
  1228. Kind: kind,
  1229. Message: upstreamMsg,
  1230. Detail: upstreamDetail,
  1231. })
  1232. if kind == "failover" {
  1233. retryable, retryDelay, retryDeadline, retryMax := grokSameAccountRetryMetadata(account, resp.StatusCode, body)
  1234. return nil, &UpstreamFailoverError{
  1235. StatusCode: resp.StatusCode,
  1236. ResponseBody: body,
  1237. ResponseHeaders: resp.Header.Clone(),
  1238. RetryableOnSameAccount: retryable,
  1239. RequestScopedTransient: retryable && resp.StatusCode == http.StatusTooManyRequests,
  1240. SameAccountRetryDelay: retryDelay,
  1241. SameAccountRetryDeadline: retryDeadline,
  1242. SameAccountRetryMax: retryMax,
  1243. }
  1244. }
  1245. MarkResponseCommitted(c)
  1246. writeGrokMediaErrorResponse(c, resp.StatusCode, grokMediaErrorType(resp.StatusCode), upstreamMsg)
  1247. return nil, fmt.Errorf("upstream error: %d %s", resp.StatusCode, upstreamMsg)
  1248. }
  1249. func grokMediaErrorType(statusCode int) string {
  1250. switch statusCode {
  1251. case http.StatusBadRequest:
  1252. return "invalid_request_error"
  1253. case http.StatusNotFound:
  1254. return "not_found_error"
  1255. case http.StatusTooManyRequests:
  1256. return "rate_limit_error"
  1257. default:
  1258. return "upstream_error"
  1259. }
  1260. }
  1261. func writeGrokMediaErrorResponse(c *gin.Context, statusCode int, errType, message string) {
  1262. if c == nil || c.Writer == nil || c.Writer.Written() {
  1263. return
  1264. }
  1265. c.JSON(statusCode, gin.H{
  1266. "error": gin.H{
  1267. "type": strings.TrimSpace(errType),
  1268. "message": strings.TrimSpace(message),
  1269. },
  1270. })
  1271. }
  1272. func writeGrokMediaResponse(c *gin.Context, resp *http.Response, body []byte, filter *responseheaders.CompiledHeaderFilter) {
  1273. if c == nil || resp == nil {
  1274. return
  1275. }
  1276. writeOpenAIPassthroughResponseHeaders(c.Writer.Header(), resp.Header, filter)
  1277. contentType := strings.TrimSpace(resp.Header.Get("Content-Type"))
  1278. if contentType == "" {
  1279. contentType = "application/json"
  1280. }
  1281. c.Data(resp.StatusCode, contentType, body)
  1282. }
  1283. func writeGrokMediaContentResponse(c *gin.Context, resp *http.Response) error {
  1284. if c == nil || resp == nil || resp.Body == nil {
  1285. return fmt.Errorf("grok media content response is incomplete")
  1286. }
  1287. for _, name := range []string{
  1288. "Content-Type",
  1289. "Content-Length",
  1290. "Content-Range",
  1291. "Accept-Ranges",
  1292. "Content-Disposition",
  1293. } {
  1294. if value := strings.TrimSpace(resp.Header.Get(name)); value != "" {
  1295. c.Header(name, value)
  1296. }
  1297. }
  1298. if strings.TrimSpace(c.Writer.Header().Get("Content-Length")) == "" && resp.ContentLength >= 0 {
  1299. c.Header("Content-Length", strconv.FormatInt(resp.ContentLength, 10))
  1300. }
  1301. if strings.TrimSpace(c.Writer.Header().Get("Content-Type")) == "" {
  1302. c.Header("Content-Type", "application/octet-stream")
  1303. }
  1304. c.Status(resp.StatusCode)
  1305. MarkResponseCommitted(c)
  1306. _, err := io.Copy(c.Writer, resp.Body)
  1307. return err
  1308. }
  1309. func rewriteGrokMediaVideoContentURLs(body []byte, requestID, proxyURL string) []byte {
  1310. if len(body) == 0 || strings.TrimSpace(requestID) == "" || strings.TrimSpace(proxyURL) == "" || !gjson.ValidBytes(body) {
  1311. return body
  1312. }
  1313. decoder := json.NewDecoder(bytes.NewReader(body))
  1314. decoder.UseNumber()
  1315. var value any
  1316. if err := decoder.Decode(&value); err != nil {
  1317. return body
  1318. }
  1319. changed := rewriteGrokMediaKnownVideoURL(&value, proxyURL)
  1320. if rewriteGrokMediaVideoContentURLValue(&value, requestID, proxyURL) {
  1321. changed = true
  1322. }
  1323. if !changed {
  1324. return body
  1325. }
  1326. rewritten, err := json.Marshal(value)
  1327. if err != nil {
  1328. return body
  1329. }
  1330. return rewritten
  1331. }
  1332. func rewriteGrokMediaKnownVideoURL(value *any, proxyURL string) bool {
  1333. if value == nil {
  1334. return false
  1335. }
  1336. root, ok := (*value).(map[string]any)
  1337. if !ok {
  1338. return false
  1339. }
  1340. video, ok := root["video"].(map[string]any)
  1341. if !ok {
  1342. return false
  1343. }
  1344. rawURL, ok := video["url"].(string)
  1345. if !ok || strings.TrimSpace(rawURL) == "" {
  1346. return false
  1347. }
  1348. video["url"] = proxyURL
  1349. return true
  1350. }
  1351. func rewriteGrokMediaVideoContentURLValue(value *any, requestID, proxyURL string) bool {
  1352. if value == nil {
  1353. return false
  1354. }
  1355. switch typed := (*value).(type) {
  1356. case map[string]any:
  1357. changed := false
  1358. for key, child := range typed {
  1359. childValue := child
  1360. if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) {
  1361. typed[key] = childValue
  1362. changed = true
  1363. }
  1364. }
  1365. return changed
  1366. case []any:
  1367. changed := false
  1368. for index, child := range typed {
  1369. childValue := child
  1370. if rewriteGrokMediaVideoContentURLValue(&childValue, requestID, proxyURL) {
  1371. typed[index] = childValue
  1372. changed = true
  1373. }
  1374. }
  1375. return changed
  1376. case string:
  1377. if isGrokMediaVideoContentURL(typed, requestID) {
  1378. *value = proxyURL
  1379. return true
  1380. }
  1381. }
  1382. return false
  1383. }
  1384. func isGrokMediaVideoContentURL(rawURL, requestID string) bool {
  1385. parsed, err := url.Parse(strings.TrimSpace(rawURL))
  1386. if err != nil || parsed.Path == "" {
  1387. return false
  1388. }
  1389. segments := strings.Split(strings.Trim(parsed.EscapedPath(), "/"), "/")
  1390. if len(segments) < 3 {
  1391. return false
  1392. }
  1393. requestID = strings.Trim(requestID, "/")
  1394. decodedID, err := url.PathUnescape(segments[len(segments)-2])
  1395. if err != nil {
  1396. return false
  1397. }
  1398. return segments[len(segments)-3] == "videos" &&
  1399. decodedID == requestID &&
  1400. segments[len(segments)-1] == "content"
  1401. }
  1402. func grokMediaContentProxyURL(c *gin.Context, requestID string) string {
  1403. if c == nil || c.Request == nil || c.Request.URL == nil || strings.TrimSpace(requestID) == "" {
  1404. return ""
  1405. }
  1406. pathPrefix := ""
  1407. if strings.HasPrefix(c.Request.URL.Path, "/v1/") {
  1408. pathPrefix = "/v1"
  1409. }
  1410. return pathPrefix + "/videos/" + url.PathEscape(strings.Trim(requestID, "/")) + "/content"
  1411. }