grok_media.go 28 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725
  1. package handler
  2. import (
  3. "context"
  4. "errors"
  5. "net/http"
  6. "strconv"
  7. "strings"
  8. "time"
  9. pkghttputil "github.com/Wei-Shaw/sub2api/internal/pkg/httputil"
  10. "github.com/Wei-Shaw/sub2api/internal/pkg/ip"
  11. "github.com/Wei-Shaw/sub2api/internal/pkg/logger"
  12. middleware2 "github.com/Wei-Shaw/sub2api/internal/server/middleware"
  13. "github.com/Wei-Shaw/sub2api/internal/service"
  14. "github.com/gin-gonic/gin"
  15. "go.uber.org/zap"
  16. )
  17. // GrokImages handles xAI image generation/editing through Grok groups.
  18. func (h *OpenAIGatewayHandler) GrokImages(c *gin.Context) {
  19. endpoint := service.GrokMediaEndpointImagesGenerations
  20. if strings.Contains(c.Request.URL.Path, "/images/edits") {
  21. endpoint = service.GrokMediaEndpointImagesEdits
  22. }
  23. h.handleGrokMedia(c, endpoint, "")
  24. }
  25. // GrokVideoGeneration handles xAI video generation through Grok groups.
  26. func (h *OpenAIGatewayHandler) GrokVideoGeneration(c *gin.Context) {
  27. h.handleGrokMedia(c, service.GrokMediaEndpointVideosGenerations, "")
  28. }
  29. // GrokVideoEdit handles asynchronous xAI video edits through Grok groups.
  30. func (h *OpenAIGatewayHandler) GrokVideoEdit(c *gin.Context) {
  31. h.handleGrokMedia(c, service.GrokMediaEndpointVideosEdits, "")
  32. }
  33. // GrokVideoExtension handles asynchronous xAI video extensions through Grok groups.
  34. func (h *OpenAIGatewayHandler) GrokVideoExtension(c *gin.Context) {
  35. h.handleGrokMedia(c, service.GrokMediaEndpointVideosExtensions, "")
  36. }
  37. // GrokVideoStatus handles xAI video status retrieval through Grok groups.
  38. func (h *OpenAIGatewayHandler) GrokVideoStatus(c *gin.Context) {
  39. h.handleGrokMedia(c, service.GrokMediaEndpointVideoStatus, c.Param("request_id"))
  40. }
  41. // GrokVideoContent proxies downloadable video content through the task's upstream account.
  42. func (h *OpenAIGatewayHandler) GrokVideoContent(c *gin.Context) {
  43. h.handleGrokMedia(c, service.GrokMediaEndpointVideoContent, c.Param("request_id"))
  44. }
  45. func (h *OpenAIGatewayHandler) handleGrokMedia(c *gin.Context, endpoint service.GrokMediaEndpoint, requestID string) {
  46. streamStarted := false
  47. defer h.recoverResponsesPanic(c, &streamStarted)
  48. requestStart := time.Now()
  49. apiKey, ok := middleware2.GetAPIKeyFromContext(c)
  50. if !ok {
  51. h.errorResponse(c, http.StatusUnauthorized, "authentication_error", "Invalid API key")
  52. return
  53. }
  54. subject, ok := middleware2.GetAuthSubjectFromContext(c)
  55. if !ok {
  56. h.errorResponse(c, http.StatusInternalServerError, "api_error", "User context not found")
  57. return
  58. }
  59. reqLog := requestLogger(
  60. c,
  61. "handler.openai_gateway.grok_media",
  62. zap.Int64("user_id", subject.UserID),
  63. zap.Int64("api_key_id", apiKey.ID),
  64. zap.Any("group_id", apiKey.GroupID),
  65. zap.String("endpoint", string(endpoint)),
  66. )
  67. if !h.ensureResponsesDependencies(c, reqLog) {
  68. return
  69. }
  70. var body []byte
  71. var err error
  72. if endpoint.RequiresRequestBody() {
  73. body, err = pkghttputil.ReadRequestBodyWithPrealloc(c.Request)
  74. if err != nil {
  75. if maxErr, ok := extractMaxBytesError(err); ok {
  76. h.errorResponse(c, http.StatusRequestEntityTooLarge, "invalid_request_error", buildBodyTooLargeMessage(maxErr.Limit))
  77. return
  78. }
  79. h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Failed to read request body")
  80. return
  81. }
  82. if len(body) == 0 {
  83. h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "Request body is empty")
  84. return
  85. }
  86. }
  87. contentType := c.GetHeader("Content-Type")
  88. requestInfo := service.ParseGrokMediaRequest(contentType, body)
  89. requestModel := requestInfo.Model
  90. routingModel := service.NormalizeGrokMediaModelForEndpoint(endpoint, requestModel, requestInfo.HasInputImage())
  91. if endpoint.IsGenerationRequest() && strings.TrimSpace(requestModel) == "" {
  92. h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "model is required")
  93. return
  94. }
  95. if endpoint.IsVideoLookupRequest() && strings.TrimSpace(requestID) == "" {
  96. h.errorResponse(c, http.StatusBadRequest, "invalid_request_error", "request_id is required")
  97. return
  98. }
  99. reqLog = reqLog.With(zap.String("model", requestModel))
  100. setOpsRequestContext(c, requestModel, false)
  101. setOpsEndpointContext(c, "", int16(service.RequestTypeSync))
  102. if endpoint.IsGenerationRequest() {
  103. if !service.GroupAllowsImageGeneration(apiKey.Group) {
  104. h.errorResponse(c, http.StatusForbidden, "permission_error", service.ImageGenerationPermissionMessage())
  105. return
  106. }
  107. if moderationBody := requestInfo.ModerationBody(); len(moderationBody) > 0 {
  108. decision := h.checkSecurityAudit(c, reqLog, apiKey, subject, service.ContentModerationProtocolOpenAIImages, requestModel, moderationBody)
  109. if decision != nil && !decision.AllowNextStage {
  110. h.openAISecurityAuditError(c, decision)
  111. return
  112. }
  113. }
  114. imageReleaseFunc, acquired := h.acquireImageGenerationSlot(c, streamStarted)
  115. if !acquired {
  116. return
  117. }
  118. if imageReleaseFunc != nil {
  119. defer imageReleaseFunc()
  120. }
  121. }
  122. if h.errorPassthroughService != nil {
  123. service.BindErrorPassthroughService(c, h.errorPassthroughService)
  124. }
  125. subscription, _ := middleware2.GetSubscriptionFromContext(c)
  126. service.SetOpsLatencyMs(c, service.OpsAuthLatencyMsKey, time.Since(requestStart).Milliseconds())
  127. userReleaseFunc, acquired := h.acquireResponsesUserSlot(c, subject.UserID, subject.Concurrency, false, &streamStarted, reqLog)
  128. if !acquired {
  129. return
  130. }
  131. if userReleaseFunc != nil {
  132. defer userReleaseFunc()
  133. }
  134. if err := h.billingCacheService.CheckBillingEligibility(c.Request.Context(), apiKey.User, apiKey, apiKey.Group, subscription, service.QuotaPlatform(c.Request.Context(), apiKey)); err != nil {
  135. reqLog.Info("grok_media.billing_eligibility_check_failed", zap.Error(err))
  136. status, code, message, retryAfter := billingErrorDetails(err)
  137. if retryAfter > 0 {
  138. c.Header("Retry-After", strconv.Itoa(retryAfter))
  139. }
  140. h.errorResponse(c, status, code, message)
  141. return
  142. }
  143. sessionSeed := body
  144. if len(sessionSeed) == 0 && strings.TrimSpace(requestID) != "" {
  145. sessionSeed = []byte(requestID)
  146. }
  147. sessionHash := h.gatewayService.GenerateExplicitSessionHash(c, sessionSeed)
  148. boundLookupAccountID := int64(0)
  149. if endpoint.IsVideoLookupRequest() {
  150. sessionHash = service.GrokMediaVideoRequestSessionHash(requestID, subject.UserID, apiKey.ID)
  151. boundLookupAccountID, err = h.gatewayService.ResolveGrokMediaVideoRequestAccount(
  152. c.Request.Context(), apiKey.GroupID, requestID, subject.UserID, apiKey.ID,
  153. )
  154. if err != nil || boundLookupAccountID <= 0 {
  155. reqLog.Info("grok_media.video_lookup_owner_binding_missing", zap.Error(err))
  156. h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
  157. return
  158. }
  159. }
  160. // Grok 媒体(图片/视频生成与视频查询)按媒体倍率计费,不在 token 利润门
  161. // 范围内:显式豁免,防止 service 层防御性装门按文本 D 误过滤媒体请求,
  162. // 也防止已计费的在途视频任务因绑定账号被门排除而查询返回伪 404。
  163. requestCtx := service.WithOpenAIProfitControlSuppressed(c.Request.Context())
  164. profitVetoCount := 0
  165. failedAccountIDs := make(map[int64]struct{})
  166. sameAccountRetryCount := make(map[int64]int)
  167. var lastFailoverErr *service.UpstreamFailoverError
  168. var oauth429FailoverState service.OpenAIOAuth429FailoverState
  169. mediaEligibilityRejected := false
  170. switchCount := 0
  171. videoCreateStartedAt := ""
  172. if isGrokVideoCreateEndpoint(endpoint) {
  173. videoCreateStartedAt = service.GrokVideoPendingCreatedAtNow()
  174. }
  175. maxAccountSwitches := h.maxAccountSwitches
  176. if maxAccountSwitches <= 0 {
  177. maxAccountSwitches = 3
  178. }
  179. routingStart := time.Now()
  180. requiredCapability := grokMediaRequiredCapability(endpoint)
  181. for {
  182. if failoverClientGone(c) {
  183. return
  184. }
  185. selection, scheduleDecision, err := h.gatewayService.SelectAccountWithSchedulerForCapability(
  186. requestCtx,
  187. apiKey.GroupID,
  188. "",
  189. sessionHash,
  190. routingModel,
  191. failedAccountIDs,
  192. service.OpenAIUpstreamTransportHTTPSSE,
  193. requiredCapability,
  194. false,
  195. false,
  196. false,
  197. service.PlatformGrok,
  198. )
  199. if err != nil {
  200. if failoverClientGone(c) {
  201. reqLog.Info("grok_media.account_select_aborted_client_disconnected", zap.Error(err))
  202. return
  203. }
  204. reqLog.Warn("grok_media.account_select_failed",
  205. zap.Error(err),
  206. zap.Int("excluded_account_count", len(failedAccountIDs)),
  207. )
  208. if endpoint.IsGenerationRequest() && errors.Is(err, service.ErrNoAvailableAccounts) &&
  209. (len(failedAccountIDs) == 0 || (mediaEligibilityRejected && lastFailoverErr == nil)) {
  210. markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
  211. h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
  212. return
  213. }
  214. if len(failedAccountIDs) == 0 {
  215. cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
  216. if !cls.ModelNotFound {
  217. markOpsRoutingCapacityLimitedIfNoAvailable(c, err)
  218. }
  219. h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
  220. return
  221. }
  222. if lastFailoverErr != nil {
  223. h.handleFailoverExhausted(c, lastFailoverErr, false)
  224. } else {
  225. h.errorResponse(c, http.StatusBadGateway, "api_error", "Upstream request failed")
  226. }
  227. return
  228. }
  229. if selection == nil || selection.Account == nil {
  230. if endpoint.IsGenerationRequest() {
  231. markOpsRoutingCapacityLimited(c)
  232. h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
  233. return
  234. }
  235. cls := classifyNoAccountErrorFromGin(c, h.gatewayService, apiKey, requestModel, routingModel, service.PlatformGrok)
  236. if !cls.ModelNotFound {
  237. markOpsRoutingCapacityLimited(c)
  238. }
  239. h.errorResponse(c, cls.Status, cls.ErrType, cls.Message)
  240. return
  241. }
  242. if boundLookupAccountID > 0 && selection.Account.ID != boundLookupAccountID {
  243. reqLog.Warn("grok_media.video_lookup_bound_account_unavailable",
  244. zap.Int64("bound_account_id", boundLookupAccountID),
  245. zap.Int64("selected_account_id", selection.Account.ID),
  246. )
  247. h.errorResponse(c, http.StatusNotFound, "not_found_error", "Video request not found")
  248. return
  249. }
  250. reqLog.Debug("grok_media.account_schedule_decision",
  251. zap.String("layer", scheduleDecision.Layer),
  252. zap.Bool("sticky_session_hit", scheduleDecision.StickySessionHit),
  253. zap.Int("candidate_count", scheduleDecision.CandidateCount),
  254. zap.Int("top_k", scheduleDecision.TopK),
  255. zap.Int64("latency_ms", scheduleDecision.LatencyMs),
  256. zap.Float64("load_skew", scheduleDecision.LoadSkew),
  257. )
  258. account := selection.Account
  259. if endpoint.IsGenerationRequest() {
  260. eligible, eligibilityReason, eligibilityErr := h.ensureGrokMediaAccountEligibility(requestCtx, account)
  261. if !eligible {
  262. mediaEligibilityRejected = true
  263. failedAccountIDs[account.ID] = struct{}{}
  264. reqLog.Warn("grok_media.account_eligibility_rejected",
  265. zap.Int64("account_id", account.ID),
  266. zap.String("reason", eligibilityReason),
  267. zap.Bool("probe_failed", eligibilityErr != nil),
  268. )
  269. if switchCount >= maxAccountSwitches {
  270. markOpsRoutingCapacityLimited(c)
  271. h.errorResponse(c, http.StatusServiceUnavailable, "grok_media_no_eligible_account", "No eligible Grok media accounts")
  272. return
  273. }
  274. switchCount++
  275. continue
  276. }
  277. }
  278. sessionHash = ensureOpenAIPoolModeSessionHash(sessionHash, account)
  279. setOpsSelectedAccount(c, account.ID, account.Platform)
  280. accountReleaseFunc, slotResult := h.acquireResponsesAccountSlot(c, apiKey.GroupID, sessionHash, selection, false, &streamStarted, reqLog)
  281. if slotResult == openAISlotAcquireProfitVetoed {
  282. // 媒体路径已显式豁免利润门(suppress 标记),此分支仅防御性兜底,
  283. // 同样受否决上限约束。
  284. if !recordOpenAIProfitVeto(failedAccountIDs, account.ID, &profitVetoCount) {
  285. h.handleOpenAIProfitVetoExhausted(c, streamStarted, reqLog, profitVetoCount)
  286. return
  287. }
  288. continue
  289. }
  290. if slotResult != openAISlotAcquireOK {
  291. return
  292. }
  293. service.SetOpsLatencyMs(c, service.OpsRoutingLatencyMsKey, time.Since(routingStart).Milliseconds())
  294. forwardStart := time.Now()
  295. writerSizeBeforeForward := c.Writer.Size()
  296. result, err := func() (*service.OpenAIForwardResult, error) {
  297. defer func() {
  298. if accountReleaseFunc != nil {
  299. accountReleaseFunc()
  300. }
  301. }()
  302. return h.gatewayService.ForwardGrokMedia(requestCtx, c, account, endpoint, requestID, body, contentType)
  303. }()
  304. forwardDurationMs := time.Since(forwardStart).Milliseconds()
  305. upstreamLatencyMs, _ := getContextInt64(c, service.OpsUpstreamLatencyMsKey)
  306. responseLatencyMs := forwardDurationMs
  307. if upstreamLatencyMs > 0 && forwardDurationMs > upstreamLatencyMs {
  308. responseLatencyMs = forwardDurationMs - upstreamLatencyMs
  309. }
  310. service.SetOpsLatencyMs(c, service.OpsResponseLatencyMsKey, responseLatencyMs)
  311. if err != nil {
  312. var failoverErr *service.UpstreamFailoverError
  313. if errors.As(err, &failoverErr) {
  314. if failoverClientGone(c) {
  315. reqLog.Info("grok_media.failover_aborted_client_disconnected",
  316. zap.Int64("account_id", account.ID),
  317. zap.Int("upstream_status", failoverErr.StatusCode),
  318. )
  319. return
  320. }
  321. if failoverErr.ShouldReportAccountScheduleFailure() {
  322. h.gatewayService.ReportOpenAIAccountScheduleResult(account, grokMediaScheduleModel(account, routingModel, nil), false, nil)
  323. }
  324. if c.Writer.Size() != writerSizeBeforeForward {
  325. h.handleFailoverExhausted(c, failoverErr, true)
  326. return
  327. }
  328. if !failoverErr.ShouldRetryNextAccount() {
  329. h.handleFailoverExhausted(c, failoverErr, false)
  330. return
  331. }
  332. if endpoint.IsVideoLookupRequest() {
  333. h.handleFailoverExhausted(c, failoverErr, false)
  334. return
  335. }
  336. if failoverErr.RetryableOnSameAccount {
  337. retryLimit := effectiveSameAccountRetryLimit(failoverErr, account)
  338. if sameAccountRetryAllowed(failoverErr, sameAccountRetryCount[account.ID], retryLimit) {
  339. sameAccountRetryCount[account.ID]++
  340. retryDelay := sameAccountRetryDelayFor(failoverErr, sameAccountRetryCount[account.ID])
  341. reqLog.Warn("grok_media.pool_mode_same_account_retry",
  342. zap.Int64("account_id", account.ID),
  343. zap.Int("upstream_status", failoverErr.StatusCode),
  344. zap.Int("retry_limit", retryLimit),
  345. zap.Int("retry_count", sameAccountRetryCount[account.ID]),
  346. zap.Duration("retry_delay", retryDelay),
  347. )
  348. select {
  349. case <-requestCtx.Done():
  350. return
  351. case <-time.After(retryDelay):
  352. }
  353. continue
  354. }
  355. }
  356. h.gatewayService.RecordOpenAIAccountSwitch()
  357. failedAccountIDs[account.ID] = struct{}{}
  358. lastFailoverErr = failoverErr
  359. if switchCount >= maxAccountSwitches {
  360. h.handleFailoverExhausted(c, failoverErr, false)
  361. return
  362. }
  363. switchCount++
  364. if h.gatewayService.ShouldStopOpenAIOAuth429Failover(account, failoverErr.StatusCode, switchCount, &oauth429FailoverState) {
  365. h.handleFailoverExhausted(c, failoverErr, false)
  366. return
  367. }
  368. reqLog.Warn("grok_media.upstream_failover_switching",
  369. zap.Int64("account_id", account.ID),
  370. zap.Int("upstream_status", failoverErr.StatusCode),
  371. zap.Int("switch_count", switchCount),
  372. zap.Int("max_switches", maxAccountSwitches),
  373. )
  374. continue
  375. }
  376. h.gatewayService.ReportOpenAIAccountScheduleResult(account, grokMediaScheduleModel(account, routingModel, nil), false, nil)
  377. if !service.IsResponseCommitted(c) && c.Writer.Size() == writerSizeBeforeForward {
  378. h.errorResponse(c, http.StatusBadGateway, "upstream_error", "Upstream request failed")
  379. }
  380. reqLog.Warn("grok_media.forward_failed",
  381. zap.Int64("account_id", account.ID),
  382. zap.Error(err),
  383. )
  384. return
  385. }
  386. h.gatewayService.ReportOpenAIAccountScheduleResult(account, grokMediaScheduleModel(account, routingModel, result), true, nil)
  387. if isGrokVideoCreateEndpoint(endpoint) && strings.TrimSpace(result.ResponseID) != "" {
  388. if err := h.gatewayService.BindGrokMediaVideoRequestAccount(
  389. requestCtx, apiKey.GroupID, result.ResponseID, subject.UserID, apiKey.ID, account.ID,
  390. ); err != nil {
  391. reqLog.Warn("grok_media.bind_video_request_account_failed",
  392. zap.Int64("account_id", account.ID),
  393. zap.String("request_id", result.ResponseID),
  394. zap.Error(err),
  395. )
  396. }
  397. // Defer billing until status polling observes video.url. Persist create-time
  398. // model/duration/resolution so status can still price if upstream omits them.
  399. // Retry once: missing pending causes silent underpricing (status omits resolution).
  400. pending := service.GrokVideoPendingBilling{
  401. Model: requestModel,
  402. BillingModel: firstNonEmptyString(result.BillingModel, requestModel),
  403. UpstreamModel: result.UpstreamModel,
  404. VideoResolution: result.VideoResolution,
  405. VideoDurationSeconds: result.VideoDurationSeconds,
  406. OriginalModel: clientRequestedModel(c, requestModel),
  407. // Wall-clock start for usage duration_ms: create accepted → first done discovery.
  408. CreatedAt: videoCreateStartedAt,
  409. }
  410. if err := h.gatewayService.StoreGrokVideoPendingBilling(requestCtx, result.ResponseID, subject.UserID, apiKey.ID, pending); err != nil {
  411. reqLog.Warn("grok_media.store_video_pending_billing_failed_retrying",
  412. zap.Int64("account_id", account.ID),
  413. zap.String("request_id", result.ResponseID),
  414. zap.Error(err),
  415. )
  416. if err2 := h.gatewayService.StoreGrokVideoPendingBilling(requestCtx, result.ResponseID, subject.UserID, apiKey.ID, pending); err2 != nil {
  417. // Response body may already be committed; completion path will fail-closed
  418. // when pending is still missing and status cannot price duration.
  419. reqLog.Error("grok_media.store_video_pending_billing_failed",
  420. zap.Int64("account_id", account.ID),
  421. zap.String("request_id", result.ResponseID),
  422. zap.Error(err2),
  423. )
  424. }
  425. }
  426. }
  427. // Status poll OR content download can observe official done+video.url.
  428. // Both paths share the same claim key so the customer is charged once.
  429. if endpoint == service.GrokMediaEndpointVideoStatus || endpoint == service.GrokMediaEndpointVideoContent {
  430. taskID := strings.TrimSpace(requestID)
  431. if billResult := prepareGrokVideoCompletionBilling(requestCtx, h, reqLog, apiKey, subject, taskID, result); billResult != nil {
  432. recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, billResult, billResult.Model, body, taskID)
  433. }
  434. } else if shouldRecordGrokMediaUsage(endpoint, requestModel, result) {
  435. recordGrokMediaUsage(c, h, reqLog, apiKey, subject, subscription, account, result, requestModel, body, requestID)
  436. }
  437. reqLog.Debug("grok_media.request_completed",
  438. zap.Int64("account_id", account.ID),
  439. zap.Int("switch_count", switchCount),
  440. )
  441. return
  442. }
  443. }
  444. func (h *OpenAIGatewayHandler) ensureGrokMediaAccountEligibility(ctx context.Context, account *service.Account) (bool, string, error) {
  445. if account == nil {
  446. return false, "missing_account", errors.New("grok media account is required")
  447. }
  448. eligible, reason := account.GrokMediaGenerationEligibility()
  449. if eligible || reason != "billing_unobserved" {
  450. return eligible, reason, nil
  451. }
  452. if h == nil || h.grokMediaEligibilityProber == nil {
  453. return false, "billing_probe_unavailable", errors.New("grok media eligibility probe is not configured")
  454. }
  455. return h.grokMediaEligibilityProber.ProbeMediaEligibility(ctx, account.ID)
  456. }
  457. func grokMediaRequiredCapability(endpoint service.GrokMediaEndpoint) service.OpenAIEndpointCapability {
  458. if endpoint.IsGenerationRequest() {
  459. return service.OpenAIEndpointCapabilityGrokMediaGeneration
  460. }
  461. return ""
  462. }
  463. func grokMediaScheduleModel(account *service.Account, routingModel string, result *service.OpenAIForwardResult) string {
  464. if result != nil && strings.TrimSpace(result.UpstreamModel) != "" {
  465. return result.UpstreamModel
  466. }
  467. if account == nil {
  468. return strings.TrimSpace(routingModel)
  469. }
  470. return account.GetMappedModel(routingModel)
  471. }
  472. func isGrokVideoCreateEndpoint(endpoint service.GrokMediaEndpoint) bool {
  473. switch endpoint {
  474. case service.GrokMediaEndpointVideosGenerations,
  475. service.GrokMediaEndpointVideosEdits,
  476. service.GrokMediaEndpointVideosExtensions:
  477. return true
  478. default:
  479. return false
  480. }
  481. }
  482. // shouldRecordGrokMediaUsage gates usage writes for immediate (image) generation.
  483. // Async video create never bills here — status polling does on official
  484. // status=done with video.url (docs.x.ai Video Generation).
  485. // Status/content polls, empty model, and failed generations with zero billable
  486. // image units never bill via this helper.
  487. func shouldRecordGrokMediaUsage(endpoint service.GrokMediaEndpoint, requestModel string, result *service.OpenAIForwardResult) bool {
  488. if result == nil {
  489. return false
  490. }
  491. if isGrokVideoCreateEndpoint(endpoint) || endpoint.IsVideoLookupRequest() {
  492. return false
  493. }
  494. if !endpoint.IsGenerationRequest() || strings.TrimSpace(requestModel) == "" {
  495. return false
  496. }
  497. return result.ImageCount > 0
  498. }
  499. // prepareGrokVideoCompletionBilling claims one-shot billing for official done+video.url
  500. // observations (status poll or content download). Duration/model prefer status body;
  501. // resolution uses create-time request (status response does not document resolution).
  502. func prepareGrokVideoCompletionBilling(
  503. ctx context.Context,
  504. h *OpenAIGatewayHandler,
  505. reqLog *zap.Logger,
  506. apiKey *service.APIKey,
  507. subject middleware2.AuthSubject,
  508. taskRequestID string,
  509. statusResult *service.OpenAIForwardResult,
  510. ) *service.OpenAIForwardResult {
  511. if h == nil || h.gatewayService == nil || apiKey == nil || statusResult == nil {
  512. return nil
  513. }
  514. // Forward already set VideoCount only when status=done && video.url (official).
  515. if statusResult.VideoCount <= 0 {
  516. return nil
  517. }
  518. taskRequestID = strings.TrimSpace(firstNonEmptyString(taskRequestID, statusResult.ResponseID))
  519. if taskRequestID == "" {
  520. return nil
  521. }
  522. // Load create-time snapshot before claim so we can fail-closed without burning the claim
  523. // when Redis lost pending and status cannot price the job.
  524. pending, loadErr := h.gatewayService.LoadGrokVideoPendingBilling(ctx, taskRequestID, subject.UserID, apiKey.ID)
  525. if loadErr != nil {
  526. reqLog.Warn("grok_media.video_pending_billing_load_failed", zap.String("request_id", taskRequestID), zap.Error(loadErr))
  527. }
  528. if pending == nil {
  529. // Status omits resolution; without pending we would silently default to 480p and underbill.
  530. // Allow billing only when official status carries duration (still may default resolution).
  531. if statusResult.VideoDurationSeconds <= 0 {
  532. reqLog.Error("grok_media.video_billing_skipped_missing_pending",
  533. zap.String("request_id", taskRequestID),
  534. zap.String("reason", "no create-time snapshot and status has no video.duration"),
  535. )
  536. return nil
  537. }
  538. reqLog.Error("grok_media.video_billing_without_pending",
  539. zap.String("request_id", taskRequestID),
  540. zap.Int("status_duration_seconds", statusResult.VideoDurationSeconds),
  541. zap.String("note", "resolution falls back to default 480p; investigate pending store failures"),
  542. )
  543. }
  544. claimed, err := h.gatewayService.ClaimGrokVideoBilling(ctx, taskRequestID, subject.UserID, apiKey.ID)
  545. if err != nil {
  546. reqLog.Warn("grok_media.video_billing_claim_failed", zap.String("request_id", taskRequestID), zap.Error(err))
  547. return nil
  548. }
  549. if !claimed {
  550. reqLog.Debug("grok_media.video_billing_already_claimed", zap.String("request_id", taskRequestID))
  551. return nil
  552. }
  553. // Re-merge with pending: resolution is request-only; model/duration fill gaps.
  554. merged := *statusResult
  555. if pending != nil {
  556. if strings.TrimSpace(merged.Model) == "" {
  557. merged.Model = firstNonEmptyString(pending.BillingModel, pending.Model, pending.OriginalModel)
  558. }
  559. if strings.TrimSpace(merged.BillingModel) == "" {
  560. merged.BillingModel = firstNonEmptyString(pending.BillingModel, pending.Model, merged.Model)
  561. }
  562. if strings.TrimSpace(merged.UpstreamModel) == "" {
  563. merged.UpstreamModel = pending.UpstreamModel
  564. }
  565. // Official status omits resolution — always prefer create request.
  566. if strings.TrimSpace(pending.VideoResolution) != "" {
  567. merged.VideoResolution = pending.VideoResolution
  568. }
  569. if merged.VideoDurationSeconds <= 0 {
  570. merged.VideoDurationSeconds = pending.VideoDurationSeconds
  571. }
  572. if strings.TrimSpace(merged.ResponseID) == "" {
  573. merged.ResponseID = taskRequestID
  574. }
  575. }
  576. if strings.TrimSpace(merged.Model) == "" {
  577. merged.Model = "grok-imagine-video"
  578. }
  579. if strings.TrimSpace(merged.BillingModel) == "" {
  580. merged.BillingModel = merged.Model
  581. }
  582. // Always force durable task id so usage_billing_dedup survives multi-poll +
  583. // context-local request ids (do not prefer empty-only fill).
  584. merged.RequestID = service.StableGrokVideoBillingRequestID(firstNonEmptyString(merged.ResponseID, taskRequestID))
  585. merged.ResponseID = firstNonEmptyString(merged.ResponseID, taskRequestID)
  586. merged.VideoCount = 1
  587. // Pure video: do not keep legacy ImageCount (avoids image-path heuristics).
  588. merged.ImageCount = 0
  589. // Official default resolution is 480p when the create request omitted it.
  590. merged.VideoResolution = service.NormalizeVideoBillingResolutionOrDefault(merged.VideoResolution)
  591. // Official default duration is 8s when neither status nor create provided it.
  592. merged.VideoDurationSeconds = service.NormalizeVideoBillingDurationSecondsOrDefault(merged.VideoDurationSeconds)
  593. // E2E latency for async video: create accept → this discovery of done+url.
  594. // Bill on discovery (status/content), not after further client polls; duration
  595. // must not be only the single discovery hop (~hundreds of ms).
  596. if pending != nil {
  597. if e2e := service.GrokVideoE2EDuration(pending.CreatedAt, time.Now()); e2e > 0 {
  598. merged.Duration = e2e
  599. }
  600. }
  601. return &merged
  602. }
  603. func firstNonEmptyString(values ...string) string {
  604. for _, v := range values {
  605. if s := strings.TrimSpace(v); s != "" {
  606. return s
  607. }
  608. }
  609. return ""
  610. }
  611. func recordGrokMediaUsage(
  612. c *gin.Context,
  613. h *OpenAIGatewayHandler,
  614. reqLog *zap.Logger,
  615. apiKey *service.APIKey,
  616. subject middleware2.AuthSubject,
  617. subscription *service.UserSubscription,
  618. account *service.Account,
  619. result *service.OpenAIForwardResult,
  620. requestModel string,
  621. body []byte,
  622. requestID string,
  623. ) {
  624. userAgent := c.GetHeader("User-Agent")
  625. clientIP := ip.GetClientIP(c)
  626. sessionID := service.ExtractClientSessionID(c)
  627. payloadForHash := body
  628. if len(payloadForHash) == 0 && strings.TrimSpace(requestID) != "" {
  629. payloadForHash = []byte(requestID)
  630. }
  631. inboundEndpoint := GetInboundEndpoint(c)
  632. upstreamEndpoint := GetUpstreamEndpoint(c, account.Platform)
  633. quotaPlatform := service.QuotaPlatform(c.Request.Context(), apiKey)
  634. // OriginalModel 记录客户端请求的模型:composite 分组下 body 已被改写为具体模型,
  635. // 公开别名需从 context 取回,与其他端点的用量归因口径一致(计费不受影响:
  636. // BillingModelSource 为空不会触发来源覆盖)。
  637. channelUsageFields := service.ChannelUsageFields{
  638. OriginalModel: clientRequestedModel(c, requestModel),
  639. ChannelMappedModel: requestModel,
  640. }
  641. // Async video: force durable task request id and release claim if billing fails.
  642. videoTaskID := ""
  643. if result != nil && result.VideoCount > 0 {
  644. videoTaskID = strings.TrimSpace(firstNonEmptyString(requestID, result.ResponseID))
  645. if stable := service.StableGrokVideoBillingRequestID(firstNonEmptyString(result.ResponseID, requestID)); stable != "" {
  646. result.RequestID = stable
  647. }
  648. // Prefer task id hash for payload fingerprint stability across status/content.
  649. if len(body) == 0 && videoTaskID != "" {
  650. payloadForHash = []byte(videoTaskID)
  651. }
  652. }
  653. h.submitOpenAIUsageRecordTask(c.Request.Context(), result, func(ctx context.Context) {
  654. if err := h.gatewayService.RecordUsage(ctx, &service.OpenAIRecordUsageInput{
  655. Result: result,
  656. APIKey: apiKey,
  657. User: apiKey.User,
  658. Account: account,
  659. Subscription: subscription,
  660. InboundEndpoint: inboundEndpoint,
  661. UpstreamEndpoint: upstreamEndpoint,
  662. UserAgent: userAgent,
  663. IPAddress: clientIP,
  664. RequestPayloadHash: service.HashUsageRequestPayload(payloadForHash),
  665. APIKeyService: h.apiKeyService,
  666. QuotaPlatform: quotaPlatform,
  667. SessionID: sessionID,
  668. ChannelUsageFields: channelUsageFields,
  669. }); err != nil {
  670. if videoTaskID != "" {
  671. if releaseErr := h.gatewayService.ReleaseGrokVideoBilling(ctx, videoTaskID, subject.UserID, apiKey.ID); releaseErr != nil {
  672. reqLog.Warn("grok_media.video_billing_claim_release_failed",
  673. zap.String("request_id", videoTaskID),
  674. zap.Error(releaseErr),
  675. )
  676. }
  677. }
  678. logger.L().With(
  679. zap.String("component", "handler.openai_gateway.grok_media"),
  680. zap.Int64("user_id", subject.UserID),
  681. zap.Int64("api_key_id", apiKey.ID),
  682. zap.Any("group_id", apiKey.GroupID),
  683. zap.String("model", requestModel),
  684. zap.Int64("account_id", account.ID),
  685. ).Error("grok_media.record_usage_failed", zap.Error(err))
  686. reqLog.Debug("grok_media.record_usage_failed", zap.Error(err))
  687. }
  688. })
  689. }